i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-08 20:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main'

Browse Source
This commit is contained in:
fawney
2026-09-02 22:08:36 +08:00
parent 5059093d29 cae9aa4134
commit e8d9877b79
461 changed files with 82784 additions and 3546 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+1 -1
View File
@@ -72,7 +72,7 @@ ADMIN_USERNAME=admin123456
# docker compose 下 app 启动前自动执行 pending migration/backfill（默认 true）
# AETHER_GATEWAY_AUTO_PREPARE_DATABASE=true
# PostgreSQL 连接池配置（默认按 CPU 自动计算；正式高并发环境可显式预算）
# PostgreSQL 连接池配置（默认每核 4 条、总池至少 32 条且最多 100 条；多实例部署应显式分配每实例预算）
# AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS=12
# AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS=80
# AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS=2048
.github/workflows/nightly.yml
+613
View File
@@ -0,0 +1,613 @@
name: Nightly Release
on:
# 02:17 Asia/Shanghai (18:17 UTC) every day.
schedule:
- cron: '17 18 * * *'
workflow_dispatch:
# Checks and builds only need read access. Publishing jobs opt into write access
# below so a failed build cannot modify the existing nightly release.
permissions:
actions: read
contents: read
# A rolling tag and image are shared by scheduled and manually retried runs.
# Keep GitHub Release immutability disabled for this repository: the tag and
# assets intentionally move after each successful daily build.
concurrency:
group: nightly-main
cancel-in-progress: false
env:
CARGO_INCREMENTAL: '0'
CARGO_PROFILE_DEV_DEBUG: '0'
CARGO_PROFILE_TEST_DEBUG: '0'
CARGO_TERM_COLOR: always
RUST_BACKTRACE: '1'
GHCR_IMAGE: ghcr.io/fawney19/aether
jobs:
source:
name: Resolve main snapshot
runs-on: ubuntu-latest
timeout-minutes: 5
outputs:
sha: ${{ steps.snapshot.outputs.sha }}
short_sha: ${{ steps.snapshot.outputs.short_sha }}
date: ${{ steps.snapshot.outputs.date }}
steps:
- name: Require main branch
id: snapshot
shell: bash
run: |
set -euo pipefail
if [[ "${GITHUB_REF}" != "refs/heads/main" ]]; then
echo "Nightly releases must run from refs/heads/main (got ${GITHUB_REF})." >&2
exit 1
fi
sha="${GITHUB_SHA}"
echo "sha=${sha}" >> "${GITHUB_OUTPUT}"
echo "short_sha=${sha:0:7}" >> "${GITHUB_OUTPUT}"
echo "date=$(date -u +'%Y-%m-%d')" >> "${GITHUB_OUTPUT}"
echo "Building main at ${sha}."
# Keep the scheduled backend coverage in one place so it cannot drift from PR CI.
rust_ci:
name: Rust CI
needs: source
uses: ./.github/workflows/rust-ci.yml
rust_extended:
name: Rust extended checks
needs: source
runs-on: ubuntu-latest
timeout-minutes: 90
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Install pinned Rust toolchain
uses: dtolnay/rust-toolchain@stable
with:
toolchain: 1.95.0
- name: Show Rust toolchain
run: rustc -Vv
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
shared-key: nightly-rust-1.95-${{ runner.os }}
workspaces: . -> target
- name: Setup sccache
uses: mozilla-actions/sccache-action@v0.0.9
- name: Check all workspace targets
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: 'true'
run: cargo check --workspace --all-targets --all-features --locked
- name: Run workspace doctests
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: 'true'
run: cargo test --workspace --all-features --doc --locked
- name: Show sccache stats
if: always()
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: 'true'
run: sccache --show-stats
frontend:
name: Frontend checks and build
needs: source
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Setup Node.js
uses: actions/setup-node@v5
with:
node-version: '22'
cache: npm
cache-dependency-path: |
frontend/package-lock.json
aether-vscodex/web/package-lock.json
# The frontend prebuild synchronizes the embedded VSCodex UI by running
# its build from a separate package. Install that package explicitly so
# vue-tsc can resolve vite/client, vitest/globals, and node types in a
# clean runner.
- name: Install VSCodex web dependencies
working-directory: aether-vscodex/web
run: npm ci
- name: Install dependencies
working-directory: frontend
run: npm ci
- name: Lint
working-directory: frontend
run: npx --no-install eslint .
- name: Type-check
working-directory: frontend
run: npm run type-check
- name: Run unit tests
working-directory: frontend
run: npm run test:run
- name: Build nightly frontend
working-directory: frontend
env:
AETHER_BUILD_VERSION: nightly-${{ needs.source.outputs.short_sha }}
AETHER_VERSION: nightly
run: npm run build
- name: Upload frontend artifact
uses: actions/upload-artifact@v5
with:
name: nightly-frontend-dist
path: frontend/dist/
if-no-files-found: error
overwrite: true
retention-days: 7
repository_health:
name: Repository health checks
needs: source
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Setup Node.js
uses: actions/setup-node@v5
with:
node-version: '22'
- name: Check generated format coverage matrix
run: python3 docs/api/generate_format_field_coverage.py --check
- name: Test pressure report checker
run: node --test tools/pressure/check_gateway_stage_report.test.js
checks:
name: Nightly check gate
runs-on: ubuntu-latest
if: ${{ always() }}
needs:
- source
- rust_ci
- rust_extended
- frontend
- repository_health
steps:
- name: Verify check jobs
shell: bash
run: |
set -euo pipefail
failed=0
echo "source=${{ needs.source.result }}"
echo "rust_ci=${{ needs.rust_ci.result }}"
echo "rust_extended=${{ needs.rust_extended.result }}"
echo "frontend=${{ needs.frontend.result }}"
echo "repository_health=${{ needs.repository_health.result }}"
for result in \
"${{ needs.source.result }}" \
"${{ needs.rust_ci.result }}" \
"${{ needs.rust_extended.result }}" \
"${{ needs.frontend.result }}" \
"${{ needs.repository_health.result }}"; do
if [[ "${result}" != "success" ]]; then
failed=1
fi
done
if [[ "${failed}" -ne 0 ]]; then
echo 'One or more nightly checks failed or were cancelled.' >&2
exit 1
fi
build:
name: Build ${{ matrix.name }}
needs: [source, checks]
if: ${{ needs.checks.result == 'success' }}
runs-on: ${{ matrix.os }}
timeout-minutes: 120
strategy:
fail-fast: false
matrix:
include:
- name: linux-amd64
target: x86_64-unknown-linux-musl
platform: linux
arch: amd64
os: ubuntu-latest
use_cross: true
- name: linux-arm64
target: aarch64-unknown-linux-musl
platform: linux
arch: arm64
os: ubuntu-latest
use_cross: true
- name: macos-amd64
target: x86_64-apple-darwin
platform: macos
arch: amd64
os: macos-15-intel
use_cross: false
- name: macos-arm64
target: aarch64-apple-darwin
platform: macos
arch: arm64
os: macos-15
use_cross: false
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Install pinned Rust toolchain
uses: dtolnay/rust-toolchain@stable
with:
toolchain: 1.95.0
targets: ${{ matrix.target }}
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
shared-key: nightly-release-${{ matrix.target }}
workspaces: . -> target
- name: Install cross
if: matrix.use_cross
uses: taiki-e/install-action@cross
- name: Build release binary
env:
AETHER_BUILD_VERSION: nightly-${{ needs.source.outputs.short_sha }}
AETHER_VERSION: nightly
AETHER_BUILD_TYPE: release
CARGO_TERM_COLOR: always
shell: bash
run: |
if [[ "${{ matrix.use_cross }}" == "true" ]]; then
cross build --release --locked -p aether-gateway --target "${{ matrix.target }}"
else
cargo build --release --locked -p aether-gateway --target "${{ matrix.target }}"
fi
- name: Upload binary artifact
uses: actions/upload-artifact@v5
with:
name: nightly-gateway-${{ matrix.platform }}-${{ matrix.arch }}
path: target/${{ matrix.target }}/release/aether-gateway
if-no-files-found: error
overwrite: true
retention-days: 7
docker:
name: Publish nightly GHCR image
needs: [source, checks, build]
if: ${{ needs.checks.result == 'success' && needs.build.result == 'success' }}
runs-on: ubuntu-latest
permissions:
actions: read
contents: read
packages: write
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Download Linux binaries and frontend
uses: actions/download-artifact@v5
with:
pattern: nightly-*
path: artifacts
merge-multiple: false
- name: Prepare Docker build context
shell: bash
run: |
set -euo pipefail
mkdir -p dist/frontend
cp artifacts/nightly-gateway-linux-amd64/aether-gateway dist/aether-gateway-amd64
cp artifacts/nightly-gateway-linux-arm64/aether-gateway dist/aether-gateway-arm64
chmod 0755 dist/aether-gateway-amd64 dist/aether-gateway-arm64
cp -R artifacts/nightly-frontend-dist/. dist/frontend/
- name: Set up QEMU
uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Log in to GHCR
uses: docker/login-action@v3
with:
registry: ghcr.io
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Build and push nightly image
uses: docker/build-push-action@v6
with:
context: .
file: ./Dockerfile.app
push: true
platforms: linux/amd64,linux/arm64
tags: |
${{ env.GHCR_IMAGE }}:nightly
${{ env.GHCR_IMAGE }}:nightly-${{ needs.source.outputs.sha }}
labels: |
org.opencontainers.image.title=Aether
org.opencontainers.image.version=nightly
org.opencontainers.image.revision=${{ needs.source.outputs.sha }}
org.opencontainers.image.source=https://github.com/${{ github.repository }}
package:
name: Package nightly archives
needs: [source, checks, build]
if: ${{ needs.checks.result == 'success' && needs.build.result == 'success' }}
runs-on: ubuntu-latest
permissions:
actions: read
contents: read
steps:
- uses: actions/checkout@v5
with:
ref: ${{ needs.source.outputs.sha }}
- name: Download nightly artifacts
uses: actions/download-artifact@v5
with:
pattern: nightly-*
path: artifacts
merge-multiple: false
- name: Build nightly release packages
shell: bash
env:
SOURCE_REF: ${{ needs.source.outputs.sha }}
run: |
set -euo pipefail
VERSION="nightly"
mkdir -p package release-assets
for platform in linux macos; do
for arch in amd64 arm64; do
bundle="aether-${VERSION}-${platform}-${arch}"
root="package/${bundle}"
mkdir -p "${root}/bin" "${root}/frontend"
install -m 0755 \
"artifacts/nightly-gateway-${platform}-${arch}/aether-gateway" \
"${root}/bin/aether-gateway"
cp -R artifacts/nightly-frontend-dist/. "${root}/frontend/"
sed \
-e "s/^SOURCE_REF=\"\${AETHER_SOURCE_REF:-main}\"/SOURCE_REF=\"\${AETHER_SOURCE_REF:-${SOURCE_REF}}\"/" \
-e "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" \
install.sh > "${root}/install.sh"
chmod 0755 "${root}/install.sh"
install -m 0755 update.sh "${root}/update.sh"
install -m 0644 docker-compose.yml "${root}/docker-compose.yml"
install -m 0644 docker-compose.single-node.yml "${root}/docker-compose.single-node.yml"
install -m 0644 .env.example "${root}/.env.example"
install -m 0755 generate_keys.sh "${root}/generate_keys.sh"
install -m 0644 README.md "${root}/README.md"
install -m 0644 LICENSE "${root}/LICENSE"
tar -C package -czf "release-assets/${bundle}.tar.gz" "${bundle}"
done
done
sed \
-e "s/^SOURCE_REF=\"\${AETHER_SOURCE_REF:-main}\"/SOURCE_REF=\"\${AETHER_SOURCE_REF:-${SOURCE_REF}}\"/" \
-e "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" \
install.sh > release-assets/install.sh
chmod 0755 release-assets/install.sh
(cd release-assets && sha256sum *.tar.gz > SHA256SUMS)
test "$(find release-assets -maxdepth 1 -name '*.tar.gz' | wc -l)" -eq 4
test "$(wc -l < release-assets/SHA256SUMS)" -eq 4
(cd release-assets && sha256sum -c SHA256SUMS)
for archive in release-assets/*.tar.gz; do
tar -tzf "${archive}" >/dev/null
done
- name: Upload nightly package artifact
uses: actions/upload-artifact@v5
with:
name: nightly-release-assets
path: release-assets/*
if-no-files-found: error
overwrite: true
retention-days: 7
github_release:
name: Publish nightly GitHub Release
needs: [source, checks, docker, package]
if: ${{ needs.checks.result == 'success' && needs.docker.result == 'success' && needs.package.result == 'success' }}
runs-on: ubuntu-latest
permissions:
actions: read
contents: write
steps:
- name: Download nightly package artifact
uses: actions/download-artifact@v5
with:
name: nightly-release-assets
path: release-assets
- name: Update rolling nightly release
shell: bash
env:
GH_TOKEN: ${{ github.token }}
REPOSITORY: ${{ github.repository }}
RELEASE_TAG: nightly
SOURCE_SHA: ${{ needs.source.outputs.sha }}
SOURCE_SHORT_SHA: ${{ needs.source.outputs.short_sha }}
RELEASE_DATE: ${{ needs.source.outputs.date }}
run: |
set -euo pipefail
release_title="Aether Nightly ${RELEASE_DATE} (${SOURCE_SHORT_SHA})"
notes_file="${RUNNER_TEMP}/nightly-release-notes.md"
cat > "${notes_file}" <<EOF
## Aether nightly
This rolling prerelease was built from [main commit ${SOURCE_SHORT_SHA}](https://github.com/${REPOSITORY}/commit/${SOURCE_SHA}).
- Source branch: main
- Source commit: ${SOURCE_SHA}
- Build date (UTC): ${RELEASE_DATE}
- Container image: ${GHCR_IMAGE}:nightly
- Commit image: ${GHCR_IMAGE}:nightly-${SOURCE_SHA}
The nightly tag and assets are replaced by the next successful daily build.
EOF
# Create a draft on the first run. Later runs repair the same rolling
# release on retry if any upload or metadata update is interrupted.
if ! gh release view "${RELEASE_TAG}" --repo "${REPOSITORY}" >/dev/null 2>&1; then
gh release create "${RELEASE_TAG}" \
--repo "${REPOSITORY}" \
--draft \
--prerelease \
--latest=false \
--target "${SOURCE_SHA}" \
--title "${release_title}" \
--notes-file "${notes_file}"
fi
# Upload archives first, then the checksum/installer metadata. This
# keeps a failed upload from leaving a checksum that describes files
# which have not reached the Release yet.
gh release upload "${RELEASE_TAG}" release-assets/*.tar.gz \
--repo "${REPOSITORY}" \
--clobber
gh release upload "${RELEASE_TAG}" \
release-assets/SHA256SUMS \
release-assets/install.sh \
--repo "${REPOSITORY}" \
--clobber
# target_commitish does not move an existing git tag. Move the ref
# only after the complete asset set is available.
if gh api "repos/${REPOSITORY}/git/ref/tags/${RELEASE_TAG}" >/dev/null 2>&1; then
gh api -X PATCH "repos/${REPOSITORY}/git/refs/tags/${RELEASE_TAG}" \
-f "sha=${SOURCE_SHA}" \
-F 'force=true' >/dev/null
else
gh api -X POST "repos/${REPOSITORY}/git/refs" \
-f "ref=refs/tags/${RELEASE_TAG}" \
-f "sha=${SOURCE_SHA}" >/dev/null
fi
gh release edit "${RELEASE_TAG}" \
--repo "${REPOSITORY}" \
--draft=false \
--prerelease \
--latest=false \
--target "${SOURCE_SHA}" \
--title "${release_title}" \
--notes-file "${notes_file}"
expected_assets=(
aether-nightly-linux-amd64.tar.gz
aether-nightly-linux-arm64.tar.gz
aether-nightly-macos-amd64.tar.gz
aether-nightly-macos-arm64.tar.gz
SHA256SUMS
install.sh
)
asset_names="$(gh release view "${RELEASE_TAG}" --repo "${REPOSITORY}" --json assets --jq '.assets[].name')"
for expected_asset in "${expected_assets[@]}"; do
if ! grep -Fxq "${expected_asset}" <<<"${asset_names}"; then
echo "Published release is missing asset ${expected_asset}." >&2
exit 1
fi
done
resolved_sha=""
for attempt in {1..10}; do
resolved_sha="$(gh api "repos/${REPOSITORY}/commits/${RELEASE_TAG}" --jq '.sha' 2>/dev/null || true)"
if [[ "${resolved_sha}" == "${SOURCE_SHA}" ]]; then
break
fi
sleep 2
done
if [[ "${resolved_sha}" != "${SOURCE_SHA}" ]]; then
echo "nightly tag resolved to ${resolved_sha}, expected ${SOURCE_SHA}." >&2
exit 1
fi
release_state="$(gh release view "${RELEASE_TAG}" --repo "${REPOSITORY}" --json isDraft,isPrerelease --jq '[.isDraft, .isPrerelease] | @tsv')"
if [[ "${release_state}" != $'false\ttrue' ]]; then
echo "nightly release has unexpected state: ${release_state}" >&2
exit 1
fi
echo "Published ${RELEASE_TAG} for ${SOURCE_SHA}."
summary:
name: Nightly summary
runs-on: ubuntu-latest
if: ${{ always() }}
needs:
- source
- rust_ci
- rust_extended
- frontend
- repository_health
- checks
- build
- docker
- package
- github_release
steps:
- name: Verify nightly pipeline
shell: bash
run: |
set -euo pipefail
failed=0
for entry in \
"source=${{ needs.source.result }}" \
"rust_ci=${{ needs.rust_ci.result }}" \
"rust_extended=${{ needs.rust_extended.result }}" \
"frontend=${{ needs.frontend.result }}" \
"repository_health=${{ needs.repository_health.result }}" \
"checks=${{ needs.checks.result }}" \
"build=${{ needs.build.result }}" \
"docker=${{ needs.docker.result }}" \
"package=${{ needs.package.result }}" \
"github_release=${{ needs.github_release.result }}"; do
echo "${entry}"
if [[ "${entry#*=}" != "success" ]]; then
failed=1
fi
done
if [[ "${failed}" -ne 0 ]]; then
echo 'Nightly pipeline did not publish a new release.' >&2
exit 1
fi
.github/workflows/rust-ci.yml
+11 -5
View File
@@ -1,6 +1,7 @@
name: Rust CI
on:
workflow_call:
push:
branches:
- master
@@ -11,6 +12,7 @@ on:
- "crates/**"
- "apps/**"
- ".github/workflows/rust-ci.yml"
- ".github/workflows/nightly.yml"
pull_request:
paths:
- "Cargo.toml"
@@ -18,11 +20,15 @@ on:
- "crates/**"
- "apps/**"
- ".github/workflows/rust-ci.yml"
- ".github/workflows/nightly.yml"
concurrency:
group: rust-ci-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
group: rust-ci-${{ github.event_name }}-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
permissions:
contents: read
env:
CARGO_INCREMENTAL: 0
CARGO_PROFILE_DEV_DEBUG: 0
@@ -200,13 +206,13 @@ jobs:
RUSTFLAGS: "-C link-arg=-fuse-ld=mold"
run: cargo nextest run -p aether-gateway --lib
- name: Test bin
- name: Test bins
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true"
RUST_MIN_STACK: "16777216"
RUSTFLAGS: "-C link-arg=-fuse-ld=mold"
run: cargo nextest run -p aether-gateway --bin aether-gateway
run: cargo nextest run -p aether-gateway --bins
- name: Show sccache stats
if: always()
@@ -387,11 +393,11 @@ jobs:
- name: Setup sccache
uses: mozilla-actions/sccache-action@v0.0.9
- name: Test scenario binaries
- name: Test scenario binaries and end-to-end suites
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true"
run: cargo test -p aether-integration-tests --bins
run: cargo test -p aether-integration-tests --bins --tests
- name: Show sccache stats
if: always()
Cargo.lock
Generated
+4
View File
@@ -325,6 +325,7 @@ dependencies = [
"axum",
"base64 0.22.1",
"bcrypt",
"brotli",
"bytes",
"chrono",
"chrono-tz",
@@ -346,6 +347,7 @@ dependencies = [
"reqwest",
"rsa",
"rustls 0.23.37",
"semver",
"serde",
"serde_json",
"sha1",
@@ -442,6 +444,7 @@ name = "aether-integration-tests"
version = "0.1.0"
dependencies = [
"aether-contracts",
"aether-crypto",
"aether-data",
"aether-data-contracts",
"aether-gateway",
@@ -458,6 +461,7 @@ dependencies = [
"sqlx",
"tokio",
"tokio-tungstenite 0.28.0",
"uuid",
]
[[package]]
Cargo.toml
+2 -1
View File
@@ -106,6 +106,7 @@ async-trait = "0.1"
axum = "0.8"
base64 = "0.22"
bcrypt = "0.16"
brotli = "8"
bytes = "1"
cbc = "0.1"
chrono = { version = "0.4", features = ["serde"] }
@@ -135,7 +136,7 @@ tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "signal
tokio-util = { version = "0.7", features = ["codec", "io-util"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
uuid = { version = "1", features = ["serde", "v4", "v5"] }
uuid = { version = "1", features = ["serde", "v4", "v5", "v7"] }
webpki-roots = "0.26"
wreq = { version = "6.0.0-rc.28", default-features = false, features = ["json", "stream", "socks", "webpki-roots", "ws"] }
wreq-util = "3.0.0-rc.10"
README.md
+15 -1
View File
@@ -113,6 +113,18 @@ cd Aether
curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash
```
### Nightly（每日 main 构建）
Nightly workflow 每天从 `main` 的固定 commit 构建并发布滚动的 GitHub Release `nightly`，同时推送多架构 GHCR 镜像 `ghcr.io/fawney19/aether:nightly`。Nightly 是预发布版本，适合验证最新代码，不保证与正式版相同的稳定性。滚动 Release 需要仓库保持关闭 GitHub Release immutability。
安装最新 nightly（Linux systemd / macOS launchd + SQLite）：
```bash
curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash -s -- --channel nightly
```
Docker Compose 用户可在部署目录的 `.env` 中设置 `APP_IMAGE=ghcr.io/fawney19/aether:nightly`，然后运行 `./update.sh` 获取下一次 nightly。二进制方式可重新执行上述安装命令升级；当前管理后台的在线更新列表只跟踪正式版/RC/Beta，不会自动提示下一次 nightly。
## 本地开发
依赖 Docker、Rust toolchain、Node.js 和 make。
@@ -143,12 +155,14 @@ Aether Tunnel 是配套的正向代理节点，部署在海外 VPS 上，为墙
- Embeddings: [OpenAI compatible `POST /v1/embeddings`](docs/api/embeddings.md)
- Rerank: [OpenAI/Jina compatible `POST /v1/rerank`](docs/api/rerank.md)
- Responses WebSocket mode: [protocol and Aether behavior](docs/WebSocket-Mode.md)
- WebSocket probes: [Codex](docs/operations/codex-responses-websocket-probe.md) · [OpenAI Responses](docs/operations/openai-responses-websocket-probe.md)
## 环境变量
- `APP_PORT`：`aether-gateway` 唯一监听端口，固定绑定 `0.0.0.0:${APP_PORT}`
- `DATABASE_URL`：数据库连接串；SQLite 例如 `sqlite:///opt/aether/data/aether.db`，Postgres 例如 `postgresql://postgres:aether@postgres:5432/aether`
- `AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS` / `AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS`：数据库连接池手动覆盖值；未配置时会自动推导，SQLite 固定 `1/1`，Postgres/MySQL 按 CPU 核心数计算并默认封顶 `100`
- `AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS` / `AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS`：数据库连接池手动覆盖值；未配置时 SQLite 固定 `1/1`，Postgres/MySQL 按每核 `4` 条自动推导，总池范围为 `32-100`。该预算按进程计算，多实例部署应按数据库连接上限显式分配
- `AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS`：单实例请求并发上限；未配置时按 CPU 自动推导（基础范围 `512-65536`），低文件描述符预算时会进一步下调
- `AETHER_GATEWAY_REQUEST_BODY_BUFFER_BUDGET_MB`：单实例同时读取和解压请求体的加权内存预算，默认 `256MB`
- `AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS`：请求体完整读取超时，默认 `120000ms`
apps/aether-gateway/Cargo.toml
+2
View File
@@ -52,6 +52,7 @@ async-trait.workspace = true
axum = { version = "0.8", features = ["ws"] }
base64.workspace = true
bcrypt.workspace = true
brotli.workspace = true
bytes.workspace = true
chrono.workspace = true
chrono-tz.workspace = true
@@ -75,6 +76,7 @@ rsa = "0.9.10"
rustls.workspace = true
serde.workspace = true
serde_json.workspace = true
semver.workspace = true
sha1 = "0.10"
sha2 = { workspace = true, features = ["oid"] }
socket2.workspace = true
apps/aether-gateway/src/ai_serving/api.rs
+14 -10
View File
@@ -50,15 +50,16 @@ pub(crate) use aether_ai_formats::api::{
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec,
resolve_gemini_sync_spec, resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_openai_embedding_sync_spec, sanitize_request_path_and_query, AiControlPlanRequest,
CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter,
ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind, LocalOpenAiImageSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec, OpenAIChatClientEmitter,
OpenAIResponsesClientEmitter, StreamingStandardTerminalObserver, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CLI_STREAM_PLAN_KIND, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
resolve_openai_embedding_sync_spec, sanitize_request_path_and_query,
sanitize_request_query_string, AiControlPlanRequest, CanonicalContentPart,
CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter, ExecutionRuntimeAuthContext,
LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, OpenAIChatClientEmitter, OpenAIResponsesClientEmitter,
StreamingStandardTerminalObserver, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
@@ -68,7 +69,10 @@ pub(crate) use aether_ai_formats::api::{
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub(crate) use aether_ai_formats::protocol::stream::CanonicalUsage as StreamingCanonicalUsage;
pub(crate) use aether_ai_formats::CODEX_RESPONSES_LITE_HEADER;
/// Codex client identity headers re-exported for out-of-crate probe binaries,
/// which must reach `aether_ai_formats` through this seam.
pub use aether_ai_formats::{CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT};
pub(crate) use aether_ai_formats::{CODEX_RESPONSES_LITE_HEADER, UPSTREAM_IS_STREAM_KEY};
pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,
apps/aether-gateway/src/ai_serving/codex_context.rs
+249
View File
@@ -0,0 +1,249 @@
use std::sync::{Arc, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_provider_transport::CodexFingerprintConvergenceContext;
use http::{request::Parts, HeaderMap};
use serde_json::Value;
use uuid::Uuid;
use crate::client_session_affinity::codex_request_signals_from_request;
#[derive(Debug, Clone)]
pub(crate) struct CodexFingerprintContextSlot(Arc<OnceLock<CodexFingerprintConvergenceContext>>);
impl Default for CodexFingerprintContextSlot {
fn default() -> Self {
Self(Arc::new(OnceLock::new()))
}
}
impl CodexFingerprintContextSlot {
fn resolve(
&self,
headers: &HeaderMap,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
self.0
.get_or_init(|| {
build_codex_fingerprint_context(headers, body_json, Uuid::now_v7().to_string())
})
.clone()
}
}
pub(crate) fn resolve_codex_fingerprint_context(
parts: &Parts,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
if let Some(context) = parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.cloned()
{
return context;
}
if let Some(slot) = parts.extensions.get::<CodexFingerprintContextSlot>() {
return slot.resolve(&parts.headers, body_json);
}
build_codex_fingerprint_context(&parts.headers, body_json, Uuid::now_v7().to_string())
}
pub(crate) fn install_codex_fingerprint_context_slot(parts: &mut Parts) {
if parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.is_none()
&& parts
.extensions
.get::<CodexFingerprintContextSlot>()
.is_none()
{
parts
.extensions
.insert(CodexFingerprintContextSlot::default());
}
}
pub(crate) fn ensure_codex_fingerprint_context(
parts: &mut Parts,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
let context = resolve_codex_fingerprint_context(parts, body_json);
if parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.is_none()
{
parts.extensions.remove::<CodexFingerprintContextSlot>();
parts.extensions.insert(context.clone());
}
context
}
pub(crate) fn attach_codex_logical_turn_context(
parts: &mut Parts,
body_json: &Value,
logical_turn_id: &str,
) -> CodexFingerprintConvergenceContext {
let context =
build_codex_fingerprint_context(&parts.headers, body_json, logical_turn_id.to_string());
parts.extensions.remove::<CodexFingerprintContextSlot>();
parts.extensions.insert(context.clone());
context
}
pub(crate) fn restore_codex_logical_turn_context(
parts: &mut Parts,
context: &CodexFingerprintConvergenceContext,
) {
parts.extensions.remove::<CodexFingerprintContextSlot>();
parts.extensions.insert(context.clone());
}
fn build_codex_fingerprint_context(
headers: &HeaderMap,
body_json: &Value,
logical_turn_id: String,
) -> CodexFingerprintConvergenceContext {
let signals = codex_request_signals_from_request(headers, Some(body_json));
let mut context =
CodexFingerprintConvergenceContext::new(logical_turn_id, current_unix_millis());
if let Some(turn_id) = signals.turn_id {
context = context.with_original_turn_id(turn_id);
}
if let Some(session_id) = signals.thread_id.or(signals.session_id) {
context = context.with_original_client_session_id(session_id);
}
if let Some(prompt_cache_key) = signals.prompt_cache_key {
context = context.with_original_prompt_cache_key(prompt_cache_key);
}
context
}
fn current_unix_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
.try_into()
.unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use http::HeaderValue;
use serde_json::json;
use super::*;
#[test]
fn request_signals_are_captured_once_for_the_logical_turn() {
let request = http::Request::builder()
.header("thread-id", "header-thread")
.body(())
.expect("request should build");
let (mut parts, _) = request.into_parts();
let body = json!({
"prompt_cache_key": "client-cache",
"client_metadata": {
"turn_id": "client-turn",
"thread_id": "body-thread"
}
});
let context = attach_codex_logical_turn_context(&mut parts, &body, "logical-turn");
assert_eq!(context.logical_turn_id(), "logical-turn");
assert_eq!(context.original_turn_id(), Some("client-turn"));
assert_eq!(context.original_client_session_id(), Some("header-thread"));
assert_eq!(context.original_prompt_cache_key(), Some("client-cache"));
assert_eq!(
parts.extensions.get::<CodexFingerprintConvergenceContext>(),
Some(&context)
);
}
#[test]
fn restored_context_wins_over_retry_request_signals() {
let original = CodexFingerprintConvergenceContext::new("logical-turn", 1234)
.with_original_turn_id("original-turn")
.with_original_client_session_id("original-thread")
.with_original_prompt_cache_key("original-cache");
let request = http::Request::builder()
.body(())
.expect("request should build");
let (mut parts, _) = request.into_parts();
parts
.headers
.insert("thread-id", HeaderValue::from_static("retry-thread"));
restore_codex_logical_turn_context(&mut parts, &original);
let resolved = resolve_codex_fingerprint_context(
&parts,
&json!({
"prompt_cache_key": "retry-cache",
"client_metadata": {"turn_id": "retry-turn"}
}),
);
assert_eq!(resolved, original);
assert_eq!(resolved.turn_started_at_unix_ms(), 1234);
}
#[test]
fn generated_context_is_persisted_for_http_replanning() {
let request = http::Request::builder()
.header("session-id", "client-session")
.body(())
.expect("request should build");
let (mut parts, _) = request.into_parts();
let body = json!({
"prompt_cache_key": "client-cache",
"client_metadata": {"turn_id": "client-turn"}
});
let first = ensure_codex_fingerprint_context(&mut parts, &body);
let second = resolve_codex_fingerprint_context(
&parts,
&json!({
"prompt_cache_key": "retry-cache",
"client_metadata": {"turn_id": "retry-turn"}
}),
);
assert_eq!(second, first);
assert_eq!(second.original_turn_id(), Some("client-turn"));
assert_eq!(second.original_prompt_cache_key(), Some("client-cache"));
}
#[test]
fn installed_slot_reuses_context_across_cloned_parts() {
let request = http::Request::builder()
.body(())
.expect("request should build");
let (mut parts, _) = request.into_parts();
install_codex_fingerprint_context_slot(&mut parts);
let cloned_parts = parts.clone();
let first = resolve_codex_fingerprint_context(
&parts,
&json!({
"prompt_cache_key": "first-cache",
"client_metadata": {"turn_id": "first-turn"}
}),
);
let second = resolve_codex_fingerprint_context(
&cloned_parts,
&json!({
"prompt_cache_key": "second-cache",
"client_metadata": {"turn_id": "second-turn"}
}),
);
assert_eq!(second, first);
assert_eq!(second.original_turn_id(), Some("first-turn"));
assert_eq!(second.original_prompt_cache_key(), Some("first-cache"));
}
}
apps/aether-gateway/src/ai_serving/mod.rs
+10 -5
View File
@@ -1,5 +1,6 @@
mod adaptation;
pub(crate) mod api;
pub(crate) mod codex_context;
mod finalize;
mod planner;
mod pure;
@@ -23,6 +24,7 @@ pub(crate) use self::finalize::internal::{
maybe_build_sync_finalize_outcome, maybe_compile_sync_finalize_response,
SyncToStreamBridgeOutcome,
};
pub(crate) use self::planner::openai_responses_reasoning_replay_policy;
pub(crate) use self::planner::{
apply_local_runtime_candidate_terminal_reason, build_gemini_stream_plan_from_decision,
build_gemini_sync_plan_from_decision, build_local_gemini_files_stream_attempt_source_for_kind,
@@ -52,17 +54,20 @@ pub(crate) use self::planner::{
build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason,
codex_model_capabilities_for_transport, extract_pool_sticky_session_token,
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
planner_is_matching_stream_request, provider_key_pool_score_id, provider_key_pool_score_scope,
read_candidate_transport_snapshot, record_local_runtime_candidate_skip_reason,
maybe_build_pinned_stream_local_same_format_provider_decision_payload,
maybe_build_responses_websocket_decision, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id,
provider_key_pool_score_scope, read_candidate_transport_snapshot,
record_local_runtime_candidate_skip_reason, resolve_provider_chat_pii_redaction,
resolve_tunnel_scheduler_affinity_context, resolve_upstream_is_stream_for_provider,
set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, validate_final_openai_provider_request,
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate,
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalExecutionAttemptSource,
LocalExecutionCandidateKind, LocalResolvedOAuthRequestAuth, PlannerAppState,
SkippedLocalExecutionCandidate,
ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
ResponsesWebSocketPinnedCandidate, SkippedLocalExecutionCandidate,
};
pub(crate) use self::pure::*;
pub(crate) use self::response_history::{
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+42 -79
View File
@@ -47,13 +47,12 @@ use crate::cache::{
use crate::clock::current_unix_ms;
use crate::dispatch::refs::dispatch_ref_for_local_candidate;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
use crate::orchestration::{ExecutionAttemptIdentity, POOL_KEY_RETRY_INDEX_STRIDE};
use crate::scheduler::candidate::is_auth_api_key_concurrency_limit_skip_reason;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
const AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET: Duration = Duration::from_millis(100);
const AUTH_API_KEY_CONCURRENCY_RETRY_DELAY: Duration = Duration::from_millis(10);
@@ -481,10 +480,6 @@ where
type ExtraData = Value;
type Error = Infallible;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
local_attempt_slot_count(&candidate.transport)
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
available_candidate_extra_data_with_dispatch_ref(candidate, &self.build_extra_data)
}
@@ -1610,7 +1605,8 @@ async fn persist_available_local_execution_candidate_at_index<F>(
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
// Exactly one attempt is materialized per candidate; same-key retries are
// derived lazily by the attempt loop after a failure.
let extra_data = ai_candidate_extra_data_with_ranking(
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
candidate.ranking.as_ref(),
@@ -1625,53 +1621,34 @@ where
Some(candidate_index),
extra_data,
);
let should_persist = should_persist_available_local_candidate(&candidate);
let mut attempts = Vec::with_capacity(attempt_slots as usize);
let mut owned_candidate = Some(candidate);
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = if should_persist_available_local_candidate(&candidate) {
state
.persist_available_local_candidate(
trace_id,
context.user_id,
context.api_key_id,
&candidate.candidate,
candidate_index,
retry_index,
generated_candidate_id.as_str(),
context.required_capabilities,
extra_data,
current_unix_ms(),
context.error_context,
)
.await
} else {
generated_candidate_id
};
for retry_index in 0..attempt_slots {
let candidate_ref = owned_candidate
.as_ref()
.expect("candidate should remain available until final retry");
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = if should_persist {
state
.persist_available_local_candidate(
trace_id,
context.user_id,
context.api_key_id,
&candidate_ref.candidate,
candidate_index,
effective_retry_index(retry_index, candidate_ref.orchestration.pool_key_index),
generated_candidate_id.as_str(),
context.required_capabilities,
extra_data.clone(),
current_unix_ms(),
context.error_context,
)
.await
} else {
generated_candidate_id
};
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
candidate_ref.clone()
};
let retry_index =
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
attempts.push(LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id,
});
}
attempts
vec![LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id,
}]
}
fn available_candidate_extra_data_with_dispatch_ref<F>(
@@ -1840,6 +1817,7 @@ fn routing_trace_for_candidate(
CandidateKind::Provider => Some(candidate.key_id.clone()),
CandidateKind::PoolGroup => None,
},
api_format: Some(candidate.endpoint_api_format.clone()),
provider_priority: candidate.provider_priority,
key_priority: candidate
.key_global_priority_for_format
@@ -1923,32 +1901,15 @@ fn build_unpersisted_local_execution_candidate_attempts(
candidate: EligibleLocalExecutionCandidate,
candidate_index: u32,
) -> VecDeque<LocalExecutionCandidateAttempt> {
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
let mut attempts = VecDeque::with_capacity(attempt_slots as usize);
let mut owned_candidate = Some(candidate);
for retry_index in 0..attempt_slots {
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
owned_candidate
.as_ref()
.expect("candidate should remain available until final retry")
.clone()
};
let retry_index =
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
attempts.push_back(LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id: Uuid::new_v4().to_string(),
});
}
attempts
// One attempt per candidate; same-key retries are derived lazily by the
// attempt loop after a failure.
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
VecDeque::from([LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id: Uuid::new_v4().to_string(),
}])
}
async fn persist_pool_group_exhaustion_skipped_candidate(
@@ -2276,6 +2237,8 @@ mod tests {
pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch: None,
// These tests cover persistence shape, not same-key retries.
sticky_key_attempts: Some(1),
},
ranking: None,
}
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
+30 -31
View File
@@ -3,7 +3,7 @@ use aether_ai_serving::{
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
AiRankingSchedulingMode,
};
use aether_routing_core::{ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode};
use aether_routing_core::ResolvedRoutingPolicy;
use async_trait::async_trait;
use tokio::sync::Mutex;
use tracing::warn;
@@ -184,35 +184,16 @@ fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedul
}
}
/// Ordering config for a request. A resolved routing policy is authoritative
/// and is never merged with legacy system-config values; without a policy the
/// effective default (system-default routing group, then legacy keys) applies.
pub(crate) async fn scheduler_ordering_config_for_routing_policy(
state: PlannerAppState<'_>,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> SchedulerOrderingConfig {
let system_config = read_scheduler_ordering_config_or_default(state).await;
match routing_policy {
Some(policy) => {
let mut config = scheduler_ordering_config_from_routing_policy(policy);
config.keep_priority_on_conversion |= system_config.keep_priority_on_conversion;
config
}
None => system_config,
}
}
fn scheduler_ordering_config_from_routing_policy(
policy: &ResolvedRoutingPolicy,
) -> SchedulerOrderingConfig {
SchedulerOrderingConfig {
priority_mode: match policy.priority_mode {
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
},
scheduling_mode: match policy.scheduling_mode {
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
},
keep_priority_on_conversion: policy.keep_priority_on_conversion,
Some(policy) => SchedulerOrderingConfig::from_routing_policy(policy),
None => read_scheduler_ordering_config_or_default(state).await,
}
}
@@ -231,14 +212,26 @@ fn routing_overlaid_candidate(
let overlaid_key_priority = match kind {
LocalExecutionCandidateKind::SingleKey => policy
.ranking_overlay
.key_priority_overrides
.get(candidate.key_id.as_str()),
.key_priority_override_matching_format(candidate.key_id.as_str(), |format| {
crate::ai_serving::api_format_alias_matches(
format,
candidate.endpoint_api_format.as_str(),
)
})
.or_else(|| {
policy
.ranking_overlay
.key_priority_overrides
.get(candidate.key_id.as_str())
.copied()
}),
LocalExecutionCandidateKind::PoolGroup => policy
.ranking_overlay
.pool_priority_overrides
.get(candidate.provider_id.as_str()),
.get(candidate.provider_id.as_str())
.copied(),
};
if let Some(overlaid_key_priority) = overlaid_key_priority.copied() {
if let Some(overlaid_key_priority) = overlaid_key_priority {
overlaid.key_internal_priority = overlaid_key_priority;
overlaid.key_global_priority_for_format = Some(overlaid_key_priority);
}
@@ -378,6 +371,7 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: aether_routing_core::RankingOverlay::default(),
mutation_plan: Default::default(),
pool_policy_overrides: BTreeMap::new(),
@@ -396,7 +390,7 @@ mod tests {
}
#[tokio::test]
async fn routing_policy_inherits_global_conversion_priority_override() {
async fn routing_policy_ignores_legacy_global_conversion_priority_override() {
let data_state = GatewayDataState::default().with_system_config_values_for_tests([(
"keep_priority_on_conversion".to_string(),
json!(true),
@@ -413,6 +407,7 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -429,7 +424,10 @@ mod tests {
ordering.scheduling_mode,
crate::scheduler::config::SchedulerSchedulingMode::FixedOrder
);
assert!(ordering.keep_priority_on_conversion);
assert!(
!ordering.keep_priority_on_conversion,
"a resolved routing policy must not inherit the legacy system-config flag"
);
}
#[test]
@@ -447,6 +445,7 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::GlobalKey,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: aether_routing_core::RankingOverlay {
pool_priority_overrides: BTreeMap::from([("provider-1".to_string(), 4)]),
key_priority_overrides: BTreeMap::from([("representative-key".to_string(), 1)]),
apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs
+13 -1
View File
@@ -23,7 +23,9 @@ use crate::ai_serving::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::stage_metrics::observe_gateway_stage_ms;
use super::candidate_ranking::rank_eligible_local_execution_candidates;
use super::candidate_ranking::{
rank_eligible_local_execution_candidates, scheduler_ordering_config_for_routing_policy,
};
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
@@ -378,8 +380,18 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
"candidate_resolution_core",
started_at.elapsed().as_millis() as u64,
);
let sticky_key_attempts = if outcome.eligible_candidates.is_empty() {
None
} else {
Some(
scheduler_ordering_config_for_routing_policy(state, routing_policy)
.await
.sticky_key_attempts,
)
};
for candidate in &mut outcome.eligible_candidates {
candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch);
candidate.orchestration.sticky_key_attempts = sticky_key_attempts;
}
(outcome.eligible_candidates, outcome.skipped_candidates)
}
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
+99 -13
View File
@@ -86,6 +86,52 @@ impl GatewayLocalCandidatePreselectionPort<'_> {
}
}
/// A Responses compaction request carries the OpenAI-only `compaction_trigger`
/// control item. It must stay on an OpenAI Responses endpoint: treating it as
/// an ordinary cross-format request would make Gemini/Claude candidates look
/// eligible and defer the inevitable lossy-conversion failure until payload
/// construction.
fn request_candidate_api_formats_for_operation(
client_api_format: &str,
require_streaming: bool,
request_operation: Option<&str>,
) -> Vec<String> {
let candidate_api_formats =
crate::ai_serving::request_candidate_api_formats(client_api_format, require_streaming)
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
restrict_candidate_api_formats_for_operation(
client_api_format,
request_operation,
candidate_api_formats,
)
}
fn restrict_candidate_api_formats_for_operation(
client_api_format: &str,
request_operation: Option<&str>,
candidate_api_formats: Vec<String>,
) -> Vec<String> {
let is_responses_compaction = request_operation.is_some_and(|operation| {
operation.eq_ignore_ascii_case(crate::ai_serving::OPENAI_RESPONSES_OPERATION_COMPACT)
});
let is_standard_responses_client =
crate::ai_serving::normalize_api_format_alias(client_api_format) == "openai:responses";
if !(is_responses_compaction && is_standard_responses_client) {
return candidate_api_formats;
}
candidate_api_formats
.into_iter()
.filter(|candidate_api_format| {
crate::ai_serving::normalize_api_format_alias(candidate_api_format)
== "openai:responses"
})
.collect()
}
#[async_trait]
impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
type Candidate = SchedulerMinimalCandidateSelectionCandidate;
@@ -128,6 +174,8 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
self.ranking_seed,
false,
self.request_operation,
self.routing_policy
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
@@ -219,11 +267,11 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
>,
GatewayError,
> {
let candidate_api_formats =
crate::ai_serving::request_candidate_api_formats(client_api_format, require_streaming)
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
let candidate_api_formats = request_candidate_api_formats_for_operation(
client_api_format,
require_streaming,
request_operation,
);
preselect_local_execution_candidates_for_api_formats_with_serving(
state,
model_directive_policy,
@@ -264,6 +312,11 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se
>,
GatewayError,
> {
let candidate_api_formats = restrict_candidate_api_formats_for_operation(
client_api_format,
request_operation,
candidate_api_formats,
);
let model_directive_routing_models = resolve_model_directive_routing_models(
model_directive_policy,
&candidate_api_formats,
@@ -362,11 +415,11 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
allow_priority_page_cache: bool,
trace_id: Option<&str>,
) -> Self {
let candidate_api_formats =
crate::ai_serving::request_candidate_api_formats(client_api_format, require_streaming)
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
let candidate_api_formats = request_candidate_api_formats_for_operation(
client_api_format,
require_streaming,
request_operation,
);
let model_directive_routing_models = resolve_model_directive_routing_models(
model_directive_policy,
&candidate_api_formats,
@@ -1240,6 +1293,9 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.then_some(self.client_session_affinity.as_ref())
.flatten(),
self.ranking_seed,
self.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let skipped_candidates = skipped_candidates
@@ -1437,6 +1493,31 @@ mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
#[test]
fn compaction_operation_excludes_non_responses_provider_formats() {
assert_eq!(
request_candidate_api_formats_for_operation("openai:responses", true, Some("compact"),),
vec!["openai:responses"]
);
assert_eq!(
request_candidate_api_formats_for_operation("openai:responses", true, None),
vec![
"openai:responses",
"openai:chat",
"claude:messages",
"gemini:generate_content"
]
);
assert_eq!(
request_candidate_api_formats_for_operation(
"openai:responses:compact",
false,
Some("compact"),
),
vec!["openai:responses:compact"]
);
}
#[derive(Default)]
struct EmptyFallbackCountingRepository {
fallback_reads: AtomicUsize,
@@ -1808,6 +1889,7 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -1871,6 +1953,7 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -2581,14 +2664,16 @@ mod tests {
candidate_repository,
)
.with_encryption_key_for_tests("development-key")
// Legacy keys deliberately disagree with the routing policy: the
// resolved policy must be the only source of scheduler ordering.
.with_system_config_values_for_tests([
(
"scheduling_mode".to_string(),
serde_json::json!("fixed_order"),
serde_json::json!("cache_affinity"),
),
(
"keep_priority_on_conversion".to_string(),
serde_json::json!(true),
serde_json::json!(false),
),
]);
let app = AppState::new()
@@ -2605,7 +2690,8 @@ mod tests {
resolved_model: "gpt-5.4-mini".to_string(),
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
keep_priority_on_conversion: true,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
+362 -6
View File
@@ -13,6 +13,7 @@ use http::{HeaderMap, HeaderName, HeaderValue};
use serde_json::{json, Value};
use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::transport::CodexFingerprintConvergenceContext;
use crate::ai_serving::{
ClientSurface, ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot,
GatewayCredentialCarrier, GatewayProviderTransportSnapshot, PlannerAppState,
@@ -55,6 +56,7 @@ pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) client_surface: Option<ClientSurface>,
pub(crate) gateway_credential_carrier: Option<GatewayCredentialCarrier>,
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
pub(crate) codex_fingerprint_context: Option<CodexFingerprintConvergenceContext>,
pub(crate) routing_policy: Option<ResolvedRoutingPolicy>,
pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>,
pub(crate) routing_context: Option<LocalRoutingRequestContext>,
@@ -100,6 +102,22 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
input: &LocalRequestedModelDecisionInput,
decision: &mut AiExecutionDecision,
transport: Option<&GatewayProviderTransportSnapshot>,
) -> Result<(), GatewayError> {
apply_provider_request_routing_policy_to_decision_with_websocket_mode(
input, decision, transport, false,
)
}
/// Applies provider-request routing mutations while retaining the transport
/// boundary of a pinned Responses WebSocket continuation. Routing rules may
/// mutate the body and therefore require a second provider-contract pass; the
/// pass must use the same explicit continuation mode as the first pass rather
/// than guessing from JSON fields.
pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_mode(
input: &LocalRequestedModelDecisionInput,
decision: &mut AiExecutionDecision,
transport: Option<&GatewayProviderTransportSnapshot>,
websocket_continuation: bool,
) -> Result<(), GatewayError> {
let provider_api_format = decision
.provider_api_format
@@ -150,6 +168,12 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
provider_api_format.as_str(),
);
}
apply_codex_fingerprint_convergence_to_decision(
input,
decision,
transport,
provider_api_format.as_str(),
);
return Ok(());
};
let provider_body_rules = decision
@@ -207,6 +231,12 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
provider_api_format.as_str(),
);
}
apply_codex_fingerprint_convergence_to_decision(
input,
decision,
transport,
provider_api_format.as_str(),
);
return Ok(());
}
if original_provider_request_body.is_none() && !policy.mutation_plan.body_patch.is_empty() {
@@ -244,9 +274,8 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
input.requested_model.as_str(),
)
});
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
&mut provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
{
let finalization = crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: context.client_api_format.as_str(),
provider_api_format: provider_api_format.as_str(),
provider_type: provider_type.as_str(),
@@ -257,9 +286,32 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
require_body_stream_field: original_provider_request_body
.as_ref()
.is_some_and(|body| body.get("stream").is_some()),
},
model_capabilities.as_ref(),
)
};
let reasoning_replay_policy = transport
.map(|transport| {
crate::ai_serving::openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
provider_model.as_str(),
)
})
.unwrap_or_default();
if websocket_continuation {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation(
&mut provider_request_body,
finalization,
model_capabilities.as_ref(),
reasoning_replay_policy,
)
} else {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut provider_request_body,
finalization,
model_capabilities.as_ref(),
reasoning_replay_policy,
)
}
}
.map_err(|violation| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("routing provider_request violates provider contract: {violation:?}"),
@@ -305,13 +357,51 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
if original_provider_request_body.is_some() {
decision.provider_request_body = Some(provider_request_body);
}
apply_codex_fingerprint_convergence_to_decision(
input,
decision,
transport,
provider_api_format.as_str(),
);
update_report_context_provider_request_mutation(decision, &policy);
Ok(())
}
fn apply_codex_fingerprint_convergence_to_decision(
input: &LocalRequestedModelDecisionInput,
decision: &mut AiExecutionDecision,
transport: Option<&GatewayProviderTransportSnapshot>,
provider_api_format: &str,
) {
let (Some(transport), Some(provider_request_body)) =
(transport, decision.provider_request_body.as_mut())
else {
return;
};
let Some(context) = input.codex_fingerprint_context.as_ref() else {
return;
};
let applied = crate::ai_serving::transport::apply_codex_fingerprint_convergence_with_context(
transport,
provider_api_format,
context,
&mut decision.provider_request_headers,
provider_request_body,
);
if applied {
decision.prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
}
}
struct GatewayAuthenticatedDecisionInputPort<'a> {
state: PlannerAppState<'a>,
now_unix_secs: u64,
auth_snapshot_override: Option<GatewayAuthApiKeySnapshot>,
model_directive_policy: &'a crate::system_features::ModelDirectivePolicySnapshot,
model_directive_base_model: Option<String>,
}
@@ -328,6 +418,17 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<
&self,
auth_context: &Self::AuthContext,
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
if let Some(snapshot) = self.auth_snapshot_override.as_ref() {
if snapshot.user_id != auth_context.user_id
|| snapshot.api_key_id != auth_context.api_key_id
{
return Err(GatewayError::Internal(
"WebSocket auth snapshot identity does not match its control decision"
.to_string(),
));
}
return Ok(Some(snapshot.clone()));
}
self.state
.read_auth_api_key_snapshot(
&auth_context.user_id,
@@ -383,6 +484,7 @@ pub(crate) fn build_local_requested_model_decision_input(
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
@@ -397,6 +499,8 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
body_json: &Value,
client_api_format: &str,
) -> Result<(), GatewayError> {
input.codex_fingerprint_context =
Some(crate::ai_serving::codex_context::resolve_codex_fingerprint_context(parts, body_json));
let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER);
let selected_group = match state.routing_group_read_repository() {
Some(repository) => {
@@ -708,6 +812,27 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
requested_model_api_format: Option<&str>,
explicit_required_capabilities: Option<&serde_json::Value>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
resolve_local_authenticated_decision_input_with_snapshot(
state,
auth_context,
None,
requested_model,
requested_model_api_format,
explicit_required_capabilities,
model_directive_policy,
)
.await
}
pub(crate) async fn resolve_local_authenticated_decision_input_with_snapshot(
state: &AppState,
auth_context: ExecutionRuntimeAuthContext,
auth_snapshot_override: Option<GatewayAuthApiKeySnapshot>,
requested_model: Option<&str>,
requested_model_api_format: Option<&str>,
explicit_required_capabilities: Option<&serde_json::Value>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
let model_directive_base_model = match (requested_model, requested_model_api_format) {
(Some(model), Some(api_format)) => model_directive_policy
@@ -719,6 +844,7 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
let port = GatewayAuthenticatedDecisionInputPort {
state: PlannerAppState::new(state),
now_unix_secs: current_unix_secs(),
auth_snapshot_override,
model_directive_policy,
model_directive_base_model,
};
@@ -947,6 +1073,7 @@ fn ensure_report_context_routing_trace(
endpoint_id: decision.endpoint_id.clone().unwrap_or_default(),
model_id,
key_id,
api_format: decision.provider_api_format.clone(),
provider_priority,
key_priority,
},
@@ -1024,6 +1151,52 @@ mod tests {
}
}
#[tokio::test]
async fn explicit_auth_snapshot_override_does_not_fall_back_to_the_planner_cache() {
// AppState::new has no auth snapshot repository. Without the explicit
// override this resolver returns None; a WebSocket strong snapshot must
// therefore be the exact value used to build the planner input.
let state = AppState::new().expect("test state should build");
let mut strong_snapshot = sample_auth_snapshot();
strong_snapshot.api_key_allowed_models = Some(vec!["gpt-live-only".to_string()]);
let resolved = resolve_local_authenticated_decision_input_with_snapshot(
&state,
sample_auth_context(),
Some(strong_snapshot.clone()),
Some("gpt-live-only"),
Some("openai:responses"),
None,
&Default::default(),
)
.await
.expect("snapshot override should resolve")
.expect("the explicit snapshot should replace the missing cached value");
assert_eq!(resolved.auth_snapshot, strong_snapshot);
}
#[tokio::test]
async fn explicit_auth_snapshot_override_rejects_an_identity_mismatch() {
let state = AppState::new().expect("test state should build");
let mut wrong_snapshot = sample_auth_snapshot();
wrong_snapshot.api_key_id = "another-key".to_string();
let error = resolve_local_authenticated_decision_input_with_snapshot(
&state,
sample_auth_context(),
Some(wrong_snapshot),
Some("gpt-live-only"),
Some("openai:responses"),
None,
&Default::default(),
)
.await
.expect_err("a snapshot for another API key must never be injected");
assert!(matches!(error, GatewayError::Internal(message) if message.contains("identity")));
}
#[tokio::test]
async fn explicit_routing_attachment_authorizes_and_caches_per_principal() {
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
@@ -1154,6 +1327,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
model_directive_policy: Default::default(),
@@ -1312,6 +1486,57 @@ mod tests {
}
}
fn sample_codex_fingerprint_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_codex_transport_with_card();
transport.provider.config = Some(json!({
"codex": {"fingerprint_convergence_enabled": true}
}));
transport.endpoint.api_format = "openai:responses".to_string();
transport.endpoint.endpoint_kind = Some("responses".to_string());
transport.key.api_formats = Some(vec!["openai:responses".to_string()]);
transport.key.decrypted_auth_config =
Some(json!({"account_id": "account-codex-1"}).to_string());
transport
}
fn sample_codex_fingerprint_decision() -> AiExecutionDecision {
let prompt_cache_key = "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3";
let mut decision = sample_decision();
decision.provider_type = Some("codex".to_string());
decision.provider_api_format = Some("openai:responses".to_string());
decision.client_api_format = Some("openai:responses".to_string());
decision.provider_request_headers.extend([
("session-id".to_string(), "spoofed-session".to_string()),
("thread-id".to_string(), "spoofed-thread".to_string()),
(
"x-codex-turn-metadata".to_string(),
json!({
"installation_id": "spoofed-installation",
"session_id": "spoofed-session",
"thread_source": "cli"
})
.to_string(),
),
]);
decision.provider_request_body = Some(json!({
"model": "gpt-5",
"input": [],
"metadata": {},
"prompt_cache_key": prompt_cache_key,
"client_metadata": {
"session_id": "spoofed-session",
"thread_id": "spoofed-thread",
"caller": "sdk",
"x-codex-turn-metadata": json!({
"installation_id": "spoofed-installation",
"session_id": "spoofed-session",
"sandbox": "workspace-write"
}).to_string()
}
}));
decision
}
fn set_provider_request_rules(
input: &mut LocalRequestedModelDecisionInput,
allowed_models: &[&str],
@@ -1351,6 +1576,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
model_directive_policy: Default::default(),
@@ -1420,6 +1646,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
@@ -1488,6 +1715,134 @@ mod tests {
);
}
#[test]
fn codex_fingerprint_convergence_runs_at_every_provider_routing_success_exit() {
let transport = sample_codex_fingerprint_transport();
let mut no_context = sample_decision_input();
no_context.routing_context = None;
let mut empty_mutation = sample_decision_input();
empty_mutation
.routing_context
.as_mut()
.expect("routing context")
.group_config_json = json!({
"allowed_models": ["gpt-5"],
"rules": []
});
let mut with_mutation = sample_decision_input();
for input in [&mut no_context, &mut empty_mutation, &mut with_mutation] {
input.codex_fingerprint_context = Some(
CodexFingerprintConvergenceContext::new(
uuid::Uuid::new_v4().to_string(),
1_756_668_000_000,
)
.with_original_client_session_id("client-session-1".to_string()),
);
}
let mut stable_identity = None;
let mut turn_ids = std::collections::BTreeSet::new();
for (exit_name, input) in [
("no_context", no_context),
("empty_mutation", empty_mutation),
("with_mutation", with_mutation),
] {
let mut decision = sample_codex_fingerprint_decision();
apply_provider_request_routing_policy_to_decision(
&input,
&mut decision,
Some(&transport),
)
.unwrap_or_else(|error| panic!("{exit_name} should converge: {error:?}"));
let session_id = decision.provider_request_headers["session-id"].clone();
let thread_id = decision.provider_request_headers["thread-id"].clone();
let installation_id =
decision.provider_request_headers["x-codex-installation-id"].clone();
let window_id = decision.provider_request_headers["x-codex-window-id"].clone();
assert_eq!(decision.provider_request_headers["session_id"], session_id);
assert_eq!(
decision.provider_request_headers["x-client-request-id"],
thread_id
);
assert_eq!(window_id, format!("{thread_id}:0"));
assert_eq!(
uuid::Uuid::parse_str(&session_id)
.expect("session UUID")
.get_version_num(),
4
);
assert_eq!(
uuid::Uuid::parse_str(&thread_id)
.expect("thread UUID")
.get_version_num(),
4
);
let body = decision
.provider_request_body
.as_ref()
.expect("request body");
assert_eq!(
decision.prompt_cache_key.as_deref(),
body.get("prompt_cache_key").and_then(Value::as_str)
);
assert_eq!(
body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
);
assert_eq!(body["client_metadata"]["session_id"], session_id);
assert_eq!(body["client_metadata"]["thread_id"], thread_id);
assert_eq!(body["client_metadata"]["caller"], "sdk");
assert_eq!(
body["client_metadata"]["x-codex-installation-id"],
installation_id
);
assert_eq!(body["client_metadata"]["x-codex-window-id"], window_id);
let header_metadata: Value =
serde_json::from_str(&decision.provider_request_headers["x-codex-turn-metadata"])
.expect("header turn metadata");
let body_metadata: Value = serde_json::from_str(
body["client_metadata"]["x-codex-turn-metadata"]
.as_str()
.expect("embedded turn metadata"),
)
.expect("embedded turn metadata JSON");
assert_eq!(header_metadata["thread_source"], "cli");
assert_eq!(body_metadata["sandbox"], "workspace-write");
assert_eq!(
header_metadata["turn_id"],
body["client_metadata"]["turn_id"]
);
assert_eq!(body_metadata["turn_id"], body["client_metadata"]["turn_id"]);
assert_eq!(
header_metadata["turn_started_at_unix_ms"],
body_metadata["turn_started_at_unix_ms"]
);
let turn_id = body["client_metadata"]["turn_id"]
.as_str()
.expect("turn ID")
.to_string();
assert_eq!(
uuid::Uuid::parse_str(&turn_id)
.expect("turn UUID")
.get_version_num(),
7
);
turn_ids.insert(turn_id);
let identity = (installation_id, session_id, thread_id);
if let Some(expected) = stable_identity.as_ref() {
assert_eq!(&identity, expected, "identity changed at {exit_name}");
} else {
stable_identity = Some(identity);
}
}
assert_eq!(turn_ids.len(), 3, "each request needs a fresh turn ID");
}
#[test]
fn provider_request_routing_policy_cannot_restore_credentials_or_aether_internal_headers() {
for header_name in [
@@ -1637,6 +1992,7 @@ mod tests {
let body = decision.provider_request_body.as_ref().expect("body");
assert!(body.get("prompt_cache_key").is_none());
assert!(decision.prompt_cache_key.is_none());
assert!(body.get("client_metadata").is_none());
assert!(!decision.provider_request_headers.contains_key("session-id"));
assert!(!decision.provider_request_headers.contains_key("thread-id"));
apps/aether-gateway/src/ai_serving/planner/mod.rs
+6 -2
View File
@@ -38,6 +38,7 @@ pub(crate) use self::common::resolve_upstream_is_stream_for_provider;
pub(crate) use self::passthrough::{
build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports,
build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports,
maybe_build_pinned_stream_local_same_format_provider_decision_payload,
};
pub(crate) use self::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
@@ -49,6 +50,7 @@ pub(crate) use self::plan_builders::{
pub(crate) use self::pool_scores::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
};
pub(crate) use self::redaction::resolve_provider_chat_pii_redaction;
pub(crate) use self::request_gzip::resolve_transport_request_encoding_policy;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
@@ -80,8 +82,10 @@ pub(crate) use self::standard::{
build_local_stream_plan_and_reports as build_standard_family_stream_plan_and_reports,
build_local_sync_attempt_source as build_standard_family_sync_attempt_source,
build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports,
codex_model_capabilities_for_transport, set_local_openai_chat_execution_exhausted_diagnostic,
validate_final_openai_provider_request,
codex_model_capabilities_for_transport, maybe_build_responses_websocket_decision,
openai_responses_reasoning_replay_policy, set_local_openai_chat_execution_exhausted_diagnostic,
validate_final_openai_provider_request, ResponsesWebSocketBodyNormalization,
ResponsesWebSocketDecision, ResponsesWebSocketPinnedCandidate,
};
pub(crate) use self::state::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs
+1
View File
@@ -8,6 +8,7 @@ pub(crate) use self::provider::{
build_local_sync_attempt_source as build_local_same_format_sync_attempt_source,
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
maybe_build_pinned_stream_local_same_format_provider_decision_payload,
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
};
apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs
+1
View File
@@ -69,6 +69,7 @@ pub(crate) use self::family::{
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
pub(crate) use self::family::{
maybe_build_pinned_stream_local_same_format_provider_decision_payload,
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
};
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs
+27
View File
@@ -89,6 +89,21 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
maybe_build_pinned_stream_local_same_format_provider_decision_payload(
state, parts, trace_id, decision, body_json, plan_kind, None,
)
.await
}
pub(crate) async fn maybe_build_pinned_stream_local_same_format_provider_decision_payload(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
pinned_candidate: Option<(&str, &str, &str)>,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let Some(spec) = resolve_stream_spec(plan_kind) else {
return Ok(None);
@@ -141,6 +156,18 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
)
.await?
{
if pinned_candidate.is_some_and(|(provider_id, endpoint_id, key_id)| {
payload.provider_id.as_deref() != Some(provider_id)
|| payload.endpoint_id.as_deref() != Some(endpoint_id)
|| payload.key_id.as_deref() != Some(key_id)
}) {
crate::orchestration::release_pool_key_lease_from_report_context(
state,
payload.report_context.as_ref(),
)
.await;
continue;
}
return Ok(Some(payload));
}
}
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs
+17 -6
View File
@@ -26,6 +26,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_api_request;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
use super::{
@@ -140,6 +141,10 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -167,20 +172,21 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
.collect(),
LocalCandidateResolutionMode::Standard,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: spec_metadata.api_format,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
spec_metadata.api_format,
provider_api_format.as_str(),
))
},
|mut skipped_candidate| {
@@ -245,6 +251,10 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
@@ -273,20 +283,21 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
.collect(),
LocalCandidateResolutionMode::Standard,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: spec_metadata.api_format,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
spec_metadata.api_format,
provider_api_format.as_str(),
))
},
|mut skipped_candidate| {
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs
+1
View File
@@ -4,6 +4,7 @@ mod payload;
mod request;
pub(crate) use self::build::{
maybe_build_pinned_stream_local_same_format_provider_decision_payload,
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
};
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
+9 -3
View File
@@ -55,8 +55,6 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
..
} = &attempt;
let candidate = &eligible.candidate;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let Some(resolved) = resolve_local_same_format_provider_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
@@ -164,6 +162,10 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
}
}
let provider_api_format = resolved.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
provider_api_format.as_str(),
);
let effective_headers = input.effective_headers(&parts.headers);
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
@@ -201,13 +203,17 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
upstream_is_stream: resolved.upstream_is_stream,
has_envelope: resolved.is_kiro || resolved.is_antigravity || resolved.is_gemini_cli,
needs_conversion: false,
needs_conversion: matches!(
conversion_mode,
crate::ai_serving::ConversionMode::Bidirectional
),
extra_fields,
}),
execution_strategy,
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
+10 -2
View File
@@ -42,7 +42,7 @@ use super::{
};
use crate::ai_serving::planner::standard::{
codex_model_capabilities_for_transport, openai_provider_request_contract_failure_extra_data,
same_format_provider_request_body_failure_extra_data,
openai_responses_reasoning_replay_policy, same_format_provider_request_body_failure_extra_data,
};
pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trace(
@@ -180,12 +180,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
}
};
let effective_headers = input.effective_headers(&parts.headers);
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
prepared.transport.provider.provider_type.as_str(),
prepared.transport.endpoint.base_url.as_str(),
prepared.mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec.api_format,
reasoning_replay_policy,
&attempt.candidate_id,
)
.await?;
@@ -205,6 +211,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
prepared.kiro_auth.as_ref(),
prepared.is_claude_code,
false,
reasoning_replay_policy,
)
else {
mark_skipped_local_same_format_provider_candidate_with_extra_data(
@@ -269,7 +276,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
source_model,
);
if let Err(violation) =
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut base_provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: spec.api_format,
@@ -285,6 +292,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
),
},
codex_model_capabilities.as_ref(),
reasoning_replay_policy,
)
{
mark_skipped_local_same_format_provider_candidate_with_extra_data(
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/body.rs
+3 -1
View File
@@ -3,7 +3,7 @@ use serde_json::Value;
use super::super::LocalSameFormatProviderSpec;
use crate::ai_serving::transport::{
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
build_same_format_provider_request_body_with_compatibility_report as build_same_format_provider_request_body_with_compatibility_report_impl,
build_same_format_provider_request_body_with_compatibility_report_and_reasoning_replay_policy as build_same_format_provider_request_body_with_compatibility_report_impl,
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
SameFormatProviderRequestBodyOutput,
};
@@ -50,6 +50,7 @@ pub(crate) fn build_same_format_provider_request_body_with_compatibility_report(
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
is_claude_code: bool,
enable_model_directives: bool,
reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
) -> Option<SameFormatProviderRequestBodyOutput> {
build_same_format_provider_request_body_with_compatibility_report_impl(
SameFormatProviderRequestBodyInput {
@@ -67,6 +68,7 @@ pub(crate) fn build_same_format_provider_request_body_with_compatibility_report(
is_claude_code,
enable_model_directives,
},
reasoning_replay_policy,
)
}
apps/aether-gateway/src/ai_serving/planner/redaction.rs
+76 -6
View File
@@ -61,12 +61,41 @@ pub(crate) fn request_identity_response_encoding_when_redacted(
}
}
/// Removes credential-bearing URL components before attaching an upstream URL
/// to a diagnostic event. Endpoint query parameters remain untouched on the
/// wire, but they can contain API keys or signed tokens and must not reach
/// logs.
pub(crate) fn sanitize_upstream_url_for_log(raw: &str) -> String {
if let Ok(mut url) = url::Url::parse(raw) {
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return "<invalid-upstream-url>".to_string();
}
let _ = url.set_username("");
let _ = url.set_password(None);
url.set_query(None);
url.set_fragment(None);
return url.to_string();
}
let suffix_offset = raw
.char_indices()
.find_map(|(offset, character)| matches!(character, '?' | '#').then_some(offset))
.unwrap_or(raw.len());
let path = &raw[..suffix_offset];
if path.starts_with('/') && !path.starts_with("//") && !path.contains('@') {
path.to_string()
} else {
"<invalid-upstream-url>".to_string()
}
}
pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
state: &AppState,
parts: &http::request::Parts,
body_json: &'a Value,
auth_context: &ExecutionRuntimeAuthContext,
client_api_format: &str,
reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
candidate_id: &str,
) -> Result<ProviderRequestRedaction<'a>, GatewayError> {
let Some(format) = ChatPiiRedactionRequestFormat::from_api_format(client_api_format) else {
@@ -75,7 +104,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
let Some(slot) = parts.extensions.get::<RedactionSessionSlot>() else {
return Ok(ProviderRequestRedaction::disabled(body_json));
};
let request_cache_key = request_redaction_cache_key(format, body_json);
let request_cache_key = request_redaction_cache_key(format, reasoning_replay_policy, body_json);
if let Some(cached) = slot.cached_request_redaction(&request_cache_key) {
crate::stage_metrics::record_chat_pii_redaction_request_cache_hit();
observe_gateway_stage_ms("chat_pii_redaction_request_cache_hit", 0);
@@ -132,7 +161,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
body_json,
format,
build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs),
MaskChatRequestOptions::runtime(),
MaskChatRequestOptions::runtime().with_reasoning_replay_policy(reasoning_replay_policy),
Some(&cache),
)
.await
@@ -165,8 +194,16 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
})
}
fn request_redaction_cache_key(format: ChatPiiRedactionRequestFormat, body_json: &Value) -> String {
format!("{format:?}:{:p}", body_json)
fn request_redaction_cache_key(
format: ChatPiiRedactionRequestFormat,
reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
body_json: &Value,
) -> String {
// A request may be attempted against providers with different replay
// contracts. Reusing a cached DeepSeek opaque decision for an ordinary
// Responses candidate (or vice versa) would either bypass masking or
// corrupt provider-owned continuation state.
format!("{format:?}:{reasoning_replay_policy:?}:{:p}", body_json)
}
fn provider_redaction_from_cached<'a>(
@@ -216,14 +253,47 @@ async fn resolve_chat_pii_redaction_feature_settings(
}
fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayError {
match error {}
match error {
RedactionMaskError::SensitiveOpaqueReasoningState => {
warn!("gateway rejected provider-bound reasoning state containing sensitive text");
GatewayError::Client {
status: http::StatusCode::BAD_REQUEST,
message: "provider-bound reasoning state contains sensitive text and cannot be safely replayed while chat PII redaction is enabled".to_string(),
}
}
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::ChatPiiRedactionFeatureSettings;
use super::{sanitize_upstream_url_for_log, ChatPiiRedactionFeatureSettings};
#[test]
fn upstream_url_log_projection_removes_all_credential_carriers() {
assert_eq!(
sanitize_upstream_url_for_log(
"https://user:password@api.example.test/v1/responses?api-version=2026-08-01&token=secret#fragment"
),
"https://api.example.test/v1/responses"
);
assert_eq!(
sanitize_upstream_url_for_log("/v1/responses?key=secret#fragment"),
"/v1/responses"
);
for invalid in [
"https://user:secret@invalid host/v1/responses",
"//user:secret@api.example.test/v1/responses?token=hidden",
"not-a-url?token=hidden",
"data:text/plain,secret",
] {
assert_eq!(
sanitize_upstream_url_for_log(invalid),
"<invalid-upstream-url>"
);
}
}
#[test]
fn chat_pii_redaction_feature_settings_only_control_enablement() {
apps/aether-gateway/src/ai_serving/planner/report_context.rs
+13 -1
View File
@@ -4,7 +4,7 @@ use aether_ai_serving::{
build_ai_execution_report_context,
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
AiExecutionReportContextParts, AiRequestOrigin,
AiExecutionReportContextParts, AiRequestOrigin, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
};
use aether_routing_core::ResolvedRoutingPolicy;
use aether_runtime_state::RuntimeLockLease;
@@ -59,6 +59,9 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
pub(crate) routing_policy: Option<&'a ResolvedRoutingPolicy>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
/// Routing policy sticky-key attempt budget; read back by the attempt
/// loop to derive same-key retries lazily.
pub(crate) sticky_key_attempts: Option<u32>,
pub(crate) client_requested_stream: bool,
pub(crate) upstream_is_stream: bool,
pub(crate) has_envelope: bool,
@@ -124,6 +127,12 @@ pub(crate) fn build_local_execution_report_context(
Value::Number(epoch.into()),
);
}
if let Some(sticky_key_attempts) = parts.sticky_key_attempts {
extra_fields.insert(
STICKY_KEY_ATTEMPTS_REPORT_FIELD.to_string(),
Value::Number(sticky_key_attempts.into()),
);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
@@ -330,6 +339,7 @@ mod tests {
client_session_affinity: Some(&client_session_affinity),
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
@@ -413,6 +423,7 @@ mod tests {
client_session_affinity: None,
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
@@ -480,6 +491,7 @@ mod tests {
client_session_affinity: None,
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs
+1
View File
@@ -109,6 +109,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream: spec_metadata.require_streaming,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs
+9
View File
@@ -26,6 +26,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
@@ -108,6 +109,10 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -181,6 +186,10 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs
+1
View File
@@ -122,6 +122,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs
+9
View File
@@ -27,6 +27,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
@@ -127,6 +128,10 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
@@ -201,6 +206,10 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs
+1
View File
@@ -90,6 +90,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs
+9
View File
@@ -29,6 +29,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalVideoCreateCandidateAttempt;
@@ -133,6 +134,10 @@ pub(super) async fn list_local_video_create_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
@@ -190,6 +195,10 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs
+251 -21
View File
@@ -9,10 +9,37 @@ pub(crate) fn is_deepseek_provider(provider_type: &str, base_url: &str) -> bool
return true;
}
let host = base_url_host(base_url);
let Some(host) = base_url_host(base_url) else {
return false;
};
host == "deepseek.com" || host.ends_with(".deepseek.com")
}
fn is_deepseek_model(provider_model: &str) -> bool {
let provider_model = provider_model.trim().to_ascii_lowercase();
let leaf = provider_model
.rsplit(['/', ':'])
.next()
.unwrap_or(provider_model.as_str());
leaf == "deepseek" || leaf.starts_with("deepseek-") || leaf.starts_with("deepseek_")
}
fn is_deepseek_upstream(provider_type: &str, base_url: &str, provider_model: &str) -> bool {
is_deepseek_provider(provider_type, base_url) || is_deepseek_model(provider_model)
}
pub(crate) fn openai_responses_reasoning_replay_policy(
provider_type: &str,
base_url: &str,
provider_model: &str,
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
if is_deepseek_upstream(provider_type, base_url, provider_model) {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
} else {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
}
}
pub(crate) fn apply_deepseek_tool_call_thinking_compat(
provider_request_body: &mut Value,
provider_type: &str,
@@ -20,7 +47,11 @@ pub(crate) fn apply_deepseek_tool_call_thinking_compat(
provider_api_format: &str,
original_request_body: Option<&Value>,
) {
if !is_deepseek_provider(provider_type, base_url) {
let provider_model = provider_request_body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default();
if !is_deepseek_upstream(provider_type, base_url, provider_model) {
return;
}
@@ -36,24 +67,33 @@ pub(crate) fn apply_deepseek_tool_call_thinking_compat(
}
}
fn base_url_host(base_url: &str) -> String {
let lower = base_url.trim().to_ascii_lowercase();
let without_scheme = lower
.split_once("://")
.map(|(_, rest)| rest)
.unwrap_or(lower.as_str());
let without_userinfo = without_scheme
.rsplit_once('@')
.map(|(_, host)| host)
.unwrap_or(without_scheme);
without_userinfo
.split(['/', '?', '#'])
.next()
.unwrap_or_default()
.split(':')
.next()
.unwrap_or_default()
.to_string()
fn base_url_host(base_url: &str) -> Option<String> {
let base_url = base_url.trim();
if base_url.is_empty() {
return None;
}
// Provider configuration historically accepted both absolute URLs and a
// bare authority/path. Use a real URL parser for both forms: hand-parsing
// userinfo with `rsplit_once('@')` can mistake an `@` in the path or query
// for the authority delimiter and classify an attacker-controlled host as
// `api.deepseek.com`.
if let Ok(parsed) = url::Url::parse(base_url) {
if let Some(host) = parsed
.host_str()
.filter(|_| matches!(parsed.scheme(), "http" | "https" | "ws" | "wss"))
{
return Some(host.to_ascii_lowercase());
}
if base_url.contains("://") {
return None;
}
}
url::Url::parse(&format!("https://{base_url}"))
.ok()?
.host_str()
.map(str::to_ascii_lowercase)
}
fn source_disables_thinking(
@@ -241,7 +281,10 @@ fn is_claude_thinking_block(block: &Value) -> bool {
mod tests {
use serde_json::json;
use super::{apply_deepseek_tool_call_thinking_compat, is_deepseek_provider};
use super::{
apply_deepseek_tool_call_thinking_compat, is_deepseek_provider,
openai_responses_reasoning_replay_policy,
};
#[test]
fn detects_deepseek_provider_by_type_or_host() {
@@ -253,10 +296,151 @@ mod tests {
"custom",
"https://api.deepseek.com/v1"
));
assert!(is_deepseek_provider("custom", "api.deepseek.com/v1"));
assert!(is_deepseek_provider("custom", "api.deepseek.com:443/v1"));
assert!(!is_deepseek_provider(
"custom",
"https://example.com/deepseek"
));
assert!(!is_deepseek_provider(
"custom",
"https://api.deepseek.com.evil.example/v1"
));
assert!(!is_deepseek_provider(
"custom",
"https://api.deepseek.com@evil.example/v1"
));
assert!(!is_deepseek_provider(
"custom",
"https://evil.example/path@api.deepseek.com/v1"
));
assert!(!is_deepseek_provider(
"custom",
"https://evil.example/?relay=@api.deepseek.com"
));
assert!(!is_deepseek_provider("custom", "ftp://api.deepseek.com/v1"));
assert_eq!(
openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"openai",
"https://api.openai.com/v1",
"gpt-5.6-sol",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"custom",
"https://api.b.ai/v1",
"deepseek-v4-flash",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"custom",
"https://api.b.ai/v1",
"not-deepseek-compatible",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
}
#[test]
fn custom_deepseek_host_preserves_production_shaped_opaque_reasoning_replay() {
let reasoning_items = (0..66)
.map(|index| {
json!({
"type": "reasoning",
"encrypted_content": format!("550e8400-e29b-41d4-a716-{index:012}"),
"content": [{
"type": "reasoning_text",
"text": format!("opaque DeepSeek reasoning {index}")
}]
})
})
.collect::<Vec<_>>();
let request = json!({
"model": "deepseek-v4-flash",
"input": reasoning_items.clone(),
"future_request_field": {"preserve": true}
});
let replay_policy = openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
);
let mut provider_body = crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy(
&request,
"openai:responses",
"deepseek-v4-flash",
"custom",
"openai:responses",
"/v1/responses",
false,
None,
None,
None,
false,
replay_policy,
)
.expect("custom DeepSeek Responses body should build");
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut provider_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: "openai:responses",
provider_api_format: "openai:responses",
provider_type: "custom",
provider_model: "deepseek-v4-flash",
source_model: "deepseek-v4-flash",
body_rules: None,
upstream_is_stream: false,
require_body_stream_field: false,
},
None,
replay_policy,
)
.expect("custom DeepSeek finalization should accept opaque reasoning replay");
assert_eq!(provider_body["input"].as_array().map(Vec::len), Some(66));
assert_eq!(provider_body["future_request_field"]["preserve"], true);
let mut deepseek = json!({"input": reasoning_items.clone()});
let mut openai = json!({"input": reasoning_items});
assert_eq!(
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy(
&mut deepseek,
"openai:responses",
openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
),
),
0
);
assert_eq!(deepseek["input"].as_array().map(Vec::len), Some(66));
assert_eq!(
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy(
&mut openai,
"openai:responses",
openai_responses_reasoning_replay_policy(
"openai",
"https://api.openai.com/v1",
"gpt-5.6-sol",
),
),
66
);
assert_eq!(openai["input"].as_array().map(Vec::len), Some(0));
}
#[test]
@@ -286,6 +470,52 @@ mod tests {
assert_eq!(body["messages"][1]["reasoning_content"], "");
}
#[test]
fn custom_relay_deepseek_model_adds_chat_thinking_compat() {
let mut body = json!({
"model": "deepseek-v4-flash",
"messages": [
{"role": "user", "content": "inspect the repository"},
{"role": "assistant", "content": null, "tool_calls": [{
"id": "call_1",
"type": "function",
"function": {"name": "inspect", "arguments": "{}"}
}]},
{"role": "tool", "tool_call_id": "call_1", "content": "done"}
]
});
apply_deepseek_tool_call_thinking_compat(
&mut body,
"custom",
"https://api.b.ai/v1",
"openai:chat",
None,
);
assert_eq!(body["thinking"]["type"], "enabled");
assert_eq!(body["messages"][1]["reasoning_content"], "");
}
#[test]
fn custom_relay_non_deepseek_model_is_not_rewritten() {
let original = json!({
"model": "not-deepseek-compatible",
"messages": [{"role": "assistant", "content": "done"}]
});
let mut body = original.clone();
apply_deepseek_tool_call_thinking_compat(
&mut body,
"custom",
"https://api.b.ai/v1",
"openai:chat",
None,
);
assert_eq!(body, original);
}
#[test]
fn openai_chat_deepseek_honors_disabled_thinking() {
let original = json!({"reasoning_effort": "none"});
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
+2
View File
@@ -142,6 +142,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
@@ -377,6 +378,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
+24 -17
View File
@@ -23,8 +23,8 @@ use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat,
codex_model_capabilities_for_transport, is_deepseek_provider,
openai_provider_request_contract_failure_extra_data, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy,
request_body_build_failure_extra_data, request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body,
@@ -364,6 +364,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
body_json,
&input.auth_context,
spec_metadata.api_format,
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds,
&attempt.candidate_id,
)
.await?;
@@ -593,18 +594,24 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
input.auth_context.api_key_id.as_str(),
)
.await?;
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
prepared_candidate.mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec_metadata.api_format,
reasoning_replay_policy,
&attempt.candidate_id,
)
.await?;
let body_json = redaction.body_json.as_ref();
let mut provider_request_body =
match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers(
match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy(
body_json,
spec_metadata.api_format,
&prepared_candidate.mapped_model,
@@ -620,6 +627,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
Some(input.auth_context.api_key_id.as_str()),
Some(effective_headers),
false,
reasoning_replay_policy,
) {
Some(body) => body,
None => {
@@ -750,7 +758,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
source_model,
);
if let Err(violation) =
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: spec_metadata.api_format,
@@ -766,6 +774,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
),
},
codex_model_capabilities.as_ref(),
reasoning_replay_policy,
)
{
mark_skipped_local_standard_candidate_with_extra_data(
@@ -1314,21 +1323,19 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
.as_object_mut()?
.insert("stream".to_string(), Value::Bool(true));
}
provider_request_body = project_openai_image_api_request_body(
&provider_request_body,
&prepared_candidate.mapped_model,
converted.operation,
crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
transport.provider.provider_type.as_str(),
Some(prepared_candidate.mapped_model.as_str()),
),
)?;
if is_codex {
provider_request_body = project_codex_openai_image_api_request_body(
provider_request_body = if is_codex {
project_codex_openai_image_api_request_body(&provider_request_body, converted.operation)?
} else {
project_openai_image_api_request_body(
&provider_request_body,
&prepared_candidate.mapped_model,
converted.operation,
)?;
}
crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
transport.provider.provider_type.as_str(),
Some(prepared_candidate.mapped_model.as_str()),
),
)?
};
let request_path = match converted.operation {
OpenAiImageOperation::Generate => "/v1/images/generations",
OpenAiImageOperation::Edit => "/v1/images/edits",
apps/aether-gateway/src/ai_serving/planner/standard/mod.rs
+10 -4
View File
@@ -18,7 +18,10 @@ pub(crate) use self::codex::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers,
codex_model_capabilities_for_transport,
};
pub(crate) use self::deepseek::{apply_deepseek_tool_call_thinking_compat, is_deepseek_provider};
pub(crate) use self::deepseek::{
apply_deepseek_tool_call_thinking_compat, is_deepseek_provider,
openai_responses_reasoning_replay_policy,
};
pub(crate) use self::family::{
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
build_local_sync_attempt_source, build_local_sync_plan_and_reports,
@@ -30,6 +33,7 @@ pub(crate) use self::normalize::{
build_cross_format_openai_responses_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_codex_model_capabilities,
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation,
build_local_openai_responses_upstream_url, validate_final_openai_provider_request,
};
pub(crate) use self::openai::{
@@ -42,13 +46,15 @@ pub(crate) use self::openai::{
build_local_openai_responses_sync_attempt_source_for_kind,
build_local_openai_responses_sync_plan_and_reports_for_kind, copy_request_number_field,
copy_request_number_field_as, map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_gemini_budget, maybe_build_stream_local_decision_payload,
map_openai_reasoning_effort_to_gemini_budget, maybe_build_responses_websocket_decision,
maybe_build_stream_local_decision_payload,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_decision_payload,
maybe_build_sync_local_openai_embedding_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, parse_openai_stop_sequences,
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
value_as_u64,
value_as_u64, ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
ResponsesWebSocketPinnedCandidate,
};
pub(crate) use crate::ai_serving::normalize_standard_request_to_openai_chat_request;
pub(crate) use crate::ai_serving::{
@@ -59,7 +65,7 @@ pub(crate) use crate::ai_serving::{
};
pub(crate) use crate::ai_serving::{
build_standard_request_body, build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy,
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content,
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,
apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs
+1
View File
@@ -15,6 +15,7 @@ pub(crate) use self::responses::{
build_cross_format_openai_responses_request_body_with_codex_model_capabilities,
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_codex_model_capabilities,
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation,
build_local_openai_responses_upstream_url,
};
pub(super) use crate::ai_serving::planner::common::{
apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs
+84 -17
View File
@@ -50,6 +50,69 @@ pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabil
request_headers: &http::HeaderMap,
model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>,
enable_model_directives: bool,
) -> Option<Value> {
build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode(
body_json,
mapped_model,
require_streaming,
force_body_stream_field,
provider_type,
provider_api_format,
body_rules,
request_headers,
model_capabilities,
enable_model_directives,
false,
)
}
/// Builds a Responses body for a pinned WebSocket continuation.
///
/// This is intentionally an additive variant of the ordinary HTTP builder.
/// The WebSocket framing layer, rather than a JSON-body heuristic, tells the
/// Codex compatibility pass that `previous_response_id` is transport state and
/// that Responses Lite `tools`/`instructions` must not be materialized into a
/// second historical input prefix.
pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
force_body_stream_field: bool,
provider_type: &str,
provider_api_format: &str,
body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>,
enable_model_directives: bool,
) -> Option<Value> {
build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode(
body_json,
mapped_model,
require_streaming,
force_body_stream_field,
provider_type,
provider_api_format,
body_rules,
request_headers,
model_capabilities,
enable_model_directives,
true,
)
}
#[allow(clippy::too_many_arguments)]
fn build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
force_body_stream_field: bool,
provider_type: &str,
provider_api_format: &str,
body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>,
enable_model_directives: bool,
websocket_continuation: bool,
) -> Option<Value> {
let provider_request_body = surface_build_local_openai_responses_request_body(
body_json,
@@ -68,23 +131,31 @@ pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabil
.get("model")
.and_then(Value::as_str)
.unwrap_or(mapped_model);
crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities(
&mut provider_request_body,
provider_type,
provider_api_format,
mapped_model,
source_model,
model_capabilities,
body_rules,
);
if websocket_continuation {
crate::ai_serving::apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities(
&mut provider_request_body,
provider_type,
provider_api_format,
mapped_model,
source_model,
model_capabilities,
body_rules,
);
} else {
crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities(
&mut provider_request_body,
provider_type,
provider_api_format,
mapped_model,
source_model,
model_capabilities,
body_rules,
);
}
apply_openai_responses_compact_special_body_edits(
&mut provider_request_body,
provider_api_format,
);
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items(
&mut provider_request_body,
provider_api_format,
);
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
@@ -176,10 +247,6 @@ pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_
&mut provider_request_body,
provider_api_format,
);
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items(
&mut provider_request_body,
provider_api_format,
);
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs
+61 -3
View File
@@ -102,6 +102,59 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
assert_eq!(provider_request_body["messages"][0]["content"], "hello");
}
#[test]
fn maps_openai_responses_additional_tools_without_message_name() {
let body_json = json!({
"model": "gpt-5",
"input": [
{
"type": "additional_tools",
"role": "developer",
"tools": [{
"type": "function",
"name": "get_weather",
"description": "Get the weather",
"parameters": {
"type": "object",
"properties": {}
}
}]
},
{
"role": "user",
"content": "What is the weather?"
}
]
});
let provider_request_body = build_cross_format_openai_responses_request_body(
&body_json,
"gpt-5-upstream",
"openai:responses",
"openai:chat",
false,
false,
"openai",
None,
None,
&http::HeaderMap::new(),
false,
)
.expect("Responses additional tools should map to a Chat request body");
assert_eq!(
provider_request_body["messages"].as_array().map(Vec::len),
Some(1)
);
assert_eq!(provider_request_body["messages"][0]["role"], "user");
assert!(provider_request_body["messages"][0].get("name").is_none());
assert_eq!(provider_request_body["tools"][0]["type"], "function");
assert_eq!(
provider_request_body["tools"][0]["function"]["name"],
"get_weather"
);
}
#[test]
fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
let body_json: Value = serde_json::from_str(
@@ -153,7 +206,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
}
#[test]
fn local_openai_responses_wrapper_strips_foreign_reasoning_item_ids() {
fn local_openai_responses_wrapper_defers_reasoning_replay_filtering() {
let body_json = json!({
"model": "gpt-5.4",
"input": [
@@ -184,9 +237,14 @@ fn local_openai_responses_wrapper_strips_foreign_reasoning_item_ids() {
let input = provider_request_body["input"]
.as_array()
.expect("input array");
assert_eq!(input.len(), 2);
// This provider-agnostic normalization layer cannot decide whether an
// id-less/foreign reasoning item is opaque state required by DeepSeek.
// The provider-aware finalization pass applies the strict or DeepSeek
// replay policy once the selected upstream base URL is known.
assert_eq!(input.len(), 3);
assert_eq!(input[0]["id"], "rs_provider_123");
assert_eq!(input[1]["type"], "message");
assert_eq!(input[1]["id"], "item_72d3bd8d367d01977ace23f1");
assert_eq!(input[2]["type"], "message");
}
#[test]
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
+1
View File
@@ -195,6 +195,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+90 -32
View File
@@ -29,8 +29,8 @@ use crate::ai_serving::planner::standard::{
apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, codex_model_capabilities_for_transport,
openai_provider_request_contract_failure_extra_data, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy,
request_body_build_failure_extra_data, request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
@@ -140,7 +140,7 @@ fn finalize_openai_chat_provider_request_body(
mapped_model,
source_model,
);
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: "openai:chat",
@@ -156,6 +156,11 @@ fn finalize_openai_chat_provider_request_body(
),
},
codex_model_capabilities.as_ref(),
openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model,
),
)
.err()
.map(|violation| {
@@ -206,6 +211,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
body_json,
&input.auth_context,
"openai:chat",
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds,
candidate_id,
)
.await?;
@@ -1417,38 +1423,20 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
return Ok(None);
};
if !is_chatgpt_web {
let Some(projected) = project_openai_image_api_request_body(
&provider_request_body,
&prepared_candidate.mapped_model,
operation,
crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
transport.provider.provider_type.as_str(),
Some(prepared_candidate.mapped_model.as_str()),
),
) else {
mark_skipped_local_openai_chat_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
"openai:chat",
provider_api_format,
let projected = if is_codex {
project_codex_openai_image_api_request_body(&provider_request_body, operation)
} else {
project_openai_image_api_request_body(
&provider_request_body,
&prepared_candidate.mapped_model,
operation,
crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
transport.provider.provider_type.as_str(),
Some(prepared_candidate.mapped_model.as_str()),
),
)
.await;
return Ok(None);
};
provider_request_body = projected;
}
if is_codex {
let Some(projected) =
project_codex_openai_image_api_request_body(&provider_request_body, operation)
else {
let Some(projected) = projected else {
mark_skipped_local_openai_chat_candidate_with_extra_data(
state,
input,
@@ -2195,6 +2183,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
@@ -2364,6 +2353,23 @@ mod tests {
eligible
}
fn sample_custom_deepseek_responses_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_gemini_cli_transport();
transport.provider.name = "deepseek".to_string();
transport.provider.provider_type = "custom".to_string();
transport.endpoint.api_format = "openai:responses".to_string();
transport.endpoint.api_family = Some("openai".to_string());
transport.endpoint.endpoint_kind = Some("responses".to_string());
transport.endpoint.base_url = "https://api.deepseek.com/v1".to_string();
transport.endpoint.custom_path = None;
transport.key.api_formats = Some(vec!["openai:responses".to_string()]);
transport.key.auth_type = "bearer".to_string();
transport.key.decrypted_api_key = "test-api-key".to_string();
transport.key.decrypted_auth_config = None;
transport.key.upstream_metadata = None;
transport
}
fn sample_custom_directive_input() -> LocalOpenAiChatDecisionInput {
let mut input = sample_input();
input.requested_model = "gpt-5.6-sol-high".to_string();
@@ -2389,6 +2395,58 @@ mod tests {
input
}
#[test]
fn responses_shaped_chat_request_preserves_deepseek_opaque_reasoning_at_finalization() {
let reasoning_items = (0..66)
.map(|index| {
json!({
"type": "reasoning",
"encrypted_content": format!("opaque-deepseek-state-{index}"),
"content": [{
"type": "reasoning_text",
"text": format!("provider thinking state {index}")
}],
"future_capability": {"preserve": true}
})
})
.collect::<Vec<_>>();
let original_body = json!({
"model": "deepseek-v4-flash",
"input": reasoning_items
});
let transport = sample_custom_deepseek_responses_transport();
let mut provider_body = build_cross_format_openai_chat_request_body(
&original_body,
"deepseek-v4-flash",
"custom",
"openai:responses",
false,
false,
None,
None,
&http::HeaderMap::new(),
false,
)
.expect("Responses-shaped chat request should build");
assert!(finalize_openai_chat_provider_request_body(
&mut provider_body,
None,
"openai:responses",
false,
false,
&original_body,
&transport,
"deepseek-v4-flash",
)
.is_none());
let input = provider_body["input"].as_array().expect("provider input");
assert_eq!(input.len(), 66);
assert_eq!(input[0]["type"], "reasoning");
assert_eq!(input[0]["content"][0]["type"], "reasoning_text");
assert_eq!(input[0]["future_capability"]["preserve"], true);
}
fn sample_alias_max_directive_input() -> LocalOpenAiChatDecisionInput {
let mut input = sample_input();
input.requested_model = "deployment-alias-max".to_string();
apps/aether-gateway/src/ai_serving/planner/standard/openai/mod.rs
+3 -1
View File
@@ -23,6 +23,8 @@ pub(crate) use responses::{
build_local_openai_responses_stream_plan_and_reports_for_kind,
build_local_openai_responses_sync_attempt_source_for_kind,
build_local_openai_responses_sync_plan_and_reports_for_kind,
maybe_build_responses_websocket_decision,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, ResponsesWebSocketBodyNormalization,
ResponsesWebSocketDecision, ResponsesWebSocketPinnedCandidate,
};
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs
+18 -4
View File
@@ -10,6 +10,7 @@ use super::super::{
AiStreamAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
@@ -233,6 +234,19 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
},
);
let log_downstream_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
event_name = "local_openai_responses_stream_plan_built",
log_type = "debug",
@@ -242,11 +256,11 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path = %parts.uri.path(),
downstream_query = ?parts.uri.query(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?payload.upstream_base_url,
decision_upstream_url = ?payload.upstream_url,
plan_url = %plan.url,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = effective_upstream_is_stream,
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs
+18 -4
View File
@@ -10,6 +10,7 @@ use super::super::{
AiSyncAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -200,6 +201,19 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
},
);
let log_downstream_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
event_name = "local_openai_responses_sync_plan_built",
log_type = "debug",
@@ -209,11 +223,11 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path = %parts.uri.path(),
downstream_query = ?parts.uri.query(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?payload.upstream_base_url,
decision_upstream_url = ?payload.upstream_url,
plan_url = %plan.url,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = payload.upstream_is_stream,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs
+8 -3
View File
@@ -5,11 +5,16 @@ mod request;
#[path = "decision/support.rs"]
mod support;
pub(super) use self::payload::maybe_build_local_openai_responses_decision_payload_for_candidate;
pub(super) use self::payload::{
maybe_build_local_openai_responses_decision_payload_for_candidate,
maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode,
};
pub(super) use self::support::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
resolve_local_openai_responses_decision_input,
resolve_local_openai_responses_decision_input_with_snapshot,
LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesCandidateAttemptSource,
LocalOpenAiResponsesDecisionInput,
};
pub(super) use crate::ai_serving::LocalOpenAiResponsesSpec;
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
+67 -19
View File
@@ -2,7 +2,8 @@ use serde_json::json;
use tracing::debug;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision_with_websocket_mode;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_native_client_envelope_name,
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
@@ -20,7 +21,10 @@ use crate::{
AiExecutionDecision, AppState, GatewayError,
};
use super::request::resolve_local_openai_responses_candidate_payload_parts;
use super::request::{
resolve_local_openai_responses_candidate_payload_parts,
resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode,
};
use super::support::{LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesDecisionInput};
use super::LocalOpenAiResponsesSpec;
@@ -32,6 +36,26 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
input: &LocalOpenAiResponsesDecisionInput,
attempt: LocalOpenAiResponsesCandidateAttempt,
spec: LocalOpenAiResponsesSpec,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode(
state, parts, trace_id, body_json, input, attempt, spec, false,
)
.await
}
/// Builds a candidate payload for a pinned WebSocket turn without changing
/// the ordinary HTTP/plan-builder path. The explicit mode is carried all the
/// way to body normalization because a JSON `type` field is not a reliable
/// transport discriminator once body rules and conversions have run.
pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
attempt: LocalOpenAiResponsesCandidateAttempt,
spec: LocalOpenAiResponsesSpec,
websocket_continuation: bool,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let attempt_identity = attempt.attempt_identity();
@@ -41,19 +65,35 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
candidate_id,
..
} = attempt;
let Some(resolved) = resolve_local_openai_responses_candidate_payload_parts(
state,
parts,
trace_id,
body_json,
input,
&eligible,
candidate_index,
&candidate_id,
spec,
)
.await?
else {
let resolved = if websocket_continuation {
resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode(
state,
parts,
trace_id,
body_json,
input,
&eligible,
candidate_index,
&candidate_id,
spec,
true,
)
.await?
} else {
resolve_local_openai_responses_candidate_payload_parts(
state,
parts,
trace_id,
body_json,
input,
&eligible,
candidate_index,
&candidate_id,
spec,
)
.await?
};
let Some(resolved) = resolved else {
return Ok(None);
};
let candidate = &eligible.candidate;
@@ -144,6 +184,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
@@ -164,6 +205,12 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
&resolved.transport,
);
let log_base_url = sanitize_upstream_url_for_log(resolved.transport.endpoint.base_url.as_str());
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(resolved.upstream_url.as_str());
debug!(
event_name = "local_openai_responses_decision_payload_built",
log_type = "debug",
@@ -180,9 +227,9 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_api_format = spec_metadata.api_format,
provider_api_format = %resolved.provider_api_format,
request_path = %parts.uri.path(),
request_query = ?parts.uri.query(),
upstream_base_url = %resolved.transport.endpoint.base_url,
upstream_url = %resolved.upstream_url,
request_query = ?log_request_query,
upstream_base_url = %log_base_url,
upstream_url = %log_upstream_url,
upstream_is_stream = resolved.upstream_is_stream,
has_envelope = resolved.envelope_name.is_some(),
"gateway built local openai responses decision payload"
@@ -243,10 +290,11 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
});
apply_provider_request_routing_policy_to_decision(
apply_provider_request_routing_policy_to_decision_with_websocket_mode(
input,
&mut decision,
Some(transport.as_ref()),
websocket_continuation,
)?;
Ok(Some(decision))
}
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
+139 -65
View File
@@ -24,6 +24,7 @@ use crate::ai_serving::planner::gemini_cli::{
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
sanitize_upstream_url_for_log,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
@@ -31,9 +32,10 @@ use crate::ai_serving::planner::standard::{
build_cross_format_openai_responses_request_body_with_codex_model_capabilities,
build_cross_format_openai_responses_upstream_url,
build_local_openai_responses_request_body_with_codex_model_capabilities,
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation,
build_local_openai_responses_upstream_url, codex_model_capabilities_for_transport,
openai_provider_request_contract_failure_extra_data, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy,
request_body_build_failure_extra_data, request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::{
@@ -203,6 +205,34 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiResponsesSpec,
) -> Result<Option<LocalOpenAiResponsesCandidatePayloadParts>, GatewayError> {
resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
spec,
false,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiResponsesSpec,
websocket_continuation: bool,
) -> Result<Option<LocalOpenAiResponsesCandidatePayloadParts>, GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
@@ -405,12 +435,18 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
input.auth_context.api_key_id.as_str(),
)
.await?;
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
body_json,
&input.auth_context,
spec_metadata.api_format,
reasoning_replay_policy,
candidate_id,
)
.await?;
@@ -437,41 +473,42 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
mapped_model.as_str(),
source_model,
);
let Some(mut base_provider_request_body) =
(if is_grok && is_grok_text_provider_api_format(provider_api_format) {
build_local_openai_responses_request_body_with_codex_model_capabilities(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
spec_metadata.api_format,
transport.endpoint.body_rules.as_ref(),
effective_headers,
codex_model_capabilities.as_ref(),
false,
)
} else if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body_with_codex_model_capabilities(
body_json,
&mapped_model,
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
},
effective_headers,
Some(input.auth_context.api_key_id.as_str()),
codex_model_capabilities.as_ref(),
false,
)
} else {
build_local_openai_responses_request_body_with_codex_model_capabilities(
let Some(mut base_provider_request_body) = (if is_grok
&& is_grok_text_provider_api_format(provider_api_format)
{
build_local_openai_responses_request_body_with_codex_model_capabilities(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
spec_metadata.api_format,
transport.endpoint.body_rules.as_ref(),
effective_headers,
codex_model_capabilities.as_ref(),
false,
)
} else if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body_with_codex_model_capabilities(
body_json,
&mapped_model,
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
},
effective_headers,
Some(input.auth_context.api_key_id.as_str()),
codex_model_capabilities.as_ref(),
false,
)
} else if websocket_continuation {
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation(
body_json,
&mapped_model,
upstream_is_stream,
@@ -487,8 +524,24 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
codex_model_capabilities.as_ref(),
false,
)
})
else {
} else {
build_local_openai_responses_request_body_with_codex_model_capabilities(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
},
effective_headers,
codex_model_capabilities.as_ref(),
false,
)
}) else {
mark_skipped_local_openai_responses_candidate_with_extra_data(
state,
input,
@@ -531,25 +584,35 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
provider_api_format,
Some(body_json),
);
if let Err(violation) =
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
let finalization = crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: spec_metadata.api_format,
provider_api_format,
provider_type: transport.provider.provider_type.as_str(),
provider_model: mapped_model.as_str(),
source_model,
body_rules: transport.endpoint.body_rules.as_ref(),
upstream_is_stream,
require_body_stream_field: request_requires_body_stream_field(
body_json,
force_body_stream_field,
),
};
let finalization_result = if websocket_continuation {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation(
&mut base_provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: spec_metadata.api_format,
provider_api_format,
provider_type: transport.provider.provider_type.as_str(),
provider_model: mapped_model.as_str(),
source_model,
body_rules: transport.endpoint.body_rules.as_ref(),
upstream_is_stream,
require_body_stream_field: request_requires_body_stream_field(
body_json,
force_body_stream_field,
),
},
finalization,
codex_model_capabilities.as_ref(),
reasoning_replay_policy,
)
{
} else {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut base_provider_request_body,
finalization,
codex_model_capabilities.as_ref(),
reasoning_replay_policy,
)
};
if let Err(violation) = finalization_result {
mark_skipped_local_openai_responses_candidate_with_extra_data(
state,
input,
@@ -804,6 +867,17 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
let log_base_url = sanitize_upstream_url_for_log(transport.endpoint.base_url.as_str());
let log_custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_upstream_url_resolved",
@@ -819,12 +893,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
base_url = %transport.endpoint.base_url,
custom_path = ?transport.endpoint.custom_path,
base_url = %log_base_url,
custom_path = ?log_custom_path,
request_path = %parts.uri.path(),
request_query = ?parts.uri.query(),
request_query = ?log_request_query,
mapped_model = %mapped_model,
upstream_url = %upstream_url,
upstream_url = %log_upstream_url,
upstream_is_stream,
"gateway resolved local openai responses upstream url"
);
@@ -1395,7 +1469,10 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
return None;
};
let operation = openai_image_operation_from_summary(&image_request_summary)?;
if !is_chatgpt_web {
if is_codex {
provider_request_body =
project_codex_openai_image_api_request_body(&provider_request_body, operation)?;
} else if !is_chatgpt_web {
provider_request_body = project_openai_image_api_request_body(
&provider_request_body,
&prepared_candidate.mapped_model,
@@ -1406,10 +1483,6 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
),
)?;
}
if is_codex {
provider_request_body =
project_codex_openai_image_api_request_body(&provider_request_body, operation)?;
}
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
@@ -1938,6 +2011,7 @@ async fn build_kiro_openai_responses_payload_parts(
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_kiro_upstream_url_resolved",
@@ -1953,7 +2027,7 @@ async fn build_kiro_openai_responses_payload_parts(
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
upstream_url = %upstream_url,
upstream_url = %log_upstream_url,
upstream_is_stream,
"gateway resolved local openai responses kiro upstream url"
);
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs
+40 -11
View File
@@ -22,6 +22,7 @@ use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::planner::decision_input::{
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
resolve_local_authenticated_decision_input_with_snapshot,
};
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
@@ -32,7 +33,8 @@ use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
openai_responses_request_operation, resolve_local_decision_execution_runtime_auth_context,
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayControlDecision,
PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
@@ -51,6 +53,21 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
resolve_local_openai_responses_decision_input_with_snapshot(
state, parts, trace_id, decision, body_json, plan_kind, None,
)
.await
}
pub(crate) async fn resolve_local_openai_responses_decision_input_with_snapshot(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
auth_snapshot_override: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
warn!(
@@ -87,16 +104,28 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
return Ok(None);
};
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
{
let resolved_input = match if let Some(auth_snapshot) = auth_snapshot_override {
resolve_local_authenticated_decision_input_with_snapshot(
state,
auth_context.clone(),
Some(auth_snapshot.clone()),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
} else {
resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
} {
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => {
warn!(
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
+893
View File
@@ -1,6 +1,63 @@
use crate::ai_serving::planner::common::endpoint_config_forces_body_stream_field;
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::codex::codex_model_capabilities_for_transport;
use crate::ai_serving::planner::standard::normalize::{
build_local_openai_responses_request_body_with_codex_model_capabilities,
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation,
};
use crate::ai_serving::planner::standard::openai_responses_reasoning_replay_policy;
use crate::ai_serving::GatewayControlDecision;
use crate::orchestration::{
codex_quota_breaker_blocks_candidate, log_codex_quota_breaker_check_failure,
responses_websocket_adapter, ResponsesWebSocketAdapter,
};
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_runtime_state::RuntimeLockLease;
use std::collections::BTreeSet;
/// Releases a scheduler pool-key lease if WebSocket planning is cancelled
/// after candidate selection but before ownership reaches the turn lifecycle.
struct ResponsesWebSocketPlanningLeaseGuard {
state: AppState,
lease: Option<RuntimeLockLease>,
}
impl ResponsesWebSocketPlanningLeaseGuard {
fn new(state: &AppState, lease: Option<&RuntimeLockLease>) -> Self {
Self {
state: state.clone(),
lease: lease.cloned(),
}
}
async fn release(mut self) {
// Keep the lease armed across the await. If the owner task is aborted
// or reaches its hard deadline while the runtime backend is stalled,
// Drop can still hand cleanup to a detached owner.
if release_responses_websocket_planning_lease(&self.state, self.lease.as_ref()).await {
self.lease = None;
}
}
fn disarm(&mut self) {
self.lease = None;
}
}
impl Drop for ResponsesWebSocketPlanningLeaseGuard {
fn drop(&mut self) {
let Some(lease) = self.lease.take() else {
return;
};
let state = self.state.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = release_responses_websocket_planning_lease(&state, Some(&lease)).await;
});
}
}
}
mod decision;
mod plans;
@@ -8,7 +65,9 @@ mod plans;
use self::decision::{
build_local_openai_responses_candidate_attempt_source,
maybe_build_local_openai_responses_decision_payload_for_candidate,
maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode,
resolve_local_openai_responses_decision_input,
resolve_local_openai_responses_decision_input_with_snapshot,
};
use self::plans::{
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
@@ -165,3 +224,837 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
Ok(None)
}
/// One eligible upstream plus the adapter that is allowed to speak to it.
///
/// The adapter is selected from the provider-scoped capability before the
/// decision leaves the planner. This prevents a public Responses socket from
/// choosing an arbitrary provider protocol after scheduling has completed.
pub(crate) struct ResponsesWebSocketDecision {
pub(crate) execution: AiExecutionDecision,
pub(crate) adapter: ResponsesWebSocketAdapter,
pub(crate) normalization: ResponsesWebSocketBodyNormalization,
/// Effective key auth after applying the endpoint API-format override.
/// Protocol companions such as Codex Live must not infer this from a URL
/// or from the presence of one particular generated header.
pub(crate) effective_auth_type: String,
}
/// The scheduler identity a continuation is allowed to reuse.
///
/// A `previous_response_id` chain cannot move to another provider connection,
/// but it still has to pass the current scheduler runtime checks on every
/// turn. The planner uses this identity as a filter rather than selecting an
/// arbitrary eligible replacement.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ResponsesWebSocketPinnedCandidate {
provider_id: String,
endpoint_id: String,
key_id: String,
}
impl ResponsesWebSocketPinnedCandidate {
pub(crate) fn new(provider_id: &str, endpoint_id: &str, key_id: &str) -> Option<Self> {
Some(Self {
provider_id: non_empty_decision_identity(Some(provider_id))?,
endpoint_id: non_empty_decision_identity(Some(endpoint_id))?,
key_id: non_empty_decision_identity(Some(key_id))?,
})
}
pub(crate) fn from_decision(decision: &AiExecutionDecision) -> Option<Self> {
Self::new(
decision.provider_id.as_deref()?,
decision.endpoint_id.as_deref()?,
decision.key_id.as_deref()?,
)
}
pub(crate) fn provider_id(&self) -> &str {
self.provider_id.as_str()
}
pub(crate) fn endpoint_id(&self) -> &str {
self.endpoint_id.as_str()
}
pub(crate) fn key_id(&self) -> &str {
self.key_id.as_str()
}
fn matches(
&self,
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
candidate.provider_id == self.provider_id
&& candidate.endpoint_id == self.endpoint_id
&& candidate.key_id == self.key_id
}
}
fn non_empty_decision_identity(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
}
/// Everything needed to re-run provider-body normalization for the candidate a
/// socket is already bound to.
///
/// A continuation turn (`previous_response_id` on the bound upstream) cannot
/// re-enter the planner, because planning selects a candidate and a different
/// key would break the response chain. Without this, such turns reached the
/// provider with only their `model` rewritten — skipping model directives,
/// endpoint body rules, and the Codex body contract that turn 1 received.
///
/// This value holds cloned scalars and JSON only: no candidate, no pool key
/// lease, no `AppState`. It cannot influence selection.
#[derive(Debug, Clone)]
pub(crate) struct ResponsesWebSocketBodyNormalization {
provider_type: String,
provider_api_format: String,
client_api_format: String,
mapped_model: String,
requested_model: String,
upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<serde_json::Value>,
request_headers: http::HeaderMap,
codex_model_capabilities: Option<crate::ai_serving::CodexResponsesModelCapabilities>,
reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
model_directive_patch: Option<serde_json::Value>,
}
impl ResponsesWebSocketBodyNormalization {
/// Builds a normalizer for a plain `openai:responses` upstream with no
/// endpoint body rules, directives or Codex capabilities, so relay tests can
/// construct a bound connection without standing up a provider snapshot.
#[cfg(test)]
pub(crate) fn for_tests(mapped_model: &str) -> Self {
Self {
provider_type: "openai".to_string(),
provider_api_format: "openai:responses".to_string(),
client_api_format: "openai:responses".to_string(),
mapped_model: mapped_model.to_string(),
requested_model: mapped_model.to_string(),
upstream_is_stream: true,
force_body_stream_field: false,
body_rules: None,
request_headers: http::HeaderMap::new(),
codex_model_capabilities: None,
reasoning_replay_policy:
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds,
model_directive_patch: None,
}
}
#[cfg(test)]
pub(crate) fn with_provider_type_for_tests(mut self, provider_type: &str) -> Self {
self.provider_type = provider_type.to_string();
self
}
#[cfg(test)]
pub(crate) fn with_body_rules_for_tests(mut self, body_rules: serde_json::Value) -> Self {
self.body_rules = Some(body_rules);
self
}
#[cfg(test)]
pub(crate) fn with_reasoning_replay_policy_for_tests(
mut self,
reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
) -> Self {
self.reasoning_replay_policy = reasoning_replay_policy;
self
}
#[cfg(test)]
pub(crate) fn with_model_directive_patch_for_tests(mut self, patch: serde_json::Value) -> Self {
self.model_directive_patch = Some(patch);
self
}
pub(crate) fn uses_codex_responses_lite(&self) -> bool {
if !self.provider_type.trim().eq_ignore_ascii_case("codex")
|| !crate::ai_serving::is_openai_responses_family_format(
self.provider_api_format.as_str(),
)
{
return false;
}
self.codex_model_capabilities
.clone()
.unwrap_or_else(|| {
crate::ai_serving::resolve_codex_responses_model_capabilities(
self.mapped_model.as_str(),
self.requested_model.as_str(),
None,
)
})
.use_responses_lite
}
pub(crate) fn reasoning_replay_policy(
&self,
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
self.reasoning_replay_policy
}
/// Returns whether an enabled endpoint body rule that applies to this
/// request owns the final value of a non-lineage WebSocket framing field.
///
/// Codex's HTTP-shaped normalization intentionally removes or rewrites a
/// few WebSocket-only fields. The framing layer may restore a value from
/// the raw client event only when an administrator rule did not handle
/// that path; otherwise the restore would silently undo the endpoint
/// policy after all request finalization had completed. Opaque lineage
/// (`previous_response_id`) is deliberately excluded by the framing layer:
/// its final value must remain the authenticated client value.
pub(crate) fn body_rules_handle_websocket_field(
&self,
client_event: &serde_json::Value,
field: &str,
) -> bool {
let Some(mut body_before_rules) =
crate::ai_serving::build_local_openai_responses_request_body_with_model_directives(
client_event,
self.mapped_model.as_str(),
self.upstream_is_stream,
false,
)
else {
// Normalization will reject the same malformed event. Keep the
// framing pass fail closed if this method is ever called alone.
return true;
};
crate::ai_serving::transport::rules::apply_local_body_rules_with_request_headers_and_track_path(
&mut body_before_rules,
self.body_rules.as_ref(),
Some(client_event),
Some(&self.request_headers),
field,
)
.unwrap_or(true)
}
pub(crate) fn has_same_responses_lite_static_contract(&self, other: &Self) -> bool {
self.provider_type
.trim()
.eq_ignore_ascii_case(other.provider_type.trim())
&& crate::ai_serving::api_format_alias_matches(
self.provider_api_format.as_str(),
other.provider_api_format.as_str(),
)
&& self.mapped_model == other.mapped_model
&& self.requested_model == other.requested_model
&& self.body_rules == other.body_rules
&& self.codex_model_capabilities == other.codex_model_capabilities
&& self.model_directive_patch == other.model_directive_patch
&& self.uses_codex_responses_lite() == other.uses_codex_responses_lite()
}
/// Produces a versioned digest of the complete body-normalization
/// contract. A continuation registry stores only this digest so a new
/// socket can fail closed when endpoint rules, model capabilities or
/// header-dependent normalization has changed, without persisting request
/// headers or other sensitive configuration.
pub(crate) fn continuation_fingerprint(&self) -> [u8; 32] {
use sha2::Digest as _;
let mut digest = sha2::Sha256::new();
digest.update(b"aether-responses-websocket-normalization-v1");
update_normalization_string_digest(&mut digest, self.provider_type.as_str());
update_normalization_string_digest(&mut digest, self.provider_api_format.as_str());
update_normalization_string_digest(&mut digest, self.client_api_format.as_str());
update_normalization_string_digest(&mut digest, self.mapped_model.as_str());
update_normalization_string_digest(&mut digest, self.requested_model.as_str());
digest.update([
u8::from(self.upstream_is_stream),
u8::from(self.force_body_stream_field),
]);
update_normalization_optional_json_digest(&mut digest, self.body_rules.as_ref());
update_normalization_body_rule_headers_digest(
&mut digest,
&self.request_headers,
self.body_rules.as_ref(),
);
update_normalization_codex_capabilities_digest(
&mut digest,
self.codex_model_capabilities.as_ref(),
);
digest.update([match self.reasoning_replay_policy {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds => 0,
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque => 1,
}]);
update_normalization_optional_json_digest(&mut digest, self.model_directive_patch.as_ref());
digest.finalize().into()
}
/// Applies the same body transformations the planner applied on the turn
/// that bound this upstream.
///
/// Mirrors the same-format branch of
/// `resolve_local_openai_responses_candidate_payload_parts`. The
/// cross-format, Kiro, Windsurf and Antigravity branches are unreachable
/// here: the WebSocket planner only returns candidates whose provider API
/// format is `openai:responses`.
///
/// Returns `None` when normalization fails. The WebSocket caller rejects
/// that turn rather than sending an unnormalized event that bypasses body
/// rules or replays a Responses Lite static prefix.
pub(crate) fn normalize_response_create(
&self,
client_event: &serde_json::Value,
) -> Option<serde_json::Value> {
use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};
let source_model = client_event
.get("model")
.and_then(serde_json::Value::as_str)
.unwrap_or(self.requested_model.as_str());
// The first response.create on a socket is a normal Responses request.
// Only a non-empty previous_response_id denotes a continuation whose
// stored history already contains the synthetic Responses Lite
// tools/instructions prefix. Keep this discriminator explicit instead
// of applying continuation edits to every socket turn.
let websocket_continuation = client_event
.get("previous_response_id")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
let require_body_stream_field =
request_requires_body_stream_field(client_event, self.force_body_stream_field);
let mut body = if websocket_continuation {
build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation(
client_event,
&self.mapped_model,
self.upstream_is_stream,
self.force_body_stream_field,
self.provider_type.as_str(),
self.provider_api_format.as_str(),
self.body_rules.as_ref(),
&self.request_headers,
self.codex_model_capabilities.as_ref(),
false,
)
} else {
build_local_openai_responses_request_body_with_codex_model_capabilities(
client_event,
&self.mapped_model,
self.upstream_is_stream,
self.force_body_stream_field,
self.provider_type.as_str(),
self.provider_api_format.as_str(),
self.body_rules.as_ref(),
&self.request_headers,
self.codex_model_capabilities.as_ref(),
false,
)
}?;
if let Some(patch) = self.model_directive_patch.as_ref() {
crate::ai_serving::apply_model_directive_mapping_patch(&mut body, patch);
// The patch is a deep merge and may reintroduce `stream`.
enforce_provider_body_stream_policy(
&mut body,
self.provider_api_format.as_str(),
self.upstream_is_stream,
require_body_stream_field,
);
}
let finalization = crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: self.client_api_format.as_str(),
provider_api_format: self.provider_api_format.as_str(),
provider_type: self.provider_type.as_str(),
provider_model: self.mapped_model.as_str(),
source_model,
body_rules: self.body_rules.as_ref(),
upstream_is_stream: self.upstream_is_stream,
require_body_stream_field,
};
let finalized = if websocket_continuation {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation(
&mut body,
finalization,
self.codex_model_capabilities.as_ref(),
self.reasoning_replay_policy,
)
} else {
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut body,
finalization,
self.codex_model_capabilities.as_ref(),
self.reasoning_replay_policy,
)
};
finalized.ok()?;
Some(body)
}
}
fn update_normalization_bytes_digest(digest: &mut sha2::Sha256, value: &[u8]) {
use sha2::Digest as _;
digest.update((value.len() as u64).to_be_bytes());
digest.update(value);
}
fn update_normalization_string_digest(digest: &mut sha2::Sha256, value: &str) {
update_normalization_bytes_digest(digest, value.as_bytes());
}
fn update_normalization_optional_string_digest(digest: &mut sha2::Sha256, value: Option<&str>) {
use sha2::Digest as _;
match value {
Some(value) => {
digest.update([1]);
update_normalization_string_digest(digest, value);
}
None => digest.update([0]),
}
}
fn update_normalization_string_vec_digest(digest: &mut sha2::Sha256, values: &[String]) {
use sha2::Digest as _;
digest.update((values.len() as u64).to_be_bytes());
for value in values {
update_normalization_string_digest(digest, value);
}
}
fn update_normalization_optional_json_digest(
digest: &mut sha2::Sha256,
value: Option<&serde_json::Value>,
) {
use sha2::Digest as _;
match value {
Some(value) => {
digest.update([1]);
update_normalization_json_digest(digest, value);
}
None => digest.update([0]),
}
}
fn update_normalization_json_digest(digest: &mut sha2::Sha256, value: &serde_json::Value) {
use serde_json::Value;
use sha2::Digest as _;
match value {
Value::Null => digest.update(b"n"),
Value::Bool(value) => digest.update(if *value { b"t" } else { b"f" }),
Value::Number(value) => {
digest.update(b"d");
update_normalization_string_digest(digest, value.to_string().as_str());
}
Value::String(value) => {
digest.update(b"s");
update_normalization_string_digest(digest, value);
}
Value::Array(values) => {
digest.update(b"[");
digest.update((values.len() as u64).to_be_bytes());
for value in values {
update_normalization_json_digest(digest, value);
}
digest.update(b"]");
}
Value::Object(values) => {
digest.update(b"{");
digest.update((values.len() as u64).to_be_bytes());
let mut keys = values.keys().collect::<Vec<_>>();
keys.sort_unstable();
for key in keys {
update_normalization_string_digest(digest, key);
update_normalization_json_digest(digest, &values[key]);
}
digest.update(b"}");
}
}
}
fn update_normalization_body_rule_headers_digest(
digest: &mut sha2::Sha256,
headers: &http::HeaderMap,
body_rules: Option<&serde_json::Value>,
) {
use sha2::Digest as _;
let dependencies =
crate::ai_serving::transport::rules::body_rules_request_header_dependencies(body_rules);
digest.update((dependencies.len() as u64).to_be_bytes());
for name in dependencies {
update_normalization_string_digest(digest, name.as_str());
let value = headers
.get(name.as_str())
.and_then(|value| value.to_str().ok())
.map(str::trim);
update_normalization_optional_string_digest(digest, value);
}
}
fn update_normalization_codex_capabilities_digest(
digest: &mut sha2::Sha256,
capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>,
) {
use sha2::Digest as _;
let Some(capabilities) = capabilities else {
digest.update([0]);
return;
};
digest.update([1]);
digest.update([
u8::from(capabilities.use_responses_lite),
u8::from(capabilities.supports_reasoning_summary_parameter),
u8::from(capabilities.supports_parallel_tool_calls),
u8::from(capabilities.support_verbosity),
]);
update_normalization_optional_string_digest(
digest,
capabilities.default_reasoning_effort.as_deref(),
);
update_normalization_optional_string_digest(
digest,
capabilities.default_reasoning_summary.as_deref(),
);
update_normalization_string_vec_digest(digest, &capabilities.supported_reasoning_efforts);
update_normalization_optional_string_digest(digest, capabilities.default_verbosity.as_deref());
update_normalization_string_vec_digest(digest, &capabilities.supported_service_tiers);
}
#[cfg(test)]
mod continuation_fingerprint_tests {
use http::HeaderValue;
use serde_json::json;
use super::ResponsesWebSocketBodyNormalization;
use crate::ai_serving::OpenAiResponsesReasoningReplayPolicy;
#[test]
fn normalization_fingerprint_is_stable_for_json_object_key_order() {
let first = ResponsesWebSocketBodyNormalization::for_tests("provider-model")
.with_model_directive_patch_for_tests(json!({"reasoning": {"effort": "high"}, "x": 1}));
let second = ResponsesWebSocketBodyNormalization::for_tests("provider-model")
.with_model_directive_patch_for_tests(json!({"x": 1, "reasoning": {"effort": "high"}}));
assert_eq!(
first.continuation_fingerprint(),
second.continuation_fingerprint()
);
}
#[test]
fn normalization_fingerprint_changes_with_effective_contract() {
let base = ResponsesWebSocketBodyNormalization::for_tests("provider-model");
let changed_policy = base.clone().with_reasoning_replay_policy_for_tests(
OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque,
);
assert_ne!(
base.continuation_fingerprint(),
changed_policy.continuation_fingerprint()
);
let changed_patch = base
.clone()
.with_model_directive_patch_for_tests(json!({"reasoning": {"effort": "low"}}));
assert_ne!(
base.continuation_fingerprint(),
changed_patch.continuation_fingerprint()
);
}
#[test]
fn normalization_fingerprint_ignores_unrelated_volatile_request_headers() {
let body_rules = json!([{
"action": "set",
"path": "store",
"value": false,
"condition": {
"source": "request_headers",
"path": "x-contract",
"op": "eq",
"value": "enabled"
}
}]);
let mut first = ResponsesWebSocketBodyNormalization::for_tests("provider-model")
.with_body_rules_for_tests(body_rules);
first
.request_headers
.insert("x-contract", HeaderValue::from_static("enabled"));
first
.request_headers
.insert("x-request-id", HeaderValue::from_static("request-1"));
first
.request_headers
.insert("cf-ray", HeaderValue::from_static("edge-1"));
let mut second = first.clone();
second
.request_headers
.insert("x-request-id", HeaderValue::from_static("request-2"));
second
.request_headers
.insert("cf-ray", HeaderValue::from_static("edge-2"));
assert_eq!(
first.continuation_fingerprint(),
second.continuation_fingerprint(),
"headers that no body-rule condition reads must not invalidate a persisted continuation"
);
}
#[test]
fn normalization_fingerprint_tracks_headers_used_by_body_rule_conditions() {
let body_rules = json!([{
"action": "set",
"path": "store",
"value": false,
"condition": {
"source": "request_headers",
"path": "X-Contract",
"op": "eq",
"value": "enabled"
}
}]);
let mut first = ResponsesWebSocketBodyNormalization::for_tests("provider-model")
.with_body_rules_for_tests(body_rules.clone());
first
.request_headers
.insert("x-contract", HeaderValue::from_static("enabled"));
let mut second = ResponsesWebSocketBodyNormalization::for_tests("provider-model")
.with_body_rules_for_tests(body_rules);
second
.request_headers
.insert("x-contract", HeaderValue::from_static("disabled"));
assert_ne!(
first.continuation_fingerprint(),
second.continuation_fingerprint(),
"a header that controls an effective body-rule condition remains part of the contract"
);
}
}
/// Builds one upstream decision for a Responses WebSocket turn. The session
/// reuses this decision for same-model turns and invokes the planner again when
/// a later `response.create` changes the public model.
pub(crate) async fn maybe_build_responses_websocket_decision(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
auth_snapshot: Option<&crate::ai_serving::GatewayAuthApiKeySnapshot>,
body_json: &serde_json::Value,
excluded_key_ids: Option<&BTreeSet<String>>,
excluded_codex_account_ids: Option<&BTreeSet<String>>,
pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>,
) -> Result<Option<ResponsesWebSocketDecision>, GatewayError> {
let Some(spec) = resolve_stream_spec(crate::ai_serving::OPENAI_RESPONSES_STREAM_PLAN_KIND)
else {
return Ok(None);
};
let Some(input) = resolve_local_openai_responses_decision_input_with_snapshot(
state,
parts,
trace_id,
decision,
body_json,
spec.decision_kind,
auth_snapshot,
)
.await?
else {
return Ok(None);
};
// The continuation discriminator belongs to the WebSocket protocol, not
// provider body rules/redaction. Capture it before the planner creates its
// effective body so a rule cannot accidentally turn a valid chain into a
// first-turn Lite normalization pass.
let websocket_continuation = body_json
.get("previous_response_id")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
let body_json = input.effective_body_json(body_json);
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
while let Some(attempt) = source.next_attempt().await? {
// `next_attempt` may return with a distributed pool-key lease. Arm a
// guard before the first await so owner-task timeout/cancellation
// cannot strand that lease until its server-side TTL expires.
let mut planning_lease = ResponsesWebSocketPlanningLeaseGuard::new(
state,
attempt.eligible.orchestration.pool_key_lease.as_ref(),
);
if pinned_candidate.is_some_and(|pinned| !pinned.matches(&attempt.eligible.candidate)) {
planning_lease.release().await;
continue;
}
if excluded_key_ids
.is_some_and(|key_ids| key_ids.contains(attempt.eligible.candidate.key_id.as_str()))
{
planning_lease.release().await;
continue;
}
let Some(adapter) = responses_websocket_adapter(
&attempt.eligible.transport.provider.provider_type,
attempt.eligible.transport.provider.config.as_ref(),
) else {
planning_lease.release().await;
continue;
};
// Captured before `attempt` is consumed so a later continuation turn can
// reproduce this candidate's body normalization without re-planning.
let transport = std::sync::Arc::clone(&attempt.eligible.transport);
let effective_auth_type =
aether_provider_transport::auth::resolve_local_auth_type_for_transport_format(
transport.as_ref(),
);
let candidate_provider_api_format = attempt.eligible.provider_api_format.clone();
let payload = match maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode(
state,
parts,
trace_id,
body_json,
&input,
attempt,
spec,
websocket_continuation,
)
.await
{
Ok(Some(payload)) => payload,
Ok(None) => {
planning_lease.release().await;
continue;
}
Err(error) => {
planning_lease.release().await;
return Err(error);
}
};
if payload
.provider_type
.as_deref()
.is_some_and(|value| value.trim().eq_ignore_ascii_case("codex"))
&& crate::orchestration::codex_account_id_from_headers(
&payload.provider_request_headers,
)
.is_some_and(|account_id| {
excluded_codex_account_ids
.is_some_and(|account_ids| account_ids.contains(account_id))
})
{
planning_lease.release().await;
continue;
}
match codex_quota_breaker_blocks_candidate(
state,
payload.provider_type.as_deref(),
payload.key_id.as_deref(),
&payload.provider_request_headers,
)
.await
{
Ok(true) => {
planning_lease.release().await;
continue;
}
Ok(false) => {}
Err(error) => log_codex_quota_breaker_check_failure(&error),
}
if payload
.provider_type
.as_deref()
.is_some_and(|value| adapter.supports_provider_type(value))
&& payload.provider_api_format.as_deref().is_some_and(|value| {
crate::ai_serving::normalize_api_format_alias(value) == "openai:responses"
})
{
let mapped_model = payload.mapped_model.clone().unwrap_or_default();
let source_model = body_json
.get("model")
.and_then(serde_json::Value::as_str)
.unwrap_or(input.requested_model.as_str());
let normalization = ResponsesWebSocketBodyNormalization {
provider_type: transport.provider.provider_type.clone(),
provider_api_format: candidate_provider_api_format.clone(),
client_api_format: local_openai_responses_spec_metadata(spec)
.api_format
.to_string(),
requested_model: input.requested_model.clone(),
upstream_is_stream: payload.upstream_is_stream,
force_body_stream_field: endpoint_config_forces_body_stream_field(
transport.endpoint.config.as_ref(),
),
body_rules: transport.endpoint.body_rules.clone(),
request_headers: input.effective_headers(&parts.headers).clone(),
codex_model_capabilities: codex_model_capabilities_for_transport(
&transport,
candidate_provider_api_format.as_str(),
mapped_model.as_str(),
source_model,
),
reasoning_replay_policy: openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model.as_str(),
),
model_directive_patch: input
.model_directive_policy
.resolve_reasoning(
candidate_provider_api_format.as_str(),
Some(&input.requested_model),
)
.mapping_patch_for_mapped_model(mapped_model.as_str())
.ok()
.flatten(),
mapped_model,
};
let decision = ResponsesWebSocketDecision {
execution: payload,
adapter,
normalization,
effective_auth_type,
};
// The decision report context now carries the lease identity. The
// WebSocket ownership layer takes over before any further await.
planning_lease.disarm();
return Ok(Some(decision));
}
planning_lease.release().await;
}
Ok(None)
}
async fn release_responses_websocket_planning_lease(
state: &AppState,
lease: Option<&RuntimeLockLease>,
) -> bool {
let Some(lease) = lease else {
return true;
};
match crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease(
state.runtime_state.as_ref(),
lease,
)
.await
{
Ok(_) => true,
Err(error) => {
tracing::warn!(
error = ?error,
"gateway Responses WebSocket planner failed to release an unused pool key lease"
);
false
}
}
}
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
+19
View File
@@ -8,9 +8,14 @@ use crate::constants::{
API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
};
use crate::scheduler::candidate::SchedulerSkippedCandidate;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::GatewayError;
impl<'a> PlannerAppState<'a> {
/// `ordering_config` is the request's routing-policy derived scheduler
/// config (see `SchedulerOrderingConfig::from_routing_policy`). `None`
/// falls back to the runtime default.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates(
self,
api_format: &str,
@@ -21,6 +26,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates(
self.app().data.as_ref(),
@@ -33,10 +39,12 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity,
now_unix_secs,
enable_model_directives,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
self,
api_format: &str,
@@ -47,6 +55,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -64,10 +73,12 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs,
enable_model_directives,
None,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_operation(
self,
api_format: &str,
@@ -79,6 +90,7 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -103,6 +115,7 @@ impl<'a> PlannerAppState<'a> {
attempt_now_unix_secs,
enable_model_directives,
request_operation,
ordering_config,
)
.await?;
@@ -123,6 +136,7 @@ impl<'a> PlannerAppState<'a> {
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
self,
api_format: &str,
@@ -132,6 +146,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -148,10 +163,12 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
self,
candidate_api_format: &str,
@@ -160,6 +177,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
@@ -176,6 +194,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot,
client_session_affinity,
attempt_now_unix_secs,
ordering_config,
)
.await?;
apps/aether-gateway/src/ai_serving/pure/mod.rs
+27 -21
View File
@@ -9,6 +9,7 @@ pub(crate) use aether_ai_formats::api::{
apply_codex_openai_responses_lite_header_with_capabilities,
apply_codex_openai_responses_special_body_edits,
apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities,
apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities,
apply_codex_openai_special_headers, apply_model_directive_mapping_patch,
apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request,
apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body,
@@ -37,6 +38,7 @@ pub(crate) use aether_ai_formats::api::{
build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives,
build_standard_request_body_with_model_directives_and_request_headers,
build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy,
calculate_kiro_context_input_tokens, canonicalize_tool_arguments,
convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
@@ -54,6 +56,8 @@ pub(crate) use aether_ai_formats::api::{
enforce_request_body_stream_field, estimate_kiro_tokens, extract_openai_text_content,
finalize_openai_provider_request,
finalize_openai_provider_request_with_codex_model_capabilities,
finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy,
finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, forbid_upstream_streaming_for_provider,
force_upstream_streaming_for_provider, gemini_request_is_image_generation,
@@ -82,10 +86,10 @@ pub(crate) use aether_ai_formats::api::{
parse_codex_auth_identity, parse_direct_request_body, parse_model_directive,
parse_model_directive_with_suffixes, parse_openai_stop_sequences,
parse_openai_tool_result_content, prepare_local_success_response_parts,
prepare_local_success_response_parts_owned, project_codex_openai_image_api_request_body,
project_openai_image_api_request_body, provider_adaptation_allows_sync_finalize_envelope,
provider_adaptation_anchor_api_format, provider_adaptation_descriptor_for_envelope,
provider_adaptation_descriptor_for_provider_type,
prepare_local_success_response_parts_owned, project_codex_catalog_model_card,
project_codex_openai_image_api_request_body, project_openai_image_api_request_body,
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format,
provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type,
provider_adaptation_requires_eventstream_accept,
provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, record_converted_response_history,
@@ -123,9 +127,9 @@ pub(crate) use aether_ai_formats::api::{
OpenAIResponsesProviderState, OpenAiImageNormalizeOptions, OpenAiImageOperation,
OpenAiImageRequestForGemini, OpenAiImageResponseFormat, OpenAiImageStreamState,
OpenAiImageSyncFinalizeProduct, OpenAiProviderRequestFinalization,
ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer,
ReasoningEffort, RequestConversionKind, ResponseHistoryRecord, ServiceTier,
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
OpenAiResponsesReasoningReplayPolicy, ProviderAdaptationDescriptor, ProviderAdaptationSurface,
ProviderPrivateStreamNormalizer, ReasoningEffort, RequestConversionKind, ResponseHistoryRecord,
ServiceTier, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,
@@ -134,14 +138,14 @@ pub(crate) use aether_ai_formats::api::{
CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND,
CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND,
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND,
CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT,
CODEX_OPENAI_IMAGE_INTERNAL_MODEL, EXECUTION_RUNTIME_STREAM_ACTION,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND,
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
CODEX_LIVE_STREAM_PLAN_KIND, CODEX_OPENAI_IMAGE_DEFAULT_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND,
GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
@@ -156,8 +160,8 @@ pub(crate) use aether_ai_formats::api::{
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
@@ -171,10 +175,12 @@ pub(crate) use aether_ai_formats::api::{
};
pub(crate) use aether_ai_formats::{
api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream,
api_format_permission_covers, intersect_api_format_allowed_lists, is_embedding_api_format,
is_rerank_api_format, openai_responses_request_operation,
openai_responses_synthetic_reasoning_item_id,
strip_incompatible_openai_responses_reasoning_items, ApiOperation, ClientSurface,
api_format_permission_covers, codex_responses_lite_tool_is_client_executed,
intersect_api_format_allowed_lists, is_embedding_api_format, is_rerank_api_format,
openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy, ApiOperation, ClientSurface,
CODEX_CLIENT_VERSION, OPENAI_RESPONSES_OPERATION_COMPACT,
};
pub(crate) fn plan_kind_matches_api_operation(
apps/aether-gateway/src/ai_serving/transport.rs
+6 -3
View File
@@ -59,7 +59,8 @@ pub(crate) mod windsurf {
}
pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_local_auth_config_header_overrides,
append_transport_diagnostics_to_value, apply_codex_fingerprint_convergence,
apply_codex_fingerprint_convergence_with_context, apply_local_auth_config_header_overrides,
apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
apply_standard_provider_request_body_rules_with_request_headers,
@@ -74,6 +75,7 @@ pub(crate) use aether_provider_transport::{
build_request_trace_proxy_value, build_same_format_provider_headers,
build_same_format_provider_request_body,
build_same_format_provider_request_body_with_compatibility_report,
build_same_format_provider_request_body_with_compatibility_report_and_reasoning_replay_policy,
build_same_format_provider_upstream_url, build_standard_plan_fallback_headers,
build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
@@ -105,8 +107,9 @@ pub(crate) use aether_provider_transport::{
supports_local_generic_oauth_request_auth_resolution,
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
transport_supports_api_operation, video_create_transport_unsupported_reason,
AnthropicCompatibilityProfile, CandidateTransportPolicyFacts, GatewayProviderTransportSnapshot,
GeminiCliRequestAuth, GeminiCliRequestAuthSupport, GeminiCliRequestAuthUnsupportedReason,
AnthropicCompatibilityProfile, CandidateTransportPolicyFacts,
CodexFingerprintConvergenceContext, GatewayProviderTransportSnapshot, GeminiCliRequestAuth,
GeminiCliRequestAuthSupport, GeminiCliRequestAuthUnsupportedReason,
GeminiCliRequestEnvelopeSupport, GeminiFilesHeadersInput, GeminiFilesRequestBodyError,
GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth,
ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
apps/aether-gateway/src/api/ai/openai.rs
+4
View File
@@ -5,9 +5,11 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("openai:rerank"),
"openai:responses" => Some("openai:responses"),
"openai:responses:compact" => Some("openai:responses:compact"),
"openai:realtime" => Some("openai:realtime"),
"openai:search" => Some("openai:search"),
"openai:image" => Some("openai:image"),
"openai:video" => Some("openai:video"),
"codex:live" => Some("codex:live"),
_ => None,
}
}
@@ -19,9 +21,11 @@ pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("/v1/rerank"),
"openai:responses" => Some("/v1/responses"),
"openai:responses:compact" => Some("/v1/responses/compact"),
"openai:realtime" => Some("/v1/realtime"),
"openai:search" => Some("/v1/alpha/search"),
"openai:image" => Some("/v1/images/generations"),
"openai:video" => Some("/v1/videos"),
"codex:live" => Some("/v1/live"),
_ => None,
}
}
apps/aether-gateway/src/api/ai/registry.rs
+145 -6
View File
@@ -1,13 +1,20 @@
use std::net::SocketAddr;
use axum::body::Body;
use axum::extract::Request;
use axum::http::{header, HeaderValue, Response, StatusCode};
use axum::routing::{any, post};
use axum::extract::ws::WebSocketUpgrade;
use axum::extract::{ConnectInfo, Request, State};
use axum::http::{header, HeaderMap, HeaderValue, Response, StatusCode, Uri};
use axum::routing::{any, get, post};
use axum::Router;
use super::{aliyun, claude, doubao, gemini, jina, openai};
use crate::api::response::build_local_http_error_response_with_request_path;
use crate::headers::extract_or_generate_trace_id;
use crate::{handlers::proxy::proxy_request, state::AppState, GatewayError};
use crate::{
handlers::proxy::{live_websocket, proxy_request, realtime_websocket, responses_websocket},
state::AppState,
GatewayError,
};
// Router registration patterns live here so AI public ingress has a single mount registry.
// They intentionally stay separate from manifest-facing route inventories in constants.rs,
@@ -18,6 +25,8 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/rerank",
"/v1/responses",
"/v1/responses/compact",
"/v1/live",
"/v1/realtime/calls",
"/v1/alpha/search",
"/v1/images/generations",
"/v1/images/edits",
@@ -51,8 +60,16 @@ const AI_ANY_ROUTE_PATTERNS: &[&str] = &[
pub(crate) fn mount_ai_routes(mut router: Router<AppState>) -> Router<AppState> {
for path in AI_POST_ROUTE_PATTERNS {
router = router.route(path, post(proxy_request));
router = if *path == "/v1/responses" {
router.route(path, get(responses_websocket).post(proxy_request))
} else if *path == "/v1/live" {
router.route(path, get(live_websocket).post(proxy_request))
} else {
router.route(path, post(proxy_request))
};
}
router = router.route("/v1/live/{call_id}", get(live_websocket));
router = router.route("/v1/realtime", get(dispatch_realtime_websocket));
for path in CLAUDE_POST_ROUTE_PATTERNS {
router = router.route(
path,
@@ -65,6 +82,61 @@ pub(crate) fn mount_ai_routes(mut router: Router<AppState>) -> Router<AppState>
router
}
async fn dispatch_realtime_websocket(
State(state): State<AppState>,
ConnectInfo(remote_addr): ConnectInfo<SocketAddr>,
ws: WebSocketUpgrade,
headers: HeaderMap,
uri: Uri,
) -> Result<Response<Body>, GatewayError> {
if realtime_query_is_codex_live(uri.query(), &headers) {
live_websocket(State(state), ConnectInfo(remote_addr), ws, headers, uri).await
} else {
realtime_websocket(State(state), ConnectInfo(remote_addr), ws, headers, uri).await
}
}
fn realtime_query_is_codex_live(query: Option<&str>, headers: &HeaderMap) -> bool {
let mut has_call_id = false;
let mut has_live_intent = false;
let mut duplicate_or_conflicting_intent = false;
for (key, value) in url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()) {
if key.eq_ignore_ascii_case("call_id") {
has_call_id = true;
} else if key.eq_ignore_ascii_case("intent") {
if has_live_intent || !value.eq_ignore_ascii_case("quicksilver") {
duplicate_or_conflicting_intent = true;
}
has_live_intent = true;
}
}
if has_live_intent {
return !duplicate_or_conflicting_intent;
}
// `call_id` is part of the ordinary OpenAI Realtime WebRTC sideband
// contract too. Without Codex's explicit v1 intent it must remain on the
// generic Realtime handler instead of being authorized as `codex:live`.
if has_call_id {
return false;
}
// Realtime v2 has no intent selector. A malformed or conflicting intent
// must not be reclassified as v2 merely because a Codex originator is
// present.
let has_model = url::form_urlencoded::parse(query.unwrap_or_default().as_bytes())
.any(|(key, value)| key.eq_ignore_ascii_case("model") && !value.trim().is_empty());
let Some(originator) = crate::headers::header_value_str(headers, "originator") else {
return false;
};
has_model
&& originator.split_whitespace().next().is_some_and(|value| {
value.eq_ignore_ascii_case("codex_cli_rs")
|| value.to_ascii_lowercase().starts_with("codex_cli_rs/")
|| value.eq_ignore_ascii_case("codex_work_desktop")
|| value.eq_ignore_ascii_case("codex_work_web")
|| value.eq_ignore_ascii_case("codex_work_mobile")
})
}
async fn claude_method_not_allowed(request: Request) -> Result<Response<Body>, GatewayError> {
let trace_id = extract_or_generate_trace_id(request.headers());
let mut response = build_local_http_error_response_with_request_path(
@@ -120,7 +192,72 @@ pub(crate) fn admin_default_body_rules_for_signature(
#[cfg(test)]
mod tests {
use super::{admin_endpoint_signature_parts, public_api_format_local_path};
use axum::http::{HeaderMap, HeaderValue};
use super::{
admin_endpoint_signature_parts, public_api_format_local_path, realtime_query_is_codex_live,
AI_POST_ROUTE_PATTERNS,
};
#[test]
fn registers_and_dispatches_realtime_live_aliases() {
assert!(AI_POST_ROUTE_PATTERNS.contains(&"/v1/realtime/calls"));
let no_headers = HeaderMap::new();
assert!(realtime_query_is_codex_live(
Some("intent=quicksilver&call_id=rtc_opaque"),
&no_headers
));
assert!(realtime_query_is_codex_live(
Some("intent=quicksilver&c%61ll_id=rtc_encoded"),
&no_headers
));
assert!(!realtime_query_is_codex_live(
Some("c%61ll_id=rtc_encoded"),
&no_headers
));
assert!(!realtime_query_is_codex_live(None, &no_headers));
assert!(!realtime_query_is_codex_live(Some("call_id="), &no_headers));
assert!(!realtime_query_is_codex_live(
Some("call_id=%20"),
&no_headers
));
assert!(realtime_query_is_codex_live(
Some("intent=quicksilver&model=gpt-realtime-1.5"),
&no_headers
));
assert!(!realtime_query_is_codex_live(
Some("model=gpt-realtime-1.5"),
&no_headers
));
assert!(!realtime_query_is_codex_live(
Some("intent=other&model=gpt-realtime-1.5"),
&no_headers
));
assert!(!realtime_query_is_codex_live(
Some("intent=quicksilver&intent=other&model=gpt-realtime-1.5"),
&no_headers
));
let mut codex_v2_headers = HeaderMap::new();
codex_v2_headers.insert("originator", HeaderValue::from_static("codex_work_desktop"));
assert!(realtime_query_is_codex_live(
Some("model=gpt-live-1-codex"),
&codex_v2_headers
));
assert!(!realtime_query_is_codex_live(
Some("call_id=rtc_ordinary&model=gpt-live-1-codex"),
&codex_v2_headers
));
assert!(!realtime_query_is_codex_live(
Some("call_id=rtc_one&call_id=rtc_two"),
&codex_v2_headers
));
let mut ordinary_headers = HeaderMap::new();
ordinary_headers.insert("originator", HeaderValue::from_static("openai-python"));
assert!(!realtime_query_is_codex_live(
Some("model=gpt-realtime-1.5"),
&ordinary_headers
));
}
#[test]
fn supports_data_api_endpoint_signatures_and_public_paths() {
@@ -148,6 +285,8 @@ mod tests {
),
("openai:rerank", "openai", "rerank", "/v1/rerank"),
("openai:search", "openai", "search", "/v1/alpha/search"),
("openai:realtime", "openai", "realtime", "/v1/realtime"),
("codex:live", "codex", "live", "/v1/live"),
("jina:rerank", "jina", "rerank", "/v1/rerank"),
] {
assert_eq!(
apps/aether-gateway/src/api/core.rs
+34
View File
@@ -50,12 +50,40 @@ pub(crate) async fn health(State(state): State<AppState>) -> impl IntoResponse {
"rejected": snapshot.rejected,
})
});
let websocket_connection_concurrency =
state
.websocket_connection_concurrency_snapshot()
.map(|snapshot| {
json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
let distributed_websocket_connection_concurrency = state
.distributed_websocket_connection_concurrency_snapshot()
.await
.ok()
.flatten()
.map(|snapshot| {
json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
Json(json!({
"status": "ok",
"component": "aether-gateway",
"control_api_enabled": true,
"request_concurrency": request_concurrency,
"distributed_request_concurrency": distributed_request_concurrency,
"websocket_connection_concurrency": websocket_connection_concurrency,
"distributed_websocket_connection_concurrency": distributed_websocket_connection_concurrency,
}))
}
@@ -113,6 +141,12 @@ pub(crate) async fn frontdoor_manifest(State(state): State<AppState>) -> impl In
"execution_runtime_configured": state.execution_runtime_configured(),
"request_concurrency_enabled": state.request_concurrency_snapshot().is_some(),
"distributed_request_concurrency_enabled": state.distributed_request_gate.is_some(),
"websocket_connection_concurrency_enabled": state
.websocket_connection_concurrency_snapshot()
.is_some(),
"distributed_websocket_connection_concurrency_enabled": state
.distributed_websocket_connection_gate
.is_some(),
"frontdoor_cors_enabled": cors_enabled,
"frontdoor_cors_allow_credentials": cors_allow_credentials,
"frontdoor_cors_allowed_origins": cors_allowed_origins,
apps/aether-gateway/src/backup/worker.rs
+3 -15
View File
@@ -85,29 +85,17 @@ async fn read_last_backup_slot(app: &AppState) -> Result<Option<String>, Gateway
#[cfg(test)]
mod tests {
use crate::backup::schedule::{BackupSchedule, BackupScheduleUnit};
use crate::task_runtime::{task_definition, TASK_KEY_SYSTEM_S3_BACKUP};
#[test]
fn backup_worker_skips_already_recorded_slot() {
let schedule = BackupSchedule {
unit: BackupScheduleUnit::Days,
interval: 1,
minute: 0,
hour: 3,
weekday: 1,
month_day: 1,
};
let now = chrono::DateTime::parse_from_rfc3339("2026-05-24T03:00:30+08:00")
.unwrap()
.with_timezone(&chrono::Utc);
let slot = schedule.due_slot(now).expect("slot should be due");
let slot = "days:2026-05-23T19:00:00Z";
assert!(super::should_start_scheduled_backup(
Some("days:2026-05-22T19:00:00Z"),
&slot
slot
));
assert!(!super::should_start_scheduled_backup(Some(&slot), &slot));
assert!(!super::should_start_scheduled_backup(Some(slot), slot));
}
#[test]
apps/aether-gateway/src/bin/aether-codex-ws-probe.rs
+124
View File
@@ -0,0 +1,124 @@
//! Credential-safe compatibility probe for the Codex Responses WebSocket path.
//!
//! This binary preserves the established Codex CLI and environment contract.
//! The common Responses WebSocket flow lives in `support/responses_ws_probe`;
//! this profile owns only Codex authentication and header requirements.
#[path = "support/responses_ws_probe.rs"]
mod responses_ws_probe;
use aether_gateway::{CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT};
use clap::Parser;
use http::header::{AUTHORIZATION, USER_AGENT};
use http::{HeaderMap, HeaderName, HeaderValue};
use responses_ws_probe::{
bearer_authorization_value, required_env, resolve_probe_url, run_profile_probe, turn_timeout,
ProbeArgs, ProbeConfig, ProbeFailure, ResponsesWebSocketProbeProfile,
};
const ACCESS_TOKEN_ENV: &str = "AETHER_CODEX_WS_PROBE_ACCESS_TOKEN";
const ACCOUNT_ID_ENV: &str = "AETHER_CODEX_WS_PROBE_ACCOUNT_ID";
const MODEL_ENV: &str = "AETHER_CODEX_WS_PROBE_MODEL";
const URL_ENV: &str = "AETHER_CODEX_WS_PROBE_URL";
#[derive(Parser)]
#[command(
name = "aether-codex-ws-probe",
about = "Verify a Codex Responses WebSocket endpoint without exposing credentials"
)]
struct Args {
/// WebSocket endpoint. If omitted, AETHER_CODEX_WS_PROBE_URL is used.
#[arg(long)]
url: Option<String>,
/// Per-turn receive timeout in seconds.
#[arg(long, default_value_t = 20, value_parser = clap::value_parser!(u64).range(1..=120))]
timeout_secs: u64,
}
impl From<Args> for ProbeArgs {
fn from(args: Args) -> Self {
Self {
url: args.url,
timeout_secs: args.timeout_secs,
}
}
}
struct CodexResponsesProbeProfile;
impl ResponsesWebSocketProbeProfile for CodexResponsesProbeProfile {
fn build_config(args: &ProbeArgs) -> Result<ProbeConfig, ProbeFailure> {
let url = resolve_probe_url(args, URL_ENV, None)?;
let access_token = required_env(ACCESS_TOKEN_ENV)?;
let account_id = required_env(ACCOUNT_ID_ENV)?;
let model = required_env(MODEL_ENV)?;
Ok(ProbeConfig::new(
url,
model,
turn_timeout(args),
handshake_headers(&access_token, &account_id)?,
Self::sent_header_names(),
))
}
fn sent_header_names() -> Vec<&'static str> {
vec![
"authorization",
"chatgpt-account-id",
"user-agent",
"originator",
]
}
}
fn handshake_headers(access_token: &str, account_id: &str) -> Result<HeaderMap, ProbeFailure> {
let account_id =
HeaderValue::from_str(account_id).map_err(|_| ProbeFailure::MissingConfiguration)?;
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, bearer_authorization_value(access_token)?);
headers.insert(HeaderName::from_static("chatgpt-account-id"), account_id);
headers.insert(
USER_AGENT,
HeaderValue::from_static(CODEX_CLIENT_USER_AGENT),
);
headers.insert(
HeaderName::from_static("originator"),
HeaderValue::from_static(CODEX_CLIENT_ORIGINATOR),
);
Ok(headers)
}
#[tokio::main]
async fn main() {
let exit_code = run_profile_probe::<CodexResponsesProbeProfile>(Args::parse().into()).await;
if exit_code != 0 {
std::process::exit(exit_code);
}
}
#[cfg(test)]
mod tests {
use http::header::{AUTHORIZATION, USER_AGENT};
use super::{handshake_headers, CodexResponsesProbeProfile, ResponsesWebSocketProbeProfile};
#[test]
fn codex_profile_keeps_its_required_handshake_headers() {
let headers =
handshake_headers("test-token", "test-account").expect("headers should build");
assert!(headers.contains_key(AUTHORIZATION));
assert!(headers.contains_key("chatgpt-account-id"));
assert!(headers.contains_key(USER_AGENT));
assert!(headers.contains_key("originator"));
assert_eq!(
CodexResponsesProbeProfile::sent_header_names(),
vec![
"authorization",
"chatgpt-account-id",
"user-agent",
"originator",
]
);
}
}
apps/aether-gateway/src/bin/aether-openai-responses-ws-probe.rs
+104
View File
@@ -0,0 +1,104 @@
//! Credential-safe compatibility probe for the official OpenAI Responses
//! WebSocket endpoint.
//!
//! This profile uses standard API-key Bearer authentication and shares the
//! protocol flow with the Codex probe without inheriting Codex-specific
//! account headers or quota assumptions.
#[path = "support/responses_ws_probe.rs"]
mod responses_ws_probe;
use clap::Parser;
use http::header::AUTHORIZATION;
use http::HeaderMap;
use responses_ws_probe::{
bearer_authorization_value, required_env, resolve_probe_url, run_profile_probe, turn_timeout,
ProbeArgs, ProbeConfig, ProbeFailure, ResponsesWebSocketProbeProfile,
};
const API_KEY_ENV: &str = "AETHER_OPENAI_WS_PROBE_API_KEY";
const MODEL_ENV: &str = "AETHER_OPENAI_WS_PROBE_MODEL";
const URL_ENV: &str = "AETHER_OPENAI_WS_PROBE_URL";
const DEFAULT_URL: &str = "wss://api.openai.com/v1/responses";
#[derive(Parser)]
#[command(
name = "aether-openai-responses-ws-probe",
about = "Verify an OpenAI Responses WebSocket endpoint without exposing credentials"
)]
struct Args {
/// WebSocket endpoint. If omitted, AETHER_OPENAI_WS_PROBE_URL or the
/// official OpenAI endpoint is used.
#[arg(long)]
url: Option<String>,
/// Per-turn receive timeout in seconds.
#[arg(long, default_value_t = 20, value_parser = clap::value_parser!(u64).range(1..=120))]
timeout_secs: u64,
}
impl From<Args> for ProbeArgs {
fn from(args: Args) -> Self {
Self {
url: args.url,
timeout_secs: args.timeout_secs,
}
}
}
struct OpenAiResponsesProbeProfile;
impl ResponsesWebSocketProbeProfile for OpenAiResponsesProbeProfile {
fn build_config(args: &ProbeArgs) -> Result<ProbeConfig, ProbeFailure> {
let url = resolve_probe_url(args, URL_ENV, Some(DEFAULT_URL))?;
let api_key = required_env(API_KEY_ENV)?;
let model = required_env(MODEL_ENV)?;
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, bearer_authorization_value(&api_key)?);
Ok(ProbeConfig::new(
url,
model,
turn_timeout(args),
headers,
Self::sent_header_names(),
))
}
fn sent_header_names() -> Vec<&'static str> {
vec!["authorization"]
}
}
#[tokio::main]
async fn main() {
let exit_code = run_profile_probe::<OpenAiResponsesProbeProfile>(Args::parse().into()).await;
if exit_code != 0 {
std::process::exit(exit_code);
}
}
#[cfg(test)]
mod tests {
use http::header::AUTHORIZATION;
use super::{
bearer_authorization_value, responses_ws_probe::parse_probe_url,
OpenAiResponsesProbeProfile, ResponsesWebSocketProbeProfile, DEFAULT_URL,
};
#[test]
fn openai_profile_exposes_only_standard_bearer_authentication() {
let authorization = bearer_authorization_value("test-key").expect("header should build");
assert_eq!(authorization.to_str().ok(), Some("Bearer test-key"));
assert_eq!(
OpenAiResponsesProbeProfile::sent_header_names(),
vec![AUTHORIZATION.as_str()]
);
}
#[test]
fn openai_profile_uses_the_official_responses_websocket_endpoint_by_default() {
let url = parse_probe_url(DEFAULT_URL).expect("default OpenAI endpoint should be valid");
assert_eq!(url.as_str(), DEFAULT_URL);
}
}
apps/aether-gateway/src/bin/support/responses_ws_probe.rs
+567
View File
@@ -0,0 +1,567 @@
//! Shared, credential-safe engine for Responses WebSocket compatibility probes.
//!
//! Provider profiles own their environment variables and handshake headers.
//! This module owns the common Responses WebSocket contract: two sequential
//! `response.create` warmups, continuation with `previous_response_id`, safe
//! event observation, and a redacted JSON report.
use std::env;
use std::time::{Duration, Instant};
use http::{HeaderMap, HeaderValue};
use serde::Serialize;
use serde_json::{json, Value};
use url::Url;
use wreq::ws::message::Message as WreqWsMessage;
const MAX_FRAME_SIZE: usize = 1 << 20;
const MAX_EVENTS_PER_TURN: usize = 16;
pub(crate) struct ProbeArgs {
pub(crate) url: Option<String>,
pub(crate) timeout_secs: u64,
}
pub(crate) struct ProbeConfig {
url: Url,
model: String,
turn_timeout: Duration,
handshake_headers: HeaderMap,
sent_header_names: Vec<&'static str>,
}
impl ProbeConfig {
pub(crate) fn new(
url: Url,
model: String,
turn_timeout: Duration,
handshake_headers: HeaderMap,
sent_header_names: Vec<&'static str>,
) -> Self {
Self {
url,
model,
turn_timeout,
handshake_headers,
sent_header_names,
}
}
}
/// A profile retains provider-specific authentication and configuration while
/// reusing one Responses protocol probe engine.
pub(crate) trait ResponsesWebSocketProbeProfile {
fn build_config(args: &ProbeArgs) -> Result<ProbeConfig, ProbeFailure>;
fn sent_header_names() -> Vec<&'static str>;
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum ProbeFailure {
MissingConfiguration,
InvalidEndpoint,
ClientBuild,
Handshake,
Upgrade,
Send,
ReceiveTimeout,
Receive,
RemoteError,
MissingResponseId,
UnexpectedFrame,
}
impl ProbeFailure {
const fn code(self) -> &'static str {
match self {
Self::MissingConfiguration => "missing_configuration",
Self::InvalidEndpoint => "invalid_endpoint",
Self::ClientBuild => "client_build_failed",
Self::Handshake => "handshake_failed",
Self::Upgrade => "upgrade_failed",
Self::Send => "send_failed",
Self::ReceiveTimeout => "receive_timeout",
Self::Receive => "receive_failed",
Self::RemoteError => "upstream_error_event",
Self::MissingResponseId => "response_id_not_observed",
Self::UnexpectedFrame => "unexpected_frame",
}
}
}
#[derive(Serialize)]
struct ProbeReport {
status: &'static str,
target_host: Option<String>,
handshake_status: Option<u16>,
sent_header_names: Vec<&'static str>,
received_header_names: Vec<String>,
observed_event_types: Vec<String>,
continuation_confirmed: bool,
elapsed_ms: u64,
error: Option<&'static str>,
}
impl ProbeReport {
fn failed(
config: Option<&ProbeConfig>,
sent_header_names: Vec<&'static str>,
started_at: Instant,
error: ProbeFailure,
) -> Self {
Self {
status: "failed",
target_host: config.and_then(target_host),
handshake_status: None,
sent_header_names,
received_header_names: Vec::new(),
observed_event_types: Vec::new(),
continuation_confirmed: false,
elapsed_ms: started_at.elapsed().as_millis() as u64,
error: Some(error.code()),
}
}
}
/// Runs a profile and returns the process exit code after emitting exactly one
/// credential-safe JSON report.
pub(crate) async fn run_profile_probe<P: ResponsesWebSocketProbeProfile>(args: ProbeArgs) -> i32 {
let started_at = Instant::now();
let config = match P::build_config(&args) {
Ok(config) => config,
Err(error) => {
print_report(&ProbeReport::failed(
None,
P::sent_header_names(),
started_at,
error,
));
return 2;
}
};
match run_probe(&config, started_at).await {
Ok(report) => {
print_report(&report);
0
}
Err(error) => {
print_report(&ProbeReport::failed(
Some(&config),
config.sent_header_names.clone(),
started_at,
error,
));
1
}
}
}
pub(crate) fn required_env(name: &str) -> Result<String, ProbeFailure> {
env::var(name)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.ok_or(ProbeFailure::MissingConfiguration)
}
pub(crate) fn resolve_probe_url(
args: &ProbeArgs,
url_env: &str,
default_url: Option<&str>,
) -> Result<Url, ProbeFailure> {
let raw_url = args
.url
.as_deref()
.map(str::to_owned)
.or_else(|| env::var(url_env).ok())
.or_else(|| default_url.map(str::to_owned));
let Some(raw_url) = raw_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Err(ProbeFailure::MissingConfiguration);
};
parse_probe_url(raw_url)
}
pub(crate) fn parse_probe_url(raw: &str) -> Result<Url, ProbeFailure> {
let url = Url::parse(raw).map_err(|_| ProbeFailure::InvalidEndpoint)?;
if !matches!(url.scheme(), "ws" | "wss")
|| url.host_str().is_none()
|| !url.username().is_empty()
|| url.password().is_some()
|| url.query().is_some()
|| url.fragment().is_some()
{
return Err(ProbeFailure::InvalidEndpoint);
}
Ok(url)
}
pub(crate) fn bearer_authorization_value(token: &str) -> Result<HeaderValue, ProbeFailure> {
HeaderValue::from_str(format!("Bearer {token}").as_str())
.map_err(|_| ProbeFailure::MissingConfiguration)
}
pub(crate) const fn turn_timeout(args: &ProbeArgs) -> Duration {
Duration::from_secs(args.timeout_secs)
}
async fn run_probe(config: &ProbeConfig, started_at: Instant) -> Result<ProbeReport, ProbeFailure> {
let client = wreq::Client::builder()
.connect_timeout(config.turn_timeout)
.timeout(config.turn_timeout)
.build()
.map_err(|_| ProbeFailure::ClientBuild)?;
let response = client
.websocket(config.url.as_str())
.headers(config.handshake_headers.clone())
.max_frame_size(MAX_FRAME_SIZE)
.max_message_size(MAX_FRAME_SIZE)
.send()
.await
.map_err(|_| ProbeFailure::Handshake)?;
let handshake_status = response.status().as_u16();
let received_header_names = response
.headers()
.keys()
.map(|name| name.as_str().to_string())
.collect();
let mut socket = response
.into_websocket()
.await
.map_err(|_| ProbeFailure::Upgrade)?;
let mut observed_event_types = Vec::new();
send_warmup(&mut socket, &config.model, None).await?;
let first_response_id =
receive_completed_response_id(&mut socket, config.turn_timeout, &mut observed_event_types)
.await?;
send_warmup(&mut socket, &config.model, Some(&first_response_id)).await?;
let _second_response_id =
receive_completed_response_id(&mut socket, config.turn_timeout, &mut observed_event_types)
.await?;
Ok(ProbeReport {
status: "passed",
target_host: target_host(config),
handshake_status: Some(handshake_status),
sent_header_names: config.sent_header_names.clone(),
received_header_names,
observed_event_types,
continuation_confirmed: true,
elapsed_ms: started_at.elapsed().as_millis() as u64,
error: None,
})
}
fn target_host(config: &ProbeConfig) -> Option<String> {
config.url.host_str().map(|host| match config.url.port() {
Some(port) => format!("{host}:{port}"),
None => host.to_string(),
})
}
async fn send_warmup(
socket: &mut wreq::ws::WebSocket,
model: &str,
previous_response_id: Option<&str>,
) -> Result<(), ProbeFailure> {
let mut event = json!({
"type": "response.create",
"model": model,
"store": false,
"generate": false,
"input": [],
"tools": [],
});
if let Some(previous_response_id) = previous_response_id {
event["previous_response_id"] = Value::String(previous_response_id.to_string());
}
socket
.send(WreqWsMessage::text(event.to_string()))
.await
.map_err(|_| ProbeFailure::Send)
}
async fn receive_completed_response_id(
socket: &mut wreq::ws::WebSocket,
timeout: Duration,
observed_event_types: &mut Vec<String>,
) -> Result<String, ProbeFailure> {
let mut response_id = None;
for _ in 0..MAX_EVENTS_PER_TURN {
let message = tokio::time::timeout(timeout, socket.recv())
.await
.map_err(|_| ProbeFailure::ReceiveTimeout)?
.ok_or(ProbeFailure::MissingResponseId)?
.map_err(|_| ProbeFailure::Receive)?;
match message {
WreqWsMessage::Text(text) => {
let event: Value = serde_json::from_str(text.as_str())
.map_err(|_| ProbeFailure::UnexpectedFrame)?;
let event_type = event
.get("type")
.and_then(Value::as_str)
.map(safe_event_label)
.unwrap_or_else(|| "unknown".to_string());
let is_remote_error = event_type == "error";
let is_completed = event_type == "response.completed";
observed_event_types.push(event_type);
if is_remote_error {
return Err(ProbeFailure::RemoteError);
}
if let Some(observed_response_id) = event
.pointer("/response/id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
response_id = Some(observed_response_id.to_string());
}
if is_completed {
return response_id.ok_or(ProbeFailure::MissingResponseId);
}
}
WreqWsMessage::Ping(_) | WreqWsMessage::Pong(_) => continue,
WreqWsMessage::Close(_) => return Err(ProbeFailure::MissingResponseId),
_ => return Err(ProbeFailure::UnexpectedFrame),
}
}
Err(ProbeFailure::MissingResponseId)
}
fn safe_event_label(value: &str) -> String {
let trimmed = value.trim();
if trimmed.is_empty()
|| trimmed.len() > 80
|| !trimmed
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
{
return "unknown".to_string();
}
trimmed.to_string()
}
fn print_report(report: &ProbeReport) {
match serde_json::to_string(report) {
Ok(json) => println!("{json}"),
Err(_) => println!("{{\"status\":\"failed\",\"error\":\"report_serialization_failed\"}}"),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::{Duration, Instant};
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::State;
use axum::http::header::AUTHORIZATION;
use axum::http::{HeaderMap, HeaderValue};
use axum::response::IntoResponse;
use axum::routing::get;
use axum::Router;
use futures_util::{SinkExt, StreamExt};
use serde_json::Value;
use tokio::sync::{oneshot, Mutex};
use super::{parse_probe_url, run_probe, ProbeConfig};
#[derive(Default)]
struct MockState {
observed: Mutex<Option<oneshot::Sender<ObservedClientMessages>>>,
}
struct ObservedClientMessages {
authorization_present: bool,
profile_header_present: bool,
second_before_first_completion: bool,
first: Value,
second: Value,
}
#[tokio::test]
async fn probe_confirms_sequential_response_continuation_without_exposing_values() {
let (url, observed, server) = spawn_mock_server().await;
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer test-token-that-must-not-be-reported"),
);
headers.insert(
"x-aether-probe-profile",
HeaderValue::from_static("test-profile-id"),
);
let config = ProbeConfig::new(
parse_probe_url(url.as_str()).expect("mock URL should be valid"),
"gpt-test".to_string(),
Duration::from_secs(2),
headers,
vec!["authorization", "x-aether-probe-profile"],
);
let report = run_probe(&config, Instant::now())
.await
.expect("probe should complete against mock server");
let client_messages = observed.await.expect("mock should observe client messages");
server.abort();
assert_eq!(report.status, "passed");
assert!(report.continuation_confirmed);
assert!(report
.observed_event_types
.contains(&"response.created".to_string()));
assert!(report
.observed_event_types
.contains(&"response.completed".to_string()));
assert!(client_messages.authorization_present);
assert!(client_messages.profile_header_present);
assert!(!client_messages.second_before_first_completion);
assert_eq!(client_messages.first["type"], "response.create");
assert_eq!(client_messages.first["generate"], false);
assert_eq!(client_messages.first["store"], false);
assert_eq!(client_messages.second["previous_response_id"], "resp-first");
let report_json = serde_json::to_string(&report).expect("report should serialize");
assert!(!report_json.contains("test-token-that-must-not-be-reported"));
assert!(!report_json.contains("test-profile-id"));
assert!(!report_json.contains("resp-first"));
}
#[test]
fn probe_url_rejects_credentials_and_query_strings() {
assert!(parse_probe_url("wss://example.test/v1/responses").is_ok());
assert!(parse_probe_url("https://example.test/v1/responses").is_err());
assert!(parse_probe_url("wss://token@example.test/v1/responses").is_err());
assert!(parse_probe_url("wss://example.test/v1/responses?token=secret").is_err());
}
async fn spawn_mock_server() -> (
String,
oneshot::Receiver<ObservedClientMessages>,
tokio::task::JoinHandle<()>,
) {
let (observed_tx, observed_rx) = oneshot::channel();
let state = Arc::new(MockState {
observed: Mutex::new(Some(observed_tx)),
});
let app = Router::new()
.route("/v1/responses", get(mock_websocket))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("mock listener should bind");
let address = listener
.local_addr()
.expect("mock listener should expose address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("mock server should run");
});
(format!("ws://{address}/v1/responses"), observed_rx, server)
}
async fn mock_websocket(
ws: WebSocketUpgrade,
State(state): State<Arc<MockState>>,
headers: HeaderMap,
) -> impl IntoResponse {
let authorization_present = headers
.get(AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.starts_with("Bearer "));
let profile_header_present = headers.contains_key("x-aether-probe-profile");
ws.on_upgrade(move |socket| async move {
serve_mock_socket(socket, state, authorization_present, profile_header_present).await;
})
}
async fn serve_mock_socket(
socket: WebSocket,
state: Arc<MockState>,
authorization_present: bool,
profile_header_present: bool,
) {
let (mut sender, mut receiver) = socket.split();
let first = receive_json(&mut receiver).await;
let _ = sender
.send(Message::Text(
serde_json::json!({
"type": "response.created",
"response": {"id": "resp-first"}
})
.to_string()
.into(),
))
.await;
let early_second = tokio::select! {
message = receiver.next() => Some(message),
_ = tokio::time::sleep(Duration::from_millis(50)) => None,
};
let second_before_first_completion = early_second.is_some();
let _ = sender
.send(Message::Text(
serde_json::json!({
"type": "response.completed",
"response": {"id": "resp-first", "status": "completed"}
})
.to_string()
.into(),
))
.await;
let second = match early_second {
Some(Some(Ok(Message::Text(text)))) => {
serde_json::from_str(text.as_str()).expect("early client message should be JSON")
}
Some(Some(Ok(_))) => panic!("expected text continuation message"),
Some(Some(Err(error))) => panic!("client message should be valid: {error}"),
Some(None) => panic!("client closed before continuation"),
None => receive_json(&mut receiver).await,
};
let _ = sender
.send(Message::Text(
serde_json::json!({
"type": "response.created",
"response": {"id": "resp-second"}
})
.to_string()
.into(),
))
.await;
let _ = sender
.send(Message::Text(
serde_json::json!({
"type": "response.completed",
"response": {"id": "resp-second", "status": "completed"}
})
.to_string()
.into(),
))
.await;
if let Some(observed) = state.observed.lock().await.take() {
let _ = observed.send(ObservedClientMessages {
authorization_present,
profile_header_present,
second_before_first_completion,
first,
second,
});
}
}
async fn receive_json(receiver: &mut futures_util::stream::SplitStream<WebSocket>) -> Value {
let message = receiver
.next()
.await
.expect("client should send a message")
.expect("client message should be valid");
let Message::Text(text) = message else {
panic!("expected text message");
};
serde_json::from_str(text.as_str()).expect("client message should be JSON")
}
}
apps/aether-gateway/src/cache/auth_runtime.rs
Vendored
-10
View File
@@ -1274,7 +1274,6 @@ mod tests {
let first_cache = Arc::clone(&cache);
let first_key = key.clone();
let first_calls = Arc::clone(&calls);
let first_started = Instant::now();
let first = tokio::spawn(async move {
first_cache
.get_or_load_once_stale_while_refreshing::<(), _, _>(
@@ -1292,7 +1291,6 @@ mod tests {
});
let follower_cache = Arc::clone(&cache);
let follower_started = Instant::now();
let follower_calls = Arc::clone(&calls);
let follower = tokio::spawn(async move {
follower_cache
@@ -1310,15 +1308,7 @@ mod tests {
});
assert_eq!(first.await.unwrap().unwrap(), Some(1));
assert!(
first_started.elapsed() < Duration::from_millis(80),
"stale value should not wait for request-path refresh"
);
assert_eq!(follower.await.unwrap().unwrap(), Some(1));
assert!(
follower_started.elapsed() < Duration::from_millis(80),
"follower should return stale value without waiting for refresh"
);
assert_eq!(calls.load(Ordering::Acquire), 0);
}
apps/aether-gateway/src/client_session_affinity.rs
+464 -25
View File
@@ -23,6 +23,21 @@ pub(crate) struct ClientSessionScope {
pub(crate) source: ClientSessionSignalSource,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct CodexRequestSignals {
pub(crate) session_id: Option<String>,
pub(crate) thread_id: Option<String>,
pub(crate) turn_id: Option<String>,
pub(crate) prompt_cache_key: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct CodexTurnMetadataSignals {
session_id: Option<String>,
thread_id: Option<String>,
turn_id: Option<String>,
}
impl ClientSessionScope {
fn new(
client_family: impl Into<String>,
@@ -136,6 +151,13 @@ pub(crate) fn client_session_scope_from_request(
.or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str()))
}
pub(crate) fn codex_request_signals_from_request(
headers: &http::HeaderMap,
body_json: Option<&Value>,
) -> CodexRequestSignals {
extract_codex_request_signals(&ClientSessionRequest { headers, body_json })
}
fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
let session_id = request
.body_json?
@@ -408,29 +430,7 @@ impl ClientSessionScopeAdapter for CodexSessionScopeAdapter {
}
fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
header_value_str(request.headers, "session-id")
.or_else(|| header_value_str(request.headers, "thread-id"))
.or_else(|| header_value_str(request.headers, "session_id"))
.or_else(|| header_value_str(request.headers, "conversation_id"))
.map(|root_session| {
ClientSessionScope::new(
self.family(),
root_session,
None,
header_value_str(request.headers, "chatgpt-account-id"),
ClientSessionSignalSource::Header,
)
})
.or_else(|| {
let body_session = GenericSessionScopeAdapter.extract_scope(request)?;
Some(ClientSessionScope::new(
self.family(),
body_session.session_id,
body_session.agent_id,
header_value_str(request.headers, "chatgpt-account-id"),
body_session.source,
))
})
codex_request_session_scope_from_request(request)
}
}
@@ -778,6 +778,162 @@ fn explicit_aether_session_scope(
))
}
fn extract_codex_request_signals(request: &ClientSessionRequest<'_>) -> CodexRequestSignals {
let body_client_metadata = request
.body_json
.and_then(|body| body.get("client_metadata"))
.and_then(Value::as_object);
let body_turn_metadata = codex_turn_metadata_signals(
body_client_metadata.and_then(|metadata| metadata.get("x-codex-turn-metadata")),
);
let header_turn_metadata = header_value_str(request.headers, "x-codex-turn-metadata")
.map(|raw| parse_codex_turn_metadata(&raw))
.unwrap_or_default();
let native_thread_id = header_value_str(request.headers, "thread-id")
.or_else(|| {
body_client_metadata
.and_then(|metadata| value_at_map_path(metadata, "thread_id"))
.map(ToOwned::to_owned)
})
.or_else(|| body_turn_metadata.thread_id.clone());
let turn_id = body_client_metadata
.and_then(|metadata| value_at_map_path(metadata, "turn_id"))
.map(ToOwned::to_owned)
.or_else(|| body_turn_metadata.turn_id.clone())
.or_else(|| {
request
.body_json
.and_then(|body| value_at_path(body, &["turn_id"]))
.map(ToOwned::to_owned)
})
.or(header_turn_metadata.turn_id);
let prompt_cache_key = request
.body_json
.and_then(|body| value_at_path(body, &["prompt_cache_key"]))
.map(ToOwned::to_owned);
let session_id =
codex_request_session_scope(request, &body_turn_metadata).map(|scope| scope.session_id);
let thread_id = native_thread_id.or_else(|| session_id.clone());
CodexRequestSignals {
session_id,
thread_id,
turn_id,
prompt_cache_key,
}
}
fn codex_request_session_scope_from_request(
request: &ClientSessionRequest<'_>,
) -> Option<ClientSessionScope> {
let body_turn_metadata = codex_turn_metadata_signals(
request
.body_json
.and_then(|body| body.get("client_metadata"))
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("x-codex-turn-metadata")),
);
codex_request_session_scope(request, &body_turn_metadata)
}
fn codex_request_session_scope(
request: &ClientSessionRequest<'_>,
body_turn_metadata: &CodexTurnMetadataSignals,
) -> Option<ClientSessionScope> {
if let Some(scope) = explicit_aether_session_scope(request, CodexSessionScopeAdapter.family()) {
return Some(scope);
}
if let Some(root_session) = header_value_str(request.headers, "session-id")
.or_else(|| header_value_str(request.headers, "thread-id"))
.or_else(|| header_value_str(request.headers, "session_id"))
.or_else(|| header_value_str(request.headers, "conversation_id"))
.or_else(|| header_value_str(request.headers, "x-session-id"))
{
return Some(codex_session_scope(
request,
root_session,
None,
ClientSessionSignalSource::Header,
));
}
let body_client_metadata = request
.body_json
.and_then(|body| body.get("client_metadata"))
.and_then(Value::as_object);
if let Some(root_session) = body_client_metadata
.and_then(|metadata| value_at_map_path(metadata, "session_id"))
.or_else(|| {
body_client_metadata.and_then(|metadata| value_at_map_path(metadata, "thread_id"))
})
.map(ToOwned::to_owned)
.or_else(|| body_turn_metadata.session_id.clone())
.or_else(|| body_turn_metadata.thread_id.clone())
{
return Some(codex_session_scope(
request,
root_session,
None,
ClientSessionSignalSource::Body,
));
}
let generic = GenericSessionScopeAdapter.extract_scope(request)?;
Some(codex_session_scope(
request,
generic.session_id,
generic.agent_id,
generic.source,
))
}
fn codex_session_scope(
request: &ClientSessionRequest<'_>,
session_id: String,
agent_id: Option<String>,
source: ClientSessionSignalSource,
) -> ClientSessionScope {
ClientSessionScope::new(
CodexSessionScopeAdapter.family(),
session_id,
agent_id,
header_value_str(request.headers, "chatgpt-account-id"),
source,
)
}
fn codex_turn_metadata_signals(value: Option<&Value>) -> CodexTurnMetadataSignals {
match value {
Some(Value::Object(metadata)) => codex_turn_metadata_signals_from_map(metadata),
Some(Value::String(raw)) => parse_codex_turn_metadata(raw),
_ => CodexTurnMetadataSignals::default(),
}
}
fn parse_codex_turn_metadata(raw: &str) -> CodexTurnMetadataSignals {
serde_json::from_str::<Map<String, Value>>(raw)
.map(|metadata| codex_turn_metadata_signals_from_map(&metadata))
.unwrap_or_default()
}
fn codex_turn_metadata_signals_from_map(metadata: &Map<String, Value>) -> CodexTurnMetadataSignals {
CodexTurnMetadataSignals {
session_id: value_at_map_path(metadata, "session_id").map(ToOwned::to_owned),
thread_id: value_at_map_path(metadata, "thread_id").map(ToOwned::to_owned),
turn_id: value_at_map_path(metadata, "turn_id").map(ToOwned::to_owned),
}
}
fn value_at_map_path<'a>(object: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn normalize_session_key(
account_hint: Option<&str>,
root_session: &str,
@@ -831,12 +987,25 @@ mod tests {
client_session_affinity_from_api_request,
client_session_affinity_from_report_context_value, client_session_affinity_from_request,
client_session_affinity_report_context_value, client_session_scope_from_request,
ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER,
codex_request_signals_from_request, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER,
AETHER_SESSION_ID_HEADER,
};
use aether_scheduler_core::ClientSessionAffinity;
use http::{HeaderMap, HeaderValue};
use http::{HeaderMap, HeaderName, HeaderValue};
use serde_json::json;
fn request_headers(values: &[(&str, &str)]) -> HeaderMap {
values
.iter()
.map(|(name, value)| {
(
HeaderName::from_bytes(name.as_bytes()).expect("valid test header name"),
HeaderValue::from_bytes(value.as_bytes()).expect("valid test header value"),
)
})
.collect()
}
#[test]
fn unknown_adapter_extracts_body_session_and_agent() {
let body = json!({
@@ -933,6 +1102,276 @@ mod tests {
);
}
#[test]
fn codex_request_signals_apply_session_precedence() {
let cases = vec![
(
request_headers(&[
(AETHER_SESSION_ID_HEADER, "aether-session"),
("session-id", "header-session"),
]),
json!({"client_metadata": {"session_id": "body-session"}}),
"aether-session",
ClientSessionSignalSource::ExplicitAetherHeader,
),
(
request_headers(&[
("session-id", "header-session"),
("thread-id", "header-thread"),
("session_id", "header-session-underscore"),
("conversation_id", "header-conversation"),
]),
json!({"client_metadata": {"session_id": "body-session"}}),
"header-session",
ClientSessionSignalSource::Header,
),
(
request_headers(&[
("thread-id", "header-thread"),
("session_id", "header-session-underscore"),
("conversation_id", "header-conversation"),
]),
json!({"client_metadata": {"session_id": "body-session"}}),
"header-thread",
ClientSessionSignalSource::Header,
),
(
request_headers(&[
("session_id", "header-session-underscore"),
("conversation_id", "header-conversation"),
]),
json!({"client_metadata": {"session_id": "body-session"}}),
"header-session-underscore",
ClientSessionSignalSource::Header,
),
(
request_headers(&[("conversation_id", "header-conversation")]),
json!({"client_metadata": {"session_id": "body-session"}}),
"header-conversation",
ClientSessionSignalSource::Header,
),
(
HeaderMap::new(),
json!({
"prompt_cache_key": "prompt-cache",
"client_metadata": {
"session_id": "body-session",
"thread_id": "body-thread",
"x-codex-turn-metadata": {
"session_id": "nested-session",
"thread_id": "nested-thread"
}
}
}),
"body-session",
ClientSessionSignalSource::Body,
),
(
HeaderMap::new(),
json!({
"prompt_cache_key": "prompt-cache",
"client_metadata": {
"thread_id": "body-thread",
"x-codex-turn-metadata": {"session_id": "nested-session"}
}
}),
"body-thread",
ClientSessionSignalSource::Body,
),
(
HeaderMap::new(),
json!({
"prompt_cache_key": "prompt-cache",
"client_metadata": {
"x-codex-turn-metadata": {
"session_id": "nested-session",
"thread_id": "nested-thread"
}
}
}),
"nested-session",
ClientSessionSignalSource::Body,
),
(
HeaderMap::new(),
json!({
"prompt_cache_key": "prompt-cache",
"client_metadata": {
"x-codex-turn-metadata": json!({
"thread_id": "nested-thread"
}).to_string()
}
}),
"nested-thread",
ClientSessionSignalSource::Body,
),
(
HeaderMap::new(),
json!({
"prompt_cache_key": "prompt-cache",
"conversation_id": "generic-conversation"
}),
"prompt-cache",
ClientSessionSignalSource::Body,
),
(
HeaderMap::new(),
json!({"metadata": {"session_id": "generic-session"}}),
"generic-session",
ClientSessionSignalSource::Body,
),
];
for (headers, body, expected_session_id, expected_source) in cases {
let signals = codex_request_signals_from_request(&headers, Some(&body));
assert_eq!(signals.session_id.as_deref(), Some(expected_session_id));
let mut codex_headers = headers;
codex_headers.insert(
http::header::USER_AGENT,
HeaderValue::from_static("codex_cli_rs/0.144.1"),
);
let scope = client_session_scope_from_request(&codex_headers, Some(&body))
.expect("Codex scope should reuse the native signal precedence");
assert_eq!(scope.client_family, "codex");
assert_eq!(scope.session_id, expected_session_id);
assert_eq!(scope.source, expected_source);
}
}
#[test]
fn codex_request_signals_extract_thread_and_prompt_cache_independently() {
let body = json!({
"prompt_cache_key": "prompt-cache",
"client_metadata": {
"thread_id": "body-thread",
"x-codex-turn-metadata": {"thread_id": "nested-thread"}
}
});
let headers = request_headers(&[("thread-id", "header-thread")]);
let header_signals = codex_request_signals_from_request(&headers, Some(&body));
assert_eq!(header_signals.thread_id.as_deref(), Some("header-thread"));
assert_eq!(
header_signals.prompt_cache_key.as_deref(),
Some("prompt-cache")
);
let body_signals = codex_request_signals_from_request(&HeaderMap::new(), Some(&body));
assert_eq!(body_signals.thread_id.as_deref(), Some("body-thread"));
let nested_body = json!({
"client_metadata": {
"x-codex-turn-metadata": json!({
"thread_id": "nested-thread"
}).to_string()
}
});
let nested_signals =
codex_request_signals_from_request(&HeaderMap::new(), Some(&nested_body));
assert_eq!(nested_signals.thread_id.as_deref(), Some("nested-thread"));
let session_only_body = json!({"client_metadata": {"session_id": "body-session"}});
let session_only_signals =
codex_request_signals_from_request(&HeaderMap::new(), Some(&session_only_body));
assert_eq!(
session_only_signals.thread_id.as_deref(),
Some("body-session")
);
}
#[test]
fn codex_request_signals_use_live_session_header() {
let headers = request_headers(&[("x-session-id", "live-session")]);
let signals = codex_request_signals_from_request(&headers, None);
assert_eq!(signals.session_id.as_deref(), Some("live-session"));
assert_eq!(signals.thread_id.as_deref(), Some("live-session"));
}
#[test]
fn codex_request_signals_prefer_responses_session_header_over_live_session_header() {
let headers = request_headers(&[
("session-id", "responses-session"),
("x-session-id", "live-session"),
]);
let signals = codex_request_signals_from_request(&headers, None);
assert_eq!(signals.session_id.as_deref(), Some("responses-session"));
assert_eq!(signals.thread_id.as_deref(), Some("responses-session"));
}
#[test]
fn codex_request_signals_apply_turn_precedence() {
let headers = request_headers(&[("x-codex-turn-metadata", r#"{"turn_id":"header-turn"}"#)]);
let direct_body = json!({
"turn_id": "top-level-turn",
"client_metadata": {
"turn_id": "body-turn",
"x-codex-turn-metadata": {"turn_id": "nested-turn"}
}
});
assert_eq!(
codex_request_signals_from_request(&headers, Some(&direct_body))
.turn_id
.as_deref(),
Some("body-turn")
);
let nested_object_body = json!({
"turn_id": "top-level-turn",
"client_metadata": {
"x-codex-turn-metadata": {"turn_id": "nested-object-turn"}
}
});
assert_eq!(
codex_request_signals_from_request(&headers, Some(&nested_object_body))
.turn_id
.as_deref(),
Some("nested-object-turn")
);
let nested_string_body = json!({
"turn_id": "top-level-turn",
"client_metadata": {
"x-codex-turn-metadata": json!({
"turn_id": "nested-string-turn"
}).to_string()
}
});
assert_eq!(
codex_request_signals_from_request(&headers, Some(&nested_string_body))
.turn_id
.as_deref(),
Some("nested-string-turn")
);
let top_level_body = json!({
"turn_id": "top-level-turn",
"client_metadata": {"x-codex-turn-metadata": "not-json"}
});
assert_eq!(
codex_request_signals_from_request(&headers, Some(&top_level_body))
.turn_id
.as_deref(),
Some("top-level-turn")
);
assert_eq!(
codex_request_signals_from_request(&headers, None)
.turn_id
.as_deref(),
Some("header-turn")
);
}
#[test]
fn codex_request_signals_ignore_client_request_id() {
let headers = request_headers(&[("x-client-request-id", "request-only-id")]);
let signals =
codex_request_signals_from_request(&headers, Some(&json!({"model": "gpt-5"})));
assert_eq!(signals, super::CodexRequestSignals::default());
}
#[test]
fn report_context_round_trips_normalized_session_affinity() {
let affinity = ClientSessionAffinity::new(
apps/aether-gateway/src/constants.rs
+4
View File
@@ -125,6 +125,10 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens",
"/v1/responses",
"/v1/responses/compact",
"/v1/realtime",
"/v1/realtime/calls",
"/v1/live",
"/v1/live/{call_id}",
"/v1/alpha/search",
"/v1/models/{model}:generateContent",
"/v1/models/{model}:streamGenerateContent",
apps/aether-gateway/src/control/auth/credentials.rs
+33 -1
View File
@@ -177,6 +177,19 @@ fn extract_trusted_auth_headers(headers: &http::HeaderMap) -> Option<GatewayTrus
})
}
#[cfg(not(test))]
pub(super) fn extract_trusted_admin_headers(
_headers: &http::HeaderMap,
) -> Option<GatewayTrustedAdminHeaders> {
// The public gateway has no authenticated upstream that is allowed to
// assert an administrator principal. `x-aether-gateway` is also emitted
// on public responses, so it cannot serve as proof that these headers were
// produced by a trusted hop. Production requests must authenticate with a
// real admin session or management bearer token instead.
None
}
#[cfg(test)]
pub(super) fn extract_trusted_admin_headers(
headers: &http::HeaderMap,
) -> Option<GatewayTrustedAdminHeaders> {
@@ -230,7 +243,7 @@ fn select_primary_credential(
if signature.starts_with("claude:") {
return select_claude_messages_credential(bundle);
}
if signature.starts_with("openai:") {
if signature.starts_with("openai:") || signature.starts_with("codex:") {
return select_openai_credential(bundle);
}
if signature.starts_with("aether:") {
@@ -478,6 +491,25 @@ mod tests {
);
}
#[test]
fn selects_codex_live_bearer_as_provider_api_key() {
let mut headers = http::HeaderMap::new();
headers.insert(
http::header::AUTHORIZATION,
"Bearer sk-codex-live".parse().unwrap(),
);
let extracted =
extract_request_credentials(&headers, &uri("/v1/live?model=gpt-live"), "codex:live");
assert_eq!(
extracted.primary,
Some(GatewayPrimaryCredential::ProviderApiKey {
raw: "sk-codex-live".to_string(),
carrier: GatewayCredentialCarrier::AuthorizationBearer,
})
);
}
#[test]
fn prefers_claude_chat_x_api_key_over_bearer() {
let mut headers = http::HeaderMap::new();
apps/aether-gateway/src/control/auth/gate.rs
+34 -5
View File
@@ -205,8 +205,8 @@ async fn available_balance_capacity_usd(
.as_ref()
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
Ok(match quota.as_ref() {
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
Some(_) if wallet_is_unlimited => None,
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)),
None if wallet_is_unlimited => None,
None => wallet_available_usd,
@@ -832,9 +832,10 @@ mod tests {
use serde_json::json;
use super::{
execution_plan_balance_capacity_rejection, execution_plan_cost_upper_bound_cache_key,
max_output_tokens_from_request, openai_request_input_is_self_contained,
output_choice_count_upper_bound, request_model_local_rejection, GatewayLocalAuthRejection,
available_balance_capacity_usd, execution_plan_balance_capacity_rejection,
execution_plan_cost_upper_bound_cache_key, max_output_tokens_from_request,
openai_request_input_is_self_contained, output_choice_count_upper_bound,
request_model_local_rejection, GatewayLocalAuthRejection,
};
use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
use crate::data::GatewayDataState;
@@ -939,6 +940,14 @@ mod tests {
fn state_with_quota_and_wallet(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
) -> AppState {
state_with_quota_context_and_wallet(quota, context, sample_wallet("user-1", 30.0))
}
fn state_with_quota_context_and_wallet(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
wallet: StoredWalletSnapshot,
) -> AppState {
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
@@ -952,7 +961,7 @@ mod tests {
AppState::new()
.expect("state should build")
.with_data_state_for_tests(data)
.with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)])
.with_auth_wallets_for_tests(vec![wallet])
}
fn state_with_model_mapping() -> AppState {
@@ -1350,6 +1359,26 @@ mod tests {
}
}
#[tokio::test]
async fn unlimited_wallet_capacity_ignores_exhausted_non_overage_quota() {
let context = billing_context_with_pricing(None, None, None, None);
let mut wallet = sample_wallet("user-1", 0.0);
wallet.limit_mode = "unlimited".to_string();
let state =
state_with_quota_context_and_wallet(quota_availability(0.0, false), context, wallet);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let auth_context = decision
.auth_context
.as_ref()
.expect("decision should include auth context");
let capacity = available_balance_capacity_usd(&state, auth_context)
.await
.expect("capacity should resolve");
assert_eq!(capacity, None);
}
#[tokio::test]
async fn positive_balance_does_not_allow_historical_invalid_processing_pricing() {
let context = billing_context_with_pricing(
apps/aether-gateway/src/control/auth/mod.rs
+3 -2
View File
@@ -11,8 +11,9 @@ pub(crate) use gate::{
should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection,
};
pub(crate) use resolution::{
refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
GatewayAdminPrincipalContext, GatewayControlAuthContext,
refresh_execution_runtime_auth_context, refresh_execution_runtime_auth_context_with_snapshot,
resolve_execution_runtime_auth_context, GatewayAdminPrincipalContext,
GatewayControlAuthContext,
};
pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution};
pub(crate) use types::GatewayCredentialCarrier;
apps/aether-gateway/src/control/auth/resolution.rs
+35 -7
View File
@@ -725,20 +725,47 @@ pub(crate) async fn refresh_execution_runtime_auth_context(
auth_context: GatewayControlAuthContext,
auth_endpoint_signature: Option<&str>,
) -> Result<GatewayControlAuthContext, GatewayError> {
refresh_execution_runtime_auth_context_with_snapshot(
state,
auth_context,
auth_endpoint_signature,
)
.await
.map(|(auth_context, _)| auth_context)
}
/// Strongly refreshes the long-lived execution authorization context and
/// returns the exact API-key snapshot that produced it.
///
/// WebSocket turns need both values: using the refreshed context for RPM and
/// balance checks while letting the planner independently read its normal
/// cache can authorize a different provider/model snapshot for up to the cache
/// TTL. Ordinary HTTP callers keep using [`refresh_execution_runtime_auth_context`].
pub(crate) async fn refresh_execution_runtime_auth_context_with_snapshot(
state: &AppState,
auth_context: GatewayControlAuthContext,
auth_endpoint_signature: Option<&str>,
) -> Result<
(
GatewayControlAuthContext,
Option<crate::ai_serving::GatewayAuthApiKeySnapshot>,
),
GatewayError,
> {
if auth_context.local_rejection.is_some() || !auth_context.access_allowed {
return Ok(auth_context);
return Ok((auth_context, None));
}
let Some(auth_endpoint_signature) = auth_endpoint_signature
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(auth_context);
return Ok((auth_context, None));
};
if !state.has_auth_api_key_reader()
|| auth_context.user_id.trim().is_empty()
|| auth_context.api_key_id.trim().is_empty()
{
return Ok(auth_context);
return Ok((auth_context, None));
}
let snapshot = {
@@ -758,19 +785,20 @@ pub(crate) async fn refresh_execution_runtime_auth_context(
denied.access_allowed = false;
denied.local_rejection = Some(GatewayLocalAuthRejection::InvalidApiKey);
denied.balance_remaining = None;
return Ok(denied);
return Ok((denied, None));
};
let wallet_access = resolve_wallet_auth_gate_uncached(state, &snapshot).await?;
Ok(build_data_backed_auth_context(
let refreshed = build_data_backed_auth_context(
state,
snapshot,
snapshot.clone(),
auth_endpoint_signature,
Some(true),
auth_context.balance_remaining,
wallet_access,
)
.await)
.await;
Ok((refreshed, Some(snapshot)))
}
fn put_cached_auth_context(
apps/aether-gateway/src/control/mod.rs
+5 -4
View File
@@ -9,10 +9,11 @@ mod route;
pub(crate) use auth::{
execution_plan_balance_capacity_rejection, extract_requested_model,
refresh_execution_runtime_auth_context, request_model_local_rejection,
resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext,
GatewayCredentialCarrier, GatewayLocalAuthRejection,
refresh_execution_runtime_auth_context, refresh_execution_runtime_auth_context_with_snapshot,
request_model_local_rejection, resolve_execution_runtime_auth_context,
should_buffer_request_for_local_auth, trusted_auth_local_rejection,
GatewayAdminPrincipalContext, GatewayControlAuthContext, GatewayCredentialCarrier,
GatewayLocalAuthRejection,
};
pub(crate) use execute::{allows_control_execute_emergency, maybe_execute_via_control};
pub(crate) use management_token_permissions::{
apps/aether-gateway/src/control/route/admin/operations_families.rs
+12
View File
@@ -797,6 +797,18 @@ pub(super) fn classify_admin_operations_family_route(
"admin:users",
false,
))
} else if method == http::Method::DELETE
&& normalized_path_no_trailing.starts_with("/api/admin/users/")
&& normalized_path_no_trailing.contains("/billing/entitlements/")
&& normalized_path_no_trailing.matches('/').count() == 7
{
Some(classified(
"admin_proxy",
"users_manage",
"revoke_user_billing_entitlement",
"admin:users",
false,
))
} else if method == http::Method::GET
&& normalized_path.starts_with("/api/admin/users/")
&& normalized_path.ends_with("/sessions")
apps/aether-gateway/src/control/route/ai.rs
+327 -1
View File
@@ -7,6 +7,7 @@ use crate::ai_serving::ApiOperation;
pub(super) fn classify_ai_public_route(
method: &http::Method,
normalized_path: &str,
query: Option<&str>,
headers: &http::HeaderMap,
) -> Option<ClassifiedRoute> {
if let Some(route) = classify_antigravity_v1internal_route(method, normalized_path) {
@@ -35,7 +36,39 @@ pub(super) fn classify_ai_public_route(
"openai:rerank",
true,
))
} else if method == http::Method::POST
} else if (method == http::Method::POST && normalized_path == "/v1/live")
|| (method == http::Method::POST
&& normalized_path == "/v1/realtime/calls"
&& realtime_query_has_codex_live_intent(query))
|| (method == http::Method::GET
&& ((normalized_path == "/v1/live" || normalized_path.starts_with("/v1/live/"))
|| (normalized_path == "/v1/realtime"
&& (realtime_query_has_codex_live_intent(query)
|| realtime_query_is_codex_v2(query, headers))))
&& is_websocket_upgrade_request(headers))
{
// Codex Live has an independent wire contract and permission surface;
// it must never be authorized as an OpenAI Responses request.
Some(classified("ai_public", "codex", "live", "codex:live", true))
} else if method == http::Method::GET
&& normalized_path == "/v1/realtime"
&& is_websocket_upgrade_request(headers)
{
// Ordinary OpenAI Realtime WebSockets retain their independent
// permission surface. The GA WebRTC call-creation endpoint is not
// handled by this WebSocket-only implementation; only Codex AVAS is
// accepted above when it carries the explicit quicksilver intent.
Some(classified(
"ai_public",
"openai",
"realtime",
"openai:realtime",
true,
))
} else if (method == http::Method::POST
|| (method == http::Method::GET
&& normalized_path == "/v1/responses"
&& is_websocket_upgrade_request(headers)))
&& matches!(normalized_path, "/v1/responses" | "/v1/responses/compact")
{
if normalized_path.ends_with("/compact") {
@@ -185,6 +218,56 @@ pub(super) fn classify_ai_public_route(
}
}
/// Codex's standalone realtime WebSocket is distinguished from the ordinary
/// OpenAI Realtime API by its `intent=quicksilver` selector. Keep this check
/// deliberately narrow: `call_id` is also used by ordinary OpenAI Realtime
/// sideband sockets, and neither it nor a model query may select the
/// `codex:live` permission surface on its own.
fn realtime_query_has_codex_live_intent(query: Option<&str>) -> bool {
let mut seen = false;
for (key, value) in url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()) {
if key.eq_ignore_ascii_case("intent") {
if seen || !value.eq_ignore_ascii_case("quicksilver") {
return false;
}
seen = true;
}
}
seen
}
/// Codex realtime v2 intentionally omits the v1 `intent=quicksilver` query
/// marker. Its default headers still carry the stable Codex originator, so
/// use that identity plus a model query to select the Codex Live permission
/// surface without stealing ordinary OpenAI Realtime sockets.
fn realtime_query_is_codex_v2(query: Option<&str>, headers: &http::HeaderMap) -> bool {
let mut has_model = false;
for (key, value) in url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()) {
// V2 has no intent marker. If one is present, let the normal
// Realtime/V1 classifier handle it instead of silently treating a
// conflicting request as Codex V2.
if key.eq_ignore_ascii_case("intent") || key.eq_ignore_ascii_case("call_id") {
return false;
}
if key.eq_ignore_ascii_case("model") && !value.trim().is_empty() {
has_model = true;
}
}
if !has_model {
return false;
}
let Some(originator) = crate::headers::header_value_str(headers, "originator") else {
return false;
};
originator.split_whitespace().next().is_some_and(|value| {
value.eq_ignore_ascii_case("codex_cli_rs")
|| value.to_ascii_lowercase().starts_with("codex_cli_rs/")
|| value.eq_ignore_ascii_case("codex_work_desktop")
|| value.eq_ignore_ascii_case("codex_work_web")
|| value.eq_ignore_ascii_case("codex_work_mobile")
})
}
fn claude_request_auth_channel(headers: &http::HeaderMap) -> &'static str {
if crate::headers::header_value_str(headers, "x-api-key").is_some()
|| crate::headers::header_value_str(headers, "api-key").is_some()
@@ -199,6 +282,24 @@ fn claude_request_auth_channel(headers: &http::HeaderMap) -> &'static str {
}
}
fn is_websocket_upgrade_request(headers: &http::HeaderMap) -> bool {
let has_upgrade_connection = headers
.get(http::header::CONNECTION)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| {
value
.split(',')
.map(str::trim)
.any(|value| value.eq_ignore_ascii_case("upgrade"))
});
let has_websocket_upgrade = headers
.get(http::header::UPGRADE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.eq_ignore_ascii_case("websocket"));
has_upgrade_connection && has_websocket_upgrade
}
fn is_gemini_operation_method(method: &http::Method, normalized_path: &str) -> bool {
method == http::Method::GET
|| (method == http::Method::POST && normalized_path.ends_with(":cancel"))
@@ -242,3 +343,228 @@ fn classify_antigravity_v1internal_route(
execution_runtime_candidate,
))
}
#[cfg(test)]
mod tests {
use axum::http::header::{CONNECTION, UPGRADE};
use axum::http::{HeaderMap, HeaderValue, Method};
use super::classify_ai_public_route;
#[test]
fn classifies_websocket_upgrade_on_responses_route() {
let mut headers = HeaderMap::new();
headers.insert(CONNECTION, HeaderValue::from_static("keep-alive, Upgrade"));
headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
let route = classify_ai_public_route(&Method::GET, "/v1/responses", None, &headers)
.expect("Responses WebSocket should be an AI public route");
assert_eq!(route.route_class, "ai_public");
assert_eq!(route.route_family, "openai");
assert_eq!(route.route_kind, "responses");
assert_eq!(route.auth_endpoint_signature, "openai:responses");
}
#[test]
fn does_not_classify_plain_get_as_responses_websocket() {
assert!(
classify_ai_public_route(&Method::GET, "/v1/responses", None, &HeaderMap::new())
.is_none()
);
}
#[test]
fn classifies_only_websocket_upgrade_on_realtime_route() {
let mut headers = HeaderMap::new();
headers.insert(CONNECTION, HeaderValue::from_static("keep-alive, Upgrade"));
headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
let route = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("model=gpt-realtime-2.1"),
&headers,
)
.expect("Realtime WebSocket should be an AI public route");
assert_eq!(route.route_class, "ai_public");
assert_eq!(route.route_family, "openai");
assert_eq!(route.route_kind, "realtime");
assert_eq!(route.auth_endpoint_signature, "openai:realtime");
assert!(route.execution_runtime_candidate);
let mut codex_v2_headers = headers.clone();
codex_v2_headers.insert("originator", HeaderValue::from_static("codex_work_desktop"));
let codex_v2 = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("model=gpt-live-1-codex"),
&codex_v2_headers,
)
.expect("Codex realtime v2 should use the Live route");
assert_eq!(codex_v2.route_family, "codex");
assert_eq!(codex_v2.auth_endpoint_signature, "codex:live");
let mut codex_cli_headers = headers.clone();
codex_cli_headers.insert("originator", HeaderValue::from_static("codex_cli_rs"));
let codex_cli_v2 = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("model=gpt-live-1-codex"),
&codex_cli_headers,
)
.expect("Codex CLI realtime v2 should use the Live route");
assert_eq!(codex_cli_v2.auth_endpoint_signature, "codex:live");
let mut ordinary_headers = headers.clone();
ordinary_headers.insert("originator", HeaderValue::from_static("openai-python"));
let ordinary = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("model=gpt-realtime-2.1"),
&ordinary_headers,
)
.expect("ordinary Realtime should remain available");
assert_eq!(ordinary.auth_endpoint_signature, "openai:realtime");
assert!(classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("model=gpt-realtime-2.1"),
&HeaderMap::new(),
)
.is_none());
assert!(classify_ai_public_route(&Method::POST, "/v1/realtime", None, &headers).is_none());
}
#[test]
fn classifies_live_http_and_websocket_routes_as_codex_live() {
let legacy_post =
classify_ai_public_route(&Method::POST, "/v1/live", None, &HeaderMap::new())
.expect("legacy Live call creation should be an AI public route");
assert_eq!(legacy_post.route_family, "codex");
assert_eq!(legacy_post.route_kind, "live");
assert_eq!(legacy_post.auth_endpoint_signature, "codex:live");
let avas_post = classify_ai_public_route(
&Method::POST,
"/v1/realtime/calls",
Some("intent=quicksilver&architecture=avas"),
&HeaderMap::new(),
)
.expect("Codex AVAS call creation should be an AI public route");
assert_eq!(avas_post.route_family, "codex");
assert_eq!(avas_post.route_kind, "live");
assert_eq!(avas_post.auth_endpoint_signature, "codex:live");
// The same endpoint is part of the ordinary OpenAI Realtime API, but
// Aether's OpenAI Realtime implementation currently supports direct
// WebSockets only. A request without Codex's explicit AVAS intent must
// neither be captured by the Codex Live planner nor be advertised as a
// supported ordinary call-create request.
for query in [None, Some("model=gpt-realtime"), Some("architecture=avas")] {
assert!(classify_ai_public_route(
&Method::POST,
"/v1/realtime/calls",
query,
&HeaderMap::new(),
)
.is_none());
}
for query in [
Some("intent=other&architecture=avas"),
Some("intent=quicksilver&intent=other&architecture=avas"),
] {
assert!(classify_ai_public_route(
&Method::POST,
"/v1/realtime/calls",
query,
&HeaderMap::new(),
)
.is_none());
}
let mut headers = HeaderMap::new();
headers.insert(CONNECTION, HeaderValue::from_static("Upgrade"));
headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
for path in ["/v1/live", "/v1/live/rtc_opaque"] {
let route = classify_ai_public_route(&Method::GET, path, None, &headers)
.expect("Live WebSocket should be an AI public route");
assert_eq!(route.route_family, "codex");
assert_eq!(route.route_kind, "live");
assert_eq!(route.auth_endpoint_signature, "codex:live");
}
assert!(classify_ai_public_route(
&Method::GET,
"/v1/live/rtc_opaque",
None,
&HeaderMap::new(),
)
.is_none());
let sideband = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("intent=quicksilver&call_id=rtc_opaque"),
&headers,
)
.expect("Realtime sideband WebSocket should be a Codex Live route");
assert_eq!(sideband.route_family, "codex");
assert_eq!(sideband.route_kind, "live");
assert_eq!(sideband.auth_endpoint_signature, "codex:live");
for query in [None, Some("model=gpt-realtime")] {
let route = classify_ai_public_route(&Method::GET, "/v1/realtime", query, &headers)
.expect("Realtime WebSocket without a call_id key should remain OpenAI Realtime");
assert_eq!(route.auth_endpoint_signature, "openai:realtime");
}
let codex_direct = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("intent=quicksilver&model=gpt-realtime-1.5"),
&headers,
)
.expect("Codex direct realtime WebSocket should use the Live route");
assert_eq!(codex_direct.route_family, "codex");
assert_eq!(codex_direct.route_kind, "live");
assert_eq!(codex_direct.auth_endpoint_signature, "codex:live");
for query in [
Some("intent=other&model=gpt-realtime-1.5"),
Some("intent=quicksilver&intent=other&model=gpt-realtime-1.5"),
] {
let route = classify_ai_public_route(&Method::GET, "/v1/realtime", query, &headers)
.expect("non-Codex realtime intent should remain OpenAI Realtime");
assert_eq!(route.auth_endpoint_signature, "openai:realtime");
}
// `call_id` is shared by ordinary OpenAI Realtime WebRTC sideband
// sockets. It must not select Codex Live without Codex's explicit
// `intent=quicksilver` signal.
for query in [
Some("call_id=rtc_ordinary"),
Some("c%61ll_id=rtc_encoded"),
Some("call_id="),
Some("call_id=%20"),
Some("call_id=rtc_one&call_id=rtc_two"),
Some("call_id=rtc_ordinary&model=gpt-realtime-1.5"),
] {
let route = classify_ai_public_route(&Method::GET, "/v1/realtime", query, &headers)
.expect("ordinary Realtime sideband should remain routable");
assert_eq!(route.route_family, "openai");
assert_eq!(route.route_kind, "realtime");
assert_eq!(route.auth_endpoint_signature, "openai:realtime");
}
let codex_sideband_with_encoded_call_id = classify_ai_public_route(
&Method::GET,
"/v1/realtime",
Some("intent=quicksilver&c%61ll_id=rtc_encoded"),
&headers,
)
.expect("encoded call_id must not hide Codex's explicit Live intent");
assert_eq!(
codex_sideband_with_encoded_call_id.auth_endpoint_signature,
"codex:live"
);
}
}
apps/aether-gateway/src/control/route/mod.rs
+2 -3
View File
@@ -212,7 +212,7 @@ pub(crate) fn classify_control_route(
.or_else(|| oauth::classify_oauth_route(method, &normalized_path))
.or_else(|| admin::classify_admin_route(method, &normalized_path))
.or_else(|| internal::classify_internal_route(method, &normalized_path))
.or_else(|| ai::classify_ai_public_route(method, &normalized_path, headers))?;
.or_else(|| ai::classify_ai_public_route(method, &normalized_path, uri.query(), headers))?;
let mut decision = classified.into_decision(normalized_path);
if let Some(signature) = decision.auth_endpoint_signature.as_deref() {
@@ -247,8 +247,7 @@ pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::Hea
let has_codex_client_version = uri.path() == "/v1/models"
&& uri.query().is_some_and(|query| {
url::form_urlencoded::parse(query.as_bytes())
.any(|(key, value)| key == "client_version" && !value.trim().is_empty())
url::form_urlencoded::parse(query.as_bytes()).any(|(key, _)| key == "client_version")
});
if has_codex_client_version {
return "openai:responses".to_string();
apps/aether-gateway/src/control/tests/admin_users.rs
+16
View File
@@ -103,6 +103,22 @@ fn classifies_admin_user_billing_routes_as_admin_proxy_route() {
Some("admin:users")
);
let revoke_uri: Uri = "/api/admin/users/user-1/billing/entitlements/entitlement-1"
.parse()
.expect("uri should parse");
let revoke = classify_control_route(&http::Method::DELETE, &revoke_uri, &headers)
.expect("route should classify");
assert_eq!(revoke.route_class.as_deref(), Some("admin_proxy"));
assert_eq!(revoke.route_family.as_deref(), Some("users_manage"));
assert_eq!(
revoke.route_kind.as_deref(),
Some("revoke_user_billing_entitlement")
);
assert_eq!(
revoke.auth_endpoint_signature.as_deref(),
Some("admin:users")
);
let context = GatewayPublicRequestContext::from_request_parts(
"trace-user-billing-grant",
&http::Method::POST,
apps/aether-gateway/src/control/tests/ai.rs
+26
View File
@@ -108,6 +108,32 @@ fn classifies_openai_chat_and_responses_separately_from_embedding() {
assert_ne!(responses.route_kind.as_deref(), Some("embedding"));
}
#[test]
fn classifies_openai_realtime_only_for_websocket_upgrades() {
let websocket_headers = headers(&[
("authorization", "Bearer sk-test"),
("connection", "keep-alive, Upgrade"),
("upgrade", "websocket"),
]);
let uri: Uri = "/v1/realtime?model=gpt-realtime"
.parse()
.expect("uri should parse");
let decision = classify_control_route(&http::Method::GET, &uri, &websocket_headers)
.expect("Realtime WebSocket route should classify");
assert_eq!(decision.route_family.as_deref(), Some("openai"));
assert_eq!(decision.route_kind.as_deref(), Some("realtime"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:realtime")
);
assert!(decision.is_execution_runtime_candidate());
let plain_headers = headers(&[("authorization", "Bearer sk-test")]);
assert!(classify_control_route(&http::Method::GET, &uri, &plain_headers).is_none());
assert!(classify_control_route(&http::Method::POST, &uri, &websocket_headers).is_none());
}
#[test]
fn classifies_openai_image_generation_and_edit_but_not_variation() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
apps/aether-gateway/src/control/tests/public_support.rs
+2 -2
View File
@@ -39,7 +39,7 @@ fn classifies_codex_models_list_with_responses_auth_signature() {
}
#[test]
fn empty_codex_client_version_keeps_standard_openai_models_signature() {
fn empty_codex_client_version_uses_responses_signature_for_bounded_fallback() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
let uri: Uri = "/v1/models?client_version="
.parse()
@@ -49,7 +49,7 @@ fn empty_codex_client_version_keeps_standard_openai_models_signature() {
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:chat")
Some("openai:responses")
);
}
apps/aether-gateway/src/data/config.rs
+29
View File
@@ -195,6 +195,35 @@ mod tests {
assert_eq!(background.pool.min_connections, 1);
}
#[test]
fn runtime_pool_split_gives_default_small_server_more_foreground_capacity() {
let config = GatewayDataConfig::from_database_config(
SqlDatabaseConfig::new(
DatabaseDriver::Postgres,
"postgres://localhost/aether",
SqlPoolConfig {
min_connections: 4,
max_connections: 32,
..SqlPoolConfig::default()
},
)
.expect("database config should be valid"),
);
let (foreground, background) = config.split_runtime_pools_with_background_max(None);
let foreground = foreground.database().expect("foreground database");
let background = background
.expect("background database config")
.database()
.expect("background database")
.clone();
assert_eq!(foreground.pool.max_connections, 26);
assert_eq!(background.pool.max_connections, 6);
assert_eq!(foreground.pool.min_connections, 4);
assert_eq!(background.pool.min_connections, 1);
}
#[test]
fn runtime_pool_split_can_be_disabled_or_degrade_for_single_connection() {
let mut database = SqlDatabaseConfig::sqlite_default();
apps/aether-gateway/src/data/state/catalog.rs
+32 -8
View File
@@ -1,14 +1,14 @@
use super::{
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
PublicHealthStatusCount, PublicHealthTimelineBucket, StoredGeminiFileMapping,
StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, PublicHealthStatusCount, PublicHealthTimelineBucket,
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredRequestCandidate, UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
};
impl GatewayDataState {
@@ -282,6 +282,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_provider_catalog_keys_by_ids_strong(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_by_ids_strong(key_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_provider_catalog_keys_by_provider_ids(
&self,
provider_ids: &[String],
@@ -561,6 +571,20 @@ impl GatewayDataState {
Ok(updated)
}
pub(crate) async fn compare_and_update_provider_catalog_key_admin_state(
&self,
update: &ProviderCatalogKeyAdminCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_update_key_admin_state(update).await,
None => Ok(false),
}?;
// Clear on both success and conflict so a retry cannot reuse the stale
// credential snapshot that lost the CAS.
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn update_provider_catalog_keys(
&self,
keys: &[StoredProviderCatalogKey],
apps/aether-gateway/src/data/state/mod.rs
+7 -7
View File
@@ -121,13 +121,13 @@ use aether_data_contracts::repository::pool_scores::{
UpsertPoolMemberScore,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
apps/aether-gateway/src/data/state/provider_catalog_cache.rs
+57
View File
@@ -334,6 +334,13 @@ impl ProviderCatalogReadRepository for CachedProviderCatalogReadRepository {
}
}
async fn list_keys_by_ids_strong(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.inner.list_keys_by_ids_strong(key_ids).await
}
async fn list_keys_by_provider_ids(
&self,
provider_ids: &[String],
@@ -538,6 +545,7 @@ fn normalize_ids(ids: &[String]) -> Vec<String> {
mod tests {
use super::*;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogWriteRepository;
fn cache() -> CachedProviderCatalogReadRepository {
CachedProviderCatalogReadRepository::new(Arc::new(
@@ -555,6 +563,55 @@ mod tests {
.expect("provider should be valid")
}
#[tokio::test]
async fn provider_catalog_strong_key_read_bypasses_fresh_cached_generation() {
let old_metadata = serde_json::json!({
"codex": {"credential_generation": "old"}
});
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should be valid");
key.upstream_metadata = Some(old_metadata.clone());
let inner = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider("provider-1")],
Vec::new(),
vec![key],
));
let cache = CachedProviderCatalogReadRepository::new(inner.clone());
let key_ids = vec!["key-1".to_string()];
let first = cache
.list_keys_by_ids(&key_ids)
.await
.expect("initial key read should succeed");
assert_eq!(first[0].upstream_metadata.as_ref(), Some(&old_metadata));
let new_metadata = serde_json::json!({
"codex": {"credential_generation": "new"}
});
assert!(inner
.upsert_key_upstream_metadata_namespace("key-1", "codex", &new_metadata["codex"], None,)
.await
.expect("inner metadata update should succeed"));
let cached = cache
.list_keys_by_ids(&key_ids)
.await
.expect("cached key read should succeed");
assert_eq!(cached[0].upstream_metadata.as_ref(), Some(&old_metadata));
let strong = cache
.list_keys_by_ids_strong(&key_ids)
.await
.expect("strong key read should succeed");
assert_eq!(strong[0].upstream_metadata.as_ref(), Some(&new_metadata));
}
#[tokio::test]
async fn provider_catalog_follower_observes_completion_before_first_poll() {
let cache = cache();
apps/aether-gateway/src/data/state/runtime.rs
+15
View File
@@ -2578,6 +2578,21 @@ impl GatewayDataState {
}
}
pub(crate) async fn revoke_user_plan_entitlement(
&self,
user_id: &str,
entitlement_id: &str,
) -> Result<AdminBillingMutationOutcome<()>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.revoke_user_plan_entitlement(user_id, entitlement_id)
.await
}
None => Ok(AdminBillingMutationOutcome::Unavailable),
}
}
pub(crate) async fn find_user_daily_quota_availability(
&self,
user_id: &str,
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+163 -33
View File
@@ -105,6 +105,7 @@ async fn schedule_pool_page_candidates(
candidates: Vec<EligibleLocalExecutionCandidate>,
sticky_session_token: Option<&str>,
effective_pool_config: Option<&AdminProviderPoolConfig>,
provider_model_name: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -128,7 +129,8 @@ async fn schedule_pool_page_candidates(
entry.1.insert(candidate.candidate.key_id.clone());
}
let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await;
let key_context_by_id =
read_pool_catalog_key_contexts_by_id(state, &candidates, provider_model_name).await;
let mut runtime_by_provider = BTreeMap::new();
let mut pool_config_by_provider = BTreeMap::new();
@@ -596,11 +598,17 @@ impl<'a> PoolKeyCursor<'a> {
return;
};
self.exhaustion_skip_recorded = true;
record_local_runtime_candidate_skip_reason(
self.state.app(),
trace_id,
self.runtime_miss_pool_exhaustion_skip_reason(),
);
if self.skip_reason_counts.is_empty() {
record_local_runtime_candidate_skip_reason(
self.state.app(),
trace_id,
"pool_group_exhausted",
);
return;
}
for reason in self.skip_reason_counts.keys() {
record_local_runtime_candidate_skip_reason(self.state.app(), trace_id, reason);
}
}
fn runtime_miss_pool_exhaustion_skip_reason(&self) -> &'static str {
@@ -746,19 +754,28 @@ impl<'a> PoolKeyCursor<'a> {
return None;
}
};
let api_format = self.group.candidate.endpoint_api_format.as_str();
rows.sort_by(|left, right| {
let left_priority = self
.routing_overlay
.as_ref()
.map_or(left.key_internal_priority, |overlay| {
overlay.key_priority(&left.key_id, left.key_internal_priority)
});
let right_priority = self
.routing_overlay
.as_ref()
.map_or(right.key_internal_priority, |overlay| {
overlay.key_priority(&right.key_id, right.key_internal_priority)
});
let left_priority = self.routing_overlay.as_ref().map_or(
left.key_internal_priority,
|overlay| {
overlay.key_priority_for_format(
&left.key_id,
api_format,
left.key_internal_priority,
)
},
);
let right_priority = self.routing_overlay.as_ref().map_or(
right.key_internal_priority,
|overlay| {
overlay.key_priority_for_format(
&right.key_id,
api_format,
right.key_internal_priority,
)
},
);
left_priority
.cmp(&right_priority)
.then(left.key_id.cmp(&right.key_id))
@@ -1034,6 +1051,7 @@ impl<'a> PoolKeyCursor<'a> {
candidates,
self.sticky_session_token.as_deref(),
self.effective_pool_config.as_ref(),
Some(self.group.candidate.selected_provider_model_name.as_str()),
)
.await;
self.record_skipped_candidates(&skipped);
@@ -1406,6 +1424,7 @@ fn pool_candidate_from_catalog_key(
async fn read_pool_catalog_key_contexts_by_id(
state: PlannerAppState<'_>,
candidates: &[EligibleLocalExecutionCandidate],
provider_model_name: Option<&str>,
) -> BTreeMap<String, PoolCatalogKeyContext> {
let mut key_ids = Vec::new();
let mut provider_type_by_key_id = BTreeMap::<String, String>::new();
@@ -1451,7 +1470,13 @@ async fn read_pool_catalog_key_contexts_by_id(
.unwrap_or_default();
(
key.id.clone(),
build_pool_catalog_key_context(state, &provider_pool_service, &key, provider_type),
build_pool_catalog_key_context(
state,
&provider_pool_service,
&key,
provider_type,
provider_model_name,
),
)
})
.collect()
@@ -1462,6 +1487,7 @@ fn build_pool_catalog_key_context(
provider_pool_service: &ProviderPoolService,
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: Option<&str>,
) -> PoolCatalogKeyContext {
let (health_score, _, _, _, _) = provider_key_health_summary(key);
let health_score = key
@@ -1480,8 +1506,12 @@ fn build_pool_catalog_key_context(
.filter(|value| value.is_finite() && *value >= 0.0);
let auth_config = parse_catalog_auth_config_json(state.app(), key);
let mut signals =
provider_pool_service.member_signals(provider_type, key, auth_config.as_ref());
let mut signals = provider_pool_service.member_signals(
provider_type,
key,
auth_config.as_ref(),
provider_model_name,
);
signals.account_blocked |= admin_provider_pool_pure::admin_pool_key_is_known_banned(key);
signals.account_blocked |=
pool_key_requires_reauth_for_scheduling(key, current_unix_ms().saturating_div(1000));
@@ -1826,17 +1856,17 @@ fn pool_key_candidate_order_for_group(
})
.collect::<Vec<_>>();
let active_presets = ProviderPoolService::with_builtin_adapters()
.normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets)
.into_iter()
.map(|preset| preset.preset)
.collect::<Vec<_>>();
.normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets);
if let Some(distribution_mode) = active_presets
.iter()
.find(|preset| pool_distribution_mode_preset(preset.as_str()))
.map(String::as_str)
.find(|preset| pool_distribution_mode_preset(preset.preset.as_str()))
{
return match distribution_mode {
"cache_affinity" => StoredPoolKeyCandidateOrder::CacheAffinity,
return match distribution_mode.preset.as_str() {
"cache_affinity" => match distribution_mode.mode.as_deref() {
Some("lru") => StoredPoolKeyCandidateOrder::Lru,
Some("single_account") => StoredPoolKeyCandidateOrder::SingleAccount,
_ => StoredPoolKeyCandidateOrder::CacheAffinity,
},
"load_balance" => StoredPoolKeyCandidateOrder::LoadBalance {
seed: pool_sort_seed(),
},
@@ -1921,11 +1951,13 @@ fn apply_pool_orchestration(
orchestration: PoolCandidateOrchestration,
) -> EligibleLocalExecutionCandidate {
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
let sticky_key_attempts = candidate.orchestration.sticky_key_attempts;
candidate.orchestration = LocalExecutionCandidateMetadata {
candidate_group_id: orchestration.candidate_group_id,
pool_key_index: orchestration.pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch,
sticky_key_attempts,
};
candidate
}
@@ -2129,6 +2161,7 @@ mod tests {
pool_key_index: Some(0),
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
}
);
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
@@ -2146,12 +2179,13 @@ mod tests {
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
}
);
}
#[test]
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
fn pool_scheduler_promotes_sticky_hit_before_lru_secondary_order() {
let key_a = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
@@ -2159,7 +2193,11 @@ mod tests {
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "cache_affinity", "enabled": true}]
"scheduling_presets": [{
"preset": "cache_affinity",
"enabled": true,
"mode": "lru"
}]
}
})),
);
@@ -2170,7 +2208,11 @@ mod tests {
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "cache_affinity", "enabled": true}]
"scheduling_presets": [{
"preset": "cache_affinity",
"enabled": true,
"mode": "lru"
}]
}
})),
);
@@ -2204,6 +2246,37 @@ mod tests {
);
}
#[test]
fn cache_affinity_secondary_modes_select_distinct_candidate_orders() {
for (mode, expected) in [
("single_account", StoredPoolKeyCandidateOrder::SingleAccount),
("lru", StoredPoolKeyCandidateOrder::Lru),
] {
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{
"preset": "cache_affinity",
"enabled": true,
"mode": mode
}]
}
})),
);
let config = pool_config_for_candidate(&group).expect("pool config should parse");
assert!(admin_provider_pool_cache_affinity_enabled(&config));
assert_eq!(
pool_key_candidate_order_for_group(&group, Some(&config)),
expected
);
}
}
#[test]
fn pool_scheduler_ignores_sticky_hit_without_cache_affinity() {
let key_a = sample_eligible_candidate(
@@ -3191,8 +3264,12 @@ mod tests {
.take_local_execution_runtime_miss_diagnostic(trace_id)
.expect("runtime miss diagnostic should exist");
assert_eq!(diagnostic.reason, "all_candidates_skipped");
assert_eq!(diagnostic.skipped_candidate_count, Some(1));
assert_eq!(diagnostic.skipped_candidate_count, Some(2));
assert_eq!(diagnostic.skip_reasons.get("pool_cooldown"), Some(&1));
assert_eq!(
diagnostic.skip_reasons.get("transport_snapshot_missing"),
Some(&1)
);
}
#[tokio::test]
@@ -4426,6 +4503,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert_eq!(context.plan_tier.as_deref(), Some("team"));
@@ -4471,6 +4549,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(!context.quota_exhausted);
@@ -4491,6 +4570,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(context.quota_exhausted);
@@ -4521,11 +4601,59 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"antigravity",
None,
);
assert!(context.quota_exhausted);
}
#[test]
fn pool_catalog_context_scopes_antigravity_exhaustion_to_requested_model() {
let mut key = sample_catalog_oauth_key("key-antigravity-model-quota");
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"exhausted": false,
"windows": [
{
"code": "model:gemini-3.1-pro-high",
"scope": "model",
"model": "gemini-3.1-pro-high",
"used_ratio": 1.0,
"is_exhausted": true
},
{
"code": "model:gemini-3-flash-agent",
"scope": "model",
"model": "gemini-3-flash-agent",
"used_ratio": 0.1,
"is_exhausted": false
}
]
}
}));
let app = app_state_with_catalog_key(key.clone());
let exhausted = build_pool_catalog_key_context(
PlannerAppState::new(&app),
&ProviderPoolService::with_builtin_adapters(),
&key,
"antigravity",
Some("gemini-3.1-pro-high"),
);
let available = build_pool_catalog_key_context(
PlannerAppState::new(&app),
&ProviderPoolService::with_builtin_adapters(),
&key,
"antigravity",
Some("gemini-3-flash-agent"),
);
assert!(exhausted.quota_exhausted);
assert!(!available.quota_exhausted);
}
#[test]
fn pool_catalog_context_marks_known_banned_account_from_metadata() {
let mut key = sample_catalog_oauth_key("key-account-banned");
@@ -4542,6 +4670,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(context.account_blocked);
@@ -4947,6 +5076,7 @@ mod tests {
priority_mode: RoutingSetPriorityMode::Provider,
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: RankingOverlay {
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
..RankingOverlay::default()
apps/aether-gateway/src/dispatch/refs.rs
+1
View File
@@ -208,6 +208,7 @@ mod tests {
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
},
ranking: None,
}
apps/aether-gateway/src/execution_runtime/admission.rs
+62
View File
@@ -0,0 +1,62 @@
//! Shared admission helpers for local upstream execution.
//!
//! The stream candidate loop and long-lived WebSocket turns both need to
//! participate in the same gateway-wide upstream execution gate. Keep the
//! provider abstraction here so tests can supply an isolated gate while
//! production callers use `AppState` directly.
use std::time::Duration;
use aether_runtime::{ConcurrencyGate, ConcurrencyPermit};
use tokio::time::timeout;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
pub(crate) const UPSTREAM_EXECUTION_GATE_NAME: &str = "gateway_upstream_execution";
pub(crate) trait UpstreamExecutionGateProvider {
fn upstream_execution_gate(&self) -> Option<&ConcurrencyGate>;
fn upstream_execution_gate_queue_budget(&self) -> Duration;
}
impl UpstreamExecutionGateProvider for AppState {
fn upstream_execution_gate(&self) -> Option<&ConcurrencyGate> {
self.upstream_execution_gate.as_deref()
}
fn upstream_execution_gate_queue_budget(&self) -> Duration {
self.frontdoor_runtime_guards.internal_gate_queue_budget
}
}
/// Acquires the shared gateway-wide upstream execution permit.
///
/// A missing gate is an intentional configuration (unlimited), so callers
/// receive `Ok(None)`. Saturation keeps the existing candidate-level
/// `AdmissionTimeout` contract used by the HTTP stream path.
pub(crate) async fn acquire_upstream_execution_gate(
state: &(impl UpstreamExecutionGateProvider + ?Sized),
trace_id: &str,
) -> Result<Option<ConcurrencyPermit>, GatewayError> {
let Some(gate) = state.upstream_execution_gate() else {
return Ok(None);
};
let budget = state.upstream_execution_gate_queue_budget();
let gate_wait_started_at = std::time::Instant::now();
match timeout(budget, gate.acquire()).await {
Ok(Ok(permit)) => {
observe_gateway_stage_ms(
"upstream_execution_gate_wait",
gate_wait_started_at.elapsed().as_millis() as u64,
);
Ok(Some(permit))
}
Ok(Err(err)) => Err(GatewayError::Internal(err.to_string())),
Err(_) => Err(GatewayError::AdmissionTimeout {
trace_id: trace_id.to_string(),
gate: UPSTREAM_EXECUTION_GATE_NAME,
queue_budget_ms: budget.as_millis() as u64,
}),
}
}
apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs
+1497
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs
+3
View File
@@ -2573,6 +2573,7 @@ fn json_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body),
body_bytes_b64: None,
@@ -2615,6 +2616,7 @@ fn bytes_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: None,
body: Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
@@ -2637,6 +2639,7 @@ fn execution_result_frame_stream(
payload: StreamFramePayload::Headers {
status_code: result.status_code,
headers: result.headers.clone(),
response_observation: result.response_observation.clone(),
},
},
StreamFrame {
apps/aether-gateway/src/execution_runtime/fallback.rs
+12
View File
@@ -480,6 +480,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -519,6 +520,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -583,6 +585,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -614,6 +617,7 @@ mod tests {
candidate_id: None,
status_code: 401,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -645,6 +649,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: Some(ExecutionError {
@@ -708,6 +713,7 @@ mod tests {
candidate_id: None,
status_code: 404,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -900,6 +906,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1022,6 +1029,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1068,6 +1076,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1174,6 +1183,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1211,6 +1221,7 @@ mod tests {
candidate_id: None,
status_code: 400,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1259,6 +1270,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
apps/aether-gateway/src/execution_runtime/grok.rs
+3
View File
@@ -841,6 +841,7 @@ fn encode_grok_headers_frame(
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: None,
},
})
}
@@ -2157,6 +2158,7 @@ fn grok_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
@@ -2220,6 +2222,7 @@ fn grok_collected_frame_stream(
"application/json".to_string()
},
)]),
response_observation: None,
},
},
StreamFrame {
apps/aether-gateway/src/execution_runtime/kiro_web_search.rs
+2
View File
@@ -279,6 +279,7 @@ fn raw_response_frame_stream(
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: None,
},
},
StreamFrame {
@@ -1449,6 +1450,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"jsonrpc": "2.0",
apps/aether-gateway/src/execution_runtime/mod.rs
+5
View File
@@ -3,6 +3,8 @@ use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
pub(crate) mod admission;
pub(crate) mod attempt_lifecycle;
mod chatgpt_web_image;
mod constants;
mod fallback;
@@ -23,6 +25,9 @@ pub(crate) mod transport;
mod transport_failure;
mod windsurf;
pub(crate) use self::admission::{
acquire_upstream_execution_gate, UpstreamExecutionGateProvider, UPSTREAM_EXECUTION_GATE_NAME,
};
pub(crate) use self::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
pub(crate) use self::constants::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
apps/aether-gateway/src/execution_runtime/oauth_retry.rs
+105 -18
View File
@@ -2,10 +2,10 @@ use aether_contracts::ExecutionPlan;
use tracing::warn;
use crate::orchestration::{
oauth_status_may_be_invalid as status_may_be_oauth_invalid,
local_failover_error_message, oauth_status_may_be_invalid as status_may_be_oauth_invalid,
oauth_status_proves_access_token_invalid as status_proves_access_token_invalid,
};
use crate::state::AgentIdentityAuthConfigFence;
use crate::state::{AgentIdentityAuthConfigFence, CodexRuntimeOAuthObservation};
use crate::{provider_transport::LocalOAuthRefreshError, AppState};
pub(crate) async fn refresh_oauth_plan_auth_for_retry(
@@ -14,6 +14,9 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
status_code: u16,
response_text: Option<&str>,
trace_id: &str,
report_context: Option<&serde_json::Value>,
request_started_at_unix_ms: Option<u64>,
request_order_id: Option<&str>,
) -> bool {
if !status_may_be_oauth_invalid(status_code, response_text) {
return false;
@@ -109,15 +112,49 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
body_excerpt,
..
}) if matches!(refresh_status_code, 400 | 401 | 403) => {
if let Err(err) = state
.persist_local_oauth_refresh_failure_state(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
)
.await
{
let observed_credential_generation =
report_context_string(report_context, "codex_credential_generation");
let runtime_invalid_message = local_failover_error_message(response_text);
let runtime_invalid_reason =
aether_admin::provider::quota::codex_runtime_invalid_reason(
status_code,
runtime_invalid_message.as_deref(),
);
let persist_result = match (request_started_at_unix_ms, request_order_id) {
(Some(request_started_at_unix_ms), Some(request_order_id))
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex") =>
{
state
.persist_local_oauth_refresh_failure_state_observed(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
CodexRuntimeOAuthObservation {
request_started_at_unix_ms,
request_order_id,
observed_credential_generation,
runtime_invalid_reason: runtime_invalid_reason.as_deref(),
},
)
.await
}
_ => {
state
.persist_local_oauth_refresh_failure_state(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
)
.await
}
};
if let Err(err) = persist_result {
warn!(
event_name = "local_oauth_retry_refresh_failure_persist_failed",
log_type = "ops",
@@ -161,6 +198,17 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
}
}
fn report_context_string<'a>(
report_context: Option<&'a serde_json::Value>,
field: &str,
) -> Option<&'a str> {
report_context
.and_then(|context| context.get(field))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
plan.headers
.iter()
@@ -209,6 +257,7 @@ mod tests {
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -312,7 +361,7 @@ mod tests {
}
#[tokio::test]
async fn auto_removes_codex_key_after_request_proven_terminal_refresh_failure() {
async fn retains_codex_key_after_request_proven_terminal_refresh_failure() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
@@ -458,16 +507,34 @@ mod tests {
401,
Some(r#"{"error":"oauth_token_invalid"}"#),
"trace-oauth-retry",
None,
Some(1_000),
Some("01900000-0000-7000-8000-000000000010"),
)
.await;
assert!(!retried);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-retry".to_string()])
.await
.expect("keys should read");
assert!(keys.is_empty());
.expect("keys should read")
.into_iter()
.next()
.expect("request-scoped refresh failure should retain the key");
let invalid_reason = stored_key
.oauth_invalid_reason
.as_deref()
.expect("combined invalid reason should persist");
assert!(invalid_reason.contains("[OAUTH_EXPIRED]"));
assert!(invalid_reason.contains("[REFRESH_FAILED]"));
assert_eq!(
stored_key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")),
Some(&json!("01900000-0000-7000-8000-000000000010"))
);
token_handle.abort();
}
@@ -619,6 +686,9 @@ mod tests {
401,
Some(r#"{"error":"invalid_token"}"#),
"trace-claude-oauth-fence-first",
None,
None,
None,
)
.await
);
@@ -647,6 +717,9 @@ mod tests {
401,
Some(r#"{"error":"invalid_token"}"#),
"trace-claude-oauth-fence-stale",
None,
None,
None,
)
.await
);
@@ -665,6 +738,7 @@ mod tests {
.expect("Claude key should load")
.pop()
.expect("Claude key should exist");
let expected_admin_replacement = admin_replacement.clone();
admin_replacement.encrypted_api_key = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -673,10 +747,23 @@ mod tests {
.expect("admin access token should encrypt"),
);
admin_replacement.expires_at_unix_secs = Some(4_102_444_800);
provider_catalog_repository
.update_key(&admin_replacement)
assert!(provider_catalog_repository
.compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate {
expected_encrypted_auth_config: expected_admin_replacement
.encrypted_auth_config
.clone(),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: expected_admin_replacement.encrypted_api_key.clone(),
auth_type: expected_admin_replacement.auth_type.clone(),
provider_id: expected_admin_replacement.provider_id.clone(),
provider_type: "claude_code".to_string(),
},
key: admin_replacement,
codex_rotation: None,
reset_oauth_runtime: true,
})
.await
.expect("admin replacement should persist");
.expect("admin replacement CAS should run"));
let admin_result = state
.force_local_oauth_refresh_entry(&stale_transport)
apps/aether-gateway/src/execution_runtime/response_header_rules.rs
+101 -4
View File
@@ -10,6 +10,10 @@ use crate::{AppState, GatewayError};
const RESPONSE_HEADER_RULES_KEY: &str = "response_header_rules";
const RESPONSE_HEADER_RULES_CAMEL_KEY: &str = "responseHeaderRules";
const PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY: &str = "provider_response_headers";
const PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY: &str = "provider_request_started_at_unix_ms";
const PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY: &str = "provider_request_order_id";
const PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY: &str =
"provider_response_headers_observed_at_unix_ms";
const RESPONSE_HEADER_RULE_PROTECTED_KEYS: &[&str] = &["content-length"];
const RESPONSE_HEADER_RULES_CACHE_TTL: Duration = Duration::from_secs(5);
@@ -98,6 +102,9 @@ pub(crate) async fn apply_endpoint_response_header_rules(
pub(crate) fn attach_provider_response_headers_to_report_context(
report_context: Option<Value>,
provider_headers: &BTreeMap<String, String>,
provider_request_started_at_unix_ms: u64,
provider_response_headers_observed_at_unix_ms: u64,
provider_request_order_id: &str,
) -> Option<Value> {
let provider_headers = serde_json::to_value(provider_headers).ok()?;
let mut object = match report_context {
@@ -105,9 +112,99 @@ pub(crate) fn attach_provider_response_headers_to_report_context(
Some(other) => Map::from_iter([("seed".to_string(), other)]),
None => Map::new(),
};
object.insert(
PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(),
provider_headers,
);
let observation_is_absent = !object.contains_key(PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY);
if observation_is_absent {
object.insert(
PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(),
provider_headers,
);
object.insert(
PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY.to_string(),
Value::from(provider_request_started_at_unix_ms),
);
object.insert(
PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY.to_string(),
Value::from(provider_response_headers_observed_at_unix_ms),
);
object.insert(
PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY.to_string(),
Value::from(provider_request_order_id),
);
}
Some(Value::Object(object))
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn provider_response_observation_is_first_write_wins() {
let first_headers =
BTreeMap::from([("x-codex-primary-used-percent".to_string(), "10".to_string())]);
let second_headers =
BTreeMap::from([("x-codex-primary-used-percent".to_string(), "20".to_string())]);
let report_context = attach_provider_response_headers_to_report_context(
Some(json!("seed-value")),
&first_headers,
100,
200,
"observation-1",
);
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&second_headers,
300,
400,
"observation-2",
)
.expect("report context should exist");
assert_eq!(report_context["seed"], json!("seed-value"));
assert_eq!(
report_context["provider_response_headers"]["x-codex-primary-used-percent"],
json!("10")
);
assert_eq!(
report_context["provider_request_started_at_unix_ms"],
json!(100)
);
assert_eq!(
report_context["provider_response_headers_observed_at_unix_ms"],
json!(200)
);
assert_eq!(
report_context["provider_request_order_id"],
json!("observation-1")
);
}
#[test]
fn provider_response_observation_does_not_complete_a_partial_triplet() {
let report_context = attach_provider_response_headers_to_report_context(
Some(json!({"provider_response_headers": {"x-existing": "1"}})),
&BTreeMap::from([("x-new".to_string(), "2".to_string())]),
300,
400,
"observation-2",
)
.expect("report context should exist");
assert_eq!(
report_context["provider_response_headers"]["x-existing"],
json!("1")
);
assert!(report_context
.get("provider_request_started_at_unix_ms")
.is_none());
assert!(report_context
.get("provider_response_headers_observed_at_unix_ms")
.is_none());
assert!(report_context.get("provider_request_order_id").is_none());
}
}
apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs
+304 -22
View File
@@ -5,6 +5,7 @@ use serde_json::Value;
use crate::execution_runtime::MAX_STREAM_PREFETCH_BYTES;
const ANTHROPIC_PRECOMMIT_MAX_WAIT: Duration = Duration::from_millis(750);
const GEMINI_PRECOMMIT_MAX_WAIT: Duration = Duration::from_millis(750);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum StreamCommitPolicy {
@@ -14,6 +15,10 @@ pub(super) enum StreamCommitPolicy {
max_bytes: usize,
max_wait: Duration,
},
FirstGeminiSemanticEvent {
max_bytes: usize,
max_wait: Duration,
},
}
impl StreamCommitPolicy {
@@ -51,6 +56,12 @@ impl StreamCommitPolicy {
max_wait: ANTHROPIC_PRECOMMIT_MAX_WAIT,
};
}
if provider_api_format.eq_ignore_ascii_case("gemini:generate_content") {
return Self::FirstGeminiSemanticEvent {
max_bytes: MAX_STREAM_PREFETCH_BYTES,
max_wait: GEMINI_PRECOMMIT_MAX_WAIT,
};
}
return Self::ResponseHeaders;
}
@@ -78,12 +89,16 @@ impl StreamCommitPolicy {
}
pub(super) const fn requires_bounded_frame_wait(self) -> bool {
matches!(self, Self::FirstAnthropicSemanticEvent { .. })
matches!(
self,
Self::FirstAnthropicSemanticEvent { .. } | Self::FirstGeminiSemanticEvent { .. }
)
}
pub(super) const fn max_precommit_wait(self) -> Option<Duration> {
match self {
Self::FirstAnthropicSemanticEvent { max_wait, .. } => Some(max_wait),
Self::FirstAnthropicSemanticEvent { max_wait, .. }
| Self::FirstGeminiSemanticEvent { max_wait, .. } => Some(max_wait),
Self::ResponseHeaders | Self::FirstClassifiedBody => None,
}
}
@@ -91,6 +106,10 @@ impl StreamCommitPolicy {
pub(super) const fn is_native_anthropic(self) -> bool {
matches!(self, Self::FirstAnthropicSemanticEvent { .. })
}
pub(super) const fn is_gemini(self) -> bool {
matches!(self, Self::FirstGeminiSemanticEvent { .. })
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -113,6 +132,7 @@ pub(super) struct StreamCommitGate {
state: StreamCommitState,
observed_bytes: usize,
anthropic: AnthropicSsePrecommitInspector,
gemini: GeminiSsePrecommitInspector,
}
impl StreamCommitGate {
@@ -127,6 +147,7 @@ impl StreamCommitGate {
state,
observed_bytes: 0,
anthropic: AnthropicSsePrecommitInspector::default(),
gemini: GeminiSsePrecommitInspector::default(),
}
}
@@ -143,21 +164,32 @@ impl StreamCommitGate {
return StreamPrecommitObservation::Commit;
}
let StreamCommitPolicy::FirstAnthropicSemanticEvent { max_bytes, .. } = self.policy else {
return StreamPrecommitObservation::Pending;
let (max_bytes, observation) = match self.policy {
StreamCommitPolicy::FirstAnthropicSemanticEvent { max_bytes, .. } => {
(max_bytes, self.anthropic.observe(chunk, max_bytes))
}
StreamCommitPolicy::FirstGeminiSemanticEvent { max_bytes, .. } => {
(max_bytes, self.gemini.observe(chunk, max_bytes))
}
StreamCommitPolicy::ResponseHeaders | StreamCommitPolicy::FirstClassifiedBody => {
return StreamPrecommitObservation::Pending;
}
};
self.observed_bytes = self.observed_bytes.saturating_add(chunk.len());
match self.anthropic.observe(chunk, max_bytes) {
AnthropicSseObservation::Pending => {}
AnthropicSseObservation::SemanticEvent => {
match observation {
SemanticSseObservation::Pending => {}
SemanticSseObservation::SemanticEvent => {
self.state = StreamCommitState::Committed;
return StreamPrecommitObservation::Commit;
}
AnthropicSseObservation::Error(body_json) => {
SemanticSseObservation::Error {
status_code,
body_json,
} => {
self.state = StreamCommitState::Terminal;
return StreamPrecommitObservation::UpstreamError {
status_code: anthropic_error_status_code(&body_json),
status_code,
body_json,
};
}
@@ -179,10 +211,10 @@ impl StreamCommitGate {
}
#[derive(Debug)]
enum AnthropicSseObservation {
enum SemanticSseObservation {
Pending,
SemanticEvent,
Error(Value),
Error { status_code: u16, body_json: Value },
}
#[derive(Debug, Default)]
@@ -191,7 +223,7 @@ struct AnthropicSsePrecommitInspector {
}
impl AnthropicSsePrecommitInspector {
fn observe(&mut self, chunk: &[u8], max_bytes: usize) -> AnthropicSseObservation {
fn observe(&mut self, chunk: &[u8], max_bytes: usize) -> SemanticSseObservation {
let remaining = max_bytes.saturating_sub(self.buffered.len());
let truncated = chunk.len() > remaining;
self.buffered
@@ -201,15 +233,44 @@ impl AnthropicSsePrecommitInspector {
let record = self.buffered[..record_end].to_vec();
self.buffered.drain(..record_end + separator_len);
match classify_anthropic_sse_record(&record) {
AnthropicSseObservation::Pending => {}
SemanticSseObservation::Pending => {}
decision => return decision,
}
}
if truncated {
AnthropicSseObservation::SemanticEvent
SemanticSseObservation::SemanticEvent
} else {
AnthropicSseObservation::Pending
SemanticSseObservation::Pending
}
}
}
#[derive(Debug, Default)]
struct GeminiSsePrecommitInspector {
buffered: Vec<u8>,
}
impl GeminiSsePrecommitInspector {
fn observe(&mut self, chunk: &[u8], max_bytes: usize) -> SemanticSseObservation {
let remaining = max_bytes.saturating_sub(self.buffered.len());
let truncated = chunk.len() > remaining;
self.buffered
.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
while let Some((record_end, separator_len)) = find_sse_record_boundary(&self.buffered) {
let record = self.buffered[..record_end].to_vec();
self.buffered.drain(..record_end + separator_len);
match classify_gemini_sse_record(&record) {
SemanticSseObservation::Pending => {}
decision => return decision,
}
}
if truncated {
SemanticSseObservation::SemanticEvent
} else {
SemanticSseObservation::Pending
}
}
}
@@ -249,9 +310,9 @@ fn next_sse_line_ending(buffer: &[u8], start: usize) -> Option<(usize, usize)> {
Some((index, ending_len))
}
fn classify_anthropic_sse_record(record: &[u8]) -> AnthropicSseObservation {
fn classify_anthropic_sse_record(record: &[u8]) -> SemanticSseObservation {
let Ok(record) = std::str::from_utf8(record) else {
return AnthropicSseObservation::Pending;
return SemanticSseObservation::Pending;
};
let normalized_record = record.replace("\r\n", "\n").replace('\r', "\n");
let mut event_type = None;
@@ -275,15 +336,18 @@ fn classify_anthropic_sse_record(record: &[u8]) -> AnthropicSseObservation {
}
}
if data.trim().is_empty() {
return AnthropicSseObservation::Pending;
return SemanticSseObservation::Pending;
}
let Ok(body_json) = serde_json::from_str::<Value>(data.trim()) else {
return AnthropicSseObservation::Pending;
return SemanticSseObservation::Pending;
};
let payload_type = body_json.get("type").and_then(Value::as_str).map(str::trim);
if event_type == Some("error") || payload_type == Some("error") {
return AnthropicSseObservation::Error(body_json);
return SemanticSseObservation::Error {
status_code: anthropic_error_status_code(&body_json),
body_json,
};
}
let semantic_type = match (event_type, payload_type) {
@@ -292,12 +356,120 @@ fn classify_anthropic_sse_record(record: &[u8]) -> AnthropicSseObservation {
_ => None,
};
if semantic_type.is_some_and(is_anthropic_semantic_event_type) {
AnthropicSseObservation::SemanticEvent
SemanticSseObservation::SemanticEvent
} else {
AnthropicSseObservation::Pending
SemanticSseObservation::Pending
}
}
fn classify_gemini_sse_record(record: &[u8]) -> SemanticSseObservation {
let Ok(record) = std::str::from_utf8(record) else {
return SemanticSseObservation::Pending;
};
let data = record
.replace("\r\n", "\n")
.replace('\r', "\n")
.lines()
.filter_map(|line| line.strip_prefix("data:").map(str::trim_start))
.collect::<Vec<_>>()
.join("\n");
if data.trim().is_empty() {
return SemanticSseObservation::Pending;
}
if data.trim() == "[DONE]" {
return SemanticSseObservation::SemanticEvent;
}
let Ok(body_json) = serde_json::from_str::<Value>(data.trim()) else {
return SemanticSseObservation::Pending;
};
let response = body_json.get("response").unwrap_or(&body_json);
let Some(candidates) = response.get("candidates").and_then(Value::as_array) else {
return SemanticSseObservation::Pending;
};
for candidate in candidates {
let finish_reason = candidate
.get("finishReason")
.or_else(|| candidate.get("finish_reason"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(finish_reason) = finish_reason.filter(|reason| {
matches!(
*reason,
"MALFORMED_FUNCTION_CALL"
| "UNEXPECTED_TOOL_CALL"
| "TOO_MANY_TOOL_CALLS"
| "MISSING_THOUGHT_SIGNATURE"
| "MALFORMED_RESPONSE"
)
}) {
let message = candidate
.get("finishMessage")
.or_else(|| candidate.get("finish_message"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("Gemini stream ended with {finish_reason}"));
return SemanticSseObservation::Error {
status_code: 502,
body_json: serde_json::json!({
"error": {
"type": "upstream_gemini_finish_error",
"code": finish_reason,
"message": message,
"upstream_status": 200
}
}),
};
}
if finish_reason.is_some() {
return SemanticSseObservation::SemanticEvent;
}
let Some(parts) = candidate
.get("content")
.and_then(|content| content.get("parts"))
.and_then(Value::as_array)
else {
continue;
};
if parts.iter().any(gemini_part_is_client_semantic) {
return SemanticSseObservation::SemanticEvent;
}
}
SemanticSseObservation::Pending
}
fn gemini_part_is_client_semantic(part: &Value) -> bool {
let Some(part) = part.as_object() else {
return true;
};
if part
.keys()
.any(|key| !matches!(key.as_str(), "text" | "thought" | "thoughtSignature"))
{
return true;
}
if part.get("thought").and_then(Value::as_bool) == Some(true) {
return false;
}
if part.keys().all(|key| key == "thoughtSignature") {
return false;
}
if part
.get("text")
.and_then(Value::as_str)
.is_some_and(|text| !text.is_empty())
{
return true;
}
false
}
fn is_anthropic_semantic_event_type(event_type: &str) -> bool {
matches!(
event_type,
@@ -346,6 +518,13 @@ mod tests {
}
}
fn gemini_policy() -> StreamCommitPolicy {
StreamCommitPolicy::FirstGeminiSemanticEvent {
max_bytes: 16_384,
max_wait: Duration::from_millis(750),
}
}
#[test]
fn policy_selects_bounded_anthropic_gate_only_for_native_same_format_sse() {
let native = StreamCommitPolicy::for_response(
@@ -384,6 +563,109 @@ mod tests {
.commits_on_response_headers());
}
#[test]
fn policy_selects_bounded_gemini_gate_for_event_streams() {
let policy = StreamCommitPolicy::for_response(
true,
Some("text/event-stream"),
"gemini:generate_content",
"openai:responses",
false,
true,
false,
);
assert!(policy.is_gemini());
assert!(policy.requires_bounded_frame_wait());
assert_eq!(
policy.max_precommit_wait(),
Some(Duration::from_millis(750))
);
}
#[test]
fn gemini_gate_waits_through_thought_and_commits_on_text() {
let mut gate = StreamCommitGate::new(gemini_policy());
let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"checking\"}]}}]}}\n\n";
let text = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"answer\"}]}}]}}\n\n";
assert_eq!(
gate.observe_provider_bytes(thought),
StreamPrecommitObservation::Pending
);
assert_eq!(
gate.observe_provider_bytes(text),
StreamPrecommitObservation::Commit
);
assert_eq!(gate.state(), StreamCommitState::Committed);
}
#[test]
fn gemini_gate_commits_on_function_call_even_with_thought_marker() {
let mut gate = StreamCommitGate::new(gemini_policy());
let tool_call = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"functionCall\":{\"name\":\"validate\",\"args\":{}}}]}}]}}\n\n";
assert_eq!(
gate.observe_provider_bytes(tool_call),
StreamPrecommitObservation::Commit
);
assert_eq!(gate.state(), StreamCommitState::Committed);
}
#[test]
fn gemini_gate_rejects_malformed_function_call_before_commit() {
let mut gate = StreamCommitGate::new(gemini_policy());
let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"calling\"}]}}]}}\n\n";
let malformed = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]},\"finishReason\":\"MALFORMED_FUNCTION_CALL\",\"finishMessage\":\"Malformed function call: Function call is empty - no input to parse.\"}]}}\n\n";
assert_eq!(
gate.observe_provider_bytes(thought),
StreamPrecommitObservation::Pending
);
let StreamPrecommitObservation::UpstreamError {
status_code,
body_json,
} = gate.observe_provider_bytes(malformed)
else {
panic!("malformed Gemini function call should fail before stream commit");
};
assert_eq!(status_code, 502);
assert_eq!(body_json["error"]["code"], "MALFORMED_FUNCTION_CALL");
assert_eq!(
body_json["error"]["message"],
"Malformed function call: Function call is empty - no input to parse."
);
assert_eq!(gate.state(), StreamCommitState::Terminal);
}
#[test]
fn gemini_gate_detects_malformed_function_call_across_chunk_boundaries() {
let malformed = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]},\"finishReason\":\"MALFORMED_FUNCTION_CALL\",\"finishMessage\":\"empty call\"}]}}\r\n\r\n";
for split in 1..malformed.len() {
let mut gate = StreamCommitGate::new(gemini_policy());
let first_observation = gate.observe_provider_bytes(&malformed[..split]);
if !matches!(
first_observation,
StreamPrecommitObservation::UpstreamError {
status_code: 502,
..
}
) {
assert_eq!(first_observation, StreamPrecommitObservation::Pending);
assert!(matches!(
gate.observe_provider_bytes(&malformed[split..]),
StreamPrecommitObservation::UpstreamError {
status_code: 502,
..
}
));
}
assert_eq!(gate.state(), StreamCommitState::Terminal);
}
}
#[test]
fn gate_detects_anthropic_error_across_every_chunk_boundary() {
let event = b"event: error\r\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}\r\n\r\n";
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+482 -36
View File
@@ -11,8 +11,8 @@ use std::time::{Duration, Instant};
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope};
use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage,
StreamFrame, StreamFramePayload,
ExecutionPlan, ExecutionResponseObservation, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, UpsertRequestCandidateRecord,
@@ -64,8 +64,14 @@ use crate::ai_serving::api::{
extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream,
maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter,
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
UPSTREAM_IS_STREAM_KEY,
};
use crate::ai_serving::is_openai_responses_family_format;
use crate::ai_serving::record_local_runtime_candidate_skip_reason;
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
};
@@ -112,15 +118,16 @@ use crate::execution_runtime::{
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, build_local_error_flow_metadata, classify_failure_disposition,
cyber_continue_failover_enabled, trace_upstream_response_body, with_error_flow_report_context,
cyber_continue_failover_enabled, spawn_local_oauth_success_effect,
trace_upstream_response_body, with_error_flow_report_context,
with_upstream_response_report_context, FailureDisposition, FailureTokenAction,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis,
LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect,
LocalPoolErrorEffect,
LocalOAuthSuccessEffect, LocalPoolErrorEffect,
};
use crate::provider_pool_demand::{
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, persist_local_request_candidate_status_record,
@@ -143,7 +150,6 @@ use crate::{
AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream";
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n";
const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024;
@@ -1249,6 +1255,9 @@ async fn execute_in_process_stream_with_oauth_retry(
retry_status_code,
response_text.as_deref(),
trace_id,
report_context,
Some(execution.response_observation.request_started_at_unix_ms),
Some(&execution.response_observation.request_order_id),
)
.await
{
@@ -2818,6 +2827,7 @@ async fn execute_stream_from_direct_passthrough(
stream_precommit_committed: _,
response,
started_at: upstream_started_at,
response_observation,
stream_first_byte_timeout,
upstream_target_permit,
} = execution;
@@ -2834,8 +2844,23 @@ async fn execute_stream_from_direct_passthrough(
let request_id = plan.request_id.clone();
let candidate_id = plan.candidate_id.clone();
let request_id_for_log = short_request_id(request_id.as_str());
let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers);
let mut report_context = attach_provider_response_headers_to_report_context(
report_context,
&headers,
response_observation.request_started_at_unix_ms,
response_observation.response_headers_observed_at_unix_ms,
&response_observation.request_order_id,
);
spawn_local_oauth_success_effect(
state.clone(),
&plan,
report_context.as_ref(),
LocalOAuthSuccessEffect {
status_code,
request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms),
request_order_id: Some(&response_observation.request_order_id),
},
);
if status_code == 200 {
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
@@ -3748,6 +3773,46 @@ async fn execute_execution_runtime_stream_inner(
plan_kind,
report_context.as_ref(),
);
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let provider_in_flight_started_at = Instant::now();
let mut provider_pool_in_flight_guard =
match acquire_provider_pool_execution_guard(state, &plan).await? {
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
record_local_runtime_candidate_skip_reason(
state,
trace_id,
"provider_key_concurrency_limit_reached",
);
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
*retry_scope = AiAttemptRetryScope::Candidate;
}
if let Some(snapshot) = request_candidate_status_snapshot.as_ref() {
record_local_request_candidate_status_snapshot(
state,
snapshot,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Skipped,
status_code: Some(http::StatusCode::TOO_MANY_REQUESTS.as_u16()),
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
error_message: Some(format!(
"provider key concurrency limit reached: {limit}"
)),
latency_ms: Some(0),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(candidate_started_unix_secs),
},
)
.await;
}
return Ok(None);
}
};
observe_gateway_stage_trace_ms(
&mut stage_trace,
"stream_provider_in_flight",
provider_in_flight_started_at.elapsed().as_millis() as u64,
);
// Inline passthrough records its lifecycle seed after upstream headers are
// available. Avoid constructing a throwaway seed on the common path.
let mut lifecycle_seed = (!defer_stream_pending_for_direct_inline)
@@ -3757,7 +3822,6 @@ async fn execute_execution_runtime_stream_inner(
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
lifecycle_pending_recorded = true;
}
let candidate_started_unix_secs = current_request_candidate_unix_ms();
if let Some(snapshot) = request_candidate_status_snapshot.clone() {
record_local_request_candidate_status_snapshot(
state,
@@ -3786,20 +3850,6 @@ async fn execute_execution_runtime_stream_inner(
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
let provider_in_flight_started_at = Instant::now();
let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan.request_id.as_str(),
plan.candidate_id.as_deref(),
key_id.as_str(),
)
.await;
observe_gateway_stage_trace_ms(
&mut stage_trace,
"stream_provider_in_flight",
provider_in_flight_started_at.elapsed().as_millis() as u64,
);
match maybe_execute_grok_stream(&plan, report_context.as_ref()).await {
Ok(Some(grok_stream)) => {
return execute_stream_from_frame_stream_with_retry_scope(
@@ -3819,6 +3869,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -3891,6 +3942,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -3963,6 +4015,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -4035,6 +4088,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -4192,6 +4246,15 @@ async fn execute_execution_runtime_stream_inner(
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
lifecycle_pending_recorded = true;
}
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&execution.headers,
execution.response_observation.request_started_at_unix_ms,
execution
.response_observation
.response_headers_observed_at_unix_ms,
&execution.response_observation.request_order_id,
);
let stream_precommit_committed = execution.stream_precommit_committed;
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
return execute_stream_from_frame_stream_with_retry_scope(
@@ -4211,6 +4274,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out,
retry_fallback_out,
None,
)
.await;
}
@@ -4327,6 +4391,15 @@ async fn execute_execution_runtime_stream_inner(
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
lifecycle_pending_recorded = true;
}
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&execution.headers,
execution.response_observation.request_started_at_unix_ms,
execution
.response_observation
.response_headers_observed_at_unix_ms,
&execution.response_observation.request_order_id,
);
let stream_precommit_committed = execution.stream_precommit_committed;
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
return execute_stream_from_frame_stream_with_retry_scope(
@@ -4346,10 +4419,13 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
let remote_request_started_at_unix_ms = current_request_candidate_unix_ms();
let remote_request_order_id = uuid::Uuid::now_v7().to_string();
let response = match post_stream_plan_to_remote_execution_runtime(
state,
remote_execution_runtime_base_url,
@@ -4431,6 +4507,12 @@ async fn execute_execution_runtime_stream_inner(
)?));
}
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
let remote_fallback_observation = ExecutionResponseObservation {
request_started_at_unix_ms: remote_request_started_at_unix_ms,
response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms,
request_order_id: remote_request_order_id,
};
let frame_stream = response
.bytes_stream()
.map_err(|err| IoError::other(err.to_string()))
@@ -4452,6 +4534,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
Some(remote_fallback_observation),
)
.await;
}
@@ -4471,13 +4554,86 @@ fn decode_stream_data_chunk(
fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
headers
.get("content-type")
.map(String::as_str)
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("content-type"))
.map(|(_, value)| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
}
fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool {
report_context
.and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}
fn response_headers_have_octet_stream_content_type(headers: &BTreeMap<String, String>) -> bool {
headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("content-type"))
.map(|(_, value)| value.as_str())
.and_then(|value| value.split(';').next())
.map(str::trim)
.is_some_and(|value| value.eq_ignore_ascii_case("application/octet-stream"))
}
fn response_headers_have_only_identity_content_encoding(
headers: &BTreeMap<String, String>,
) -> bool {
headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("content-encoding"))
.map(|(_, value)| value.as_str())
.is_none_or(|value| {
value
.split(',')
.map(str::trim)
.all(|coding| coding.is_empty() || coding.eq_ignore_ascii_case("identity"))
})
}
fn plan_kind_uses_text_event_stream(plan_kind: &str) -> bool {
matches!(
plan_kind,
OPENAI_CHAT_STREAM_PLAN_KIND
| OPENAI_RESPONSES_STREAM_PLAN_KIND
| OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND
| OPENAI_IMAGE_STREAM_PLAN_KIND
| CLAUDE_CHAT_STREAM_PLAN_KIND
| CLAUDE_CLI_STREAM_PLAN_KIND
| GEMINI_CHAT_STREAM_PLAN_KIND
| GEMINI_CLI_STREAM_PLAN_KIND
| GEMINI_INTERACTIONS_STREAM_PLAN_KIND
)
}
fn should_normalize_declared_stream_response_headers(
plan_kind: &str,
status_code: u16,
headers: &BTreeMap<String, String>,
report_context: Option<&Value>,
) -> bool {
plan_kind_uses_text_event_stream(plan_kind)
&& (200..300).contains(&status_code)
&& report_context_upstream_is_stream(report_context)
&& response_headers_have_octet_stream_content_type(headers)
&& response_headers_have_only_identity_content_encoding(headers)
&& !headers
.keys()
.any(|name| name.eq_ignore_ascii_case("content-length"))
}
fn normalize_declared_stream_response_headers(headers: &mut BTreeMap<String, String>) {
headers.retain(|name, _| {
!name.eq_ignore_ascii_case("content-encoding")
&& !name.eq_ignore_ascii_case("content-length")
&& !name.eq_ignore_ascii_case("content-type")
});
headers.insert("content-type".to_string(), "text/event-stream".to_string());
}
fn parse_prefetched_sync_json_body(body: &[u8]) -> Option<Value> {
let stripped = strip_utf8_bom_and_ws(body);
serde_json::from_slice::<Value>(stripped).ok()
@@ -5481,6 +5637,7 @@ async fn execute_stream_from_frame_stream(
in_flight_guard,
None,
None,
None,
)
.await
}
@@ -5503,6 +5660,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
in_flight_guard: Option<ProviderPoolInFlightGuard>,
mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
fallback_response_observation: Option<ExecutionResponseObservation>,
) -> Result<Option<Response<Body>>, GatewayError> {
let request_id = plan.request_id.as_str();
let request_id_for_log = short_request_id(request_id);
@@ -5535,14 +5693,37 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let StreamFramePayload::Headers {
status_code,
mut headers,
response_observation,
} = first_frame.payload
else {
return Err(GatewayError::Internal(
"execution runtime stream must start with headers frame".to_string(),
));
};
let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers);
let response_observation = response_observation
.or(fallback_response_observation)
.unwrap_or(ExecutionResponseObservation {
request_started_at_unix_ms: candidate_started_unix_secs,
response_headers_observed_at_unix_ms: current_request_candidate_unix_ms(),
request_order_id: uuid::Uuid::now_v7().to_string(),
});
let mut report_context = attach_provider_response_headers_to_report_context(
report_context,
&headers,
response_observation.request_started_at_unix_ms,
response_observation.response_headers_observed_at_unix_ms,
&response_observation.request_order_id,
);
spawn_local_oauth_success_effect(
state.clone(),
&plan,
report_context.as_ref(),
LocalOAuthSuccessEffect {
status_code,
request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms),
request_order_id: Some(&response_observation.request_order_id),
},
);
if status_code == 200 {
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
@@ -5964,6 +6145,32 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers.insert("content-type".to_string(), "text/event-stream".to_string());
}
let upstream_content_type = upstream_headers.get("content-type").map(String::as_str);
let normalized_declared_stream_headers = private_stream_normalizer.is_none()
&& local_stream_rewriter.is_none()
&& should_normalize_declared_stream_response_headers(
plan_kind,
status_code,
&upstream_headers,
report_context.as_ref(),
);
if normalized_declared_stream_headers {
normalize_declared_stream_response_headers(&mut headers);
debug!(
event_name = "execution_runtime_stream_content_type_corrected",
log_type = "debug",
trace_id = %trace_id,
request_id = %request_id_for_log,
candidate_id = ?candidate_id,
plan_kind,
provider_name,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
model_name,
candidate_index = candidate_index.as_str(),
upstream_content_type = upstream_content_type.unwrap_or("-"),
"gateway normalized declared upstream stream response headers for the client"
);
}
let prefetch_for_cyber_failover =
is_openai_responses_family_format(plan.provider_api_format.as_str())
&& cyber_continue_failover_enabled(state).await;
@@ -6289,7 +6496,9 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
}
}
let inspection = if stream_commit_policy.is_native_anthropic() {
let inspection = if stream_commit_policy.is_native_anthropic()
|| stream_commit_policy.is_gemini()
{
StreamPrefetchInspection::NeedMore
} else {
inspect_prefetched_stream_body(
@@ -6651,8 +6860,9 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let native_anthropic_stream_for_report = stream_commit_policy.is_native_anthropic();
let plan_for_report = plan;
let emit_passthrough_sse_terminal_error = (skip_direct_finalize_prefetch
|| stream_commit_policy.is_native_anthropic())
&& response_headers_indicate_sse(&upstream_headers)
|| stream_commit_policy.is_native_anthropic()
|| normalized_declared_stream_headers)
&& (response_headers_indicate_sse(&upstream_headers) || normalized_declared_stream_headers)
&& !is_openai_image_stream_for_report;
let plan_kind_for_report = plan_kind.to_string();
let stream_started_at_for_report = stream_started_at;
@@ -8014,20 +8224,23 @@ mod tests {
execute_execution_runtime_stream, execute_in_process_stream_with_oauth_retry,
execute_stream_from_frame_stream, execute_stream_from_frame_stream_with_retry_scope,
maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
parse_direct_passthrough_mode, prefetch_direct_stream_error_body,
prefetched_openai_responses_body_has_output_boundary,
normalize_declared_stream_response_headers, parse_direct_passthrough_mode,
prefetch_direct_stream_error_body, prefetched_openai_responses_body_has_output_boundary,
record_sync_terminal_usage_with_handoff,
record_sync_terminal_usage_with_handoff_after_spawn,
resolve_provider_stream_error_status_code, select_direct_anthropic_prefetch_wait,
should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream,
should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done,
stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish,
should_limit_direct_finalize_prefetch, should_normalize_declared_stream_response_headers,
should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
stream_chunk_contains_sse_done, stream_requires_observed_terminal_event,
stream_terminal_summary_missing_observed_finish,
stream_terminal_summary_missing_observed_finish_with_requirement,
stream_terminal_summary_represents_failure_with_requirement,
ClientVisibleStreamCompletionTracker, DirectPassthroughFinalizer,
DirectPassthroughFinalizerCore, DirectPassthroughInlineBodyState, DirectPassthroughMode,
PostStopFrameReadBudget, PostStopLimitedStreamReader, ProviderStreamErrorInspection,
ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL,
ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL,
};
use crate::control::GatewayControlDecision;
use crate::stage_metrics::RequestStageTrace;
@@ -8310,6 +8523,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8389,6 +8603,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8431,6 +8646,7 @@ mod tests {
None,
Some(&mut retry_scope),
None,
None,
)
.await
.expect("prefetch transport execution should resolve");
@@ -8480,6 +8696,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8522,6 +8739,7 @@ mod tests {
None,
Some(&mut retry_scope),
None,
None,
)
.await
.expect("prefetch HTTP status execution should resolve");
@@ -8567,6 +8785,36 @@ mod tests {
}
}
fn antigravity_gemini_stream_plan(request_id: &str) -> ExecutionPlan {
ExecutionPlan {
request_id: request_id.to_string(),
candidate_id: Some(format!("candidate-{request_id}")),
provider_name: Some("antigravity".to_string()),
provider_id: format!("provider-{request_id}"),
endpoint_id: format!("endpoint-{request_id}"),
key_id: format!("key-{request_id}"),
method: "POST".to_string(),
url: "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent".to_string(),
headers: BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("accept".to_string(), "text/event-stream".to_string()),
]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gemini-3.7-flash-tiered",
"contents": [{"role": "user", "parts": [{"text": "validate"}]}]
})),
stream: true,
client_api_format: "openai:responses".to_string(),
provider_api_format: "gemini:generate_content".to_string(),
model_name: Some("gemini-3.7-flash-tiered".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
}
}
struct StreamDropFlag(Arc<AtomicBool>);
impl Drop for StreamDropFlag {
@@ -8680,6 +8928,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
for chunk in chunks {
@@ -8725,6 +8974,7 @@ mod tests {
None,
Some(&mut retry_scope),
Some(&mut fallback_response),
None,
)
.await
.expect("native Anthropic stream execution should succeed");
@@ -9364,6 +9614,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -9413,6 +9664,7 @@ mod tests {
None,
None,
None,
None,
),
)
.await
@@ -9850,6 +10102,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
}
@@ -10310,6 +10563,92 @@ mod tests {
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
}
#[tokio::test]
async fn malformed_antigravity_function_call_retries_before_stream_commit() {
let request_id = "req-antigravity-malformed-function-call";
let plan = antigravity_gemini_stream_plan(request_id);
let provider_catalog = provider_catalog_for_plan(
&plan,
Some(json!({
"failover_rules": {
"continue_status_codes": [502]
}
})),
);
let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data_state);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
for chunk in [
r#"data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":"Validating the document."}]} }],"modelVersion":"gemini-3.7-flash-tiered"}}
"#,
r#"data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"signature","text":""}]},"finishReason":"MALFORMED_FUNCTION_CALL","finishMessage":"Malformed function call: Function call is empty - no input to parse."}],"modelVersion":"gemini-3.7-flash-tiered"}}
"#,
] {
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: None,
text: Some(chunk.to_string()),
},
}));
}
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame::eof()));
}
.boxed();
let mut retry_scope = AiAttemptRetryScope::Provider;
let response = execute_stream_from_frame_stream_with_retry_scope(
&state,
plan,
"trace-antigravity-malformed-function-call",
&test_decision(),
OPENAI_RESPONSES_STREAM_PLAN_KIND,
Some("openai_responses_stream_success".to_string()),
Some(json!({
"request_id": request_id,
"candidate_id": format!("candidate-{request_id}"),
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "gemini:generate_content",
"client_api_format": "openai:responses",
"needs_conversion": true
})),
crate::clock::current_unix_ms(),
Instant::now(),
RequestStageTrace::from_env(),
true,
frame_stream,
false,
None,
Some(&mut retry_scope),
None,
None,
)
.await
.expect("malformed Antigravity stream should resolve through failover");
assert!(response.is_none());
assert_eq!(retry_scope, AiAttemptRetryScope::Candidate);
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),
@@ -11532,6 +11871,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -11660,6 +12000,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -11969,6 +12310,109 @@ mod tests {
));
}
#[test]
fn declared_stream_response_headers_are_normalized_without_body_inspection() {
let mut headers = BTreeMap::from([
(
"Content-Type".to_string(),
"Application/Octet-Stream; charset=binary".to_string(),
),
("Content-Encoding".to_string(), "identity".to_string()),
("x-upstream-header".to_string(), "preserved".to_string()),
]);
assert!(should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&headers,
Some(&json!({"upstream_is_stream": true})),
));
headers.insert("Content-Length".to_string(), "4096".to_string());
normalize_declared_stream_response_headers(&mut headers);
assert_eq!(
headers.get("content-type").map(String::as_str),
Some("text/event-stream")
);
assert!(!headers
.keys()
.any(|name| name.eq_ignore_ascii_case("content-encoding")));
assert!(!headers
.keys()
.any(|name| name.eq_ignore_ascii_case("content-length")));
assert_eq!(
headers.get("x-upstream-header").map(String::as_str),
Some("preserved")
);
}
#[test]
fn declared_stream_header_normalization_requires_success_and_stream_context() {
let headers = BTreeMap::from([(
"content-type".to_string(),
"application/octet-stream".to_string(),
)]);
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
500,
&headers,
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&headers,
Some(&json!({"upstream_is_stream": false})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string(),)]),
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&BTreeMap::from([("content-type".to_string(), "application/json".to_string(),)]),
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&BTreeMap::from([("content-type".to_string(), "text/plain".to_string(),)]),
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&BTreeMap::from([
(
"content-type".to_string(),
"application/octet-stream".to_string(),
),
("content-encoding".to_string(), "gzip".to_string()),
]),
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
OPENAI_CHAT_STREAM_PLAN_KIND,
200,
&BTreeMap::from([
(
"content-type".to_string(),
"application/octet-stream".to_string(),
),
("content-length".to_string(), "128".to_string()),
]),
Some(&json!({"upstream_is_stream": true})),
));
assert!(!should_normalize_declared_stream_response_headers(
GEMINI_FILES_DOWNLOAD_PLAN_KIND,
200,
&headers,
Some(&json!({"upstream_is_stream": true})),
));
}
#[test]
fn skips_prefetch_for_event_streams_even_when_cross_format_or_rewritten() {
assert!(should_skip_direct_finalize_prefetch(
@@ -12384,6 +12828,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -13700,6 +14145,7 @@ mod tests {
Some(json!({
"provider_api_format": "openai:responses",
"client_api_format": "openai:responses",
"upstream_is_stream": true,
})),
)
.await
apps/aether-gateway/src/execution_runtime/stream_pump.rs
+71 -41
View File
@@ -4,8 +4,9 @@ use std::io::Error as IoError;
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionResponseObservation,
ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame, StreamFramePayload,
StreamFrameType,
};
use async_stream::stream;
use axum::body::Bytes;
@@ -44,6 +45,7 @@ pub(crate) fn build_direct_execution_frame_stream(
stream_precommit_committed: _,
response,
started_at,
response_observation,
stream_first_byte_timeout,
upstream_target_permit,
} = execution;
@@ -108,7 +110,11 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
match encode_headers_frame(status_code, response_headers) {
match encode_headers_frame(
status_code,
response_headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -153,7 +159,11 @@ pub(crate) fn build_direct_execution_frame_stream(
upstream_bytes,
first_byte_timeout,
}) => {
match encode_headers_frame(status_code, original_headers) {
match encode_headers_frame(
status_code,
original_headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -192,7 +202,11 @@ pub(crate) fn build_direct_execution_frame_stream(
return;
}
match encode_headers_frame(status_code, headers) {
match encode_headers_frame(
status_code,
headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -611,12 +625,14 @@ pub(crate) fn build_direct_execution_frame_stream(
fn encode_headers_frame(
status_code: u16,
headers: BTreeMap<String, String>,
response_observation: &ExecutionResponseObservation,
) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: Some(response_observation.clone()),
},
})
}
@@ -1606,43 +1622,47 @@ mod tests {
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/responses",
post(|| async {
let body = serde_json::json!({
"id": "resp_sync_bridge_123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg_sync_bridge_123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from buffered JSON stream",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
});
let mut response = axum::http::Response::new(Body::from(
serde_json::to_vec(&body).expect("json should encode"),
));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response
}),
);
axum::serve(listener, app)
let (mut socket, _) = listener.accept().await.expect("client should connect");
let mut request = [0_u8; 4096];
let _ = socket
.read(&mut request)
.await
.expect("server should start");
.expect("request should read");
let body = serde_json::to_vec(&serde_json::json!({
"id": "resp_sync_bridge_123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg_sync_bridge_123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from buffered JSON stream",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}))
.expect("json should encode");
socket
.write_all(
format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n",
body.len()
)
.as_bytes(),
)
.await
.expect("headers should write");
socket.flush().await.expect("headers should flush");
tokio::time::sleep(Duration::from_millis(75)).await;
socket.write_all(&body).await.expect("body should write");
});
let runtime = DirectSyncExecutionRuntime::new();
@@ -1678,6 +1698,12 @@ mod tests {
})
.await
.expect("stream execution should succeed");
let expected_observation = execution.response_observation.clone();
assert!(
expected_observation.response_headers_observed_at_unix_ms
>= expected_observation.request_started_at_unix_ms
);
assert!(!expected_observation.request_order_id.is_empty());
let frames = build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
@@ -1691,6 +1717,10 @@ mod tests {
let header_frame: Value =
serde_json::from_str(&frames[0]).expect("headers frame should parse");
let encoded_observation: aether_contracts::ExecutionResponseObservation =
serde_json::from_value(header_frame["payload"]["response_observation"].clone())
.expect("headers frame should retain the response observation");
assert_eq!(encoded_observation, expected_observation);
assert_eq!(
header_frame
.get("payload")
apps/aether-gateway/src/execution_runtime/sync/execution.rs
+296 -34
View File
@@ -5,8 +5,8 @@ use std::time::{Duration, Instant};
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY};
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
ExecutionTelemetry,
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan,
ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
@@ -34,6 +34,7 @@ use crate::ai_serving::api::{
implicit_sync_finalize_report_kind, maybe_build_sync_finalize_outcome, LocalCoreSyncErrorKind,
LocalCoreSyncFinalizeOutcome,
};
use crate::ai_serving::record_local_runtime_candidate_skip_reason;
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
build_client_response_from_parts_with_mutator,
@@ -55,8 +56,9 @@ use crate::execution_runtime::submission::{
resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::transport::{
append_upstream_response_body_chunk, build_execution_response_body, build_request_body,
collect_response_headers, decode_response_body_bytes, execution_response_body_mode,
append_upstream_response_body_chunk_with_limit, build_execution_response_body,
build_request_body, collect_response_headers, decode_response_body_bytes_with_limit,
execution_plan_response_body_limit_bytes, execution_response_body_mode,
format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error,
response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError,
@@ -70,13 +72,16 @@ use crate::execution_runtime::{
};
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, build_local_error_flow_metadata, trace_upstream_response_body,
with_error_flow_report_context, with_upstream_response_report_context,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
apply_local_execution_effect, build_local_error_flow_metadata,
spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context,
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect,
};
use crate::provider_pool_demand::{
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission,
};
use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
record_local_request_candidate_status, record_local_request_candidate_status_snapshot,
@@ -1379,7 +1384,19 @@ async fn execute_direct_sync_runtime_candidate(
candidate_started_unix_ms,
event.status_code,
event.ttfb_ms,
)
);
spawn_local_oauth_success_effect(
state_for_response_started.clone(),
plan,
report_context,
LocalOAuthSuccessEffect {
status_code: event.status_code,
request_started_at_unix_ms: Some(
event.response_observation.request_started_at_unix_ms,
),
request_order_id: Some(&event.response_observation.request_order_id),
},
);
})
.await
.map_err(SyncExecutionFailure::from_transport);
@@ -1478,17 +1495,31 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
) -> Result<ExecutionResult, SyncExecutionFailure> {
let request_body = build_request_body(plan).map_err(SyncExecutionFailure::from_transport)?;
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
let started_at = Instant::now();
let mut progress =
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
progress.record_connecting().await;
let request_started_at_unix_ms = current_request_candidate_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let response = send_request(plan, request_body)
.await
.map_err(SyncExecutionFailure::from_transport)?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let response_headers_observed_at_unix_ms = current_request_candidate_unix_ms();
let status_code = response.status_code();
let headers = response.headers();
spawn_local_oauth_success_effect(
state.clone(),
plan,
report_context,
LocalOAuthSuccessEffect {
status_code,
request_started_at_unix_ms: Some(request_started_at_unix_ms),
request_order_id: Some(&request_order_id),
},
);
progress.record_response_started(status_code, ttfb_ms).await;
let mut body_bytes = Vec::new();
@@ -1503,8 +1534,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
),
)
})?;
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
.map_err(SyncExecutionFailure::from_transport)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)
.map_err(SyncExecutionFailure::from_transport)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
@@ -1521,8 +1556,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
)),
)
})?;
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
.map_err(SyncExecutionFailure::from_transport)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)
.map_err(SyncExecutionFailure::from_transport)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
@@ -1539,8 +1578,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
),
)
})?;
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
.map_err(SyncExecutionFailure::from_transport)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)
.map_err(SyncExecutionFailure::from_transport)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
@@ -1549,8 +1592,9 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
}
}
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
.map_err(SyncExecutionFailure::from_transport)?;
let decoded_body_bytes =
decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)
.map_err(SyncExecutionFailure::from_transport)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let upstream_bytes = body_bytes.len() as u64;
progress.finish(status_code, elapsed_ms).await;
@@ -1569,6 +1613,11 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: Some(ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
}),
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: Some(ttfb_ms),
@@ -1949,6 +1998,37 @@ async fn execute_execution_runtime_sync_impl(
.unwrap_or_else(|| "-".to_string());
let candidate_started_at = Instant::now();
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let _provider_pool_in_flight_guard = match acquire_provider_pool_execution_guard(state, &plan)
.await?
{
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
record_local_runtime_candidate_skip_reason(
state,
trace_id,
"provider_key_concurrency_limit_reached",
);
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
*retry_scope = AiAttemptRetryScope::Candidate;
}
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Skipped,
status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()),
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
error_message: Some(format!("provider key concurrency limit reached: {limit}")),
latency_ms: Some(0),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(candidate_started_unix_secs),
},
)
.await;
return Ok(None);
}
};
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let usage_data = state.usage_lifecycle_data_state().as_ref().clone();
state
@@ -1978,14 +2058,6 @@ async fn execute_execution_runtime_sync_impl(
candidate_started_at,
);
let result = (async {
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan_request_id.as_str(),
plan_candidate_id.as_deref(),
key_id.as_str(),
)
.await;
record_sync_execution_active(
state,
&plan,
@@ -2461,6 +2533,16 @@ async fn execute_execution_runtime_sync_impl(
};
let mut candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result);
let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms();
let mut provider_response_observation =
result
.response_observation
.clone()
.unwrap_or(ExecutionResponseObservation {
request_started_at_unix_ms: candidate_started_unix_secs,
response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms,
request_order_id: uuid::Uuid::now_v7().to_string(),
});
let mut oauth_retry_attempted = false;
let (
result_error_type,
@@ -2473,6 +2555,18 @@ async fn execute_execution_runtime_sync_impl(
local_failover_response_text,
local_failover_analysis,
) = loop {
spawn_local_oauth_success_effect(
state.clone(),
&plan,
report_context.as_ref(),
LocalOAuthSuccessEffect {
status_code: result.status_code,
request_started_at_unix_ms: Some(
provider_response_observation.request_started_at_unix_ms,
),
request_order_id: Some(&provider_response_observation.request_order_id),
},
);
let result_latency_ms = result
.telemetry
.as_ref()
@@ -2534,10 +2628,15 @@ async fn execute_execution_runtime_sync_impl(
result.status_code,
local_failover_response_text.as_deref(),
trace_id,
report_context.as_ref(),
Some(provider_response_observation.request_started_at_unix_ms),
Some(&provider_response_observation.request_order_id),
)
.await
{
oauth_retry_attempted = true;
let retry_started_at_unix_ms = current_request_candidate_unix_ms();
let retry_request_order_id = uuid::Uuid::now_v7().to_string();
match crate::execution_runtime::execute_execution_runtime_sync_plan(
state,
Some(trace_id),
@@ -2546,6 +2645,16 @@ async fn execute_execution_runtime_sync_impl(
.await
{
Ok(retry_result) => {
let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms();
provider_response_observation = retry_result
.response_observation
.clone()
.unwrap_or(ExecutionResponseObservation {
request_started_at_unix_ms: retry_started_at_unix_ms,
response_headers_observed_at_unix_ms:
retry_response_observed_at_unix_ms,
request_order_id: retry_request_order_id,
});
candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms(
candidate_started_at,
@@ -2594,6 +2703,13 @@ async fn execute_execution_runtime_sync_impl(
local_failover_analysis,
);
};
let mut report_context = attach_provider_response_headers_to_report_context(
report_context,
&headers,
provider_response_observation.request_started_at_unix_ms,
provider_response_observation.response_headers_observed_at_unix_ms,
&provider_response_observation.request_order_id,
);
if result.status_code >= 400 {
apply_local_execution_effect(
state,
@@ -2739,8 +2855,6 @@ async fn execute_execution_runtime_sync_impl(
}
let status_code = result.status_code;
let has_body_bytes = body_base64.is_some();
let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers);
if (200..300).contains(&status_code) {
seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
@@ -3191,7 +3305,14 @@ fn maybe_build_implicit_sync_finalize_outcome(
body_base64: &Option<String>,
telemetry: &Option<ExecutionTelemetry>,
) -> Result<Option<ImplicitSyncFinalizeOutcome>, GatewayError> {
if status_code >= 400 || body_json.is_some() || body_base64.is_none() {
let needs_conversion = report_context
.as_ref()
.and_then(|value| value.get("needs_conversion"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let has_captured_stream_body = body_json.is_none() && body_base64.is_some();
let has_cross_format_sync_body = needs_conversion && body_json.is_some();
if status_code >= 400 || (!has_captured_stream_body && !has_cross_format_sync_body) {
return Ok(None);
}
@@ -3231,6 +3352,8 @@ async fn execute_sync_via_remote_execution_runtime(
candidate_started_unix_secs: u64,
candidate_started_at: Instant,
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
let remote_request_started_at_unix_ms = current_request_candidate_unix_ms();
let remote_request_order_id = uuid::Uuid::now_v7().to_string();
let response = match post_sync_plan_to_remote_execution_runtime(
state,
remote_execution_runtime_base_url,
@@ -3299,11 +3422,19 @@ async fn execute_sync_via_remote_execution_runtime(
));
}
response
.json()
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
let mut result = response
.json::<ExecutionResult>()
.await
.map(RemoteSyncFallbackOutcome::Executed)
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
result
.response_observation
.get_or_insert(ExecutionResponseObservation {
request_started_at_unix_ms: remote_request_started_at_unix_ms,
response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms,
request_order_id: remote_request_order_id,
});
Ok(RemoteSyncFallbackOutcome::Executed(result))
}
#[cfg(test)]
@@ -3363,6 +3494,137 @@ mod tests {
.with_execution_runtime_candidate(true)
}
#[tokio::test]
async fn implicit_sync_finalize_converts_chat_json_to_namespaced_responses() {
let decision = GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
)
.with_execution_runtime_candidate(true);
let report_context = Some(json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:responses",
"needs_conversion": true,
"mapped_model": "qwen-upstream",
"original_request_body": {
"model": "qwen",
"tools": [{
"type": "namespace",
"name": "mcp__vulnerability_report",
"description": "reporting tools",
"tools": [{
"type": "function",
"name": "vulnerability_report",
"description": "write the confirmed report",
"parameters": {
"type": "object",
"properties": {
"report_path": {"type": "string"}
},
"required": ["report_path"]
},
"strict": true
}]
}]
}
}));
let provider_body = Some(json!({
"id": "chatcmpl_namespace_sync",
"object": "chat.completion",
"created": 1_777_777_777,
"model": "qwen-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": null,
"tool_calls": [{
"id": "call_report_1",
"type": "function",
"function": {
"name": "vulnerability_report",
"arguments": "{\"report_path\":\"reports/sql-001-c1.md\"}"
}
}]
},
"finish_reason": "tool_calls"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 4,
"total_tokens": 14
}
}));
let implicit = maybe_build_implicit_sync_finalize_outcome(
"trace-namespace-sync",
&decision,
"openai_responses_sync",
&report_context,
StatusCode::OK.as_u16(),
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
&provider_body,
&None,
&None,
)
.expect("cross-format sync JSON finalize should not error")
.expect("cross-format sync JSON should be finalized");
let response_body = axum::body::to_bytes(implicit.outcome.response.into_body(), usize::MAX)
.await
.expect("response body should read");
let response_json: Value =
serde_json::from_slice(&response_body).expect("response body should be JSON");
assert_eq!(response_json["object"], "response");
assert!(response_json.get("choices").is_none());
assert_eq!(response_json["output"][0]["type"], "function_call");
assert_eq!(response_json["output"][0]["name"], "vulnerability_report");
assert_eq!(
response_json["output"][0]["namespace"],
"mcp__vulnerability_report"
);
assert_eq!(response_json["output"][0]["call_id"], "call_report_1");
}
#[test]
fn implicit_sync_finalize_leaves_same_format_json_on_passthrough_path() {
let report_context = Some(json!({
"provider_api_format": "openai:responses",
"client_api_format": "openai:responses",
"needs_conversion": false
}));
let body_json = Some(json!({
"id": "resp_same_format",
"object": "response",
"status": "completed",
"output": []
}));
let outcome = maybe_build_implicit_sync_finalize_outcome(
"trace-same-format-sync",
&GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
),
"openai_responses_sync",
&report_context,
StatusCode::OK.as_u16(),
&BTreeMap::new(),
&body_json,
&None,
&None,
)
.expect("same-format sync JSON guard should not error");
assert!(outcome.is_none());
}
#[tokio::test]
async fn oversized_upstream_response_builds_claude_502_retry_fallback() {
let mut plan = test_openai_image_plan(false);
apps/aether-gateway/src/execution_runtime/transport.rs
+522 -33
View File
@@ -9,18 +9,19 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex, OnceLock, RwLock as StdRwLock}
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, ExecutionTelemetry, ProxySnapshot,
ResolvedTransportProfile, ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE,
TRANSPORT_HTTP_MODE_HTTP1_ONLY,
ExecutionPlan, ExecutionResponseBodyMode, ExecutionResponseObservation, ExecutionResult,
ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile, ResponseBody,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, EXECUTION_RESPONSE_BODY_MODE_HEADER,
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS,
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
};
use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation;
use aether_http::{apply_http_client_config, HttpClientConfig};
use aether_runtime::{MetricKind, MetricSample};
use axum::body::Bytes;
use base64::Engine as _;
use brotli::Decompressor as BrotliDecoder;
use flate2::read::{DeflateDecoder, GzDecoder};
use flate2::write::GzEncoder;
use flate2::Compression;
@@ -62,6 +63,10 @@ const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000;
const DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS: u64 = 1_200_000;
const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1;
const EXECUTION_RESPONSE_BODY_LIMIT_HEADER: &str = "x-aether-execution-response-body-limit-bytes";
const DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024;
const MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024;
const MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024;
const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS";
const DIRECT_REQWEST_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS";
const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str =
@@ -607,6 +612,56 @@ impl std::fmt::Display for UpstreamResponseBodyPhase {
}
}
pub(crate) fn with_upstream_response_body_limit(
plan: &ExecutionPlan,
limit_bytes: usize,
) -> ExecutionPlan {
let mut bounded_plan = plan.clone();
bounded_plan
.headers
.retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER));
bounded_plan.headers.insert(
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(),
normalize_scoped_response_body_limit(limit_bytes)
.unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
.to_string(),
);
bounded_plan
}
pub(crate) fn execution_plan_response_body_limit_bytes(plan: &ExecutionPlan) -> usize {
effective_response_body_limit_bytes(
execution_transport_header_value(&plan.headers, EXECUTION_RESPONSE_BODY_LIMIT_HEADER),
crate::headers::max_internal_buffered_body_bytes(),
)
}
fn effective_response_body_limit_bytes(
raw_scoped_limit: Option<&str>,
global_limit: usize,
) -> usize {
let Some(raw_scoped_limit) = raw_scoped_limit else {
return global_limit;
};
parse_scoped_response_body_limit(raw_scoped_limit)
.unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
.min(global_limit)
}
fn parse_scoped_response_body_limit(value: &str) -> Option<usize> {
let raw_limit = value.trim().parse::<u64>().ok()?;
usize::try_from(raw_limit)
.ok()
.and_then(normalize_scoped_response_body_limit)
}
fn normalize_scoped_response_body_limit(limit_bytes: usize) -> Option<usize> {
(limit_bytes > 0).then_some(limit_bytes.clamp(
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
))
}
pub(crate) fn append_upstream_response_body_chunk(
body: &mut Vec<u8>,
chunk: &[u8],
@@ -618,7 +673,7 @@ pub(crate) fn append_upstream_response_body_chunk(
)
}
fn append_upstream_response_body_chunk_with_limit(
pub(crate) fn append_upstream_response_body_chunk_with_limit(
body: &mut Vec<u8>,
chunk: &[u8],
limit_bytes: usize,
@@ -691,6 +746,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) stream_precommit_committed: bool,
pub(crate) response: DirectUpstreamResponse,
pub(crate) started_at: Instant,
pub(crate) response_observation: ExecutionResponseObservation,
pub(crate) stream_first_byte_timeout: Option<Duration>,
pub(crate) upstream_target_permit: Option<UpstreamTargetAdmissionPermit>,
}
@@ -699,6 +755,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) struct DirectSyncResponseStarted {
pub(crate) status_code: u16,
pub(crate) ttfb_ms: u64,
pub(crate) response_observation: ExecutionResponseObservation,
}
impl DirectSyncExecutionRuntime {
@@ -722,20 +779,35 @@ impl DirectSyncExecutionRuntime {
F: FnOnce(DirectSyncResponseStarted),
{
let body_bytes = build_request_body(plan)?;
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
with_non_stream_total_timeout(plan, async move {
let response = send_request_inner(plan, body_bytes, false).await?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let status_code = response.status_code();
let headers = response.headers();
let response_observation = ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
};
on_response_started(DirectSyncResponseStarted {
status_code,
ttfb_ms,
response_observation: response_observation.clone(),
});
let (body_bytes, stream_ttfb_ms) =
response.bytes_with_stream_timeout(plan, started_at).await?;
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
let (body_bytes, stream_ttfb_ms) = response
.bytes_with_stream_timeout(plan, started_at, response_body_limit_bytes)
.await?;
let decoded_body_bytes = decode_response_body_bytes_with_limit(
&headers,
&body_bytes,
response_body_limit_bytes,
)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let upstream_bytes = body_bytes.len() as u64;
@@ -752,6 +824,7 @@ impl DirectSyncExecutionRuntime {
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: Some(response_observation),
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)),
@@ -776,6 +849,8 @@ impl DirectSyncExecutionRuntime {
);
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let response = send_request(plan, body_bytes).await?;
observe_gateway_stage_ms(
"direct_send_headers",
@@ -783,6 +858,7 @@ impl DirectSyncExecutionRuntime {
);
let status_code = response.status_code();
let headers = response.headers();
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let stream_summary_report_context = build_stream_summary_report_context(plan);
@@ -797,6 +873,11 @@ impl DirectSyncExecutionRuntime {
stream_precommit_committed: false,
response: response.into_direct_upstream_response(),
started_at,
response_observation: ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
upstream_target_permit: None,
})
@@ -834,7 +915,7 @@ pub(crate) async fn execute_sync_plan_with_report_context(
}
if resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).is_some() {
return execute_sync_plan_via_local_tunnel(state, plan)
return execute_sync_plan_via_local_tunnel(state, plan, report_context)
.await
.map_err(|err| GatewayError::Internal(err.to_string()));
}
@@ -857,7 +938,24 @@ pub(crate) async fn execute_sync_plan_with_report_context(
Ok(None) => {}
Err(err) => return Err(GatewayError::Internal(err.to_string())),
}
match DirectSyncExecutionRuntime::new().execute_sync(plan).await {
let state_for_response_started = state.clone();
match DirectSyncExecutionRuntime::new()
.execute_sync_with_response_started(plan, move |event| {
crate::orchestration::spawn_local_oauth_success_effect(
state_for_response_started,
plan,
report_context,
crate::orchestration::LocalOAuthSuccessEffect {
status_code: event.status_code,
request_started_at_unix_ms: Some(
event.response_observation.request_started_at_unix_ms,
),
request_order_id: Some(&event.response_observation.request_order_id),
},
);
})
.await
{
Ok(result) => {
record_manual_proxy_request_outcome(state, plan, result.status_code).await;
Ok(result)
@@ -889,6 +987,8 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
plan.body.body_bytes_b64.is_some(),
)?;
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let response = state
.tunnel
.open_direct_relay_stream(
@@ -900,6 +1000,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
Ok(Some(DirectUpstreamStreamExecution {
request_id: plan.request_id.clone(),
@@ -912,6 +1013,11 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
stream_precommit_committed: false,
response: DirectUpstreamResponse::LocalTunnel(response),
started_at,
response_observation: ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
upstream_target_permit: None,
}))
@@ -991,13 +1097,19 @@ fn manual_proxy_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
async fn execute_sync_plan_via_local_tunnel(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
with_non_stream_total_timeout(plan, execute_sync_plan_via_local_tunnel_inner(state, plan)).await
with_non_stream_total_timeout(
plan,
execute_sync_plan_via_local_tunnel_inner(state, plan, report_context),
)
.await
}
async fn execute_sync_plan_via_local_tunnel_inner(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
let node_id = resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).ok_or_else(|| {
ExecutionRuntimeTransportError::RelayError("local tunnel node unavailable".to_string())
@@ -1007,6 +1119,7 @@ async fn execute_sync_plan_via_local_tunnel_inner(
}
let body_bytes = build_request_body(plan)?;
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
let transport_controls = resolve_execution_transport_controls(&plan.headers);
let headers = build_request_headers(
&plan.headers,
@@ -1030,6 +1143,8 @@ async fn execute_sync_plan_via_local_tunnel_inner(
"gateway execution runtime local tunnel request prepared"
);
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let mut response = state
.tunnel
.open_direct_relay_stream(
@@ -1040,12 +1155,30 @@ async fn execute_sync_plan_via_local_tunnel_inner(
.await
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
let response_observation = ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
};
crate::orchestration::spawn_local_oauth_success_effect(
state.clone(),
plan,
report_context,
crate::orchestration::LocalOAuthSuccessEffect {
status_code,
request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms),
request_order_id: Some(&response_observation.request_order_id),
},
);
let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-");
let (body_bytes, stream_ttfb_ms) =
collect_local_tunnel_response_body(response, plan, started_at).await?;
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
collect_local_tunnel_response_body(response, plan, started_at, response_body_limit_bytes)
.await?;
let decoded_body_bytes =
decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let upstream_bytes = body_bytes.len() as u64;
if status_code >= 400 {
@@ -1095,6 +1228,7 @@ async fn execute_sync_plan_via_local_tunnel_inner(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: Some(response_observation),
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)),
@@ -1109,6 +1243,7 @@ async fn collect_local_tunnel_response_body(
mut response: tunnel::DirectRelayResponse,
plan: &ExecutionPlan,
started_at: Instant,
response_body_limit_bytes: usize,
) -> Result<(Vec<u8>, Option<u64>), ExecutionRuntimeTransportError> {
let mut body_bytes = Vec::new();
let mut first_byte_ms = None;
@@ -1131,7 +1266,11 @@ async fn collect_local_tunnel_response_body(
if plan.stream && first_byte_ms.is_none() && !chunk.is_empty() {
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)?;
}
Ok((body_bytes, first_byte_ms))
@@ -1292,20 +1431,28 @@ impl DirectHttpResponse {
}
pub(crate) async fn bytes(self) -> Result<Bytes, ExecutionRuntimeTransportError> {
self.bytes_with_limit(crate::headers::max_internal_buffered_body_bytes())
.await
}
async fn bytes_with_limit(
self,
response_body_limit_bytes: usize,
) -> Result<Bytes, ExecutionRuntimeTransportError> {
let started_at = Instant::now();
match self {
DirectHttpResponse::Reqwest(response) => {
collect_reqwest_stream_body(response, started_at, None)
collect_reqwest_stream_body(response, started_at, None, response_body_limit_bytes)
.await
.map(|(body, _)| body)
}
DirectHttpResponse::HyperH2c(response) => {
collect_hyper_stream_body(response, started_at, None)
collect_hyper_stream_body(response, started_at, None, response_body_limit_bytes)
.await
.map(|(body, _)| body)
}
DirectHttpResponse::BrowserWreq(response) => {
collect_wreq_stream_body(response, started_at, None)
collect_wreq_stream_body(response, started_at, None, response_body_limit_bytes)
.await
.map(|(body, _)| body)
}
@@ -1316,21 +1463,43 @@ impl DirectHttpResponse {
self,
plan: &ExecutionPlan,
started_at: Instant,
response_body_limit_bytes: usize,
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
if !plan.stream {
return self.bytes().await.map(|bytes| (bytes, None));
return self
.bytes_with_limit(response_body_limit_bytes)
.await
.map(|bytes| (bytes, None));
}
let first_byte_timeout = resolve_stream_first_byte_timeout(plan);
match self {
DirectHttpResponse::Reqwest(response) => {
collect_reqwest_stream_body(response, started_at, first_byte_timeout).await
collect_reqwest_stream_body(
response,
started_at,
first_byte_timeout,
response_body_limit_bytes,
)
.await
}
DirectHttpResponse::HyperH2c(response) => {
collect_hyper_stream_body(response, started_at, first_byte_timeout).await
collect_hyper_stream_body(
response,
started_at,
first_byte_timeout,
response_body_limit_bytes,
)
.await
}
DirectHttpResponse::BrowserWreq(response) => {
collect_wreq_stream_body(response, started_at, first_byte_timeout).await
collect_wreq_stream_body(
response,
started_at,
first_byte_timeout,
response_body_limit_bytes,
)
.await
}
}
}
@@ -1376,6 +1545,7 @@ async fn collect_reqwest_stream_body(
response: reqwest::Response,
started_at: Instant,
first_byte_timeout: Option<Duration>,
response_body_limit_bytes: usize,
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
let mut stream = response.bytes_stream();
let mut body_bytes = Vec::new();
@@ -1396,7 +1566,11 @@ async fn collect_reqwest_stream_body(
if first_byte_ms.is_none() && !chunk.is_empty() {
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)?;
}
Ok((Bytes::from(body_bytes), first_byte_ms))
@@ -1406,6 +1580,7 @@ async fn collect_hyper_stream_body(
response: hyper::Response<HyperIncomingBody>,
started_at: Instant,
first_byte_timeout: Option<Duration>,
response_body_limit_bytes: usize,
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
let mut stream = response.into_body().into_data_stream();
let mut body_bytes = Vec::new();
@@ -1426,7 +1601,11 @@ async fn collect_hyper_stream_body(
if first_byte_ms.is_none() && !chunk.is_empty() {
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)?;
}
Ok((Bytes::from(body_bytes), first_byte_ms))
@@ -1436,6 +1615,7 @@ async fn collect_wreq_stream_body(
response: wreq::Response,
started_at: Instant,
first_byte_timeout: Option<Duration>,
response_body_limit_bytes: usize,
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
let mut stream = response.bytes_stream();
let mut body_bytes = Vec::new();
@@ -1456,7 +1636,11 @@ async fn collect_wreq_stream_body(
if first_byte_ms.is_none() && !chunk.is_empty() {
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
append_upstream_response_body_chunk_with_limit(
&mut body_bytes,
&chunk,
response_body_limit_bytes,
)?;
}
Ok((Bytes::from(body_bytes), first_byte_ms))
@@ -2308,10 +2492,31 @@ async fn send_via_tunnel_relay(
error_kind = %kind,
"gateway execution runtime tunnel relay returned relay error"
);
let message = response
.text()
.await
.unwrap_or_else(|_| format!("hub relay error: {kind}"));
let response_headers = collect_response_headers(response.headers());
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
let (wire_body, _) =
collect_reqwest_stream_body(response, Instant::now(), None, response_body_limit_bytes)
.await
.map_err(|error| {
ExecutionRuntimeTransportError::RelayError(format!(
"hub relay error: {kind}: bounded error body read failed: {error}"
))
})?;
let decoded_body = decode_response_body_bytes_with_limit(
&response_headers,
&wire_body,
response_body_limit_bytes,
)
.map_err(|error| {
ExecutionRuntimeTransportError::RelayError(format!(
"hub relay error: {kind}: bounded error body decode failed: {error}"
))
})?;
let message = if decoded_body.is_empty() {
format!("hub relay error: {kind}")
} else {
String::from_utf8_lossy(decoded_body.as_ref()).into_owned()
};
return Err(ExecutionRuntimeTransportError::RelayError(message));
}
@@ -3833,6 +4038,7 @@ pub(crate) fn build_request_headers(
|| normalized_key == EXECUTION_REQUEST_HTTP1_ONLY_HEADER
|| normalized_key == EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER
|| normalized_key == EXECUTION_RESPONSE_BODY_MODE_HEADER
|| normalized_key == EXECUTION_RESPONSE_BODY_LIMIT_HEADER
{
continue;
}
@@ -3981,7 +4187,7 @@ pub(crate) fn decode_response_body_bytes<'a>(
)
}
fn decode_response_body_bytes_with_limit<'a>(
pub(crate) fn decode_response_body_bytes_with_limit<'a>(
headers: &BTreeMap<String, String>,
body_bytes: &'a [u8],
limit_bytes: usize,
@@ -4003,6 +4209,11 @@ fn decode_response_body_bytes_with_limit<'a>(
read_upstream_response_decoder_with_limit("deflate", &mut decoder, limit_bytes)
.map(Cow::Owned)
}
Some("br") => {
let mut decoder = BrotliDecoder::new(body_bytes, 4_096);
read_upstream_response_decoder_with_limit("br", &mut decoder, limit_bytes)
.map(Cow::Owned)
}
_ => Ok(Cow::Borrowed(body_bytes)),
}
}
@@ -4122,12 +4333,16 @@ mod tests {
use super::{
append_upstream_response_body_chunk_with_limit, build_browser_wreq_client, build_client,
build_direct_tunnel_request_meta, build_execution_response_body, build_request_headers,
decode_response_body_bytes_with_limit, execute_sync_plan, execution_response_body_mode,
decode_response_body_bytes_with_limit, effective_response_body_limit_bytes,
execute_sync_plan, execution_plan_response_body_limit_bytes, execution_response_body_mode,
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
record_manual_proxy_request_success, record_manual_proxy_stream_error,
resolve_execution_transport_controls, resolve_non_stream_total_timeout,
resolve_stream_first_byte_timeout, response_body_is_json, DirectSyncExecutionRuntime,
resolve_stream_first_byte_timeout, response_body_is_json,
with_upstream_response_body_limit, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError, ExecutionTransportControls, UpstreamResponseBodyPhase,
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES, EXECUTION_RESPONSE_BODY_LIMIT_HEADER,
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES, MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
};
use crate::constants::{
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
@@ -4182,6 +4397,162 @@ mod tests {
assert!(!materialized.contains_key("x-aether-future-control"));
}
#[test]
fn scoped_response_body_limit_injection_preserves_transport_profile_and_extra() {
let mut plan = tunnel_timeout_plan(false);
let original_profile = ResolvedTransportProfile {
profile_id: "existing-profile".into(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(),
http_mode: TRANSPORT_HTTP_MODE_HTTP1_ONLY.into(),
pool_scope: "provider".into(),
header_fingerprint: Some(json!({"user_agent": "existing"})),
extra: Some(json!({"existing": {"nested": true}})),
};
plan.transport_profile = Some(original_profile.clone());
let bounded_plan =
with_upstream_response_body_limit(&plan, DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES);
assert_eq!(plan.transport_profile, Some(original_profile.clone()));
assert_eq!(bounded_plan.transport_profile, Some(original_profile));
assert_eq!(
bounded_plan
.headers
.get(EXECUTION_RESPONSE_BODY_LIMIT_HEADER)
.and_then(|value| value.parse::<usize>().ok()),
Some(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
);
assert_eq!(
execution_plan_response_body_limit_bytes(&bounded_plan),
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
);
let unprofiled_plan = tunnel_timeout_plan(false);
let bounded_unprofiled_plan = with_upstream_response_body_limit(
&unprofiled_plan,
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
);
assert!(unprofiled_plan.transport_profile.is_none());
assert!(bounded_unprofiled_plan.transport_profile.is_none());
assert_eq!(
execution_plan_response_body_limit_bytes(&bounded_unprofiled_plan),
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
);
let mut shadowed_plan = tunnel_timeout_plan(false);
shadowed_plan.headers.insert(
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_ascii_uppercase(),
"65536".to_string(),
);
let bounded_shadowed_plan = with_upstream_response_body_limit(
&shadowed_plan,
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
);
assert_eq!(
bounded_shadowed_plan
.headers
.keys()
.filter(|name| name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER))
.count(),
1
);
}
#[test]
fn scoped_response_body_limit_parsing_rejects_invalid_values_and_clamps_bounds() {
let scoped_plan = |raw_limit: &str| {
let mut plan = tunnel_timeout_plan(false);
plan.headers.insert(
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(),
raw_limit.to_string(),
);
plan
};
for invalid in ["0", "-1", "1.5", "", "invalid"] {
assert_eq!(
execution_plan_response_body_limit_bytes(&scoped_plan(invalid)),
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
);
}
assert_eq!(
execution_plan_response_body_limit_bytes(&scoped_plan("1")),
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES
);
assert_eq!(
execution_plan_response_body_limit_bytes(&scoped_plan(
&(MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES as u64 + 1).to_string()
)),
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES
);
assert_eq!(
execution_plan_response_body_limit_bytes(&scoped_plan("1048576")),
1_048_576
);
assert_eq!(
effective_response_body_limit_bytes(
Some(&(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES * 2).to_string()),
1024 * 1024,
),
1024 * 1024,
"a scoped limit must never raise the operator's global cap"
);
}
#[test]
fn scoped_response_body_wire_limit_rejects_overflow() {
let bounded_plan = with_upstream_response_body_limit(
&tunnel_timeout_plan(false),
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
);
let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan);
let mut body = vec![b'x'; limit_bytes];
let error =
append_upstream_response_body_chunk_with_limit(&mut body, b"overflow", limit_bytes)
.expect_err("wire body above the plan-scoped limit should fail");
assert!(matches!(
error,
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
phase: UpstreamResponseBodyPhase::Wire,
limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
}
));
}
#[test]
fn scoped_response_body_limit_rejects_gzip_bomb_after_wire_check() {
let bounded_plan = with_upstream_response_body_limit(
&tunnel_timeout_plan(false),
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
);
let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan);
let payload = vec![b'x'; limit_bytes + 1];
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
encoder
.write_all(&payload)
.expect("gzip payload should encode");
let encoded = encoder.finish().expect("gzip payload should finish");
assert!(encoded.len() < limit_bytes);
let mut wire_body = Vec::new();
append_upstream_response_body_chunk_with_limit(&mut wire_body, &encoded, limit_bytes)
.expect("compressed wire body should fit within the plan-scoped limit");
let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]);
let error = decode_response_body_bytes_with_limit(&headers, &wire_body, limit_bytes)
.expect_err("decoded body above the plan-scoped limit should fail");
assert!(matches!(
error,
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
phase: UpstreamResponseBodyPhase::Decoded,
limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
}
));
}
#[test]
fn upstream_response_wire_limit_allows_exact_body_and_rejects_next_byte() {
let mut body = Vec::new();
@@ -5484,6 +5855,100 @@ mod tests {
);
}
#[tokio::test]
async fn direct_sync_execution_runtime_preserves_gemini_tool_config_on_wire() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let captured_body = Arc::new(Mutex::new(None));
let captured_body_for_handler = Arc::clone(&captured_body);
let app = Router::new().route(
"/generate",
post(move |body: Bytes| {
let captured_body = Arc::clone(&captured_body_for_handler);
async move {
*captured_body
.lock()
.expect("capture lock should not be poisoned") = Some(body.to_vec());
Json(json!({"ok": true}))
}
}),
);
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("test server should run");
});
let result = DirectSyncExecutionRuntime::new()
.execute_sync(&ExecutionPlan {
request_id: "req-gemini-tool-config-wire".into(),
candidate_id: Some("cand-gemini-tool-config-wire".into()),
provider_name: Some("google".into()),
provider_id: "prov-gemini-tool-config-wire".into(),
endpoint_id: "ep-gemini-tool-config-wire".into(),
key_id: "key-gemini-tool-config-wire".into(),
method: "POST".into(),
url: format!("http://{addr}/generate"),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gemini-3-flash-preview",
"contents": [{
"role": "user",
"parts": [{"text": "Search, then save the result."}]
}],
"tools": [
{"googleSearch": {}},
{"functionDeclarations": [{
"name": "save_result",
"parameters": {
"type": "OBJECT",
"properties": {"result": {"type": "STRING"}}
}
}]}
],
"toolConfig": {
"includeServerSideToolInvocations": true,
"functionCallingConfig": {"mode": "ANY"}
}
})),
stream: false,
client_api_format: "openai:responses".into(),
provider_api_format: "gemini:generate_content".into(),
model_name: Some("gemini-3-flash-preview".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
..ExecutionTimeouts::default()
}),
})
.await
.expect("sync execution should succeed");
server.abort();
assert_eq!(result.status_code, 200);
let body = captured_body
.lock()
.expect("capture lock should not be poisoned")
.take()
.and_then(|body| serde_json::from_slice::<serde_json::Value>(&body).ok())
.expect("upstream should receive a JSON body");
assert_eq!(
body["toolConfig"]["includeServerSideToolInvocations"],
json!(true)
);
assert_eq!(body["toolConfig"]["functionCallingConfig"]["mode"], "ANY");
assert!(body["toolConfig"]
.get("include_server_side_tool_invocations")
.is_none());
}
#[tokio::test]
async fn direct_sync_execution_runtime_applies_non_stream_total_timeout_to_body() {
let listener = crate::test_support::bind_loopback_listener()
@@ -5605,6 +6070,8 @@ mod tests {
)
.await
.expect("headers should write");
socket.flush().await.expect("headers should flush");
tokio::time::sleep(std::time::Duration::from_millis(40)).await;
socket
.write_all(b"b\r\ndata: one\n\n\r\n")
.await
@@ -5634,12 +6101,34 @@ mod tests {
let body = result
.body
.clone()
.and_then(|body| body.body_bytes_b64)
.and_then(|body| base64::engine::general_purpose::STANDARD.decode(body).ok())
.expect("stream body should be captured as bytes");
let body = String::from_utf8(body).expect("stream body should be utf8");
assert!(body.contains("data: one"));
assert!(body.contains("data: two"));
let observation = result
.response_observation
.expect("stream sync execution should preserve header observation");
let telemetry = result
.telemetry
.expect("stream sync execution should include telemetry");
let ttfb_ms = telemetry
.ttfb_ms
.expect("stream sync execution should measure the first body byte");
assert!(
observation.response_headers_observed_at_unix_ms
>= observation.request_started_at_unix_ms
);
assert!(
observation
.response_headers_observed_at_unix_ms
.saturating_sub(observation.request_started_at_unix_ms)
< ttfb_ms,
"header observation must not be derived from body-byte ttfb"
);
assert!(!observation.request_order_id.is_empty());
}
#[tokio::test]
apps/aether-gateway/src/execution_runtime/windsurf.rs
+2
View File
@@ -266,6 +266,7 @@ pub(crate) async fn maybe_execute_windsurf_sync(
candidate_id: prepared.candidate_id,
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
@@ -527,6 +528,7 @@ fn build_windsurf_stream_frame_stream(
("cache-control".to_string(), "no-cache".to_string()),
("content-type".to_string(), "text/event-stream".to_string()),
]),
response_observation: None,
},
});
apps/aether-gateway/src/executor/candidate_loop.rs
+36 -52
View File
@@ -20,9 +20,11 @@ use crate::ai_serving::LocalExecutionAttemptSource;
use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{
build_transport_error_stop_response, execute_execution_runtime_stream_with_retry_scope,
acquire_upstream_execution_gate, build_transport_error_stop_response,
execute_execution_runtime_stream_with_retry_scope,
execute_execution_runtime_sync_with_retry_scope,
mark_stream_candidate_watchdog_terminal_started, StreamCandidateWatchdogProgress,
UpstreamExecutionGateProvider, UPSTREAM_EXECUTION_GATE_NAME,
};
use crate::executor::{
build_local_execution_exhaustion, mark_deferred_upstream_response, LocalExecutionRequestOutcome,
@@ -43,7 +45,6 @@ use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
const DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS: u64 = 30_000;
const UPSTREAM_EXECUTION_GATE_NAME: &str = "gateway_upstream_execution";
const UPSTREAM_TARGET_GATE_NAME: &str = "gateway_upstream_target";
const UPSTREAM_EXECUTION_GATE_HOLD_STREAM_RESPONSE_ENV: &str =
"AETHER_GATEWAY_UPSTREAM_EXECUTION_GATE_HOLD_STREAM_RESPONSE";
@@ -251,6 +252,10 @@ where
Ok(())
}
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
}
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
record_provider_transfer_attempt_failed(
self.state,
@@ -787,18 +792,31 @@ where
{
let mut last_attempted = None;
let mut fallback_response = None;
// A same-key retry derived after a candidate-scoped failure runs before
// the source is asked for the next candidate.
let mut pending_same_key_retry: Option<Attempt> = None;
loop {
let next_started_at = std::time::Instant::now();
let next_attempt =
next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout)
let attempt = match pending_same_key_retry.take() {
Some(attempt) => attempt,
None => {
let next_started_at = std::time::Instant::now();
let next_attempt = next_execution_attempt_with_timeout(
source,
trace_id,
plan_kind,
planning_timeout,
)
.await?;
observe_gateway_stage_ms(
"stream_candidate_next",
next_started_at.elapsed().as_millis() as u64,
);
let Some(attempt) = next_attempt else {
break;
observe_gateway_stage_ms(
"stream_candidate_next",
next_started_at.elapsed().as_millis() as u64,
);
let Some(attempt) = next_attempt else {
break;
};
attempt
}
};
if port.should_skip_attempt(&attempt).await? {
let provider_id = attempt.execution_plan().provider_id.clone();
@@ -838,6 +856,9 @@ where
if attempt_fallback_response.is_some() {
fallback_response = attempt_fallback_response;
}
if scope == AiAttemptRetryScope::Candidate {
pending_same_key_retry = port.next_same_key_retry(&attempt).await?;
}
apply_attempt_retry_scope(source, &attempt, scope).await?;
}
}
@@ -951,6 +972,10 @@ where
Ok(())
}
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
}
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
record_provider_transfer_attempt_failed(
self.state,
@@ -1612,47 +1637,6 @@ fn hold_response_upstream_execution_permit(
Response::from_parts(parts, Body::from_stream(stream))
}
trait UpstreamExecutionGateProvider {
fn upstream_execution_gate(&self) -> Option<&aether_runtime::ConcurrencyGate>;
fn upstream_execution_gate_queue_budget(&self) -> Duration;
}
impl UpstreamExecutionGateProvider for AppState {
fn upstream_execution_gate(&self) -> Option<&aether_runtime::ConcurrencyGate> {
self.upstream_execution_gate.as_deref()
}
fn upstream_execution_gate_queue_budget(&self) -> Duration {
self.frontdoor_runtime_guards.internal_gate_queue_budget
}
}
async fn acquire_upstream_execution_gate(
state: &(impl UpstreamExecutionGateProvider + ?Sized),
trace_id: &str,
) -> Result<Option<ConcurrencyPermit>, GatewayError> {
let Some(gate) = state.upstream_execution_gate() else {
return Ok(None);
};
let budget = state.upstream_execution_gate_queue_budget();
let gate_wait_started_at = std::time::Instant::now();
match timeout(budget, gate.acquire()).await {
Ok(Ok(permit)) => {
observe_gateway_stage_ms(
"upstream_execution_gate_wait",
gate_wait_started_at.elapsed().as_millis() as u64,
);
Ok(Some(permit))
}
Ok(Err(err)) => Err(GatewayError::Internal(err.to_string())),
Err(_) => Err(GatewayError::AdmissionTimeout {
trace_id: trace_id.to_string(),
gate: UPSTREAM_EXECUTION_GATE_NAME,
queue_budget_ms: budget.as_millis() as u64,
}),
}
}
pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
state: &AppState,
remaining: Vec<T>,
apps/aether-gateway/src/executor/orchestration.rs
+106
View File
@@ -1663,6 +1663,22 @@ mod tests {
candidate_index: u32,
endpoint_id: &str,
candidate_id: &str,
) -> AiSyncAttempt {
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
candidate_index,
endpoint_id,
candidate_id,
1,
)
}
/// `sticky_key_attempts` is pinned so these tests exercise candidate
/// failover; the default same-key retry is covered separately.
fn test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
candidate_index: u32,
endpoint_id: &str,
candidate_id: &str,
sticky_key_attempts: u32,
) -> AiSyncAttempt {
AiSyncAttempt {
plan: test_openai_image_heartbeat_plan(endpoint_id, candidate_id),
@@ -1670,6 +1686,7 @@ mod tests {
report_context: Some(json!({
"candidate_index": candidate_index,
"retry_index": 0,
"sticky_key_attempts": sticky_key_attempts,
})),
}
}
@@ -1687,6 +1704,7 @@ mod tests {
CONTENT_TYPE.as_str().to_string(),
"application/json".to_string(),
)]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
@@ -1827,6 +1845,9 @@ mod tests {
report_context: Some(json!({
"candidate_index": candidate_index,
"retry_index": 0,
// Pin to a single attempt so this helper exercises candidate
// failover rather than the default same-key retry.
"sticky_key_attempts": 1,
"client_api_format": client_api_format,
"provider_api_format": client_api_format,
})),
@@ -1979,6 +2000,90 @@ mod tests {
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
}
#[tokio::test]
async fn openai_image_sync_heartbeat_retries_sticky_key_lazily_before_failover() {
let seen_plans = Arc::new(std::sync::Mutex::new(Vec::<(String, Option<String>)>::new()));
let seen_plans_for_override = Arc::clone(&seen_plans);
let state = AppState::new()
.expect("state should build")
.with_execution_runtime_sync_override_for_tests(move |plan| {
seen_plans_for_override
.lock()
.expect("mutex should lock")
.push((plan.endpoint_id.clone(), plan.candidate_id.clone()));
if plan.endpoint_id == "endpoint-retry" {
Ok(test_openai_image_execution_result(
plan,
StatusCode::TOO_MANY_REQUESTS.as_u16(),
json!({"error": {"message": "retry this candidate"}}),
))
} else {
Ok(test_openai_image_execution_result(
plan,
StatusCode::OK.as_u16(),
json!({"data": [{"b64_json": "second-candidate"}]}),
))
}
});
// Three total attempts on the sticky key; only one attempt is
// materialized up front, the other two are derived after each failure.
let attempts = vec![
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
0,
"endpoint-retry",
"candidate-retry",
3,
),
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
1,
"endpoint-success",
"candidate-success",
3,
),
];
let outcome = execute_openai_image_sync_heartbeat_attempts(
state,
"/v1/images/generations".to_string(),
"trace-image-heartbeat-sticky-retry".to_string(),
test_openai_image_heartbeat_decision(),
TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(),
attempts,
ProviderTransferTracker::default(),
Instant::now(),
)
.await
.expect("heartbeat attempts should execute");
let LocalExecutionRequestOutcome::Responded(response) = outcome else {
panic!("second candidate should return a response");
};
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
let seen_plans = seen_plans.lock().expect("mutex should lock").clone();
assert_eq!(
seen_plans
.iter()
.map(|(endpoint_id, _)| endpoint_id.as_str())
.collect::<Vec<_>>(),
[
"endpoint-retry",
"endpoint-retry",
"endpoint-retry",
"endpoint-success"
]
);
let sticky_candidate_ids = seen_plans[..3]
.iter()
.map(|(_, candidate_id)| candidate_id.clone())
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
sticky_candidate_ids.len(),
3,
"each derived same-key retry must carry a fresh candidate id"
);
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
}
#[tokio::test]
async fn openai_image_sync_heartbeat_honors_provider_transfer_limit() {
let call_count = Arc::new(AtomicUsize::new(0));
@@ -2012,6 +2117,7 @@ mod tests {
attempt.report_context = Some(json!({
"candidate_index": index,
"retry_index": 0,
"sticky_key_attempts": 1,
"local_failover_policy": {
"max_transfer_count": 1,
"max_transfer_timeout_seconds": 0
apps/aether-gateway/src/executor/outcome.rs
+16
View File
@@ -111,6 +111,22 @@ impl LocalExecutionRuntimeMissContext {
})
}
pub(crate) fn all_candidates_skipped_for_reasons(&self, reasons: &[&str]) -> bool {
if reasons.is_empty() || self.candidate_contexts.is_empty() {
return false;
}
self.candidate_contexts.iter().all(|candidate| {
candidate.candidate.status == RequestCandidateStatus::Skipped
&& candidate
.candidate
.skip_reason
.as_deref()
.map(str::trim)
.is_some_and(|value| reasons.contains(&value))
})
}
pub(crate) fn candidate_summary(&self) -> Option<String> {
const MAX_ITEMS: usize = 5;
apps/aether-gateway/src/frontdoor_loop_guard.rs
+26 -1
View File
@@ -41,12 +41,16 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
| "/v1/rerank"
| "/v1/responses"
| "/v1/responses/compact"
| "/v1/realtime"
| "/v1/realtime/calls"
| "/v1/live"
| "/v1/alpha/search"
| "/v1beta/files"
| "/upload/v1beta/files"
| "/v1beta/operations"
| "/v1/videos"
) || path.starts_with("/v1/videos/")
) || path.starts_with("/v1/live/")
|| path.starts_with("/v1/videos/")
|| path.starts_with("/v1beta/files/")
|| path.starts_with("/v1beta/operations/")
|| path.starts_with("/v1internal:")
@@ -141,3 +145,24 @@ fn normalize_host_for_frontdoor_loop_guard(host: &str) -> String {
fn is_loopbackish_host(host: &str) -> bool {
matches!(host, "localhost" | "127.0.0.1" | "::1" | "0.0.0.0" | "::")
}
#[cfg(test)]
mod tests {
use super::{
frontdoor_self_loop_public_ai_path, gateway_frontdoor_self_loop_guard_matches_with_port,
};
#[test]
fn realtime_is_protected_from_frontdoor_self_loops() {
assert!(frontdoor_self_loop_public_ai_path("/v1/realtime"));
assert!(frontdoor_self_loop_public_ai_path("/v1/realtime/calls"));
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
8084,
"ws://127.0.0.1:8084/v1/realtime?model=gpt-realtime"
));
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
8084,
"wss://localhost:8084/v1/realtime?model=gpt-realtime"
));
}
}
apps/aether-gateway/src/handlers/admin/model/routing.rs
+7 -20
View File
@@ -67,27 +67,14 @@ pub(crate) async fn build_admin_global_model_routing_payload(
.push(key);
}
let scheduling_mode = state
.read_system_config_json_value("scheduling_mode")
// Effective default scheduling: system-default routing group first, then
// legacy system-config keys.
let ordering_config = crate::scheduler::config::read_scheduler_ordering_config(state.app())
.await
.ok()
.flatten()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "cache_affinity".to_string());
let priority_mode = state
.read_system_config_json_value("provider_priority_mode")
.await
.ok()
.flatten()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "provider".to_string());
let keep_priority_on_conversion = state
.read_system_config_json_value("keep_priority_on_conversion")
.await
.ok()
.flatten()
.and_then(|value| value.as_bool())
.unwrap_or(false);
.unwrap_or_default();
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
let priority_mode = ordering_config.priority_mode_str().to_string();
let keep_priority_on_conversion = ordering_config.keep_priority_on_conversion;
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())

Some files were not shown because too many files have changed in this diff Show More