Compare commits
125 Commits
feat/47-ph
...
fix/126-li
| Author | SHA1 | Date | |
|---|---|---|---|
|
98d095099a
|
|||
| 1773df34ad | |||
|
687324745f
|
|||
|
66632db578
|
|||
| abac5cb303 | |||
|
daeab6be12
|
|||
|
2a366a1c17
|
|||
|
4c3a2735b0
|
|||
| 7454ec7744 | |||
|
d3ef9513f8
|
|||
|
ec7a5b750a
|
|||
|
7f2fcc3527
|
|||
|
9ac1a1586c
|
|||
|
05b6c70a3a
|
|||
| 9c037c5a0b | |||
|
cb6ddbbd98
|
|||
| e824ea3274 | |||
|
9139b07d71
|
|||
|
4762537225
|
|||
| 9b635bfe5c | |||
|
28aa06b768
|
|||
| fc95faaed6 | |||
|
2df829ed73
|
|||
| f611e3479b | |||
|
41ec05e40c
|
|||
|
3242cce39b
|
|||
| 35041f11de | |||
|
40fecb2cf7
|
|||
| 8d50918a1a | |||
|
29a7054c23
|
|||
|
7489a49d24
|
|||
| 341f97f705 | |||
|
fc56eafbf7
|
|||
|
a157c8fa27
|
|||
|
201cc54a7b
|
|||
| 781d92d9a3 | |||
|
9bf13f09dd
|
|||
|
63fb8c6e76
|
|||
|
a1426f177c
|
|||
| a9f5ac7ab4 | |||
|
0126144a93
|
|||
|
869033d08e
|
|||
|
7c5011e5d3
|
|||
|
d1366614f9
|
|||
| 1884535542 | |||
|
140f9124ee
|
|||
| 3b77dbfa32 | |||
|
842caf5d4d
|
|||
| a488ade675 | |||
|
d472b6428a
|
|||
| 0488768afc | |||
|
209bf58207
|
|||
| 1ea7425585 | |||
|
a415454962
|
|||
| 06a36566d1 | |||
| 4cb52e3144 | |||
|
afc1f7a706
|
|||
|
6f956dfda3
|
|||
|
6e0f15c888
|
|||
| 66eb9f558f | |||
|
f96a2e7ed3
|
|||
| b17b555a3d | |||
|
13daf95514
|
|||
| 319b01e0b2 | |||
|
6731adca51
|
|||
|
7e11a7688c
|
|||
|
5600575ba2
|
|||
|
bc7476bf1b
|
|||
|
5a8f6bc7b3
|
|||
|
452d7d9b3d
|
|||
|
21eb211d6a
|
|||
|
508b326bf7
|
|||
|
0de99a8cc7
|
|||
|
f4117224fc
|
|||
|
ce29e0c171
|
|||
|
1bf3348c8c
|
|||
|
7c12b9ea98
|
|||
|
c596519dbd
|
|||
|
a6b1fdc33d
|
|||
|
8dd82776f1
|
|||
|
8600d4fbf2
|
|||
|
7a6f252fe0
|
|||
|
bb0d1e51b8
|
|||
|
2348cc2234
|
|||
|
f2ba12bbc5
|
|||
|
a9d7382be8
|
|||
|
d94c62c143
|
|||
|
cb9e7c7c2e
|
|||
|
2604b9f134
|
|||
|
178e3092d5
|
|||
|
46befde4cd
|
|||
|
cf87e156c5
|
|||
|
79073170ec
|
|||
|
71106afaf1
|
|||
|
d2dcdd6ebb
|
|||
|
222c2a6116
|
|||
|
1115bb0942
|
|||
|
63f578cb15
|
|||
|
76c90fa993
|
|||
|
7984d27553
|
|||
|
43ffffdccb
|
|||
|
5fd7736abd
|
|||
|
03fd4960c3
|
|||
|
5ed6bc3390
|
|||
|
cabec1d08a
|
|||
|
881fc85a4c
|
|||
|
b2ed20b55a
|
|||
|
bee27e9b9c
|
|||
|
87d9c291ce
|
|||
|
d4742467e0
|
|||
|
e7f7e376fc
|
|||
|
3b9a6e37f6
|
|||
|
526b662c5e
|
|||
|
db7e373b90
|
|||
|
5c1623a817
|
|||
|
3b60dd7a31
|
|||
|
4feaaf1cfb
|
|||
|
057bc71e80
|
|||
|
dd31c3cd49
|
|||
|
c83f1eb98c
|
|||
|
a60c9f1075
|
|||
|
b2bd86bfa5
|
|||
|
cdf87284af
|
|||
|
4f16b8c541
|
|||
|
486d7e9a8f
|
@@ -66,6 +66,7 @@ jobs:
|
|||||||
build_cortex: ${{ steps.changes.outputs.build_cortex }}
|
build_cortex: ${{ steps.changes.outputs.build_cortex }}
|
||||||
build_neuron: ${{ steps.changes.outputs.build_neuron }}
|
build_neuron: ${{ steps.changes.outputs.build_neuron }}
|
||||||
build_bench: ${{ steps.changes.outputs.build_bench }}
|
build_bench: ${{ steps.changes.outputs.build_bench }}
|
||||||
|
build_upstream: ${{ steps.changes.outputs.build_upstream }}
|
||||||
check_rust: ${{ steps.changes.outputs.check_rust }}
|
check_rust: ${{ steps.changes.outputs.check_rust }}
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
@@ -86,7 +87,14 @@ jobs:
|
|||||||
# rpmvercmp ranks digit-prefixed segments above alpha ones.
|
# rpmvercmp ranks digit-prefixed segments above alpha ones.
|
||||||
# The SHA stays only as a debug identifier; sort order is
|
# The SHA stays only as a debug identifier; sort order is
|
||||||
# decided entirely by the timestamp.
|
# decided entirely by the timestamp.
|
||||||
COMMIT_TIMESTAMP=$(git log -1 --format=%cd --date=format:%Y%m%d%H%M%S HEAD)
|
# format-local + TZ=UTC0 renders the stamp in UTC regardless
|
||||||
|
# of the committer's recorded timezone. Plain `format:` uses
|
||||||
|
# each commit's own TZ offset — local commits (UTC+3) and
|
||||||
|
# Gitea server-side merge commits (UTC) interleaved
|
||||||
|
# non-monotonically, letting an older build out-rank newer
|
||||||
|
# ones in RPM EVR comparison (the 2026-07-01 "235959" stamp
|
||||||
|
# that froze the fleet on a stale build).
|
||||||
|
COMMIT_TIMESTAMP=$(TZ=UTC0 git log -1 --format=%cd --date=format-local:%Y%m%d%H%M%S HEAD)
|
||||||
RELEASE="0.1.${COMMIT_TIMESTAMP}.git${SHORT_SHA}"
|
RELEASE="0.1.${COMMIT_TIMESTAMP}.git${SHORT_SHA}"
|
||||||
echo "version=${VERSION}" >> "$GITHUB_OUTPUT"
|
echo "version=${VERSION}" >> "$GITHUB_OUTPUT"
|
||||||
echo "release=${RELEASE}" >> "$GITHUB_OUTPUT"
|
echo "release=${RELEASE}" >> "$GITHUB_OUTPUT"
|
||||||
@@ -104,6 +112,7 @@ jobs:
|
|||||||
BUILD_CORTEX=true
|
BUILD_CORTEX=true
|
||||||
BUILD_NEURON=true
|
BUILD_NEURON=true
|
||||||
BUILD_BENCH=true
|
BUILD_BENCH=true
|
||||||
|
BUILD_UPSTREAM=true
|
||||||
CHECK_RUST=true
|
CHECK_RUST=true
|
||||||
|
|
||||||
if [ "${GITHUB_EVENT_NAME}" = "push" ]; then
|
if [ "${GITHUB_EVENT_NAME}" = "push" ]; then
|
||||||
@@ -149,6 +158,7 @@ jobs:
|
|||||||
NEURON_RE='^crates/neuron/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-neuron-prerelease\.spec$|^data/neuron|^neuron\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
NEURON_RE='^crates/neuron/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-neuron-prerelease\.spec$|^data/neuron|^neuron\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
||||||
CORTEX_RE='^crates/cortex-gateway/|^crates/cortex-cli/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/cortex-prerelease\.spec$|^data/cortex|^cortex\.example\.toml$|^models\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
CORTEX_RE='^crates/cortex-gateway/|^crates/cortex-cli/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/cortex-prerelease\.spec$|^data/cortex|^cortex\.example\.toml$|^models\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
||||||
BENCH_RE='^crates/helexa-bench/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-bench-prerelease\.spec$|^data/helexa-bench|^helexa-bench\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
BENCH_RE='^crates/helexa-bench/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-bench-prerelease\.spec$|^data/helexa-bench|^helexa-bench\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
||||||
|
UPSTREAM_RE='^crates/helexa-upstream/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-upstream-prerelease\.spec$|^data/helexa-upstream|^helexa-upstream\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
|
||||||
# Any Rust change (incl. crates not packaged here, e.g.
|
# Any Rust change (incl. crates not packaged here, e.g.
|
||||||
# helexa-acp) still needs lint+test on main.
|
# helexa-acp) still needs lint+test on main.
|
||||||
RUST_RE='\.rs$|^crates/|Cargo\.toml$|^Cargo\.lock$'
|
RUST_RE='\.rs$|^crates/|Cargo\.toml$|^Cargo\.lock$'
|
||||||
@@ -156,10 +166,12 @@ jobs:
|
|||||||
CORTEX_BASE=$(base_for cortex)
|
CORTEX_BASE=$(base_for cortex)
|
||||||
NEURON_BASE=$(base_for helexa-neuron-blackwell)
|
NEURON_BASE=$(base_for helexa-neuron-blackwell)
|
||||||
BENCH_BASE=$(base_for helexa-bench)
|
BENCH_BASE=$(base_for helexa-bench)
|
||||||
|
UPSTREAM_BASE=$(base_for helexa-upstream)
|
||||||
BUILD_CORTEX=$(decide "$CORTEX_BASE" "$CORTEX_RE")
|
BUILD_CORTEX=$(decide "$CORTEX_BASE" "$CORTEX_RE")
|
||||||
BUILD_NEURON=$(decide "$NEURON_BASE" "$NEURON_RE")
|
BUILD_NEURON=$(decide "$NEURON_BASE" "$NEURON_RE")
|
||||||
BUILD_BENCH=$(decide "$BENCH_BASE" "$BENCH_RE")
|
BUILD_BENCH=$(decide "$BENCH_BASE" "$BENCH_RE")
|
||||||
if [ "$BUILD_CORTEX" = "true" ] || [ "$BUILD_NEURON" = "true" ] || [ "$BUILD_BENCH" = "true" ]; then
|
BUILD_UPSTREAM=$(decide "$UPSTREAM_BASE" "$UPSTREAM_RE")
|
||||||
|
if [ "$BUILD_CORTEX" = "true" ] || [ "$BUILD_NEURON" = "true" ] || [ "$BUILD_BENCH" = "true" ] || [ "$BUILD_UPSTREAM" = "true" ]; then
|
||||||
CHECK_RUST=true
|
CHECK_RUST=true
|
||||||
else
|
else
|
||||||
CHECK_RUST=$(decide "$CORTEX_BASE" "$RUST_RE")
|
CHECK_RUST=$(decide "$CORTEX_BASE" "$RUST_RE")
|
||||||
@@ -170,8 +182,9 @@ jobs:
|
|||||||
echo "build_cortex=${BUILD_CORTEX}" >> "$GITHUB_OUTPUT"
|
echo "build_cortex=${BUILD_CORTEX}" >> "$GITHUB_OUTPUT"
|
||||||
echo "build_neuron=${BUILD_NEURON}" >> "$GITHUB_OUTPUT"
|
echo "build_neuron=${BUILD_NEURON}" >> "$GITHUB_OUTPUT"
|
||||||
echo "build_bench=${BUILD_BENCH}" >> "$GITHUB_OUTPUT"
|
echo "build_bench=${BUILD_BENCH}" >> "$GITHUB_OUTPUT"
|
||||||
|
echo "build_upstream=${BUILD_UPSTREAM}" >> "$GITHUB_OUTPUT"
|
||||||
echo "check_rust=${CHECK_RUST}" >> "$GITHUB_OUTPUT"
|
echo "check_rust=${CHECK_RUST}" >> "$GITHUB_OUTPUT"
|
||||||
echo "### change detection: build_cortex=${BUILD_CORTEX} build_neuron=${BUILD_NEURON} build_bench=${BUILD_BENCH} check_rust=${CHECK_RUST}"
|
echo "### change detection: build_cortex=${BUILD_CORTEX} build_neuron=${BUILD_NEURON} build_bench=${BUILD_BENCH} build_upstream=${BUILD_UPSTREAM} check_rust=${CHECK_RUST}"
|
||||||
|
|
||||||
# fmt + clippy + test moved here from ci.yml for main pushes so the
|
# fmt + clippy + test moved here from ci.yml for main pushes so the
|
||||||
# two workflows stop queueing against each other (ci.yml's checks
|
# two workflows stop queueing against each other (ci.yml's checks
|
||||||
@@ -303,6 +316,45 @@ jobs:
|
|||||||
path: artifacts/helexa-bench
|
path: artifacts/helexa-bench
|
||||||
retention-days: 1
|
retention-days: 1
|
||||||
|
|
||||||
|
build-upstream:
|
||||||
|
name: Build helexa-upstream binary
|
||||||
|
timeout-minutes: 25
|
||||||
|
needs: prepare
|
||||||
|
if: needs.prepare.outputs.build_upstream == 'true'
|
||||||
|
# Pure-Rust, non-CUDA binary — same runner as cortex/bench.
|
||||||
|
runs-on: rust
|
||||||
|
env:
|
||||||
|
RUSTC_WRAPPER: sccache
|
||||||
|
SCCACHE_BUCKET: sccache
|
||||||
|
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
|
||||||
|
SCCACHE_REGION: auto
|
||||||
|
SCCACHE_S3_USE_SSL: "false"
|
||||||
|
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
|
||||||
|
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
|
||||||
|
# helexa-upstream uses the sqlx runtime query API (no compile-time
|
||||||
|
# query macros), so it builds without a database or a .sqlx cache.
|
||||||
|
# Set OFFLINE defensively so a stray macro can never reach for a DB.
|
||||||
|
SQLX_OFFLINE: "true"
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
with:
|
||||||
|
ref: ${{ inputs.ref }}
|
||||||
|
|
||||||
|
- name: Build helexa-upstream (release, sccache escalation)
|
||||||
|
run: script/ci-cargo-escalate.sh cargo build --release -p helexa-upstream
|
||||||
|
|
||||||
|
- name: Stage binary
|
||||||
|
run: |
|
||||||
|
mkdir --parents artifacts
|
||||||
|
cp target/release/helexa-upstream artifacts/helexa-upstream
|
||||||
|
./artifacts/helexa-upstream --version || true
|
||||||
|
|
||||||
|
- uses: actions/upload-artifact@v3
|
||||||
|
with:
|
||||||
|
name: upstream-fc43
|
||||||
|
path: artifacts/helexa-upstream
|
||||||
|
retention-days: 1
|
||||||
|
|
||||||
build-neuron:
|
build-neuron:
|
||||||
name: Build neuron-${{ matrix.flavour }}
|
name: Build neuron-${{ matrix.flavour }}
|
||||||
timeout-minutes: 35
|
timeout-minutes: 35
|
||||||
@@ -332,7 +384,10 @@ jobs:
|
|||||||
cuda_home: /usr/local/cuda-13.0
|
cuda_home: /usr/local/cuda-13.0
|
||||||
build_jobs: 8
|
build_jobs: 8
|
||||||
nvcc_threads: 4
|
nvcc_threads: 4
|
||||||
cargo_features: "cuda cudnn"
|
# flash-attn on blackwell first (#95): beast carries the
|
||||||
|
# agentic prefill pain; ada/ampere follow once the win is
|
||||||
|
# measured. NEURON_FLASH_ATTN=0 is the runtime rollback.
|
||||||
|
cargo_features: "cuda cudnn flash-attn"
|
||||||
runs-on: ${{ matrix.runner }}
|
runs-on: ${{ matrix.runner }}
|
||||||
env:
|
env:
|
||||||
SCCACHE_BUCKET: sccache
|
SCCACHE_BUCKET: sccache
|
||||||
@@ -459,6 +514,44 @@ jobs:
|
|||||||
path: ~/rpmbuild/RPMS/x86_64/*.rpm
|
path: ~/rpmbuild/RPMS/x86_64/*.rpm
|
||||||
retention-days: 7
|
retention-days: 7
|
||||||
|
|
||||||
|
package-upstream:
|
||||||
|
name: Package helexa-upstream RPM
|
||||||
|
timeout-minutes: 20
|
||||||
|
needs: [prepare, build-upstream]
|
||||||
|
runs-on: rpm
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
with:
|
||||||
|
ref: ${{ inputs.ref }}
|
||||||
|
|
||||||
|
- uses: actions/download-artifact@v3
|
||||||
|
with:
|
||||||
|
name: upstream-fc43
|
||||||
|
path: artifacts/
|
||||||
|
|
||||||
|
- name: Build RPM
|
||||||
|
run: |
|
||||||
|
set -eux
|
||||||
|
rm -f ~/.rpmmacros
|
||||||
|
rpmdev-setuptree
|
||||||
|
cp artifacts/helexa-upstream ~/rpmbuild/SOURCES/
|
||||||
|
cp data/helexa-upstream.service ~/rpmbuild/SOURCES/
|
||||||
|
cp data/helexa-upstream-sysusers.conf ~/rpmbuild/SOURCES/
|
||||||
|
cp data/helexa-upstream-firewalld.xml ~/rpmbuild/SOURCES/
|
||||||
|
cp helexa-upstream.example.toml ~/rpmbuild/SOURCES/
|
||||||
|
cp LICENSE ~/rpmbuild/SOURCES/
|
||||||
|
rpmbuild -bb rpm/helexa-upstream-prerelease.spec \
|
||||||
|
--define "upstream_version ${{ needs.prepare.outputs.version }}" \
|
||||||
|
--define "upstream_prerelease ${{ needs.prepare.outputs.release }}" \
|
||||||
|
--undefine dist \
|
||||||
|
--define "dist .fc43"
|
||||||
|
|
||||||
|
- uses: actions/upload-artifact@v3
|
||||||
|
with:
|
||||||
|
name: rpm-upstream-fc43
|
||||||
|
path: ~/rpmbuild/RPMS/x86_64/*.rpm
|
||||||
|
retention-days: 7
|
||||||
|
|
||||||
package-neuron:
|
package-neuron:
|
||||||
name: Package helexa-neuron-${{ matrix.flavour }} RPM
|
name: Package helexa-neuron-${{ matrix.flavour }} RPM
|
||||||
timeout-minutes: 20
|
timeout-minutes: 20
|
||||||
@@ -508,7 +601,7 @@ jobs:
|
|||||||
publish:
|
publish:
|
||||||
name: Publish to rpm.lair.cafe (unstable)
|
name: Publish to rpm.lair.cafe (unstable)
|
||||||
timeout-minutes: 25
|
timeout-minutes: 25
|
||||||
needs: [lint, test, package-cortex, package-neuron, package-bench]
|
needs: [lint, test, package-cortex, package-neuron, package-bench, package-upstream]
|
||||||
# Runs when at least one package was built and nothing failed.
|
# Runs when at least one package was built and nothing failed.
|
||||||
# lint/test may be skipped (docs-only refs never get here because
|
# lint/test may be skipped (docs-only refs never get here because
|
||||||
# no packages build), but a real failure in any blocks the
|
# no packages build), but a real failure in any blocks the
|
||||||
@@ -518,10 +611,11 @@ jobs:
|
|||||||
!cancelled()
|
!cancelled()
|
||||||
&& (needs.lint.result == 'success' || needs.lint.result == 'skipped')
|
&& (needs.lint.result == 'success' || needs.lint.result == 'skipped')
|
||||||
&& (needs.test.result == 'success' || needs.test.result == 'skipped')
|
&& (needs.test.result == 'success' || needs.test.result == 'skipped')
|
||||||
&& (needs.package-cortex.result == 'success' || needs.package-neuron.result == 'success' || needs.package-bench.result == 'success')
|
&& (needs.package-cortex.result == 'success' || needs.package-neuron.result == 'success' || needs.package-bench.result == 'success' || needs.package-upstream.result == 'success')
|
||||||
&& needs.package-cortex.result != 'failure'
|
&& needs.package-cortex.result != 'failure'
|
||||||
&& needs.package-neuron.result != 'failure'
|
&& needs.package-neuron.result != 'failure'
|
||||||
&& needs.package-bench.result != 'failure'
|
&& needs.package-bench.result != 'failure'
|
||||||
|
&& needs.package-upstream.result != 'failure'
|
||||||
}}
|
}}
|
||||||
runs-on: rpm
|
runs-on: rpm
|
||||||
concurrency:
|
concurrency:
|
||||||
|
|||||||
@@ -82,7 +82,9 @@ jobs:
|
|||||||
# see commit history).
|
# see commit history).
|
||||||
cuda-check:
|
cuda-check:
|
||||||
name: CUDA type-check
|
name: CUDA type-check
|
||||||
timeout-minutes: 35
|
# flash-attn kernel compilation dominates the first uncached run;
|
||||||
|
# sccache + the cargo target cache absorb it afterwards.
|
||||||
|
timeout-minutes: 70
|
||||||
runs-on: cuda-13.0
|
runs-on: cuda-13.0
|
||||||
# The workflow-level env sets `RUSTC_WRAPPER: sccache`
|
# The workflow-level env sets `RUSTC_WRAPPER: sccache`
|
||||||
# unconditionally, which hard-fails cargo if the CUDA image
|
# unconditionally, which hard-fails cargo if the CUDA image
|
||||||
@@ -115,7 +117,7 @@ jobs:
|
|||||||
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
|
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
|
||||||
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
|
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
|
||||||
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
|
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
|
||||||
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda --all-targets
|
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda,flash-attn --all-targets
|
||||||
|
|
||||||
srpm-cortex:
|
srpm-cortex:
|
||||||
name: Build cortex SRPM
|
name: Build cortex SRPM
|
||||||
|
|||||||
@@ -274,19 +274,52 @@ jobs:
|
|||||||
|
|
||||||
echo "LLM probe against ${model}"
|
echo "LLM probe against ${model}"
|
||||||
probe_body=$(printf '{"model":"%s","messages":[{"role":"user","content":"Reply with exactly one word: pineapple"}],"max_tokens":512,"temperature":0}' "${model}")
|
probe_body=$(printf '{"model":"%s","messages":[{"role":"user","content":"Reply with exactly one word: pineapple"}],"max_tokens":512,"temperature":0}' "${model}")
|
||||||
resp=$(curl -fsS --max-time 180 -H "content-type: application/json" \
|
# The probe races real traffic: a deploy publishes a new build
|
||||||
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || {
|
# SHA, which is exactly what triggers helexa-bench to re-sweep
|
||||||
echo "FAIL: probe request errored"
|
# every scenario against this neuron — long capability-probe
|
||||||
exit 1
|
# generations can hold the batch-1 model for minutes. Admission
|
||||||
}
|
# control answers concurrent requests with 429/503 +
|
||||||
if printf %s "${resp}" | grep -qi pineapple; then
|
# Retry-After (#53/#63); treat those as "busy, try again", not
|
||||||
echo "LLM probe passed"
|
# deploy failure. Overall budget ~6 min.
|
||||||
else
|
probe_deadline=$(( $(date +%s) + 360 ))
|
||||||
echo "FAIL: probe response missing expected token"
|
attempt=0
|
||||||
printf %s "${resp}" | head -c 2000
|
while :; do
|
||||||
echo
|
attempt=$(( attempt + 1 ))
|
||||||
exit 1
|
hdrs=$(mktemp)
|
||||||
fi
|
resp=$(curl -sS --max-time 180 -D "${hdrs}" \
|
||||||
|
-H "content-type: application/json" \
|
||||||
|
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || resp=""
|
||||||
|
status=$(awk 'toupper($1) ~ /^HTTP/ {code=$2} END {print code}' "${hdrs}")
|
||||||
|
retry_after=$(awk 'tolower($1) == "retry-after:" {print $2+0; exit}' "${hdrs}")
|
||||||
|
rm -f "${hdrs}"
|
||||||
|
if [ "${status}" = "200" ]; then
|
||||||
|
if printf %s "${resp}" | grep -qi pineapple; then
|
||||||
|
echo "LLM probe passed (attempt ${attempt})"
|
||||||
|
break
|
||||||
|
fi
|
||||||
|
echo "FAIL: probe response missing expected token"
|
||||||
|
printf %s "${resp}" | head -c 2000
|
||||||
|
echo
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
case "${status}" in
|
||||||
|
429|503) ;; # busy/queue-full — retryable
|
||||||
|
*)
|
||||||
|
echo "FAIL: probe request errored (HTTP ${status:-none})"
|
||||||
|
printf %s "${resp}" | head -c 2000
|
||||||
|
echo
|
||||||
|
exit 1
|
||||||
|
;;
|
||||||
|
esac
|
||||||
|
if [ "$(date +%s)" -ge "${probe_deadline}" ]; then
|
||||||
|
echo "FAIL: model still busy (HTTP ${status}) after probe budget"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
wait_s=${retry_after:-15}
|
||||||
|
[ "${wait_s}" -ge 5 ] 2>/dev/null || wait_s=15
|
||||||
|
echo "model busy (HTTP ${status}), retrying in ${wait_s}s (attempt ${attempt})"
|
||||||
|
sleep "${wait_s}"
|
||||||
|
done
|
||||||
DEPLOY
|
DEPLOY
|
||||||
|
|
||||||
- name: Ensure firewalld allows helexa-neuron
|
- name: Ensure firewalld allows helexa-neuron
|
||||||
|
|||||||
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1,6 +1,9 @@
|
|||||||
/target
|
/target
|
||||||
/bench/node_modules
|
/bench/node_modules
|
||||||
/bench/dist
|
/bench/dist
|
||||||
|
/helexa.ai/node_modules
|
||||||
|
/helexa.ai/dist
|
||||||
|
helexa.ai/.env.local
|
||||||
*.swp
|
*.swp
|
||||||
*.swo
|
*.swo
|
||||||
.idea/
|
.idea/
|
||||||
|
|||||||
26
CLAUDE.md
26
CLAUDE.md
@@ -185,6 +185,32 @@ Run these locally before pushing. `cargo fmt --all` fixes formatting
|
|||||||
automatically. Clippy warnings must be resolved, not suppressed with
|
automatically. Clippy warnings must be resolved, not suppressed with
|
||||||
`#[allow(...)]` unless there is a clear rationale.
|
`#[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
|
## Environment
|
||||||
|
|
||||||
- Targets Fedora 43 (systemd, SELinux enforcing)
|
- Targets Fedora 43 (systemd, SELinux enforcing)
|
||||||
|
|||||||
874
Cargo.lock
generated
874
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -7,6 +7,9 @@ members = [
|
|||||||
"crates/neuron",
|
"crates/neuron",
|
||||||
"crates/helexa-acp",
|
"crates/helexa-acp",
|
||||||
"crates/helexa-bench",
|
"crates/helexa-bench",
|
||||||
|
"crates/helexa-router",
|
||||||
|
"crates/helexa-stream",
|
||||||
|
"crates/helexa-upstream",
|
||||||
]
|
]
|
||||||
|
|
||||||
[workspace.package]
|
[workspace.package]
|
||||||
|
|||||||
@@ -18,6 +18,35 @@ db_path = "/var/lib/helexa-bench/bench.sqlite"
|
|||||||
[scenarios]
|
[scenarios]
|
||||||
prompt_sizes = [128, 4096]
|
prompt_sizes = [128, 4096]
|
||||||
max_tokens = 256
|
max_tokens = 256
|
||||||
|
# Concurrency / agentic-load scenarios (#89), enabled so the 27B baseline
|
||||||
|
# carries p95-under-concurrency data before the F3 A/B gate (#94) needs
|
||||||
|
# it for comparison. Levels mirror the real a0/hermes/opencode fan-out.
|
||||||
|
concurrency_levels = [2, 4, 8]
|
||||||
|
concurrency_prompt_tokens = 512
|
||||||
|
|
||||||
|
# Capability probes (#91) — the reasoning/planning axis the speed
|
||||||
|
# scenarios miss; scored manually via `helexa-bench score` (O7). Enabled
|
||||||
|
# for the same reason: the F3 gate compares 80B-A3B variants against the
|
||||||
|
# 27B on planning quality, so the 27B needs scored artifacts first.
|
||||||
|
[[scenarios.capability_probes]]
|
||||||
|
name = "rust-plan"
|
||||||
|
max_tokens = 4096
|
||||||
|
prompt = """
|
||||||
|
Write an implementation plan for adding rate limiting to an Axum service.
|
||||||
|
Honor existing conventions, call out trade-offs, and sequence the work.
|
||||||
|
"""
|
||||||
|
|
||||||
|
[[scenarios.capability_probes]]
|
||||||
|
name = "debug-reason"
|
||||||
|
max_tokens = 4096
|
||||||
|
prompt = """
|
||||||
|
A Rust axum server streams SSE responses through a reverse proxy. Clients
|
||||||
|
report that streams stall for exactly 60 seconds and then resume, but only
|
||||||
|
when response chunks are small and infrequent. Curling the backend directly
|
||||||
|
never stalls. List the most likely causes in order of probability, explain
|
||||||
|
the mechanism behind each, and describe the smallest experiment that would
|
||||||
|
confirm or eliminate each cause.
|
||||||
|
"""
|
||||||
|
|
||||||
# Read-only JSON API consumed by the bench UI (hosted separately) and for
|
# Read-only JSON API consumed by the bench UI (hosted separately) and for
|
||||||
# programmatic access. Served alongside the sweep loop.
|
# programmatic access. Served alongside the sweep loop.
|
||||||
|
|||||||
@@ -16,6 +16,14 @@ name = "candle"
|
|||||||
|
|
||||||
[harness.candle]
|
[harness.candle]
|
||||||
|
|
||||||
|
# Batched decode engine (#98): up to 8 concurrent text streams
|
||||||
|
# multiplex through one lockstep (B,1) forward per decode step.
|
||||||
|
# max_in_flight is both the admission bound and the engine's slot
|
||||||
|
# count. NEURON_BATCHING=0 (systemd drop-in) is the kill switch;
|
||||||
|
# removing this section reverts to batch-1 serialization.
|
||||||
|
[harness.candle.admission]
|
||||||
|
max_in_flight = 8
|
||||||
|
|
||||||
[[default_models]]
|
[[default_models]]
|
||||||
model_id = "Qwen/Qwen3.6-27B"
|
model_id = "Qwen/Qwen3.6-27B"
|
||||||
harness = "candle"
|
harness = "candle"
|
||||||
|
|||||||
@@ -90,3 +90,20 @@ account_id = "operator"
|
|||||||
key_id = "infra"
|
key_id = "infra"
|
||||||
# No hard_cap → uncapped operator infra key (own fleet, own use). Still
|
# No hard_cap → uncapped operator infra key (own fleet, own use). Still
|
||||||
# metered for visibility.
|
# metered for visibility.
|
||||||
|
|
||||||
|
# -- Upstream (helexa mesh) entitlements client (#57) --------------------
|
||||||
|
# When enabled, a bearer key NOT found in [[entitlements.keys]] above is
|
||||||
|
# validated against the helexa-upstream authority (mesh accounts), and its
|
||||||
|
# budget is reserved/settled there. Operator-local keys (incl. the infra
|
||||||
|
# key) never leave this process. Fail-closed: if upstream is unreachable a
|
||||||
|
# request is refused (503 + Retry-After), never served un-authorized.
|
||||||
|
# Disabled by default — a standalone operator runs purely local.
|
||||||
|
[upstream]
|
||||||
|
enabled = false
|
||||||
|
# url = "https://upstream.helexa.ai"
|
||||||
|
# Shared client bearer this cortex presents (maps to an operator_id
|
||||||
|
# upstream). Override via CORTEX_UPSTREAM__BEARER in prod.
|
||||||
|
# bearer = "replace-with-operator-client-secret"
|
||||||
|
# timeout_secs = 5
|
||||||
|
# How often to flush served-usage counters to upstream for reconciliation (#58).
|
||||||
|
# served_usage_report_interval_secs = 60
|
||||||
|
|||||||
@@ -22,6 +22,43 @@ pub struct GatewayConfig {
|
|||||||
/// setups keep working until keys are configured.
|
/// setups keep working until keys are configured.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub entitlements: EntitlementsConfig,
|
pub entitlements: EntitlementsConfig,
|
||||||
|
/// helexa-upstream client (#57). When enabled, keys not found in the
|
||||||
|
/// local `[entitlements]` config are validated against the mesh
|
||||||
|
/// authority, and budget is reserved/settled there. Disabled by default
|
||||||
|
/// — a single operator runs purely local.
|
||||||
|
#[serde(default)]
|
||||||
|
pub upstream: UpstreamClientConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[upstream]` — the helexa-upstream authority client (#57). Locally
|
||||||
|
/// unrecognised bearer keys are resolved against `url`'s `/authz/v1` surface
|
||||||
|
/// (mesh accounts); local keys (operator + infra) never leave the process.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct UpstreamClientConfig {
|
||||||
|
/// Enable the upstream fallthrough. Off → purely local entitlements.
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
/// Base URL of helexa-upstream (e.g. "https://upstream.helexa.ai").
|
||||||
|
#[serde(default)]
|
||||||
|
pub url: String,
|
||||||
|
/// Shared client bearer this cortex presents to `/authz/v1` (maps to an
|
||||||
|
/// operator_id upstream). Sent as `Authorization: Bearer <bearer>`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub bearer: String,
|
||||||
|
/// Per-call timeout (seconds) to upstream.
|
||||||
|
#[serde(default = "default_upstream_timeout")]
|
||||||
|
pub timeout_secs: u64,
|
||||||
|
/// How often (seconds) to flush served-usage counters to upstream for
|
||||||
|
/// reconciliation (#58).
|
||||||
|
#[serde(default = "default_served_usage_interval")]
|
||||||
|
pub served_usage_report_interval_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_upstream_timeout() -> u64 {
|
||||||
|
5
|
||||||
|
}
|
||||||
|
fn default_served_usage_interval() -> u64 {
|
||||||
|
60
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `[entitlements]` — the local/static [`crate::entitlements::EntitlementProvider`]
|
/// `[entitlements]` — the local/static [`crate::entitlements::EntitlementProvider`]
|
||||||
@@ -129,6 +166,7 @@ impl Default for GatewayConfig {
|
|||||||
neurons: vec![],
|
neurons: vec![],
|
||||||
models_config: default_models_path(),
|
models_config: default_models_path(),
|
||||||
entitlements: EntitlementsConfig::default(),
|
entitlements: EntitlementsConfig::default(),
|
||||||
|
upstream: UpstreamClientConfig::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -68,6 +68,57 @@ pub struct HealthResponse {
|
|||||||
pub devices: Vec<DeviceHealth>,
|
pub devices: Vec<DeviceHealth>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub activation: ActivationStatus,
|
pub activation: ActivationStatus,
|
||||||
|
/// Per-model admission load (#53): how many requests are running vs.
|
||||||
|
/// queued on each loaded model right now. Cortex's load-aware router
|
||||||
|
/// (#55) reads this to spread traffic across replicas and to propagate
|
||||||
|
/// honest backpressure. `#[serde(default)]` keeps older gateways/neurons
|
||||||
|
/// interoperable (absent → empty → treated as no load info).
|
||||||
|
#[serde(default)]
|
||||||
|
pub models: Vec<ModelLoad>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Live admission load for one loaded model (#53).
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ModelLoad {
|
||||||
|
pub id: String,
|
||||||
|
/// Requests currently running (batch-1 → 0 or 1).
|
||||||
|
pub in_flight: usize,
|
||||||
|
/// Requests waiting in the bounded admission queue.
|
||||||
|
pub queue_depth: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod health_load_tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn health_response_without_models_field_still_deserializes() {
|
||||||
|
// A pre-#53 neuron's /health payload omits `models`; the gateway
|
||||||
|
// must still parse it (serde default → empty).
|
||||||
|
let json = r#"{"uptime_secs":42,"devices":[]}"#;
|
||||||
|
let resp: HealthResponse = serde_json::from_str(json).expect("back-compat parse");
|
||||||
|
assert_eq!(resp.uptime_secs, 42);
|
||||||
|
assert!(resp.models.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn health_response_round_trips_model_load() {
|
||||||
|
let resp = HealthResponse {
|
||||||
|
uptime_secs: 1,
|
||||||
|
devices: vec![],
|
||||||
|
activation: ActivationStatus::default(),
|
||||||
|
models: vec![ModelLoad {
|
||||||
|
id: "Qwen/Qwen3.6-27B".into(),
|
||||||
|
in_flight: 1,
|
||||||
|
queue_depth: 3,
|
||||||
|
}],
|
||||||
|
};
|
||||||
|
let s = serde_json::to_string(&resp).unwrap();
|
||||||
|
let back: HealthResponse = serde_json::from_str(&s).unwrap();
|
||||||
|
assert_eq!(back.models.len(), 1);
|
||||||
|
assert_eq!(back.models[0].in_flight, 1);
|
||||||
|
assert_eq!(back.models[0].queue_depth, 3);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// High-level activation state of the neuron daemon. The HTTP listener
|
/// High-level activation state of the neuron daemon. The HTTP listener
|
||||||
|
|||||||
@@ -23,6 +23,14 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// Internal header carrying the resolved account id from cortex to neuron.
|
||||||
|
/// neuron trusts these over the WireGuard link (#54); cortex **strips** any
|
||||||
|
/// client-supplied copy before stamping the authoritative value, so a client
|
||||||
|
/// can never assert a principal directly.
|
||||||
|
pub const HEADER_ACCOUNT_ID: &str = "x-helexa-account-id";
|
||||||
|
/// Internal header carrying the resolved key id from cortex to neuron.
|
||||||
|
pub const HEADER_KEY_ID: &str = "x-helexa-key-id";
|
||||||
|
|
||||||
/// Who a request is for. Resolved once at the edge from the bearer key and
|
/// Who a request is for. Resolved once at the edge from the bearer key and
|
||||||
/// carried through the request context. `account_id` is the billable owner
|
/// carried through the request context. `account_id` is the billable owner
|
||||||
/// (spendable at any operator, by decision); `key_id` identifies the
|
/// (spendable at any operator, by decision); `key_id` identifies the
|
||||||
@@ -73,12 +81,19 @@ pub struct BudgetSnapshot {
|
|||||||
pub reserved: u64,
|
pub reserved: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Authentication failure — the bearer key could not be resolved. Maps to
|
/// Authentication failure — the bearer key could not be resolved.
|
||||||
/// `401 invalid_api_key` (#49/#63).
|
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
pub enum AuthError {
|
pub enum AuthError {
|
||||||
|
/// The key is genuinely unknown → `401 invalid_api_key` (#49/#63).
|
||||||
#[error("invalid or unknown API key")]
|
#[error("invalid or unknown API key")]
|
||||||
InvalidKey,
|
InvalidKey,
|
||||||
|
/// The authority that could resolve the key is unreachable (e.g. the
|
||||||
|
/// helexa-upstream client failed, #57). Fail **closed** but distinctly:
|
||||||
|
/// a transient outage must surface as `503 service_unavailable` +
|
||||||
|
/// `Retry-After`, never `401` — a real key must not be rejected as
|
||||||
|
/// invalid during an upstream blip.
|
||||||
|
#[error("entitlement authority unavailable; retry in {retry_after_secs}s")]
|
||||||
|
Unavailable { retry_after_secs: u64 },
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Why a reservation was refused. Carries enough for the caller to build the
|
/// Why a reservation was refused. Carries enough for the caller to build the
|
||||||
|
|||||||
@@ -54,10 +54,26 @@ pub struct ModelLimit {
|
|||||||
pub output: usize,
|
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
|
/// `cost` is sourced from the operator's `models.toml` catalogue profile and
|
||||||
/// optional — set when the backend supports a prefix-cache discount tier.
|
/// 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)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct ModelCost {
|
pub struct ModelCost {
|
||||||
/// USD per 1M input (prompt) tokens.
|
/// 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.
|
/// `None` when neither the catalogue nor the loaded model can provide it.
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub limit: Option<ModelLimit>,
|
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")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub cost: Option<ModelCost>,
|
pub cost: Option<ModelCost>,
|
||||||
/// `true` when the model's tokenizer contains recognised tool-call
|
/// `true` when the model's tokenizer contains recognised tool-call
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
use crate::discovery::{ActivationStatus, DiscoveryResponse};
|
use crate::discovery::{ActivationStatus, DiscoveryResponse, ModelLoad};
|
||||||
use crate::harness::{ModelCost, ModelLimit};
|
use crate::harness::{ModelCost, ModelLimit};
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
@@ -27,6 +27,17 @@ pub struct NodeState {
|
|||||||
/// to synthesize `Loading` locations so clients see a catalogued
|
/// to synthesize `Loading` locations so clients see a catalogued
|
||||||
/// model that's mid-prewarm as "loading", not "missing".
|
/// model that's mid-prewarm as "loading", not "missing".
|
||||||
pub activation: Option<ActivationStatus>,
|
pub activation: Option<ActivationStatus>,
|
||||||
|
/// Last-seen per-model admission load from this neuron's `/health`
|
||||||
|
/// (#53), keyed by model id. The router (#55) reads it to pick the
|
||||||
|
/// least-busy replica when a model is loaded on more than one neuron.
|
||||||
|
/// Empty until the first /health poll reports load.
|
||||||
|
pub model_load: HashMap<String, ModelLoad>,
|
||||||
|
/// Consecutive failed `/models` polls. The poller marks a node
|
||||||
|
/// unhealthy only once this crosses a threshold, so a single transient
|
||||||
|
/// miss (e.g. a neuron momentarily slow to answer while busy) doesn't
|
||||||
|
/// yank the node — and all its models — out of routing. Reset to 0 on
|
||||||
|
/// any successful poll.
|
||||||
|
pub consecutive_poll_failures: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A model registered on a node, with its runtime status.
|
/// A model registered on a node, with its runtime status.
|
||||||
@@ -125,7 +136,9 @@ pub struct CortexModelEntry {
|
|||||||
/// at load time. `None` when neither source provides it.
|
/// at load time. `None` when neither source provides it.
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub limit: Option<ModelLimit>,
|
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")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub cost: Option<ModelCost>,
|
pub cost: Option<ModelCost>,
|
||||||
/// `true` when any neuron reports this model supports tool calls.
|
/// `true` when any neuron reports this model supports tool calls.
|
||||||
@@ -134,6 +147,39 @@ pub struct CortexModelEntry {
|
|||||||
/// `true` when any neuron reports this model supports reasoning tokens.
|
/// `true` when any neuron reports this model supports reasoning tokens.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub reasoning: bool,
|
pub reasoning: bool,
|
||||||
|
// ── Flat ecosystem context-window fields (issue #78) ──────
|
||||||
|
// Duplicates of `limit` under the flat, vLLM-convention key names
|
||||||
|
// (`max_model_len` et al.) that OpenAI-ecosystem clients (Hermes
|
||||||
|
// Agent, vLLM tooling) probe for — they cannot see `limit.context`.
|
||||||
|
// Additive: `limit` stays the opencode-oriented source of truth.
|
||||||
|
// Derived, never set directly — call [`sync_flat_limit`] after the
|
||||||
|
// final `limit` value is known. Omitted (not `0`) when the window
|
||||||
|
// is unknown; absent-vs-zero is load-bearing, as with `cost`.
|
||||||
|
//
|
||||||
|
// [`sync_flat_limit`]: CortexModelEntry::sync_flat_limit
|
||||||
|
/// Served max-seq-len in tokens — mirrors `limit.context`.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub max_model_len: Option<usize>,
|
||||||
|
/// Usable input budget — mirrors `limit.input` when present.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub max_input_tokens: Option<usize>,
|
||||||
|
/// Maximum generation tokens — mirrors `limit.output`.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub max_output_tokens: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CortexModelEntry {
|
||||||
|
/// Re-derive the flat ecosystem fields (#78) from `limit`.
|
||||||
|
///
|
||||||
|
/// Must run after the final `limit` is known (post merge/tightening),
|
||||||
|
/// immediately before serialization. Fully overwrites: a `None` limit
|
||||||
|
/// clears the flat fields, so stale values can't survive a merge that
|
||||||
|
/// dropped the limit.
|
||||||
|
pub fn sync_flat_limit(&mut self) {
|
||||||
|
self.max_model_len = self.limit.as_ref().map(|l| l.context);
|
||||||
|
self.max_input_tokens = self.limit.as_ref().and_then(|l| l.input);
|
||||||
|
self.max_output_tokens = self.limit.as_ref().map(|l| l.output);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
|||||||
@@ -116,6 +116,23 @@ pub struct Usage {
|
|||||||
/// prompt caching lands (#11); `None` until then.
|
/// prompt caching lands (#11); `None` until then.
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
||||||
|
/// helexa extension (non-OpenAI): server-measured prefill/decode
|
||||||
|
/// timing, so the bench harness can compute true prefill vs decode
|
||||||
|
/// tok/s instead of inferring both from client-side SSE arrival
|
||||||
|
/// (#85). Additive and optional — standard OpenAI clients ignore
|
||||||
|
/// it; cortex forwards usage verbatim so it survives proxying.
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub helexa_timing: Option<HelexaTiming>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// helexa extension carried on [`Usage::helexa_timing`]. Mirrors
|
||||||
|
/// neuron's internal `FinishTiming`. All fields are server-measured;
|
||||||
|
/// `prefill_tokens` is the prefill-rate denominator.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct HelexaTiming {
|
||||||
|
pub prefill_ms: u64,
|
||||||
|
pub decode_ms: u64,
|
||||||
|
pub prefill_tokens: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sub-counts of `Usage::completion_tokens`.
|
/// Sub-counts of `Usage::completion_tokens`.
|
||||||
|
|||||||
@@ -66,14 +66,48 @@ pub struct ResponsesRequest {
|
|||||||
pub extra: Value,
|
pub extra: Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `input` is either a single string or an array of typed items.
|
/// `input` is either a single string or an array of items.
|
||||||
/// `#[serde(untagged)]` so the wire shape `"input": "hi"` and
|
/// `#[serde(untagged)]` so the wire shape `"input": "hi"` and
|
||||||
/// `"input": [{...}]` both deserialize.
|
/// `"input": [{...}]` both deserialize.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
#[serde(untagged)]
|
#[serde(untagged)]
|
||||||
pub enum ResponsesInput {
|
pub enum ResponsesInput {
|
||||||
Text(String),
|
Text(String),
|
||||||
Items(Vec<ResponsesInputItem>),
|
Items(Vec<ResponsesInputElement>),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One element of an `input` array.
|
||||||
|
///
|
||||||
|
/// OpenAI's Responses API accepts three shapes here, and real clients
|
||||||
|
/// use all of them — most notably agent-zero (via litellm), which
|
||||||
|
/// sends the bare "easy message" form. We must tolerate every shape,
|
||||||
|
/// because `input` is an `#[serde(untagged)]` array: a single element
|
||||||
|
/// that matches no variant fails the *entire* request with a 422
|
||||||
|
/// (`did not match any variant of untagged enum ResponsesInput`).
|
||||||
|
///
|
||||||
|
/// 1. [`Self::Typed`] — an item carrying an explicit `"type"`
|
||||||
|
/// discriminant (`message`, `function_call`, `function_call_output`,
|
||||||
|
/// `reasoning`).
|
||||||
|
/// 2. [`Self::EasyMessage`] — a bare `{role, content}` with **no**
|
||||||
|
/// `type` field. This is OpenAI's `EasyInputMessage` and what
|
||||||
|
/// litellm emits for every turn. `content` is optional so an
|
||||||
|
/// assistant turn carrying only tool calls (`content: null`) still
|
||||||
|
/// parses.
|
||||||
|
/// 3. [`Self::Other`] — anything else, captured as raw JSON and
|
||||||
|
/// dropped during translation. This is the forward-compat escape
|
||||||
|
/// hatch that mirrors [`ResponsesRequest::extra`] at the item
|
||||||
|
/// level: an unmodeled item type can never again reject the whole
|
||||||
|
/// request.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
pub enum ResponsesInputElement {
|
||||||
|
Typed(ResponsesInputItem),
|
||||||
|
EasyMessage {
|
||||||
|
role: String,
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
content: Option<ResponsesMessageContent>,
|
||||||
|
},
|
||||||
|
Other(Value),
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -91,8 +125,11 @@ pub enum ResponsesInputItem {
|
|||||||
name: String,
|
name: String,
|
||||||
arguments: String,
|
arguments: String,
|
||||||
},
|
},
|
||||||
/// User is feeding a tool result back into the model.
|
/// User is feeding a tool result back into the model. `output`
|
||||||
FunctionCallOutput { call_id: String, output: String },
|
/// is a `Value` because OpenAI allows it to be either a plain
|
||||||
|
/// string or an array of content parts; the translator renders
|
||||||
|
/// either form to text rather than losing the tool result.
|
||||||
|
FunctionCallOutput { call_id: String, output: Value },
|
||||||
/// Reasoning items emitted by o-series models. Accepted but
|
/// Reasoning items emitted by o-series models. Accepted but
|
||||||
/// not forwarded to the model — neuron's candle path doesn't
|
/// not forwarded to the model — neuron's candle path doesn't
|
||||||
/// surface reasoning separately yet.
|
/// surface reasoning separately yet.
|
||||||
@@ -132,6 +169,11 @@ pub enum ResponsesContentPart {
|
|||||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||||
annotations: Vec<Value>,
|
annotations: Vec<Value>,
|
||||||
},
|
},
|
||||||
|
/// Any content-part type we don't model (e.g. `refusal`, audio).
|
||||||
|
/// Captured as a unit so an unknown part can't reject the whole
|
||||||
|
/// request; dropped during translation.
|
||||||
|
#[serde(other)]
|
||||||
|
Unknown,
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Response (non-streaming) ─────────────────────────────────────────
|
// ── Response (non-streaming) ─────────────────────────────────────────
|
||||||
@@ -277,20 +319,116 @@ mod tests {
|
|||||||
ResponsesInput::Items(items) => {
|
ResponsesInput::Items(items) => {
|
||||||
assert_eq!(items.len(), 1);
|
assert_eq!(items.len(), 1);
|
||||||
match &items[0] {
|
match &items[0] {
|
||||||
ResponsesInputItem::Message { role, content } => {
|
ResponsesInputElement::Typed(ResponsesInputItem::Message { role, content }) => {
|
||||||
assert_eq!(role, "user");
|
assert_eq!(role, "user");
|
||||||
match content {
|
match content {
|
||||||
ResponsesMessageContent::Text(t) => assert_eq!(t, "hi"),
|
ResponsesMessageContent::Text(t) => assert_eq!(t, "hi"),
|
||||||
other => panic!("expected Text content, got {other:?}"),
|
other => panic!("expected Text content, got {other:?}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
other => panic!("expected Message item, got {other:?}"),
|
other => panic!("expected typed Message item, got {other:?}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
other => panic!("expected Items, got {other:?}"),
|
other => panic!("expected Items, got {other:?}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn deserialises_bare_easy_message_without_type() {
|
||||||
|
// The shape agent-zero (via litellm) actually sends: `input`
|
||||||
|
// items are bare `{role, content}` with NO `type` field. This
|
||||||
|
// is the exact payload that was returning 422.
|
||||||
|
let raw = r#"{
|
||||||
|
"model": "Qwen/Qwen3.6-27B",
|
||||||
|
"store": true,
|
||||||
|
"tools": [{"type": "function", "name": "x", "description": "d", "parameters": {}}],
|
||||||
|
"input": [
|
||||||
|
{"role": "system", "content": "you are helpful"},
|
||||||
|
{"role": "assistant", "content": "{\"tool_name\":\"response\"}"},
|
||||||
|
{"role": "user", "content": "hi"}
|
||||||
|
]
|
||||||
|
}"#;
|
||||||
|
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||||
|
let items = match req.input {
|
||||||
|
ResponsesInput::Items(i) => i,
|
||||||
|
other => panic!("expected Items, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(items.len(), 3);
|
||||||
|
for el in &items {
|
||||||
|
assert!(
|
||||||
|
matches!(el, ResponsesInputElement::EasyMessage { .. }),
|
||||||
|
"expected EasyMessage, got {el:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
// `tools` / `store` ride through `extra`, not `input`.
|
||||||
|
assert!(req.extra.get("tools").is_some());
|
||||||
|
assert_eq!(req.extra.get("store"), Some(&Value::Bool(true)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn tolerates_null_content_and_unknown_item_types() {
|
||||||
|
// An assistant turn carrying only tool calls has `content: null`;
|
||||||
|
// and a future/unmodeled item type must not 422 the request.
|
||||||
|
let raw = r#"{
|
||||||
|
"model": "m",
|
||||||
|
"input": [
|
||||||
|
{"role": "assistant", "content": null},
|
||||||
|
{"type": "item_reference", "id": "abc"},
|
||||||
|
{"type": "function_call_output", "call_id": "c1",
|
||||||
|
"output": [{"type": "output_text", "text": "result"}]},
|
||||||
|
{"role": "user", "content": "go"}
|
||||||
|
]
|
||||||
|
}"#;
|
||||||
|
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||||
|
let items = match req.input {
|
||||||
|
ResponsesInput::Items(i) => i,
|
||||||
|
other => panic!("expected Items, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(items.len(), 4);
|
||||||
|
assert!(matches!(
|
||||||
|
&items[0],
|
||||||
|
ResponsesInputElement::EasyMessage { content: None, .. }
|
||||||
|
));
|
||||||
|
assert!(matches!(&items[1], ResponsesInputElement::Other(_)));
|
||||||
|
assert!(matches!(
|
||||||
|
&items[2],
|
||||||
|
ResponsesInputElement::Typed(ResponsesInputItem::FunctionCallOutput { .. })
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
&items[3],
|
||||||
|
ResponsesInputElement::EasyMessage { .. }
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn tolerates_unknown_content_part_type() {
|
||||||
|
// A `refusal` (or any unmodeled) content part must parse, not 422.
|
||||||
|
let raw = r#"{
|
||||||
|
"model": "m",
|
||||||
|
"input": [
|
||||||
|
{"role": "assistant", "content": [
|
||||||
|
{"type": "refusal", "refusal": "no"},
|
||||||
|
{"type": "output_text", "text": "ok"}
|
||||||
|
]}
|
||||||
|
]
|
||||||
|
}"#;
|
||||||
|
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||||
|
let items = match req.input {
|
||||||
|
ResponsesInput::Items(i) => i,
|
||||||
|
other => panic!("expected Items, got {other:?}"),
|
||||||
|
};
|
||||||
|
let parts = match &items[0] {
|
||||||
|
ResponsesInputElement::EasyMessage {
|
||||||
|
content: Some(ResponsesMessageContent::Parts(p)),
|
||||||
|
..
|
||||||
|
} => p,
|
||||||
|
other => panic!("expected EasyMessage with Parts, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(parts.len(), 2);
|
||||||
|
assert!(matches!(&parts[0], ResponsesContentPart::Unknown));
|
||||||
|
assert!(matches!(&parts[1], ResponsesContentPart::OutputText { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn deserialises_input_with_image() {
|
fn deserialises_input_with_image() {
|
||||||
let raw = r#"{
|
let raw = r#"{
|
||||||
@@ -308,10 +446,10 @@ mod tests {
|
|||||||
other => panic!("expected Items, got {other:?}"),
|
other => panic!("expected Items, got {other:?}"),
|
||||||
};
|
};
|
||||||
let parts = match &items[0] {
|
let parts = match &items[0] {
|
||||||
ResponsesInputItem::Message {
|
ResponsesInputElement::Typed(ResponsesInputItem::Message {
|
||||||
content: ResponsesMessageContent::Parts(p),
|
content: ResponsesMessageContent::Parts(p),
|
||||||
..
|
..
|
||||||
} => p,
|
}) => p,
|
||||||
other => panic!("expected Parts, got {other:?}"),
|
other => panic!("expected Parts, got {other:?}"),
|
||||||
};
|
};
|
||||||
assert_eq!(parts.len(), 2);
|
assert_eq!(parts.len(), 2);
|
||||||
|
|||||||
@@ -400,6 +400,7 @@ pub fn openai_to_anthropic(resp: ChatCompletionResponse) -> MessagesResponse {
|
|||||||
total_tokens: 0,
|
total_tokens: 0,
|
||||||
completion_tokens_details: None,
|
completion_tokens_details: None,
|
||||||
prompt_tokens_details: None,
|
prompt_tokens_details: None,
|
||||||
|
helexa_timing: None,
|
||||||
});
|
});
|
||||||
|
|
||||||
MessagesResponse {
|
MessagesResponse {
|
||||||
@@ -772,6 +773,7 @@ mod stream_tests {
|
|||||||
total_tokens: 267,
|
total_tokens: 267,
|
||||||
completion_tokens_details: None,
|
completion_tokens_details: None,
|
||||||
prompt_tokens_details: None,
|
prompt_tokens_details: None,
|
||||||
|
helexa_timing: None,
|
||||||
});
|
});
|
||||||
t.on_chunk(&usage_chunk);
|
t.on_chunk(&usage_chunk);
|
||||||
let fin = t.finish();
|
let fin = t.finish();
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ license.workspace = true
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
cortex-core.workspace = true
|
cortex-core.workspace = true
|
||||||
|
helexa-stream = { path = "../helexa-stream" }
|
||||||
async-trait.workspace = true
|
async-trait.workspace = true
|
||||||
tokio.workspace = true
|
tokio.workspace = true
|
||||||
axum.workspace = true
|
axum.workspace = true
|
||||||
|
|||||||
@@ -32,6 +32,8 @@ pub async fn stream_translated(
|
|||||||
openai_body: axum::body::Bytes,
|
openai_body: axum::body::Bytes,
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
node_name: &str,
|
node_name: &str,
|
||||||
|
inbound_headers: &axum::http::HeaderMap,
|
||||||
|
usage_sink: Option<crate::metering::UsageSink>,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
let url = format!("{endpoint}/v1/chat/completions");
|
let url = format!("{endpoint}/v1/chat/completions");
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
@@ -42,13 +44,14 @@ pub async fn stream_translated(
|
|||||||
"proxying streaming request (anthropic SSE translation)"
|
"proxying streaming request (anthropic SSE translation)"
|
||||||
);
|
);
|
||||||
|
|
||||||
let upstream = match client
|
let request = crate::auth::forward_principal_headers(
|
||||||
.post(&url)
|
client
|
||||||
.header("content-type", "application/json")
|
.post(&url)
|
||||||
.body(openai_body)
|
.header("content-type", "application/json")
|
||||||
.send()
|
.body(openai_body),
|
||||||
.await
|
inbound_headers,
|
||||||
{
|
);
|
||||||
|
let upstream = match request.send().await {
|
||||||
Ok(r) => r,
|
Ok(r) => r,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
@@ -94,6 +97,10 @@ pub async fn stream_translated(
|
|||||||
let mut saw_tool_call = false;
|
let mut saw_tool_call = false;
|
||||||
let mut last_finish: Option<String> = None;
|
let mut last_finish: Option<String> = None;
|
||||||
let mut frames = 0u64;
|
let mut frames = 0u64;
|
||||||
|
// Engine-truth usage for metering (#51), scanned from the upstream
|
||||||
|
// frames (neuron emits a final `usage` object on the stream, #48).
|
||||||
|
let mut usage_prompt = 0u64;
|
||||||
|
let mut usage_completion = 0u64;
|
||||||
|
|
||||||
'outer: while let Some(block) = upstream.next().await {
|
'outer: while let Some(block) = upstream.next().await {
|
||||||
let block = match block {
|
let block = match block {
|
||||||
@@ -121,6 +128,15 @@ pub async fn stream_translated(
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
tracing::trace!(node = %node, frame = %data, "anthropic stream: upstream frame");
|
tracing::trace!(node = %node, frame = %data, "anthropic stream: upstream frame");
|
||||||
|
// Capture usage for metering before translation — the
|
||||||
|
// usage object rides on a late frame (often after the
|
||||||
|
// last content delta).
|
||||||
|
if let Some(p) = crate::proxy::last_count_for(data, "prompt_tokens") {
|
||||||
|
usage_prompt = p;
|
||||||
|
}
|
||||||
|
if let Some(c) = crate::proxy::last_count_for(data, "completion_tokens") {
|
||||||
|
usage_completion = c;
|
||||||
|
}
|
||||||
let Ok(chunk) = serde_json::from_str::<ChatCompletionChunk>(data) else {
|
let Ok(chunk) = serde_json::from_str::<ChatCompletionChunk>(data) else {
|
||||||
tracing::debug!(node = %node, "anthropic stream: unparsable upstream frame skipped");
|
tracing::debug!(node = %node, "anthropic stream: unparsable upstream frame skipped");
|
||||||
continue;
|
continue;
|
||||||
@@ -162,6 +178,14 @@ pub async fn stream_translated(
|
|||||||
terminated = done,
|
terminated = done,
|
||||||
"anthropic stream complete"
|
"anthropic stream complete"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
// Settle metering with the observed usage (#51). Runs on every exit
|
||||||
|
// path of the pump — clean end, early break, or upstream error — so
|
||||||
|
// the reservation is always resolved. `(0, 0)` when no usage frame
|
||||||
|
// was seen, which releases without recording spend.
|
||||||
|
if let Some(sink) = usage_sink {
|
||||||
|
sink(usage_prompt, usage_completion);
|
||||||
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
Response::builder()
|
Response::builder()
|
||||||
|
|||||||
144
crates/cortex-gateway/src/auth.rs
Normal file
144
crates/cortex-gateway/src/auth.rs
Normal file
@@ -0,0 +1,144 @@
|
|||||||
|
//! API-key authentication + principal resolution (#49).
|
||||||
|
//!
|
||||||
|
//! Identity rides standard bearer auth only — `Authorization: Bearer <key>`
|
||||||
|
//! — which is what keeps every tier OpenAI-compatible by construction (no
|
||||||
|
//! custom required headers or body fields, per #47). The middleware resolves
|
||||||
|
//! the key to a [`Principal`] via the [`EntitlementProvider`], carries it in
|
||||||
|
//! the request extensions for cortex-side metering/enforcement (#51/#52), and
|
||||||
|
//! stamps it as internal headers on the request so it reaches neuron, which
|
||||||
|
//! trusts cortex's assertion over WireGuard (#54).
|
||||||
|
//!
|
||||||
|
//! Anti-spoofing: any client-supplied principal header is **stripped** before
|
||||||
|
//! the authoritative value is stamped, so a client can never assert a
|
||||||
|
//! principal it didn't authenticate as.
|
||||||
|
//!
|
||||||
|
//! Rejection contract (#63): missing key under `require_auth`, or any present
|
||||||
|
//! but unresolvable key, yields `401 invalid_api_key` in the #60 envelope.
|
||||||
|
|
||||||
|
use crate::error::envelope_response;
|
||||||
|
use crate::state::CortexState;
|
||||||
|
use axum::extract::{Request, State};
|
||||||
|
use axum::http::header::AUTHORIZATION;
|
||||||
|
use axum::http::{HeaderMap, HeaderValue};
|
||||||
|
use axum::middleware::Next;
|
||||||
|
use axum::response::Response;
|
||||||
|
use cortex_core::entitlements::{AuthError, HEADER_ACCOUNT_ID, HEADER_KEY_ID};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Endpoints that never require auth: liveness/readiness probes. Everything
|
||||||
|
/// else flows through resolution.
|
||||||
|
fn is_public(path: &str) -> bool {
|
||||||
|
path == "/health" || path == "/"
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract the bearer token from an `Authorization` header value, if present
|
||||||
|
/// and well-formed. Scheme match is case-insensitive per RFC 7235.
|
||||||
|
fn parse_bearer(headers: &HeaderMap) -> Option<String> {
|
||||||
|
let raw = headers.get(AUTHORIZATION)?.to_str().ok()?;
|
||||||
|
let (scheme, token) = raw.split_once(' ')?;
|
||||||
|
if scheme.eq_ignore_ascii_case("bearer") {
|
||||||
|
let token = token.trim();
|
||||||
|
(!token.is_empty()).then(|| token.to_string())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Axum middleware: resolve the bearer key, attach the principal, stamp the
|
||||||
|
/// internal headers. Wired in `build_app` via `from_fn_with_state`.
|
||||||
|
pub async fn require_principal(
|
||||||
|
State(fleet): State<Arc<CortexState>>,
|
||||||
|
mut req: Request,
|
||||||
|
next: Next,
|
||||||
|
) -> Response {
|
||||||
|
if is_public(req.uri().path()) {
|
||||||
|
return next.run(req).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Anti-spoof: drop any client-supplied principal headers up front.
|
||||||
|
{
|
||||||
|
let headers = req.headers_mut();
|
||||||
|
headers.remove(HEADER_ACCOUNT_ID);
|
||||||
|
headers.remove(HEADER_KEY_ID);
|
||||||
|
}
|
||||||
|
|
||||||
|
match parse_bearer(req.headers()) {
|
||||||
|
Some(key) => match fleet.entitlements.resolve(&key).await {
|
||||||
|
Ok(principal) => {
|
||||||
|
// Stamp the authoritative principal for neuron. Account/key
|
||||||
|
// ids come from operator config, so they're valid header
|
||||||
|
// values; guard anyway and skip a malformed one rather than
|
||||||
|
// panic.
|
||||||
|
if let (Ok(account), Ok(key_id)) = (
|
||||||
|
HeaderValue::from_str(&principal.account_id),
|
||||||
|
HeaderValue::from_str(&principal.key_id),
|
||||||
|
) {
|
||||||
|
let headers = req.headers_mut();
|
||||||
|
headers.insert(HEADER_ACCOUNT_ID, account);
|
||||||
|
headers.insert(HEADER_KEY_ID, key_id);
|
||||||
|
}
|
||||||
|
// Carry the typed principal for cortex-side metering (#51)
|
||||||
|
// and budget enforcement (#52).
|
||||||
|
req.extensions_mut().insert(principal);
|
||||||
|
next.run(req).await
|
||||||
|
}
|
||||||
|
// The entitlement authority is unreachable (upstream client
|
||||||
|
// blip, #57). Fail **closed but distinct**: a transient outage
|
||||||
|
// must not reject a real key as `401 invalid_api_key` — it's a
|
||||||
|
// retryable `503`. This holds regardless of require_auth: we
|
||||||
|
// can't safely serve a key we couldn't authorize.
|
||||||
|
Err(AuthError::Unavailable { retry_after_secs }) => {
|
||||||
|
envelope_response(OpenAiError::service_unavailable(
|
||||||
|
"entitlement authority temporarily unavailable",
|
||||||
|
Some(retry_after_secs),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
// A genuinely unrecognized key only hard-fails when auth is
|
||||||
|
// *required*. In allow-anonymous mode (the default) we IGNORE it
|
||||||
|
// and serve unauthenticated — otherwise the placeholder keys that
|
||||||
|
// OpenAI-compatible clients send by default (opencode, Open WebUI,
|
||||||
|
// Agent Zero, litellm) would all break though the operator never
|
||||||
|
// opted into auth. Pre-#49 the bearer was never inspected; this
|
||||||
|
// preserves that for require_auth=false.
|
||||||
|
Err(AuthError::InvalidKey) => {
|
||||||
|
if fleet.require_auth {
|
||||||
|
unauthorized("invalid API key")
|
||||||
|
} else {
|
||||||
|
tracing::debug!(
|
||||||
|
"ignoring unrecognized bearer token (require_auth=false): serving anonymously"
|
||||||
|
);
|
||||||
|
next.run(req).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
None => {
|
||||||
|
if fleet.require_auth {
|
||||||
|
unauthorized("missing API key; supply 'Authorization: Bearer <key>'")
|
||||||
|
} else {
|
||||||
|
next.run(req).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `401 invalid_api_key` in the standard envelope (#63).
|
||||||
|
fn unauthorized(message: &str) -> Response {
|
||||||
|
envelope_response(OpenAiError::invalid_api_key(message))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Copy the cortex-stamped principal headers from an inbound [`HeaderMap`]
|
||||||
|
/// onto an outbound reqwest builder. Used by the Anthropic proxy paths,
|
||||||
|
/// which construct their own upstream requests instead of going through
|
||||||
|
/// [`crate::proxy::forward_request`] (which forwards all headers verbatim).
|
||||||
|
pub fn forward_principal_headers(
|
||||||
|
mut builder: reqwest::RequestBuilder,
|
||||||
|
headers: &HeaderMap,
|
||||||
|
) -> reqwest::RequestBuilder {
|
||||||
|
for name in [HEADER_ACCOUNT_ID, HEADER_KEY_ID] {
|
||||||
|
if let Some(value) = headers.get(name) {
|
||||||
|
builder = builder.header(name, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
}
|
||||||
112
crates/cortex-gateway/src/entitlements_chain.rs
Normal file
112
crates/cortex-gateway/src/entitlements_chain.rs
Normal file
@@ -0,0 +1,112 @@
|
|||||||
|
//! Chained entitlement provider (#57): operator-local keys first, mesh
|
||||||
|
//! upstream for everything else.
|
||||||
|
//!
|
||||||
|
//! `resolve` tries the [`LocalEntitlementProvider`] (operator + infra keys —
|
||||||
|
//! never a network hop); only a locally-unknown key falls through to
|
||||||
|
//! [`UpstreamEntitlementProvider`]. Because the local provider treats an
|
||||||
|
//! unconfigured principal as uncapped, reserve/settle/release/snapshot must
|
||||||
|
//! **not** blindly hit local — they dispatch to whichever backend resolved
|
||||||
|
//! that account, remembered in a map keyed by `account_id` (populated at
|
||||||
|
//! resolve time).
|
||||||
|
|
||||||
|
use crate::entitlements_local::LocalEntitlementProvider;
|
||||||
|
use crate::entitlements_upstream::UpstreamEntitlementProvider;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use cortex_core::entitlements::{
|
||||||
|
AuthError, BudgetError, BudgetSnapshot, EntitlementProvider, Principal, Reservation,
|
||||||
|
};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
|
enum Backend {
|
||||||
|
Local,
|
||||||
|
Upstream,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ChainedEntitlementProvider {
|
||||||
|
local: LocalEntitlementProvider,
|
||||||
|
upstream: UpstreamEntitlementProvider,
|
||||||
|
/// account_id → which backend owns it, learned at resolve time.
|
||||||
|
backends: RwLock<HashMap<String, Backend>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChainedEntitlementProvider {
|
||||||
|
pub fn new(local: LocalEntitlementProvider, upstream: UpstreamEntitlementProvider) -> Self {
|
||||||
|
Self {
|
||||||
|
local,
|
||||||
|
upstream,
|
||||||
|
backends: RwLock::new(HashMap::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record(&self, account_id: &str, backend: Backend) {
|
||||||
|
self.backends
|
||||||
|
.write()
|
||||||
|
.await
|
||||||
|
.insert(account_id.to_string(), backend);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The backend that owns `account_id`. Defaults to `Upstream` for an
|
||||||
|
/// account never resolved this process-lifetime (a resolve always
|
||||||
|
/// precedes reserve in a request, so this is just a safe fallback —
|
||||||
|
/// upstream fails closed if the account is bogus).
|
||||||
|
async fn backend_for(&self, account_id: &str) -> Backend {
|
||||||
|
self.backends
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.get(account_id)
|
||||||
|
.copied()
|
||||||
|
.unwrap_or(Backend::Upstream)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl EntitlementProvider for ChainedEntitlementProvider {
|
||||||
|
async fn resolve(&self, api_key: &str) -> Result<Principal, AuthError> {
|
||||||
|
match self.local.resolve(api_key).await {
|
||||||
|
Ok(p) => {
|
||||||
|
self.record(&p.account_id, Backend::Local).await;
|
||||||
|
Ok(p)
|
||||||
|
}
|
||||||
|
Err(AuthError::InvalidKey) => {
|
||||||
|
let p = self.upstream.resolve(api_key).await?;
|
||||||
|
self.record(&p.account_id, Backend::Upstream).await;
|
||||||
|
Ok(p)
|
||||||
|
}
|
||||||
|
Err(e) => Err(e),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn reserve(
|
||||||
|
&self,
|
||||||
|
principal: &Principal,
|
||||||
|
max_tokens: u64,
|
||||||
|
) -> Result<Reservation, BudgetError> {
|
||||||
|
match self.backend_for(&principal.account_id).await {
|
||||||
|
Backend::Local => self.local.reserve(principal, max_tokens).await,
|
||||||
|
Backend::Upstream => self.upstream.reserve(principal, max_tokens).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn settle(&self, reservation: Reservation, actual_tokens: u64) {
|
||||||
|
match self.backend_for(&reservation.principal.account_id).await {
|
||||||
|
Backend::Local => self.local.settle(reservation, actual_tokens).await,
|
||||||
|
Backend::Upstream => self.upstream.settle(reservation, actual_tokens).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn release(&self, reservation: Reservation) {
|
||||||
|
match self.backend_for(&reservation.principal.account_id).await {
|
||||||
|
Backend::Local => self.local.release(reservation).await,
|
||||||
|
Backend::Upstream => self.upstream.release(reservation).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn snapshot(&self, principal: &Principal) -> Option<BudgetSnapshot> {
|
||||||
|
match self.backend_for(&principal.account_id).await {
|
||||||
|
Backend::Local => self.local.snapshot(principal).await,
|
||||||
|
Backend::Upstream => self.upstream.snapshot(principal).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
246
crates/cortex-gateway/src/entitlements_upstream.rs
Normal file
246
crates/cortex-gateway/src/entitlements_upstream.rs
Normal file
@@ -0,0 +1,246 @@
|
|||||||
|
//! helexa-upstream client (#57): an [`EntitlementProvider`] that resolves
|
||||||
|
//! keys and reserves/settles budget against the mesh authority's
|
||||||
|
//! `/authz/v1` surface (B2). It is "just another impl of the trait" — cortex
|
||||||
|
//! enforcement (`auth.rs`, `metering.rs`) is unchanged.
|
||||||
|
//!
|
||||||
|
//! **Fail closed.** When upstream is unreachable, `resolve` returns
|
||||||
|
//! [`AuthError::Unavailable`] (→ `503`, never `401`) and `reserve` refuses
|
||||||
|
//! with a retryable [`BudgetError::RateLimited`] — a request is never served
|
||||||
|
//! on an un-authorized key, and a real key is never rejected as invalid
|
||||||
|
//! during a blip.
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use cortex_core::config::UpstreamClientConfig;
|
||||||
|
use cortex_core::entitlements::{
|
||||||
|
AuthError, BudgetError, BudgetSnapshot, EntitlementProvider, Principal, Reservation,
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
/// Retry-After (seconds) advertised when we fail closed on an upstream
|
||||||
|
/// outage.
|
||||||
|
const FAIL_CLOSED_RETRY_SECS: u64 = 5;
|
||||||
|
|
||||||
|
pub struct UpstreamEntitlementProvider {
|
||||||
|
client: reqwest::Client,
|
||||||
|
base_url: String,
|
||||||
|
bearer: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct PrincipalDto {
|
||||||
|
account_id: String,
|
||||||
|
key_id: String,
|
||||||
|
}
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct SnapshotDto {
|
||||||
|
hard_cap: Option<u64>,
|
||||||
|
spent: u64,
|
||||||
|
reserved: u64,
|
||||||
|
}
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ResolveResp {
|
||||||
|
principal: PrincipalDto,
|
||||||
|
#[allow(dead_code)]
|
||||||
|
snapshot: Option<SnapshotDto>,
|
||||||
|
}
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ReserveResp {
|
||||||
|
reservation_id: Option<i64>,
|
||||||
|
rejected: Option<Rejection>,
|
||||||
|
}
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||||
|
enum Rejection {
|
||||||
|
InsufficientQuota {
|
||||||
|
requested: u64,
|
||||||
|
available: u64,
|
||||||
|
},
|
||||||
|
RateLimited {
|
||||||
|
requested: u64,
|
||||||
|
available: u64,
|
||||||
|
retry_after_secs: u64,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
impl UpstreamEntitlementProvider {
|
||||||
|
pub fn new(cfg: &UpstreamClientConfig) -> Self {
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(Duration::from_secs(cfg.timeout_secs))
|
||||||
|
.build()
|
||||||
|
.expect("failed to build upstream HTTP client");
|
||||||
|
Self {
|
||||||
|
client,
|
||||||
|
base_url: cfg.url.trim_end_matches('/').to_string(),
|
||||||
|
bearer: cfg.bearer.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn url(&self, path: &str) -> String {
|
||||||
|
format!("{}{}", self.base_url, path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl EntitlementProvider for UpstreamEntitlementProvider {
|
||||||
|
async fn resolve(&self, api_key: &str) -> Result<Principal, AuthError> {
|
||||||
|
let resp = self
|
||||||
|
.client
|
||||||
|
.post(self.url("/authz/v1/resolve"))
|
||||||
|
.bearer_auth(&self.bearer)
|
||||||
|
.json(&serde_json::json!({ "api_key": api_key }))
|
||||||
|
.send()
|
||||||
|
.await;
|
||||||
|
let resp = match resp {
|
||||||
|
Ok(r) => r,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(error = %e, "upstream resolve unreachable; failing closed");
|
||||||
|
return Err(AuthError::Unavailable {
|
||||||
|
retry_after_secs: FAIL_CLOSED_RETRY_SECS,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if resp.status().as_u16() == 401 {
|
||||||
|
return Err(AuthError::InvalidKey);
|
||||||
|
}
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return Err(AuthError::Unavailable {
|
||||||
|
retry_after_secs: FAIL_CLOSED_RETRY_SECS,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
match resp.json::<ResolveResp>().await {
|
||||||
|
Ok(r) => Ok(Principal {
|
||||||
|
account_id: r.principal.account_id,
|
||||||
|
key_id: r.principal.key_id,
|
||||||
|
}),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(error = %e, "upstream resolve: bad body; failing closed");
|
||||||
|
Err(AuthError::Unavailable {
|
||||||
|
retry_after_secs: FAIL_CLOSED_RETRY_SECS,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn reserve(
|
||||||
|
&self,
|
||||||
|
principal: &Principal,
|
||||||
|
max_tokens: u64,
|
||||||
|
) -> Result<Reservation, BudgetError> {
|
||||||
|
let fail_closed = || BudgetError::RateLimited {
|
||||||
|
requested: max_tokens,
|
||||||
|
available: 0,
|
||||||
|
retry_after_secs: FAIL_CLOSED_RETRY_SECS,
|
||||||
|
};
|
||||||
|
let resp = self
|
||||||
|
.client
|
||||||
|
.post(self.url("/authz/v1/reserve"))
|
||||||
|
.bearer_auth(&self.bearer)
|
||||||
|
.json(&serde_json::json!({
|
||||||
|
"account_id": principal.account_id,
|
||||||
|
"key_id": principal.key_id,
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await;
|
||||||
|
let resp = match resp {
|
||||||
|
Ok(r) if r.status().is_success() => r,
|
||||||
|
Ok(r) => {
|
||||||
|
tracing::warn!(status = %r.status(), "upstream reserve non-2xx; failing closed");
|
||||||
|
return Err(fail_closed());
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(error = %e, "upstream reserve unreachable; failing closed");
|
||||||
|
return Err(fail_closed());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match resp.json::<ReserveResp>().await {
|
||||||
|
Ok(ReserveResp {
|
||||||
|
reservation_id: Some(id),
|
||||||
|
..
|
||||||
|
}) => Ok(Reservation {
|
||||||
|
id: id as u64,
|
||||||
|
principal: principal.clone(),
|
||||||
|
reserved: max_tokens,
|
||||||
|
}),
|
||||||
|
Ok(ReserveResp {
|
||||||
|
rejected:
|
||||||
|
Some(Rejection::InsufficientQuota {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
}),
|
||||||
|
..
|
||||||
|
}) => Err(BudgetError::InsufficientQuota {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
}),
|
||||||
|
Ok(ReserveResp {
|
||||||
|
rejected:
|
||||||
|
Some(Rejection::RateLimited {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
retry_after_secs,
|
||||||
|
}),
|
||||||
|
..
|
||||||
|
}) => Err(BudgetError::RateLimited {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
retry_after_secs,
|
||||||
|
}),
|
||||||
|
_ => Err(fail_closed()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn settle(&self, reservation: Reservation, actual_tokens: u64) {
|
||||||
|
// Best-effort; a lost settle is reaped by the upstream sweeper (B2).
|
||||||
|
let _ = self
|
||||||
|
.client
|
||||||
|
.post(self.url("/authz/v1/settle"))
|
||||||
|
.bearer_auth(&self.bearer)
|
||||||
|
.json(&serde_json::json!({
|
||||||
|
"reservation_id": reservation.id as i64,
|
||||||
|
"actual_tokens": actual_tokens,
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.inspect_err(
|
||||||
|
|e| tracing::warn!(error = %e, "upstream settle failed (sweeper will reap)"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn release(&self, reservation: Reservation) {
|
||||||
|
let _ = self
|
||||||
|
.client
|
||||||
|
.post(self.url("/authz/v1/release"))
|
||||||
|
.bearer_auth(&self.bearer)
|
||||||
|
.json(&serde_json::json!({ "reservation_id": reservation.id as i64 }))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.inspect_err(
|
||||||
|
|e| tracing::warn!(error = %e, "upstream release failed (sweeper will reap)"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn snapshot(&self, principal: &Principal) -> Option<BudgetSnapshot> {
|
||||||
|
let resp = self
|
||||||
|
.client
|
||||||
|
.post(self.url("/authz/v1/snapshot"))
|
||||||
|
.bearer_auth(&self.bearer)
|
||||||
|
.json(&serde_json::json!({
|
||||||
|
"account_id": principal.account_id,
|
||||||
|
"key_id": principal.key_id,
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.ok()?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let dto = resp.json::<SnapshotDto>().await.ok()?;
|
||||||
|
Some(BudgetSnapshot {
|
||||||
|
hard_cap: dto.hard_cap,
|
||||||
|
spent: dto.spent,
|
||||||
|
reserved: dto.reserved,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -190,7 +190,7 @@ async fn completions(
|
|||||||
/// `POST /v1/messages` — accept Anthropic format, translate, proxy, translate back.
|
/// `POST /v1/messages` — accept Anthropic format, translate, proxy, translate back.
|
||||||
async fn anthropic_messages(
|
async fn anthropic_messages(
|
||||||
State(fleet): State<Arc<CortexState>>,
|
State(fleet): State<Arc<CortexState>>,
|
||||||
_headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
body: Bytes,
|
body: Bytes,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
// Parse as Anthropic request.
|
// Parse as Anthropic request.
|
||||||
@@ -306,6 +306,33 @@ async fn anthropic_messages(
|
|||||||
}
|
}
|
||||||
let start = Instant::now();
|
let start = Instant::now();
|
||||||
|
|
||||||
|
// Per-request metering + budget enforcement (#51/#52), same lifecycle as
|
||||||
|
// the OpenAI paths. Estimate from the translated OpenAI body (what neuron
|
||||||
|
// sees). Refuse over-cap before dispatch via the #63 envelope; otherwise
|
||||||
|
// build the sink consumed by whichever branch runs below.
|
||||||
|
let usage_sink = match crate::metering::principal_from_headers(&headers) {
|
||||||
|
Some(principal) => {
|
||||||
|
let advertised =
|
||||||
|
advertised_output_limit(&fleet, &route.node_name, &route.resolved_model_id).await;
|
||||||
|
let max_tokens = crate::metering::reservation_estimate(&openai_body, advertised);
|
||||||
|
match crate::metering::reserve_or_reject(
|
||||||
|
Arc::clone(&fleet.entitlements),
|
||||||
|
&principal,
|
||||||
|
max_tokens,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(guard) => Some(crate::metering::usage_sink(
|
||||||
|
principal,
|
||||||
|
guard,
|
||||||
|
std::sync::Arc::clone(&fleet.served_usage),
|
||||||
|
)),
|
||||||
|
Err(env) => return crate::error::envelope_response(env),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
if is_streaming {
|
if is_streaming {
|
||||||
// Anthropic SSE translation (#24): upstream speaks OpenAI SSE;
|
// Anthropic SSE translation (#24): upstream speaks OpenAI SSE;
|
||||||
// re-frame it event-by-event into Anthropic's message_start /
|
// re-frame it event-by-event into Anthropic's message_start /
|
||||||
@@ -316,6 +343,8 @@ async fn anthropic_messages(
|
|||||||
openai_body,
|
openai_body,
|
||||||
&model_id,
|
&model_id,
|
||||||
&route.node_name,
|
&route.node_name,
|
||||||
|
&headers,
|
||||||
|
usage_sink,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
metrics::histogram!("cortex_request_duration_seconds", &labels)
|
metrics::histogram!("cortex_request_duration_seconds", &labels)
|
||||||
@@ -335,13 +364,16 @@ async fn anthropic_messages(
|
|||||||
cold_start = route.cold_start,
|
cold_start = route.cold_start,
|
||||||
"proxying request"
|
"proxying request"
|
||||||
);
|
);
|
||||||
let upstream_resp = fleet
|
let upstream_resp = crate::auth::forward_principal_headers(
|
||||||
.http_client
|
fleet
|
||||||
.post(&target_url)
|
.http_client
|
||||||
.body(openai_body)
|
.post(&target_url)
|
||||||
.header("content-type", "application/json")
|
.body(openai_body)
|
||||||
.send()
|
.header("content-type", "application/json"),
|
||||||
.await;
|
&headers,
|
||||||
|
)
|
||||||
|
.send()
|
||||||
|
.await;
|
||||||
|
|
||||||
let upstream_resp = match upstream_resp {
|
let upstream_resp = match upstream_resp {
|
||||||
Ok(r) => r,
|
Ok(r) => r,
|
||||||
@@ -437,6 +469,15 @@ async fn anthropic_messages(
|
|||||||
|
|
||||||
metrics::histogram!("cortex_request_duration_seconds", &labels)
|
metrics::histogram!("cortex_request_duration_seconds", &labels)
|
||||||
.record(start.elapsed().as_secs_f64());
|
.record(start.elapsed().as_secs_f64());
|
||||||
|
// Settle metering with the upstream usage (#51). Scanned from the
|
||||||
|
// raw body — same engine-truth source as the streaming path — so we
|
||||||
|
// don't depend on the typed usage struct's optionality.
|
||||||
|
if let Some(sink) = usage_sink {
|
||||||
|
let tail = String::from_utf8_lossy(&body_bytes);
|
||||||
|
let prompt = proxy::last_count_for(&tail, "prompt_tokens").unwrap_or(0);
|
||||||
|
let completion = proxy::last_count_for(&tail, "completion_tokens").unwrap_or(0);
|
||||||
|
sink(prompt, completion);
|
||||||
|
}
|
||||||
// Did the model actually produce a structured tool call, or just
|
// Did the model actually produce a structured tool call, or just
|
||||||
// text? This is the single most useful signal for "is tool
|
// text? This is the single most useful signal for "is tool
|
||||||
// calling working end-to-end" — a `false` here alongside a
|
// calling working end-to-end" — a `false` here alongside a
|
||||||
@@ -540,6 +581,11 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
|||||||
// Runtime-detected — will be OR-ed in Pass 2 from neuron data.
|
// Runtime-detected — will be OR-ed in Pass 2 from neuron data.
|
||||||
tool_call: false,
|
tool_call: false,
|
||||||
reasoning: false,
|
reasoning: false,
|
||||||
|
// Flat #78 fields are derived from `limit` in the final
|
||||||
|
// sync pass, once merging is done.
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -597,6 +643,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
|||||||
cost: None,
|
cost: None,
|
||||||
tool_call: entry.tool_call,
|
tool_call: entry.tool_call,
|
||||||
reasoning: entry.reasoning,
|
reasoning: entry.reasoning,
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -653,6 +702,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
|||||||
cost: None,
|
cost: None,
|
||||||
tool_call: false,
|
tool_call: false,
|
||||||
reasoning: false,
|
reasoning: false,
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -687,11 +739,24 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
|||||||
cost: target_entry.cost.clone(),
|
cost: target_entry.cost.clone(),
|
||||||
tool_call: target_entry.tool_call,
|
tool_call: target_entry.tool_call,
|
||||||
reasoning: target_entry.reasoning,
|
reasoning: target_entry.reasoning,
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let data: Vec<Value> = entries.values().map(|e| json!(e)).collect();
|
// Final pass: derive the flat ecosystem context-window fields (#78)
|
||||||
|
// from each entry's now-settled `limit`, so vLLM-convention clients
|
||||||
|
// (Hermes Agent et al.) can read the window without knowing helexa's
|
||||||
|
// `limit` schema.
|
||||||
|
let data: Vec<Value> = entries
|
||||||
|
.values_mut()
|
||||||
|
.map(|e| {
|
||||||
|
e.sync_flat_limit();
|
||||||
|
json!(e)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
Json(json!({
|
Json(json!({
|
||||||
"object": "list",
|
"object": "list",
|
||||||
"data": data,
|
"data": data,
|
||||||
@@ -724,6 +789,19 @@ async fn proxy_with_metrics(
|
|||||||
body: Bytes,
|
body: Bytes,
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
|
// Fail-fast prompt pre-validation (#56): refuse a prompt that already
|
||||||
|
// exceeds the model's advertised context window *before* dispatching to
|
||||||
|
// neuron — the same `400 context_length_exceeded` neuron would emit on
|
||||||
|
// overflow, just earlier and without burning a cold-load/queue slot.
|
||||||
|
// cortex has no tokenizer, so the estimate under-counts and neuron stays
|
||||||
|
// the exact wall; we only catch gross overages (the A0 failure mode).
|
||||||
|
if let Some(context) = advertised_context(fleet, &route.node_name, model_id).await {
|
||||||
|
let est = estimate_prompt_tokens(&body);
|
||||||
|
if est > context {
|
||||||
|
return context_length_exceeded_response(context, est, &headers);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let labels = [
|
let labels = [
|
||||||
("model", model_id.to_string()),
|
("model", model_id.to_string()),
|
||||||
("node", route.node_name.clone()),
|
("node", route.node_name.clone()),
|
||||||
@@ -734,9 +812,46 @@ async fn proxy_with_metrics(
|
|||||||
metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
|
metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Per-request metering + budget enforcement (#51/#52): reconstruct the
|
||||||
|
// principal from the middleware-stamped headers, reserve the request's
|
||||||
|
// upper-bound cost (prompt estimate + max output), and build the
|
||||||
|
// completion sink that settles actual spend when the response finishes.
|
||||||
|
// A reservation over the hard cap is refused *before* dispatch with the
|
||||||
|
// #63 envelope. Anonymous requests skip all of this. Must happen before
|
||||||
|
// `headers`/`body` are moved into the proxy.
|
||||||
|
let usage_sink = match crate::metering::principal_from_headers(&headers) {
|
||||||
|
Some(principal) => {
|
||||||
|
let advertised = advertised_output_limit(fleet, &route.node_name, model_id).await;
|
||||||
|
let max_tokens = crate::metering::reservation_estimate(&body, advertised);
|
||||||
|
match crate::metering::reserve_or_reject(
|
||||||
|
Arc::clone(&fleet.entitlements),
|
||||||
|
&principal,
|
||||||
|
max_tokens,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(guard) => Some(crate::metering::usage_sink(
|
||||||
|
principal,
|
||||||
|
guard,
|
||||||
|
std::sync::Arc::clone(&fleet.served_usage),
|
||||||
|
)),
|
||||||
|
Err(env) => return crate::error::envelope_response(env),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
let start = Instant::now();
|
let start = Instant::now();
|
||||||
let result =
|
let result = proxy::forward_request(
|
||||||
proxy::forward_request(&fleet.http_client, route, path, headers, body, model_id).await;
|
&fleet.http_client,
|
||||||
|
route,
|
||||||
|
path,
|
||||||
|
headers,
|
||||||
|
body,
|
||||||
|
model_id,
|
||||||
|
usage_sink,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
let duration = start.elapsed();
|
let duration = start.elapsed();
|
||||||
|
|
||||||
match result {
|
match result {
|
||||||
@@ -755,6 +870,117 @@ async fn proxy_with_metrics(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The model's advertised `limit.output` (#62) on a given node, used as the
|
||||||
|
/// default output budget for budget reservations (#52) when the request
|
||||||
|
/// omits `max_(completion_)tokens`. `None` when the node/model/limit is
|
||||||
|
/// unknown — callers fall back to [`crate::metering::FALLBACK_MAX_OUTPUT`].
|
||||||
|
async fn advertised_output_limit(
|
||||||
|
fleet: &CortexState,
|
||||||
|
node_name: &str,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Option<u64> {
|
||||||
|
let nodes = fleet.nodes.read().await;
|
||||||
|
nodes
|
||||||
|
.get(node_name)?
|
||||||
|
.models
|
||||||
|
.get(model_id)?
|
||||||
|
.limit
|
||||||
|
.as_ref()
|
||||||
|
.map(|l| l.output as u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The model's advertised hard context window (`limit.context`, #62/#67) on a
|
||||||
|
/// node, used for fail-fast prompt pre-validation (#56). `None` when no limit
|
||||||
|
/// is known — pre-validation is then skipped and neuron remains the wall.
|
||||||
|
async fn advertised_context(fleet: &CortexState, node_name: &str, model_id: &str) -> Option<u64> {
|
||||||
|
let nodes = fleet.nodes.read().await;
|
||||||
|
nodes
|
||||||
|
.get(node_name)?
|
||||||
|
.models
|
||||||
|
.get(model_id)?
|
||||||
|
.limit
|
||||||
|
.as_ref()
|
||||||
|
.map(|l| l.context as u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Conservative prompt-token estimate (~4 chars/token over message text).
|
||||||
|
/// cortex has no tokenizer; under-counting is the safe direction — we only
|
||||||
|
/// pre-reject gross overages (#56), and neuron enforces the exact wall.
|
||||||
|
fn estimate_prompt_tokens(body: &[u8]) -> u64 {
|
||||||
|
let Ok(v) = serde_json::from_slice::<Value>(body) else {
|
||||||
|
return (body.len() as u64 / 4).max(1);
|
||||||
|
};
|
||||||
|
let mut chars = 0usize;
|
||||||
|
if let Some(messages) = v.get("messages").and_then(Value::as_array) {
|
||||||
|
for m in messages {
|
||||||
|
match m.get("content") {
|
||||||
|
Some(Value::String(s)) => chars += s.len(),
|
||||||
|
Some(Value::Array(parts)) => {
|
||||||
|
for p in parts {
|
||||||
|
if let Some(t) = p.get("text").and_then(Value::as_str) {
|
||||||
|
chars += t.len();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
chars += 8; // rough per-message role/formatting overhead
|
||||||
|
}
|
||||||
|
} else if let Some(prompt) = v.get("prompt").and_then(Value::as_str) {
|
||||||
|
chars += prompt.len(); // legacy /v1/completions
|
||||||
|
} else {
|
||||||
|
return (body.len() as u64 / 4).max(1);
|
||||||
|
}
|
||||||
|
(chars as u64 / 4).max(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Client-specific, advisory guidance for an over-long prompt (#56),
|
||||||
|
/// fingerprinted from `User-Agent`. Strictly advisory: it rides the
|
||||||
|
/// `X-Helexa-Advice` header only, never the error envelope, and behaviour
|
||||||
|
/// never depends on it. Unknown clients get nothing.
|
||||||
|
fn client_advice(headers: &HeaderMap) -> Option<&'static str> {
|
||||||
|
let ua = headers
|
||||||
|
.get(axum::http::header::USER_AGENT)?
|
||||||
|
.to_str()
|
||||||
|
.ok()?
|
||||||
|
.to_ascii_lowercase();
|
||||||
|
if ua.contains("litellm") {
|
||||||
|
Some(
|
||||||
|
"litellm forwards the full context; lower the configured context window or enable client-side compaction",
|
||||||
|
)
|
||||||
|
} else if ua.contains("agent-zero") || ua.contains("agent zero") {
|
||||||
|
Some("reduce the conversation/context size or summarize earlier turns before resending")
|
||||||
|
} else if ua.contains("zed") {
|
||||||
|
Some("reduce the assistant context window in Zed's settings")
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `400 context_length_exceeded` for an over-long prompt caught at the edge
|
||||||
|
/// (#56), in the #60 envelope — the same shape neuron emits on overflow, so
|
||||||
|
/// clients (opencode auto-compacts) handle it identically. Attaches the
|
||||||
|
/// advisory `X-Helexa-Advice` header for fingerprinted clients.
|
||||||
|
fn context_length_exceeded_response(
|
||||||
|
context: u64,
|
||||||
|
prompt_est: u64,
|
||||||
|
headers: &HeaderMap,
|
||||||
|
) -> Response {
|
||||||
|
let env = OpenAiError::context_length_exceeded(format!(
|
||||||
|
"This model's maximum context length is {context} tokens. Your request is \
|
||||||
|
estimated at ~{prompt_est} tokens. Please reduce the length of the messages."
|
||||||
|
))
|
||||||
|
.with_extra("max", json!(context))
|
||||||
|
.with_extra("estimated_prompt_tokens", json!(prompt_est));
|
||||||
|
let mut response = crate::error::envelope_response(env);
|
||||||
|
if let Some(advice) = client_advice(headers)
|
||||||
|
&& let Ok(value) = axum::http::HeaderValue::from_str(advice)
|
||||||
|
{
|
||||||
|
response.headers_mut().insert("x-helexa-advice", value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
}
|
||||||
|
|
||||||
/// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
|
/// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
|
||||||
async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
|
async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
|
||||||
let mut nodes = fleet.nodes.write().await;
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
|||||||
@@ -1,25 +1,41 @@
|
|||||||
pub mod anthropic_sse;
|
pub mod anthropic_sse;
|
||||||
|
pub mod auth;
|
||||||
|
pub mod entitlements_chain;
|
||||||
pub mod entitlements_local;
|
pub mod entitlements_local;
|
||||||
|
pub mod entitlements_upstream;
|
||||||
pub mod error;
|
pub mod error;
|
||||||
pub mod evictor;
|
pub mod evictor;
|
||||||
pub mod handlers;
|
pub mod handlers;
|
||||||
|
pub mod metering;
|
||||||
pub mod metrics;
|
pub mod metrics;
|
||||||
pub mod poller;
|
pub mod poller;
|
||||||
pub mod proxy;
|
pub mod proxy;
|
||||||
pub mod router;
|
pub mod router;
|
||||||
|
pub mod served_usage;
|
||||||
pub mod state;
|
pub mod state;
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
|
use axum::middleware::from_fn_with_state;
|
||||||
use cortex_core::config::GatewayConfig;
|
use cortex_core::config::GatewayConfig;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tower_http::cors::CorsLayer;
|
use tower_http::cors::CorsLayer;
|
||||||
use tower_http::trace::TraceLayer;
|
use tower_http::trace::TraceLayer;
|
||||||
|
|
||||||
/// Build the Axum application router with all routes wired up.
|
/// Build the Axum application router with all routes wired up.
|
||||||
|
///
|
||||||
|
/// Layer order (outermost first): trace → CORS → auth → handlers. CORS is
|
||||||
|
/// outer to auth so preflight `OPTIONS` short-circuits before resolution;
|
||||||
|
/// auth (`require_principal`) resolves the bearer key, attaches the
|
||||||
|
/// principal, and stamps the internal principal headers before any handler
|
||||||
|
/// runs.
|
||||||
pub fn build_app(fleet: Arc<state::CortexState>) -> Router {
|
pub fn build_app(fleet: Arc<state::CortexState>) -> Router {
|
||||||
Router::new()
|
Router::new()
|
||||||
.merge(handlers::api_routes())
|
.merge(handlers::api_routes())
|
||||||
|
.layer(from_fn_with_state(
|
||||||
|
Arc::clone(&fleet),
|
||||||
|
auth::require_principal,
|
||||||
|
))
|
||||||
.layer(CorsLayer::permissive())
|
.layer(CorsLayer::permissive())
|
||||||
.layer(TraceLayer::new_for_http())
|
.layer(TraceLayer::new_for_http())
|
||||||
.with_state(fleet)
|
.with_state(fleet)
|
||||||
@@ -42,6 +58,28 @@ pub async fn run(config: GatewayConfig) -> Result<()> {
|
|||||||
evictor::eviction_loop(evictor_fleet).await;
|
evictor::eviction_loop(evictor_fleet).await;
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Served-usage reporter (#58): when this operator is part of the mesh,
|
||||||
|
// periodically flush absolute per-principal served-token counters to
|
||||||
|
// upstream for reconciliation.
|
||||||
|
if config.upstream.enabled {
|
||||||
|
let su_fleet = Arc::clone(&fleet);
|
||||||
|
let url = config.upstream.url.clone();
|
||||||
|
let bearer = config.upstream.bearer.clone();
|
||||||
|
let interval =
|
||||||
|
std::time::Duration::from_secs(config.upstream.served_usage_report_interval_secs);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
tokio::time::sleep(interval).await;
|
||||||
|
let rows = su_fleet.served_usage.snapshot();
|
||||||
|
if let Err(e) =
|
||||||
|
served_usage::report(&su_fleet.http_client, &url, &bearer, &rows).await
|
||||||
|
{
|
||||||
|
tracing::warn!(error = %e, "served-usage report failed (will retry)");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
let app = build_app(Arc::clone(&fleet));
|
let app = build_app(Arc::clone(&fleet));
|
||||||
|
|
||||||
let listen_addr = config.gateway.listen.parse::<std::net::SocketAddr>()?;
|
let listen_addr = config.gateway.listen.parse::<std::net::SocketAddr>()?;
|
||||||
|
|||||||
231
crates/cortex-gateway/src/metering.rs
Normal file
231
crates/cortex-gateway/src/metering.rs
Normal file
@@ -0,0 +1,231 @@
|
|||||||
|
//! Per-request token metering (#51).
|
||||||
|
//!
|
||||||
|
//! Captures the real `(prompt, completion)` usage of every request and feeds
|
||||||
|
//! it to two places: the [`EntitlementProvider`] spend ledger (via
|
||||||
|
//! reserve→settle) and per-principal Prometheus counters. The principal is
|
||||||
|
//! reconstructed from the internal headers the auth middleware stamped (#49),
|
||||||
|
//! so this works uniformly across every proxy path without threading the
|
||||||
|
//! typed principal through each handler.
|
||||||
|
//!
|
||||||
|
//! The reserve→settle lifecycle is established here but, in this phase,
|
||||||
|
//! reserves **zero** tokens — metering only, no enforcement. Budget
|
||||||
|
//! enforcement (#52) flips the reserved amount to the real
|
||||||
|
//! `prompt + max_output` and handles the [`BudgetError`] rejection; the
|
||||||
|
//! settle/release plumbing is identical, so that change is localized.
|
||||||
|
//!
|
||||||
|
//! [`ReservationGuard`] makes leaks impossible: settling records actual
|
||||||
|
//! spend and releases the unused remainder; dropping a guard that was never
|
||||||
|
//! settled releases the whole reservation. So an early return, error path,
|
||||||
|
//! or dropped stream can't strand a reservation.
|
||||||
|
|
||||||
|
use axum::http::HeaderMap;
|
||||||
|
use cortex_core::entitlements::{
|
||||||
|
BudgetError, EntitlementProvider, HEADER_ACCOUNT_ID, HEADER_KEY_ID, Principal,
|
||||||
|
};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Fallback output-token budget when neither the request nor the model's
|
||||||
|
/// advertised limit gives one. Bounds the reservation so a capped key is
|
||||||
|
/// still gated even on under-specified requests (#52).
|
||||||
|
pub const FALLBACK_MAX_OUTPUT: u64 = 4096;
|
||||||
|
|
||||||
|
/// Invoked exactly once at request completion with best-effort
|
||||||
|
/// `(prompt_tokens, completion_tokens)`. When no usage could be observed
|
||||||
|
/// (e.g. a pre-dispatch failure or a dropped stream) it is dropped unused —
|
||||||
|
/// which releases the held reservation via [`ReservationGuard`]'s `Drop`.
|
||||||
|
pub type UsageSink = Box<dyn FnOnce(u64, u64) + Send>;
|
||||||
|
|
||||||
|
/// Reconstruct the principal from the cortex-stamped internal headers. The
|
||||||
|
/// auth middleware strips any client copy and stamps the authoritative value,
|
||||||
|
/// so these headers are trustworthy within cortex. `None` for anonymous
|
||||||
|
/// (unauthenticated) requests.
|
||||||
|
pub fn principal_from_headers(headers: &HeaderMap) -> Option<Principal> {
|
||||||
|
let account_id = headers.get(HEADER_ACCOUNT_ID)?.to_str().ok()?.to_string();
|
||||||
|
let key_id = headers.get(HEADER_KEY_ID)?.to_str().ok()?.to_string();
|
||||||
|
Some(Principal { account_id, key_id })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Emit per-principal spend counters (#51). Labelled by account/key only —
|
||||||
|
/// both are operator-bounded, so cardinality is controlled.
|
||||||
|
pub fn record_spend(principal: &Principal, prompt: u64, completion: u64) {
|
||||||
|
let labels = [
|
||||||
|
("account", principal.account_id.clone()),
|
||||||
|
("key", principal.key_id.clone()),
|
||||||
|
];
|
||||||
|
metrics::counter!("cortex_spend_tokens_total", &labels).increment(prompt + completion);
|
||||||
|
metrics::counter!("cortex_spend_prompt_tokens_total", &labels).increment(prompt);
|
||||||
|
metrics::counter!("cortex_spend_completion_tokens_total", &labels).increment(completion);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Holds a budget reservation for the life of a request. [`settle`] records
|
||||||
|
/// actual spend and releases the remainder; an un-settled guard releases the
|
||||||
|
/// whole reservation when dropped. Anonymous requests carry an empty guard,
|
||||||
|
/// where every operation is a no-op.
|
||||||
|
///
|
||||||
|
/// [`settle`]: ReservationGuard::settle
|
||||||
|
pub struct ReservationGuard {
|
||||||
|
provider: Arc<dyn EntitlementProvider>,
|
||||||
|
reservation: Option<cortex_core::entitlements::Reservation>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ReservationGuard {
|
||||||
|
/// An empty guard for an anonymous request — no reservation to resolve.
|
||||||
|
pub fn anonymous(provider: Arc<dyn EntitlementProvider>) -> Self {
|
||||||
|
Self {
|
||||||
|
provider,
|
||||||
|
reservation: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Wrap an already-acquired reservation.
|
||||||
|
fn held(
|
||||||
|
provider: Arc<dyn EntitlementProvider>,
|
||||||
|
reservation: cortex_core::entitlements::Reservation,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
provider,
|
||||||
|
reservation: Some(reservation),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Settle with the tokens actually consumed, disarming the drop-release.
|
||||||
|
/// Spawns the (fast, in-process for the local provider) settle so the
|
||||||
|
/// caller — which may be a sync stream-completion callback — needn't
|
||||||
|
/// await.
|
||||||
|
pub fn settle(mut self, actual_tokens: u64) {
|
||||||
|
if let Some(reservation) = self.reservation.take() {
|
||||||
|
let provider = Arc::clone(&self.provider);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
provider.settle(reservation, actual_tokens).await;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ReservationGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if let Some(reservation) = self.reservation.take() {
|
||||||
|
let provider = Arc::clone(&self.provider);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
provider.release(reservation).await;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build the completion sink for an authenticated request: record spend and
|
||||||
|
/// settle the reservation with the observed total. Dropping it unused (no
|
||||||
|
/// usage observed) releases the reservation via the guard.
|
||||||
|
pub fn usage_sink(
|
||||||
|
principal: Principal,
|
||||||
|
guard: ReservationGuard,
|
||||||
|
served_usage: std::sync::Arc<crate::served_usage::ServedUsage>,
|
||||||
|
) -> UsageSink {
|
||||||
|
Box::new(move |prompt, completion| {
|
||||||
|
record_spend(&principal, prompt, completion);
|
||||||
|
// Per-principal served-usage tally for #58 reconciliation. Recorded
|
||||||
|
// for every metered (authenticated) request; the flush task reports
|
||||||
|
// it to upstream when the operator is part of the mesh.
|
||||||
|
served_usage.add(
|
||||||
|
&principal.account_id,
|
||||||
|
&principal.key_id,
|
||||||
|
prompt + completion,
|
||||||
|
);
|
||||||
|
guard.settle(prompt + completion);
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reserve the request's upper-bound token cost for the principal, refusing
|
||||||
|
/// *before* dispatch if it would exceed the hard cap (#52). On success
|
||||||
|
/// returns a guard the caller settles with actual usage; on refusal returns
|
||||||
|
/// the #63 envelope (`rate_limit_exceeded` + `Retry-After` for a resetting
|
||||||
|
/// window, `insufficient_quota` for a hard balance — never `402`).
|
||||||
|
pub async fn reserve_or_reject(
|
||||||
|
provider: Arc<dyn EntitlementProvider>,
|
||||||
|
principal: &Principal,
|
||||||
|
max_tokens: u64,
|
||||||
|
) -> Result<ReservationGuard, OpenAiError> {
|
||||||
|
match provider.reserve(principal, max_tokens).await {
|
||||||
|
Ok(reservation) => Ok(ReservationGuard::held(provider, reservation)),
|
||||||
|
Err(err) => Err(budget_error_to_envelope(err)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Map a [`BudgetError`] to the #63 envelope. The provider chose the window
|
||||||
|
/// semantics; this only translates them to HTTP.
|
||||||
|
fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
|
||||||
|
match err {
|
||||||
|
BudgetError::RateLimited {
|
||||||
|
retry_after_secs, ..
|
||||||
|
} => OpenAiError::rate_limit_exceeded(err.to_string(), retry_after_secs),
|
||||||
|
BudgetError::InsufficientQuota { .. } => OpenAiError::insufficient_quota(err.to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Upper-bound tokens to reserve for a request (#52): an over-estimate of
|
||||||
|
/// the prompt plus the maximum output. `advertised_output` is the model's
|
||||||
|
/// `limit.output` (#62), used when the request omits `max_(completion_)tokens`.
|
||||||
|
/// Over-reserving is safe — settle corrects spend to the actual usage.
|
||||||
|
pub fn reservation_estimate(body: &[u8], advertised_output: Option<u64>) -> u64 {
|
||||||
|
let max_output = requested_max_output(body)
|
||||||
|
.or(advertised_output)
|
||||||
|
.unwrap_or(FALLBACK_MAX_OUTPUT);
|
||||||
|
estimate_prompt_tokens(body).saturating_add(max_output)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The client's requested output cap, from `max_completion_tokens` (or the
|
||||||
|
/// legacy `max_tokens`). `None` when unspecified.
|
||||||
|
fn requested_max_output(body: &[u8]) -> Option<u64> {
|
||||||
|
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||||
|
v.get("max_completion_tokens")
|
||||||
|
.or_else(|| v.get("max_tokens"))
|
||||||
|
.and_then(serde_json::Value::as_u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rough prompt-token estimate at ~4 chars/token over the whole body. cortex
|
||||||
|
/// has no tokenizer; JSON overhead makes this a conservative over-estimate,
|
||||||
|
/// and neuron remains the exact context wall (#56/#60). Settle reconciles to
|
||||||
|
/// the real usage afterward.
|
||||||
|
fn estimate_prompt_tokens(body: &[u8]) -> u64 {
|
||||||
|
(body.len() as u64 / 4).max(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn requested_max_output_prefers_max_completion_tokens() {
|
||||||
|
let body = br#"{"model":"m","max_completion_tokens":256,"max_tokens":99}"#;
|
||||||
|
assert_eq!(requested_max_output(body), Some(256));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn requested_max_output_falls_back_to_legacy_max_tokens() {
|
||||||
|
let body = br#"{"model":"m","max_tokens":128}"#;
|
||||||
|
assert_eq!(requested_max_output(body), Some(128));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn estimate_uses_requested_output_when_present() {
|
||||||
|
// Requested output dominates; prompt estimate is small for a tiny body.
|
||||||
|
let body = br#"{"model":"m","max_tokens":1000}"#;
|
||||||
|
let est = reservation_estimate(body, Some(8192));
|
||||||
|
assert!(est >= 1000 && est < 1100, "est was {est}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn estimate_uses_advertised_output_when_request_omits_it() {
|
||||||
|
let body = br#"{"model":"m","messages":[]}"#;
|
||||||
|
let est = reservation_estimate(body, Some(8192));
|
||||||
|
assert!(est >= 8192, "est was {est}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn estimate_falls_back_when_nothing_advertised() {
|
||||||
|
let body = br#"{"model":"m"}"#;
|
||||||
|
let est = reservation_estimate(body, None);
|
||||||
|
assert!(est >= FALLBACK_MAX_OUTPUT, "est was {est}");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -63,4 +63,16 @@ fn describe_metrics() {
|
|||||||
"cortex_cold_starts_total",
|
"cortex_cold_starts_total",
|
||||||
"Total number of cold-start model loads"
|
"Total number of cold-start model loads"
|
||||||
);
|
);
|
||||||
|
metrics::describe_counter!(
|
||||||
|
"cortex_spend_tokens_total",
|
||||||
|
"Total metered tokens (prompt + completion) per principal, labelled by account/key (#51)"
|
||||||
|
);
|
||||||
|
metrics::describe_counter!(
|
||||||
|
"cortex_spend_prompt_tokens_total",
|
||||||
|
"Metered prompt tokens per principal, labelled by account/key (#51)"
|
||||||
|
);
|
||||||
|
metrics::describe_counter!(
|
||||||
|
"cortex_spend_completion_tokens_total",
|
||||||
|
"Metered completion tokens per principal, labelled by account/key (#51)"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,12 +5,29 @@ use crate::state::CortexState;
|
|||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||||
use cortex_core::harness::ModelInfo;
|
use cortex_core::harness::ModelInfo;
|
||||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
use cortex_core::node::{ModelEntry, ModelStatus, NodeState};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
const POLL_INTERVAL: Duration = Duration::from_secs(10);
|
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.
|
/// Runs forever, polling all neurons on a fixed interval.
|
||||||
pub async fn poll_loop(fleet: Arc<CortexState>) {
|
pub async fn poll_loop(fleet: Arc<CortexState>) {
|
||||||
loop {
|
loop {
|
||||||
@@ -138,13 +155,14 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
|
|||||||
// Remove models no longer reported by the neuron.
|
// Remove models no longer reported by the neuron.
|
||||||
node.models.retain(|id, _| seen.contains(id));
|
node.models.retain(|id, _| seen.contains(id));
|
||||||
|
|
||||||
|
node.consecutive_poll_failures = 0;
|
||||||
node.healthy = true;
|
node.healthy = true;
|
||||||
node.last_poll = Some(Utc::now());
|
node.last_poll = Some(Utc::now());
|
||||||
tracing::debug!(node = name, models = models.len(), "poll ok");
|
tracing::debug!(node = name, models = models.len(), "poll ok");
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!(node = name, error = %e, "failed to parse /models response");
|
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(),
|
status = %resp.status(),
|
||||||
"neuron returned non-success status"
|
"neuron returned non-success status"
|
||||||
);
|
);
|
||||||
node.healthy = false;
|
record_poll_failure(node);
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!(node = name, error = %e, "failed to reach neuron");
|
tracing::warn!(node = name, error = %e, "failed to reach neuron");
|
||||||
node.healthy = false;
|
record_poll_failure(node);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -200,6 +218,9 @@ async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
|
|||||||
let mut nodes = fleet.nodes.write().await;
|
let mut nodes = fleet.nodes.write().await;
|
||||||
if let Some(node) = nodes.get_mut(name) {
|
if let Some(node) = nodes.get_mut(name) {
|
||||||
node.activation = Some(h.activation);
|
node.activation = Some(h.activation);
|
||||||
|
// Per-model admission load (#53) → keyed by id for the
|
||||||
|
// load-aware router (#55).
|
||||||
|
node.model_load = h.models.into_iter().map(|m| (m.id.clone(), m)).collect();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
|||||||
@@ -1,21 +1,27 @@
|
|||||||
//! Streaming HTTP reverse proxy to neuron backends.
|
//! Streaming HTTP reverse proxy to neuron backends.
|
||||||
//!
|
//!
|
||||||
//! For streaming requests, SSE chunks are forwarded as they arrive.
|
//! The streaming *mechanism* — forward an SSE body chunk-for-chunk without
|
||||||
//! The proxy captures timing information for metrics but does not
|
//! buffering, observing the bytes for metrics — lives in the shared
|
||||||
//! buffer the full response.
|
//! [`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 crate::router::RouteDecision;
|
||||||
use anyhow::Result;
|
use axum::http::HeaderMap;
|
||||||
use axum::body::Body;
|
use axum::http::StatusCode;
|
||||||
use axum::http::{HeaderMap, StatusCode};
|
|
||||||
use axum::response::{IntoResponse, Response};
|
use axum::response::{IntoResponse, Response};
|
||||||
use futures::Stream;
|
use helexa_stream::{BodyTail, ChunkObserver, StreamError};
|
||||||
use futures::stream::BoxStream;
|
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use std::pin::Pin;
|
|
||||||
use std::task::{Context, Poll};
|
|
||||||
use std::time::Instant;
|
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.
|
/// Proxy a request body to the resolved backend node and stream the response.
|
||||||
///
|
///
|
||||||
/// Logging contract: every call emits exactly one structured event at
|
/// Logging contract: every call emits exactly one structured event at
|
||||||
@@ -31,6 +37,7 @@ pub async fn forward_request(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
body: bytes::Bytes,
|
body: bytes::Bytes,
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
|
usage_sink: Option<crate::metering::UsageSink>,
|
||||||
) -> Result<Response, ProxyError> {
|
) -> Result<Response, ProxyError> {
|
||||||
let request_start = Instant::now();
|
let request_start = Instant::now();
|
||||||
let url = format!("{}{}", route.endpoint, path);
|
let url = format!("{}{}", route.endpoint, path);
|
||||||
@@ -41,66 +48,41 @@ pub async fn forward_request(
|
|||||||
"proxying 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.
|
let response = helexa_stream::forward_streaming(client, &url, headers, body, observer)
|
||||||
for (key, value) in headers.iter() {
|
.await
|
||||||
if key == "host" || key == "content-length" {
|
.map_err(|e| {
|
||||||
continue; // reqwest sets these
|
match &e {
|
||||||
}
|
StreamError::Upstream(err) => tracing::warn!(
|
||||||
req_builder = req_builder.header(key, value);
|
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 {
|
if !response.status().is_success() {
|
||||||
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() {
|
|
||||||
// Streaming body — can't snippet without breaking the stream
|
// Streaming body — can't snippet without breaking the stream
|
||||||
// pass-through. Log status + URL; the client still gets the
|
// pass-through. Log status + URL; the client still gets the
|
||||||
// upstream status, just without the leaked body.
|
// upstream status, just without the leaked body.
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
node = %route.node_name,
|
node = %route.node_name,
|
||||||
url = %url,
|
url = %url,
|
||||||
status = upstream_status.as_u16(),
|
status = response.status().as_u16(),
|
||||||
"proxy: upstream returned non-2xx"
|
"proxy: upstream returned non-2xx"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
Ok(response)
|
||||||
|
|
||||||
let resp_headers = upstream_resp.headers().clone();
|
|
||||||
let stream = TokenMetricsStream::new(
|
|
||||||
Box::pin(upstream_resp.bytes_stream()),
|
|
||||||
TokenMetrics::new(model_id, &route.node_name, request_start),
|
|
||||||
);
|
|
||||||
|
|
||||||
let body = Body::from_stream(stream);
|
|
||||||
|
|
||||||
let mut response = Response::builder().status(status);
|
|
||||||
for (key, value) in resp_headers.iter() {
|
|
||||||
response = response.header(key, value);
|
|
||||||
}
|
|
||||||
|
|
||||||
response.body(body).map_err(|e| {
|
|
||||||
tracing::warn!(
|
|
||||||
node = %route.node_name,
|
|
||||||
url = %url,
|
|
||||||
error = %e,
|
|
||||||
"proxy: failed to build response"
|
|
||||||
);
|
|
||||||
ProxyError::ResponseBuild(e.to_string())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
@@ -111,6 +93,15 @@ pub enum ProxyError {
|
|||||||
ResponseBuild(String),
|
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 {
|
impl IntoResponse for ProxyError {
|
||||||
fn into_response(self) -> Response {
|
fn into_response(self) -> Response {
|
||||||
let (status, code, message) = match &self {
|
let (status, code, message) = match &self {
|
||||||
@@ -138,9 +129,10 @@ impl IntoResponse for ProxyError {
|
|||||||
//
|
//
|
||||||
// The proxy never buffers or re-serialises the upstream body — chunks
|
// The proxy never buffers or re-serialises the upstream body — chunks
|
||||||
// are forwarded verbatim. For metrics it observes each chunk's arrival
|
// 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
|
// time and keeps a bounded tail of the body text (via the shared
|
||||||
// OpenAI `usage` object (present on the last SSE chunk and on
|
// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
|
||||||
// non-streaming JSON bodies alike) yields engine-truth token counts.
|
// (present on the last SSE chunk and on non-streaming JSON bodies alike)
|
||||||
|
// yields engine-truth token counts.
|
||||||
//
|
//
|
||||||
// Emitted per request, labelled {model, node}:
|
// Emitted per request, labelled {model, node}:
|
||||||
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
||||||
@@ -154,42 +146,29 @@ impl IntoResponse for ProxyError {
|
|||||||
/// non-streaming bodies.
|
/// non-streaming bodies.
|
||||||
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
||||||
|
|
||||||
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
/// cortex's [`ChunkObserver`]: per-request token metrics plus the
|
||||||
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
/// per-principal reservation settle. Drives cortex policy over the shared
|
||||||
/// The quoted-needle form means `completion_tokens` never matches
|
/// streaming mechanism.
|
||||||
/// `completion_tokens_details`.
|
struct CortexMetrics {
|
||||||
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 {
|
|
||||||
labels: [(&'static str, String); 2],
|
labels: [(&'static str, String); 2],
|
||||||
request_start: Instant,
|
request_start: Instant,
|
||||||
first_chunk: Option<Instant>,
|
first_chunk: Option<Instant>,
|
||||||
last_chunk: Option<Instant>,
|
last_chunk: Option<Instant>,
|
||||||
tail: String,
|
tail: BodyTail,
|
||||||
finished: bool,
|
finished: bool,
|
||||||
|
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
|
||||||
|
/// with the observed `(prompt, completion)` so the reservation can be
|
||||||
|
/// settled and spend recorded. `None` for anonymous requests.
|
||||||
|
usage_sink: Option<crate::metering::UsageSink>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TokenMetrics {
|
impl CortexMetrics {
|
||||||
fn new(model_id: &str, node_name: &str, request_start: Instant) -> Self {
|
fn new(
|
||||||
|
model_id: &str,
|
||||||
|
node_name: &str,
|
||||||
|
request_start: Instant,
|
||||||
|
usage_sink: Option<crate::metering::UsageSink>,
|
||||||
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
labels: [
|
labels: [
|
||||||
("model", model_id.to_string()),
|
("model", model_id.to_string()),
|
||||||
@@ -198,25 +177,19 @@ impl TokenMetrics {
|
|||||||
request_start,
|
request_start,
|
||||||
first_chunk: None,
|
first_chunk: None,
|
||||||
last_chunk: None,
|
last_chunk: None,
|
||||||
tail: String::new(),
|
tail: BodyTail::new(TAIL_CAP_BYTES),
|
||||||
finished: false,
|
finished: false,
|
||||||
|
usage_sink,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkObserver for CortexMetrics {
|
||||||
fn observe(&mut self, chunk: &[u8]) {
|
fn observe(&mut self, chunk: &[u8]) {
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
self.first_chunk.get_or_insert(now);
|
self.first_chunk.get_or_insert(now);
|
||||||
self.last_chunk = Some(now);
|
self.last_chunk = Some(now);
|
||||||
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
self.tail.push(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);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Emit the metrics exactly once — called on clean stream end and
|
/// Emit the metrics exactly once — called on clean stream end and
|
||||||
@@ -227,130 +200,45 @@ impl TokenMetrics {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
self.finished = true;
|
self.finished = true;
|
||||||
let Some(first) = self.first_chunk else {
|
|
||||||
return; // no body ever arrived — nothing to record
|
|
||||||
};
|
|
||||||
let ttft = first.duration_since(self.request_start).as_secs_f64();
|
|
||||||
metrics::histogram!("cortex_time_to_first_token_seconds", &self.labels).record(ttft);
|
|
||||||
|
|
||||||
if let Some(prompt) = last_count_for(&self.tail, "prompt_tokens") {
|
let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
|
||||||
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
|
let completion = last_count_for(self.tail.as_str(), "completion_tokens");
|
||||||
}
|
|
||||||
let Some(completion) = last_count_for(&self.tail, "completion_tokens") else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
if completion == 0 {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
metrics::counter!("cortex_completion_tokens_total", &self.labels).increment(completion);
|
|
||||||
|
|
||||||
let last = self.last_chunk.unwrap_or(first);
|
// Per-model metrics — only when body chunks actually arrived.
|
||||||
let decode_window = last.duration_since(first).as_secs_f64();
|
if let Some(first) = self.first_chunk {
|
||||||
// Streaming: rate over the decode window (first→last chunk).
|
let ttft = first.duration_since(self.request_start).as_secs_f64();
|
||||||
// Non-streaming bodies arrive as ~one chunk (window ≈ 0), where
|
metrics::histogram!("cortex_time_to_first_token_seconds", &self.labels).record(ttft);
|
||||||
// the only honest denominator is the full request duration.
|
|
||||||
let secs = if decode_window >= 0.1 {
|
|
||||||
decode_window
|
|
||||||
} else {
|
|
||||||
last.duration_since(self.request_start).as_secs_f64()
|
|
||||||
};
|
|
||||||
if secs > 0.0 {
|
|
||||||
metrics::histogram!("cortex_tokens_per_second", &self.labels)
|
|
||||||
.record(completion as f64 / secs);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
|
if let Some(prompt) = prompt {
|
||||||
/// clean end-of-stream; the Drop impl covers client disconnects.
|
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
|
||||||
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))),
|
if let Some(completion) = completion.filter(|c| *c > 0) {
|
||||||
Poll::Ready(None) => {
|
metrics::counter!("cortex_completion_tokens_total", &self.labels)
|
||||||
this.metrics.finish();
|
.increment(completion);
|
||||||
Poll::Ready(None)
|
|
||||||
|
let last = self.last_chunk.unwrap_or(first);
|
||||||
|
let decode_window = last.duration_since(first).as_secs_f64();
|
||||||
|
// Streaming: rate over the decode window (first→last chunk).
|
||||||
|
// Non-streaming bodies arrive as ~one chunk (window ≈ 0),
|
||||||
|
// where the only honest denominator is the full request
|
||||||
|
// duration.
|
||||||
|
let secs = if decode_window >= 0.1 {
|
||||||
|
decode_window
|
||||||
|
} else {
|
||||||
|
last.duration_since(self.request_start).as_secs_f64()
|
||||||
|
};
|
||||||
|
if secs > 0.0 {
|
||||||
|
metrics::histogram!("cortex_tokens_per_second", &self.labels)
|
||||||
|
.record(completion as f64 / secs);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Poll::Pending => Poll::Pending,
|
}
|
||||||
|
|
||||||
|
// Per-principal metering + reservation settle (#51). Always runs so
|
||||||
|
// the reservation is resolved even when no usage/body was observed
|
||||||
|
// (sink with (0, 0) → settle 0 → release).
|
||||||
|
if let Some(sink) = self.usage_sink.take() {
|
||||||
|
sink(prompt.unwrap_or(0), completion.unwrap_or(0));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for TokenMetricsStream {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
self.metrics.finish();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::last_count_for;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extracts_counts_from_final_sse_usage_chunk() {
|
|
||||||
let tail = concat!(
|
|
||||||
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
|
||||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
|
||||||
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
|
||||||
"data: [DONE]\n\n"
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn extracts_counts_from_non_streaming_body() {
|
|
||||||
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
|
||||||
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
|
||||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn ignores_details_variants_and_takes_last_occurrence() {
|
|
||||||
// completion_tokens_details must not shadow completion_tokens,
|
|
||||||
// and the LAST usage object wins (matters when content echoes
|
|
||||||
// a usage-shaped string earlier in the stream).
|
|
||||||
let tail = concat!(
|
|
||||||
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
|
||||||
"data: {\"usage\":{\"completion_tokens\":99,",
|
|
||||||
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn absent_keys_yield_none() {
|
|
||||||
assert_eq!(
|
|
||||||
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
|
||||||
None
|
|
||||||
);
|
|
||||||
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
|
||||||
// key present but non-numeric value
|
|
||||||
assert_eq!(
|
|
||||||
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
|
||||||
None
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -50,6 +50,10 @@ pub enum RouteError {
|
|||||||
"model '{model_id}' is in the catalogue but no healthy neuron's topology satisfies its constraints"
|
"model '{model_id}' is in the catalogue but no healthy neuron's topology satisfies its constraints"
|
||||||
)]
|
)]
|
||||||
NoFeasibleNeuron { model_id: String },
|
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}")]
|
#[error("cold-load of '{model_id}' on '{node}' failed: {message}")]
|
||||||
ColdLoadFailed {
|
ColdLoadFailed {
|
||||||
model_id: String,
|
model_id: String,
|
||||||
@@ -68,7 +72,9 @@ impl RouteError {
|
|||||||
/// safe to retry the same request); everything else is 404.
|
/// safe to retry the same request); everything else is 404.
|
||||||
pub fn http_status(&self) -> u16 {
|
pub fn http_status(&self) -> u16 {
|
||||||
match self {
|
match self {
|
||||||
RouteError::NoHealthyNodes | RouteError::ModelRecovering { .. } => 503,
|
RouteError::NoHealthyNodes
|
||||||
|
| RouteError::ModelRecovering { .. }
|
||||||
|
| RouteError::FeasibleNodeUnhealthy { .. } => 503,
|
||||||
_ => 404,
|
_ => 404,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -81,7 +87,8 @@ impl RouteError {
|
|||||||
| RouteError::EndpointResolveFailed(_, _)
|
| RouteError::EndpointResolveFailed(_, _)
|
||||||
| RouteError::NoFeasibleNeuron { .. }
|
| RouteError::NoFeasibleNeuron { .. }
|
||||||
| RouteError::ColdLoadFailed { .. }
|
| RouteError::ColdLoadFailed { .. }
|
||||||
| RouteError::ModelRecovering { .. } => "api_error",
|
| RouteError::ModelRecovering { .. }
|
||||||
|
| RouteError::FeasibleNodeUnhealthy { .. } => "api_error",
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -94,6 +101,7 @@ impl RouteError {
|
|||||||
RouteError::NoFeasibleNeuron { .. } => "service_unavailable",
|
RouteError::NoFeasibleNeuron { .. } => "service_unavailable",
|
||||||
RouteError::ColdLoadFailed { .. } => "service_unavailable",
|
RouteError::ColdLoadFailed { .. } => "service_unavailable",
|
||||||
RouteError::ModelRecovering { .. } => "service_unavailable",
|
RouteError::ModelRecovering { .. } => "service_unavailable",
|
||||||
|
RouteError::FeasibleNodeUnhealthy { .. } => "service_unavailable",
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -105,6 +113,7 @@ impl RouteError {
|
|||||||
pub fn retry_after_secs(&self) -> Option<u64> {
|
pub fn retry_after_secs(&self) -> Option<u64> {
|
||||||
match self {
|
match self {
|
||||||
RouteError::ModelRecovering { .. } => Some(2),
|
RouteError::ModelRecovering { .. } => Some(2),
|
||||||
|
RouteError::FeasibleNodeUnhealthy { .. } => Some(3),
|
||||||
RouteError::NoHealthyNodes => Some(5),
|
RouteError::NoHealthyNodes => Some(5),
|
||||||
_ => None,
|
_ => None,
|
||||||
}
|
}
|
||||||
@@ -132,7 +141,9 @@ pub async fn resolve(
|
|||||||
// Snapshot loaded / unloaded / recovering state from the poller cache.
|
// Snapshot loaded / unloaded / recovering state from the poller cache.
|
||||||
let (loaded_route, unloaded_route, recovering_node, any_healthy) = {
|
let (loaded_route, unloaded_route, recovering_node, any_healthy) = {
|
||||||
let nodes = fleet.nodes.read().await;
|
let nodes = fleet.nodes.read().await;
|
||||||
let mut loaded_route = None;
|
// All healthy nodes with the model loaded, each with its current
|
||||||
|
// admission load (#53) so we can pick the least-busy replica (#55).
|
||||||
|
let mut loaded_candidates: Vec<(String, String, usize)> = Vec::new();
|
||||||
let mut unloaded_route = None;
|
let mut unloaded_route = None;
|
||||||
let mut recovering_node = None;
|
let mut recovering_node = None;
|
||||||
let mut any_healthy = false;
|
let mut any_healthy = false;
|
||||||
@@ -144,8 +155,15 @@ pub async fn resolve(
|
|||||||
if let Some(entry) = node.models.get(model_id) {
|
if let Some(entry) = node.models.get(model_id) {
|
||||||
match entry.status {
|
match entry.status {
|
||||||
ModelStatus::Loaded | ModelStatus::Reloading => {
|
ModelStatus::Loaded | ModelStatus::Reloading => {
|
||||||
loaded_route = Some((node.name.clone(), node.endpoint.clone(), false));
|
// Least-busy score: in-flight + queued from the
|
||||||
break;
|
// neuron's last /health (#53). Unknown load (no poll
|
||||||
|
// yet) scores 0 so the replica stays eligible.
|
||||||
|
let score = node
|
||||||
|
.model_load
|
||||||
|
.get(model_id)
|
||||||
|
.map(|l| l.in_flight + l.queue_depth)
|
||||||
|
.unwrap_or(0);
|
||||||
|
loaded_candidates.push((node.name.clone(), node.endpoint.clone(), score));
|
||||||
}
|
}
|
||||||
ModelStatus::Unloaded => {
|
ModelStatus::Unloaded => {
|
||||||
if unloaded_route.is_none() {
|
if unloaded_route.is_none() {
|
||||||
@@ -175,6 +193,12 @@ pub async fn resolve(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// Pick the least-busy loaded replica; ties break by node name for
|
||||||
|
// deterministic routing. `false` = not a cold start.
|
||||||
|
let loaded_route = loaded_candidates
|
||||||
|
.into_iter()
|
||||||
|
.min_by(|a, b| a.2.cmp(&b.2).then_with(|| a.0.cmp(&b.0)))
|
||||||
|
.map(|(name, endpoint, _score)| (name, endpoint, false));
|
||||||
(loaded_route, unloaded_route, recovering_node, any_healthy)
|
(loaded_route, unloaded_route, recovering_node, any_healthy)
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -237,11 +261,32 @@ async fn pick_feasible_neuron(
|
|||||||
b.2.cmp(&a.2) // pinned first (true > false)
|
b.2.cmp(&a.2) // pinned first (true > false)
|
||||||
.then(a.0.cmp(&b.0))
|
.then(a.0.cmp(&b.0))
|
||||||
});
|
});
|
||||||
let pick = candidates.into_iter().next();
|
if let Some((n, e, _)) = candidates.into_iter().next() {
|
||||||
pick.map(|(n, e, _)| (n, e))
|
return Ok((n, e));
|
||||||
.ok_or_else(|| RouteError::NoFeasibleNeuron {
|
}
|
||||||
|
|
||||||
|
// 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(),
|
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,
|
/// Issue `POST {endpoint}/models/load` for this profile on this neuron,
|
||||||
|
|||||||
105
crates/cortex-gateway/src/served_usage.rs
Normal file
105
crates/cortex-gateway/src/served_usage.rs
Normal file
@@ -0,0 +1,105 @@
|
|||||||
|
//! Served-usage ledger (#58): cortex meters, per principal and per UTC day,
|
||||||
|
//! the tokens it has served on behalf of mesh accounts, and periodically
|
||||||
|
//! reports **absolute** cumulative counters to helexa-upstream for
|
||||||
|
//! reconciliation (operators are compensated for served tokens).
|
||||||
|
//!
|
||||||
|
//! Counters are cumulative-since-process-start for the current period;
|
||||||
|
//! upstream upserts them monotonically (GREATEST), so re-sending the same
|
||||||
|
//! value is idempotent and a flush that races another is harmless. (A
|
||||||
|
//! process restart resets the in-memory counter; the monotonic upsert keeps
|
||||||
|
//! upstream from regressing — at most it under-counts the restarted window,
|
||||||
|
//! acceptable for beta. One cortex per operator token is assumed.)
|
||||||
|
|
||||||
|
use serde::Serialize;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||||
|
pub struct ServedRow {
|
||||||
|
pub account_id: String,
|
||||||
|
pub key_id: String,
|
||||||
|
pub period: String, // YYYY-MM-DD (UTC)
|
||||||
|
pub served_tokens: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
pub struct ServedUsage {
|
||||||
|
inner: Mutex<HashMap<(String, String, String), u64>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServedUsage {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add served tokens for a principal in today's (UTC) period.
|
||||||
|
pub fn add(&self, account_id: &str, key_id: &str, tokens: u64) {
|
||||||
|
if tokens == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let period = chrono::Utc::now().format("%Y-%m-%d").to_string();
|
||||||
|
let mut m = self.inner.lock().expect("served-usage lock");
|
||||||
|
*m.entry((account_id.to_string(), key_id.to_string(), period))
|
||||||
|
.or_insert(0) += tokens;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Absolute cumulative counters, for a flush to upstream.
|
||||||
|
pub fn snapshot(&self) -> Vec<ServedRow> {
|
||||||
|
let m = self.inner.lock().expect("served-usage lock");
|
||||||
|
m.iter()
|
||||||
|
.map(|((account_id, key_id, period), &served_tokens)| ServedRow {
|
||||||
|
account_id: account_id.clone(),
|
||||||
|
key_id: key_id.clone(),
|
||||||
|
period: period.clone(),
|
||||||
|
served_tokens,
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST the absolute counters to upstream's `/authz/v1/served-usage`.
|
||||||
|
pub async fn report(
|
||||||
|
client: &reqwest::Client,
|
||||||
|
base_url: &str,
|
||||||
|
bearer: &str,
|
||||||
|
rows: &[ServedRow],
|
||||||
|
) -> Result<(), reqwest::Error> {
|
||||||
|
if rows.is_empty() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let url = format!("{}/authz/v1/served-usage", base_url.trim_end_matches('/'));
|
||||||
|
client
|
||||||
|
.post(url)
|
||||||
|
.bearer_auth(bearer)
|
||||||
|
.json(&serde_json::json!({ "rows": rows }))
|
||||||
|
.send()
|
||||||
|
.await?
|
||||||
|
.error_for_status()?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn accumulates_per_principal_and_period() {
|
||||||
|
let su = ServedUsage::new();
|
||||||
|
su.add("acct", "key", 10);
|
||||||
|
su.add("acct", "key", 5);
|
||||||
|
su.add("acct", "other", 7);
|
||||||
|
su.add("acct", "key", 0); // no-op
|
||||||
|
let mut rows = su.snapshot();
|
||||||
|
rows.sort_by(|a, b| a.key_id.cmp(&b.key_id));
|
||||||
|
assert_eq!(rows.len(), 2);
|
||||||
|
let key_row = rows.iter().find(|r| r.key_id == "key").unwrap();
|
||||||
|
assert_eq!(key_row.served_tokens, 15);
|
||||||
|
assert_eq!(
|
||||||
|
rows.iter()
|
||||||
|
.find(|r| r.key_id == "other")
|
||||||
|
.unwrap()
|
||||||
|
.served_tokens,
|
||||||
|
7
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,6 @@
|
|||||||
|
use crate::entitlements_chain::ChainedEntitlementProvider;
|
||||||
use crate::entitlements_local::LocalEntitlementProvider;
|
use crate::entitlements_local::LocalEntitlementProvider;
|
||||||
|
use crate::entitlements_upstream::UpstreamEntitlementProvider;
|
||||||
use cortex_core::catalogue::ModelCatalogue;
|
use cortex_core::catalogue::ModelCatalogue;
|
||||||
use cortex_core::config::{EvictionSettings, GatewayConfig, NeuronEndpoint};
|
use cortex_core::config::{EvictionSettings, GatewayConfig, NeuronEndpoint};
|
||||||
use cortex_core::entitlements::EntitlementProvider;
|
use cortex_core::entitlements::EntitlementProvider;
|
||||||
@@ -20,6 +22,9 @@ pub struct CortexState {
|
|||||||
/// Whether to reject unauthenticated requests (#49). Read by the auth
|
/// Whether to reject unauthenticated requests (#49). Read by the auth
|
||||||
/// middleware once it lands.
|
/// middleware once it lands.
|
||||||
pub require_auth: bool,
|
pub require_auth: bool,
|
||||||
|
/// Per-principal served-token tally (#58), reported to upstream for
|
||||||
|
/// operator reconciliation by the flush task when upstream is enabled.
|
||||||
|
pub served_usage: Arc<crate::served_usage::ServedUsage>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl CortexState {
|
impl CortexState {
|
||||||
@@ -37,14 +42,28 @@ impl CortexState {
|
|||||||
last_poll: None,
|
last_poll: None,
|
||||||
discovery: None,
|
discovery: None,
|
||||||
activation: None,
|
activation: None,
|
||||||
|
model_load: HashMap::new(),
|
||||||
|
consecutive_poll_failures: 0,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let catalogue = ModelCatalogue::load(&config.models_config);
|
let catalogue = ModelCatalogue::load(&config.models_config);
|
||||||
|
|
||||||
let entitlements: Arc<dyn EntitlementProvider> =
|
// Local provider always handles operator + infra keys. When the
|
||||||
Arc::new(LocalEntitlementProvider::from_config(&config.entitlements));
|
// upstream client is enabled (#57), wrap it in the chain so locally
|
||||||
|
// unknown keys fall through to the mesh authority; otherwise stay
|
||||||
|
// purely local.
|
||||||
|
let local = LocalEntitlementProvider::from_config(&config.entitlements);
|
||||||
|
let entitlements: Arc<dyn EntitlementProvider> = if config.upstream.enabled {
|
||||||
|
tracing::info!(url = %config.upstream.url, "upstream entitlement client enabled");
|
||||||
|
Arc::new(ChainedEntitlementProvider::new(
|
||||||
|
local,
|
||||||
|
UpstreamEntitlementProvider::new(&config.upstream),
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
Arc::new(local)
|
||||||
|
};
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
nodes: RwLock::new(nodes),
|
nodes: RwLock::new(nodes),
|
||||||
@@ -57,6 +76,7 @@ impl CortexState {
|
|||||||
.expect("failed to build HTTP client"),
|
.expect("failed to build HTTP client"),
|
||||||
entitlements,
|
entitlements,
|
||||||
require_auth: config.entitlements.require_auth,
|
require_auth: config.entitlements.require_auth,
|
||||||
|
served_usage: Arc::new(crate::served_usage::ServedUsage::new()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -57,6 +57,7 @@ async fn test_alias_resolves_in_chat_completions() {
|
|||||||
}],
|
}],
|
||||||
models_config: models_path.to_string_lossy().to_string(),
|
models_config: models_path.to_string_lossy().to_string(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -143,6 +144,7 @@ async fn test_aliases_surface_in_v1_models() {
|
|||||||
}],
|
}],
|
||||||
models_config: models_path.to_string_lossy().to_string(),
|
models_config: models_path.to_string_lossy().to_string(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -232,6 +234,7 @@ async fn test_alias_falls_through_for_unmapped_model() {
|
|||||||
}],
|
}],
|
||||||
models_config: models_path.to_string_lossy().to_string(),
|
models_config: models_path.to_string_lossy().to_string(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
|||||||
273
crates/cortex-gateway/tests/auth.rs
Normal file
273
crates/cortex-gateway/tests/auth.rs
Normal file
@@ -0,0 +1,273 @@
|
|||||||
|
//! Integration tests for API-key auth + principal resolution (#49).
|
||||||
|
//!
|
||||||
|
//! Verifies the #63 rejection contract (401 invalid_api_key via the #60
|
||||||
|
//! envelope) and that an authenticated request reaches neuron carrying the
|
||||||
|
//! internal principal headers — while a client-supplied principal header is
|
||||||
|
//! stripped (anti-spoofing).
|
||||||
|
|
||||||
|
use axum::Json;
|
||||||
|
use axum::extract::Path;
|
||||||
|
use axum::http::HeaderMap;
|
||||||
|
use axum::routing::{get, post};
|
||||||
|
use cortex_core::config::{
|
||||||
|
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
|
||||||
|
GatewaySettings, NeuronEndpoint,
|
||||||
|
};
|
||||||
|
use cortex_core::entitlements::{CapWindow, HEADER_ACCOUNT_ID, HEADER_KEY_ID};
|
||||||
|
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||||
|
use cortex_gateway::state::CortexState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// What the mock neuron observed on the inbound `/v1/chat/completions`
|
||||||
|
/// request: the principal headers cortex stamped (or didn't).
|
||||||
|
#[derive(Default)]
|
||||||
|
struct Seen {
|
||||||
|
account_id: Option<String>,
|
||||||
|
key_id: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a mock neuron that records the principal headers it receives and
|
||||||
|
/// returns a trivial chat completion. Returns (base_url, observed).
|
||||||
|
async fn spawn_capturing_neuron() -> (String, Arc<Mutex<Seen>>) {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let base_url = format!("http://{addr}");
|
||||||
|
let inference_url = base_url.clone();
|
||||||
|
let seen: Arc<Mutex<Seen>> = Arc::new(Mutex::new(Seen::default()));
|
||||||
|
let sink = Arc::clone(&seen);
|
||||||
|
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route(
|
||||||
|
"/models/{model_id}/endpoint",
|
||||||
|
get(move |Path(_): Path<String>| {
|
||||||
|
let url = inference_url.clone();
|
||||||
|
async move { Json(json!({ "url": url })) }
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/v1/chat/completions",
|
||||||
|
post(move |headers: HeaderMap, Json(body): Json<Value>| {
|
||||||
|
let sink = Arc::clone(&sink);
|
||||||
|
async move {
|
||||||
|
{
|
||||||
|
let mut s = sink.lock().unwrap();
|
||||||
|
s.account_id = headers
|
||||||
|
.get(HEADER_ACCOUNT_ID)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.map(str::to_string);
|
||||||
|
s.key_id = headers
|
||||||
|
.get(HEADER_KEY_ID)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.map(str::to_string);
|
||||||
|
}
|
||||||
|
let model = body.get("model").and_then(Value::as_str).unwrap_or("m");
|
||||||
|
Json(json!({
|
||||||
|
"id": "chatcmpl-auth-001",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"created": 1700000000_u64,
|
||||||
|
"model": model,
|
||||||
|
"choices": [{
|
||||||
|
"index": 0,
|
||||||
|
"message": {"role": "assistant", "content": "ok"},
|
||||||
|
"finish_reason": "stop"
|
||||||
|
}],
|
||||||
|
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.with_state(());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
(base_url, seen)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a gateway with the given entitlements config, a single neuron, and
|
||||||
|
/// `test-model` seeded as loaded (build_app spawns no poller).
|
||||||
|
async fn spawn_gateway(neuron_url: &str, entitlements: EntitlementsConfig) -> String {
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![NeuronEndpoint {
|
||||||
|
name: "mock-node".into(),
|
||||||
|
endpoint: neuron_url.to_string(),
|
||||||
|
}],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements,
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
{
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let node = nodes.get_mut("mock-node").unwrap();
|
||||||
|
node.healthy = true;
|
||||||
|
node.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn one_key_config(require_auth: bool) -> EntitlementsConfig {
|
||||||
|
EntitlementsConfig {
|
||||||
|
require_auth,
|
||||||
|
keys: vec![ApiKeyConfig {
|
||||||
|
key: "sk-good".into(),
|
||||||
|
account_id: "acct-1".into(),
|
||||||
|
key_id: Some("key-1".into()),
|
||||||
|
hard_cap: None,
|
||||||
|
window: CapWindow::Balance,
|
||||||
|
}],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chat_body() -> Value {
|
||||||
|
json!({
|
||||||
|
"model": "test-model",
|
||||||
|
"messages": [{"role": "user", "content": "hi"}]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_key_when_required_is_401_invalid_api_key() {
|
||||||
|
let (neuron, _seen) = spawn_capturing_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.json(&chat_body())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::UNAUTHORIZED);
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "invalid_api_key");
|
||||||
|
assert_eq!(body["error"]["type"], "invalid_request_error");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unrecognized_key_is_ignored_when_auth_not_required() {
|
||||||
|
let (neuron, seen) = spawn_capturing_neuron().await;
|
||||||
|
// allow-anonymous mode: a placeholder/unknown bearer (as opencode,
|
||||||
|
// Open WebUI, Agent Zero, litellm all send by default) must NOT be
|
||||||
|
// rejected — it's ignored and the request is served anonymously.
|
||||||
|
let gateway = spawn_gateway(&neuron, one_key_config(false)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-dummy-placeholder")
|
||||||
|
.json(&chat_body())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
// Served, but anonymous — no principal stamped from the bogus key.
|
||||||
|
assert!(seen.lock().unwrap().account_id.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn invalid_key_is_401_when_auth_required() {
|
||||||
|
let (neuron, seen) = spawn_capturing_neuron().await;
|
||||||
|
// With auth required, a present-but-wrong credential is rejected.
|
||||||
|
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-wrong")
|
||||||
|
.json(&chat_body())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::UNAUTHORIZED);
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "invalid_api_key");
|
||||||
|
// Rejected before dispatch — neuron never saw the request.
|
||||||
|
assert!(seen.lock().unwrap().account_id.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn valid_key_reaches_neuron_with_principal_headers() {
|
||||||
|
let (neuron, seen) = spawn_capturing_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-good")
|
||||||
|
// A spoofed principal header must be stripped, not forwarded.
|
||||||
|
.header(HEADER_ACCOUNT_ID, "attacker")
|
||||||
|
.json(&chat_body())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
let s = seen.lock().unwrap();
|
||||||
|
assert_eq!(s.account_id.as_deref(), Some("acct-1"));
|
||||||
|
assert_eq!(s.key_id.as_deref(), Some("key-1"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn anonymous_allowed_when_auth_not_required() {
|
||||||
|
let (neuron, seen) = spawn_capturing_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, EntitlementsConfig::default()).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.json(&chat_body())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
// No principal resolved → no principal headers stamped.
|
||||||
|
let s = seen.lock().unwrap();
|
||||||
|
assert!(s.account_id.is_none());
|
||||||
|
assert!(s.key_id.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn health_is_public_even_when_auth_required() {
|
||||||
|
let (neuron, _seen) = spawn_capturing_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, one_key_config(true)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.get(format!("{gateway}/health"))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
}
|
||||||
254
crates/cortex-gateway/tests/budget_enforcement.rs
Normal file
254
crates/cortex-gateway/tests/budget_enforcement.rs
Normal file
@@ -0,0 +1,254 @@
|
|||||||
|
//! Integration tests for budget enforcement (#52) — the A0 seatbelt.
|
||||||
|
//!
|
||||||
|
//! A reservation over the key's hard cap is refused *before* neuron is hit,
|
||||||
|
//! with the #63 code matching the cap-window semantics (rate_limit_exceeded
|
||||||
|
//! + Retry-After for a resetting window, insufficient_quota for a hard
|
||||||
|
//! balance). Spend never exceeds the cap. No 402, ever.
|
||||||
|
|
||||||
|
use axum::Json;
|
||||||
|
use axum::extract::Path;
|
||||||
|
use axum::routing::{get, post};
|
||||||
|
use cortex_core::config::{
|
||||||
|
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
|
||||||
|
GatewaySettings, NeuronEndpoint,
|
||||||
|
};
|
||||||
|
use cortex_core::entitlements::{CapWindow, Principal};
|
||||||
|
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||||
|
use cortex_gateway::state::CortexState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Mock neuron with a hit counter on the inference path, so a test can prove
|
||||||
|
/// a request was (or wasn't) dispatched.
|
||||||
|
async fn spawn_counting_neuron() -> (String, Arc<AtomicU64>) {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let base_url = format!("http://{addr}");
|
||||||
|
let inference_url = base_url.clone();
|
||||||
|
let hits = Arc::new(AtomicU64::new(0));
|
||||||
|
let sink = Arc::clone(&hits);
|
||||||
|
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route(
|
||||||
|
"/models/{model_id}/endpoint",
|
||||||
|
get(move |Path(_): Path<String>| {
|
||||||
|
let url = inference_url.clone();
|
||||||
|
async move { Json(json!({ "url": url })) }
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/v1/chat/completions",
|
||||||
|
post(move |Json(body): Json<Value>| {
|
||||||
|
let sink = Arc::clone(&sink);
|
||||||
|
async move {
|
||||||
|
sink.fetch_add(1, Ordering::SeqCst);
|
||||||
|
let model = body.get("model").and_then(Value::as_str).unwrap_or("m");
|
||||||
|
Json(json!({
|
||||||
|
"id": "chatcmpl-budget",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"created": 1700000000_u64,
|
||||||
|
"model": model,
|
||||||
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||||
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(base_url, hits)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_gateway(neuron_url: &str, key: ApiKeyConfig) -> (Arc<CortexState>, String) {
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![NeuronEndpoint {
|
||||||
|
name: "mock-node".into(),
|
||||||
|
endpoint: neuron_url.to_string(),
|
||||||
|
}],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements: EntitlementsConfig {
|
||||||
|
require_auth: true,
|
||||||
|
keys: vec![key],
|
||||||
|
},
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
{
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let node = nodes.get_mut("mock-node").unwrap();
|
||||||
|
node.healthy = true;
|
||||||
|
node.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(fleet, format!("http://{addr}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn key(window: CapWindow, hard_cap: u64) -> ApiKeyConfig {
|
||||||
|
ApiKeyConfig {
|
||||||
|
key: "sk-cap".into(),
|
||||||
|
account_id: "acct-cap".into(),
|
||||||
|
key_id: Some("key-cap".into()),
|
||||||
|
hard_cap: Some(hard_cap),
|
||||||
|
window,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chat(max_tokens: u64) -> Value {
|
||||||
|
json!({
|
||||||
|
"model": "test-model",
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
"messages": [{"role": "user", "content": "hi"}]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn balance_over_cap_is_429_insufficient_quota_before_dispatch() {
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
// Cap far below a single request's reservation (max_tokens 1000).
|
||||||
|
let (_fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 10)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-cap")
|
||||||
|
.json(&chat(1000))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
// Hard balance → no Retry-After.
|
||||||
|
assert!(resp.headers().get(reqwest::header::RETRY_AFTER).is_none());
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "insufficient_quota");
|
||||||
|
// Refused before dispatch — neuron never saw it.
|
||||||
|
assert_eq!(hits.load(Ordering::SeqCst), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rolling_over_cap_is_429_rate_limited_with_retry_after() {
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
let (_fleet, gateway) =
|
||||||
|
spawn_gateway(&neuron, key(CapWindow::Rolling { seconds: 3600 }, 10)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-cap")
|
||||||
|
.json(&chat(1000))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
let retry = resp
|
||||||
|
.headers()
|
||||||
|
.get(reqwest::header::RETRY_AFTER)
|
||||||
|
.expect("rolling-window rejection must carry Retry-After");
|
||||||
|
assert!(retry.to_str().unwrap().parse::<u64>().unwrap() >= 1);
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||||
|
assert_eq!(hits.load(Ordering::SeqCst), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn within_cap_is_served() {
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
let (_fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 1_000_000)).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-cap")
|
||||||
|
.json(&chat(50))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
assert_eq!(hits.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a0_seatbelt_caps_a_runaway_fan_out() {
|
||||||
|
// An Agent-Zero-style key with a modest cap: a burst of requests drains
|
||||||
|
// it, then further requests are refused — the account stops draining and
|
||||||
|
// spend never exceeds the cap.
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
let (fleet, gateway) = spawn_gateway(&neuron, key(CapWindow::Balance, 100)).await;
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
|
||||||
|
let mut ok = 0;
|
||||||
|
let mut refused = 0;
|
||||||
|
for _ in 0..20 {
|
||||||
|
let resp = client
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth("sk-cap")
|
||||||
|
.json(&chat(20))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
match resp.status() {
|
||||||
|
reqwest::StatusCode::OK => {
|
||||||
|
ok += 1;
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
}
|
||||||
|
reqwest::StatusCode::TOO_MANY_REQUESTS => {
|
||||||
|
refused += 1;
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "insufficient_quota");
|
||||||
|
}
|
||||||
|
other => panic!("unexpected status {other}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
assert!(ok >= 1, "some requests should be served");
|
||||||
|
assert!(refused >= 1, "the cap must eventually refuse the fan-out");
|
||||||
|
assert_eq!(
|
||||||
|
hits.load(Ordering::SeqCst),
|
||||||
|
ok,
|
||||||
|
"refused requests never dispatched"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Spend never exceeded the hard cap (reservation prevents overshoot).
|
||||||
|
// Poll briefly for in-flight settles to land.
|
||||||
|
let principal = Principal {
|
||||||
|
account_id: "acct-cap".into(),
|
||||||
|
key_id: "key-cap".into(),
|
||||||
|
};
|
||||||
|
for _ in 0..50 {
|
||||||
|
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
|
||||||
|
if snap.reserved == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
||||||
|
}
|
||||||
|
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
|
||||||
|
assert!(snap.spent <= 100, "spent {} exceeded cap", snap.spent);
|
||||||
|
}
|
||||||
@@ -430,6 +430,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc<CortexState>, Stri
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
|||||||
@@ -89,6 +89,7 @@ async fn error_response_no_healthy_nodes() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
let fleet = Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
||||||
|
|||||||
@@ -72,6 +72,7 @@ fn make_fleet(endpoint: &str, defrag_after: u32) -> Arc<CortexState> {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
Arc::new(CortexState::from_config(&config))
|
Arc::new(CortexState::from_config(&config))
|
||||||
}
|
}
|
||||||
|
|||||||
125
crates/cortex-gateway/tests/feasibility_routing.rs
Normal file
125
crates/cortex-gateway/tests/feasibility_routing.rs
Normal file
@@ -0,0 +1,125 @@
|
|||||||
|
//! 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(),
|
||||||
|
upstream: 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);
|
||||||
|
}
|
||||||
190
crates/cortex-gateway/tests/load_routing.rs
Normal file
190
crates/cortex-gateway/tests/load_routing.rs
Normal file
@@ -0,0 +1,190 @@
|
|||||||
|
//! Load-aware routing across replicas (#55).
|
||||||
|
//!
|
||||||
|
//! When a model is loaded on more than one healthy neuron, the router picks
|
||||||
|
//! the least-busy replica using the per-model admission load each neuron
|
||||||
|
//! reports on `GET /health` (#53), rather than always taking the first.
|
||||||
|
|
||||||
|
mod common;
|
||||||
|
|
||||||
|
use axum::Json;
|
||||||
|
use axum::extract::Path;
|
||||||
|
use axum::http::{StatusCode, header};
|
||||||
|
use axum::response::IntoResponse;
|
||||||
|
use axum::routing::{get, post};
|
||||||
|
use cortex_core::config::{
|
||||||
|
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||||
|
};
|
||||||
|
use cortex_core::discovery::ModelLoad;
|
||||||
|
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||||
|
use cortex_gateway::state::CortexState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Seed a node as healthy with `test-model` loaded and a given admission load.
|
||||||
|
async fn seed_loaded(fleet: &CortexState, node: &str, in_flight: usize, queue_depth: usize) {
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let n = nodes.get_mut(node).expect("node exists");
|
||||||
|
n.healthy = true;
|
||||||
|
n.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
n.model_load.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelLoad {
|
||||||
|
id: "test-model".into(),
|
||||||
|
in_flight,
|
||||||
|
queue_depth,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a gateway state over two mock neurons (no poller; we seed state).
|
||||||
|
async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexState> {
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![
|
||||||
|
NeuronEndpoint {
|
||||||
|
name: "node-a".into(),
|
||||||
|
endpoint: endpoint_a.to_string(),
|
||||||
|
},
|
||||||
|
NeuronEndpoint {
|
||||||
|
name: "node-b".into(),
|
||||||
|
endpoint: endpoint_b.to_string(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
Arc::new(CortexState::from_config(&config))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn routes_to_least_busy_replica() {
|
||||||
|
let neuron_a = common::spawn_mock_neuron().await;
|
||||||
|
let neuron_b = common::spawn_mock_neuron().await;
|
||||||
|
let fleet = two_neuron_fleet(&neuron_a, &neuron_b).await;
|
||||||
|
|
||||||
|
// A is busy (1 running + 3 queued), B is idle.
|
||||||
|
seed_loaded(&fleet, "node-a", 1, 3).await;
|
||||||
|
seed_loaded(&fleet, "node-b", 0, 0).await;
|
||||||
|
|
||||||
|
let route = cortex_gateway::router::resolve(&fleet, "test-model")
|
||||||
|
.await
|
||||||
|
.expect("model is loaded on both nodes");
|
||||||
|
assert_eq!(route.node_name, "node-b", "should pick the idle replica");
|
||||||
|
|
||||||
|
// Flip the load: now B is the busy one.
|
||||||
|
seed_loaded(&fleet, "node-a", 0, 0).await;
|
||||||
|
seed_loaded(&fleet, "node-b", 1, 5).await;
|
||||||
|
let route = cortex_gateway::router::resolve(&fleet, "test-model")
|
||||||
|
.await
|
||||||
|
.expect("still loaded");
|
||||||
|
assert_eq!(route.node_name, "node-a", "should follow the lighter load");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mock neuron whose inference endpoint always returns a #63 backpressure
|
||||||
|
/// envelope (503 + Retry-After) — simulating a saturated neuron.
|
||||||
|
async fn spawn_busy_neuron() -> String {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let base_url = format!("http://{addr}");
|
||||||
|
let inference_url = base_url.clone();
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route(
|
||||||
|
"/models/{model_id}/endpoint",
|
||||||
|
get(move |Path(_): Path<String>| {
|
||||||
|
let url = inference_url.clone();
|
||||||
|
async move { Json(json!({ "url": url })) }
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/v1/chat/completions",
|
||||||
|
post(|| async {
|
||||||
|
let body = json!({"error": {
|
||||||
|
"message": "model is busy (admission queue full); retry shortly",
|
||||||
|
"type": "rate_limit_error",
|
||||||
|
"code": "rate_limit_exceeded",
|
||||||
|
"param": null
|
||||||
|
}});
|
||||||
|
(
|
||||||
|
StatusCode::SERVICE_UNAVAILABLE,
|
||||||
|
[(header::RETRY_AFTER, "6")],
|
||||||
|
Json(body),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
base_url
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn neuron_backpressure_is_propagated_intact() {
|
||||||
|
// A saturated neuron's 503 + Retry-After + envelope must reach the client
|
||||||
|
// verbatim — not unwrapped, remapped, or stripped (#55 / #63).
|
||||||
|
let neuron = spawn_busy_neuron().await;
|
||||||
|
let fleet = two_neuron_fleet(&neuron, &neuron).await;
|
||||||
|
seed_loaded(&fleet, "node-a", 1, 8).await;
|
||||||
|
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("http://{addr}/v1/chat/completions"))
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
assert_eq!(
|
||||||
|
resp.headers()
|
||||||
|
.get(reqwest::header::RETRY_AFTER)
|
||||||
|
.and_then(|v| v.to_str().ok()),
|
||||||
|
Some("6"),
|
||||||
|
"Retry-After must survive the proxy"
|
||||||
|
);
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ties_break_deterministically_by_name() {
|
||||||
|
let neuron_a = common::spawn_mock_neuron().await;
|
||||||
|
let neuron_b = common::spawn_mock_neuron().await;
|
||||||
|
let fleet = two_neuron_fleet(&neuron_a, &neuron_b).await;
|
||||||
|
|
||||||
|
// Equal load on both → stable pick (lowest node name).
|
||||||
|
seed_loaded(&fleet, "node-a", 0, 0).await;
|
||||||
|
seed_loaded(&fleet, "node-b", 0, 0).await;
|
||||||
|
|
||||||
|
let route = cortex_gateway::router::resolve(&fleet, "test-model")
|
||||||
|
.await
|
||||||
|
.expect("loaded");
|
||||||
|
assert_eq!(route.node_name, "node-a", "ties break by name");
|
||||||
|
}
|
||||||
209
crates/cortex-gateway/tests/metering.rs
Normal file
209
crates/cortex-gateway/tests/metering.rs
Normal file
@@ -0,0 +1,209 @@
|
|||||||
|
//! Integration tests for per-request token metering (#51).
|
||||||
|
//!
|
||||||
|
//! Drives authenticated requests through the gateway to a mock neuron that
|
||||||
|
//! reports a fixed `usage` object, then asserts the EntitlementProvider's
|
||||||
|
//! spend ledger reflects cumulative per-key spend and that reservations
|
||||||
|
//! settle to actual (no outstanding reserved tokens once requests complete).
|
||||||
|
|
||||||
|
mod common;
|
||||||
|
|
||||||
|
use cortex_core::config::{
|
||||||
|
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
|
||||||
|
GatewaySettings, NeuronEndpoint,
|
||||||
|
};
|
||||||
|
use cortex_core::entitlements::{CapWindow, Principal};
|
||||||
|
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||||
|
use cortex_gateway::state::CortexState;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
const ACCOUNT: &str = "acct-meter";
|
||||||
|
const KEY_ID: &str = "key-meter";
|
||||||
|
const BEARER: &str = "sk-meter";
|
||||||
|
|
||||||
|
/// The mock neuron (common::spawn_mock_neuron) reports this fixed usage on
|
||||||
|
/// every chat completion.
|
||||||
|
const PROMPT_PER_REQ: u64 = 10;
|
||||||
|
const COMPLETION_PER_REQ: u64 = 5;
|
||||||
|
|
||||||
|
async fn spawn_metered_gateway(neuron_url: &str) -> (Arc<CortexState>, String) {
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![NeuronEndpoint {
|
||||||
|
name: "mock-node".into(),
|
||||||
|
endpoint: neuron_url.to_string(),
|
||||||
|
}],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements: EntitlementsConfig {
|
||||||
|
require_auth: true,
|
||||||
|
keys: vec![ApiKeyConfig {
|
||||||
|
key: BEARER.into(),
|
||||||
|
account_id: ACCOUNT.into(),
|
||||||
|
key_id: Some(KEY_ID.into()),
|
||||||
|
hard_cap: Some(1_000_000),
|
||||||
|
window: CapWindow::Balance,
|
||||||
|
}],
|
||||||
|
},
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
{
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let node = nodes.get_mut("mock-node").unwrap();
|
||||||
|
node.healthy = true;
|
||||||
|
node.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(fleet, format!("http://{addr}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn principal() -> Principal {
|
||||||
|
Principal {
|
||||||
|
account_id: ACCOUNT.into(),
|
||||||
|
key_id: KEY_ID.into(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Poll the provider ledger until settled spend reaches `expected` (settle
|
||||||
|
/// runs in a spawned task after the response stream finishes) or time out.
|
||||||
|
async fn await_spent(fleet: &CortexState, expected: u64) -> u64 {
|
||||||
|
let principal = principal();
|
||||||
|
for _ in 0..100 {
|
||||||
|
let snap = fleet.entitlements.snapshot(&principal).await.unwrap();
|
||||||
|
if snap.spent >= expected {
|
||||||
|
return snap.spent;
|
||||||
|
}
|
||||||
|
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||||
|
}
|
||||||
|
fleet.entitlements.snapshot(&principal).await.unwrap().spent
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cumulative_spend_is_metered_per_key() {
|
||||||
|
let neuron = common::spawn_mock_neuron().await;
|
||||||
|
let (fleet, gateway) = spawn_metered_gateway(&neuron).await;
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
|
||||||
|
const N: u64 = 3;
|
||||||
|
for _ in 0..N {
|
||||||
|
let resp = client
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.bearer_auth(BEARER)
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
// Drain the body so the response stream finishes and metering settles.
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
let expected = N * (PROMPT_PER_REQ + COMPLETION_PER_REQ);
|
||||||
|
let spent = await_spent(&fleet, expected).await;
|
||||||
|
assert_eq!(
|
||||||
|
spent, expected,
|
||||||
|
"ledger must reflect cumulative per-key spend"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Reservations settled to actual — nothing left outstanding.
|
||||||
|
let snap = fleet.entitlements.snapshot(&principal()).await.unwrap();
|
||||||
|
assert_eq!(snap.reserved, 0, "all reservations must settle/release");
|
||||||
|
assert_eq!(snap.hard_cap, Some(1_000_000));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn anonymous_request_records_no_spend() {
|
||||||
|
// require_auth=false so the unauthenticated request is served, but with
|
||||||
|
// no principal it must not touch any ledger.
|
||||||
|
let neuron = common::spawn_mock_neuron().await;
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![NeuronEndpoint {
|
||||||
|
name: "mock-node".into(),
|
||||||
|
endpoint: neuron.clone(),
|
||||||
|
}],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements: EntitlementsConfig::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
{
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let node = nodes.get_mut("mock-node").unwrap();
|
||||||
|
node.healthy = true;
|
||||||
|
node.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("http://{addr}/v1/chat/completions"))
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
|
||||||
|
// An unconfigured principal has a zeroed snapshot — nothing was metered.
|
||||||
|
let snap = fleet
|
||||||
|
.entitlements
|
||||||
|
.snapshot(&Principal {
|
||||||
|
account_id: "nobody".into(),
|
||||||
|
key_id: "nobody".into(),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(snap.spent, 0);
|
||||||
|
}
|
||||||
132
crates/cortex-gateway/tests/model_cost.rs
Normal file
132
crates/cortex-gateway/tests/model_cost.rs
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
//! 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(),
|
||||||
|
upstream: 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);
|
||||||
|
}
|
||||||
@@ -8,8 +8,11 @@
|
|||||||
//! - `cost` from the catalogue profile (operator-set pricing).
|
//! - `cost` from the catalogue profile (operator-set pricing).
|
||||||
//! - `tool_call` / `reasoning` from the neuron's runtime detection (OR-ed in)
|
//! - `tool_call` / `reasoning` from the neuron's runtime detection (OR-ed in)
|
||||||
//!
|
//!
|
||||||
//! Also a regression guard for the removal of `max_model_len` — the misnamed,
|
//! Also asserts the flat, vLLM-convention duplicates (`max_model_len`,
|
||||||
//! unconsumed vLLM-ism that this contract replaces.
|
//! `max_input_tokens`, `max_output_tokens`) mirror `limit` (#78): the
|
||||||
|
//! earlier removal of `max_model_len` as "unconsumed" was wrong — Hermes
|
||||||
|
//! Agent (and the wider OpenAI client ecosystem) probes those flat keys
|
||||||
|
//! and cannot see `limit.context`.
|
||||||
|
|
||||||
use cortex_core::config::{
|
use cortex_core::config::{
|
||||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||||
@@ -55,6 +58,7 @@ capabilities = ["text"]
|
|||||||
}],
|
}],
|
||||||
models_config: cat_path.to_string_lossy().into_owned(),
|
models_config: cat_path.to_string_lossy().into_owned(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -84,6 +88,21 @@ capabilities = ["text"]
|
|||||||
}),
|
}),
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
// A model with no derivable limit: the flat #78 fields must be
|
||||||
|
// OMITTED (absent-vs-zero is load-bearing), never 0 or a guess.
|
||||||
|
node.models.insert(
|
||||||
|
"no-limit-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "no-limit-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: None,
|
||||||
|
capabilities: vec!["text".into()],
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: None,
|
||||||
|
},
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
@@ -122,11 +141,26 @@ capabilities = ["text"]
|
|||||||
assert_eq!(entry["tool_call"], true);
|
assert_eq!(entry["tool_call"], true);
|
||||||
assert_eq!(entry["reasoning"], true);
|
assert_eq!(entry["reasoning"], true);
|
||||||
|
|
||||||
// Regression guard: the removed, unconsumed vLLM-ism must not reappear.
|
// Flat ecosystem duplicates (#78) mirror the advertised `limit` so
|
||||||
assert!(
|
// vLLM-convention probes (Hermes Agent) auto-detect the window.
|
||||||
entry.get("max_model_len").is_none(),
|
assert_eq!(entry["max_model_len"], 49152);
|
||||||
"max_model_len was removed; /v1/models must not advertise it"
|
assert_eq!(entry["max_input_tokens"], 40960);
|
||||||
);
|
assert_eq!(entry["max_output_tokens"], 8192);
|
||||||
|
|
||||||
|
// No limit → flat fields omitted entirely, never 0 or a guess.
|
||||||
|
let unknown = body["data"]
|
||||||
|
.as_array()
|
||||||
|
.unwrap()
|
||||||
|
.iter()
|
||||||
|
.find(|m| m["id"] == "no-limit-model")
|
||||||
|
.expect("no-limit-model present in /v1/models");
|
||||||
|
assert!(unknown.get("limit").is_none());
|
||||||
|
for key in ["max_model_len", "max_input_tokens", "max_output_tokens"] {
|
||||||
|
assert!(
|
||||||
|
unknown.get(key).is_none(),
|
||||||
|
"{key} must be omitted when the window is unknown"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let _ = std::fs::remove_file(&cat_path);
|
let _ = std::fs::remove_file(&cat_path);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ async fn test_poller_discovers_models() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -84,6 +85,7 @@ async fn test_poller_updates_gateway_models_endpoint() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -156,6 +158,7 @@ async fn test_models_endpoint_unions_capabilities_across_nodes() {
|
|||||||
],
|
],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -219,6 +222,7 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -228,10 +232,26 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
|
|||||||
nodes.get_mut("dead-node").unwrap().healthy = true;
|
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;
|
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;
|
// It flips unhealthy only after POLL_FAILURE_THRESHOLD (3) consecutive
|
||||||
assert!(!nodes.get("dead-node").unwrap().healthy);
|
// 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]
|
#[tokio::test]
|
||||||
@@ -257,6 +277,7 @@ async fn test_poller_removes_stale_models() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -288,6 +309,7 @@ async fn test_poller_removes_stale_models() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet2 = Arc::new(CortexState::from_config(&config2));
|
let fleet2 = Arc::new(CortexState::from_config(&config2));
|
||||||
@@ -370,6 +392,7 @@ async fn test_poller_captures_activation_from_health() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
@@ -415,6 +438,7 @@ async fn test_poller_parses_recovering_status() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let fleet = Arc::new(CortexState::from_config(&config));
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
|||||||
175
crates/cortex-gateway/tests/prompt_prevalidation.rs
Normal file
175
crates/cortex-gateway/tests/prompt_prevalidation.rs
Normal file
@@ -0,0 +1,175 @@
|
|||||||
|
//! Fail-fast prompt pre-validation + advisory client hints (#56).
|
||||||
|
//!
|
||||||
|
//! cortex refuses a prompt that already exceeds the model's advertised
|
||||||
|
//! context window before dispatching to neuron — the same #60
|
||||||
|
//! `context_length_exceeded` envelope neuron would emit, just earlier — and
|
||||||
|
//! attaches an advisory `X-Helexa-Advice` header for fingerprinted clients.
|
||||||
|
|
||||||
|
use axum::Json;
|
||||||
|
use axum::extract::Path;
|
||||||
|
use axum::routing::{get, post};
|
||||||
|
use cortex_core::config::{
|
||||||
|
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||||
|
};
|
||||||
|
use cortex_core::harness::ModelLimit;
|
||||||
|
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||||
|
use cortex_gateway::state::CortexState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Mock neuron with a hit counter, so a test can prove a request was (or
|
||||||
|
/// wasn't) dispatched past the gateway's pre-validation.
|
||||||
|
async fn spawn_counting_neuron() -> (String, Arc<AtomicU64>) {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
let base_url = format!("http://{addr}");
|
||||||
|
let inference_url = base_url.clone();
|
||||||
|
let hits = Arc::new(AtomicU64::new(0));
|
||||||
|
let sink = Arc::clone(&hits);
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route(
|
||||||
|
"/models/{model_id}/endpoint",
|
||||||
|
get(move |Path(_): Path<String>| {
|
||||||
|
let url = inference_url.clone();
|
||||||
|
async move { Json(json!({ "url": url })) }
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/v1/chat/completions",
|
||||||
|
post(move || {
|
||||||
|
let sink = Arc::clone(&sink);
|
||||||
|
async move {
|
||||||
|
sink.fetch_add(1, Ordering::SeqCst);
|
||||||
|
Json(json!({
|
||||||
|
"id": "c", "object": "chat.completion", "created": 1_700_000_000_u64,
|
||||||
|
"model": "test-model",
|
||||||
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||||
|
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(base_url, hits)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Gateway over one neuron with `test-model` loaded and a tiny advertised
|
||||||
|
/// context window (so a modest prompt overflows it).
|
||||||
|
async fn spawn_gateway(neuron: &str, context: usize) -> String {
|
||||||
|
let config = GatewayConfig {
|
||||||
|
gateway: GatewaySettings {
|
||||||
|
listen: "127.0.0.1:0".into(),
|
||||||
|
metrics_listen: "127.0.0.1:0".into(),
|
||||||
|
},
|
||||||
|
eviction: EvictionSettings {
|
||||||
|
strategy: EvictionStrategy::Lru,
|
||||||
|
defrag_after_cycles: 0,
|
||||||
|
},
|
||||||
|
neurons: vec![NeuronEndpoint {
|
||||||
|
name: "mock-node".into(),
|
||||||
|
endpoint: neuron.to_string(),
|
||||||
|
}],
|
||||||
|
models_config: "/dev/null".into(),
|
||||||
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
|
};
|
||||||
|
let fleet = Arc::new(CortexState::from_config(&config));
|
||||||
|
{
|
||||||
|
let mut nodes = fleet.nodes.write().await;
|
||||||
|
let n = nodes.get_mut("mock-node").unwrap();
|
||||||
|
n.healthy = true;
|
||||||
|
n.models.insert(
|
||||||
|
"test-model".into(),
|
||||||
|
ModelEntry {
|
||||||
|
id: "test-model".into(),
|
||||||
|
status: ModelStatus::Loaded,
|
||||||
|
last_accessed: None,
|
||||||
|
vram_estimate_mb: Some(8000),
|
||||||
|
capabilities: Vec::new(),
|
||||||
|
tool_call: false,
|
||||||
|
reasoning: false,
|
||||||
|
limit: Some(ModelLimit {
|
||||||
|
context,
|
||||||
|
input: None,
|
||||||
|
output: 16,
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn over_long_prompt_is_rejected_before_dispatch() {
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, 50).await; // tiny 50-token window
|
||||||
|
|
||||||
|
// ~1200 chars → ~300 est tokens, well over 50.
|
||||||
|
let big = "word ".repeat(240);
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.header("user-agent", "litellm/1.0")
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": big}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST);
|
||||||
|
// Advisory hint for the fingerprinted client (header only, never body).
|
||||||
|
assert!(
|
||||||
|
resp.headers().get("x-helexa-advice").is_some(),
|
||||||
|
"litellm should get advice"
|
||||||
|
);
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "context_length_exceeded");
|
||||||
|
assert_eq!(body["error"]["max"], 50);
|
||||||
|
// Refused at the edge — neuron never saw it.
|
||||||
|
assert_eq!(hits.load(Ordering::SeqCst), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn within_context_passes_through() {
|
||||||
|
let (neuron, hits) = spawn_counting_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, 4096).await;
|
||||||
|
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
let _ = resp.bytes().await.unwrap();
|
||||||
|
assert_eq!(hits.load(Ordering::SeqCst), 1, "served by neuron");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unknown_client_gets_no_advice_header() {
|
||||||
|
let (neuron, _hits) = spawn_counting_neuron().await;
|
||||||
|
let gateway = spawn_gateway(&neuron, 50).await;
|
||||||
|
|
||||||
|
let big = "word ".repeat(240);
|
||||||
|
let resp = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway}/v1/chat/completions"))
|
||||||
|
// no/unknown User-Agent → no advice, but still a clean 400
|
||||||
|
.json(&json!({"model": "test-model", "messages": [{"role": "user", "content": big}]}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST);
|
||||||
|
assert!(resp.headers().get("x-helexa-advice").is_none());
|
||||||
|
let body: Value = resp.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "context_length_exceeded");
|
||||||
|
}
|
||||||
@@ -118,6 +118,7 @@ async fn test_no_healthy_nodes() {
|
|||||||
}],
|
}],
|
||||||
models_config: "/dev/null".into(),
|
models_config: "/dev/null".into(),
|
||||||
entitlements: Default::default(),
|
entitlements: Default::default(),
|
||||||
|
upstream: Default::default(),
|
||||||
};
|
};
|
||||||
let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
|
||||||
|
|
||||||
|
|||||||
106
crates/cortex-gateway/tests/upstream_chain.rs
Normal file
106
crates/cortex-gateway/tests/upstream_chain.rs
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
//! B3: the chained entitlement provider (local → upstream) and fail-closed
|
||||||
|
//! semantics, exercised against a mock helexa-upstream `/authz/v1`.
|
||||||
|
|
||||||
|
use axum::{Json, Router, routing::post};
|
||||||
|
use cortex_core::config::{ApiKeyConfig, EntitlementsConfig, UpstreamClientConfig};
|
||||||
|
use cortex_core::entitlements::{AuthError, EntitlementProvider};
|
||||||
|
use cortex_gateway::entitlements_chain::ChainedEntitlementProvider;
|
||||||
|
use cortex_gateway::entitlements_local::LocalEntitlementProvider;
|
||||||
|
use cortex_gateway::entitlements_upstream::UpstreamEntitlementProvider;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Mock upstream: `mesh-key` resolves to a mesh account; anything else 401.
|
||||||
|
/// reserve always grants reservation 1.
|
||||||
|
async fn spawn_mock_upstream() -> String {
|
||||||
|
async fn resolve(Json(body): Json<Value>) -> axum::response::Response {
|
||||||
|
use axum::response::IntoResponse;
|
||||||
|
if body["api_key"] == "mesh-key" {
|
||||||
|
Json(json!({"principal": {"account_id": "mesh-acct", "key_id": "mesh-key-1"}}))
|
||||||
|
.into_response()
|
||||||
|
} else {
|
||||||
|
(
|
||||||
|
axum::http::StatusCode::UNAUTHORIZED,
|
||||||
|
Json(json!({"error": {"code": "invalid_api_key"}})),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
async fn reserve() -> Json<Value> {
|
||||||
|
Json(json!({ "reservation_id": 1 }))
|
||||||
|
}
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/authz/v1/resolve", post(resolve))
|
||||||
|
.route("/authz/v1/reserve", post(reserve));
|
||||||
|
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 local_with_key() -> LocalEntitlementProvider {
|
||||||
|
let cfg = EntitlementsConfig {
|
||||||
|
require_auth: false,
|
||||||
|
keys: vec![ApiKeyConfig {
|
||||||
|
key: "local-key".into(),
|
||||||
|
account_id: "op".into(),
|
||||||
|
key_id: None,
|
||||||
|
hard_cap: None,
|
||||||
|
window: Default::default(),
|
||||||
|
}],
|
||||||
|
};
|
||||||
|
LocalEntitlementProvider::from_config(&cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chain(local: LocalEntitlementProvider, url: &str) -> ChainedEntitlementProvider {
|
||||||
|
let upstream = UpstreamEntitlementProvider::new(&UpstreamClientConfig {
|
||||||
|
enabled: true,
|
||||||
|
url: url.to_string(),
|
||||||
|
bearer: "client-secret".into(),
|
||||||
|
timeout_secs: 5,
|
||||||
|
served_usage_report_interval_secs: 60,
|
||||||
|
});
|
||||||
|
ChainedEntitlementProvider::new(local, upstream)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn local_key_resolves_locally() {
|
||||||
|
let url = spawn_mock_upstream().await;
|
||||||
|
let c = chain(local_with_key(), &url);
|
||||||
|
let p = c.resolve("local-key").await.expect("local resolves");
|
||||||
|
assert_eq!(p.account_id, "op");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unknown_key_falls_through_to_upstream() {
|
||||||
|
let url = spawn_mock_upstream().await;
|
||||||
|
let c = chain(local_with_key(), &url);
|
||||||
|
let p = c.resolve("mesh-key").await.expect("upstream resolves");
|
||||||
|
assert_eq!(p.account_id, "mesh-acct");
|
||||||
|
assert_eq!(p.key_id, "mesh-key-1");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unknown_everywhere_is_invalid_key() {
|
||||||
|
let url = spawn_mock_upstream().await;
|
||||||
|
let c = chain(local_with_key(), &url);
|
||||||
|
match c.resolve("nope").await {
|
||||||
|
Err(AuthError::InvalidKey) => {}
|
||||||
|
other => panic!("expected InvalidKey, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn upstream_unreachable_fails_closed_as_unavailable() {
|
||||||
|
// No mock — point at a dead port. A locally-unknown key must surface
|
||||||
|
// Unavailable (→ 503), never InvalidKey (→ 401).
|
||||||
|
let c = chain(local_with_key(), "http://127.0.0.1:1");
|
||||||
|
match c.resolve("some-mesh-key").await {
|
||||||
|
Err(AuthError::Unavailable { retry_after_secs }) => assert!(retry_after_secs > 0),
|
||||||
|
other => panic!("expected Unavailable, got {other:?}"),
|
||||||
|
}
|
||||||
|
// A local key still resolves without touching upstream.
|
||||||
|
assert_eq!(c.resolve("local-key").await.unwrap().account_id, "op");
|
||||||
|
}
|
||||||
@@ -35,6 +35,9 @@ pub fn api_routes(state: ApiState) -> Router {
|
|||||||
.route("/api/health", get(health))
|
.route("/api/health", get(health))
|
||||||
.route("/api/dimensions", get(dimensions))
|
.route("/api/dimensions", get(dimensions))
|
||||||
.route("/api/summary", get(summary))
|
.route("/api/summary", get(summary))
|
||||||
|
.route("/api/scaling", get(scaling))
|
||||||
|
.route("/api/swap", get(swap))
|
||||||
|
.route("/api/capability", get(capability))
|
||||||
.route("/api/series", get(series))
|
.route("/api/series", get(series))
|
||||||
.route("/api/runs", get(runs))
|
.route("/api/runs", get(runs))
|
||||||
.layer(CorsLayer::permissive())
|
.layer(CorsLayer::permissive())
|
||||||
@@ -73,6 +76,30 @@ async fn summary(
|
|||||||
store.summary().map(Json).map_err(err500)
|
store.summary().map(Json).map_err(err500)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Context-length scaling curves per (target, model) — prefill & decode
|
||||||
|
/// tok/s vs context, with decode-flatness (#88).
|
||||||
|
async fn scaling(
|
||||||
|
State(s): State<ApiState>,
|
||||||
|
) -> Result<Json<Vec<crate::store::ScalingCurve>>, ApiError> {
|
||||||
|
let store = s.lock().await;
|
||||||
|
store.scaling().map(Json).map_err(err500)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cold-load / model-swap costs per (target, model) — reload latency + cold
|
||||||
|
/// first-request (#90).
|
||||||
|
async fn swap(State(s): State<ApiState>) -> Result<Json<Vec<crate::store::SwapCost>>, ApiError> {
|
||||||
|
let store = s.lock().await;
|
||||||
|
store.swap_costs().map(Json).map_err(err500)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Capability-probe runs — stored artifacts + quality scores (#91).
|
||||||
|
async fn capability(
|
||||||
|
State(s): State<ApiState>,
|
||||||
|
) -> Result<Json<Vec<crate::store::CapabilityRun>>, ApiError> {
|
||||||
|
let store = s.lock().await;
|
||||||
|
store.capability_runs(false).map(Json).map_err(err500)
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct SeriesQuery {
|
struct SeriesQuery {
|
||||||
/// Optional — when omitted the store resolves the host serving this model.
|
/// Optional — when omitted the store resolves the host serving this model.
|
||||||
|
|||||||
@@ -3,10 +3,10 @@
|
|||||||
//! `openai` targets use the OpenAI-compatible surface (preliminary).
|
//! `openai` targets use the OpenAI-compatible surface (preliminary).
|
||||||
|
|
||||||
use crate::config::{TargetConfig, TargetKind};
|
use crate::config::{TargetConfig, TargetKind};
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result, anyhow};
|
||||||
use cortex_core::build_info::BuildInfo;
|
use cortex_core::build_info::BuildInfo;
|
||||||
use cortex_core::discovery::DiscoveryResponse;
|
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||||
use cortex_core::harness::ModelInfo;
|
use cortex_core::harness::{ModelInfo, ModelSpec};
|
||||||
use cortex_core::openai::ModelsResponse;
|
use cortex_core::openai::ModelsResponse;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
@@ -94,6 +94,84 @@ impl TargetClient {
|
|||||||
Ok(Some(disco))
|
Ok(Some(disco))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Runtime device health (neuron only): per-GPU VRAM used/free,
|
||||||
|
/// utilization, and temperature from `GET /health`. Bench samples this
|
||||||
|
/// around each measured run to record VRAM high-water + GPU telemetry
|
||||||
|
/// (#87). Returns `Ok(None)` for non-neuron targets; a soft `Ok(None)`
|
||||||
|
/// (not an error) on transport failure so a flaky `/health` never fails
|
||||||
|
/// a measurement.
|
||||||
|
pub async fn fetch_health(&self, target: &TargetConfig) -> Result<Option<HealthResponse>> {
|
||||||
|
if target.kind != TargetKind::Neuron {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
let base = target.endpoint.trim_end_matches('/');
|
||||||
|
let health = self
|
||||||
|
.http
|
||||||
|
.get(format!("{base}/health"))
|
||||||
|
.timeout(META_TIMEOUT)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.context("GET /health")?
|
||||||
|
.error_for_status()
|
||||||
|
.context("GET /health status")?
|
||||||
|
.json::<HealthResponse>()
|
||||||
|
.await
|
||||||
|
.context("decoding /health")?;
|
||||||
|
Ok(Some(health))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unload a model (neuron only): `POST /models/unload {model_id}`.
|
||||||
|
/// Used by the deliberate swap-cost measurement (#90), never the sweep.
|
||||||
|
pub async fn unload_model(&self, target: &TargetConfig, model_id: &str) -> Result<()> {
|
||||||
|
let base = target.endpoint.trim_end_matches('/');
|
||||||
|
self.http
|
||||||
|
.post(format!("{base}/models/unload"))
|
||||||
|
.json(&serde_json::json!({ "model_id": model_id }))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.context("POST /models/unload")?
|
||||||
|
.error_for_status()
|
||||||
|
.context("POST /models/unload status")?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a model from a spec (neuron only): `POST /models/load`. neuron
|
||||||
|
/// returns synchronously once loaded, so the call duration is the reload
|
||||||
|
/// cost the swap-cost measurement records (#90).
|
||||||
|
pub async fn load_model(&self, target: &TargetConfig, spec: &ModelSpec) -> Result<()> {
|
||||||
|
let base = target.endpoint.trim_end_matches('/');
|
||||||
|
self.http
|
||||||
|
.post(format!("{base}/models/load"))
|
||||||
|
.json(spec)
|
||||||
|
// A cold load can take tens of seconds; use the full request
|
||||||
|
// timeout rather than the short metadata one.
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.context("POST /models/load")?
|
||||||
|
.error_for_status()
|
||||||
|
.context("POST /models/load status")?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reconstruct a reload [`ModelSpec`] from a model's `/models` entry.
|
||||||
|
/// Tensor-parallel is inferred from the device count; `quant` is left
|
||||||
|
/// `None` for neuron to resolve from the catalogue / its prior load.
|
||||||
|
pub fn spec_from_info(info: &ModelInfo) -> Result<ModelSpec> {
|
||||||
|
if info.devices.is_empty() {
|
||||||
|
return Err(anyhow!(
|
||||||
|
"model '{}' reports no devices; cannot reconstruct a load spec",
|
||||||
|
info.id
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(ModelSpec {
|
||||||
|
model_id: info.id.clone(),
|
||||||
|
harness: info.harness.clone(),
|
||||||
|
quant: None,
|
||||||
|
tensor_parallel: (info.devices.len() > 1).then_some(info.devices.len() as u32),
|
||||||
|
devices: Some(info.devices.clone()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// Warm models — those ready to serve without a cold load.
|
/// Warm models — those ready to serve without a cold load.
|
||||||
///
|
///
|
||||||
/// Neuron: `GET /models` filtered to `status == "loaded"` (skips
|
/// Neuron: `GET /models` filtered to `status == "loaded"` (skips
|
||||||
|
|||||||
@@ -104,6 +104,37 @@ pub struct ScenarioConfig {
|
|||||||
/// Max generated tokens per request.
|
/// Max generated tokens per request.
|
||||||
#[serde(default = "default_max_tokens")]
|
#[serde(default = "default_max_tokens")]
|
||||||
pub max_tokens: u64,
|
pub max_tokens: u64,
|
||||||
|
/// Concurrency levels (#89): one `concurrency:<n>` scenario per entry,
|
||||||
|
/// each firing N simultaneous streams. Defaults to empty (opt-in) — a
|
||||||
|
/// burst puts real load on a serving fleet, so operators enable it
|
||||||
|
/// deliberately, e.g. `concurrency_levels = [2, 4, 8]`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub concurrency_levels: Vec<u32>,
|
||||||
|
/// Approximate prompt size (tokens) used by the concurrency scenarios.
|
||||||
|
#[serde(default = "default_concurrency_prompt_tokens")]
|
||||||
|
pub concurrency_prompt_tokens: u32,
|
||||||
|
/// Capability probes (#91): one `capability:<name>` scenario per entry,
|
||||||
|
/// each running a fixed prompt and storing the full output artifact for
|
||||||
|
/// quality scoring (manual now, LLM-judge later). Defaults to empty
|
||||||
|
/// (opt-in) — these generate long outputs and exist to compare reasoning
|
||||||
|
/// quality across models, not to run on every sweep by default.
|
||||||
|
#[serde(default)]
|
||||||
|
pub capability_probes: Vec<CapabilityProbe>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One capability probe: a named prompt whose output is stored and scored
|
||||||
|
/// for quality (#91). The probe is deterministic (temperature 0) so the
|
||||||
|
/// same model+build produces a stable artifact to score.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CapabilityProbe {
|
||||||
|
/// Stable id fragment — the scenario becomes `capability:<name>`.
|
||||||
|
pub name: String,
|
||||||
|
/// The prompt sent verbatim as the user message.
|
||||||
|
pub prompt: String,
|
||||||
|
/// Generation budget for the probe (planning answers are long; the
|
||||||
|
/// default 256 is too small). Defaults to 2048.
|
||||||
|
#[serde(default = "default_capability_max_tokens")]
|
||||||
|
pub max_tokens: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for ScenarioConfig {
|
impl Default for ScenarioConfig {
|
||||||
@@ -111,6 +142,9 @@ impl Default for ScenarioConfig {
|
|||||||
ScenarioConfig {
|
ScenarioConfig {
|
||||||
prompt_sizes: default_prompt_sizes(),
|
prompt_sizes: default_prompt_sizes(),
|
||||||
max_tokens: default_max_tokens(),
|
max_tokens: default_max_tokens(),
|
||||||
|
concurrency_levels: Vec::new(),
|
||||||
|
concurrency_prompt_tokens: default_concurrency_prompt_tokens(),
|
||||||
|
capability_probes: Vec::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -187,6 +221,12 @@ fn default_prompt_sizes() -> Vec<u32> {
|
|||||||
fn default_max_tokens() -> u64 {
|
fn default_max_tokens() -> u64 {
|
||||||
256
|
256
|
||||||
}
|
}
|
||||||
|
fn default_concurrency_prompt_tokens() -> u32 {
|
||||||
|
512
|
||||||
|
}
|
||||||
|
fn default_capability_max_tokens() -> u64 {
|
||||||
|
2048
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
// Jail's closure must return figment::Result; the large-Err type is
|
// Jail's closure must return figment::Result; the large-Err type is
|
||||||
|
|||||||
@@ -43,6 +43,33 @@ enum Command {
|
|||||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||||
config: String,
|
config: String,
|
||||||
},
|
},
|
||||||
|
/// Measure cold-load / model-swap cost (#90): for each neuron target's
|
||||||
|
/// warm models, unload → time reload → time a cold first request, recorded
|
||||||
|
/// under scenario "swap". DELIBERATE — takes each model offline for its
|
||||||
|
/// reload, so run it in a maintenance window, not against live traffic.
|
||||||
|
SwapCost {
|
||||||
|
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||||
|
config: String,
|
||||||
|
},
|
||||||
|
/// Attach a quality score to a capability-probe run (#91). Find run ids
|
||||||
|
/// with `report --capability`. `--scorer` records who scored it
|
||||||
|
/// (defaults to "manual"); a future LLM-judge would set e.g. "llm:…".
|
||||||
|
Score {
|
||||||
|
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||||
|
config: String,
|
||||||
|
/// Override the SQLite path (skips reading the config file).
|
||||||
|
#[arg(long)]
|
||||||
|
db: Option<String>,
|
||||||
|
/// The run id to score.
|
||||||
|
#[arg(long)]
|
||||||
|
id: i64,
|
||||||
|
/// The quality score to attach (scale is the operator's rubric).
|
||||||
|
#[arg(long)]
|
||||||
|
score: f64,
|
||||||
|
/// Who/what produced the score.
|
||||||
|
#[arg(long, default_value = "manual")]
|
||||||
|
scorer: String,
|
||||||
|
},
|
||||||
/// Render recorded results. Uses `--db` if given, else the db_path
|
/// Render recorded results. Uses `--db` if given, else the db_path
|
||||||
/// from `--config`.
|
/// from `--config`.
|
||||||
Report {
|
Report {
|
||||||
@@ -54,6 +81,19 @@ enum Command {
|
|||||||
/// Output format.
|
/// Output format.
|
||||||
#[arg(long, default_value = "md")]
|
#[arg(long, default_value = "md")]
|
||||||
format: Format,
|
format: Format,
|
||||||
|
/// Render the context-length scaling view (prefill & decode tok/s
|
||||||
|
/// vs context per model, with decode-flatness) instead of the flat
|
||||||
|
/// results table (#88).
|
||||||
|
#[arg(long)]
|
||||||
|
scaling: bool,
|
||||||
|
/// Render the cold-load / model-swap cost view (#90) instead of the
|
||||||
|
/// flat results table.
|
||||||
|
#[arg(long)]
|
||||||
|
swap: bool,
|
||||||
|
/// Render the capability-probe view (#91): stored artifacts + quality
|
||||||
|
/// scores, with per-model median.
|
||||||
|
#[arg(long)]
|
||||||
|
capability: bool,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -122,16 +162,79 @@ async fn run(cli: Cli) -> Result<()> {
|
|||||||
);
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
Command::Report { config, db, format } => {
|
Command::SwapCost { config } => {
|
||||||
|
let cfg = load_config(&config)?;
|
||||||
|
require_targets(&cfg)?;
|
||||||
|
let sweeper = Sweeper::new(cfg)?;
|
||||||
|
tracing::warn!(
|
||||||
|
"swap-cost: cycling each warm model (unload → reload → cold request); models go offline during reload"
|
||||||
|
);
|
||||||
|
let summary = sweeper.swap_cost_once().await?;
|
||||||
|
tracing::info!(
|
||||||
|
measured = summary.measured,
|
||||||
|
failed = summary.failed,
|
||||||
|
unreachable = summary.targets_unreachable,
|
||||||
|
"swap-cost measurement complete"
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Command::Score {
|
||||||
|
config,
|
||||||
|
db,
|
||||||
|
id,
|
||||||
|
score,
|
||||||
|
scorer,
|
||||||
|
} => {
|
||||||
let db_path = match db {
|
let db_path = match db {
|
||||||
Some(p) => p,
|
Some(p) => p,
|
||||||
None => load_config(&config)?.bench.db_path,
|
None => load_config(&config)?.bench.db_path,
|
||||||
};
|
};
|
||||||
let store = Store::open(&db_path)?;
|
let store = Store::open(&db_path)?;
|
||||||
let rows = store.report_rows()?;
|
match store.set_score(id, score, &scorer)? {
|
||||||
let rendered = match format {
|
0 => anyhow::bail!("no run with id {id}"),
|
||||||
Format::Md => report::render_markdown(&rows),
|
_ => {
|
||||||
Format::Json => report::render_json(&rows)?,
|
println!("scored run {id}: {score} ({scorer})");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Command::Report {
|
||||||
|
config,
|
||||||
|
db,
|
||||||
|
format,
|
||||||
|
scaling,
|
||||||
|
swap,
|
||||||
|
capability,
|
||||||
|
} => {
|
||||||
|
let db_path = match db {
|
||||||
|
Some(p) => p,
|
||||||
|
None => load_config(&config)?.bench.db_path,
|
||||||
|
};
|
||||||
|
let store = Store::open(&db_path)?;
|
||||||
|
let rendered = if capability {
|
||||||
|
let runs = store.capability_runs(false)?;
|
||||||
|
match format {
|
||||||
|
Format::Md => report::render_capability_markdown(&runs),
|
||||||
|
Format::Json => report::render_capability_json(&runs)?,
|
||||||
|
}
|
||||||
|
} else if swap {
|
||||||
|
let costs = store.swap_costs()?;
|
||||||
|
match format {
|
||||||
|
Format::Md => report::render_swap_markdown(&costs),
|
||||||
|
Format::Json => report::render_swap_json(&costs)?,
|
||||||
|
}
|
||||||
|
} else if scaling {
|
||||||
|
let curves = store.scaling()?;
|
||||||
|
match format {
|
||||||
|
Format::Md => report::render_scaling_markdown(&curves),
|
||||||
|
Format::Json => report::render_scaling_json(&curves)?,
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
let rows = store.report_rows()?;
|
||||||
|
match format {
|
||||||
|
Format::Md => report::render_markdown(&rows),
|
||||||
|
Format::Json => report::render_json(&rows)?,
|
||||||
|
}
|
||||||
};
|
};
|
||||||
println!("{rendered}");
|
println!("{rendered}");
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|||||||
@@ -3,28 +3,36 @@
|
|||||||
//! doc: engine, model, prompt tok, TTFT (s), decode tok/s, total (s),
|
//! doc: engine, model, prompt tok, TTFT (s), decode tok/s, total (s),
|
||||||
//! plus the build SHA each cell was measured against.
|
//! plus the build SHA each cell was measured against.
|
||||||
|
|
||||||
use crate::store::ReportRow;
|
use crate::store::{CapabilityRun, ReportRow, ScalingCurve, SwapCost};
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
|
|
||||||
pub fn render_markdown(rows: &[ReportRow]) -> String {
|
pub fn render_markdown(rows: &[ReportRow]) -> String {
|
||||||
let mut out = String::new();
|
let mut out = String::new();
|
||||||
out.push_str(
|
out.push_str(
|
||||||
"| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) | build | n |\n",
|
"| engine | model | prompt tok | prefill tok/s | TTFT (s) | TTFT p95 | \
|
||||||
|
decode tok/s | total (s) | total p95 | VRAM (GB) | conc | queue ms | rej | build | n |\n",
|
||||||
);
|
);
|
||||||
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
|
out.push_str("|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---|---:|\n");
|
||||||
for r in rows {
|
for r in rows {
|
||||||
let ptok = r
|
let ptok = r
|
||||||
.prompt_tokens
|
.prompt_tokens
|
||||||
.map(|t| t.to_string())
|
.map(|t| t.to_string())
|
||||||
.unwrap_or_else(|| format!("~{}", r.prompt_size_approx));
|
.unwrap_or_else(|| format!("~{}", r.prompt_size_approx));
|
||||||
out.push_str(&format!(
|
out.push_str(&format!(
|
||||||
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
"| {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||||
r.target_name,
|
r.target_name,
|
||||||
r.model_id,
|
r.model_id,
|
||||||
ptok,
|
ptok,
|
||||||
|
fmt_opt(r.prefill_tps_median, 1),
|
||||||
fmt_opt(r.ttft_s_median, 3),
|
fmt_opt(r.ttft_s_median, 3),
|
||||||
|
fmt_opt(r.ttft_s_p95, 3),
|
||||||
fmt_opt(r.decode_tps_median, 1),
|
fmt_opt(r.decode_tps_median, 1),
|
||||||
fmt_opt(r.total_s_median, 3),
|
fmt_opt(r.total_s_median, 3),
|
||||||
|
fmt_opt(r.total_s_p95, 3),
|
||||||
|
fmt_vram(r.vram_used_mb_median, r.vram_total_mb),
|
||||||
|
fmt_u64(r.concurrency),
|
||||||
|
fmt_opt(r.queue_wait_ms_median, 0),
|
||||||
|
fmt_opt(r.rejected_median, 0),
|
||||||
r.git_sha,
|
r.git_sha,
|
||||||
r.samples,
|
r.samples,
|
||||||
));
|
));
|
||||||
@@ -43,8 +51,23 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
|||||||
"prompt_size_approx": r.prompt_size_approx,
|
"prompt_size_approx": r.prompt_size_approx,
|
||||||
"prompt_tokens": r.prompt_tokens,
|
"prompt_tokens": r.prompt_tokens,
|
||||||
"ttft_s_median": r.ttft_s_median,
|
"ttft_s_median": r.ttft_s_median,
|
||||||
|
"ttft_s_p95": r.ttft_s_p95,
|
||||||
|
"ttft_s_p99": r.ttft_s_p99,
|
||||||
"decode_tps_median": r.decode_tps_median,
|
"decode_tps_median": r.decode_tps_median,
|
||||||
"total_s_median": r.total_s_median,
|
"total_s_median": r.total_s_median,
|
||||||
|
"total_s_p95": r.total_s_p95,
|
||||||
|
"total_s_p99": r.total_s_p99,
|
||||||
|
"prefill_ms_median": r.prefill_ms_median,
|
||||||
|
"decode_ms_median": r.decode_ms_median,
|
||||||
|
"prefill_tps_median": r.prefill_tps_median,
|
||||||
|
"vram_used_mb_median": r.vram_used_mb_median,
|
||||||
|
"vram_total_mb": r.vram_total_mb,
|
||||||
|
"gpu_util_pct_median": r.gpu_util_pct_median,
|
||||||
|
"gpu_temp_c_median": r.gpu_temp_c_median,
|
||||||
|
"concurrency": r.concurrency,
|
||||||
|
"ttft_p95_load_s": r.ttft_p95_load_s,
|
||||||
|
"queue_wait_ms_median": r.queue_wait_ms_median,
|
||||||
|
"rejected_median": r.rejected_median,
|
||||||
"git_sha": r.git_sha,
|
"git_sha": r.git_sha,
|
||||||
"samples": r.samples,
|
"samples": r.samples,
|
||||||
"gpu": r.gpu,
|
"gpu": r.gpu,
|
||||||
@@ -54,6 +77,179 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
|||||||
Ok(serde_json::to_string_pretty(&arr)?)
|
Ok(serde_json::to_string_pretty(&arr)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Context-length scaling view (#88): one block per (target, model) with
|
||||||
|
/// prefill & decode tok/s vs context, then the decode-flatness verdict.
|
||||||
|
pub fn render_scaling_markdown(curves: &[ScalingCurve]) -> String {
|
||||||
|
let mut out = String::new();
|
||||||
|
for c in curves {
|
||||||
|
let gpu = c.gpu.as_deref().unwrap_or("");
|
||||||
|
out.push_str(&format!(
|
||||||
|
"### {} · {} (`{}`{})\n\n",
|
||||||
|
c.target_name,
|
||||||
|
c.model_id,
|
||||||
|
c.git_sha,
|
||||||
|
if gpu.is_empty() {
|
||||||
|
String::new()
|
||||||
|
} else {
|
||||||
|
format!(", {gpu}")
|
||||||
|
},
|
||||||
|
));
|
||||||
|
out.push_str("| ctx tok | prefill tok/s | decode tok/s | n |\n");
|
||||||
|
out.push_str("|---:|---:|---:|---:|\n");
|
||||||
|
for p in &c.points {
|
||||||
|
let ctx = p
|
||||||
|
.prompt_tokens
|
||||||
|
.map(|t| t.to_string())
|
||||||
|
.unwrap_or_else(|| format!("~{}", p.prompt_size));
|
||||||
|
out.push_str(&format!(
|
||||||
|
"| {} | {} | {} | {} |\n",
|
||||||
|
ctx,
|
||||||
|
fmt_opt(p.prefill_tps, 1),
|
||||||
|
fmt_opt(p.decode_tps, 1),
|
||||||
|
p.samples,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
match c.decode_flatness {
|
||||||
|
Some(f) => out.push_str(&format!(
|
||||||
|
"\ndecode flatness: {f:.2} — decode tok/s {} across the context range \
|
||||||
|
({})\n\n",
|
||||||
|
if f >= 0.9 {
|
||||||
|
"holds"
|
||||||
|
} else if f >= 0.7 {
|
||||||
|
"softens"
|
||||||
|
} else {
|
||||||
|
"drops sharply"
|
||||||
|
},
|
||||||
|
if f >= 0.9 {
|
||||||
|
"Gated-DeltaNet O(1) decode confirmed"
|
||||||
|
} else {
|
||||||
|
"investigate where it breaks"
|
||||||
|
},
|
||||||
|
)),
|
||||||
|
None => out.push_str("\ndecode flatness: — (need ≥2 context points)\n\n"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn render_scaling_json(curves: &[ScalingCurve]) -> Result<String> {
|
||||||
|
Ok(serde_json::to_string_pretty(curves)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cold-load / model-swap cost view (#90): reload latency + cold
|
||||||
|
/// first-request per model.
|
||||||
|
pub fn render_swap_markdown(costs: &[SwapCost]) -> String {
|
||||||
|
let mut out = String::new();
|
||||||
|
out.push_str(
|
||||||
|
"| engine | model | unload (s) | reload (s) | cold TTFT (s) | cold total (s) | build | n |\n",
|
||||||
|
);
|
||||||
|
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
|
||||||
|
for c in costs {
|
||||||
|
out.push_str(&format!(
|
||||||
|
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||||
|
c.target_name,
|
||||||
|
c.model_id,
|
||||||
|
fmt_ms_as_s(c.unload_ms_median),
|
||||||
|
fmt_ms_as_s(c.load_ms_median),
|
||||||
|
fmt_opt(c.cold_ttft_s_median, 3),
|
||||||
|
fmt_opt(c.cold_total_s_median, 3),
|
||||||
|
c.git_sha,
|
||||||
|
c.samples,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn render_swap_json(costs: &[SwapCost]) -> Result<String> {
|
||||||
|
Ok(serde_json::to_string_pretty(costs)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Capability-probe view (#91): per (model, probe) the median quality score
|
||||||
|
/// (the A/B number), then each run's id, score, and an artifact snippet so
|
||||||
|
/// unscored runs can be located and scored (`helexa-bench score --id …`).
|
||||||
|
pub fn render_capability_markdown(runs: &[CapabilityRun]) -> String {
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
let mut groups: BTreeMap<(String, String, String), Vec<&CapabilityRun>> = BTreeMap::new();
|
||||||
|
for r in runs {
|
||||||
|
groups
|
||||||
|
.entry((
|
||||||
|
r.target_name.clone(),
|
||||||
|
r.model_id.clone(),
|
||||||
|
r.scenario_id.clone(),
|
||||||
|
))
|
||||||
|
.or_default()
|
||||||
|
.push(r);
|
||||||
|
}
|
||||||
|
let mut out = String::new();
|
||||||
|
for ((target, model, scenario), rs) in groups {
|
||||||
|
let scores: Vec<f64> = rs.iter().filter_map(|r| r.quality_score).collect();
|
||||||
|
let median = median_slice(&scores);
|
||||||
|
out.push_str(&format!(
|
||||||
|
"### {target} · {model} · {scenario} — median score {} ({}/{} scored)\n\n",
|
||||||
|
median
|
||||||
|
.map(|m| format!("{m:.1}"))
|
||||||
|
.unwrap_or_else(|| "—".into()),
|
||||||
|
scores.len(),
|
||||||
|
rs.len(),
|
||||||
|
));
|
||||||
|
out.push_str("| run | score | scorer | build | artifact (snippet) |\n");
|
||||||
|
out.push_str("|---:|---:|---|---|---|\n");
|
||||||
|
for r in rs {
|
||||||
|
out.push_str(&format!(
|
||||||
|
"| {} | {} | {} | `{}` | {} |\n",
|
||||||
|
r.id,
|
||||||
|
r.quality_score
|
||||||
|
.map(|s| format!("{s:.1}"))
|
||||||
|
.unwrap_or_else(|| "—".into()),
|
||||||
|
r.scorer.as_deref().unwrap_or("—"),
|
||||||
|
r.git_sha,
|
||||||
|
snippet(r.artifact.as_deref()),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
out.push('\n');
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn render_capability_json(runs: &[CapabilityRun]) -> Result<String> {
|
||||||
|
Ok(serde_json::to_string_pretty(runs)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// First ~80 chars of an artifact on one line, for the table cell.
|
||||||
|
fn snippet(artifact: Option<&str>) -> String {
|
||||||
|
match artifact {
|
||||||
|
Some(a) => {
|
||||||
|
let one_line: String = a.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||||
|
let trimmed: String = one_line.chars().take(80).collect();
|
||||||
|
if one_line.chars().count() > 80 {
|
||||||
|
format!("{trimmed}…")
|
||||||
|
} else {
|
||||||
|
trimmed
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => "—".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn median_slice(v: &[f64]) -> Option<f64> {
|
||||||
|
if v.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let mut s = v.to_vec();
|
||||||
|
s.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
let lo = (s.len() - 1) / 2;
|
||||||
|
let hi = s.len() / 2;
|
||||||
|
Some((s[lo] + s[hi]) / 2.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Milliseconds rendered as seconds (reload costs read naturally in s).
|
||||||
|
fn fmt_ms_as_s(ms: Option<f64>) -> String {
|
||||||
|
match ms {
|
||||||
|
Some(x) => format!("{:.2}", x / 1000.0),
|
||||||
|
None => "—".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
||||||
match v {
|
match v {
|
||||||
Some(x) => format!("{x:.places$}"),
|
Some(x) => format!("{x:.places$}"),
|
||||||
@@ -61,9 +257,113 @@ fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Integer cell (concurrency width); `—` when unset (non-concurrency rows).
|
||||||
|
fn fmt_u64(v: Option<u64>) -> String {
|
||||||
|
match v {
|
||||||
|
Some(x) => x.to_string(),
|
||||||
|
None => "—".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `used/total` in GB (e.g. `42.0/64.0`) — the headroom-at-a-glance cell.
|
||||||
|
/// `used` alone if the node total is unknown; `—` if no telemetry.
|
||||||
|
fn fmt_vram(used_mb: Option<f64>, total_mb: Option<u64>) -> String {
|
||||||
|
match (used_mb, total_mb) {
|
||||||
|
(Some(u), Some(t)) => format!("{:.1}/{:.1}", u / 1024.0, t as f64 / 1024.0),
|
||||||
|
(Some(u), None) => format!("{:.1}", u / 1024.0),
|
||||||
|
_ => "—".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::store::{ScalingCurve, ScalingPoint};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn capability_markdown_groups_with_median_and_snippet() {
|
||||||
|
let runs = vec![
|
||||||
|
CapabilityRun {
|
||||||
|
id: 7,
|
||||||
|
ts: "t".into(),
|
||||||
|
target_name: "beast".into(),
|
||||||
|
model_id: "m".into(),
|
||||||
|
scenario_id: "capability:plan".into(),
|
||||||
|
git_sha: "abc".into(),
|
||||||
|
quality_score: Some(8.0),
|
||||||
|
scorer: Some("manual".into()),
|
||||||
|
artifact: Some("A detailed plan with trade-offs and sequencing.".into()),
|
||||||
|
},
|
||||||
|
CapabilityRun {
|
||||||
|
id: 8,
|
||||||
|
ts: "t".into(),
|
||||||
|
target_name: "beast".into(),
|
||||||
|
model_id: "m".into(),
|
||||||
|
scenario_id: "capability:plan".into(),
|
||||||
|
git_sha: "abc".into(),
|
||||||
|
quality_score: Some(6.0),
|
||||||
|
scorer: Some("manual".into()),
|
||||||
|
artifact: Some("Shorter plan.".into()),
|
||||||
|
},
|
||||||
|
];
|
||||||
|
let md = render_capability_markdown(&runs);
|
||||||
|
assert!(md.contains("capability:plan"));
|
||||||
|
assert!(md.contains("median score 7.0")); // median(8,6)
|
||||||
|
assert!(md.contains("trade-offs"));
|
||||||
|
assert!(md.contains("| 7 |"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn swap_markdown_renders_reload_and_cold_costs() {
|
||||||
|
let costs = vec![SwapCost {
|
||||||
|
target_name: "beast".into(),
|
||||||
|
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||||
|
git_sha: "abc1234".into(),
|
||||||
|
gpu: Some("2× RTX 5090".into()),
|
||||||
|
unload_ms_median: Some(320.0),
|
||||||
|
load_ms_median: Some(25000.0),
|
||||||
|
cold_ttft_s_median: Some(2.5),
|
||||||
|
cold_total_s_median: Some(5.0),
|
||||||
|
samples: 3,
|
||||||
|
}];
|
||||||
|
let md = render_swap_markdown(&costs);
|
||||||
|
assert!(md.contains("reload (s)"));
|
||||||
|
assert!(md.contains("beast"));
|
||||||
|
assert!(md.contains("25.00")); // 25000 ms → 25.00 s
|
||||||
|
assert!(md.contains("2.500"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn scaling_markdown_renders_curve_and_flatness() {
|
||||||
|
let curves = vec![ScalingCurve {
|
||||||
|
target_name: "beast".into(),
|
||||||
|
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||||
|
git_sha: "abc1234".into(),
|
||||||
|
gpu: Some("2× RTX 5090".into()),
|
||||||
|
points: vec![
|
||||||
|
ScalingPoint {
|
||||||
|
prompt_size: 128,
|
||||||
|
prompt_tokens: Some(130),
|
||||||
|
prefill_tps: Some(900.0),
|
||||||
|
decode_tps: Some(50.0),
|
||||||
|
samples: 5,
|
||||||
|
},
|
||||||
|
ScalingPoint {
|
||||||
|
prompt_size: 4096,
|
||||||
|
prompt_tokens: Some(4100),
|
||||||
|
prefill_tps: Some(2800.0),
|
||||||
|
decode_tps: Some(48.0),
|
||||||
|
samples: 5,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
decode_flatness: Some(0.96),
|
||||||
|
}];
|
||||||
|
let md = render_scaling_markdown(&curves);
|
||||||
|
assert!(md.contains("### beast · Qwen/Qwen3.6-27B"));
|
||||||
|
assert!(md.contains("ctx tok"));
|
||||||
|
assert!(md.contains("decode flatness: 0.96"));
|
||||||
|
assert!(md.contains("holds"));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn markdown_has_header_and_row() {
|
fn markdown_has_header_and_row() {
|
||||||
@@ -77,14 +377,36 @@ mod tests {
|
|||||||
ttft_s_median: Some(0.123),
|
ttft_s_median: Some(0.123),
|
||||||
decode_tps_median: Some(45.6),
|
decode_tps_median: Some(45.6),
|
||||||
total_s_median: Some(1.234),
|
total_s_median: Some(1.234),
|
||||||
|
ttft_s_p95: Some(0.222),
|
||||||
|
ttft_s_p99: Some(0.250),
|
||||||
|
total_s_p95: Some(1.5),
|
||||||
|
total_s_p99: Some(1.6),
|
||||||
|
prefill_ms_median: Some(120.0),
|
||||||
|
decode_ms_median: Some(1100.0),
|
||||||
|
prefill_tps_median: Some(1066.7),
|
||||||
|
vram_used_mb_median: Some(43008.0),
|
||||||
|
vram_total_mb: Some(65536),
|
||||||
|
gpu_util_pct_median: Some(89.0),
|
||||||
|
gpu_temp_c_median: Some(64.0),
|
||||||
|
concurrency: None,
|
||||||
|
ttft_p95_load_s: None,
|
||||||
|
queue_wait_ms_median: None,
|
||||||
|
rejected_median: None,
|
||||||
samples: 5,
|
samples: 5,
|
||||||
gpu: Some("2× RTX 5090".into()),
|
gpu: Some("2× RTX 5090".into()),
|
||||||
}];
|
}];
|
||||||
let md = render_markdown(&rows);
|
let md = render_markdown(&rows);
|
||||||
assert!(md.contains("| engine |"));
|
assert!(md.contains("| engine |"));
|
||||||
|
assert!(md.contains("prefill tok/s"));
|
||||||
|
assert!(md.contains("VRAM (GB)"));
|
||||||
|
assert!(md.contains("conc"));
|
||||||
assert!(md.contains("beast"));
|
assert!(md.contains("beast"));
|
||||||
assert!(md.contains("`30d50d6`"));
|
assert!(md.contains("`30d50d6`"));
|
||||||
assert!(md.contains("0.123"));
|
assert!(md.contains("0.123"));
|
||||||
|
// p95 column rendered.
|
||||||
|
assert!(md.contains("0.222"));
|
||||||
|
// VRAM used/total in GB (43008/65536 MiB → 42.0/64.0).
|
||||||
|
assert!(md.contains("42.0/64.0"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -99,6 +421,21 @@ mod tests {
|
|||||||
ttft_s_median: Some(0.1),
|
ttft_s_median: Some(0.1),
|
||||||
decode_tps_median: None,
|
decode_tps_median: None,
|
||||||
total_s_median: Some(0.5),
|
total_s_median: Some(0.5),
|
||||||
|
ttft_s_p95: Some(0.1),
|
||||||
|
ttft_s_p99: Some(0.1),
|
||||||
|
total_s_p95: Some(0.5),
|
||||||
|
total_s_p99: Some(0.5),
|
||||||
|
prefill_ms_median: None,
|
||||||
|
decode_ms_median: None,
|
||||||
|
prefill_tps_median: None,
|
||||||
|
vram_used_mb_median: None,
|
||||||
|
vram_total_mb: None,
|
||||||
|
gpu_util_pct_median: None,
|
||||||
|
gpu_temp_c_median: None,
|
||||||
|
concurrency: None,
|
||||||
|
ttft_p95_load_s: None,
|
||||||
|
queue_wait_ms_median: None,
|
||||||
|
rejected_median: None,
|
||||||
samples: 1,
|
samples: 1,
|
||||||
gpu: None,
|
gpu: None,
|
||||||
}];
|
}];
|
||||||
|
|||||||
@@ -62,6 +62,41 @@ pub struct ScenarioMetrics {
|
|||||||
pub prompt_tokens: Option<u64>,
|
pub prompt_tokens: Option<u64>,
|
||||||
/// Completion tokens: from `usage` when present, else content-chunk count.
|
/// Completion tokens: from `usage` when present, else content-chunk count.
|
||||||
pub completion_tokens: u64,
|
pub completion_tokens: u64,
|
||||||
|
/// Server-measured prefill duration (ms), from the `usage.helexa_timing`
|
||||||
|
/// extension (#85). `None` when the server didn't emit it (external
|
||||||
|
/// engines, non-instrumented paths). The honest prefill-phase number,
|
||||||
|
/// distinct from client-observed `ttft_s` which also includes request
|
||||||
|
/// setup + first-byte network latency.
|
||||||
|
pub prefill_ms: Option<u64>,
|
||||||
|
/// Server-measured decode duration (ms), from `usage.helexa_timing`.
|
||||||
|
pub decode_ms: Option<u64>,
|
||||||
|
/// Tokens submitted to prefill — the denominator for prefill tok/s.
|
||||||
|
pub prefill_tokens: Option<u64>,
|
||||||
|
// ── Concurrency / agentic-load fields (#89) ──────────────────────────
|
||||||
|
// Set only by the concurrency scenario, which fans out N simultaneous
|
||||||
|
// streams to characterize the real a0/hermes/opencode workload that
|
||||||
|
// batch-1 single-request measurement can't see. `None` for single
|
||||||
|
// requests. For a concurrency burst, the inherited fields carry the
|
||||||
|
// aggregate: `ttft_s` = median TTFT across streams, `decode_tps` = node
|
||||||
|
// throughput (total tokens / burst window), `total_s` = burst wall-clock,
|
||||||
|
// `completion_tokens` = total across streams.
|
||||||
|
/// Number of simultaneous streams in the burst (the cell dimension).
|
||||||
|
pub concurrency: Option<u32>,
|
||||||
|
/// p95 of per-stream TTFT within the burst — the tail under simultaneous
|
||||||
|
/// load, where batch-1 serialization actually hurts.
|
||||||
|
pub ttft_p95_s: Option<f64>,
|
||||||
|
/// Median per-stream admission queue-wait (ms), approximated as
|
||||||
|
/// `ttft − prefill_ms` (#85): on a batch-1 server, later streams wait for
|
||||||
|
/// earlier ones, so TTFT inflates while server prefill stays constant —
|
||||||
|
/// the gap is the wait. `None` if streams didn't report `helexa_timing`.
|
||||||
|
pub queue_wait_ms_median: Option<f64>,
|
||||||
|
/// Streams shed by admission control (HTTP 429/503) during the burst —
|
||||||
|
/// honest backpressure, not silent failures.
|
||||||
|
pub rejected: Option<u32>,
|
||||||
|
/// Full generated text, captured only by the capability probe (#91) so
|
||||||
|
/// the output can be quality-scored later (manual or LLM-judge). `None`
|
||||||
|
/// for latency/throughput scenarios, which discard the text.
|
||||||
|
pub artifact: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -85,10 +120,13 @@ pub trait Scenario: Send + Sync {
|
|||||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build the active scenario set from config. One chat-latency scenario
|
/// Build the active scenario set from config: one chat-latency scenario per
|
||||||
/// per configured prompt size.
|
/// prompt size, plus one concurrency scenario per configured level (#89).
|
||||||
|
/// Concurrency levels default to empty (opt-in), since a burst puts real
|
||||||
|
/// simultaneous load on a serving fleet — operators enable it deliberately.
|
||||||
pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||||
cfg.prompt_sizes
|
let mut scenarios: Vec<Box<dyn Scenario>> = cfg
|
||||||
|
.prompt_sizes
|
||||||
.iter()
|
.iter()
|
||||||
.map(|&size| {
|
.map(|&size| {
|
||||||
Box::new(ChatLatencyScenario {
|
Box::new(ChatLatencyScenario {
|
||||||
@@ -96,7 +134,46 @@ pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
|||||||
approx_prompt_tokens: size,
|
approx_prompt_tokens: size,
|
||||||
}) as Box<dyn Scenario>
|
}) as Box<dyn Scenario>
|
||||||
})
|
})
|
||||||
.collect()
|
.collect();
|
||||||
|
for &n in &cfg.concurrency_levels {
|
||||||
|
scenarios.push(Box::new(ConcurrencyScenario {
|
||||||
|
id: format!("concurrency:{n}"),
|
||||||
|
concurrency: n,
|
||||||
|
approx_prompt_tokens: cfg.concurrency_prompt_tokens,
|
||||||
|
}) as Box<dyn Scenario>);
|
||||||
|
}
|
||||||
|
for probe in &cfg.capability_probes {
|
||||||
|
scenarios.push(Box::new(CapabilityScenario {
|
||||||
|
id: format!("capability:{}", probe.name),
|
||||||
|
prompt: probe.prompt.clone(),
|
||||||
|
max_tokens: probe.max_tokens,
|
||||||
|
}) as Box<dyn Scenario>);
|
||||||
|
}
|
||||||
|
scenarios
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single small streamed request, timed like a chat-latency run. Used by
|
||||||
|
/// the swap-cost measurement (#90) to capture the cold first-request latency
|
||||||
|
/// straight after a reload. Reuses the shared SSE-timing core.
|
||||||
|
pub async fn cold_probe(ctx: &RunCtx<'_>) -> Result<ScenarioMetrics> {
|
||||||
|
let prompt = build_prompt(128);
|
||||||
|
let payload = chat_payload(ctx, &prompt);
|
||||||
|
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload))
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow!("cold probe timed out after {:?}", ctx.timeout))?
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The chat-completions request body shared by the latency and concurrency
|
||||||
|
/// scenarios — streamed, deterministic (temperature 0), usage included.
|
||||||
|
fn chat_payload(ctx: &RunCtx, prompt: &str) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"model": ctx.model_id,
|
||||||
|
"messages": [{"role": "user", "content": prompt}],
|
||||||
|
"max_tokens": ctx.max_tokens,
|
||||||
|
"temperature": 0,
|
||||||
|
"stream": true,
|
||||||
|
"stream_options": {"include_usage": true},
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Streamed single-request chat-completions latency probe — the batch-1
|
/// Streamed single-request chat-completions latency probe — the batch-1
|
||||||
@@ -118,15 +195,7 @@ impl Scenario for ChatLatencyScenario {
|
|||||||
|
|
||||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||||
let prompt = build_prompt(self.approx_prompt_tokens);
|
let prompt = build_prompt(self.approx_prompt_tokens);
|
||||||
let payload = json!({
|
let payload = chat_payload(ctx, &prompt);
|
||||||
"model": ctx.model_id,
|
|
||||||
"messages": [{"role": "user", "content": prompt}],
|
|
||||||
"max_tokens": ctx.max_tokens,
|
|
||||||
"temperature": 0,
|
|
||||||
"stream": true,
|
|
||||||
"stream_options": {"include_usage": true},
|
|
||||||
});
|
|
||||||
|
|
||||||
let fut = stream_and_measure(ctx, &payload);
|
let fut = stream_and_measure(ctx, &payload);
|
||||||
tokio::time::timeout(ctx.timeout, fut)
|
tokio::time::timeout(ctx.timeout, fut)
|
||||||
.await
|
.await
|
||||||
@@ -134,11 +203,182 @@ impl Scenario for ChatLatencyScenario {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Fan-out load probe: fire `concurrency` identical streams at once and
|
||||||
|
/// measure how the fleet behaves under simultaneous pressure (#89). This is
|
||||||
|
/// the only scenario that exercises the real a0/hermes/opencode pattern —
|
||||||
|
/// many agentic requests per user turn — which batch-1 single-request
|
||||||
|
/// timing cannot characterize. On a batch-1 serialized server, aggregate
|
||||||
|
/// throughput stays ~flat while TTFT/queue-wait inflate with `concurrency`;
|
||||||
|
/// that gap is the evidence for/against continuous batching.
|
||||||
|
pub struct ConcurrencyScenario {
|
||||||
|
id: String,
|
||||||
|
concurrency: u32,
|
||||||
|
approx_prompt_tokens: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Scenario for ConcurrencyScenario {
|
||||||
|
fn id(&self) -> &str {
|
||||||
|
&self.id
|
||||||
|
}
|
||||||
|
|
||||||
|
fn prompt_size(&self) -> u32 {
|
||||||
|
self.approx_prompt_tokens
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||||
|
let prompt = build_prompt(self.approx_prompt_tokens);
|
||||||
|
let payload = chat_payload(ctx, &prompt);
|
||||||
|
|
||||||
|
// Fire all streams at once; each is independently timed and capped by
|
||||||
|
// the per-request timeout so one hung stream can't stall the burst.
|
||||||
|
let burst_start = Instant::now();
|
||||||
|
let futs = (0..self.concurrency).map(|_| async {
|
||||||
|
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload)).await
|
||||||
|
});
|
||||||
|
let results = futures::future::join_all(futs).await;
|
||||||
|
let burst_window = burst_start.elapsed().as_secs_f64();
|
||||||
|
|
||||||
|
let mut streams: Vec<ScenarioMetrics> = Vec::new();
|
||||||
|
let mut rejected: u32 = 0;
|
||||||
|
for r in results {
|
||||||
|
match r {
|
||||||
|
Ok(Ok(m)) => streams.push(m),
|
||||||
|
// Admission backpressure (429/503) is shed load, counted
|
||||||
|
// separately from genuine failures/timeouts.
|
||||||
|
Ok(Err(e)) if is_admission_reject(&e) => rejected += 1,
|
||||||
|
Ok(Err(_)) | Err(_) => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if streams.is_empty() {
|
||||||
|
return Err(anyhow!(
|
||||||
|
"all {} concurrent streams failed ({rejected} shed by admission)",
|
||||||
|
self.concurrency
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let total_tokens: u64 = streams.iter().map(|m| m.completion_tokens).sum();
|
||||||
|
let ttfts: Vec<f64> = streams.iter().map(|m| m.ttft_s).collect();
|
||||||
|
// queue-wait ≈ TTFT − server prefill (#85); only for streams that
|
||||||
|
// reported helexa_timing.
|
||||||
|
let queue_waits: Vec<f64> = streams
|
||||||
|
.iter()
|
||||||
|
.filter_map(|m| {
|
||||||
|
m.prefill_ms
|
||||||
|
.map(|p| (m.ttft_s * 1000.0 - p as f64).max(0.0))
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
// Aggregate decode throughput across the whole node for the burst.
|
||||||
|
let aggregate_tps = if burst_window > 0.0 {
|
||||||
|
Some(total_tokens as f64 / burst_window)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(ScenarioMetrics {
|
||||||
|
ttft_s: median(&ttfts).unwrap_or(0.0),
|
||||||
|
decode_tps: aggregate_tps,
|
||||||
|
total_s: burst_window,
|
||||||
|
prompt_tokens: streams.iter().find_map(|m| m.prompt_tokens),
|
||||||
|
completion_tokens: total_tokens,
|
||||||
|
prefill_ms: None,
|
||||||
|
decode_ms: None,
|
||||||
|
prefill_tokens: None,
|
||||||
|
concurrency: Some(self.concurrency),
|
||||||
|
ttft_p95_s: percentile(&ttfts, 95.0),
|
||||||
|
queue_wait_ms_median: median(&queue_waits),
|
||||||
|
rejected: Some(rejected),
|
||||||
|
artifact: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Quality probe (#91): runs a fixed prompt and stores the full generated
|
||||||
|
/// text as an artifact for later scoring (manual now, LLM-judge later). The
|
||||||
|
/// point is to compare reasoning/planning quality across models — the axis
|
||||||
|
/// speed-only scenarios miss — so the frontier A/B (F3) picks on capability,
|
||||||
|
/// not just throughput.
|
||||||
|
pub struct CapabilityScenario {
|
||||||
|
id: String,
|
||||||
|
prompt: String,
|
||||||
|
max_tokens: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Scenario for CapabilityScenario {
|
||||||
|
fn id(&self) -> &str {
|
||||||
|
&self.id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Capability probes have no synthetic prompt-token target; the cell is
|
||||||
|
/// keyed by the scenario id alone.
|
||||||
|
fn prompt_size(&self) -> u32 {
|
||||||
|
0
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||||
|
let payload = json!({
|
||||||
|
"model": ctx.model_id,
|
||||||
|
"messages": [{"role": "user", "content": self.prompt}],
|
||||||
|
"max_tokens": self.max_tokens,
|
||||||
|
"temperature": 0,
|
||||||
|
"stream": true,
|
||||||
|
"stream_options": {"include_usage": true},
|
||||||
|
});
|
||||||
|
let fut = stream_and_measure_inner(ctx, &payload, true);
|
||||||
|
tokio::time::timeout(ctx.timeout, fut)
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow!("capability probe timed out after {:?}", ctx.timeout))?
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether a stream error was admission backpressure (HTTP 429/503) rather
|
||||||
|
/// than a genuine failure. `stream_and_measure` renders the upstream status
|
||||||
|
/// into the error string, so a substring check is sufficient.
|
||||||
|
fn is_admission_reject(e: &anyhow::Error) -> bool {
|
||||||
|
let s = e.to_string();
|
||||||
|
s.contains("429") || s.contains("503")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Median of a slice (sorted copy). `None` if empty.
|
||||||
|
fn median(values: &[f64]) -> Option<f64> {
|
||||||
|
if values.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let mut v = values.to_vec();
|
||||||
|
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
let lo = (v.len() - 1) / 2;
|
||||||
|
let hi = v.len() / 2;
|
||||||
|
Some((v[lo] + v[hi]) / 2.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Nearest-rank percentile of a slice (`p` in 0..=100). `None` if empty.
|
||||||
|
fn percentile(values: &[f64], p: f64) -> Option<f64> {
|
||||||
|
if values.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let mut v = values.to_vec();
|
||||||
|
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||||
|
Some(v[rank.clamp(1, v.len()) - 1])
|
||||||
|
}
|
||||||
|
|
||||||
/// The SSE-timing core, ported from `bench.py::one_run`. Kept free of the
|
/// The SSE-timing core, ported from `bench.py::one_run`. Kept free of the
|
||||||
/// `Scenario` trait so it's unit-testable against a mock byte stream.
|
/// `Scenario` trait so it's unit-testable against a mock byte stream.
|
||||||
async fn stream_and_measure(
|
async fn stream_and_measure(
|
||||||
ctx: &RunCtx<'_>,
|
ctx: &RunCtx<'_>,
|
||||||
payload: &serde_json::Value,
|
payload: &serde_json::Value,
|
||||||
|
) -> Result<ScenarioMetrics> {
|
||||||
|
stream_and_measure_inner(ctx, payload, false).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// As [`stream_and_measure`] but accumulates the full visible text when
|
||||||
|
/// `capture_text` is set — used by the capability probe (#91) to store the
|
||||||
|
/// generated artifact for later quality scoring.
|
||||||
|
async fn stream_and_measure_inner(
|
||||||
|
ctx: &RunCtx<'_>,
|
||||||
|
payload: &serde_json::Value,
|
||||||
|
capture_text: bool,
|
||||||
) -> Result<ScenarioMetrics> {
|
) -> Result<ScenarioMetrics> {
|
||||||
let start = Instant::now();
|
let start = Instant::now();
|
||||||
let resp = ctx
|
let resp = ctx
|
||||||
@@ -160,6 +400,10 @@ async fn stream_and_measure(
|
|||||||
let mut chunk_count: u64 = 0;
|
let mut chunk_count: u64 = 0;
|
||||||
let mut prompt_tokens: Option<u64> = None;
|
let mut prompt_tokens: Option<u64> = None;
|
||||||
let mut completion_tokens: Option<u64> = None;
|
let mut completion_tokens: Option<u64> = None;
|
||||||
|
let mut prefill_ms: Option<u64> = None;
|
||||||
|
let mut decode_ms: Option<u64> = None;
|
||||||
|
let mut prefill_tokens: Option<u64> = None;
|
||||||
|
let mut captured = String::new();
|
||||||
|
|
||||||
while let Some(event) = stream.next().await {
|
while let Some(event) = stream.next().await {
|
||||||
let event = event.context("reading SSE stream")?;
|
let event = event.context("reading SSE stream")?;
|
||||||
@@ -172,46 +416,86 @@ async fn stream_and_measure(
|
|||||||
Ok(c) => c,
|
Ok(c) => c,
|
||||||
Err(_) => continue, // tolerate non-JSON keepalive frames
|
Err(_) => continue, // tolerate non-JSON keepalive frames
|
||||||
};
|
};
|
||||||
if let Some(choice) = chunk.choices.first()
|
if let Some(choice) = chunk.choices.first() {
|
||||||
&& choice
|
// Liveness counts ANY generated delta (#117). Thinking
|
||||||
|
// models (Qwen3-Next-Thinking, Qwen3 with thinking on)
|
||||||
|
// stream `reasoning_content` first — sometimes for their
|
||||||
|
// entire budget — and a content-only view misread that as
|
||||||
|
// a dead stream ("no content chunks received") while also
|
||||||
|
// producing impossible client-side rates (reasoning-
|
||||||
|
// inclusive token counts over a visible-content-only
|
||||||
|
// window; observed: "244 tok/s" on a 3060). For
|
||||||
|
// non-thinking models the first delta IS content, so
|
||||||
|
// `ttft_s` semantics are unchanged for them.
|
||||||
|
let content = choice
|
||||||
.delta
|
.delta
|
||||||
.get("content")
|
.get("content")
|
||||||
.and_then(|c| c.as_str())
|
.and_then(|c| c.as_str())
|
||||||
.is_some_and(|s| !s.is_empty())
|
.filter(|c| !c.is_empty());
|
||||||
{
|
let reasoning = choice
|
||||||
if first.is_none() {
|
.delta
|
||||||
first = Some(now);
|
.get("reasoning_content")
|
||||||
|
.and_then(|c| c.as_str())
|
||||||
|
.filter(|c| !c.is_empty());
|
||||||
|
if content.is_some() || reasoning.is_some() {
|
||||||
|
if first.is_none() {
|
||||||
|
first = Some(now);
|
||||||
|
}
|
||||||
|
last = Some(now);
|
||||||
|
chunk_count += 1;
|
||||||
|
}
|
||||||
|
if capture_text && let Some(text) = content {
|
||||||
|
captured.push_str(text);
|
||||||
}
|
}
|
||||||
last = Some(now);
|
|
||||||
chunk_count += 1;
|
|
||||||
}
|
}
|
||||||
if let Some(usage) = chunk.usage {
|
if let Some(usage) = chunk.usage {
|
||||||
prompt_tokens = Some(usage.prompt_tokens);
|
prompt_tokens = Some(usage.prompt_tokens);
|
||||||
completion_tokens = Some(usage.completion_tokens);
|
completion_tokens = Some(usage.completion_tokens);
|
||||||
|
if let Some(t) = usage.helexa_timing {
|
||||||
|
prefill_ms = Some(t.prefill_ms);
|
||||||
|
decode_ms = Some(t.decode_ms);
|
||||||
|
prefill_tokens = Some(t.prefill_tokens);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let end = Instant::now();
|
let end = Instant::now();
|
||||||
|
|
||||||
let first = first.ok_or_else(|| anyhow!("no content chunks received"))?;
|
let first = first.ok_or_else(|| anyhow!("no generated chunks received"))?;
|
||||||
|
|
||||||
// neuron emits one SSE chunk per visible token, so chunk_count is an
|
// neuron emits one SSE chunk per generated token, so chunk_count is
|
||||||
// engine-truth count when no usage frame is sent.
|
// an engine-truth count when no usage frame is sent.
|
||||||
let tokens = completion_tokens.filter(|&t| t > 0).unwrap_or(chunk_count);
|
let tokens = completion_tokens.filter(|&t| t > 0).unwrap_or(chunk_count);
|
||||||
// decode rate is only meaningful over a real inter-chunk window.
|
// Decode rate: prefer the server-measured split (#85) — it counts
|
||||||
|
// every generated token over the actual decode window, immune to
|
||||||
|
// reasoning-suppression frame mismatches. Fall back to the client
|
||||||
|
// inter-chunk window with the CHUNK count (same frame) — never
|
||||||
|
// usage.completion_tokens over the chunk window, which mixes a
|
||||||
|
// reasoning-inclusive numerator with a visible-only denominator.
|
||||||
let window = last
|
let window = last
|
||||||
.filter(|&l| l > first)
|
.filter(|&l| l > first)
|
||||||
.map(|l| (l - first).as_secs_f64())
|
.map(|l| (l - first).as_secs_f64())
|
||||||
.unwrap_or(0.0);
|
.unwrap_or(0.0);
|
||||||
|
let decode_tps = match decode_ms {
|
||||||
|
Some(ms) if ms > 200 && tokens > 0 => Some(tokens as f64 / (ms as f64 / 1000.0)),
|
||||||
|
_ if window > 0.2 => Some(chunk_count as f64 / window),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
Ok(ScenarioMetrics {
|
Ok(ScenarioMetrics {
|
||||||
ttft_s: (first - start).as_secs_f64(),
|
ttft_s: (first - start).as_secs_f64(),
|
||||||
decode_tps: if window > 0.2 {
|
decode_tps,
|
||||||
Some(tokens as f64 / window)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
},
|
|
||||||
total_s: (end - start).as_secs_f64(),
|
total_s: (end - start).as_secs_f64(),
|
||||||
prompt_tokens,
|
prompt_tokens,
|
||||||
completion_tokens: tokens,
|
completion_tokens: tokens,
|
||||||
|
prefill_ms,
|
||||||
|
decode_ms,
|
||||||
|
prefill_tokens,
|
||||||
|
// Concurrency fields unset on the single-request path; the
|
||||||
|
// concurrency scenario builds its own aggregate (#89).
|
||||||
|
concurrency: None,
|
||||||
|
ttft_p95_s: None,
|
||||||
|
queue_wait_ms_median: None,
|
||||||
|
rejected: None,
|
||||||
|
artifact: if capture_text { Some(captured) } else { None },
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -229,6 +513,54 @@ mod tests {
|
|||||||
assert!(small.ends_with("/no_think"));
|
assert!(small.ends_with("/no_think"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn median_and_percentile_basics() {
|
||||||
|
assert_eq!(median(&[3.0, 1.0, 2.0]), Some(2.0));
|
||||||
|
assert_eq!(median(&[]), None);
|
||||||
|
let v = [1.0, 2.0, 3.0, 4.0, 5.0];
|
||||||
|
assert_eq!(percentile(&v, 50.0), Some(3.0));
|
||||||
|
assert_eq!(percentile(&v, 95.0), Some(5.0)); // nearest-rank → max with n=5
|
||||||
|
assert_eq!(percentile(&[], 95.0), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn admission_rejects_detected_by_status() {
|
||||||
|
assert!(is_admission_reject(&anyhow!(
|
||||||
|
"upstream returned 429 Too Many Requests"
|
||||||
|
)));
|
||||||
|
assert!(is_admission_reject(&anyhow!(
|
||||||
|
"upstream returned 503 Service Unavailable"
|
||||||
|
)));
|
||||||
|
assert!(!is_admission_reject(&anyhow!(
|
||||||
|
"upstream returned 500 Internal"
|
||||||
|
)));
|
||||||
|
assert!(!is_admission_reject(&anyhow!("connection refused")));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn concurrency_scenarios_built_from_config() {
|
||||||
|
use crate::config::{CapabilityProbe, ScenarioConfig};
|
||||||
|
let cfg = ScenarioConfig {
|
||||||
|
prompt_sizes: vec![128],
|
||||||
|
max_tokens: 64,
|
||||||
|
concurrency_levels: vec![2, 8],
|
||||||
|
concurrency_prompt_tokens: 512,
|
||||||
|
capability_probes: vec![CapabilityProbe {
|
||||||
|
name: "plan".into(),
|
||||||
|
prompt: "Write a plan.".into(),
|
||||||
|
max_tokens: 2048,
|
||||||
|
}],
|
||||||
|
};
|
||||||
|
let ids: Vec<String> = build_scenarios(&cfg)
|
||||||
|
.iter()
|
||||||
|
.map(|s| s.id().to_string())
|
||||||
|
.collect();
|
||||||
|
assert!(ids.contains(&"chat:128".to_string()));
|
||||||
|
assert!(ids.contains(&"concurrency:2".to_string()));
|
||||||
|
assert!(ids.contains(&"concurrency:8".to_string()));
|
||||||
|
assert!(ids.contains(&"capability:plan".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn prompt_floor_for_tiny_targets() {
|
fn prompt_floor_for_tiny_targets() {
|
||||||
// max(approx,16) floor means even 0 yields a non-trivial prompt.
|
// max(approx,16) floor means even 0 yields a non-trivial prompt.
|
||||||
|
|||||||
@@ -51,6 +51,35 @@ pub struct RunRecord {
|
|||||||
pub decode_tps: Option<f64>,
|
pub decode_tps: Option<f64>,
|
||||||
pub total_s: Option<f64>,
|
pub total_s: Option<f64>,
|
||||||
pub completion_tokens: Option<u64>,
|
pub completion_tokens: Option<u64>,
|
||||||
|
// server-measured prefill/decode split (#85), null on engines/paths
|
||||||
|
// that don't emit `usage.helexa_timing`.
|
||||||
|
pub prefill_ms: Option<u64>,
|
||||||
|
pub decode_ms: Option<u64>,
|
||||||
|
pub prefill_tokens: Option<u64>,
|
||||||
|
// GPU telemetry sampled from /health around the run (#87), null for
|
||||||
|
// non-neuron targets or when /health was unreachable. vram_used_mb is
|
||||||
|
// the node sum; util/temp are the hottest single device.
|
||||||
|
pub vram_used_mb: Option<u64>,
|
||||||
|
pub gpu_util_pct: Option<u32>,
|
||||||
|
pub gpu_temp_c: Option<u32>,
|
||||||
|
// concurrency / agentic-load burst metrics (#89), null for single-request
|
||||||
|
// scenarios. For a burst, ttft_s/decode_tps/total_s carry the aggregate.
|
||||||
|
pub concurrency: Option<u32>,
|
||||||
|
pub ttft_p95_s: Option<f64>,
|
||||||
|
pub queue_wait_ms: Option<f64>,
|
||||||
|
pub rejected: Option<u32>,
|
||||||
|
// cold-load / model-swap cost (#90), set only by the deliberate
|
||||||
|
// `swap-cost` measurement (scenario_id = "swap"). The other metric
|
||||||
|
// fields carry the cold first-request after reload.
|
||||||
|
pub swap_unload_ms: Option<u64>,
|
||||||
|
pub swap_load_ms: Option<u64>,
|
||||||
|
// capability probe (#91): the full generated text (scenario_id =
|
||||||
|
// "capability:<name>"), scored for quality later. `quality_score` /
|
||||||
|
// `scorer` are null at insert — set by `score` (manual) or a future
|
||||||
|
// LLM-judge.
|
||||||
|
pub artifact: Option<String>,
|
||||||
|
pub quality_score: Option<f64>,
|
||||||
|
pub scorer: Option<String>,
|
||||||
// outcome
|
// outcome
|
||||||
pub ok: bool,
|
pub ok: bool,
|
||||||
pub error: Option<String>,
|
pub error: Option<String>,
|
||||||
@@ -123,6 +152,21 @@ impl Store {
|
|||||||
decode_tps REAL,
|
decode_tps REAL,
|
||||||
total_s REAL,
|
total_s REAL,
|
||||||
completion_tokens INTEGER,
|
completion_tokens INTEGER,
|
||||||
|
prefill_ms INTEGER,
|
||||||
|
decode_ms INTEGER,
|
||||||
|
prefill_tokens INTEGER,
|
||||||
|
vram_used_mb INTEGER,
|
||||||
|
gpu_util_pct INTEGER,
|
||||||
|
gpu_temp_c INTEGER,
|
||||||
|
concurrency INTEGER,
|
||||||
|
ttft_p95_s REAL,
|
||||||
|
queue_wait_ms REAL,
|
||||||
|
rejected INTEGER,
|
||||||
|
swap_unload_ms INTEGER,
|
||||||
|
swap_load_ms INTEGER,
|
||||||
|
artifact TEXT,
|
||||||
|
quality_score REAL,
|
||||||
|
scorer TEXT,
|
||||||
ok INTEGER NOT NULL,
|
ok INTEGER NOT NULL,
|
||||||
error TEXT
|
error TEXT
|
||||||
);
|
);
|
||||||
@@ -133,6 +177,51 @@ impl Store {
|
|||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.context("initialising sqlite schema")?;
|
.context("initialising sqlite schema")?;
|
||||||
|
// Additive migrations for DBs created before a column existed.
|
||||||
|
// `CREATE TABLE IF NOT EXISTS` above only seeds fresh DBs; existing
|
||||||
|
// ones need the columns backfilled (as NULL) so older rows coexist
|
||||||
|
// with new metrics. There is no migration framework — each entry is
|
||||||
|
// an idempotent "add if missing".
|
||||||
|
Self::ensure_columns(
|
||||||
|
conn,
|
||||||
|
"runs",
|
||||||
|
&[
|
||||||
|
("prefill_ms", "INTEGER"),
|
||||||
|
("decode_ms", "INTEGER"),
|
||||||
|
("prefill_tokens", "INTEGER"),
|
||||||
|
("vram_used_mb", "INTEGER"),
|
||||||
|
("gpu_util_pct", "INTEGER"),
|
||||||
|
("gpu_temp_c", "INTEGER"),
|
||||||
|
("concurrency", "INTEGER"),
|
||||||
|
("ttft_p95_s", "REAL"),
|
||||||
|
("queue_wait_ms", "REAL"),
|
||||||
|
("rejected", "INTEGER"),
|
||||||
|
("swap_unload_ms", "INTEGER"),
|
||||||
|
("swap_load_ms", "INTEGER"),
|
||||||
|
("artifact", "TEXT"),
|
||||||
|
("quality_score", "REAL"),
|
||||||
|
("scorer", "TEXT"),
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add any of `columns` that the table is missing (`ALTER TABLE ADD
|
||||||
|
/// COLUMN`). Idempotent: existing columns are read from
|
||||||
|
/// `PRAGMA table_info` and skipped, so this is safe to run on every open.
|
||||||
|
fn ensure_columns(conn: &Connection, table: &str, columns: &[(&str, &str)]) -> Result<()> {
|
||||||
|
let mut existing = std::collections::HashSet::new();
|
||||||
|
let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
|
||||||
|
let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
|
||||||
|
for name in names {
|
||||||
|
existing.insert(name?);
|
||||||
|
}
|
||||||
|
for (name, ty) in columns {
|
||||||
|
if !existing.contains(*name) {
|
||||||
|
conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {name} {ty};"))
|
||||||
|
.with_context(|| format!("adding column {table}.{name}"))?;
|
||||||
|
}
|
||||||
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -166,6 +255,11 @@ impl Store {
|
|||||||
model_id, harness, capabilities_json, devices_json,
|
model_id, harness, capabilities_json, devices_json,
|
||||||
scenario_id, prompt_size_approx, prompt_tokens_actual, max_tokens,
|
scenario_id, prompt_size_approx, prompt_tokens_actual, max_tokens,
|
||||||
ttft_s, decode_tps, total_s, completion_tokens,
|
ttft_s, decode_tps, total_s, completion_tokens,
|
||||||
|
prefill_ms, decode_ms, prefill_tokens,
|
||||||
|
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||||
|
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||||
|
swap_unload_ms, swap_load_ms,
|
||||||
|
artifact, quality_score, scorer,
|
||||||
ok, error
|
ok, error
|
||||||
) VALUES (
|
) VALUES (
|
||||||
?1, ?2, ?3, ?4,
|
?1, ?2, ?3, ?4,
|
||||||
@@ -176,7 +270,12 @@ impl Store {
|
|||||||
?20, ?21, ?22, ?23,
|
?20, ?21, ?22, ?23,
|
||||||
?24, ?25, ?26, ?27,
|
?24, ?25, ?26, ?27,
|
||||||
?28, ?29, ?30, ?31,
|
?28, ?29, ?30, ?31,
|
||||||
?32, ?33
|
?32, ?33, ?34,
|
||||||
|
?35, ?36, ?37,
|
||||||
|
?38, ?39, ?40, ?41,
|
||||||
|
?42, ?43,
|
||||||
|
?44, ?45, ?46,
|
||||||
|
?47, ?48
|
||||||
)",
|
)",
|
||||||
params![
|
params![
|
||||||
r.ts,
|
r.ts,
|
||||||
@@ -210,6 +309,21 @@ impl Store {
|
|||||||
r.decode_tps,
|
r.decode_tps,
|
||||||
r.total_s,
|
r.total_s,
|
||||||
r.completion_tokens,
|
r.completion_tokens,
|
||||||
|
r.prefill_ms,
|
||||||
|
r.decode_ms,
|
||||||
|
r.prefill_tokens,
|
||||||
|
r.vram_used_mb,
|
||||||
|
r.gpu_util_pct,
|
||||||
|
r.gpu_temp_c,
|
||||||
|
r.concurrency,
|
||||||
|
r.ttft_p95_s,
|
||||||
|
r.queue_wait_ms,
|
||||||
|
r.rejected,
|
||||||
|
r.swap_unload_ms,
|
||||||
|
r.swap_load_ms,
|
||||||
|
r.artifact,
|
||||||
|
r.quality_score,
|
||||||
|
r.scorer,
|
||||||
r.ok as i64,
|
r.ok as i64,
|
||||||
r.error,
|
r.error,
|
||||||
],
|
],
|
||||||
@@ -224,7 +338,10 @@ impl Store {
|
|||||||
// successful run, then median that SHA's samples.
|
// successful run, then median that SHA's samples.
|
||||||
let mut stmt = self.conn.prepare(
|
let mut stmt = self.conn.prepare(
|
||||||
"SELECT target_name, model_id, scenario_id, prompt_size_approx, git_sha,
|
"SELECT target_name, model_id, scenario_id, prompt_size_approx, git_sha,
|
||||||
ttft_s, decode_tps, total_s, prompt_tokens_actual, gpus_json
|
ttft_s, decode_tps, total_s, prompt_tokens_actual, gpus_json,
|
||||||
|
prefill_ms, decode_ms, prefill_tokens,
|
||||||
|
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||||
|
concurrency, ttft_p95_s, queue_wait_ms, rejected
|
||||||
FROM runs
|
FROM runs
|
||||||
WHERE ok=1
|
WHERE ok=1
|
||||||
ORDER BY target_name, model_id, scenario_id, id",
|
ORDER BY target_name, model_id, scenario_id, id",
|
||||||
@@ -241,12 +358,188 @@ impl Store {
|
|||||||
total_s: row.get(7)?,
|
total_s: row.get(7)?,
|
||||||
prompt_tokens_actual: row.get(8)?,
|
prompt_tokens_actual: row.get(8)?,
|
||||||
gpus_json: row.get(9)?,
|
gpus_json: row.get(9)?,
|
||||||
|
prefill_ms: row.get(10)?,
|
||||||
|
decode_ms: row.get(11)?,
|
||||||
|
prefill_tokens: row.get(12)?,
|
||||||
|
vram_used_mb: row.get(13)?,
|
||||||
|
gpu_util_pct: row.get(14)?,
|
||||||
|
gpu_temp_c: row.get(15)?,
|
||||||
|
concurrency: row.get(16)?,
|
||||||
|
ttft_p95_s: row.get(17)?,
|
||||||
|
queue_wait_ms: row.get(18)?,
|
||||||
|
rejected: row.get(19)?,
|
||||||
})
|
})
|
||||||
})?;
|
})?;
|
||||||
let raws: Vec<RawRow> = rows.collect::<rusqlite::Result<_>>()?;
|
let raws: Vec<RawRow> = rows.collect::<rusqlite::Result<_>>()?;
|
||||||
Ok(aggregate(raws))
|
Ok(aggregate(raws))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Context-length scaling curves (#88): per (target, model), the latest
|
||||||
|
/// build's `chat:<n>` cells pivoted by prompt size into prefill & decode
|
||||||
|
/// tok/s vs context. The headline is `decode_flatness` — decode tok/s at
|
||||||
|
/// the largest context divided by the smallest. Near 1.0 confirms the
|
||||||
|
/// Gated-DeltaNet O(1)-in-sequence-length decode; a sharp drop locates
|
||||||
|
/// where the model stops scaling for free.
|
||||||
|
pub fn scaling(&self) -> Result<Vec<ScalingCurve>> {
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
// Reuse the already-aggregated report cells; the chat:<n> rows are
|
||||||
|
// exactly the per-context measurement points.
|
||||||
|
let mut by_model: BTreeMap<(String, String), Vec<ReportRow>> = BTreeMap::new();
|
||||||
|
for r in self.report_rows()? {
|
||||||
|
if r.scenario_id.starts_with("chat:") {
|
||||||
|
by_model
|
||||||
|
.entry((r.target_name.clone(), r.model_id.clone()))
|
||||||
|
.or_default()
|
||||||
|
.push(r);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut out = Vec::new();
|
||||||
|
for ((target_name, model_id), mut rows) in by_model {
|
||||||
|
rows.sort_by_key(|r| r.prompt_size_approx);
|
||||||
|
let points: Vec<ScalingPoint> = rows
|
||||||
|
.iter()
|
||||||
|
.map(|r| ScalingPoint {
|
||||||
|
prompt_size: r.prompt_size_approx,
|
||||||
|
prompt_tokens: r.prompt_tokens,
|
||||||
|
prefill_tps: r.prefill_tps_median,
|
||||||
|
decode_tps: r.decode_tps_median,
|
||||||
|
samples: r.samples,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
// Flatness across the smallest→largest points that both have a
|
||||||
|
// decode rate (skips cells where the decode window was too short).
|
||||||
|
let with_decode: Vec<&ScalingPoint> =
|
||||||
|
points.iter().filter(|p| p.decode_tps.is_some()).collect();
|
||||||
|
let decode_flatness = match (with_decode.first(), with_decode.last()) {
|
||||||
|
(Some(lo), Some(hi)) if with_decode.len() >= 2 => {
|
||||||
|
match (lo.decode_tps, hi.decode_tps) {
|
||||||
|
(Some(a), Some(b)) if a > 0.0 => Some(b / a),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
out.push(ScalingCurve {
|
||||||
|
target_name,
|
||||||
|
model_id,
|
||||||
|
git_sha: rows.first().map(|r| r.git_sha.clone()).unwrap_or_default(),
|
||||||
|
gpu: rows.iter().find_map(|r| r.gpu.clone()),
|
||||||
|
points,
|
||||||
|
decode_flatness,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cold-load / model-swap costs (#90): per (target, model) at the latest
|
||||||
|
/// build, the median unload→reload latency and the cold first-request
|
||||||
|
/// latency after reload (the `scenario_id = "swap"` rows).
|
||||||
|
pub fn swap_costs(&self) -> Result<Vec<SwapCost>> {
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
let mut stmt = self.conn.prepare(
|
||||||
|
"SELECT target_name, model_id, git_sha, gpus_json,
|
||||||
|
swap_unload_ms, swap_load_ms, ttft_s, total_s
|
||||||
|
FROM runs
|
||||||
|
WHERE ok=1 AND scenario_id='swap'
|
||||||
|
ORDER BY target_name, model_id, id",
|
||||||
|
)?;
|
||||||
|
struct Raw {
|
||||||
|
target: String,
|
||||||
|
model: String,
|
||||||
|
sha: String,
|
||||||
|
gpus_json: Option<String>,
|
||||||
|
unload_ms: Option<f64>,
|
||||||
|
load_ms: Option<f64>,
|
||||||
|
ttft_s: Option<f64>,
|
||||||
|
total_s: Option<f64>,
|
||||||
|
}
|
||||||
|
let raws: Vec<Raw> = stmt
|
||||||
|
.query_map([], |r| {
|
||||||
|
Ok(Raw {
|
||||||
|
target: r.get(0)?,
|
||||||
|
model: r.get(1)?,
|
||||||
|
sha: r.get(2)?,
|
||||||
|
gpus_json: r.get(3)?,
|
||||||
|
unload_ms: r.get::<_, Option<i64>>(4)?.map(|v| v as f64),
|
||||||
|
load_ms: r.get::<_, Option<i64>>(5)?.map(|v| v as f64),
|
||||||
|
ttft_s: r.get(6)?,
|
||||||
|
total_s: r.get(7)?,
|
||||||
|
})
|
||||||
|
})?
|
||||||
|
.collect::<rusqlite::Result<_>>()?;
|
||||||
|
let mut by: BTreeMap<(String, String), Vec<Raw>> = BTreeMap::new();
|
||||||
|
for r in raws {
|
||||||
|
by.entry((r.target.clone(), r.model.clone()))
|
||||||
|
.or_default()
|
||||||
|
.push(r);
|
||||||
|
}
|
||||||
|
let mut out = Vec::new();
|
||||||
|
for ((target, model), rows) in by {
|
||||||
|
let latest = rows.last().map(|r| r.sha.clone()).unwrap_or_default();
|
||||||
|
let cell: Vec<&Raw> = rows.iter().filter(|r| r.sha == latest).collect();
|
||||||
|
out.push(SwapCost {
|
||||||
|
target_name: target,
|
||||||
|
model_id: model,
|
||||||
|
git_sha: latest,
|
||||||
|
gpu: cell
|
||||||
|
.iter()
|
||||||
|
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_label)),
|
||||||
|
unload_ms_median: median(cell.iter().filter_map(|r| r.unload_ms)),
|
||||||
|
load_ms_median: median(cell.iter().filter_map(|r| r.load_ms)),
|
||||||
|
cold_ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
||||||
|
cold_total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
||||||
|
samples: cell.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Capability-probe runs (#91): the stored output artifacts and their
|
||||||
|
/// quality scores. `unscored_only` filters to runs awaiting a score —
|
||||||
|
/// the worklist for manual scoring or a future LLM-judge.
|
||||||
|
pub fn capability_runs(&self, unscored_only: bool) -> Result<Vec<CapabilityRun>> {
|
||||||
|
let sql = format!(
|
||||||
|
"SELECT id, ts, target_name, model_id, scenario_id, git_sha,
|
||||||
|
quality_score, scorer, artifact
|
||||||
|
FROM runs
|
||||||
|
WHERE ok=1 AND scenario_id LIKE 'capability:%'{}
|
||||||
|
ORDER BY id DESC",
|
||||||
|
if unscored_only {
|
||||||
|
" AND quality_score IS NULL"
|
||||||
|
} else {
|
||||||
|
""
|
||||||
|
},
|
||||||
|
);
|
||||||
|
let mut stmt = self.conn.prepare(&sql)?;
|
||||||
|
let rows = stmt
|
||||||
|
.query_map([], |r| {
|
||||||
|
Ok(CapabilityRun {
|
||||||
|
id: r.get(0)?,
|
||||||
|
ts: r.get(1)?,
|
||||||
|
target_name: r.get(2)?,
|
||||||
|
model_id: r.get(3)?,
|
||||||
|
scenario_id: r.get(4)?,
|
||||||
|
git_sha: r.get(5)?,
|
||||||
|
quality_score: r.get(6)?,
|
||||||
|
scorer: r.get(7)?,
|
||||||
|
artifact: r.get(8)?,
|
||||||
|
})
|
||||||
|
})?
|
||||||
|
.collect::<rusqlite::Result<_>>()?;
|
||||||
|
Ok(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Attach a quality score to one capability run (#91). `scorer` records
|
||||||
|
/// who/what scored it ("manual", "llm:claude-…"). Returns the row count
|
||||||
|
/// updated (0 if the id doesn't exist).
|
||||||
|
pub fn set_score(&self, run_id: i64, score: f64, scorer: &str) -> Result<usize> {
|
||||||
|
let n = self.conn.execute(
|
||||||
|
"UPDATE runs SET quality_score=?1, scorer=?2 WHERE id=?3",
|
||||||
|
params![score, scorer, run_id],
|
||||||
|
)?;
|
||||||
|
Ok(n)
|
||||||
|
}
|
||||||
|
|
||||||
// ── Read API surface (consumed by api.rs) ─────────────────────────
|
// ── Read API surface (consumed by api.rs) ─────────────────────────
|
||||||
|
|
||||||
/// Total recorded runs (for `/api/health`).
|
/// Total recorded runs (for `/api/health`).
|
||||||
@@ -379,7 +672,10 @@ impl Store {
|
|||||||
"SELECT id, ts, target_name, hostname, git_sha, build_timestamp, package_version,
|
"SELECT id, ts, target_name, hostname, git_sha, build_timestamp, package_version,
|
||||||
model_id, harness, scenario_id, prompt_size_approx, prompt_tokens_actual,
|
model_id, harness, scenario_id, prompt_size_approx, prompt_tokens_actual,
|
||||||
max_tokens, ttft_s, decode_tps, total_s, completion_tokens, ok, error,
|
max_tokens, ttft_s, decode_tps, total_s, completion_tokens, ok, error,
|
||||||
gpus_json
|
gpus_json, prefill_ms, decode_ms, prefill_tokens,
|
||||||
|
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||||
|
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||||
|
swap_unload_ms, swap_load_ms
|
||||||
FROM runs",
|
FROM runs",
|
||||||
);
|
);
|
||||||
let mut conds: Vec<String> = Vec::new();
|
let mut conds: Vec<String> = Vec::new();
|
||||||
@@ -435,6 +731,18 @@ impl Store {
|
|||||||
completion_tokens: r.get(16)?,
|
completion_tokens: r.get(16)?,
|
||||||
ok: r.get::<_, i64>(17)? != 0,
|
ok: r.get::<_, i64>(17)? != 0,
|
||||||
error: r.get(18)?,
|
error: r.get(18)?,
|
||||||
|
prefill_ms: r.get(20)?,
|
||||||
|
decode_ms: r.get(21)?,
|
||||||
|
prefill_tokens: r.get(22)?,
|
||||||
|
vram_used_mb: r.get(23)?,
|
||||||
|
gpu_util_pct: r.get(24)?,
|
||||||
|
gpu_temp_c: r.get(25)?,
|
||||||
|
concurrency: r.get(26)?,
|
||||||
|
ttft_p95_s: r.get(27)?,
|
||||||
|
queue_wait_ms: r.get(28)?,
|
||||||
|
rejected: r.get(29)?,
|
||||||
|
swap_unload_ms: r.get(30)?,
|
||||||
|
swap_load_ms: r.get(31)?,
|
||||||
})
|
})
|
||||||
})?
|
})?
|
||||||
.collect::<rusqlite::Result<_>>()?;
|
.collect::<rusqlite::Result<_>>()?;
|
||||||
@@ -554,6 +862,18 @@ pub struct RunRow {
|
|||||||
pub decode_tps: Option<f64>,
|
pub decode_tps: Option<f64>,
|
||||||
pub total_s: Option<f64>,
|
pub total_s: Option<f64>,
|
||||||
pub completion_tokens: Option<u64>,
|
pub completion_tokens: Option<u64>,
|
||||||
|
pub prefill_ms: Option<u64>,
|
||||||
|
pub decode_ms: Option<u64>,
|
||||||
|
pub prefill_tokens: Option<u64>,
|
||||||
|
pub vram_used_mb: Option<u64>,
|
||||||
|
pub gpu_util_pct: Option<u64>,
|
||||||
|
pub gpu_temp_c: Option<u64>,
|
||||||
|
pub concurrency: Option<u64>,
|
||||||
|
pub ttft_p95_s: Option<f64>,
|
||||||
|
pub queue_wait_ms: Option<f64>,
|
||||||
|
pub rejected: Option<u64>,
|
||||||
|
pub swap_unload_ms: Option<u64>,
|
||||||
|
pub swap_load_ms: Option<u64>,
|
||||||
pub ok: bool,
|
pub ok: bool,
|
||||||
pub error: Option<String>,
|
pub error: Option<String>,
|
||||||
}
|
}
|
||||||
@@ -569,6 +889,16 @@ struct RawRow {
|
|||||||
total_s: Option<f64>,
|
total_s: Option<f64>,
|
||||||
prompt_tokens_actual: Option<u64>,
|
prompt_tokens_actual: Option<u64>,
|
||||||
gpus_json: Option<String>,
|
gpus_json: Option<String>,
|
||||||
|
prefill_ms: Option<u64>,
|
||||||
|
decode_ms: Option<u64>,
|
||||||
|
prefill_tokens: Option<u64>,
|
||||||
|
vram_used_mb: Option<u64>,
|
||||||
|
gpu_util_pct: Option<u64>,
|
||||||
|
gpu_temp_c: Option<u64>,
|
||||||
|
concurrency: Option<u64>,
|
||||||
|
ttft_p95_s: Option<f64>,
|
||||||
|
queue_wait_ms: Option<f64>,
|
||||||
|
rejected: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// An aggregated cell ready for the report table.
|
/// An aggregated cell ready for the report table.
|
||||||
@@ -583,11 +913,100 @@ pub struct ReportRow {
|
|||||||
pub ttft_s_median: Option<f64>,
|
pub ttft_s_median: Option<f64>,
|
||||||
pub decode_tps_median: Option<f64>,
|
pub decode_tps_median: Option<f64>,
|
||||||
pub total_s_median: Option<f64>,
|
pub total_s_median: Option<f64>,
|
||||||
|
/// Latency tail percentiles — where batch-1 pain actually shows up, and
|
||||||
|
/// invisible behind a bare median. p95/p99 nearest-rank; with few
|
||||||
|
/// samples they collapse toward the max (honest, not interpolated).
|
||||||
|
pub ttft_s_p95: Option<f64>,
|
||||||
|
pub ttft_s_p99: Option<f64>,
|
||||||
|
pub total_s_p95: Option<f64>,
|
||||||
|
pub total_s_p99: Option<f64>,
|
||||||
|
/// Server-measured prefill/decode split (#85). `prefill_tps_median` is
|
||||||
|
/// the true prompt-encoding rate (prefill_tokens / prefill_ms),
|
||||||
|
/// complementing `decode_tps_median` (the generation rate).
|
||||||
|
pub prefill_ms_median: Option<f64>,
|
||||||
|
pub decode_ms_median: Option<f64>,
|
||||||
|
pub prefill_tps_median: Option<f64>,
|
||||||
|
/// GPU telemetry sampled from /health around the run (#87). `vram_used_mb`
|
||||||
|
/// is the node sum; `vram_total_mb` (from discovery) lets the report show
|
||||||
|
/// real headroom — the "2/3 used" hunch as a number. util/temp are the
|
||||||
|
/// hottest device. All `None` for non-neuron targets.
|
||||||
|
pub vram_used_mb_median: Option<f64>,
|
||||||
|
pub vram_total_mb: Option<u64>,
|
||||||
|
pub gpu_util_pct_median: Option<f64>,
|
||||||
|
pub gpu_temp_c_median: Option<f64>,
|
||||||
|
/// Concurrency / agentic-load burst metrics (#89). `concurrency` is the
|
||||||
|
/// burst width (constant per cell). `ttft_p95_load_s` is the within-burst
|
||||||
|
/// TTFT tail; `queue_wait_ms_median` the admission wait; `rejected_median`
|
||||||
|
/// the per-burst shed count. All `None` for non-concurrency scenarios.
|
||||||
|
pub concurrency: Option<u64>,
|
||||||
|
pub ttft_p95_load_s: Option<f64>,
|
||||||
|
pub queue_wait_ms_median: Option<f64>,
|
||||||
|
pub rejected_median: Option<f64>,
|
||||||
pub samples: usize,
|
pub samples: usize,
|
||||||
/// Public-facing resource name (the host's GPU(s)), e.g. "2× RTX 5090".
|
/// Public-facing resource name (the host's GPU(s)), e.g. "2× RTX 5090".
|
||||||
pub gpu: Option<String>,
|
pub gpu: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// One context-length scaling curve for a (target, model) at its latest
|
||||||
|
/// build — the points ordered by prompt size, plus the decode-flatness
|
||||||
|
/// summary (#88).
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct ScalingCurve {
|
||||||
|
pub target_name: String,
|
||||||
|
pub model_id: String,
|
||||||
|
pub git_sha: String,
|
||||||
|
pub gpu: Option<String>,
|
||||||
|
pub points: Vec<ScalingPoint>,
|
||||||
|
/// decode tok/s at the largest context ÷ at the smallest. ~1.0 = flat
|
||||||
|
/// (GDN O(1) decode); <1 quantifies the drop-off. `None` with <2 points.
|
||||||
|
pub decode_flatness: Option<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cold-load / model-swap cost for a (target, model) at its latest build
|
||||||
|
/// (#90): the reload latency and the cold first-request after it.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct SwapCost {
|
||||||
|
pub target_name: String,
|
||||||
|
pub model_id: String,
|
||||||
|
pub git_sha: String,
|
||||||
|
pub gpu: Option<String>,
|
||||||
|
/// Time to free the model (`POST /models/unload`), ms.
|
||||||
|
pub unload_ms_median: Option<f64>,
|
||||||
|
/// Time to reload it (`POST /models/load`, synchronous), ms — the
|
||||||
|
/// headline swap cost feeding the vision cold-swap policy (F4e).
|
||||||
|
pub load_ms_median: Option<f64>,
|
||||||
|
/// TTFT of the first request after reload (cold caches), seconds.
|
||||||
|
pub cold_ttft_s_median: Option<f64>,
|
||||||
|
/// Total wall-clock of that cold first request, seconds.
|
||||||
|
pub cold_total_s_median: Option<f64>,
|
||||||
|
pub samples: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One capability-probe run (#91): the stored output artifact plus its
|
||||||
|
/// quality score (null until scored manually or by a future LLM-judge).
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct CapabilityRun {
|
||||||
|
pub id: i64,
|
||||||
|
pub ts: String,
|
||||||
|
pub target_name: String,
|
||||||
|
pub model_id: String,
|
||||||
|
pub scenario_id: String,
|
||||||
|
pub git_sha: String,
|
||||||
|
pub quality_score: Option<f64>,
|
||||||
|
pub scorer: Option<String>,
|
||||||
|
pub artifact: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One point on a [`ScalingCurve`]: the throughput at a given context size.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
|
pub struct ScalingPoint {
|
||||||
|
pub prompt_size: u32,
|
||||||
|
pub prompt_tokens: Option<u64>,
|
||||||
|
pub prefill_tps: Option<f64>,
|
||||||
|
pub decode_tps: Option<f64>,
|
||||||
|
pub samples: usize,
|
||||||
|
}
|
||||||
|
|
||||||
/// Group by (target, model, scenario), keep only the latest SHA's rows
|
/// Group by (target, model, scenario), keep only the latest SHA's rows
|
||||||
/// (latest = the SHA of the last-inserted row, since input is id-ordered),
|
/// (latest = the SHA of the last-inserted row, since input is id-ordered),
|
||||||
/// and median each metric.
|
/// and median each metric.
|
||||||
@@ -611,6 +1030,11 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
|||||||
let latest_sha = rows.last().map(|r| r.git_sha.clone()).unwrap_or_default();
|
let latest_sha = rows.last().map(|r| r.git_sha.clone()).unwrap_or_default();
|
||||||
let cell: Vec<&RawRow> = rows.iter().filter(|r| r.git_sha == latest_sha).collect();
|
let cell: Vec<&RawRow> = rows.iter().filter(|r| r.git_sha == latest_sha).collect();
|
||||||
let prompt_size_approx = cell.first().map(|r| r.prompt_size_approx).unwrap_or(0);
|
let prompt_size_approx = cell.first().map(|r| r.prompt_size_approx).unwrap_or(0);
|
||||||
|
// Per-row prefill tok/s, derived from the server-measured split.
|
||||||
|
let prefill_tps = |r: &&RawRow| match (r.prefill_tokens, r.prefill_ms) {
|
||||||
|
(Some(tok), Some(ms)) if ms > 0 => Some(tok as f64 * 1000.0 / ms as f64),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
out.push(ReportRow {
|
out.push(ReportRow {
|
||||||
target_name,
|
target_name,
|
||||||
model_id,
|
model_id,
|
||||||
@@ -621,6 +1045,27 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
|||||||
ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
||||||
decode_tps_median: median(cell.iter().filter_map(|r| r.decode_tps)),
|
decode_tps_median: median(cell.iter().filter_map(|r| r.decode_tps)),
|
||||||
total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
||||||
|
ttft_s_p95: percentile(cell.iter().filter_map(|r| r.ttft_s), 95.0),
|
||||||
|
ttft_s_p99: percentile(cell.iter().filter_map(|r| r.ttft_s), 99.0),
|
||||||
|
total_s_p95: percentile(cell.iter().filter_map(|r| r.total_s), 95.0),
|
||||||
|
total_s_p99: percentile(cell.iter().filter_map(|r| r.total_s), 99.0),
|
||||||
|
prefill_ms_median: median(cell.iter().filter_map(|r| r.prefill_ms.map(|m| m as f64))),
|
||||||
|
decode_ms_median: median(cell.iter().filter_map(|r| r.decode_ms.map(|m| m as f64))),
|
||||||
|
prefill_tps_median: median(cell.iter().filter_map(prefill_tps)),
|
||||||
|
vram_used_mb_median: median(
|
||||||
|
cell.iter().filter_map(|r| r.vram_used_mb.map(|v| v as f64)),
|
||||||
|
),
|
||||||
|
vram_total_mb: cell
|
||||||
|
.iter()
|
||||||
|
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_total_vram_mb)),
|
||||||
|
gpu_util_pct_median: median(
|
||||||
|
cell.iter().filter_map(|r| r.gpu_util_pct.map(|v| v as f64)),
|
||||||
|
),
|
||||||
|
gpu_temp_c_median: median(cell.iter().filter_map(|r| r.gpu_temp_c.map(|v| v as f64))),
|
||||||
|
concurrency: cell.iter().find_map(|r| r.concurrency),
|
||||||
|
ttft_p95_load_s: median(cell.iter().filter_map(|r| r.ttft_p95_s)),
|
||||||
|
queue_wait_ms_median: median(cell.iter().filter_map(|r| r.queue_wait_ms)),
|
||||||
|
rejected_median: median(cell.iter().filter_map(|r| r.rejected.map(|v| v as f64))),
|
||||||
samples: cell.len(),
|
samples: cell.len(),
|
||||||
gpu: cell
|
gpu: cell
|
||||||
.iter()
|
.iter()
|
||||||
@@ -630,6 +1075,19 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
|||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Node total VRAM in MB, summed across the devices in a run's stored
|
||||||
|
/// `gpus_json` (the discovery `DeviceInfo` list, each with `vram_total_mb`).
|
||||||
|
/// Pairs with the sampled `vram_used_mb` to report real headroom (#87).
|
||||||
|
/// `None` when empty/absent or no device declares a total.
|
||||||
|
fn gpu_total_vram_mb(gpus_json: &str) -> Option<u64> {
|
||||||
|
let devices: Vec<serde_json::Value> = serde_json::from_str(gpus_json).ok()?;
|
||||||
|
let total: u64 = devices
|
||||||
|
.iter()
|
||||||
|
.filter_map(|d| d.get("vram_total_mb").and_then(|v| v.as_u64()))
|
||||||
|
.sum();
|
||||||
|
(total > 0).then_some(total)
|
||||||
|
}
|
||||||
|
|
||||||
/// Compact GPU label from a run's stored `gpus_json` (the discovery device
|
/// Compact GPU label from a run's stored `gpus_json` (the discovery device
|
||||||
/// list) — e.g. "2× RTX 5090", "RTX 4090". `None` when empty/absent. Used
|
/// list) — e.g. "2× RTX 5090", "RTX 4090". `None` when empty/absent. Used
|
||||||
/// as the public-facing resource name in place of internal hostnames.
|
/// as the public-facing resource name in place of internal hostnames.
|
||||||
@@ -680,6 +1138,22 @@ fn median(values: impl Iterator<Item = f64>) -> Option<f64> {
|
|||||||
Some((v[lo] + v[hi]) / 2.0)
|
Some((v[lo] + v[hi]) / 2.0)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Nearest-rank percentile (`p` in 0..=100). Chosen over interpolation
|
||||||
|
/// because bench cells hold only a handful of samples: with n=5, p95/p99
|
||||||
|
/// resolve to the max, which honestly says "this is the worst we saw"
|
||||||
|
/// rather than inventing a value between samples we never observed.
|
||||||
|
fn percentile(values: impl Iterator<Item = f64>, p: f64) -> Option<f64> {
|
||||||
|
let mut v: Vec<f64> = values.collect();
|
||||||
|
if v.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
// rank = ceil(p/100 * n), clamped to [1, n]; index is rank-1.
|
||||||
|
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||||
|
let idx = rank.clamp(1, v.len()) - 1;
|
||||||
|
Some(v[idx])
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -717,6 +1191,21 @@ mod tests {
|
|||||||
decode_tps: Some(50.0),
|
decode_tps: Some(50.0),
|
||||||
total_s: Some(1.0),
|
total_s: Some(1.0),
|
||||||
completion_tokens: Some(50),
|
completion_tokens: Some(50),
|
||||||
|
prefill_ms: Some(200),
|
||||||
|
decode_ms: Some(1000),
|
||||||
|
prefill_tokens: Some(130),
|
||||||
|
vram_used_mb: Some(42000),
|
||||||
|
gpu_util_pct: Some(88),
|
||||||
|
gpu_temp_c: Some(64),
|
||||||
|
concurrency: None,
|
||||||
|
ttft_p95_s: None,
|
||||||
|
queue_wait_ms: None,
|
||||||
|
rejected: None,
|
||||||
|
swap_unload_ms: None,
|
||||||
|
swap_load_ms: None,
|
||||||
|
artifact: None,
|
||||||
|
quality_score: None,
|
||||||
|
scorer: None,
|
||||||
ok,
|
ok,
|
||||||
error: if ok { None } else { Some("boom".into()) },
|
error: if ok { None } else { Some("boom".into()) },
|
||||||
}
|
}
|
||||||
@@ -755,6 +1244,197 @@ mod tests {
|
|||||||
assert!((rows[0].ttft_s_median.unwrap() - 0.3).abs() < 1e-9);
|
assert!((rows[0].ttft_s_median.unwrap() - 0.3).abs() < 1e-9);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn report_surfaces_percentiles_and_prefill_split() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
// Five samples on one cell with spread TTFT so percentiles differ
|
||||||
|
// from the median, plus a server-measured prefill/decode split.
|
||||||
|
for (i, ttft) in [0.10, 0.12, 0.14, 0.16, 0.50].iter().enumerate() {
|
||||||
|
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||||
|
r.ttft_s = Some(*ttft);
|
||||||
|
r.total_s = Some(ttft + 1.0);
|
||||||
|
r.prefill_ms = Some(200 + i as u64);
|
||||||
|
r.prefill_tokens = Some(400);
|
||||||
|
s.insert_run(&r).unwrap();
|
||||||
|
}
|
||||||
|
let rows = s.report_rows().unwrap();
|
||||||
|
assert_eq!(rows.len(), 1);
|
||||||
|
let row = &rows[0];
|
||||||
|
assert_eq!(row.samples, 5);
|
||||||
|
// p50 is the middle value; p95/p99 (nearest-rank, n=5) hit the max.
|
||||||
|
assert!((row.ttft_s_median.unwrap() - 0.14).abs() < 1e-9);
|
||||||
|
assert!((row.ttft_s_p95.unwrap() - 0.50).abs() < 1e-9);
|
||||||
|
assert!((row.ttft_s_p99.unwrap() - 0.50).abs() < 1e-9);
|
||||||
|
// prefill tok/s = 400 tok / ~0.2 s ≈ 2000 tok/s.
|
||||||
|
assert!(row.prefill_tps_median.unwrap() > 1900.0);
|
||||||
|
assert!(row.prefill_ms_median.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn report_surfaces_vram_and_gpu_telemetry() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||||
|
// Node total VRAM from discovery devices → headroom denominator.
|
||||||
|
r.gpus_json =
|
||||||
|
Some(r#"[{"name":"RTX 5090","vram_total_mb":32000},{"name":"RTX 5090","vram_total_mb":32000}]"#.into());
|
||||||
|
r.vram_used_mb = Some(42000);
|
||||||
|
r.gpu_util_pct = Some(90);
|
||||||
|
r.gpu_temp_c = Some(66);
|
||||||
|
s.insert_run(&r).unwrap();
|
||||||
|
let rows = s.report_rows().unwrap();
|
||||||
|
let row = &rows[0];
|
||||||
|
assert_eq!(row.vram_used_mb_median, Some(42000.0));
|
||||||
|
assert_eq!(row.vram_total_mb, Some(64000)); // 2× 32000
|
||||||
|
assert_eq!(row.gpu_util_pct_median, Some(90.0));
|
||||||
|
assert_eq!(row.gpu_temp_c_median, Some(66.0));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn report_surfaces_concurrency_burst_metrics() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
// Two concurrency:8 burst-runs with shed load and a queue-wait tail.
|
||||||
|
for (qw, rej) in [(120.0, 1u32), (180.0, 3u32)] {
|
||||||
|
let mut r = rec("beast", "sha", "m", "concurrency:8", true);
|
||||||
|
r.concurrency = Some(8);
|
||||||
|
r.ttft_p95_s = Some(0.9);
|
||||||
|
r.queue_wait_ms = Some(qw);
|
||||||
|
r.rejected = Some(rej);
|
||||||
|
s.insert_run(&r).unwrap();
|
||||||
|
}
|
||||||
|
let row = s
|
||||||
|
.report_rows()
|
||||||
|
.unwrap()
|
||||||
|
.into_iter()
|
||||||
|
.find(|r| r.scenario_id == "concurrency:8")
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(row.concurrency, Some(8));
|
||||||
|
assert_eq!(row.queue_wait_ms_median, Some(150.0)); // median(120,180)
|
||||||
|
assert_eq!(row.rejected_median, Some(2.0)); // median(1,3)
|
||||||
|
assert_eq!(row.ttft_p95_load_s, Some(0.9));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn capability_runs_store_artifact_and_accept_scores() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
let mut r = rec("beast", "sha", "m", "capability:plan", true);
|
||||||
|
r.artifact = Some("a thorough implementation plan...".into());
|
||||||
|
s.insert_run(&r).unwrap();
|
||||||
|
// A non-capability row must not appear in capability_runs.
|
||||||
|
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let unscored = s.capability_runs(true).unwrap();
|
||||||
|
assert_eq!(unscored.len(), 1);
|
||||||
|
let id = unscored[0].id;
|
||||||
|
assert!(unscored[0].artifact.as_deref().unwrap().contains("plan"));
|
||||||
|
assert!(unscored[0].quality_score.is_none());
|
||||||
|
|
||||||
|
// Score it; it then drops off the unscored worklist.
|
||||||
|
assert_eq!(s.set_score(id, 7.5, "manual").unwrap(), 1);
|
||||||
|
assert!(s.capability_runs(true).unwrap().is_empty());
|
||||||
|
let all = s.capability_runs(false).unwrap();
|
||||||
|
assert_eq!(all[0].quality_score, Some(7.5));
|
||||||
|
assert_eq!(all[0].scorer.as_deref(), Some("manual"));
|
||||||
|
// Unknown id updates nothing.
|
||||||
|
assert_eq!(s.set_score(999_999, 1.0, "manual").unwrap(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn swap_costs_pivots_swap_rows() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
for (unload, load) in [(300u64, 24000u64), (340, 26000)] {
|
||||||
|
let mut r = rec("beast", "sha", "m", "swap", true);
|
||||||
|
r.swap_unload_ms = Some(unload);
|
||||||
|
r.swap_load_ms = Some(load);
|
||||||
|
r.ttft_s = Some(2.5); // cold first-request
|
||||||
|
r.total_s = Some(5.0);
|
||||||
|
s.insert_run(&r).unwrap();
|
||||||
|
}
|
||||||
|
// A non-swap row must be ignored by swap_costs.
|
||||||
|
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||||
|
.unwrap();
|
||||||
|
let costs = s.swap_costs().unwrap();
|
||||||
|
assert_eq!(costs.len(), 1);
|
||||||
|
let c = &costs[0];
|
||||||
|
assert_eq!(c.unload_ms_median, Some(320.0)); // median(300,340)
|
||||||
|
assert_eq!(c.load_ms_median, Some(25000.0)); // median(24000,26000)
|
||||||
|
assert_eq!(c.cold_ttft_s_median, Some(2.5));
|
||||||
|
assert_eq!(c.samples, 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn scaling_pivots_chat_cells_and_computes_flatness() {
|
||||||
|
let s = Store::open_in_memory().unwrap();
|
||||||
|
// Two context points for one model: decode tok/s 50 @128, 45 @4096.
|
||||||
|
let mut small = rec("beast", "sha", "m", "chat:128", true);
|
||||||
|
small.prompt_size_approx = 128;
|
||||||
|
small.decode_tps = Some(50.0);
|
||||||
|
small.prefill_ms = Some(100);
|
||||||
|
small.prefill_tokens = Some(128);
|
||||||
|
s.insert_run(&small).unwrap();
|
||||||
|
let mut big = rec("beast", "sha", "m", "chat:4096", true);
|
||||||
|
big.prompt_size_approx = 4096;
|
||||||
|
big.decode_tps = Some(45.0);
|
||||||
|
big.prefill_ms = Some(1000);
|
||||||
|
big.prefill_tokens = Some(4096);
|
||||||
|
s.insert_run(&big).unwrap();
|
||||||
|
// A concurrency cell must NOT leak into the scaling curve.
|
||||||
|
let mut conc = rec("beast", "sha", "m", "concurrency:8", true);
|
||||||
|
conc.concurrency = Some(8);
|
||||||
|
s.insert_run(&conc).unwrap();
|
||||||
|
|
||||||
|
let curves = s.scaling().unwrap();
|
||||||
|
assert_eq!(curves.len(), 1);
|
||||||
|
let c = &curves[0];
|
||||||
|
assert_eq!(c.points.len(), 2); // only the two chat:<n> points
|
||||||
|
assert_eq!(c.points[0].prompt_size, 128); // ordered ascending
|
||||||
|
assert_eq!(c.points[1].prompt_size, 4096);
|
||||||
|
// flatness = decode@largest / decode@smallest = 45/50 = 0.9
|
||||||
|
assert!((c.decode_flatness.unwrap() - 0.9).abs() < 1e-9);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gpu_total_vram_sums_devices() {
|
||||||
|
let j = r#"[{"vram_total_mb":32000},{"vram_total_mb":32000}]"#;
|
||||||
|
assert_eq!(gpu_total_vram_mb(j), Some(64000));
|
||||||
|
assert_eq!(gpu_total_vram_mb("[]"), None);
|
||||||
|
assert_eq!(gpu_total_vram_mb("not json"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn percentile_nearest_rank() {
|
||||||
|
let vals = || [1.0, 2.0, 3.0, 4.0, 5.0].into_iter();
|
||||||
|
assert_eq!(percentile(vals(), 50.0), Some(3.0));
|
||||||
|
assert_eq!(percentile(vals(), 95.0), Some(5.0));
|
||||||
|
assert_eq!(percentile(vals(), 99.0), Some(5.0));
|
||||||
|
assert_eq!(percentile(std::iter::empty(), 95.0), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn migration_is_idempotent_and_backfills() {
|
||||||
|
// A DB whose `runs` table predates the prefill columns: create the
|
||||||
|
// pre-#85 shape, insert a row, then run ensure_columns twice.
|
||||||
|
let conn = Connection::open_in_memory().unwrap();
|
||||||
|
conn.execute_batch(
|
||||||
|
"CREATE TABLE runs (id INTEGER PRIMARY KEY, ttft_s REAL);
|
||||||
|
INSERT INTO runs (ttft_s) VALUES (0.1);",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
for _ in 0..2 {
|
||||||
|
Store::ensure_columns(
|
||||||
|
&conn,
|
||||||
|
"runs",
|
||||||
|
&[("prefill_ms", "INTEGER"), ("decode_ms", "INTEGER")],
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
// Columns now exist and the old row reads them back as NULL.
|
||||||
|
let got: Option<i64> = conn
|
||||||
|
.query_row("SELECT prefill_ms FROM runs", [], |r| r.get(0))
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(got, None);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn gpu_label_formats() {
|
fn gpu_label_formats() {
|
||||||
let two = r#"[{"name":"NVIDIA GeForce RTX 5090"},{"name":"NVIDIA GeForce RTX 5090"}]"#;
|
let two = r#"[{"name":"NVIDIA GeForce RTX 5090"},{"name":"NVIDIA GeForce RTX 5090"}]"#;
|
||||||
|
|||||||
@@ -9,11 +9,11 @@
|
|||||||
|
|
||||||
use crate::client::TargetClient;
|
use crate::client::TargetClient;
|
||||||
use crate::config::{BenchConfig, TargetConfig, TargetKind};
|
use crate::config::{BenchConfig, TargetConfig, TargetKind};
|
||||||
use crate::scenario::{RunCtx, build_scenarios};
|
use crate::scenario::{RunCtx, ScenarioMetrics, build_scenarios};
|
||||||
use crate::store::{RunRecord, Store};
|
use crate::store::{RunRecord, Store};
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use cortex_core::build_info::BuildInfo;
|
use cortex_core::build_info::BuildInfo;
|
||||||
use cortex_core::discovery::DiscoveryResponse;
|
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||||
use cortex_core::harness::ModelInfo;
|
use cortex_core::harness::ModelInfo;
|
||||||
|
|
||||||
/// helexa-bench's own build version.
|
/// helexa-bench's own build version.
|
||||||
@@ -38,6 +38,38 @@ pub struct SweepSummary {
|
|||||||
pub targets_unreachable: usize,
|
pub targets_unreachable: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Node-level GPU telemetry folded from one `/health` snapshot (#87):
|
||||||
|
/// VRAM used summed across the node's devices, and the hottest/busiest
|
||||||
|
/// single device for utilization and temperature.
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
struct HealthAgg {
|
||||||
|
vram_used_mb: u64,
|
||||||
|
gpu_util_pct: u32,
|
||||||
|
gpu_temp_c: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cold-load / model-swap timing for one measure_swap cycle (#90).
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
struct SwapTiming {
|
||||||
|
unload_ms: u64,
|
||||||
|
load_ms: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HealthAgg {
|
||||||
|
fn from_health(h: &HealthResponse) -> Self {
|
||||||
|
HealthAgg {
|
||||||
|
vram_used_mb: h.devices.iter().map(|d| d.vram_used_mb).sum(),
|
||||||
|
gpu_util_pct: h
|
||||||
|
.devices
|
||||||
|
.iter()
|
||||||
|
.map(|d| d.utilization_pct)
|
||||||
|
.max()
|
||||||
|
.unwrap_or(0),
|
||||||
|
gpu_temp_c: h.devices.iter().map(|d| d.temp_c).max().unwrap_or(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub struct Sweeper {
|
pub struct Sweeper {
|
||||||
cfg: BenchConfig,
|
cfg: BenchConfig,
|
||||||
client: TargetClient,
|
client: TargetClient,
|
||||||
@@ -72,6 +104,105 @@ impl Sweeper {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Deliberate cold-load / model-swap cost measurement (#90), invoked by
|
||||||
|
/// the `swap-cost` subcommand — **never** the continuous sweep. For each
|
||||||
|
/// neuron target and each currently-warm model: unload it, time the
|
||||||
|
/// reload, then time a cold first request. This takes the model offline
|
||||||
|
/// for the reload, so it is an explicit operator action (maintenance
|
||||||
|
/// window), recorded under `scenario_id = "swap"`.
|
||||||
|
pub async fn swap_cost_once(&self) -> Result<SweepSummary> {
|
||||||
|
let mut summary = SweepSummary::default();
|
||||||
|
for target in &self.cfg.targets {
|
||||||
|
if target.kind != TargetKind::Neuron {
|
||||||
|
continue; // load/unload is a neuron-native operation
|
||||||
|
}
|
||||||
|
let build = match self.client.fetch_version(target).await {
|
||||||
|
Ok(b) => b,
|
||||||
|
Err(e) => {
|
||||||
|
summary.targets_unreachable += 1;
|
||||||
|
tracing::warn!(target = %target.name, error = %format!("{e:#}"), "swap: target unreachable");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let discovery = self.client.fetch_discovery(target).await.unwrap_or(None);
|
||||||
|
let models = self.client.warm_models(target).await.unwrap_or_default();
|
||||||
|
for model in &models {
|
||||||
|
match self
|
||||||
|
.measure_swap(target, &build, discovery.as_ref(), model)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(()) => summary.measured += 1,
|
||||||
|
Err(e) => {
|
||||||
|
summary.failed += 1;
|
||||||
|
tracing::warn!(target = %target.name, model = %model.id, error = %format!("{e:#}"), "swap: measurement failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(summary)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unload → timed reload → timed cold first request for one model.
|
||||||
|
async fn measure_swap(
|
||||||
|
&self,
|
||||||
|
target: &TargetConfig,
|
||||||
|
build: &BuildInfo,
|
||||||
|
discovery: Option<&DiscoveryResponse>,
|
||||||
|
model: &ModelInfo,
|
||||||
|
) -> Result<()> {
|
||||||
|
let spec = TargetClient::spec_from_info(model)?;
|
||||||
|
tracing::warn!(target = %target.name, model = %model.id, "swap: unloading (model goes offline until reload)");
|
||||||
|
|
||||||
|
let t0 = std::time::Instant::now();
|
||||||
|
self.client.unload_model(target, &model.id).await?;
|
||||||
|
let unload_ms = t0.elapsed().as_millis() as u64;
|
||||||
|
|
||||||
|
let t1 = std::time::Instant::now();
|
||||||
|
self.client.load_model(target, &spec).await?;
|
||||||
|
let load_ms = t1.elapsed().as_millis() as u64;
|
||||||
|
tracing::info!(target = %target.name, model = %model.id, unload_ms, load_ms, "swap: reloaded; measuring cold first request");
|
||||||
|
|
||||||
|
// Cold first request — caches empty straight after the load.
|
||||||
|
let ctx = RunCtx {
|
||||||
|
client: self.client.http(),
|
||||||
|
chat_url: self.client.chat_url(target),
|
||||||
|
model_id: model.id.clone(),
|
||||||
|
max_tokens: self.cfg.scenarios.max_tokens,
|
||||||
|
timeout: self.cfg.bench.request_timeout(),
|
||||||
|
};
|
||||||
|
let cold = crate::scenario::cold_probe(&ctx).await;
|
||||||
|
let swap = SwapTiming { unload_ms, load_ms };
|
||||||
|
let rec = match &cold {
|
||||||
|
Ok(m) => self.build_record(
|
||||||
|
target,
|
||||||
|
build,
|
||||||
|
discovery,
|
||||||
|
model,
|
||||||
|
"swap",
|
||||||
|
0,
|
||||||
|
Ok(m),
|
||||||
|
None,
|
||||||
|
Some(swap),
|
||||||
|
),
|
||||||
|
Err(e) => {
|
||||||
|
let msg = format!("{e:#}");
|
||||||
|
self.build_record(
|
||||||
|
target,
|
||||||
|
build,
|
||||||
|
discovery,
|
||||||
|
model,
|
||||||
|
"swap",
|
||||||
|
0,
|
||||||
|
Err(&msg),
|
||||||
|
None,
|
||||||
|
Some(swap),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
self.store.insert_run(&rec)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// One full pass over all targets.
|
/// One full pass over all targets.
|
||||||
pub async fn run_once(&self) -> Result<SweepSummary> {
|
pub async fn run_once(&self) -> Result<SweepSummary> {
|
||||||
let mut summary = SweepSummary::default();
|
let mut summary = SweepSummary::default();
|
||||||
@@ -131,7 +262,21 @@ impl Sweeper {
|
|||||||
}
|
}
|
||||||
|
|
||||||
for i in 0..need {
|
for i in 0..need {
|
||||||
match scenario.run(&ctx).await {
|
let result = scenario.run(&ctx).await;
|
||||||
|
// Sample GPU telemetry right after the run, while the
|
||||||
|
// model is loaded and decode VRAM is at its recent peak
|
||||||
|
// (#87). neuron's /health is ~5s-cached, so this is a
|
||||||
|
// coarse high-water proxy, not an instantaneous peak — but
|
||||||
|
// it's the headroom signal we can read over the wire. A
|
||||||
|
// flaky /health degrades to None, never a failed run.
|
||||||
|
let health = self
|
||||||
|
.client
|
||||||
|
.fetch_health(target)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.map(|h| HealthAgg::from_health(&h));
|
||||||
|
match result {
|
||||||
Ok(m) => {
|
Ok(m) => {
|
||||||
let rec = self.build_record(
|
let rec = self.build_record(
|
||||||
target,
|
target,
|
||||||
@@ -141,6 +286,8 @@ impl Sweeper {
|
|||||||
scenario.id(),
|
scenario.id(),
|
||||||
scenario.prompt_size(),
|
scenario.prompt_size(),
|
||||||
Ok(&m),
|
Ok(&m),
|
||||||
|
health,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
self.store.insert_run(&rec)?;
|
self.store.insert_run(&rec)?;
|
||||||
summary.measured += 1;
|
summary.measured += 1;
|
||||||
@@ -160,6 +307,8 @@ impl Sweeper {
|
|||||||
scenario.id(),
|
scenario.id(),
|
||||||
scenario.prompt_size(),
|
scenario.prompt_size(),
|
||||||
Err(&msg),
|
Err(&msg),
|
||||||
|
health,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
self.store.insert_run(&rec)?;
|
self.store.insert_run(&rec)?;
|
||||||
summary.failed += 1;
|
summary.failed += 1;
|
||||||
@@ -186,19 +335,14 @@ impl Sweeper {
|
|||||||
scenario_id: &str,
|
scenario_id: &str,
|
||||||
prompt_size: u32,
|
prompt_size: u32,
|
||||||
result: Result<&crate::scenario::ScenarioMetrics, &str>,
|
result: Result<&crate::scenario::ScenarioMetrics, &str>,
|
||||||
|
health: Option<HealthAgg>,
|
||||||
|
swap: Option<SwapTiming>,
|
||||||
) -> RunRecord {
|
) -> RunRecord {
|
||||||
let (ok, error, ttft, decode, total, prompt_tokens, completion) = match result {
|
let (m, error): (Option<&ScenarioMetrics>, Option<String>) = match result {
|
||||||
Ok(m) => (
|
Ok(m) => (Some(m), None),
|
||||||
true,
|
Err(e) => (None, Some(e.to_string())),
|
||||||
None,
|
|
||||||
Some(m.ttft_s),
|
|
||||||
m.decode_tps,
|
|
||||||
Some(m.total_s),
|
|
||||||
m.prompt_tokens,
|
|
||||||
Some(m.completion_tokens),
|
|
||||||
),
|
|
||||||
Err(e) => (false, Some(e.to_string()), None, None, None, None, None),
|
|
||||||
};
|
};
|
||||||
|
let ok = m.is_some();
|
||||||
|
|
||||||
RunRecord {
|
RunRecord {
|
||||||
ts: chrono::Utc::now().to_rfc3339(),
|
ts: chrono::Utc::now().to_rfc3339(),
|
||||||
@@ -230,12 +374,29 @@ impl Sweeper {
|
|||||||
.unwrap_or_else(|_| "[]".to_string()),
|
.unwrap_or_else(|_| "[]".to_string()),
|
||||||
scenario_id: scenario_id.to_string(),
|
scenario_id: scenario_id.to_string(),
|
||||||
prompt_size_approx: prompt_size,
|
prompt_size_approx: prompt_size,
|
||||||
prompt_tokens_actual: prompt_tokens,
|
prompt_tokens_actual: m.and_then(|m| m.prompt_tokens),
|
||||||
max_tokens: self.cfg.scenarios.max_tokens,
|
max_tokens: self.cfg.scenarios.max_tokens,
|
||||||
ttft_s: ttft,
|
ttft_s: m.map(|m| m.ttft_s),
|
||||||
decode_tps: decode,
|
decode_tps: m.and_then(|m| m.decode_tps),
|
||||||
total_s: total,
|
total_s: m.map(|m| m.total_s),
|
||||||
completion_tokens: completion,
|
completion_tokens: m.map(|m| m.completion_tokens),
|
||||||
|
prefill_ms: m.and_then(|m| m.prefill_ms),
|
||||||
|
decode_ms: m.and_then(|m| m.decode_ms),
|
||||||
|
prefill_tokens: m.and_then(|m| m.prefill_tokens),
|
||||||
|
vram_used_mb: health.map(|h| h.vram_used_mb),
|
||||||
|
gpu_util_pct: health.map(|h| h.gpu_util_pct),
|
||||||
|
gpu_temp_c: health.map(|h| h.gpu_temp_c),
|
||||||
|
concurrency: m.and_then(|m| m.concurrency),
|
||||||
|
ttft_p95_s: m.and_then(|m| m.ttft_p95_s),
|
||||||
|
queue_wait_ms: m.and_then(|m| m.queue_wait_ms_median),
|
||||||
|
rejected: m.and_then(|m| m.rejected),
|
||||||
|
swap_unload_ms: swap.map(|s| s.unload_ms),
|
||||||
|
swap_load_ms: swap.map(|s| s.load_ms),
|
||||||
|
// Capability artifact (#91); score/scorer are attached later by
|
||||||
|
// the `score` subcommand or a future LLM-judge.
|
||||||
|
artifact: m.and_then(|m| m.artifact.clone()),
|
||||||
|
quality_score: None,
|
||||||
|
scorer: None,
|
||||||
ok,
|
ok,
|
||||||
error,
|
error,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -46,6 +46,21 @@ fn rec(
|
|||||||
decode_tps: if ok { Some(30.0) } else { None },
|
decode_tps: if ok { Some(30.0) } else { None },
|
||||||
total_s: if ok { Some(2.0) } else { None },
|
total_s: if ok { Some(2.0) } else { None },
|
||||||
completion_tokens: if ok { Some(60) } else { None },
|
completion_tokens: if ok { Some(60) } else { None },
|
||||||
|
prefill_ms: if ok { Some(150) } else { None },
|
||||||
|
decode_ms: if ok { Some(1800) } else { None },
|
||||||
|
prefill_tokens: if ok { Some(130) } else { None },
|
||||||
|
vram_used_mb: if ok { Some(42000) } else { None },
|
||||||
|
gpu_util_pct: if ok { Some(85) } else { None },
|
||||||
|
gpu_temp_c: if ok { Some(63) } else { None },
|
||||||
|
concurrency: None,
|
||||||
|
ttft_p95_s: None,
|
||||||
|
queue_wait_ms: None,
|
||||||
|
rejected: None,
|
||||||
|
swap_unload_ms: None,
|
||||||
|
swap_load_ms: None,
|
||||||
|
artifact: None,
|
||||||
|
quality_score: None,
|
||||||
|
scorer: None,
|
||||||
ok,
|
ok,
|
||||||
error: if ok { None } else { Some("boom".into()) },
|
error: if ok { None } else { Some("boom".into()) },
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -89,6 +89,9 @@ fn config_for(endpoint: String, db_path: String) -> BenchConfig {
|
|||||||
scenarios: ScenarioConfig {
|
scenarios: ScenarioConfig {
|
||||||
prompt_sizes: vec![128], // single scenario keeps assertions simple
|
prompt_sizes: vec![128], // single scenario keeps assertions simple
|
||||||
max_tokens: 16,
|
max_tokens: 16,
|
||||||
|
concurrency_levels: Vec::new(),
|
||||||
|
concurrency_prompt_tokens: 512,
|
||||||
|
capability_probes: Vec::new(),
|
||||||
},
|
},
|
||||||
api: Default::default(),
|
api: Default::default(),
|
||||||
targets: vec![TargetConfig {
|
targets: vec![TargetConfig {
|
||||||
|
|||||||
41
crates/helexa-router/Cargo.toml
Normal file
41
crates/helexa-router/Cargo.toml
Normal file
@@ -0,0 +1,41 @@
|
|||||||
|
[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 }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
clap = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
tracing-subscriber = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
# Jail (isolated cwd + env) for config tests.
|
||||||
|
figment = { workspace = true, features = ["test"] }
|
||||||
|
# Self-signed cert generation + a minimal HTTPS server for the outbound
|
||||||
|
# TLS-pinning tests (#74).
|
||||||
|
rcgen = "0.13"
|
||||||
|
rustls = "0.23"
|
||||||
|
tokio-rustls = "0.26"
|
||||||
260
crates/helexa-router/src/catalogue.rs
Normal file
260
crates/helexa-router/src/catalogue.rs
Normal file
@@ -0,0 +1,260 @@
|
|||||||
|
//! 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));
|
||||||
|
// Re-derive the flat ecosystem fields (#78) from the merged (tightest)
|
||||||
|
// limit — the values deserialized from each cortex are per-operator and
|
||||||
|
// may not match the federation-wide merge.
|
||||||
|
for e in &mut out {
|
||||||
|
e.sync_flat_limit();
|
||||||
|
}
|
||||||
|
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,
|
||||||
|
// Derived from `limit` by the final sync pass in aggregate_models.
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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,
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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);
|
||||||
|
// Flat #78 fields re-derived from the merged (tightest) limit.
|
||||||
|
assert_eq!(out[0].max_model_len, Some(16_384));
|
||||||
|
assert_eq!(out[0].max_input_tokens, None);
|
||||||
|
assert_eq!(out[0].max_output_tokens, Some(4096));
|
||||||
|
}
|
||||||
|
}
|
||||||
100
crates/helexa-router/src/config.rs
Normal file
100
crates/helexa-router/src/config.rs
Normal file
@@ -0,0 +1,100 @@
|
|||||||
|
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>,
|
||||||
|
/// Path to a PEM trust anchor that **enrols** this cortex (#74): the
|
||||||
|
/// expected CA (or self-signed cert) the cortex's TLS cert must chain
|
||||||
|
/// to. When set on an `https://` endpoint, the router builds a client
|
||||||
|
/// that trusts **only** this anchor (platform roots disabled), so the
|
||||||
|
/// outbound router→cortex hop — which carries the client's bearer —
|
||||||
|
/// reaches a cert the router was told to expect, and a rogue endpoint
|
||||||
|
/// presenting any other (even publicly-valid) cert is rejected at the
|
||||||
|
/// TLS handshake. A rejected handshake surfaces as a connection error,
|
||||||
|
/// which the poller (#72) already treats as unreachable → excluded.
|
||||||
|
///
|
||||||
|
/// `None` → standard platform-root validation (use for cortexes behind
|
||||||
|
/// a publicly-trusted cert, or plaintext `http://` on a private network
|
||||||
|
/// where the WireGuard mesh is the trust boundary).
|
||||||
|
#[serde(default)]
|
||||||
|
pub tls_ca: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RouterConfig {
|
||||||
|
/// 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![],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
221
crates/helexa-router/src/dispatch.rs
Normal file
221
crates/helexa-router/src/dispatch.rs
Normal file
@@ -0,0 +1,221 @@
|
|||||||
|
//! Capacity-aware dispatch (#73) — the router's data path.
|
||||||
|
//!
|
||||||
|
//! Given an inbound request's `model`, pick a reachable cortex that can
|
||||||
|
//! serve it (preferring warm/loaded, region-affine, higher-headroom),
|
||||||
|
//! forward the client's bearer **unchanged** (auth stays at cortex), and
|
||||||
|
//! stream the response back verbatim via the shared [`helexa_stream`]
|
||||||
|
//! module. Cortex's #63-shaped rejections (`429 rate_limit_exceeded`,
|
||||||
|
//! `400 context_length_exceeded`, …) pass through untouched. Transport
|
||||||
|
//! failures fail over to the next feasible cortex; a genuine HTTP response —
|
||||||
|
//! any status — is returned as-is and never retried away.
|
||||||
|
//!
|
||||||
|
//! The router holds **no entitlement logic**: it routes on capacity, not
|
||||||
|
//! budget.
|
||||||
|
|
||||||
|
use crate::config::CortexEndpoint;
|
||||||
|
use crate::error::envelope_response;
|
||||||
|
use crate::state::RouterState;
|
||||||
|
use axum::body::Bytes;
|
||||||
|
use axum::http::HeaderMap;
|
||||||
|
use axum::response::Response;
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
use helexa_stream::{ChunkObserver, StreamError};
|
||||||
|
use std::cmp::Reverse;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Retry-After hint (seconds) on the router's own transient rejections.
|
||||||
|
const RETRY_AFTER_SECS: u64 = 5;
|
||||||
|
|
||||||
|
/// Outcome of choosing where to send a request.
|
||||||
|
#[derive(Debug, PartialEq, Eq)]
|
||||||
|
pub enum Selection {
|
||||||
|
/// Feasible reachable cortexes, best-first (failover order).
|
||||||
|
Candidates(Vec<CortexEndpoint>),
|
||||||
|
/// Some cortex knows the model but none are reachable right now → 503.
|
||||||
|
NoReachableCapacity,
|
||||||
|
/// No configured cortex serves the model at all → 404.
|
||||||
|
UnknownModel,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rank the reachable cortexes that can serve `model`, best-first.
|
||||||
|
///
|
||||||
|
/// Ordering (each a tie-break for the next): loaded/warm before cold-loadable
|
||||||
|
/// · region match before not · more healthy nodes before fewer · name for
|
||||||
|
/// determinism.
|
||||||
|
pub async fn select_cortexes(state: &RouterState, model: &str) -> Selection {
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
let by_name: HashMap<&str, &CortexEndpoint> = state
|
||||||
|
.cortexes
|
||||||
|
.iter()
|
||||||
|
.map(|c| (c.name.as_str(), c))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut ranked: Vec<Ranked> = Vec::new();
|
||||||
|
let mut known_anywhere = false;
|
||||||
|
|
||||||
|
for (name, t) in topo.iter() {
|
||||||
|
let Some(entry) = t.models.get(model) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
if !crate::state::entry_feasible(entry) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Known even via an unreachable cortex's last-good poll — lets us
|
||||||
|
// tell "temporarily down" (503) from "nobody serves it" (404).
|
||||||
|
known_anywhere = true;
|
||||||
|
if !t.reachable {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let Some(ep) = by_name.get(name.as_str()) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let region_match = match (&state.region, &ep.region) {
|
||||||
|
(Some(r), Some(cr)) => r == cr,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
ranked.push(Ranked {
|
||||||
|
loaded: entry.loaded,
|
||||||
|
region_match,
|
||||||
|
healthy_nodes: t.healthy_nodes,
|
||||||
|
endpoint: (*ep).clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
if ranked.is_empty() {
|
||||||
|
return if known_anywhere {
|
||||||
|
Selection::NoReachableCapacity
|
||||||
|
} else {
|
||||||
|
Selection::UnknownModel
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
ranked.sort_by(|a, b| {
|
||||||
|
// false < true, so negate the "good" booleans to sort good first.
|
||||||
|
(
|
||||||
|
!a.loaded,
|
||||||
|
!a.region_match,
|
||||||
|
Reverse(a.healthy_nodes),
|
||||||
|
&a.endpoint.name,
|
||||||
|
)
|
||||||
|
.cmp(&(
|
||||||
|
!b.loaded,
|
||||||
|
!b.region_match,
|
||||||
|
Reverse(b.healthy_nodes),
|
||||||
|
&b.endpoint.name,
|
||||||
|
))
|
||||||
|
});
|
||||||
|
|
||||||
|
Selection::Candidates(ranked.into_iter().map(|r| r.endpoint).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Ranked {
|
||||||
|
loaded: bool,
|
||||||
|
region_match: bool,
|
||||||
|
healthy_nodes: u32,
|
||||||
|
endpoint: CortexEndpoint,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Proxy an inbound inference request to a capacity-bearing cortex.
|
||||||
|
///
|
||||||
|
/// `path` is the inference path to forward to (same on the cortex, e.g.
|
||||||
|
/// `/v1/chat/completions`). The body is parsed only to extract `model`.
|
||||||
|
pub async fn dispatch(
|
||||||
|
state: &RouterState,
|
||||||
|
path: &str,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
) -> Response {
|
||||||
|
let Some(model) = extract_model(&body) else {
|
||||||
|
return envelope_response(OpenAiError::new(
|
||||||
|
400,
|
||||||
|
"invalid_request_error",
|
||||||
|
"missing_model_field",
|
||||||
|
"missing 'model' field in request body",
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
|
let candidates = match select_cortexes(state, &model).await {
|
||||||
|
Selection::Candidates(c) => c,
|
||||||
|
Selection::UnknownModel => {
|
||||||
|
return envelope_response(
|
||||||
|
OpenAiError::new(
|
||||||
|
404,
|
||||||
|
"invalid_request_error",
|
||||||
|
"model_not_found",
|
||||||
|
format!("no operator serves model '{model}'"),
|
||||||
|
)
|
||||||
|
.with_param("model"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Selection::NoReachableCapacity => {
|
||||||
|
return envelope_response(OpenAiError::service_unavailable(
|
||||||
|
format!("model '{model}' is temporarily unavailable on all operators"),
|
||||||
|
Some(RETRY_AFTER_SECS),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Try candidates in order, failing over only on transport errors. A
|
||||||
|
// genuine HTTP response (any status — including cortex's #63 429/400)
|
||||||
|
// is returned verbatim and never retried away.
|
||||||
|
for ep in &candidates {
|
||||||
|
// A candidate whose pinned TLS client failed to build (#74) is
|
||||||
|
// disabled — skip it and fail over, same as an unreachable cortex.
|
||||||
|
let Some(client) = state.client_for(&ep.name) else {
|
||||||
|
tracing::warn!(cortex = %ep.name, "no TLS client (disabled); skipping candidate");
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let url = format!("{}{}", ep.endpoint, path);
|
||||||
|
tracing::info!(cortex = %ep.name, url = %url, model = %model, "dispatching");
|
||||||
|
match helexa_stream::forward_streaming(
|
||||||
|
client,
|
||||||
|
&url,
|
||||||
|
headers.clone(),
|
||||||
|
body.clone(),
|
||||||
|
NoopObserver,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(resp) => return resp,
|
||||||
|
Err(StreamError::Upstream(e)) => {
|
||||||
|
tracing::warn!(
|
||||||
|
cortex = %ep.name,
|
||||||
|
url = %url,
|
||||||
|
error = %e,
|
||||||
|
"cortex unreachable; failing over"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(StreamError::ResponseBuild(msg)) => {
|
||||||
|
tracing::error!(cortex = %ep.name, error = %msg, "failed to build proxied response");
|
||||||
|
return envelope_response(OpenAiError::without_code(
|
||||||
|
500,
|
||||||
|
"api_error",
|
||||||
|
"failed to build proxied response",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Every feasible cortex failed to connect.
|
||||||
|
tracing::warn!(model = %model, tried = candidates.len(), "all feasible operators unreachable");
|
||||||
|
envelope_response(OpenAiError::service_unavailable(
|
||||||
|
format!("all operators able to serve '{model}' are unreachable"),
|
||||||
|
Some(RETRY_AFTER_SECS),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pull the `model` field out of a request body without re-serialising it.
|
||||||
|
fn extract_model(body: &Bytes) -> Option<String> {
|
||||||
|
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||||
|
v.get("model")?.as_str().map(str::to_string)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The router proxies bytes verbatim and keeps no per-request policy, so it
|
||||||
|
/// needs no observation hooks. (Token metrics/metering stay at cortex.)
|
||||||
|
struct NoopObserver;
|
||||||
|
|
||||||
|
impl ChunkObserver for NoopObserver {
|
||||||
|
fn observe(&mut self, _chunk: &[u8]) {}
|
||||||
|
fn finish(&mut self) {}
|
||||||
|
}
|
||||||
27
crates/helexa-router/src/error.rs
Normal file
27
crates/helexa-router/src/error.rs
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
//! Router adapter from the shared, axum-agnostic
|
||||||
|
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
|
||||||
|
//! [`Response`], setting `Retry-After` when the envelope carries one.
|
||||||
|
//!
|
||||||
|
//! cortex-core owns the envelope shape; this is the only place the router
|
||||||
|
//! crosses from that data into axum. Mirrors cortex-gateway's adapter so
|
||||||
|
//! the router's own rejections (no feasible operator, all unreachable) are
|
||||||
|
//! the same #63-shaped envelopes clients already understand — distinct from
|
||||||
|
//! cortex's rejections, which the router proxies through verbatim.
|
||||||
|
|
||||||
|
use axum::http::{HeaderValue, StatusCode, header};
|
||||||
|
use axum::response::{IntoResponse, Json, Response};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
|
||||||
|
/// Render an [`OpenAiError`] as an axum response (status + JSON envelope +
|
||||||
|
/// optional `Retry-After`).
|
||||||
|
pub fn envelope_response(err: OpenAiError) -> Response {
|
||||||
|
let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||||
|
let retry_after = err.retry_after_secs;
|
||||||
|
let mut response = (status, Json(err.body())).into_response();
|
||||||
|
if let Some(secs) = retry_after
|
||||||
|
&& let Ok(value) = HeaderValue::from_str(&secs.to_string())
|
||||||
|
{
|
||||||
|
response.headers_mut().insert(header::RETRY_AFTER, value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
}
|
||||||
89
crates/helexa-router/src/handlers.rs
Normal file
89
crates/helexa-router/src/handlers.rs
Normal 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 }))
|
||||||
|
}
|
||||||
60
crates/helexa-router/src/lib.rs
Normal file
60
crates/helexa-router/src/lib.rs
Normal 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(())
|
||||||
|
}
|
||||||
52
crates/helexa-router/src/main.rs
Normal file
52
crates/helexa-router/src/main.rs
Normal 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(())
|
||||||
|
}
|
||||||
150
crates/helexa-router/src/poller.rs
Normal file
150
crates/helexa-router/src/poller.rs
Normal file
@@ -0,0 +1,150 @@
|
|||||||
|
//! Background poller that refreshes the multi-operator topology (#72).
|
||||||
|
//!
|
||||||
|
//! The same pattern as cortex↔neuron, one tier up: periodically poll each
|
||||||
|
//! configured cortex's `GET /v1/models` (catalogue × topology feasibility +
|
||||||
|
//! loaded state) and `GET /health` (coarse node-health/load), building the
|
||||||
|
//! live map the dispatcher (#73) routes on. An unreachable or erroring
|
||||||
|
//! cortex is debounced over [`POLL_FAILURE_THRESHOLD`] consecutive misses,
|
||||||
|
//! then flipped unhealthy and excluded from routing; it recovers on the
|
||||||
|
//! next successful poll.
|
||||||
|
|
||||||
|
use crate::state::RouterState;
|
||||||
|
use chrono::Utc;
|
||||||
|
use cortex_core::node::CortexModelEntry;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
/// Per-cortex HTTP timeout for each poll request.
|
||||||
|
const POLL_TIMEOUT: Duration = Duration::from_secs(5);
|
||||||
|
|
||||||
|
/// Consecutive failed polls before a cortex is marked unreachable. Mirrors
|
||||||
|
/// cortex's neuron-poll debounce: a single blip (a busy cortex briefly slow
|
||||||
|
/// to answer) can't yank it — and all its models — out of routing.
|
||||||
|
pub const POLL_FAILURE_THRESHOLD: u32 = 3;
|
||||||
|
|
||||||
|
/// cortex's `/v1/models` envelope — `{ "object": "list", "data": [...] }`.
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct ModelsEnvelope {
|
||||||
|
#[serde(default)]
|
||||||
|
data: Vec<CortexModelEntry>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The subset of cortex's `/health` the router reads.
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct CortexHealth {
|
||||||
|
nodes: CortexHealthNodes,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct CortexHealthNodes {
|
||||||
|
healthy: u32,
|
||||||
|
total: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run forever, polling all cortexes on the configured interval.
|
||||||
|
pub async fn poll_loop(state: std::sync::Arc<RouterState>) {
|
||||||
|
loop {
|
||||||
|
poll_once(&state).await;
|
||||||
|
tokio::time::sleep(state.poll_interval).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Poll every configured cortex once. Public for testing.
|
||||||
|
pub async fn poll_once(state: &RouterState) {
|
||||||
|
for cortex in &state.cortexes {
|
||||||
|
poll_cortex(state, &cortex.name, &cortex.endpoint).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Poll one cortex: refresh its model map from `/v1/models`, then its node
|
||||||
|
/// health from `/health`. A `/v1/models` failure debounces toward
|
||||||
|
/// unreachable; the `/health` poll is best-effort and never flips
|
||||||
|
/// reachability on its own (a cortex serving `/v1/models` is routable even
|
||||||
|
/// if `/health` momentarily isn't).
|
||||||
|
async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
|
||||||
|
// A cortex whose pinned TLS client failed to build (#74) is disabled:
|
||||||
|
// there is no client to poll with, so it stays unreachable.
|
||||||
|
let Some(client) = state.client_for(name) else {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
if let Some(entry) = topo.get_mut(name) {
|
||||||
|
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||||
|
entry.reachable = false;
|
||||||
|
}
|
||||||
|
tracing::warn!(cortex = name, "no TLS client (disabled); skipping poll");
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
let models = fetch_models(client, endpoint).await;
|
||||||
|
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
let Some(entry) = topo.get_mut(name) else {
|
||||||
|
return; // not a configured cortex (shouldn't happen)
|
||||||
|
};
|
||||||
|
|
||||||
|
match models {
|
||||||
|
Ok(models) => {
|
||||||
|
entry.models = models.into_iter().map(|m| (m.id.clone(), m)).collect();
|
||||||
|
entry.reachable = true;
|
||||||
|
entry.consecutive_failures = 0;
|
||||||
|
entry.last_poll = Some(Utc::now());
|
||||||
|
tracing::debug!(cortex = name, models = entry.models.len(), "poll ok");
|
||||||
|
}
|
||||||
|
Err(reason) => {
|
||||||
|
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||||
|
if entry.consecutive_failures >= POLL_FAILURE_THRESHOLD {
|
||||||
|
entry.reachable = false;
|
||||||
|
}
|
||||||
|
tracing::warn!(
|
||||||
|
cortex = name,
|
||||||
|
failures = entry.consecutive_failures,
|
||||||
|
reachable = entry.reachable,
|
||||||
|
reason,
|
||||||
|
"cortex poll failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
drop(topo);
|
||||||
|
|
||||||
|
// Best-effort health (node counts). Never flips reachability.
|
||||||
|
if let Some((healthy, total)) = fetch_health(client, endpoint).await {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
if let Some(entry) = topo.get_mut(name) {
|
||||||
|
entry.healthy_nodes = healthy;
|
||||||
|
entry.total_nodes = total;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET `/v1/models`, returning the parsed entries or a short failure reason.
|
||||||
|
async fn fetch_models(
|
||||||
|
client: &reqwest::Client,
|
||||||
|
endpoint: &str,
|
||||||
|
) -> Result<Vec<CortexModelEntry>, &'static str> {
|
||||||
|
let url = format!("{endpoint}/v1/models");
|
||||||
|
let resp = client
|
||||||
|
.get(&url)
|
||||||
|
.timeout(POLL_TIMEOUT)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|_| "unreachable")?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return Err("non-success status");
|
||||||
|
}
|
||||||
|
let envelope = resp
|
||||||
|
.json::<ModelsEnvelope>()
|
||||||
|
.await
|
||||||
|
.map_err(|_| "bad json")?;
|
||||||
|
Ok(envelope.data)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET `/health`, returning `(healthy, total)` node counts. `None` on any
|
||||||
|
/// failure — the caller leaves the previous counts in place.
|
||||||
|
async fn fetch_health(client: &reqwest::Client, endpoint: &str) -> Option<(u32, u32)> {
|
||||||
|
let url = format!("{endpoint}/health");
|
||||||
|
let resp = client.get(&url).timeout(POLL_TIMEOUT).send().await.ok()?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let health = resp.json::<CortexHealth>().await.ok()?;
|
||||||
|
Some((health.nodes.healthy, health.nodes.total))
|
||||||
|
}
|
||||||
144
crates/helexa-router/src/state.rs
Normal file
144
crates/helexa-router/src/state.rs
Normal file
@@ -0,0 +1,144 @@
|
|||||||
|
use crate::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use cortex_core::node::CortexModelEntry;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// Shared router state: 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>,
|
||||||
|
/// Per-cortex HTTP client, keyed by cortex name (#74). A cortex enrolled
|
||||||
|
/// with a `tls_ca` gets a client that trusts only that anchor; others
|
||||||
|
/// get a default client. A cortex whose `tls_ca` failed to load is
|
||||||
|
/// **absent** here — `client_for` returns `None` and it is never
|
||||||
|
/// polled or routed to (fail closed: a misconfigured pin must not
|
||||||
|
/// silently fall back to unpinned TLS).
|
||||||
|
clients: HashMap<String, reqwest::Client>,
|
||||||
|
/// This router instance's region, for dispatch geo affinity (#73).
|
||||||
|
pub region: Option<String>,
|
||||||
|
/// How often the poller refreshes the topology.
|
||||||
|
pub poll_interval: Duration,
|
||||||
|
/// Live per-cortex topology, keyed by cortex name. Pre-populated from
|
||||||
|
/// config (every configured cortex present, `reachable = false`) so the
|
||||||
|
/// poller and handlers always find an entry; the poller flips
|
||||||
|
/// reachability and fills the model map.
|
||||||
|
pub topology: RwLock<HashMap<String, CortexTopology>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Live view of one downstream cortex, refreshed each poll.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct CortexTopology {
|
||||||
|
/// Whether the cortex is currently routable. Flipped `false` only after
|
||||||
|
/// [`crate::poller::POLL_FAILURE_THRESHOLD`] consecutive failed polls
|
||||||
|
/// (debounces transient blips); restored on the next successful poll.
|
||||||
|
pub reachable: bool,
|
||||||
|
/// Consecutive failed polls; reset to 0 on success.
|
||||||
|
pub consecutive_failures: u32,
|
||||||
|
/// Timestamp of the last successful poll.
|
||||||
|
pub last_poll: Option<DateTime<Utc>>,
|
||||||
|
/// Healthy / total neuron counts from the cortex's `/health` (coarse
|
||||||
|
/// load signal; #73 refines headroom). 0/0 until first health poll.
|
||||||
|
pub healthy_nodes: u32,
|
||||||
|
pub total_nodes: u32,
|
||||||
|
/// The cortex's full `/v1/models` entries, keyed by model id. Stored
|
||||||
|
/// whole (not distilled to a loaded/feasible bool) so the federation
|
||||||
|
/// catalogue (#75) can preserve per-model `limit`/`cost`/capabilities.
|
||||||
|
pub models: HashMap<String, CortexModelEntry>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether a cortex can serve this model — loaded now, or feasible to
|
||||||
|
/// cold-load (its catalogue × topology says some neuron can host it).
|
||||||
|
pub fn entry_feasible(entry: &CortexModelEntry) -> bool {
|
||||||
|
entry.loaded || !entry.feasible_on.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RouterState {
|
||||||
|
pub fn from_config(config: &RouterConfig) -> Self {
|
||||||
|
let topology = config
|
||||||
|
.cortexes
|
||||||
|
.iter()
|
||||||
|
.map(|c| (c.name.clone(), CortexTopology::default()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// One client per cortex. A `tls_ca` that fails to load omits the
|
||||||
|
// cortex from the map (fail closed) rather than degrading to an
|
||||||
|
// unpinned client.
|
||||||
|
let mut clients = HashMap::new();
|
||||||
|
for c in &config.cortexes {
|
||||||
|
match build_client(c.tls_ca.as_deref()) {
|
||||||
|
Ok(client) => {
|
||||||
|
clients.insert(c.name.clone(), client);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(
|
||||||
|
cortex = %c.name,
|
||||||
|
tls_ca = c.tls_ca.as_deref().unwrap_or(""),
|
||||||
|
error = %e,
|
||||||
|
"failed to build pinned TLS client; cortex disabled (fail closed)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Self {
|
||||||
|
cortexes: config.cortexes.clone(),
|
||||||
|
clients,
|
||||||
|
region: config.router.region.clone(),
|
||||||
|
poll_interval: Duration::from_secs(config.router.poll_interval_secs),
|
||||||
|
topology: RwLock::new(topology),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The HTTP client to use for `name`, or `None` if the cortex is
|
||||||
|
/// disabled (its `tls_ca` failed to load). Callers must treat `None` as
|
||||||
|
/// "not routable / not pollable".
|
||||||
|
pub fn client_for(&self, name: &str) -> Option<&reqwest::Client> {
|
||||||
|
self.clients.get(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Names of reachable cortexes that can serve `model_id` (loaded or
|
||||||
|
/// feasible to cold-load). Groundwork for capacity-aware dispatch (#73);
|
||||||
|
/// unreachable cortexes are excluded by construction.
|
||||||
|
pub async fn cortexes_serving(&self, model_id: &str) -> Vec<String> {
|
||||||
|
let topo = self.topology.read().await;
|
||||||
|
topo.iter()
|
||||||
|
.filter(|(_, t)| t.reachable)
|
||||||
|
.filter(|(_, t)| t.models.get(model_id).is_some_and(entry_feasible))
|
||||||
|
.map(|(name, _)| name.clone())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a cortex HTTP client. With `tls_ca` set, the client trusts **only**
|
||||||
|
/// that PEM anchor (platform roots disabled) — pinning the router→cortex hop
|
||||||
|
/// to an enrolled cert (#74). Without it, standard platform-root validation.
|
||||||
|
pub fn build_client(tls_ca: Option<&str>) -> Result<reqwest::Client, BuildClientError> {
|
||||||
|
let mut builder = reqwest::Client::builder();
|
||||||
|
if let Some(path) = tls_ca {
|
||||||
|
let pem = std::fs::read(path).map_err(|e| BuildClientError::Read(path.to_string(), e))?;
|
||||||
|
let cert = reqwest::Certificate::from_pem(&pem).map_err(BuildClientError::Parse)?;
|
||||||
|
builder = builder
|
||||||
|
.tls_built_in_root_certs(false)
|
||||||
|
.add_root_certificate(cert);
|
||||||
|
}
|
||||||
|
builder.build().map_err(BuildClientError::Build)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Why a cortex's pinned client could not be built (→ cortex disabled).
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum BuildClientError {
|
||||||
|
#[error("reading TLS anchor '{0}'")]
|
||||||
|
Read(String, #[source] std::io::Error),
|
||||||
|
#[error("parsing TLS anchor PEM")]
|
||||||
|
Parse(#[source] reqwest::Error),
|
||||||
|
#[error("building HTTP client")]
|
||||||
|
Build(#[source] reqwest::Error),
|
||||||
|
}
|
||||||
132
crates/helexa-router/tests/catalogue.rs
Normal file
132
crates/helexa-router/tests/catalogue.rs
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
//! End-to-end federation-catalogue test for #75: poll two mock cortexes
|
||||||
|
//! that overlap on a model, then `GET /v1/models` on the router and verify
|
||||||
|
//! the deduped union with merged availability and preserved limit/cost.
|
||||||
|
|
||||||
|
use axum::Router;
|
||||||
|
use axum::routing::get;
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::poll_once;
|
||||||
|
use helexa_router::state::RouterState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Spawn a mock cortex serving the given `/v1/models` `data` array.
|
||||||
|
async fn spawn_cortex(models: Value) -> String {
|
||||||
|
let models = Arc::new(models);
|
||||||
|
let app = Router::new()
|
||||||
|
.route(
|
||||||
|
"/v1/models",
|
||||||
|
get({
|
||||||
|
let models = Arc::clone(&models);
|
||||||
|
move || {
|
||||||
|
let models = Arc::clone(&models);
|
||||||
|
async move { axum::Json(json!({ "object": "list", "data": &*models })) }
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/health",
|
||||||
|
get(|| async { axum::Json(json!({"status":"ok","nodes":{"healthy":1,"total":1}})) }),
|
||||||
|
);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn the router (with poller) wired to the given cortex endpoints, and
|
||||||
|
/// poll once synchronously so the topology is populated before we query.
|
||||||
|
async fn spawn_router(cortexes: Vec<CortexEndpoint>) -> String {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = Arc::new(RouterState::from_config(&cfg));
|
||||||
|
poll_once(&state).await; // deterministic: fill topology now
|
||||||
|
|
||||||
|
let app = helexa_router::build_app(Arc::clone(&state));
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model(id: &str, loaded: bool, feasible_on: &[&str], ctx: u64, input_cost: f64) -> Value {
|
||||||
|
json!({
|
||||||
|
"id": id,
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": loaded,
|
||||||
|
"feasible_on": feasible_on,
|
||||||
|
"locations": [],
|
||||||
|
"limit": { "context": ctx, "output": 4096 },
|
||||||
|
"cost": { "input": input_cost, "output": input_cost * 3.0 }
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn federation_catalogue_dedupes_and_preserves_limit_cost() {
|
||||||
|
// cortex A: "shared" loaded (ctx 32768, $0.50) + "only-a" loaded.
|
||||||
|
let a = spawn_cortex(json!([
|
||||||
|
model("shared", true, &["beast"], 32_768, 0.50),
|
||||||
|
model("only-a", true, &["beast"], 8_192, 1.00),
|
||||||
|
]))
|
||||||
|
.await;
|
||||||
|
// cortex B: "shared" cold-loadable, tighter ctx (16384), cheaper ($0.20).
|
||||||
|
let b = spawn_cortex(json!([model("shared", false, &["benjy"], 16_384, 0.20)])).await;
|
||||||
|
|
||||||
|
let router = spawn_router(vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "op-a".into(),
|
||||||
|
endpoint: a,
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "op-b".into(),
|
||||||
|
endpoint: b,
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
])
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let body: Value = reqwest::get(format!("{router}/v1/models"))
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(body["object"], "list");
|
||||||
|
let data = body["data"].as_array().unwrap();
|
||||||
|
// Deduped union: "shared" once + "only-a".
|
||||||
|
assert_eq!(data.len(), 2);
|
||||||
|
|
||||||
|
let shared = data.iter().find(|m| m["id"] == "shared").unwrap();
|
||||||
|
// Loaded somewhere (op-a) → loaded.
|
||||||
|
assert_eq!(shared["loaded"], true);
|
||||||
|
// feasible_on re-tiered to operator names, both present, sorted.
|
||||||
|
let feasible: Vec<&str> = shared["feasible_on"]
|
||||||
|
.as_array()
|
||||||
|
.unwrap()
|
||||||
|
.iter()
|
||||||
|
.map(|v| v.as_str().unwrap())
|
||||||
|
.collect();
|
||||||
|
assert_eq!(feasible, vec!["op-a", "op-b"]);
|
||||||
|
// Tightest limit (16384) and cheapest cost ($0.20) win.
|
||||||
|
assert_eq!(shared["limit"]["context"], 16_384);
|
||||||
|
assert_eq!(shared["cost"]["input"], 0.20);
|
||||||
|
// Loaded location named by operator, no neuron VRAM leaked.
|
||||||
|
let locs = shared["locations"].as_array().unwrap();
|
||||||
|
assert_eq!(locs.len(), 1);
|
||||||
|
assert_eq!(locs[0]["node"], "op-a");
|
||||||
|
|
||||||
|
assert!(data.iter().any(|m| m["id"] == "only-a"));
|
||||||
|
}
|
||||||
304
crates/helexa-router/tests/dispatch.rs
Normal file
304
crates/helexa-router/tests/dispatch.rs
Normal file
@@ -0,0 +1,304 @@
|
|||||||
|
//! 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,
|
||||||
|
max_model_len: None,
|
||||||
|
max_input_tokens: None,
|
||||||
|
max_output_tokens: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const MODEL: &str = "Qwen/Qwen3-Coder-30B";
|
||||||
|
|
||||||
|
// ── Mock cortex backend ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Behaviour of a mock cortex, carried in axum State.
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct MockCortex {
|
||||||
|
/// Identifies which cortex answered, echoed in the 200 body.
|
||||||
|
name: &'static str,
|
||||||
|
/// When true, return a genuine #63-shaped `429 rate_limit_exceeded`.
|
||||||
|
rate_limited: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap) -> Response {
|
||||||
|
if m.rate_limited {
|
||||||
|
return (
|
||||||
|
StatusCode::TOO_MANY_REQUESTS,
|
||||||
|
Json(json!({"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down","param":null}})),
|
||||||
|
)
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
let auth = headers
|
||||||
|
.get("authorization")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
Json(json!({ "served_by": m.name, "auth_seen": auth })).into_response()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_cortex(mock: MockCortex) -> String {
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/v1/chat/completions", post(mock_handler))
|
||||||
|
.with_state(mock);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ok_cortex(name: &'static str) -> MockCortex {
|
||||||
|
MockCortex {
|
||||||
|
name,
|
||||||
|
rate_limited: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Helpers to build state with a hand-set topology ──────────────────
|
||||||
|
|
||||||
|
fn state_with(cortexes: Vec<CortexEndpoint>, region: Option<String>) -> RouterState {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let mut state = RouterState::from_config(&cfg);
|
||||||
|
state.region = region;
|
||||||
|
state
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Overwrite the topology entry for `name` so tests control reachability and
|
||||||
|
/// model serveability directly (no live poll).
|
||||||
|
async fn set_topology(
|
||||||
|
state: &RouterState,
|
||||||
|
name: &str,
|
||||||
|
reachable: bool,
|
||||||
|
loaded: bool,
|
||||||
|
feasible: bool,
|
||||||
|
healthy_nodes: u32,
|
||||||
|
) {
|
||||||
|
let mut topo = state.topology.write().await;
|
||||||
|
let mut models = HashMap::new();
|
||||||
|
models.insert(MODEL.to_string(), model_entry(loaded, feasible));
|
||||||
|
topo.insert(
|
||||||
|
name.to_string(),
|
||||||
|
CortexTopology {
|
||||||
|
reachable,
|
||||||
|
consecutive_failures: 0,
|
||||||
|
last_poll: None,
|
||||||
|
healthy_nodes,
|
||||||
|
total_nodes: healthy_nodes,
|
||||||
|
models,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ep(name: &str, endpoint: &str, region: Option<&str>) -> CortexEndpoint {
|
||||||
|
CortexEndpoint {
|
||||||
|
name: name.into(),
|
||||||
|
endpoint: endpoint.into(),
|
||||||
|
region: region.map(str::to_string),
|
||||||
|
tls_ca: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn chat_body() -> Bytes {
|
||||||
|
Bytes::from(format!("{{\"model\":\"{MODEL}\",\"stream\":false}}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn body_json(resp: Response) -> (StatusCode, Value) {
|
||||||
|
let status = resp.status();
|
||||||
|
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let v = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
|
||||||
|
(status, v)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Tests ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn routes_to_serving_cortex_and_forwards_bearer() {
|
||||||
|
let url = spawn_cortex(ok_cortex("c1")).await;
|
||||||
|
let state = state_with(vec![ep("c1", &url, None)], None);
|
||||||
|
set_topology(&state, "c1", true, true, true, 2).await;
|
||||||
|
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert("authorization", "Bearer sk-test-123".parse().unwrap());
|
||||||
|
|
||||||
|
let resp = dispatch(&state, "/v1/chat/completions", headers, chat_body()).await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
assert_eq!(status, StatusCode::OK);
|
||||||
|
assert_eq!(body["served_by"], "c1");
|
||||||
|
// Bearer reached the cortex unchanged.
|
||||||
|
assert_eq!(body["auth_seen"], "Bearer sk-test-123");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cortex_429_passes_through_and_is_not_retried() {
|
||||||
|
// c1 (ranked first: loaded) returns a genuine 429; c2 would return 200.
|
||||||
|
let c1 = spawn_cortex(MockCortex {
|
||||||
|
name: "c1",
|
||||||
|
rate_limited: true,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
let c2 = spawn_cortex(ok_cortex("c2")).await;
|
||||||
|
let state = state_with(vec![ep("c1", &c1, None), ep("c2", &c2, None)], None);
|
||||||
|
// Both reachable + loaded; c1 has more headroom so it ranks first.
|
||||||
|
set_topology(&state, "c1", true, true, true, 5).await;
|
||||||
|
set_topology(&state, "c2", true, true, true, 1).await;
|
||||||
|
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
// The genuine 4xx is returned verbatim — NOT retried to c2.
|
||||||
|
assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||||
|
assert!(body.get("served_by").is_none(), "must not have hit c2");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn fails_over_to_next_cortex_on_transport_error() {
|
||||||
|
// c_dead ranks first (more headroom) but its endpoint is a closed port;
|
||||||
|
// c_live is the fallback. The router must fail over and c_live serves.
|
||||||
|
let live = spawn_cortex(ok_cortex("c_live")).await;
|
||||||
|
let state = state_with(
|
||||||
|
vec![
|
||||||
|
ep("c_dead", "http://127.0.0.1:1", None),
|
||||||
|
ep("c_live", &live, None),
|
||||||
|
],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
set_topology(&state, "c_dead", true, true, true, 9).await;
|
||||||
|
set_topology(&state, "c_live", true, true, true, 1).await;
|
||||||
|
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
|
||||||
|
assert_eq!(status, StatusCode::OK);
|
||||||
|
assert_eq!(body["served_by"], "c_live");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unknown_model_is_404() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
// Topology has no entry for MODEL at all.
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||||
|
assert_eq!(body["error"]["code"], "model_not_found");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn known_but_all_unreachable_is_503() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
// Cortex knows the model (from a prior good poll) but is now unreachable.
|
||||||
|
set_topology(&state, "c1", false, true, true, 2).await;
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
chat_body(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
assert_eq!(body["error"]["code"], "service_unavailable");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_model_field_is_400() {
|
||||||
|
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||||
|
let resp = dispatch(
|
||||||
|
&state,
|
||||||
|
"/v1/chat/completions",
|
||||||
|
HeaderMap::new(),
|
||||||
|
Bytes::from_static(b"{\"messages\":[]}"),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let (status, body) = body_json(resp).await;
|
||||||
|
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||||
|
assert_eq!(body["error"]["code"], "missing_model_field");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ranking_prefers_loaded_then_region_then_headroom() {
|
||||||
|
// Router is in eu-west. Candidates:
|
||||||
|
// warm-eu : loaded, region match, 1 node → best
|
||||||
|
// warm-us : loaded, no region, 9 nodes
|
||||||
|
// cold-eu : feasible only, region match → worst (cold)
|
||||||
|
let state = state_with(
|
||||||
|
vec![
|
||||||
|
ep("warm-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||||
|
ep("warm-us", "http://127.0.0.1:1", Some("us-east")),
|
||||||
|
ep("cold-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||||
|
],
|
||||||
|
Some("eu-west".into()),
|
||||||
|
);
|
||||||
|
set_topology(&state, "warm-eu", true, true, true, 1).await;
|
||||||
|
set_topology(&state, "warm-us", true, true, true, 9).await;
|
||||||
|
set_topology(&state, "cold-eu", true, false, true, 5).await;
|
||||||
|
|
||||||
|
let Selection::Candidates(order) = select_cortexes(&state, MODEL).await else {
|
||||||
|
panic!("expected candidates");
|
||||||
|
};
|
||||||
|
let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
|
||||||
|
assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
|
||||||
|
}
|
||||||
97
crates/helexa-router/tests/skeleton.rs
Normal file
97
crates/helexa-router/tests/skeleton.rs
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
//! 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,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "b".into(),
|
||||||
|
endpoint: "https://b.example.com".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: 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(())
|
||||||
|
});
|
||||||
|
}
|
||||||
210
crates/helexa-router/tests/tls.rs
Normal file
210
crates/helexa-router/tests/tls.rs
Normal file
@@ -0,0 +1,210 @@
|
|||||||
|
//! Outbound TLS pinning tests for #74.
|
||||||
|
//!
|
||||||
|
//! Proves the router, as a TLS client to cortexes, reaches a cortex
|
||||||
|
//! presenting its **enrolled** cert and rejects one presenting an
|
||||||
|
//! unexpected (or untrusted) cert — and that a rejected handshake flows
|
||||||
|
//! through the existing reachability path (#72) to exclude the cortex.
|
||||||
|
//!
|
||||||
|
//! A minimal `tokio-rustls` HTTPS server presents a self-signed cert; the
|
||||||
|
//! router's `reqwest` client (native-tls) validates against the PEM anchor
|
||||||
|
//! enrolled in config. Server (rustls) and client (native-tls) interoperate
|
||||||
|
//! at the protocol level — what matters is the trust decision.
|
||||||
|
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::poll_once;
|
||||||
|
use helexa_router::state::{RouterState, build_client};
|
||||||
|
use std::io::Write;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tokio_rustls::TlsAcceptor;
|
||||||
|
|
||||||
|
/// A self-signed cert: PEM (for the reqwest pin file) + DER cert/key (for
|
||||||
|
/// the rustls server).
|
||||||
|
struct TestCert {
|
||||||
|
cert_pem: String,
|
||||||
|
cert_der: rustls::pki_types::CertificateDer<'static>,
|
||||||
|
key_der: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_cert() -> TestCert {
|
||||||
|
let key = rcgen::generate_simple_self_signed(vec!["127.0.0.1".to_string()]).unwrap();
|
||||||
|
TestCert {
|
||||||
|
cert_pem: key.cert.pem(),
|
||||||
|
cert_der: key.cert.der().clone(),
|
||||||
|
key_der: key.key_pair.serialize_der(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write a cert PEM to a unique temp file (named by `tag`) and return the
|
||||||
|
/// path. `tag` is caller-unique (we use the bound port), so no randomness.
|
||||||
|
fn write_pem(tag: &str, pem: &str) -> String {
|
||||||
|
let path = std::env::temp_dir().join(format!("helexa-router-tls-{tag}.pem"));
|
||||||
|
let mut f = std::fs::File::create(&path).unwrap();
|
||||||
|
f.write_all(pem.as_bytes()).unwrap();
|
||||||
|
path.to_string_lossy().into_owned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a minimal HTTPS server presenting `cert`, answering every request
|
||||||
|
/// with a canned `/v1/models`-shaped 200. Returns its `https://` base URL.
|
||||||
|
async fn spawn_https(cert: &TestCert) -> String {
|
||||||
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||||
|
|
||||||
|
let key = rustls::pki_types::PrivateKeyDer::Pkcs8(rustls::pki_types::PrivatePkcs8KeyDer::from(
|
||||||
|
cert.key_der.clone(),
|
||||||
|
));
|
||||||
|
let config = rustls::ServerConfig::builder()
|
||||||
|
.with_no_client_auth()
|
||||||
|
.with_single_cert(vec![cert.cert_der.clone()], key)
|
||||||
|
.unwrap();
|
||||||
|
let acceptor = TlsAcceptor::from(Arc::new(config));
|
||||||
|
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
let Ok((stream, _)) = listener.accept().await else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let acceptor = acceptor.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Ok(mut tls) = acceptor.accept(stream).await {
|
||||||
|
let mut buf = [0u8; 2048];
|
||||||
|
let _ = tls.read(&mut buf).await; // consume request line/headers
|
||||||
|
let body = "{\"object\":\"list\",\"data\":[]}";
|
||||||
|
let resp = format!(
|
||||||
|
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||||
|
body.len(),
|
||||||
|
body
|
||||||
|
);
|
||||||
|
let _ = tls.write_all(resp.as_bytes()).await;
|
||||||
|
let _ = tls.shutdown().await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
format!("https://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tag_for(url: &str) -> String {
|
||||||
|
url.rsplit(':').next().unwrap_or("0").to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn pinned_client_accepts_enrolled_cert_and_rejects_others() {
|
||||||
|
let server_cert = make_cert();
|
||||||
|
let other_cert = make_cert();
|
||||||
|
let url = spawn_https(&server_cert).await;
|
||||||
|
let tag = tag_for(&url);
|
||||||
|
|
||||||
|
let good_pin = write_pem(&format!("{tag}-good"), &server_cert.cert_pem);
|
||||||
|
let bad_pin = write_pem(&format!("{tag}-bad"), &other_cert.cert_pem);
|
||||||
|
|
||||||
|
// Enrolled with the server's own cert → handshake trusted → 200.
|
||||||
|
let good = build_client(Some(&good_pin)).unwrap();
|
||||||
|
let resp = good.get(format!("{url}/v1/models")).send().await;
|
||||||
|
assert!(resp.is_ok(), "enrolled cert must be accepted: {resp:?}");
|
||||||
|
assert_eq!(resp.unwrap().status(), 200);
|
||||||
|
|
||||||
|
// Enrolled with a different cert → server's cert is unexpected → reject.
|
||||||
|
let bad = build_client(Some(&bad_pin)).unwrap();
|
||||||
|
assert!(
|
||||||
|
bad.get(format!("{url}/v1/models")).send().await.is_err(),
|
||||||
|
"unexpected cert must be rejected"
|
||||||
|
);
|
||||||
|
|
||||||
|
// No enrollment (default platform roots) → self-signed cert untrusted.
|
||||||
|
let default = build_client(None).unwrap();
|
||||||
|
assert!(
|
||||||
|
default
|
||||||
|
.get(format!("{url}/v1/models"))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.is_err(),
|
||||||
|
"un-enrolled self-signed cert must be rejected by default roots"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn poller_excludes_cortex_with_unexpected_cert() {
|
||||||
|
let server_cert = make_cert();
|
||||||
|
let other_cert = make_cert();
|
||||||
|
let url = spawn_https(&server_cert).await;
|
||||||
|
let tag = tag_for(&url);
|
||||||
|
|
||||||
|
let good_pin = write_pem(&format!("{tag}-pgood"), &server_cert.cert_pem);
|
||||||
|
let bad_pin = write_pem(&format!("{tag}-pbad"), &other_cert.cert_pem);
|
||||||
|
|
||||||
|
// Cortex A enrolled correctly → reachable. Cortex B enrolled with the
|
||||||
|
// wrong cert → TLS handshake fails → excluded.
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "good".into(),
|
||||||
|
endpoint: url.clone(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some(good_pin),
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "bad".into(),
|
||||||
|
endpoint: url.clone(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some(bad_pin),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = RouterState::from_config(&cfg);
|
||||||
|
poll_once(&state).await;
|
||||||
|
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(
|
||||||
|
topo["good"].reachable,
|
||||||
|
"correctly-enrolled cortex reachable"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!topo["bad"].reachable,
|
||||||
|
"cortex presenting an unexpected cert is excluded"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn misconfigured_pin_disables_cortex_fail_closed() {
|
||||||
|
// A `tls_ca` pointing at a nonexistent file must NOT fall back to an
|
||||||
|
// unpinned client — the cortex is disabled entirely.
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "broken".into(),
|
||||||
|
endpoint: "https://127.0.0.1:1".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: Some("/no/such/anchor.pem".into()),
|
||||||
|
},
|
||||||
|
CortexEndpoint {
|
||||||
|
name: "plain".into(),
|
||||||
|
endpoint: "http://127.0.0.1:1".into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let state = RouterState::from_config(&cfg);
|
||||||
|
assert!(
|
||||||
|
state.client_for("broken").is_none(),
|
||||||
|
"a cortex with an unloadable pin is disabled (fail closed)"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
state.client_for("plain").is_some(),
|
||||||
|
"an un-pinned cortex still gets a client"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn build_client_rejects_garbage_pem() {
|
||||||
|
let path = write_pem(
|
||||||
|
"garbage",
|
||||||
|
"-----BEGIN CERTIFICATE-----\nnope\n-----END CERTIFICATE-----",
|
||||||
|
);
|
||||||
|
assert!(build_client(Some(&path)).is_err());
|
||||||
|
}
|
||||||
172
crates/helexa-router/tests/topology.rs
Normal file
172
crates/helexa-router/tests/topology.rs
Normal file
@@ -0,0 +1,172 @@
|
|||||||
|
//! Topology-poller acceptance tests for #72: the router maintains a live
|
||||||
|
//! map of which cortexes serve which models, marks an unreachable/erroring
|
||||||
|
//! cortex unhealthy and excludes it from routing, and recovers it once
|
||||||
|
//! reachable again.
|
||||||
|
|
||||||
|
use axum::extract::State;
|
||||||
|
use axum::http::StatusCode;
|
||||||
|
use axum::routing::get;
|
||||||
|
use axum::{Json, Router};
|
||||||
|
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||||
|
use helexa_router::poller::{POLL_FAILURE_THRESHOLD, poll_once};
|
||||||
|
use helexa_router::state::{RouterState, entry_feasible};
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Shared "is this mock cortex up?" flag, toggled by tests to simulate
|
||||||
|
/// outage and recovery.
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct MockState {
|
||||||
|
up: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_models(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||||
|
if !s.up.load(Ordering::SeqCst) {
|
||||||
|
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
}
|
||||||
|
Ok(Json(json!({
|
||||||
|
"object": "list",
|
||||||
|
"data": [
|
||||||
|
{
|
||||||
|
"id": "Qwen/Qwen3-Coder-30B",
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": true,
|
||||||
|
"feasible_on": ["beast"],
|
||||||
|
"locations": [{"node": "beast", "status": "loaded", "vram_estimate_mb": 19000}]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": "Qwen/Qwen3-VL-8B",
|
||||||
|
"object": "model",
|
||||||
|
"created": 0,
|
||||||
|
"owned_by": "helexa",
|
||||||
|
"loaded": false,
|
||||||
|
"feasible_on": ["beast"],
|
||||||
|
"locations": []
|
||||||
|
}
|
||||||
|
]
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mock_health(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||||
|
if !s.up.load(Ordering::SeqCst) {
|
||||||
|
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
}
|
||||||
|
Ok(Json(json!({
|
||||||
|
"status": "ok",
|
||||||
|
"nodes": { "healthy": 2, "total": 3 }
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawn a mock cortex; returns (base_url, up_flag).
|
||||||
|
async fn spawn_mock_cortex() -> (String, Arc<AtomicBool>) {
|
||||||
|
let up = Arc::new(AtomicBool::new(true));
|
||||||
|
let state = MockState { up: up.clone() };
|
||||||
|
let app = Router::new()
|
||||||
|
.route("/v1/models", get(mock_models))
|
||||||
|
.route("/health", get(mock_health))
|
||||||
|
.with_state(state);
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
(format!("http://{addr}"), up)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state_for(name: &str, endpoint: &str) -> RouterState {
|
||||||
|
let cfg = RouterConfig {
|
||||||
|
cortexes: vec![CortexEndpoint {
|
||||||
|
name: name.into(),
|
||||||
|
endpoint: endpoint.into(),
|
||||||
|
region: None,
|
||||||
|
tls_ca: None,
|
||||||
|
}],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
RouterState::from_config(&cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn poll_builds_live_topology() {
|
||||||
|
let (base, _up) = spawn_mock_cortex().await;
|
||||||
|
let state = state_for("c1", &base);
|
||||||
|
|
||||||
|
poll_once(&state).await;
|
||||||
|
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
let c1 = topo.get("c1").expect("cortex present");
|
||||||
|
assert!(c1.reachable, "should be reachable after a good poll");
|
||||||
|
assert_eq!(c1.consecutive_failures, 0);
|
||||||
|
assert!(c1.last_poll.is_some());
|
||||||
|
assert_eq!((c1.healthy_nodes, c1.total_nodes), (2, 3));
|
||||||
|
|
||||||
|
// Loaded model: loaded + feasible. Catalogue-only model: feasible only
|
||||||
|
// (not loaded, but feasible_on non-empty).
|
||||||
|
let coder = c1.models.get("Qwen/Qwen3-Coder-30B").unwrap();
|
||||||
|
assert!(coder.loaded && entry_feasible(coder));
|
||||||
|
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
|
||||||
|
assert!(!vl.loaded && entry_feasible(vl));
|
||||||
|
drop(topo);
|
||||||
|
|
||||||
|
// The routing helper sees both serveable models on the reachable cortex.
|
||||||
|
assert_eq!(
|
||||||
|
state.cortexes_serving("Qwen/Qwen3-VL-8B").await,
|
||||||
|
vec!["c1".to_string()]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unreachable_cortex_excluded_then_recovers() {
|
||||||
|
let (base, up) = spawn_mock_cortex().await;
|
||||||
|
let state = state_for("c1", &base);
|
||||||
|
|
||||||
|
// Healthy first.
|
||||||
|
poll_once(&state).await;
|
||||||
|
assert!(state.topology.read().await["c1"].reachable);
|
||||||
|
|
||||||
|
// Take it down. The first failures debounce (stay reachable) until the
|
||||||
|
// threshold; only then is it excluded.
|
||||||
|
up.store(false, Ordering::SeqCst);
|
||||||
|
for i in 1..POLL_FAILURE_THRESHOLD {
|
||||||
|
poll_once(&state).await;
|
||||||
|
assert!(
|
||||||
|
state.topology.read().await["c1"].reachable,
|
||||||
|
"still reachable after {i} failure(s) (below threshold)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
poll_once(&state).await; // crosses the threshold
|
||||||
|
{
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(!topo["c1"].reachable, "excluded after threshold failures");
|
||||||
|
assert!(topo["c1"].consecutive_failures >= POLL_FAILURE_THRESHOLD);
|
||||||
|
}
|
||||||
|
// Excluded from routing.
|
||||||
|
assert!(
|
||||||
|
state
|
||||||
|
.cortexes_serving("Qwen/Qwen3-Coder-30B")
|
||||||
|
.await
|
||||||
|
.is_empty()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Bring it back: the next successful poll restores it.
|
||||||
|
up.store(true, Ordering::SeqCst);
|
||||||
|
poll_once(&state).await;
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(topo["c1"].reachable, "recovered after a good poll");
|
||||||
|
assert_eq!(topo["c1"].consecutive_failures, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unconfigured_endpoint_is_unreachable() {
|
||||||
|
// Nothing listening on this port → polls fail; below threshold it stays
|
||||||
|
// at its initial unreachable state, and never panics.
|
||||||
|
let state = state_for("dead", "http://127.0.0.1:1");
|
||||||
|
poll_once(&state).await;
|
||||||
|
let topo = state.topology.read().await;
|
||||||
|
assert!(!topo["dead"].reachable);
|
||||||
|
assert_eq!(topo["dead"].consecutive_failures, 1);
|
||||||
|
}
|
||||||
21
crates/helexa-stream/Cargo.toml
Normal file
21
crates/helexa-stream/Cargo.toml
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
[package]
|
||||||
|
name = "helexa-stream"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
repository.workspace = true
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
name = "helexa_stream"
|
||||||
|
path = "src/lib.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
axum = { workspace = true }
|
||||||
|
reqwest = { workspace = true }
|
||||||
|
futures = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
tokio = { workspace = true }
|
||||||
|
tokio-stream = { workspace = true }
|
||||||
|
async-stream = "0.3"
|
||||||
290
crates/helexa-stream/src/lib.rs
Normal file
290
crates/helexa-stream/src/lib.rs
Normal file
@@ -0,0 +1,290 @@
|
|||||||
|
//! Shared streaming reverse-proxy mechanism (#71).
|
||||||
|
//!
|
||||||
|
//! cortex and helexa-router both need to proxy an OpenAI/Anthropic SSE
|
||||||
|
//! response from a downstream backend **verbatim** — chunks forwarded as
|
||||||
|
//! they arrive, never buffering the full body — while observing the bytes
|
||||||
|
//! for metrics/metering. This crate owns that mechanism so there is one
|
||||||
|
//! implementation, not one per tier.
|
||||||
|
//!
|
||||||
|
//! The split is mechanism vs policy:
|
||||||
|
//!
|
||||||
|
//! - **Mechanism (here):** [`forward_streaming`] POSTs to a backend and
|
||||||
|
//! streams the response body back through an [`ObservedStream`], which
|
||||||
|
//! feeds every chunk to a caller-supplied [`ChunkObserver`] and calls
|
||||||
|
//! [`ChunkObserver::finish`] exactly once on clean end-of-stream or on
|
||||||
|
//! drop (client disconnect mid-stream). [`BodyTail`] and
|
||||||
|
//! [`last_count_for`] are the reusable pieces an observer uses to pull
|
||||||
|
//! the trailing OpenAI `usage` object out of the streamed bytes.
|
||||||
|
//! - **Policy (caller):** what to *do* with the observed bytes — which
|
||||||
|
//! metric names to emit, which labels, whether to settle a per-principal
|
||||||
|
//! reservation — lives in the consumer's `ChunkObserver` impl, not here.
|
||||||
|
//!
|
||||||
|
//! The proxy is status-agnostic: a non-2xx upstream response (e.g. a
|
||||||
|
//! cortex `429 rate_limit_exceeded`) is streamed back with its status and
|
||||||
|
//! headers intact, so honest backpressure reaches the client unchanged.
|
||||||
|
//! Only a network failure or a malformed response build is an error.
|
||||||
|
|
||||||
|
use axum::body::{Body, Bytes};
|
||||||
|
use axum::http::{HeaderMap, StatusCode};
|
||||||
|
use axum::response::Response;
|
||||||
|
use futures::Stream;
|
||||||
|
use futures::stream::BoxStream;
|
||||||
|
use reqwest::Client;
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
|
/// Observes the bytes of a streamed proxy response without altering them.
|
||||||
|
///
|
||||||
|
/// `observe` is called for each forwarded chunk; `finish` is called
|
||||||
|
/// exactly once — on clean end-of-stream or on drop — and implementations
|
||||||
|
/// must be idempotent (the [`ObservedStream`] guards against a double call,
|
||||||
|
/// but a `finish` that runs side effects should still self-guard).
|
||||||
|
pub trait ChunkObserver: Send + Unpin + 'static {
|
||||||
|
/// A body chunk has been forwarded downstream. The slice is the exact
|
||||||
|
/// bytes the client receives.
|
||||||
|
fn observe(&mut self, chunk: &[u8]);
|
||||||
|
|
||||||
|
/// The stream has ended (cleanly or via client disconnect). Called once.
|
||||||
|
fn finish(&mut self);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A bounded accumulator for the tail of a streamed body.
|
||||||
|
///
|
||||||
|
/// The OpenAI `usage` object rides on the final SSE chunk (and sits at the
|
||||||
|
/// end of a non-streaming JSON body), so retaining a generous tail is
|
||||||
|
/// enough to recover token counts via [`last_count_for`]; the cap bounds
|
||||||
|
/// memory on huge bodies. Appends are char-boundary-safe.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct BodyTail {
|
||||||
|
tail: String,
|
||||||
|
cap: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BodyTail {
|
||||||
|
/// Create a tail retaining at most `cap` bytes.
|
||||||
|
pub fn new(cap: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
tail: String::new(),
|
||||||
|
cap,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a chunk, trimming from the front past the cap. When trimming,
|
||||||
|
/// the newest half is kept (the usage object is always at the very end).
|
||||||
|
pub fn push(&mut self, chunk: &[u8]) {
|
||||||
|
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||||
|
if self.tail.len() > self.cap {
|
||||||
|
let mut cut = self.tail.len() - self.cap / 2;
|
||||||
|
while !self.tail.is_char_boundary(cut) {
|
||||||
|
cut += 1;
|
||||||
|
}
|
||||||
|
self.tail.drain(..cut);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The retained tail text.
|
||||||
|
pub fn as_str(&self) -> &str {
|
||||||
|
&self.tail
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
||||||
|
///
|
||||||
|
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
||||||
|
/// The quoted-needle form means `completion_tokens` never matches
|
||||||
|
/// `completion_tokens_details`, and taking the last occurrence means the
|
||||||
|
/// final `usage` object wins even if content earlier in the stream echoed
|
||||||
|
/// a usage-shaped string.
|
||||||
|
pub fn last_count_for(tail: &str, key: &str) -> Option<u64> {
|
||||||
|
let needle = format!("\"{key}\"");
|
||||||
|
let mut result = None;
|
||||||
|
for (idx, _) in tail.match_indices(&needle) {
|
||||||
|
let rest = tail[idx + needle.len()..].trim_start();
|
||||||
|
let Some(rest) = rest.strip_prefix(':') else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let rest = rest.trim_start();
|
||||||
|
let digits: &str = &rest[..rest
|
||||||
|
.char_indices()
|
||||||
|
.find(|(_, c)| !c.is_ascii_digit())
|
||||||
|
.map(|(i, _)| i)
|
||||||
|
.unwrap_or(rest.len())];
|
||||||
|
if let Ok(v) = digits.parse::<u64>() {
|
||||||
|
result = Some(v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Error from [`forward_streaming`]. Distinguishes a network/transport
|
||||||
|
/// failure reaching the backend from a failure assembling the downstream
|
||||||
|
/// response. A non-2xx upstream *status* is not an error — it is streamed
|
||||||
|
/// through verbatim.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum StreamError {
|
||||||
|
#[error("upstream request failed")]
|
||||||
|
Upstream(reqwest::Error),
|
||||||
|
#[error("failed to build response")]
|
||||||
|
ResponseBuild(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST `body` to `url` and stream the response back verbatim through
|
||||||
|
/// `observer`.
|
||||||
|
///
|
||||||
|
/// Request headers are forwarded except `host` / `content-length` (reqwest
|
||||||
|
/// sets these). The returned [`Response`] carries the upstream status and
|
||||||
|
/// headers unchanged — including non-2xx — with a body that streams the
|
||||||
|
/// upstream bytes chunk-for-chunk, feeding each chunk to `observer`.
|
||||||
|
pub async fn forward_streaming<O: ChunkObserver>(
|
||||||
|
client: &Client,
|
||||||
|
url: &str,
|
||||||
|
headers: HeaderMap,
|
||||||
|
body: Bytes,
|
||||||
|
observer: O,
|
||||||
|
) -> Result<Response, StreamError> {
|
||||||
|
let mut req_builder = client.post(url).body(body);
|
||||||
|
for (key, value) in headers.iter() {
|
||||||
|
if key == "host" || key == "content-length" {
|
||||||
|
continue; // reqwest sets these
|
||||||
|
}
|
||||||
|
req_builder = req_builder.header(key, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
let upstream = req_builder.send().await.map_err(StreamError::Upstream)?;
|
||||||
|
|
||||||
|
let status =
|
||||||
|
StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||||
|
let resp_headers = upstream.headers().clone();
|
||||||
|
|
||||||
|
let stream = ObservedStream::new(Box::pin(upstream.bytes_stream()), observer);
|
||||||
|
let body = Body::from_stream(stream);
|
||||||
|
|
||||||
|
let mut response = Response::builder().status(status);
|
||||||
|
for (key, value) in resp_headers.iter() {
|
||||||
|
response = response.header(key, value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
.body(body)
|
||||||
|
.map_err(|e| StreamError::ResponseBuild(e.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pass-through stream wrapper that feeds a [`ChunkObserver`]. Forwards
|
||||||
|
/// each chunk verbatim, calls `observe` per chunk, and `finish` once on
|
||||||
|
/// clean end-of-stream; the `Drop` impl covers client disconnects.
|
||||||
|
pub struct ObservedStream<O: ChunkObserver> {
|
||||||
|
inner: BoxStream<'static, Result<Bytes, reqwest::Error>>,
|
||||||
|
observer: O,
|
||||||
|
finished: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> ObservedStream<O> {
|
||||||
|
/// Wrap a byte stream with an observer.
|
||||||
|
pub fn new(inner: BoxStream<'static, Result<Bytes, reqwest::Error>>, observer: O) -> Self {
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
observer,
|
||||||
|
finished: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish(&mut self) {
|
||||||
|
if self.finished {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.finished = true;
|
||||||
|
self.observer.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> Stream for ObservedStream<O> {
|
||||||
|
type Item = Result<Bytes, reqwest::Error>;
|
||||||
|
|
||||||
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
match this.inner.as_mut().poll_next(cx) {
|
||||||
|
Poll::Ready(Some(Ok(chunk))) => {
|
||||||
|
this.observer.observe(&chunk);
|
||||||
|
Poll::Ready(Some(Ok(chunk)))
|
||||||
|
}
|
||||||
|
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
||||||
|
Poll::Ready(None) => {
|
||||||
|
this.finish();
|
||||||
|
Poll::Ready(None)
|
||||||
|
}
|
||||||
|
Poll::Pending => Poll::Pending,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<O: ChunkObserver> Drop for ObservedStream<O> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extracts_counts_from_final_sse_usage_chunk() {
|
||||||
|
let tail = concat!(
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
||||||
|
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
||||||
|
"data: [DONE]\n\n"
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extracts_counts_from_non_streaming_body() {
|
||||||
|
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
||||||
|
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
||||||
|
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ignores_details_variants_and_takes_last_occurrence() {
|
||||||
|
// completion_tokens_details must not shadow completion_tokens,
|
||||||
|
// and the LAST usage object wins (matters when content echoes
|
||||||
|
// a usage-shaped string earlier in the stream).
|
||||||
|
let tail = concat!(
|
||||||
|
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
||||||
|
"data: {\"usage\":{\"completion_tokens\":99,",
|
||||||
|
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn absent_keys_yield_none() {
|
||||||
|
assert_eq!(
|
||||||
|
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
||||||
|
// key present but non-numeric value
|
||||||
|
assert_eq!(
|
||||||
|
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn body_tail_retains_usage_after_cap_trim() {
|
||||||
|
// Cap small enough that the filler forces several front-trims, but
|
||||||
|
// (as in production, where cap ≫ the usage object) large enough that
|
||||||
|
// the trailing usage object survives the newest-half retention.
|
||||||
|
let mut tail = BodyTail::new(512);
|
||||||
|
for _ in 0..100 {
|
||||||
|
tail.push(b"data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n");
|
||||||
|
}
|
||||||
|
assert!(tail.as_str().len() <= 512, "cap must bound the tail");
|
||||||
|
tail.push(b"data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":9}}\n\n");
|
||||||
|
assert_eq!(last_count_for(tail.as_str(), "prompt_tokens"), Some(5));
|
||||||
|
assert_eq!(last_count_for(tail.as_str(), "completion_tokens"), Some(9));
|
||||||
|
}
|
||||||
|
}
|
||||||
162
crates/helexa-stream/tests/streaming.rs
Normal file
162
crates/helexa-stream/tests/streaming.rs
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
//! Integration tests for the shared streaming proxy (#71): proves a backend
|
||||||
|
//! SSE response is forwarded chunk-for-chunk (no buffering), the observer
|
||||||
|
//! sees every byte and finishes once, and non-2xx is streamed through with
|
||||||
|
//! its status intact — the behaviours both cortex and helexa-router rely on.
|
||||||
|
|
||||||
|
use axum::Router;
|
||||||
|
use axum::body::Body;
|
||||||
|
use axum::http::{HeaderMap, StatusCode};
|
||||||
|
use axum::response::Response;
|
||||||
|
use axum::routing::post;
|
||||||
|
use helexa_stream::{BodyTail, ChunkObserver, forward_streaming, last_count_for};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
|
/// Observer that records what it saw, for assertions.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct RecordingObserver {
|
||||||
|
inner: Arc<Mutex<Recorded>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct Recorded {
|
||||||
|
chunks: usize,
|
||||||
|
finished: usize,
|
||||||
|
tail: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkObserver for RecordingObserver {
|
||||||
|
fn observe(&mut self, chunk: &[u8]) {
|
||||||
|
let mut r = self.inner.lock().unwrap();
|
||||||
|
r.chunks += 1;
|
||||||
|
r.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||||
|
}
|
||||||
|
fn finish(&mut self) {
|
||||||
|
self.inner.lock().unwrap().finished += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mock backend that streams 5 SSE chunks with 30ms gaps, then a usage
|
||||||
|
/// chunk and `[DONE]`.
|
||||||
|
async fn sse_handler() -> Response {
|
||||||
|
let chunks: Vec<&'static str> = vec![
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"d\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"e\"}}]}\n\n",
|
||||||
|
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}}\n\n",
|
||||||
|
"data: [DONE]\n\n",
|
||||||
|
];
|
||||||
|
let stream = async_stream::stream! {
|
||||||
|
for c in chunks {
|
||||||
|
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||||
|
yield Ok::<_, std::io::Error>(axum::body::Bytes::from_static(c.as_bytes()));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Response::new(Body::from_stream(stream))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn rate_limited_handler() -> Response {
|
||||||
|
Response::builder()
|
||||||
|
.status(StatusCode::TOO_MANY_REQUESTS)
|
||||||
|
.body(Body::from("{\"error\":{\"type\":\"rate_limit_exceeded\"}}"))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn spawn_backend(router: Router) -> String {
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, router).await.unwrap();
|
||||||
|
});
|
||||||
|
format!("http://{addr}")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn streams_chunks_incrementally_and_observes_usage() {
|
||||||
|
let base = spawn_backend(Router::new().route("/v1/chat/completions", post(sse_handler))).await;
|
||||||
|
let observer = RecordingObserver::default();
|
||||||
|
let probe = observer.clone();
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = forward_streaming(
|
||||||
|
&client,
|
||||||
|
&format!("{base}/v1/chat/completions"),
|
||||||
|
HeaderMap::new(),
|
||||||
|
axum::body::Bytes::from_static(b"{\"model\":\"x\",\"stream\":true}"),
|
||||||
|
observer,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("forward ok");
|
||||||
|
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
// Read the proxied body as a stream, timestamping arrivals.
|
||||||
|
let mut body = resp.into_body().into_data_stream();
|
||||||
|
let mut arrivals: Vec<Instant> = Vec::new();
|
||||||
|
let mut collected = String::new();
|
||||||
|
use futures::StreamExt;
|
||||||
|
while let Some(item) = body.next().await {
|
||||||
|
let bytes = item.unwrap();
|
||||||
|
arrivals.push(Instant::now());
|
||||||
|
collected.push_str(&String::from_utf8_lossy(&bytes));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Incremental delivery: first and last chunk are meaningfully apart
|
||||||
|
// (5×30ms gaps), proving no full-response buffering.
|
||||||
|
let spread = *arrivals.last().unwrap() - arrivals[0];
|
||||||
|
assert!(
|
||||||
|
spread >= Duration::from_millis(100),
|
||||||
|
"expected incremental delivery, spread was {spread:?}"
|
||||||
|
);
|
||||||
|
|
||||||
|
// The client received the terminator and the usage object verbatim.
|
||||||
|
assert!(collected.contains("data: [DONE]"));
|
||||||
|
|
||||||
|
// The observer saw the bytes and finished exactly once.
|
||||||
|
let r = probe.inner.lock().unwrap();
|
||||||
|
assert!(r.chunks >= 5, "observer saw {} chunks", r.chunks);
|
||||||
|
assert_eq!(r.finished, 1, "finish must run exactly once");
|
||||||
|
assert_eq!(last_count_for(&r.tail, "prompt_tokens"), Some(11));
|
||||||
|
assert_eq!(last_count_for(&r.tail, "completion_tokens"), Some(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn non_2xx_is_streamed_through_verbatim() {
|
||||||
|
let base =
|
||||||
|
spawn_backend(Router::new().route("/v1/chat/completions", post(rate_limited_handler)))
|
||||||
|
.await;
|
||||||
|
let observer = RecordingObserver::default();
|
||||||
|
let probe = observer.clone();
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = forward_streaming(
|
||||||
|
&client,
|
||||||
|
&format!("{base}/v1/chat/completions"),
|
||||||
|
HeaderMap::new(),
|
||||||
|
axum::body::Bytes::new(),
|
||||||
|
observer,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("forward ok");
|
||||||
|
|
||||||
|
// Backpressure status reaches the client unchanged.
|
||||||
|
assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert!(String::from_utf8_lossy(&body).contains("rate_limit_exceeded"));
|
||||||
|
|
||||||
|
// finish still runs once even with a tiny non-streaming body.
|
||||||
|
assert_eq!(probe.inner.lock().unwrap().finished, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn body_tail_smoke() {
|
||||||
|
let mut tail = BodyTail::new(128);
|
||||||
|
tail.push(b"hello ");
|
||||||
|
tail.push(b"world");
|
||||||
|
assert_eq!(tail.as_str(), "hello world");
|
||||||
|
}
|
||||||
63
crates/helexa-upstream/Cargo.toml
Normal file
63
crates/helexa-upstream/Cargo.toml
Normal file
@@ -0,0 +1,63 @@
|
|||||||
|
[package]
|
||||||
|
name = "helexa-upstream"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
repository.workspace = true
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "helexa-upstream"
|
||||||
|
path = "src/main.rs"
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
name = "helexa_upstream"
|
||||||
|
path = "src/lib.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
tokio = { workspace = true }
|
||||||
|
axum = { workspace = true }
|
||||||
|
tower-http = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
figment = { workspace = true }
|
||||||
|
anyhow = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
clap = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
tracing-subscriber = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|
||||||
|
# PostgreSQL — the mesh authority's system of record. Runtime query API
|
||||||
|
# (not the compile-time `query!` macros) so the crate builds in CI without a
|
||||||
|
# live database or a committed `.sqlx` offline cache; correctness is covered
|
||||||
|
# by the gated integration tests. (Macro adoption is a later refinement once
|
||||||
|
# a dev DB + offline cache exist.)
|
||||||
|
sqlx = { version = "0.8", default-features = false, features = [
|
||||||
|
"runtime-tokio",
|
||||||
|
"tls-rustls",
|
||||||
|
"postgres",
|
||||||
|
"macros",
|
||||||
|
"migrate",
|
||||||
|
"chrono",
|
||||||
|
"uuid",
|
||||||
|
] }
|
||||||
|
uuid = { version = "1", features = ["v4", "serde"] }
|
||||||
|
sha2 = "0.10"
|
||||||
|
subtle = "2.6"
|
||||||
|
# Web auth (B4): argon2id password hashing, JWT sessions, CSPRNG secrets,
|
||||||
|
# transactional email.
|
||||||
|
argon2 = "0.5"
|
||||||
|
jsonwebtoken = "9"
|
||||||
|
rand = "0.8"
|
||||||
|
lettre = { version = "0.11", default-features = false, features = [
|
||||||
|
"tokio1-rustls-tls",
|
||||||
|
"smtp-transport",
|
||||||
|
"builder",
|
||||||
|
] }
|
||||||
|
|
||||||
|
# cortex-core for the shared #63 OpenAiError envelope on the authz surface.
|
||||||
|
cortex-core = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
figment = { workspace = true, features = ["test"] }
|
||||||
|
reqwest = { workspace = true }
|
||||||
137
crates/helexa-upstream/migrations/0001_init.sql
Normal file
137
crates/helexa-upstream/migrations/0001_init.sql
Normal file
@@ -0,0 +1,137 @@
|
|||||||
|
-- helexa-upstream initial schema (#59): accounts, keys, ledger, top-up
|
||||||
|
-- codes, served-usage. The mesh-level authority's system of record.
|
||||||
|
--
|
||||||
|
-- Token amounts are BIGINT (i64) throughout; the cortex EntitlementProvider
|
||||||
|
-- carries u64 but mesh allocations sit comfortably inside i64 and Postgres
|
||||||
|
-- has no unsigned type.
|
||||||
|
|
||||||
|
CREATE EXTENSION IF NOT EXISTS citext;
|
||||||
|
CREATE EXTENSION IF NOT EXISTS pgcrypto;
|
||||||
|
|
||||||
|
-- ── Users (web auth: email + password) ──────────────────────────────
|
||||||
|
CREATE TABLE users (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
email CITEXT NOT NULL UNIQUE,
|
||||||
|
password_hash TEXT NOT NULL, -- argon2id PHC string
|
||||||
|
email_verified BOOLEAN NOT NULL DEFAULT FALSE,
|
||||||
|
-- Browser fingerprint captured at registration (#abuse). Best-effort,
|
||||||
|
-- client-supplied; the primary signal for silent multi-account
|
||||||
|
-- detection. NULL when the client could not produce one.
|
||||||
|
registration_fingerprint TEXT,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||||
|
);
|
||||||
|
CREATE INDEX users_registration_fingerprint_idx
|
||||||
|
ON users (registration_fingerprint)
|
||||||
|
WHERE registration_fingerprint IS NOT NULL;
|
||||||
|
|
||||||
|
-- Single-use email tokens for verification and password reset. Only the
|
||||||
|
-- sha256 of the emailed secret is stored.
|
||||||
|
CREATE TABLE email_tokens (
|
||||||
|
token_hash BYTEA PRIMARY KEY,
|
||||||
|
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
|
kind TEXT NOT NULL CHECK (kind IN ('verify', 'reset')),
|
||||||
|
expires_at TIMESTAMPTZ NOT NULL,
|
||||||
|
consumed_at TIMESTAMPTZ
|
||||||
|
);
|
||||||
|
CREATE INDEX email_tokens_user_idx ON email_tokens (user_id);
|
||||||
|
|
||||||
|
-- ── Accounts (the billable allocation ledger) ───────────────────────
|
||||||
|
CREATE TABLE accounts (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
owner_user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
|
allocation_total BIGINT NOT NULL DEFAULT 0,
|
||||||
|
allocation_spent BIGINT NOT NULL DEFAULT 0,
|
||||||
|
allocation_reserved BIGINT NOT NULL DEFAULT 0,
|
||||||
|
-- 'deactivated' is the SILENT abuse flag: keys stop authorizing but no
|
||||||
|
-- surface ever tells the user why (see resolve → 401).
|
||||||
|
status TEXT NOT NULL DEFAULT 'active'
|
||||||
|
CHECK (status IN ('active', 'deactivated')),
|
||||||
|
-- This account shares a registration fingerprint with >= 1 other.
|
||||||
|
fingerprint_flagged BOOLEAN NOT NULL DEFAULT FALSE,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||||
|
-- The no-overshoot backstop to the atomic reserve UPDATE.
|
||||||
|
CONSTRAINT accounts_no_overshoot
|
||||||
|
CHECK (allocation_spent + allocation_reserved <= allocation_total),
|
||||||
|
CONSTRAINT accounts_nonneg
|
||||||
|
CHECK (allocation_spent >= 0 AND allocation_reserved >= 0)
|
||||||
|
);
|
||||||
|
CREATE INDEX accounts_owner_idx ON accounts (owner_user_id);
|
||||||
|
|
||||||
|
-- ── API keys (Principal.key_id = api_keys.id) ───────────────────────
|
||||||
|
CREATE TABLE api_keys (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
|
||||||
|
key_hash BYTEA NOT NULL, -- sha256(raw key)
|
||||||
|
key_prefix TEXT NOT NULL, -- non-secret display prefix
|
||||||
|
label TEXT NOT NULL DEFAULT '',
|
||||||
|
status TEXT NOT NULL DEFAULT 'active'
|
||||||
|
CHECK (status IN ('active', 'archived')),
|
||||||
|
-- Per-key sub-cap: 'hardcap' = absolute tokens; 'percent' = % of the
|
||||||
|
-- account's allocation_total (resolved to an absolute at reserve time).
|
||||||
|
limit_kind TEXT NOT NULL DEFAULT 'percent'
|
||||||
|
CHECK (limit_kind IN ('percent', 'hardcap')),
|
||||||
|
limit_value BIGINT NOT NULL DEFAULT 100,
|
||||||
|
-- serde of cortex_core::entitlements::CapWindow (Balance | Rolling).
|
||||||
|
cap_window JSONB NOT NULL DEFAULT '{"kind":"balance"}'::jsonb,
|
||||||
|
-- Per-key running ledger (mirrors the account ledger; Balance semantics
|
||||||
|
-- in this migration — rolling-window reset lands with the authz API).
|
||||||
|
key_spent BIGINT NOT NULL DEFAULT 0,
|
||||||
|
key_reserved BIGINT NOT NULL DEFAULT 0,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||||
|
CONSTRAINT api_keys_key_nonneg
|
||||||
|
CHECK (key_spent >= 0 AND key_reserved >= 0)
|
||||||
|
);
|
||||||
|
-- A raw key resolves only while active; the hash is unique among active keys.
|
||||||
|
CREATE UNIQUE INDEX api_keys_active_hash_idx
|
||||||
|
ON api_keys (key_hash) WHERE status = 'active';
|
||||||
|
CREATE INDEX api_keys_account_idx ON api_keys (account_id);
|
||||||
|
|
||||||
|
-- ── Reservations (reserve → settle/release) ─────────────────────────
|
||||||
|
-- id is BIGSERIAL so it maps to the cortex Reservation.id (u64) verbatim,
|
||||||
|
-- with the Postgres sequence as the sole global authority.
|
||||||
|
CREATE TABLE reservations (
|
||||||
|
id BIGSERIAL PRIMARY KEY,
|
||||||
|
account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
|
||||||
|
key_id UUID NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
|
||||||
|
reserved BIGINT NOT NULL,
|
||||||
|
actual BIGINT,
|
||||||
|
state TEXT NOT NULL DEFAULT 'open'
|
||||||
|
CHECK (state IN ('open', 'settled', 'released')),
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||||
|
settled_at TIMESTAMPTZ
|
||||||
|
);
|
||||||
|
-- The sweeper scans open reservations by age.
|
||||||
|
CREATE INDEX reservations_open_idx
|
||||||
|
ON reservations (created_at) WHERE state = 'open';
|
||||||
|
|
||||||
|
-- ── Top-up codes (hybrid allocation) ────────────────────────────────
|
||||||
|
CREATE TABLE top_up_codes (
|
||||||
|
code_hash BYTEA PRIMARY KEY, -- sha256(raw code)
|
||||||
|
value BIGINT NOT NULL, -- tokens this code grants
|
||||||
|
denomination TEXT, -- human label (e.g. "small")
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||||
|
redeemed_by UUID REFERENCES accounts(id) ON DELETE SET NULL,
|
||||||
|
redeemed_at TIMESTAMPTZ
|
||||||
|
);
|
||||||
|
|
||||||
|
-- ── Served-usage ledger (#58 reconciliation) ────────────────────────
|
||||||
|
-- Absolute per-(operator, account, key, period) served tokens, upserted by
|
||||||
|
-- each cortex; reconciliation rolls these up for operator compensation.
|
||||||
|
CREATE TABLE served_usage (
|
||||||
|
operator_id TEXT NOT NULL,
|
||||||
|
account_id UUID NOT NULL,
|
||||||
|
key_id UUID NOT NULL,
|
||||||
|
period DATE NOT NULL,
|
||||||
|
served_tokens BIGINT NOT NULL DEFAULT 0,
|
||||||
|
reconciled_at TIMESTAMPTZ,
|
||||||
|
PRIMARY KEY (operator_id, account_id, key_id, period)
|
||||||
|
);
|
||||||
|
|
||||||
|
-- ── Web sessions (DB-backed; alt/complement to stateless JWT) ───────
|
||||||
|
CREATE TABLE sessions (
|
||||||
|
token_hash BYTEA PRIMARY KEY, -- sha256(session token)
|
||||||
|
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
|
expires_at TIMESTAMPTZ NOT NULL,
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||||
|
);
|
||||||
|
CREATE INDEX sessions_user_idx ON sessions (user_id);
|
||||||
340
crates/helexa-upstream/src/authz.rs
Normal file
340
crates/helexa-upstream/src/authz.rs
Normal file
@@ -0,0 +1,340 @@
|
|||||||
|
//! `/authz/v1` — the machine surface cortex's `UpstreamEntitlementProvider`
|
||||||
|
//! (#57) consumes. It mirrors the `cortex_core::entitlements::EntitlementProvider`
|
||||||
|
//! trait 1:1 (resolve / reserve / settle / release / snapshot) over the B1
|
||||||
|
//! ledger.
|
||||||
|
//!
|
||||||
|
//! Contract notes for the cortex client:
|
||||||
|
//! - A **non-2xx** response means the authority could not give an
|
||||||
|
//! authoritative answer (bad caller auth, malformed request, server
|
||||||
|
//! error) → the client should **fail closed**.
|
||||||
|
//! - `reserve` returns **200** whether granted or budget-refused: the body
|
||||||
|
//! carries either `reservation_id` or a `rejected` discriminant. A budget
|
||||||
|
//! refusal is an authoritative answer, not a transport failure.
|
||||||
|
//! - Rejections that are genuinely auth failures use the #63 `OpenAiError`
|
||||||
|
//! envelope so they can be surfaced verbatim.
|
||||||
|
|
||||||
|
use crate::crypto::sha256;
|
||||||
|
use crate::error::envelope_response;
|
||||||
|
use crate::ledger::{self, LedgerError};
|
||||||
|
use crate::state::AppState;
|
||||||
|
use axum::extract::{Request, State};
|
||||||
|
use axum::http::{StatusCode, header};
|
||||||
|
use axum::middleware::Next;
|
||||||
|
use axum::response::{IntoResponse, Response};
|
||||||
|
use axum::routing::post;
|
||||||
|
use axum::{Extension, Json, Router};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use subtle::ConstantTimeEq;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// The operator a validated client bearer identifies (served-usage
|
||||||
|
/// attribution, #58). Inserted into request extensions by [`client_auth`].
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct OperatorId(pub String);
|
||||||
|
|
||||||
|
/// Build the `/authz/v1` router with the client-auth layer applied.
|
||||||
|
pub fn router(state: &AppState) -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route("/authz/v1/resolve", post(resolve))
|
||||||
|
.route("/authz/v1/reserve", post(reserve))
|
||||||
|
.route("/authz/v1/settle", post(settle))
|
||||||
|
.route("/authz/v1/release", post(release))
|
||||||
|
.route("/authz/v1/snapshot", post(snapshot))
|
||||||
|
.route("/authz/v1/served-usage", post(served_usage))
|
||||||
|
.layer(axum::middleware::from_fn_with_state(
|
||||||
|
state.clone(),
|
||||||
|
client_auth,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── client auth (shared bearer → operator_id) ───────────────────────
|
||||||
|
|
||||||
|
/// Validate the caller's `Authorization: Bearer` against the configured
|
||||||
|
/// client tokens (constant-time) and stamp the `operator_id`. When no tokens
|
||||||
|
/// are configured the surface is open (dev) and a synthetic operator is
|
||||||
|
/// used.
|
||||||
|
async fn client_auth(State(state): State<AppState>, mut req: Request, next: Next) -> Response {
|
||||||
|
let tokens = &state.config.client_auth.tokens;
|
||||||
|
if tokens.is_empty() {
|
||||||
|
req.extensions_mut().insert(OperatorId("dev".into()));
|
||||||
|
return next.run(req).await;
|
||||||
|
}
|
||||||
|
let presented = req
|
||||||
|
.headers()
|
||||||
|
.get(header::AUTHORIZATION)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.and_then(|v| v.strip_prefix("Bearer "))
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or("");
|
||||||
|
let matched = tokens
|
||||||
|
.iter()
|
||||||
|
.find(|t| t.token.as_bytes().ct_eq(presented.as_bytes()).into());
|
||||||
|
match matched {
|
||||||
|
Some(t) => {
|
||||||
|
req.extensions_mut()
|
||||||
|
.insert(OperatorId(t.operator_id.clone()));
|
||||||
|
next.run(req).await
|
||||||
|
}
|
||||||
|
None => envelope_response(OpenAiError::invalid_api_key(
|
||||||
|
"missing or invalid client credentials",
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── DTOs ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ResolveReq {
|
||||||
|
api_key: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct PrincipalDto {
|
||||||
|
account_id: String,
|
||||||
|
key_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct SnapshotDto {
|
||||||
|
hard_cap: Option<i64>,
|
||||||
|
spent: i64,
|
||||||
|
reserved: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
struct ResolveResp {
|
||||||
|
principal: PrincipalDto,
|
||||||
|
snapshot: SnapshotDto,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ReserveReq {
|
||||||
|
account_id: String,
|
||||||
|
key_id: String,
|
||||||
|
max_tokens: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Default)]
|
||||||
|
struct ReserveResp {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
reservation_id: Option<i64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
rejected: Option<Rejection>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||||
|
enum Rejection {
|
||||||
|
InsufficientQuota {
|
||||||
|
requested: i64,
|
||||||
|
available: i64,
|
||||||
|
},
|
||||||
|
// Part of the frozen wire contract so the cortex client (#57) can map it
|
||||||
|
// without a later breaking change. Not yet constructed: the B1 ledger
|
||||||
|
// implements Balance caps only; rolling-window key sub-caps (which yield
|
||||||
|
// this) land in a follow-up.
|
||||||
|
#[allow(dead_code)]
|
||||||
|
RateLimited {
|
||||||
|
requested: i64,
|
||||||
|
available: i64,
|
||||||
|
retry_after_secs: u64,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct SettleReq {
|
||||||
|
reservation_id: i64,
|
||||||
|
actual_tokens: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ReservationRef {
|
||||||
|
reservation_id: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct SnapshotReq {
|
||||||
|
account_id: String,
|
||||||
|
key_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── handlers ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// `POST /authz/v1/resolve` — bearer key → principal + snapshot, or
|
||||||
|
/// `401 invalid_api_key` (also for a deactivated account: no clue).
|
||||||
|
async fn resolve(State(state): State<AppState>, Json(req): Json<ResolveReq>) -> Response {
|
||||||
|
match ledger::resolve_key(&state.pool, &sha256(&req.api_key)).await {
|
||||||
|
Ok(Some(p)) => Json(ResolveResp {
|
||||||
|
principal: PrincipalDto {
|
||||||
|
account_id: p.account_id.to_string(),
|
||||||
|
key_id: p.key_id.to_string(),
|
||||||
|
},
|
||||||
|
snapshot: SnapshotDto {
|
||||||
|
hard_cap: Some(p.hard_cap),
|
||||||
|
spent: p.key_spent,
|
||||||
|
reserved: p.key_reserved,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
.into_response(),
|
||||||
|
Ok(None) => envelope_response(OpenAiError::invalid_api_key("invalid or unknown API key")),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(error = %e, "resolve query failed");
|
||||||
|
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /authz/v1/reserve` — 200 with `reservation_id` (granted) or
|
||||||
|
/// `rejected` (budget). Non-2xx only for bad input / server error.
|
||||||
|
async fn reserve(State(state): State<AppState>, Json(req): Json<ReserveReq>) -> Response {
|
||||||
|
let (Ok(account_id), Ok(key_id)) = (
|
||||||
|
Uuid::parse_str(&req.account_id),
|
||||||
|
Uuid::parse_str(&req.key_id),
|
||||||
|
) else {
|
||||||
|
return bad_request("account_id and key_id must be UUIDs");
|
||||||
|
};
|
||||||
|
match ledger::reserve(&state.pool, account_id, key_id, req.max_tokens).await {
|
||||||
|
Ok(reservation_id) => Json(ReserveResp {
|
||||||
|
reservation_id: Some(reservation_id),
|
||||||
|
rejected: None,
|
||||||
|
})
|
||||||
|
.into_response(),
|
||||||
|
Err(LedgerError::InsufficientQuota {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
}) => Json(ReserveResp {
|
||||||
|
reservation_id: None,
|
||||||
|
rejected: Some(Rejection::InsufficientQuota {
|
||||||
|
requested,
|
||||||
|
available,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
.into_response(),
|
||||||
|
Err(LedgerError::AccountNotFound | LedgerError::KeyNotFound) => {
|
||||||
|
// Resolve succeeded earlier; the principal vanished (archived /
|
||||||
|
// deactivated). Treat as no budget — fail closed at the client.
|
||||||
|
Json(ReserveResp {
|
||||||
|
reservation_id: None,
|
||||||
|
rejected: Some(Rejection::InsufficientQuota {
|
||||||
|
requested: req.max_tokens,
|
||||||
|
available: 0,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
Err(LedgerError::Db(e)) => {
|
||||||
|
tracing::error!(error = %e, "reserve failed");
|
||||||
|
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /authz/v1/settle` — idempotent; `204`.
|
||||||
|
async fn settle(State(state): State<AppState>, Json(req): Json<SettleReq>) -> Response {
|
||||||
|
match ledger::settle(&state.pool, req.reservation_id, req.actual_tokens).await {
|
||||||
|
Ok(()) => StatusCode::NO_CONTENT.into_response(),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(error = %e, "settle failed");
|
||||||
|
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /authz/v1/release` — idempotent; `204`.
|
||||||
|
async fn release(State(state): State<AppState>, Json(req): Json<ReservationRef>) -> Response {
|
||||||
|
match ledger::release(&state.pool, req.reservation_id).await {
|
||||||
|
Ok(()) => StatusCode::NO_CONTENT.into_response(),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(error = %e, "release failed");
|
||||||
|
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /authz/v1/snapshot` — `{hard_cap, spent, reserved}` or `404`.
|
||||||
|
async fn snapshot(State(state): State<AppState>, Json(req): Json<SnapshotReq>) -> Response {
|
||||||
|
let (Ok(account_id), Ok(key_id)) = (
|
||||||
|
Uuid::parse_str(&req.account_id),
|
||||||
|
Uuid::parse_str(&req.key_id),
|
||||||
|
) else {
|
||||||
|
return bad_request("account_id and key_id must be UUIDs");
|
||||||
|
};
|
||||||
|
match ledger::snapshot(&state.pool, account_id, key_id).await {
|
||||||
|
Ok(Some((hard_cap, spent, reserved))) => Json(SnapshotDto {
|
||||||
|
hard_cap: Some(hard_cap),
|
||||||
|
spent,
|
||||||
|
reserved,
|
||||||
|
})
|
||||||
|
.into_response(),
|
||||||
|
Ok(None) => StatusCode::NOT_FOUND.into_response(),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(error = %e, "snapshot failed");
|
||||||
|
envelope_response(OpenAiError::service_unavailable("authority error", Some(5)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── served-usage report (#58) ───────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ServedUsageReport {
|
||||||
|
rows: Vec<ServedUsageRow>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ServedUsageRow {
|
||||||
|
account_id: String,
|
||||||
|
key_id: String,
|
||||||
|
period: String, // YYYY-MM-DD
|
||||||
|
served_tokens: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /authz/v1/served-usage` — a cortex reports the absolute served-token
|
||||||
|
/// counters it has accrued for the current period. Upsert is monotonic
|
||||||
|
/// (`GREATEST`) so re-sends and races are idempotent and never regress.
|
||||||
|
/// `operator_id` comes from the validated client bearer (request extension).
|
||||||
|
async fn served_usage(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(operator): Extension<OperatorId>,
|
||||||
|
Json(req): Json<ServedUsageReport>,
|
||||||
|
) -> Response {
|
||||||
|
for row in &req.rows {
|
||||||
|
let (Ok(account_id), Ok(key_id)) = (
|
||||||
|
Uuid::parse_str(&row.account_id),
|
||||||
|
Uuid::parse_str(&row.key_id),
|
||||||
|
) else {
|
||||||
|
continue; // skip malformed ids rather than fail the whole batch
|
||||||
|
};
|
||||||
|
let Ok(period) = chrono::NaiveDate::parse_from_str(&row.period, "%Y-%m-%d") else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let res = sqlx::query(
|
||||||
|
"INSERT INTO served_usage (operator_id, account_id, key_id, period, served_tokens) \
|
||||||
|
VALUES ($1, $2, $3, $4, $5) \
|
||||||
|
ON CONFLICT (operator_id, account_id, key_id, period) \
|
||||||
|
DO UPDATE SET served_tokens = GREATEST(served_usage.served_tokens, EXCLUDED.served_tokens)",
|
||||||
|
)
|
||||||
|
.bind(&operator.0)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(period)
|
||||||
|
.bind(row.served_tokens.max(0))
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await;
|
||||||
|
if let Err(e) = res {
|
||||||
|
tracing::error!(error = %e, "served-usage upsert failed");
|
||||||
|
return envelope_response(OpenAiError::service_unavailable("authority error", Some(5)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
StatusCode::NO_CONTENT.into_response()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bad_request(msg: &str) -> Response {
|
||||||
|
envelope_response(OpenAiError::new(
|
||||||
|
400,
|
||||||
|
"invalid_request_error",
|
||||||
|
"invalid_request",
|
||||||
|
msg,
|
||||||
|
))
|
||||||
|
}
|
||||||
261
crates/helexa-upstream/src/config.rs
Normal file
261
crates/helexa-upstream/src/config.rs
Normal file
@@ -0,0 +1,261 @@
|
|||||||
|
//! helexa-upstream configuration: loaded from `helexa-upstream.toml` with
|
||||||
|
//! figment, `UPSTREAM_`-prefixed env overrides (mirrors the cortex/router
|
||||||
|
//! convention, e.g. `UPSTREAM_SERVER__LISTEN`, `UPSTREAM_DB__URL`).
|
||||||
|
|
||||||
|
use figment::{
|
||||||
|
Figment,
|
||||||
|
providers::{Env, Format, Toml},
|
||||||
|
};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct UpstreamConfig {
|
||||||
|
#[serde(default)]
|
||||||
|
pub server: ServerSettings,
|
||||||
|
pub db: DbSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub grant: GrantSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub abuse: AbuseSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub client_auth: ClientAuthSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub authz: AuthzSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub auth: AuthSettings,
|
||||||
|
#[serde(default)]
|
||||||
|
pub email: EmailSettings,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[auth]` — web-session signing + token lifetimes (B4). Web sessions are
|
||||||
|
/// JWTs, distinct from inference API keys.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AuthSettings {
|
||||||
|
/// HMAC secret for signing session JWTs. MUST be overridden in prod
|
||||||
|
/// (env `UPSTREAM_AUTH__JWT_SECRET`); the default is dev-only.
|
||||||
|
#[serde(default = "default_jwt_secret")]
|
||||||
|
pub jwt_secret: String,
|
||||||
|
/// Session token lifetime (seconds).
|
||||||
|
#[serde(default = "default_session_ttl")]
|
||||||
|
pub session_ttl_secs: u64,
|
||||||
|
/// Email verification / password-reset token lifetime (seconds).
|
||||||
|
#[serde(default = "default_email_token_ttl")]
|
||||||
|
pub email_token_ttl_secs: u64,
|
||||||
|
/// Public base URL of the frontend, used to build verify/reset links.
|
||||||
|
#[serde(default = "default_app_base_url")]
|
||||||
|
pub app_base_url: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AuthSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
jwt_secret: default_jwt_secret(),
|
||||||
|
session_ttl_secs: default_session_ttl(),
|
||||||
|
email_token_ttl_secs: default_email_token_ttl(),
|
||||||
|
app_base_url: default_app_base_url(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[email]` — transactional email transport for verify/reset.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct EmailSettings {
|
||||||
|
/// `"log"` (dev — logs the link) or `"smtp"`.
|
||||||
|
#[serde(default = "default_email_provider")]
|
||||||
|
pub provider: String,
|
||||||
|
/// SMTP relay URL (e.g. "smtp://user:pass@host:587") when provider=smtp.
|
||||||
|
#[serde(default)]
|
||||||
|
pub smtp_url: Option<String>,
|
||||||
|
/// `From:` address.
|
||||||
|
#[serde(default = "default_from_addr")]
|
||||||
|
pub from_addr: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for EmailSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
provider: default_email_provider(),
|
||||||
|
smtp_url: None,
|
||||||
|
from_addr: default_from_addr(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[client_auth]` — credentials operators' cortexes present to `/authz/v1`.
|
||||||
|
/// Each token maps to an `operator_id` (served-usage attribution, #58). This
|
||||||
|
/// transport credential is distinct from end-user API keys (which ride in
|
||||||
|
/// the `resolve` body). v2 adds mTLS.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ClientAuthSettings {
|
||||||
|
/// When empty the authz surface is **open** (dev only; logged at warn).
|
||||||
|
#[serde(default)]
|
||||||
|
pub tokens: Vec<ClientToken>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ClientToken {
|
||||||
|
/// Shared bearer a cortex presents.
|
||||||
|
pub token: String,
|
||||||
|
/// Operator this token identifies.
|
||||||
|
pub operator_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[authz]` — reservation lifecycle knobs.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AuthzSettings {
|
||||||
|
/// Open reservations older than this are swept (released), self-healing
|
||||||
|
/// a reservation whose settle/release from cortex was lost.
|
||||||
|
#[serde(default = "default_reservation_ttl")]
|
||||||
|
pub reservation_ttl_secs: u64,
|
||||||
|
/// How often the sweeper runs.
|
||||||
|
#[serde(default = "default_sweep_interval")]
|
||||||
|
pub sweep_interval_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AuthzSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
reservation_ttl_secs: default_reservation_ttl(),
|
||||||
|
sweep_interval_secs: default_sweep_interval(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ServerSettings {
|
||||||
|
/// Address to listen on (e.g. "0.0.0.0:8090"). Plaintext — edge nginx
|
||||||
|
/// terminates TLS, consistent with the rest of the stack.
|
||||||
|
#[serde(default = "default_listen")]
|
||||||
|
pub listen: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ServerSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
listen: default_listen(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct DbSettings {
|
||||||
|
/// PostgreSQL connection URL (e.g. "postgres://user:pass@host/helexa").
|
||||||
|
pub url: String,
|
||||||
|
/// Max pool connections.
|
||||||
|
#[serde(default = "default_max_connections")]
|
||||||
|
pub max_connections: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[grant]` — the flat free token grant every email-verified account
|
||||||
|
/// receives (the floor of the hybrid allocation model; top-up codes extend
|
||||||
|
/// it).
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct GrantSettings {
|
||||||
|
#[serde(default = "default_free_grant")]
|
||||||
|
pub free_token_grant: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for GrantSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
free_token_grant: default_free_grant(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[abuse]` — silent multi-account abuse detection. When at least
|
||||||
|
/// `fingerprint_account_threshold` accounts share one registration
|
||||||
|
/// fingerprint, all of them are silently deactivated (no notice to the
|
||||||
|
/// user; deactivation only surfaces as ordinary inference rejections).
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AbuseSettings {
|
||||||
|
#[serde(default = "default_fingerprint_threshold")]
|
||||||
|
pub fingerprint_account_threshold: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AbuseSettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
fingerprint_account_threshold: default_fingerprint_threshold(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_listen() -> String {
|
||||||
|
"0.0.0.0:8090".into()
|
||||||
|
}
|
||||||
|
fn default_max_connections() -> u32 {
|
||||||
|
16
|
||||||
|
}
|
||||||
|
fn default_free_grant() -> i64 {
|
||||||
|
1_000_000
|
||||||
|
}
|
||||||
|
fn default_fingerprint_threshold() -> i64 {
|
||||||
|
5
|
||||||
|
}
|
||||||
|
fn default_reservation_ttl() -> u64 {
|
||||||
|
120
|
||||||
|
}
|
||||||
|
fn default_sweep_interval() -> u64 {
|
||||||
|
60
|
||||||
|
}
|
||||||
|
fn default_jwt_secret() -> String {
|
||||||
|
"dev-insecure-change-me".into()
|
||||||
|
}
|
||||||
|
fn default_session_ttl() -> u64 {
|
||||||
|
7 * 24 * 3600
|
||||||
|
}
|
||||||
|
fn default_email_token_ttl() -> u64 {
|
||||||
|
24 * 3600
|
||||||
|
}
|
||||||
|
fn default_app_base_url() -> String {
|
||||||
|
"http://localhost:5173".into()
|
||||||
|
}
|
||||||
|
fn default_email_provider() -> String {
|
||||||
|
"log".into()
|
||||||
|
}
|
||||||
|
fn default_from_addr() -> String {
|
||||||
|
"helexa <no-reply@helexa.ai>".into()
|
||||||
|
}
|
||||||
|
|
||||||
|
impl UpstreamConfig {
|
||||||
|
/// Load from a TOML file with `UPSTREAM_`-prefixed env overrides
|
||||||
|
/// (`__` nesting separator).
|
||||||
|
pub fn load(path: impl AsRef<Path>) -> Result<Self, Box<figment::Error>> {
|
||||||
|
Figment::new()
|
||||||
|
.merge(Toml::file(path))
|
||||||
|
.merge(Env::prefixed("UPSTREAM_").split("__"))
|
||||||
|
.extract()
|
||||||
|
.map_err(Box::new)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[allow(clippy::result_large_err)]
|
||||||
|
fn loads_toml_with_env_override_and_defaults() {
|
||||||
|
figment::Jail::expect_with(|jail| {
|
||||||
|
jail.create_file(
|
||||||
|
"helexa-upstream.toml",
|
||||||
|
r#"
|
||||||
|
[db]
|
||||||
|
url = "postgres://localhost/helexa"
|
||||||
|
"#,
|
||||||
|
)?;
|
||||||
|
jail.set_env("UPSTREAM_SERVER__LISTEN", "127.0.0.1:9099");
|
||||||
|
|
||||||
|
let cfg = UpstreamConfig::load("helexa-upstream.toml").expect("load");
|
||||||
|
assert_eq!(cfg.server.listen, "127.0.0.1:9099");
|
||||||
|
assert_eq!(cfg.db.url, "postgres://localhost/helexa");
|
||||||
|
// Defaults applied when sections omitted.
|
||||||
|
assert_eq!(cfg.grant.free_token_grant, 1_000_000);
|
||||||
|
assert_eq!(cfg.abuse.fingerprint_account_threshold, 5);
|
||||||
|
assert_eq!(cfg.db.max_connections, 16);
|
||||||
|
Ok(())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
113
crates/helexa-upstream/src/crypto.rs
Normal file
113
crates/helexa-upstream/src/crypto.rs
Normal file
@@ -0,0 +1,113 @@
|
|||||||
|
//! Hashing + secret-generation helpers.
|
||||||
|
//!
|
||||||
|
//! - **Passwords** (low-entropy) → argon2id PHC strings.
|
||||||
|
//! - **API keys / top-up codes / email + session tokens** (high-entropy
|
||||||
|
//! secrets minted here) → stored only as their sha256; sha256 is the fast,
|
||||||
|
//! sufficient choice for high-entropy material.
|
||||||
|
|
||||||
|
use argon2::Argon2;
|
||||||
|
use argon2::password_hash::rand_core::OsRng as ArgonOsRng;
|
||||||
|
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
|
||||||
|
use rand::RngCore;
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
|
||||||
|
/// sha256 of `input`, as raw bytes (matches the `BYTEA` columns).
|
||||||
|
pub fn sha256(input: &str) -> Vec<u8> {
|
||||||
|
let mut h = Sha256::new();
|
||||||
|
h.update(input.as_bytes());
|
||||||
|
h.finalize().to_vec()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hash a password with argon2id, returning a PHC string for storage.
|
||||||
|
pub fn hash_password(password: &str) -> Result<String, argon2::password_hash::Error> {
|
||||||
|
let salt = SaltString::generate(&mut ArgonOsRng);
|
||||||
|
Ok(Argon2::default()
|
||||||
|
.hash_password(password.as_bytes(), &salt)?
|
||||||
|
.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Verify a password against a stored PHC hash. `false` on any mismatch or
|
||||||
|
/// malformed hash (never panics).
|
||||||
|
pub fn verify_password(password: &str, phc: &str) -> bool {
|
||||||
|
match PasswordHash::new(phc) {
|
||||||
|
Ok(parsed) => Argon2::default()
|
||||||
|
.verify_password(password.as_bytes(), &parsed)
|
||||||
|
.is_ok(),
|
||||||
|
Err(_) => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A fresh URL-safe high-entropy secret (256 bits) for email/session/reset
|
||||||
|
/// tokens. The caller stores only `sha256` of this and emails/returns the
|
||||||
|
/// raw value.
|
||||||
|
pub fn random_token() -> String {
|
||||||
|
let mut bytes = [0u8; 32];
|
||||||
|
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||||
|
base62(&bytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mint a new API key: `(raw, prefix)`. `raw` is shown to the user once;
|
||||||
|
/// only `sha256(raw)` is stored. The prefix is a non-secret display tag.
|
||||||
|
pub fn generate_api_key() -> (String, String) {
|
||||||
|
let mut bytes = [0u8; 32];
|
||||||
|
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||||
|
let raw = format!("sk-helexa-{}", base62(&bytes));
|
||||||
|
// Non-secret prefix for the dashboard list (scheme + first few chars).
|
||||||
|
let prefix: String = raw.chars().take(14).collect();
|
||||||
|
(raw, prefix)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// base62 encode (0-9A-Za-z) — URL/clipboard friendly, no padding.
|
||||||
|
fn base62(bytes: &[u8]) -> String {
|
||||||
|
const ALPHABET: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
|
||||||
|
// Treat the bytes as a big-endian integer and base62 it. 32 bytes → ~43
|
||||||
|
// chars. Simple repeated-division over a big-uint built from the bytes.
|
||||||
|
let mut digits: Vec<u8> = vec![0];
|
||||||
|
for &byte in bytes {
|
||||||
|
let mut carry = byte as u32;
|
||||||
|
for d in digits.iter_mut() {
|
||||||
|
let v = (*d as u32) * 256 + carry;
|
||||||
|
*d = (v % 62) as u8;
|
||||||
|
carry = v / 62;
|
||||||
|
}
|
||||||
|
while carry > 0 {
|
||||||
|
digits.push((carry % 62) as u8);
|
||||||
|
carry /= 62;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
digits
|
||||||
|
.iter()
|
||||||
|
.rev()
|
||||||
|
.map(|&d| ALPHABET[d as usize] as char)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn password_round_trips_and_rejects_wrong() {
|
||||||
|
let phc = hash_password("correct horse").unwrap();
|
||||||
|
assert!(verify_password("correct horse", &phc));
|
||||||
|
assert!(!verify_password("wrong", &phc));
|
||||||
|
assert!(!verify_password("correct horse", "not-a-phc-string"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn api_key_has_scheme_prefix_and_unique_body() {
|
||||||
|
let (raw, prefix) = generate_api_key();
|
||||||
|
assert!(raw.starts_with("sk-helexa-"));
|
||||||
|
assert!(prefix.starts_with("sk-helexa-"));
|
||||||
|
let (raw2, _) = generate_api_key();
|
||||||
|
assert_ne!(raw, raw2, "keys are unique");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn random_tokens_are_unique_and_nonempty() {
|
||||||
|
let a = random_token();
|
||||||
|
let b = random_token();
|
||||||
|
assert!(!a.is_empty());
|
||||||
|
assert_ne!(a, b);
|
||||||
|
}
|
||||||
|
}
|
||||||
20
crates/helexa-upstream/src/db.rs
Normal file
20
crates/helexa-upstream/src/db.rs
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
//! PostgreSQL pool + embedded migrations.
|
||||||
|
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use sqlx::postgres::{PgPool, PgPoolOptions};
|
||||||
|
|
||||||
|
/// Connect to Postgres and run embedded migrations (`./migrations`).
|
||||||
|
pub async fn connect_and_migrate(url: &str, max_connections: u32) -> Result<PgPool> {
|
||||||
|
let pool = PgPoolOptions::new()
|
||||||
|
.max_connections(max_connections)
|
||||||
|
.connect(url)
|
||||||
|
.await
|
||||||
|
.with_context(|| "connecting to PostgreSQL")?;
|
||||||
|
|
||||||
|
sqlx::migrate!("./migrations")
|
||||||
|
.run(&pool)
|
||||||
|
.await
|
||||||
|
.with_context(|| "running migrations")?;
|
||||||
|
|
||||||
|
Ok(pool)
|
||||||
|
}
|
||||||
64
crates/helexa-upstream/src/email.rs
Normal file
64
crates/helexa-upstream/src/email.rs
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
//! Transactional email for verification + password-reset links.
|
||||||
|
//!
|
||||||
|
//! Two transports: `Log` (dev — writes the link to the log so flows are
|
||||||
|
//! testable without a relay) and `Smtp` (lettre over rustls). Built from
|
||||||
|
//! `[email]` config.
|
||||||
|
|
||||||
|
use crate::config::EmailSettings;
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use lettre::message::Mailbox;
|
||||||
|
use lettre::{AsyncSmtpTransport, AsyncTransport, Message, Tokio1Executor};
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub enum EmailSender {
|
||||||
|
/// Dev: log the message instead of sending.
|
||||||
|
Log { from: String },
|
||||||
|
Smtp {
|
||||||
|
from: String,
|
||||||
|
transport: AsyncSmtpTransport<Tokio1Executor>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EmailSender {
|
||||||
|
pub fn from_config(cfg: &EmailSettings) -> Result<Self> {
|
||||||
|
match cfg.provider.as_str() {
|
||||||
|
"smtp" => {
|
||||||
|
let url = cfg
|
||||||
|
.smtp_url
|
||||||
|
.as_deref()
|
||||||
|
.context("[email].smtp_url required when provider = \"smtp\"")?;
|
||||||
|
let transport = AsyncSmtpTransport::<Tokio1Executor>::from_url(url)
|
||||||
|
.context("parsing [email].smtp_url")?
|
||||||
|
.build();
|
||||||
|
Ok(EmailSender::Smtp {
|
||||||
|
from: cfg.from_addr.clone(),
|
||||||
|
transport,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
_ => Ok(EmailSender::Log {
|
||||||
|
from: cfg.from_addr.clone(),
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Send a plaintext email. Errors are returned but the caller treats
|
||||||
|
/// send failures as non-fatal to the request (the user can re-request).
|
||||||
|
pub async fn send(&self, to: &str, subject: &str, body: &str) -> Result<()> {
|
||||||
|
match self {
|
||||||
|
EmailSender::Log { from } => {
|
||||||
|
tracing::info!(%from, %to, %subject, body, "EMAIL (log transport)");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
EmailSender::Smtp { from, transport } => {
|
||||||
|
let msg = Message::builder()
|
||||||
|
.from(from.parse::<Mailbox>().context("parsing from_addr")?)
|
||||||
|
.to(to.parse::<Mailbox>().context("parsing recipient")?)
|
||||||
|
.subject(subject)
|
||||||
|
.body(body.to_string())
|
||||||
|
.context("building message")?;
|
||||||
|
transport.send(msg).await.context("sending email")?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
21
crates/helexa-upstream/src/error.rs
Normal file
21
crates/helexa-upstream/src/error.rs
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
//! Adapter from the shared, axum-agnostic
|
||||||
|
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
|
||||||
|
//! response, with `Retry-After`. The `/authz/v1` surface speaks the #63
|
||||||
|
//! envelope so cortex (an OpenAI-compatible proxy) can forward rejections
|
||||||
|
//! verbatim. (The future `/web/v1` surface uses a plain JSON error shape.)
|
||||||
|
|
||||||
|
use axum::http::{HeaderValue, StatusCode, header};
|
||||||
|
use axum::response::{IntoResponse, Json, Response};
|
||||||
|
use cortex_core::error_envelope::OpenAiError;
|
||||||
|
|
||||||
|
pub fn envelope_response(err: OpenAiError) -> Response {
|
||||||
|
let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||||
|
let retry_after = err.retry_after_secs;
|
||||||
|
let mut response = (status, Json(err.body())).into_response();
|
||||||
|
if let Some(secs) = retry_after
|
||||||
|
&& let Ok(value) = HeaderValue::from_str(&secs.to_string())
|
||||||
|
{
|
||||||
|
response.headers_mut().insert(header::RETRY_AFTER, value);
|
||||||
|
}
|
||||||
|
response
|
||||||
|
}
|
||||||
21
crates/helexa-upstream/src/handlers.rs
Normal file
21
crates/helexa-upstream/src/handlers.rs
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
//! HTTP handlers. B1 ships `/health`; the authz (`/authz/v1`) and web
|
||||||
|
//! (`/web/v1`) surfaces land in later phases.
|
||||||
|
|
||||||
|
use crate::state::AppState;
|
||||||
|
use axum::{Json, Router, extract::State, routing::get};
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route("/health", get(health))
|
||||||
|
.route("/", get(health))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `GET /health` — liveness + a database round-trip (`SELECT 1`).
|
||||||
|
async fn health(State(state): State<AppState>) -> Json<Value> {
|
||||||
|
let db_ok = sqlx::query("SELECT 1").execute(&state.pool).await.is_ok();
|
||||||
|
Json(json!({
|
||||||
|
"status": if db_ok { "ok" } else { "degraded" },
|
||||||
|
"db": if db_ok { "ok" } else { "unreachable" },
|
||||||
|
}))
|
||||||
|
}
|
||||||
338
crates/helexa-upstream/src/ledger.rs
Normal file
338
crates/helexa-upstream/src/ledger.rs
Normal file
@@ -0,0 +1,338 @@
|
|||||||
|
//! The allocation ledger: reserve → settle/release with the no-overshoot
|
||||||
|
//! guarantee enforced by a row-locked transaction.
|
||||||
|
//!
|
||||||
|
//! Each reserve takes `SELECT … FOR UPDATE` on the account (and key) row, so
|
||||||
|
//! concurrent reserves from many cortexes serialize and `spent + reserved`
|
||||||
|
//! can never exceed the effective cap. The `accounts_no_overshoot` CHECK is
|
||||||
|
//! the DB-level backstop. Settle/release are idempotent (they only act on a
|
||||||
|
//! reservation still in `open`).
|
||||||
|
//!
|
||||||
|
//! Per-key effective cap = `min(resolved key cap, remaining account
|
||||||
|
//! allocation)`. The key cap is resolved from its `limit_kind`:
|
||||||
|
//! `hardcap` → the value verbatim; `percent` → that % of the account's
|
||||||
|
//! `allocation_total`.
|
||||||
|
//!
|
||||||
|
//! Cap-window semantics: this module implements **Balance** (non-resetting)
|
||||||
|
//! caps. Rolling-window key sub-caps (and the `RateLimited` rejection that
|
||||||
|
//! rides them) land with the authz API (B2); today an over-cap is always
|
||||||
|
//! `InsufficientQuota`.
|
||||||
|
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// A bearer key resolved to its principal + a budget snapshot.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ResolvedPrincipal {
|
||||||
|
pub account_id: Uuid,
|
||||||
|
pub key_id: Uuid,
|
||||||
|
/// Effective per-key absolute cap (the key sub-cap; the account cap
|
||||||
|
/// still binds at reserve time).
|
||||||
|
pub hard_cap: i64,
|
||||||
|
pub key_spent: i64,
|
||||||
|
pub key_reserved: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve a key by its `sha256` hash to its principal, or `None` when the
|
||||||
|
/// key is unknown/archived **or its account is deactivated** (the silent
|
||||||
|
/// abuse flag — indistinguishable from an unknown key, by design: no clue).
|
||||||
|
pub async fn resolve_key(
|
||||||
|
pool: &PgPool,
|
||||||
|
key_hash: &[u8],
|
||||||
|
) -> Result<Option<ResolvedPrincipal>, sqlx::Error> {
|
||||||
|
let row = sqlx::query(
|
||||||
|
"SELECT k.id AS key_id, k.account_id, k.limit_kind, k.limit_value, \
|
||||||
|
k.key_spent, k.key_reserved, a.allocation_total \
|
||||||
|
FROM api_keys k JOIN accounts a ON a.id = k.account_id \
|
||||||
|
WHERE k.key_hash = $1 AND k.status = 'active' AND a.status = 'active'",
|
||||||
|
)
|
||||||
|
.bind(key_hash)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(row.map(|r| {
|
||||||
|
let total: i64 = sqlx::Row::get(&r, "allocation_total");
|
||||||
|
let limit_kind: String = sqlx::Row::get(&r, "limit_kind");
|
||||||
|
let limit_value: i64 = sqlx::Row::get(&r, "limit_value");
|
||||||
|
ResolvedPrincipal {
|
||||||
|
account_id: sqlx::Row::get(&r, "account_id"),
|
||||||
|
key_id: sqlx::Row::get(&r, "key_id"),
|
||||||
|
hard_cap: resolve_abs_cap(&limit_kind, limit_value, total),
|
||||||
|
key_spent: sqlx::Row::get(&r, "key_spent"),
|
||||||
|
key_reserved: sqlx::Row::get(&r, "key_reserved"),
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-key budget snapshot `(hard_cap, spent, reserved)`, or `None` if the
|
||||||
|
/// key/account isn't an active pair.
|
||||||
|
pub async fn snapshot(
|
||||||
|
pool: &PgPool,
|
||||||
|
account_id: Uuid,
|
||||||
|
key_id: Uuid,
|
||||||
|
) -> Result<Option<(i64, i64, i64)>, sqlx::Error> {
|
||||||
|
let row = sqlx::query(
|
||||||
|
"SELECT k.limit_kind, k.limit_value, k.key_spent, k.key_reserved, a.allocation_total \
|
||||||
|
FROM api_keys k JOIN accounts a ON a.id = k.account_id \
|
||||||
|
WHERE k.id = $1 AND k.account_id = $2 AND k.status = 'active' AND a.status = 'active'",
|
||||||
|
)
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(row.map(|r| {
|
||||||
|
let total: i64 = sqlx::Row::get(&r, "allocation_total");
|
||||||
|
let limit_kind: String = sqlx::Row::get(&r, "limit_kind");
|
||||||
|
let limit_value: i64 = sqlx::Row::get(&r, "limit_value");
|
||||||
|
let cap = resolve_abs_cap(&limit_kind, limit_value, total);
|
||||||
|
(
|
||||||
|
cap,
|
||||||
|
sqlx::Row::get::<i64, _>(&r, "key_spent"),
|
||||||
|
sqlx::Row::get::<i64, _>(&r, "key_reserved"),
|
||||||
|
)
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Release every `open` reservation older than `max_age_secs`, returning
|
||||||
|
/// each one's reserved tokens to its account and key in a single statement.
|
||||||
|
/// The lost-settle self-heal. Returns the number swept.
|
||||||
|
pub async fn sweep_stale(pool: &PgPool, max_age_secs: i64) -> Result<u64, sqlx::Error> {
|
||||||
|
// Data-modifying CTEs: release stale rows, then fold their reserved sums
|
||||||
|
// back into accounts and api_keys. All in one atomic statement.
|
||||||
|
let result = sqlx::query(
|
||||||
|
"WITH stale AS ( \
|
||||||
|
UPDATE reservations SET state = 'released', settled_at = now() \
|
||||||
|
WHERE state = 'open' AND created_at < now() - make_interval(secs => $1) \
|
||||||
|
RETURNING account_id, key_id, reserved \
|
||||||
|
), acct AS ( \
|
||||||
|
UPDATE accounts a SET allocation_reserved = allocation_reserved - s.total \
|
||||||
|
FROM (SELECT account_id, SUM(reserved) AS total FROM stale GROUP BY account_id) s \
|
||||||
|
WHERE a.id = s.account_id \
|
||||||
|
) \
|
||||||
|
UPDATE api_keys k SET key_reserved = key_reserved - s.total \
|
||||||
|
FROM (SELECT key_id, SUM(reserved) AS total FROM stale GROUP BY key_id) s \
|
||||||
|
WHERE k.id = s.key_id",
|
||||||
|
)
|
||||||
|
.bind(max_age_secs as f64)
|
||||||
|
.execute(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(result.rows_affected())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve a key's per-key cap to an absolute token count.
|
||||||
|
///
|
||||||
|
/// `percent` is `floor(allocation_total * limit_value / 100)`; `hardcap` is
|
||||||
|
/// `limit_value` verbatim. Computed in i128 to avoid overflow, floored at 0.
|
||||||
|
pub fn resolve_abs_cap(limit_kind: &str, limit_value: i64, allocation_total: i64) -> i64 {
|
||||||
|
let cap = match limit_kind {
|
||||||
|
"percent" => (allocation_total as i128 * limit_value as i128) / 100,
|
||||||
|
_ => limit_value as i128, // "hardcap" (and any unknown → treat as absolute)
|
||||||
|
};
|
||||||
|
cap.clamp(0, i64::MAX as i128) as i64
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum LedgerError {
|
||||||
|
#[error("account not found")]
|
||||||
|
AccountNotFound,
|
||||||
|
#[error("api key not found or not active")]
|
||||||
|
KeyNotFound,
|
||||||
|
/// Account balance or a Balance-window key sub-cap is exhausted.
|
||||||
|
#[error("insufficient quota: requested {requested}, available {available}")]
|
||||||
|
InsufficientQuota { requested: i64, available: i64 },
|
||||||
|
#[error(transparent)]
|
||||||
|
Db(#[from] sqlx::Error),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reserve `max_tokens` against `account_id`/`key_id`. Returns the
|
||||||
|
/// reservation id (the `BIGSERIAL`, mapped to the cortex `Reservation.id`).
|
||||||
|
pub async fn reserve(
|
||||||
|
pool: &PgPool,
|
||||||
|
account_id: Uuid,
|
||||||
|
key_id: Uuid,
|
||||||
|
max_tokens: i64,
|
||||||
|
) -> Result<i64, LedgerError> {
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
|
||||||
|
// Lock the account row — serializes concurrent reserves on this account.
|
||||||
|
let acct = sqlx::query(
|
||||||
|
"SELECT allocation_total, allocation_spent, allocation_reserved \
|
||||||
|
FROM accounts WHERE id = $1 AND status = 'active' FOR UPDATE",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_optional(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let Some(acct) = acct else {
|
||||||
|
return Err(LedgerError::AccountNotFound);
|
||||||
|
};
|
||||||
|
let total: i64 = sqlx::Row::get(&acct, "allocation_total");
|
||||||
|
let spent: i64 = sqlx::Row::get(&acct, "allocation_spent");
|
||||||
|
let reserved: i64 = sqlx::Row::get(&acct, "allocation_reserved");
|
||||||
|
let account_avail = total - spent - reserved;
|
||||||
|
|
||||||
|
// Lock the key row and resolve its absolute sub-cap.
|
||||||
|
let key = sqlx::query(
|
||||||
|
"SELECT limit_kind, limit_value, key_spent, key_reserved \
|
||||||
|
FROM api_keys WHERE id = $1 AND account_id = $2 AND status = 'active' FOR UPDATE",
|
||||||
|
)
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_optional(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let Some(key) = key else {
|
||||||
|
return Err(LedgerError::KeyNotFound);
|
||||||
|
};
|
||||||
|
let limit_kind: String = sqlx::Row::get(&key, "limit_kind");
|
||||||
|
let limit_value: i64 = sqlx::Row::get(&key, "limit_value");
|
||||||
|
let key_spent: i64 = sqlx::Row::get(&key, "key_spent");
|
||||||
|
let key_reserved: i64 = sqlx::Row::get(&key, "key_reserved");
|
||||||
|
let key_cap = resolve_abs_cap(&limit_kind, limit_value, total);
|
||||||
|
let key_avail = key_cap - key_spent - key_reserved;
|
||||||
|
|
||||||
|
let available = account_avail.min(key_avail).max(0);
|
||||||
|
if max_tokens > available {
|
||||||
|
// tx rolls back on drop
|
||||||
|
return Err(LedgerError::InsufficientQuota {
|
||||||
|
requested: max_tokens,
|
||||||
|
available,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let id: i64 = sqlx::Row::get(
|
||||||
|
&sqlx::query(
|
||||||
|
"INSERT INTO reservations (account_id, key_id, reserved, state) \
|
||||||
|
VALUES ($1, $2, $3, 'open') RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(max_tokens)
|
||||||
|
.fetch_one(&mut *tx)
|
||||||
|
.await?,
|
||||||
|
"id",
|
||||||
|
);
|
||||||
|
sqlx::query("UPDATE accounts SET allocation_reserved = allocation_reserved + $1 WHERE id = $2")
|
||||||
|
.bind(max_tokens)
|
||||||
|
.bind(account_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
sqlx::query("UPDATE api_keys SET key_reserved = key_reserved + $1 WHERE id = $2")
|
||||||
|
.bind(max_tokens)
|
||||||
|
.bind(key_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
tx.commit().await?;
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Settle a reservation with the actual tokens used (clamped to
|
||||||
|
/// `[0, reserved]`). Idempotent: a second settle (or settle after release)
|
||||||
|
/// is a no-op.
|
||||||
|
pub async fn settle(
|
||||||
|
pool: &PgPool,
|
||||||
|
reservation_id: i64,
|
||||||
|
actual_tokens: i64,
|
||||||
|
) -> Result<(), LedgerError> {
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
let row = sqlx::query(
|
||||||
|
"UPDATE reservations SET state = 'settled', settled_at = now(), \
|
||||||
|
actual = LEAST(GREATEST($2, 0), reserved) \
|
||||||
|
WHERE id = $1 AND state = 'open' \
|
||||||
|
RETURNING reserved, account_id, key_id, actual",
|
||||||
|
)
|
||||||
|
.bind(reservation_id)
|
||||||
|
.bind(actual_tokens)
|
||||||
|
.fetch_optional(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = row else {
|
||||||
|
return Ok(()); // already settled/released, or unknown → idempotent no-op
|
||||||
|
};
|
||||||
|
let reserved: i64 = sqlx::Row::get(&row, "reserved");
|
||||||
|
let actual: i64 = sqlx::Row::get(&row, "actual");
|
||||||
|
let account_id: Uuid = sqlx::Row::get(&row, "account_id");
|
||||||
|
let key_id: Uuid = sqlx::Row::get(&row, "key_id");
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE accounts SET allocation_reserved = allocation_reserved - $1, \
|
||||||
|
allocation_spent = allocation_spent + $2 WHERE id = $3",
|
||||||
|
)
|
||||||
|
.bind(reserved)
|
||||||
|
.bind(actual)
|
||||||
|
.bind(account_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE api_keys SET key_reserved = key_reserved - $1, key_spent = key_spent + $2 WHERE id = $3",
|
||||||
|
)
|
||||||
|
.bind(reserved)
|
||||||
|
.bind(actual)
|
||||||
|
.bind(key_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
tx.commit().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Release a reservation, returning its full reserved amount to the
|
||||||
|
/// allocation. Idempotent.
|
||||||
|
pub async fn release(pool: &PgPool, reservation_id: i64) -> Result<(), LedgerError> {
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
let row = sqlx::query(
|
||||||
|
"UPDATE reservations SET state = 'released', settled_at = now() \
|
||||||
|
WHERE id = $1 AND state = 'open' \
|
||||||
|
RETURNING reserved, account_id, key_id",
|
||||||
|
)
|
||||||
|
.bind(reservation_id)
|
||||||
|
.fetch_optional(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = row else {
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
|
let reserved: i64 = sqlx::Row::get(&row, "reserved");
|
||||||
|
let account_id: Uuid = sqlx::Row::get(&row, "account_id");
|
||||||
|
let key_id: Uuid = sqlx::Row::get(&row, "key_id");
|
||||||
|
|
||||||
|
sqlx::query("UPDATE accounts SET allocation_reserved = allocation_reserved - $1 WHERE id = $2")
|
||||||
|
.bind(reserved)
|
||||||
|
.bind(account_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
sqlx::query("UPDATE api_keys SET key_reserved = key_reserved - $1 WHERE id = $2")
|
||||||
|
.bind(reserved)
|
||||||
|
.bind(key_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
tx.commit().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::resolve_abs_cap;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hardcap_is_verbatim() {
|
||||||
|
assert_eq!(resolve_abs_cap("hardcap", 50_000, 1_000_000), 50_000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn percent_is_fraction_of_allocation() {
|
||||||
|
assert_eq!(resolve_abs_cap("percent", 25, 1_000_000), 250_000);
|
||||||
|
assert_eq!(resolve_abs_cap("percent", 100, 1_000_000), 1_000_000);
|
||||||
|
// floor
|
||||||
|
assert_eq!(resolve_abs_cap("percent", 33, 10), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn percent_does_not_overflow_on_large_allocation() {
|
||||||
|
// total * value would overflow i64 if not widened to i128.
|
||||||
|
let cap = resolve_abs_cap("percent", 100, i64::MAX);
|
||||||
|
assert_eq!(cap, i64::MAX);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn negative_or_zero_clamps_to_zero() {
|
||||||
|
assert_eq!(resolve_abs_cap("hardcap", -5, 100), 0);
|
||||||
|
assert_eq!(resolve_abs_cap("percent", 0, 1_000_000), 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
90
crates/helexa-upstream/src/lib.rs
Normal file
90
crates/helexa-upstream/src/lib.rs
Normal file
@@ -0,0 +1,90 @@
|
|||||||
|
//! helexa-upstream — the mesh-level account/authorization authority (#59).
|
||||||
|
//!
|
||||||
|
//! The clearing house above cortex: it issues accounts and API keys, holds
|
||||||
|
//! the real token-allocation ledger, authorizes inference in real time
|
||||||
|
//! (reserve → settle, fail-closed), and tracks served usage for operator
|
||||||
|
//! reconciliation. cortex's `UpstreamEntitlementProvider` (#57) is a client
|
||||||
|
//! of the `/authz/v1` surface; the helexa.ai frontend is a client of the
|
||||||
|
//! `/web/v1` surface.
|
||||||
|
//!
|
||||||
|
//! Landed so far: B1 — schema + reserve→settle [`ledger`] (no-overshoot) +
|
||||||
|
//! `/health`. B2 — the `/authz/v1` [`authz`] surface (resolve/reserve/
|
||||||
|
//! settle/release/snapshot) with shared-bearer client auth and a
|
||||||
|
//! stale-reservation sweeper.
|
||||||
|
|
||||||
|
pub mod authz;
|
||||||
|
pub mod config;
|
||||||
|
pub mod crypto;
|
||||||
|
pub mod db;
|
||||||
|
pub mod email;
|
||||||
|
pub mod error;
|
||||||
|
pub mod handlers;
|
||||||
|
pub mod ledger;
|
||||||
|
pub mod reconcile;
|
||||||
|
pub mod state;
|
||||||
|
pub mod topup;
|
||||||
|
pub mod web;
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use config::UpstreamConfig;
|
||||||
|
use email::EmailSender;
|
||||||
|
use state::AppState;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tower_http::cors::CorsLayer;
|
||||||
|
use tower_http::trace::TraceLayer;
|
||||||
|
|
||||||
|
/// Build the axum application.
|
||||||
|
pub fn build_app(state: AppState) -> axum::Router {
|
||||||
|
axum::Router::new()
|
||||||
|
.merge(handlers::routes())
|
||||||
|
.merge(authz::router(&state))
|
||||||
|
.merge(web::router(&state))
|
||||||
|
// The /web/v1 surface is called cross-origin by the browser SPA in
|
||||||
|
// dev; same-origin behind nginx in prod. Permissive is fine — these
|
||||||
|
// endpoints authenticate via bearer/JWT, not cookies.
|
||||||
|
.layer(CorsLayer::permissive())
|
||||||
|
.layer(TraceLayer::new_for_http())
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start the service: connect Postgres, run migrations, spawn the
|
||||||
|
/// reservation sweeper, bind the listener.
|
||||||
|
pub async fn run(config: UpstreamConfig) -> Result<()> {
|
||||||
|
let pool = db::connect_and_migrate(&config.db.url, config.db.max_connections).await?;
|
||||||
|
let email = EmailSender::from_config(&config.email)?;
|
||||||
|
let listen = config.server.listen.clone();
|
||||||
|
let state = AppState::new(pool, config, email);
|
||||||
|
|
||||||
|
if state.config.client_auth.tokens.is_empty() {
|
||||||
|
tracing::warn!(
|
||||||
|
"no [client_auth] tokens configured — the /authz/v1 surface is OPEN (dev only)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stale-reservation sweeper: releases open reservations whose
|
||||||
|
// settle/release from cortex was lost, self-healing allocation_reserved.
|
||||||
|
spawn_sweeper(&state);
|
||||||
|
|
||||||
|
let addr = listen.parse::<std::net::SocketAddr>()?;
|
||||||
|
tracing::info!("helexa-upstream listening on {addr}");
|
||||||
|
|
||||||
|
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||||
|
axum::serve(listener, build_app(state)).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn spawn_sweeper(state: &AppState) {
|
||||||
|
let pool = state.pool.clone();
|
||||||
|
let ttl = state.config.authz.reservation_ttl_secs as i64;
|
||||||
|
let interval = Duration::from_secs(state.config.authz.sweep_interval_secs);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
tokio::time::sleep(interval).await;
|
||||||
|
match ledger::sweep_stale(&pool, ttl).await {
|
||||||
|
Ok(n) if n > 0 => tracing::info!(swept = n, "released stale reservations"),
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(e) => tracing::warn!(error = %e, "reservation sweep failed"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
99
crates/helexa-upstream/src/main.rs
Normal file
99
crates/helexa-upstream/src/main.rs
Normal file
@@ -0,0 +1,99 @@
|
|||||||
|
use anyhow::Result;
|
||||||
|
use clap::{Parser, Subcommand};
|
||||||
|
use helexa_upstream::config::UpstreamConfig;
|
||||||
|
use tracing_subscriber::EnvFilter;
|
||||||
|
|
||||||
|
#[derive(Parser)]
|
||||||
|
#[command(name = "helexa-upstream")]
|
||||||
|
#[command(about = "Mesh-level account & authorization authority for helexa")]
|
||||||
|
#[command(version)]
|
||||||
|
struct Cli {
|
||||||
|
#[command(subcommand)]
|
||||||
|
command: Commands,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Subcommand)]
|
||||||
|
enum Commands {
|
||||||
|
/// Start the upstream server.
|
||||||
|
Serve {
|
||||||
|
/// Path to the config file.
|
||||||
|
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||||
|
config: String,
|
||||||
|
},
|
||||||
|
/// Mint single-use top-up codes and print them (one per line). The raw
|
||||||
|
/// codes are shown only here — only their hash is stored. (The future
|
||||||
|
/// faucet bot calls the same path.)
|
||||||
|
Mint {
|
||||||
|
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||||
|
config: String,
|
||||||
|
/// Tokens each code grants.
|
||||||
|
#[arg(long)]
|
||||||
|
value: i64,
|
||||||
|
/// How many codes to mint.
|
||||||
|
#[arg(long, default_value_t = 1)]
|
||||||
|
count: u32,
|
||||||
|
/// Optional human label (e.g. "small", "beta-launch").
|
||||||
|
#[arg(long)]
|
||||||
|
denomination: Option<String>,
|
||||||
|
},
|
||||||
|
/// Roll up not-yet-reconciled served usage per operator/period (#58),
|
||||||
|
/// stamp it reconciled, and print the totals. Payout is out of scope.
|
||||||
|
Reconcile {
|
||||||
|
#[arg(short, long, default_value = "helexa-upstream.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_upstream=debug")),
|
||||||
|
)
|
||||||
|
.init();
|
||||||
|
|
||||||
|
let cli = Cli::parse();
|
||||||
|
|
||||||
|
match cli.command {
|
||||||
|
Commands::Serve { config } => {
|
||||||
|
let cfg = UpstreamConfig::load(&config)
|
||||||
|
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||||
|
tracing::info!(listen = %cfg.server.listen, "starting helexa-upstream");
|
||||||
|
helexa_upstream::run(cfg).await?;
|
||||||
|
}
|
||||||
|
Commands::Mint {
|
||||||
|
config,
|
||||||
|
value,
|
||||||
|
count,
|
||||||
|
denomination,
|
||||||
|
} => {
|
||||||
|
let cfg = UpstreamConfig::load(&config)
|
||||||
|
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||||
|
let pool =
|
||||||
|
helexa_upstream::db::connect_and_migrate(&cfg.db.url, cfg.db.max_connections)
|
||||||
|
.await?;
|
||||||
|
let codes =
|
||||||
|
helexa_upstream::topup::mint(&pool, value, count, denomination.as_deref()).await?;
|
||||||
|
// Raw codes to stdout (one per line) for the operator to distribute;
|
||||||
|
// logs/diagnostics go to stderr via tracing.
|
||||||
|
for code in codes {
|
||||||
|
println!("{code}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Commands::Reconcile { config } => {
|
||||||
|
let cfg = UpstreamConfig::load(&config)
|
||||||
|
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||||
|
let pool =
|
||||||
|
helexa_upstream::db::connect_and_migrate(&cfg.db.url, cfg.db.max_connections)
|
||||||
|
.await?;
|
||||||
|
let rollup = helexa_upstream::reconcile::reconcile(&pool).await?;
|
||||||
|
for r in &rollup {
|
||||||
|
println!("{}\t{}\t{}", r.operator_id, r.period, r.total_served_tokens);
|
||||||
|
}
|
||||||
|
tracing::info!(operators_periods = rollup.len(), "reconciliation complete");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
43
crates/helexa-upstream/src/reconcile.rs
Normal file
43
crates/helexa-upstream/src/reconcile.rs
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
//! Reconciliation rollup (#58): aggregate the served-usage ledger per
|
||||||
|
//! operator and period for operator compensation, stamping rows
|
||||||
|
//! `reconciled_at` so each window is settled once. The payout mechanism
|
||||||
|
//! itself is out of scope — this produces the authoritative per-operator
|
||||||
|
//! totals a settlement process consumes.
|
||||||
|
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct RollupRow {
|
||||||
|
pub operator_id: String,
|
||||||
|
pub period: chrono::NaiveDate,
|
||||||
|
pub total_served_tokens: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Roll up all not-yet-reconciled served-usage into per-(operator, period)
|
||||||
|
/// totals, then stamp those rows `reconciled_at`. Returns the rollup.
|
||||||
|
/// Idempotent: a second run finds nothing unreconciled and returns empty.
|
||||||
|
pub async fn reconcile(pool: &PgPool) -> Result<Vec<RollupRow>, sqlx::Error> {
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
let rows = sqlx::query(
|
||||||
|
// SUM(bigint) is numeric in Postgres — cast back to bigint for i64.
|
||||||
|
"SELECT operator_id, period, SUM(served_tokens)::bigint AS total \
|
||||||
|
FROM served_usage WHERE reconciled_at IS NULL \
|
||||||
|
GROUP BY operator_id, period ORDER BY operator_id, period",
|
||||||
|
)
|
||||||
|
.fetch_all(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let rollup: Vec<RollupRow> = rows
|
||||||
|
.iter()
|
||||||
|
.map(|r| RollupRow {
|
||||||
|
operator_id: r.get("operator_id"),
|
||||||
|
period: r.get("period"),
|
||||||
|
total_served_tokens: r.get::<i64, _>("total"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
sqlx::query("UPDATE served_usage SET reconciled_at = now() WHERE reconciled_at IS NULL")
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
tx.commit().await?;
|
||||||
|
Ok(rollup)
|
||||||
|
}
|
||||||
23
crates/helexa-upstream/src/state.rs
Normal file
23
crates/helexa-upstream/src/state.rs
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
//! Shared application state.
|
||||||
|
|
||||||
|
use crate::config::UpstreamConfig;
|
||||||
|
use crate::email::EmailSender;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AppState {
|
||||||
|
pub pool: PgPool,
|
||||||
|
pub config: Arc<UpstreamConfig>,
|
||||||
|
pub email: EmailSender,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AppState {
|
||||||
|
pub fn new(pool: PgPool, config: UpstreamConfig, email: EmailSender) -> Self {
|
||||||
|
Self {
|
||||||
|
pool,
|
||||||
|
config: Arc::new(config),
|
||||||
|
email,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
82
crates/helexa-upstream/src/topup.rs
Normal file
82
crates/helexa-upstream/src/topup.rs
Normal file
@@ -0,0 +1,82 @@
|
|||||||
|
//! Single-use top-up codes (#B5) — the second half of the hybrid allocation
|
||||||
|
//! model. Each code grants `value` tokens to the account that redeems it,
|
||||||
|
//! raising `accounts.allocation_total`. Minting codes is operator/CLI side
|
||||||
|
//! (the future faucet bot calls the same `mint` path); redemption is a
|
||||||
|
//! `/web/v1` action.
|
||||||
|
//!
|
||||||
|
//! Security: only `sha256(code)` is stored. Redemption is **timing-safe and
|
||||||
|
//! single-use** — a conditional `UPDATE … WHERE redeemed_by IS NULL` does
|
||||||
|
//! the claim atomically (concurrent double-redeem → exactly one winner), and
|
||||||
|
//! a not-found code and an already-redeemed code return the **same** generic
|
||||||
|
//! failure with the same code path (no oracle for "valid but spent").
|
||||||
|
|
||||||
|
use crate::crypto::{random_token, sha256};
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum TopUpError {
|
||||||
|
/// Code unknown OR already redeemed — deliberately indistinguishable.
|
||||||
|
#[error("invalid or already-redeemed code")]
|
||||||
|
Invalid,
|
||||||
|
#[error(transparent)]
|
||||||
|
Db(#[from] sqlx::Error),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Redeem `raw_code` for `account_id`, raising the account's
|
||||||
|
/// `allocation_total` by the code's value. Returns the new total.
|
||||||
|
pub async fn redeem(pool: &PgPool, account_id: Uuid, raw_code: &str) -> Result<i64, TopUpError> {
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
// Atomic single-use claim. `redeemed_by IS NULL` is the guarantee: under
|
||||||
|
// concurrent redemption exactly one UPDATE touches the row.
|
||||||
|
let claimed = sqlx::query(
|
||||||
|
"UPDATE top_up_codes SET redeemed_by = $1, redeemed_at = now() \
|
||||||
|
WHERE code_hash = $2 AND redeemed_by IS NULL RETURNING value",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(sha256(raw_code))
|
||||||
|
.fetch_optional(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = claimed else {
|
||||||
|
// Not found or already redeemed — same path, same error.
|
||||||
|
return Err(TopUpError::Invalid);
|
||||||
|
};
|
||||||
|
let value: i64 = row.get("value");
|
||||||
|
let new_total: i64 = sqlx::query(
|
||||||
|
"UPDATE accounts SET allocation_total = allocation_total + $1 WHERE id = $2 \
|
||||||
|
RETURNING allocation_total",
|
||||||
|
)
|
||||||
|
.bind(value)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_one(&mut *tx)
|
||||||
|
.await?
|
||||||
|
.get("allocation_total");
|
||||||
|
tx.commit().await?;
|
||||||
|
Ok(new_total)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mint `count` codes each worth `value` tokens, optionally tagged with a
|
||||||
|
/// `denomination` label. Returns the raw codes (shown once — only their
|
||||||
|
/// hash is stored). The CLI prints these; the future faucet bot calls this.
|
||||||
|
pub async fn mint(
|
||||||
|
pool: &PgPool,
|
||||||
|
value: i64,
|
||||||
|
count: u32,
|
||||||
|
denomination: Option<&str>,
|
||||||
|
) -> Result<Vec<String>, sqlx::Error> {
|
||||||
|
let mut codes = Vec::with_capacity(count as usize);
|
||||||
|
for _ in 0..count {
|
||||||
|
let raw = format!("helexa-topup-{}", random_token());
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO top_up_codes (code_hash, value, denomination) VALUES ($1, $2, $3)",
|
||||||
|
)
|
||||||
|
.bind(sha256(&raw))
|
||||||
|
.bind(value)
|
||||||
|
.bind(denomination)
|
||||||
|
.execute(pool)
|
||||||
|
.await?;
|
||||||
|
codes.push(raw);
|
||||||
|
}
|
||||||
|
Ok(codes)
|
||||||
|
}
|
||||||
594
crates/helexa-upstream/src/web.rs
Normal file
594
crates/helexa-upstream/src/web.rs
Normal file
@@ -0,0 +1,594 @@
|
|||||||
|
//! `/web/v1` — the human-facing account API the helexa.ai frontend (#F4)
|
||||||
|
//! consumes: email+password auth (register / verify / login / reset),
|
||||||
|
//! API-key CRUD with per-key limits, and the account balance. Web sessions
|
||||||
|
//! are JWTs, **distinct** from inference API keys.
|
||||||
|
//!
|
||||||
|
//! Errors use a plain JSON shape `{ "error": { "message", "code" } }` (web
|
||||||
|
//! clients, not OpenAI clients — the #63 envelope is the authz surface).
|
||||||
|
//!
|
||||||
|
//! Silent fingerprint abuse (no clue to the abuser): registration captures
|
||||||
|
//! the browser fingerprint and always succeeds; when ≥ threshold accounts
|
||||||
|
//! share one fingerprint, all are silently `deactivated` (keys then resolve
|
||||||
|
//! as ordinary `401`s at the authz surface — never a "banned" signal).
|
||||||
|
|
||||||
|
use crate::crypto::{generate_api_key, hash_password, random_token, sha256, verify_password};
|
||||||
|
use crate::state::AppState;
|
||||||
|
use axum::extract::{Path, Request, State};
|
||||||
|
use axum::http::{StatusCode, header};
|
||||||
|
use axum::middleware::Next;
|
||||||
|
use axum::response::{IntoResponse, Json, Response};
|
||||||
|
use axum::routing::{get, post};
|
||||||
|
use axum::{Extension, Router};
|
||||||
|
use chrono::{DateTime, Duration, Utc};
|
||||||
|
use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::json;
|
||||||
|
use sqlx::Row;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
pub fn router(state: &AppState) -> Router<AppState> {
|
||||||
|
let protected = Router::new()
|
||||||
|
.route("/web/v1/account", get(account))
|
||||||
|
.route("/web/v1/keys", get(list_keys).post(create_key))
|
||||||
|
.route("/web/v1/keys/{id}/archive", post(archive_key))
|
||||||
|
.route(
|
||||||
|
"/web/v1/keys/{id}/limit",
|
||||||
|
axum::routing::patch(update_key_limit),
|
||||||
|
)
|
||||||
|
.route("/web/v1/redeem", post(redeem))
|
||||||
|
.layer(axum::middleware::from_fn_with_state(
|
||||||
|
state.clone(),
|
||||||
|
require_session,
|
||||||
|
));
|
||||||
|
|
||||||
|
Router::new()
|
||||||
|
.route("/web/v1/register", post(register))
|
||||||
|
.route("/web/v1/verify", post(verify))
|
||||||
|
.route("/web/v1/login", post(login))
|
||||||
|
.route("/web/v1/password-reset/request", post(reset_request))
|
||||||
|
.route("/web/v1/password-reset/confirm", post(reset_confirm))
|
||||||
|
.merge(protected)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── errors ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
enum WebError {
|
||||||
|
BadRequest(&'static str),
|
||||||
|
Unauthorized,
|
||||||
|
Internal,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IntoResponse for WebError {
|
||||||
|
fn into_response(self) -> Response {
|
||||||
|
let (status, code, message) = match self {
|
||||||
|
WebError::BadRequest(m) => (StatusCode::BAD_REQUEST, "bad_request", m),
|
||||||
|
WebError::Unauthorized => (StatusCode::UNAUTHORIZED, "unauthorized", "unauthorized"),
|
||||||
|
WebError::Internal => (
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
"internal_error",
|
||||||
|
"internal error",
|
||||||
|
),
|
||||||
|
};
|
||||||
|
(
|
||||||
|
status,
|
||||||
|
Json(json!({"error": {"message": message, "code": code}})),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<sqlx::Error> for WebError {
|
||||||
|
fn from(e: sqlx::Error) -> Self {
|
||||||
|
tracing::error!(error = %e, "web db error");
|
||||||
|
WebError::Internal
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type WebResult<T> = Result<T, WebError>;
|
||||||
|
|
||||||
|
// ── sessions (JWT) ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize)]
|
||||||
|
struct Claims {
|
||||||
|
sub: String, // user id
|
||||||
|
exp: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mint_session(state: &AppState, user_id: Uuid) -> WebResult<String> {
|
||||||
|
let exp = (Utc::now() + Duration::seconds(state.config.auth.session_ttl_secs as i64))
|
||||||
|
.timestamp() as usize;
|
||||||
|
let claims = Claims {
|
||||||
|
sub: user_id.to_string(),
|
||||||
|
exp,
|
||||||
|
};
|
||||||
|
encode(
|
||||||
|
&Header::default(),
|
||||||
|
&claims,
|
||||||
|
&EncodingKey::from_secret(state.config.auth.jwt_secret.as_bytes()),
|
||||||
|
)
|
||||||
|
.map_err(|_| WebError::Internal)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Authenticated user id, injected by [`require_session`].
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct AuthUser(Uuid);
|
||||||
|
|
||||||
|
async fn require_session(State(state): State<AppState>, mut req: Request, next: Next) -> Response {
|
||||||
|
let token = req
|
||||||
|
.headers()
|
||||||
|
.get(header::AUTHORIZATION)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.and_then(|v| v.strip_prefix("Bearer "))
|
||||||
|
.map(str::trim);
|
||||||
|
let Some(token) = token else {
|
||||||
|
return WebError::Unauthorized.into_response();
|
||||||
|
};
|
||||||
|
let decoded = decode::<Claims>(
|
||||||
|
token,
|
||||||
|
&DecodingKey::from_secret(state.config.auth.jwt_secret.as_bytes()),
|
||||||
|
&Validation::default(),
|
||||||
|
);
|
||||||
|
match decoded
|
||||||
|
.ok()
|
||||||
|
.and_then(|d| Uuid::parse_str(&d.claims.sub).ok())
|
||||||
|
{
|
||||||
|
Some(uid) => {
|
||||||
|
req.extensions_mut().insert(AuthUser(uid));
|
||||||
|
next.run(req).await
|
||||||
|
}
|
||||||
|
None => WebError::Unauthorized.into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The caller's single account id.
|
||||||
|
async fn account_id_for(state: &AppState, user_id: Uuid) -> WebResult<Uuid> {
|
||||||
|
let row = sqlx::query("SELECT id FROM accounts WHERE owner_user_id = $1")
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?;
|
||||||
|
row.map(|r| r.get::<Uuid, _>("id"))
|
||||||
|
.ok_or(WebError::Internal)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── auth lifecycle ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct RegisterReq {
|
||||||
|
email: String,
|
||||||
|
password: String,
|
||||||
|
#[serde(default)]
|
||||||
|
fingerprint: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/register` — always returns `202`, regardless of whether the
|
||||||
|
/// email was new, already taken, or fingerprint-flagged (no enumeration, no
|
||||||
|
/// abuse clue).
|
||||||
|
async fn register(State(state): State<AppState>, Json(req): Json<RegisterReq>) -> Response {
|
||||||
|
match register_inner(&state, req).await {
|
||||||
|
Ok(()) | Err(WebError::BadRequest(_)) => {}
|
||||||
|
Err(e) => return e.into_response(),
|
||||||
|
}
|
||||||
|
// Generic 202 whatever happened above (except hard server errors).
|
||||||
|
StatusCode::ACCEPTED.into_response()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn register_inner(state: &AppState, req: RegisterReq) -> WebResult<()> {
|
||||||
|
if !req.email.contains('@') {
|
||||||
|
return Err(WebError::BadRequest("invalid email"));
|
||||||
|
}
|
||||||
|
if req.password.len() < 8 {
|
||||||
|
return Err(WebError::BadRequest("password too short (min 8)"));
|
||||||
|
}
|
||||||
|
let phc = hash_password(&req.password).map_err(|_| WebError::Internal)?;
|
||||||
|
|
||||||
|
// Insert the user; a duplicate email silently no-ops (no enumeration).
|
||||||
|
let user_id: Option<Uuid> = sqlx::query(
|
||||||
|
"INSERT INTO users (email, password_hash, registration_fingerprint) \
|
||||||
|
VALUES ($1, $2, $3) ON CONFLICT (email) DO NOTHING RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(&req.email)
|
||||||
|
.bind(&phc)
|
||||||
|
.bind(&req.fingerprint)
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?
|
||||||
|
.map(|r| r.get("id"));
|
||||||
|
|
||||||
|
let Some(user_id) = user_id else {
|
||||||
|
return Ok(()); // email already registered — say nothing
|
||||||
|
};
|
||||||
|
|
||||||
|
// Account with the flat free grant.
|
||||||
|
sqlx::query("INSERT INTO accounts (owner_user_id, allocation_total) VALUES ($1, $2)")
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(state.config.grant.free_token_grant)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Silent fingerprint abuse handling.
|
||||||
|
if let Some(fp) = req.fingerprint.as_deref().filter(|f| !f.is_empty()) {
|
||||||
|
apply_fingerprint_policy(state, fp).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Email verification link.
|
||||||
|
let token = random_token();
|
||||||
|
let expires: DateTime<Utc> =
|
||||||
|
Utc::now() + Duration::seconds(state.config.auth.email_token_ttl_secs as i64);
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||||
|
VALUES ($1, $2, 'verify', $3)",
|
||||||
|
)
|
||||||
|
.bind(sha256(&token))
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(expires)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let link = format!("{}/verify?token={token}", state.config.auth.app_base_url);
|
||||||
|
let _ = state
|
||||||
|
.email
|
||||||
|
.send(
|
||||||
|
&req.email,
|
||||||
|
"Verify your helexa account",
|
||||||
|
&format!("Welcome to helexa. Verify your email:\n\n{link}\n"),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Count accounts sharing `fp`; flag them, and silently deactivate all once
|
||||||
|
/// the count reaches the configured threshold. No response difference — the
|
||||||
|
/// abuser gets no signal.
|
||||||
|
async fn apply_fingerprint_policy(state: &AppState, fp: &str) -> WebResult<()> {
|
||||||
|
let count: i64 =
|
||||||
|
sqlx::query_scalar("SELECT count(*) FROM users WHERE registration_fingerprint = $1")
|
||||||
|
.bind(fp)
|
||||||
|
.fetch_one(&state.pool)
|
||||||
|
.await?;
|
||||||
|
if count > 1 {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE accounts SET fingerprint_flagged = true \
|
||||||
|
WHERE owner_user_id IN (SELECT id FROM users WHERE registration_fingerprint = $1)",
|
||||||
|
)
|
||||||
|
.bind(fp)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
if count >= state.config.abuse.fingerprint_account_threshold {
|
||||||
|
let res = sqlx::query(
|
||||||
|
"UPDATE accounts SET status = 'deactivated' \
|
||||||
|
WHERE owner_user_id IN (SELECT id FROM users WHERE registration_fingerprint = $1)",
|
||||||
|
)
|
||||||
|
.bind(fp)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
tracing::warn!(
|
||||||
|
fingerprint = fp,
|
||||||
|
accounts = res.rows_affected(),
|
||||||
|
"silently deactivated fingerprint-abusing accounts"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct TokenReq {
|
||||||
|
token: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/verify` — consume a verification token, mark verified.
|
||||||
|
async fn verify(State(state): State<AppState>, Json(req): Json<TokenReq>) -> WebResult<Response> {
|
||||||
|
let row = sqlx::query(
|
||||||
|
"UPDATE email_tokens SET consumed_at = now() \
|
||||||
|
WHERE token_hash = $1 AND kind = 'verify' AND consumed_at IS NULL AND expires_at > now() \
|
||||||
|
RETURNING user_id",
|
||||||
|
)
|
||||||
|
.bind(sha256(&req.token))
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = row else {
|
||||||
|
return Err(WebError::BadRequest("invalid or expired token"));
|
||||||
|
};
|
||||||
|
let user_id: Uuid = row.get("user_id");
|
||||||
|
sqlx::query("UPDATE users SET email_verified = true WHERE id = $1")
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
Ok(StatusCode::OK.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct LoginReq {
|
||||||
|
email: String,
|
||||||
|
password: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/login` — verify password + email-verified → session JWT.
|
||||||
|
async fn login(State(state): State<AppState>, Json(req): Json<LoginReq>) -> WebResult<Response> {
|
||||||
|
let row = sqlx::query("SELECT id, password_hash, email_verified FROM users WHERE email = $1")
|
||||||
|
.bind(&req.email)
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?;
|
||||||
|
// Generic 401 for every failure mode (no enumeration).
|
||||||
|
let Some(row) = row else {
|
||||||
|
return Err(WebError::Unauthorized);
|
||||||
|
};
|
||||||
|
let phc: String = row.get("password_hash");
|
||||||
|
let verified: bool = row.get("email_verified");
|
||||||
|
if !verify_password(&req.password, &phc) || !verified {
|
||||||
|
return Err(WebError::Unauthorized);
|
||||||
|
}
|
||||||
|
let user_id: Uuid = row.get("id");
|
||||||
|
let token = mint_session(&state, user_id)?;
|
||||||
|
Ok(Json(json!({
|
||||||
|
"token": token,
|
||||||
|
"expires_in": state.config.auth.session_ttl_secs,
|
||||||
|
}))
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct EmailReq {
|
||||||
|
email: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/password-reset/request` — always `202` (no enumeration);
|
||||||
|
/// mints + emails a reset token only if the account exists.
|
||||||
|
async fn reset_request(State(state): State<AppState>, Json(req): Json<EmailReq>) -> Response {
|
||||||
|
// The inner only ever yields `Internal` (DB failure); a missing email is
|
||||||
|
// Ok(()) so there's no enumeration. Surface 500 on a real error, else 202.
|
||||||
|
match reset_request_inner(&state, &req.email).await {
|
||||||
|
Ok(()) => StatusCode::ACCEPTED.into_response(),
|
||||||
|
Err(e) => e.into_response(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn reset_request_inner(state: &AppState, email: &str) -> WebResult<()> {
|
||||||
|
let row = sqlx::query("SELECT id FROM users WHERE email = $1")
|
||||||
|
.bind(email)
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = row else { return Ok(()) };
|
||||||
|
let user_id: Uuid = row.get("id");
|
||||||
|
let token = random_token();
|
||||||
|
let expires: DateTime<Utc> =
|
||||||
|
Utc::now() + Duration::seconds(state.config.auth.email_token_ttl_secs as i64);
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||||
|
VALUES ($1, $2, 'reset', $3)",
|
||||||
|
)
|
||||||
|
.bind(sha256(&token))
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(expires)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
let link = format!("{}/reset?token={token}", state.config.auth.app_base_url);
|
||||||
|
let _ = state
|
||||||
|
.email
|
||||||
|
.send(
|
||||||
|
email,
|
||||||
|
"Reset your helexa password",
|
||||||
|
&format!("Reset your password:\n\n{link}\n"),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ResetConfirmReq {
|
||||||
|
token: String,
|
||||||
|
new_password: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/password-reset/confirm` — consume reset token, rotate hash.
|
||||||
|
async fn reset_confirm(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(req): Json<ResetConfirmReq>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
if req.new_password.len() < 8 {
|
||||||
|
return Err(WebError::BadRequest("password too short (min 8)"));
|
||||||
|
}
|
||||||
|
let row = sqlx::query(
|
||||||
|
"UPDATE email_tokens SET consumed_at = now() \
|
||||||
|
WHERE token_hash = $1 AND kind = 'reset' AND consumed_at IS NULL AND expires_at > now() \
|
||||||
|
RETURNING user_id",
|
||||||
|
)
|
||||||
|
.bind(sha256(&req.token))
|
||||||
|
.fetch_optional(&state.pool)
|
||||||
|
.await?;
|
||||||
|
let Some(row) = row else {
|
||||||
|
return Err(WebError::BadRequest("invalid or expired token"));
|
||||||
|
};
|
||||||
|
let user_id: Uuid = row.get("user_id");
|
||||||
|
let phc = hash_password(&req.new_password).map_err(|_| WebError::Internal)?;
|
||||||
|
sqlx::query("UPDATE users SET password_hash = $1 WHERE id = $2")
|
||||||
|
.bind(phc)
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
Ok(StatusCode::OK.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── account + keys (protected) ──────────────────────────────────────
|
||||||
|
|
||||||
|
async fn account(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
let row = sqlx::query(
|
||||||
|
"SELECT allocation_total, allocation_spent, allocation_reserved FROM accounts WHERE id = $1",
|
||||||
|
)
|
||||||
|
.bind(acct)
|
||||||
|
.fetch_one(&state.pool)
|
||||||
|
.await?;
|
||||||
|
Ok(Json(json!({
|
||||||
|
"account_id": acct.to_string(),
|
||||||
|
"allocation_total": row.get::<i64, _>("allocation_total"),
|
||||||
|
"allocation_spent": row.get::<i64, _>("allocation_spent"),
|
||||||
|
"allocation_reserved": row.get::<i64, _>("allocation_reserved"),
|
||||||
|
}))
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_keys(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
let rows = sqlx::query(
|
||||||
|
"SELECT id, key_prefix, label, status, limit_kind, limit_value, key_spent, key_reserved, \
|
||||||
|
created_at \
|
||||||
|
FROM api_keys WHERE account_id = $1 ORDER BY created_at DESC",
|
||||||
|
)
|
||||||
|
.bind(acct)
|
||||||
|
.fetch_all(&state.pool)
|
||||||
|
.await?;
|
||||||
|
let keys: Vec<_> = rows
|
||||||
|
.iter()
|
||||||
|
.map(|r| {
|
||||||
|
json!({
|
||||||
|
"id": r.get::<Uuid, _>("id").to_string(),
|
||||||
|
"prefix": r.get::<String, _>("key_prefix"),
|
||||||
|
"label": r.get::<String, _>("label"),
|
||||||
|
"status": r.get::<String, _>("status"),
|
||||||
|
"limit_kind": r.get::<String, _>("limit_kind"),
|
||||||
|
"limit_value": r.get::<i64, _>("limit_value"),
|
||||||
|
"spent": r.get::<i64, _>("key_spent"),
|
||||||
|
"reserved": r.get::<i64, _>("key_reserved"),
|
||||||
|
"created_at": r.get::<DateTime<Utc>, _>("created_at").to_rfc3339(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Ok(Json(json!({ "keys": keys })).into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CreateKeyReq {
|
||||||
|
#[serde(default)]
|
||||||
|
label: String,
|
||||||
|
/// "percent" | "hardcap" (default percent=100 → full allocation).
|
||||||
|
#[serde(default)]
|
||||||
|
limit_kind: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
limit_value: Option<i64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
Json(req): Json<CreateKeyReq>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
let limit_kind = match req.limit_kind.as_deref() {
|
||||||
|
Some("hardcap") => "hardcap",
|
||||||
|
_ => "percent",
|
||||||
|
};
|
||||||
|
let limit_value = req.limit_value.unwrap_or(100).max(0);
|
||||||
|
let (raw, prefix) = generate_api_key();
|
||||||
|
let id: Uuid = sqlx::query(
|
||||||
|
"INSERT INTO api_keys (account_id, key_hash, key_prefix, label, limit_kind, limit_value) \
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(acct)
|
||||||
|
.bind(sha256(&raw))
|
||||||
|
.bind(&prefix)
|
||||||
|
.bind(&req.label)
|
||||||
|
.bind(limit_kind)
|
||||||
|
.bind(limit_value)
|
||||||
|
.fetch_one(&state.pool)
|
||||||
|
.await?
|
||||||
|
.get("id");
|
||||||
|
// The raw key is shown exactly once.
|
||||||
|
Ok((
|
||||||
|
StatusCode::CREATED,
|
||||||
|
Json(json!({
|
||||||
|
"id": id.to_string(),
|
||||||
|
"key": raw,
|
||||||
|
"prefix": prefix,
|
||||||
|
"limit_kind": limit_kind,
|
||||||
|
"limit_value": limit_value,
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn archive_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
Path(id): Path<Uuid>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
let res = sqlx::query(
|
||||||
|
"UPDATE api_keys SET status = 'archived' WHERE id = $1 AND account_id = $2 AND status = 'active'",
|
||||||
|
)
|
||||||
|
.bind(id)
|
||||||
|
.bind(acct)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
if res.rows_affected() == 0 {
|
||||||
|
return Err(WebError::BadRequest("no such active key"));
|
||||||
|
}
|
||||||
|
Ok(StatusCode::NO_CONTENT.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct UpdateLimitReq {
|
||||||
|
limit_kind: String,
|
||||||
|
limit_value: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn update_key_limit(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
Path(id): Path<Uuid>,
|
||||||
|
Json(req): Json<UpdateLimitReq>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
if req.limit_kind != "percent" && req.limit_kind != "hardcap" {
|
||||||
|
return Err(WebError::BadRequest(
|
||||||
|
"limit_kind must be percent or hardcap",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if req.limit_value < 0 {
|
||||||
|
return Err(WebError::BadRequest("limit_value must be >= 0"));
|
||||||
|
}
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
let res = sqlx::query(
|
||||||
|
"UPDATE api_keys SET limit_kind = $1, limit_value = $2 WHERE id = $3 AND account_id = $4",
|
||||||
|
)
|
||||||
|
.bind(&req.limit_kind)
|
||||||
|
.bind(req.limit_value)
|
||||||
|
.bind(id)
|
||||||
|
.bind(acct)
|
||||||
|
.execute(&state.pool)
|
||||||
|
.await?;
|
||||||
|
if res.rows_affected() == 0 {
|
||||||
|
return Err(WebError::BadRequest("no such key"));
|
||||||
|
}
|
||||||
|
Ok(StatusCode::NO_CONTENT.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct RedeemReq {
|
||||||
|
code: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `POST /web/v1/redeem` — redeem a single-use top-up code, raising the
|
||||||
|
/// account's allocation. Returns the new total. Generic 400 for an invalid
|
||||||
|
/// or already-redeemed code (no oracle).
|
||||||
|
async fn redeem(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(user): Extension<AuthUser>,
|
||||||
|
Json(req): Json<RedeemReq>,
|
||||||
|
) -> WebResult<Response> {
|
||||||
|
let acct = account_id_for(&state, user.0).await?;
|
||||||
|
match crate::topup::redeem(&state.pool, acct, &req.code).await {
|
||||||
|
Ok(new_total) => Ok(Json(json!({ "allocation_total": new_total })).into_response()),
|
||||||
|
Err(crate::topup::TopUpError::Invalid) => {
|
||||||
|
Err(WebError::BadRequest("invalid or already-redeemed code"))
|
||||||
|
}
|
||||||
|
Err(crate::topup::TopUpError::Db(e)) => {
|
||||||
|
tracing::error!(error = %e, "redeem db error");
|
||||||
|
Err(WebError::Internal)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
243
crates/helexa-upstream/tests/authz_pg.rs
Normal file
243
crates/helexa-upstream/tests/authz_pg.rs
Normal file
@@ -0,0 +1,243 @@
|
|||||||
|
//! Integration tests for the `/authz/v1` surface against a real Postgres,
|
||||||
|
//! driving the built axum app over HTTP. Gated on `UPSTREAM_TEST_DATABASE_URL`
|
||||||
|
//! (skips cleanly when unset, so CI stays green without a DB):
|
||||||
|
//!
|
||||||
|
//! UPSTREAM_TEST_DATABASE_URL=postgres://helexa:helexa@localhost/helexa_test \
|
||||||
|
//! cargo test -p helexa-upstream --test authz_pg
|
||||||
|
|
||||||
|
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||||
|
use helexa_upstream::crypto::sha256;
|
||||||
|
use helexa_upstream::db::connect_and_migrate;
|
||||||
|
use helexa_upstream::state::AppState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use sqlx::Executor;
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
const CLIENT_TOKEN: &str = "test-operator-token";
|
||||||
|
|
||||||
|
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||||
|
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||||
|
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||||
|
return None;
|
||||||
|
};
|
||||||
|
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||||
|
|
||||||
|
let mut config = UpstreamConfig {
|
||||||
|
server: Default::default(),
|
||||||
|
db: helexa_upstream::config::DbSettings {
|
||||||
|
url,
|
||||||
|
max_connections: 16,
|
||||||
|
},
|
||||||
|
grant: Default::default(),
|
||||||
|
abuse: Default::default(),
|
||||||
|
client_auth: Default::default(),
|
||||||
|
authz: Default::default(),
|
||||||
|
auth: Default::default(),
|
||||||
|
email: Default::default(),
|
||||||
|
};
|
||||||
|
config.client_auth.tokens.push(ClientToken {
|
||||||
|
token: CLIENT_TOKEN.into(),
|
||||||
|
operator_id: "op-test".into(),
|
||||||
|
});
|
||||||
|
|
||||||
|
let email = helexa_upstream::email::EmailSender::from_config(&config.email).unwrap();
|
||||||
|
let state = AppState::new(pool.clone(), config, email);
|
||||||
|
let app = helexa_upstream::build_app(state);
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
Some((format!("http://{addr}"), pool))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Seed an account with `total` allocation and an active key with raw value
|
||||||
|
/// `raw` (percent=100). Optionally deactivate the account. Returns
|
||||||
|
/// (account_id, key_id).
|
||||||
|
async fn seed_key(pool: &PgPool, total: i64, raw: &str, deactivated: bool) -> (Uuid, Uuid) {
|
||||||
|
let user_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO users (email, password_hash, email_verified) VALUES ($1,'x',true) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(format!("u-{}@t.local", Uuid::new_v4())),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let status = if deactivated { "deactivated" } else { "active" };
|
||||||
|
let account_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO accounts (owner_user_id, allocation_total, status) VALUES ($1,$2,$3) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(total)
|
||||||
|
.bind(status),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let key_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||||
|
VALUES ($1,$2,'sk-test','percent',100) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(sha256(raw)),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
(account_id, key_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn client() -> reqwest::Client {
|
||||||
|
reqwest::Client::new()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn post(
|
||||||
|
c: &reqwest::Client,
|
||||||
|
url: String,
|
||||||
|
body: Value,
|
||||||
|
bearer: Option<&str>,
|
||||||
|
) -> reqwest::Response {
|
||||||
|
let mut req = c.post(url).json(&body);
|
||||||
|
if let Some(b) = bearer {
|
||||||
|
req = req.bearer_auth(b);
|
||||||
|
}
|
||||||
|
req.send().await.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolve_reserve_settle_round_trip() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("resolve_reserve_settle_round_trip").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let raw = format!("sk-{}", Uuid::new_v4());
|
||||||
|
let (account_id, key_id) = seed_key(&pool, 1000, &raw, false).await;
|
||||||
|
let c = client();
|
||||||
|
|
||||||
|
// resolve
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200);
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
assert_eq!(body["principal"]["account_id"], account_id.to_string());
|
||||||
|
assert_eq!(body["principal"]["key_id"], key_id.to_string());
|
||||||
|
assert_eq!(body["snapshot"]["hard_cap"], 1000);
|
||||||
|
|
||||||
|
// reserve 400
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/reserve"),
|
||||||
|
json!({"account_id": account_id, "key_id": key_id, "max_tokens": 400}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200);
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
let rid = body["reservation_id"].as_i64().expect("granted");
|
||||||
|
|
||||||
|
// settle 150
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/settle"),
|
||||||
|
json!({"reservation_id": rid, "actual_tokens": 150}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 204);
|
||||||
|
|
||||||
|
// snapshot reflects spend
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/snapshot"),
|
||||||
|
json!({"account_id": account_id, "key_id": key_id}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
assert_eq!(body["spent"], 150);
|
||||||
|
assert_eq!(body["reserved"], 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn over_cap_reserve_is_rejected_not_errored() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("over_cap_reserve_is_rejected_not_errored").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let raw = format!("sk-{}", Uuid::new_v4());
|
||||||
|
let (account_id, key_id) = seed_key(&pool, 100, &raw, false).await;
|
||||||
|
let c = client();
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/reserve"),
|
||||||
|
json!({"account_id": account_id, "key_id": key_id, "max_tokens": 999}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200, "budget refusal is an authoritative 200");
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
assert!(body["reservation_id"].is_null());
|
||||||
|
assert_eq!(body["rejected"]["kind"], "insufficient_quota");
|
||||||
|
assert_eq!(body["rejected"]["available"], 100);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn deactivated_account_resolves_as_invalid_no_clue() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("deactivated_account_resolves_as_invalid_no_clue").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let raw = format!("sk-{}", Uuid::new_v4());
|
||||||
|
seed_key(&pool, 1000, &raw, true).await; // deactivated
|
||||||
|
let c = client();
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
// Indistinguishable from an unknown key.
|
||||||
|
assert_eq!(r.status(), 401);
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "invalid_api_key");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_client_auth_is_401_before_db() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("missing_client_auth_is_401_before_db").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let raw = format!("sk-{}", Uuid::new_v4());
|
||||||
|
seed_key(&pool, 1000, &raw, false).await;
|
||||||
|
let c = client();
|
||||||
|
// No bearer → rejected by client_auth.
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 401);
|
||||||
|
// Wrong bearer → also rejected.
|
||||||
|
let r = post(
|
||||||
|
&c,
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw}),
|
||||||
|
Some("wrong"),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 401);
|
||||||
|
}
|
||||||
204
crates/helexa-upstream/tests/ledger_pg.rs
Normal file
204
crates/helexa-upstream/tests/ledger_pg.rs
Normal file
@@ -0,0 +1,204 @@
|
|||||||
|
//! Integration tests for the allocation ledger against a real PostgreSQL.
|
||||||
|
//!
|
||||||
|
//! Gated on `UPSTREAM_TEST_DATABASE_URL` — when unset (CI's generic runner,
|
||||||
|
//! local builds without a DB), every test logs a skip and returns, so
|
||||||
|
//! `cargo test --workspace` stays green without Postgres. Point the env var
|
||||||
|
//! at a throwaway database to exercise the no-overshoot guarantee and
|
||||||
|
//! settle/release idempotency:
|
||||||
|
//!
|
||||||
|
//! UPSTREAM_TEST_DATABASE_URL=postgres://helexa:helexa@localhost/helexa_test \
|
||||||
|
//! cargo test -p helexa-upstream --test ledger_pg
|
||||||
|
|
||||||
|
use helexa_upstream::db::connect_and_migrate;
|
||||||
|
use helexa_upstream::ledger::{self, LedgerError};
|
||||||
|
use sqlx::Executor;
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// Returns a migrated pool, or `None` (with a skip log) when the env var is
|
||||||
|
/// unset.
|
||||||
|
async fn pool_or_skip(test: &str) -> Option<PgPool> {
|
||||||
|
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||||
|
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||||
|
return None;
|
||||||
|
};
|
||||||
|
Some(
|
||||||
|
connect_and_migrate(&url, 16)
|
||||||
|
.await
|
||||||
|
.expect("connect + migrate"),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Seed a verified user + account (with `total` allocation) + an active key
|
||||||
|
/// (percent=100 so the account cap binds). Returns (account_id, key_id).
|
||||||
|
async fn seed(pool: &PgPool, total: i64) -> (Uuid, Uuid) {
|
||||||
|
let user_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO users (email, password_hash, email_verified) \
|
||||||
|
VALUES ($1, 'x', true) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(format!("u-{}@test.local", Uuid::new_v4())),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let account_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO accounts (owner_user_id, allocation_total) \
|
||||||
|
VALUES ($1, $2) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(user_id)
|
||||||
|
.bind(total),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let key_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||||
|
VALUES ($1, $2, 'sk-test', 'percent', 100) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(Uuid::new_v4().as_bytes().to_vec()),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
(account_id, key_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn account_cols(pool: &PgPool, account_id: Uuid) -> (i64, i64) {
|
||||||
|
let row = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query("SELECT allocation_spent, allocation_reserved FROM accounts WHERE id = $1")
|
||||||
|
.bind(account_id),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
(row.get("allocation_spent"), row.get("allocation_reserved"))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn concurrent_reserves_never_overshoot() {
|
||||||
|
let Some(pool) = pool_or_skip("concurrent_reserves_never_overshoot").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
// Allocation admits exactly 5 reservations of 100 (cap 500).
|
||||||
|
let (account_id, key_id) = seed(&pool, 500).await;
|
||||||
|
|
||||||
|
let mut handles = Vec::new();
|
||||||
|
for _ in 0..20 {
|
||||||
|
let pool = pool.clone();
|
||||||
|
handles.push(tokio::spawn(async move {
|
||||||
|
ledger::reserve(&pool, account_id, key_id, 100).await
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
let mut ok = 0;
|
||||||
|
let mut quota = 0;
|
||||||
|
for h in handles {
|
||||||
|
match h.await.unwrap() {
|
||||||
|
Ok(_) => ok += 1,
|
||||||
|
Err(LedgerError::InsufficientQuota { .. }) => quota += 1,
|
||||||
|
Err(e) => panic!("unexpected error: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert_eq!(ok, 5, "exactly 5 reserves of 100 fit in a 500 allocation");
|
||||||
|
assert_eq!(quota, 15);
|
||||||
|
|
||||||
|
let (spent, reserved) = account_cols(&pool, account_id).await;
|
||||||
|
assert_eq!(spent, 0);
|
||||||
|
assert_eq!(reserved, 500, "reserved exactly the cap, never over");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn settle_is_idempotent_and_reconciles_spend() {
|
||||||
|
let Some(pool) = pool_or_skip("settle_is_idempotent_and_reconciles_spend").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let (account_id, key_id) = seed(&pool, 1000).await;
|
||||||
|
let rid = ledger::reserve(&pool, account_id, key_id, 400)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Settle actual=150 (< reserved 400): spent=150, reserved back to 0.
|
||||||
|
ledger::settle(&pool, rid, 150).await.unwrap();
|
||||||
|
let (spent, reserved) = account_cols(&pool, account_id).await;
|
||||||
|
assert_eq!((spent, reserved), (150, 0));
|
||||||
|
|
||||||
|
// Second settle is a no-op.
|
||||||
|
ledger::settle(&pool, rid, 999).await.unwrap();
|
||||||
|
let (spent2, reserved2) = account_cols(&pool, account_id).await;
|
||||||
|
assert_eq!((spent2, reserved2), (150, 0), "settle is idempotent");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn release_returns_reservation_and_is_idempotent() {
|
||||||
|
let Some(pool) = pool_or_skip("release_returns_reservation_and_is_idempotent").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let (account_id, key_id) = seed(&pool, 1000).await;
|
||||||
|
let rid = ledger::reserve(&pool, account_id, key_id, 300)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(account_cols(&pool, account_id).await, (0, 300));
|
||||||
|
|
||||||
|
ledger::release(&pool, rid).await.unwrap();
|
||||||
|
assert_eq!(account_cols(&pool, account_id).await, (0, 0));
|
||||||
|
// Idempotent; settle-after-release also a no-op.
|
||||||
|
ledger::release(&pool, rid).await.unwrap();
|
||||||
|
ledger::settle(&pool, rid, 100).await.unwrap();
|
||||||
|
assert_eq!(account_cols(&pool, account_id).await, (0, 0));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn hardcap_key_subcap_binds_below_account() {
|
||||||
|
let Some(pool) = pool_or_skip("hardcap_key_subcap_binds_below_account").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
// Account has 1000 but the key is hard-capped at 200.
|
||||||
|
let user_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO users (email, password_hash, email_verified) \
|
||||||
|
VALUES ($1, 'x', true) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(format!("u-{}@test.local", Uuid::new_v4())),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let account_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO accounts (owner_user_id, allocation_total) VALUES ($1, 1000) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(user_id),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let key_id: Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO api_keys (account_id, key_hash, key_prefix, limit_kind, limit_value) \
|
||||||
|
VALUES ($1, $2, 'sk-test', 'hardcap', 200) RETURNING id",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(Uuid::new_v4().as_bytes().to_vec()),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
|
||||||
|
ledger::reserve(&pool, account_id, key_id, 200)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
match ledger::reserve(&pool, account_id, key_id, 1).await {
|
||||||
|
Err(LedgerError::InsufficientQuota { available, .. }) => assert_eq!(available, 0),
|
||||||
|
other => panic!("expected InsufficientQuota, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
116
crates/helexa-upstream/tests/served_usage_pg.rs
Normal file
116
crates/helexa-upstream/tests/served_usage_pg.rs
Normal file
@@ -0,0 +1,116 @@
|
|||||||
|
//! Integration test for the served-usage report (#58): the idempotent,
|
||||||
|
//! monotonic upsert and the reconcile rollup. Gated on
|
||||||
|
//! UPSTREAM_TEST_DATABASE_URL (skips cleanly when unset).
|
||||||
|
|
||||||
|
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||||
|
use helexa_upstream::db::connect_and_migrate;
|
||||||
|
use helexa_upstream::email::EmailSender;
|
||||||
|
use helexa_upstream::reconcile::reconcile;
|
||||||
|
use helexa_upstream::state::AppState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
const CLIENT_TOKEN: &str = "su-test-token";
|
||||||
|
const OPERATOR: &str = "op-su-test";
|
||||||
|
|
||||||
|
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||||
|
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||||
|
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||||
|
return None;
|
||||||
|
};
|
||||||
|
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||||
|
let mut config = UpstreamConfig {
|
||||||
|
server: Default::default(),
|
||||||
|
db: helexa_upstream::config::DbSettings {
|
||||||
|
url,
|
||||||
|
max_connections: 16,
|
||||||
|
},
|
||||||
|
grant: Default::default(),
|
||||||
|
abuse: Default::default(),
|
||||||
|
client_auth: Default::default(),
|
||||||
|
authz: Default::default(),
|
||||||
|
auth: Default::default(),
|
||||||
|
email: Default::default(),
|
||||||
|
};
|
||||||
|
config.client_auth.tokens.push(ClientToken {
|
||||||
|
token: CLIENT_TOKEN.into(),
|
||||||
|
operator_id: OPERATOR.into(),
|
||||||
|
});
|
||||||
|
let email = EmailSender::from_config(&config.email).unwrap();
|
||||||
|
let state = AppState::new(pool.clone(), config, email);
|
||||||
|
let app = helexa_upstream::build_app(state);
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
Some((format!("http://{addr}"), pool))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn report(base: &str, rows: Value) -> u16 {
|
||||||
|
reqwest::Client::new()
|
||||||
|
.post(format!("{base}/authz/v1/served-usage"))
|
||||||
|
.bearer_auth(CLIENT_TOKEN)
|
||||||
|
.json(&json!({ "rows": rows }))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.status()
|
||||||
|
.as_u16()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn stored(pool: &PgPool, account: Uuid, key: Uuid) -> i64 {
|
||||||
|
sqlx::query(
|
||||||
|
"SELECT served_tokens FROM served_usage WHERE operator_id = $1 AND account_id = $2 AND key_id = $3",
|
||||||
|
)
|
||||||
|
.bind(OPERATOR)
|
||||||
|
.bind(account)
|
||||||
|
.bind(key)
|
||||||
|
.fetch_one(pool)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("served_tokens")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn served_usage_upsert_is_monotonic_and_reconciles() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("served_usage_upsert_is_monotonic_and_reconciles").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let account = Uuid::new_v4();
|
||||||
|
let key = Uuid::new_v4();
|
||||||
|
let period = "2026-06-23";
|
||||||
|
let row = |n: i64| json!([{"account_id": account, "key_id": key, "period": period, "served_tokens": n}]);
|
||||||
|
|
||||||
|
// First report.
|
||||||
|
assert_eq!(report(&base, row(100)).await, 204);
|
||||||
|
assert_eq!(stored(&pool, account, key).await, 100);
|
||||||
|
|
||||||
|
// Re-send a higher absolute value → advances.
|
||||||
|
assert_eq!(report(&base, row(250)).await, 204);
|
||||||
|
assert_eq!(stored(&pool, account, key).await, 250);
|
||||||
|
|
||||||
|
// A lower value (e.g. a restarted cortex) must NOT regress (GREATEST).
|
||||||
|
assert_eq!(report(&base, row(50)).await, 204);
|
||||||
|
assert_eq!(stored(&pool, account, key).await, 250);
|
||||||
|
|
||||||
|
// Re-sending the same value is idempotent.
|
||||||
|
assert_eq!(report(&base, row(250)).await, 204);
|
||||||
|
assert_eq!(stored(&pool, account, key).await, 250);
|
||||||
|
|
||||||
|
// Reconcile rolls it up and stamps reconciled_at; a second run is empty.
|
||||||
|
let rollup = reconcile(&pool).await.unwrap();
|
||||||
|
let mine = rollup
|
||||||
|
.iter()
|
||||||
|
.find(|r| r.operator_id == OPERATOR)
|
||||||
|
.expect("operator in rollup");
|
||||||
|
assert!(mine.total_served_tokens >= 250);
|
||||||
|
let again = reconcile(&pool).await.unwrap();
|
||||||
|
assert!(
|
||||||
|
again.iter().all(|r| r.operator_id != OPERATOR),
|
||||||
|
"already reconciled"
|
||||||
|
);
|
||||||
|
}
|
||||||
425
crates/helexa-upstream/tests/web_pg.rs
Normal file
425
crates/helexa-upstream/tests/web_pg.rs
Normal file
@@ -0,0 +1,425 @@
|
|||||||
|
//! Integration tests for the `/web/v1` account API + the silent fingerprint
|
||||||
|
//! abuse policy, driving the built app over HTTP against a real Postgres.
|
||||||
|
//! Gated on `UPSTREAM_TEST_DATABASE_URL` (skips cleanly when unset).
|
||||||
|
|
||||||
|
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||||
|
use helexa_upstream::crypto::sha256;
|
||||||
|
use helexa_upstream::db::connect_and_migrate;
|
||||||
|
use helexa_upstream::email::EmailSender;
|
||||||
|
use helexa_upstream::state::AppState;
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use sqlx::Executor;
|
||||||
|
use sqlx::Row;
|
||||||
|
use sqlx::postgres::PgPool;
|
||||||
|
|
||||||
|
const CLIENT_TOKEN: &str = "web-test-operator-token";
|
||||||
|
|
||||||
|
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||||
|
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||||
|
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||||
|
return None;
|
||||||
|
};
|
||||||
|
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||||
|
let mut config = UpstreamConfig {
|
||||||
|
server: Default::default(),
|
||||||
|
db: helexa_upstream::config::DbSettings {
|
||||||
|
url,
|
||||||
|
max_connections: 16,
|
||||||
|
},
|
||||||
|
grant: Default::default(),
|
||||||
|
abuse: Default::default(),
|
||||||
|
client_auth: Default::default(),
|
||||||
|
authz: Default::default(),
|
||||||
|
auth: Default::default(),
|
||||||
|
email: Default::default(), // Log transport
|
||||||
|
};
|
||||||
|
config.client_auth.tokens.push(ClientToken {
|
||||||
|
token: CLIENT_TOKEN.into(),
|
||||||
|
operator_id: "op-web".into(),
|
||||||
|
});
|
||||||
|
let email = EmailSender::from_config(&config.email).unwrap();
|
||||||
|
let state = AppState::new(pool.clone(), config, email);
|
||||||
|
let app = helexa_upstream::build_app(state);
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
axum::serve(listener, app).await.unwrap();
|
||||||
|
});
|
||||||
|
Some((format!("http://{addr}"), pool))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn unique_email() -> String {
|
||||||
|
format!("u-{}@test.local", uuid::Uuid::new_v4())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn post(url: String, body: Value, bearer: Option<&str>) -> reqwest::Response {
|
||||||
|
let c = reqwest::Client::new();
|
||||||
|
let mut req = c.post(url).json(&body);
|
||||||
|
if let Some(b) = bearer {
|
||||||
|
req = req.bearer_auth(b);
|
||||||
|
}
|
||||||
|
req.send().await.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn verify_endpoint_consumes_token_once() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("verify_endpoint_consumes_token_once").await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let email = unique_email();
|
||||||
|
// Register, then mint a verify token directly (the raw token is only in
|
||||||
|
// the email; here we insert a known one to drive the endpoint).
|
||||||
|
assert_eq!(
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/register"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.status(),
|
||||||
|
202
|
||||||
|
);
|
||||||
|
let user_id: uuid::Uuid = pool
|
||||||
|
.fetch_one(sqlx::query("SELECT id FROM users WHERE email = $1").bind(&email))
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let raw = "verify-raw-token-xyz";
|
||||||
|
pool.execute(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO email_tokens (token_hash, user_id, kind, expires_at) \
|
||||||
|
VALUES ($1, $2, 'verify', now() + interval '1 hour')",
|
||||||
|
)
|
||||||
|
.bind(sha256(raw))
|
||||||
|
.bind(user_id),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
post(format!("{base}/web/v1/verify"), json!({"token": raw}), None)
|
||||||
|
.await
|
||||||
|
.status(),
|
||||||
|
200
|
||||||
|
);
|
||||||
|
// Consumed → second attempt fails.
|
||||||
|
assert_eq!(
|
||||||
|
post(format!("{base}/web/v1/verify"), json!({"token": raw}), None)
|
||||||
|
.await
|
||||||
|
.status(),
|
||||||
|
400
|
||||||
|
);
|
||||||
|
|
||||||
|
let verified: bool = pool
|
||||||
|
.fetch_one(sqlx::query("SELECT email_verified FROM users WHERE id = $1").bind(user_id))
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("email_verified");
|
||||||
|
assert!(verified);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn account_lifecycle_and_key_resolves_then_archives() {
|
||||||
|
let Some((base, pool)) =
|
||||||
|
spawn_or_skip("account_lifecycle_and_key_resolves_then_archives").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let email = unique_email();
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/register"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
// Bypass the email step for the login/key portion.
|
||||||
|
pool.execute(
|
||||||
|
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// login → session JWT
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/login"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200);
|
||||||
|
let token = r.json::<Value>().await.unwrap()["token"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
// create key (raw shown once)
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/keys"),
|
||||||
|
json!({"label": "laptop"}),
|
||||||
|
Some(&token),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 201);
|
||||||
|
let body: Value = r.json().await.unwrap();
|
||||||
|
let raw_key = body["key"].as_str().unwrap().to_string();
|
||||||
|
let key_id = body["id"].as_str().unwrap().to_string();
|
||||||
|
assert!(raw_key.starts_with("sk-helexa-"));
|
||||||
|
|
||||||
|
// account balance reflects the free grant
|
||||||
|
let r = reqwest::Client::new()
|
||||||
|
.get(format!("{base}/web/v1/account"))
|
||||||
|
.bearer_auth(&token)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
r.json::<Value>().await.unwrap()["allocation_total"],
|
||||||
|
1_000_000
|
||||||
|
);
|
||||||
|
|
||||||
|
// list keys shows the prefix, never the raw secret
|
||||||
|
let r = reqwest::Client::new()
|
||||||
|
.get(format!("{base}/web/v1/keys"))
|
||||||
|
.bearer_auth(&token)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let listed = r.json::<Value>().await.unwrap();
|
||||||
|
let k = &listed["keys"][0];
|
||||||
|
assert_eq!(k["id"], key_id);
|
||||||
|
assert!(k.get("key").is_none(), "raw secret never listed");
|
||||||
|
assert!(k["prefix"].as_str().unwrap().starts_with("sk-helexa-"));
|
||||||
|
|
||||||
|
// the key authorizes at the authz surface
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw_key}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200);
|
||||||
|
|
||||||
|
// archive → the key no longer resolves
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/keys/{key_id}/archive"),
|
||||||
|
json!({}),
|
||||||
|
Some(&token),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 204);
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw_key}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 401);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn fingerprint_abuse_silently_deactivates_all_no_clue() {
|
||||||
|
let Some((base, pool)) =
|
||||||
|
spawn_or_skip("fingerprint_abuse_silently_deactivates_all_no_clue").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let fp = format!("fp-{}", uuid::Uuid::new_v4());
|
||||||
|
|
||||||
|
// 5 registrations sharing one fingerprint — every one returns a normal 202.
|
||||||
|
let mut emails = Vec::new();
|
||||||
|
for _ in 0..5 {
|
||||||
|
let email = unique_email();
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/register"),
|
||||||
|
json!({"email": email, "password": "password123", "fingerprint": fp}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 202, "registration always looks successful");
|
||||||
|
emails.push(email);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Silent effect: all 5 accounts are deactivated + flagged.
|
||||||
|
let (deactivated, flagged): (i64, i64) = {
|
||||||
|
let row = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"SELECT \
|
||||||
|
count(*) FILTER (WHERE a.status = 'deactivated') AS d, \
|
||||||
|
count(*) FILTER (WHERE a.fingerprint_flagged) AS f \
|
||||||
|
FROM accounts a JOIN users u ON u.id = a.owner_user_id \
|
||||||
|
WHERE u.registration_fingerprint = $1",
|
||||||
|
)
|
||||||
|
.bind(&fp),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
(row.get("d"), row.get("f"))
|
||||||
|
};
|
||||||
|
assert_eq!(deactivated, 5, "all sharing accounts silently deactivated");
|
||||||
|
assert_eq!(flagged, 5);
|
||||||
|
|
||||||
|
// No clue at the authz surface: a key on a deactivated account resolves
|
||||||
|
// as an ordinary 401, indistinguishable from an unknown key.
|
||||||
|
let acct: uuid::Uuid = pool
|
||||||
|
.fetch_one(
|
||||||
|
sqlx::query(
|
||||||
|
"SELECT a.id FROM accounts a JOIN users u ON u.id = a.owner_user_id \
|
||||||
|
WHERE u.registration_fingerprint = $1 LIMIT 1",
|
||||||
|
)
|
||||||
|
.bind(&fp),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.get("id");
|
||||||
|
let raw = "sk-helexa-deactivated-probe";
|
||||||
|
pool.execute(
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO api_keys (account_id, key_hash, key_prefix) VALUES ($1, $2, 'sk-helexa-')",
|
||||||
|
)
|
||||||
|
.bind(acct)
|
||||||
|
.bind(sha256(raw)),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/authz/v1/resolve"),
|
||||||
|
json!({"api_key": raw}),
|
||||||
|
Some(CLIENT_TOKEN),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(
|
||||||
|
r.status(),
|
||||||
|
401,
|
||||||
|
"deactivated account's key looks like any invalid key"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn topup_redeem_raises_allocation_single_use() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("topup_redeem_raises_allocation_single_use").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let email = unique_email();
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/register"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
pool.execute(
|
||||||
|
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let token = post(
|
||||||
|
format!("{base}/web/v1/login"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.json::<Value>()
|
||||||
|
.await
|
||||||
|
.unwrap()["token"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
// Mint a code worth 500_000 (mint path used by the CLI/faucet).
|
||||||
|
let codes = helexa_upstream::topup::mint(&pool, 500_000, 1, Some("test"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let code = &codes[0];
|
||||||
|
|
||||||
|
// Redeem → allocation_total rises from the 1_000_000 free grant.
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/redeem"),
|
||||||
|
json!({"code": code}),
|
||||||
|
Some(&token),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 200);
|
||||||
|
assert_eq!(
|
||||||
|
r.json::<Value>().await.unwrap()["allocation_total"],
|
||||||
|
1_500_000
|
||||||
|
);
|
||||||
|
|
||||||
|
// Single-use: a second redemption fails generically (no oracle).
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/redeem"),
|
||||||
|
json!({"code": code}),
|
||||||
|
Some(&token),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 400);
|
||||||
|
|
||||||
|
// Unknown code: same generic 400.
|
||||||
|
let r = post(
|
||||||
|
format!("{base}/web/v1/redeem"),
|
||||||
|
json!({"code": "helexa-topup-does-not-exist"}),
|
||||||
|
Some(&token),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(r.status(), 400);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn topup_concurrent_double_redeem_one_winner() {
|
||||||
|
let Some((base, pool)) = spawn_or_skip("topup_concurrent_double_redeem_one_winner").await
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
// Two verified accounts.
|
||||||
|
let mut tokens = Vec::new();
|
||||||
|
for _ in 0..2 {
|
||||||
|
let email = unique_email();
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/register"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
pool.execute(
|
||||||
|
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let t = post(
|
||||||
|
format!("{base}/web/v1/login"),
|
||||||
|
json!({"email": email, "password": "password123"}),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.json::<Value>()
|
||||||
|
.await
|
||||||
|
.unwrap()["token"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.to_string();
|
||||||
|
tokens.push(t);
|
||||||
|
}
|
||||||
|
let code = helexa_upstream::topup::mint(&pool, 100, 1, None)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.remove(0);
|
||||||
|
|
||||||
|
// Both accounts race to redeem the same code; exactly one wins.
|
||||||
|
let (a, b) = tokio::join!(
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/redeem"),
|
||||||
|
json!({"code": code}),
|
||||||
|
Some(&tokens[0])
|
||||||
|
),
|
||||||
|
post(
|
||||||
|
format!("{base}/web/v1/redeem"),
|
||||||
|
json!({"code": code}),
|
||||||
|
Some(&tokens[1])
|
||||||
|
),
|
||||||
|
);
|
||||||
|
let wins = [a.status(), b.status()]
|
||||||
|
.iter()
|
||||||
|
.filter(|s| s.as_u16() == 200)
|
||||||
|
.count();
|
||||||
|
assert_eq!(wins, 1, "exactly one redemption wins the single-use code");
|
||||||
|
}
|
||||||
@@ -38,6 +38,7 @@ cudnn = [
|
|||||||
flash-attn = [
|
flash-attn = [
|
||||||
"cuda",
|
"cuda",
|
||||||
"candle-transformers/flash-attn",
|
"candle-transformers/flash-attn",
|
||||||
|
"dep:candle-flash-attn",
|
||||||
]
|
]
|
||||||
# Reserved for GPU-only integration tests in later stages.
|
# Reserved for GPU-only integration tests in later stages.
|
||||||
cuda-integration = ["cuda"]
|
cuda-integration = ["cuda"]
|
||||||
@@ -71,6 +72,10 @@ rayon = "1"
|
|||||||
candle-core = "0.10.2"
|
candle-core = "0.10.2"
|
||||||
candle-nn = "0.10.2"
|
candle-nn = "0.10.2"
|
||||||
candle-transformers = "0.10.2"
|
candle-transformers = "0.10.2"
|
||||||
|
# Direct dependency so `candle_flash_attn::flash_attn` is nameable in
|
||||||
|
# the attention core (#95); the candle-transformers feature alone only
|
||||||
|
# enables flash inside ITS models.
|
||||||
|
candle-flash-attn = { version = "0.10.2", optional = true }
|
||||||
# Direct dep on cudarc (matching candle's transitive version) so the
|
# Direct dep on cudarc (matching candle's transitive version) so the
|
||||||
# TP worker pool can call cudarc::nccl::{Comm, Id} directly. Gated on
|
# TP worker pool can call cudarc::nccl::{Comm, Id} directly. Gated on
|
||||||
# the `cuda` feature; same toolchain requirement as candle's CUDA path.
|
# the `cuda` feature; same toolchain requirement as candle's CUDA path.
|
||||||
|
|||||||
@@ -13,9 +13,10 @@ use axum::response::sse::{Event, KeepAlive, Sse};
|
|||||||
use axum::response::{IntoResponse, Json};
|
use axum::response::{IntoResponse, Json};
|
||||||
use axum::routing::{get, post};
|
use axum::routing::{get, post};
|
||||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||||
|
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID};
|
||||||
use cortex_core::harness::ModelSpec;
|
use cortex_core::harness::ModelSpec;
|
||||||
use cortex_core::openai::{ChatCompletionRequest, MessageContent};
|
use cortex_core::openai::{ChatCompletionRequest, MessageContent};
|
||||||
use cortex_core::responses::{ResponsesRequest, ResponsesUsage};
|
use cortex_core::responses::{OutputTokensDetails, ResponsesRequest, ResponsesUsage};
|
||||||
use futures::stream::{self, StreamExt};
|
use futures::stream::{self, StreamExt};
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use std::convert::Infallible;
|
use std::convert::Infallible;
|
||||||
@@ -71,6 +72,12 @@ async fn health_handler(State(state): State<Arc<NeuronState>>) -> Json<HealthRes
|
|||||||
// know about activation lifecycle.
|
// know about activation lifecycle.
|
||||||
let mut snapshot = state.health_cache.snapshot().await;
|
let mut snapshot = state.health_cache.snapshot().await;
|
||||||
snapshot.activation = state.activation.snapshot().await;
|
snapshot.activation = state.activation.snapshot().await;
|
||||||
|
// Per-model admission load (#53) — read live from the candle harness so
|
||||||
|
// cortex's load-aware router (#55) can spread traffic and propagate
|
||||||
|
// backpressure. Absent when no candle harness is present.
|
||||||
|
if let Some(candle) = &state.candle {
|
||||||
|
snapshot.models = candle.load_snapshot().await;
|
||||||
|
}
|
||||||
Json(snapshot)
|
Json(snapshot)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -228,6 +235,17 @@ fn default_enable_thinking(req: &mut ChatCompletionRequest, include_thinking: bo
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The request's principal for fair-share admission (#54), reconstructed
|
||||||
|
/// from the internal headers cortex stamps (#49). cortex strips any
|
||||||
|
/// client-supplied copy and asserts the authoritative value, so over the
|
||||||
|
/// trusted WireGuard link these are safe to key fair-share on. `None` for an
|
||||||
|
/// unauthenticated/direct request — exempt from the per-principal cap.
|
||||||
|
fn principal_key(headers: &axum::http::HeaderMap) -> Option<String> {
|
||||||
|
let account = headers.get(HEADER_ACCOUNT_ID)?.to_str().ok()?;
|
||||||
|
let key = headers.get(HEADER_KEY_ID)?.to_str().ok()?;
|
||||||
|
Some(format!("{account}/{key}"))
|
||||||
|
}
|
||||||
|
|
||||||
/// OpenAI-compatible chat completions. Dispatches to streaming SSE when
|
/// OpenAI-compatible chat completions. Dispatches to streaming SSE when
|
||||||
/// `stream: true` is set on the request; otherwise returns a single
|
/// `stream: true` is set on the request; otherwise returns a single
|
||||||
/// `ChatCompletionResponse`.
|
/// `ChatCompletionResponse`.
|
||||||
@@ -271,8 +289,14 @@ async fn chat_completions(
|
|||||||
// true`) keep reasoning on.
|
// true`) keep reasoning on.
|
||||||
default_enable_thinking(&mut req, include_thinking);
|
default_enable_thinking(&mut req, include_thinking);
|
||||||
|
|
||||||
|
// Fair-share admission principal (#54), from cortex's stamped headers.
|
||||||
|
let principal = principal_key(&headers);
|
||||||
|
|
||||||
if req.stream.unwrap_or(false) {
|
if req.stream.unwrap_or(false) {
|
||||||
match candle.chat_completion_stream_with(req, chat_config).await {
|
match candle
|
||||||
|
.chat_completion_stream_with(req, chat_config, principal)
|
||||||
|
.await
|
||||||
|
{
|
||||||
Ok(rx) => {
|
Ok(rx) => {
|
||||||
// Each chunk → one SSE `data: {json}` line. After the
|
// Each chunk → one SSE `data: {json}` line. After the
|
||||||
// channel closes, append the OpenAI [DONE] terminator.
|
// channel closes, append the OpenAI [DONE] terminator.
|
||||||
@@ -289,7 +313,7 @@ async fn chat_completions(
|
|||||||
Err(e) => inference_error_response(e),
|
Err(e) => inference_error_response(e),
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
match candle.chat_completion(req).await {
|
match candle.chat_completion(req, principal).await {
|
||||||
Ok(resp) => Json(resp).into_response(),
|
Ok(resp) => Json(resp).into_response(),
|
||||||
Err(e) => inference_error_response(e),
|
Err(e) => inference_error_response(e),
|
||||||
}
|
}
|
||||||
@@ -302,6 +326,7 @@ async fn chat_completions(
|
|||||||
/// event stream into the Responses event family.
|
/// event stream into the Responses event family.
|
||||||
async fn responses(
|
async fn responses(
|
||||||
State(state): State<Arc<NeuronState>>,
|
State(state): State<Arc<NeuronState>>,
|
||||||
|
headers: axum::http::HeaderMap,
|
||||||
Json(req): Json<ResponsesRequest>,
|
Json(req): Json<ResponsesRequest>,
|
||||||
) -> impl IntoResponse {
|
) -> impl IntoResponse {
|
||||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||||
@@ -336,9 +361,12 @@ async fn responses(
|
|||||||
};
|
};
|
||||||
chat_req.stream = Some(stream_requested);
|
chat_req.stream = Some(stream_requested);
|
||||||
|
|
||||||
|
// Fair-share admission principal (#54), from cortex's stamped headers.
|
||||||
|
let principal = principal_key(&headers);
|
||||||
|
|
||||||
if stream_requested {
|
if stream_requested {
|
||||||
match candle
|
match candle
|
||||||
.responses_stream(chat_req, response_id, message_item_id)
|
.responses_stream(chat_req, response_id, message_item_id, principal)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(rx) => {
|
Ok(rx) => {
|
||||||
@@ -362,7 +390,7 @@ async fn responses(
|
|||||||
// and translate the result. We don't currently re-tokenise
|
// and translate the result. We don't currently re-tokenise
|
||||||
// to compute usage; the harness returns it via the chat
|
// to compute usage; the harness returns it via the chat
|
||||||
// response and we pass it through.
|
// response and we pass it through.
|
||||||
match candle.chat_completion(chat_req).await {
|
match candle.chat_completion(chat_req, principal).await {
|
||||||
Ok(chat_resp) => {
|
Ok(chat_resp) => {
|
||||||
// Extract the assistant text (chat completions
|
// Extract the assistant text (chat completions
|
||||||
// always emits one choice on the candle path).
|
// always emits one choice on the candle path).
|
||||||
@@ -390,8 +418,14 @@ async fn responses(
|
|||||||
input_tokens: u.prompt_tokens,
|
input_tokens: u.prompt_tokens,
|
||||||
output_tokens: u.completion_tokens,
|
output_tokens: u.completion_tokens,
|
||||||
total_tokens: u.prompt_tokens + u.completion_tokens,
|
total_tokens: u.prompt_tokens + u.completion_tokens,
|
||||||
// Non-streaming reasoning accounting deferred (#64).
|
// Carry the reasoning sub-count through from the chat
|
||||||
output_tokens_details: None,
|
// usage — the non-streaming path now splits off the
|
||||||
|
// `<think>` span and counts it (see `split_off_reasoning`).
|
||||||
|
output_tokens_details: u.completion_tokens_details.as_ref().map(|d| {
|
||||||
|
OutputTokensDetails {
|
||||||
|
reasoning_tokens: d.reasoning_tokens,
|
||||||
|
}
|
||||||
|
}),
|
||||||
input_tokens_details: None,
|
input_tokens_details: None,
|
||||||
});
|
});
|
||||||
let meta = openai_responses::ResponseMeta {
|
let meta = openai_responses::ResponseMeta {
|
||||||
@@ -486,6 +520,24 @@ fn inference_error_response(err: InferenceError) -> axum::response::Response {
|
|||||||
"template_render_failed",
|
"template_render_failed",
|
||||||
format!("chat template could not render this request: {detail}"),
|
format!("chat template could not render this request: {detail}"),
|
||||||
),
|
),
|
||||||
|
// Admission control refused on load (#53): a fast, retryable "busy"
|
||||||
|
// signal. 503 (service busy) + Retry-After; opencode/AI SDK back off.
|
||||||
|
InferenceError::Overloaded { retry_after_secs } => OpenAiError::new(
|
||||||
|
503,
|
||||||
|
"rate_limit_error",
|
||||||
|
"rate_limit_exceeded",
|
||||||
|
"model is busy (admission queue full); retry shortly",
|
||||||
|
)
|
||||||
|
.with_retry_after(retry_after_secs),
|
||||||
|
// Per-principal fair-share cap (#54): 429 rate_limit_exceeded +
|
||||||
|
// Retry-After — the caller is sending too many concurrent requests.
|
||||||
|
InferenceError::PerPrincipalLimit { retry_after_secs } => OpenAiError::new(
|
||||||
|
429,
|
||||||
|
"rate_limit_error",
|
||||||
|
"rate_limit_exceeded",
|
||||||
|
"too many concurrent requests for this key; retry shortly",
|
||||||
|
)
|
||||||
|
.with_retry_after(retry_after_secs),
|
||||||
InferenceError::Other(e) => OpenAiError::without_code(500, "api_error", format!("{e:#}")),
|
InferenceError::Other(e) => OpenAiError::without_code(500, "api_error", format!("{e:#}")),
|
||||||
};
|
};
|
||||||
envelope_response(env)
|
envelope_response(env)
|
||||||
@@ -660,6 +712,26 @@ mod error_envelope_tests {
|
|||||||
assert_eq!(error["required_mb"], 8_192);
|
assert_eq!(error["required_mb"], 8_192);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn overloaded_is_503_rate_limited_with_retry_after() {
|
||||||
|
// Admission rejection (#53) → fast, retryable backpressure.
|
||||||
|
let resp = inference_error_response(InferenceError::Overloaded {
|
||||||
|
retry_after_secs: 7,
|
||||||
|
});
|
||||||
|
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
let retry = resp
|
||||||
|
.headers()
|
||||||
|
.get(axum::http::header::RETRY_AFTER)
|
||||||
|
.expect("admission rejection must advertise Retry-After");
|
||||||
|
assert_eq!(retry.to_str().unwrap(), "7");
|
||||||
|
|
||||||
|
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let body: Value = serde_json::from_slice(&bytes).unwrap();
|
||||||
|
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn insufficient_vram_carries_retry_after() {
|
async fn insufficient_vram_carries_retry_after() {
|
||||||
// Transient 503 — VRAM frees as in-flight requests finish, so the
|
// Transient 503 — VRAM frees as in-flight requests finish, so the
|
||||||
|
|||||||
@@ -85,6 +85,68 @@ pub struct CandleHarnessConfig {
|
|||||||
/// `/models`, and enforces it. These knobs tune that derivation.
|
/// `/models`, and enforces it. These knobs tune that derivation.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub context_limit: ContextLimitConfig,
|
pub context_limit: ContextLimitConfig,
|
||||||
|
|
||||||
|
/// Admission control (#53): bounds the per-model wait queue so a busy
|
||||||
|
/// model returns a fast, retryable `429`/`503` instead of stalling new
|
||||||
|
/// requests until their client times out.
|
||||||
|
#[serde(default)]
|
||||||
|
pub admission: AdmissionConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `[harness.candle.admission]` settings (#53).
|
||||||
|
///
|
||||||
|
/// Inference is batch-1, so `max_in_flight` is 1 in practice; the queue
|
||||||
|
/// (`max_queue_depth`) absorbs short bursts, and `max_wait_secs` caps how
|
||||||
|
/// long a queued request waits before it's refused with backpressure.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AdmissionConfig {
|
||||||
|
/// Concurrent running requests per model. Batch-1 inference → 1.
|
||||||
|
#[serde(default = "default_admission_max_in_flight")]
|
||||||
|
pub max_in_flight: usize,
|
||||||
|
/// Queued (waiting) requests allowed beyond the in-flight one. The
|
||||||
|
/// `(max_in_flight + max_queue_depth + 1)`-th request is refused
|
||||||
|
/// immediately with `429`/`503` + `Retry-After`.
|
||||||
|
#[serde(default = "default_admission_max_queue_depth")]
|
||||||
|
pub max_queue_depth: usize,
|
||||||
|
/// Maximum seconds a queued request waits for the in-flight slot before
|
||||||
|
/// it is refused (turns the old ~300s client-side hang into a fast,
|
||||||
|
/// honest signal).
|
||||||
|
#[serde(default = "default_admission_max_wait_secs")]
|
||||||
|
pub max_wait_secs: u64,
|
||||||
|
/// Per-principal fair-share cap (#54): max in-flight + queued requests
|
||||||
|
/// for any single principal (resolved from the `x-helexa-*` headers
|
||||||
|
/// cortex stamps), so one client can't monopolize the queue while others
|
||||||
|
/// wait. Over-cap → `429 rate_limit_exceeded` + `Retry-After`. `0`
|
||||||
|
/// disables the cap; anonymous requests are always exempt.
|
||||||
|
#[serde(default = "default_admission_max_per_principal")]
|
||||||
|
pub max_per_principal: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AdmissionConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_in_flight: default_admission_max_in_flight(),
|
||||||
|
max_queue_depth: default_admission_max_queue_depth(),
|
||||||
|
max_wait_secs: default_admission_max_wait_secs(),
|
||||||
|
max_per_principal: default_admission_max_per_principal(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_admission_max_in_flight() -> usize {
|
||||||
|
1
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_admission_max_queue_depth() -> usize {
|
||||||
|
8
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_admission_max_wait_secs() -> u64 {
|
||||||
|
30
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_admission_max_per_principal() -> usize {
|
||||||
|
2
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `[harness.candle.prefix_cache]` settings.
|
/// `[harness.candle.prefix_cache]` settings.
|
||||||
|
|||||||
373
crates/neuron/src/harness/admission.rs
Normal file
373
crates/neuron/src/harness/admission.rs
Normal file
@@ -0,0 +1,373 @@
|
|||||||
|
//! Per-model admission control (#53).
|
||||||
|
//!
|
||||||
|
//! Inference against a loaded model is batch-1: one request runs at a time,
|
||||||
|
//! serialized by the model's `inference_lock` (single-GPU) / `pool` mutex
|
||||||
|
//! (TP). Before this, the wait for that lock was an **unbounded FIFO of
|
||||||
|
//! mutex waiters with no timeout** — a busy model made every new request
|
||||||
|
//! hang until its client gave up (~300s) with an opaque error.
|
||||||
|
//!
|
||||||
|
//! [`AdmissionController`] replaces that implicit unbounded wait with an
|
||||||
|
//! explicit bounded scheduler: at most `max_in_flight` running (1, batch-1)
|
||||||
|
//! plus a bounded queue of `max_queue_depth` waiters, each waiting at most
|
||||||
|
//! `max_wait`. When the queue is full or the wait elapses, the request is
|
||||||
|
//! rejected *immediately* — an honest, fast, retryable "busy" signal
|
||||||
|
//! (`429`/`503` + `Retry-After` per #63) instead of a silent stall.
|
||||||
|
//!
|
||||||
|
//! The controller is pure async (no CUDA), so the inference paths just call
|
||||||
|
//! [`AdmissionController::enter`] before taking the inference lock and hold
|
||||||
|
//! the returned [`AdmissionPermit`] for the request's lifetime. Its counters
|
||||||
|
//! ([`in_flight`](AdmissionController::in_flight) /
|
||||||
|
//! [`queue_depth`](AdmissionController::queue_depth)) are lock-free, so
|
||||||
|
//! `/health` can read live load without contending with inference.
|
||||||
|
|
||||||
|
use crate::config::AdmissionConfig;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||||
|
|
||||||
|
/// Why admission was refused. All map to the #63 backpressure envelope
|
||||||
|
/// (`rate_limit_exceeded` + `Retry-After`); they differ in cause (and HTTP
|
||||||
|
/// status — load → `503`, per-principal → `429`).
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
pub enum AdmissionRejection {
|
||||||
|
/// The bounded wait queue was already full (server-side load).
|
||||||
|
QueueFull { retry_after_secs: u64 },
|
||||||
|
/// A queue slot was taken but the in-flight slot didn't free within
|
||||||
|
/// `max_wait` (server-side load).
|
||||||
|
Timeout { retry_after_secs: u64 },
|
||||||
|
/// This principal already has `max_per_principal` requests in flight or
|
||||||
|
/// queued (#54 fair-share) — one principal can't monopolize the model.
|
||||||
|
PrincipalCap { retry_after_secs: u64 },
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AdmissionRejection {
|
||||||
|
pub fn retry_after_secs(&self) -> u64 {
|
||||||
|
match self {
|
||||||
|
AdmissionRejection::QueueFull { retry_after_secs }
|
||||||
|
| AdmissionRejection::Timeout { retry_after_secs }
|
||||||
|
| AdmissionRejection::PrincipalCap { retry_after_secs } => *retry_after_secs,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Admission accounting, mutated under a brief lock (never held across an
|
||||||
|
/// await). `pending` is queued + in-flight overall; `per_principal` is the
|
||||||
|
/// same count keyed by principal for fair-share (#54).
|
||||||
|
#[derive(Default, Debug)]
|
||||||
|
struct AdmissionState {
|
||||||
|
pending: usize,
|
||||||
|
per_principal: HashMap<String, usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Bounded batch-1 scheduler for one loaded model, with per-principal
|
||||||
|
/// fair-share.
|
||||||
|
pub struct AdmissionController {
|
||||||
|
/// In-flight slots — `max_in_flight` permits (1 for batch-1).
|
||||||
|
slots: Arc<Semaphore>,
|
||||||
|
/// Queued + in-flight accounting (overall + per principal).
|
||||||
|
state: Arc<Mutex<AdmissionState>>,
|
||||||
|
/// `max_in_flight + max_queue_depth` — the overall rejection threshold.
|
||||||
|
max_pending: usize,
|
||||||
|
/// Max in-flight + queued for any single principal (#54). `0` disables.
|
||||||
|
max_per_principal: usize,
|
||||||
|
max_in_flight: usize,
|
||||||
|
max_wait: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AdmissionController {
|
||||||
|
pub fn new(cfg: &AdmissionConfig) -> Self {
|
||||||
|
// A controller with zero in-flight slots would deadlock; clamp.
|
||||||
|
let max_in_flight = cfg.max_in_flight.max(1);
|
||||||
|
Self {
|
||||||
|
slots: Arc::new(Semaphore::new(max_in_flight)),
|
||||||
|
state: Arc::new(Mutex::new(AdmissionState::default())),
|
||||||
|
max_pending: max_in_flight + cfg.max_queue_depth,
|
||||||
|
max_per_principal: cfg.max_per_principal,
|
||||||
|
max_in_flight,
|
||||||
|
max_wait: Duration::from_secs(cfg.max_wait_secs),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Admit a request for `principal` (`None` = anonymous, exempt from the
|
||||||
|
/// per-principal cap). Reserves a queue slot — fast-rejecting if the
|
||||||
|
/// overall queue is full or the principal is over its fair-share cap —
|
||||||
|
/// then waits up to `max_wait` for an in-flight slot. The returned permit
|
||||||
|
/// must be held for the request's lifetime; dropping it frees the slots.
|
||||||
|
///
|
||||||
|
/// CANCELLATION SAFETY: the semaphore wait below is where a client
|
||||||
|
/// disconnect lands — axum drops the request future mid-await. The
|
||||||
|
/// reservation therefore lives in a RAII [`PendingReservation`] taken
|
||||||
|
/// BEFORE the await: if this future is dropped while queued, the
|
||||||
|
/// guard's Drop rolls the counts back. (The original version
|
||||||
|
/// incremented raw counters and only decremented on the timeout
|
||||||
|
/// branch — every abandoned wait leaked a `pending` + per-principal
|
||||||
|
/// slot, ratcheting the model into a permanent instant-429 state
|
||||||
|
/// under client retry storms. Observed live 2026-07-02:
|
||||||
|
/// `queue_depth: 1` pinned on an idle model.)
|
||||||
|
pub async fn enter(
|
||||||
|
&self,
|
||||||
|
principal: Option<&str>,
|
||||||
|
) -> Result<AdmissionPermit, AdmissionRejection> {
|
||||||
|
// Decision + reservation under one brief lock so concurrent callers
|
||||||
|
// can't both slip past the thresholds. No await is held here.
|
||||||
|
let reservation = {
|
||||||
|
let mut st = self.state.lock().expect("admission state poisoned");
|
||||||
|
if st.pending >= self.max_pending {
|
||||||
|
return Err(AdmissionRejection::QueueFull {
|
||||||
|
retry_after_secs: self.retry_hint(st.pending),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if let Some(p) = principal
|
||||||
|
&& self.max_per_principal > 0
|
||||||
|
&& st.per_principal.get(p).copied().unwrap_or(0) >= self.max_per_principal
|
||||||
|
{
|
||||||
|
return Err(AdmissionRejection::PrincipalCap {
|
||||||
|
retry_after_secs: self.retry_hint(st.pending),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
st.pending += 1;
|
||||||
|
if let Some(p) = principal {
|
||||||
|
*st.per_principal.entry(p.to_string()).or_insert(0) += 1;
|
||||||
|
}
|
||||||
|
PendingReservation {
|
||||||
|
state: Arc::clone(&self.state),
|
||||||
|
principal: principal.map(str::to_string),
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match tokio::time::timeout(self.max_wait, Arc::clone(&self.slots).acquire_owned()).await {
|
||||||
|
Ok(Ok(permit)) => Ok(AdmissionPermit {
|
||||||
|
_permit: permit,
|
||||||
|
_reservation: reservation,
|
||||||
|
}),
|
||||||
|
// Semaphore is never closed; treat a closed/elapsed wait the
|
||||||
|
// same. `reservation` drops here, rolling back the counts.
|
||||||
|
Ok(Err(_)) | Err(_) => Err(AdmissionRejection::Timeout {
|
||||||
|
retry_after_secs: self.retry_hint(self.max_pending),
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Requests currently running (holding an in-flight slot).
|
||||||
|
pub fn in_flight(&self) -> usize {
|
||||||
|
self.max_in_flight
|
||||||
|
.saturating_sub(self.slots.available_permits())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Requests waiting for an in-flight slot.
|
||||||
|
pub fn queue_depth(&self) -> usize {
|
||||||
|
let pending = self.state.lock().expect("admission state poisoned").pending;
|
||||||
|
pending.saturating_sub(self.in_flight())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rough `Retry-After`: scale with how backed-up the model is, clamped to
|
||||||
|
/// a sane band. Without per-request timing this is a heuristic, but it
|
||||||
|
/// gives well-behaved clients (opencode/AI SDK) a sensible backoff.
|
||||||
|
fn retry_hint(&self, pending: usize) -> u64 {
|
||||||
|
let queued = pending.saturating_sub(self.max_in_flight) as u64;
|
||||||
|
((queued + 1) * 2).clamp(1, 120)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decrement (and prune at zero) a principal's outstanding count.
|
||||||
|
fn decrement_principal(map: &mut HashMap<String, usize>, principal: Option<&str>) {
|
||||||
|
if let Some(p) = principal
|
||||||
|
&& let Some(count) = map.get_mut(p)
|
||||||
|
{
|
||||||
|
*count -= 1;
|
||||||
|
if *count == 0 {
|
||||||
|
map.remove(p);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RAII accounting for one reserved slot (queued or in-flight): decrements
|
||||||
|
/// `pending` and the principal's fair-share count on drop, whichever way
|
||||||
|
/// the reservation ends — admitted-and-finished, wait timeout, or the
|
||||||
|
/// caller's future being dropped mid-queue (client disconnect).
|
||||||
|
#[derive(Debug)]
|
||||||
|
struct PendingReservation {
|
||||||
|
state: Arc<Mutex<AdmissionState>>,
|
||||||
|
principal: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for PendingReservation {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let mut st = self.state.lock().expect("admission state poisoned");
|
||||||
|
st.pending = st.pending.saturating_sub(1);
|
||||||
|
decrement_principal(&mut st.per_principal, self.principal.as_deref());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Held for a request's lifetime; frees the in-flight slot (semaphore
|
||||||
|
/// permit) and the queue + fair-share accounting (reservation) on drop.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct AdmissionPermit {
|
||||||
|
_permit: OwnedSemaphorePermit,
|
||||||
|
_reservation: PendingReservation,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// Config with the per-principal cap disabled (0) — most tests exercise
|
||||||
|
/// the overall queue with anonymous (`None`) callers.
|
||||||
|
fn cfg(max_in_flight: usize, max_queue_depth: usize, max_wait_secs: u64) -> AdmissionConfig {
|
||||||
|
AdmissionConfig {
|
||||||
|
max_in_flight,
|
||||||
|
max_queue_depth,
|
||||||
|
max_wait_secs,
|
||||||
|
max_per_principal: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn admits_up_to_in_flight_and_reports_load() {
|
||||||
|
let ctrl = AdmissionController::new(&cfg(1, 4, 30));
|
||||||
|
assert_eq!(ctrl.in_flight(), 0);
|
||||||
|
let p = ctrl.enter(None).await.expect("first admits");
|
||||||
|
assert_eq!(ctrl.in_flight(), 1);
|
||||||
|
assert_eq!(ctrl.queue_depth(), 0);
|
||||||
|
drop(p);
|
||||||
|
assert_eq!(ctrl.in_flight(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rejects_when_queue_full() {
|
||||||
|
// 1 in-flight + 1 queue slot = capacity 2; the 3rd is refused fast.
|
||||||
|
let ctrl = Arc::new(AdmissionController::new(&cfg(1, 1, 30)));
|
||||||
|
let _running = ctrl.enter(None).await.expect("admit running");
|
||||||
|
|
||||||
|
// Fill the single queue slot with a waiter that parks on the semaphore.
|
||||||
|
let ctrl2 = Arc::clone(&ctrl);
|
||||||
|
let waiter = tokio::spawn(async move { ctrl2.enter(None).await.map(|p| drop(p)) });
|
||||||
|
// Give the waiter a moment to occupy the queue slot.
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
assert_eq!(ctrl.queue_depth(), 1);
|
||||||
|
|
||||||
|
// Queue full → immediate QueueFull with a Retry-After hint.
|
||||||
|
match ctrl.enter(None).await {
|
||||||
|
Err(AdmissionRejection::QueueFull { retry_after_secs }) => {
|
||||||
|
assert!(retry_after_secs >= 1)
|
||||||
|
}
|
||||||
|
other => panic!("expected QueueFull, got {other:?}"),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Release the runner so the parked waiter can proceed and finish.
|
||||||
|
drop(_running);
|
||||||
|
waiter.await.unwrap().unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rejects_on_wait_timeout() {
|
||||||
|
// Zero queue depth + a runner holding the only slot → a second
|
||||||
|
// request can't even queue, so it's QueueFull, not Timeout. Use a
|
||||||
|
// queue of 1 and a tiny max_wait to exercise the timeout path.
|
||||||
|
let ctrl = Arc::new(AdmissionController::new(&cfg(1, 1, 0)));
|
||||||
|
let _running = ctrl.enter(None).await.expect("admit running");
|
||||||
|
// max_wait 0 → the queued request times out almost immediately.
|
||||||
|
match ctrl.enter(None).await {
|
||||||
|
Err(AdmissionRejection::Timeout { .. }) => {}
|
||||||
|
other => panic!("expected Timeout, got {other:?}"),
|
||||||
|
}
|
||||||
|
// The timed-out request released its queue slot.
|
||||||
|
assert_eq!(ctrl.queue_depth(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn per_principal_cap_protects_other_principals() {
|
||||||
|
// Generous overall queue, but each principal capped at 1 in-flight+
|
||||||
|
// queued. Principal A holds the running slot; A's second request is
|
||||||
|
// refused (PrincipalCap) rather than occupying the queue, so B's
|
||||||
|
// single request still gets a queue slot and proceeds.
|
||||||
|
let cfg = AdmissionConfig {
|
||||||
|
max_in_flight: 1,
|
||||||
|
max_queue_depth: 8,
|
||||||
|
max_wait_secs: 30,
|
||||||
|
max_per_principal: 1,
|
||||||
|
};
|
||||||
|
let ctrl = Arc::new(AdmissionController::new(&cfg));
|
||||||
|
|
||||||
|
let _a1 = ctrl.enter(Some("acct-a/key-a")).await.expect("A admits");
|
||||||
|
|
||||||
|
// A is over its fair-share cap → fast PrincipalCap, no queue slot taken.
|
||||||
|
match ctrl.enter(Some("acct-a/key-a")).await {
|
||||||
|
Err(AdmissionRejection::PrincipalCap { retry_after_secs }) => {
|
||||||
|
assert!(retry_after_secs >= 1)
|
||||||
|
}
|
||||||
|
other => panic!("expected PrincipalCap, got {other:?}"),
|
||||||
|
}
|
||||||
|
|
||||||
|
// B (a different principal) is admitted to the queue and proceeds
|
||||||
|
// once A releases — it was never stuck behind A's backlog.
|
||||||
|
let ctrl2 = Arc::clone(&ctrl);
|
||||||
|
let b = tokio::spawn(async move { ctrl2.enter(Some("acct-b/key-b")).await.map(drop) });
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
assert_eq!(ctrl.queue_depth(), 1, "B is queued, not rejected");
|
||||||
|
drop(_a1);
|
||||||
|
b.await.unwrap().expect("B is served after A releases");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Regression for the 2026-07-02 retry-storm incident: a client that
|
||||||
|
/// disconnects while QUEUED drops the `enter()` future mid-await.
|
||||||
|
/// The reservation must roll back — the original implementation
|
||||||
|
/// leaked `pending` + the per-principal count on this path, pinning
|
||||||
|
/// the model in a permanent instant-429 state.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cancelled_queued_waiter_rolls_back_accounting() {
|
||||||
|
let cfg = AdmissionConfig {
|
||||||
|
max_in_flight: 1,
|
||||||
|
max_queue_depth: 2,
|
||||||
|
max_wait_secs: 30,
|
||||||
|
// Cap 3 lets the runner + both waiters coexist; if the two
|
||||||
|
// cancelled waiters leaked their counts, the principal would
|
||||||
|
// sit at 3 == cap and the post-cancel enter below would hit
|
||||||
|
// PrincipalCap instead of queueing.
|
||||||
|
max_per_principal: 3,
|
||||||
|
};
|
||||||
|
let ctrl = Arc::new(AdmissionController::new(&cfg));
|
||||||
|
let running = ctrl.enter(Some("acct/key")).await.expect("admit running");
|
||||||
|
|
||||||
|
// Two waiters from the same principal park in the queue…
|
||||||
|
let mut waiters = Vec::new();
|
||||||
|
for _ in 0..2 {
|
||||||
|
let c = Arc::clone(&ctrl);
|
||||||
|
waiters.push(tokio::spawn(async move {
|
||||||
|
c.enter(Some("acct/key")).await.map(drop)
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
assert_eq!(ctrl.queue_depth(), 2);
|
||||||
|
|
||||||
|
// …and both clients vanish (abort = the dropped request future).
|
||||||
|
for w in &waiters {
|
||||||
|
w.abort();
|
||||||
|
}
|
||||||
|
for w in waiters {
|
||||||
|
let _ = w.await;
|
||||||
|
}
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
ctrl.queue_depth(),
|
||||||
|
0,
|
||||||
|
"cancelled waiters must not leak queue slots"
|
||||||
|
);
|
||||||
|
|
||||||
|
// The principal's fair-share count must also be clean: with the
|
||||||
|
// runner still holding 1 of its cap of 3, a new request from the
|
||||||
|
// same principal queues instead of hitting PrincipalCap (which a
|
||||||
|
// leak of the two cancelled counts would trigger).
|
||||||
|
let c = Arc::clone(&ctrl);
|
||||||
|
let retry = tokio::spawn(async move { c.enter(Some("acct/key")).await.map(drop) });
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
assert_eq!(ctrl.queue_depth(), 1, "post-cancel request queues normally");
|
||||||
|
drop(running);
|
||||||
|
retry
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.expect("post-cancel request is served — no leaked principal count");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -22,6 +22,7 @@ use super::TextConfig;
|
|||||||
use super::full_attn::Qwen3_5Attention;
|
use super::full_attn::Qwen3_5Attention;
|
||||||
use super::linear_attn::GatedDeltaNet;
|
use super::linear_attn::GatedDeltaNet;
|
||||||
use super::mlp::Qwen3_5MLP;
|
use super::mlp::Qwen3_5MLP;
|
||||||
|
use super::moe::Qwen3_5MoeBlock;
|
||||||
use super::rmsnorm::Qwen3_5RmsNorm;
|
use super::rmsnorm::Qwen3_5RmsNorm;
|
||||||
use super::rope::RotaryEmbedding;
|
use super::rope::RotaryEmbedding;
|
||||||
use super::snapshot::LayerKvSnapshot;
|
use super::snapshot::LayerKvSnapshot;
|
||||||
@@ -35,10 +36,27 @@ enum AttentionKind {
|
|||||||
Linear(GatedDeltaNet),
|
Linear(GatedDeltaNet),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The FFN slot: dense SwiGLU (Qwen3.6) or the high-sparsity MoE block
|
||||||
|
/// (qwen3_next 80B-A3B family, #92), selected per layer by
|
||||||
|
/// [`TextConfig::layer_uses_moe`].
|
||||||
|
enum MlpKind {
|
||||||
|
Dense(Qwen3_5MLP),
|
||||||
|
Moe(Qwen3_5MoeBlock),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Module for MlpKind {
|
||||||
|
fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||||
|
match self {
|
||||||
|
MlpKind::Dense(mlp) => mlp.forward(x),
|
||||||
|
MlpKind::Moe(moe) => moe.forward(x),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub struct Qwen3_5DecoderLayer {
|
pub struct Qwen3_5DecoderLayer {
|
||||||
input_layernorm: Qwen3_5RmsNorm,
|
input_layernorm: Qwen3_5RmsNorm,
|
||||||
post_attention_layernorm: Qwen3_5RmsNorm,
|
post_attention_layernorm: Qwen3_5RmsNorm,
|
||||||
mlp: Qwen3_5MLP,
|
mlp: MlpKind,
|
||||||
attention: AttentionKind,
|
attention: AttentionKind,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -73,7 +91,11 @@ impl Qwen3_5DecoderLayer {
|
|||||||
),
|
),
|
||||||
};
|
};
|
||||||
|
|
||||||
let mlp = Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?;
|
let mlp = if cfg.layer_uses_moe(layer_idx) {
|
||||||
|
MlpKind::Moe(Qwen3_5MoeBlock::load(cfg, &vb.pp("mlp"))?)
|
||||||
|
} else {
|
||||||
|
MlpKind::Dense(Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?)
|
||||||
|
};
|
||||||
let input_layernorm =
|
let input_layernorm =
|
||||||
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
||||||
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user