mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b11e1ecf4b | ||
|
|
a34d3ca7bc | ||
|
|
28f2879a05 | ||
|
|
2474259c22 | ||
|
|
79c2c5948f | ||
|
|
3c4bd8757f | ||
|
|
06f3443e56 | ||
|
|
88fd21f169 | ||
|
|
521201c926 | ||
|
|
be2123ea30 | ||
|
|
4b6d369225 | ||
|
|
76b81a6ceb | ||
|
|
29c8b679ba | ||
|
|
4091a0c2d8 | ||
|
|
f0ccc981fb | ||
|
|
d01ea7f775 | ||
|
|
fd749ea8f8 | ||
|
|
f0543b637e | ||
|
|
549bb9f62b | ||
|
|
9e57f78bab | ||
|
|
73c0ecdc4c | ||
|
|
bfd05722d5 | ||
|
|
9dc2ea4a2f | ||
|
|
6fb27e869e | ||
|
|
81c58b1cdc | ||
|
|
c8dc2e293a | ||
|
|
25081de913 | ||
|
|
6eb0917ff1 | ||
|
|
72d892d5e3 | ||
|
|
0b2da3087b | ||
|
|
5391fdb31c | ||
|
|
24036f1f84 | ||
|
|
b043fced0c | ||
|
|
40a2b70bf2 | ||
|
|
1e6ac77d3c | ||
|
|
3340934944 | ||
|
|
37a368de32 | ||
|
|
79cd8ac965 | ||
|
|
57c6b95596 | ||
|
|
a4f1843876 | ||
|
|
096e48621a | ||
|
|
36828b1a7e | ||
|
|
05ec1a4d38 | ||
|
|
9808e5939c | ||
|
|
a58911c127 | ||
|
|
0b9e080f1b | ||
|
|
f218e31c72 | ||
|
|
83f2f55e4a | ||
|
|
0384e39289 | ||
|
|
9036288d69 | ||
|
|
77b18eec68 | ||
|
|
3f7d473bea | ||
|
|
99cf39c76d | ||
|
|
f4d0576e6b | ||
|
|
2d9510d9e3 | ||
|
|
36d6e26291 | ||
|
|
eff522f2b6 | ||
|
|
65a175d48c | ||
|
|
f1872df293 | ||
|
|
71d87c5ab9 | ||
|
|
e23689ae22 | ||
|
|
4e527feb38 | ||
|
|
537c749b73 | ||
|
|
28b89200f7 | ||
|
|
a8f5c5a54e | ||
|
|
d94bca33dd | ||
|
|
9b33e60df2 | ||
|
|
1e99c5037d | ||
|
|
26cf170f33 | ||
|
|
8d54171771 |
@@ -34,6 +34,8 @@ updates:
|
||||
day: monday
|
||||
time: '09:20'
|
||||
timezone: Asia/Shanghai
|
||||
cooldown:
|
||||
default-days: 1
|
||||
groups:
|
||||
minor-and-patch:
|
||||
update-types:
|
||||
|
||||
@@ -19,21 +19,24 @@ jobs:
|
||||
env:
|
||||
COMPOSE_PROJECT_NAME: cpr-ci-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
CPR_IMAGE: codex-proxy-rs:ci-${{ inputs.ref }}
|
||||
CPR_ADMIN_DEFAULT_PASSWORD: 'admin @:/?#%2F ci password'
|
||||
CPR_POSTGRES_PASSWORD: 'pg @:/?#%2F ci'
|
||||
CPR_REDIS_PASSWORD: 'redis @:/?#%2F ci'
|
||||
CPR_SMOKE_PORT: '18080'
|
||||
TRIVY_QUIET: 'true'
|
||||
TRIVY_SKIP_VERSION_CHECK: 'true'
|
||||
steps:
|
||||
- name: Define isolated runtime paths
|
||||
- name: Define isolated runtime values
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
admin_password="$(openssl rand -hex 24)"
|
||||
postgres_password="$(openssl rand -hex 24)"
|
||||
redis_password="$(openssl rand -hex 24)"
|
||||
echo "::add-mask::$admin_password"
|
||||
echo "::add-mask::$postgres_password"
|
||||
echo "::add-mask::$redis_password"
|
||||
{
|
||||
echo "CPR_CONFIG_FILE=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/config.yaml"
|
||||
echo "CPR_DATA_DIR=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/data"
|
||||
echo "CPR_LOG_DIR=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/logs"
|
||||
echo "DIAGNOSTICS_DIR=$RUNNER_TEMP/codex-proxy-rs-diagnostics"
|
||||
echo "CI_ADMIN_PASSWORD=$admin_password"
|
||||
echo "CI_POSTGRES_PASSWORD=$postgres_password"
|
||||
echo "CI_REDIS_PASSWORD=$redis_password"
|
||||
} >> "$GITHUB_ENV"
|
||||
|
||||
- name: Checkout
|
||||
@@ -94,19 +97,35 @@ jobs:
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
runtime_dir="$(dirname "$CPR_CONFIG_FILE")"
|
||||
mkdir -p "$runtime_dir" "$DIAGNOSTICS_DIR"
|
||||
sudo install -d -m 0770 -o "$(id -u)" -g 10001 "$CPR_DATA_DIR" "$CPR_LOG_DIR"
|
||||
cp deploy/config.example.yaml "$CPR_CONFIG_FILE"
|
||||
sudo chown "$(id -u):10001" "$CPR_CONFIG_FILE"
|
||||
chmod 0640 "$CPR_CONFIG_FILE"
|
||||
umask 0077
|
||||
# Bind roots hand off ownership to the image-specific database users.
|
||||
# Sticky directories keep that initialization compatible with user-namespaced runners.
|
||||
install -d -m 1777 .runtime/postgres .runtime/redis
|
||||
mkdir -p "$DIAGNOSTICS_DIR"
|
||||
sudo install -d -m 0770 -o "$(id -u)" -g 10001 \
|
||||
.runtime/data .runtime/logs
|
||||
ruby <<'RUBY'
|
||||
path = "deploy/config.example.yaml"
|
||||
content = File.read(path)
|
||||
replacements = {
|
||||
"password: &postgres_password ''" => "password: &postgres_password '#{ENV.fetch("CI_POSTGRES_PASSWORD")}'",
|
||||
"password: &redis_password ''" => "password: &redis_password '#{ENV.fetch("CI_REDIS_PASSWORD")}'",
|
||||
"default_password: ''" => "default_password: '#{ENV.fetch("CI_ADMIN_PASSWORD")}'",
|
||||
}
|
||||
replacements.each do |placeholder, configured|
|
||||
raise "missing config placeholder: #{placeholder}" unless content.sub!(placeholder, configured)
|
||||
end
|
||||
File.write("deploy/config.yaml", content)
|
||||
RUBY
|
||||
sudo chown "$(id -u):10001" deploy/config.yaml
|
||||
chmod 0640 deploy/config.yaml
|
||||
|
||||
- name: Validate Docker Compose
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml config --quiet
|
||||
-f deploy/compose.yaml config --quiet
|
||||
|
||||
- name: Configure Redis host memory overcommit
|
||||
shell: bash
|
||||
@@ -120,23 +139,25 @@ jobs:
|
||||
run: |
|
||||
set -euo pipefail
|
||||
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
|
||||
-f deploy/compose.yaml)
|
||||
"${compose[@]}" up -d --no-build --wait postgres redis
|
||||
"${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
|
||||
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
|
||||
psql -At -U codex_proxy -d codex_proxy -c "SHOW TimeZone;" | tr -d '\r')" = Asia/Shanghai
|
||||
"${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
|
||||
psql -v ON_ERROR_STOP=1 -U codex_proxy -d codex_proxy -c \
|
||||
"CREATE TABLE cpr_ci_persistence_probe (value text PRIMARY KEY); INSERT INTO cpr_ci_persistence_probe VALUES ('persisted');"
|
||||
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
|
||||
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
|
||||
redis-cli SET cpr:ci:persistence-probe persisted
|
||||
"${compose[@]}" up -d --no-build --wait --force-recreate postgres redis
|
||||
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
|
||||
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
|
||||
psql -At -U codex_proxy -d codex_proxy -c \
|
||||
"SELECT value FROM cpr_ci_persistence_probe;")" = persisted
|
||||
test "$("${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
|
||||
test "$("${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
|
||||
redis-cli --raw GET cpr:ci:persistence-probe | tr -d '\r')" = persisted
|
||||
"${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
|
||||
"${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
|
||||
psql -v ON_ERROR_STOP=1 -U codex_proxy -d codex_proxy -c \
|
||||
"DROP TABLE cpr_ci_persistence_probe;"
|
||||
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
|
||||
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
|
||||
redis-cli DEL cpr:ci:persistence-probe
|
||||
|
||||
- name: Start and verify application
|
||||
@@ -144,22 +165,24 @@ jobs:
|
||||
run: |
|
||||
set -euo pipefail
|
||||
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
|
||||
-f deploy/compose.yaml)
|
||||
"${compose[@]}" up -d --no-build --wait codex-proxy-rs
|
||||
status="$(curl --silent --show-error --output /dev/null --write-out '%{http_code}' \
|
||||
"http://127.0.0.1:${CPR_SMOKE_PORT}/healthz")"
|
||||
"http://127.0.0.1:8080/healthz")"
|
||||
test "$status" = 204
|
||||
test "$("${compose[@]}" exec -T codex-proxy-rs id -u | tr -d '\r')" = 10001
|
||||
test "$("${compose[@]}" exec -T codex-proxy-rs id -g | tr -d '\r')" = 10001
|
||||
test "$("${compose[@]}" exec -T codex-proxy-rs date +%z | tr -d '\r')" = +0800
|
||||
test "$("${compose[@]}" exec -T codex-proxy-rs cat /etc/timezone | tr -d '\r')" = Asia/Shanghai
|
||||
"${compose[@]}" exec -T codex-proxy-rs sh -c \
|
||||
'touch /app/data/.ci-write /app/logs/.ci-write && rm /app/data/.ci-write /app/logs/.ci-write'
|
||||
'touch /app/.runtime/data/.ci-write /app/.runtime/logs/.ci-write && rm /app/.runtime/data/.ci-write /app/.runtime/logs/.ci-write'
|
||||
|
||||
- name: Verify clean logs and graceful shutdown
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
|
||||
-f deploy/compose.yaml)
|
||||
"${compose[@]}" stop --timeout 30 codex-proxy-rs
|
||||
container_id="$("${compose[@]}" ps -a -q codex-proxy-rs)"
|
||||
test -n "$container_id"
|
||||
@@ -191,7 +214,7 @@ jobs:
|
||||
mkdir -p "$DIAGNOSTICS_DIR"
|
||||
raw_dir="$(mktemp -d)"
|
||||
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
|
||||
-f deploy/compose.yaml)
|
||||
"${compose[@]}" config > "$raw_dir/compose-config.yaml" 2>&1
|
||||
"${compose[@]}" ps -a > "$raw_dir/compose-ps.txt" 2>&1
|
||||
"${compose[@]}" logs --no-color --timestamps > "$raw_dir/compose.log" 2>&1
|
||||
@@ -203,9 +226,9 @@ jobs:
|
||||
raw_dir = ENV.fetch("RAW_DIR")
|
||||
output_dir = ENV.fetch("OUTPUT_DIR")
|
||||
secrets = [
|
||||
ENV.fetch("CPR_ADMIN_DEFAULT_PASSWORD"),
|
||||
ENV.fetch("CPR_POSTGRES_PASSWORD"),
|
||||
ENV.fetch("CPR_REDIS_PASSWORD"),
|
||||
ENV.fetch("CI_ADMIN_PASSWORD"),
|
||||
ENV.fetch("CI_POSTGRES_PASSWORD"),
|
||||
ENV.fetch("CI_REDIS_PASSWORD"),
|
||||
]
|
||||
encoded_secrets = secrets.map do |secret|
|
||||
secret.encode("UTF-8").bytes.map do |byte|
|
||||
@@ -239,6 +262,5 @@ jobs:
|
||||
shell: bash
|
||||
run: |
|
||||
docker compose --project-name "$COMPOSE_PROJECT_NAME" \
|
||||
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml \
|
||||
down --volumes --remove-orphans || true
|
||||
sudo rm -rf "$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT"
|
||||
-f deploy/compose.yaml down --remove-orphans || true
|
||||
sudo rm -rf "$GITHUB_WORKSPACE/.runtime" "$GITHUB_WORKSPACE/deploy/config.yaml"
|
||||
|
||||
@@ -12,19 +12,29 @@ on:
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
run_backend:
|
||||
description: Run backend quality checks
|
||||
required: false
|
||||
default: true
|
||||
type: boolean
|
||||
run_frontend:
|
||||
description: Run frontend quality checks
|
||||
required: false
|
||||
default: true
|
||||
type: boolean
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
backend:
|
||||
if: inputs.run_backend
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
CPR_TEST_DATABASE_URL: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy
|
||||
CPR_TEST_REDIS_URL: redis://127.0.0.1:6379
|
||||
RUSTFLAGS: -D warnings
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:18-bookworm@sha256:d9c83446333daec3f0588cc709adb80c26090b7f9f0f7ec8d43c243385d79818
|
||||
@@ -56,7 +66,7 @@ jobs:
|
||||
ref: ${{ inputs.ref }}
|
||||
|
||||
- name: Setup Rust
|
||||
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
|
||||
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
|
||||
with:
|
||||
components: rustfmt,clippy
|
||||
|
||||
@@ -71,13 +81,14 @@ jobs:
|
||||
|
||||
- name: Rust clippy
|
||||
working-directory: backend
|
||||
run: cargo clippy --all-targets --all-features --locked -- -D warnings
|
||||
run: cargo clippy --all-targets --all-features --locked
|
||||
|
||||
- name: Rust tests
|
||||
working-directory: backend
|
||||
run: cargo test --test main --locked
|
||||
|
||||
frontend:
|
||||
if: inputs.run_frontend
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
@@ -94,7 +105,7 @@ jobs:
|
||||
version: 11.7.0
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
|
||||
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
|
||||
@@ -3,8 +3,10 @@ name: CI
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- '**'
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
@@ -18,6 +20,8 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
outputs:
|
||||
backend: ${{ steps.filter.outputs.backend }}
|
||||
frontend: ${{ steps.filter.outputs.frontend }}
|
||||
container: ${{ steps.filter.outputs.container }}
|
||||
steps:
|
||||
- name: Checkout
|
||||
@@ -26,11 +30,18 @@ jobs:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Detect container changes
|
||||
- name: Detect affected checks
|
||||
id: filter
|
||||
uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2
|
||||
with:
|
||||
filters: |
|
||||
backend:
|
||||
- 'backend/**'
|
||||
- 'release/version.yaml'
|
||||
- '.github/workflows/**'
|
||||
frontend:
|
||||
- 'frontend/**'
|
||||
- '.github/workflows/**'
|
||||
container:
|
||||
- 'backend/**'
|
||||
- 'frontend/**'
|
||||
@@ -41,12 +52,15 @@ jobs:
|
||||
- '.github/workflows/ci.yml'
|
||||
|
||||
quality:
|
||||
needs: changes
|
||||
uses: ./.github/workflows/_quality.yml
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
run_backend: ${{ needs.changes.outputs.backend == 'true' }}
|
||||
run_frontend: ${{ needs.changes.outputs.frontend == 'true' }}
|
||||
|
||||
container:
|
||||
needs: [changes, quality]
|
||||
needs: changes
|
||||
if: needs.changes.outputs.container == 'true'
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
@@ -165,7 +165,7 @@ jobs:
|
||||
|
||||
build-docker-binary:
|
||||
name: build-docker-binary (${{ matrix.os }}-${{ matrix.arch }})
|
||||
needs: [release-info, quality]
|
||||
needs: release-info
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix: ${{ fromJson(needs.release-info.outputs.docker_matrix) }}
|
||||
@@ -205,7 +205,7 @@ jobs:
|
||||
-v "$HOME/.cargo:/cargo" \
|
||||
-v "$PWD/..:/workspace" \
|
||||
-w /workspace/backend \
|
||||
rust:1.95-bookworm@sha256:6258907abe69656e41cd992e0b705cdcfabcbbe3db374f92ed2d47121282d4a1 \
|
||||
rust:1.97-bookworm@sha256:7d0723df719e7f213b69dc7c8c595985c3f4b060cfbee4f7bc0e347a86fe3b6a \
|
||||
cargo build --release --locked --target "${{ matrix.target }}" --bin codex-proxy-rs
|
||||
cp "target/${{ matrix.target }}/release/${{ matrix.binary }}" "../dist/binaries/${{ matrix.os }}-${{ matrix.arch }}/${{ matrix.binary }}"
|
||||
working-directory: backend
|
||||
@@ -225,7 +225,7 @@ jobs:
|
||||
|
||||
build-binary:
|
||||
name: build-binary (${{ matrix.os }}-${{ matrix.arch }})
|
||||
needs: [release-info, quality]
|
||||
needs: release-info
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix: ${{ fromJson(needs.release-info.outputs.non_docker_matrix) }}
|
||||
@@ -239,7 +239,7 @@ jobs:
|
||||
ref: ${{ needs.release-info.outputs.git_sha }}
|
||||
|
||||
- name: Setup Rust
|
||||
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
|
||||
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
|
||||
with:
|
||||
targets: ${{ matrix.target }}
|
||||
|
||||
@@ -621,7 +621,7 @@ jobs:
|
||||
cosign sign --yes "${images[@]}"
|
||||
|
||||
- name: Publish GitHub Release
|
||||
uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3.0.1
|
||||
uses: softprops/action-gh-release@3d0d9888cb7fd7b750713d6e236d1fcb99157228 # v3.0.2
|
||||
with:
|
||||
tag_name: ${{ needs.release-info.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
|
||||
@@ -3,8 +3,10 @@ name: Security Scan
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- '**'
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
schedule:
|
||||
- cron: '0 3 * * 1'
|
||||
|
||||
@@ -26,7 +28,7 @@ jobs:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
|
||||
- name: Setup Rust
|
||||
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
|
||||
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
|
||||
|
||||
- name: Cache Rust
|
||||
uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
|
||||
@@ -55,7 +57,7 @@ jobs:
|
||||
version: 11.7.0
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
|
||||
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
|
||||
+30
-32
@@ -1,43 +1,41 @@
|
||||
# Rust build output.
|
||||
# Rust
|
||||
/backend/target/
|
||||
/target/
|
||||
/dist/
|
||||
/package/
|
||||
/checksums.txt
|
||||
/codex-proxy-rs_*.tar.gz
|
||||
/codex-proxy-rs_*.zip
|
||||
*.profraw
|
||||
*.profdata
|
||||
|
||||
# Frontend dependencies and build output.
|
||||
# Frontend
|
||||
/frontend/node_modules/
|
||||
/frontend/dist/
|
||||
.pnpm-store
|
||||
/frontend/.vite/
|
||||
/.pnpm-store/
|
||||
/frontend/.pnpm-store/
|
||||
*.tsbuildinfo
|
||||
|
||||
# Application runtime state.
|
||||
*.runtime
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
*.db
|
||||
*.db-shm
|
||||
*.db-wal
|
||||
# Release artifacts
|
||||
/dist/
|
||||
|
||||
# Runtime data
|
||||
/.runtime/
|
||||
*.log
|
||||
|
||||
# Local configuration.
|
||||
/local.yaml
|
||||
/local.yml
|
||||
/.env
|
||||
/config.yaml
|
||||
/backend/config.yaml
|
||||
/deploy/.env
|
||||
# Local configuration and secrets
|
||||
/deploy/config.yaml
|
||||
|
||||
# Local workspaces and tooling state.
|
||||
# Editors and operating systems
|
||||
/.idea/
|
||||
/.vscode/
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Local tools and worktrees
|
||||
/.codegraph/
|
||||
/.codex-ws-audit/
|
||||
/.agents/
|
||||
/.reasonix/
|
||||
/.worktrees/
|
||||
/worktrees/
|
||||
/.codex-ws-audit/
|
||||
/.codegraph/
|
||||
*.tmp
|
||||
*.swp
|
||||
*~
|
||||
*.codegraph
|
||||
/.agents/
|
||||
/.cargo/
|
||||
/.reasonix/
|
||||
*.tmp
|
||||
|
||||
@@ -1,125 +1,110 @@
|
||||
<!-- prettier-ignore -->
|
||||
<div align="center">
|
||||
<h1>Codex Proxy RS</h1>
|
||||
<p>OpenAI Responses 兼容的 Codex 多账号网关</p>
|
||||
<p>
|
||||
<a href="https://www.rust-lang.org/"><img alt="Rust 1.95" src="https://img.shields.io/badge/Rust-1.95-000?style=flat-square&logo=rust" /></a>
|
||||
<a href="https://github.com/tokio-rs/axum"><img alt="Axum 0.8" src="https://img.shields.io/badge/Axum-0.8-2f7d95?style=flat-square" /></a>
|
||||
<a href="https://vuejs.org/"><img alt="Vue 3.5" src="https://img.shields.io/badge/Vue-3.5-42b883?style=flat-square&logo=vuedotjs&logoColor=white" /></a>
|
||||
<a href="https://vite.dev/"><img alt="Vite 8" src="https://img.shields.io/badge/Vite-8-646cff?style=flat-square&logo=vite&logoColor=white" /></a>
|
||||
<a href="https://www.postgresql.org/"><img alt="PostgreSQL 18" src="https://img.shields.io/badge/PostgreSQL-18-4169e1?style=flat-square&logo=postgresql&logoColor=white" /></a>
|
||||
<a href="https://redis.io/"><img alt="Redis 8" src="https://img.shields.io/badge/Redis-8-ff4438?style=flat-square&logo=redis&logoColor=white" /></a>
|
||||
<a href="https://opensource.org/license/mit"><img alt="MIT License" src="https://img.shields.io/badge/License-MIT-blue?style=flat-square" /></a>
|
||||
</p>
|
||||
<p>
|
||||
<a href="#快速部署">快速部署</a>
|
||||
·
|
||||
<a href="#客户端接入">客户端接入</a>
|
||||
·
|
||||
<a href="docs/architecture.md">架构说明</a>
|
||||
</p>
|
||||
|
||||
<img src="frontend/public/favicon.svg" alt="Codex Proxy RS" width="80" height="80" />
|
||||
|
||||
# Codex Proxy RS
|
||||
|
||||
**基于 Rust 的高性能 Codex 多账号网关代理**
|
||||
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](https://opensource.org/license/mit)
|
||||
|
||||
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署文档](deploy/README.md) · [架构](docs/architecture.md)
|
||||
|
||||
</div>
|
||||
|
||||
Codex Proxy RS 为多个 Codex 账号提供统一的 Responses API。客户端使用同一个服务地址和 API Key;账号调度、并发限制、额度检查、故障换号和会话续接都在服务端完成。
|
||||
> [!NOTE]
|
||||
> 只支持 OpenAI Responses API,不提供 `/v1/chat/completions`。
|
||||
|
||||
## 功能
|
||||
## 能力
|
||||
|
||||
- 兼容 `/v1/responses`、SSE 流式响应、Review、Compact 和模型查询。
|
||||
- 支持 HTTP/SSE 与 WebSocket 上游连接。
|
||||
- 提供 `smart`、额度重置优先、轮询和粘性四种账号调度策略。
|
||||
- 单个账号发生额度、认证、模型或传输故障时,继续使用账号池中的其他可用账号。
|
||||
- 使用会话亲和和受限历史重放处理 `previous_response_not_found` 与跨账号续接。
|
||||
- 为每个账号派生稳定且隔离的 installation、session、thread、turn 和 cache 标识。
|
||||
- 提供账号、API Key、模型映射、运行参数、用量、错误记录和系统更新管理界面。
|
||||
- 使用 PostgreSQL 保存权威业务数据,使用 Redis 保存会话、租约、模型快照和会话亲和数据。
|
||||
- 内置日志轮转、遥测保留期清理和 PostgreSQL/Redis 健康检查。
|
||||
| 领域 | 实现 |
|
||||
| -------- | ------------------------------------------------ |
|
||||
| 协议 | Responses HTTP、SSE、官方 WebSocket |
|
||||
| 传输 | 会话级 WebSocket 池、共享 HTTP/2 后备 |
|
||||
| 调度 | 智能调度、额度优先、轮询、粘滞 |
|
||||
| 故障处理 | 额度、认证、封禁、模型与传输故障切换 |
|
||||
| 身份 | 账号身份隔离、稳定 installation ID、会话拓扑透传 |
|
||||
| 管理 | 账号、客户端 Key、模型别名、运行设置与在线更新 |
|
||||
| 遥测 | Token、费用、缓存、TTFT、延迟与错误 |
|
||||
|
||||
代理只补充账号凭据、上游请求头和账号作用域身份,并在必要时执行会话恢复;其余请求与响应尽量保持上游语义。
|
||||
## 快速开始
|
||||
|
||||
## 快速部署
|
||||
需要 Docker Engine 与 Docker Compose Plugin。
|
||||
|
||||
需要 Docker Engine 和 Docker Compose Plugin。下面的 Compose 默认使用:
|
||||
|
||||
- 应用镜像:`ghcr.io/zyycn/codex-proxy-rs:latest`
|
||||
- PostgreSQL 18
|
||||
- Redis 8
|
||||
- 管理端地址:`http://127.0.0.1:8080`
|
||||
|
||||
### 1. 获取项目
|
||||
### 1. 准备
|
||||
|
||||
```bash
|
||||
git clone https://github.com/zyycn/codex-proxy-rs.git
|
||||
cd codex-proxy-rs
|
||||
|
||||
mkdir -p .runtime/data .runtime/logs
|
||||
install -d -m 0750 .runtime/postgres .runtime/redis
|
||||
cp deploy/config.example.yaml deploy/config.yaml
|
||||
sudo chown "$(id -u):10001" deploy/config.yaml
|
||||
chmod 0640 deploy/config.yaml
|
||||
```
|
||||
|
||||
### 2. 创建配置
|
||||
分别执行三次:
|
||||
|
||||
```bash
|
||||
mkdir -p .runtime/data .runtime/logs
|
||||
cp deploy/config.example.yaml .runtime/config.yaml
|
||||
cp deploy/.env.example deploy/.env
|
||||
chmod 0600 deploy/.env
|
||||
openssl rand -hex 24
|
||||
```
|
||||
|
||||
编辑 `deploy/.env`:
|
||||
将结果写入 `deploy/config.yaml`:
|
||||
|
||||
```dotenv
|
||||
CPR_ADMIN_DEFAULT_PASSWORD='<管理员密码,至少 12 个字符>'
|
||||
CPR_POSTGRES_PASSWORD='<PostgreSQL 密码>'
|
||||
CPR_REDIS_PASSWORD='<Redis 密码>'
|
||||
```
|
||||
| 配置 | 约束 |
|
||||
| ------------------------------ | ------------------------ |
|
||||
| `x-cpr.database.password` | 48 位十六进制 |
|
||||
| `x-cpr.redis.password` | 48 位十六进制 |
|
||||
| `x-cpr.admin.default_password` | 至少 12 位,不能包含 `$` |
|
||||
|
||||
`deploy/.env` 已被 Git 忽略。应用会安全地把 PostgreSQL 和 Redis 密码写入运行时连接 URL,不需要把真实密码填入 `.runtime/config.yaml`。
|
||||
|
||||
Linux 宿主机需要保证容器用户 `10001:10001` 可以读写数据和日志目录:
|
||||
Linux 需要允许容器组写入运行目录:
|
||||
|
||||
```bash
|
||||
sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
|
||||
chmod 0770 .runtime/data .runtime/logs
|
||||
sudo chown "$(id -u):10001" .runtime/config.yaml
|
||||
chmod 0640 .runtime/config.yaml
|
||||
```
|
||||
|
||||
### 3. 启动
|
||||
### 2. 启动
|
||||
|
||||
```bash
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml config
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml pull
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d --no-build
|
||||
```
|
||||
|
||||
检查服务:
|
||||
|
||||
```bash
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml ps
|
||||
docker compose -f deploy/compose.yaml config --quiet
|
||||
docker compose -f deploy/compose.yaml pull
|
||||
docker compose -f deploy/compose.yaml up -d --no-build
|
||||
curl -i http://127.0.0.1:8080/healthz
|
||||
```
|
||||
|
||||
`/healthz` 正常时返回 `204 No Content`。随后打开 <http://127.0.0.1:8080>:
|
||||
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址:
|
||||
`http://127.0.0.1:8080`。
|
||||
|
||||
- 用户名:`admin@cpr.local`
|
||||
- 密码:`deploy/.env` 中的 `CPR_ADMIN_DEFAULT_PASSWORD`
|
||||
### 3. 初始化
|
||||
|
||||
首次登录后导入账号,在“API 密钥”页面创建客户端 Key。
|
||||
1. 使用 `admin@cpr.local` 与初始密码登录。
|
||||
2. 添加 Codex 账号并执行连接测试。
|
||||
3. 创建 `sk_...` 客户端 Key。
|
||||
4. 复制 Codex CLI 配置,或导入 CCSwitch。
|
||||
|
||||
### 本地构建镜像
|
||||
> [!NOTE]
|
||||
> `admin.default_password` 只在首次创建管理员时生效。
|
||||
|
||||
需要验证当前源码时,不拉取 GHCR 镜像,直接构建:
|
||||
|
||||
```bash
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml build codex-proxy-rs
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d
|
||||
```
|
||||
完整部署、权限和密码轮换规则见 [deploy/README.md](deploy/README.md)。
|
||||
|
||||
## 客户端接入
|
||||
|
||||
OpenAI 兼容客户端填写:
|
||||
|
||||
- Base URL:`http://127.0.0.1:8080/v1`
|
||||
- API Key:管理端创建的 `sk_...` 客户端 Key
|
||||
|
||||
Responses 请求示例:
|
||||
| 配置 | 值 |
|
||||
| -------- | ---------------------------------------- |
|
||||
| Base URL | `http://127.0.0.1:8080/v1` |
|
||||
| API Key | 管理端创建的 `sk_...` |
|
||||
| 鉴权 | `Authorization: Bearer <client-api-key>` |
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/models \
|
||||
-H 'Authorization: Bearer <client-api-key>'
|
||||
|
||||
curl http://127.0.0.1:8080/v1/responses \
|
||||
-H 'Authorization: Bearer <client-api-key>' \
|
||||
-H 'Content-Type: application/json' \
|
||||
@@ -130,98 +115,46 @@ curl http://127.0.0.1:8080/v1/responses \
|
||||
}'
|
||||
```
|
||||
|
||||
可用接口:
|
||||
<details>
|
||||
<summary>API 路由</summary>
|
||||
|
||||
| 路由 | 说明 |
|
||||
| --- | --- |
|
||||
| `POST /v1/responses` | Responses 非流式或 SSE 流式请求 |
|
||||
| `POST /v1/responses/review` | Review 请求 |
|
||||
| `POST /v1/responses/compact` | Compact 请求 |
|
||||
| `GET /v1/models` | 模型列表 |
|
||||
| `GET /v1/models/catalog` | 模型目录 |
|
||||
| `GET /v1/models/{model_id}` | 模型详情 |
|
||||
| `GET /v1/models/{model_id}/info` | 模型运行信息 |
|
||||
| 路由 | 用途 |
|
||||
| --------------------------- | ------------------- |
|
||||
| `POST /v1/responses` | JSON 或 SSE |
|
||||
| `GET /v1/responses` | Responses WebSocket |
|
||||
| `POST /v1/responses/review` | Review |
|
||||
| `GET /v1/models` | 模型列表 |
|
||||
| `GET /v1/models/catalog` | 模型目录 |
|
||||
| `GET /v1/models/{model_id}` | 模型详情 |
|
||||
|
||||
Compose 默认只监听宿主机 `127.0.0.1`。需要从其他设备访问时,应在前面配置 HTTPS 反向代理,不要直接暴露 PostgreSQL 和 Redis 端口。
|
||||
所有 `/v1/*` 路由都需要客户端 API Key。
|
||||
|
||||
## 管理端
|
||||
</details>
|
||||
|
||||
管理端提供以下页面:
|
||||
|
||||
- Dashboard:请求趋势、服务状态、账号概览和最近请求。
|
||||
- 账号:导入、导出、OAuth、连接测试、quota 和 token 刷新。
|
||||
- API 密钥:创建、停用、删除和客户端接入信息。
|
||||
- 用量:请求记录、token、延迟、模型分布和运维错误。
|
||||
- 设置:模型别名、账号调度策略、单账号并发、请求间隔和刷新参数。
|
||||
- 系统:版本检查、更新、重启和回滚。
|
||||
|
||||
管理端登录会话保存在 Redis,并按 `admin.session_ttl_minutes` 自动过期。
|
||||
|
||||
## 配置
|
||||
|
||||
启动配置位于 `.runtime/config.yaml`,模板见 [deploy/config.example.yaml](deploy/config.example.yaml)。
|
||||
|
||||
| 配置段 | 用途 |
|
||||
| --- | --- |
|
||||
| `server` | HTTP 监听地址和端口 |
|
||||
| `api` | Codex 上游地址 |
|
||||
| `database`、`redis` | PostgreSQL 与 Redis 连接 |
|
||||
| `quota` | quota 刷新周期和耗尽账号过滤 |
|
||||
| `tls` | 上游 HTTP 协议偏好 |
|
||||
| `ws_pool` | WebSocket 连接池与首事件超时 |
|
||||
| `fingerprint` | Codex Desktop 请求指纹 |
|
||||
| `admin` | 默认管理员和登录会话 TTL |
|
||||
| `logging` | 日志级别、目录、大小和保留期 |
|
||||
| `telemetry` | 请求事实记录开关 |
|
||||
|
||||
模型别名、调度策略、单账号并发、请求间隔和刷新参数由管理端保存到 PostgreSQL,可在运行中更新。请求事实、运维错误和聚合时间桶由后台任务按各自保留期自动清理。
|
||||
|
||||
## 数据与升级
|
||||
|
||||
| 数据 | 默认位置 |
|
||||
| --- | --- |
|
||||
| 账号、API Key、设置、用量和错误记录 | Compose 命名卷 `postgres-data` |
|
||||
| 管理会话、租约、模型快照和会话亲和 | Compose 命名卷 `redis-data` |
|
||||
| 身份派生密钥和更新状态 | `.runtime/data` |
|
||||
| 文件日志 | `.runtime/logs` |
|
||||
| 启动配置 | `.runtime/config.yaml` |
|
||||
| Docker 密钥 | `deploy/.env` |
|
||||
|
||||
普通升级只替换应用镜像:
|
||||
## 运维
|
||||
|
||||
```bash
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml pull codex-proxy-rs
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d --no-build
|
||||
# 升级
|
||||
docker compose -f deploy/compose.yaml pull codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d --no-build
|
||||
|
||||
# 日志
|
||||
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
|
||||
|
||||
# 从源码构建
|
||||
docker compose -f deploy/compose.yaml build codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d
|
||||
```
|
||||
|
||||
更新镜像不会删除数据库卷,但仍应在升级前备份 PostgreSQL 和 `.runtime/data`。
|
||||
> [!IMPORTANT]
|
||||
> `.runtime/` 保存数据库、Redis、身份密钥和日志。删除该目录会永久清除运行状态。
|
||||
|
||||
不要使用 `docker compose down -v` 做普通升级,该命令会删除 PostgreSQL 与 Redis 命名卷。
|
||||
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
|
||||
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
|
||||
|
||||
`.runtime/data/identity_hmac_secret` 必须长期保留。丢失该文件会改变全部账号的 installation ID 和账号作用域身份。
|
||||
## 文档
|
||||
|
||||
## 查看状态
|
||||
|
||||
```bash
|
||||
# 容器状态
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml ps
|
||||
|
||||
# 应用日志
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml logs -f codex-proxy-rs
|
||||
|
||||
# PostgreSQL 和 Redis 日志
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml logs postgres redis
|
||||
|
||||
# 展开并检查最终 Compose 配置
|
||||
docker compose --env-file deploy/.env -f deploy/docker-compose.yml config
|
||||
```
|
||||
|
||||
`/healthz` 返回 `503` 时,先检查 PostgreSQL 和 Redis 的健康状态。容器反复重启时,再检查 `.runtime/config.yaml`、`deploy/.env` 和 `.runtime/data`、`.runtime/logs` 的权限。
|
||||
|
||||
## 架构
|
||||
|
||||
模块边界、请求链路、换号条件、previous response 恢复、存储归属和后台任务见 [架构说明](docs/architecture.md)。
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
- [部署](deploy/README.md)
|
||||
- [架构](docs/architecture.md)
|
||||
- [Release](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
- [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
|
||||
Generated
+31
-69
@@ -359,18 +359,19 @@ dependencies = [
|
||||
"bytes",
|
||||
"chrono",
|
||||
"config",
|
||||
"dirs",
|
||||
"flate2",
|
||||
"futures",
|
||||
"hex",
|
||||
"hmac 0.12.1",
|
||||
"hmac",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"indexmap",
|
||||
"ipnet",
|
||||
"percent-encoding",
|
||||
"rand 0.10.2",
|
||||
"rand_core 0.6.4",
|
||||
"redis",
|
||||
"reqwest",
|
||||
"roxmltree",
|
||||
"rustls",
|
||||
"rustls-native-certs",
|
||||
"rustls-pki-types",
|
||||
@@ -378,9 +379,10 @@ dependencies = [
|
||||
"semver",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"sha2 0.11.0",
|
||||
"sqlx",
|
||||
"subtle",
|
||||
"syn",
|
||||
"tar",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
@@ -453,6 +455,12 @@ dependencies = [
|
||||
"yaml-rust2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
|
||||
|
||||
[[package]]
|
||||
name = "cookie"
|
||||
version = "0.18.1"
|
||||
@@ -640,31 +648,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
|
||||
dependencies = [
|
||||
"block-buffer 0.12.1",
|
||||
"const-oid",
|
||||
"crypto-common 0.2.2",
|
||||
"ctutils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dirs"
|
||||
version = "6.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e"
|
||||
dependencies = [
|
||||
"dirs-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dirs-sys"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"option-ext",
|
||||
"redox_users",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "displaydoc"
|
||||
version = "0.2.6"
|
||||
@@ -1062,16 +1050,7 @@ version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018"
|
||||
dependencies = [
|
||||
"hmac 0.13.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"hmac",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1388,22 +1367,12 @@ version = "0.2.186"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "libredox"
|
||||
version = "0.1.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libsqlite3-sys"
|
||||
version = "0.37.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
@@ -1561,12 +1530,6 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
|
||||
|
||||
[[package]]
|
||||
name = "option-ext"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
|
||||
|
||||
[[package]]
|
||||
name = "parking"
|
||||
version = "2.2.1"
|
||||
@@ -1823,9 +1786,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2fa6f8e4b491d7a8ef3a9550a4d71969bd0064f46e32b8dbbcc7fc60dad94fed"
|
||||
checksum = "3cb5358643f48330db5c78856982faf39c93ba50c60d682b43d0344a3b649769"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
@@ -1857,17 +1820,6 @@ dependencies = [
|
||||
"bitflags",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_users"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac"
|
||||
dependencies = [
|
||||
"getrandom 0.2.17",
|
||||
"libredox",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex"
|
||||
version = "1.13.0"
|
||||
@@ -1955,6 +1907,15 @@ dependencies = [
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "roxmltree"
|
||||
version = "0.21.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1964b10c76125c36f8afe190065a4bf9a87bf324842c05701330bba9f1cacbb"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "2.1.3"
|
||||
@@ -1976,9 +1937,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.23.36"
|
||||
version = "0.23.42"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b"
|
||||
checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138"
|
||||
dependencies = [
|
||||
"once_cell",
|
||||
"ring",
|
||||
@@ -2054,6 +2015,7 @@ version = "0.10.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
@@ -2402,7 +2364,7 @@ dependencies = [
|
||||
"futures-util",
|
||||
"hex",
|
||||
"hkdf",
|
||||
"hmac 0.13.0",
|
||||
"hmac",
|
||||
"itoa",
|
||||
"log",
|
||||
"md-5",
|
||||
@@ -2931,9 +2893,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.23.4"
|
||||
version = "1.23.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53"
|
||||
checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"js-sys",
|
||||
|
||||
+12
-11
@@ -1,9 +1,9 @@
|
||||
[package]
|
||||
name = "codex-proxy-rs"
|
||||
edition = "2021"
|
||||
edition = "2024"
|
||||
license = "MIT"
|
||||
publish = false
|
||||
rust-version = "1.95"
|
||||
rust-version = "1.97"
|
||||
|
||||
[dependencies]
|
||||
argon2 = "0.5.3"
|
||||
@@ -13,29 +13,30 @@ base64 = "0.22.1"
|
||||
bytes = "1.11.1"
|
||||
chrono = { version = "0.4.45", features = ["serde", "clock"] }
|
||||
config = { version = "0.15.23", default-features = false, features = ["yaml"] }
|
||||
dirs = "6.0.0"
|
||||
flate2 = "1.1.9"
|
||||
futures = "0.3.32"
|
||||
hex = "0.4.3"
|
||||
hmac = "0.12.1"
|
||||
hmac = "0.13.0"
|
||||
hyper = { version = "1.10.1", features = ["http1", "server"] }
|
||||
hyper-util = { version = "0.1.20", features = ["tokio"] }
|
||||
indexmap = { version = "2.0", features = ["serde"] }
|
||||
ipnet = "2.12.0"
|
||||
percent-encoding = "2.3.2"
|
||||
rand = "0.10.1"
|
||||
# Enables password_hash::rand_core::OsRng for argon2 0.5.
|
||||
rand_core = { version = "0.6.4", features = ["getrandom"] }
|
||||
redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] }
|
||||
roxmltree = "0.21.1"
|
||||
# Pinned to the Codex Desktop 26.707.41301 HTTP stack; re-audit a newer Desktop before upgrading.
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
|
||||
rustls = { version = "=0.23.36", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustls = { version = "=0.23.42", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustls-native-certs = "0.8.4"
|
||||
rustls-pki-types = "1.14.1"
|
||||
secrecy = "0.10.3"
|
||||
secrecy = { version = "0.10.3", features = ["serde"] }
|
||||
semver = "1.0.27"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = { version = "1.0.150", features = ["preserve_order"] }
|
||||
sha2 = "0.10.9"
|
||||
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "sqlite-bundled", "chrono", "json"] }
|
||||
sha2 = "0.11.0"
|
||||
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "json"] }
|
||||
subtle = "2.6.1"
|
||||
tar = "0.4.44"
|
||||
thiserror = "2.0.18"
|
||||
@@ -53,6 +54,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
|
||||
|
||||
[dev-dependencies]
|
||||
sqlx = { version = "0.9.0", default-features = false, features = ["derive"] }
|
||||
syn = { version = "2.0.118", features = ["full", "visit"] }
|
||||
tempfile = "3.27.0"
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["test-util"] }
|
||||
wiremock = "0.6.5"
|
||||
@@ -62,8 +64,7 @@ tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungsteni
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f" }
|
||||
|
||||
[lints.rust]
|
||||
future_incompatible = { level = "warn", priority = -1 }
|
||||
nonstandard_style = { level = "deny", priority = -1 }
|
||||
warnings = "deny"
|
||||
unsafe_code = "forbid"
|
||||
|
||||
[lints.clippy]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
[toolchain]
|
||||
channel = "1.95"
|
||||
channel = "1.97"
|
||||
components = ["clippy", "rustfmt"]
|
||||
profile = "minimal"
|
||||
components = ["rustfmt", "clippy"]
|
||||
|
||||
@@ -33,12 +33,12 @@ pub async fn health_check_accounts(
|
||||
"stagger_ms must be between 500 and 30000",
|
||||
));
|
||||
}
|
||||
if let Some(concurrency) = payload.concurrency {
|
||||
if !(1..=10).contains(&concurrency) {
|
||||
return Err(AdminError::bad_request(
|
||||
"concurrency must be between 1 and 10",
|
||||
));
|
||||
}
|
||||
if let Some(concurrency) = payload.concurrency
|
||||
&& !(1..=10).contains(&concurrency)
|
||||
{
|
||||
return Err(AdminError::bad_request(
|
||||
"concurrency must be between 1 and 10",
|
||||
));
|
||||
}
|
||||
match state
|
||||
.services
|
||||
|
||||
@@ -1,56 +1,46 @@
|
||||
//! 账号管理 HTTP 处理器。
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
body::Body,
|
||||
extract::{Query, State},
|
||||
http::{
|
||||
header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE},
|
||||
StatusCode,
|
||||
header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE},
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::Utc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::{
|
||||
response::{
|
||||
parse_editable_update, AdminEnvelope, AdminError, AdminResponse, BatchDeleteData,
|
||||
EditableUpdateMessages, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE,
|
||||
ADMIN_OK_CODE, ADMIN_OK_MESSAGE, AdminEnvelope, AdminError, AdminResponse,
|
||||
BatchDeleteData, EditableUpdateMessages, PageMeta, parse_editable_update,
|
||||
},
|
||||
session::AdminAuth,
|
||||
},
|
||||
api::AppState,
|
||||
fleet::{
|
||||
account::AccountStatus,
|
||||
account::{AccountListSort, AccountSortField, AccountStatus},
|
||||
manage::{
|
||||
quota_view::{
|
||||
quota_data, AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage,
|
||||
},
|
||||
AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult,
|
||||
AccountUpdate, ManagedAccount, OAuthExchangeInput,
|
||||
},
|
||||
refresh::token_refresh_status_eligible,
|
||||
},
|
||||
infra::{
|
||||
format::{
|
||||
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
|
||||
format_billing_amount, format_percent, format_plain_number, format_tokens,
|
||||
nonnegative_i64_to_u64,
|
||||
},
|
||||
json::{clamp_limit, clamp_page, total_pages},
|
||||
json::{SortDirection, clamp_limit, clamp_page, total_pages},
|
||||
time::{china_datetime, china_relative_time, china_rfc3339},
|
||||
},
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageRecord,
|
||||
billing,
|
||||
buckets::query::{ModelUsageWindow, UsageBucketWindow},
|
||||
},
|
||||
telemetry::account_usage::query::AccountUsageRecord,
|
||||
};
|
||||
|
||||
const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200;
|
||||
const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts";
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -59,6 +49,9 @@ pub(crate) struct AccountsQuery {
|
||||
page: Option<u32>,
|
||||
page_size: Option<u32>,
|
||||
search: Option<String>,
|
||||
status: Option<String>,
|
||||
sort_by: Option<String>,
|
||||
sort_direction: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -191,21 +184,14 @@ impl AdminAccountData {
|
||||
struct AdminAccountUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
empty_response_count: i64,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
reasoning_tokens: i64,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
image_input_tokens: i64,
|
||||
image_output_tokens: i64,
|
||||
image_tokens_display: String,
|
||||
image_request_count: i64,
|
||||
image_request_failed_count: i64,
|
||||
created_tokens: i64,
|
||||
created_tokens_display: String,
|
||||
read_tokens: i64,
|
||||
@@ -221,16 +207,10 @@ impl AdminAccountUsageData {
|
||||
models: Vec<AdminAccountModelUsageData>,
|
||||
) -> Self {
|
||||
let request_count = usage.map_or(0, |usage| usage.window_request_count);
|
||||
let empty_response_count = 0;
|
||||
let input_tokens = usage.map_or(0, |usage| usage.window_input_tokens);
|
||||
let output_tokens = usage.map_or(0, |usage| usage.window_output_tokens);
|
||||
let cached_tokens = usage.map_or(0, |usage| usage.window_cached_tokens);
|
||||
let reasoning_tokens = 0;
|
||||
let total_tokens = input_tokens.saturating_add(output_tokens);
|
||||
let image_input_tokens = 0;
|
||||
let image_output_tokens = 0;
|
||||
let image_request_count = 0;
|
||||
let image_request_failed_count = 0;
|
||||
let created_tokens = input_tokens.saturating_sub(cached_tokens);
|
||||
let read_tokens = cached_tokens;
|
||||
let last_used_at = usage.and_then(|usage| usage.last_used_at);
|
||||
@@ -238,23 +218,14 @@ impl AdminAccountUsageData {
|
||||
Self {
|
||||
request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(request_count)),
|
||||
empty_response_count,
|
||||
input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(input_tokens)),
|
||||
output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(output_tokens)),
|
||||
cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(cached_tokens)),
|
||||
reasoning_tokens,
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
image_input_tokens,
|
||||
image_output_tokens,
|
||||
image_tokens_display: format_tokens(nonnegative_i64_to_u64(
|
||||
image_input_tokens + image_output_tokens,
|
||||
)),
|
||||
image_request_count,
|
||||
image_request_failed_count,
|
||||
created_tokens,
|
||||
created_tokens_display: format_tokens(nonnegative_i64_to_u64(created_tokens)),
|
||||
read_tokens,
|
||||
@@ -282,65 +253,32 @@ struct AdminAccountModelUsageData {
|
||||
cached_tokens_display: String,
|
||||
total_tokens: u64,
|
||||
total_tokens_display: String,
|
||||
total_cost_usd: f64,
|
||||
total_cost_usd_display: String,
|
||||
billing_amount_usd: f64,
|
||||
billing_amount_usd_display: String,
|
||||
last_used_at: Option<String>,
|
||||
last_used_at_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct AccountListStats {
|
||||
usage_by_account: HashMap<String, AccountUsageRecord>,
|
||||
quota_by_account: HashMap<String, AccountQuotaData>,
|
||||
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
|
||||
refreshing_account_ids: HashSet<String>,
|
||||
}
|
||||
|
||||
struct AccountModelUsageRecord {
|
||||
account_id: String,
|
||||
model: String,
|
||||
request_count: i64,
|
||||
error_count: i64,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
cached_tokens: i64,
|
||||
total_cost_usd: f64,
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct AccountQuotaWindowSelection {
|
||||
account_id: String,
|
||||
window: AccountQuotaUsageWindow,
|
||||
}
|
||||
|
||||
impl AccountListStats {
|
||||
fn data_for(&self, account: ManagedAccount) -> AdminAccountData {
|
||||
let account_id = account.id.clone();
|
||||
let token_refreshing = token_refresh_status_eligible(account.status)
|
||||
&& self.refreshing_account_ids.contains(&account_id);
|
||||
AdminAccountData::from_parts(
|
||||
account,
|
||||
self.usage_by_account.get(&account_id),
|
||||
self.quota_by_account.get(&account_id).cloned(),
|
||||
self.models_by_account
|
||||
.get(&account_id)
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
token_refreshing,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountSummaryData {
|
||||
total: u64,
|
||||
active: u64,
|
||||
high_usage: u64,
|
||||
quota_exhausted: u64,
|
||||
attention: u64,
|
||||
}
|
||||
|
||||
impl From<AccountSummary> for AdminAccountSummaryData {
|
||||
fn from(summary: AccountSummary) -> Self {
|
||||
Self {
|
||||
total: summary.total,
|
||||
active: summary.active,
|
||||
quota_exhausted: summary.quota_exhausted,
|
||||
attention: summary.attention,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountPageData {
|
||||
@@ -505,34 +443,35 @@ pub(crate) async fn accounts(
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let page = clamp_page(params.page.unwrap_or(1));
|
||||
let page_size = clamp_limit(params.page_size.unwrap_or(50));
|
||||
let quota_by_account = quota_snapshots_by_account(&state).await;
|
||||
let summary = account_summary_data(&state, "a_by_account).await;
|
||||
|
||||
match state
|
||||
let status = account_status_filter(params.status)?;
|
||||
let sort = account_list_sort(params.sort_by, params.sort_direction)?;
|
||||
let result = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.list_page(page, page_size, params.search)
|
||||
.account_list
|
||||
.query(AccountListQuery {
|
||||
page,
|
||||
page_size,
|
||||
search: params.search,
|
||||
status,
|
||||
sort,
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|item| stats.data_for(item))
|
||||
.collect(),
|
||||
total: page.total,
|
||||
page: page.page,
|
||||
page_size: page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::ok(page, summary),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(account_error(&error)),
|
||||
}
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: result
|
||||
.page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(account_list_item_data)
|
||||
.collect(),
|
||||
total: result.page.total,
|
||||
page: result.page.page,
|
||||
page_size: result.page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::ok(page, result.summary.into()),
|
||||
))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts`
|
||||
@@ -560,7 +499,8 @@ mod import_routes;
|
||||
mod lifecycle_routes;
|
||||
mod oauth_routes;
|
||||
mod probe_routes;
|
||||
mod query;
|
||||
pub(crate) mod query;
|
||||
mod quota_presenter;
|
||||
mod quota_routes;
|
||||
|
||||
pub(crate) use export_routes::*;
|
||||
@@ -569,4 +509,5 @@ pub(crate) use lifecycle_routes::*;
|
||||
pub(crate) use oauth_routes::*;
|
||||
pub(crate) use probe_routes::*;
|
||||
use query::*;
|
||||
use quota_presenter::*;
|
||||
pub(crate) use quota_routes::*;
|
||||
|
||||
@@ -1,4 +1,11 @@
|
||||
use super::*;
|
||||
use std::convert::Infallible;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::StreamExt;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::fleet::manage::AccountTestEvent;
|
||||
|
||||
/// `GET /api/admin/accounts/test?id=...&modelId=...`
|
||||
pub async fn test_account_connection(
|
||||
@@ -21,6 +28,12 @@ pub async fn test_account_connection(
|
||||
.test_connection_stream(&account_id, model)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?;
|
||||
let stream = stream.map(|event| {
|
||||
Ok::<Bytes, Infallible>(Bytes::from(format!(
|
||||
"data: {}\n\n",
|
||||
account_test_event_json(event)
|
||||
)))
|
||||
});
|
||||
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
@@ -32,6 +45,43 @@ pub async fn test_account_connection(
|
||||
.map_err(|_| AdminError::internal("Failed to build account test stream"))
|
||||
}
|
||||
|
||||
fn account_test_event_json(event: AccountTestEvent) -> Value {
|
||||
match event {
|
||||
AccountTestEvent::Started { model } => json!({
|
||||
"type": "test_start",
|
||||
"model": model,
|
||||
"text": "正在连接 Codex Responses"
|
||||
}),
|
||||
AccountTestEvent::Request { payload } => json!({
|
||||
"type": "request",
|
||||
"payload": payload
|
||||
}),
|
||||
AccountTestEvent::Content { text } => json!({ "type": "content", "text": text }),
|
||||
AccountTestEvent::Complete { account_status } => account_test_terminal_event(
|
||||
json!({ "type": "test_complete", "success": true }),
|
||||
account_status,
|
||||
),
|
||||
AccountTestEvent::Error {
|
||||
error,
|
||||
account_status,
|
||||
} => {
|
||||
account_test_terminal_event(json!({ "type": "error", "error": error }), account_status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn account_test_terminal_event(mut event: Value, account_status: Option<AccountStatus>) -> Value {
|
||||
if let Some(account_status) = account_status
|
||||
&& let Some(event) = event.as_object_mut()
|
||||
{
|
||||
event.insert(
|
||||
"accountStatus".to_string(),
|
||||
Value::String(account_status.as_str().to_string()),
|
||||
);
|
||||
}
|
||||
event
|
||||
}
|
||||
|
||||
/// `GET /api/admin/accounts/models?id=...`
|
||||
pub async fn account_models(
|
||||
State(state): State<AppState>,
|
||||
@@ -60,3 +110,29 @@ pub async fn account_models(
|
||||
Err(error) => Err(account_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts/models`
|
||||
pub async fn refresh_account_models(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Json(payload): Json<AccountActionRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let models = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.refresh_account_models(&payload.id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountModelsData {
|
||||
models: models
|
||||
.into_iter()
|
||||
.map(|model| AccountModelData {
|
||||
id: model.id,
|
||||
label: model.label,
|
||||
})
|
||||
.collect(),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,211 @@
|
||||
//! typed quota read model 的管理端 presenter。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::{
|
||||
api::admin::accounts_routes::query::{AccountQuotaReadModel, AccountQuotaWindowLocalUsage},
|
||||
fleet::quota::{QuotaWindowKind, QuotaWindowRole},
|
||||
infra::{
|
||||
format::{
|
||||
format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
|
||||
},
|
||||
time::{china_datetime, china_relative_time},
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct AccountQuotaData {
|
||||
refreshed_at_display: String,
|
||||
windows: Vec<AccountQuotaWindowData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowData {
|
||||
key: String,
|
||||
group: &'static str,
|
||||
window_seconds: Option<u64>,
|
||||
label_display: String,
|
||||
used_percent: Option<f64>,
|
||||
used_percent_display: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
local_usage: Option<AccountQuotaWindowLocalUsageData>,
|
||||
reset_at_display: String,
|
||||
window_used_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowLocalUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
impl From<AccountQuotaWindowLocalUsage> for AccountQuotaWindowLocalUsageData {
|
||||
fn from(usage: AccountQuotaWindowLocalUsage) -> Self {
|
||||
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
|
||||
Self {
|
||||
request_count: usage.request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
|
||||
input_tokens: usage.input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
|
||||
output_tokens: usage.output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AccountQuotaData {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
refreshed_at_display: "-".to_string(),
|
||||
windows: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn quota_data(read: AccountQuotaReadModel) -> AccountQuotaData {
|
||||
let now = Utc::now();
|
||||
AccountQuotaData {
|
||||
refreshed_at_display: china_relative_time(read.fetched_at, now),
|
||||
windows: read
|
||||
.windows
|
||||
.into_iter()
|
||||
.map(|read| {
|
||||
let window = read.quota.window;
|
||||
let window_seconds = window.window_seconds();
|
||||
let reset_at = window.reset_datetime();
|
||||
let used_percent = window.used_percent;
|
||||
AccountQuotaWindowData {
|
||||
key: read.quota.key,
|
||||
group: quota_window_group(window.kind),
|
||||
window_seconds,
|
||||
label_display: quota_window_label(
|
||||
read.quota.source.as_str(),
|
||||
read.quota.limit_name.as_deref(),
|
||||
read.quota.metered_feature.as_deref(),
|
||||
read.quota.role,
|
||||
window.kind,
|
||||
window_seconds,
|
||||
),
|
||||
used_percent,
|
||||
used_percent_display: used_percent
|
||||
.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
local_usage: read.local_usage.map(Into::into),
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
window_used_display: quota_window_used_display(reset_at, window_seconds, now),
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_group(kind: QuotaWindowKind) -> &'static str {
|
||||
match kind {
|
||||
QuotaWindowKind::Monthly => "monthly",
|
||||
QuotaWindowKind::ShortTerm | QuotaWindowKind::Weekly => "shortTerm",
|
||||
QuotaWindowKind::Other => "other",
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_label(
|
||||
source: &str,
|
||||
limit_name: Option<&str>,
|
||||
metered_feature: Option<&str>,
|
||||
role: QuotaWindowRole,
|
||||
kind: QuotaWindowKind,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
if role == QuotaWindowRole::Monthly {
|
||||
return "月限额".to_string();
|
||||
}
|
||||
let base = match kind {
|
||||
QuotaWindowKind::ShortTerm => "5小时限额".to_string(),
|
||||
QuotaWindowKind::Weekly => "周限额".to_string(),
|
||||
QuotaWindowKind::Monthly => "月限额".to_string(),
|
||||
QuotaWindowKind::Other => custom_window_label(window_seconds),
|
||||
};
|
||||
let label = limit_name
|
||||
.or(metered_feature)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty() && !value.eq_ignore_ascii_case("codex"));
|
||||
if source == "core" || label.is_none() {
|
||||
return base;
|
||||
}
|
||||
let label = label.unwrap_or_default();
|
||||
let label = if is_review_limit_label(label) {
|
||||
"代码审查"
|
||||
} else {
|
||||
label
|
||||
};
|
||||
format!("{label} · {base}")
|
||||
}
|
||||
|
||||
fn custom_window_label(window_seconds: Option<u64>) -> String {
|
||||
let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
|
||||
return "额度".to_string();
|
||||
};
|
||||
if seconds % 86_400 == 0 {
|
||||
format!("{}天限额", seconds / 86_400)
|
||||
} else if seconds % 3_600 == 0 {
|
||||
format!("{}小时限额", seconds / 3_600)
|
||||
} else {
|
||||
format!("{}分钟限额", seconds.div_ceil(60))
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_used_display(
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
window_seconds: Option<u64>,
|
||||
now: DateTime<Utc>,
|
||||
) -> String {
|
||||
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
|
||||
return "-".to_string();
|
||||
};
|
||||
let remaining = reset_at
|
||||
.signed_duration_since(now)
|
||||
.num_seconds()
|
||||
.max(0)
|
||||
.cast_unsigned();
|
||||
let used = window_seconds.saturating_sub(remaining);
|
||||
format!(
|
||||
"{} / {}",
|
||||
format_duration_days(used),
|
||||
format_duration_days(window_seconds)
|
||||
)
|
||||
}
|
||||
|
||||
fn format_duration_days(seconds: u64) -> String {
|
||||
let days = seconds as f64 / 86_400.0;
|
||||
if days >= 1.0 {
|
||||
format!("{days:.1}d")
|
||||
} else {
|
||||
format!("{:.1}h", seconds as f64 / 3_600.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn is_review_limit_label(value: &str) -> bool {
|
||||
let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_");
|
||||
normalized == "review"
|
||||
|| normalized == "code_review"
|
||||
|| normalized == "codex_review"
|
||||
|| normalized == "codex_code_review"
|
||||
|| normalized.contains("code_review")
|
||||
|| normalized.contains("codex_review")
|
||||
}
|
||||
@@ -14,13 +14,25 @@ pub async fn account_quota(
|
||||
.await
|
||||
{
|
||||
Ok(data) => {
|
||||
let quota = data.get("quota").cloned().unwrap_or(Value::Null);
|
||||
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
|
||||
let quota_json = quota.to_string();
|
||||
let mut quota_data = quota_data("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
let quota = serde_json::to_value(&data.quota)
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let raw = data.raw;
|
||||
let quota_read = AccountQuotaReadModel::from_snapshot(data.quota, Some(Utc::now()));
|
||||
let account = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.get(&account_id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?
|
||||
.ok_or_else(account_not_found)?;
|
||||
let item = state
|
||||
.services
|
||||
.account_list
|
||||
.enrich_account(account, quota_read)
|
||||
.await
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let quota_data = item.quota.clone().map(quota_data).unwrap_or_default();
|
||||
let account = account_list_item_data(item);
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
|
||||
@@ -39,23 +51,3 @@ pub async fn account_quota(
|
||||
Err(e) => Err(account_error(&e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AccountQuotaData,
|
||||
) -> Result<AdminAccountData, AdminError> {
|
||||
let Some(account) = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.get(account_id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?
|
||||
else {
|
||||
return Err(account_not_found());
|
||||
};
|
||||
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
|
||||
let stats =
|
||||
account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?;
|
||||
Ok(stats.data_for(account))
|
||||
}
|
||||
|
||||
@@ -1,21 +1,21 @@
|
||||
//! 管理端认证路由。
|
||||
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
Extension, Json,
|
||||
extract::State,
|
||||
http::{HeaderMap, HeaderValue, StatusCode, header::SET_COOKIE},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
api::{
|
||||
AppState,
|
||||
admin::{
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
session::{admin_session_cookie, ADMIN_SESSION_COOKIE},
|
||||
session::{ADMIN_SESSION_COOKIE, admin_session_cookie},
|
||||
},
|
||||
middleware::request_id::ClientIp,
|
||||
AppState,
|
||||
},
|
||||
auth::types::SessionError,
|
||||
infra::time::china_rfc3339,
|
||||
@@ -107,15 +107,14 @@ pub(crate) async fn logout(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Response, AdminError> {
|
||||
if let Some(session_id) = admin_session_cookie(&headers) {
|
||||
if let Err(error) = state
|
||||
if let Some(session_id) = admin_session_cookie(&headers)
|
||||
&& let Err(error) = state
|
||||
.services
|
||||
.admin_sessions
|
||||
.delete_session(&session_id)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %error, "failed to revoke admin session during logout");
|
||||
}
|
||||
{
|
||||
tracing::warn!(error = %error, "Failed to revoke admin session during logout");
|
||||
}
|
||||
|
||||
let mut response = AdminResponse::new(
|
||||
|
||||
+210
-126
@@ -1,21 +1,25 @@
|
||||
//! Dashboard 遥测聚合与展示模型。
|
||||
//! Dashboard 遥测事实的 HTTP presenter。
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::infra::{
|
||||
format::{
|
||||
format_compact_number, format_cost, format_duration_ms, format_percent, format_rate,
|
||||
format_tokens, nonnegative_i64_to_u64,
|
||||
format_billing_amount, format_compact_number, format_duration_ms, format_percent,
|
||||
format_rate, format_tokens, nonnegative_i64_to_u64,
|
||||
},
|
||||
time::{china_day_start, china_hour, china_hour_start, china_quarter_hour_start},
|
||||
time::{china_day_start, china_quarter_hour_start},
|
||||
};
|
||||
|
||||
use super::account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket};
|
||||
use crate::telemetry::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
|
||||
use crate::telemetry::usage::insights::RequestHealthTimeBucket;
|
||||
|
||||
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
|
||||
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
|
||||
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
|
||||
const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10;
|
||||
const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3;
|
||||
const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 99.0;
|
||||
const TREND_SLOT_MINUTES: i64 = 15;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct DashboardAccountCounts {
|
||||
@@ -39,7 +43,35 @@ pub struct DashboardHealthTimelineData {
|
||||
title: String,
|
||||
description: String,
|
||||
reliability_display: String,
|
||||
points: String,
|
||||
status: DashboardHealthTimelineStatus,
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
points: Vec<DashboardHealthTimelinePointData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DashboardHealthTimelinePointData {
|
||||
time: String,
|
||||
status: DashboardHealthTimelineStatus,
|
||||
reliability_display: String,
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum DashboardHealthTimelineStatus {
|
||||
Future,
|
||||
NoData,
|
||||
Unavailable,
|
||||
Unstable,
|
||||
LowSample,
|
||||
Stable,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -78,7 +110,7 @@ struct DashboardTokenCardData {
|
||||
today_tokens_value: u64,
|
||||
yesterday_tokens_value: u64,
|
||||
total_tokens: String,
|
||||
today_cost_usd: String,
|
||||
total_billing_amount_usd: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -117,13 +149,13 @@ struct DashboardTrendPointData {
|
||||
errors: String,
|
||||
errors_value: u64,
|
||||
latency: String,
|
||||
latency_value: u64,
|
||||
latency_value: Option<u64>,
|
||||
max_latency: String,
|
||||
max_latency_value: u64,
|
||||
max_latency_value: Option<u64>,
|
||||
min_latency: String,
|
||||
min_latency_value: u64,
|
||||
min_latency_value: Option<u64>,
|
||||
success_rate: String,
|
||||
success_rate_value: f64,
|
||||
success_rate_value: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -147,6 +179,14 @@ struct UsageWindow {
|
||||
min_first_token_bucket_latency: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct HealthWindow {
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
}
|
||||
|
||||
impl UsageWindow {
|
||||
fn tokens(self) -> u64 {
|
||||
self.input_tokens + self.output_tokens
|
||||
@@ -173,25 +213,25 @@ impl UsageWindow {
|
||||
pub fn dashboard_cards(
|
||||
accounts: DashboardAccountCounts,
|
||||
buckets: &[AccountUsageTimeBucket],
|
||||
lifetime_usage: &AccountUsageSummary,
|
||||
retained_usage: &RetainedUsageSummary,
|
||||
) -> DashboardCardsData {
|
||||
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
|
||||
dashboard_cards_at(accounts, buckets, retained_usage, Utc::now())
|
||||
}
|
||||
|
||||
pub fn dashboard_cards_at(
|
||||
accounts: DashboardAccountCounts,
|
||||
buckets: &[AccountUsageTimeBucket],
|
||||
lifetime_usage: &AccountUsageSummary,
|
||||
retained_usage: &RetainedUsageSummary,
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardCardsData {
|
||||
let today_start = china_day_start(now);
|
||||
let yesterday_start = today_start - Duration::days(1);
|
||||
let today = usage_window(buckets, today_start, now);
|
||||
let yesterday = usage_window(buckets, yesterday_start, today_start);
|
||||
let today_cost = cost_window(buckets, today_start, now).unwrap_or(0.0);
|
||||
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
|
||||
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
|
||||
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
|
||||
let total_requests = nonnegative_i64_to_u64(retained_usage.totals.request_count);
|
||||
let total_tokens = nonnegative_i64_to_u64(retained_usage.totals.input_tokens)
|
||||
.saturating_add(nonnegative_i64_to_u64(retained_usage.totals.output_tokens));
|
||||
let total_cached_tokens = nonnegative_i64_to_u64(retained_usage.totals.cached_tokens);
|
||||
|
||||
DashboardCardsData {
|
||||
accounts: DashboardAccountsCardData {
|
||||
@@ -213,13 +253,13 @@ pub fn dashboard_cards_at(
|
||||
today_tokens_value: today.tokens(),
|
||||
yesterday_tokens_value: yesterday.tokens(),
|
||||
total_tokens: format_tokens(total_tokens),
|
||||
today_cost_usd: format_cost(today_cost),
|
||||
total_billing_amount_usd: format_billing_amount(retained_usage.billing_amount_usd),
|
||||
},
|
||||
cache: DashboardCacheCardData {
|
||||
today_hit_rate: format_rate(today.cache_hit_rate()),
|
||||
today_hit_rate_value: today.cache_hit_rate(),
|
||||
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
|
||||
total_hit_rate: format_rate(summary_cache_hit_rate(lifetime_usage)),
|
||||
total_hit_rate: format_rate(summary_cache_hit_rate(retained_usage)),
|
||||
total_cached_tokens: format_tokens(total_cached_tokens),
|
||||
average_first_token_latency_ms: format_optional_duration_ms(
|
||||
today.avg_first_token_latency(),
|
||||
@@ -240,41 +280,51 @@ pub fn dashboard_trend_data_at(
|
||||
kind: DashboardTrendKind,
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardTrendData {
|
||||
let current_hour = china_hour_start(now);
|
||||
let current_slot = china_quarter_hour_start(now);
|
||||
let start = china_day_start(now);
|
||||
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
|
||||
let mut buckets = (0..=elapsed_hours)
|
||||
.map(|index| (start + Duration::hours(index), UsageWindow::default()))
|
||||
let elapsed_slots =
|
||||
current_slot.signed_duration_since(start).num_minutes() / TREND_SLOT_MINUTES;
|
||||
let mut buckets = (0..=elapsed_slots)
|
||||
.map(|index| {
|
||||
(
|
||||
start + Duration::minutes(TREND_SLOT_MINUTES * index),
|
||||
UsageWindow::default(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for record in records {
|
||||
if record.bucket_start < start || record.bucket_start > now {
|
||||
continue;
|
||||
}
|
||||
let record_hour = china_hour_start(record.bucket_start);
|
||||
if let Some((_, bucket)) = buckets
|
||||
.iter_mut()
|
||||
.find(|(bucket_start, _)| *bucket_start == record_hour)
|
||||
{
|
||||
apply_bucket(bucket, record);
|
||||
}
|
||||
let record_slot = china_quarter_hour_start(record.bucket_start);
|
||||
let slot_index =
|
||||
record_slot.signed_duration_since(start).num_minutes() / TREND_SLOT_MINUTES;
|
||||
let Ok(slot_index) = usize::try_from(slot_index) else {
|
||||
continue;
|
||||
};
|
||||
let Some((_, bucket)) = buckets.get_mut(slot_index) else {
|
||||
continue;
|
||||
};
|
||||
apply_bucket(bucket, record);
|
||||
}
|
||||
|
||||
let points = buckets
|
||||
.iter()
|
||||
.map(|(bucket_start, bucket)| {
|
||||
let latency = bucket.avg_first_token_latency().unwrap_or(0);
|
||||
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
|
||||
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
|
||||
let latency = bucket.avg_first_token_latency();
|
||||
let max_latency = bucket.max_first_token_bucket_latency();
|
||||
let min_latency = bucket.min_first_token_bucket_latency();
|
||||
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
|
||||
let success_rate = if bucket.requests > 0 {
|
||||
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
|
||||
let success_rate = (bucket.requests > 0).then(|| {
|
||||
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
|
||||
* 1000.0)
|
||||
.round()
|
||||
/ 10.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
});
|
||||
let elapsed_minutes = bucket_start.signed_duration_since(start).num_minutes();
|
||||
DashboardTrendPointData {
|
||||
time: format!("{:02}", china_hour(bucket_start)),
|
||||
time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60),
|
||||
requests: format_compact_number(bucket.requests),
|
||||
requests_value: bucket.requests,
|
||||
input_tokens: format_tokens(bucket.input_tokens),
|
||||
@@ -287,13 +337,15 @@ pub fn dashboard_trend_data_at(
|
||||
tokens_value: bucket.tokens(),
|
||||
errors: format_compact_number(bucket.errors),
|
||||
errors_value: bucket.errors,
|
||||
latency: format_optional_duration_ms(Some(latency)),
|
||||
latency: format_optional_duration_ms(latency),
|
||||
latency_value: latency,
|
||||
max_latency: format_optional_duration_ms(Some(max_latency)),
|
||||
max_latency: format_optional_duration_ms(max_latency),
|
||||
max_latency_value: max_latency,
|
||||
min_latency: format_optional_duration_ms(Some(min_latency)),
|
||||
min_latency: format_optional_duration_ms(min_latency),
|
||||
min_latency_value: min_latency,
|
||||
success_rate: format_percent(success_rate),
|
||||
success_rate: success_rate
|
||||
.map(format_percent)
|
||||
.unwrap_or_else(|| "—".to_string()),
|
||||
success_rate_value: success_rate,
|
||||
}
|
||||
})
|
||||
@@ -307,13 +359,13 @@ pub fn dashboard_trend_data_at(
|
||||
}
|
||||
|
||||
pub fn dashboard_health_timeline_data(
|
||||
buckets: &[AccountUsageTimeBucket],
|
||||
buckets: &[RequestHealthTimeBucket],
|
||||
) -> DashboardHealthTimelineData {
|
||||
dashboard_health_timeline_data_at(buckets, Utc::now())
|
||||
}
|
||||
|
||||
pub fn dashboard_health_timeline_data_at(
|
||||
records: &[AccountUsageTimeBucket],
|
||||
records: &[RequestHealthTimeBucket],
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardHealthTimelineData {
|
||||
let current_slot = china_quarter_hour_start(now);
|
||||
@@ -322,7 +374,7 @@ pub fn dashboard_health_timeline_data_at(
|
||||
.map(|index| {
|
||||
(
|
||||
start + Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index),
|
||||
UsageWindow::default(),
|
||||
HealthWindow::default(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -335,36 +387,62 @@ pub fn dashboard_health_timeline_data_at(
|
||||
.iter_mut()
|
||||
.find(|(bucket_start, _)| *bucket_start == record_slot)
|
||||
{
|
||||
apply_bucket(bucket, record);
|
||||
bucket.success_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(record.success_requests);
|
||||
bucket.failed_requests = bucket
|
||||
.failed_requests
|
||||
.saturating_add(record.failed_requests);
|
||||
bucket.cancelled_requests = bucket
|
||||
.cancelled_requests
|
||||
.saturating_add(record.cancelled_requests);
|
||||
bucket.caller_error_requests = bucket
|
||||
.caller_error_requests
|
||||
.saturating_add(record.caller_error_requests);
|
||||
}
|
||||
}
|
||||
|
||||
let (requests, errors) = buckets
|
||||
let totals = buckets
|
||||
.iter()
|
||||
.filter(|(bucket_start, _)| *bucket_start <= current_slot)
|
||||
.fold((0, 0), |(requests, errors), (_, bucket)| {
|
||||
(requests + bucket.requests, errors + bucket.errors)
|
||||
.fold(HealthWindow::default(), |mut totals, (_, bucket)| {
|
||||
totals.success_requests = totals
|
||||
.success_requests
|
||||
.saturating_add(bucket.success_requests);
|
||||
totals.failed_requests = totals
|
||||
.failed_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
totals.cancelled_requests = totals
|
||||
.cancelled_requests
|
||||
.saturating_add(bucket.cancelled_requests);
|
||||
totals.caller_error_requests = totals
|
||||
.caller_error_requests
|
||||
.saturating_add(bucket.caller_error_requests);
|
||||
totals
|
||||
});
|
||||
let reliability = (requests > 0)
|
||||
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
|
||||
DashboardHealthTimelineData {
|
||||
title: "请求健康时间线".to_string(),
|
||||
description: "请求可靠性".to_string(),
|
||||
reliability_display: reliability
|
||||
.map(|value| format!("{value:.1}%"))
|
||||
.unwrap_or_else(|| "-".to_string()),
|
||||
description: "有效请求可用性".to_string(),
|
||||
reliability_display: format_health_reliability(totals),
|
||||
status: health_status(totals, false),
|
||||
success_requests: totals.success_requests,
|
||||
failed_requests: totals.failed_requests,
|
||||
cancelled_requests: totals.cancelled_requests,
|
||||
caller_error_requests: totals.caller_error_requests,
|
||||
points: buckets
|
||||
.into_iter()
|
||||
.map(|(bucket_start, bucket)| {
|
||||
let success_rate = if bucket.requests > 0 {
|
||||
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
|
||||
* 1000.0)
|
||||
.round()
|
||||
/ 10.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
health_tone(bucket, success_rate, bucket_start > current_slot)
|
||||
.enumerate()
|
||||
.map(|(index, (bucket_start, bucket))| {
|
||||
let elapsed_minutes = index as i64 * HEALTH_TIMELINE_SLOT_MINUTES;
|
||||
DashboardHealthTimelinePointData {
|
||||
time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60),
|
||||
status: health_status(bucket, bucket_start > current_slot),
|
||||
reliability_display: format_health_reliability(bucket),
|
||||
success_requests: bucket.success_requests,
|
||||
failed_requests: bucket.failed_requests,
|
||||
cancelled_requests: bucket.cancelled_requests,
|
||||
caller_error_requests: bucket.caller_error_requests,
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
@@ -379,52 +457,48 @@ fn trend_summary(
|
||||
trend_summary_item(
|
||||
kind,
|
||||
"输入",
|
||||
points.iter().map(|p| p.input_tokens_value).sum(),
|
||||
Some(points.iter().map(|p| p.input_tokens_value).sum()),
|
||||
None,
|
||||
),
|
||||
trend_summary_item(
|
||||
kind,
|
||||
"输出",
|
||||
points.iter().map(|p| p.output_tokens_value).sum(),
|
||||
Some(points.iter().map(|p| p.output_tokens_value).sum()),
|
||||
None,
|
||||
),
|
||||
trend_summary_item(
|
||||
kind,
|
||||
"缓存",
|
||||
points.iter().map(|p| p.cached_tokens_value).sum(),
|
||||
Some(points.iter().map(|p| p.cached_tokens_value).sum()),
|
||||
None,
|
||||
),
|
||||
],
|
||||
DashboardTrendKind::Latency => {
|
||||
let samples = points
|
||||
let (latency_sum, latency_count) = points
|
||||
.iter()
|
||||
.filter(|point| point.latency_value > 0)
|
||||
.collect::<Vec<_>>();
|
||||
let average = if samples.is_empty() {
|
||||
0
|
||||
} else {
|
||||
samples.iter().map(|point| point.latency_value).sum::<u64>() / samples.len() as u64
|
||||
};
|
||||
.filter_map(|point| point.latency_value)
|
||||
.fold((0_u64, 0_u64), |(sum, count), latency| {
|
||||
(sum + latency, count + 1)
|
||||
});
|
||||
let average = (latency_count > 0).then(|| latency_sum / latency_count);
|
||||
vec![
|
||||
trend_summary_item(kind, "平均", average, None),
|
||||
trend_summary_item(
|
||||
kind,
|
||||
"最高",
|
||||
samples
|
||||
points
|
||||
.iter()
|
||||
.map(|p| p.max_latency_value)
|
||||
.max()
|
||||
.unwrap_or(0),
|
||||
.filter_map(|point| point.max_latency_value)
|
||||
.max(),
|
||||
None,
|
||||
),
|
||||
trend_summary_item(
|
||||
kind,
|
||||
"最低",
|
||||
samples
|
||||
points
|
||||
.iter()
|
||||
.filter_map(|p| (p.min_latency_value > 0).then_some(p.min_latency_value))
|
||||
.min()
|
||||
.unwrap_or(0),
|
||||
.filter_map(|point| point.min_latency_value)
|
||||
.min(),
|
||||
None,
|
||||
),
|
||||
]
|
||||
@@ -432,12 +506,13 @@ fn trend_summary(
|
||||
DashboardTrendKind::Errors => {
|
||||
let errors = points.iter().map(|point| point.errors_value).sum::<u64>();
|
||||
let requests = points.iter().map(|point| point.requests_value).sum::<u64>();
|
||||
let success_rate = (requests > 0)
|
||||
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
|
||||
let success_rate = (requests > 0).then(|| {
|
||||
(requests.saturating_sub(errors) as f64 / requests as f64 * 1000.0).round() / 10.0
|
||||
});
|
||||
vec![
|
||||
trend_summary_item(kind, "错误数", errors, None),
|
||||
trend_summary_item(kind, "成功率", 0, success_rate),
|
||||
trend_summary_item(kind, "总请求", requests, None),
|
||||
trend_summary_item(kind, "错误数", Some(errors), None),
|
||||
trend_summary_item(kind, "成功率", None, success_rate),
|
||||
trend_summary_item(kind, "总请求", Some(requests), None),
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -446,16 +521,18 @@ fn trend_summary(
|
||||
fn trend_summary_item(
|
||||
kind: DashboardTrendKind,
|
||||
label: &str,
|
||||
value: u64,
|
||||
value: Option<u64>,
|
||||
ratio: Option<f64>,
|
||||
) -> DashboardTrendSummaryData {
|
||||
DashboardTrendSummaryData {
|
||||
label: label.to_string(),
|
||||
value: match kind {
|
||||
DashboardTrendKind::Usage => format_tokens(value),
|
||||
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
|
||||
DashboardTrendKind::Errors => format_compact_number(value),
|
||||
},
|
||||
value: value
|
||||
.map(|value| match kind {
|
||||
DashboardTrendKind::Usage => format_tokens(value),
|
||||
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
|
||||
DashboardTrendKind::Errors => format_compact_number(value),
|
||||
})
|
||||
.unwrap_or_else(|| "—".to_string()),
|
||||
ratio: ratio.map(format_percent),
|
||||
}
|
||||
}
|
||||
@@ -474,19 +551,6 @@ fn usage_window(
|
||||
})
|
||||
}
|
||||
|
||||
fn cost_window(
|
||||
records: &[AccountUsageTimeBucket],
|
||||
start: DateTime<Utc>,
|
||||
end: DateTime<Utc>,
|
||||
) -> Option<f64> {
|
||||
let costs = records
|
||||
.iter()
|
||||
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
|
||||
.filter_map(|record| record.cost_usd)
|
||||
.collect::<Vec<_>>();
|
||||
(!costs.is_empty()).then(|| costs.into_iter().sum())
|
||||
}
|
||||
|
||||
fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
|
||||
window.requests += nonnegative_i64_to_u64(record.request_count);
|
||||
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
|
||||
@@ -511,27 +575,47 @@ fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option<u64> {
|
||||
.filter(|latency| *latency > 0)
|
||||
}
|
||||
|
||||
fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char {
|
||||
let successes = bucket.requests.saturating_sub(bucket.errors);
|
||||
fn health_status(bucket: HealthWindow, is_future: bool) -> DashboardHealthTimelineStatus {
|
||||
let eligible_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
if is_future {
|
||||
'0'
|
||||
} else if bucket.requests == 0 {
|
||||
'1'
|
||||
} else if successes == 0 {
|
||||
'2'
|
||||
} else if success_rate < 90.0 {
|
||||
'3'
|
||||
} else if successes < HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD {
|
||||
'4'
|
||||
DashboardHealthTimelineStatus::Future
|
||||
} else if eligible_requests == 0 {
|
||||
DashboardHealthTimelineStatus::NoData
|
||||
} else if bucket.success_requests == 0
|
||||
&& bucket.failed_requests >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD
|
||||
{
|
||||
DashboardHealthTimelineStatus::Unavailable
|
||||
} else if eligible_requests < HEALTH_TIMELINE_MIN_SAMPLE_SIZE {
|
||||
DashboardHealthTimelineStatus::LowSample
|
||||
} else if health_reliability(bucket)
|
||||
.is_some_and(|reliability| reliability < HEALTH_TIMELINE_STABLE_RELIABILITY)
|
||||
{
|
||||
DashboardHealthTimelineStatus::Unstable
|
||||
} else {
|
||||
'5'
|
||||
DashboardHealthTimelineStatus::Stable
|
||||
}
|
||||
}
|
||||
|
||||
fn summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
|
||||
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
|
||||
fn health_reliability(bucket: HealthWindow) -> Option<f64> {
|
||||
let eligible_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
(eligible_requests > 0)
|
||||
.then(|| bucket.success_requests as f64 / eligible_requests as f64 * 100.0)
|
||||
}
|
||||
|
||||
fn format_health_reliability(bucket: HealthWindow) -> String {
|
||||
health_reliability(bucket)
|
||||
.map(|value| format!("{value:.1}%"))
|
||||
.unwrap_or_else(|| "-".to_string())
|
||||
}
|
||||
|
||||
fn summary_cache_hit_rate(summary: &RetainedUsageSummary) -> Option<f64> {
|
||||
let input_tokens = nonnegative_i64_to_u64(summary.totals.input_tokens);
|
||||
(input_tokens > 0)
|
||||
.then(|| nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
|
||||
.then(|| nonnegative_i64_to_u64(summary.totals.cached_tokens) as f64 / input_tokens as f64)
|
||||
}
|
||||
|
||||
fn format_optional_duration_ms(value: Option<u64>) -> String {
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,26 +1,28 @@
|
||||
//! 管理端 v1 接口访问 Key 处理器(列表、创建、生命周期)。
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::response::{
|
||||
parse_editable_update, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse,
|
||||
BatchDeleteData, EditableUpdateMessages,
|
||||
AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, BatchDeleteData,
|
||||
EditableUpdateMessages, parse_editable_update,
|
||||
},
|
||||
api::admin::session::AdminAuth,
|
||||
api::AppState,
|
||||
infra::{
|
||||
json::{clamp_limit, clamp_page, NumberedPage},
|
||||
json::{NumberedPage, SortDirection, clamp_limit, clamp_page},
|
||||
time::{china_relative_time_str, china_rfc3339_str},
|
||||
},
|
||||
keys::types::{KeyManageError, ManagedClientApiKey},
|
||||
keys::types::{
|
||||
ClientApiKeyListSort, ClientApiKeySortField, KeyManageError, ManagedClientApiKey,
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -29,6 +31,8 @@ pub(crate) struct ApiKeysQuery {
|
||||
page: Option<u32>,
|
||||
page_size: Option<u32>,
|
||||
search: Option<String>,
|
||||
sort_by: Option<String>,
|
||||
sort_direction: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -120,10 +124,11 @@ pub(crate) async fn api_keys(
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let page = clamp_page(query.page.unwrap_or(1));
|
||||
let page_size = clamp_limit(query.page_size.unwrap_or(50));
|
||||
let sort = client_api_key_list_sort(query.sort_by, query.sort_direction)?;
|
||||
match state
|
||||
.services
|
||||
.admin_client_keys
|
||||
.list_page(page, page_size, query.search)
|
||||
.list_page(page, page_size, query.search, sort)
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
@@ -142,6 +147,31 @@ pub(crate) async fn api_keys(
|
||||
}
|
||||
}
|
||||
|
||||
fn client_api_key_list_sort(
|
||||
sort_by: Option<String>,
|
||||
sort_direction: Option<String>,
|
||||
) -> Result<Option<ClientApiKeyListSort>, AdminError> {
|
||||
let (sort_by, sort_direction) = match (sort_by, sort_direction) {
|
||||
(None, None) => return Ok(None),
|
||||
(Some(sort_by), Some(sort_direction)) => (sort_by, sort_direction),
|
||||
_ => {
|
||||
return Err(AdminError::bad_request(
|
||||
"Client API key sort field and direction must be provided together",
|
||||
));
|
||||
}
|
||||
};
|
||||
let field = match sort_by.trim() {
|
||||
"name" => ClientApiKeySortField::Name,
|
||||
"enabled" => ClientApiKeySortField::Enabled,
|
||||
"createdAt" => ClientApiKeySortField::CreatedAt,
|
||||
"lastUsedAt" => ClientApiKeySortField::LastUsedAt,
|
||||
_ => return Err(AdminError::bad_request("Invalid client API key sort field")),
|
||||
};
|
||||
let direction = SortDirection::parse(&sort_direction)
|
||||
.ok_or_else(|| AdminError::bad_request("Invalid client API key sort direction"))?;
|
||||
Ok(Some(ClientApiKeyListSort { field, direction }))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/keys`
|
||||
pub(crate) async fn create_api_key(
|
||||
State(state): State<AppState>,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
pub mod accounts_routes;
|
||||
pub mod auth_routes;
|
||||
mod dashboard_presenter;
|
||||
pub mod dashboard_routes;
|
||||
pub mod keys_routes;
|
||||
pub mod ops_routes;
|
||||
|
||||
@@ -8,16 +8,19 @@ use axum::{
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::{
|
||||
response::{AdminError, AdminPageEnvelope, AdminResponse},
|
||||
session::AdminAuth,
|
||||
},
|
||||
api::AppState,
|
||||
infra::{
|
||||
json::{clamp_limit, clamp_page, NumberedPage},
|
||||
json::{NumberedPage, clamp_limit, clamp_page},
|
||||
time::china_datetime,
|
||||
},
|
||||
telemetry::{ops::query::OpsQueryError, ops::store::OpsErrorFilter, ops::types::OpsErrorLog},
|
||||
telemetry::{
|
||||
ops::query::OpsQueryError,
|
||||
ops::types::{OpsErrorFilter, OpsErrorLog},
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
//! 管理端响应封装与 From 转换。
|
||||
|
||||
use crate::infra::json::{total_pages, NumberedPage};
|
||||
use crate::infra::json::{NumberedPage, total_pages};
|
||||
use axum::{
|
||||
Json,
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
@@ -16,7 +16,6 @@ const ADMIN_MALFORMED_JSON_CODE: u32 = 40000;
|
||||
const ADMIN_BAD_REQUEST_CODE: u32 = 40001;
|
||||
const ADMIN_INVALID_TIME_RANGE_CODE: u32 = 40002;
|
||||
const ADMIN_INVALID_MODEL_SOURCE_CODE: u32 = 40003;
|
||||
const ADMIN_INVALID_ENDPOINT_SOURCE_CODE: u32 = 40004;
|
||||
const ADMIN_SESSION_REQUIRED_CODE: u32 = 40101;
|
||||
const ADMIN_INVALID_CREDENTIALS_CODE: u32 = 40102;
|
||||
const ADMIN_INVALID_API_KEY_CODE: u32 = 40103;
|
||||
@@ -68,14 +67,6 @@ impl AdminError {
|
||||
)
|
||||
}
|
||||
|
||||
pub fn invalid_endpoint_source(message: impl Into<String>) -> Self {
|
||||
Self::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
ADMIN_INVALID_ENDPOINT_SOURCE_CODE,
|
||||
message,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn admin_session_required() -> Self {
|
||||
Self::new(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
//! 管理端路由。
|
||||
|
||||
use axum::{
|
||||
http::{header, HeaderValue},
|
||||
Router,
|
||||
http::{HeaderValue, header},
|
||||
middleware,
|
||||
response::Response,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
|
||||
use crate::api::AppState;
|
||||
@@ -14,7 +14,8 @@ use super::{
|
||||
accounts_routes::{
|
||||
account_models, account_quota, accounts, batch_delete_accounts, create_account,
|
||||
export_accounts, health_check_accounts, import_accounts, oauth_authorize_account,
|
||||
oauth_exchange_account, refresh_account, test_account_connection, update_account,
|
||||
oauth_exchange_account, refresh_account, refresh_account_models, test_account_connection,
|
||||
update_account,
|
||||
},
|
||||
auth_routes::{login, logout, session_status},
|
||||
dashboard_routes::{dashboard_summary, dashboard_trend},
|
||||
@@ -30,7 +31,8 @@ use super::{
|
||||
},
|
||||
usage_routes::{
|
||||
clear_usage_records, usage_record_detail, usage_records,
|
||||
usage_records_endpoint_distribution, usage_records_latency_trend,
|
||||
usage_records_endpoint_distribution, usage_records_insights_diagnostics,
|
||||
usage_records_insights_overview, usage_records_latency_trend,
|
||||
usage_records_model_distribution, usage_records_summary, usage_records_token_trend,
|
||||
},
|
||||
};
|
||||
@@ -44,7 +46,7 @@ pub fn router() -> Router<AppState> {
|
||||
.route("/api/admin/settings", get(settings).post(update_settings))
|
||||
.route(
|
||||
"/api/admin/settings/admin-api-key",
|
||||
get(admin_api_key_status).delete(delete_admin_api_key),
|
||||
get(admin_api_key_status).post(delete_admin_api_key),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/settings/admin-api-key/regenerate",
|
||||
@@ -71,7 +73,10 @@ pub fn router() -> Router<AppState> {
|
||||
post(oauth_exchange_account),
|
||||
)
|
||||
.route("/api/admin/accounts/test", get(test_account_connection))
|
||||
.route("/api/admin/accounts/models", get(account_models))
|
||||
.route(
|
||||
"/api/admin/accounts/models",
|
||||
get(account_models).post(refresh_account_models),
|
||||
)
|
||||
.route("/api/admin/accounts/delete", post(batch_delete_accounts))
|
||||
.route("/api/admin/accounts/update", post(update_account))
|
||||
.route("/api/admin/accounts/refresh", post(refresh_account))
|
||||
@@ -101,6 +106,14 @@ pub fn router() -> Router<AppState> {
|
||||
"/api/admin/usage/records/insights/latency-trend",
|
||||
get(usage_records_latency_trend),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/usage/records/insights/overview",
|
||||
get(usage_records_insights_overview),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/usage/records/insights/diagnostics",
|
||||
get(usage_records_insights_diagnostics),
|
||||
)
|
||||
.route("/api/admin/usage/records/delete", post(clear_usage_records))
|
||||
.route("/api/admin/usage/records/detail", get(usage_record_detail))
|
||||
.route("/api/admin/ops/errors", get(ops_errors))
|
||||
|
||||
@@ -2,10 +2,10 @@
|
||||
|
||||
use axum::{
|
||||
extract::FromRequestParts,
|
||||
http::{request::Parts, HeaderMap},
|
||||
http::{HeaderMap, request::Parts},
|
||||
};
|
||||
|
||||
use crate::api::{admin::response::AdminError, AppState};
|
||||
use crate::api::{AppState, admin::response::AdminError};
|
||||
|
||||
pub(super) const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session";
|
||||
|
||||
|
||||
@@ -6,11 +6,11 @@ use serde_json::Value;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::{
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
session::AdminAuth,
|
||||
},
|
||||
api::AppState,
|
||||
settings::{SettingsError, SettingsPatch, SettingsSnapshot, SettingsValidationError},
|
||||
};
|
||||
|
||||
@@ -115,7 +115,7 @@ pub(crate) async fn regenerate_admin_api_key(
|
||||
))
|
||||
}
|
||||
|
||||
/// `DELETE /api/admin/settings/admin-api-key`
|
||||
/// `POST /api/admin/settings/admin-api-key`
|
||||
pub(crate) async fn delete_admin_api_key(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
|
||||
@@ -6,24 +6,20 @@ use axum::{
|
||||
extract::{Json, Query, State},
|
||||
http::StatusCode,
|
||||
response::{
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
IntoResponse,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
},
|
||||
};
|
||||
use futures::Stream;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::{
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
session::AdminAuth,
|
||||
},
|
||||
api::AppState,
|
||||
update::{
|
||||
service::{self, RestartAction},
|
||||
state::operation_id,
|
||||
types::UpdateError,
|
||||
},
|
||||
update::{service::RestartAction, state::operation_id, types::UpdateError},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -38,20 +34,30 @@ pub(crate) struct UpdateRequest {
|
||||
target_version: String,
|
||||
}
|
||||
|
||||
pub(crate) async fn version(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
|
||||
pub(crate) async fn version(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(service::version_data().await),
|
||||
AdminEnvelope::ok(state.services.system_update.version_data().await),
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_detail(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Query(query): Query<UpdateDetailQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(service::update_detail(query.refresh.unwrap_or(false)).await),
|
||||
AdminEnvelope::ok(
|
||||
state
|
||||
.services
|
||||
.system_update
|
||||
.update_detail(query.refresh.unwrap_or(false))
|
||||
.await,
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -59,29 +65,25 @@ pub(crate) async fn update_event_stream(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, AdminError> {
|
||||
let receiver = service::subscribe_update_events();
|
||||
let shutdown = state.services.process_control.subscribe_shutdown();
|
||||
let receiver = state.services.system_update.subscribe_update_events();
|
||||
let stream = futures::stream::unfold(
|
||||
(receiver, shutdown, false),
|
||||
|(mut receiver, mut shutdown, close_after_send)| async move {
|
||||
(receiver, false),
|
||||
|(mut receiver, close_after_send)| async move {
|
||||
if close_after_send {
|
||||
return None;
|
||||
}
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.recv() => return None,
|
||||
received = receiver.recv() => match received {
|
||||
Ok(message) => {
|
||||
let close_after_send = message.is_terminal();
|
||||
let id = message.id().to_string();
|
||||
let data = serde_json::to_string(&message)
|
||||
.unwrap_or_else(|_| "{}".to_string());
|
||||
let event = Event::default().event("update").id(id).data(data);
|
||||
return Some((Ok(event), (receiver, shutdown, close_after_send)));
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
|
||||
match receiver.recv().await {
|
||||
Ok(message) => {
|
||||
let close_after_send = message.is_terminal();
|
||||
let id = message.id().to_string();
|
||||
let data =
|
||||
serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string());
|
||||
let event = Event::default().event("update").id(id).data(data);
|
||||
return Some((Ok(event), (receiver, close_after_send)));
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -90,11 +92,15 @@ pub(crate) async fn update_event_stream(
|
||||
}
|
||||
|
||||
pub(crate) async fn perform_update(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
payload: Option<Json<UpdateRequest>>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let target = payload.map(|Json(payload)| payload.target_version);
|
||||
let result = service::perform_update(target)
|
||||
let result = state
|
||||
.services
|
||||
.system_update
|
||||
.perform_update(target)
|
||||
.await
|
||||
.map_err(update_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
@@ -103,16 +109,31 @@ pub(crate) async fn perform_update(
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_status(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
|
||||
let status = service::update_status().map_err(update_error)?;
|
||||
pub(crate) async fn update_status(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let status = state
|
||||
.services
|
||||
.system_update
|
||||
.update_status()
|
||||
.map_err(update_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(status),
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn rollback(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
|
||||
let operation_id = service::rollback().await.map_err(update_error)?;
|
||||
pub(crate) async fn rollback(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let operation_id = state
|
||||
.services
|
||||
.system_update
|
||||
.rollback()
|
||||
.await
|
||||
.map_err(update_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({
|
||||
@@ -127,7 +148,11 @@ pub(crate) async fn restart(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let plan = service::restart_plan().map_err(update_error)?;
|
||||
let plan = state
|
||||
.services
|
||||
.system_update
|
||||
.restart_plan()
|
||||
.map_err(update_error)?;
|
||||
let message = plan.message;
|
||||
let process_control = state.services.process_control.clone();
|
||||
tokio::spawn(async move {
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 使用记录 HTTP 处理器。
|
||||
//! 使用统计与请求记录 HTTP 处理器。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
@@ -8,28 +8,30 @@ use axum::{
|
||||
response::IntoResponse,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
api::AppState,
|
||||
api::admin::{
|
||||
response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
|
||||
session::AdminAuth,
|
||||
},
|
||||
api::AppState,
|
||||
infra::{
|
||||
format::{
|
||||
format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64,
|
||||
format_multiplier, format_number, format_percent, format_token_price, format_tokens,
|
||||
format_billing_amount, format_compact_number, format_duration_ms,
|
||||
format_duration_ms_f64, format_multiplier, format_number, format_percent,
|
||||
format_token_price, format_tokens,
|
||||
},
|
||||
json::{clamp_limit, clamp_page, NumberedPage},
|
||||
json::{NumberedPage, clamp_limit, clamp_page},
|
||||
time::china_datetime,
|
||||
},
|
||||
telemetry::{
|
||||
usage::insights::{UsageDiagnosticsDimension, default_time_range},
|
||||
usage::query::{
|
||||
usage_record_cost_details, UsageQueryError, UsageQueryFilter, UsageRecordBreakdown,
|
||||
UsageRecordEndpointSource, UsageRecordModelSource, UsageRecordSummary,
|
||||
UsageRecordTrendPoint,
|
||||
UsageQueryError, UsageQueryFilter, UsageRecordBreakdown, UsageRecordModelSource,
|
||||
UsageRecordSummary, UsageRecordTrendPoint, usage_record_billing,
|
||||
},
|
||||
usage::types::{metadata_string, UsageRecord},
|
||||
usage::types::{UsageRecord, metadata_string},
|
||||
},
|
||||
};
|
||||
|
||||
@@ -63,6 +65,21 @@ pub(crate) struct UsageRecordDistributionQuery {
|
||||
source: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct UsageInsightsQuery {
|
||||
start_time: Option<String>,
|
||||
end_time: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct UsageDiagnosticsQuery {
|
||||
#[serde(flatten)]
|
||||
range: UsageInsightsQuery,
|
||||
dimension: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct UsageRecordDetailQuery {
|
||||
@@ -87,8 +104,12 @@ pub(crate) struct UsageRecordData {
|
||||
client_ip: Option<String>,
|
||||
user_agent: Option<String>,
|
||||
reasoning_effort: Option<String>,
|
||||
reasoning_preset: Option<String>,
|
||||
compact: bool,
|
||||
request_kind: Option<String>,
|
||||
subagent_kind: Option<String>,
|
||||
token_details: UsageRecordTokenDetailsData,
|
||||
cost_details: Option<UsageRecordCostDetailsData>,
|
||||
billing: Option<UsageRecordBillingData>,
|
||||
first_token_latency_ms: Option<i64>,
|
||||
first_token_latency_ms_display: String,
|
||||
latency_ms_display: String,
|
||||
@@ -100,37 +121,43 @@ pub(crate) struct UsageRecordTokenDetailsData {
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
reasoning_tokens: u64,
|
||||
total_tokens: u64,
|
||||
input_tokens_display: String,
|
||||
output_tokens_display: String,
|
||||
cached_tokens_display: String,
|
||||
cache_write_tokens_display: String,
|
||||
reasoning_tokens_display: String,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct UsageRecordCostDetailsData {
|
||||
input_cost: f64,
|
||||
output_cost: f64,
|
||||
cache_read_cost: f64,
|
||||
total_cost: f64,
|
||||
original_cost: f64,
|
||||
pub(crate) struct UsageRecordBillingData {
|
||||
input_amount: f64,
|
||||
output_amount: f64,
|
||||
cache_read_amount: f64,
|
||||
cache_write_amount: f64,
|
||||
standard_amount: f64,
|
||||
total_amount: f64,
|
||||
input_price_per_mtoken: f64,
|
||||
output_price_per_mtoken: f64,
|
||||
cache_read_price_per_mtoken: f64,
|
||||
cache_write_price_per_mtoken: f64,
|
||||
service_tier: Option<String>,
|
||||
service_tier_display: String,
|
||||
multiplier: f64,
|
||||
input_cost_display: String,
|
||||
output_cost_display: String,
|
||||
cache_read_cost_display: String,
|
||||
total_cost_display: String,
|
||||
original_cost_display: String,
|
||||
input_amount_display: String,
|
||||
output_amount_display: String,
|
||||
cache_read_amount_display: String,
|
||||
cache_write_amount_display: String,
|
||||
standard_amount_display: String,
|
||||
total_amount_display: String,
|
||||
input_price_display: String,
|
||||
output_price_display: String,
|
||||
cache_read_price_display: String,
|
||||
cache_write_price_display: String,
|
||||
multiplier_display: String,
|
||||
}
|
||||
|
||||
@@ -141,6 +168,7 @@ struct UsageRecordSummaryData {
|
||||
input_tokens: String,
|
||||
output_tokens: String,
|
||||
cached_tokens: String,
|
||||
cache_write_tokens: String,
|
||||
total_tokens: String,
|
||||
average_latency_ms: String,
|
||||
}
|
||||
@@ -154,13 +182,14 @@ struct UsageRecordBreakdownData {
|
||||
input_tokens: String,
|
||||
output_tokens: String,
|
||||
cached_tokens: String,
|
||||
cache_write_tokens: String,
|
||||
total_tokens: String,
|
||||
total_tokens_value: u64,
|
||||
total_tokens_total: String,
|
||||
total_tokens_total_value: u64,
|
||||
cost: String,
|
||||
actual_cost: String,
|
||||
account_cost: String,
|
||||
standard_billing_amount: String,
|
||||
actual_billing_amount: String,
|
||||
account_billing_amount: String,
|
||||
average_latency_ms: String,
|
||||
}
|
||||
|
||||
@@ -174,12 +203,14 @@ struct UsageRecordTrendPointData {
|
||||
output_tokens_value: u64,
|
||||
cached_tokens: String,
|
||||
cached_tokens_value: u64,
|
||||
cache_write_tokens: String,
|
||||
cache_write_tokens_value: u64,
|
||||
total_tokens: String,
|
||||
total_tokens_value: u64,
|
||||
cache_hit_rate: String,
|
||||
cache_hit_rate_value: f64,
|
||||
cost: String,
|
||||
actual_cost: String,
|
||||
standard_billing_amount: String,
|
||||
actual_billing_amount: String,
|
||||
average_latency_ms: String,
|
||||
average_latency_ms_value: Option<f64>,
|
||||
}
|
||||
@@ -264,14 +295,13 @@ pub(crate) async fn usage_records_model_distribution(
|
||||
pub(crate) async fn usage_records_endpoint_distribution(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Query(query): Query<UsageRecordDistributionQuery>,
|
||||
Query(query): Query<UsageRecordsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let source = endpoint_source_from_query(query.source)?;
|
||||
let filter = filter_from_query(query.records)?;
|
||||
let filter = filter_from_query(query)?;
|
||||
match state
|
||||
.services
|
||||
.usage_records
|
||||
.endpoint_distribution(filter, source)
|
||||
.endpoint_distribution(filter)
|
||||
.await
|
||||
{
|
||||
Ok(items) => Ok(AdminResponse::new(
|
||||
@@ -300,6 +330,39 @@ pub(crate) async fn usage_records_latency_trend(
|
||||
usage_records_trend_response(state, query).await
|
||||
}
|
||||
|
||||
/// `GET /api/admin/usage/records/insights/overview`
|
||||
pub(crate) async fn usage_records_insights_overview(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Query(query): Query<UsageInsightsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let (start, end) = insights_time_range(query)?;
|
||||
state
|
||||
.services
|
||||
.usage_records
|
||||
.insights_overview(start, end)
|
||||
.await
|
||||
.map(|insights| AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(insights)))
|
||||
.map_err(|error| log_error(&error))
|
||||
}
|
||||
|
||||
/// `GET /api/admin/usage/records/insights/diagnostics`
|
||||
pub(crate) async fn usage_records_insights_diagnostics(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Query(query): Query<UsageDiagnosticsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let dimension = diagnostics_dimension(query.dimension)?;
|
||||
let (start, end) = insights_time_range(query.range)?;
|
||||
state
|
||||
.services
|
||||
.usage_records
|
||||
.insights_diagnostics(start, end, dimension)
|
||||
.await
|
||||
.map(|insights| AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(insights)))
|
||||
.map_err(|error| log_error(&error))
|
||||
}
|
||||
|
||||
async fn usage_records_trend_response(
|
||||
state: AppState,
|
||||
query: UsageRecordsQuery,
|
||||
@@ -390,19 +453,6 @@ fn model_source_from_query(value: Option<String>) -> Result<UsageRecordModelSour
|
||||
}
|
||||
}
|
||||
|
||||
fn endpoint_source_from_query(
|
||||
value: Option<String>,
|
||||
) -> Result<UsageRecordEndpointSource, AdminError> {
|
||||
match non_empty(value).as_deref().unwrap_or("inbound") {
|
||||
"inbound" => Ok(UsageRecordEndpointSource::Inbound),
|
||||
"upstream" => Ok(UsageRecordEndpointSource::Upstream),
|
||||
"path" => Ok(UsageRecordEndpointSource::Path),
|
||||
_ => Err(AdminError::invalid_endpoint_source(
|
||||
"Invalid endpoint distribution source",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_datetime(
|
||||
value: Option<String>,
|
||||
) -> Result<Option<chrono::DateTime<chrono::Utc>>, AdminError> {
|
||||
@@ -414,6 +464,25 @@ fn optional_datetime(
|
||||
.map_err(|_| AdminError::invalid_time_range("Invalid time range"))
|
||||
}
|
||||
|
||||
fn insights_time_range(
|
||||
query: UsageInsightsQuery,
|
||||
) -> Result<(chrono::DateTime<chrono::Utc>, chrono::DateTime<chrono::Utc>), AdminError> {
|
||||
let end = optional_datetime(query.end_time)?.unwrap_or_else(chrono::Utc::now);
|
||||
let start = optional_datetime(query.start_time)?.unwrap_or_else(|| default_time_range(end).0);
|
||||
if start >= end {
|
||||
return Err(AdminError::invalid_time_range(
|
||||
"Start time must be earlier than end time",
|
||||
));
|
||||
}
|
||||
Ok((start, end))
|
||||
}
|
||||
|
||||
fn diagnostics_dimension(value: Option<String>) -> Result<UsageDiagnosticsDimension, AdminError> {
|
||||
let value = non_empty(value).unwrap_or_else(|| "model".to_string());
|
||||
UsageDiagnosticsDimension::parse(&value)
|
||||
.ok_or_else(|| AdminError::bad_request("Invalid diagnostics dimension"))
|
||||
}
|
||||
|
||||
fn log_error(error: &UsageQueryError) -> AdminError {
|
||||
match error {
|
||||
UsageQueryError::List | UsageQueryError::Get | UsageQueryError::Clear => {
|
||||
@@ -452,8 +521,16 @@ fn usage_record_data(
|
||||
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
|
||||
let reasoning_effort =
|
||||
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
|
||||
let reasoning_preset = metadata_string(&record.metadata, &["reasoningPreset"]);
|
||||
let compact = record
|
||||
.metadata
|
||||
.get("compact")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
let request_kind = metadata_string(&record.metadata, &["requestKind", "request_kind"]);
|
||||
let subagent_kind = metadata_string(&record.metadata, &["subagentKind", "subagent_kind"]);
|
||||
let token_details = usage_token_details(&record);
|
||||
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
|
||||
let billing = usage_billing(&record, upstream_model.as_deref(), &token_details);
|
||||
let first_token_latency_ms = record.first_token_ms;
|
||||
let first_token_latency_ms_display = format_duration_ms(first_token_latency_ms);
|
||||
let latency_ms_display = format_duration_ms(record.latency_ms);
|
||||
@@ -466,8 +543,12 @@ fn usage_record_data(
|
||||
client_ip,
|
||||
user_agent,
|
||||
reasoning_effort,
|
||||
reasoning_preset,
|
||||
compact,
|
||||
request_kind,
|
||||
subagent_kind,
|
||||
token_details,
|
||||
cost_details,
|
||||
billing,
|
||||
first_token_latency_ms,
|
||||
first_token_latency_ms_display,
|
||||
latency_ms_display,
|
||||
@@ -503,6 +584,10 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
|
||||
.cached_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.unwrap_or(0);
|
||||
let cache_write_tokens: u64 = record
|
||||
.cache_write_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.unwrap_or(0);
|
||||
let reasoning_tokens: u64 = record
|
||||
.reasoning_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
@@ -513,50 +598,57 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
reasoning_tokens,
|
||||
total_tokens,
|
||||
input_tokens_display: format_number(input_tokens),
|
||||
output_tokens_display: format_number(output_tokens),
|
||||
cached_tokens_display: format_compact_number(cached_tokens),
|
||||
cache_write_tokens_display: format_compact_number(cache_write_tokens),
|
||||
reasoning_tokens_display: format_number(reasoning_tokens),
|
||||
total_tokens_display: format_number(total_tokens),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_cost_details(
|
||||
pub(crate) fn usage_billing(
|
||||
record: &UsageRecord,
|
||||
upstream_model: Option<&str>,
|
||||
tokens: &UsageRecordTokenDetailsData,
|
||||
) -> Option<UsageRecordCostDetailsData> {
|
||||
let cost = usage_record_cost_details(
|
||||
) -> Option<UsageRecordBillingData> {
|
||||
let billing = usage_record_billing(
|
||||
record,
|
||||
upstream_model,
|
||||
tokens.input_tokens,
|
||||
tokens.output_tokens,
|
||||
tokens.cached_tokens,
|
||||
tokens.cache_write_tokens,
|
||||
)?;
|
||||
|
||||
Some(UsageRecordCostDetailsData {
|
||||
input_cost: cost.input_cost,
|
||||
output_cost: cost.output_cost,
|
||||
cache_read_cost: cost.cache_read_cost,
|
||||
total_cost: cost.total_cost,
|
||||
original_cost: cost.original_cost,
|
||||
input_price_per_mtoken: cost.input_price_per_mtoken,
|
||||
output_price_per_mtoken: cost.output_price_per_mtoken,
|
||||
cache_read_price_per_mtoken: cost.cache_read_price_per_mtoken,
|
||||
service_tier: cost.service_tier,
|
||||
service_tier_display: cost.service_tier_display,
|
||||
multiplier: cost.multiplier,
|
||||
input_cost_display: format_cost(cost.input_cost),
|
||||
output_cost_display: format_cost(cost.output_cost),
|
||||
cache_read_cost_display: format_cost(cost.cache_read_cost),
|
||||
total_cost_display: format_cost(cost.total_cost),
|
||||
original_cost_display: format_cost(cost.original_cost),
|
||||
input_price_display: format_token_price(cost.input_price_per_mtoken),
|
||||
output_price_display: format_token_price(cost.output_price_per_mtoken),
|
||||
cache_read_price_display: format_token_price(cost.cache_read_price_per_mtoken),
|
||||
multiplier_display: format_multiplier(cost.multiplier),
|
||||
Some(UsageRecordBillingData {
|
||||
input_amount: billing.input_amount,
|
||||
output_amount: billing.output_amount,
|
||||
cache_read_amount: billing.cache_read_amount,
|
||||
cache_write_amount: billing.cache_write_amount,
|
||||
standard_amount: billing.standard_amount,
|
||||
total_amount: billing.total_amount,
|
||||
input_price_per_mtoken: billing.input_price_per_mtoken,
|
||||
output_price_per_mtoken: billing.output_price_per_mtoken,
|
||||
cache_read_price_per_mtoken: billing.cache_read_price_per_mtoken,
|
||||
cache_write_price_per_mtoken: billing.cache_write_price_per_mtoken,
|
||||
service_tier: billing.service_tier,
|
||||
service_tier_display: billing.service_tier_display,
|
||||
multiplier: billing.multiplier,
|
||||
input_amount_display: format_billing_amount(billing.input_amount),
|
||||
output_amount_display: format_billing_amount(billing.output_amount),
|
||||
cache_read_amount_display: format_billing_amount(billing.cache_read_amount),
|
||||
cache_write_amount_display: format_billing_amount(billing.cache_write_amount),
|
||||
standard_amount_display: format_billing_amount(billing.standard_amount),
|
||||
total_amount_display: format_billing_amount(billing.total_amount),
|
||||
input_price_display: format_token_price(billing.input_price_per_mtoken),
|
||||
output_price_display: format_token_price(billing.output_price_per_mtoken),
|
||||
cache_read_price_display: format_token_price(billing.cache_read_price_per_mtoken),
|
||||
cache_write_price_display: format_token_price(billing.cache_write_price_per_mtoken),
|
||||
multiplier_display: format_multiplier(billing.multiplier),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -567,6 +659,7 @@ impl From<UsageRecordSummary> for UsageRecordSummaryData {
|
||||
input_tokens: format_tokens(summary.input_tokens),
|
||||
output_tokens: format_tokens(summary.output_tokens),
|
||||
cached_tokens: format_tokens(summary.cached_tokens),
|
||||
cache_write_tokens: format_tokens(summary.cache_write_tokens),
|
||||
total_tokens: format_tokens(summary.total_tokens),
|
||||
average_latency_ms: format_duration_ms_f64(summary.average_latency_ms),
|
||||
}
|
||||
@@ -588,13 +681,14 @@ impl UsageRecordBreakdownData {
|
||||
input_tokens: format_tokens(item.input_tokens),
|
||||
output_tokens: format_tokens(item.output_tokens),
|
||||
cached_tokens: format_tokens(item.cached_tokens),
|
||||
cache_write_tokens: format_tokens(item.cache_write_tokens),
|
||||
total_tokens: format_tokens(item.total_tokens),
|
||||
total_tokens_value: item.total_tokens,
|
||||
total_tokens_total: format_tokens(total_tokens_total),
|
||||
total_tokens_total_value: total_tokens_total,
|
||||
cost: format_cost(item.cost),
|
||||
actual_cost: format_cost(item.actual_cost),
|
||||
account_cost: format_cost(item.account_cost),
|
||||
standard_billing_amount: format_billing_amount(item.standard_billing_amount),
|
||||
actual_billing_amount: format_billing_amount(item.actual_billing_amount),
|
||||
account_billing_amount: format_billing_amount(item.account_billing_amount),
|
||||
average_latency_ms: format_duration_ms_f64(item.average_latency_ms),
|
||||
}
|
||||
}
|
||||
@@ -613,9 +707,9 @@ fn usage_record_breakdown_items(items: Vec<UsageRecordBreakdown>) -> Vec<UsageRe
|
||||
|
||||
impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
|
||||
fn from(point: UsageRecordTrendPoint) -> Self {
|
||||
let prompt_tokens = point.input_tokens.saturating_add(point.cached_tokens);
|
||||
let cache_hit_rate_value = if prompt_tokens > 0 {
|
||||
(point.cached_tokens as f64 / prompt_tokens as f64 * 100.0).round()
|
||||
let cache_hit_rate_value = if point.input_tokens > 0 {
|
||||
(point.cached_tokens.min(point.input_tokens) as f64 / point.input_tokens as f64 * 100.0)
|
||||
.round()
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
@@ -627,12 +721,14 @@ impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
|
||||
output_tokens_value: point.output_tokens,
|
||||
cached_tokens: format_tokens(point.cached_tokens),
|
||||
cached_tokens_value: point.cached_tokens,
|
||||
cache_write_tokens: format_tokens(point.cache_write_tokens),
|
||||
cache_write_tokens_value: point.cache_write_tokens,
|
||||
total_tokens: format_tokens(point.total_tokens),
|
||||
total_tokens_value: point.total_tokens,
|
||||
cache_hit_rate: format_percent(cache_hit_rate_value),
|
||||
cache_hit_rate_value,
|
||||
cost: format_cost(point.cost),
|
||||
actual_cost: format_cost(point.actual_cost),
|
||||
standard_billing_amount: format_billing_amount(point.standard_billing_amount),
|
||||
actual_billing_amount: format_billing_amount(point.actual_billing_amount),
|
||||
average_latency_ms: format_duration_ms_f64(point.average_latency_ms),
|
||||
average_latency_ms_value: point.average_latency_ms,
|
||||
}
|
||||
|
||||
@@ -4,10 +4,10 @@ use std::{
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tower::service_fn;
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
//!
|
||||
//! 提供 Bearer token 提取与客户端 API key 校验。
|
||||
|
||||
use axum::http::{header::AUTHORIZATION, HeaderMap};
|
||||
use axum::http::{HeaderMap, header::AUTHORIZATION};
|
||||
|
||||
use crate::api::AppState;
|
||||
|
||||
@@ -93,13 +93,13 @@ fn log_client_api_key_auth_failure(error: &ClientApiKeyAuthError) {
|
||||
tracing::warn!(
|
||||
auth_failure = error.reason(),
|
||||
error = %message,
|
||||
"client api key authorization failed"
|
||||
"Client API key authorization failed"
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
tracing::info!(
|
||||
auth_failure = error.reason(),
|
||||
"client api key authorization failed"
|
||||
"Client API key authorization failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,18 @@
|
||||
//! OpenAI 路由共享错误响应。
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{
|
||||
dispatch::errors::{DispatchFailureClass, ResponseDispatchError},
|
||||
upstream::openai::protocol::sse::response_failed_sse_event as encode_response_failed_sse_event,
|
||||
upstream::openai::protocol::{
|
||||
responses::ResponsesSseFailure,
|
||||
sse::response_failed_sse_event as encode_response_failed_sse_event,
|
||||
},
|
||||
upstream::openai::transport::CodexClientError,
|
||||
};
|
||||
|
||||
@@ -47,6 +50,13 @@ struct OpenAiErrorKind {
|
||||
code: &'static str,
|
||||
}
|
||||
|
||||
struct EncodedClientFailure<'a> {
|
||||
status: StatusCode,
|
||||
kind: OpenAiErrorKind,
|
||||
code: &'a str,
|
||||
message: &'a str,
|
||||
}
|
||||
|
||||
const NO_AVAILABLE_ACCOUNTS_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
error_type: "server_error",
|
||||
code: "no_available_accounts",
|
||||
@@ -73,7 +83,7 @@ const MODEL_NOT_FOUND_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
};
|
||||
const HISTORY_UNAVAILABLE_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
error_type: "invalid_request_error",
|
||||
code: "previous_response_unavailable",
|
||||
code: "previous_response_not_found",
|
||||
};
|
||||
const CONTINUATION_BUSY_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
error_type: "server_error",
|
||||
@@ -117,12 +127,20 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon
|
||||
|
||||
pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> Response {
|
||||
if let ResponseDispatchError::Upstream(CodexClientError::Upstream { status, body, .. }) = error
|
||||
&& status.is_client_error()
|
||||
&& let Ok(body) = serde_json::from_str::<Value>(body)
|
||||
{
|
||||
if status.is_client_error() {
|
||||
if let Ok(body) = serde_json::from_str::<Value>(body) {
|
||||
return (*status, Json(body)).into_response();
|
||||
}
|
||||
}
|
||||
return (*status, Json(body)).into_response();
|
||||
}
|
||||
|
||||
if let Some(failure) = response_dispatch_client_failure(error) {
|
||||
return openai_error_response(
|
||||
failure.status,
|
||||
failure.message,
|
||||
failure.kind.error_type,
|
||||
failure.code,
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
match error {
|
||||
@@ -134,12 +152,101 @@ pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> R
|
||||
}
|
||||
|
||||
pub fn responses_stream_dispatch_failed_sse_event(error: &ResponseDispatchError) -> String {
|
||||
if let Some(failure) = response_dispatch_client_failure(error) {
|
||||
return encode_response_failed_sse_event(
|
||||
failure.kind.error_type,
|
||||
failure.code,
|
||||
failure.message,
|
||||
);
|
||||
}
|
||||
let http_error =
|
||||
response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream);
|
||||
let kind = response_dispatch_error_kind(error);
|
||||
responses_failed_sse_event(kind, &http_error.message)
|
||||
}
|
||||
|
||||
pub(super) fn responses_websocket_dispatch_error_event(
|
||||
error: &ResponseDispatchError,
|
||||
request_id: &str,
|
||||
) -> String {
|
||||
if let Some(failure) = response_dispatch_client_failure(error) {
|
||||
return responses_websocket_error_event(
|
||||
failure.status,
|
||||
failure.kind.error_type,
|
||||
failure.code,
|
||||
failure.message,
|
||||
request_id,
|
||||
);
|
||||
}
|
||||
if let Some(event) = upstream_client_websocket_error_event(error, request_id) {
|
||||
return event;
|
||||
}
|
||||
let http_error =
|
||||
response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream);
|
||||
let kind = response_dispatch_error_kind(error);
|
||||
responses_websocket_error_event(
|
||||
http_error.status,
|
||||
kind.error_type,
|
||||
kind.code,
|
||||
&http_error.message,
|
||||
request_id,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn responses_websocket_error_event(
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
request_id: &str,
|
||||
) -> String {
|
||||
json!({
|
||||
"type": "error",
|
||||
"status": status.as_u16(),
|
||||
"error": {
|
||||
"type": error_type,
|
||||
"code": code,
|
||||
"message": message,
|
||||
},
|
||||
"headers": {
|
||||
"x-request-id": request_id,
|
||||
},
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn upstream_client_websocket_error_event(
|
||||
error: &ResponseDispatchError,
|
||||
request_id: &str,
|
||||
) -> Option<String> {
|
||||
let ResponseDispatchError::Upstream(CodexClientError::Upstream { status, body, .. }) = error
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
if !status.is_client_error() {
|
||||
return None;
|
||||
}
|
||||
let Value::Object(mut payload) = serde_json::from_str::<Value>(body).ok()? else {
|
||||
return None;
|
||||
};
|
||||
if !payload.get("error").is_some_and(Value::is_object) {
|
||||
return None;
|
||||
}
|
||||
payload.insert("type".to_string(), Value::String("error".to_string()));
|
||||
payload.insert("status".to_string(), Value::Number(status.as_u16().into()));
|
||||
let headers = payload
|
||||
.entry("headers".to_string())
|
||||
.or_insert_with(|| Value::Object(serde_json::Map::new()));
|
||||
if !headers.is_object() {
|
||||
*headers = Value::Object(serde_json::Map::new());
|
||||
}
|
||||
headers.as_object_mut()?.insert(
|
||||
"x-request-id".to_string(),
|
||||
Value::String(request_id.to_string()),
|
||||
);
|
||||
Some(Value::Object(payload).to_string())
|
||||
}
|
||||
|
||||
fn responses_no_available_accounts_response() -> (StatusCode, Json<Value>) {
|
||||
let kind = NO_AVAILABLE_ACCOUNTS_ERROR;
|
||||
(
|
||||
@@ -242,6 +349,33 @@ fn responses_error_kind_for_status(status: StatusCode) -> OpenAiErrorKind {
|
||||
}
|
||||
}
|
||||
|
||||
fn response_dispatch_client_failure(
|
||||
error: &ResponseDispatchError,
|
||||
) -> Option<EncodedClientFailure<'_>> {
|
||||
let client_failure = match error {
|
||||
ResponseDispatchError::Failed(failure) => failure,
|
||||
_ => return None,
|
||||
};
|
||||
let failure = client_failure.exposed_upstream()?;
|
||||
Some(encoded_client_failure(
|
||||
failure,
|
||||
status_from_u16(client_failure.status_code(), StatusCode::BAD_GATEWAY),
|
||||
))
|
||||
}
|
||||
|
||||
fn encoded_client_failure(
|
||||
failure: &ResponsesSseFailure,
|
||||
status: StatusCode,
|
||||
) -> EncodedClientFailure<'_> {
|
||||
let kind = responses_error_kind_for_status(status);
|
||||
EncodedClientFailure {
|
||||
status,
|
||||
kind,
|
||||
code: failure.upstream_code.as_deref().unwrap_or(kind.code),
|
||||
message: &failure.message,
|
||||
}
|
||||
}
|
||||
|
||||
fn response_dispatch_message(
|
||||
error: &ResponseDispatchError,
|
||||
style: ResponseDispatchMessageStyle,
|
||||
|
||||
@@ -5,4 +5,3 @@ pub mod errors;
|
||||
pub mod models;
|
||||
pub mod responses;
|
||||
pub mod router;
|
||||
pub mod sse;
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
//! OpenAI 模型列表以及模型处理器。
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{
|
||||
api::client::auth::authorize_client_api_key, api::AppState, models::catalog::ModelCatalog,
|
||||
api::AppState, api::client::auth::authorize_client_api_key, models::catalog::ModelCatalog,
|
||||
};
|
||||
|
||||
use super::errors::{missing_client_api_key_response, model_not_found_response};
|
||||
|
||||
@@ -1,25 +1,34 @@
|
||||
//! OpenAI Responses API 透明代理入口:解析原始 JSON body,提取调度元数据,
|
||||
//! 做最小 patch 后交给调度层,请求语义原样透传上游。
|
||||
|
||||
mod sse;
|
||||
mod websocket;
|
||||
|
||||
pub use websocket::responses_websocket;
|
||||
|
||||
use axum::{
|
||||
Extension, Json,
|
||||
body::{Body, Bytes},
|
||||
extract::State,
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
Extension, Json,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
api::middleware::request_id::{ClientIp, RequestId},
|
||||
api::AppState,
|
||||
api::middleware::request_id::{ClientIp, RequestId},
|
||||
dispatch::{errors::ResponseDispatchError, service::ResponseDispatchStream},
|
||||
upstream::openai::protocol::{
|
||||
responses::{CodexCompactRequest, CodexResponsesRequest},
|
||||
responses::{
|
||||
CodexResponsesRequest, WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY,
|
||||
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
|
||||
},
|
||||
sse::sse_body_has_done,
|
||||
},
|
||||
};
|
||||
|
||||
use self::sse::{done_sse_frame, event_stream_response as sse_event_stream_response};
|
||||
use super::{
|
||||
errors::{
|
||||
invalid_responses_request_response, missing_client_api_key_response,
|
||||
@@ -27,7 +36,6 @@ use super::{
|
||||
responses_dispatch_error_response_ref, responses_stream_dispatch_failed_sse_event,
|
||||
},
|
||||
models::model_catalog_for_state,
|
||||
sse::{done_sse_frame, event_stream_response as sse_event_stream_response},
|
||||
};
|
||||
|
||||
const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
|
||||
@@ -101,32 +109,54 @@ pub fn build_codex_request(
|
||||
None => {}
|
||||
}
|
||||
apply_context_header_fields(&mut request, headers);
|
||||
apply_identity_fields(&mut request, headers);
|
||||
apply_request_semantic_headers(&mut request, headers);
|
||||
apply_session_fields(&mut request, headers);
|
||||
request.client_user_agent = user_agent_from_headers(headers);
|
||||
request
|
||||
}
|
||||
|
||||
fn apply_identity_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
|
||||
request.client_conversation_id = request.prompt_cache_key().map(ToString::to_string);
|
||||
request.client_session_id = identity_context_string(request.body(), "session_id")
|
||||
.or_else(|| header_string(headers, "session-id"));
|
||||
request.client_thread_id = identity_context_string(request.body(), "thread_id")
|
||||
fn apply_request_semantic_headers(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
|
||||
request.responses_lite = header_string(headers, X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER)
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_metadata()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| {
|
||||
metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY)
|
||||
})
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
});
|
||||
request.memgen_request = header_string(headers, X_OPENAI_MEMGEN_REQUEST_HEADER);
|
||||
}
|
||||
|
||||
fn apply_session_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
|
||||
request.client_conversation_id = session_context_string(request.body(), "conversation_id")
|
||||
.or_else(|| header_string(headers, "conversation-id"))
|
||||
.or_else(|| header_string(headers, "conversation_id"))
|
||||
.or_else(|| request.prompt_cache_key().map(ToString::to_string));
|
||||
request.client_session_id = session_context_string(request.body(), "session_id")
|
||||
.or_else(|| header_string(headers, "session-id"))
|
||||
.or_else(|| header_string(headers, "session_id"));
|
||||
request.client_thread_id = session_context_string(request.body(), "thread_id")
|
||||
.or_else(|| header_string(headers, "thread-id"));
|
||||
request.client_request_id = identity_context_string(request.body(), "x-client-request-id")
|
||||
request.client_request_id = session_context_string(request.body(), "x-client-request-id")
|
||||
.or_else(|| header_string(headers, "x-client-request-id"));
|
||||
request.client_turn_id = identity_context_string(request.body(), "turn_id")
|
||||
request.client_turn_id = session_context_string(request.body(), "turn_id")
|
||||
.or_else(|| header_string(headers, "x-codex-turn-id"));
|
||||
request.codex_window_id = request.codex_window_id.take().or_else(|| {
|
||||
identity_context_string(request.body(), "x-codex-window-id")
|
||||
session_context_string(request.body(), "x-codex-window-id")
|
||||
.or_else(|| header_string(headers, "x-codex-window-id"))
|
||||
});
|
||||
request.parent_thread_id = request.parent_thread_id.take().or_else(|| {
|
||||
identity_context_string(request.body(), "x-codex-parent-thread-id")
|
||||
session_context_string(request.body(), "x-codex-parent-thread-id")
|
||||
.or_else(|| header_string(headers, "x-codex-parent-thread-id"))
|
||||
});
|
||||
}
|
||||
|
||||
fn identity_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
|
||||
fn session_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
|
||||
body_context_string(body, key).or_else(|| {
|
||||
body.get("client_metadata")
|
||||
.and_then(Value::as_object)
|
||||
@@ -138,49 +168,6 @@ fn identity_context_string(body: &Map<String, Value>, key: &str) -> Option<Strin
|
||||
})
|
||||
}
|
||||
|
||||
/// 从客户端原始 Responses JSON body 构造 Codex compact 请求。
|
||||
///
|
||||
/// compact 端点只返回一次性 JSON(代理按非流式读取),故剥离 `stream` 这一
|
||||
/// transport 控制字段,避免上游返回代理无法解析的 SSE 形态。业务语义字段
|
||||
/// (`reasoning`/`text`/`store`/`prompt_cache_key`/`previous_response_id`/
|
||||
/// `include`/`client_metadata`/未知字段)一律原样透传。
|
||||
pub fn build_compact_request(
|
||||
mut body: Map<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
) -> CodexCompactRequest {
|
||||
let client_session_id = identity_context_string(&body, "session_id")
|
||||
.or_else(|| header_string(headers, "session-id"));
|
||||
let client_thread_id =
|
||||
identity_context_string(&body, "thread_id").or_else(|| header_string(headers, "thread-id"));
|
||||
let client_request_id = identity_context_string(&body, "x-client-request-id")
|
||||
.or_else(|| header_string(headers, "x-client-request-id"));
|
||||
let client_turn_id = identity_context_string(&body, "turn_id")
|
||||
.or_else(|| header_string(headers, "x-codex-turn-id"));
|
||||
let client_window_id = identity_context_string(&body, "x-codex-window-id")
|
||||
.or_else(|| header_string(headers, "x-codex-window-id"));
|
||||
let client_parent_thread_id = identity_context_string(&body, "x-codex-parent-thread-id")
|
||||
.or_else(|| header_string(headers, "x-codex-parent-thread-id"));
|
||||
body.remove(COMPACT_STREAM_KEY);
|
||||
for key in TRANSPORT_ONLY_KEYS {
|
||||
body.remove(key);
|
||||
}
|
||||
CodexCompactRequest {
|
||||
body,
|
||||
client_ip: None,
|
||||
client_user_agent: user_agent_from_headers(headers),
|
||||
client_api_key_id: None,
|
||||
client_session_id,
|
||||
client_thread_id,
|
||||
client_request_id,
|
||||
client_turn_id,
|
||||
client_window_id,
|
||||
client_parent_thread_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// compact 端点不支持流式,转发前剥离 `stream`(transport 控制字段,非业务语义)。
|
||||
const COMPACT_STREAM_KEY: &str = "stream";
|
||||
|
||||
fn inject_subagent_metadata(body: &mut Map<String, Value>, subagent: &str) {
|
||||
let metadata = body
|
||||
.entry("client_metadata")
|
||||
@@ -286,24 +273,6 @@ pub async fn review_responses(
|
||||
.await
|
||||
}
|
||||
|
||||
/// `POST /v1/responses/compact`
|
||||
pub async fn compact_responses(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
client_ip: Option<Extension<ClientIp>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
handle_compact_responses(
|
||||
state,
|
||||
request_id,
|
||||
client_ip.map(|Extension(client_ip)| client_ip.as_str().to_string()),
|
||||
headers,
|
||||
body,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn handle_responses(
|
||||
state: AppState,
|
||||
request_id: RequestId,
|
||||
@@ -322,29 +291,39 @@ async fn handle_responses(
|
||||
let Some(body) = parse_responses_body(&body) else {
|
||||
return invalid_responses_request_response().into_response();
|
||||
};
|
||||
let model = match validated_responses_model(&state, request_model(&body)).await {
|
||||
Ok(model) => model,
|
||||
Err(ResponsesModelValidationError::InvalidRequest) => {
|
||||
let prepared = match prepare_responses_request(
|
||||
&state,
|
||||
body,
|
||||
&headers,
|
||||
forced_subagent,
|
||||
client_ip,
|
||||
client_api_key_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(ResponsesRequestValidationError::InvalidRequest) => {
|
||||
return invalid_responses_request_response().into_response();
|
||||
}
|
||||
Err(ResponsesModelValidationError::ModelNotFound) => {
|
||||
Err(ResponsesRequestValidationError::ModelNotFound) => {
|
||||
return model_not_found_response().into_response();
|
||||
}
|
||||
};
|
||||
let mut codex_request = build_codex_request(body, &headers, forced_subagent);
|
||||
codex_request.client_ip = client_ip;
|
||||
codex_request.client_api_key_id = Some(client_api_key_id);
|
||||
let PreparedResponsesRequest {
|
||||
mut request,
|
||||
requested_model,
|
||||
} = prepared;
|
||||
// 客户端未显式指定 transport 时默认偏好 WebSocket。
|
||||
if !codex_request.force_http_sse {
|
||||
codex_request.use_websocket = true;
|
||||
if !request.force_http_sse {
|
||||
request.use_websocket = true;
|
||||
}
|
||||
let stream = codex_request.stream();
|
||||
let stream = request.stream();
|
||||
|
||||
if stream {
|
||||
return match state
|
||||
.services
|
||||
.responses
|
||||
.stream(request_id.as_str(), route, codex_request, &model)
|
||||
.stream(request_id.as_str(), route, request, &requested_model)
|
||||
.await
|
||||
{
|
||||
Ok(stream) => live_event_stream_response(stream),
|
||||
@@ -355,50 +334,7 @@ async fn handle_responses(
|
||||
match state
|
||||
.services
|
||||
.responses
|
||||
.complete(request_id.as_str(), route, codex_request, &model)
|
||||
.await
|
||||
{
|
||||
Ok(result) => apply_safe_response_headers(
|
||||
(StatusCode::OK, Json(result.body)).into_response(),
|
||||
result.response_headers,
|
||||
),
|
||||
Err(error) => responses_dispatch_error_response(error),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_compact_responses(
|
||||
state: AppState,
|
||||
request_id: RequestId,
|
||||
client_ip: Option<String>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let Some(client_api_key_id) =
|
||||
crate::api::client::auth::authorized_client_api_key_id(&state, &headers).await
|
||||
else {
|
||||
return missing_client_api_key_response().into_response();
|
||||
};
|
||||
|
||||
let Some(body) = parse_responses_body(&body) else {
|
||||
return invalid_responses_request_response().into_response();
|
||||
};
|
||||
let model = match validated_responses_model(&state, request_model(&body)).await {
|
||||
Ok(model) => model,
|
||||
Err(ResponsesModelValidationError::InvalidRequest) => {
|
||||
return invalid_responses_request_response().into_response();
|
||||
}
|
||||
Err(ResponsesModelValidationError::ModelNotFound) => {
|
||||
return model_not_found_response().into_response();
|
||||
}
|
||||
};
|
||||
let mut compact_request = build_compact_request(body, &headers);
|
||||
compact_request.client_ip = client_ip;
|
||||
compact_request.client_api_key_id = Some(client_api_key_id);
|
||||
|
||||
match state
|
||||
.services
|
||||
.responses
|
||||
.compact(request_id.as_str(), compact_request, &model)
|
||||
.complete(request_id.as_str(), route, request, &requested_model)
|
||||
.await
|
||||
{
|
||||
Ok(result) => apply_safe_response_headers(
|
||||
@@ -423,23 +359,47 @@ fn request_model(body: &Map<String, Value>) -> &str {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(super) struct PreparedResponsesRequest {
|
||||
pub(super) request: CodexResponsesRequest,
|
||||
pub(super) requested_model: String,
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_responses_request(
|
||||
state: &AppState,
|
||||
body: Map<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
forced_subagent: Option<&str>,
|
||||
client_ip: Option<String>,
|
||||
client_api_key_id: String,
|
||||
) -> Result<PreparedResponsesRequest, ResponsesRequestValidationError> {
|
||||
let requested_model = validated_responses_model(state, request_model(&body)).await?;
|
||||
let mut request = build_codex_request(body, headers, forced_subagent);
|
||||
request.client_ip = client_ip;
|
||||
request.client_api_key_id = Some(client_api_key_id);
|
||||
Ok(PreparedResponsesRequest {
|
||||
request,
|
||||
requested_model,
|
||||
})
|
||||
}
|
||||
|
||||
async fn validated_responses_model(
|
||||
state: &AppState,
|
||||
raw_model: &str,
|
||||
) -> Result<String, ResponsesModelValidationError> {
|
||||
) -> Result<String, ResponsesRequestValidationError> {
|
||||
let model = raw_model.trim();
|
||||
if model.is_empty() {
|
||||
return Err(ResponsesModelValidationError::InvalidRequest);
|
||||
return Err(ResponsesRequestValidationError::InvalidRequest);
|
||||
}
|
||||
let catalog = model_catalog_for_state(state).await;
|
||||
if catalog.is_recognized_model_name(model) {
|
||||
if catalog.is_empty() || catalog.is_recognized_model_name(model) {
|
||||
Ok(model.to_string())
|
||||
} else {
|
||||
Err(ResponsesModelValidationError::ModelNotFound)
|
||||
Err(ResponsesRequestValidationError::ModelNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
enum ResponsesModelValidationError {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum ResponsesRequestValidationError {
|
||||
InvalidRequest,
|
||||
ModelNotFound,
|
||||
}
|
||||
@@ -449,15 +409,11 @@ pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
|
||||
}
|
||||
|
||||
fn openai_subagent_from_headers(headers: &HeaderMap) -> Option<String> {
|
||||
header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| {
|
||||
if matches!(
|
||||
header_string(headers, OPENAI_SUBAGENT_HEADER).filter(|value| {
|
||||
matches!(
|
||||
value.as_str(),
|
||||
"review" | "compact" | "memory_consolidation" | "collab_spawn"
|
||||
) {
|
||||
Some(value)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,15 +5,15 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{
|
||||
header::{CACHE_CONTROL, CONTENT_TYPE},
|
||||
StatusCode,
|
||||
header::{CACHE_CONTROL, CONTENT_TYPE},
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
|
||||
use crate::upstream::openai::protocol::sse::DONE_SSE_FRAME;
|
||||
|
||||
use super::errors::openai_error_response;
|
||||
use crate::api::client::errors::openai_error_response;
|
||||
|
||||
/// 编码 OpenAI 流结束标记。
|
||||
pub fn done_sse_frame() -> &'static str {
|
||||
@@ -0,0 +1,596 @@
|
||||
//! 官方 Responses WebSocket 入站协议。
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Request, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use hyper::upgrade::Upgraded;
|
||||
use hyper_util::rt::TokioIo;
|
||||
use serde_json::{Map, Value, json};
|
||||
use thiserror::Error;
|
||||
use tokio_tungstenite::{
|
||||
WebSocketStream,
|
||||
tungstenite::{
|
||||
Error as WebSocketError, Message,
|
||||
handshake::server::create_response_with_body,
|
||||
protocol::{CloseFrame, Role, WebSocketConfig, frame::coding::CloseCode},
|
||||
},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
api::{AppState, middleware::request_id::RequestId},
|
||||
dispatch::{
|
||||
service::{ConnectionTranscriptFacts, ResponseDispatchStream},
|
||||
transport::canonical::CanonicalResponseEvent,
|
||||
},
|
||||
upstream::openai::protocol::websocket::is_terminal_websocket_event,
|
||||
};
|
||||
|
||||
use super::{PreparedResponsesRequest, ResponsesRequestValidationError, prepare_responses_request};
|
||||
use crate::api::client::{
|
||||
auth::authorized_client_api_key_id,
|
||||
errors::{
|
||||
missing_client_api_key_response, openai_error_response,
|
||||
responses_websocket_dispatch_error_event, responses_websocket_error_event,
|
||||
},
|
||||
router::MAX_CLIENT_REQUEST_BODY_BYTES,
|
||||
};
|
||||
|
||||
const RESPONSES_ROUTE: &str = "/v1/responses";
|
||||
const INVALID_REQUEST_MESSAGE: &str = "Invalid responses request";
|
||||
const MODEL_NOT_FOUND_MESSAGE: &str = "Model not found";
|
||||
const INVALID_STREAM_MESSAGE: &str = "Responses WebSocket requests require stream=true";
|
||||
const INVALID_UPSTREAM_EVENT_MESSAGE: &str = "Invalid upstream Responses event";
|
||||
const STREAM_DISCONNECTED_MESSAGE: &str = "Upstream stream closed before response.completed";
|
||||
|
||||
type ResponsesWebSocket = WebSocketStream<TokioIo<Upgraded>>;
|
||||
|
||||
/// 将已认证的 `GET /v1/responses` 请求升级为官方 Responses WebSocket 协议。
|
||||
pub async fn responses_websocket(State(state): State<AppState>, mut request: Request) -> Response {
|
||||
let headers = request.headers().clone();
|
||||
let Some(client_api_key_id) = authorized_client_api_key_id(&state, &headers).await else {
|
||||
return missing_client_api_key_response().into_response();
|
||||
};
|
||||
let connection_request_id = request
|
||||
.extensions()
|
||||
.get::<RequestId>()
|
||||
.cloned()
|
||||
.unwrap_or_else(RequestId::generate);
|
||||
let client_ip = request
|
||||
.extensions()
|
||||
.get::<crate::api::middleware::request_id::ClientIp>()
|
||||
.map(|client_ip| client_ip.as_str().to_string());
|
||||
let connection_id = connection_request_id.as_str().to_string();
|
||||
let response = match create_response_with_body(&request, Body::empty) {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %connection_id,
|
||||
error = %error,
|
||||
"Invalid Responses WebSocket upgrade request"
|
||||
);
|
||||
return openai_error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Invalid Responses WebSocket upgrade request",
|
||||
"invalid_request_error",
|
||||
"invalid_websocket_upgrade",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
let upgrade = hyper::upgrade::on(&mut request);
|
||||
let connection_drain = state.services.connection_drain.clone();
|
||||
connection_drain.spawn(async move {
|
||||
match upgrade.await {
|
||||
Ok(upgraded) => {
|
||||
let socket = WebSocketStream::from_raw_socket(
|
||||
TokioIo::new(upgraded),
|
||||
Role::Server,
|
||||
Some(downstream_websocket_config()),
|
||||
)
|
||||
.await;
|
||||
serve_responses_websocket(
|
||||
socket,
|
||||
ResponsesWebSocketSession {
|
||||
state,
|
||||
headers,
|
||||
client_ip,
|
||||
client_api_key_id,
|
||||
connection_id,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
websocket_connection_id = %connection_id,
|
||||
error = %error,
|
||||
"Responses WebSocket upgrade failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
response
|
||||
}
|
||||
|
||||
fn downstream_websocket_config() -> WebSocketConfig {
|
||||
let mut config = WebSocketConfig::default();
|
||||
config.max_message_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES);
|
||||
config.max_frame_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES);
|
||||
config
|
||||
}
|
||||
|
||||
struct ResponsesWebSocketSession {
|
||||
state: AppState,
|
||||
headers: HeaderMap,
|
||||
client_ip: Option<String>,
|
||||
client_api_key_id: String,
|
||||
connection_id: String,
|
||||
}
|
||||
|
||||
struct ConnectionTranscriptCollector {
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
completed_items: Vec<Value>,
|
||||
completed: Option<ConnectionTranscriptFacts>,
|
||||
}
|
||||
|
||||
enum ForwardDispatchOutcome {
|
||||
Continue(Option<ConnectionTranscriptFacts>),
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
impl ConnectionTranscriptCollector {
|
||||
fn new(
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
completed_items: Vec::new(),
|
||||
completed: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe(&mut self, event: &CanonicalResponseEvent) {
|
||||
match event.event_type() {
|
||||
Some("response.output_item.done") => {
|
||||
if let Some(item) = event.data().get("item") {
|
||||
self.completed_items.push(item.clone());
|
||||
}
|
||||
}
|
||||
Some("response.completed") => {
|
||||
let Some(response_id) = event
|
||||
.data()
|
||||
.pointer("/response/id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let output = event
|
||||
.data()
|
||||
.pointer("/response/output")
|
||||
.and_then(Value::as_array)
|
||||
.filter(|output| !output.is_empty())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| self.completed_items.clone());
|
||||
self.completed = Some(ConnectionTranscriptFacts::new(
|
||||
self.account_id.clone(),
|
||||
self.request_input.clone(),
|
||||
self.continued_from_previous_response,
|
||||
response_id,
|
||||
output,
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve_responses_websocket(
|
||||
mut socket: ResponsesWebSocket,
|
||||
session: ResponsesWebSocketSession,
|
||||
) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
"Responses WebSocket connected"
|
||||
);
|
||||
let mut request_count = 0u64;
|
||||
let mut replay_snapshot = session
|
||||
.state
|
||||
.services
|
||||
.responses
|
||||
.connection_replay_snapshot();
|
||||
|
||||
while let Some(message) = socket.next().await {
|
||||
let payload = match message {
|
||||
Ok(Message::Text(payload)) => payload,
|
||||
Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue,
|
||||
Ok(Message::Close(_)) => break,
|
||||
Ok(Message::Binary(_)) => {
|
||||
close_for_protocol_violation(
|
||||
&mut socket,
|
||||
"Responses WebSocket accepts text frames only",
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
error = %error,
|
||||
"Responses WebSocket receive failed"
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
request_count = request_count.saturating_add(1);
|
||||
let request_id = RequestId::generate();
|
||||
let body = match parse_response_create(payload.as_str()) {
|
||||
Ok(body) => body,
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
request_id = %request_id.as_str(),
|
||||
error = %error,
|
||||
"Responses WebSocket request rejected"
|
||||
);
|
||||
if !send_error(
|
||||
&mut socket,
|
||||
StatusCode::BAD_REQUEST,
|
||||
"invalid_request_error",
|
||||
"invalid_request",
|
||||
&error.to_string(),
|
||||
request_id.as_str(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let mut prepared = match prepare_responses_request(
|
||||
&session.state,
|
||||
body,
|
||||
&session.headers,
|
||||
None,
|
||||
session.client_ip.clone(),
|
||||
session.client_api_key_id.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
let (status, code, message) = match error {
|
||||
ResponsesRequestValidationError::InvalidRequest => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"invalid_request",
|
||||
INVALID_REQUEST_MESSAGE,
|
||||
),
|
||||
ResponsesRequestValidationError::ModelNotFound => (
|
||||
StatusCode::NOT_FOUND,
|
||||
"model_not_found",
|
||||
MODEL_NOT_FOUND_MESSAGE,
|
||||
),
|
||||
};
|
||||
if !send_error(
|
||||
&mut socket,
|
||||
status,
|
||||
"invalid_request_error",
|
||||
code,
|
||||
message,
|
||||
request_id.as_str(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let replay_plan = session
|
||||
.state
|
||||
.services
|
||||
.responses
|
||||
.prepare_connection_replay(&replay_snapshot, &mut prepared.request);
|
||||
|
||||
match dispatch_response_create(&mut socket, &session, request_id, prepared).await {
|
||||
ForwardDispatchOutcome::Continue(Some(transcript)) => {
|
||||
session.state.services.responses.commit_connection_replay(
|
||||
&mut replay_snapshot,
|
||||
replay_plan,
|
||||
transcript,
|
||||
);
|
||||
}
|
||||
ForwardDispatchOutcome::Continue(None) => {}
|
||||
ForwardDispatchOutcome::Disconnect => break,
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
request_count,
|
||||
"Responses WebSocket disconnected"
|
||||
);
|
||||
}
|
||||
|
||||
async fn dispatch_response_create(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
session: &ResponsesWebSocketSession,
|
||||
request_id: RequestId,
|
||||
prepared: PreparedResponsesRequest,
|
||||
) -> ForwardDispatchOutcome {
|
||||
let PreparedResponsesRequest {
|
||||
mut request,
|
||||
requested_model,
|
||||
} = prepared;
|
||||
if !request.force_http_sse {
|
||||
request.use_websocket = true;
|
||||
}
|
||||
let stream = session
|
||||
.state
|
||||
.services
|
||||
.responses
|
||||
.stream(
|
||||
request_id.as_str(),
|
||||
RESPONSES_ROUTE,
|
||||
request,
|
||||
&requested_model,
|
||||
)
|
||||
.await;
|
||||
let stream = match stream {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
let event = responses_websocket_dispatch_error_event(&error, request_id.as_str());
|
||||
return if send_text(socket, event).await {
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
};
|
||||
forward_dispatch_stream(socket, stream, request_id.as_str()).await
|
||||
}
|
||||
|
||||
async fn forward_dispatch_stream(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
stream: ResponseDispatchStream,
|
||||
request_id: &str,
|
||||
) -> ForwardDispatchOutcome {
|
||||
let ResponseDispatchStream {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
body,
|
||||
mut canonical_events,
|
||||
response_headers,
|
||||
} = stream;
|
||||
if !send_text(
|
||||
socket,
|
||||
response_metadata_event(request_id, &response_headers),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
|
||||
let mut body = body;
|
||||
let mut terminal_sent = false;
|
||||
let mut completed = ConnectionTranscriptCollector::new(
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
);
|
||||
loop {
|
||||
let next = if terminal_sent {
|
||||
ActiveResponseInput::Upstream(body.next().await)
|
||||
} else {
|
||||
tokio::select! {
|
||||
next = body.next() => ActiveResponseInput::Upstream(next),
|
||||
message = socket.next() => ActiveResponseInput::Downstream(message),
|
||||
}
|
||||
};
|
||||
match next {
|
||||
ActiveResponseInput::Upstream(Some(Ok(chunk))) => {
|
||||
let Some(events) = canonical_events.recv().await else {
|
||||
tracing::warn!(
|
||||
request_id,
|
||||
"Canonical Responses event stream closed before body stream"
|
||||
);
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
};
|
||||
drop(chunk);
|
||||
if !forward_canonical_events(
|
||||
socket,
|
||||
events,
|
||||
request_id,
|
||||
&mut terminal_sent,
|
||||
&mut completed,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
ActiveResponseInput::Upstream(Some(Err(error))) => {
|
||||
tracing::warn!(request_id, error = %error, "Responses dispatch stream failed");
|
||||
if terminal_sent {
|
||||
return ForwardDispatchOutcome::Continue(completed.completed);
|
||||
}
|
||||
return if send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
"upstream_error",
|
||||
STREAM_DISCONNECTED_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
ActiveResponseInput::Upstream(None) => {
|
||||
if terminal_sent {
|
||||
return ForwardDispatchOutcome::Continue(completed.completed);
|
||||
}
|
||||
return if send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
"stream_disconnected",
|
||||
STREAM_DISCONNECTED_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Ok(
|
||||
Message::Ping(_) | Message::Pong(_) | Message::Frame(_),
|
||||
))) => {}
|
||||
ActiveResponseInput::Downstream(Some(Ok(Message::Close(_))) | None) => {
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Ok(Message::Text(_) | Message::Binary(_)))) => {
|
||||
close_for_protocol_violation(
|
||||
socket,
|
||||
"Only one response.create may be active per connection",
|
||||
)
|
||||
.await;
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Err(error))) => {
|
||||
tracing::info!(request_id, error = %error, "Responses WebSocket client disconnected during response");
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ActiveResponseInput {
|
||||
Upstream(Option<Result<bytes::Bytes, crate::dispatch::errors::ResponseDispatchStreamError>>),
|
||||
Downstream(Option<Result<Message, WebSocketError>>),
|
||||
}
|
||||
|
||||
async fn forward_canonical_events(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
events: Vec<CanonicalResponseEvent>,
|
||||
request_id: &str,
|
||||
terminal_sent: &mut bool,
|
||||
completed: &mut ConnectionTranscriptCollector,
|
||||
) -> bool {
|
||||
for event in events {
|
||||
let Some(event_type) = event.event_type() else {
|
||||
tracing::warn!(
|
||||
request_id,
|
||||
"Responses SSE event does not contain a valid JSON type"
|
||||
);
|
||||
let _ = send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
"invalid_upstream_response",
|
||||
INVALID_UPSTREAM_EVENT_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await;
|
||||
return false;
|
||||
};
|
||||
completed.observe(&event);
|
||||
if !send_text(socket, event.data().to_string()).await {
|
||||
return false;
|
||||
}
|
||||
*terminal_sent |= is_terminal_websocket_event(event_type);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn response_metadata_event(request_id: &str, response_headers: &[(String, String)]) -> String {
|
||||
let mut headers = response_headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.clone(), Value::String(value.clone())))
|
||||
.collect::<Map<String, Value>>();
|
||||
headers.insert(
|
||||
"x-request-id".to_string(),
|
||||
Value::String(request_id.to_string()),
|
||||
);
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": headers,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn parse_response_create(payload: &str) -> Result<Map<String, Value>, ResponseCreateFrameError> {
|
||||
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)? else {
|
||||
return Err(ResponseCreateFrameError::ExpectedObject);
|
||||
};
|
||||
match body.remove("type") {
|
||||
Some(Value::String(message_type)) if message_type == "response.create" => {}
|
||||
_ => return Err(ResponseCreateFrameError::UnsupportedType),
|
||||
}
|
||||
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
|
||||
return Err(ResponseCreateFrameError::StreamingRequired);
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
enum ResponseCreateFrameError {
|
||||
#[error("Invalid response.create JSON: {0}")]
|
||||
InvalidJson(#[from] serde_json::Error),
|
||||
#[error("A response.create payload must be a JSON object")]
|
||||
ExpectedObject,
|
||||
#[error("Unsupported Responses WebSocket message type")]
|
||||
UnsupportedType,
|
||||
#[error("{INVALID_STREAM_MESSAGE}")]
|
||||
StreamingRequired,
|
||||
}
|
||||
|
||||
async fn send_error(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
request_id: &str,
|
||||
) -> bool {
|
||||
send_text(
|
||||
socket,
|
||||
responses_websocket_error_event(status, error_type, code, message, request_id),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_text(socket: &mut ResponsesWebSocket, payload: String) -> bool {
|
||||
if let Err(error) = socket.send(Message::Text(payload.into())).await {
|
||||
tracing::info!(error = %error, "Responses WebSocket send failed");
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn close_for_protocol_violation(socket: &mut ResponsesWebSocket, reason: &'static str) {
|
||||
let _ = socket
|
||||
.send(Message::Close(Some(CloseFrame {
|
||||
code: CloseCode::Policy,
|
||||
reason: reason.into(),
|
||||
})))
|
||||
.await;
|
||||
}
|
||||
@@ -1,16 +1,16 @@
|
||||
//! OpenAI 兼容 API 路由。
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
extract::DefaultBodyLimit,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
|
||||
use crate::api::AppState;
|
||||
|
||||
use super::{
|
||||
models::{model_catalog, model_detail, model_info, models},
|
||||
responses::{compact_responses, responses, review_responses},
|
||||
responses::{responses, responses_websocket, review_responses},
|
||||
};
|
||||
|
||||
pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024;
|
||||
@@ -18,9 +18,8 @@ pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024;
|
||||
/// 构造 OpenAI 兼容 API 路由。
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/responses", get(responses_websocket).post(responses))
|
||||
.route("/v1/responses/review", post(review_responses))
|
||||
.route("/v1/responses/compact", post(compact_responses))
|
||||
.route("/v1/models", get(models))
|
||||
.route("/v1/models/catalog", get(model_catalog))
|
||||
.route("/v1/models/{model_id}", get(model_detail))
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
//! 进程关闭时统一排空入站 HTTP 流与 WebSocket 任务。
|
||||
|
||||
use std::future::Future;
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Request, State},
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use tokio_util::{sync::CancellationToken, task::TaskTracker};
|
||||
|
||||
/// 入站连接排空控制器。
|
||||
#[derive(Clone)]
|
||||
pub struct ConnectionDrain {
|
||||
cancellation: CancellationToken,
|
||||
connection_tasks: TaskTracker,
|
||||
}
|
||||
|
||||
impl Default for ConnectionDrain {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cancellation: CancellationToken::new(),
|
||||
connection_tasks: TaskTracker::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionDrain {
|
||||
/// 为下层长生命周期执行器提供与连接排空一致的取消信号。
|
||||
pub(crate) fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancellation.clone()
|
||||
}
|
||||
|
||||
/// 注册需要随进程关闭而终止的长连接任务。
|
||||
pub fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) {
|
||||
let cancellation = self.cancellation.clone();
|
||||
drop(self.connection_tasks.spawn(async move {
|
||||
cancellation.run_until_cancelled(future).await;
|
||||
}));
|
||||
}
|
||||
|
||||
/// 停止所有受管连接,并返回关闭前仍在运行的长连接任务数。
|
||||
pub fn begin_shutdown(&self) -> usize {
|
||||
let active_websocket_connections = self.connection_tasks.len();
|
||||
self.connection_tasks.close();
|
||||
self.cancellation.cancel();
|
||||
active_websocket_connections
|
||||
}
|
||||
|
||||
/// 等待所有受管长连接任务退出。
|
||||
pub async fn wait(&self) {
|
||||
self.connection_tasks.wait().await;
|
||||
}
|
||||
|
||||
/// 是否已进入连接排空阶段。
|
||||
pub fn is_shutting_down(&self) -> bool {
|
||||
self.cancellation.is_cancelled()
|
||||
}
|
||||
}
|
||||
|
||||
/// 将所有 HTTP 响应体接入统一关闭信号。
|
||||
pub async fn drain_response_body(
|
||||
State(connection_drain): State<ConnectionDrain>,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let response = next.run(request).await;
|
||||
let (parts, body) = response.into_parts();
|
||||
let shutdown = connection_drain.cancellation.clone().cancelled_owned();
|
||||
let body = Body::from_stream(body.into_data_stream().take_until(shutdown));
|
||||
Response::from_parts(parts, body)
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
//! HTTP 中间件 —— 请求 ID 与 Tracing。
|
||||
//! HTTP 中间件 —— 连接排空、请求 ID 与 Tracing。
|
||||
|
||||
pub mod connection_drain;
|
||||
pub mod request_id;
|
||||
pub mod trace;
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use axum::{
|
||||
extract::{connect_info::ConnectInfo, Request},
|
||||
extract::{Request, connect_info::ConnectInfo},
|
||||
http::{HeaderMap, HeaderValue},
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
|
||||
@@ -10,7 +10,7 @@ use tower_http::{
|
||||
classify::{ServerErrorsAsFailures, ServerErrorsFailureClass, SharedClassifier},
|
||||
trace::TraceLayer,
|
||||
};
|
||||
use tracing::{debug, info, warn, Span};
|
||||
use tracing::{Span, debug, info, warn};
|
||||
|
||||
use crate::api::middleware::request_id::{ClientIp, RequestId};
|
||||
|
||||
@@ -48,7 +48,7 @@ fn make_http_span(request: &Request<Body>) -> Span {
|
||||
}
|
||||
|
||||
fn on_http_request(_request: &Request<Body>, _span: &Span) {
|
||||
debug!("received HTTP request");
|
||||
debug!("Received HTTP request");
|
||||
}
|
||||
|
||||
fn on_http_response(response: &Response<Body>, latency: Duration, _span: &Span) {
|
||||
@@ -58,7 +58,7 @@ fn on_http_response(response: &Response<Body>, latency: Duration, _span: &Span)
|
||||
info!(
|
||||
status = response.status().as_u16(),
|
||||
latency_ms = latency.as_millis(),
|
||||
"completed HTTP request"
|
||||
"Completed HTTP request"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -70,7 +70,7 @@ fn on_http_failure(failure_class: ServerErrorsFailureClass, latency: Duration, _
|
||||
warn!(
|
||||
failure = %failure_class,
|
||||
latency_ms = latency.as_millis(),
|
||||
"failed HTTP request"
|
||||
"Failed HTTP request"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
+14
-17
@@ -6,35 +6,34 @@ use std::{
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
|
||||
use axum::{Router, extract::State, http::StatusCode, middleware, routing::get};
|
||||
|
||||
use crate::{
|
||||
api::{
|
||||
admin, assets, client,
|
||||
admin::{
|
||||
self, accounts_routes::query::AccountListQueryService,
|
||||
dashboard_routes::DashboardQueryService,
|
||||
},
|
||||
assets, client,
|
||||
middleware::{request_id::attach_request_id, trace::http_trace_layer},
|
||||
},
|
||||
auth::service::SessionService,
|
||||
dispatch::{affinity::SessionAffinityService, service::ResponseDispatchService},
|
||||
fleet::{
|
||||
manage::AccountManageService, pool::AccountPoolService, refresh::TokenRefreshService,
|
||||
store::AccountStore,
|
||||
},
|
||||
fleet::{manage::AccountManageService, pool::AccountPoolService},
|
||||
keys::{manage::KeyManageService, service::KeyVerifier},
|
||||
models::service::ModelService,
|
||||
settings::service::SettingsService,
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageQueryService, ops::query::OpsQueryService,
|
||||
usage::query::UsageQueryService,
|
||||
},
|
||||
upstream::openai::{fingerprint::RuntimeFingerprint, token_client::OpenAiTokenClient},
|
||||
telemetry::{ops::query::OpsQueryService, usage::query::UsageQueryService},
|
||||
update::service::SystemUpdateService,
|
||||
};
|
||||
|
||||
/// HTTP API 所需的领域服务集合。
|
||||
#[derive(Clone)]
|
||||
pub struct ApiServices {
|
||||
pub(crate) account_list: Arc<AccountListQueryService>,
|
||||
pub(crate) dashboard: Arc<DashboardQueryService>,
|
||||
pub health_probe: Arc<dyn HealthProbe>,
|
||||
pub models: Arc<ModelService>,
|
||||
pub accounts: Arc<dyn AccountStore>,
|
||||
pub client_keys: Arc<KeyVerifier>,
|
||||
pub admin_client_keys: Arc<KeyManageService>,
|
||||
pub admin_sessions: Arc<SessionService>,
|
||||
@@ -42,20 +41,18 @@ pub struct ApiServices {
|
||||
pub admin_accounts: Arc<AccountManageService>,
|
||||
pub usage_records: Arc<UsageQueryService>,
|
||||
pub ops_errors: Arc<OpsQueryService>,
|
||||
pub usage: Arc<AccountUsageQueryService>,
|
||||
pub account_pool: Arc<AccountPoolService>,
|
||||
pub token_refresh: Arc<TokenRefreshService<OpenAiTokenClient>>,
|
||||
pub responses: Arc<ResponseDispatchService>,
|
||||
pub session_affinity: Arc<SessionAffinityService>,
|
||||
pub fingerprint: RuntimeFingerprint,
|
||||
pub process_control: Arc<dyn ProcessControl>,
|
||||
pub system_update: Arc<SystemUpdateService>,
|
||||
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
|
||||
}
|
||||
|
||||
/// 由 bootstrap 实现的进程生命周期控制端口。
|
||||
pub trait ProcessControl: Send + Sync + 'static {
|
||||
fn request_shutdown(&self);
|
||||
fn request_restart(&self, executable_path: PathBuf);
|
||||
fn subscribe_shutdown(&self) -> tokio::sync::broadcast::Receiver<()>;
|
||||
}
|
||||
|
||||
/// 由 bootstrap 实现的基础设施健康检查端口。
|
||||
@@ -94,7 +91,7 @@ pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
|
||||
|
||||
async fn healthz(State(state): State<AppState>) -> StatusCode {
|
||||
if let Err(error) = state.services.health_probe.check().await {
|
||||
tracing::warn!(error, "health check failed");
|
||||
tracing::warn!(error, "Health check failed");
|
||||
return StatusCode::SERVICE_UNAVAILABLE;
|
||||
}
|
||||
StatusCode::NO_CONTENT
|
||||
|
||||
@@ -17,6 +17,7 @@ use super::{
|
||||
const LOGIN_FAILURE_WINDOW: Duration = Duration::minutes(15);
|
||||
const LOGIN_LOCK_DURATION: Duration = Duration::minutes(15);
|
||||
const LOGIN_MAX_FAILURES: u32 = 5;
|
||||
const LOGIN_FAILURE_SOURCE_LIMIT: usize = 10_000;
|
||||
|
||||
/// 管理员会话服务。
|
||||
#[derive(Clone)]
|
||||
@@ -134,6 +135,9 @@ impl SessionService {
|
||||
let login_locked = {
|
||||
let mut failures = self.login_failures.lock().await;
|
||||
retain_active_login_failures(&mut failures, now);
|
||||
if !failures.contains_key(source) && failures.len() >= LOGIN_FAILURE_SOURCE_LIMIT {
|
||||
return Err(SessionError::LoginThrottled);
|
||||
}
|
||||
failures
|
||||
.get(source)
|
||||
.and_then(|state| state.locked_until)
|
||||
|
||||
+537
-304
@@ -1,9 +1,23 @@
|
||||
use std::{collections::BTreeMap, env, path::Path};
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
env,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use chrono::{DateTime, Utc};
|
||||
use secrecy::{ExposeSecret, SecretBox};
|
||||
use serde::{Deserialize, de::IgnoredAny};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use url::Url;
|
||||
|
||||
const CONFIG_SCHEMA_VERSION: u32 = 1;
|
||||
const CONFIG_RELATIVE_PATH: &str = "deploy/config.yaml";
|
||||
const SERVER_HOST_ENV: &str = "CPR_SERVER_HOST";
|
||||
const SERVER_PORT_ENV: &str = "CPR_SERVER_PORT";
|
||||
const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL";
|
||||
const REDIS_URL_ENV: &str = "CPR_REDIS_URL";
|
||||
const SERVICE_PASSWORD_HEX_LENGTH: usize = 48;
|
||||
|
||||
const WEAK_ADMIN_PASSWORDS: &[&str] = &[
|
||||
"",
|
||||
"admin",
|
||||
@@ -15,49 +29,230 @@ const WEAK_ADMIN_PASSWORDS: &[&str] = &[
|
||||
"codex-proxy-rs",
|
||||
];
|
||||
|
||||
/// 应用运行总配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
type SecretValue = SecretBox<String>;
|
||||
|
||||
/// 已加载并完成校验的启动配置。
|
||||
///
|
||||
/// 数据库、Redis 和管理员初始化密码只在启动阶段保留,并在释放时清零。
|
||||
#[derive(Debug)]
|
||||
pub struct BootstrapConfig {
|
||||
app: AppConfig,
|
||||
database_url: SecretValue,
|
||||
redis_url: SecretValue,
|
||||
admin_default_password: SecretValue,
|
||||
}
|
||||
|
||||
impl BootstrapConfig {
|
||||
/// 从当前目录或父目录中的 `deploy/config.yaml` 加载配置。
|
||||
pub fn load() -> Result<Self, ConfigError> {
|
||||
let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?;
|
||||
let config_path = discover_config_path(¤t_directory)?;
|
||||
let overrides = TopologyOverrides::from_environment()?;
|
||||
Self::load_from_path_with_overrides(config_path, overrides)
|
||||
}
|
||||
|
||||
/// 从指定 YAML 文件加载配置,不应用容器拓扑覆盖。
|
||||
pub fn load_from_path(config_path: impl AsRef<Path>) -> Result<Self, ConfigError> {
|
||||
Self::load_from_path_with_overrides(config_path, TopologyOverrides::default())
|
||||
}
|
||||
|
||||
/// 从指定 YAML 文件加载配置,并应用明确的容器拓扑覆盖。
|
||||
pub fn load_from_path_with_overrides(
|
||||
config_path: impl AsRef<Path>,
|
||||
overrides: TopologyOverrides,
|
||||
) -> Result<Self, ConfigError> {
|
||||
let config_path = absolute_path(config_path.as_ref())?;
|
||||
let config_directory = config_path.parent().ok_or(ConfigError::InvalidConfigPath)?;
|
||||
|
||||
let document: ConfigDocument = ::config::Config::builder()
|
||||
.add_source(::config::File::from(config_path.as_path()).required(true))
|
||||
.build()
|
||||
.and_then(::config::Config::try_deserialize)
|
||||
.map_err(|_| ConfigError::InvalidDocument {
|
||||
path: config_path.clone(),
|
||||
})?;
|
||||
|
||||
let ConfigDocument {
|
||||
cpr: mut file_config,
|
||||
_services: _,
|
||||
} = document;
|
||||
if file_config.schema_version != CONFIG_SCHEMA_VERSION {
|
||||
return Err(ConfigError::UnsupportedSchemaVersion);
|
||||
}
|
||||
|
||||
overrides.apply(&mut file_config);
|
||||
validate_file_config(&file_config)?;
|
||||
|
||||
resolve_relative_path(config_directory, &mut file_config.runtime.data_directory);
|
||||
resolve_relative_path(config_directory, &mut file_config.logging.file.directory);
|
||||
|
||||
let database_url = connection_url_with_password(
|
||||
&file_config.database.url,
|
||||
&file_config.database.password,
|
||||
"database.url",
|
||||
)?;
|
||||
let redis_url = connection_url_with_password(
|
||||
&file_config.redis.url,
|
||||
&file_config.redis.password,
|
||||
"redis.url",
|
||||
)?;
|
||||
|
||||
let FileAppConfig {
|
||||
schema_version: _,
|
||||
server,
|
||||
api,
|
||||
database,
|
||||
redis,
|
||||
runtime,
|
||||
quota,
|
||||
ws_pool,
|
||||
wire_profile,
|
||||
admin,
|
||||
logging,
|
||||
telemetry,
|
||||
} = file_config;
|
||||
|
||||
let app = AppConfig {
|
||||
server,
|
||||
api,
|
||||
model_aliases: BTreeMap::new(),
|
||||
auth: AuthConfig::default(),
|
||||
quota,
|
||||
database: DatabaseConfig { url: database.url },
|
||||
redis: RedisConfig { url: redis.url },
|
||||
runtime,
|
||||
ws_pool,
|
||||
wire_profile,
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: admin.session_ttl_minutes,
|
||||
default_username: admin.default_username,
|
||||
},
|
||||
logging,
|
||||
telemetry,
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
app,
|
||||
database_url,
|
||||
redis_url,
|
||||
admin_default_password: admin.default_password,
|
||||
})
|
||||
}
|
||||
|
||||
/// 返回不包含启动密码的应用运行配置。
|
||||
pub fn app(&self) -> &AppConfig {
|
||||
&self.app
|
||||
}
|
||||
|
||||
/// 返回已安全注入密码的 PostgreSQL 连接地址。
|
||||
pub fn database_url(&self) -> &str {
|
||||
self.database_url.expose_secret()
|
||||
}
|
||||
|
||||
/// 返回已安全注入密码的 Redis 连接地址。
|
||||
pub fn redis_url(&self) -> &str {
|
||||
self.redis_url.expose_secret()
|
||||
}
|
||||
|
||||
pub(crate) fn into_parts(self) -> BootstrapConfigParts {
|
||||
BootstrapConfigParts {
|
||||
app: self.app,
|
||||
database_url: self.database_url,
|
||||
redis_url: self.redis_url,
|
||||
admin_default_password: self.admin_default_password,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct BootstrapConfigParts {
|
||||
pub(crate) app: AppConfig,
|
||||
pub(crate) database_url: SecretValue,
|
||||
pub(crate) redis_url: SecretValue,
|
||||
pub(crate) admin_default_password: SecretValue,
|
||||
}
|
||||
|
||||
/// Docker Compose 对容器内部网络拓扑的固定覆盖。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct TopologyOverrides {
|
||||
/// 容器内监听地址。
|
||||
pub server_host: Option<String>,
|
||||
/// 容器内监听端口。
|
||||
pub server_port: Option<u16>,
|
||||
/// 容器内 PostgreSQL 地址,不包含密码。
|
||||
pub database_url: Option<String>,
|
||||
/// 容器内 Redis 地址,不包含密码。
|
||||
pub redis_url: Option<String>,
|
||||
}
|
||||
|
||||
impl TopologyOverrides {
|
||||
fn from_environment() -> Result<Self, ConfigError> {
|
||||
let server_port = optional_environment_value(SERVER_PORT_ENV)?
|
||||
.map(|value| {
|
||||
value
|
||||
.parse::<u16>()
|
||||
.ok()
|
||||
.filter(|port| *port > 0)
|
||||
.ok_or(ConfigError::InvalidTopologyOverride(SERVER_PORT_ENV))
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(Self {
|
||||
server_host: optional_environment_value(SERVER_HOST_ENV)?,
|
||||
server_port,
|
||||
database_url: optional_environment_value(DATABASE_URL_ENV)?,
|
||||
redis_url: optional_environment_value(REDIS_URL_ENV)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn apply(self, config: &mut FileAppConfig) {
|
||||
if let Some(host) = self.server_host {
|
||||
config.server.host = host;
|
||||
}
|
||||
if let Some(port) = self.server_port {
|
||||
config.server.port = port;
|
||||
}
|
||||
if let Some(url) = self.database_url {
|
||||
config.database.url = url;
|
||||
}
|
||||
if let Some(url) = self.redis_url {
|
||||
config.redis.url = url;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 应用运行总配置,不包含任何启动密码。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AppConfig {
|
||||
/// HTTP 服务配置。
|
||||
pub server: ServerConfig,
|
||||
/// API 地址配置。
|
||||
pub api: ApiConfig,
|
||||
/// 运行时模型别名,由数据库设置加载,配置文件不承载该字段。
|
||||
#[serde(skip)]
|
||||
pub model_aliases: BTreeMap<String, String>,
|
||||
/// 认证与刷新配置,由固定默认值和数据库运行时设置承载,配置文件不承载该字段。
|
||||
#[serde(skip)]
|
||||
pub auth: AuthConfig,
|
||||
/// 配额配置。
|
||||
#[serde(default)]
|
||||
pub quota: QuotaConfig,
|
||||
/// 数据库配置。
|
||||
/// PostgreSQL 基础连接地址,不包含密码。
|
||||
pub database: DatabaseConfig,
|
||||
/// Redis 运行态存储配置。
|
||||
/// Redis 基础连接地址,不包含密码。
|
||||
pub redis: RedisConfig,
|
||||
/// TLS 偏好配置。
|
||||
#[serde(default)]
|
||||
pub tls: TlsConfig,
|
||||
/// 本地持久运行目录。
|
||||
pub runtime: RuntimePathsConfig,
|
||||
/// WebSocket 连接池启动设置。
|
||||
#[serde(default)]
|
||||
pub ws_pool: WebSocketPoolSettings,
|
||||
/// 上游请求指纹默认配置。
|
||||
#[serde(default)]
|
||||
pub fingerprint: FingerprintConfig,
|
||||
/// 管理员初始化配置。
|
||||
#[serde(default)]
|
||||
/// 经审计固定的 Codex Desktop 上游请求画像。
|
||||
pub wire_profile: WireProfileConfig,
|
||||
/// 管理员初始化配置,不包含密码。
|
||||
pub admin: AdminConfig,
|
||||
/// 日志配置。
|
||||
#[serde(default)]
|
||||
pub logging: LoggingConfig,
|
||||
/// PostgreSQL 遥测事实记录配置。
|
||||
#[serde(default)]
|
||||
pub telemetry: TelemetryConfig,
|
||||
}
|
||||
|
||||
/// HTTP 监听配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ServerConfig {
|
||||
/// 监听主机。
|
||||
@@ -67,14 +262,15 @@ pub struct ServerConfig {
|
||||
}
|
||||
|
||||
/// 上游 API 地址配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ApiConfig {
|
||||
/// 上游基础 URL。
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
/// 认证、轮换与 token 续期配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
pub struct AuthConfig {
|
||||
/// 提前刷新秒数。
|
||||
pub refresh_margin_seconds: u64,
|
||||
@@ -113,58 +309,39 @@ impl Default for AuthConfig {
|
||||
}
|
||||
|
||||
/// 配额刷新与跳过配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct QuotaConfig {
|
||||
/// 配额刷新周期。
|
||||
#[serde(default = "default_quota_refresh_interval_minutes")]
|
||||
pub refresh_interval_minutes: u64,
|
||||
/// 是否跳过已耗尽配额的账号。
|
||||
#[serde(default = "default_skip_exhausted")]
|
||||
pub skip_exhausted: bool,
|
||||
}
|
||||
|
||||
impl Default for QuotaConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
refresh_interval_minutes: default_quota_refresh_interval_minutes(),
|
||||
skip_exhausted: default_skip_exhausted(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn default_quota_refresh_interval_minutes() -> u64 {
|
||||
5
|
||||
}
|
||||
|
||||
fn default_skip_exhausted() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// 数据库连接配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DatabaseConfig {
|
||||
/// PostgreSQL 连接 URL。
|
||||
/// PostgreSQL 基础连接 URL,不包含密码。
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
/// Redis 连接配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisConfig {
|
||||
/// Redis 连接 URL。
|
||||
/// Redis 基础连接 URL,不包含密码。
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
/// TLS/HTTP 协议偏好配置。
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub struct TlsConfig {
|
||||
/// 是否强制 HTTP/1.1。
|
||||
pub force_http11: bool,
|
||||
/// 本地持久运行目录。
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct RuntimePathsConfig {
|
||||
/// 持久身份密钥与更新状态目录。
|
||||
pub data_directory: PathBuf,
|
||||
}
|
||||
|
||||
/// WebSocket 连接池启动设置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct WebSocketPoolSettings {
|
||||
/// 是否启用连接池。
|
||||
@@ -173,179 +350,52 @@ pub struct WebSocketPoolSettings {
|
||||
pub max_age_ms: u64,
|
||||
/// 单账号最大连接数。
|
||||
pub max_per_account: usize,
|
||||
/// 所有账号合计最大连接数。
|
||||
pub max_total: usize,
|
||||
/// 所有账号合计最大并发建连数。
|
||||
pub max_connecting: usize,
|
||||
/// 首个真实输出(首事件)到达前的绝对超时;`0` 表示禁用。
|
||||
///
|
||||
/// 从发出 `response.create` 起算,覆盖建连/发送/上游排队/思考的全程,直到收到
|
||||
/// 首个真实内容帧(`response.created`/`response.in_progress` 等纯生命周期帧不计)。
|
||||
/// 超时判定为连接落到病态上游后端,丢弃并换新连接重试。
|
||||
#[serde(default = "default_ws_initial_event_timeout_ms")]
|
||||
pub initial_event_timeout_ms: u64,
|
||||
}
|
||||
|
||||
fn default_ws_initial_event_timeout_ms() -> u64 {
|
||||
20_000
|
||||
}
|
||||
|
||||
impl Default for WebSocketPoolSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
max_age_ms: 3_300_000,
|
||||
max_per_account: 8,
|
||||
initial_event_timeout_ms: default_ws_initial_event_timeout_ms(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 上游请求指纹默认配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
/// 经审计固定的 Codex Desktop 上游请求画像配置。
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct FingerprintConfig {
|
||||
/// 客户端来源名。
|
||||
pub struct WireProfileConfig {
|
||||
/// `originator` 请求头及 User-Agent 产品名。
|
||||
pub originator: String,
|
||||
/// 应用版本。
|
||||
pub app_version: String,
|
||||
/// 构建号。
|
||||
pub build_number: String,
|
||||
/// 平台名。
|
||||
pub platform: String,
|
||||
/// 架构名。
|
||||
/// bundled Codex Core 版本。
|
||||
pub codex_version: String,
|
||||
/// Desktop 应用版本。
|
||||
pub desktop_version: String,
|
||||
/// Desktop 制品构建号。
|
||||
pub desktop_build: String,
|
||||
/// Codex Core UA 中的目标操作系统类型。
|
||||
pub os_type: String,
|
||||
/// Codex Core UA 中的目标操作系统版本。
|
||||
pub os_version: String,
|
||||
/// Codex Core UA 中的目标架构。
|
||||
pub arch: String,
|
||||
/// Chromium 主版本。
|
||||
pub chromium_version: String,
|
||||
/// User-Agent 模板。
|
||||
pub user_agent_template: String,
|
||||
/// 默认请求头。
|
||||
pub default_headers: Vec<FingerprintHeaderConfig>,
|
||||
/// 请求头排序优先级。
|
||||
pub header_order: Vec<String>,
|
||||
/// Codex Core UA 中的终端标记。
|
||||
pub terminal: String,
|
||||
/// 此画像最后一次经制品与源码核验的时间。
|
||||
pub verified_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl Default for FingerprintConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
originator: "Codex Desktop".to_string(),
|
||||
app_version: "26.519.81530".to_string(),
|
||||
build_number: "3178".to_string(),
|
||||
platform: "darwin".to_string(),
|
||||
arch: "arm64".to_string(),
|
||||
chromium_version: "146".to_string(),
|
||||
user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(),
|
||||
default_headers: vec![
|
||||
FingerprintHeaderConfig::new("Accept-Encoding", "gzip, deflate, br, zstd"),
|
||||
FingerprintHeaderConfig::new("Accept-Language", "en-US,en;q=0.9"),
|
||||
FingerprintHeaderConfig::new("sec-ch-ua-mobile", "?0"),
|
||||
FingerprintHeaderConfig::new("sec-ch-ua-platform", "\"macOS\""),
|
||||
FingerprintHeaderConfig::new("sec-fetch-site", "same-origin"),
|
||||
FingerprintHeaderConfig::new("sec-fetch-mode", "cors"),
|
||||
FingerprintHeaderConfig::new("sec-fetch-dest", "empty"),
|
||||
],
|
||||
header_order: vec![
|
||||
"authorization".to_string(),
|
||||
"chatgpt-account-id".to_string(),
|
||||
"originator".to_string(),
|
||||
"x-openai-internal-codex-residency".to_string(),
|
||||
"x-client-request-id".to_string(),
|
||||
"x-codex-installation-id".to_string(),
|
||||
"x-codex-turn-state".to_string(),
|
||||
"openai-beta".to_string(),
|
||||
"user-agent".to_string(),
|
||||
"sec-ch-ua".to_string(),
|
||||
"sec-ch-ua-mobile".to_string(),
|
||||
"sec-ch-ua-platform".to_string(),
|
||||
"accept-encoding".to_string(),
|
||||
"accept-language".to_string(),
|
||||
"sec-fetch-site".to_string(),
|
||||
"sec-fetch-mode".to_string(),
|
||||
"sec-fetch-dest".to_string(),
|
||||
"content-type".to_string(),
|
||||
"accept".to_string(),
|
||||
"cookie".to_string(),
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 指纹默认请求头配置项。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct FingerprintHeaderConfig {
|
||||
/// 请求头名称。
|
||||
pub name: String,
|
||||
/// 请求头值。
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
impl FingerprintHeaderConfig {
|
||||
fn new(name: &str, value: &str) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理员初始化与会话配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
/// 管理员初始化与会话配置,不包含初始化密码。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminConfig {
|
||||
/// 会话有效期(分钟)。
|
||||
pub session_ttl_minutes: u64,
|
||||
/// 默认管理员用户名(首次启动时创建)
|
||||
#[serde(default = "default_admin_username")]
|
||||
/// 默认管理员用户名(首次启动时创建)。
|
||||
pub default_username: String,
|
||||
/// 默认管理员密码(首次启动时创建)
|
||||
#[serde(default = "default_admin_password")]
|
||||
pub default_password: String,
|
||||
}
|
||||
|
||||
impl Default for AdminConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
session_ttl_minutes: 1440,
|
||||
default_username: default_admin_username(),
|
||||
default_password: default_admin_password(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AdminConfig {
|
||||
/// 校验首次启动创建管理员时使用的默认密码。
|
||||
pub fn validate_default_password(&self) -> Result<(), AdminConfigError> {
|
||||
let password = self.default_password.trim();
|
||||
if password.len() < 12 {
|
||||
return Err(AdminConfigError::WeakDefaultPassword);
|
||||
}
|
||||
if WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str()) {
|
||||
return Err(AdminConfigError::WeakDefaultPassword);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理员配置校验错误。
|
||||
#[derive(Debug, Clone, Copy, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum AdminConfigError {
|
||||
/// 默认管理员密码为空、过短或属于常见弱口令。
|
||||
#[error(
|
||||
"admin.default_password must be set to a non-default password with at least 12 characters"
|
||||
)]
|
||||
WeakDefaultPassword,
|
||||
}
|
||||
|
||||
fn default_admin_username() -> String {
|
||||
"admin".to_string()
|
||||
}
|
||||
|
||||
fn default_admin_password() -> String {
|
||||
String::new()
|
||||
}
|
||||
|
||||
/// 应用结构化日志配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LoggingConfig {
|
||||
/// 默认日志级别;`RUST_LOG` 存在时优先使用环境变量。
|
||||
/// 默认日志级别;进程 `RUST_LOG` 可作临时覆盖。
|
||||
pub level: String,
|
||||
/// 是否写入标准输出。
|
||||
pub stdout: bool,
|
||||
@@ -353,24 +403,14 @@ pub struct LoggingConfig {
|
||||
pub file: FileLoggingConfig,
|
||||
}
|
||||
|
||||
impl Default for LoggingConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
level: "info".to_string(),
|
||||
stdout: true,
|
||||
file: FileLoggingConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 应用结构化文件日志配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct FileLoggingConfig {
|
||||
/// 是否写入轮转文件。
|
||||
pub enabled: bool,
|
||||
/// 日志目录。
|
||||
pub directory: String,
|
||||
pub directory: PathBuf,
|
||||
/// 按中国自然日计算的保留天数。
|
||||
pub retention_days: usize,
|
||||
/// 单文件大小上限,单位 MiB。
|
||||
@@ -379,104 +419,297 @@ pub struct FileLoggingConfig {
|
||||
pub max_files: usize,
|
||||
}
|
||||
|
||||
impl Default for FileLoggingConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
directory: ".runtime/logs".to_string(),
|
||||
retention_days: 14,
|
||||
max_file_size_mb: 20,
|
||||
max_files: 20,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// PostgreSQL 遥测事实记录配置。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct TelemetryConfig {
|
||||
/// 是否记录成功与失败代理事实。
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
impl Default for TelemetryConfig {
|
||||
fn default() -> Self {
|
||||
Self { enabled: true }
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ConfigDocument {
|
||||
#[serde(rename = "x-cpr")]
|
||||
cpr: FileAppConfig,
|
||||
#[serde(rename = "services")]
|
||||
_services: IgnoredAny,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct FileAppConfig {
|
||||
schema_version: u32,
|
||||
server: ServerConfig,
|
||||
api: ApiConfig,
|
||||
database: FileDatabaseConfig,
|
||||
redis: FileRedisConfig,
|
||||
runtime: RuntimePathsConfig,
|
||||
quota: QuotaConfig,
|
||||
ws_pool: WebSocketPoolSettings,
|
||||
wire_profile: WireProfileConfig,
|
||||
admin: FileAdminConfig,
|
||||
logging: LoggingConfig,
|
||||
telemetry: TelemetryConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct FileDatabaseConfig {
|
||||
url: String,
|
||||
password: SecretValue,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct FileRedisConfig {
|
||||
url: String,
|
||||
password: SecretValue,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct FileAdminConfig {
|
||||
session_ttl_minutes: u64,
|
||||
default_username: String,
|
||||
default_password: SecretValue,
|
||||
}
|
||||
|
||||
fn discover_config_path(start: &Path) -> Result<PathBuf, ConfigError> {
|
||||
start
|
||||
.ancestors()
|
||||
.map(|directory| directory.join(CONFIG_RELATIVE_PATH))
|
||||
.find(|candidate| candidate.is_file())
|
||||
.ok_or(ConfigError::ConfigFileNotFound)
|
||||
}
|
||||
|
||||
fn absolute_path(path: &Path) -> Result<PathBuf, ConfigError> {
|
||||
if path.is_absolute() {
|
||||
return Ok(path.to_path_buf());
|
||||
}
|
||||
let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?;
|
||||
Ok(current_directory.join(path))
|
||||
}
|
||||
|
||||
fn resolve_relative_path(base: &Path, path: &mut PathBuf) {
|
||||
if path.is_relative() {
|
||||
*path = base.join(&*path);
|
||||
}
|
||||
}
|
||||
|
||||
const CONFIG_FILE_ENV: &str = "CPR_CONFIG_FILE";
|
||||
const ADMIN_DEFAULT_PASSWORD_ENV: &str = "CPR_ADMIN_DEFAULT_PASSWORD";
|
||||
const POSTGRES_PASSWORD_ENV: &str = "CPR_POSTGRES_PASSWORD";
|
||||
const REDIS_PASSWORD_ENV: &str = "CPR_REDIS_PASSWORD";
|
||||
|
||||
impl AppConfig {
|
||||
/// 从 `CPR_CONFIG_FILE` 或当前目录 `config.yaml` 加载配置文件,并应用运行时密码。
|
||||
pub fn load() -> Result<Self, ::config::ConfigError> {
|
||||
let config_file = env::var(CONFIG_FILE_ENV)
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
match config_file {
|
||||
Some(path) => load_file(path),
|
||||
None => Self::load_from_dir("."),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从指定目录加载配置文件,并应用运行时密码。
|
||||
pub fn load_from_dir(config_dir: impl AsRef<Path>) -> Result<Self, ::config::ConfigError> {
|
||||
load_file(config_dir.as_ref().join("config.yaml"))
|
||||
fn optional_environment_value(name: &'static str) -> Result<Option<String>, ConfigError> {
|
||||
match env::var(name) {
|
||||
Ok(value) if value.trim().is_empty() => Err(ConfigError::InvalidTopologyOverride(name)),
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(env::VarError::NotPresent) => Ok(None),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(ConfigError::InvalidTopologyOverride(name)),
|
||||
}
|
||||
}
|
||||
|
||||
fn load_file(config_file: impl AsRef<Path>) -> Result<AppConfig, ::config::ConfigError> {
|
||||
let mut app_config: AppConfig = ::config::Config::builder()
|
||||
.add_source(::config::File::from(config_file.as_ref()).required(true))
|
||||
.build()?
|
||||
.try_deserialize()?;
|
||||
apply_password_override(
|
||||
&mut app_config.database.url,
|
||||
"database.url",
|
||||
POSTGRES_PASSWORD_ENV,
|
||||
fn validate_file_config(config: &FileAppConfig) -> Result<(), ConfigError> {
|
||||
validate_nonempty("server.host", &config.server.host)?;
|
||||
if config.server.port == 0 {
|
||||
return Err(ConfigError::InvalidField("server.port"));
|
||||
}
|
||||
validate_url(
|
||||
"api.base_url",
|
||||
&config.api.base_url,
|
||||
&["http", "https"],
|
||||
false,
|
||||
)?;
|
||||
apply_password_override(&mut app_config.redis.url, "redis.url", REDIS_PASSWORD_ENV)?;
|
||||
if let Some(password) = password_override(ADMIN_DEFAULT_PASSWORD_ENV)? {
|
||||
app_config.admin.default_password = password;
|
||||
validate_url(
|
||||
"database.url",
|
||||
&config.database.url,
|
||||
&["postgres", "postgresql"],
|
||||
true,
|
||||
)?;
|
||||
validate_url("redis.url", &config.redis.url, &["redis", "rediss"], true)?;
|
||||
validate_service_password("database.password", &config.database.password)?;
|
||||
validate_service_password("redis.password", &config.redis.password)?;
|
||||
validate_admin_password(config.admin.default_password.expose_secret())?;
|
||||
validate_nonempty("admin.default_username", &config.admin.default_username)?;
|
||||
if config.admin.session_ttl_minutes == 0 {
|
||||
return Err(ConfigError::InvalidField("admin.session_ttl_minutes"));
|
||||
}
|
||||
Ok(app_config)
|
||||
}
|
||||
|
||||
fn apply_password_override(
|
||||
raw_url: &mut String,
|
||||
config_key: &str,
|
||||
env_name: &str,
|
||||
) -> Result<(), ::config::ConfigError> {
|
||||
let Some(password) = password_override(env_name)? else {
|
||||
return Ok(());
|
||||
};
|
||||
let mut url = Url::parse(raw_url).map_err(|_| {
|
||||
::config::ConfigError::Message(format!(
|
||||
"{config_key} must be a valid URL when {env_name} is set"
|
||||
))
|
||||
})?;
|
||||
let encoded_password = utf8_percent_encode(&password, NON_ALPHANUMERIC).to_string();
|
||||
url.set_password(Some(&encoded_password)).map_err(|()| {
|
||||
::config::ConfigError::Message(format!(
|
||||
"{config_key} does not support a password override from {env_name}"
|
||||
))
|
||||
})?;
|
||||
*raw_url = url.into();
|
||||
if config.runtime.data_directory.as_os_str().is_empty() {
|
||||
return Err(ConfigError::InvalidField("runtime.data_directory"));
|
||||
}
|
||||
if config.quota.refresh_interval_minutes == 0 {
|
||||
return Err(ConfigError::InvalidField("quota.refresh_interval_minutes"));
|
||||
}
|
||||
if config.ws_pool.max_age_ms == 0 {
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_age_ms"));
|
||||
}
|
||||
if config.ws_pool.max_per_account == 0 {
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_per_account"));
|
||||
}
|
||||
if config.ws_pool.max_total == 0 {
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_total"));
|
||||
}
|
||||
if config.ws_pool.max_connecting == 0
|
||||
|| config.ws_pool.max_connecting > config.ws_pool.max_total
|
||||
{
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_connecting"));
|
||||
}
|
||||
validate_wire_profile(&config.wire_profile)?;
|
||||
validate_logging(&config.logging)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn password_override(env_name: &str) -> Result<Option<String>, ::config::ConfigError> {
|
||||
match env::var(env_name) {
|
||||
Ok(password) if password.is_empty() => Err(::config::ConfigError::Message(format!(
|
||||
"{env_name} must not be empty"
|
||||
))),
|
||||
Ok(password) => Ok(Some(password)),
|
||||
Err(env::VarError::NotPresent) => Ok(None),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(::config::ConfigError::Message(format!(
|
||||
"{env_name} must contain valid UTF-8"
|
||||
))),
|
||||
fn validate_url(
|
||||
field: &'static str,
|
||||
raw_url: &str,
|
||||
allowed_schemes: &[&str],
|
||||
reject_password: bool,
|
||||
) -> Result<(), ConfigError> {
|
||||
let url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?;
|
||||
if !allowed_schemes.contains(&url.scheme()) || url.host_str().is_none() {
|
||||
return Err(ConfigError::InvalidField(field));
|
||||
}
|
||||
if reject_password && url.password().is_some() {
|
||||
return Err(ConfigError::PasswordInUrl(field));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_service_password(
|
||||
field: &'static str,
|
||||
password: &SecretValue,
|
||||
) -> Result<(), ConfigError> {
|
||||
let password = password.expose_secret();
|
||||
if password.len() != SERVICE_PASSWORD_HEX_LENGTH
|
||||
|| !password.bytes().all(|byte| byte.is_ascii_hexdigit())
|
||||
{
|
||||
return Err(ConfigError::InvalidServicePassword(field));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_admin_password(password: &str) -> Result<(), ConfigError> {
|
||||
let password = password.trim();
|
||||
if password.len() < 12
|
||||
|| password.contains('$')
|
||||
|| WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str())
|
||||
{
|
||||
return Err(ConfigError::WeakAdminPassword);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_wire_profile(config: &WireProfileConfig) -> Result<(), ConfigError> {
|
||||
for (field, value) in [
|
||||
("wire_profile.originator", config.originator.as_str()),
|
||||
("wire_profile.codex_version", config.codex_version.as_str()),
|
||||
(
|
||||
"wire_profile.desktop_version",
|
||||
config.desktop_version.as_str(),
|
||||
),
|
||||
("wire_profile.desktop_build", config.desktop_build.as_str()),
|
||||
("wire_profile.os_type", config.os_type.as_str()),
|
||||
("wire_profile.os_version", config.os_version.as_str()),
|
||||
("wire_profile.arch", config.arch.as_str()),
|
||||
("wire_profile.terminal", config.terminal.as_str()),
|
||||
] {
|
||||
validate_nonempty(field, value)?;
|
||||
}
|
||||
|
||||
if semver::Version::parse(&config.codex_version).is_err() {
|
||||
return Err(ConfigError::InvalidField("wire_profile.codex_version"));
|
||||
}
|
||||
if !is_numeric_dotted_version(&config.desktop_version) {
|
||||
return Err(ConfigError::InvalidField("wire_profile.desktop_version"));
|
||||
}
|
||||
if !config
|
||||
.desktop_build
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit())
|
||||
{
|
||||
return Err(ConfigError::InvalidField("wire_profile.desktop_build"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_numeric_dotted_version(value: &str) -> bool {
|
||||
let mut parts = value.split('.');
|
||||
let valid_parts = parts
|
||||
.by_ref()
|
||||
.filter(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit()))
|
||||
.count();
|
||||
valid_parts >= 2 && valid_parts == value.split('.').count()
|
||||
}
|
||||
|
||||
fn validate_logging(config: &LoggingConfig) -> Result<(), ConfigError> {
|
||||
EnvFilter::try_new(&config.level).map_err(|_| ConfigError::InvalidField("logging.level"))?;
|
||||
if !config.stdout && !config.file.enabled {
|
||||
return Err(ConfigError::InvalidField("logging"));
|
||||
}
|
||||
if config.file.directory.as_os_str().is_empty() {
|
||||
return Err(ConfigError::InvalidField("logging.file.directory"));
|
||||
}
|
||||
if config.file.retention_days == 0 {
|
||||
return Err(ConfigError::InvalidField("logging.file.retention_days"));
|
||||
}
|
||||
if config.file.max_file_size_mb == 0 {
|
||||
return Err(ConfigError::InvalidField("logging.file.max_file_size_mb"));
|
||||
}
|
||||
if config.file.max_files == 0 {
|
||||
return Err(ConfigError::InvalidField("logging.file.max_files"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_nonempty(field: &'static str, value: &str) -> Result<(), ConfigError> {
|
||||
if value.trim().is_empty() {
|
||||
return Err(ConfigError::InvalidField(field));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn connection_url_with_password(
|
||||
raw_url: &str,
|
||||
password: &SecretValue,
|
||||
field: &'static str,
|
||||
) -> Result<SecretValue, ConfigError> {
|
||||
let mut url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?;
|
||||
url.set_password(Some(password.expose_secret()))
|
||||
.map_err(|()| ConfigError::InvalidField(field))?;
|
||||
Ok(SecretBox::new(Box::new(url.into())))
|
||||
}
|
||||
|
||||
/// 启动配置加载错误。错误文本不会包含配置值或连接地址。
|
||||
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum ConfigError {
|
||||
/// 无法确定当前工作目录。
|
||||
#[error("failed to determine the current working directory")]
|
||||
CurrentDirectory,
|
||||
/// 从当前目录及父目录找不到配置文件。
|
||||
#[error("deploy/config.yaml was not found in the current directory or its parents")]
|
||||
ConfigFileNotFound,
|
||||
/// 配置路径没有父目录。
|
||||
#[error("configuration file path is invalid")]
|
||||
InvalidConfigPath,
|
||||
/// YAML 无法读取、解析或不符合结构;不透传解析器文本以避免泄露密码。
|
||||
#[error("failed to load configuration file `{path}`; check its permissions and YAML structure")]
|
||||
InvalidDocument {
|
||||
/// 配置文件路径。
|
||||
path: PathBuf,
|
||||
},
|
||||
/// 配置格式版本不受支持。
|
||||
#[error("x-cpr.schema_version must be 1")]
|
||||
UnsupportedSchemaVersion,
|
||||
/// 字段值非法。
|
||||
#[error("configuration field `{0}` is invalid")]
|
||||
InvalidField(&'static str),
|
||||
/// 连接 URL 不允许内嵌密码。
|
||||
#[error("configuration field `{0}` must not contain a password; use its password field")]
|
||||
PasswordInUrl(&'static str),
|
||||
/// PostgreSQL 或 Redis 密码格式非法。
|
||||
#[error("configuration field `{0}` must be exactly 48 hexadecimal characters")]
|
||||
InvalidServicePassword(&'static str),
|
||||
/// 管理员初始化密码非法。
|
||||
#[error("admin.default_password must be strong, at least 12 characters, and contain no `$`")]
|
||||
WeakAdminPassword,
|
||||
/// Compose 拓扑覆盖非法。
|
||||
#[error("container topology override `{0}` is invalid")]
|
||||
InvalidTopologyOverride(&'static str),
|
||||
}
|
||||
|
||||
@@ -1,460 +0,0 @@
|
||||
use futures::TryStreamExt;
|
||||
use sqlx::{types::Json, Postgres, Row, SqlitePool, Transaction};
|
||||
|
||||
use crate::infra::identity::hash_credential;
|
||||
|
||||
use super::{
|
||||
parse_cookie_expiration, parse_json, parse_optional_json, parse_optional_timestamp,
|
||||
parse_timestamp, ImportSqliteError, ImportSqliteReport,
|
||||
};
|
||||
|
||||
pub(super) async fn import_core_tables(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
import_admin_users(source, target, report).await?;
|
||||
import_client_api_keys(source, target, report).await?;
|
||||
import_runtime_settings(source, target, report).await?;
|
||||
import_accounts(source, target, report).await?;
|
||||
import_account_usage(source, target, report).await?;
|
||||
import_account_cookies(source, target, report).await?;
|
||||
import_fingerprints(source, target, report).await?;
|
||||
import_fingerprint_history(source, target, report).await
|
||||
}
|
||||
|
||||
async fn import_admin_users(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, password_hash, created_at, updated_at from admin_users order by id",
|
||||
)
|
||||
.fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at)
|
||||
values ($1, $2, $3, $4)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<String, _>("password_hash")?)
|
||||
.bind(parse_timestamp(
|
||||
"admin_users",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.bind(parse_timestamp(
|
||||
"admin_users",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("admin_users");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_client_api_keys(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, name, prefix, key, label, enabled, created_at, last_used_at
|
||||
from client_api_keys order by id",
|
||||
)
|
||||
.fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys
|
||||
(id, name, prefix, key, label, enabled, created_at, last_used_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<String, _>("name")?)
|
||||
.bind(row.try_get::<String, _>("prefix")?)
|
||||
.bind(row.try_get::<String, _>("key")?)
|
||||
.bind(row.try_get::<Option<String>, _>("label")?)
|
||||
.bind(row.try_get::<i64, _>("enabled")? == 1)
|
||||
.bind(parse_timestamp(
|
||||
"client_api_keys",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"client_api_keys",
|
||||
"last_used_at",
|
||||
row.try_get("last_used_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("client_api_keys");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_runtime_settings(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
|
||||
max_concurrent_per_account, request_interval_ms, rotation_strategy,
|
||||
admin_api_key, updated_at
|
||||
from runtime_settings order by id",
|
||||
)
|
||||
.fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let aliases: String = row.try_get("model_aliases_json")?;
|
||||
let admin_key = row
|
||||
.try_get::<Option<String>, _>("admin_api_key")?
|
||||
.and_then(|key| (!key.trim().is_empty()).then(|| hash_credential(&key)));
|
||||
sqlx::query(
|
||||
"insert into runtime_settings
|
||||
(id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
|
||||
max_concurrent_per_account, request_interval_ms, rotation_strategy,
|
||||
admin_api_key_hash, usage_retention_days, ops_error_retention_days,
|
||||
bucket_retention_days, updated_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8, 30, 30, 90, $9)",
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("id")?)
|
||||
.bind(Json(parse_json(
|
||||
"runtime_settings",
|
||||
"model_aliases_json",
|
||||
&aliases,
|
||||
)?))
|
||||
.bind(row.try_get::<i64, _>("refresh_margin_seconds")?)
|
||||
.bind(row.try_get::<i64, _>("refresh_concurrency")?)
|
||||
.bind(row.try_get::<i64, _>("max_concurrent_per_account")?)
|
||||
.bind(row.try_get::<i64, _>("request_interval_ms")?)
|
||||
.bind(row.try_get::<String, _>("rotation_strategy")?)
|
||||
.bind(admin_key)
|
||||
.bind(parse_timestamp(
|
||||
"runtime_settings",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("runtime_settings");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_accounts(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
|
||||
access_token, refresh_token, access_token_expires_at, next_refresh_at,
|
||||
status, quota_json, quota_fetched_at, quota_limit_reached,
|
||||
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
|
||||
added_at, updated_at
|
||||
from accounts order by id",
|
||||
)
|
||||
.fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let source_status = row.try_get::<String, _>("status")?;
|
||||
let (status, next_refresh_at) = if source_status == "refreshing" {
|
||||
report.add_normalized("accounts.refreshing_to_expired");
|
||||
("expired", None)
|
||||
} else {
|
||||
(
|
||||
source_status.as_str(),
|
||||
parse_optional_timestamp(
|
||||
"accounts",
|
||||
"next_refresh_at",
|
||||
row.try_get("next_refresh_at")?,
|
||||
)?,
|
||||
)
|
||||
};
|
||||
let quota = parse_optional_json(
|
||||
"accounts",
|
||||
"quota_json",
|
||||
row.try_get::<Option<String>, _>("quota_json")?,
|
||||
)?
|
||||
.map(Json);
|
||||
sqlx::query(
|
||||
"insert into accounts
|
||||
(id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
|
||||
access_token, refresh_token, access_token_expires_at, next_refresh_at,
|
||||
status, quota_json, quota_fetched_at, quota_limit_reached,
|
||||
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
|
||||
added_at, updated_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
|
||||
$14, $15, $16, $17, $18, $19)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("email")?)
|
||||
.bind(row.try_get::<Option<String>, _>("chatgpt_account_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("chatgpt_user_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("label")?)
|
||||
.bind(row.try_get::<Option<String>, _>("plan_type")?)
|
||||
.bind(row.try_get::<String, _>("access_token")?)
|
||||
.bind(row.try_get::<Option<String>, _>("refresh_token")?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"accounts",
|
||||
"access_token_expires_at",
|
||||
row.try_get("access_token_expires_at")?,
|
||||
)?)
|
||||
.bind(next_refresh_at)
|
||||
.bind(status)
|
||||
.bind(quota)
|
||||
.bind(parse_optional_timestamp(
|
||||
"accounts",
|
||||
"quota_fetched_at",
|
||||
row.try_get("quota_fetched_at")?,
|
||||
)?)
|
||||
.bind(row.try_get::<i64, _>("quota_limit_reached")? == 1)
|
||||
.bind(row.try_get::<i64, _>("quota_verify_required")? == 1)
|
||||
.bind(parse_optional_timestamp(
|
||||
"accounts",
|
||||
"quota_cooldown_until",
|
||||
row.try_get("quota_cooldown_until")?,
|
||||
)?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"accounts",
|
||||
"cloudflare_cooldown_until",
|
||||
row.try_get("cloudflare_cooldown_until")?,
|
||||
)?)
|
||||
.bind(parse_timestamp(
|
||||
"accounts",
|
||||
"added_at",
|
||||
row.try_get("added_at")?,
|
||||
)?)
|
||||
.bind(parse_timestamp(
|
||||
"accounts",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("accounts");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_account_usage(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query("select * from account_usage order by account_id").fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
sqlx::query(
|
||||
"insert into account_usage
|
||||
(account_id, request_count, empty_response_count, input_tokens, output_tokens,
|
||||
cached_tokens, reasoning_tokens, total_tokens, image_input_tokens,
|
||||
image_output_tokens, image_request_count, image_request_failed_count,
|
||||
window_request_count, window_input_tokens, window_output_tokens,
|
||||
window_cached_tokens, window_image_input_tokens, window_image_output_tokens,
|
||||
window_image_request_count, window_image_request_failed_count,
|
||||
window_started_at, window_reset_at, limit_window_seconds, last_used_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12,
|
||||
$13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("account_id")?)
|
||||
.bind(row.try_get::<i64, _>("request_count")?)
|
||||
.bind(row.try_get::<i64, _>("empty_response_count")?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cached_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("reasoning_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("total_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("image_input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("image_output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("image_request_count")?)
|
||||
.bind(row.try_get::<i64, _>("image_request_failed_count")?)
|
||||
.bind(row.try_get::<i64, _>("window_request_count")?)
|
||||
.bind(row.try_get::<i64, _>("window_input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("window_output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("window_cached_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("window_image_input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("window_image_output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("window_image_request_count")?)
|
||||
.bind(row.try_get::<i64, _>("window_image_request_failed_count")?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"account_usage",
|
||||
"window_started_at",
|
||||
row.try_get("window_started_at")?,
|
||||
)?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"account_usage",
|
||||
"window_reset_at",
|
||||
row.try_get("window_reset_at")?,
|
||||
)?)
|
||||
.bind(row.try_get::<Option<i64>, _>("limit_window_seconds")?)
|
||||
.bind(parse_optional_timestamp(
|
||||
"account_usage",
|
||||
"last_used_at",
|
||||
row.try_get("last_used_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("account_usage");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_account_cookies(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, account_id, domain, name, value, path, expires_at, updated_at
|
||||
from account_cookies order by id",
|
||||
)
|
||||
.fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let expires_at = row
|
||||
.try_get::<Option<String>, _>("expires_at")?
|
||||
.and_then(|value| match parse_cookie_expiration(&value) {
|
||||
Some(timestamp) => Some(timestamp),
|
||||
None => {
|
||||
report.add_cookie_expiration_parse_failure();
|
||||
None
|
||||
}
|
||||
});
|
||||
sqlx::query(
|
||||
"insert into account_cookies
|
||||
(id, account_id, domain, name, value, path, expires_at, updated_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<String, _>("account_id")?)
|
||||
.bind(row.try_get::<String, _>("domain")?)
|
||||
.bind(row.try_get::<String, _>("name")?)
|
||||
.bind(row.try_get::<String, _>("value")?)
|
||||
.bind(row.try_get::<String, _>("path")?)
|
||||
.bind(expires_at)
|
||||
.bind(parse_timestamp(
|
||||
"account_cookies",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("account_cookies");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_fingerprints(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query("select * from fingerprints order by id").fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let default_headers: String = row.try_get("default_headers_json")?;
|
||||
let header_order: String = row.try_get("header_order_json")?;
|
||||
sqlx::query(
|
||||
"insert into fingerprints
|
||||
(id, originator, app_version, build_number, platform, arch, chromium_version,
|
||||
user_agent_template, default_headers_json, header_order_json, source,
|
||||
created_at, updated_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<String, _>("originator")?)
|
||||
.bind(row.try_get::<String, _>("app_version")?)
|
||||
.bind(row.try_get::<String, _>("build_number")?)
|
||||
.bind(row.try_get::<String, _>("platform")?)
|
||||
.bind(row.try_get::<String, _>("arch")?)
|
||||
.bind(row.try_get::<String, _>("chromium_version")?)
|
||||
.bind(row.try_get::<String, _>("user_agent_template")?)
|
||||
.bind(Json(parse_json(
|
||||
"fingerprints",
|
||||
"default_headers_json",
|
||||
&default_headers,
|
||||
)?))
|
||||
.bind(Json(parse_json(
|
||||
"fingerprints",
|
||||
"header_order_json",
|
||||
&header_order,
|
||||
)?))
|
||||
.bind(row.try_get::<String, _>("source")?)
|
||||
.bind(parse_timestamp(
|
||||
"fingerprints",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.bind(parse_timestamp(
|
||||
"fingerprints",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("fingerprints");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_fingerprint_history(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows =
|
||||
sqlx::query("select * from fingerprint_update_history order by id").fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let manifest = parse_optional_json(
|
||||
"fingerprint_update_history",
|
||||
"manifest_json",
|
||||
row.try_get::<Option<String>, _>("manifest_json")?,
|
||||
)?
|
||||
.map(Json);
|
||||
sqlx::query(
|
||||
"insert into fingerprint_update_history
|
||||
(id, current_fingerprint_id, app_version, build_number, chromium_version,
|
||||
source, manifest_json, created_at)
|
||||
values ($1, $2, $3, $4, $5, $6, $7, $8)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<String, _>("current_fingerprint_id")?)
|
||||
.bind(row.try_get::<String, _>("app_version")?)
|
||||
.bind(row.try_get::<String, _>("build_number")?)
|
||||
.bind(row.try_get::<Option<String>, _>("chromium_version")?)
|
||||
.bind(row.try_get::<String, _>("source")?)
|
||||
.bind(manifest)
|
||||
.bind(parse_timestamp(
|
||||
"fingerprint_update_history",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("fingerprint_update_history");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn count_discarded_runtime_tables(
|
||||
source: &SqlitePool,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let (sessions, affinities, leases, snapshots, model_usage): (i64, i64, i64, i64, i64) =
|
||||
sqlx::query_as(
|
||||
"select
|
||||
(select count(*) from admin_sessions),
|
||||
(select count(*) from session_affinities),
|
||||
(select count(*) from account_refresh_leases),
|
||||
(select count(*) from model_plan_snapshots),
|
||||
(select count(*) from account_model_usage)",
|
||||
)
|
||||
.fetch_one(source)
|
||||
.await?;
|
||||
report.add_discarded("admin_sessions", sessions.max(0) as u64);
|
||||
report.add_discarded("session_affinities", affinities.max(0) as u64);
|
||||
report.add_discarded("account_refresh_leases", leases.max(0) as u64);
|
||||
report.add_discarded("model_plan_snapshots", snapshots.max(0) as u64);
|
||||
report.add_discarded("account_model_usage", model_usage.max(0) as u64);
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,232 +0,0 @@
|
||||
//! 一次性 SQLite v3 到 PostgreSQL 终态库导入。
|
||||
|
||||
mod core;
|
||||
mod telemetry;
|
||||
|
||||
use std::{collections::BTreeMap, fmt, path::Path};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::Value;
|
||||
use sqlx::{
|
||||
sqlite::{SqliteConnectOptions, SqlitePoolOptions},
|
||||
PgPool, Postgres, SqlitePool, Transaction,
|
||||
};
|
||||
use thiserror::Error;
|
||||
|
||||
const SQLITE_SCHEMA_VERSION: i64 = 3;
|
||||
|
||||
/// SQLite 历史库导入结果。
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct ImportSqliteReport {
|
||||
imported: BTreeMap<&'static str, u64>,
|
||||
normalized: BTreeMap<&'static str, u64>,
|
||||
discarded: BTreeMap<&'static str, u64>,
|
||||
cookie_expiration_parse_failures: u64,
|
||||
}
|
||||
|
||||
impl ImportSqliteReport {
|
||||
fn add_imported(&mut self, table: &'static str) {
|
||||
*self.imported.entry(table).or_default() += 1;
|
||||
}
|
||||
|
||||
fn add_discarded(&mut self, reason: &'static str, count: u64) {
|
||||
if count > 0 {
|
||||
*self.discarded.entry(reason).or_default() += count;
|
||||
}
|
||||
}
|
||||
|
||||
fn add_normalized(&mut self, reason: &'static str) {
|
||||
*self.normalized.entry(reason).or_default() += 1;
|
||||
}
|
||||
|
||||
fn add_cookie_expiration_parse_failure(&mut self) {
|
||||
self.cookie_expiration_parse_failures += 1;
|
||||
}
|
||||
|
||||
/// 返回指定目标表导入的行数。
|
||||
pub fn imported_rows(&self, table: &str) -> u64 {
|
||||
self.imported.get(table).copied().unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 返回指定原因丢弃的行数。
|
||||
pub fn discarded_rows(&self, reason: &str) -> u64 {
|
||||
self.discarded.get(reason).copied().unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 返回指定规则规范化的行数。
|
||||
pub fn normalized_rows(&self, reason: &str) -> u64 {
|
||||
self.normalized.get(reason).copied().unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 返回无法解析并降级为会话 Cookie 的过期时间数量。
|
||||
pub fn cookie_expiration_parse_failures(&self) -> u64 {
|
||||
self.cookie_expiration_parse_failures
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ImportSqliteReport {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
writeln!(formatter, "SQLite v3 import completed")?;
|
||||
for (table, count) in &self.imported {
|
||||
writeln!(formatter, " imported {table}: {count}")?;
|
||||
}
|
||||
for (reason, count) in &self.normalized {
|
||||
writeln!(formatter, " normalized {reason}: {count}")?;
|
||||
}
|
||||
for (reason, count) in &self.discarded {
|
||||
writeln!(formatter, " discarded {reason}: {count}")?;
|
||||
}
|
||||
write!(
|
||||
formatter,
|
||||
" cookie expiration parse failures: {}",
|
||||
self.cookie_expiration_parse_failures
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// SQLite 历史库导入错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ImportSqliteError {
|
||||
#[error("database operation failed: {0}")]
|
||||
Database(#[from] sqlx::Error),
|
||||
#[error("unsupported SQLite schema version {actual:?}; expected 3")]
|
||||
UnsupportedSourceVersion { actual: Option<i64> },
|
||||
#[error("target PostgreSQL database is not empty")]
|
||||
TargetNotEmpty,
|
||||
#[error("invalid timestamp in {table}.{column}: {value}")]
|
||||
InvalidTimestamp {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
value: String,
|
||||
},
|
||||
#[error("invalid JSON in {table}.{column}: {source}")]
|
||||
InvalidJson {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
#[source]
|
||||
source: serde_json::Error,
|
||||
},
|
||||
}
|
||||
|
||||
/// 将只读 SQLite v3 历史库一次性导入空的 PostgreSQL 终态库。
|
||||
pub async fn import_sqlite(
|
||||
target: &PgPool,
|
||||
source_path: impl AsRef<Path>,
|
||||
) -> Result<ImportSqliteReport, ImportSqliteError> {
|
||||
let options = SqliteConnectOptions::new()
|
||||
.filename(source_path.as_ref())
|
||||
.read_only(true)
|
||||
.create_if_missing(false);
|
||||
let source = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect_with(options)
|
||||
.await?;
|
||||
validate_source_version(&source).await?;
|
||||
|
||||
let mut transaction = target.begin().await?;
|
||||
reject_nonempty_target(&mut transaction).await?;
|
||||
|
||||
let mut report = ImportSqliteReport::default();
|
||||
core::import_core_tables(&source, &mut transaction, &mut report).await?;
|
||||
telemetry::import_telemetry_tables(&source, &mut transaction, &mut report).await?;
|
||||
core::count_discarded_runtime_tables(&source, &mut report).await?;
|
||||
|
||||
transaction.commit().await?;
|
||||
source.close().await;
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
async fn validate_source_version(source: &SqlitePool) -> Result<(), ImportSqliteError> {
|
||||
let (actual,): (Option<i64>,) = sqlx::query_as("select max(version) from schema_migrations")
|
||||
.fetch_one(source)
|
||||
.await?;
|
||||
if actual != Some(SQLITE_SCHEMA_VERSION) {
|
||||
return Err(ImportSqliteError::UnsupportedSourceVersion { actual });
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn reject_nonempty_target(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let (rows,): (i64,) = sqlx::query_as(
|
||||
"select
|
||||
(select count(*) from admin_users)
|
||||
+ (select count(*) from client_api_keys)
|
||||
+ (select count(*) from runtime_settings)
|
||||
+ (select count(*) from accounts)
|
||||
+ (select count(*) from account_usage)
|
||||
+ (select count(*) from usage_records)
|
||||
+ (select count(*) from ops_error_logs)
|
||||
+ (select count(*) from request_time_buckets)
|
||||
+ (select count(*) from account_cookies)
|
||||
+ (select count(*) from fingerprints)
|
||||
+ (select count(*) from fingerprint_update_history)",
|
||||
)
|
||||
.fetch_one(&mut **transaction)
|
||||
.await?;
|
||||
if rows != 0 {
|
||||
return Err(ImportSqliteError::TargetNotEmpty);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_timestamp(
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
value: String,
|
||||
) -> Result<DateTime<Utc>, ImportSqliteError> {
|
||||
DateTime::parse_from_rfc3339(&value)
|
||||
.map(|timestamp| timestamp.with_timezone(&Utc))
|
||||
.map_err(|_| ImportSqliteError::InvalidTimestamp {
|
||||
table,
|
||||
column,
|
||||
value,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_timestamp(
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
value: Option<String>,
|
||||
) -> Result<Option<DateTime<Utc>>, ImportSqliteError> {
|
||||
value
|
||||
.map(|value| parse_timestamp(table, column, value))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_cookie_expiration(value: &str) -> Option<DateTime<Utc>> {
|
||||
DateTime::parse_from_rfc3339(value)
|
||||
.or_else(|_| DateTime::parse_from_rfc2822(value))
|
||||
.ok()
|
||||
.map(|timestamp| timestamp.with_timezone(&Utc))
|
||||
}
|
||||
|
||||
fn parse_json(
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
value: &str,
|
||||
) -> Result<Value, ImportSqliteError> {
|
||||
serde_json::from_str(value).map_err(|source| ImportSqliteError::InvalidJson {
|
||||
table,
|
||||
column,
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_json(
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
value: Option<String>,
|
||||
) -> Result<Option<Value>, ImportSqliteError> {
|
||||
value
|
||||
.map(|value| parse_json(table, column, &value))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn normalized_dimension(value: Option<String>) -> String {
|
||||
value
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| "__unknown__".to_string())
|
||||
}
|
||||
@@ -1,330 +0,0 @@
|
||||
use futures::TryStreamExt;
|
||||
use serde_json::Value;
|
||||
use sqlx::{types::Json, Postgres, Row, SqlitePool, Transaction};
|
||||
|
||||
use super::{
|
||||
normalized_dimension, parse_json, parse_timestamp, ImportSqliteError, ImportSqliteReport,
|
||||
};
|
||||
|
||||
pub(super) async fn import_telemetry_tables(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
import_mixed_usage_records(source, target, report).await?;
|
||||
import_ops_error_logs(source, target, report).await?;
|
||||
import_request_time_buckets(source, target, report).await
|
||||
}
|
||||
|
||||
async fn import_mixed_usage_records(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows = sqlx::query(
|
||||
"select id, request_id, kind, level, account_id, route, model, status_code,
|
||||
transport, attempt_index, upstream_status_code, failure_class,
|
||||
response_id, upstream_request_id, latency_ms, message, metadata_json,
|
||||
created_at
|
||||
from usage_records order by created_at, id",
|
||||
)
|
||||
.fetch(source);
|
||||
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let level: String = row.try_get("level")?;
|
||||
let account_id: Option<String> = row.try_get("account_id")?;
|
||||
let model: Option<String> = row.try_get("model")?;
|
||||
let status_code: Option<i64> = row.try_get("status_code")?;
|
||||
let metadata_text: String = row.try_get("metadata_json")?;
|
||||
let metadata = parse_json("usage_records", "metadata_json", &metadata_text)?;
|
||||
|
||||
if level == "error" {
|
||||
insert_usage_error(target, &row, metadata).await?;
|
||||
report.add_imported("ops_error_logs");
|
||||
continue;
|
||||
}
|
||||
|
||||
let valid_account = account_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
let valid_model = model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
let valid_status = status_code.is_some_and(|status| (200..=399).contains(&status));
|
||||
if !valid_account || !valid_model || !valid_status {
|
||||
report.add_discarded("usage_record_noise", 1);
|
||||
continue;
|
||||
}
|
||||
|
||||
insert_success_record(
|
||||
target,
|
||||
&row,
|
||||
account_id.expect("validated account id"),
|
||||
model.expect("validated model"),
|
||||
status_code.expect("validated status code"),
|
||||
metadata,
|
||||
)
|
||||
.await?;
|
||||
report.add_imported("usage_records");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_success_record(
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
account_id: String,
|
||||
model: String,
|
||||
status_code: i64,
|
||||
mut metadata: Value,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let requested_model = metadata_string(&metadata, &["requestedModel", "requested_model"]);
|
||||
let upstream_model = metadata_string(&metadata, &["upstreamModel", "upstream_model"]);
|
||||
let service_tier = metadata_string(&metadata, &["serviceTier", "service_tier"]);
|
||||
let first_token_ms = metadata_nonnegative(&metadata, &["firstTokenMs", "first_token_ms"]);
|
||||
let input_tokens = usage_number(&metadata, &["inputTokens", "input_tokens"]);
|
||||
let output_tokens = usage_number(&metadata, &["outputTokens", "output_tokens"]);
|
||||
let cached_tokens = usage_number(&metadata, &["cachedTokens", "cached_tokens"]);
|
||||
let reasoning_tokens = usage_number(&metadata, &["reasoningTokens", "reasoning_tokens"]);
|
||||
strip_promoted_success_fields(&mut metadata);
|
||||
|
||||
sqlx::query(
|
||||
"insert into usage_records
|
||||
(id, request_id, client_api_key_id, kind, route, provider, account_id, model,
|
||||
requested_model, upstream_model, service_tier, status_code, transport,
|
||||
attempt_index, response_id, upstream_request_id, latency_ms, first_token_ms,
|
||||
input_tokens, output_tokens, cached_tokens, reasoning_tokens, message,
|
||||
metadata_json, created_at)
|
||||
values ($1, $2, null, $3, $4, 'openai', $5, $6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("request_id")?)
|
||||
.bind(row.try_get::<String, _>("kind")?)
|
||||
.bind(row.try_get::<Option<String>, _>("route")?)
|
||||
.bind(account_id)
|
||||
.bind(model)
|
||||
.bind(requested_model)
|
||||
.bind(upstream_model)
|
||||
.bind(service_tier)
|
||||
.bind(status_code as i32)
|
||||
.bind(row.try_get::<Option<String>, _>("transport")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
|
||||
.bind(row.try_get::<Option<String>, _>("response_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
|
||||
.bind(first_token_ms)
|
||||
.bind(input_tokens)
|
||||
.bind(output_tokens)
|
||||
.bind(cached_tokens)
|
||||
.bind(reasoning_tokens)
|
||||
.bind(row.try_get::<String, _>("message")?)
|
||||
.bind(Json(metadata))
|
||||
.bind(parse_timestamp(
|
||||
"usage_records",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_usage_error(
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
metadata: Value,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
sqlx::query(
|
||||
"insert into ops_error_logs
|
||||
(id, request_id, client_api_key_id, kind, provider, account_id, route, model,
|
||||
status_code, client_status_code, upstream_status_code, transport, attempt_index,
|
||||
failure_class, response_id, upstream_request_id, latency_ms, message,
|
||||
metadata_json, created_at)
|
||||
values ($1, $2, null, $3, 'openai', $4, $5, $6, $7, null, $8, $9,
|
||||
$10, $11, $12, $13, $14, $15, $16, $17)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("request_id")?)
|
||||
.bind(row.try_get::<String, _>("kind")?)
|
||||
.bind(row.try_get::<Option<String>, _>("account_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("route")?)
|
||||
.bind(row.try_get::<Option<String>, _>("model")?)
|
||||
.bind(optional_status(row.try_get("status_code")?))
|
||||
.bind(optional_status(row.try_get("upstream_status_code")?))
|
||||
.bind(row.try_get::<Option<String>, _>("transport")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
|
||||
.bind(row.try_get::<Option<String>, _>("failure_class")?)
|
||||
.bind(row.try_get::<Option<String>, _>("response_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
|
||||
.bind(row.try_get::<String, _>("message")?)
|
||||
.bind(Json(metadata))
|
||||
.bind(parse_timestamp(
|
||||
"usage_records",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_ops_error_logs(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows =
|
||||
sqlx::query("select * from ops_error_logs order by created_at, id").fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let metadata_text: String = row.try_get("metadata_json")?;
|
||||
sqlx::query(
|
||||
"insert into ops_error_logs
|
||||
(id, request_id, client_api_key_id, kind, provider, account_id, route, model,
|
||||
status_code, client_status_code, upstream_status_code, transport, attempt_index,
|
||||
failure_class, response_id, upstream_request_id, latency_ms, message,
|
||||
metadata_json, created_at)
|
||||
values ($1, $2, null, $3, 'openai', $4, $5, $6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15, $16, $17, $18)",
|
||||
)
|
||||
.bind(row.try_get::<String, _>("id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("request_id")?)
|
||||
.bind(row.try_get::<String, _>("kind")?)
|
||||
.bind(row.try_get::<Option<String>, _>("account_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("route")?)
|
||||
.bind(row.try_get::<Option<String>, _>("model")?)
|
||||
.bind(optional_status(row.try_get("status_code")?))
|
||||
.bind(optional_status(row.try_get("client_status_code")?))
|
||||
.bind(optional_status(row.try_get("upstream_status_code")?))
|
||||
.bind(row.try_get::<Option<String>, _>("transport")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
|
||||
.bind(row.try_get::<Option<String>, _>("failure_class")?)
|
||||
.bind(row.try_get::<Option<String>, _>("response_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
|
||||
.bind(row.try_get::<String, _>("message")?)
|
||||
.bind(Json(parse_json(
|
||||
"ops_error_logs",
|
||||
"metadata_json",
|
||||
&metadata_text,
|
||||
)?))
|
||||
.bind(parse_timestamp(
|
||||
"ops_error_logs",
|
||||
"created_at",
|
||||
row.try_get("created_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("ops_error_logs");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_request_time_buckets(
|
||||
source: &SqlitePool,
|
||||
target: &mut Transaction<'_, Postgres>,
|
||||
report: &mut ImportSqliteReport,
|
||||
) -> Result<(), ImportSqliteError> {
|
||||
let mut rows =
|
||||
sqlx::query("select * from usage_time_buckets order by bucket_start").fetch(source);
|
||||
while let Some(row) = rows.try_next().await? {
|
||||
let request_count: i64 = row.try_get("request_count")?;
|
||||
let error_count: i64 = row.try_get("error_count")?;
|
||||
let min_latency_ms: i64 = row.try_get("min_latency_ms")?;
|
||||
sqlx::query(
|
||||
"insert into request_time_buckets
|
||||
(bucket_start, provider, account_id, model, service_tier, success_count,
|
||||
error_count, input_tokens, output_tokens, cached_tokens,
|
||||
first_token_latency_sum, first_token_latency_count, latency_sum,
|
||||
latency_count, max_latency_ms, min_latency_ms, updated_at)
|
||||
values ($1, 'openai', $2, $3, $4, $5, $6, $7, $8, $9, $10, $11,
|
||||
$12, $13, $14, $15, $16)",
|
||||
)
|
||||
.bind(parse_timestamp(
|
||||
"usage_time_buckets",
|
||||
"bucket_start",
|
||||
row.try_get("bucket_start")?,
|
||||
)?)
|
||||
.bind(normalized_dimension(row.try_get("account_id")?))
|
||||
.bind(normalized_dimension(row.try_get("model")?))
|
||||
.bind(normalized_dimension(row.try_get("service_tier")?))
|
||||
.bind(request_count.saturating_sub(error_count).max(0))
|
||||
.bind(error_count)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cached_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("first_token_latency_sum")?)
|
||||
.bind(row.try_get::<i64, _>("first_token_latency_count")?)
|
||||
.bind(row.try_get::<i64, _>("latency_sum")?)
|
||||
.bind(row.try_get::<i64, _>("latency_count")?)
|
||||
.bind(row.try_get::<i64, _>("max_latency_ms")?)
|
||||
.bind((min_latency_ms != 0).then_some(min_latency_ms))
|
||||
.bind(parse_timestamp(
|
||||
"usage_time_buckets",
|
||||
"updated_at",
|
||||
row.try_get("updated_at")?,
|
||||
)?)
|
||||
.execute(&mut **target)
|
||||
.await?;
|
||||
report.add_imported("request_time_buckets");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn optional_status(status: Option<i64>) -> Option<i32> {
|
||||
status.and_then(|status| i32::try_from(status).ok())
|
||||
}
|
||||
|
||||
fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
|
||||
keys.iter().find_map(|key| {
|
||||
metadata
|
||||
.get(*key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
})
|
||||
}
|
||||
|
||||
fn metadata_nonnegative(metadata: &Value, keys: &[&str]) -> Option<i64> {
|
||||
keys.iter()
|
||||
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
|
||||
.filter(|value| *value >= 0)
|
||||
}
|
||||
|
||||
fn usage_number(metadata: &Value, keys: &[&str]) -> Option<i64> {
|
||||
metadata
|
||||
.get("usage")
|
||||
.and_then(|usage| metadata_nonnegative(usage, keys))
|
||||
.or_else(|| metadata_nonnegative(metadata, keys))
|
||||
}
|
||||
|
||||
fn strip_promoted_success_fields(metadata: &mut Value) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"usage",
|
||||
"route",
|
||||
"apiKind",
|
||||
"requestedModel",
|
||||
"requested_model",
|
||||
"upstreamModel",
|
||||
"upstream_model",
|
||||
"serviceTier",
|
||||
"service_tier",
|
||||
"statusCode",
|
||||
"transport",
|
||||
"attemptIndex",
|
||||
"responseId",
|
||||
"upstreamRequestId",
|
||||
"openaiRequestId",
|
||||
"firstTokenMs",
|
||||
"first_token_ms",
|
||||
] {
|
||||
object.remove(key);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,6 @@
|
||||
//! 进程装配、共享状态与后台任务。
|
||||
|
||||
pub mod config;
|
||||
pub mod import_sqlite;
|
||||
pub mod services;
|
||||
pub mod shutdown;
|
||||
pub mod state;
|
||||
|
||||
@@ -3,13 +3,22 @@
|
||||
use std::{env, error::Error, future::IntoFuture, sync::Arc, time::Duration};
|
||||
|
||||
use axum::Router;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::{
|
||||
api::admin::{
|
||||
accounts_routes::query::{AccountListQueryService, RefreshActivityQuery},
|
||||
dashboard_routes::{
|
||||
DashboardDesktopRelease, DashboardDesktopReleaseQuery, DashboardDesktopReleaseSnapshot,
|
||||
DashboardQueryService, DashboardQueryServiceParts, DashboardWireProfile,
|
||||
DashboardWireProfileQuery,
|
||||
},
|
||||
},
|
||||
auth::{
|
||||
service::SessionService,
|
||||
store::{PgAdminUserStore, RedisAdminSessionStore},
|
||||
},
|
||||
bootstrap::config::AppConfig,
|
||||
bootstrap::config::{AppConfig, BootstrapConfig},
|
||||
dispatch::{
|
||||
affinity::{RedisSessionAffinityStore, SessionAffinityService},
|
||||
service::{ResponseDispatchService, ResponseDispatchServiceParts},
|
||||
@@ -17,42 +26,49 @@ use crate::{
|
||||
fleet::{
|
||||
cookies::PgCookieStore,
|
||||
manage::{AccountManageService, AccountManageServiceParts},
|
||||
pool::{AccountPoolOptions, AccountPoolService, AccountPoolStaticSettings},
|
||||
pool::{
|
||||
AccountPoolOptions, AccountPoolRuntimeOptions, AccountPoolService,
|
||||
AccountPoolStaticSettings, RotationStrategy,
|
||||
},
|
||||
refresh::{
|
||||
RedisRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, TokenRefreshService,
|
||||
},
|
||||
store::{AccountStore, PgAccountStore},
|
||||
usage::AccountUsageStore,
|
||||
},
|
||||
infra::{
|
||||
database::connect,
|
||||
identity::AccountPseudonymizer,
|
||||
logging::{init_tracing, LogError, LogGuard, RotationConfig, TracingConfig},
|
||||
logging::{LogError, LogGuard, RotationConfig, TracingConfig, init_tracing},
|
||||
paths::{ensure_data_dir, load_or_create_identity_secret},
|
||||
redis::RedisConnection,
|
||||
},
|
||||
keys::{manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore},
|
||||
models::{
|
||||
gateway::ModelCatalogSource,
|
||||
service::{ModelService, ModelServiceError},
|
||||
store::{ModelSnapshotStore, RedisModelSnapshotStore},
|
||||
types::ModelConfig,
|
||||
},
|
||||
settings::{service::SettingsService, SettingsSnapshot},
|
||||
settings::{SettingsSnapshot, service::SettingsService},
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageQueryService,
|
||||
account_usage::store::{AccountUsageStore, PgAccountUsageStore},
|
||||
account_usage::store::PgAccountUsageStore,
|
||||
buckets::store::PgRequestBucketStore,
|
||||
ops::query::OpsQueryService,
|
||||
ops::store::PgOpsErrorLogStore,
|
||||
recorder::Recorder,
|
||||
usage::query::UsageQueryService,
|
||||
usage::store::{PgUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
|
||||
usage::store::{DEFAULT_USAGE_RECORD_CAPTURE_BODY, PgUsageRecordStore},
|
||||
},
|
||||
update::service::SystemUpdateService,
|
||||
upstream::openai::{
|
||||
fingerprint::{Fingerprint, PgFingerprintStore, RuntimeFingerprint},
|
||||
token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig},
|
||||
desktop_release::DesktopReleaseStatus,
|
||||
profile::{CodexWireProfile, CodexWireProfileState},
|
||||
token_client::{OpenAiTokenClient, TokenClientConfig, default_openai_token_client},
|
||||
transport::{
|
||||
build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient,
|
||||
CodexWebSocketPool, CodexWebSocketPoolConfig,
|
||||
CodexBackendClient, CodexWebSocketPool, CodexWebSocketPoolConfig, build_reqwest_client,
|
||||
tls::CustomCaError,
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -78,7 +94,6 @@ pub struct BackgroundTaskStores {
|
||||
pub admin_users: PgAdminUserStore,
|
||||
pub admin_sessions: RedisAdminSessionStore,
|
||||
pub cookies: PgCookieStore,
|
||||
pub fingerprints: PgFingerprintStore,
|
||||
pub session_affinity: RedisSessionAffinityStore,
|
||||
pub refresh_leases: RedisRefreshLeaseStore,
|
||||
pub model_snapshots: RedisModelSnapshotStore,
|
||||
@@ -115,8 +130,11 @@ pub struct Services {
|
||||
pub session_affinity: Arc<SessionAffinityService>,
|
||||
pub codex: Arc<CodexBackendClient>,
|
||||
pub websocket_pool: Option<Arc<CodexWebSocketPool>>,
|
||||
pub fingerprint: RuntimeFingerprint,
|
||||
pub wire_profile: CodexWireProfileState,
|
||||
pub desktop_release: DesktopReleaseStatus,
|
||||
pub account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
pub system_update: Arc<SystemUpdateService>,
|
||||
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
|
||||
pub background_tasks: BackgroundTaskStores,
|
||||
pub(crate) account_pool_static: AccountPoolStaticSettings,
|
||||
}
|
||||
@@ -137,25 +155,64 @@ impl UsageRecordOptions {
|
||||
}
|
||||
}
|
||||
|
||||
impl DashboardWireProfileQuery for CodexWireProfileState {
|
||||
fn snapshot(&self) -> DashboardWireProfile {
|
||||
let profile = CodexWireProfileState::snapshot(self);
|
||||
let user_agent = profile.user_agent();
|
||||
DashboardWireProfile {
|
||||
originator: profile.originator,
|
||||
codex_version: profile.codex_version,
|
||||
desktop_version: profile.desktop_version,
|
||||
desktop_build: profile.desktop_build,
|
||||
os_type: profile.os_type,
|
||||
os_version: profile.os_version,
|
||||
arch: profile.arch,
|
||||
terminal: profile.terminal,
|
||||
user_agent,
|
||||
verified_at: profile.verified_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DashboardDesktopReleaseQuery for DesktopReleaseStatus {
|
||||
fn snapshot(&self) -> DashboardDesktopReleaseSnapshot {
|
||||
let snapshot = DesktopReleaseStatus::snapshot(self);
|
||||
DashboardDesktopReleaseSnapshot {
|
||||
checked_at: snapshot.checked_at,
|
||||
latest: snapshot.latest.map(|release| DashboardDesktopRelease {
|
||||
version: release.version,
|
||||
build: release.build,
|
||||
published_at: release.published_at,
|
||||
minimum_system_version: release.minimum_system_version,
|
||||
hardware_requirements: release.hardware_requirements,
|
||||
download_url: release.download_url,
|
||||
download_size: release.download_size,
|
||||
signature_present: release.signature_present,
|
||||
}),
|
||||
last_error: snapshot.last_error,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Services {
|
||||
pub fn new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
) -> Self {
|
||||
Self::try_new(config, stores, fingerprint)
|
||||
Self::try_new(config, stores, wire_profile)
|
||||
.expect("failed to build runtime services with configured TLS transport")
|
||||
}
|
||||
|
||||
pub fn try_new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
Self::try_with_usage_record_options(
|
||||
config,
|
||||
stores,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
UsageRecordOptions::from_config(config),
|
||||
)
|
||||
}
|
||||
@@ -163,7 +220,7 @@ impl Services {
|
||||
pub fn try_with_usage_record_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
usage_record_options: UsageRecordOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
let mut identity_secret = [0u8; 32];
|
||||
@@ -171,7 +228,7 @@ impl Services {
|
||||
Self::try_with_identity_secret_and_usage_record_options(
|
||||
config,
|
||||
stores,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
identity_secret,
|
||||
usage_record_options,
|
||||
)
|
||||
@@ -180,7 +237,7 @@ impl Services {
|
||||
fn try_with_identity_secret_and_usage_record_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
identity_secret: [u8; 32],
|
||||
usage_record_options: UsageRecordOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
@@ -198,6 +255,8 @@ impl Services {
|
||||
enabled: true,
|
||||
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
|
||||
max_per_account: config.ws_pool.max_per_account,
|
||||
max_total: config.ws_pool.max_total,
|
||||
max_connecting: config.ws_pool.max_connecting,
|
||||
initial_event_timeout: (config.ws_pool.initial_event_timeout_ms > 0).then(|| {
|
||||
std::time::Duration::from_millis(config.ws_pool.initial_event_timeout_ms)
|
||||
}),
|
||||
@@ -208,9 +267,9 @@ impl Services {
|
||||
let websocket_initial_event_timeout = (config.ws_pool.initial_event_timeout_ms > 0)
|
||||
.then(|| std::time::Duration::from_millis(config.ws_pool.initial_event_timeout_ms));
|
||||
let client = CodexBackendClient::new(
|
||||
build_reqwest_client(config.tls.force_http11)?,
|
||||
build_reqwest_client()?,
|
||||
config.api.base_url.clone(),
|
||||
fingerprint.clone(),
|
||||
wire_profile.clone(),
|
||||
)
|
||||
.with_websocket_initial_event_timeout(websocket_initial_event_timeout);
|
||||
if let Some(pool) = &websocket_pool {
|
||||
@@ -262,9 +321,10 @@ impl Services {
|
||||
refresh_policy.clone(),
|
||||
token_client.clone(),
|
||||
)
|
||||
.with_refresh_lease_store(stores.refresh_leases.clone()),
|
||||
.with_refresh_lease_store(stores.refresh_leases.clone())
|
||||
.with_account_pool_sync(account_pool.clone()),
|
||||
);
|
||||
let upstream_client: Arc<dyn CodexModelCatalogClient> = codex.clone();
|
||||
let upstream_client: Arc<dyn ModelCatalogSource> = codex.clone();
|
||||
let store: Arc<dyn ModelSnapshotStore> = Arc::new(stores.model_snapshots.clone());
|
||||
let models = Arc::new(ModelService::new(
|
||||
ModelConfig {
|
||||
@@ -276,7 +336,7 @@ impl Services {
|
||||
let admin_accounts = Arc::new(AccountManageService::new(AccountManageServiceParts {
|
||||
store: stores.accounts.clone(),
|
||||
cookies: stores.cookies.clone(),
|
||||
codex: codex.clone(),
|
||||
upstream: codex.clone(),
|
||||
models: models.clone(),
|
||||
account_pool: account_pool.clone(),
|
||||
token_refresher: Arc::new(token_client),
|
||||
@@ -293,8 +353,10 @@ impl Services {
|
||||
let session_affinity =
|
||||
Arc::new(SessionAffinityService::new(stores.session_affinity.clone()));
|
||||
|
||||
let cloudflare_recovery =
|
||||
crate::dispatch::recovery::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
|
||||
let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default();
|
||||
let cloudflare_recovery = crate::dispatch::controllers::cloudflare::CloudflareRecovery::new(
|
||||
stores.cookies.clone(),
|
||||
);
|
||||
let responses = Arc::new(ResponseDispatchService::new(ResponseDispatchServiceParts {
|
||||
account_pool: account_pool.clone(),
|
||||
models: models.clone(),
|
||||
@@ -303,7 +365,11 @@ impl Services {
|
||||
recorder,
|
||||
account_identity,
|
||||
cloudflare: cloudflare_recovery,
|
||||
shutdown: connection_drain.cancellation_token(),
|
||||
}));
|
||||
let system_update = Arc::new(SystemUpdateService::from_env());
|
||||
|
||||
let desktop_release = DesktopReleaseStatus::default();
|
||||
|
||||
Ok(Self {
|
||||
database,
|
||||
@@ -325,8 +391,11 @@ impl Services {
|
||||
session_affinity,
|
||||
codex,
|
||||
websocket_pool,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
desktop_release,
|
||||
account_pseudonymizer,
|
||||
system_update,
|
||||
connection_drain,
|
||||
background_tasks: stores,
|
||||
account_pool_static,
|
||||
})
|
||||
@@ -341,13 +410,30 @@ impl Services {
|
||||
|
||||
impl From<&Services> for crate::api::router::ApiServices {
|
||||
fn from(services: &Services) -> Self {
|
||||
let refresh_activity = services.token_refresh.clone() as Arc<dyn RefreshActivityQuery>;
|
||||
let account_list = Arc::new(AccountListQueryService::new(
|
||||
services.admin_accounts.clone(),
|
||||
services.usage.clone(),
|
||||
refresh_activity.clone(),
|
||||
));
|
||||
let dashboard = Arc::new(DashboardQueryService::new(DashboardQueryServiceParts {
|
||||
accounts: services.accounts.clone(),
|
||||
usage: services.usage.clone(),
|
||||
usage_records: services.usage_records.clone(),
|
||||
account_pool: services.account_pool.clone(),
|
||||
refresh_activity,
|
||||
settings: services.settings.clone(),
|
||||
wire_profile: Arc::new(services.wire_profile.clone()),
|
||||
desktop_release: Arc::new(services.desktop_release.clone()),
|
||||
}));
|
||||
Self {
|
||||
account_list,
|
||||
dashboard,
|
||||
health_probe: Arc::new(crate::bootstrap::state::RuntimeHealthProbe::new(
|
||||
services.database.clone(),
|
||||
services.redis.clone(),
|
||||
)),
|
||||
models: services.models.clone(),
|
||||
accounts: services.accounts.clone(),
|
||||
client_keys: services.client_keys.clone(),
|
||||
admin_client_keys: services.admin_client_keys.clone(),
|
||||
admin_sessions: services.admin_sessions.clone(),
|
||||
@@ -355,13 +441,12 @@ impl From<&Services> for crate::api::router::ApiServices {
|
||||
admin_accounts: services.admin_accounts.clone(),
|
||||
usage_records: services.usage_records.clone(),
|
||||
ops_errors: services.ops_errors.clone(),
|
||||
usage: services.usage.clone(),
|
||||
account_pool: services.account_pool.clone(),
|
||||
token_refresh: services.token_refresh.clone(),
|
||||
responses: services.responses.clone(),
|
||||
session_affinity: services.session_affinity.clone(),
|
||||
fingerprint: services.fingerprint.clone(),
|
||||
process_control: Arc::new(crate::bootstrap::shutdown::RuntimeProcessControl),
|
||||
system_update: services.system_update.clone(),
|
||||
connection_drain: services.connection_drain.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -407,44 +492,62 @@ pub fn apply_settings_to_config(config: &mut AppConfig, settings: &SettingsSnaps
|
||||
config.auth.rotation_strategy = settings.rotation_strategy.clone();
|
||||
}
|
||||
|
||||
pub fn fingerprint_from_config(
|
||||
config: &crate::bootstrap::config::FingerprintConfig,
|
||||
) -> Fingerprint {
|
||||
Fingerprint {
|
||||
pub fn wire_profile_from_config(
|
||||
config: &crate::bootstrap::config::WireProfileConfig,
|
||||
) -> CodexWireProfile {
|
||||
CodexWireProfile {
|
||||
originator: config.originator.clone(),
|
||||
app_version: config.app_version.clone(),
|
||||
build_number: config.build_number.clone(),
|
||||
platform: config.platform.clone(),
|
||||
codex_version: config.codex_version.clone(),
|
||||
desktop_version: config.desktop_version.clone(),
|
||||
desktop_build: config.desktop_build.clone(),
|
||||
os_type: config.os_type.clone(),
|
||||
os_version: config.os_version.clone(),
|
||||
arch: config.arch.clone(),
|
||||
chromium_version: config.chromium_version.clone(),
|
||||
user_agent_template: config.user_agent_template.clone(),
|
||||
default_headers: config
|
||||
.default_headers
|
||||
.iter()
|
||||
.map(|header| (header.name.clone(), header.value.clone()))
|
||||
.collect(),
|
||||
header_order: config.header_order.clone(),
|
||||
updated_at: None,
|
||||
terminal: config.terminal.clone(),
|
||||
verified_at: config.verified_at,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
|
||||
AccountPoolStaticSettings {
|
||||
let static_settings = AccountPoolStaticSettings {
|
||||
skip_quota_limited: config.quota.skip_exhausted,
|
||||
tier_priority: config.auth.tier_priority.clone(),
|
||||
}
|
||||
.pool_options(&settings_snapshot_from_config(config))
|
||||
};
|
||||
account_pool_runtime_options(&settings_snapshot_from_config(config), &static_settings).options
|
||||
}
|
||||
|
||||
pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
pub fn account_pool_runtime_options(
|
||||
settings: &SettingsSnapshot,
|
||||
static_settings: &AccountPoolStaticSettings,
|
||||
) -> AccountPoolRuntimeOptions {
|
||||
AccountPoolRuntimeOptions {
|
||||
options: AccountPoolOptions {
|
||||
max_concurrent_per_account: settings.max_concurrent_per_account,
|
||||
rotation_strategy: match settings.rotation_strategy.as_str() {
|
||||
"smart" => RotationStrategy::Smart,
|
||||
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
|
||||
"round_robin" => RotationStrategy::RoundRobin,
|
||||
"sticky" => RotationStrategy::Sticky,
|
||||
_ => RotationStrategy::Smart,
|
||||
},
|
||||
skip_quota_limited: static_settings.skip_quota_limited,
|
||||
tier_priority: static_settings.tier_priority.clone(),
|
||||
..AccountPoolOptions::default()
|
||||
},
|
||||
request_interval_ms: settings.request_interval_ms,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
wait_for_scheduled_restart().await;
|
||||
|
||||
let config = AppConfig::load()?;
|
||||
let host = config.server.host.clone();
|
||||
let port = config.server.port;
|
||||
let log_guard = init_logging(&config)?;
|
||||
let host = config.app().server.host.clone();
|
||||
let port = config.app().server.port;
|
||||
let log_guard = init_logging(config.app())?;
|
||||
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
|
||||
let (router, task_coordinator) = build_router(config).await?;
|
||||
let (router, task_coordinator, connection_drain) = build_router(config).await?;
|
||||
let mut task_coordinator = Some(task_coordinator);
|
||||
let mut restart_path = None;
|
||||
|
||||
let serve_result = {
|
||||
let (graceful_shutdown_tx, graceful_shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
@@ -459,66 +562,109 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
tokio::pin!(serve);
|
||||
tokio::select! {
|
||||
result = &mut serve => result,
|
||||
() = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
let _ = graceful_shutdown_tx.send(());
|
||||
tokio::select! {
|
||||
result = &mut serve => result,
|
||||
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
|
||||
tracing::warn!(
|
||||
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
|
||||
"HTTP graceful shutdown timed out; cancelling remaining connections"
|
||||
);
|
||||
action = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
match action {
|
||||
crate::bootstrap::shutdown::ShutdownAction::Restart(executable_path) => {
|
||||
restart_path = Some(executable_path);
|
||||
Ok(())
|
||||
}
|
||||
() = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
tracing::warn!("received a second shutdown signal; cancelling remaining connections");
|
||||
Ok(())
|
||||
crate::bootstrap::shutdown::ShutdownAction::Graceful => {
|
||||
let _ = graceful_shutdown_tx.send(());
|
||||
let active_websocket_connections = connection_drain.begin_shutdown();
|
||||
tracing::info!(
|
||||
active_websocket_connections,
|
||||
"正在排空入站连接并关闭运行时任务"
|
||||
);
|
||||
let coordinator = task_coordinator.take();
|
||||
let connection_drain = connection_drain.clone();
|
||||
let runtime_shutdown = async move {
|
||||
let shutdown_tasks = async move {
|
||||
if let Some(coordinator) = coordinator {
|
||||
coordinator.shutdown().await;
|
||||
}
|
||||
};
|
||||
tokio::join!(shutdown_tasks, connection_drain.wait());
|
||||
};
|
||||
let drain_runtime = async {
|
||||
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
|
||||
serve_result
|
||||
};
|
||||
tokio::pin!(drain_runtime);
|
||||
tokio::select! {
|
||||
result = &mut drain_runtime => result,
|
||||
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
|
||||
tracing::warn!(
|
||||
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
|
||||
"HTTP graceful shutdown timed out; cancelling remaining connections"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
_ = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
task_coordinator.shutdown().await;
|
||||
serve_result?;
|
||||
|
||||
if let Some(executable_path) = crate::bootstrap::shutdown::restart_executable_path() {
|
||||
if let Some(executable_path) = restart_path {
|
||||
tracing::info!(
|
||||
executable = %executable_path.display(),
|
||||
"正在以更新后的二进制替换当前进程"
|
||||
);
|
||||
drop(log_guard);
|
||||
exec_replacement_process(executable_path)?;
|
||||
return exec_replacement_process(executable_path);
|
||||
}
|
||||
|
||||
if let Some(coordinator) = task_coordinator.take() {
|
||||
let active_websocket_connections = connection_drain.begin_shutdown();
|
||||
tracing::info!(
|
||||
active_websocket_connections,
|
||||
"HTTP 服务已退出,正在关闭其余运行时任务"
|
||||
);
|
||||
tokio::join!(coordinator.shutdown(), connection_drain.wait());
|
||||
}
|
||||
serve_result?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const HTTP_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
|
||||
|
||||
async fn build_router(
|
||||
mut config: AppConfig,
|
||||
bootstrap_config: BootstrapConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Router,
|
||||
crate::bootstrap::tasks::coordinator::TaskCoordinator,
|
||||
crate::api::middleware::connection_drain::ConnectionDrain,
|
||||
),
|
||||
Box<dyn Error + Send + Sync>,
|
||||
> {
|
||||
let pool = connect(&config.database.url).await?;
|
||||
let redis = RedisConnection::connect(&config.redis.url, "cpr").await?;
|
||||
let crate::bootstrap::config::BootstrapConfigParts {
|
||||
app: mut config,
|
||||
database_url,
|
||||
redis_url,
|
||||
admin_default_password,
|
||||
} = bootstrap_config.into_parts();
|
||||
|
||||
let pool = connect(database_url.expose_secret()).await?;
|
||||
drop(database_url);
|
||||
let redis = RedisConnection::connect(redis_url.expose_secret(), "cpr").await?;
|
||||
drop(redis_url);
|
||||
let settings =
|
||||
SettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool).await?;
|
||||
apply_settings_to_config(&mut config, &settings);
|
||||
|
||||
let fingerprint_store = PgFingerprintStore::new(pool.clone());
|
||||
let stores = BackgroundTaskStores {
|
||||
redis: redis.clone(),
|
||||
accounts: PgAccountStore::new(pool.clone()),
|
||||
admin_users: PgAdminUserStore::new(pool.clone()),
|
||||
admin_sessions: RedisAdminSessionStore::new(redis.clone()),
|
||||
cookies: PgCookieStore::new(pool.clone()),
|
||||
fingerprints: fingerprint_store.clone(),
|
||||
session_affinity: RedisSessionAffinityStore::new(redis.clone()),
|
||||
refresh_leases: RedisRefreshLeaseStore::new(redis.clone()),
|
||||
model_snapshots: RedisModelSnapshotStore::new(redis.clone()),
|
||||
@@ -529,20 +675,14 @@ async fn build_router(
|
||||
request_buckets: PgRequestBucketStore::new(pool),
|
||||
};
|
||||
|
||||
config.admin.validate_default_password()?;
|
||||
let default_admin_password = config.admin.default_password.clone();
|
||||
let data_dir = ensure_data_dir()?;
|
||||
let data_dir = ensure_data_dir(&config.runtime.data_directory)?;
|
||||
let identity_secret = load_or_create_identity_secret(&data_dir)?;
|
||||
let default_fingerprint = fingerprint_from_config(&config.fingerprint);
|
||||
let runtime_fingerprint = fingerprint_store
|
||||
.ensure_current_seed(&default_fingerprint)
|
||||
.await?;
|
||||
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
|
||||
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(config.clone());
|
||||
let wire_profile = CodexWireProfileState::new(wire_profile_from_config(&config.wire_profile));
|
||||
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(&config);
|
||||
let services = Services::try_with_identity_secret_and_usage_record_options(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint,
|
||||
wire_profile,
|
||||
identity_secret,
|
||||
UsageRecordOptions::from_config(&config),
|
||||
)?;
|
||||
@@ -550,8 +690,9 @@ async fn build_router(
|
||||
|
||||
let created_default_admin = services
|
||||
.admin_sessions
|
||||
.ensure_default_admin(&default_admin_password)
|
||||
.ensure_default_admin(admin_default_password.expose_secret())
|
||||
.await?;
|
||||
drop(admin_default_password);
|
||||
tracing::info!(
|
||||
created = created_default_admin,
|
||||
"默认管理员账号已完成初始化检查"
|
||||
@@ -566,11 +707,15 @@ async fn build_router(
|
||||
let task_coordinator =
|
||||
crate::bootstrap::tasks::coordinator::TaskCoordinator::start(&runtime_config, &services);
|
||||
let app_state = crate::api::AppState::from(&services);
|
||||
let connection_drain = services.connection_drain.clone();
|
||||
let router = crate::api::router::router().with_state(app_state).layer(
|
||||
axum::middleware::from_fn_with_state(
|
||||
connection_drain.clone(),
|
||||
crate::api::middleware::connection_drain::drain_response_body,
|
||||
),
|
||||
);
|
||||
|
||||
Ok((
|
||||
crate::api::router::router().with_state(app_state),
|
||||
task_coordinator,
|
||||
))
|
||||
Ok((router, task_coordinator, connection_drain))
|
||||
}
|
||||
|
||||
fn init_logging(config: &AppConfig) -> Result<LogGuard, LogError> {
|
||||
|
||||
@@ -1,19 +1,17 @@
|
||||
//! 优雅关闭信号处理。
|
||||
//! 进程关闭与重启信号处理。
|
||||
|
||||
use std::{
|
||||
path::PathBuf,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
OnceLock,
|
||||
},
|
||||
};
|
||||
use std::{path::PathBuf, sync::OnceLock};
|
||||
|
||||
use tokio::signal;
|
||||
use tokio::sync::broadcast::{self, Receiver};
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
|
||||
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);
|
||||
static RESTART_EXECUTABLE_PATH: OnceLock<PathBuf> = OnceLock::new();
|
||||
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<ShutdownAction>> = OnceLock::new();
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum ShutdownAction {
|
||||
Graceful,
|
||||
Restart(PathBuf),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct RuntimeProcessControl;
|
||||
@@ -24,16 +22,12 @@ impl crate::api::router::ProcessControl for RuntimeProcessControl {
|
||||
}
|
||||
|
||||
fn request_restart(&self, executable_path: PathBuf) {
|
||||
request_process_restart(executable_path);
|
||||
}
|
||||
|
||||
fn subscribe_shutdown(&self) -> Receiver<()> {
|
||||
shutdown_subscription()
|
||||
request_immediate_restart(executable_path);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等待进程关闭信号(Ctrl+C 或 SIGTERM)。
|
||||
pub async fn shutdown_signal() {
|
||||
/// 等待进程关闭或重启信号。
|
||||
pub(crate) async fn shutdown_signal() -> ShutdownAction {
|
||||
let mut requested = shutdown_sender().subscribe();
|
||||
let ctrl_c = async {
|
||||
signal::ctrl_c()
|
||||
@@ -53,38 +47,25 @@ pub async fn shutdown_signal() {
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = requested.recv() => {},
|
||||
() = ctrl_c => {},
|
||||
() = terminate => {},
|
||||
action = requested.recv() => match action {
|
||||
Ok(action) => action,
|
||||
Err(_) => ShutdownAction::Graceful,
|
||||
},
|
||||
() = ctrl_c => ShutdownAction::Graceful,
|
||||
() = terminate => ShutdownAction::Graceful,
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求进程按优雅关闭路径退出。
|
||||
pub fn request_shutdown() {
|
||||
let _ = shutdown_sender().send(());
|
||||
let _ = shutdown_sender().send(ShutdownAction::Graceful);
|
||||
}
|
||||
|
||||
/// 订阅进程关闭请求。
|
||||
pub fn shutdown_subscription() -> Receiver<()> {
|
||||
shutdown_sender().subscribe()
|
||||
fn request_immediate_restart(executable_path: PathBuf) {
|
||||
let _ = shutdown_sender().send(ShutdownAction::Restart(executable_path));
|
||||
}
|
||||
|
||||
/// 请求进程完成优雅关闭后以新二进制替换当前进程。
|
||||
pub fn request_process_restart(executable_path: PathBuf) {
|
||||
let _ = RESTART_EXECUTABLE_PATH.set(executable_path);
|
||||
RESTART_REQUESTED.store(true, Ordering::SeqCst);
|
||||
request_shutdown();
|
||||
}
|
||||
|
||||
/// 获取待 exec 的重启目标。
|
||||
pub fn restart_executable_path() -> Option<PathBuf> {
|
||||
RESTART_REQUESTED
|
||||
.load(Ordering::SeqCst)
|
||||
.then(|| RESTART_EXECUTABLE_PATH.get().cloned())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
fn shutdown_sender() -> &'static broadcast::Sender<()> {
|
||||
fn shutdown_sender() -> &'static broadcast::Sender<ShutdownAction> {
|
||||
SHUTDOWN_REQUESTS.get_or_init(|| {
|
||||
let (sender, _receiver) = broadcast::channel(16);
|
||||
sender
|
||||
|
||||
@@ -9,7 +9,6 @@ use crate::{
|
||||
/// 运行时配置镜像(从 AppConfig 衍生)。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RuntimeConfig {
|
||||
pub admin: crate::bootstrap::config::AdminConfig,
|
||||
pub auth: crate::bootstrap::config::AuthConfig,
|
||||
pub quota: crate::bootstrap::config::QuotaConfig,
|
||||
}
|
||||
@@ -40,12 +39,11 @@ impl HealthProbe for RuntimeHealthProbe {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AppConfig> for RuntimeConfig {
|
||||
fn from(config: AppConfig) -> Self {
|
||||
impl From<&AppConfig> for RuntimeConfig {
|
||||
fn from(config: &AppConfig) -> Self {
|
||||
Self {
|
||||
admin: config.admin.clone(),
|
||||
auth: config.auth.clone(),
|
||||
quota: config.quota,
|
||||
quota: config.quota.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use tracing::{debug, info, warn};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
/// 支持按时间清理过期记录的存储。
|
||||
|
||||
@@ -1,19 +1,13 @@
|
||||
//! 后台任务协调器。
|
||||
|
||||
use std::{future::Future, path::PathBuf, sync::Arc, time::Duration};
|
||||
use std::{future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use super::{
|
||||
cookie_cleanup::CookieCleanupTask,
|
||||
fingerprint_update::{
|
||||
start_fingerprint_update_task, CODEX_DESKTOP_APPCAST_URL,
|
||||
DEFAULT_EXTRACTED_FINGERPRINT_PATH,
|
||||
},
|
||||
model_refresh::ModelRefreshTask,
|
||||
quota_refresh::QuotaRefreshTask,
|
||||
retention_trim::RetentionTrimTask,
|
||||
token_refresh::TokenRefreshTask,
|
||||
cookie_cleanup::CookieCleanupTask, desktop_release_update::start_desktop_release_update_task,
|
||||
model_refresh::ModelRefreshTask, quota_refresh::QuotaRefreshTask,
|
||||
retention_trim::RetentionTrimTask, token_refresh::TokenRefreshTask,
|
||||
};
|
||||
|
||||
/// 后台任务协调器。
|
||||
@@ -46,7 +40,6 @@ impl TaskCoordinator {
|
||||
stores.usage_records.clone(),
|
||||
stores.ops_errors.clone(),
|
||||
stores.request_buckets.clone(),
|
||||
stores.fingerprints.clone(),
|
||||
)
|
||||
.start(),
|
||||
);
|
||||
@@ -82,6 +75,7 @@ impl TaskCoordinator {
|
||||
stores.accounts.clone(),
|
||||
stores.account_usage.clone(),
|
||||
services.codex.clone(),
|
||||
services.account_pool.clone(),
|
||||
services.account_pseudonymizer.clone(),
|
||||
config
|
||||
.quota
|
||||
@@ -91,20 +85,16 @@ impl TaskCoordinator {
|
||||
30 * 60,
|
||||
)
|
||||
.with_cookie_store(stores.cookies.clone())
|
||||
.with_account_pool(services.account_pool.clone())
|
||||
.start(),
|
||||
);
|
||||
coordinator.push("fingerprint_update", {
|
||||
let fingerprint = services.fingerprint.snapshot();
|
||||
start_fingerprint_update_task(
|
||||
stores.fingerprints.clone(),
|
||||
services.fingerprint.clone(),
|
||||
CODEX_DESKTOP_APPCAST_URL.to_string(),
|
||||
PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH),
|
||||
fingerprint.app_version,
|
||||
fingerprint.build_number,
|
||||
)
|
||||
});
|
||||
coordinator.push(
|
||||
"desktop_release_update",
|
||||
start_desktop_release_update_task(
|
||||
services.desktop_release.clone(),
|
||||
services.wire_profile.clone(),
|
||||
crate::upstream::openai::desktop_release::CODEX_DESKTOP_APPCAST_URL.to_string(),
|
||||
),
|
||||
);
|
||||
if let Some(pool) = &services.websocket_pool {
|
||||
coordinator.push(
|
||||
"websocket_pool",
|
||||
@@ -118,30 +108,52 @@ impl TaskCoordinator {
|
||||
/// 启动各领域的设置订阅任务。
|
||||
pub fn start_settings_subscriptions(services: &crate::bootstrap::services::Services) -> Self {
|
||||
let mut coordinator = TaskCoordinator::default();
|
||||
let mut settings = services.settings.subscribe();
|
||||
let initial = settings.borrow().clone();
|
||||
let (model_tx, model_rx) = tokio::sync::watch::channel(crate::models::types::ModelConfig {
|
||||
model_aliases: initial.model_aliases.clone(),
|
||||
});
|
||||
let (pool_tx, pool_rx) =
|
||||
tokio::sync::watch::channel(crate::bootstrap::services::account_pool_runtime_options(
|
||||
&initial,
|
||||
&services.account_pool_static,
|
||||
));
|
||||
let (refresh_tx, refresh_rx) =
|
||||
tokio::sync::watch::channel(crate::fleet::refresh::RefreshPolicy {
|
||||
refresh_margin_seconds: initial.refresh_margin_seconds,
|
||||
refresh_concurrency: initial.refresh_concurrency,
|
||||
});
|
||||
coordinator.push(
|
||||
"model_settings",
|
||||
SchedulerHandle::spawn(
|
||||
services
|
||||
.models
|
||||
.clone()
|
||||
.subscribe_settings(services.settings.subscribe()),
|
||||
),
|
||||
SchedulerHandle::spawn(services.models.clone().subscribe_config(model_rx)),
|
||||
);
|
||||
coordinator.push(
|
||||
"account_pool_settings",
|
||||
SchedulerHandle::spawn(services.account_pool.clone().subscribe_settings(
|
||||
services.settings.subscribe(),
|
||||
services.account_pool_static.clone(),
|
||||
)),
|
||||
SchedulerHandle::spawn(services.account_pool.clone().subscribe_options(pool_rx)),
|
||||
);
|
||||
coordinator.push(
|
||||
"refresh_policy_settings",
|
||||
SchedulerHandle::spawn(
|
||||
services
|
||||
.refresh_policy
|
||||
.clone()
|
||||
.subscribe_settings(services.settings.subscribe()),
|
||||
),
|
||||
SchedulerHandle::spawn(services.refresh_policy.clone().subscribe_policy(refresh_rx)),
|
||||
);
|
||||
let static_settings = services.account_pool_static.clone();
|
||||
coordinator.push(
|
||||
"settings_fanout",
|
||||
SchedulerHandle::spawn(async move {
|
||||
while settings.changed().await.is_ok() {
|
||||
let snapshot = settings.borrow_and_update().clone();
|
||||
model_tx.send_replace(crate::models::types::ModelConfig {
|
||||
model_aliases: snapshot.model_aliases.clone(),
|
||||
});
|
||||
pool_tx.send_replace(crate::bootstrap::services::account_pool_runtime_options(
|
||||
&snapshot,
|
||||
&static_settings,
|
||||
));
|
||||
refresh_tx.send_replace(crate::fleet::refresh::RefreshPolicy {
|
||||
refresh_margin_seconds: snapshot.refresh_margin_seconds,
|
||||
refresh_concurrency: snapshot.refresh_concurrency,
|
||||
});
|
||||
}
|
||||
}),
|
||||
);
|
||||
coordinator
|
||||
}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
//! Codex Desktop 官方发布观测任务接线。
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::openai::desktop_release::{
|
||||
APPCAST_POLL_INTERVAL, DesktopReleaseChecker, DesktopReleaseStatus,
|
||||
};
|
||||
use crate::upstream::openai::profile::CodexWireProfileState;
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
/// 启动 Desktop 官方发布检查后台任务。
|
||||
pub fn start_desktop_release_update_task(
|
||||
status: DesktopReleaseStatus,
|
||||
wire_profile: CodexWireProfileState,
|
||||
appcast_url: String,
|
||||
) -> SchedulerHandle {
|
||||
let checker = DesktopReleaseChecker::with_client(
|
||||
reqwest::Client::new(),
|
||||
appcast_url,
|
||||
status,
|
||||
wire_profile,
|
||||
);
|
||||
spawn_periodic_task(
|
||||
DesktopReleaseUpdateTask {
|
||||
checker,
|
||||
first_tick: true,
|
||||
},
|
||||
PeriodicTaskConfig::new(
|
||||
APPCAST_POLL_INTERVAL.as_secs(),
|
||||
"Codex Desktop 发布检查器已启动",
|
||||
"Codex Desktop 发布检查器已关闭",
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
struct DesktopReleaseUpdateTask {
|
||||
checker: DesktopReleaseChecker,
|
||||
first_tick: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for DesktopReleaseUpdateTask {
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let first_tick = std::mem::take(&mut self.first_tick);
|
||||
match self.checker.check_and_record().await {
|
||||
Ok(release) => info!(
|
||||
version = %release.version,
|
||||
build = %release.build,
|
||||
"Codex Desktop 最新发布信息已刷新并已同步请求身份"
|
||||
),
|
||||
Err(error) if first_tick => {
|
||||
warn!(error = %error, "Codex Desktop 首次发布检查失败");
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error, "Codex Desktop 定期发布检查失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,85 +0,0 @@
|
||||
//! 指纹更新任务接线。
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::openai::fingerprint::{
|
||||
PgFingerprintStore, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
|
||||
};
|
||||
|
||||
/// Codex Desktop 官方 appcast 地址。
|
||||
pub const CODEX_DESKTOP_APPCAST_URL: &str =
|
||||
"https://persistent.oaistatic.com/codex-app-prod/appcast.xml";
|
||||
/// 本地可选完整指纹提取文件路径。
|
||||
pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint.json";
|
||||
|
||||
/// 启动指纹自动更新后台任务。
|
||||
pub fn start_fingerprint_update_task(
|
||||
store: PgFingerprintStore,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
appcast_url: String,
|
||||
extracted_fingerprint_path: PathBuf,
|
||||
current_version: String,
|
||||
current_build: String,
|
||||
) -> SchedulerHandle {
|
||||
let update_checker = UpdateChecker::with_client(
|
||||
store,
|
||||
reqwest::Client::new(),
|
||||
appcast_url,
|
||||
extracted_fingerprint_path,
|
||||
current_version,
|
||||
current_build,
|
||||
);
|
||||
spawn_periodic_task(
|
||||
FingerprintUpdateTask {
|
||||
checker: update_checker,
|
||||
runtime_fingerprint,
|
||||
first_tick: true,
|
||||
},
|
||||
PeriodicTaskConfig::new(
|
||||
APPCAST_POLL_INTERVAL.as_secs(),
|
||||
"fingerprint 后台版本检查器已启动",
|
||||
"fingerprint 后台版本检查器已关闭",
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
struct FingerprintUpdateTask {
|
||||
checker: UpdateChecker,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
first_tick: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for FingerprintUpdateTask {
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let first_tick = std::mem::take(&mut self.first_tick);
|
||||
match self.checker.check_and_apply_update().await {
|
||||
Ok(Some(fingerprint)) => {
|
||||
let app_version = fingerprint.app_version.clone();
|
||||
let build_number = fingerprint.build_number.clone();
|
||||
self.runtime_fingerprint.replace(fingerprint);
|
||||
info!(
|
||||
app_version = %app_version,
|
||||
build_number = %build_number,
|
||||
"fingerprint 运行时版本已更新"
|
||||
);
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(error) => {
|
||||
if first_tick {
|
||||
warn!(error = %error, "fingerprint 首次版本检查失败");
|
||||
} else {
|
||||
warn!(error = %error, "fingerprint 定期版本检查失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@
|
||||
mod cleanup;
|
||||
pub mod cookie_cleanup;
|
||||
pub mod coordinator;
|
||||
pub mod fingerprint_update;
|
||||
pub mod desktop_release_update;
|
||||
pub mod model_refresh;
|
||||
mod periodic;
|
||||
pub mod quota_refresh;
|
||||
|
||||
@@ -13,23 +13,18 @@ use crate::models::service::{
|
||||
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup},
|
||||
};
|
||||
use super::coordinator::SchedulerHandle;
|
||||
|
||||
/// 模型刷新任务接线器。
|
||||
pub struct ModelRefreshTask {
|
||||
model_service: Arc<ModelService>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
refresh_interval_secs: u64,
|
||||
initial_delay_ms: u64,
|
||||
retry_delay_ms: u64,
|
||||
max_retries: u32,
|
||||
account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
}
|
||||
|
||||
const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600;
|
||||
const INITIAL_DELAY_MS: u64 = 1000;
|
||||
const RETRY_DELAY_MS: u64 = 10_000;
|
||||
const MAX_RETRIES: u32 = 12;
|
||||
@@ -44,7 +39,6 @@ impl ModelRefreshTask {
|
||||
Self {
|
||||
model_service,
|
||||
account_pool,
|
||||
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
|
||||
initial_delay_ms: INITIAL_DELAY_MS,
|
||||
retry_delay_ms: RETRY_DELAY_MS,
|
||||
max_retries: MAX_RETRIES,
|
||||
@@ -52,15 +46,13 @@ impl ModelRefreshTask {
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动后台刷新任务。
|
||||
/// 启动一次性首次刷新任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.refresh_interval_secs,
|
||||
"模型刷新任务已启动",
|
||||
"模型刷新任务正在关闭",
|
||||
)
|
||||
.wait_first_interval();
|
||||
spawn_periodic_task(self, config)
|
||||
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
|
||||
let handle = tokio::spawn(async move {
|
||||
self.run_initial_refresh(&mut shutdown_rx).await;
|
||||
});
|
||||
SchedulerHandle::new(shutdown_tx, handle)
|
||||
}
|
||||
|
||||
/// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。
|
||||
@@ -115,84 +107,52 @@ impl ModelRefreshTask {
|
||||
.await;
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for ModelRefreshTask {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(async move {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
}
|
||||
|
||||
let mut attempt = 0;
|
||||
let mut has_fetched_once = false;
|
||||
let mut deferred_until_accounts_available = false;
|
||||
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
has_fetched_once = true;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
attempt += 1;
|
||||
if matches!(error, ModelServiceError::NoAccounts) {
|
||||
info!("首次模型刷新跳过:没有可用账户,将进入周期重试");
|
||||
deferred_until_accounts_available = true;
|
||||
break;
|
||||
}
|
||||
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !has_fetched_once && !deferred_until_accounts_available {
|
||||
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
|
||||
}
|
||||
|
||||
PeriodicTaskStartup::Continue
|
||||
})
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let mut attempt = 0;
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(_) => {}
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(ModelServiceError::NoAccounts) => {
|
||||
info!("模型刷新跳过:没有可用账户");
|
||||
info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发");
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = ?error, "刷新模型列表失败");
|
||||
attempt += 1;
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle;
|
||||
|
||||
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
|
||||
/// 周期任务进入循环前的决策。
|
||||
pub(crate) enum PeriodicTaskStartup {
|
||||
/// 继续进入周期循环。
|
||||
Continue,
|
||||
/// 任务已自行结束,不再进入周期循环。
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// 周期任务调度配置。
|
||||
pub(crate) struct PeriodicTaskConfig {
|
||||
interval_secs: u64,
|
||||
started_message: &'static str,
|
||||
stopped_message: &'static str,
|
||||
run_immediately: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskConfig {
|
||||
@@ -36,25 +27,12 @@ impl PeriodicTaskConfig {
|
||||
interval_secs,
|
||||
started_message,
|
||||
stopped_message,
|
||||
run_immediately: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn wait_first_interval(mut self) -> Self {
|
||||
self.run_immediately = false;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 周期任务行为。
|
||||
pub(crate) trait PeriodicTaskRunner: Send + 'static {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
_shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> TaskFuture<'_, ()>;
|
||||
|
||||
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
|
||||
@@ -75,17 +53,7 @@ where
|
||||
"{}", config.started_message
|
||||
);
|
||||
|
||||
if matches!(
|
||||
task.before_loop(&mut shutdown_rx).await,
|
||||
PeriodicTaskStartup::Stop
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut ticker = interval(Duration::from_secs(config.interval_secs));
|
||||
if !config.run_immediately {
|
||||
ticker.tick().await;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
|
||||
@@ -10,7 +10,8 @@ use crate::{
|
||||
cookies::PgCookieStore,
|
||||
pool::AccountPoolService,
|
||||
quota::{QuotaRefreshService, QuotaRefreshSummary},
|
||||
store::PgAccountStore,
|
||||
store::{AccountStore, PgAccountStore},
|
||||
usage::AccountUsageStore,
|
||||
},
|
||||
infra::identity::AccountPseudonymizer,
|
||||
telemetry::account_usage::store::PgAccountUsageStore,
|
||||
@@ -19,7 +20,7 @@ use crate::{
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
/// 主动配额刷新后台任务。
|
||||
@@ -35,15 +36,17 @@ impl QuotaRefreshTask {
|
||||
store: PgAccountStore,
|
||||
usage_store: PgAccountUsageStore,
|
||||
codex: Arc<CodexBackendClient>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
interval_secs: u64,
|
||||
min_refresh_interval_secs: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
service: QuotaRefreshService::with_min_refresh_interval_secs(
|
||||
store,
|
||||
usage_store,
|
||||
Arc::new(store) as Arc<dyn AccountStore>,
|
||||
Arc::new(usage_store) as Arc<dyn AccountUsageStore>,
|
||||
codex,
|
||||
account_pool,
|
||||
account_pseudonymizer,
|
||||
min_refresh_interval_secs,
|
||||
),
|
||||
@@ -58,18 +61,12 @@ impl QuotaRefreshTask {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置运行时账号池,用于刷新后同步内存调度状态。
|
||||
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
|
||||
self.service = self.service.with_account_pool(account_pool);
|
||||
self
|
||||
}
|
||||
|
||||
/// 启动后台刷新任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.interval_secs,
|
||||
"quota 刷新任务已启动",
|
||||
"quota 刷新任务已关闭",
|
||||
"Quota 刷新任务已启动",
|
||||
"Quota 刷新任务已关闭",
|
||||
);
|
||||
spawn_periodic_task(self, config)
|
||||
}
|
||||
@@ -96,7 +93,7 @@ impl PeriodicTaskRunner for QuotaRefreshTask {
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error, "quota 刷新任务失败");
|
||||
warn!(error = %error, "Quota 刷新任务失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
@@ -107,6 +104,6 @@ fn log_refreshed_summary(summary: QuotaRefreshSummary) {
|
||||
info!(
|
||||
refreshed = summary.refreshed,
|
||||
failed = summary.failed,
|
||||
"quota 刷新任务完成"
|
||||
"Quota 刷新任务完成"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -3,27 +3,22 @@
|
||||
use chrono::Utc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::{
|
||||
telemetry::{
|
||||
buckets::store::PgRequestBucketStore, ops::store::PgOpsErrorLogStore,
|
||||
usage::store::PgUsageRecordStore,
|
||||
},
|
||||
upstream::openai::fingerprint::PgFingerprintStore,
|
||||
use crate::telemetry::{
|
||||
buckets::store::PgRequestBucketStore, ops::store::PgOpsErrorLogStore,
|
||||
usage::store::PgUsageRecordStore,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, TaskFuture},
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, TaskFuture, spawn_periodic_task},
|
||||
};
|
||||
|
||||
const RETENTION_TRIM_INTERVAL_SECS: u64 = 60 * 60;
|
||||
const FINGERPRINT_HISTORY_LIMIT: std::num::NonZeroU32 = std::num::NonZeroU32::new(100).unwrap();
|
||||
|
||||
pub struct RetentionTrimTask {
|
||||
usage_records: PgUsageRecordStore,
|
||||
ops_errors: PgOpsErrorLogStore,
|
||||
buckets: PgRequestBucketStore,
|
||||
fingerprints: PgFingerprintStore,
|
||||
}
|
||||
|
||||
impl RetentionTrimTask {
|
||||
@@ -31,13 +26,11 @@ impl RetentionTrimTask {
|
||||
usage_records: PgUsageRecordStore,
|
||||
ops_errors: PgOpsErrorLogStore,
|
||||
buckets: PgRequestBucketStore,
|
||||
fingerprints: PgFingerprintStore,
|
||||
) -> Self {
|
||||
Self {
|
||||
usage_records,
|
||||
ops_errors,
|
||||
buckets,
|
||||
fingerprints,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -70,17 +63,6 @@ impl RetentionTrimTask {
|
||||
Ok(_) => {}
|
||||
Err(error) => warn!(error = %error, "清理请求时间桶失败"),
|
||||
}
|
||||
match self
|
||||
.fingerprints
|
||||
.trim_update_history_to_limit(FINGERPRINT_HISTORY_LIMIT)
|
||||
.await
|
||||
{
|
||||
Ok(deleted) if deleted > 0 => {
|
||||
info!(deleted, "已清理超限指纹更新历史")
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(error) => warn!(error = %error, "清理超限指纹更新历史失败"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,15 +2,15 @@
|
||||
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::fleet::refresh::TokenRefresher;
|
||||
use crate::fleet::{
|
||||
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, TokenRefreshService},
|
||||
store::PgAccountStore,
|
||||
};
|
||||
use crate::upstream::openai::token_client::TokenRefresher;
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
const DEFAULT_INTERVAL_SECS: u64 = 60;
|
||||
@@ -54,8 +54,8 @@ where
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.interval_secs,
|
||||
"token 刷新任务已启动",
|
||||
"token 刷新任务已关闭",
|
||||
"Token 刷新任务已启动",
|
||||
"Token 刷新任务已关闭",
|
||||
);
|
||||
spawn_periodic_task(self, config)
|
||||
}
|
||||
@@ -73,7 +73,7 @@ where
|
||||
scheduled = summary.scheduled,
|
||||
immediate = summary.immediate,
|
||||
replaced = summary.replaced,
|
||||
"token 刷新定时器已调度"
|
||||
"Token 刷新定时器已调度"
|
||||
);
|
||||
}
|
||||
Ok(summary) => {
|
||||
@@ -84,7 +84,7 @@ where
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error, "token 刷新定时器调度失败");
|
||||
warn!(error = %error, "Token 刷新定时器调度失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -1,10 +1,59 @@
|
||||
//! 客户端身份、本地会话索引与账号作用域 wire 身份隔离。
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
dispatch::controllers::history::{CrossAccountReplay, sanitize_cross_account_input},
|
||||
infra::identity::AccountPseudonymizer,
|
||||
upstream::openai::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
|
||||
upstream::openai::protocol::responses::CodexResponsesRequest,
|
||||
};
|
||||
|
||||
const TURN_STATE_METADATA_KEY: &str = "x-codex-turn-state";
|
||||
const TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
|
||||
const INSTALLATION_IDENTITY_KEYS: [&str; 3] = [
|
||||
"installation_id",
|
||||
"installationId",
|
||||
"x-codex-installation-id",
|
||||
];
|
||||
const CROSS_ACCOUNT_IDENTITY_KEYS: [&str; 26] = [
|
||||
"authorization",
|
||||
"Authorization",
|
||||
"cookie",
|
||||
"Cookie",
|
||||
"chatgpt-account-id",
|
||||
"chatgpt_account_id",
|
||||
"chatgptAccountId",
|
||||
"account_id",
|
||||
"accountId",
|
||||
"user_id",
|
||||
"userId",
|
||||
"chatgpt_user_id",
|
||||
"chatgptUserId",
|
||||
"access_token",
|
||||
"accessToken",
|
||||
"session_token",
|
||||
"sessionToken",
|
||||
"refresh_token",
|
||||
"refreshToken",
|
||||
"id_token",
|
||||
"idToken",
|
||||
"token",
|
||||
"cookies",
|
||||
"cookie_header",
|
||||
"cookieHeader",
|
||||
"cf_clearance",
|
||||
];
|
||||
const CROSS_ACCOUNT_OPAQUE_STATE_KEYS: [&str; 8] = [
|
||||
"turnState",
|
||||
"turn_state",
|
||||
"x-codex-turn-state",
|
||||
"previous_response_id",
|
||||
"previousResponseId",
|
||||
"response_id",
|
||||
"responseId",
|
||||
"conversation",
|
||||
];
|
||||
|
||||
/// 持久化密钥驱动的账号身份隔离服务。
|
||||
#[derive(Clone)]
|
||||
pub struct AccountIdentityService {
|
||||
@@ -21,25 +70,48 @@ impl std::fmt::Debug for AccountIdentityService {
|
||||
|
||||
/// 单次账号 attempt 的 wire 身份。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountScopedIdentity {
|
||||
pub prompt_cache_key: Option<String>,
|
||||
pub session_id: Option<String>,
|
||||
pub thread_id: Option<String>,
|
||||
pub client_request_id: String,
|
||||
pub turn_id: Option<String>,
|
||||
pub window_id: Option<String>,
|
||||
pub parent_thread_id: Option<String>,
|
||||
pub(in crate::dispatch) struct AccountScopedIdentity {
|
||||
pub installation_id: String,
|
||||
}
|
||||
|
||||
struct IdentityValues<'a> {
|
||||
prompt_cache_key: Option<&'a str>,
|
||||
session_id: Option<&'a str>,
|
||||
thread_id: Option<&'a str>,
|
||||
client_request_id: Option<&'a str>,
|
||||
turn_id: Option<&'a str>,
|
||||
window_id: Option<&'a str>,
|
||||
parent_thread_id: Option<&'a str>,
|
||||
/// 单个请求内账号绑定 opaque 状态的 owner。
|
||||
pub(in crate::dispatch) struct AccountIdentityScope {
|
||||
account_bound_state_owner_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AccountIdentityScope {
|
||||
pub(in crate::dispatch) fn new(preferred_account_id: Option<String>) -> Self {
|
||||
Self {
|
||||
account_bound_state_owner_id: preferred_account_id,
|
||||
}
|
||||
}
|
||||
|
||||
fn bind_and_is_cross_account(&mut self, account_id: &str) -> bool {
|
||||
let owner = self
|
||||
.account_bound_state_owner_id
|
||||
.get_or_insert_with(|| account_id.to_string());
|
||||
owner != account_id
|
||||
}
|
||||
}
|
||||
|
||||
/// 已按当前账号重建 wire 身份、可安全交给上游 transport 的请求。
|
||||
pub(in crate::dispatch) struct AccountScopedRequest {
|
||||
request: CodexResponsesRequest,
|
||||
identity: AccountScopedIdentity,
|
||||
}
|
||||
|
||||
impl AccountScopedRequest {
|
||||
pub(in crate::dispatch) fn request(&self) -> &CodexResponsesRequest {
|
||||
&self.request
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn identity(&self) -> &AccountScopedIdentity {
|
||||
&self.identity
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn into_request(self) -> CodexResponsesRequest {
|
||||
self.request
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountIdentityService {
|
||||
@@ -63,109 +135,218 @@ impl AccountIdentityService {
|
||||
.map(|anchor| format!("lc_{}", self.pseudonym("local-conversation", None, anchor)));
|
||||
}
|
||||
|
||||
pub fn scope(
|
||||
/// 从请求语义为单次账号 attempt 重建完整 wire 身份。
|
||||
pub(in crate::dispatch) fn scope_request(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
scope: &mut AccountIdentityScope,
|
||||
mut request: CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
proxy_request_id: &str,
|
||||
) -> AccountScopedIdentity {
|
||||
self.scope_values(
|
||||
account_id,
|
||||
proxy_request_id,
|
||||
IdentityValues {
|
||||
prompt_cache_key: request.prompt_cache_key(),
|
||||
session_id: request.client_session_id.as_deref(),
|
||||
thread_id: request.client_thread_id.as_deref(),
|
||||
client_request_id: request.client_request_id.as_deref(),
|
||||
turn_id: request.client_turn_id.as_deref(),
|
||||
window_id: request.codex_window_id.as_deref(),
|
||||
parent_thread_id: request.parent_thread_id.as_deref(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub fn scope_compact(
|
||||
&self,
|
||||
request: &CodexCompactRequest,
|
||||
account_id: &str,
|
||||
proxy_request_id: &str,
|
||||
) -> AccountScopedIdentity {
|
||||
self.scope_values(
|
||||
account_id,
|
||||
proxy_request_id,
|
||||
IdentityValues {
|
||||
prompt_cache_key: request.prompt_cache_key(),
|
||||
session_id: request.client_session_id.as_deref(),
|
||||
thread_id: request.client_thread_id.as_deref(),
|
||||
client_request_id: request.client_request_id.as_deref(),
|
||||
turn_id: request.client_turn_id.as_deref(),
|
||||
window_id: request.client_window_id.as_deref(),
|
||||
parent_thread_id: request.client_parent_thread_id.as_deref(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub fn scope_auxiliary(
|
||||
&self,
|
||||
account_id: &str,
|
||||
proxy_request_id: &str,
|
||||
) -> AccountScopedIdentity {
|
||||
AccountScopedIdentity {
|
||||
client_request_id: format!(
|
||||
"cr_{}",
|
||||
self.pseudonym("client-request-id", Some(account_id), proxy_request_id)
|
||||
),
|
||||
installation_id: self.pseudonymizer.installation_id(account_id),
|
||||
..AccountScopedIdentity::default()
|
||||
cross_account_replay: Option<&CrossAccountReplay>,
|
||||
) -> AccountScopedRequest {
|
||||
if let Some(replay) = cross_account_replay {
|
||||
replay.validate_target(account_id);
|
||||
}
|
||||
}
|
||||
|
||||
fn scope_values(
|
||||
&self,
|
||||
account_id: &str,
|
||||
proxy_request_id: &str,
|
||||
values: IdentityValues<'_>,
|
||||
) -> AccountScopedIdentity {
|
||||
AccountScopedIdentity {
|
||||
prompt_cache_key: self.scope_optional(
|
||||
"prompt-cache-key",
|
||||
account_id,
|
||||
values.prompt_cache_key,
|
||||
),
|
||||
session_id: self.scope_optional("session-id", account_id, values.session_id),
|
||||
thread_id: self.scope_optional("thread-id", account_id, values.thread_id),
|
||||
client_request_id: format!(
|
||||
"cr_{}",
|
||||
self.pseudonym(
|
||||
"client-request-id",
|
||||
Some(account_id),
|
||||
values.client_request_id.unwrap_or(proxy_request_id),
|
||||
)
|
||||
),
|
||||
turn_id: self.scope_optional("turn-id", account_id, values.turn_id),
|
||||
window_id: self.scope_optional("window-id", account_id, values.window_id),
|
||||
parent_thread_id: self.scope_optional(
|
||||
"parent-thread-id",
|
||||
account_id,
|
||||
values.parent_thread_id,
|
||||
),
|
||||
let cross_account =
|
||||
cross_account_replay.is_some() || scope.bind_and_is_cross_account(account_id);
|
||||
let identity = AccountScopedIdentity {
|
||||
installation_id: self.pseudonymizer.installation_id(account_id),
|
||||
};
|
||||
let client_metadata_turn_state = client_metadata_string(&request, TURN_STATE_METADATA_KEY);
|
||||
let preserve_turn_state = !cross_account
|
||||
&& (request.turn_state.is_some() || client_metadata_turn_state.is_some());
|
||||
let turn_state = if preserve_turn_state {
|
||||
request.turn_state.clone()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let client_metadata_turn_state = if preserve_turn_state {
|
||||
client_metadata_turn_state
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let turn_metadata = request
|
||||
.turn_metadata
|
||||
.as_deref()
|
||||
.and_then(|raw| Self::scope_turn_metadata(raw, &identity, cross_account));
|
||||
let client_metadata_turn_metadata = client_metadata_string(&request, TURN_METADATA_KEY)
|
||||
.and_then(|raw| Self::scope_turn_metadata(&raw, &identity, cross_account));
|
||||
|
||||
if cross_account {
|
||||
sanitize_cross_account_input(&mut request);
|
||||
remove_cross_account_request_state(&mut request);
|
||||
}
|
||||
|
||||
apply_identity_to_request(
|
||||
&mut request,
|
||||
&identity,
|
||||
turn_state.as_deref(),
|
||||
turn_metadata.as_deref(),
|
||||
client_metadata_turn_state.as_deref(),
|
||||
client_metadata_turn_metadata.as_deref(),
|
||||
cross_account,
|
||||
);
|
||||
request.turn_state = turn_state;
|
||||
request.turn_metadata = turn_metadata;
|
||||
|
||||
AccountScopedRequest { request, identity }
|
||||
}
|
||||
|
||||
fn scope_optional(
|
||||
&self,
|
||||
domain: &str,
|
||||
account_id: &str,
|
||||
value: Option<&str>,
|
||||
fn scope_turn_metadata(
|
||||
raw: &str,
|
||||
identity: &AccountScopedIdentity,
|
||||
cross_account: bool,
|
||||
) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| format!("wi_{}", self.pseudonym(domain, Some(account_id), value)))
|
||||
let Ok(Value::Object(mut metadata)) = serde_json::from_str::<Value>(raw) else {
|
||||
return (!cross_account).then(|| raw.to_string());
|
||||
};
|
||||
|
||||
if cross_account {
|
||||
remove_cross_account_metadata_state(&mut metadata);
|
||||
} else if !INSTALLATION_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.any(|key| metadata.contains_key(key))
|
||||
{
|
||||
return Some(raw.to_string());
|
||||
}
|
||||
|
||||
for key in INSTALLATION_IDENTITY_KEYS {
|
||||
replace_existing_metadata_field(&mut metadata, key, Some(&identity.installation_id));
|
||||
}
|
||||
|
||||
match serde_json::to_string(&metadata) {
|
||||
Ok(metadata) => Some(metadata),
|
||||
Err(_) if cross_account => None,
|
||||
Err(_) => Some(raw.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn pseudonym(&self, domain: &str, account_id: Option<&str>, value: &str) -> String {
|
||||
self.pseudonymizer.scoped(domain, account_id, value)
|
||||
}
|
||||
}
|
||||
|
||||
fn client_metadata_string(request: &CodexResponsesRequest, key: &str) -> Option<String> {
|
||||
request
|
||||
.client_metadata()?
|
||||
.as_object()?
|
||||
.get(key)?
|
||||
.as_str()
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn apply_identity_to_request(
|
||||
request: &mut CodexResponsesRequest,
|
||||
identity: &AccountScopedIdentity,
|
||||
turn_state: Option<&str>,
|
||||
turn_metadata: Option<&str>,
|
||||
client_metadata_turn_state: Option<&str>,
|
||||
client_metadata_turn_metadata: Option<&str>,
|
||||
cross_account: bool,
|
||||
) {
|
||||
for (key, value) in [
|
||||
(
|
||||
"x-codex-installation-id",
|
||||
Some(identity.installation_id.as_str()),
|
||||
),
|
||||
("installation_id", Some(identity.installation_id.as_str())),
|
||||
("installationId", Some(identity.installation_id.as_str())),
|
||||
("turnState", turn_state),
|
||||
("x-codex-turn-state", turn_state),
|
||||
("turnMetadata", turn_metadata),
|
||||
("x-codex-turn-metadata", turn_metadata),
|
||||
] {
|
||||
request.replace_existing_identity_field(key, value);
|
||||
}
|
||||
|
||||
let client_metadata = scoped_client_metadata(
|
||||
request.client_metadata(),
|
||||
identity,
|
||||
client_metadata_turn_state,
|
||||
client_metadata_turn_metadata,
|
||||
cross_account,
|
||||
);
|
||||
request.set_client_metadata(client_metadata);
|
||||
}
|
||||
|
||||
fn scoped_client_metadata(
|
||||
client_metadata: Option<&Value>,
|
||||
identity: &AccountScopedIdentity,
|
||||
turn_state: Option<&str>,
|
||||
turn_metadata: Option<&str>,
|
||||
cross_account: bool,
|
||||
) -> Option<Value> {
|
||||
let client_metadata = client_metadata?;
|
||||
let Value::Object(mut metadata) = client_metadata.clone() else {
|
||||
return (!cross_account).then(|| client_metadata.clone());
|
||||
};
|
||||
|
||||
if cross_account {
|
||||
remove_cross_account_metadata_state(&mut metadata);
|
||||
}
|
||||
|
||||
for (key, value) in [
|
||||
(
|
||||
"x-codex-installation-id",
|
||||
Some(identity.installation_id.as_str()),
|
||||
),
|
||||
(TURN_STATE_METADATA_KEY, turn_state),
|
||||
(TURN_METADATA_KEY, turn_metadata),
|
||||
] {
|
||||
replace_metadata_field(&mut metadata, key, value);
|
||||
}
|
||||
replace_existing_metadata_field(
|
||||
&mut metadata,
|
||||
"installation_id",
|
||||
Some(&identity.installation_id),
|
||||
);
|
||||
replace_existing_metadata_field(
|
||||
&mut metadata,
|
||||
"installationId",
|
||||
Some(&identity.installation_id),
|
||||
);
|
||||
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
fn remove_cross_account_request_state(request: &mut CodexResponsesRequest) {
|
||||
for key in CROSS_ACCOUNT_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
|
||||
{
|
||||
request.replace_existing_identity_field(key, None);
|
||||
}
|
||||
for key in ["turn_metadata", TURN_METADATA_KEY] {
|
||||
request.replace_existing_identity_field(key, None);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_cross_account_metadata_state(metadata: &mut Map<String, Value>) {
|
||||
for key in CROSS_ACCOUNT_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
|
||||
.chain(["turnMetadata", "turn_metadata", TURN_METADATA_KEY])
|
||||
{
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_existing_metadata_field(
|
||||
metadata: &mut Map<String, Value>,
|
||||
key: &str,
|
||||
value: Option<&str>,
|
||||
) {
|
||||
if metadata.contains_key(key) {
|
||||
replace_metadata_field(metadata, key, value);
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_metadata_field(metadata: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
match value.filter(|value| !value.trim().is_empty()) {
|
||||
Some(value) => {
|
||||
metadata.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
None => {
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,8 @@ pub mod service;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
|
||||
pub use identity::{AccountIdentityService, AccountScopedIdentity};
|
||||
pub use identity::AccountIdentityService;
|
||||
pub(in crate::dispatch) use identity::{AccountIdentityScope, AccountScopedRequest};
|
||||
pub use resolve::*;
|
||||
pub use service::*;
|
||||
pub use store::*;
|
||||
|
||||
@@ -1,108 +1,7 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::{SessionAffinityEntry, SessionAffinityService},
|
||||
recovery::history::HistoryRecoveryPlan,
|
||||
service::ResponseDispatchService,
|
||||
},
|
||||
upstream::openai::protocol::{
|
||||
events::TokenUsage,
|
||||
responses::{completed_response_metadata, CodexResponsesRequest},
|
||||
},
|
||||
};
|
||||
|
||||
impl ResponseDispatchService {
|
||||
pub(in crate::dispatch) async fn preferred_account_id_for_request(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
history: &HistoryRecoveryPlan,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
if let Some(account_id) = history.preferred_account_id() {
|
||||
return Some(account_id.to_string());
|
||||
}
|
||||
if request.previous_response_id().is_some() {
|
||||
return None;
|
||||
}
|
||||
self.conversation_account_id(request, now).await
|
||||
}
|
||||
|
||||
async fn conversation_account_id(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
let conversation_id = request
|
||||
.local_conversation_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let variant_hash = compute_variant_hash(request);
|
||||
self.session_affinity
|
||||
.lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn record_response_affinity(
|
||||
session_affinity: &SessionAffinityService,
|
||||
history: &HistoryRecoveryPlan,
|
||||
original_request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
body: &str,
|
||||
turn_state: Option<String>,
|
||||
usage: Option<TokenUsage>,
|
||||
) {
|
||||
let metadata = match completed_response_metadata(body) {
|
||||
Ok(Some(metadata)) => metadata,
|
||||
Ok(None) => return,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
"failed to parse completed response metadata for session affinity"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let conversation_id = history
|
||||
.conversation_id(original_request)
|
||||
.unwrap_or(&metadata.response_id)
|
||||
.to_string();
|
||||
let variant_hash = compute_variant_hash(original_request);
|
||||
let entry = SessionAffinityEntry {
|
||||
account_id: account_id.to_string(),
|
||||
conversation_id,
|
||||
turn_state: turn_state
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| original_request.turn_state.clone()),
|
||||
instructions_hash: Some(hash_instructions(Some(original_request.instructions()))),
|
||||
input_tokens: usage.map(|usage| usage.input_tokens),
|
||||
function_call_ids: metadata.function_call_ids,
|
||||
variant_hash: Some(variant_hash),
|
||||
continuation_scope: if original_request.store() {
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::Persisted
|
||||
} else {
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal
|
||||
},
|
||||
replay: history.completed_replay(original_request, &metadata.output),
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
if let Err(error) = session_affinity
|
||||
.record(metadata.response_id.clone(), entry)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
response_id = %metadata.response_id,
|
||||
account_id = %account_id,
|
||||
"failed to record session affinity"
|
||||
);
|
||||
}
|
||||
}
|
||||
use crate::upstream::openai::protocol::responses::CodexResponsesRequest;
|
||||
|
||||
// ====================================================================
|
||||
// Variant identity helpers
|
||||
@@ -146,12 +45,11 @@ fn build_variant_identity(request: &CodexResponsesRequest) -> Option<String> {
|
||||
if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) {
|
||||
parts.push(format!("window:{window_id}"));
|
||||
}
|
||||
if request.explicit_prompt_cache_key
|
||||
|| non_empty_str(request.client_conversation_id.as_deref()).is_some()
|
||||
if (request.explicit_prompt_cache_key
|
||||
|| non_empty_str(request.client_conversation_id.as_deref()).is_some())
|
||||
&& let Some(anchor) = derive_stable_conversation_key(request)
|
||||
{
|
||||
if let Some(anchor) = derive_stable_conversation_key(request) {
|
||||
parts.push(format!("anchor:{anchor}"));
|
||||
}
|
||||
parts.push(format!("anchor:{anchor}"));
|
||||
}
|
||||
|
||||
(!parts.is_empty()).then(|| parts.join("\0"))
|
||||
|
||||
@@ -3,15 +3,14 @@
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::redis::RedisConnection;
|
||||
|
||||
use super::{
|
||||
store::{RedisSessionAffinityStore, RedisSessionAffinityStoreError},
|
||||
types::SessionAffinityEntry,
|
||||
};
|
||||
|
||||
const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60;
|
||||
pub const MAX_REPLAY_SNAPSHOT_BYTES: u64 = 2 * 1024 * 1024;
|
||||
pub const MAX_REPLAY_SESSION_BYTES: u64 = 16 * 1024 * 1024;
|
||||
pub const MAX_REPLAY_DEPTH: u16 = 128;
|
||||
|
||||
/// 运行时会话亲和性服务。
|
||||
#[derive(Clone)]
|
||||
@@ -49,36 +48,6 @@ impl SessionAffinityService {
|
||||
self.entry(response_id, now).await
|
||||
}
|
||||
|
||||
pub async fn replay_input(
|
||||
&self,
|
||||
response_id: &str,
|
||||
entry: &SessionAffinityEntry,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<Vec<serde_json::Value>> {
|
||||
match self
|
||||
.store
|
||||
.replay_input(
|
||||
response_id,
|
||||
entry,
|
||||
now,
|
||||
self.ttl,
|
||||
MAX_REPLAY_DEPTH,
|
||||
MAX_REPLAY_SESSION_BYTES,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(input) => input,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
error = %error,
|
||||
"failed to reconstruct Redis response replay history"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn lookup_latest_account_by_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
@@ -98,7 +67,7 @@ impl SessionAffinityService {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
error = %error,
|
||||
"failed to remove Redis session affinity"
|
||||
"Failed to remove Redis session affinity"
|
||||
);
|
||||
false
|
||||
}
|
||||
@@ -109,17 +78,21 @@ impl SessionAffinityService {
|
||||
match self.store.forget_account(account_id).await {
|
||||
Ok(forgotten) => forgotten > 0,
|
||||
Err(error) => {
|
||||
tracing::warn!(account_id, error = %error, "failed to remove account affinities");
|
||||
tracing::warn!(account_id, error = %error, "Failed to remove account affinities");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn redis_connection(&self) -> RedisConnection {
|
||||
self.store.redis_connection()
|
||||
}
|
||||
|
||||
async fn entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<SessionAffinityEntry> {
|
||||
match self.store.get(response_id, now, self.ttl).await {
|
||||
Ok(entry) => entry,
|
||||
Err(error) => {
|
||||
tracing::warn!(response_id, error = %error, "failed to read Redis session affinity");
|
||||
tracing::warn!(response_id, error = %error, "Failed to read Redis session affinity");
|
||||
None
|
||||
}
|
||||
}
|
||||
@@ -139,7 +112,7 @@ impl SessionAffinityService {
|
||||
{
|
||||
Ok(entry) => entry,
|
||||
Err(error) => {
|
||||
tracing::warn!(conversation_id, error = %error, "failed to query Redis affinity index");
|
||||
tracing::warn!(conversation_id, error = %error, "Failed to query Redis affinity index");
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,277 @@
|
||||
//! Redis 会话亲和存储。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use redis::AsyncCommands;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::redis::RedisConnection;
|
||||
|
||||
use super::types::SessionAffinityEntry;
|
||||
|
||||
/// 单会话保留的最近响应元数据数量。
|
||||
pub const MAX_CONVERSATION_AFFINITIES: usize = 128;
|
||||
/// 单账号保留的最近响应索引数量。
|
||||
pub const MAX_ACCOUNT_AFFINITIES: usize = 16_384;
|
||||
/// Redis 全局保留的最近响应元数据数量。
|
||||
pub const MAX_GLOBAL_AFFINITIES: usize = 65_536;
|
||||
/// Redis 全局响应元数据 JSON 字节预算。
|
||||
pub const MAX_GLOBAL_AFFINITY_BYTES: u64 = 128 * 1024 * 1024;
|
||||
|
||||
const MAX_AFFINITY_METADATA_BYTES: usize = 64 * 1024;
|
||||
const GLOBAL_STALE_CLEANUP_BATCH: usize = 256;
|
||||
const AFFINITY_VERSION: &str = "affinity:v3";
|
||||
|
||||
const UPSERT_SCRIPT: &str = r#"
|
||||
local response_id = ARGV[1]
|
||||
local metadata_json = ARGV[2]
|
||||
local score = tonumber(ARGV[3])
|
||||
local expires_at = tonumber(ARGV[4])
|
||||
local ttl_seconds = tonumber(ARGV[5])
|
||||
local cutoff_ms = tonumber(ARGV[6])
|
||||
local response_prefix = ARGV[7]
|
||||
local conversation_prefix = ARGV[8]
|
||||
local account_prefix = ARGV[9]
|
||||
local max_conversation = tonumber(ARGV[10])
|
||||
local max_account = tonumber(ARGV[11])
|
||||
local max_global = tonumber(ARGV[12])
|
||||
local max_global_bytes = tonumber(ARGV[13])
|
||||
local stale_cleanup_batch = tonumber(ARGV[14])
|
||||
|
||||
local current = cjson.decode(metadata_json)
|
||||
local conversation_id = current.conversationId
|
||||
local account_id = current.accountId
|
||||
|
||||
local function decrement_bytes(size)
|
||||
if size <= 0 then
|
||||
return
|
||||
end
|
||||
local remaining = redis.call('DECRBY', KEYS[5], size)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[5])
|
||||
end
|
||||
end
|
||||
|
||||
local function read_info(id)
|
||||
local raw = redis.call('HGET', KEYS[6], id)
|
||||
if not raw then
|
||||
return nil
|
||||
end
|
||||
local ok, info = pcall(cjson.decode, raw)
|
||||
if not ok then
|
||||
redis.call('HDEL', KEYS[6], id)
|
||||
return nil
|
||||
end
|
||||
return info
|
||||
end
|
||||
|
||||
local function remove_response(id)
|
||||
local response_key = response_prefix .. id
|
||||
local info = read_info(id)
|
||||
local metadata = nil
|
||||
local raw = redis.call('GET', response_key)
|
||||
if raw then
|
||||
local ok, decoded = pcall(cjson.decode, raw)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
|
||||
local stored_conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local stored_account_id = info and info.accountId or metadata and metadata.accountId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[4], id)
|
||||
redis.call('HDEL', KEYS[6], id)
|
||||
redis.call('ZREM', KEYS[2], id)
|
||||
redis.call('ZREM', KEYS[3], id)
|
||||
decrement_bytes(metadata_bytes)
|
||||
if stored_conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. stored_conversation_id, id)
|
||||
end
|
||||
if stored_account_id then
|
||||
redis.call('ZREM', account_prefix .. stored_account_id, id)
|
||||
end
|
||||
end
|
||||
|
||||
local function trim_index(key, max_members)
|
||||
local excess = redis.call('ZCARD', key) - max_members
|
||||
if excess <= 0 then
|
||||
return
|
||||
end
|
||||
local ids = redis.call('ZRANGE', key, 0, excess - 1)
|
||||
for _, id in ipairs(ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
end
|
||||
|
||||
local previous_info = read_info(response_id)
|
||||
if previous_info then
|
||||
if previous_info.conversationId and previous_info.conversationId ~= conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. previous_info.conversationId, response_id)
|
||||
end
|
||||
if previous_info.accountId and previous_info.accountId ~= account_id then
|
||||
redis.call('ZREM', account_prefix .. previous_info.accountId, response_id)
|
||||
end
|
||||
end
|
||||
local previous_bytes = previous_info and tonumber(previous_info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', KEYS[1])
|
||||
local metadata_bytes = string.len(metadata_json)
|
||||
|
||||
redis.call('SET', KEYS[1], metadata_json, 'EXAT', expires_at)
|
||||
redis.call('ZADD', KEYS[2], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[2], ttl_seconds)
|
||||
redis.call('ZADD', KEYS[3], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[3], ttl_seconds)
|
||||
redis.call('ZADD', KEYS[4], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[4], ttl_seconds)
|
||||
redis.call('INCRBY', KEYS[5], metadata_bytes - previous_bytes)
|
||||
redis.call('EXPIRE', KEYS[5], ttl_seconds)
|
||||
redis.call('HSET', KEYS[6], response_id, cjson.encode({
|
||||
conversationId = conversation_id,
|
||||
accountId = account_id,
|
||||
metadataBytes = metadata_bytes,
|
||||
}))
|
||||
redis.call('EXPIRE', KEYS[6], ttl_seconds)
|
||||
|
||||
local stale_conversation_ids = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', cutoff_ms)
|
||||
for _, id in ipairs(stale_conversation_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
local stale_account_ids = redis.call(
|
||||
'ZRANGEBYSCORE', KEYS[3], '-inf', cutoff_ms, 'LIMIT', 0, stale_cleanup_batch
|
||||
)
|
||||
for _, id in ipairs(stale_account_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
local stale_global_ids = redis.call(
|
||||
'ZRANGEBYSCORE', KEYS[4], '-inf', cutoff_ms, 'LIMIT', 0, stale_cleanup_batch
|
||||
)
|
||||
for _, id in ipairs(stale_global_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
|
||||
trim_index(KEYS[2], max_conversation)
|
||||
trim_index(KEYS[3], max_account)
|
||||
trim_index(KEYS[4], max_global)
|
||||
|
||||
while tonumber(redis.call('GET', KEYS[5]) or '0') > max_global_bytes do
|
||||
local oldest = redis.call('ZRANGE', KEYS[4], 0, 0)
|
||||
if #oldest == 0 then
|
||||
redis.call('DEL', KEYS[5])
|
||||
break
|
||||
end
|
||||
remove_response(oldest[1])
|
||||
end
|
||||
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const FORGET_SCRIPT: &str = r#"
|
||||
local response_id = ARGV[1]
|
||||
local response_prefix = ARGV[2]
|
||||
local conversation_prefix = ARGV[3]
|
||||
local account_prefix = ARGV[4]
|
||||
local response_key = response_prefix .. response_id
|
||||
local info_json = redis.call('HGET', KEYS[3], response_id)
|
||||
local metadata_json = redis.call('GET', response_key)
|
||||
if not info_json and not metadata_json then
|
||||
return 0
|
||||
end
|
||||
|
||||
local info = nil
|
||||
if info_json then
|
||||
local ok, decoded = pcall(cjson.decode, info_json)
|
||||
if ok then
|
||||
info = decoded
|
||||
end
|
||||
end
|
||||
local metadata = nil
|
||||
if metadata_json then
|
||||
local ok, decoded = pcall(cjson.decode, metadata_json)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
local conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local account_id = info and info.accountId or metadata and metadata.accountId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[1], response_id)
|
||||
redis.call('HDEL', KEYS[3], response_id)
|
||||
if metadata_bytes > 0 then
|
||||
local remaining = redis.call('DECRBY', KEYS[2], metadata_bytes)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[2])
|
||||
end
|
||||
end
|
||||
if conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. conversation_id, response_id)
|
||||
end
|
||||
if account_id then
|
||||
redis.call('ZREM', account_prefix .. account_id, response_id)
|
||||
end
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const FORGET_ACCOUNT_SCRIPT: &str = r#"
|
||||
local account_id = ARGV[1]
|
||||
local response_prefix = ARGV[2]
|
||||
local conversation_prefix = ARGV[3]
|
||||
local account_prefix = ARGV[4]
|
||||
local account_key = account_prefix .. account_id
|
||||
local response_ids = redis.call('ZRANGE', account_key, 0, -1)
|
||||
|
||||
local function decrement_bytes(size)
|
||||
if size <= 0 then
|
||||
return
|
||||
end
|
||||
local remaining = redis.call('DECRBY', KEYS[3], size)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[3])
|
||||
end
|
||||
end
|
||||
|
||||
for _, response_id in ipairs(response_ids) do
|
||||
local response_key = response_prefix .. response_id
|
||||
local info_json = redis.call('HGET', KEYS[4], response_id)
|
||||
local metadata_json = redis.call('GET', response_key)
|
||||
local info = nil
|
||||
if info_json then
|
||||
local ok, decoded = pcall(cjson.decode, info_json)
|
||||
if ok then
|
||||
info = decoded
|
||||
end
|
||||
end
|
||||
local metadata = nil
|
||||
if metadata_json then
|
||||
local ok, decoded = pcall(cjson.decode, metadata_json)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
local stored_account_id = info and info.accountId or metadata and metadata.accountId
|
||||
if stored_account_id == account_id then
|
||||
local conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[2], response_id)
|
||||
redis.call('HDEL', KEYS[4], response_id)
|
||||
decrement_bytes(metadata_bytes)
|
||||
if conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. conversation_id, response_id)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
redis.call('DEL', account_key)
|
||||
return #response_ids
|
||||
"#;
|
||||
|
||||
/// Redis 会话亲和性存储。
|
||||
#[derive(Clone)]
|
||||
pub struct RedisSessionAffinityStore {
|
||||
@@ -24,6 +285,8 @@ pub enum RedisSessionAffinityStoreError {
|
||||
Redis(#[from] redis::RedisError),
|
||||
#[error("invalid Redis session affinity value: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("session affinity metadata is too large: {bytes} bytes exceeds {max} bytes")]
|
||||
MetadataTooLarge { bytes: usize, max: usize },
|
||||
}
|
||||
|
||||
pub type RedisSessionAffinityStoreResult<T> = Result<T, RedisSessionAffinityStoreError>;
|
||||
@@ -33,13 +296,16 @@ impl RedisSessionAffinityStore {
|
||||
Self { redis }
|
||||
}
|
||||
|
||||
pub(crate) fn redis_connection(&self) -> RedisConnection {
|
||||
self.redis.clone()
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
response_id: &str,
|
||||
entry: &SessionAffinityEntry,
|
||||
ttl: Duration,
|
||||
) -> RedisSessionAffinityStoreResult<()> {
|
||||
let value = serde_json::to_string(entry)?;
|
||||
let expires_at = entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
@@ -48,50 +314,40 @@ impl RedisSessionAffinityStore {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let response_key = self.response_key(response_id);
|
||||
let conversation_key = self.conversation_key(&entry.conversation_id);
|
||||
let account_key = self.account_key(&entry.account_id);
|
||||
let metadata_json = serde_json::to_string(entry)?;
|
||||
if metadata_json.len() > MAX_AFFINITY_METADATA_BYTES {
|
||||
return Err(RedisSessionAffinityStoreError::MetadataTooLarge {
|
||||
bytes: metadata_json.len(),
|
||||
max: MAX_AFFINITY_METADATA_BYTES,
|
||||
});
|
||||
}
|
||||
|
||||
let ttl_seconds = ttl.num_seconds().max(1);
|
||||
let cutoff_ms = Utc::now().timestamp_millis() - ttl.num_milliseconds();
|
||||
let mut connection = self.redis.manager();
|
||||
let mut transaction = redis::pipe();
|
||||
transaction
|
||||
.atomic()
|
||||
.cmd("SET")
|
||||
.arg(response_key)
|
||||
.arg(value)
|
||||
.arg("EXAT")
|
||||
let _: i64 = redis::Script::new(UPSERT_SCRIPT)
|
||||
.key(self.response_key(response_id))
|
||||
.key(self.conversation_key(&entry.conversation_id))
|
||||
.key(self.account_key(&entry.account_id))
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(response_id)
|
||||
.arg(metadata_json)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(expires_at.timestamp())
|
||||
.ignore()
|
||||
.cmd("ZADD")
|
||||
.arg(&conversation_key)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("EXPIRE")
|
||||
.arg(&conversation_key)
|
||||
.arg(ttl_seconds)
|
||||
.ignore()
|
||||
.cmd("ZADD")
|
||||
.arg(&account_key)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("EXPIRE")
|
||||
.arg(&account_key)
|
||||
.arg(ttl_seconds)
|
||||
.ignore()
|
||||
.cmd("ZREMRANGEBYSCORE")
|
||||
.arg(conversation_key)
|
||||
.arg("-inf")
|
||||
.arg(cutoff_ms)
|
||||
.ignore()
|
||||
.cmd("ZREMRANGEBYSCORE")
|
||||
.arg(&account_key)
|
||||
.arg("-inf")
|
||||
.arg(cutoff_ms)
|
||||
.ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.arg(MAX_CONVERSATION_AFFINITIES)
|
||||
.arg(MAX_ACCOUNT_AFFINITIES)
|
||||
.arg(MAX_GLOBAL_AFFINITIES)
|
||||
.arg(MAX_GLOBAL_AFFINITY_BYTES)
|
||||
.arg(GLOBAL_STALE_CLEANUP_BATCH)
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -128,163 +384,14 @@ impl RedisSessionAffinityStore {
|
||||
) -> RedisSessionAffinityStoreResult<Option<(String, SessionAffinityEntry)>> {
|
||||
let conversation_key = self.conversation_key(conversation_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = connection.zrevrange(&conversation_key, 0, -1).await?;
|
||||
for response_id in response_ids {
|
||||
let Some(entry) = self.get(&response_id, now, ttl).await? else {
|
||||
let _: usize = connection.zrem(&conversation_key, &response_id).await?;
|
||||
continue;
|
||||
};
|
||||
if entry.conversation_id != conversation_id {
|
||||
let _: usize = connection.zrem(&conversation_key, &response_id).await?;
|
||||
continue;
|
||||
}
|
||||
if variant_hash.is_some_and(|expected| entry.variant_hash.as_deref() != Some(expected))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if max_age.is_some_and(|max_age| now.signed_duration_since(entry.created_at) > max_age)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
return Ok(Some((response_id, entry)));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub async fn replay_input(
|
||||
&self,
|
||||
response_id: &str,
|
||||
head: &SessionAffinityEntry,
|
||||
now: DateTime<Utc>,
|
||||
ttl: Duration,
|
||||
max_depth: u16,
|
||||
max_total_bytes: u64,
|
||||
) -> RedisSessionAffinityStoreResult<Option<Vec<Value>>> {
|
||||
let Some(head_snapshot) = head.replay.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if head_snapshot.depth == 0
|
||||
|| head_snapshot.depth > max_depth
|
||||
|| head_snapshot.total_bytes > max_total_bytes
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let conversation_key = self.conversation_key(&head.conversation_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = redis::cmd("ZREVRANGEBYSCORE")
|
||||
let response_ids: Vec<String> = redis::cmd("ZREVRANGE")
|
||||
.arg(&conversation_key)
|
||||
.arg(head.created_at.timestamp_millis())
|
||||
.arg("-inf")
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(max_depth)
|
||||
.arg(MAX_CONVERSATION_AFFINITIES.saturating_sub(1))
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
if !response_ids
|
||||
.iter()
|
||||
.any(|candidate| candidate == response_id)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_keys = response_ids
|
||||
.iter()
|
||||
.map(|candidate| self.response_key(candidate))
|
||||
.collect::<Vec<_>>();
|
||||
let values: Vec<Option<String>> = redis::cmd("MGET")
|
||||
.arg(&response_keys)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
let mut entries = HashMap::with_capacity(values.len());
|
||||
for (candidate, value) in response_ids.into_iter().zip(values) {
|
||||
let Some(value) = value else {
|
||||
continue;
|
||||
};
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
if entry.conversation_id == head.conversation_id
|
||||
&& entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
.is_some_and(|expires_at| expires_at > now)
|
||||
{
|
||||
entries.insert(candidate, entry);
|
||||
}
|
||||
}
|
||||
|
||||
let mut current_id = response_id.to_string();
|
||||
let mut expected_depth = head_snapshot.depth;
|
||||
let mut expected_total_bytes = head_snapshot.total_bytes;
|
||||
let mut nodes = Vec::with_capacity(usize::from(expected_depth));
|
||||
loop {
|
||||
let Some(entry) = entries.get(¤t_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(snapshot) = entry.replay.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.depth != expected_depth || snapshot.total_bytes != expected_total_bytes {
|
||||
return Ok(None);
|
||||
}
|
||||
let node_bytes = replay_node_bytes(snapshot)?;
|
||||
if node_bytes > expected_total_bytes {
|
||||
return Ok(None);
|
||||
}
|
||||
nodes.push((snapshot.turn_input.clone(), snapshot.turn_output.clone()));
|
||||
|
||||
match snapshot.parent_response_id.as_deref() {
|
||||
None if expected_depth == 1 && expected_total_bytes == node_bytes => break,
|
||||
Some(parent_response_id) if expected_depth > 1 => {
|
||||
expected_depth -= 1;
|
||||
expected_total_bytes -= node_bytes;
|
||||
current_id = parent_response_id.to_string();
|
||||
}
|
||||
None | Some(_) => return Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
nodes.reverse();
|
||||
let mut full_input = Vec::new();
|
||||
for (mut input, mut output) in nodes {
|
||||
full_input.append(&mut input);
|
||||
full_input.append(&mut output);
|
||||
}
|
||||
Ok(Some(full_input))
|
||||
}
|
||||
|
||||
pub async fn forget(&self, response_id: &str) -> RedisSessionAffinityStoreResult<bool> {
|
||||
let response_key = self.response_key(response_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let value: Option<String> = connection.get(&response_key).await?;
|
||||
let Some(value) = value else {
|
||||
return Ok(false);
|
||||
};
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
let mut transaction = redis::pipe();
|
||||
transaction
|
||||
.atomic()
|
||||
.cmd("DEL")
|
||||
.arg(response_key)
|
||||
.ignore()
|
||||
.cmd("ZREM")
|
||||
.arg(self.conversation_key(&entry.conversation_id))
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("ZREM")
|
||||
.arg(self.account_key(&entry.account_id))
|
||||
.arg(response_id)
|
||||
.ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub async fn forget_account(&self, account_id: &str) -> RedisSessionAffinityStoreResult<u64> {
|
||||
let account_key = self.account_key(account_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = connection.zrange(&account_key, 0, -1).await?;
|
||||
if response_ids.is_empty() {
|
||||
let _: usize = connection.del(account_key).await?;
|
||||
return Ok(0);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_keys = response_ids
|
||||
@@ -292,67 +399,105 @@ impl RedisSessionAffinityStore {
|
||||
.map(|response_id| self.response_key(response_id))
|
||||
.collect::<Vec<_>>();
|
||||
let values: Vec<Option<String>> = redis::cmd("MGET")
|
||||
.arg(&response_keys)
|
||||
.arg(response_keys)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
let mut conversation_members = Vec::with_capacity(response_ids.len());
|
||||
for (response_id, value) in response_ids.iter().zip(values) {
|
||||
let mut stale_ids = Vec::new();
|
||||
let mut result = None;
|
||||
for (response_id, value) in response_ids.into_iter().zip(values) {
|
||||
let Some(value) = value else {
|
||||
stale_ids.push(response_id);
|
||||
continue;
|
||||
};
|
||||
match serde_json::from_str::<SessionAffinityEntry>(&value) {
|
||||
Ok(entry) if entry.account_id == account_id => {
|
||||
conversation_members.push((response_id, entry.conversation_id));
|
||||
}
|
||||
Ok(_) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
account_id,
|
||||
"account affinity index points to another account"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
account_id,
|
||||
error = %error,
|
||||
"invalid account affinity entry removed during account cleanup"
|
||||
);
|
||||
}
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
if entry.conversation_id != conversation_id
|
||||
|| entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
.is_none_or(|expires_at| expires_at <= now)
|
||||
{
|
||||
stale_ids.push(response_id);
|
||||
continue;
|
||||
}
|
||||
if variant_hash.is_some_and(|expected| entry.variant_hash.as_deref() != Some(expected))
|
||||
|| max_age
|
||||
.is_some_and(|max_age| now.signed_duration_since(entry.created_at) > max_age)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
result = Some((response_id, entry));
|
||||
break;
|
||||
}
|
||||
if !stale_ids.is_empty() {
|
||||
let _: usize = connection.zrem(conversation_key, stale_ids).await?;
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
let mut transaction = redis::pipe();
|
||||
transaction.atomic().cmd("DEL").arg(&response_keys).ignore();
|
||||
for (response_id, conversation_id) in &conversation_members {
|
||||
transaction
|
||||
.cmd("ZREM")
|
||||
.arg(self.conversation_key(conversation_id))
|
||||
.arg(response_id)
|
||||
.ignore();
|
||||
}
|
||||
transaction.cmd("DEL").arg(account_key).ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
Ok(response_ids.len() as u64)
|
||||
pub async fn forget(&self, response_id: &str) -> RedisSessionAffinityStoreResult<bool> {
|
||||
let mut connection = self.redis.manager();
|
||||
let forgotten: i64 = redis::Script::new(FORGET_SCRIPT)
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(response_id)
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(forgotten > 0)
|
||||
}
|
||||
|
||||
pub async fn forget_account(&self, account_id: &str) -> RedisSessionAffinityStoreResult<u64> {
|
||||
let mut connection = self.redis.manager();
|
||||
let forgotten: u64 = redis::Script::new(FORGET_ACCOUNT_SCRIPT)
|
||||
.key(self.account_key(account_id))
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(account_id)
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(forgotten)
|
||||
}
|
||||
|
||||
fn response_key(&self, response_id: &str) -> String {
|
||||
self.redis.key(&format!("affinity:v2:resp:{response_id}"))
|
||||
format!("{}{response_id}", self.response_key_prefix())
|
||||
}
|
||||
|
||||
fn conversation_key(&self, conversation_id: &str) -> String {
|
||||
self.redis
|
||||
.key(&format!("affinity:v2:conv:{conversation_id}"))
|
||||
format!("{}{conversation_id}", self.conversation_key_prefix())
|
||||
}
|
||||
|
||||
fn account_key(&self, account_id: &str) -> String {
|
||||
self.redis.key(&format!("affinity:v2:account:{account_id}"))
|
||||
format!("{}{account_id}", self.account_key_prefix())
|
||||
}
|
||||
|
||||
fn global_index_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:index"))
|
||||
}
|
||||
|
||||
fn global_bytes_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:bytes"))
|
||||
}
|
||||
|
||||
fn global_info_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:info"))
|
||||
}
|
||||
|
||||
fn response_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:resp:"))
|
||||
}
|
||||
|
||||
fn conversation_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:conv:"))
|
||||
}
|
||||
|
||||
fn account_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:account:"))
|
||||
}
|
||||
}
|
||||
|
||||
fn replay_node_bytes(
|
||||
snapshot: &super::types::ResponseReplaySnapshot,
|
||||
) -> Result<u64, serde_json::Error> {
|
||||
serde_json::to_vec(&(&snapshot.turn_input, &snapshot.turn_output))
|
||||
.map(|value| value.len() as u64)
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::upstream::openai::protocol::responses::PreviousResponseScope;
|
||||
|
||||
@@ -18,22 +17,5 @@ pub struct SessionAffinityEntry {
|
||||
pub function_call_ids: Vec<String>,
|
||||
pub variant_hash: Option<String>,
|
||||
pub continuation_scope: PreviousResponseScope,
|
||||
pub replay: Option<ResponseReplaySnapshot>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 单轮 completed response 的有界、无凭据重放增量。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ResponseReplaySnapshot {
|
||||
/// 上一轮由本代理记录的 completed response。
|
||||
pub parent_response_id: Option<String>,
|
||||
/// 本轮客户端新增的输入,已剥离账号绑定字段。
|
||||
pub turn_input: Vec<Value>,
|
||||
/// 本轮上游输出,已剥离账号绑定字段和上游对象 ID。
|
||||
pub turn_output: Vec<Value>,
|
||||
/// 从会话起点到本节点的增量节点数。
|
||||
pub depth: u16,
|
||||
/// 从会话起点到本节点的 JSON 编码累计字节数。
|
||||
pub total_bytes: u64,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
//! 账号级失败分类、隔离副作用与换号决策 owner。
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
controllers::ControllerFailureFact,
|
||||
errors::ClientFailure,
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
|
||||
},
|
||||
},
|
||||
fleet::{
|
||||
account_failure::{
|
||||
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
|
||||
classify_account_failure,
|
||||
},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::{UpstreamFailureFacts, UpstreamFailureKind},
|
||||
protocol::responses::ResponsesSseFailure,
|
||||
transport::CodexBackendClient,
|
||||
},
|
||||
};
|
||||
|
||||
pub(super) struct AccountFailureController;
|
||||
|
||||
pub(super) struct ClassifiedFailure {
|
||||
exhaustion: AccountExhaustionRecord,
|
||||
effect: Option<AccountStateEffect>,
|
||||
}
|
||||
|
||||
impl AccountFailureController {
|
||||
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedFailure> {
|
||||
let account_id = observation.account.as_ref()?.id.as_str();
|
||||
classify(observation, account_id)
|
||||
}
|
||||
|
||||
pub(super) fn classify_failure(
|
||||
account_id: &str,
|
||||
failure: ControllerFailureFact<'_>,
|
||||
) -> Option<ClassifiedFailure> {
|
||||
classify_failure(account_id, failure)
|
||||
}
|
||||
|
||||
pub(super) async fn apply_effect(
|
||||
account_pool: &AccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
classified: &ClassifiedFailure,
|
||||
) {
|
||||
let Some(account_id) = classified.exhaustion.account_id.as_deref() else {
|
||||
return;
|
||||
};
|
||||
if let Some(effect) = &classified.effect {
|
||||
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decision(
|
||||
observation: &AttemptObservation,
|
||||
classified: ClassifiedFailure,
|
||||
) -> AttemptDecision {
|
||||
if observation.routing.can_retry_next_candidate {
|
||||
AttemptDecision::RetryNextCandidate {
|
||||
exhaustion: Some(classified.exhaustion),
|
||||
on_exhaustion: None,
|
||||
}
|
||||
} else {
|
||||
AttemptDecision::Return(
|
||||
crate::dispatch::lifecycle::contract::AttemptReturnKind::Observed,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn is_retryable_transport(observation: &AttemptObservation) -> bool {
|
||||
let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else {
|
||||
return false;
|
||||
};
|
||||
matches!(
|
||||
facts.kind,
|
||||
UpstreamFailureKind::HttpConnect
|
||||
| UpstreamFailureKind::HttpTimeout
|
||||
| UpstreamFailureKind::HttpTransport
|
||||
| UpstreamFailureKind::StreamIdle
|
||||
| UpstreamFailureKind::WebSocketTransport
|
||||
| UpstreamFailureKind::WebSocketTimeout
|
||||
) || facts
|
||||
.status_code
|
||||
.is_some_and(|status| (500..=599).contains(&status))
|
||||
}
|
||||
|
||||
pub(super) fn client_failure(
|
||||
classified: &ClassifiedFailure,
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> ClientFailure {
|
||||
let status_code =
|
||||
classified
|
||||
.exhaustion
|
||||
.status_code
|
||||
.unwrap_or(match classified.exhaustion.kind {
|
||||
ExhaustedAccountKind::ModelUnsupported => 400,
|
||||
ExhaustedAccountKind::Expired | ExhaustedAccountKind::Disabled => 401,
|
||||
ExhaustedAccountKind::Banned => 403,
|
||||
ExhaustedAccountKind::QuotaExhausted
|
||||
| ExhaustedAccountKind::RateLimited
|
||||
| ExhaustedAccountKind::CloudflareChallenge
|
||||
| ExhaustedAccountKind::CloudflarePathBlocked
|
||||
| ExhaustedAccountKind::UpstreamUnavailable => 502,
|
||||
});
|
||||
ClientFailure::new(failure.clone(), status_code, false)
|
||||
}
|
||||
|
||||
pub(super) fn unowned_client_failure(failure: &ResponsesSseFailure) -> Option<ClientFailure> {
|
||||
let code = failure
|
||||
.upstream_code
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
if code.contains("invalid_request")
|
||||
|| code.contains("not_found")
|
||||
|| code.contains("context_window")
|
||||
|| code == "context_length_exceeded"
|
||||
|| code.contains("invalid_prompt")
|
||||
|| code.contains("bad_request")
|
||||
{
|
||||
return Some(ClientFailure::new(failure.clone(), 400, false));
|
||||
}
|
||||
if matches!(code.as_str(), "server_is_overloaded" | "slow_down")
|
||||
|| code.contains("server_overloaded")
|
||||
{
|
||||
return Some(ClientFailure::new(failure.clone(), 503, true));
|
||||
}
|
||||
if code == "usage_not_included" {
|
||||
return Some(ClientFailure::new(failure.clone(), 403, false));
|
||||
}
|
||||
classify_account_failure(&super::failure_observation::from_response(failure)).map(
|
||||
|classified| {
|
||||
let status = match classified.kind {
|
||||
AccountFailureKind::ModelUnsupported => 400,
|
||||
AccountFailureKind::Expired => 401,
|
||||
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
|
||||
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
|
||||
};
|
||||
ClientFailure::new(failure.clone(), status, false)
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn classify(observation: &AttemptObservation, account_id: &str) -> Option<ClassifiedFailure> {
|
||||
match &observation.kind {
|
||||
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty) => {
|
||||
Some(ClassifiedFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::UpstreamUnavailable,
|
||||
"upstream response did not include visible output",
|
||||
),
|
||||
effect: None,
|
||||
})
|
||||
}
|
||||
_ => ControllerFailureFact::from_attempt(observation)
|
||||
.and_then(|failure| classify_failure(account_id, failure)),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_failure(
|
||||
account_id: &str,
|
||||
failure: ControllerFailureFact<'_>,
|
||||
) -> Option<ClassifiedFailure> {
|
||||
match failure {
|
||||
ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts),
|
||||
ControllerFailureFact::Response(failure) => classify_sse(account_id, failure),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<ClassifiedFailure> {
|
||||
classified_failure(
|
||||
account_id,
|
||||
&facts.body,
|
||||
facts.status_code,
|
||||
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<ClassifiedFailure> {
|
||||
classified_failure(
|
||||
account_id,
|
||||
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
|
||||
failure.explicit_status_code,
|
||||
classify_account_failure(&super::failure_observation::from_response(failure))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classified_failure(
|
||||
account_id: &str,
|
||||
message: &str,
|
||||
status_code: Option<u16>,
|
||||
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
|
||||
) -> Option<ClassifiedFailure> {
|
||||
let kind = match classified.kind {
|
||||
AccountFailureKind::ModelUnsupported => ExhaustedAccountKind::ModelUnsupported,
|
||||
AccountFailureKind::Expired => ExhaustedAccountKind::Expired,
|
||||
AccountFailureKind::Disabled => ExhaustedAccountKind::Disabled,
|
||||
AccountFailureKind::Banned => ExhaustedAccountKind::Banned,
|
||||
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => return None,
|
||||
};
|
||||
let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message);
|
||||
if let Some(status_code) = status_code {
|
||||
exhaustion = exhaustion.with_status_code(status_code);
|
||||
}
|
||||
Some(ClassifiedFailure {
|
||||
exhaustion,
|
||||
effect: classified.effect,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
//! 会话亲和性路由与写入的唯一 feature owner。
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::dispatch::affinity::{
|
||||
SessionAffinityEntry, SessionAffinityService, compute_variant_hash, hash_instructions,
|
||||
};
|
||||
use crate::upstream::openai::protocol::{
|
||||
events::TokenUsage,
|
||||
responses::{CodexResponsesRequest, PreviousResponseScope, completed_response_metadata},
|
||||
};
|
||||
|
||||
const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
|
||||
pub(super) struct AffinityController;
|
||||
|
||||
pub(super) struct ResponseExit<'a> {
|
||||
pub affinity: &'a SessionAffinityService,
|
||||
pub conversation_id: Option<&'a str>,
|
||||
pub request: &'a CodexResponsesRequest,
|
||||
pub account_id: &'a str,
|
||||
pub body: &'a str,
|
||||
pub turn_state: Option<String>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
pub continuation_scope: PreviousResponseScope,
|
||||
}
|
||||
|
||||
pub(super) struct StreamExit<'a> {
|
||||
pub response: ResponseExit<'a>,
|
||||
pub completed: bool,
|
||||
}
|
||||
|
||||
impl AffinityController {
|
||||
pub(super) async fn preferred_account_id(
|
||||
affinity: &SessionAffinityService,
|
||||
request: &CodexResponsesRequest,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
if request.previous_response_id().is_some() {
|
||||
return None;
|
||||
}
|
||||
let conversation_id = request
|
||||
.local_conversation_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let variant_hash = compute_variant_hash(request);
|
||||
match timeout(
|
||||
REDIS_BEST_EFFORT_TIMEOUT,
|
||||
affinity.lookup_latest_account_by_conversation(
|
||||
conversation_id,
|
||||
None,
|
||||
Some(&variant_hash),
|
||||
now,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(account_id) => account_id,
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
conversation_id,
|
||||
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
|
||||
"Timed out reading conversation affinity; continuing without preference"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn leave_complete(exit: ResponseExit<'_>) {
|
||||
record_response_affinity(exit).await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
|
||||
if !exit.completed {
|
||||
return;
|
||||
}
|
||||
record_response_affinity(exit.response).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_response_affinity(exit: ResponseExit<'_>) {
|
||||
let ResponseExit {
|
||||
affinity,
|
||||
conversation_id,
|
||||
request,
|
||||
account_id,
|
||||
body,
|
||||
turn_state,
|
||||
usage,
|
||||
continuation_scope,
|
||||
} = exit;
|
||||
let metadata = match completed_response_metadata(body) {
|
||||
Ok(Some(metadata)) => metadata,
|
||||
Ok(None) => return,
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "Failed to parse completed response metadata for session affinity");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let entry = SessionAffinityEntry {
|
||||
account_id: account_id.to_string(),
|
||||
conversation_id: conversation_id.unwrap_or(&metadata.response_id).to_string(),
|
||||
turn_state: turn_state
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| request.turn_state.clone()),
|
||||
instructions_hash: Some(hash_instructions(Some(request.instructions()))),
|
||||
input_tokens: usage.map(|usage| usage.input_tokens),
|
||||
function_call_ids: metadata.function_call_ids,
|
||||
variant_hash: Some(compute_variant_hash(request)),
|
||||
continuation_scope,
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
match timeout(
|
||||
REDIS_BEST_EFFORT_TIMEOUT,
|
||||
affinity.record(metadata.response_id.clone(), entry),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(error = %error, response_id = %metadata.response_id, account_id, "Failed to record session affinity");
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
response_id = %metadata.response_id,
|
||||
account_id,
|
||||
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
|
||||
"Timed out recording session affinity"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
//! Cloudflare challenge、路径封禁与 Cookie 生命周期的唯一 feature owner。
|
||||
|
||||
use chrono::Utc;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
|
||||
},
|
||||
lifecycle::stream::StreamTerminal,
|
||||
},
|
||||
fleet::{
|
||||
account::AccountStatus,
|
||||
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
upstream::openai::failure::UpstreamFailureFacts,
|
||||
};
|
||||
|
||||
pub(super) struct CloudflareController;
|
||||
|
||||
pub(super) struct StreamExit<'a> {
|
||||
pub recovery: &'a CloudflareRecovery,
|
||||
pub account_id: &'a str,
|
||||
pub set_cookie_headers: &'a [String],
|
||||
pub terminal: &'a StreamTerminal,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum CloudflareFailureKind {
|
||||
Challenge,
|
||||
PathBlocked,
|
||||
}
|
||||
|
||||
pub(super) struct CloudflareFailure {
|
||||
kind: CloudflareFailureKind,
|
||||
exhaustion: AccountExhaustionRecord,
|
||||
}
|
||||
|
||||
impl CloudflareController {
|
||||
pub(super) fn classify(observation: &AttemptObservation) -> Option<CloudflareFailure> {
|
||||
classify(observation)
|
||||
}
|
||||
|
||||
pub(super) async fn apply_effect(
|
||||
recovery: &CloudflareRecovery,
|
||||
account_pool: &AccountPoolService,
|
||||
failure: &CloudflareFailure,
|
||||
) {
|
||||
match failure.kind {
|
||||
CloudflareFailureKind::Challenge => {
|
||||
let Some(account_id) = failure.exhaustion.account_id.as_deref() else {
|
||||
return;
|
||||
};
|
||||
recovery.apply_challenge(account_pool, account_id).await;
|
||||
}
|
||||
CloudflareFailureKind::PathBlocked => {
|
||||
let Some(account_id) = failure.exhaustion.account_id.as_deref() else {
|
||||
return;
|
||||
};
|
||||
recovery.apply_path_block(account_pool, account_id).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decision(
|
||||
observation: &AttemptObservation,
|
||||
failure: CloudflareFailure,
|
||||
) -> AttemptDecision {
|
||||
if observation.routing.can_retry_next_candidate {
|
||||
return AttemptDecision::RetryNextCandidate {
|
||||
exhaustion: Some(failure.exhaustion),
|
||||
on_exhaustion: None,
|
||||
};
|
||||
}
|
||||
AttemptDecision::Return(AttemptReturnKind::Observed)
|
||||
}
|
||||
|
||||
pub(super) async fn leave_complete(recovery: &CloudflareRecovery, account_id: &str) {
|
||||
recovery.reset_account_recovery(account_id).await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
|
||||
exit.recovery
|
||||
.capture_set_cookie_headers(exit.account_id, exit.set_cookie_headers)
|
||||
.await;
|
||||
if matches!(
|
||||
exit.terminal,
|
||||
StreamTerminal::Completed { .. } | StreamTerminal::Incomplete { .. }
|
||||
) {
|
||||
exit.recovery.reset_account_recovery(exit.account_id).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn classify(observation: &AttemptObservation) -> Option<CloudflareFailure> {
|
||||
let account_id = observation.account.as_ref()?.id.as_str();
|
||||
let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else {
|
||||
return None;
|
||||
};
|
||||
classify_upstream(account_id, facts)
|
||||
}
|
||||
|
||||
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<CloudflareFailure> {
|
||||
if facts.status_code == Some(403) && is_challenge(&facts.body) {
|
||||
return Some(CloudflareFailure {
|
||||
kind: CloudflareFailureKind::Challenge,
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::CloudflareChallenge,
|
||||
"Upstream blocked the request (Cloudflare challenge)",
|
||||
),
|
||||
});
|
||||
}
|
||||
if facts.status_code == Some(404) && facts.body.trim().is_empty() {
|
||||
return Some(CloudflareFailure {
|
||||
kind: CloudflareFailureKind::PathBlocked,
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::CloudflarePathBlocked,
|
||||
"Upstream blocked the request (Cloudflare path-block)",
|
||||
),
|
||||
});
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn is_challenge(value: &str) -> bool {
|
||||
let value = value.to_ascii_lowercase();
|
||||
value.contains("cf-mitigated")
|
||||
|| value.contains("cf-chl-bypass")
|
||||
|| value.contains("_cf_chl")
|
||||
|| value.contains("cf_chl")
|
||||
|| value.contains("attention required")
|
||||
|| value.contains("just a moment")
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CloudflareRecovery {
|
||||
path_block_tracker: CloudflarePathBlockTracker,
|
||||
challenge_tracker: CloudflareChallengeCooldownTracker,
|
||||
cookie_store: PgCookieStore,
|
||||
}
|
||||
|
||||
impl CloudflareRecovery {
|
||||
pub fn new(cookie_store: PgCookieStore) -> Self {
|
||||
Self {
|
||||
path_block_tracker: CloudflarePathBlockTracker::default(),
|
||||
challenge_tracker: CloudflareChallengeCooldownTracker::default(),
|
||||
cookie_store,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option<String> {
|
||||
self.cookie_store
|
||||
.cookie_header_for_request(account_id, "chatgpt.com", path)
|
||||
.await
|
||||
.ok()?
|
||||
}
|
||||
|
||||
pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) {
|
||||
for header in headers {
|
||||
if let Err(error) = self
|
||||
.cookie_store
|
||||
.capture_set_cookie(account_id, header)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"Failed to persist upstream set-cookie header"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) {
|
||||
let now = Utc::now();
|
||||
let cooldown = self
|
||||
.challenge_tracker
|
||||
.record_challenge(account_id, now)
|
||||
.await;
|
||||
account_pool
|
||||
.set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until)
|
||||
.await;
|
||||
if let Err(error) = self
|
||||
.cookie_store
|
||||
.expire_account_cookies_at(account_id, cooldown.cooldown_until)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"Failed to persist Cloudflare challenge cookie cleanup deadline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) {
|
||||
self.delete_account_cookies(account_id, "Cloudflare path-block")
|
||||
.await;
|
||||
let now = Utc::now();
|
||||
self.path_block_tracker
|
||||
.record_path_block(account_id, now)
|
||||
.await;
|
||||
if self
|
||||
.path_block_tracker
|
||||
.should_disable(account_id, now)
|
||||
.await
|
||||
{
|
||||
account_pool
|
||||
.set_status(account_id, AccountStatus::Disabled)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn reset_account_recovery(&self, account_id: &str) {
|
||||
self.path_block_tracker.reset(account_id).await;
|
||||
self.challenge_tracker.reset(account_id).await;
|
||||
}
|
||||
|
||||
async fn delete_account_cookies(&self, account_id: &str, reason: &str) {
|
||||
delete_account_cookies(&self.cookie_store, account_id, reason).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) {
|
||||
if let Err(error) = cookie_store.delete_account_cookies(account_id).await {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
reason,
|
||||
error = %error,
|
||||
"Failed to delete account cookies after Cloudflare recovery signal"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
//! `cyber_policy` 会话级生命周期控制器。
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc, time::Duration as StdDuration};
|
||||
|
||||
use chrono::Duration;
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::time::timeout;
|
||||
|
||||
mod store;
|
||||
mod types;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::SessionAffinityService,
|
||||
controllers::ControllerFailureFact,
|
||||
errors::ClientFailure,
|
||||
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::UpstreamFailureFacts,
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
},
|
||||
};
|
||||
|
||||
use self::{store::Store, types::SessionState};
|
||||
|
||||
const SESSION_STATE_TTL: Duration = Duration::hours(1);
|
||||
const STATE_IO_TIMEOUT: StdDuration = StdDuration::from_millis(100);
|
||||
|
||||
/// 单次请求使用的 `cyber_policy` 会话路由计划。
|
||||
pub(in crate::dispatch) struct CyberPolicyScope {
|
||||
session_key: Option<String>,
|
||||
state: SessionState,
|
||||
}
|
||||
|
||||
impl CyberPolicyScope {
|
||||
pub(in crate::dispatch) fn excluded_account_ids(&self) -> BTreeSet<String> {
|
||||
self.state.failed_account_ids.iter().cloned().collect()
|
||||
}
|
||||
|
||||
fn session_key(&self) -> Option<&str> {
|
||||
self.session_key.as_deref()
|
||||
}
|
||||
|
||||
fn has_failures(&self) -> bool {
|
||||
!self.state.failed_account_ids.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// 识别失败、维护会话状态,并为下一次请求生成账号排除集合。
|
||||
#[derive(Clone)]
|
||||
pub(in crate::dispatch) struct CyberPolicyController {
|
||||
store: Store,
|
||||
}
|
||||
|
||||
/// 已由 Cyber owner 认领的失败;调用方只能请求 owner 产生 effect 或 decision。
|
||||
pub(super) struct ClassifiedCyberPolicyFailure<'a> {
|
||||
fact: ControllerFailureFact<'a>,
|
||||
}
|
||||
|
||||
impl CyberPolicyController {
|
||||
pub(in crate::dispatch) fn new(session_affinity: Arc<SessionAffinityService>) -> Self {
|
||||
Self {
|
||||
store: Store::new(session_affinity.redis_connection()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn classify(
|
||||
fact: ControllerFailureFact<'_>,
|
||||
) -> Option<ClassifiedCyberPolicyFailure<'_>> {
|
||||
let owned = match fact {
|
||||
ControllerFailureFact::Upstream(facts) => {
|
||||
facts
|
||||
.status_code
|
||||
.is_some_and(|status| (400..=499).contains(&status))
|
||||
&& is_cyber_policy_code(facts.code.as_deref())
|
||||
}
|
||||
ControllerFailureFact::Response(failure) => is_cyber_policy_failure(failure),
|
||||
};
|
||||
owned.then_some(ClassifiedCyberPolicyFailure { fact })
|
||||
}
|
||||
|
||||
pub(super) fn client_failure(classified: &ClassifiedCyberPolicyFailure<'_>) -> ClientFailure {
|
||||
match classified.fact {
|
||||
ControllerFailureFact::Upstream(facts) => ClientFailure::new(
|
||||
failure_from_upstream_facts(facts),
|
||||
facts.status_code.unwrap_or(400),
|
||||
true,
|
||||
),
|
||||
ControllerFailureFact::Response(failure) => {
|
||||
ClientFailure::new(failure.clone(), 400, true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decision(
|
||||
observation: &AttemptObservation,
|
||||
classified: ClassifiedCyberPolicyFailure<'_>,
|
||||
) -> AttemptDecision {
|
||||
let client_failure = Self::client_failure(&classified);
|
||||
if observation.routing.can_retry_next_candidate {
|
||||
AttemptDecision::RetryNextCandidate {
|
||||
exhaustion: None,
|
||||
on_exhaustion: Some(client_failure),
|
||||
}
|
||||
} else {
|
||||
AttemptDecision::Return(AttemptReturnKind::Failed(client_failure))
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn prepare(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
) -> CyberPolicyScope {
|
||||
let Some(session_key) = session_key(request) else {
|
||||
return CyberPolicyScope {
|
||||
session_key: None,
|
||||
state: SessionState::default(),
|
||||
};
|
||||
};
|
||||
let state = match timeout(STATE_IO_TIMEOUT, self.store.load(&session_key)).await {
|
||||
Ok(Ok(state)) => state.unwrap_or_default(),
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(error = %error, "Failed to read cyber policy session state");
|
||||
SessionState::default()
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!("Timed out reading cyber policy session state");
|
||||
SessionState::default()
|
||||
}
|
||||
};
|
||||
CyberPolicyScope {
|
||||
session_key: Some(session_key),
|
||||
state,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn exclude_account(&self, plan: &CyberPolicyScope, account_id: &str) {
|
||||
let Some(session_key) = plan.session_key() else {
|
||||
return;
|
||||
};
|
||||
match timeout(
|
||||
STATE_IO_TIMEOUT,
|
||||
self.store
|
||||
.record_failure(session_key, account_id, SESSION_STATE_TTL),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(state)) => tracing::warn!(
|
||||
account_id,
|
||||
excluded_account_count = state.failed_account_ids.len(),
|
||||
"Excluded cyber policy account from this session"
|
||||
),
|
||||
Ok(Err(error)) => tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"Failed to record cyber policy session state"
|
||||
),
|
||||
Err(_) => tracing::warn!(account_id, "Timed out recording cyber policy session state"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn observe_success(&self, plan: &CyberPolicyScope) {
|
||||
if !plan.has_failures() {
|
||||
return;
|
||||
}
|
||||
let Some(session_key) = plan.session_key() else {
|
||||
return;
|
||||
};
|
||||
match timeout(
|
||||
STATE_IO_TIMEOUT,
|
||||
self.store.clear(session_key, &plan.state.revision),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(error = %error, "Failed to clear cyber policy session state");
|
||||
}
|
||||
Err(_) => tracing::warn!("Timed out clearing cyber policy session state"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn session_key(request: &CodexResponsesRequest) -> Option<String> {
|
||||
if request.previous_response_id().is_some() {
|
||||
return None;
|
||||
}
|
||||
let api_key_id = non_empty(request.client_api_key_id.as_deref())?;
|
||||
let explicit_session_id = non_empty(request.client_session_id.as_deref())
|
||||
.or_else(|| non_empty(request.client_conversation_id.as_deref()))
|
||||
.or_else(|| {
|
||||
request
|
||||
.explicit_prompt_cache_key
|
||||
.then(|| non_empty(request.prompt_cache_key()))
|
||||
.flatten()
|
||||
})?;
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(b"cyber-policy-session\0");
|
||||
hasher.update(api_key_id.as_bytes());
|
||||
hasher.update(b"\0");
|
||||
hasher.update(explicit_session_id.as_bytes());
|
||||
Some(hex::encode(hasher.finalize()))
|
||||
}
|
||||
|
||||
fn is_cyber_policy_code(code: Option<&str>) -> bool {
|
||||
code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy"))
|
||||
}
|
||||
|
||||
fn failure_from_upstream_facts(facts: &UpstreamFailureFacts) -> ResponsesSseFailure {
|
||||
ResponsesSseFailure {
|
||||
event: "error".to_string(),
|
||||
message: facts.message.clone(),
|
||||
upstream_code: facts.code.clone(),
|
||||
upstream_type: None,
|
||||
explicit_status_code: facts.status_code,
|
||||
retry_after_seconds: facts.retry_after_seconds,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn is_cyber_policy_failure(failure: &ResponsesSseFailure) -> bool {
|
||||
is_cyber_policy_code(failure.upstream_code.as_deref())
|
||||
}
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
use chrono::Duration;
|
||||
use redis::AsyncCommands;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::redis::RedisConnection;
|
||||
|
||||
use super::types::SessionState;
|
||||
|
||||
const STATE_VERSION: &str = "cyber-policy:v2";
|
||||
|
||||
const RECORD_FAILURE_SCRIPT: &str = r#"
|
||||
local raw = redis.call('GET', KEYS[1])
|
||||
local state = { failedAccountIds = {}, revision = '' }
|
||||
if raw then
|
||||
local ok, decoded = pcall(cjson.decode, raw)
|
||||
if ok and type(decoded) == 'table' and type(decoded.failedAccountIds) == 'table' then
|
||||
state = decoded
|
||||
end
|
||||
end
|
||||
|
||||
local account_id = ARGV[1]
|
||||
local known = false
|
||||
for _, failed_account_id in ipairs(state.failedAccountIds) do
|
||||
if failed_account_id == account_id then
|
||||
known = true
|
||||
break
|
||||
end
|
||||
end
|
||||
|
||||
if not known then
|
||||
table.insert(state.failedAccountIds, account_id)
|
||||
end
|
||||
|
||||
state.revision = ARGV[3]
|
||||
|
||||
local encoded = cjson.encode(state)
|
||||
redis.call('SET', KEYS[1], encoded, 'EX', tonumber(ARGV[2]))
|
||||
return encoded
|
||||
"#;
|
||||
|
||||
const CLEAR_STATE_SCRIPT: &str = r#"
|
||||
local raw = redis.call('GET', KEYS[1])
|
||||
if not raw then
|
||||
return 0
|
||||
end
|
||||
local ok, state = pcall(cjson.decode, raw)
|
||||
if not ok or type(state) ~= 'table' then
|
||||
return 0
|
||||
end
|
||||
if tostring(state.revision or '') ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
return redis.call('DEL', KEYS[1])
|
||||
"#;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Store {
|
||||
redis: RedisConnection,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum StoreError {
|
||||
#[error("Redis cyber policy operation failed: {0}")]
|
||||
Redis(#[from] redis::RedisError),
|
||||
#[error("invalid Redis cyber policy value: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
}
|
||||
|
||||
impl Store {
|
||||
pub fn new(redis: RedisConnection) -> Self {
|
||||
Self { redis }
|
||||
}
|
||||
|
||||
pub async fn load(&self, session_key: &str) -> Result<Option<SessionState>, StoreError> {
|
||||
let mut connection = self.redis.manager();
|
||||
let value: Option<String> = connection.get(self.key(session_key)).await?;
|
||||
value
|
||||
.map(|value| serde_json::from_str(&value))
|
||||
.transpose()
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub async fn record_failure(
|
||||
&self,
|
||||
session_key: &str,
|
||||
account_id: &str,
|
||||
ttl: Duration,
|
||||
) -> Result<SessionState, StoreError> {
|
||||
let mut connection = self.redis.manager();
|
||||
let revision = uuid::Uuid::new_v4().to_string();
|
||||
let value: String = redis::Script::new(RECORD_FAILURE_SCRIPT)
|
||||
.key(self.key(session_key))
|
||||
.arg(account_id)
|
||||
.arg(ttl.num_seconds().max(1))
|
||||
.arg(revision)
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(serde_json::from_str(&value)?)
|
||||
}
|
||||
|
||||
pub async fn clear(
|
||||
&self,
|
||||
session_key: &str,
|
||||
expected_revision: &str,
|
||||
) -> Result<bool, StoreError> {
|
||||
let mut connection = self.redis.manager();
|
||||
let removed: usize = redis::Script::new(CLEAR_STATE_SCRIPT)
|
||||
.key(self.key(session_key))
|
||||
.arg(expected_revision)
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(removed > 0)
|
||||
}
|
||||
|
||||
fn key(&self, session_key: &str) -> String {
|
||||
self.redis
|
||||
.key(&format!("{STATE_VERSION}:session:{session_key}"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionState {
|
||||
pub failed_account_ids: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub revision: String,
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
//! 上游 typed failure 到 fleet 账号事实的边界映射。
|
||||
|
||||
use crate::{
|
||||
fleet::account_gateway::AccountFailureObservation,
|
||||
upstream::openai::{failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure},
|
||||
};
|
||||
|
||||
pub(super) fn from_upstream(facts: &UpstreamFailureFacts) -> AccountFailureObservation {
|
||||
AccountFailureObservation {
|
||||
status_code: facts.status_code,
|
||||
code: facts.code.clone(),
|
||||
error_type: facts.error_type.clone(),
|
||||
identity_authorization_error: facts.identity_authorization_error.clone(),
|
||||
identity_error_code: facts.identity_error_code.clone(),
|
||||
message: facts.message.clone(),
|
||||
body: facts.body.clone(),
|
||||
retry_after_seconds: facts.retry_after_seconds,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn from_response(failure: &ResponsesSseFailure) -> AccountFailureObservation {
|
||||
AccountFailureObservation {
|
||||
status_code: failure.explicit_status_code,
|
||||
code: failure.upstream_code.clone(),
|
||||
error_type: failure.upstream_type.clone(),
|
||||
message: failure.message.clone(),
|
||||
body: failure.message.clone(),
|
||||
retry_after_seconds: failure.retry_after_seconds,
|
||||
..AccountFailureObservation::default()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,626 @@
|
||||
//! previous response 所有权、重放与恢复规则的唯一 owner。
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::{SessionAffinityEntry, SessionAffinityService},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
|
||||
AttemptRoutingFacts, CompleteResponseFacts, PinnedCandidateAcquireFailureKind,
|
||||
},
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::{UpstreamFailureFacts, UpstreamFailureKind},
|
||||
protocol::responses::{
|
||||
CodexResponsesRequest, LocalReplayItem, LocalReplayTranscript, PreviousResponseScope,
|
||||
StreamCommitPolicy,
|
||||
},
|
||||
transport::CodexBackendTransport,
|
||||
transport::websocket::PreviousResponseUnavailableReason,
|
||||
},
|
||||
};
|
||||
|
||||
const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
const CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE: &str =
|
||||
"previous response history cannot be sent to another account";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum HistorySource {
|
||||
None,
|
||||
Managed,
|
||||
ConnectionLocal,
|
||||
ExternalUnknown,
|
||||
}
|
||||
|
||||
/// 单个请求内的 previous response 状态;策略只由 `HistoryController` 解释。
|
||||
pub(super) struct HistoryState {
|
||||
source: HistorySource,
|
||||
managed_entry: Option<SessionAffinityEntry>,
|
||||
local_replay_transcript: Option<Arc<LocalReplayTranscript>>,
|
||||
local_replay_account_id: Option<String>,
|
||||
external_attempted_account_id: Option<String>,
|
||||
replay_owner_account: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(in crate::dispatch) struct ConnectionReplaySnapshot {
|
||||
last_response_id: Option<String>,
|
||||
account_id: Option<String>,
|
||||
transcript: Option<Arc<LocalReplayTranscript>>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct ConnectionReplayPlan {
|
||||
update: ConnectionReplayUpdate,
|
||||
}
|
||||
|
||||
enum ConnectionReplayUpdate {
|
||||
Replace,
|
||||
Append { history: Arc<LocalReplayTranscript> },
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// 已确认来源和目标账号不同的完整 transcript 重放边界。
|
||||
pub(in crate::dispatch) struct CrossAccountReplay {
|
||||
source_account_id: String,
|
||||
target_account_id: String,
|
||||
}
|
||||
|
||||
impl CrossAccountReplay {
|
||||
fn between(source_account_id: &str, target_account_id: &str) -> Option<Self> {
|
||||
(source_account_id != target_account_id).then(|| Self {
|
||||
source_account_id: source_account_id.to_string(),
|
||||
target_account_id: target_account_id.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn validate_target(&self, account_id: &str) {
|
||||
debug_assert_ne!(self.source_account_id, self.target_account_id);
|
||||
debug_assert_eq!(self.target_account_id, account_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct PreparedHistoryRequest {
|
||||
pub(super) request: CodexResponsesRequest,
|
||||
pub(super) cross_account_replay: Option<CrossAccountReplay>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct HistoryController;
|
||||
|
||||
pub(in crate::dispatch) fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) {
|
||||
if request.input().is_empty() {
|
||||
return;
|
||||
}
|
||||
let input = request
|
||||
.input()
|
||||
.iter()
|
||||
.cloned()
|
||||
.filter_map(sanitize_cross_account_item)
|
||||
.collect();
|
||||
request.set_input(input);
|
||||
}
|
||||
|
||||
impl HistoryController {
|
||||
pub(in crate::dispatch) fn new_connection_replay_snapshot() -> ConnectionReplaySnapshot {
|
||||
ConnectionReplaySnapshot::default()
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn prepare_connection_replay(
|
||||
snapshot: &ConnectionReplaySnapshot,
|
||||
request: &mut CodexResponsesRequest,
|
||||
) -> ConnectionReplayPlan {
|
||||
let update = match request.previous_response_id() {
|
||||
None => {
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Replace
|
||||
}
|
||||
Some(previous_response_id)
|
||||
if snapshot.last_response_id.as_deref() == Some(previous_response_id) =>
|
||||
{
|
||||
if let (Some(account_id), Some(history)) =
|
||||
(snapshot.account_id.as_ref(), snapshot.transcript.as_ref())
|
||||
{
|
||||
request.local_replay_transcript = Some(Arc::clone(history));
|
||||
request.local_replay_account_id = Some(account_id.clone());
|
||||
ConnectionReplayUpdate::Append {
|
||||
history: Arc::clone(history),
|
||||
}
|
||||
} else {
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Unavailable
|
||||
}
|
||||
}
|
||||
Some(_) => {
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Unavailable
|
||||
}
|
||||
};
|
||||
request.local_replay_available = !matches!(update, ConnectionReplayUpdate::Unavailable);
|
||||
ConnectionReplayPlan { update }
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn commit_connection_replay(
|
||||
snapshot: &mut ConnectionReplaySnapshot,
|
||||
plan: ConnectionReplayPlan,
|
||||
response_id: String,
|
||||
output: Vec<Value>,
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
) {
|
||||
snapshot.last_response_id = Some(response_id);
|
||||
snapshot.account_id = Some(account_id.clone());
|
||||
match plan.update {
|
||||
ConnectionReplayUpdate::Replace => {
|
||||
snapshot.transcript = Some(Arc::new(transcript_for_turn(
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
)));
|
||||
}
|
||||
ConnectionReplayUpdate::Append { mut history } => {
|
||||
if continued_from_previous_response {
|
||||
append_turn(
|
||||
Arc::make_mut(&mut history),
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
);
|
||||
snapshot.transcript = Some(history);
|
||||
} else {
|
||||
snapshot.transcript = Some(Arc::new(transcript_for_turn(
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
)));
|
||||
}
|
||||
}
|
||||
ConnectionReplayUpdate::Unavailable => snapshot.transcript = None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn enter(
|
||||
affinity: &SessionAffinityService,
|
||||
request: &CodexResponsesRequest,
|
||||
) -> HistoryState {
|
||||
let Some(previous_response_id) = request.previous_response_id() else {
|
||||
return HistoryState::empty();
|
||||
};
|
||||
let managed_entry = match timeout(
|
||||
REDIS_BEST_EFFORT_TIMEOUT,
|
||||
affinity.lookup(previous_response_id, Utc::now()),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(entry) => entry,
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
previous_response_id,
|
||||
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
|
||||
"Timed out reading previous response affinity; continuing without managed history"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let (local_replay_transcript, local_replay_account_id) = match (
|
||||
request.local_replay_transcript.clone(),
|
||||
request.local_replay_account_id.clone(),
|
||||
) {
|
||||
(Some(transcript), Some(account_id)) => (Some(transcript), Some(account_id)),
|
||||
(None, None) => (None, None),
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
previous_response_id,
|
||||
"Ignoring incomplete connection-local replay provenance"
|
||||
);
|
||||
(None, None)
|
||||
}
|
||||
};
|
||||
let has_complete_local_replay = local_replay_transcript.is_some();
|
||||
let source = match (&managed_entry, has_complete_local_replay) {
|
||||
(Some(_), _) => HistorySource::Managed,
|
||||
(None, true) => HistorySource::ConnectionLocal,
|
||||
(None, false) => HistorySource::ExternalUnknown,
|
||||
};
|
||||
HistoryState {
|
||||
source,
|
||||
managed_entry,
|
||||
local_replay_transcript,
|
||||
local_replay_account_id,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn preferred_account_id(state: &HistoryState) -> Option<&str> {
|
||||
state.local_replay_account_id.as_deref().or_else(|| {
|
||||
state
|
||||
.managed_entry
|
||||
.as_ref()
|
||||
.map(|entry| entry.account_id.as_str())
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn prepare_attempt(
|
||||
state: &mut HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Result<PreparedHistoryRequest, String> {
|
||||
let prepared = match state.source {
|
||||
HistorySource::None => Some(PreparedHistoryRequest::plain(request.clone())),
|
||||
HistorySource::ExternalUnknown => {
|
||||
match state.external_attempted_account_id.as_deref() {
|
||||
None => {
|
||||
state.external_attempted_account_id = Some(account_id.to_string());
|
||||
Some(PreparedHistoryRequest::plain(external_history_request(
|
||||
request,
|
||||
)))
|
||||
}
|
||||
Some(first_account_id) if first_account_id == account_id => Some(
|
||||
PreparedHistoryRequest::plain(external_history_request(request)),
|
||||
),
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
HistorySource::ConnectionLocal => full_replay_from_state(state, request, account_id),
|
||||
HistorySource::Managed => prepare_managed_attempt(state, request, account_id),
|
||||
};
|
||||
prepared.ok_or_else(|| CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE.to_string())
|
||||
}
|
||||
|
||||
pub(super) fn prepare_same_account_retry(state: &mut HistoryState, account_id: &str) -> bool {
|
||||
if !Self::can_retry_same_account(state, account_id) {
|
||||
return false;
|
||||
}
|
||||
state.replay_owner_account = true;
|
||||
true
|
||||
}
|
||||
|
||||
pub(super) fn routing_facts(
|
||||
state: &HistoryState,
|
||||
account_id: Option<&str>,
|
||||
) -> AttemptRoutingFacts {
|
||||
AttemptRoutingFacts {
|
||||
external_origin: state.source == HistorySource::ExternalUnknown,
|
||||
can_retry_same_account: account_id
|
||||
.is_some_and(|account_id| Self::can_retry_same_account(state, account_id)),
|
||||
can_retry_next_candidate: match state.source {
|
||||
HistorySource::None | HistorySource::ConnectionLocal => true,
|
||||
HistorySource::Managed => state.local_replay_transcript.is_some(),
|
||||
HistorySource::ExternalUnknown => false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn conversation_id<'a>(
|
||||
state: &'a HistoryState,
|
||||
request: &'a CodexResponsesRequest,
|
||||
) -> Option<&'a str> {
|
||||
state
|
||||
.managed_entry
|
||||
.as_ref()
|
||||
.map(|entry| entry.conversation_id.as_str())
|
||||
.or_else(|| non_empty(request.local_conversation_id.as_deref()))
|
||||
}
|
||||
|
||||
pub(super) fn decide(observation: &AttemptObservation) -> Option<AttemptDecision> {
|
||||
if let AttemptObservationKind::PinnedCandidateUnavailable { kind, .. } = observation.kind {
|
||||
return Some(match kind {
|
||||
PinnedCandidateAcquireFailureKind::Busy => {
|
||||
AttemptDecision::Return(AttemptReturnKind::ContinuationBusy)
|
||||
}
|
||||
PinnedCandidateAcquireFailureKind::Unavailable => {
|
||||
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
|
||||
message: "the account that owns the previous response is unavailable for same-account recovery".to_string(),
|
||||
})
|
||||
}
|
||||
});
|
||||
}
|
||||
let detail = match &observation.kind {
|
||||
AttemptObservationKind::UpstreamFailure(facts) if is_continuation_busy(facts) => {
|
||||
if observation.routing.can_retry_same_account {
|
||||
return Some(AttemptDecision::RetrySameAccount);
|
||||
}
|
||||
return Some(AttemptDecision::Return(AttemptReturnKind::ContinuationBusy));
|
||||
}
|
||||
AttemptObservationKind::UpstreamFailure(facts) if is_history_failure(facts) => {
|
||||
facts.body.clone()
|
||||
}
|
||||
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Failed(failure))
|
||||
| AttemptObservationKind::StreamFailure(failure)
|
||||
if is_history_code(failure.upstream_code.as_deref()) =>
|
||||
{
|
||||
crate::dispatch::failure::sse::sse_failure_error_body(failure)
|
||||
}
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
if observation.routing.can_retry_same_account {
|
||||
return Some(AttemptDecision::RetrySameAccount);
|
||||
}
|
||||
if observation.routing.external_origin {
|
||||
return Some(AttemptDecision::Return(AttemptReturnKind::Observed));
|
||||
}
|
||||
Some(AttemptDecision::Return(
|
||||
AttemptReturnKind::RouteUnavailable { message: detail },
|
||||
))
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn continuation_scope(
|
||||
request: &CodexResponsesRequest,
|
||||
transport: CodexBackendTransport,
|
||||
connection_local_continuation: bool,
|
||||
) -> PreviousResponseScope {
|
||||
if request.store() {
|
||||
PreviousResponseScope::Persisted
|
||||
} else if transport == CodexBackendTransport::WebSocket && connection_local_continuation {
|
||||
PreviousResponseScope::ConnectionLocal
|
||||
} else if request.local_replay_available {
|
||||
PreviousResponseScope::ReplayRequired
|
||||
} else {
|
||||
PreviousResponseScope::Unavailable
|
||||
}
|
||||
}
|
||||
|
||||
fn can_retry_same_account(state: &HistoryState, account_id: &str) -> bool {
|
||||
let Some(entry) = state.managed_entry.as_ref() else {
|
||||
return false;
|
||||
};
|
||||
entry.account_id == account_id
|
||||
&& state.local_replay_transcript.is_some()
|
||||
&& !state.replay_owner_account
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_managed_attempt(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
let entry = state.managed_entry.as_ref()?;
|
||||
if entry.account_id != account_id || state.replay_owner_account {
|
||||
return replay_for_conversation(state, request, account_id, &entry.conversation_id);
|
||||
}
|
||||
|
||||
match entry.continuation_scope {
|
||||
PreviousResponseScope::Persisted
|
||||
| PreviousResponseScope::ConnectionLocal
|
||||
| PreviousResponseScope::ExternalUnknown => {
|
||||
let mut prepared = request.clone();
|
||||
prepared.local_conversation_id = Some(entry.conversation_id.clone());
|
||||
prepared.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
prepared.previous_response_scope = Some(entry.continuation_scope);
|
||||
if prepared.turn_state.as_deref().is_none_or(str::is_empty) {
|
||||
prepared.turn_state.clone_from(&entry.turn_state);
|
||||
}
|
||||
Some(PreparedHistoryRequest::plain(prepared))
|
||||
}
|
||||
PreviousResponseScope::ReplayRequired => {
|
||||
replay_for_conversation(state, request, account_id, &entry.conversation_id)
|
||||
}
|
||||
PreviousResponseScope::Unavailable => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn replay_for_conversation(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
conversation_id: &str,
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
full_replay_from_state(state, request, account_id).map(|mut prepared| {
|
||||
prepared.request.local_conversation_id = Some(conversation_id.to_string());
|
||||
prepared
|
||||
})
|
||||
}
|
||||
|
||||
fn full_replay_from_state(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
let transcript = state.local_replay_transcript.as_deref()?;
|
||||
let source_account_id = state.local_replay_account_id.as_deref()?;
|
||||
Some(full_replay_request(
|
||||
request,
|
||||
transcript,
|
||||
source_account_id,
|
||||
account_id,
|
||||
))
|
||||
}
|
||||
|
||||
impl HistoryState {
|
||||
fn empty() -> Self {
|
||||
Self {
|
||||
source: HistorySource::None,
|
||||
managed_entry: None,
|
||||
local_replay_transcript: None,
|
||||
local_replay_account_id: None,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest {
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope =
|
||||
Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown);
|
||||
request
|
||||
}
|
||||
|
||||
fn full_replay_request(
|
||||
original: &CodexResponsesRequest,
|
||||
transcript: &LocalReplayTranscript,
|
||||
source_account_id: &str,
|
||||
target_account_id: &str,
|
||||
) -> PreparedHistoryRequest {
|
||||
let cross_account_replay = CrossAccountReplay::between(source_account_id, target_account_id);
|
||||
let mut request = original.clone();
|
||||
let mut input = replay_input_for_account(transcript, target_account_id);
|
||||
input.reserve(original.input().len());
|
||||
input.extend_from_slice(original.input());
|
||||
request.set_input(input);
|
||||
request.set_previous_response_id(None);
|
||||
request.turn_state = None;
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
PreparedHistoryRequest {
|
||||
request,
|
||||
cross_account_replay,
|
||||
}
|
||||
}
|
||||
|
||||
impl PreparedHistoryRequest {
|
||||
fn plain(request: CodexResponsesRequest) -> Self {
|
||||
Self {
|
||||
request,
|
||||
cross_account_replay: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_local_replay(request: &mut CodexResponsesRequest) {
|
||||
request.local_replay_transcript = None;
|
||||
request.local_replay_account_id = None;
|
||||
}
|
||||
|
||||
fn transcript_for_turn(
|
||||
input: Vec<Value>,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) -> LocalReplayTranscript {
|
||||
let mut transcript = LocalReplayTranscript::default();
|
||||
append_client_input(&mut transcript, input);
|
||||
append_account_output(&mut transcript, output, account_id);
|
||||
transcript
|
||||
}
|
||||
|
||||
fn append_turn(
|
||||
transcript: &mut LocalReplayTranscript,
|
||||
input: Vec<Value>,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) {
|
||||
project_transcript_to_account(transcript, account_id);
|
||||
append_client_input(transcript, input);
|
||||
append_account_output(transcript, output, account_id);
|
||||
}
|
||||
|
||||
fn append_client_input(transcript: &mut LocalReplayTranscript, input: Vec<Value>) {
|
||||
transcript
|
||||
.items
|
||||
.extend(input.into_iter().map(LocalReplayItem::ClientInput));
|
||||
}
|
||||
|
||||
fn append_account_output(
|
||||
transcript: &mut LocalReplayTranscript,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) {
|
||||
transcript.items.extend(
|
||||
output
|
||||
.into_iter()
|
||||
.map(|item| LocalReplayItem::AccountOutput {
|
||||
account_id: account_id.to_string(),
|
||||
item,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn project_transcript_to_account(transcript: &mut LocalReplayTranscript, account_id: &str) {
|
||||
transcript.items = transcript
|
||||
.items
|
||||
.drain(..)
|
||||
.filter_map(|item| match item {
|
||||
LocalReplayItem::AccountOutput {
|
||||
account_id: owner,
|
||||
item,
|
||||
} if owner != account_id => {
|
||||
sanitize_cross_account_item(item).map(LocalReplayItem::SanitizedOutput)
|
||||
}
|
||||
item => Some(item),
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
fn replay_input_for_account(transcript: &LocalReplayTranscript, account_id: &str) -> Vec<Value> {
|
||||
transcript
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
LocalReplayItem::ClientInput(value) | LocalReplayItem::SanitizedOutput(value) => {
|
||||
Some(value.clone())
|
||||
}
|
||||
LocalReplayItem::AccountOutput {
|
||||
account_id: owner,
|
||||
item,
|
||||
} if owner == account_id => Some(without_output_id(item.clone())),
|
||||
LocalReplayItem::AccountOutput { item, .. } => {
|
||||
sanitize_cross_account_item(item.clone())
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn without_output_id(mut item: Value) -> Value {
|
||||
if let Value::Object(object) = &mut item {
|
||||
object.remove("id");
|
||||
}
|
||||
item
|
||||
}
|
||||
|
||||
fn sanitize_cross_account_item(mut item: Value) -> Option<Value> {
|
||||
if let Value::Object(object) = &mut item {
|
||||
if matches!(
|
||||
object.get("type").and_then(Value::as_str),
|
||||
Some("compaction" | "compaction_summary" | "context_compaction")
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
object.remove("id");
|
||||
object.remove("encrypted_content");
|
||||
}
|
||||
Some(item)
|
||||
}
|
||||
|
||||
fn is_continuation_busy(facts: &UpstreamFailureFacts) -> bool {
|
||||
matches!(
|
||||
facts.kind,
|
||||
UpstreamFailureKind::ContinuationUnavailable(
|
||||
PreviousResponseUnavailableReason::ConnectionBusy
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
fn is_history_failure(facts: &UpstreamFailureFacts) -> bool {
|
||||
facts
|
||||
.code
|
||||
.as_deref()
|
||||
.is_some_and(|code| is_history_code(Some(code)))
|
||||
|| matches!(
|
||||
facts.kind,
|
||||
UpstreamFailureKind::ContinuationUnavailable(reason)
|
||||
if reason != PreviousResponseUnavailableReason::ConnectionBusy
|
||||
)
|
||||
}
|
||||
|
||||
fn is_history_code(code: Option<&str>) -> bool {
|
||||
matches!(
|
||||
code,
|
||||
Some(
|
||||
"previous_response_not_found"
|
||||
| "invalid_encrypted_content"
|
||||
| "missing_tool_output"
|
||||
| "no_tool_output"
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,129 @@
|
||||
//! 配额失败规则 owner。
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
controllers::ControllerFailureFact,
|
||||
errors::ClientFailure,
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
|
||||
},
|
||||
fleet::{
|
||||
account_failure::{
|
||||
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
|
||||
classify_account_failure,
|
||||
},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure,
|
||||
transport::CodexBackendClient,
|
||||
},
|
||||
};
|
||||
|
||||
pub(super) struct QuotaController;
|
||||
|
||||
impl QuotaController {
|
||||
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
|
||||
classify(observation)
|
||||
}
|
||||
|
||||
pub(super) fn classify_failure(
|
||||
account_id: &str,
|
||||
failure: ControllerFailureFact<'_>,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
classify_failure(account_id, failure)
|
||||
}
|
||||
|
||||
pub(super) fn decision(
|
||||
observation: &AttemptObservation,
|
||||
classified: ClassifiedQuotaFailure,
|
||||
) -> AttemptDecision {
|
||||
if observation.routing.can_retry_next_candidate {
|
||||
return AttemptDecision::RetryNextCandidate {
|
||||
exhaustion: Some(classified.exhaustion),
|
||||
on_exhaustion: None,
|
||||
};
|
||||
}
|
||||
AttemptDecision::Return(AttemptReturnKind::Observed)
|
||||
}
|
||||
|
||||
pub(super) async fn apply_effect(
|
||||
account_pool: &AccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
classified: &ClassifiedQuotaFailure,
|
||||
) {
|
||||
let Some(account_id) = classified.exhaustion.account_id.as_deref() else {
|
||||
return;
|
||||
};
|
||||
if let Some(effect) = &classified.effect {
|
||||
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn client_failure(failure: &ResponsesSseFailure) -> ClientFailure {
|
||||
ClientFailure::new(failure.clone(), 429, false)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct ClassifiedQuotaFailure {
|
||||
exhaustion: AccountExhaustionRecord,
|
||||
effect: Option<AccountStateEffect>,
|
||||
}
|
||||
|
||||
fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
|
||||
let account_id = observation.account.as_ref()?.id.as_str();
|
||||
ControllerFailureFact::from_attempt(observation)
|
||||
.and_then(|failure| classify_failure(account_id, failure))
|
||||
}
|
||||
|
||||
fn classify_failure(
|
||||
account_id: &str,
|
||||
failure: ControllerFailureFact<'_>,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
match failure {
|
||||
ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts),
|
||||
ControllerFailureFact::Response(failure) => classify_stream(account_id, failure),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_upstream(
|
||||
account_id: &str,
|
||||
facts: &UpstreamFailureFacts,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
classified_quota_failure(
|
||||
account_id,
|
||||
&facts.body,
|
||||
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classify_stream(
|
||||
account_id: &str,
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
|
||||
classified_quota_failure(
|
||||
account_id,
|
||||
&body,
|
||||
classify_account_failure(&super::failure_observation::from_response(failure))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classified_quota_failure(
|
||||
account_id: &str,
|
||||
message: &str,
|
||||
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
let kind = match classified.kind {
|
||||
AccountFailureKind::RateLimited => ExhaustedAccountKind::RateLimited,
|
||||
AccountFailureKind::QuotaExhausted => ExhaustedAccountKind::QuotaExhausted,
|
||||
AccountFailureKind::ModelUnsupported
|
||||
| AccountFailureKind::Expired
|
||||
| AccountFailureKind::Disabled
|
||||
| AccountFailureKind::Banned => return None,
|
||||
};
|
||||
Some(ClassifiedQuotaFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(account_id, kind, message),
|
||||
effect: classified.effect,
|
||||
})
|
||||
}
|
||||
+501
-126
@@ -1,42 +1,44 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
dispatch::errors::{
|
||||
backend_transport_name, insert_dispatch_error_metadata,
|
||||
insert_upstream_diagnostics_metadata, upstream_error_http_status, DispatchErrorMetadata,
|
||||
DispatchErrorMetadata, backend_transport_name, insert_dispatch_error_metadata,
|
||||
insert_upstream_diagnostics_metadata, upstream_error_http_status,
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
telemetry::{
|
||||
ops::types::OpsErrorLog,
|
||||
recorder::{
|
||||
enrich_event_route_metadata, enrich_usage_record_identity,
|
||||
event_kind as response_event_kind, reasoning_effort_from_request,
|
||||
record_dispatch_error_event, DispatchErrorLogRecord, Recorder,
|
||||
recorder::Recorder,
|
||||
usage::types::{
|
||||
UsageRecord, UsageRecordLevel, metadata_i64, metadata_service_tier,
|
||||
metadata_string as metadata_string_any,
|
||||
},
|
||||
usage::types::UsageRecordLevel,
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::events::TokenUsage,
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
transport::{
|
||||
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
|
||||
WebSocketPoolDecision,
|
||||
CodexBackendTransport, CodexClientError, CodexResponseMetadata,
|
||||
CodexUpstreamDiagnostics, WebSocketPoolDecision,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
use super::{
|
||||
const RESPONSE_EVENT_KIND: &str = "v1.response";
|
||||
const RESPONSES_API_KIND: &str = "responses";
|
||||
|
||||
use crate::dispatch::{
|
||||
errors::{
|
||||
dispatch_error_metadata, enrich_response_dispatch_error_metadata, ResponseDispatchError,
|
||||
},
|
||||
stream::{
|
||||
live::{latest_response_id, LiveResponseStreamContext},
|
||||
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
|
||||
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
|
||||
ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata,
|
||||
},
|
||||
failure::sse::stream_failure_metadata,
|
||||
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
|
||||
};
|
||||
|
||||
use super::StreamContext;
|
||||
|
||||
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
|
||||
pub(super) recorder: &'a Recorder,
|
||||
pub(super) request_id: &'a str,
|
||||
@@ -64,6 +66,7 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
|
||||
pub(super) transport: CodexBackendTransport,
|
||||
pub(super) request: &'a CodexResponsesRequest,
|
||||
pub(super) failure: &'a ResponsesSseFailure,
|
||||
pub(super) status_code: i64,
|
||||
pub(super) diagnostics: &'a CodexUpstreamDiagnostics,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
pub(super) prefetched: &'a [u8],
|
||||
@@ -85,12 +88,117 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> {
|
||||
pub(super) error: &'a ResponseDispatchError,
|
||||
}
|
||||
|
||||
pub(super) struct ResponseDispatchErrorDetails<'a> {
|
||||
pub(super) struct ResponseUsageEventRecord<'a> {
|
||||
pub(super) recorder: &'a Recorder,
|
||||
pub(super) request_id: &'a str,
|
||||
pub(super) client_api_key_id: Option<&'a str>,
|
||||
pub(super) account_id: Option<&'a str>,
|
||||
pub(super) stream: bool,
|
||||
pub(super) compact: bool,
|
||||
pub(super) transport: Option<&'a str>,
|
||||
pub(super) account_id: &'a str,
|
||||
pub(super) route: &'a str,
|
||||
pub(super) model: &'a str,
|
||||
pub(super) requested_model: Option<&'a str>,
|
||||
pub(super) client_ip: Option<&'a str>,
|
||||
pub(super) client_user_agent: Option<&'a str>,
|
||||
pub(super) reasoning_effort: Option<&'a str>,
|
||||
pub(super) service_tier: Option<&'a str>,
|
||||
pub(super) started_at: Instant,
|
||||
pub(super) status_code: i64,
|
||||
pub(super) message: &'a str,
|
||||
pub(super) usage: Option<TokenUsage>,
|
||||
pub(super) metadata: Value,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
}
|
||||
|
||||
pub(super) struct LiveResponseStreamEventRecord<'a> {
|
||||
pub(super) context: &'a StreamContext<'a>,
|
||||
pub(super) status_code: i64,
|
||||
pub(super) level: UsageRecordLevel,
|
||||
pub(super) message: &'a str,
|
||||
pub(super) usage: Option<TokenUsage>,
|
||||
pub(super) metadata: Value,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
pub(super) body: &'a str,
|
||||
}
|
||||
|
||||
struct DispatchErrorLogRecord<'a> {
|
||||
recorder: &'a Recorder,
|
||||
request_id: &'a str,
|
||||
client_api_key_id: Option<&'a str>,
|
||||
provider: Option<&'a str>,
|
||||
account_id: Option<&'a str>,
|
||||
route: &'a str,
|
||||
model: &'a str,
|
||||
started_at: Instant,
|
||||
status_code: i64,
|
||||
message: &'a str,
|
||||
metadata: Value,
|
||||
}
|
||||
|
||||
pub(super) async fn record_response_event(record: ResponseUsageEventRecord<'_>) {
|
||||
let mut metadata = record.metadata;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
|
||||
let mut event = UsageRecord::new(
|
||||
RESPONSE_EVENT_KIND,
|
||||
record.message,
|
||||
record.account_id,
|
||||
record.model,
|
||||
record.status_code,
|
||||
);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
|
||||
event.route = Some(record.route.to_string());
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
event.requested_model = normalized_string(record.requested_model.or(Some(record.model)));
|
||||
event.upstream_model = normalized_string(Some(record.model));
|
||||
event.service_tier = normalized_string(record.service_tier);
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
record.requested_model,
|
||||
record.model,
|
||||
record.client_ip,
|
||||
record.client_user_agent,
|
||||
record.reasoning_effort,
|
||||
record.service_tier,
|
||||
);
|
||||
if !record.rate_limit_headers.is_empty()
|
||||
&& let Some(object) = metadata.as_object_mut()
|
||||
{
|
||||
object.insert(
|
||||
"rateLimitHeaders".to_string(),
|
||||
serde_json::json!(record.rate_limit_headers),
|
||||
);
|
||||
}
|
||||
lift_success_fact_fields(&mut event, &mut metadata, record.usage);
|
||||
event.metadata = metadata;
|
||||
|
||||
if let Err(error) = record.recorder.record_usage(event).await {
|
||||
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record response event");
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
|
||||
let mut metadata = record.metadata;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, record.message);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
|
||||
event.provider = record.provider.map(ToString::to_string);
|
||||
event.account_id = record.account_id.map(ToString::to_string);
|
||||
event.route = Some(record.route.to_string());
|
||||
event.model = Some(record.model.to_string());
|
||||
event.status_code = Some(record.status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
|
||||
if let Err(error) = record.recorder.record_error(event).await {
|
||||
tracing::error!(
|
||||
account_id = record.account_id.unwrap_or(""),
|
||||
error = %error,
|
||||
"Failed to record dispatch error event"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn record_response_dispatch_error_event(
|
||||
@@ -177,16 +285,18 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
reasoning_effort_from_request(record.request),
|
||||
record.request.service_tier(),
|
||||
);
|
||||
enrich_response_request_semantics(&mut metadata, record.request);
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
if let Err(error) = record.recorder.record_error(event).await {
|
||||
tracing::error!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
|
||||
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record upstream error event");
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
record: ResponseStreamFailureEventRecord<'_>,
|
||||
) {
|
||||
let event_status_code = status_code_for_stream_failure(record.failure);
|
||||
let event_status_code = record.status_code;
|
||||
let mut metadata = stream_failure_metadata(record.failure, None);
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
if record.transport == CodexBackendTransport::WebSocket {
|
||||
@@ -202,11 +312,13 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
diagnostics_status_code(record.diagnostics),
|
||||
);
|
||||
insert_response_request_summary_object(object, record.request, record.transport);
|
||||
object.insert("requestBody".to_string(), json!(record.request));
|
||||
object.insert(
|
||||
"responseBody".to_string(),
|
||||
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
|
||||
);
|
||||
if record.recorder.captures_body() {
|
||||
object.insert("requestBody".to_string(), json!(record.request));
|
||||
object.insert(
|
||||
"responseBody".to_string(),
|
||||
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
|
||||
);
|
||||
}
|
||||
insert_upstream_diagnostics_metadata(object, record.diagnostics);
|
||||
insert_response_trace_metadata_object(
|
||||
object,
|
||||
@@ -224,10 +336,8 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
reasoning_effort_from_request(record.request),
|
||||
record.request.service_tier(),
|
||||
);
|
||||
let mut event = OpsErrorLog::new(
|
||||
response_event_kind(record.route),
|
||||
"v1 responses stream failed",
|
||||
);
|
||||
enrich_response_request_semantics(&mut metadata, record.request);
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, "v1 responses stream failed");
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.request.client_api_key_id.clone();
|
||||
event.provider = Some("openai".to_string());
|
||||
@@ -236,17 +346,215 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
event.model = Some(record.model.to_string());
|
||||
event.status_code = Some(event_status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
if !record.rate_limit_headers.is_empty() {
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert(
|
||||
"rateLimitHeaders".to_string(),
|
||||
serde_json::json!(record.rate_limit_headers),
|
||||
);
|
||||
}
|
||||
if !record.rate_limit_headers.is_empty()
|
||||
&& let Some(object) = metadata.as_object_mut()
|
||||
{
|
||||
object.insert(
|
||||
"rateLimitHeaders".to_string(),
|
||||
serde_json::json!(record.rate_limit_headers),
|
||||
);
|
||||
}
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
if let Err(error) = record.recorder.record_error(event).await {
|
||||
tracing::error!(account_id = %record.account_id, error = %error, "failed to record prefetched stream error event");
|
||||
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record prefetched stream error event");
|
||||
}
|
||||
}
|
||||
|
||||
fn lift_error_fact_fields(event: &mut OpsErrorLog) {
|
||||
event.client_status_code = event
|
||||
.client_status_code
|
||||
.or_else(|| metadata_i64(&event.metadata, &["clientStatusCode", "client_status_code"]));
|
||||
event.upstream_status_code = event
|
||||
.upstream_status_code
|
||||
.or_else(|| metadata_i64(&event.metadata, &["upstreamStatusCode", "upstreamStatus"]));
|
||||
event.transport = event
|
||||
.transport
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["transport"]));
|
||||
event.attempt_index = event
|
||||
.attempt_index
|
||||
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"]));
|
||||
event.failure_class = event
|
||||
.failure_class
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["failureClass", "failure_class"]));
|
||||
event.response_id = event
|
||||
.response_id
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["responseId", "response_id"]));
|
||||
event.upstream_request_id = event.upstream_request_id.take().or_else(|| {
|
||||
metadata_string_any(
|
||||
&event.metadata,
|
||||
&[
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
],
|
||||
)
|
||||
});
|
||||
event.service_tier = event
|
||||
.service_tier
|
||||
.take()
|
||||
.or_else(|| metadata_service_tier(&event.metadata).map(ToString::to_string));
|
||||
|
||||
let Some(metadata) = event.metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"clientStatusCode",
|
||||
"client_status_code",
|
||||
"upstreamStatusCode",
|
||||
"upstreamStatus",
|
||||
"transport",
|
||||
"attemptIndex",
|
||||
"attempt_index",
|
||||
"failureClass",
|
||||
"failure_class",
|
||||
"responseId",
|
||||
"response_id",
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
"serviceTier",
|
||||
] {
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn lift_success_fact_fields(
|
||||
event: &mut UsageRecord,
|
||||
metadata: &mut Value,
|
||||
usage: Option<TokenUsage>,
|
||||
) {
|
||||
event.transport = metadata_string(metadata, "transport");
|
||||
event.attempt_index = metadata_nonnegative_i64(metadata, "attemptIndex");
|
||||
event.response_id = metadata_string(metadata, "responseId");
|
||||
event.upstream_request_id = metadata_string(metadata, "upstreamRequestId")
|
||||
.or_else(|| metadata_string(metadata, "openaiRequestId"));
|
||||
event.first_token_ms = metadata_nonnegative_i64(metadata, "firstTokenMs");
|
||||
if let Some(usage) = usage {
|
||||
event.input_tokens = i64::try_from(usage.input_tokens).ok();
|
||||
event.output_tokens = i64::try_from(usage.output_tokens).ok();
|
||||
event.cached_tokens = i64::try_from(usage.cached_tokens).ok();
|
||||
event.cache_write_tokens = i64::try_from(usage.cache_write_tokens).ok();
|
||||
event.reasoning_tokens = i64::try_from(usage.reasoning_tokens).ok();
|
||||
}
|
||||
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"usage",
|
||||
"route",
|
||||
"apiKind",
|
||||
"requestedModel",
|
||||
"upstreamModel",
|
||||
"serviceTier",
|
||||
"statusCode",
|
||||
"transport",
|
||||
"attemptIndex",
|
||||
"responseId",
|
||||
"upstreamRequestId",
|
||||
"openaiRequestId",
|
||||
"firstTokenMs",
|
||||
] {
|
||||
object.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata_string(metadata: &Value, field: &str) -> Option<String> {
|
||||
metadata
|
||||
.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn metadata_nonnegative_i64(metadata: &Value, field: &str) -> Option<i64> {
|
||||
metadata
|
||||
.get(field)
|
||||
.and_then(Value::as_i64)
|
||||
.filter(|value| *value >= 0)
|
||||
}
|
||||
|
||||
fn normalized_string(value: Option<&str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
object
|
||||
.entry("route".to_string())
|
||||
.or_insert_with(|| Value::String(route.to_string()));
|
||||
object
|
||||
.entry("apiKind".to_string())
|
||||
.or_insert_with(|| Value::String(RESPONSES_API_KIND.to_string()));
|
||||
}
|
||||
|
||||
fn enrich_usage_record_identity(
|
||||
metadata: &mut Value,
|
||||
requested_model: Option<&str>,
|
||||
upstream_model: &str,
|
||||
client_ip: Option<&str>,
|
||||
client_user_agent: Option<&str>,
|
||||
reasoning_effort: Option<&str>,
|
||||
service_tier: Option<&str>,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let upstream_model = upstream_model.trim();
|
||||
let requested_model = requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(upstream_model);
|
||||
object.insert(
|
||||
"requestedModel".to_string(),
|
||||
Value::String(requested_model.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"upstreamModel".to_string(),
|
||||
Value::String(upstream_model.to_string()),
|
||||
);
|
||||
insert_trimmed_string(object, "clientIp", client_ip);
|
||||
insert_trimmed_string(object, "userAgent", client_user_agent);
|
||||
insert_trimmed_string(object, "reasoningEffort", reasoning_effort);
|
||||
insert_trimmed_string(object, "serviceTier", service_tier);
|
||||
}
|
||||
|
||||
pub(super) fn enrich_response_request_semantics(
|
||||
metadata: &mut Value,
|
||||
request: &CodexResponsesRequest,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let semantics = request.semantics();
|
||||
object.insert("compact".to_string(), Value::Bool(semantics.compact));
|
||||
insert_trimmed_string(object, "requestKind", semantics.request_kind.as_deref());
|
||||
insert_trimmed_string(object, "subagentKind", semantics.subagent_kind.as_deref());
|
||||
insert_trimmed_string(object, "reasoningPreset", semantics.reasoning_preset);
|
||||
}
|
||||
|
||||
pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
|
||||
request
|
||||
.reasoning()?
|
||||
.get("effort")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn insert_trimmed_string(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
object.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,8 +568,16 @@ fn ensure_stream_metadata_flag(metadata: &mut Value) {
|
||||
.or_insert(Value::Bool(true));
|
||||
}
|
||||
|
||||
fn insert_optional_i64(object: &mut serde_json::Map<String, Value>, key: &str, value: Option<i64>) {
|
||||
if let Some(value) = value {
|
||||
object
|
||||
.entry(key.to_string())
|
||||
.or_insert_with(|| Value::Number(value.into()));
|
||||
}
|
||||
}
|
||||
|
||||
fn enrich_live_response_stream_metadata(
|
||||
context: &LiveResponseStreamContext,
|
||||
context: &StreamContext<'_>,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
metadata: &mut Value,
|
||||
status_code: i64,
|
||||
@@ -283,21 +599,58 @@ fn enrich_live_response_stream_metadata(
|
||||
.entry("websocketPool".to_string())
|
||||
.or_insert_with(|| decision.metadata_value());
|
||||
}
|
||||
if let Some(decision) = context.transport_metrics.decision {
|
||||
object
|
||||
.entry("transportDecision".to_string())
|
||||
.or_insert_with(|| Value::String(decision.as_str().to_string()));
|
||||
}
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"wsConnectMs",
|
||||
context.transport_metrics.ws_connect_ms,
|
||||
);
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"transportDecisionWaitMs",
|
||||
context.transport_metrics.transport_decision_wait_ms,
|
||||
);
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"upstreamHeadersMs",
|
||||
context.transport_metrics.upstream_headers_ms,
|
||||
);
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"firstEventMs",
|
||||
context.transport_metrics.first_event_ms,
|
||||
);
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"openaiProcessingMs",
|
||||
openai_processing_ms(context.response_metadata),
|
||||
);
|
||||
if let Some(http_version) = context.transport_metrics.http_version.as_deref() {
|
||||
object
|
||||
.entry("httpVersion".to_string())
|
||||
.or_insert_with(|| Value::String(http_version.to_string()));
|
||||
}
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
status_code,
|
||||
200,
|
||||
diagnostics_status_code(&context.diagnostics),
|
||||
diagnostics_status_code(context.diagnostics),
|
||||
);
|
||||
insert_response_request_summary_object(object, &context.request, context.transport);
|
||||
insert_upstream_diagnostics_metadata(object, &context.diagnostics);
|
||||
insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt));
|
||||
object
|
||||
.entry("requestBody".to_string())
|
||||
.or_insert_with(|| serde_json::json!(context.request));
|
||||
object
|
||||
.entry("responseBody".to_string())
|
||||
.or_insert_with(|| Value::String(body.to_string()));
|
||||
insert_response_request_summary_object(object, context.request, context.transport);
|
||||
insert_upstream_diagnostics_metadata(object, context.diagnostics);
|
||||
insert_response_trace_metadata_object(object, context.attempts, Some(context.attempt));
|
||||
if context.recorder.captures_body() {
|
||||
object
|
||||
.entry("requestBody".to_string())
|
||||
.or_insert_with(|| serde_json::json!(context.request));
|
||||
object
|
||||
.entry("responseBody".to_string())
|
||||
.or_insert_with(|| Value::String(body.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn insert_response_upstream_diagnostics(
|
||||
@@ -409,10 +762,15 @@ fn response_request_summary(
|
||||
let body = request.body();
|
||||
let input = body.get("input");
|
||||
let tools = body.get("tools");
|
||||
let semantics = request.semantics();
|
||||
json!({
|
||||
"model": request.model(),
|
||||
"stream": request.stream(),
|
||||
"store": request.store(),
|
||||
"compact": semantics.compact,
|
||||
"requestKind": semantics.request_kind,
|
||||
"subagentKind": semantics.subagent_kind,
|
||||
"reasoningPreset": semantics.reasoning_preset,
|
||||
"transport": backend_transport_name(transport),
|
||||
"inputType": json_value_kind(input),
|
||||
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
|
||||
@@ -463,20 +821,22 @@ fn upstream_failure_client_status_code(upstream_status: i64) -> i64 {
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn record_live_response_stream_event(
|
||||
context: &LiveResponseStreamContext,
|
||||
status_code: i64,
|
||||
level: UsageRecordLevel,
|
||||
message: &str,
|
||||
mut metadata: Value,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
body: &str,
|
||||
) {
|
||||
pub(super) async fn record_live_response_stream_event(record: LiveResponseStreamEventRecord<'_>) {
|
||||
let LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
level,
|
||||
message,
|
||||
usage,
|
||||
mut metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
} = record;
|
||||
let effective_model = context
|
||||
.response_metadata
|
||||
.effective_model
|
||||
.as_deref()
|
||||
.unwrap_or(&context.display_model);
|
||||
.unwrap_or(context.display_model);
|
||||
ensure_stream_metadata_flag(&mut metadata);
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert(
|
||||
@@ -496,7 +856,7 @@ pub(super) async fn record_live_response_stream_event(
|
||||
Value::Bool(context.response_metadata.reasoning_included),
|
||||
);
|
||||
}
|
||||
enrich_event_route_metadata(&mut metadata, &context.route);
|
||||
enrich_event_route_metadata(&mut metadata, context.route);
|
||||
enrich_live_response_stream_metadata(
|
||||
context,
|
||||
rate_limit_headers,
|
||||
@@ -504,52 +864,56 @@ pub(super) async fn record_live_response_stream_event(
|
||||
status_code,
|
||||
body,
|
||||
);
|
||||
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
|
||||
if !is_success_usage_event(level, status_code) {
|
||||
ensure_stream_failure_class(&mut metadata);
|
||||
}
|
||||
log_live_response_stream_finalized(context, status_code, level, message, &metadata);
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
Some(&context.requested_model),
|
||||
Some(context.requested_model),
|
||||
effective_model,
|
||||
context.client_ip.as_deref(),
|
||||
context.request.client_ip.as_deref(),
|
||||
context.request.client_user_agent.as_deref(),
|
||||
reasoning_effort_from_request(&context.request),
|
||||
reasoning_effort_from_request(context.request),
|
||||
context.request.service_tier(),
|
||||
);
|
||||
enrich_response_request_semantics(&mut metadata, context.request);
|
||||
|
||||
if is_success_usage_event(level, status_code) {
|
||||
crate::telemetry::recorder::record_response_event(
|
||||
crate::telemetry::usage::types::ResponseUsageRecord {
|
||||
recorder: &context.recorder,
|
||||
request_id: &context.request_id,
|
||||
client_api_key_id: context.request.client_api_key_id.as_deref(),
|
||||
account_id: &context.account_id,
|
||||
route: &context.route,
|
||||
model: effective_model,
|
||||
requested_model: Some(&context.requested_model),
|
||||
client_ip: context.client_ip.as_deref(),
|
||||
client_user_agent: context.request.client_user_agent.as_deref(),
|
||||
reasoning_effort: reasoning_effort_from_request(&context.request),
|
||||
service_tier: context.request.service_tier(),
|
||||
started_at: context.started_at,
|
||||
status_code,
|
||||
message,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
},
|
||||
)
|
||||
record_response_event(ResponseUsageEventRecord {
|
||||
recorder: context.recorder,
|
||||
request_id: context.request_id,
|
||||
client_api_key_id: context.request.client_api_key_id.as_deref(),
|
||||
account_id: context.account_id,
|
||||
route: context.route,
|
||||
model: effective_model,
|
||||
requested_model: Some(context.requested_model),
|
||||
client_ip: context.request.client_ip.as_deref(),
|
||||
client_user_agent: context.request.client_user_agent.as_deref(),
|
||||
reasoning_effort: reasoning_effort_from_request(context.request),
|
||||
service_tier: context.request.service_tier(),
|
||||
started_at: context.started_at,
|
||||
status_code,
|
||||
message,
|
||||
usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
})
|
||||
.await;
|
||||
} else {
|
||||
let mut event = OpsErrorLog::new(response_event_kind(&context.route), message);
|
||||
event.request_id = Some(context.request_id.clone());
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, message);
|
||||
event.request_id = Some(context.request_id.to_string());
|
||||
event.client_api_key_id = context.request.client_api_key_id.clone();
|
||||
event.provider = Some("openai".to_string());
|
||||
event.account_id = Some(context.account_id.clone());
|
||||
event.route = Some(context.route.clone());
|
||||
event.account_id = Some(context.account_id.to_string());
|
||||
event.route = Some(context.route.to_string());
|
||||
event.model = Some(effective_model.to_string());
|
||||
event.status_code = Some(status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
if let Err(error) = context.recorder.record_error(event).await {
|
||||
tracing::error!(account_id = %context.account_id, error = %error, "failed to record live response stream error event");
|
||||
tracing::error!(account_id = %context.account_id, error = %error, "Failed to record live response stream error event");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -558,24 +922,31 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
|
||||
level != UsageRecordLevel::Error && status_code < 400
|
||||
}
|
||||
|
||||
fn ensure_stream_failure_class(metadata: &mut Value) {
|
||||
let failure_class = ["failureClass", "upstreamCode", "terminal"]
|
||||
.into_iter()
|
||||
.find_map(|field| metadata_string_field(metadata, field))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string);
|
||||
let (Some(failure_class), Some(object)) = (failure_class, metadata.as_object_mut()) else {
|
||||
return;
|
||||
};
|
||||
object.insert("failureClass".to_string(), Value::String(failure_class));
|
||||
}
|
||||
|
||||
fn log_live_response_stream_finalized(
|
||||
context: &LiveResponseStreamContext,
|
||||
context: &StreamContext<'_>,
|
||||
status_code: i64,
|
||||
level: UsageRecordLevel,
|
||||
message: &str,
|
||||
metadata: &Value,
|
||||
body: &str,
|
||||
) {
|
||||
let response_id = metadata_string_field(metadata, "responseId")
|
||||
.map(ToString::to_string)
|
||||
.or_else(|| latest_response_id(body));
|
||||
let response_id = metadata_string_field(metadata, "responseId").map(ToString::to_string);
|
||||
let first_token_ms = metadata.get("firstTokenMs").and_then(Value::as_i64);
|
||||
let websocket_pool_kind = context
|
||||
.websocket_pool_decision
|
||||
.map(WebSocketPoolDecision::kind);
|
||||
let websocket_pool_reason = context
|
||||
.websocket_pool_decision
|
||||
.and_then(WebSocketPoolDecision::reason);
|
||||
let completed = metadata.get("completed").and_then(Value::as_bool);
|
||||
let failed = metadata.get("failed").and_then(Value::as_bool);
|
||||
let upstream_code = metadata_string_field(metadata, "upstreamCode");
|
||||
@@ -596,7 +967,6 @@ fn log_live_response_stream_finalized(
|
||||
event_message = %message,
|
||||
transport = %backend_transport_name(context.transport),
|
||||
websocket_pool_kind = ?websocket_pool_kind,
|
||||
websocket_pool_reason = ?websocket_pool_reason,
|
||||
response_id = response_id.as_deref().unwrap_or(""),
|
||||
first_token_ms = ?first_token_ms,
|
||||
latency_ms = elapsed_millis_i64(context.started_at),
|
||||
@@ -606,7 +976,7 @@ fn log_live_response_stream_finalized(
|
||||
failure_class = ?failure_class,
|
||||
failure_source = ?failure_source,
|
||||
failure_detail = ?failure_detail,
|
||||
"live response stream finalized"
|
||||
"Live response stream finalized"
|
||||
);
|
||||
};
|
||||
}
|
||||
@@ -631,35 +1001,40 @@ fn metadata_string_field<'a>(metadata: &'a Value, field: &str) -> Option<&'a str
|
||||
metadata.get(field).and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub(super) async fn live_response_rate_limit_headers(
|
||||
context: &LiveResponseStreamContext,
|
||||
) -> Vec<(String, String)> {
|
||||
let mut headers = context.rate_limit_headers.clone();
|
||||
if let Some(updates) = &context.rate_limit_header_updates {
|
||||
headers.extend(updates.lock().await.iter().cloned());
|
||||
pub(super) fn insert_output_timing_ms(
|
||||
metadata: &mut Value,
|
||||
first_token_ms: Option<i64>,
|
||||
first_reasoning_ms: Option<i64>,
|
||||
first_text_ms: Option<i64>,
|
||||
) {
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
insert_optional_i64(object, "firstTokenMs", first_token_ms);
|
||||
insert_optional_i64(object, "firstReasoningMs", first_reasoning_ms);
|
||||
insert_optional_i64(object, "firstTextMs", first_text_ms);
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
pub(super) async fn live_response_turn_state(
|
||||
context: &LiveResponseStreamContext,
|
||||
) -> Option<String> {
|
||||
if let Some(update) = &context.turn_state_update {
|
||||
return update.lock().await.clone();
|
||||
}
|
||||
context.turn_state.clone()
|
||||
}
|
||||
|
||||
pub(super) fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option<i64>) {
|
||||
let Some(first_token_ms) = first_token_ms else {
|
||||
pub(super) fn insert_openai_processing_ms(
|
||||
metadata: &mut Value,
|
||||
response_metadata: &CodexResponseMetadata,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert(
|
||||
"firstTokenMs".to_string(),
|
||||
Value::Number(first_token_ms.into()),
|
||||
);
|
||||
}
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"openaiProcessingMs",
|
||||
openai_processing_ms(response_metadata),
|
||||
);
|
||||
}
|
||||
|
||||
fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option<i64> {
|
||||
response_metadata
|
||||
.client_headers
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms"))
|
||||
.and_then(|(_, value)| value.parse::<i64>().ok())
|
||||
.filter(|value| *value >= 0)
|
||||
}
|
||||
|
||||
pub(super) fn insert_websocket_pool_decision(
|
||||
@@ -0,0 +1,415 @@
|
||||
//! 请求终态遥测的唯一 feature owner。
|
||||
|
||||
mod events;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
controllers::telemetry::events::{
|
||||
LiveResponseStreamEventRecord, ResponseUpstreamErrorEventRecord,
|
||||
ResponseUsageEventRecord, enrich_response_request_semantics,
|
||||
insert_openai_processing_ms, insert_output_timing_ms, insert_response_status_metadata,
|
||||
insert_response_trace_metadata, insert_response_upstream_diagnostics,
|
||||
insert_websocket_pool_decision, reasoning_effort_from_request,
|
||||
record_live_response_stream_event, record_response_event,
|
||||
record_response_upstream_error_event,
|
||||
},
|
||||
controllers::{
|
||||
AttemptUpstreamErrorContext, CompleteExit, DispatchErrorObservation,
|
||||
PrefetchedStreamFailureObservation,
|
||||
},
|
||||
errors::{ResponseDispatchError, backend_transport_name},
|
||||
failure::sse::{
|
||||
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, stream_failure_metadata,
|
||||
stream_failure_source, synthetic_stream_disconnected_detail,
|
||||
},
|
||||
lifecycle::{
|
||||
stream::{StreamSummary, StreamTerminal},
|
||||
trace::ResponseDispatchAttempt,
|
||||
},
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
telemetry::{recorder::Recorder, usage::types::UsageRecordLevel},
|
||||
upstream::openai::{
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
transport::{
|
||||
CodexBackendTransport, CodexResponseMetadata, CodexTransportMetrics,
|
||||
CodexUpstreamDiagnostics, WebSocketPoolDecision,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
pub(super) struct TelemetryController;
|
||||
|
||||
pub(super) struct StreamContext<'a> {
|
||||
pub recorder: &'a Recorder,
|
||||
pub account_id: &'a str,
|
||||
pub request_id: &'a str,
|
||||
pub route: &'a str,
|
||||
pub display_model: &'a str,
|
||||
pub requested_model: &'a str,
|
||||
pub request: &'a CodexResponsesRequest,
|
||||
pub transport: CodexBackendTransport,
|
||||
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
pub diagnostics: &'a CodexUpstreamDiagnostics,
|
||||
pub response_metadata: &'a CodexResponseMetadata,
|
||||
pub transport_metrics: &'a CodexTransportMetrics,
|
||||
pub attempt: &'a ResponseDispatchAttempt,
|
||||
pub attempts: &'a [ResponseDispatchAttempt],
|
||||
pub started_at: std::time::Instant,
|
||||
}
|
||||
|
||||
pub(super) struct StreamExit<'a> {
|
||||
pub context: StreamContext<'a>,
|
||||
pub summary: &'a StreamSummary,
|
||||
pub rate_limit_headers: &'a [(String, String)],
|
||||
pub failure_status: Option<i64>,
|
||||
pub body: &'a str,
|
||||
}
|
||||
|
||||
impl TelemetryController {
|
||||
pub(super) async fn observe_dispatch_error(
|
||||
recorder: &Recorder,
|
||||
observation: DispatchErrorObservation<'_>,
|
||||
error: &ResponseDispatchError,
|
||||
) {
|
||||
events::record_response_dispatch_error_event(events::ResponseDispatchErrorEventRecord {
|
||||
recorder,
|
||||
request_id: observation.request_id,
|
||||
client_api_key_id: observation.client_api_key_id,
|
||||
account_id: observation.account_id,
|
||||
route: observation.route,
|
||||
model: observation.model,
|
||||
started_at: observation.started_at,
|
||||
stream: observation.stream,
|
||||
compact: observation.compact,
|
||||
transport: observation.transport,
|
||||
error,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn observe_prefetched_stream_failure(
|
||||
recorder: &Recorder,
|
||||
observation: PrefetchedStreamFailureObservation<'_>,
|
||||
) {
|
||||
events::record_prefetched_response_stream_failure_event(
|
||||
events::ResponseStreamFailureEventRecord {
|
||||
recorder,
|
||||
request_id: observation.request_id,
|
||||
account_id: observation.account_id,
|
||||
route: observation.route,
|
||||
model: observation.model,
|
||||
requested_model: observation.requested_model,
|
||||
started_at: observation.started_at,
|
||||
transport: observation.transport,
|
||||
request: observation.request,
|
||||
failure: observation.failure,
|
||||
status_code: i64::from(observation.error.client_http_status_code()),
|
||||
diagnostics: observation.diagnostics,
|
||||
rate_limit_headers: observation.rate_limit_headers,
|
||||
prefetched: observation.prefetched,
|
||||
trace: observation.trace,
|
||||
attempt: observation.attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn observe_upstream_error(
|
||||
recorder: &Recorder,
|
||||
context: &AttemptUpstreamErrorContext<'_>,
|
||||
) {
|
||||
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
|
||||
recorder,
|
||||
request_id: context.request_id,
|
||||
account_id: &context.account.id,
|
||||
account_email: context.account.email.as_deref(),
|
||||
route: context.route,
|
||||
model: context.model,
|
||||
started_at: context.started_at,
|
||||
stream: context.stream,
|
||||
transport: context.transport,
|
||||
request: context.request,
|
||||
error: context.error,
|
||||
trace: context.trace,
|
||||
attempt: Some(context.attempt),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_complete(exit: &CompleteExit<'_>) {
|
||||
let effective_model = exit
|
||||
.response
|
||||
.response_metadata
|
||||
.effective_model
|
||||
.as_deref()
|
||||
.unwrap_or(exit.display_model);
|
||||
let first_token_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_token_ms,
|
||||
);
|
||||
let first_reasoning_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_reasoning_ms,
|
||||
);
|
||||
let first_text_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_text_ms,
|
||||
);
|
||||
let mut metadata = serde_json::json!({
|
||||
"responseId": exit.response_id,
|
||||
"stream": false,
|
||||
"completed": exit.completed,
|
||||
"incomplete": !exit.completed,
|
||||
"transport": backend_transport_name(exit.response.transport),
|
||||
"usage": exit.response.usage,
|
||||
"effectiveModel": effective_model,
|
||||
"modelsEtag": exit.response.response_metadata.models_etag.as_deref(),
|
||||
"reasoningIncluded": exit.response.response_metadata.reasoning_included,
|
||||
"transportDecision": exit.response.transport_metrics.decision.map(|decision| decision.as_str()),
|
||||
"wsConnectMs": exit.response.transport_metrics.ws_connect_ms,
|
||||
"transportDecisionWaitMs": exit.response.transport_metrics.transport_decision_wait_ms,
|
||||
"upstreamHeadersMs": exit.response.transport_metrics.upstream_headers_ms,
|
||||
"firstEventMs": exit.response.transport_metrics.first_event_ms,
|
||||
"httpVersion": exit.response.transport_metrics.http_version.as_deref(),
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
first_token_ms,
|
||||
first_reasoning_ms,
|
||||
first_text_ms,
|
||||
);
|
||||
insert_openai_processing_ms(&mut metadata, &exit.response.response_metadata);
|
||||
insert_response_status_metadata(
|
||||
&mut metadata,
|
||||
200,
|
||||
200,
|
||||
exit.response.diagnostics.status_code.map(i64::from),
|
||||
);
|
||||
insert_response_upstream_diagnostics(&mut metadata, &exit.response.diagnostics);
|
||||
insert_response_trace_metadata(&mut metadata, exit.trace, Some(exit.attempt));
|
||||
insert_websocket_pool_decision(&mut metadata, exit.response.websocket_pool_decision);
|
||||
enrich_response_request_semantics(&mut metadata, exit.request);
|
||||
record_response_event(ResponseUsageEventRecord {
|
||||
recorder: exit.recorder,
|
||||
request_id: exit.request_id,
|
||||
client_api_key_id: exit.request.client_api_key_id.as_deref(),
|
||||
account_id: exit.account_id,
|
||||
route: exit.route,
|
||||
model: effective_model,
|
||||
requested_model: Some(exit.requested_model),
|
||||
client_ip: exit.request.client_ip.as_deref(),
|
||||
client_user_agent: exit.request.client_user_agent.as_deref(),
|
||||
reasoning_effort: reasoning_effort_from_request(exit.request),
|
||||
service_tier: exit.request.service_tier(),
|
||||
started_at: exit.started_at,
|
||||
status_code: 200,
|
||||
message: if exit.completed {
|
||||
"v1 responses completed"
|
||||
} else {
|
||||
"v1 responses incomplete"
|
||||
},
|
||||
usage: exit.response.usage,
|
||||
metadata,
|
||||
rate_limit_headers: &exit.response.rate_limit_headers,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
|
||||
let context = &exit.context;
|
||||
let summary = exit.summary;
|
||||
let rate_limit_headers = exit.rate_limit_headers;
|
||||
let body = exit.body;
|
||||
match &summary.terminal {
|
||||
StreamTerminal::Completed { response } => {
|
||||
record_success(context, summary, response, true, rate_limit_headers, body).await;
|
||||
}
|
||||
StreamTerminal::Incomplete { response } => {
|
||||
record_success(context, summary, response, false, rate_limit_headers, body).await;
|
||||
}
|
||||
StreamTerminal::Failed { failure } => {
|
||||
record_failure(
|
||||
context,
|
||||
summary,
|
||||
failure,
|
||||
exit.failure_status.unwrap_or(502),
|
||||
rate_limit_headers,
|
||||
body,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
terminal @ (StreamTerminal::UpstreamClosed
|
||||
| StreamTerminal::UpstreamError { .. }
|
||||
| StreamTerminal::ProtocolError { .. }
|
||||
| StreamTerminal::CaptureLimitExceeded) => {
|
||||
let failure = synthetic_failure(terminal);
|
||||
record_failure(context, summary, &failure, 502, rate_limit_headers, body).await;
|
||||
}
|
||||
terminal @ (StreamTerminal::Cancelled
|
||||
| StreamTerminal::DownstreamClosed
|
||||
| StreamTerminal::Shutdown) => {
|
||||
let status_code = if matches!(terminal, StreamTerminal::Shutdown) {
|
||||
503
|
||||
} else {
|
||||
499
|
||||
};
|
||||
let mut metadata = serde_json::json!({
|
||||
"stream": true,
|
||||
"cancelled": true,
|
||||
"terminal": terminal.name(),
|
||||
"responseId": summary.last_response_id,
|
||||
"firstEventMs": summary.first_event_ms,
|
||||
"usage": summary.usage,
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
level: UsageRecordLevel::Warn,
|
||||
message: "v1 responses stream cancelled",
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn request_timing_ms(
|
||||
request_started_at: Instant,
|
||||
attempt_started_at: Instant,
|
||||
attempt_first_token_ms: Option<i64>,
|
||||
) -> Option<i64> {
|
||||
let attempt_first_token_ms = u64::try_from(attempt_first_token_ms?).ok()?;
|
||||
let output_at =
|
||||
attempt_started_at.checked_add(Duration::from_millis(attempt_first_token_ms))?;
|
||||
i64::try_from(
|
||||
output_at
|
||||
.saturating_duration_since(request_started_at)
|
||||
.as_millis(),
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
async fn record_success(
|
||||
context: &StreamContext<'_>,
|
||||
summary: &StreamSummary,
|
||||
response: &Value,
|
||||
completed: bool,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
body: &str,
|
||||
) {
|
||||
let mut metadata = serde_json::json!({
|
||||
"stream": true,
|
||||
"completed": completed,
|
||||
"incomplete": !completed,
|
||||
"responseId": response.get("id").and_then(Value::as_str),
|
||||
"firstEventMs": summary.first_event_ms,
|
||||
"usage": summary.usage,
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code: 200,
|
||||
level: UsageRecordLevel::Info,
|
||||
message: if completed {
|
||||
"v1 responses stream completed"
|
||||
} else {
|
||||
"v1 responses stream incomplete"
|
||||
},
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
fn synthetic_failure(terminal: &StreamTerminal) -> ResponsesSseFailure {
|
||||
let detail = terminal.synthetic_failure_detail().unwrap_or_default();
|
||||
ResponsesSseFailure {
|
||||
event: "response.failed".to_string(),
|
||||
message: if detail.trim().is_empty() {
|
||||
STREAM_DISCONNECTED_MESSAGE.to_string()
|
||||
} else {
|
||||
format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}")
|
||||
},
|
||||
upstream_code: Some(STREAM_DISCONNECTED_CODE.to_string()),
|
||||
upstream_type: Some("server_error".to_string()),
|
||||
explicit_status_code: None,
|
||||
retry_after_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_failure(
|
||||
context: &StreamContext<'_>,
|
||||
summary: &StreamSummary,
|
||||
failure: &ResponsesSseFailure,
|
||||
status_code: i64,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
body: &str,
|
||||
) {
|
||||
let failure_source = stream_failure_source(failure);
|
||||
let failure_detail = synthetic_stream_disconnected_detail(failure);
|
||||
tracing::warn!(
|
||||
account_id = %context.account_id,
|
||||
request_id = %context.request_id,
|
||||
response_id = summary.last_response_id.as_deref().unwrap_or(""),
|
||||
transport = %backend_transport_name(context.transport),
|
||||
websocket_pool_kind = ?context.websocket_pool_decision.map(WebSocketPoolDecision::kind),
|
||||
first_token_ms = ?summary.first_token_ms,
|
||||
first_event_ms = summary.first_event_ms,
|
||||
latency_ms = elapsed_millis_i64(context.started_at),
|
||||
event = %failure.event,
|
||||
code = ?failure.upstream_code.as_deref(),
|
||||
failure_source,
|
||||
failure_detail = ?failure_detail.as_deref(),
|
||||
"Live upstream stream ended with response.failed"
|
||||
);
|
||||
let mut metadata = stream_failure_metadata(failure, summary.usage);
|
||||
metadata["responseId"] = summary
|
||||
.last_response_id
|
||||
.as_ref()
|
||||
.map_or(Value::Null, |response_id| {
|
||||
Value::String(response_id.clone())
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
metadata["firstEventMs"] = Value::Number(summary.first_event_ms.into());
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
level: UsageRecordLevel::Error,
|
||||
message: "v1 responses stream failed",
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
//! 配额响应头与 token 用量写入的唯一 feature owner。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::lifecycle::contract::{
|
||||
AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
|
||||
},
|
||||
fleet::{
|
||||
pool::AccountPoolService,
|
||||
quota::{
|
||||
QuotaCreditsObservation, QuotaLimitObservation, QuotaObservation,
|
||||
QuotaWindowObservation,
|
||||
},
|
||||
usage::ResponseUsage,
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::{
|
||||
events::{RateLimitWindow, TokenUsage, parse_rate_limit_headers},
|
||||
responses::CodexResponsesRequest,
|
||||
},
|
||||
transport::{CodexRateLimitHeaderUpdates, CodexTurnStateUpdate},
|
||||
},
|
||||
};
|
||||
|
||||
pub(super) struct UsageScope {
|
||||
image_generation_requested: bool,
|
||||
}
|
||||
|
||||
pub(super) struct StreamUsageExit {
|
||||
pub(super) rate_limit_headers: Vec<(String, String)>,
|
||||
pub(super) turn_state: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) struct StreamExit<'a> {
|
||||
pub account_pool: &'a AccountPoolService,
|
||||
pub account_id: &'a str,
|
||||
pub account_plan_type: Option<&'a str>,
|
||||
pub rate_limit_headers: &'a [(String, String)],
|
||||
pub rate_limit_header_updates: Option<&'a CodexRateLimitHeaderUpdates>,
|
||||
pub turn_state_update: Option<&'a CodexTurnStateUpdate>,
|
||||
pub turn_state: Option<&'a str>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
pub(super) struct UsageController;
|
||||
|
||||
impl UsageController {
|
||||
pub(super) fn enter(request: &CodexResponsesRequest) -> UsageScope {
|
||||
UsageScope {
|
||||
image_generation_requested: request.expects_image_generation(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn observe_attempt(
|
||||
account_pool: &AccountPoolService,
|
||||
scope: &UsageScope,
|
||||
observation: &AttemptObservation,
|
||||
) {
|
||||
if !matches!(
|
||||
observation.kind,
|
||||
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
let Some(account_id) = observation
|
||||
.account
|
||||
.as_ref()
|
||||
.map(|account| account.id.as_str())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
account_pool
|
||||
.record_empty_response_attempt(account_id, scope.image_generation_requested)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_complete(
|
||||
account_pool: &AccountPoolService,
|
||||
account_id: &str,
|
||||
usage: Option<TokenUsage>,
|
||||
image_generation_requested: bool,
|
||||
) {
|
||||
if let Some(usage) = usage {
|
||||
account_pool
|
||||
.record_response_usage(
|
||||
account_id,
|
||||
response_usage(usage),
|
||||
image_generation_requested,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn sync_passive_quota(
|
||||
account_pool: &AccountPoolService,
|
||||
account_id: &str,
|
||||
plan_type: Option<&str>,
|
||||
headers: &[(String, String)],
|
||||
) {
|
||||
let Some(observation) = quota_observation(headers) else {
|
||||
return;
|
||||
};
|
||||
account_pool
|
||||
.sync_passive_quota_observation_for_account(account_id, plan_type, &observation)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn leave_stream(exit: StreamExit<'_>) -> StreamUsageExit {
|
||||
let mut rate_limit_headers = exit.rate_limit_headers.to_vec();
|
||||
if let Some(updates) = exit.rate_limit_header_updates {
|
||||
rate_limit_headers.extend(updates.lock().await.iter().cloned());
|
||||
}
|
||||
Self::sync_passive_quota(
|
||||
exit.account_pool,
|
||||
exit.account_id,
|
||||
exit.account_plan_type,
|
||||
&rate_limit_headers,
|
||||
)
|
||||
.await;
|
||||
let turn_state = if let Some(update) = exit.turn_state_update {
|
||||
update.lock().await.clone()
|
||||
} else {
|
||||
exit.turn_state.map(ToString::to_string)
|
||||
};
|
||||
if let Some(usage) = exit.usage {
|
||||
exit.account_pool
|
||||
.record_response_usage(exit.account_id, response_usage(usage), false)
|
||||
.await;
|
||||
}
|
||||
StreamUsageExit {
|
||||
rate_limit_headers,
|
||||
turn_state,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn response_usage(usage: TokenUsage) -> ResponseUsage {
|
||||
ResponseUsage {
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
reasoning_tokens: usage.reasoning_tokens,
|
||||
total_tokens: usage.total_tokens,
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_observation(headers: &[(String, String)]) -> Option<QuotaObservation> {
|
||||
let parsed = parse_rate_limit_headers(headers)?;
|
||||
Some(QuotaObservation {
|
||||
limits: parsed
|
||||
.limits
|
||||
.into_iter()
|
||||
.map(|(limit_id, details)| {
|
||||
(
|
||||
limit_id,
|
||||
QuotaLimitObservation {
|
||||
limit_id: details.limit_id,
|
||||
limit_name: details.limit_name,
|
||||
metered_feature: None,
|
||||
allowed: details.allowed,
|
||||
limit_reached: details.limit_reached,
|
||||
primary: details.primary.map(quota_window_observation),
|
||||
secondary: details.secondary.map(quota_window_observation),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
active_limit: parsed.active_limit,
|
||||
credits: parsed.credits.map(|credits| QuotaCreditsObservation {
|
||||
has_credits: credits.has_credits,
|
||||
unlimited: credits.unlimited,
|
||||
overage_limit_reached: false,
|
||||
balance: credits.balance.map(Value::String),
|
||||
}),
|
||||
spend_control: None,
|
||||
plan_type: parsed.plan_type,
|
||||
promo_message: parsed.promo_message,
|
||||
rate_limit_reached_type: parsed.rate_limit_reached_type,
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_window_observation(window: RateLimitWindow) -> QuotaWindowObservation {
|
||||
QuotaWindowObservation {
|
||||
used_percent: window.used_percent,
|
||||
window_minutes: window.window_minutes,
|
||||
reset_at: window.reset_at,
|
||||
used: None,
|
||||
limit: None,
|
||||
}
|
||||
}
|
||||
+49
-135
@@ -1,13 +1,12 @@
|
||||
//! Shared upstream error classification for dispatch routes.
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde_json::{json, Map, Value};
|
||||
use serde_json::{Map, Value, json};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
recovery::exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
|
||||
stream::sse_failure::{sse_failure_error_body, stream_failure_http_status},
|
||||
dispatch::failure::{
|
||||
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
|
||||
sse::sse_failure_error_body,
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::{responses::ResponsesSseFailure, sse::SseError},
|
||||
@@ -15,6 +14,41 @@ use crate::{
|
||||
},
|
||||
};
|
||||
|
||||
/// controller 已经完成分类、可直接交给 API 编码的客户端失败事实。
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ClientFailure {
|
||||
failure: ResponsesSseFailure,
|
||||
status_code: u16,
|
||||
expose_upstream: bool,
|
||||
}
|
||||
|
||||
impl ClientFailure {
|
||||
pub(in crate::dispatch) fn new(
|
||||
failure: ResponsesSseFailure,
|
||||
status_code: u16,
|
||||
expose_upstream: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
failure,
|
||||
status_code,
|
||||
expose_upstream,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn failure(&self) -> &ResponsesSseFailure {
|
||||
&self.failure
|
||||
}
|
||||
|
||||
pub(crate) fn status_code(&self) -> u16 {
|
||||
self.status_code
|
||||
}
|
||||
|
||||
/// 只有 controller 明确允许时,API 才透传上游 code/message。
|
||||
pub(crate) fn exposed_upstream(&self) -> Option<&ResponsesSseFailure> {
|
||||
self.expose_upstream.then_some(&self.failure)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct DispatchErrorMetadata {
|
||||
pub failure_class: DispatchFailureClass,
|
||||
@@ -188,77 +222,7 @@ pub(crate) fn upstream_error_diagnostics(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, .. } if status_code_is_rate_limited(status.as_u16())
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_retryable_upstream_5xx_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, .. }
|
||||
if status_code_is_transient_upstream(status.as_u16())
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_retryable_account_transport_error(error: &CodexClientError) -> bool {
|
||||
match error {
|
||||
CodexClientError::Http(error) => error.is_connect() || error.is_timeout(),
|
||||
CodexClientError::WebSocket(
|
||||
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::Transport(_)
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ClosedBeforeTerminal
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. }
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::InitialEventTimeout { .. },
|
||||
) => true,
|
||||
_ => is_retryable_upstream_5xx_error(error),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_quota_exhausted_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, .. } if status_code_is_quota_exhausted(status.as_u16())
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_model_unsupported_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if status.is_client_error()
|
||||
&& !matches!(status.as_u16(), 401 | 402 | 403 | 404 | 429)
|
||||
&& is_model_unsupported_signal(body)
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_history_recovery_upstream_error(error: &CodexClientError) -> bool {
|
||||
upstream_error_code(error).is_some_and(|code| is_history_recovery_code(&code))
|
||||
|| matches!(
|
||||
error,
|
||||
CodexClientError::WebSocket(
|
||||
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable { reason }
|
||||
) if !matches!(
|
||||
reason,
|
||||
crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn is_continuation_busy_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::WebSocket(
|
||||
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable {
|
||||
reason: crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy,
|
||||
}
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_body(error: &CodexClientError) -> String {
|
||||
fn upstream_error_body(error: &CodexClientError) -> String {
|
||||
match error {
|
||||
CodexClientError::Upstream { body, .. } => body.clone(),
|
||||
error => error.to_string(),
|
||||
@@ -274,52 +238,13 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_model_unsupported_signal(value: &str) -> bool {
|
||||
let value = value.to_ascii_lowercase();
|
||||
value.contains("model_not_supported")
|
||||
|| value.contains("model_not_available")
|
||||
|| (value.contains("model")
|
||||
&& (value.contains("not supported")
|
||||
|| value.contains("not available")
|
||||
|| value.contains("not_supported")
|
||||
|| value.contains("not_available")))
|
||||
}
|
||||
|
||||
pub(crate) fn is_history_recovery_code(code: &str) -> bool {
|
||||
matches!(
|
||||
code,
|
||||
"previous_response_not_found"
|
||||
| "invalid_encrypted_content"
|
||||
| "missing_tool_output"
|
||||
| "no_tool_output"
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_code(error: &CodexClientError) -> Option<String> {
|
||||
let CodexClientError::Upstream { body, .. } = error else {
|
||||
return None;
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(body).ok()?;
|
||||
value
|
||||
.pointer("/response/error/code")
|
||||
.or_else(|| value.pointer("/error/code"))
|
||||
.or_else(|| value.get("code"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
pub(crate) fn rate_limit_cooldown_until(
|
||||
error: &CodexClientError,
|
||||
now: DateTime<Utc>,
|
||||
) -> DateTime<Utc> {
|
||||
let retry_after_seconds = match error {
|
||||
pub(crate) fn upstream_error_rate_limit_headers(error: &CodexClientError) -> Vec<(String, String)> {
|
||||
match error {
|
||||
CodexClientError::Upstream {
|
||||
retry_after_seconds,
|
||||
..
|
||||
} => retry_after_seconds.unwrap_or(60),
|
||||
_ => 60,
|
||||
};
|
||||
now + Duration::seconds(retry_after_seconds.min(i64::MAX as u64) as i64)
|
||||
rate_limit_headers, ..
|
||||
} => rate_limit_headers.clone(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 {
|
||||
@@ -336,17 +261,6 @@ pub(crate) fn backend_transport_name(transport: CodexBackendTransport) -> &'stat
|
||||
}
|
||||
}
|
||||
|
||||
fn status_code_is_rate_limited(status_code: u16) -> bool {
|
||||
status_code == 429
|
||||
}
|
||||
|
||||
fn status_code_is_quota_exhausted(status_code: u16) -> bool {
|
||||
status_code == 402
|
||||
}
|
||||
|
||||
fn status_code_is_transient_upstream(status_code: u16) -> bool {
|
||||
matches!(status_code, 500..=599)
|
||||
}
|
||||
/// Responses 调度错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ResponseDispatchError {
|
||||
@@ -411,7 +325,7 @@ pub enum ResponseDispatchError {
|
||||
#[error("upstream response did not include visible output")]
|
||||
EmptyUpstreamResponse,
|
||||
#[error("upstream response failed: {0:?}")]
|
||||
Failed(ResponsesSseFailure),
|
||||
Failed(ClientFailure),
|
||||
}
|
||||
|
||||
impl ResponseDispatchError {
|
||||
@@ -471,7 +385,7 @@ impl ResponseDispatchError {
|
||||
| Self::EmptyUpstreamResponse => 502,
|
||||
Self::UpstreamUnavailable { .. } => 503,
|
||||
Self::ContinuationBusy => 503,
|
||||
Self::Failed(failure) => stream_failure_http_status(failure),
|
||||
Self::Failed(failure) => failure.status_code(),
|
||||
Self::ModelUnsupported { .. } | Self::HistoryUnavailable { .. } => 400,
|
||||
Self::Upstream(error) => upstream_error_http_status(error),
|
||||
}
|
||||
@@ -567,7 +481,7 @@ impl ResponseDispatchError {
|
||||
Self::Failed(failure) => DispatchErrorMetadata {
|
||||
failure_class: DispatchFailureClass::ResponseFailed,
|
||||
exhausted_count: None,
|
||||
upstream_error: Some(sse_failure_error_body(failure)),
|
||||
upstream_error: Some(sse_failure_error_body(failure.failure())),
|
||||
upstream_status: None,
|
||||
diagnostics: None,
|
||||
},
|
||||
|
||||
+40
-97
@@ -1,8 +1,6 @@
|
||||
//! 账号耗尽状态聚合。
|
||||
|
||||
use crate::fleet::account::AccountStatus;
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum ExhaustedAccountKind {
|
||||
QuotaExhausted,
|
||||
RateLimited,
|
||||
@@ -15,6 +13,35 @@ pub(crate) enum ExhaustedAccountKind {
|
||||
UpstreamUnavailable,
|
||||
}
|
||||
|
||||
/// Controller 交给 attempt runner 聚合的单账号耗尽事实。
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct AccountExhaustionRecord {
|
||||
pub account_id: Option<String>,
|
||||
pub kind: ExhaustedAccountKind,
|
||||
pub upstream_error: String,
|
||||
pub status_code: Option<u16>,
|
||||
}
|
||||
|
||||
impl AccountExhaustionRecord {
|
||||
pub(in crate::dispatch) fn new(
|
||||
account_id: impl Into<String>,
|
||||
kind: ExhaustedAccountKind,
|
||||
upstream_error: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
account_id: Some(account_id.into()),
|
||||
kind,
|
||||
upstream_error: upstream_error.into(),
|
||||
status_code: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn with_status_code(mut self, status_code: u16) -> Self {
|
||||
self.status_code = Some(status_code);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl ExhaustedAccountKind {
|
||||
pub(crate) fn message_reason(self) -> &'static str {
|
||||
match self {
|
||||
@@ -31,7 +58,7 @@ impl ExhaustedAccountKind {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(crate) struct ExhaustedAccount {
|
||||
pub kind: ExhaustedAccountKind,
|
||||
pub count: usize,
|
||||
@@ -69,100 +96,12 @@ struct ExhaustedCounter {
|
||||
}
|
||||
|
||||
impl AccountExhaustionTracker {
|
||||
pub(crate) fn last_account_id(&self) -> Option<&str> {
|
||||
self.last_account_id.as_deref()
|
||||
}
|
||||
|
||||
pub(crate) fn record_rate_limited(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
pub(in crate::dispatch) fn record_exhaustion(&mut self, record: AccountExhaustionRecord) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::RateLimited,
|
||||
upstream_error,
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn record_quota_exhausted(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::QuotaExhausted,
|
||||
upstream_error,
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn record_auth_failure(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
account_status: AccountStatus,
|
||||
upstream_error: impl Into<String>,
|
||||
status_code: Option<u16>,
|
||||
) {
|
||||
let kind = match account_status {
|
||||
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
|
||||
AccountStatus::Banned => ExhaustedAccountKind::Banned,
|
||||
_ => ExhaustedAccountKind::Expired,
|
||||
};
|
||||
self.record(account_id, kind, upstream_error, status_code);
|
||||
}
|
||||
|
||||
pub(crate) fn record_cloudflare_challenge(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::CloudflareChallenge,
|
||||
upstream_error,
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn record_cloudflare_path_blocked(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::CloudflarePathBlocked,
|
||||
upstream_error,
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn record_model_unsupported(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::ModelUnsupported,
|
||||
upstream_error,
|
||||
None,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn record_upstream_unavailable(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
upstream_error: impl Into<String>,
|
||||
) {
|
||||
self.record(
|
||||
account_id,
|
||||
ExhaustedAccountKind::UpstreamUnavailable,
|
||||
upstream_error,
|
||||
None,
|
||||
record.account_id.as_deref(),
|
||||
record.kind,
|
||||
record.upstream_error,
|
||||
record.status_code,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -177,6 +116,10 @@ impl AccountExhaustionTracker {
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn last_account_id(&self) -> Option<&str> {
|
||||
self.last_account_id.as_deref()
|
||||
}
|
||||
|
||||
fn record(
|
||||
&mut self,
|
||||
account_id: Option<&str>,
|
||||
@@ -0,0 +1,4 @@
|
||||
//! 已分类账号失败的聚合、状态副作用和 SSE 终态事实。
|
||||
|
||||
pub(crate) mod exhaustion;
|
||||
pub(in crate::dispatch) mod sse;
|
||||
@@ -0,0 +1,78 @@
|
||||
//! Responses SSE 失败事实和流终态元数据。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure};
|
||||
|
||||
pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
|
||||
pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str =
|
||||
"Upstream stream closed before response.completed";
|
||||
|
||||
pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
|
||||
match failure.upstream_code.as_deref() {
|
||||
Some(code) => serde_json::json!({
|
||||
"error": {
|
||||
"code": code,
|
||||
"message": failure.message.as_str(),
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
None => failure.message.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn stream_failure_metadata(
|
||||
failure: &ResponsesSseFailure,
|
||||
usage: Option<TokenUsage>,
|
||||
) -> Value {
|
||||
let mut metadata = serde_json::json!({
|
||||
"stream": true,
|
||||
"failed": true,
|
||||
"failureEvent": failure.event,
|
||||
"failureMessage": failure.message,
|
||||
"upstreamCode": failure.upstream_code,
|
||||
"usage": usage,
|
||||
});
|
||||
enrich_stream_failure_source_metadata(&mut metadata, failure);
|
||||
metadata
|
||||
}
|
||||
|
||||
fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
object.insert(
|
||||
"failureSource".to_string(),
|
||||
Value::String(stream_failure_source(failure).to_string()),
|
||||
);
|
||||
if let Some(detail) = synthetic_stream_disconnected_detail(failure) {
|
||||
object.insert("synthetic".to_string(), Value::Bool(true));
|
||||
if !detail.is_empty() {
|
||||
object.insert("failureDetail".to_string(), Value::String(detail));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
|
||||
if synthetic_stream_disconnected_detail(failure).is_some() {
|
||||
"proxy"
|
||||
} else {
|
||||
"upstream"
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn synthetic_stream_disconnected_detail(
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> Option<String> {
|
||||
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
|
||||
return None;
|
||||
}
|
||||
let detail = failure
|
||||
.message
|
||||
.strip_prefix(STREAM_DISCONNECTED_MESSAGE)?
|
||||
.strip_prefix(": ")
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_string();
|
||||
Some(detail)
|
||||
}
|
||||
@@ -0,0 +1,959 @@
|
||||
//! 单账号 attempt 的 stepwise 执行器。
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::AccountIdentityService,
|
||||
controllers::{
|
||||
AttemptRoutePreparation, AttemptRoutePreparationContext, AttemptUpstreamErrorContext,
|
||||
ControllerRequestScope, ControllerSet,
|
||||
},
|
||||
errors::{ClientFailure, ResponseDispatchError},
|
||||
failure::exhaustion::AccountExhaustionTracker,
|
||||
routing::candidates::AccountAttemptLedger,
|
||||
transport::{
|
||||
account::{
|
||||
AccountUpstreamContext, create_response_stream_with_account,
|
||||
create_response_with_account, prepare_response_transport_with_account,
|
||||
},
|
||||
canonical::{CanonicalStreamDecoder, normalize_complete_response},
|
||||
prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit},
|
||||
},
|
||||
},
|
||||
fleet::{
|
||||
account::Account,
|
||||
pool::{AccountCandidateLease, AccountPoolService},
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::upstream_failure_facts,
|
||||
protocol::responses::{
|
||||
CodexResponsesRequest, CollectedResponse, TransportRequirement, transport_requirement,
|
||||
},
|
||||
transport::{CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport},
|
||||
},
|
||||
};
|
||||
|
||||
use super::{
|
||||
contract::{
|
||||
AttemptAccountFacts, AttemptApplyOutcome, AttemptContractError, AttemptDecision,
|
||||
AttemptObservation, AttemptObservationKind, AttemptRejection, AttemptReturnKind,
|
||||
CandidateLedgerFacts, CompleteResponseFacts, EstablishedAttemptContext,
|
||||
EstablishedComplete, EstablishedCompleteBody, EstablishedResponse, EstablishedStream,
|
||||
PendingAttempt, PendingCompleteResponse, PendingProtocolFailure, PendingStreamResponse,
|
||||
PinnedCandidateAcquireFailureKind, ProtocolFailureKind,
|
||||
},
|
||||
trace::ResponseDispatchTrace,
|
||||
};
|
||||
|
||||
pub(in crate::dispatch) enum AttemptMode {
|
||||
Complete { tuple_schema: Option<Value> },
|
||||
Stream { tuple_schema: Option<Value> },
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct AttemptRunnerDependencies<'a> {
|
||||
pub account_pool: &'a AccountPoolService,
|
||||
pub codex: &'a CodexBackendClient,
|
||||
pub controllers: &'a ControllerSet,
|
||||
pub account_identity: &'a AccountIdentityService,
|
||||
pub request_id: &'a str,
|
||||
pub route: &'a str,
|
||||
pub requested_model: &'a str,
|
||||
pub started_at: Instant,
|
||||
}
|
||||
|
||||
/// 尚未越过提交边界的 attempt capability;只有该类型接受 retry decision。
|
||||
pub(in crate::dispatch) struct OpenAttempt<'runner, 'dependencies> {
|
||||
runner: &'runner mut AttemptRunner<'dependencies>,
|
||||
observation: Box<AttemptObservation>,
|
||||
}
|
||||
|
||||
impl OpenAttempt<'_, '_> {
|
||||
pub(in crate::dispatch) fn observation(&self) -> &AttemptObservation {
|
||||
&self.observation
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn controller_scope(
|
||||
&self,
|
||||
) -> Result<&ControllerRequestScope, AttemptContractError> {
|
||||
self.runner.controller_scope()
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn apply(
|
||||
self,
|
||||
decision: AttemptDecision,
|
||||
) -> Result<AttemptApplyOutcome, AttemptContractError> {
|
||||
self.runner.apply_open(decision).await
|
||||
}
|
||||
}
|
||||
|
||||
/// 已越过提交边界的 attempt capability;该类型只允许建立 stream。
|
||||
pub(in crate::dispatch) struct CommittedAttempt {
|
||||
established: EstablishedResponse,
|
||||
}
|
||||
|
||||
impl CommittedAttempt {
|
||||
pub(in crate::dispatch) fn accept(self) -> EstablishedResponse {
|
||||
self.established
|
||||
}
|
||||
}
|
||||
|
||||
/// `next` 产出的提交 typestate;Committed 分支没有 decision/retry API。
|
||||
pub(in crate::dispatch) enum AttemptStep<'runner, 'dependencies> {
|
||||
Open(OpenAttempt<'runner, 'dependencies>),
|
||||
Committed(CommittedAttempt),
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct AttemptRunner<'a> {
|
||||
dependencies: AttemptRunnerDependencies<'a>,
|
||||
mode: AttemptMode,
|
||||
request: Option<CodexResponsesRequest>,
|
||||
controller_scope: Option<ControllerRequestScope>,
|
||||
candidates: AccountAttemptLedger,
|
||||
exhaustion: AccountExhaustionTracker,
|
||||
on_exhaustion: Option<ClientFailure>,
|
||||
trace: Option<ResponseDispatchTrace>,
|
||||
required_account_id: Option<String>,
|
||||
pending: Option<PendingAttempt>,
|
||||
}
|
||||
|
||||
impl<'a> AttemptRunner<'a> {
|
||||
pub(in crate::dispatch) fn new(
|
||||
dependencies: AttemptRunnerDependencies<'a>,
|
||||
mode: AttemptMode,
|
||||
request: CodexResponsesRequest,
|
||||
controller_scope: ControllerRequestScope,
|
||||
candidates: AccountAttemptLedger,
|
||||
) -> Self {
|
||||
Self {
|
||||
dependencies,
|
||||
mode,
|
||||
request: Some(request),
|
||||
controller_scope: Some(controller_scope),
|
||||
candidates,
|
||||
exhaustion: AccountExhaustionTracker::default(),
|
||||
on_exhaustion: None,
|
||||
trace: Some(ResponseDispatchTrace::default()),
|
||||
required_account_id: None,
|
||||
pending: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn next(
|
||||
&mut self,
|
||||
) -> Result<AttemptStep<'_, 'a>, AttemptContractError> {
|
||||
if self.pending.is_some() {
|
||||
return Err(AttemptContractError::DecisionRequired);
|
||||
}
|
||||
|
||||
let request = self
|
||||
.request
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let acquired = if let Some(account_id) = self.required_account_id.take() {
|
||||
match self
|
||||
.dependencies
|
||||
.account_pool
|
||||
.acquire_candidate(request.model(), &account_id, chrono::Utc::now())
|
||||
.await
|
||||
{
|
||||
AccountCandidateLease::Acquired(acquired) => Some(*acquired),
|
||||
AccountCandidateLease::Busy => {
|
||||
self.pending = Some(PendingAttempt::PinnedCandidateUnavailable {
|
||||
account_id,
|
||||
kind: PinnedCandidateAcquireFailureKind::Busy,
|
||||
});
|
||||
return self.step();
|
||||
}
|
||||
AccountCandidateLease::Unavailable => {
|
||||
self.pending = Some(PendingAttempt::PinnedCandidateUnavailable {
|
||||
account_id,
|
||||
kind: PinnedCandidateAcquireFailureKind::Unavailable,
|
||||
});
|
||||
return self.step();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.candidates
|
||||
.acquire_next(self.dependencies.account_pool)
|
||||
.await
|
||||
};
|
||||
let Some(acquired) = acquired else {
|
||||
self.pending = Some(PendingAttempt::NoCandidate);
|
||||
return self.step();
|
||||
};
|
||||
|
||||
let account = acquired.account.clone();
|
||||
let controller_scope = self
|
||||
.controller_scope
|
||||
.as_mut()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let prepared_account = self
|
||||
.dependencies
|
||||
.controllers
|
||||
.prepare_attempt_account(acquired)
|
||||
.await;
|
||||
let acquired = prepared_account.lease;
|
||||
let cookie_header = prepared_account.cookie_header;
|
||||
|
||||
let account_scoped_request = match self.dependencies.controllers.prepare_attempt_route(
|
||||
controller_scope,
|
||||
request,
|
||||
AttemptRoutePreparationContext {
|
||||
account_identity: self.dependencies.account_identity,
|
||||
account_id: &account.id,
|
||||
},
|
||||
) {
|
||||
AttemptRoutePreparation::Ready(request) => *request,
|
||||
AttemptRoutePreparation::Unavailable { message } => {
|
||||
acquired.release_without_usage().await;
|
||||
self.pending = Some(PendingAttempt::RoutePreparationRejected { account, message });
|
||||
return self.step();
|
||||
}
|
||||
};
|
||||
|
||||
let attempt = self
|
||||
.trace
|
||||
.as_mut()
|
||||
.ok_or(AttemptContractError::Terminal)?
|
||||
.start_attempt(&account.id);
|
||||
let upstream_context = AccountUpstreamContext {
|
||||
codex: self.dependencies.codex,
|
||||
request_id: self.dependencies.request_id,
|
||||
account: &account,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
};
|
||||
let ((), prepared_transport) = tokio::join!(
|
||||
self.dependencies
|
||||
.account_pool
|
||||
.wait_for_request_interval(&acquired),
|
||||
prepare_response_transport_with_account(upstream_context, &account_scoped_request),
|
||||
);
|
||||
let prepared_transport = match prepared_transport {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
let attempt_request = account_scoped_request.into_request();
|
||||
acquired.complete().await;
|
||||
let transport = error
|
||||
.transport()
|
||||
.unwrap_or_else(|| intended_transport(&attempt_request));
|
||||
observe_upstream_error(
|
||||
self.dependencies.controllers,
|
||||
AttemptUpstreamErrorContext {
|
||||
request_id: self.dependencies.request_id,
|
||||
account: &account,
|
||||
route: self.dependencies.route,
|
||||
model: self.dependencies.requested_model,
|
||||
started_at: self.dependencies.started_at,
|
||||
stream: matches!(self.mode, AttemptMode::Stream { .. }),
|
||||
transport,
|
||||
request: &attempt_request,
|
||||
error: &error,
|
||||
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
|
||||
attempt: &attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
self.pending = Some(PendingAttempt::UpstreamFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error,
|
||||
});
|
||||
return self.step();
|
||||
}
|
||||
};
|
||||
|
||||
match &self.mode {
|
||||
AttemptMode::Complete { tuple_schema } => {
|
||||
let result = create_response_with_account(
|
||||
upstream_context,
|
||||
&account_scoped_request,
|
||||
prepared_transport,
|
||||
attempt.started_at(),
|
||||
)
|
||||
.await;
|
||||
let attempt_request = account_scoped_request.into_request();
|
||||
acquired.complete().await;
|
||||
self.pending = Some(match result {
|
||||
Ok(response) => {
|
||||
self.dependencies
|
||||
.controllers
|
||||
.observe_complete_upstream(
|
||||
&account,
|
||||
&response.set_cookie_headers,
|
||||
&response.rate_limit_headers,
|
||||
)
|
||||
.await;
|
||||
match normalize_complete_response(&response.body, tuple_schema.as_ref()) {
|
||||
Ok(collected) => PendingAttempt::CompleteResponse(Box::new(
|
||||
PendingCompleteResponse {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
response,
|
||||
collected,
|
||||
},
|
||||
)),
|
||||
Err(error) => PendingAttempt::ProtocolFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport: response.transport,
|
||||
error: PendingProtocolFailure::InvalidSse(error),
|
||||
},
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let transport = error
|
||||
.transport()
|
||||
.unwrap_or_else(|| intended_transport(request));
|
||||
observe_upstream_error(
|
||||
self.dependencies.controllers,
|
||||
AttemptUpstreamErrorContext {
|
||||
request_id: self.dependencies.request_id,
|
||||
account: &account,
|
||||
route: self.dependencies.route,
|
||||
model: self.dependencies.requested_model,
|
||||
started_at: self.dependencies.started_at,
|
||||
stream: matches!(self.mode, AttemptMode::Stream { .. }),
|
||||
transport,
|
||||
request: &attempt_request,
|
||||
error: &error,
|
||||
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
|
||||
attempt: &attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
PendingAttempt::UpstreamFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
AttemptMode::Stream { tuple_schema } => {
|
||||
let result = create_response_stream_with_account(
|
||||
upstream_context,
|
||||
&account_scoped_request,
|
||||
prepared_transport,
|
||||
)
|
||||
.await;
|
||||
let attempt_request = account_scoped_request.into_request();
|
||||
self.pending = Some(match result {
|
||||
Err(error) => {
|
||||
acquired.complete().await;
|
||||
let transport = error
|
||||
.transport()
|
||||
.unwrap_or_else(|| intended_transport(request));
|
||||
observe_upstream_error(
|
||||
self.dependencies.controllers,
|
||||
AttemptUpstreamErrorContext {
|
||||
request_id: self.dependencies.request_id,
|
||||
account: &account,
|
||||
route: self.dependencies.route,
|
||||
model: self.dependencies.requested_model,
|
||||
started_at: self.dependencies.started_at,
|
||||
stream: matches!(self.mode, AttemptMode::Stream { .. }),
|
||||
transport,
|
||||
request: &attempt_request,
|
||||
error: &error,
|
||||
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
|
||||
attempt: &attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
PendingAttempt::UpstreamFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error,
|
||||
}
|
||||
}
|
||||
Ok(response) => {
|
||||
let CodexBackendStreamingResponse {
|
||||
body,
|
||||
transport,
|
||||
turn_state,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_header_updates,
|
||||
turn_state_update,
|
||||
websocket_pool_decision,
|
||||
diagnostics,
|
||||
response_metadata,
|
||||
transport_metrics,
|
||||
connection_local_continuation,
|
||||
} = response;
|
||||
let mut decoder = CanonicalStreamDecoder::new(tuple_schema.clone());
|
||||
match prefetch_until_commit(
|
||||
body,
|
||||
&mut decoder,
|
||||
attempt_request.stream_commit_policy,
|
||||
attempt.started_at(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Err(StreamPrefetchError::Upstream(error)) => {
|
||||
acquired.complete().await;
|
||||
observe_upstream_error(
|
||||
self.dependencies.controllers,
|
||||
AttemptUpstreamErrorContext {
|
||||
request_id: self.dependencies.request_id,
|
||||
account: &account,
|
||||
route: self.dependencies.route,
|
||||
model: self.dependencies.requested_model,
|
||||
started_at: self.dependencies.started_at,
|
||||
stream: matches!(self.mode, AttemptMode::Stream { .. }),
|
||||
transport,
|
||||
request: &attempt_request,
|
||||
error: &error,
|
||||
trace: self
|
||||
.trace
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)?,
|
||||
attempt: &attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
PendingAttempt::UpstreamFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error,
|
||||
}
|
||||
}
|
||||
Err(StreamPrefetchError::Empty) => {
|
||||
acquired.complete().await;
|
||||
PendingAttempt::ProtocolFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error: PendingProtocolFailure::EmptyStream,
|
||||
}
|
||||
}
|
||||
Err(StreamPrefetchError::NoCommitBoundary) => {
|
||||
acquired.complete().await;
|
||||
PendingAttempt::ProtocolFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error: PendingProtocolFailure::NoCommitBoundary,
|
||||
}
|
||||
}
|
||||
Err(StreamPrefetchError::InvalidSse(error)) => {
|
||||
acquired.complete().await;
|
||||
PendingAttempt::ProtocolFailure {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
error: PendingProtocolFailure::InvalidSse(error),
|
||||
}
|
||||
}
|
||||
Ok(PrefetchedStream {
|
||||
bytes: prefetched,
|
||||
initial_batch,
|
||||
body,
|
||||
first_event_ms,
|
||||
}) => {
|
||||
let first_failure = initial_batch.first_failure();
|
||||
PendingAttempt::StreamResponse(Box::new(PendingStreamResponse {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
lease: acquired,
|
||||
response: CodexBackendStreamingResponse {
|
||||
body,
|
||||
transport,
|
||||
turn_state,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_header_updates,
|
||||
turn_state_update,
|
||||
websocket_pool_decision,
|
||||
diagnostics,
|
||||
response_metadata,
|
||||
transport_metrics,
|
||||
connection_local_continuation,
|
||||
},
|
||||
prefetched,
|
||||
decoder,
|
||||
initial_batch,
|
||||
first_failure,
|
||||
first_event_ms,
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
self.step()
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn controller_scope(
|
||||
&self,
|
||||
) -> Result<&ControllerRequestScope, AttemptContractError> {
|
||||
self.controller_scope
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)
|
||||
}
|
||||
|
||||
async fn apply_open(
|
||||
&mut self,
|
||||
decision: AttemptDecision,
|
||||
) -> Result<AttemptApplyOutcome, AttemptContractError> {
|
||||
let pending = self
|
||||
.pending
|
||||
.take()
|
||||
.ok_or(AttemptContractError::DecisionRequired)?;
|
||||
match decision {
|
||||
AttemptDecision::Accept => self
|
||||
.establish(pending)
|
||||
.map(AttemptApplyOutcome::Established),
|
||||
AttemptDecision::RetrySameAccount => {
|
||||
let account_id = pending_account(&pending)
|
||||
.ok_or_else(|| invalid("retry_same", &pending))?
|
||||
.to_string();
|
||||
if !self.dependencies.controllers.prepare_same_account_retry(
|
||||
self.controller_scope
|
||||
.as_mut()
|
||||
.ok_or(AttemptContractError::Terminal)?,
|
||||
&account_id,
|
||||
) {
|
||||
return Err(invalid("retry_same", &pending));
|
||||
}
|
||||
release_pending_stream(pending).await;
|
||||
self.required_account_id = Some(account_id);
|
||||
Ok(AttemptApplyOutcome::Continue)
|
||||
}
|
||||
AttemptDecision::RetryNextCandidate {
|
||||
exhaustion,
|
||||
on_exhaustion,
|
||||
} => {
|
||||
release_pending_stream(pending).await;
|
||||
self.on_exhaustion = on_exhaustion;
|
||||
if let Some(exhaustion) = exhaustion {
|
||||
self.exhaustion.record_exhaustion(exhaustion);
|
||||
}
|
||||
Ok(AttemptApplyOutcome::Continue)
|
||||
}
|
||||
AttemptDecision::Return(kind) => self.reject(pending, kind).await,
|
||||
}
|
||||
}
|
||||
|
||||
fn step(&mut self) -> Result<AttemptStep<'_, 'a>, AttemptContractError> {
|
||||
let kind = {
|
||||
let pending = self
|
||||
.pending
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::DecisionRequired)?;
|
||||
observation_kind(pending, &self.candidates, &self.exhaustion)
|
||||
};
|
||||
let Some(kind) = kind else {
|
||||
let pending = self
|
||||
.pending
|
||||
.take()
|
||||
.ok_or(AttemptContractError::DecisionRequired)?;
|
||||
return Ok(AttemptStep::Committed(CommittedAttempt {
|
||||
established: self.establish(pending)?,
|
||||
}));
|
||||
};
|
||||
let observation = self.observation(kind)?;
|
||||
Ok(AttemptStep::Open(OpenAttempt {
|
||||
runner: self,
|
||||
observation: Box::new(observation),
|
||||
}))
|
||||
}
|
||||
|
||||
fn observation(
|
||||
&self,
|
||||
kind: AttemptObservationKind,
|
||||
) -> Result<AttemptObservation, AttemptContractError> {
|
||||
let pending = self
|
||||
.pending
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::DecisionRequired)?;
|
||||
let request = self
|
||||
.request
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let controller_scope = self
|
||||
.controller_scope
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let account = pending_account_value(pending);
|
||||
let account_id = account.map(|account| account.id.as_str());
|
||||
Ok(AttemptObservation {
|
||||
account: account.map(AttemptAccountFacts::from),
|
||||
attempt: pending_attempt(pending).cloned(),
|
||||
transport: pending_transport(pending).unwrap_or_else(|| intended_transport(request)),
|
||||
routing: self
|
||||
.dependencies
|
||||
.controllers
|
||||
.attempt_routing_facts(controller_scope, account_id),
|
||||
kind,
|
||||
})
|
||||
}
|
||||
|
||||
fn establish(
|
||||
&mut self,
|
||||
pending: PendingAttempt,
|
||||
) -> Result<EstablishedResponse, AttemptContractError> {
|
||||
let established = match pending {
|
||||
PendingAttempt::CompleteResponse(complete) => {
|
||||
let PendingCompleteResponse {
|
||||
account,
|
||||
attempt,
|
||||
response,
|
||||
collected,
|
||||
..
|
||||
} = *complete;
|
||||
let body = EstablishedCompleteBody::try_from(collected).map_err(|()| {
|
||||
AttemptContractError::InvalidDecision {
|
||||
decision: "establish",
|
||||
observation: "non_successful_complete_response",
|
||||
}
|
||||
})?;
|
||||
EstablishedResponse::Complete(Box::new(EstablishedComplete {
|
||||
context: self.take_established_context(account, attempt)?,
|
||||
response,
|
||||
body,
|
||||
}))
|
||||
}
|
||||
PendingAttempt::StreamResponse(stream) => {
|
||||
let PendingStreamResponse {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
lease,
|
||||
response,
|
||||
decoder,
|
||||
initial_batch,
|
||||
first_event_ms,
|
||||
..
|
||||
} = *stream;
|
||||
EstablishedResponse::Stream(Box::new(EstablishedStream {
|
||||
context: self.take_established_context(account, attempt)?,
|
||||
attempt_request,
|
||||
lease,
|
||||
response,
|
||||
decoder,
|
||||
initial_batch,
|
||||
first_event_ms,
|
||||
}))
|
||||
}
|
||||
other => return Err(invalid("establish", &other)),
|
||||
};
|
||||
Ok(established)
|
||||
}
|
||||
|
||||
async fn reject(
|
||||
&mut self,
|
||||
pending: PendingAttempt,
|
||||
kind: AttemptReturnKind,
|
||||
) -> Result<AttemptApplyOutcome, AttemptContractError> {
|
||||
let account = pending_account_value(&pending).cloned();
|
||||
let account_id = account
|
||||
.as_ref()
|
||||
.map(|account| account.id.clone())
|
||||
.or_else(|| self.exhaustion.last_account_id().map(ToString::to_string));
|
||||
let attempt_request = pending_attempt_request(&pending).cloned();
|
||||
let attempt = pending_attempt(&pending).cloned();
|
||||
let request = self
|
||||
.request
|
||||
.as_ref()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let transport = pending_transport(&pending).unwrap_or_else(|| intended_transport(request));
|
||||
let stream_failure = match &pending {
|
||||
PendingAttempt::StreamResponse(stream) => {
|
||||
stream.first_failure.as_ref().map(|failure| {
|
||||
super::contract::RejectedStreamFailure {
|
||||
failure: failure.clone(),
|
||||
prefetched: stream.prefetched.clone(),
|
||||
diagnostics: stream.response.diagnostics.clone(),
|
||||
rate_limit_headers: stream.response.rate_limit_headers.clone(),
|
||||
}
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let error =
|
||||
rejection_error(pending, kind, &self.exhaustion, self.on_exhaustion.take()).await?;
|
||||
let request = self.request.take().ok_or(AttemptContractError::Terminal)?;
|
||||
self.controller_scope
|
||||
.take()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let trace = self.trace.take().ok_or(AttemptContractError::Terminal)?;
|
||||
Ok(AttemptApplyOutcome::Rejected(Box::new(AttemptRejection {
|
||||
request,
|
||||
trace,
|
||||
account_id,
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
transport,
|
||||
stream_failure,
|
||||
error,
|
||||
})))
|
||||
}
|
||||
|
||||
fn take_established_context(
|
||||
&mut self,
|
||||
account: Account,
|
||||
attempt: super::trace::ResponseDispatchAttempt,
|
||||
) -> Result<EstablishedAttemptContext, AttemptContractError> {
|
||||
Ok(EstablishedAttemptContext {
|
||||
request: self.request.take().ok_or(AttemptContractError::Terminal)?,
|
||||
controller_scope: self
|
||||
.controller_scope
|
||||
.take()
|
||||
.ok_or(AttemptContractError::Terminal)?,
|
||||
trace: self.trace.take().ok_or(AttemptContractError::Terminal)?,
|
||||
account,
|
||||
attempt,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn observe_upstream_error(
|
||||
controllers: &ControllerSet,
|
||||
context: AttemptUpstreamErrorContext<'_>,
|
||||
) {
|
||||
controllers.observe_upstream_error(context).await;
|
||||
}
|
||||
|
||||
fn observation_kind(
|
||||
pending: &PendingAttempt,
|
||||
candidates: &AccountAttemptLedger,
|
||||
exhaustion: &AccountExhaustionTracker,
|
||||
) -> Option<AttemptObservationKind> {
|
||||
Some(match pending {
|
||||
PendingAttempt::NoCandidate => AttemptObservationKind::NoCandidate {
|
||||
ledger: CandidateLedgerFacts {
|
||||
candidates: candidates.candidate_count(),
|
||||
attempted: candidates.attempted_count(),
|
||||
state_excluded: candidates.state_excluded_count(),
|
||||
},
|
||||
last_exhausted: exhaustion.last_exhausted(),
|
||||
},
|
||||
PendingAttempt::PinnedCandidateUnavailable { account_id, kind } => {
|
||||
AttemptObservationKind::PinnedCandidateUnavailable {
|
||||
account_id: account_id.clone(),
|
||||
kind: *kind,
|
||||
}
|
||||
}
|
||||
PendingAttempt::RoutePreparationRejected { message, .. } => {
|
||||
AttemptObservationKind::RoutePreparationRejected {
|
||||
message: message.clone(),
|
||||
}
|
||||
}
|
||||
PendingAttempt::UpstreamFailure { error, .. } => {
|
||||
AttemptObservationKind::UpstreamFailure(upstream_failure_facts(error))
|
||||
}
|
||||
PendingAttempt::ProtocolFailure { error, .. } => AttemptObservationKind::ProtocolFailure {
|
||||
kind: match error {
|
||||
PendingProtocolFailure::InvalidSse(_) => ProtocolFailureKind::InvalidSse,
|
||||
PendingProtocolFailure::EmptyStream => ProtocolFailureKind::EmptyStream,
|
||||
PendingProtocolFailure::NoCommitBoundary => ProtocolFailureKind::NoCommitBoundary,
|
||||
},
|
||||
message: match error {
|
||||
PendingProtocolFailure::InvalidSse(error) => error.to_string(),
|
||||
PendingProtocolFailure::EmptyStream => {
|
||||
"upstream stream ended without data".to_string()
|
||||
}
|
||||
PendingProtocolFailure::NoCommitBoundary => {
|
||||
"upstream stream ended before output or a terminal event".to_string()
|
||||
}
|
||||
},
|
||||
},
|
||||
PendingAttempt::CompleteResponse(complete) => {
|
||||
AttemptObservationKind::CompleteResponse(match &complete.collected {
|
||||
CollectedResponse::Completed(_) => CompleteResponseFacts::Completed,
|
||||
CollectedResponse::Incomplete(_) => CompleteResponseFacts::Incomplete,
|
||||
CollectedResponse::Failed(failure) => {
|
||||
CompleteResponseFacts::Failed(failure.clone())
|
||||
}
|
||||
CollectedResponse::MissingCompleted => CompleteResponseFacts::MissingCompleted,
|
||||
CollectedResponse::Empty => CompleteResponseFacts::Empty,
|
||||
})
|
||||
}
|
||||
PendingAttempt::StreamResponse(stream) => {
|
||||
return stream
|
||||
.first_failure
|
||||
.clone()
|
||||
.map(AttemptObservationKind::StreamFailure);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn release_pending_stream(pending: PendingAttempt) {
|
||||
if let PendingAttempt::StreamResponse(stream) = pending {
|
||||
stream.lease.complete().await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn rejection_error(
|
||||
pending: PendingAttempt,
|
||||
kind: AttemptReturnKind,
|
||||
exhaustion: &AccountExhaustionTracker,
|
||||
on_exhaustion: Option<ClientFailure>,
|
||||
) -> Result<ResponseDispatchError, AttemptContractError> {
|
||||
match kind {
|
||||
AttemptReturnKind::Failed(failure) => {
|
||||
release_pending_stream(pending).await;
|
||||
Ok(ResponseDispatchError::Failed(failure))
|
||||
}
|
||||
AttemptReturnKind::ContinuationBusy => Ok(ResponseDispatchError::ContinuationBusy),
|
||||
AttemptReturnKind::RouteUnavailable { message } => {
|
||||
Ok(ResponseDispatchError::HistoryUnavailable {
|
||||
upstream_error: message,
|
||||
})
|
||||
}
|
||||
AttemptReturnKind::Observed => match pending {
|
||||
PendingAttempt::NoCandidate => Ok(on_exhaustion.map_or_else(
|
||||
|| {
|
||||
exhaustion
|
||||
.last_exhausted()
|
||||
.map(ResponseDispatchError::from_exhausted_account)
|
||||
.unwrap_or(ResponseDispatchError::NoActiveAccount)
|
||||
},
|
||||
ResponseDispatchError::Failed,
|
||||
)),
|
||||
PendingAttempt::RoutePreparationRejected { message, .. } => {
|
||||
Ok(ResponseDispatchError::HistoryUnavailable {
|
||||
upstream_error: message,
|
||||
})
|
||||
}
|
||||
PendingAttempt::PinnedCandidateUnavailable { .. } => {
|
||||
Err(invalid("return", &PendingAttempt::NoCandidate))
|
||||
}
|
||||
PendingAttempt::UpstreamFailure { error, .. } => {
|
||||
Ok(ResponseDispatchError::Upstream(error))
|
||||
}
|
||||
PendingAttempt::ProtocolFailure { error, .. } => Ok(match error {
|
||||
PendingProtocolFailure::InvalidSse(error) => {
|
||||
ResponseDispatchError::InvalidSse(error)
|
||||
}
|
||||
PendingProtocolFailure::EmptyStream => ResponseDispatchError::EmptyUpstreamResponse,
|
||||
PendingProtocolFailure::NoCommitBoundary => ResponseDispatchError::MissingCompleted,
|
||||
}),
|
||||
PendingAttempt::CompleteResponse(complete) => match complete.collected {
|
||||
CollectedResponse::Failed(_) => Err(invalid(
|
||||
"return_untyped_failure",
|
||||
&PendingAttempt::NoCandidate,
|
||||
)),
|
||||
CollectedResponse::MissingCompleted => Ok(ResponseDispatchError::MissingCompleted),
|
||||
CollectedResponse::Empty => Ok(ResponseDispatchError::EmptyUpstreamResponse),
|
||||
CollectedResponse::Completed(_) | CollectedResponse::Incomplete(_) => {
|
||||
Err(invalid("return", &PendingAttempt::NoCandidate))
|
||||
}
|
||||
},
|
||||
PendingAttempt::StreamResponse(stream) => match stream.first_failure {
|
||||
Some(failure) => {
|
||||
stream.lease.complete().await;
|
||||
let _ = failure;
|
||||
Err(invalid(
|
||||
"return_untyped_failure",
|
||||
&PendingAttempt::NoCandidate,
|
||||
))
|
||||
}
|
||||
None => Err(invalid("return", &PendingAttempt::NoCandidate)),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn pending_account(pending: &PendingAttempt) -> Option<&str> {
|
||||
pending_account_value(pending).map(|account| account.id.as_str())
|
||||
}
|
||||
|
||||
fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> {
|
||||
match pending {
|
||||
PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None,
|
||||
PendingAttempt::RoutePreparationRejected { account, .. }
|
||||
| PendingAttempt::UpstreamFailure { account, .. }
|
||||
| PendingAttempt::ProtocolFailure { account, .. } => Some(account),
|
||||
PendingAttempt::CompleteResponse(complete) => Some(&complete.account),
|
||||
PendingAttempt::StreamResponse(stream) => Some(&stream.account),
|
||||
}
|
||||
}
|
||||
|
||||
fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDispatchAttempt> {
|
||||
match pending {
|
||||
PendingAttempt::UpstreamFailure { attempt, .. }
|
||||
| PendingAttempt::ProtocolFailure { attempt, .. } => Some(attempt),
|
||||
PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt),
|
||||
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRequest> {
|
||||
match pending {
|
||||
PendingAttempt::UpstreamFailure {
|
||||
attempt_request, ..
|
||||
}
|
||||
| PendingAttempt::ProtocolFailure {
|
||||
attempt_request, ..
|
||||
} => Some(attempt_request),
|
||||
PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt_request),
|
||||
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn pending_transport(
|
||||
pending: &PendingAttempt,
|
||||
) -> Option<crate::upstream::openai::transport::CodexBackendTransport> {
|
||||
match pending {
|
||||
PendingAttempt::UpstreamFailure { transport, .. }
|
||||
| PendingAttempt::ProtocolFailure { transport, .. } => Some(*transport),
|
||||
PendingAttempt::CompleteResponse(complete) => Some(complete.response.transport),
|
||||
PendingAttempt::StreamResponse(stream) => Some(stream.response.transport),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn intended_transport(request: &CodexResponsesRequest) -> CodexBackendTransport {
|
||||
match transport_requirement(request) {
|
||||
TransportRequirement::HttpRequired => CodexBackendTransport::HttpSse,
|
||||
TransportRequirement::ExplicitWebSocketWarmup
|
||||
| TransportRequirement::ExactWebSocketContinuation
|
||||
| TransportRequirement::PersistedContinuation
|
||||
| TransportRequirement::ExternalUnknown
|
||||
| TransportRequirement::NewChain => CodexBackendTransport::WebSocket,
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractError {
|
||||
AttemptContractError::InvalidDecision {
|
||||
decision,
|
||||
observation: match pending {
|
||||
PendingAttempt::NoCandidate => "no_candidate",
|
||||
PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable",
|
||||
PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected",
|
||||
PendingAttempt::UpstreamFailure { .. } => "upstream_failure",
|
||||
PendingAttempt::ProtocolFailure { .. } => "protocol_failure",
|
||||
PendingAttempt::CompleteResponse(_) => "complete_response",
|
||||
PendingAttempt::StreamResponse(_) => "stream_response",
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
//! Attempt lifecycle 的类型化 observation、decision 与结果。
|
||||
|
||||
use bytes::Bytes;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
errors::{ClientFailure, ResponseDispatchError},
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount},
|
||||
transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
|
||||
},
|
||||
fleet::{account::Account, pool::AccountLease},
|
||||
upstream::openai::{
|
||||
failure::UpstreamFailureFacts,
|
||||
protocol::{
|
||||
responses::{CodexResponsesRequest, CollectedResponse, ResponsesSseFailure},
|
||||
sse::SseError,
|
||||
},
|
||||
transport::{
|
||||
CodexBackendResponse, CodexBackendStreamingResponse, CodexBackendTransport,
|
||||
CodexClientError, CodexUpstreamDiagnostics,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
use super::trace::{ResponseDispatchAttempt, ResponseDispatchTrace};
|
||||
|
||||
/// Request/Attempt/Stream 共同消费的最终结果类别。
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum FinalOutcome {
|
||||
Completed,
|
||||
Incomplete,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct AttemptAccountFacts {
|
||||
pub id: String,
|
||||
pub email: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
}
|
||||
|
||||
impl From<&Account> for AttemptAccountFacts {
|
||||
fn from(account: &Account) -> Self {
|
||||
Self {
|
||||
id: account.id.clone(),
|
||||
email: account.email.clone(),
|
||||
plan_type: account.plan_type.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct AttemptRoutingFacts {
|
||||
pub external_origin: bool,
|
||||
pub can_retry_same_account: bool,
|
||||
pub can_retry_next_candidate: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct CandidateLedgerFacts {
|
||||
pub candidates: usize,
|
||||
pub attempted: usize,
|
||||
pub state_excluded: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum CompleteResponseFacts {
|
||||
Completed,
|
||||
Incomplete,
|
||||
Failed(ResponsesSseFailure),
|
||||
MissingCompleted,
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum ProtocolFailureKind {
|
||||
InvalidSse,
|
||||
EmptyStream,
|
||||
NoCommitBoundary,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum PinnedCandidateAcquireFailureKind {
|
||||
Busy,
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum AttemptObservationKind {
|
||||
NoCandidate {
|
||||
ledger: CandidateLedgerFacts,
|
||||
last_exhausted: Option<ExhaustedAccount>,
|
||||
},
|
||||
RoutePreparationRejected {
|
||||
message: String,
|
||||
},
|
||||
PinnedCandidateUnavailable {
|
||||
account_id: String,
|
||||
kind: PinnedCandidateAcquireFailureKind,
|
||||
},
|
||||
UpstreamFailure(UpstreamFailureFacts),
|
||||
ProtocolFailure {
|
||||
kind: ProtocolFailureKind,
|
||||
message: String,
|
||||
},
|
||||
CompleteResponse(CompleteResponseFacts),
|
||||
StreamFailure(ResponsesSseFailure),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct AttemptObservation {
|
||||
pub account: Option<AttemptAccountFacts>,
|
||||
pub attempt: Option<ResponseDispatchAttempt>,
|
||||
pub transport: CodexBackendTransport,
|
||||
pub routing: AttemptRoutingFacts,
|
||||
pub kind: AttemptObservationKind,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum AttemptDecision {
|
||||
Accept,
|
||||
RetrySameAccount,
|
||||
RetryNextCandidate {
|
||||
exhaustion: Option<AccountExhaustionRecord>,
|
||||
on_exhaustion: Option<ClientFailure>,
|
||||
},
|
||||
Return(AttemptReturnKind),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum AttemptReturnKind {
|
||||
Observed,
|
||||
Failed(ClientFailure),
|
||||
ContinuationBusy,
|
||||
RouteUnavailable { message: String },
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct EstablishedAttemptContext {
|
||||
pub request: CodexResponsesRequest,
|
||||
pub controller_scope: crate::dispatch::controllers::ControllerRequestScope,
|
||||
pub trace: ResponseDispatchTrace,
|
||||
pub account: Account,
|
||||
pub attempt: ResponseDispatchAttempt,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct EstablishedComplete {
|
||||
pub context: EstablishedAttemptContext,
|
||||
pub response: CodexBackendResponse,
|
||||
pub body: EstablishedCompleteBody,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) enum EstablishedCompleteBody {
|
||||
Completed(serde_json::Value),
|
||||
Incomplete(serde_json::Value),
|
||||
}
|
||||
|
||||
impl TryFrom<CollectedResponse> for EstablishedCompleteBody {
|
||||
type Error = ();
|
||||
|
||||
fn try_from(value: CollectedResponse) -> Result<Self, Self::Error> {
|
||||
match value {
|
||||
CollectedResponse::Completed(body) => Ok(Self::Completed(body)),
|
||||
CollectedResponse::Incomplete(body) => Ok(Self::Incomplete(body)),
|
||||
CollectedResponse::Failed(_)
|
||||
| CollectedResponse::MissingCompleted
|
||||
| CollectedResponse::Empty => Err(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct EstablishedStream {
|
||||
pub context: EstablishedAttemptContext,
|
||||
pub attempt_request: CodexResponsesRequest,
|
||||
pub lease: AccountLease,
|
||||
pub response: CodexBackendStreamingResponse,
|
||||
pub decoder: CanonicalStreamDecoder,
|
||||
pub initial_batch: CanonicalStreamBatch,
|
||||
pub first_event_ms: i64,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) enum EstablishedResponse {
|
||||
Complete(Box<EstablishedComplete>),
|
||||
Stream(Box<EstablishedStream>),
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct AttemptRejection {
|
||||
pub request: CodexResponsesRequest,
|
||||
pub trace: ResponseDispatchTrace,
|
||||
pub account_id: Option<String>,
|
||||
pub account: Option<Account>,
|
||||
pub attempt_request: Option<CodexResponsesRequest>,
|
||||
pub attempt: Option<ResponseDispatchAttempt>,
|
||||
pub transport: CodexBackendTransport,
|
||||
pub stream_failure: Option<RejectedStreamFailure>,
|
||||
pub error: ResponseDispatchError,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct RejectedStreamFailure {
|
||||
pub failure: ResponsesSseFailure,
|
||||
pub prefetched: Bytes,
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) enum AttemptApplyOutcome {
|
||||
Continue,
|
||||
Established(EstablishedResponse),
|
||||
Rejected(Box<AttemptRejection>),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(in crate::dispatch) enum AttemptContractError {
|
||||
#[error("the previous attempt observation has not been decided")]
|
||||
DecisionRequired,
|
||||
#[error("attempt runner has already reached a terminal outcome")]
|
||||
Terminal,
|
||||
#[error("decision `{decision}` is invalid for observation `{observation}`")]
|
||||
InvalidDecision {
|
||||
decision: &'static str,
|
||||
observation: &'static str,
|
||||
},
|
||||
}
|
||||
|
||||
pub(super) enum PendingProtocolFailure {
|
||||
InvalidSse(SseError),
|
||||
EmptyStream,
|
||||
NoCommitBoundary,
|
||||
}
|
||||
|
||||
pub(super) struct PendingStreamResponse {
|
||||
pub account: Account,
|
||||
pub attempt_request: CodexResponsesRequest,
|
||||
pub attempt: ResponseDispatchAttempt,
|
||||
pub lease: AccountLease,
|
||||
pub response: CodexBackendStreamingResponse,
|
||||
pub prefetched: Bytes,
|
||||
pub decoder: CanonicalStreamDecoder,
|
||||
pub initial_batch: CanonicalStreamBatch,
|
||||
pub first_failure: Option<ResponsesSseFailure>,
|
||||
pub first_event_ms: i64,
|
||||
}
|
||||
|
||||
pub(super) struct PendingCompleteResponse {
|
||||
pub account: Account,
|
||||
pub attempt_request: CodexResponsesRequest,
|
||||
pub attempt: ResponseDispatchAttempt,
|
||||
pub response: CodexBackendResponse,
|
||||
pub collected: CollectedResponse,
|
||||
}
|
||||
|
||||
pub(super) enum PendingAttempt {
|
||||
NoCandidate,
|
||||
PinnedCandidateUnavailable {
|
||||
account_id: String,
|
||||
kind: PinnedCandidateAcquireFailureKind,
|
||||
},
|
||||
RoutePreparationRejected {
|
||||
account: Account,
|
||||
message: String,
|
||||
},
|
||||
UpstreamFailure {
|
||||
account: Account,
|
||||
attempt_request: CodexResponsesRequest,
|
||||
attempt: ResponseDispatchAttempt,
|
||||
transport: CodexBackendTransport,
|
||||
error: CodexClientError,
|
||||
},
|
||||
ProtocolFailure {
|
||||
account: Account,
|
||||
attempt_request: CodexResponsesRequest,
|
||||
attempt: ResponseDispatchAttempt,
|
||||
transport: CodexBackendTransport,
|
||||
error: PendingProtocolFailure,
|
||||
},
|
||||
CompleteResponse(Box<PendingCompleteResponse>),
|
||||
StreamResponse(Box<PendingStreamResponse>),
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
//! Responses live stream 的唯一结算出口。
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use bytes::Bytes;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
controllers::{ControllerRequestScope, ControllerSet, StreamControllerContext},
|
||||
errors::ResponseDispatchStreamError,
|
||||
failure::sse::{STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE},
|
||||
lifecycle::stream::StreamSummary,
|
||||
transport::canonical::{CanonicalResponseChunk, CanonicalResponseEvent},
|
||||
},
|
||||
fleet::pool::AccountLease,
|
||||
upstream::openai::protocol::sse::{
|
||||
DONE_SSE_FRAME, encode_sse_event, response_failed_sse_data_with_id,
|
||||
},
|
||||
};
|
||||
|
||||
const LEASE_FINALIZE_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
|
||||
/// 消费 stream context,类型层面保证每个 live loop 只能结算一次。
|
||||
pub(in crate::dispatch) struct StreamFinalizer {
|
||||
controllers: ControllerSet,
|
||||
controller_scope: ControllerRequestScope,
|
||||
controller_context: StreamControllerContext,
|
||||
request_input: Vec<serde_json::Value>,
|
||||
continued_from_previous_response: bool,
|
||||
account_lease: AccountLease,
|
||||
}
|
||||
|
||||
impl StreamFinalizer {
|
||||
pub(in crate::dispatch) fn new(
|
||||
controllers: ControllerSet,
|
||||
controller_scope: ControllerRequestScope,
|
||||
controller_context: StreamControllerContext,
|
||||
request_input: Vec<serde_json::Value>,
|
||||
continued_from_previous_response: bool,
|
||||
account_lease: AccountLease,
|
||||
) -> Self {
|
||||
Self {
|
||||
controllers,
|
||||
controller_scope,
|
||||
controller_context,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
account_lease,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn response_headers(&self) -> &[(String, String)] {
|
||||
&self.controller_context.response_metadata.client_headers
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn started_at(&self) -> Instant {
|
||||
self.controller_context.started_at
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn account_id(&self) -> &str {
|
||||
&self.controller_context.account_id
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn request_input(&self) -> &[serde_json::Value] {
|
||||
&self.request_input
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn continued_from_previous_response(&self) -> bool {
|
||||
self.continued_from_previous_response
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn finalize(
|
||||
self,
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
|
||||
mut summary: StreamSummary,
|
||||
) {
|
||||
let Self {
|
||||
controllers,
|
||||
controller_scope,
|
||||
controller_context,
|
||||
account_lease,
|
||||
..
|
||||
} = self;
|
||||
let body = String::from_utf8_lossy(&summary.body).into_owned();
|
||||
controllers
|
||||
.leave_stream(controller_scope, controller_context, &summary, &body)
|
||||
.await;
|
||||
let mut lease_completion = Box::pin(account_lease.complete());
|
||||
if timeout(LEASE_FINALIZE_TIMEOUT, lease_completion.as_mut())
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
tracing::warn!(
|
||||
timeout_ms = LEASE_FINALIZE_TIMEOUT.as_millis(),
|
||||
"Response account lease is continuing finalization in the background"
|
||||
);
|
||||
// 超时只约束客户端收尾,不取消已经开始的 slot 释放与请求用量持久化。
|
||||
drop(tokio::spawn(lease_completion));
|
||||
}
|
||||
finish_client_stream(sender, event_sender, &mut summary).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn finish_client_stream(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
|
||||
summary: &mut StreamSummary,
|
||||
) {
|
||||
if !summary.terminal.should_finish_client_stream() {
|
||||
return;
|
||||
}
|
||||
if flush_terminal_chunks(sender, event_sender, &mut summary.terminal_chunks)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
if let Some(detail) = summary.terminal.synthetic_failure_detail() {
|
||||
if append_separator_if_needed(sender, event_sender, &mut summary.body)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
let failure = response_failed_sse_data_with_id(
|
||||
summary.last_response_id.as_deref(),
|
||||
"server_error",
|
||||
STREAM_DISCONNECTED_CODE,
|
||||
&synthetic_failure_message(detail),
|
||||
);
|
||||
let bytes = Bytes::from(encode_sse_event("response.failed", &failure.to_string()));
|
||||
if append_client_chunk(
|
||||
sender,
|
||||
event_sender,
|
||||
&mut summary.body,
|
||||
CanonicalResponseChunk::new(
|
||||
bytes,
|
||||
vec![CanonicalResponseEvent::proxy_failure(failure)],
|
||||
),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if !body_has_done(&summary.body) {
|
||||
let _ = append_client_chunk(
|
||||
sender,
|
||||
event_sender,
|
||||
&mut summary.body,
|
||||
CanonicalResponseChunk::new(Bytes::from_static(DONE_SSE_FRAME.as_bytes()), Vec::new()),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn flush_terminal_chunks(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
|
||||
terminal_chunks: &mut Vec<CanonicalResponseChunk>,
|
||||
) -> Result<(), ()> {
|
||||
for chunk in terminal_chunks.drain(..) {
|
||||
let (bytes, events) = chunk.into_parts();
|
||||
let _ = event_sender.send(events);
|
||||
sender.send(Ok(bytes)).await.map_err(|_| ())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn append_separator_if_needed(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
|
||||
body: &mut Vec<u8>,
|
||||
) -> Result<(), ()> {
|
||||
let separator = if body.is_empty()
|
||||
|| body.ends_with(b"\n\n")
|
||||
|| body.ends_with(b"\r\n\r\n")
|
||||
|| body.ends_with(b"\r\r")
|
||||
{
|
||||
return Ok(());
|
||||
} else if body.ends_with(b"\n") || body.ends_with(b"\r") {
|
||||
Bytes::from_static(b"\n")
|
||||
} else {
|
||||
Bytes::from_static(b"\n\n")
|
||||
};
|
||||
append_client_chunk(
|
||||
sender,
|
||||
event_sender,
|
||||
body,
|
||||
CanonicalResponseChunk::new(separator, Vec::new()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn append_client_chunk(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
|
||||
body: &mut Vec<u8>,
|
||||
chunk: CanonicalResponseChunk,
|
||||
) -> Result<(), ()> {
|
||||
body.extend_from_slice(chunk.bytes());
|
||||
let (bytes, events) = chunk.into_parts();
|
||||
let _ = event_sender.send(events);
|
||||
sender.send(Ok(bytes)).await.map_err(|_| ())
|
||||
}
|
||||
|
||||
fn body_has_done(body: &[u8]) -> bool {
|
||||
String::from_utf8_lossy(body)
|
||||
.trim_end_matches(['\r', '\n'])
|
||||
.ends_with(DONE_SSE_FRAME.trim_end_matches(['\r', '\n']))
|
||||
}
|
||||
|
||||
fn synthetic_failure_message(detail: &str) -> String {
|
||||
if detail.trim().is_empty() {
|
||||
STREAM_DISCONNECTED_MESSAGE.to_string()
|
||||
} else {
|
||||
format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
//! `v1/*` 请求的生命周期基础类型与入口准备。
|
||||
|
||||
pub(in crate::dispatch) mod attempt;
|
||||
pub(in crate::dispatch) mod contract;
|
||||
pub(in crate::dispatch) mod finalizer;
|
||||
pub(in crate::dispatch) mod pipeline;
|
||||
pub(in crate::dispatch) mod request;
|
||||
pub(in crate::dispatch) mod stream;
|
||||
pub(in crate::dispatch) mod trace;
|
||||
|
||||
pub(in crate::dispatch) use request::{
|
||||
RequestContext, RequestEnterDependencies, RequestMode, enter_request,
|
||||
};
|
||||
@@ -0,0 +1,53 @@
|
||||
//! Attempt 生命周期的唯一静态 pipeline。
|
||||
|
||||
use crate::{
|
||||
dispatch::{controllers::ControllerRequestScope, routing::candidates::AccountAttemptLedger},
|
||||
upstream::openai::protocol::responses::CodexResponsesRequest,
|
||||
};
|
||||
|
||||
use super::{
|
||||
attempt::{AttemptMode, AttemptRunner, AttemptRunnerDependencies, AttemptStep},
|
||||
contract::{AttemptApplyOutcome, AttemptContractError, AttemptRejection, EstablishedResponse},
|
||||
};
|
||||
|
||||
pub(in crate::dispatch) enum AttemptPipelineOutcome {
|
||||
Established(EstablishedResponse),
|
||||
Rejected(Box<AttemptRejection>),
|
||||
}
|
||||
|
||||
/// 执行完整的 attempt/retry/commit 状态机。
|
||||
///
|
||||
/// mode 适配器只能消费最终结果,不能自行复制 retry loop 或越过 commit 边界。
|
||||
pub(in crate::dispatch) async fn run_attempt_pipeline(
|
||||
dependencies: AttemptRunnerDependencies<'_>,
|
||||
mode: AttemptMode,
|
||||
request: CodexResponsesRequest,
|
||||
controller_scope: ControllerRequestScope,
|
||||
candidates: AccountAttemptLedger,
|
||||
) -> Result<AttemptPipelineOutcome, AttemptContractError> {
|
||||
let controllers = dependencies.controllers;
|
||||
let mut runner = AttemptRunner::new(dependencies, mode, request, controller_scope, candidates);
|
||||
|
||||
loop {
|
||||
let outcome = match runner.next().await? {
|
||||
AttemptStep::Open(attempt) => {
|
||||
let decision = controllers
|
||||
.handle_attempt(attempt.controller_scope()?, attempt.observation())
|
||||
.await;
|
||||
attempt.apply(decision).await?
|
||||
}
|
||||
AttemptStep::Committed(attempt) => {
|
||||
return Ok(AttemptPipelineOutcome::Established(attempt.accept()));
|
||||
}
|
||||
};
|
||||
match outcome {
|
||||
AttemptApplyOutcome::Continue => {}
|
||||
AttemptApplyOutcome::Established(established) => {
|
||||
return Ok(AttemptPipelineOutcome::Established(established));
|
||||
}
|
||||
AttemptApplyOutcome::Rejected(rejection) => {
|
||||
return Ok(AttemptPipelineOutcome::Rejected(rejection));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
//! Responses 请求进入 dispatch 的共享生命周期阶段。
|
||||
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::{AccountIdentityService, prepare_variant_identity},
|
||||
controllers::{ControllerEnter, ControllerRequestScope, ControllerSet},
|
||||
routing::candidates::AccountAttemptLedger,
|
||||
},
|
||||
fleet::pool::{AccountAcquireRequest, AccountPoolService},
|
||||
models::service::ModelService,
|
||||
upstream::openai::protocol::responses::CodexResponsesRequest,
|
||||
};
|
||||
|
||||
/// Responses 请求的执行模式。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum RequestMode {
|
||||
Complete,
|
||||
Stream,
|
||||
}
|
||||
|
||||
impl RequestMode {
|
||||
fn stream(self) -> bool {
|
||||
matches!(self, Self::Stream)
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求 enter 阶段所需的窄依赖集合。
|
||||
pub(in crate::dispatch) struct RequestEnterDependencies<'a> {
|
||||
pub account_pool: &'a AccountPoolService,
|
||||
pub models: &'a ModelService,
|
||||
pub account_identity: &'a AccountIdentityService,
|
||||
pub controllers: ControllerSet,
|
||||
}
|
||||
|
||||
/// 已完成请求级 enter 的 Responses 上下文。
|
||||
///
|
||||
/// 该类型只承载请求、会话恢复、会话策略和候选快照;账号 attempt
|
||||
/// 与流生命周期由后续 pipeline 接管。
|
||||
pub(in crate::dispatch) struct RequestContext {
|
||||
pub request: CodexResponsesRequest,
|
||||
pub display_model: String,
|
||||
pub compact: bool,
|
||||
pub tuple_schema: Option<Value>,
|
||||
pub image_generation_requested: bool,
|
||||
pub controllers: ControllerSet,
|
||||
pub controller_scope: ControllerRequestScope,
|
||||
pub candidates: AccountAttemptLedger,
|
||||
}
|
||||
|
||||
/// 执行所有 Responses 模式共享的请求 enter 阶段。
|
||||
pub(in crate::dispatch) async fn enter_request(
|
||||
dependencies: RequestEnterDependencies<'_>,
|
||||
mut request: CodexResponsesRequest,
|
||||
requested_model: &str,
|
||||
mode: RequestMode,
|
||||
) -> RequestContext {
|
||||
let catalog = dependencies.models.catalog().await;
|
||||
let display_model = catalog.resolve_model_id(requested_model);
|
||||
request.set_model(display_model.clone());
|
||||
debug_assert_eq!(request.stream(), mode.stream());
|
||||
|
||||
let compact = request.semantics().compact;
|
||||
let tuple_schema = request.tuple_schema.clone();
|
||||
let image_generation_requested = request.expects_image_generation();
|
||||
let now = Utc::now();
|
||||
prepare_variant_identity(&mut request);
|
||||
dependencies
|
||||
.account_identity
|
||||
.prepare_local_identity(&mut request);
|
||||
let controllers = dependencies.controllers;
|
||||
let ControllerEnter {
|
||||
scope,
|
||||
preferred_account_id,
|
||||
excluded_account_ids,
|
||||
} = controllers.enter(&request, now).await;
|
||||
|
||||
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
|
||||
if let Some(preferred_account_id) = preferred_account_id {
|
||||
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
|
||||
}
|
||||
let candidates = AccountAttemptLedger::freeze(
|
||||
dependencies.account_pool,
|
||||
&acquire_request,
|
||||
&excluded_account_ids,
|
||||
)
|
||||
.await;
|
||||
|
||||
RequestContext {
|
||||
request,
|
||||
display_model,
|
||||
compact,
|
||||
tuple_schema,
|
||||
image_generation_requested,
|
||||
controllers,
|
||||
controller_scope: scope,
|
||||
candidates,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
//! Responses live stream 的规范事件与终态。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::transport::canonical::{CanonicalResponseChunk, CanonicalStreamTerminal},
|
||||
upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure},
|
||||
};
|
||||
|
||||
/// live stream 离开事件循环的唯一终态。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum StreamTerminal {
|
||||
Completed { response: Value },
|
||||
Incomplete { response: Value },
|
||||
Failed { failure: ResponsesSseFailure },
|
||||
UpstreamClosed,
|
||||
UpstreamError { detail: String },
|
||||
ProtocolError { detail: String },
|
||||
CaptureLimitExceeded,
|
||||
Cancelled,
|
||||
DownstreamClosed,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// live loop 交给唯一 finalizer 的完整事实快照。
|
||||
pub(in crate::dispatch) struct StreamSummary {
|
||||
pub terminal: StreamTerminal,
|
||||
pub body: Vec<u8>,
|
||||
pub terminal_chunks: Vec<CanonicalResponseChunk>,
|
||||
pub first_token_ms: Option<i64>,
|
||||
pub first_reasoning_ms: Option<i64>,
|
||||
pub first_text_ms: Option<i64>,
|
||||
pub first_event_ms: i64,
|
||||
pub usage: Option<TokenUsage>,
|
||||
pub last_response_id: Option<String>,
|
||||
}
|
||||
|
||||
impl From<CanonicalStreamTerminal> for StreamTerminal {
|
||||
fn from(terminal: CanonicalStreamTerminal) -> Self {
|
||||
match terminal {
|
||||
CanonicalStreamTerminal::Completed(response) => Self::Completed { response },
|
||||
CanonicalStreamTerminal::Incomplete(response) => Self::Incomplete { response },
|
||||
CanonicalStreamTerminal::Failed(failure) => Self::Failed { failure },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamTerminal {
|
||||
/// 只有上游没有给出业务终态时,代理才需要合成 `response.failed`。
|
||||
pub(in crate::dispatch) fn synthetic_failure_detail(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::UpstreamClosed => Some(""),
|
||||
Self::UpstreamError { detail } | Self::ProtocolError { detail } => Some(detail),
|
||||
Self::CaptureLimitExceeded => {
|
||||
Some("upstream response exceeded the 16 MiB proxy capture limit")
|
||||
}
|
||||
Self::Completed { .. }
|
||||
| Self::Incomplete { .. }
|
||||
| Self::Failed { .. }
|
||||
| Self::Cancelled
|
||||
| Self::DownstreamClosed
|
||||
| Self::Shutdown => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn should_finish_client_stream(&self) -> bool {
|
||||
!matches!(
|
||||
self,
|
||||
Self::Cancelled | Self::DownstreamClosed | Self::Shutdown
|
||||
)
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Completed { .. } => "completed",
|
||||
Self::Incomplete { .. } => "incomplete",
|
||||
Self::Failed { .. } => "failed",
|
||||
Self::UpstreamClosed => "upstream_closed",
|
||||
Self::UpstreamError { .. } => "upstream_error",
|
||||
Self::ProtocolError { .. } => "protocol_error",
|
||||
Self::CaptureLimitExceeded => "capture_limit_exceeded",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::DownstreamClosed => "downstream_closed",
|
||||
Self::Shutdown => "shutdown",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,6 @@
|
||||
//! Request-level trace state for Responses dispatch.
|
||||
//! Responses 请求级账号尝试轨迹。
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(in crate::dispatch) struct ResponseDispatchTrace {
|
||||
@@ -10,6 +12,7 @@ pub(in crate::dispatch) struct ResponseDispatchTrace {
|
||||
pub(in crate::dispatch) struct ResponseDispatchAttempt {
|
||||
index: i64,
|
||||
account_id: String,
|
||||
started_at: Instant,
|
||||
}
|
||||
|
||||
impl ResponseDispatchTrace {
|
||||
@@ -20,6 +23,7 @@ impl ResponseDispatchTrace {
|
||||
let attempt = ResponseDispatchAttempt {
|
||||
index: self.next_attempt_index,
|
||||
account_id: account_id.to_string(),
|
||||
started_at: Instant::now(),
|
||||
};
|
||||
self.next_attempt_index += 1;
|
||||
self.attempts.push(attempt.clone());
|
||||
@@ -39,4 +43,8 @@ impl ResponseDispatchAttempt {
|
||||
pub(in crate::dispatch) fn account_id(&self) -> &str {
|
||||
&self.account_id
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn started_at(&self) -> Instant {
|
||||
self.started_at
|
||||
}
|
||||
}
|
||||
@@ -1,10 +1,11 @@
|
||||
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
|
||||
|
||||
pub mod affinity;
|
||||
mod attempts;
|
||||
pub(crate) mod controllers;
|
||||
pub mod errors;
|
||||
mod recording;
|
||||
pub mod recovery;
|
||||
pub(crate) mod failure;
|
||||
pub(crate) mod lifecycle;
|
||||
pub(crate) mod routing;
|
||||
pub(crate) mod service;
|
||||
pub(crate) mod stream;
|
||||
pub mod upstream_call;
|
||||
pub(crate) mod transport;
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
//! 可换号上游失败的账号风险剥离。
|
||||
|
||||
use chrono::Utc;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
errors::{
|
||||
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
|
||||
is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body,
|
||||
upstream_error_http_status,
|
||||
},
|
||||
recovery::{
|
||||
auth::{auth_failure_account_status, is_auth_upstream_error},
|
||||
cloudflare::{
|
||||
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
|
||||
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
|
||||
CloudflareRecovery,
|
||||
},
|
||||
exhaustion::AccountExhaustionTracker,
|
||||
},
|
||||
stream::sse_failure::{
|
||||
auth_sse_failure_account_status, is_auth_sse_failure, is_model_unsupported_sse_failure,
|
||||
is_quota_exhausted_sse_failure, sse_failure_error_body, stream_failure_http_status,
|
||||
},
|
||||
},
|
||||
fleet::{account::AccountStatus, pool::AccountPoolService},
|
||||
upstream::openai::{
|
||||
protocol::responses::ResponsesSseFailure,
|
||||
transport::{is_banned_upstream_error, CodexClientError},
|
||||
},
|
||||
};
|
||||
|
||||
/// 处理需要隔离当前账号并交回调度器选择下一候选的上游失败。
|
||||
///
|
||||
/// 本函数不选择下一个账号,也不实现固定轮换策略;调用方继续使用请求级候选账本,
|
||||
/// 因而始终服从当前运行时调度策略。
|
||||
pub(crate) async fn isolate_rotatable_account_failure(
|
||||
account_pool: &AccountPoolService,
|
||||
cloudflare: &CloudflareRecovery,
|
||||
exhausted_accounts: &mut AccountExhaustionTracker,
|
||||
account_id: &str,
|
||||
error: &CodexClientError,
|
||||
) -> bool {
|
||||
if is_rate_limit_upstream_error(error) {
|
||||
exhausted_accounts.record_rate_limited(Some(account_id), upstream_error_body(error));
|
||||
account_pool
|
||||
.mark_quota_limited_until(account_id, rate_limit_cooldown_until(error, Utc::now()))
|
||||
.await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_quota_exhausted_upstream_error(error) {
|
||||
exhausted_accounts.record_quota_exhausted(Some(account_id), upstream_error_body(error));
|
||||
account_pool
|
||||
.set_status(account_id, AccountStatus::QuotaExhausted)
|
||||
.await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_auth_upstream_error(error) {
|
||||
let account_status = auth_failure_account_status(error);
|
||||
exhausted_accounts.record_auth_failure(
|
||||
Some(account_id),
|
||||
account_status,
|
||||
upstream_error_body(error),
|
||||
Some(upstream_error_http_status(error)),
|
||||
);
|
||||
account_pool.set_status(account_id, account_status).await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_cloudflare_challenge_upstream_error(error) {
|
||||
exhausted_accounts
|
||||
.record_cloudflare_challenge(Some(account_id), cloudflare_challenge_error_message());
|
||||
cloudflare.apply_challenge(account_pool, account_id).await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_cloudflare_path_block_upstream_error(error) {
|
||||
exhausted_accounts.record_cloudflare_path_blocked(
|
||||
Some(account_id),
|
||||
cloudflare_path_block_error_message(),
|
||||
);
|
||||
cloudflare.apply_path_block(account_pool, account_id).await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_model_unsupported_upstream_error(error) {
|
||||
exhausted_accounts.record_model_unsupported(Some(account_id), upstream_error_body(error));
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_banned_upstream_error(error) {
|
||||
exhausted_accounts.record_auth_failure(
|
||||
Some(account_id),
|
||||
AccountStatus::Banned,
|
||||
upstream_error_body(error),
|
||||
Some(upstream_error_http_status(error)),
|
||||
);
|
||||
account_pool
|
||||
.set_status(account_id, AccountStatus::Banned)
|
||||
.await;
|
||||
return true;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// 处理 SSE `response.failed` 中需要隔离当前账号的失败。
|
||||
pub(crate) async fn isolate_sse_account_failure(
|
||||
account_pool: &AccountPoolService,
|
||||
exhausted_accounts: &mut AccountExhaustionTracker,
|
||||
account_id: &str,
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> bool {
|
||||
if is_model_unsupported_sse_failure(failure) {
|
||||
exhausted_accounts
|
||||
.record_model_unsupported(Some(account_id), sse_failure_error_body(failure));
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_quota_exhausted_sse_failure(failure) {
|
||||
exhausted_accounts.record_quota_exhausted(Some(account_id), failure.message.clone());
|
||||
account_pool
|
||||
.set_status(account_id, AccountStatus::QuotaExhausted)
|
||||
.await;
|
||||
return true;
|
||||
}
|
||||
|
||||
if is_auth_sse_failure(failure) {
|
||||
let account_status = auth_sse_failure_account_status(failure);
|
||||
exhausted_accounts.record_auth_failure(
|
||||
Some(account_id),
|
||||
account_status,
|
||||
sse_failure_error_body(failure),
|
||||
Some(stream_failure_http_status(failure)),
|
||||
);
|
||||
account_pool.set_status(account_id, account_status).await;
|
||||
return true;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
@@ -1,22 +0,0 @@
|
||||
//! 上游认证失败分类。
|
||||
|
||||
use crate::{
|
||||
fleet::account::AccountStatus,
|
||||
upstream::openai::transport::{is_banned_auth_signal, CodexClientError},
|
||||
};
|
||||
|
||||
pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, .. } if status.as_u16() == 401
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus {
|
||||
match error {
|
||||
CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => {
|
||||
AccountStatus::Banned
|
||||
}
|
||||
_ => AccountStatus::Expired,
|
||||
}
|
||||
}
|
||||
@@ -1,146 +0,0 @@
|
||||
//! Cloudflare challenge and path-block recovery shared by dispatch routes.
|
||||
|
||||
use chrono::Utc;
|
||||
|
||||
use crate::{
|
||||
fleet::{
|
||||
account::AccountStatus,
|
||||
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
upstream::openai::transport::CodexClientError,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CloudflareRecovery {
|
||||
path_block_tracker: CloudflarePathBlockTracker,
|
||||
challenge_tracker: CloudflareChallengeCooldownTracker,
|
||||
cookie_store: PgCookieStore,
|
||||
}
|
||||
|
||||
impl CloudflareRecovery {
|
||||
pub fn new(cookie_store: PgCookieStore) -> Self {
|
||||
Self {
|
||||
path_block_tracker: CloudflarePathBlockTracker::default(),
|
||||
challenge_tracker: CloudflareChallengeCooldownTracker::default(),
|
||||
cookie_store,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option<String> {
|
||||
self.cookie_store
|
||||
.cookie_header_for_request(account_id, "chatgpt.com", path)
|
||||
.await
|
||||
.ok()?
|
||||
}
|
||||
|
||||
pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) {
|
||||
for header in headers {
|
||||
if let Err(error) = self
|
||||
.cookie_store
|
||||
.capture_set_cookie(account_id, header)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist upstream set-cookie header"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) {
|
||||
let now = Utc::now();
|
||||
let cooldown = self
|
||||
.challenge_tracker
|
||||
.record_challenge(account_id, now)
|
||||
.await;
|
||||
account_pool
|
||||
.set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until)
|
||||
.await;
|
||||
if let Err(error) = self
|
||||
.cookie_store
|
||||
.expire_account_cookies_at(account_id, cooldown.cooldown_until)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist Cloudflare challenge cookie cleanup deadline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) {
|
||||
self.delete_account_cookies(account_id, "Cloudflare path-block")
|
||||
.await;
|
||||
let now = Utc::now();
|
||||
self.path_block_tracker
|
||||
.record_path_block(account_id, now)
|
||||
.await;
|
||||
if self
|
||||
.path_block_tracker
|
||||
.should_disable(account_id, now)
|
||||
.await
|
||||
{
|
||||
account_pool
|
||||
.set_status(account_id, AccountStatus::Disabled)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn reset_account_recovery(&self, account_id: &str) {
|
||||
self.path_block_tracker.reset(account_id).await;
|
||||
self.challenge_tracker.reset(account_id).await;
|
||||
}
|
||||
|
||||
async fn delete_account_cookies(&self, account_id: &str, reason: &str) {
|
||||
delete_account_cookies(&self.cookie_store, account_id, reason).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) {
|
||||
if let Err(error) = cookie_store.delete_account_cookies(account_id).await {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
reason,
|
||||
error = %error,
|
||||
"failed to delete account cookies after Cloudflare recovery signal"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_cloudflare_challenge_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if status.as_u16() == 403 && is_cloudflare_challenge_signal(body)
|
||||
)
|
||||
}
|
||||
|
||||
pub fn is_cloudflare_path_block_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if status.as_u16() == 404 && body.trim().is_empty()
|
||||
)
|
||||
}
|
||||
|
||||
pub fn cloudflare_challenge_error_message() -> &'static str {
|
||||
"Upstream blocked the request (Cloudflare challenge)"
|
||||
}
|
||||
|
||||
pub fn cloudflare_path_block_error_message() -> &'static str {
|
||||
"Upstream blocked the request (Cloudflare path-block)"
|
||||
}
|
||||
|
||||
fn is_cloudflare_challenge_signal(value: &str) -> bool {
|
||||
let value = value.to_ascii_lowercase();
|
||||
value.contains("cf-mitigated")
|
||||
|| value.contains("cf-chl-bypass")
|
||||
|| value.contains("_cf_chl")
|
||||
|| value.contains("cf_chl")
|
||||
|| value.contains("attention required")
|
||||
|| value.contains("just a moment")
|
||||
}
|
||||
@@ -1,293 +0,0 @@
|
||||
//! previous response 所有权与完整上下文重放。
|
||||
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::affinity::{
|
||||
service::{MAX_REPLAY_DEPTH, MAX_REPLAY_SESSION_BYTES, MAX_REPLAY_SNAPSHOT_BYTES},
|
||||
types::ResponseReplaySnapshot,
|
||||
SessionAffinityEntry, SessionAffinityService,
|
||||
},
|
||||
upstream::openai::protocol::responses::{CodexResponsesRequest, StreamCommitPolicy},
|
||||
};
|
||||
|
||||
/// 当前请求的 previous response 来源。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum HistorySource {
|
||||
None,
|
||||
Managed,
|
||||
ExternalUnknown,
|
||||
}
|
||||
|
||||
/// 请求级 previous response 恢复计划。
|
||||
pub(in crate::dispatch) struct HistoryRecoveryPlan {
|
||||
source: HistorySource,
|
||||
managed_entry: Option<SessionAffinityEntry>,
|
||||
replay_input: Option<Vec<Value>>,
|
||||
external_attempted_account_id: Option<String>,
|
||||
replay_owner_account: bool,
|
||||
}
|
||||
|
||||
impl HistoryRecoveryPlan {
|
||||
pub(in crate::dispatch) async fn load(
|
||||
session_affinity: &SessionAffinityService,
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Self {
|
||||
let Some(previous_response_id) = request.previous_response_id() else {
|
||||
return Self {
|
||||
source: HistorySource::None,
|
||||
managed_entry: None,
|
||||
replay_input: None,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
};
|
||||
};
|
||||
let managed_entry = session_affinity
|
||||
.lookup(previous_response_id, Utc::now())
|
||||
.await;
|
||||
let replay_input = match managed_entry.as_ref() {
|
||||
Some(entry) => {
|
||||
session_affinity
|
||||
.replay_input(previous_response_id, entry, Utc::now())
|
||||
.await
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
Self {
|
||||
source: if managed_entry.is_some() {
|
||||
HistorySource::Managed
|
||||
} else {
|
||||
HistorySource::ExternalUnknown
|
||||
},
|
||||
managed_entry,
|
||||
replay_input,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn preferred_account_id(&self) -> Option<&str> {
|
||||
self.managed_entry
|
||||
.as_ref()
|
||||
.map(|entry| entry.account_id.as_str())
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn prepare_attempt(
|
||||
&mut self,
|
||||
original: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Option<CodexResponsesRequest> {
|
||||
match self.source {
|
||||
HistorySource::None => Some(original.clone()),
|
||||
HistorySource::ExternalUnknown => match self.external_attempted_account_id.as_deref() {
|
||||
None => {
|
||||
self.external_attempted_account_id = Some(account_id.to_string());
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope = Some(
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown,
|
||||
);
|
||||
Some(request)
|
||||
}
|
||||
Some(first_account_id) if first_account_id == account_id => {
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope = Some(
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown,
|
||||
);
|
||||
Some(request)
|
||||
}
|
||||
Some(_) => None,
|
||||
},
|
||||
HistorySource::Managed => {
|
||||
let entry = self.managed_entry.as_ref()?;
|
||||
if entry.account_id == account_id && !self.replay_owner_account {
|
||||
let mut request = original.clone();
|
||||
request.local_conversation_id = Some(entry.conversation_id.clone());
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope = Some(entry.continuation_scope);
|
||||
if request.turn_state.as_deref().is_none_or(str::is_empty) {
|
||||
request.turn_state.clone_from(&entry.turn_state);
|
||||
}
|
||||
Some(request)
|
||||
} else {
|
||||
self.replay_input.as_ref().map(|replay_input| {
|
||||
let mut request = full_replay_request(original, replay_input);
|
||||
request.local_conversation_id = Some(entry.conversation_id.clone());
|
||||
request
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// previous ID 在原连接不可用时,切换为同账号完整重放。
|
||||
pub(in crate::dispatch) fn recover_managed_history(&mut self, account_id: &str) -> bool {
|
||||
let Some(entry) = self.managed_entry.as_ref() else {
|
||||
return false;
|
||||
};
|
||||
if entry.account_id != account_id
|
||||
|| self.replay_input.is_none()
|
||||
|| self.replay_owner_account
|
||||
{
|
||||
return false;
|
||||
}
|
||||
self.replay_owner_account = true;
|
||||
true
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn can_failover(&self) -> bool {
|
||||
!matches!(self.source, HistorySource::ExternalUnknown)
|
||||
&& self
|
||||
.managed_entry
|
||||
.as_ref()
|
||||
.is_none_or(|_| self.replay_input.is_some())
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn is_external_unknown(&self) -> bool {
|
||||
self.source == HistorySource::ExternalUnknown
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn completed_replay(
|
||||
&self,
|
||||
original: &CodexResponsesRequest,
|
||||
output: &[Value],
|
||||
) -> Option<ResponseReplaySnapshot> {
|
||||
let (parent_response_id, parent_depth, parent_bytes) = match self.source {
|
||||
HistorySource::None => (None, 0, 0),
|
||||
HistorySource::Managed => {
|
||||
self.replay_input.as_ref()?;
|
||||
let replay = self.managed_entry.as_ref()?.replay.as_ref()?;
|
||||
(
|
||||
original.previous_response_id().map(ToString::to_string),
|
||||
replay.depth,
|
||||
replay.total_bytes,
|
||||
)
|
||||
}
|
||||
HistorySource::ExternalUnknown => return None,
|
||||
};
|
||||
let depth = parent_depth.checked_add(1)?;
|
||||
if depth > MAX_REPLAY_DEPTH {
|
||||
tracing::warn!(depth, "response replay depth limit reached");
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut turn_input = original.input().to_vec();
|
||||
let mut turn_output = output.to_vec();
|
||||
sanitize_replay_items(&mut turn_input);
|
||||
sanitize_replay_items(&mut turn_output);
|
||||
let Some(node_bytes) = limited_replay_json_bytes(&turn_input, &turn_output) else {
|
||||
tracing::warn!("response replay snapshot byte limit reached");
|
||||
return None;
|
||||
};
|
||||
let total_bytes = parent_bytes.checked_add(node_bytes)?;
|
||||
if total_bytes > MAX_REPLAY_SESSION_BYTES {
|
||||
tracing::warn!(
|
||||
node_bytes,
|
||||
total_bytes,
|
||||
"response replay session byte limit reached"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
Some(ResponseReplaySnapshot {
|
||||
parent_response_id,
|
||||
turn_input,
|
||||
turn_output,
|
||||
depth,
|
||||
total_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn conversation_id<'a>(
|
||||
&'a self,
|
||||
original: &'a CodexResponsesRequest,
|
||||
) -> Option<&'a str> {
|
||||
self.managed_entry
|
||||
.as_ref()
|
||||
.map(|entry| entry.conversation_id.as_str())
|
||||
.or_else(|| non_empty(original.local_conversation_id.as_deref()))
|
||||
}
|
||||
}
|
||||
|
||||
fn full_replay_request(
|
||||
original: &CodexResponsesRequest,
|
||||
replay_input: &[Value],
|
||||
) -> CodexResponsesRequest {
|
||||
let mut request = original.clone();
|
||||
let mut input = replay_input.to_vec();
|
||||
input.extend_from_slice(original.input());
|
||||
sanitize_replay_items(&mut input);
|
||||
request.set_input(input);
|
||||
request.set_previous_response_id(None);
|
||||
request.turn_state = None;
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request
|
||||
}
|
||||
|
||||
fn sanitize_replay_items(values: &mut [Value]) {
|
||||
for value in values {
|
||||
match value {
|
||||
Value::Array(values) => sanitize_replay_items(values),
|
||||
Value::Object(object) => {
|
||||
object.remove("id");
|
||||
object.remove("encrypted_content");
|
||||
for value in object.values_mut() {
|
||||
strip_encrypted_content(std::slice::from_mut(value));
|
||||
}
|
||||
}
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_encrypted_content(values: &mut [Value]) {
|
||||
for value in values {
|
||||
match value {
|
||||
Value::Array(values) => strip_encrypted_content(values),
|
||||
Value::Object(object) => {
|
||||
object.remove("encrypted_content");
|
||||
for value in object.values_mut() {
|
||||
strip_encrypted_content(std::slice::from_mut(value));
|
||||
}
|
||||
}
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn limited_replay_json_bytes(input: &[Value], output: &[Value]) -> Option<u64> {
|
||||
let mut writer = ReplaySizeWriter::default();
|
||||
serde_json::to_writer(&mut writer, &(input, output)).ok()?;
|
||||
Some(writer.written)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReplaySizeWriter {
|
||||
written: u64,
|
||||
}
|
||||
|
||||
impl std::io::Write for ReplaySizeWriter {
|
||||
fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
|
||||
let next = self
|
||||
.written
|
||||
.checked_add(buffer.len() as u64)
|
||||
.filter(|next| *next <= MAX_REPLAY_SNAPSHOT_BYTES)
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::FileTooLarge,
|
||||
"response replay snapshot byte limit reached",
|
||||
)
|
||||
})?;
|
||||
self.written = next;
|
||||
Ok(buffer.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
//! 调度失败恢复策略。
|
||||
|
||||
pub(crate) mod account_failure;
|
||||
pub(crate) mod auth;
|
||||
pub mod cloudflare;
|
||||
pub(crate) mod exhaustion;
|
||||
pub(crate) mod history;
|
||||
@@ -3,11 +3,14 @@
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::time::{Duration, Instant, timeout_at};
|
||||
|
||||
use crate::fleet::pool::{
|
||||
AccountAcquireRequest, AccountCandidateLease, AccountLease, AccountPoolService,
|
||||
};
|
||||
|
||||
const ACCOUNT_CANDIDATE_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
/// 请求开始时冻结的完整候选集,以及候选的最终处理状态。
|
||||
pub(in crate::dispatch) struct AccountAttemptLedger {
|
||||
model: String,
|
||||
@@ -16,14 +19,17 @@ pub(in crate::dispatch) struct AccountAttemptLedger {
|
||||
busy: VecDeque<String>,
|
||||
attempted: BTreeSet<String>,
|
||||
state_excluded: BTreeSet<String>,
|
||||
busy_wait_deadline: Instant,
|
||||
}
|
||||
|
||||
impl AccountAttemptLedger {
|
||||
pub(in crate::dispatch) async fn freeze(
|
||||
account_pool: &AccountPoolService,
|
||||
request: &AccountAcquireRequest,
|
||||
excluded_account_ids: &BTreeSet<String>,
|
||||
) -> Self {
|
||||
let candidate_ids = account_pool.candidate_snapshot(request).await;
|
||||
let mut candidate_ids = account_pool.candidate_snapshot(request).await;
|
||||
candidate_ids.retain(|account_id| !excluded_account_ids.contains(account_id));
|
||||
Self {
|
||||
model: request.model.clone(),
|
||||
pending: candidate_ids.iter().cloned().collect(),
|
||||
@@ -31,6 +37,7 @@ impl AccountAttemptLedger {
|
||||
busy: VecDeque::new(),
|
||||
attempted: BTreeSet::new(),
|
||||
state_excluded: BTreeSet::new(),
|
||||
busy_wait_deadline: Instant::now() + ACCOUNT_CANDIDATE_WAIT_TIMEOUT,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,7 +70,22 @@ impl AccountAttemptLedger {
|
||||
return None;
|
||||
}
|
||||
|
||||
account_pool.wait_for_candidate_change().await;
|
||||
if timeout_at(
|
||||
self.busy_wait_deadline,
|
||||
account_pool.wait_for_candidate_change(),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
tracing::warn!(
|
||||
candidate_count = self.candidate_ids.len(),
|
||||
attempted_count = self.attempted.len(),
|
||||
busy_count = self.busy.len(),
|
||||
timeout_ms = ACCOUNT_CANDIDATE_WAIT_TIMEOUT.as_millis(),
|
||||
"Timed out waiting for a busy account candidate"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
std::mem::swap(&mut self.pending, &mut self.busy);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user