diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..a9d3ab0 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,8 @@ +# The usage ledger builds SQLite in-tree (rusqlite "bundled"), which needs a +# C compiler. cargo-dist's release CI builds aarch64-unknown-linux-musl on +# GitHub's arm64 runners with only `musl-tools` installed, whose compiler is +# `musl-gcc`; the cc crate defaults to `aarch64-linux-musl-gcc` for that +# target, which musl-tools does not provide. Point it at musl-gcc, the same +# compiler the x86_64-unknown-linux-musl build already uses. +[env] +CC_aarch64_unknown_linux_musl = "musl-gcc" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e95eb05..fefde3c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,5 +42,7 @@ jobs: with: toolchain: 1.96.1 - uses: Swatinem/rust-cache@v2 - - run: cargo install cargo-shear --locked + # Pinned: cargo-shear >= 1.14 pulls ra_ap_* crates that require + # rustc 1.98, above the 1.96.1 toolchain pinned above. Bump both together. + - run: cargo install cargo-shear --locked --version 1.13.0 - run: cargo shear diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..1dfcd0f --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,296 @@ +# This file was autogenerated by dist: https://axodotdev.github.io/cargo-dist +# +# Copyright 2022-2024, axodotdev +# SPDX-License-Identifier: MIT or Apache-2.0 +# +# CI that: +# +# * checks for a Git Tag that looks like a release +# * builds artifacts with dist (archives, installers, hashes) +# * uploads those artifacts to temporary workflow zip +# * on success, uploads the artifacts to a GitHub Release +# +# Note that the GitHub Release will be created with a generated +# title/body based on your changelogs. + +name: Release +permissions: + "contents": "write" + +# This task will run whenever you push a git tag that looks like a version +# like "1.0.0", "v0.1.0-prerelease.1", "my-app/0.1.0", "releases/v1.0.0", etc. +# Various formats will be parsed into a VERSION and an optional PACKAGE_NAME, where +# PACKAGE_NAME must be the name of a Cargo package in your workspace, and VERSION +# must be a Cargo-style SemVer Version (must have at least major.minor.patch). +# +# If PACKAGE_NAME is specified, then the announcement will be for that +# package (erroring out if it doesn't have the given version or isn't dist-able). +# +# If PACKAGE_NAME isn't specified, then the announcement will be for all +# (dist-able) packages in the workspace with that version (this mode is +# intended for workspaces with only one dist-able package, or with all dist-able +# packages versioned/released in lockstep). +# +# If you push multiple tags at once, separate instances of this workflow will +# spin up, creating an independent announcement for each one. However, GitHub +# will hard limit this to 3 tags per commit, as it will assume more tags is a +# mistake. +# +# If there's a prerelease-style suffix to the version, then the release(s) +# will be marked as a prerelease. +on: + pull_request: + push: + tags: + - '**[0-9]+.[0-9]+.[0-9]+*' + +jobs: + # Run 'dist plan' (or host) to determine what tasks we need to do + plan: + runs-on: "ubuntu-22.04" + outputs: + val: ${{ steps.plan.outputs.manifest }} + tag: ${{ !github.event.pull_request && github.ref_name || '' }} + tag-flag: ${{ !github.event.pull_request && format('--tag={0}', github.ref_name) || '' }} + publishing: ${{ !github.event.pull_request }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install dist + # we specify bash to get pipefail; it guards against the `curl` command + # failing. otherwise `sh` won't catch that `curl` returned non-0 + shell: bash + run: "curl --proto '=https' --tlsv1.2 -LsSf https://github.com/axodotdev/cargo-dist/releases/download/v0.32.0/cargo-dist-installer.sh | sh" + - name: Cache dist + uses: actions/upload-artifact@v7 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/dist + # sure would be cool if github gave us proper conditionals... + # so here's a doubly-nested ternary-via-truthiness to try to provide the best possible + # functionality based on whether this is a pull_request, and whether it's from a fork. + # (PRs run on the *source* but secrets are usually on the *target* -- that's *good* + # but also really annoying to build CI around when it needs secrets to work right.) + - id: plan + run: | + dist ${{ (!github.event.pull_request && format('host --steps=create --tag={0}', github.ref_name)) || 'plan' }} --output-format=json > plan-dist-manifest.json + echo "dist ran successfully" + cat plan-dist-manifest.json + echo "manifest=$(jq -c "." plan-dist-manifest.json)" >> "$GITHUB_OUTPUT" + - name: "Upload dist-manifest.json" + uses: actions/upload-artifact@v7 + with: + name: artifacts-plan-dist-manifest + path: plan-dist-manifest.json + + # Build and packages all the platform-specific things + build-local-artifacts: + name: build-local-artifacts (${{ join(matrix.targets, ', ') }}) + # Let the initial task tell us to not run (currently very blunt) + needs: + - plan + if: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix.include != null && (needs.plan.outputs.publishing == 'true' || fromJson(needs.plan.outputs.val).ci.github.pr_run_mode == 'upload') }} + strategy: + fail-fast: false + # Target platforms/runners are computed by dist in create-release. + # Each member of the matrix has the following arguments: + # + # - runner: the github runner + # - dist-args: cli flags to pass to dist + # - install-dist: expression to run to install dist on the runner + # + # Typically there will be: + # - 1 "global" task that builds universal installers + # - N "local" tasks that build each platform's binaries and platform-specific installers + matrix: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix }} + runs-on: ${{ matrix.runner }} + container: ${{ matrix.container && matrix.container.image || null }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + BUILD_MANIFEST_NAME: target/distrib/${{ join(matrix.targets, '-') }}-dist-manifest.json + steps: + - name: enable windows longpaths + run: | + git config --global core.longpaths true + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install Rust non-interactively if not already installed + if: ${{ matrix.container }} + run: | + if ! command -v cargo > /dev/null 2>&1; then + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + echo "$HOME/.cargo/bin" >> $GITHUB_PATH + fi + - name: Install dist + run: ${{ matrix.install_dist.run }} + # Get the dist-manifest + - name: Fetch local artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - name: Install dependencies + run: | + ${{ matrix.packages_install }} + - name: Build artifacts + run: | + # Actually do builds and make zips and whatnot + dist build ${{ needs.plan.outputs.tag-flag }} --print=linkage --output-format=json ${{ matrix.dist_args }} > dist-manifest.json + echo "dist ran successfully" + - id: cargo-dist + name: Post-build + # We force bash here just because github makes it really hard to get values up + # to "real" actions without writing to env-vars, and writing to env-vars has + # inconsistent syntax between shell and powershell. + shell: bash + run: | + # Parse out what we just built and upload it to scratch storage + echo "paths<> "$GITHUB_OUTPUT" + dist print-upload-files-from-manifest --manifest dist-manifest.json >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + cp dist-manifest.json "$BUILD_MANIFEST_NAME" + - name: "Upload artifacts" + uses: actions/upload-artifact@v7 + with: + name: artifacts-build-local-${{ join(matrix.targets, '_') }} + path: | + ${{ steps.cargo-dist.outputs.paths }} + ${{ env.BUILD_MANIFEST_NAME }} + + # Build and package all the platform-agnostic(ish) things + build-global-artifacts: + needs: + - plan + - build-local-artifacts + runs-on: "ubuntu-22.04" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + BUILD_MANIFEST_NAME: target/distrib/global-dist-manifest.json + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install cached dist + uses: actions/download-artifact@v8 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/ + - run: chmod +x ~/.cargo/bin/dist + # Get all the local artifacts for the global tasks to use (for e.g. checksums) + - name: Fetch local artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - id: cargo-dist + shell: bash + run: | + dist build ${{ needs.plan.outputs.tag-flag }} --output-format=json "--artifacts=global" > dist-manifest.json + echo "dist ran successfully" + + # Parse out what we just built and upload it to scratch storage + echo "paths<> "$GITHUB_OUTPUT" + jq --raw-output ".upload_files[]" dist-manifest.json >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + cp dist-manifest.json "$BUILD_MANIFEST_NAME" + - name: "Upload artifacts" + uses: actions/upload-artifact@v7 + with: + name: artifacts-build-global + path: | + ${{ steps.cargo-dist.outputs.paths }} + ${{ env.BUILD_MANIFEST_NAME }} + # Determines if we should publish/announce + host: + needs: + - plan + - build-local-artifacts + - build-global-artifacts + # Only run if we're "publishing", and only if plan, local and global didn't fail (skipped is fine) + if: ${{ always() && needs.plan.result == 'success' && needs.plan.outputs.publishing == 'true' && (needs.build-global-artifacts.result == 'skipped' || needs.build-global-artifacts.result == 'success') && (needs.build-local-artifacts.result == 'skipped' || needs.build-local-artifacts.result == 'success') }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + runs-on: "ubuntu-22.04" + outputs: + val: ${{ steps.host.outputs.manifest }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install cached dist + uses: actions/download-artifact@v8 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/ + - run: chmod +x ~/.cargo/bin/dist + # Fetch artifacts from scratch-storage + - name: Fetch artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - id: host + shell: bash + run: | + dist host ${{ needs.plan.outputs.tag-flag }} --steps=upload --steps=release --output-format=json > dist-manifest.json + echo "artifacts uploaded and released successfully" + cat dist-manifest.json + echo "manifest=$(jq -c "." dist-manifest.json)" >> "$GITHUB_OUTPUT" + - name: "Upload dist-manifest.json" + uses: actions/upload-artifact@v7 + with: + # Overwrite the previous copy + name: artifacts-dist-manifest + path: dist-manifest.json + # Create a GitHub Release while uploading all files to it + - name: "Download GitHub Artifacts" + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: artifacts + merge-multiple: true + - name: Cleanup + run: | + # Remove the granular manifests + rm -f artifacts/*-dist-manifest.json + - name: Create GitHub Release + env: + PRERELEASE_FLAG: "${{ fromJson(steps.host.outputs.manifest).announcement_is_prerelease && '--prerelease' || '' }}" + ANNOUNCEMENT_TITLE: "${{ fromJson(steps.host.outputs.manifest).announcement_title }}" + ANNOUNCEMENT_BODY: "${{ fromJson(steps.host.outputs.manifest).announcement_github_body }}" + RELEASE_COMMIT: "${{ github.sha }}" + run: | + # Write and read notes from a file to avoid quoting breaking things + echo "$ANNOUNCEMENT_BODY" > $RUNNER_TEMP/notes.txt + + gh release create "${{ needs.plan.outputs.tag }}" --target "$RELEASE_COMMIT" $PRERELEASE_FLAG --title "$ANNOUNCEMENT_TITLE" --notes-file "$RUNNER_TEMP/notes.txt" artifacts/* + + announce: + needs: + - plan + - host + # use "always() && ..." to allow us to wait for all publish jobs while + # still allowing individual publish jobs to skip themselves (for prereleases). + # "host" however must run to completion, no skipping allowed! + if: ${{ always() && needs.host.result == 'success' }} + runs-on: "ubuntu-22.04" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive diff --git a/.gitignore b/.gitignore index 96f35bd..db36c96 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,7 @@ !/ref/REFS .playwright-mcp + +.z-proflow/ +.z-flashflow/ +.zcode/ diff --git a/Cargo.lock b/Cargo.lock index cb4e128..6feb413 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -111,6 +111,18 @@ dependencies = [ "rustversion", ] +[[package]] +name = "argon2" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "134c52ddac6d63c576bef8168db10c83c49c26444ecbc68060fef078925a901c" +dependencies = [ + "base64ct", + "blake2", + "cpufeatures 0.3.1", + "password-hash", +] + [[package]] name = "arrayvec" version = "0.7.6" @@ -205,6 +217,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.5.3" @@ -238,6 +256,15 @@ dependencies = [ "wyz", ] +[[package]] +name = "blake2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -247,6 +274,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "borsh" version = "1.7.0" @@ -393,6 +429,12 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "colorchoice" version = "1.0.5" @@ -445,6 +487,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -470,6 +521,24 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03bb0e1cc970d482d121d9a1744999169b69a07470b3d644a7894e53fcaf4574" +dependencies = [ + "cmov", +] + [[package]] name = "darling" version = "0.20.11" @@ -566,8 +635,19 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -618,6 +698,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fancy-regex" version = "0.13.0" @@ -1749,6 +1841,18 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash 0.2.0", +] + +[[package]] +name = "hashlink" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb" +dependencies = [ + "hashbrown 0.17.1", +] [[package]] name = "heapless" @@ -1827,6 +1931,15 @@ dependencies = [ "serde", ] +[[package]] +name = "hybrid-array" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27f864f10dfb56725ce5ce5472bc52252c8f93a4ab86327122cebf62c5f59a17" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -2155,6 +2268,17 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsqlite3-sys" +version = "0.38.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -2183,8 +2307,11 @@ dependencies = [ "insta", "libc", "llm-proxy-core", + "llm-proxy-protocol", "llm-proxy-provider", "llm-proxy-server", + "llm-proxy-storage", + "secrecy", "socket2", "tempfile", "time", @@ -2198,6 +2325,25 @@ dependencies = [ [[package]] name = "llm-proxy-api" version = "0.1.0" +dependencies = [ + "axum", + "humantime", + "llm-proxy-protocol", + "llm-proxy-storage", + "rusqlite", + "rust_decimal", + "serde", + "serde_json", + "sha2", + "tempfile", + "thiserror", + "time", + "tokio", + "tower", + "tower-http", + "tracing", + "tracing-subscriber", +] [[package]] name = "llm-proxy-core" @@ -2252,9 +2398,11 @@ dependencies = [ name = "llm-proxy-server" version = "0.1.0" dependencies = [ + "argon2", "axum", "bytes", "futures", + "llm-proxy-api", "llm-proxy-core", "llm-proxy-protocol", "llm-proxy-provider", @@ -2265,6 +2413,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "tempfile", "thiserror", "time", "tokio", @@ -2274,6 +2423,7 @@ dependencies = [ "tower", "tower-http", "tracing", + "tracing-subscriber", "uuid", ] @@ -2282,9 +2432,15 @@ name = "llm-proxy-storage" version = "0.1.0" dependencies = [ "llm-proxy-protocol", + "rusqlite", + "rust_decimal", "serde", "serde_json", + "tempfile", + "thiserror", "time", + "tokio", + "tracing", ] [[package]] @@ -2504,12 +2660,33 @@ dependencies = [ "windows-link", ] +[[package]] +name = "password-hash" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aab41826031698d6ffcd9cff78ef56ef998e39dc7e5067cdfebe373842d4723b" +dependencies = [ + "getrandom 0.4.2", + "phc", +] + [[package]] name = "percent-encoding" version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phc" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44dc769b75f93afdddd8c7fa12d685292ddeff1e66f7f0f3a234cf1818afe892" +dependencies = [ + "base64ct", + "ctutils", + "getrandom 0.4.2", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2897,6 +3074,31 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "rsqlite-vfs" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" +dependencies = [ + "hashbrown 0.16.1", + "thiserror", +] + +[[package]] +name = "rusqlite" +version = "0.40.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", + "sqlite-wasm-rs", +] + [[package]] name = "rust_decimal" version = "1.42.1" @@ -3150,8 +3352,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -3160,7 +3362,7 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89f599ac0c323ebb1c6082821a54962b839832b03984598375bff3975b804423" dependencies = [ - "digest", + "digest 0.10.7", "sha1", ] @@ -3171,8 +3373,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -3256,6 +3458,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "sqlite-wasm-rs" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75" +dependencies = [ + "cc", + "js-sys", + "rsqlite-vfs", + "wasm-bindgen", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index d454d12..3ca0631 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,7 @@ llm-proxy-core = { path = "crates/llm-proxy-core" } llm-proxy-server = { path = "crates/llm-proxy-server" } # Async / runtime -tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "time"] } +tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } tokio-util = { version = "0.7", features = ["rt"] } # Web @@ -50,6 +50,8 @@ anyhow = "1" # Secrets secrecy = { version = "0.10", features = ["serde"] } +# Password hashing KDF for client-key verifiers (PHC strings) +argon2 = "0.6" # Tokenization tiktoken-rs = "0.6" @@ -57,6 +59,10 @@ tiktoken-rs = "0.6" # Money (cost/pricing) -- fixed-point, never float rust_decimal = "1" +# Persistence (usage ledger). "bundled" compiles SQLite in-tree so no +# system libsqlite3 is needed on deploy targets. +rusqlite = { version = "0.40", features = ["bundled"] } + # Observability tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json"] } @@ -161,3 +167,8 @@ literal_string_with_formatting_args = "allow" lto = "thin" codegen-units = 1 strip = "symbols" + +# The profile that 'dist' will build with (cargo-dist passes --profile dist) +[profile.dist] +inherits = "release" +lto = "thin" diff --git a/apps/llm-proxy/Cargo.toml b/apps/llm-proxy/Cargo.toml index 0c595c7..25984f5 100644 --- a/apps/llm-proxy/Cargo.toml +++ b/apps/llm-proxy/Cargo.toml @@ -25,6 +25,7 @@ workspace = true llm-proxy-server = { workspace = true } llm-proxy-core = { workspace = true } llm-proxy-provider = { path = "../../crates/llm-proxy-provider" } +llm-proxy-storage = { path = "../../crates/llm-proxy-storage" } anyhow = { workspace = true } axum = { workspace = true } @@ -43,5 +44,7 @@ vergen-gix = { workspace = true } [dev-dependencies] insta = { version = "1.42", features = ["yaml", "filters"] } +llm-proxy-protocol = { path = "../../crates/llm-proxy-protocol" } +secrecy = { workspace = true } tempfile = "3" toml = { workspace = true } diff --git a/apps/llm-proxy/build.rs b/apps/llm-proxy/build.rs index 888ec71..32f2111 100644 --- a/apps/llm-proxy/build.rs +++ b/apps/llm-proxy/build.rs @@ -4,15 +4,9 @@ use vergen_gix::{Emitter, GixBuilder}; fn main() -> Result<(), Box> { - // VERGEN_GIT_SHA (short) let gix = GixBuilder::default().sha(true).build()?; - // vergen-gix emits its own cargo:rustc-env instructions. We emit the - // TARGET triple separately below since vergen-gix does not expose it. - // The ordering (vergen first, manual TARGET second) is safe because each - // `println!("cargo:rustc-env=...")` call sets an independent env var. Emitter::default().add_instructions(&gix)?.emit()?; - // TARGET is available in build scripts (but not in the compiled crate). println!( "cargo:rustc-env=VERGEN_CARGO_TARGET_TRIPLE={}", std::env::var("TARGET")? diff --git a/apps/llm-proxy/src/cli.rs b/apps/llm-proxy/src/cli.rs index 9900b0a..66be06d 100644 --- a/apps/llm-proxy/src/cli.rs +++ b/apps/llm-proxy/src/cli.rs @@ -86,8 +86,7 @@ pub enum AutostartAction { #[arg(short, long)] port: Option, /// Overwrite an existing launchd plist / .desktop entry instead of - /// aborting. Without this flag `enable` refuses to clobber a - /// hand-edited unit (audit GAP-MED-4). + /// aborting. #[arg(long)] force: bool, }, diff --git a/apps/llm-proxy/src/commands/autostart.rs b/apps/llm-proxy/src/commands/autostart.rs index 2bc19e6..31be136 100644 --- a/apps/llm-proxy/src/commands/autostart.rs +++ b/apps/llm-proxy/src/commands/autostart.rs @@ -1,7 +1,3 @@ -//! `autostart` command implementations. -//! -//! Manages auto-start on login: enable, disable, and status. - use std::path::Path; use anyhow::{Context, Result, bail}; @@ -19,11 +15,8 @@ use crate::platform::format_plist; /// Run the `autostart enable` command. /// -/// If the launchd plist (macOS) or `.desktop` entry (Linux) already exists, -/// the command aborts unless `force` is set, mirroring `cmd_init`'s overwrite -/// guard. This prevents a hand-edited unit (`KeepAlive`, custom -/// `EnvironmentVariables`, `StartInterval`, …) being silently rewritten on the -/// next `enable` (audit GAP-MED-4). +/// Aborts when the launchd plist (macOS) or `.desktop` entry (Linux) already +/// exists, unless `force` is set. pub fn cmd_autostart_enable( config_path: Option<&Path>, port: Option, @@ -34,10 +27,6 @@ pub fn cmd_autostart_enable( .with_context(|| format!("creating directory {}", dir.display()))?; let exe = std::env::current_exe().with_context(|| "resolving current executable")?; - // Build args as strings for plist/desktop entry display. Lossy conversion is - // acceptable here because these are display-only strings that are XML-escaped - // by format_plist(). Non-UTF-8 paths are rare and the lossy replacement is - // sufficient for auto-start purposes. let mut args = vec![exe.to_string_lossy().to_string(), "serve".to_owned()]; if let Some(p) = config_path { @@ -64,7 +53,6 @@ pub fn cmd_autostart_enable( std::fs::write(&plist_path, plist_content) .with_context(|| format!("writing {}", plist_path.display()))?; - // Set restrictive permissions on the plist file. set_private_permissions(&plist_path)?; println!("created launchd plist at {}", plist_path.display()); @@ -81,7 +69,6 @@ pub fn cmd_autostart_enable( require_overwrite_ok(&desktop_path, "desktop entry", force)?; std::fs::write(&desktop_path, desktop_content) .with_context(|| format!("writing {}", desktop_path.display()))?; - // Set restrictive permissions on the desktop entry. set_private_permissions(&desktop_path)?; println!("created desktop entry at {}", desktop_path.display()); } @@ -180,12 +167,6 @@ pub fn cmd_autostart_status() -> Result<()> { Ok(()) } -/// Reject an overwrite of an existing unit file unless `force` is set. -/// -/// Returns `Ok(())` when the path does not exist or `force` is set, otherwise -/// an error naming the file and offering `--force`. Extracted as a pure -/// function so the clobber guard is unit-testable without resolving the real -/// `$HOME`-derived plist/desktop path (audit GAP-MED-4). fn require_overwrite_ok(path: &std::path::Path, kind: &str, force: bool) -> Result<()> { if path.exists() && !force { bail!( diff --git a/apps/llm-proxy/src/commands/init.rs b/apps/llm-proxy/src/commands/init.rs index 5fe64c9..932f655 100644 --- a/apps/llm-proxy/src/commands/init.rs +++ b/apps/llm-proxy/src/commands/init.rs @@ -1,7 +1,3 @@ -//! `init` command implementation. -//! -//! Creates default config files and provider templates. - use std::path::{Path, PathBuf}; use anyhow::{Context, Result, bail}; @@ -13,30 +9,10 @@ use crate::paths::{config_dir, providers_dir_for}; use crate::permissions::{create_private_file, set_private_dir_permissions}; /// Run the `init` command against the default config directory. -/// -/// Thin wrapper around [`cmd_init_at`] using [`config_dir`] as the base, kept -/// as the zero-arg entry point wired into the CLI dispatch. pub fn cmd_init() -> Result<()> { cmd_init_at(&config_dir()) } -/// Run the `init` command against an explicit base directory. -/// -/// Creates, under `base_dir`: -/// - `config.toml` -/// - `providers/opencode-go.toml` -/// - `providers/opencode-zen.toml` -/// -/// Does NOT write: -/// - fallback/scenario JSON -/// - sampling/tool/reasoning/cache/stream overrides -/// -/// Accepting the base directory as a parameter (rather than hard-coding -/// `config_dir()`) lets the full control flow -- the overwrite guard, -/// `create_dir_all` + `set_private_dir_permissions` ordering, and the -/// sequential born-private provider writes -- be exercised against a -/// tempdir in tests instead of the process-global config directory (audit -/// finding app-cli:init-test-does-not-exercise-cmd-init). pub(super) fn cmd_init_at(base_dir: &Path) -> Result<()> { let config_path: PathBuf = base_dir.join("config.toml"); if config_path.exists() { @@ -47,21 +23,17 @@ pub(super) fn cmd_init_at(base_dir: &Path) -> Result<()> { .with_context(|| format!("creating directory {}", base_dir.display()))?; set_private_dir_permissions(base_dir)?; - // Write main config.toml with restrictive permissions (contains env var references). create_private_file(&config_path, DEFAULT_CONFIG_TOML.as_bytes())?; println!("created config at {}", config_path.display()); - // Create providers directory. let providers_dir = providers_dir_for(&config_path); std::fs::create_dir_all(&providers_dir) .with_context(|| format!("creating directory {}", providers_dir.display()))?; - // Write opencode-go provider with restrictive permissions. let go_path = providers_dir.join("opencode-go.toml"); create_private_file(&go_path, DEFAULT_PROVIDER_OPENCODE_GO.as_bytes())?; println!("created provider config at {}", go_path.display()); - // Write opencode-zen provider with restrictive permissions. let zen_path = providers_dir.join("opencode-zen.toml"); create_private_file(&zen_path, DEFAULT_PROVIDER_OPENCODE_ZEN.as_bytes())?; println!("created provider config at {}", zen_path.display()); @@ -81,14 +53,6 @@ pub(super) fn cmd_init_at(base_dir: &Path) -> Result<()> { mod tests { use super::cmd_init_at; - /// Run the real `cmd_init_at` against a tempdir and verify it creates all - /// expected files with valid TOML contents and born-private permissions. - /// - /// Unlike the previous version of this test (which manually `std::fs::write`d - /// the files and never invoked `cmd_init`), this exercises the actual - /// command flow: the overwrite guard, `create_dir_all` + - /// `set_private_dir_permissions`, and the sequential `create_private_file` - /// provider writes. #[test] fn cmd_init_creates_files_with_permissions() { let dir = tempfile::tempdir().unwrap(); @@ -103,7 +67,6 @@ mod tests { assert!(go_path.exists(), "opencode-go.toml was created"); assert!(zen_path.exists(), "opencode-zen.toml was created"); - // Verify file contents are valid TOML with the expected structure. let main_content = std::fs::read_to_string(&config_path).unwrap(); let parsed: toml::Value = toml::from_str(&main_content).expect("generated config.toml must be valid TOML"); @@ -121,8 +84,6 @@ mod tests { let parsed: toml::Value = toml::from_str(&zen_content).unwrap(); assert_eq!(parsed["provider"]["name"].as_str(), Some("opencode-zen")); - // Verify restrictive permissions on Unix via the real create_private_file - // path (the atomic 0600 open), not a separate set_private_permissions call. #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; @@ -139,10 +100,6 @@ mod tests { } } - /// Verify `cmd_init_at` bails when a config already exists and does NOT - /// overwrite the existing file. This exercises the real overwrite guard - /// at the top of the command (previously a tautology that only asserted a - /// file exists right after writing it). #[test] fn cmd_init_rejects_existing_config() { let dir = tempfile::tempdir().unwrap(); @@ -158,15 +115,12 @@ mod tests { "error should mention the config already exists, got: {msg}" ); - // The existing file must be left untouched. assert_eq!( std::fs::read_to_string(&config_path).unwrap(), "existing", "cmd_init_at must not overwrite an existing config" ); - // And the provider files must NOT have been created either, since the - // guard bails before any write. let go_path = dir.path().join("providers").join("opencode-go.toml"); assert!( !go_path.exists(), @@ -174,9 +128,6 @@ mod tests { ); } - /// A second `cmd_init_at` on the same (now-populated) directory must bail - /// rather than silently clobbering the operator's config. Guards the - /// "re-running init" UX path end-to-end. #[test] fn cmd_init_is_idempotent_reject() { let dir = tempfile::tempdir().unwrap(); diff --git a/apps/llm-proxy/src/commands/mod.rs b/apps/llm-proxy/src/commands/mod.rs index eb07363..907e4cb 100644 --- a/apps/llm-proxy/src/commands/mod.rs +++ b/apps/llm-proxy/src/commands/mod.rs @@ -1,6 +1,4 @@ //! Command implementations. -//! -//! Each command is in its own submodule for clarity. mod autostart; mod init; diff --git a/apps/llm-proxy/src/commands/models.rs b/apps/llm-proxy/src/commands/models.rs index 46d5051..a2ca73d 100644 --- a/apps/llm-proxy/src/commands/models.rs +++ b/apps/llm-proxy/src/commands/models.rs @@ -1,7 +1,3 @@ -//! `models` command implementation. -//! -//! Lists provider routes and static, cached, or live model catalogs. - use std::path::PathBuf; use anyhow::{Context, Result, bail}; @@ -15,7 +11,6 @@ use llm_proxy_server::{ModelCatalogService, write_catalog_atomic}; use crate::config_validation::validate_toml_extension; use crate::paths::{providers_dir_for, resolve_config}; -/// Print the configured route adapters for a provider. fn print_provider_routes(provider: &ProviderConfig) { let has_routes = provider.routes.chat_completions.is_some() || provider.routes.messages.is_some(); @@ -31,7 +26,6 @@ fn print_provider_routes(provider: &ProviderConfig) { } } -/// Build the on-disk cache file for a freshly discovered catalog. fn discovered_catalog_file( provider: &ProviderConfig, discovered: Vec, @@ -40,10 +34,6 @@ fn discovered_catalog_file( catalog: CatalogFileMetadata { provider: provider.name.clone(), source: "live".to_owned(), - // Rfc3339 formatting of a UTC OffsetDateTime is infallible in - // practice, but the time crate models it as fallible. Propagate a - // fallback epoch timestamp instead of panicking after a successful - // live fetch, mirroring catalog_service::now_rfc3339. generated_at: time::OffsetDateTime::now_utc() .format(&time::format_description::well_known::Rfc3339) .unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned()), @@ -52,7 +42,6 @@ fn discovered_catalog_file( } } -/// Print the resolved model catalog for a provider. fn print_model_entries(entries: &[StaticModelCatalogEntry]) { println!(" models:"); if entries.is_empty() { @@ -76,10 +65,8 @@ pub async fn cmd_models( ) -> Result<()> { let path = resolve_config(config_path.as_deref()); - // Reject JSON config and non-TOML extensions. validate_toml_extension(&path)?; - // Load providers from the providers/ directory next to the config. let providers_dir = providers_dir_for(&path); let registry = if providers_dir.exists() { ProviderRegistry::load_from_dir(&providers_dir) @@ -121,7 +108,6 @@ pub async fn cmd_models( for provider in &providers { println!(" {}:", provider.name); - // Show routes. print_provider_routes(provider); let catalog_config = provider.catalog.clone().unwrap_or_default(); @@ -132,8 +118,6 @@ pub async fn cmd_models( provider.name ); } - // Built once and only on the live path (avoids a reqwest client - // build when --live is not passed). let discovery = discovery.get_or_insert_with(DiscoveryClient::default); match discovery.discover(provider).await { Ok(discovered) => { diff --git a/apps/llm-proxy/src/commands/serve.rs b/apps/llm-proxy/src/commands/serve.rs index 53fffc7..9447fc2 100644 --- a/apps/llm-proxy/src/commands/serve.rs +++ b/apps/llm-proxy/src/commands/serve.rs @@ -1,14 +1,10 @@ -//! `serve` command implementation. -//! -//! Starts the proxy server in foreground or background (daemon) mode. - use std::path::{Path, PathBuf}; use anyhow::{Context, Result, bail}; use axum::serve::{Listener, ListenerExt}; use llm_proxy_core::LogFormat; use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; -use llm_proxy_server::{build_router, shutdown_signal}; +use llm_proxy_server::{admin_router, build_router, shutdown_signal}; use tokio::net::TcpListener; use tracing::info; @@ -26,17 +22,8 @@ pub async fn cmd_serve( background: bool, daemonize: bool, ) -> Result<()> { - // Initialize tracing before the daemonization check so that the parent - // process can log during spawn_daemon(). The daemon child reinitializes - // its own subscriber; calling init_tracing twice is safe (the second - // call is a no-op since the global subscriber is already set). init_tracing(&LogFormat::from_env()); - // If background mode requested, spawn self as child with --daemonize. - // Run the cheap, pure config checks in the parent first so a clearly-bad - // --config (wrong extension / unresolvable path) produces a normal error - // rather than a spawned-then-immediately-died daemon. Full config loading - // (HTTP client build, TOML parse) stays in the child. if background && !daemonize { let path = resolve_config(config_path.as_deref()); validate_toml_extension(&path)?; @@ -47,26 +34,16 @@ pub async fn cmd_serve( validate_toml_extension(&path)?; - // Build shared infrastructure. let adapter_registry = ProviderAdapterRegistry::builtin(); - // Build the HTTP client fallibly. `reqwest::Client::builder().build()` can - // fail at startup on an invalid `HTTP_PROXY`/`HTTPS_PROXY` env value or a - // TLS-init failure; using `try_new` (audit GAP-LOW-13) lets us surface that - // as a contextualized error instead of panicking the daemon. let proxy_client = ProxyClient::try_new().with_context( || "building proxy HTTP client (check HTTP_PROXY/HTTPS_PROXY and TLS config)", )?; - // Load TOML config and build state before writing the PID file. This - // ensures a bad config does not leave a stale PID file on disk. let state = load_toml_state(&path, adapter_registry, proxy_client, port_override)?; - // Attempt atomic PID file creation. If it fails because the file already - // exists, check whether the existing owner is still alive before removing - // the stale file and retrying. This closes the TOCTOU window between the - // stale-check and the write. + let admin = bind_admin_listener(&state).await?; + if let Err(first_err) = write_pid() { - // Atomic create_new failed -- file exists. Check if the owner is stale. if let Some(existing_pid) = read_pid()? { if is_process_running(existing_pid) { bail!( @@ -78,7 +55,6 @@ pub async fn cmd_serve( info!("stale PID file found (PID {}), cleaning up", existing_pid); } remove_pid()?; - // Retry -- this time the file is gone so create_new should succeed. write_pid().with_context(|| { format!( "failed to create PID file after removing stale entry: {}", @@ -87,19 +63,11 @@ pub async fn cmd_serve( })?; } - // Ensure PID file is cleaned up on shutdown. The two removal sites below - // are mutually exclusive along any path (the deadline branch is followed - // by process::exit), so remove_pid() runs exactly once per shutdown. let bind_addr = state.bind_address(); let shutdown_deadline = state.shutdown_timeout(); + let event_bus = std::sync::Arc::clone(state.event_bus()); let app = build_router(state); - // `tap_io` runs on every accepted connection (audit MEDIUM-2 option A). - // `TCP_NODELAY` is a per-connection socket option that is *not* inherited - // from the listening socket, so it must be applied to each accepted - // `TcpStream` here. Disabling Nagle flushes small SSE chunks to the client - // without the ~40ms coalescing delay, which matters for streaming token - // output. `SO_KEEPALIVE` is set on the listener in `bind_listener`. let listener = bind_listener(bind_addr).await?.tap_io(|tcp_stream| { if let Err(error) = tcp_stream.set_nodelay(true) { tracing::trace!(%error, "failed to set TCP_NODELAY on accepted connection"); @@ -111,47 +79,33 @@ pub async fn cmd_serve( "llm-proxy listening" ); - // Serve with graceful shutdown bounded by a hard deadline. - // - // axum 0.8's `with_graceful_shutdown` waits indefinitely for in-flight - // connections to drain. Because the proxy serves long-lived streaming - // responses an upstream can hold open arbitrarily, one stuck stream could - // otherwise pin the process forever on SIGTERM/SIGINT. We wrap the serve - // future in `tokio::time::timeout`: once the deadline elapses after the - // signal, stuck connections are dropped and we force-exit so an - // orchestrator (Kubernetes/LB) can recycle the pod. `shutdown_timeout == 0` - // disables the deadline (drain forever) for operators who want that. let serve = axum::serve( listener, app.into_make_service_with_connect_info::(), ) .with_graceful_shutdown(shutdown_signal()); - let result = if shutdown_deadline.is_zero() { - serve.await - } else { - match tokio::time::timeout(shutdown_deadline, serve).await { - Ok(r) => r, - Err(_) => { - tracing::error!( - deadline = ?shutdown_deadline, - "graceful-shutdown deadline elapsed; forcing exit with in-flight connections dropped" - ); - // Run PID cleanup best-effort, then exit. The deadline exists - // precisely to escape stuck connections, so we do not wait for - // them; remove_pid is fast, synchronous filesystem I/O. - if let Err(e) = remove_pid() { - tracing::warn!(error = %e, "failed to remove PID file during shutdown"); - } - std::process::exit(1); + let admin_serve = async { + if let Some((listener, router)) = admin { + if let Err(error) = axum::serve(listener, router.into_make_service()) + .with_graceful_shutdown(shutdown_signal()) + .await + { + tracing::error!(%error, "admin usage listener failed"); } } }; + let result = serve_with_deadline(serve.into_future(), admin_serve, shutdown_deadline).await; + if let Err(e) = remove_pid() { tracing::warn!(error = %e, "failed to remove PID file during shutdown"); } + if let Err(close_err) = tokio::task::spawn_blocking(move || event_bus.close()).await { + tracing::warn!(error = %close_err, "usage-ledger close task failed to join"); + } + match result { Ok(()) => { info!("server stopped cleanly"); @@ -164,26 +118,66 @@ pub async fn cmd_serve( } } -/// Bind the listening TCP socket with connection-hardening options. -/// -/// On Unix the socket is created via `socket2` so we can enable `SO_KEEPALIVE` -/// with a bounded idle interval: half-open or stalled peers (a dropped VPN, a -/// NAT timeout, an idle slowloris connection) are probed and reaped instead of -/// pinning a connection forever. Accepted connections inherit the listener's -/// keepalive setting on Linux and macOS. -/// -/// `SO_REUSEADDR` is set so a quick restart after a crash can rebind while the -/// previous socket is in `TIME_WAIT`. -/// -/// # Slowloris note -/// -/// TCP keepalive does not by itself stop a client that drip-feeds request -/// *headers* one byte at a time: `TimeoutLayer` only bounds the fully-buffered -/// handler body, and axum 0.8 exposes no per-connection read-header timeout. -/// For internet-facing deployments, place this proxy behind a reverse proxy -/// (nginx, envoy, a cloud LB) that imposes a read-header timeout and -/// connection cap. For the typical loopback / developer-facing deployment the -/// risk is negligible. +async fn serve_with_deadline( + serve: F, + admin_serve: A, + shutdown_deadline: std::time::Duration, +) -> std::io::Result<()> +where + F: std::future::Future>, + A: std::future::Future, +{ + if shutdown_deadline.is_zero() { + let (main, ()) = tokio::join!(serve, admin_serve); + return main; + } + match tokio::time::timeout(shutdown_deadline, async { + tokio::join!(serve, admin_serve) + }) + .await + { + Ok((main, ())) => main, + Err(_) => { + tracing::error!( + deadline = ?shutdown_deadline, + "graceful-shutdown deadline elapsed; forcing exit with in-flight connections dropped" + ); + if let Err(e) = remove_pid() { + tracing::warn!(error = %e, "failed to remove PID file during shutdown"); + } + std::process::exit(1); + } + } +} + +async fn bind_admin_listener( + state: &llm_proxy_server::AppState, +) -> Result> { + let Some(posture) = state.app_config().admin_posture() else { + return Ok(None); + }; + if !state.app_config().server.dedicated_admin_listener() { + tracing::info!( + posture = posture.name(), + "admin usage API served on the main listener only; no dedicated socket" + ); + return Ok(None); + } + let Some(addr) = state.app_config().server.effective_admin_bind() else { + bail!("admin API enabled but no admin bind could be derived"); + }; + let Some(router) = admin_router(state) else { + bail!("admin API enabled but no /usage router could be built"); + }; + let listener = bind_listener(addr).await?; + info!( + addr = %addr, + posture = posture.name(), + "admin usage API listening (read-only, behind the admin token)" + ); + Ok(Some((listener, router))) +} + async fn bind_listener(bind_addr: std::net::SocketAddr) -> Result { #[cfg(unix)] { @@ -204,11 +198,9 @@ async fn bind_listener(bind_addr: std::net::SocketAddr) -> Result { socket .listen(1024) .with_context(|| format!("listening on {bind_addr}"))?; - // tokio's I/O driver requires a non-blocking socket. socket .set_nonblocking(true) .with_context(|| "setting nonblocking on listener")?; - // Convert the configured socket2 socket into a tokio listener. let std_listener = std::net::TcpListener::from(socket); Ok(TcpListener::from_std(std_listener)?) } @@ -220,7 +212,6 @@ async fn bind_listener(bind_addr: std::net::SocketAddr) -> Result { } } -/// Spawn the current binary as a background daemon. fn spawn_daemon(config_path: Option<&Path>, port_override: Option) -> Result<()> { let exe = std::env::current_exe().with_context(|| "resolving current executable")?; let mut cmd = std::process::Command::new(&exe); @@ -233,21 +224,11 @@ fn spawn_daemon(config_path: Option<&Path>, port_override: Option) -> Resul cmd.arg("--port").arg(p.to_string()); } - // Detach from terminal. #[cfg(unix)] { use std::os::unix::process::CommandExt; cmd.process_group(0); - // Redirect stdout/stderr to a log file. - // NOTE: The log file grows unbounded. For long-running daemons, set up - // external log rotation (e.g., logrotate on Linux, newsyslog on macOS) - // or pipe to a log aggregation service. let log_path = config_dir().join("llm-proxy.log"); - // Open the log file born-private (0600) via `open_private_append`, - // avoiding the create-then-chmod TOCTOU window where the file briefly - // existed world-readable under the process umask. Daemon logs may - // reference requests/config and (per GAP-MED-1) redaction can leak key - // suffixes, so the file must never be readable by other local users. let log_file = crate::permissions::open_private_append(&log_path) .with_context(|| format!("opening log file {}", log_path.display()))?; cmd.stdout( @@ -270,19 +251,7 @@ fn spawn_daemon(config_path: Option<&Path>, port_override: Option) -> Resul let pid = child.id(); - // Write the PID file before detaching so callers can reliably find it. - // - // A failed write here is fatal for the background path: `cmd_stop` would - // print "no PID file found -- server not running" and be unable to stop - // the daemon, and a second `serve --background` is only guarded by the - // child's own create_new PID-file check (serve.rs:63) rather than this - // parent write. Rather than hand the user an unmanageable daemon, kill - // the just-spawned child and bail so they get a clear, recoverable error - // (audit finding app-cli:serve-daemon-pid-write-failure-orphan). if let Err(e) = write_pid_value(pid) { - // Best-effort cleanup: send SIGKILL so the child cannot continue, then - // reap it. We do not propagate a kill/reap failure over the original - // PID-write error -- the original cause is what the user needs to see. if let Err(kill_err) = child.kill() { tracing::error!( error = %kill_err, @@ -304,13 +273,9 @@ fn spawn_daemon(config_path: Option<&Path>, port_override: Option) -> Resul ); } - // Confirm the child is still alive before reporting success. std::thread::sleep(std::time::Duration::from_millis(100)); match child.try_wait() { Ok(Some(status)) => { - // The child died before it could manage its own PID file (it only - // reaches its write_pid() after load_toml_state succeeds), so remove - // the one we wrote to avoid leaving a stale entry on disk. if let Err(e) = crate::pid::remove_pid() { tracing::warn!(error = %e, "failed to remove PID file after daemon child exited immediately"); } @@ -322,7 +287,7 @@ fn spawn_daemon(config_path: Option<&Path>, port_override: Option) -> Resul .unwrap_or_else(|| "signal".to_owned()) ); } - Ok(None) => { /* child is running */ } + Ok(None) => {} Err(e) => { bail!("failed to check daemon child status: {}", e); } diff --git a/apps/llm-proxy/src/commands/status.rs b/apps/llm-proxy/src/commands/status.rs index 8936050..434baf8 100644 --- a/apps/llm-proxy/src/commands/status.rs +++ b/apps/llm-proxy/src/commands/status.rs @@ -1,10 +1,3 @@ -//! `status` command implementation. -//! -//! Checks if the proxy server is running by reading the PID file. -//! -//! Parameterised over a [`CommandPaths`] bundle so the full branching is -//! exercisable against a tempdir-backed manager in tests (audit MEDIUM-6). - use anyhow::Result; use llm_proxy_core::load_app_config; @@ -13,16 +6,14 @@ use crate::pid::{is_process_running, read_pid_from}; /// Run the `status` command against the given paths. /// -/// The `listen` address is read from `paths.config_path()`. If the server was -/// started with `--config` or `LLM_PROXY_CONFIG` pointing elsewhere, the -/// displayed address may not match the actual bind address. +/// The displayed `listen` address is read from `paths.config_path()` and may +/// not match the actual bind address. pub fn cmd_status(paths: &CommandPaths) -> Result<()> { let mgr = paths.pid_manager(); match read_pid_from(mgr)? { Some(pid) => { if is_process_running(pid) { println!("server is running (PID {pid})"); - // Try to load actual bind address from config. let config_path = paths.config_path(); if config_path.exists() { if let Ok(cfg) = load_app_config(config_path) { diff --git a/apps/llm-proxy/src/commands/stop.rs b/apps/llm-proxy/src/commands/stop.rs index d1d43d8..4f39fd5 100644 --- a/apps/llm-proxy/src/commands/stop.rs +++ b/apps/llm-proxy/src/commands/stop.rs @@ -1,11 +1,3 @@ -//! `stop` command implementation. -//! -//! Stops a running proxy server by sending SIGTERM (graceful) then SIGKILL (forced). -//! -//! The command logic is parameterised over a [`CommandPaths`] bundle so the -//! full escalation (stale cleanup, no-file, SIGTERM→poll→SIGKILL) can be -//! exercised against a tempdir-backed manager in tests (audit MEDIUM-6). - use std::time::Duration; use anyhow::{Result, bail}; @@ -13,30 +5,14 @@ use anyhow::{Result, bail}; use crate::paths::CommandPaths; use crate::pid::{STOP_MAX_POLLS, STOP_POLL_INTERVAL_MS, is_process_running, read_pid_from}; -/// Outcome of a best-effort re-read of the PID file immediately before -/// signaling. Used to narrow the TOCTOU / PID-reuse window in `cmd_stop`. -/// -/// Only meaningful on Unix, where `cmd_stop` actually signals processes; gated -/// to match its `#[cfg(unix)]` call sites so non-Unix builds do not see a -/// dead-code warning. #[cfg(unix)] #[derive(Debug)] enum PidOwnership { - /// The PID file still holds exactly `expected`. StillOurs, - /// The PID file now holds a different PID -- refuse to signal `expected`. Changed(u32), - /// The PID file is gone, unreadable, or corrupt -- treat the process as no - /// longer owned/tracked. Gone, } -/// Re-read the PID file and classify it relative to `expected`. -/// -/// Any error from `read_pid_from` (missing file, unreadable, corrupt contents, -/// invalid PID) is collapsed to [`PidOwnership::Gone`] rather than propagated, -/// because this is the *verification* read whose sole job is to decide whether -/// it is still safe to signal `expected`. #[cfg(unix)] fn classify_pid_ownership(mgr: &llm_proxy_core::PidManager, expected: u32) -> PidOwnership { match read_pid_from(mgr) { @@ -49,8 +25,6 @@ fn classify_pid_ownership(mgr: &llm_proxy_core::PidManager, expected: u32) -> Pi /// Run the `stop` command against the given paths using the production poll /// cadence ([`STOP_MAX_POLLS`] × [`STOP_POLL_INTERVAL_MS`]). -/// -/// Thin wrapper around [`cmd_stop_with_timing`] for the real CLI invocation. pub fn cmd_stop(paths: &CommandPaths) -> Result<()> { cmd_stop_with_timing( paths, @@ -62,10 +36,7 @@ pub fn cmd_stop(paths: &CommandPaths) -> Result<()> { /// Run the `stop` command with an explicit graceful-shutdown poll cadence. /// /// `max_polls` × `poll_interval` bounds how long the command waits for the -/// process to exit after SIGTERM before escalating to SIGKILL. The configurable -/// cadence exists so tests can drive the full escalation path quickly against a -/// tempdir-backed manager without waiting the production 10 s budget (audit -/// MEDIUM-6). +/// process to exit after SIGTERM before escalating to SIGKILL. pub fn cmd_stop_with_timing( paths: &CommandPaths, max_polls: u32, @@ -88,28 +59,8 @@ pub fn cmd_stop_with_timing( #[cfg(unix)] { - // Narrow the TOCTOU / PID-reuse window before signaling. - // - // After the initial `read_pid_from` + `is_process_running` we may - // have spent time doing I/O; re-read the PID file immediately - // before sending the signal and refuse to proceed unless the file - // still holds the *same* PID. If the daemon exited and cleaned up - // its PID file (or the file was rewritten with a different PID), - // we no longer own this PID and must not signal it. - // - // NOTE: this is NOT a full PID-identity check. `is_process_running` - // only probes existence via `kill(pid, 0)`, and this re-read only - // confirms the PID file still agrees with the PID we resolved. If - // the OS recycled the *exact* PID for an unrelated process AND the - // stale PID file was never updated, `stop` could still signal a - // bystander. A robust fix needs platform-specific identity - // verification (Linux `/proc//comm`/`cmdline` or process - // start-time; macOS `libproc` `proc_name`), which is intentionally - // out of scope here to avoid pulling extra platform deps. The - // re-read below removes the most common races (PID file cleared on - // clean shutdown, or rewritten by a new `serve`). match classify_pid_ownership(mgr, pid) { - PidOwnership::StillOurs => { /* unchanged */ } + PidOwnership::StillOurs => {} PidOwnership::Changed(other) => { bail!( "PID file changed while stopping (was {pid}, now {other}); \ @@ -148,7 +99,6 @@ pub fn cmd_stop_with_timing( println!("sent SIGTERM to PID {pid}"); } - // Wait up to `max_polls * poll_interval` for graceful shutdown. for _ in 0..max_polls { if !is_process_running(pid) { mgr.remove_pid()?; @@ -158,7 +108,6 @@ pub fn cmd_stop_with_timing( std::thread::sleep(poll_interval); } - // Force kill if still running after graceful period. println!( "WARNING: PID {pid} did not exit within {} seconds, escalating to SIGKILL", max_polls as u64 * poll_interval.as_millis() as u64 / 1000 @@ -166,13 +115,8 @@ pub fn cmd_stop_with_timing( #[cfg(unix)] { - // Re-verify ownership immediately before the forceful SIGKILL, - // which is the most dangerous signal: the polling loop above may - // have run for several seconds, during which the PID file could - // have been cleared or rewritten. See the SIGTERM block above for - // the rationale and the residual PID-reuse caveat. match classify_pid_ownership(mgr, pid) { - PidOwnership::StillOurs => { /* unchanged */ } + PidOwnership::StillOurs => {} PidOwnership::Changed(other) => { bail!( "PID file changed during graceful-shutdown wait \ @@ -181,8 +125,6 @@ pub fn cmd_stop_with_timing( ); } PidOwnership::Gone => { - // The process exited (and cleaned up its PID file) during - // the wait -- no escalation needed. println!("server stopped during graceful-shutdown wait (PID {pid} gone)"); return Ok(()); } diff --git a/apps/llm-proxy/src/commands/validate.rs b/apps/llm-proxy/src/commands/validate.rs index dbebf90..50cdfbf 100644 --- a/apps/llm-proxy/src/commands/validate.rs +++ b/apps/llm-proxy/src/commands/validate.rs @@ -1,7 +1,3 @@ -//! `validate` command implementation. -//! -//! Validates TOML config and prints provider route table. - use std::path::Path; use anyhow::{Context, Result, bail}; @@ -19,7 +15,6 @@ use crate::state::catalog_enforcement_warning; pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { let path = resolve_config(config_path); - // Reject JSON config and non-TOML extensions. validate_toml_extension(&path)?; println!("validating config: {}", path.display()); @@ -27,7 +22,6 @@ pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { let app_config: AppConfig = load_app_config(&path) .with_context(|| format!("loading TOML config from {}", path.display()))?; - // Print settings summary. println!(); println!("=== Server Configuration ==="); println!(" bind: {}", app_config.server.bind); @@ -48,7 +42,6 @@ pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { ); println!(" server_name: {}", app_config.server.server_name); - // Load and validate providers. let providers_dir = providers_dir_for(&path); let registry = if providers_dir.exists() { ProviderRegistry::load_from_dir(&providers_dir) @@ -58,7 +51,6 @@ pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { .with_context(|| "building empty provider registry")? }; - // Validate provider protocols against builtin adapters. let adapter_registry = ProviderAdapterRegistry::builtin(); let known_protocols = adapter_registry.protocol_names(); registry @@ -87,7 +79,6 @@ pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { } } - // Print provider route table: route_kind -> adapter/protocol println!(); println!("=== Provider Routes ==="); if registry.is_empty() { @@ -121,8 +112,6 @@ pub fn cmd_validate(config_path: Option<&Path>) -> Result<()> { Ok(()) } -/// Print one provider's route table, returning one message per route whose -/// adapter could not be resolved. fn print_provider_routes(provider: &llm_proxy_core::ProviderConfig) -> Vec { println!("{}:", provider.name); diff --git a/apps/llm-proxy/src/config_validation.rs b/apps/llm-proxy/src/config_validation.rs index 96b3834..09fdfaa 100644 --- a/apps/llm-proxy/src/config_validation.rs +++ b/apps/llm-proxy/src/config_validation.rs @@ -58,8 +58,6 @@ mod tests { assert!(result.is_err()); } - // `is_toml_config` matrix, split into one assertion per test so a single - // regression points at the exact input shape that broke (see LOW-20). #[test] fn is_toml_config_accepts_lowercase_toml() { assert!(is_toml_config(Path::new("a.toml"))); diff --git a/apps/llm-proxy/src/defaults.rs b/apps/llm-proxy/src/defaults.rs index 865b73c..64b32aa 100644 --- a/apps/llm-proxy/src/defaults.rs +++ b/apps/llm-proxy/src/defaults.rs @@ -3,16 +3,10 @@ //! Contains the default TOML config content and provider TOML templates //! generated by `llm-proxy init`. -/// Default main config TOML content generated by `llm-proxy init`. +/// Default main config TOML content written by `llm-proxy init`. /// -/// Contains server settings only. Provider routing is configured in -/// individual provider TOML files under `providers/`. Does NOT include: -/// - sampling overrides (temperature, top_p, etc.) -/// - tool overrides -/// - reasoning/thinking overrides -/// - cache settings -/// - stream overrides -/// - fallback/scenario logic +/// Server settings only; provider routing lives in the per-provider files +/// under `providers/`. pub const DEFAULT_CONFIG_TOML: &str = r#"[server] bind = "127.0.0.1:3456" request_timeout = "300s" diff --git a/apps/llm-proxy/src/lib.rs b/apps/llm-proxy/src/lib.rs index 97aadb9..79bdb9f 100644 --- a/apps/llm-proxy/src/lib.rs +++ b/apps/llm-proxy/src/lib.rs @@ -1,16 +1,5 @@ -//! `llm-proxy` application library. -//! -//! Multi-command CLI for the LLM proxy server. Uses TOML config exclusively. -//! -//! Commands: -//! -//! - `serve` Start the proxy server (foreground or daemon) -//! - `stop` Stop a running daemon -//! - `status` Check if the server is running -//! - `init` Create default TOML config and provider files -//! - `validate` Validate TOML config and print provider route table -//! - `models` List provider-scoped models -//! - `autostart` Manage auto-start on login (enable / disable / status) +//! `llm-proxy` application library: a multi-command CLI for the LLM proxy +//! server, using TOML config exclusively. pub mod cli; pub mod commands; diff --git a/apps/llm-proxy/src/paths.rs b/apps/llm-proxy/src/paths.rs index fd8ac0b..3751486 100644 --- a/apps/llm-proxy/src/paths.rs +++ b/apps/llm-proxy/src/paths.rs @@ -1,17 +1,12 @@ -//! Path resolution helpers. -//! -//! Provides functions for resolving config directories, file paths, -//! provider directories, and platform-specific autostart paths. -//! Includes a minimal `home_env` module for home directory resolution. +//! Path resolution helpers: config directories, file paths, provider +//! directories, and platform-specific autostart paths. use std::path::{Path, PathBuf}; /// Config directory: `~/.config/llm-proxy/`. /// -/// Falls back to `./.config/llm-proxy/` (relative to current working directory) -/// when `$HOME` and `$USERPROFILE` are both unset. This is unlikely in practice -/// but avoids a panic. A warning is emitted to stderr in that case since the -/// resolved path depends on the current working directory. +/// Falls back to a path relative to the current working directory (with a +/// stderr warning) when `$HOME` and `$USERPROFILE` are both unset. pub fn config_dir() -> PathBuf { let home = home_env::home_dir(); if home.is_none() { @@ -25,14 +20,8 @@ pub fn config_dir() -> PathBuf { .join("llm-proxy") } -/// Injectable path bundle used by the `stop`/`status` commands (audit -/// MEDIUM-6). -/// -/// Carries the [`llm_proxy_core::PidManager`] that resolves the PID file plus the config file -/// path that `status` reads for the `listen` address. Bundling them lets the -/// full command branching — stale-file cleanup, "no PID file found", and the -/// complete SIGTERM→poll→SIGKILL escalation — be exercised against a -/// tempdir-backed manager in tests instead of the process-global `config_dir`. +/// Path bundle for the `stop`/`status` commands: the PID file manager plus the +/// config file `status` reads for the `listen` address. #[derive(Debug, Clone)] pub struct CommandPaths { pid_manager: llm_proxy_core::PidManager, @@ -62,8 +51,7 @@ impl CommandPaths { } } -/// Build [`CommandPaths`] for the default config directory + default config -/// path (the values used by the real CLI). +/// Build [`CommandPaths`] for the default config directory and config path. #[must_use] pub fn default_command_paths() -> CommandPaths { CommandPaths::new( @@ -82,18 +70,12 @@ pub fn pid_file_path() -> PathBuf { config_dir().join("llm-proxy.pid") } -/// Resolve the config file path. -/// -/// Priority: -/// 1. Explicit CLI `--config` argument -/// 2. `$LLM_PROXY_CONFIG` env var -/// 3. Default path `~/.config/llm-proxy/config.toml` +/// Resolve the config file path: explicit `--config`, then `$LLM_PROXY_CONFIG`, +/// then `~/.config/llm-proxy/config.toml`. pub fn resolve_config(cli_path: Option<&Path>) -> PathBuf { - // 1. Explicit CLI path always wins. if let Some(p) = cli_path { return p.to_path_buf(); } - // 2. Preferred env var. if let Ok(p) = std::env::var("LLM_PROXY_CONFIG") { if !p.is_empty() { return PathBuf::from(p); @@ -102,11 +84,8 @@ pub fn resolve_config(cli_path: Option<&Path>) -> PathBuf { default_config_path() } -/// Returns the providers directory path, which is always `providers/` -/// next to the config file. -/// -/// Falls back to `./providers` if the config path has no parent directory -/// (e.g. bare filename `config.toml` with no directory component). +/// Returns the providers directory path: always `providers/` next to the +/// config file, or `./providers` when the config path has no parent directory. pub fn providers_dir_for(config_path: &Path) -> PathBuf { config_path .parent() @@ -138,20 +117,9 @@ pub fn linux_autostart_dir() -> PathBuf { .unwrap_or_else(|| PathBuf::from(".config/autostart")) } -// --------------------------------------------------------------------------- -// home_env helper (minimal, no external dep) -// --------------------------------------------------------------------------- - -/// Minimal home-directory resolution without depending on the `dirs` crate. -/// -/// Checks `$HOME` (macOS/Linux) and `$USERPROFILE` (Windows) only. -/// Does not fall back to `getpwuid()` or XDG defaults. This is intentional: -/// the `dirs` crate is avoided to minimize dependency tree, and these two -/// env vars cover all standard configurations. pub(crate) mod home_env { use std::path::PathBuf; - /// Returns the user's home directory. pub(crate) fn home_dir() -> Option { std::env::var("HOME") .or_else(|_| std::env::var("USERPROFILE")) @@ -173,9 +141,6 @@ mod tests { #[test] fn resolve_config_uses_default_when_no_args() { - // Without LLM_PROXY_CONFIG set, should return default path. - // We can't unset env vars in this test, so just verify the function - // returns a path ending in config.toml. let result = resolve_config(None); assert!(result.to_string_lossy().ends_with("config.toml")); } diff --git a/apps/llm-proxy/src/permissions.rs b/apps/llm-proxy/src/permissions.rs index d88113c..7743aa9 100644 --- a/apps/llm-proxy/src/permissions.rs +++ b/apps/llm-proxy/src/permissions.rs @@ -9,8 +9,7 @@ use anyhow::{Context, Result}; /// Set owner-only read/write permissions (0600) on a file. /// -/// Used for config and autostart files that reference API keys via -/// environment variables. On non-Unix platforms this is a no-op. +/// No-op on non-Unix platforms. pub fn set_private_permissions(path: &std::path::Path) -> Result<()> { #[cfg(unix)] { @@ -27,8 +26,7 @@ pub fn set_private_permissions(path: &std::path::Path) -> Result<()> { /// Set owner-only permissions (0700) on a directory. /// -/// Used for the config directory which contains provider TOML files with -/// API key references. On non-Unix platforms this is a no-op. +/// No-op on non-Unix platforms. pub fn set_private_dir_permissions(path: &std::path::Path) -> Result<()> { #[cfg(unix)] { @@ -45,9 +43,8 @@ pub fn set_private_dir_permissions(path: &std::path::Path) -> Result<()> { /// Create a new file with restrictive permissions (0600) and write `content`. /// -/// On Unix, uses `OpenOptions` with `mode(0o600)` to set permissions atomically -/// at creation time, avoiding a window where the file exists with default umask -/// permissions. Falls back to create-then-chmod on non-Unix. +/// On Unix the 0600 mode is applied atomically at creation; non-Unix falls +/// back to create-then-chmod. pub fn create_private_file(path: &std::path::Path, content: &[u8]) -> Result<()> { #[cfg(unix)] { @@ -72,15 +69,10 @@ pub fn create_private_file(path: &std::path::Path, content: &[u8]) -> Result<()> Ok(()) } -/// Open a file for appending, creating it if absent, with owner-only (0600) -/// permissions. +/// Open (or create) a file for appending with owner-only (0600) permissions. /// -/// Unlike [`create_private_file`] this uses `create(true)` (not `create_new`) -/// so it is suitable for append-mode sinks such as the daemon log file. On -/// Unix the 0600 mode is applied atomically at creation via `OpenOptionsExt`, -/// avoiding the create-then-chmod TOCTOU window where the file briefly exists -/// under the process umask (typically 0644, world-readable). On non-Unix it -/// falls back to create-then-chmod. +/// Unlike [`create_private_file`] this succeeds when the file exists (append +/// sinks such as the daemon log). Atomic 0600 at creation on Unix. pub fn open_private_append(path: &std::path::Path) -> Result { #[cfg(unix)] { @@ -139,14 +131,12 @@ mod tests { use std::os::unix::fs::PermissionsExt; let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("daemon.log"); - // New file: must be created owner-only (no umask window). { let mut f = open_private_append(&path).unwrap(); writeln!(f, "first line").unwrap(); } let mode = path.metadata().unwrap().permissions().mode(); assert_eq!(mode & 0o777, 0o600, "log file must be born 0600"); - // Re-open (append to existing): must still succeed and keep appending. { let mut f = open_private_append(&path).unwrap(); writeln!(f, "second line").unwrap(); diff --git a/apps/llm-proxy/src/pid.rs b/apps/llm-proxy/src/pid.rs index 66b1c08..aab6160 100644 --- a/apps/llm-proxy/src/pid.rs +++ b/apps/llm-proxy/src/pid.rs @@ -1,8 +1,7 @@ //! PID file management helpers. //! -//! Delegates to `llm_proxy_core::PidManager` for the core PID file operations -//! and adds validation (reject PID 0, overflow checks) and restrictive -//! file permissions. +//! Delegates to `llm_proxy_core::PidManager` and adds PID validation +//! (reject 0, overflow) plus restrictive file permissions. use std::io::Write; @@ -18,21 +17,14 @@ pub const STOP_MAX_POLLS: u32 = 50; /// Interval between polls when waiting for graceful shutdown. pub const STOP_POLL_INTERVAL_MS: u64 = 200; -/// Build a [`llm_proxy_core::PidManager`] for the default config directory. pub(crate) fn pid_manager() -> llm_proxy_core::PidManager { llm_proxy_core::PidManager::new(config_dir()) } /// Read the PID from the PID file via the given manager. /// -/// Returns `None` if the file does not exist. Validates that the parsed PID -/// is non-zero and fits in `i32` to catch corrupt/stale PID files and prevent -/// undefined behavior when passing the PID to `libc::kill()`. -/// -/// This is the injectable variant: callers that already hold a [`llm_proxy_core::PidManager`] -/// (e.g. `cmd_stop`/`cmd_status` operating on a tempdir-backed manager) use it -/// directly instead of rebuilding one from the global config dir (audit -/// MEDIUM-6). +/// Returns `None` if the file does not exist. Rejects PID 0 and PIDs above +/// `i32::MAX` as corrupt. pub fn read_pid_from(mgr: &llm_proxy_core::PidManager) -> Result> { let pid = mgr.read_pid()?; if let Some(p) = pid { @@ -46,22 +38,17 @@ pub fn read_pid_from(mgr: &llm_proxy_core::PidManager) -> Result> { Ok(pid) } -/// Read the PID from the PID file. Returns `None` if the file does not exist. +/// Read the PID from the default PID file. Returns `None` if it does not exist. /// -/// Validates that the parsed PID is non-zero and fits in `i32` to catch -/// corrupt/stale PID files and prevent undefined behavior when passing the -/// PID to `libc::kill()`. +/// Rejects PID 0 and PIDs above `i32::MAX` as corrupt. pub fn read_pid() -> Result> { read_pid_from(&pid_manager()) } /// Atomically create the PID file with exclusive ownership. /// -/// Uses `create_new(true)` so the file is created exclusively -- if another -/// process already created it, the call fails and we return the OS error. -/// This eliminates the TOCTOU race between checking for a stale PID and -/// writing the new one. After creation, restrictive permissions (0600) are -/// applied. +/// Fails if the file already exists (`create_new`); the file is written with +/// 0600 permissions. pub fn write_pid() -> Result<()> { let dir = config_dir(); std::fs::create_dir_all(&dir) @@ -75,16 +62,12 @@ pub fn write_pid() -> Result<()> { .open(&path) .with_context(|| format!("creating PID file {}", path.display()))?; write!(f, "{pid}").with_context(|| format!("writing PID file {}", path.display()))?; - // Set restrictive permissions on the PID file. set_private_permissions(&path)?; info!(pid, "wrote PID file"); Ok(()) } -/// Write a specific PID value to the PID file. -/// -/// Used by the daemon parent to write the child PID before detaching. -/// Applies restrictive permissions (0600) after creation. +/// Write a specific PID value (e.g. a daemon child's PID) to a new PID file. pub fn write_pid_value(pid: u32) -> Result<()> { let dir = config_dir(); std::fs::create_dir_all(&dir) @@ -104,17 +87,13 @@ pub fn write_pid_value(pid: u32) -> Result<()> { /// Remove the PID file. pub fn remove_pid() -> Result<()> { - // `PidManager::remove_pid` returns a typed `Result<(), PidError>`; convert - // into the app layer's `anyhow::Result` via `?` (anyhow implements - // `From` because `PidError: std::error::Error + Send + Sync`). Ok(pid_manager().remove_pid()?) } /// Check if a process with the given PID is running. /// -/// Delegates to [`llm_proxy_core::PidManager::is_process_running`]. Returns `false` for -/// PID 0 (which would otherwise signal the entire process group on Unix) -/// and on non-Unix platforms where process signaling is unavailable. +/// Returns `false` for PID 0 (which would signal the entire process group on +/// Unix) and on non-Unix platforms. pub fn is_process_running(pid: u32) -> bool { if pid == 0 { return false; @@ -132,9 +111,6 @@ mod tests { let pid_path = dir.path().join("llm-proxy.pid"); std::fs::write(&pid_path, "0").unwrap(); - // Exercise the injectable wrapper that performs the zero-PID guard - // (the production path used by status/stop). A PID file of "0" must be - // rejected rather than returned as Some(0). let mgr = llm_proxy_core::PidManager::new(dir.path()); let result = read_pid_from(&mgr); assert!( @@ -151,14 +127,11 @@ mod tests { #[test] fn is_process_running_nonexistent_pid() { - // PID 299999999 is extremely unlikely to exist. assert!(!is_process_running(299_999_999)); } #[test] fn is_process_running_rejects_pid_zero() { - // PID 0 should always return false to avoid signaling the entire - // process group on Unix. assert!(!is_process_running(0)); } } diff --git a/apps/llm-proxy/src/platform.rs b/apps/llm-proxy/src/platform.rs index 2defb52..39692d6 100644 --- a/apps/llm-proxy/src/platform.rs +++ b/apps/llm-proxy/src/platform.rs @@ -3,14 +3,13 @@ //! Provides plist generation (macOS), desktop entry generation (Linux), //! and associated escaping utilities. +#[cfg(target_os = "macos")] use crate::paths::config_dir; /// Generate a macOS launchd plist for auto-start. /// -/// Each argument is placed in its own `` element in the -/// `ProgramArguments` array, avoiding shell interpretation entirely. -/// Argument values are XML-escaped to prevent malformed plists when the -/// executable path contains XML special characters (<, >, &, ", '). +/// Each argument becomes its own XML-escaped `` in `ProgramArguments`, +/// avoiding shell interpretation entirely. #[cfg(target_os = "macos")] pub fn format_plist(args: &[String]) -> String { let arg_elements: String = args @@ -47,10 +46,8 @@ pub fn format_plist(args: &[String]) -> String { /// Generate a Linux XDG desktop entry for auto-start. /// -/// Each argument is individually escaped per the Desktop Entry Specification -/// for the `Exec` key, then joined with spaces. This avoids shell injection -/// by properly handling spaces, quotes, dollar signs, backticks, backslashes, -/// and other shell metacharacters in paths. +/// Each argument is escaped per the Desktop Entry Specification for the `Exec` +/// key, avoiding shell interpretation entirely. #[cfg(target_os = "linux")] pub fn format_desktop_entry(args: &[String]) -> String { let escaped_args: String = args @@ -69,13 +66,8 @@ Hidden=false ) } -/// Escape a single argument for use in a Desktop Entry `Exec` key. -/// -/// Per the XDG Desktop Entry Specification, the following characters must be -/// escaped in Exec values: space, tab, newline, backslash, double-quote, -/// single-quote, dollar, backtick, tilde, hash, percent, ampersand, asterisk, -/// parentheses, vertical bar, semicolon, less-than, greater-than, question mark, -/// square brackets, curly braces. +/// Escape a single argument for use in a Desktop Entry `Exec` key, per the +/// XDG Desktop Entry Specification. #[cfg(target_os = "linux")] pub fn desktop_exec_escape(arg: &str) -> String { let mut out = String::with_capacity(arg.len()); @@ -134,11 +126,7 @@ pub fn xml_escape(s: &str) -> String { mod tests { use super::*; - /// Escape the regex metacharacters in a literal string. - /// - /// Used to build an `insta::Settings::add_filter` regex that matches the - /// verbatim log path (which contains regex-special characters like `.`) - /// so it can be redacted to `[log_path]` in snapshots. + #[cfg(target_os = "macos")] fn regex_escape(literal: &str) -> String { let mut out = String::with_capacity(literal.len()); for ch in literal.chars() { @@ -150,8 +138,6 @@ mod tests { out } - // --- xml_escape tests --- - #[test] fn xml_escape_handles_all_entities() { assert_eq!(xml_escape("<>&\"'"), "<>&"'"); @@ -164,19 +150,15 @@ mod tests { #[test] fn xml_escape_xss_vectors() { - // CDATA injection attempt assert_eq!(xml_escape("]]>"), "<?xml version='1.0'?>" ); - // Script tag injection assert_eq!( xml_escape(""), "<script>alert('xss')</script>" ); - // Null byte assert_eq!(xml_escape("a\x00b"), "a\x00b"); } @@ -190,20 +172,6 @@ mod tests { assert_eq!(xml_escape("<<<"), "<<<"); } - // --- format_plist tests (macOS only) --- - // - // The full-document test (`format_plist_basic`) snapshots the rendered - // plist via insta so structural regressions (a dropped ``, a missing - // ``) fail loudly instead of slipping past substring checks. The - // generated plist embeds a platform- and home-dependent absolute log path - // (`config_dir()/llm-proxy.log`), which would otherwise make the snapshot - // non-reproducible across machines, so the test binds an `insta::Settings` - // scope that filters the concrete log path down to a `[log_path]` - // placeholder before snapshotting. The two companion tests below keep - // targeted `contains()` assertions for the XML-escaping contract - // (positive/negative containment), which reads more clearly there than a - // snapshot would. - #[cfg(target_os = "macos")] #[test] fn format_plist_basic() { @@ -213,9 +181,6 @@ mod tests { .collect(); let plist = format_plist(&args); - // Redact the machine-specific log path so the snapshot is portable. - // `add_filter` takes a regex, so the path's regex metacharacters (e.g. - // the dots in `llm-proxy.log`) are escaped first. let log_path = config_dir().join("llm-proxy.log").display().to_string(); let filter = format!("({})", regex_escape(&log_path)); let mut settings = insta::Settings::clone_current(); @@ -249,13 +214,10 @@ mod tests { .collect(); let plist = format_plist(&args); - // XML special characters in the path should be escaped. assert!(!plist.contains("")); assert!(plist.contains("<special>")); } - // --- desktop_exec_escape tests --- - #[cfg(target_os = "linux")] #[test] fn desktop_exec_escape_spaces() { @@ -269,10 +231,13 @@ mod tests { #[test] fn desktop_exec_escape_shell_metacharacters() { let escaped = desktop_exec_escape("$HOME/`echo test`"); - assert!(!escaped.contains('$')); - assert!(!escaped.contains('`')); assert!(escaped.contains("\\$")); assert!(escaped.contains("\\`")); + let unescaped = escaped.replace("\\$", "").replace("\\`", ""); + assert!( + !unescaped.contains('$') && !unescaped.contains('`'), + "metacharacters must always be escaped, got: {escaped}" + ); } #[cfg(target_os = "linux")] @@ -284,13 +249,6 @@ mod tests { ); } - // --- format_desktop_entry full-document snapshot (Linux only) --- - // - // Unlike the plist, the desktop entry contains no platform-dependent - // absolute paths, so the full rendered document can be snapshotted - // directly. This catches structural regressions (e.g. a dropped - // `X-GNOME-Autostart-enabled` key) that substring checks would miss. - #[cfg(target_os = "linux")] #[test] fn format_desktop_entry_basic() { diff --git a/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap b/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap new file mode 100644 index 0000000..67aeae1 --- /dev/null +++ b/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap @@ -0,0 +1,11 @@ +--- +source: apps/llm-proxy/src/platform.rs +assertion_line: 302 +expression: entry +--- +[Desktop Entry] +Type=Application +Name=LLM Proxy +Exec=/usr/bin/llm-proxy serve +X-GNOME-Autostart-enabled=true +Hidden=false diff --git a/apps/llm-proxy/src/state.rs b/apps/llm-proxy/src/state.rs index da90357..fa12bcf 100644 --- a/apps/llm-proxy/src/state.rs +++ b/apps/llm-proxy/src/state.rs @@ -15,16 +15,6 @@ use tracing::info; use crate::paths::providers_dir_for; -/// Load TOML config files and build [`AppState`] in TOML new-runtime mode. -/// -/// Loads the main `AppConfig` from `path`, then scans a `providers/` directory -/// next to `path` for provider TOML files. Uses [`ProviderRegistry::load_from_dir`] -/// which handles non-regular-file filtering, deterministic sorted ordering, and -/// duplicate-name detection. Protocol validation is performed separately against -/// the builtin adapter set after loading. -/// -/// Takes ownership of `adapter_registry` and `proxy_client` since they are -/// consumed by [`AppState::new`] and not reused after this call. pub(crate) fn load_toml_state( path: &Path, adapter_registry: ProviderAdapterRegistry, @@ -34,16 +24,12 @@ pub(crate) fn load_toml_state( let mut app_config: AppConfig = llm_proxy_core::load_app_config(path) .with_context(|| format!("loading TOML config from {}", path.display()))?; - // Apply CLI port override. - // In-place mutation is safe here: the config is consumed by AppState::new - // below and is not shared with any other caller. if let Some(p) = port_override { let mut bind = app_config.server.bind; bind.set_port(p); app_config.server.bind = bind; } - // Load provider files from providers/ directory next to the main config. let providers_dir = providers_dir_for(path); let registry = if providers_dir.exists() { ProviderRegistry::load_from_dir(&providers_dir) @@ -53,7 +39,6 @@ pub(crate) fn load_toml_state( .with_context(|| "building empty provider registry")? }; - // Validate all provider protocols against the builtin adapter set. let known_protocols = adapter_registry.protocol_names(); registry .validate_protocols(known_protocols) @@ -66,12 +51,55 @@ pub(crate) fn load_toml_state( } } + let client_auth = llm_proxy_server::ClientAuth::from_config(&app_config.auth) + .with_context(|| "building client auth from [auth] config")?; + + let mut egress_clients = std::collections::HashMap::new(); + for provider in registry.iter() { + if let Some(egress) = &provider.egress { + let client = ProxyClient::try_with_egress_proxy(&egress.proxy).with_context(|| { + format!( + "building egress proxy client for provider \"{}\" from [provider.egress]", + provider.name + ) + })?; + info!( + provider = %provider.name, + scheme = egress.proxy.split("://").next().unwrap_or_default(), + "routing upstream traffic through the configured egress proxy" + ); + egress_clients.insert(provider.name.clone(), std::sync::Arc::new(client)); + } + } + + let event_bus: std::sync::Arc = match &app_config.storage { + Some(storage) => { + let bus = llm_proxy_storage::SqliteEventBus::open( + &storage.path, + llm_proxy_storage::RetentionPolicy { + max_age: storage.max_age, + max_rows: storage.max_rows, + prune_every: storage.prune_interval, + }, + ) + .with_context(|| format!("opening usage ledger at {}", storage.path.display()))?; + info!( + path = %storage.path.display(), + "persisting usage events to the SQLite ledger" + ); + std::sync::Arc::new(bus) + } + None => std::sync::Arc::new(llm_proxy_storage::NoopBus), + }; + info!( config = %path.display(), providers = registry.len(), "loaded TOML config" ); + warn_admin_posture(&app_config); + Ok(AppState::new_with_catalog_dir( app_config, registry, @@ -79,7 +107,45 @@ pub(crate) fn load_toml_state( proxy_client, build_info(), Some(cache_dir), - )) + ) + .with_client_auth(std::sync::Arc::new(client_auth)) + .with_egress_clients(egress_clients) + .with_event_bus(event_bus)) +} + +/// Emit the startup warnings naming the admin `/usage` mount posture. +/// +/// Warns when the admin surface is mounted on the main listener or when the +/// dedicated admin listener binds a non-loopback address. +pub fn warn_admin_posture(config: &AppConfig) { + use llm_proxy_core::AdminMountPosture; + match config.admin_posture() { + Some(AdminMountPosture::MainListener) => { + tracing::warn!( + posture = AdminMountPosture::MainListener.name(), + main_bind = %config.server.bind, + "ADMIN SECURITY POSTURE: /usage is also served on the main listener, \ + protected only by the admin token — every network that reaches the \ + main listener reaches the admin API. Set server.admin_on_main_listener \ + = false to return to the loopback-only dedicated admin listener." + ); + } + Some(AdminMountPosture::DedicatedListener) => { + if let Some(bind) = config.server.effective_admin_bind() + && !bind.ip().is_loopback() + { + tracing::warn!( + posture = AdminMountPosture::DedicatedListener.name(), + admin_bind = %bind, + "ADMIN SECURITY POSTURE: the dedicated admin listener binds a \ + non-loopback address; any peer that reaches it can attempt the \ + admin token. Point server.admin_bind back at loopback unless a \ + tunnel requires otherwise." + ); + } + } + None => {} + } } /// Check if a provider with catalog enforcement enabled has no usable models. @@ -124,25 +190,9 @@ pub fn build_info() -> BuildInfo { /// Initialize the tracing subscriber. /// -/// Uses `RUST_LOG` env var with fallback to "info". The output format is -/// selected by `log_format`: [`LogFormat::Plain`] (default) for human-readable -/// logs, [`LogFormat::Json`] for structured JSON (one object per line). The -/// format is read from the `RUST_LOG_FORMAT` env var at init via -/// [`LogFormat::from_env`], before the TOML config is loaded (tracing is needed -/// during config loading itself). To control log verbosity, set `RUST_LOG`. +/// Reads `RUST_LOG` (falling back to `"info"`) and the output format from +/// `RUST_LOG_FORMAT` via [`LogFormat::from_env`]. pub fn init_tracing(log_format: &LogFormat) { - // Delegates to [`build_log_layer`] so the format-selection logic - // (`.json()` vs plain) is exercised by `build_log_layer`'s own tests and - // is not duplicated here. The public signature is unchanged. - // - // The filter is attached *to the boxed layer* (via `Layer::with_filter`) - // rather than to the registry first. `build_log_layer` returns a type-erased - // `Box>`; chaining `.with(filter).with(boxed_layer)` - // would require the erased layer to implement `Layer>`, which the `Box>` blanket impl does not provide - // (its `S` is fixed to `Registry`). Filtering the layer instead yields a - // concrete `Filtered<..>` type that implements `Layer` and - // composes cleanly with `registry().with(...)`. use tracing_subscriber::Layer; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; @@ -151,27 +201,11 @@ pub fn init_tracing(log_format: &LogFormat) { tracing_subscriber::registry().with(layer).init(); } -/// Build the [`EnvFilter`] for tracing. -/// -/// Falls back to `"info"` when `RUST_LOG` is unset. `"info"` is a known-valid -/// filter string, so `parse_lossy` is safe here (it never panics). fn build_env_filter() -> tracing_subscriber::EnvFilter { use tracing_subscriber::EnvFilter; EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::builder().parse_lossy("info")) } -/// Build the tracing-subscriber fmt layer for `log_format`, writing to the sink -/// produced by `make_writer`. -/// -/// Extracted from [`init_tracing`] so the format selection -/// (plain vs `.json()`) is independently testable: a test passes a buffer-backed -/// [`MakeWriter`] here, installs the composed subscriber as the thread-local -/// default, emits a [`tracing::event!`], and asserts on the captured bytes — -/// without needing to call the global, process-wide `init()` (which can only -/// run once and would race with other tests). -/// -/// Returns a boxed, type-erased layer so callers (and tests) do not have to -/// name the concrete `Layer` generic over the writer type. fn build_log_layer( log_format: &LogFormat, make_writer: W, @@ -262,45 +296,12 @@ enforce = true assert!(warning.is_none()); } - // ----------------------------------------------------------------------- - // Log format tests (drive the extracted `build_log_layer` helper) - // ----------------------------------------------------------------------- - // - // These tests install the composed subscriber as the thread-local default - // via `Subscriber::set_default` (a scoped guard), NOT via the global - // `init()`. `init()` can only run once per process and would race with - // other tests; `set_default` scopes the subscriber to this thread only, so - // a buffer `MakeWriter` captures exactly the events emitted here. This is - // the standard `tracing-subscriber` pattern for isolated format tests. - - /// Emit a known `tracing::event!` while `subscriber` is the thread-local - /// default, returning the captured bytes. fn capture_event_for(log_format: &LogFormat) -> Vec { use std::sync::{Arc, Mutex}; let buf = Arc::new(Mutex::new(Vec::::new())); - // `BufferWriter` is a tiny newtype that implements - // `tracing_subscriber::fmt::MakeWriter` by locking the shared buffer - // and returning a `MutexGuard`-backed writer. The layer clones the - // `BufferWriter` (cheap `Arc` clone) internally; we keep the original - // `Arc` to read the captured bytes after the event is emitted. let layer = build_log_layer(log_format, BufferWriter(Arc::clone(&buf))); - // `SubscriberExt::with` must be in scope to compose the layer onto the - // registry. `set_default` (not the one-shot global `init`) scopes the - // subscriber to this thread only, so a buffer `MakeWriter` captures - // exactly the events emitted here and tests do not race with each other. use tracing_subscriber::layer::SubscriberExt; let subscriber = tracing_subscriber::registry().with(layer); - // Scope the subscriber to this thread only. `set_default` (from the - // `tracing` crate) takes the subscriber by value and returns a guard - // that restores the prior subscriber on drop. This avoids the global - // `init()` (one-shot, process-wide, would race with other tests). - // - // The fmt layer (and the subscriber+guard that own it) clones the - // `BufferWriter`/`Arc` internally and may outlive the `event!` call, so - // this helper does NOT depend on `Arc::try_unwrap` to recover the - // buffer (that would race the layer's drop and panic spuriously). The - // captured bytes are read through a shared `lock()` on the live `Arc` - // instead, which is correct regardless of how many clones exist. { let _guard = tracing::subscriber::set_default(subscriber); tracing::event!( @@ -313,15 +314,6 @@ enforce = true buf.lock().expect("buffer lock not poisoned").clone() } - /// `MakeWriter` implementation backed by a shared, lockable byte buffer. - /// - /// `tracing_subscriber::fmt::MakeWriter` is not implemented for - /// `Arc>>` directly (the `Arc` blanket impl requires - /// `W: MakeWriter`, and `Mutex>` is not a writer), so this newtype - /// provides the impl explicitly. Each `make_writer` call returns a - /// `MutexGuard<'static>`-shaped writer that appends bytes to the shared - /// buffer; the `MutexGuard` is `Send` and `Sync` and lives as long as the - /// borrowed `Arc` (which the layer clones and holds for `'static`). #[derive(Clone)] struct BufferWriter(std::sync::Arc>>); @@ -333,8 +325,6 @@ enforce = true } } - /// Write handle for [`BufferWriter`], backed by a `MutexGuard` over the - /// shared buffer. Implements `Write` by appending directly. struct BufferWriterGuard<'a>(std::sync::MutexGuard<'a, Vec>); impl<'a> std::io::Write for BufferWriterGuard<'a> { @@ -348,11 +338,6 @@ enforce = true } } - /// Minimal structural check that `line` is a single JSON object (starts with - /// `{`, ends with `}`). A full `serde_json` parse is intentionally avoided - /// to keep this test off a `serde_json` dev-dependency in the app crate; the - /// brace check plus key-content assertions below are sufficient to - /// distinguish the JSON formatter from the human-readable plain formatter. fn looks_like_json_object(line: &str) -> bool { let trimmed = line.trim(); trimmed.starts_with('{') && trimmed.ends_with('}') @@ -366,7 +351,6 @@ enforce = true !captured.is_empty(), "a tracing::event! must produce output under LogFormat::Json" ); - // Each non-empty line must be a single JSON object (one per event). for line in captured.lines() { if line.trim().is_empty() { continue; @@ -376,10 +360,6 @@ enforce = true "LogFormat::Json line must be a JSON object, got: {line}" ); } - // The JSON formatter emits `level` and the `fields.message` payload as - // structured keys. Assert both are present to prove our `event!` - // reached the JSON formatter and the fields were serialized (not - // dropped by a future edit to `build_log_layer`). assert!( captured.contains("\"level\"") && captured.contains("INFO"), "JSON log must carry the structured `level` field, got: {captured}" @@ -398,10 +378,6 @@ enforce = true !captured.is_empty(), "a tracing::event! must produce output under LogFormat::Plain" ); - // Take the first non-empty line. Under Plain it is human-readable text - // like `2026-... INFO log_format_test: ...` and must NOT be a JSON - // object — if it were, the format selection in `build_log_layer` would - // be broken (Plain behaving like Json). let first = captured .lines() .find(|line| !line.trim().is_empty()) @@ -410,10 +386,303 @@ enforce = true !looks_like_json_object(first), "Plain log line must NOT be a JSON object, got: {first}" ); - // Sanity: the probe message still reached the formatter. assert!( captured.contains("log-format-probe"), "Plain log must still carry the emitted message, got: {captured}" ); } + + fn write_main_config( + dir: &std::path::Path, + storage_section: Option<&str>, + ) -> std::path::PathBuf { + let path = dir.join("config.toml"); + let storage = storage_section.unwrap_or_default(); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +{storage}"# + ), + ) + .unwrap(); + path + } + + #[tokio::test] + async fn load_toml_state_wires_sqlite_ledger_when_storage_configured() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("usage.db"); + let storage = format!( + " +[storage] +path = \"{}\" +", + db_path.display() + ); + let path = write_main_config(dir.path(), Some(&storage)); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state with [storage]"); + + let debug = format!("{state:?}"); + assert!( + debug.contains("SqliteEventBus"), + "[storage] must wire the SQLite ledger, got: {debug}" + ); + + let event = + llm_proxy_storage::ProxyEvent::RequestReceived(llm_proxy_storage::RequestReceived { + request_id: "wiring-1".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "laptop".to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "digest".to_owned(), + }); + let bus = std::sync::Arc::clone(state.event_bus()); + bus.emit(&event); + let closer = std::sync::Arc::clone(&bus); + tokio::task::spawn_blocking(move || llm_proxy_storage::EventBus::close(&*closer)) + .await + .expect("close"); + + let db = std::fs::read(&db_path).expect("database file exists after close"); + assert!( + !db.is_empty(), + "database file must hold the committed schema and rows" + ); + } + + #[tokio::test] + async fn load_toml_state_keeps_noop_bus_without_storage() { + let dir = tempfile::tempdir().unwrap(); + let path = write_main_config(dir.path(), None); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state without [storage]"); + + let debug = format!("{state:?}"); + assert!( + debug.contains("NoopBus"), + "absent [storage] must keep the no-op bus (today's behavior), got: {debug}" + ); + llm_proxy_storage::EventBus::close(state.event_bus().as_ref()); + } + + fn write_provider_file(dir: &std::path::Path, egress: Option<&str>) { + let egress_section = egress.unwrap_or_default(); + std::fs::create_dir_all(dir.join("providers")).unwrap(); + std::fs::write( + dir.join("providers/egress-test.toml"), + format!( + r#" +[provider] +name = "egress-test" +api_key = "test-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions" + +[provider.routes] +chat_completions = "chat" +{egress_section}"# + ), + ) + .unwrap(); + } + + #[tokio::test] + async fn load_toml_state_wires_egress_client_only_for_configured_providers() { + let dir = tempfile::tempdir().unwrap(); + write_provider_file( + dir.path(), + Some( + "[provider.egress]\nproxy = \"socks5h://egress-user:egress-pass@127.0.0.1:1080\"\n", + ), + ); + std::fs::write( + dir.path().join("providers/direct.toml"), + r#" +[provider] +name = "direct" +api_key = "test-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions" + +[provider.routes] +chat_completions = "chat" +"#, + ) + .unwrap(); + let path = write_main_config(dir.path(), None); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state builds with an egress provider"); + + let debug = format!("{state:?}"); + let egress_map = debug + .split("egress_clients: ") + .nth(1) + .and_then(|rest| rest.split(')').next()) + .unwrap_or_default(); + assert!( + egress_map.contains("egress-test"), + "the egress provider must have a wired client: {debug}" + ); + assert!( + !egress_map.contains("direct"), + "providers without egress must not appear in the egress map: {debug}" + ); + assert!( + !debug.contains("egress-user") && !debug.contains("egress-pass"), + "proxy credentials must never print in the state Debug: {debug}" + ); + } + + #[tokio::test] + async fn load_toml_state_rejects_egress_url_the_transport_cannot_use() { + let dir = tempfile::tempdir().unwrap(); + write_provider_file( + dir.path(), + Some("[provider.egress]\nproxy = \"http://[::1:8080\"\n"), + ); + let path = write_main_config(dir.path(), None); + + let err = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect_err("unusable egress proxy must be a startup error"); + + let message = format!("{err:#}"); + assert!( + message.contains("egress"), + "the error must point at the egress wiring: {message}" + ); + assert!( + message.contains("egress-test"), + "the error must name the provider: {message}" + ); + } + + fn admin_app_config( + on_main: bool, + admin_bind: Option, + ) -> llm_proxy_core::AppConfig { + use llm_proxy_core::{AdminAuthConfig, ServerConfig, StorageConfig}; + llm_proxy_core::AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: std::time::Duration::from_secs(300), + shutdown_timeout: std::time::Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "posture-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: std::time::Duration::ZERO, + log_format: Default::default(), + admin_bind, + admin_on_main_listener: on_main, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients: std::collections::HashMap::new(), + admin: Some(AdminAuthConfig { + token: secrecy::SecretString::from("admin-secret"), + }), + }, + storage: Some(StorageConfig { + path: "usage.db".into(), + max_age: None, + max_rows: None, + prune_interval: std::time::Duration::from_secs(3600), + }), + } + } + + fn captured_posture_warnings(config: &llm_proxy_core::AppConfig) -> Vec { + let buf = std::sync::Arc::new(std::sync::Mutex::new(Vec::::new())); + let layer = build_log_layer(&LogFormat::Json, BufferWriter(std::sync::Arc::clone(&buf))); + use tracing_subscriber::layer::SubscriberExt; + let subscriber = tracing_subscriber::registry().with(layer); + { + let _guard = tracing::subscriber::set_default(subscriber); + warn_admin_posture(config); + } + buf.lock().expect("buffer lock not poisoned").clone() + } + + #[test] + fn main_listener_optin_is_warned_about_loudly_by_name() { + let config = admin_app_config(true, None); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + captured.contains("admin-on-main-listener"), + "the posture must be named for grepping: {captured}" + ); + assert!( + captured.contains("ADMIN SECURITY POSTURE"), + "the warning must be loud and identifiable: {captured}" + ); + } + + #[test] + fn dedicated_loopback_posture_stays_quiet() { + let config = admin_app_config(false, None); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + !captured.contains("ADMIN SECURITY POSTURE"), + "the default loopback posture needs no warning: {captured}" + ); + } + + #[test] + fn dedicated_non_loopback_bind_is_warned_about() { + let config = admin_app_config(false, Some("10.0.0.5:9900".parse().unwrap())); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + captured.contains("non-loopback"), + "a non-loopback admin listener must be called out: {captured}" + ); + } } diff --git a/apps/llm-proxy/tests/cli_init.rs b/apps/llm-proxy/tests/cli_init.rs index db1ee19..13d3df5 100644 --- a/apps/llm-proxy/tests/cli_init.rs +++ b/apps/llm-proxy/tests/cli_init.rs @@ -13,21 +13,14 @@ mod tests { DEFAULT_CONFIG_TOML, DEFAULT_PROVIDER_OPENCODE_GO, DEFAULT_PROVIDER_OPENCODE_ZEN, }; - /// Serializes tests that mutate the process-global `$HOME` env var so they do - /// not race with each other when run in parallel. Every such test must hold - /// this lock for its whole body and restore `$HOME` before releasing it. static ENV_LOCK: Mutex<()> = Mutex::new(()); - /// RAII guard that sets `$HOME` to `temp_home` on construction and restores the - /// prior value on drop. Must be constructed while holding [`ENV_LOCK`]. struct HomeGuard { prior: Option, } impl HomeGuard { fn new(temp_home: &std::path::Path) -> Self { - // Callers hold ENV_LOCK, so no other test mutates or reads HOME - // concurrently. The prior value is always restored on drop. let prior = std::env::var_os("HOME").map(|v| v.to_string_lossy().into_owned()); #[expect( unsafe_code, @@ -93,8 +86,6 @@ mod tests { fn init_writes_files_to_tempdir() { let dir = tempfile::tempdir().expect("tempdir"); - // Simulate the file-creation logic of cmd_init without calling the real - // function (which hard-codes paths via config_dir()). let config_path = dir.path().join("config.toml"); let providers_dir = dir.path().join("providers"); fs::create_dir_all(&providers_dir).expect("providers dir"); @@ -115,7 +106,6 @@ mod tests { assert!(providers_dir.join("opencode-go.toml").exists()); assert!(providers_dir.join("opencode-zen.toml").exists()); - // Verify parsed contents match. let main: toml::Value = toml::from_str(&fs::read_to_string(&config_path).unwrap()).unwrap(); assert!(main.get("server").is_some()); @@ -130,19 +120,11 @@ mod tests { assert_eq!(zen["provider"]["name"].as_str(), Some("opencode-zen")); } - /// `cmd_init` must refuse to overwrite an existing config and leave the file - /// contents untouched. - /// - /// `cmd_init` resolves its config directory via `$HOME` (see `paths::config_dir`). - /// We point `$HOME` at a tempdir, pre-create the config file, invoke the real - /// `cmd_init`, and assert it returns `Err` with the "already exists" message - /// without overwriting the file. `$HOME` mutation is serialized via `ENV_LOCK`. #[test] fn init_rejects_existing_config() { let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let dir = tempfile::tempdir().expect("tempdir"); - // cmd_init writes to $HOME/.config/llm-proxy/config.toml. let config_subdir = dir.path().join(".config").join("llm-proxy"); fs::create_dir_all(&config_subdir).expect("create config dir"); let config_path = config_subdir.join("config.toml"); @@ -150,18 +132,15 @@ mod tests { let _home = HomeGuard::new(dir.path()); - // The real command must bail because the config already exists. let result = cmd_init(); let err = result.expect_err("cmd_init should reject an existing config"); - // The bail message names the config path and the "already exists" condition. let msg = format!("{err:#}"); assert!( msg.contains("already exists"), "error should mention the config already existing, got: {msg}" ); - // The existing file must NOT have been overwritten. let on_disk = fs::read_to_string(&config_path).expect("read config"); assert_eq!( on_disk, "existing", diff --git a/apps/llm-proxy/tests/cli_parsing.rs b/apps/llm-proxy/tests/cli_parsing.rs index 7ab4563..cd3e764 100644 --- a/apps/llm-proxy/tests/cli_parsing.rs +++ b/apps/llm-proxy/tests/cli_parsing.rs @@ -103,7 +103,6 @@ mod tests { #[test] fn parse_models_write_catalog_requires_live() { - // --write-catalog requires --live; clap should reject this. let result = Cli::try_parse_from(["llm-proxy", "models", "--write-catalog"]); result.unwrap_err(); } diff --git a/apps/llm-proxy/tests/cli_serve.rs b/apps/llm-proxy/tests/cli_serve.rs index c2586d2..5b3c2fc 100644 --- a/apps/llm-proxy/tests/cli_serve.rs +++ b/apps/llm-proxy/tests/cli_serve.rs @@ -1,10 +1,6 @@ -//! Integration tests for the `serve` command. -//! -//! Starting the server binds a port and blocks, so these tests only cover the -//! pre-flight config validation that `serve` runs before it would bind. The -//! full `validate_toml_extension` accept/reject matrix is exercised as per-case -//! unit tests in `src/config_validation.rs` (see LOW-20); this file keeps a -//! single smoke test confirming the error shape `serve` surfaces. +//! Integration tests for the `serve` command: pre-flight config validation +//! only (starting the server binds and blocks). The full extension matrix +//! lives in `src/config_validation.rs`. #[cfg(test)] mod tests { diff --git a/apps/llm-proxy/tests/cli_status.rs b/apps/llm-proxy/tests/cli_status.rs index bfec052..4a296f7 100644 --- a/apps/llm-proxy/tests/cli_status.rs +++ b/apps/llm-proxy/tests/cli_status.rs @@ -1,8 +1,5 @@ -//! Integration tests for the `status` command. -//! -//! Drives `cmd_status` directly against a tempdir-backed [`CommandPaths`], -//! covering the "no PID file", stale-PID, running-process, and config-display -//! branches (audit MEDIUM-6). +//! Integration tests for the `status` command, driven against a tempdir-backed +//! [`CommandPaths`]. #[cfg(test)] mod tests { @@ -14,8 +11,6 @@ mod tests { use llm_proxy_app::paths::CommandPaths; use llm_proxy_core::PidManager; - /// Build a [`CommandPaths`] whose PID file and config path both live under - /// `dir`. fn tempdir_paths(dir: &Path) -> CommandPaths { CommandPaths::new(PidManager::new(dir), dir.join("config.toml")) } @@ -50,16 +45,12 @@ mod tests { cmd_status(&paths).expect("running process -> Ok"); - // Reap the child so it does not outlive the test. child.kill().expect("kill sleep child"); child.wait().expect("reap sleep child"); } #[test] fn status_with_present_config_is_ok() { - // `cmd_status` reads the config to display the `listen` address. With a - // valid config present it hits the parse-success arm; the command - // returns Ok either way (parse failures are reported, not propagated). let dir = tempfile::tempdir().unwrap(); let paths = tempdir_paths(dir.path()); diff --git a/apps/llm-proxy/tests/cli_stop.rs b/apps/llm-proxy/tests/cli_stop.rs index 20b9953..74412c1 100644 --- a/apps/llm-proxy/tests/cli_stop.rs +++ b/apps/llm-proxy/tests/cli_stop.rs @@ -1,10 +1,6 @@ -//! Integration tests for the `stop` command. -//! -//! Drives `cmd_stop` / `cmd_stop_with_timing` directly against a tempdir-backed -//! [`CommandPaths`], covering every branch the audit flagged as uncovered -//! (audit MEDIUM-6): the "no PID file found" path, stale-PID cleanup, the -//! graceful SIGTERM shutdown of a live process, and the full SIGTERM→poll→SIGKILL -//! escalation when SIGTERM is ignored. +//! Integration tests for the `stop` command, driven against a tempdir-backed +//! [`CommandPaths`]: no PID file, stale-PID cleanup, graceful SIGTERM, and the +//! SIGKILL escalation. #[cfg(test)] mod tests { @@ -16,15 +12,12 @@ mod tests { use llm_proxy_app::paths::CommandPaths; use llm_proxy_core::PidManager; - /// Build a [`CommandPaths`] whose PID file and config path both live under - /// `dir` (the config path is a sibling of the PID file, never written here). fn tempdir_paths(dir: &Path) -> CommandPaths { CommandPaths::new(PidManager::new(dir), dir.join("config.toml")) } #[cfg(unix)] mod unix { - //! On Unix, every `cmd_stop` branch is exercised end-to-end. use super::*; use std::process::Command; @@ -45,7 +38,6 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let paths = tempdir_paths(dir.path()); - // 299_999_999 is effectively never a live PID -> "stale" branch. std::fs::write(paths.pid_manager().pid_file(), "299999999").unwrap(); cmd_stop(&paths).expect("stale PID -> Ok + cleanup"); @@ -60,7 +52,6 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let paths = tempdir_paths(dir.path()); - // `sleep` exits on SIGTERM by default -> exercises the graceful branch. let mut child = Command::new("sleep").arg("30").spawn().unwrap(); std::fs::write(paths.pid_manager().pid_file(), child.id().to_string()).unwrap(); @@ -83,10 +74,6 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let paths = tempdir_paths(dir.path()); - // Child ignores SIGTERM and keeps running -> survives the poll loop and - // is force-killed with SIGKILL. `sleep 30 & wait` keeps the recorded - // shell PID alive (the shell, not `sleep`, owns the PID file) and - // holding the SIGTERM trap so the shell is not replaced by exec. let mut child = Command::new("sh") .arg("-c") .arg("trap '' TERM; sleep 30 & wait") @@ -110,7 +97,6 @@ mod tests { #[cfg(not(unix))] mod non_unix { - //! `stop` is unsupported off-Unix; the command must bail cleanly. use super::*; #[test] diff --git a/apps/llm-proxy/tests/cli_validate.rs b/apps/llm-proxy/tests/cli_validate.rs index f2cbf48..1fd72bb 100644 --- a/apps/llm-proxy/tests/cli_validate.rs +++ b/apps/llm-proxy/tests/cli_validate.rs @@ -1,9 +1,5 @@ -//! Integration tests for the `validate` command. -//! -//! Covers end-to-end `cmd_validate` behaviour: valid TOML, invalid TOML, and -//! the wrong-extension error surfacing through the command. The pure extension -//! accept/reject matrix is exercised as per-case unit tests in -//! `src/config_validation.rs` (see LOW-20). +//! Integration tests for the `validate` command: valid TOML, invalid TOML, +//! wrong extension, and the admin/auth validation rules. #[cfg(test)] mod tests { @@ -49,4 +45,184 @@ mod tests { "error should mention unsupported extension: {msg}" ); } + + #[test] + fn validate_rejects_non_loopback_bind_without_client_keys() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "0.0.0.0:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!( + "{:#}", + result.expect_err("non-loopback anonymous bind must be rejected") + ); + assert!( + msg.contains("not loopback"), + "error should explain the loopback restriction: {msg}" + ); + } + + #[test] + fn validate_accepts_client_keys_on_non_loopback_bind() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "0.0.0.0:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=8,t=1,p=1$MDEyMzQ1Njc4OWFiY2RlZg$eEnjdfBL8hfvU2QuGOpow8Yq2Jyrr8GXKYgtAAsscSs" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + assert!( + result.is_ok(), + "client-keys on a non-loopback bind should validate: {result:?}" + ); + } + + #[test] + fn validate_accepts_admin_api_with_storage() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth.admin] +token = "admin-secret" + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + assert!( + result.is_ok(), + "admin + storage should validate: {result:?}" + ); + } + + #[test] + fn validate_rejects_admin_token_without_storage() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth.admin] +token = "admin-secret" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!( + "{:#}", + result.expect_err("admin without storage must be rejected") + ); + assert!( + msg.contains("[auth.admin] requires"), + "error should point at the missing ledger: {msg}" + ); + } + + #[test] + fn validate_rejects_admin_on_main_listener_without_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +admin_on_main_listener = true + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!("{:#}", result.expect_err("dead opt-in must be rejected")); + assert!( + msg.contains("admin_on_main_listener"), + "error should name the dead flag: {msg}" + ); + } + + #[test] + fn validate_rejects_admin_bind_equal_to_main_bind() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +admin_bind = "127.0.0.1:3456" + +[auth.admin] +token = "admin-secret" + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!("{:#}", result.expect_err("same-socket bind rejected")); + assert!( + msg.contains("must differ from server.bind"), + "error should explain the socket conflict: {msg}" + ); + } } diff --git a/config.toml.example b/config.toml.example index 858ffcc..e79526f 100644 --- a/config.toml.example +++ b/config.toml.example @@ -20,3 +20,56 @@ trust_forwarded_headers = false # Window for rejecting identical path/body requests. "0ms" disables dedup. dedup_window = "0ms" server_name = "llm-proxy" +# Dedicated listener for the read-only admin /usage API. Only read when +# [auth.admin] is configured. Defaults to "127.0.0.1" on the port after +# `bind` (loopback-only); set it explicitly to point at a tunnel interface +# or a different port. Must differ from `bind`. +# admin_bind = "127.0.0.1:3457" +# Also serve /usage on the main listener behind the same admin token. +# Off by default; enabling it names a distinct security posture +# ("admin-on-main-listener") that startup warns about loudly: the admin +# surface then reaches every network the main listener does. With no +# explicit admin_bind set, this posture serves /usage on the main listener +# only (no dedicated socket) — the option for setups where a second port +# is not available. +# admin_on_main_listener = false + +# Client authentication for proxy routes. +# +# mode = "none" (default): no client authentication. Only permitted while +# [server].bind is a loopback address; the server refuses to start on a +# non-loopback bind without client keys. +[auth] +mode = "none" +# mode = "client-keys": every proxy request must present +# "Authorization: Bearer ". Named keys map key ids to argon2id +# password hashes in PHC format -- never plaintext secrets. Generate one with +# e.g. `printf 'my-secret' | argon2 'some-salt' -id -e`. +# The key id (not the secret) appears in usage events and rate-limit keys. +# +# [auth] +# mode = "client-keys" +# +# [auth.clients] +# laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$..." + +# Usage ledger persistence. Absent (the default) disables persistence: +# usage events go to an in-memory no-op sink and nothing is written to disk. +# +# [storage] +# SQLite database file. Parent directories are created at startup. +# path = "${HOME}/.local/state/llm-proxy/usage.db" +# Delete ledger rows older than this age; absent keeps rows forever. +# max_age = "30d" +# Keep at most this many rows; the oldest rows are deleted when the cap is +# exceeded. Absent means unlimited. +# max_rows = 1000000 +# How often retention pruning runs. Default "1h"; must be at least "1s". +# prune_interval = "1h" + +# Admin API for the usage ledger. Absent (the default) disables the /usage +# surface entirely. Requires [storage]; every /usage route on every mount +# demands "Authorization: Bearer " (compared in constant time). +# +# [auth.admin] +# token = "${ADMIN_TOKEN}" diff --git a/crates/llm-proxy-api/Cargo.toml b/crates/llm-proxy-api/Cargo.toml index 906abd0..fbbec00 100644 --- a/crates/llm-proxy-api/Cargo.toml +++ b/crates/llm-proxy-api/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition.workspace = true rust-version.workspace = true license.workspace = true -description = "Reserved public HTTP API surface for llm-proxy; a placeholder until the handler layer lands." +description = "Read-only authenticated /usage query surface for the llm-proxy usage ledger." repository.workspace = true readme.workspace = true keywords.workspace = true @@ -12,3 +12,24 @@ categories.workspace = true [lints] workspace = true + +[dependencies] +axum = { workspace = true } +serde = { workspace = true } +thiserror = { workspace = true } +time = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } +rusqlite = { workspace = true } +rust_decimal = { workspace = true } +sha2 = "0.10" +humantime = "2.3" +tower-http = { version = "0.6", features = ["trace", "catch-panic", "set-header"] } + +[dev-dependencies] +serde_json = { workspace = true } +tempfile = "3" +tower = { workspace = true, features = ["util"] } +tracing-subscriber = { workspace = true } +llm-proxy-storage = { path = "../llm-proxy-storage" } +llm-proxy-protocol = { path = "../llm-proxy-protocol" } diff --git a/crates/llm-proxy-api/src/auth.rs b/crates/llm-proxy-api/src/auth.rs new file mode 100644 index 0000000..4de1633 --- /dev/null +++ b/crates/llm-proxy-api/src/auth.rs @@ -0,0 +1,146 @@ +use axum::extract::State; +use axum::http::{StatusCode, header}; +use axum::response::{IntoResponse, Response}; +use axum::{Json, middleware::Next}; +use serde::Serialize; +use sha2::{Digest, Sha256}; + +use crate::routes::UsageState; + +#[derive(Clone)] +pub(crate) struct AdminAuth { + expected_digest: [u8; 32], +} + +impl std::fmt::Debug for AdminAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AdminAuth") + .field("token", &"[REDACTED]") + .finish() + } +} + +impl AdminAuth { + pub(crate) fn new(token: &str) -> Self { + Self { + expected_digest: Sha256::digest(token.as_bytes()).into(), + } + } + + pub(crate) fn accepts(&self, presented: Option<&str>) -> bool { + let Some(token) = presented else { + return false; + }; + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + constant_time_eq(&digest, &self.expected_digest) + } +} + +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + +fn bearer_token(headers: &axum::http::HeaderMap) -> Option<&str> { + let value = headers.get(header::AUTHORIZATION)?; + let header_value = value.to_str().ok()?; + let mut parts = header_value.splitn(2, ' '); + match (parts.next(), parts.next()) { + (Some(scheme), Some(token)) if scheme.eq_ignore_ascii_case("bearer") => { + let token = token.trim(); + (!token.is_empty()).then_some(token) + } + _ => None, + } +} + +#[derive(Serialize)] +struct UnauthorizedBody { + error: UnauthorizedError, +} + +#[derive(Serialize)] +struct UnauthorizedError { + r#type: &'static str, + message: &'static str, +} + +pub(crate) async fn require_admin( + State(state): State, + req: axum::extract::Request, + next: Next, +) -> Response { + if state.admin.accepts(bearer_token(req.headers())) { + return next.run(req).await; + } + ( + StatusCode::UNAUTHORIZED, + Json(UnauthorizedBody { + error: UnauthorizedError { + r#type: "unauthorized", + message: "missing or invalid admin token", + }, + }), + ) + .into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn accepts_only_the_configured_token() { + let auth = AdminAuth::new("admin-secret"); + assert!(auth.accepts(Some("admin-secret"))); + assert!(!auth.accepts(Some("admin-secret2"))); + assert!(!auth.accepts(Some(""))); + assert!(!auth.accepts(None)); + } + + #[test] + fn debug_never_prints_the_digest() { + let auth = AdminAuth::new("admin-secret"); + let debug = format!("{auth:?}"); + assert!( + !debug.contains("admin-secret"), + "Debug must not expose the token, got: {debug}" + ); + } + + #[test] + fn bearer_token_extraction() { + let mut headers = axum::http::HeaderMap::new(); + assert_eq!(bearer_token(&headers), None); + headers.insert( + header::AUTHORIZATION, + "Bearer tok".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), Some("tok")); + headers.insert( + header::AUTHORIZATION, + "bearer tok".parse().expect("valid header value"), + ); + assert_eq!( + bearer_token(&headers), + Some("tok"), + "scheme matching is case-insensitive" + ); + headers.insert( + header::AUTHORIZATION, + "Basic dXNlcjpwYXNz".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), None); + headers.insert( + header::AUTHORIZATION, + "Bearer ".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), None, "empty token is no token"); + } +} diff --git a/crates/llm-proxy-api/src/lib.rs b/crates/llm-proxy-api/src/lib.rs index 67de5c4..7664e92 100644 --- a/crates/llm-proxy-api/src/lib.rs +++ b/crates/llm-proxy-api/src/lib.rs @@ -1,30 +1,11 @@ -//! llm-proxy-api: public API surface for the LLM proxy. -//! -//! This crate is intentionally empty in v1. It reserves the namespace for a -//! future public HTTP API layer. See `docs/protocol.md` for the implemented -//! protocol boundaries. +//! Read-only, admin-authenticated `/usage` query surface over the proxy's +//! usage ledger: every route is a `GET` behind `Authorization: Bearer +//! ` and no response carries request bodies. -/// Placeholder module to reserve the public API namespace. -/// -/// Remove this module and replace with real types once implementation -/// begins. Tracking: . -pub mod placeholder { - /// Sentinel type indicating this crate has not yet been implemented. - /// - /// This type exists so that downstream crates can reference - /// `llm_proxy_api::placeholder::NotImplemented` in feature-gated - /// code without pulling in a fully implemented crate. - #[derive(Debug, Clone, Copy, Default)] - pub struct NotImplemented; -} +#![deny(missing_docs)] -// Ensure the crate compiles as a valid library target. -#[cfg(test)] -mod tests { - use super::placeholder::NotImplemented; +mod auth; +mod reader; +mod routes; - #[test] - fn placeholder_type_exists() { - let _ = NotImplemented; - } -} +pub use routes::UsageApi; diff --git a/crates/llm-proxy-api/src/reader.rs b/crates/llm-proxy-api/src/reader.rs new file mode 100644 index 0000000..b370dee --- /dev/null +++ b/crates/llm-proxy-api/src/reader.rs @@ -0,0 +1,494 @@ +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use rusqlite::Connection; +use rusqlite::OpenFlags; +use rusqlite::types::Value as SqlValue; +use serde::Serialize; +use time::OffsetDateTime; + +const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); + +pub(crate) const MAX_PAGE_LIMIT: i64 = 1000; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum UsageQueryError { + #[error("usage ledger query failed: {0}")] + Sqlite(#[from] rusqlite::Error), + #[error("usage ledger cost value is not a valid decimal: {0}")] + Cost(String), + #[error("usage ledger query task failed: {0}")] + Task(String), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum GroupBy { + Provider, + Model, + Caller, +} + +impl GroupBy { + const fn column(self) -> &'static str { + match self { + Self::Provider => "provider", + Self::Model => "model_requested", + Self::Caller => "caller", + } + } + + pub(crate) fn parse(value: &str) -> Option { + match value { + "provider" => Some(Self::Provider), + "model" => Some(Self::Model), + "caller" => Some(Self::Caller), + _ => None, + } + } + + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Provider => "provider", + Self::Model => "model", + Self::Caller => "caller", + } + } +} + +#[derive(Debug, Clone, Default, Serialize)] +pub(crate) struct TokenCounts { + pub input: i64, + pub output: i64, + pub reasoning: i64, + pub cache_read: i64, + pub cache_creation: i64, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct CostTotals { + pub input: String, + pub output: String, + pub cache_read: String, + pub cache_creation: String, + pub reasoning: String, + pub total: String, +} + +impl CostTotals { + fn zeroed() -> Self { + Self { + input: "0".to_owned(), + output: "0".to_owned(), + cache_read: "0".to_owned(), + cache_creation: "0".to_owned(), + reasoning: "0".to_owned(), + total: "0".to_owned(), + } + } +} + +#[derive(Debug, Serialize)] +pub(crate) struct SummaryGroup { + pub key: Option, + pub requests: i64, + pub errors: i64, + pub tokens: TokenCounts, + pub cost: CostTotals, +} + +#[derive(Debug, Serialize)] +pub(crate) struct RequestRow { + pub request_id: String, + pub received_at_ms: Option, + pub completed_at_ms: Option, + pub caller: Option, + pub provider: Option, + pub route_kind: Option, + pub client_protocol: Option, + pub streaming: Option, + pub model_requested: Option, + pub model_upstream: Option, + pub tokens: Option, + pub usage_provenance: Option, + pub cost: Option, + pub stop_reason: Option, + pub error_kind: Option, + pub http_status: Option, + pub latency_ms: Option, + pub body_hash: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct RequestsPage { + pub rows: Vec, + pub next_cursor: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct CallerSummary { + pub caller: Option, + pub requests: i64, + pub errors: i64, + pub last_seen_ms: Option, +} + +#[derive(Debug, Default)] +pub(crate) struct RequestFilter { + pub since_ms: Option, + pub until_ms: Option, + pub provider: Option, + pub limit: i64, + pub cursor: Option<(i64, String)>, +} + +#[derive(Debug, Clone)] +pub(crate) struct UsageReader { + path: PathBuf, +} + +fn offset_to_unix_millis(ts: OffsetDateTime) -> i64 { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) +} + +pub(crate) fn now_unix_millis() -> i64 { + offset_to_unix_millis(OffsetDateTime::now_utc()) +} + +impl UsageReader { + pub(crate) fn new(path: impl Into) -> Self { + Self { path: path.into() } + } + + pub(crate) async fn summary( + &self, + since_ms: i64, + group: GroupBy, + ) -> Result, UsageQueryError> { + let path = self.path.clone(); + run_blocking(move || summary_sync(&path, since_ms, group)).await + } + + pub(crate) async fn requests( + &self, + filter: RequestFilter, + ) -> Result { + let path = self.path.clone(); + run_blocking(move || requests_sync(&path, &filter)).await + } + + pub(crate) async fn callers(&self) -> Result, UsageQueryError> { + let path = self.path.clone(); + run_blocking(move || callers_sync(&path)).await + } + + fn open(path: &Path) -> Result { + let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?; + conn.busy_timeout(SQLITE_BUSY_TIMEOUT)?; + Ok(conn) + } +} + +async fn run_blocking(f: F) -> Result +where + F: FnOnce() -> Result + Send + 'static, + T: Send + 'static, +{ + tokio::task::spawn_blocking(f) + .await + .map_err(|error| UsageQueryError::Task(error.to_string()))? +} + +fn parse_cost(value: Option) -> Result { + match value { + None => Ok(rust_decimal::Decimal::ZERO), + Some(text) => rust_decimal::Decimal::from_str_exact(&text) + .map_err(|_parse_error| UsageQueryError::Cost(text)), + } +} + +fn summary_sync( + path: &Path, + since_ms: i64, + group: GroupBy, +) -> Result, UsageQueryError> { + let conn = UsageReader::open(path)?; + let key = group.column(); + let sql = format!( + "SELECT {key} AS k, COUNT(*), \ + SUM(CASE WHEN error_kind IS NOT NULL THEN 1 ELSE 0 END), \ + COALESCE(SUM(COALESCE(input_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(output_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(reasoning_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(cache_read_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(cache_creation_tokens, 0)), 0) \ + FROM requests \ + WHERE COALESCE(received_at, completed_at) >= ?1 \ + GROUP BY {key} ORDER BY 2 DESC" + ); + let mut groups: Vec = Vec::new(); + { + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([since_ms], |row| { + Ok(SummaryGroup { + key: row.get(0)?, + requests: row.get(1)?, + errors: row.get::<_, Option>(2)?.unwrap_or(0), + tokens: TokenCounts { + input: row.get(3)?, + output: row.get(4)?, + reasoning: row.get(5)?, + cache_read: row.get(6)?, + cache_creation: row.get(7)?, + }, + cost: CostTotals::zeroed(), + }) + })?; + for row in rows { + groups.push(row?); + } + } + let costs = sum_costs_by_group(&conn, key, since_ms)?; + for group_row in &mut groups { + if let Some(total) = costs.get(&group_row.key) { + group_row.cost = cost_totals_from(*total); + } + } + Ok(groups) +} + +fn sum_costs_by_group( + conn: &Connection, + key: &str, + since_ms: i64, +) -> Result, [rust_decimal::Decimal; 6]>, UsageQueryError> { + let sql = format!( + "SELECT {key} AS k, cost_input, cost_output, cost_cache_read, \ + cost_cache_creation, cost_reasoning, cost_total \ + FROM requests \ + WHERE COALESCE(received_at, completed_at) >= ?1 AND cost_total IS NOT NULL" + ); + let mut sums: HashMap, [rust_decimal::Decimal; 6]> = HashMap::new(); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([since_ms], |row| { + Ok(( + row.get::<_, Option>(0)?, + row.get::<_, Option>(1)?, + row.get::<_, Option>(2)?, + row.get::<_, Option>(3)?, + row.get::<_, Option>(4)?, + row.get::<_, Option>(5)?, + row.get::<_, Option>(6)?, + )) + })?; + for row in rows { + let (key, input, output, cache_read, cache_creation, reasoning, total) = row?; + let entry = sums.entry(key).or_insert([rust_decimal::Decimal::ZERO; 6]); + let parts = [ + parse_cost(input)?, + parse_cost(output)?, + parse_cost(cache_read)?, + parse_cost(cache_creation)?, + parse_cost(reasoning)?, + parse_cost(total)?, + ]; + for (slot, part) in entry.iter_mut().zip(parts) { + *slot += part; + } + } + Ok(sums) +} + +fn cost_totals_from(sums: [rust_decimal::Decimal; 6]) -> CostTotals { + let show = |d: rust_decimal::Decimal| d.normalize().to_string(); + CostTotals { + input: show(sums[0]), + output: show(sums[1]), + cache_read: show(sums[2]), + cache_creation: show(sums[3]), + reasoning: show(sums[4]), + total: show(sums[5]), + } +} + +fn map_request_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<(RequestRow, i64)> { + let input: Option = row.get(10)?; + let output: Option = row.get(11)?; + let reasoning: Option = row.get(12)?; + let cache_read: Option = row.get(13)?; + let cache_creation: Option = row.get(14)?; + let cost_input: Option = row.get(16)?; + let cost_output: Option = row.get(17)?; + let cost_cache_creation: Option = row.get(18)?; + let cost_cache_read: Option = row.get(19)?; + let cost_reasoning: Option = row.get(20)?; + let cost_total: Option = row.get(21)?; + let streaming: Option = row.get(7)?; + let tokens = input.map(|input| TokenCounts { + input, + output: output.unwrap_or(0), + reasoning: reasoning.unwrap_or(0), + cache_read: cache_read.unwrap_or(0), + cache_creation: cache_creation.unwrap_or(0), + }); + let cost = cost_total.map(|total| CostTotals { + input: cost_input.unwrap_or_else(|| "0".to_owned()), + output: cost_output.unwrap_or_else(|| "0".to_owned()), + cache_read: cost_cache_read.unwrap_or_else(|| "0".to_owned()), + cache_creation: cost_cache_creation.unwrap_or_else(|| "0".to_owned()), + reasoning: cost_reasoning.unwrap_or_else(|| "0".to_owned()), + total, + }); + let row_out = RequestRow { + request_id: row.get(0)?, + received_at_ms: row.get(1)?, + completed_at_ms: row.get(2)?, + caller: row.get(3)?, + provider: row.get(4)?, + route_kind: row.get(5)?, + client_protocol: row.get(6)?, + streaming: streaming.map(|flag| flag != 0), + model_requested: row.get(8)?, + model_upstream: row.get(9)?, + tokens, + usage_provenance: row.get(15)?, + cost, + stop_reason: row.get(22)?, + error_kind: row.get(23)?, + http_status: row.get(24)?, + latency_ms: row.get(25)?, + body_hash: row.get(26)?, + }; + let ts: Option = row.get(27)?; + Ok((row_out, ts.unwrap_or(i64::MIN))) +} + +fn requests_sync(path: &Path, filter: &RequestFilter) -> Result { + let conn = UsageReader::open(path)?; + let mut clauses: Vec<&'static str> = Vec::new(); + let mut params: Vec = Vec::new(); + if let Some(since) = filter.since_ms { + clauses.push("COALESCE(received_at, completed_at) >= ?"); + params.push(SqlValue::Integer(since)); + } + if let Some(until) = filter.until_ms { + clauses.push("COALESCE(received_at, completed_at) < ?"); + params.push(SqlValue::Integer(until)); + } + if let Some(provider) = &filter.provider { + clauses.push("provider = ?"); + params.push(SqlValue::Text(provider.clone())); + } + if let Some((cursor_ts, cursor_id)) = &filter.cursor { + clauses.push( + "(COALESCE(received_at, completed_at) < ? \ + OR (COALESCE(received_at, completed_at) = ? AND request_id < ?))", + ); + params.push(SqlValue::Integer(*cursor_ts)); + params.push(SqlValue::Integer(*cursor_ts)); + params.push(SqlValue::Text(cursor_id.clone())); + } + let where_sql = if clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", clauses.join(" AND ")) + }; + let sql = format!( + "SELECT request_id, received_at, completed_at, caller, provider, route_kind, \ + client_protocol, streaming, model_requested, model_upstream, input_tokens, \ + output_tokens, reasoning_tokens, cache_read_tokens, cache_creation_tokens, \ + usage_provenance, cost_input, cost_output, cost_cache_creation, cost_cache_read, \ + cost_reasoning, cost_total, stop_reason, error_kind, http_status, latency_ms, \ + body_hash, COALESCE(received_at, completed_at) AS ts \ + FROM requests {where_sql} \ + ORDER BY ts DESC, request_id DESC \ + LIMIT ?" + ); + params.push(SqlValue::Integer(filter.limit + 1)); + let mut fetched: Vec<(RequestRow, i64)> = Vec::new(); + { + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map(rusqlite::params_from_iter(params), map_request_row)?; + for row in rows { + fetched.push(row?); + } + } + let has_more = fetched.len() > usize::try_from(filter.limit).unwrap_or(0); + if has_more { + fetched.truncate(usize::try_from(filter.limit).unwrap_or(0)); + } + let next_cursor = if has_more { + fetched + .last() + .map(|(row, ts)| format!("{ts}:{}", row.request_id)) + } else { + None + }; + Ok(RequestsPage { + rows: fetched.into_iter().map(|(row, _ts)| row).collect(), + next_cursor, + }) +} + +fn callers_sync(path: &Path) -> Result, UsageQueryError> { + let conn = UsageReader::open(path)?; + let mut stmt = conn.prepare( + "SELECT caller, COUNT(*), \ + SUM(CASE WHEN error_kind IS NOT NULL THEN 1 ELSE 0 END), \ + MAX(COALESCE(received_at, completed_at)) \ + FROM requests GROUP BY caller ORDER BY 4 DESC NULLS LAST", + )?; + let rows = stmt.query_map([], |row| { + Ok(CallerSummary { + caller: row.get(0)?, + requests: row.get(1)?, + errors: row.get::<_, Option>(2)?.unwrap_or(0), + last_seen_ms: row.get(3)?, + }) + })?; + let mut callers = Vec::new(); + for row in rows { + callers.push(row?); + } + Ok(callers) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn read_only_connections_refuse_writes() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("usage.db"); + { + let writer = Connection::open(&db).expect("create db"); + writer + .execute_batch("CREATE TABLE t (a INTEGER);") + .expect("seed table"); + } + let reader = UsageReader::open(&db).expect("open read-only"); + let attempt = reader.execute("INSERT INTO t (a) VALUES (1)", []); + assert!( + attempt.is_err(), + "SQLITE_OPEN_READ_ONLY must refuse writes at the connection level" + ); + let count: i64 = reader + .query_row("SELECT COUNT(*) FROM t", [], |row| row.get(0)) + .expect("read on the read-only connection"); + assert_eq!(count, 0); + } + + #[test] + fn group_by_parses_known_values_and_rejects_unknown() { + assert_eq!(GroupBy::parse("provider"), Some(GroupBy::Provider)); + assert_eq!(GroupBy::parse("model"), Some(GroupBy::Model)); + assert_eq!(GroupBy::parse("caller"), Some(GroupBy::Caller)); + assert_eq!(GroupBy::parse("team"), None); + assert_eq!(GroupBy::Provider.as_str(), "provider"); + assert_eq!(GroupBy::Model.as_str(), "model"); + assert_eq!(GroupBy::Caller.as_str(), "caller"); + } +} diff --git a/crates/llm-proxy-api/src/routes.rs b/crates/llm-proxy-api/src/routes.rs new file mode 100644 index 0000000..1173b32 --- /dev/null +++ b/crates/llm-proxy-api/src/routes.rs @@ -0,0 +1,274 @@ +use std::path::PathBuf; + +use axum::Json; +use axum::extract::{Query, State}; +use axum::http::{HeaderValue, StatusCode, header}; +use axum::middleware::from_fn_with_state; +use axum::response::{IntoResponse, Response}; +use axum::routing::get; +use serde::Deserialize; +use serde::Serialize; +use tower_http::catch_panic::CatchPanicLayer; +use tower_http::set_header::SetResponseHeaderLayer; +use tower_http::trace::TraceLayer; + +use crate::auth::{AdminAuth, require_admin}; +use crate::reader::{GroupBy, MAX_PAGE_LIMIT, RequestFilter, UsageReader, now_unix_millis}; + +const DEFAULT_WINDOW: std::time::Duration = std::time::Duration::from_secs(24 * 3600); + +#[derive(Clone, Debug)] +pub(crate) struct UsageState { + pub(crate) reader: UsageReader, + pub(crate) admin: AdminAuth, +} + +/// Read-only, admin-authenticated `/usage` query surface over the usage +/// ledger: every route is a `GET` behind the admin bearer token, and no +/// response carries request bodies. +#[derive(Debug)] +pub struct UsageApi { + state: UsageState, +} + +impl UsageApi { + /// Build the `/usage` surface over the ledger at `db_path`, guarded by + /// `admin_token`. + #[must_use] + pub fn new(db_path: impl Into, admin_token: &str) -> Self { + Self { + state: UsageState { + reader: UsageReader::new(db_path), + admin: AdminAuth::new(admin_token), + }, + } + } + + /// The `/usage` router with state applied, ready to serve or merge; admin + /// middleware and `nosniff`/panic-catch/trace layers are already applied. + #[must_use = "the Router must be used with a server or merged into one"] + pub fn router(self) -> axum::Router { + let state = self.state; + axum::Router::new() + .route("/usage/summary", get(summary)) + .route("/usage/requests", get(requests)) + .route("/usage/keys", get(keys)) + .layer(from_fn_with_state(state.clone(), require_admin)) + .layer(TraceLayer::new_for_http()) + .layer(SetResponseHeaderLayer::if_not_present( + header::X_CONTENT_TYPE_OPTIONS, + HeaderValue::from_static("nosniff"), + )) + .layer(CatchPanicLayer::new()) + .with_state(state) + } +} + +#[derive(Debug)] +pub(crate) enum UsageApiError { + InvalidParam(String), + Query(crate::reader::UsageQueryError), +} + +#[derive(Serialize)] +struct ErrorBody { + error: ErrorDetail, +} + +#[derive(Serialize)] +struct ErrorDetail { + r#type: &'static str, + message: String, +} + +impl IntoResponse for UsageApiError { + fn into_response(self) -> Response { + match self { + Self::InvalidParam(message) => ( + StatusCode::BAD_REQUEST, + Json(ErrorBody { + error: ErrorDetail { + r#type: "invalid_request", + message, + }, + }), + ) + .into_response(), + Self::Query(error) => { + tracing::error!(error = %error, "usage ledger query failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(ErrorBody { + error: ErrorDetail { + r#type: "internal_error", + message: "usage ledger query failed".to_owned(), + }, + }), + ) + .into_response() + } + } + } +} + +#[derive(Deserialize)] +struct SummaryParams { + window: Option, + group_by: Option, +} + +#[derive(Serialize)] +struct SummaryBody { + since_ms: i64, + until_ms: i64, + group_by: &'static str, + groups: Vec, +} + +async fn summary( + State(state): State, + Query(params): Query, +) -> Result, UsageApiError> { + let until_ms = now_unix_millis(); + let window_ms = match ¶ms.window { + None => DEFAULT_WINDOW.as_millis(), + Some(raw) => { + let parsed = humantime::parse_duration(raw).map_err(|_parse_error| { + UsageApiError::InvalidParam(format!( + "window must be a duration like \"1h\" or \"30m\", got {raw:?}" + )) + })?; + if parsed.is_zero() { + return Err(UsageApiError::InvalidParam( + "window must be a positive duration".to_owned(), + )); + } + parsed.as_millis() + } + }; + let since_ms = until_ms.saturating_sub(i64::try_from(window_ms).unwrap_or(i64::MAX)); + let group = match params.group_by.as_deref() { + None | Some("provider") => GroupBy::Provider, + Some(other) => GroupBy::parse(other).ok_or_else(|| { + UsageApiError::InvalidParam(format!( + "group_by must be one of provider, model, caller; got {other:?}" + )) + })?, + }; + let groups = state + .reader + .summary(since_ms, group) + .await + .map_err(UsageApiError::Query)?; + Ok(Json(SummaryBody { + since_ms, + until_ms, + group_by: group.as_str(), + groups, + })) +} + +#[derive(Deserialize)] +struct RequestsParams { + since: Option, + until: Option, + provider: Option, + limit: Option, + cursor: Option, +} + +#[derive(Serialize)] +struct RequestsBody { + rows: Vec, + next_cursor: Option, +} + +fn parse_rfc3339_ms(name: &str, raw: &str) -> Result { + use time::format_description::well_known::Rfc3339; + time::OffsetDateTime::parse(raw, &Rfc3339) + .map(|ts| { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) + }) + .map_err(|_parse_error| { + UsageApiError::InvalidParam(format!( + "{name} must be an RFC 3339 timestamp like 2026-01-02T03:04:05Z, got {raw:?}" + )) + }) +} + +fn parse_cursor(raw: &str) -> Result<(i64, String), UsageApiError> { + let Some((ts, id)) = raw.split_once(':') else { + return Err(UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + )); + }; + let ts: i64 = ts.parse().map_err(|_parse_error| { + UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + ) + })?; + if id.is_empty() { + return Err(UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + )); + } + Ok((ts, id.to_owned())) +} + +async fn requests( + State(state): State, + Query(params): Query, +) -> Result, UsageApiError> { + let since_ms = match ¶ms.since { + Some(raw) => Some(parse_rfc3339_ms("since", raw)?), + None => None, + }; + let until_ms = match ¶ms.until { + Some(raw) => Some(parse_rfc3339_ms("until", raw)?), + None => None, + }; + if let (Some(since), Some(until)) = (since_ms, until_ms) + && since >= until + { + return Err(UsageApiError::InvalidParam( + "since must be before until".to_owned(), + )); + } + let limit = params.limit.unwrap_or(100); + if !(1..=MAX_PAGE_LIMIT).contains(&limit) { + return Err(UsageApiError::InvalidParam(format!( + "limit must be between 1 and {MAX_PAGE_LIMIT}, got {limit}" + ))); + } + let cursor = match ¶ms.cursor { + Some(raw) => Some(parse_cursor(raw)?), + None => None, + }; + let page = state + .reader + .requests(RequestFilter { + since_ms, + until_ms, + provider: params.provider, + limit, + cursor, + }) + .await + .map_err(UsageApiError::Query)?; + Ok(Json(RequestsBody { + rows: page.rows, + next_cursor: page.next_cursor, + })) +} + +#[derive(Serialize)] +struct KeysBody { + callers: Vec, +} + +async fn keys(State(state): State) -> Result, UsageApiError> { + let callers = state.reader.callers().await.map_err(UsageApiError::Query)?; + Ok(Json(KeysBody { callers })) +} diff --git a/crates/llm-proxy-api/tests/secret_abuse.rs b/crates/llm-proxy-api/tests/secret_abuse.rs new file mode 100644 index 0000000..083dee2 --- /dev/null +++ b/crates/llm-proxy-api/tests/secret_abuse.rs @@ -0,0 +1,197 @@ +//! Secret-abuse checks for the `/usage` surface: no credential reaches the +//! captured tracing output or the raw ledger bytes, and every negative scan +//! first asserts seeded ledger content is findable. + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::sync::{Arc, Mutex, MutexGuard}; + use std::time::Duration; + + use axum::Router; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use llm_proxy_storage::{EventBus, ProxyEvent, RequestReceived, SqliteEventBus}; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "prx6-api-admin-token-0a3f5e91"; + const WRONG_TOKEN: &str = "prx6-api-attacker-guess-2c4d"; + const CLIENT_TOKEN: &str = "prx6-api-client-token-7e8f"; + const CALLER_MARKER: &str = "prx6-seeded-caller-marker-4451"; + + async fn seeded_ledger(dir: &tempfile::TempDir) -> PathBuf { + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens"); + bus.emit(&ProxyEvent::RequestReceived(RequestReceived { + request_id: "prx6-seed".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: CALLER_MARKER.to_owned(), + provider: "prov-seed".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "prx6-seed-digest".to_owned(), + })); + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joins"); + path + } + + fn usage_app(path: &Path) -> Router { + llm_proxy_api::UsageApi::new(path.to_owned(), ADMIN_TOKEN).router() + } + + fn get(uri: &str, token: Option<&str>) -> Request { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + builder.body(Body::empty()).expect("request builds") + } + + async fn drive_all_outcomes(path: &Path) { + for uri in ["/usage/summary", "/usage/requests", "/usage/keys"] { + for token in [ + Some(ADMIN_TOKEN), + Some(WRONG_TOKEN), + Some(CLIENT_TOKEN), + None, + ] { + let (status, body) = send(usage_app(path).clone(), get(uri, token)).await; + let expected = if token == Some(ADMIN_TOKEN) { + StatusCode::OK + } else { + StatusCode::UNAUTHORIZED + }; + assert_eq!(status, expected, "{uri} with {token:?}: {body}"); + } + } + } + + async fn send(app: Router, request: Request) -> (StatusCode, String) { + let response = app.oneshot(request).await.expect("oneshot completes"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body reads"); + (status, String::from_utf8_lossy(&bytes).into_owned()) + } + + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + struct TraceCapture { + buffer: Arc>>, + } + + static CAPTURE_BUFFER: std::sync::OnceLock>>> = std::sync::OnceLock::new(); + static CAPTURE_INSTALLED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + + impl TraceCapture { + fn install() -> Self { + use tracing_subscriber::layer::SubscriberExt; + let buffer = CAPTURE_BUFFER.get_or_init(|| Arc::new(Mutex::new(Vec::::new()))); + CAPTURE_INSTALLED.get_or_init(|| { + let layer = tracing_subscriber::fmt::layer() + .with_writer(BufferWriter(Arc::clone(buffer))) + .with_ansi(false); + tracing::subscriber::set_global_default(tracing_subscriber::registry().with(layer)) + .expect("the global capture subscriber installs exactly once"); + }); + Self { + buffer: Arc::clone(buffer), + } + } + + fn text(&self) -> String { + let bytes = self.buffer.lock().expect("buffer lock").clone(); + String::from_utf8_lossy(&bytes).into_owned() + } + } + + #[derive(Clone)] + struct BufferWriter(Arc>>); + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufferWriter { + type Writer = BufferWriterGuard<'a>; + + fn make_writer(&'a self) -> Self::Writer { + BufferWriterGuard(self.0.lock().expect("buffer lock")) + } + } + + struct BufferWriterGuard<'a>(MutexGuard<'a, Vec>); + + impl std::io::Write for BufferWriterGuard<'_> { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn no_credential_ever_lands_in_the_raw_ledger_file_bytes() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = seeded_ledger(&dir).await; + drive_all_outcomes(&path).await; + + let bytes = ledger_bytes(&path); + let text = String::from_utf8_lossy(&bytes); + assert!( + text.contains(CALLER_MARKER) && text.contains("prx6-seed-digest"), + "control: seeded persisted values must be findable in the ledger bytes" + ); + for secret in [ADMIN_TOKEN, WRONG_TOKEN, CLIENT_TOKEN] { + assert!( + !text.contains(secret), + "credential {secret} leaked into the raw ledger file bytes" + ); + } + } + + #[tokio::test] + async fn admin_token_never_appears_in_captured_usage_tracing() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = seeded_ledger(&dir).await; + + let capture = TraceCapture::install(); + drive_all_outcomes(&path).await; + let captured = capture.text(); + assert!( + captured.contains("started processing request"), + "control: usage request traffic must land in the capture, got: {captured}" + ); + for secret in [ADMIN_TOKEN, WRONG_TOKEN, CLIENT_TOKEN] { + assert!( + !captured.contains(secret), + "credential {secret} leaked into tracing output: {captured}" + ); + } + } +} diff --git a/crates/llm-proxy-api/tests/usage_api.rs b/crates/llm-proxy-api/tests/usage_api.rs new file mode 100644 index 0000000..12f0d09 --- /dev/null +++ b/crates/llm-proxy-api/tests/usage_api.rs @@ -0,0 +1,612 @@ +//! Integration tests for the `/usage` query surface over a real SQLite +//! ledger seeded through `SqliteEventBus`. + +#[cfg(test)] +mod tests { + use axum::body::Body; + use axum::http::Request; + use axum::http::StatusCode; + use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage, UsageProvenance}; + use llm_proxy_storage::{ + EventBus, ProxyEvent, RequestReceived, ResponseCompleted, ResponseFailed, SqliteEventBus, + }; + use rust_decimal::Decimal; + use serde_json::Value; + use time::OffsetDateTime; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "test-admin-token"; + + struct Ledger { + _dir: tempfile::TempDir, + path: std::path::PathBuf, + t0: OffsetDateTime, + } + + fn received_event( + id: &str, + caller: &str, + provider: &str, + model: &str, + ts: OffsetDateTime, + ) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: id.to_owned(), + timestamp: ts, + caller: caller.to_owned(), + provider: provider.to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: model.to_owned(), + upstream: None, + }, + streaming: true, + body_hash: format!("sha256-digest-{id}"), + }) + } + + fn completed_event( + id: &str, + provider: &str, + model: &str, + ts: OffsetDateTime, + total: &str, + ) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: id.to_owned(), + timestamp: ts, + provider: provider.to_owned(), + upstream_message_id: None, + model: ModelRef { + requested: model.to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 100, + output_tokens: 50, + reasoning_tokens: Some(7), + cache_creation_input_tokens: Some(30), + cache_read_input_tokens: Some(200), + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from_str_exact(total).expect("decimal fixture"), + output: Decimal::ZERO, + cache_creation: Decimal::ZERO, + cache_read: Decimal::ZERO, + reasoning: Decimal::ZERO, + total: Decimal::from_str_exact(total).expect("decimal fixture"), + }), + stop_reason: StopReason::EndTurn, + latency_ms: 4321, + }) + } + + fn failed_event(id: &str, provider: Option<&str>, ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: id.to_owned(), + timestamp: ts, + provider: provider.map(str::to_owned), + model: None, + error_kind: "UpstreamError".to_owned(), + message: "sanitized failure text".to_owned(), + http_status: 502, + latency_ms: 99, + }) + } + + async fn seed() -> Ledger { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: std::time::Duration::from_secs(3600), + }, + ) + .expect("open ledger"); + + let t0 = OffsetDateTime::now_utc().saturating_sub(time::Duration::seconds(60)); + let second = std::time::Duration::from_secs(1); + let events = [ + received_event("req-1", "laptop", "prov-a", "claude-3", t0), + completed_event( + "req-1", + "prov-a", + "claude-3", + t0 + second * 4, + "0.1000000000000000000000001", + ), + received_event("req-2", "laptop", "prov-a", "gpt-4", t0 + second), + completed_event( + "req-2", + "prov-a", + "gpt-4", + t0 + second * 5, + "0.2000000000000000000000002", + ), + received_event("req-3", "phone", "prov-b", "claude-3", t0 + second * 2), + failed_event("req-3", Some("prov-b"), t0 + second * 6), + failed_event("orphan-1", None, t0 + second * 3), + received_event( + "old-1", + "laptop", + "prov-old", + "claude-3", + t0.saturating_sub(time::Duration::seconds(7200)), + ), + completed_event( + "old-1", + "prov-old", + "claude-3", + t0.saturating_sub(time::Duration::seconds(7190)), + "0.5", + ), + ]; + for event in &events { + bus.emit(event); + } + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + Ledger { + _dir: dir, + path, + t0, + } + } + + fn app_for(path: &std::path::Path) -> axum::Router { + llm_proxy_api::UsageApi::new(path.to_owned(), ADMIN_TOKEN).router() + } + + async fn get_json(app: axum::Router, uri: &str, token: Option<&str>) -> (StatusCode, Value) { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + let response = app + .oneshot(builder.body(Body::empty()).expect("request builds")) + .await + .expect("oneshot"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + let json = serde_json::from_slice(&bytes).expect("body is JSON"); + (status, json) + } + + #[tokio::test] + async fn usage_routes_refuse_missing_and_wrong_tokens_without_echoing_them() { + let ledger = seed().await; + for path in ["/usage/summary", "/usage/requests", "/usage/keys"] { + let (status, body) = get_json(app_for(&ledger.path), path, None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED, "{path} without token"); + assert_eq!( + body["error"]["message"], "missing or invalid admin token", + "uniform 401 body for {path}" + ); + + let (status, body_text) = { + let (status, body) = + get_json(app_for(&ledger.path), path, Some("wrong-token")).await; + (status, body.to_string()) + }; + assert_eq!(status, StatusCode::UNAUTHORIZED, "{path} with wrong token"); + assert!( + !body_text.contains("wrong-token") && !body_text.contains(ADMIN_TOKEN), + "the 401 body must not echo any presented or configured token: {body_text}" + ); + + let (status, _) = get_json(app_for(&ledger.path), path, Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK, "{path} with the admin token"); + } + } + + #[tokio::test] + async fn summary_groups_by_provider_with_exact_decimal_costs() { + let ledger = seed().await; + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=provider", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + + let groups = body["groups"].as_array().expect("groups array"); + let find = |key: &str| { + groups + .iter() + .find(|g| g["key"] == key) + .unwrap_or_else(|| panic!("group {key} missing: {groups:?}")) + }; + + let prov_a = find("prov-a"); + assert_eq!(prov_a["requests"], 2); + assert_eq!(prov_a["errors"], 0, "both prov-a requests completed"); + assert_eq!( + prov_a["tokens"]["input"], 200, + "100 input tokens per request" + ); + assert_eq!(prov_a["tokens"]["output"], 100); + assert_eq!(prov_a["tokens"]["reasoning"], 14); + assert_eq!(prov_a["tokens"]["cache_read"], 400); + assert_eq!(prov_a["tokens"]["cache_creation"], 60); + assert_eq!( + prov_a["cost"]["total"], "0.3000000000000000000000003", + "cost sums must be exact decimals — a float sum of these two values \ + yields 0.30000000000000004" + ); + assert_eq!( + prov_a["cost"]["input"], "0.3000000000000000000000003", + "input-cost column sums exactly too" + ); + + let prov_b = find("prov-b"); + assert_eq!(prov_b["requests"], 1); + assert_eq!(prov_b["errors"], 1); + + let null_group = groups + .iter() + .find(|g| g["key"].is_null()) + .expect("failure-only rows group under a null provider"); + assert_eq!(null_group["requests"], 1); + assert_eq!(null_group["errors"], 1); + + assert_eq!(body["group_by"], "provider"); + } + + #[tokio::test] + async fn summary_groups_by_caller_and_model() { + let ledger = seed().await; + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=caller", + Some(ADMIN_TOKEN), + ) + .await; + let groups = &body["groups"]; + let laptop = groups + .as_array() + .expect("groups") + .iter() + .find(|g| g["key"] == "laptop") + .expect("laptop group"); + assert_eq!( + laptop["requests"], 3, + "req-1, req-2 and the two-hour-old row" + ); + assert_eq!(laptop["errors"], 0); + + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=model", + Some(ADMIN_TOKEN), + ) + .await; + let model_groups = body["groups"].as_array().expect("groups"); + let claude = model_groups + .iter() + .find(|g| g["key"] == "claude-3") + .expect("claude-3 group"); + assert_eq!(claude["requests"], 3, "req-1, req-3 and old-1"); + let gpt = model_groups + .iter() + .find(|g| g["key"] == "gpt-4") + .expect("gpt-4 group"); + assert_eq!(gpt["requests"], 1); + } + + #[tokio::test] + async fn summary_window_excludes_rows_older_than_it() { + let ledger = seed().await; + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=1h", + Some(ADMIN_TOKEN), + ) + .await; + let groups = body["groups"].as_array().expect("groups"); + assert!( + groups.iter().all(|g| g["key"] != "prov-old"), + "the two-hour-old row must fall outside a 1h window: {groups:?}" + ); + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=4h", + Some(ADMIN_TOKEN), + ) + .await; + let groups = body["groups"].as_array().expect("groups"); + assert!( + groups.iter().any(|g| g["key"] == "prov-old"), + "a 4h window covers the two-hour-old row" + ); + + let (status, _) = get_json( + app_for(&ledger.path), + "/usage/summary?window=7d", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK, "\"7d\" must be an accepted window"); + } + + #[tokio::test] + async fn summary_rejects_bad_window_and_group_by() { + let ledger = seed().await; + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=abc", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body["error"]["message"].to_string().contains("window")); + + let (status, _) = get_json( + app_for(&ledger.path), + "/usage/summary?window=0s", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!( + status, + StatusCode::BAD_REQUEST, + "zero window is not a window" + ); + + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=team", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body["error"]["message"].to_string().contains("group_by")); + } + + #[tokio::test] + async fn requests_pages_newest_first_with_cursor_and_provider_filter() { + let ledger = seed().await; + let (status, page1) = get_json( + app_for(&ledger.path), + "/usage/requests?limit=2", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + let rows1 = page1["rows"].as_array().expect("rows"); + assert_eq!(rows1.len(), 2, "limit is honored"); + let first_id = rows1[0]["request_id"].as_str().expect("id"); + assert_eq!( + first_id, "orphan-1", + "newest first: orphan-1 completed latest among the seeded rows" + ); + let cursor = page1["next_cursor"].as_str().expect("cursor present"); + assert_eq!(rows1[1]["request_id"], "req-3"); + + let (status, page2) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?limit=10&cursor={cursor}"), + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + let rows2 = page2["rows"].as_array().expect("rows"); + assert!( + rows2 + .iter() + .all(|r| r["request_id"] != "orphan-1" && r["request_id"] != "req-3"), + "cursor paging must not repeat rows: {rows2:?}" + ); + + let (_, filtered) = get_json( + app_for(&ledger.path), + "/usage/requests?provider=prov-a", + Some(ADMIN_TOKEN), + ) + .await; + let filtered_rows = filtered["rows"].as_array().expect("rows"); + assert!( + !filtered_rows.is_empty() && filtered_rows.iter().all(|r| r["provider"] == "prov-a"), + "provider filter: {filtered_rows:?}" + ); + } + + #[tokio::test] + async fn requests_honors_since_and_until_bounds() { + let ledger = seed().await; + let until = (ledger.t0 + std::time::Duration::from_secs(3)) + .format(&time::format_description::well_known::Rfc3339) + .expect("rfc3339"); + let (_, body) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?until={until}"), + Some(ADMIN_TOKEN), + ) + .await; + let ids: Vec<&str> = body["rows"] + .as_array() + .expect("rows") + .iter() + .filter_map(|r| r["request_id"].as_str()) + .collect(); + assert_eq!( + ids, + vec!["req-3", "req-2", "req-1", "old-1"], + "until is exclusive on the ledger time axis (orphan-1 at t0+3s is out)" + ); + + let since = (ledger.t0 + std::time::Duration::from_secs(1)) + .format(&time::format_description::well_known::Rfc3339) + .expect("rfc3339"); + let (_, body) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?since={since}"), + Some(ADMIN_TOKEN), + ) + .await; + let ids: Vec<&str> = body["rows"] + .as_array() + .expect("rows") + .iter() + .filter_map(|r| r["request_id"].as_str()) + .collect(); + assert_eq!( + ids, + vec!["orphan-1", "req-3", "req-2"], + "since is inclusive" + ); + } + + #[tokio::test] + async fn requests_rejects_invalid_parameters() { + let ledger = seed().await; + let cases = [ + "/usage/requests?limit=0", + "/usage/requests?limit=1001", + "/usage/requests?limit=-5", + "/usage/requests?since=not-a-time", + "/usage/requests?until=not-a-time", + "/usage/requests?cursor=garbage", + "/usage/requests?cursor=123:", + "/usage/requests?since=2026-01-02T03:04:05Z&until=2026-01-01T00:00:00Z", + ]; + for uri in cases { + let (status, body) = get_json(app_for(&ledger.path), uri, Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}: {body}"); + } + } + + #[tokio::test] + async fn requests_rows_carry_digests_but_never_bodies() { + let ledger = seed().await; + let (_, body) = get_json(app_for(&ledger.path), "/usage/requests", Some(ADMIN_TOKEN)).await; + let text = body.to_string(); + assert!( + text.contains("sha256-digest-"), + "rows expose the body digest for correlation" + ); + for row in body["rows"].as_array().expect("rows") { + let keys: Vec<&str> = row + .as_object() + .expect("row object") + .keys() + .map(String::as_str) + .collect(); + assert!( + !keys + .iter() + .any(|k| *k == "body" || *k == "request_body" || *k == "message"), + "raw rows must never carry request bodies or sanitized messages: {keys:?}" + ); + } + } + + #[tokio::test] + async fn keys_lists_callers_with_counts_and_last_seen() { + let ledger = seed().await; + let (status, body) = + get_json(app_for(&ledger.path), "/usage/keys", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + let callers = body["callers"].as_array().expect("callers"); + let find = |name: &str| { + callers + .iter() + .find(|c| c["caller"] == name) + .unwrap_or_else(|| panic!("caller {name} missing: {callers:?}")) + }; + let laptop = find("laptop"); + assert_eq!(laptop["requests"], 3, "req-1, req-2, old-1"); + assert_eq!(laptop["errors"], 0); + let phone = find("phone"); + assert_eq!(phone["requests"], 1); + assert_eq!(phone["errors"], 1); + assert!(laptop["last_seen_ms"].is_u64(), "last_seen_ms is populated"); + + let dated: Vec = callers + .iter() + .filter_map(|c| c["last_seen_ms"].as_i64()) + .collect(); + assert!( + !dated.is_empty(), + "callers with ledger activity must carry last_seen_ms: {callers:?}" + ); + assert_eq!( + dated, + { + let mut sorted = dated.clone(); + sorted.sort_by(|a, b| b.cmp(a)); + sorted + }, + "callers appear most-recently-seen first" + ); + let order: Vec<&str> = callers + .iter() + .map(|c| c["caller"].as_str().unwrap_or("")) + .collect(); + assert_eq!( + order, + vec!["", "phone", "laptop"], + "caller ordering follows last-seen timestamps, null caller included" + ); + } + + #[tokio::test] + async fn summary_serves_while_the_ledger_writer_is_active() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("usage.db"); + let bus = std::sync::Arc::new( + SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: std::time::Duration::from_secs(3600), + }, + ) + .expect("open ledger"), + ); + bus.emit(&received_event( + "live-1", + "laptop", + "prov-live", + "claude-3", + OffsetDateTime::now_utc(), + )); + + let mut seen = false; + for _ in 0..200 { + let (status, body) = get_json( + app_for(&path), + "/usage/summary?group_by=provider", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK, "query against live writer"); + let found = body["groups"] + .as_array() + .is_some_and(|g| g.iter().any(|x| x["key"] == "prov-live")); + if found { + seen = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!( + seen, + "the row must become visible while the writer is active" + ); + let closer = std::sync::Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + } +} diff --git a/crates/llm-proxy-core/src/catalog.rs b/crates/llm-proxy-core/src/catalog.rs index c10a931..fbd07b3 100644 --- a/crates/llm-proxy-core/src/catalog.rs +++ b/crates/llm-proxy-core/src/catalog.rs @@ -14,13 +14,8 @@ pub struct CatalogFileMetadata { /// Catalog source, normally `live`. pub source: String, /// Best-effort timestamp string (normally RFC 3339) recording when the - /// catalog was generated. - /// - /// This is an opaque, operator-controlled cache value that is **not** - /// validated at parse time. Format validity is checked only where the - /// value is consumed (e.g. staleness comparison in `catalog_service`); a - /// malformed value parses to `None` there and the entry is treated as - /// stale. Do not rely on this field being a well-formed RFC 3339 string. + /// catalog was generated. Opaque and **not** validated at parse time: + /// consumers treat a malformed value as stale. pub generated_at: String, /// Discovered model entries. #[serde(default)] @@ -63,10 +58,8 @@ pub fn parse_catalog_file(raw: &str) -> Result { /// Merge static and discovered catalog entries according to provider config. /// -/// In `Hybrid` mode, static entries override discovered entries that share the -/// same `id` (static metadata wins). Each input is borrowed, so each retained -/// entry is cloned exactly once — no second clone of the `id` for a throwaway -/// map key. +/// In `Hybrid` mode, static entries override discovered entries that share +/// the same `id` (static metadata wins). /// /// ``` /// use std::time::Duration; @@ -96,11 +89,6 @@ pub fn merge_catalog( config: &ProviderCatalogConfig, discovered: &[StaticModelCatalogEntry], ) -> Vec { - // Deduplicate by id via a `Vec` instead of a `BTreeMap` so the - // id is not cloned a second time purely to serve as a discarded map key - // (the value already carries its own `id`). The catalog is bounded and - // built only on the cold refresh path, so the linear override lookup is - // cheaper than the redundant allocation it replaces. let mut entries: Vec = Vec::new(); if matches!( @@ -117,7 +105,6 @@ pub fn merge_catalog( ) { for model in &config.models { if let Some(existing) = entries.iter_mut().find(|entry| entry.id == model.id) { - // Static metadata overrides the discovered entry for this id. *existing = model.clone(); } else { entries.push(model.clone()); @@ -125,15 +112,6 @@ pub fn merge_catalog( } } - // Restore the deterministic by-id ordering that the previous `BTreeMap` - // dedup provided as a side effect of keyed insertion. A stable, sorted - // model list is an observable contract of the `/v1/models` endpoint - // (clients paginate and snapshot-test on it; `first_id`/`last_id` depend on - // a fixed order), so the clone-elimination refactor above must not drop it. - // Sorting the owned Vec in place by the already-present `id` field avoids - // the second `id.clone()` that the `BTreeMap` key required, so - // this keeps the GAP-LOW-7 benefit without sacrificing sorted output. Ids - // are unique after dedup, so the unstable variant is safe. entries.sort_unstable_by(|a, b| a.id.cmp(&b.id)); entries .into_iter() @@ -164,9 +142,6 @@ pub fn merge_catalog( /// ``` #[must_use] pub fn model_allowed(model: &str, allow: &[String], deny: &[String]) -> bool { - // Evaluate each `allow` pattern against `model` exactly once; both - // `allowed` and `explicitly_allowed` are derived from this single pass so - // `glob_matches` (which takes the `GLOB_CACHE` mutex) isn't re-run. let matched: Vec = allow .iter() .map(|pattern| glob_matches(pattern, model)) @@ -180,20 +155,8 @@ pub fn model_allowed(model: &str, allow: &[String], deny: &[String]) -> bool { allowed && (!denied || explicitly_allowed) } -/// Cache of compiled glob-to-regex patterns. -/// -/// Avoids recompiling the same glob pattern on every call to [`glob_matches`]. -/// Uses `LazyLock` for thread-safe one-time initialization per unique pattern. static GLOB_CACHE: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); -/// Match `value` against a simple glob `pattern` supporting only `*` and `?`. -/// -/// Characters that are special in regex (`.`, `+`, `(`, etc.) are escaped -/// via [`regex::escape`] so they match literally. Invalid patterns silently -/// return `false`. -/// -/// Compiled regexes are cached for reuse so the same pattern is only -/// compiled once per process. fn glob_matches(pattern: &str, value: &str) -> bool { let mut cache = GLOB_CACHE.lock().unwrap_or_else(|e| e.into_inner()); if let Some((_, regex)) = cache.iter().find(|(p, _)| p == pattern) { @@ -266,13 +229,11 @@ mod tests { &config(ProviderCatalogMode::Hybrid), &[model("shared", Some("discovered")), model("other", None)], ); - // The static entry should override the discovered entry for "shared". let shared = merged .iter() .find(|m| m.id == "shared") .expect("shared model should be present"); assert_eq!(shared.display_name.as_deref(), Some("static")); - // The discovered-only "other" should also be present. let other = merged .iter() .find(|m| m.id == "other") @@ -282,11 +243,6 @@ mod tests { #[test] fn merge_catalog_sorts_output_by_id() { - // Regression guard for the sorted-by-id contract: the `/v1/models` - // endpoint paginates and snapshot-tests on a stable, lexicographic - // model order. An earlier BTreeMap dedup provided this for free; the - // Vec dedup that replaced it does not, so the sort is now explicit. - // Feed models in deliberately non-sorted insertion order. let cfg = ProviderCatalogConfig { mode: ProviderCatalogMode::Discovered, enforce: false, @@ -311,10 +267,6 @@ mod tests { #[test] fn merge_catalog_sorts_after_static_override() { - // The static-override path updates a discovered entry in place (keeping - // its discovered-list position) rather than re-appending, so the final - // sort must run AFTER the override. Verify a static entry that overrides - // a discovered one still lands in its correct sorted position. let cfg = ProviderCatalogConfig { mode: ProviderCatalogMode::Hybrid, enforce: false, @@ -367,7 +319,6 @@ mod tests { assert_eq!(merged.len(), 2); assert!(merged.iter().any(|m| m.id == "discovered-a")); assert!(merged.iter().any(|m| m.id == "discovered-b")); - // Static "shared" model should NOT be present in Discovered mode. assert!(!merged.iter().any(|m| m.id == "shared")); } @@ -381,28 +332,22 @@ mod tests { model("shared", Some("discovered")), ], ); - // Only static entries should be present. assert_eq!(merged.len(), 1); assert_eq!(merged[0].id, "shared"); - // The static display_name should win since only static entries are used. assert_eq!(merged[0].display_name.as_deref(), Some("static")); } #[test] fn glob_matches_special_regex_characters() { - // Dots should be treated literally, not as regex wildcards. assert!(glob_matches("model.v2", "model.v2")); assert!(!glob_matches("model.v2", "modelXv2")); - // Plus signs, parentheses should be literal. assert!(glob_matches("model+v2", "model+v2")); assert!(!glob_matches("model+v2", "modelvv2")); - // Parentheses. assert!(glob_matches("model(test)", "model(test)")); } #[test] fn glob_matches_slashes_in_model_ids() { - // Real-world model IDs with slashes and dots. assert!(glob_matches( "accounts/fireworks/models/deepseek-v3p1", "accounts/fireworks/models/deepseek-v3p1" diff --git a/crates/llm-proxy-core/src/env_interpolate.rs b/crates/llm-proxy-core/src/env_interpolate.rs index 22576f7..e17eabe 100644 --- a/crates/llm-proxy-core/src/env_interpolate.rs +++ b/crates/llm-proxy-core/src/env_interpolate.rs @@ -1,46 +1,25 @@ //! Shared environment-variable interpolation for configuration files. -//! -//! Provides a single [`crate::env_interpolate::interpolate_env_vars`] function and the compiled regex -//! used by the TOML provider config ([`crate::provider_config`]). use regex::Regex; -/// The compiled regex for `${VAR_NAME}` patterns. -/// -/// Shared across all callers so the regex is compiled at most once per process. -/// Uses `LazyLock` for idiomatic one-time initialization. static ENV_VAR_RE: std::sync::LazyLock = std::sync::LazyLock::new(|| { - // INVARIANT: this regex is a compile-time constant that is syntactically valid. - // Failure here indicates a programming error in the regex literal, not a - // runtime condition. Regex::new(r"\$\{([A-Za-z0-9_]+)\}").expect("env var regex is valid") }); -/// Return the shared compiled regex for `${VAR_NAME}` patterns. pub(crate) fn env_var_regex() -> &'static Regex { &ENV_VAR_RE } /// Replace `${ENV_VAR}` patterns in `input` with the value of the -/// corresponding environment variable. -/// -/// Unset variables are left as-is (the `${VAR}` pattern remains in the output). +/// corresponding environment variable; unset variables are left as-is. /// /// # Trust boundary /// -/// Values are substituted as raw text. Callers must ensure environment -/// variables contain no TOML-breaking characters if the output will be -/// parsed as TOML. A compromised or misconfigured environment variable -/// could inject arbitrary TOML structure. This is acceptable because the -/// process owner controls the environment. -/// -/// # Whole-file interpolation +/// Values are substituted as raw text, so a misconfigured environment +/// variable can inject arbitrary TOML structure into the output. /// -/// This function operates on the entire raw file content before TOML/JSON -/// parsing. This means any `${VAR}` pattern anywhere in the file will be -/// interpolated, including in `api_key`, `endpoint`, `name`, and other string -/// fields. Single-braced patterns like `{model}` are **not** interpolated -/// (they use `{` not `${`). +/// Interpolation runs on the whole raw file content before parsing; +/// `{model}`-style single-brace patterns are **not** interpolated. pub fn interpolate_env_vars(input: &str) -> String { let re = env_var_regex(); re.replace_all(input, |caps: ®ex::Captures<'_>| { @@ -58,11 +37,8 @@ pub fn find_unresolved_env_var(s: &str) -> Option { re.captures(s).map(|caps| caps[1].to_owned()) } -/// Find all `${VAR}` patterns in `s` and return them as `(var_name, resolved_value)` pairs. -/// -/// This is used to track which env vars were interpolated so that empty -/// resolutions can be detected and reported distinctly from literal empty -/// strings. +/// Find all `${VAR}` patterns in `s` and return them as `(var_name, +/// resolved_value)` pairs; the value is `None` when the variable is unset. pub fn find_env_var_refs(s: &str) -> Vec<(String, Option)> { let re = env_var_regex(); re.captures_iter(s) @@ -74,10 +50,6 @@ pub fn find_env_var_refs(s: &str) -> Vec<(String, Option)> { .collect() } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -109,8 +81,6 @@ mod tests { assert_eq!(find_unresolved_env_var("no vars here"), None); } - // -- Non-matching patterns are left as-is ------------------------------------ - #[test] fn empty_braces_not_interpolated() { let result = interpolate_env_vars("key = ${}"); @@ -141,14 +111,9 @@ mod tests { #[test] fn numeric_only_var_not_interpolated() { let result = interpolate_env_vars("key = ${123}"); - // ${123} does match [A-Za-z0-9_]+ since digits are allowed. - // This is documented behavior: the regex allows numeric-only names. - // If 123 is not set, it stays as-is. assert_eq!(result, "key = ${123}"); } - // -- Multiple interpolations in a single string ------------------------- - #[test] fn multiple_vars_in_single_string() { let _lock = crate::test_support::TestEnvLock::acquire(); @@ -168,8 +133,6 @@ mod tests { assert_eq!(result, "yes ${_LLM_PROXY_NEVER_EXISTS_MIX_B}"); } - // -- Env var set to empty string ----------------------------------------- - #[test] fn var_set_to_empty_string_replaces_with_empty() { let _lock = crate::test_support::TestEnvLock::acquire(); @@ -178,8 +141,6 @@ mod tests { assert_eq!(result, "key = "); } - // -- Adjacent and edge-case patterns ------------------------------------- - #[test] fn adjacent_patterns() { let _lock = crate::test_support::TestEnvLock::acquire(); @@ -192,14 +153,9 @@ mod tests { #[test] fn dollar_sign_near_pattern() { let result = interpolate_env_vars("$$${_LLM_PROXY_NEVER_EXISTS_999}$$"); - // The inner ${...} is left as-is; surrounding $ characters are literal. - // The regex matches ${_LLM_PROXY_NEVER_EXISTS_999} and replaces it with - // itself (since it is unresolved), leaving the surrounding $$ intact. assert_eq!(result, "$$${_LLM_PROXY_NEVER_EXISTS_999}$$"); } - // -- find_env_var_refs tests --------------------------------------------- - #[test] fn find_env_var_refs_multiple_matches() { let _lock = crate::test_support::TestEnvLock::acquire(); @@ -209,7 +165,6 @@ mod tests { assert_eq!(refs[0].0, "_LLM_PROXY_TEST_REFS_A"); assert_eq!(refs[0].1, Some("val_a".to_owned())); assert_eq!(refs[1].0, "_LLM_PROXY_TEST_REFS_B"); - // Not set, so resolved value is None. assert_eq!(refs[1].1, None); } @@ -220,7 +175,6 @@ mod tests { let refs = find_env_var_refs("${_LLM_PROXY_TEST_REFS_EMPTY}"); assert_eq!(refs.len(), 1); assert_eq!(refs[0].0, "_LLM_PROXY_TEST_REFS_EMPTY"); - // Set but empty: Some(""). assert_eq!(refs[0].1, Some(String::new())); } diff --git a/crates/llm-proxy-core/src/error.rs b/crates/llm-proxy-core/src/error.rs index ec545bf..b661560 100644 --- a/crates/llm-proxy-core/src/error.rs +++ b/crates/llm-proxy-core/src/error.rs @@ -23,24 +23,20 @@ pub enum CoreError { /// Config validation failed. /// /// The structured [`ConfigValidationError`] is carried typed (not boxed - /// behind `dyn Error`) so callers can match on the concrete variant - /// directly via [`CoreError::validation_error`] instead of downcasting - /// through a trait object (GAP-LOW-11). + /// behind `dyn Error`); match it via [`CoreError::validation_error`]. #[error("config validation error: {message}")] ConfigValidation { /// Human-readable description of the validation failure. message: String, - /// The structured validation error, when this failure originated from a - /// typed [`ConfigValidationError`]. `None` for ad-hoc message-only - /// validations such as the HIGH-2 request-timeout guard. + /// The structured validation error, when this failure originated from + /// a typed [`ConfigValidationError`]. `None` for message-only + /// validations. #[source] source: Option, }, /// Provider registry construction or protocol validation failed. /// - /// Carries only a message: every provider-resolution failure in this crate - /// is an ad-hoc diagnostic (duplicate provider name, unknown protocol), so - /// there is no typed source to erase (GAP-LOW-11). + /// Carries only a diagnostic message; there is no typed source. #[error("provider resolution error: {message}")] ProviderResolution { /// Human-readable description of the resolution failure. @@ -67,9 +63,6 @@ impl CoreError { /// Return the structured validation error when this is a typed /// [`CoreError::ConfigValidation`], otherwise `None`. /// - /// Callers inspect the concrete [`ConfigValidationError`] variant directly - /// — no `dyn Error` downcast required (GAP-LOW-11). - /// /// ``` /// use llm_proxy_core::{CoreError, provider_config::ConfigValidationError}; /// @@ -145,21 +138,15 @@ mod tests { fn config_validation_from_config_validation_error() { let validation_err = crate::provider_config::ConfigValidationError::EmptyProviderName; let core_err: CoreError = validation_err.into(); - // The typed variant is recoverable directly — no `dyn Error` downcast - // (GAP-LOW-11). assert_eq! on the structured variant replaces the brittle - // `to_string().contains(...)` substring match the old test used (GAP-LOW-12). assert_eq!( core_err.validation_error(), Some(&crate::provider_config::ConfigValidationError::EmptyProviderName), ); - // The source chain still surfaces the typed error for consumers that walk it. assert!(core_err.source().is_some()); } #[test] fn config_validation_preserves_fielded_variant() { - // A non-unit (fielded) variant must round-trip through CoreError so callers - // can match its fields without downcasting (GAP-LOW-11 / GAP-LOW-12). let validation_err = crate::provider_config::ConfigValidationError::EmptyApiKey { provider: "my-provider".to_owned(), }; diff --git a/crates/llm-proxy-core/src/lib.rs b/crates/llm-proxy-core/src/lib.rs index bdcb108..048429d 100644 --- a/crates/llm-proxy-core/src/lib.rs +++ b/crates/llm-proxy-core/src/lib.rs @@ -1,30 +1,6 @@ -//! Core types shared across all `llm-proxy` crates. -//! -//! Provides: -//! -//! - **Configuration**: the TOML [`AppConfig`]/[`ProviderConfig`] types with -//! env-var interpolation and validation. -//! - **Routing**: [`ProviderRouteKind`] and [`ProviderRoutesConfig`] for -//! provider-based route resolution. -//! - **Registry**: [`ProviderRegistry`] and [`ProviderAdapterTargetConfig`] for -//! resolving provider routes to concrete adapter configurations. -//! - **Metrics**: the runtime [`Metrics`] collector with counters and latency -//! histograms. -//! - **Token counting**: [`Counter`] and [`MessageContent`] for usage tracking. -//! - **PID management**: [`PidManager`] for daemon mode. -//! - **Errors**: the crate-level [`CoreError`] enum. -//! -//! # Crate layout -//! -//! ```text -//! provider_config - TOML provider/app config types -//! provider_registry - Provider registry and adapter target resolution -//! env_interpolate - Shared ${ENV_VAR} interpolation -//! token - Token counting -//! metrics - Runtime metrics -//! pid - PID file management -//! error - Crate-level errors -//! ``` +//! Core types shared across all `llm-proxy` crates: configuration, routing, +//! the provider registry, metrics, token counting, PID management, and +//! errors. #![deny(missing_docs)] @@ -56,10 +32,11 @@ pub use error::CoreError; pub use metrics::{Metrics, Snapshot}; pub use pid::{PidError, PidManager}; pub use provider_config::{ - AppConfig, AuthStyle, ConfigValidationError, LogFormat, ModelId, ModelPricing, - ProviderAdapterConfig, ProviderCatalogConfig, ProviderCatalogMode, ProviderConfig, - ProviderDiscoveryConfig, ProviderDiscoveryKind, ProviderFile, ProviderRouteKind, - ProviderRoutesConfig, ServerConfig, StaticModelCatalogEntry, load_app_config, + ANONYMOUS_CALLER, AdminAuthConfig, AdminMountPosture, AppConfig, AuthConfig, AuthMode, + AuthStyle, ConfigValidationError, LogFormat, ModelId, ModelPricing, ProviderAdapterConfig, + ProviderCatalogConfig, ProviderCatalogMode, ProviderConfig, ProviderDiscoveryConfig, + ProviderDiscoveryKind, ProviderEgressConfig, ProviderFile, ProviderRouteKind, + ProviderRoutesConfig, ServerConfig, StaticModelCatalogEntry, StorageConfig, load_app_config, load_provider_config, validate_provider_config, }; pub use provider_registry::{ diff --git a/crates/llm-proxy-core/src/metrics.rs b/crates/llm-proxy-core/src/metrics.rs index 8904be0..af07c6b 100644 --- a/crates/llm-proxy-core/src/metrics.rs +++ b/crates/llm-proxy-core/src/metrics.rs @@ -3,56 +3,20 @@ //! All counters are lock-free ([`std::sync::atomic::AtomicI64`]). The only //! mutex-guarded state is the latency ring-buffer (last 1 000 samples) and //! the per-provider-model request counter map. -//! -//! # Ordering -//! -//! All atomic operations use [`std::sync::atomic::Ordering::Relaxed`] which is correct here -//! because the counters are monotonically increasing, used only for -//! observability, and no synchronization invariants depend on their -//! relative ordering. use std::collections::{HashMap, VecDeque}; use std::sync::Mutex; use std::sync::atomic::{AtomicI64, Ordering}; use std::time::Duration; -/// Maximum number of latency samples retained in the ring-buffer. const LATENCY_CAP: usize = 1000; -/// Maximum number of distinct `"{provider}:{model}"` keys retained in -/// [`Metrics::model_counts`]. -/// -/// The `model` axis is attacker-controlled (it arrives from the inbound client -/// request and, when catalog enforcement is off -- the default --, is not -/// canonicalized before keying). Without a cap, a client sending a unique model -/// string per request could grow this map without bound for the lifetime of the -/// process. Once this many distinct keys exist, any further new pair is folded -/// into the [`MODEL_COUNTS_OVERFLOW`] aggregate bucket instead of allocating a -/// new entry, bounding memory while still reporting the total request count. -/// The hard bound on map size is therefore [`MODEL_COUNTS_CAP`] real keys plus -/// at most one lazily-created overflow bucket. const MODEL_COUNTS_CAP: usize = 1024; -/// Aggregate bucket key used once [`MODEL_COUNTS_CAP`] distinct keys have been -/// recorded. Chosen so it cannot collide with a real `"{provider}:{model}"` -/// key: the provider segment of a real key is URL-path-validated against -/// `^[a-z0-9_-]+$`, so a leading `~` (not in that class) makes this key -/// structurally unreachable from any request input. (A prior sentinel -/// `__other__:` was collidable because `_` is a legal provider-name character, -/// so a configured provider literally named `__other__` serving an empty model -/// would alias the overflow bucket.) const MODEL_COUNTS_OVERFLOW: &str = "~overflow:"; -/// Separator used in the composite metrics key: `"{provider}:{model}"`. -/// -/// The `:` delimiter is chosen because provider names are restricted to -/// `[a-z0-9_-]+` and model IDs cannot contain `:`, preventing key collisions. const KEY_SEPARATOR: &str = ":"; -// --------------------------------------------------------------------------- -// Metrics -// --------------------------------------------------------------------------- - /// Thread-safe metrics collector. /// /// Internally uses [`AtomicI64`] counters for hot-path increments and @@ -110,12 +74,6 @@ impl Metrics { self.requests_success.fetch_add(1, Ordering::Relaxed); self.upstream_calls.fetch_add(1, Ordering::Relaxed); - // Store latency sample (ring-buffer), then release the latencies lock - // before touching model_counts so the two independent maps are not - // serialized by a single held guard. Recover from poison to preserve - // data from panicked threads. Renaming the inner binding from `buf` to - // `latencies` avoids shadowing the thread-local `buf` in the closure - // below (audit LOW-9). { let mut latencies = self.latencies.lock().unwrap_or_else(|e| e.into_inner()); if latencies.len() >= LATENCY_CAP { @@ -124,13 +82,6 @@ impl Metrics { latencies.push_back(latency); } - // Bump per-provider-model counter. - // - // The composite key is built into a thread-local reusable buffer so no - // fresh `String` is heap-allocated per call. On the hot path (provider/ - // model pair already tracked) the buffer is only borrowed for a `get` - // lookup; the cold insert path clones the buffer contents into an owned - // key exactly once per distinct pair (audit LOW-7). thread_local! { static KEY_BUF: std::cell::RefCell = std::cell::RefCell::new(String::with_capacity(64)); @@ -145,14 +96,8 @@ impl Metrics { if let Some(count) = map.get_mut(buf.as_str()) { *count += 1; } else if map.len() < MODEL_COUNTS_CAP { - // Cold path: first sighting of this pair — pay for one owned key. *map.entry(buf.clone()).or_insert(0) += 1; } else { - // The map is at capacity (defending against attacker-controlled, - // unbounded model strings -- see [`MODEL_COUNTS_CAP`]). Fold this - // and any further distinct pairs into the aggregate overflow - // bucket instead of allocating a new key. The bucket itself is - // lazily created the first time it is needed. *map.entry(MODEL_COUNTS_OVERFLOW.to_owned()).or_insert(0) += 1; } }); @@ -160,13 +105,10 @@ impl Metrics { /// Record a failed request. /// - /// Increments `requests_failed` and `upstream_calls`. - /// - /// Reserve this for genuine upstream/decode failures. A client-initiated - /// disconnect (the client went away mid-stream, or the timeout layer tore - /// the connection down) is **not** a failure of the upstream call and must - /// be recorded via [`Metrics::record_client_cancel`] instead, so that the - /// error-rate SLO is not corrupted by users pressing "stop" (audit MEDIUM-5). + /// Increments `requests_failed` and `upstream_calls`. Reserve this for + /// genuine upstream/decode failures: record client-initiated disconnects + /// via [`Metrics::record_client_cancel`] instead, so that "stop" presses + /// and timeout teardowns do not corrupt the error-rate SLO. pub fn record_failure(&self) { self.requests_failed.fetch_add(1, Ordering::Relaxed); self.upstream_calls.fetch_add(1, Ordering::Relaxed); @@ -175,8 +117,7 @@ impl Metrics { /// Record a client-initiated cancellation. /// /// Distinct from [`Metrics::record_failure`]: this does not touch - /// `requests_failed` or `upstream_calls`, so client "stop" presses and - /// timeout teardowns do not inflate the error-rate SLO (audit MEDIUM-5). + /// `requests_failed` or `upstream_calls`. pub fn record_client_cancel(&self) { self.client_cancelled.fetch_add(1, Ordering::Relaxed); } @@ -193,10 +134,9 @@ impl Metrics { /// Take an approximate point-in-time snapshot of all counters. /// - /// Individual fields are each loaded atomically, but the overall - /// snapshot is **approximately** consistent rather than transactional: - /// between loading the first counter and the last, other threads may - /// have mutated state. This is acceptable for observability use cases. + /// Individual fields are each loaded atomically, but the snapshot is not + /// transactional: other threads may mutate state between the first and + /// last load. pub fn get_snapshot(&self) -> Snapshot { let requests_received = self.requests_received.load(Ordering::Relaxed); let requests_streamed = self.requests_streamed.load(Ordering::Relaxed); @@ -244,14 +184,10 @@ impl Default for Metrics { } } -// --------------------------------------------------------------------------- -// Snapshot -// --------------------------------------------------------------------------- - /// An approximate point-in-time copy of all metric counters. /// -/// Individual fields are each loaded atomically but the snapshot as a whole -/// is not a transactional view. This is acceptable for observability. +/// Each field is loaded atomically, but the snapshot as a whole is not a +/// transactional view. #[derive(Debug, Clone, Default, PartialEq)] pub struct Snapshot { /// Total requests received (streaming + non-streaming). @@ -269,8 +205,7 @@ pub struct Snapshot { /// Requests that were deduplicated before reaching upstream. pub deduplicated: i64, /// Requests whose client disconnected (or were torn down by the timeout - /// layer) mid-stream. Tracked separately from `requests_failed` so client - /// "stop" actions do not corrupt the error-rate SLO (audit MEDIUM-5). + /// layer) mid-stream; tracked separately from `requests_failed`. pub client_cancelled: i64, /// Collected latency samples (up to 1 000 entries). pub latencies: Vec, @@ -304,21 +239,7 @@ impl Snapshot { } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Return the latency at the given percentile (0..=100). -/// -/// Uses the "exclusive" interpolation method common in monitoring systems: -/// `rank = pct/100 * (n + 1)`, then clamped to `[0, n-1]`. -/// -/// Returns [`Duration::ZERO`] for an empty slice. fn percentile(samples: &[Duration], pct: f64) -> Duration { - // Clamp the input unconditionally so an out-of-range `pct` (e.g. a caller - // bug passing 150.0 or a negative) cannot silently yield the wrong sample. - // This clamp is the hard guard that survives in both debug and release - // builds; `calculate_percentile_clamps_out_of_range_input` exercises it. let pct = pct.clamp(0.0, 100.0); if samples.is_empty() { @@ -329,9 +250,7 @@ fn percentile(samples: &[Duration], pct: f64) -> Duration { sorted.sort(); let n = sorted.len(); - // Nearest-rank (exclusive) method: rank = pct/100 * (n + 1), 1-based. let rank = (pct / 100.0) * (n as f64 + 1.0); - // Convert to 0-based index, clamped to valid range. #[expect( clippy::cast_sign_loss, reason = "`pct` is clamped to 0.0..=100.0 and `n >= 1`, so `rank = pct/100 * (n+1)` is always >= 0.0 and the cast cannot lose a sign" @@ -340,10 +259,6 @@ fn percentile(samples: &[Duration], pct: f64) -> Duration { sorted[idx] } -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - #[cfg(test)] mod tests { use super::*; @@ -456,8 +371,6 @@ mod tests { let snap = m.get_snapshot(); assert_eq!(snap.latencies.len(), LATENCY_CAP); - // The oldest 50 samples should have been evicted; the buffer should - // start at sample index 50. assert_eq!(snap.latencies[0], Duration::from_millis(50)); assert_eq!(snap.latencies[LATENCY_CAP - 1], Duration::from_millis(1049)); } @@ -471,7 +384,6 @@ mod tests { #[test] fn p95_p99_on_known_samples() { - // 100 samples from 0ms..99ms. let latencies: Vec = (0..100).map(Duration::from_millis).collect(); let snap = Snapshot { @@ -479,9 +391,7 @@ mod tests { ..Snapshot::default() }; - // p95: rank = 0.95 * (100+1) = 95.95 -> floor 95 -> idx 94 -> 94ms assert_eq!(snap.calculate_p95(), Duration::from_millis(94)); - // p99: rank = 0.99 * (100+1) = 99.99 -> floor 99 -> idx 98 -> 98ms assert_eq!(snap.calculate_p99(), Duration::from_millis(98)); } @@ -534,7 +444,6 @@ mod tests { } let snap = m.get_snapshot(); - // 4 threads x 500 iterations = 2000 requests received. assert_eq!(snap.requests_received, 2000); assert_eq!(snap.requests_success + snap.requests_failed, 2000); } @@ -577,9 +486,7 @@ mod tests { latencies: vec![Duration::from_millis(10), Duration::from_millis(100)], ..Snapshot::default() }; - // p95: rank = 0.95 * 3 = 2.85 -> floor 2 -> idx 1 -> 100ms assert_eq!(snap.calculate_p95(), Duration::from_millis(100)); - // p99: rank = 0.99 * 3 = 2.97 -> floor 2 -> idx 1 -> 100ms assert_eq!(snap.calculate_p99(), Duration::from_millis(100)); } @@ -594,16 +501,12 @@ mod tests { #[test] fn calculate_percentile_clamps_out_of_range_input() { - // Out-of-range `pct` must be clamped to [0, 100] rather than silently - // returning the wrong sample (the debug_assert compiles out in release). let snap = Snapshot { latencies: (0..100).map(Duration::from_millis).collect(), ..Snapshot::default() }; - // pct > 100 clamps to 100 -> max sample (99ms). assert_eq!(snap.calculate_percentile(150.0), Duration::from_millis(99)); assert_eq!(snap.calculate_percentile(100.0), Duration::from_millis(99)); - // pct < 0 clamps to 0 -> min sample (0ms). assert_eq!(snap.calculate_percentile(-42.0), Duration::from_millis(0)); assert_eq!(snap.calculate_percentile(0.0), Duration::from_millis(0)); } @@ -642,8 +545,6 @@ mod tests { #[test] fn key_separator_does_not_collide_with_provider_names() { - // Provider names are restricted to [a-z0-9_-]+ so `:` separator - // cannot appear in provider or model names. let m = Metrics::new(); m.record_success("provider-a", "model-x", Duration::from_millis(10)); m.record_success("provider-a", "model-y", Duration::from_millis(20)); @@ -655,14 +556,8 @@ mod tests { #[test] fn model_counts_capped_with_overflow_bucket() { - // The model axis is attacker-controlled, so the distinct-key map must be - // bounded. Once MODEL_COUNTS_CAP distinct pairs have been recorded, any - // further new pair is folded into the `~overflow:` aggregate bucket - // instead of allocating a new entry. The bound is MODEL_COUNTS_CAP real - // keys plus at most one lazily-created overflow bucket. let m = Metrics::new(); - // Fill the map to capacity with distinct pairs. for i in 0..MODEL_COUNTS_CAP { m.record_success("p", &format!("model-{i}"), Duration::from_micros(i as u64)); } @@ -677,8 +572,6 @@ mod tests { "overflow bucket must not exist before capacity is reached" ); - // Further distinct pairs spill into the aggregate bucket. The bucket is - // created on first spill, so the map grows by exactly one entry. m.record_success("p", "overflow-1", Duration::from_micros(0)); m.record_success("p", "overflow-2", Duration::from_micros(0)); let snap = m.get_snapshot(); @@ -693,7 +586,6 @@ mod tests { "spilled counts should accumulate in the overflow bucket" ); - // More distinct pairs must NOT grow the map further. m.record_success("p", "overflow-3", Duration::from_micros(0)); let snap = m.get_snapshot(); assert_eq!( @@ -703,8 +595,6 @@ mod tests { ); assert_eq!(snap.model_counts.get(MODEL_COUNTS_OVERFLOW), Some(&3)); - // An already-tracked pair must still increment its own entry (hot path), - // not the overflow bucket. m.record_success("p", "model-0", Duration::from_micros(0)); let snap = m.get_snapshot(); assert_eq!( diff --git a/crates/llm-proxy-core/src/pid.rs b/crates/llm-proxy-core/src/pid.rs index 918ba25..69c1335 100644 --- a/crates/llm-proxy-core/src/pid.rs +++ b/crates/llm-proxy-core/src/pid.rs @@ -1,8 +1,7 @@ //! PID file management for daemon mode. //! -//! Provides [`PidManager`] which handles creating, reading, and removing PID -//! files used to track a running daemon process. The config directory is -//! configurable so that tests can point at a temporary directory. +//! [`PidManager`] handles creating, reading, and removing the PID file used +//! to track a running daemon process. use std::io::Write; use std::num::ParseIntError; @@ -10,10 +9,8 @@ use std::path::{Path, PathBuf}; /// Errors returned by [`PidManager`] operations. /// -/// Typed (rather than `anyhow::Error`) so that callers of this public library -/// API can inspect or branch on the failure mode. Each variant carries the -/// filesystem path involved so the error stays actionable without the caller -/// having to supply that context (GAP-LOW-14). +/// Typed so callers can branch on the failure mode; each variant carries the +/// filesystem path involved. #[derive(Debug, thiserror::Error)] pub enum PidError { /// A filesystem I/O error occurred while reading, writing, or removing a @@ -22,8 +19,7 @@ pub enum PidError { Io { /// The filesystem path involved in the failed operation. path: PathBuf, - /// The underlying I/O error. (Named `source` so thiserror uses it as the - /// [`std::error::Error::source`] chain root.) + /// The underlying I/O error. source: std::io::Error, }, /// The PID file contents could not be parsed as a `u32`. @@ -99,7 +95,6 @@ impl PidManager { })?; let pid = std::process::id(); - // Write to a temporary file first, then rename atomically. let tmp_path = self.pid_file.with_extension("pid.tmp"); { #[cfg(unix)] @@ -177,7 +172,6 @@ impl PidManager { true } else { let err = std::io::Error::last_os_error(); - // EPERM means the process exists but we lack permission. err.raw_os_error() != Some(libc::ESRCH) } } @@ -189,30 +183,21 @@ impl PidManager { } } -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - #[cfg(test)] mod tests { use super::{PidError, PidManager}; - /// TestWritePIDAndGetPID_RoundTrip: writing a PID and reading it back - /// should return the same value. #[test] fn write_and_read_pid_roundtrip() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); - // Write the PID of the current test process. mgr.write_pid().expect("write_pid"); let pid = mgr.read_pid().expect("read_pid").expect("pid should exist"); assert_eq!(pid, std::process::id()); } - /// TestGetPID_MissingFile: reading from a non-existent PID file should - /// return `Ok(None)`. #[test] fn read_pid_missing_file() { let dir = tempfile::tempdir().expect("create temp dir"); @@ -222,15 +207,11 @@ mod tests { assert!(result.is_none(), "expected None for missing PID file"); } - /// TestGetPID_InvalidContent: reading a PID file with non-numeric content - /// should return a typed [`PidError::Parse`] carrying the offending path - /// (GAP-LOW-14: the error is inspectable, not erased into `anyhow::Error`). #[test] fn read_pid_invalid_content() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); - // Write garbage to the PID file. std::fs::write(mgr.pid_file(), "not-a-number").expect("write garbage"); match mgr.read_pid() { @@ -240,22 +221,17 @@ mod tests { *mgr.pid_file(), "typed error should carry the PID file path" ); - let _ = source; // ParseIntError present; not asserting its exact wording + let _ = source; } other => panic!("expected PidError::Parse, got {other:?}"), } } - /// TestReadPID_IoErrorVariant: a read failure other than "not found" should - /// surface as a typed [`PidError::Io`] with the path, so callers can branch - /// on the failure mode (GAP-LOW-14). #[test] fn read_pid_io_error_is_typed() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); - // Replace the expected PID file path with a directory, so reading it as - // a file yields an I/O error (not NotFound). std::fs::create_dir(mgr.pid_file()).expect("create dir at pid path"); match mgr.read_pid() { @@ -270,8 +246,6 @@ mod tests { } } - /// TestIsProcessRunning_CurrentProcess: the current process PID should be - /// reported as running. #[test] fn is_process_running_current_process() { let pid = std::process::id(); @@ -281,11 +255,8 @@ mod tests { ); } - /// TestIsProcessRunning_NonexistentPID: a very large PID that almost - /// certainly does not exist should be reported as not running. #[test] fn is_process_running_nonexistent_pid() { - // PID 299999999 is extremely unlikely to exist on any system. let pid = 299_999_999; assert!( !PidManager::is_process_running(pid), @@ -293,13 +264,11 @@ mod tests { ); } - /// TestRemovePID: removing the PID file should delete it from disk. #[test] fn remove_pid_deletes_file() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); - // Write then remove. mgr.write_pid().expect("write_pid"); assert!(mgr.pid_file().exists(), "PID file should exist after write"); mgr.remove_pid().expect("remove_pid"); @@ -309,26 +278,20 @@ mod tests { ); } - /// TestRemovePID_MissingFile: removing a non-existent PID file should - /// succeed silently. #[test] fn remove_pid_succeeds_when_missing() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); - // No PID file was created. remove_pid should succeed. mgr.remove_pid() .expect("remove_pid on missing file should succeed"); } - /// TestAtomicWrite: verify the PID file has content immediately after - /// write_pid (the temp-file + rename strategy ensures no partial reads). #[test] fn write_pid_is_atomic() { let dir = tempfile::tempdir().expect("create temp dir"); let mgr = PidManager::new(dir.path()); mgr.write_pid().expect("write_pid"); - // File should exist and contain the PID. let content = std::fs::read_to_string(mgr.pid_file()).expect("read PID file"); let pid: u32 = content.trim().parse().expect("parse PID"); assert_eq!(pid, std::process::id()); diff --git a/crates/llm-proxy-core/src/provider_config.rs b/crates/llm-proxy-core/src/provider_config.rs index e4a4994..14fe1ec 100644 --- a/crates/llm-proxy-core/src/provider_config.rs +++ b/crates/llm-proxy-core/src/provider_config.rs @@ -1,13 +1,10 @@ //! TOML provider configuration types, parsing, env-var interpolation, and -//! validation for the routing system. -//! -//! This is the sole configuration system for the proxy, loaded from the main -//! `config.toml` file. The main file contains server settings; provider files -//! contain provider-scoped routing and catalog configuration. +//! validation. The main `config.toml` holds server settings; provider files +//! hold provider-scoped routing and catalog configuration. use std::collections::HashMap; use std::net::SocketAddr; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; @@ -18,23 +15,61 @@ use serde::{Deserialize, Serialize}; use crate::env_interpolate::{find_env_var_refs, find_unresolved_env_var, interpolate_env_vars}; use crate::error::CoreError; -// --------------------------------------------------------------------------- -// Main config: AppConfig -// --------------------------------------------------------------------------- - /// Top-level TOML application configuration. /// -/// Loaded from the main `config.toml`, which contains server settings only. +/// Loaded from the main `config.toml`, which contains server and auth +/// settings. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct AppConfig { /// Server bind and operational settings. pub server: ServerConfig, + /// Client authentication settings (`[auth]` section). + #[serde(default)] + pub auth: AuthConfig, + /// Usage-ledger persistence settings (`[storage]` section). Absent keeps + /// the in-memory no-op event bus (nothing is persisted). + #[serde(default)] + pub storage: Option, +} + +/// Named security postures for where the admin `/usage` API is reachable. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminMountPosture { + /// Default: `/usage` is served only on the dedicated + /// [`ServerConfig::admin_bind`] listener (loopback unless configured + /// otherwise); the main listener has no admin surface at all. + DedicatedListener, + /// Opt-in (`server.admin_on_main_listener = true`): `/usage` is also + /// served on the main `bind` listener behind the same admin token, so + /// the admin surface reaches every network the main listener does. + MainListener, +} + +impl AdminMountPosture { + /// Stable name for this posture, used in startup logs. + #[must_use] + pub const fn name(self) -> &'static str { + match self { + Self::DedicatedListener => "dedicated-admin-listener", + Self::MainListener => "admin-on-main-listener", + } + } } -// --------------------------------------------------------------------------- -// ServerConfig -// --------------------------------------------------------------------------- +impl AppConfig { + /// The admin `/usage` mount posture, or `None` when no + /// `[auth.admin]` token is configured (no admin surface anywhere). + #[must_use] + pub fn admin_posture(&self) -> Option { + let posture = if self.server.admin_on_main_listener { + AdminMountPosture::MainListener + } else { + AdminMountPosture::DedicatedListener + }; + self.auth.admin.as_ref().map(|_| posture) + } +} /// Log output format for the tracing subscriber. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] @@ -51,7 +86,7 @@ impl LogFormat { /// Read the format from the `RUST_LOG_FORMAT` env var (case-insensitive). /// /// Defaults to [`LogFormat::Plain`] when unset or unrecognized. Read at - /// subscriber init (before config load), consistent with `RUST_LOG`. + /// subscriber init, before config load. pub fn from_env() -> Self { match std::env::var("RUST_LOG_FORMAT") .unwrap_or_default() @@ -76,20 +111,16 @@ pub struct ServerConfig { /// Hard deadline for graceful shutdown. /// /// After receiving SIGINT/SIGTERM the server stops accepting new - /// connections and waits this long for in-flight requests to drain. - /// Because the proxy serves long-lived streaming responses that an - /// upstream can hold open arbitrarily, an unbounded drain would let a - /// single stuck stream pin the process indefinitely. If the deadline - /// elapses the process force-exits so orchestrators (Kubernetes/LB) - /// can recycle the pod. `0s` disables the hard deadline (drain - /// forever) but is strongly discouraged. + /// connections and waits this long for in-flight requests to drain, then + /// force-exits so orchestrators can recycle the process. `0s` disables + /// the hard deadline (drain forever) but is strongly discouraged. #[serde(default = "default_shutdown_timeout", with = "humantime_serde")] pub shutdown_timeout: Duration, /// Log level: `"trace"`, `"debug"`, `"info"`, `"warn"`, `"error"`. pub log_level: String, /// Log output format for the tracing subscriber: `"plain"` (default) or - /// `"json"`. Applied at subscriber init from the `RUST_LOG_FORMAT` env var - /// (see [`LogFormat::from_env`]); this field documents the intended default. + /// `"json"`. The effective format is read from `RUST_LOG_FORMAT` at + /// subscriber init (see [`LogFormat::from_env`]); this field documents intent. #[serde(default)] pub log_format: LogFormat, /// Reserved hot-reload flag. Configuration loading rejects `true`. @@ -116,31 +147,125 @@ pub struct ServerConfig { pub dedup_window: Duration, /// Public server name reported in version/health endpoints. pub server_name: String, - /// Cross-origin origins permitted to call the proxy from a browser - /// (audit LOW-12). + /// Cross-origin origins permitted to call the proxy from a browser. /// /// `None` (the default) installs **no** CORS layer, so browsers enforce a /// same-origin policy and this server-to-server proxy is not unexpectedly /// reachable from arbitrary web origins. Setting a list installs a /// restrictive `CorsLayer` that echoes `Access-Control-Allow-Origin` only - /// for these exact origins -- replacing a former `very_permissive` policy - /// that admitted every origin. + /// for these exact origins. /// /// Each entry should be an absolute origin (`scheme://host[:port]`), /// e.g. `"https://app.example.com"`. Entries that fail to parse as a /// header value are skipped (with a warning) when the router is built. #[serde(default)] pub allowed_origins: Option>, + /// Dedicated listener bind for the admin `/usage` API. Only read when + /// `[auth.admin]` is configured; the default is `127.0.0.1` on the port + /// after `bind`'s port (see [`ServerConfig::effective_admin_bind`]). + #[serde(default)] + pub admin_bind: Option, + /// Also serve `/usage` on the main listener behind the same admin token. + /// + /// Opt-in; selects the [`AdminMountPosture::MainListener`] posture, which + /// the serve path announces with a loud startup warning: the admin + /// surface then reaches every network the main listener does. With no + /// explicit `admin_bind` this posture serves `/usage` on the main + /// listener ONLY — no dedicated socket is attempted. + #[serde(default)] + pub admin_on_main_listener: bool, +} + +impl ServerConfig { + /// Whether the dedicated admin listener binds at all. + /// + /// An explicit `admin_bind` always binds one. With none set, the + /// default loopback listener binds unless `admin_on_main_listener` + /// took over. + #[must_use] + pub fn dedicated_admin_listener(&self) -> bool { + self.admin_bind.is_some() || !self.admin_on_main_listener + } + + /// Bind address for the dedicated admin listener. + /// + /// `admin_bind` when configured; otherwise the default: `127.0.0.1` on + /// the port after `bind`'s port. `None` when that default port cannot be + /// derived (`bind` on port 65535) — `load_app_config` rejects that + /// combination whenever the dedicated listener would bind the default. + #[must_use] + pub fn effective_admin_bind(&self) -> Option { + match self.admin_bind { + Some(addr) => Some(addr), + None => self + .bind + .port() + .checked_add(1) + .map(|port| SocketAddr::new(std::net::IpAddr::from([127, 0, 0, 1]), port)), + } + } } const fn default_rate_limit_rpm() -> u32 { 100 } -/// Default graceful-shutdown hard deadline (30 seconds). +/// Client authentication mode for proxy routes. +/// +/// Serialized in TOML as `"none"` / `"client-keys"`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum AuthMode { + /// No client authentication: every request is attributed to the + /// `"anonymous"` caller. Only permitted while the server binds loopback; + /// `load_app_config` rejects a non-loopback bind in this mode. + #[default] + None, + /// Requests to proxy routes must present + /// `Authorization: Bearer ` matching a secret hashed under + /// `[auth.clients]`. + ClientKeys, +} + +/// Client authentication configuration (`[auth]` section of `config.toml`). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AuthConfig { + /// Authentication mode. Defaults to [`AuthMode::None`]. + #[serde(default)] + pub mode: AuthMode, + /// Key id to hashed client secret, under `[auth.clients]`. + /// + /// Each value is a PHC-format password-hashing string (argon2id), never a + /// plaintext secret. Wrapped in [`SecretString`] so the value is redacted + /// in `Debug` output. `skip_serializing` keeps the hashes out of every + /// `Serialize` path, so serializing then deserializing a config fails on + /// this field — by design. + #[serde(default, skip_serializing)] + pub clients: HashMap, + /// Admin API authentication (`[auth.admin]` section). When present, the + /// read-only `/usage` query surface is enabled and every `/usage` route + /// on every mount requires its bearer token. + #[serde(default)] + pub admin: Option, +} + +/// Admin API configuration (`[auth.admin]` section of `config.toml`). /// -/// Chosen to comfortably exceed the maximum upstream request timeout while -/// bounding how long a stuck streaming connection can delay process recycle. +/// A single bearer token guarding every `/usage` route. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AdminAuthConfig { + /// The admin bearer token, presented as `Authorization: Bearer `. + /// + /// Compared in constant time against a SHA-256 digest of the configured + /// token, so the plaintext never needs to be retained in memory. Wrapped + /// in [`SecretString`] and `skip_serializing`, keeping it out of `Debug` + /// output and every `Serialize` path. + #[serde(skip_serializing)] + pub token: SecretString, +} + const fn default_shutdown_timeout() -> Duration { Duration::from_secs(30) } @@ -149,9 +274,72 @@ const fn default_dedup_window() -> Duration { Duration::ZERO } -// --------------------------------------------------------------------------- -// ProviderFile -// --------------------------------------------------------------------------- +const fn default_prune_interval() -> Duration { + Duration::from_secs(3600) +} + +/// Usage-ledger persistence configuration (`[storage]` section of +/// `config.toml`). +/// +/// When present, every proxied request is persisted to a SQLite ledger by the +/// storage crate's event bus. When absent, the in-memory no-op bus is used and +/// nothing is written to disk. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StorageConfig { + /// Path to the SQLite database file. Parent directories are created on + /// startup. + pub path: PathBuf, + /// Delete ledger rows older than this age. Absent means rows are kept + /// forever. Must be positive when present. + #[serde(default, with = "humantime_serde")] + pub max_age: Option, + /// Keep at most this many ledger rows; the oldest rows are deleted when + /// the cap is exceeded. Absent means unlimited. Must be at least 1 when + /// present. + #[serde(default)] + pub max_rows: Option, + /// How often the retention pruning pass runs. Defaults to 1 hour. Must be + /// at least 1 second (the pruning timer rejects a zero period). + #[serde(default = "default_prune_interval", with = "humantime_serde")] + pub prune_interval: Duration, +} + +fn validate_storage_config(storage: &StorageConfig) -> Result<(), CoreError> { + if storage.path.as_os_str().is_empty() { + return Err(CoreError::ConfigValidation { + message: "storage.path must not be empty".to_owned(), + source: None, + }); + } + if storage.max_age == Some(Duration::ZERO) { + return Err(CoreError::ConfigValidation { + message: "storage.max_age must be a positive duration; a zero age would \ + delete every ledger row on every pruning pass" + .to_owned(), + source: None, + }); + } + if storage.max_rows == Some(0) { + return Err(CoreError::ConfigValidation { + message: "storage.max_rows must be at least 1; a cap of zero would keep \ + no rows at all" + .to_owned(), + source: None, + }); + } + const MIN_PRUNE_INTERVAL: Duration = Duration::from_secs(1); + if storage.prune_interval < MIN_PRUNE_INTERVAL { + return Err(CoreError::ConfigValidation { + message: format!( + "storage.prune_interval must be at least 1 second, but is {:?}", + storage.prune_interval + ), + source: None, + }); + } + Ok(()) +} /// Wrapper for a single provider TOML file. /// @@ -164,10 +352,6 @@ pub struct ProviderFile { pub provider: ProviderConfig, } -// --------------------------------------------------------------------------- -// ProviderConfig (custom Debug to redact api_key) -// --------------------------------------------------------------------------- - /// Provider configuration loaded from a dedicated TOML file. /// /// The `api_key` field is redacted in [`Debug`] output and excluded from @@ -183,15 +367,10 @@ pub struct ProviderConfig { /// this field is constructed, so the resolved value is what gets wrapped). /// /// Wrapped in [`SecretString`]: redacted in [`Debug`] and [`Display`] output - /// and zeroized on drop. Not serialized (`skip_serializing`) -- use explicit - /// methods if you need to write a config file. - /// - /// WHY `skip_serializing`: prevents credentials from leaking through - /// serialization paths (e.g. debug logs, API responses, config dumps). - /// Round-tripping a `ProviderFile` through serialize/deserialize will - /// FAIL because `skip_serializing` omits `api_key` and `api_key` is a - /// required field. This is intentional -- credentials should never - /// round-trip through serialization. + /// and zeroized on drop. Not serialized (`skip_serializing`): credentials + /// never leak through serialization paths, so round-tripping a + /// `ProviderFile` through serialize/deserialize FAILS because `api_key` + /// is a required field. This is intentional. #[serde(skip_serializing)] pub api_key: SecretString, /// Authentication header style. @@ -224,17 +403,54 @@ pub struct ProviderConfig { /// the key is the upstream model id. Defaults to empty (no cost accounting). #[serde(default)] pub pricing: HashMap, + /// Upstream egress (network exit) configuration. When present, this + /// provider's upstream requests are sent through the configured forward + /// proxy instead of a direct connection; absent means direct egress. + #[serde(default)] + pub egress: Option, +} + +/// Per-provider forward-proxy egress configuration (`[provider.egress]`). +/// +/// The proxy URL scheme must be one of `socks5h://`, `socks5://`, `http://`, +/// or `https://`; anything else is rejected at config load, never silently +/// ignored. `socks5h://` resolves upstream hostnames at the proxy (the exit +/// point's DNS view); `socks5://` resolves them locally. +/// +/// The URL may carry userinfo (`http://user:pass@host:port`), which is passed +/// to the transport as proxy authentication. [`Debug`] redacts the userinfo; +/// only `scheme://host:port` is ever printed. +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ProviderEgressConfig { + /// Forward proxy URL (`socks5h://host:port`, `socks5://`, `http://`, or + /// `https://`). Supports `${ENV_VAR}` interpolation like other string + /// fields. The port is optional and defaults per scheme (1080 for socks). + pub proxy: String, +} + +impl std::fmt::Debug for ProviderEgressConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ProviderEgressConfig") + .field("proxy", &redacted_proxy_url(&self.proxy)) + .finish() + } } -// --------------------------------------------------------------------------- -// ModelId / ModelPricing -// --------------------------------------------------------------------------- +fn redacted_proxy_url(url: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return "".to_owned(); + }; + let authority = rest.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + let after_authority = rest.strip_prefix(authority).unwrap_or_default(); + format!("{scheme}://{host}{after_authority}") +} /// Provider-local model identifier used as a pricing key. /// /// Transparent newtype over `String` so it serializes as a plain string in -/// TOML/JSON. Intentionally minimal: existing `String` model fields elsewhere -/// are not migrated to this type. +/// TOML/JSON. #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(transparent)] pub struct ModelId( @@ -281,24 +497,10 @@ pub struct ModelPricing { pub reasoning: Decimal, } -/// Upper bound for any per-token USD price accepted by config validation. -/// -/// Every real per-token price is far below $1; values above $1000/token are -/// treated as gross misconfigurations (stray exponents, per-1M-token prices -/// entered verbatim, misplaced decimals) and rejected at load time. See -/// [`validate_provider_config`]. `compute_cost` uses checked arithmetic on top -/// of this, so the proxy can never panic on overflow regardless. -/// -/// Returned by a function rather than a `const` because `Decimal` has no -/// const constructor at this rust_decimal version (`Decimal::new` is not const). fn max_per_token_price() -> Decimal { Decimal::new(1_000, 0) } -// --------------------------------------------------------------------------- -// AuthStyle -// --------------------------------------------------------------------------- - /// How the API key is sent in HTTP headers. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] @@ -313,16 +515,6 @@ pub enum AuthStyle { Both, } -// --------------------------------------------------------------------------- -// ProviderAdapterConfig -// --------------------------------------------------------------------------- - -/// Serialize an `Arc>` by value (LOW-5). -/// -/// serde only implements `Serialize`/`Deserialize` for `Arc` behind its `rc` -/// feature, which this crate deliberately leaves OFF (it is a footgun: it -/// deserializes `Rc`/`Arc` by reference-sharing). These helpers give the shared -/// headers map ordinary by-value (de)serialization without enabling that feature. fn serialize_arc_headers( headers: &Arc>, serializer: S, @@ -333,7 +525,6 @@ where headers.as_ref().serialize(serializer) } -/// Deserialize a headers map into a fresh `Arc` (by value, not reference-shared). fn deserialize_arc_headers<'de, D>( deserializer: D, ) -> Result>, D::Error> @@ -360,11 +551,9 @@ pub struct ProviderAdapterConfig { pub endpoint: String, /// Static headers to include in upstream requests (e.g. anthropic-version). /// - /// Wrapped in [`Arc`] because adapter headers are fixed at config-load time - /// and shared across every request that resolves to this adapter; resolving - /// a route then bumps the refcount instead of deep-cloning the map (LOW-5). - /// Serde deserializes this by value (wrapping the parsed `HashMap` in - /// `Arc::new`) without the `rc` feature. + /// Wrapped in [`Arc`]: adapter headers are fixed at config-load time and + /// shared across every request that resolves to this adapter, so downstream + /// code that needs to mutate must materialize an owned copy. #[serde( default, serialize_with = "serialize_arc_headers", @@ -389,10 +578,6 @@ pub(crate) fn endpoint_without_query(endpoint: &str) -> &str { .map_or(endpoint, |(base, _query)| base) } -// --------------------------------------------------------------------------- -// ProviderRouteKind & ProviderRoutesConfig -// --------------------------------------------------------------------------- - /// The kind of client-facing route being requested. /// /// Each variant corresponds to an inbound endpoint that the proxy exposes. @@ -434,15 +619,7 @@ impl ProviderRoutesConfig { } } -// --------------------------------------------------------------------------- -// ProviderDiscoveryKind & ProviderDiscoveryConfig -// --------------------------------------------------------------------------- - /// Supported discovery backends for fetching upstream model lists. -/// -/// Provider-specific discovery kinds are only justified when authentication, -/// pagination, or response shape differs materially from the OpenAI-compatible -/// default. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum ProviderDiscoveryKind { @@ -524,10 +701,6 @@ impl std::fmt::Debug for ProviderDiscoveryConfig { } } -// --------------------------------------------------------------------------- -// ProviderCatalogMode & StaticModelCatalogEntry & ProviderCatalogConfig -// --------------------------------------------------------------------------- - /// How the model catalog is assembled for a provider. /// /// - `Static`: only use operator-defined `[[provider.catalog.models]]` entries. @@ -613,7 +786,7 @@ pub struct ProviderCatalogConfig { } fn default_cache_ttl() -> Duration { - Duration::from_secs(24 * 60 * 60) // 24 hours + Duration::from_secs(24 * 60 * 60) } fn default_allow_all() -> Vec { @@ -633,10 +806,6 @@ impl Default for ProviderCatalogConfig { } } -// --------------------------------------------------------------------------- -// Validation error -// --------------------------------------------------------------------------- - /// Errors produced during config validation. /// /// This enum is `#[non_exhaustive]` so validation can evolve without breaking @@ -738,6 +907,18 @@ pub enum ConfigValidationError { /// The endpoint URL that has an invalid scheme. endpoint: String, }, + /// The `[provider.egress]` proxy URL is not usable as a forward proxy. + #[error( + "provider \"{provider}\": [provider.egress] proxy URL is invalid ({reason}): \"{proxy}\"" + )] + InvalidEgressProxy { + /// Provider name. + provider: String, + /// The rejected proxy URL. + proxy: String, + /// What is wrong with it. + reason: &'static str, + }, /// An adapter header name or value contains CR or LF characters. #[error("provider \"{provider}\": adapter \"{adapter}\" header contains CR/LF characters")] HeaderContainsCrlf { @@ -829,10 +1010,6 @@ pub enum ConfigValidationError { }, } -// --------------------------------------------------------------------------- -// Validation helpers -// --------------------------------------------------------------------------- - /// Validate a single [`ProviderConfig`]. /// /// Checks: @@ -880,6 +1057,7 @@ pub enum ConfigValidationError { /// api_key: SecretString::from("secret"), /// auth_style: AuthStyle::Bearer, /// passthrough_auth: false, +/// egress: None, /// adapters: HashMap::from([( /// "chat".to_owned(), /// ProviderAdapterConfig { @@ -912,6 +1090,7 @@ pub enum ConfigValidationError { /// api_key: SecretString::from("secret"), /// auth_style: AuthStyle::Bearer, /// passthrough_auth: false, +/// egress: None, /// adapters: HashMap::new(), /// routes: ProviderRoutesConfig::default(), /// model_aliases: HashMap::new(), @@ -935,9 +1114,8 @@ pub fn validate_provider_config( validate_pricing(provider, name)?; validate_adapters(provider, name, known_protocols)?; validate_discovery(provider, name)?; + validate_egress(provider, name)?; - // Both the catalog and the model aliases reject unsafe model IDs, but only - // when an adapter interpolates `{model}` into a Gemini endpoint URL. let uses_gemini_model_template = provider.adapters.values().any(|adapter| { adapter.protocol == "gemini_generate_content" && adapter.endpoint.contains("{model}") }); @@ -948,7 +1126,6 @@ pub fn validate_provider_config( Ok(()) } -/// Provider name must be non-empty and a URL-safe slug of at most 64 chars. fn validate_name(name: &str) -> Result<(), ConfigValidationError> { if name.trim().is_empty() { return Err(ConfigValidationError::EmptyProviderName); @@ -965,10 +1142,6 @@ fn validate_name(name: &str) -> Result<(), ConfigValidationError> { Ok(()) } -/// Checks the configured API key. -/// -/// Skipped entirely for passthrough auth: the upstream credential comes from -/// the inbound client request, so the configured key may be empty or unset. fn validate_api_key(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { if provider.passthrough_auth { return Ok(()); @@ -987,8 +1160,6 @@ fn validate_api_key(provider: &ProviderConfig, name: &str) -> Result<(), ConfigV Ok(()) } -/// Pricing keys must be non-empty and every price must be non-negative and -/// within the per-token upper bound. fn validate_pricing(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { for (model_id, pricing) in &provider.pricing { if model_id.as_str().is_empty() { @@ -1012,13 +1183,6 @@ fn validate_pricing(provider: &ProviderConfig, name: &str) -> Result<(), ConfigV message: format!("negative {field} price"), }); } - // Upper bound: every real per-token USD price is well under $1, so - // any value above $1000/token is a gross misconfiguration (a stray - // exponent, a per-1M-token price entered verbatim, or a misplaced - // decimal). Rejecting it here fails fast at config-load instead of - // letting a bad entry produce nonsensical costs. `compute_cost` - // additionally uses checked arithmetic so even a value that slipped - // past this bound can never panic on overflow. let max_price = max_per_token_price(); if value > max_price { return Err(ConfigValidationError::InvalidPricing { @@ -1036,9 +1200,6 @@ fn validate_pricing(provider: &ProviderConfig, name: &str) -> Result<(), ConfigV Ok(()) } -/// Every adapter needs a name, a protocol, and an HTTP(S) endpoint, plus -/// header checks. Non-HTTP schemes are rejected to prevent SSRF via crafted -/// endpoints. fn validate_adapters( provider: &ProviderConfig, name: &str, @@ -1084,8 +1245,6 @@ fn validate_adapters( Ok(()) } -/// Discovery, when configured, needs an HTTP(S) endpoint, valid headers, and -/// non-zero limits. fn validate_discovery(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { let Some(discovery) = &provider.discovery else { return Ok(()); @@ -1120,8 +1279,6 @@ fn validate_discovery(provider: &ProviderConfig, name: &str) -> Result<(), Confi Ok(()) } -/// Catalog entries need a non-empty ID, and a Gemini-safe one when an adapter -/// interpolates the model into its endpoint URL. fn validate_catalog( provider: &ProviderConfig, name: &str, @@ -1147,7 +1304,6 @@ fn validate_catalog( Ok(()) } -/// Each configured route must name a non-empty adapter that actually exists. fn validate_routes(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { for (route_kind, adapter_name) in [ ( @@ -1177,8 +1333,6 @@ fn validate_routes(provider: &ProviderConfig, name: &str) -> Result<(), ConfigVa Ok(()) } -/// Alias keys and targets must be non-empty, and targets Gemini-safe when an -/// adapter interpolates the model into its endpoint URL. fn validate_model_aliases( provider: &ProviderConfig, name: &str, @@ -1207,11 +1361,6 @@ fn validate_model_aliases( Ok(()) } -/// Forbidden HTTP header names that providers must not set via static headers. -/// -/// These headers are managed by the HTTP transport layer and must not be -/// overridden by provider configuration, as doing so could break request -/// routing, authentication, or connection handling. const FORBIDDEN_HEADERS: &[&str] = &[ "authorization", "cookie", @@ -1224,22 +1373,11 @@ const FORBIDDEN_HEADERS: &[&str] = &[ "transfer-encoding", ]; -/// Check whether an endpoint URL has a valid HTTP(S) scheme. -/// -/// Rejects non-HTTP schemes (e.g. `file://`, `ftp://`) to prevent SSRF -/// attacks via crafted endpoint URLs. fn has_valid_http_scheme(endpoint: &str) -> bool { - // Schemes are case-insensitive per RFC 3986 (e.g. `HTTP://`, `Https://`), - // and reqwest accepts them. Lowercase only the leading scheme region - // (everything up to the first `://`, capped at the known scheme length) so - // the prefix test is case-insensitive without allocating a full lowercase - // copy of the endpoint or pulling the `url` crate into llm-proxy-core. let scheme_region = scheme_region(endpoint); scheme_region.eq_ignore_ascii_case("http") || scheme_region.eq_ignore_ascii_case("https") } -/// Return the leading scheme region of `endpoint` (the text before the first -/// `://`, or the empty string when no `://` separator is present). fn scheme_region(endpoint: &str) -> &str { match endpoint.find("://") { Some(idx) => &endpoint[..idx], @@ -1247,10 +1385,54 @@ fn scheme_region(endpoint: &str) -> &str { } } -/// Validate a map of static headers for CRLF injection and forbidden header names. -/// -/// Shared between adapter header validation and discovery header validation -/// to avoid duplicating the same checks. +const EGRESS_PROXY_SCHEMES: [&str; 4] = ["socks5h", "socks5", "http", "https"]; + +fn validate_egress(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { + let Some(egress) = &provider.egress else { + return Ok(()); + }; + let url = egress.proxy.trim(); + if url.is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "empty proxy URL", + }); + } + let scheme = scheme_region(url); + if scheme.is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "no URL scheme; expected socks5h://, socks5://, http://, or https://", + }); + } + if !EGRESS_PROXY_SCHEMES + .iter() + .any(|accepted| scheme.eq_ignore_ascii_case(accepted)) + { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "unsupported scheme; expected socks5h://, socks5://, http://, or https://", + }); + } + let authority = url + .strip_prefix(scheme) + .and_then(|after_scheme| after_scheme.strip_prefix("://")) + .unwrap_or_default(); + let authority = authority.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + if host.trim().is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "no proxy host", + }); + } + Ok(()) +} + fn validate_headers( headers: &HashMap, provider: &str, @@ -1295,10 +1477,6 @@ fn is_safe_gemini_model_id(model: &str) -> bool { .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_')) } -// --------------------------------------------------------------------------- -// TOML loading -// --------------------------------------------------------------------------- - /// Load and parse a main [`AppConfig`] TOML file. /// /// The raw file content is first processed for `${ENV_VAR}` interpolation, @@ -1328,7 +1506,6 @@ pub fn load_app_config(path: impl AsRef) -> Result { source, })?; - // Check for env vars that resolve to empty, consistent with provider config. for (var_name, resolved) in find_env_var_refs(&raw) { if let Some(value) = resolved { if value.is_empty() { @@ -1344,10 +1521,6 @@ pub fn load_app_config(path: impl AsRef) -> Result { let interpolated = interpolate_env_vars(&raw); - // After interpolation, check for any unresolved env vars (set but not - // present in the environment). This ensures `${MISSING_VAR}` in server - // fields like `bind` or `server_name` is caught rather than silently - // passed through as a literal string. if let Some(var) = find_unresolved_env_var(&interpolated) { return Err(CoreError::ConfigValidation { message: format!("app config contains unresolvable environment variable \"{var}\""), @@ -1364,12 +1537,6 @@ pub fn load_app_config(path: impl AsRef) -> Result { source: None, }); } - // Reject a zero or sub-second request timeout. tower-http's `TimeoutLayer` - // builds `tokio::time::sleep(self.timeout)` on every call; `sleep(ZERO)` is - // immediately ready, so `request_timeout = "0s"` would return 408 for every - // API request before any handler runs. A sub-second timeout is almost - // always a misconfiguration (e.g. mistyped units) and silently bricks the - // API the same way. Require at least one second. (HIGH-2 / LOW-10.) const MIN_REQUEST_TIMEOUT: Duration = Duration::from_secs(1); if cfg.server.request_timeout < MIN_REQUEST_TIMEOUT { return Err(CoreError::ConfigValidation { @@ -1382,12 +1549,6 @@ pub fn load_app_config(path: impl AsRef) -> Result { source: None, }); } - // Reject a dedup window whose millisecond magnitude does not fit u64. The - // deduplicator stores the window as a `u64` millisecond count (see - // `AppState::new_with_catalog_dir` in llm-proxy-server), so an out-of-range - // `Duration` would otherwise silently fall back to the `unwrap_or` default - // rather than being reported. `Duration::MAX` is the only realistic value - // that trips this, and it indicates a misconfiguration. if u64::try_from(cfg.server.dedup_window.as_millis()).is_err() { return Err(CoreError::ConfigValidation { message: format!( @@ -1398,9 +1559,126 @@ pub fn load_app_config(path: impl AsRef) -> Result { source: None, }); } + validate_auth_config(&cfg.auth, cfg.server.bind)?; + if let Some(storage) = cfg.storage.as_ref() { + validate_storage_config(storage)?; + } + validate_admin_mount(&cfg)?; Ok(cfg) } +fn validate_admin_mount(cfg: &AppConfig) -> Result<(), CoreError> { + let Some(admin) = cfg.auth.admin.as_ref() else { + if cfg.server.admin_bind.is_some() { + return Err(CoreError::ConfigValidation { + message: "server.admin_bind is configured but [auth.admin] is not; \ + set [auth.admin] token or remove server.admin_bind" + .to_owned(), + source: None, + }); + } + if cfg.server.admin_on_main_listener { + return Err(CoreError::ConfigValidation { + message: "server.admin_on_main_listener = true but [auth.admin] is not \ + configured; there would be no admin token to guard /usage" + .to_owned(), + source: None, + }); + } + return Ok(()); + }; + if admin.token.expose_secret().is_empty() { + return Err(CoreError::ConfigValidation { + message: "[auth.admin] token must not be empty".to_owned(), + source: None, + }); + } + if cfg.storage.is_none() { + return Err(CoreError::ConfigValidation { + message: "[auth.admin] requires the [storage] usage ledger; configure \ + [storage] or remove [auth.admin]" + .to_owned(), + source: None, + }); + } + if let Some(admin_bind) = cfg.server.admin_bind { + if admin_bind == cfg.server.bind { + return Err(CoreError::ConfigValidation { + message: format!( + "server.admin_bind {admin_bind} must differ from server.bind; \ + the admin listener needs its own socket" + ), + source: None, + }); + } + } else if !cfg.server.admin_on_main_listener && cfg.server.bind.port() == u16::MAX { + return Err(CoreError::ConfigValidation { + message: "server.bind port 65535 leaves no distinct port for the default \ + admin listener; set server.admin_bind explicitly or enable \ + server.admin_on_main_listener" + .to_owned(), + source: None, + }); + } + Ok(()) +} + +/// Reserved caller name used for unauthenticated requests under +/// [`AuthMode::None`]; a client key id must not collide with it or event +/// attribution becomes ambiguous. +pub const ANONYMOUS_CALLER: &str = "anonymous"; + +fn validate_auth_config(auth: &AuthConfig, bind: SocketAddr) -> Result<(), CoreError> { + let clients_empty = auth.clients.is_empty(); + match auth.mode { + AuthMode::None if !clients_empty => { + return Err(CoreError::ConfigValidation { + message: "auth.mode = \"none\" but [auth.clients] is non-empty; \ + set auth.mode = \"client-keys\" to enforce the configured keys" + .to_owned(), + source: None, + }); + } + AuthMode::ClientKeys if clients_empty => { + return Err(CoreError::ConfigValidation { + message: "auth.mode = \"client-keys\" requires at least one entry \ + under [auth.clients]" + .to_owned(), + source: None, + }); + } + AuthMode::None if !bind.ip().is_loopback() => { + return Err(CoreError::ConfigValidation { + message: format!( + "server.bind {bind} is not loopback and no client keys are \ + configured; anonymous auth is only permitted on loopback binds. \ + Configure [auth.clients] and set auth.mode = \"client-keys\"" + ), + source: None, + }); + } + AuthMode::None | AuthMode::ClientKeys => {} + } + for key_id in auth.clients.keys() { + if key_id.is_empty() { + return Err(CoreError::ConfigValidation { + message: "[auth.clients] contains an empty key id".to_owned(), + source: None, + }); + } + if key_id == ANONYMOUS_CALLER { + return Err(CoreError::ConfigValidation { + message: format!( + "[auth.clients] key id \"{ANONYMOUS_CALLER}\" is reserved for \ + unattributed requests; choose a different key id" + ), + source: None, + }); + } + } + Ok(()) +} + /// Load and parse a provider [`ProviderFile`] TOML file. /// /// The raw file content is first processed for `${ENV_VAR}` interpolation, @@ -1436,26 +1714,10 @@ pub fn load_provider_config( source, })?; - // Before interpolation, check if any env vars referenced in the raw TOML - // resolve to empty strings. This must happen before interpolation replaces - // the ${VAR} pattern, otherwise we lose the diagnostic info about which - // env var was empty. - // - // WHY: We scan the raw TOML string with the shared env-var regex instead of - // parsing the TOML twice. This avoids a redundant full TOML parse just to - // extract the api_key value for the empty-env-var check. for (var_name, resolved) in find_env_var_refs(&raw) { match resolved { - None => { - // Env var is not set -- interpolation will leave ${VAR} as-is. - // The unresolved-env-var check below (after interpolation) will - // catch this for all fields, not just api_key. - } + None => {} Some(value) if value.is_empty() => { - // Provider name is not yet known (no parse has occurred), so we - // include the file path as context instead. Carried as a typed - // `CoreError::ConfigValidation` source via `From` (GAP-LOW-11), - // so this failure is variant-recoverable by callers. return Err(ConfigValidationError::EmptyEnvVar { provider: format!("(file: {})", path.display()), var: var_name, @@ -1466,11 +1728,6 @@ pub fn load_provider_config( } } - // Now perform full interpolation on the raw content and check for any - // unresolved env vars that remain. This catches ${MISSING_VAR} patterns - // in *any* field (endpoint, protocol, name), not just api_key. - // validate_provider_config only checks api_key for unresolved patterns, - // so a whole-file check is needed here to match load_app_config behavior. let interpolated = interpolate_env_vars(&raw); if let Some(var) = find_unresolved_env_var(&interpolated) { return Err(CoreError::ConfigValidation { @@ -1481,17 +1738,10 @@ pub fn load_provider_config( }); } let file: ProviderFile = toml::from_str(&interpolated).map_err(CoreError::ConfigParse)?; - // The typed `ConfigValidationError` flows into `CoreError::ConfigValidation` - // as a first-class typed source via `From` (GAP-LOW-11), not boxed behind - // `dyn Error`, so callers can match on the concrete variant. validate_provider_config(&file.provider, known_protocols).map_err(CoreError::from)?; Ok(file.provider) } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -1521,13 +1771,11 @@ server_name = "test" assert_eq!(config.server.rate_limit_rpm, 100); assert!(!config.server.trust_forwarded_headers); assert_eq!(config.server.dedup_window, Duration::ZERO); - // shutdown_timeout omits from the TOML above and must default to 30s (HIGH-1). assert_eq!(config.server.shutdown_timeout, Duration::from_secs(30)); } #[test] fn app_config_accepts_explicit_shutdown_timeout() { - // Operators may override the 30s default; the value round-trips verbatim. let config: AppConfig = toml::from_str( r#" [server] @@ -1571,7 +1819,6 @@ server_name = "test" #[test] fn load_app_config_rejects_zero_request_timeout() { - // HIGH-2: a zero timeout makes every API request 408 before any handler. let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("config.toml"); std::fs::write( @@ -1598,7 +1845,6 @@ server_name = "test" #[test] fn load_app_config_rejects_subsecond_request_timeout() { - // Sub-second timeouts are almost always a mistyped unit and brick the API. let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("config.toml"); std::fs::write( @@ -1625,9 +1871,6 @@ server_name = "test" #[test] fn load_app_config_rejects_dedup_window_exceeding_u64_millis() { - // A dedup window whose millisecond magnitude does not fit u64 (only - // realistically `Duration::MAX`) cannot be stored by the deduplicator - // and must be rejected at load time rather than silently clamped. let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("config.toml"); std::fs::write( @@ -1652,16 +1895,178 @@ server_name = "test" ); } + fn config_dir_with_auth(bind: &str, auth_body: &str) -> tempfile::TempDir { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "{bind}" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +{auth_body} +"# + ), + ) + .expect("write config file"); + directory + } + + #[test] + fn auth_absent_section_defaults_to_none_on_loopback() { + let dir = config_dir_with_auth("127.0.0.1:3456", ""); + let config = + load_app_config(dir.path().join("config.toml")).expect("default auth accepted"); + assert_eq!(config.auth.mode, AuthMode::None); + assert!(config.auth.clients.is_empty()); + } + + #[test] + fn auth_none_on_non_loopback_bind_is_rejected() { + for bind in ["0.0.0.0:3456", "192.168.1.4:3456", "[::]:3456"] { + let dir = config_dir_with_auth(bind, ""); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("non-loopback none must be rejected"); + assert!( + error.to_string().contains("not loopback"), + "bind {bind}: got: {error}" + ); + } + } + + #[test] + fn auth_none_on_ipv6_loopback_bind_is_accepted() { + let dir = config_dir_with_auth("[::1]:3456", ""); + load_app_config(dir.path().join("config.toml")).expect("ipv6 loopback accepted"); + } + + #[test] + fn auth_none_with_configured_clients_is_rejected() { + let dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "none" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("none mode with clients is contradictory"); + assert!( + error.to_string().contains("[auth.clients] is non-empty"), + "got: {error}" + ); + } + + #[test] + fn auth_client_keys_without_clients_is_rejected() { + let dir = config_dir_with_auth( + "0.0.0.0:3456", + r#" +[auth] +mode = "client-keys" +"#, + ); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("client-keys mode needs at least one key"); + assert!( + error + .to_string() + .contains("requires at least one entry under [auth.clients]"), + "got: {error}" + ); + } + + #[test] + fn auth_client_keys_with_a_client_accepts_any_bind() { + let dir = config_dir_with_auth( + "0.0.0.0:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let config = load_app_config(dir.path().join("config.toml")) + .expect("client-keys on non-loopback accepted"); + assert_eq!(config.auth.mode, AuthMode::ClientKeys); + assert_eq!(config.auth.clients.len(), 1); + } + + #[test] + fn auth_rejects_reserved_and_empty_key_ids() { + let reserved_dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +anonymous = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let reserved = load_app_config(reserved_dir.path().join("config.toml")) + .expect_err("reserved key id must be rejected"); + assert!(reserved.to_string().contains("reserved"), "got: {reserved}"); + + let empty_dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +"" = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let empty = load_app_config(empty_dir.path().join("config.toml")) + .expect_err("empty key id must be rejected"); + assert!(empty.to_string().contains("empty key id"), "got: {empty}"); + } + + #[test] + fn auth_clients_do_not_leak_through_debug() { + let config: AppConfig = toml::from_str( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hashvalue" +"#, + ) + .expect("app config"); + let debug = format!("{config:?}"); + assert!( + !debug.contains("hashvalue"), + "client hash must be redacted in Debug, got: {debug}" + ); + assert!(debug.contains("[REDACTED]"), "got: {debug}"); + } + #[test] fn has_valid_http_scheme_is_case_insensitive() { - // Schemes are case-insensitive per RFC 3986 and reqwest accepts mixed - // case; the validator must not reject `HTTP://` / `Https://`. assert!(has_valid_http_scheme("http://host")); assert!(has_valid_http_scheme("https://host")); assert!(has_valid_http_scheme("HTTP://host")); assert!(has_valid_http_scheme("Https://host")); assert!(has_valid_http_scheme("HtTp://host/path")); - // Non-HTTP schemes are still rejected. assert!(!has_valid_http_scheme("file://host")); assert!(!has_valid_http_scheme("ftp://host")); assert!(!has_valid_http_scheme("host")); @@ -1685,18 +2090,192 @@ chat_completion = "chat" } #[test] - fn provider_debug_redacts_header_values_and_endpoint_query() { - let config: ProviderFile = toml::from_str( + fn egress_section_parses_and_absent_section_means_none() { + let with_egress: ProviderFile = toml::from_str( r#" [provider] name = "example" -api_key = "secret-api-key" +api_key = "key" auth_style = "bearer" -[provider.adapters.chat] -protocol = "openai_chat_completions" -endpoint = "https://example.com/v1/chat/completions?key=adapter-query-secret" - +[provider.egress] +proxy = "socks5h://127.0.0.1:1080" +"#, + ) + .expect("provider with egress parses"); + assert_eq!( + with_egress.provider.egress, + Some(ProviderEgressConfig { + proxy: "socks5h://127.0.0.1:1080".to_owned(), + }) + ); + + let without_egress: ProviderFile = toml::from_str( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" +"#, + ) + .expect("provider without egress parses"); + assert_eq!(without_egress.provider.egress, None); + } + + #[test] + fn egress_section_rejects_unknown_fields() { + let result = toml::from_str::( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "socks5h://127.0.0.1:1080" +bind_out = "198.51.100.10" +"#, + ); + result.unwrap_err(); + } + + #[test] + fn egress_debug_redacts_proxy_userinfo_but_keeps_the_host() { + let egress = ProviderEgressConfig { + proxy: "http://proxy-user:proxy-pass@proxy.internal:3128".to_owned(), + }; + let debug = format!("{egress:?}"); + assert!( + !debug.contains("proxy-user") && !debug.contains("proxy-pass"), + "proxy credentials must never print: {debug}" + ); + assert!( + debug.contains("proxy.internal:3128"), + "the credential-free host must stay visible for diagnostics: {debug}" + ); + let raw = ProviderEgressConfig { + proxy: "user:pass@host".to_owned(), + }; + assert!( + !format!("{raw:?}").contains("user:pass"), + "unvalidated URLs must not be echoed: {:?}", + raw + ); + } + + #[test] + fn validate_accepts_each_supported_egress_proxy_scheme() { + for url in [ + "socks5h://tunnel.example.com:1080", + "socks5://127.0.0.1:1080", + "http://proxy.internal:3128", + "https://proxy.internal:8443", + "SOCKS5H://tunnel.example.com:1080", + "http://user:pass@proxy.internal:3128", + "socks5h://tunnel.example.com", + ] { + let file: ProviderFile = toml::from_str(&format!( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "{url}" +"# + )) + .expect("parses"); + assert_eq!( + validate_provider_config(&file.provider, None), + Ok(()), + "egress proxy {url} must validate" + ); + } + } + + #[test] + fn validate_rejects_unusable_egress_proxy_urls() { + for (url, reason_fragment) in [ + ("ftp://proxy:21", "unsupported scheme"), + ("socks4://127.0.0.1:1080", "unsupported scheme"), + ("127.0.0.1:1080", "no URL scheme"), + ("proxy.internal:3128", "no URL scheme"), + ("socks5h://", "no proxy host"), + ("http:///path", "no proxy host"), + ("", "empty proxy URL"), + (" ", "empty proxy URL"), + ] { + let file: ProviderFile = toml::from_str(&format!( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "{url}" +"# + )) + .expect("parses"); + let err = validate_provider_config(&file.provider, None) + .expect_err("egress proxy must be rejected"); + match &err { + ConfigValidationError::InvalidEgressProxy { + provider, reason, .. + } => { + assert_eq!(provider, "example"); + assert!( + reason.contains(reason_fragment), + "reason {reason:?} must name {reason_fragment:?} for {url:?}" + ); + } + other => panic!("expected InvalidEgressProxy for {url:?}, got {other:?}"), + } + let message = err.to_string(); + assert!(message.contains("example")); + assert!(message.contains("[provider.egress]")); + } + } + + #[test] + fn load_provider_config_rejects_invalid_egress_at_load_time() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("bad-egress.toml"); + std::fs::write( + &path, + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "ftp://proxy.example.com:21" +"#, + ) + .expect("write provider toml"); + + let err = load_provider_config(&path, None).expect_err("load must reject bad egress"); + assert!( + err.to_string().contains("egress"), + "error must point at the egress section: {err}" + ); + } + + #[test] + fn provider_debug_redacts_header_values_and_endpoint_query() { + let config: ProviderFile = toml::from_str( + r#" +[provider] +name = "example" +api_key = "secret-api-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions?key=adapter-query-secret" + [provider.adapters.chat.headers] x-secret = "adapter-secret" @@ -1722,9 +2301,6 @@ x-discovery-secret = "super-sensitive-value" #[test] fn env_interpolated_api_key_loads_into_secret_string() { - // Production secret-loading path: `${VAR}` in the raw TOML is resolved - // by env interpolation *before* serde constructs the `SecretString`, so - // the resolved value is what gets wrapped and `expose_secret()` returns. let _lock = crate::test_support::TestEnvLock::acquire(); let _guard = crate::test_support::EnvVarGuard::set("_LLM_PROXY_TEST_SECRET_KEY", "sk-resolved-007"); @@ -1914,8 +2490,6 @@ id = "models/gemini-2.5-pro" validate_provider_config(&provider.provider, None).unwrap(); } - // -- Adapter header CRLF validation --------------------------------------- - #[test] fn adapter_rejects_header_with_crlf() { let mut provider: ProviderFile = toml::from_str( @@ -1970,8 +2544,6 @@ endpoint = "https://example.com/v1/models" )); } - // -- Empty adapter name, protocol, model alias key/target ----------------- - #[test] fn validate_rejects_empty_adapter_name() { let provider: ProviderFile = toml::from_str( @@ -1987,7 +2559,6 @@ endpoint = "https://example.com/v1" "#, ) .expect("provider config"); - // The TOML key "" is deserialized as an empty string. assert!(matches!( validate_provider_config(&provider.provider, None), Err(ConfigValidationError::EmptyAdapterName { .. }) @@ -2001,6 +2572,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2028,6 +2600,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2055,6 +2628,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2078,10 +2652,11 @@ endpoint = "https://example.com/v1" #[test] fn validate_rejects_provider_name_exceeding_length_limit() { let provider = ProviderConfig { - name: "a".repeat(65), // 65 chars exceeds the 64-char limit + name: "a".repeat(65), api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2098,10 +2673,11 @@ endpoint = "https://example.com/v1" #[test] fn validate_accepts_provider_name_at_max_length() { let provider = ProviderConfig { - name: "a".repeat(64), // exactly 64 chars + name: "a".repeat(64), api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2109,14 +2685,9 @@ endpoint = "https://example.com/v1" catalog: None, pricing: Default::default(), }; - // Should not fail on name validation (will fail on empty api_key - // if key is empty, but our key is "key"). - // Actually no adapters/routes means it passes. validate_provider_config(&provider, None).unwrap(); } - // -- Pricing / cost ------------------------------------------------------ - #[test] fn model_pricing_parses_from_toml_with_quoted_decimals() { let raw = r#" @@ -2147,8 +2718,6 @@ reasoning = "0.000010" #[test] fn model_pricing_unquoted_float_is_rejected() { - // An unquoted decimal is parsed by TOML as f64; rust_decimal rejects - // floats, so deserialization fails. Operators MUST quote decimals. let raw = r#" [provider] name = "test" @@ -2167,6 +2736,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2193,6 +2763,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2208,15 +2779,12 @@ input = 0.0000025 #[test] fn validate_rejects_implausibly_large_pricing() { - // A per-token USD price above the MAX_PER_TOKEN_PRICE bound is a gross - // misconfiguration (a stray exponent, a per-1M-token price entered - // verbatim, or a misplaced decimal) and is rejected at load time rather - // than allowed to produce nonsensical costs. let provider = ProviderConfig { name: "test".to_owned(), api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2224,7 +2792,6 @@ input = 0.0000025 catalog: None, pricing: HashMap::from([( ModelId::new("gpt-4o"), - // $1000000/token: parses cleanly as a Decimal but exceeds the bound. ModelPricing { input: "1000000".parse().unwrap(), ..ModelPricing::default() @@ -2244,13 +2811,12 @@ input = 0.0000025 #[test] fn validate_accepts_realistic_pricing_below_bound() { - // A realistic per-token price (a few $/1M tokens -> tiny $/token) is - // accepted. Guards against the upper bound being too tight. let provider = ProviderConfig { name: "test".to_owned(), api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2271,17 +2837,14 @@ input = 0.0000025 #[test] fn log_format_from_env_reads_rust_log_format() { let _lock = crate::test_support::TestEnvLock::acquire(); - // "json" (case-insensitive) -> Json. { let _g = crate::test_support::EnvVarGuard::set("RUST_LOG_FORMAT", "JSON"); assert_eq!(LogFormat::from_env(), LogFormat::Json); } - // Unrecognized value -> Plain. { let _g = crate::test_support::EnvVarGuard::set("RUST_LOG_FORMAT", "xml"); assert_eq!(LogFormat::from_env(), LogFormat::Plain); } - // Empty string -> Plain (same as unset). { let _g = crate::test_support::EnvVarGuard::set("RUST_LOG_FORMAT", ""); assert_eq!(LogFormat::from_env(), LogFormat::Plain); @@ -2290,13 +2853,12 @@ input = 0.0000025 #[test] fn validate_accepts_empty_api_key_when_passthrough_auth() { - // A passthrough-auth provider gets its upstream credential from the - // inbound client request, so an empty/unset api_key must be allowed. let provider = ProviderConfig { name: "test".to_owned(), api_key: SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: true, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2312,12 +2874,12 @@ input = 0.0000025 #[test] fn validate_rejects_empty_api_key_when_not_passthrough() { - // Non-passthrough providers still require a non-empty api_key. let provider = ProviderConfig { name: "test".to_owned(), api_key: SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2330,4 +2892,369 @@ input = 0.0000025 Err(ConfigValidationError::EmptyApiKey { .. }) )); } + + fn write_app_config(extra_storage: &str) -> (tempfile::TempDir, std::path::PathBuf) { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +{extra_storage}"# + ), + ) + .unwrap(); + (directory, path) + } + + #[test] + fn storage_section_defaults_to_none_when_absent() { + let (_dir, path) = write_app_config(""); + let config = load_app_config(&path).expect("valid config without [storage]"); + assert!( + config.storage.is_none(), + "absent [storage] must keep the no-op event bus" + ); + } + + #[test] + fn storage_section_parses_and_applies_defaults() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +"#, + ); + let config = load_app_config(&path).expect("valid [storage]"); + let storage = config.storage.expect("storage present"); + assert_eq!(storage.path, std::path::PathBuf::from("/tmp/usage.db")); + assert_eq!(storage.max_age, None, "absent max_age keeps rows forever"); + assert_eq!(storage.max_rows, None, "absent max_rows is unlimited"); + assert_eq!( + storage.prune_interval, + Duration::from_secs(3600), + "absent prune_interval defaults to 1h" + ); + } + + #[test] + fn storage_section_parses_full_retention() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "30d" +max_rows = 500_000 +prune_interval = "15m" +"#, + ); + let config = load_app_config(&path).expect("valid [storage]"); + let storage = config.storage.expect("storage present"); + assert_eq!(storage.max_age, Some(Duration::from_secs(30 * 24 * 3600))); + assert_eq!(storage.max_rows, Some(500_000)); + assert_eq!(storage.prune_interval, Duration::from_secs(15 * 60)); + } + + #[test] + fn storage_zero_prune_interval_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +prune_interval = "0s" +"#, + ); + let error = load_app_config(&path).expect_err("zero prune interval rejected"); + assert!( + error.to_string().contains("storage.prune_interval"), + "got: {error}" + ); + } + + #[test] + fn storage_subsecond_prune_interval_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "1d" +prune_interval = "500ms" +"#, + ); + assert!( + load_app_config(&path).is_err(), + "sub-second prune interval must be rejected (timer panics on zero/near-zero)" + ); + } + + #[test] + fn storage_zero_max_age_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "0s" +"#, + ); + let error = load_app_config(&path).expect_err("zero max_age rejected"); + assert!( + error.to_string().contains("storage.max_age"), + "got: {error}" + ); + } + + #[test] + fn storage_zero_max_rows_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_rows = 0 +"#, + ); + let error = load_app_config(&path).expect_err("zero max_rows rejected"); + assert!( + error.to_string().contains("storage.max_rows"), + "got: {error}" + ); + } + + #[test] + fn storage_empty_path_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "" +"#, + ); + let error = load_app_config(&path).expect_err("empty path rejected"); + assert!(error.to_string().contains("storage.path"), "got: {error}"); + } + + #[test] + fn storage_unknown_field_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +db_host = "postgres.example" +"#, + ); + assert!( + load_app_config(&path).is_err(), + "deny_unknown_fields must reject unrecognized [storage] keys" + ); + } + + const STORAGE_SECTION: &str = "\n[storage]\npath = \"/tmp/usage.db\"\n"; + + #[test] + fn admin_section_parses_token_and_defaults_off() { + let (_dir, path) = write_app_config(&format!( + "{STORAGE_SECTION}\n[auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("valid [auth.admin]"); + assert!( + config + .auth + .admin + .as_ref() + .is_some_and(|a| a.token.expose_secret() == "admin-secret"), + "admin token must parse into the SecretString" + ); + assert!( + !config.server.admin_on_main_listener, + "main-listener mount must be opt-in" + ); + assert_eq!( + config.server.effective_admin_bind(), + Some("127.0.0.1:3457".parse().unwrap()), + "default admin bind is loopback on the port after the main bind" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::DedicatedListener), + "the default posture is the dedicated listener" + ); + } + + #[test] + fn admin_explicit_bind_and_main_listener_opt_in_parse() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\nadmin_on_main_listener = true\n{STORAGE_SECTION}\n\ + [auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("valid admin config"); + assert_eq!( + config.server.effective_admin_bind(), + Some("127.0.0.1:9900".parse().unwrap()), + "explicit admin_bind wins over the derived default" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::MainListener), + "admin_on_main_listener selects the main-listener posture" + ); + assert_eq!( + AdminMountPosture::MainListener.name(), + "admin-on-main-listener", + "posture names are stable for log grepping" + ); + } + + #[test] + fn admin_without_storage_is_rejected() { + let (_dir, path) = write_app_config("\n[auth.admin]\ntoken = \"admin-secret\"\n"); + let error = load_app_config(&path).expect_err("admin requires [storage]"); + assert!( + error.to_string().contains("[auth.admin] requires"), + "got: {error}" + ); + } + + #[test] + fn admin_empty_token_is_rejected() { + let (_dir, path) = + write_app_config(&format!("{STORAGE_SECTION}\n[auth.admin]\ntoken = \"\"\n")); + let error = load_app_config(&path).expect_err("empty admin token rejected"); + assert!( + error.to_string().contains("token must not be empty"), + "got: {error}" + ); + } + + #[test] + fn admin_bind_equal_to_main_bind_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:3456\"\n{STORAGE_SECTION}\n[auth.admin]\n\ + token = \"admin-secret\"\n" + )); + let error = load_app_config(&path).expect_err("same-socket bind rejected"); + assert!( + error.to_string().contains("must differ from server.bind"), + "got: {error}" + ); + } + + #[test] + fn admin_default_port_underviable_on_65535_is_rejected() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + "[server]\nbind = \"127.0.0.1:65535\"\nrequest_timeout = \"300s\"\n\ + log_level = \"info\"\nhot_reload = false\nserver_name = \"t\"\n{STORAGE_SECTION}\n\ + \n[auth.admin]\ntoken = \"admin-secret\"\n" + ), + ) + .unwrap(); + let error = load_app_config(&path).expect_err("no derivable admin port"); + assert!(error.to_string().contains("65535"), "got: {error}"); + } + + #[test] + fn admin_main_listener_optin_needs_no_derivable_port() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + "[server]\nbind = \"127.0.0.1:65535\"\nrequest_timeout = \"300s\"\n\ + log_level = \"info\"\nhot_reload = false\nserver_name = \"t\"\n\ + admin_on_main_listener = true\n{STORAGE_SECTION}\n\ + \n[auth.admin]\ntoken = \"admin-secret\"\n" + ), + ) + .unwrap(); + let config = load_app_config(&path).expect("main-listener posture needs no port"); + assert!( + !config.server.dedicated_admin_listener(), + "the opt-in without an explicit admin_bind must not bind a dedicated socket" + ); + } + + #[test] + fn explicit_admin_bind_keeps_the_dedicated_listener_under_the_optin() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\nadmin_on_main_listener = true\n{STORAGE_SECTION}\n\ + [auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("both mounts are legitimate"); + assert!( + config.server.dedicated_admin_listener(), + "an explicit admin_bind always binds the dedicated listener" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::MainListener), + "and the main-listener mount is on top of it" + ); + } + + #[test] + fn admin_bind_without_admin_section_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\n{STORAGE_SECTION}" + )); + let error = load_app_config(&path).expect_err("dead admin_bind rejected"); + assert!( + error + .to_string() + .contains("server.admin_bind is configured"), + "got: {error}" + ); + } + + #[test] + fn admin_main_listener_flag_without_admin_section_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_on_main_listener = true\n{STORAGE_SECTION}" + )); + let error = load_app_config(&path).expect_err("dead flag rejected"); + assert!( + error.to_string().contains("admin_on_main_listener"), + "got: {error}" + ); + } + + #[test] + fn admin_unknown_field_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "{STORAGE_SECTION}\n[auth.admin]\ntoken = \"admin-secret\"\nrotates = \"weekly\"\n" + )); + assert!( + load_app_config(&path).is_err(), + "deny_unknown_fields must reject unrecognized [auth.admin] keys" + ); + } + + #[test] + fn admin_posture_is_none_without_admin_section() { + let (_dir, path) = write_app_config(STORAGE_SECTION); + let config = load_app_config(&path).expect("valid config"); + assert_eq!( + config.admin_posture(), + None, + "no [auth.admin] means no admin surface anywhere" + ); + } + + #[test] + fn admin_auth_config_debug_redacts_token() { + let admin = AdminAuthConfig { + token: SecretString::from("super-secret-admin-token"), + }; + let debug = format!("{admin:?}"); + assert!( + !debug.contains("super-secret-admin-token"), + "Debug must redact the admin token, got: {debug}" + ); + } } diff --git a/crates/llm-proxy-core/src/provider_registry.rs b/crates/llm-proxy-core/src/provider_registry.rs index 977e0dd..4ded318 100644 --- a/crates/llm-proxy-core/src/provider_registry.rs +++ b/crates/llm-proxy-core/src/provider_registry.rs @@ -1,22 +1,11 @@ //! Provider registry: load provider TOML files and resolve route targets to //! adapter configurations. //! -//! # Provider-based routing -//! -//! ```text -//! provider_name + route_kind -> provider config + route -> adapter name -//! adapter_name -> protocol + endpoint + headers -//! requested_model -> model_aliases -> upstream_model -//! ``` -//! -//! The registry does **not**: -//! -//! - translate protocol fields -//! - parse response or stream chunks -//! - build final protocol-specific URL paths -//! - inspect messages or content -//! - mutate sampling, tools, metadata, reasoning, cache, or stream intent -//! - infer protocol families from model names +//! Resolution chain: provider name and route kind select a provider config and +//! adapter; the requested model is mapped through `model_aliases` to the +//! upstream model. The registry does not translate protocol fields, parse +//! responses or stream chunks, build final protocol-specific URLs, inspect +//! messages, or infer protocol families from model names. use std::collections::{HashMap, HashSet}; use std::path::Path; @@ -64,10 +53,6 @@ pub enum ProviderRouteResolutionError { }, } -// --------------------------------------------------------------------------- -// ProviderAdapterTargetConfig -// --------------------------------------------------------------------------- - /// Fully resolved adapter target produced by the provider registry. /// /// Contains everything the provider adapter layer needs to build an upstream @@ -97,12 +82,11 @@ pub struct ProviderAdapterTargetConfig { pub requested_model: String, /// The model name to send to the upstream provider. pub upstream_model: String, - /// Optional static headers from the adapter config (e.g. `anthropic-version`). - /// - /// `Arc`-shared with the source [`crate::provider_config::ProviderAdapterConfig`] so producing a - /// resolved target per request is a refcount bump, not a `HashMap` clone - /// (LOW-5). Downstream code that needs to mutate (e.g. defaulting the - /// `anthropic-version` header) materializes an owned copy. + /// Optional static headers from the adapter config (e.g. + /// `anthropic-version`), `Arc`-shared with the source + /// [`crate::provider_config::ProviderAdapterConfig`] so resolving a route + /// per request is a refcount bump, not a map clone. Downstream code that + /// needs to mutate materializes an owned copy. pub headers: Arc>, } @@ -122,10 +106,6 @@ impl std::fmt::Debug for ProviderAdapterTargetConfig { } } -// --------------------------------------------------------------------------- -// ProviderRegistry -// --------------------------------------------------------------------------- - /// Loads provider TOML files and resolves provider routes to /// [`ProviderAdapterTargetConfig`]s. /// @@ -171,12 +151,8 @@ impl ProviderRegistry { })?; let mut providers: HashMap = HashMap::new(); - // Track which file each provider name was first loaded from, so that - // duplicate-name errors can report *both* file paths for quick diagnosis. let mut source_paths: HashMap = HashMap::new(); - // Collect and sort entries by filename for deterministic ordering across - // platforms (filesystem order is not guaranteed). let mut sorted_entries: Vec<_> = entries .collect::, _>>() @@ -189,10 +165,6 @@ impl ProviderRegistry { for entry in sorted_entries { let file_path = entry.path(); - // Skip symlinks to prevent a malicious actor with write access to - // the providers directory from tricking the proxy into reading - // arbitrary files. This check uses symlink_metadata which does not - // follow symlinks, unlike entry.file_type() which does. if file_path.is_symlink() { tracing::warn!( path = %file_path.display(), @@ -201,7 +173,6 @@ impl ProviderRegistry { continue; } - // Skip non-regular files (directories, pipes, sockets, etc.). if let Ok(ft) = entry.file_type() { if !ft.is_file() { tracing::debug!( @@ -250,10 +221,6 @@ impl ProviderRegistry { /// Build a registry from an existing map of provider configs. /// - /// Returns an error if any two configs share the same provider name - /// (this should not happen if the caller built the map correctly, but - /// we validate defensively). - /// /// # Errors /// /// Returns [`CoreError::ProviderResolution`] if duplicate provider names @@ -296,10 +263,6 @@ impl ProviderRegistry { &self, known_protocols: impl IntoIterator>, ) -> Result<(), CoreError> { - // Note: This allocates a String per protocol name. The set is tiny - // (4 built-in protocols) and this runs once at startup, so the cost - // is negligible. A HashSet is used because the input iterator - // owns its items and we cannot borrow from them. let known: HashSet = known_protocols .into_iter() .map(|s| s.as_ref().to_owned()) @@ -345,14 +308,12 @@ impl ProviderRegistry { route_kind: ProviderRouteKind, requested_model: &str, ) -> Result { - // 1. Look up the provider. let provider = self.providers.get(provider_name).ok_or_else(|| { ProviderRouteResolutionError::UnknownProvider { provider: provider_name.to_owned(), } })?; - // 2. Look up the route to find the adapter name. let adapter_name = provider.routes.get(route_kind).ok_or_else(|| { let route_name = match route_kind { ProviderRouteKind::ChatCompletions => "chat_completions", @@ -364,7 +325,6 @@ impl ProviderRegistry { } })?; - // 3. Look up the adapter to get protocol, endpoint, and headers. let adapter_cfg = provider.adapters.get(adapter_name).ok_or_else(|| { let available = if provider.adapters.is_empty() { "(none)".to_owned() @@ -387,7 +347,6 @@ impl ProviderRegistry { } })?; - // 4. Apply provider-local model alias if present. let upstream_model = provider .model_aliases .get(requested_model) @@ -467,13 +426,8 @@ impl<'a> IntoIterator for &'a ProviderRegistry { } } -/// Maximum number of adapter names shown in error messages before truncation. const MAX_ADAPTERS_IN_ERROR: usize = 5; -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -486,6 +440,7 @@ mod tests { api_key: SecretString::from("api-secret"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -525,6 +480,7 @@ mod tests { api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -552,14 +508,11 @@ mod tests { )]), }; let registry = ProviderRegistry::from_providers(vec![provider]).expect("registry"); - // Pricing is keyed by the UPSTREAM (alias-resolved) model id. let p = registry .pricing_for("example", "real-upstream-model") .expect("pricing present"); assert_eq!(p.input.to_string(), "0.001"); - // The client-facing alias is NOT a pricing key. assert!(registry.pricing_for("example", "alias-name").is_none()); - // Unknown provider / model -> None. assert!( registry .pricing_for("other", "real-upstream-model") @@ -590,6 +543,7 @@ mod tests { api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -612,9 +566,6 @@ mod tests { )); } - // -- load_from_dir tests --------------------------------------------------- - - /// Helper: write a minimal valid provider TOML file into `dir`. fn write_provider(dir: &std::path::Path, filename: &str, name: &str) { let content = format!( r#" @@ -683,31 +634,26 @@ chat_completions = "chat" let dir = tempfile::tempdir().expect("temp dir"); write_provider(dir.path(), "real.toml", "real"); - // Create a symlink to the real file. #[cfg(unix)] { let link = dir.path().join("linked.toml"); std::os::unix::fs::symlink(dir.path().join("real.toml"), &link) .expect("create symlink"); } - // On non-Unix, skip symlink creation (test still passes; just fewer entries). let registry = ProviderRegistry::load_from_dir(dir.path()).expect("load_from_dir"); - // Should only have "real" -- the symlink must be skipped. assert!(registry.get("real").is_some()); assert_eq!(registry.len(), 1, "symlink should be skipped"); } - // -- MissingAdapter test --------------------------------------------------- - #[test] fn resolve_provider_route_returns_missing_adapter_error() { - // Build a provider with a route that points to a non-existent adapter. let provider = ProviderConfig { name: "broken".to_owned(), api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, - adapters: HashMap::new(), // no adapters defined + egress: None, + adapters: HashMap::new(), routes: ProviderRoutesConfig { chat_completions: Some("nonexistent_adapter".to_owned()), messages: None, @@ -730,8 +676,6 @@ chat_completions = "chat" )); } - // -- from_providers duplicate detection ----------------------------------- - #[test] fn from_providers_rejects_duplicate_names() { let provider_a = ProviderConfig { @@ -739,6 +683,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -758,8 +703,6 @@ chat_completions = "chat" assert!(err.to_string().contains("duplicate provider name")); } - // -- len() / is_empty() / get() ------------------------------------------- - #[test] fn empty_registry_is_empty_and_zero_len() { let registry = ProviderRegistry::from_providers(Vec::new()).expect("registry"); @@ -775,6 +718,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -789,8 +733,6 @@ chat_completions = "chat" assert!(registry.get("nonexistent").is_none()); } - // -- catalog_models() ----------------------------------------------------- - #[test] fn catalog_models_returns_none_for_unknown_provider() { let registry = ProviderRegistry::from_providers(Vec::new()).expect("registry"); @@ -804,6 +746,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -823,6 +766,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -832,12 +776,9 @@ chat_completions = "chat" }; let registry = ProviderRegistry::from_providers(vec![provider]).expect("registry"); let models = registry.catalog_models("test").expect("provider exists"); - // Default catalog has no models. assert!(models.is_empty()); } - // -- resolve without alias (fallback) ------------------------------------- - #[test] fn resolve_without_alias_returns_requested_model() { let provider = ProviderConfig { @@ -845,6 +786,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -870,8 +812,6 @@ chat_completions = "chat" assert_eq!(target.requested_model, "gpt-4o"); } - // -- IntoIterator --------------------------------------------------------- - #[test] fn into_iterator_works() { let provider_a = ProviderConfig { @@ -879,6 +819,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -898,11 +839,8 @@ chat_completions = "chat" assert!(names.contains(&"b")); } - // -- Adapter list truncation in MissingAdapter error ---------------------- - #[test] fn missing_adapter_error_truncates_long_adapter_list() { - // Build a provider with >5 adapters and a route pointing to a non-existent one. let mut adapters = HashMap::new(); for i in 0..7 { adapters.insert( @@ -919,6 +857,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters, routes: ProviderRoutesConfig { chat_completions: Some("nonexistent".to_owned()), diff --git a/crates/llm-proxy-core/src/test_support.rs b/crates/llm-proxy-core/src/test_support.rs index 021ad76..05584a8 100644 --- a/crates/llm-proxy-core/src/test_support.rs +++ b/crates/llm-proxy-core/src/test_support.rs @@ -1,32 +1,23 @@ //! Shared test utilities for `llm-proxy-core`. //! -//! Provides a single crate-level mutex (`TEST_ENV_LOCK`) that serialises all -//! environment-variable-mutating tests across modules (currently used by -//! [`crate::env_interpolate`] and [`crate::provider_config`]), eliminating -//! the race condition that occurred when each module used its own independent -//! mutex. +//! [`TEST_ENV_LOCK`] serialises all environment-variable-mutating tests +//! across modules. use std::sync::Mutex; -/// Crate-level mutex for serialising tests that mutate environment variables. -/// -/// Both `config::tests` and `provider_config::tests` acquire this lock before -/// setting or removing env vars, preventing the race condition where one -/// module's `EnvVarGuard` removes a variable that the other module's test -/// expects to be present. +/// Crate-level mutex serialising tests that mutate environment variables. /// /// # Safety invariant /// -/// The mutex guarantees exclusive access, so `unsafe { std::env::set_var(...) }` -/// and `unsafe { std::env::remove_var(...) }` calls within the lock scope are -/// safe: no other test thread can concurrently read or write the same variable. +/// The mutex guarantees exclusive access, so the `unsafe` +/// `std::env::set_var` / `std::env::remove_var` calls within the lock scope +/// are safe: no other test thread can concurrently read or write the +/// environment. pub static TEST_ENV_LOCK: Mutex<()> = Mutex::new(()); -/// RAII wrapper that holds the test env mutex lock. +/// RAII wrapper that holds [`TEST_ENV_LOCK`]; the lock is released on drop. /// -/// Acquire via [`TestEnvLock::acquire`]. The lock is released when this value -/// is dropped. This is a test utility for serialising environment-variable-mutating -/// tests to prevent race conditions. +/// Acquire via [`TestEnvLock::acquire`]. #[derive(Debug)] pub struct TestEnvLock { _guard: std::sync::MutexGuard<'static, ()>, @@ -43,38 +34,23 @@ impl TestEnvLock { /// RAII guard that saves an environment variable on creation and restores it /// (or removes it) on drop. /// -/// This is a test utility for safely setting/removing environment variables -/// in unit tests. It ensures the variable is always restored to its original -/// state when the guard is dropped, even if the test panics. -/// /// # Lock discipline (load-bearing invariant) /// -/// Both [`EnvVarGuard::set`] and [`EnvVarGuard::remove`] perform `unsafe` +/// [`EnvVarGuard::set`] and [`EnvVarGuard::remove`] perform `unsafe` /// `std::env::set_var` / `remove_var` calls, and [`Drop`] restores the -/// original value with the same `unsafe` calls. These are only sound while -/// no other thread is touching the process environment, i.e. while a -/// [`TestEnvLock`] is held. -/// -/// `EnvVarGuard` is returned by value with **no lifetime tie** to the lock, -/// so this invariant is enforced by *convention*, not by the borrow checker. -/// Every existing call site respects it by binding the lock and the guard in -/// the same scope: +/// original value with the same calls. These are only sound while no other +/// thread is touching the process environment, i.e. while a [`TestEnvLock`] +/// is held. The guard carries no lifetime tie to the lock, so this invariant +/// is enforced by convention: **callers must not let an [`EnvVarGuard`] +/// outlive the [`TestEnvLock`] that authorises it** (e.g. by returning it +/// from a helper or storing it in a struct) — bind both in the same scope so +/// the guard drops before the lock. /// /// ```no_run /// # use llm_proxy_core::test_support::{TestEnvLock, EnvVarGuard}; -/// let _lock = TestEnvLock::acquire(); // acquired first ... -/// let _guard = EnvVarGuard::set("K", "v"); // ... so the guard is dropped -/// // *before* the lock (reverse -/// // drop order), keeping the -/// // restore inside the lock scope. +/// let _lock = TestEnvLock::acquire(); // acquired first ... +/// let _guard = EnvVarGuard::set("K", "v"); // ... so the guard drops first. /// ``` -/// -/// **Callers must not let an `EnvVarGuard` outlive the `TestEnvLock` that -/// authorises it** (e.g. by returning it from a helper, storing it in a -/// struct, or moving it into a `'static` context). Doing so would make the -/// restore-on-drop run unsynchronised against concurrent env access. If such -/// a use case is needed, restructure the guard to own the `MutexGuard` (see -/// the TODO in `Drop` below) so the compiler enforces the invariant. #[derive(Debug)] pub struct EnvVarGuard { key: String, @@ -148,13 +124,7 @@ impl EnvVarGuard { /// [`TestEnvLock`] is still held. This is guaranteed today solely by Rust's /// reverse field-drop order — the guard is always dropped before the lock /// binding it shares a scope with — because [`EnvVarGuard`] carries no -/// lifetime tie to the lock (see the type-level "Lock discipline" note). -/// -/// TODO(test-support): for full defence-in-depth, restructure `EnvVarGuard` -/// to *own* the `MutexGuard` (have `set`/`remove` consume a `TestEnvLock`), -/// so the compiler proves the restore happens inside the lock. That change -/// requires updating every call site in `env_interpolate` and -/// `provider_config` tests and is deferred to a dedicated refactor. +/// lifetime tie to the lock. impl Drop for EnvVarGuard { #[expect( unsafe_code, diff --git a/crates/llm-proxy-core/src/token/counter.rs b/crates/llm-proxy-core/src/token/counter.rs index 77ffc00..334a4cc 100644 --- a/crates/llm-proxy-core/src/token/counter.rs +++ b/crates/llm-proxy-core/src/token/counter.rs @@ -1,16 +1,13 @@ //! Token counter with model-aware BPE dispatch. //! -//! [`Counter`] is the public front door for token counting. For known OpenAI -//! model ids it dispatches to a real BPE tokenizer (via -//! [`super::tiktoken`]); for everything else it falls back to a character -//! heuristic (~4 chars/token). The BPE ranks are bundled at compile time, so -//! counting is offline-safe and the heuristic is always available as a -//! last resort. +//! [`Counter`] is the public front door for token counting: a real BPE +//! tokenizer (via [`super::tiktoken`]) for known OpenAI model ids, a +//! character heuristic (~4 chars/token) for everything else. BPE ranks are +//! bundled at compile time, so counting is offline-safe. //! //! The per-message overhead mirrors `tiktoken-rs`'s //! `num_tokens_from_messages` convention (4 tokens/message for the `gpt-3.5` -//! family, 3 otherwise, plus a trailing +3) so unit-test expectations stay -//! aligned with the upstream tokenizer. +//! family, 3 otherwise, plus a trailing +3). use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -21,9 +18,7 @@ use crate::token::tiktoken::{self, Encoding, HeuristicTokenizer, TiktokenTokeniz #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct MessageContent { /// The role of the message author (e.g. `"user"`, `"assistant"`, `"system"`). - /// - /// Role validation is the caller's responsibility. The counter does not - /// enforce valid role strings. + /// Validation is the caller's responsibility. pub role: String, /// The text content of the message. pub content: String, @@ -41,23 +36,17 @@ impl MessageContent { /// Model-aware token counter. /// -/// Holds the heuristic backend and a lazily-populated cache of loaded BPE -/// tokenizers keyed by [`Encoding`]. The cache uses interior mutability so a -/// `&Counter` (shared via `Arc` in `AppState`) can load encodings on first use. -/// Cloning a `Counter` shares the cache (the `Mutex` lives behind an `Arc`). +/// Holds a heuristic backend plus a lazily-populated cache of BPE tokenizers +/// keyed by [`Encoding`]; the cache uses interior mutability so a `&Counter` +/// can load encodings on first use. Cloning a `Counter` shares the cache. #[derive(Clone, Default)] pub struct Counter { - /// Heuristic fallback for unknown models and load failures. heuristic: HeuristicTokenizer, - /// Lazily-loaded BPE tokenizers, keyed by encoding. Cheap to share: clone - /// is an `Arc` refcount bump. cache: Arc>>>, } impl std::fmt::Debug for Counter { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // Report how many encodings are loaded rather than dumping the (large) - // BPE tables or fighting `Mutex`'s `Debug` bound. let loaded = self.cache.lock().map(|m| m.len()).unwrap_or(0); f.debug_struct("Counter") .field("backend", &"bpe+heuristic") @@ -82,7 +71,6 @@ impl Counter { match tiktoken::encoding_for_model(model) { Some(encoding) => match self.bpe_for(encoding) { Some(tokenizer) => tokenizer.count_tokens(text), - // Known model but the ranks could not be loaded: heuristic. None => self.heuristic.count_tokens(text), }, None => self.heuristic.count_tokens(text), @@ -92,13 +80,11 @@ impl Counter { /// Get the cached BPE tokenizer for an encoding, loading it lazily on first /// use. Returns `None` only if the ranks could not be loaded. fn bpe_for(&self, encoding: Encoding) -> Option> { - // Fast path: already loaded. if let Ok(cache) = self.cache.lock() { if let Some(tokenizer) = cache.get(&encoding).cloned() { return Some(tokenizer); } } - // Slow path: load outside the lock, then insert. let tokenizer = Arc::new(TiktokenTokenizer::new(encoding.load()?)); if let Ok(mut cache) = self.cache.lock() { cache.insert(encoding, Arc::clone(&tokenizer)); @@ -108,8 +94,7 @@ impl Counter { /// Estimate the total token count for a chat completion request. /// - /// The formula follows `tiktoken-rs`'s `num_tokens_from_messages` - /// convention rather than any separate "Go reference": + /// Follows `tiktoken-rs`'s `num_tokens_from_messages` convention: /// /// ```text /// tokens_per_message = if model.starts_with("gpt-3.5") { 4 } else { 3 } @@ -118,22 +103,16 @@ impl Counter { /// + 3 // trailing priming tokens /// ``` /// - /// Each piece is counted with the model-appropriate backend (BPE for known - /// OpenAI models, heuristic otherwise). The system prompt is folded into - /// the message loop as a single `"system"`-role message when it is - /// non-empty, so it contributes `tokens_per_message` exactly the way - /// tiktoken-rs folds a system message into its iteration. When `system` is - /// empty no per-message overhead is added for it. + /// The system prompt is folded into the message loop as a single + /// `"system"`-role message when it is non-empty; when `system` is empty no + /// per-message overhead is added for it. /// - /// Note: `system` is a flat string. If the upstream protocol supports - /// multiple system blocks, the caller must pre-concatenate them. - /// Role validation is the caller's responsibility. + /// `system` is a flat string: if the upstream protocol supports multiple + /// system blocks, the caller must pre-concatenate them. Role validation is + /// the caller's responsibility. pub fn count_messages(&self, model: &str, system: &str, messages: &[MessageContent]) -> usize { - // tiktoken-rs convention (see `num_tokens_from_messages`): - // 4 tokens/message for gpt-3.5*, 3 for everything else. const GPT_35_PER_MESSAGE: usize = 4; const DEFAULT_PER_MESSAGE: usize = 3; - // Trailing +3: every reply is primed with `<|start|>assistant<|message|>`. const TRAILING_TOKENS: usize = 3; let tokens_per_message = if model.starts_with("gpt-3.5") { @@ -142,10 +121,6 @@ impl Counter { DEFAULT_PER_MESSAGE }; - // The system prompt is treated as one message in the loop (matching - // tiktoken-rs, which iterates system alongside the rest). Empty system - // contributes nothing -- no bogus per-message overhead for a missing - // block. let system_overhead = if system.is_empty() { 0 } else { @@ -171,8 +146,6 @@ impl Counter { mod tests { use super::*; - /// An unknown model id that matches no [`Encoding`] prefix, so `Counter` - /// routes it to the heuristic backend. const HEURISTIC_MODEL: &str = "unknown-test-model"; #[test] @@ -184,36 +157,30 @@ mod tests { #[test] fn count_tokens_single_char() { let counter = Counter::new(); - // 1 char / 4 = 0, but minimum is 1 for non-empty. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, "a"), 1); } #[test] fn count_tokens_short_text() { let counter = Counter::new(); - // "hello" = 5 chars, 5/4 = 1. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, "hello"), 1); } #[test] fn count_tokens_medium_text() { let counter = Counter::new(); - // 20 chars / 4 = 5 tokens. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, &"a".repeat(20)), 5); } #[test] fn count_tokens_longer_text() { let counter = Counter::new(); - // 100 chars / 4 = 25 tokens. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, &"a".repeat(100)), 25); } #[test] fn count_messages_no_messages() { let counter = Counter::new(); - // system="test" folded in: per_msg(3) + role("system"=6/4=1) + content("test"=4/4=1) = 5 - // 0 messages, +3 trailing = 5 + 0 + 3 = 8 let total = counter.count_messages(HEURISTIC_MODEL, "test", &[]); assert_eq!(total, 8); } @@ -221,7 +188,6 @@ mod tests { #[test] fn count_messages_empty_system() { let counter = Counter::new(); - // Empty system -> no system overhead. 0 messages. +3 trailing = 3. let total = counter.count_messages(HEURISTIC_MODEL, "", &[]); assert_eq!(total, 3); } @@ -230,9 +196,6 @@ mod tests { fn count_messages_single_user_message() { let counter = Counter::new(); let messages = vec![MessageContent::new("user", "hello")]; - // empty system -> 0 - // msg: per_msg(3) + role("user"=4/4=1) + content("hello"=5/4=1) = 5 - // +3 trailing = 0 + 5 + 3 = 8 let total = counter.count_messages(HEURISTIC_MODEL, "", &messages); assert_eq!(total, 8); } @@ -246,12 +209,6 @@ mod tests { MessageContent::new("assistant", "I'm doing well, thank you!"), ]; let total = counter.count_messages(HEURISTIC_MODEL, "Be helpful", &messages); - // tokens_per_message = 3 (HEURISTIC_MODEL is not gpt-3.5). - // system="Be helpful"(10/4=2): per_msg(3) + role("system"=6/4=1) + content(2) = 6 - // msg0: per_msg(3) + role("system"=1) + content="You are a helpful assistant."(27/4=6) = 10 - // msg1: per_msg(3) + role("user"=4/4=1) + content="Hello, how are you?"(20/4=5) = 9 - // msg2: per_msg(3) + role("assistant"=9/4=2) + content="I'm doing well, thank you!"(27/4=6) = 11 - // +3 trailing = 6 + 10 + 9 + 11 + 3 = 39 assert_eq!(total, 39); } @@ -259,9 +216,6 @@ mod tests { fn count_messages_with_system_prompt() { let counter = Counter::new(); let messages = vec![MessageContent::new("user", "Hi")]; - // system="You are helpful"(15/4=3): per_msg(3) + role("system"=1) + content(3) = 7 - // msg: per_msg(3) + role("user"=4/4=1) + content("Hi"=2/4=0->1) = 5 - // +3 trailing = 7 + 5 + 3 = 15 let total = counter.count_messages(HEURISTIC_MODEL, "You are helpful", &messages); assert_eq!(total, 15); } @@ -283,9 +237,6 @@ mod tests { fn count_messages_empty_role() { let counter = Counter::new(); let messages = vec![MessageContent::new("", "hello")]; - // empty system -> 0 - // msg: per_msg(3) + role("") + content("hello"=5/4=1) = 3 + 0 + 1 = 4 - // +3 trailing = 0 + 4 + 3 = 7 let total = counter.count_messages(HEURISTIC_MODEL, "", &messages); assert_eq!(total, 7); } @@ -294,9 +245,6 @@ mod tests { fn count_messages_empty_content() { let counter = Counter::new(); let messages = vec![MessageContent::new("user", "")]; - // empty system -> 0 - // msg: per_msg(3) + role("user"=4/4=1) + content("") = 3 + 1 + 0 = 4 - // +3 trailing = 0 + 4 + 3 = 7 let total = counter.count_messages(HEURISTIC_MODEL, "", &messages); assert_eq!(total, 7); } @@ -304,7 +252,6 @@ mod tests { #[test] fn count_tokens_whitespace_only() { let counter = Counter::new(); - // Whitespace-only string is non-empty, so minimum 1 token. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, " "), 1); assert_eq!(counter.count_tokens(HEURISTIC_MODEL, "\t\n"), 1); } @@ -312,14 +259,11 @@ mod tests { #[test] fn count_tokens_multibyte_unicode() { let counter = Counter::new(); - // Emoji: U+1F600 is 4 bytes but 1 code point. 1/4 = 0 -> max(0,1) = 1. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, "\u{1F600}"), 1); - // CJK string: 4 characters (each 3 bytes in UTF-8). 4/4 = 1 token. assert_eq!( counter.count_tokens(HEURISTIC_MODEL, "\u{4F60}\u{597D}\u{4E16}\u{754C}"), 1 ); - // 8 CJK characters: 8/4 = 2 tokens. let cjk_8: String = "\u{4F60}\u{597D}\u{4E16}\u{754C}\u{4F60}\u{597D}\u{4E16}\u{754C}".to_owned(); assert_eq!(counter.count_tokens(HEURISTIC_MODEL, &cjk_8), 2); @@ -328,18 +272,14 @@ mod tests { #[test] fn count_tokens_mixed_ascii_unicode() { let counter = Counter::new(); - // "hello" (5 chars) + emoji (1 char) = 6 chars. 6/4 = 1. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, "hello\u{1F600}"), 1); } - // -- BPE dispatch (real tiktoken) ---------------------------------------- - #[test] fn count_tokens_unknown_model_uses_heuristic() { let counter = Counter::new(); let text = "The quick brown fox jumps over the lazy dog."; let heuristic = (text.chars().count() / 4).max(1); - // Unknown model -> heuristic. assert_eq!(counter.count_tokens(HEURISTIC_MODEL, text), heuristic); } @@ -348,24 +288,18 @@ mod tests { let counter = Counter::new(); let text = "The quick brown fox jumps over the lazy dog."; - // gpt-4o -> o200k_base. The count must match tiktoken's own encoding of - // the same text under o200k_base (this verifies dispatch + selection, - // not a hardcoded magic number). let o200k = tiktoken_rs::o200k_base().expect("o200k_base loads"); assert_eq!( counter.count_tokens("gpt-4o", text), o200k.encode_with_special_tokens(text).len() ); - // gpt-4 -> cl100k_base. let cl100k = tiktoken_rs::cl100k_base().expect("cl100k_base loads"); assert_eq!( counter.count_tokens("gpt-4", text), cl100k.encode_with_special_tokens(text).len() ); - // BPE must differ from the heuristic for this text (otherwise the test - // would not prove BPE is active). let heuristic = (text.chars().count() / 4).max(1); assert_ne!( counter.count_tokens("gpt-4o", text), @@ -376,9 +310,6 @@ mod tests { #[test] fn count_tokens_bpe_encoding_is_cached() { - // A second count for the same encoding should reuse the cached BPE - // (loaded_encodings grows to 1, not 2, and a second distinct encoding - // grows it to 2). let counter = Counter::new(); let _ = counter.count_tokens("gpt-4o", "hello world"); let _ = counter.count_tokens("gpt-4o-2024-08-06", "another"); @@ -393,13 +324,6 @@ mod tests { assert_eq!(loaded, 2, "gpt-4 adds the cl100k_base encoding"); } - // -- tok-02: golden parity with tiktoken-rs ----------------------------- - - /// Fixed multi-message fixture shared with tiktoken-rs's - /// `num_tokens_from_messages`. The proxy's `count_messages` must match the - /// upstream tokenizer to within a small tolerance (the plan allows <= 3 - /// absolute; since we use the same BPE and the same overhead convention, - /// equality should hold exactly when `name` is absent). fn golden_messages() -> Vec { vec![ MessageContent::new( @@ -435,9 +359,6 @@ mod tests { ] } - /// The system prompt passed separately to the proxy must equal the - /// `content` of the system message tiktoken-rs sees (otherwise the two - /// formulas would not be comparable). const GOLDEN_SYSTEM: &str = "Be concise and correct."; #[test] @@ -450,9 +371,6 @@ mod tests { let theirs = tiktoken_rs::num_tokens_from_messages("gpt-4o", &tiktoken_msgs) .expect("gpt-4o resolves to a supported chat tokenizer"); - // `num_tokens_from_messages` does not model the separate top-level - // `system` string the proxy exposes. We account for it by handing the - // same text to tiktoken-rs as an extra leading system message. let mut with_system = golden_tiktoken_messages(); with_system.insert( 0, @@ -466,15 +384,11 @@ mod tests { let theirs_with_system = tiktoken_rs::num_tokens_from_messages("gpt-4o", &with_system) .expect("gpt-4o resolves to a supported chat tokenizer"); - // Same BPE + same overhead convention -> exact equality is expected; - // the plan's <= 3 tolerance is the documented safety margin. assert!( (ours as isize - theirs_with_system as isize).abs() <= 3, "gpt-4o mismatch: ours={ours}, tiktoken_baseline_no_system={theirs}, \ tiktoken_with_system={theirs_with_system}" ); - // And the system-less baseline must be strictly less than our value - // (we added a real system message). assert!( ours > theirs, "gpt-4o: adding a system prompt must increase the count (ours={ours}, baseline={theirs})" @@ -483,7 +397,6 @@ mod tests { #[test] fn count_messages_matches_tiktoken_for_gpt35_turbo() { - // gpt-3.5-turbo exercises the `tokens_per_message = 4` branch. let counter = Counter::new(); let messages = golden_messages(); @@ -507,10 +420,6 @@ mod tests { ); } - /// gpt-3.5-turbo uses 4 tokens/message while other cl100k models (e.g. - /// gpt-4) use 3. A fixture with at least one real message must therefore - /// count strictly higher for gpt-3.5 than for gpt-4. This directly - /// exercises the `model.starts_with("gpt-3.5")` branch. #[test] fn count_messages_gpt35_overhead_is_higher_than_gpt4() { let counter = Counter::new(); @@ -518,8 +427,6 @@ mod tests { let gpt35 = counter.count_messages("gpt-3.5-turbo", "Be helpful", &messages); let gpt4 = counter.count_messages("gpt-4", "Be helpful", &messages); - // system counts as 1 message + the user message = 2 messages, so gpt-3.5 - // (4/msg) is +2 ahead of gpt-4 (3/msg). assert_eq!( gpt35 - gpt4, 2, @@ -527,23 +434,17 @@ mod tests { ); } - // -- NON-ZST: Counter is Clone + Debug + non-zero-sized ------------------ - #[test] fn counter_is_clone_and_debug() { let counter = Counter::new(); let cloned = counter.clone(); - // Debug must not panic and must mention the backend. let s = format!("{counter:?}"); assert!(s.contains("bpe+heuristic"), "debug output: {s}"); - // Cloned counter is independently usable. assert!(format!("{cloned:?}").contains("bpe+heuristic")); } #[test] fn counter_is_not_zero_sized() { - // A ZST counter (e.g. if the cache/heuristic fields were dropped) would - // be a regression: it would mean the lazily-loaded BPE cache is gone. assert!( std::mem::size_of::() > 0, "Counter must not be zero-sized; it carries an Arc>>" diff --git a/crates/llm-proxy-core/src/token/mod.rs b/crates/llm-proxy-core/src/token/mod.rs index a5256b3..b0b0725 100644 --- a/crates/llm-proxy-core/src/token/mod.rs +++ b/crates/llm-proxy-core/src/token/mod.rs @@ -2,9 +2,8 @@ //! //! Provides a [`Counter`] that counts tokens with a real BPE tokenizer //! ([`tiktoken`]) for known OpenAI model ids and falls back to a character -//! heuristic (~4 chars/token) for everything else. The BPE ranks are bundled at -//! compile time, so counting is offline-safe and the heuristic is always -//! available as a last resort. +//! heuristic (~4 chars/token) for everything else. BPE ranks are bundled at +//! compile time, so counting is offline-safe. pub mod counter; pub mod tiktoken; diff --git a/crates/llm-proxy-core/src/token/tiktoken.rs b/crates/llm-proxy-core/src/token/tiktoken.rs index dbea508..9bd6b9d 100644 --- a/crates/llm-proxy-core/src/token/tiktoken.rs +++ b/crates/llm-proxy-core/src/token/tiktoken.rs @@ -1,13 +1,9 @@ //! Real BPE token-counting backends via `tiktoken-rs`. //! -//! `tiktoken-rs` bundles the OpenAI BPE rank files (`cl100k_base`, -//! `o200k_base`, `p50k_base`, `r50k_base`) at compile time via `include_bytes!`, -//! so there is no network fetch at runtime and the proxy is fully offline-safe. -//! The binary grows ~2-4 MB per encoding; ranks are gzip-compressed and -//! decompressed on first use. -//! -//! [`Counter`](super::Counter) remains the public front door. It dispatches to -//! a [`TiktokenTokenizer`] for known OpenAI model ids and falls back to a +//! The BPE rank tables for the supported OpenAI encodings are bundled at +//! compile time, so counting requires no network access and works offline. +//! [`Counter`](super::Counter) remains the public front door: it dispatches +//! to a [`TiktokenTokenizer`] for known OpenAI model ids and falls back to a //! [`HeuristicTokenizer`] (~4 chars/token) for everything else, so the //! `/v1/messages/count_tokens` route never hard-fails on an unfamiliar model. @@ -15,10 +11,6 @@ use std::sync::Arc; use tiktoken_rs::CoreBPE; -/// Approximate number of characters per token for the heuristic backend. -/// -/// Shared with [`super::counter`] so the heuristic path keeps the same -/// approximation the proxy has always used. pub(crate) const CHARS_PER_TOKEN: usize = 4; /// Token-counting backend. @@ -51,9 +43,6 @@ impl Tokenizer for HeuristicTokenizer { } } -/// Heuristic token estimate: ~`CHARS_PER_TOKEN` characters per token, counting -/// Unicode code points. Returns 0 for empty input and a minimum of 1 for any -/// non-empty string. pub(crate) fn heuristic_count(text: &str) -> usize { if text.is_empty() { return 0; @@ -63,9 +52,7 @@ pub(crate) fn heuristic_count(text: &str) -> usize { /// BPE backend using a specific tiktoken encoding. /// -/// `CoreBPE` is `Send + Sync` in `tiktoken-rs` 0.6 (its fields are -/// `HashMap`/`Vec`). Wrapped in `Arc` so cloning the tokenizer is a -/// refcount bump. +/// Wrapped in [`Arc`] so cloning the tokenizer is a refcount bump. #[derive(Clone)] pub struct TiktokenTokenizer { bpe: Arc, @@ -73,14 +60,11 @@ pub struct TiktokenTokenizer { impl std::fmt::Debug for TiktokenTokenizer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // Do not dump the (large) rank tables; the encoding identity is what - // matters for diagnostics. f.debug_struct("TiktokenTokenizer").finish_non_exhaustive() } } impl TiktokenTokenizer { - /// Wrap an already-loaded BPE. pub(crate) fn new(bpe: Arc) -> Self { Self { bpe } } @@ -91,26 +75,13 @@ impl Tokenizer for TiktokenTokenizer { if text.is_empty() { return 0; } - // `encode_with_special_tokens` is the counting convention used by - // tiktoken-rs's own message-token counter. For typical input that does - // not contain literal special-token strings it is equivalent to - // `encode_ordinary`. - // - // tiktoken-rs 0.6 internally `.unwrap()`s the fancy-regex match - // results (see `vendor_tiktoken`'s `find_from_pos(...).unwrap()` and - // `mat.unwrap()`). On pathological input those regexes can return - // `Err` (catastrophic backtracking / stack overflow), which would - // `.unwrap()`-panic the worker task handling - // `/v1/messages/count_tokens`. The plan guarantees the route never - // hard-fails on the tokenizer, so we catch the panic and fall back to - // the heuristic instead. + // tiktoken-rs 0.6 internally unwraps fancy-regex match results, so a + // pathological input can panic inside the BPE encode. Catch the panic + // and fall back to the heuristic so the route never hard-fails. let bpe = std::panic::AssertUnwindSafe(&self.bpe); match std::panic::catch_unwind(|| bpe.encode_with_special_tokens(text)) { Ok(encoded) => encoded.len(), Err(panic_payload) => { - // Best-effort description of the panic payload; fancy-regex's - // `Err`-unwrap produces a string-like payload, but anything is - // possible so stay defensive. let msg = panic_payload .downcast_ref::<&'static str>() .copied() @@ -144,12 +115,6 @@ pub enum Encoding { } impl Encoding { - /// Load the BPE ranks for this encoding. - /// - /// Ranks are bundled at compile time, so this is effectively infallible in - /// practice; the `Result` is handled defensively. Returns `None` (after a - /// one-line `warn!`) if the ranks somehow cannot be loaded, so the caller - /// can fall back to the heuristic rather than failing the request. pub(crate) fn load(self) -> Option> { let result = match self { Encoding::O200kBase => tiktoken_rs::o200k_base(), @@ -184,12 +149,6 @@ impl Encoding { /// authoritative count for billing; the local counter is only an estimate for /// the `/v1/messages/count_tokens` route and pre-flight guardrails. pub fn encoding_for_model(model: &str) -> Option { - // `o200k_base` family. Prefix order matters: every branch here is checked - // BEFORE the `gpt-4` cl100k branch below, because models like `gpt-4.1`, - // `gpt-4.5`, and `gpt-4o` all start with `gpt-4` but use o200k_base. This - // mirrors OpenAI's authoritative `MODEL_PREFIX_TO_ENCODING` table - // (tiktoken/model.py): o200k_base covers `gpt-4o`, `gpt-4.1`, `gpt-4.5`, - // `chatgpt-4o-latest`, and the `o1`/`o3`/`o4` reasoning families. if model.starts_with("gpt-4o") || model.starts_with("gpt-4.1") || model.starts_with("gpt-4.5") @@ -200,13 +159,9 @@ pub fn encoding_for_model(model: &str) -> Option { { return Some(Encoding::O200kBase); } - // `cl100k_base` family. if model.starts_with("gpt-4") || model.starts_with("gpt-3.5") || model.starts_with("gpt-35") { return Some(Encoding::Cl100kBase); } - // `p50k_base` family. Per OpenAI's `encoding_for_model`: - // `text-davinci-003`, `text-davinci-002`, `code-davinci-*`, and - // `code-cushman-*` use p50k_base. if model.starts_with("text-davinci-003") || model.starts_with("text-davinci-002") || model.starts_with("code-davinci") @@ -214,10 +169,6 @@ pub fn encoding_for_model(model: &str) -> Option { { return Some(Encoding::P50kBase); } - // `r50k_base` family. The original GPT-3 base models (`text-davinci-001`, - // `text-curie-001`, `text-babbage-001`, `text-ada-001`, and their bare - // names `davinci` / `curie` / `babbage` / `ada`) map to r50k_base, not - // p50k_base. if model.starts_with("text-davinci-001") || model.starts_with("text-curie") || model.starts_with("text-babbage") @@ -238,7 +189,6 @@ mod tests { #[test] fn encoding_for_model_maps_known_families() { - // o200k_base family. assert_eq!(encoding_for_model("gpt-4o"), Some(Encoding::O200kBase)); assert_eq!( encoding_for_model("gpt-4o-2024-08-06"), @@ -247,10 +197,6 @@ mod tests { assert_eq!(encoding_for_model("o1-preview"), Some(Encoding::O200kBase)); assert_eq!(encoding_for_model("o1-mini"), Some(Encoding::O200kBase)); assert_eq!(encoding_for_model("o3-mini"), Some(Encoding::O200kBase)); - // o200k_base: current OpenAI models whose ids start with `gpt-4` but - // use o200k_base. These MUST be matched before the cl100k `gpt-4` branch - // below -- a regression here silently mis-tokenizes the most common - // production models on `/v1/messages/count_tokens`. assert_eq!(encoding_for_model("gpt-4.1"), Some(Encoding::O200kBase)); assert_eq!( encoding_for_model("gpt-4.1-mini"), @@ -265,8 +211,6 @@ mod tests { encoding_for_model("chatgpt-4o-latest"), Some(Encoding::O200kBase) ); - // cl100k_base family. Note: must NOT match the gpt-4o/gpt-4.1/gpt-4.5 - // prefixes above. assert_eq!(encoding_for_model("gpt-4"), Some(Encoding::Cl100kBase)); assert_eq!( encoding_for_model("gpt-4-turbo"), @@ -276,7 +220,6 @@ mod tests { encoding_for_model("gpt-3.5-turbo"), Some(Encoding::Cl100kBase) ); - // p50k_base family. Only -003/-002 and code-*; NOT -001 or bare names. assert_eq!( encoding_for_model("text-davinci-003"), Some(Encoding::P50kBase) @@ -293,7 +236,6 @@ mod tests { encoding_for_model("code-cushman-001"), Some(Encoding::P50kBase) ); - // r50k_base family: original GPT-3 base models. assert_eq!( encoding_for_model("text-davinci-001"), Some(Encoding::R50kBase) @@ -309,7 +251,6 @@ mod tests { assert_eq!(encoding_for_model("text-ada-001"), Some(Encoding::R50kBase)); assert_eq!(encoding_for_model("davinci"), Some(Encoding::R50kBase)); assert_eq!(encoding_for_model("curie"), Some(Encoding::R50kBase)); - // Unknown -> None (heuristic). assert_eq!(encoding_for_model("claude-sonnet-4-6"), None); assert_eq!(encoding_for_model("gemini-2.5-pro"), None); assert_eq!( @@ -322,56 +263,34 @@ mod tests { fn heuristic_count_basic() { assert_eq!(heuristic_count(""), 0); assert_eq!(heuristic_count("a"), 1); - assert_eq!(heuristic_count("hello"), 1); // 5/4 = 1 - assert_eq!(heuristic_count(&"a".repeat(20)), 5); // 20/4 = 5 - // CJK: 4 code points -> 1 token. + assert_eq!(heuristic_count("hello"), 1); + assert_eq!(heuristic_count(&"a".repeat(20)), 5); assert_eq!(heuristic_count("\u{4F60}\u{597D}\u{4E16}\u{754C}"), 1); } - // -- tok-01: catch_unwind fallback -------------------------------------- - - /// The heuristic fallback used when BPE encoding panics is exercised - /// directly here. It must be infallible and never panic. #[test] fn heuristic_fallback_is_infallible() { assert!(heuristic_count("hello") > 0); assert_eq!(heuristic_count(""), 0); - // Pathological-looking input must still produce a sane count. let nasty = "\u{0}".repeat(1000); assert!(heuristic_count(&nasty) > 0); } - /// `TiktokenTokenizer::count_tokens` must never panic on a real BPE, - /// even on unusual input. This exercises the catch_unwind wrapper on the - /// happy path (proving the wrapper itself does not interfere). #[test] fn tiktoken_count_tokens_never_panics() { let bpe = Arc::new(tiktoken_rs::cl100k_base().expect("cl100k_base loads")); let tok = TiktokenTokenizer::new(bpe); - // Normal input. let n = tok.count_tokens("The quick brown fox"); assert!(n > 0, "real BPE count should be positive"); - // Empty input. assert_eq!(tok.count_tokens(""), 0); - // Unusual but valid UTF-8 (control chars, long repetition). The public - // API contract is: return a usize, do not panic. let weird = format!("{}{}", "\u{0}\u{1}".repeat(500), "hello world"); let _weird_count: usize = tok.count_tokens(&weird); } - /// If a panic is somehow triggered inside `count_tokens`, it is caught and - /// the heuristic value is returned. We cannot easily force fancy-regex's - /// internal unwrap to fire deterministically, but we can prove the catch - /// path works by panicking inside a stand-in closure that mirrors the - /// wrapper's structure. This guards the catch_unwind plumbing against - /// regressions (e.g. someone removing it). #[test] fn catch_unwind_returns_fallback_on_panic_shape() { - // Mirror of the wrapper's payload-extraction logic, fed a real panic. let payload = std::panic::catch_unwind(|| panic!("boom")); assert!(payload.is_err(), "panic should be captured"); - // The wrapper turns Err(payload) into heuristic_count(text); verify the - // extraction branch compiles and behaves for both payload kinds. let err = payload.unwrap_err(); let msg_static = err .downcast_ref::<&'static str>() @@ -379,7 +298,6 @@ mod tests { .or_else(|| err.downcast_ref::().map(String::as_str)) .unwrap_or(""); assert_eq!(msg_static, "boom"); - // And the heuristic value the wrapper would have returned. assert!(heuristic_count("hello") > 0); } } diff --git a/crates/llm-proxy-protocol/src/anthropic.rs b/crates/llm-proxy-protocol/src/anthropic.rs index 32fdaac..9922249 100644 --- a/crates/llm-proxy-protocol/src/anthropic.rs +++ b/crates/llm-proxy-protocol/src/anthropic.rs @@ -1,16 +1,9 @@ //! Anthropic API types for the Messages API. //! //! Reference: -//! -//! This module contains request/response types for the Anthropic Messages API, -//! including support for streaming via Server-Sent Events. use serde::{Deserialize, Serialize, Serializer}; -// --------------------------------------------------------------------------- -// MessageRequest -// --------------------------------------------------------------------------- - /// A request to the Anthropic Messages API. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -43,35 +36,25 @@ pub struct MessageRequest { /// Extended thinking configuration. #[serde(default, skip_serializing_if = "Option::is_none")] pub thinking: Option, - /// Controls which (if any) tool the model must call. - /// - /// Accepts `{"type": "auto"}`, `{"type": "any"}`, or - /// `{"type": "tool", "name": "..."}` as defined by the Anthropic - /// Messages API. + /// Controls which (if any) tool the model must call. Accepts + /// `{"type": "auto"}`, `{"type": "any"}`, or `{"type": "tool", "name": "..."}` + /// as defined by the Anthropic Messages API. #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_choice: Option, } impl MessageRequest { - /// Extracts the system prompt text from the `system` field. - /// - /// Anthropic accepts `system` as either a plain string or an array of - /// [`SystemContentBlock`]: - /// - /// ```json - /// "system": "You are helpful" - /// "system": [{"type":"text","text":"You are helpful","cache_control":...}] - /// ``` + /// Extracts the system prompt text from the `system` field, which may be + /// a plain string or an array of [`SystemContentBlock`]s; array text blocks + /// are joined with `\n`, and anything unparseable falls back to raw JSON. #[must_use] pub fn system_text(&self) -> String { let Some(value) = &self.system else { return String::new(); }; - // Try plain string first. if let Some(s) = value.as_str() { return s.to_owned(); } - // Try array of SystemContentBlock. if let Some(arr) = value.as_array() { let parts: Vec = arr .iter() @@ -83,14 +66,11 @@ impl MessageRequest { return parts.join("\n"); } } - // Fallback: raw JSON string. value.to_string() } - /// Validates that required fields are present. - /// - /// Returns `Ok(())` if all required fields are present, or a - /// `ProtocolError::InvalidRequest` with a description of the problem. + /// Validates that required fields are present; returns + /// `ProtocolError::InvalidRequest` describing the first problem otherwise. pub fn validate(&self) -> Result<(), crate::client::ProtocolError> { if self.model.is_empty() { return Err(crate::client::ProtocolError::InvalidRequest( @@ -112,10 +92,6 @@ impl MessageRequest { } } -// --------------------------------------------------------------------------- -// SystemContentBlock -// --------------------------------------------------------------------------- - /// A content block inside the `system` array. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -131,10 +107,6 @@ pub struct SystemContentBlock { pub cache_control: Option, } -// --------------------------------------------------------------------------- -// CacheControl -// --------------------------------------------------------------------------- - /// Cache control directives for prompt caching. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -144,20 +116,9 @@ pub struct CacheControl { pub r#type: String, } -// --------------------------------------------------------------------------- -// Metadata -// --------------------------------------------------------------------------- - -/// Optional metadata attached to a [`MessageRequest`]. -/// -/// Unknown metadata fields beyond `user_id` are preserved in the `extra` -/// flattened map so they can be forwarded to `RequestMetadata.raw` during -/// decode rather than silently dropped. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because this struct -/// uses `#[serde(flatten)]` to capture unknown fields. The flattened `extra` -/// map may grow large with arbitrary client metadata; callers that need to -/// bound payload size should check `extra.len()` before forwarding. +/// Optional metadata attached to a [`MessageRequest`]; unknown fields are +/// preserved in `extra` (via `#[serde(flatten)]`, so `deny_unknown_fields` +/// is not set) and forwarded to `RequestMetadata.raw`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Metadata { /// An external identifier for the end-user. @@ -168,20 +129,9 @@ pub struct Metadata { pub extra: serde_json::Map, } -// --------------------------------------------------------------------------- -// Message -// --------------------------------------------------------------------------- - -/// A single message in the conversation. -/// -/// The `content` field can be either a plain string or an array of -/// [`ContentBlock`]. -/// -/// Note: `deny_unknown_fields` is intentionally omitted here even though the -/// parent `MessageRequest` has it. New Anthropic message-level fields (e.g. a -/// future `metadata` field on individual messages) should be tolerated rather -/// than causing deserialization failures at this layer. The parent struct's -/// `deny_unknown_fields` already guards against truly unknown top-level fields. +/// A single message in the conversation; `content` is either a plain string +/// or an array of [`ContentBlock`]s. Unknown fields are tolerated (no +/// `deny_unknown_fields`) so future Anthropic fields decode cleanly. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] pub struct Message { @@ -192,21 +142,14 @@ pub struct Message { } impl Message { - /// Parses the message `content` into a list of [`ContentBlock`]. - /// - /// Handles both a plain string and an array-of-blocks representation: - /// - /// ```json - /// "content": "hello" - /// "content": [{"type":"text","text":"hello"}] - /// ``` + /// Parses the message `content` into a list of [`ContentBlock`]s; accepts + /// both `"content": "text"` and `"content": [block, ...]` wire shapes. #[must_use] #[expect(deprecated, reason = "constructs ContentBlock with output: None")] pub fn content_blocks(&self) -> Vec { if self.content.is_null() { return Vec::new(); } - // Try plain string first. if let Some(s) = self.content.as_str() { return vec![ContentBlock { r#type: "text".to_owned(), @@ -225,7 +168,6 @@ impl Message { data: None, }]; } - // Try array of content blocks. if let Some(arr) = self.content.as_array() { let mut blocks = Vec::with_capacity(arr.len()); for item in arr { @@ -249,32 +191,9 @@ impl Message { } } -// --------------------------------------------------------------------------- -// ContentBlock -// --------------------------------------------------------------------------- - -/// A block of content within a message. -/// -/// The `r#type` field determines which other fields are populated: -/// -/// | Type | Populated fields | -/// |----------------------|---------------------------------------------------| -/// | `"text"` | `text` | -/// | `"tool_use"` | `id`, `name`, `input` | -/// | `"tool_result"` | `tool_use_id`, `content`, `is_error` | -/// | `"thinking"` | `thinking`, `signature` | -/// | `"redacted_thinking"`| `data` | -/// | `"image"` | `source` | -/// -/// Note: `deny_unknown_fields` is intentionally omitted because this struct -/// receives polymorphic content blocks from external JSON. Unknown block types -/// must be tolerated so that future Anthropic API additions do not break -/// deserialization. The `Serialize` impl handles unknown types via the -/// catch-all arm. -/// -/// Note: `#[non_exhaustive]` is applied to allow future field additions -/// without breaking downstream code, since this struct represents an external -/// protocol and all fields are public. +/// A block of content within a message; `r#type` selects the populated +/// fields (`"text"`, `"tool_use"`, `"tool_result"`, `"thinking"`, +/// `"redacted_thinking"`, `"image"`). Unknown types must be tolerated. #[derive(Debug, Clone, Deserialize)] #[non_exhaustive] pub struct ContentBlock { @@ -458,10 +377,8 @@ impl ContentBlock { } } - /// Returns the appropriate tool ID for this block. - /// - /// For `"tool_result"` blocks returns `tool_use_id`; for all others - /// (notably `"tool_use"`) returns `id`. + /// Returns the tool ID for this block: `tool_use_id` for `"tool_result"`, + /// `id` for every other type. #[must_use] pub fn get_tool_id(&self) -> &str { if self.r#type == "tool_result" { @@ -471,13 +388,10 @@ impl ContentBlock { } } - /// Extracts text from a `"tool_result"` block's `content` field. - /// - /// The `content` field can be a plain string or an array of content - /// blocks. Falls back to the deprecated `output` field. + /// Extracts text from a `"tool_result"` block's `content` field (string + /// or array of text blocks), falling back to the deprecated `output`. #[must_use] pub fn text_content(&self) -> String { - // Try the content field first. if let Some(ref val) = self.content { if let Some(s) = val.as_str() { return s.to_owned(); @@ -494,7 +408,6 @@ impl ContentBlock { } } } - // Fallback to the deprecated output field. #[expect( deprecated, reason = "reads the deprecated `output` field as a text fallback" @@ -509,14 +422,6 @@ impl ContentBlock { } } -/// Serializes an unknown / future content block type by emitting every non-`None` -/// field. -/// -/// This clones every field to construct an intermediate `AllFields` struct. -/// The clone cost is acceptable because (a) unknown block types are rare, -/// and (b) the approach keeps the Serialize impl maintainable by leveraging -/// derive rather than a manual `SerializeMap` that conditionally includes -/// each field. fn serialize_unknown_block(block: &ContentBlock, serializer: S) -> Result where S: Serializer, @@ -637,11 +542,6 @@ impl Serialize for ContentBlock { map.end() } "image" => { - // When `source` is `None` we emit an empty ImageSource placeholder. - // This produces structurally valid but semantically meaningless JSON. - // The `new_image` constructor requires a source, so this branch is only - // reachable for blocks constructed with Default or deserialized from - // malformed input. let mut map = serializer.serialize_map(Some(2))?; map.serialize_entry("type", &self.r#type)?; map.serialize_entry( @@ -659,10 +559,6 @@ impl Serialize for ContentBlock { } } -// --------------------------------------------------------------------------- -// ImageSource -// --------------------------------------------------------------------------- - /// An image source embedded in a content block. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -676,10 +572,6 @@ pub struct ImageSource { pub data: String, } -// --------------------------------------------------------------------------- -// Tool -// --------------------------------------------------------------------------- - /// A tool definition for function calling. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -693,10 +585,6 @@ pub struct Tool { pub input_schema: serde_json::Value, } -// --------------------------------------------------------------------------- -// ToolResult -// --------------------------------------------------------------------------- - /// The result of executing a tool call. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -710,22 +598,8 @@ pub struct ToolResult { pub is_error: Option, } -// --------------------------------------------------------------------------- -// MessageResponse -// --------------------------------------------------------------------------- - -/// A response from the Anthropic Messages API. -/// -/// This struct is both constructed for client-facing output AND deserialized -/// from upstream Anthropic SSE events (specifically `message_start` events -/// contain a full `MessageResponse` snapshot). `Deserialize` is required for -/// the SSE decode path; `Serialize` is required for the encode path. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because this is an -/// outbound-only type -- the proxy constructs it internally and serializes -/// it to clients. It is never deserialized from external input. Omitting -/// `deny_unknown_fields` is consistent with other outbound response types -/// (`ChatCompletionResponse`, `ResponsesResponse`, `GeminiResponse`). +/// A response from the Anthropic Messages API, built by the encode path +/// and deserialized from upstream `message_start` SSE snapshots. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] pub struct MessageResponse { @@ -750,15 +624,8 @@ pub struct MessageResponse { pub usage: Usage, } -// --------------------------------------------------------------------------- -// Usage -// --------------------------------------------------------------------------- - -/// Token usage statistics returned with every response. -/// -/// Uses `i32` for token counts for consistency with the Anthropic API wire -/// format (JSON integers) and the core `Usage` type. The ~2.1 billion limit -/// is well beyond any realistic single request. +/// Token usage statistics returned with every response; counts are `i32`, +/// matching the Anthropic JSON integer wire format. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] pub struct Usage { @@ -774,10 +641,6 @@ pub struct Usage { pub cache_read_input_tokens: Option, } -// --------------------------------------------------------------------------- -// ContentBlockDelta -// --------------------------------------------------------------------------- - /// A streaming delta event for a content block. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ContentBlockDelta { @@ -790,19 +653,13 @@ pub struct ContentBlockDelta { pub delta: Delta, } -// --------------------------------------------------------------------------- -// Delta -// --------------------------------------------------------------------------- - /// A partial update in a streaming response. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] pub struct Delta { /// Delta type (e.g. `"text_delta"`, `"thinking_delta"`, `"input_json_delta"`). - /// - /// Kept as `Option` with `#[serde(default)]` for forward-compatibility - /// with future Anthropic delta types. In practice the Anthropic API always - /// includes this field; `None` indicates a malformed or unknown delta. + /// `Option` with `#[serde(default)]` for forward compatibility; `None` + /// indicates a malformed or unknown delta. #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")] pub r#type: Option, /// Partial text content. @@ -822,10 +679,6 @@ pub struct Delta { pub stop_sequence: Option, } -// --------------------------------------------------------------------------- -// MessageEvent -// --------------------------------------------------------------------------- - /// A Server-Sent Event from the Anthropic streaming API. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] @@ -854,10 +707,6 @@ pub struct MessageEvent { pub error: Option, } -// --------------------------------------------------------------------------- -// ApiError -// --------------------------------------------------------------------------- - /// An error returned by the Anthropic API. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -869,10 +718,6 @@ pub struct ApiError { pub message: String, } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] #[expect( deprecated, @@ -881,8 +726,6 @@ pub struct ApiError { mod tests { use super::*; - // -- system_text -------------------------------------------------------- - #[test] fn system_text_none_when_absent() { let req = MessageRequest { @@ -950,8 +793,6 @@ mod tests { assert_eq!(req.system_text(), "You are helpful. \nBe concise."); } - // -- content_blocks ----------------------------------------------------- - #[test] fn content_blocks_from_string() { let msg = Message { @@ -992,8 +833,6 @@ mod tests { assert!(msg.content_blocks().is_empty()); } - // -- validate ----------------------------------------------------------- - #[test] fn validate_ok() { let req = MessageRequest { @@ -1060,8 +899,6 @@ mod tests { ); } - // -- ContentBlock custom serialize -------------------------------------- - #[test] fn serialize_text_block() { let block = ContentBlock { @@ -1189,8 +1026,6 @@ mod tests { assert_eq!(json.as_object().unwrap().len(), 3); } - // -- text_content ------------------------------------------------------- - #[test] fn text_content_from_string_content() { let block = ContentBlock { @@ -1258,8 +1093,6 @@ mod tests { assert_eq!(block.text_content(), "legacy output"); } - // -- get_tool_id -------------------------------------------------------- - #[test] fn get_tool_id_tool_result() { let block = ContentBlock { @@ -1280,8 +1113,6 @@ mod tests { assert_eq!(block.get_tool_id(), "tu_99"); } - // -- helper to build a minimal block for tests -------------------------- - fn make_minimal_block() -> ContentBlock { ContentBlock { r#type: String::new(), @@ -1301,10 +1132,6 @@ mod tests { } } - // -- text_content: empty string content ----------------------------------- - - /// Characterization: `text_content()` returns an empty string when - /// `content` is `Some(Value::String(""))`. #[test] fn text_content_from_empty_string_content() { let block = ContentBlock { diff --git a/crates/llm-proxy-protocol/src/client/anthropic.rs b/crates/llm-proxy-protocol/src/client/anthropic.rs index b8fbaba..2332230 100644 --- a/crates/llm-proxy-protocol/src/client/anthropic.rs +++ b/crates/llm-proxy-protocol/src/client/anthropic.rs @@ -1,20 +1,5 @@ -//! Anthropic Messages client protocol adapter. -//! -//! Translates between Anthropic Messages API wire format and the normalised -//! core protocol types. This adapter only knows about Anthropic Messages -//! and the core types -- it never imports provider, server, or config code. -//! -//! ## Scope guardrails -//! -//! This module must **not** import or call: -//! - `llm_proxy_provider` -//! - `llm_proxy_server` -//! - core config or routing modules -//! - endpoint classification helpers -//! - scenario or fallback code -//! -//! The `transformer/` module was removed in Phase 11; the guard test at the -//! bottom of this file remains as a regression safety net. +//! Anthropic Messages client protocol adapter: translates between the +//! Anthropic wire format and the normalised core protocol types. use crate::anthropic::{ self, ApiError, CacheControl as AnthropicCacheControl, ContentBlock, Delta, Message, @@ -27,10 +12,6 @@ use crate::core::{ ModelRef, ProviderHints, RequestMetadata, SamplingOptions, StopReason, Usage, }; -// --------------------------------------------------------------------------- -// decode_request -// --------------------------------------------------------------------------- - /// Decode an Anthropic [`MessageRequest`] into a normalised [`CoreRequest`]. pub fn decode_request(req: MessageRequest) -> Result { if req.model.is_empty() { @@ -116,7 +97,6 @@ fn decode_system(system: Option) -> Vec { } Some(serde_json::Value::Array(arr)) => { let mut result = Vec::new(); - // We own `system`, so iterate owned array items without cloning. for item in arr { let Ok(block) = serde_json::from_value::(item) else { tracing::warn!( @@ -125,12 +105,7 @@ fn decode_system(system: Option) -> Vec { continue; }; if block.r#type != "text" { - // Non-text system blocks (e.g. future image blocks) are not - // yet supported -- log a warning so they are not silently lost. - // Truncate block_type to limit log output from client input. let bt = block.r#type.as_str(); - // Char-boundary-safe truncation: avoids panics on - // multi-byte UTF-8 characters in client-supplied data. let truncated = crate::util::truncate_str_safe(bt, 64); tracing::warn!( block_type = truncated, @@ -147,8 +122,6 @@ fn decode_system(system: Option) -> Vec { } result } - // Non-string, non-array system values (e.g. a bare number) cannot be - // represented as system content -- return empty rather than erroring. _ => Vec::new(), } } @@ -171,23 +144,15 @@ fn decode_message(msg: &Message) -> Result { Ok(CoreMessage { role, content }) } -/// Decodes the `content` array of a `tool_result` block. -/// -/// Inner blocks that fail to deserialize, or whose type is unknown, are logged -/// and skipped rather than failing the whole `tool_result`. If nothing survives -/// decoding, an empty text block is synthesized so the tool result is not lost. fn decode_tool_result_blocks(arr: &[serde_json::Value]) -> Result, ProtocolError> { let mut blocks = Vec::with_capacity(arr.len()); for item in arr { let Ok(cb) = serde_json::from_value::(item.clone()) else { continue; }; - // Recursively decode each inner block. Errors from unknown block types - // are logged and skipped rather than failing the entire tool_result decode. match decode_content_block(cb) { Ok(core) => blocks.push(core), Err(ProtocolError::Decode(msg)) => { - // Char-boundary-safe truncation for log output. let truncated = crate::util::truncate_str_safe(&msg, 64); tracing::warn!( decode_error = truncated, @@ -210,21 +175,17 @@ fn decode_tool_result_blocks(arr: &[serde_json::Value]) -> Result Result { let is_error = block.is_error.unwrap_or(false); let inner_content = if let Some(ref content_val) = block.content { match content_val { - // Plain string content. serde_json::Value::String(s) => { vec![CoreContent::Text { text: s.clone(), cache: None, }] } - // Array of content blocks -- iterate and decode each one. serde_json::Value::Array(arr) => decode_tool_result_blocks(arr)?, - // Fallback: try text_content() for any other shape. _ => { let text = block.text_content(); if text.is_empty() { @@ -235,8 +196,6 @@ fn decode_tool_result_block(block: ContentBlock) -> Result Result { - // Unknown block types cannot be safely represented -- return an error - // so the caller knows data was lost rather than silently creating - // an empty text block. Truncate the block_type to limit log output - // from potentially malicious client input. Char-boundary-safe - // truncation avoids panics on multi-byte UTF-8. let truncated = crate::util::truncate_str_safe(other, 64); tracing::warn!( block_type = truncated, @@ -353,21 +307,9 @@ fn decode_tool_choice(value: serde_json::Value) -> CoreToolChoice { } } -// --------------------------------------------------------------------------- -// encode_response -// --------------------------------------------------------------------------- - -/// Encode a [`CoreResponse`] into an Anthropic [`MessageResponse`]. -/// -/// Content blocks that cannot be represented in the Anthropic wire format are -/// handled according to the plan's two-path rule: -/// -/// - **Safe to drop** (Document, Audio, Video): the block is logged with -/// `tracing::warn!` and silently omitted. `encode_content_block` signals this -/// via `Err(ProtocolError::EncodeSkippable(..))`. -/// -/// - **Must propagate** (Refusal): returns `Err(ProtocolError::Encode(..))` -/// because silently dropping a refusal would change the response semantics. +/// Encode a [`CoreResponse`] into an Anthropic [`MessageResponse`]. Blocks +/// the format cannot represent are dropped with a warning +/// (`EncodeSkippable`); a `CoreContent::Refusal` is a hard `Encode` error. pub fn encode_response(resp: CoreResponse) -> Result { if !resp.provider_meta.is_empty() { tracing::warn!( @@ -385,10 +327,7 @@ pub fn encode_response(resp: CoreResponse) -> Result content.push(encoded), - Err(ProtocolError::EncodeSkippable(_msg)) => { - // Safe-to-drop block type for this protocol -- skip it. - // The warning was already logged in encode_content_block. - } + Err(ProtocolError::EncodeSkippable(_msg)) => {} Err(other) => return Err(other), } } @@ -465,9 +404,6 @@ fn encode_content_block(content: CoreContent) -> Result { - // Document blocks are not supported in Anthropic responses. The - // omission is safe (the data was never user-visible in this context), - // so we signal via EncodeSkippable so encode_response can skip it. tracing::warn!( "Anthropic Messages protocol does not natively support document blocks in responses; omitting" ); @@ -492,8 +428,6 @@ fn encode_content_block(content: CoreContent) -> Result { - // Anthropic natively supports redacted_thinking -- encode as proper - // redacted_thinking content block type. let data_str = match data { serde_json::Value::String(s) => s, other => other.to_string(), @@ -501,10 +435,6 @@ fn encode_content_block(content: CoreContent) -> Result { - // Anthropic does not have a native refusal block type. Per the plan, - // return ProtocolError::Encode rather than silently converting to text. - // Log only the length to avoid writing potentially sensitive refusal - // content into log output. tracing::warn!( refusal_text_len = text.len(), "Anthropic Messages protocol has no refusal field; cannot encode" @@ -516,16 +446,6 @@ fn encode_content_block(content: CoreContent) -> Result `"stop_sequence"` (natively supported by Anthropic) -/// - `Refusal` -> `"end_turn"` with warning (Anthropic has no refusal stop reason) -/// - `Error` -> `"end_turn"` with warning (Anthropic has no error stop reason) -/// - `Unknown` -> `"end_turn"` (safe default) fn encode_stop_reason(reason: StopReason) -> String { match reason { StopReason::EndTurn => "end_turn".to_owned(), @@ -571,11 +491,6 @@ fn encode_usage(usage: &Usage) -> anthropic::Usage { } } -/// Builds the `content_block` payload for a `content_block_start` event. -/// -/// Returns `None` for kinds that must not produce a `content_block_start`: -/// `ToolUse` is emitted by `ToolCallStart`, which carries the actual id/name, -/// and the remaining kinds are unsupported. fn encode_content_start_block(index: usize, kind: ContentKind) -> Option { match kind { ContentKind::Text => Some(ContentBlock::new_text(String::new())), @@ -588,9 +503,6 @@ fn encode_content_start_block(index: usize, kind: ContentKind) -> Option { - // ContentStart for unsupported kinds (Image, Document, Audio, - // Video, ToolResult, Refusal) -- skip entirely rather than - // emitting a misleading text block start. tracing::warn!( kind = ?other, "unsupported ContentKind in Anthropic stream encode; skipping content_block_start" @@ -600,24 +512,7 @@ fn encode_content_start_block(index: usize, kind: ContentKind) -> Option MessageEvent { - // Length cap: truncate the error message to 1024 characters to prevent - // excessively large SSE payloads from upstream errors. Char-boundary-safe - // truncation avoids panics on multi-byte UTF-8 characters in - // upstream/provider-derived error messages (consistent with the other - // truncation sites in this file). let msg = error.message(); let capped_msg = if msg.len() > 1024 { tracing::warn!( @@ -642,41 +537,22 @@ fn encode_error_event(error: &CoreStreamError) -> MessageEvent { } } -// --------------------------------------------------------------------------- -// StreamEncoder -// --------------------------------------------------------------------------- - /// Stateful encoder that converts [`CoreEvent`] values into Anthropic SSE -/// [`MessageEvent`] values. -/// -/// The encoder buffers the latest `UsageDelta` and coalesces it with the -/// terminal `MessageStop` into a single `message_delta` event carrying both -/// `usage` and `delta.stop_reason`/`stop_sequence`, followed by one -/// `message_stop`. This matches native Anthropic behavior. +/// [`MessageEvent`] values, buffering the latest `UsageDelta` and coalescing +/// it into the terminal `message_delta` + `message_stop` pair. #[derive(Debug)] pub struct StreamEncoder { - /// The message ID for this stream. msg_id: String, - /// The model name to report. model: String, - /// Buffered usage from the latest UsageDelta. pending_usage: Option, - /// Whether message_start has been emitted. started: bool, - /// Whether the terminal message_delta + message_stop have been emitted. finished: bool, - /// Indices of text/thinking blocks that received content_block_start - /// but have not yet been closed. Native Anthropic always emits - /// content_block_stop after every block; we synthesize stops for - /// text/thinking at MessageStop since CoreEvent has no ContentStop. open_text_blocks: Vec, } impl StreamEncoder { - /// Create a new stream encoder. - /// - /// `msg_id` and `model` are typically taken from the first - /// `CoreEvent::MessageStart`. + /// Create a new stream encoder; `msg_id` and `model` are typically taken + /// from the first `CoreEvent::MessageStart`. pub fn new(msg_id: String, model: String) -> Self { Self { msg_id, @@ -688,7 +564,6 @@ impl StreamEncoder { } } - /// Records the stream's id/model and builds the `message_start` event. fn encode_message_start(&mut self, id: Option, model: String) -> MessageEvent { if let Some(id) = id { self.msg_id = id; @@ -721,12 +596,6 @@ impl StreamEncoder { } } - /// Emits the terminal `message_delta` + `message_stop` pair, first closing - /// any text/thinking blocks that were opened but never closed. - /// - /// Native Anthropic emits `content_block_stop` for every block; without the - /// synthesized stops, strict clients that pair start/stop events mis-frame - /// the stream. Tool blocks are closed by `ToolCallStop`. fn encode_message_stop( &mut self, stop_reason: StopReason, @@ -785,6 +654,13 @@ impl StreamEncoder { } /// Encode a single [`CoreEvent`] into zero or more Anthropic [`MessageEvent`]s. + #[expect( + clippy::too_many_lines, + reason = "one match arm per CoreEvent variant: the Anthropic Messages SSE \ + wire shape is a single translation table, and splitting the arms \ + into helpers would separate each event's field mapping from the \ + frame layout it produces" + )] pub fn encode_event(&mut self, event: CoreEvent) -> Result, ProtocolError> { if self.finished { tracing::trace!( @@ -803,7 +679,6 @@ impl StreamEncoder { CoreEvent::ContentStart { index, kind } => { let Some(block) = encode_content_start_block(index, kind) else { - // Unsupported or duplicate-producing kind: emit nothing. return Ok(events); }; events.push(MessageEvent { @@ -904,7 +779,6 @@ impl StreamEncoder { } CoreEvent::UsageDelta { usage } => { - // Buffer usage; it will be emitted with the terminal event. self.pending_usage = Some(encode_usage(&usage)); } @@ -935,23 +809,15 @@ impl StreamEncoder { Ok(events) } - /// Mark the encoder as finished without emitting synthetic terminal events. - /// - /// This is used after an in-band error event has been emitted: the error - /// event itself is the terminal event, so subsequent `finish()` calls - /// return an empty vec rather than a synthetic message_delta/message_stop pair. + /// Mark the encoder as finished without emitting synthetic terminal + /// events; later `finish()` calls return an empty vec. pub fn mark_finished(&mut self) { self.finished = true; } - /// Flush any remaining buffered events (e.g. if the stream was terminated - /// without a `MessageStop`). - /// - /// When the stream terminated abnormally (no explicit `MessageStop`), the - /// stop reason is set to `"end_turn"` with a `tracing::warn`, as this is the - /// least misleading stop reason for an interrupted stream. Always emits a - /// `message_delta` + `message_stop` pair so the Anthropic client receives a - /// proper terminal sequence. + /// Flush remaining buffered events when the stream ended without a + /// `MessageStop`: emits a synthetic `message_delta` (stop reason + /// `"end_turn"`, with a warning) + `message_stop` pair. pub fn finish(&mut self) -> Result, ProtocolError> { if self.finished { return Ok(Vec::new()); @@ -980,10 +846,6 @@ impl StreamEncoder { text: None, thinking: None, partial_json: None, - // Use end_turn as the least misleading stop reason for abnormal - // termination. max_tokens has a specific semantic meaning that - // would cause clients to incorrectly believe the model ran out - // of tokens. stop_reason: Some("end_turn".to_owned()), stop_sequence: None, }), @@ -1016,18 +878,12 @@ fn encode_error_kind(kind: &CoreStreamErrorKind) -> String { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; use crate::anthropic::Tool; use crate::core::UsageProvenance; - // -- helpers ------------------------------------------------------------ - fn make_anthropic_request() -> MessageRequest { MessageRequest { model: "claude-sonnet-4-20250514".into(), @@ -1047,8 +903,6 @@ mod tests { } } - // -- decode tests ------------------------------------------------------- - #[test] fn plain_text_request_decodes_to_core() { let req = make_anthropic_request(); @@ -1319,8 +1173,6 @@ mod tests { assert!(!core2.stream); } - // -- encode tests ------------------------------------------------------- - #[test] fn core_text_response_encodes_to_route_response() { let resp = CoreResponse { @@ -1468,8 +1320,6 @@ mod tests { assert!(out.id.starts_with("msg_")); } - // -- streaming encode tests --------------------------------------------- - #[test] fn streaming_text_event_mapping() { let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); @@ -1552,12 +1402,8 @@ mod tests { #[test] fn streaming_content_start_tool_use_defers_to_tool_call_start() { - // ContentStart { kind: ToolUse } must NOT emit content_block_start - // because ToolCallStart carries the actual id/name and will emit it. - // Emitting both would produce a duplicate content_block_start. let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); - // ContentStart for ToolUse should produce zero events. let content_start_events = enc .encode_event(CoreEvent::ContentStart { index: 0, @@ -1570,7 +1416,6 @@ mod tests { content_start_events.len() ); - // ToolCallStart should produce exactly one content_block_start. let tool_start_events = enc .encode_event(CoreEvent::ToolCallStart { index: 0, @@ -1593,10 +1438,6 @@ mod tests { #[test] fn streaming_text_block_emits_content_block_stop() { - // Text/thinking blocks receive content_block_start and content_block_delta - // but never content_block_stop (only ToolCallStop emits one). The encoder - // must synthesize a stop at MessageStop, before message_delta, so strict - // clients that pair start/stop events do not mis-frame the stream. let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); enc.encode_event(CoreEvent::MessageStart { id: None, @@ -1623,13 +1464,11 @@ mod tests { }) .unwrap(); - // A content_block_stop for index 0 must be present. assert!( stop.iter() .any(|e| e.r#type == "content_block_stop" && e.index == Some(0)) ); - // ...and it must precede the message_delta event. let md = stop .iter() .position(|e| e.r#type == "message_delta") @@ -1645,13 +1484,11 @@ mod tests { fn streaming_usage_coalesced_with_message_stop() { let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); - // Send UsageDelta first. enc.encode_event(CoreEvent::UsageDelta { usage: Usage::provider_reported(50, 100), }) .unwrap(); - // Then send MessageStop -- should coalesce. let events = enc .encode_event(CoreEvent::MessageStop { stop_reason: StopReason::EndTurn, @@ -1659,7 +1496,6 @@ mod tests { }) .unwrap(); - // Should produce exactly 2 events: message_delta (with usage + stop) and message_stop. assert_eq!(events.len(), 2); assert_eq!(events[0].r#type, "message_delta"); assert!(events[0].usage.is_some()); @@ -1734,7 +1570,6 @@ mod tests { let remaining = enc.finish().unwrap(); assert_eq!(remaining.len(), 2); assert_eq!(remaining[0].r#type, "message_delta"); - // Abnormal termination uses end_turn as the least misleading stop reason. assert_eq!( remaining[0].delta.as_ref().unwrap().stop_reason.as_deref(), Some("end_turn") @@ -1745,8 +1580,6 @@ mod tests { #[test] fn streaming_finish_without_pending_usage_still_emits_terminal() { let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); - // No UsageDelta sent -- finish() should still emit terminal events - // with synthetic zero usage. let remaining = enc.finish().unwrap(); assert_eq!(remaining.len(), 2); assert_eq!(remaining[0].r#type, "message_delta"); @@ -1777,14 +1610,8 @@ mod tests { assert!(events.is_empty()); } - // -- every CoreEvent variant handled ------------------------------------ - #[test] fn every_core_event_variant_maps_or_errors_intentionally() { - // Test all ContentKind variants for ContentStart. The stream encoder - // handles Text and Thinking explicitly (emitting content_block_start); - // ToolUse is deferred to ToolCallStart; other kinds are skipped with a - // warning (returning empty events, not an error). let content_kinds: Vec = vec![ ContentKind::Text, ContentKind::Thinking, @@ -1805,7 +1632,6 @@ mod tests { assert!(result.is_ok(), "ContentStart {:?} should not error", kind); } - // Test all other CoreEvent variants. let variants: Vec = vec![ CoreEvent::MessageStart { id: None, @@ -1851,8 +1677,6 @@ mod tests { } } - // -- unsupported content variants --------------------------------------- - #[test] fn encode_image_produces_image_block() { let content = CoreContent::Image { @@ -1911,7 +1735,6 @@ mod tests { #[test] fn encode_response_omits_unsupported_blocks() { - // encode_response filters out blocks that return ProtocolError::Encode. let resp = CoreResponse { id: Some("msg_123".into()), model: ModelRef { @@ -1934,13 +1757,10 @@ mod tests { cost: None, }; let out = encode_response(resp).unwrap(); - // Document block should be omitted; only text remains. assert_eq!(out.content.len(), 1); assert_eq!(out.content[0].r#type, "text"); } - // -- metadata / provider hints preserved -------------------------------- - #[test] fn metadata_raw_preserved_in_decode() { let mut req = make_anthropic_request(); @@ -2000,8 +1820,6 @@ mod tests { assert!(matches!(result.unwrap_err(), ProtocolError::Decode(_))); } - // -- tool definitions --------------------------------------------------- - #[test] fn tools_decoded_to_core() { let mut req = make_anthropic_request(); @@ -2016,17 +1834,8 @@ mod tests { assert_eq!(core.tools[0].description.as_deref(), Some("Get weather")); } - // -- source guard ------------------------------------------------------- - #[test] fn client_anthropic_does_not_import_forbidden_modules() { - // Source guard: this module must not import llm_proxy_provider, - // llm_proxy_server, core config/routing, endpoint classification, - // scenario/fallback code, or transformer/*. The crate dependency - // graph prevents most of these at compile time (those crates are not - // dependencies of llm-proxy-protocol). This runtime check provides - // a secondary defense by verifying that the module source does not - // contain forbidden import patterns in `use` statements. let source = include_str!("anthropic.rs"); for line in source.lines() { let trimmed = line.trim(); @@ -2058,8 +1867,6 @@ mod tests { } } - // -- edge-case tests ---------------------------------------------------- - #[test] fn decode_empty_model_returns_error() { let mut req = make_anthropic_request(); @@ -2082,7 +1889,6 @@ mod tests { #[test] fn decode_null_optional_fields() { - // Ensure that explicitly-null optional fields are handled gracefully. let mut req = make_anthropic_request(); req.system = None; req.temperature = None; @@ -2121,7 +1927,6 @@ mod tests { cost: None, }; let result = encode_response(resp); - // Refusal should propagate as ProtocolError::Encode (not skippable). assert!(matches!(result, Err(ProtocolError::Encode(_)))); } @@ -2153,7 +1958,6 @@ mod tests { cost: None, }; let out = encode_response(resp).unwrap(); - // Document is dropped; both text blocks are preserved. assert_eq!(out.content.len(), 2); assert_eq!(out.content[0].text.as_deref(), Some("hello")); assert_eq!(out.content[1].text.as_deref(), Some(" world")); @@ -2161,9 +1965,6 @@ mod tests { #[test] fn encode_error_does_not_leak_secret_in_message() { - // This test documents the trust boundary: the error message from - // CoreStreamError is forwarded verbatim into the client-facing SSE - // event. Sanitization must happen at CoreStreamError::new() time. let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); let secret_msg = "api_key=sk-12345-secret"; let events = enc @@ -2172,22 +1973,11 @@ mod tests { }) .unwrap(); assert_eq!(events[0].error.as_ref().unwrap().message, secret_msg); - // The defense-in-depth contract requires that provider adapters - // sanitize the message before constructing CoreStreamError. } #[test] fn encode_error_truncates_multibyte_message_without_panicking() { - // Regression: the >1024-byte branch of the CoreEvent::Error arm used to - // slice with `msg[..1024]`, which panics if byte 1024 lands inside a - // multi-byte UTF-8 code point. Build a message longer than 1024 bytes - // whose byte offset 1024 falls mid-code-point deterministically, then - // assert the encoder truncates at a char boundary instead of panicking. - // - // 1 ASCII char (1 byte) + 257 four-byte chars ('🚀') = 1 + 1028 = 1029 - // bytes. Byte 1024 is the 3rd byte of the 257th emoji (offset 1023 in - // the emoji run = 3 bytes into that code point), i.e. mid-code-point. - let emoji = "🚀"; // U+1F680, 4 bytes in UTF-8 + let emoji = "🚀"; assert_eq!(emoji.len(), 4); let msg: String = std::iter::once('_') .chain(std::iter::repeat_n(emoji, 257).flat_map(|s| s.chars())) @@ -2202,7 +1992,6 @@ mod tests { ); let mut enc = StreamEncoder::new("msg_1".into(), "m".into()); - // Must not panic. let events = enc .encode_event(CoreEvent::Error { error: CoreStreamError::new(CoreStreamErrorKind::Upstream, msg.clone()), @@ -2214,12 +2003,10 @@ mod tests { capped.len() <= 1024, "truncated message must not exceed the cap" ); - // The result is a valid char boundary (and valid UTF-8 by construction). assert!( msg.starts_with(capped.as_str()), "truncated message must be a prefix of the original" ); - // Sanity: we actually exercised the truncation branch. assert!(capped.len() < msg.len()); } diff --git a/crates/llm-proxy-protocol/src/client/mod.rs b/crates/llm-proxy-protocol/src/client/mod.rs index ae4854e..2296767 100644 --- a/crates/llm-proxy-protocol/src/client/mod.rs +++ b/crates/llm-proxy-protocol/src/client/mod.rs @@ -1,50 +1,12 @@ -//! Client protocol adapters. -//! -//! Each adapter translates between a specific client wire format and the -//! normalised core protocol types. Adapters do **not** import or call -//! provider code, server code, core config/routing, endpoint classification, -//! scenario/fallback code, or the transformer module. -//! -//! ## Design rules -//! -//! * Each adapter only knows: `its protocol <-> Core types`. -//! * No adapter may special-case another protocol's hints or wire format. -//! * HTTP framing (SSE, status codes) is a route concern, not an adapter concern. +//! Client protocol adapters: each translates between one client wire format +//! and the normalised core protocol types, never another adapter's format. pub mod anthropic; pub mod openai_chat; -// --------------------------------------------------------------------------- -// ProtocolError -// --------------------------------------------------------------------------- - -/// Errors returned by client protocol adapters. -/// -/// Every adapter returns this type so that route handlers can map errors to -/// appropriate HTTP responses without depending on adapter-specific error types. -/// -/// # HTTP status mapping -/// -/// Route handlers should translate `ProtocolError` variants into HTTP responses -/// as follows: -/// -/// - `InvalidRequest` -> 400 Bad Request -/// - `Encode` -> 500 Internal Server Error (or 502 Bad Gateway if the cause -/// is upstream data that cannot be represented in the client protocol) -/// - `Decode` -> 400 Bad Request (malformed client input) -/// -/// `ProtocolError` does not carry an HTTP status code directly because the -/// route handler may need to override the status based on context (e.g. -/// streaming vs non-streaming, request-phase vs response-phase). -/// -/// # Error source chain -/// -/// `ProtocolError` variants hold only `String` messages, not source errors. -/// This is intentional: the protocol adapter layer consumes underlying -/// serde/JSON errors during translation and converts them to human-readable -/// strings. The original error type information and backtraces are lost at -/// this boundary. If richer error chains are needed in the future, consider -/// wrapping the source error with `#[source]` or using `Box`. +/// Errors returned by client protocol adapters. Route handlers map +/// `InvalidRequest`/`Decode` to 400 and `Encode` to 500 (502 for +/// unrepresentable upstream data); no status is carried, so context wins. #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum ProtocolError { diff --git a/crates/llm-proxy-protocol/src/client/openai_chat.rs b/crates/llm-proxy-protocol/src/client/openai_chat.rs index 78e82e9..96edf57 100644 --- a/crates/llm-proxy-protocol/src/client/openai_chat.rs +++ b/crates/llm-proxy-protocol/src/client/openai_chat.rs @@ -1,20 +1,5 @@ -//! OpenAI Chat Completions client protocol adapter. -//! -//! Translates between OpenAI Chat Completions API wire format and the normalised -//! core protocol types. This adapter only knows about OpenAI Chat Completions -//! and the core types -- it never imports provider, server, or config code. -//! -//! ## Scope guardrails -//! -//! This module must **not** import or call: -//! - `llm_proxy_provider` -//! - `llm_proxy_server` -//! - core config or routing modules -//! - endpoint classification helpers -//! - scenario or fallback code -//! -//! The `transformer/` module was removed in Phase 11; the guard test at the -//! bottom of this file remains as a regression safety net. +//! OpenAI Chat Completions client protocol adapter: translates between the +//! OpenAI Chat wire format and the normalised core protocol types. use std::sync::Arc; @@ -33,16 +18,6 @@ use crate::openai::{ FunctionCall, ToolCall, UsageInfo, }; -// --------------------------------------------------------------------------- -// Shared helpers -// --------------------------------------------------------------------------- - -/// Return the current Unix timestamp in seconds. -/// -/// The `u64 -> i64` cast is safe because `u64::MAX` corresponds to a date -/// ~584 billion years in the future, which is unreachable in any real timeline. -/// `unwrap_or_default()` produces 0 if the system clock is before the Unix epoch, -/// which maps cleanly to `i64`. fn unix_timestamp_secs() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -53,11 +28,6 @@ fn unix_timestamp_secs() -> i64 { .as_secs() as i64 } -/// Clamp a tool-call index from `usize` to the OpenAI wire `i32` range. -/// -/// OpenAI represents tool-call indices as `i32` inside JSON chunks. If a -/// provider ever emits an index larger than `i32::MAX` we clamp it and emit a -/// warning so the stream is not broken by an overflow panic. fn clamp_tool_index(index: usize) -> i32 { let capped = index.min(i32::MAX as usize); let clamped = capped as i32; @@ -71,8 +41,6 @@ fn clamp_tool_index(index: usize) -> i32 { clamped } -/// Parses a tool-call `arguments` JSON string, warning and yielding `None` -/// when the provider emitted something that is not valid JSON. fn parse_tool_call_arguments(arguments: &str) -> Option { match serde_json::from_str::(arguments) { Ok(value) => Some(value), @@ -86,11 +54,13 @@ fn parse_tool_call_arguments(arguments: &str) -> Option { } } -// --------------------------------------------------------------------------- -// decode_request -// --------------------------------------------------------------------------- - /// Decode an OpenAI [`ChatCompletionRequest`] into a normalised [`CoreRequest`]. +#[expect( + clippy::too_many_lines, + reason = "field-by-field wire decode of the OpenAI Chat request shape, kept \ + in one place so the mapping can be audited against the API schema \ + without jumping between helpers" +)] pub fn decode_request(req: ChatCompletionRequest) -> Result { if req.model.is_empty() { return Err(ProtocolError::InvalidRequest("model is required".into())); @@ -133,19 +103,9 @@ pub fn decode_request(req: ChatCompletionRequest) -> Result { - // Ordering: Thinking -> ToolUse -> Text. OpenAI's wire format - // has reasoning_content, tool_calls, and content as parallel - // fields with no defined ordering. This ordering is a - // deterministic choice that puts thinking before tool calls - // and text content. - // - // Read content_text before other fields are moved so the borrow - // checker is satisfied. let assistant_text = msg.content_text(); let mut content = Vec::new(); if let Some(thinking) = msg.reasoning_content { - // OpenAI returns Some("") for reasoning_content on non-reasoning - // models; treat empty string as absent. if !thinking.is_empty() { content.push(CoreContent::Thinking { text: thinking, @@ -187,10 +147,6 @@ pub fn decode_request(req: ChatCompletionRequest) -> Result { - // OpenAI tool messages carry an optional `is_error` flag (some - // providers spell it `status`, handled via a serde alias on - // ChatMessage). Default to false when absent, matching the - // Anthropic adapter's `block.is_error.unwrap_or(false)`. let is_error = msg.is_error.unwrap_or(false); let tool_text = msg.content_text(); let tool_use_id = msg.tool_call_id.ok_or_else(|| { @@ -238,14 +194,6 @@ pub fn decode_request(req: ChatCompletionRequest) -> Result no stop; a non-string item in a stop array, or a - // top-level stop value that is not a string/array/null, is a client - // error (rejected) rather than silently dropped. Previously this path - // warn-logged and kept the surviving items, and treated an empty array - // as `Some(vec![])` instead of `None` — both divergences from the core - // wire-format rule. The shared `json_type_name` helper yields - // byte-identical error messages to `deserialize_stop`. match s { serde_json::Value::Null => {} serde_json::Value::String(st) => stop = Some(vec![st]), @@ -262,7 +210,6 @@ pub fn decode_request(req: ChatCompletionRequest) -> Result CoreToolChoice { } } -// --------------------------------------------------------------------------- -// encode_response -// --------------------------------------------------------------------------- - -/// Encode a [`CoreResponse`] into an OpenAI [`ChatCompletionResponse`]. -/// -/// # Content handling -/// -/// - Multiple `CoreContent::Text` blocks are concatenated into a single string -/// without separator, since OpenAI's `content` field is a single string. -/// - Multiple `CoreContent::Thinking` blocks are concatenated with a newline -/// separator into a single string for `reasoning_content`. -/// - `stop_sequence` is silently dropped since OpenAI has no native field for it. -/// - `provider_meta` is silently dropped (for provider adapters only). +/// Encode a [`CoreResponse`] into an OpenAI [`ChatCompletionResponse`]: +/// text blocks concatenate without separator and thinking blocks join with +/// `\n`; `stop_sequence` and `provider_meta` are dropped. +#[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the OpenAI Chat response shape, kept \ + in one place so the mapping can be audited against the API schema \ + without jumping between helpers" +)] pub fn encode_response(resp: CoreResponse) -> Result { if !resp.provider_meta.is_empty() { tracing::warn!( @@ -397,14 +339,9 @@ pub fn encode_response(resp: CoreResponse) -> Result { - // Multiple text blocks are concatenated without separator. content_text.push_str(&text); } CoreContent::Thinking { text, .. } => { - // Concatenate multiple Thinking blocks with a newline separator - // rather than silently overwriting earlier blocks. OpenAI exposes - // reasoning_content as a single string, so concatenation is the - // least-lossy approach. if reasoning.is_some() { tracing::warn!( "multiple Thinking blocks in response; concatenating with newline separator" @@ -448,26 +385,16 @@ pub fn encode_response(resp: CoreResponse) -> Result { - // Tool results should not appear in assistant responses -- this - // likely indicates a logic error upstream. tracing::warn!( "ToolResult in response content is unexpected for OpenAI Chat encode; dropping" ); } CoreContent::RedactedThinking { .. } => { - // RedactedThinking cannot be represented in OpenAI Chat. - // Unlike Document/Audio/Video (safe to drop), RedactedThinking - // signals that thinking occurred but is intentionally withheld. - // Dropping it silently would lose that signal, so we warn and - // skip rather than returning Err, so that any previously encoded - // blocks (text, tool_calls) are preserved. tracing::warn!( "OpenAI Chat Completions does not support redacted thinking; omitting block" ); } CoreContent::Refusal { text } => { - // OpenAI has a native refusal field on assistant messages. - // Encode the refusal text into the dedicated field. refusal_text = Some(text); } } @@ -489,9 +416,6 @@ pub fn encode_response(resp: CoreResponse) -> Result Result `"content_filter"` (OpenAI uses content_filter for refused outputs) -/// - `StopSequence` -> `"stop"` (falls under "other normal stop") -/// - `Error` -> `"stop"` with warning (no native error finish reason in OpenAI) -/// - `Unknown` -> `"stop"` (safe default, "other normal stop") fn encode_finish_reason(reason: StopReason) -> String { match reason { StopReason::EndTurn => "stop".to_owned(), @@ -540,12 +455,6 @@ fn encode_finish_reason(reason: StopReason) -> String { } } -/// Encode core [`Usage`] into an OpenAI [`UsageInfo`]. -/// -/// Known limitation: `reasoning_tokens` are dropped because [`UsageInfo`] does not -/// have a `completion_tokens_details` field. If a provider reports non-zero -/// reasoning tokens, they are logged as a warning. To fix this, add a -/// `completion_tokens_details` field to `UsageInfo` mirroring the OpenAI API shape. fn encode_usage(usage: &Usage) -> UsageInfo { if let Some(rt) = usage.reasoning_tokens { if rt != 0 { @@ -559,60 +468,27 @@ fn encode_usage(usage: &Usage) -> UsageInfo { prompt_tokens: usage.input_tokens, completion_tokens: usage.output_tokens, total_tokens: usage.input_tokens.saturating_add(usage.output_tokens), - // Cache token mapping: - // - // cache_creation_input_tokens -> prompt_cache_miss_tokens - // cache_read_input_tokens -> prompt_cache_hit_tokens - // - // This is an intentional semantic mapping: cache_creation means the prompt was - // NOT in the cache (a miss), and the cache was populated as a side effect. - // Therefore cache_creation_input_tokens represents the cache-miss cost. - // - // Operators correlating upstream provider usage with downstream client usage - // should be aware that these fields are inverses of each other: - // - Upstream "creation" = Downstream "miss" (same tokens, different label) - // - Upstream "read" = Downstream "hit" (same tokens, same label) prompt_cache_hit_tokens: usage.cache_read_input_tokens, prompt_cache_miss_tokens: usage.cache_creation_input_tokens, } } -// --------------------------------------------------------------------------- -// StreamEncoder -// --------------------------------------------------------------------------- - /// Stateful encoder that converts [`CoreEvent`] values into OpenAI -/// [`ChatCompletionChunk`] values. -/// -/// The encoder is constructed with a stable completion `id`, `model`, -/// `created` timestamp, and the `include_usage` flag (from stream_options). -/// Every `chat.completion.chunk` reuses that `id`/`model`/`created`. -/// -/// `id` and `model` are stored as `Arc` rather than `String` so that -/// cloning the encoder (and any future sharing of these values) is a cheap -/// refcount bump instead of a heap allocation. The per-chunk owned `String` -/// required by [`ChatCompletionChunk`] is materialised once in `Self::make_chunk`. +/// [`ChatCompletionChunk`] values, reusing a stable `id`/`model`/`created` +/// across chunks (stored as `Arc` so cloning is a refcount bump). #[derive(Debug)] pub struct StreamEncoder { - /// Stable completion ID across all chunks. id: Arc, - /// Model name. model: Arc, - /// Unix timestamp. created: i64, - /// Whether to emit the usage chunk at the end. include_usage: bool, - /// Whether the terminal chunk has been emitted. finished: bool, - /// Buffered usage from the latest UsageDelta. pending_usage: Option, } impl StreamEncoder { - /// Create a new stream encoder. - /// - /// The route constructs these args from the first `CoreEvent::MessageStart` - /// and the decoded request. + /// Create a new stream encoder; the route builds these args from the first + /// `CoreEvent::MessageStart` and the decoded request. pub fn new(id: String, model: String, created: i64, include_usage: bool) -> Self { Self { id: id.into(), @@ -625,6 +501,13 @@ impl StreamEncoder { } /// Encode a single [`CoreEvent`] into zero or more [`ChatCompletionChunk`]s. + #[expect( + clippy::too_many_lines, + reason = "one match arm per CoreEvent variant: the OpenAI Chat SSE chunk \ + wire shape is a single translation table, and splitting the arms \ + into helpers would separate each event's field mapping from the \ + chunk layout it produces" + )] pub fn encode_event( &mut self, event: CoreEvent, @@ -637,9 +520,6 @@ impl StreamEncoder { match event { CoreEvent::MessageStart { .. } => { - // OpenAI does not have an explicit start event -- the first - // chunk implicitly starts the response. - // Emit a role chunk. chunks.push(self.make_chunk(Choice { index: 0, message: None, @@ -660,18 +540,12 @@ impl StreamEncoder { CoreEvent::ContentStart { index, kind } => { match kind { - ContentKind::ToolUse => { - // Tool use start is handled by ToolCallStart. - } - ContentKind::Text | ContentKind::Thinking => { - // Text/thinking start is implicit; deltas follow. - } + ContentKind::ToolUse => {} + ContentKind::Text | ContentKind::Thinking => {} _ => { tracing::warn!(kind = ?kind, "unsupported ContentKind in OpenAI stream"); } } - // OpenAI does not use content block indices in the same way Anthropic - // does -- the chunk-level position is implicit. Discard the index. let _ = index; } @@ -769,10 +643,7 @@ impl StreamEncoder { })); } - CoreEvent::ToolCallStop { .. } => { - // No explicit event needed in OpenAI; the finish_reason - // on the terminal chunk signals tool completion. - } + CoreEvent::ToolCallStop { .. } => {} CoreEvent::UsageDelta { usage } => { self.pending_usage = Some(encode_usage(&usage)); @@ -797,7 +668,6 @@ impl StreamEncoder { }), })); - // If include_usage is set, emit a final usage chunk. if self.include_usage { let usage = self.pending_usage.take().unwrap_or(UsageInfo { prompt_tokens: 0, @@ -813,13 +683,6 @@ impl StreamEncoder { } CoreEvent::Error { error } => { - // Defense-in-depth note: the error message is embedded in the - // ProtocolError which may be logged or sent as a client response. - // Sanitization of secrets happens at CoreStreamError::new() - // construction time in the provider adapter. Only the error kind - // is used here to avoid leaking unsanitized messages through logs. - // Using Display (lowercase) instead of Debug (PascalCase) for a - // stable, human-readable format. self.finished = true; return Err(ProtocolError::Encode(format!( "stream error: {}", @@ -827,21 +690,15 @@ impl StreamEncoder { ))); } - CoreEvent::Ping => { - // No-op for OpenAI unless the route chooses heartbeat. - } + CoreEvent::Ping => {} } Ok(chunks) } - /// Flush any remaining buffered events. - /// - /// When the stream terminated abnormally (no prior `MessageStop`), emits a - /// synthetic finish_reason=`"stop"` chunk so the client receives a proper - /// terminal signal. If `include_usage` is set, also emits a final usage chunk. - /// This is analogous to how the Anthropic encoder emits synthetic terminal - /// events in `finish()`. + /// Flush remaining buffered events when the stream ended without a + /// `MessageStop`: emits a synthetic `finish_reason = "stop"` chunk (with + /// a warning), plus the final usage chunk when `include_usage` is set. pub fn finish(&mut self) -> Result, ProtocolError> { if self.finished { return Ok(Vec::new()); @@ -854,7 +711,6 @@ impl StreamEncoder { let mut chunks = Vec::new(); - // Emit a synthetic finish chunk. chunks.push(self.make_chunk(Choice { index: 0, message: None, @@ -872,7 +728,6 @@ impl StreamEncoder { }), })); - // If include_usage is set, emit a final usage chunk. if self.include_usage { let usage = self.pending_usage.take().unwrap_or(UsageInfo { prompt_tokens: 0, @@ -889,20 +744,6 @@ impl StreamEncoder { fn make_chunk(&self, choice: Choice) -> ChatCompletionChunk { ChatCompletionChunk { - // GAP-LOW-7 (evaluated, deferred by design): the per-chunk - // allocation of `id`/`model` persists here because - // `ChatCompletionChunk` is ALSO a deserialized *input* type - // (upstream provider SSE chunks are decoded into it at - // `adapter/openai_chat.rs`), so its `id`/`model` must stay owned - // `String`. Eliminating the allocation would require either serde's - // `rc` feature workspace-wide (a deliberate footgun: transparent - // `Arc`/`Rc` deserialization) or a parallel output-only chunk type - // that must stay wire-compatible with this one — both - // disproportionate for a GAP-LOW on a network-bound path, where - // upstream LLM latency dominates two small allocations per chunk. - // The encoder stores `id`/`model` as `Arc` so that any future - // cloning of the encoder is a refcount bump; the per-chunk owned - // `String` remains the single unavoidable allocation. id: self.id.to_string(), object: "chat.completion.chunk".to_owned(), created: self.created, @@ -912,9 +753,6 @@ impl StreamEncoder { } } - /// Build the terminal usage-only chunk. Shares the same `id`/`model` - /// materialisation path as `Self::make_chunk` to keep the two chunk - /// shapes consistent. fn make_usage_chunk(&self, usage: UsageInfo) -> ChatCompletionChunk { ChatCompletionChunk { id: self.id.to_string(), @@ -926,54 +764,26 @@ impl StreamEncoder { } } - /// Mark the encoder as finished so that subsequent `finish()` calls return - /// empty events. This is used after emitting an in-band error event to - /// prevent the finalization path from emitting synthetic terminal events. + /// Mark the encoder as finished so subsequent `finish()` calls return + /// empty; used after emitting an in-band error. pub fn mark_finished(&mut self) { self.finished = true; } } -// --------------------------------------------------------------------------- -// stream_error_sse_payload -// --------------------------------------------------------------------------- - -/// Build the OpenAI Chat stream-error SSE payload for a [`CoreEvent::Error`]. -/// -/// [`StreamEncoder::encode_event`] deliberately returns `Err` for -/// `CoreEvent::Error` (the OpenAI Chat wire format has no dedicated error -/// chunk type), so the normal encode path yields no SSE event for an error. -/// When an error is the *leading* event of a stream, dropping it silently -/// would leave the first-byte boundary uncrossed and surface a misleading -/// "empty stream" 502 to the client. This helper synthesizes the error as an -/// ordinary OpenAI error-object `data:` line so it can cross the boundary as a -/// real in-band SSE event (200 + error object + `[DONE]`), mirroring the -/// post-first-byte `emit_stream_error` path in the server crate. -/// -/// The `error.kind` selects the OpenAI `type` string so the client can -/// classify the failure (e.g. `rate_limit` vs `server_error`); the sanitized -/// `error.message()` is used verbatim (it was sanitized at -/// [`CoreStreamError::new`] construction time in the adapter). -/// -/// Returns the JSON string suitable for a `data: ` SSE line, or `None` -/// if serialization fails (treated as a dropped event by the caller). The -/// caller is responsible for emitting the trailing `data: [DONE]` terminator. +/// Build the OpenAI stream-error payload as an error-object `data:` line so +/// a *leading* error crosses the first-byte boundary. Returns the JSON, or +/// `None` if serialization fails; the caller emits the `data: [DONE]`. pub fn stream_error_sse_payload(error: &crate::core::CoreStreamError) -> Option { let error_type = match error.kind { crate::core::CoreStreamErrorKind::InvalidRequest => "invalid_request_error", crate::core::CoreStreamErrorKind::Authentication => "authentication_error", crate::core::CoreStreamErrorKind::Permission => "permission_error", crate::core::CoreStreamErrorKind::RateLimit => "rate_limit_error", - // Upstream provider failures and internal proxy errors both surface as - // server-side errors on the OpenAI wire format. crate::core::CoreStreamErrorKind::Upstream | crate::core::CoreStreamErrorKind::Internal => { "server_error" } }; - // Reuse the same typed structs as the HTTP error path so both stay - // consistent at compile time. Built inline (rather than via the server - // crate's `openai_stream_error_json_with_type`) to keep this adapter free - // of server/config dependencies (see the module scope guardrails). #[derive(serde::Serialize)] #[serde(deny_unknown_fields)] struct ErrorBody { @@ -997,18 +807,12 @@ pub fn stream_error_sse_payload(error: &crate::core::CoreStreamError) -> Option< serde_json::to_string(&body).ok() } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; use crate::core::UsageProvenance; use crate::openai::{FunctionDef, ToolDef}; - // -- helpers ------------------------------------------------------------ - fn make_openai_request() -> ChatCompletionRequest { ChatCompletionRequest { model: "gpt-4o".into(), @@ -1057,8 +861,6 @@ mod tests { } } - // -- decode tests ------------------------------------------------------- - #[test] fn plain_text_request_decodes_to_core() { let req = make_openai_request(); @@ -1095,7 +897,6 @@ mod tests { CoreContent::Text { text, .. } => assert_eq!(text, "You are helpful"), _ => panic!("expected Text"), } - // System message should NOT appear in messages. assert_eq!(core.messages.len(), 1); assert_eq!(core.messages[0].role, CoreRole::User); } @@ -1328,8 +1129,6 @@ mod tests { #[test] fn empty_user_content_with_cache_control_preserves_cache() { - // When a user message has empty content but carries cache_control, - // the cache_control must not be silently dropped. let mut req = make_openai_request(); req.messages[0].content = serde_json::Value::String(String::new()); req.messages[0].cache_control = Some(OpenAICacheControl { @@ -1356,8 +1155,6 @@ mod tests { #[test] fn empty_system_content_with_cache_control_preserves_cache() { - // When a system message has empty content but carries cache_control, - // the cache_control must not be silently dropped. let mut req = make_openai_request(); req.messages.insert( 0, @@ -1429,8 +1226,6 @@ mod tests { assert_eq!(core.tools[0].name, "get_weather"); } - // -- encode tests ------------------------------------------------------- - #[test] fn core_text_response_encodes_to_route_response() { let resp = make_core_response(); @@ -1537,8 +1332,6 @@ mod tests { assert_eq!(msg.content, serde_json::Value::String("answer".into())); } - // -- streaming encode tests --------------------------------------------- - #[test] fn streaming_text_event_mapping() { let mut enc = StreamEncoder::new("chatcmpl-1".into(), "gpt-4o".into(), 1000, false); @@ -1639,7 +1432,6 @@ mod tests { stop_sequence: None, }) .unwrap(); - // Two chunks: finish + usage. assert_eq!(chunks.len(), 2); assert_eq!(chunks[0].choices.len(), 1); assert!(chunks[1].usage.is_some()); @@ -1659,7 +1451,6 @@ mod tests { stop_sequence: None, }) .unwrap(); - // Only the finish chunk, no usage chunk. assert_eq!(chunks.len(), 1); } @@ -1669,19 +1460,12 @@ mod tests { let result = enc.encode_event(CoreEvent::Error { error: CoreStreamError::new(CoreStreamErrorKind::RateLimit, "too many requests".into()), }); - // Error events now return Err so the route handler can emit a proper - // error response, rather than leaking error text into content delta. assert!(result.is_err()); assert!(matches!(result, Err(ProtocolError::Encode(_)))); } #[test] fn stream_error_sse_payload_builds_typed_error_object() { - // The OpenAI StreamEncoder drops CoreEvent::Error (returns Err). The - // route handler uses this helper to synthesize the error as an ordinary - // OpenAI error-object `data:` line so a *leading* error crosses the - // first-byte boundary instead of surfacing a misleading 502 - // (audit LOW, stream-ids-01). let payload = stream_error_sse_payload(&CoreStreamError::new( CoreStreamErrorKind::RateLimit, "too many requests".into(), @@ -1689,10 +1473,8 @@ mod tests { .expect("payload serializes"); let json: serde_json::Value = serde_json::from_str(&payload).unwrap(); assert_eq!(json["error"]["message"], "too many requests"); - // RateLimit maps to the OpenAI rate-limit error type. assert_eq!(json["error"]["type"], "rate_limit_error"); assert!(json["error"]["code"].is_null()); - // Must NOT carry Anthropic-shaped fields. assert!(json.get("type").is_none() || json["type"].is_null()); } @@ -1747,7 +1529,6 @@ mod tests { #[test] fn streaming_finish_without_prior_events_emits_synthetic_terminal() { let mut enc = StreamEncoder::new("chatcmpl-1".into(), "gpt-4o".into(), 1000, false); - // No events sent at all -- finish() should emit a synthetic terminal chunk. let remaining = enc.finish().unwrap(); assert_eq!(remaining.len(), 1); assert_eq!( @@ -1756,13 +1537,8 @@ mod tests { ); } - // -- every CoreEvent variant handled ------------------------------------ - #[test] fn every_core_event_variant_maps_or_errors_intentionally() { - // Test all ContentKind variants for ContentStart. The OpenAI stream - // encoder handles Text, Thinking, and ToolUse explicitly; other kinds - // produce a warning log and no events. let content_kinds: Vec = vec![ ContentKind::Text, ContentKind::Thinking, @@ -1783,7 +1559,6 @@ mod tests { assert!(result.is_ok(), "ContentStart {:?} should not error", kind); } - // Test all other CoreEvent variants. let variants: Vec = vec![ CoreEvent::MessageStart { id: None, @@ -1827,8 +1602,6 @@ mod tests { let result = enc.encode_event(event.clone()); match event { CoreEvent::Error { .. } => { - // Error events intentionally return Err so the route handler - // can emit a proper error response. assert!(result.is_err(), "Error event should return Err"); } _ => { @@ -1838,8 +1611,6 @@ mod tests { } } - // -- unsupported content variants --------------------------------------- - #[test] fn encode_image_drops_with_warning() { let resp = CoreResponse { @@ -1882,7 +1653,6 @@ mod tests { }; let out = encode_response(resp).unwrap(); let msg = out.choices[0].message.as_ref().unwrap(); - // Refusal text goes into the native refusal field, not content. assert_eq!(msg.refusal.as_deref(), Some("I cannot help")); assert!(msg.content_is_empty()); } @@ -1977,8 +1747,6 @@ mod tests { assert!(msg.content_is_empty()); } - // -- metadata / provider hints ------------------------------------------ - #[test] fn provider_hints_raw_preserved_in_decode() { let mut req = make_openai_request(); @@ -2039,7 +1807,6 @@ mod tests { cost: None, }; let out = encode_response(resp).unwrap(); - // RedactedThinking is skipped (not an error), so content is empty. let msg = out.choices[0].message.as_ref().unwrap(); assert!(msg.content_is_empty()); } @@ -2069,7 +1836,6 @@ mod tests { }; let out = encode_response(resp).unwrap(); let msg = out.choices[0].message.as_ref().unwrap(); - // Prior Text block is preserved even though RedactedThinking follows. assert_eq!(msg.content, serde_json::Value::String("hello".into())); } @@ -2113,17 +1879,8 @@ mod tests { assert!(out.id.starts_with("chatcmpl-")); } - // -- source guard ------------------------------------------------------- - #[test] fn client_openai_chat_does_not_import_forbidden_modules() { - // Source guard: this module must not import llm_proxy_provider, - // llm_proxy_server, core config/routing, endpoint classification, - // scenario/fallback code, or transformer/*. The crate dependency - // graph prevents most of these at compile time (those crates are not - // dependencies of llm-proxy-protocol). This runtime check provides - // a secondary defense by verifying that the module source does not - // contain forbidden import patterns in `use` statements. let source = include_str!("openai_chat.rs"); for line in source.lines() { let trimmed = line.trim(); @@ -2155,8 +1912,6 @@ mod tests { } } - // -- edge-case tests ---------------------------------------------------- - #[test] fn decode_empty_model_returns_error() { let mut req = make_openai_request(); @@ -2214,8 +1969,6 @@ mod tests { cost: None, }; let out = encode_response(resp).unwrap(); - // stop_sequence is dropped (OpenAI has no field for it), but the - // stop_reason is mapped to "stop". assert_eq!(out.choices[0].finish_reason.as_deref(), Some("stop")); } @@ -2245,7 +1998,6 @@ mod tests { }; let out = encode_response(resp).unwrap(); let msg = out.choices[0].message.as_ref().unwrap(); - // Multiple Thinking blocks are concatenated with newline separator. assert_eq!(msg.reasoning_content.as_deref(), Some("first\nsecond")); } @@ -2281,7 +2033,6 @@ mod tests { }) .unwrap(); let remaining = enc.finish().unwrap(); - // Should produce finish chunk + usage chunk. assert_eq!(remaining.len(), 2); assert_eq!( remaining[0].choices[0].finish_reason.as_deref(), @@ -2291,8 +2042,6 @@ mod tests { assert_eq!(remaining[1].usage.as_ref().unwrap().prompt_tokens, 50); } - // -- tool_call validation tests ------------------------------------------- - #[test] fn tool_call_missing_id_returns_error() { let mut req = make_openai_request(); @@ -2302,7 +2051,7 @@ mod tests { reasoning_content: None, tool_calls: vec![ToolCall { index: Some(0), - id: None, // missing id + id: None, r#type: Some("function".into()), function: Some(FunctionCall { name: Some("get_weather".into()), @@ -2338,7 +2087,7 @@ mod tests { index: Some(0), id: Some("call_1".into()), r#type: Some("function".into()), - function: None, // missing function + function: None, }], name: None, tool_call_id: None, @@ -2370,7 +2119,7 @@ mod tests { id: Some("call_1".into()), r#type: Some("function".into()), function: Some(FunctionCall { - name: None, // missing name + name: None, arguments: Some("{}".into()), }), }], @@ -2401,7 +2150,7 @@ mod tests { reasoning_content: None, tool_calls: vec![], name: None, - tool_call_id: None, // missing tool_call_id + tool_call_id: None, cache_control: None, refusal: None, is_error: None, @@ -2443,7 +2192,6 @@ mod tests { let core = decode_request(req).unwrap(); match &core.messages[1].content[0] { CoreContent::ToolUse { input, .. } => { - // Malformed arguments should be replaced with an empty object. assert_eq!( input, &serde_json::json!({}), @@ -2456,9 +2204,6 @@ mod tests { #[test] fn tool_message_is_error_decodes_to_core() { - // A tool message carrying `is_error: Some(true)` must surface in the - // decoded CoreContent::ToolResult, mirroring the Anthropic adapter's - // `tool_result_with_is_error_decodes_to_core` regression test. let mut req = make_openai_request(); req.messages.push(ChatMessage { role: "tool".into(), @@ -2477,7 +2222,6 @@ mod tests { _ => panic!("expected ToolResult"), } - // An absent `is_error` field must default to false (no semantic flip). let mut req = make_openai_request(); req.messages.push(ChatMessage { role: "tool".into(), @@ -2496,8 +2240,6 @@ mod tests { _ => panic!("expected ToolResult"), } - // Some OpenAI-compatible providers spell the flag `status`; the serde - // alias on ChatMessage must deserialize it into `is_error`. let json = serde_json::json!({ "model": "gpt-4o", "messages": [ @@ -2520,11 +2262,6 @@ mod tests { #[test] fn assistant_with_reasoning_tool_calls_and_text_preserves_order() { - // Locks the documented Thinking -> ToolUse -> Text ordering invariant - // (see the comment in the "assistant" arm of decode_request). The OpenAI - // wire format carries reasoning_content, tool_calls, and content as - // parallel fields with no defined ordering; a refactor of the early-read - // borrow-checker workaround must not silently reorder these. let mut req = make_openai_request(); req.messages.push(ChatMessage { role: "assistant".into(), @@ -2562,8 +2299,6 @@ mod tests { ); } - // -- stop field parsing tests --------------------------------------------- - #[test] fn stop_as_string_decodes_to_single_element_vec() { let mut req = make_openai_request(); @@ -2574,7 +2309,6 @@ mod tests { #[test] fn stop_as_null_decodes_to_none() { - // Converges with core::deserialize_stop: explicit null means "no stop". let mut req = make_openai_request(); req.stop = Some(serde_json::Value::Null); let core = decode_request(req).unwrap(); @@ -2583,8 +2317,6 @@ mod tests { #[test] fn stop_as_empty_array_decodes_to_none() { - // Converges with core::deserialize_stop: an empty stop array is "no - // stop", not Some(vec![]). let mut req = make_openai_request(); req.stop = Some(serde_json::json!([])); let core = decode_request(req).unwrap(); @@ -2593,9 +2325,6 @@ mod tests { #[test] fn stop_as_non_string_non_array_is_rejected() { - // GAP-LOW-8: a top-level stop value that is not a string/array/null is - // a client error (matching core::deserialize_stop), not silently - // dropped as it was before. let mut req = make_openai_request(); req.stop = Some(serde_json::json!(42)); let err = decode_request(req).unwrap_err(); @@ -2612,8 +2341,6 @@ mod tests { #[test] fn stop_array_with_non_string_items_is_rejected() { - // GAP-LOW-8: a non-string item in a stop array is a client error - // (matching core::deserialize_stop), not filtered out as it was before. let mut req = make_openai_request(); req.stop = Some(serde_json::json!(["STOP", 123, "END"])); let err = decode_request(req).unwrap_err(); diff --git a/crates/llm-proxy-protocol/src/core.rs b/crates/llm-proxy-protocol/src/core.rs index ee8635e..dca40ef 100644 --- a/crates/llm-proxy-protocol/src/core.rs +++ b/crates/llm-proxy-protocol/src/core.rs @@ -1,58 +1,12 @@ -//! Core protocol types for the LLM proxy. -//! -//! This module defines the normalized internal representation that sits between -//! client protocol adapters and provider adapters. Every adapter translates -//! **to** or **from** these types -- never directly to another adapter's wire -//! format. -//! -//! ## Design rules -//! -//! * Chat-endpoint focused. Do not add embeddings, image generation, audio -//! endpoints, rerank, files, or batch families here. -//! * Chat content still preserves image/document/audio/video blocks because -//! those can appear inside chat requests and responses. -//! * The stream error type is `CoreStreamError` (not `CoreError`) to avoid -//! colliding with the config / registry error owned by the core crate. -//! -//! ## Opaque field safety -//! -//! Several fields hold opaque JSON data that passes through the proxy without -//! interpretation: [`RequestMetadata::raw`], [`ProviderHints::raw`], -//! [`CoreResponse::provider_meta`], [`SamplingOptions::thinking`], and -//! [`CoreToolChoice::Raw`]. These opaque fields use a **redacting `Debug` -//! implementation** that shows only the JSON type and approximate size (e.g. -//! `Object(3 keys)`, `String(42 bytes)`) rather than printing values verbatim. -//! This is a defense-in-depth measure: adapters **must not** store secrets (API -//! keys, bearer tokens, etc.) in any of these fields. Strip credentials before -//! placing data into these fields. -//! -//! ## Equality semantics -//! -//! Types that contain [`serde_json::Value`] fields derive `PartialEq` but **not** -//! `Eq`, because `serde_json::Value` uses `f64` internally and does not implement -//! `Eq`. Equality comparisons on such types are partial: if any JSON value field -//! contains NaN, the comparison may return `false` unpredictably. This primarily -//! affects test assertions; production code should not rely on exact equality of -//! JSON-heavy types. -//! -//! ## Forward compatibility -//! -//! All public enums are annotated with `#[non_exhaustive]`, so adding new variants -//! is not a semver-breaking change. Core structs use `#[serde(deny_unknown_fields)]` -//! to ensure that unknown fields cause a deserialization error rather than being -//! silently dropped. Wire-type structs in `openai.rs` may use `#[serde(flatten)]` -//! with an extra `Map` for passthrough instead. +//! Normalised core protocol types between client and provider adapters; +//! adapters translate to/from these, never to each other's wire formats. +//! Opaque JSON fields redact their `Debug` output — do not store secrets. use std::fmt; use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; -// --------------------------------------------------------------------------- -// OpaqueJsonRef -- redacting Debug wrapper for borrowed serde_json::Value -// --------------------------------------------------------------------------- - -/// A reference wrapper for redacting Debug output of a borrowed `Value`. struct OpaqueJsonRef<'a>(&'a serde_json::Value); impl fmt::Debug for OpaqueJsonRef<'_> { @@ -61,16 +15,6 @@ impl fmt::Debug for OpaqueJsonRef<'_> { } } -/// Formats a [`serde_json::Value`] for debug output without revealing content. -/// -/// Bool and Number values are shown without their actual values for consistency -/// with the redaction policy: a malicious payload could encode secret fragments -/// in number values or boolean field names. -/// -/// Note: `String(N bytes)` reveals the string length. This is an intentional -/// trade-off: hiding the length would make debugging harder (e.g. diagnosing -/// empty vs populated fields). If this becomes a security concern, change to -/// `String(_)`. fn redact_value(val: &serde_json::Value, f: &mut fmt::Formatter<'_>) -> fmt::Result { match val { serde_json::Value::Null => f.write_str("Null"), @@ -82,7 +26,6 @@ fn redact_value(val: &serde_json::Value, f: &mut fmt::Formatter<'_>) -> fmt::Res } } -/// Formats a [`serde_json::Map`] for debug output without revealing content. fn redact_map( map: &serde_json::Map, f: &mut fmt::Formatter<'_>, @@ -90,16 +33,8 @@ fn redact_map( write!(f, "Object({} keys)", map.len()) } -// --------------------------------------------------------------------------- -// ModelRef -// --------------------------------------------------------------------------- - -/// A model identifier that tracks both the client-facing name and the resolved -/// upstream name. -/// -/// `requested` is always the model the client asked for. Provider resolution -/// may populate `upstream` with the actual upstream model id. Adapters must -/// **not** replace `requested`. +/// A model identifier tracking both the client-facing name and the resolved +/// upstream name; adapters must not replace `requested`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct ModelRef { @@ -110,32 +45,21 @@ pub struct ModelRef { } impl ModelRef { - /// Returns the effective model identifier for the provider. - /// - /// If `upstream` is set it is used; otherwise `requested` is used. + /// Returns the effective model identifier: `upstream` if set, else `requested`. pub fn provider_model(&self) -> &str { self.upstream.as_deref().unwrap_or(&self.requested) } } -// --------------------------------------------------------------------------- -// CoreRequest -// --------------------------------------------------------------------------- - -/// A normalized chat request. -/// -/// This is the canonical internal representation of a client's intent. -/// Client adapters decode their wire format into this shape; provider adapters -/// encode it into the provider's wire format. +/// A normalized chat request: the canonical internal representation that +/// client adapters decode into and provider adapters encode from. #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct CoreRequest { /// Model identifier (client-facing + optional upstream override). pub model: ModelRef, - /// System / developer instructions (canonical location). - /// - /// Client adapters should normalise system-role messages into this field - /// where the client protocol allows it. + /// System / developer instructions (canonical location); client adapters + /// normalise system-role messages into this field where possible. pub system: Vec, /// Conversation turns after system content has been separated. pub messages: Vec, @@ -169,10 +93,6 @@ impl fmt::Debug for CoreRequest { } } -// --------------------------------------------------------------------------- -// CoreMessage / CoreRole -// --------------------------------------------------------------------------- - /// A single conversation turn. #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -207,22 +127,9 @@ pub enum CoreRole { Tool, } -// --------------------------------------------------------------------------- -// CoreContent -// --------------------------------------------------------------------------- - -/// A single content block inside a message or system prompt. -/// -/// Variants cover text, media, tool interactions, and extended thinking. -/// Each variant is a struct-like enum arm so that fields are named and -/// self-documenting. -/// -/// Note: `deny_unknown_fields` is applied to ensure that serialization -/// round-trips are lossless. If a provider adds new fields to a content -/// block variant that the proxy does not model, those fields will be -/// rejected at deserialization. This is intentional for the core protocol -/// type: provider-specific passthrough should use `provider_hints` or the -/// extra maps on wire-format types rather than extending core variants. +/// A single content block inside a message or system prompt: text, media, +/// tool interactions, and extended thinking. Unknown variant fields are +/// rejected on decode; passthrough belongs in `provider_hints`/extra maps. #[non_exhaustive] #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -276,20 +183,14 @@ pub enum CoreContent { Thinking { /// The thinking text. text: String, - /// Optional cryptographic signature (e.g. Anthropic-style). - /// - /// This field contains security-sensitive data. It is redacted in - /// `Debug` output to `[REDACTED]` as a defense-in-depth measure. + /// Optional cryptographic signature (e.g. Anthropic-style); redacted to + /// `[REDACTED]` in `Debug` output. signature: Option, }, /// Redacted thinking block whose plaintext is not available. RedactedThinking { - /// Opaque redacted data. - /// - /// Uses `serde_json::Value` rather than `String` because providers may - /// encode the redacted data in different forms (base64 string, JSON object, - /// etc.). The Anthropic adapter always converts non-string values to a JSON - /// string during serialization. + /// Opaque redacted data (`serde_json::Value` since providers encode it + /// differently); the Anthropic adapter converts non-strings to strings. data: serde_json::Value, }, /// A model refusal to answer. @@ -355,10 +256,6 @@ impl fmt::Debug for CoreContent { } } -// --------------------------------------------------------------------------- -// CacheControl -// --------------------------------------------------------------------------- - /// Cache-control directive attached to a content block. #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -375,10 +272,8 @@ impl fmt::Debug for CacheControl { } } -/// Known cache control type variants. -/// -/// Uses an enum with a catch-all `Other` variant so that unknown values from -/// future provider extensions are preserved rather than rejected. +/// Known cache control type variants; unknown provider values are preserved +/// in `Other(String)` rather than rejected. #[non_exhaustive] #[derive(Debug, Clone, PartialEq, Eq)] pub enum CacheControlType { @@ -429,10 +324,6 @@ impl<'de> Deserialize<'de> for CacheControlType { } } -// --------------------------------------------------------------------------- -// CoreTool / CoreToolChoice -// --------------------------------------------------------------------------- - /// A tool definition available to the model. #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -441,15 +332,9 @@ pub struct CoreTool { pub name: String, /// Human-readable description. pub description: Option, - /// JSON Schema describing the tool's parameters. - /// - /// Tool parameter schemas must be JSON objects. The type system cannot - /// encode that constraint, so it is enforced on the wire decode path: - /// [`Deserialize`] rejects any `input_schema` that is not a JSON object - /// (see [`deserialize_input_schema_object`]). In-crate adapters that build - /// a [`CoreTool`] with the struct literal should supply a JSON object for - /// this field; adapters that translate an upstream schema default to an - /// empty object when the upstream value is absent. + /// JSON Schema describing the tool's parameters; must be a JSON object, + /// enforced on decode via [`deserialize_input_schema_object`]. Struct-literal + /// builders supply an object (empty when upstream omits it). #[serde(deserialize_with = "deserialize_input_schema_object")] pub input_schema: serde_json::Value, } @@ -479,11 +364,8 @@ pub enum CoreToolChoice { /// Name of the tool to call. name: String, }, - /// Temporary preservation of a client tool-choice shape that has no - /// canonical variant yet. - /// - /// Must not contain unrelated wire nesting. Adapter tests must either - /// round-trip it intentionally or move it into metadata / provider hints. + /// A client tool-choice shape with no canonical variant yet; adapter tests + /// must round-trip it intentionally or move it into metadata/hints. Raw(serde_json::Value), } @@ -499,87 +381,42 @@ impl fmt::Debug for CoreToolChoice { } } -// --------------------------------------------------------------------------- -// SamplingOptions -// --------------------------------------------------------------------------- - -/// Sampling parameters that control generation behaviour. -/// -/// The decode boundary (i.e. the `Deserialize` impl) enforces the universal -/// sampling invariants so they cannot be bypassed: -/// - `temperature`, if present, is finite and `>= 0`. -/// - `top_p`, if present, is finite and in `[0, 1]`. -/// -/// `max_tokens` is intentionally **not** range-checked here: a value of `0` -/// is a legitimate sentinel for some providers (e.g. OpenAI) and the proxy -/// forwards it for the upstream to accept or reject. See -/// [`SamplingOptions::max_tokens`] for details. +/// Sampling parameters. The decode boundary enforces `temperature` finite +/// `>= 0` and `top_p` finite in `[0, 1]`; `max_tokens` is not range-checked +/// (a `0` sentinel is legitimate for some providers; the upstream decides). #[derive(Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct SamplingOptions { - /// Sampling temperature. - /// - /// Must be finite and `>= 0`. Both constraints are enforced symmetrically - /// on the decode path (via [`deserialize_temperature`]) and the encode path - /// (via [`serialize_finite_f64`]). Non-finite values are rejected because - /// `serde_json` silently converts them to `null`, which would silently lose - /// the value during a serialize -> deserialize round-trip. + /// Sampling temperature; must be finite and `>= 0`, enforced on both the + /// decode and encode paths (serde_json would turn non-finite into `null`). #[serde( default, serialize_with = "serialize_finite_f64", deserialize_with = "deserialize_temperature" )] pub temperature: Option, - /// Nucleus sampling parameter. - /// - /// Must be finite and in `[0, 1]`. See [`Self::temperature`] for the - /// finite-value handling rationale; the bounds are enforced symmetrically - /// via [`deserialize_top_p`] and [`serialize_finite_f64`]. + /// Nucleus sampling parameter; must be finite and in `[0, 1]`, enforced + /// symmetrically on the decode and encode paths. #[serde( default, serialize_with = "serialize_finite_f64", deserialize_with = "deserialize_top_p" )] pub top_p: Option, - /// Maximum number of tokens to generate. - /// - /// Not range-checked at the core decode boundary. The semantics of a `0` - /// or negative value vary by provider (Anthropic treats `<= 0` as invalid - /// in [`crate::anthropic`] validation; OpenAI accepts `max_tokens >= 0`), - /// so the proxy forwards the value and lets the upstream decide. This - /// mirrors the integration-test contract that `max_tokens=0` must not - /// trigger an internal error. + /// Maximum number of tokens to generate. Not range-checked at the core + /// decode boundary; providers treat `<= 0` differently, so forward it. pub max_tokens: Option, - /// Stop sequences. - /// - /// OpenAI's Chat Completions API allows `stop` to be a bare string - /// (`"STOP"`), an array (`["STOP"]`), or null. Anthropic always uses an - /// array. This field normalises all forms into `Vec` at - /// deserialization time: a bare string becomes a single-element vec. + /// Stop sequences; normalised at deserialization into `Vec` — a + /// bare OpenAI string becomes a single-element vec. #[serde(default, deserialize_with = "deserialize_stop")] pub stop: Option>, /// Reasoning effort level (e.g. `"low"`, `"medium"`, `"high"`). pub reasoning_effort: Option, - /// Extended thinking configuration (provider-specific JSON). - /// - /// This is intentionally opaque. Expected shapes vary by provider, e.g. - /// Anthropic: `{"type": "enabled", "budget_tokens": N}`. Because the field - /// accepts arbitrary client-supplied JSON (and a malicious payload could - /// embed secret fragments in field names or values), it is redacted in - /// [`fmt::Debug`] output via [`OpaqueJsonRef`] as a defense-in-depth - /// measure. Do not place genuinely secret data here and assume it will be - /// printed. + /// Extended thinking configuration (provider-specific JSON, e.g. Anthropic + /// `{"type": "enabled", "budget_tokens": N}`); redacted in `Debug` output. pub thinking: Option, } -/// Serializes an `Option`, rejecting non-finite values (NaN, Infinity). -/// -/// This is the symmetric encode-side guard for the finite-value invariants on -/// [`SamplingOptions::temperature`] and [`SamplingOptions::top_p`]. Without -/// it, an in-memory-constructed `CoreRequest` (e.g. one an adapter mutated -/// after decode) with a non-finite sampling value would serialize to `null` -/// via `serde_json`'s default float handling and silently lose the value on -/// any serialize -> deserialize round-trip through the proxy. fn serialize_finite_f64(val: &Option, serializer: S) -> Result where S: serde::Serializer, @@ -593,12 +430,6 @@ where } } -/// Deserializer for [`SamplingOptions::temperature`]. -/// -/// Enforces the universal sampling invariants at the point of ingestion: -/// the value must be finite and `>= 0`. `serde_json` silently converts -/// non-finite floats to `null`, which would otherwise become `None` on -/// round-trip; rejecting here keeps the invariant from being bypassed. fn deserialize_temperature<'de, D>(deserializer: D) -> Result, D::Error> where D: serde::Deserializer<'de>, @@ -615,12 +446,6 @@ where } } -/// Deserializer for [`SamplingOptions::top_p`]. -/// -/// Enforces the universal sampling invariants at the point of ingestion: -/// the value must be finite and in `[0, 1]`. See [`deserialize_temperature`] -/// for why finiteness is enforced here rather than left to a separate -/// validation step. fn deserialize_top_p<'de, D>(deserializer: D) -> Result, D::Error> where D: serde::Deserializer<'de>, @@ -637,13 +462,6 @@ where } } -/// Deserializer for [`CoreTool::input_schema`] that rejects non-object values. -/// -/// Tool parameter schemas must be JSON objects. The type system cannot encode -/// that constraint on a [`serde_json::Value`] field, so the invariant is -/// enforced here on the wire decode path rather than in an easy-to-bypass -/// constructor (mirroring the [`deserialize_temperature`] / [`deserialize_top_p`] -/// convention of baking invariants into the deserializer). fn deserialize_input_schema_object<'de, D>(deserializer: D) -> Result where D: serde::Deserializer<'de>, @@ -657,9 +475,6 @@ where Ok(v) } -/// Custom deserializer for `SamplingOptions::stop` that accepts a bare string, -/// an array of strings, or null. This normalises the OpenAI wire format where -/// `stop` can be polymorphic. fn deserialize_stop<'de, D>(deserializer: D) -> Result>, D::Error> where D: serde::Deserializer<'de>, @@ -696,12 +511,6 @@ where } } -/// Returns a human-readable name for a JSON value type. -/// -/// `pub(crate)` so the per-provider client decoders (e.g. OpenAI Chat's -/// `decode_request`) can produce byte-identical stop-array error messages to -/// the canonical `deserialize_stop`, keeping the two wire-format rules from -/// diverging again (audit GAP-LOW-8). pub(crate) fn json_type_name(val: &serde_json::Value) -> &'static str { match val { serde_json::Value::Null => "null", @@ -713,20 +522,6 @@ pub(crate) fn json_type_name(val: &serde_json::Value) -> &'static str { } } -impl SamplingOptions { - // The universal sampling invariants (temperature finite & >= 0, top_p - // finite & in [0, 1]) are enforced directly in the field deserializers - // ([`deserialize_temperature`], [`deserialize_top_p`]) so they cannot be - // bypassed. `max_tokens` range validation is deliberately left to the - // provider adapter / upstream (see the field doc); it is not a core-decode - // invariant. - // - // There is intentionally no `validate()` method here: a previous version - // had one, but it was never wired into any request-handling path, which - // gave a false sense of correctness. Baking the checks into the - // deserializer removes that trap. -} - impl fmt::Debug for SamplingOptions { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("SamplingOptions") @@ -740,10 +535,6 @@ impl fmt::Debug for SamplingOptions { } } -// --------------------------------------------------------------------------- -// RequestMetadata -// --------------------------------------------------------------------------- - /// Caller metadata attached to a request. #[derive(Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -763,7 +554,6 @@ impl fmt::Debug for RequestMetadata { } } -/// Helper that wraps a Map reference for redacted Debug output. struct RedactedMap<'a>(&'a serde_json::Map); fn redact_map_helper(map: &serde_json::Map) -> RedactedMap<'_> { @@ -776,23 +566,9 @@ impl fmt::Debug for RedactedMap<'_> { } } -// --------------------------------------------------------------------------- -// ProviderHints -// --------------------------------------------------------------------------- - -/// Opaque hints consumed only by provider adapters. -/// -/// Provider hints carry client-side metadata that influences how a provider -/// adapter encodes a request, but which does not belong in the canonical -/// `CoreRequest` fields. Examples include: -/// -/// - `stream_options` (OpenAI Chat: `{"include_usage": true}`) -/// - ` thinking` configuration passthrough -/// - provider-specific flags like `reasoning_effort` mapping hints -/// -/// **Architectural rule:** No adapter may special-case another protocol's -/// hints. Client adapters write hints; provider adapters read only the -/// hints they understand. Unknown keys must be silently ignored. +/// Opaque hints consumed only by provider adapters (e.g. OpenAI +/// `stream_options`). No adapter may special-case another protocol's hints; +/// unknown keys must be silently ignored. #[derive(Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct ProviderHints { @@ -808,10 +584,6 @@ impl fmt::Debug for ProviderHints { } } -// --------------------------------------------------------------------------- -// CoreResponse -// --------------------------------------------------------------------------- - /// A normalized non-streaming chat response. #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -828,19 +600,13 @@ pub struct CoreResponse { pub stop_sequence: Option, /// Token usage statistics. pub usage: Usage, - /// Provider-specific metadata (warnings, extra fields, etc.). - /// - /// Uses a `Map` (not bare `Value`) so that the field serializes as `{}` - /// when empty rather than `null`, distinguishing "no metadata" from - /// "null metadata". + /// Provider-specific metadata; a `Map` (not bare `Value`) so it serializes + /// as `{}` when empty rather than `null`. #[serde(default)] pub provider_meta: serde_json::Map, - /// Computed USD cost for this request, if pricing was configured for the - /// model. `None` when no pricing is configured; in that case the field is - /// omitted from serialization so the wire shape is byte-identical to a - /// build without cost accounting. Not surfaced on the Anthropic/OpenAI - /// client wire (those formats have no cost field); available to the event - /// log and the future API crate. + /// Computed USD cost, `None` (and omitted from serialization) when no + /// pricing is configured. Not on the Anthropic/OpenAI client wire; + /// available to the event log and API crate. #[serde(default, skip_serializing_if = "Option::is_none")] pub cost: Option, } @@ -860,18 +626,8 @@ impl fmt::Debug for CoreResponse { } } -// --------------------------------------------------------------------------- -// Cost -// --------------------------------------------------------------------------- - -/// Computed cost for one request, in USD. -/// -/// All fields are [`Decimal`] to avoid floating-point money. Derived from -/// [`Usage`] and a per-token price table; the multiplication lives in the -/// server crate (the first crate that sees both [`Usage`] and the core -/// `ModelPricing` type). A request for a model with no configured pricing -/// simply has no [`Cost`] (the carrier is `Option`), so responses are -/// unchanged when pricing is absent. +/// Computed cost for one request, in USD, derived from [`Usage`] and a +/// per-token price table; `Decimal` throughout to avoid floating-point money. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct Cost { @@ -889,10 +645,6 @@ pub struct Cost { pub total: Decimal, } -// --------------------------------------------------------------------------- -// StopReason -// --------------------------------------------------------------------------- - /// Why the model stopped generating. #[non_exhaustive] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -910,31 +662,14 @@ pub enum StopReason { /// An error occurred. Error, /// The stop reason was not recognised by the adapter's finish-reason map. - /// - /// The original provider value is carried in the inner string so it - /// survives the normalize → core → encode pipeline and remains available in - /// diagnostics: the encoder logs it when mapping to a native finish reason. - /// This mirrors [`CacheControlType::Other`], which solves the same - /// problem for cache-control types. The sentinel `"absent"` indicates the - /// provider omitted the finish-reason field entirely rather than sending an - /// unmapped value. + /// The original provider value is carried so it survives into diagnostics; + /// the sentinel `"absent"` means the provider omitted the field entirely. Unknown(String), } -// --------------------------------------------------------------------------- -// Usage / UsageProvenance -// --------------------------------------------------------------------------- - -/// Token usage statistics. -/// -/// Must carry provenance. Never report fabricated or zero-filled usage as -/// [`UsageProvenance::ProviderReported`]; absent or zero-filled provider -/// usage is [`UsageProvenance::SyntheticZero`] / [`UsageProvenance::Unknown`]. -/// -/// Uses `i32` for token counts for consistency with provider wire formats -/// (JSON integers) and the Anthropic/OpenAI API types. The ~2.1 billion limit -/// is well beyond any realistic single request. If aggregated or batch usage -/// is needed in the future, consider widening to `i64`. +/// Token usage statistics; must carry [`UsageProvenance`] — never report +/// fabricated or zero-filled usage as `ProviderReported`. Counts are `i32`, +/// matching provider wire formats. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct Usage { @@ -953,11 +688,8 @@ pub struct Usage { } impl Usage { - /// Creates a synthetic-zero usage with all token counts set to zero. - /// - /// Use this when the provider did not report any usage data. - /// The provenance is set to [`UsageProvenance::SyntheticZero`] to - /// distinguish synthesised zeros from genuine provider-reported zeros. + /// Creates a synthetic-zero usage; provenance is + /// [`UsageProvenance::SyntheticZero`], distinguishing it from provider zeros. pub fn synthetic_zero() -> Self { Self { input_tokens: 0, @@ -967,10 +699,8 @@ impl Usage { } } - /// Creates a provider-reported usage with the given token counts. - /// - /// The provenance is set to [`UsageProvenance::ProviderReported`]. - /// Callers must ensure these numbers actually came from the provider. + /// Creates a provider-reported usage; callers must ensure the numbers + /// actually came from the provider. pub fn provider_reported(input_tokens: i32, output_tokens: i32) -> Self { Self { input_tokens, @@ -994,18 +724,9 @@ pub enum UsageProvenance { Unknown, } -// --------------------------------------------------------------------------- -// CoreEvent (streaming) -// --------------------------------------------------------------------------- - -/// A single normalised streaming event. -/// -/// V1 streaming carries only `Text`, `Thinking`, and `ToolUse` incrementally -/// (text/thinking/tool-call deltas). `Image`, `Document`, `Audio`, and -/// `Video` have no per-delta stream event; if a provider streams them the -/// adapter buffers the block and emits it through the non-stream -/// [`CoreResponse`] content path, or rejects it. Stream encoders therefore -/// never receive deltas for those kinds. +/// A single normalised streaming event. V1 streams only text, thinking, and +/// tool-use incrementally; image/document/audio/video have no per-delta +/// event and travel via the non-stream [`CoreResponse`] path or are rejected. #[non_exhaustive] #[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -1135,10 +856,6 @@ impl fmt::Debug for CoreEvent { } } -// --------------------------------------------------------------------------- -// ContentKind -// --------------------------------------------------------------------------- - /// Discriminator for content block types. #[non_exhaustive] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -1149,12 +866,9 @@ pub enum ContentKind { Thinking, /// Tool invocation. ToolUse, - /// Tool result. - /// - /// Note: `ToolResult` has no corresponding streaming delta event. It is - /// used only for `ContentStart` signalling in non-standard cases where a - /// tool result block begins in a stream. The actual content is carried - /// through the non-stream `CoreResponse` content path. + /// Tool result. Has no streaming delta event; used only for `ContentStart` + /// signalling — actual content travels through the non-stream + /// [`CoreResponse`] content path. ToolResult, /// Image. Image, @@ -1168,55 +882,28 @@ pub enum ContentKind { Refusal, } -// --------------------------------------------------------------------------- -// CoreStreamError / CoreStreamErrorKind -// --------------------------------------------------------------------------- - -/// An error that occurred during streaming. -/// -/// Named `CoreStreamError` (not `CoreError`) to avoid colliding with the -/// config/registry error owned by the core crate. The two are unrelated -/// types in different crates. -/// -/// # Message safety -/// -/// The `message` field is intentionally `pub(crate)` to enforce that all -/// construction goes through [`CoreStreamError::new()`], which documents the -/// sanitization contract. Adapters must strip API keys, bearer tokens, and -/// other secrets before passing the message to `new()`. -/// -/// The `Debug` impl redacts the message to its character count. The `Display` -/// impl shows only the error kind, not the message content. The `Serialize` -/// impl includes the message verbatim for internal transport between proxy -/// components -- do **not** serialize `CoreStreamError` directly into -/// client-facing API responses. +/// An error that occurred during streaming. The `message` is `pub(crate)` +/// and must be sanitized (strip API keys/tokens) before +/// [`CoreStreamError::new`]; `Debug`/`Display` redact it — never serialize it to clients. #[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct CoreStreamError { /// Category of the error. pub kind: CoreStreamErrorKind, - /// Human-readable error message (sanitized by the constructor). #[serde(skip_serializing, default)] pub(crate) message: String, } impl CoreStreamError { - /// Constructs a new `CoreStreamError` with the given kind and message. - /// - /// Callers **must** sanitize `message` before passing it: strip API keys, - /// bearer tokens, and other secrets from upstream provider error text. + /// Constructs a new `CoreStreamError`; callers **must** sanitize `message` + /// (strip API keys, bearer tokens, and other secrets) before passing it. pub fn new(kind: CoreStreamErrorKind, message: String) -> Self { Self { kind, message } } - /// Returns the error message. - /// - /// The message was sanitized at construction time by the adapter that - /// created this error. Note that the `message` field is - /// `#[serde(skip_serializing)]`, so values reconstructed via `Deserialize` - /// (e.g. after being transported between proxy components) return an empty - /// string here -- the message is only present on instances built directly - /// via [`new`](CoreStreamError::new). + /// Returns the error message. Because the field is + /// `#[serde(skip_serializing)]`, instances rebuilt via `Deserialize` + /// return `""`; only [`new`](CoreStreamError::new) instances carry it. pub fn message(&self) -> &str { &self.message } @@ -1226,7 +913,6 @@ impl fmt::Debug for CoreStreamError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CoreStreamError") .field("kind", &self.kind) - // Truncate message to avoid leaking sensitive error details. .field("message", &format_args!("{} chars", self.message.len())) .finish() } @@ -1234,28 +920,14 @@ impl fmt::Debug for CoreStreamError { impl fmt::Display for CoreStreamError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Intentionally do NOT include self.message here to avoid leaking - // unsanitized upstream error text through Display (used by - // tracing::error!("{err}"), error chains, etc.). Use - // `error.message()` explicitly if the full message is needed. write!(f, "{}", self.kind) } } impl std::error::Error for CoreStreamError {} -// NOTE: CoreStreamError uses a manual Error impl rather than thiserror derive -// because it also carries custom Serialize/Deserialize, a redacting Debug that -// hides the message field, and a Display that intentionally omits the message -// for security. thiserror's `#[error()]` attribute would conflict with the -// hand-rolled Display, and the derive macro does not support the serde + -// redacting-Debug combination used here. - -/// Category of a stream error. -/// -/// Each variant maps to a canonical HTTP status code via -/// [`CoreStreamErrorKind::http_status()`]. Adapters should use this mapping -/// rather than inventing their own, so that all providers produce consistent -/// HTTP responses. + +/// Category of a stream error; use [`CoreStreamErrorKind::http_status`] +/// for the canonical HTTP mapping rather than adapter-specific codes. #[non_exhaustive] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum CoreStreamErrorKind { @@ -1275,14 +947,6 @@ pub enum CoreStreamErrorKind { impl CoreStreamErrorKind { /// Returns the canonical HTTP status code for this error kind. - /// - /// Mapping: - /// - `InvalidRequest` -> 400 - /// - `Authentication` -> 401 - /// - `Permission` -> 403 - /// - `RateLimit` -> 429 - /// - `Upstream` -> 502 - /// - `Internal` -> 500 pub fn http_status(&self) -> u16 { match self { Self::InvalidRequest => 400, @@ -1308,10 +972,6 @@ impl fmt::Display for CoreStreamErrorKind { } } -// =========================================================================== -// Unit tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -1384,11 +1044,6 @@ mod tests { } } - // The manual `Debug` impl for `CoreResponse` (unlike a derived one) is - // hand-maintained, so a future edit could silently drop the `cost` field - // from its output. Since `cost` is the only audit-relevant field surfaced - // for observability/pricing diagnostics, pin its presence in the `{:?}` - // output here. #[test] fn core_response_debug_includes_cost_when_present() { use std::str::FromStr; @@ -1408,20 +1063,15 @@ mod tests { usage: Usage::default(), provider_meta: serde_json::Map::new(), cost: Some(Cost { - // Only set `total`; the rest come from `Default` so this test is - // robust to field reordering/renaming of the per-component fields. total: Decimal::from_str("0.003").unwrap(), ..Cost::default() }), }; let debug = format!("{resp:?}"); - // The field name must appear (manual Debug could drop it silently). assert!( debug.contains("cost:"), "CoreResponse Debug must include the `cost` field, got: {debug}" ); - // And the value must render (Decimal Debug is non-empty), proving the - // field carries the `Some(Cost)` payload rather than being stubbed. assert!( debug.contains("0.003"), "CoreResponse Debug must surface the total cost value, got: {debug}" @@ -1430,9 +1080,6 @@ mod tests { #[test] fn core_response_debug_sensible_when_cost_absent() { - // When no pricing is configured, `cost` is `None`. The manual Debug - // must still render the field name (so the field is observable as - // absent, not silently elided) with a `None` value. let resp = CoreResponse { id: None, model: ModelRef { @@ -1611,7 +1258,6 @@ mod tests { provider_hints: ProviderHints::default(), }; assert_eq!(req.system.len(), 2); - // System content stays in `system`, not in messages. assert_eq!(req.messages.len(), 1); assert_eq!(req.messages[0].role, CoreRole::User); } @@ -1647,7 +1293,6 @@ mod tests { }; let json = serde_json::to_string(&req).expect("serialize CoreRequest"); let back: CoreRequest = serde_json::from_str(&json).expect("deserialize CoreRequest"); - // With PartialEq on CoreRequest, we can compare the whole struct. assert_eq!(req, back); } @@ -1685,7 +1330,6 @@ mod tests { }; let json = serde_json::to_string(&resp).expect("serialize CoreResponse"); let back: CoreResponse = serde_json::from_str(&json).expect("deserialize CoreResponse"); - // Full structural equality catches any serialization drift. assert_eq!(resp, back); } @@ -1758,13 +1402,10 @@ mod tests { let json = serde_json::to_string(&events).expect("serialize CoreEvent vec"); let back: Vec = serde_json::from_str(&json).expect("deserialize CoreEvent vec"); - // Compare all events except the Error variant, whose `message` field - // uses #[serde(skip_serializing)] and will be empty after round-trip. for (i, (original, round_tripped)) in events.iter().zip(back.iter()).enumerate() { match (original, round_tripped) { (CoreEvent::Error { error: orig }, CoreEvent::Error { error: rt }) => { assert_eq!(orig.kind, rt.kind, "kind mismatch at event {i}"); - // message is skip_serializing, so it comes back empty. assert!( rt.message.is_empty(), "message should be empty after round-trip at event {i}" @@ -1816,13 +1457,8 @@ mod tests { assert_eq!(back.provenance, UsageProvenance::ProviderReported); } - // ======================================================================= - // New edge-case and coverage tests (audit round 1) - // ======================================================================= - #[test] fn cache_control_type_known_and_unknown_variants() { - // Test serialization (takes ownership) then From conversion on a fresh value. let json = serde_json::to_string(&CacheControlType::Ephemeral).unwrap(); assert_eq!(json, "\"ephemeral\""); let back: CacheControlType = serde_json::from_str(&json).unwrap(); @@ -1904,11 +1540,7 @@ mod tests { #[test] fn core_request_rejects_malformed_json() { - let cases = vec![ - "not json at all", - "{", - r#"{"model": 123}"#, // model should be object, not number - ]; + let cases = vec!["not json at all", "{", r#"{"model": 123}"#]; for bad in cases { assert!( serde_json::from_str::(bad).is_err(), @@ -1919,11 +1551,7 @@ mod tests { #[test] fn core_response_rejects_malformed_json() { - let cases = vec![ - "not json at all", - "{", - r#"{"id": 123}"#, // id should be string or null - ]; + let cases = vec!["not json at all", "{", r#"{"id": 123}"#]; for bad in cases { assert!( serde_json::from_str::(bad).is_err(), @@ -2002,9 +1630,7 @@ mod tests { #[test] fn core_request_rejects_wrong_field_types() { let cases = vec![ - // model.requested should be string, not number r#"{"model": {"requested": 123}, "messages": [], "system": [], "tools": [], "sampling": {}, "stream": false, "metadata": {}, "provider_hints": {}}"#, - // stream should be bool, not string r#"{"model": {"requested": "m"}, "messages": [], "system": [], "tools": [], "sampling": {}, "stream": "yes", "metadata": {}, "provider_hints": {}}"#, ]; for bad in cases { @@ -2169,8 +1795,6 @@ mod tests { let json = serde_json::to_string(&err).expect("serialize CoreStreamError"); let back: CoreStreamError = serde_json::from_str(&json).expect("deserialize CoreStreamError"); - // kind round-trips; message is #[serde(skip_serializing)] so it - // comes back empty. assert_eq!(err.kind, back.kind); assert!(back.message.is_empty()); } @@ -2341,13 +1965,8 @@ mod tests { assert_eq!(req, back); } - // ======================================================================= - // Audit round 2 -- edge-case and coverage tests - // ======================================================================= - #[test] fn sampling_options_stop_deserializes_bare_string() { - // OpenAI allows stop: "STOP" as a bare string. let json = r#"{"stop": "STOP"}"#; let opts: SamplingOptions = serde_json::from_str(json).unwrap(); assert_eq!(opts.stop, Some(vec!["STOP".to_owned()])); @@ -2375,8 +1994,6 @@ mod tests { #[test] fn sampling_options_rejects_out_of_range_sampling_values() { - // temperature >= 0 and top_p in [0, 1] are enforced at the decode - // boundary so a client cannot smuggle them through to the upstream. let bad_temp = r#"{"temperature": -0.5}"#; assert!( serde_json::from_str::(bad_temp).is_err(), @@ -2392,8 +2009,6 @@ mod tests { serde_json::from_str::(bad_top_p_high).is_err(), "top_p > 1 should be rejected at decode time" ); - // max_tokens is intentionally NOT range-checked here (provider-specific - // semantics): a 0 or negative value must still decode successfully. let zero_max = r#"{"max_tokens": 0}"#; let opts: SamplingOptions = serde_json::from_str(zero_max).expect("max_tokens=0 must decode"); @@ -2402,9 +2017,6 @@ mod tests { #[test] fn sampling_options_nan_temperature_is_rejected_on_serialize_and_deserialize() { - // The encode side guards against NaN so an in-memory-constructed value - // (e.g. one an adapter mutated after decode) cannot silently lose data - // via serde_json's NaN -> null conversion during a round-trip. let opts = SamplingOptions { temperature: Some(f64::NAN), ..Default::default() @@ -2413,11 +2025,7 @@ mod tests { serde_json::to_string(&opts).is_err(), "serializing NaN temperature should error, not silently emit null" ); - // The decode side rejects NaN at ingestion. let bad = r#"{"temperature": NaN}"#; - // serde_json rejects bare `NaN` tokens at the lexer level already; the - // deserializer guard is the second line of defence for any float that - // parses as non-finite. assert!( serde_json::from_str::(bad).is_err(), "NaN temperature should not decode" @@ -2450,7 +2058,6 @@ mod tests { #[test] fn core_request_large_messages_round_trips() { - // Build a request with many messages to exercise large-input handling. let messages: Vec = (0..1000) .map(|i| CoreMessage { role: CoreRole::User, @@ -2481,8 +2088,6 @@ mod tests { #[test] fn deeply_nested_tool_result_round_trips() { - // Build a deeply nested ToolResult chain. Depth kept at 20 to stay - // within serde_json's default recursion limit (128). let depth = 20; let mut content = CoreContent::Text { text: "leaf".into(), @@ -2521,7 +2126,6 @@ mod tests { #[test] fn core_stream_error_display_and_error_traits() { let err = CoreStreamError::new(CoreStreamErrorKind::RateLimit, "too many requests".into()); - // Display shows the kind, not the raw message (defense-in-depth). let display = format!("{err}"); assert!( display.contains("rate_limit"), @@ -2532,7 +2136,6 @@ mod tests { "Display should not leak message" ); - // Error trait let _: &dyn std::error::Error = &err; } @@ -2687,10 +2290,6 @@ mod tests { assert!(!debug.contains("secret"), "Debug should redact Raw value"); } - // ======================================================================= - // Audit round 3 -- deny_unknown_fields, missing round-trips, edge cases - // ======================================================================= - #[test] fn core_request_rejects_unknown_fields() { let json = r#"{ @@ -2754,13 +2353,11 @@ mod tests { #[test] fn sampling_options_stop_rejects_array_with_non_strings() { - // Array containing a number should be rejected. let json = r#"{"stop": ["STOP", 42]}"#; assert!( serde_json::from_str::(json).is_err(), "should reject array with non-string items" ); - // Array containing a boolean should be rejected. let json2 = r#"{"stop": [true]}"#; assert!( serde_json::from_str::(json2).is_err(), @@ -2825,13 +2422,10 @@ mod tests { #[test] fn bool_and_number_redacted_in_debug() { - // Bool and Number values should not show their actual values in Debug. let content = CoreContent::Image { source: serde_json::json!({"flag": true, "count": 42}), }; let debug = format!("{content:?}"); - // The outer Object type is shown but Bool/Number values inside are - // not printed verbatim -- redact_value hides inner details. assert!( debug.contains("Object"), "should show Object type for source" @@ -2840,19 +2434,16 @@ mod tests { #[test] fn core_content_rejects_unknown_fields_in_variant() { - // Text variant with an extra field should be rejected, not silently dropped. let json = r#"{"Text": {"text": "hi", "secret_extra": "leaked"}}"#; assert!( serde_json::from_str::(json).is_err(), "CoreContent::Text should reject unknown field" ); - // ToolUse variant with an extra field. let json2 = r#"{"ToolUse": {"id": "1", "name": "fn", "input": {}, "extra": true}}"#; assert!( serde_json::from_str::(json2).is_err(), "CoreContent::ToolUse should reject unknown field" ); - // ToolResult variant with an extra field. let json3 = r#"{"ToolResult": {"tool_use_id": "1", "content": [], "is_error": false, "extra": 1}}"#; assert!( @@ -2863,19 +2454,16 @@ mod tests { #[test] fn core_event_rejects_unknown_fields_in_variant() { - // TextDelta with an extra field. let json = r#"{"TextDelta": {"index": 0, "text": "hi", "extra": true}}"#; assert!( serde_json::from_str::(json).is_err(), "CoreEvent::TextDelta should reject unknown field" ); - // MessageStart with an extra field. let json2 = r#"{"MessageStart": {"id": null, "model": {"requested": "m"}, "extra": 1}}"#; assert!( serde_json::from_str::(json2).is_err(), "CoreEvent::MessageStart should reject unknown field" ); - // MessageStop with an extra field. let json3 = r#"{"MessageStop": {"stop_reason": "EndTurn", "stop_sequence": null, "extra": true}}"#; assert!( diff --git a/crates/llm-proxy-protocol/src/lib.rs b/crates/llm-proxy-protocol/src/lib.rs index c939084..637f175 100644 --- a/crates/llm-proxy-protocol/src/lib.rs +++ b/crates/llm-proxy-protocol/src/lib.rs @@ -1,8 +1,6 @@ -//! llm-proxy-protocol: wire-format types for the LLM proxy. -//! -//! Contains request/response types for the Anthropic Messages API, -//! OpenAI Chat Completions API, Responses API, Google Gemini API, -//! and the normalised core protocol types. +//! llm-proxy-protocol: wire-format types for the LLM proxy — Anthropic +//! Messages, OpenAI Chat Completions, Responses API, Gemini API, and the +//! normalised core protocol types. #![deny(missing_docs)] diff --git a/crates/llm-proxy-protocol/src/openai.rs b/crates/llm-proxy-protocol/src/openai.rs index 75d9d85..771e70a 100644 --- a/crates/llm-proxy-protocol/src/openai.rs +++ b/crates/llm-proxy-protocol/src/openai.rs @@ -1,27 +1,12 @@ -//! OpenAI Chat Completions API types. -//! -//! Covers request/response structures for the `/v1/chat/completions` endpoint, -//! including streaming chunks and error payloads. -//! +//! OpenAI Chat Completions API types for `/v1/chat/completions`, including +//! streaming chunks and error payloads. //! Reference: use serde::{Deserialize, Serialize}; -// --------------------------------------------------------------------------- -// Cache control (shared with the Anthropic module) -// --------------------------------------------------------------------------- - -/// Cache-control directive attached to a message or content block. -/// -/// This is a separate type from [`crate::core::CacheControl`] by design: -/// the OpenAI wire format uses a plain string for the `type` field, while -/// the core type uses a proper [`CacheControlType`](crate::core::CacheControlType) -/// enum with an `Other(String)` catch-all. The two modules represent different -/// layers (wire format vs. canonical internal representation) and duplicating -/// the struct keeps the serde concerns cleanly separated. -/// -/// Anthropic models accept `cache_control` on messages; the proxy -/// preserves the annotation when translating between formats. +/// Cache-control directive attached to a message or content block; distinct +/// from [`crate::core::CacheControl`] because the OpenAI wire `type` field is +/// a plain string rather than a typed enum. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct CacheControl { /// Discriminator – typically `"ephemeral"`. @@ -33,10 +18,6 @@ impl CacheControl { pub const EPHEMERAL: &str = "ephemeral"; } -// --------------------------------------------------------------------------- -// Streaming options -// --------------------------------------------------------------------------- - /// Controls extra metadata returned alongside streaming chunks. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct StreamOptions { @@ -45,18 +26,11 @@ pub struct StreamOptions { pub include_usage: Option, } -// --------------------------------------------------------------------------- -// Tool definitions -// --------------------------------------------------------------------------- - /// A tool (function) that the model may invoke during generation. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ToolDef { - /// Always `"function"` for function-calling tools. - /// - /// Kept as a bare `String` for serde round-tripping with providers that may - /// introduce new tool types in the future. The adapter does not validate this - /// value; callers that need to enforce `"function"` should check at decode time. + /// Always `"function"` for function-calling tools; kept as a bare `String` + /// for serde round-tripping with providers that add new tool types. pub r#type: String, /// Schema of the function the model can call. pub function: FunctionDef, @@ -75,22 +49,12 @@ pub struct FunctionDef { pub parameters: Option, } -// --------------------------------------------------------------------------- -// Function / tool calls inside messages -// --------------------------------------------------------------------------- - -/// A single function invocation produced by the model. -/// -/// In streaming delta chunks both `id` and `function` may be absent; -/// the server sends partial `ToolCall`s that the client merges by `index`. +/// A single function invocation produced by the model; in streaming deltas +/// `id`/`function` may be absent and chunks merge by `index`. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ToolCall { - /// Position inside the `tool_calls` array (streaming deltas only). - /// - /// Uses `i32` to match the OpenAI wire format (JSON integers). The OpenAI - /// spec defines this as a non-negative integer; negative values are - /// semantically meaningless. The stream encoder clamps `usize` core indices - /// to `i32::MAX` when converting. See `clamp_tool_index` in the adapter. + /// Position inside the `tool_calls` array (streaming deltas only); `i32` + /// matching the OpenAI wire format, clamped from `usize` by the encoder. #[serde(skip_serializing_if = "Option::is_none")] pub index: Option, /// Stable identifier for this tool call (absent in deltas). @@ -104,12 +68,8 @@ pub struct ToolCall { pub function: Option, } -/// Name + arguments of a function call. -/// -/// Both fields are `Option` to accommodate streaming deltas where the -/// server sends partial `ToolCall`s that the client merges by `index`. For -/// non-streaming (complete) responses, callers should validate that `name` and -/// `arguments` are present. +/// Name + arguments of a function call; both `Option` to accommodate +/// streaming deltas that merge by `index`. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct FunctionCall { /// The function name chosen by the model. @@ -120,44 +80,22 @@ pub struct FunctionCall { pub arguments: Option, } -// --------------------------------------------------------------------------- -// Chat message -// --------------------------------------------------------------------------- - -/// A single message in a conversation. -/// -/// For assistant messages that carry tool calls, `content` may be empty -/// while `tool_calls` is populated. For tool-result messages, `tool_call_id` -/// identifies the call being answered. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because this struct is -/// used in streaming deltas where fields may vary between providers. Unknown -/// fields are silently ignored rather than causing parse failures. +/// A single message in a conversation. Assistant messages may carry +/// `tool_calls` with empty `content`; unknown fields are ignored (no +/// `deny_unknown_fields`) because streaming deltas vary by provider. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatMessage { - /// `"system"`, `"user"`, `"assistant"`, or `"tool"`. - /// - /// Defaults to empty when absent (streaming deltas often omit this). - /// Validation of known role values happens at decode time in the adapter - /// rather than at the type level, to tolerate provider-specific extensions. + /// `"system"`, `"user"`, `"assistant"`, or `"tool"`; empty when absent + /// (streaming deltas omit it). Known-role validation happens at decode. #[serde(default)] pub role: String, - /// The content body of the message. - /// - /// OpenAI allows `content` to be either a plain string or an array of - /// structured content parts (e.g. `[{"type":"text","text":"..."}, - /// {"type":"image_url","image_url":{...}}]`). Using `serde_json::Value` - /// lets us accept both forms without rejecting multimodal requests. - /// - /// Defaults to `Value::String("")` when absent (streaming deltas often - /// omit this). + /// The content body: a plain string or an array of structured content + /// parts (e.g. text/`image_url`); `serde_json::Value` accepts both forms. + /// Defaults to `""` when absent (streaming deltas often omit it). #[serde(default = "default_content")] pub content: serde_json::Value, - /// Chain-of-thought text emitted by reasoning models. - /// - /// Some OpenAI-compatible providers (e.g. DeepSeek) use the `reasoning` - /// field name instead of `reasoning_content`. The `serde` alias ensures - /// both names deserialize into this field. + /// Chain-of-thought text from reasoning models; the `reasoning` alias + /// covers providers (e.g. DeepSeek) that use that field name. #[serde(alias = "reasoning", skip_serializing_if = "Option::is_none")] pub reasoning_content: Option, /// Tool calls emitted by an assistant message. @@ -183,15 +121,9 @@ pub struct ChatMessage { } impl ChatMessage { - /// Extract the text content as a `String`. - /// - /// Handles both forms accepted by the OpenAI Chat Completions API: - /// - `Value::String(s)` — returned directly. - /// - `Value::Array(parts)` — concatenates all `"type":"text"` parts' `"text"` - /// fields. Non-text parts (e.g. `image_url`) are silently skipped. - /// - `Value::Null` — returns an empty string. - /// - /// Any other value type returns an empty string with a warning. + /// Extract the text content as a `String`: `Value::String` returned + /// directly, array parts concatenate their `"type":"text"` `"text"` + /// fields, anything else yields `""` (with a warning). pub fn content_text(&self) -> String { match &self.content { serde_json::Value::String(s) => s.clone(), @@ -218,11 +150,8 @@ impl ChatMessage { } } - /// Returns `true` when `content` is semantically empty. - /// - /// This is `true` when the value is `Null`, an empty string (`""`), or an - /// empty array. It is the replacement for the old `msg.content.is_empty()` - /// checks used throughout the codebase. + /// Returns `true` when `content` is `Null`, an empty string, or an + /// empty array. pub fn content_is_empty(&self) -> bool { match &self.content { serde_json::Value::String(s) => s.is_empty(), @@ -233,22 +162,13 @@ impl ChatMessage { } } -/// Default value for [`ChatMessage::content`]: an empty JSON string. fn default_content() -> serde_json::Value { serde_json::Value::String(String::new()) } -// --------------------------------------------------------------------------- -// Chat completion request -// --------------------------------------------------------------------------- - -/// Request body for `POST /v1/chat/completions`. -/// -/// This type serves a dual purpose: it is both deserialized from inbound client -/// requests (via `openai_chat::decode_request`) and constructed internally for -/// outbound provider calls (via `transform_request()`). Unknown fields from -/// clients are captured via the `extra` flattened map rather than silently -/// dropped. +/// Request body for `POST /v1/chat/completions`, deserialized from inbound +/// client requests and constructed internally for outbound provider calls; +/// unknown client fields are captured in `extra` rather than dropped. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionRequest { /// Model identifier (e.g. `"gpt-4o"`, `"glm-5.1"`). @@ -264,15 +184,9 @@ pub struct ChatCompletionRequest { /// Nucleus sampling threshold. #[serde(skip_serializing_if = "Option::is_none")] pub top_p: Option, - /// Upper bound on tokens generated in the response. - /// - /// OpenAI renamed `max_tokens` to `max_completion_tokens` in later API - /// versions. Both names are accepted via the serde alias; the canonical - /// field name follows the newer convention. - /// - /// Uses `i32` for consistency with provider wire formats. Negative values - /// are semantically invalid; the adapter should validate `max_tokens >= 0` - /// during decode. + /// Upper bound on tokens generated in the response. OpenAI renamed + /// `max_tokens` to `max_completion_tokens`; both are accepted (serde + /// alias) and the canonical field follows the newer convention. #[serde( rename = "max_completion_tokens", alias = "max_tokens", @@ -291,12 +205,9 @@ pub struct ChatCompletionRequest { /// Controls which (if any) tool the model must call. #[serde(skip_serializing_if = "Option::is_none")] pub tool_choice: Option, - /// Stop sequences (string or array of strings). - /// - /// Kept as `Option` (rather than `Option>`) - /// to preserve the raw OpenAI wire format for passthrough. The decode - /// function in `openai_chat::decode_request` normalises this into - /// `SamplingOptions.stop` using the same logic as `core::deserialize_stop`. + /// Stop sequences (string or array of strings), kept as raw + /// `serde_json::Value` for passthrough; `openai_chat::decode_request` + /// normalises it into `SamplingOptions.stop`. #[serde(skip_serializing_if = "Option::is_none")] pub stop: Option, /// Extra streaming metadata flags. @@ -308,43 +219,23 @@ pub struct ChatCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub user: Option, /// Catch-all for unrecognized client fields, preserved in - /// `RequestMetadata.raw` during decode so nothing is silently dropped. - /// - /// Because `#[serde(flatten)]` and `#[serde(deny_unknown_fields)]` are - /// mutually exclusive, this struct cannot use `deny_unknown_fields`. - /// Fields that look like typos of known fields (e.g. `"temperatur"` instead - /// of `"temperature"`) will end up here silently. The decode function - /// should warn on such cases if desired. + /// `RequestMetadata.raw` during decode. `#[serde(flatten)]` precludes + /// `deny_unknown_fields`, so typo'd fields land here silently. #[serde(flatten)] pub extra: serde_json::Map, } -// --------------------------------------------------------------------------- -// Token usage -// --------------------------------------------------------------------------- - -/// Token usage statistics returned by the API. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because providers may -/// return additional usage fields (e.g. `prompt_tokens_details`) that the -/// proxy does not model. Unknown fields are silently ignored rather than -/// causing parse failures. -/// -/// `Default` is derived so the [`ChatCompletionResponse`] `usage` field can use -/// `#[serde(default)]`: some OpenAI-compatible providers omit `usage` entirely -/// on 2xx responses, and absent usage should yield zero counts rather than a -/// decode failure. +/// Token usage statistics returned by the API; unknown provider fields are +/// ignored, and `Default` lets an absent or `null` `usage` on a 2xx decode +/// as zero counts instead of failing. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] pub struct UsageInfo { /// Tokens consumed by the prompt. pub prompt_tokens: i32, /// Tokens produced by the completion. pub completion_tokens: i32, - /// `prompt_tokens + completion_tokens`. - /// - /// The proxy recomputes this via `saturating_add` during encoding rather than - /// passing through the provider's value, ensuring consistency. If a provider - /// reports a different total, the proxy's computed value takes precedence. + /// `prompt_tokens + completion_tokens`; recomputed via `saturating_add` + /// during encoding, so the proxy's value wins over the provider's. pub total_tokens: i32, /// Prompt tokens served from a cache (Anthropic / provider-specific). #[serde(skip_serializing_if = "Option::is_none")] @@ -354,10 +245,6 @@ pub struct UsageInfo { pub prompt_cache_miss_tokens: Option, } -/// Deserialize `T`, treating a JSON `null` as `T::default()`. Combined with -/// `#[serde(default)]` (which covers an absent field), this lets the -/// [`ChatCompletionResponse`] `usage` field decode cleanly whether the provider -/// omits it or sends `null` — both yield zero counts instead of a 502. fn deserialize_default_on_null<'de, D, T>(deserializer: D) -> Result where D: serde::Deserializer<'de>, @@ -367,16 +254,8 @@ where Ok(Option::::deserialize(deserializer)?.unwrap_or_default()) } -// --------------------------------------------------------------------------- -// Choice -// --------------------------------------------------------------------------- - -/// A single completion alternative inside a response or chunk. -/// -/// The `message` and `delta` fields are mutually exclusive in practice: -/// non-streaming responses use `message`, streaming chunks use `delta`. -/// This is not enforced at the type level to keep the struct compatible with -/// both streaming and non-streaming JSON shapes. +/// A single completion alternative inside a response or chunk; `message` +/// for non-streaming, `delta` for streaming chunks. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct Choice { /// Position of this choice in the array. @@ -392,18 +271,9 @@ pub struct Choice { pub delta: Option, } -// --------------------------------------------------------------------------- -// Non-streaming response -// --------------------------------------------------------------------------- - -/// Response body for a non-streaming Chat Completions call. -/// -/// Note: `deny_unknown_fields` is intentionally omitted on upstream response -/// types. Providers may add fields (e.g. `service_tier`, `system_fingerprint`) -/// that the proxy does not model. Without `deny_unknown_fields`, these fields -/// are silently ignored during deserialization rather than causing a parse -/// failure that would surface as a 502 to the client. Request types retain -/// `deny_unknown_fields` for strict inbound validation. +/// Response body for a non-streaming Chat Completions call. Unknown fields +/// are ignored rather than failing decode into a 502 (providers add fields +/// like `service_tier`); request types retain strict validation instead. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionResponse { /// Unique completion identifier. @@ -416,36 +286,19 @@ pub struct ChatCompletionResponse { pub model: String, /// Ordered list of completion alternatives. pub choices: Vec, - /// Token usage for this request. - /// - /// Some OpenAI-compatible providers (e.g. certain OpenRouter backends) omit - /// `usage` entirely or send it as `null` on 2xx responses. The - /// `#[serde(default)]` + `deserialize_default_on_null` combination maps both - /// cases to zero token counts so the response decodes normally instead of - /// failing into a 502 "provider response decode error". + /// Token usage for this request. Some providers omit `usage` or send + /// `null` on 2xx; `#[serde(default)]` + `deserialize_default_on_null` + /// map both to zero token counts instead of a decode failure. #[serde(default, deserialize_with = "deserialize_default_on_null")] pub usage: UsageInfo, - /// Catch-all for provider-specific response fields (e.g. `service_tier`, - /// `system_fingerprint`) that the proxy does not model explicitly. - /// - /// Because this struct uses `#[serde(flatten)]` here, `deny_unknown_fields` - /// cannot also be applied. Unknown fields are captured rather than rejected. + /// Catch-all for provider-specific response fields; `#[serde(flatten)]` + /// precludes `deny_unknown_fields`, so unknown fields are captured. #[serde(flatten)] pub extra: serde_json::Map, } -// --------------------------------------------------------------------------- -// Streaming chunk -// --------------------------------------------------------------------------- - -/// A single Server-Sent Events chunk for a streaming completion. -/// -/// Note: `deny_unknown_fields` is intentionally omitted here because upstream -/// providers may add new fields to streaming chunks at any time. With -/// `deny_unknown_fields`, serde would reject unknown fields, causing the chunk -/// to be silently dropped. Without it, unknown fields are simply ignored, -/// avoiding silent data loss. Request types (e.g. `ChatCompletionRequest`) -/// retain `deny_unknown_fields` because the proxy controls their construction. +/// A single Server-Sent Events chunk for a streaming completion; unknown +/// fields are ignored because upstream providers add fields at any time. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionChunk { /// Unique completion identifier (stable across all chunks). @@ -463,10 +316,6 @@ pub struct ChatCompletionChunk { pub usage: Option, } -// --------------------------------------------------------------------------- -// Errors -// --------------------------------------------------------------------------- - /// Top-level error envelope returned by OpenAI-compatible APIs. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ErrorResponse { @@ -487,10 +336,6 @@ pub struct ErrorDetails { pub code: Option, } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -537,7 +382,6 @@ mod tests { #[test] fn chat_message_reasoning_alias() { - // Verify that the "reasoning" alias deserializes into reasoning_content. let json = r#"{"role":"assistant","reasoning":"thinking..."}"#; let msg: ChatMessage = serde_json::from_str(json).unwrap(); assert_eq!(msg.reasoning_content.as_deref(), Some("thinking...")); diff --git a/crates/llm-proxy-protocol/src/util.rs b/crates/llm-proxy-protocol/src/util.rs index c0029d7..0dfe77c 100644 --- a/crates/llm-proxy-protocol/src/util.rs +++ b/crates/llm-proxy-protocol/src/util.rs @@ -1,12 +1,8 @@ //! Shared utility functions for the protocol crate. -//! -//! Contains small helper functions used internally by the protocol adapters. - -/// Truncate `s` to at most `max_bytes` bytes on a UTF-8 char boundary. -/// -/// Returns `s` unchanged if it fits within `max_bytes`. Otherwise backs up -/// to the nearest char boundary at or before `max_bytes`, so the result may -/// be up to 3 bytes shorter than `max_bytes`. Never panics on multi-byte input. + +/// Truncate `s` to at most `max_bytes` bytes on a UTF-8 char boundary, +/// returning it unchanged if it fits; the result may be up to 3 bytes +/// shorter than `max_bytes`. Never panics on multi-byte input. pub fn truncate_str_safe(s: &str, max_bytes: usize) -> &str { if s.len() <= max_bytes { s @@ -19,25 +15,13 @@ pub fn truncate_str_safe(s: &str, max_bytes: usize) -> &str { } } -/// Truncate a string to `max_len` bytes, appending `suffix` if truncation occurs. -/// -/// The final string is at most `max_len` bytes (the suffix is included within -/// this budget). Respects UTF-8 char boundaries to prevent panics on multi-byte -/// characters. -/// -/// When a multi-byte character straddles the truncation boundary, the function -/// backs up to the previous character boundary, so the output may be shorter -/// than `max_len` by up to 3 bytes (the maximum UTF-8 character width). -/// -/// If `suffix.len() >= max_len`, the suffix alone would exceed the budget, so -/// the original string is returned untruncated (the caller's intent cannot be -/// satisfied safely). +/// Truncate a string to `max_len` bytes, appending `suffix` if truncation +/// occurs (suffix included in the budget, UTF-8 boundaries respected). If +/// `suffix.len() >= max_len`, the original string is returned untruncated. pub fn truncate_with_suffix(s: &str, max_len: usize, suffix: &str) -> String { if s.len() <= max_len { return s.to_owned(); } - // If the suffix alone would exceed or exactly fill the budget, we cannot - // produce a meaningful truncated result -- return the original string. if suffix.len() >= max_len { return s.to_owned(); } @@ -67,8 +51,6 @@ mod tests { #[test] fn truncate_str_safe_respects_utf8_boundaries() { - // "é" is 2 bytes in UTF-8: 6 chars = 12 bytes. Cutting at 7 bytes - // would split a char, so we back up to 6. let input = "éééééé"; let result = truncate_str_safe(input, 7); assert!(result.len() <= 7); @@ -85,8 +67,6 @@ mod tests { #[test] fn truncate_str_safe_multibyte_boundary_split() { - // "αβγδε" -- each char 2 bytes = 10 bytes total. max_bytes=3 splits - // β (bytes 2-3), so we back up to byte 2 -> "α". let input = "αβγδε"; assert_eq!(truncate_str_safe(input, 3), "α"); } @@ -103,8 +83,7 @@ mod tests { #[test] fn respects_utf8_boundaries() { - // "é" is 2 bytes in UTF-8 - let input = "éééééé"; // 12 bytes + let input = "éééééé"; let result = truncate_with_suffix(input, 7, "..."); assert!(result.len() <= 7); assert!(result.ends_with("...")); @@ -112,13 +91,11 @@ mod tests { #[test] fn suffix_equals_max_len_returns_original() { - // suffix.len() == max_len: cannot fit any content, return original assert_eq!(truncate_with_suffix("hello", 3, "..."), "hello"); } #[test] fn suffix_exceeds_max_len_returns_original() { - // suffix.len() > max_len: suffix alone exceeds budget, return original assert_eq!(truncate_with_suffix("hello", 2, "..."), "hello"); } @@ -139,21 +116,15 @@ mod tests { #[test] fn max_len_equals_suffix_len_plus_one() { - // "abcdef" is 6 bytes, max_len=4, suffix="..." (3 bytes) -> max_content=1 - // Only 1 byte of content fits before the suffix. assert_eq!(truncate_with_suffix("abcdef", 4, "..."), "a..."); } #[test] fn utf8_boundary_backup_produces_valid_string() { - // Input: "αβγδε" (each char 2 bytes = 10 bytes total) - // max_len=7, suffix="..." (3 bytes) -> max_content=4 - // byte 4 is mid-character in "αβ" (α=0-1, β=2-3, γ=4-5) so end backs to 4 let input = "αβγδε"; let result = truncate_with_suffix(input, 7, "..."); assert!(result.len() <= 7); assert!(result.ends_with("...")); - // Content before suffix should be valid UTF-8 let content = &result[..result.len() - 3]; assert!(content.is_char_boundary(content.len())); } diff --git a/crates/llm-proxy-protocol/src/zen.rs b/crates/llm-proxy-protocol/src/zen.rs index a01eeba..d498181 100644 --- a/crates/llm-proxy-protocol/src/zen.rs +++ b/crates/llm-proxy-protocol/src/zen.rs @@ -1,26 +1,11 @@ -//! Responses API and Gemini API request/response types. -//! -//! These types model the wire format for two LLM provider families: -//! -//! - **Responses API** – the OpenAI-style Responses endpoint used for -//! chat completions with optional tool use and streaming. -//! - **Gemini API** – the Google Gemini generate-content endpoint, -//! including streaming chunk support. -//! -//! Every struct derives `Debug`, `Clone`, `Serialize`, and `Deserialize` -//! so they can be serialized/deserialized directly via `serde_json`. +//! Responses API (OpenAI-style) and Gemini API (Google generate-content) +//! wire-format request/response types, including streaming chunk support. use serde::{Deserialize, Serialize}; -// --------------------------------------------------------------------------- -// Responses API types -// --------------------------------------------------------------------------- - -/// Top-level request body for the Responses API. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because this is an -/// outbound-only type -- the proxy constructs it internally via -/// `transform_to_responses()` and sends it to upstream providers. +/// Top-level request body for the Responses API; outbound-only (the proxy +/// constructs it via `transform_to_responses()`), so unknown-field +/// rejection is not applied. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponsesRequest { /// Model identifier (e.g. `"gpt-4o"`). @@ -68,11 +53,8 @@ pub struct ResponsesReasoning { pub effort: Option, } -/// Non-streaming response from the Responses API. -/// -/// Note: `deny_unknown_fields` is intentionally omitted on upstream response -/// types. Providers may add fields that the proxy does not model; unknown -/// fields are silently ignored rather than causing parse failures. +/// Non-streaming response from the Responses API; unknown provider fields +/// are silently ignored rather than failing decode. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponsesResponse { /// Unique response identifier. @@ -131,18 +113,9 @@ pub struct ResponsesUsage { pub output_tokens: i32, } -/// A single chunk in a streaming Responses API response. -/// -/// Note: `deny_unknown_fields` is intentionally omitted on streaming chunk -/// types. Upstream providers may add new event fields at any time; unknown -/// fields are silently ignored rather than causing chunk drops. The `delta` -/// field carries incremental text for `response.output_text.delta` events. -/// For `response.function_call_arguments.delta` events, the delta carries -/// incremental JSON arguments for tool calls. Parallel function calls have -/// their argument deltas interleaved and disambiguated by `output_index`; the -/// Responses API provider adapter tracks open blocks per `output_index` and -/// emits the full `CoreEvent` tool call lifecycle -/// (`ToolCallStart`/`ToolCallDelta`/`ToolCallStop`) for each. +/// A single chunk in a streaming Responses API response; `delta` carries +/// incremental text or tool-call JSON arguments, interleaved across parallel +/// calls and disambiguated by `output_index`. Unknown fields are ignored. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponsesChunk { /// Chunk type discriminator. @@ -158,30 +131,15 @@ pub struct ResponsesChunk { pub usage: Option, /// Error object present on `response.failed` events. pub error: Option, - /// Upstream output-item index carried on `response.output_item.added`, - /// `response.function_call_arguments.delta`, and - /// `response.function_call_arguments.done` events. - /// - /// The Responses API interleaves parallel function-call argument deltas and - /// disambiguates them via this index. Defaults to `None` so legacy chunks - /// that predate the field (and non-indexed events like `response.created`) - /// still parse. When absent on a function-call event the adapter falls back - /// to treating it as a single serial call. + /// Upstream output-item index carried on function-call events; parallel + /// function-call argument deltas are disambiguated via this index. + /// Defaults to `None` so legacy/non-indexed chunks still parse. #[serde(default)] pub output_index: Option, } -// --------------------------------------------------------------------------- -// Gemini API types -// --------------------------------------------------------------------------- - -/// Top-level request body for the Gemini generate-content endpoint. -/// -/// Note: `deny_unknown_fields` is intentionally omitted for outbound-only -/// request structs. The proxy constructs these internally and never -/// deserializes them from external input, so unknown field protection -/// provides no benefit and would be overly restrictive if the proxy adds -/// passthrough fields in the future. +/// Top-level request body for the Gemini generate-content endpoint; +/// outbound-only, so unknown-field rejection is not applied. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiRequest { /// Conversation contents to send to the model. @@ -190,12 +148,9 @@ pub struct GeminiRequest { pub generation_config: Option, /// Tools available to the model. pub tools: Vec, - /// When `Some(true)` the server streams chunked responses. - /// - /// Note: The Gemini API typically accepts stream as a query parameter - /// (`?alt=sse`) rather than a body field. This field is kept for internal - /// bookkeeping by the adapter, which should translate it to the appropriate - /// query parameter when constructing the HTTP request. + /// When `Some(true)` the server streams chunked responses; the Gemini API + /// usually takes streaming as the `?alt=sse` query parameter, which the + /// adapter constructs from this field. pub stream: Option, } @@ -204,21 +159,14 @@ pub struct GeminiRequest { pub struct GeminiContent { /// Speaker role (e.g. `"user"`, `"model"`). pub role: String, - /// Ordered content parts within this message. - /// - /// Defaults to an empty vec when omitted. The Gemini API may return a - /// candidate with no `parts` (e.g. empty content blocks), so this field is - /// deserialized with a default rather than treated as required. + /// Ordered content parts within this message; defaults to empty when + /// omitted (Gemini may return candidates with no `parts`). #[serde(default)] pub parts: Vec, } -/// A single content part inside a Gemini message. -/// -/// Note: `deny_unknown_fields` is intentionally omitted because Gemini may -/// return additional part types (e.g. `inline_data`, `file_data`) that the -/// proxy does not model yet. Unknown fields are silently ignored rather than -/// causing parse failures. +/// A single content part inside a Gemini message; unknown fields are +/// ignored since Gemini may return part types the proxy does not model. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] pub struct GeminiPart { @@ -313,11 +261,8 @@ pub struct GeminiFunctionDeclaration { pub parameters: Option, } -/// Non-streaming response from the Gemini generate-content endpoint. -/// -/// Note: `deny_unknown_fields` is intentionally omitted on upstream response -/// types. Providers may add fields that the proxy does not model; unknown -/// fields are silently ignored rather than causing parse failures. +/// Non-streaming response from the Gemini generate-content endpoint; +/// unknown provider fields are silently ignored rather than failing decode. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiResponse { /// Candidate completions. @@ -330,13 +275,9 @@ pub struct GeminiResponse { /// A single candidate in a Gemini response. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiCandidate { - /// The generated content. - /// - /// This is optional because real Gemini API responses routinely omit - /// `content` entirely for candidates blocked by `SAFETY`/`RECITATION` or - /// for prompt-feedback-only responses. A missing `content` is treated as - /// producing no text parts; the stop reason is derived from - /// `finish_reason` instead. + /// The generated content; optional because Gemini omits it for candidates + /// blocked by `SAFETY`/`RECITATION`. Missing `content` yields no text + /// parts; derive the stop reason from `finish_reason`. #[serde(default, skip_serializing_if = "Option::is_none")] pub content: Option, /// Reason the model stopped generating (e.g. `"STOP"`). @@ -345,9 +286,6 @@ pub struct GeminiCandidate { } /// Token usage metadata reported by the Gemini API. -/// -/// Uses manual `#[serde(rename)]` attributes for consistency with other Gemini -/// types in this module, rather than `rename_all = "camelCase"`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiUsage { /// Number of tokens in the prompt. @@ -361,11 +299,8 @@ pub struct GeminiUsage { pub total_token_count: i32, } -/// A single chunk in a streaming Gemini response. -/// -/// Note: `deny_unknown_fields` is intentionally omitted on streaming chunk -/// types. Upstream providers may add new fields at any time; unknown fields -/// are silently ignored rather than causing chunk drops. +/// A single chunk in a streaming Gemini response; unknown fields are +/// silently ignored rather than causing chunk drops. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiStreamChunk { /// Candidate completions in this chunk. @@ -375,10 +310,6 @@ pub struct GeminiStreamChunk { pub usage_metadata: Option, } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -488,7 +419,6 @@ mod tests { #[test] fn responses_chunk_output_index_defaulted_when_absent() { - // Legacy chunks that predate `output_index` must still parse into None. let json = r#"{"type":"response.created","id":"resp_1"}"#; let back: ResponsesChunk = serde_json::from_str(json).unwrap(); assert!(back.output_index.is_none()); diff --git a/crates/llm-proxy-protocol/tests/core_exports.rs b/crates/llm-proxy-protocol/tests/core_exports.rs index e361fa5..2432d27 100644 --- a/crates/llm-proxy-protocol/tests/core_exports.rs +++ b/crates/llm-proxy-protocol/tests/core_exports.rs @@ -1,9 +1,5 @@ -//! Integration smoke test: verifies that the core module is publicly -//! exported and the primary types can be imported, constructed, and -//! serialized. -//! -//! This prevents the Phase 1 gate from passing if `core.rs` or -//! `pub mod core;` is missing. +//! Smoke test: the core module is publicly exported and its primary types +//! construct and round-trip through serde. #[cfg(test)] mod tests { @@ -15,7 +11,6 @@ mod tests { Usage, UsageProvenance, }; - /// Trivial compile-time check: the types exist and are constructible. #[test] fn core_types_are_exported() { let _request: CoreRequest; @@ -31,7 +26,6 @@ mod tests { let _provenance: UsageProvenance; } - /// Stronger check: construct a minimal CoreRequest and verify serialization. #[test] fn core_request_serializes() { let req = CoreRequest { @@ -64,7 +58,6 @@ mod tests { assert_eq!(back.model.requested, "test-model"); } - /// Verify CacheControlType enum is exported and works. #[test] fn cache_control_type_is_exported() { let cc = CacheControl { @@ -76,7 +69,6 @@ mod tests { ); } - /// Verify that CoreEvent variants can be serialized and deserialized. #[test] fn core_event_round_trips() { let events = vec![ @@ -102,7 +94,6 @@ mod tests { } } - /// Verify that CoreRequest round-trips with full structural equality. #[test] fn core_request_full_round_trip() { let req = CoreRequest { @@ -137,11 +128,6 @@ mod tests { ); } - /// Verify that CoreResponse round-trips with full structural equality. - /// - /// Analogous to `core_request_full_round_trip` but exercises the response type - /// with diverse content variants (Text, ToolUse, Thinking) and all CoreResponse - /// fields including usage, stop_reason, stop_sequence, and provider_meta. #[test] fn core_response_full_round_trip() { let mut provider_meta = serde_json::Map::new(); diff --git a/crates/llm-proxy-protocol/tests/fixture_tests.rs b/crates/llm-proxy-protocol/tests/fixture_tests.rs index 4c41ae1..f5e6567 100644 --- a/crates/llm-proxy-protocol/tests/fixture_tests.rs +++ b/crates/llm-proxy-protocol/tests/fixture_tests.rs @@ -1,22 +1,6 @@ -//! Fixture-driven integration tests for client protocol adapters. -//! -//! Loads each fixture's input.json, decodes it to core, compares with core.json, -//! then encodes a response from core and compares with output.json. -//! -//! # Fixture format: client vs provider -//! -//! **Client-side fixtures** (this file) use serde's externally-tagged enum format -//! in `core-events.json` (e.g. `{"TextDelta": {"index": 0, "text": "Hello"}}`). -//! This matches how `CoreEvent` serializes with default serde derive. -//! -//! **Provider-side fixtures** (in `llm-proxy-provider/tests/fixture_tests.rs`) use -//! a human-readable format with a `"type"` field per event (e.g. -//! `{"type": "text_delta", "text": "Hello"}`). This format is compared against the -//! actual adapter output via `event_variant_name()`. -//! -//! The two formats serve different purposes: client fixtures test serde round-trips, -//! while provider fixtures test adapter decode output against human-readable expected -//! values. This is a deliberate design choice. +//! Fixture-driven integration tests for the client protocol adapters: decode +//! `input.json` to core and compare with `core.json`, then encode a +//! `CoreResponse` and compare key fields with `output.json`. #[cfg(test)] mod tests { @@ -31,13 +15,8 @@ mod tests { use std::fs; use std::path::Path; - // --------------------------------------------------------------------------- - // Helpers - // --------------------------------------------------------------------------- - const FIXTURE_ROOT: &str = "tests/fixtures"; - /// Read a fixture file as a parsed JSON value. fn read_fixture(dir: &Path, name: &str) -> serde_json::Value { let path = dir.join(name); let content = fs::read_to_string(&path) @@ -46,15 +25,12 @@ mod tests { .unwrap_or_else(|e| panic!("failed to parse {} as JSON: {e}", path.display())) } - /// Read a fixture file as raw string. fn read_fixture_raw(dir: &Path, name: &str) -> String { let path = dir.join(name); fs::read_to_string(&path) .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display())) } - /// Verify that decoding an input produces a core request that matches core.json. - /// Returns the decoded core request for further testing. fn assert_decode_matches_core( adapter: &str, input_json: &serde_json::Value, @@ -81,24 +57,20 @@ mod tests { ); } - /// Verify that a malformed input produces a ProtocolError. - /// - /// Both serde deserialization failures AND semantic validation errors - /// (e.g. empty model, empty messages) are treated as expected errors. fn assert_malformed_returns_error(adapter: &str, input_json: &serde_json::Value) { let result = match adapter { "anthropic" => { let req: AnthropicMessageRequest = match serde_json::from_value(input_json.clone()) { Ok(r) => r, - Err(_) => return, // serde parse failure is also a valid error + Err(_) => return, }; anthropic_adapter::decode_request(req) } "openai_chat" => { let req: ChatCompletionRequest = match serde_json::from_value(input_json.clone()) { Ok(r) => r, - Err(_) => return, // serde parse failure is also a valid error + Err(_) => return, }; openai_adapter::decode_request(req) } @@ -110,8 +82,6 @@ mod tests { ); } - /// Compare the `tool_calls` array of an encoded OpenAI message against the - /// expected message. A no-op when the expected message has no `tool_calls`. fn assert_tool_calls_match(out_m: &serde_json::Value, exp_m: &serde_json::Value) { if exp_m.get("tool_calls").is_none() { return; @@ -138,7 +108,6 @@ mod tests { } } - /// Verify that encoding a CoreResponse produces output that matches output.json. fn assert_encode_matches_output( adapter: &str, response: CoreResponse, @@ -150,7 +119,6 @@ mod tests { .expect("anthropic encode_response should succeed"); let serialized = serde_json::to_value(&out) .expect("serialized Anthropic response should be valid JSON"); - // Compare key fields (id, type, content, stop_reason, model) assert_eq!( serialized.get("type").and_then(|v| v.as_str()), output_json.get("type").and_then(|v| v.as_str()), @@ -190,7 +158,6 @@ mod tests { exp_choices.and_then(|c| c.get("finish_reason")), "finish_reason mismatch" ); - // Compare message content let out_msg = out_choices.and_then(|c| c.get("message")); let exp_msg = exp_choices.and_then(|c| c.get("message")); if let (Some(out_m), Some(exp_m)) = (out_msg, exp_msg) { @@ -199,10 +166,8 @@ mod tests { exp_m.get("content"), "message content mismatch" ); - // Compare tool_calls if present assert_tool_calls_match(out_m, exp_m); } - // Compare model if present in expected output if output_json.get("model").is_some() { assert_eq!( serialized.get("model"), @@ -210,7 +175,6 @@ mod tests { "model mismatch" ); } - // Compare usage if present in expected output if let Some(exp_usage) = output_json.get("usage") { let out_usage = serialized.get("usage"); assert!( @@ -235,17 +199,10 @@ mod tests { }; } - /// Run the decode-then-encode golden test for a non-stream fixture. fn run_non_stream_fixture(adapter: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter).join(case); let input = read_fixture(&dir, "input.json"); - // Special handling for malformed cases -- only input.json is needed. - // - // Note: this relies on the exact directory name "malformed". If multiple - // malformed variants are needed in the future (e.g. "malformed-empty-messages", - // "malformed-no-model"), either use a marker file like `malformed.flag` inside - // the fixture directory, or explicitly list all malformed variant names here. if case == "malformed" { assert_malformed_returns_error(adapter, &input); return; @@ -255,11 +212,6 @@ mod tests { assert_decode_matches_core(adapter, &input, &core); - // Encode direction: load the independently-authored core-response.json - // (normalized CoreResponse), encode it through the adapter, and compare - // key fields with output.json. This breaks the tautology where the - // CoreResponse was previously derived from the same output.json it was - // compared against. if dir.join("output.json").exists() { let output = read_fixture(&dir, "output.json"); let core_response_raw = read_fixture_raw(&dir, "core-response.json"); @@ -269,10 +221,6 @@ mod tests { } } - // --------------------------------------------------------------------------- - // Anthropic non-stream fixtures - // --------------------------------------------------------------------------- - #[test] fn anthropic_plain_text_request_fixture() { run_non_stream_fixture("anthropic", "plain-text-request"); @@ -323,10 +271,6 @@ mod tests { run_non_stream_fixture("anthropic", "malformed"); } - // --------------------------------------------------------------------------- - // OpenAI Chat non-stream fixtures - // --------------------------------------------------------------------------- - #[test] fn openai_plain_text_request_fixture() { run_non_stream_fixture("openai_chat", "plain-text-request"); @@ -377,10 +321,6 @@ mod tests { run_non_stream_fixture("openai_chat", "malformed"); } - // --------------------------------------------------------------------------- - // Fixture file completeness checks - // --------------------------------------------------------------------------- - #[test] fn all_anthropic_non_stream_fixtures_have_required_files() { let cases = [ @@ -455,11 +395,6 @@ mod tests { } } - /// Returns the list of all streaming fixture (adapter, case) pairs. - /// - /// Used by `all_streaming_fixtures_have_required_files`, - /// `streaming_core_events_json_is_valid`, `streaming_sse_fixtures_are_well_formed`, - /// and `streaming_encode_round_trip` to avoid duplicating the case list. fn all_streaming_cases() -> Vec<(&'static str, &'static str)> { vec![ ("anthropic", "streaming-text"), @@ -477,14 +412,6 @@ mod tests { ] } - /// Verify that all streaming fixtures have the required files. - /// - /// Client streaming fixtures require: - /// - `input.sse`: Optional reference file showing the corresponding provider wire - /// format. Not consumed by any test, but serves as documentation of the provider - /// stream that would produce these CoreEvents. - /// - `core-events.json`: The CoreEvent sequence to encode through the StreamEncoder. - /// - `output.sse`: The expected SSE output from the StreamEncoder. #[test] fn all_streaming_fixtures_have_required_files() { for (adapter, case) in &all_streaming_cases() { @@ -504,17 +431,6 @@ mod tests { } } - // --------------------------------------------------------------------------- - // Unified coverage-matrix test - // --------------------------------------------------------------------------- - - /// Non-stream fixture cases required for every client adapter. - /// Each entry is (case_name, required_files). - /// - /// Universal cases apply to all client adapters. Additional cases from the - /// plan (multiple-messages, refusal, redacted-thinking, multimedia content, - /// sampling-intent, model-mapping, stop-sequence) are listed as TODO items - /// for the next audit round. fn required_client_non_stream_cases() -> Vec<(&'static str, Vec<&'static str>)> { vec![ ( @@ -599,13 +515,9 @@ mod tests { ], ), ("malformed", vec!["input.json"]), - // TODO: Additional cases for the next audit round: - // multiple-messages, refusal, redacted-thinking, image/document/audio/video - // content, sampling-intent, model-mapping, stop-sequence. ] } - /// Streaming fixture cases required for every client adapter. fn required_client_streaming_cases() -> Vec<(&'static str, Vec<&'static str>)> { vec![ ( @@ -635,8 +547,6 @@ mod tests { ] } - /// Append `adapter/case/file` for every required fixture file that is - /// absent on disk. fn collect_missing_fixtures( adapter: &str, cases: &[(&str, Vec<&str>)], @@ -653,18 +563,13 @@ mod tests { } } - /// Unified coverage-matrix test that enforces required fixture cases for - /// all client adapters. Mirrors the provider-side coverage-matrix pattern. - /// `cargo test` must fail when a required fixture is missing. #[test] fn coverage_matrix_all_required_client_fixtures_exist() { let adapters = ["anthropic", "openai_chat"]; let mut missing: Vec = Vec::new(); for adapter in &adapters { - // Non-stream cases collect_missing_fixtures(adapter, &required_client_non_stream_cases(), &mut missing); - // Streaming cases collect_missing_fixtures(adapter, &required_client_streaming_cases(), &mut missing); } @@ -675,57 +580,9 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Streaming fixture validation tests - // --------------------------------------------------------------------------- - - // TODO(streaming-audit): The streaming fixture tests below only validate - // structure (core-events.json deserializes, SSE files contain valid SSE lines, - // encoder produces some valid output). They do NOT perform strict byte-level - // comparison of the encoded output.sse against the fixture's output.sse. This - // gap exists because: - // 1. The StreamEncoder may emit additional framing events (e.g. - // content_block_start/stop) not explicitly listed in the minimal - // core-events.json representation. - // 2. A strict comparison would require normalizing SSE whitespace and event - // ordering, or building a SSE parser that extracts structured events from - // both the actual output and output.sse for semantic comparison. - // 3. The streaming decode path (input.sse -> CoreEvents) is not yet - // implemented for client adapters (only provider adapters have - // StreamDecoder). - // Adding strict output.sse comparison is deferred to a future audit round. - - /// Tracked placeholder for the client-side streaming **decode** round-trip. - /// - /// The client adapters in `llm_proxy_protocol::client` currently expose only a - /// `StreamEncoder` (see `client::anthropic` and `client::openai_chat`); there - /// is no client `StreamDecoder` analogous to the provider crate's - /// `ProviderStreamDecoder`. So parsing of upstream `input.sse` frames into - /// `CoreEvent`s — the path a real proxy client would exercise on - /// partial/malformed/abruptly-truncated streams — is unimplemented and - /// therefore untested at the protocol layer. - /// - /// This test is kept as a tracked `#[ignore]` stub so the gap is visible in - /// `cargo test` output (run with `--ignored` to see it) and is not forgotten - /// when the decoder lands. When a client `StreamDecoder` is implemented, - /// un-ignore this test and mirror the provider crate's - /// `run_stream_decode_fixture` (see `llm-proxy-provider/tests/fixture_tests.rs`): - /// - /// 1. Parse each `(adapter, case)` `input.sse` into SSE frames. - /// 2. Feed the frames through the new client `StreamDecoder` and collect the - /// emitted `CoreEvent`s. - /// 3. Compare the decoded events' variant + fields against - /// `core-events.json` (semantic comparison, not byte-for-byte, to allow - /// for framing/normalization differences). - /// 4. Add a dedicated **malformed / abrupt-truncation** case: feed a frame - /// split mid-UTF-8 and a stream cut off before the terminal event, and - /// assert the decoder returns a structured error rather than panicking - /// or silently dropping the partial frame. - #[ignore = "client StreamDecoder (input.sse -> CoreEvents) not yet implemented; see doc comment"] + #[ignore = "client StreamDecoder (input.sse -> CoreEvents) not yet implemented"] #[test] fn streaming_decode_round_trip() { - // This stub intentionally fails if run, so it cannot silently rot into a - // false-pass: un-ignore only after implementing the decoder. panic!( "client streaming decode path is unimplemented: parse input.sse through \ a client StreamDecoder and compare CoreEvents against core-events.json \ @@ -734,12 +591,6 @@ mod tests { ); } - /// Verify that the core-events.json fixture can be deserialized into CoreEvent - /// values. This is a basic validation that streaming fixtures are well-formed. - /// Full encode/decode round-trip tests (parsing input.sse through a - /// StreamDecoder, encoding CoreEvents through a StreamEncoder, comparing with - /// output.sse) are deferred to a later audit cycle. The current tests - /// validate fixture structure only, not adapter behavior. #[test] fn streaming_core_events_json_is_valid() { for (adapter, case) in &all_streaming_cases() { @@ -755,15 +606,11 @@ mod tests { } } - /// Verify that input.sse and output.sse fixture files contain valid - /// SSE-formatted data when non-empty. Empty files are allowed for cases where - /// the protocol produces no output (e.g. OpenAI ping is a no-op). #[test] fn streaming_sse_fixtures_are_well_formed() { for (adapter, case) in &all_streaming_cases() { let dir = Path::new(FIXTURE_ROOT).join(adapter).join(case); - // input.sse should always have content (it represents client input). let input = read_fixture_raw(&dir, "input.sse"); assert!( !input.is_empty(), @@ -777,8 +624,6 @@ mod tests { "{adapter}/{case}/input.sse should contain SSE-formatted lines" ); - // output.sse may be empty or whitespace-only when the protocol produces - // no output for the given input (e.g. OpenAI Ping is a no-op). let output = read_fixture_raw(&dir, "output.sse"); let output_trimmed = output.trim(); if !output_trimmed.is_empty() { @@ -793,22 +638,6 @@ mod tests { } } - // --------------------------------------------------------------------------- - // Streaming encode round-trip tests - // --------------------------------------------------------------------------- - - /// Streaming round-trip test: parse core-events.json into CoreEvent values, - /// encode them through the client adapter's StreamEncoder, format the output - /// as SSE, and verify that encoding succeeds and produces valid SSE output. - /// - /// This test validates that the encode path produces structurally valid SSE - /// output for the given CoreEvent sequence. It does NOT compare exact byte - /// output with output.sse because the encoder may emit additional framing - /// events (e.g. content_block_start/stop) that are not explicit in the - /// minimal core-events.json representation. - /// - /// The test covers error and ping events in addition to text/tool/usage/thinking - /// events to ensure the encode path handles all CoreEvent variants. #[test] fn streaming_encode_round_trip() { for (adapter, case) in &all_streaming_cases() { @@ -820,7 +649,6 @@ mod tests { panic!("failed to parse {adapter}/{case}/core-events.json: {e}") }); - // Extract msg_id and model from MessageStart event (or use defaults). let default_id = if *adapter == "anthropic" { "msg_default".to_owned() } else { @@ -849,8 +677,6 @@ mod tests { } } - /// Encode `events` through the Anthropic `StreamEncoder` and verify that - /// every emitted event is structurally valid. fn encode_anthropic_stream( adapter: &str, case: &str, @@ -885,14 +711,6 @@ mod tests { ); } - /// Encode `events` through the OpenAI Chat `StreamEncoder` and verify that - /// every emitted chunk is structurally valid. - /// - /// NOTE: created is hard-coded to 1000 and include_usage is false. This means - /// the streaming encode test does not exercise the include_usage=true code - /// path (which emits a final usage chunk). - /// TODO: Derive include_usage from the fixture data (e.g. set to true for - /// streaming-usage cases) to ensure the usage-reporting path is tested. fn encode_openai_stream( adapter: &str, case: &str, @@ -904,13 +722,7 @@ mod tests { let mut total_data_lines = 0; for event in events { - // OpenAI encoder returns Err for CoreEvent::Error by design -- errors - // are handled at the transport level (HTTP status codes), not in stream - // encoding. We accept the Err gracefully and verify that non-error - // events still encode correctly. let Ok(chunks) = encoder.encode_event(event.clone()) else { - // Expected for CoreEvent::Error on OpenAI Chat. - // Verify this was indeed an error event. assert!( matches!(event, CoreEvent::Error { .. }), "{adapter}/{case}: encode_event failed for non-error event" @@ -931,8 +743,6 @@ mod tests { total_data_lines += 1; } - // Verify at least some chunks were produced, unless the fixture is an - // error-only case where OpenAI intentionally returns Err. let is_error_only = events .iter() .all(|e| matches!(e, CoreEvent::Error { .. } | CoreEvent::Ping)); @@ -944,7 +754,6 @@ mod tests { } } - /// Verify an Anthropic MessageEvent is valid JSON with a known event type. fn verify_anthropic_event_json(adapter: &str, case: &str, me: &MessageEvent) { let json = serde_json::to_string(me).unwrap_or_else(|e| { panic!("failed to serialize MessageEvent for {adapter}/{case}: {e}") @@ -970,7 +779,6 @@ mod tests { ); } - /// Verify an OpenAI ChatCompletionChunk is valid JSON. fn verify_openai_chunk_json(adapter: &str, case: &str, chunk: &ChatCompletionChunk) { let json = serde_json::to_string(chunk).unwrap_or_else(|e| { panic!("failed to serialize ChatCompletionChunk for {adapter}/{case}: {e}") diff --git a/crates/llm-proxy-provider/Cargo.toml b/crates/llm-proxy-provider/Cargo.toml index f9a8b6a..3a0fb0e 100644 --- a/crates/llm-proxy-provider/Cargo.toml +++ b/crates/llm-proxy-provider/Cargo.toml @@ -18,7 +18,9 @@ llm-proxy-core = { workspace = true } llm-proxy-protocol = { path = "../llm-proxy-protocol" } bytes = "1" form_urlencoded = "1" -reqwest = { version = "0.12", features = ["json", "stream", "rustls-tls"] } +# "socks" enables socks5/socks5h forward-proxy support for [provider.egress] +# (verified against docs.rs/reqwest 0.12 and empirically in transport.rs tests). +reqwest = { version = "0.12", features = ["json", "stream", "rustls-tls", "socks"] } futures = "0.3" tokio = { workspace = true } serde_json = { workspace = true } @@ -30,4 +32,5 @@ secrecy = { workspace = true } [dev-dependencies] axum = { workspace = true, features = ["macros", "json"] } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +# io-util: read_exact/write_all for the in-test SOCKS5 handshake server. +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "io-util"] } diff --git a/crates/llm-proxy-provider/src/adapter/anthropic.rs b/crates/llm-proxy-provider/src/adapter/anthropic.rs index 4e8d7c1..dbde901 100644 --- a/crates/llm-proxy-provider/src/adapter/anthropic.rs +++ b/crates/llm-proxy-provider/src/adapter/anthropic.rs @@ -1,12 +1,5 @@ -//! Anthropic Messages API provider adapter. -//! -//! Translates between core types and the Anthropic Messages wire format: -//! -//! ```text -//! CoreRequest -> MessageRequest -//! MessageResponse -> CoreResponse -//! MessageEvent stream -> CoreEvent stream -//! ``` +//! Anthropic Messages API provider adapter: translates core requests, +//! responses, and event streams to and from the Anthropic wire format. use std::sync::LazyLock; @@ -23,58 +16,17 @@ use super::{ use crate::error::ProviderError; use crate::sse::SseFrame; -// --------------------------------------------------------------------------- -// Tool name sanitization (collision-safe, reversible) -// --------------------------------------------------------------------------- - -/// Anthropic requires max_tokens but the core type makes it optional. -/// Use this default when the caller does not specify one. Chosen to be -/// large enough for most use cases while staying within typical model limits. const ANTHROPIC_DEFAULT_MAX_TOKENS: i32 = 4096; -/// Extract a mutable JSON object from a `serde_json::Value`. -/// -/// The `serde_json::json!()` macro always produces either an object, array, -/// or primitive value. In all call sites below the macro literal is an object -/// (`{ ... }`), so this unwrap is safe. The function exists so the safety -/// argument is documented in one place rather than repeated at each site. #[inline] fn json_object(val: &mut serde_json::Value) -> &mut serde_json::Map { val.as_object_mut() .expect("json! macro with {{...}} always produces a JSON object") } -// The sanitize/desanitize approach uses a sentinel prefix `__llmp_` to mark -// names that were actually rewritten. Only names carrying this sentinel are -// decoded during desanitization, which prevents false-positive decoding of -// tool names that naturally contain `_0xHH_` patterns (e.g. `parse_0xff_value`). -// -// NOTE: The `__llmp_` sentinel is NOT collision-safe against tool names that -// literally start with `__llmp_` (e.g. an original tool named -// `__llmp_foo_0x2e_bar`). Such a name would be falsely treated as encoded and -// decoded back incorrectly. This is accepted as a negligible-probability edge -// case: real tool names virtually never start with a double-underscore reserved -// prefix, and the only consequence is a mangled name on round-trip, not a -// security issue. - -// Anthropic `tool_use_id` must match `^[A-Za-z0-9_]{0,256}$`. -// NOTE: The regex pattern is a compile-time constant that is syntactically -// valid, so compilation cannot fail at runtime. static INVALID_TOOL_USE_ID_CHAR: LazyLock = LazyLock::new(|| regex::Regex::new(r"[^A-Za-z0-9_]").expect("valid regex")); -/// Sanitize a tool name for the Anthropic API using collision-safe, reversible encoding. -/// -/// Disallowed characters are encoded as `_0xHH_` where HH is the hex byte value. -/// This ensures two different tool names (e.g. `my.tool` and `my_tool`) produce -/// different sanitized names (`my_0x2e_tool` vs `my_tool`), avoiding collisions. -/// -/// When encoding is applied, a sentinel prefix `__llmp_` is prepended so that -/// [`desanitize_tool_name`] can distinguish genuinely encoded names from tool names -/// that naturally contain `_0xHH_` patterns (e.g. `parse_0xff_value`). -/// -/// Truncates to 128 chars if needed (avoiding partial `_0xHH_` sequences). -/// Returns the sanitized name and whether it was actually rewritten. fn sanitize_tool_name(name: &str) -> (String, bool) { use std::fmt::Write; let sentinel = "__llmp_"; @@ -86,44 +38,32 @@ fn sanitize_tool_name(name: &str) -> (String, bool) { if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' { result.push(ch); } else { - // Encode as _0xHH_ for each byte of the character using a - // stack-allocated buffer instead of allocating a String. let mut buf = [0u8; 4]; let encoded = ch.encode_utf8(&mut buf); for &byte in encoded.as_bytes() { - // NOTE: write! to a String cannot fail (the fmt::Write - // impl for String is infallible). write!(result, "_0x{:02x}_", byte).expect("write! to String is infallible"); } changed = true; } } - // If encoding was applied, prepend the sentinel. if changed { result.insert_str(0, sentinel); } - // Truncate to max_len if needed, avoiding partial _0xHH_ sequences. if result.len() > max_len { let mut end = max_len; while !result.is_char_boundary(end) && end > 0 { end -= 1; } - // Scan backward: if we cut in the middle of a `_0xHH_` sequence, - // truncate before the opening `_` instead. if end >= 5 { - // Check if we're inside a potential _0xHH_ pattern. if let Some(pos) = result[..end].rfind("_0x") { - let seq_end = pos + 6; // _0xHH_ is 6 bytes + let seq_end = pos + 6; if seq_end > end { - // The _0xHH_ sequence straddles the cut point. - // Truncate before the opening `_` instead. end = pos; } } } - // Final char-boundary check after adjustment. while !result.is_char_boundary(end) && end > 0 { end -= 1; } @@ -133,7 +73,6 @@ fn sanitize_tool_name(name: &str) -> (String, bool) { } } - // Guarantee non-empty. if result.is_empty() { return ("tool".to_owned(), name != "tool"); } @@ -141,29 +80,13 @@ fn sanitize_tool_name(name: &str) -> (String, bool) { (result, changed) } -/// Reverse a sanitized tool name back to the original by decoding `_0xHH_` patterns. -/// -/// Uses a sentinel prefix `__llmp_` to distinguish genuinely encoded names from -/// tool names that happen to contain `_0x` naturally (e.g. `parse_0xff_value`). -/// Only names that carry the `__llmp_` sentinel are decoded; names without it -/// are returned as-is, preventing false-positive decoding. -/// -/// # Collision edge case (accepted) -/// -/// The `__llmp_` sentinel is NOT collision-safe against tool names that -/// literally start with `__llmp_`. An original name such as -/// `__llmp_foo_0x2e_bar` would be misinterpreted as encoded and decoded back -/// incorrectly. This is accepted as a negligible-probability edge case; real -/// tool names virtually never start with this reserved double-underscore prefix. fn desanitize_tool_name(name: &str) -> std::borrow::Cow<'_, str> { const SENTINEL: &str = "__llmp_"; - // Only decode names that carry the sentinel prefix. if !name.starts_with(SENTINEL) { return std::borrow::Cow::Borrowed(name); } - // Strip the sentinel before decoding. let stripped = &name[SENTINEL.len()..]; if !stripped.contains("_0x") { return std::borrow::Cow::Owned(stripped.to_owned()); @@ -174,44 +97,34 @@ fn desanitize_tool_name(name: &str) -> std::borrow::Cow<'_, str> { let mut i = 0; while i < bytes.len() { - // Pattern: _0xHH_ (6 bytes at offsets i..=i+5). - // Need at least 6 bytes remaining: i + 5 must be a valid index. if bytes[i] == b'_' && i + 5 < bytes.len() && bytes[i + 1] == b'0' && bytes[i + 2] == b'x' { let hex_hi = i + 3; let hex_lo = i + 4; let closing = i + 5; if bytes[closing] == b'_' { - // Parse the two hex digits. if let Some(byte_val) = hex_byte(bytes[hex_hi], bytes[hex_lo]) { result.push(byte_val); - i = closing + 1; // skip past the closing '_' + i = closing + 1; continue; } } } - // Not a valid pattern, push the byte as-is. result.push(bytes[i]); i += 1; } - // Convert accumulated bytes back to a UTF-8 string. match String::from_utf8(result) { Ok(s) => std::borrow::Cow::Owned(s), - Err(e) => { - // Fallback: lossy conversion for safety. - std::borrow::Cow::Owned(String::from_utf8_lossy(e.as_bytes()).into_owned()) - } + Err(e) => std::borrow::Cow::Owned(String::from_utf8_lossy(e.as_bytes()).into_owned()), } } -/// Parse two ASCII hex digits into a byte value. fn hex_byte(hi: u8, lo: u8) -> Option { let hi_val = hex_digit(hi)?; let lo_val = hex_digit(lo)?; Some(hi_val << 4 | lo_val) } -/// Parse a single ASCII hex digit into its numeric value. const fn hex_digit(b: u8) -> Option { match b { b'0'..=b'9' => Some(b - b'0'), @@ -221,9 +134,6 @@ const fn hex_digit(b: u8) -> Option { } } -/// Sanitize a tool_use_id for the Anthropic API. -/// -/// Replaces disallowed characters with underscores, truncates to 256 chars. fn sanitize_tool_use_id(id: &str) -> String { let sanitized: String = INVALID_TOOL_USE_ID_CHAR.replace_all(id, "_").into_owned(); if sanitized.len() > 256 { @@ -239,10 +149,6 @@ fn sanitize_tool_use_id(id: &str) -> String { } } -// --------------------------------------------------------------------------- -// Adapter struct -// --------------------------------------------------------------------------- - /// Adapter for the Anthropic Messages API. #[derive(Debug, Clone, Default)] pub struct AnthropicAdapter; @@ -255,43 +161,17 @@ impl AnthropicAdapter { } } -// --------------------------------------------------------------------------- -// Stream decoder -// --------------------------------------------------------------------------- - -/// Stateful stream decoder for Anthropic SSE frames. -/// -/// ## CoreEvent variants emitted -/// -/// - `MessageStart` -- on `message_start` -/// - `ContentStart` -- on `content_block_start` (Text, Thinking, or ToolUse) -/// - `TextDelta` -- on `content_block_delta` with `text_delta` -/// - `ThinkingDelta` -- on `content_block_delta` with `thinking_delta` -/// - `ToolCallStart` -- on `content_block_start` with `tool_use` -/// - `ToolCallDelta` -- on `content_block_delta` with `input_json_delta` -/// - `ToolCallStop` -- on `content_block_stop` for tool_use blocks -/// - `UsageDelta` -- on `message_delta` with usage -/// - `MessageStop` -- on `message_delta` with stop_reason or `message_stop` -/// - `Ping` -- on `ping` -/// - `Error` -- on `error` -/// -/// Intentionally never emitted: none (all variants are covered). +/// Stateful stream decoder translating Anthropic SSE frames into `CoreEvent`s. #[derive(Debug)] pub struct AnthropicStreamDecoder { model_ref: ModelRef, started: bool, current_block_kind: ContentKind, tool_blocks: Vec, - /// Tracks which tool blocks have already received a `ToolCallStop` via - /// `content_block_stop`, so `finish()` does not emit duplicates. tool_blocks_closed: Vec, stop_sent: bool, } -/// Returns the tool-block `field` value, or a sentinel when upstream omitted it. -/// -/// Anthropic is expected to always send `id` and `name` on a `tool_use` block; -/// a missing value is logged and replaced so the stream stays decodable. fn tool_field_or_sentinel(value: Option<&str>, field: &str, sentinel: &str) -> String { match value { Some(v) => v.to_owned(), @@ -305,10 +185,6 @@ fn tool_field_or_sentinel(value: Option<&str>, field: &str, sentinel: &str) -> S } } -/// Translates a `content_block_delta` event into the matching core delta event. -/// -/// Empty payloads are dropped: they carry no information and would otherwise -/// surface as zero-length deltas downstream. fn push_content_block_delta(event: &MessageEvent, events: &mut Vec) { let idx = event.index.unwrap_or(0); let Some(ref delta) = event.delta else { @@ -349,7 +225,6 @@ fn push_content_block_delta(event: &MessageEvent, events: &mut Vec) { } } -/// Translates an Anthropic `error` event into a `CoreEvent::Error`. fn push_stream_error(event: &MessageEvent, events: &mut Vec) { let Some(ref err) = event.error else { return; @@ -367,7 +242,6 @@ fn push_stream_error(event: &MessageEvent, events: &mut Vec) { } impl AnthropicStreamDecoder { - /// Handles `content_block_start`, tracking the block kind and tool blocks. fn push_content_block_start(&mut self, event: &MessageEvent, events: &mut Vec) { let idx = event.index.unwrap_or(0); let Some(ref block) = event.content_block else { @@ -387,8 +261,6 @@ impl AnthropicStreamDecoder { self.tool_blocks_closed.push(false); let id = tool_field_or_sentinel(block.id.as_deref(), "id", ""); let name = tool_field_or_sentinel(block.name.as_deref(), "name", ""); - // Reverse-map sanitized tool names back to originals, - // same as the non-streaming decode path. let original_name = desanitize_tool_name(&name).into_owned(); events.push(CoreEvent::ToolCallStart { index: idx, @@ -396,7 +268,6 @@ impl AnthropicStreamDecoder { name: original_name, }); } - // "text", plus any unknown block type, emit a generic ContentStart. _ => { self.current_block_kind = ContentKind::Text; events.push(CoreEvent::ContentStart { @@ -407,12 +278,10 @@ impl AnthropicStreamDecoder { } } - /// Handles `content_block_stop`, closing an open tool block. fn push_content_block_stop(&mut self, event: &MessageEvent, events: &mut Vec) { let idx = event.index.unwrap_or(0); if self.current_block_kind == ContentKind::ToolUse { events.push(CoreEvent::ToolCallStop { index: idx }); - // Mark this tool block as closed so finish() won't re-emit. let closed_slot = self .tool_blocks .iter() @@ -421,10 +290,6 @@ impl AnthropicStreamDecoder { if let Some(closed) = closed_slot { *closed = true; } else { - // Index not found in tracked tool blocks -- the tool call - // was never opened via a content_block_start event, so - // there is nothing to mark as closed. Log the mismatch - // for observability. tracing::warn!( idx, "Anthropic: tool block index not found in content_block_stop; \ @@ -435,11 +300,7 @@ impl AnthropicStreamDecoder { self.current_block_kind = ContentKind::Text; } - /// Handles `message_delta`, emitting usage before any terminal stop event. fn push_message_delta(&mut self, event: &MessageEvent, events: &mut Vec) { - // Usage from message_delta -- emitted BEFORE MessageStop so that - // downstream consumers see UsageDelta before the terminal event, - // consistent with the plan's documented event ordering. if let Some(ref usage) = event.usage { events.push(CoreEvent::UsageDelta { usage: build_anthropic_usage( @@ -450,7 +311,6 @@ impl AnthropicStreamDecoder { ), }); } - // Stop reason from message_delta -- emitted AFTER UsageDelta. if let Some(ref delta) = event.delta && let Some(ref reason) = delta.stop_reason && !reason.is_empty() @@ -531,9 +391,7 @@ impl ProviderStreamDecoder for AnthropicStreamDecoder { "error" => { push_stream_error(&event, &mut events); } - _ => { - // Unknown event type; skip. - } + _ => {} } Ok(events) @@ -550,7 +408,6 @@ impl ProviderStreamDecoder for AnthropicStreamDecoder { }); } - // Close any tool blocks that were NOT already closed by content_block_stop. for (i, &idx) in self.tool_blocks.iter().enumerate() { let closed = self.tool_blocks_closed.get(i).copied().unwrap_or(false); if !closed { @@ -574,10 +431,6 @@ impl ProviderStreamDecoder for AnthropicStreamDecoder { } } -// --------------------------------------------------------------------------- -// AnthropicAdapter impl -// --------------------------------------------------------------------------- - impl AnthropicAdapter { /// Encode a core request into an Anthropic Messages request. pub fn encode_request( @@ -590,7 +443,6 @@ impl AnthropicAdapter { let tools = encode_tools(&core.tools); let tool_choice = core.tool_choice.as_ref().and_then(encode_tool_choice); - // Build the full request as JSON to avoid #[non_exhaustive] struct literal issues. let mut req = serde_json::json!({ "model": target.upstream_model, "max_tokens": core.sampling.max_tokens.unwrap_or(ANTHROPIC_DEFAULT_MAX_TOKENS), @@ -614,7 +466,6 @@ impl AnthropicAdapter { if let Some(top_p) = core.sampling.top_p { obj.insert("top_p".to_owned(), serde_json::json!(top_p)); } - // Forward stop sequences. if let Some(ref stop) = core.sampling.stop && !stop.is_empty() { @@ -633,8 +484,6 @@ impl AnthropicAdapter { if let Some(ref thinking) = core.sampling.thinking { obj.insert("thinking".to_owned(), thinking.clone()); } - // Forward reasoning_effort if present (Anthropic supports this in - // extended thinking mode). if let Some(ref effort) = core.sampling.reasoning_effort { obj.insert( "reasoning_effort".to_owned(), @@ -680,13 +529,10 @@ impl AnthropicAdapter { }); } "tool_use" => { - // Lazily build the empty object only on the `None` path so the - // common tool-calling happy path avoids a heap allocation (MEDIUM-4). let input = block .input .clone() .unwrap_or_else(|| serde_json::Value::Object(serde_json::Map::new())); - // Reverse-map sanitized tool names back to originals. let name = block.name.clone().unwrap_or_else(|| { tracing::warn!("Anthropic: tool_use block missing name in response decode"); "".to_owned() @@ -703,7 +549,6 @@ impl AnthropicAdapter { }); } _ => { - // Unknown block type; store as text for safety. let text = block.text.clone().unwrap_or_default(); if !text.is_empty() { content.push(CoreContent::Text { text, cache: None }); @@ -712,7 +557,6 @@ impl AnthropicAdapter { } } - // Guarantee at least one content block. if content.is_empty() { content.push(CoreContent::Text { text: String::new(), @@ -724,8 +568,6 @@ impl AnthropicAdapter { .stop_reason .as_deref() .map(map_anthropic_stop_reason) - // Provider omitted `stop_reason`; record that it was absent rather - // than unmapped so the diagnostic is unambiguous. .unwrap_or(StopReason::Unknown("absent".to_owned())); let usage = build_anthropic_usage( @@ -764,11 +606,6 @@ impl AnthropicAdapter { } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Map an Anthropic stop reason to a core StopReason. fn map_anthropic_stop_reason(reason: &str) -> StopReason { match reason { "end_turn" => StopReason::EndTurn, @@ -780,7 +617,6 @@ fn map_anthropic_stop_reason(reason: &str) -> StopReason { } } -/// Build usage from Anthropic-style usage info. fn build_anthropic_usage( input_tokens: i32, output_tokens: i32, @@ -797,7 +633,6 @@ fn build_anthropic_usage( } } -/// Encode content blocks for a user message. fn encode_content_blocks(content: &[CoreContent]) -> serde_json::Value { let blocks: Vec = content .iter() @@ -858,11 +693,6 @@ fn encode_content_blocks(content: &[CoreContent]) -> serde_json::Value { serde_json::Value::Array(blocks) } -/// Encode content blocks for an assistant message. -/// Builds the Anthropic `system` field from core system content. -/// -/// Returns `None` when there is nothing to send; non-`Text` blocks are dropped -/// with a warning because Anthropic's system field only accepts text blocks. fn encode_system(system: &[CoreContent]) -> Option { if system.is_empty() { return None; @@ -902,11 +732,6 @@ fn encode_system(system: &[CoreContent]) -> Option { } } -/// Builds the Anthropic `messages` array from core messages. -/// -/// Anthropic has no system role inside `messages` -- a system-role message -/// reaching here means the client adapter did not extract it, so it is dropped -/// with a warning. Tool results and unknown future roles are sent as `user`. fn encode_messages(core_messages: &[CoreMessage]) -> Vec { let mut messages = Vec::with_capacity(core_messages.len()); for msg in core_messages { @@ -926,9 +751,6 @@ fn encode_messages(core_messages: &[CoreMessage]) -> Vec { })); } CoreRole::System => { - // System messages should have been extracted into the - // `system` field by the client adapter. If one reaches - // here it means the client adapter did not separate it. tracing::warn!( role = "system", "Anthropic adapter received system-role message in messages; \ @@ -936,7 +758,6 @@ fn encode_messages(core_messages: &[CoreMessage]) -> Vec { ); } CoreRole::Tool => { - // Tool results come through user role in Anthropic. let content = encode_content_blocks(&msg.content); messages.push(serde_json::json!({ "role": "user", @@ -944,7 +765,6 @@ fn encode_messages(core_messages: &[CoreMessage]) -> Vec { })); } _ => { - // Handle future CoreRole variants. let content = encode_content_blocks(&msg.content); messages.push(serde_json::json!({ "role": "user", @@ -956,13 +776,10 @@ fn encode_messages(core_messages: &[CoreMessage]) -> Vec { messages } -/// Builds the Anthropic `tools` array, sanitizing tool names and coercing a -/// non-object `input_schema` to an empty object schema. fn encode_tools(tools: &[CoreTool]) -> Vec { tools .iter() .map(|t| { - // Coerce non-object input_schema to a default object schema. let schema = if t.input_schema.is_object() { t.input_schema.clone() } else { @@ -984,10 +801,6 @@ fn encode_tools(tools: &[CoreTool]) -> Vec { .collect() } -/// Encodes the core tool choice into Anthropic's `tool_choice` field. -/// -/// Unknown variants return `None` (the field is omitted) rather than JSON -/// null, which the Anthropic API would reject. fn encode_tool_choice(tc: &CoreToolChoice) -> Option { match tc { CoreToolChoice::Auto => Some(serde_json::json!({"type": "auto"})), @@ -1057,10 +870,6 @@ fn encode_assistant_content(content: &[CoreContent]) -> serde_json::Value { serde_json::Value::Array(blocks) } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -1097,8 +906,6 @@ mod tests { } } - // -- Tool name sanitization tests ---------------------------------------- - #[test] fn sanitize_tool_name_valid() { let (name, changed) = sanitize_tool_name("get_weather"); @@ -1111,7 +918,6 @@ mod tests { #[test] fn sanitize_tool_name_replaces_dots_collision_safe() { - // "my.tool.name" should NOT collide with "my_tool_name". let (name1, changed1) = sanitize_tool_name("my.tool.name"); let (name2, changed2) = sanitize_tool_name("my_tool_name"); assert!(changed1); @@ -1148,12 +954,9 @@ mod tests { fn sanitize_tool_name_special_chars() { let (name, changed) = sanitize_tool_name("get weather!@#"); assert!(changed); - // Spaces and special chars should be encoded, not simply replaced with _. assert!(name.contains("0x")); } - // -- Tool use ID sanitization tests ---------------------------------------- - #[test] fn sanitize_tool_use_id_valid() { assert_eq!(sanitize_tool_use_id("toolu_123"), "toolu_123"); @@ -1169,8 +972,6 @@ mod tests { assert_eq!(sanitize_tool_use_id(""), "tool_result"); } - // -- Encode tests -------------------------------------------------------- - #[test] fn encode_text_request() { let core = make_core_request(vec![CoreMessage { @@ -1232,7 +1033,6 @@ mod tests { let body: serde_json::Value = serde_json::from_slice(&proxy_req.body).unwrap(); let tools = body["tools"].as_array().unwrap(); assert_eq!(tools.len(), 1); - // Dots are encoded collision-safe as _0x2e_ assert!(tools[0]["name"].as_str().unwrap().contains("0x2e")); } @@ -1312,8 +1112,6 @@ mod tests { assert_eq!(body["thinking"]["budget_tokens"], 5000); } - // -- Decode tests -------------------------------------------------------- - #[test] fn decode_text_response() { let target = make_target(); @@ -1443,8 +1241,6 @@ mod tests { ); } - // -- Streaming tests ----------------------------------------------------- - #[test] fn stream_text_decoding() { let target = make_target(); @@ -1640,8 +1436,6 @@ mod tests { ))); } - // -- Stop reason mapping tests ------------------------------------------- - #[test] fn anthropic_stop_reason_mappings() { assert_eq!(map_anthropic_stop_reason("end_turn"), StopReason::EndTurn); @@ -1655,15 +1449,12 @@ mod tests { StopReason::StopSequence ); assert_eq!(map_anthropic_stop_reason("refusal"), StopReason::Refusal); - // Unmapped values are preserved in Unknown(String) (GAP-LOW-9). assert_eq!( map_anthropic_stop_reason("unknown"), StopReason::Unknown("unknown".to_owned()) ); } - // -- Missing tests: encode stop sequences, input_schema coercion ---------- - #[test] fn encode_stop_sequences() { let mut core = make_core_request(vec![CoreMessage { @@ -1728,8 +1519,6 @@ mod tests { assert_eq!(schema["type"], "object"); } - // -- Missing tests: decode redacted thinking, stop_sequence, malformed frame -- - #[test] fn decode_redacted_thinking_response() { let target = make_target(); @@ -1795,15 +1584,12 @@ mod tests { assert!(events.is_empty()); } - // -- Missing test: usage-only chunk classification ------------------------ - #[test] fn stream_usage_only_chunk() { let target = make_target(); let adapter = AnthropicAdapter::new(); let mut decoder = adapter.new_stream_decoder(&target); - // Send a message_delta with usage but no text content. let frames = vec![ make_frame( r#"{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}"#, @@ -1818,7 +1604,6 @@ mod tests { all_events.extend(decoder.decode_frame(frame).unwrap()); } - // Usage should NOT be dropped just because there was no text. assert!( all_events .iter() @@ -1831,8 +1616,6 @@ mod tests { ); } - // -- Missing test: full lifecycle event ordering -------------------------- - #[test] fn stream_full_lifecycle_ordering() { let target = make_target(); @@ -1860,7 +1643,6 @@ mod tests { all_events.extend(decoder.decode_frame(frame).unwrap()); } - // Verify ordering: MessageStart < ContentStart < TextDelta < MessageStop. let msg_start_idx = all_events .iter() .position(|e| matches!(e, CoreEvent::MessageStart { .. })) @@ -1892,8 +1674,6 @@ mod tests { ); } - // -- Additional missing tests ---------------------------------------------- - #[test] fn encode_tool_result_as_tool_result_block() { let core = make_core_request(vec![CoreMessage { @@ -2102,7 +1882,6 @@ mod tests { let frame = make_frame(r#"{"type":"some_new_event","data":"whatever"}"#); let events = decoder.decode_frame(&frame).unwrap(); - // Unknown event types produce no events (they are skipped). assert!(events.is_empty()); } @@ -2134,13 +1913,9 @@ mod tests { } } - // -- Additional tests: sentinel, false-positive, upstream model, malformed, etc -- - #[test] fn desanitize_sentinel_prevents_false_positive() { - // A tool name that naturally contains _0xHH_ should NOT be decoded. let natural_name = "parse_0xff_value"; - // This name was never sanitized, so it should be returned as-is. assert_eq!(desanitize_tool_name(natural_name), natural_name); } @@ -2209,7 +1984,6 @@ mod tests { let target = make_target(); let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: serde_json::Value = serde_json::from_slice(&proxy_req.body).unwrap(); - // Should produce a valid request with empty messages array. assert_eq!(body["messages"].as_array().unwrap().len(), 0); } @@ -2259,7 +2033,6 @@ mod tests { let adapter = AnthropicAdapter::new(); let mut decoder = adapter.new_stream_decoder(&target); - // A chunk that only has message_stop with no prior content. let frames = vec![ make_frame( r#"{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}"#, @@ -2287,8 +2060,6 @@ mod tests { #[test] fn stream_no_duplicate_tool_call_stop() { - // Verify that finish() does NOT emit duplicate ToolCallStop for tool - // blocks that already received ToolCallStop during decode_frame(). let target = make_target(); let adapter = AnthropicAdapter::new(); let mut decoder = adapter.new_stream_decoder(&target); @@ -2312,7 +2083,6 @@ mod tests { } all_events.extend(decoder.finish().unwrap()); - // Count ToolCallStop events -- should be exactly 1. let tool_call_stop_count = all_events .iter() .filter(|e| matches!(e, CoreEvent::ToolCallStop { .. })) @@ -2323,8 +2093,6 @@ mod tests { ); } - // -- Source guard --------------------------------------------------------- - #[test] fn adapter_source_no_forbidden_imports() { let source = include_str!("anthropic.rs"); @@ -2342,8 +2110,6 @@ mod tests { ); } - // -- Helpers -------------------------------------------------------------- - fn make_frame(data: &str) -> SseFrame { SseFrame { event: None, diff --git a/crates/llm-proxy-provider/src/adapter/gemini.rs b/crates/llm-proxy-provider/src/adapter/gemini.rs index b2a0382..b13198a 100644 --- a/crates/llm-proxy-provider/src/adapter/gemini.rs +++ b/crates/llm-proxy-provider/src/adapter/gemini.rs @@ -1,16 +1,5 @@ -//! Google Gemini GenerateContent API provider adapter. -//! -//! Translates between core types and the Gemini wire format: -//! -//! ```text -//! CoreRequest -> GeminiRequest -//! GeminiResponse -> CoreResponse -//! GeminiStreamChunk stream -> CoreEvent stream -//! ``` -//! -//! The stream decoder emits the following [`CoreEvent`] variants: -//! `MessageStart`, `ContentStart`, `TextDelta`, `ContentStop` (implicit), -//! `ToolCallStart`, `ToolCallDelta`, `ToolCallStop`, `UsageDelta`, `MessageStop`. +//! Google Gemini GenerateContent API provider adapter: translates core +//! requests, responses, and event streams to and from the Gemini wire format. use std::collections::HashMap; @@ -30,45 +19,19 @@ use super::{ use crate::error::ProviderError; use crate::sse::SseFrame; -// --------------------------------------------------------------------------- -// Adapter struct -// --------------------------------------------------------------------------- - /// Adapter for the Google Gemini GenerateContent API. #[derive(Debug, Clone, Default)] pub struct GeminiAdapter; -// --------------------------------------------------------------------------- -// Stream decoder -// --------------------------------------------------------------------------- - -/// Stateful stream decoder for Gemini SSE frames. -/// -/// ## CoreEvent variants emitted -/// -/// - `MessageStart` -- on the first chunk received -/// - `ContentStart` -- on the first text delta -/// - `TextDelta` -- on text parts -/// - `ToolCallStart` -- on function_call parts -/// - `ToolCallDelta` -- on function_call arguments -/// - `ToolCallStop` -- immediately after each function_call part -/// - `UsageDelta` -- on usage-only chunks or final chunks with usage -/// - `MessageStop` -- on finish_reason or stream end -/// -/// Intentionally never emitted: `Ping` (Gemini has no heartbeat), `ThinkingDelta` -/// (Gemini thinking is not streamed as deltas in the current API), `Error` -/// (Gemini errors are handled at the transport level, not in stream decoding). +/// Stateful stream decoder translating Gemini SSE frames into `CoreEvent`s. #[derive(Debug)] pub struct GeminiStreamDecoder { model_ref: ModelRef, started: bool, content_index: usize, content_started: bool, - /// Maps Gemini candidate index to tool content block index. tool_blocks: Vec, - /// Whether each tool block has already received a ToolCallStop during decode_frame. tool_blocks_closed: Vec, - /// Monotonically increasing counter for generating unique tool use IDs. tool_id_counter: usize, stop_sent: bool, } @@ -99,7 +62,6 @@ impl ProviderStreamDecoder for GeminiStreamDecoder { }); } - // Usage-only chunk (no candidates). if chunk.candidates.is_empty() { if let Some(ref usage) = chunk.usage_metadata { events.push(CoreEvent::UsageDelta { @@ -111,10 +73,6 @@ impl ProviderStreamDecoder for GeminiStreamDecoder { let candidate = &chunk.candidates[0]; - // Process parts. A candidate may omit `content` entirely (e.g. for - // SAFETY/RECITATION-blocked responses or prompt-feedback-only - // candidates); treat a missing/empty content as producing no parts and - // rely on `finish_reason`/`usage_metadata` below for the stop/usage. let parts: &[GeminiPart] = candidate .content .as_ref() @@ -127,13 +85,11 @@ impl ProviderStreamDecoder for GeminiStreamDecoder { self.push_text_delta(text, &mut events); } - // Handle function call parts. if let Some(ref function_call) = part.function_call { self.push_function_call(function_call, &mut events); } } - // Handle finish reason. if let Some(ref reason) = candidate.finish_reason { if !reason.is_empty() && !self.stop_sent { self.close_content_if_open(); @@ -170,7 +126,6 @@ impl ProviderStreamDecoder for GeminiStreamDecoder { self.close_content_if_open(); - // Emit ToolCallStop for any tool blocks that were not closed during decode_frame(). for (&idx, &closed) in self.tool_blocks.iter().zip(&self.tool_blocks_closed) { if !closed { events.push(CoreEvent::ToolCallStop { index: idx }); @@ -194,7 +149,6 @@ impl ProviderStreamDecoder for GeminiStreamDecoder { } impl GeminiStreamDecoder { - /// Emits a text delta, opening a text content block first if none is open. fn push_text_delta(&mut self, text: &str, events: &mut Vec) { if !self.content_started { self.content_started = true; @@ -209,10 +163,6 @@ impl GeminiStreamDecoder { }); } - /// Emits the start/delta/stop triple for one Gemini `functionCall` part. - /// - /// Gemini delivers function calls whole rather than as argument deltas, so - /// the block is opened and closed within this call. fn push_function_call( &mut self, function_call: &GeminiFunctionCall, @@ -248,7 +198,6 @@ impl GeminiStreamDecoder { } events.push(CoreEvent::ToolCallStop { index: block_idx }); - // Mark this tool block as closed so finish() does not emit a duplicate. if let Some(last) = self.tool_blocks_closed.last_mut() { *last = true; } @@ -263,15 +212,6 @@ impl GeminiStreamDecoder { } } -// --------------------------------------------------------------------------- -// GeminiAdapter impl -// --------------------------------------------------------------------------- - -/// Encodes one core content block into zero or more Gemini parts. -/// -/// Content Gemini cannot represent is dropped with a warning; an unresolvable -/// `tool_use_id` is an error because a guessed function name is rejected -/// upstream with an opaque 400. fn encode_content_part( c: &CoreContent, tool_name_map: &HashMap, @@ -302,17 +242,7 @@ fn encode_content_part( } let response_val = encode_tool_result_payload(tool_use_id, result_content); - // The name field must match the function name from the original - // call. Gemini uses it to correlate the response with the function - // declaration, and the map was built from the ToolUse blocks in the - // conversation history. let Some(fn_name) = tool_name_map.get(tool_use_id) else { - // The lookup missed because the ToolUse block is absent from the - // request being encoded (truncated/pruned history) or the id is - // foreign (e.g. an Anthropic `toolu_...` id round-tripping - // through Gemini). Guessing a name only ever produced an invalid - // one that Gemini rejects with a confusing 400, so return a - // clear error instead. return Err(ProviderError::InvalidConfig(format!( "Gemini: cannot resolve function name for tool_use_id \ \"{tool_use_id}\"; the corresponding ToolUse block was \ @@ -331,10 +261,6 @@ fn encode_content_part( Ok(()) } -/// Flattens tool-result content into Gemini's string-only `functionResponse`. -/// -/// Non-text blocks are dropped with a warning: the wire format carries a single -/// string `result`. fn encode_tool_result_payload( tool_use_id: &str, result_content: &[CoreContent], @@ -360,6 +286,12 @@ fn encode_tool_result_payload( impl GeminiAdapter { /// Encode a core request into a Gemini GenerateContent request. + #[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the Gemini GenerateContent request \ + shape, kept in one place so the mapping can be audited against \ + the API schema without jumping between helpers" + )] pub fn encode_request( &self, core: &CoreRequest, @@ -367,26 +299,14 @@ impl GeminiAdapter { ) -> Result { let mut contents = Vec::new(); - // Pre-build a tool-id-to-name map for O(1) lookups during tool result - // encoding, instead of scanning all messages for each tool result. let tool_name_map = build_tool_name_map(core); - // System instructions are handled via generation config or prepended. - // Gemini does not have a dedicated system field in the basic request. - for msg in &core.messages { let role = match msg.role { CoreRole::User => "user", CoreRole::Assistant => "model", - // Gemini has no system role; user is the standard mapping per - // Gemini's documentation for injecting system instructions. CoreRole::System => "user", - // Gemini's API requires tool results to be in "user" role - // messages with functionResponse parts. This is the canonical - // mapping per Gemini's documentation. CoreRole::Tool => "user", - // Future CoreRole variants are mapped to "user" as a safe default. - // Update this match when new variants are added to CoreRole. _ => "user", }; @@ -403,20 +323,7 @@ impl GeminiAdapter { } } - // Prepend system prompt as user message if present. - // NOTE: A synthetic "Understood." model acknowledgment is inserted - // immediately after the system prompt. Gemini requires the conversation - // to start with a user turn followed by a model turn, so this dummy - // acknowledgment satisfies that constraint. It does not affect model - // behavior. - // - // DESIGN DECISION: If both `core.system` and `CoreRole::System` messages - // are present, both are included. The prepended system prompt comes first, - // followed by CoreRole::System messages mapped to "user" role. This may - // create duplicate system content, but deduplication is deferred to a - // the caller typically provides one or the other. if !core.system.is_empty() { - // Warn about non-text system content blocks (consistent with other adapters). for c in &core.system { if !matches!(c, CoreContent::Text { .. }) { tracing::warn!( @@ -447,15 +354,6 @@ impl GeminiAdapter { parts: vec![GeminiPart::text(system_text)], }, ); - // Synthetic model acknowledgment (see NOTE above). - // - // Only insert when the turn immediately following the system - // prompt is NOT already a "model" turn. Gemini requires the - // conversation to start with a user turn followed by a model - // turn; if the first real message is already a model turn - // (e.g. an Assistant-first conversation), inserting a synthetic - // one would produce two consecutive model turns, which Gemini - // rejects. let next_is_model = contents.get(1).is_some_and(|c| c.role == "model"); if !next_is_model { contents.insert( @@ -503,11 +401,6 @@ impl GeminiAdapter { .and_then(|s| if s.is_empty() { None } else { Some(s.clone()) }), }; - // Forward tool_choice if present. - // Note: Gemini supports `tool_config.function_calling_config` for - // controlling tool choice, but the wire types don't model it yet. - // Per the plan's lossy translation rules, emit a warning since the - // omission is safe but potentially impactful. if core.tool_choice.is_some() { tracing::warn!( tool_choice = ?core.tool_choice, @@ -516,7 +409,6 @@ impl GeminiAdapter { ); } - // Warn about fields that Gemini cannot forward. if core.sampling.reasoning_effort.is_some() { tracing::warn!( "Gemini: reasoning_effort specified but not forwarded; \ @@ -561,10 +453,6 @@ impl GeminiAdapter { let mut content = Vec::new(); let mut tool_id_counter = 0usize; - // A candidate may omit `content` entirely (e.g. for SAFETY/RECITATION- - // blocked responses or prompt-feedback-only candidates). Treat a - // missing/empty content as producing no parts; `finish_reason` still - // drives the stop reason and `content` is back-filled below if empty. let parts: &[GeminiPart] = candidate .content .as_ref() @@ -580,8 +468,6 @@ impl GeminiAdapter { } } if let Some(ref function_call) = part.function_call { - // Lazily build the empty object only on the `None` path so the - // common tool-calling happy path avoids a heap allocation (MEDIUM-4). let input = function_call .args .clone() @@ -595,7 +481,6 @@ impl GeminiAdapter { } } - // Guarantee at least one content block. if content.is_empty() { content.push(CoreContent::Text { text: String::new(), @@ -613,8 +498,6 @@ impl GeminiAdapter { .unwrap_or(if has_tool_use { StopReason::ToolUse } else { - // Provider omitted `finish_reason`; record that it was absent - // rather than unmapped so the diagnostic is unambiguous. StopReason::Unknown("absent".to_owned()) }); @@ -651,14 +534,6 @@ impl GeminiAdapter { } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Look up the original function name for a tool_use_id by searching ToolUse -/// content blocks in the conversation messages. This is needed because Gemini -/// requires the `function_response.name` to match the original function -/// declaration, but the core `ToolResult` type only carries `tool_use_id`. fn build_tool_name_map(core: &CoreRequest) -> HashMap { let mut map = HashMap::new(); for msg in &core.messages { @@ -671,7 +546,6 @@ fn build_tool_name_map(core: &CoreRequest) -> HashMap { map } -/// Build usage from Gemini usage metadata. fn build_gemini_usage(usage: &GeminiUsage) -> Usage { Usage { input_tokens: usage.prompt_token_count, @@ -683,10 +557,6 @@ fn build_gemini_usage(usage: &GeminiUsage) -> Usage { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -725,8 +595,6 @@ mod tests { } } - // -- Encode tests -------------------------------------------------------- - #[test] fn encode_text_request() { let core = make_core_request(vec![CoreMessage { @@ -764,7 +632,6 @@ mod tests { let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: GeminiRequest = serde_json::from_slice(&proxy_req.body).unwrap(); - // System prompt should be prepended as user + model acknowledgment. assert_eq!(body.contents[0].role, "user"); assert_eq!( body.contents[0].parts[0].text, @@ -775,10 +642,6 @@ mod tests { #[test] fn encode_system_prompt_no_double_model_when_assistant_first() { - // Regression for gemini.rs:416 -- when a system prompt is present AND - // the first real message is already a model/assistant turn, the - // synthetic "Understood." model turn must NOT be inserted (it would - // create two consecutive model turns, which Gemini rejects). let mut core = make_core_request(vec![CoreMessage { role: CoreRole::Assistant, content: vec![CoreContent::Text { @@ -795,12 +658,8 @@ mod tests { let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: GeminiRequest = serde_json::from_slice(&proxy_req.body).unwrap(); - // contents[0] = system prompt as user. assert_eq!(body.contents[0].role, "user"); - // contents[1] must be the real Assistant message (model), NOT the - // synthetic "Understood." turn. assert_eq!(body.contents[1].role, "model"); - // No synthetic "Understood." text anywhere. let has_understood = body.contents.iter().any(|c| { c.parts .iter() @@ -810,7 +669,6 @@ mod tests { !has_understood, "synthetic Understood. turn must be suppressed when the next turn is already model" ); - // And no two consecutive model turns. for pair in body.contents.windows(2) { assert!( !(pair[0].role == "model" && pair[1].role == "model"), @@ -887,8 +745,6 @@ mod tests { ); } - // -- Decode tests -------------------------------------------------------- - #[test] fn decode_text_response() { let target = make_target(); @@ -990,8 +846,6 @@ mod tests { assert_eq!(core_resp.model.requested, "my-gemini"); } - // -- Streaming tests ----------------------------------------------------- - #[test] fn stream_text_decoding() { let target = make_target(); @@ -1063,8 +917,6 @@ mod tests { ); } - // -- Streaming tests: tool call, malformed frame, usage ------------------- - #[test] fn stream_tool_call_decoding() { let target = make_target(); @@ -1111,7 +963,6 @@ mod tests { let frame = make_frame("not valid json {{{"); let events = decoder.decode_frame(&frame).unwrap(); - // Malformed JSON is silently skipped; no events emitted at all. assert!( events.is_empty(), "malformed frame should produce no events" @@ -1129,7 +980,6 @@ mod tests { ); let events = decoder.decode_frame(&frame).unwrap(); - // Should emit MessageStart + UsageDelta. assert!( events .iter() @@ -1143,14 +993,8 @@ mod tests { ); } - // -- Encode tests: tool_result, tool_use, top_p, input_schema ------------ - #[test] fn encode_tool_result_as_function_response() { - // The function name for a functionResponse is resolved from the - // matching ToolUse block in the conversation history (audit - // gemini-tool-result-name-heuristic). A bare ToolResult without the - // preceding ToolUse is now a hard error, so include the call. let core = make_core_request(vec![ CoreMessage { role: CoreRole::Assistant, @@ -1177,15 +1021,11 @@ mod tests { let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: GeminiRequest = serde_json::from_slice(&proxy_req.body).unwrap(); - // The assistant ToolUse and the tool ToolResult produce their own - // contents entries (one function_call, one function_response). Find the - // function_response across the contents. let fr = body .contents .iter() .find_map(|c| c.parts.iter().find_map(|p| p.function_response.clone())) .expect("ToolResult must be encoded as a functionResponse"); - // Name is the real function name from the ToolUse, not a guessed id. assert_eq!(fr.name, "get_weather"); } @@ -1259,8 +1099,6 @@ mod tests { assert_eq!(params["type"], "object"); } - // -- Usage mapping test --------------------------------------------------- - #[test] fn usage_mapping() { let usage = GeminiUsage { @@ -1274,8 +1112,6 @@ mod tests { assert_eq!(core_usage.provenance, UsageProvenance::ProviderReported); } - // -- Full lifecycle ordering test ----------------------------------------- - #[test] fn stream_full_lifecycle_ordering() { let target = make_target(); @@ -1296,7 +1132,6 @@ mod tests { all_events.extend(decoder.decode_frame(frame).unwrap()); } - // Verify lifecycle ordering: MessageStart before everything else. let start_idx = all_events .iter() .position(|e| matches!(e, CoreEvent::MessageStart { .. })) @@ -1328,8 +1163,6 @@ mod tests { ); } - // -- Additional missing tests ---------------------------------------------- - #[test] fn encode_stop_sequences_forwarded() { let mut core = make_core_request(vec![CoreMessage { @@ -1380,8 +1213,6 @@ mod tests { let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: GeminiRequest = serde_json::from_slice(&proxy_req.body).unwrap(); - // The ToolResult should be encoded as a function_response with the - // correct function name (looked up from the ToolUse block). let fr = body .contents .iter() @@ -1398,7 +1229,6 @@ mod tests { let frame = make_frame(r#"{"candidates":[],"usageMetadata":null}"#); let events = decoder.decode_frame(&frame).unwrap(); - // Only MessageStart should be emitted (first chunk), no content events. assert!( events .iter() @@ -1472,14 +1302,12 @@ mod tests { let adapter = GeminiAdapter; let mut decoder = adapter.new_stream_decoder(&target); - // Simulate receiving a function call chunk but no explicit finish. let frame = make_frame( r#"{"candidates":[{"content":{"role":"model","parts":[{"function_call":{"name":"get_weather","args":{"city":"SF"}}}]},"finishReason":null}],"usageMetadata":null}"#, ); decoder.decode_frame(&frame).unwrap(); let events = decoder.finish().unwrap(); - // finish() should emit ToolCallStop for the unclosed block and infer ToolUse. let msg_stop = events .iter() .find(|e| matches!(e, CoreEvent::MessageStop { .. })); @@ -1491,7 +1319,6 @@ mod tests { #[test] fn encode_reasoning_effort_omitted() { - // Gemini does not forward reasoning_effort; verify no panic. let mut core = make_core_request(vec![CoreMessage { role: CoreRole::User, content: vec![CoreContent::Text { @@ -1505,7 +1332,6 @@ mod tests { let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: GeminiRequest = serde_json::from_slice(&proxy_req.body).unwrap(); - // Gemini generation_config should not have reasoning_effort. let config = body.generation_config.unwrap(); assert_eq!(config.temperature, None); } @@ -1543,8 +1369,6 @@ mod tests { } } - // -- Additional tests: upstream model alias, stream flag, malformed, etc -- - #[test] fn encode_uses_upstream_model() { let mut target = make_target(); @@ -1558,7 +1382,6 @@ mod tests { }]); let adapter = GeminiAdapter; let proxy_req = adapter.encode_request(&core, &target).unwrap(); - // The URL should contain the upstream model name. assert!(proxy_req.url.contains("gemini-2.5-flash-preview-05-20")); } @@ -1588,7 +1411,6 @@ mod tests { let target = make_target(); let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: serde_json::Value = serde_json::from_slice(&proxy_req.body).unwrap(); - // Should produce a valid request with empty contents array. assert_eq!(body["contents"].as_array().unwrap().len(), 0); } @@ -1624,7 +1446,6 @@ mod tests { } all_events.extend(decoder.finish().unwrap()); - // Count ToolCallStop events -- should be exactly 1. let tool_call_stop_count = all_events .iter() .filter(|e| matches!(e, CoreEvent::ToolCallStop { .. })) @@ -1635,8 +1456,6 @@ mod tests { ); } - // -- Source guard --------------------------------------------------------- - #[test] fn adapter_source_no_forbidden_imports() { let source = include_str!("gemini.rs"); @@ -1654,8 +1473,6 @@ mod tests { ); } - // -- Helpers -------------------------------------------------------------- - fn make_frame(data: &str) -> SseFrame { SseFrame { event: None, diff --git a/crates/llm-proxy-provider/src/adapter/mod.rs b/crates/llm-proxy-provider/src/adapter/mod.rs index ef48c08..551eeae 100644 --- a/crates/llm-proxy-provider/src/adapter/mod.rs +++ b/crates/llm-proxy-provider/src/adapter/mod.rs @@ -1,18 +1,5 @@ -//! Provider protocol adapters: -//! [`llm_proxy_protocol::core::CoreRequest`] -> provider wire format and back. -//! -//! Each adapter handles exactly one provider protocol: -//! -//! - [`OpenAiChatAdapter`] -- OpenAI Chat Completions API -//! - [`AnthropicAdapter`] -- Anthropic Messages API -//! - [`ResponsesAdapter`] -- OpenAI Responses API -//! - [`GeminiAdapter`] -- Google Gemini GenerateContent API -//! -//! Adapters translate between the normalized core types -//! ([`llm_proxy_protocol::core::CoreRequest`], -//! [`llm_proxy_protocol::core::CoreResponse`], -//! [`llm_proxy_protocol::core::CoreEvent`]) and the provider-specific wire -//! types. They never import client adapters, route handlers, or server state. +//! Provider protocol adapters: translate `llm_proxy_protocol::core` request, +//! response, and event types to and from provider-specific wire formats. pub mod anthropic; pub mod gemini; @@ -38,13 +25,8 @@ use crate::error::ProviderError; use crate::sse::SseFrame; use crate::transport::{AuthHeaders, ProxyRequest}; -/// Anthropic API version header value used when the caller does not supply one. const ANTHROPIC_VERSION: &str = "2023-06-01"; -// --------------------------------------------------------------------------- -// ProviderProtocol -// --------------------------------------------------------------------------- - /// Identifies a provider protocol family. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[non_exhaustive] @@ -60,10 +42,8 @@ pub enum ProviderProtocol { } impl ProviderProtocol { - /// Returns the canonical string name for this protocol. - /// - /// Names use snake_case to match the provider TOML convention - /// (e.g. `openai_chat_completions`, `anthropic_messages`). + /// Returns the canonical snake_case name for this protocol, matching the + /// provider TOML convention (e.g. `openai_chat_completions`, `anthropic_messages`). #[must_use] pub fn name(self) -> &'static str { match self { @@ -74,9 +54,7 @@ impl ProviderProtocol { } } - /// Parse a protocol name string (case-insensitive). - /// - /// Accepts canonical snake_case protocol names. + /// Parses a protocol name (case-insensitive); accepts canonical snake_case names. #[must_use] pub fn parse(name: &str) -> Option { match name.to_ascii_lowercase().as_str() { @@ -89,14 +67,8 @@ impl ProviderProtocol { } } -// --------------------------------------------------------------------------- -// ProviderAdapterTarget -// --------------------------------------------------------------------------- - /// Everything an adapter needs to encode a request for a specific provider. -/// -/// The `api_key` field is redacted in [`fmt::Debug`] output so that -/// `tracing::debug!(?target)` or snapshot output never leaks the secret. +/// The `api_key` field is redacted in [`fmt::Debug`] output so logs never leak it. #[derive(Clone)] pub struct ProviderAdapterTarget { /// Logical provider name (for logging/metrics). @@ -115,12 +87,8 @@ pub struct ProviderAdapterTarget { pub requested_model: String, /// The model to send upstream (may differ due to aliasing). pub upstream_model: String, - /// Optional static headers from the adapter config (e.g. `anthropic-version`). - /// - /// `Arc`-shared from the registry so constructing this target is a refcount - /// bump, not a `HashMap` clone (LOW-5). Mutation (defaulting the - /// `anthropic-version` header) materializes an owned copy via - /// `(*target.headers).clone()` in `build_proxy_request`. + /// Optional static headers from the adapter config (e.g. `anthropic-version`), + /// `Arc`-shared so building a target is a refcount bump, not a map clone. pub headers: Arc>, } @@ -142,14 +110,7 @@ impl fmt::Debug for ProviderAdapterTarget { pub(crate) use crate::transport::endpoint_without_query; -// --------------------------------------------------------------------------- -// ProviderAdapter enum -// --------------------------------------------------------------------------- - /// Enum-dispatched provider adapter. -/// -/// Each variant wraps a concrete adapter that knows how to translate between -/// core types and its provider-specific wire format. #[derive(Debug, Clone)] #[non_exhaustive] pub enum ProviderAdapter { @@ -181,10 +142,6 @@ impl ProviderAdapter { core: &CoreRequest, target: &ProviderAdapterTarget, ) -> Result { - // Log non-empty provider_hints so they are not silently ignored. - // Downgraded to debug since this fires on every request when clients - // consistently send hints -- which is expected behavior. If a specific - // adapter adds hint support, only truly unsupported hints should warn. if !core.provider_hints.raw.is_empty() { tracing::debug!( protocol = %target.protocol.name(), @@ -213,12 +170,8 @@ impl ProviderAdapter { } } - /// Create a new stream decoder for this adapter. - /// - /// Returns a stack-resident [`ProviderStreamDecoderKind`] enum rather than - /// a boxed trait object: the decoder set is closed (4 impls), so an enum - /// avoids the per-stream heap allocation and lets `decode_frame`/`finish` - /// static-dispatch instead of going through a vtable per frame (GAP-LOW-1). + /// Create a new stream decoder for this adapter, returned as a + /// stack-resident [`ProviderStreamDecoderKind`]. pub fn new_stream_decoder(&self, target: &ProviderAdapterTarget) -> ProviderStreamDecoderKind { match self { Self::OpenAiChat(a) => { @@ -235,14 +188,8 @@ impl ProviderAdapter { } } -// --------------------------------------------------------------------------- -// ProviderStreamDecoder trait -// --------------------------------------------------------------------------- - -/// Decodes already-framed SSE events from an upstream provider into core events. -/// -/// The decoder receives parsed [`SseFrame`] objects from [`SseFramer`](crate::sse::SseFramer). -/// It must not parse raw network chunks. +/// Decodes already-framed SSE events from an upstream provider into core +/// events; must not parse raw network chunks. pub trait ProviderStreamDecoder: fmt::Debug + Send { /// Decode a single SSE frame into zero or more core events. fn decode_frame(&mut self, frame: &SseFrame) -> Result, ProviderError>; @@ -251,13 +198,8 @@ pub trait ProviderStreamDecoder: fmt::Debug + Send { fn finish(&mut self) -> Result, ProviderError>; } -/// Owned, stack-resident stream decoder covering the closed 4-adapter set. -/// -/// Replaces `Box` so that a provider stream -/// no longer heap-allocates its decoder and so `decode_frame`/`finish` calls -/// static-dispatch (and inline) instead of going through a vtable on every -/// SSE frame (GAP-LOW-1). [`ProviderAdapter::new_stream_decoder`] constructs -/// the appropriate variant directly. +/// Owned, stack-resident stream decoder covering the closed 4-adapter set; +/// `decode_frame`/`finish` static-dispatch instead of going through a vtable. #[derive(Debug)] pub enum ProviderStreamDecoderKind { /// OpenAI Chat Completions stream decoder. @@ -290,10 +232,6 @@ impl ProviderStreamDecoder for ProviderStreamDecoderKind { } } -// --------------------------------------------------------------------------- -// ProviderAdapterRegistry -// --------------------------------------------------------------------------- - /// Registry of available provider adapters, keyed by protocol. #[derive(Debug, Clone)] pub struct ProviderAdapterRegistry { @@ -342,26 +280,6 @@ impl ProviderAdapterRegistry { } } -// --------------------------------------------------------------------------- -// Shared helpers used by multiple adapters -// --------------------------------------------------------------------------- - -/// Expand URL template placeholders. -/// -/// Currently supports `{model}` -> `upstream_model`. -/// -/// Validates that the model name contains only safe characters. Gemini resource -/// names may contain one leading `models/` segment; all other slashes are -/// rejected. -/// -/// # Path canonicalization -/// -/// The expanded URL is **not** canonicalized (no `..` resolution, no double-slash -/// collapse). This is acceptable because the model name validation above rejects -/// `..` and `/`, so the only way `..` can appear in the URL is via the template -/// itself, which is a static config value controlled by the operator. If user- -/// controlled input is ever interpolated into URLs beyond the model name, path -/// canonicalization must be added here. pub(crate) fn expand_url_template( template: &str, target: &ProviderAdapterTarget, @@ -371,15 +289,11 @@ pub(crate) fn expand_url_template( }; let configured_model = &target.upstream_model; - // When the template uses "models/{model}", strip the "models/" prefix - // from the configured model name to avoid "models/models/..." in the URL. let has_models_prefix = template[..model_pos].ends_with("models/"); let model = configured_model .strip_prefix("models/") .filter(|_| has_models_prefix) .unwrap_or(configured_model); - // Reject model names containing path traversal or other unsafe characters. - // Additionally reject `.` and `..` exactly (path traversal patterns). if model == ".." || model == "." { return Err(ProviderError::InvalidConfig(format!( "upstream_model {:?} is a path traversal pattern; refusing to interpolate into URL", @@ -395,7 +309,6 @@ pub(crate) fn expand_url_template( model ))); } - // Single-pass replacement using the position found above. let mut result = String::with_capacity(template.len() + model.len()); result.push_str(&template[..model_pos]); result.push_str(model); @@ -412,25 +325,16 @@ fn is_safe_model_component(model: &str) -> bool { .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_')) } -/// Map a finish reason string from OpenAI-compatible providers to a core StopReason. pub(crate) fn map_openai_finish_reason(reason: &str) -> StopReason { match reason { "stop" => StopReason::EndTurn, "length" => StopReason::MaxTokens, "tool_calls" | "tool_use" => StopReason::ToolUse, "content_filter" => StopReason::Refusal, - // Preserve the unmapped provider value so it is visible in diagnostics - // (the encoder logs it) rather than silently discarded. _ => StopReason::Unknown(reason.to_owned()), } } -/// Map a Gemini finish reason string to a core StopReason. -/// -/// SAFETY and RECITATION indicate content was filtered by safety systems. -/// These map to `StopReason::Refusal` because the model refused to complete -/// the response due to content policy, which is semantically closer to a -/// refusal than a normal end-of-turn. pub(crate) fn map_gemini_finish_reason(reason: &str) -> StopReason { match reason { "STOP" => StopReason::EndTurn, @@ -440,7 +344,6 @@ pub(crate) fn map_gemini_finish_reason(reason: &str) -> StopReason { } } -/// Build usage from OpenAI-style usage info, accounting for cache tokens. pub(crate) fn build_usage_from_openai( prompt_tokens: i32, completion_tokens: i32, @@ -449,12 +352,6 @@ pub(crate) fn build_usage_from_openai( ) -> Usage { let cache_hit_i64 = cache_hit.unwrap_or(0) as i64; let cache_miss_i64 = cache_miss.unwrap_or(0) as i64; - // OpenAI's `prompt_tokens` is INCLUSIVE of cached tokens (both cache-hit and - // cache-creation/miss contributions are counted within it), whereas the core - // `input_tokens` field is EXCLUSIVE of cache tokens (which are reported - // separately via `cache_read_input_tokens` / `cache_creation_input_tokens`). - // Subtract the cache contributions here so the core total does not - // double-count them. let input = (prompt_tokens as i64) - cache_hit_i64 - cache_miss_i64; let input_tokens = i32::try_from(input.max(0)).unwrap_or(i32::MAX); @@ -468,16 +365,12 @@ pub(crate) fn build_usage_from_openai( } } -/// Build a `ProxyRequest` from the adapter's encoded JSON body and target info. pub(crate) fn build_proxy_request( body: Vec, target: &ProviderAdapterTarget, stream: bool, url: String, ) -> ProxyRequest { - // Materialize an owned copy of the shared adapter headers so we can default - // `anthropic-version`; this is the only per-request clone of the headers map - // (the registry → target hand-off is an `Arc` refcount bump, LOW-5). let mut extra_headers = (*target.headers).clone(); if target.protocol == ProviderProtocol::AnthropicMessages && !extra_headers @@ -499,7 +392,6 @@ pub(crate) fn build_proxy_request( } } -/// Build the response ModelRef preserving the client-facing requested model. pub(crate) fn response_model_ref(target: &ProviderAdapterTarget) -> ModelRef { let upstream = if target.upstream_model != target.requested_model { Some(target.upstream_model.clone()) @@ -512,11 +404,6 @@ pub(crate) fn response_model_ref(target: &ProviderAdapterTarget) -> ModelRef { } } -/// Truncate a string to `max_len` bytes, respecting UTF-8 char boundaries. -/// -/// If the string is longer than `max_len`, finds the nearest char boundary at -/// or before `max_len`. This prevents panics when slicing strings containing -/// multi-byte UTF-8 characters (e.g. CJK text, emoji). pub(crate) fn truncate_str_safe(s: &str, max_len: usize) -> &str { if s.len() <= max_len { s @@ -529,16 +416,10 @@ pub(crate) fn truncate_str_safe(s: &str, max_len: usize) -> &str { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; - // -- ProviderProtocol name/parse ----------------------------------------- - #[test] fn protocol_name_roundtrip() { let protocols = [ @@ -574,8 +455,6 @@ mod tests { assert_eq!(ProviderProtocol::parse(""), None); } - // -- Registry ----------------------------------------------------------- - #[test] fn registry_builtin_has_all_protocols() { let reg = ProviderAdapterRegistry::builtin(); @@ -609,8 +488,6 @@ mod tests { assert!(!reg.has_protocol_name("anthropic")); } - // -- Adapter protocol method --------------------------------------------- - #[test] fn adapter_protocol_dispatch() { let reg = ProviderAdapterRegistry::builtin(); @@ -640,8 +517,6 @@ mod tests { ); } - // -- URL template expansion ---------------------------------------------- - #[test] fn url_template_expands_model() { let target = make_target(ProviderProtocol::GeminiGenerateContent); @@ -755,8 +630,6 @@ mod tests { assert!(!debug.contains("query-secret")); } - // -- Finish reason mapping ----------------------------------------------- - #[test] fn openai_finish_reason_mappings() { assert_eq!(map_openai_finish_reason("stop"), StopReason::EndTurn); @@ -767,7 +640,6 @@ mod tests { map_openai_finish_reason("content_filter"), StopReason::Refusal ); - // Unmapped values are preserved in Unknown(String) (GAP-LOW-9). assert_eq!( map_openai_finish_reason("unknown"), StopReason::Unknown("unknown".to_owned()) @@ -783,19 +655,16 @@ mod tests { ); assert_eq!(map_gemini_finish_reason("SAFETY"), StopReason::Refusal); assert_eq!(map_gemini_finish_reason("RECITATION"), StopReason::Refusal); - // Unmapped values are preserved in Unknown(String) (GAP-LOW-9). assert_eq!( map_gemini_finish_reason("other"), StopReason::Unknown("other".to_owned()) ); } - // -- Usage building ------------------------------------------------------ - #[test] fn usage_from_openai_basic() { let usage = build_usage_from_openai(100, 50, Some(10), Some(20)); - assert_eq!(usage.input_tokens, 70); // 100 - 10 - 20 + assert_eq!(usage.input_tokens, 70); assert_eq!(usage.output_tokens, 50); assert_eq!(usage.cache_read_input_tokens, Some(10)); assert_eq!(usage.cache_creation_input_tokens, Some(20)); @@ -815,8 +684,6 @@ mod tests { assert_eq!(usage.input_tokens, 0); } - // -- Response model ref -------------------------------------------------- - #[test] fn response_model_ref_no_alias() { let target = ProviderAdapterTarget { @@ -853,8 +720,6 @@ mod tests { assert_eq!(mr.upstream.as_deref(), Some("gpt-4o-2024-08-06")); } - // -- Source guard --------------------------------------------------------- - #[test] fn adapter_source_no_forbidden_imports() { let source = include_str!("mod.rs"); @@ -872,8 +737,6 @@ mod tests { ); } - // -- truncate_str_safe --------------------------------------------------- - #[test] fn truncate_str_safe_ascii_within_limit() { assert_eq!(truncate_str_safe("hello", 10), "hello"); @@ -886,12 +749,9 @@ mod tests { #[test] fn truncate_str_safe_multibyte_no_panic() { - // Japanese 'あ' is 3 bytes. If max_len lands mid-character, we must - // back up to the previous char boundary. - let s = "あいうえお"; // 15 bytes total - // max_len=4 lands inside 'い' (bytes 3-5), so we back up to byte 3. + let s = "あいうえお"; let truncated = truncate_str_safe(s, 4); - assert_eq!(truncated, "あ"); // Only first char (bytes 0-2) fits + assert_eq!(truncated, "あ"); assert!(truncated.len() <= 4); } @@ -930,8 +790,6 @@ mod tests { assert!(!request.extra_headers.contains_key("anthropic-version")); } - // -- Helpers -------------------------------------------------------------- - fn make_target(protocol: ProviderProtocol) -> ProviderAdapterTarget { let (endpoint, auth_style) = match protocol { ProviderProtocol::OpenAiChatCompletions => ( diff --git a/crates/llm-proxy-provider/src/adapter/openai_chat.rs b/crates/llm-proxy-provider/src/adapter/openai_chat.rs index c547c70..b7923b2 100644 --- a/crates/llm-proxy-provider/src/adapter/openai_chat.rs +++ b/crates/llm-proxy-provider/src/adapter/openai_chat.rs @@ -1,12 +1,5 @@ -//! OpenAI Chat Completions provider adapter. -//! -//! Translates between core types and the OpenAI Chat Completions wire format: -//! -//! ```text -//! CoreRequest -> ChatCompletionRequest -//! ChatCompletionResponse -> CoreResponse -//! ChatCompletionChunk stream -> CoreEvent stream -//! ``` +//! OpenAI Chat Completions provider adapter: translates core requests, +//! responses, and event streams to and from the Chat Completions wire format. use std::collections::HashMap; @@ -26,35 +19,11 @@ use super::{ use crate::error::ProviderError; use crate::sse::SseFrame; -// --------------------------------------------------------------------------- -// Adapter struct -// --------------------------------------------------------------------------- - /// Adapter for the OpenAI Chat Completions API. #[derive(Debug, Clone, Default)] pub struct OpenAiChatAdapter; -// --------------------------------------------------------------------------- -// Stream decoder -// --------------------------------------------------------------------------- - -/// Stateful stream decoder for OpenAI Chat Completions SSE frames. -/// -/// ## CoreEvent variants emitted -/// -/// - `MessageStart` -- on the first chunk received -/// - `ContentStart` -- on first text or reasoning delta -/// - `TextDelta` -- on `delta.content` -/// - `ThinkingDelta` -- on `delta.reasoning_content` (or `delta.reasoning`) -/// - `ToolCallStart` -- on new tool calls in `delta.tool_calls` -/// - `ToolCallDelta` -- on tool call argument deltas -/// - `ToolCallStop` -- on stream end for unclosed tool blocks -/// - `UsageDelta` -- on usage-only chunks or final chunks with usage -/// - `MessageStop` -- on `finish_reason` or stream end -/// -/// Intentionally never emitted: `Ping` (OpenAI Chat has no ping mechanism), -/// `Error` (OpenAI errors are handled at the transport level, not in stream -/// decoding). +/// Stateful stream decoder translating OpenAI Chat Completions SSE frames into `CoreEvent`s. #[derive(Debug)] pub struct OpenAiChatStreamDecoder { model_ref: ModelRef, @@ -62,7 +31,6 @@ pub struct OpenAiChatStreamDecoder { content_index: usize, content_started: bool, reasoning_started: bool, - /// Maps OpenAI tool-call index to our content block index. tool_blocks: HashMap, stop_sent: bool, } @@ -93,7 +61,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { }); } - // Usage-only chunk (no choices). if chunk.choices.is_empty() { if let Some(ref usage) = chunk.usage { events.push(CoreEvent::UsageDelta { @@ -110,13 +77,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { let choice = &chunk.choices[0]; - // Handle reasoning content (both `reasoning_content` and `reasoning`). - // The serde alias on ChatMessage ensures both field names deserialize - // into `reasoning_content`. - // - // Empty-string reasoning deltas are intentionally skipped: OpenAI sends - // them as keep-alives with no semantic meaning, and usage-only / - // finish-reason-only chunks are handled in separate code paths below. if let Some(reasoning) = choice .delta .as_ref() @@ -138,12 +98,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { } } - // Handle text content deltas. - // - // Empty-string content deltas are intentionally skipped: OpenAI sends - // them as keep-alives with no semantic meaning. The finish_reason and - // usage are handled in separate code paths, so this truthy check does not - // drop meaningful events. if let Some(ref delta) = choice.delta { if !delta.content_is_empty() { self.close_reasoning_if_open(); @@ -161,7 +115,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { } } - // Handle tool call deltas. if let Some(ref delta) = choice.delta && !delta.tool_calls.is_empty() { @@ -169,7 +122,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { self.push_tool_call_deltas(&delta.tool_calls, &mut events); } - // Handle finish reason. if let Some(ref reason) = choice.finish_reason { if !reason.is_empty() && !self.stop_sent { self.close_content_if_open(); @@ -231,7 +183,6 @@ impl ProviderStreamDecoder for OpenAiChatStreamDecoder { } impl OpenAiChatStreamDecoder { - /// Close any open text content block. Reserved for future ContentStop emission. fn close_text_if_open(&mut self) { if self.content_started { self.content_started = false; @@ -239,7 +190,6 @@ impl OpenAiChatStreamDecoder { } } - /// Close any open reasoning content block. Reserved for future ContentStop emission. fn close_reasoning_if_open(&mut self) { if self.reasoning_started { self.reasoning_started = false; @@ -247,7 +197,6 @@ impl OpenAiChatStreamDecoder { } } - /// Close any open text or reasoning content blocks. fn close_content_if_open(&mut self) { if self.content_started { self.content_started = false; @@ -259,18 +208,8 @@ impl OpenAiChatStreamDecoder { } } - /// Translates a chunk's `tool_calls` array into core tool-call events. - /// - /// Opens a block the first time an upstream index is seen and routes - /// argument fragments to the block that index owns, so parallel calls stay - /// separated. fn push_tool_call_deltas(&mut self, tool_calls: &[ToolCall], events: &mut Vec) { for tc in tool_calls { - // The wire `index` is `Option`; negative values are - // semantically meaningless but a malformed/malicious chunk could - // carry one. An unchecked `as usize` would wrap (e.g. `-1` -> - // `usize::MAX`), producing phantom tool blocks with huge indices. - // Validate before casting: negatives fall back to 0 and are logged. let oi = tc .index .and_then(|i| usize::try_from(i).ok()) @@ -284,7 +223,6 @@ impl OpenAiChatStreamDecoder { 0 }); - // New tool call? if !self.tool_blocks.contains_key(&oi) { let func_name = tc .function @@ -316,7 +254,6 @@ impl OpenAiChatStreamDecoder { self.content_index += 1; } - // Argument delta. if let Some(ref func) = tc.function && let Some(ref args) = func.arguments && !args.is_empty() @@ -339,13 +276,6 @@ impl OpenAiChatStreamDecoder { } } -// --------------------------------------------------------------------------- -// OpenAiChatAdapter impl -// --------------------------------------------------------------------------- - -/// Folds one assistant content block into the outgoing tool calls / reasoning. -/// -/// Blocks the Chat Completions schema cannot carry are dropped with a warning. fn encode_assistant_content( content: &CoreContent, tool_calls: &mut Vec, @@ -381,10 +311,6 @@ fn encode_assistant_content( } } -/// Appends one `role: "tool"` message per tool result in `content`. -/// -/// Chat Completions carries tool results as standalone messages rather than as -/// content blocks on the user turn. fn push_tool_result_messages(content: &[CoreContent], messages: &mut Vec) { for c in content { let CoreContent::ToolResult { @@ -411,6 +337,12 @@ fn push_tool_result_messages(content: &[CoreContent], messages: &mut Vec { @@ -479,13 +410,6 @@ impl OpenAiChatAdapter { encode_assistant_content(content, &mut tool_calls, &mut reasoning_content); } - // When an assistant turn carries only tool calls (no text), - // emit `content: null` rather than `content: ""`. The - // OpenAI Chat Completions API tolerates both, but some - // strict OpenAI-compatible upstreams (e.g. Ollama) treat an - // empty string differently from null/absent and drop out of - // tool-calling mode. `null` is the conventionally-correct - // representation and is accepted by both OpenAI and Ollama. let content = if text.is_empty() && !tool_calls.is_empty() { serde_json::Value::Null } else { @@ -550,7 +474,6 @@ impl OpenAiChatAdapter { }); } _ => { - // Handle future CoreRole variants as user messages. let text = collect_text(&msg.content); if !text.is_empty() { messages.push(ChatMessage { @@ -589,7 +512,6 @@ impl OpenAiChatAdapter { }) .collect(); - // Tool choice -- omit (None) for unknown variants instead of sending null. let tool_choice = core.tool_choice.as_ref().and_then(|tc| match tc { CoreToolChoice::Auto => Some(serde_json::json!({"type": "auto"})), CoreToolChoice::Any => Some(serde_json::json!({"type": "required"})), @@ -620,8 +542,6 @@ impl OpenAiChatAdapter { if v.len() == 1 { serde_json::Value::String(v[0].clone()) } else { - // Manual serialization avoids an infallible expect. - // Vec always serializes to a JSON array of strings. serde_json::Value::Array( v.iter() .map(|s| serde_json::Value::String(s.clone())) @@ -634,8 +554,6 @@ impl OpenAiChatAdapter { extra: serde_json::Map::new(), }; - // Propagate the client's include_usage preference to the upstream request. - // If the client didn't request usage, we don't ask upstream for it either. if core.stream { let include_usage = core .provider_hints @@ -678,7 +596,6 @@ impl OpenAiChatAdapter { let mut content = Vec::new(); - // Reasoning -> Thinking (typically comes first from the provider). if let Some(ref reasoning) = msg.reasoning_content { if !reasoning.is_empty() { content.push(CoreContent::Thinking { @@ -688,7 +605,6 @@ impl OpenAiChatAdapter { } } - // Text content (may appear before or after tool calls). if !msg.content_is_empty() { content.push(CoreContent::Text { text: msg.content_text(), @@ -696,7 +612,6 @@ impl OpenAiChatAdapter { }); } - // Refusal. if let Some(ref refusal) = msg.refusal { if !refusal.is_empty() { content.push(CoreContent::Refusal { @@ -705,7 +620,6 @@ impl OpenAiChatAdapter { } } - // Tool calls -> ToolUse. for tc in &msg.tool_calls { let raw_args = tc.function.as_ref().and_then(|f| f.arguments.as_ref()); let input = raw_args @@ -736,7 +650,6 @@ impl OpenAiChatAdapter { }); } - // Guarantee at least one content block. if content.is_empty() { content.push(CoreContent::Text { text: String::new(), @@ -748,8 +661,6 @@ impl OpenAiChatAdapter { .finish_reason .as_deref() .map(map_openai_finish_reason) - // Provider omitted `finish_reason`; record that it was absent - // rather than unmapped so the diagnostic is unambiguous. .unwrap_or(StopReason::Unknown("absent".to_owned())); let usage = build_usage_from_openai( @@ -785,15 +696,6 @@ impl OpenAiChatAdapter { } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Collect all text from a list of content blocks. -/// -/// Non-text blocks are silently omitted, but a warning is logged since this -/// means data is being dropped. Per the plan's lossy translation rules, every -/// dropped feature must be either rejected, warned, or preserved opaquely. fn collect_text(content: &[CoreContent]) -> String { let mut text = String::new(); for c in content { @@ -817,10 +719,6 @@ fn collect_text(content: &[CoreContent]) -> String { text } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -859,8 +757,6 @@ mod tests { } } - // -- encode_request tests ------------------------------------------------ - #[test] fn encode_text_request() { let core = make_core_request(vec![CoreMessage { @@ -1007,7 +903,6 @@ mod tests { }], }]); core.stream = true; - // Simulate the client requesting usage via provider_hints. core.provider_hints.raw.insert( "stream_options".into(), serde_json::json!({"include_usage": true}), @@ -1033,7 +928,6 @@ mod tests { }], }]); core.stream = true; - // No provider_hints set -- include_usage should default to false. let adapter = OpenAiChatAdapter; let target = make_target(); let proxy_req = adapter.encode_request(&core, &target).unwrap(); @@ -1088,8 +982,6 @@ mod tests { assert_eq!(body.stop, Some(serde_json::json!("END"))); } - // -- decode_response tests ----------------------------------------------- - #[test] fn decode_text_response() { let target = make_target(); @@ -1239,8 +1131,6 @@ mod tests { assert_eq!(core_resp.model.requested, "gpt-4o"); } - // -- Streaming tests ----------------------------------------------------- - #[test] fn stream_text_decoding() { let target = make_target(); @@ -1331,9 +1221,6 @@ mod tests { #[test] fn stream_negative_tool_index_clamped() { - // A malformed/malicious chunk carrying a negative `index` must not wrap - // to usize::MAX and create a phantom tool block. It should be clamped - // to 0 and still produce a valid ToolCallStart at index 0. let target = make_target(); let adapter = OpenAiChatAdapter; let mut decoder = adapter.new_stream_decoder(&target); @@ -1343,8 +1230,6 @@ mod tests { ); let events = decoder.decode_frame(&frame).unwrap(); - // Exactly one ToolCallStart, and its content_index is 0 (not a wrapped - // huge value). let starts: Vec<_> = events .iter() .filter_map(|e| match e { @@ -1359,9 +1244,6 @@ mod tests { #[test] fn encode_assistant_only_tool_calls_emits_null_content() { - // An assistant turn with only ToolUse blocks (no text) must serialize - // `content: null`, not `content: ""`, for compatibility with strict - // OpenAI-compatible upstreams. let core = make_core_request(vec![CoreMessage { role: CoreRole::Assistant, content: vec![CoreContent::ToolUse { @@ -1386,8 +1268,6 @@ mod tests { #[test] fn encode_assistant_with_text_keeps_string_content() { - // An assistant turn with real text must keep `content` as the string - // (regression guard for the null-content fix). let core = make_core_request(vec![CoreMessage { role: CoreRole::Assistant, content: vec![ @@ -1446,8 +1326,6 @@ mod tests { assert_has_event(&events, "MessageStop"); } - // -- Missing tests ------------------------------------------------------- - #[test] fn decode_thinking_response() { let target = make_target(); @@ -1494,7 +1372,6 @@ mod tests { #[test] fn decode_reasoning_alias_field() { - // Verify that the `reasoning` field alias works via serde. let json = r#"{"role":"assistant","content":"hi","reasoning":"thinking tokens"}"#; let msg: ChatMessage = serde_json::from_str(json).unwrap(); assert_eq!(msg.reasoning_content, Some("thinking tokens".to_owned())); @@ -1598,7 +1475,6 @@ mod tests { #[test] fn stream_reasoning_alias() { - // Verify that streaming chunks with `reasoning` (not `reasoning_content`) work. let target = make_target(); let adapter = OpenAiChatAdapter; let mut decoder = adapter.new_stream_decoder(&target); @@ -1615,8 +1491,6 @@ mod tests { ); } - // -- Additional missing tests ---------------------------------------------- - #[test] fn encode_reasoning_effort_forwarded() { let mut core = make_core_request(vec![CoreMessage { @@ -1693,7 +1567,6 @@ mod tests { #[test] fn encode_tool_choice_none_omitted() { - // Verify that unknown tool_choice variants produce None (omitted). let mut core = make_core_request(vec![CoreMessage { role: CoreRole::User, content: vec![CoreContent::Text { @@ -1701,7 +1574,6 @@ mod tests { cache: None, }], }]); - // Raw null is a valid variant, but we can test with Auto to ensure it's Some. core.tool_choice = Some(CoreToolChoice::Auto); let adapter = OpenAiChatAdapter; let target = make_target(); @@ -1736,7 +1608,6 @@ mod tests { #[test] fn decode_preserves_content_order_thinking_before_text() { - // Verify that Thinking appears before Text (provider order). let target = make_target(); let resp = llm_proxy_protocol::openai::ChatCompletionResponse { id: "test".into(), @@ -1772,7 +1643,6 @@ mod tests { let adapter = OpenAiChatAdapter; let core_resp = adapter.decode_response(&bytes, &target).unwrap(); - // Thinking must come before Text. let thinking_idx = core_resp .content .iter() @@ -1845,7 +1715,6 @@ mod tests { #[test] fn stream_finish_reason_only_chunk() { - // Test a chunk that only has finish_reason with no prior content. let target = make_target(); let adapter = OpenAiChatAdapter; let mut decoder = adapter.new_stream_decoder(&target); @@ -1873,8 +1742,6 @@ mod tests { ))); } - // -- Source guard --------------------------------------------------------- - #[test] fn adapter_source_no_forbidden_imports() { let source = include_str!("openai_chat.rs"); @@ -1892,8 +1759,6 @@ mod tests { ); } - // -- Helpers -------------------------------------------------------------- - fn make_frame(data: &str) -> crate::sse::SseFrame { crate::sse::SseFrame { event: None, diff --git a/crates/llm-proxy-provider/src/adapter/responses.rs b/crates/llm-proxy-provider/src/adapter/responses.rs index 0839ce6..cdd76fa 100644 --- a/crates/llm-proxy-provider/src/adapter/responses.rs +++ b/crates/llm-proxy-provider/src/adapter/responses.rs @@ -1,12 +1,5 @@ -//! OpenAI Responses API provider adapter. -//! -//! Translates between core types and the Responses API wire format: -//! -//! ```text -//! CoreRequest -> ResponsesRequest -//! ResponsesResponse -> CoreResponse -//! ResponsesChunk stream -> CoreEvent stream -//! ``` +//! OpenAI Responses API provider adapter: translates core requests, responses, +//! and event streams to and from the OpenAI Responses wire format. use std::collections::HashMap; @@ -26,56 +19,30 @@ use super::{ use crate::error::ProviderError; use crate::sse::SseFrame; -// --------------------------------------------------------------------------- -// Adapter struct -// --------------------------------------------------------------------------- - /// Adapter for the OpenAI Responses API. #[derive(Debug, Clone, Default)] pub struct ResponsesAdapter; -// --------------------------------------------------------------------------- -// Stream decoder -// --------------------------------------------------------------------------- - -/// Stateful stream decoder for Responses API SSE frames. -/// -/// ## CoreEvent variants emitted -/// -/// - `MessageStart` -- on `response.created` or `response.in_progress` -/// - `ContentStart` -- on `response.output_text.delta` (text) or -/// `response.output_item.added` with `function_call` (tool use) -/// - `TextDelta` -- on `response.output_text.delta` -/// - `ToolCallStart` -- on `response.output_item.added` with `function_call` -/// (one per upstream `output_index`, supporting parallel calls) -/// - `ToolCallDelta` -- on `response.function_call_arguments.delta` -/// (routed to the open block matching the chunk's `output_index`) -/// - `ToolCallStop` -- on `response.function_call_arguments.done`; any blocks -/// still open are closed on `response.completed`/`response.done`/`finish()` -/// - `UsageDelta` -- on `response.completed` or `response.done` with usage -/// - `MessageStop` -- on `response.completed` or `response.done` -/// - `Error` -- on `response.failed` -/// -/// Intentionally never emitted: `Ping`, `ThinkingDelta` (the Responses API -/// does not produce heartbeat or thinking-stream events). +/// Stateful stream decoder translating OpenAI Responses API SSE frames into `CoreEvent`s. #[derive(Debug)] pub struct ResponsesStreamDecoder { model_ref: ModelRef, started: bool, content_index: usize, content_started: bool, - /// Open tool-call blocks keyed by the upstream `output_index`, mapping to - /// the core `content_index` assigned when the block started. The Responses - /// API interleaves parallel function-call argument deltas and - /// disambiguates them via `output_index`, so the decoder must track one - /// block per concurrent call rather than a single open call. tool_blocks: HashMap, - /// Whether any tool call was seen during this stream. saw_tool_call: bool, stop_sent: bool, } impl ProviderStreamDecoder for ResponsesStreamDecoder { + #[expect( + clippy::too_many_lines, + reason = "one match arm per Responses API stream event: the wire shape \ + is a single translation table, and splitting the arms into \ + helpers would separate each event's field mapping from the \ + frame layout it produces" + )] fn decode_frame(&mut self, frame: &SseFrame) -> Result, ProviderError> { let data = frame.data.trim(); if data.is_empty() || data == "[DONE]" { @@ -141,9 +108,6 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { && !delta.is_empty() { let oi = chunk.output_index.unwrap_or(0); - // If no block is open for this output_index (e.g. a missed - // response.output_item.added event), synthesize a start - // keyed by it before routing the delta. let block_idx = self.ensure_tool_block( oi, "no prior response.output_item.added event", @@ -157,9 +121,6 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { } "response.function_call_arguments.done" => { let oi = chunk.output_index.unwrap_or(0); - // If no block is open for this output_index (e.g. a missed - // response.output_item.added event), emit a synthetic start so - // the stop has a matching ToolCallStart, then close it. self.ensure_tool_block( oi, "no prior response.output_item.added event before arguments.done", @@ -173,16 +134,12 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { self.close_content_if_open(); self.close_tool_blocks(&mut events); - // Extract usage. if let Some(ref usage) = chunk.usage { events.push(CoreEvent::UsageDelta { usage: build_responses_usage(usage), }); } - // Extract stop reason from output if present, falling back to - // saw_tool_call for streams where function_call output items - // may not be present on the response.completed event. let stop_reason = self.infer_stop_reason(chunk.output.as_deref()); if !self.stop_sent { @@ -219,7 +176,6 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { .and_then(|e| e.get("message").and_then(|m| m.as_str())) .unwrap_or("response failed") .to_owned(); - // Map error code to more specific CoreStreamErrorKind when available. let kind = chunk .error .as_ref() @@ -244,9 +200,7 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { error: llm_proxy_protocol::core::CoreStreamError::new(kind, message), }); } - _ => { - // Unknown event type; skip. - } + _ => {} } Ok(events) @@ -265,11 +219,6 @@ impl ProviderStreamDecoder for ResponsesStreamDecoder { self.close_content_if_open(); - // Flush any unclosed tool calls (e.g. a stream that emitted - // ToolCallStart/arguments.delta but ended without arguments.done or - // response.completed), so downstream consumers receive a balanced - // ToolCallStart/ToolCallStop pair. Mirrors the Gemini/OpenAI-chat - // decoders' unclosed-tool-block handling in finish(). self.close_tool_blocks(&mut events); if !self.stop_sent { @@ -297,12 +246,6 @@ impl ResponsesStreamDecoder { } } - /// Emit `ToolCallStop` for every still-open tool block (sorted by - /// `content_index` for a deterministic, ascending close order) and clear - /// the map. Mirrors the Gemini/OpenAI-chat decoders' `close_tool_blocks` - /// so every `ToolCallStart` is guaranteed a matching `ToolCallStop`. - /// Opens a text content block when a `message` output first carries - /// `output_text` and none is currently open. fn start_text_block_if_needed( &mut self, output: &ResponsesOutput, @@ -325,10 +268,6 @@ impl ResponsesStreamDecoder { } } - /// Opens a tool-call block for `oi` unless one is already open. - /// - /// Parallel calls are disambiguated by the upstream `output_index`, so one - /// block is tracked per concurrent call. fn start_function_call( &mut self, oi: usize, @@ -358,9 +297,6 @@ impl ResponsesStreamDecoder { }); } - /// Returns the content index of the tool block for `oi`, synthesizing a - /// `ToolCallStart` (logged with `warn_message`) when the upstream never sent - /// a `response.output_item.added` for it. fn ensure_tool_block( &mut self, oi: usize, @@ -398,10 +334,6 @@ impl ResponsesStreamDecoder { self.tool_blocks.clear(); } - /// Infer the stop reason from output items and whether a tool call was seen. - /// - /// Checks for `function_call` output items in the chunk. If none are found, - /// falls back to the `saw_tool_call` flag that was set during streaming. fn infer_stop_reason(&self, outputs: Option<&[ResponsesOutput]>) -> StopReason { if let Some(outs) = outputs { if outs.iter().any(|o| o.r#type == "function_call") { @@ -416,15 +348,6 @@ impl ResponsesStreamDecoder { } } -// --------------------------------------------------------------------------- -// ResponsesAdapter impl -// --------------------------------------------------------------------------- - -/// Folds one core content block into the Responses input item list. -/// -/// Text accumulates in `text_parts` so consecutive text blocks collapse into a -/// single input item; a tool block flushes that buffer first so a message never -/// emits duplicate role entries. fn encode_content_item<'a>( c: &'a CoreContent, role: &str, @@ -445,7 +368,6 @@ fn encode_content_item<'a>( name, input: tool_input, } => { - // Flush pending text as its own input item first. if !text_parts.is_empty() { input.push(ResponsesInput { role: role.to_owned(), @@ -454,8 +376,6 @@ fn encode_content_item<'a>( text_parts.clear(); } - // The Responses API represents prior tool calls as input items with - // type "function_call". let arguments = serde_json::to_string(tool_input).unwrap_or_else(|e| { tracing::warn!( error = %e, @@ -511,10 +431,6 @@ fn encode_content_item<'a>( } } -/// Collects the text of every content block on a `message` output. -/// -/// Unknown block types still contribute their `text` field, if any, rather than -/// being dropped silently. fn push_message_text(output: &ResponsesOutput, content: &mut Vec) { for block in output.content.iter().flatten() { if block.r#type != "output_text" { @@ -534,6 +450,12 @@ fn push_message_text(output: &ResponsesOutput, content: &mut Vec) { impl ResponsesAdapter { /// Encode a core request into a Responses API request. + #[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the Responses API request shape, \ + kept in one place so the mapping can be audited against the \ + API schema without jumping between helpers" + )] pub fn encode_request( &self, core: &CoreRequest, @@ -541,9 +463,7 @@ impl ResponsesAdapter { ) -> Result { let mut input = Vec::new(); - // System prompt as first developer message. if !core.system.is_empty() { - // Warn about non-text system content blocks (consistent with other adapters). for c in &core.system { if !matches!(c, CoreContent::Text { .. }) { tracing::warn!( @@ -574,31 +494,21 @@ impl ResponsesAdapter { } } - // Conversation messages. for msg in &core.messages { let role = match msg.role { CoreRole::User => "user", CoreRole::Assistant => "assistant", CoreRole::System => "developer", CoreRole::Tool => "user", - // Future CoreRole variants are mapped to "user" as a safe default. - // Update this match when new variants are added to CoreRole. _ => "user", }; - // Build content from all supported types, not just text. let mut text_parts = Vec::new(); for c in &msg.content { encode_content_item(c, role, &mut text_parts, &mut input); } - // Flush any remaining text content as a separate same-role input - // item. Trailing text after a tool block (e.g. an assistant message - // shaped [ToolUse, Text] or [ToolResult, Text]) must be preserved, - // and the Responses API permits multiple input items with the same - // role, so flush unconditionally rather than gating on the absence - // of function_call/function_call_output. if !text_parts.is_empty() { let text: String = text_parts.join(""); input.push(ResponsesInput { @@ -608,7 +518,6 @@ impl ResponsesAdapter { } } - // Tools. let tools: Vec = core .tools .iter() @@ -624,7 +533,6 @@ impl ResponsesAdapter { }) .collect(); - // Reasoning effort. let reasoning = core .sampling .reasoning_effort @@ -642,7 +550,6 @@ impl ResponsesAdapter { tool_choice: None, }; - // Warn about sampling fields that the Responses API does not support. if core.sampling.temperature.is_some() { tracing::warn!( "Responses: temperature specified but not forwarded; \ @@ -674,8 +581,6 @@ impl ResponsesAdapter { ); } - // Forward tool_choice if present. Omit unknown variants instead of - // sending null (per the plan's lossy translation rules). if let Some(ref tc) = core.tool_choice { req.tool_choice = match tc { CoreToolChoice::Auto => Some(serde_json::json!({"type": "auto"})), @@ -716,8 +621,6 @@ impl ResponsesAdapter { } "function_call" => { let input = output - // Lazily build the empty object only on the `None`/parse-fail path - // so the common tool-calling happy path avoids a heap allocation (MEDIUM-4). .arguments .as_ref() .and_then(|args| serde_json::from_str::(args).ok()) @@ -738,7 +641,6 @@ impl ResponsesAdapter { } } - // Guarantee at least one content block. if content.is_empty() { content.push(CoreContent::Text { text: String::new(), @@ -784,11 +686,6 @@ impl ResponsesAdapter { } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Build usage from Responses API usage. fn build_responses_usage(usage: &ResponsesUsage) -> Usage { Usage { input_tokens: usage.input_tokens, @@ -800,10 +697,6 @@ fn build_responses_usage(usage: &ResponsesUsage) -> Usage { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -842,8 +735,6 @@ mod tests { } } - // -- Encode tests -------------------------------------------------------- - #[test] fn encode_text_request() { let core = make_core_request(vec![CoreMessage { @@ -945,8 +836,6 @@ mod tests { assert!(proxy_req.stream); } - // -- Decode tests -------------------------------------------------------- - #[test] fn decode_text_response() { let target = make_target(); @@ -1056,8 +945,6 @@ mod tests { assert_eq!(core_resp.model.requested, "my-model"); } - // -- Streaming tests ----------------------------------------------------- - #[test] fn stream_text_decoding() { let target = make_target(); @@ -1129,8 +1016,6 @@ mod tests { ); } - // -- Stream tool call / malformed frame / response.failed tests ----------- - #[test] fn stream_tool_call_decoding() { let target = make_target(); @@ -1183,26 +1068,18 @@ mod tests { #[test] fn stream_parallel_tool_calls_partitioned_by_output_index() { - // The Responses API interleaves parallel function-call argument deltas - // and disambiguates them via `output_index`. The decoder must emit two - // distinct ToolCallStart events with different indices and route each - // delta to the correct content_index so the two argument streams are - // partitioned rather than merged into one invalid JSON blob. let target = make_target(); let adapter = ResponsesAdapter; let mut decoder = adapter.new_stream_decoder(&target); let frames = vec![ make_frame(r#"{"type":"response.created","id":"resp_1"}"#), - // First parallel call (output_index 0). make_frame( r#"{"type":"response.output_item.added","output_index":0,"output":[{"type":"function_call","call_id":"call_a","name":"get_weather"}]}"#, ), - // Second parallel call (output_index 1) before the first completes. make_frame( r#"{"type":"response.output_item.added","output_index":1,"output":[{"type":"function_call","call_id":"call_b","name":"get_time"}]}"#, ), - // Interleaved argument deltas. make_frame( r#"{"type":"response.function_call_arguments.delta","output_index":0,"delta":"{\"city\":"}"#, ), @@ -1215,7 +1092,6 @@ mod tests { make_frame( r#"{"type":"response.function_call_arguments.delta","output_index":1,"delta":"\"UTC\"}"}"#, ), - // Close the two calls out of arrival order (call_b first). make_frame(r#"{"type":"response.function_call_arguments.done","output_index":1}"#), make_frame(r#"{"type":"response.function_call_arguments.done","output_index":0}"#), make_frame( @@ -1228,7 +1104,6 @@ mod tests { all_events.extend(decoder.decode_frame(frame).unwrap()); } - // Two distinct ToolCallStart events with different indices. let starts: Vec<_> = all_events .iter() .filter_map(|e| match e { @@ -1242,7 +1117,6 @@ mod tests { let indices: Vec<_> = starts.iter().map(|(i, _, _)| *i).collect(); assert_eq!(indices, vec![0, 1], "expected distinct ascending indices"); - // The argument deltas must be partitioned per content_index, not merged. let mut args_by_index: std::collections::HashMap = std::collections::HashMap::new(); for e in &all_events { @@ -1259,7 +1133,6 @@ mod tests { "output_index 1 args must be partitioned" ); - // Each started block gets a matching stop. let stops: Vec<_> = all_events .iter() .filter_map(|e| match e { @@ -1279,7 +1152,6 @@ mod tests { let frame = make_frame("not valid json {{{"); let events = decoder.decode_frame(&frame).unwrap(); - // Malformed frame should return empty (no MessageStart yet since type is unknown). assert!(events.is_empty()); } @@ -1299,8 +1171,6 @@ mod tests { ); } - // -- Encode: tool_choice forwarding test ---------------------------------- - #[test] fn encode_tool_choice_forwarded() { let mut core = make_core_request(vec![CoreMessage { @@ -1321,8 +1191,6 @@ mod tests { assert_eq!(tc["type"], "auto"); } - // -- Additional missing tests ---------------------------------------------- - #[test] fn encode_tool_use_as_function_call() { let core = make_core_request(vec![CoreMessage { @@ -1412,7 +1280,6 @@ mod tests { all_events.extend(decoder.decode_frame(frame).unwrap()); } - // Stop reason should be ToolUse since we saw a function_call. let msg_stop = all_events .iter() .find(|e| matches!(e, CoreEvent::MessageStop { .. })); @@ -1439,7 +1306,6 @@ mod tests { let adapter = ResponsesAdapter; let mut decoder = adapter.new_stream_decoder(&target); - // Simulate receiving a tool call event but no explicit done/completed. let frames = vec![ make_frame(r#"{"type":"response.created","id":"resp_1"}"#), make_frame( @@ -1519,8 +1385,6 @@ mod tests { } } - // -- Additional tests: upstream model alias, malformed decode, etc -- - #[test] fn encode_uses_upstream_model() { let mut target = make_target(); @@ -1545,7 +1409,6 @@ mod tests { let target = make_target(); let proxy_req = adapter.encode_request(&core, &target).unwrap(); let body: serde_json::Value = serde_json::from_slice(&proxy_req.body).unwrap(); - // Should produce a valid request with empty input array. assert_eq!(body["input"].as_array().unwrap().len(), 0); } @@ -1630,8 +1493,6 @@ mod tests { assert_eq!(core_resp.stop_reason, StopReason::MaxTokens); } - // -- Source guard --------------------------------------------------------- - #[test] fn adapter_source_no_forbidden_imports() { let source = include_str!("responses.rs"); @@ -1649,8 +1510,6 @@ mod tests { ); } - // -- Helpers -------------------------------------------------------------- - fn make_frame(data: &str) -> SseFrame { SseFrame { event: None, diff --git a/crates/llm-proxy-provider/src/discovery.rs b/crates/llm-proxy-provider/src/discovery.rs index 7f7f952..da3b519 100644 --- a/crates/llm-proxy-provider/src/discovery.rs +++ b/crates/llm-proxy-provider/src/discovery.rs @@ -16,12 +16,6 @@ use crate::ProviderError; const ANTHROPIC_VERSION_HEADER: &str = "anthropic-version"; const DEFAULT_ANTHROPIC_VERSION: &str = "2023-06-01"; -/// Maximum number of bytes streamed from a discovery error body before it is -/// passed to [`ProviderError::api`]. -/// -/// Streamed (not buffered-then-truncated) so an adversarial discovery endpoint -/// returning a large error body cannot spike transient memory (audit -/// discovery-error-body-full-read). const MAX_DISCOVERY_ERROR_BODY_BYTES: usize = 2048; /// GET-capable HTTP client for provider model discovery. @@ -31,11 +25,8 @@ pub struct DiscoveryClient { } impl DiscoveryClient { - /// Build a discovery client with short, independent timeouts. - /// - /// # Errors - /// - /// Returns an error if the HTTP client cannot be constructed. + /// Build a discovery client with short, independent timeouts; errors if the + /// HTTP client cannot be constructed. pub fn try_new() -> Result { let http = Client::builder() .connect_timeout(Duration::from_secs(10)) @@ -48,13 +39,6 @@ impl DiscoveryClient { if !matches!(next.scheme(), "http" | "https") { return attempt.stop(); } - // SSRF / credential-disclosure guard: reqwest strips - // `Authorization` on cross-host redirects, but it does NOT strip - // custom auth headers such as `x-api-key` / `x-goog-api-key` - // (the credentials used by Anthropic and Gemini discovery). - // Refuse to follow a redirect whose host/port/scheme differs from - // the original request's so a compromised or MITM discovery - // endpoint cannot replay the provider API key to an attacker host. if let Some(origin) = attempt.previous().first() { let cross_host = next.host_str() != origin.host_str() || next.port_or_known_default() != origin.port_or_known_default() @@ -69,12 +53,8 @@ impl DiscoveryClient { Ok(Self { http }) } - /// Fetch and normalize the configured provider model list. - /// - /// # Errors - /// - /// Returns an error for missing discovery config, invalid responses, - /// pagination failures, or transport errors. + /// Fetch and normalize the configured provider model list; errors on missing + /// discovery config, invalid responses, pagination, or transport failures. pub async fn discover( &self, provider: &ProviderConfig, @@ -91,11 +71,6 @@ impl DiscoveryClient { })?; let mut models = Vec::new(); let mut dropped = 0usize; - // Track the previous pagination token so a non-advancing token (the - // upstream returns the same `nextPageToken` / `after_id` every page) is - // detected and pagination stops early with a warning, instead of - // fetching the same page up to `max_pages` times (audit - // discovery-nonadvancing-pagination-token). let mut previous_token: Option = None; for _ in 0..discovery.max_pages { @@ -111,10 +86,6 @@ impl DiscoveryClient { let Some((name, token)) = next_page(discovery.kind, &value) else { models.sort_by(|a, b| a.id.cmp(&b.id)); models.dedup_by(|a, b| a.id == b.id); - // Surface the aggregate count of records dropped across all - // pages so a silently partial catalog is observable in the - // success path (GAP-LOW-6). Logged at `info!` only when - // non-zero to avoid noise on clean responses. if dropped > 0 { tracing::info!( provider = %provider.name, @@ -126,12 +97,6 @@ impl DiscoveryClient { return Ok(models); }; - // Stop early if the upstream returned a token identical to the one - // we just used: feeding it back would re-fetch the same page in a - // tight loop until `max_pages`. Sort/dedup the (possibly partial) - // catalog and return it rather than treating this as a hard error, - // since the records gathered so far are valid (audit - // discovery-nonadvancing-pagination-token). if previous_token.as_deref() == Some(token.as_str()) { tracing::warn!( provider = %provider.name, @@ -163,12 +128,6 @@ impl DiscoveryClient { ProviderError::InvalidConfig("missing discovery configuration".to_owned()) })?; - // Bounded retry with fixed backoff for transient failures (GAP-LOW-4): - // a single 5xx or transport hiccup (TLS reset, connection reset, - // timeout) mid-pagination would otherwise abort the whole catalog - // refresh for a provider. Only server-side / transport errors are - // retried — 4xx and parse errors are surfaced immediately. The total - // extra wall-clock is bounded by `MAX_TRANSIENT_RETRIES * BACKOFF`. let mut last_error = None; for attempt in 0..=Self::MAX_TRANSIENT_RETRIES { if attempt > 0 { @@ -191,34 +150,17 @@ impl DiscoveryClient { Err(error) => return Err(error), } } - // All retries exhausted; surface the last transient error. Err(last_error.expect("retry loop ran at least once")) } - /// Upper bound on the number of retries after the first attempt for a - /// transient discovery failure. - /// - /// Exposed (`pub`) so callers and tests can reason about the total - /// worst-case probe count for a persistently-failing endpoint, which is - /// `1 + MAX_TRANSIENT_RETRIES`. + /// Upper bound on retries after the first attempt for a transient failure; + /// worst-case probe count is `1 + MAX_TRANSIENT_RETRIES`. pub const MAX_TRANSIENT_RETRIES: u32 = 2; - /// Fixed backoff applied before each retry of a transient discovery - /// failure. Intentionally constant (no jitter) because discovery is an - /// infrequent background refresh, not a fan-out request hot path. const RETRY_BACKOFF: Duration = Duration::from_millis(500); - /// Classify a [`ProviderError`] as transient (worth retrying). - /// - /// Returns `true` for transport-level failures (`Http`, including TLS - /// resets, connection resets, and timeouts) and `Api` responses with a - /// server-side 5xx status code. Client-side 4xx errors, parse failures, - /// and configuration errors are non-transient. fn is_transient(error: &ProviderError) -> bool { match error { - // Redirect-policy violations (e.g. too many redirects, or a - // cross-host redirect refused by the custom policy) are - // deterministic: retrying just re-runs the same redirect chain. ProviderError::Http { redirect: true, .. } => false, ProviderError::Http { .. } => true, ProviderError::Api { status, .. } => (500..600).contains(status), @@ -233,8 +175,6 @@ impl DiscoveryClient { url: &Url, ) -> RequestBuilder { let mut request = self.http.get(url.clone()); - // Transport boundary (discovery bypasses `AuthHeaders`): the secret is - // exposed only here to build the reqwest auth header(s). let api_key: &str = provider.api_key.expose_secret(); request = match provider.auth_style { AuthStyle::Bearer => request.bearer_auth(api_key), @@ -256,8 +196,6 @@ impl DiscoveryClient { request } - /// Execute a single discovery request attempt: send, validate status, and - /// stream the body up to `max_response_bytes`, then parse as JSON. async fn send_once( &self, request: RequestBuilder, @@ -266,10 +204,6 @@ impl DiscoveryClient { let response = request.send().await?; let status = response.status(); if !status.is_success() { - // Stream only a small prefix of the error body for diagnostics, - // mirroring the success path's bounded read, so a misbehaving - // discovery endpoint returning a large error body cannot spike - // transient memory (audit discovery-error-body-full-read). let body_prefix = crate::error::read_error_body_bounded(response, MAX_DISCOVERY_ERROR_BODY_BYTES) .await; @@ -309,9 +243,6 @@ fn parse_models( "discovery response has an invalid top-level model list".to_owned(), )); }; - // Count of records dropped this page (no id / empty id, or a Gemini model - // lacking generateContent support). Surfaced as an aggregate by `discover` - // in the success path so a silent partial catalog is observable (GAP-LOW-6). let mut dropped = 0usize; for record in records { let id = record @@ -360,19 +291,6 @@ fn parse_models( } fn set_query_parameter(url: &mut Url, name: &str, value: &str) { - // Single-pass query rebuild (audit GAP-LOW-5). - // - // Snapshot the existing query string once (a single allocation of the whole - // query), clear it, then re-append every pair except the one named `name`, - // and finally append the new value. `form_urlencoded::parse` yields - // `Cow` borrowed from the snapshot, so no per-pair `into_owned()` is - // needed -- the prior implementation allocated two owned `String`s per pair - // (one for the key, one for the value) solely to detach them from `url`'s - // read borrow before the mutable `query_pairs_mut` borrow began. Parsing the - // snapshot sidesteps that borrow conflict with a single allocation. - // - // All existing pairs keyed `name` are dropped and exactly one new pair is - // appended, matching the prior behavior. let snapshot = url.query().unwrap_or("").to_owned(); url.set_query(None); let mut serializer = url.query_pairs_mut(); @@ -419,21 +337,8 @@ fn next_page(kind: ProviderDiscoveryKind, value: &Value) -> Option<(&'static str } impl Default for DiscoveryClient { - /// Returns a default discovery client. - /// - /// # Lint expectation - /// - /// The `expect` below is deliberate: this is the infallible-in-practice - /// convenience constructor (no custom TLS backend), with [`Self::try_new`] - /// as the fallible alternative -- mirroring the convention established for - /// `ProxyClient::new` in `transport.rs` (audit LOW-2). `#[expect]` makes the - /// deliberate use compile-time-checked. - /// - /// # Panics - /// - /// Panics if the underlying `reqwest::Client` cannot be constructed (e.g. - /// due to a TLS backend initialization failure). Use [`DiscoveryClient::try_new`] - /// for a fallible constructor. + /// Returns a default discovery client; panics if the HTTP client cannot be + /// constructed — use [`DiscoveryClient::try_new`] for the fallible form. #[expect( clippy::expect_used, reason = "Default is infallible in practice; try_new is the fallible constructor" @@ -488,6 +393,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::XApiKey, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -620,10 +526,6 @@ mod tests { #[tokio::test] async fn discovery_stops_on_non_advancing_pagination_token() { - // A Gemini-style endpoint that always returns the SAME nextPageToken - // must be detected: pagination stops early after the second identical - // token instead of fetching the same page up to `max_pages` times - // (audit discovery-nonadvancing-pagination-token). async fn static_token_models(State(calls): State>) -> Json { calls.fetch_add(1, Ordering::SeqCst); Json(json!({ @@ -653,14 +555,11 @@ mod tests { .await .expect("non-advancing token should stop early, not error"); - // Exactly two pages fetched: the first yields the token, the second - // returns the identical token and trips the non-advancing guard. assert_eq!( calls.load(Ordering::SeqCst), 2, "pagination should stop after detecting the repeated token" ); - // The catalog gathered so far is valid and deduplicated. assert_eq!(models.len(), 1); assert_eq!(models[0].id, "models/gemini-pro"); } @@ -682,8 +581,6 @@ mod tests { .collect::>(), ["a", "z"] ); - // GAP-LOW-6: the aggregate dropped-count must be surfaced from the - // success path; the malformed `{"missing": "id"}` record is counted. assert_eq!(dropped, 1); } @@ -724,10 +621,6 @@ mod tests { ); } - // -- discovery auth: secret is used on the wire AND redacted from Debug -- - - /// Build a discovery provider with an explicit api key and auth style, so - /// tests can prove each auth style places the exposed secret on the wire. fn discovery_provider_with_auth( endpoint: String, auth_style: AuthStyle, @@ -738,6 +631,7 @@ mod tests { api_key: secrecy::SecretString::from(api_key), auth_style, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -754,10 +648,6 @@ mod tests { } } - /// Test-server handler that echoes the received auth header value back as - /// the single model id, so the test can assert which credential reached the - /// wire. Configured via `EchoAuthConfig` state (the header name to read and - /// an optional prefix to strip, e.g. the `Bearer ` scheme prefix). #[derive(Clone)] struct EchoAuthConfig { header_name: String, @@ -779,11 +669,6 @@ mod tests { Json(json!({ "data": [{ "id": value }] })) } - /// The discovery auth path exposes the secret ONLY to build the reqwest - /// auth header(s); each auth style must place the configured key on the - /// wire (proving the secret IS used), while the ProviderConfig Debug output - /// must never contain the raw key (proving redaction). Parameterized over - /// the three auth styles the discovery builder supports. async fn run_discovery_auth_on_the_wire_and_redacted( auth_style: AuthStyle, header_name: &'static str, @@ -803,9 +688,6 @@ mod tests { .await; let provider = discovery_provider_with_auth(endpoint, auth_style, SECRET); - // (a) The exposed secret must reach the wire: the echo handler returns - // the received auth-header value as the model id, which must equal the - // configured key. let models = DiscoveryClient::default() .discover(&provider) .await @@ -820,7 +702,6 @@ mod tests { "the configured api key must be sent on the wire as the auth header" ); - // (b) The ProviderConfig Debug output must NOT leak the raw key. let debug = format!("{provider:?}"); assert!( !debug.contains(SECRET), @@ -859,8 +740,6 @@ mod tests { #[tokio::test] async fn discovery_both_auth_uses_secret_on_wire_and_redacts_debug() { - // AuthStyle::Both sets Bearer + x-api-key; assert the Bearer side lands - // on the wire (the x-api-key side is covered by the XApiKey test). run_discovery_auth_on_the_wire_and_redacted( AuthStyle::Both, "authorization", diff --git a/crates/llm-proxy-provider/src/error.rs b/crates/llm-proxy-provider/src/error.rs index 9bbad23..9fe7324 100644 --- a/crates/llm-proxy-provider/src/error.rs +++ b/crates/llm-proxy-provider/src/error.rs @@ -1,15 +1,7 @@ //! Provider error types shared across transport, adapters, and routes. -//! -//! Extracted so that the transport layer and provider adapters can depend on -//! [`ProviderError`] without coupling to any specific client implementation. -// --------------------------------------------------------------------------- -// Provider error -// --------------------------------------------------------------------------- - -/// Errors produced by provider transport and adapter operations. -/// -/// This enum is non-exhaustive; downstream code should include a catch-all arm. +/// Errors produced by provider transport and adapter operations; the enum is +/// non-exhaustive, so downstream match arms must include a catch-all. #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum ProviderError { @@ -17,29 +9,21 @@ pub enum ProviderError { #[error("failed to marshal request: {0}")] Serialize(#[from] serde_json::Error), - /// The HTTP request failed at the transport level. - /// - /// The message is stored without the request URL so query credentials - /// cannot leak through logs, CLI output, or downstream responses. + /// The HTTP request failed at the transport level; the message is stored + /// without the request URL so query credentials cannot leak through logs. #[error("request failed: {message}")] Http { /// URL-free transport error text. message: String, /// Whether reqwest classified the failure as a timeout. timeout: bool, - /// Whether reqwest classified the failure as a redirect-policy - /// violation (e.g. too many redirects, or a cross-host redirect refused - /// by a custom policy). Redirect-policy failures are deterministic, so - /// [`DiscoveryClient::is_transient`] must not retry them. + /// Whether reqwest classified the failure as a redirect-policy violation; + /// deterministic, so [`DiscoveryClient::is_transient`] must not retry it. redirect: bool, }, - /// The upstream API returned an error status code. - /// - /// The `body` field is truncated to the internal maximum error body length - /// (512 bytes) and stripped of common API key patterns at construction time - /// so that `Display` output (used in `warn!()` / `error!()` logging) never - /// contains full key material. + /// The upstream API returned an error status code; `body` is truncated and + /// stripped of API-key patterns so `Display` output never leaks key material. #[error("API error {status}: {body}")] Api { /// HTTP status code. @@ -52,10 +36,8 @@ pub enum ProviderError { #[error("invalid UTF-8 in stream: {0}")] Utf8(#[from] std::str::Utf8Error), - /// SSE framing error (malformed frame structure, unexpected stream termination). - /// - /// Constructed by [`SseFramer`](crate::sse::SseFramer) when it detects - /// malformed SSE structure such as unterminated frames after stream end. + /// SSE framing error (malformed frame structure, unexpected stream + /// termination), constructed by [`SseFramer`](crate::sse::SseFramer). #[error("SSE framing error: {0}")] SseFraming(String), @@ -80,17 +62,6 @@ impl From for ProviderError { } } -/// Read at most `cap` bytes of an upstream error body via a byte stream. -/// -/// Rather than calling `Response::text()`/`bytes()` (which buffer the entire -/// body) and truncating after the fact, this streams `bytes_stream()` and stops -/// as soon as `cap` bytes have accumulated. A read error is rendered as a -/// placeholder so the caller still gets a usable error body (audit -/// transport-unbounded-error-body-read / discovery-error-body-full-read). -/// -/// Shared by the request transport ([`crate::transport`]) and the discovery -/// client ([`crate::discovery`]); the cap is larger than the final sanitization -/// truncation so the most diagnostic prefix is retained for logging. pub(crate) async fn read_error_body_bounded(resp: reqwest::Response, cap: usize) -> String { use futures::TryStreamExt; let mut stream = resp.bytes_stream(); @@ -102,8 +73,6 @@ pub(crate) async fn read_error_body_bounded(resp: reqwest::Response, cap: usize) if remaining == 0 { break; } - // Multi-byte safe: only append whole bytes up to the cap; never - // slice a &str at a raw byte offset. let take = bytes.len().min(remaining); buf.extend_from_slice(&bytes[..take]); if buf.len() >= cap { @@ -117,134 +86,49 @@ pub(crate) async fn read_error_body_bounded(resp: reqwest::Response, cap: usize) String::from_utf8_lossy(&buf).into_owned() } -/// Maximum length for upstream API error bodies stored in [`ProviderError::Api`]. -/// -/// After truncation the string is at most `MAX_API_ERROR_BODY_LEN` bytes long -/// (the `...[truncated]` suffix is included within this budget). pub(crate) const MAX_API_ERROR_BODY_LEN: usize = 512; -/// Suffix appended when a body exceeds the truncation limit. const TRUNCATED_SUFFIX: &str = "...[truncated]"; -/// Truncate a string to `max_len` bytes, appending `suffix` if truncation occurs. -/// -/// Delegates to the shared implementation in `llm_proxy_protocol::util`. fn truncate_with_suffix(s: &str, max_len: usize, suffix: &str) -> String { llm_proxy_protocol::util::truncate_with_suffix(s, max_len, suffix) } -/// Redaction patterns compiled once via `LazyLock`. -/// -/// Each pattern matches a known API-key prefix followed by enough token -/// characters to be a real key (10–20+). This avoids false positives on short -/// substrings like `sk-` that appear in ordinary words (e.g. "desk-area"). static REDACTION_PATTERNS: std::sync::LazyLock> = std::sync::LazyLock::new(|| { - // Character class covering the URL/JSON-safe punctuation real tokens - // carry: `.` (dotted OpenAI restricted keys, JWT segment separators), - // `-`, `_`, `/`, `+`, `=` (base64). Delimiters — quotes, braces, - // brackets, commas, whitespace — are deliberately excluded so a - // trailing delimiter adjacent to the token is never consumed into the - // redaction. Prior to GAP-MED-1 this class stopped at `.`/`/`/`+`/`=`, - // so a dotted key or JWT was only redacted up to its first dot and the - // secret-bearing suffix leaked to the client. const C: &str = r"A-Za-z0-9_.\-/+="; - // Order matters only for intent-documentation: every pattern runs a - // full `replace_all` pass over the body, and `***` is inert, so the - // final output is the union of all matches regardless of order. - // Longer/more-specific patterns are listed first. [ - // OpenAI restricted (dotted) keys: sk-proj-XXXX.YYYY format!(r"sk-proj-[{C}]{{10,}}"), - // Anthropic keys: sk-ant-api03-XXXXX format!(r"sk-ant-api03-[{C}]{{10,}}"), - // Anthropic keys: sk-ant-XXXXX format!(r"sk-ant-[{C}]{{10,}}"), - // OpenAI keys: sk-live-XXXXX (hyphen form) format!(r"sk-live-[{C}]{{10,}}"), - // OpenAI keys: sk-test-XXXXX (hyphen form) format!(r"sk-test-[{C}]{{10,}}"), - // OpenAI keys: sk_live_XXXXX (underscore form) format!(r"sk_live_[{C}]{{10,}}"), - // OpenAI keys: sk_test_XXXXX (underscore form) format!(r"sk_test_[{C}]{{10,}}"), - // Generic sk- prefix with enough trailing chars to look like a key format!(r"sk-[{C}]{{20,}}"), - // Google API keys: AIza followed by 30+ token chars format!(r"AIza[{C}]{{30,}}"), - // Generic key- prefix with enough trailing chars format!(r"key-[{C}]{{20,}}"), - // JWT-shaped Bearer tokens: Bearer eyJ....body.sig (lowered bar so - // even a partially-echoed JWT fragment is caught). format!(r"Bearer eyJ[{C}]{{10,}}"), - // Bearer token values echoed in error responses format!(r"Bearer [{C}]{{20,}}"), - // Generic key= assignment patterns (key=VALUE with 20+ chars) format!(r"(?i)key=[{C}]{{20,}}"), - // Generic token= assignment patterns (token=VALUE with 20+ chars) format!(r"(?i)token=[{C}]{{20,}}"), ] .iter() - // SAFETY: every pattern is built from a static string literal (via - // format! with a compile-time constant char class), so the resulting - // regex is known valid and `Regex::new` cannot fail here. .map(|pat| regex::Regex::new(pat).expect("invalid redaction regex")) .collect() }); -/// Sanitize an upstream API error body: strip common key patterns and -/// truncate to [`MAX_API_ERROR_BODY_LEN`]. -/// -/// # Ordering rationale -/// -/// Redaction runs BEFORE truncation. This means if a key pattern appears in -/// the portion that would be truncated away, the pattern survives because -/// truncation happens after redaction. This is acceptable because: -/// 1. The truncation limit (512 bytes) is large enough to cover the vast -/// majority of real error messages. -/// 2. The redaction patterns target key *prefixes* (e.g. `sk-ant-`) which -/// typically appear near the start of error bodies, not near the end. -/// 3. Even if a partial key survives in the truncated tail, the key is -/// incomplete and unlikely to be the full secret. -/// -/// # Covered patterns -/// -/// - OpenAI keys: `sk-live-...`, `sk-test-...`, `sk_live_...`, `sk_test_...`, -/// `sk-proj-...` (dotted restricted keys), generic `sk-...` (only when -/// followed by 20+ token chars). -/// - Anthropic keys: `sk-ant-api03-...`, `sk-ant-...` -/// - Google API keys: `AIza...` (only when followed by 30+ chars) -/// - Generic key prefixes: `key-...` (only when followed by 20+ chars) -/// - Generic assignment patterns: `key=...`, `token=...` (case-insensitive, -/// only when followed by 20+ chars) -/// - Bearer token values: `Bearer ...` (only when followed by 20+ chars) and -/// JWT-shaped `Bearer eyJ....body.sig` -/// -/// The token character class includes `.`/`/`/`+`/`=` so dotted keys and JWTs -/// are redacted whole rather than only up to their first `.` (audit GAP-MED-1). -/// -/// Uses regex-based matching to avoid false-positive redaction of short -/// substrings like `sk-` or `key-` that appear in ordinary words. pub(crate) fn sanitize_api_error_body(mut body: String) -> String { - // Redact API key patterns (regex-based, avoids false positives on short substrings). for re in REDACTION_PATTERNS.iter() { body = re.replace_all(&body, "***").into_owned(); } - // Truncate if the sanitized body exceeds the limit. - // Uses the same truncate-with-suffix logic as - // `llm_proxy_server::routes::error_response::truncate_with_suffix`. - // Both implementations must stay in sync. if body.len() > MAX_API_ERROR_BODY_LEN { body = truncate_with_suffix(&body, MAX_API_ERROR_BODY_LEN, TRUNCATED_SUFFIX); } body } -// --------------------------------------------------------------------------- -// ProviderError convenience constructors -// --------------------------------------------------------------------------- - impl ProviderError { /// Return whether this is an HTTP timeout. #[must_use] @@ -253,9 +137,6 @@ impl ProviderError { } /// Construct a [`ProviderError::Api`] with automatic body sanitization. - /// - /// Encapsulates the sanitization call so callers never need to remember - /// to call the internal `sanitize_api_error_body` function manually. pub fn api(status: u16, body_text: String) -> Self { Self::Api { status, @@ -263,10 +144,8 @@ impl ProviderError { } } - /// Classify the HTTP status code of a [`ProviderError::Api`] into a - /// [`llm_proxy_protocol::core::CoreStreamErrorKind`]. - /// - /// Returns `None` for non-`Api` variants. + /// Classify the HTTP status of a [`ProviderError::Api`] into a + /// [`llm_proxy_protocol::core::CoreStreamErrorKind`]; `None` for non-`Api`. pub fn api_error_kind(&self) -> Option { match self { Self::Api { status, .. } => Some(match *status { @@ -275,9 +154,6 @@ impl ProviderError { 403 => llm_proxy_protocol::core::CoreStreamErrorKind::Permission, 429 => llm_proxy_protocol::core::CoreStreamErrorKind::RateLimit, 500 | 502 | 503 => llm_proxy_protocol::core::CoreStreamErrorKind::Upstream, - // All remaining status codes (including 4xx not listed above) - // are classified as Upstream errors. This is intentional since - // ProviderError::Api is only constructed for status >= 400. _ => llm_proxy_protocol::core::CoreStreamErrorKind::Upstream, }), _ => None, @@ -285,18 +161,10 @@ impl ProviderError { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; - // ----------------------------------------------------------------------- - // sanitize_api_error_body tests - // ----------------------------------------------------------------------- - #[test] fn sanitize_body_under_limit_is_unchanged() { let body = "simple error message".to_owned(); @@ -313,12 +181,10 @@ mod tests { #[test] fn sanitize_body_truncation_exact_length() { - // Exactly at the limit should NOT be truncated. let body = "a".repeat(MAX_API_ERROR_BODY_LEN); let result = sanitize_api_error_body(body.clone()); assert_eq!(result, body); - // One over the limit should be truncated. let body = "a".repeat(MAX_API_ERROR_BODY_LEN + 1); let result = sanitize_api_error_body(body); assert!(result.ends_with("...[truncated]")); @@ -327,11 +193,8 @@ mod tests { #[test] fn sanitize_body_multibyte_char_at_boundary_does_not_panic() { - // Japanese characters are 3 bytes each in UTF-8. - // Build a body whose byte length straddles the truncation boundary. - let body = "あ".repeat(200); // 600 bytes, exceeds 512 + let body = "あ".repeat(200); let result = sanitize_api_error_body(body); - // Must not panic and must be valid UTF-8. assert!(result.len() <= MAX_API_ERROR_BODY_LEN); assert!(result.ends_with("...[truncated]")); } @@ -402,7 +265,6 @@ mod tests { #[test] fn sanitize_does_not_redact_short_sk_substring() { - // "desk-area" should NOT be redacted -- `sk-` only matches with 20+ trailing chars. let body = "the desk-area is reserved".to_owned(); let result = sanitize_api_error_body(body.clone()); assert_eq!(result, body); @@ -410,7 +272,6 @@ mod tests { #[test] fn sanitize_does_not_redact_short_key_substring() { - // "key-value store" should NOT be redacted. let body = "key-value store".to_owned(); let result = sanitize_api_error_body(body.clone()); assert_eq!(result, body); @@ -418,7 +279,6 @@ mod tests { #[test] fn sanitize_does_not_redact_short_aiza_substring() { - // "AIza" alone without 30+ trailing chars should not be redacted. let body = "The area is at AIza Lane".to_owned(); let result = sanitize_api_error_body(body.clone()); assert_eq!(result, body); @@ -426,8 +286,6 @@ mod tests { #[test] fn sanitize_redacts_bearer_token_in_error_body() { - // If an upstream API echoes the Authorization header value back in its - // error response, the Bearer prefix + token must be redacted. let body = r#"error: invalid key Bearer sk-ant-api03-abc123def456ghi789jkl012"#.to_owned(); let result = sanitize_api_error_body(body); assert!( @@ -480,7 +338,6 @@ mod tests { #[test] fn sanitize_does_not_redact_short_key_assignment() { - // "key=short" should NOT be redacted -- under 20 chars. let body = "error: key=shortvalue".to_owned(); let result = sanitize_api_error_body(body.clone()); assert_eq!(result, body); @@ -488,22 +345,11 @@ mod tests { #[test] fn redaction_patterns_all_compile_successfully() { - // Verify every regex in REDACTION_PATTERNS compiles without panicking. - // This is a regression guard: if a pattern is edited to be invalid, - // the LazyLock would panic at first use in production. This test - // exercises the initialization path explicitly. let _ = &*REDACTION_PATTERNS; } - // ----------------------------------------------------------------------- - // GAP-MED-1: dotted keys / JWTs must redact whole (no suffix leak) - // ----------------------------------------------------------------------- - #[test] fn sanitize_redacts_dotted_openai_proj_key_with_no_suffix_leak() { - // OpenAI restricted key echoed in a dotted format. The pre-GAP-MED-1 - // char class stopped at `.` so only `sk-proj-AbCd...` matched and the - // `.T3BlbkFJ...` suffix leaked. The whole token must now redact. let body = "invalid key: sk-proj-AbCdEfGh1234567890.T3BlbkFJabc123def456".to_owned(); let result = sanitize_api_error_body(body); assert!( @@ -515,9 +361,6 @@ mod tests { #[test] fn sanitize_redacts_jwt_bearer_token_with_no_suffix_leak() { - // A `Bearer ` is three dot-separated base64url segments. The - // pre-GAP-MED-1 class stopped at `.` so only the header segment - // matched and the payload + signature leaked. let jwt = concat!( "Bearer eyJhbGciOiJIUzI1NiJ9.", "eyJzdWIiOiIxMjM0NTY3ODkwIn0.", @@ -532,10 +375,6 @@ mod tests { #[test] fn sanitize_redacts_legacy_dotted_key_with_no_suffix_leak() { - // A legacy dotted key (no recognized prefix) is still caught by the - // broadened generic `key=` / `sk-` patterns as long as it carries a - // known prefix; here the `key=` assignment form redacts the dotted - // value whole. let body = "error: key=abCdEf.1234567890abcdefghij.KLuMnO".to_owned(); let result = sanitize_api_error_body(body); assert!( @@ -547,9 +386,6 @@ mod tests { #[test] fn sanitize_does_not_consume_trailing_delimiter_after_token() { - // A token immediately followed by a quote/brace must redact the token - // but leave the delimiter in place — the char class excludes `"` `}` - // etc. so the suffix cannot be swallowed. let body = r#"{"error":"key sk-ant-abc123def456ghi789jkl012mno345pqr"}"#.to_owned(); let result = sanitize_api_error_body(body); assert!( diff --git a/crates/llm-proxy-provider/src/lib.rs b/crates/llm-proxy-provider/src/lib.rs index 8c640e8..87c54a0 100644 --- a/crates/llm-proxy-provider/src/lib.rs +++ b/crates/llm-proxy-provider/src/lib.rs @@ -1,7 +1,5 @@ -//! llm-proxy-provider: upstream provider adapter framework. -//! -//! Manages protocol adapters for upstream LLM providers with connection pooling, -//! streaming support, and provider-specific encoding/decoding. +//! llm-proxy-provider: upstream provider adapter framework — protocol adapters, +//! connection pooling, and streaming for upstream LLM providers. #![deny(missing_docs)] diff --git a/crates/llm-proxy-provider/src/sse.rs b/crates/llm-proxy-provider/src/sse.rs index b9bd7fe..33cb0cf 100644 --- a/crates/llm-proxy-provider/src/sse.rs +++ b/crates/llm-proxy-provider/src/sse.rs @@ -1,22 +1,12 @@ -//! Buffered SSE parser for upstream provider streams. -//! -//! Provider streams can split a single SSE frame across multiple TCP chunks, -//! and some providers emit comments, event names, ids, or multi-line `data:` -//! fields. This module handles those edge cases so that route handlers and -//! adapters never need to split raw HTTP byte chunks with `str::lines()`. +//! Buffered SSE parser for upstream provider streams: reassembles frames split +//! across TCP chunks and handles comments, event/id fields, and multi-line data. use std::str; use crate::error::ProviderError; -// --------------------------------------------------------------------------- -// SseFrame -// --------------------------------------------------------------------------- - -/// A parsed SSE frame. -/// -/// Emitted only after a blank line boundary or when [`SseFramer::finish`] is -/// called. +/// A parsed SSE frame, emitted only after a blank line boundary or when +/// [`SseFramer::finish`] is called. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SseFrame { /// Optional `event:` field value. @@ -27,53 +17,20 @@ pub struct SseFrame { pub data: String, } -// --------------------------------------------------------------------------- -// SseFramer -// --------------------------------------------------------------------------- - -/// Default ceiling on the framer's internal byte buffer and on the size of a -/// single in-progress frame's accumulated `data:` payload. -/// -/// A real SSE line or frame is tiny (a JSON event is at most a few KiB). The -/// framer drains complete lines as they arrive, so `buffer` only holds bytes -/// between newlines; if a single line (no `\n`) exceeds this cap the upstream -/// is dribbling non-terminated bytes and we fail fast instead of growing the -/// buffer toward OOM. Likewise a `data:` frame that accumulates beyond this cap -/// without a blank-line boundary is rejected. See audit GAP-MED-2. -pub const DEFAULT_MAX_BUFFER_BYTES: usize = 1024 * 1024; // 1 MiB - -/// Buffered SSE frame parser. -/// -/// Call [`SseFramer::push_chunk`] with raw bytes as they arrive from the -/// upstream. Call [`SseFramer::finish`] when the stream ends to emit any -/// trailing partial frame. -/// -/// Both the raw byte buffer and an in-progress frame's `data:` payload are -/// capped at `max_buffer_bytes` (default [`DEFAULT_MAX_BUFFER_BYTES`]); an -/// upstream that would grow either beyond the cap yields -/// [`ProviderError::SseFraming`] instead of unbounded allocation (audit -/// GAP-MED-2). +/// Default ceiling on the framer's byte buffer and on a single frame's +/// accumulated `data:` payload; exceeding either yields [`ProviderError::SseFraming`]. +pub const DEFAULT_MAX_BUFFER_BYTES: usize = 1024 * 1024; + +/// Buffered SSE frame parser: feed raw bytes via [`SseFramer::push_chunk`] as +/// they arrive; call [`SseFramer::finish`] at stream end to flush a partial frame. #[derive(Debug)] pub struct SseFramer { - /// Buffered raw bytes. Stores partial UTF-8 sequences that arrived split - /// across TCP chunks so that multi-byte characters (CJK, emoji) are not - /// rejected by an intermediate `str::from_utf8` call. buffer: Vec, - /// Ceiling on `buffer` and on the accumulated `data:` frame payload. max_buffer_bytes: usize, - /// Current frame being assembled. current_event: Option, current_id: Option, current_data_lines: Vec, - /// Running byte length of the strings in `current_data_lines` (the join - /// separators are negligible against `max_buffer_bytes`). Reset whenever a - /// frame is taken. current_data_bytes: usize, - /// Whether a leading UTF-8 BOM (U+FEFF, bytes `EF BB BF`) has been stripped - /// from the start of the stream. The WHATWG SSE spec mandates stripping one - /// leading BOM; without it a first line of `\u{FEFF}data: ...` parses the - /// field as `\u{FEFF}data`, matches no known field, and is silently - /// dropped, losing the first frame (audit sse-no-bom-strip). bom_checked: bool, } @@ -84,10 +41,8 @@ impl SseFramer { Self::with_max_buffer_bytes(DEFAULT_MAX_BUFFER_BYTES) } - /// Create a new empty framer with a custom buffer/frame cap. - /// - /// `max_buffer_bytes` bounds both a single unterminated line (the raw - /// buffer) and the accumulated `data:` payload of a single frame. + /// Create a new empty framer with a custom cap bounding both a single + /// unterminated line and the accumulated `data:` payload of a single frame. pub fn with_max_buffer_bytes(max_buffer_bytes: usize) -> Self { Self { buffer: Vec::new(), @@ -100,35 +55,16 @@ impl SseFramer { } } - /// Feed a raw byte chunk and return any completed frames. - /// - /// The chunk may contain partial lines or partial UTF-8 sequences. - /// Completed frames are emitted only when a blank line boundary is - /// encountered. Returns an error if the accumulated buffer (after appending - /// the chunk) contains invalid UTF-8 that cannot be explained by a partial - /// multi-byte sequence at the tail, or if the buffer exceeds - /// `max_buffer_bytes` (a misbehaving upstream dribbling non-newline bytes). + /// Feed a raw byte chunk (possibly partial lines/UTF-8) and return frames + /// completed at a blank-line boundary; errors on invalid UTF-8 or cap overrun. pub fn push_chunk(&mut self, chunk: &[u8]) -> Result, ProviderError> { self.buffer.extend_from_slice(chunk); - // Strip a single leading UTF-8 BOM (U+FEFF, bytes `EF BB BF`) once at - // the start of the stream, per the WHATWG SSE spec. A deferred - // decision is needed when the buffer holds only a *prefix* of the BOM - // (e.g. just `EF`): we cannot yet tell whether it is the start of a BOM - // or the first byte of real content, so we wait for more bytes. As soon - // as a non-matching byte arrives (or the full BOM is present) the - // decision is final and `bom_checked` is set so this never runs again - // (audit sse-no-bom-strip). if !self.bom_checked { const BOM: &[u8; 3] = &[0xEF, 0xBB, 0xBF]; if self.buffer.starts_with(BOM) { - // Full BOM present at the very start: strip it. self.buffer.drain(..BOM.len()); } - // Decide whether the leading bytes can still grow into a BOM. If - // the buffer is a strict prefix of the BOM (and nothing more), - // defer: wait for the next chunk to disambiguate. Otherwise the - // leading byte(s) cannot be a BOM, so the check is final. let still_prefix = (1..BOM.len()) .any(|n| self.buffer.len() == n && BOM.starts_with(&self.buffer[..n])); if !still_prefix { @@ -146,20 +82,14 @@ impl SseFramer { self.drain_buffer() } - /// Flush any remaining buffered data as a final frame. - /// - /// Call this when the upstream stream ends. If there is partial data in the - /// buffer, it is emitted as a frame. + /// Flush remaining buffered data as a final frame; call when the upstream + /// stream ends, emitting any trailing partial frame. pub fn finish(&mut self) -> Result, ProviderError> { let mut frames = Vec::new(); if !self.buffer.is_empty() { - // Process any complete lines first (up to the last newline). - // This avoids UTF-8 decode failures from partial multi-byte - // sequences that may appear after the last newline. self.flush_complete_lines(&mut frames)?; - // Handle any remaining bytes after the last newline (no trailing newline). if !self.buffer.is_empty() { let text = str::from_utf8(&self.buffer).map_err(ProviderError::from)?; let line = text.trim_end_matches('\r').to_owned(); @@ -174,22 +104,12 @@ impl SseFramer { Ok(frames) } - // ----------------------------------------------------------------------- - // Internal - // ----------------------------------------------------------------------- - - /// Drains the buffer through its last newline, emitting any frames it - /// completes. - /// - /// Stopping at a newline boundary keeps a multi-byte UTF-8 character that - /// straddles two chunks intact. fn flush_complete_lines(&mut self, frames: &mut Vec) -> Result<(), ProviderError> { let Some(last_nl) = self.buffer.iter().rposition(|&b| b == b'\n') else { return Ok(()); }; let complete = self.buffer[..last_nl + 1].to_vec(); self.buffer.drain(..last_nl + 1); - // Complete lines ending at a newline boundary are valid UTF-8. let text = str::from_utf8(&complete).map_err(ProviderError::from)?; for line in text.lines() { let trimmed = line.trim_end_matches('\r'); @@ -202,38 +122,22 @@ impl SseFramer { Ok(()) } - /// Parse as many complete lines as possible from the buffer. - /// - /// Handles the case where a multi-byte UTF-8 character is split across - /// two TCP chunks: we only decode up to the last newline, keeping any - /// trailing bytes (which may be an incomplete UTF-8 sequence) in the - /// buffer for the next chunk. fn drain_buffer(&mut self) -> Result, ProviderError> { let mut frames = Vec::new(); while let Some(nl_pos) = self.buffer.iter().position(|&b| b == b'\n') { - // Find the last non-\r byte in the line (before the newline). - // This lets us trim trailing \r bytes using a slice reference - // instead of allocating a separate Vec. let trimmed_end = match self.buffer[..nl_pos].iter().rposition(|&b| b != b'\r') { - Some(pos) => pos + 1, // exclusive end - None => 0, // Line was all \r characters. + Some(pos) => pos + 1, + None => 0, }; - // Decode the trimmed line as UTF-8 and convert to an owned string - // so we release the immutable borrow on self.buffer before draining. - // The line is guaranteed to be valid UTF-8 because it ends at a \n - // boundary which is a single byte; any partial multi-byte UTF-8 - // sequence would not contain \n. let line = str::from_utf8(&self.buffer[..trimmed_end]) .map_err(ProviderError::from)? .to_owned(); - // Remove the line + newline from the buffer. self.buffer.drain(..nl_pos + 1); if line.is_empty() { - // Blank line = frame boundary. if let Some(frame) = self.take_current_frame() { frames.push(frame); } @@ -245,17 +149,13 @@ impl SseFramer { Ok(frames) } - /// Process a single non-blank SSE line. fn process_line(&mut self, line: &str) -> Result<(), ProviderError> { - // Ignore comments. if line.starts_with(':') { return Ok(()); } - // Parse field:value or field: value (space after colon is optional). if let Some(colon_pos) = line.find(':') { let field = &line[..colon_pos]; - // Skip the optional single space after the colon. let value = &line[colon_pos + 1..]; let value = value.strip_prefix(' ').unwrap_or(value); @@ -269,12 +169,9 @@ impl SseFramer { "data" => { self.push_data(value)?; } - // Ignore unknown fields per spec. _ => {} } } else { - // A field with no colon and no value (e.g. just "data" with no ":"). - // Per SSE spec, this is a field with an empty value. match line { "event" => { self.current_event = Some(String::new()); @@ -291,9 +188,6 @@ impl SseFramer { Ok(()) } - /// Append a `data:` value to the in-progress frame, enforcing the frame - /// cap so a misbehaving upstream emitting an unbounded multi-line `data:` - /// frame (no blank-line boundary) cannot grow the frame without bound. fn push_data(&mut self, value: &str) -> Result<(), ProviderError> { let next = self.current_data_bytes.saturating_add(value.len()); if next > self.max_buffer_bytes { @@ -309,7 +203,6 @@ impl SseFramer { Ok(()) } - /// Take the current accumulated frame if it has any data. fn take_current_frame(&mut self) -> Option { let data_lines = std::mem::take(&mut self.current_data_lines); self.current_data_bytes = 0; @@ -332,10 +225,6 @@ impl Default for SseFramer { } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -344,7 +233,6 @@ mod tests { fn sse_framer_handles_partial_frames_split_across_chunks() { let mut framer = SseFramer::new(); - // First chunk has a partial frame. let frames = framer.push_chunk(b"data: hel").unwrap(); assert!(frames.is_empty(), "No complete frame yet"); @@ -401,11 +289,9 @@ mod tests { fn sse_framer_emits_trailing_frame_from_finish() { let mut framer = SseFramer::new(); - // Push data without a trailing blank line. let frames = framer.push_chunk(b"data: partial").unwrap(); assert!(frames.is_empty()); - // Finish should emit the trailing frame. let frames = framer.finish().unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "partial"); @@ -439,13 +325,11 @@ mod tests { fn sse_framer_rejects_invalid_utf8() { let mut framer = SseFramer::new(); - // Invalid UTF-8 bytes buffered (no newline, so drain_buffer finds nothing). let result = framer.push_chunk(&[0xFF, 0xFE]); assert!( result.is_ok(), "push_chunk should buffer raw bytes without error" ); - // The error surfaces when finish() tries to decode the buffer as UTF-8. let result = framer.finish(); assert!( result.is_err(), @@ -461,8 +345,6 @@ mod tests { fn sse_framer_rejects_invalid_utf8_at_newline() { let mut framer = SseFramer::new(); - // Invalid UTF-8 followed by a newline: drain_buffer will try to decode - // the line and fail. let result = framer.push_chunk(&[0xFF, 0xFE, b'\n']); assert!( result.is_err(), @@ -523,16 +405,11 @@ mod tests { assert_eq!(frames[0].event.as_deref(), Some("msg1")); assert_eq!(frames[0].id.as_deref(), Some("1")); assert_eq!(frames[0].data, "hello"); - // Second frame has no event or id. assert!(frames[1].event.is_none()); assert!(frames[1].id.is_none()); assert_eq!(frames[1].data, "world"); } - // ----------------------------------------------------------------------- - // Edge case tests - // ----------------------------------------------------------------------- - #[test] fn sse_framer_empty_input_returns_no_frames() { let mut framer = SseFramer::new(); @@ -596,12 +473,10 @@ mod tests { fn sse_framer_finish_after_complete_stream_is_empty() { let mut framer = SseFramer::new(); - // Push a complete frame (terminated by blank line). let frames = framer.push_chunk(b"data: hello\n\n").unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "hello"); - // Calling finish() again should return empty. let frames = framer.finish().unwrap(); assert!( frames.is_empty(), @@ -616,7 +491,6 @@ mod tests { let frames = framer.push_chunk(b"retry: 5000\ndata: hello\n\n").unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "hello"); - // The retry field is silently ignored per SSE spec. } #[test] @@ -634,32 +508,27 @@ mod tests { fn sse_framer_finish_trims_trailing_cr() { let mut framer = SseFramer::new(); - // Simulate a stream ending with 'data: hello\r' (no trailing \n). let frames = framer.push_chunk(b"data: hello\r").unwrap(); assert!(frames.is_empty()); let frames = framer.finish().unwrap(); assert_eq!(frames.len(), 1); - // The trailing \r must be trimmed, not preserved in the data. assert_eq!(frames[0].data, "hello"); } #[test] fn sse_framer_multibyte_utf8_split_across_chunks() { - // The character 'あ' is 3 bytes in UTF-8: 0xE3 0x81 0x82. let full_char = "あ"; let full_bytes = full_char.as_bytes(); assert_eq!(full_bytes.len(), 3); let mut framer = SseFramer::new(); - // First chunk: "data: " + first 2 bytes of あ. let mut chunk1 = b"data: ".to_vec(); chunk1.extend_from_slice(&full_bytes[..2]); let frames = framer.push_chunk(&chunk1).unwrap(); assert!(frames.is_empty(), "No complete frame yet"); - // Second chunk: remaining 1 byte of あ + "\n\n". let mut chunk2 = full_bytes[2..].to_vec(); chunk2.extend_from_slice(b"\n\n"); let frames = framer.push_chunk(&chunk2).unwrap(); @@ -672,9 +541,6 @@ mod tests { #[test] fn sse_framer_data_newline_then_finish() { - // "data: hello\n" (newline-terminated line but no blank line after it) - // push_chunk drains the line but doesn't emit a frame; finish() should - // emit it. let mut framer = SseFramer::new(); let frames = framer.push_chunk(b"data: hello\n").unwrap(); assert!(frames.is_empty(), "No blank line yet, no frame"); @@ -686,11 +552,8 @@ mod tests { #[test] fn sse_framer_reusable_after_finish() { - // Verify the framer is in a clean state after finish() so it can be - // reused for a second stream. let mut framer = SseFramer::new(); - // First round. let frames = framer.push_chunk(b"data: first\n\n").unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "first"); @@ -698,20 +561,13 @@ mod tests { let frames = framer.finish().unwrap(); assert!(frames.is_empty(), "Nothing left after complete stream"); - // Second round -- framer should be clean. let frames = framer.push_chunk(b"data: second\n\n").unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "second"); } - // ----------------------------------------------------------------------- - // Bounded-buffer / bounded-frame protections (audit GAP-MED-2) - // ----------------------------------------------------------------------- - #[test] fn sse_framer_rejects_unterminated_line_exceeding_cap() { - // A single unterminated line (no newline) larger than the cap must - // error instead of growing the buffer without bound. let mut framer = SseFramer::with_max_buffer_bytes(64); let big = b"x".repeat(65); let err = framer.push_chunk(&big).unwrap_err(); @@ -724,15 +580,10 @@ mod tests { #[test] fn sse_framer_rejects_oversized_data_frame_exceeding_cap() { - // Many `data:` lines without a blank-line boundary must trip the frame - // cap rather than accumulating without bound. let mut framer = SseFramer::with_max_buffer_bytes(64); - // Each line is "data: aaa..." (8 bytes of payload) repeated until the - // accumulated frame payload exceeds 64 bytes. let line = b"data: aaaaaaaa\n"; // 8 payload bytes, no blank line let mut last_ok = true; for _ in 0..16 { - // 8 * 16 = 128 > 64, so this must eventually error. let res = framer.push_chunk(line); match res { Ok(_) => last_ok = true, @@ -755,9 +606,6 @@ mod tests { #[test] fn sse_framer_resets_data_byte_counter_after_frame_boundary() { - // After a blank-line boundary the per-frame byte counter resets, so a - // stream of many small legitimate frames (each under the cap) parses - // fine even though their total exceeds the cap. let mut framer = SseFramer::with_max_buffer_bytes(32); for i in 0..50 { let chunk = format!("data: item{i}\n\n"); @@ -769,18 +617,11 @@ mod tests { #[test] fn sse_framer_default_cap_is_one_mebibyte() { - // Sanity: the public default matches the documented constant. assert_eq!(DEFAULT_MAX_BUFFER_BYTES, 1024 * 1024); } - // ----------------------------------------------------------------------- - // Leading UTF-8 BOM stripping (audit sse-no-bom-strip) - // ----------------------------------------------------------------------- - #[test] fn sse_framer_strips_single_leading_bom() { - // U+FEFF as UTF-8 is EF BB BF. A BOM-prefixed first data line must - // parse the field as `data`, not `\u{FEFF}data`. let mut framer = SseFramer::new(); let mut input = vec![0xEF, 0xBB, 0xBF]; input.extend_from_slice(b"data: hello\n\n"); @@ -791,8 +632,6 @@ mod tests { #[test] fn sse_framer_strips_bom_split_across_chunks() { - // The BOM arrives split across two chunks; the first data line must - // still parse correctly. let mut framer = SseFramer::new(); let frames = framer.push_chunk(&[0xEF, 0xBB]).unwrap(); assert!(frames.is_empty(), "no complete frame yet"); @@ -805,11 +644,8 @@ mod tests { #[test] fn sse_framer_strips_only_one_bom() { - // Only one leading BOM is stripped; a second BOM in-stream is content - // that survives into the data field (it is valid UTF-8). let mut framer = SseFramer::new(); let mut input = vec![0xEF, 0xBB, 0xBF]; - // "data: \u{FEFF}hello\n\n" — the BOM inside the data value is kept. input.extend_from_slice(b"data: "); input.extend_from_slice(&[0xEF, 0xBB, 0xBF]); input.extend_from_slice(b"hello\n\n"); @@ -820,17 +656,12 @@ mod tests { #[test] fn sse_framer_no_bom_preserves_leading_content() { - // No BOM present: the first byte is real content and must survive. let mut framer = SseFramer::new(); let frames = framer.push_chunk(b"data: hello\n\n").unwrap(); assert_eq!(frames.len(), 1); assert_eq!(frames[0].data, "hello"); } - // ----------------------------------------------------------------------- - // Source guard: sse.rs must not import forbidden types - // ----------------------------------------------------------------------- - #[test] fn sse_source_no_forbidden_imports() { let source = include_str!("sse.rs"); diff --git a/crates/llm-proxy-provider/src/transport.rs b/crates/llm-proxy-provider/src/transport.rs index 5aa29c4..6c7451b 100644 --- a/crates/llm-proxy-provider/src/transport.rs +++ b/crates/llm-proxy-provider/src/transport.rs @@ -1,14 +1,5 @@ -//! Protocol-neutral HTTP transport for upstream LLM providers. -//! -//! This module provides a thin transport layer that sends prepared bytes to a -//! prepared URL with prepared auth headers. It knows nothing about protocol -//! names, model IDs, or provider adapters. -//! -//! ## Neutrality guardrails -//! -//! The transport layer must not import or mention endpoint classification, -//! provider model names, or protocol crate types. It only sends prepared -//! bytes to a prepared URL with prepared auth headers. +//! Protocol-neutral HTTP transport: sends prepared bytes to a prepared URL +//! with prepared auth headers, knowing nothing about providers or protocols. use std::fmt; use std::pin::Pin; @@ -24,35 +15,14 @@ use tokio::time::Sleep; use crate::error::ProviderError; -/// Default time-to-first-byte (response headers) deadline for upstream requests. -/// -/// This bounds the unprotected window between connection establishment and the -/// receipt of response headers on both [`ProxyClient::send`] and -/// [`ProxyClient::send_stream`]. `connect_timeout` only covers the TCP+TLS -/// handshake, so without this an upstream that accepts the connection but never -/// returns headers could hold the call open until the downstream tower-http -/// timeout (~300 s) fires. The streaming *body* is left under the existing -/// per-chunk idle timeout ([`DEFAULT_STREAM_IDLE_TIMEOUT`]); this deadline only -/// protects the headers phase (audit transport-no-upstream-timeout). +/// Default time-to-first-byte (response headers) deadline for upstream requests; +/// bounds the headers phase without capping the streaming body. pub const DEFAULT_UPSTREAM_HEADERS_TIMEOUT: Duration = Duration::from_secs(60); -/// Maximum number of bytes read from an upstream error body before it is passed -/// to [`ProviderError::api`]. -/// -/// Streamed (not buffered-then-truncated) so a malicious or grossly misbehaving -/// upstream returning a multi-gigabyte error response cannot spike transient -/// memory (audit transport-unbounded-error-body-read). Mirrors the 2 KiB prefix -/// cap used by the discovery client. pub(crate) const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 2048; -// --------------------------------------------------------------------------- -// AuthHeaders -// --------------------------------------------------------------------------- - -/// Authentication headers for an upstream request. -/// -/// The `api_key` field is redacted in [`fmt::Debug`] output so that -/// `tracing::debug!(?auth)` or snapshot output never leaks the secret. +/// Authentication headers for an upstream request; `api_key` is redacted in +/// [`fmt::Debug`] output so logs and snapshots never leak the secret. #[derive(Clone, Debug)] #[non_exhaustive] pub struct AuthHeaders { @@ -63,16 +33,8 @@ pub struct AuthHeaders { pub api_key: SecretString, } -// --------------------------------------------------------------------------- -// ProxyRequest -// --------------------------------------------------------------------------- - -/// A protocol-neutral request to send to an upstream provider. -/// -/// The transport sends `body` byte-for-byte. It does not serialize typed -/// request structs. -/// -/// The `auth.api_key` field is redacted in [`fmt::Debug`] output. +/// A protocol-neutral request to an upstream provider; the transport sends +/// `body` byte-for-byte and does not serialize typed request structs. #[derive(Clone)] #[non_exhaustive] pub struct ProxyRequest { @@ -82,27 +44,10 @@ pub struct ProxyRequest { pub auth: AuthHeaders, /// Raw request body bytes. pub body: Vec, - /// Whether this is a streaming request (informational for logging/metrics). - /// - /// **This field is NOT read by [`ProxyClient`].** The caller must pick the - /// correct method: [`ProxyClient::send`] for non-streaming requests and - /// [`ProxyClient::send_stream`] for streaming requests. This field exists - /// solely for structured logging and future metrics. - /// - /// LOW-1 (evaluated, deferred by design): the audit's softened analysis - /// confirms the two production call sites (`core_pipeline.rs`) pick - /// `send`/`send_stream` statically from routing — there is no runtime - /// `if req.stream` dispatch, so the field cannot cause a wrong-method call. - /// Promoting the documented invariant to a compile-time one would require - /// type-state (`ProxyRequest` / ``) or an enum - /// return, but `core.stream` is a *runtime* `bool`, so pure type-state is - /// infeasible without splitting the encode path too. Rated low-impact / - /// high-churn by the audit itself; the documentation invariant is retained. + /// Informational only: NOT read by [`ProxyClient`], which the caller picks + /// (`send` vs `send_stream`); exists for structured logging and metrics. pub stream: bool, - /// Static adapter headers from config (e.g. `anthropic-version`). - /// - /// These are validated at config-load time and applied to every upstream - /// request in addition to auth headers. + /// Static adapter headers from config (e.g. `anthropic-version`), applied to every upstream request. pub extra_headers: std::collections::HashMap, } @@ -123,40 +68,30 @@ pub(crate) fn endpoint_without_query(endpoint: &str) -> &str { .map_or(endpoint, |(base, _query)| base) } -// --------------------------------------------------------------------------- -// ProxyClient -// --------------------------------------------------------------------------- - -/// Default per-chunk idle read timeout for streaming upstream responses. -/// -/// Legitimate SSE providers send heartbeats / keep-alive bytes (e.g. -/// Anthropic's `event: ping`, OpenAI's immediate first chunk) well within this -/// window, so a stream that yields no byte for longer than this is a dead or -/// wedged upstream. A stalled provider can no longer hold the connection open -/// indefinitely (audit GAP-MED-2). +/// Default per-chunk idle read timeout for streaming responses; a stream +/// silent for longer than this is a dead or wedged upstream. pub const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(60); -/// Protocol-neutral HTTP client for upstream LLM providers. -/// -/// Owns a connection-pooled [`reqwest::Client`] and sends [`ProxyRequest`] -/// instances without any knowledge of protocol names or model IDs. -#[derive(Debug, Clone)] +/// Protocol-neutral HTTP client for upstream LLM providers, owning a +/// connection-pooled [`reqwest::Client`]. +#[derive(Clone)] #[must_use = "ProxyClient does nothing until send/send_stream is called"] pub struct ProxyClient { http: reqwest::Client, - /// Per-chunk idle read timeout applied to streaming responses. If the - /// upstream yields no body byte within this window the stream errors so a - /// stalled provider cannot hold the connection open (audit GAP-MED-2). stream_idle_timeout: Duration, - /// Time-to-first-byte (response headers) deadline applied to every upstream - /// request. Bounds the unprotected window between connection establishment - /// and headers receipt so an upstream that accepts the connection but never - /// returns headers cannot hold the call open indefinitely. The streaming - /// *body* remains under [`Self::stream_idle_timeout`]; this deadline does - /// not cap a legitimate long token stream (audit transport-no-upstream-timeout). upstream_headers_timeout: Duration, } +impl fmt::Debug for ProxyClient { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ProxyClient") + .field("http", &"") + .field("stream_idle_timeout", &self.stream_idle_timeout) + .field("upstream_headers_timeout", &self.upstream_headers_timeout) + .finish() + } +} + impl Default for ProxyClient { fn default() -> Self { Self::new() @@ -164,28 +99,8 @@ impl Default for ProxyClient { } impl ProxyClient { - /// Creates a new transport client with sensible connection pool defaults. - /// - /// This is a convenience wrapper around [`Self::try_new`] that unwraps the - /// result. The current configuration (no custom TLS backend, no proxy env - /// validation at build time) makes `reqwest::Client::build()` infallible - /// in practice. If a future configuration change makes this fallible, - /// callers should switch to [`Self::try_new`]. - /// - /// # Lint expectation (audit LOW-2) - /// - /// The `expect` below is deliberate: this is the documented infallible-in- - /// practice convenience constructor, with [`Self::try_new`] offered as the - /// fallible alternative. `#[expect(clippy::expect_used)]` acknowledges the - /// intentional `.expect` at this site. `clippy::expect_used` is a - /// `clippy::restriction` lint, disabled across the workspace, so a plain - /// `#[allow(...)]` here would be inert (it suppresses a never-enabled lint). - /// `#[expect(...)]` instead makes the deliberate use compile-time-checked: - /// the expectation is satisfied because the `.expect` is present, and if the - /// `.expect` is ever removed the expectation becomes unfulfilled and fails - /// the build — surfacing the change for review instead of leaving a silent, - /// inert `#[allow]`. `#[expect]` does not emit the lint as a warning; it - /// only errors when the lint does not fire. + /// Convenience wrapper around [`Self::try_new`] for the default + /// (infallible-in-practice) configuration; `try_new` is the fallible form. #[expect( clippy::expect_used, reason = "documented infallible-in-practice constructor; try_new is the fallible alternative" @@ -194,17 +109,10 @@ impl ProxyClient { Self::try_new().expect("failed to build reqwest client with default configuration") } - /// Creates a new transport client, returning an error if the HTTP client - /// cannot be built. - /// - /// Use this when configuration may make client construction fallible - /// (e.g. custom TLS backends, proxy environment validation). + /// Creates a transport client, returning an error if the HTTP client cannot + /// be built (custom TLS backends, proxy environment validation). pub fn try_new() -> Result { - let http = reqwest::Client::builder() - .pool_max_idle_per_host(20) - .pool_idle_timeout(Duration::from_secs(90)) - .connect_timeout(Duration::from_secs(10)) - .build()?; + let http = pooled_http_builder().build()?; Ok(Self { http, @@ -213,36 +121,34 @@ impl ProxyClient { }) } - /// Override the per-chunk idle read timeout applied to streaming responses. - /// - /// Returns the client for chaining. The default is - /// [`DEFAULT_STREAM_IDLE_TIMEOUT`] (60 s); legitimate SSE providers - /// heartbeat well within that window. Raise it only if a specific upstream - /// is known to pause longer between bytes (audit GAP-MED-2). + /// Creates a client routed through a forward proxy with the same pool and + /// timeouts as [`Self::try_new`]; `proxy_url` must be socks5h/socks5/http(s) with a host. + pub fn try_with_egress_proxy(proxy_url: &str) -> Result { + validate_egress_proxy_url(proxy_url)?; + let proxy = reqwest::Proxy::all(proxy_url)?; + let http = pooled_http_builder().proxy(proxy).build()?; + + Ok(Self { + http, + stream_idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + upstream_headers_timeout: DEFAULT_UPSTREAM_HEADERS_TIMEOUT, + }) + } + + /// Override the per-chunk idle read timeout for streaming responses; + /// the default is [`DEFAULT_STREAM_IDLE_TIMEOUT`]. pub fn with_stream_idle_timeout(mut self, idle: Duration) -> Self { self.stream_idle_timeout = idle; self } - /// Override the time-to-first-byte (response headers) deadline for upstream - /// requests. - /// - /// Returns the client for chaining. The default is - /// [`DEFAULT_UPSTREAM_HEADERS_TIMEOUT`] (60 s). This bounds the headers - /// receipt phase without capping the streaming body, which remains under - /// the per-chunk idle timeout (audit transport-no-upstream-timeout). + /// Override the time-to-first-byte deadline for upstream requests; the + /// default is [`DEFAULT_UPSTREAM_HEADERS_TIMEOUT`]. Does not cap the body. pub fn with_upstream_headers_timeout(mut self, deadline: Duration) -> Self { self.upstream_headers_timeout = deadline; self } - /// Build the common `reqwest` request skeleton shared by [`send`](Self::send) - /// and [`send_stream`](Self::send_stream). - /// - /// Sets `Content-Type: application/json`, conditionally sets - /// `Accept: text/event-stream`, applies auth, and applies any extra headers. - /// The caller attaches the body (`.body(req.body)`) afterwards so ownership of - /// `req.body` stays with the caller. fn build_request( &self, req: &ProxyRequest, @@ -263,20 +169,10 @@ impl ProxyClient { } /// Sends a non-streaming request and returns the response body bytes. - /// - /// - Always sets `Content-Type: application/json`. - /// - Does **not** set `Accept: text/event-stream`. - /// - Returns [`ProviderError::Api`] for HTTP status >= 400. + /// Sets `Content-Type: application/json` but not `Accept: text/event-stream`. pub async fn send(&self, req: ProxyRequest) -> Result, ProviderError> { let builder = self.build_request(&req, false); - // Bound the headers-receipt (time-to-first-byte) phase so an upstream - // that accepts the connection but never returns headers cannot hold the - // call open until the downstream tower-http timeout. The body read in - // `check_status` remains under the underlying read deadline; a global - // client `.timeout()` is intentionally avoided so legitimate long - // (non-streaming) generations are not killed (audit - // transport-no-upstream-timeout). let resp = tokio::time::timeout(self.upstream_headers_timeout, builder.body(req.body).send()) .await @@ -294,22 +190,8 @@ impl ProxyClient { check_status(resp).await } - /// Sends a streaming request and returns a byte stream. - /// - /// # Cancel safety - /// - /// The returned stream owns the `reqwest::Response` via `bytes_stream()`. - /// Dropping the stream drops the underlying connection, aborting the - /// upstream request. No detached task is spawned that outlives the - /// consumer. - /// - /// - Always sets `Content-Type: application/json`. - /// - Sets `Accept: text/event-stream`. - /// - Returns [`ProviderError::Api`] for HTTP status >= 400 **before** - /// any byte stream is exposed. - /// - Stream item errors are wrapped as [`ProviderError`]; consumers never - /// see raw `reqwest::Error`. - /// - Dropping the returned stream aborts the in-flight upstream request. + /// Sends a streaming request; returns [`ProviderError::Api`] for status >= 400 + /// before any bytes are exposed. Dropping the returned stream aborts the request. pub async fn send_stream( &self, req: ProxyRequest, @@ -319,11 +201,6 @@ impl ProxyClient { > { let builder = self.build_request(&req, true); - // Bound only the headers-receipt (time-to-first-byte) phase. The - // streaming *body* is left under the per-chunk idle timeout applied - // below via `IdleTimeoutStream`; a global client `.timeout()` would - // wrongly kill legitimate long token streams (audit - // transport-no-upstream-timeout). let resp = tokio::time::timeout(self.upstream_headers_timeout, builder.body(req.body).send()) .await @@ -340,25 +217,13 @@ impl ProxyClient { if resp.status().as_u16() >= 400 { let status = resp.status().as_u16(); - // Stream the error body with a small cap rather than buffering it - // in full via `.text()`, so a misbehaving upstream returning a huge - // error response cannot spike transient memory before - // `ProviderError::api` truncates it (audit - // transport-unbounded-error-body-read). let body_text = crate::error::read_error_body_bounded(resp, MAX_UPSTREAM_ERROR_BODY_BYTES).await; return Err(ProviderError::api(status, body_text)); } - // `bytes_stream()` consumes the Response and returns an owned stream; - // dropping the stream drops the underlying connection, aborting the - // upstream request. Do not spawn a detached task that outlives the - // consumer. let stream = resp.bytes_stream(); let mapped = stream.map_err(ProviderError::from); - // Wrap with a per-chunk idle read timeout so a stalled upstream (one - // that opens the connection but then sends no bytes) cannot hold the - // proxy's connection open indefinitely (audit GAP-MED-2). let inner: Pin> + Send>> = Box::pin(mapped); Ok(Box::pin(IdleTimeoutStream::new( @@ -368,28 +233,8 @@ impl ProxyClient { } } -// --------------------------------------------------------------------------- -// IdleTimeoutStream -// --------------------------------------------------------------------------- - -/// A byte-stream wrapper enforcing a per-chunk idle read timeout. -/// -/// If the inner stream does not yield an item within `idle` of the previous one -/// (or of subscription, for the first chunk), the next poll returns -/// [`ProviderError::Http`] with `timeout: true` so the consumer can close the -/// connection. This prevents a stalled upstream — one that opens the -/// connection then sends no bytes — from holding the proxy's connection open -/// forever (audit GAP-MED-2). -/// -/// An *idle* (not *overall*) timeout is used deliberately: a legitimate long -/// LLM token stream can run for minutes and must not be killed merely for -/// being long, only for going silent. Legitimate SSE providers heartbeat -/// within the idle window. struct IdleTimeoutStream { inner: Pin> + Send>>, - /// Timer armed for `idle` after the most recent chunk (or at construction). - /// `Sleep` is `!Unpin`, so it is boxed; `Pin>` is `Unpin`, keeping - /// the whole wrapper `Unpin` and the manual `Stream` impl safe. sleep: Pin>, idle: Duration, } @@ -411,59 +256,68 @@ impl Stream for IdleTimeoutStream { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Every field is a `Pin>` or `Duration` (all `Unpin`), so the - // wrapper itself is `Unpin` and `Pin<&mut Self>::get_mut` is sound. let this = self.get_mut(); match this.inner.as_mut().poll_next(cx) { Poll::Ready(Some(item)) => { - // Restart the idle window for the next chunk. Re-arm the existing - // `Sleep` in place to avoid a heap allocation and timer - // register/deregister per chunk on the streaming path. this.sleep .as_mut() .reset(tokio::time::Instant::now() + this.idle); Poll::Ready(Some(item)) } Poll::Ready(None) => Poll::Ready(None), - Poll::Pending => { - // No chunk yet — is the idle window elapsed? - match this.sleep.as_mut().poll(cx) { - Poll::Ready(()) => Poll::Ready(Some(Err(ProviderError::Http { - message: format!( - "upstream stream stalled: no data within {}s \ + Poll::Pending => match this.sleep.as_mut().poll(cx) { + Poll::Ready(()) => Poll::Ready(Some(Err(ProviderError::Http { + message: format!( + "upstream stream stalled: no data within {}s \ (stream idle timeout); aborting to avoid an \ open-ended stall (audit GAP-MED-2)", - this.idle.as_secs() - ), - timeout: true, - redirect: false, - }))), - Poll::Pending => Poll::Pending, - } - } + this.idle.as_secs() + ), + timeout: true, + redirect: false, + }))), + Poll::Pending => Poll::Pending, + }, } } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Apply authentication headers to a request builder based on [`AuthStyle`]. -/// -/// # Security note on `AuthStyle::Both` -/// -/// When `AuthStyle::Both` is used, the API key is sent in **both** the -/// `Authorization: Bearer` and `x-api-key` headers simultaneously. This -/// doubles the exposure surface of the secret in the wire data. This mode -/// exists for compatibility with providers that accept either header (e.g., -/// Anthropic). Where possible, prefer `AuthStyle::Bearer` or -/// `AuthStyle::XApiKey` to send the key in only one header. +fn pooled_http_builder() -> reqwest::ClientBuilder { + reqwest::Client::builder() + .pool_max_idle_per_host(20) + .pool_idle_timeout(Duration::from_secs(90)) + .connect_timeout(Duration::from_secs(10)) +} + +const EGRESS_PROXY_SCHEMES: [&str; 4] = ["socks5h", "socks5", "http", "https"]; + +fn validate_egress_proxy_url(proxy_url: &str) -> Result<(), ProviderError> { + let url = proxy_url.trim(); + let (scheme, rest) = url.split_once("://").ok_or_else(|| { + ProviderError::InvalidConfig(format!( + "egress proxy URL has no scheme (expected socks5h://, socks5://, http://, or https://): {url:?}" + )) + })?; + if !EGRESS_PROXY_SCHEMES + .iter() + .any(|accepted| scheme.eq_ignore_ascii_case(accepted)) + { + return Err(ProviderError::InvalidConfig(format!( + "egress proxy scheme must be socks5h://, socks5://, http://, or https://, got {scheme}://" + ))); + } + let authority = rest.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + if host.is_empty() { + return Err(ProviderError::InvalidConfig(format!( + "egress proxy URL has no host: {url:?}" + ))); + } + Ok(()) +} + fn apply_auth(mut builder: reqwest::RequestBuilder, auth: &AuthHeaders) -> reqwest::RequestBuilder { - // Transport boundary: the secret is exposed here solely to build the - // reqwest header value. The borrowed `&str` never escapes this function - // except into the header value. let api_key: &str = auth.api_key.expose_secret(); match auth.style { AuthStyle::Bearer => { @@ -483,14 +337,9 @@ fn apply_auth(mut builder: reqwest::RequestBuilder, auth: &AuthHeaders) -> reqwe builder } -/// Check response status and return body bytes or a [`ProviderError::Api`]. async fn check_status(resp: reqwest::Response) -> Result, ProviderError> { if resp.status().as_u16() >= 400 { let status = resp.status().as_u16(); - // Stream the error body with a small cap rather than buffering it in - // full via `.text()`, so a misbehaving upstream returning a huge error - // response cannot spike transient memory before `ProviderError::api` - // truncates it (audit transport-unbounded-error-body-read). let body_text = crate::error::read_error_body_bounded(resp, MAX_UPSTREAM_ERROR_BODY_BYTES).await; return Err(ProviderError::api(status, body_text)); @@ -499,18 +348,10 @@ async fn check_status(resp: reqwest::Response) -> Result, ProviderError> Ok(body.to_vec()) } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; - // ----------------------------------------------------------------------- - // Debug redaction tests - // ----------------------------------------------------------------------- - #[test] fn auth_headers_debug_redacts_api_key() { let auth = AuthHeaders { @@ -556,10 +397,6 @@ mod tests { assert!(!debug_output.contains("query-secret")); } - // ----------------------------------------------------------------------- - // Transport integration tests (axum test server) - // ----------------------------------------------------------------------- - use axum::Router; use axum::body::Body; use axum::extract::State; @@ -569,11 +406,9 @@ mod tests { use futures::StreamExt; use tokio::net::TcpListener; - /// A simple axum handler that echoes back the request body and headers. async fn echo_handler(headers: HeaderMap, body: bytes::Bytes) -> impl IntoResponse { let mut response_parts = Vec::new(); - // Echo all headers as "name: value" pairs. for (name, value) in &headers { response_parts.push(format!( "{}: {}", @@ -588,11 +423,9 @@ mod tests { (StatusCode::OK, response_parts.join("\n")) } - /// Handler that echoes request headers as a streaming SSE response. async fn echo_stream_handler(headers: HeaderMap, body: bytes::Bytes) -> impl IntoResponse { let mut events = Vec::new(); - // Echo Content-Type if let Some(ct) = headers.get("content-type") { events.push(format!( "data: {{\"content-type\": \"{}\"}}\n\n", @@ -600,7 +433,6 @@ mod tests { )); } - // Echo Authorization if let Some(auth) = headers.get("authorization") { events.push(format!( "data: {{\"authorization\": \"{}\"}}\n\n", @@ -608,7 +440,6 @@ mod tests { )); } - // Echo x-api-key if let Some(key) = headers.get("x-api-key") { events.push(format!( "data: {{\"x-api-key\": \"{}\"}}\n\n", @@ -616,7 +447,6 @@ mod tests { )); } - // Echo Accept if let Some(accept) = headers.get("accept") { events.push(format!( "data: {{\"accept\": \"{}\"}}\n\n", @@ -624,11 +454,9 @@ mod tests { )); } - // Echo body let body_str = String::from_utf8_lossy(&body); events.push(format!("data: {{\"body\": \"{}\"}}\n\n", body_str)); - // Terminal frame events.push("data: [DONE]\n\n".to_owned()); let stream = futures::stream::iter(events); @@ -641,7 +469,6 @@ mod tests { ) } - /// Handler that returns a 400 error with a JSON body. async fn error_handler(State(code): State) -> impl IntoResponse { let body = format!(r#"{{"error": "upstream error", "status": {}}}"#, code); ( @@ -650,12 +477,6 @@ mod tests { ) } - /// Wait for a test server to be ready by polling its TCP port. - /// - /// Replaces the previous fixed `tokio::time::sleep(Duration::from_millis(50))` - /// with a deterministic readiness check that retries until the server - /// accepts a connection or the deadline elapses (audit LOW-19). This is the - /// same pattern used in `tests/chat_completions.rs` and `tests/core_pipeline.rs`. async fn wait_for_ready(addr: std::net::SocketAddr) { let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { @@ -669,7 +490,6 @@ mod tests { } } - /// Spin up an ephemeral axum server, return its base URL once it is ready. async fn start_test_server(routes: Router) -> String { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -679,8 +499,6 @@ mod tests { axum::serve(listener, routes).await.unwrap(); }); - // Wait deterministically for the server to accept connections instead - // of a fixed sleep (audit LOW-19). wait_for_ready(addr).await; base @@ -803,7 +621,6 @@ mod tests { "Response must contain x-goog-api-key header: {}", text ); - // Must NOT contain Bearer or x-api-key. assert!( !text.contains("authorization:"), "XGoogleApiKey must not set Authorization header: {}", @@ -871,8 +688,6 @@ mod tests { extra_headers: extra, }; - // The streaming request should succeed (status < 400) because the - // extra headers are applied before the body is sent. let result = client.send_stream(req).await; assert!( result.is_ok(), @@ -908,7 +723,6 @@ mod tests { #[tokio::test] async fn non_json_body_sent_byte_for_byte() { - // Verify that arbitrary non-JSON bytes pass through unchanged. let app = Router::new().route("/test", post(echo_handler)); let base = start_test_server(app).await; @@ -927,17 +741,12 @@ mod tests { let resp = client.send(req).await.unwrap(); let resp_bytes = resp; - // The echo handler uses String::from_utf8_lossy, so check that the - // lossy-converted body contains the expected bytes. let text = String::from_utf8_lossy(&resp_bytes); - // The raw bytes 0x00, 0x01, 0x02 will appear literally; 0xFF, 0xFE - // become the replacement char. Verify the body was received. assert!( text.contains("body:"), "Response must echo body field: {}", text ); - // Verify the exact raw bytes were received by checking the first few. let body_prefix = b"body: \x00\x01\x02"; assert!( resp_bytes @@ -1030,7 +839,6 @@ mod tests { let mut stream = client.send_stream(req).await.unwrap(); - // Collect all events and look for the Accept header echo. let mut found_accept = false; while let Some(item) = stream.next().await { let chunk = item.unwrap(); @@ -1099,10 +907,6 @@ mod tests { ); } - /// Streams data slowly so a test can drop the response mid-stream. - /// - /// Each chunk sent bumps the shared counter, so the test can observe that - /// the upstream stops producing once the consumer goes away. async fn slow_handler( axum::extract::State(counter): axum::extract::State< std::sync::Arc, @@ -1161,16 +965,9 @@ mod tests { { let mut stream = client.send_stream(req).await.unwrap(); - // Read one item to confirm stream is working. let _first = stream.next().await; - // Stream is dropped here when it goes out of scope. } - // Poll the counter until it stabilizes (two consecutive identical - // samples) so the assertion sees the post-abort steady state instead - // of racing the handler's last in-flight chunk. Replaces the previous - // fixed 200 ms sleep + timing-dependent `count < 50` assertion with a - // deterministic wait and a deterministic property (audit LOW-19). let mut count = counter.load(Ordering::SeqCst); let poll_deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { @@ -1181,18 +978,10 @@ mod tests { } count = next; if tokio::time::Instant::now() >= poll_deadline { - // Still incrementing after 5 s — the abort did not propagate; - // fall through to the assertion which will then fail loudly. break; } } - // The deterministic property: aborting the stream must stop the - // handler short of its 100-chunk limit. The handler stops emitting at - // `i >= 100`, so reaching 100 means the stream ran to completion and - // the drop did NOT abort upstream. Any value < 100 proves the abort - // took effect (the exact count depends on scheduling and is not - // itself the property under test). assert!( count < 100, "Dropping the stream should abort upstream before it runs to \ @@ -1201,8 +990,6 @@ mod tests { ); } - /// Handler that yields one chunk then stalls (sends no further byte and - /// keeps the connection open) to exercise the per-chunk idle timeout. async fn stalling_handler() -> axum::response::Response { let stream = futures::stream::unfold(false, |sent| async move { if !sent { @@ -1211,7 +998,6 @@ mod tests { true, )) } else { - // Stall far longer than any reasonable idle timeout. tokio::time::sleep(Duration::from_secs(120)).await; None } @@ -1243,12 +1029,9 @@ mod tests { let mut stream = client.send_stream(req).await.unwrap(); - // First chunk arrives promptly (well within the idle window). let first = stream.next().await.expect("first chunk should arrive"); assert!(first.is_ok(), "first chunk should be Ok: {:?}", first); - // Bound the test so a regression (timeout never fires) fails fast - // instead of hanging the suite. let second = tokio::time::timeout(Duration::from_secs(5), stream.next()).await; match second { Ok(Some(Err(e))) => { @@ -1264,15 +1047,9 @@ mod tests { } } - // ----------------------------------------------------------------------- - // Source guard checks - // ----------------------------------------------------------------------- - - /// Handler that sends partial data then drops the connection (upstream disconnect). async fn partial_disconnect_handler() -> axum::response::Response { let stream = futures::stream::unfold(0u32, |i| async move { if i >= 2 { - // Stop sending -- simulates upstream dropping the connection. None } else { let chunk = format!("data: chunk {}\n\n", i); @@ -1313,12 +1090,10 @@ mod tests { match item { Ok(bytes) => chunks.push(bytes), Err(e) => { - // Mid-stream errors are also acceptable. panic!("unexpected mid-stream error: {:?}", e); } } } - // Stream should have ended with the two chunks the server sent. assert!( !chunks.is_empty(), "Expected at least one chunk before upstream disconnect" @@ -1329,12 +1104,7 @@ mod tests { assert!(all_text.contains("chunk 1"), "Should have received chunk 1"); } - /// Handler that sends some data then returns an error body mid-stream. async fn mid_stream_error_handler() -> axum::response::Response { - // We simulate a mid-stream error by sending two successful chunks - // then ending the stream. The consumer should see the data. - // (A true mid-stream TCP error is hard to simulate with axum; - // this test verifies the stream terminates correctly.) let stream = futures::stream::iter(vec![ Ok::<_, std::convert::Infallible>(bytes::Bytes::from("data: first\n\n")), Ok::<_, std::convert::Infallible>(bytes::Bytes::from("data: second\n\n")), @@ -1391,7 +1161,7 @@ mod tests { let err = client.send(req).await.unwrap_err(); match err { - ProviderError::Http { .. } => {} // expected + ProviderError::Http { .. } => {} other => panic!( "expected ProviderError::Http for empty URL, got: {:?}", other @@ -1401,16 +1171,12 @@ mod tests { #[tokio::test] async fn large_body_sent_successfully() { - // Use a body size that exceeds axum's default 2 MB limit would reject, - // but is still reasonable. We configure the test server with a larger - // body limit to verify the transport does not buffer or truncate. let app = axum::Router::new() .route("/test", post(echo_handler)) .layer(axum::extract::DefaultBodyLimit::max(10 * 1024 * 1024)); let base = start_test_server(app).await; let client = ProxyClient::new(); - // 5 MB body let large_body = "x".repeat(5 * 1024 * 1024); let req = ProxyRequest { url: format!("{}/test", base), @@ -1431,16 +1197,8 @@ mod tests { ); } - // ----------------------------------------------------------------------- - // Source guard checks (original) - // ----------------------------------------------------------------------- - - /// Return only the production (non-test) portion of the source file. - /// The guard tests check for forbidden identifiers in production code, - /// not in the test assertions that mention those identifiers by name. fn prod_source() -> &'static str { let source = include_str!("transport.rs"); - // Split at the test module boundary. source .split_once("#[cfg(test)]") .map(|(prod, _)| prod) @@ -1498,9 +1256,211 @@ mod tests { ); } - // ----------------------------------------------------------------------- - // Edge case tests - // ----------------------------------------------------------------------- + #[test] + fn egress_proxy_client_builds_for_each_supported_scheme() { + for url in [ + "socks5h://tunnel.example.com:1080", + "socks5://127.0.0.1:1080", + "http://proxy.internal:3128", + "https://proxy.internal:8443", + "SOCKS5H://tunnel.example.com:1080", + "http://user:pass@proxy.internal:3128", + "socks5h://tunnel.example.com", + ] { + let client = ProxyClient::try_with_egress_proxy(url) + .unwrap_or_else(|e| panic!("egress proxy {url} must build: {e}")); + assert_eq!( + client.stream_idle_timeout, DEFAULT_STREAM_IDLE_TIMEOUT, + "egress client must keep the default stream idle timeout" + ); + assert_eq!( + client.upstream_headers_timeout, DEFAULT_UPSTREAM_HEADERS_TIMEOUT, + "egress client must keep the default headers timeout" + ); + } + } + + #[test] + fn egress_proxy_client_rejects_unsupported_or_malformed_urls() { + for url in [ + "ftp://proxy:21", + "socks4://127.0.0.1:1080", + "127.0.0.1:1080", + "socks5h://", + "", + ] { + let err = ProxyClient::try_with_egress_proxy(url) + .expect_err(&format!("egress proxy {url:?} must be rejected")); + assert!( + matches!(err, ProviderError::InvalidConfig(_)), + "egress proxy {url:?} must be InvalidConfig, got {err:?}" + ); + } + } + + async fn absolute_target_echo_handler(uri: axum::http::Uri) -> String { + uri.to_string() + } + + #[test] + fn egress_proxy_client_debug_does_not_print_the_proxy() { + let client = ProxyClient::try_with_egress_proxy("socks5h://debug-secret-host.invalid:1080") + .expect("builds"); + let debug = format!("{client:?}"); + assert!( + !debug.contains("debug-secret-host.invalid"), + "ProxyClient Debug must not print the egress proxy: {debug}" + ); + } + + #[tokio::test] + async fn egress_proxy_client_sends_requests_through_the_proxy() { + let app = Router::new().route("/v1/test", post(absolute_target_echo_handler)); + let proxy_base = start_test_server(app).await; + + let client = + ProxyClient::try_with_egress_proxy(&proxy_base).expect("http egress proxy builds"); + let req = ProxyRequest { + url: "http://egress-upstream.invalid/v1/test".to_owned(), + auth: AuthHeaders { + style: AuthStyle::Bearer, + api_key: SecretString::from("key"), + }, + body: br#"{"hello":"world"}"#.to_vec(), + stream: false, + extra_headers: std::collections::HashMap::new(), + }; + + let resp = client.send(req).await.expect("request via proxy succeeds"); + let text = String::from_utf8_lossy(&resp); + assert!( + text.contains("http://egress-upstream.invalid/v1/test"), + "the proxy must receive the absolute-form upstream target, got: {text}" + ); + } + + async fn read_exact(stream: &mut tokio::net::TcpStream, buf: &mut [u8]) -> std::io::Result<()> { + use tokio::io::AsyncReadExt; + let mut filled = 0; + while filled < buf.len() { + let n = stream.read(&mut buf[filled..]).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "socks test client closed early", + )); + } + filled += n; + } + Ok(()) + } + + async fn serve_one_socks5_connect( + listener: tokio::net::TcpListener, + upstream: std::net::SocketAddr, + destination_tx: tokio::sync::oneshot::Sender, + ) { + use tokio::io::AsyncWriteExt; + let (mut sock, _peer) = listener.accept().await.expect("socks client connects"); + + let mut head = [0u8; 2]; + read_exact(&mut sock, &mut head).await.expect("greeting"); + let mut methods = vec![0u8; head[1] as usize]; + read_exact(&mut sock, &mut methods).await.expect("methods"); + sock.write_all(&[0x05, 0x00]).await.expect("auth reply"); + + let mut req_head = [0u8; 4]; + read_exact(&mut sock, &mut req_head) + .await + .expect("request head"); + assert_eq!(req_head[1], 0x01, "expected CONNECT"); + let addr_len = match req_head[3] { + 0x01 => 4, + 0x04 => 16, + 0x03 => { + let mut len = [0u8; 1]; + read_exact(&mut sock, &mut len) + .await + .expect("domain length"); + len[0] as usize + } + other => panic!("unknown SOCKS address type {other}"), + }; + let mut addr_bytes = vec![0u8; addr_len]; + read_exact(&mut sock, &mut addr_bytes) + .await + .expect("address"); + let mut port = [0u8; 2]; + read_exact(&mut sock, &mut port).await.expect("port"); + let destination = if req_head[3] == 0x03 { + String::from_utf8(addr_bytes).expect("domain is UTF-8") + } else { + format!("{}-byte address", addr_bytes.len()) + }; + + let up_port = upstream.port().to_be_bytes(); + sock.write_all(&[0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, up_port[0], up_port[1]]) + .await + .expect("connect reply"); + + destination_tx + .send(destination) + .expect("test holds the receiver"); + + let up = tokio::net::TcpStream::connect(upstream) + .await + .expect("connect to upstream"); + let (mut client_read, mut client_write) = sock.into_split(); + let (mut up_read, mut up_write) = up.into_split(); + let a = tokio::spawn(async move { tokio::io::copy(&mut client_read, &mut up_write).await }); + let b = tokio::spawn(async move { tokio::io::copy(&mut up_read, &mut client_write).await }); + let (_, _) = tokio::join!(a, b); + } + + #[tokio::test] + async fn socks5h_egress_delivers_the_hostname_to_the_proxy_tunnel() { + let app = Router::new().route("/socks-test", post(echo_handler)); + let upstream_base = start_test_server(app).await; + let upstream_addr: std::net::SocketAddr = upstream_base + .trim_start_matches("http://") + .parse() + .expect("upstream addr"); + + let socks_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let socks_addr = socks_listener.local_addr().unwrap(); + let (destination_tx, destination_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(serve_one_socks5_connect( + socks_listener, + upstream_addr, + destination_tx, + )); + + let client = ProxyClient::try_with_egress_proxy(&format!("socks5h://{socks_addr}")) + .expect("socks5h egress proxy builds"); + let req = ProxyRequest { + url: "http://egress-upstream.invalid/socks-test".to_owned(), + auth: AuthHeaders { + style: AuthStyle::Bearer, + api_key: SecretString::from("key"), + }, + body: br#"{"via":"socks"}"#.to_vec(), + stream: false, + extra_headers: std::collections::HashMap::new(), + }; + + let resp = client.send(req).await.expect("request via socks tunnel"); + let text = String::from_utf8_lossy(&resp); + assert!( + text.contains(r#""via":"socks""#), + "the echo upstream must see the body through the tunnel, got: {text}" + ); + + let destination = destination_rx.await.expect("destination"); + assert_eq!( + destination, "egress-upstream.invalid", + "socks5h must deliver the unresolved hostname to the proxy" + ); + } #[tokio::test] async fn empty_body_sent_successfully() { @@ -1530,11 +1490,8 @@ mod tests { #[tokio::test] async fn unreachable_url_returns_http_error() { - // With connect_timeout(10s) on ProxyClient, this should fail within - // ~10 seconds rather than waiting for the OS TCP timeout (120+ s). let client = ProxyClient::new(); let req = ProxyRequest { - // Non-routable address guarantees a connection failure. url: "http://192.0.2.1:1/test".to_owned(), auth: AuthHeaders { style: AuthStyle::Bearer, @@ -1545,7 +1502,6 @@ mod tests { extra_headers: std::collections::HashMap::new(), }; - // Bound the test to 15 seconds as a safety net. let result = tokio::time::timeout(Duration::from_secs(15), client.send(req)).await; let err = match result { Ok(Err(e)) => e, @@ -1553,7 +1509,7 @@ mod tests { Err(_) => panic!("test timed out waiting for connection failure"), }; match err { - ProviderError::Http { .. } => {} // expected + ProviderError::Http { .. } => {} other => panic!( "expected ProviderError::Http for unreachable URL, got: {:?}", other diff --git a/crates/llm-proxy-provider/tests/fixture_tests.rs b/crates/llm-proxy-provider/tests/fixture_tests.rs index 6e6c3ad..af41d23 100644 --- a/crates/llm-proxy-provider/tests/fixture_tests.rs +++ b/crates/llm-proxy-provider/tests/fixture_tests.rs @@ -1,9 +1,5 @@ -//! Fixture-driven integration tests for provider protocol adapters. -//! -//! Loads each fixture's input JSON, runs encode/decode against the adapter, -//! and compares the result with the expected output JSON. For stream fixtures, -//! parses SSE input, feeds through the decoder, and compares events with -//! core-events.json. +//! Fixture-driven integration tests: run adapter encode/decode against JSON +//! and SSE fixtures and compare with the expected outputs. #[cfg(test)] mod tests { @@ -21,26 +17,13 @@ mod tests { use std::fs; use std::path::Path; - // --------------------------------------------------------------------------- - // Constants - // --------------------------------------------------------------------------- - const FIXTURE_ROOT: &str = "tests/fixtures"; - // --------------------------------------------------------------------------- - // Adapter names - // --------------------------------------------------------------------------- - const OPENAI_CHAT: &str = "openai_chat"; const ANTHROPIC: &str = "anthropic"; const RESPONSES: &str = "responses"; const GEMINI: &str = "gemini"; - // --------------------------------------------------------------------------- - // Helpers - // --------------------------------------------------------------------------- - - /// Read a fixture file as a parsed JSON value. fn read_fixture(dir: &Path, name: &str) -> serde_json::Value { let path = dir.join(name); let content = fs::read_to_string(&path) @@ -49,14 +32,12 @@ mod tests { .unwrap_or_else(|e| panic!("failed to parse {} as JSON: {e}", path.display())) } - /// Read a fixture file as raw string. fn read_fixture_raw(dir: &Path, name: &str) -> String { let path = dir.join(name); fs::read_to_string(&path) .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display())) } - /// Build a default ProviderAdapterTarget for the given protocol. fn make_target(protocol: ProviderProtocol) -> ProviderAdapterTarget { let (endpoint, auth_style, model): (String, AuthStyle, String) = match protocol { ProviderProtocol::OpenAiChatCompletions => ( @@ -99,7 +80,6 @@ mod tests { } } - /// Get the ProviderAdapter for the given name. fn get_adapter(name: &str) -> ProviderAdapter { match name { OPENAI_CHAT => ProviderAdapter::OpenAiChat(OpenAiChatAdapter), @@ -110,7 +90,6 @@ mod tests { } } - /// Get the protocol for the given adapter name. fn get_protocol(name: &str) -> ProviderProtocol { match name { OPENAI_CHAT => ProviderProtocol::OpenAiChatCompletions, @@ -121,18 +100,11 @@ mod tests { } } - /// Build a CoreRequest from the human-readable core.json fixture format. - /// - /// The fixture format uses shorthand conventions: - /// - `"model": "gpt-4o"` -> `ModelRef { requested: "gpt-4o", upstream: None }` - /// - `"tool_choice": "auto"` -> `CoreToolChoice::Auto` - /// - `"tool_choice": null` -> `None` fn build_core_request(core_json: &serde_json::Value) -> CoreRequest { let obj = core_json .as_object() .expect("core.json must be a JSON object"); - // Model let model = match obj.get("model") { Some(v) if v.is_string() => ModelRef { requested: v.as_str().unwrap().to_owned(), @@ -144,28 +116,24 @@ mod tests { _ => panic!("core.json missing 'model' field"), }; - // System let system: Vec = obj .get("system") .and_then(|v| v.as_array()) .map(|arr| arr.iter().map(parse_core_content).collect()) .unwrap_or_default(); - // Messages let messages: Vec = obj .get("messages") .and_then(|v| v.as_array()) .map(|arr| arr.iter().map(parse_core_message).collect()) .unwrap_or_default(); - // Tools let tools: Vec = obj .get("tools") .and_then(|v| v.as_array()) .map(|arr| arr.iter().map(parse_core_tool).collect()) .unwrap_or_default(); - // Tool choice let tool_choice = obj.get("tool_choice").and_then(|v| match v { serde_json::Value::Null => None, serde_json::Value::String(s) => Some(parse_tool_choice_string(s)), @@ -175,10 +143,8 @@ mod tests { _ => None, }); - // Sampling let sampling: SamplingOptions = obj.get("sampling").map(parse_sampling).unwrap_or_default(); - // Stream let stream = obj.get("stream").and_then(|v| v.as_bool()).unwrap_or(false); CoreRequest { @@ -194,7 +160,6 @@ mod tests { } } - /// Parse a CoreContent from JSON. fn parse_core_content(v: &serde_json::Value) -> CoreContent { let obj = v.as_object().expect("content must be a JSON object"); match obj.get("type").and_then(|t| t.as_str()).unwrap_or("") { @@ -290,7 +255,6 @@ mod tests { } } - /// Parse a CoreMessage from JSON. fn parse_core_message(v: &serde_json::Value) -> CoreMessage { let obj = v.as_object().expect("message must be a JSON object"); let role = match obj.get("role").and_then(|r| r.as_str()).unwrap_or("user") { @@ -311,7 +275,6 @@ mod tests { CoreMessage { role, content } } - /// Parse a CoreTool from JSON. fn parse_core_tool(v: &serde_json::Value) -> CoreTool { let obj = v.as_object().expect("tool must be a JSON object"); CoreTool { @@ -331,7 +294,6 @@ mod tests { } } - /// Parse a tool_choice string value. fn parse_tool_choice_string(s: &str) -> CoreToolChoice { match s { "auto" => CoreToolChoice::Auto, @@ -341,7 +303,6 @@ mod tests { } } - /// Parse SamplingOptions from JSON. fn parse_sampling(v: &serde_json::Value) -> SamplingOptions { let obj = v .as_object() @@ -384,7 +345,6 @@ mod tests { } } - /// Parse SSE text into a list of SseFrame values. fn parse_sse_frames(sse_text: &str) -> Vec { let mut frames = Vec::new(); let mut current_data = String::new(); @@ -393,7 +353,6 @@ mod tests { for line in sse_text.lines() { let trimmed = line.trim(); if trimmed.is_empty() { - // End of event if !current_data.is_empty() { frames.push(SseFrame { event: current_event.take(), @@ -410,10 +369,8 @@ mod tests { } else if let Some(event) = trimmed.strip_prefix("event:") { current_event = Some(event.trim_start().to_owned()); } - // Ignore other SSE fields (id:, retry:, comments) } - // Handle last event if file doesn't end with blank line if !current_data.is_empty() { frames.push(SseFrame { event: current_event, @@ -425,8 +382,6 @@ mod tests { frames } - /// Normalize a stop reason string to PascalCase for comparison. - /// Accepts both snake_case ("end_turn") and PascalCase ("EndTurn"). fn normalize_stop_reason(s: &str) -> &'static str { match s { "end_turn" | "EndTurn" => "EndTurn", @@ -440,8 +395,6 @@ mod tests { } } - /// Extract the variant key from an externally-tagged CoreContent JSON value. - /// E.g., `{"Text": {...}}` returns "Text". fn extract_content_variant(v: &serde_json::Value) -> String { v.as_object() .and_then(|obj| { @@ -466,19 +419,6 @@ mod tests { .unwrap_or_else(|| "Unknown".to_owned()) } - // --------------------------------------------------------------------------- - // Provider request encode tests - // --------------------------------------------------------------------------- - - /// Test that encoding a CoreRequest produces valid provider wire format. - /// - /// Verifies the adapter successfully encodes the request and the output - /// matches key fields from output.json. The output.json fixture serves - /// both as documentation and as a structural comparison target. - /// Flattens a wire `content` field to plain text. - /// - /// Provider formats allow either a bare string or an array of content - /// blocks; only `text` blocks contribute. fn json_text_content(v: &serde_json::Value) -> String { if let Some(text) = v.as_str() { return text.to_owned(); @@ -494,7 +434,6 @@ mod tests { .join("") } - /// Asserts one encoded message matches the fixture's expected message. fn assert_encoded_message( adapter_name: &str, case: &str, @@ -512,9 +451,6 @@ mod tests { let Some(act_content) = act_msg.get("content") else { panic!("{adapter_name}/{case}: messages[{i}] missing content field"); }; - // A string expectation accepts either a matching string or an array - // whose text blocks concatenate to the same string; anything else - // must match exactly. if exp_content.is_string() { assert_eq!( json_text_content(act_content), @@ -562,6 +498,12 @@ mod tests { } } + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's expected wire output field by field \ + (messages, tools, sampling); splitting the walk would scatter \ + the fixture's assertions across helpers" + )] fn run_encode_request_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let core_json = read_fixture(&dir, "core.json"); @@ -581,14 +523,11 @@ mod tests { panic!("encoded body is not valid JSON for {adapter_name}/{case}: {e}") }); - // The adapter must produce a JSON object. assert!( actual.is_object(), "{adapter_name}/{case}: encoded body is not a JSON object" ); - // Adapter must include the target model somewhere -- either in the request - // body or in the URL (Gemini puts the model in the URL path). let model_str = target.upstream_model.as_str(); let body_str = actual.to_string(); let has_model = body_str.contains(model_str) || proxy_req.url.contains(model_str); @@ -598,8 +537,6 @@ mod tests { proxy_req.url ); - // Compare key structural fields from output.json against the encoded body. - // Messages/system/tools/sampling fields are compared when present in output.json. if let Some(expected_msgs) = output_json.get("messages").and_then(|v| v.as_array()) { let actual_msgs = actual.get("messages").and_then(|v| v.as_array()); let actual_msgs = actual_msgs.map(|a| a.as_slice()).unwrap_or(&[]); @@ -615,7 +552,6 @@ mod tests { } } - // Compare system prompt content if present in output.json. if let Some(expected_sys) = output_json.get("system") { assert_eq!( actual.get("system"), @@ -624,15 +560,6 @@ mod tests { ); } - // Compare model field if present in output.json. - // Verify that the encoded body or URL contains a model name. The exact - // value comes from the test target's upstream_model, which may differ - // from the fixture's expected model, so we only check presence. - if output_json.get("model").and_then(|v| v.as_str()).is_some() { - // Model presence already verified by the has_model assertion above. - } - - // Compare tools count and key fields if present in output.json. if let Some(expected_tools) = output_json.get("tools").and_then(|v| v.as_array()) { let actual_tools = actual .get("tools") @@ -646,13 +573,11 @@ mod tests { ); for (i, exp_tool) in expected_tools.iter().enumerate() { let act_tool = &actual_tools[i]; - // Compare tool name assert_eq!( act_tool.get("name"), exp_tool.get("name"), "{adapter_name}/{case}: tools[{i}].name mismatch" ); - // Compare tool description if present in expected if exp_tool.get("description").is_some() { assert_eq!( act_tool.get("description"), @@ -660,7 +585,6 @@ mod tests { "{adapter_name}/{case}: tools[{i}].description mismatch" ); } - // Compare tool input_schema if present in expected if exp_tool.get("input_schema").is_some() || exp_tool.get("parameters").is_some() { let act_schema = act_tool .get("input_schema") @@ -676,13 +600,9 @@ mod tests { } } - // Compare sampling fields if present in output.json. - // Adapters omit fields when they are null (serde skip_serializing_if), - // so we treat `None` (absent) and `Some(Null)` as equivalent. for field in &["temperature", "top_p", "max_tokens"] { if let Some(expected_val) = output_json.get(*field) { let actual_val = actual.get(*field); - // If expected is null, accept either absent or null in actual. if expected_val.is_null() { assert!( actual_val.is_none_or(|v| v.is_null()), @@ -691,7 +611,6 @@ mod tests { } else if let Some(av) = actual_val { assert_eq!(av, expected_val, "{adapter_name}/{case}: {field} mismatch"); } else { - // expected is non-null but actual is absent panic!( "{adapter_name}/{case}: {field} expected {expected_val} but field is absent in encoded body" ); @@ -700,11 +619,12 @@ mod tests { } } - // --------------------------------------------------------------------------- - // Provider response decode tests - // --------------------------------------------------------------------------- - - /// Test that decoding a provider response produces the expected output.json fields. + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's decoded response field by field (content \ + blocks, usage, stop reason); splitting the walk would scatter \ + the fixture's assertions across helpers" + )] fn run_decode_response_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let input_json = read_fixture(&dir, "input.json"); @@ -725,8 +645,6 @@ mod tests { panic!("failed to serialize CoreResponse for {adapter_name}/{case}: {e}") }); - // Compare stop_reason (normalize to PascalCase for comparison since - // CoreResponse serializes with default serde). let actual_stop = actual .get("stop_reason") .cloned() @@ -735,7 +653,6 @@ mod tests { .get("stop_reason") .cloned() .unwrap_or(serde_json::Value::Null); - // Normalize both to PascalCase for comparison let actual_stop_str = normalize_stop_reason(actual_stop.as_str().unwrap_or("")); let expected_stop_str = normalize_stop_reason(expected_stop.as_str().unwrap_or("")); assert_eq!( @@ -743,7 +660,6 @@ mod tests { "{adapter_name}/{case}: stop_reason mismatch (actual: {actual_stop}, expected: {expected_stop})" ); - // Compare id if present in output.json if let Some(expected_id) = output_json.get("id") { assert_eq!( actual.get("id"), @@ -752,19 +668,12 @@ mod tests { ); } - // Compare model if present in output.json. - // Adapters may override the model with the target's upstream_model, so - // we only verify that the model field exists and has a non-empty - // requested value. Direct string comparison is skipped because the - // decoded model depends on ProviderAdapterTarget configuration. if let Some(expected_str) = output_json.get("model").and_then(|m| m.as_str()) && let Some(am_req) = actual .get("model") .and_then(|am| am.get("requested")) .and_then(|v| v.as_str()) { - // Verify both are non-empty; the actual model may be overridden by - // the adapter's target configuration. assert!( !expected_str.is_empty() || am_req.is_empty(), "{adapter_name}/{case}: expected model is empty but actual is '{am_req}'" @@ -775,7 +684,6 @@ mod tests { ); } - // Compare stop_sequence if present in output.json if let Some(expected_ss) = output_json.get("stop_sequence") { assert_eq!( actual.get("stop_sequence"), @@ -784,15 +692,6 @@ mod tests { ); } - // Compare content blocks structurally. - // CoreContent serializes as externally-tagged enum, e.g.: - // {"Text": {"text": "...", "cache": null}} - // {"ToolUse": {"id": "...", "name": "...", "input": {...}}} - // {"Thinking": {"text": "...", "signature": null}} - // The fixture output.json uses a more human-readable format: - // {"type": "text", "text": "..."} - // {"type": "tool_use", "id": "...", "name": "...", "input": {...}} - // So we extract the actual variant from the serde key and compare. let actual_content = actual .get("content") .and_then(|v| v.as_array()) @@ -815,11 +714,8 @@ mod tests { .zip(expected_content.iter()) .enumerate() { - // Actual: externally-tagged serde, extract the key (e.g. "Text", "ToolUse") let actual_key = extract_content_variant(ac); - // Expected: human-readable "type" field (e.g. "text", "tool_use") let expected_type = ec.get("type").and_then(|v| v.as_str()).unwrap_or("unknown"); - // Normalize both to PascalCase for comparison let actual_normalized = match actual_key.as_str() { "Text" => "text", "ToolUse" => "tool_use", @@ -837,12 +733,10 @@ mod tests { actual_normalized, expected_type, "{adapter_name}/{case}: content[{i}] type mismatch" ); - // Get the inner object from the actual serde format let inner = ac .get(&actual_key) .cloned() .unwrap_or(serde_json::Value::Object(Default::default())); - // For text blocks, compare text if expected_type == "text" { assert_eq!( inner.get("text"), @@ -850,14 +744,12 @@ mod tests { "{adapter_name}/{case}: content[{i}] text mismatch" ); } - // For thinking blocks, compare text and optionally signature if expected_type == "thinking" { assert_eq!( inner.get("text"), ec.get("text"), "{adapter_name}/{case}: content[{i}] thinking text mismatch" ); - // Compare signature if present in expected output if ec.get("signature").is_some() { assert_eq!( inner.get("signature"), @@ -866,7 +758,6 @@ mod tests { ); } } - // For tool_use blocks, compare id, name, and input if expected_type == "tool_use" { assert_eq!( inner.get("name"), @@ -878,7 +769,6 @@ mod tests { ec.get("input"), "{adapter_name}/{case}: content[{i}] tool_use input mismatch" ); - // Compare id if present in expected output if ec.get("id").is_some() { assert_eq!( inner.get("id"), @@ -887,7 +777,6 @@ mod tests { ); } } - // For refusal blocks, compare text if expected_type == "refusal" { assert_eq!( inner.get("text"), @@ -895,7 +784,6 @@ mod tests { "{adapter_name}/{case}: content[{i}] refusal text mismatch" ); } - // For redacted_thinking blocks, compare data if present if expected_type == "redacted_thinking" && ec.get("data").is_some() { assert_eq!( inner.get("data"), @@ -905,7 +793,6 @@ mod tests { } } - // Compare usage token counts let actual_usage = actual .get("usage") .cloned() @@ -925,7 +812,6 @@ mod tests { "{adapter_name}/{case}: usage.output_tokens mismatch" ); - // Compare extended usage fields if present in output.json if expected_usage.get("reasoning_tokens").is_some() { assert_eq!( actual_usage.get("reasoning_tokens"), @@ -949,16 +835,12 @@ mod tests { } } - // --------------------------------------------------------------------------- - // Provider stream decode tests - // --------------------------------------------------------------------------- - - /// Test that decoding a provider SSE stream produces events matching core-events.json. - /// - /// The core-events.json file uses a human-readable format with a `"type"` field - /// per event (e.g., `{"type": "message_start", ...}`). The test compares: - /// 1. Event counts must match. - /// 2. Each event's variant kind (from the `"type"` field) must match. + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's expected SSE event sequence frame by \ + frame; splitting the walk would scatter the fixture's \ + assertions across helpers" + )] fn run_stream_decode_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let sse_input = read_fixture_raw(&dir, "input.sse"); @@ -979,19 +861,16 @@ mod tests { all_events.extend(events); } - // Flush remaining state let remaining = decoder .finish() .unwrap_or_else(|e| panic!("finish failed for {adapter_name}/{case}: {e}")); all_events.extend(remaining); - // Parse expected event types from the fixture file. let expected_raw: Vec = serde_json::from_str(&core_events_raw) .unwrap_or_else(|e| { panic!("failed to parse core-events.json for {adapter_name}/{case}: {e}") }); - // Compare event counts. assert_eq!( all_events.len(), expected_raw.len(), @@ -1000,7 +879,6 @@ mod tests { expected_raw.len() ); - // Compare each event's variant kind with the fixture's "type" field. for (i, (actual, expected)) in all_events.iter().zip(expected_raw.iter()).enumerate() { let actual_kind = event_variant_name(actual); let expected_type = expected @@ -1012,7 +890,6 @@ mod tests { "{adapter_name}/{case}: event[{i}] variant mismatch" ); - // Compare key event field values based on variant type. match actual { CoreEvent::TextDelta { index, text } => { if let Some(exp_text) = expected.get("text") { @@ -1095,7 +972,6 @@ mod tests { } => { if let Some(exp_sr) = expected.get("stop_reason") { let actual_sr = format!("{stop_reason:?}"); - // Normalize: serde serializes PascalCase, fixture may use either let expected_sr = normalize_stop_reason(exp_sr.as_str().unwrap_or("")); let actual_normalized = normalize_stop_reason(&actual_sr); assert_eq!( @@ -1162,9 +1038,6 @@ mod tests { ); } if let Some(exp_kind) = expected.get("kind").and_then(|v| v.as_str()) { - // Use an explicit match on ContentKind to produce the expected - // string, rather than relying on Debug formatting which could - // change if the Debug impl is modified. let actual_kind = match kind { llm_proxy_protocol::core::ContentKind::Text => "text", llm_proxy_protocol::core::ContentKind::ToolUse => "tool_use", @@ -1187,21 +1060,12 @@ mod tests { ); } } - // Ping -- variant-only comparison is sufficient. CoreEvent::Ping => {} - // #[non_exhaustive] requires a wildcard arm. New variants should - // get explicit match arms above for field-level validation. _ => {} } } } - /// Get the variant name of a CoreEvent for comparison. - /// - /// This match explicitly handles every known CoreEvent variant. Because - /// CoreEvent is `#[non_exhaustive]`, a wildcard arm is required, but it - /// panics so that any new variant introduced in the protocol crate is - /// caught immediately during test execution. fn event_variant_name(event: &CoreEvent) -> &'static str { match event { CoreEvent::MessageStart { .. } => "message_start", @@ -1215,8 +1079,6 @@ mod tests { CoreEvent::MessageStop { .. } => "message_stop", CoreEvent::Error { .. } => "error", CoreEvent::Ping => "ping", - // #[non_exhaustive] requires a wildcard arm. This will panic if a - // new CoreEvent variant is introduced without updating this match. _ => panic!( "event_variant_name: unknown CoreEvent variant. \ Add an explicit match arm for the new variant." @@ -1224,10 +1086,6 @@ mod tests { } } - // --------------------------------------------------------------------------- - // OpenAI Chat provider fixtures - // --------------------------------------------------------------------------- - #[test] fn openai_chat_encode_text_request() { run_encode_request_fixture(OPENAI_CHAT, "text_request"); @@ -1314,10 +1172,6 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Anthropic provider fixtures - // --------------------------------------------------------------------------- - #[test] fn anthropic_encode_text_request() { run_encode_request_fixture(ANTHROPIC, "text_request"); @@ -1414,10 +1268,6 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Responses provider fixtures - // --------------------------------------------------------------------------- - #[test] fn responses_encode_text_request() { run_encode_request_fixture(RESPONSES, "text_request"); @@ -1504,10 +1354,6 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Gemini provider fixtures - // --------------------------------------------------------------------------- - #[test] fn gemini_encode_text_request() { run_encode_request_fixture(GEMINI, "text_request"); @@ -1589,15 +1435,6 @@ mod tests { ); } - // =========================================================================== - // Coverage-matrix completeness test - // =========================================================================== - - /// Fixture cases required for every provider adapter (request encode direction). - /// Format: (case_name, required_files) - /// - /// Universal cases apply to all adapters. Adapter-specific extras are handled - /// in the `coverage_matrix_all_required_fixtures_exist` match block below. fn required_encode_request_cases() -> Vec<(&'static str, Vec<&'static str>)> { vec![ ("text_request", vec!["core.json", "output.json"]), @@ -1605,16 +1442,9 @@ mod tests { ("tool_request", vec!["core.json", "output.json"]), ("tool_result_request", vec!["core.json", "output.json"]), ("thinking_request", vec!["core.json", "output.json"]), - // Additional cases from the plan's 'Minimum fixture cases' section. - // These are not yet enforced as hard requirements (fixtures may not - // exist yet) but are listed here as a TODO for the next audit round: - // multiple_messages, tool_choice, refusal, image/document/audio/video - // content, cache_marker, sampling preservation, model_mapping, - // stop_sequence, unsupported core content, malformed client request. ] } - /// Fixture cases required for every provider adapter (response decode direction). fn required_decode_response_cases() -> Vec<(&'static str, Vec<&'static str>)> { vec![ ("text_response", vec!["input.json", "output.json"]), @@ -1622,29 +1452,18 @@ mod tests { ("thinking_response", vec!["input.json", "output.json"]), ("stop_reason_response", vec!["input.json", "output.json"]), ("malformed_response", vec!["input.json"]), - // Additional cases planned for next audit round: - // redacted_thinking, refusal, usage_mapping, stop_sequence_mapping, - // unsupported provider fields, lossy translation warning/provider_meta. ] } - /// Fixture cases required for every provider adapter (stream decode direction). fn required_stream_cases() -> Vec<(&'static str, Vec<&'static str>)> { vec![ ("stream_text", vec!["input.sse", "core-events.json"]), ("stream_tool", vec!["input.sse", "core-events.json"]), ("stream_usage", vec!["input.sse", "core-events.json"]), ("stream_malformed", vec!["input.sse", "core-events.json"]), - // Additional cases planned for next audit round: - // stream_thinking (adapter-specific -- see match block below), - // stream_error (adapter-specific -- see match block below), - // stream_ping (adapter-specific -- see match block below), - // partial-tool-call-buffering, upstream-disconnect, - // unknown-provider-events. ] } - /// Records every required fixture file that is absent on disk. fn push_missing_files( adapter: &str, cases: &[(&'static str, Vec<&'static str>)], @@ -1660,25 +1479,11 @@ mod tests { } } - /// Records missing fixtures for the stream cases only some adapters support. - /// - /// Each adapter documents which extra streaming cases it supports; an - /// adapter that cannot produce a given event (e.g. Gemini and - /// `ThinkingDelta`) documents the exclusion here rather than omitting it - /// silently. fn push_missing_adapter_specific(adapter: &str, missing: &mut Vec) { let extra_cases: &[&str] = match adapter { - // Anthropic supports ping, error, and thinking SSE events. ANTHROPIC => &["stream_error", "stream_ping", "stream_thinking"], - // OpenAI Chat supports thinking (reasoning) streaming. It emits - // neither Ping (no such SSE event) nor Error (upstream errors are - // HTTP status codes, handled at the transport level). OPENAI_CHAT => &["stream_thinking"], - // Responses supports error SSE events, but produces neither - // ThinkingDelta nor Ping. RESPONSES => &["stream_error"], - // Gemini's streaming tool calls are covered by the universal - // stream_tool case; it emits no thinking, error, or ping events. _ => &[], }; let cases: Vec<(&str, Vec<&str>)> = extra_cases @@ -1711,10 +1516,6 @@ mod tests { ); } - /// Asserts every `.json` fixture present for these cases parses. - /// - /// `skip_malformed` exempts `malformed_response/input.json`, which is - /// invalid on purpose. fn assert_json_fixtures_parse( adapter: &str, cases: &[(&'static str, Vec<&'static str>)], @@ -1736,8 +1537,6 @@ mod tests { } } - /// Asserts each stream case's `core-events.json` parses and its `input.sse` - /// is non-empty. fn assert_stream_fixtures_parse(adapter: &str, cases: &[(&'static str, Vec<&'static str>)]) { for (case, _) in cases { let dir = Path::new(FIXTURE_ROOT).join(adapter).join(case); @@ -1762,7 +1561,6 @@ mod tests { let adapters = [OPENAI_CHAT, ANTHROPIC, RESPONSES, GEMINI]; for adapter in &adapters { - // `malformed_response/input.json` is intentionally invalid JSON. assert_json_fixtures_parse(adapter, &required_encode_request_cases(), false); assert_json_fixtures_parse(adapter, &required_decode_response_cases(), true); assert_stream_fixtures_parse(adapter, &required_stream_cases()); diff --git a/crates/llm-proxy-server/Cargo.toml b/crates/llm-proxy-server/Cargo.toml index 8547f84..71944e0 100644 --- a/crates/llm-proxy-server/Cargo.toml +++ b/crates/llm-proxy-server/Cargo.toml @@ -18,6 +18,7 @@ llm-proxy-core = { workspace = true } llm-proxy-protocol = { path = "../llm-proxy-protocol" } llm-proxy-provider = { path = "../llm-proxy-provider" } llm-proxy-storage = { path = "../llm-proxy-storage" } +llm-proxy-api = { path = "../llm-proxy-api" } axum = { workspace = true, features = ["macros", "json"] } tokio = { workspace = true, features = ["fs"] } @@ -34,6 +35,7 @@ bytes = "1" futures = "0.3" regex = { workspace = true } sha2 = "0.10" +argon2 = { workspace = true } toml = { workspace = true } time = { workspace = true } secrecy = { workspace = true } @@ -42,3 +44,5 @@ rust_decimal = { workspace = true } [dev-dependencies] tower = { workspace = true, features = ["util"] } llm-proxy-storage = { path = "../llm-proxy-storage", features = ["test-utils"] } +tempfile = "3" +tracing-subscriber = { workspace = true } diff --git a/crates/llm-proxy-server/src/catalog_service.rs b/crates/llm-proxy-server/src/catalog_service.rs index 7f2b411..ccda2fc 100644 --- a/crates/llm-proxy-server/src/catalog_service.rs +++ b/crates/llm-proxy-server/src/catalog_service.rs @@ -14,33 +14,17 @@ use tokio::sync::{Mutex as AsyncMutex, RwLock}; #[derive(Debug, Clone)] struct RefreshOutcome { completed_at: Instant, - /// Preserved so single-flight followers receive the same error variant - /// (and therefore the same HTTP status via `map_catalog_error`) as the - /// leader. Only the variants `map_catalog_error` discriminates on are - /// preserved; everything else already collapses to `RouteError::Internal`, - /// so storing it as `InvalidConfig` is lossless. error: Option, } -/// Category-preserving snapshot of a [`ProviderError`] produced by a refresh. -/// -/// `ProviderError` is not `Clone` (it carries a `serde_json::Error`), so the -/// single-flight outcome cache stores this reconstruction instead of the full -/// error. Only the variants [`map_catalog_error`](crate::routes) discriminates -/// on are preserved by variant; all other variants map to `RouteError::Internal` -/// regardless, so they collapse to [`ProviderError::InvalidConfig`] losslessly. #[derive(Debug, Clone)] enum PreservedError { - /// Original was `ProviderError::Http`. Http { message: String, timeout: bool }, - /// Original was `ProviderError::Api`. Api { status: u16, body: String }, - /// Original was any other variant; only the `Display` string is retained. Other(String), } impl PreservedError { - /// Snapshot the category-relevant fields of a [`ProviderError`]. fn from_provider_error(error: &ProviderError) -> Self { match error { ProviderError::Http { @@ -57,8 +41,6 @@ impl PreservedError { } } - /// Reconstruct a [`ProviderError`] that maps to the same HTTP status as the - /// original (the `redirect` flag is irrelevant to status mapping). fn to_provider_error(&self) -> ProviderError { match self { Self::Http { message, timeout } => ProviderError::Http { @@ -75,11 +57,8 @@ impl PreservedError { } } -/// Cached, generation-tagged catalog membership index (audit LOW-9). #[derive(Clone, Debug)] struct IndexedCatalog { - /// The `cache_generations` value this index was built from; a mismatch - /// against the live generation marks the index stale and forces a rebuild. generation: u64, ids: Arc>, } @@ -89,14 +68,7 @@ struct IndexedCatalog { pub struct ModelCatalogService { cache_dir: Option, cache: RwLock>, - /// Bumped every time the discovered `cache` entry for a provider is - /// replaced, so `id_index` can detect staleness without holding a lock - /// across the network refresh (audit LOW-9). cache_generations: Mutex>, - /// Cached O(1) membership index of catalog model ids, keyed by provider - /// name (audit LOW-9). Built from the merged catalog; tagged with the - /// `cache` generation it reflects so a concurrent refresh forces a rebuild - /// instead of serving a stale set. id_index: RwLock>, refresh_locks: Mutex>>>, refresh_outcomes: Mutex>, @@ -120,8 +92,6 @@ impl ModelCatalogService { /// Return the merged, filtered catalog, optionally refreshing it live. /// - /// # Errors - /// /// Returns an error when a live refresh fails without stale/static fallback, /// or when a cache file is unsafe or malformed. pub async fn catalog( @@ -164,46 +134,20 @@ impl ModelCatalogService { )) } - /// O(1) catalog-enforcement membership check (audit LOW-9). - /// - /// Replaces the per-request O(N) linear scan over the merged catalog that - /// `resolve_target` used to perform via `catalog()` + a `Vec::iter().any` - /// check. Backed by a cached `HashSet` of catalog model ids (built by - /// `id_index_for`), rebuilt automatically when the provider's discovered - /// catalog is refreshed or loaded from disk. - /// - /// `protocol` follows the same contract as the former linear scan: the - /// `"gemini_generate_content"` protocol strips a leading `models/` prefix - /// from the request and accepts either spelling of the catalog id, so a - /// request for `"gemini-2.5-pro"` matches a catalog id of - /// `"models/gemini-2.5-pro"` (and vice-versa); every other protocol - /// requires an exact id match. - /// - /// # Errors - /// - /// Returns an error when the on-disk catalog cache cannot be loaded - /// (mirrors the failure mode of [`Self::catalog`]). + /// O(1) catalog-enforcement membership check for `upstream_model`: + /// `"gemini_generate_content"` strips a leading `models/` and accepts + /// either spelling; other protocols require an exact id match. pub async fn contains_model( &self, provider: &ProviderConfig, upstream_model: &str, protocol: &str, ) -> Result { - // Load the on-disk cache first so the index reflects the latest - // persisted catalog. Mirrors the load prelude of `catalog()`. self.load_disk_cache(provider).await?; let ids = self.id_index_for(provider).await?; Ok(model_id_present(&ids, upstream_model, protocol)) } - /// Return the cached membership index for `provider`, rebuilding it from - /// the merged catalog when absent or stale (audit LOW-9). - /// - /// Staleness is tracked via [`Self::current_generation`]: the index is - /// tagged with the generation it was built from, and a mismatch against the - /// live generation (advanced by every discovered-cache mutation) forces a - /// rebuild. The rebuild re-checks the generation before storing, so a - /// concurrent refresh can never leave a stale index behind. async fn id_index_for( &self, provider: &ProviderConfig, @@ -211,7 +155,6 @@ impl ModelCatalogService { let name = provider.name.as_str(); let gen_before = self.current_generation(name); - // Fast path: a current-generation index already exists. { let guard = self.id_index.read().await; if let Some(indexed) = guard.get(name) { @@ -221,8 +164,6 @@ impl ModelCatalogService { } } - // Slow path: build from the merged catalog (config allow/deny/static - // applied to the discovered models), exactly as `catalog()` returns. let config = provider.catalog.as_ref().cloned().unwrap_or_default(); let discovered = self .cache @@ -234,9 +175,6 @@ impl ModelCatalogService { let merged = merge_catalog(&config, &discovered); let ids = Arc::new(normalized_ids(&merged)); - // Store only if the discovered cache did not change under us. If a - // concurrent refresh advanced the generation, drop this build and let - // the next caller rebuild — a stale index is never persisted. let gen_after = self.current_generation(name); if gen_before == gen_after { let mut guard = self.id_index.write().await; @@ -251,8 +189,6 @@ impl ModelCatalogService { Ok(ids) } - /// Current generation of the discovered cache for `provider` (0 when never - /// populated). Advanced by [`Self::bump_cache_generation`]. fn current_generation(&self, provider: &str) -> u64 { self.cache_generations .lock() @@ -262,8 +198,6 @@ impl ModelCatalogService { .unwrap_or(0) } - /// Advance the generation for `provider`, marking any cached membership - /// index stale. Called after every mutation of the discovered cache. fn bump_cache_generation(&self, provider: &str) { let mut guard = self .cache_generations @@ -295,13 +229,9 @@ impl ModelCatalogService { .filter(|outcome| outcome.completed_at >= requested_at) .cloned() { - return outcome.error.map_or(Ok(()), |preserved| { - // Reconstruct the typed error so single-flight followers receive - // the same error variant -- and therefore the same HTTP status - // via `map_catalog_error` -- as the leader (e.g. an upstream - // timeout stays a 504, not a 500 Internal). - Err(preserved.to_provider_error()) - }); + return outcome + .error + .map_or(Ok(()), |preserved| Err(preserved.to_provider_error())); } let models = match self.discovery.discover(provider).await { @@ -347,8 +277,6 @@ impl ModelCatalogService { } } self.cache.write().await.insert(provider.name.clone(), file); - // The discovered catalog changed; advance the generation so any cached - // membership index is rebuilt on next access (audit LOW-9). self.bump_cache_generation(&provider.name); self.record_refresh_outcome(&provider.name, None)?; Ok(()) @@ -379,8 +307,6 @@ impl ModelCatalogService { .catalog .as_ref() .map_or_else(std::time::Duration::default, |c| c.cache_ttl); - // Clamp the TTL seconds to i64 range to avoid silent wrapping on - // absurdly large values. In practice, TTL is always well within range. let ttl_secs = i64::try_from(ttl.as_secs()).unwrap_or(i64::MAX); self.cache .read() @@ -404,23 +330,11 @@ impl ModelCatalogService { )); } let Some(cache_dir) = &self.cache_dir else { - // No disk cache: nothing to load. return Ok(()); }; - // Fast path: if this provider's catalog is already in memory, return - // without taking the write lock at all, avoiding contention on the - // common (already-loaded) request path. if self.cache.read().await.contains_key(&provider.name) { return Ok(()); } - // Cold path: stat + read + parse OUTSIDE any lock. The write lock is - // global across all providers, so holding it across disk I/O would - // serialize every concurrent catalog()/contains_model() call for every - // other provider (head-of-line blocking). Two concurrent first-load - // callers for the same provider may both read and parse the same file; - // that is harmless -- the second insert overwrites byte-identical - // content, and the extra bump_cache_generation only forces one cheap - // id_index rebuild (which re-checks the generation before storing). let path = cache_dir.join(format!("{}.toml", provider.name)); let metadata = match tokio::fs::symlink_metadata(&path).await { Ok(metadata) => metadata, @@ -448,9 +362,6 @@ impl ModelCatalogService { path.display() ))); } - // Critical section: re-check (another caller may have populated it) and - // insert. Bump the generation after dropping the guard so it does not - // run under the write lock. { let mut cache = self.cache.write().await; if cache.contains_key(&provider.name) { @@ -458,16 +369,11 @@ impl ModelCatalogService { } cache.insert(provider.name.clone(), file); } - // The discovered catalog changed; advance the generation so any cached - // membership index is rebuilt on next access (audit LOW-9). self.bump_cache_generation(&provider.name); Ok(()) } } -/// Build a `HashSet` of raw catalog model ids from a merged catalog slice -/// (audit LOW-9). The gemini `models/`-prefix normalization is handled at -/// lookup time in [`model_id_present`], so this stores only the verbatim ids. fn normalized_ids(merged: &[StaticModelCatalogEntry]) -> HashSet { let mut ids = HashSet::with_capacity(merged.len()); for entry in merged { @@ -476,21 +382,11 @@ fn normalized_ids(merged: &[StaticModelCatalogEntry]) -> HashSet { ids } -/// O(1) membership check against a catalog id set, preserving the exact -/// semantics of the former linear `catalog_contains_model` scan (audit LOW-9): -/// -/// - `"gemini_generate_content"` strips a leading `models/` from the request -/// and matches either the bare id or the `models/`-prefixed form present in -/// the catalog, so a request for `"gemini-2.5-pro"` matches a catalog id of -/// `"models/gemini-2.5-pro"` and vice-versa. -/// - Every other protocol requires an exact id match (no prefix stripping). fn model_id_present(ids: &HashSet, upstream_model: &str, protocol: &str) -> bool { if protocol == "gemini_generate_content" { let requested = upstream_model .strip_prefix("models/") .unwrap_or(upstream_model); - // The catalog may store either the bare id or the `models/`-prefixed - // form; accept either spelling (mirrors the former strip-prefix scan). return ids.contains(requested) || { let prefixed = format!("models/{requested}"); ids.contains(prefixed.as_str()) @@ -503,18 +399,6 @@ fn invalid_cache(error: &std::io::Error) -> ProviderError { ProviderError::InvalidConfig(format!("catalog cache I/O failed: {error}")) } -/// Format the current UTC time as an RFC 3339 string. -/// -/// # Errors -/// -/// Returns an error if formatting the timestamp fails. Propagated (rather -/// than falling back to the Unix epoch) because an epoch sentinel would make -/// the cache perpetually stale: `cache_is_fresh` computes -/// `now_utc() - generated_at`, and an epoch `generated_at` yields a huge -/// positive age, forcing a network refresh on every `catalog()` request. -/// Failing the refresh loudly is the safe direction — the caller already -/// handles refresh failures with a stale/static fallback (audit: -/// now-rfc3339-epoch-fallback). fn now_rfc3339() -> Result { time::OffsetDateTime::now_utc() .format(&time::format_description::well_known::Rfc3339) @@ -583,7 +467,6 @@ pub fn write_catalog_atomic(cache_dir: &Path, file: &CatalogFile) -> Result<(), Ok(()) } -/// Maximum allowed length for a provider slug. const MAX_PROVIDER_SLUG_LEN: usize = 128; fn is_safe_provider_slug(provider: &str) -> bool { @@ -623,6 +506,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -722,12 +606,6 @@ mod tests { .unwrap() .expect_err("every single-flight follower should observe the refresh failure"); } - // GAP-LOW-4: discovery retries transient (5xx/transport) failures up to - // `MAX_TRANSIENT_RETRIES` times after the first attempt, so a - // persistently-failing endpoint is probed `1 + MAX_TRANSIENT_RETRIES` - // times. Single-flight still holds — only one refresh actually reaches - // `discover`; the others reuse the cached failure outcome — so the - // count is exactly the per-call probe budget, not 8× it. assert_eq!( count.load(Ordering::SeqCst), 1 + DiscoveryClient::MAX_TRANSIENT_RETRIES as usize @@ -743,12 +621,6 @@ mod tests { let models = service.catalog(&provider(endpoint), true).await.unwrap(); assert_eq!(models.len(), 1); assert_eq!(models[0].id, "cached-model"); - // GAP-LOW-4: discovery retries transient (5xx/transport) failures up to - // `MAX_TRANSIENT_RETRIES` times after the first attempt, so a - // persistently-failing endpoint is probed `1 + MAX_TRANSIENT_RETRIES` - // times. Single-flight still holds — only one refresh actually reaches - // `discover`; the others reuse the cached failure outcome — so the - // count is exactly the per-call probe budget, not 8× it. assert_eq!( count.load(Ordering::SeqCst), 1 + DiscoveryClient::MAX_TRANSIENT_RETRIES as usize @@ -771,26 +643,19 @@ mod tests { let result = service.catalog(&provider, true).await; result.expect_err("catalog refresh should fail when the endpoint persistently errors"); - // GAP-LOW-4: discovery retries transient (5xx/transport) failures up to - // `MAX_TRANSIENT_RETRIES` times after the first attempt, so a - // persistently-failing endpoint is probed `1 + MAX_TRANSIENT_RETRIES` - // times. Single-flight still holds — only one refresh actually reaches - // `discover`; the others reuse the cached failure outcome — so the - // count is exactly the per-call probe budget, not 8× it. assert_eq!( count.load(Ordering::SeqCst), 1 + DiscoveryClient::MAX_TRANSIENT_RETRIES as usize ); } - // -- load_disk_cache validation edge cases --------------------------------- - fn minimal_provider(name: &str) -> ProviderConfig { ProviderConfig { name: name.to_owned(), api_key: secrecy::SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -805,7 +670,6 @@ mod tests { let directory = temp_dir(); std::fs::create_dir_all(&directory).unwrap(); - // Write a real cache file, then replace it with a symlink. let real = directory.join("real.toml"); std::fs::write(&real, "not used").unwrap(); let link = directory.join("symlink-provider.toml"); @@ -831,7 +695,6 @@ mod tests { let directory = temp_dir(); std::fs::create_dir_all(&directory).unwrap(); - // Create a FIFO (named pipe) which is neither a regular file nor a symlink. let fifo = directory.join("fifo-provider.toml"); #[cfg(unix)] { @@ -861,7 +724,6 @@ mod tests { let directory = temp_dir(); std::fs::create_dir_all(&directory).unwrap(); - // Write a cache file whose internal provider name differs from the filename. let mut file = catalog_file("2026-06-10T00:00:00Z"); file.catalog.provider = "other-provider".to_owned(); let raw = toml::to_string_pretty(&file).unwrap(); @@ -906,8 +768,6 @@ mod tests { std::fs::remove_dir_all(directory).unwrap(); } - // -- is_safe_provider_slug tests ------------------------------------------- - #[test] fn safe_slug_accepts_valid_names() { assert!(is_safe_provider_slug("openai")); @@ -946,8 +806,6 @@ mod tests { ); } - // -- cache_is_fresh logic tests -------------------------------------------- - #[tokio::test] async fn cache_is_fresh_returns_true_within_ttl() { let directory = temp_dir(); @@ -970,7 +828,6 @@ mod tests { #[tokio::test] async fn cache_is_fresh_returns_false_for_expired() { let directory = temp_dir(); - // generated_at far in the past (well beyond any reasonable TTL) write_catalog_atomic(&directory, &catalog_file("2020-01-01T00:00:00Z")).unwrap(); let service = ModelCatalogService::new(Some(directory.clone())); @@ -994,8 +851,6 @@ mod tests { ); } - // -- cache_is_fresh boundary tests (findings 4, 10) ------------------------ - #[tokio::test] async fn cache_is_fresh_with_zero_ttl_is_always_stale() { let directory = temp_dir(); @@ -1019,7 +874,6 @@ mod tests { #[tokio::test] async fn cache_is_fresh_with_future_timestamp_is_stale() { let directory = temp_dir(); - // A timestamp far in the future (clock skew) let future = time::OffsetDateTime::now_utc() + time::Duration::hours(1); let future_str = future .format(&time::format_description::well_known::Rfc3339) @@ -1053,8 +907,6 @@ mod tests { std::fs::remove_dir_all(directory).unwrap(); } - // -- Concurrent successful refresh test (finding 13) ----------------------- - async fn succeeding_server(model_ids: Vec<&str>) -> (String, Arc) { let count = Arc::new(AtomicUsize::new(0)); let models_json = model_ids @@ -1121,16 +973,13 @@ mod tests { ); } - // -- contains_model O(1) index (audit LOW-9) ------------------------------- - - /// A provider whose catalog is purely `Static` (no discovery), so the - /// membership index is built lazily from config and never invalidated. fn static_catalog_provider(name: &str, models: Vec) -> ProviderConfig { ProviderConfig { name: name.to_owned(), api_key: secrecy::SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -1178,8 +1027,6 @@ mod tests { #[tokio::test] async fn contains_model_gemini_prefix_normalization() { - // Catalog stores the `models/`-prefixed form; gemini requests may omit - // the prefix, non-gemini requests must not strip it. let provider = static_catalog_provider("gemini-prov", vec![model_entry("models/gemini-2.5-pro")]); let service = ModelCatalogService::new(None); @@ -1208,7 +1055,6 @@ mod tests { .unwrap(), "gemini: unrelated model must be rejected" ); - // Non-gemini protocol must NOT strip the prefix — exact match only. assert!( service .contains_model(&provider, "models/gemini-2.5-pro", "openai_chat") @@ -1227,18 +1073,13 @@ mod tests { #[tokio::test] async fn contains_model_index_invalidated_after_refresh() { - // Seed the disk cache with the default "cached-model" so the first - // membership check builds an index from it. let directory = temp_dir(); write_catalog_atomic(&directory, &catalog_file("2026-06-10T00:00:00Z")).unwrap(); let (endpoint, _count) = succeeding_server(vec!["new-model"]).await; let mut provider = provider(endpoint); - // Zero TTL forces the next live catalog() call to refresh, which - // replaces the discovered cache and bumps the generation. provider.catalog.as_mut().unwrap().cache_ttl = Duration::ZERO; let service = ModelCatalogService::new(Some(directory.clone())); - // First check loads the disk cache ("cached-model") and caches the index. assert!( service .contains_model(&provider, "cached-model", "openai_chat") @@ -1247,8 +1088,6 @@ mod tests { "seeded model must be present before refresh" ); - // A live refresh replaces the discovered catalog with "new-model" and - // bumps the generation, forcing the index to rebuild. service.catalog(&provider, true).await.unwrap(); assert!( diff --git a/crates/llm-proxy-server/src/lib.rs b/crates/llm-proxy-server/src/lib.rs index 642c4b7..b49b294 100644 --- a/crates/llm-proxy-server/src/lib.rs +++ b/crates/llm-proxy-server/src/lib.rs @@ -1,7 +1,4 @@ //! The axum wiring for `llm-proxy`. -//! -//! Exposes [`build_router`] which assembles the full HTTP router -//! with middleware, routes, and shared state. #![deny(missing_docs)] @@ -17,22 +14,29 @@ pub mod shutdown; pub mod state; pub use catalog_service::{ModelCatalogService, write_catalog_atomic}; -pub use middleware::{RateLimiter, RequestDeduplicator, RequestId, RequestIdGenerator}; +pub use middleware::{ + CallerId, ClientAuth, ClientAuthError, RateLimiter, RequestDeduplicator, RequestId, + RequestIdGenerator, +}; pub use shutdown::shutdown_signal; pub use state::{AppState, BuildInfo}; use axum::Router; +use secrecy::ExposeSecret; /// Build the application router for the given state. -/// -/// Carries an explicit `#[must_use]` (audit LOW-2): `axum::Router` is not -/// itself `#[must_use]`, so this attribute is the single, non-redundant -/// reminder that the returned router must be driven by a server. No -/// `clippy::double_must_use` annotation is warranted — that lint does not fire -/// here (verified: an `#[expect(clippy::double_must_use)]` was unfulfilled), -/// because only the function, not the return type, is `#[must_use]`. The -/// former inert `#[allow(...)]` that LOW-2 flagged has been removed. #[must_use = "the Router must be used with a server"] pub fn build_router(state: AppState) -> Router { routes::router(state) } + +/// Build the admin `/usage` router behind the `[auth.admin]` token. +/// +/// Returns `None` when no `[auth.admin]` token or `[storage]` is configured; +/// every route is a read-only `GET`. +#[must_use = "the Router must be used with a server"] +pub fn admin_router(state: &AppState) -> Option { + let admin = state.app_config.auth.admin.as_ref()?; + let storage = state.app_config.storage.as_ref()?; + Some(llm_proxy_api::UsageApi::new(storage.path.clone(), admin.token.expose_secret()).router()) +} diff --git a/crates/llm-proxy-server/src/middleware.rs b/crates/llm-proxy-server/src/middleware.rs index ceaf3a7..ef470f4 100644 --- a/crates/llm-proxy-server/src/middleware.rs +++ b/crates/llm-proxy-server/src/middleware.rs @@ -1,7 +1,4 @@ //! Middleware components for the LLM proxy server. -//! -//! Provides request deduplication, rate limiting, request ID generation, -//! and client IP extraction. use std::collections::HashMap; use std::convert::Infallible; @@ -12,49 +9,23 @@ use std::time::Instant; use axum::extract::{ConnectInfo, FromRequestParts}; use axum::http::request::Parts; +use llm_proxy_core::{ANONYMOUS_CALLER, AuthConfig, AuthMode}; +use secrecy::{ExposeSecret, SecretString}; use sha2::{Digest, Sha256}; -// --------------------------------------------------------------------------- -// RequestDeduplicator -// --------------------------------------------------------------------------- +use argon2::PasswordVerifier; -/// Maximum number of in-flight deduplication entries. Prevents unbounded -/// memory growth from unique-body requests. const MAX_DEDUP_ENTRIES: usize = 10_000; -/// Maximum body size (in bytes) eligible for dedup tracking. Bodies larger -/// than this are never hashed or tracked to avoid expensive SHA-256 computation -/// on large payloads and prevent memory pressure from storing many hash entries. -/// 64 KiB is chosen as a reasonable threshold: most LLM API requests with -/// moderate context fit within this limit, while very large file uploads or -/// multi-turn conversations with huge contexts are excluded from dedup. const MAX_DEDUP_BODY_SIZE: usize = 64 * 1024; -/// Deduplicates requests based on a SHA-256 hash of the request path and body. +/// Deduplicates requests by SHA-256 hash of the request path and body. /// -/// Tracks in-flight request hashes with a configurable deduplication window. -/// The standalone type defaults to 500 ms; server configuration defaults to -/// disabled. If a second request with the same path and body arrives during a -/// nonzero window, the duplicate is rejected. -/// -/// The request path is included in the hash so that the same body sent to -/// different protocol endpoints (e.g. `/v1/messages` vs `/v1/chat/completions`) -/// is not treated as a duplicate. -/// -/// # Why `std::sync::Mutex`? -/// -/// Uses `std::sync::Mutex` rather than `tokio::sync::Mutex` because the -/// critical section is purely in-memory HashMap operations (insert, remove, -/// retain) that complete in microseconds. `tokio::sync::Mutex` would add -/// unnecessary async overhead for a non-I/O-bound lock. The mutex is never -/// held across `.await` points. +/// A second identical path+body inside the window (default 500 ms; 0 +/// disables) is rejected. #[derive(Debug)] pub struct RequestDeduplicator { - /// SHA-256 raw digest -> insertion time. Keyed on the 32-byte digest - /// (`[u8; 32]` is `Copy + Eq + Hash`) rather than a 64-char hex `String` - /// to avoid a per-request allocation and formatting pass (audit LOW-8). in_flight: Mutex>, - /// Deduplication window in milliseconds. window_ms: u64, } @@ -74,28 +45,20 @@ impl RequestDeduplicator { /// Check whether this request body is a duplicate. /// - /// Returns `true` if the request is a duplicate (should be rejected), - /// `false` if it is new (should be processed). - /// - /// # Deprecation - /// - /// Prefer [`Self::is_duplicate_with_path`] which includes the request path in the - /// hash to distinguish requests to different protocol endpoints. + /// Deprecated: prefer [`Self::is_duplicate_with_path`], which includes the + /// request path in the hash. #[deprecated(note = "use is_duplicate_with_path instead")] pub fn is_duplicate(&self, body: &[u8]) -> bool { self.is_duplicate_with_path("", body) } /// Check whether this request is a duplicate, keyed by path and body. - /// - /// Including the path distinguishes requests to different protocol - /// endpoints (e.g. `/v1/messages` vs `/v1/chat/completions`) that - /// happen to carry the same body. + /// The path distinguishes requests to different protocol endpoints + /// that carry the same body. pub fn is_duplicate_with_path(&self, path: &str, body: &[u8]) -> bool { if self.window_ms == 0 { return false; } - // Skip dedup for large bodies to avoid expensive hashing and memory pressure. if body.len() > MAX_DEDUP_BODY_SIZE { return false; } @@ -107,26 +70,18 @@ impl RequestDeduplicator { e.into_inner() }); - // Prune expired entries. map.retain(|_, t| now.duration_since(*t).as_millis() < self.window_ms as u128); if map.contains_key(&hash) { return true; } - // Skip dedup tracking when at capacity to bound memory usage. if map.len() < MAX_DEDUP_ENTRIES { map.insert(hash, now); } false } - /// Compute the SHA-256 digest of the request path and body. - /// - /// Returns the raw 32-byte digest directly (rather than a 64-char hex - /// `String`) since `[u8; 32]` is `Copy + Eq + Hash` and serves as the - /// `in_flight` key without any formatting pass or heap allocation - /// (audit LOW-8). fn hash_path_body(path: &str, body: &[u8]) -> [u8; 32] { let mut hasher = Sha256::new(); hasher.update(path.as_bytes()); @@ -141,22 +96,12 @@ impl Default for RequestDeduplicator { } } -// --------------------------------------------------------------------------- -// RateLimiter -// --------------------------------------------------------------------------- - -/// Maximum number of per-IP rate-limit buckets. Prevents unbounded memory -/// growth from spoofed or diverse source IPs. const MAX_RATE_LIMIT_BUCKETS: usize = 100_000; -/// Per-client token bucket for rate limiting. #[derive(Debug)] struct ClientTokenBucket { - /// Remaining tokens. tokens: f64, - /// Maximum tokens. max_tokens: f64, - /// Timestamp of last refill. last_refill: Instant, } @@ -169,7 +114,6 @@ impl ClientTokenBucket { } } - /// Try to consume one token. Returns `true` if allowed. fn try_consume(&mut self, refill_rate: f64) -> bool { let now = Instant::now(); let elapsed = now.duration_since(self.last_refill).as_secs_f64(); @@ -186,44 +130,22 @@ impl ClientTokenBucket { } /// Per-IP rate limiter using token buckets. -/// -/// Each client IP gets its own bucket allowing `max_requests_per_minute` -/// requests per minute. Pruning of stale entries occurs periodically -/// (every `PRUNE_INTERVAL` requests) rather than on every call to reduce -/// per-request overhead under high load. #[derive(Debug)] pub struct RateLimiter { - /// Per-IP token buckets. buckets: Mutex>, - /// Maximum requests per minute per client. max_requests_per_minute: f64, - /// Counter for periodic pruning. Pruning runs every `PRUNE_INTERVAL` requests. prune_counter: AtomicU64, } -/// Number of requests between prune sweeps. Trade-off: lower values prune -/// more aggressively but add overhead; higher values reduce overhead but -/// allow stale entries to linger longer. const PRUNE_INTERVAL: u64 = 128; -/// Idle eviction threshold for per-IP rate-limit buckets (5 minutes). -/// -/// Buckets whose last refill was more than this duration ago are evicted -/// during periodic pruning to prevent unbounded memory growth. const RATE_LIMITER_IDLE_EVICT_SECS: u64 = 300; impl RateLimiter { /// Create a new rate limiter. /// - /// `max_requests_per_minute` is the maximum number of requests allowed - /// per client IP per minute. `0` disables rate limiting. - /// - /// # Panics - /// - /// Does not panic, but extremely large values (e.g. `u32::MAX`) create - /// per-IP buckets with billions of tokens. Operators should choose - /// reasonable values (e.g. 100--100_000 RPM) to avoid pathological - /// memory usage from unbounded per-IP bucket allocation. + /// `max_requests_per_minute` is the per-client-IP request budget per + /// minute; `0` disables rate limiting. pub fn new(max_requests_per_minute: u32) -> Self { Self { buckets: Mutex::new(HashMap::new()), @@ -244,8 +166,6 @@ impl RateLimiter { e.into_inner() }); - // Periodic pruning: evict stale entries every PRUNE_INTERVAL requests - // to prevent unbounded memory growth without paying the cost on every call. if self.prune_counter.fetch_add(1, Ordering::Relaxed) % PRUNE_INTERVAL == 0 { let now = std::time::Instant::now(); buckets.retain(|_, bucket| { @@ -254,20 +174,12 @@ impl RateLimiter { }); } - // Refill rate: tokens per second. let refill_rate = self.max_requests_per_minute / 60.0; - // Single-lookup hot path: for a known IP we mutate the existing bucket - // in place without allocating a key `String`. Only on a miss do we - // perform the capacity check and a second (entry) probe that allocates. - // This avoids the redundant `contains_key` + `entry` double hash probe - // and the unconditional `to_owned()` on every call (audit LOW-6). if let Some(bucket) = buckets.get_mut(client_ip) { return bucket.try_consume(refill_rate); } - // New IP: reject if inserting it would exceed the bucket cap. Bounds - // memory growth from spoofed or diverse source IPs. if buckets.len() >= MAX_RATE_LIMIT_BUCKETS { return false; } @@ -278,6 +190,28 @@ impl RateLimiter { bucket.try_consume(refill_rate) } + + /// Whether a request from `client_ip` would currently be allowed, + /// **without consuming a token**. + #[must_use] + pub fn would_allow(&self, client_ip: &str) -> bool { + if self.max_requests_per_minute == 0.0 { + return true; + } + let buckets = self.buckets.lock().unwrap_or_else(|e| { + tracing::warn!("rate limiter mutex poisoned; recovering from panic"); + e.into_inner() + }); + let Some(bucket) = buckets.get(client_ip) else { + return true; + }; + let elapsed = Instant::now() + .duration_since(bucket.last_refill) + .as_secs_f64(); + let refill_rate = self.max_requests_per_minute / 60.0; + let tokens = (bucket.tokens + elapsed * refill_rate).min(bucket.max_tokens); + tokens >= 1.0 + } } impl Default for RateLimiter { @@ -286,10 +220,6 @@ impl Default for RateLimiter { } } -// --------------------------------------------------------------------------- -// RequestIdGenerator -// --------------------------------------------------------------------------- - /// Generates unique request IDs in the format `req-{unix}-{counter}`. #[derive(Debug)] pub struct RequestIdGenerator { @@ -304,15 +234,10 @@ impl RequestIdGenerator { } } - /// Generate the next request ID. - /// - /// The ID format is `req-{unix_seconds}-{monotonic_counter}`. + /// Generate the next request ID, `req-{unix_seconds}-{monotonic_counter}`. /// - /// # Clock fallback - /// - /// If `SystemTime::now()` is before `UNIX_EPOCH` (nearly impossible on - /// modern systems), the timestamp falls back to 0. The monotonic counter - /// still guarantees uniqueness within the process. + /// A pre-epoch clock falls back to timestamp `0`; the counter still + /// guarantees uniqueness within the process. pub fn next_id(&self) -> String { let unix = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -321,9 +246,6 @@ impl RequestIdGenerator { tracing::error!(error = %e, "system clock is before UNIX epoch; using 0 as request ID timestamp"); 0 }); - // Ordering::Relaxed is sufficient: the counter only needs to be - // monotonic within this process. No other memory depends on the - // counter value for synchronization. let counter = self.counter.fetch_add(1, Ordering::Relaxed); format!("req-{unix}-{counter}") } @@ -336,12 +258,6 @@ impl Default for RequestIdGenerator { } /// Typed request identifier carried in request extensions. -/// -/// Inserted by an outermost `middleware::from_fn` layer so that the -/// `TraceLayer` span (`make_span_with`) and the route handler share a single -/// id. This keeps the id that appears in the `x-request-id` response header and -/// in handler log events identical to the id stamped on the request's tracing -/// span, restoring end-to-end span/event correlation (audit MEDIUM-1). #[derive(Debug, Clone)] pub struct RequestId(pub String); @@ -351,10 +267,6 @@ impl std::fmt::Display for RequestId { } } -// --------------------------------------------------------------------------- -// Client IP extraction -// --------------------------------------------------------------------------- - /// Optional peer socket metadata for handlers that also run in router-only tests. #[derive(Clone, Copy, Debug)] pub struct OptionalConnectInfo(pub Option); @@ -377,31 +289,15 @@ where } } -/// Extract the client IP from trusted forwarding headers or connection info. -/// -/// When `trust_forwarded_headers` is true, the leftmost value from -/// `X-Forwarded-For` or the value from `X-Real-Ip` is used, but only after -/// validating it parses as a legitimate [`std::net::IpAddr`]. Invalid or -/// spoofed values are silently ignored and the connection-info fallback is -/// used instead. -/// -/// # Trust model -/// -/// The **leftmost** `X-Forwarded-For` entry is trusted as the client address. -/// This is only correct behind a **single** reverse proxy that strips or -/// overwrites any inbound `X-Forwarded-For` before appending its own hop. With -/// multiple untrusted hops the leftmost entry is attacker-controllable (a -/// client may prepend arbitrary IPs), so the extracted address must not be -/// trusted for authentication, rate-limit bypass, or audit logging. The IpAddr -/// parse check defends against malformed garbage, not against well-formed -/// spoofed values. +/// Client IP: the leftmost `X-Forwarded-For`/`X-Real-Ip` value when +/// `trust_forwarded_headers` is set and it parses as an IP, else the +/// connection info, else `"unknown"`. Safe only behind a single proxy. pub fn get_client_ip( headers: &axum::http::HeaderMap, connect_info: Option<&SocketAddr>, trust_forwarded_headers: bool, ) -> String { if trust_forwarded_headers { - // Check X-Forwarded-For first (leftmost IP). if let Some(xff) = headers.get("x-forwarded-for") && let Ok(val) = xff.to_str() && let Some(ip) = val.split(',').next() @@ -412,7 +308,6 @@ pub fn get_client_ip( } } - // Check X-Real-Ip. if let Some(xri) = headers.get("x-real-ip") && let Ok(val) = xri.to_str() { @@ -423,7 +318,6 @@ pub fn get_client_ip( } } - // Fall back to connection info. if let Some(ci) = connect_info { return ci.ip().to_string(); } @@ -431,15 +325,505 @@ pub fn get_client_ip( "unknown".to_owned() } -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- +struct ClientVerifier { + key_id: String, + phc: SecretString, +} + +struct CachedDigest { + digest: [u8; 32], + key_id: String, +} + +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + +/// Errors raised while building [`ClientAuth`] from configuration. +#[derive(Debug, thiserror::Error)] +pub enum ClientAuthError { + /// `auth.mode = "none"` was combined with configured client keys. + #[error("auth.mode is \"none\" but [auth.clients] is non-empty")] + ModeNoneWithClients, + /// `auth.mode = "client-keys"` had no client keys configured. + #[error("auth.mode is \"client-keys\" but [auth.clients] is empty")] + ClientKeysWithoutClients, + /// A key id was empty or the reserved `anonymous` name. + #[error("client key id \"{0}\" is empty or reserved")] + InvalidKeyId(String), + /// A verifier string is not a parseable PHC argon2 hash. + /// + /// The error carries the key id, never the hash string itself. + #[error("client key \"{key_id}\" has an invalid PHC verifier: {source}")] + InvalidVerifier { + /// The key id whose verifier failed to parse. + key_id: String, + /// The underlying parse error. + source: argon2::password_hash::Error, + }, +} + +/// Client-key authentication state: the configured verifiers plus a digest +/// cache for verified keys. Holds hashed verifiers only, never plaintext. +/// +/// Manual `Debug` prints the mode and key ids, never the verifier strings. +pub struct ClientAuth { + enforced: bool, + verifiers: Vec, + cache: Mutex>, +} + +impl std::fmt::Debug for ClientAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClientAuth") + .field("enforced", &self.enforced) + .field( + "client_key_ids", + &self + .verifiers + .iter() + .map(|v| v.key_id.as_str()) + .collect::>(), + ) + .finish() + } +} + +impl ClientAuth { + /// Build authentication state from the `[auth]` config section. + /// + /// Returns [`ClientAuthError`] on contradictory mode/client combinations, + /// reserved or empty key ids, or an unparseable verifier string. + pub fn from_config(config: &AuthConfig) -> Result { + match config.mode { + AuthMode::None if !config.clients.is_empty() => { + return Err(ClientAuthError::ModeNoneWithClients); + } + AuthMode::ClientKeys if config.clients.is_empty() => { + return Err(ClientAuthError::ClientKeysWithoutClients); + } + AuthMode::None | AuthMode::ClientKeys => {} + } + let mut verifiers = Vec::with_capacity(config.clients.len()); + for (key_id, phc) in &config.clients { + if key_id.is_empty() || key_id == ANONYMOUS_CALLER { + return Err(ClientAuthError::InvalidKeyId(key_id.clone())); + } + if let Err(source) = argon2::password_hash::phc::PasswordHash::new(phc.expose_secret()) + { + return Err(ClientAuthError::InvalidVerifier { + key_id: key_id.clone(), + source: source.into(), + }); + } + verifiers.push(ClientVerifier { + key_id: key_id.clone(), + phc: SecretString::from(phc.expose_secret().to_owned()), + }); + } + verifiers.sort_by(|a, b| a.key_id.cmp(&b.key_id)); + Ok(Self { + enforced: matches!(config.mode, AuthMode::ClientKeys), + verifiers, + cache: Mutex::new(Vec::new()), + }) + } + + /// Fail-closed authenticator: enforcement on, nothing can authenticate. + /// + /// Used when config is invalid at construction time so every proxied + /// request is rejected rather than silently downgraded to no-auth. + #[must_use] + pub fn fail_closed() -> Self { + Self { + enforced: true, + verifiers: Vec::new(), + cache: Mutex::new(Vec::new()), + } + } + + /// Whether proxy routes require a valid client key. + #[must_use] + pub fn is_enforced(&self) -> bool { + self.enforced + } + + /// Number of configured client keys (0 when not enforced). + #[must_use] + pub fn client_count(&self) -> usize { + self.verifiers.len() + } + + fn cached_key_id(&self, digest: &[u8; 32]) -> Option { + let cache = self.cache.lock().unwrap_or_else(|e| { + tracing::warn!("client auth digest cache poisoned; recovering from panic"); + e.into_inner() + }); + cache + .iter() + .find(|entry| constant_time_eq(&entry.digest, digest)) + .map(|entry| entry.key_id.clone()) + } + + async fn verify_uncached(&self, token: &str) -> Option { + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + let verifiers: Vec<(String, String)> = self + .verifiers + .iter() + .map(|v| (v.key_id.clone(), v.phc.expose_secret().to_owned())) + .collect(); + let token = token.to_owned(); + let matched = tokio::task::spawn_blocking(move || { + let argon2 = argon2::Argon2::default(); + for (key_id, phc) in &verifiers { + if let Ok(hash) = argon2::password_hash::phc::PasswordHash::new(phc) + && argon2.verify_password(token.as_bytes(), &hash).is_ok() + { + return Some(key_id.clone()); + } + } + None + }) + .await + .map_err(|error| { + tracing::warn!(%error, "client auth verification task failed"); + error + }) + .ok() + .flatten(); + if let Some(key_id) = &matched { + let mut cache = self.cache.lock().unwrap_or_else(|e| { + tracing::warn!("client auth digest cache poisoned; recovering from panic"); + e.into_inner() + }); + if !cache + .iter() + .any(|entry| constant_time_eq(&entry.digest, &digest)) + { + cache.push(CachedDigest { + digest, + key_id: key_id.clone(), + }); + } + } + matched + } + + pub(crate) async fn authenticate(&self, token: &str) -> Option { + if !self.enforced || self.verifiers.is_empty() { + return None; + } + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + if let Some(key_id) = self.cached_key_id(&digest) { + return Some(key_id); + } + self.verify_uncached(token).await + } +} + +pub(crate) fn extract_bearer_token(headers: &axum::http::HeaderMap) -> Option { + let value = headers.get(axum::http::header::AUTHORIZATION)?; + let header = value.to_str().ok()?; + let mut parts = header.splitn(2, ' '); + match (parts.next(), parts.next()) { + (Some(scheme), Some(token)) if scheme.eq_ignore_ascii_case("bearer") => { + let token = token.trim(); + if token.is_empty() { + None + } else { + Some(SecretString::from(token.to_owned())) + } + } + _ => None, + } +} + +/// Typed caller identity carried in request extensions. +/// +/// The authenticated key id when client keys are enforced, or +/// [`CallerId::anonymous`] under the loopback no-auth exemption. +#[derive(Debug, Clone)] +pub struct CallerId { + key_id: Option, +} + +impl CallerId { + /// The unauthenticated caller used under the loopback no-auth exemption. + #[must_use] + pub fn anonymous() -> Self { + Self { key_id: None } + } + + /// An authenticated caller identified by `key_id`. + #[must_use] + pub fn authenticated(key_id: impl Into) -> Self { + Self { + key_id: Some(key_id.into()), + } + } + + /// Whether a client key authenticated this request. + #[must_use] + pub fn is_authenticated(&self) -> bool { + self.key_id.is_some() + } + + /// The attribution string for events and logs: the key id, or + /// `"anonymous"` when unauthenticated. + #[must_use] + pub fn as_str(&self) -> &str { + self.key_id.as_deref().unwrap_or(ANONYMOUS_CALLER) + } +} + +impl std::fmt::Display for CallerId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RequestMeta { + pub(crate) request_id: RequestId, + pub(crate) caller: CallerId, + pub(crate) connect_info: Option, +} + +impl FromRequestParts for RequestMeta +where + S: Send + Sync, +{ + type Rejection = Infallible; + + fn from_request_parts( + parts: &mut Parts, + _state: &S, + ) -> impl Future> + Send { + let request_id = parts + .extensions + .get::() + .cloned() + .unwrap_or_else(|| RequestId("unknown".to_owned())); + let caller = parts + .extensions + .get::() + .cloned() + .unwrap_or_else(CallerId::anonymous); + let connect_info = parts + .extensions + .get::>() + .map(|connect_info| connect_info.0); + async move { + Ok(Self { + request_id, + caller, + connect_info, + }) + } + } +} #[cfg(test)] mod tests { use super::*; - // -- RequestDeduplicator --------------------------------------------------- + fn cheap_phc(token: &str, salt: &[u8]) -> String { + let params = argon2::Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = argon2::Argon2::new( + argon2::Algorithm::Argon2id, + argon2::Version::V0x13, + params.clone(), + ); + use argon2::CustomizedPasswordHasher; + hasher + .hash_password_customized( + token.as_bytes(), + salt, + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string() + } + + fn auth_config_with(phc: &str) -> llm_proxy_core::AuthConfig { + let mut clients = HashMap::new(); + clients.insert("laptop".to_owned(), SecretString::from(phc.to_owned())); + llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + } + } + + #[test] + fn client_auth_from_config_rejects_mode_contradictions() { + let none_with_clients = { + let mut clients = HashMap::new(); + clients.insert( + "laptop".to_owned(), + SecretString::from(cheap_phc("secret", b"0123456789abcdef")), + ); + llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients, + admin: None, + } + }; + assert!(matches!( + ClientAuth::from_config(&none_with_clients), + Err(ClientAuthError::ModeNoneWithClients) + )); + + let keys_without_clients = llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients: HashMap::new(), + admin: None, + }; + assert!(matches!( + ClientAuth::from_config(&keys_without_clients), + Err(ClientAuthError::ClientKeysWithoutClients) + )); + } + + #[test] + fn client_auth_from_config_rejects_bad_key_ids_and_hashes() { + for bad_id in ["", "anonymous"] { + let mut config = auth_config_with(&cheap_phc("secret", b"0123456789abcdef")); + config.clients.remove("laptop"); + config + .clients + .insert(bad_id.to_owned(), SecretString::from("x".to_owned())); + assert!( + matches!( + ClientAuth::from_config(&config), + Err(ClientAuthError::InvalidKeyId(id)) if id == bad_id + ), + "key id {bad_id:?} must be rejected" + ); + } + + let config = auth_config_with("not-a-phc-string"); + let error = ClientAuth::from_config(&config).expect_err("malformed verifier rejected"); + assert!( + error.to_string().contains("laptop"), + "error names the key id, got: {error}" + ); + assert!( + !error.to_string().contains("not-a-phc-string"), + "error must not echo the verifier, got: {error}" + ); + } + + #[tokio::test] + async fn client_auth_authenticates_a_configured_key_and_rejects_others() { + let secret = cheap_phc("the-real-secret", b"0123456789abcdef"); + let auth = ClientAuth::from_config(&auth_config_with(&secret)).expect("valid config"); + + let matched = auth.authenticate("the-real-secret").await; + assert_eq!(matched.as_deref(), Some("laptop")); + assert_eq!( + auth.authenticate("the-real-secret").await.as_deref(), + Some("laptop") + ); + assert_eq!(auth.authenticate("wrong-secret").await, None); + assert_eq!(auth.authenticate("").await, None); + } + + #[tokio::test] + async fn client_auth_picks_the_matching_key_among_many() { + let mut clients = HashMap::new(); + for (id, token) in [("alpha", "one"), ("beta", "two"), ("gamma", "three")] { + clients.insert( + id.to_owned(), + SecretString::from(cheap_phc(token, b"0123456789abcdef")), + ); + } + let config = llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }; + let auth = ClientAuth::from_config(&config).expect("valid config"); + assert_eq!(auth.authenticate("two").await.as_deref(), Some("beta")); + assert_eq!(auth.authenticate("three").await.as_deref(), Some("gamma")); + assert_eq!(auth.authenticate("four").await, None); + } + + #[test] + fn client_auth_none_mode_is_not_enforced() { + let auth = + ClientAuth::from_config(&llm_proxy_core::AuthConfig::default()).expect("valid config"); + assert!(!auth.is_enforced()); + } + + #[test] + fn client_auth_debug_shows_key_ids_but_never_verifiers() { + let secret = cheap_phc("the-real-secret", b"0123456789abcdef"); + let auth = ClientAuth::from_config(&auth_config_with(&secret)).expect("valid config"); + let debug = format!("{auth:?}"); + assert!(debug.contains("laptop"), "got: {debug}"); + assert!( + !debug.contains("argon2id") && !debug.contains(&secret), + "verifier string must not appear in Debug, got: {debug}" + ); + } + + #[test] + fn constant_time_eq_behaves_as_equality() { + assert!(constant_time_eq(b"abc", b"abc")); + assert!(!constant_time_eq(b"abc", b"abd")); + assert!(!constant_time_eq(b"abc", b"ab")); + } + + #[test] + fn bearer_token_extracts_case_insensitive_scheme_and_trims() { + use axum::http::HeaderMap; + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer abc123".parse().unwrap()); + assert_eq!( + extract_bearer_token(&headers).map(|t| t.expose_secret().to_owned()), + Some("abc123".to_owned()) + ); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "bearer abc123 ".parse().unwrap()); + assert_eq!( + extract_bearer_token(&headers).map(|t| t.expose_secret().to_owned()), + Some("abc123".to_owned()) + ); + } + + #[test] + fn bearer_token_rejects_missing_other_scheme_and_empty() { + use axum::http::HeaderMap; + assert!(extract_bearer_token(&HeaderMap::new()).is_none()); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Basic dXNlcjpwYXNz".parse().unwrap()); + assert!(extract_bearer_token(&headers).is_none()); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer ".parse().unwrap()); + assert!(extract_bearer_token(&headers).is_none()); + } + + #[test] + fn caller_id_attribution_strings() { + assert_eq!(CallerId::anonymous().as_str(), "anonymous"); + assert_eq!(CallerId::authenticated("laptop").as_str(), "laptop"); + assert!(!CallerId::anonymous().is_authenticated()); + assert!(CallerId::authenticated("laptop").is_authenticated()); + assert_eq!(CallerId::authenticated("laptop").to_string(), "laptop"); + } #[test] fn dedup_new_request_is_not_duplicate() { @@ -474,12 +858,9 @@ mod tests { let generator = RequestIdGenerator::new(); let id1 = generator.next_id(); let id2 = generator.next_id(); - // The counter portion should be different. assert_ne!(id1, id2); } - // -- RateLimiter ----------------------------------------------------------- - #[test] fn rate_limiter_allows_under_limit() { let limiter = RateLimiter::new(100); @@ -494,17 +875,13 @@ mod tests { for _ in 0..5 { assert!(limiter.is_allowed("10.0.0.1")); } - // Different IP should still be allowed. assert!(limiter.is_allowed("10.0.0.2")); } - // -- RequestDeduplicator with path ----------------------------------------- - #[test] fn dedup_same_body_different_path_is_not_duplicate() { let dedup = RequestDeduplicator::new(); assert!(!dedup.is_duplicate_with_path("/v1/messages", b"hello")); - // Same body, different path -> not a duplicate. assert!(!dedup.is_duplicate_with_path("/v1/chat/completions", b"hello")); } @@ -529,20 +906,14 @@ mod tests { assert!(!dedup.is_duplicate_with_path("", b"hello")); } - /// Exercise the deprecated path-less shim so the public entry point cannot - /// break silently. `#[expect(deprecated)]` (rather than `#[allow]`) surfaces - /// an `unfulfilled_lint_expectation` warning the day the shim is removed. #[test] #[expect(deprecated, reason = "is_duplicate shim is exercised directly here")] fn dedup_path_less_shim_still_works() { let dedup = RequestDeduplicator::new(); assert!(!dedup.is_duplicate(b"hello")); - // Identical body seen twice under the empty path is still a duplicate. assert!(dedup.is_duplicate(b"hello")); } - // -- get_client_ip --------------------------------------------------------- - #[test] fn client_ip_trust_forwarded_headers_false_ignores_headers() { use axum::http::HeaderMap; @@ -580,9 +951,49 @@ mod tests { } } + #[test] + fn would_allow_is_false_once_the_peer_bucket_is_drained() { + let limiter = RateLimiter::new(1); + assert!( + limiter.is_allowed("10.0.0.9"), + "first request consumes the only token" + ); + assert!( + !limiter.would_allow("10.0.0.9"), + "drained peer must be denied by the peek" + ); + assert!(limiter.would_allow("10.0.0.8")); + } + + #[test] + fn would_allow_reports_true_for_unknown_peer_and_disabled_limiter() { + let limiter = RateLimiter::new(100); + assert!( + limiter.would_allow("10.0.0.7"), + "unknown peer starts with a full bucket" + ); + + let disabled = RateLimiter::new(0); + for _ in 0..50 { + assert!(disabled.is_allowed("10.0.0.7")); + } + assert!( + disabled.would_allow("10.0.0.7"), + "a disabled limiter never denies" + ); + } + + #[test] + fn would_allow_does_not_consume_tokens() { + let limiter = RateLimiter::new(1); + assert!(limiter.is_allowed("10.0.0.6")); + assert!(!limiter.would_allow("10.0.0.6")); + assert!(!limiter.would_allow("10.0.0.6")); + assert!(!limiter.is_allowed("10.0.0.6")); + } + #[test] fn rate_limiter_rejects_requests_at_the_limit() { - // With 1 RPM, the first request should be allowed and the second rejected. let limiter = RateLimiter::new(1); assert!( limiter.is_allowed("10.0.0.1"), @@ -597,8 +1008,7 @@ mod tests { #[test] fn dedup_large_body_is_not_tracked() { let dedup = RequestDeduplicator::with_window_ms(5000); - let large_body = vec![b'x'; 65 * 1024]; // 65 KiB > 64 KiB cap - // Large bodies should never be tracked as duplicates. + let large_body = vec![b'x'; 65 * 1024]; assert!(!dedup.is_duplicate_with_path("/v1/messages", &large_body)); assert!( !dedup.is_duplicate_with_path("/v1/messages", &large_body), @@ -609,7 +1019,7 @@ mod tests { #[test] fn dedup_body_at_cap_is_still_tracked() { let dedup = RequestDeduplicator::with_window_ms(5000); - let body_at_cap = vec![b'x'; 64 * 1024]; // Exactly at 64 KiB cap + let body_at_cap = vec![b'x'; 64 * 1024]; assert!(!dedup.is_duplicate_with_path("/v1/messages", &body_at_cap)); assert!( dedup.is_duplicate_with_path("/v1/messages", &body_at_cap), @@ -617,8 +1027,6 @@ mod tests { ); } - // -- Multi-value X-Forwarded-For (finding 24) ------------------------------ - #[test] fn client_ip_multi_value_xff_picks_leftmost() { use axum::http::HeaderMap; @@ -635,8 +1043,6 @@ mod tests { ); } - // -- get_client_ip returns 'unknown' (finding 26) -------------------------- - #[test] fn client_ip_returns_unknown_when_no_connect_info_and_no_headers() { use axum::http::HeaderMap; diff --git a/crates/llm-proxy-server/src/routes/chat.rs b/crates/llm-proxy-server/src/routes/chat.rs index 07d620e..0658a16 100644 --- a/crates/llm-proxy-server/src/routes/chat.rs +++ b/crates/llm-proxy-server/src/routes/chat.rs @@ -1,51 +1,27 @@ -//! `/providers/{provider}/v1/chat/completions` handler using the core pipeline. -//! -//! The handler parses the incoming OpenAI `ChatCompletionRequest`, decodes it -//! into a `CoreRequest` via the OpenAI Chat -//! client adapter, and then dispatches through the shared core pipeline -//! ([`handle_core_once`] or [`handle_core_stream`]). -//! -//! Architecture boundary: route handlers do not perform scenario detection, -//! endpoint classification, fallback routing, or provider-specific streaming. - use std::sync::Arc; use axum::Json; use axum::body::Body; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use llm_proxy_core::ProviderRouteKind; use llm_proxy_protocol::client::openai_chat; use llm_proxy_protocol::openai::ChatCompletionRequest; use tracing::{info, warn}; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; use super::error_response::{ClientProtocol, RouteError, route_error_response}; - -/// POST `/providers/{provider}/v1/chat/completions` -/// -/// Accepts an OpenAI-format [`ChatCompletionRequest`], decodes it through the -/// core pipeline, and returns an OpenAI-shaped response. pub(crate) async fn handle_chat_completions( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { - // The wrapper only renders: it does NOT emit ResponseFailed. Each EARLY-GATE - // error return inside the inner function emits exactly one ResponseFailed - // itself (the sites before the core pipeline dispatch), and the core pipeline - // emits its own single ResponseFailed for failures that occur inside it. A - // blanket emit here would double-count every pipeline failure (the pipeline - // emits AND returns Err via `?`), so the wrapper must stay emit-free - // (audit route-responsefailed-gaps regression). - match handle_chat_completions_inner(state, req_id, provider, connect_info, headers, body).await - { + match handle_chat_completions_inner(state, meta, provider, headers, body).await { Ok(response) => response, Err(error) => { warn!(error = %error, "request failed"); @@ -53,30 +29,16 @@ pub(crate) async fn handle_chat_completions( } } } - -/// Inner handler that returns `Result` so errors can be mapped uniformly. -/// -/// ResponseFailed emission discipline (audit route-responsefailed-gaps): -/// every EARLY-GATE error return (the sites BEFORE the core pipeline dispatch) -/// emits exactly one ResponseFailed here. Failures INSIDE the core pipeline -/// emit their own single ResponseFailed internally and return Err, which this -/// function propagates without re-emitting. The outer wrapper renders only. async fn handle_chat_completions_inner( state: AppState, - req_id: RequestId, + meta: RequestMeta, provider: String, - connect_info: Option, headers: HeaderMap, body: axum::body::Bytes, ) -> Result, RouteError> { let start = std::time::Instant::now(); let event_bus = Arc::clone(&state.event_bus); - // The request id is consumed by prepare_request below; snapshot it once so - // every early-gate emit references the same id. - let request_id = req_id.0.clone(); - // Every early gate emits exactly one ResponseFailed with the same - // provider/request/start triple; bind it once instead of repeating the - // six-argument call at each gate. + let request_id = meta.request_id.0.clone(); let emit_failed = |e: &RouteError| { core_pipeline::emit_response_failed( &event_bus, @@ -87,14 +49,8 @@ async fn handle_chat_completions_inner( start, ); }; - // Input validation gates, ordered cheapest-first to avoid charging - // rate-limit/dedup budget for malformed requests. Mirrors the ordering - // established in `token_count.rs`: provider name -> existence -> - // content-type -> rate-limit/dedup -> JSON parse. core_pipeline::validate_provider_name(&provider).inspect_err(&emit_failed)?; if state.providers().get(&provider).is_none() { - // Unknown-provider early gate: emit exactly one ResponseFailed before - // returning. provider is known from the path (audit route-responsefailed-gaps). let err = RouteError::UnknownProvider(provider.clone()); emit_failed(&err); return Err(err); @@ -102,20 +58,8 @@ async fn handle_chat_completions_inner( core_pipeline::validate_json_content_type(&headers).inspect_err(&emit_failed)?; let request_path = format!("/providers/{provider}/v1/chat/completions"); - let ctx = core_pipeline::prepare_request( - &state, - req_id.0, - &headers, - connect_info.as_ref(), - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; - - // Parse the OpenAI ChatCompletionRequest via axum's `Json` helper so the - // `JsonRejection` taxonomy (syntax vs data error) is preserved and mapped - // to a precise 400 body, rather than collapsing both into one opaque - // "invalid JSON" string. + let ctx = core_pipeline::prepare_request(&state, &meta, &headers, &body, &request_path) + .inspect_err(&emit_failed)?; let req: ChatCompletionRequest = match Json::::from_bytes(&body) { Ok(Json(value)) => value, Err(rejection) => { @@ -124,11 +68,6 @@ async fn handle_chat_completions_inner( return Err(err); } }; - - // Decode the OpenAI Chat request into a core request. - // Note: ChatCompletionRequest does not have a separate validate() method - // (unlike the Anthropic handler). All validation is performed inside - // decode_request: it checks for non-empty model and non-empty messages. let core = openai_chat::decode_request(req) .map_err(core_pipeline::protocol_error_to_route) .inspect_err(|e| emit_failed(e))?; @@ -141,38 +80,19 @@ async fn handle_chat_completions_inner( streaming = is_streaming, "decoded OpenAI Chat request into CoreRequest" ); - - // Extract the inbound client auth token (used by passthrough-auth providers). let inbound_auth = core_pipeline::extract_inbound_auth(&headers); + let dispatch = core_pipeline::CoreDispatch { + provider_name: &provider, + route_kind: ProviderRouteKind::ChatCompletions, + client_protocol: ClientProtocol::OpenAiChat, + }; if is_streaming { - core_pipeline::handle_core_stream( - state, - ctx, - &provider, - ProviderRouteKind::ChatCompletions, - core, - ClientProtocol::OpenAiChat, - inbound_auth, - ) - .await + core_pipeline::handle_core_stream(state, ctx, dispatch, core, inbound_auth).await } else { - core_pipeline::handle_core_once( - state, - ctx, - &provider, - ProviderRouteKind::ChatCompletions, - core, - ClientProtocol::OpenAiChat, - inbound_auth, - ) - .await + core_pipeline::handle_core_once(state, ctx, dispatch, core, inbound_auth).await } } - -/// Map an axum `JsonRejection` to a `RouteError`, preserving the distinction -/// between a syntactically invalid JSON body and a body that is valid JSON but -/// does not fit the target type. fn json_rejection_to_route_error(rejection: axum::extract::rejection::JsonRejection) -> RouteError { use axum::extract::rejection::JsonRejection; match rejection { @@ -185,20 +105,12 @@ fn json_rejection_to_route_error(rejection: axum::extract::rejection::JsonReject JsonRejection::MissingJsonContentType(_) => { RouteError::InvalidRequest("missing application/json Content-Type".to_owned()) } - // BytesRejection and any future non_exhaustive variant: surface a - // generic but informative invalid-request body. other => RouteError::InvalidRequest(format!("could not read request body: {other}")), } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { - // -- Architecture source guard --------------------------------------------- - #[test] fn source_guard_enforces_route_boundaries() { let source = include_str!("chat.rs"); diff --git a/crates/llm-proxy-server/src/routes/core_pipeline.rs b/crates/llm-proxy-server/src/routes/core_pipeline.rs index 731bc20..145f40f 100644 --- a/crates/llm-proxy-server/src/routes/core_pipeline.rs +++ b/crates/llm-proxy-server/src/routes/core_pipeline.rs @@ -1,10 +1,3 @@ -//! Shared core pipeline for route handlers. -//! -//! Provides [`prepare_request`] (rate-limit, dedup, request-ID generation), -//! [`handle_core_once`] (non-streaming), and [`handle_core_stream`] (streaming) -//! used by `/providers/{provider}/v1/messages` and -//! `/providers/{provider}/v1/chat/completions`. - use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; @@ -41,37 +34,16 @@ use super::error_response::{ AuthOwner, ClientProtocol, PROVIDER_DECODE_CLIENT_MESSAGE, RouteError, extract_error_fields, openai_stream_error_json_with_type, route_error_variant_name, truncate_with_suffix, }; - -// --------------------------------------------------------------------------- -// ClientStreamEncoder — protocol-agnostic stream encoder wrapper -// --------------------------------------------------------------------------- - -/// Protocol-agnostic wrapper around the Anthropic and OpenAI stream encoders. -/// -/// Both encoders share the same interface (`encode_event`, `finish`, -/// `mark_finished`) but produce different output types. This enum dispatches -/// to the correct encoder based on the client protocol. enum ClientStreamEncoder { Anthropic(AnthropicStreamEncoder), OpenAi(OpenAiStreamEncoder), } - -/// Output produced by a client stream encoder. Each variant carries the data -/// needed to construct an SSE `Event` for that protocol. enum ClientEncodedEvent { - /// Anthropic SSE: requires `event:` line (e.g. `event: message_start`) and - /// JSON data. The `event_type` is the SSE event name. Anthropic { event_type: String, json: String }, - /// OpenAI SSE: just `data: `. No event type field. The caller must - /// also emit `data: [DONE]` as the terminal frame after the stream ends. OpenAi { json: String }, } impl ClientStreamEncoder { - /// Create a new encoder for the given protocol. - /// - /// For Anthropic, `msg_id` and `model` are used directly. - /// For OpenAI, `include_usage` is extracted from `provider_hints` if present. fn new(protocol: ClientProtocol, msg_id: String, model: String, core: &CoreRequest) -> Self { match protocol { ClientProtocol::Anthropic => { @@ -85,12 +57,6 @@ impl ClientStreamEncoder { .and_then(|v| v.get("include_usage")) .and_then(|v| v.as_bool()) .unwrap_or(false); - // The u64 -> i64 cast is safe: u64::MAX corresponds to a date - // ~584 billion years in the future, unreachable in any real timeline. - // The unwrap_or_default() handles the unlikely case of a system - // clock set before UNIX_EPOCH (e.g. clock skew after boot) by - // falling back to 0, which produces a valid but incorrect timestamp. - // This is acceptable because `created` is informational only. let created = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() @@ -104,8 +70,6 @@ impl ClientStreamEncoder { } } } - - /// Encode a single [`CoreEvent`] into zero or more protocol-specific events. fn encode_event( &mut self, event: CoreEvent, @@ -121,8 +85,6 @@ impl ClientStreamEncoder { } } } - - /// Flush any remaining buffered events (synthetic terminal if needed). fn finish( &mut self, ) -> Result, llm_proxy_protocol::client::ProtocolError> { @@ -137,8 +99,6 @@ impl ClientStreamEncoder { } } } - - /// Mark the encoder as finished so subsequent `finish()` returns empty. fn mark_finished(&mut self) { match self { Self::Anthropic(enc) => enc.mark_finished(), @@ -146,10 +106,6 @@ impl ClientStreamEncoder { } } } - -/// Serialize Anthropic SSE events, dropping any that fail to serialize. -/// -/// `phase` is used in log messages (e.g. "encode" vs "finish") for context. fn wrap_anthropic_events( events: Vec, phase: &str, @@ -158,9 +114,6 @@ fn wrap_anthropic_events( .into_iter() .filter_map(|me| match serde_json::to_string(&me) { Ok(json) => { - // Serialize first (borrowing &me); the borrow ends when the - // owned String returns, so we can move r#type out of `me` - // instead of cloning it on the per-token encode hot path. let event_type = me.r#type; Some(ClientEncodedEvent::Anthropic { event_type, json }) } @@ -171,10 +124,6 @@ fn wrap_anthropic_events( }) .collect() } - -/// Serialize OpenAI SSE chunks, dropping any that fail to serialize. -/// -/// `phase` is used in log messages (e.g. "encode" vs "finish") for context. fn wrap_openai_events(chunks: Vec, phase: &str) -> Vec { chunks .into_iter() @@ -187,8 +136,6 @@ fn wrap_openai_events(chunks: Vec, phase: &str) -> Vec Event { match encoded { ClientEncodedEvent::Anthropic { event_type, json } => { @@ -197,64 +144,37 @@ fn client_event_to_sse(encoded: ClientEncodedEvent) -> Event { ClientEncodedEvent::OpenAi { json } => Event::default().data(json), } } - -/// Build the terminal `[DONE]` event for OpenAI Chat streams. fn openai_done_event() -> Event { Event::default().data("[DONE]") } - -// --------------------------------------------------------------------------- -// RequestContext -// --------------------------------------------------------------------------- - -/// Pre-resolved per-request context shared between non-streaming and streaming -/// paths. #[derive(Debug)] pub(crate) struct RequestContext { - /// Unique request ID for tracing and the `x-request-id` header. pub(crate) request_id: String, - /// Instant when the handler was entered (for latency metrics). pub(crate) start: Instant, - /// SHA-256 of the raw client request body, for the `RequestReceived.body_hash` - /// correlation key. Computed in [`prepare_request`] over the raw wire bytes - /// (where they first arrive), per the plan's locked decision, so it is shared - /// across the inference routes and `/v1/messages/count_tokens`. + pub(crate) caller: String, pub(crate) body_hash: String, } - -// --------------------------------------------------------------------------- -// prepare_request -// --------------------------------------------------------------------------- - -/// Perform pre-flight checks and build a [`RequestContext`]. -/// -/// Returns `Err(RouteError)` early for rate-limited or duplicate requests so -/// the handler can immediately produce an error response without entering the -/// core pipeline. -/// -/// `request_id` is minted by the outermost `inject_request_id` layer and passed -/// in by the handler so the id attached to the `x-request-id` header matches -/// the id stamped on the request's tracing span (audit MEDIUM-1). pub(crate) fn prepare_request( state: &AppState, - request_id: String, + meta: &crate::middleware::RequestMeta, headers: &HeaderMap, - connect_info: Option<&std::net::SocketAddr>, body: &[u8], path: &str, ) -> Result { + let caller = &meta.caller; + let connect_info = meta.connect_info.as_ref(); let trust_forwarded_headers = state.trust_forwarded_headers(); let client_ip = get_client_ip(headers, connect_info, trust_forwarded_headers); - // Security: loopback bypass is intentional for local development and testing. - // When `trust_forwarded_headers` is false, the connection-info IP is guaranteed - // to be the actual TCP peer (not spoofable via headers). This means only - // processes on the same machine can bypass rate limiting. When - // `trust_forwarded_headers` is true, the loopback check is disabled because - // the peer IP may be spoofed by an intermediate reverse proxy. - let direct_loopback = - !trust_forwarded_headers && connect_info.is_some_and(|address| address.ip().is_loopback()); - - if !direct_loopback && !state.rate_limiter.is_allowed(&client_ip) { + let direct_loopback = !trust_forwarded_headers + && !caller.is_authenticated() + && connect_info.is_some_and(|address| address.ip().is_loopback()); + + let rate_limit_key = if caller.is_authenticated() { + caller.as_str().to_owned() + } else { + client_ip + }; + if !direct_loopback && !state.rate_limiter.is_allowed(&rate_limit_key) { state.metrics.record_rate_limited(); return Err(RouteError::RateLimited); } @@ -265,32 +185,16 @@ pub(crate) fn prepare_request( } Ok(RequestContext { - request_id, + request_id: meta.request_id.0.clone(), start: Instant::now(), + caller: caller.as_str().to_owned(), body_hash: body_hash_of_bytes(body), }) } - -// --------------------------------------------------------------------------- -// resolve target helper -// --------------------------------------------------------------------------- - -/// Maximum length for a provider name in URL path parameters. -/// -/// Prevents regex DoS from arbitrarily long names. 128 characters is generous -/// enough for any realistic provider identifier. const MAX_PROVIDER_NAME_LEN: usize = 128; - -/// Regex for validating provider names in URL paths. -/// -/// Provider names must be ASCII lowercase letters, digits, hyphens, and -/// underscores. This ensures URL-safe slugs that cannot be confused with path -/// traversal or encoding attacks. static PROVIDER_NAME_REGEX: std::sync::LazyLock = std::sync::LazyLock::new(|| { regex::Regex::new(r"^[a-z0-9_-]+$").expect("valid provider name regex") }); - -/// Validate that a provider name from a URL path contains only safe characters. pub(super) fn validate_provider_name(name: &str) -> Result<(), RouteError> { if name.is_empty() { return Err(RouteError::InvalidProviderName( @@ -310,16 +214,6 @@ pub(super) fn validate_provider_name(name: &str) -> Result<(), RouteError> { } Ok(()) } - -/// Reject a request whose `Content-Type` is not JSON before attempting to parse -/// the body, so a non-JSON body yields a precise "expected application/json" -/// 400 rather than a misleading "invalid JSON" parse error (audit LOW-30). -/// -/// Handlers accept `axum::body::Bytes` (so the raw body can be hashed for -/// dedup), which bypasses the `Json` extractor's built-in content-type check. -/// This restores the spirit of that check uniformly across every JSON POST -/// handler (`chat`, `messages`, `token_count`). A missing `Content-Type` is -/// tolerated for backwards compatibility, matching the prior behaviour. pub(super) fn validate_json_content_type(headers: &HeaderMap) -> Result<(), RouteError> { let Some(value) = headers.get(header::CONTENT_TYPE) else { return Ok(()); @@ -327,8 +221,6 @@ pub(super) fn validate_json_content_type(headers: &HeaderMap) -> Result<(), Rout let Ok(ct) = value.to_str() else { return Ok(()); }; - // Accept `application/json` and any `+json` suffix (e.g. - // `application/vnd.api+json`). Strip any `; charset=...` parameters first. let essence = ct .split(';') .next() @@ -343,28 +235,10 @@ pub(super) fn validate_json_content_type(headers: &HeaderMap) -> Result<(), Rout } Ok(()) } - -/// Validate a provider name extracted from the URL path. -/// -/// This is a thin wrapper around [`validate_provider_name`] used at handler -/// entry points as defense-in-depth. Axum's `Path` extractor does not -/// perform character validation, so this ensures the provider name is safe -/// before it is used in format strings or lookups. fn validate_provider_name_ref(provider_name: &str) -> Result<&str, RouteError> { validate_provider_name(provider_name)?; Ok(provider_name) } - -/// Resolve a provider route through the provider registry into a tuple of -/// (`ProviderAdapterTarget`, `ProviderAdapter`) ready for encoding. -/// -/// Lookup chain: -/// -/// ```text -/// provider_name + route_kind -> provider config + route -> adapter name -/// adapter_name -> protocol + endpoint + headers -/// requested_model -> model_aliases -> upstream_model -/// ``` async fn resolve_target( state: &AppState, provider_name: &str, @@ -377,18 +251,12 @@ async fn resolve_target( let adapter_target_config = providers .resolve_provider_route(provider_name, route_kind, &core.model.requested) .map_err(map_provider_route_error)?; - - // Catalog enforcement: fetch the provider once and check both the - // enforcement flag and the model catalog in a single lookup. if let Some(provider) = providers.get(provider_name) { if provider .catalog .as_ref() .is_some_and(|catalog| catalog.enforce) { - // O(1) membership check backed by a cached id index inside - // ModelCatalogService (audit LOW-9); replaces a per-request O(N) - // linear scan over the merged catalog. let allowed = state .model_catalogs() .contains_model( @@ -422,15 +290,7 @@ async fn resolve_target( adapter_target_config.protocol )) })? - // ProviderAdapter variants are zero-sized unit structs; clone is a - // trivial copy. (The Arc sits one level out, on - // `provider_adapters: Arc`; `get()` returns a - // borrowed `&ProviderAdapter`, and this `.clone()` copies only the - // zero-sized enum value out from under the borrow — no refcount bump.) .clone(); - - // Resolve the upstream api_key: passthrough-auth providers use the client's - // inbound token (401 if absent); others use the configured static key. let api_key = if providers .get(provider_name) .is_some_and(|p| p.passthrough_auth) @@ -469,55 +329,24 @@ fn map_provider_route_error(error: ProviderRouteResolutionError) -> RouteError { error => RouteError::Internal(error.to_string()), } } - -// --------------------------------------------------------------------------- -// handle_core_once -// --------------------------------------------------------------------------- - -/// Canonical string for a [`ProviderRouteKind`] (used in the event log). fn route_kind_str(kind: ProviderRouteKind) -> &'static str { match kind { ProviderRouteKind::ChatCompletions => "chat_completions", ProviderRouteKind::Messages => "messages", } } - -/// Canonical string for a [`ClientProtocol`] (used in the event log). fn client_protocol_str(protocol: ClientProtocol) -> &'static str { match protocol { ClientProtocol::OpenAiChat => "openai_chat", ClientProtocol::Anthropic => "anthropic", } } - -/// Stable correlation hash of the RAW client request body, hex-encoded. -/// -/// SHA-256 of the raw wire bytes (`&[u8]`) exactly as the client sent them. The -/// `/v1/messages/count_tokens` route uses the same input, so a request has the -/// same `body_hash` correlation key across every route that emits a -/// `RequestReceived`. The raw body itself is never stored -- only this digest. -/// -/// Shared with `token_count.rs` so the two routes agree on the correlation key -/// (audit finding: body_hash was hashed over the normalized `CoreRequest` on the -/// inference routes but raw bytes on token_count, breaking cross-route -/// correlation). pub(crate) fn body_hash_of_bytes(body: &[u8]) -> String { use sha2::{Digest, Sha256}; let mut hasher = Sha256::new(); hasher.update(body); format!("{:x}", hasher.finalize()) } - -/// Emit a structured `ResponseFailed` event for a route error. Borrows the -/// error so the caller can still return it; derives HTTP status + client-facing -/// error kind from the canonical [`extract_error_fields`] mapping. -/// -/// When the caller does not supply a `provider` (the early-gate failure path, -/// where the request never reached a resolved provider), the provider name is -/// recovered from [`RouteError::UnknownProvider`] so the event still records -/// which provider name was probed. This keeps the unknown-provider failure -/// attributable in the event log without each early-gate caller having to -/// thread the name through separately. pub(crate) fn emit_response_failed( event_bus: &Arc, request_id: &str, @@ -531,11 +360,6 @@ pub(crate) fn emit_response_failed( _ => None, }); let (status, _client_error_type, message) = extract_error_fields(error); - // Log the (already key/URL-redacted) upstream body exactly once here. This - // is the single canonical failure-recording site, guaranteed to run once - // per failure; emitting it inside `extract_error_fields` instead would - // double-warn because that helper is also called from `route_error_response` - // on the render path. if let RouteError::Upstream { status: upstream_status, body, @@ -550,10 +374,6 @@ pub(crate) fn emit_response_failed( "upstream error mapped for the client; real body logged here only" ); } - // `error_kind` records the Rust variant name (the documented event-log - // discriminant), NOT the client-facing protocol error-type string, so - // distinct variants stay separable in the event log (audit finding: - // error_kind not variant name). let error_kind = route_error_variant_name(error); event_bus.emit(&ProxyEvent::ResponseFailed(ResponseFailed { request_id: request_id.to_owned(), @@ -566,18 +386,12 @@ pub(crate) fn emit_response_failed( latency_ms: start.elapsed().as_millis().try_into().unwrap_or(u64::MAX), })); } - -/// Capture lifecycle fields from a [`CoreEvent`] into the stream context -/// (upstream message id, cumulative usage, stop reason). Called on every event -/// before it is consumed by `encode_core_event`, so the values survive to the -/// stream-completion `ResponseCompleted` emission. fn capture_lifecycle(ctx: &mut StreamContext, event: &CoreEvent) { match event { CoreEvent::MessageStart { id: Some(id), .. } => { ctx.upstream_message_id = Some(id.clone()); } CoreEvent::UsageDelta { usage } => { - // Providers send a cumulative running total, so replace (do not sum). ctx.pending_usage = Some(usage.clone()); } CoreEvent::MessageStop { stop_reason, .. } => { @@ -586,18 +400,9 @@ fn capture_lifecycle(ctx: &mut StreamContext, event: &CoreEvent) { _ => {} } } - -/// Extract the client's inbound auth token from request headers, for passthrough -/// providers. Checks `Authorization: Bearer `, then `x-api-key`, then -/// `x-goog-api-key`. The raw token is wrapped in [`SecretString`] so it is -/// redacted in logs/Debug like any other credential. pub(crate) fn extract_inbound_auth(headers: &HeaderMap) -> Option { if let Some(value) = headers.get(header::AUTHORIZATION) { if let Ok(s) = value.to_str() { - // The auth scheme is case-insensitive (RFC 7235): "Bearer", - // "bearer", "BEARER", "bEaReR" are all valid. Split on the first - // space to separate the scheme from the token and compare the - // scheme case-insensitively, rather than matching only two casings. let mut parts = s.splitn(2, ' '); match (parts.next(), parts.next()) { (Some(scheme), Some(token)) if scheme.eq_ignore_ascii_case("bearer") => { @@ -621,40 +426,13 @@ pub(crate) fn extract_inbound_auth(headers: &HeaderMap) -> Option } None } - -/// Compute the USD cost of a request from its [`Usage`] and a per-token price -/// table. -/// -/// Pure domain math. Lives in the server crate (not on the protocol [`Cost`] -/// type) because [`Usage`] is in the protocol crate and [`ModelPricing`] in the -/// core crate, and the server is the first crate that depends on both. -/// -/// # Overflow safety -/// -/// `rust_decimal`'s `*`/`+` operators PANIC on overflow. Token counts (up to -/// ~2.1e9) multiplied by a price, then summed across five components, can -/// exceed `Decimal`'s 96-bit range, which would panic the request handler. We -/// use checked arithmetic and saturate to `Decimal::MAX` instead, so a bad -/// (or just very large) price/config can never crash a cost-bearing request. -/// Config validation also bounds prices (see `MAX_PER_TOKEN_PRICE`); the -/// checked math here is the belt-and-suspenders guard. -/// -/// # Rounding -/// -/// Per the plan's locked decision 5 ("rounded to the nearest microcent on -/// output only"), each component is rounded to 6 decimal places (a microcent -/// of USD) after the multiply. The total is the checked sum of the rounded -/// components, so it is always exactly their sum (no separate rounding drift). fn compute_cost(usage: &Usage, pricing: &ModelPricing) -> Cost { - /// Multiply a token count by a per-token price, saturating on overflow, - /// then round to the nearest microcent (6 dp). fn priced(count: i32, price: Decimal) -> Decimal { Decimal::from(count) .checked_mul(price) .unwrap_or(Decimal::MAX) .round_dp(6) } - /// Checked sum that saturates to `Decimal::MAX` on overflow. fn sum(a: Decimal, b: Decimal) -> Decimal { a.checked_add(b).unwrap_or(Decimal::MAX) } @@ -683,30 +461,34 @@ fn compute_cost(usage: &Usage, pricing: &ModelPricing) -> Cost { total, } } - -/// Non-streaming core pipeline. -/// -/// ```text -/// CoreRequest -> encode -> send -> decode -> client encode -> HTTP response -/// ``` +#[derive(Debug, Clone, Copy)] +pub(crate) struct CoreDispatch<'a> { + pub(crate) provider_name: &'a str, + pub(crate) route_kind: ProviderRouteKind, + pub(crate) client_protocol: ClientProtocol, +} +#[expect( + clippy::too_many_lines, + reason = "the pipeline stages run in request order and each stage's \ + metric + ResponseFailed emission sits next to the call that can \ + fail; extracting stages would separate the failure bookkeeping \ + from the step it belongs to" +)] pub(crate) async fn handle_core_once( state: AppState, ctx: RequestContext, - provider_name: &str, - route_kind: ProviderRouteKind, + dispatch: CoreDispatch<'_>, core: CoreRequest, - client_protocol: ClientProtocol, inbound_auth: Option, ) -> Result, RouteError> { + let CoreDispatch { + provider_name, + route_kind, + client_protocol, + } = dispatch; validate_provider_name_ref(provider_name)?; state.metrics.record_request(false); - - // Whether the client's own token is forwarded to the upstream - // (passthrough_auth). Captured once here so error classification - // (map_provider_error) can mark auth/permission failures as client-owned - // and pass 401/403 through verbatim instead of collapsing to 502 - // (audit LOW, passthrough-auth-401-collapse). let passthrough_auth = state .providers() .get(provider_name) @@ -719,13 +501,12 @@ pub(crate) async fn handle_core_once( streaming = false, "processing request" ); - - // Emit a structured RequestReceived event. state .event_bus .emit(&ProxyEvent::RequestReceived(RequestReceived { request_id: ctx.request_id.clone(), timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), provider: provider_name.to_owned(), route_kind: route_kind_str(route_kind).to_owned(), client_protocol: client_protocol_str(client_protocol).to_owned(), @@ -760,11 +541,6 @@ pub(crate) async fn handle_core_once( upstream_model = %target.upstream_model, "routed to provider" ); - - // Encode the core request into a provider-specific HTTP request. - // All encode errors map to Internal per the plan's error behavior spec: - // the core pipeline has already validated the request at this point, so - // any encode failure is a proxy/adapter issue, not a client error. let proxy_req: ProxyRequest = adapter.encode_request(&core, &target).map_err(|e| { let err = RouteError::Internal(format!("encode error: {e}")); state.metrics.record_failure(); @@ -779,21 +555,23 @@ pub(crate) async fn handle_core_once( err })?; - let response_bytes = state.proxy_client.send(proxy_req).await.map_err(|e| { - let err = map_provider_error(&e, passthrough_auth); - state.metrics.record_failure(); - emit_response_failed( - &state.event_bus, - &ctx.request_id, - Some(provider_name), - Some(&core.model), - &err, - ctx.start, - ); - err - })?; - - // Decode the provider response into a CoreResponse. + let response_bytes = state + .upstream_client_for(&target.provider_name) + .send(proxy_req) + .await + .map_err(|e| { + let err = map_provider_error(&e, passthrough_auth); + state.metrics.record_failure(); + emit_response_failed( + &state.event_bus, + &ctx.request_id, + Some(provider_name), + Some(&core.model), + &err, + ctx.start, + ); + err + })?; let mut core_resp = adapter .decode_response(&response_bytes, &target) .map_err(|e| { @@ -809,40 +587,21 @@ pub(crate) async fn handle_core_once( ); err })?; - - // Encode into the client-specific response. let latency = ctx.start.elapsed(); - - // Attach computed cost when pricing is configured for the (alias-resolved) - // upstream model. The price is cloned out of the registry before the - // `core_resp` move into the client encoder below. let pricing = state .providers() .pricing_for(&target.provider_name, &target.upstream_model) .cloned(); - // Only attach a cost when the usage was genuinely reported by the provider. - // Synthetic/Unknown usage (zero-filled because the provider omitted it) - // would otherwise produce a Some(zero-Cost) indistinguishable from a truly - // free request; leaving cost None for unprovenanced usage keeps the event - // honest (audit INFO, cost-on-synthetic-zero). let cost = if core_resp.usage.provenance == UsageProvenance::ProviderReported { pricing.map(|p| compute_cost(&core_resp.usage, &p)) } else { None }; core_resp.cost = cost.clone(); - // Capture the ResponseCompleted payload before `core_resp` is moved into the - // client encoder; the event is emitted AFTER a successful encode so we never - // log "completed" for a request that then fails client encoding. let resp_usage = core_resp.usage.clone(); let resp_model = core_resp.model.clone(); let resp_upstream_id = core_resp.id.clone(); let resp_stop_reason = core_resp.stop_reason.clone(); - - // Client-encode failures map to Internal. They must record a failure metric - // AND emit a ResponseFailed event (mirroring the decode/send arms above), - // and must NOT be counted as a success. `record_success` is therefore - // deferred until AFTER this match succeeds (audit LOW, eventlog-04). let response_body = match client_protocol { ClientProtocol::Anthropic => { let msg_resp = anthropic::encode_response(core_resp).map_err(|e| { @@ -901,16 +660,9 @@ pub(crate) async fn handle_core_once( })? } }; - - // Client encoding succeeded: record the success metric now (after the - // encode match above), so an encode failure is never counted as success - // (audit LOW, eventlog-04). The stream path already defers this to - // `stream_succeeded`. state .metrics .record_success(&target.provider_name, &target.upstream_model, latency); - - // Client encoding succeeded: emit the ResponseCompleted event. state .event_bus .emit(&ProxyEvent::ResponseCompleted(ResponseCompleted { @@ -933,17 +685,12 @@ pub(crate) async fn handle_core_once( ); let mut response = (StatusCode::OK, Body::from(response_body)).into_response(); - - // Insert the request ID header safely -- the ID is dynamically generated - // so we use from_str with a fallback rather than expect/unwrap. response.headers_mut().insert( "x-request-id", ctx.request_id .parse() .unwrap_or_else(|_| axum::http::HeaderValue::from_static("unknown")), ); - - // Set content-type explicitly (defense-in-depth). response.headers_mut().insert( header::CONTENT_TYPE, axum::http::HeaderValue::from_static("application/json"), @@ -951,65 +698,32 @@ pub(crate) async fn handle_core_once( Ok(response) } - -// --------------------------------------------------------------------------- -// handle_core_stream -// --------------------------------------------------------------------------- - -/// Heartbeat interval for SSE streams. -/// -/// Chosen as 3 seconds to keep Anthropic SDK clients from timing out on -/// long-running streams while avoiding excessive bandwidth overhead. -/// TODO(future): Make this configurable via AppState or TOML config. const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(3); - -/// Message sent from the spawned stream task to the handler during the -/// first-byte probe phase. This enables the handler to return a proper -/// HTTP error status if the stream fails before emitting any data. enum FirstByteResult { - /// At least one SSE event was produced by the stream. The stream is now - /// live and all subsequent events/errors must be delivered as in-band - /// SSE events (the HTTP status is already committed as 200 OK). FirstEvent(Event), - /// The stream failed before producing any event. The handler should - /// return an HTTP error response (e.g. 502 Bad Gateway) instead of - /// committing a 200 OK SSE response. PreStreamError(RouteError), } - -/// Streaming core pipeline. -/// -/// ```text -/// CoreRequest -> encode -> send_stream -> SseFramer -> ProviderStreamDecoder -/// -> StreamEncoder -> SSE response -/// ``` -/// -/// ## first_byte_sent tracking -/// -/// The plan requires that errors before the first data byte be returned as -/// HTTP errors (e.g. 502), not as in-band SSE error events. This is -/// implemented by probing for the first event before committing the HTTP 200 -/// response: the spawned task sends the first event (or an error) back to -/// the handler via a oneshot channel, and the handler decides whether to -/// return an HTTP error or an SSE response. +#[expect( + clippy::too_many_lines, + reason = "mirrors handle_core_once for the streaming path: stage order is \ + the audit trail and each stage's failure bookkeeping stays next \ + to the failing call" +)] pub(crate) async fn handle_core_stream( state: AppState, ctx: RequestContext, - provider_name: &str, - route_kind: ProviderRouteKind, + dispatch: CoreDispatch<'_>, core: CoreRequest, - client_protocol: ClientProtocol, inbound_auth: Option, ) -> Result, RouteError> { + let CoreDispatch { + provider_name, + route_kind, + client_protocol, + } = dispatch; validate_provider_name_ref(provider_name)?; state.metrics.record_request(true); - - // Whether the client's own token is forwarded to the upstream - // (passthrough_auth). Captured once here so error classification - // (map_provider_error, including inside the spawned stream task) can mark - // auth/permission failures as client-owned and pass 401/403 through verbatim - // instead of collapsing to 502 (audit LOW, passthrough-auth-401-collapse). let passthrough_auth = state .providers() .get(provider_name) @@ -1022,17 +736,12 @@ pub(crate) async fn handle_core_stream( streaming = true, "processing streaming request" ); - - // Emit a structured RequestReceived event BEFORE resolve_target/send_stream - // so a failure in either step still has a preceding RequestReceived (the - // non-stream path emits at the same point). Without this, an early-return - // from resolve_target/send_stream would emit an orphaned ResponseFailed with - // no matching RequestReceived (audit LOW, eventlog-03). state .event_bus .emit(&ProxyEvent::RequestReceived(RequestReceived { request_id: ctx.request_id.clone(), timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), provider: provider_name.to_owned(), route_kind: route_kind_str(route_kind).to_owned(), client_protocol: client_protocol_str(client_protocol).to_owned(), @@ -1067,10 +776,6 @@ pub(crate) async fn handle_core_stream( upstream_model = %target.upstream_model, "routed to provider (streaming)" ); - - // Encode the core request into a provider-specific HTTP request. - // All encode errors map to Internal per the plan's error behavior spec: - // the core pipeline has already validated the request at this point. let proxy_req: ProxyRequest = adapter.encode_request(&core, &target).map_err(|e| { let err = RouteError::Internal(format!("encode error: {e}")); state.metrics.record_failure(); @@ -1084,10 +789,8 @@ pub(crate) async fn handle_core_stream( ); err })?; - - // Open the streaming connection. let byte_stream = state - .proxy_client + .upstream_client_for(&target.provider_name) .send_stream(proxy_req) .await .map_err(|e| { @@ -1106,41 +809,16 @@ pub(crate) async fn handle_core_stream( let provider_decoder = adapter.new_stream_decoder(&target); let sse_framer = SseFramer::new(); - - // The client stream encoder is constructed *inside* the spawned stream - // task (see `build_sse_output_stream`), after buffering the first decoded - // `CoreEvent`. This lets the encoder be seeded with the upstream provider's - // real message ID (from `CoreEvent::MessageStart`) instead of a synthetic - // one, so clients tracking message IDs see the upstream's own ID rather than - // a proxy-generated one. - - // ctx is consumed after this point. request_id is cloned once for the - // spawned task and once for the response header (both are needed). let request_id = ctx.request_id; let upstream_model = target.upstream_model.clone(); - - // Pricing for the resolved upstream model (None => no cost on the event). let pricing = state .providers() .pricing_for(&target.provider_name, &target.upstream_model) .cloned(); let event_bus = Arc::clone(&state.event_bus); - - // Capture fields needed for a stream ResponseFailed event, before `core` - // and `ctx.start` are moved into the StreamContext below. let requested_model = core.model.clone(); let stream_start = ctx.start; - - // Build the output SSE stream with first-byte tracking. let (first_byte_tx, first_byte_rx) = tokio::sync::oneshot::channel::(); - // Shared "stream failure already recorded" guard. A pre-first-byte panic is - // observed by TWO independent watchers: this handler (the `first_byte_rx` - // `RecvError` path, since the panicked task drops the sender) AND the - // detached supervisor (the `JoinError` path). Without coordination each - // would emit its own `ResponseFailed` and bump the failure metric for the - // same request. The first watcher to run claims the flag via `swap`; the - // other skips. Exactly one terminal event + one failure metric per request - // (audit finding: stream-task panic double-ResponseFailed). let failure_recorded = Arc::new(AtomicBool::new(false)); let output_stream = build_sse_output_stream( byte_stream, @@ -1167,16 +845,10 @@ pub(crate) async fn handle_core_stream( }, Arc::clone(&failure_recorded), ); - - // Wait for the first event (or a pre-stream error). This is the - // first_byte_sent boundary: errors before this point become HTTP - // errors; errors after this point become in-band SSE error events. let first_event = first_byte_rx.await.map_err(|error| { let err = RouteError::Internal(format!( "stream task exited before first event ({error}: possible panic, cancellation, or empty stream)" )); - // Only record/emit if the supervisor has not already done so for this - // request (see `failure_recorded` above). if !failure_recorded.swap(true, Ordering::SeqCst) { state.metrics.record_failure(); emit_response_failed( @@ -1193,9 +865,6 @@ pub(crate) async fn handle_core_stream( match first_event { FirstByteResult::PreStreamError(route_error) => { - // Stream failed before emitting any data. Emit a ResponseFailed - // event (covers every in-task pre-stream failure funnelled through - // the first-byte channel), then return as an HTTP error. emit_response_failed( &state.event_bus, &request_id, @@ -1207,8 +876,6 @@ pub(crate) async fn handle_core_stream( Err(route_error) } FirstByteResult::FirstEvent(event) => { - // First event received. Commit HTTP 200 and start streaming. - // Prepend the first event to the output stream. let stream = futures::stream::once(async move { Ok::<_, std::convert::Infallible>(event) }) .chain(output_stream.map(Ok::<_, std::convert::Infallible>)); @@ -1217,23 +884,14 @@ pub(crate) async fn handle_core_stream( let response = sse.into_response(); let (mut parts, body) = response.into_parts(); - - // Add custom headers. - // Defense-in-depth: explicitly set Content-Type even though - // Sse::into_response() sets it internally. This guards against - // axum version changes and makes the intent explicit. parts.headers.insert( header::CONTENT_TYPE, axum::http::HeaderValue::from_static("text/event-stream"), ); - // Prevent proxy/CDN caching of streaming responses. parts.headers.insert( header::CACHE_CONTROL, axum::http::HeaderValue::from_static("no-cache"), ); - // Defense-in-depth: explicitly request persistent connection for - // reverse-proxy scenarios. axum adds this automatically for HTTP/1.1 - // but some reverse proxies strip it unless explicitly set. parts.headers.insert( header::CONNECTION, axum::http::HeaderValue::from_static("keep-alive"), @@ -1259,68 +917,31 @@ pub(crate) async fn handle_core_stream( } } } - -// --------------------------------------------------------------------------- -// SSE output stream builder -// --------------------------------------------------------------------------- - -/// Context for tracking metrics during a streaming response. #[derive(Debug)] struct StreamMetrics { - /// Shared metrics recorder for success/failure/latency tracking. metrics: Arc, - /// Provider name used as part of the composite metrics key. provider_name: String, - /// Upstream model name used as part of the composite metrics key. upstream_model: String, - /// Instant when the handler was entered (for latency measurement). start: Instant, } - -/// Mutable context used inside the spawned stream task. -/// -/// Groups the state that was previously passed as individual parameters to -/// `build_sse_output_stream`, reducing the function's parameter count from 9 -/// to a single struct plus the byte stream. struct StreamContext { provider_decoder: ProviderStreamDecoderKind, sse_framer: SseFramer, client_protocol: ClientProtocol, - /// The normalized core request. Held so the client stream encoder can be - /// constructed inside the stream task -- after buffering the first event to - /// extract the upstream message ID -- with access to `model.requested` and - /// provider hints. core: CoreRequest, request_id: String, stream_metrics: StreamMetrics, first_byte_tx: Option>, first_byte_sent: bool, - /// Per-token pricing for the resolved upstream model, for stream cost. pricing: Option, - /// Whether the client's own token was forwarded (passthrough_auth). Drives - /// 401/403 pass-through vs 502 collapse in in-band/pre-stream error - /// classification (audit LOW, passthrough-auth-401-collapse). passthrough_auth: bool, - /// Event sink handle (the task emits ResponseCompleted at stream end). event_bus: Arc, - /// Latest cumulative usage seen via `CoreEvent::UsageDelta` (replaced, not - /// summed -- providers send a running total). pending_usage: Option, - /// Stop reason captured from `CoreEvent::MessageStop`. stop_reason: Option, - /// Upstream message id captured from `CoreEvent::MessageStart`. upstream_message_id: Option, } impl StreamContext { - /// Emit a single SSE event, handling the first-byte boundary. - /// - /// Before the first event is sent, events go through the `first_byte_tx` - /// channel so the handler can return a proper HTTP status. After the first - /// event, events go directly to the `tx` channel. - /// - /// Returns `true` if the event was successfully delivered (or the stream - /// should continue), `false` if the task should exit. async fn emit_event(&mut self, event: Event, tx: &tokio::sync::mpsc::Sender) -> bool { if !self.first_byte_sent { if let Some(fb_tx) = self.first_byte_tx.take() { @@ -1328,24 +949,9 @@ impl StreamContext { self.first_byte_sent = true; return true; } - // Receiver dropped -- handler is gone (client disconnected or - // handler timed out). The failure is unactionable because the - // handler has already returned, so we just record the metric - // and stop the stream task. - // - // Client-initiated disconnect, not an upstream failure: - // record a cancellation so it does not inflate the error-rate - // SLO (audit MEDIUM-5). self.stream_metrics.metrics.record_client_cancel(); return false; } - // Unreachable today: `first_byte_tx` is consumed only by a - // successful first-event send (which sets `first_byte_sent`) or by - // `send_pre_stream_error`, and every `send_pre_stream_error` call - // site is immediately followed by a `return`/`break`. Guard it - // explicitly so a future refactor cannot silently drop a - // client-visible event while reporting success: stop the task - // loudly instead of dropping the event and returning `true`. warn!( request_id = %self.request_id, "emit_event: first-byte channel already consumed but no first \ @@ -1353,23 +959,12 @@ impl StreamContext { ); return false; } else if tx.send(event).await.is_err() { - // Channel receiver dropped -- the SSE handler task has exited - // (client disconnect or timeout). No further events can be - // delivered, so stop the stream task. - // - // Client disconnect, not an upstream failure (audit MEDIUM-5). Emit - // a terminal `ClientCancelled` event so this `RequestReceived` is - // not left orphaned in the event log (audit finding: buffering vs - // main-loop disconnect event inconsistency -- a pre-first-byte - // disconnect already emits via the PreStreamError path). self.stream_metrics.metrics.record_client_cancel(); self.emit_client_cancelled(); return false; } true } - - /// Emit a pre-stream error through the first-byte channel. fn send_pre_stream_error(&mut self, error: RouteError) -> bool { if let Some(fb_tx) = self.first_byte_tx.take() { fb_tx.send(FirstByteResult::PreStreamError(error)).is_ok() @@ -1377,18 +972,6 @@ impl StreamContext { false } } - - /// Emit a terminal `ResponseFailed` recording that the client disconnected - /// after the first byte (mid-stream abort). - /// - /// This closes the streaming lifecycle: every `RequestReceived` for a stream - /// that started gets a matching terminal event, so a client-aborted stream - /// is attributable instead of orphaned. `error_kind` is `"ClientCancelled"`, - /// distinct from genuine upstream failures, so operators can filter these - /// out of error-rate analysis. The HTTP status is 200 (already committed at - /// the first byte); the `error_kind` carries the cancellation semantics. - /// Idempotent callers are fine, but this is intended to be the single - /// terminal emit for a disconnect path. fn emit_client_cancelled(&self) { self.event_bus .emit(&ProxyEvent::ResponseFailed(ResponseFailed { @@ -1409,11 +992,6 @@ impl StreamContext { })); } } - -/// Emit the core events produced by a finalization step. -/// -/// Returns `false` when the client disconnected, in which case the caller must -/// abandon the stream task. async fn emit_finalization_events( ctx: &mut StreamContext, encoder: &mut ClientStreamEncoder, @@ -1430,12 +1008,6 @@ async fn emit_finalization_events( } true } - -/// Flush any trailing partial SSE frame left in the framer and emit the events -/// it decodes to. -/// -/// Returns `false` when the client disconnected, in which case the caller must -/// abandon the stream task. async fn finalize_sse_framer( ctx: &mut StreamContext, encoder: &mut ClientStreamEncoder, @@ -1470,12 +1042,6 @@ async fn finalize_sse_framer( } true } - -/// Emit the client encoder's trailing events, plus the OpenAI `[DONE]` -/// terminator when the client protocol requires it. -/// -/// Returns `false` when the client disconnected, in which case the caller must -/// abandon the stream task. async fn emit_encoder_finish( ctx: &mut StreamContext, tx: &tokio::sync::mpsc::Sender, @@ -1487,8 +1053,6 @@ async fn emit_encoder_finish( return false; } } - - // For OpenAI Chat, emit the [DONE] terminator after all chunks. if matches!(ctx.client_protocol, ClientProtocol::OpenAiChat) && !ctx.emit_event(openai_done_event(), tx).await { @@ -1497,14 +1061,13 @@ async fn emit_encoder_finish( true } - -/// Build a stream that converts provider byte chunks into client SSE events. -/// -/// The `first_byte_tx` channel is used to signal the first-byte boundary: -/// the spawned task sends either `FirstByteResult::FirstEvent` (once the -/// first SSE event is produced) or `FirstByteResult::PreStreamError` (if -/// the stream fails before any event). After the first event is sent, the -/// channel is dropped and all subsequent errors become in-band SSE events. +#[expect( + clippy::too_many_lines, + reason = "one streaming state machine: the select-loop's cancellation, \ + decode, first-byte, and disconnect transitions read as a unit; \ + splitting the arms into helpers would scatter the state they \ + mutate" +)] fn build_sse_output_stream( byte_stream: std::pin::Pin< Box< @@ -1516,25 +1079,9 @@ fn build_sse_output_stream( ctx: StreamContext, failure_recorded: Arc, ) -> BoxStream<'static, Event> { - // SSE channel buffer size. 256 events provides ~256 KB of headroom (typical - // SSE events are ~1 KB). If the client reads slowly and the buffer fills, - // backpressure is applied naturally via the tokio mpsc channel. A larger - // buffer reduces the chance of blocking the stream processing task at the - // cost of more memory per concurrent stream. This value could be made - // configurable via AppState in a future release if tuning is needed. let (tx, rx) = tokio::sync::mpsc::channel::(256); let cancel = tokio_util::sync::CancellationToken::new(); let cancel_clone = cancel.clone(); - - // Capture the fields needed to observe a panic in the spawned stream task. - // CatchPanicLayer (mod.rs) wraps only the request future, not detached - // tasks, so a panic *after* the first byte crosses the boundary is - // otherwise swallowed: `tx` drops, the SSE stream ends mid-flight, and no - // ResponseFailed event / failure metric is recorded. We capture the - // JoinHandle below and supervise it so a JoinError (panic) emits a - // ResponseFailed event + records a failure (audit LOW, detached spawn). - // The span instruments the spawned future so the runtime's panic log also - // carries the request id. let supervise_request_id = ctx.request_id.clone(); let supervise_event_bus = Arc::clone(&ctx.event_bus); let supervise_metrics = Arc::clone(&ctx.stream_metrics.metrics); @@ -1551,26 +1098,11 @@ fn build_sse_output_stream( let mut stream = byte_stream; let mut ctx = ctx; let mut stream_succeeded = false; - // Track whether the stream terminated due to an error so we can - // skip the finalization path that follows the main loop. let mut stream_errored = false; - - // Buffer the first decoded event batch so the client stream encoder can - // be seeded with the upstream provider's real message ID (from - // `CoreEvent::MessageStart`) instead of a synthetic one. This runs - // *inside* the spawned task so the cancellation `select!` covers the - // first-frame wait: a client disconnect during buffering aborts cleanly - // instead of blocking on the upstream. - // - // `CoreEvent`s accumulate across frames/chunks until the first non-empty - // batch arrives (empty decodes -- e.g. `:keepalive` comments -- are - // skipped), then the encoder is constructed and the buffered events are - // replayed through it. No events are lost or reordered. let mut pending: Vec = Vec::new(); 'buffer: loop { tokio::select! { _ = cancel_clone.cancelled() => { - // Client disconnected before the first event. ctx.stream_metrics.metrics.record_client_cancel(); ctx.send_pre_stream_error( RouteError::Internal("client disconnected before first byte".to_owned()), @@ -1629,10 +1161,6 @@ fn build_sse_output_stream( } } } - - // Seed the encoder with the real upstream message ID when the buffered - // batch contained a `MessageStart { id: Some(..) }`; otherwise fall back - // to the protocol-conventional synthetic id. let msg_id = extract_upstream_message_id(&pending).unwrap_or_else(|| match ctx.client_protocol { ClientProtocol::OpenAiChat => format!("chatcmpl-{}", uuid::Uuid::new_v4()), @@ -1644,26 +1172,8 @@ fn build_sse_output_stream( ctx.core.model.requested.clone(), &ctx.core, ); - - // Replay the buffered events. The first one crosses the first-byte - // boundary (committing HTTP 200 via the `first_byte` channel); the rest - // flow directly to the output channel. - // - // Special case: a leading `CoreEvent::Error` on the OpenAI path. The - // OpenAI StreamEncoder returns `Err` for `CoreEvent::Error`, so - // `encode_core_event` yields no SSE event and the error would be - // silently dropped -- leaving `first_byte_sent` false and surfacing a - // misleading "empty stream" 502 (audit LOW, stream-ids-01). When that - // happens, synthesize the OpenAI error-object SSE event from the - // CoreStreamError so it crosses the boundary as a real in-band event - // (200 + error + [DONE]), mirroring the post-first-byte - // `emit_stream_error` path. Anthropic encodes `Error` fine, so this - // branch is OpenAI-only. for core_event in pending.drain(..) { capture_lifecycle(&mut ctx, &core_event); - // Capture the error before `encode_core_event` consumes the event, - // so we can synthesize an SSE error frame if the OpenAI encoder - // drops it (it returns `Err` for `CoreEvent::Error`). let leading_openai_error = if matches!(ctx.client_protocol, ClientProtocol::OpenAiChat) { match &core_event { CoreEvent::Error { error } => Some(error.clone()), @@ -1676,27 +1186,18 @@ fn build_sse_output_stream( if client_events.is_empty() && let Some(stream_error) = leading_openai_error { - // If serialization failed there is nothing else to do for - // this event; continue the replay loop. let Some(json_str) = openai_chat::stream_error_sse_payload(&stream_error) else { continue; }; if !ctx.emit_event(Event::default().data(json_str), &tx).await { return; } - // Emit the OpenAI `[DONE]` terminator after the error. if !ctx.emit_event(openai_done_event(), &tx).await { return; } encoder.mark_finished(); stream_succeeded = false; stream_errored = true; - // Record the failure and emit a ResponseFailed event so - // the in-band error is reflected in metrics/event-log, - // mirroring the post-first-byte in-band error path. The - // HTTP status is already committed as 200 (the error - // crossed the boundary in-band), but the event still - // classifies the failure for operators. ctx.stream_metrics.metrics.record_failure(); let route_error = RouteError::ProviderDecode(stream_error.message().to_owned()); emit_response_failed( @@ -1715,31 +1216,16 @@ fn build_sse_output_stream( } } } - - // Skip the main consumption loop when the replay already terminated - // the stream with an in-band error (leading CoreEvent::Error on the - // OpenAI path). The error + [DONE] have already crossed the boundary; - // reading more from the upstream would only append noise after the - // terminal error (audit LOW, stream-ids-01). if !stream_errored { 'outer: loop { tokio::select! { _ = cancel_clone.cancelled() => { - // Client disconnected; abort upstream stream. This is a - // client-initiated cancellation, not an upstream failure: - // record it as a cancellation so it does not inflate the - // error-rate SLO (audit MEDIUM-5). ctx.stream_metrics.metrics.record_client_cancel(); if !ctx.first_byte_sent { ctx.send_pre_stream_error( RouteError::Internal("client disconnected before first byte".to_owned()), ); } else { - // Post-first-byte disconnect: emit a terminal - // `ClientCancelled` event so the lifecycle closes - // symmetrically with the pre-first-byte path (which - // emits via PreStreamError). Audit finding: buffering - // vs main-loop disconnect event inconsistency. ctx.emit_client_cancelled(); } return; @@ -1854,11 +1340,6 @@ fn build_sse_output_stream( map_provider_error(&e, ctx.passthrough_auth), ); } else { - // In-band stream error: run through the full - // sanitizer (key-redact + URL-redact + truncate) - // so a future error string carrying an API key - // is redacted, matching the `Api` HTTP path - // (audit GAP-LOW-2). let sanitized = fully_sanitize_upstream_error(&e.to_string()); emit_stream_error( &mut encoder, @@ -1867,16 +1348,6 @@ fn build_sse_output_stream( &sanitized, ).await; ctx.stream_metrics.metrics.record_failure(); - // Classify the in-band error through - // `map_provider_error` (the same helper the - // non-stream and pre-stream paths use) so the - // ResponseFailed.http_status matches the - // non-stream classification: an upstream `Api` - // error preserves its real status (429/400/413 - // pass through; 5xx -> 502), and an in-band - // reqwest timeout is attributed as an upstream - // timeout (504) rather than a generic 502 - // (audit LOW, eventlog-inband-stream-httpstatus). let route_error = map_provider_error(&e, ctx.passthrough_auth); emit_response_failed( &ctx.event_bus, @@ -1891,32 +1362,17 @@ fn build_sse_output_stream( break 'outer; } None => { - // Stream ended with no events at all. - // If first byte was never sent, this is a normal - // completion of an empty stream -- the handler will - // see the channel close and respond accordingly. break 'outer; } } } } } - } // end `if !stream_errored` guard around the main consumption loop - - // Only run finalization when the stream completed normally (not errored) - // AND actually produced output (first_byte_sent). An empty / never-started - // upstream stream must NOT run finalization: the provider decoder's - // finish() would synthesize MessageStart+MessageStop, cross the first-byte - // boundary, and commit HTTP 200. Skipping finalization lets such a stream - // fall through to the PreStreamError -> 502 path below. + } if !stream_errored && ctx.first_byte_sent { - // Finalize: call sse_framer.finish() first to flush any trailing partial - // SSE frame that arrived without a terminating blank line. if !finalize_sse_framer(&mut ctx, &mut encoder, &tx).await { return; } - - // Then finalize: call provider decoder finish() to emit any remaining events. match ctx.provider_decoder.finish() { Ok(final_events) => { if !emit_finalization_events(&mut ctx, &mut encoder, &tx, final_events).await { @@ -1931,8 +1387,6 @@ fn build_sse_output_stream( ); } } - - // Emit any remaining client encoder events (synthetic terminal if needed). match encoder.finish() { Ok(final_encoded_events) => { if !emit_encoder_finish(&mut ctx, &tx, final_encoded_events).await { @@ -1950,26 +1404,16 @@ fn build_sse_output_stream( } } } - - // Record metrics after stream completes. if stream_succeeded { ctx.stream_metrics.metrics.record_success( &ctx.stream_metrics.provider_name, &ctx.stream_metrics.upstream_model, ctx.stream_metrics.start.elapsed(), ); - // Emit the stream ResponseCompleted event with accumulated lifecycle - // data (usage, cost, stop reason, upstream message id). let usage = ctx .pending_usage .clone() .unwrap_or_else(Usage::synthetic_zero); - // Only attach a cost when the usage was genuinely reported by the - // provider. Synthetic/Unknown usage (the `unwrap_or_else` zero-fill - // when no `UsageDelta` arrived) would otherwise produce a - // Some(zero-Cost) indistinguishable from a truly free request; - // leaving cost None for unprovenanced usage keeps the event honest - // (audit INFO, cost-on-synthetic-zero). Mirrors handle_core_once. let cost = if usage.provenance == UsageProvenance::ProviderReported { ctx.pricing.as_ref().map(|p| compute_cost(&usage, p)) } else { @@ -1994,12 +1438,6 @@ fn build_sse_output_stream( .unwrap_or(u64::MAX), })); } - - // If the stream ended without ever sending a first byte (empty stream - // with no errors), send a PreStreamError with a descriptive message. - // The handler will return this as a 502 Bad Gateway to the client. - // This is preferable to silently dropping the channel (which produces - // a misleading "stream task panicked" error). if !ctx.first_byte_sent && let Some(tx) = ctx.first_byte_tx.take() && tx @@ -2013,16 +1451,6 @@ fn build_sse_output_stream( } .instrument(stream_span), ); - - // Supervise the stream task so a panic *after* the first byte is not - // swallowed. Pre-first-byte panics surface to the client through the - // `first_byte_rx` RecvError path (mapped to RouteError::Internal in the - // handler); post-first-byte panics drop `tx` and silently end the stream - // with no event/metric. This supervisor closes that observability gap: - // on a JoinError it emits a ResponseFailed event and records a failure so - // operators can diagnose the panic (audit LOW, detached spawn). The stream - // task itself is not blocked -- it terminates independently and the - // supervisor merely observes its exit. tokio::spawn(async move { match stream_task.await { Ok(()) => {} @@ -2042,15 +1470,6 @@ fn build_sse_output_stream( panicked, "detached stream task did not complete cleanly" ); - // Only record/emit if the handler's `first_byte_rx` RecvError - // path has not already done so for this same pre-first-byte - // panic. The shared `failure_recorded` flag is claimed by the - // first watcher to run; the second skips, so a pre-first-byte - // panic produces exactly one `ResponseFailed` + one failure - // metric (audit finding: stream-task panic double-ResponseFailed). - // Post-first-byte panics are only ever observed here (the handler - // is no longer awaiting `first_byte_rx`), so the flag is still - // false and the supervisor emits. if !failure_recorded.swap(true, Ordering::SeqCst) { supervise_metrics.record_failure(); emit_response_failed( @@ -2065,8 +1484,6 @@ fn build_sse_output_stream( } } }); - - // Wrap the receiver so that dropping it cancels the spawned task. let rx_stream = tokio_stream::wrappers::ReceiverStream::new(rx); let cancel_guard = cancel.drop_guard(); rx_stream @@ -2076,34 +1493,12 @@ fn build_sse_output_stream( }) .boxed() } - -/// Extract the upstream message ID from the first `MessageStart` event in a -/// batch, if any carries one. -/// -/// Used to seed the client stream encoder with the provider's real message ID -/// instead of a synthetic one. Only `MessageStart { id: Some(..) }` contributes; -/// earlier events in the batch (e.g. `Ping`) are skipped, so a -/// `Ping`-then-`MessageStart` opening still surfaces the real id. fn extract_upstream_message_id(events: &[CoreEvent]) -> Option { events.iter().find_map(|event| match event { CoreEvent::MessageStart { id: Some(id), .. } => Some(id.clone()), _ => None, }) } - -/// Encode a single [`CoreEvent`] using the appropriate client stream encoder. -/// -/// # Silently dropped events -/// -/// If encoding fails (e.g. the encoder encounters an invalid state), the event -/// is logged at warn level and silently dropped. This prevents a single -/// malformed event from terminating the entire SSE stream. The trade-off is -/// that clients may miss events without explicit notification. This is -/// acceptable because: -/// 1. Encoding failures indicate a proxy bug, not a transient client issue. -/// 2. The stream will still terminate normally (with `finish()`) rather than -/// abruptly mid-frame. -/// 3. The `warn!` log ensures operators can diagnose the root cause. fn encode_core_event(client_encoder: &mut ClientStreamEncoder, event: CoreEvent) -> Vec { match client_encoder.encode_event(event) { Ok(encoded_events) => encoded_events @@ -2116,17 +1511,6 @@ fn encode_core_event(client_encoder: &mut ClientStreamEncoder, event: CoreEvent) } } } - -/// Emit an error event into the stream. -/// -/// For OpenAI Chat, the StreamEncoder returns `Err` for `CoreEvent::Error`. -/// Rather than silently dropping the error, we construct a raw SSE error -/// event directly so the client sees an error indication before `[DONE]`. -/// -/// After the error event, the encoder's `finished` flag is set to `true` so -/// that any subsequent `finish()` call returns an empty vec. The error event -/// itself is the terminal event -- no synthetic message_delta/message_stop -/// pair should follow an error. async fn emit_stream_error( client_encoder: &mut ClientStreamEncoder, tx: &tokio::sync::mpsc::Sender, @@ -2141,14 +1525,6 @@ async fn emit_stream_error( let events = encode_core_event(client_encoder, error_event); if events.is_empty() && matches!(client_protocol, ClientProtocol::OpenAiChat) { - // The OpenAI StreamEncoder returns Err for CoreEvent::Error, which - // causes encode_core_event to produce an empty Vec. Construct an SSE - // error event using the same typed structs as the HTTP error path so - // both paths stay consistent at compile time. - // - // Use "server_error" as the error type to match the HTTP 500 path's - // convention in openai_error_response, since in-band stream errors are - // always upstream/internal issues. if let Some(json_str) = openai_stream_error_json_with_type(message, "server_error") && tx.send(Event::default().data(json_str)).await.is_err() { @@ -2161,43 +1537,17 @@ async fn emit_stream_error( } } } - - // For OpenAI Chat, emit the [DONE] terminator after an error event. if matches!(client_protocol, ClientProtocol::OpenAiChat) && tx.send(openai_done_event()).await.is_err() { return; } - - // Mark the encoder as finished so that subsequent finish() calls return - // empty events. The error event itself is the terminal event -- we do NOT - // emit synthetic message_delta/message_stop after an error. client_encoder.mark_finished(); } - -// --------------------------------------------------------------------------- -// Error mapping -// --------------------------------------------------------------------------- - -/// Map a [`ProviderError`](llm_proxy_provider::error::ProviderError) to a -/// [`RouteError`]. -/// -/// Upstream error bodies are sanitized before being forwarded to clients: -/// - `Api` variants: already sanitized at construction time. -/// - Non-`Api` variants (Http, Serialize, etc.): run through -/// [`fully_sanitize_upstream_error`], which applies the canonical -/// key-redacting sanitizer (via [`ProviderError::api`]) followed by URL -/// redaction and truncation, so neither API keys nor upstream hostnames/URL -/// paths leak to the client (audit GAP-LOW-2). pub(crate) fn map_provider_error( e: &llm_proxy_provider::error::ProviderError, passthrough_auth: bool, ) -> RouteError { - // Whose credentials authenticated with the upstream: under passthrough_auth - // the client's own token is forwarded, so auth/permission failures are - // client-owned; otherwise they are operator-owned. This drives whether - // 401/403 pass through or collapse to 502 in extract_error_fields - // (audit LOW, passthrough-auth-401-collapse). let auth_owner = if passthrough_auth { AuthOwner::Client } else { @@ -2219,16 +1569,10 @@ pub(crate) fn map_provider_error( } } _ => { - // Check for reqwest timeout specifically so we can return 504 - // instead of the generic 502 Bad Gateway. if e.is_timeout() { let sanitized = fully_sanitize_upstream_error(&e.to_string()); return RouteError::UpstreamTimeout(sanitized); } - // Sanitize non-Api error messages through the full sanitizer - // (key-redact + URL-redact + truncate) so a future error string - // carrying an API key is redacted, matching the `Api` branch above - // (audit GAP-LOW-2). let sanitized = fully_sanitize_upstream_error(&e.to_string()); RouteError::Upstream { status: StatusCode::BAD_GATEWAY, @@ -2238,105 +1582,34 @@ pub(crate) fn map_provider_error( } } } - -/// Maximum length for sanitized upstream error bodies before passing to -/// `truncate_error_body()` in the error response encoder. const MAX_SANITIZE_LEN: usize = 512; const SANITIZE_SUFFIX: &str = "...[truncated]"; - -/// Compiled regex for URL redaction, initialized once via `LazyLock`. static URL_REDACT_REGEX: std::sync::LazyLock = std::sync::LazyLock::new(|| { regex::Regex::new(r"https?://\S+").expect("valid URL redaction regex") }); - -/// Sanitize non-Api upstream error bodies to prevent information leakage. -/// -/// Non-Api errors (Http, Serialize, Utf8, etc.) are converted via -/// `e.to_string()`, which for reqwest errors can contain full URLs including -/// hostnames and paths. This function replaces such details with generic -/// messages while preserving enough context for debugging. pub(super) fn sanitize_upstream_error_body(msg: &str) -> String { - // Truncate to MAX_SANITIZE_LEN as a safety net. The suffix is included - // within this budget. The actual error body in RouteError::Upstream is - // further truncated by truncate_error_body() in the error response encoder. let truncated = truncate_with_suffix(msg, MAX_SANITIZE_LEN, SANITIZE_SUFFIX); - - // Redact URL-like patterns that may contain hostnames/paths. - // Matches http:// or https:// followed by any non-whitespace chars. URL_REDACT_REGEX .replace_all(&truncated, "[url-redacted]") .into_owned() } - -/// Sanitize a raw non-`Api` upstream error string through the **full** sanitizer -/// pipeline: the canonical key-redacting sanitizer plus URL redaction and -/// truncation. -/// -/// This is the defense-in-depth counterpart to `sanitize_upstream_error_body`: -/// that function only redacts URLs and truncates, so a future error string that -/// happens to carry an API key (e.g. a provider echoing a `Bearer` token or an -/// `sk-...` value in a transport message) would reach the client un-redacted. -/// Routing the divergent error paths (streaming in-band errors and non-`Api` -/// HTTP errors) through this function keeps them consistent with the `Api` -/// branch, which is already key-redacted at construction time -/// (`ProviderError::api` → `sanitize_api_error_body`). -/// -/// # Implementation -/// -/// The canonical key-redacting sanitizer lives in -/// `llm_proxy_provider::error::sanitize_api_error_body` and is `pub(crate)` to -/// the provider crate, so it cannot be called directly. Its only public entry -/// point is the [`ProviderError::api`] constructor, whose doc comment states it -/// "encapsulates the sanitization call so callers never need to remember to -/// call the internal `sanitize_api_error_body` function manually." We therefore -/// round-trip the message through that constructor (with a sentinel status that -/// is never surfaced — only the `body` is read back) to obtain a key-redacted, -/// truncated string, then run the result through [`sanitize_upstream_error_body`] -/// for URL redaction. Both truncation passes are idempotent (the second is a -/// no-op because the body is already within `MAX_SANITIZE_LEN`). pub(super) fn fully_sanitize_upstream_error(msg: &str) -> String { - // Apply the canonical key-redaction + truncation via the public `api` - // constructor. The status code is a sentinel: only the sanitized `body` - // field is extracted below, and this function never produces a `RouteError` - // whose status depends on it. let key_redacted = match llm_proxy_provider::error::ProviderError::api(0, msg.to_owned()) { llm_proxy_provider::error::ProviderError::Api { body, .. } => body, - // `api()` is guaranteed to construct the `Api` variant, but - // `ProviderError` is `#[non_exhaustive]`, so a catch-all keeps the - // match exhaustive against future variants without silently changing - // behaviour. Fall back to the URL+truncate sanitizer. _ => return sanitize_upstream_error_body(msg), }; - // Layer URL redaction on top (truncation is idempotent here). sanitize_upstream_error_body(&key_redacted) } - -/// Map a [`ProtocolError`](llm_proxy_protocol::client::ProtocolError) to a -/// [`RouteError`]. -/// -/// Per the `ProtocolError` doc comments: -/// - `InvalidRequest` -> 400 Bad Request (bad client input) -/// - `Decode` -> 400 Bad Request (malformed client input) -/// - `Encode` -> 500 Internal Server Error -/// - `EncodeSkippable` -> mapped as non-fatal (skipped), so 500 if it -/// reaches this path. pub(crate) fn protocol_error_to_route(e: llm_proxy_protocol::client::ProtocolError) -> RouteError { match e { llm_proxy_protocol::client::ProtocolError::InvalidRequest(msg) => { RouteError::InvalidRequest(msg) } - llm_proxy_protocol::client::ProtocolError::Decode(msg) => { - // Per ProtocolError doc: Decode -> 400 Bad Request (malformed client input). - RouteError::InvalidRequest(msg) - } + llm_proxy_protocol::client::ProtocolError::Decode(msg) => RouteError::InvalidRequest(msg), other => RouteError::Internal(other.to_string()), } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -2362,14 +1635,6 @@ mod tests { reasoning: "0.000003".parse().unwrap(), }; let cost = compute_cost(&usage, &pricing); - // input: 1000 * 0.0000015 = 0.0015 - // output: 500 * 0.000003 = 0.0015 - // cache_creation: 200 * 0.000001875 = 0.000375 - // cache_read: None -> 0 - // reasoning: 50 * 0.000003 = 0.00015 - // total = 0.0015 + 0.0015 + 0.000375 + 0 + 0.00015 = 0.003525 - // `.normalize().to_string()` strips trailing zeros so the expected - // strings are scale-independent. assert_eq!(cost.input.normalize().to_string(), "0.0015"); assert_eq!(cost.output.normalize().to_string(), "0.0015"); assert_eq!(cost.cache_creation.normalize().to_string(), "0.000375"); @@ -2380,9 +1645,6 @@ mod tests { #[test] fn compute_cost_rounds_to_microcent() { - // Per the plan's locked decision 5, cost is rounded to the nearest - // microcent (6 dp) on output. A price that would otherwise carry more - // than 6 dp (0.000001875 * 3 tokens = 0.000005625) is rounded. use llm_proxy_core::ModelPricing; use llm_proxy_protocol::core::Usage; let usage = Usage { @@ -2395,19 +1657,12 @@ mod tests { ..ModelPricing::default() }; let cost = compute_cost(&usage, &pricing); - // 3 * 0.000001875 = 0.000005625 -> rounded to 6 dp -> 0.000006 (half up). assert_eq!(cost.input.normalize().to_string(), "0.000006"); - // total == input (the only nonzero component). assert_eq!(cost.total.normalize().to_string(), "0.000006"); } #[test] fn compute_cost_saturates_on_overflow_and_never_panics() { - // `rust_decimal`'s `*`/`+` panic on overflow; compute_cost must use - // checked arithmetic so a pathological (count, price) pair saturates to - // Decimal::MAX instead of crashing the request. The validation bound - // prevents such prices from reaching compute_cost via config, but this - // guards the function itself against any caller. use llm_proxy_core::ModelPricing; use llm_proxy_protocol::core::Usage; let usage = Usage { @@ -2425,19 +1680,10 @@ mod tests { cache_read: "999999999999999999999".parse().unwrap(), reasoning: "999999999999999999999".parse().unwrap(), }; - // Must not panic. let cost = compute_cost(&usage, &pricing); assert_eq!(cost.input, Decimal::MAX); assert_eq!(cost.total, Decimal::MAX); } - - /// Regression guard for the 502-collapse secret/schema leak (Cluster A) and - /// the value-oriented source-guard gap (finding #8). The storage crate's - /// source-guard only checks that secret-bearing KEY NAMES never appear as - /// JSON keys; it does not check that a secret VALUE is absent from the - /// `message` string. This test routes an upstream error carrying a - /// realistic API-key-shaped value through `emit_response_failed` and asserts - /// the secret appears NOWHERE in the serialized event JSON. #[test] fn emit_response_failed_redacts_secret_values_from_message() { use llm_proxy_storage::RecordingBus; @@ -2467,9 +1713,7 @@ mod tests { _ => None, }) .expect("a ResponseFailed was emitted"); - // error_kind is the RouteError variant name (Cluster C guard). assert_eq!(recorded.error_kind, "Upstream"); - // The secret value must not appear in the message... assert!( !recorded.message.contains("sk-ant-api03-"), "secret prefix leaked into message: {}", @@ -2480,18 +1724,12 @@ mod tests { "key material leaked into message: {}", recorded.message ); - // ...nor anywhere in the serialized event JSON (the storage key-name - // guard would miss this; this is the value-oriented guard). let json = serde_json::to_string(&ProxyEvent::ResponseFailed(recorded)).unwrap(); assert!( !json.contains("sk-ant-api03-") && !json.contains("ABCDEFGHIJKLMNOP"), "secret value leaked into event JSON: {json}" ); } - - /// The passthrough-auth Authorization scheme is case-insensitive (RFC 7235): - /// "Bearer", "bearer", "BEARER", "bEaReR" must all be accepted. Guards the - /// Cluster F fix that replaced a two-casing prefix match. #[test] fn extract_inbound_auth_bearer_scheme_is_case_insensitive() { use secrecy::ExposeSecret; @@ -2504,25 +1742,16 @@ mod tests { assert_eq!(extracted_token("bearer abc123").as_deref(), Some("abc123")); assert_eq!(extracted_token("BEARER abc123").as_deref(), Some("abc123")); assert_eq!(extracted_token("BeArEr abc123").as_deref(), Some("abc123")); - // The token itself is trimmed; surrounding spaces do not leak in. assert_eq!( extracted_token("Bearer abc123 ").as_deref(), Some("abc123") ); - // A non-Bearer scheme is not extracted via Authorization. assert_eq!(extracted_token("Basic dXNlcjpwYXNz"), None); - // An empty token after the scheme yields None (falls through to the - // x-api-key / x-goog-api-key headers, which are absent here). assert_eq!(extracted_token("Bearer "), None); } #[test] fn body_hash_is_stable_and_distinguishes_requests() { - // `body_hash` is computed over the raw client body bytes, matching the - // `/v1/messages/count_tokens` route, so identical raw bodies hash - // equally and distinct bodies differ. Two bodies that normalize to the - // same CoreRequest but differ in whitespace/byte order hash differently - // because the correlation key tracks the raw wire payload. let a1 = br#"{"model":"gpt-4o","messages":[]}"#; let a2 = br#"{"model":"gpt-4o","messages":[]}"#; let b = br#"{"model":"gpt-3.5-turbo","messages":[]}"#; @@ -2536,8 +1765,6 @@ mod tests { body_hash_of_bytes(b), "different raw bodies must produce different hashes" ); - // Whitespace-sensitive (raw bytes, not normalized): the same JSON with - // a space after the colon hashes differently. let spaced = br#"{"model": "gpt-4o","messages":[]}"#; assert_ne!( body_hash_of_bytes(a1), @@ -2563,7 +1790,11 @@ mod tests { dedup_window: std::time::Duration::ZERO, server_name: "test".to_owned(), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"); @@ -2582,8 +1813,6 @@ mod tests { ) } - // -- protocol_error_to_route ------------------------------------------------ - #[test] fn protocol_error_invalid_request_maps_to_route_invalid_request() { let err = llm_proxy_protocol::client::ProtocolError::InvalidRequest("bad".into()); @@ -2612,60 +1841,104 @@ mod tests { fn direct_loopback_bypasses_rate_limit() { let state = state_with_operational_config(1, false); let headers = HeaderMap::new(); - let address = "127.0.0.1:12345".parse().unwrap(); + let meta = test_meta( + crate::middleware::CallerId::anonymous(), + Some("127.0.0.1:12345"), + ); - prepare_request( - &state, - "req-test".to_owned(), - &headers, - Some(&address), - b"one", - "/test", - ) - .expect("request under the rate limit should be admitted"); - prepare_request( - &state, - "req-test".to_owned(), - &headers, - Some(&address), - b"two", - "/test", - ) - .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &headers, b"one", "/test") + .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &headers, b"two", "/test") + .expect("request under the rate limit should be admitted"); + } + fn test_meta( + caller: crate::middleware::CallerId, + peer: Option<&str>, + ) -> crate::middleware::RequestMeta { + crate::middleware::RequestMeta { + request_id: crate::middleware::RequestId("req-test".to_owned()), + caller, + connect_info: peer.map(|addr| addr.parse().unwrap()), + } } #[test] fn untrusted_forwarded_headers_cannot_bypass_rate_limit() { let state = state_with_operational_config(1, false); - let address = "198.51.100.10:12345".parse().unwrap(); + let meta = test_meta( + crate::middleware::CallerId::anonymous(), + Some("198.51.100.10:12345"), + ); let mut first_headers = HeaderMap::new(); first_headers.insert("x-forwarded-for", "203.0.113.1".parse().unwrap()); let mut second_headers = HeaderMap::new(); second_headers.insert("x-forwarded-for", "203.0.113.2".parse().unwrap()); - prepare_request( - &state, - "req-test".to_owned(), - &first_headers, - Some(&address), - b"one", - "/test", - ) - .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &first_headers, b"one", "/test") + .expect("request under the rate limit should be admitted"); assert!(matches!( - prepare_request( - &state, - "req-test".to_owned(), - &second_headers, - Some(&address), - b"two", - "/test" - ), + prepare_request(&state, &meta, &second_headers, b"two", "/test"), Err(RouteError::RateLimited) )); } - // -- map_provider_error ----------------------------------------------------- + #[test] + fn authenticated_caller_is_limited_by_key_id_across_addresses() { + let state = state_with_operational_config(1, false); + let caller = crate::middleware::CallerId::authenticated("laptop"); + let first_meta = test_meta(caller.clone(), Some("198.51.100.10:12345")); + let second_meta = test_meta(caller, Some("198.51.100.11:12345")); + let headers = HeaderMap::new(); + + prepare_request(&state, &first_meta, &headers, b"one", "/test") + .expect("first request under the rate limit should be admitted"); + assert!(matches!( + prepare_request(&state, &second_meta, &headers, b"two", "/test"), + Err(RouteError::RateLimited) + )); + } + + #[test] + fn authenticated_loopback_caller_is_still_rate_limited() { + let state = state_with_operational_config(1, false); + let meta = test_meta( + crate::middleware::CallerId::authenticated("laptop"), + Some("127.0.0.1:12345"), + ); + let headers = HeaderMap::new(); + + prepare_request(&state, &meta, &headers, b"one", "/test") + .expect("first request under the rate limit should be admitted"); + assert!(matches!( + prepare_request(&state, &meta, &headers, b"two", "/test"), + Err(RouteError::RateLimited) + )); + } + + #[test] + fn request_context_carries_caller_attribution() { + let state = state_with_operational_config(100, false); + let headers = HeaderMap::new(); + let ctx = prepare_request( + &state, + &test_meta(crate::middleware::CallerId::anonymous(), None), + &headers, + b"", + "/test", + ) + .expect("request admitted"); + assert_eq!(ctx.caller, "anonymous"); + + let ctx = prepare_request( + &state, + &test_meta(crate::middleware::CallerId::authenticated("laptop"), None), + &headers, + b"", + "/test", + ) + .expect("request admitted"); + assert_eq!(ctx.caller, "laptop"); + } #[test] fn provider_route_error_maps_unsupported_route_without_message_parsing() { @@ -2706,7 +1979,6 @@ mod tests { #[test] fn provider_api_error_marks_client_owner_under_passthrough() { - // passthrough_auth=true => auth_owner is Client so 401/403 pass through. let err = llm_proxy_provider::error::ProviderError::api(401, "unauthorized".into()); let route_err = map_provider_error(&err, true); match route_err { @@ -2722,7 +1994,6 @@ mod tests { #[test] fn provider_api_error_marks_operator_owner_by_default() { - // passthrough_auth=false => auth_owner is Operator so 401 collapses to 502. let err = llm_proxy_provider::error::ProviderError::api(401, "unauthorized".into()); let route_err = map_provider_error(&err, false); match route_err { @@ -2736,8 +2007,6 @@ mod tests { } } - // -- resolve_target with unknown provider ----------------------------------- - #[tokio::test] async fn resolve_target_unknown_provider_returns_unknown_provider() { use llm_proxy_core::AppConfig; @@ -2757,7 +2026,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = @@ -2806,8 +2079,6 @@ mod tests { } } - // -- resolve_target with empty registry returns unknown provider ------------ - #[tokio::test] async fn resolve_target_empty_registry_returns_unknown_provider() { use llm_proxy_core::AppConfig; @@ -2827,7 +2098,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = @@ -2876,8 +2151,6 @@ mod tests { } } - // -- encode_core_event smoke test ------------------------------------------- - #[test] fn encode_core_event_ping_produces_event() { use llm_proxy_protocol::core::CoreEvent; @@ -2886,13 +2159,9 @@ mod tests { "test-model".to_owned(), )); let events = encode_core_event(&mut encoder, CoreEvent::Ping); - // Ping may or may not produce output depending on the encoder impl. - // The important thing is it doesn't panic. let _ = events; } - // -- sanitize_upstream_error_body tests ------------------------------------ - #[test] fn sanitize_removes_urls() { let msg = @@ -2922,16 +2191,8 @@ mod tests { assert_eq!(sanitized, msg); } - // -- fully_sanitize_upstream_error tests (audit GAP-LOW-2) ------------------ - // - // The divergent error paths (streaming in-band, non-Api timeout/generic) - // used to run only `sanitize_upstream_error_body` (URL-redact + truncate), - // which never strips API-key patterns. `fully_sanitize_upstream_error` - // routes them through the canonical key-redacting sanitizer as well. - #[test] fn fully_sanitize_redacts_api_key_patterns() { - // A long sk-ant-... value (40 trailing token chars) is a real key shape. let msg = "upstream auth failed for key sk-ant-api03-ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890abcd"; let sanitized = fully_sanitize_upstream_error(msg); @@ -2951,7 +2212,6 @@ mod tests { #[test] fn fully_sanitize_redacts_urls() { - // URL redaction (from sanitize_upstream_error_body) must still apply. let msg = "connection refused to https://api.openai.com/v1/chat/completions"; let sanitized = fully_sanitize_upstream_error(msg); assert!( diff --git a/crates/llm-proxy-server/src/routes/error_response.rs b/crates/llm-proxy-server/src/routes/error_response.rs index 4ab1fd7..3d07927 100644 --- a/crates/llm-proxy-server/src/routes/error_response.rs +++ b/crates/llm-proxy-server/src/routes/error_response.rs @@ -1,171 +1,59 @@ -//! Protocol-aware error encoding for route handlers. -//! -//! Defines [`RouteError`] as the single internal error model used by the core -//! pipeline, and [`route_error_response`] to encode it into a client-specific -//! HTTP response. Each route passes its [`ClientProtocol`] so the error shape -//! matches what the client expects (Anthropic JSON envelope, OpenAI error, etc.). -//! -//! # Security: error message sanitization -//! -//! Internal error messages (config details, adapter names, protocol names, etc.) -//! are NEVER sent to clients. Instead, generic messages are used in the HTTP -//! response body, while the actual messages are logged server-side only. This -//! prevents information disclosure about the proxy's internal architecture. -//! -//! Upstream error bodies are sanitized by the provider error layer and further -//! truncated here. For codes that pass through to the client (400/413/429 and -//! client-owned 401/403), the provider-specific message is replaced with a -//! generic proxy message and the real body is logged server-side only, so the -//! upstream's own schema never reaches the client. Provider decode errors are -//! similarly replaced with generic messages to avoid leaking upstream response -//! fragments. - use axum::body::Body; use axum::http::{HeaderValue, StatusCode, header}; use axum::response::{IntoResponse, Response}; use serde::Serialize; - -// --------------------------------------------------------------------------- -// ClientProtocol -// --------------------------------------------------------------------------- - -/// Identifies the client-facing protocol for error encoding. -/// -/// Each route knows which protocol its callers speak. The error response module -/// uses this to produce the correct JSON envelope. -/// -/// Rate limiting, deduplication, routing, and provider failures are represented -/// by [`RouteError`] and encoded according to this protocol selection. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub(super) enum ClientProtocol { - /// Anthropic Messages API (`/v1/messages`). Anthropic, - /// OpenAI Chat Completions API (`/v1/chat/completions`). OpenAiChat, } - -// --------------------------------------------------------------------------- -// RouteError -// --------------------------------------------------------------------------- - -/// Who owns the credentials used to authenticate with the upstream provider. -/// -/// This distinguishes the two operating modes of the proxy and drives how -/// upstream auth/permission failures (401/403) are surfaced to the client: -/// -/// - [`AuthOwner::Operator`] (default, managed-key mode): the proxy forwards -/// *its own* configured upstream credentials. An upstream 401/403 therefore -/// means the *operator's* key is expired/revoked/scoped wrong -- an -/// infrastructure problem the client cannot fix -- so it is collapsed to -/// 502 Bad Gateway (see [`map_upstream_status`]). -/// - [`AuthOwner::Client`] (`passthrough_auth = true`): the proxy forwards the -/// *client's* own inbound token. An upstream 401/403 is then the client's -/// credential problem and is passed through verbatim so the client can act -/// on it (audit LOW, passthrough-auth-401-collapse). -/// -/// Non-auth codes (400/413/429) are always client-owned and pass through -/// regardless of this value; 404/5xx always collapse to 502. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub(super) enum AuthOwner { - /// The proxy's own (operator) credentials are forwarded (managed-key mode). #[default] Operator, - /// The client's own credentials are forwarded (`passthrough_auth`). Client, } - -/// Internal error type for the core pipeline. -/// -/// Route handlers return `Result`. The [`route_error_response`] -/// function encodes this into a protocol-specific HTTP response. -/// -/// # Security note -/// -/// Some variants carry internal details (Internal, ProviderDecode, Upstream). -/// The error response encoder sanitizes these before sending them to clients. #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub(super) enum RouteError { - /// Bad client input (malformed JSON, missing fields). #[error("invalid request: {0}")] InvalidRequest(String), - /// The requested model is excluded by provider catalog enforcement. #[error("model not allowed: {0}")] ModelNotAllowed(String), - /// The requested provider name is not registered. #[error("unknown provider: {0}")] UnknownProvider(String), - /// The route kind is not supported by the named provider. #[error("unsupported route: {0}")] UnsupportedRoute(String), - /// The provider name in the URL path contains invalid characters. #[error("invalid provider name: {0}")] InvalidProviderName(String), - /// Upstream provider returned an error status. #[error("upstream error: {status}")] Upstream { - /// HTTP status code from the upstream (or a synthetic one). status: StatusCode, - /// Sanitized error body. body: String, - /// Whose credentials authenticated with the upstream. Drives whether - /// auth/permission codes (401/403) pass through to the client - /// ([`AuthOwner::Client`]) or collapse to 502 ([`AuthOwner::Operator`]). auth_owner: AuthOwner, }, - /// Upstream provider timed out. #[error("upstream timeout: {0}")] UpstreamTimeout(String), - /// Provider adapter failed to decode the upstream response. #[error("provider decode error: {0}")] ProviderDecode(String), - /// Internal server error (misconfiguration, missing state). - /// - /// **Security note:** The message is for server-side logging only. The error - /// response encoder discards it and sends a generic `INTERNAL_ERROR_CLIENT_MESSAGE` - /// to clients. Any future error handling path must NOT format this variant's - /// message directly into client responses. #[error("internal error: {0}")] Internal(String), - /// Client has exceeded its rate limit. #[error("rate limit exceeded")] RateLimited, - /// Duplicate request detected by the deduplicator. #[error("duplicate request")] Conflict, - /// The requested path does not match any mounted route. #[error("not found")] NotFound, - /// The request exceeded the configured server timeout (408). - /// - /// Produced by the per-route timeout layer, then normalised into a - /// protocol-shaped JSON body by the outermost response-normaliser layer - /// (audit MEDIUM-7). Distinct from [`RouteError::UpstreamTimeout`], which - /// is a 504 for an upstream call that exceeded its own deadline. #[error("request timeout")] RequestTimeout, - /// The request body exceeded the maximum allowed size (413). #[error("request body too large")] PayloadTooLarge, - /// The request used an HTTP method the matched route does not allow (405). - /// - /// Produced by the router-wide `method_not_allowed_fallback` so that axum's - /// bare 405 (empty body) is rendered through the same protocol-shaped JSON - /// envelope as 404/408/413 (audit LOW-27). #[error("method not allowed")] MethodNotAllowed, - /// The client did not provide an auth token that a passthrough-auth provider - /// (`passthrough_auth = true`) requires. Maps to 401. #[error("missing client auth token for passthrough provider")] Unauthorized, } - -// --------------------------------------------------------------------------- -// JSON envelope types -// --------------------------------------------------------------------------- - -/// Anthropic-shaped error body. #[derive(Serialize)] #[serde(deny_unknown_fields)] struct AnthropicErrorBody { @@ -179,11 +67,6 @@ struct AnthropicErrorDetail { r#type: String, message: String, } - -/// OpenAI Chat-shaped error body. -/// -/// Uses a typed struct instead of `serde_json::json!()` for compile-time -/// field validation and consistency with the Anthropic path. #[derive(Serialize)] #[serde(deny_unknown_fields)] struct OpenAiErrorBody { @@ -195,56 +78,18 @@ struct OpenAiErrorBody { struct OpenAiErrorDetail { message: String, r#type: String, - // Always serializes as `null` to match the OpenAI wire format: - // `{"error":{"message":"...","type":"...","code":null}}`. - // Using `serde_json::Value::Null` is more idiomatic than `Option<()>`. code: serde_json::Value, } - -// --------------------------------------------------------------------------- -// route_error_response -// --------------------------------------------------------------------------- - -/// Encode a [`RouteError`] into a protocol-specific HTTP response. -/// -/// For `ClientProtocol::Anthropic`, the response body is an Anthropic-shaped -/// JSON envelope with `{"type":"error","error":{...}}`. pub(super) fn route_error_response(protocol: ClientProtocol, error: &RouteError) -> Response { match protocol { ClientProtocol::Anthropic => anthropic_error_response(error), ClientProtocol::OpenAiChat => openai_error_response(error), } } - -/// Generic message used for 500 Internal Server Error responses. -/// The actual internal error details are logged server-side only. pub(crate) const INTERNAL_ERROR_CLIENT_MESSAGE: &str = "internal server error"; - -/// Generic message used for provider decode error responses. -/// The actual decode error details are logged server-side only. pub(crate) const PROVIDER_DECODE_CLIENT_MESSAGE: &str = "provider response decode error"; - -/// Generic message used for upstream 4xx errors that pass through to the -/// client. The provider's own error envelope is logged server-side only and -/// never forwarded, so provider-specific schema (field names, model -/// identifiers, request echoes) does not leak (audit LOW, upstream-400-body-leak). pub(crate) const UPSTREAM_PASSTHROUGH_CLIENT_MESSAGE: &str = "upstream rejected the request"; - -/// Generic message used for upstream errors that collapse to 502 (every status -/// that is not in the pass-through allowlist: 404/408/422/5xx, and -/// operator-owned 401/403). Like the pass-through branch, the real (already -/// key/URL-redacted) upstream body is logged server-side only and a generic -/// message is returned, so the provider's own schema never reaches the client -/// envelope or the persisted `ResponseFailed.message` (audit finding: -/// 502-collapse forwarded a URL-only-redacted body carrying provider schema). pub(crate) const UPSTREAM_COLLAPSE_CLIENT_MESSAGE: &str = "upstream service error"; - -/// Build an Anthropic-shaped error response. -/// -/// Internal error messages are sanitized: `Internal` and `ProviderDecode` -/// variants use generic messages in the response body to prevent information -/// disclosure. The actual messages are available through `RouteError::Display` -/// for server-side logging before this function is called. fn anthropic_error_response(error: &RouteError) -> Response { let (status, error_type, message) = extract_error_fields(error); @@ -255,10 +100,6 @@ fn anthropic_error_response(error: &RouteError) -> Response { message, }, }; - - // axum::Json already sets Content-Type: application/json in its - // IntoResponse implementation. The explicit insert below is - // defense-in-depth to ensure the header is always present. let mut response = (status, axum::Json(body)).into_response(); response.headers_mut().insert( header::CONTENT_TYPE, @@ -266,35 +107,8 @@ fn anthropic_error_response(error: &RouteError) -> Response { ); response } - -/// Build an OpenAI Chat-shaped error response. -/// -/// The error envelope follows the OpenAI error shape: -/// `{"error":{"message":"...","type":"invalid_request_error","code":null}}`. -/// -/// Internal error messages are sanitized: `Internal` and `ProviderDecode` -/// variants use generic messages in the response body to prevent information -/// disclosure. -/// -/// # Two-stage error mapping -/// -/// 1. [`extract_error_fields`] first converts the [`RouteError`] into a -/// (status, error_type, message) tuple. Upstream errors go through -/// [`map_upstream_status`] which converts most upstream codes to 502. -/// 2. This function then overrides `error_type` for specific status codes -/// to match OpenAI conventions: 500 becomes `"server_error"` and 404 -/// becomes `"invalid_request_error"`. -/// -/// # Why 502 Bad Gateway for upstream errors? -/// -/// When the upstream provider returns 4xx/5xx, the proxy maps most of these -/// to 502 Bad Gateway (except 429 which is passed through). This signals to -/// the client that the failure is between the proxy and the upstream, not a -/// client error. The original upstream status is preserved in server-side -/// logs for debugging. fn openai_error_response(error: &RouteError) -> Response { let (status, error_type, message) = extract_error_fields(error); - // OpenAI uses "server_error" for internal errors instead of "api_error". let error_type = match status { StatusCode::INTERNAL_SERVER_ERROR => "server_error", StatusCode::NOT_FOUND => "invalid_request_error", @@ -308,9 +122,6 @@ fn openai_error_response(error: &RouteError) -> Response { code: serde_json::Value::Null, }, }; - - // axum::Json already sets Content-Type: application/json. The explicit - // insert below is defense-in-depth. let mut response = (status, axum::Json(body)).into_response(); response.headers_mut().insert( header::CONTENT_TYPE, @@ -318,12 +129,6 @@ fn openai_error_response(error: &RouteError) -> Response { ); response } - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -/// Build an OpenAI-shaped SSE error JSON string with a custom error type. pub(super) fn openai_stream_error_json_with_type( message: &str, error_type: &str, @@ -337,20 +142,6 @@ pub(super) fn openai_stream_error_json_with_type( }; serde_json::to_string(&body).ok() } - -/// Return the Rust variant name of a [`RouteError`]. -/// -/// The structured event-log contract documents `ResponseFailed.error_kind` as -/// "from the `RouteError` variant name" (see the `ResponseFailed` doc in -/// `llm-proxy-storage`), so operators querying the event log by failure class -/// match a stable discriminant. This is deliberately distinct from the -/// client-facing protocol error-type string returned by -/// [`extract_error_fields`] (e.g. `"api_error"`, `"not_found_error"`): several -/// unrelated variants share an envelope string, so the variant name is the only -/// value that keeps `UnknownProvider`, `NotFound`, and `RateLimited` -/// separable downstream (audit finding: error_kind not variant name). -/// -/// Returns a `&'static str` matching the variant identifier exactly. pub(crate) fn route_error_variant_name(error: &RouteError) -> &'static str { match error { RouteError::InvalidRequest(_) => "InvalidRequest", @@ -371,12 +162,6 @@ pub(crate) fn route_error_variant_name(error: &RouteError) -> &'static str { RouteError::Unauthorized => "Unauthorized", } } - -/// Extract the (status, error_type, message) tuple from a RouteError. -/// -/// Shared between `anthropic_error_response` and `openai_error_response` to -/// avoid duplicating the match arms. Each caller wraps the tuple in its own -/// JSON envelope. pub(crate) fn extract_error_fields(error: &RouteError) -> (StatusCode, &'static str, String) { match error { RouteError::InvalidRequest(msg) => ( @@ -410,14 +195,6 @@ pub(crate) fn extract_error_fields(error: &RouteError) -> (StatusCode, &'static auth_owner, } => { let mapped = map_upstream_status(*status, *auth_owner); - // Both branches return a generic message to the client (and to the - // persisted `ResponseFailed.message`), so the upstream's own schema - // (field names, model identifiers, request echoes) never reaches - // the client envelope or the event log. The real upstream body -- - // already key/URL-redacted at construction -- is logged server-side - // once in `emit_response_failed`, NOT here: this function is also - // called from `route_error_response` on the render path, so logging - // here would double-warn every upstream failure. let message = if mapped == *status { UPSTREAM_PASSTHROUGH_CLIENT_MESSAGE } else { @@ -477,53 +254,11 @@ pub(crate) fn extract_error_fields(error: &RouteError) -> (StatusCode, &'static ), } } - -/// Map an upstream HTTP status to the status we return to the client. -/// -/// Most upstream >= 400 errors become 502 Bad Gateway because they indicate a -/// problem between the proxy and the upstream provider rather than a client -/// error. A conservative allowlist of clearly-client-owned codes is passed -/// through verbatim so the client can act on the failure: -/// - `400 Bad Request` -- the upstream rejected the *content* of a request -/// the proxy faithfully forwarded, which is almost always the client's -/// responsibility. -/// - `413 Payload Too Large` -- the upstream rejected the body size the -/// client sent (mirrors our own `DefaultBodyLimit` 413). -/// - `429 Too Many Requests` -- propagated so clients can implement their -/// own back-off. -/// -/// Auth/permission codes (`401 Unauthorized`, `403 Forbidden`) depend on -/// [`AuthOwner`]: -/// - [`AuthOwner::Operator`] (managed-key mode, the default): the proxy -/// forwards *its own* credentials, so an upstream 401/403 almost always -/// means the *operator's* key is expired/revoked/scoped wrong -- an -/// infrastructure problem the client cannot fix. These collapse to 502 so -/// an operator/infrastructure issue is not misattributed to the client. -/// - [`AuthOwner::Client`] (`passthrough_auth`): the proxy forwards the -/// *client's* own token, so a 401/403 is the client's credential problem -/// and passes through verbatim (audit LOW, passthrough-auth-401-collapse). -/// -/// Other ambiguous or server-side codes (404, 408, 422, 5xx, etc.) also remain -/// 502: they either hint at proxy/provider configuration (404 model-not-found) -/// or are not unambiguously the client's fault, so collapsing them keeps the -/// client contract uniform and avoids leaking provider-specific schemas. The -/// original upstream status is always preserved in server-side logs for operators. -/// -/// Upstream statuses that may indicate a configuration issue (401/403/404 under -/// [`AuthOwner::Operator`]) are logged at warn level so operators can -/// distinguish them in logs. fn map_upstream_status(upstream: StatusCode, auth_owner: AuthOwner) -> StatusCode { match upstream.as_u16() { - // Always-client-owned codes: forward verbatim. 400 | 413 | 429 => upstream, - // Auth/permission codes pass through only when the client's own - // credentials were forwarded; otherwise collapse to 502 (operator-side). 401 | 403 if auth_owner == AuthOwner::Client => upstream, code => { - // Auth/permission (401/403 when operator-owned) and not-found (404) - // statuses usually indicate an operator configuration problem - // (bad/expired upstream key, unknown model) rather than a transient - // upstream failure; log them at warn so operators can diagnose. if matches!(code, 401 | 403 | 404) { tracing::warn!( upstream_status = code, @@ -532,8 +267,6 @@ fn map_upstream_status(upstream: StatusCode, auth_owner: AuthOwner) -> StatusCod issue (auth, forbidden, or not found); mapping to 502 for client" ); } else if code < 400 { - // Informational (1xx), success (2xx), or redirect (3xx) codes should - // never reach the error path. Log at debug level for diagnostics. tracing::debug!( upstream_status = code, "unexpected 1xx/2xx/3xx upstream status in error path; mapping to 502" @@ -543,38 +276,19 @@ fn map_upstream_status(upstream: StatusCode, auth_owner: AuthOwner) -> StatusCod } } } - -/// Maximum length for error messages returned to clients. const MAX_ERROR_MESSAGE_LEN: usize = 512; - -/// Suffix appended when an error body is truncated. const TRUNCATED_SUFFIX: &str = "...[truncated]"; - -/// Truncate an error body to a safe length for client responses. -/// -/// Delegates to the shared implementation in `llm_proxy_protocol::util`. pub(crate) fn truncate_with_suffix(s: &str, max_len: usize, suffix: &str) -> String { llm_proxy_protocol::util::truncate_with_suffix(s, max_len, suffix) } - -/// Truncate an error body to a safe length for client responses. -/// -/// The final string is at most `MAX_ERROR_MESSAGE_LEN` bytes (the suffix is -/// included within this budget). fn truncate_error_body(body: &str) -> String { truncate_with_suffix(body, MAX_ERROR_MESSAGE_LEN, TRUNCATED_SUFFIX) } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; - // -- Anthropic error encoding ---------------------------------------------- - #[test] fn anthropic_invalid_request_returns_400() { let err = RouteError::InvalidRequest("bad input".into()); @@ -639,8 +353,6 @@ mod tests { assert_eq!(response.status(), StatusCode::CONFLICT); } - // -- Internal error message sanitization ----------------------------------- - #[tokio::test] async fn internal_error_message_is_sanitized_in_response_body() { let err = RouteError::Internal("secret config detail: /etc/proxy.toml".into()); @@ -675,8 +387,6 @@ mod tests { ); } - // -- OpenAI error encoding ------------------------------------------------- - #[test] fn openai_invalid_request_returns_400() { let err = RouteError::InvalidRequest("bad input".into()); @@ -753,7 +463,6 @@ mod tests { assert_eq!(json["error"]["type"], "invalid_request_error"); assert_eq!(json["error"]["message"], "bad input"); assert!(json["error"]["code"].is_null(), "code must be null"); - // Must NOT have Anthropic-shaped fields. assert!( json.get("type").is_none() || json["type"].is_null(), "must not have Anthropic 'type' field" @@ -780,10 +489,6 @@ mod tests { #[tokio::test] async fn openai_upstream_500_body_has_correct_structure() { - // A 500 collapses to 502. The client must see a GENERIC message -- the - // upstream's own body (which may carry provider-specific schema, model - // identifiers, or request echoes) is logged server-side only and never - // reaches the client envelope (audit finding: 502-collapse message leak). let err = RouteError::Upstream { status: StatusCode::INTERNAL_SERVER_ERROR, body: "model 'gpt-foo' is overloaded: request echoed here".into(), @@ -798,7 +503,6 @@ mod tests { assert_eq!(json["error"]["type"], "api_error"); assert_eq!(json["error"]["message"], UPSTREAM_COLLAPSE_CLIENT_MESSAGE); assert!(json["error"]["code"].is_null()); - // The provider-specific schema/body must NOT reach the client. let message = json["error"]["message"].as_str().unwrap(); assert!(!message.contains("gpt-foo")); assert!(!message.contains("overloaded")); @@ -807,10 +511,6 @@ mod tests { #[tokio::test] async fn upstream_passthrough_400_uses_generic_message() { - // A pass-through code (400) forwards a GENERIC message to the client so - // the upstream's own schema (field names, model identifiers, request - // echoes) never leaks. The real upstream body is logged server-side - // only (audit LOW, upstream-400-body-leak). let err = RouteError::Upstream { status: StatusCode::BAD_REQUEST, body: r#"{"error":{"message":"model 'gpt-foo' does not exist","type":"invalid_request_error"}}"#.into(), @@ -824,15 +524,12 @@ mod tests { let json: serde_json::Value = serde_json::from_slice(&body).expect("valid JSON"); let message = json["error"]["message"].as_str().unwrap(); assert_eq!(message, UPSTREAM_PASSTHROUGH_CLIENT_MESSAGE); - // The provider-specific schema/model name must NOT reach the client. assert!(!message.contains("gpt-foo")); assert!(!message.contains("invalid_request_error")); } #[tokio::test] async fn upstream_passthrough_401_client_uses_generic_message_and_passes_through() { - // Under passthrough_auth a 401 is client-owned and passes through, but - // the body is still the generic message (not the provider envelope). let err = RouteError::Upstream { status: StatusCode::UNAUTHORIZED, body: "invalid api key sk-leaked".into(), @@ -851,7 +548,6 @@ mod tests { #[tokio::test] async fn upstream_401_operator_collapses_to_502() { - // Managed-key mode: a 401 collapses to 502 (operator's key is bad). let err = RouteError::Upstream { status: StatusCode::UNAUTHORIZED, body: "invalid api key".into(), @@ -933,8 +629,6 @@ mod tests { assert_eq!(ct, "application/json"); } - // -- MEDIUM-7: timeout / payload-too-large variants ----------------------- - #[test] fn request_timeout_returns_408() { let response = route_error_response(ClientProtocol::Anthropic, &RouteError::RequestTimeout); @@ -978,8 +672,6 @@ mod tests { ); } - // -- LOW-27: MethodNotAllowed --------------------------------------------- - #[test] fn method_not_allowed_returns_405() { let response = @@ -1017,15 +709,6 @@ mod tests { ); } - // -- route_error_variant_name ---------------------------------------------- - - // The event-log `ResponseFailed.error_kind` is documented as the RouteError - // variant name (a stable discriminant), NOT the client-facing protocol - // error-type string. Pin a representative set so distinct variants that - // share an envelope string (e.g. UnknownProvider vs NotFound, both - // "not_found_error" to the client) stay separable downstream - // (audit finding: error_kind not variant name). - #[test] fn route_error_variant_name_returns_variant_identifier() { assert_eq!( @@ -1057,22 +740,12 @@ mod tests { }), "Upstream" ); - // The client-facing type string is a DIFFERENT value (proving the two - // are not accidentally the same source). let (status, client_type, _msg) = extract_error_fields(&RouteError::UnknownProvider("p".into())); assert_eq!(status, StatusCode::NOT_FOUND); assert_ne!(client_type, "UnknownProvider"); } - // -- map_upstream_status --------------------------------------------------- - - // A conservative allowlist of clearly-client-owned upstream codes is - // forwarded verbatim so clients can act on their own bad request / - // size / rate-limit failures (audit LOW, upstream 4xx mapping). Auth - // codes (401/403) collapse to 502 for operator-owned credentials and pass - // through only for client-owned (passthrough_auth) credentials. - #[test] fn upstream_400_passes_through() { assert_eq!( @@ -1083,9 +756,6 @@ mod tests { #[test] fn upstream_401_operator_maps_to_502() { - // In managed-key mode an upstream 401 means the operator's key is bad, - // not the client's credentials -- collapse to 502 so the client is not - // misattributed an operator/infrastructure issue. assert_eq!( map_upstream_status(StatusCode::UNAUTHORIZED, AuthOwner::Operator), StatusCode::BAD_GATEWAY @@ -1094,9 +764,6 @@ mod tests { #[test] fn upstream_401_client_passes_through() { - // Under passthrough_auth the client's own token is forwarded, so an - // upstream 401 is the client's credential problem and passes through - // verbatim (audit LOW, passthrough-auth-401-collapse). assert_eq!( map_upstream_status(StatusCode::UNAUTHORIZED, AuthOwner::Client), StatusCode::UNAUTHORIZED @@ -1137,9 +804,6 @@ mod tests { #[test] fn upstream_404_maps_to_502() { - // Ambiguous (model-not-found vs unknown path): kept as 502 to avoid - // leaking provider-specific schemas; logged at warn server-side. 404 is - // never client-owned, so it collapses regardless of auth owner. assert_eq!( map_upstream_status(StatusCode::NOT_FOUND, AuthOwner::Client), StatusCode::BAD_GATEWAY @@ -1154,8 +818,6 @@ mod tests { ); } - // -- truncate_error_body --------------------------------------------------- - #[test] fn truncate_short_body_unchanged() { let body = "short error"; @@ -1167,7 +829,6 @@ mod tests { let body = "x".repeat(600); let truncated = truncate_error_body(&body); assert!(truncated.ends_with("...[truncated]")); - // The suffix is included within the MAX_ERROR_MESSAGE_LEN budget. assert!( truncated.len() <= MAX_ERROR_MESSAGE_LEN, "truncated body must not exceed MAX_ERROR_MESSAGE_LEN, got {}", @@ -1177,11 +838,8 @@ mod tests { #[test] fn truncate_exact_length_is_unchanged() { - // Exactly at the limit should NOT be truncated. let body = "a".repeat(MAX_ERROR_MESSAGE_LEN); assert_eq!(truncate_error_body(&body), body); - - // One over the limit should be truncated. let body = "a".repeat(MAX_ERROR_MESSAGE_LEN + 1); let truncated = truncate_error_body(&body); assert!(truncated.ends_with("...[truncated]")); @@ -1190,15 +848,12 @@ mod tests { #[test] fn truncate_multibyte_at_boundary_does_not_panic() { - // Japanese characters are 3 bytes each in UTF-8. - let body = "あ".repeat(200); // 600 bytes, exceeds 512 + let body = "あ".repeat(200); let truncated = truncate_error_body(&body); assert!(truncated.len() <= MAX_ERROR_MESSAGE_LEN); assert!(truncated.ends_with("...[truncated]")); } - // -- Response body content-type -------------------------------------------- - #[test] fn anthropic_error_has_json_content_type() { let err = RouteError::InvalidRequest("test".into()); @@ -1210,8 +865,6 @@ mod tests { assert_eq!(ct, "application/json"); } - // -- Response body JSON structure ----------------------------------------- - #[tokio::test] async fn anthropic_invalid_request_body_has_correct_structure() { let err = RouteError::InvalidRequest("bad input".into()); @@ -1245,9 +898,6 @@ mod tests { #[tokio::test] async fn anthropic_upstream_500_body_has_correct_structure() { - // A 500 collapses to 502 with a GENERIC message; the upstream body is - // logged server-side only and never reaches the client (mirror of the - // OpenAI test; audit finding: 502-collapse message leak). let err = RouteError::Upstream { status: StatusCode::INTERNAL_SERVER_ERROR, body: "model 'gpt-foo' is overloaded: request echoed here".into(), @@ -1312,16 +962,12 @@ mod tests { ); } - // -- ClientProtocol equality ----------------------------------------------- - #[test] fn client_protocol_equality() { assert_eq!(ClientProtocol::Anthropic, ClientProtocol::Anthropic); assert_ne!(ClientProtocol::Anthropic, ClientProtocol::OpenAiChat); } - // -- openai_stream_error_json_with_type (finding 55) ----------------------- - #[test] fn stream_error_json_produces_valid_structure() { let json_str = openai_stream_error_json_with_type("test error", "server_error"); diff --git a/crates/llm-proxy-server/src/routes/health.rs b/crates/llm-proxy-server/src/routes/health.rs index 0511262..021b46d 100644 --- a/crates/llm-proxy-server/src/routes/health.rs +++ b/crates/llm-proxy-server/src/routes/health.rs @@ -2,30 +2,11 @@ use axum::{Json, extract::State, http::StatusCode}; use serde::Serialize; use crate::state::AppState; - -/// Health check response body. -/// -/// Deliberately trivial (`{status, service}`): `/health` is the unauthenticated -/// liveness probe hammered by load balancers, so it must not expose operational -/// telemetry (audit LOW-13). Aggregate counters live on the in-process -/// [`Metrics`](llm_proxy_core::Metrics) snapshot for observability tooling, not -/// on this public route; a future authenticated admin or Prometheus endpoint -/// could surface them without putting them on the liveness probe. -/// -/// Note: `#[serde(deny_unknown_fields)]` is intentionally omitted because these -/// structs are outbound-only (`Serialize`, never `Deserialize` from external -/// input). The attribute has no runtime effect on Serialize-only types. #[derive(Serialize)] pub(crate) struct HealthBody { status: &'static str, service: String, } - -/// Liveness probe. -/// -/// Returns a minimal `{status, service}` body. No operational metrics are -/// exposed: the endpoint is unauthenticated and serves as a liveness probe, so -/// it must not leak telemetry (audit LOW-13). Query parameters are ignored. pub(crate) async fn health(State(state): State) -> (StatusCode, Json) { ( StatusCode::OK, @@ -35,45 +16,20 @@ pub(crate) async fn health(State(state): State) -> (StatusCode, Json (StatusCode, Json) { (StatusCode::OK, Json(ReadyBody { status: "ready" })) } - -/// Version response body. -/// -/// Note: `#[serde(deny_unknown_fields)]` is intentionally omitted (Serialize-only type). #[derive(Serialize)] pub(crate) struct VersionBody { - /// Public server name, from `config.server_name`. name: String, version: &'static str, target: &'static str, git_sha: &'static str, } - -/// Build metadata for ops/debugging. `name` comes from `config.server_name` -/// (via `state.server_name()`) so operators can distinguish deployments; the -/// rest is compile-time build info. `BuildInfo.name` (the Cargo package name) -/// is not used here because it is a static constant that does not vary between -/// deployments. pub(crate) async fn version(State(state): State) -> (StatusCode, Json) { let build = state.build_info(); ( diff --git a/crates/llm-proxy-server/src/routes/messages.rs b/crates/llm-proxy-server/src/routes/messages.rs index 5a98ea1..f718ba8 100644 --- a/crates/llm-proxy-server/src/routes/messages.rs +++ b/crates/llm-proxy-server/src/routes/messages.rs @@ -1,50 +1,27 @@ -//! `/providers/{provider}/v1/messages` handler using the core pipeline. -//! -//! The handler parses the incoming Anthropic `MessageRequest`, decodes it into -//! a `CoreRequest` via the Anthropic client -//! adapter, and then dispatches through the shared core pipeline -//! ([`handle_core_once`] or [`handle_core_stream`]). -//! -//! Architecture boundary: route handlers do not perform scenario detection, -//! endpoint classification, fallback routing, or provider-specific streaming. - use std::sync::Arc; use axum::Json; use axum::body::Body; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use llm_proxy_core::ProviderRouteKind; use llm_proxy_protocol::anthropic::MessageRequest; use llm_proxy_protocol::client::anthropic; use tracing::{info, warn}; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; use super::error_response::{ClientProtocol, RouteError, route_error_response}; - -/// POST `/providers/{provider}/v1/messages` -/// -/// Accepts an Anthropic-format [`MessageRequest`], decodes it through the -/// core pipeline, and returns an Anthropic-shaped response. pub(crate) async fn handle_messages( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { - // The wrapper only renders: it does NOT emit ResponseFailed. Each EARLY-GATE - // error return inside the inner function emits exactly one ResponseFailed - // itself (the sites before the core pipeline dispatch), and the core pipeline - // emits its own single ResponseFailed for failures that occur inside it. A - // blanket emit here would double-count every pipeline failure (the pipeline - // emits AND returns Err via `?`), so the wrapper must stay emit-free - // (audit route-responsefailed-gaps regression). - match handle_messages_inner(state, req_id, provider, connect_info, headers, body).await { + match handle_messages_inner(state, meta, provider, headers, body).await { Ok(response) => response, Err(error) => { warn!(error = %error, "request failed"); @@ -52,30 +29,16 @@ pub(crate) async fn handle_messages( } } } - -/// Inner handler that returns `Result` so errors can be mapped uniformly. -/// -/// ResponseFailed emission discipline (audit route-responsefailed-gaps): -/// every EARLY-GATE error return (the sites BEFORE the core pipeline dispatch) -/// emits exactly one ResponseFailed here. Failures INSIDE the core pipeline -/// emit their own single ResponseFailed internally and return Err, which this -/// function propagates without re-emitting. The outer wrapper renders only. async fn handle_messages_inner( state: AppState, - req_id: RequestId, + meta: RequestMeta, provider: String, - connect_info: Option, headers: HeaderMap, body: axum::body::Bytes, ) -> Result, RouteError> { let start = std::time::Instant::now(); let event_bus = Arc::clone(&state.event_bus); - // The request id is consumed by prepare_request below; snapshot it once so - // every early-gate emit references the same id. - let request_id = req_id.0.clone(); - // Every early gate emits exactly one ResponseFailed with the same - // provider/request/start triple; bind it once instead of repeating the - // six-argument call at each gate. + let request_id = meta.request_id.0.clone(); let emit_failed = |e: &RouteError| { core_pipeline::emit_response_failed( &event_bus, @@ -86,40 +49,16 @@ async fn handle_messages_inner( start, ); }; - // Input validation gate (mirrors token_count.rs): validate the provider - // name and confirm the provider exists BEFORE charging rate-limit/dedup - // counters in `prepare_request`. This ensures malformed probes (bad path, - // unknown provider) receive an immediate 400/404 without burning rate-limit - // budget or surfacing as a misleading 429/409 (audit LOW-29). core_pipeline::validate_provider_name(&provider).inspect_err(&emit_failed)?; if state.providers().get(&provider).is_none() { - // Unknown-provider early gate: emit exactly one ResponseFailed before - // returning (audit route-responsefailed-gaps). let err = RouteError::UnknownProvider(provider.clone()); emit_failed(&err); return Err(err); } - - // Reject non-JSON Content-Type before parsing, so a wrong media type is not - // misreported as a JSON syntax error (audit LOW-30). core_pipeline::validate_json_content_type(&headers).inspect_err(&emit_failed)?; - - // Pre-flight: rate limit, dedup, request ID. let request_path = format!("/providers/{provider}/v1/messages"); - let ctx = core_pipeline::prepare_request( - &state, - req_id.0, - &headers, - connect_info.as_ref(), - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; - - // Parse the Anthropic MessageRequest via the axum Json helper so that - // syntax-vs-data parse failures are distinguished into precise messages - // instead of collapsing both into one opaque "invalid JSON" string - // (audit LOW-28). + let ctx = core_pipeline::prepare_request(&state, &meta, &headers, &body, &request_path) + .inspect_err(&emit_failed)?; let req: MessageRequest = match Json::::from_bytes(&body) { Ok(json) => json.0, Err(rejection) => { @@ -128,15 +67,9 @@ async fn handle_messages_inner( return Err(err); } }; - - // Defense-in-depth: validate() checks for empty model/messages before - // decode_request also validates the same fields. This catches issues - // early with a clearer error message. req.validate() .map_err(|e| RouteError::InvalidRequest(e.to_string())) .inspect_err(&emit_failed)?; - - // Decode the Anthropic request into a core request. let core = anthropic::decode_request(req) .map_err(core_pipeline::protocol_error_to_route) .inspect_err(&emit_failed)?; @@ -149,43 +82,19 @@ async fn handle_messages_inner( streaming = is_streaming, "decoded Anthropic request into CoreRequest" ); - - // Extract the inbound client auth token (used by passthrough-auth providers). let inbound_auth = core_pipeline::extract_inbound_auth(&headers); + let dispatch = core_pipeline::CoreDispatch { + provider_name: &provider, + route_kind: ProviderRouteKind::Messages, + client_protocol: ClientProtocol::Anthropic, + }; if is_streaming { - core_pipeline::handle_core_stream( - state, - ctx, - &provider, - ProviderRouteKind::Messages, - core, - ClientProtocol::Anthropic, - inbound_auth, - ) - .await + core_pipeline::handle_core_stream(state, ctx, dispatch, core, inbound_auth).await } else { - core_pipeline::handle_core_once( - state, - ctx, - &provider, - ProviderRouteKind::Messages, - core, - ClientProtocol::Anthropic, - inbound_auth, - ) - .await + core_pipeline::handle_core_once(state, ctx, dispatch, core, inbound_auth).await } } - -/// Map an axum [`JsonRejection`] into a [`RouteError::InvalidRequest`] with a -/// distinct message for syntax errors (malformed JSON) versus data errors -/// (well-formed JSON that failed deserialization). -/// -/// `from_bytes` cannot produce `MissingJsonContentType` (the content-type -/// header is only inspected by the `FromRequest` extractor), and `BytesRejection` -/// only arises from the request extraction path — neither is reachable here, so -/// the wildcard arm covers any future `#[non_exhaustive]` variant defensively. fn json_rejection_to_route_error(rejection: axum::extract::rejection::JsonRejection) -> RouteError { match rejection { axum::extract::rejection::JsonRejection::JsonSyntaxError(e) => { @@ -198,14 +107,8 @@ fn json_rejection_to_route_error(rejection: axum::extract::rejection::JsonReject } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { - // -- Architecture source guard --------------------------------------------- - #[test] fn source_guard_enforces_route_boundaries() { let source = include_str!("messages.rs"); diff --git a/crates/llm-proxy-server/src/routes/mod.rs b/crates/llm-proxy-server/src/routes/mod.rs index 1b67671..b0c59fe 100644 --- a/crates/llm-proxy-server/src/routes/mod.rs +++ b/crates/llm-proxy-server/src/routes/mod.rs @@ -16,9 +16,10 @@ use tower_http::{ }; use tracing::info_span; -use crate::middleware::{RequestId, RequestIdGenerator}; +use crate::middleware::{CallerId, RequestId, RequestIdGenerator, extract_bearer_token}; use crate::state::AppState; use llm_proxy_storage::EventBus; +use secrecy::ExposeSecret; mod chat; mod core_pipeline; @@ -33,66 +34,18 @@ use health::{health, ready, version}; use messages::handle_messages; use models::handle_models; use token_count::count_tokens; - -/// Maximum request body size for API routes (32 MiB). -/// -/// This limit applies to all `/providers/{provider}/v1/*` routes via -/// `DefaultBodyLimit`. The value is chosen to accommodate large multi-turn -/// conversations with tool-use payloads while protecting against resource -/// exhaustion. Requests exceeding this limit receive a 413 Payload Too Large -/// response before any JSON parsing or upstream forwarding occurs. -/// -/// Note: This limit does NOT apply to lightweight health/readiness/version -/// endpoints, which accept no request body. const MAX_BODY_BYTES: usize = 32 * 1024 * 1024; - -/// Maximum body size to drain for 404 fallback responses. -/// -/// Deliberately smaller than [`MAX_BODY_BYTES`] because 404 responses are -/// never processed -- the body is only drained to clean up the connection. -/// 1 KiB is enough to read any reasonable probe or small payload while -/// avoiding wasting memory on large accidental uploads to wrong paths. const NOT_FOUND_BODY_DRAIN_LIMIT: usize = 1024; -/// Build the full router. -/// -/// Middleware order (outermost to innermost per group): -/// -/// **Lightweight routes** (`/health`, `/ready`, `/version`): -/// 1. `TraceLayer` -- logs every request and response, measures latency. -/// -/// **API routes** (`/providers/{provider}/v1/*`): -/// 1. `DefaultBodyLimit` -- caps the request body for JSON extractors (runs -/// first so oversized payloads are rejected before the timeout starts). -/// 2. `TraceLayer` -- logs every request and response, measures latency. -/// 3. `TimeoutLayer` -- cancels requests exceeding `config.request_timeout`. -/// -/// `TimeoutLayer` is scoped to provider API routes only. Lightweight health/readiness -/// endpoints respond instantly and must not be subject to a 408 timeout, -/// which would confuse orchestrators (Kubernetes, load balancers). -/// -/// For streaming handlers, `TimeoutLayer` covers work through creation of the -/// HTTP response, including the first upstream event. Once Axum returns the -/// streaming response body, this layer no longer times subsequent SSE events. +/// Build the full HTTP router: `/health`, `/ready`, `/version`, provider API +/// routes under `/providers/{provider}/v1/*`, and the 404/405 fallbacks. The +/// admin `/usage` surface merges in only under the `admin_on_main_listener` opt-in. pub fn router(state: AppState) -> Router { let timeout = state.request_timeout(); - // Clone the shared request-id generator before `state` is moved into the - // router. The outermost `inject_request_id` layer generates one id per - // request and stamps it into extensions so the tracing span, the handler, - // and the `x-request-id` header all share a single id (audit MEDIUM-1). let id_gen = Arc::clone(&state.request_id_gen); - // Clone the shared event bus before `state` is moved into the router, so the - // error-normalisation layer can emit a `ResponseFailed` for 408/413 - // rejections produced by the timeout/body-limit layers (audit finding: layer - // fallbacks emit no terminal event). let event_bus = Arc::clone(&state.event_bus); - - // Span factory shared by both routers: it reads the id injected by the - // outermost layer so every request span carries `request_id`, matching the - // id that later appears in handler log events and the response header. + let state_arc = state.clone(); let trace = TraceLayer::new_for_http().make_span_with(make_request_span); - - // Lightweight routes: request-id injection + tracing only. let lightweight = Router::new() .route("/health", get(health)) .route("/ready", get(ready)) @@ -102,27 +55,9 @@ pub fn router(state: AppState) -> Router { .layer(from_fn_with_state(Arc::clone(&id_gen), inject_request_id)) .layer(trace.clone()), ); - - // API routes: request-id injection, error normalisation, body limit, - // tracing, timeout. - // - // Layer order (outermost first): - // 1. `inject_request_id` -- generates the id and stamps it into extensions - // so the span (step 4) and the handler see the same id. - // 2. `normalize_error_responses` -- rewrites the 408 (timeout) and 413 - // (body-limit) rejections emitted by layers below into protocol-shaped - // JSON envelopes carrying `x-request-id` (audit MEDIUM-7). It runs - // outside the body-limit/timeout layers so their rejections pass - // through it on the way out. - // 3. `DefaultBodyLimit` -- caps the body the JSON extractors can buffer. - // 4. `TraceLayer` -- opens the request span (carrying the request id). - // 5. `TimeoutLayer` -- cancels handler work exceeding `request_timeout`. - // - // Because the timeout and body-limit layers emit non-JSON, header-less - // rejections, step 2 normalises their final responses into the same - // Anthropic/OpenAI-shaped JSON schema every handler uses. let api_middleware = ServiceBuilder::new() .layer(from_fn_with_state(id_gen, inject_request_id)) + .layer(from_fn_with_state(state_arc, require_client_auth)) .layer(from_fn_with_state(event_bus, normalize_error_responses)) .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .layer(trace) @@ -148,52 +83,31 @@ pub fn router(state: AppState) -> Router { .merge(lightweight) .merge(api) .fallback(not_found) - // Protocol-aware 405: shape method-not-allowed responses through the - // same error envelope as 404/408/413 instead of axum's empty body - // (audit LOW-27). .method_not_allowed_fallback(method_not_allowed); - - // Config-gated CORS (audit LOW-12). A `CorsLayer` is installed only when - // the operator explicitly lists cross-origin callers, and then only for - // those exact origins. This replaces a former `CorsLayer::very_permissive()` - // that admitted every origin -- the opposite of the finding's intent. The - // default (`server.allowed_origins` unset) applies no CORS layer, so - // browsers enforce a same-origin policy and this server-to-server proxy is - // not unexpectedly reachable from arbitrary web origins. + let usage_on_main = match ( + crate::admin_router(&state), + state.app_config.server.admin_on_main_listener, + ) { + (Some(usage), true) => Some(usage), + _ => None, + }; let router = match cors_layer_for(state.allowed_origins()) { Some(cors) => router.layer(cors), None => router, }; - router - // Global, outermost middleware (audit LOW-11, LOW-12) -- applied to the - // merged router so they cover every route and both fallbacks. CORS sits - // just inside these so a preflight still receives the headers below and - // panic protection. - // 1. `SetResponseHeaderLayer` -- `X-Content-Type-Options: nosniff` - // on every response so JSON error bodies are not MIME-sniffed - // into executable types (LOW-12). - // 2. `CatchPanicLayer` -- outermost (LOW-11): converts a panic in - // any inner layer or handler into a 500 instead of dropping the - // connection (which would otherwise present as a reset to the - // client and bypass `TraceLayer`'s response-span logging). + let router: Router = router .layer(SetResponseHeaderLayer::if_not_present( header::X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"), )) .layer(CatchPanicLayer::new()) - .with_state(state) + .with_state(state); + match usage_on_main { + Some(usage) => router.merge(usage), + None => router, + } } - -/// Build a restrictive `CorsLayer` for an explicit allow-origin list, or return -/// `None` when cross-origin access is disabled (audit LOW-12). -/// -/// `None` (no origins configured, or every entry fails to parse as a header -/// value) means "install no CORS layer": browsers then enforce a same-origin -/// default, which is the safe baseline for a server-to-server proxy. When -/// origins are configured the layer echoes `Access-Control-Allow-Origin` only -/// for those exact origins and answers preflight with the proxy's actual -/// methods plus the request headers the LLM SDKs it fronts commonly send. fn cors_layer_for(origins: Option<&[String]>) -> Option { let origins = origins?; let parsed: Vec = origins @@ -209,12 +123,7 @@ fn cors_layer_for(origins: Option<&[String]>) -> Option { if parsed.is_empty() { return None; } - // The methods the proxy actually serves; OPTIONS is covered by the CORS - // preflight the layer itself synthesizes, but is listed so the preflight - // response advertises it as permitted. let methods = [Method::GET, Method::POST, Method::OPTIONS]; - // Request headers the LLM SDKs the proxy fronts commonly send. Kept to a - // curated set rather than `Any` so the policy stays restrictive. let headers = [ header::CONTENT_TYPE, header::AUTHORIZATION, @@ -233,13 +142,6 @@ fn cors_layer_for(origins: Option<&[String]>) -> Option { .allow_headers(headers), ) } - -/// Outermost middleware: mint one request id per request and store it in -/// extensions as a [`RequestId`]. -/// -/// Running this before `TraceLayer` lets [`make_request_span`] stamp the same -/// id onto the span that the handler later emits in events and attaches to the -/// `x-request-id` header, restoring end-to-end correlation (audit MEDIUM-1). async fn inject_request_id( State(id_gen): State>, mut req: Request, @@ -249,18 +151,73 @@ async fn inject_request_id( req.extensions_mut().insert(RequestId(id)); next.run(req).await } +async fn require_client_auth(State(state): State, req: Request, next: Next) -> Response { + if !state.client_auth.is_enforced() { + let mut req = req; + req.extensions_mut().insert(CallerId::anonymous()); + return next.run(req).await; + } + let connect_info = req + .extensions() + .get::>() + .map(|ci| ci.0); + let client_ip = crate::middleware::get_client_ip( + req.headers(), + connect_info.as_ref(), + state.trust_forwarded_headers(), + ); + + let authenticated = match extract_bearer_token(req.headers()) { + Some(token) => { + if !state.rate_limiter.would_allow(&client_ip) { + state.metrics.record_rate_limited(); + None + } else { + state.client_auth.authenticate(token.expose_secret()).await + } + } + None => None, + }; -/// Build the per-request tracing span, carrying the injected request id. -/// -/// `make_span_with` runs after [`inject_request_id`] has stamped the id into -/// extensions, so the span records the same id that flows to the handler. If -/// the id is somehow absent (e.g. a request that bypassed the injection -/// layer), the span still opens with an empty id rather than panicking. -/// -/// Only the request **path** is recorded, never the full URI (audit GAP-LOW-3). -/// Several LLM SDKs fall back to authenticating via `?api_key=`/`?key=` query -/// parameters; logging the full URI would leak those secrets into the tracing -/// span and any downstream log sink. `req.uri().path()` strips the query. + let mut req = req; + match authenticated { + Some(key_id) => { + tracing::debug!(caller = %key_id, "client authenticated"); + req.extensions_mut().insert(CallerId::authenticated(key_id)); + next.run(req).await + } + None => { + if !state.rate_limiter.is_allowed(&client_ip) { + state.metrics.record_rate_limited(); + } + let request_id = req + .extensions() + .get::() + .cloned() + .unwrap_or_else(|| RequestId("unknown".to_owned())); + let path = req.uri().path().to_owned(); + if let Err(e) = axum::body::to_bytes(req.into_body(), NOT_FOUND_BODY_DRAIN_LIMIT).await + { + tracing::trace!(error = %e, "draining body of auth-rejected request failed"); + } + let error = error_response::RouteError::Unauthorized; + core_pipeline::emit_response_failed( + &state.event_bus, + &request_id.0, + None, + None, + &error, + std::time::Instant::now(), + ); + let mut response = + error_response::route_error_response(protocol_for_path(&path), &error); + if let Ok(value) = request_id.0.parse() { + response.headers_mut().insert("x-request-id", value); + } + response + } + } +} fn make_request_span(req: &Request) -> tracing::Span { let id = req.extensions().get::(); info_span!( @@ -270,23 +227,11 @@ fn make_request_span(req: &Request) -> tracing::Span { path = %req.uri().path(), ) } - -/// Normalise the non-JSON error bodies emitted by the timeout and body-limit -/// layers into protocol-shaped JSON (audit MEDIUM-7). -/// -/// `TimeoutLayer` returns a bare 408 with an empty body, and `DefaultBodyLimit` -/// surfaces a 413 as a plain-text `BytesRejection` from the extractors. Neither -/// carries the `x-request-id` header (that is attached inside the handlers). -/// This layer intercepts both statuses, rebuilds them via -/// [`error_response::route_error_response`] with the correct protocol envelope, -/// and stamps the `x-request-id` header from extensions so the 408/413 paths -/// match the schema every handler returns. async fn normalize_error_responses( State(event_bus): State>, req: Request, next: Next, ) -> Response { - // Capture everything we need before the request is consumed by `next`. let start = Instant::now(); let protocol = protocol_for_path(req.uri().path()); let request_id = req.extensions().get::().cloned(); @@ -296,9 +241,6 @@ async fn normalize_error_responses( if status != StatusCode::REQUEST_TIMEOUT && status != StatusCode::PAYLOAD_TOO_LARGE { return response; } - - // Drain the (empty or small) layer-produced body so the connection is not - // left half-written; the replacement body is what the client sees. if let Err(e) = axum::body::to_bytes(response.into_body(), NOT_FOUND_BODY_DRAIN_LIMIT).await { tracing::trace!(error = %e, "draining layer rejection body failed"); } @@ -308,14 +250,6 @@ async fn normalize_error_responses( } else { error_response::RouteError::PayloadTooLarge }; - - // Emit a `ResponseFailed` so a 408 timeout or 413 oversize-body rejection - // is attributable in the structured event log. These rejections are produced - // by the timeout/body-limit layers BELOW this one, so the request handler - // never runs and would otherwise leave no terminal event. `inject_request_id` - // (layer 1, outer) has already stamped the id, so it is present here - // (audit finding: layer fallbacks emit no terminal event). provider/model - // are unknown at this layer (pre-dispatch). let id_str = request_id .as_ref() .map(|id| id.0.as_str()) @@ -330,15 +264,6 @@ async fn normalize_error_responses( } rewritten } - -/// Select the client protocol for a request path, reusing the same prefix -/// heuristic as [`not_found`]. -/// -/// The timeout / body-limit layers sit above routing, so on rejection they -/// cannot know which protocol's envelope to render. This helper mirrors the -/// 404 heuristic (`/v1/chat/...` -> OpenAI, otherwise Anthropic) so the -/// normalised 408/413 bodies stay consistent with the handler the request -/// *would* have reached. fn protocol_for_path(path: &str) -> error_response::ClientProtocol { let is_openai_chat_path = path.contains("/v1/chat/completions") || path.contains("/v1/chat/edits"); @@ -350,26 +275,6 @@ fn protocol_for_path(path: &str) -> error_response::ClientProtocol { } async fn not_found(State(state): State, req: Request) -> impl IntoResponse { - // Protocol-aware 404: return OpenAI-shaped errors for provider-scoped chat - // routes and any path under /v1/chat/ (future OpenAI chat sub-routes), and - // Anthropic-shaped errors for everything else. - // - // Boundary note: the path check covers future routes - // like /v1/chat/edits. If a non-OpenAI protocol is ever mounted under - // /v1/chat/, this heuristic must be updated. - // - // NOTE: This heuristic is fragile and depends on path-prefix conventions - // rather than a registered route registry. A more robust approach would - // be to maintain a route-to-protocol mapping that is consulted by the - // fallback handler. For now, the prefix-based approach is sufficient - // because only two protocols are mounted and their paths are disjoint. - - // Capture the request id BEFORE the body is consumed so the response - // carries the same correlation header every handler-emitted status does. - // The id is injected by `inject_request_id`; for a request that bypassed - // that layer (e.g. a fully-unknown path reaching the top-level fallback - // outside the sub-router wrapping), the extension is absent and we fall - // back to "unknown" (audit MEDIUM-1). let start = Instant::now(); let request_id = req .extensions() @@ -378,22 +283,11 @@ async fn not_found(State(state): State, req: Request) -> impl IntoResp .unwrap_or_else(|| RequestId("unknown".to_owned())); let path = req.uri().path().to_owned(); - - // Drain the body to ensure the connection is cleaned up promptly. - // We attempt to read up to NOT_FOUND_BODY_DRAIN_LIMIT bytes. If the body - // is larger, the remaining bytes are discarded when the request body is - // dropped. This is acceptable because: (a) the 404 response will cause - // the client to close the connection, and (b) HTTP/2 multiplexing does - // not have the same head-of-line blocking concern as HTTP/1.1 keep-alive. if let Err(e) = axum::body::to_bytes(req.into_body(), NOT_FOUND_BODY_DRAIN_LIMIT).await { tracing::trace!(error = %e, "body drain in 404 handler failed"); } let error = error_response::RouteError::NotFound; - // Emit a `ResponseFailed` so a 404 is attributable in the structured event - // log (audit finding: layer fallbacks emit no terminal event). provider/model - // are None (pre-dispatch); the request id falls back to "unknown" for paths - // that were never assigned one. core_pipeline::emit_response_failed(&state.event_bus, &request_id.0, None, None, &error, start); let protocol = protocol_for_path(&path); @@ -403,28 +297,7 @@ async fn not_found(State(state): State, req: Request) -> impl IntoResp } response } - -/// Protocol-aware 405 Method Not Allowed fallback (audit LOW-27). -/// -/// A request that hits a registered path with a disallowed method (e.g. GET to -/// a POST-only route) does **not** reach [`not_found`] -- axum's per-route -/// `MethodRouter` returns a stock 405 with an empty body. Without this -/// fallback the client receives an inconsistent, non-protocol-shaped response. -/// -/// This handler mirrors [`not_found`]: it drains the body, selects the client -/// protocol via the same path-prefix heuristic, and renders a -/// protocol-shaped 405 envelope. Axum still sets the `Allow` header -/// automatically unless the response sets it, which it does not, so the -/// header is preserved. -/// -/// Delegates to [`error_response::route_error_response`] with -/// [`error_response::RouteError::MethodNotAllowed`] so the 405 body is produced -/// by the single shared encoder every other status code uses (audit LOW-27), -/// rather than a bespoke inline envelope. async fn method_not_allowed(State(state): State, req: Request) -> impl IntoResponse { - // Capture the request id BEFORE the body is consumed so the 405 response - // carries the same correlation header every handler-emitted status does, - // matching [`not_found`] (audit MEDIUM-1). let start = Instant::now(); let request_id = req .extensions() @@ -433,16 +306,11 @@ async fn method_not_allowed(State(state): State, req: Request) -> impl .unwrap_or_else(|| RequestId("unknown".to_owned())); let path = req.uri().path().to_owned(); - - // Drain the body for the same connection-cleanup reason as [`not_found`]. if let Err(e) = axum::body::to_bytes(req.into_body(), NOT_FOUND_BODY_DRAIN_LIMIT).await { tracing::trace!(error = %e, "body drain in 405 handler failed"); } let error = error_response::RouteError::MethodNotAllowed; - // Emit a `ResponseFailed` so a 405 is attributable in the structured event - // log (audit finding: layer fallbacks emit no terminal event), mirroring - // [`not_found`]. core_pipeline::emit_response_failed(&state.event_bus, &request_id.0, None, None, &error, start); let protocol = protocol_for_path(&path); @@ -470,8 +338,6 @@ mod tests { #[test] fn cors_layer_absent_when_every_origin_is_invalid() { - // Newlines and NUL are not legal header-value bytes, so every entry is - // skipped and no layer is installed (audit LOW-12). let bad = vec!["bad\norigin".to_owned(), "\0".to_owned()]; assert!(cors_layer_for(Some(&bad)).is_none()); } @@ -491,3 +357,115 @@ mod tests { assert!(cors_layer_for(Some(&origins)).is_some()); } } +#[cfg(test)] +mod auth_gate_tests { + use super::*; + + use axum::body::Body; + use axum::http::Request; + use tower::ServiceExt; + fn auth_router(rate_limit_rpm: u32) -> (Router, AppState) { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + b"sekrit-token", + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + + let mut clients = std::collections::HashMap::new(); + clients.insert("laptop".to_owned(), secrecy::SecretString::from(phc)); + let app_config = llm_proxy_core::AppConfig { + server: llm_proxy_core::ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: std::time::Duration::from_secs(300), + shutdown_timeout: std::time::Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "gate-test".to_owned(), + rate_limit_rpm, + trust_forwarded_headers: false, + dedup_window: std::time::Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }, + storage: None, + }; + + let state = crate::state::AppState::new( + app_config, + llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"), + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::new(), + crate::state::BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ); + + let router = Router::new() + .route( + "/providers/x/v1/messages", + post(|| async { StatusCode::OK }), + ) + .layer(from_fn_with_state( + state.clone(), + super::require_client_auth, + )); + (router, state) + } + + fn bearer_request(token: &str) -> Request { + Request::builder() + .method("POST") + .uri("/providers/x/v1/messages") + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() + } + + #[tokio::test] + async fn valid_token_passes_when_the_peer_has_budget() { + let (app, _state) = auth_router(100); + let response = app.oneshot(bearer_request("sekrit-token")).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn exhausted_failure_budget_gates_verification_before_argon2() { + let (app, state) = auth_router(1); + let first = app + .clone() + .oneshot(bearer_request("wrong-token")) + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::UNAUTHORIZED); + assert_eq!(state.metrics.get_snapshot().rate_limited, 0); + let second = app.oneshot(bearer_request("sekrit-token")).await.unwrap(); + assert_eq!( + second.status(), + StatusCode::UNAUTHORIZED, + "a peer with an exhausted failure budget must be gated before verification" + ); + assert!( + state.metrics.get_snapshot().rate_limited >= 1, + "the gate records a limiter rejection" + ); + } +} diff --git a/crates/llm-proxy-server/src/routes/models.rs b/crates/llm-proxy-server/src/routes/models.rs index 0314d93..f70249a 100644 --- a/crates/llm-proxy-server/src/routes/models.rs +++ b/crates/llm-proxy-server/src/routes/models.rs @@ -1,10 +1,3 @@ -//! `GET /providers/{provider}/v1/models` handler. -//! -//! Returns the merged static and discovered model catalog for a provider. -//! -//! The response is a normalized superset model card that includes both OpenAI -//! and Anthropic fields so either SDK can consume it. - use std::sync::Arc; use axum::body::Body; @@ -20,61 +13,25 @@ use crate::state::AppState; use super::core_pipeline; use super::error_response::{AuthOwner, ClientProtocol, RouteError, route_error_response}; - -// --------------------------------------------------------------------------- -// Response types -// --------------------------------------------------------------------------- - -/// Top-level response for the models endpoint. -/// -/// Follows the OpenAI `/v1/models` response shape as the base, with Anthropic -/// pagination fields added as a superset. -/// -/// Note: `#[serde(deny_unknown_fields)]` is intentionally omitted because this -/// is a Serialize-only type (never deserialized from external input). #[derive(Serialize)] struct ModelsResponse { - /// OpenAI-compatible: always `"list"`. object: &'static str, - /// OpenAI-compatible: the list of model cards. data: Vec, - /// Anthropic-compatible: whether more results exist. has_more: bool, - /// Anthropic-compatible: ID of the first model in the list. first_id: Option, - /// Anthropic-compatible: ID of the last model in the list. last_id: Option, } - -/// A normalized model card containing both OpenAI and Anthropic fields. -/// -/// Common SDKs should be able to ignore fields they do not use. -/// -/// Note: `#[serde(deny_unknown_fields)]` is intentionally omitted because this -/// is a Serialize-only type (never deserialized from external input). #[derive(Serialize)] struct ModelCard { - /// OpenAI: model identifier. id: String, - /// OpenAI: always `"model"`. object: &'static str, - /// OpenAI: Unix timestamp of creation. Set to the epoch (0) as a sentinel - /// value because the static catalog entries do not carry a creation timestamp. - /// A future improvement should propagate the discovered_at or generated_at - /// field from the catalog metadata. created: u64, - /// OpenAI: owner identifier (provider name). owned_by: String, - /// Anthropic: always `"model"`. #[serde(rename = "type")] model_type: &'static str, - /// Anthropic: human-readable display name. display_name: Option, - /// Anthropic-compatible creation timestamp, absent when unknown. created_at: Option, - /// The route kinds this model supports (non-standard extension). supports: Vec<&'static str>, - /// Maximum context length in tokens (non-standard extension). context_length: Option, } @@ -82,32 +39,11 @@ struct ModelCard { pub(super) struct ModelsQuery { refresh: Option, } - -// --------------------------------------------------------------------------- -// Handler -// --------------------------------------------------------------------------- - -/// `GET /providers/{provider}/v1/models` -/// -/// Returns the configured catalog, optionally refreshing discovery with -/// `?refresh=live`. When no catalog is configured, returns an empty list. -/// -/// The query string is parsed inside the handler (rather than via a -/// `Query` extractor) so a malformed query string -- e.g. a -/// duplicate `?refresh=a&refresh=b`, which axum's extractor rejects with a -/// plain-text 400 -- is mapped into the same OpenAI-shaped JSON envelope every -/// other status on this route returns, instead of an inconsistent bare-text -/// rejection that also lacks the `x-request-id` header. pub(crate) async fn handle_models( State(state): State, Path(provider): Path, req: axum::extract::Request, ) -> Response { - // The wrapper only renders: it does NOT emit ResponseFailed. /models has no - // core-pipeline dispatch, so each early-gate error return inside the inner - // function emits exactly one ResponseFailed itself. A blanket emit here - // would double-count, so the wrapper stays emit-free - // (audit route-responsefailed-gaps regression). match handle_models_inner(&state, &provider, req).await { Ok(response) => response, Err(error) => { @@ -116,16 +52,6 @@ pub(crate) async fn handle_models( } } } - -/// Inner handler that returns `Result` so errors can be mapped uniformly. -/// -/// ResponseFailed emission discipline (audit route-responsefailed-gaps): /models -/// has no core-pipeline dispatch, so each early-gate error return emits exactly -/// one ResponseFailed here. /models has no request-id extension (it does not -/// pass through the inference request-id middleware), so a stable sentinel -/// ("models") correlates the event to this endpoint; the provider path segment -/// is carried in the event. model is None because /models lists all models and -/// is not bound to a single ModelRef. async fn handle_models_inner( state: &AppState, provider_name: &str, @@ -134,10 +60,6 @@ async fn handle_models_inner( let start = std::time::Instant::now(); let event_bus = Arc::clone(&state.event_bus); let request_id = "models"; - // Surface the axum QueryRejection (a plain-text 400) through the shared - // JSON envelope so malformed-query responses are consistent with the rest - // of the API. Parsed inside the handler (not via a Query extractor) so the - // rejection is mapped into RouteError and emits one ResponseFailed. let query = Query::::try_from_uri(req.uri()) .map(|query| query.0) .map_err(|rejection| { @@ -166,7 +88,6 @@ async fn handle_models_inner( start, ) })?; - // Look up the provider. Returns 404 if not found. let provider = state .providers() .get(provider_name) @@ -197,11 +118,6 @@ async fn handle_models_inner( start, ) })?; - - // Build model cards from static entries. - // `entries` is fully owned (returned by value), so derive the pagination - // IDs from the owned Vec first, then consume it with `into_iter()` so - // every per-entry field moves into its `ModelCard` with zero clones. let first_id = entries.first().map(|m| m.id.clone()); let last_id = entries.last().map(|m| m.id.clone()); @@ -280,16 +196,10 @@ fn map_catalog_error(error: ProviderError) -> RouteError { | ProviderError::SseFraming(_) | ProviderError::EmptyResponse(_) | ProviderError::InvalidConfig(_) => RouteError::Internal(error.to_string()), - // ProviderError is #[non_exhaustive] so a wildcard arm is required - // to handle future variants added to the enum. _ => RouteError::Internal(error.to_string()), } } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { use super::*; @@ -298,11 +208,6 @@ mod tests { StaticModelCatalogEntry, }; use std::collections::HashMap; - - /// Build a GET request for the models endpoint with an optional query - /// suffix. Used so the in-handler query parser exercised by unit tests - /// receives a real `axum::extract::Request` (the `Query` extractor reads - /// the URI). fn models_request(query: &str) -> axum::extract::Request { let uri = if query.is_empty() { "/v1/models".to_owned() @@ -315,14 +220,13 @@ mod tests { .body(Body::empty()) .expect("request") } - - /// Build a minimal AppState with a provider that has a static catalog. fn build_state_with_catalog() -> AppState { let provider = ProviderConfig { name: "test-provider".to_owned(), api_key: secrecy::SecretString::from("sk-test"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -370,7 +274,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, llm_proxy_provider::ProviderAdapterRegistry::builtin(), @@ -383,14 +291,13 @@ mod tests { }, ) } - - /// Build a minimal AppState with a provider that has NO catalog. fn build_state_without_catalog() -> AppState { let provider = ProviderConfig { name: "no-catalog".to_owned(), api_key: secrecy::SecretString::from("sk-test"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -415,7 +322,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, llm_proxy_provider::ProviderAdapterRegistry::builtin(), @@ -445,22 +356,16 @@ mod tests { assert_eq!(json["object"], "list"); assert_eq!(json["data"].as_array().unwrap().len(), 2); assert_eq!(json["has_more"], false); - - // Catalog output is sorted by model ID for deterministic responses. let m1 = &json["data"][0]; assert_eq!(m1["id"], "claude-sonnet-4"); assert_eq!(m1["object"], "model"); assert_eq!(m1["owned_by"], "test-provider"); assert_eq!(m1["display_name"], "Claude Sonnet 4"); assert_eq!(m1["context_length"], 200000); - - // Second model let m2 = &json["data"][1]; assert_eq!(m2["id"], "deepseek-v3.1"); assert_eq!(m2["display_name"], "DeepSeek V3.1"); assert_eq!(m2["context_length"], 131072); - - // Pagination fields assert_eq!(json["first_id"], "claude-sonnet-4"); assert_eq!(json["last_id"], "deepseek-v3.1"); } diff --git a/crates/llm-proxy-server/src/routes/token_count.rs b/crates/llm-proxy-server/src/routes/token_count.rs index c3f17e7..039ed84 100644 --- a/crates/llm-proxy-server/src/routes/token_count.rs +++ b/crates/llm-proxy-server/src/routes/token_count.rs @@ -1,14 +1,8 @@ -//! `/providers/{provider}/v1/messages/count_tokens` handler. -//! -//! Accepts an Anthropic-format request, decodes it through the Anthropic client -//! adapter into a `CoreRequest`, and estimates the token count from the core -//! representation. - use std::sync::Arc; use axum::body::Body; use axum::extract::rejection::JsonRejection; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use axum::response::IntoResponse; use llm_proxy_core::MessageContent; @@ -18,65 +12,24 @@ use llm_proxy_protocol::core::{StopReason, Usage}; use llm_proxy_storage::{ProxyEvent, RequestReceived, ResponseCompleted}; use serde::Serialize; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; use super::error_response::{ClientProtocol, RouteError, route_error_response}; - -/// Response body for the token count endpoint. -/// -/// **Note:** The `input_tokens` value is computed locally with the model's own -/// BPE tokenizer for known OpenAI models (`gpt-4o`, `gpt-4`, `gpt-3.5`, etc.), -/// and falls back to a character-level heuristic (~4 chars/token) for any model -/// the proxy does not have a tokenizer for. It is therefore exact for known -/// OpenAI text but still an estimate for other providers and for non-text -/// content. It intentionally excludes tool definitions, non-text content blocks -/// (images, documents), and tool-result content. For non-OpenAI providers the -/// provider-reported usage (returned in the actual response) remains the -/// authoritative count for billing; this estimate is for pre-flight size gating. #[derive(Serialize)] #[serde(deny_unknown_fields)] pub(crate) struct TokenCountResponse { input_tokens: usize, } - -/// POST `/providers/{provider}/v1/messages/count_tokens` -/// -/// Accepts an Anthropic-format request, decodes it through the Anthropic client -/// adapter into a `CoreRequest`, and estimates the token count using the -/// configured tokenizer: the model's own BPE tokenizer for known OpenAI models, -/// falling back to the ~4 chars/token heuristic otherwise. -/// -/// This handler uses the core pipeline decode (`anthropic::decode_request`) -/// rather than direct field access on `MessageRequest`, so the requested model -/// id (`core.model.requested`) is available to select the right tokenizer. pub(crate) async fn count_tokens( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { - let start = std::time::Instant::now(); - // The wrapper only renders: it does NOT emit ResponseFailed. token_count - // has no core-pipeline dispatch, so each early-gate error return inside the - // inner function emits exactly one ResponseFailed itself, and the success - // path emits exactly one ResponseCompleted. A blanket emit here would - // double-count, so the wrapper stays emit-free - // (audit route-responsefailed-gaps regression). - match count_tokens_inner( - &state, - req_id, - &provider, - connect_info.as_ref(), - &headers, - body, - start, - ) - .await - { + match count_tokens_inner(&state, meta, &provider, &headers, body).await { Ok(response) => response, Err(error) => { tracing::warn!(error = %error, "request failed"); @@ -84,10 +37,6 @@ pub(crate) async fn count_tokens( } } } - -/// Render an axum [`JsonRejection`] as a human-readable message, preserving -/// the distinction between a syntactically invalid JSON body and a body that is -/// valid JSON but does not fit the target type (audit LOW-28). fn json_rejection_message(rejection: &JsonRejection) -> String { match rejection { JsonRejection::JsonSyntaxError(e) => format!("invalid JSON syntax: {e}"), @@ -98,112 +47,9 @@ fn json_rejection_message(rejection: &JsonRejection) -> String { other => format!("could not read request body: {other}"), } } - -async fn count_tokens_inner( - state: &AppState, - req_id: RequestId, - provider: &str, - connect_info: Option<&std::net::SocketAddr>, - headers: &HeaderMap, - body: axum::body::Bytes, - start: std::time::Instant, -) -> Result, RouteError> { - // Pre-flight: rate limit, dedup, request ID. token_count has no core-pipeline - // dispatch, so each EARLY-GATE error return below emits exactly one - // ResponseFailed itself, and the success path emits exactly one - // ResponseCompleted (audit route-responsefailed-gaps). - let event_bus = Arc::clone(&state.event_bus); - // The request id is consumed by prepare_request below; snapshot it once so - // every early-gate emit references the same id. - let request_id = req_id.0.clone(); - // Every early gate emits exactly one ResponseFailed with the same - // provider/request/start triple; bind it once instead of repeating the - // six-argument call at each gate. - let emit_failed = |e: &RouteError| { - core_pipeline::emit_response_failed( - &event_bus, - &request_id, - Some(provider), - None, - e, - start, - ); - }; - core_pipeline::validate_provider_name(provider).inspect_err(&emit_failed)?; - if state.providers().get(provider).is_none() { - let err = RouteError::UnknownProvider(provider.to_owned()); - emit_failed(&err); - return Err(err); - } - // Reject non-JSON Content-Type before parsing, so a wrong media type is not - // misreported as a JSON syntax error (audit LOW-30). - core_pipeline::validate_json_content_type(headers).inspect_err(&emit_failed)?; - let request_path = format!("/providers/{provider}/v1/messages/count_tokens"); - let ctx = core_pipeline::prepare_request( - state, - req_id.0, - headers, - connect_info, - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; - // Parse and validate the Anthropic MessageRequest. - // - // Go through `axum::Json::from_bytes` (rather than `serde_json::from_slice`) - // so the `JsonRejection` taxonomy is preserved: syntax errors and data - // (deserialization) errors are surfaced with distinct messages instead of - // being collapsed into a single opaque "invalid JSON" string. - let req: MessageRequest = match axum::Json::::from_bytes(&body) { - Ok(axum::Json(value)) => value, - Err(rejection) => { - let err = RouteError::InvalidRequest(json_rejection_message(&rejection)); - emit_failed(&err); - return Err(err); - } - }; - - req.validate() - .map_err(|e| RouteError::InvalidRequest(e.to_string())) - .inspect_err(&emit_failed)?; - - // Decode through the Anthropic client adapter to get a CoreRequest. - // This validates the request shape and normalises it. - let core = anthropic::decode_request(req) - .map_err(core_pipeline::protocol_error_to_route) - .inspect_err(&emit_failed)?; - - // Emit a RequestReceived event for the token-count request. token_count is a - // real proxy request (it decodes client input and consumes rate-limit/dedup - // budget) so it is logged like any other route. route_kind/client_protocol - // are marked "count_tokens" / "anthropic" so consumers can distinguish it - // from inference requests (audit eventlog-early-gates). - state - .event_bus - .emit(&ProxyEvent::RequestReceived(RequestReceived { - request_id: ctx.request_id.clone(), - timestamp: time::OffsetDateTime::now_utc(), - provider: provider.to_owned(), - route_kind: "count_tokens".to_owned(), - client_protocol: "anthropic".to_owned(), - model: core.model.clone(), - streaming: false, - body_hash: super::core_pipeline::body_hash_of_bytes(&body), - })); - - // Extract text content from core messages for token counting. - // - // NOTE: This estimate intentionally only counts text from `CoreContent::Text` - // blocks. The following are NOT counted: - // - Tool definitions (`core.tools`) -- their name, description, and - // input_schema text fields are excluded. - // - Non-text content blocks (images, documents, etc.). - // - Tool result content blocks. - // - System prompt non-text blocks. - // - // TODO(future): Extend counting to include tool definitions, tool_use blocks, - // and tool_result blocks for a more accurate estimate. - // Collect system text without intermediate Vec allocation. +fn extract_countable_text( + core: &llm_proxy_protocol::core::CoreRequest, +) -> (String, Vec) { let system_text: String = core .system .iter() @@ -228,10 +74,6 @@ async fn count_tokens_inner( _ => None, }) .collect(); - // Non-text content blocks (images, tool_use, thinking) produce - // empty text here. This is a known limitation of the heuristic - // counter. Log at debug level so operators can detect requests - // where the estimate may be significantly off. if text.is_empty() && !msg.content.is_empty() { tracing::debug!( role = ?msg.role, @@ -243,11 +85,6 @@ async fn count_tokens_inner( llm_proxy_protocol::core::CoreRole::User => "user", llm_proxy_protocol::core::CoreRole::Assistant => "assistant", llm_proxy_protocol::core::CoreRole::System => "system", - // The Anthropic Messages API does not have a 'tool' role. - // Tool results are carried in 'user' role messages with - // tool_result content blocks. If a CoreRole::Tool reaches - // here, treat it as 'user' since tool_result content is - // typically nested under the user role in Anthropic. llm_proxy_protocol::core::CoreRole::Tool => { tracing::debug!( "CoreRole::Tool mapped to 'user' for token counting; \ @@ -255,9 +92,6 @@ async fn count_tokens_inner( ); "user" } - // All known variants are handled above. This catch-all exists - // for forward compatibility. Unknown roles are treated as - // 'user' since that is the least-lossy mapping. _ => { tracing::warn!( role = ?msg.role, @@ -270,12 +104,68 @@ async fn count_tokens_inner( }) .collect(); - // BPE tokenization is CPU-bound (fancy-regex driven) and the request body - // is attacker-controlled up to MAX_BODY_BYTES (32 MiB). Running it on the - // tokio worker thread stalls every other future polled on that worker for - // the full duration of the encode, so offload it to the blocking pool. - // `Counter::clone` is a cheap refcount bump; `system_text` and `messages` - // are already owned here and are the last use on this path. + (system_text, messages) +} + +async fn count_tokens_inner( + state: &AppState, + meta: RequestMeta, + provider: &str, + headers: &HeaderMap, + body: axum::body::Bytes, +) -> Result, RouteError> { + let start = std::time::Instant::now(); + let event_bus = Arc::clone(&state.event_bus); + let request_id = meta.request_id.0.clone(); + let emit_failed = |e: &RouteError| { + core_pipeline::emit_response_failed( + &event_bus, + &request_id, + Some(provider), + None, + e, + start, + ); + }; + core_pipeline::validate_provider_name(provider).inspect_err(&emit_failed)?; + if state.providers().get(provider).is_none() { + let err = RouteError::UnknownProvider(provider.to_owned()); + emit_failed(&err); + return Err(err); + } + core_pipeline::validate_json_content_type(headers).inspect_err(&emit_failed)?; + let request_path = format!("/providers/{provider}/v1/messages/count_tokens"); + let ctx = core_pipeline::prepare_request(state, &meta, headers, &body, &request_path) + .inspect_err(&emit_failed)?; + let req: MessageRequest = match axum::Json::::from_bytes(&body) { + Ok(axum::Json(value)) => value, + Err(rejection) => { + let err = RouteError::InvalidRequest(json_rejection_message(&rejection)); + emit_failed(&err); + return Err(err); + } + }; + + req.validate() + .map_err(|e| RouteError::InvalidRequest(e.to_string())) + .inspect_err(&emit_failed)?; + let core = anthropic::decode_request(req) + .map_err(core_pipeline::protocol_error_to_route) + .inspect_err(&emit_failed)?; + state + .event_bus + .emit(&ProxyEvent::RequestReceived(RequestReceived { + request_id: ctx.request_id.clone(), + timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), + provider: provider.to_owned(), + route_kind: "count_tokens".to_owned(), + client_protocol: "anthropic".to_owned(), + model: core.model.clone(), + streaming: false, + body_hash: super::core_pipeline::body_hash_of_bytes(&body), + })); + let (system_text, messages) = extract_countable_text(&core); let counter = state.token_counter.clone(); let model = core.model.requested.clone(); let count = tokio::task::spawn_blocking(move || { @@ -283,12 +173,6 @@ async fn count_tokens_inner( }) .await .map_err(|e| RouteError::Internal(format!("token count task failed: {e}")))?; - - // Emit a ResponseCompleted event. token_count performs no upstream inference - // so there is no provider usage, cost, or stop reason; usage is a synthetic - // zero with provenance SyntheticZero and stop_reason is EndTurn so the event - // shape stays uniform with the inference routes. The computed token estimate - // travels in the HTTP response body, not the event (audit eventlog-early-gates). let latency = start.elapsed(); state .event_bus @@ -309,9 +193,6 @@ async fn count_tokens_inner( }); let mut response = response.into_response(); - - // Defense-in-depth: explicitly set Content-Type even though - // axum::Json's IntoResponse already sets it internally. response.headers_mut().insert( axum::http::header::CONTENT_TYPE, axum::http::HeaderValue::from_static("application/json"), @@ -327,10 +208,6 @@ async fn count_tokens_inner( Ok(response) } -// =========================================================================== -// Tests -// =========================================================================== - #[cfg(test)] mod tests { #[test] diff --git a/crates/llm-proxy-server/src/shutdown.rs b/crates/llm-proxy-server/src/shutdown.rs index b9baf84..6e00cb7 100644 --- a/crates/llm-proxy-server/src/shutdown.rs +++ b/crates/llm-proxy-server/src/shutdown.rs @@ -2,17 +2,8 @@ use tokio::signal; /// Wait for SIGINT (Ctrl-C) or SIGTERM (Unix only). /// -/// Handler-install failures degrade **fail-safe** rather than hang. If a signal -/// handler cannot be installed, that branch resolves immediately (after logging) -/// instead of parking on [`std::future::pending`]. Otherwise `tokio::select!` -/// could be left waiting on two futures that will both never fire, which would -/// hang a headless deployment (Kubernetes/LB pod recycling via SIGTERM) whose -/// SIGTERM handler failed to install and where no human can send Ctrl-C — the -/// graceful drain would then never begin (audit GAP-LOW-15). -/// -/// Resolving a failed branch is preferred over hanging: a process that cannot -/// install its primary shutdown signal is in a degraded state, and a supervisor -/// (systemd/k8s) will restart it cleanly. +/// A handler install failure resolves (fail-safe) rather than parking on +/// [`std::future::pending`], so graceful shutdown proceeds instead of hanging. pub async fn shutdown_signal() { let ctrl_c = async { match signal::ctrl_c().await { @@ -36,22 +27,14 @@ pub async fn shutdown_signal() { "failed to install SIGTERM handler; beginning graceful shutdown" ), } - // An install failure falls through here (resolving the branch) instead of - // parking on pending(), so tokio::select! below is never left waiting on - // two futures that will both never fire (audit GAP-LOW-15). }; #[cfg(not(unix))] let terminate = async { - // No SIGTERM on non-Unix; Ctrl-C is the sole signal and always resolves - // (Ok on signal, Err on install failure), so this branch is never the - // sole deciding future and the pending park is harmless. std::future::pending::<()>().await; }; tokio::select! { - // Both branches are biased to resolve: a real signal OR an install - // failure completes the select and lets graceful shutdown proceed. _ = ctrl_c => {} _ = terminate => {} } diff --git a/crates/llm-proxy-server/src/state.rs b/crates/llm-proxy-server/src/state.rs index d439402..d2a91ad 100644 --- a/crates/llm-proxy-server/src/state.rs +++ b/crates/llm-proxy-server/src/state.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -6,7 +7,7 @@ use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; use llm_proxy_storage::{EventBus, NoopBus}; use crate::ModelCatalogService; -use crate::middleware::{RateLimiter, RequestDeduplicator, RequestIdGenerator}; +use crate::middleware::{ClientAuth, RateLimiter, RequestDeduplicator, RequestIdGenerator}; /// Static information about this binary. #[derive(Debug, Clone)] @@ -21,96 +22,37 @@ pub struct BuildInfo { pub git_sha: &'static str, } -// --------------------------------------------------------------------------- -// AppState -// --------------------------------------------------------------------------- - /// Application state shared with every handler. /// -/// Cheap to clone: nearly all fields are `Arc`-wrapped so cloning only bumps -/// reference counts. `token_counter` is a non-ZST `Counter` stored inline, but -/// its lazily-loaded BPE cache is held in an inner `Arc`, so cloning `Counter` -/// is a refcount bump and every `AppState` clone shares the same loaded BPE -/// encodings. Required by axum's `State` extractor. -/// -/// # Security note -/// -/// Manual `Debug` impl is provided to ensure `api_key` in nested config types -/// is never leaked through debug formatting (e.g. in error logs). -/// -/// The transitive redaction chain for Debug is: -/// - `AppState` -> `app_config` field: `AppConfig` has no secrets (safe). -/// - `AppState` -> `providers` field: `ProviderRegistry` -> `ProviderConfig`, -/// whose `api_key` is a `secrecy::SecretString`; `ProviderConfig` derives -/// `Debug` and the key is redacted by `SecretString`'s own impl. Safe. -/// - `AppState` -> `provider_adapters` field: `ProviderAdapterRegistry` -> -/// `ProviderAdapter` types whose `api_key` is `SecretString`; the adapter -/// config/target types keep a *manual* `Debug` that additionally redacts the -/// endpoint query string and auth-header *values* (not just the key). Safe. -/// - `AppState` -> remaining fields: `BuildInfo`, `proxy_client`, `Counter`, -/// `Metrics`, `RateLimiter`, `RequestDeduplicator`, `RequestIdGenerator`, -/// `model_catalogs`, `event_bus` -- none hold secrets. Safe. -/// -/// **Maintenance note:** Secret-backed fields use `secrecy::SecretString`, which -/// redacts in `Debug`/`Display` regardless of whether the containing type -/// derives or hand-writes `Debug`. The remaining risk is non-`SecretString` -/// sensitive data (e.g. endpoint query strings, raw header values): those still -/// rely on the manual `Debug` impls on the adapter config/target types. If such -/// a field is added to a type that derives `Debug`, audit it. The test -/// `app_state_debug_does_not_leak_api_key` provides regression coverage. +/// Cheap to clone (`Arc`-shared fields); the manual `Debug` impl never emits +/// secrets (regression-tested by `app_state_debug_does_not_leak_api_key`). #[derive(Clone)] #[non_exhaustive] pub struct AppState { - /// TOML application config. pub(crate) app_config: Arc, - /// Provider registry loaded from TOML provider files. pub(crate) providers: Arc, - /// Provider adapter registry (always present). pub(crate) provider_adapters: Arc, - /// Protocol-neutral HTTP transport client. pub(crate) proxy_client: Arc, - /// Build identifier reported by `/version`. + pub(crate) egress_clients: Arc>>, pub(crate) build: Arc, - /// Token counter for estimating token usage. - /// - /// `Counter` dispatches to a real BPE tokenizer (tiktoken) for known OpenAI - /// models and a character heuristic otherwise. It holds a lazily-populated, - /// `Arc`-shared cache of loaded BPE encodings behind a `Mutex`, so cloning - /// the counter via `AppState: Clone` is a cheap refcount bump that shares - /// the cache across handlers. pub(crate) token_counter: Counter, - /// Runtime metrics collector. pub(crate) metrics: Arc, - /// Per-IP rate limiter. pub(crate) rate_limiter: Arc, - /// Request deduplicator for idempotent request handling. pub(crate) request_dedup: Arc, - /// Request ID generator. pub(crate) request_id_gen: Arc, - /// Mutable provider model catalog cache and discovery coordinator. pub(crate) model_catalogs: Arc, - /// Structured request/response event sink. Defaults to [`NoopBus`]; a real - /// persisted backend lands in a later phase. pub(crate) event_bus: Arc, + pub(crate) client_auth: Arc, } impl std::fmt::Debug for AppState { - // SECURITY: This manual Debug impl delegates through Arc to each field's - // Debug. The redaction chain documented on the struct must be maintained. - // The test `app_state_debug_does_not_leak_api_key` provides regression - // coverage. See struct-level doc comment for the full chain. - // - // NOTE: A `redacted_debug!()` macro could reduce boilerplate for types - // with sensitive fields, but would add a procedural-macro dependency. - // The manual impl is preferred for now because: (a) AppState has few - // fields, (b) the redaction chain is explicitly documented and tested, - // and (c) a macro would obscure the security-critical delegation paths. fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AppState") .field("app_config", &self.app_config) .field("providers", &self.providers) .field("provider_adapters", &self.provider_adapters) .field("proxy_client", &self.proxy_client) + .field("egress_clients", &self.egress_clients) .field("build", &self.build) .field("token_counter", &self.token_counter) .field("metrics", &self.metrics) @@ -119,6 +61,7 @@ impl std::fmt::Debug for AppState { .field("request_id_gen", &self.request_id_gen) .field("model_catalogs", &self.model_catalogs) .field("event_bus", &self.event_bus) + .field("client_auth", &self.client_auth) .finish() } } @@ -154,16 +97,24 @@ impl AppState { catalog_dir: Option, ) -> Self { let rate_limit_rpm = app_config.server.rate_limit_rpm; - // A dedup window whose millis do not fit u64 (Duration::MAX) cannot - // meaningfully be represented; fall back to 0, which disables dedup, - // matching the documented "0s disables deduplication" contract. let dedup_window_ms = u64::try_from(app_config.server.dedup_window.as_millis()).unwrap_or(0); + let client_auth = match ClientAuth::from_config(&app_config.auth) { + Ok(auth) => Arc::new(auth), + Err(error) => { + tracing::error!( + error = %error, + "invalid [auth] config; failing closed (all proxy requests will be rejected)" + ); + Arc::new(ClientAuth::fail_closed()) + } + }; Self { app_config: Arc::new(app_config), providers: Arc::new(providers), provider_adapters: Arc::new(provider_adapters), proxy_client: Arc::new(proxy_client), + egress_clients: Arc::new(HashMap::new()), build: Arc::new(build), token_counter: Counter::new(), metrics: Arc::new(Metrics::new()), @@ -172,12 +123,31 @@ impl AppState { request_id_gen: Arc::new(RequestIdGenerator::new()), model_catalogs: Arc::new(ModelCatalogService::new(catalog_dir)), event_bus: Arc::new(NoopBus) as Arc, + client_auth, } } - /// Replace the event bus (builder-style). Used by tests to inject a - /// [`RecordingBus`](llm_proxy_storage::RecordingBus); production wires the - /// real bus in `cmd_serve` (still `NoopBus` for now). + /// Replace the client-auth state (builder-style). + #[must_use] + pub fn with_client_auth(mut self, client_auth: Arc) -> Self { + self.client_auth = client_auth; + self + } + + /// Replace the per-provider egress clients (builder-style). + /// + /// An empty map (the default) keeps every provider on the shared + /// `proxy_client`. + #[must_use] + pub fn with_egress_clients( + mut self, + egress_clients: HashMap>, + ) -> Self { + self.egress_clients = Arc::new(egress_clients); + self + } + + /// Replace the event bus (builder-style). #[must_use] pub fn with_event_bus(mut self, event_bus: Arc) -> Self { self.event_bus = event_bus; @@ -207,11 +177,9 @@ impl AppState { self.app_config.server.trust_forwarded_headers } - /// Configured CORS allow-origin list, if any (audit LOW-12). - /// - /// `None` means cross-origin requests are not permitted: no `CorsLayer` is - /// installed and browsers enforce a same-origin default. `Some(origins)` - /// are the exact origins a restrictive `CorsLayer` will echo. + /// Configured CORS allow-origin list, if any. `None` permits no + /// cross-origin requests (no `CorsLayer` installed); `Some(origins)` are + /// the exact origins the `CorsLayer` echoes. pub fn allowed_origins(&self) -> Option<&[String]> { self.app_config.server.allowed_origins.as_deref() } @@ -226,11 +194,40 @@ impl AppState { &self.providers } + /// The transport client for an upstream request to `provider`: its wired + /// `[provider.egress]` client, or the shared default client otherwise. + /// An egress-configured-but-unwired provider logs a loud warning per request. + pub fn upstream_client_for(&self, provider: &str) -> &ProxyClient { + if let Some(client) = self.egress_clients.get(provider) { + return client; + } + if self + .providers + .get(provider) + .is_some_and(|config| config.egress.is_some()) + { + tracing::warn!( + provider, + "provider has [provider.egress] configured but no egress client was wired; \ + its upstream traffic is leaving directly, not through the configured proxy" + ); + } + &self.proxy_client + } + /// Access the provider model catalog service. pub fn model_catalogs(&self) -> &ModelCatalogService { &self.model_catalogs } + /// Access the request/response event sink. + /// + /// Holds the bus so the serve path can [`EventBus::close`] it (flush the + /// usage ledger) after the HTTP server has stopped. + pub fn event_bus(&self) -> &Arc { + &self.event_bus + } + /// Return the bind address the server should listen on. pub fn bind_address(&self) -> std::net::SocketAddr { self.app_config.server.bind @@ -245,10 +242,8 @@ impl AppState { #[cfg(test)] mod tests { use super::*; - use std::collections::HashMap; use std::net::SocketAddr; - /// Helper: build a minimal TOML AppConfig for tests. fn make_app_config() -> AppConfig { use llm_proxy_core::ServerConfig; AppConfig { @@ -264,11 +259,14 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, } } - /// Helper: build a minimal ProviderRegistry for tests. fn make_provider_registry() -> ProviderRegistry { ProviderRegistry::from_providers(vec![]).expect("empty registry") } @@ -282,8 +280,6 @@ mod tests { } } - // -- AppState construction ------------------------------------------------ - #[test] fn app_state_construction() { let state = AppState::new( @@ -298,7 +294,59 @@ mod tests { assert_eq!(state.shutdown_timeout(), Duration::from_secs(30)); } - // -- Helper methods ------------------------------------------------------- + #[test] + fn upstream_client_for_returns_the_shared_client_without_egress() { + let state = AppState::new( + make_app_config(), + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ); + for provider in ["opencode-go", "opencode-zen", "anything-else"] { + assert!( + std::ptr::eq( + state.upstream_client_for(provider), + state.proxy_client.as_ref() + ), + "provider {provider} without egress must use the shared client" + ); + } + } + + #[test] + fn upstream_client_for_returns_the_wired_client_only_for_that_provider() { + let egress_client = Arc::new( + ProxyClient::try_with_egress_proxy("socks5h://127.0.0.1:1080") + .expect("egress proxy client builds"), + ); + let state = AppState::new( + make_app_config(), + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ) + .with_egress_clients(HashMap::from([( + "opencode-go".to_owned(), + Arc::clone(&egress_client), + )])); + + assert!( + std::ptr::eq( + state.upstream_client_for("opencode-go"), + egress_client.as_ref() + ), + "the provider with egress must use its wired client" + ); + assert!( + std::ptr::eq( + state.upstream_client_for("opencode-zen"), + state.proxy_client.as_ref() + ), + "providers without egress must keep the shared client" + ); + } #[test] fn helpers_return_app_config_values() { @@ -313,8 +361,6 @@ mod tests { assert_eq!(state.request_timeout(), Duration::from_secs(300)); } - // -- Debug does not leak API keys ----------------------------------------- - #[test] fn app_state_debug_does_not_leak_api_key() { use llm_proxy_core::{AuthStyle, ProviderConfig, ProviderRoutesConfig}; @@ -323,6 +369,7 @@ mod tests { api_key: secrecy::SecretString::from("sk-secret-key-99999"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -351,9 +398,6 @@ mod tests { #[test] fn app_state_debug_covers_all_fields() { - // Regression test: verify that the Debug impl actually outputs all - // field names. If a field is added to AppState but missing from the - // manual Debug impl, this test catches it. let state = AppState::new( make_app_config(), make_provider_registry(), @@ -362,7 +406,6 @@ mod tests { make_build_info(), ); let debug_output = format!("{:?}", state); - // All fields that should appear in the Debug output. for field in [ "app_config", "providers", @@ -376,6 +419,7 @@ mod tests { "request_id_gen", "model_catalogs", "event_bus", + "client_auth", ] { assert!( debug_output.contains(field), @@ -384,8 +428,6 @@ mod tests { } } - // -- Provider adapters are populated -------------------------------------- - #[test] fn provider_adapters_are_populated() { let state = AppState::new( @@ -395,15 +437,12 @@ mod tests { ProxyClient::new(), make_build_info(), ); - // Use >= 4 rather than == 4 so adding new builtin adapters does not - // break this test (the intent is to verify adapters are populated). let protocol_names = state.provider_adapters.protocol_names(); assert!( protocol_names.len() >= 4, "expected at least 4 builtin adapters, got {len}", len = protocol_names.len() ); - // Verify specific known protocol names are present for stronger coverage. assert!( protocol_names.contains(&"openai_chat_completions"), "missing openai_chat_completions adapter" @@ -414,8 +453,6 @@ mod tests { ); } - // -- TOML provider protocol validation against builtin ------------------- - #[test] fn toml_provider_protocol_validation_against_builtin() { use llm_proxy_core::{ @@ -426,6 +463,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -458,8 +496,6 @@ mod tests { assert!(result.is_ok(), "all protocols are builtin, got: {result:?}"); } - // -- TOML unknown protocol fails validation against builtin --------------- - #[test] fn toml_unknown_protocol_fails_validation_against_builtin() { use llm_proxy_core::{ @@ -470,6 +506,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -499,12 +536,8 @@ mod tests { ); } - // -- Config extension dispatch logic ------------------------------------- - #[test] fn config_extension_dispatch_classifies_correctly() { - /// Mimics the match arm logic from cmd_serve to verify extension - /// classification without requiring a full integration test. fn classify(path: &std::path::Path) -> &'static str { match path.extension().and_then(|e| e.to_str()).unwrap_or("") { "toml" => "toml", @@ -521,16 +554,12 @@ mod tests { assert_eq!(classify(std::path::Path::new("config.YAML")), "unsupported"); } - // -- Send + Sync static assertion for AppState --------------------------- - #[test] fn app_state_is_send_sync() { fn check() {} check::(); } - // -- bind_address() tests ------------------------------------------------ - #[test] fn bind_address_returns_configured_bind() { let mut cfg = make_app_config(); @@ -548,8 +577,6 @@ mod tests { ); } - // -- Clone shares Arc references ----------------------------------------- - #[test] fn clone_shares_arc_references() { let state = AppState::new( @@ -568,28 +595,20 @@ mod tests { Arc::ptr_eq(&state.build, &cloned.build), "cloned AppState should share the same Arc" ); - // NOTE: `token_counter` is an inline `Counter` (not `Arc`), so - // there is no outer `Arc` to compare with `Arc::ptr_eq`. It is - // intentionally excluded from the Arc-sharing assertions above. The - // loaded BPE cache is still shared: `Counter` holds its cache in an - // inner `Arc>>`, so this `cloned` state reuses the - // same loaded encodings as the original. } - // -- Multi-provider TOML validation -------------------------------------- - #[test] fn multi_provider_toml_mixed_protocols() { use llm_proxy_core::{ AuthStyle, ProviderAdapterConfig, ProviderConfig, ProviderRoutesConfig, }; - // Provider with valid protocols let good = ProviderConfig { name: "good".to_owned(), api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -609,12 +628,12 @@ mod tests { pricing: Default::default(), }; - // Provider with invalid protocol let bad = ProviderConfig { name: "bad".to_owned(), api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -643,16 +662,8 @@ mod tests { ); } - // -- Boundary values for helpers ----------------------------------------- - #[test] fn zero_duration_timeout_is_returned() { - // AppState is a dependency-injection container: request_timeout() is a - // pure passthrough accessor and performs no validation. Zero/sub-second - // timeouts are rejected earlier, at config load, in `load_app_config` - // (see `llm-proxy-core` `validate_*` — HIGH-2 / LOW-10). This test - // documents the accessor's faithfulness to whatever ServerConfig it is - // handed, not an endorsement of `request_timeout = "0s"` as valid config. let mut cfg = make_app_config(); cfg.server.request_timeout = Duration::ZERO; let state = AppState::new( @@ -679,8 +690,6 @@ mod tests { assert_eq!(state.server_name(), ""); } - // -- Unicode server_name --------------------------------------------------- - #[test] fn unicode_server_name_is_returned_verbatim() { let mut cfg = make_app_config(); @@ -709,4 +718,65 @@ mod tests { ); assert_eq!(state.server_name(), long_name); } + + #[test] + fn invalid_auth_config_fails_closed() { + let mut cfg = make_app_config(); + cfg.auth.mode = llm_proxy_core::AuthMode::ClientKeys; + cfg.auth.clients.insert( + "laptop".to_owned(), + secrecy::SecretString::from("not-a-phc-string"), + ); + let state = AppState::new( + cfg, + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ); + assert!( + state.client_auth.is_enforced(), + "invalid auth config must fail closed, not open" + ); + assert_eq!(state.client_auth.client_count(), 0); + } + + #[test] + fn app_state_debug_does_not_leak_client_verifiers() { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + b"leakable-secret-tail", + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + + let mut cfg = make_app_config(); + cfg.auth.mode = llm_proxy_core::AuthMode::ClientKeys; + cfg.auth + .clients + .insert("laptop".to_owned(), secrecy::SecretString::from(phc)); + let client_auth = + crate::middleware::ClientAuth::from_config(&cfg.auth).expect("valid verifier"); + let state = AppState::new( + cfg, + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ) + .with_client_auth(Arc::new(client_auth)); + let debug = format!("{state:?}"); + assert!( + !debug.contains("leakable-secret-tail"), + "client secret/verifier must not appear in Debug, got: {debug}" + ); + assert!(debug.contains("laptop"), "key id may appear: {debug}"); + } } diff --git a/crates/llm-proxy-server/tests/admin_usage.rs b/crates/llm-proxy-server/tests/admin_usage.rs new file mode 100644 index 0000000..2146549 --- /dev/null +++ b/crates/llm-proxy-server/tests/admin_usage.rs @@ -0,0 +1,210 @@ +//! Router-level tests for the admin `/usage` surface on both mounts: the +//! dedicated admin listener router and the opt-in main-listener mount. + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use axum::body::Body; + use axum::http::Request; + use axum::http::StatusCode; + use llm_proxy_core::{AppConfig, ServerConfig, StorageConfig}; + use llm_proxy_server::{AppState, BuildInfo, admin_router, build_router}; + use llm_proxy_storage::{EventBus, SqliteEventBus}; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "integration-admin-token"; + + async fn seeded_ledger_path(dir: &tempfile::TempDir) -> std::path::PathBuf { + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("open ledger"); + bus.emit(&llm_proxy_storage::ProxyEvent::RequestReceived( + llm_proxy_storage::RequestReceived { + request_id: "srv-1".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "laptop".to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "digest".to_owned(), + }, + )); + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + path + } + + fn make_config( + db_path: std::path::PathBuf, + admin: Option, + on_main_listener: bool, + ) -> AppConfig { + AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: Duration::from_secs(300), + shutdown_timeout: Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "admin-usage-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: on_main_listener, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients: std::collections::HashMap::new(), + admin, + }, + storage: Some(StorageConfig { + path: db_path, + max_age: None, + max_rows: None, + prune_interval: Duration::from_secs(3600), + }), + } + } + + fn make_state(config: AppConfig) -> AppState { + AppState::new( + config, + llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"), + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::new(), + BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ) + } + + async fn get(app: axum::Router, uri: &str, token: Option<&str>) -> (StatusCode, String) { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + let response = app + .oneshot(builder.body(Body::empty()).expect("request builds")) + .await + .expect("oneshot"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + (status, String::from_utf8_lossy(&bytes).into_owned()) + } + + #[tokio::test] + async fn main_listener_has_no_usage_surface_by_default() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + false, + )); + let (status, _) = get(build_router(state), "/usage/summary", Some(ADMIN_TOKEN)).await; + assert_eq!( + status, + StatusCode::NOT_FOUND, + "without the opt-in, the main listener must not serve /usage even to the admin token" + ); + } + + #[tokio::test] + async fn main_listener_optin_serves_usage_behind_the_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + true, + )); + + let (status, body) = get(build_router(state.clone()), "/usage/summary", None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!( + !body.contains(ADMIN_TOKEN), + "401 must not echo the configured token" + ); + + let (status, body) = get( + build_router(state.clone()), + "/usage/summary", + Some("wrong-token"), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!( + !body.contains("wrong-token") && !body.contains(ADMIN_TOKEN), + "401 must not echo any token: {body}" + ); + + let (status, body) = get(build_router(state), "/usage/summary", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + assert!( + body.contains("prov-a"), + "the seeded provider group is served: {body}" + ); + } + + #[tokio::test] + async fn dedicated_admin_router_serves_usage_behind_the_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + false, + )); + + let admin = admin_router(&state).expect("admin router with [auth.admin]"); + let (status, _) = get(admin.clone(), "/usage/keys", None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + + let (status, body) = get(admin, "/usage/keys", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + assert!( + body.contains("laptop"), + "the seeded caller is listed: {body}" + ); + } + + #[tokio::test] + async fn admin_router_is_none_without_admin_config() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config(path, None, false)); + assert!( + admin_router(&state).is_none(), + "no [auth.admin] means no admin surface anywhere" + ); + } +} diff --git a/crates/llm-proxy-server/tests/chat_completions.rs b/crates/llm-proxy-server/tests/chat_completions.rs index 4eb7a82..4d7d645 100644 --- a/crates/llm-proxy-server/tests/chat_completions.rs +++ b/crates/llm-proxy-server/tests/chat_completions.rs @@ -1,8 +1,6 @@ -//! Integration tests for Phase 9: /v1/chat/completions endpoint. -//! -//! These tests exercise the OpenAI Chat Completions route through the axum -//! router with mock upstream servers, verifying that the full pipeline -//! (OpenAI decode -> core pipeline -> OpenAI encode) works end-to-end. +//! Integration tests for the `/v1/chat/completions` route: mock upstreams +//! driven through the router to verify the OpenAI decode -> core pipeline +//! -> OpenAI encode path end-to-end. #[cfg(test)] mod tests { @@ -27,15 +25,6 @@ mod tests { use serde_json::{Value, json}; use tower::ServiceExt; - // --------------------------------------------------------------------------- - // Mock server readiness helper - // --------------------------------------------------------------------------- - - /// Wait for a mock server to be ready by polling its TCP port. - /// - /// Replaces `tokio::time::sleep(Duration::from_millis(50))` with a - /// deterministic readiness check that retries until the server accepts - /// a connection or the timeout elapses. async fn wait_for_ready(addr: std::net::SocketAddr) { let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { @@ -49,19 +38,8 @@ mod tests { } } - /// Spawn a mock axum server on `listener`, returning immediately while the - /// server runs in the background. - /// - /// Unlike a bare `tokio::spawn(async move { axum::serve(..).await })` whose - /// `JoinHandle` is dropped (silently swallowing any panic in the mock server - /// task — audit `testing-audit:mock-server-joinhandle-swallow`), this wrapper - /// catches a panic in the serve future, prints it to stderr, and aborts the - /// process. A mock-handler bug therefore surfaces loudly as a test failure - /// instead of the connection just closing mid-stream with no diagnostic. fn spawn_mock_serve(listener: tokio::net::TcpListener, app: Router) { tokio::spawn(async move { - // std::panic::catch_unwind requires UnwindSafe; the router/listener - // capture is fine for a mock that owns them exclusively. let result = std::panic::AssertUnwindSafe(axum::serve(listener, app).into_future()) .catch_unwind() .await; @@ -69,8 +47,6 @@ mod tests { Ok(Ok(())) => {} Ok(Err(e)) => eprintln!("mock axum::serve error: {e}"), Err(panic) => { - // Surface the panic payload, then abort so the test binary fails - // rather than continuing with a dead mock. let msg = panic .downcast_ref::() .map(|s| s.as_str()) @@ -83,11 +59,6 @@ mod tests { }); } - // --------------------------------------------------------------------------- - // Mock upstream server helpers - // --------------------------------------------------------------------------- - - /// A canned OpenAI Chat Completions API success response. fn openai_chat_success_response() -> Vec { serde_json::to_vec(&json!({ "id": "chatcmpl-test123", @@ -111,7 +82,6 @@ mod tests { .unwrap() } - /// A canned OpenAI Chat Completions API response with tool calls. fn openai_chat_tool_call_response() -> Vec { serde_json::to_vec(&json!({ "id": "chatcmpl-tool123", @@ -143,8 +113,6 @@ mod tests { .unwrap() } - /// Build AppState in TOML mode with a single provider of the given protocol - /// routing to the given mock endpoint. fn state_with_provider( mock_endpoint: &str, protocol: &str, @@ -160,7 +128,6 @@ mod tests { ) } - /// Like [`state_with_provider`] but with a configurable `request_timeout`. fn state_with_provider_timeout( mock_endpoint: &str, protocol: &str, @@ -168,8 +135,6 @@ mod tests { _model_name: &str, request_timeout: Duration, ) -> AppState { - // Build routes based on protocol so the provider-based routing can resolve. - // Anthropic providers also support chat_completions for cross-protocol tests. let routes = if protocol == "anthropic_messages" { llm_proxy_core::ProviderRoutesConfig { messages: Some(adapter_name.to_owned()), @@ -187,6 +152,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -221,7 +187,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -238,12 +208,10 @@ mod tests { ) } - /// Convenience: AppState with an OpenAI Chat provider. fn state_with_openai_chat_provider(mock_endpoint: &str) -> AppState { state_with_provider(mock_endpoint, "openai_chat_completions", "chat", "gpt-4o") } - /// Convenience: AppState with an Anthropic provider (for cross-protocol tests). fn state_with_anthropic_provider(mock_endpoint: &str) -> AppState { state_with_provider( mock_endpoint, @@ -253,7 +221,6 @@ mod tests { ) } - /// Spawn a local mock axum server returning a canned response body. async fn spawn_mock_server(response_body: Vec, content_type: &str) -> String { let ct = content_type.to_owned(); let app = Router::new().route( @@ -276,25 +243,18 @@ mod tests { ) } - /// Spawn a local mock axum server returning OpenAI Chat non-streaming response. async fn spawn_mock_openai_chat_non_stream() -> String { spawn_mock_server(openai_chat_success_response(), "application/json").await } - /// Spawn a local mock axum server returning OpenAI Chat tool call response. async fn spawn_mock_openai_chat_tool_call() -> String { spawn_mock_server(openai_chat_tool_call_response(), "application/json").await } - /// Spawn a mock server that accepts the connection but never responds, holding - /// the upstream call open so the proxy's per-route `TimeoutLayer` fires. Used to - /// exercise the 408 normalisation path (audit MEDIUM-7). async fn spawn_hanging_mock() -> String { let app = Router::new().route( "/{*path}", post(|| async move { - // Sleep well past any short test timeout so the caller's - // TimeoutLayer always fires first. tokio::time::sleep(Duration::from_secs(30)).await; StatusCode::OK }), @@ -306,7 +266,6 @@ mod tests { format!("http://{addr}/providers/mock-provider/v1/chat/completions") } - /// Spawn a mock server returning OpenAI Chat streaming SSE events. async fn spawn_mock_openai_chat_stream() -> String { let app = Router::new().route( "/{*path}", @@ -340,7 +299,6 @@ mod tests { ) } - /// Spawn a mock server that returns HTTP 500. async fn spawn_mock_500() -> String { let app = Router::new().route( "/{*path}", @@ -364,7 +322,6 @@ mod tests { ) } - /// Spawn a mock server that records the received request body. async fn spawn_mock_with_body_capture( response_body: Vec, ) -> (String, std::sync::Arc>>>) { @@ -400,7 +357,6 @@ mod tests { ) } - /// Build an AppState with empty routing table (for error tests). fn empty_state() -> AppState { let app_config = AppConfig { server: ServerConfig { @@ -415,7 +371,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( @@ -432,24 +392,13 @@ mod tests { ) } - /// A state that registers `mock-provider` (with an unreachable adapter) for - /// validation-path tests. - /// - /// Tests that assert a malformed body or a missing-field request yields an - /// OpenAI-shaped 400 (`empty_body_*`, `wrong_field_types_*`, - /// `missing_model_*`, `missing_messages_*`, `invalid_utf8_body_*`) must reach - /// the JSON-parsing/decode path. Since LOW-29 reordered the handler to run the - /// provider-existence gate *before* JSON parsing, these requests now need a - /// registered provider to get past that gate — `empty_state()` (no providers) - /// would instead surface `UnknownProvider` (404). The adapter endpoint below is - /// never contacted: every one of these requests fails at validation, well - /// before any upstream dispatch. fn state_with_mock_provider_no_upstream() -> AppState { let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -485,7 +434,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, @@ -501,7 +454,6 @@ mod tests { ) } - /// Build a request to /v1/chat/completions. fn chat_request(body: &str) -> Request { Request::builder() .method("POST") @@ -511,7 +463,6 @@ mod tests { .unwrap() } - /// Build a standard OpenAI Chat request body. fn make_chat_body(model: &str, stream: bool) -> String { json!({ "model": model, @@ -521,20 +472,12 @@ mod tests { .to_string() } - // =========================================================================== - // Route mounting - // =========================================================================== - - /// route is mounted, no longer 404 #[tokio::test] async fn route_is_mounted_no_longer_404() { let mock_url = spawn_mock_openai_chat_non_stream().await; let app = build_router(state_with_openai_chat_provider(&mock_url)); let body = make_chat_body("gpt-4o", false); let resp = app.oneshot(chat_request(&body)).await.unwrap(); - // The route is now mounted -- should not return 404. - // With a valid provider, the request should succeed (200) or fail with - // a non-404 error. assert_ne!( resp.status(), StatusCode::NOT_FOUND, @@ -542,11 +485,6 @@ mod tests { ); } - // =========================================================================== - // Non-streaming tests - // =========================================================================== - - /// non-streaming text request returns OpenAI-shaped response #[tokio::test] async fn non_streaming_text_request_returns_openai_shaped_response() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -564,7 +502,6 @@ mod tests { ) .unwrap(); - // Assert the OpenAI response shape. assert_eq!( resp_body["object"], "chat.completion", "object must be chat.completion" @@ -573,7 +510,6 @@ mod tests { assert!(resp_body["model"].is_string(), "model must be a string"); assert!(resp_body["created"].is_number(), "created must be a number"); - // Assert choices. let choices = resp_body["choices"] .as_array() .expect("choices must be array"); @@ -581,7 +517,6 @@ mod tests { assert_eq!(choices[0]["index"], 0); assert_eq!(choices[0]["finish_reason"], "stop"); - // Assert message content. let message = &choices[0]["message"]; assert_eq!(message["role"], "assistant"); assert!( @@ -592,15 +527,12 @@ mod tests { "response should contain translated text" ); - // Assert usage. let usage = &resp_body["usage"]; assert_eq!(usage["prompt_tokens"], 10); assert_eq!(usage["completion_tokens"], 5); assert_eq!(usage["total_tokens"], 15); } - /// non-streaming response asserts object='chat.completion', id, model, - /// choices[].finish_reason, text content, and usage #[tokio::test] async fn non_streaming_response_shape_detailed() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -631,7 +563,6 @@ mod tests { assert!(resp_body["usage"]["total_tokens"].is_number()); } - /// tool-call core response returns tool_calls #[tokio::test] async fn tool_call_response_returns_tool_calls() { let mock_url = spawn_mock_openai_chat_tool_call().await; @@ -669,8 +600,6 @@ mod tests { assert_eq!(tool_calls[0]["function"]["name"], "get_weather"); } - /// route preserves temperature, top-p, max tokens, tools, tool choice, metadata, - /// stream, stream options/provider hints, reasoning/thinking, and cache markers #[tokio::test] async fn route_preserves_fields_through_core() { let (mock_url, captured_body) = @@ -709,7 +638,6 @@ mod tests { let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Check the upstream received the correct fields. tokio::time::sleep(Duration::from_millis(100)).await; let guard = captured_body.lock().await; let upstream_body = guard @@ -718,43 +646,24 @@ mod tests { let upstream_json: Value = serde_json::from_slice(upstream_body).expect("upstream body should be valid JSON"); - // The OpenAI adapter should have preserved these fields. assert_eq!(upstream_json["model"], "gpt-4o"); - // max_tokens serializes as max_completion_tokens per OpenAI API convention assert_eq!(upstream_json["max_completion_tokens"], 256); assert_eq!(upstream_json["temperature"], 0.7); assert_eq!(upstream_json["top_p"], 0.9); - // Tools and tool_choice should be preserved. assert!(upstream_json["tools"].is_array(), "tools must be preserved"); assert_eq!(upstream_json["tools"][0]["function"]["name"], "get_weather"); assert_eq!(upstream_json["tool_choice"]["type"], "auto"); - // reasoning_effort should be preserved. assert_eq!(upstream_json["reasoning_effort"], "high"); - // user (metadata) should be preserved. assert_eq!(upstream_json["user"], "test-user-123"); - // thinking should be preserved. assert_eq!(upstream_json["thinking"]["type"], "enabled"); assert_eq!(upstream_json["thinking"]["budget_tokens"], 10000); - // Note: stream_options is only forwarded to upstream when stream=true, - // because the provider adapter only reads provider_hints for streaming - // requests. This is expected behavior -- stream_options is a streaming- - // specific hint. The cache_control on the system message is decoded into - // the core system array but the OpenAI provider adapter does not re-emit - // it in the outbound request, which is correct for OpenAI-to-OpenAI - // passthrough (OpenAI does not have a native cache_control field). } - // =========================================================================== - // Error tests - // =========================================================================== - - /// unknown provider returns OpenAI-shaped 404 #[tokio::test] async fn unknown_model_returns_openai_shaped_400() { let app = build_router(empty_state()); let body = make_chat_body("nonexistent-model", false); let resp = app.oneshot(chat_request(&body)).await.unwrap(); - // With provider-based routing, empty state means the provider is unknown -> 404. assert_eq!(resp.status(), StatusCode::NOT_FOUND); let resp_body: Value = serde_json::from_slice( @@ -764,20 +673,17 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert_eq!(resp_body["error"]["type"], "invalid_request_error"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Must NOT be Anthropic-shaped. assert!( !resp_body.as_object().unwrap().contains_key("type"), "must not have Anthropic type field" ); } - /// invalid JSON/client decode failure returns OpenAI-shaped 400 #[tokio::test] async fn invalid_json_returns_openai_shaped_400() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -800,7 +706,6 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert_eq!(resp_body["error"]["type"], "invalid_request_error"); assert!( @@ -812,14 +717,12 @@ mod tests { ); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Must NOT be Anthropic-shaped. assert!( !resp_body.as_object().unwrap().contains_key("type"), "must not have Anthropic type field" ); } - /// upstream failure returns OpenAI-shaped 502 #[tokio::test] async fn upstream_failure_returns_openai_shaped_502() { let mock_url = spawn_mock_500().await; @@ -837,18 +740,14 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert_eq!(resp_body["error"]["type"], "api_error"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - /// provider decode failure returns OpenAI-shaped 502 - /// (sending Anthropic-shaped response to an OpenAI Chat endpoint causes decode failure) #[tokio::test] async fn provider_decode_failure_returns_openai_shaped_502() { - // Use an OpenAI Chat provider but mock returns malformed response. let malformed = serde_json::to_vec(&json!({"not": "a valid response"})).unwrap(); let mock_url = spawn_mock_server(malformed, "application/json").await; let state = state_with_openai_chat_provider(&mock_url); @@ -865,17 +764,11 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert_eq!(resp_body["error"]["type"], "api_error"); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - // =========================================================================== - // Streaming tests - // =========================================================================== - - /// streaming response emits chat.completion.chunk #[tokio::test] async fn streaming_response_emits_chat_completion_chunk() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -891,14 +784,12 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Must contain chat.completion.chunk objects. assert!( text.contains("chat.completion.chunk"), "must emit chat.completion.chunk objects" ); } - /// streaming response sets SSE content type #[tokio::test] async fn streaming_response_sets_sse_content_type() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -921,7 +812,6 @@ mod tests { ); } - /// stream: true uses shared streaming pipeline #[tokio::test] async fn stream_true_uses_shared_streaming_pipeline() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -932,7 +822,6 @@ mod tests { let resp = app.oneshot(chat_request(&body)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Verify the request ID header is present (set by shared pipeline). let request_id = resp.headers().get("x-request-id"); assert!( request_id.is_some(), @@ -944,21 +833,20 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Verify text content came through the core pipeline. assert!( text.contains("Hi from OpenAI!"), "expected translated text delta in SSE output, got: {text}" ); } - /// streaming tool-call start/delta/stop maps to choices[].delta.tool_calls - /// Verifies the full start/delta/stop sequence: start has id + function.name, - /// delta has function.arguments with partial JSON, and correct index values. + #[expect( + clippy::too_many_lines, + reason = "one end-to-end streaming scenario: the mock feed, the SSE \ + consumption, and the start/delta/stop assertions read as a \ + single narrative" + )] #[tokio::test] async fn streaming_tool_call_maps_to_delta_tool_calls() { - // This test uses an Anthropic provider that returns tool call events. - // The core pipeline translates them, and the OpenAI client encoder maps - // them to choices[].delta.tool_calls. let app = Router::new().route( "/{*path}", post(|| async move { @@ -1000,8 +888,6 @@ mod tests { wait_for_ready(addr).await; let mock_url = format!("http://{}/providers/mock-provider/v1/messages", addr); - // Use Anthropic provider (which supports tool call events in SSE), - // but call the OpenAI chat/completions endpoint. let state = state_with_anthropic_provider(&mock_url); let app = build_router(state); @@ -1018,15 +904,12 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Parse the SSE output into structured chunks and verify tool_call mapping. let chunks: Vec = text .lines() .filter(|l| l.starts_with("data: ") && !l.contains("[DONE]")) .filter_map(|l| serde_json::from_str(l.trim_start_matches("data: ")).ok()) .collect(); - // Find the tool_call start chunk: should have delta.tool_calls[0].id and - // delta.tool_calls[0].function.name let start_chunk = chunks .iter() .find(|c| { @@ -1057,9 +940,6 @@ mod tests { "tool_call index must be 0" ); - // Find the tool_call delta chunk: should have delta.tool_calls[0].function.arguments - // The provider decoder aggregates partial JSON deltas, so the delta chunk may - // contain the full or partial arguments string. let delta_chunk = chunks .iter() .find(|c| { @@ -1076,14 +956,12 @@ mod tests { let tc_delta = &delta_chunk["choices"][0]["delta"]["tool_calls"][0]; let args = tc_delta["function"]["arguments"].as_str().unwrap_or(""); - // The arguments may be a partial or full JSON string containing city/SF data. assert!( !args.is_empty() || tc_delta["function"].get("arguments").is_some(), "tool_call delta must have a function.arguments field (may be empty string for start chunk)" ); } - /// streaming usage maps when requested #[tokio::test] async fn streaming_usage_maps_when_requested() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -1106,15 +984,11 @@ mod tests { assert!(text.contains("[DONE]"), "stream must end with [DONE]"); - // When include_usage is set, the final chunk should have a usage object. - // The mock doesn't send usage data, so the encoder emits zero usage -- - // but the "usage" key must still be present in the output. assert!( text.contains("\"usage\""), "stream with include_usage=true should contain a usage key, got: {text}" ); - // Parse chunks and verify the usage chunk has the expected fields. let chunks: Vec = text .lines() .filter(|l| l.starts_with("data: ") && !l.contains("[DONE]")) @@ -1137,7 +1011,6 @@ mod tests { ); } - /// streaming stop reason maps to finish_reason #[tokio::test] async fn streaming_stop_reason_maps_to_finish_reason() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -1153,14 +1026,12 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // The mock sends finish_reason: "stop", which should appear in the output. assert!( text.contains("stop"), "stream must contain finish_reason: stop" ); } - /// stream errors become OpenAI-shaped stream errors or route errors #[tokio::test] async fn stream_error_before_first_byte_returns_http_502() { let mock_url = spawn_mock_500().await; @@ -1182,12 +1053,10 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - /// streaming response ends with [DONE] #[tokio::test] async fn streaming_response_ends_with_done() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -1203,20 +1072,17 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Must end with data: [DONE]. assert!( text.contains("data: [DONE]"), "stream must end with data: [DONE], got: {text}" ); } - /// no Anthropic error envelope appears on this route #[tokio::test] async fn no_anthropic_error_envelope_on_openai_route() { let app = build_router(empty_state()); let body = make_chat_body("nonexistent-model", false); let resp = app.oneshot(chat_request(&body)).await.unwrap(); - // With provider-based routing, empty state -> UnknownProvider -> 404. assert_eq!(resp.status(), StatusCode::NOT_FOUND); let resp_body: Value = serde_json::from_slice( @@ -1226,7 +1092,6 @@ mod tests { ) .unwrap(); - // Must NOT have Anthropic-shaped error fields. assert!( !resp_body.as_object().unwrap().contains_key("type"), "must not have Anthropic 'type' field at top level" @@ -1238,18 +1103,11 @@ mod tests { .contains("not_found_error"), "must not use Anthropic error types" ); - // Must have OpenAI error shape. assert!(resp_body["error"].is_object()); assert_eq!(resp_body["error"]["type"], "invalid_request_error"); assert!(resp_body["error"]["code"].is_null()); } - // =========================================================================== - // Source guard test (integration level) - // =========================================================================== - - /// source guard confirms routes/chat.rs respects architecture boundaries, - /// endpoint-classifier, fallback, or direct-provider execution path #[test] fn source_guard_chat_rs_architecture_boundaries() { let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); @@ -1292,17 +1150,9 @@ mod tests { } } - // =========================================================================== - // Protocol-aware 404 tests - // =========================================================================== - - /// unmatched /v1/chat/* path returns OpenAI-shaped 404 #[tokio::test] async fn not_found_openai_path_returns_openai_shaped_error() { let app = build_router(empty_state()); - // Use a path that contains the OpenAI chat completions segment but has an - // extra suffix so it does NOT match the mounted route. The path still - // matches the OpenAI protocol heuristic in the 404 handler. let req = Request::builder() .method("POST") .uri("/providers/mock-provider/v1/chat/completions/extra") @@ -1319,23 +1169,19 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Must NOT be Anthropic-shaped. assert!( resp_body.get("type").is_none() || resp_body["type"].is_null(), "must not have Anthropic 'type' field" ); } - /// unmatched /v1/messages path returns Anthropic-shaped 404 #[tokio::test] async fn not_found_anthropic_path_returns_anthropic_shaped_error() { let app = build_router(empty_state()); - // Use a path under /v1/messages that does NOT match any mounted route. let req = Request::builder() .method("POST") .uri("/providers/mock-provider/v1/messages/typo") @@ -1352,16 +1198,10 @@ mod tests { ) .unwrap(); - // Must be Anthropic error shape. assert_eq!(resp_body["type"], "error"); assert!(resp_body["error"]["type"].is_string()); } - // =========================================================================== - // Edge-case tests - // =========================================================================== - - /// empty request body returns OpenAI-shaped 400 #[tokio::test] async fn empty_body_returns_openai_shaped_400() { let app = build_router(state_with_mock_provider_no_upstream()); @@ -1381,13 +1221,11 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - /// request with wrong field types returns OpenAI-shaped 400 #[tokio::test] async fn wrong_field_types_returns_openai_shaped_400() { let app = build_router(state_with_mock_provider_no_upstream()); @@ -1409,13 +1247,11 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - /// request missing required 'model' field returns OpenAI-shaped 400 #[tokio::test] async fn missing_model_returns_openai_shaped_400() { let app = build_router(state_with_mock_provider_no_upstream()); @@ -1437,11 +1273,9 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Message should indicate model is required (decode_request checks this). let msg = resp_body["error"]["message"].as_str().unwrap(); assert!( msg.to_lowercase().contains("model"), @@ -1449,7 +1283,6 @@ mod tests { ); } - /// request missing required 'messages' field returns OpenAI-shaped 400 #[tokio::test] async fn missing_messages_returns_openai_shaped_400() { let app = build_router(state_with_mock_provider_no_upstream()); @@ -1469,11 +1302,9 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Message should indicate messages is required. let msg = resp_body["error"]["message"].as_str().unwrap(); assert!( msg.to_lowercase().contains("messages"), @@ -1481,7 +1312,6 @@ mod tests { ); } - /// request with invalid UTF-8 bytes returns OpenAI-shaped 400 #[tokio::test] async fn invalid_utf8_body_returns_openai_shaped_400() { let app = build_router(state_with_mock_provider_no_upstream()); @@ -1501,13 +1331,11 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); } - /// boundary value for max_tokens (0) is accepted #[tokio::test] async fn max_tokens_zero_is_accepted() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -1521,7 +1349,6 @@ mod tests { "stream": false }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // Should be accepted (the upstream decides whether to reject max_tokens=0). assert_ne!( resp.status(), StatusCode::INTERNAL_SERVER_ERROR, @@ -1529,15 +1356,9 @@ mod tests { ); } - // =========================================================================== - // Missing tests from audit round 2 - // =========================================================================== - - /// oversized body (> 32 MiB) returns 413 Payload Too Large #[tokio::test] async fn chat_completions_oversized_body_returns_payload_too_large() { let app = build_router(empty_state()); - // 33 MiB body (exceeds the 32 MiB limit). let oversized_body = "X".repeat(33 * 1024 * 1024); let req = Request::builder() .method("POST") @@ -1548,8 +1369,6 @@ mod tests { let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); - // MEDIUM-7: the 413 is normalised into an OpenAI-shaped JSON envelope - // (not plain text) and carries `x-request-id`. assert_eq!( resp.headers().get("content-type").unwrap(), "application/json" @@ -1567,12 +1386,9 @@ mod tests { assert!(json["error"]["code"].is_null()); } - /// MEDIUM-7: a request that exceeds the per-route timeout returns a JSON 408 - /// (not an empty body) carrying an `x-request-id`, matching the protocol schema. #[tokio::test] async fn timeout_returns_normalised_json_408() { let mock_url = spawn_hanging_mock().await; - // 50 ms timeout: the hanging upstream holds the handler past it. let state = state_with_provider_timeout( &mock_url, "openai_chat_completions", @@ -1616,20 +1432,11 @@ mod tests { assert!(json["error"]["code"].is_null()); } - /// A malformed SSE data frame is silently skipped by the Anthropic provider - /// adapter (it returns `Ok(vec![])` on bad JSON, see `decode_frame`), so it does - /// NOT drive the in-band error path. This test verifies the stream still - /// completes gracefully (`[DONE]`) when a mid-stream frame is unparseable, and - /// that the valid events around it are delivered. For a genuine in-band error - /// assertion see `in_band_stream_error_emits_openai_error_chunk`. #[tokio::test] async fn malformed_frame_is_skipped_and_stream_completes() { - // Spawn a mock that sends one valid chunk then an invalid/malformed SSE event. let app = Router::new().route( "/{*path}", post(|| async move { - // First, send a valid message_start event via Anthropic SSE format - // (since we're testing with an Anthropic provider mock). let events = vec![ Event::default() .event("message_start") @@ -1637,7 +1444,6 @@ mod tests { Event::default() .event("content_block_delta") .data(r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}"#), - // Then send a malformed event that the Anthropic adapter skips. Event::default() .event("content_block_delta") .data(r#"this is not valid JSON for the provider decoder"#), @@ -1666,7 +1472,6 @@ mod tests { "stream": true }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // The first byte was sent successfully, so HTTP status should be 200. assert_eq!(resp.status(), StatusCode::OK); let bytes = axum::body::to_bytes(resp.into_body(), 1024 * 1024) @@ -1674,43 +1479,21 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // The malformed frame is skipped, so the stream completes with [DONE]. assert!( text.contains("[DONE]"), "stream must end with [DONE] even when a frame is skipped, got: {text}" ); } - /// A genuine in-band (post-first-byte) error emits an OpenAI-shaped error chunk - /// followed by `[DONE]`. - /// - /// Drives the real error path (audit `in-band-error-event-path-untested`): - /// after a valid `message_start` crosses the first byte (committing HTTP 200), - /// the mock sends a chunk that the SSE framer rejects — a "data:" line whose - /// payload is invalid UTF-8 (a lone `0xFF` byte). `SseFramer::push_chunk` - /// returns `Err(ProviderError::Utf8)` for that line, which is the - /// post-first-byte framing-error branch in `build_sse_output_stream`; for an - /// OpenAI Chat client that branch calls `emit_stream_error`, emitting an OpenAI - /// error JSON chunk (`server_error`) and then terminating with `[DONE]`. - /// - /// This is deterministic regardless of how reqwest buffers the body: invalid - /// UTF-8 inside a drained line is a hard framer error, not silently tolerated - /// like malformed JSON (which the Anthropic adapter skips) and not sensitive to - /// chunk boundaries like an over-long line. #[tokio::test] async fn in_band_stream_error_emits_openai_error_chunk() { use bytes::Bytes; - // A valid Anthropic message_start frame (terminated by a blank line) so the - // first byte crosses and HTTP 200 is committed. let first_frame = Bytes::from( "event: message_start\n\ data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_err\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-sonnet-4-6\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}\n\n" .to_owned(), ); - // A second chunk whose "data:" line payload is invalid UTF-8 (a lone - // continuation byte `0xFF`), terminated by a blank line so the framer - // drains it and hits `str::from_utf8` -> Err. let bad_frame = Bytes::from(vec![b'd', b'a', b't', b'a', b':', b' ', 0xFF, b'\n', b'\n']); let app = Router::new().route( @@ -1719,10 +1502,6 @@ mod tests { let first = first_frame.clone(); let bad = bad_frame.clone(); async move { - // Yield the valid frame, then the malformed-UTF-8 frame. The - // proxy's SSE framer reports the second as Err once the first - // byte has crossed, deterministically driving the in-band - // error path. let body_stream = futures::stream::iter(vec![ Ok::(first), Ok::(bad), @@ -1750,7 +1529,6 @@ mod tests { "stream": true }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // The first byte was sent successfully, so HTTP status should be 200. assert_eq!(resp.status(), StatusCode::OK); let bytes = axum::body::to_bytes(resp.into_body(), 1024 * 1024) @@ -1758,23 +1536,18 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // The in-band error must surface as an OpenAI-shaped error chunk with the - // `server_error` type (set by emit_stream_error for OpenAI Chat). assert!( text.contains("\"type\":\"server_error\""), "in-band error must emit an OpenAI error chunk, got: {text}" ); - // ...and the stream must still terminate with [DONE]. assert!( text.contains("[DONE]"), "stream must end with [DONE] after the in-band error, got: {text}" ); } - /// upstream disconnect mid-stream completes with synthetic terminal #[tokio::test] async fn upstream_disconnect_completes_with_synthetic_terminal() { - // Spawn a mock that sends partial events then disconnects (no message_stop). let app = Router::new().route( "/{*path}", post(|| async move { @@ -1785,7 +1558,6 @@ mod tests { Event::default() .event("content_block_delta") .data(r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#), - // No message_stop -- simulates upstream disconnect. ]; let stream = futures::stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); let sse = Sse::new(stream).keep_alive(KeepAlive::default()); @@ -1818,19 +1590,16 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // The stream should end with [DONE] even though upstream disconnected. assert!( text.contains("[DONE]"), "stream must end with [DONE] after upstream disconnect, got: {text}" ); - // A synthetic finish_reason should be present. assert!( text.contains("stop"), "stream should contain synthetic finish_reason: stop" ); } - /// request with only system messages (no user message) does not panic #[tokio::test] async fn system_only_messages_does_not_panic() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -1844,7 +1613,6 @@ mod tests { ] }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // Should not panic. It may succeed (upstream accepts it) or return an error. assert_ne!( resp.status(), StatusCode::INTERNAL_SERVER_ERROR, @@ -1852,14 +1620,6 @@ mod tests { ); } - // =========================================================================== - // Missing tests from audit round 3 - // =========================================================================== - - /// Unknown model with empty routing table returns 400 Bad Request with - /// OpenAI-shaped error envelope (audit round 3 coverage). This exercises - /// the core pipeline error path when the model routing table has no matching - /// entry, confirming the OpenAI error shape for the chat completions route. #[tokio::test] async fn unknown_model_empty_routing_table_openai_shaped_400() { let app = build_router(empty_state()); @@ -1870,7 +1630,6 @@ mod tests { }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // Empty routing table -> unknown provider -> 404 Not Found. assert_eq!( resp.status(), StatusCode::NOT_FOUND, @@ -1884,26 +1643,22 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); assert!(resp_body["error"]["code"].is_null(), "code must be null"); - // Must NOT be Anthropic-shaped (no top-level "type" field). assert!( resp_body.get("type").is_none() || resp_body["type"].is_null(), "must not have Anthropic 'type' field" ); } - /// Large messages array (1000 messages) is handled without errors. #[tokio::test] async fn large_messages_array_is_handled() { let mock_url = spawn_mock_openai_chat_non_stream().await; let state = state_with_openai_chat_provider(&mock_url); let app = build_router(state); - // Build a request with 1000 messages. let mut messages = Vec::new(); for i in 0..1000 { messages.push(json!({ @@ -1916,7 +1671,6 @@ mod tests { "messages": messages }); let resp = app.oneshot(chat_request(&body.to_string())).await.unwrap(); - // Should not panic or return an internal server error. assert_ne!( resp.status(), StatusCode::INTERNAL_SERVER_ERROR, @@ -1924,13 +1678,6 @@ mod tests { ); } - // =========================================================================== - // Bare route 404 tests (locked decision: bare /v1/* not registered) - // =========================================================================== - - /// Bare `/v1/chat/completions` (no provider segment) must return 404 with - /// an OpenAI-shaped error envelope. This verifies locked decision #2: - /// "Bare /v1/* API routes are not registered." #[tokio::test] async fn bare_v1_chat_completions_returns_404_openai_shaped() { let app = build_router(empty_state()); @@ -1952,13 +1699,10 @@ mod tests { ) .unwrap(); - // Must be OpenAI error shape. assert!(resp_body["error"].is_object(), "must have error object"); assert!(resp_body["error"]["message"].is_string()); } - /// Bare `/v1/messages` (no provider segment) must return 404 with - /// an Anthropic-shaped error envelope. #[tokio::test] async fn bare_v1_messages_returns_404_anthropic_shaped() { let app = build_router(empty_state()); @@ -1980,14 +1724,12 @@ mod tests { ) .unwrap(); - // Must be Anthropic error shape (has "type" and "error" fields). assert!( resp_body["type"].is_string(), "must have Anthropic 'type' field" ); } - /// Provider name with invalid characters (uppercase, dots) must return 400. #[tokio::test] async fn invalid_provider_name_returns_400() { let app = build_router(empty_state()); @@ -2007,10 +1749,14 @@ mod tests { ); } - /// Same model name routed to different providers via URL path. + #[expect( + clippy::too_many_lines, + reason = "one end-to-end routing scenario: two provider registries, two \ + dispatches, and their divergent assertions read as a single \ + narrative" + )] #[tokio::test] async fn same_model_routes_to_different_providers() { - // Create two providers both supporting chat_completions. let mut adapters_a = HashMap::new(); adapters_a.insert( "chat".to_owned(), @@ -2025,6 +1771,7 @@ mod tests { api_key: secrecy::SecretString::from("key-a"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: adapters_a, routes: llm_proxy_core::ProviderRoutesConfig { chat_completions: Some("chat".to_owned()), @@ -2050,6 +1797,7 @@ mod tests { api_key: secrecy::SecretString::from("key-b"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: adapters_b, routes: llm_proxy_core::ProviderRoutesConfig { chat_completions: Some("chat".to_owned()), @@ -2076,7 +1824,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2092,7 +1844,6 @@ mod tests { let body = serde_json::json!({"model": "gpt-4o", "messages": []}).to_string(); - // Request to provider-a — should resolve (not 404 for unknown provider). let req_a = Request::builder() .method("POST") .uri("/providers/provider-a/v1/chat/completions") @@ -2100,15 +1851,12 @@ mod tests { .body(Body::from(body.clone())) .unwrap(); let resp_a = app.clone().oneshot(req_a).await.unwrap(); - // Provider-a doesn't have a real upstream, so we expect an upstream error - // (502), NOT 404 unknown provider. assert_ne!( resp_a.status(), StatusCode::NOT_FOUND, "provider-a should resolve, not 404" ); - // Request to provider-b — same model, different provider. let req_b = Request::builder() .method("POST") .uri("/providers/provider-b/v1/chat/completions") diff --git a/crates/llm-proxy-server/tests/core_pipeline.rs b/crates/llm-proxy-server/tests/core_pipeline.rs index dc5dfde..41d855c 100644 --- a/crates/llm-proxy-server/tests/core_pipeline.rs +++ b/crates/llm-proxy-server/tests/core_pipeline.rs @@ -1,9 +1,6 @@ -//! Integration tests for the core pipeline (Phase 8). -//! -//! These tests exercise the full request path through the axum router with a -//! local mock upstream server. Each test configures AppState with a provider route -//! pointing at the mock server, which returns canned provider-specific responses -//! that are translated through the core pipeline into Anthropic-shaped responses. +//! Integration tests for the core pipeline: the full request path through +//! the axum router against a local mock upstream, with canned provider +//! responses translated into Anthropic-shaped responses. #[cfg(test)] mod tests { @@ -32,15 +29,6 @@ mod tests { use serde_json::{Value, json}; use tower::ServiceExt; - // --------------------------------------------------------------------------- - // Mock server readiness helper - // --------------------------------------------------------------------------- - - /// Wait for a mock server to be ready by polling its TCP port. - /// - /// Replaces `tokio::time::sleep(Duration::from_millis(50))` with a - /// deterministic readiness check that retries until the server accepts - /// a connection or the timeout elapses. async fn wait_for_ready(addr: std::net::SocketAddr) { let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { @@ -54,19 +42,8 @@ mod tests { } } - /// Spawn a mock axum server on `listener`, returning immediately while the - /// server runs in the background. - /// - /// Unlike a bare `tokio::spawn(async move { axum::serve(..).await })` whose - /// `JoinHandle` is dropped (silently swallowing any panic in the mock server - /// task — audit `testing-audit:mock-server-joinhandle-swallow`), this wrapper - /// catches a panic in the serve future, prints it to stderr, and aborts the - /// process. A mock-handler bug therefore surfaces loudly as a test failure - /// instead of the connection just closing mid-stream with no diagnostic. fn spawn_mock_serve(listener: tokio::net::TcpListener, app: Router) { tokio::spawn(async move { - // std::panic::catch_unwind requires UnwindSafe; the router/listener - // capture is fine for a mock that owns them exclusively. let result = std::panic::AssertUnwindSafe(axum::serve(listener, app).into_future()) .catch_unwind() .await; @@ -74,8 +51,6 @@ mod tests { Ok(Ok(())) => {} Ok(Err(e)) => eprintln!("mock axum::serve error: {e}"), Err(panic) => { - // Surface the panic payload, then abort so the test binary fails - // rather than continuing with a dead mock. let msg = panic .downcast_ref::() .map(|s| s.as_str()) @@ -88,11 +63,6 @@ mod tests { }); } - // --------------------------------------------------------------------------- - // Mock upstream server helpers - // --------------------------------------------------------------------------- - - /// A canned Anthropic Messages API success response. fn anthropic_success_response() -> Vec { serde_json::to_vec(&json!({ "id": "msg_mock123", @@ -107,7 +77,6 @@ mod tests { .unwrap() } - /// A canned OpenAI Chat Completions API success response. fn openai_chat_success_response() -> Vec { serde_json::to_vec(&json!({ "id": "chatcmpl-test123", @@ -131,7 +100,6 @@ mod tests { .unwrap() } - /// A canned OpenAI Responses API success response. fn openai_responses_success_response() -> Vec { serde_json::to_vec(&json!({ "id": "resp_test123", @@ -155,7 +123,6 @@ mod tests { .unwrap() } - /// A canned Gemini generateContent API success response. fn gemini_success_response() -> Vec { serde_json::to_vec(&json!({ "candidates": [{ @@ -174,8 +141,6 @@ mod tests { .unwrap() } - /// Build AppState in TOML mode with a single provider of the given protocol - /// routing to the given mock endpoint. fn state_with_provider( mock_endpoint: &str, protocol: &str, @@ -183,7 +148,6 @@ mod tests { _model_name: &str, pricing: HashMap, ) -> AppState { - // Build routes based on protocol so the provider-based routing can resolve. let routes = match protocol { "anthropic_messages" => llm_proxy_core::ProviderRoutesConfig { messages: Some(adapter_name.to_owned()), @@ -204,6 +168,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -238,7 +203,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -255,7 +224,6 @@ mod tests { ) } - /// Convenience: AppState with an Anthropic provider. fn state_with_anthropic_provider(mock_endpoint: &str) -> AppState { state_with_provider( mock_endpoint, @@ -266,8 +234,6 @@ mod tests { ) } - /// Convenience: AppState with an Anthropic provider and a pricing table for the - /// upstream model, so cost computation is exercised end-to-end. fn state_with_anthropic_provider_priced( mock_endpoint: &str, pricing: HashMap, @@ -281,7 +247,6 @@ mod tests { ) } - /// Convenience: AppState with an OpenAI Chat provider. fn state_with_openai_chat_provider(mock_endpoint: &str) -> AppState { state_with_provider( mock_endpoint, @@ -292,7 +257,6 @@ mod tests { ) } - /// Convenience: AppState with an OpenAI Responses provider. fn state_with_openai_responses_provider(mock_endpoint: &str) -> AppState { state_with_provider( mock_endpoint, @@ -303,7 +267,6 @@ mod tests { ) } - /// Convenience: AppState with a Gemini provider. fn state_with_gemini_provider(mock_endpoint: &str) -> AppState { state_with_provider( mock_endpoint, @@ -314,9 +277,6 @@ mod tests { ) } - /// Spawn a local mock axum server returning a canned response body. - /// The handler receives the response body as bytes and returns it with the - /// given content type. async fn spawn_mock_server(response_body: Vec, content_type: &str) -> String { let ct = content_type.to_owned(); let app = Router::new().route( @@ -336,29 +296,22 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a local mock axum server returning canned Anthropic non-streaming - /// responses. Returns the base URL. async fn spawn_mock_anthropic_non_stream() -> String { spawn_mock_server(anthropic_success_response(), "application/json").await } - /// Spawn a local mock axum server returning OpenAI Chat non-streaming response. async fn spawn_mock_openai_chat_non_stream() -> String { spawn_mock_server(openai_chat_success_response(), "application/json").await } - /// Spawn a local mock axum server returning OpenAI Responses non-streaming response. async fn spawn_mock_openai_responses_non_stream() -> String { spawn_mock_server(openai_responses_success_response(), "application/json").await } - /// Spawn a local mock axum server returning Gemini non-streaming response. async fn spawn_mock_gemini_non_stream() -> String { spawn_mock_server(gemini_success_response(), "application/json").await } - /// Spawn a local mock axum server returning canned Anthropic streaming SSE - /// events. async fn spawn_mock_anthropic_stream() -> String { let app = Router::new().route( "/{*path}", @@ -399,7 +352,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server returning OpenAI Chat streaming SSE events. async fn spawn_mock_openai_chat_stream() -> String { let app = Router::new().route( "/{*path}", @@ -430,7 +382,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server returning OpenAI Responses streaming SSE events. async fn spawn_mock_openai_responses_stream() -> String { let app = Router::new().route( "/{*path}", @@ -465,7 +416,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server returning Gemini streaming SSE events. async fn spawn_mock_gemini_stream() -> String { let app = Router::new().route( "/{*path}", @@ -490,7 +440,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server that returns HTTP 500. async fn spawn_mock_500() -> String { let app = Router::new().route( "/{*path}", @@ -511,14 +460,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server that returns an SSE stream where a valid - /// `content_block_delta` carrying the marker text `"good-delta"` is followed by - /// a malformed JSON event, then a final `message_stop`. - /// - /// The Anthropic provider adapter silently skips the malformed frame (returns - /// `Ok(vec![])`), so the valid delta BEFORE it is delivered and the stream - /// completes normally. The marker text lets tests assert that valid events - /// survive a skipped malformed frame (audit `duplicate-malformed-stream-tests`). async fn spawn_mock_malformed_stream() -> String { let app = Router::new().route( "/{*path}", @@ -527,14 +468,12 @@ mod tests { Event::default() .event("message_start") .data(r#"{"type":"message_start","message":{"id":"msg_mock","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}"#), - // A valid delta carrying a marker so tests can assert it is delivered. Event::default() .event("content_block_start") .data(r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#), Event::default() .event("content_block_delta") .data(r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"good-delta"}}"#), - // Malformed event: invalid JSON, silently skipped by the adapter. Event::default() .event("content_block_delta") .data("this is not valid json {{{"), @@ -558,13 +497,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server that returns a genuine Anthropic `type: error` event - /// mid-stream after a valid `message_start`. - /// - /// Unlike `spawn_mock_malformed_stream` (whose malformed JSON is silently - /// skipped), this drives the real in-band error path: the Anthropic adapter maps - /// `event: error` to `CoreEvent::Error`, which the client encoder surfaces as an - /// `event: error` SSE frame (audit `in-band-error-event-path-untested`). async fn spawn_mock_in_band_error_stream() -> String { let app = Router::new().route( "/{*path}", @@ -573,8 +505,6 @@ mod tests { Event::default() .event("message_start") .data(r#"{"type":"message_start","message":{"id":"msg_err","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}"#), - // A genuine protocol-level error event the decoder maps to - // CoreEvent::Error (anthropic.rs "error" arm). Event::default() .event("error") .data(r#"{"type":"error","error":{"type":"overloaded_error","message":"Too many requests"}}"#), @@ -598,13 +528,10 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server that accepts a request, sends a few events, then - /// abruptly closes the connection (simulating upstream disconnect). async fn spawn_mock_disconnect_stream() -> String { let app = Router::new().route( "/{*path}", post(|| async move { - // Return only message_start, then stop (simulating disconnect). let events = vec![ Event::default() .event("message_start") @@ -612,7 +539,6 @@ mod tests { Event::default() .event("content_block_start") .data(r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#), - // No message_stop -- stream is truncated. ]; let stream = futures::stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); let sse = Sse::new(stream).keep_alive(KeepAlive::default()); @@ -630,8 +556,6 @@ mod tests { format!("http://{}/providers/mock-provider/v1/messages", addr) } - /// Spawn a mock server that captures whether it received a request. - /// Returns (base_url, request_received_flag). async fn spawn_mock_with_request_tracker(response_body: Vec) -> (String, Arc) { let received = Arc::new(AtomicBool::new(false)); let received_clone = Arc::clone(&received); @@ -659,8 +583,6 @@ mod tests { ) } - /// Spawn a mock server that records the received request body. - /// Returns (base_url, received_body_arc). async fn spawn_mock_with_body_capture( response_body: Vec, ) -> (String, Arc>>>) { @@ -703,7 +625,6 @@ mod tests { .to_string() } - /// Build a request to /providers/mock-provider/v1/messages. fn messages_request(body: &str) -> Request { Request::builder() .method("POST") @@ -713,11 +634,6 @@ mod tests { .unwrap() } - // =========================================================================== - // Non-streaming tests: Anthropic provider - // =========================================================================== - - /// configured Anthropic provider returns Anthropic response through core pipeline #[tokio::test] async fn anthropic_provider_returns_anthropic_response() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -744,7 +660,6 @@ mod tests { ); } - /// request ID header is present on success #[tokio::test] async fn request_id_header_present_on_success() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -764,7 +679,6 @@ mod tests { assert!(!id_str.is_empty(), "x-request-id must not be empty"); } - /// upstream 500 returns 502 Bad Gateway #[tokio::test] async fn upstream_500_returns_502() { let mock_url = spawn_mock_500().await; @@ -785,8 +699,6 @@ mod tests { assert_eq!(resp_body["error"]["type"], "api_error"); } - /// unknown model is passed through to upstream (provider-based routing does not - /// validate model names locally; the upstream provider decides) #[tokio::test] async fn unknown_model_is_passed_through() { let (mock_url, received) = @@ -796,11 +708,8 @@ mod tests { let body = make_messages_body("nonexistent-model", false); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // With provider-based routing, the model name is passed through to the - // upstream provider. The mock returns 200, so we get a successful response. assert_eq!(resp.status(), StatusCode::OK); - // Verify upstream WAS called (model name is passed through). tokio::time::sleep(Duration::from_millis(100)).await; assert!( received.load(Ordering::SeqCst), @@ -808,11 +717,6 @@ mod tests { ); } - // =========================================================================== - // Non-streaming tests: OpenAI Chat provider - // =========================================================================== - - /// configured OpenAI Chat provider returns Anthropic response #[tokio::test] async fn openai_chat_provider_returns_anthropic_response() { let mock_url = spawn_mock_openai_chat_non_stream().await; @@ -829,7 +733,6 @@ mod tests { .unwrap(), ) .unwrap(); - // Response must be Anthropic-shaped. assert_eq!(resp_body["type"], "message"); assert_eq!(resp_body["role"], "assistant"); assert!( @@ -841,11 +744,6 @@ mod tests { ); } - // =========================================================================== - // Non-streaming tests: OpenAI Responses provider - // =========================================================================== - - /// configured Responses provider returns Anthropic response #[tokio::test] async fn openai_responses_provider_returns_anthropic_response() { let mock_url = spawn_mock_openai_responses_non_stream().await; @@ -862,7 +760,6 @@ mod tests { .unwrap(), ) .unwrap(); - // Response must be Anthropic-shaped. assert_eq!(resp_body["type"], "message"); assert_eq!(resp_body["role"], "assistant"); assert!( @@ -874,11 +771,6 @@ mod tests { ); } - // =========================================================================== - // Non-streaming tests: Gemini provider - // =========================================================================== - - /// configured Gemini provider returns Anthropic response #[tokio::test] async fn gemini_provider_returns_anthropic_response() { let mock_url = spawn_mock_gemini_non_stream().await; @@ -895,7 +787,6 @@ mod tests { .unwrap(), ) .unwrap(); - // Response must be Anthropic-shaped. assert_eq!(resp_body["type"], "message"); assert_eq!(resp_body["role"], "assistant"); assert!( @@ -907,13 +798,6 @@ mod tests { ); } - // =========================================================================== - // Streaming tests: Anthropic provider - // =========================================================================== - - /// Collect the full SSE response body as a UTF-8 string for an Anthropic - /// streaming request against the given mock URL. Shared by the focused - /// streaming tests below so each test can assert a single behavior. async fn collect_anthropic_stream_body(mock_url: &str) -> String { let state = state_with_anthropic_provider(mock_url); let app = build_router(state); @@ -926,7 +810,6 @@ mod tests { String::from_utf8(bytes.to_vec()).unwrap() } - /// stream:true Anthropic provider returns HTTP 200 with an SSE content-type #[tokio::test] async fn stream_anthropic_provider_returns_sse_content_type() { let mock_url = spawn_mock_anthropic_stream().await; @@ -949,8 +832,6 @@ mod tests { ); } - /// stream:true Anthropic provider includes an x-request-id header on the - /// stream response #[tokio::test] async fn stream_anthropic_provider_includes_request_id() { let mock_url = spawn_mock_anthropic_stream().await; @@ -968,8 +849,6 @@ mod tests { ); } - /// stream:true Anthropic provider emits the core SSE protocol events - /// (message_start, content_block_delta, message_stop) #[tokio::test] async fn stream_anthropic_provider_emits_sse_protocol_events() { let mock_url = spawn_mock_anthropic_stream().await; @@ -989,7 +868,6 @@ mod tests { ); } - /// stream:true Anthropic provider forwards the translated text delta #[tokio::test] async fn stream_anthropic_provider_forwards_text_delta() { let mock_url = spawn_mock_anthropic_stream().await; @@ -1001,26 +879,16 @@ mod tests { ); } - /// stream:true Anthropic provider surfaces the upstream provider's real message - /// id (from `message_start`) rather than a proxy-generated synthetic id. #[tokio::test] async fn stream_anthropic_provider_surfaces_upstream_message_id() { let mock_url = spawn_mock_anthropic_stream().await; let text = collect_anthropic_stream_body(&mock_url).await; - // The mock emits `message_start` with id "msg_mock_stream"; the client SSE - // must carry that real id, not a synthetic `msg_{uuid}`. assert!( text.contains("msg_mock_stream"), "client SSE must surface the upstream message id; got: {text}" ); } - // =========================================================================== - // Streaming tests: OpenAI Chat provider - // =========================================================================== - - /// Collect the full SSE response body for an OpenAI Chat streaming request - /// against the given mock URL. Shared by the focused streaming tests below. async fn collect_openai_chat_stream_body(mock_url: &str) -> String { let state = state_with_openai_chat_provider(mock_url); let app = build_router(state); @@ -1033,34 +901,16 @@ mod tests { String::from_utf8(bytes.to_vec()).unwrap() } - /// stream:true OpenAI Chat provider surfaces the upstream provider's real chunk - /// id rather than a proxy-generated synthetic `chatcmpl-{uuid}`. - /// - /// The OpenAI client encoder uses only the id passed at construction (it - /// ignores `CoreEvent::MessageStart`'s id), so this only succeeds because the - /// first event is buffered to extract the upstream id *before* the encoder is - /// built. This is the distinguishing case for the buffer-first-event change: - /// before it, OpenAI streaming always emitted a synthetic id. #[tokio::test] async fn stream_openai_chat_provider_surfaces_upstream_message_id() { let mock_url = spawn_mock_openai_chat_stream().await; let text = collect_openai_chat_stream_body(&mock_url).await; - // The mock emits chunks with id "chatcmpl-stream"; the client SSE must carry - // that real upstream id, not a synthetic chatcmpl-{uuid}. assert!( text.contains("chatcmpl-stream"), "client SSE must surface the upstream chunk id; got: {text}" ); } - /// The OpenAI Chat streaming path must populate - /// `ResponseCompleted.upstream_message_id` on the event bus. The OpenAI Chat - /// client encoder ignores `CoreEvent::MessageStart`'s id (it uses only the id - /// passed at construction), so the upstream id reaches the response solely via - /// the buffer-first-event extraction. The SSE-body test above checks the wire; - /// this guards the persistence sink, which is a distinct surface (plan Test - /// Plan: "Stream request emits ResponseCompleted with the real - /// upstream_message_id from Step 1" -- the OpenAI path was uncovered). #[tokio::test] async fn event_bus_records_openai_chat_stream_upstream_message_id() { use std::sync::Arc; @@ -1076,8 +926,6 @@ mod tests { let body = make_messages_body("gpt-4o", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Drain the SSE body so the spawned task runs to completion and emits the - // stream ResponseCompleted (it fires at the end of the task). let _ = axum::body::to_bytes(resp.into_body(), 1024 * 1024) .await .unwrap(); @@ -1090,7 +938,6 @@ mod tests { _ => None, }) .expect("ResponseCompleted emitted for OpenAI Chat stream"); - // The OpenAI stream mock seeds every chunk with id "chatcmpl-stream". assert_eq!( completed.upstream_message_id.as_deref(), Some("chatcmpl-stream"), @@ -1098,9 +945,6 @@ mod tests { ); } - /// A successful non-stream request emits `RequestReceived` then - /// `ResponseCompleted` on the event bus, with the provider-reported usage and - /// upstream message id on the completed event, and no api-key leakage. #[tokio::test] async fn event_bus_records_request_received_and_response_completed() { use std::sync::Arc; @@ -1133,8 +977,6 @@ mod tests { assert_eq!(rc.upstream_message_id.as_deref(), Some("msg_mock123")); assert_eq!(rc.usage.input_tokens, 10); assert_eq!(rc.usage.output_tokens, 5); - // No pricing configured for this provider -> cost must be None (the - // "no pricing -> Cost = None" contract from the plan's Test Plan). assert!( rc.cost.is_none(), "cost must be None when no pricing is configured, got {:?}", @@ -1144,7 +986,6 @@ mod tests { other => panic!("expected ResponseCompleted, got {other:?}"), } - // Defense-in-depth: no api-key fragments in the serialized event stream. let json = serde_json::to_string(&events).expect("serialize events"); assert!( !json.contains("test-key") && !json.contains("sk-"), @@ -1152,10 +993,6 @@ mod tests { ); } - /// When pricing IS configured for the upstream model, a successful non-stream - /// request attaches a computed `Some(Cost)` to `ResponseCompleted` -- the - /// end-to-end pricing->usage->cost->event wiring (audit finding: cost pipeline - /// never exercised end-to-end). #[tokio::test] async fn priced_provider_attaches_cost_to_response_completed() { use std::sync::Arc; @@ -1193,10 +1030,6 @@ mod tests { _ => None, }) .expect("ResponseCompleted emitted"); - // Mock usage is input=10, output=5. - // input = 10 * 0.000003 = 0.00003 - // output = 5 * 0.000015 = 0.000075 - // total = 0.000105 let cost = rc .cost .expect("cost must be Some when pricing is configured"); @@ -1205,10 +1038,6 @@ mod tests { assert_eq!(cost.total.normalize().to_string(), "0.000105"); } - /// The first upstream chunk may decode to MULTIPLE CoreEvents (message_start - /// immediately followed by content_block_delta). The buffered first-event - /// replay must emit ALL of them exactly once -- none lost, none duplicated - /// (plan requirement; audit finding: multi-event first-frame test missing). #[tokio::test] async fn stream_multi_event_first_frame_emits_all_without_loss_or_duplication() { let body = anthropic_sse_body("msg_multi", ""); @@ -1235,8 +1064,6 @@ mod tests { ); } - /// A successful streaming request emits `RequestReceived` (streaming=true) then - /// `ResponseCompleted`, with the upstream message id captured from MessageStart. #[tokio::test] async fn event_bus_records_stream_request_and_response_completed() { use std::sync::Arc; @@ -1252,9 +1079,6 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Drain the SSE body so the spawned task runs to completion and emits the - // stream ResponseCompleted (it fires at the end of the task, before the - // output channel is dropped). let _ = axum::body::to_bytes(resp.into_body(), 1024 * 1024) .await .unwrap(); @@ -1278,15 +1102,174 @@ mod tests { _ => None, }) .expect("ResponseCompleted emitted for stream"); - // The Anthropic stream mock seeds message_start with id "msg_mock_stream". assert_eq!( completed.upstream_message_id.as_deref(), Some("msg_mock_stream") ); } - /// A request for an unknown provider emits `ResponseFailed` with the mapped - /// HTTP status (404) and the provider name from the URL path. + fn client_auth_for(key_id: &str, token: &str) -> llm_proxy_server::ClientAuth { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + token.as_bytes(), + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + let mut clients = HashMap::new(); + clients.insert(key_id.to_owned(), secrecy::SecretString::from(phc)); + llm_proxy_server::ClientAuth::from_config(&llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }) + .expect("valid client auth") + } + + #[tokio::test] + async fn enforced_client_keys_reject_missing_bearer_with_sanitized_401() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let body = make_messages_body("claude-sonnet-4-6", false); + let resp = app.oneshot(messages_request(&body)).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + assert!( + resp.headers().get("x-request-id").is_some(), + "401 must carry the request id" + ); + let body = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert!( + !text.contains("sekrit-token"), + "401 body must not echo any secret material, got: {text}" + ); + } + + #[tokio::test] + async fn enforced_client_keys_reject_wrong_bearer_without_echoing_it() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let request = Request::builder() + .method("POST") + .uri("/providers/mock-provider/v1/messages") + .header("content-type", "application/json") + .header("authorization", "Bearer attacker-guess") + .body(Body::from(make_messages_body("claude-sonnet-4-6", false))) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let body = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert!( + !text.contains("attacker-guess"), + "401 body must not echo the attempted token, got: {text}" + ); + } + + #[tokio::test] + async fn valid_bearer_passes_and_attributes_caller_on_event() { + use std::sync::Arc; + + use llm_proxy_storage::{ProxyEvent, RecordingBus}; + + let mock_url = spawn_mock_anthropic_non_stream().await; + let bus = Arc::new(RecordingBus::new()); + let state = state_with_anthropic_provider(&mock_url) + .with_event_bus(Arc::::clone(&bus)) + .with_client_auth(std::sync::Arc::new(client_auth_for( + "laptop", + "sekrit-token", + ))); + let app = build_router(state); + + let request = Request::builder() + .method("POST") + .uri("/providers/mock-provider/v1/messages") + .header("content-type", "application/json") + .header("authorization", "bearer sekrit-token") + .body(Body::from(make_messages_body("claude-sonnet-4-6", false))) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let events = bus.snapshot(); + let received = events + .iter() + .find_map(|e| match e { + ProxyEvent::RequestReceived(r) => Some(r), + _ => None, + }) + .expect("RequestReceived emitted"); + assert_eq!(received.caller, "laptop"); + let event_json = + serde_json::to_string(&ProxyEvent::RequestReceived(received.clone())).unwrap(); + assert!( + !event_json.contains("sekrit-token"), + "client secret must never appear on the event, got: {event_json}" + ); + } + + #[tokio::test] + async fn no_auth_mode_attributes_anonymous_caller() { + use std::sync::Arc; + + use llm_proxy_storage::{ProxyEvent, RecordingBus}; + + let mock_url = spawn_mock_anthropic_non_stream().await; + let bus = Arc::new(RecordingBus::new()); + let state = state_with_anthropic_provider(&mock_url) + .with_event_bus(Arc::::clone(&bus)); + let app = build_router(state); + + let body = make_messages_body("claude-sonnet-4-6", false); + let resp = app.oneshot(messages_request(&body)).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let events = bus.snapshot(); + let received = events + .iter() + .find_map(|e| match e { + ProxyEvent::RequestReceived(r) => Some(r), + _ => None, + }) + .expect("RequestReceived emitted"); + assert_eq!(received.caller, "anonymous"); + } + + #[tokio::test] + async fn enforced_client_keys_leave_health_open() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let request = Request::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + #[tokio::test] async fn event_bus_records_response_failed_for_unknown_provider() { use std::sync::Arc; @@ -1299,7 +1282,6 @@ mod tests { .with_event_bus(Arc::::clone(&bus)); let app = build_router(state); - // Target a provider that is not registered. let body = make_messages_body("claude-sonnet-4-6", false); let req = Request::builder() .method("POST") @@ -1322,10 +1304,6 @@ mod tests { assert_eq!(failed.provider.as_deref(), Some("does-not-exist")); } - /// An empty upstream stream (no events at all) returns 502, not 200 -- the - /// pre-existing bug where finalization synthesized a terminal and committed - /// HTTP 200. Also asserts the event log gets a ResponseFailed, not a - /// ResponseCompleted, for the empty stream. #[tokio::test] async fn empty_upstream_stream_returns_502_not_200() { use std::sync::Arc; @@ -1361,12 +1339,6 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Upstream message ID edge-case tests (plan Test Plan §"Upstream message IDs") - // --------------------------------------------------------------------------- - - /// Build a raw Anthropic SSE body with a configurable message-start id and an - /// optional leading prefix ("" / "ping" / "keepalive"). fn anthropic_sse_body(message_start_id: &str, prefix: &str) -> Vec { let mut s = String::new(); match prefix { @@ -1401,7 +1373,6 @@ mod tests { s.into_bytes() } - /// MessageStart with an empty id falls back to a synthetic `msg_`-prefixed id. #[tokio::test] async fn stream_empty_message_start_id_falls_back_to_synthetic() { let body = anthropic_sse_body("", ""); @@ -1413,8 +1384,6 @@ mod tests { ); } - /// A stream that opens with a Ping event then MessageStart (same chunk) still - /// surfaces the real upstream message id. #[tokio::test] async fn stream_ping_then_message_start_surfaces_real_id() { let body = anthropic_sse_body("msg_ping_real", "ping"); @@ -1426,8 +1395,6 @@ mod tests { ); } - /// A stream that opens with a `:keepalive` comment (empty decode) then - /// MessageStart still surfaces the real upstream message id. #[tokio::test] async fn stream_keepalive_then_message_start_surfaces_real_id() { let body = anthropic_sse_body("msg_ka_real", "keepalive"); @@ -1439,7 +1406,6 @@ mod tests { ); } - /// stream:true OpenAI Chat provider returns HTTP 200 with an SSE content-type #[tokio::test] async fn stream_openai_chat_provider_returns_sse_content_type() { let mock_url = spawn_mock_openai_chat_stream().await; @@ -1462,14 +1428,11 @@ mod tests { ); } - /// stream:true OpenAI Chat provider emits Anthropic-shaped SSE protocol - /// events (message_start, message_stop) #[tokio::test] async fn stream_openai_chat_provider_emits_sse_protocol_events() { let mock_url = spawn_mock_openai_chat_stream().await; let text = collect_openai_chat_stream_body(&mock_url).await; - // The output must be Anthropic-shaped SSE events. assert!( text.contains("event: message_start"), "missing message_start event" @@ -1480,26 +1443,17 @@ mod tests { ); } - /// stream:true OpenAI Chat provider forwards the translated text delta #[tokio::test] async fn stream_openai_chat_provider_forwards_translated_text_delta() { let mock_url = spawn_mock_openai_chat_stream().await; let text = collect_openai_chat_stream_body(&mock_url).await; - // Verify text content was translated from OpenAI to Anthropic SSE format. assert!( text.contains("Hi from OpenAI!"), "expected translated text delta in SSE output, got: {text}" ); } - // =========================================================================== - // Streaming tests: OpenAI Responses provider - // =========================================================================== - - /// Collect the full SSE response body for an OpenAI Responses streaming - /// request against the given mock URL. Shared by the focused streaming tests - /// below. async fn collect_openai_responses_stream_body(mock_url: &str) -> String { let state = state_with_openai_responses_provider(mock_url); let app = build_router(state); @@ -1512,7 +1466,6 @@ mod tests { String::from_utf8(bytes.to_vec()).unwrap() } - /// stream:true Responses provider returns HTTP 200 with an SSE content-type #[tokio::test] async fn stream_openai_responses_provider_returns_sse_content_type() { let mock_url = spawn_mock_openai_responses_stream().await; @@ -1535,14 +1488,11 @@ mod tests { ); } - /// stream:true Responses provider emits Anthropic-shaped SSE protocol events - /// (message_start, message_stop) #[tokio::test] async fn stream_openai_responses_provider_emits_sse_protocol_events() { let mock_url = spawn_mock_openai_responses_stream().await; let text = collect_openai_responses_stream_body(&mock_url).await; - // The output must be Anthropic-shaped SSE events. assert!( text.contains("event: message_start"), "missing message_start event" @@ -1553,25 +1503,17 @@ mod tests { ); } - /// stream:true Responses provider forwards the translated text delta #[tokio::test] async fn stream_openai_responses_provider_forwards_translated_text_delta() { let mock_url = spawn_mock_openai_responses_stream().await; let text = collect_openai_responses_stream_body(&mock_url).await; - // Verify text content was translated. assert!( text.contains("Hi from Responses!"), "expected translated text delta in SSE output, got: {text}" ); } - // =========================================================================== - // Streaming tests: Gemini provider - // =========================================================================== - - /// Collect the full SSE response body for a Gemini streaming request against - /// the given mock URL. Shared by the focused streaming tests below. async fn collect_gemini_stream_body(mock_url: &str) -> String { let state = state_with_gemini_provider(mock_url); let app = build_router(state); @@ -1584,7 +1526,6 @@ mod tests { String::from_utf8(bytes.to_vec()).unwrap() } - /// stream:true Gemini provider returns HTTP 200 with an SSE content-type #[tokio::test] async fn stream_gemini_provider_returns_sse_content_type() { let mock_url = spawn_mock_gemini_stream().await; @@ -1607,14 +1548,11 @@ mod tests { ); } - /// stream:true Gemini provider emits Anthropic-shaped SSE protocol events - /// (message_start, message_stop) #[tokio::test] async fn stream_gemini_provider_emits_sse_protocol_events() { let mock_url = spawn_mock_gemini_stream().await; let text = collect_gemini_stream_body(&mock_url).await; - // The output must be Anthropic-shaped SSE events. assert!( text.contains("event: message_start"), "missing message_start event" @@ -1625,24 +1563,17 @@ mod tests { ); } - /// stream:true Gemini provider forwards the translated text delta #[tokio::test] async fn stream_gemini_provider_forwards_translated_text_delta() { let mock_url = spawn_mock_gemini_stream().await; let text = collect_gemini_stream_body(&mock_url).await; - // Verify text content was translated. assert!( text.contains("Hi from Gemini!"), "expected translated text delta in SSE output, got: {text}" ); } - // =========================================================================== - // Stream error behavior tests - // =========================================================================== - - /// stream_error_before_first_byte_returns_http_502 (via upstream 500 on stream request) #[tokio::test] async fn stream_error_before_first_byte_returns_http_502() { let mock_url = spawn_mock_500().await; @@ -1651,9 +1582,6 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // The upstream returns 500, which send_stream propagates as an error - // before the SSE response is committed (since the HTTP connection itself - // failed). This should map to 502 Bad Gateway. assert_eq!( resp.status(), StatusCode::BAD_GATEWAY, @@ -1661,14 +1589,6 @@ mod tests { ); } - /// A malformed SSE data frame is silently skipped by the Anthropic provider - /// adapter (`decode_frame` returns `Ok(vec![])` on bad JSON), so it does NOT - /// drive the in-band error path. This test verifies the distinct condition that - /// a VALID delta delivered BEFORE the malformed frame survives to the client, - /// and that the stream completes gracefully with `message_stop` rather than - /// panicking. The genuine in-band error path is covered separately by - /// `in_band_upstream_error_emits_anthropic_error_event` - /// (audit `duplicate-malformed-stream-tests`). #[tokio::test] async fn malformed_frame_is_skipped_and_valid_events_delivered() { let mock_url = spawn_mock_malformed_stream().await; @@ -1677,7 +1597,6 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // The HTTP status is 200 since SSE was committed. assert_eq!(resp.status(), StatusCode::OK); let bytes = axum::body::to_bytes(resp.into_body(), 1024 * 1024) @@ -1685,31 +1604,22 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Verify that the stream started with a message_start event. assert!( text.contains("event: message_start"), "stream should start with message_start" ); - // The valid delta BEFORE the malformed frame must be delivered to the - // client (it carries the marker text from spawn_mock_malformed_stream). assert!( text.contains("good-delta"), "valid events before a malformed frame must be delivered, got: {text}" ); - // The stream should complete with a message_stop (the malformed frame is - // silently skipped by the Anthropic adapter, not treated as a fatal error). assert!( text.contains("event: message_stop"), "stream should complete with message_stop even after malformed data" ); } - /// upstream disconnect: when the upstream stream ends prematurely (no - /// message_stop from upstream), the client encoder's finish() emits a - /// synthetic terminal event. The stream should complete gracefully with - /// a message_stop event. #[tokio::test] async fn upstream_disconnect_completes_with_synthetic_terminal() { let mock_url = spawn_mock_disconnect_stream().await; @@ -1725,30 +1635,17 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Stream should have started with message_start. assert!( text.contains("event: message_start"), "stream should start with message_start" ); - // The stream should complete with a message_stop event generated by - // the client encoder's finish() method, even though the upstream - // didn't send one. assert!( text.contains("event: message_stop"), "stream should have a synthetic message_stop after premature upstream end" ); } - /// A genuine in-band upstream error surfaces as an Anthropic-shaped - /// `event: error` frame to the client. - /// - /// Distinct from `malformed_frame_is_skipped_and_valid_events_delivered` - /// (which verifies a malformed JSON frame is skipped): here the mock emits a - /// real protocol-level Anthropic `event: error`, which the provider decoder - /// maps to `CoreEvent::Error` and the client encoder forwards as an - /// `event: error` SSE frame (audit `in-band-error-event-path-untested`, - /// `duplicate-malformed-stream-tests`). #[tokio::test] async fn in_band_upstream_error_emits_anthropic_error_event() { let mock_url = spawn_mock_in_band_error_stream().await; @@ -1764,43 +1661,23 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Stream should have started with message_start. assert!( text.contains("event: message_start"), "stream should start with message_start" ); - // The in-band error must surface as an Anthropic `event: error` frame. assert!( text.contains("event: error"), "in-band upstream error must emit an event: error frame, got: {text}" ); - // The error payload carries the upstream error message. assert!( text.contains("Too many requests"), "event: error should carry the upstream error message, got: {text}" ); } - // --------------------------------------------------------------------------- - // Client-disconnect / mid-stream cancellation - // --------------------------------------------------------------------------- - - /// A mock upstream body stream that yields exactly one valid Anthropic SSE - /// `message_start` frame, then parks forever (`Poll::Pending`). Its `Drop` impl - /// records cancellation so a test can observe that the proxy tore the upstream - /// down when the CLIENT disconnected mid-stream. - /// - /// This is the deterministic equivalent of a "slow-dripping" upstream: no - /// `sleep` is involved. The stream just never produces a second frame, so the - /// only way the proxy's spawned task exits after the first byte is the - /// `cancel.cancelled()` branch in `build_sse_output_stream` (audit - /// `client-disconnect-cancellation-untested`). struct SlowDripStream { - /// First SSE frame to deliver; consumed on the first poll. first: Option, - /// Set to `true` when this stream is dropped (i.e. the upstream connection - /// was torn down because the proxy cancelled its read). canceled: Arc, } @@ -1823,8 +1700,6 @@ mod tests { if let Some(chunk) = self.first.take() { std::task::Poll::Ready(Some(Ok(chunk))) } else { - // Park forever: the proxy will keep this future alive until the - // client disconnects and the CancellationToken fires. std::task::Poll::Pending } } @@ -1836,28 +1711,10 @@ mod tests { } } - /// When the client drops the response body mid-stream (after the first event), - /// the proxy must cancel the upstream read rather than leak it. - /// - /// This exercises the `cancel.cancelled()` branch wired via the response-body - /// drop-guard in `build_sse_output_stream`. The proxy reads one upstream frame - /// (crossing the first byte), the test then drops the response body, and the - /// proxy's spawned task is expected to abort the upstream stream — observed - /// here deterministically via the `SlowDripStream` Drop flag (no `sleep`-based - /// polling). A bounded `tokio::time::timeout` guards against a regression that - /// fails to propagate cancellation. - /// - /// NOTE: this asserts upstream teardown + no panic + exactly one upstream - /// request. It does NOT directly assert `Metrics::record_client_cancel()` (which - /// the cancel branch also invokes) because that field is `pub(crate)` and not - /// exposed on `AppState`; a future change that exposes a cancellation counter - /// should add that assertion here. #[tokio::test] async fn client_disconnect_cancels_upstream_stream() { use futures::StreamExt; - // Valid Anthropic message_start frame; the proxy crosses the first byte once - // it is decoded and the encoder emits the first client event. let first_frame = bytes::Bytes::from( "event: message_start\n\ data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_cancel\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-sonnet-4-6\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}\n\n" @@ -1898,27 +1755,18 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // The first byte crossed, so the SSE is committed as HTTP 200. assert_eq!(resp.status(), StatusCode::OK); - // Read exactly one chunk from the response body (the first client event), - // then DROP the body mid-stream — this is the client disconnect. let mut stream = resp.into_body().into_data_stream(); let first_chunk = stream.next().await; assert!( first_chunk.is_some(), "client should receive at least the first SSE chunk" ); - // Drop the response body future (client disconnect). The proxy's - // drop-guard must fire the CancellationToken and abort the upstream stream. drop(stream); - // The upstream read is torn down synchronously with the body drop. Bound the - // wait deterministically with a timeout (not a sleep-poll loop). let teardown = tokio::time::timeout(Duration::from_secs(2), async { while !upstream_canceled.load(Ordering::SeqCst) { - // Yield to let the spawned proxy task observe the cancel and drop - // the upstream reqwest stream (which drops SlowDripStream). tokio::task::yield_now().await; } }) @@ -1933,16 +1781,12 @@ mod tests { upstream_canceled.load(Ordering::SeqCst), "upstream SlowDripStream must have been dropped on client disconnect" ); - // The mock must have received exactly one upstream request. assert!( request_received.load(Ordering::SeqCst), "mock upstream should have received the request" ); } - /// Stream terminal event is emitted after provider decoder finish(). - /// Verifies that the finalization order follows: consume chunks -> SSE frames - /// -> decoder events -> client encoder -> finish() -> terminal events. #[tokio::test] async fn stream_terminal_event_emitted_after_decoder_finish() { let mock_url = spawn_mock_disconnect_stream().await; @@ -1958,8 +1802,6 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // Verify the finalization sequence: - // 1. Real events from upstream arrive first (message_start, content_block_start). assert!( text.contains("event: message_start"), "should have message_start from upstream" @@ -1969,8 +1811,6 @@ mod tests { "should have content_block_start from upstream" ); - // 2. The synthetic terminal events (from finish()) appear after real events. - // Find positions to verify ordering. let start_pos = text.find("event: message_start").expect("message_start"); let stop_pos = text.find("event: message_stop").expect("message_stop"); assert!( @@ -1978,7 +1818,6 @@ mod tests { "message_stop must appear after message_start" ); - // 3. message_delta (with stop_reason) must appear before message_stop. let delta_pos = text.find("event: message_delta").expect("message_delta"); assert!( delta_pos < stop_pos, @@ -1990,23 +1829,15 @@ mod tests { ); } - // =========================================================================== - // Rate limit and dedup tests - // =========================================================================== - - /// Rate-limited request returns 429 Too Many Requests. - /// - /// Uses a state configured with rpm=1 so that the second request is - /// deterministically rejected. #[tokio::test] async fn rate_limited_request_returns_429() { let mock_url = spawn_mock_anthropic_non_stream().await; - // Build state with rpm=1 so the second request is deterministically rejected. let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2043,7 +1874,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2057,7 +1892,6 @@ mod tests { ); let app = build_router(state); - // First request should be allowed (rpm=1, 1 token available). let body1 = json!({ "model": "claude-sonnet-4-6", "messages": [{ "role": "user", "content": "hello one" }], @@ -2068,15 +1902,12 @@ mod tests { .oneshot(messages_request(&body1.to_string())) .await .unwrap(); - // The first request may succeed (200) or fail for non-rate-limit reasons. - // It must NOT be 429. assert_ne!( resp1.status(), StatusCode::TOO_MANY_REQUESTS, "first request should not be rate-limited" ); - // Second request should be rate-limited (token bucket depleted). let body2 = json!({ "model": "claude-sonnet-4-6", "messages": [{ "role": "user", "content": "hello two" }], @@ -2102,20 +1933,15 @@ mod tests { assert_eq!(resp_body["error"]["type"], "rate_limit_error"); } - /// Duplicate request returns 409 Conflict. - /// - /// Uses a large dedup window (60 s) so the second request is deterministically - /// caught as a duplicate. #[tokio::test] async fn duplicate_request_returns_409() { let mock_url = spawn_mock_anthropic_non_stream().await; - // Build state with a 60-second dedup window so the duplicate is guaranteed - // to be caught, and rpm=100 so rate limiting does not interfere. let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2152,7 +1978,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_secs(60), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2168,7 +1998,6 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", false); - // First request should succeed or fail for non-dedup reasons. let resp1 = app.clone().oneshot(messages_request(&body)).await.unwrap(); assert_ne!( resp1.status(), @@ -2176,7 +2005,6 @@ mod tests { "first request should not be a duplicate" ); - // Second request with the same body and path should be deduplicated. let resp2 = app.oneshot(messages_request(&body)).await.unwrap(); assert_eq!( resp2.status(), @@ -2199,11 +2027,6 @@ mod tests { ); } - // =========================================================================== - // Token count tests - // =========================================================================== - - /// Token count with tool definitions returns a positive count (tools excluded by design). #[tokio::test] async fn token_count_with_tools_returns_positive_count() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2239,15 +2062,12 @@ mod tests { .unwrap(), ) .unwrap(); - // Token count should be positive but may be underestimated since tools - // are intentionally excluded from the estimate. assert!( resp_body["input_tokens"].as_u64().unwrap() > 0, "token count should be positive even when tools are present" ); } - /// Token count with non-text content returns a positive count (non-text excluded by design). #[tokio::test] async fn token_count_with_non_text_content_returns_positive_count() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2285,7 +2105,6 @@ mod tests { ); } - /// Token count with tool_result content returns a positive count (excluded by design). #[tokio::test] async fn token_count_with_tool_result_returns_positive_count() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2321,7 +2140,6 @@ mod tests { ); } - /// Token count endpoint works with the provider-configured application state. #[tokio::test] async fn token_count_endpoint_works_with_provider_state() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2343,11 +2161,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::OK); } - // =========================================================================== - // Error response shape tests - // =========================================================================== - - /// Invalid JSON returns 400 with Anthropic-shaped error. #[tokio::test] async fn invalid_json_returns_400_anthropic_shape() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2372,7 +2185,6 @@ mod tests { assert_eq!(body["error"]["type"], "invalid_request_error"); } - /// Provider-scoped messages routing works with the configured application state. #[tokio::test] async fn provider_messages_route_works() { let mock_url = spawn_mock_anthropic_non_stream().await; @@ -2381,13 +2193,9 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", false); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // Should succeed since we have a valid mock upstream. assert_eq!(resp.status(), StatusCode::OK); } - /// Route preserves key fields through core pipeline round-trip. - /// Sends a request with optional fields (system, temperature, max_tokens) - /// and verifies the upstream mock receives them in the encoded request. #[tokio::test] async fn route_preserves_fields_through_core() { let (mock_url, captured_body) = @@ -2410,7 +2218,6 @@ mod tests { .unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Check the upstream received the correct fields. tokio::time::sleep(Duration::from_millis(100)).await; let guard = captured_body.lock().await; let upstream_body = guard @@ -2419,7 +2226,6 @@ mod tests { let upstream_json: Value = serde_json::from_slice(upstream_body).expect("upstream body should be valid JSON"); - // The Anthropic adapter should have preserved these fields. assert_eq!(upstream_json["model"], "claude-sonnet-4-6"); assert!( upstream_json["system"].is_string() || upstream_json["system"].is_array(), @@ -2428,8 +2234,6 @@ mod tests { assert_eq!(upstream_json["max_tokens"], 128); assert_eq!(upstream_json["temperature"], 0.7); assert_eq!(upstream_json["top_p"], 0.9); - // The adapter may omit `stream` when false (default), so accept either - // `false` or absent (Null). let stream_val = &upstream_json["stream"]; assert!( stream_val.is_null() || stream_val == &json!(false), @@ -2437,18 +2241,13 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Provider name validation tests (findings 103, 104) - // --------------------------------------------------------------------------- - - /// Simple state helper for provider name validation tests. Does not need - /// a live mock server since the tests exercise the validation layer only. fn state_for_validation_tests() -> AppState { let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2485,7 +2284,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2521,7 +2324,6 @@ mod tests { .unwrap(), ) .unwrap(); - // Anthropic-shaped error. assert_eq!(resp_body["type"], "error"); assert_eq!(resp_body["error"]["type"], "not_found_error"); } @@ -2585,14 +2387,9 @@ mod tests { let req = Request::builder() .method("POST") .uri("/providers/mock-provider/v1/messages") - // Intentionally omit content-type header. .body(Body::from(body.to_string())) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // Without content-type, the request may still be processed since - // axum Bytes extractor does not require content-type. The request - // should either succeed (and fail at the upstream call) or return - // an error status -- either way it should not panic. assert!( resp.status().is_client_error() || resp.status().is_server_error() @@ -2600,16 +2397,6 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Error/status path event-log ordering (audit LOW, eventlog-03 / eventlog-04) - // --------------------------------------------------------------------------- - - /// A streaming request that fails at the upstream (500) must emit - /// `RequestReceived` BEFORE `ResponseFailed`. Before the eventlog-03 fix the - /// streaming `RequestReceived` was emitted after `resolve_target`/`send_stream`, - /// so an early-return produced an orphaned `ResponseFailed` with no preceding - /// `RequestReceived`. The RequestReceived is now emitted before any upstream - /// call, mirroring the non-stream path. #[tokio::test] async fn stream_emits_request_received_before_response_failed_on_upstream_failure() { use std::sync::Arc; @@ -2624,7 +2411,6 @@ mod tests { let body = make_messages_body("claude-sonnet-4-6", true); let resp = app.oneshot(messages_request(&body)).await.unwrap(); - // Drain the body so the spawned task runs to completion. axum::body::to_bytes(resp.into_body(), 1024 * 1024) .await .expect("draining the error body should succeed"); @@ -2650,10 +2436,6 @@ mod tests { ); } - /// A non-stream upstream 500 emits `ResponseFailed` (not `ResponseCompleted`) - /// and is counted as a failure -- never as a success. This guards the - /// eventlog-04 invariant that a failure path never records success, and that - /// the cost/ResponseCompleted path is skipped on failure. #[tokio::test] async fn non_stream_upstream_500_emits_response_failed_not_completed() { use std::sync::Arc; @@ -2687,24 +2469,13 @@ mod tests { assert_eq!(failed.http_status, 502); } - // =========================================================================== - // stream-ids-01: leading CoreEvent::Error on the OpenAI streaming path - // =========================================================================== - - /// A provider whose `chat_completions` route is served by an Anthropic-speaking - /// adapter. The client hits `/v1/chat/completions` (so `client_protocol` is - /// `OpenAiChat`), while the upstream speaks the Anthropic Messages protocol. - /// This is the only configuration where a leading upstream `error` event can - /// surface as a leading `CoreEvent::Error` on the OpenAI client path: the - /// Anthropic adapter decodes the upstream `error` SSE frame into - /// `CoreEvent::Error`, and the OpenAI client encoder drops it (returns `Err`), - /// which is exactly the case the stream-ids-01 synthesis handles. fn state_with_openai_chat_client_anthropic_upstream(mock_endpoint: &str) -> AppState { let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2741,7 +2512,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, @@ -2757,9 +2532,6 @@ mod tests { ) } - /// Spawn a mock upstream whose FIRST (and only) SSE frame is an Anthropic - /// `error` event. No preceding `message_start` is emitted, so the first decoded - /// `CoreEvent` is `CoreEvent::Error`. async fn spawn_mock_anthropic_leading_error_stream() -> String { let app = Router::new().route( "/{*path}", @@ -2784,7 +2556,6 @@ mod tests { format!("http://{addr}/v1/messages") } - /// Build a minimal OpenAI Chat Completions streaming request body. fn make_chat_completions_stream_body(model: &str) -> String { json!({ "model": model, @@ -2794,7 +2565,6 @@ mod tests { .to_string() } - /// Build a POST request to the OpenAI Chat Completions client route. fn chat_completions_stream_request(body: &str) -> Request { Request::builder() .method("POST") @@ -2804,12 +2574,6 @@ mod tests { .unwrap() } - /// `stream-ids-01`: when the FIRST decoded upstream event on an OpenAI Chat - /// STREAMING request (`ClientProtocol::OpenAiChat`) is an error -- with no - /// preceding valid event -- the client must receive HTTP 200 with an in-band - /// SSE error chunk, NOT a 502 "empty stream" body. The replay loop synthesizes - /// the OpenAI error-object `data:` line plus the `data: [DONE]` terminator so - /// the error crosses the first-byte boundary as a real in-band event. #[tokio::test] async fn stream_openai_chat_leading_error_returns_200_in_band_sse_error() { let mock_url = spawn_mock_anthropic_leading_error_stream().await; @@ -2822,9 +2586,6 @@ mod tests { .await .unwrap(); - // The leading error must cross the first-byte boundary as an in-band SSE - // event, committing HTTP 200 -- not surface as a pre-stream 502 "empty - // stream" body. assert_eq!( resp.status(), StatusCode::OK, @@ -2847,9 +2608,6 @@ mod tests { .unwrap(); let text = String::from_utf8(bytes.to_vec()).unwrap(); - // The in-band error must be an OpenAI-shaped error object on a `data:` line. - // `overloaded_error` maps to `CoreStreamErrorKind::Upstream`, which the OpenAI - // payload renders as `type: "server_error"`. assert!( text.contains("\"error\"") && text.contains("\"type\""), "in-band SSE error must carry an OpenAI error object; got: {text}" @@ -2863,13 +2621,11 @@ mod tests { "the sanitized error message must be surfaced in-band; got: {text}" ); - // The OpenAI `[DONE]` terminator must follow the error chunk. assert!( text.contains("[DONE]"), "the OpenAI [DONE] terminator must follow the in-band error; got: {text}" ); - // It must NOT be the pre-stream "empty stream" 502 JSON body. assert!( !text.contains("empty stream"), "must not surface the misleading empty-stream 502 body; got: {text}" diff --git a/crates/llm-proxy-server/tests/integration.rs b/crates/llm-proxy-server/tests/integration.rs index 9f1ea4c..ceee502 100644 --- a/crates/llm-proxy-server/tests/integration.rs +++ b/crates/llm-proxy-server/tests/integration.rs @@ -21,7 +21,6 @@ mod tests { use serde_json::{Value, json}; use tower::ServiceExt; - /// Build AppState in TOML mode for integration testing. fn state() -> AppState { let app_config = AppConfig { server: ServerConfig { @@ -36,7 +35,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( @@ -53,16 +56,13 @@ mod tests { ) } - /// Build AppState with a single "mock-provider" registered so that provider-based - /// routing can resolve. The provider has an `anthropic_messages` adapter with a - /// dummy endpoint. This is used by tests that send valid JSON past the parsing - /// layer and need the request to reach the core pipeline. fn state_with_provider() -> AppState { let provider = ProviderConfig { name: "mock-provider".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -99,7 +99,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -171,16 +175,9 @@ mod tests { assert_eq!(body["git_sha"], "test"); } - /// Messages route accepts requests without x-api-key header. - /// - /// The proxy does not enforce authentication itself -- it forwards the - /// configured provider API key to the upstream. The upstream is unreachable - /// (dummy endpoint), so the response is a 502 Bad Gateway, not a 404 (route - /// is registered) or a 401 (the proxy does not check auth headers). #[tokio::test] async fn messages_without_auth_header_reaches_upstream() { let app = build_router(state_with_provider()); - // Minimal Anthropic-shaped request without x-api-key. let body = json!({ "model": "claude-sonnet-4-6", "messages": [{ "role": "user", "content": "hi" }], @@ -193,15 +190,9 @@ mod tests { .body(Body::from(body.to_string())) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // The route is registered (not 404) and the proxy forwards the request - // using its configured API key. The upstream is unreachable, so we get - // 502 Bad Gateway (or 500 for a routing error). let status = resp.status(); assert_ne!(status, StatusCode::NOT_FOUND, "route must be registered"); - let valid_statuses = [ - StatusCode::BAD_GATEWAY, // 502 - upstream unreachable - StatusCode::INTERNAL_SERVER_ERROR, // 500 - routing error - ]; + let valid_statuses = [StatusCode::BAD_GATEWAY, StatusCode::INTERNAL_SERVER_ERROR]; assert!( valid_statuses.contains(&status), "expected downstream error for unreachable upstream without auth header, got {status}" @@ -219,18 +210,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::NOT_FOUND); } - // -- Route guardrails (Phase 0 characterization) --------------------------- - - /// Phase 0 guardrail: POST /v1/messages with valid Anthropic JSON must parse - /// and validate the request shape. This test verifies that the route is - /// registered and the request body parses successfully. Any failure must be - /// an upstream or auth error -- never a bad-JSON or request-shape error. - /// - /// Note: The assertion on `error_type != "invalid_request_error"` when status - /// is 400 is a best-effort guard, not a guarantee. The Anthropic API itself - /// returns `invalid_request_error` for valid JSON with semantic issues. The - /// test specifically checks that the request parses and routes without a - /// JSON/shape error, not that no 400 can ever occur. #[tokio::test] async fn messages_valid_json_parses_and_validates() { let app = build_router(state_with_provider()); @@ -247,7 +226,6 @@ mod tests { .unwrap(); let resp = app.oneshot(req).await.unwrap(); let status = resp.status(); - // Parse response body to verify it is valid JSON. let _resp_body: Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 64 * 1024) .await @@ -255,21 +233,16 @@ mod tests { ) .unwrap(); - // Route must be registered (not 404). assert_ne!(status, StatusCode::NOT_FOUND, "route must be registered"); - // The request body is valid Anthropic JSON, so it must parse and route - // successfully. With an empty routing table, the model is unknown so - // 400 is returned. With a real routing table, downstream errors are: - // 502 (upstream), 500 (routing), 401 (auth), 429 (rate limit), 409 (duplicate). let valid_statuses = [ - StatusCode::BAD_REQUEST, // 400 - unknown model (empty routing table) - StatusCode::BAD_GATEWAY, // 502 - upstream failure - StatusCode::INTERNAL_SERVER_ERROR, // 500 - routing error - StatusCode::UNAUTHORIZED, // 401 - auth failure - StatusCode::TOO_MANY_REQUESTS, // 429 - rate limit - StatusCode::CONFLICT, // 409 - duplicate - StatusCode::OK, // 200 - success (unlikely without real upstream) + StatusCode::BAD_REQUEST, + StatusCode::BAD_GATEWAY, + StatusCode::INTERNAL_SERVER_ERROR, + StatusCode::UNAUTHORIZED, + StatusCode::TOO_MANY_REQUESTS, + StatusCode::CONFLICT, + StatusCode::OK, ]; assert!( valid_statuses.contains(&status), @@ -277,8 +250,6 @@ mod tests { ); } - /// Phase 0 guardrail: POST /v1/messages with invalid (non-JSON) body returns - /// 400 with invalid_request_error (bad JSON parse). #[tokio::test] async fn messages_invalid_json_returns_bad_request() { let app = build_router(state_with_provider()); @@ -306,8 +277,6 @@ mod tests { ); } - /// Phase 0 guardrail: POST /v1/messages with valid JSON but missing required - /// fields (no messages) returns 400 with invalid_request_error. #[tokio::test] async fn messages_valid_json_missing_fields_returns_validation_error() { let app = build_router(state_with_provider()); @@ -330,8 +299,6 @@ mod tests { ) .unwrap(); assert_eq!(resp_body["error"]["type"], "invalid_request_error"); - // The missing field triggers a serde deserialization error, which is - // reported as "invalid JSON: missing field `messages`". assert!( resp_body["error"]["message"] .as_str() @@ -341,11 +308,6 @@ mod tests { ); } - /// Phase 9 guardrail: POST /v1/chat/completions route is mounted and tested - /// comprehensively in tests/chat_completions.rs. The removed chat_echo test that - /// was here has been removed since its single assertion (not-404) is a subset - /// of the `route_is_mounted_no_longer_404` test in that file. - #[tokio::test] async fn count_tokens_returns_estimate() { let app = build_router(state_with_provider()); @@ -368,14 +330,9 @@ mod tests { .unwrap(), ) .unwrap(); - // Heuristic counter should return a positive token count. assert!(body["input_tokens"].as_u64().unwrap() > 0); } - // -- Phase 0: Additional characterization tests ----------------------------- - - /// Phase 0: not_found returns a JSON error envelope matching the Anthropic - /// error format, not plain text. #[tokio::test] async fn not_found_returns_json_error_envelope() { let app = build_router(state()); @@ -395,8 +352,6 @@ mod tests { assert_eq!(body["error"]["type"], "not_found_error"); } - /// Phase 0: POST /v1/messages with wrong types in JSON fields (e.g. max_tokens - /// as a string) returns 400 with invalid_request_error. #[tokio::test] async fn messages_wrong_field_type_returns_bad_request() { let app = build_router(state_with_provider()); @@ -422,8 +377,6 @@ mod tests { assert_eq!(resp_body["error"]["type"], "invalid_request_error"); } - /// Phase 0: POST /v1/messages with unknown fields returns 400 since - /// MessageRequest now uses deny_unknown_fields. #[tokio::test] async fn messages_unknown_fields_returns_bad_request() { let app = build_router(state_with_provider()); @@ -450,12 +403,9 @@ mod tests { assert_eq!(resp_body["error"]["type"], "invalid_request_error"); } - /// Phase 0: POST /v1/messages with invalid UTF-8 bytes returns 400. - /// The server must not panic when receiving binary/non-UTF-8 body data. #[tokio::test] async fn messages_invalid_utf8_returns_bad_request() { let app = build_router(state()); - // Raw non-UTF-8 bytes. let req = Request::builder() .method("POST") .uri("/providers/mock-provider/v1/messages") @@ -463,7 +413,6 @@ mod tests { .body(Body::from(vec![0x80, 0x81, 0x82])) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // The server should return a 400-level error, not panic (5xx). assert!( resp.status().is_client_error(), "expected 4xx for invalid UTF-8, got {}", @@ -471,7 +420,6 @@ mod tests { ); } - /// Phase 0: POST /v1/messages with empty body (0 bytes) returns 400. #[tokio::test] async fn messages_empty_body_returns_bad_request() { let app = build_router(state_with_provider()); @@ -492,11 +440,9 @@ mod tests { assert_eq!(body["error"]["type"], "invalid_request_error"); } - /// Phase 0: POST /v1/messages with oversized body (> 32 MiB) returns 413. #[tokio::test] async fn messages_oversized_body_returns_payload_too_large() { let app = build_router(state()); - // 33 MiB body (exceeds the 32 MiB limit). let oversized_body = "X".repeat(33 * 1024 * 1024); let req = Request::builder() .method("POST") @@ -507,9 +453,6 @@ mod tests { let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); - // MEDIUM-7: the body-limit 413 is normalised into a protocol-shaped JSON - // envelope (not axum's default plain-text rejection) and carries an - // `x-request-id` header, matching every handler-produced error. assert_eq!( resp.headers().get("content-type").unwrap(), "application/json" @@ -530,9 +473,6 @@ mod tests { ); } - // -- TOML mode integration tests ------------------------------------------- - - /// TOML mode: /health returns ok with server name from AppConfig. #[tokio::test] async fn toml_health_returns_ok() { let app = build_router(state()); @@ -549,11 +489,9 @@ mod tests { ) .unwrap(); assert_eq!(body["status"], "ok"); - // server_name comes from AppConfig. assert_eq!(body["service"], "test-proxy"); } - /// TOML mode: /version returns build info with AppConfig server name. #[tokio::test] async fn toml_version_returns_app_config_server_name() { let app = build_router(state()); @@ -575,7 +513,6 @@ mod tests { assert_eq!(body["git_sha"], "test"); } - /// TOML mode: /ready returns ready. #[tokio::test] async fn toml_ready_returns_ready() { let app = build_router(state()); @@ -594,10 +531,6 @@ mod tests { assert_eq!(body["status"], "ready"); } - /// TOML mode: POST /providers/{provider}/v1/messages with a registered provider passes the - /// request through to the upstream adapter. The model name is forwarded as-is, so - /// an unknown model results in a downstream error (502) rather than a 400, because - /// the proxy no longer validates model names locally. #[tokio::test] async fn toml_messages_passes_through_to_upstream() { let app = build_router(state_with_provider()); @@ -613,27 +546,19 @@ mod tests { .body(Body::from(body.to_string())) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // With provider-based routing, the model is passed through to the upstream - // adapter. The dummy endpoint refuses connections, so we get a 502 Bad Gateway - // (or 500 internal error), not a 400 for unknown model. let status = resp.status(); assert_ne!( status, StatusCode::NOT_FOUND, "provider is registered, route must resolve" ); - // The upstream is unreachable so downstream error is expected. - let valid_statuses = [ - StatusCode::BAD_GATEWAY, // 502 - upstream failure - StatusCode::INTERNAL_SERVER_ERROR, // 500 - routing error - ]; + let valid_statuses = [StatusCode::BAD_GATEWAY, StatusCode::INTERNAL_SERVER_ERROR]; assert!( valid_statuses.contains(&status), "expected downstream error for unreachable upstream, got {status}" ); } - /// POST /providers/{provider}/v1/messages/count_tokens works with provider state. #[tokio::test] async fn toml_count_tokens_returns_estimate() { let app = build_router(state_with_provider()); @@ -659,12 +584,6 @@ mod tests { assert!(resp_body["input_tokens"].as_u64().unwrap() > 0); } - // --------------------------------------------------------------------------- - // Additional integration tests for edge cases (findings 113, 114, 115, 117, 119) - // --------------------------------------------------------------------------- - - /// GET on POST-only messages route returns a protocol-shaped 405 (finding 115, - /// audit LOW-27): Anthropic envelope body, not axum's empty 405. #[tokio::test] async fn get_on_messages_route_returns_405() { let app = build_router(state_with_provider()); @@ -686,8 +605,6 @@ mod tests { assert_eq!(body["error"]["message"], "method not allowed"); } - /// GET on POST-only chat completions route returns a protocol-shaped 405 - /// (finding 115, audit LOW-27): OpenAI envelope body. #[tokio::test] async fn get_on_chat_completions_route_returns_405() { let app = build_router(state_with_provider()); @@ -706,7 +623,6 @@ mod tests { assert!(body["error"]["code"].is_null()); } - /// POST on GET-only health route returns 405 (finding 115). #[tokio::test] async fn post_on_health_route_returns_405() { let app = build_router(state()); @@ -719,7 +635,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED); } - /// Invalid provider name on messages route returns 400 (finding 113). #[tokio::test] async fn invalid_provider_name_on_messages_returns_400() { let app = build_router(state()); @@ -735,8 +650,6 @@ mod tests { .body(Body::from(body.to_string())) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // axum URL-decodes the path, so "../.." becomes the actual path segment. - // Either way it should be rejected. assert!( resp.status() == StatusCode::BAD_REQUEST || resp.status() == StatusCode::NOT_FOUND, "expected 400 or 404 for path traversal provider name, got {}", @@ -744,7 +657,6 @@ mod tests { ); } - /// count_tokens with invalid JSON returns 400 (finding 114). #[tokio::test] async fn count_tokens_invalid_json_returns_400() { let app = build_router(state_with_provider()); @@ -758,7 +670,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::BAD_REQUEST); } - /// count_tokens with empty body returns 400 (finding 114). #[tokio::test] async fn count_tokens_empty_body_returns_400() { let app = build_router(state_with_provider()); @@ -772,7 +683,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::BAD_REQUEST); } - /// count_tokens with unknown provider returns 404 (finding 114). #[tokio::test] async fn count_tokens_unknown_provider_returns_404() { let app = build_router(state()); @@ -791,10 +701,6 @@ mod tests { assert_eq!(resp.status(), StatusCode::NOT_FOUND); } - /// Health endpoint never exposes operational metrics, even when the caller opts - /// in via the legacy `?metrics=true` query (LOW-13): the route is - /// unauthenticated, so the query parameter is now ignored and the body is - /// always the trivial `{status, service}` liveness probe. #[tokio::test] async fn health_never_exposes_metrics_even_with_opt_in_query() { let app = build_router(state()); @@ -815,13 +721,10 @@ mod tests { "/health must not include 'metrics' even with ?metrics=true \ (unauthenticated endpoint, LOW-13); got: {body}" ); - // The trivial body still carries status + service. assert_eq!(body["status"], "ok"); assert!(body["service"].is_string()); } - /// Health endpoint omits operational metrics by default (LOW-13): the route is - /// unauthenticated, so a bare probe must not leak operational counters. #[tokio::test] async fn health_omits_metrics_by_default() { let app = build_router(state()); @@ -841,12 +744,10 @@ mod tests { body.get("metrics").is_none(), "bare /health must not include 'metrics' (unauthenticated endpoint); got: {body}" ); - // The minimal body still carries status + service. assert_eq!(body["status"], "ok"); assert!(body["service"].is_string()); } - /// count_tokens response has correct JSON shape (finding 119). #[tokio::test] async fn count_tokens_response_has_correct_shape() { let app = build_router(state_with_provider()); @@ -863,7 +764,6 @@ mod tests { .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Verify Content-Type is application/json. let ct = resp .headers() .get("content-type") @@ -877,19 +777,12 @@ mod tests { .unwrap(); assert!(resp_body["input_tokens"].is_number()); assert!(resp_body["input_tokens"].as_u64().unwrap() > 0); - // Verify only the expected field is present. assert!( resp_body.as_object().unwrap().len() == 1, "count_tokens response should only have input_tokens field" ); } - // --------------------------------------------------------------------------- - // CORS (audit LOW-12) - // --------------------------------------------------------------------------- - - /// Build AppState whose `server.allowed_origins` is set to `origins` (audit - /// LOW-12). Mirrors [`state`] but installs an explicit cross-origin allow list. fn state_with_origins(origins: Vec) -> AppState { let app_config = AppConfig { server: ServerConfig { @@ -904,7 +797,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( @@ -921,8 +818,6 @@ mod tests { ) } - /// A configured origin is echoed on a CORS preflight, and the permitted - /// method/headers are advertised (audit LOW-12). #[tokio::test] async fn cors_preflight_echoes_configured_origin() { let app = build_router(state_with_origins(vec![ @@ -944,7 +839,6 @@ mod tests { .expect("allow-origin echoed on preflight"), "https://app.example.com" ); - // The advertised methods must be the curated set, not "*" (LOW-12). let allow_methods = resp .headers() .get("access-control-allow-methods") @@ -956,13 +850,8 @@ mod tests { assert!(!allow_methods.contains('*')); } - /// With no `allowed_origins` configured, no CORS layer is installed and a - /// browser preflight is not answered with an allow-origin echo (audit LOW-12). - /// This is the regression guard: the old `CorsLayer::very_permissive()` admitted - /// every origin by default; the new default must not. #[tokio::test] async fn cors_absent_when_origins_unset() { - // `state()` builds with `allowed_origins: None`. let app = build_router(state()); let req = Request::builder() .method("OPTIONS") @@ -979,32 +868,14 @@ mod tests { ); } - // --------------------------------------------------------------------------- - // Event-log coverage for early-gate failures (audit eventlog-early-gates) - // --------------------------------------------------------------------------- - // - // Each inference handler emits ResponseFailed exactly once per failed request: - // early-gate failures (unknown provider, bad content-type, rate-limit, conflict, - // malformed JSON, decode failure) emit inside the inner function, and failures - // INSIDE the core pipeline emit inside the pipeline (which also returns Err). - // The outer wrapper renders only -- it must NOT re-emit, or a pipeline failure - // would be double-counted. These tests wire a RecordingBus into AppState and - // assert the event counts so a regression (double-emit, or a dropped early-gate) - // is caught (audit route-responsefailed-gaps). - use llm_proxy_storage::{EventBus, ProxyEvent, RecordingBus}; - /// Build a provider-backed AppState whose event bus is a [`RecordingBus`], - /// returning both the state and the bus so the test can snapshot emitted - /// events. Mirrors [`state_with_provider`] but injects the recording sink. fn state_with_provider_and_bus() -> (AppState, Arc) { let bus = Arc::new(RecordingBus::new()); let state = state_with_provider().with_event_bus(Arc::clone(&bus) as Arc); (state, bus) } - /// Count recorded events of a given variant by name ("request_received", - /// "response_completed", "response_failed"). fn count_events(bus: &RecordingBus, kind: &str) -> usize { bus.snapshot() .iter() @@ -1019,8 +890,6 @@ mod tests { .count() } - /// The single ResponseFailed emitted for a malformed-JSON early gate carries - /// http_status 400 and is emitted exactly once (audit eventlog-early-gates). #[tokio::test] async fn malformed_json_emits_one_response_failed_with_400() { let (state, bus) = state_with_provider_and_bus(); @@ -1039,7 +908,6 @@ mod tests { 1, "malformed JSON must emit exactly one ResponseFailed" ); - // The emitted event must carry the rendered HTTP status (400). let failed = bus .snapshot() .into_iter() @@ -1052,9 +920,6 @@ mod tests { assert_eq!(failed.http_status, 400); } - /// An unknown-provider request emits exactly ONE ResponseFailed (not two) after - /// the centralized wrapper replaced the per-site UnknownProvider emit - /// (audit eventlog-early-gates regression guard). #[tokio::test] async fn unknown_provider_emits_exactly_one_response_failed() { let (state, bus) = state_with_provider_and_bus(); @@ -1080,9 +945,6 @@ mod tests { ); } - /// A duplicate request within the dedup window hits the Conflict early gate and - /// emits exactly one ResponseFailed carrying http_status 409. This covers the - /// RateLimited/Conflict branch of the centralized emit (audit eventlog-early-gates). #[tokio::test] async fn duplicate_request_emits_one_response_failed_with_409() { let (state, bus) = state_with_provider_and_bus(); @@ -1100,20 +962,10 @@ mod tests { .body(Body::from(body.to_string())) .unwrap() }; - // First request consumes the dedup slot (it will fail downstream at 502 and - // emit its own pipeline-level ResponseFailed; that is expected and distinct - // from the Conflict gate emission under test). let _ = app.clone().oneshot(make_req()).await.unwrap(); - // Second identical request within the dedup window is rejected as a Conflict - // at the prepare_request early gate, before any upstream call. That early - // gate emits exactly one ResponseFailed for THIS request. let resp = app.oneshot(make_req()).await.unwrap(); assert_eq!(resp.status(), StatusCode::CONFLICT); - // Exactly one ResponseFailed carries the Conflict's 409 status: the - // early-gate emit fired once for the dedup rejection. (The first request's - // downstream failure emits a 502 ResponseFailed, which is a separate, - // expected event and is excluded by the status filter.) let conflicts: Vec<_> = bus .snapshot() .into_iter() @@ -1128,11 +980,6 @@ mod tests { "dedup Conflict must emit exactly one 409 ResponseFailed" ); - // Direct per-request count: the first request's downstream failure must also - // have emitted exactly ONE ResponseFailed (no double-emit from a pipeline - // failure re-emitting at the wrapper). Both requests share the bus, so - // confirm the total ResponseFailed count is exactly 2 (one per request), - // not more (audit route-responsefailed-gaps regression). assert_eq!( count_events(&bus, "response_failed"), 2, @@ -1140,9 +987,6 @@ mod tests { ); } - /// token_count success emits RequestReceived + ResponseCompleted (token_count - /// previously emitted NO events; it now logs like any other route) - /// (audit eventlog-early-gates). #[tokio::test] async fn token_count_success_emits_request_received_and_completed() { let (state, bus) = state_with_provider_and_bus(); @@ -1178,8 +1022,6 @@ mod tests { ); } - /// token_count failure emits exactly one ResponseFailed (token_count previously - /// emitted NO events on failure) (audit eventlog-early-gates). #[tokio::test] async fn token_count_failure_emits_one_response_failed() { let (state, bus) = state_with_provider_and_bus(); @@ -1200,9 +1042,6 @@ mod tests { ); } - /// /models emits exactly one ResponseFailed on an unknown-provider error, for - /// consistency with the rest of the API (audit eventlog-early-gates). /models is - /// a listing endpoint so it carries no RequestReceived/ResponseCompleted pair. #[tokio::test] async fn models_unknown_provider_emits_one_response_failed() { let (state, bus) = state_with_provider_and_bus(); @@ -1222,12 +1061,6 @@ mod tests { ); } - /// A failure INSIDE the core pipeline (unreachable upstream -> 502) emits - /// exactly ONE ResponseFailed for the request, not two. The core pipeline emits - /// its own ResponseFailed on the upstream failure AND returns Err; the route - /// handler must NOT re-emit at the wrapper (regression guard for - /// route-responsefailed-gaps). Only one request hits this bus, so the count is a - /// direct per-request assertion. #[tokio::test] async fn upstream_failure_inside_pipeline_emits_one_response_failed() { let (state, bus) = state_with_provider_and_bus(); @@ -1244,9 +1077,6 @@ mod tests { .body(Body::from(body.to_string())) .unwrap(); let resp = app.oneshot(req).await.unwrap(); - // The mock provider points at an unreachable endpoint, so the request - // fails INSIDE the core pipeline (upstream send) and the client gets a 502 - // (or 500 routing error), proving the failure occurred past the early gates. let status = resp.status(); let valid_statuses = [StatusCode::BAD_GATEWAY, StatusCode::INTERNAL_SERVER_ERROR]; assert!( @@ -1261,9 +1091,6 @@ mod tests { ); } - /// Same regression guard for the OpenAI chat-completions route: an - /// inside-pipeline failure emits exactly one ResponseFailed - /// (audit route-responsefailed-gaps). #[tokio::test] async fn upstream_failure_inside_pipeline_chat_emits_one_response_failed() { let (state, bus) = state_with_provider_and_bus(); diff --git a/crates/llm-proxy-server/tests/openrouter_e2e.rs b/crates/llm-proxy-server/tests/openrouter_e2e.rs index c3b9e09..61b8040 100644 --- a/crates/llm-proxy-server/tests/openrouter_e2e.rs +++ b/crates/llm-proxy-server/tests/openrouter_e2e.rs @@ -1,34 +1,6 @@ -//! End-to-end verification of the proxy's *goals* against faithful OpenRouter -//! response shapes. -//! -//! OpenRouter () speaks the OpenAI Chat Completions -//! wire protocol, but its real responses carry OpenRouter-specific extras that a -//! vanilla OpenAI fixture does not: `id`s prefixed `gen-…`, a top-level -//! `provider` routing field, a `cost` member inside `usage`, `system_fingerprint`, -//! `native_finish_reason`, `logprobs`, and its own `{"error":{"code","message", -//! "metadata"}}` envelope on 402/429. This suite drives the *real* in-process -//! proxy pipeline (`build_router` + `AppState`) against a mock upstream that -//! returns those exact shapes, so the assertions validate behaviour against -//! reality rather than hand-tidied fixtures. -//! -//! Each test is mapped to a documented goal of the proxy: -//! -//! | Goal | Test | -//! |---------------------------------------|-------------------------------------------------------| -//! | Protocol hub: OpenAI client passthrough | `openai_client_*_translates_to_openai_shape` | -//! | Cross-protocol translation (headline) | `anthropic_client_*_translates_to_anthropic_shape` | -//! | Tool-call fidelity (OpenAI client) | `tool_calls_preserved_for_openai_client` | -//! | Tool-call fidelity (cross-protocol) | `tool_calls_translate_to_anthropic_tool_use_blocks` | -//! | Robustness to provider-specific extras | `openrouter_specific_extras_do_not_break_conversion`| -//! | Streaming fidelity (OpenAI client) | `openai_client_stream_ends_with_done` | -//! | Streaming fidelity (cross-protocol) | `anthropic_client_stream_cross_protocol` | -//! | OpenRouter header injection | `openrouter_app_headers_forwarded_upstream` | -//! | Error fidelity (429 passthrough) | `openrouter_429_surfaces_as_client_shaped_error` | -//! | Live network validation (opt-in) | `live_openrouter_round_trip` (`#[ignore]`) | -//! -//! The first nine run offline and deterministically in CI. The last is -//! `#[ignore]`-gated on `OPENROUTER_API_KEY` so the real path can be re-validated -//! on demand without key material or network in CI. +//! End-to-end verification against faithful OpenRouter response shapes +//! (`gen-` ids, top-level `provider`, `usage.cost`, `native_finish_reason`, +//! its own error envelope) through the real in-process proxy pipeline. #[cfg(test)] mod tests { @@ -54,11 +26,6 @@ mod tests { use serde_json::{Value, json}; use tower::ServiceExt; - // =========================================================================== - // Mock-server plumbing (mirrors the established harness in core_pipeline.rs) - // =========================================================================== - - /// Poll the mock server's TCP port until it accepts a connection or 5 s elapse. async fn wait_for_ready(addr: std::net::SocketAddr) { let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { @@ -72,8 +39,6 @@ mod tests { } } - /// Spawn a mock axum server, surfacing a serve-task panic loudly (see the audit - /// note on `mock-server-joinhandle-swallow` in core_pipeline.rs). fn spawn_mock_serve(listener: tokio::net::TcpListener, app: Router) { tokio::spawn(async move { let result = std::panic::AssertUnwindSafe(axum::serve(listener, app).into_future()) @@ -95,12 +60,6 @@ mod tests { }); } - /// A canned non-streaming OpenRouter Chat Completions response. - /// - /// Deliberately rich in OpenRouter-specific extras (`gen-` id, top-level - /// `provider`, `cost` inside `usage`, `system_fingerprint`, - /// `native_finish_reason`, `logprobs`, `refusal`) so that the "robustness to - /// provider variance" goal is exercised by every test that consumes it. fn openrouter_non_stream_response() -> Vec { serde_json::to_vec(&json!({ "id": "gen-1719000000-AbCdEf", @@ -130,7 +89,6 @@ mod tests { .unwrap() } - /// A canned non-streaming OpenRouter tool-call response. fn openrouter_tool_call_response() -> Vec { serde_json::to_vec(&json!({ "id": "gen-toolcall-001", @@ -165,7 +123,6 @@ mod tests { .unwrap() } - /// A canned OpenRouter rate-limit (429) error envelope. fn openrouter_429_body() -> Vec { serde_json::to_vec(&json!({ "error": { @@ -180,7 +137,6 @@ mod tests { .unwrap() } - /// Spawn a mock returning a canned JSON body for any POST path. async fn spawn_mock(body: Vec) -> String { let app = Router::new().route( "/{*path}", @@ -195,7 +151,6 @@ mod tests { spawn_and_url(app).await } - /// Spawn a mock returning a canned status + body (for error envelopes). async fn spawn_mock_status(status: StatusCode, body: Vec) -> String { let app = Router::new().route( "/{*path}", @@ -210,8 +165,6 @@ mod tests { spawn_and_url(app).await } - /// Spawn a mock returning OpenRouter-style streaming chunks (with `provider`, - /// `native_finish_reason`, and a final usage chunk carrying `cost`). async fn spawn_mock_openrouter_stream() -> String { let app = Router::new().route( "/{*path}", @@ -235,9 +188,6 @@ mod tests { spawn_and_url(app).await } - /// Spawn a mock that captures the request headers, returning the canned body. - /// Used to prove OpenRouter app headers (`HTTP-Referer`, `X-Title`) reach the - /// upstream. async fn spawn_mock_with_header_capture( body: Vec, ) -> (String, Arc>>) { @@ -266,9 +216,6 @@ mod tests { (url, captured) } - /// Spawn a mock that captures the request body the proxy forwards upstream, - /// returning the canned response. Used to prove the request encode path drops - /// no input content. async fn spawn_mock_with_body_capture( body: Vec, ) -> (String, Arc>>) { @@ -297,8 +244,6 @@ mod tests { (url, captured) } - /// Bind a catch-all mock on an ephemeral port, spawn it, wait for readiness, - /// and return its base URL. async fn spawn_and_url(app: Router) -> String { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -307,21 +252,13 @@ mod tests { format!("http://{addr}/v1/chat/completions") } - // =========================================================================== - // AppState + request builders - // =========================================================================== - - /// Build an AppState whose single provider `openrouter` speaks the OpenAI Chat - /// Completions protocol at `mock_endpoint`. Both client routes - /// (`chat_completions` and `messages`) resolve to that one adapter, which is - /// what enables cross-protocol translation: an Anthropic Messages client and an - /// OpenAI Chat client both target the same OpenAI/OpenRouter upstream. fn openrouter_state(mock_endpoint: &str, extra_headers: HashMap) -> AppState { let provider = ProviderConfig { name: "openrouter".to_owned(), api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -334,7 +271,6 @@ mod tests { ); m }, - // An OpenAI-protocol provider serves both client protocols. routes: llm_proxy_core::ProviderRoutesConfig { chat_completions: Some("chat".to_owned()), messages: Some("chat".to_owned()), @@ -359,7 +295,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -413,7 +353,6 @@ mod tests { .to_string() } - /// Read a non-streaming JSON response body. async fn read_json(resp: axum::response::Response) -> Value { let bytes = axum::body::to_bytes(resp.into_body(), 64 * 1024) .await @@ -421,7 +360,6 @@ mod tests { serde_json::from_slice(&bytes).unwrap() } - /// Collect a streaming response body as a UTF-8 string. async fn read_stream_text(resp: axum::response::Response) -> String { let bytes = axum::body::to_bytes(resp.into_body(), 1024 * 1024) .await @@ -429,13 +367,6 @@ mod tests { String::from_utf8(bytes.to_vec()).unwrap() } - // =========================================================================== - // Goal: protocol hub — OpenAI client passthrough to an OpenAI/OpenRouter upstream - // =========================================================================== - - /// An OpenAI Chat client talking to the proxy receives an OpenAI-shaped - /// response back: `object=chat.completion`, translated text, normalised - /// `finish_reason`, and usage with OpenAI token field names. #[tokio::test] async fn openai_client_to_openrouter_translates_to_openai_shape() { let mock_url = spawn_mock(openrouter_non_stream_response()).await; @@ -463,20 +394,10 @@ mod tests { "text must survive the round trip: {}", body ); - // Usage keeps OpenAI field names. assert_eq!(body["usage"]["prompt_tokens"], 12); assert_eq!(body["usage"]["completion_tokens"], 8); } - // =========================================================================== - // Goal: cross-protocol translation (the headline feature) - // =========================================================================== - - /// An Anthropic Messages client talking to the proxy — whose upstream is an - /// OpenAI/OpenRouter endpoint — receives an Anthropic-shaped response: the - /// upstream OpenAI body is decoded to the core type and re-encoded for the - /// Anthropic client. Text, `stop_reason` (`stop` -> `end_turn`), and usage - /// (`completion_tokens` -> `output_tokens`) must all translate. #[tokio::test] async fn anthropic_client_to_openrouter_translates_to_anthropic_shape() { let mock_url = spawn_mock(openrouter_non_stream_response()).await; @@ -493,10 +414,8 @@ mod tests { ); let body = read_json(resp).await; - // Anthropic envelope. assert_eq!(body["type"], "message"); assert_eq!(body["role"], "assistant"); - // Text translated through the OpenAI -> core -> Anthropic path. assert!( body["content"][0]["text"] .as_str() @@ -505,19 +424,11 @@ mod tests { "text must survive cross-protocol translation: {}", body ); - // OpenAI finish_reason "stop" normalises to Anthropic "end_turn". assert_eq!(body["stop_reason"], "end_turn"); - // Usage field names translate to the Anthropic convention. assert_eq!(body["usage"]["input_tokens"], 12); assert_eq!(body["usage"]["output_tokens"], 8); } - // =========================================================================== - // Goal: tool-call fidelity - // =========================================================================== - - /// An OpenAI client receives upstream `tool_calls` intact (id, function name, - /// JSON arguments), with `finish_reason=tool_calls`. #[tokio::test] async fn tool_calls_preserved_for_openai_client() { let mock_url = spawn_mock(openrouter_tool_call_response()).await; @@ -555,9 +466,6 @@ mod tests { ); } - /// An Anthropic client receives upstream OpenAI tool calls translated into - /// Anthropic `tool_use` content blocks: `type=tool_use`, the function `name`, - /// and parsed `input` (arguments string -> object). #[tokio::test] async fn tool_calls_translate_to_anthropic_tool_use_blocks() { let mock_url = spawn_mock(openrouter_tool_call_response()).await; @@ -582,25 +490,13 @@ mod tests { .expect("at least one content block must be a tool_use") .clone(); assert_eq!(tool_use["name"], "get_weather"); - // The OpenAI arguments JSON string is parsed into an Anthropic input object. assert_eq!(tool_use["input"]["city"], "San Francisco"); } - // =========================================================================== - // Goal: robustness to provider-specific response fields - // =========================================================================== - - /// OpenRouter responses carry fields the proxy does not model (`gen-` ids, - /// top-level `provider`, `cost` inside `usage`, `system_fingerprint`, - /// `native_finish_reason`, `logprobs`, `refusal`). The OpenAI response structs - /// deliberately omit `deny_unknown_fields` (see the doc on `ChatCompletionResponse` - /// in `llm-proxy-protocol/src/openai.rs`), so these must not surface as a 502 to - /// either client protocol. This test pins that contract. #[tokio::test] async fn openrouter_specific_extras_do_not_break_conversion() { let mock_url = spawn_mock(openrouter_non_stream_response()).await; - // OpenAI client. let app = build_router(openrouter_state(&mock_url, HashMap::new())); let resp = app .oneshot(openai_chat_request(&make_chat_body(false))) @@ -617,7 +513,6 @@ mod tests { let v: Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["object"], "chat.completion"); - // Anthropic client (cross-protocol) — same upstream body, different decode. let app = build_router(openrouter_state(&mock_url, HashMap::new())); let resp = app .oneshot(anthropic_messages_request(&make_messages_body(false))) @@ -632,12 +527,6 @@ mod tests { assert_eq!(v["type"], "message"); } - // =========================================================================== - // Goal: streaming fidelity - // =========================================================================== - - /// An OpenAI client streaming request yields `chat.completion.chunk` frames and - /// terminates with `[DONE]`, with the translated text present in the deltas. #[tokio::test] async fn openai_client_stream_ends_with_done() { let mock_url = spawn_mock_openrouter_stream().await; @@ -664,10 +553,6 @@ mod tests { ); } - /// An Anthropic client streaming against an OpenAI/OpenRouter upstream is the - /// cross-protocol streaming gold path: OpenAI chunks are decoded to core events - /// and re-emitted as Anthropic SSE events (`message_start`, - /// `content_block_delta`, `message_stop`). #[tokio::test] async fn anthropic_client_stream_cross_protocol() { let mock_url = spawn_mock_openrouter_stream().await; @@ -698,13 +583,6 @@ mod tests { ); } - // =========================================================================== - // Goal: OpenRouter app-header injection - // =========================================================================== - - /// Headers declared on the adapter (`HTTP-Referer`, `X-Title` — what OpenRouter - /// expects from integrating apps) are forwarded on the upstream request, while - /// the bearer API key is applied via the configured `auth_style`. #[tokio::test] async fn openrouter_app_headers_forwarded_upstream() { let (mock_url, captured) = @@ -721,7 +599,6 @@ mod tests { .unwrap(); assert_eq!(resp.status(), StatusCode::OK); - // Give the background mock handler a moment to record the request. tokio::time::sleep(Duration::from_millis(100)).await; let g = captured.lock().unwrap(); let upstream_headers = g @@ -747,28 +624,11 @@ mod tests { ); } - // =========================================================================== - // Goal: error fidelity — upstream 429 passes through, client-shaped - // =========================================================================== - - /// An OpenRouter 429 surfaces to *each* client protocol as a client-shaped - /// error with the 429 *status* passed through — not collapsed to a 502. Two - /// finer points are pinned here, both from `routes/error_response.rs`: - /// - /// - The error `type` is `"api_error"`, **not** `"rate_limit_error"`. The - /// `rate_limit_error` label is reserved for the proxy's *own* rate limiter - /// (`RouteError::RateLimited`); an *upstream* 429 is `RouteError::Upstream`, - /// which keeps status 429 (the one code `map_upstream_status` passes through) - /// but is typed generically. - /// - The upstream's raw error body (OpenRouter's `provider_name`, `raw`, etc.) - /// is logged server-side only and **must not** reach the client envelope — - /// the client gets a sanitized generic message. #[tokio::test] async fn openrouter_429_surfaces_as_client_shaped_error() { let mock_url = spawn_mock_status(StatusCode::TOO_MANY_REQUESTS, openrouter_429_body()).await; - // Anthropic client (cross-protocol error translation). let app = build_router(openrouter_state(&mock_url, HashMap::new())); let resp = app .oneshot(anthropic_messages_request(&make_messages_body(false))) @@ -790,7 +650,6 @@ mod tests { "upstream raw error text must not leak into the client envelope: {msg}" ); - // OpenAI client. let app = build_router(openrouter_state(&mock_url, HashMap::new())); let resp = app .oneshot(openai_chat_request(&make_chat_body(false))) @@ -812,15 +671,6 @@ mod tests { ); } - // =========================================================================== - // Goal: content/token fidelity — nothing is stripped in either direction - // =========================================================================== - - /// Input fidelity: the request the proxy forwards upstream must carry every - /// piece of content the client sent — the system prompt, every message turn, - /// and verbatim text including unicode/emoji/CJK. If the encode path dropped - /// any token, one of these distinctive fragments would be missing from the - /// captured upstream body. (Anthropic-client -> OpenAI-upstream encode path.) #[tokio::test] async fn input_content_is_preserved_on_encode() { let (mock_url, captured) = @@ -874,11 +724,6 @@ mod tests { ); } - /// Output fidelity: a rich OpenRouter response (long multi-paragraph text with - /// unicode/emoji, two tool calls, and explicit usage) must convert to an - /// Anthropic response that preserves ALL of it — full verbatim text, every - /// tool_use block, and exact usage token counts. If the decode path stripped - /// anything, a fragment, a tool_use block, or a token count would be wrong. #[tokio::test] async fn output_content_is_preserved_on_decode() { let long_text = "Para one: the quick brown fox. Para two: café résumé 🚀 你好世界. \ @@ -920,7 +765,6 @@ mod tests { let v = read_json(resp).await; let joined = serde_json::to_string(&v["content"]).expect("content serializes"); - // Full verbatim text — every paragraph and the unicode/emoji runs. for fragment in [ "Para one: the quick brown fox", "café résumé 🚀 你好世界", @@ -931,7 +775,6 @@ mod tests { "decode stripped output text fragment {fragment:?}\ncontent: {joined}" ); } - // Both tool calls survive as Anthropic tool_use blocks. let tool_uses = v["content"] .as_array() .map(|blocks| blocks.iter().filter(|b| b["type"] == "tool_use").count()) @@ -940,7 +783,6 @@ mod tests { tool_uses, 2, "both tool calls must survive decode: {joined}" ); - // Usage token counts pass through unchanged (OpenAI -> Anthropic field names). assert_eq!( v["usage"]["input_tokens"], 77, "input token count must be preserved" @@ -951,12 +793,6 @@ mod tests { ); } - /// Robustness: real OpenRouter (and other OpenAI-compatible) 2xx responses can - /// omit `usage` entirely (e.g. certain Venice-backed free models). Since - /// `ChatCompletionResponse.usage` is `#[serde(default)]`, an absent `usage` - /// decodes with zero token counts instead of failing the whole response into a - /// 502 "provider response decode error". This previously reproduced the - /// intermittent live 502 for `meta-llama/llama-3.2-3b-instruct:free`. #[tokio::test] async fn response_missing_usage_decodes_with_zero_usage() { let body = serde_json::to_vec(&json!({ @@ -970,7 +806,6 @@ mod tests { "message": { "role": "assistant", "content": "pong" }, "finish_reason": "stop" }] - // NOTE: deliberately no "usage" field — some OpenRouter backends omit it. })) .unwrap(); @@ -995,8 +830,6 @@ mod tests { ); } - /// Variant where the provider sends `"usage": null` instead of omitting the - /// field. Must also decode with zero usage, not 502. #[tokio::test] async fn response_with_null_usage_decodes_with_zero_usage() { let body = serde_json::to_vec(&json!({ @@ -1031,16 +864,6 @@ mod tests { assert_eq!(v["usage"]["output_tokens"], 0); } - // =========================================================================== - // Goal: live network validation (opt-in) - // =========================================================================== - - /// Default candidate free models. OpenRouter rate-limits its free tier - /// aggressively, so the live test tries several and requires at least one to - /// succeed. Override with `OPENROUTER_MODELS=m1,m2,...` (or a single - /// `OPENROUTER_MODEL=m`). Names are current as of authoring time; OpenRouter - /// rotates free models, so a missing one is reported and skipped rather than - /// failing the run. const DEFAULT_FREE_MODELS: &[&str] = &[ "meta-llama/llama-3.2-3b-instruct:free", "liquid/lfm-2.5-1.2b-instruct:free", @@ -1052,9 +875,6 @@ mod tests { "cohere/north-mini-code:free", ]; - /// Build an AppState whose single `openrouter` provider targets the real - /// OpenRouter Chat Completions endpoint, applying the `HTTP-Referer` / `X-Title` - /// headers OpenRouter expects from integrating apps. fn live_state(api_key: &str) -> AppState { let mut headers = HashMap::new(); headers.insert( @@ -1068,6 +888,7 @@ mod tests { api_key: secrecy::SecretString::from(api_key), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -1103,7 +924,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, @@ -1119,10 +944,6 @@ mod tests { ) } - /// Issue one live Anthropic Messages request, retrying with backoff while - /// the upstream answers 429. - /// - /// Returns the final status and raw response body. async fn live_request_with_backoff( app: &axum::Router, model: &str, @@ -1155,22 +976,6 @@ mod tests { (status, bytes) } - /// Real end-to-end round trips against through the proxy. - /// - /// Skipped unless `OPENROUTER_API_KEY` is set; run with: - /// - /// ```text - /// OPENROUTER_API_KEY=sk-or-... \ - /// cargo test -p llm-proxy-server --test openrouter_e2e \ - /// live_openrouter_round_trip -- --ignored --nocapture - /// ``` - /// - /// Tries several free models (OpenRouter rate-limits the free tier, so some - /// will 429) and requires at least one to complete a full Anthropic-client → - /// OpenRouter cross-protocol conversion. A 502 from any model is a hard failure - /// (it would mean the proxy could not convert a real OpenRouter response); - /// 429/400/404/etc. are reported but tolerated. Override the set with - /// `OPENROUTER_MODELS=a,b,c` or pin a single `OPENROUTER_MODEL`. #[tokio::test] #[ignore = "requires network + OPENROUTER_API_KEY; run with --ignored"] async fn live_openrouter_round_trip() { @@ -1195,9 +1000,6 @@ mod tests { let app = build_router(live_state(&api_key)); - // OpenRouter rate-limits the free tier per-minute (and per-day for some - // models), so each model is retried with backoff on 429. Override with - // OPENROUTER_MAX_ATTEMPTS / OPENROUTER_BACKOFF_MS. let max_attempts = std::env::var("OPENROUTER_MAX_ATTEMPTS") .ok() .and_then(|s| s.parse().ok()) @@ -1214,17 +1016,9 @@ mod tests { let mut other = 0usize; let mut upstream_failures = 0usize; let mut first_ok: Option = None; - // A 502 is only a *proxy* bug when it is a ProviderDecode - // ("provider response decode error"); a 502 carrying "upstream service - // error" is an upstream 5xx the proxy correctly collapsed (the upstream - // provider's fault — common with free-tier overload) and is tolerated. - // Asserted empty at the end; run with RUST_LOG=llm_proxy_server=debug to - // see the underlying serde cause of any real decode failure. let mut decode_failures: Vec<(String, String)> = Vec::new(); for model in &candidates { - // Anthropic Messages client -> OpenRouter (OpenAI) upstream: the - // full cross-protocol translation against the real service. let body = json!({ "model": model, "messages": [{ "role": "user", "content": "Reply with exactly the word: pong" }], @@ -1238,10 +1032,6 @@ mod tests { eprintln!("[{model}] final status={status} body={text}"); - // Distinguish the two 502 causes. Only a ProviderDecode ("provider - // response decode error") is the proxy's fault; an "upstream service - // error" 502 is an upstream 5xx the proxy correctly collapsed and is - // tolerated. match status { StatusCode::OK => { let v: Value = serde_json::from_slice(&bytes) @@ -1295,12 +1085,6 @@ mod tests { } } - /// Live output fidelity: a real model is asked to emit a short, distinctive - /// multi-token sequence (six planet names). The converted Anthropic response - /// must contain ALL of them — if the decode/re-encode path dropped any output - /// token, a word would be missing. Reasoning-only responses (no text block, - /// common when a thinking model is truncated) are skipped rather than failed. - /// Retries rate-limited models once. Skipped without `OPENROUTER_API_KEY`. #[tokio::test] #[ignore = "requires network + OPENROUTER_API_KEY; run with --ignored"] async fn live_output_completeness_no_token_loss() { @@ -1336,7 +1120,6 @@ mod tests { "stream": false }); - // One retry on rate limit, with a short backoff. let (status, bytes) = live_request_with_backoff( &app, model, @@ -1352,8 +1135,6 @@ mod tests { let full = String::from_utf8_lossy(&bytes); let v: Value = serde_json::from_slice(&bytes).expect("OK body must be valid JSON"); - // Scan only text blocks — ignore tool_use ids / thinking that may carry - // unrelated tokens. let text_only: String = v["content"] .as_array() .map(|blocks| { diff --git a/crates/llm-proxy-server/tests/secret_abuse.rs b/crates/llm-proxy-server/tests/secret_abuse.rs new file mode 100644 index 0000000..a949510 --- /dev/null +++ b/crates/llm-proxy-server/tests/secret_abuse.rs @@ -0,0 +1,807 @@ +//! Secret-abuse family: secret material must never be observable. Every test +//! drives the real surfaces (router, admin router, client-auth, ledger, +//! `/usage`) and pairs each negative scan with a positive control. + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::path::{Path, PathBuf}; + use std::sync::{Arc, Mutex, MutexGuard}; + use std::time::Duration; + + use axum::Router; + use axum::body::Body; + use axum::http::{Request, StatusCode, header}; + use axum::routing::post; + use llm_proxy_core::{ + AdminAuthConfig, AppConfig, AuthMode, AuthStyle, ProviderAdapterConfig, ProviderConfig, + ProviderRegistry, ProviderRoutesConfig, ServerConfig, StorageConfig, + }; + use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; + use llm_proxy_server::{AppState, BuildInfo, ClientAuth, admin_router, build_router}; + use llm_proxy_storage::{EventBus, ProxyEvent, RecordingBus, SqliteEventBus}; + use tower::ServiceExt; + + const LAPTOP_TOKEN: &str = "prx6-laptop-client-token-4f1a9c73"; + const PHONE_TOKEN: &str = "prx6-phone-client-token-b27e03d1"; + const ADMIN_TOKEN: &str = "prx6-admin-token-9d5c48ea"; + const WRONG_TOKEN: &str = "prx6-attacker-guess-00000011"; + const STRAY_TOKEN: &str = "prx6-stray-bearer-under-noauth-77c1"; + const PROVIDER_KEY: &str = "prx6-provider-upstream-key-6b8f2e45"; + const BODY_MARKER: &str = "PRX6-BODY-MARKER-never-persist-me"; + const SECRETS: &[&str] = &[ + LAPTOP_TOKEN, + PHONE_TOKEN, + ADMIN_TOKEN, + WRONG_TOKEN, + STRAY_TOKEN, + PROVIDER_KEY, + BODY_MARKER, + ]; + const KEY_ID_LAPTOP: &str = "laptop"; + const KEY_ID_PHONE: &str = "phone"; + + const ANTHROPIC_SUCCESS: &str = concat!( + r#"{"id":"msg_prx6","type":"message","role":"assistant","content":"#, + r#"[{"type":"text","text":"prx6-mock-reply"}],"model":"claude-sonnet-4-6","#, + r#""stop_reason":"end_turn","stop_sequence":null,"#, + r#""usage":{"input_tokens":10,"output_tokens":5}}"# + ); + + fn cheap_phc(token: &str, salt: &[u8]) -> String { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid cheap test params"); + Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()) + .hash_password_customized( + token.as_bytes(), + salt, + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string() + } + + fn client_auth(pairs: &[(&str, &str)]) -> (ClientAuth, Vec) { + let mut clients = HashMap::new(); + let mut phcs = Vec::new(); + for (index, (key_id, token)) in pairs.iter().enumerate() { + let salt = format!("{index:016x}"); + let phc = cheap_phc(token, salt.as_bytes()); + clients.insert( + (*key_id).to_owned(), + secrecy::SecretString::from(phc.clone()), + ); + phcs.push(phc); + } + let auth = ClientAuth::from_config(&llm_proxy_core::AuthConfig { + mode: AuthMode::ClientKeys, + clients, + admin: None, + }) + .expect("valid client auth"); + (auth, phcs) + } + + fn admin_config() -> AdminAuthConfig { + AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + } + } + + fn base_state( + mock_endpoint: &str, + storage_path: Option, + admin: Option, + on_main: bool, + ) -> AppState { + let provider = ProviderConfig { + name: "mock-provider".to_owned(), + api_key: secrecy::SecretString::from(PROVIDER_KEY), + auth_style: AuthStyle::Bearer, + passthrough_auth: false, + egress: None, + adapters: { + let mut adapters = HashMap::new(); + adapters.insert( + "main".to_owned(), + ProviderAdapterConfig { + protocol: "anthropic_messages".to_owned(), + endpoint: mock_endpoint.to_owned(), + headers: Arc::new(HashMap::new()), + }, + ); + adapters + }, + routes: ProviderRoutesConfig { + messages: Some("main".to_owned()), + chat_completions: None, + }, + model_aliases: HashMap::new(), + discovery: None, + catalog: None, + pricing: HashMap::new(), + }; + + AppState::new( + AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().expect("bind parses"), + request_timeout: Duration::from_secs(300), + shutdown_timeout: Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "secret-abuse-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: on_main, + }, + auth: llm_proxy_core::AuthConfig { + mode: AuthMode::None, + clients: HashMap::new(), + admin, + }, + storage: storage_path.map(|path| StorageConfig { + path, + max_age: None, + max_rows: None, + prune_interval: Duration::from_secs(3600), + }), + }, + ProviderRegistry::from_providers(vec![provider]).expect("registry builds"), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ) + } + + fn marked_body(marker: &str) -> String { + serde_json::json!({ + "model": "claude-sonnet-4-6", + "messages": [{ "role": "user", "content": format!("hello {marker}") }], + "max_tokens": 64, + "stream": false + }) + .to_string() + } + + fn post_messages(token: Option<&str>, body: &str) -> Request { + request_with_token("POST", "/providers/mock-provider/v1/messages", token, body) + } + + fn get_usage(uri: &str, token: Option<&str>) -> Request { + request_with_token("GET", uri, token, "") + } + + fn request_with_token( + method: &str, + uri: &str, + token: Option<&str>, + body: &str, + ) -> Request { + let mut builder = Request::builder() + .method(method) + .uri(uri) + .header("content-type", "application/json"); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + builder + .body(Body::from(body.to_owned())) + .expect("request builds") + } + + async fn send(app: Router, request: Request) -> (StatusCode, String, Option) { + let response = app.oneshot(request).await.expect("oneshot completes"); + let status = response.status(); + let request_id = response + .headers() + .get("x-request-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body reads"); + ( + status, + String::from_utf8_lossy(&bytes).into_owned(), + request_id, + ) + } + + fn assert_free_of_secrets(context: &str, text: &str) { + for secret in SECRETS { + assert!( + !text.contains(secret), + "{context} leaked secret material ({secret}): {text}" + ); + } + } + + fn hex_sha256(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(bytes); + format!("{:x}", hasher.finalize()) + } + + fn events_for_request(bus: &RecordingBus, request_id: &str) -> Vec { + bus.snapshot() + .iter() + .filter_map(|event| { + let json = serde_json::to_string(event).expect("event serializes"); + json.contains(request_id).then_some(json) + }) + .collect() + } + + fn open_ledger(path: &Path) -> SqliteEventBus { + SqliteEventBus::open( + path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens") + } + + async fn close_ledger(bus: Arc) { + tokio::task::spawn_blocking(move || bus.close()) + .await + .expect("close joins"); + } + + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + async fn wait_ready(addr: std::net::SocketAddr) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + if tokio::net::TcpStream::connect(addr).await.is_ok() { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "mock upstream at {addr} never became ready" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + async fn spawn_anthropic_mock() -> String { + let app = Router::new().route( + "/{*path}", + post(|| async { + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "application/json")], + ANTHROPIC_SUCCESS, + ) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("mock binds"); + let addr = listener.local_addr().expect("mock address"); + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app).await { + eprintln!("mock upstream serve error: {error}"); + } + }); + wait_ready(addr).await; + format!("http://{addr}/providers/mock-provider/v1/messages") + } + + async fn spawn_echo_mock() -> String { + let app = Router::new().route( + "/{*path}", + post(|body: axum::body::Bytes| async move { + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "application/json")], + body, + ) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("mock binds"); + let addr = listener.local_addr().expect("mock address"); + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app).await { + eprintln!("mock upstream serve error: {error}"); + } + }); + wait_ready(addr).await; + format!("http://{addr}/providers/mock-provider/v1/messages") + } + + struct TraceCapture { + buffer: Arc>>, + } + + static CAPTURE_BUFFER: std::sync::OnceLock>>> = std::sync::OnceLock::new(); + static CAPTURE_INSTALLED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + + impl TraceCapture { + fn install() -> Self { + use tracing_subscriber::layer::SubscriberExt; + let buffer = CAPTURE_BUFFER.get_or_init(|| Arc::new(Mutex::new(Vec::::new()))); + CAPTURE_INSTALLED.get_or_init(|| { + let layer = tracing_subscriber::fmt::layer() + .with_writer(BufferWriter(Arc::clone(buffer))) + .with_ansi(false); + tracing::subscriber::set_global_default(tracing_subscriber::registry().with(layer)) + .expect("the global capture subscriber installs exactly once"); + }); + Self { + buffer: Arc::clone(buffer), + } + } + + fn text(&self) -> String { + let bytes = self.buffer.lock().expect("buffer lock").clone(); + String::from_utf8_lossy(&bytes).into_owned() + } + } + + #[derive(Clone)] + struct BufferWriter(Arc>>); + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufferWriter { + type Writer = BufferWriterGuard<'a>; + + fn make_writer(&'a self) -> Self::Writer { + BufferWriterGuard(self.0.lock().expect("buffer lock")) + } + } + + struct BufferWriterGuard<'a>(MutexGuard<'a, Vec>); + + impl std::io::Write for BufferWriterGuard<'_> { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn client_and_admin_tokens_never_reach_events_tracing_or_responses() { + let mock_url = spawn_anthropic_mock().await; + let (auth, phcs) = + client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN), (KEY_ID_PHONE, PHONE_TOKEN)]); + let bus = Arc::new(RecordingBus::new()); + let state = base_state(&mock_url, None, None, false) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state.clone()); + + let traffic: [(Option<&str>, StatusCode); 5] = [ + (Some(LAPTOP_TOKEN), StatusCode::OK), + (Some(PHONE_TOKEN), StatusCode::OK), + (Some(WRONG_TOKEN), StatusCode::UNAUTHORIZED), + (None, StatusCode::UNAUTHORIZED), + (Some(ADMIN_TOKEN), StatusCode::UNAUTHORIZED), + ]; + + let capture = TraceCapture::install(); + for (token, expected) in traffic { + let (status, body, _) = + send(app.clone(), post_messages(token, &marked_body(BODY_MARKER))).await; + assert_eq!( + status, expected, + "unexpected status for token {token:?}: {body}" + ); + assert_free_of_secrets(&format!("response body for {token:?}"), &body); + for phc in &phcs { + assert!( + !body.contains(phc.as_str()), + "response body leaked the argon2 verifier: {body}" + ); + } + } + + let events = bus.snapshot(); + assert!( + events.len() >= 6, + "traffic must produce events (got {})", + events.len() + ); + for event in &events { + let json = serde_json::to_string(event).expect("event serializes"); + assert_free_of_secrets("recorded event", &json); + for phc in &phcs { + assert!( + !json.contains(phc.as_str()), + "event leaked the argon2 verifier: {json}" + ); + } + } + let callers: Vec<&str> = events + .iter() + .filter_map(|event| match event { + ProxyEvent::RequestReceived(received) => Some(received.caller.as_str()), + _ => None, + }) + .collect(); + assert!(callers.contains(&KEY_ID_LAPTOP), "callers: {callers:?}"); + assert!(callers.contains(&KEY_ID_PHONE), "callers: {callers:?}"); + + let captured = capture.text(); + assert!( + captured.contains("client authenticated"), + "control: debug-level request traffic must land in the capture, got: {captured}" + ); + assert_free_of_secrets("tracing output", &captured); + for phc in &phcs { + assert!( + !captured.contains(phc.as_str()), + "tracing output leaked the argon2 verifier: {captured}" + ); + } + + let state_debug = format!("{state:?}"); + assert!( + !state_debug.contains(LAPTOP_TOKEN) && !state_debug.contains(PHONE_TOKEN), + "AppState Debug leaked a client token: {state_debug}" + ); + for phc in &phcs { + assert!( + !state_debug.contains(phc.as_str()), + "AppState Debug leaked a verifier: {state_debug}" + ); + } + } + + #[tokio::test] + async fn secrets_absent_from_raw_ledger_file_bytes() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let mock_url = spawn_anthropic_mock().await; + let (auth, phcs) = + client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN), (KEY_ID_PHONE, PHONE_TOKEN)]); + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state); + + let body = marked_body(BODY_MARKER); + let digest = hex_sha256(body.as_bytes()); + let (status, response, _) = + send(app.clone(), post_messages(Some(LAPTOP_TOKEN), &body)).await; + assert_eq!(status, StatusCode::OK, "valid client request: {response}"); + let (status, _, _) = send(app.clone(), post_messages(Some(WRONG_TOKEN), &body)).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + let (status, _, _) = + send(app.clone(), get_usage("/usage/requests", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK); + let (status, _, _) = send( + app.clone(), + get_usage("/usage/requests", Some(LAPTOP_TOKEN)), + ) + .await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "client keys are not admin" + ); + + close_ledger(Arc::clone(&bus)).await; + let raw = ledger_bytes(&db); + assert!(!raw.is_empty(), "the ledger file must exist on disk"); + + let text = String::from_utf8_lossy(&raw); + for expected in [KEY_ID_LAPTOP, "mock-provider", digest.as_str()] { + assert!( + text.contains(expected), + "control: persisted value {expected} must be findable in the ledger bytes" + ); + } + assert_free_of_secrets("raw ledger file bytes", &text); + for phc in &phcs { + assert!( + !text.contains(phc.as_str()), + "ledger bytes leaked an argon2 verifier" + ); + } + } + + #[tokio::test] + async fn usage_serves_the_body_digest_but_never_the_body_content() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let mock_url = spawn_echo_mock().await; + let (auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let main_app = build_router(state.clone()); + let dedicated = admin_router(&state).expect("admin router for the dedicated mount"); + + let body = marked_body(BODY_MARKER); + let digest = hex_sha256(body.as_bytes()); + let (status, response, _) = + send(main_app.clone(), post_messages(Some(LAPTOP_TOKEN), &body)).await; + assert_ne!( + status, + StatusCode::OK, + "a reflected request-shaped body must not decode as a valid response" + ); + assert_free_of_secrets("client-facing error response", &response); + + close_ledger(Arc::clone(&bus)).await; + for (mount, app) in [("dedicated", dedicated), ("main-opt-in", main_app)] { + for uri in [ + "/usage/requests?limit=1000", + "/usage/summary?group_by=caller", + "/usage/keys", + ] { + let (status, body, _) = send(app.clone(), get_usage(uri, Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK, "{mount} {uri}: {body}"); + assert_free_of_secrets(&format!("{mount} {uri} response"), &body); + assert!( + !body.contains(BODY_MARKER), + "{mount} {uri} must never carry request-body content: {body}" + ); + } + + let (_, body, _) = send( + app.clone(), + get_usage("/usage/requests?limit=1000", Some(ADMIN_TOKEN)), + ) + .await; + let rows = serde_json::from_str::(&body) + .expect("rows JSON") + .get("rows") + .cloned() + .unwrap_or_default(); + let digested: Vec<&str> = rows + .as_array() + .map(|rows| { + rows.iter() + .filter_map(|row| row["body_hash"].as_str()) + .collect() + }) + .unwrap_or_default(); + assert_eq!( + digested, + vec![digest.as_str()], + "{mount}: the request row must carry the sha256 of the actual body, not the body" + ); + assert!( + body.contains(KEY_ID_LAPTOP), + "{mount}: caller attribution must be visible to the admin" + ); + } + } + + #[tokio::test] + async fn admin_surface_refuses_non_admin_credentials_on_both_mounts() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let bus = Arc::new(open_ledger(&db)); + bus.emit(&ProxyEvent::RequestReceived( + llm_proxy_storage::RequestReceived { + request_id: "prx6-seed".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "prx6-seed-caller".to_owned(), + provider: "mock-provider".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "seed-digest".to_owned(), + }, + )); + close_ledger(Arc::clone(&bus)).await; + + let (auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let state = base_state("http://127.0.0.1:9", Some(db), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)); + let mounts = [ + ("dedicated", admin_router(&state).expect("admin router")), + ("main-opt-in", build_router(state)), + ]; + + for (mount, app) in mounts { + let (missing_status, missing_body, _) = + send(app.clone(), get_usage("/usage/summary", None)).await; + assert_eq!( + missing_status, + StatusCode::UNAUTHORIZED, + "{mount}: missing credentials must be 401 (not 403/5xx)" + ); + + for token in [WRONG_TOKEN, LAPTOP_TOKEN] { + let (status, body, _) = + send(app.clone(), get_usage("/usage/summary", Some(token))).await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "{mount}: {token:?} must not authenticate to the admin surface" + ); + assert_free_of_secrets(&format!("{mount} rejection for {token:?}"), &body); + assert_eq!( + body, missing_body, + "{mount}: every rejection carries the identical uniform body" + ); + } + + let (status, _, _) = + send(app.clone(), get_usage("/usage/summary", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK, "{mount}: the admin token serves"); + } + + let (proxy_auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let proxy_state = base_state("http://127.0.0.1:9", None, None, false) + .with_client_auth(Arc::new(proxy_auth)); + let (status, body, _) = send( + build_router(proxy_state), + post_messages(Some(ADMIN_TOKEN), &marked_body(BODY_MARKER)), + ) + .await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "the admin token must not authenticate proxy routes" + ); + assert_free_of_secrets("proxy-surface rejection of the admin token", &body); + } + + #[tokio::test] + async fn one_callers_key_id_is_never_observable_by_another_caller() { + const ALPHA_ID: &str = "prx6-iso-alpha"; + const BETA_ID: &str = "prx6-iso-beta"; + const ALPHA_TOKEN: &str = "prx6-iso-alpha-token-3d5e"; + const BETA_TOKEN: &str = "prx6-iso-beta-token-8f1b"; + + let mock_url = spawn_anthropic_mock().await; + let (auth, _phcs) = client_auth(&[(ALPHA_ID, ALPHA_TOKEN), (BETA_ID, BETA_TOKEN)]); + let bus = Arc::new(RecordingBus::new()); + let state = base_state(&mock_url, None, None, false) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state); + + let (status, _, _) = send( + app.clone(), + post_messages(Some(ALPHA_TOKEN), &marked_body("alpha traffic")), + ) + .await; + assert_eq!(status, StatusCode::OK); + + let capture = TraceCapture::install(); + let (status, body, request_id) = send( + app.clone(), + post_messages(Some(BETA_TOKEN), &marked_body("beta traffic")), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + !body.contains(ALPHA_ID), + "beta's response names alpha's key id: {body}" + ); + + let request_id = request_id.expect("x-request-id header present"); + let captured = capture.text(); + let correlated: Vec<&str> = captured + .lines() + .filter(|line| line.contains(&request_id) || line.contains(BETA_ID)) + .collect(); + assert!( + !correlated.is_empty(), + "control: beta's traffic must produce log lines" + ); + for line in &correlated { + assert!( + !line.contains(ALPHA_ID), + "a log line for beta's request names alpha's key id: {line}" + ); + } + + let beta_event_json = events_for_request(&bus, &request_id); + assert!(beta_event_json.len() >= 2, "beta's request emitted events"); + for json in &beta_event_json { + assert!( + !json.contains(ALPHA_ID), + "an event for beta's request names alpha's key id: {json}" + ); + } + let beta_callers: Vec = bus + .snapshot() + .into_iter() + .filter_map(|event| match event { + ProxyEvent::RequestReceived(received) if received.request_id == request_id => { + Some(received.caller) + } + _ => None, + }) + .collect(); + assert_eq!( + beta_callers, + vec![BETA_ID.to_owned()], + "self-attribution intact" + ); + } + + #[tokio::test] + async fn anonymous_attribution_survives_the_loopback_exemption() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let mock_url = spawn_anthropic_mock().await; + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), false) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state.clone()); + + let body = marked_body(BODY_MARKER); + let (status, response, _) = + send(app.clone(), post_messages(Some(STRAY_TOKEN), &body)).await; + assert_eq!(status, StatusCode::OK); + assert_free_of_secrets("response under the loopback exemption", &response); + + close_ledger(Arc::clone(&bus)).await; + let admin = admin_router(&state).expect("admin router"); + let (status, keys, _) = + send(admin.clone(), get_usage("/usage/keys", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK); + assert!( + keys.contains("anonymous"), + "attribution survives as the anonymous caller: {keys}" + ); + assert_free_of_secrets("/usage/keys under the loopback exemption", &keys); + assert!( + !keys.contains(STRAY_TOKEN), + "the stray bearer must not become the caller: {keys}" + ); + + let (status, rows, _) = send( + admin, + get_usage("/usage/requests?limit=1000", Some(ADMIN_TOKEN)), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + rows.contains("anonymous") && rows.contains(&hex_sha256(body.as_bytes())), + "the anonymous row carries the digest: {rows}" + ); + assert_free_of_secrets("/usage/requests under the loopback exemption", &rows); + + let bytes = ledger_bytes(&db); + let text = String::from_utf8_lossy(&bytes); + assert!( + text.contains("anonymous"), + "control: the anonymous caller is persisted" + ); + assert_free_of_secrets("ledger bytes under the loopback exemption", &text); + } +} diff --git a/crates/llm-proxy-storage/Cargo.toml b/crates/llm-proxy-storage/Cargo.toml index 0ce7f1a..ca1a376 100644 --- a/crates/llm-proxy-storage/Cargo.toml +++ b/crates/llm-proxy-storage/Cargo.toml @@ -20,8 +20,14 @@ test-utils = [] [dependencies] llm-proxy-protocol = { path = "../llm-proxy-protocol" } +rusqlite = { workspace = true } serde = { workspace = true } +thiserror = { workspace = true } time = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } [dev-dependencies] +rust_decimal = { workspace = true } serde_json = { workspace = true } +tempfile = "3" diff --git a/crates/llm-proxy-storage/src/lib.rs b/crates/llm-proxy-storage/src/lib.rs index a6b47c9..e115794 100644 --- a/crates/llm-proxy-storage/src/lib.rs +++ b/crates/llm-proxy-storage/src/lib.rs @@ -1,13 +1,6 @@ -//! Structured request/response event log surface. -//! -//! The persistence backend (SQLite, etc.) lands in a later phase. This crate -//! defines the event types and the [`EventBus`] trait that backends implement, -//! plus a [`NoopBus`] default and a [`RecordingBus`] for tests. -//! -//! Events are emitted at request/response boundaries in the server crate. A -//! real persisted backend will implement [`EventBus`]; for now [`NoopBus`] is -//! the production sink and [`RecordingBus`] (behind the `test-utils` cargo -//! feature) lets integration tests assert on the event sequence. +//! Event types and the [`EventBus`] trait for the per-request usage ledger. +//! Backends: [`NoopBus`] (default), [`RecordingBus`] (behind the `test-utils` +//! feature), and [`SqliteEventBus`] (persistent SQLite ledger). #![deny(missing_docs)] @@ -15,14 +8,11 @@ use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage}; use serde::{Deserialize, Serialize}; use time::OffsetDateTime; -// --------------------------------------------------------------------------- -// ProxyEvent -// --------------------------------------------------------------------------- +mod sqlite; + +pub use sqlite::{RetentionPolicy, SqliteEventBus, StorageError, run_retention_pass}; /// One request/response lifecycle event, ready to persist or emit. -/// -/// Internally tagged via `kind` so JSON/TOML consumers can discriminate without -/// a wrapper object. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum ProxyEvent { @@ -34,10 +24,6 @@ pub enum ProxyEvent { ResponseFailed(ResponseFailed), } -// --------------------------------------------------------------------------- -// RequestReceived -// --------------------------------------------------------------------------- - /// Event payload for a received request. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RequestReceived { @@ -46,6 +32,9 @@ pub struct RequestReceived { /// When the request was received (RFC 3339). #[serde(with = "time::serde::rfc3339")] pub timestamp: OffsetDateTime, + /// Authenticated client key id, or `"anonymous"` under the loopback + /// no-auth exemption. Client secrets never appear on events. + pub caller: String, /// Provider name from the URL path. pub provider: String, /// Route kind (`"chat_completions"` or `"messages"`). @@ -61,10 +50,6 @@ pub struct RequestReceived { pub body_hash: String, } -// --------------------------------------------------------------------------- -// ResponseCompleted -// --------------------------------------------------------------------------- - /// Event payload for a completed response. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponseCompleted { @@ -75,8 +60,7 @@ pub struct ResponseCompleted { pub timestamp: OffsetDateTime, /// Provider name. pub provider: String, - /// Upstream message id, if available from `MessageStart` or - /// `CoreResponse.id`. + /// Upstream message id, when the provider supplied one. pub upstream_message_id: Option, /// Requested and upstream model. pub model: ModelRef, @@ -90,15 +74,8 @@ pub struct ResponseCompleted { pub latency_ms: u64, } -// --------------------------------------------------------------------------- -// ResponseFailed -// --------------------------------------------------------------------------- - -/// Event payload for a failed response. -/// -/// `model` and `provider` are optional because early failures (unknown -/// provider, JSON parse error, rate limit) can occur before the model or -/// provider is confirmed. +/// Event payload for a failed response; `model` and `provider` are optional +/// because early failures can occur before either is confirmed. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponseFailed { /// Proxy-generated unique request identifier. @@ -112,14 +89,9 @@ pub struct ResponseFailed { pub model: Option, /// Error kind string from the `RouteError` variant name. pub error_kind: String, - /// Sanitized error message; API keys and URLs are always redacted. - /// - /// For upstream errors whose status is passed through to the client - /// (400/413/429, and client-owned 401/403) this field holds a generic - /// proxy message so the provider's own schema (field names, model - /// identifiers, request echoes) never reaches the client; the real - /// upstream body is logged server-side only. For statuses that collapse to - /// 502, this field contains the truncated, key/URL-redacted upstream body. + /// Sanitized error message; API keys and URLs are always redacted, and + /// the provider's own response schema never reaches the client — statuses + /// passed through hold a generic proxy message, 502s a truncated body. pub message: String, /// HTTP status code returned to the client. pub http_status: u16, @@ -127,21 +99,16 @@ pub struct ResponseFailed { pub latency_ms: u64, } -// --------------------------------------------------------------------------- -// EventBus trait + implementations -// --------------------------------------------------------------------------- - /// In-process event sink. -/// -/// Implementations: [`NoopBus`], [`RecordingBus`] (tests), and the future -/// SQLite backend. -/// -/// `Debug` is a super-trait so that `Arc` can be formatted in -/// `AppState`'s manual `Debug` impl. pub trait EventBus: Send + Sync + std::fmt::Debug { /// Emit one event. Must not block the calling task; backends do IO on a /// dedicated writer task. fn emit(&self, event: &ProxyEvent); + + /// Signal that no further events will be emitted and flush anything + /// queued; called once at process shutdown. Implementations must stay + /// bounded and may block briefly; wrap in `spawn_blocking` from async code. + fn close(&self) {} } /// Default no-op sink. @@ -152,13 +119,12 @@ impl EventBus for NoopBus { fn emit(&self, _event: &ProxyEvent) {} } -/// Test-only sink that records events for assertions. Available behind the -/// `test-utils` cargo feature (NOT `#[cfg(test)]`, which is stripped when the -/// crate is compiled as a dependency). +/// Test-only sink that records events for assertions; available behind the +/// `test-utils` cargo feature. #[cfg(any(test, feature = "test-utils"))] #[derive(Debug, Default)] pub struct RecordingBus { - /// Recorded events, guarded by a mutex for interior mutability. + /// Recorded events. pub events: std::sync::Mutex>, } @@ -169,11 +135,8 @@ impl RecordingBus { Self::default() } - /// Snapshot the recorded events (cloned under the lock). - /// - /// Recovers from mutex poison rather than returning an empty `Vec`, so a - /// poisoned bus still reflects the events recorded before the panic - /// (consistent with the recover-poison pattern used in `metrics.rs`). + /// Snapshot the recorded events (cloned under the lock); recovers from + /// mutex poison rather than dropping events recorded before the panic. pub fn snapshot(&self) -> Vec { self.events .lock() @@ -185,9 +148,6 @@ impl RecordingBus { #[cfg(any(test, feature = "test-utils"))] impl EventBus for RecordingBus { fn emit(&self, event: &ProxyEvent) { - // Recover from mutex poison so a poisoned bus still records the event - // rather than silently dropping it (recover-poison pattern, consistent - // with `metrics.rs` and `catalog_service.rs`). self.events .lock() .unwrap_or_else(|e| e.into_inner()) @@ -204,6 +164,7 @@ mod tests { RequestReceived { request_id: "req-1".to_owned(), timestamp: OffsetDateTime::now_utc(), + caller: "anonymous".to_owned(), provider: "p".to_owned(), route_kind: "messages".to_owned(), client_protocol: "anthropic".to_owned(), @@ -219,7 +180,6 @@ mod tests { #[test] fn noop_bus_emit_does_not_panic_or_record() { let bus = NoopBus; - // NoopBus has nowhere to record; just ensure emit is callable. let event = ProxyEvent::RequestReceived(sample_request_received()); bus.emit(&event); } @@ -260,14 +220,6 @@ mod tests { use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; - // Pin a known timestamp instead of `now_utc()` so the format assertions - // are deterministic and the round-trip can be checked exactly. The - // nanos value (`123_456_789`) is deliberately chosen to be nonzero at - // every precision group (ms/us/ns): `time`'s RFC 3339 serializer - // trims trailing zero groups, so a value like `123_000_000` would be - // emitted as milliseconds (`.123`, 3 digits) rather than nanoseconds. - // Forcing nonzero nanos guarantees the full 9-digit fractional component - // is emitted, so the precision assertion below is meaningful. let timestamp = OffsetDateTime::from_unix_timestamp(1_700_000_000) .expect("valid unix timestamp") .replace_nanosecond(123_456_789) @@ -275,6 +227,7 @@ mod tests { let event = ProxyEvent::RequestReceived(RequestReceived { request_id: "req-1".to_owned(), timestamp, + caller: "anonymous".to_owned(), provider: "p".to_owned(), route_kind: "messages".to_owned(), client_protocol: "anthropic".to_owned(), @@ -286,33 +239,20 @@ mod tests { body_hash: "abc123".to_owned(), }); let json = serde_json::to_string(&event).expect("serialize"); - // Internally tagged by `kind`. assert!(json.contains(r#""kind":"request_received""#), "got: {json}"); - // Extract the serialized timestamp string from the JSON and validate a - // full RFC 3339 shape rather than only asserting `contains('T')`. The - // previous weaker assertion would pass for any value with a literal - // 'T' anywhere, masking a malformed/non-RFC-3339 timestamp. let extracted = extract_json_string_field(&json, "timestamp") .expect("timestamp field present and a JSON string"); - // Must end in 'Z' (UTC) or a numeric offset like +00:00 / -05:00. assert!( extracted.ends_with('Z') || is_numeric_offset_suffix(&extracted), "timestamp must end in Z or a numeric offset, got: {extracted}" ); - // Full RFC 3339 round-trip: parsing the extracted string back into an - // OffsetDateTime must succeed and reproduce the original instant. let parsed = OffsetDateTime::parse(&extracted, &Rfc3339).expect("timestamp parses as RFC 3339"); assert_eq!( parsed, timestamp, "timestamp must round-trip through RFC 3339; got: {extracted}" ); - // Precision: `time::serde::rfc3339` emits a fractional-seconds component - // and trims trailing zero groups (3 digits for whole-ms, 6 for whole-us, - // 9 otherwise). Our pinned nanos are nonzero at every group, so the full - // 9-digit nanosecond component is emitted. Confirm the subsecond part is - // present and carries nanosecond precision. let frac = extracted .split('.') .nth(1) @@ -329,9 +269,7 @@ mod tests { ); } - /// Return true if `s` ends in an RFC 3339 numeric offset like `+00:00`. fn is_numeric_offset_suffix(s: &str) -> bool { - // A numeric offset is `[+-]HH:MM` (6 chars). if s.len() < 6 { return false; } @@ -343,12 +281,6 @@ mod tests { tail[1..].chars().all(|c| c.is_ascii_digit() || c == ':') } - /// Extract the string value of a top-level JSON string field from a - /// serialized `ProxyEvent` payload. - /// - /// Returns `None` if the field is absent or not a JSON string. Honors the - /// JSON `\"` escape so an embedded quote does not prematurely terminate the - /// value. fn extract_json_string_field(json: &str, field: &str) -> Option { let needle = format!("\"{field}\":\""); let start = json.find(&needle)? + needle.len(); @@ -377,22 +309,8 @@ mod tests { Some(json[start..end].to_owned()) } - /// Source guard: `ProxyEvent` has no secret-carrying field by construction - /// (the API key lives in `ProviderConfig`, never on an event). This test - /// guards against future regressions where a field named like a secret - /// (`api_key`, `secret`, `token`, `bearer`, `authorization`, `password`) - /// is accidentally added to one of the event payloads. - /// - /// It serializes all three variants and asserts none of the forbidden - /// secret-bearing field names appear as a JSON key. It also populates every - /// string content field with a sentinel and asserts the JSON only ever - /// surfaces it under one of the *known* field names (proving no new, - /// unknown string field was silently introduced). #[test] fn proxy_event_has_no_secret_field_by_construction() { - // Field-name fragments that would indicate a secret-bearing field if - // they appeared as a JSON key. None of these are legitimate `ProxyEvent` - // field names. const FORBIDDEN_KEY_FRAGMENTS: &[&str] = &[ "api_key", "apikey", @@ -404,15 +322,12 @@ mod tests { "passwd", ]; - // A sentinel placed in every legitimate string content field. It SHOULD - // appear in the serialized JSON (under the known field names) — the - // assertions below check it never appears under a forbidden field name. const SENTINEL: &str = "ZZ-sentinel-content-ZZ"; - // --- RequestReceived --- let req = ProxyEvent::RequestReceived(RequestReceived { request_id: SENTINEL.to_owned(), timestamp: OffsetDateTime::now_utc(), + caller: SENTINEL.to_owned(), provider: SENTINEL.to_owned(), route_kind: SENTINEL.to_owned(), client_protocol: SENTINEL.to_owned(), @@ -426,7 +341,6 @@ mod tests { let req_json = serde_json::to_string(&req).expect("serialize RequestReceived"); assert_no_forbidden_secret_key(&req_json, FORBIDDEN_KEY_FRAGMENTS); - // --- ResponseCompleted --- let done = ProxyEvent::ResponseCompleted(ResponseCompleted { request_id: SENTINEL.to_owned(), timestamp: OffsetDateTime::now_utc(), @@ -444,7 +358,6 @@ mod tests { let done_json = serde_json::to_string(&done).expect("serialize ResponseCompleted"); assert_no_forbidden_secret_key(&done_json, FORBIDDEN_KEY_FRAGMENTS); - // --- ResponseFailed --- let failed = ProxyEvent::ResponseFailed(ResponseFailed { request_id: SENTINEL.to_owned(), timestamp: OffsetDateTime::now_utc(), @@ -454,9 +367,6 @@ mod tests { upstream: Some(SENTINEL.to_owned()), }), error_kind: SENTINEL.to_owned(), - // The `message` field is sanitized error text; place the sentinel - // here too so the guard proves it surfaces only as a value, never - // under a forbidden key. message: SENTINEL.to_owned(), http_status: 429, latency_ms: 2, @@ -465,17 +375,8 @@ mod tests { assert_no_forbidden_secret_key(&failed_json, FORBIDDEN_KEY_FRAGMENTS); } - /// Assert that none of `fragments` appear as a JSON key in `json`. - /// - /// A JSON key is always serialized as `"fragment"` immediately followed by - /// `:` (modulo whitespace, which `serde_json::to_string` does not emit), so - /// checking for `"fragment":` avoids false positives where a fragment shows - /// up inside a legitimate string value. fn assert_no_forbidden_secret_key(json: &str, fragments: &[&str]) { for frag in fragments { - // Match the fragment as a JSON key: a `"` then the fragment then `"` - // then optional whitespace then `:`. serde_json emits no spaces, so - // the common shape is `"frag":`. let key_shape_quoted = format!("\"{frag}\":"); assert!( !json.contains(&key_shape_quoted), @@ -499,7 +400,6 @@ mod tests { let json = serde_json::to_string(&event).expect("serialize"); let back: ProxyEvent = serde_json::from_str(&json).expect("deserialize"); assert!(matches!(back, ProxyEvent::ResponseFailed(_))); - // Sanity: UsageProvenance must be referenced so it is not flagged unused. let _ = UsageProvenance::Unknown; } } diff --git a/crates/llm-proxy-storage/src/sqlite.rs b/crates/llm-proxy-storage/src/sqlite.rs new file mode 100644 index 0000000..ed4a669 --- /dev/null +++ b/crates/llm-proxy-storage/src/sqlite.rs @@ -0,0 +1,1206 @@ +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::{Duration, Instant}; + +use llm_proxy_protocol::core::{StopReason, UsageProvenance}; +use rusqlite::Connection; +use time::OffsetDateTime; +use tokio::sync::mpsc::{self, Sender}; +use tokio::sync::watch; + +use crate::{EventBus, ProxyEvent}; + +const DEFAULT_CHANNEL_CAPACITY: usize = 1024; +const WRITE_BATCH_MAX: usize = 256; +const CLOSE_DRAIN_TIMEOUT: Duration = Duration::from_secs(10); +const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); + +/// Errors from the SQLite usage ledger. +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + /// Failed to create the ledger file's parent directory. + #[error("creating usage-ledger parent directory {path}: {source}")] + CreateDir { + /// Parent directory that could not be created. + path: String, + /// Underlying I/O error. + source: std::io::Error, + }, + /// Failed to open or initialize the SQLite database. + #[error("opening usage ledger at {path}: {source}")] + Open { + /// Database path. + path: String, + /// Underlying rusqlite error. + source: rusqlite::Error, + }, + /// A write or retention operation on an open ledger failed. + #[error("usage ledger operation failed: {0}")] + Sqlite(#[from] rusqlite::Error), +} + +/// Retention bounds applied to the ledger on a timer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RetentionPolicy { + /// Delete rows older than this age. `None` keeps rows forever. + pub max_age: Option, + /// Keep at most this many rows; the oldest rows are deleted when the cap + /// is exceeded. `None` is unlimited. + pub max_rows: Option, + /// How often the pruning pass runs; must be at least 1 second (a zero + /// period is rejected by the pruning timer). + pub prune_every: Duration, +} + +impl RetentionPolicy { + fn is_bounded(&self) -> bool { + self.max_age.is_some() || self.max_rows.is_some() + } +} + +/// Run one retention pass, deleting rows past `policy`'s bounds; returns the +/// number of rows deleted. Age pruning covers failure-only rows (no +/// `received_at`) via `COALESCE(received_at, completed_at)`. +pub fn run_retention_pass( + conn: &mut Connection, + policy: &RetentionPolicy, +) -> Result { + if !policy.is_bounded() { + return Ok(0); + } + let mut deleted: u64 = 0; + if let Some(max_age) = policy.max_age { + let cutoff = unix_millis(OffsetDateTime::now_utc()) + .saturating_sub(i64::try_from(max_age.as_millis()).unwrap_or(i64::MAX)); + let n = conn.execute( + "DELETE FROM requests WHERE COALESCE(received_at, completed_at) < ?1", + rusqlite::params![cutoff], + )?; + deleted += u64::try_from(n).unwrap_or(u64::MAX); + } + if let Some(max_rows) = policy.max_rows { + let count: i64 = conn.query_row("SELECT COUNT(*) FROM requests", [], |row| row.get(0))?; + let excess = u64::try_from(count) + .unwrap_or(u64::MAX) + .saturating_sub(max_rows); + if excess > 0 { + let limit = i64::try_from(excess).unwrap_or(i64::MAX); + let n = conn.execute( + "DELETE FROM requests WHERE request_id IN ( + SELECT request_id FROM requests + ORDER BY COALESCE(received_at, completed_at) ASC + LIMIT ?1 + )", + rusqlite::params![limit], + )?; + deleted += u64::try_from(n).unwrap_or(u64::MAX); + } + } + Ok(deleted) +} + +const SCHEMA_SQL: &str = "\ +CREATE TABLE IF NOT EXISTS requests ( + request_id TEXT PRIMARY KEY, + received_at INTEGER, + completed_at INTEGER, + caller TEXT, + provider TEXT, + route_kind TEXT, + client_protocol TEXT, + streaming INTEGER, + model_requested TEXT, + model_upstream TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + reasoning_tokens INTEGER, + cache_read_tokens INTEGER, + cache_creation_tokens INTEGER, + usage_provenance TEXT, + cost_input TEXT, + cost_output TEXT, + cost_cache_creation TEXT, + cost_cache_read TEXT, + cost_reasoning TEXT, + cost_total TEXT, + stop_reason TEXT, + error_kind TEXT, + http_status INTEGER, + latency_ms INTEGER, + body_hash TEXT +); +CREATE INDEX IF NOT EXISTS idx_requests_timestamp + ON requests (received_at); +CREATE INDEX IF NOT EXISTS idx_requests_provider_timestamp + ON requests (provider, received_at); +CREATE INDEX IF NOT EXISTS idx_requests_caller_timestamp + ON requests (caller, received_at); +"; + +fn open_connection(path: &Path) -> Result { + let open_err = |source: rusqlite::Error| StorageError::Open { + path: path.display().to_string(), + source, + }; + let conn = Connection::open(path).map_err(open_err)?; + conn.pragma_update(None, "journal_mode", "WAL") + .map_err(open_err)?; + conn.pragma_update(None, "synchronous", "NORMAL") + .map_err(open_err)?; + conn.busy_timeout(SQLITE_BUSY_TIMEOUT).map_err(open_err)?; + conn.execute_batch(SCHEMA_SQL).map_err(open_err)?; + Ok(conn) +} + +fn unix_millis(ts: OffsetDateTime) -> i64 { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) +} + +fn stop_reason_label(reason: &StopReason) -> String { + match reason { + StopReason::EndTurn => "end_turn".to_owned(), + StopReason::MaxTokens => "max_tokens".to_owned(), + StopReason::ToolUse => "tool_use".to_owned(), + StopReason::StopSequence => "stop_sequence".to_owned(), + StopReason::Refusal => "refusal".to_owned(), + StopReason::Error => "error".to_owned(), + StopReason::Unknown(inner) => format!("unknown:{inner}"), + _ => "other".to_owned(), + } +} + +fn provenance_label(provenance: &UsageProvenance) -> &'static str { + match provenance { + UsageProvenance::ProviderReported => "provider_reported", + UsageProvenance::SyntheticZero => "synthetic_zero", + UsageProvenance::Unknown => "unknown", + _ => "other", + } +} + +fn write_event(conn: &Connection, event: &ProxyEvent) -> Result<(), StorageError> { + match event { + ProxyEvent::RequestReceived(e) => insert_request_received(conn, e), + ProxyEvent::ResponseCompleted(e) => upsert_response_completed(conn, e), + ProxyEvent::ResponseFailed(e) => upsert_response_failed(conn, e), + } +} + +fn insert_request_received( + conn: &Connection, + e: &crate::RequestReceived, +) -> Result<(), StorageError> { + conn.execute( + "INSERT OR IGNORE INTO requests ( + request_id, received_at, caller, provider, route_kind, + client_protocol, streaming, model_requested, model_upstream, + body_hash + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.caller, + e.provider, + e.route_kind, + e.client_protocol, + e.streaming, + e.model.requested, + e.model.upstream, + e.body_hash, + ], + )?; + Ok(()) +} + +fn upsert_response_completed( + conn: &Connection, + e: &crate::ResponseCompleted, +) -> Result<(), StorageError> { + conn.execute( + "INSERT INTO requests ( + request_id, completed_at, provider, model_requested, + model_upstream, input_tokens, output_tokens, reasoning_tokens, + cache_read_tokens, cache_creation_tokens, usage_provenance, + cost_input, cost_output, cost_cache_creation, cost_cache_read, + cost_reasoning, cost_total, stop_reason, latency_ms + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, + ?13, ?14, ?15, ?16, ?17, ?18, ?19) + ON CONFLICT(request_id) DO UPDATE SET + completed_at = excluded.completed_at, + provider = excluded.provider, + model_requested = excluded.model_requested, + model_upstream = excluded.model_upstream, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + reasoning_tokens = excluded.reasoning_tokens, + cache_read_tokens = excluded.cache_read_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + usage_provenance = excluded.usage_provenance, + cost_input = excluded.cost_input, + cost_output = excluded.cost_output, + cost_cache_creation = excluded.cost_cache_creation, + cost_cache_read = excluded.cost_cache_read, + cost_reasoning = excluded.cost_reasoning, + cost_total = excluded.cost_total, + stop_reason = excluded.stop_reason, + latency_ms = excluded.latency_ms", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.provider, + e.model.requested, + e.model.upstream, + e.usage.input_tokens, + e.usage.output_tokens, + e.usage.reasoning_tokens, + e.usage.cache_read_input_tokens, + e.usage.cache_creation_input_tokens, + provenance_label(&e.usage.provenance), + e.cost.as_ref().map(|c| c.input.to_string()), + e.cost.as_ref().map(|c| c.output.to_string()), + e.cost.as_ref().map(|c| c.cache_creation.to_string()), + e.cost.as_ref().map(|c| c.cache_read.to_string()), + e.cost.as_ref().map(|c| c.reasoning.to_string()), + e.cost.as_ref().map(|c| c.total.to_string()), + stop_reason_label(&e.stop_reason), + i64::try_from(e.latency_ms).unwrap_or(i64::MAX), + ], + )?; + Ok(()) +} + +fn upsert_response_failed( + conn: &Connection, + e: &crate::ResponseFailed, +) -> Result<(), StorageError> { + conn.execute( + "INSERT INTO requests ( + request_id, completed_at, provider, model_requested, + model_upstream, error_kind, http_status, latency_ms + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(request_id) DO UPDATE SET + completed_at = excluded.completed_at, + provider = COALESCE(excluded.provider, requests.provider), + model_requested = COALESCE( + excluded.model_requested, requests.model_requested), + model_upstream = COALESCE( + excluded.model_upstream, requests.model_upstream), + error_kind = excluded.error_kind, + http_status = excluded.http_status, + latency_ms = excluded.latency_ms", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.provider, + e.model.as_ref().map(|m| m.requested.clone()), + e.model.as_ref().and_then(|m| m.upstream.clone()), + e.error_kind, + i64::from(e.http_status), + i64::try_from(e.latency_ms).unwrap_or(i64::MAX), + ], + )?; + Ok(()) +} + +enum WriterOp { + Write { events: Vec }, + Prune { policy: RetentionPolicy }, +} + +fn run_op(conn: &mut Connection, op: WriterOp) -> Result { + match op { + WriterOp::Write { events } => { + let tx = conn.transaction()?; + for event in &events { + write_event(&tx, event)?; + } + tx.commit()?; + Ok(u64::try_from(events.len()).unwrap_or(u64::MAX)) + } + WriterOp::Prune { policy } => run_retention_pass(conn, &policy), + } +} + +async fn exec_on_conn( + path: &Path, + conn: Connection, + op: WriterOp, +) -> Option<(Connection, Result)> { + let path = path.to_path_buf(); + let joined = tokio::task::spawn_blocking(move || { + let mut conn = conn; + let result = run_op(&mut conn, op); + (conn, result) + }) + .await; + match joined { + Ok((conn, result)) => Some((conn, result)), + Err(join_error) => { + tracing::error!( + error = %join_error, + "usage-ledger writer panicked; reopening the database" + ); + let reopen_path = path.clone(); + let reopened = tokio::task::spawn_blocking(move || open_connection(&reopen_path)).await; + match reopened { + Ok(Ok(conn)) => Some((conn, Ok(0))), + Ok(Err(error)) => { + tracing::error!( + error = %error, + path = %path.display(), + "failed to reopen usage ledger; writer stopping" + ); + None + } + Err(join_error) => { + tracing::error!( + error = %join_error, + "reopening the usage ledger panicked; writer stopping" + ); + None + } + } + } + } +} + +type DoneFlag = Arc<(Mutex, Condvar)>; + +fn signal_done(done: &DoneFlag) { + let mut guard = done.0.lock().unwrap_or_else(|e| e.into_inner()); + *guard = true; + drop(guard); + done.1.notify_all(); +} + +async fn writer_task( + path: PathBuf, + conn: Connection, + mut rx: mpsc::Receiver, + mut shutdown: watch::Receiver, + retention: RetentionPolicy, + done: DoneFlag, +) { + let pruning = retention.is_bounded(); + let mut prune = tokio::time::interval(retention.prune_every.max(Duration::from_secs(1))); + prune.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut conn_slot = Some(conn); + let mut buf: Vec = Vec::with_capacity(WRITE_BATCH_MAX); + + loop { + if *shutdown.borrow_and_update() { + break; + } + tokio::select! { + biased; + _ = shutdown.changed() => { + if *shutdown.borrow() { + break; + } + } + received = rx.recv_many(&mut buf, WRITE_BATCH_MAX) => { + if received == 0 { + break; + } + let Some(conn) = conn_slot.take() else { break }; + match exec_on_conn( + &path, + conn, + WriterOp::Write { events: std::mem::take(&mut buf) }, + ) + .await + { + Some((c, result)) => { + conn_slot = Some(c); + if let Err(error) = result { + tracing::error!( + error = %error, + "usage-ledger batch write failed; batch events lost" + ); + } + } + None => break, + } + } + _ = prune.tick(), if pruning => { + if let Some(conn) = conn_slot.take() { + match exec_on_conn( + &path, + conn, + WriterOp::Prune { policy: retention }, + ) + .await + { + Some((c, result)) => { + conn_slot = Some(c); + match result { + Ok(rows) if rows > 0 => tracing::info!( + rows, + "pruned usage-ledger rows past retention bounds" + ), + Ok(_) => {} + Err(error) => tracing::error!( + error = %error, + "usage-ledger retention pass failed" + ), + } + } + None => break, + } + } + } + } + } + + while let Ok(event) = rx.try_recv() { + buf.push(event); + } + if !buf.is_empty() { + if let Some(conn) = conn_slot.take() { + match exec_on_conn(&path, conn, WriterOp::Write { events: buf }).await { + Some((_conn, Ok(_))) => {} + Some((_conn, Err(error))) => { + tracing::error!(error = %error, "usage-ledger final drain failed"); + } + None => {} + } + } + } + + signal_done(&done); +} + +/// SQLite-backed [`EventBus`]: persists every event to one `requests` ledger. +/// `emit` never blocks — a full channel drops the event and increments +/// [`SqliteEventBus::dropped_events`]. Open inside a tokio runtime. +#[derive(Debug)] +pub struct SqliteEventBus { + tx: Sender, + shutdown_tx: Mutex>, + done: DoneFlag, + dropped: AtomicU64, + path: PathBuf, +} + +impl SqliteEventBus { + /// Open (or create) the ledger at `path` and start its writer task (parent + /// directories, WAL mode, and schema created as needed). Blocking + /// filesystem work: call at startup, not on the hot path. + pub fn open(path: impl AsRef, retention: RetentionPolicy) -> Result { + Self::open_with_capacity(path, retention, DEFAULT_CHANNEL_CAPACITY) + } + + fn open_with_capacity( + path: impl AsRef, + retention: RetentionPolicy, + capacity: usize, + ) -> Result { + let path = path.as_ref(); + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent).map_err(|source| StorageError::CreateDir { + path: parent.display().to_string(), + source, + })?; + } + } + let conn = open_connection(path)?; + let (tx, rx) = mpsc::channel(capacity.max(1)); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let done: DoneFlag = Arc::new((Mutex::new(false), Condvar::new())); + tokio::spawn(writer_task( + path.to_path_buf(), + conn, + rx, + shutdown_rx, + retention, + Arc::clone(&done), + )); + Ok(Self { + tx, + shutdown_tx: Mutex::new(shutdown_tx), + done, + dropped: AtomicU64::new(0), + path: path.to_path_buf(), + }) + } + + #[cfg(test)] + fn unconnected(capacity: usize) -> Self { + let (tx, rx) = mpsc::channel(capacity.max(1)); + let (shutdown_tx, _shutdown_rx) = watch::channel(false); + Box::leak(Box::new(rx)); + Self { + tx, + shutdown_tx: Mutex::new(shutdown_tx), + done: Arc::new((Mutex::new(false), Condvar::new())), + dropped: AtomicU64::new(0), + path: PathBuf::new(), + } + } + + /// Number of events dropped so far (channel full, or writer gone). + #[must_use] + pub fn dropped_events(&self) -> u64 { + self.dropped.load(Ordering::Relaxed) + } + + /// Path of the backing database file. + #[must_use] + pub fn path(&self) -> &Path { + &self.path + } +} + +impl EventBus for SqliteEventBus { + fn emit(&self, event: &ProxyEvent) { + match self.tx.try_send(event.clone()) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + let dropped = self.dropped.fetch_add(1, Ordering::Relaxed) + 1; + tracing::debug!( + dropped, + "usage-ledger channel full; event dropped (never backpressure a response)" + ); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + let dropped = self.dropped.fetch_add(1, Ordering::Relaxed) + 1; + tracing::debug!(dropped, "usage-ledger writer is gone; event dropped"); + } + } + } + + fn close(&self) { + if let Ok(shutdown_tx) = self.shutdown_tx.lock() { + if shutdown_tx.send(true).is_err() { + tracing::debug!("usage-ledger writer had already exited"); + } + } + let deadline = Instant::now() + CLOSE_DRAIN_TIMEOUT; + let mut guard = self.done.0.lock().unwrap_or_else(|e| e.into_inner()); + while !*guard { + let now = Instant::now(); + if now >= deadline { + tracing::warn!( + timeout_secs = CLOSE_DRAIN_TIMEOUT.as_secs(), + "usage-ledger close timed out; queued events may not be flushed" + ); + return; + } + let (next, timeout) = self + .done + .1 + .wait_timeout(guard, deadline - now) + .unwrap_or_else(|e| e.into_inner()); + guard = next; + if timeout.timed_out() { + continue; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{RequestReceived, ResponseCompleted, ResponseFailed}; + use llm_proxy_protocol::core::{Cost, ModelRef, Usage}; + use rust_decimal::Decimal; + + const RECEIVED_MS: i64 = 1_700_000_000_000; + const COMPLETED_MS: i64 = 1_700_000_004_000; + + fn unbounded() -> RetentionPolicy { + RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + } + } + + fn received_event(id: &str, caller: &str) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(RECEIVED_MS / 1000) + .expect("valid unix timestamp"), + caller: caller.to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: Some("claude-3-upstream".to_owned()), + }, + streaming: true, + body_hash: "sha256-digest-only".to_owned(), + }) + } + + fn completed_event(id: &str) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(COMPLETED_MS / 1000) + .expect("valid unix timestamp"), + provider: "prov-a".to_owned(), + upstream_message_id: Some("msg_7".to_owned()), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 100, + output_tokens: 50, + reasoning_tokens: Some(7), + cache_creation_input_tokens: Some(30), + cache_read_input_tokens: Some(200), + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from_str_exact("0.000123456789012345678901234") + .expect("decimal fits"), + output: Decimal::from_str_exact("2.25").expect("decimal fits"), + cache_creation: Decimal::ZERO, + cache_read: Decimal::from_str_exact("0.0000005").expect("decimal fits"), + reasoning: Decimal::ZERO, + total: Decimal::from_str_exact("2.250123956789012345678901234") + .expect("decimal fits"), + }), + stop_reason: StopReason::Unknown("weird-finish".to_owned()), + latency_ms: 4321, + }) + } + + fn failed_event(id: &str, provider: Option<&str>) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(COMPLETED_MS / 1000) + .expect("valid unix timestamp"), + provider: provider.map(str::to_owned), + model: None, + error_kind: "UpstreamError".to_owned(), + message: "sanitized failure text".to_owned(), + http_status: 502, + latency_ms: 99, + }) + } + + async fn wait_until(description: &str, mut cond: impl FnMut() -> bool) { + for _ in 0..200 { + if cond() { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("condition not met within 2s: {description}"); + } + + fn reader(path: &Path) -> Connection { + Connection::open(path).expect("open ledger for assertions") + } + + fn scalar(conn: &Connection, sql: &str) -> i64 { + conn.query_row(sql, [], |row| row.get(0)) + .expect("scalar query") + } + + fn text_column(conn: &Connection, sql: &str) -> Option { + conn.query_row(sql, [], |row| row.get::<_, Option>(0)) + .expect("text query") + } + + #[tokio::test] + async fn open_enables_wal_and_creates_schema_and_indexes() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + let bus = Arc::new(bus); + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + let mode: String = conn + .query_row("PRAGMA journal_mode", [], |row| row.get(0)) + .expect("journal mode"); + assert_eq!(mode, "wal", "ledger must run in WAL mode"); + + let indexes: Vec = { + let mut stmt = conn + .prepare("SELECT name FROM sqlite_master WHERE type = 'index'") + .expect("prepare index query"); + let rows = stmt + .query_map([], |row| row.get::<_, String>(0)) + .expect("query indexes"); + let mut names = Vec::new(); + for row in rows { + names.push(row.expect("index row")); + } + names + }; + for expected in [ + "idx_requests_timestamp", + "idx_requests_provider_timestamp", + "idx_requests_caller_timestamp", + ] { + assert!( + indexes.contains(&expected.to_owned()), + "missing index {expected}; got {indexes:?}" + ); + } + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 0, + "a fresh ledger starts empty" + ); + } + + #[tokio::test] + async fn request_received_persists_row_with_caller_and_body_hash() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-1", "laptop")); + let conn = reader(&db); + wait_until("request row appears", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'req-1'", + ) == 1 + }) + .await; + + let row = conn + .query_row( + "SELECT received_at, caller, provider, route_kind, client_protocol, \ + streaming, model_requested, model_upstream, body_hash \ + FROM requests WHERE request_id = 'req-1'", + [], + |row| { + Ok(( + row.get::<_, i64>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + row.get::<_, String>(3)?, + row.get::<_, String>(4)?, + row.get::<_, i64>(5)?, + row.get::<_, String>(6)?, + row.get::<_, Option>(7)?, + row.get::<_, String>(8)?, + )) + }, + ) + .expect("fetch row"); + assert_eq!(row.0, RECEIVED_MS); + assert_eq!(row.1, "laptop", "caller (PRX-01 key id) must persist"); + assert_eq!(row.2, "prov-a"); + assert_eq!(row.3, "messages"); + assert_eq!(row.4, "anthropic"); + assert_eq!(row.5, 1, "streaming flag stored as 1"); + assert_eq!(row.6, "claude-3"); + assert_eq!(row.7.as_deref(), Some("claude-3-upstream")); + assert_eq!(row.8, "sha256-digest-only"); + let has_body_column: i64 = scalar( + &conn, + "SELECT COUNT(*) FROM pragma_table_info('requests') WHERE name = 'body'", + ); + assert_eq!(has_body_column, 0, "no request-body column may exist"); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn completed_event_fills_the_same_row_with_usage() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-2", "laptop")); + bus.emit(&completed_event("req-2")); + let conn = reader(&db); + wait_until("completed row is filled", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id = 'req-2' AND completed_at IS NOT NULL", + ) == 1 + }) + .await; + + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 1, + "one request, one row: the terminal event updates in place" + ); + assert_eq!( + scalar( + &conn, + "SELECT input_tokens FROM requests WHERE request_id = 'req-2'" + ), + 100 + ); + assert_eq!( + scalar( + &conn, + "SELECT output_tokens FROM requests WHERE request_id = 'req-2'" + ), + 50 + ); + assert_eq!( + scalar( + &conn, + "SELECT reasoning_tokens FROM requests WHERE request_id = 'req-2'" + ), + 7 + ); + assert_eq!( + scalar( + &conn, + "SELECT cache_read_tokens FROM requests WHERE request_id = 'req-2'" + ), + 200 + ); + assert_eq!( + scalar( + &conn, + "SELECT cache_creation_tokens FROM requests WHERE request_id = 'req-2'" + ), + 30 + ); + assert_eq!( + text_column( + &conn, + "SELECT usage_provenance FROM requests WHERE request_id = 'req-2'" + ) + .as_deref(), + Some("provider_reported") + ); + assert_eq!( + scalar( + &conn, + "SELECT received_at FROM requests WHERE request_id = 'req-2'" + ), + RECEIVED_MS, + "the terminal event must not clobber received_at" + ); + assert_eq!( + scalar( + &conn, + "SELECT completed_at FROM requests WHERE request_id = 'req-2'" + ), + COMPLETED_MS + ); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn completed_event_persists_exact_cost_and_stop_reason() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-cost", "laptop")); + bus.emit(&completed_event("req-cost")); + let conn = reader(&db); + wait_until("completed row carries its cost", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id = 'req-cost' AND cost_total IS NOT NULL", + ) == 1 + }) + .await; + + let cost_input_text = text_column( + &conn, + "SELECT cost_input FROM requests WHERE request_id = 'req-cost'", + ) + .expect("cost_input stored"); + assert_eq!( + Decimal::from_str_exact(&cost_input_text).expect("cost parses back"), + Decimal::from_str_exact("0.000123456789012345678901234").expect("original"), + "cost must round-trip exactly through its stored text" + ); + assert_eq!( + text_column( + &conn, + "SELECT cost_total FROM requests WHERE request_id = 'req-cost'" + ) + .as_deref(), + Some("2.250123956789012345678901234") + ); + assert_eq!( + text_column( + &conn, + "SELECT stop_reason FROM requests WHERE request_id = 'req-cost'" + ) + .as_deref(), + Some("unknown:weird-finish"), + "unknown stop reasons keep their provider value" + ); + assert_eq!( + scalar( + &conn, + "SELECT latency_ms FROM requests WHERE request_id = 'req-cost'" + ), + 4321 + ); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn failed_event_upserts_and_preserves_known_fields() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + + bus.emit(&received_event("req-3", "laptop")); + bus.emit(&failed_event("req-3", None)); + bus.emit(&failed_event("orphan-1", Some("prov-b"))); + + let conn = reader(&db); + wait_until("both failure rows exist", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id IN ('req-3', 'orphan-1') AND error_kind IS NOT NULL", + ) == 2 + }) + .await; + + assert_eq!( + text_column( + &conn, + "SELECT provider FROM requests WHERE request_id = 'req-3'" + ) + .as_deref(), + Some("prov-a"), + "a None provider on failure must not NULL out the stored provider" + ); + assert_eq!( + text_column( + &conn, + "SELECT error_kind FROM requests WHERE request_id = 'req-3'" + ) + .as_deref(), + Some("UpstreamError") + ); + assert_eq!( + scalar( + &conn, + "SELECT http_status FROM requests WHERE request_id = 'req-3'" + ), + 502 + ); + + let orphan_received: Option = conn + .query_row( + "SELECT received_at FROM requests WHERE request_id = 'orphan-1'", + [], + |row| row.get(0), + ) + .expect("orphan row"); + assert_eq!( + orphan_received, None, + "a failure-only row has no received timestamp to report" + ); + assert_eq!( + scalar( + &conn, + "SELECT completed_at FROM requests WHERE request_id = 'orphan-1'" + ), + COMPLETED_MS + ); + assert_eq!( + text_column( + &conn, + "SELECT provider FROM requests WHERE request_id = 'orphan-1'" + ) + .as_deref(), + Some("prov-b") + ); + let has_message_column: i64 = scalar( + &conn, + "SELECT COUNT(*) FROM pragma_table_info('requests') WHERE name = 'message'", + ); + assert_eq!(has_message_column, 0); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[test] + fn emit_counts_drops_when_channel_is_full_and_never_blocks() { + let bus = SqliteEventBus::unconnected(8); + for i in 0..20 { + bus.emit(&received_event(&format!("req-{i}"), "laptop")); + } + assert_eq!(bus.dropped_events(), 12, "20 emits into capacity 8 drop 12"); + } + + #[tokio::test] + async fn every_emitted_event_is_persisted_or_counted_never_lost() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = Arc::new( + SqliteEventBus::open_with_capacity(&db, unbounded(), 32).expect("open ledger"), + ); + for i in 0..300 { + bus.emit(&received_event(&format!("r{i}"), "laptop")); + } + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + let persisted = scalar(&conn, "SELECT COUNT(*) FROM requests"); + assert_eq!( + persisted + i64::try_from(bus.dropped_events()).expect("u64 fits i64"), + 300, + "every event is either persisted or counted as dropped — never lost silently" + ); + } + + #[tokio::test] + async fn close_flushes_everything_queued_before_returning() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = Arc::new(SqliteEventBus::open(&db, unbounded()).expect("open ledger")); + for i in 0..25 { + bus.emit(&received_event(&format!("f{i}"), "laptop")); + } + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 25, + "close() returns only after the queue is fully drained" + ); + assert_eq!(bus.dropped_events(), 0); + } + + #[tokio::test] + async fn retention_pass_prunes_by_age_then_row_cap() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + for i in 0..5 { + bus.emit(&received_event(&format!("p{i}"), "laptop")); + } + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let mut conn = Connection::open(&db).expect("open ledger"); + wait_until("all five rows visible", || { + scalar(&conn, "SELECT COUNT(*) FROM requests") == 5 + }) + .await; + let now_ms = unix_millis(OffsetDateTime::now_utc()); + for (id, ts) in [ + ("p0", 1000_i64), + ("p1", 2000), + ("p2", now_ms - 60_000), + ("p3", now_ms - 50_000), + ("p4", now_ms - 40_000), + ] { + conn.execute( + "UPDATE requests SET received_at = ?1 WHERE request_id = ?2", + rusqlite::params![ts, id], + ) + .expect("backdate row"); + } + + let deleted = run_retention_pass( + &mut conn, + &RetentionPolicy { + max_age: Some(Duration::from_secs(3600)), + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("age pass"); + assert_eq!(deleted, 2, "the two ancient rows are pruned by age"); + assert_eq!(scalar(&conn, "SELECT COUNT(*) FROM requests"), 3); + + let deleted = run_retention_pass( + &mut conn, + &RetentionPolicy { + max_age: None, + max_rows: Some(1), + prune_every: Duration::from_secs(3600), + }, + ) + .expect("row-cap pass"); + assert_eq!(deleted, 2, "the two oldest of the remaining rows go"); + assert_eq!( + text_column(&conn, "SELECT request_id FROM requests"), + Some("p4".to_owned()), + "the newest row survives the cap" + ); + + let deleted = run_retention_pass(&mut conn, &unbounded()).expect("no-op pass"); + assert_eq!(deleted, 0); + } + + #[tokio::test] + async fn retention_timer_prunes_old_rows_without_manual_invocation() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let _bus = SqliteEventBus::open( + &db, + RetentionPolicy { + max_age: Some(Duration::from_secs(3600)), + max_rows: None, + prune_every: Duration::from_secs(1), + }, + ) + .expect("open ledger"); + let ProxyEvent::RequestReceived(mut fresh) = received_event("stale", "laptop") else { + panic!("builder returns RequestReceived"); + }; + fresh.timestamp = OffsetDateTime::now_utc(); + _bus.emit(&ProxyEvent::RequestReceived(fresh)); + + let conn = reader(&db); + wait_until("row appears", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'stale'", + ) == 1 + }) + .await; + conn.execute( + "UPDATE requests SET received_at = 1000 WHERE request_id = 'stale'", + [], + ) + .expect("backdate row"); + wait_until("timer prunes the stale row", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'stale'", + ) == 0 + }) + .await; + } +} diff --git a/crates/llm-proxy-storage/tests/secret_abuse.rs b/crates/llm-proxy-storage/tests/secret_abuse.rs new file mode 100644 index 0000000..769a495 --- /dev/null +++ b/crates/llm-proxy-storage/tests/secret_abuse.rs @@ -0,0 +1,133 @@ +//! Secret-abuse family for the SQLite ledger: whatever an event's sanitized +//! message holds, it is never persisted — asserted against the raw bytes of +//! the ledger file and its WAL/SHM sidecars. + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::time::Duration; + + use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage, UsageProvenance}; + use llm_proxy_storage::{ + EventBus, ProxyEvent, RequestReceived, ResponseCompleted, ResponseFailed, SqliteEventBus, + }; + use rust_decimal::Decimal; + use time::OffsetDateTime; + + const CALLER_MARKER: &str = "prx6-storage-caller-77aa01"; + const DIGEST_MARKER: &str = "prx6-storage-digest-88bb02"; + const STOP_LABEL_CONTROL: &str = "end_turn"; + const MESSAGE_MARKER: &str = "prx6-storage-message-99cc03"; + + fn received(ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: "prx6-row-1".to_owned(), + timestamp: ts, + caller: CALLER_MARKER.to_owned(), + provider: "prov".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: DIGEST_MARKER.to_owned(), + }) + } + + fn completed(ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: "prx6-row-1".to_owned(), + timestamp: ts, + provider: "prov".to_owned(), + upstream_message_id: Some("msg_upstream".to_owned()), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 11, + output_tokens: 7, + reasoning_tokens: None, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from(11), + output: Decimal::ZERO, + cache_creation: Decimal::ZERO, + cache_read: Decimal::ZERO, + reasoning: Decimal::ZERO, + total: Decimal::from(11), + }), + stop_reason: StopReason::EndTurn, + latency_ms: 42, + }) + } + + fn failed(request_id: &str, ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: request_id.to_owned(), + timestamp: ts, + provider: Some("prov".to_owned()), + model: None, + error_kind: "UpstreamError".to_owned(), + message: format!("upstream failed after echoing {MESSAGE_MARKER}"), + http_status: 502, + latency_ms: 5, + }) + } + + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + #[tokio::test] + async fn raw_ledger_bytes_persist_columns_but_never_event_messages() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let bus = SqliteEventBus::open( + &db, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens"); + + let t0 = OffsetDateTime::now_utc(); + bus.emit(&received(t0)); + bus.emit(&failed("prx6-row-1", t0)); + bus.emit(&completed(t0)); + bus.emit(&failed("prx6-orphan", t0)); + + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joins"); + + let bytes = ledger_bytes(&db); + let text = String::from_utf8_lossy(&bytes); + for control in [CALLER_MARKER, DIGEST_MARKER, STOP_LABEL_CONTROL] { + assert!( + text.contains(control), + "control: persisted value {control} must be findable in the raw bytes" + ); + } + assert!( + !text.contains(MESSAGE_MARKER), + "event message text (upstream echoes, error fragments) leaked into the \ + raw ledger file bytes" + ); + } +} diff --git a/dist-workspace.toml b/dist-workspace.toml new file mode 100644 index 0000000..8e2573d --- /dev/null +++ b/dist-workspace.toml @@ -0,0 +1,19 @@ +[workspace] +members = ["cargo:."] + +# Config for 'dist' (https://github.com/axodotdev/cargo-dist) +[dist] +# The preferred dist version to use in CI (Cargo.toml SemVer syntax) +cargo-dist-version = "0.32.0" +# CI backends to support +ci = "github" +# The installers to generate for each app +installers = ["shell"] +# Target platforms to build apps for (Rust target-triple syntax) +targets = ["x86_64-unknown-linux-musl", "aarch64-unknown-linux-musl"] +# Path that installers should place binaries in +install-path = "CARGO_HOME" +# Where to host releases +hosting = "github" +# Whether to install an updater program +install-updater = false diff --git a/providers/opencode-go.toml.example b/providers/opencode-go.toml.example index f670984..6a6e0ce 100644 --- a/providers/opencode-go.toml.example +++ b/providers/opencode-go.toml.example @@ -45,3 +45,11 @@ messages = "anthropic" # cache_creation = "0.000003" # cache_read = "0.00000125" # reasoning = "0.000010" + +# Route this provider's upstream traffic through a forward proxy. +# Allowed schemes: socks5h:// (hostnames resolved at the proxy), socks5://, +# http://, https://. The URL may carry userinfo (proxy authentication). +# Absent [provider.egress] = direct connection, as before. +# +# [provider.egress] +# proxy = "socks5h://127.0.0.1:1080" diff --git a/providers/opencode-zen.toml.example b/providers/opencode-zen.toml.example index ce481eb..ec4d66a 100644 --- a/providers/opencode-zen.toml.example +++ b/providers/opencode-zen.toml.example @@ -53,3 +53,11 @@ messages = "anthropic" # cache_creation = "0.000003" # cache_read = "0.00000125" # reasoning = "0.000010" + +# Route this provider's upstream traffic through a forward proxy. +# Allowed schemes: socks5h:// (hostnames resolved at the proxy), socks5://, +# http://, https://. The URL may carry userinfo (proxy authentication). +# Absent [provider.egress] = direct connection, as before. +# +# [provider.egress] +# proxy = "socks5h://127.0.0.1:1080"