Compare commits

...
15 Commits
Author SHA1 Message Date
ZyyandCodex 90dd5b5a57 chore: bump version to 3.3.32
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 05:15:12 +08:00
ZyyandCodex 5ad9378b33 docs(release): prepare 3.3.32 release notes
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 05:15:08 +08:00
ZyyandCodex 4a413f5bb7 fix(budgets): remove manual charge reconciliation
Settle only observed USD costs and treat missing costs as zero without
blocking client keys. Preserve known retry costs and idempotent settlement.
Remove reconciliation controls, endpoints, pending records, and obsolete
schema fields; migrate legacy unknown charges to zero without changing
recorded usage.

Display daily and weekly amounts with two decimal places while retaining
full precision for accounting and tooltips. Distinguish missing WebSocket
closing handshakes from TCP resets in connection diagnostics.

Validation: workspace tests with PostgreSQL/Redis, targeted regressions,
strict Clippy, rustfmt, frontend lint/build, and migration checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 05:13:06 +08:00
ZyyandCodex 0bab566d04 fix(xai): correct subscription plans and Codex tool handling
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 04:49:22 +08:00
ZyyandCodex 9b19cf083a feat(release): publish release notes from Markdown
Require committed, non-empty notes matching the target release version before
creating a tag. Read notes from the tagged source and place them before the
installation guide and generated release changelog.

Add the 3.3.31 feature summary and keep maintainer instructions in the release
script help and repository skill instead of the deployment README.

Validation: shell syntax and workflow lint passed. Nine isolated checks cover
local publication, stale or missing notes, and release-body rendering.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 03:40:00 +08:00
ZyyandCodex 8b22788746 chore: bump version to 3.3.31
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 03:34:04 +08:00
ZyyandCodex e973fc516e feat(accounts): separate setup from account import
Share account settings across import, single-account editing, and batch editing.
Commit credentials and import settings in the same transaction. Improve modal
actions and fix the flash at the end of light/dark theme transitions.

Reauthenticate each execution against the current key policy so existing
WebSocket connections apply updated RPM, concurrency limits, and enabled state.
Add regression coverage for policy updates and unavailable snapshots.

Validation: all 1,950 backend tests, workspace Clippy, frontend lint/typecheck/build,
migration checks, Compose validation, and workflow lint passed. Completed 24 real
requests with a free account using gpt-5.6-luna to verify daily/weekly budgets,
RPM, concurrency, reconciliation, WebSocket policy updates, and budget persistence
across restarts and usage-log cleanup.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 03:26:48 +08:00
ZyyandCodex 12ecbf574a fix(accounts): fill missing plans from provider quota snapshots
Use explicit plan types from OpenAI quota snapshots to fill missing or
unknown account plans before building account and dashboard display names.
Preserve known plan subtypes and reuse cached quota without extra requests
or changes to persisted credential facts.

Cover quota projection and account list/detail fallback behavior, and
document the backend display contract.

Validation: gateway-admin, OpenAI, and xAI tests (1049 passed), strict
Clippy for affected backend crates, gateway build, and diff whitespace
checks passed. A real account JSON import and upstream quota refresh
confirmed free/Free in list, detail, and cached quota responses.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 02:27:45 +08:00
ZyyandCodex 443bd6ad17 feat: integrate account proxies and daily key budgets (#58)
Enforce daily and seven-day client key budgets with a durable charge ledger,
and route provider requests through each account's configured proxy.

Preserve backend plan display names and Codex personal access token imports.
Keep PAT validation on the selected proxy, preserve credential refreshes
when account egress changes, and include the upstream SOCKS IPv6 fix.

Use the PR migration layout, keep the account proxy and key forms aligned
with existing controls, and document the ownership and settlement rules
in the architecture guide.

Validation: 1942 workspace tests passed with isolated PostgreSQL/Redis;
strict Clippy, rustfmt, frontend lint/build, and migration checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 02:07:28 +08:00
ZyyandCodex c795098a87 feat(openai): support Codex personal access token imports
Accept at- personal access tokens through existing AT and JSON imports,
including the official personal_access_token field. Verify identity with
auth whoami instead of trusting document metadata, preserve optional email,
and exclude PAT credentials from OAuth refresh scheduling.

Expose explicit static PAT validation errors without leaking credentials
or raw upstream responses. Cover rejection, invalid identity, timeouts,
redirect protection, export round trips, and unchanged JWT import behavior.

Validation: OpenAI provider tests (562 passed), gateway-admin tests
(100 passed), strict Clippy for both crates with all targets and features,
workspace rustfmt, and diff whitespace checks passed.
Upstream behavior was tested with mocks; no live PAT was used.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 02:02:51 +08:00
ZyyandCodex 1f0009341d fix(accounts): resolve plan display names in the backend
Map OpenAI plan identifiers to official desktop labels in the provider.
Expose planTypeDisplay in account and dashboard responses while preserving
raw planType values, and render backend labels directly in the frontend.
Show missing plans as unknown and retain raw identifiers in badge tooltips.

Validation: Admin/API tests (368 passed), OpenAI admin tests (29 passed),
strict Clippy for affected backend crates with all targets and features,
frontend production build, targeted ESLint, and diff whitespace checks.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 02:02:51 +08:00
Codex 1359eaa894 feat: add account egress and client key USD budgets
Import sub2api proxy bindings, enforce durable daily and seven-day budgets alongside existing key concurrency, and integrate matching admin controls and audited reconciliation.
2026-09-10 23:41:35 +08:00
ZyyandCodex 74318af570 fix(routing): report missing models with explicit errors
Return model_not_found with HTTP 404 when all scoped provider catalogs
confirm the requested or mapped model is absent. Explain missing mapping
targets separately and preserve unavailable-provider errors for circuit
blocks, empty account scopes, and unsupported operations.

Cover missing models, aliases, empty and unknown catalogs, provider
fallback, and streaming and buffered HTTP error responses.

Validation: gateway-core and gateway-api tests (541 passed), strict Clippy
for both crates with all targets and features, workspace rustfmt, and diff
whitespace checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-10 22:23:42 +08:00
Zyy ddc10d7c1a docs: add Apache 2.0 license and community link 2026-09-10 16:14:44 +08:00
ZyyandCodex 6bc2189e6f refactor(frontend): streamline quota reset dialog and spacing
Show remaining usage limits and an always-visible reset credit list with per-credit reset actions. Use backend-provided credit titles and quota window labels while preserving confirmation and ambiguous-result recovery.

Display credit expiration in UTC+8 without a timezone suffix. Balance modal title and reset section spacing, and remove inline baseline whitespace from account quota summary cells without changing row heights or the borderless design.

Validation: browser layout and quota popover interaction checks, expiration timezone checks, frontend typecheck, targeted ESLint, and diff whitespace checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-10 15:33:43 +08:00
182 changed files with 7990 additions and 2176 deletions
+31 -9
View File
@@ -88,6 +88,13 @@ jobs:
fail("Release tag #{release_tag} does not match #{version_path} tag #{expected_tag}") unless release_tag == expected_tag
fail("#{version_path}: version #{version} is not a valid semver") unless version.match?(SEMVER)
notes_path = "release/notes.md"
fail("#{notes_path}: release notes are required") unless File.file?(notes_path)
notes = File.readlines(notes_path)
unless notes.first&.strip == "# #{release_tag}" && !notes.drop(1).join.strip.empty?
fail("#{notes_path}: expected non-empty notes starting with '# #{release_tag}'")
end
platforms = platforms_data["platforms"]
fail("#{platforms_path}: platforms must be a non-empty array") unless platforms.is_a?(Array) && !platforms.empty?
@@ -643,15 +650,10 @@ jobs:
cosign sign --yes "${images[@]}"
- name: Publish GitHub Release
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
with:
tag_name: ${{ needs.release-info.outputs.tag }}
generate_release_notes: true
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
body: |
- name: Prepare release notes
shell: bash
env:
INSTALLATION_NOTES: |
> [!WARNING]
> **Upgrading across major versions?** Online upgrades across major versions are not supported. Deploy the new major fresh with a new `.runtime/` data directory, then re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
@@ -672,6 +674,26 @@ jobs:
- [GitHub Repository](https://github.com/${{ github.repository }})
- [Installation Guide](https://github.com/${{ github.repository }}/blob/main/deploy/README.md)
run: |
python3 <<'PY'
import os
from pathlib import Path
notes = "\n".join(Path("release/notes.md").read_text().splitlines()[1:]).strip()
installation = os.environ["INSTALLATION_NOTES"].strip()
output = Path(os.environ["RUNNER_TEMP"]) / "release-notes.md"
output.write_text(f"{notes}\n\n{installation}\n")
PY
- name: Publish GitHub Release
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
with:
tag_name: ${{ needs.release-info.outputs.tag }}
generate_release_notes: true
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
body_path: ${{ runner.temp }}/release-notes.md
fail_on_unmatched_files: true
files: |
dist/release/*
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+12 -3
View File
@@ -10,8 +10,9 @@
[![CI](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/releases)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg?style=flat-square)](LICENSE)
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档)
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档) · [社区](#社区) · [许可证](#许可证)
</div>
@@ -20,8 +21,8 @@
## 功能
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看。
- **访问控制**:通过账号分组分配密钥权限,设置并发和请求频率。
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看;支持独立出站代理及 sub2api OpenAI 账号代理导入。
- **访问控制**:通过账号分组分配密钥权限,设置日/周 USD 限额、并发和请求频率。
- **客户端接入**:支持 Codex CLI、桌面端与 Responses API 客户端,提供 HTTP 流式响应和 WebSocket。
- **请求诊断**:查看请求记录、用量、费用、延迟和错误详情。
- **部署运维**:Docker Compose 部署、S3/R2 数据库备份、同大版本在线更新与回滚。
@@ -122,3 +123,11 @@ curl http://127.0.0.1:8080/v1/models \
- [系统架构](docs/architecture.md)
- [管理端主题](docs/theme.md)
- [数据库迁移](backend/migrations/README.md)
## 社区
感谢 [LINUX DO](https://linux.do) 社区提供开放、友善的技术交流平台。
## 许可证
本项目基于 [Apache License 2.0](LICENSE) 开源。
+9 -7
View File
@@ -1191,7 +1191,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -1488,6 +1488,7 @@ dependencies = [
"sha2 0.11.0",
"thiserror",
"tracing",
"url",
"uuid",
]
@@ -1877,10 +1878,9 @@ dependencies = [
[[package]]
name = "hyper-util"
version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
source = "git+https://github.com/hyperium/hyper-util?rev=191b1acf0a6b18e807e4abf5a6d4a7c47979ecd6#191b1acf0a6b18e807e4abf5a6d4a7c47979ecd6"
dependencies = [
"base64 0.22.1",
"base64 0.23.1",
"bytes",
"futures-channel",
"futures-util",
@@ -2610,6 +2610,7 @@ dependencies = [
"hex",
"hyper",
"hyper-rustls",
"hyper-util",
"reqwest",
"roxmltree",
"rustls",
@@ -2627,6 +2628,7 @@ dependencies = [
"tokio-rustls",
"tokio-tungstenite 0.28.0",
"tokio-util",
"tower-service",
"tracing",
"tracing-subscriber",
"tungstenite 0.27.0",
@@ -2735,7 +2737,7 @@ dependencies = [
"once_cell",
"socket2",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -3005,7 +3007,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -3707,7 +3709,7 @@ dependencies = [
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
+5 -1
View File
@@ -15,7 +15,7 @@ resolver = "3"
[workspace.package]
version = "0.0.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
@@ -38,3 +38,7 @@ strip = "symbols"
lto = "thin"
codegen-units = 1
panic = "abort"
[patch.crates-io]
# Upstream SOCKS IPv6 fix; remove once a release includes hyperium/hyper-util#302.
hyper-util = { git = "https://github.com/hyperium/hyper-util", rev = "191b1acf0a6b18e807e4abf5a6d4a7c47979ecd6" }
@@ -18,6 +18,15 @@ pub use gateway_core::account::{
QuotaState, resolve_account_status,
};
/// 导入时统一应用的账号调度与分组设置;缺省时保留原有导入语义。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountImportSettings {
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
}
/// 账号列表排序字段。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountSortField {
@@ -89,6 +98,7 @@ pub struct AccountRecord {
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
@@ -211,6 +221,7 @@ pub struct UpdateAccount {
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
}
/// 账号更新结果。
@@ -228,6 +239,7 @@ pub struct BatchUpdateAccounts {
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
}
/// 批量账号更新结果。
@@ -5,6 +5,7 @@ use std::{fmt, num::NonZeroU16};
use chrono::{DateTime, Utc};
use gateway_core::{
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
policy::{ClientApiKeyId, RateLimits},
routing::{AccountGroupId, ProviderKind},
};
@@ -93,6 +94,7 @@ pub struct ClientKeyRecord {
pub prefix: String,
pub enabled: bool,
pub limits: RateLimits,
pub budget: ClientBudgetStatus,
pub last_used_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
@@ -146,6 +148,7 @@ pub struct CreateClientKey {
pub label: Option<String>,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
pub budget: ClientBudgetLimits,
}
/// 管理用例生成 ID 与明文后的持久化命令。
@@ -156,6 +159,7 @@ pub struct NewClientKey {
pub label: Option<String>,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
pub budget: ClientBudgetLimits,
pub plaintext: String,
}
@@ -179,6 +183,8 @@ pub struct UpdateClientKey {
pub label: Option<String>,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
pub daily_limit_usd: Option<gateway_core::metering::Decimal>,
pub weekly_limit_usd: Option<gateway_core::metering::Decimal>,
}
/// 修改 Client Key 启用状态。
@@ -471,6 +471,8 @@ pub struct DashboardAccountUsage {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
/// Admin 按 Provider 补齐的套餐展示名称。
pub plan_type_display: Option<String>,
pub request_count: u64,
pub success_count: u64,
pub input_tokens: Option<u64>,
@@ -16,7 +16,9 @@ use uuid::Uuid;
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{AccountRecord, AccountSummary, AccountUsage, CredentialState},
accounts::{
AccountImportSettings, AccountRecord, AccountSummary, AccountUsage, CredentialState,
},
};
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
@@ -106,6 +108,7 @@ pub struct CredentialDetails {
/// Provider 正式文档批量导入命令。
pub struct ImportCredentials {
pub settings: Option<AccountImportSettings>,
pub context: MutationContext,
pub document: ProviderDocument,
}
@@ -144,6 +147,7 @@ impl fmt::Debug for PrepareCredentialImport {
/// Provider 已验证、可由 Store 原子创建的一份 credential。
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedCredentialCreate {
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub account_id: ProviderAccountId,
pub provider_kind: ProviderKind,
pub name: String,
@@ -171,6 +175,7 @@ pub struct PreparedCredentialImport {
/// Admin 交给 Store 的导入事务命令。
#[derive(Debug, Clone, PartialEq)]
pub struct CredentialImportCommit {
pub settings: Option<AccountImportSettings>,
pub prepared: PreparedCredentialImport,
}
@@ -249,6 +254,7 @@ pub struct PendingAuthorizationMutation {
provider_kind: ProviderKind,
target: AuthorizationMutationTarget,
owner_binding: AuthorizationOwnerBinding,
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
}
impl PendingAuthorizationMutation {
@@ -262,6 +268,7 @@ impl PendingAuthorizationMutation {
provider_kind,
target,
owner_binding,
outbound_proxy: None,
}
}
@@ -279,6 +286,19 @@ impl PendingAuthorizationMutation {
pub const fn owner_binding(&self) -> &AuthorizationOwnerBinding {
&self.owner_binding
}
#[must_use]
pub fn with_outbound_proxy(
mut self,
proxy: Option<gateway_core::account::OutboundProxy>,
) -> Self {
self.outbound_proxy = proxy;
self
}
pub fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
self.outbound_proxy.as_ref()
}
}
impl PendingAuthorizationMutation {
@@ -300,7 +320,7 @@ impl PendingAuthorizationMutation {
AuthorizationOwner::AdminApiKey => serde_json::json!({"kind": "admin_api_key"}),
AuthorizationOwner::System => serde_json::json!({"kind": "system"}),
};
Map::from_iter([
let mut document = Map::from_iter([
(
"provider_kind".to_owned(),
Value::String(self.provider_kind.as_str().to_owned()),
@@ -311,7 +331,14 @@ impl PendingAuthorizationMutation {
"started_request_id".to_owned(),
Value::String(self.owner_binding.started_request_id().to_owned()),
),
])
]);
if let Some(proxy) = &self.outbound_proxy {
document.insert(
"outbound_proxy_url".to_owned(),
Value::String(proxy.expose_url().to_owned()),
);
}
document
}
/// 恢复 adapter 已确认版本为 v1 的中立事务字段。
@@ -341,6 +368,12 @@ impl PendingAuthorizationMutation {
StoredAuthorizationOwnerV1::AdminApiKey => AuthorizationOwner::AdminApiKey,
StoredAuthorizationOwnerV1::System => AuthorizationOwner::System,
};
let proxy = document
.outbound_proxy_url
.as_deref()
.map(gateway_core::account::OutboundProxy::parse)
.transpose()
.map_err(|_| invalid())?;
Ok(Self::new(
provider_kind,
target,
@@ -348,7 +381,8 @@ impl PendingAuthorizationMutation {
owner,
started_request_id: document.started_request_id,
},
))
)
.with_outbound_proxy(proxy))
}
}
@@ -359,6 +393,8 @@ struct StoredAuthorizationMutationV1 {
target: StoredAuthorizationTargetV1,
owner: StoredAuthorizationOwnerV1,
started_request_id: String,
#[serde(default)]
outbound_proxy_url: Option<String>,
}
#[derive(Deserialize)]
@@ -382,6 +418,7 @@ pub struct StartAuthorization {
pub context: MutationContext,
pub name: String,
pub reauthorization: Option<ProviderAccountId>,
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
}
/// OAuth 流程启动结果。
@@ -447,6 +484,7 @@ pub enum AuthorizationCredentialCommit {
/// Admin 交给 Store 的 OAuth 原子事务命令。
#[derive(Debug, Clone, PartialEq)]
pub struct AuthorizationCommit {
pub settings: Option<AccountImportSettings>,
pub pending: PendingAuthorizationMutation,
pub credential: AuthorizationCredentialCommit,
}
@@ -481,7 +519,10 @@ impl PreparedAuthorizationCommit {
self
}
pub(crate) fn into_commit(self) -> AuthorizationCommitSettlement {
pub(crate) fn into_commit(
self,
settings: Option<AccountImportSettings>,
) -> AuthorizationCommitSettlement {
let (credential, guard) = match self.credential {
PreparedAuthorizationCredential::Create(credential) => {
(AuthorizationCredentialCommit::Create(credential), None)
@@ -496,6 +537,7 @@ impl PreparedAuthorizationCommit {
};
AuthorizationCommitSettlement {
command: AuthorizationCommit {
settings,
pending: self.pending,
credential,
},
@@ -509,7 +551,7 @@ impl PreparedAuthorizationCommit {
credential_guard,
authorization_guard,
..
} = self.into_commit();
} = self.into_commit(None);
drop(credential_guard);
if let Some(guard) = authorization_guard {
guard.abort().await?;
@@ -521,6 +563,7 @@ impl PreparedAuthorizationCommit {
/// 完成 Provider OAuth 流程。
#[derive(Clone, PartialEq, Eq)]
pub struct CompleteAuthorization {
pub settings: Option<AccountImportSettings>,
pub context: MutationContext,
pub flow_id: String,
pub callback_url: String,
@@ -718,6 +761,8 @@ pub struct ProviderQuotaRequest {
/// Provider 已解析的 quota 结果及其不透明差异字段。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderQuota {
/// 上游额度响应明确提供的套餐,可用于补全账号展示。
pub plan_type: Option<String>,
pub observed_at: Option<DateTime<Utc>>,
pub refresh_token_expires_at: Option<DateTime<Utc>>,
pub windows: Vec<ProviderQuotaWindow>,
@@ -726,6 +771,14 @@ pub struct ProviderQuota {
pub provider_data: Option<ProviderDocument>,
}
/// 空值和 `unknown` 代表未提供套餐;新的套餐标识仍按明确值保留。
pub(crate) fn explicit_plan_type(value: Option<&str>) -> Option<&str> {
value.filter(|value| {
let value = value.trim();
!value.is_empty() && !value.eq_ignore_ascii_case("unknown")
})
}
/// Provider 官方个人资料中的累计摘要。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderProfileStatisticsSummary {
@@ -840,6 +893,13 @@ pub struct ProviderResetCreditResult {
}
impl ProviderQuota {
/// 保留账号已有的套餐子类型,仅在缺失时使用上游额度快照补全。
pub(crate) fn fill_missing_plan_type(&self, account_plan_type: &mut Option<String>) {
if explicit_plan_type(account_plan_type.as_deref()).is_none() {
*account_plan_type = explicit_plan_type(self.plan_type.as_deref()).map(str::to_owned);
}
}
/// 返回账号展开面板使用的当前额度窗口用量。
///
/// Provider 已提供账号级本地用量时,固定窗口和无固定重置点的滚动窗口
@@ -1019,6 +1079,8 @@ impl fmt::Debug for ProviderExport {
#[derive(Debug, Clone, PartialEq)]
pub struct AccountDirectoryItem {
pub account: AccountRecord,
/// Provider 提供的套餐展示名称;未识别到套餐时为空。
pub plan_type_display: Option<String>,
pub projection: gateway_core::account::AccountStatusProjection,
pub usage: Option<super::accounts::AccountUsage>,
pub quota: ProviderQuota,
@@ -18,7 +18,7 @@ use crate::model::provider_credentials::{
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, ProviderModels,
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
ProviderResetCreditResult, ProviderResetCredits,
ProviderResetCreditResult, ProviderResetCredits, explicit_plan_type,
};
/// Provider 管理失败的稳定分类。
@@ -40,11 +40,13 @@ pub enum ProviderAdminErrorKind {
///
/// `message` 是 Provider 局部诊断,可能包含原始上游正文;通用管理用例不得自动把它作为公开文案。
/// 只有明确拥有原始诊断合同的调用方才能读取,`Debug` 始终只记录是否存在。
/// `public_message` 则是明确标记为可公开的静态提示,不允许携带动态上游材料。
#[derive(Clone, PartialEq, Eq, thiserror::Error)]
#[error("provider admin operation failed: {kind:?}")]
pub struct ProviderAdminError {
kind: ProviderAdminErrorKind,
message: Option<String>,
public_message: Option<&'static str>,
}
impl std::fmt::Debug for ProviderAdminError {
@@ -53,6 +55,7 @@ impl std::fmt::Debug for ProviderAdminError {
.debug_struct("ProviderAdminError")
.field("kind", &self.kind)
.field("message", &self.message.as_ref().map(|_| "<redacted>"))
.field("public_message", &self.public_message)
.finish()
}
}
@@ -63,6 +66,7 @@ impl ProviderAdminError {
Self {
kind,
message: None,
public_message: None,
}
}
@@ -72,6 +76,18 @@ impl ProviderAdminError {
self
}
/// Provider 明确允许公开的静态提示;不得承载上游响应或凭据材料。
#[must_use]
pub const fn with_public_message(mut self, message: &'static str) -> Self {
self.public_message = Some(message);
self
}
#[must_use]
pub const fn public_message(&self) -> Option<&'static str> {
self.public_message
}
#[must_use]
pub const fn kind(&self) -> ProviderAdminErrorKind {
self.kind
@@ -91,6 +107,11 @@ impl ProviderAdminError {
pub trait ProviderAdmin: Send + Sync {
fn provider_kind(&self) -> &ProviderKind;
/// 将原始套餐值投影为展示名称;默认保留未知 Provider 的原始名称。
fn plan_type_display(&self, plan_type: &str) -> String {
plan_type.to_owned()
}
/// 账号已经由控制面提交为不可调度状态,释放 Provider 持有的账号级运行时资源。
///
/// 无账号级运行时资源的 Provider 不需要执行额外操作。该通知发生在 Store 事务
@@ -230,6 +251,22 @@ impl ProviderAdminRegistry {
.ok_or_else(|| ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
}
/// 账号页和 Dashboard 共用的套餐展示投影,不修改持久化的原始套餐值。
pub(crate) fn plan_type_display(
&self,
provider_kind: &str,
plan_type: Option<&str>,
) -> Option<String> {
let plan_type = explicit_plan_type(plan_type)?.trim();
let provider = ProviderKind::new(provider_kind.to_owned())
.ok()
.and_then(|kind| self.providers.get(&kind));
Some(provider.map_or_else(
|| plan_type.to_owned(),
|provider| provider.plan_type_display(plan_type),
))
}
/// 返回所有已注册 Provider 的 Dashboard 上游身份画像。
pub fn dashboard_wire_profiles(&self) -> Vec<DashboardWireProfile> {
self.providers
@@ -269,7 +269,7 @@ impl DefaultAccountsService {
})
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
let stored = if refresh_quota {
let mut stored = if refresh_quota {
self.load_account(account_id).await?
} else {
stored
@@ -280,7 +280,12 @@ impl DefaultAccountsService {
)
.await?;
let usage = quota.representative_window_usage().cloned();
quota.fill_missing_plan_type(&mut stored.account.plan_type);
Ok(AccountDirectoryItem {
plan_type_display: self.providers.plan_type_display(
stored.account.provider_kind.as_str(),
stored.account.plan_type.as_deref(),
),
projection: stored.projection,
usage,
account: stored.account,
@@ -366,9 +371,14 @@ impl AccountsService for DefaultAccountsService {
.items
.into_iter()
.zip(quotas)
.map(|(item, quota)| {
.map(|(mut item, quota)| {
let usage = quota.representative_window_usage().cloned();
quota.fill_missing_plan_type(&mut item.account.plan_type);
AccountDirectoryItem {
plan_type_display: self.providers.plan_type_display(
item.account.provider_kind.as_str(),
item.account.plan_type.as_deref(),
),
usage,
account: item.account,
projection: item.projection,
@@ -735,6 +745,7 @@ fn map_reset_credits_error_after_refresh(
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
fn empty_quota() -> ProviderQuota {
ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
@@ -99,6 +99,7 @@ impl ClientKeyService for DefaultClientKeyService {
label: command.label,
group_ids: command.group_ids,
limits: command.limits,
budget: command.budget,
plaintext: plaintext.clone(),
},
context,
@@ -83,7 +83,7 @@ fn map_provider_error(
AdminErrorKind::Internal => "服务内部错误",
_ => "Provider 操作失败",
};
AdminError::new(kind, message)
AdminError::new(kind, error.public_message().unwrap_or(message))
}
async fn publish_committed(
@@ -115,9 +115,13 @@ async fn pending_authorization(
command: &StartAuthorization,
resource: &'static str,
) -> Result<PendingAuthorizationMutation, AdminError> {
let mut proxy = command.outbound_proxy.clone();
let target = match &command.reauthorization {
Some(account_id) => {
required_credential(accounts, provider_kind, account_id, resource).await?;
proxy = required_credential(accounts, provider_kind, account_id, resource)
.await?
.credential
.outbound_proxy;
AuthorizationMutationTarget::Reauthorize {
account_id: account_id.clone(),
}
@@ -132,7 +136,8 @@ async fn pending_authorization(
crate::model::provider_credentials::AuthorizationOwnerBinding::from_context(
&command.context,
),
))
)
.with_outbound_proxy(proxy))
}
fn validate_prepared_import(
@@ -217,14 +222,24 @@ async fn validate_authorization_commit(
async fn commit_authorization(
accounts: &dyn AccountStore,
prepared: PreparedAuthorizationCommit,
settings: Option<crate::model::accounts::AccountImportSettings>,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
if settings.is_some()
&& matches!(
&prepared.credential,
PreparedAuthorizationCredential::Reauthorize(_)
)
{
prepared.abort().await?;
return Err(AdminError::invalid("重新授权不能修改账号设置"));
}
let crate::model::provider_credentials::AuthorizationCommitSettlement {
command,
credential_guard,
authorization_guard,
} = prepared.into_commit();
} = prepared.into_commit(settings);
match accounts.commit_authorization(command, context).await {
Ok(result) => {
if let Some(guard) = credential_guard {
@@ -164,6 +164,11 @@ impl DefaultObservabilityService {
self.enrich_list_billing(&mut observation.recent_requests);
self.enrich_dashboard_quotas(&mut observation.account_usage)
.await;
for account in &mut observation.account_usage {
account.plan_type_display = self
.providers
.plan_type_display(&account.provider_kind, account.plan_type.as_deref());
}
let today_start = china_day_start(observation.range.end);
let yesterday_start = today_start - Duration::days(1);
let today =
@@ -770,6 +775,7 @@ impl DefaultObservabilityService {
let quotas = futures::future::join_all(quota_reads).await;
for (account, quota) in accounts.iter_mut().zip(quotas) {
if let Some(mut quota) = quota {
quota.fill_missing_plan_type(&mut account.plan_type);
quota.apply_limit_reached_display();
account.quota_used_percent = quota.representative_used_percent();
account.quota_window = quota.representative_window().cloned();
@@ -142,7 +142,13 @@ impl OpenAiService for DefaultOpenAiService {
)?;
let result = self
.accounts
.commit_credential_import(CredentialImportCommit { prepared }, &context)
.commit_credential_import(
CredentialImportCommit {
prepared,
settings: command.settings,
},
&context,
)
.await
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
self.provider
@@ -173,9 +179,10 @@ impl OpenAiService for DefaultOpenAiService {
async fn complete_authorization(
&self,
command: CompleteAuthorization,
mut command: CompleteAuthorization,
) -> Result<CredentialMutationResult, AdminError> {
let context = command.context.clone();
let settings = command.settings.take();
let prepared = self
.provider
.complete_authorization(command)
@@ -191,6 +198,7 @@ impl OpenAiService for DefaultOpenAiService {
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
settings,
&context,
"OpenAI authorization",
)
@@ -95,7 +95,13 @@ impl XaiService for DefaultXaiService {
)?;
let result = self
.accounts
.commit_credential_import(CredentialImportCommit { prepared }, &context)
.commit_credential_import(
CredentialImportCommit {
prepared,
settings: command.settings,
},
&context,
)
.await
.map_err(|error| map_store_error(error, "xAI credential import"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
@@ -132,9 +138,10 @@ impl XaiService for DefaultXaiService {
async fn complete_authorization(
&self,
command: CompleteAuthorization,
mut command: CompleteAuthorization,
) -> Result<CredentialMutationResult, AdminError> {
let context = command.context.clone();
let settings = command.settings.take();
let prepared = self
.provider
.complete_authorization(command)
@@ -150,6 +157,7 @@ impl XaiService for DefaultXaiService {
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
settings,
&context,
"xAI authorization",
)
@@ -121,6 +121,7 @@ fn provider_document_debug_should_not_expose_opaque_material() {
#[test]
fn representative_quota_should_prefer_short_window_and_highest_usage() {
let quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![
@@ -141,6 +142,7 @@ fn representative_quota_should_prefer_account_wide_window_over_model_specific_wi
let mut model_specific = quota_window("shortTerm", Some(18_000), Some(64.0));
model_specific.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
let quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![
@@ -177,6 +179,7 @@ fn representative_window_usage_should_accept_account_wide_rolling_window_without
models: Vec::new(),
};
let quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
@@ -204,6 +207,7 @@ fn representative_window_usage_should_accept_account_wide_rolling_window_without
#[test]
fn exhausted_quota_should_project_full_usage_to_only_the_representative_window() {
let mut quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![
@@ -231,6 +235,7 @@ fn exhausted_quota_should_preserve_the_provider_identified_reached_window() {
let mut reached = quota_window("monthly", Some(2_592_000), Some(98.0));
reached.limit_reached = true;
let mut quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![reached, quota_window("shortTerm", Some(18_000), Some(95.0))],
@@ -10,6 +10,7 @@ use serde_json::{Value, json};
#[test]
fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
let mut quota = ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
@@ -105,6 +105,18 @@ impl FakeProviderAdmin {
Some(ProviderAdminError::new(kind).with_message(message));
}
pub(super) fn fail_next_with_public_message(
&self,
kind: ProviderAdminErrorKind,
message: &'static str,
) {
*self.failure.lock().expect("provider failure") = Some(
ProviderAdminError::new(kind)
.with_message("upstream body containing a secret token")
.with_public_message(message),
);
}
pub(super) fn fail_next_quota(&self, kind: ProviderAdminErrorKind) {
*self.quota_failure.lock().expect("provider quota failure") = Some(kind);
}
@@ -201,6 +213,10 @@ impl ProviderAdmin for FakeProviderAdmin {
&self.kind
}
fn plan_type_display(&self, plan_type: &str) -> String {
format!("{} display: {plan_type}", self.kind)
}
async fn account_unavailable(&self, _: &ProviderAccountId) {
self.record("provider.account_unavailable");
}
@@ -425,6 +441,7 @@ pub(super) struct FakeAccountStore {
account_after_probe: Mutex<Option<AccountRecord>>,
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
import_settings: Mutex<Vec<Option<gateway_admin::model::accounts::AccountImportSettings>>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
quota_window_queries: Mutex<Vec<AccountUsageWindowQuery>>,
}
@@ -441,6 +458,7 @@ impl FakeAccountStore {
account_after_probe: Mutex::new(None),
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
import_settings: Mutex::new(Vec::new()),
quota_window_usage: Mutex::new(Vec::new()),
quota_window_queries: Mutex::new(Vec::new()),
})
@@ -450,6 +468,15 @@ impl FakeAccountStore {
*self.fail_commit.lock().expect("store failure") = true;
}
pub(super) fn import_settings(
&self,
) -> Vec<Option<gateway_admin::model::accounts::AccountImportSettings>> {
self.import_settings
.lock()
.expect("import settings")
.clone()
}
pub(super) fn audit_requests(&self) -> Vec<String> {
self.audit_requests.lock().expect("audit requests").clone()
}
@@ -656,6 +683,10 @@ impl AccountStore for FakeAccountStore {
command: CredentialImportCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialImportResult> {
self.import_settings
.lock()
.expect("import settings")
.push(command.settings);
self.record("store.commit_import");
self.record_context(context);
self.require_commit()?;
@@ -675,6 +706,10 @@ impl AccountStore for FakeAccountStore {
command: AuthorizationCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult> {
self.import_settings
.lock()
.expect("import settings")
.push(command.settings);
self.record("store.commit_authorization");
self.record_context(context);
self.require_commit()?;
@@ -1213,6 +1248,7 @@ async fn accounts_update_should_commit_then_release_disabled_account_and_publish
.update(
&context("update-request"),
UpdateAccount {
outbound_proxy: None,
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
@@ -1249,6 +1285,7 @@ async fn accounts_update_should_not_notify_provider_when_store_commit_fails() {
.update(
&context("update-failure"),
UpdateAccount {
outbound_proxy: None,
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
@@ -1290,6 +1327,7 @@ async fn accounts_batch_update_should_commit_once_and_notify_each_provider() {
.batch_update(
&context("batch-update-request"),
BatchUpdateAccounts {
outbound_proxy: None,
account_ids: vec!["acct_openai".to_owned(), "acct_xai".to_owned()],
enabled: false,
concurrency_limit: None,
@@ -1320,9 +1358,11 @@ async fn accounts_batch_update_should_commit_once_and_notify_each_provider() {
#[tokio::test]
async fn accounts_list_should_return_complete_directory_semantics() {
let provider = FakeProviderAdmin::new("openai", events());
let store = FakeAccountStore::new("openai", events());
let page = accounts_service(provider, store)
.await
let mut stored = account_record("openai");
stored.plan_type = Some(" self_serve_business_prolite ".to_owned());
let store = FakeAccountStore::with_account(stored, events());
let services = accounts_service(provider, store).await;
let page = services
.accounts()
.list(AccountListQuery {
page: 1,
@@ -1344,6 +1384,83 @@ async fn accounts_list_should_return_complete_directory_semantics() {
account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
assert_eq!(
(
account.account.plan_type.as_deref(),
account.plan_type_display.as_deref(),
),
(
Some(" self_serve_business_prolite "),
Some("openai display: self_serve_business_prolite"),
)
);
let detail = services
.accounts()
.quota(
&ProviderAccountId::new("acct_test").expect("account ID"),
false,
)
.await
.expect("account detail");
assert_eq!(detail.plan_type_display, account.plan_type_display);
assert_eq!(detail.account.plan_type, account.account.plan_type);
}
#[tokio::test]
async fn accounts_should_fill_missing_plan_from_quota_without_overriding_known_subtypes() {
for (stored_plan, quota_plan, expected) in [
(None, Some("free"), Some("free")),
(Some(" "), Some("free"), Some("free")),
(Some("unknown"), Some("free"), Some("free")),
(
Some("self_serve_business_prolite"),
Some("team"),
Some("self_serve_business_prolite"),
),
(None, None, None),
] {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
plan_type: quota_plan.map(str::to_owned),
..empty_quota()
});
let mut stored = account_record("openai");
stored.plan_type = stored_plan.map(str::to_owned);
let store = FakeAccountStore::with_account(stored, events());
let services = accounts_service(provider, store).await;
let page = services
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.expect("account list");
let account = page.items.first().expect("account item");
assert_eq!(account.account.plan_type.as_deref(), expected);
assert_eq!(
account.plan_type_display,
expected.map(|plan| format!("openai display: {plan}"))
);
for refresh in [false, true] {
let detail = services
.accounts()
.quota(
&ProviderAccountId::new("acct_test").expect("account ID"),
refresh,
)
.await
.expect("account quota detail");
assert_eq!(detail.account.plan_type.as_deref(), expected);
assert_eq!(detail.plan_type_display, account.plan_type_display);
}
}
}
#[tokio::test]
@@ -1351,6 +1468,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
let events = events();
let openai = FakeProviderAdmin::new("openai", events.clone());
openai.set_quota(ProviderQuota {
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
@@ -1482,6 +1600,7 @@ async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
@@ -1529,6 +1648,7 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
@@ -1577,6 +1697,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::hours(1);
provider.set_quota(ProviderQuota {
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
@@ -1643,6 +1764,7 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::days(7);
provider.set_quota(ProviderQuota {
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![
@@ -1897,6 +2019,7 @@ pub(super) fn document() -> ProviderDocument {
fn empty_quota() -> ProviderQuota {
ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
@@ -1931,6 +2054,7 @@ fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage {
pub(super) fn account_record(kind: &str) -> AccountRecord {
let now = Utc::now();
AccountRecord {
outbound_proxy: None,
id: "acct_test".to_owned(),
provider_kind: ProviderKind::new(kind).expect("provider kind"),
groups: Vec::new(),
@@ -1972,6 +2096,7 @@ fn prepared_create_with_id(
) -> PreparedCredentialCreate {
let now = Utc::now();
PreparedCredentialCreate {
outbound_proxy: None,
account_id: ProviderAccountId::new(account_id).expect("prepared account ID"),
provider_kind,
name: name.to_owned(),
@@ -2098,3 +2223,17 @@ fn store_unavailable() -> AdminStoreError {
fn unsupported() -> ProviderAdminError {
ProviderAdminError::new(ProviderAdminErrorKind::Unsupported)
}
pub(super) fn import_settings() -> gateway_admin::model::accounts::AccountImportSettings {
gateway_admin::model::accounts::AccountImportSettings {
enabled: false,
concurrency_limit: Some(
gateway_core::account::AccountConcurrencyLimit::new(3).expect("concurrency"),
),
weight: gateway_core::account::AccountWeight::new(7).expect("weight"),
group_ids: vec![
gateway_core::routing::AccountGroupId::new("grp_00000000000000000000000000000091")
.expect("group ID"),
],
}
}
@@ -50,6 +50,7 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
services
.openai()
.import_document(ImportCredentials {
settings: Some(super::accounts::import_settings()),
context: context("import-openai"),
document: document(),
})
@@ -75,9 +76,51 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
rolling_usage: None,
}]
);
assert_eq!(
store.import_settings(),
[Some(super::accounts::import_settings())]
);
assert_eq!(store.audit_requests(), ["import-openai"]);
}
#[tokio::test]
async fn openai_import_should_expose_only_explicit_public_errors_without_committing() {
use gateway_admin::model::AdminErrorKind;
for (provider_kind, admin_kind) in [
(ProviderAdminErrorKind::Invalid, AdminErrorKind::Invalid),
(
ProviderAdminErrorKind::Unavailable,
AdminErrorKind::Unavailable,
),
(
ProviderAdminErrorKind::BadGateway,
AdminErrorKind::BadGateway,
),
] {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider
.fail_next_with_public_message(provider_kind, "Codex PAT 验证失败,请检查令牌后重试");
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
let error = services
.openai()
.import_document(ImportCredentials {
settings: None,
context: context("import-openai-pat-failure"),
document: document(),
})
.await
.expect_err("validation must fail before committing accounts");
assert_eq!(error.kind(), admin_kind);
assert_eq!(error.to_string(), "Codex PAT 验证失败,请检查令牌后重试");
assert!(!format!("{error:?}").contains("secret token"));
assert_eq!(recorded(&events), ["provider.prepare_import"]);
assert!(store.audit_requests().is_empty());
}
}
#[tokio::test]
async fn openai_import_should_refresh_quota_for_every_imported_account() {
let events = events();
@@ -89,6 +132,7 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
let result = services
.openai()
.import_document(ImportCredentials {
settings: None,
context: context("import-openai-batch"),
document: document(),
})
@@ -132,6 +176,7 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
let result = services
.openai()
.import_document(ImportCredentials {
settings: None,
context: context("import-openai-quota-failure"),
document: document(),
})
@@ -157,6 +202,7 @@ async fn openai_authorization_create_should_observe_initial_quota() {
services
.openai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("oauth-start-openai-create"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
@@ -166,6 +212,7 @@ async fn openai_authorization_create_should_observe_initial_quota() {
services
.openai()
.complete_authorization(CompleteAuthorization {
settings: Some(super::accounts::import_settings()),
context: context("oauth-start-openai-create"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
@@ -193,6 +240,10 @@ async fn openai_authorization_create_should_observe_initial_quota() {
rolling_usage: None,
}]
);
assert_eq!(
store.import_settings(),
[Some(super::accounts::import_settings())]
);
assert_eq!(store.audit_requests(), ["oauth-start-openai-create"]);
}
@@ -207,6 +258,7 @@ async fn openai_authorization_create_should_remain_successful_when_initial_quota
services
.openai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("oauth-start-openai-quota-failure"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
@@ -216,6 +268,7 @@ async fn openai_authorization_create_should_remain_successful_when_initial_quota
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
settings: None,
context: context("oauth-start-openai-quota-failure"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
@@ -238,6 +291,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() {
store.fail_next_commit();
let services = service(provider.clone(), store).await;
let command = || CompleteAuthorization {
settings: None,
context: context("oauth-store-retry"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
@@ -246,6 +300,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() {
services
.openai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("oauth-store-retry"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
@@ -293,6 +348,7 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
services
.openai()
.import_document(ImportCredentials {
settings: None,
context: context("import-openai-error"),
document: document(),
})
@@ -311,6 +367,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances
services
.openai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("oauth-start-openai"),
name: "reauthorize".to_owned(),
reauthorization: Some(ProviderAccountId::new("acct_test").expect("account ID")),
@@ -336,6 +393,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
settings: None,
context: context("oauth-complete-openai"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
@@ -387,3 +445,34 @@ fn deletion(account_id: &str) -> CredentialDeletion {
account_ids: vec![ProviderAccountId::new(account_id).expect("account ID")],
}
}
#[tokio::test]
async fn reauthorization_with_import_settings_releases_claim_without_committing() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
services
.openai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("reauthorization-settings"),
name: "existing account".to_owned(),
reauthorization: Some(ProviderAccountId::new("acct_test").expect("account ID")),
})
.await
.expect("start reauthorization");
let error = services
.openai()
.complete_authorization(CompleteAuthorization {
settings: Some(super::accounts::import_settings()),
context: context("reauthorization-settings"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
})
.await
.expect_err("reauthorization must preserve settings");
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::Invalid);
assert!(!recorded(&events).contains(&"store.commit_authorization"));
assert!(store.import_settings().is_empty());
}
@@ -20,6 +20,7 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
services
.xai()
.import_document(ImportCredentials {
settings: Some(super::accounts::import_settings()),
context: context("import-xai"),
document: document(),
})
@@ -42,6 +43,10 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
rolling_usage: None,
}]
);
assert_eq!(
store.import_settings(),
[Some(super::accounts::import_settings())]
);
assert_eq!(store.audit_requests(), ["import-xai"]);
}
@@ -56,6 +61,7 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
services
.xai()
.import_document(ImportCredentials {
settings: None,
context: context("import-xai-batch"),
document: document(),
})
@@ -90,6 +96,7 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
let result = services
.xai()
.import_document(ImportCredentials {
settings: None,
context: context("import-xai-quota-failure"),
document: document(),
})
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-api"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
@@ -18,16 +18,51 @@ impl AccountProvider {
}
}
/// 导入统一设置,复用编辑账号的调度和分组约束。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountImportSettingsRequest {
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
pub concurrency_limit: Option<u64>,
pub weight: u64,
pub group_ids: Vec<String>,
}
impl AccountImportSettingsRequest {
fn validate(&self) -> Result<(), WireValidationError> {
parse_concurrency_limit(self.concurrency_limit)?;
parse_account_weight(self.weight)?;
validate_wire_group_ids(&self.group_ids)?;
Ok(())
}
fn into_settings(
self,
) -> Result<gateway_admin::model::accounts::AccountImportSettings, WireValidationError> {
Ok(gateway_admin::model::accounts::AccountImportSettings {
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
weight: parse_account_weight(self.weight)?,
group_ids: validate_wire_group_ids(&self.group_ids)?,
})
}
}
/// Provider-owned 账号导入请求;公共 API 不解释 `data` 内部字段。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountImportRequest {
pub settings: Option<AccountImportSettingsRequest>,
pub provider: String,
pub data: Value,
}
impl AccountImportRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if let Some(settings) = &self.settings {
settings.validate()?;
}
AccountProvider::parse(&self.provider)?;
if !self.data.is_object()
|| serde_json::to_vec(&self.data)
@@ -47,6 +82,10 @@ impl AccountImportRequest {
Ok((
provider,
ImportCredentials {
settings: self
.settings
.map(AccountImportSettingsRequest::into_settings)
.transpose()?,
context,
document: provider_document(self.data, "data")?,
},
@@ -60,6 +99,7 @@ pub struct StartAccountAuthorizationRequest {
pub provider: String,
pub name: String,
pub account_id: Option<String>,
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
}
impl StartAccountAuthorizationRequest {
@@ -89,6 +129,7 @@ impl StartAccountAuthorizationRequest {
context,
name: self.name,
reauthorization,
outbound_proxy: self.outbound_proxy_url.and_then(|value| value.0),
},
))
}
@@ -97,6 +138,7 @@ impl StartAccountAuthorizationRequest {
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CompleteAccountAuthorizationRequest {
pub settings: Option<AccountImportSettingsRequest>,
pub provider: String,
pub flow_id: String,
pub callback_url: String,
@@ -104,6 +146,9 @@ pub struct CompleteAccountAuthorizationRequest {
impl CompleteAccountAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if let Some(settings) = &self.settings {
settings.validate()?;
}
let provider = AccountProvider::parse(&self.provider)?;
match provider {
AccountProvider::OpenAi => {
@@ -128,6 +173,10 @@ impl CompleteAccountAuthorizationRequest {
Ok((
provider,
CompleteAuthorization {
settings: self
.settings
.map(AccountImportSettingsRequest::into_settings)
.transpose()?,
context,
flow_id: self.flow_id,
callback_url: self.callback_url,
@@ -139,6 +188,7 @@ impl CompleteAccountAuthorizationRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateAccountRequest {
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
pub account_id: String,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
@@ -159,6 +209,7 @@ impl UpdateAccountRequest {
pub(super) fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
self.validate()?;
Ok(UpdateAccount {
outbound_proxy: self.outbound_proxy_url.map(|value| value.0),
account_id: self.account_id,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
@@ -59,6 +59,7 @@ pub(super) fn account_models_data(result: ProviderModels) -> AccountModelsData {
pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
let AccountDirectoryItem {
account,
plan_type_display,
projection,
usage,
quota,
@@ -89,6 +90,7 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
user_id: account.upstream_user_id,
label: None,
plan_type: account.plan_type,
plan_type_display: plan_type_display.unwrap_or_else(|| "未知套餐".to_owned()),
authentication_kind: account.authentication_kind,
has_refresh_token: account.has_refresh_token,
status,
@@ -99,6 +101,10 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
enabled: account.enabled,
concurrency_limit: account.concurrency_limit.map(|limit| limit.get()),
weight: account.weight.get(),
outbound_proxy_endpoint: account
.outbound_proxy
.as_ref()
.map(|proxy| proxy.endpoint()),
access_token_expires_at: expires_at,
access_token_expires_at_display: account
.access_token_expires_at
@@ -2,6 +2,21 @@
use super::*;
#[derive(Debug, Clone)]
pub struct AccountProxyUpdate(pub Option<gateway_core::account::OutboundProxy>);
impl<'de> Deserialize<'de> for AccountProxyUpdate {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let value = String::deserialize(deserializer)?;
if value.is_empty() {
return Ok(Self(None));
}
gateway_core::account::OutboundProxy::parse(&value)
.map(|proxy| Self(Some(proxy)))
.map_err(serde::de::Error::custom)
}
}
/// 账号列表查询参数。
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
@@ -19,6 +34,7 @@ pub struct ListQuery {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct BatchUpdateAccountsRequest {
pub outbound_proxy_url: Option<AccountProxyUpdate>,
pub account_ids: Vec<String>,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
@@ -68,6 +84,7 @@ impl BatchUpdateAccountsRequest {
pub(super) fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
self.validate()?;
Ok(BatchUpdateAccounts {
outbound_proxy: self.outbound_proxy_url.map(|value| value.0),
account_ids: self.account_ids,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
@@ -192,6 +209,7 @@ pub struct AccountSummaryView {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountView {
pub outbound_proxy_endpoint: Option<String>,
pub id: String,
pub name: String,
pub provider: String,
@@ -202,6 +220,8 @@ pub struct AccountView {
pub user_id: Option<String>,
pub label: Option<String>,
pub plan_type: Option<String>,
/// 后端生成的套餐展示名称;缺失套餐时为“未知套餐”。
pub plan_type_display: String,
pub authentication_kind: String,
pub has_refresh_token: bool,
pub status: String,
@@ -11,6 +11,8 @@ use gateway_admin::model::client_keys::{
CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey,
};
use gateway_core::{
engine::budget::ClientBudgetLimits,
metering::Decimal,
policy::{ClientApiKeyId, RateLimits},
routing::AccountGroupId,
};
@@ -33,6 +35,15 @@ const MAX_CURSOR_BYTES: usize = 512;
const MAX_SEARCH_BYTES: usize = 256;
const DEFAULT_PAGE_SIZE: u16 = 50;
fn parse_budget(
value: Option<String>,
field: &'static str,
) -> Result<Option<Decimal>, WireValidationError> {
value
.map(|value| value.parse().map_err(|_| WireValidationError::new(field)))
.transpose()
}
/// Client Key 列表查询。
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
@@ -140,6 +151,8 @@ pub struct CreateClientKeyRequest {
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
daily_limit_usd: Option<String>,
weekly_limit_usd: Option<String>,
}
impl CreateClientKeyRequest {
@@ -154,6 +167,11 @@ impl CreateClientKeyRequest {
name: self.name,
label: self.label,
group_ids,
budget: ClientBudgetLimits {
daily_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?.unwrap_or_default(),
weekly_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")?
.unwrap_or_default(),
},
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -172,6 +190,8 @@ pub struct UpdateClientKeyRequest {
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
daily_limit_usd: Option<String>,
weekly_limit_usd: Option<String>,
}
impl UpdateClientKeyRequest {
@@ -188,6 +208,8 @@ impl UpdateClientKeyRequest {
name: self.name,
label: self.label,
group_ids,
daily_limit_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?,
weekly_limit_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")?,
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -247,6 +269,12 @@ pub struct ClientKeyView {
enabled: bool,
max_concurrency: u64,
requests_per_minute: u64,
daily_limit_usd: String,
weekly_limit_usd: String,
daily_used_usd: String,
weekly_used_usd: String,
daily_resets_at: Option<DateTime<Utc>>,
weekly_resets_at: Option<DateTime<Utc>>,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
last_used_at: Option<DateTime<Utc>>,
@@ -292,6 +320,12 @@ impl From<ClientKeyRecord> for ClientKeyView {
enabled: record.enabled,
max_concurrency: record.limits.max_concurrency,
requests_per_minute: record.limits.requests_per_minute,
daily_limit_usd: record.budget.limits.daily_usd.canonical(),
weekly_limit_usd: record.budget.limits.weekly_usd.canonical(),
daily_used_usd: record.budget.daily_used_usd.canonical(),
weekly_used_usd: record.budget.weekly_used_usd.canonical(),
daily_resets_at: record.budget.daily_resets_at.map(DateTime::from),
weekly_resets_at: record.budget.weekly_resets_at.map(DateTime::from),
created_at: record.created_at,
updated_at: record.updated_at,
last_used_at: record.last_used_at,
@@ -900,6 +900,10 @@ pub(crate) fn dashboard_view(
.clone()
.unwrap_or_else(|| credential.name.clone()),
plan_type: credential.plan_type.clone(),
plan_type_display: credential
.plan_type_display
.clone()
.unwrap_or_else(|| "未知套餐".to_owned()),
tokens: credential
.total_tokens
.map_or_else(|| "—".to_owned(), format_compact_number),
@@ -373,6 +373,7 @@ pub struct DashboardAccountUsageView {
pub authentication_kind: String,
pub email: String,
pub plan_type: Option<String>,
pub plan_type_display: String,
pub tokens: String,
pub request_count: u64,
pub request_buckets: Vec<DashboardAccountRequestBucketView>,
+16 -2
View File
@@ -143,13 +143,27 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
/// Gateway 错误的 OpenAI HTTP 表达。
pub fn gateway_error_response(error: &GatewayError) -> Response {
let (status, default_type, default_code) = gateway_error_contract(error.kind());
openai_error_response(
let mut response = openai_error_response(
status,
error.client_message(),
error.client_error_type().unwrap_or(default_type),
error.client_error_code().unwrap_or(default_code),
)
.into_response()
.into_response();
if let Some(delay) = error.retry_after()
&& let Ok(value) = HeaderValue::from_str(
&delay
.as_secs()
.saturating_add(u64::from(delay.subsec_nanos() > 0))
.max(1)
.to_string(),
)
{
response
.headers_mut()
.insert(axum::http::header::RETRY_AFTER, value);
}
response
}
/// 在下游尚未提交时优先交付 Provider 的原始 HTTP 失败响应。
@@ -756,3 +756,60 @@ mod actions {
);
}
}
mod import_settings {
use gateway_api::admin::accounts::{AccountImportRequest, CompleteAccountAuthorizationRequest};
use serde_json::json;
#[test]
fn import_and_oauth_apply_the_same_settings_validation() {
for (field, value) in [
("concurrencyLimit", json!(0)),
("concurrencyLimit", json!(4_294_967_296_u64)),
("weight", json!(0)),
("weight", json!(101)),
("groupIds", json!(["invalid-group"])),
] {
let mut settings =
json!({"enabled": false, "concurrencyLimit": null, "weight": 1, "groupIds": []});
settings[field] = value;
let import: AccountImportRequest = serde_json::from_value(
json!({"provider": "openai", "data": {}, "settings": settings}),
)
.expect("import request");
let oauth: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({"provider": "xai", "flowId": "flow-test", "callbackUrl": "code", "settings": settings})).expect("OAuth request");
assert_eq!(
import.validate().expect_err("invalid settings").field(),
field
);
assert_eq!(
oauth.validate().expect_err("invalid settings").field(),
field
);
}
}
#[test]
fn import_settings_require_a_complete_explicit_configuration() {
let settings =
json!({"enabled": false, "concurrencyLimit": null, "weight": 100, "groupIds": []});
let request: AccountImportRequest =
serde_json::from_value(json!({"provider": "openai", "data": {}, "settings": settings}))
.expect("settings");
assert!(request.validate().is_ok());
for field in ["enabled", "concurrencyLimit", "weight", "groupIds"] {
let mut incomplete = settings.clone();
incomplete
.as_object_mut()
.expect("settings object")
.remove(field);
assert!(
serde_json::from_value::<AccountImportRequest>(
json!({"provider": "openai", "data": {}, "settings": incomplete})
)
.is_err(),
"missing {field}"
);
}
}
}
@@ -10,6 +10,46 @@ use serde_json::json;
use super::{AdminTestFixture, AdminTestState};
#[test]
fn budget_inputs_preserve_decimal_precision_and_omitted_updates() {
let payload = json!({"name": "budget", "groupIds": [], "maxConcurrency": 3,
"requestsPerMinute": 0, "dailyLimitUsd": "0.1234567891", "weeklyLimitUsd": "15"});
let command = serde_json::from_value::<CreateClientKeyRequest>(payload.clone())
.unwrap()
.into_command()
.unwrap();
assert_eq!(command.budget.daily_usd.canonical(), "0.1234567891");
assert_eq!(command.budget.weekly_usd.canonical(), "15");
assert_eq!(command.limits.max_concurrency, 3);
for field in ["dailyLimitUsd", "weeklyLimitUsd"] {
for invalid in ["-1", "NaN", "1e3", "10000000000", "0.00000000001", ""] {
let mut payload = payload.clone();
payload[field] = json!(invalid);
assert_eq!(
serde_json::from_value::<CreateClientKeyRequest>(payload)
.unwrap()
.into_command()
.unwrap_err()
.field(),
field
);
}
}
let mut update = payload;
update["id"] = json!("key_budget");
update.as_object_mut().unwrap().remove("dailyLimitUsd");
update["weeklyLimitUsd"] = json!("0");
let command = serde_json::from_value::<UpdateClientKeyRequest>(update)
.unwrap()
.into_command()
.unwrap();
assert_eq!(command.daily_limit_usd, None);
assert_eq!(
command.weekly_limit_usd,
Some(gateway_core::metering::Decimal::ZERO)
);
}
#[test]
fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
let unknown = serde_json::from_value::<ListClientKeysQuery>(json!({ "other": true }));
@@ -207,6 +247,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
.single()
.expect("valid time");
let view = ClientKeyView::from(gateway_admin::model::client_keys::ClientKeyRecord {
budget: Default::default(),
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
name: "visible".to_owned(),
label: None,
@@ -238,6 +279,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
assert!(list["items"][0].get("tokensPerMinute").is_none());
assert!(list["items"][0].get("policyJson").is_none());
assert_eq!(list["items"][0]["dailyLimitUsd"], "0");
assert_eq!(list["items"][0]["weeklyLimitUsd"], "0");
assert_eq!(list["items"][0]["dailyUsedUsd"], "0");
assert!(list["items"][0].get("unresolvedRequests").is_none());
assert_eq!(
DateTime::parse_from_rfc3339(
list["items"][0]["lastUsedAt"]
@@ -640,6 +640,7 @@ impl ClientKeyStore for MemoryClientKeyStore {
let now = Utc::now();
Ok(Some(ClientKeySecret::new(
ClientKeyRecord {
budget: Default::default(),
id: id.clone(),
name: "revealed".to_owned(),
label: None,
@@ -64,6 +64,7 @@ fn dashboard_account_usage_should_keep_daily_request_timeline() {
authentication_kind: "oauth".to_owned(),
email: "account@example.com".to_owned(),
plan_type: Some("free".to_owned()),
plan_type_display: "Free".to_owned(),
tokens: "—".to_owned(),
request_count: 3,
request_buckets: vec![DashboardAccountRequestBucketView {
@@ -11,6 +11,21 @@ use gateway_api::openai::error::{
openai_error_response,
};
#[tokio::test]
async fn key_budget_errors_preserve_limit_code_and_retry_after() {
for code in ["key_daily_budget_exceeded", "key_weekly_budget_exceeded"] {
let error = GatewayError::new(GatewayErrorKind::RateLimited, "key budget exhausted")
.with_client_code(code)
.with_retry_after(std::time::Duration::from_millis(1501));
let response = gateway_error_response(&error);
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(response.headers()["retry-after"], "2");
let body: serde_json::Value =
serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap();
assert_eq!(body["error"]["code"], code);
}
}
#[test]
fn invalid_request_error_should_map_to_openai_bad_request() {
assert_eq!(
@@ -433,3 +448,151 @@ async fn continuation_recovery_should_preserve_the_official_retry_signal() {
})
);
}
mod model_routing {
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode},
};
use futures::future::BoxFuture;
use gateway_core::engine::execution::{
DefaultExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::engine::provider::ProviderRegistry;
use gateway_core::routing::{ProviderKind, RuntimeSnapshot};
use gateway_core::runtime::RuntimeSnapshotHandle;
use serde_json::{Value, json};
use tower::ServiceExt;
use crate::openai::{
IgnoredClientApiKeyUsage, UnusedAdmissions, UnusedContinuation, UnusedExecutionStore,
};
struct TestCircuits {
blocked: bool,
}
impl ProviderCircuitPort for TestCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async move {
Ok(if self.blocked {
ProviderCircuitDecision::BlockedUntil(
SystemTime::now() + Duration::from_secs(30),
)
} else {
ProviderCircuitDecision::Allow
})
})
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
}
}
async fn request_model(
snapshot: RuntimeSnapshot,
model: &str,
stream: bool,
blocked: bool,
) -> (StatusCode, Value) {
let execution = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(TestCircuits { blocked }),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
);
let response = crate::openai::api_router(Arc::new(execution))
.await
.oneshot(
Request::post("/v1/responses")
.header("authorization", "Bearer sk_model_routing")
.header("content-type", "application/json")
.body(Body::from(
json!({"model": model, "input": "hello", "stream": stream}).to_string(),
))
.expect("request"),
)
.await
.expect("response");
let status = response.status();
let body = to_bytes(response.into_body(), 16 * 1024)
.await
.expect("error body");
(status, serde_json::from_slice(&body).expect("error JSON"))
}
#[tokio::test]
async fn missing_models_should_return_404_with_a_specific_explanation() {
for (model, mappings, message) in [
(
"missing-model",
BTreeMap::new(),
"the requested model was not found in the provider catalogs available to this API key; check the model name",
),
(
"model-a",
BTreeMap::from([("model-a".to_owned(), "missing-target".to_owned())]),
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping",
),
] {
for stream in [false, true] {
let snapshot = crate::openai::snapshot("sk_model_routing", "openai")
.with_model_mappings(mappings.clone());
let response = request_model(snapshot, model, stream, false).await;
assert_eq!(
response,
(
StatusCode::NOT_FOUND,
json!({"error": {
"type": "invalid_request_error",
"code": "model_not_found",
"message": message,
}}),
),
"model={model}, stream={stream}",
);
}
}
}
#[tokio::test]
async fn an_existing_model_with_a_blocked_provider_should_remain_service_unavailable() {
let response = request_model(
crate::openai::snapshot("sk_model_routing", "openai"),
"model-a",
true,
true,
)
.await;
assert_eq!(
(response.0, response.1["error"]["code"].as_str()),
(
StatusCode::SERVICE_UNAVAILABLE,
Some("no_available_provider")
),
);
}
}
+2 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-core"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
@@ -24,6 +24,7 @@ semver = "1.0.27"
thiserror = "2.0.18"
tracing = "0.1.44"
uuid = { version = "1.23.3", features = ["v7"] }
url = "2.5.8"
[lints]
workspace = true
@@ -2,6 +2,8 @@
mod error;
mod model;
mod proxy;
pub use proxy::{InvalidOutboundProxy, OutboundProxy};
pub mod scope;
mod selection;
mod store;
@@ -654,7 +654,7 @@ const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
}
}
/// 不含 secret 的账号持久事实。
/// 账号持久事实;代理认证信息只通过显式 secret accessor 读取。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
id: ProviderAccountId,
@@ -676,6 +676,7 @@ pub struct ProviderAccount {
access_token_expires_at: Option<SystemTime>,
next_refresh_at: Option<SystemTime>,
has_refresh_token: bool,
outbound_proxy: Option<super::OutboundProxy>,
}
impl ProviderAccount {
@@ -710,6 +711,7 @@ impl ProviderAccount {
access_token_expires_at,
next_refresh_at: None,
has_refresh_token: false,
outbound_proxy: None,
}
}
@@ -726,6 +728,17 @@ impl ProviderAccount {
self
}
#[must_use]
pub fn with_outbound_proxy(mut self, proxy: Option<super::OutboundProxy>) -> Self {
self.outbound_proxy = proxy;
self
}
#[must_use]
pub const fn outbound_proxy(&self) -> Option<&super::OutboundProxy> {
self.outbound_proxy.as_ref()
}
#[must_use]
pub fn with_account_facts(
mut self,
@@ -0,0 +1,49 @@
//! Explicit account egress. Credentials never appear in Debug or ordinary admin projections.
use std::fmt;
use url::Url;
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct OutboundProxy(Url);
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("invalid outbound proxy; expected http, https, socks5 or socks5h URL with a host and port")]
pub struct InvalidOutboundProxy;
impl OutboundProxy {
pub fn parse(value: &str) -> Result<Self, InvalidOutboundProxy> {
if value.len() > 4096 || value.chars().any(char::is_control) {
return Err(InvalidOutboundProxy);
}
let url = Url::parse(value).map_err(|_| InvalidOutboundProxy)?;
if !matches!(url.scheme(), "http" | "https" | "socks5" | "socks5h")
|| url.host_str().is_none()
|| url.port_or_known_default().is_none_or(|port| port == 0)
|| !matches!(url.path(), "" | "/")
|| url.query().is_some()
|| url.fragment().is_some()
{
return Err(InvalidOutboundProxy);
}
Ok(Self(url))
}
#[must_use]
pub fn expose_url(&self) -> &str {
self.0.as_str()
}
#[must_use]
pub fn endpoint(&self) -> String {
let mut url = self.0.clone();
let _ = url.set_username("");
let _ = url.set_password(None);
url.to_string()
}
}
impl fmt::Debug for OutboundProxy {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("OutboundProxy(<redacted>)")
}
}
@@ -0,0 +1,52 @@
//! 按 Key 持久化 USD 限额与已取得费用,不依赖尽力写入的请求观测。
use std::time::SystemTime;
use futures::future::BoxFuture;
use crate::{error::GatewayError, metering::Decimal, policy::ClientApiKeyId};
use super::ModelRequestId;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ClientBudgetLimits {
pub daily_usd: Decimal,
pub weekly_usd: Decimal,
}
impl ClientBudgetLimits {
#[must_use]
pub fn is_limited(self) -> bool {
self.daily_usd != Decimal::ZERO || self.weekly_usd != Decimal::ZERO
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ClientBudgetStatus {
pub limits: ClientBudgetLimits,
pub daily_used_usd: Decimal,
pub weekly_used_usd: Decimal,
pub daily_resets_at: Option<SystemTime>,
pub weekly_resets_at: Option<SystemTime>,
}
#[derive(Debug, Clone)]
pub struct ClientBudgetCharge {
pub key_id: ClientApiKeyId,
pub request_id: ModelRequestId,
/// 已取得的 USD 费用,包含重试;缺少费用的尝试按零累计。
pub amount_usd: Decimal,
pub completed_at: SystemTime,
}
#[derive(Debug, Clone, Copy, thiserror::Error)]
#[error("client budget store is unavailable")]
pub struct ClientBudgetError;
pub trait ClientBudgetPort: Send + Sync {
/// 原子检查当前限额与已用金额,不创建预扣费或待结算记录。
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>>;
/// 按网关请求 ID 幂等累计已取得费用;写入失败由 Store 重试。
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>;
}
@@ -30,6 +30,7 @@ use crate::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
};
use crate::lifecycle::CancellationToken;
use crate::metering::Decimal;
use crate::operation::{Operation, ProviderSessionState};
use crate::routing::RoutingPlan;
use futures::future::Fuse;
@@ -166,6 +167,8 @@ where
request_id,
client_api_key_ref,
observation: ResponseObservation::new(timing_started_at),
budget_prior_attempts_usd: Decimal::ZERO,
budget_attempt_already_counted: false,
trace,
deadline,
deadline_timer: Delay::new(
@@ -197,6 +200,7 @@ where
delivery_pending: false,
upstream_complete: false,
finalized: false,
finalized_at: None,
image_generation_requested,
last_retryable_failure: None,
last_retryable_failure_events: Vec::new(),
@@ -253,6 +257,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
request_id: ModelRequestId,
client_api_key_ref: crate::policy::ClientApiKeyId,
observation: ResponseObservation,
budget_prior_attempts_usd: Decimal,
budget_attempt_already_counted: bool,
trace: TraceContext,
deadline: SystemTime,
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
@@ -288,6 +294,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
delivery_pending: bool,
upstream_complete: bool,
finalized: bool,
finalized_at: Option<SystemTime>,
image_generation_requested: bool,
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
last_retryable_failure: Option<ProviderError>,
@@ -475,6 +482,35 @@ where
self.finalized
}
#[must_use]
pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge {
// 超出数据库可表示范围时保留最大金额,避免溢出后误记为零。
let amount_usd = self
.budget_prior_attempts_usd
.checked_add(self.budget_attempt_usd())
.unwrap_or(Decimal::MAX);
super::budget::ClientBudgetCharge {
key_id: self.client_api_key_ref.clone(),
request_id: self.request_id.clone(),
amount_usd,
completed_at: self.finalized_at.unwrap_or_else(SystemTime::now),
}
}
fn budget_attempt_usd(&self) -> Decimal {
if self.budget_attempt_already_counted {
return Decimal::ZERO;
}
// Key 只累计已取得的 USD 费用;缺少费用的请求按零结算。
// 发送状态仍用于判断重放是否安全并保留在诊断中,不能据此推断存在欠费。
self.observation
.cost
.total()
.filter(|money| money.currency().as_str() == "USD")
.map(crate::metering::Money::amount)
.unwrap_or(Decimal::ZERO)
}
/// 返回截至当前已完成的实际上游调用结果。
///
/// 调用方可保存已消费下标;该切片在会话生命周期内只追加、不改写。
@@ -800,6 +836,7 @@ where
let mut events = std::mem::take(&mut self.last_retryable_failure_events);
if !events.is_empty() {
self.observe_atomic_terminal_events(&mut events);
self.budget_attempt_already_counted = true;
return Ok(Some(PullOutcome::TerminalFailure {
events,
error: last_failure,
@@ -1143,6 +1180,11 @@ where
"sendState": format!("{attempt_send_state:?}"),
}));
if retryable {
for event in &atomic_client_events {
for fact in event.canonical_facts() {
self.observation.observe_event(fact);
}
}
// 普通 clone 只保留稳定事实;原始 wire/HTTP response 由 request-local
// 所有权保留到下一次 attempt 成功,或最终空选路时返回客户端。
let persistence_error = self.request_persisted.then(|| error.clone());
@@ -1317,6 +1359,12 @@ where
}
fn reset_uncommitted_observations(&mut self) {
// 响应观测按尝试隔离;被丢弃尝试中已取得的费用仍计入 Key。
self.budget_prior_attempts_usd = self
.budget_prior_attempts_usd
.checked_add(self.budget_attempt_usd())
.unwrap_or(Decimal::MAX);
self.budget_attempt_already_counted = false;
self.observation.reset_for_attempt();
self.upstream_complete = false;
}
@@ -1326,6 +1374,7 @@ where
return Ok(());
}
let completed_at = SystemTime::now();
self.finalized_at = Some(completed_at);
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
@@ -1505,11 +1554,12 @@ where
"sendState": format!("{:?}", finalization.send_state), "attempts": self.attempts,
}),
);
let completed_at = SystemTime::now();
self.finalized_at = Some(completed_at);
if !self.request_persisted {
self.finalized = true;
return Ok(());
}
let completed_at = SystemTime::now();
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
@@ -14,6 +14,7 @@ use uuid::Uuid;
use crate::engine::admission::{
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
};
use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort};
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
NativeContinuationStoreErrorKind, PreviousResponseId,
@@ -256,6 +257,7 @@ pub struct DefaultExecutionService {
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
budget: Option<Arc<dyn ClientBudgetPort>>,
}
impl DefaultExecutionService {
@@ -283,9 +285,16 @@ impl DefaultExecutionService {
circuits,
continuation,
client_api_key_usage,
budget: None,
}
}
#[must_use]
pub fn with_budget(mut self, budget: Arc<dyn ClientBudgetPort>) -> Self {
self.budget = Some(budget);
self
}
async fn start_inner(&self, request: StartExecution) -> Result<StartedExecution, GatewayError> {
let StartExecution {
client,
@@ -325,9 +334,18 @@ impl DefaultExecutionService {
&self,
mut request: PendingStartExecution,
) -> Result<StartedExecution, GatewayError> {
request.client.policy.authorize().map_err(|_| {
GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled")
})?;
// 长连接每次执行都重新鉴权并冻结当前策略,确保限额和授权变更对新请求生效。
request.client = self
.authenticate(request.client.policy.plaintext_key().expose_for_auth())
.map_err(|error| match error {
ClientAuthenticationError::InvalidKey => {
GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid")
}
ClientAuthenticationError::SnapshotUnavailable => GatewayError::new(
GatewayErrorKind::Internal,
"runtime snapshot is unavailable",
),
})?;
let started_at = SystemTime::now();
let deadline_at = started_at
.checked_add(MODEL_REQUEST_DEADLINE)
@@ -504,6 +522,12 @@ impl DefaultExecutionService {
client_api_key_id: client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
if let Some(budget) = &self.budget
&& let Err(error) = budget.admit(client.policy.key_id().clone()).await
{
admission.release().await;
return Err(error);
}
let observation = plan
.candidates()
.first()
@@ -552,6 +576,18 @@ impl DefaultExecutionService {
{
Ok(core) => core,
Err(error) => {
if let Some(budget) = &self.budget {
settle_budget(
budget.as_ref(),
ClientBudgetCharge {
key_id: client.policy.key_id().clone(),
request_id: request_id.clone(),
amount_usd: crate::metering::Decimal::ZERO,
completed_at: SystemTime::now(),
},
)
.await;
}
admission.release().await;
return Err(gateway_error_from_engine(&error));
}
@@ -565,6 +601,7 @@ impl DefaultExecutionService {
admission,
Arc::clone(&self.circuits),
Arc::clone(&self.continuation),
self.budget.clone(),
)),
})
}
@@ -922,6 +959,12 @@ struct AdmissionLease {
model_request_id: ModelRequestId,
}
async fn settle_budget(port: &dyn ClientBudgetPort, charge: ClientBudgetCharge) {
if let Err(error) = port.settle(charge).await {
tracing::error!(%error, "Client budget settlement failed; storage will retry on the next request");
}
}
impl AdmissionLease {
async fn release(self) {
if let Err(error) = self
@@ -941,6 +984,7 @@ struct DefaultExecutionSession {
continuation: Arc<dyn NativeContinuationPort>,
observed_provider_outcomes: usize,
continuation_recorded: bool,
budget: Option<Arc<dyn ClientBudgetPort>>,
}
impl DefaultExecutionSession {
@@ -949,6 +993,7 @@ impl DefaultExecutionSession {
admission: AdmissionLease,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
budget: Option<Arc<dyn ClientBudgetPort>>,
) -> Self {
Self {
core: Some(core),
@@ -957,6 +1002,7 @@ impl DefaultExecutionSession {
continuation,
observed_provider_outcomes: 0,
continuation_recorded: false,
budget,
}
}
@@ -973,6 +1019,11 @@ impl DefaultExecutionSession {
.is_some_and(ResponseExecutionSession::is_finalized)
&& let Some(admission) = self.admission.take()
{
if let Some(budget) = self.budget.take()
&& let Some(core) = &self.core
{
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
admission.release().await;
}
}
@@ -1023,6 +1074,9 @@ impl DefaultExecutionSession {
.get(self.observed_provider_outcomes..)
.unwrap_or_default();
publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await;
if let Some(budget) = self.budget.take() {
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
if let Some(admission) = self.admission.take() {
admission.release().await;
}
@@ -1213,6 +1267,17 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError {
match error {
crate::validation::RoutingError::ModelNotFound {
model,
mapped_model,
} => GatewayError::new(
GatewayErrorKind::ModelNotFound,
if model == mapped_model {
"the requested model was not found in the provider catalogs available to this API key; check the model name"
} else {
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping"
},
),
crate::validation::RoutingError::NoCapableProvider { .. }
| crate::validation::RoutingError::NoCapableProviderEndpoint { .. }
| crate::validation::RoutingError::EmptyAccountScope => GatewayError::new(
@@ -1,6 +1,7 @@
//! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。
pub mod admission;
pub mod budget;
pub mod continuation;
pub mod coordinator;
pub mod execution;
+36
View File
@@ -975,6 +975,13 @@ pub struct GatewayError {
message: &'static str,
diagnostic: Option<Box<ProviderDiagnostic>>,
client_visible_upstream_error: Option<ClientVisibleUpstreamError>,
client_details: Option<Box<GatewayClientErrorDetails>>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct GatewayClientErrorDetails {
code: Option<&'static str>,
retry_after: Option<Duration>,
}
impl GatewayError {
@@ -986,9 +993,33 @@ impl GatewayError {
message,
diagnostic: None,
client_visible_upstream_error: None,
client_details: None,
}
}
#[must_use]
pub fn with_client_code(mut self, code: &'static str) -> Self {
self.client_details
.get_or_insert_with(Default::default)
.code = Some(code);
self
}
#[must_use]
pub fn with_retry_after(mut self, delay: Duration) -> Self {
self.client_details
.get_or_insert_with(Default::default)
.retry_after = Some(delay);
self
}
#[must_use]
pub fn retry_after(&self) -> Option<Duration> {
self.client_details
.as_ref()
.and_then(|details| details.retry_after)
}
/// 将 Provider 错误归一为客户端无关错误。
#[must_use]
pub fn from_provider(error: &ProviderError) -> Self {
@@ -1103,6 +1134,11 @@ impl GatewayError {
self.client_visible_upstream_error
.as_ref()
.and_then(ClientVisibleUpstreamError::code)
.or_else(|| {
self.client_details
.as_ref()
.and_then(|details| details.code)
})
}
}
+14 -2
View File
@@ -52,6 +52,7 @@ pub struct CoreStorePorts {
snapshots: Arc<dyn SnapshotStorePort>,
snapshot_subscriptions: Arc<dyn SnapshotSubscriptionPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
budget: Option<Arc<dyn engine::budget::ClientBudgetPort>>,
}
impl CoreStorePorts {
@@ -79,8 +80,15 @@ impl CoreStorePorts {
snapshots,
snapshot_subscriptions,
client_api_key_usage,
budget: None,
}
}
#[must_use]
pub fn with_budget(mut self, budget: Arc<dyn engine::budget::ClientBudgetPort>) -> Self {
self.budget = Some(budget);
self
}
}
pub struct CoreBundle {
@@ -145,7 +153,7 @@ pub async fn initialize(
Arc::clone(&ports.snapshot_subscriptions),
));
let worker_contributions = publisher.worker_contributions()?;
let service = Arc::new(DefaultExecutionService::new(
let mut service = DefaultExecutionService::new(
snapshots.clone(),
ports.execution,
providers,
@@ -153,7 +161,11 @@ pub async fn initialize(
ports.circuits,
ports.continuation,
ports.client_api_key_usage,
));
);
if let Some(budget) = ports.budget {
service = service.with_budget(budget);
}
let service = Arc::new(service);
let execution: Arc<dyn ExecutionService> = service.clone();
let account_probe: Arc<dyn AccountProbe> = service;
let snapshot_control: Arc<dyn SnapshotControl> = publisher;
@@ -14,6 +14,8 @@ pub struct Decimal(u128);
impl Decimal {
pub const ZERO: Self = Self(0);
/// 数据库可表示的最大非负金额。
pub const MAX: Self = Self(MAX_SCALED_DECIMAL);
/// 从按十位小数缩放的整数创建。
///
@@ -778,6 +778,19 @@ impl RuntimeSnapshot {
}
if candidates.is_empty() {
// 模型存在性必须包含被熔断的 Provider;目录未知时不能断言模型不存在。
let mut scoped_providers = providers.intersection(&self.providers).peekable();
if scoped_providers.peek().is_some()
&& scoped_providers.all(|provider| {
self.known_provider_catalogs.contains(provider)
&& !self.contains_public_model_for_provider(public_model, provider)
})
{
return Err(RoutingError::ModelNotFound {
model: public_model.as_str().to_owned(),
mapped_model: self.mapped_model(public_model.as_str()),
});
}
return Err(RoutingError::NoCapableProvider {
model: public_model.as_str().to_owned(),
});
@@ -85,6 +85,16 @@ pub enum RoutingError {
/// 实体 ID。
id: String,
},
/// 账号范围内的 Provider 目录均已知,且均没有请求或映射后的模型。
#[error(
"model `{model}` (mapped to `{mapped_model}`) was not found in the scoped provider catalogs"
)]
ModelNotFound {
/// 客户端提交的模型名称。
model: String,
/// 应用模型映射后的名称。
mapped_model: String,
},
/// 固定平台内没有可执行本次请求的 Provider。
#[error("no provider can execute model `{model}`")]
NoCapableProvider {
@@ -1,3 +1,4 @@
mod proxy;
mod selection;
mod selection_diagnostics;
@@ -0,0 +1,33 @@
use gateway_core::account::OutboundProxy;
#[test]
fn proxy_endpoints_support_explicit_schemes_and_redact_credentials() {
for scheme in ["http", "https", "socks5", "socks5h"] {
let proxy =
OutboundProxy::parse(&format!("{scheme}://user:p%40ss%3Aword@[::1]:1080")).unwrap();
assert!(proxy.expose_url().contains("p%40ss%3Aword"));
assert!(!proxy.endpoint().contains("user"));
assert!(!proxy.endpoint().contains("p%40ss"));
assert!(!format!("{proxy:?}").contains("user"));
assert!(proxy.endpoint().contains("[::1]:1080"));
}
}
#[test]
fn invalid_proxy_never_silently_becomes_direct() {
for url in [
"",
"localhost:1080",
"file:///etc/passwd",
"ftp://host:21",
"http://host:0",
"socks5://host",
"http://host:80/path",
"http://host:80?secret=1",
"http://host:80#fragment",
"http://user:secret@host:80\n",
] {
let error = OutboundProxy::parse(url).unwrap_err();
assert!(!error.to_string().contains("secret"));
}
}
@@ -2158,6 +2158,146 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() {
let state = store.state.lock().expect("store lock");
assert_eq!(state.finalizations[0].cost_source, CostSource::Unavailable);
assert_eq!(state.finalizations[0].cost_ticks, None);
assert_eq!(
session.budget_charge().amount_usd.scaled(),
999,
"missing final usage cannot erase a charged retry"
);
}
#[test]
fn key_budget_accumulates_only_observed_retry_costs() {
for (first_cost, first_send, expected) in [
(Some(999), UpstreamSendState::Sent, 1122),
(
Some(gateway_core::metering::Decimal::MAX.scaled()),
UpstreamSendState::Sent,
gateway_core::metering::Decimal::MAX.scaled(),
),
(None, UpstreamSendState::Sent, 123),
(None, UpstreamSendState::NotSent, 123),
] {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut first = Vec::new();
if first_send != UpstreamSendState::NotSent {
first.push(Ok(GatewayEvent::Started(ResponseMeta::new(
"resp_first",
"gpt-5",
))));
}
if let Some(cost) = first_cost {
first.push(Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
)));
}
first.push(Err(ProviderError::new(
ProviderErrorKind::Unavailable,
first_send,
)
.with_replay_safe()));
let (coordinator, _, _) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: first,
},
Script::Stream {
account_id: "acct_second",
items: vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"resp_done",
"gpt-5",
))),
Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(123).unwrap(),
)),
Ok(GatewayEvent::Completed(ResponseMeta::new(
"resp_done",
"gpt-5",
))),
],
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let charge = session.budget_charge();
assert_eq!(charge.amount_usd.scaled(), expected);
assert_eq!(
session.budget_charge().completed_at,
charge.completed_at,
"final settlement time stays stable"
);
}
}
#[test]
fn key_budget_settles_post_send_failures_from_observed_cost_without_replaying() {
for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] {
for cost in [None, Some(123)] {
for committed in [false, true] {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut items = vec![Ok(GatewayEvent::Started(ResponseMeta::new(
"response-interrupted",
"gpt-5",
)))];
if let Some(cost) = cost {
items.push(Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
)));
}
items.push(Err(ProviderError::new(
ProviderErrorKind::Transport,
send_state,
)));
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items,
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
if committed {
block_on(session.next_event()).unwrap().unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
loop {
match block_on(session.next_event()) {
Ok(Some(_)) => {}
Err(EngineError::Provider(_)) => break,
other => panic!("expected terminal transport error, got {other:?}"),
}
}
} else {
assert!(matches!(
block_on(session.collect_uncommitted()),
Err(EngineError::Provider(_))
));
}
assert!(session.is_finalized());
assert_eq!(
session.budget_charge().amount_usd.scaled(),
cost.unwrap_or(0),
"send={send_state:?}, committed={committed}, cost={cost:?}"
);
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
}
}
}
}
#[test]
@@ -2688,10 +2828,16 @@ fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once()
fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut failure = bare_atomic_response_failed("response-only-failed");
let mut events = failure.take_atomic_client_events();
events.push(ProviderEvent::canonical(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(500).unwrap(),
)));
let failure = failure.with_atomic_client_events(events);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_only",
items: vec![Err(bare_atomic_response_failed("response-only-failed"))],
items: vec![Err(failure)],
},
Script::Error(ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
@@ -2731,6 +2877,11 @@ fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire()
));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
assert_eq!(
session.budget_charge().amount_usd.scaled(),
500,
"restoring a discarded failure must not charge its usage twice"
);
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(state.intermediate_failures, 1);
@@ -2868,6 +3019,7 @@ fn ambiguous_send_state_stops_retry() {
let finalization = &state.finalizations[0];
assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
assert_eq!(finalization.send_state, UpstreamSendState::Ambiguous);
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
assert_eq!(finalization.attempt_count, 1);
assert!(!finalization.committed);
assert!(finalization.latency_ms.is_some());
@@ -3342,6 +3494,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
assert_eq!(state.commits, 0);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
assert!(!state.finalizations[0].committed);
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
}
#[test]
@@ -7,6 +7,231 @@ use std::{
time::{Duration, Instant, SystemTime},
};
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
use gateway_core::error::GatewayError;
#[derive(Default)]
struct Admissions {
active: Arc<AtomicBool>,
limits: Mutex<Vec<RateLimits>>,
}
impl ClientAdmissionPort for Admissions {
fn admit(
&self,
request: ClientAdmissionRequest,
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async move {
self.limits.lock().unwrap().push(request.limits);
assert!(!self.active.swap(true, Ordering::SeqCst));
Ok(ClientAdmissionDecision::Granted)
})
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
}
}
struct Budget {
reject: bool,
active: Arc<AtomicBool>,
charges: Mutex<Vec<ClientBudgetCharge>>,
}
impl ClientBudgetPort for Budget {
fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
Box::pin(async {
assert!(self.active.load(Ordering::SeqCst));
if self.reject {
Err(
GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted")
.with_client_code("key_daily_budget_exceeded"),
)
} else {
Ok(())
}
})
}
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
Box::pin(async move {
assert!(
self.active.load(Ordering::SeqCst),
"settle before releasing the concurrent request slot"
);
self.charges.lock().unwrap().push(charge);
Ok(())
})
}
}
fn service(admissions: Arc<Admissions>, budget: Arc<Budget>) -> DefaultExecutionService {
DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
.with_budget(budget)
}
fn request(service: &DefaultExecutionService, transport: ClientTransport) -> StartExecution {
StartExecution {
client: service.authenticate("sk_start_test").unwrap(),
public_model: PublicModelId::new("gpt-start").unwrap(),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport,
stream: transport != ClientTransport::HttpJson,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}
}
#[test]
fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
reject: true,
active: admissions.active.clone(),
charges: Mutex::default(),
});
let service = service(admissions.clone(), budget.clone());
for transport in [
ClientTransport::HttpJson,
ClientTransport::HttpSse,
ClientTransport::WebSocket,
] {
let result = block_on(service.start(request(&service, transport)));
assert!(
matches!(result, Err(error) if error.client_error_code() == Some("key_daily_budget_exceeded"))
);
assert!(!admissions.active.load(Ordering::SeqCst));
assert!(budget.charges.lock().unwrap().is_empty());
}
}
#[test]
fn reused_client_uses_updated_limits_for_each_execution() {
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
let admissions = Arc::new(Admissions::default());
let service = DefaultExecutionService::new(
snapshots.clone(),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service.authenticate("sk_start_test").unwrap();
let limited = RateLimits {
max_concurrency: 1,
requests_per_minute: 1,
};
for (revision, limits) in [(2, limited), (3, RateLimits::unlimited())] {
snapshots.publish(start_snapshot_with_policy(revision, true, limits));
let mut next = request(&service, ClientTransport::WebSocket);
next.client = client.clone();
let started = block_on(service.start(next)).expect("new execution");
block_on(started.session.detach_finalize());
}
assert_eq!(
*admissions.limits.lock().unwrap(),
vec![limited, RateLimits::unlimited()]
);
}
#[test]
fn reused_client_cannot_start_after_key_disable_or_snapshot_suspension() {
for suspend in [false, true] {
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
let admissions = Arc::new(Admissions::default());
let service = DefaultExecutionService::new(
snapshots.clone(),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let next = request(&service, ClientTransport::WebSocket);
if suspend {
snapshots.suspend();
} else {
snapshots.publish(start_snapshot_with_policy(
2,
false,
RateLimits::unlimited(),
));
}
let result = block_on(service.start(next));
let expected = if suspend {
GatewayErrorKind::Internal
} else {
GatewayErrorKind::Unauthorized
};
assert!(matches!(result, Err(error) if error.kind() == expected));
assert!(admissions.limits.lock().unwrap().is_empty());
}
}
#[test]
fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all_transports() {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
reject: false,
active: admissions.active.clone(),
charges: Mutex::default(),
});
let service = service(admissions.clone(), budget.clone());
for transport in [
ClientTransport::HttpJson,
ClientTransport::HttpSse,
ClientTransport::WebSocket,
] {
for detach in [true, false] {
let mut started = block_on(service.start(request(&service, transport))).unwrap();
assert!(admissions.active.load(Ordering::SeqCst));
if detach {
block_on(started.session.detach_finalize());
} else {
assert!(block_on(started.session.next_event()).is_err());
block_on(started.session.detach_finalize());
}
assert!(!admissions.active.load(Ordering::SeqCst));
}
}
let charges = budget.charges.lock().unwrap();
assert_eq!(
charges.len(),
6,
"one settlement per request, including detached finalized sessions"
);
assert!(
charges
.iter()
.all(|charge| charge.amount_usd == gateway_core::metering::Decimal::ZERO)
);
}
use async_trait::async_trait;
use bytes::Bytes;
use futures::{executor::block_on, future::BoxFuture};
@@ -495,7 +720,13 @@ fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
panic!("a known provider catalog must be authoritative for model availability");
};
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
assert_eq!(
(error.kind(), error.client_message()),
(
GatewayErrorKind::ModelNotFound,
"the requested model was not found in the provider catalogs available to this API key; check the model name",
)
);
}
#[test]
@@ -917,11 +1148,15 @@ fn client_snapshot() -> RuntimeSnapshot {
}
fn start_snapshot() -> RuntimeSnapshot {
start_snapshot_with_policy(1, true, RateLimits::unlimited())
}
fn start_snapshot_with_policy(revision: u64, enabled: bool, limits: RateLimits) -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
ConfigRevision::new(revision).expect("config revision"),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
std::num::NonZeroU32::new(1).expect("concurrency"),
@@ -937,8 +1172,8 @@ fn start_snapshot() -> RuntimeSnapshot {
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
account_scope(&provider, "acct_start"),
true,
RateLimits::unlimited(),
enabled,
limits,
)],
)
.expect("start snapshot")
+127 -13
View File
@@ -326,12 +326,123 @@ fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
)
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "gpt-5.4".to_owned(),
mapped_model: "gpt-5.5".to_owned(),
}
);
}
#[test]
fn missing_model_should_be_reported_when_all_scoped_catalogs_are_known() {
let snapshot = snapshot();
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect_err("neither provider publishes the requested model");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "missing-model".to_owned(),
mapped_model: "missing-model".to_owned(),
}
);
}
#[test]
fn missing_model_in_one_provider_should_still_route_to_another() {
let snapshot = snapshot();
let plan = snapshot
.plan(
&PublicModelId::new("grok-4.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("xAI publishes the model missing from OpenAI");
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
}
#[test]
fn blocked_model_provider_should_not_be_misreported_as_model_not_found() {
let snapshot = snapshot();
let error = snapshot
.plan(
&PublicModelId::new("gpt-5.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
..RoutingContext::default()
},
)
.expect_err("OpenAI is blocked and xAI does not publish this model");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn blocked_unknown_catalog_should_not_prove_model_absence() {
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![
ProviderKind::new("openai").expect("provider"),
ProviderKind::new("xai").expect("provider"),
],
vec![model("xai", "grok-4.5", capabilities())],
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory());
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
..RoutingContext::default()
},
)
.expect_err("the blocked provider's catalog cannot establish model absence");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn empty_account_scope_should_not_be_misreported_as_model_not_found() {
let snapshot = snapshot();
let empty_scope = Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::default()),
ClientRoutingScope::all_accounts(),
));
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
empty_scope,
&RoutingContext::default(),
)
.expect_err("the API key has no accounts");
assert_eq!(error, gateway_core::error::RoutingError::EmptyAccountScope);
}
#[test]
fn provider_endpoint_plan_should_not_consult_or_publish_the_text_model_catalog() {
let snapshot = snapshot();
@@ -463,19 +574,22 @@ fn known_unsupported_operation_should_not_be_bypassed() {
)
.expect("snapshot");
assert!(
snapshot
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.is_err()
);
let error = snapshot
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.expect_err("the model exists but cannot perform this operation");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
@@ -56,10 +56,13 @@ struct PublishingCatalogProvider {
queries: AtomicUsize,
}
struct UnavailableCatalogProvider;
enum TestCatalogProvider {
Unavailable,
Empty,
}
#[async_trait]
impl Provider for UnavailableCatalogProvider {
impl Provider for TestCatalogProvider {
fn name(&self) -> &'static str {
"alpha"
}
@@ -71,10 +74,13 @@ impl Provider for UnavailableCatalogProvider {
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
match self {
Self::Unavailable => Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Self::Empty => Ok(Vec::new()),
}
}
async fn execute(
@@ -145,7 +151,7 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
#[test]
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
let providers =
ProviderRegistry::new([Arc::new(UnavailableCatalogProvider) as Arc<dyn Provider>])
ProviderRegistry::new([Arc::new(TestCatalogProvider::Unavailable) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
@@ -162,6 +168,35 @@ fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
assert_eq!(snapshot.client_policies().count(), 1);
}
#[test]
fn known_empty_catalog_should_report_model_not_found() {
let providers =
ProviderRegistry::new([Arc::new(TestCatalogProvider::Empty) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let snapshot = block_on(compiler.compile()).expect("compile empty catalog");
let error = snapshot
.plan(
&PublicModelId::new("public-model").expect("model"),
&super::operation(),
snapshot.all_account_scope(),
&gateway_core::routing::RoutingContext {
required_provider: Some(ProviderKind::new("alpha").expect("provider")),
..gateway_core::routing::RoutingContext::default()
},
)
.expect_err("a successfully published empty catalog proves model absence");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "public-model".to_owned(),
mapped_model: "upstream-model".to_owned(),
},
);
}
#[test]
fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
let provider = Arc::new(PublishingCatalogProvider {
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-protocol"
version = "0.0.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-store"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
+2 -1
View File
@@ -171,7 +171,8 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
)?),
),
Arc::new(client_key_usage),
);
)
.with_budget(Arc::new(postgres::PgClientBudgetStore::new(pool.clone())));
let provider_ports = ProviderStorePorts::new(
account_store,
@@ -167,7 +167,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
"select membership.account_group_id,
account.id, account.provider_kind, account.name, account.email,
account.upstream_user_id, account.upstream_account_id, account.plan_type,
account.authentication_kind, account.credential_revision,
account.authentication_kind, account.credential_revision, account.outbound_proxy_url,
account.has_refresh_token, account.access_token_expires_at,
account.next_refresh_at, account.enabled, account.concurrency_limit,
account.weight, account.credential_state, account.quota_access_state,
@@ -0,0 +1,275 @@
//! 按 Key 串行检查限额,并幂等累计已取得的 USD 费用。
use std::{collections::BTreeMap, sync::Mutex, time::Duration};
use chrono::{DateTime, Utc};
use futures::future::BoxFuture;
use gateway_core::{
engine::budget::{
ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, ClientBudgetPort,
ClientBudgetStatus,
},
error::{GatewayError, GatewayErrorKind},
metering::Decimal,
policy::ClientApiKeyId,
};
use sqlx::{PgPool, Postgres, Row, Transaction};
use crate::{StoreResult, postgres_unavailable};
pub struct PgClientBudgetStore {
pool: PgPool,
retry: Mutex<BTreeMap<String, ClientBudgetCharge>>,
}
impl PgClientBudgetStore {
#[must_use]
pub fn new(pool: PgPool) -> Self {
Self {
pool,
retry: Mutex::new(BTreeMap::new()),
}
}
async fn admit_inner(&self, key_id: ClientApiKeyId) -> Result<(), GatewayError> {
// 短暂存储故障后按原金额重试;进程退出丢失的费用不转成人工核账或阻断 Key。
let retries = self
.retry
.lock()
.map_err(|_| unavailable())?
.values()
.filter(|charge| charge.key_id == key_id)
.cloned()
.collect::<Vec<_>>();
for charge in retries {
self.settle(charge).await.map_err(|_| unavailable())?;
}
let mut tx = self.pool.begin().await.map_err(|_| unavailable())?;
let row = sqlx::query(
"select daily_limit_usd::text, weekly_limit_usd::text, enabled
from client_api_keys where id = $1 for update",
)
.bind(key_id.as_str())
.fetch_optional(&mut *tx)
.await
.map_err(|_| unavailable())?
.ok_or_else(|| {
GatewayError::new(
GatewayErrorKind::Unauthorized,
"client API key no longer exists",
)
})?;
if !row.get::<bool, _>("enabled") {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"client API key is disabled",
));
}
let limits = ClientBudgetLimits {
daily_usd: row
.get::<String, _>("daily_limit_usd")
.parse()
.map_err(|_| unavailable())?,
weekly_usd: row
.get::<String, _>("weekly_limit_usd")
.parse()
.map_err(|_| unavailable())?,
};
let now = Utc::now();
advance_windows(&mut tx, key_id.as_str(), now)
.await
.map_err(|_| unavailable())?;
if limits.is_limited() {
let window = sqlx::query(
"select daily_used_usd::text, weekly_used_usd::text, daily_end, weekly_end
from client_key_budget_windows where client_api_key_id = $1",
)
.bind(key_id.as_str())
.fetch_one(&mut *tx)
.await
.map_err(|_| unavailable())?;
let daily: Decimal = window
.get::<String, _>("daily_used_usd")
.parse()
.map_err(|_| unavailable())?;
let weekly: Decimal = window
.get::<String, _>("weekly_used_usd")
.parse()
.map_err(|_| unavailable())?;
let daily_exceeded = limits.daily_usd != Decimal::ZERO && daily >= limits.daily_usd;
let weekly_exceeded = limits.weekly_usd != Decimal::ZERO && weekly >= limits.weekly_usd;
if daily_exceeded || weekly_exceeded {
let daily_end: DateTime<Utc> = window.get("daily_end");
let weekly_end: DateTime<Utc> = window.get("weekly_end");
let reset = if weekly_exceeded {
weekly_end
} else {
daily_end
};
let retry = (reset - now).to_std().unwrap_or(Duration::from_secs(1));
return Err(GatewayError::new(
GatewayErrorKind::RateLimited,
"client API key budget is exhausted",
)
.with_client_code(if weekly_exceeded {
"key_weekly_budget_exceeded"
} else {
"key_daily_budget_exceeded"
})
.with_retry_after(retry));
}
}
tx.commit().await.map_err(|_| unavailable())
}
async fn settle_inner(&self, charge: &ClientBudgetCharge) -> Result<(), ClientBudgetError> {
let mut tx = self.pool.begin().await.map_err(|_| ClientBudgetError)?;
// 与准入统一先锁 Key,再写窗口和费用,串行化同一 Key 的并发结算。
let key = sqlx::query_scalar::<_, String>(
"select id from client_api_keys where id = $1 for update",
)
.bind(charge.key_id.as_str())
.fetch_optional(&mut *tx)
.await
.map_err(|_| ClientBudgetError)?;
let Some(key) = key else { return Ok(()) }; // 删除 Key 时也会删除其费用记录。
settle_in_transaction(&mut tx, &key, charge)
.await
.map_err(|_| ClientBudgetError)?;
tx.commit().await.map_err(|_| ClientBudgetError)
}
}
async fn settle_in_transaction(
tx: &mut Transaction<'_, Postgres>,
key: &str,
charge: &ClientBudgetCharge,
) -> Result<(), sqlx::Error> {
advance_windows(tx, key, Utc::now()).await?;
// 仅在请求结束时写入费用;请求 ID 冲突时不重复累计。
let changed = sqlx::query(
"insert into client_key_charge_events (request_id, client_api_key_id, amount_usd, completed_at)
values ($1, $2, $3::text::numeric, $4)
on conflict (request_id) do nothing",
)
.bind(charge.request_id.as_str())
.bind(key)
.bind(charge.amount_usd.canonical())
.bind(DateTime::<Utc>::from(charge.completed_at))
.execute(&mut **tx)
.await?
.rows_affected();
if changed == 1 {
sqlx::query("update client_key_budget_windows set
daily_used_usd = daily_used_usd + case when $3 >= daily_start and $3 < daily_end then $2::text::numeric else 0 end,
weekly_used_usd = weekly_used_usd + case when $3 >= weekly_start and $3 < weekly_end then $2::text::numeric else 0 end
where client_api_key_id = $1")
.bind(key).bind(charge.amount_usd.canonical()).bind(DateTime::<Utc>::from(charge.completed_at))
.execute(&mut **tx).await?;
}
Ok(())
}
impl ClientBudgetPort for PgClientBudgetStore {
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
Box::pin(async move { self.admit_inner(key_id).await })
}
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
Box::pin(async move {
let result = self.settle_inner(&charge).await;
let mut retry = self.retry.lock().map_err(|_| ClientBudgetError)?;
if result.is_err() {
retry.insert(charge.request_id.as_str().to_owned(), charge);
} else {
retry.remove(charge.request_id.as_str());
}
result
})
}
}
async fn advance_windows(
tx: &mut Transaction<'_, Postgres>,
key: &str,
now: DateTime<Utc>,
) -> Result<(), sqlx::Error> {
sqlx::query("insert into client_key_budget_windows
(client_api_key_id, daily_start, daily_end, weekly_start, weekly_end)
select $1, day, day + interval '24 hours', day, day + interval '168 hours'
from (select date_trunc('day', $2::timestamptz at time zone 'Asia/Shanghai') at time zone 'Asia/Shanghai' as day) d
on conflict (client_api_key_id) do update set
daily_start = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_start else client_key_budget_windows.daily_start end,
daily_end = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_end else client_key_budget_windows.daily_end end,
daily_used_usd = case when client_key_budget_windows.daily_end <= $2 then 0 else client_key_budget_windows.daily_used_usd end,
weekly_start = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_start else client_key_budget_windows.weekly_start end,
weekly_end = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_end else client_key_budget_windows.weekly_end end,
weekly_used_usd = case when client_key_budget_windows.weekly_end <= $2 then 0 else client_key_budget_windows.weekly_used_usd end")
.bind(key).bind(now).execute(&mut **tx).await?;
Ok(())
}
pub(super) async fn load_client_key_budgets(
pool: &PgPool,
records: &mut [super::ClientApiKeyRecord],
) -> StoreResult<()> {
if records.is_empty() {
return Ok(());
}
let ids = records
.iter()
.map(|record| record.id.as_str())
.collect::<Vec<_>>();
let rows = sqlx::query(
"select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text,
(case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used,
(case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used,
case when w.daily_end > now() then w.daily_end end as daily_end,
case when w.weekly_end > now() then w.weekly_end end as weekly_end
from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id
where k.id = any($1)",
)
.bind(ids)
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load client budgets"))?;
let mut budgets = BTreeMap::new();
for row in rows {
let parse = |field| -> StoreResult<Decimal> {
row.get::<String, _>(field)
.parse()
.map_err(|_| postgres_unavailable("decode client budget"))
};
budgets.insert(
row.get::<String, _>("id"),
ClientBudgetStatus {
limits: ClientBudgetLimits {
daily_usd: parse("daily_limit_usd")?,
weekly_usd: parse("weekly_limit_usd")?,
},
daily_used_usd: parse("daily_used")?,
weekly_used_usd: parse("weekly_used")?,
daily_resets_at: row
.get::<Option<DateTime<Utc>>, _>("daily_end")
.map(Into::into),
weekly_resets_at: row
.get::<Option<DateTime<Utc>>, _>("weekly_end")
.map(Into::into),
},
);
}
for record in records {
record.budget = budgets
.remove(&record.id)
.ok_or_else(|| postgres_unavailable("load client budget policy"))?;
}
Ok(())
}
fn unavailable() -> GatewayError {
GatewayError::new(
GatewayErrorKind::ProviderInfrastructureUnavailable,
"client budget service is temporarily unavailable",
)
.with_client_code("key_budget_unavailable")
}
@@ -27,8 +27,10 @@ use gateway_admin::{
ports::store::{AdminStoreResult, ClientKeyStore},
};
use gateway_core::{
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
engine::execution::ClientApiKeyUsageSink,
lifecycle::CancellationToken,
metering::Decimal,
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
routing::{AccountGroupId, ProviderKind},
task::{DaemonTask, WorkerTaskError},
@@ -114,6 +116,7 @@ pub struct ClientApiKeyRecord {
pub enabled: bool,
pub max_concurrency: u64,
pub requests_per_minute: u64,
pub budget: ClientBudgetStatus,
pub last_used_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
@@ -271,6 +274,7 @@ pub struct NewClientApiKey {
pub key: String,
pub max_concurrency: u64,
pub requests_per_minute: u64,
pub budget: ClientBudgetLimits,
}
impl fmt::Debug for NewClientApiKey {
@@ -301,6 +305,8 @@ pub struct UpdateClientApiKeyDetails {
pub group_ids: Vec<String>,
pub max_concurrency: u64,
pub requests_per_minute: u64,
pub daily_limit_usd: Option<Decimal>,
pub weekly_limit_usd: Option<Decimal>,
}
impl UpdateClientApiKeyDetails {
@@ -383,6 +389,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
None
};
load_client_key_memberships(&self.pool, &mut items).await?;
super::client_budgets::load_client_key_budgets(&self.pool, &mut items).await?;
Ok(ClientApiKeyPage {
items,
total,
@@ -406,7 +413,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
require_nonempty(ENTITY, "id", id)?;
sqlx::query(
let record = sqlx::query(
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
@@ -445,7 +452,16 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
.map_err(|_| postgres_unavailable("get client API key"))?
.as_ref()
.map(client_record_from_row)
.transpose()
.transpose()?;
let Some(mut record) = record else {
return Ok(None);
};
super::client_budgets::load_client_key_budgets(
&self.pool,
std::slice::from_mut(&mut record),
)
.await?;
Ok(Some(record))
}
async fn touch_client_api_keys(
@@ -719,6 +735,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
.map(|id| id.as_str().to_owned())
.collect(),
key: command.plaintext,
budget: command.budget,
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
},
@@ -735,6 +752,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
"enabled",
"max_concurrency",
"requests_per_minute",
"daily_limit_usd",
"weekly_limit_usd",
]
.into_iter()
.map(str::to_owned)
@@ -766,6 +785,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
.collect(),
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
daily_limit_usd: command.daily_limit_usd,
weekly_limit_usd: command.weekly_limit_usd,
},
mutation_audit(
context,
@@ -778,6 +799,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
"group_ids",
"max_concurrency",
"requests_per_minute",
"daily_limit_usd",
"weekly_limit_usd",
]
.into_iter()
.map(str::to_owned)
@@ -942,6 +965,7 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<Admin
max_concurrency: record.max_concurrency,
requests_per_minute: record.requests_per_minute,
},
budget: record.budget,
last_used_at: record.last_used_at,
created_at: record.created_at,
updated_at: record.updated_at,
@@ -956,8 +980,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
sqlx::query(
"insert into client_api_keys (
id, name, label, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
last_used_at, created_at, updated_at, daily_limit_usd, weekly_limit_usd
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric)",
)
.bind(&key.id)
.bind(&key.name)
@@ -965,6 +989,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
.bind(&key.key)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.bind(key.budget.daily_usd.canonical())
.bind(key.budget.weekly_usd.canonical())
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
@@ -980,7 +1006,9 @@ pub(crate) async fn update_client_api_key_in_transaction(
let result = sqlx::query(
"update client_api_keys
set name = $2, label = $3, max_concurrency = $4,
requests_per_minute = $5, updated_at = now()
requests_per_minute = $5, updated_at = now(),
daily_limit_usd = coalesce($6::text::numeric, daily_limit_usd),
weekly_limit_usd = coalesce($7::text::numeric, weekly_limit_usd)
where id = $1",
)
.bind(&key.id)
@@ -988,6 +1016,8 @@ pub(crate) async fn update_client_api_key_in_transaction(
.bind(&key.label)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.bind(key.daily_limit_usd.map(|amount| amount.canonical()))
.bind(key.weekly_limit_usd.map(|amount| amount.canonical()))
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
@@ -1142,6 +1172,7 @@ fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiK
row.try_get("requests_per_minute")
.map_err(|_| invalid("invalid requests per minute"))?,
)?,
budget: ClientBudgetStatus::default(),
last_used_at: row
.try_get("last_used_at")
.map_err(|_| invalid("invalid last used at"))?,
@@ -15,6 +15,7 @@ mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
mod client_budgets;
mod client_keys;
mod execution;
mod execution_buffer;
@@ -30,6 +31,7 @@ pub use account_groups::*;
pub use admin_security_audit::*;
pub use admission_recovery::*;
pub use backup::*;
pub use client_budgets::PgClientBudgetStore;
pub use client_keys::*;
pub use execution::*;
pub use execution_buffer::*;
@@ -293,6 +293,7 @@ pub(crate) fn admin_dashboard_account_usage(
name: usage.name,
email: usage.email,
plan_type: usage.plan_type,
plan_type_display: None,
request_count: usage.request_count,
success_count: usage.success_count,
input_tokens: usage.input_tokens,
@@ -213,6 +213,7 @@ impl PgAdminAccountStore {
async fn commit_prepared_import(
&self,
prepared: PreparedCredentialImport,
settings: Option<AccountImportSettings>,
context: &MutationContext,
action: &str,
) -> AdminStoreResult<CredentialImportResult> {
@@ -223,9 +224,15 @@ impl PgAdminAccountStore {
.map(prepared_account)
.collect::<StoreResult<Vec<_>>>()
.map_err(|error| admin_store_error(ENTITY, error))?;
let mut changed_fields = vec!["credentials".to_owned()];
if settings.is_some() {
changed_fields
.extend(["enabled", "concurrency_limit", "weight", "group_ids"].map(str::to_owned));
}
let imported = self
.accounts
.import_provider_accounts(ImportProviderAccounts {
settings,
scope: ProviderAccountAdminScope {
provider_kind: provider_kind.clone(),
},
@@ -235,7 +242,7 @@ impl PgAdminAccountStore {
action,
"provider_account",
&provider_kind,
vec!["credentials".to_owned()],
changed_fields,
),
})
.await
@@ -615,8 +622,13 @@ impl AccountStore for PgAdminAccountStore {
command: CredentialImportCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialImportResult> {
self.commit_prepared_import(command.prepared, context, "import_document")
.await
self.commit_prepared_import(
command.prepared,
command.settings,
context,
"import_document",
)
.await
}
async fn commit_authorization(
@@ -635,6 +647,7 @@ impl AccountStore for PgAdminAccountStore {
provider_kind: credential.provider_kind.clone(),
credentials: vec![credential],
},
command.settings,
context,
"authorize",
)
@@ -666,6 +679,13 @@ impl AccountStore for PgAdminAccountStore {
})
}
AuthorizationCredentialCommit::Reauthorize(prepared) => {
if command.settings.is_some() {
return Err(AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"reauthorization cannot change account settings",
));
}
self.commit_prepared_rotation(prepared, context, "reauthorize")
.await
}
@@ -702,6 +722,15 @@ impl AccountStore for PgAdminAccountStore {
"invalid provider account ID",
)
})?;
let mut changed_fields = vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
];
if command.outbound_proxy.is_some() {
changed_fields.push("outbound_proxy".to_owned());
}
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
@@ -710,17 +739,13 @@ impl AccountStore for PgAdminAccountStore {
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
outbound_proxy: command.outbound_proxy,
audit: mutation_audit(
context,
"update",
"provider_account",
&command.account_id,
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
changed_fields,
),
})
.await
@@ -790,6 +815,15 @@ impl AccountStore for PgAdminAccountStore {
} else {
"provider_accounts".to_owned()
};
let mut changed_fields = vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
];
if command.outbound_proxy.is_some() {
changed_fields.push("outbound_proxy".to_owned());
}
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
@@ -798,17 +832,13 @@ impl AccountStore for PgAdminAccountStore {
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
outbound_proxy: command.outbound_proxy,
audit: mutation_audit(
context,
"batch_update",
"provider_account",
&audit_target,
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
changed_fields,
),
})
.await
@@ -116,7 +116,7 @@ pub(crate) async fn load_admin_account_page(
from runtime_settings
where id = 1
)
select a.id, a.provider_kind, a.name, a.email, a.upstream_user_id,
select a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.email, a.upstream_user_id,
a.upstream_account_id, a.plan_type, a.authentication_kind,
a.credential_revision, a.has_refresh_token, a.access_token_expires_at,
a.next_refresh_at, a.enabled, a.concurrency_limit, a.weight,
@@ -24,6 +24,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
)
.map_err(core_store_error)?;
self.insert_provider_account(NewProviderAccount {
outbound_proxy: account.account.outbound_proxy().cloned(),
id: account.account.id().as_str().to_owned(),
provider_kind: account.account.provider().as_str().to_owned(),
name: account.account.name().to_owned(),
@@ -6,6 +6,7 @@ pub(crate) fn admin_account_record(
summary: ProviderAccountSummary,
) -> AdminStoreResult<AccountRecord> {
Ok(AccountRecord {
outbound_proxy: summary.outbound_proxy,
id: summary.id,
provider_kind: ProviderKind::new(summary.provider_kind).map_err(|_| {
AdminStoreError::new(
@@ -42,6 +43,7 @@ pub(crate) fn prepared_account(
credential: PreparedCredentialCreate,
) -> StoreResult<NewProviderAccount> {
Ok(NewProviderAccount {
outbound_proxy: credential.outbound_proxy,
id: credential.account_id.as_str().to_owned(),
provider_kind: credential.provider_kind.as_str().to_owned(),
name: credential.name,
@@ -9,12 +9,13 @@ use gateway_admin::{
MutationContext, Revision as AdminRevision,
account_groups::AccountGroupRef,
accounts::{
AccountCost, AccountGroupFilter, AccountListQuery as AdminAccountListQuery,
AccountModelUsage, AccountPage, AccountPageItem, AccountRecord, AccountRequestBucket,
AccountSort as AdminAccountSort, AccountSortField as AdminAccountSortField,
AccountSummary, AccountUpdateResult, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, AccountsUpdateResult, BatchUpdateAccounts, DeleteAccounts,
SortDirection as AdminSortDirection, UpdateAccount,
AccountCost, AccountGroupFilter, AccountImportSettings,
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountPageItem, AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountSummary, AccountUpdateResult,
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, SortDirection as AdminSortDirection,
UpdateAccount,
},
observability::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
@@ -108,7 +108,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
include_disabled: bool,
) -> StoreResult<Vec<ProviderAccountSummary>> {
let rows = sqlx::query(
"select id, provider_kind, name, email, upstream_user_id,
"select outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
credential_observed_at, quota_access_state, quota_evidence,
@@ -135,13 +135,13 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
};
sqlx::query(
"insert into provider_accounts (
id, provider_kind, name, email, upstream_user_id,
outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
concurrency_limit, weight, credential_state, provider_quota_json,
credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$18, $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)",
)
@@ -162,6 +162,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.bind(account.outbound_proxy.as_ref().map(|proxy| proxy.expose_url()))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("insert provider account"))?;
@@ -447,6 +448,29 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
account_ids
.push(upsert_provider_account_in_transaction(&mut transaction, account).await?);
}
if let Some(settings) = &command.settings {
let unique_ids = account_ids
.iter()
.cloned()
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
update_provider_accounts_scheduling_in_transaction(
&mut transaction,
&unique_ids,
settings.enabled,
settings.concurrency_limit,
settings.weight,
None,
)
.await?;
replace_account_group_assignments_in_transaction(
&mut transaction,
&unique_ids,
&settings.group_ids,
)
.await?;
}
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
.await?;
Ok(ProviderAccountAdminImport {
@@ -524,6 +548,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
command.enabled,
command.concurrency_limit,
command.weight,
command.outbound_proxy.as_ref(),
)
.await?;
replace_account_group_assignments_in_transaction(
@@ -678,13 +703,13 @@ pub(crate) async fn upsert_provider_account_in_transaction(
};
let imported_id = sqlx::query_scalar::<_, String>(
"insert into provider_accounts (
id, provider_kind, name, email, upstream_user_id,
outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
concurrency_limit, weight, credential_state, provider_quota_json,
credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$18, $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)
on conflict (
@@ -697,6 +722,7 @@ pub(crate) async fn upsert_provider_account_in_transaction(
plan_type = excluded.plan_type,
authentication_kind = excluded.authentication_kind,
provider_credentials_json = excluded.provider_credentials_json,
outbound_proxy_url = coalesce(excluded.outbound_proxy_url, provider_accounts.outbound_proxy_url),
credential_revision = provider_accounts.credential_revision + 1,
has_refresh_token = excluded.has_refresh_token,
access_token_expires_at = excluded.access_token_expires_at,
@@ -732,6 +758,7 @@ pub(crate) async fn upsert_provider_account_in_transaction(
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.bind(account.outbound_proxy.as_ref().map(|proxy| proxy.expose_url()))
.fetch_optional(&mut **transaction)
.await
.map_err(|error| {
@@ -838,10 +865,12 @@ pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
enabled: bool,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
outbound_proxy: Option<&Option<gateway_core::account::OutboundProxy>>,
) -> StoreResult<()> {
let updated = sqlx::query_scalar::<_, String>(
"update provider_accounts
set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now()
set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = greatest(now(), updated_at),
outbound_proxy_url = case when $5 then $6 else outbound_proxy_url end
where id = any($1::text[])
returning id",
)
@@ -849,6 +878,12 @@ pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
.bind(enabled)
.bind(concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(outbound_proxy.is_some())
.bind(
outbound_proxy
.and_then(Option::as_ref)
.map(|proxy| proxy.expose_url()),
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))?
@@ -937,7 +972,7 @@ fn validate_batch_update_account_ids(account_ids: &[String]) -> StoreResult<()>
Ok(())
}
fn validate_batch_update_group_ids(group_ids: &[AccountGroupId]) -> StoreResult<()> {
pub(super) fn validate_batch_update_group_ids(group_ids: &[AccountGroupId]) -> StoreResult<()> {
const MAX_BATCH_UPDATE_GROUPS: usize = 1000;
if group_ids.len() > MAX_BATCH_UPDATE_GROUPS {
return Err(invalid("account batch update contains too many group IDs"));
@@ -120,6 +120,7 @@ pub(crate) fn parse_error_reason(value: Option<String>) -> StoreResult<Option<Ac
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountSummary {
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub id: String,
pub provider_kind: String,
pub name: String,
@@ -167,6 +168,7 @@ impl fmt::Debug for ProviderAccountRecord {
#[derive(Clone)]
pub struct NewProviderAccount {
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub id: String,
pub provider_kind: String,
pub name: String,
@@ -230,6 +232,7 @@ impl NewProviderAccount {
#[derive(Clone)]
pub struct ImportProviderAccounts {
pub settings: Option<AccountImportSettings>,
pub scope: ProviderAccountAdminScope,
pub accounts: Vec<NewProviderAccount>,
pub audit: AdminAuditEvent,
@@ -261,6 +264,9 @@ impl fmt::Debug for ImportProviderAccounts {
impl ImportProviderAccounts {
pub fn validate(&self) -> StoreResult<()> {
if let Some(settings) = &self.settings {
repository::validate_batch_update_group_ids(&settings.group_ids)?;
}
self.scope.validate()?;
if self.accounts.is_empty() || self.accounts.len() > MAX_ADMIN_IMPORT_BATCH {
return Err(invalid(
@@ -307,6 +313,7 @@ impl fmt::Debug for RotateProviderAccount {
#[derive(Debug, Clone)]
pub struct BatchUpdateProviderAccountsAdmin {
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
pub account_ids: Vec<String>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
@@ -372,7 +379,7 @@ impl ProviderAccountStateUpdate {
}
}
pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
pub(crate) const ACCOUNT_SELECT: &str = "select outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
@@ -380,7 +387,7 @@ pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email,
credential_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts where id = $1";
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id,
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
@@ -390,7 +397,7 @@ pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name,
where id = any($1::text[]) and provider_kind = $2
order by id";
pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select outbound_proxy_url, id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
@@ -465,6 +472,7 @@ pub(crate) fn core_account_from_summary(
summary.last_error_message,
)
.with_scheduling(summary.concurrency_limit, summary.weight)
.with_outbound_proxy(summary.outbound_proxy)
.with_refresh_schedule(
summary.has_refresh_token,
summary.next_refresh_at.map(Into::into),
@@ -523,6 +531,12 @@ pub(crate) fn account_summary_from_row(
.and_then(AccountWeight::new)
.ok_or_else(|| invalid("invalid weight"))?;
Ok(ProviderAccountSummary {
outbound_proxy: get::<Option<String>>(&row, "outbound_proxy_url")?
.map(|url| {
gateway_core::account::OutboundProxy::parse(&url)
.map_err(|_| invalid("invalid outbound proxy"))
})
.transpose()?,
id: get(&row, "id")?,
provider_kind: get(&row, "provider_kind")?,
name: get(&row, "name")?,
@@ -168,6 +168,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
let (scope_revision, widened) = keys
.update_client_key(
UpdateClientKey {
daily_limit_usd: None,
weekly_limit_usd: None,
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
@@ -191,6 +193,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
let (restricted_revision, restricted) = keys
.update_client_key(
UpdateClientKey {
daily_limit_usd: None,
weekly_limit_usd: None,
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
@@ -311,6 +315,7 @@ async fn group_costs_should_include_statusless_websocket_but_reject_statusless_h
fn new_key(id: &str, group_ids: Vec<AccountGroupId>) -> NewClientKey {
let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k'));
NewClientKey {
budget: Default::default(),
id: client_key_id(id),
name: id.to_owned(),
label: None,
@@ -0,0 +1,337 @@
use std::time::SystemTime;
use chrono::{DateTime, Utc};
use gateway_admin::{
model::{MutationActor, MutationContext, client_keys::UpdateClientKey},
ports::store::ClientKeyStore as _,
};
use gateway_core::{
engine::{
ModelRequestId,
budget::{ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus},
},
error::GatewayErrorKind,
policy::{ClientApiKeyId, RateLimits},
};
use gateway_store::postgres::{
ClientApiKeyRepository as _, PgAdminClientKeyStore, PgClientApiKeyRepository,
PgClientBudgetStore,
};
use super::TestDatabase;
fn key_id(key: &str) -> ClientApiKeyId {
ClientApiKeyId::new(key).unwrap()
}
fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge {
ClientBudgetCharge {
key_id: key_id(key),
request_id: ModelRequestId::new(format!("req_{request}")).unwrap(),
amount_usd: amount.parse().unwrap(),
completed_at: SystemTime::now(),
}
}
async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) {
sqlx::query(
"insert into client_api_keys (id, name, key, daily_limit_usd, weekly_limit_usd, created_at, updated_at)
values ($1, $1, $2, $3::text::numeric, $4::text::numeric, now(), now())",
)
.bind(key)
.bind(format!("sk_{key:a<43}"))
.bind(daily)
.bind(weekly)
.execute(&database.pool)
.await
.unwrap();
}
async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus {
PgClientApiKeyRepository::new(database.pool.clone())
.get_client_api_key(key)
.await
.unwrap()
.unwrap()
.budget
}
fn context() -> MutationContext {
MutationContext {
actor: MutationActor::System,
request_id: "budget-test".to_owned(),
}
}
#[tokio::test]
async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_instances() {
let Some(database) = TestDatabase::create("budgets_exact").await else {
return;
};
seed(&database, "day", "0.3", "2").await;
seed(&database, "week", "2", "0.2").await;
let first = PgClientBudgetStore::new(database.pool.clone());
let second = PgClientBudgetStore::new(database.pool.clone());
for (key, prefix, amount) in [("day", "d", "0.1"), ("week", "w", "0.1")] {
for _ in 0..3 {
// 已准入请求可完成并超过限额,不预占估算费用。
first.admit(key_id(key)).await.unwrap();
}
for index in 0..3 {
let id = format!("{prefix}-{index}");
let (a, b) = tokio::join!(
first.settle(charge(key, &id, amount)),
second.settle(charge(key, &id, amount))
);
a.unwrap();
b.unwrap();
}
}
let day = status(&database, "day").await;
assert_eq!(day.daily_used_usd.canonical(), "0.3");
assert_eq!(day.weekly_used_usd.canonical(), "0.3");
let day_error = second.admit(key_id("day")).await.unwrap_err();
assert_eq!(day_error.kind(), GatewayErrorKind::RateLimited);
assert_eq!(
day_error.client_error_code(),
Some("key_daily_budget_exceeded")
);
assert!(day_error.retry_after().is_some());
let week_error = first.admit(key_id("week")).await.unwrap_err();
assert_eq!(
week_error.client_error_code(),
Some("key_weekly_budget_exceeded")
);
let event_count: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
.fetch_one(&database.pool)
.await
.unwrap();
assert_eq!(event_count, 6, "rejected admission must not create charges");
database.close().await;
}
#[tokio::test]
async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlement() {
let Some(database) = TestDatabase::create("budgets_windows").await else {
return;
};
seed(&database, "key", "1", "2").await;
let store = PgClientBudgetStore::new(database.pool.clone());
store.admit(key_id("key")).await.unwrap();
let (day_start, day_end, week_end): (DateTime<Utc>, DateTime<Utc>, DateTime<Utc>) =
sqlx::query_as("select daily_start, daily_end, weekly_end from client_key_budget_windows where client_api_key_id = 'key'")
.fetch_one(&database.pool).await.unwrap();
assert_eq!(day_start.timestamp().rem_euclid(86400), 16 * 3600);
assert_eq!((day_end - day_start).num_hours(), 24);
assert_eq!((week_end - day_start).num_hours(), 168);
store.settle(charge("key", "first", "1")).await.unwrap();
sqlx::query("update client_key_budget_windows set daily_start = daily_start - interval '1 day', daily_end = daily_start")
.execute(&database.pool).await.unwrap();
store.admit(key_id("key")).await.unwrap();
assert_eq!(
status(&database, "key").await.daily_used_usd.canonical(),
"0"
);
assert_eq!(
status(&database, "key").await.weekly_used_usd.canonical(),
"1"
);
store.settle(charge("key", "second", "1")).await.unwrap();
sqlx::query("update client_key_budget_windows set daily_end = now() - interval '1 second', weekly_end = now() - interval '1 second'")
.execute(&database.pool).await.unwrap();
let virtual_reset = status(&database, "key").await;
assert_eq!(virtual_reset.daily_used_usd.canonical(), "0");
assert_eq!(virtual_reset.weekly_used_usd.canonical(), "0");
store.admit(key_id("key")).await.unwrap();
let old = ClientBudgetCharge {
completed_at: (day_start - chrono::Duration::seconds(1)).into(),
..charge("key", "after-reset", "0.9")
};
store.settle(old).await.unwrap();
let reset = status(&database, "key").await;
assert_eq!(reset.daily_used_usd.canonical(), "0");
assert_eq!(reset.weekly_used_usd.canonical(), "0");
assert!(reset.weekly_resets_at.is_some());
database.close().await;
}
#[tokio::test]
async fn zero_cost_and_interrupted_requests_never_block_limited_keys() {
let Some(database) = TestDatabase::create("budgets_zero").await else {
return;
};
seed(&database, "key", "1", "5").await;
let store = PgClientBudgetStore::new(database.pool.clone());
store.admit(key_id("key")).await.unwrap();
store.settle(charge("key", "no-cost", "0")).await.unwrap();
// 模拟准入后进程退出,没有费用可结算;重启后仍应允许同一 Key 使用。
store.admit(key_id("key")).await.unwrap();
let restarted = PgClientBudgetStore::new(database.pool.clone());
restarted.admit(key_id("key")).await.unwrap();
let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
.fetch_one(&database.pool)
.await
.unwrap();
assert_eq!(events, 1, "admission must not leave pending charges");
assert_eq!(
status(&database, "key").await.daily_used_usd.canonical(),
"0"
);
restarted
.settle(charge("key", "known", "0.4"))
.await
.unwrap();
assert_eq!(
status(&database, "key").await.daily_used_usd.canonical(),
"0.4"
);
database.close().await;
}
#[tokio::test]
async fn transient_settlement_failure_rolls_back_and_retries_exact_cost_before_admission() {
let Some(database) = TestDatabase::create("budgets_retry").await else {
return;
};
seed(&database, "key", "1", "5").await;
let store = PgClientBudgetStore::new(database.pool.clone());
store.admit(key_id("key")).await.unwrap();
// 在费用事件插入后让窗口写入失败,验证整个事务回滚。
sqlx::raw_sql(
"create function reject_budget_update() returns trigger language plpgsql as $$
begin
if new.daily_used_usd > 0 then raise exception 'temporary test failure'; end if;
return new;
end $$;
create trigger reject_budget_update before update on client_key_budget_windows
for each row execute function reject_budget_update();",
)
.execute(&database.pool)
.await
.unwrap();
assert!(store.settle(charge("key", "retry", "1.25")).await.is_err());
let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
.fetch_one(&database.pool)
.await
.unwrap();
assert_eq!(events, 0);
assert_eq!(
status(&database, "key").await.daily_used_usd.canonical(),
"0"
);
sqlx::query("drop trigger reject_budget_update on client_key_budget_windows")
.execute(&database.pool)
.await
.unwrap();
let error = store.admit(key_id("key")).await.unwrap_err();
assert_eq!(error.client_error_code(), Some("key_daily_budget_exceeded"));
store.settle(charge("key", "retry", "1.25")).await.unwrap();
assert_eq!(
status(&database, "key").await.daily_used_usd.canonical(),
"1.25"
);
database.close().await;
}
#[tokio::test]
async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() {
let Some(database) = TestDatabase::create("budgets_policy").await else {
return;
};
seed(&database, "key", "0", "0").await;
let store = PgClientBudgetStore::new(database.pool.clone());
let admin = PgAdminClientKeyStore::new(database.pool.clone());
store.admit(key_id("key")).await.unwrap();
store
.settle(charge("key", "unlimited", "2.75"))
.await
.unwrap();
let update = UpdateClientKey {
id: ClientApiKeyId::new("key").unwrap(),
name: "key".to_owned(),
label: None,
group_ids: vec![],
limits: RateLimits {
max_concurrency: 3,
requests_per_minute: 0,
},
daily_limit_usd: Some("2".parse().unwrap()),
weekly_limit_usd: Some("10".parse().unwrap()),
};
admin
.update_client_key(update.clone(), &context())
.await
.unwrap();
admin
.update_client_key(
UpdateClientKey {
daily_limit_usd: None,
weekly_limit_usd: None,
..update.clone()
},
&context(),
)
.await
.unwrap();
let current = status(&database, "key").await;
assert_eq!(current.limits.daily_usd.canonical(), "2");
assert_eq!(current.limits.weekly_usd.canonical(), "10");
assert_eq!(current.daily_used_usd.canonical(), "2.75");
assert_eq!(
store
.admit(key_id("key"))
.await
.unwrap_err()
.client_error_code(),
Some("key_daily_budget_exceeded")
);
admin
.update_client_key(
UpdateClientKey {
daily_limit_usd: Some("0".parse().unwrap()),
weekly_limit_usd: None,
..update
},
&context(),
)
.await
.unwrap();
store.admit(key_id("key")).await.unwrap();
sqlx::query("update client_api_keys set enabled = false where id = 'key'")
.execute(&database.pool)
.await
.unwrap();
assert_eq!(
store.admit(key_id("key")).await.unwrap_err().kind(),
GatewayErrorKind::PolicyDenied
);
sqlx::query("delete from client_api_keys where id = 'key'")
.execute(&database.pool)
.await
.unwrap();
assert_eq!(
store.admit(key_id("key")).await.unwrap_err().kind(),
GatewayErrorKind::Unauthorized
);
store.settle(charge("key", "allowed", "1")).await.unwrap();
database.close().await;
}
#[tokio::test]
async fn budget_database_outage_fails_closed() {
let Some(database) = TestDatabase::create("budgets_outage").await else {
return;
};
let store = PgClientBudgetStore::new(database.pool.clone());
database.pool.close().await;
assert_eq!(
store
.admit(key_id("key"))
.await
.unwrap_err()
.client_error_code(),
Some("key_budget_unavailable")
);
assert!(store.settle(charge("key", "offline", "1")).await.is_err());
database.close().await;
}
@@ -24,6 +24,7 @@ use super::TestDatabase;
#[test]
fn client_key_requires_the_frozen_plaintext_format() {
let key = NewClientApiKey {
budget: Default::default(),
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
@@ -382,6 +383,7 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
fn client_key_debug_redacts_plaintext() {
let secret = format!("sk_{}", "s".repeat(43));
let key = NewClientApiKey {
budget: Default::default(),
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
@@ -14,6 +14,7 @@ mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
mod client_budgets;
mod client_keys;
mod execution;
mod execution_buffer;
@@ -230,6 +231,8 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
"backup_settings",
"client_api_key_groups",
"client_api_keys",
"client_key_budget_windows",
"client_key_charge_events",
"model_requests",
"ops_events",
"provider_accounts",
@@ -1124,6 +1124,7 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
let result = store
.update_account(
UpdateAccount {
outbound_proxy: None,
account_id: "acct_terminal_mutation".to_owned(),
enabled: false,
concurrency_limit: None,
@@ -1184,6 +1185,78 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
database.close().await;
}
#[tokio::test]
async fn account_proxy_edits_preserve_credentials_and_clear_egress_without_audit_secrets() {
let Some(database) = TestDatabase::create("account_proxy").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let mut seed = account("acct_proxy", "proxy-user");
seed.outbound_proxy = Some(
gateway_core::account::OutboundProxy::parse("http://initial:secret@127.0.0.1:18080")
.unwrap(),
);
repository.insert_provider_account(seed).await.unwrap();
let store = admin_account_store(&database.pool);
let context = MutationContext {
actor: MutationActor::System,
request_id: "proxy-edit".to_owned(),
};
let command = UpdateAccount {
account_id: "acct_proxy".to_owned(),
enabled: true,
concurrency_limit: None,
weight: gateway_core::account::AccountWeight::DEFAULT,
group_ids: vec![],
outbound_proxy: None,
};
store
.update_account(command.clone(), &context)
.await
.unwrap();
let read = || async {
sqlx::query_as::<_, (Option<String>, i64)>("select outbound_proxy_url, credential_revision from provider_accounts where id = 'acct_proxy'").fetch_one(&database.pool).await.unwrap()
};
assert_eq!(
read().await,
(Some("http://initial:secret@127.0.0.1:18080/".to_owned()), 1)
);
store
.update_account(
UpdateAccount {
outbound_proxy: Some(Some(
gateway_core::account::OutboundProxy::parse(
"socks5h://next:new-secret@127.0.0.1:1080",
)
.unwrap(),
)),
..command.clone()
},
&context,
)
.await
.unwrap();
assert_eq!(read().await.1, 1);
store
.update_account(
UpdateAccount {
outbound_proxy: Some(None),
..command
},
&context,
)
.await
.unwrap();
assert_eq!(read().await, (None, 1));
let audits: Vec<serde_json::Value> =
sqlx::query_scalar("select to_jsonb(a) from admin_audit_events a")
.fetch_all(&database.pool)
.await
.unwrap();
assert!(!serde_json::to_string(&audits).unwrap().contains("secret"));
database.close().await;
}
#[tokio::test]
async fn account_recovery_resets_status_facts_without_changing_credentials_or_scheduling() {
const GROUP_ID: &str = "grp_00000000000000000000000000000070";
@@ -1342,6 +1415,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever
let result = store
.batch_update_accounts(
BatchUpdateAccounts {
outbound_proxy: None,
account_ids: account_ids.clone(),
enabled: false,
concurrency_limit: gateway_core::account::AccountConcurrencyLimit::new(7),
@@ -1377,6 +1451,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever
store
.batch_update_accounts(
BatchUpdateAccounts {
outbound_proxy: None,
account_ids: account_ids.clone(),
enabled: true,
concurrency_limit: None,
@@ -1470,6 +1545,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
};
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![account("acct_admin_upsert", "user-admin-upsert")],
audit: audit("audit_admin_upsert_create", "import", "acct_admin_upsert"),
@@ -1492,6 +1568,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
updated.credential_state = CredentialState::Banned;
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![updated],
audit: audit("audit_admin_upsert_update", "import", "acct_admin_upsert"),
@@ -1526,6 +1603,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope,
accounts: vec![account("acct_admin_upsert", "another-user")],
audit: audit("audit_admin_upsert_rebind", "import", "acct_admin_upsert"),
@@ -1563,6 +1641,12 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
let result = admin_account_store(&database.pool)
.commit_authorization(
AuthorizationCommit {
settings: Some(gateway_admin::model::accounts::AccountImportSettings {
enabled: false,
concurrency_limit: None,
weight: gateway_core::account::AccountWeight::new(9).expect("weight"),
group_ids: Vec::new(),
}),
pending: PendingAuthorizationMutation::new(
provider_kind.clone(),
AuthorizationMutationTarget::Create {
@@ -1571,6 +1655,7 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
AuthorizationOwnerBinding::from_context(&context),
),
credential: AuthorizationCredentialCommit::Create(PreparedCredentialCreate {
outbound_proxy: None,
account_id: ProviderAccountId::new("acct_authorization_candidate")
.expect("candidate account ID"),
provider_kind,
@@ -1603,6 +1688,8 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
),
("acct_authorization_existing", Some(2)),
);
let settings: (bool, Option<i64>, i16) = sqlx::query_as("select enabled, concurrency_limit, weight from provider_accounts where id = 'acct_authorization_existing'").fetch_one(&database.pool).await.expect("OAuth settings");
assert_eq!(settings, (false, None, 9));
database.close().await;
}
@@ -1614,6 +1701,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
let repository = PgProviderAccountRepository::new(database.pool.clone());
repository
.insert_provider_account(NewProviderAccount {
outbound_proxy: None,
id: "acct_core_refresh".to_owned(),
provider_kind: "xai".to_owned(),
name: "before refresh".to_owned(),
@@ -1795,6 +1883,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![
account("acct_admin_a", "user-admin-a"),
@@ -1822,6 +1911,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![
account("acct_admin_transient", "user-admin-transient"),
@@ -1947,6 +2037,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
let revision = repository
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
outbound_proxy: None,
account_ids: vec!["acct_admin_a".to_owned()],
enabled: false,
concurrency_limit: None,
@@ -2009,6 +2100,7 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
};
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![account("acct_grouped_import", "user-grouped-import")],
audit: audit("audit_grouped_import", "import", "acct_grouped_import"),
@@ -2042,6 +2134,7 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
.expect("set account scheduling before reimport");
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![account("acct_reimport_candidate", "user-grouped-import")],
audit: audit("audit_grouped_reimport", "import", "acct_grouped_import"),
@@ -2109,6 +2202,7 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
};
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![account("acct_rotation_quota", "user-rotation-quota")],
audit: audit(
@@ -2250,6 +2344,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
disabled.credential_state = CredentialState::Expired;
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
scope: scope.clone(),
accounts: vec![disabled],
audit: audit(
@@ -2305,6 +2400,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
NewProviderAccount {
outbound_proxy: None,
id: id.to_owned(),
provider_kind: "openai".to_owned(),
name: id.to_owned(),
@@ -2473,3 +2569,167 @@ fn imported_ready_account_passes_validation() {
let imported = account("acct_ready_import", "user-ready-import");
assert!(imported.validate().is_ok());
}
#[tokio::test]
async fn proxy_edit_preserves_an_inflight_token_refresh() {
let Some(database) = TestDatabase::create("review_proxy_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let id = ProviderAccountId::new("acct_review_proxy").unwrap();
let mut seed = account(id.as_str(), "review-proxy-user");
seed.has_refresh_token = true;
repository.insert_provider_account(seed).await.unwrap();
let loaded = repository.load_current_credential(&id).await.unwrap();
let refreshed = CredentialCasUpdate::new(
id.clone(),
loaded.account.revision(),
ProviderAccountUpdate {
account_id: id.clone(),
name: loaded.account.name().to_owned(),
email: loaded.account.email().map(str::to_owned),
plan_type: loaded.account.plan_type().map(str::to_owned),
},
plaintext_credential("refreshed-marker"),
true,
Some(SystemTime::now() + Duration::from_secs(3600)),
None,
)
.unwrap();
admin_account_store(&database.pool)
.update_account(
UpdateAccount {
account_id: id.as_str().to_owned(),
enabled: true,
concurrency_limit: None,
weight: gateway_core::account::AccountWeight::DEFAULT,
group_ids: vec![],
outbound_proxy: Some(Some(
gateway_core::account::OutboundProxy::parse("http://127.0.0.1:18080").unwrap(),
)),
},
&MutationContext {
actor: MutationActor::System,
request_id: "review-proxy-change".to_owned(),
},
)
.await
.unwrap();
let result = repository
.compare_and_swap_credential(refreshed)
.await
.unwrap();
let saved = repository.load_current_credential(&id).await.unwrap();
let saved_new_token =
saved.credential.expose_to_provider()["access_token"] == "refreshed-marker";
database.close().await;
assert!(
matches!(result, CredentialCasOutcome::Updated(_)),
"Changing only egress must preserve a refresh of unchanged OAuth credentials"
);
assert!(saved_new_token);
assert_eq!(
saved.account.outbound_proxy().unwrap().expose_url(),
"http://127.0.0.1:18080/"
);
}
#[tokio::test]
async fn account_import_settings_apply_atomically_to_new_and_existing_identities() {
use gateway_admin::model::accounts::AccountImportSettings;
use gateway_core::account::{AccountConcurrencyLimit, AccountWeight};
const GROUP_ID: &str = "grp_00000000000000000000000000000091";
let Some(database) = TestDatabase::create("import_settings").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
repository
.insert_provider_account(account("acct_existing_settings", "existing-settings-user"))
.await
.expect("seed account");
sqlx::query(
"insert into account_groups (id, name, color, created_at, updated_at) values ($1, 'Import settings', '#2563EBFF', now(), now())",
)
.bind(GROUP_ID)
.execute(&database.pool)
.await
.expect("seed group");
let settings = AccountImportSettings {
enabled: false,
concurrency_limit: Some(AccountConcurrencyLimit::new(3).expect("concurrency")),
weight: AccountWeight::new(7).expect("weight"),
group_ids: vec![AccountGroupId::new(GROUP_ID).expect("group ID")],
};
let result = repository
.import_provider_accounts(ImportProviderAccounts {
settings: Some(settings.clone()),
scope: ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
},
accounts: vec![
account("acct_new_settings", "new-settings-user"),
account("acct_existing_candidate", "existing-settings-user"),
account("acct_duplicate_candidate", "existing-settings-user"),
],
audit: audit("audit_import_settings", "import", "provider_accounts"),
})
.await
.expect("import with settings");
assert_eq!(result.config_revision.get(), 2);
for id in ["acct_new_settings", "acct_existing_settings"] {
let row: (bool, Option<i64>, i16) = sqlx::query_as(
"select enabled, concurrency_limit, weight from provider_accounts where id = $1",
)
.bind(id)
.fetch_one(&database.pool)
.await
.expect("saved settings");
assert_eq!(row, (false, Some(3), 7));
assert_eq!(account_group_ids(&database.pool, id).await, [GROUP_ID]);
}
let before = repository
.load_provider_account("acct_existing_settings")
.await
.expect("existing account");
let failed = repository
.import_provider_accounts(ImportProviderAccounts {
settings: Some(AccountImportSettings {
group_ids: vec![
AccountGroupId::new("grp_00000000000000000000000000000092")
.expect("missing group"),
],
..settings
}),
scope: ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
},
accounts: vec![
account("acct_rollback_settings", "rollback-settings-user"),
account("acct_rollback_candidate", "existing-settings-user"),
],
audit: audit(
"audit_import_settings_failed",
"import",
"provider_accounts",
),
})
.await;
assert!(failed.is_err());
assert_eq!(current_revision(&database.pool).await, 2);
assert_eq!(
account_count(&database.pool, "acct_rollback_settings").await,
0
);
assert_eq!(
repository
.load_provider_account("acct_existing_settings")
.await
.expect("unchanged account"),
before
);
assert_eq!(
account_group_ids(&database.pool, "acct_existing_settings").await,
[GROUP_ID]
);
database.close().await;
}
+3 -1
View File
@@ -26,8 +26,9 @@ hex = "0.4.3"
aws-lc-rs = "=1.18.1"
hyper = { version = "=1.11.1", default-features = false }
hyper-rustls = "=0.27.7"
hyper-util = { version = "0.1.20", features = ["client-legacy", "http1", "tokio"] }
# 与官方 Codex Desktop 一致:reqwest 不启用压缩特性,上游请求不带 accept-encoding。
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2", "socks"] }
roxmltree = "0.21.1"
semver = "1.0.27"
# 与官方 rustls 路径一致:aws-lc 支持 P-521,优先提供 X25519MLKEM768。
@@ -44,6 +45,7 @@ thiserror = "2.0.18"
tokio = { version = "=1.53.1", default-features = false, features = ["fs", "macros", "rt", "sync", "time"] }
tokio-rustls = { version = "=0.26.4", default-features = false }
tokio-util = { version = "0.7.18", features = ["rt"] }
tower-service = "0.3.3"
# 官方 Codex WebSocket 需要上游尚未发布的 proxy 与 permessage-deflate 实现,固定到已审计提交。
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "0e5b2d73aa18dd9f0a50ee9ff199d5aef7594186", features = ["proxy", "rustls-tls-native-roots"] }
tracing = "0.1.44"
+60 -2
View File
@@ -166,6 +166,30 @@ impl ProviderAdmin for OpenAiAdminProvider {
&self.provider_kind
}
fn plan_type_display(&self, plan_type: &str) -> String {
// 与官方 Desktop 的套餐名称映射一致;原始子类型继续由 plan_type 保留。
match plan_type.to_ascii_lowercase().as_str() {
"free" | "free_workspace" | "guest" => "Free",
"go" => "Go",
"plus" => "Plus",
"pro" | "prolite" => "Pro",
"team" | "self_serve_business_prolite" | "self_serve_business_usage_based" => {
"Business"
}
"business"
| "ent26"
| "enterprise_cbp_automation"
| "enterprise_cbp_usage_based"
| "enterprise"
| "hc" => "Enterprise",
"edu" | "education" => "Edu",
"edu_plus" => "Edu Plus",
"edu_pro" => "Edu Pro",
_ => plan_type,
}
.to_owned()
}
async fn account_unavailable(&self, account_id: &ProviderAccountId) {
self.websocket_pool.evict_account(account_id.as_str()).await;
}
@@ -595,6 +619,7 @@ fn prepared_create(
credential,
} = prepared;
Ok(PreparedCredentialCreate {
outbound_proxy: account.outbound_proxy().cloned(),
account_id: account.id().clone(),
provider_kind: account.provider().clone(),
name: account.name().to_owned(),
@@ -698,6 +723,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
revision,
account.access_token_expires_at.map(SystemTime::from),
)
.with_outbound_proxy(account.outbound_proxy.clone())
.with_profile(
account.email.clone(),
account.upstream_account_id.clone(),
@@ -765,6 +791,7 @@ fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) ->
fn empty_quota() -> ProviderQuota {
ProviderQuota {
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
@@ -825,6 +852,7 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
// 在窗口全部过期后继续维持限流。
let limit_reached = quota_windows_limit_reached(&windows);
ProviderQuota {
plan_type: snapshot.plan_type().map(str::to_owned),
observed_at: Some(DateTime::<Utc>::from(snapshot.observed_at())),
refresh_token_expires_at: None,
windows,
@@ -1282,10 +1310,29 @@ fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError
}
fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdminError {
use crate::credential::token_client::PersonalAccessTokenError;
use CodexCredentialAdminError as Error;
let upstream_message = error.upstream_message().map(ToOwned::to_owned);
let public_message = match &error {
Error::PersonalAccessToken(error) => Some(match error {
PersonalAccessTokenError::InvalidToken => {
"Codex PAT 格式无效:应为 at- 开头的完整令牌,不能包含空白或控制字符"
}
PersonalAccessTokenError::Rejected => {
"OpenAI 拒绝了 Codex PAT:令牌可能无效、已过期、已撤销或没有访问权限"
}
PersonalAccessTokenError::Unavailable => "暂时无法向 OpenAI 验证 Codex PAT,请稍后重试",
PersonalAccessTokenError::InvalidResponse => {
"OpenAI 返回的 Codex PAT 身份资料不完整或格式无效,请稍后重试"
}
}),
_ => None,
};
let kind = match error {
Error::InvalidInput
Error::PersonalAccessToken(
PersonalAccessTokenError::InvalidToken | PersonalAccessTokenError::Rejected,
)
| Error::InvalidInput
| Error::InvalidCredential
| Error::MissingRefreshToken
| Error::RefreshRejected { .. }
@@ -1294,9 +1341,17 @@ fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdmin
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous { .. } => {
ProviderAdminErrorKind::Conflict
}
Error::RefreshUnavailable => ProviderAdminErrorKind::Unavailable,
Error::PersonalAccessToken(PersonalAccessTokenError::InvalidResponse) => {
ProviderAdminErrorKind::BadGateway
}
Error::PersonalAccessToken(PersonalAccessTokenError::Unavailable)
| Error::RefreshUnavailable => ProviderAdminErrorKind::Unavailable,
};
let error = provider_admin_error(kind);
let error = match public_message {
Some(message) => error.with_public_message(message),
None => error,
};
match upstream_message {
Some(message) => error.with_message(message),
None => error,
@@ -1305,6 +1360,9 @@ fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdmin
const fn credential_admin_error_code(error: &CodexCredentialAdminError) -> &'static str {
match error {
CodexCredentialAdminError::PersonalAccessToken(_) => {
"personal_access_token_validation_failed"
}
CodexCredentialAdminError::InvalidInput => "invalid_input",
CodexCredentialAdminError::InvalidCredential => "invalid_credential",
CodexCredentialAdminError::NotFound => "not_found",
@@ -26,7 +26,9 @@ use thiserror::Error;
use super::recovery_log::{CodexOAuthRecoveryOperation, record_oauth_recovery};
use super::security::CodexCredentialCodec;
use super::token_client::{RefreshFailure, TokenRefresher};
use super::token_client::{
OpenAiTokenClient, PersonalAccessTokenError, RefreshFailure, TokenRefresher,
};
use super::types::{
CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile, CodexCredentialData,
CodexCredentialPrincipal, CodexOAuthMetadata, CodexOAuthSecret, parse_access_token_expiration,
@@ -48,8 +50,8 @@ pub struct ImportCodexOAuthCredential {
/// OAuth credential 的最小创建输入。
///
/// metadata 来自 ID token/access token 的官方本地 payload 解析;不包含签名、
/// issuer、audience、nonce 或 token 验证。
/// OAuth metadata 来自 ID token/access token 的本地 payload 解析;PAT metadata
/// 来自 auth whoami 验证。两者均不从导入文档信任身份字段。
pub(crate) struct UnresolvedCodexOAuthCredential {
pub(crate) account_id: String,
pub(crate) name: String,
@@ -122,6 +124,7 @@ struct ParsedCodexImportAccount {
name: Option<String>,
email: Option<String>,
authentication: ParsedCodexAuthentication,
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
}
struct ParsedCodexAuthentication {
@@ -226,6 +229,8 @@ struct CodexCprExportCommon {
status: &'static str,
added_at: String,
updated_at: String,
#[serde(skip_serializing_if = "Option::is_none")]
outbound_proxy_url: Option<String>,
}
#[derive(Serialize)]
@@ -321,6 +326,8 @@ impl fmt::Debug for PreparedCodexCredentialRotationGuard {
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum CodexCredentialAdminError {
#[error(transparent)]
PersonalAccessToken(#[from] PersonalAccessTokenError),
#[error("Codex account input is invalid")]
InvalidInput,
#[error("Codex credential JSON is invalid")]
@@ -348,7 +355,8 @@ impl CodexCredentialAdminError {
Self::RefreshRejected { message }
| Self::AccountBanned { message }
| Self::RefreshAmbiguous { message } => message.as_deref(),
Self::InvalidInput
Self::PersonalAccessToken(_)
| Self::InvalidInput
| Self::InvalidCredential
| Self::NotFound
| Self::MissingRefreshToken
@@ -414,7 +422,7 @@ impl CodexCredentialAdmin {
/// 为已取得 OAuth access token、但资料尚未补全的账号创建最小记录。
///
/// ID token payload 只用于本地字段投影;创建路径绝不发起 usage/profile 请求。
/// 接受调用方已取得的身份投影;本方法不发起 usage/profile 请求。
pub(crate) fn prepare_unresolved_oauth(
&self,
input: UnresolvedCodexOAuthCredential,
@@ -502,6 +510,9 @@ impl CodexCredentialAdmin {
status: cpr_status(&account),
added_at: china_rfc3339(item.added_at),
updated_at: china_rfc3339(item.updated_at),
outbound_proxy_url: account
.outbound_proxy()
.map(|proxy| proxy.expose_url().to_owned()),
};
let exported = match data {
CodexCredentialData::OAuth(data) => {
@@ -662,6 +673,7 @@ impl CodexCredentialAdmin {
/// 有状态的 Codex 手工刷新边界;消费调用方刚读取的当前 credential 并准备 CAS。
pub struct CodexCredentialAdminService {
refresher: Arc<dyn TokenRefresher>,
personal_access_token_client: Option<Arc<OpenAiTokenClient>>,
leases: Arc<dyn ProviderLeasePort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
}
@@ -671,6 +683,10 @@ impl fmt::Debug for CodexCredentialAdminService {
formatter
.debug_struct("CodexCredentialAdminService")
.field("refresher", &"TokenRefresher")
.field(
"personal_access_token_client",
&self.personal_access_token_client.is_some(),
)
.field("leases", &"ProviderLeasePort")
.field("runtime_policy", &"ProviderRuntimePolicyPort")
.finish()
@@ -685,11 +701,19 @@ impl CodexCredentialAdminService {
) -> Self {
Self {
refresher,
personal_access_token_client: None,
leases,
runtime_policy,
}
}
/// 在生产组装时复用 OAuth auth client,为 at- 导入启用上游身份验证。
#[must_use]
pub fn with_personal_access_token_client(mut self, client: Arc<OpenAiTokenClient>) -> Self {
self.personal_access_token_client = Some(client);
self
}
/// 官方 RT exchange;结果由 App 在同一 revision/audit 事务中提交。
pub async fn manual_refresh(
&self,
@@ -743,7 +767,10 @@ impl CodexCredentialAdminService {
};
let tokens = self
.refresher
.refresh(refresh_token.expose_secret())
.refresh_with_proxy(
refresh_token.expose_secret(),
current.account.outbound_proxy(),
)
.await
.inspect_err(|error| log_manual_refresh_failure(&account_id, error))
.map_err(map_refresh_failure)?;
@@ -793,6 +820,7 @@ impl CodexCredentialAdminService {
///
/// OAuth 导入先取得 access token(直接提供或 RT exchange),再按官方
/// `parse_chatgpt_jwt_claims` 从 ID token/access token 本地投影账号资料。
/// at- PAT 使用 whoami 取得身份,并丢弃不适用的 RT、ID token 和刷新计划。
pub async fn prepare_import_document(
&self,
payload: Value,
@@ -813,34 +841,56 @@ impl CodexCredentialAdminService {
let account_id = format!("acct_{}", uuid::Uuid::now_v7().simple());
let typed_account_id = ProviderAccountId::new(account_id.clone())
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
let (secret, access_token_expires_at) = self
let (mut secret, mut access_token_expires_at) = self
.resolve_import_tokens(
&typed_account_id,
candidate.authentication.access_token.clone(),
candidate.authentication.refresh_token.clone(),
candidate.authentication.id_token.clone(),
candidate.outbound_proxy.as_ref(),
)
.await?;
// 与官方 `parse_chatgpt_jwt_claims` 一致:先读 ID token,缺失字段
// 再由 access token 补齐。解析失败不阻塞 token 入库。
let id_metadata = secret
.id_token
.as_ref()
.and_then(|token| parse_chatgpt_jwt_claims(token.expose_secret()).ok())
.unwrap_or_default();
let access_metadata =
parse_chatgpt_jwt_claims(secret.access_token.expose_secret()).unwrap_or_default();
let metadata = CodexOAuthMetadata {
email: id_metadata.email.or(access_metadata.email),
chatgpt_plan_type: id_metadata
.chatgpt_plan_type
.or(access_metadata.chatgpt_plan_type),
chatgpt_user_id: id_metadata
.chatgpt_user_id
.or(access_metadata.chatgpt_user_id),
chatgpt_account_id: id_metadata
.chatgpt_account_id
.or(access_metadata.chatgpt_account_id),
let metadata = if secret
.access_token
.expose_secret()
.trim()
.starts_with("at-")
{
let token = secret.access_token.expose_secret().trim();
let metadata = self
.personal_access_token_client
.as_ref()
.ok_or(PersonalAccessTokenError::Unavailable)?
.personal_access_token_metadata(token, candidate.outbound_proxy.as_ref())
.await?;
secret = CodexOAuthSecret {
access_token: SecretString::from(token),
refresh_token: None,
id_token: None,
};
access_token_expires_at = None;
metadata
} else {
// 普通 OAuth 仍按 ID token、access token 的顺序投影,不调用 whoami。
let id_metadata = secret
.id_token
.as_ref()
.and_then(|token| parse_chatgpt_jwt_claims(token.expose_secret()).ok())
.unwrap_or_default();
let access_metadata = parse_chatgpt_jwt_claims(secret.access_token.expose_secret())
.unwrap_or_default();
CodexOAuthMetadata {
email: id_metadata.email.or(access_metadata.email),
chatgpt_plan_type: id_metadata
.chatgpt_plan_type
.or(access_metadata.chatgpt_plan_type),
chatgpt_user_id: id_metadata
.chatgpt_user_id
.or(access_metadata.chatgpt_user_id),
chatgpt_account_id: id_metadata
.chatgpt_account_id
.or(access_metadata.chatgpt_account_id),
}
};
let prepared =
CodexCredentialAdmin.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
@@ -859,7 +909,12 @@ impl CodexCredentialAdminService {
next_refresh_at: None,
enabled: true,
})?;
accounts.push(prepared);
accounts.push(NewProviderAccount {
account: prepared
.account
.with_outbound_proxy(candidate.outbound_proxy),
credential: prepared.credential,
});
}
Ok(PreparedCodexAccountImport { accounts })
}
@@ -870,6 +925,7 @@ impl CodexCredentialAdminService {
access_token: Option<String>,
refresh_token: Option<String>,
id_token: Option<String>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<(CodexOAuthSecret, Option<DateTime<Utc>>), CodexCredentialAdminError> {
let id_token = id_token.map(SecretString::from);
if let Some(access_token) = access_token {
@@ -889,7 +945,7 @@ impl CodexCredentialAdminService {
let refresh_token = refresh_token.ok_or(CodexCredentialAdminError::InvalidCredential)?;
let tokens = self
.refresher
.refresh(&refresh_token)
.refresh_with_proxy(&refresh_token, proxy)
.await
.map_err(map_refresh_failure)?;
let access_token = tokens
@@ -986,12 +1042,18 @@ fn log_manual_refresh_failure(account_id: &ProviderAccountId, error: &RefreshFai
fn parse_import_document(
payload: &Value,
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let payload = payload
.get("data")
.filter(|data| data.get("accounts").is_some())
.unwrap_or(payload);
let mut accounts = Vec::new();
for value in import_account_values(payload)? {
if !is_openai_oauth_candidate(value) {
continue;
}
accounts.push(parse_oauth_import_account(value)?);
let mut account = parse_oauth_import_account(value)?;
account.outbound_proxy = import_proxy(payload, value)?;
accounts.push(account);
}
Ok(accounts)
}
@@ -1004,9 +1066,89 @@ fn parse_oauth_import_account(
name: first_string(value, &["name", "label"]),
email: first_string(value, &["email"]).or_else(|| first_string(credentials, &["email"])),
authentication: parse_oauth_import_tokens(value)?,
outbound_proxy: None,
})
}
fn import_proxy(
payload: &Value,
account: &Value,
) -> Result<Option<gateway_core::account::OutboundProxy>, CodexCredentialAdminError> {
let invalid = || CodexCredentialAdminError::InvalidInput;
if let Some(value) = account
.get("outbound_proxy_url")
.or_else(|| account.get("outboundProxyUrl"))
.filter(|value| !value.is_null())
{
let text = value.as_str().ok_or_else(invalid)?;
if text.is_empty() {
return Ok(None);
}
return gateway_core::account::OutboundProxy::parse(text)
.map(Some)
.map_err(|_| invalid());
}
let Some(key) = account.get("proxy_key").filter(|value| !value.is_null()) else {
return Ok(None);
};
let key = key
.as_str()
.filter(|key| !key.is_empty())
.ok_or_else(invalid)?;
let proxies = payload
.get("proxies")
.and_then(Value::as_array)
.ok_or_else(invalid)?;
let mut matches = proxies
.iter()
.filter(|proxy| proxy.get("proxy_key").and_then(Value::as_str) == Some(key));
let proxy = matches.next().ok_or_else(invalid)?;
if matches.next().is_some()
|| proxy
.get("status")
.and_then(Value::as_str)
.is_some_and(|value| value != "active")
|| proxy
.get("expires_at")
.is_some_and(|value| !value.is_null())
|| proxy
.get("fallback_mode")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty() && value != "none")
{
return Err(invalid());
}
let text = |field| proxy.get(field).and_then(Value::as_str).ok_or_else(invalid);
let scheme = text("protocol")?;
// sub2api resolves SOCKS targets at the proxy.
let scheme = if scheme == "socks5" {
"socks5h"
} else {
scheme
};
let mut url = url::Url::parse(&format!("{scheme}://localhost")).map_err(|_| invalid())?;
let host = text("host")?;
if let Ok(address) = host.parse::<std::net::IpAddr>() {
url.set_ip_host(address).map_err(|_| invalid())?;
} else {
url.set_host(Some(host)).map_err(|_| invalid())?;
}
let port = proxy
.get("port")
.and_then(Value::as_u64)
.and_then(|port| u16::try_from(port).ok())
.filter(|port| *port != 0)
.ok_or_else(invalid)?;
url.set_port(Some(port)).map_err(|_| invalid())?;
url.set_username(proxy.get("username").and_then(Value::as_str).unwrap_or(""))
.map_err(|_| invalid())?;
url.set_password(proxy.get("password").and_then(Value::as_str))
.map_err(|_| invalid())?;
gateway_core::account::OutboundProxy::parse(url.as_str())
.map(Some)
.map_err(|_| invalid())
}
fn import_account_values(payload: &Value) -> Result<Vec<&Value>, CodexCredentialAdminError> {
if let Some(accounts) = payload.get("accounts") {
return accounts
@@ -1032,7 +1174,10 @@ fn parse_oauth_import_tokens(
Value::Object(object) => {
for (key, value) in object {
let token = match key.as_str() {
"accessToken" | "access_token" => &mut access_token,
"accessToken"
| "access_token"
| "personal_access_token"
| "personalAccessToken" => &mut access_token,
"refreshToken" | "refresh_token" => &mut refresh_token,
"idToken" | "id_token" => &mut id_token,
_ => {
@@ -436,6 +436,10 @@ impl CodexCredentialCatalogService {
.authorization_header()
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
let result = client
.for_account(account)
.map_err(|error| CodexCredentialCatalogError::Upstream {
detail: error.to_string(),
})?
.fetch_models_with_context(CodexRequestContext::auxiliary(
authorization.expose_secret(),
account.upstream_account_id(),
@@ -524,7 +524,16 @@ impl CodexOAuthAdminService {
.transpose()?
.flatten();
let (mut secret, id_token) = self
.exchange_pending(&pending, callback_url, fallback_refresh_token)
.exchange_pending(
&pending,
callback_url,
fallback_refresh_token,
current
.as_ref()
.map_or(pending.mutation.outbound_proxy(), |current| {
current.account.outbound_proxy()
}),
)
.await?;
let mutation = pending.mutation.clone();
let access_token_expires_at =
@@ -545,20 +554,25 @@ impl CodexOAuthAdminService {
)
} else {
let account_id = format!("acct_{}", Uuid::now_v7().simple());
CompletedCodexOAuthCredential::Create(
self.credentials
.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
account_id,
name: pending.name,
installation_id: pending.installation_id,
secret,
metadata,
access_token_expires_at,
next_refresh_at: None,
enabled: true,
})
.map_err(map_admin_error)?,
)
let prepared = self
.credentials
.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
account_id,
name: pending.name,
installation_id: pending.installation_id,
secret,
metadata,
access_token_expires_at,
next_refresh_at: None,
enabled: true,
})
.map_err(map_admin_error)?;
CompletedCodexOAuthCredential::Create(gateway_core::account::NewProviderAccount {
account: prepared
.account
.with_outbound_proxy(mutation.outbound_proxy().cloned()),
credential: prepared.credential,
})
};
Ok((mutation, credential))
}
@@ -666,6 +680,7 @@ impl CodexOAuthAdminService {
pending: &CodexPendingAuthorization,
callback_url: &str,
fallback_refresh_token: Option<SecretString>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
let (code, callback_state) = callback_parts(callback_url)?;
if !constant_time_equal(
@@ -676,10 +691,13 @@ impl CodexOAuthAdminService {
}
let tokens = self
.exchanger
.exchange_authorization_code(AuthorizationCodeGrant {
code,
code_verifier: pending.code_verifier.clone(),
})
.exchange_with_proxy(
AuthorizationCodeGrant {
code,
code_verifier: pending.code_verifier.clone(),
},
proxy,
)
.await
.map_err(map_exchange_error)?;
let mut secret = tokens.secret;
@@ -133,13 +133,16 @@ impl CodexCredentialProfileService {
&self,
account_id: &ProviderAccountId,
) -> Result<CodexProfileStatistics, CodexProfileStatisticsError> {
let (authorization, upstream_account_id) = self.account_authentication(account_id).await?;
let (authorization, upstream_account_id, account) =
self.account_authentication(account_id).await?;
let request_id = format!("profile_statistics_{}", Uuid::now_v7().simple());
let statistics = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
)
.for_account(&account)
.map_err(map_client_error)?
.fetch_profile_statistics(CodexRequestContext::auxiliary(
authorization.expose_secret(),
upstream_account_id.as_deref(),
@@ -156,7 +159,14 @@ impl CodexCredentialProfileService {
async fn account_authentication(
&self,
account_id: &ProviderAccountId,
) -> Result<(SecretString, Option<String>), CodexProfileStatisticsError> {
) -> Result<
(
SecretString,
Option<String>,
gateway_core::account::ProviderAccount,
),
CodexProfileStatisticsError,
> {
let account = self
.repository
.store()
@@ -187,6 +197,7 @@ impl CodexCredentialProfileService {
Ok((
authorization,
account.upstream_account_id().map(str::to_owned),
account,
))
}
@@ -202,9 +213,19 @@ impl CodexCredentialProfileService {
.image_url
.ok_or(CodexProfileAvatarError::Missing)?,
};
let (authorization, upstream_account_id) = self.account_authentication(account_id).await?;
let (authorization, upstream_account_id, account) =
self.account_authentication(account_id).await?;
let http = if account.outbound_proxy().is_some() {
crate::transport::client::build_account_http_client(
account.id().as_str(),
account.outbound_proxy(),
)
.map_err(|_| CodexProfileAvatarError::TransportUnavailable)?
} else {
self.http.clone()
};
fetch_profile_avatar(
&self.http,
&http,
&self.base_url,
&self.profile.snapshot(),
&source,
@@ -1225,6 +1225,8 @@ async fn list_reset_credits_once(
.authorization_header()
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
client
.for_account(&prepared.account)
.map_err(ResetCreditAttemptError::Upstream)?
.list_rate_limit_reset_credits(CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
@@ -1248,6 +1250,8 @@ async fn consume_reset_credit_once(
.authorization_header()
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
client
.for_account(&prepared.account)
.map_err(ResetCreditAttemptError::Upstream)?
.consume_rate_limit_reset_credit(
CodexRequestContext::auxiliary(
authorization.expose_secret(),
@@ -1329,6 +1333,8 @@ async fn fetch_usage_once(
.map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?;
let request_id = format!("quota_{}", Uuid::now_v7().simple());
client
.for_account(&prepared.account)
.map_err(CodexQuotaFetchAttemptError::Upstream)?
.fetch_usage(CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
@@ -113,6 +113,7 @@ pub struct CodexAccountQuotaSnapshot {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
observed_at: SystemTime,
plan_type: Option<String>,
fact: CodexQuotaFact,
quota: QuotaState,
windows: Vec<CodexQuotaWindow>,
@@ -135,6 +136,12 @@ impl CodexAccountQuotaSnapshot {
self.observed_at
}
/// 上游额度响应明确返回的套餐;缺失时不推断账号级别。
#[must_use]
pub fn plan_type(&self) -> Option<&str> {
self.plan_type.as_deref()
}
#[must_use]
pub const fn fact(&self) -> CodexQuotaFact {
self.fact
@@ -339,6 +346,10 @@ pub(crate) fn parse_account_quota_snapshot(
account_id,
credential_revision,
observed_at,
plan_type: object
.get("plan_type")
.and_then(Value::as_str)
.map(str::to_owned),
fact,
quota,
windows,
@@ -288,7 +288,10 @@ impl CodexCredentialRefreshService {
.refresh_token
.as_ref()
.ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?;
let refresh_result = self.refresher.refresh(refresh_token.expose_secret()).await;
let refresh_result = self
.refresher
.refresh_with_proxy(refresh_token.expose_secret(), due.account.outbound_proxy())
.await;
if recovery_window_exhausted && let Err(failure) = &refresh_result {
let message = failure.message().map(str::to_owned);
let upstream = failure.upstream();
@@ -1,4 +1,6 @@
//! OpenAI token 续期 Reqwest 适配器。
//! OpenAI OAuth token exchange 与 Codex PAT 验证的 Reqwest 适配器。
use super::types::CodexOAuthMetadata;
use crate::transport::{
headers::build_codex_profile_headers,
@@ -22,6 +24,33 @@ pub const OFFICIAL_CODEX_OAUTH_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
pub const OFFICIAL_CODEX_TOKEN_ENDPOINT: &str = "https://auth.openai.com/oauth/token";
/// Codex Desktop loopback callback;管理员复制完整回调 URL 交回固定 complete API。
pub const OFFICIAL_CODEX_REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
const PERSONAL_ACCESS_TOKEN_WHOAMI_PATH: &str = "/api/accounts/v1/user-auth-credential/whoami";
/// PAT 验证失败;不保留令牌、响应体或可能包含秘密的底层 HTTP 错误。
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum PersonalAccessTokenError {
#[error("Codex PAT must be a non-empty at- token without whitespace or control characters")]
InvalidToken,
#[error(
"Codex PAT was rejected by OpenAI; it may be invalid, expired, revoked, or lack permission"
)]
Rejected,
#[error("Codex PAT validation is unavailable; retry later")]
Unavailable,
#[error("OpenAI returned an invalid Codex PAT identity response")]
InvalidResponse,
}
// 与 Codex personal_access_token.rs 一致:email 可缺失,其余身份字段必填。
#[derive(Deserialize)]
struct PersonalAccessTokenResponse {
email: Option<String>,
chatgpt_user_id: String,
chatgpt_account_id: String,
chatgpt_plan_type: String,
#[serde(rename = "chatgpt_account_is_fedramp")]
_chatgpt_account_is_fedramp: bool,
}
/// Token 刷新成功后得到的认证材料。
#[derive(Clone)]
@@ -178,6 +207,22 @@ impl fmt::Debug for RefreshUpstreamFailure {
#[async_trait]
pub trait TokenRefresher: Send + Sync + 'static {
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
async fn refresh_with_proxy(
&self,
refresh_token: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<TokenPair, RefreshFailure> {
if proxy.is_some() {
return Err(proxy_refresh_failure());
}
self.refresh(refresh_token).await
}
}
fn proxy_refresh_failure() -> RefreshFailure {
RefreshFailure::RetryableTransport {
message: "account OAuth egress unavailable".to_owned(),
}
}
/// Authorization Code + PKCE 的一次性 grant。
@@ -232,6 +277,16 @@ pub trait AuthorizationCodeExchanger: Send + Sync + 'static {
&self,
grant: AuthorizationCodeGrant,
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError>;
async fn exchange_with_proxy(
&self,
grant: AuthorizationCodeGrant,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
if proxy.is_some() {
return Err(AuthorizationCodeExchangeError::Unavailable);
}
self.exchange_authorization_code(grant).await
}
}
/// OpenAI token 续期客户端配置。
@@ -257,6 +312,24 @@ pub struct OpenAiTokenClient {
pub struct TokenClientBuildError;
impl OpenAiTokenClient {
fn with_proxy(
&self,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Self, TokenClientBuildError> {
let Some(proxy) = proxy else {
return Ok(self.clone());
};
let builder = Client::builder()
.use_rustls_tls()
.no_proxy()
.proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| TokenClientBuildError)?)
.redirect(Policy::none())
.connect_timeout(TOKEN_CONNECT_TIMEOUT)
.timeout(TOKEN_REQUEST_TIMEOUT);
let client =
build_reqwest_client_with_custom_ca(builder).map_err(|_| TokenClientBuildError)?;
Ok(Self::new(client, self.config.clone(), self.profile.clone()))
}
/// 共享运行时画像;刷新时取快照,授权码交换仍使用 raw auth 请求。
pub fn new(client: Client, config: TokenClientConfig, profile: CodexWireProfileState) -> Self {
Self {
@@ -265,6 +338,73 @@ impl OpenAiTokenClient {
profile,
}
}
pub(crate) async fn personal_access_token_metadata(
&self,
access_token: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<CodexOAuthMetadata, PersonalAccessTokenError> {
if !access_token.starts_with("at-")
|| access_token.len() <= 3
|| access_token.len() > MAX_OAUTH_RESPONSE_BYTES
|| access_token
.chars()
.any(|ch| ch.is_whitespace() || ch.is_control())
{
return Err(PersonalAccessTokenError::InvalidToken);
}
// 复用固定 auth origin 及其 TLS/超时/禁止重定向策略;导入文档不能指定验证地址。
let mut endpoint = reqwest::Url::parse(&self.config.token_endpoint)
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
endpoint.set_path(PERSONAL_ACCESS_TOKEN_WHOAMI_PATH);
endpoint.set_query(None);
endpoint.set_fragment(None);
let headers = build_codex_profile_headers(&self.profile.snapshot())
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
let client = self
.with_proxy(proxy)
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
let response = client
.client
.get(endpoint)
.headers(headers)
.bearer_auth(access_token)
.header(reqwest::header::ACCEPT, "application/json")
.send()
.await
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
match response.status() {
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
return Err(PersonalAccessTokenError::Rejected);
}
status if !status.is_success() => return Err(PersonalAccessTokenError::Unavailable),
_ => {}
}
let (_, body) = read_bounded_response(response)
.await
.map_err(|_| PersonalAccessTokenError::InvalidResponse)?;
let identity: PersonalAccessTokenResponse =
serde_json::from_slice(&body).map_err(|_| PersonalAccessTokenError::InvalidResponse)?;
if [
&identity.chatgpt_user_id,
&identity.chatgpt_account_id,
&identity.chatgpt_plan_type,
]
.iter()
.any(|value| value.trim().is_empty() || value.chars().any(char::is_control))
{
return Err(PersonalAccessTokenError::InvalidResponse);
}
Ok(CodexOAuthMetadata {
email: identity
.email
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty()),
chatgpt_user_id: Some(identity.chatgpt_user_id.trim().to_owned()),
chatgpt_account_id: Some(identity.chatgpt_account_id.trim().to_owned()),
chatgpt_plan_type: Some(identity.chatgpt_plan_type.trim().to_owned()),
})
}
}
/// 构建禁止 redirect 且无自动重试的 Codex token client。
@@ -357,6 +497,16 @@ struct AuthorizationCodeResponse {
#[async_trait]
impl TokenRefresher for OpenAiTokenClient {
async fn refresh_with_proxy(
&self,
refresh_token: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<TokenPair, RefreshFailure> {
self.with_proxy(proxy)
.map_err(|_| proxy_refresh_failure())?
.refresh(refresh_token)
.await
}
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
let headers = build_codex_profile_headers(&self.profile.snapshot()).map_err(|_| {
RefreshFailure::Transport {
@@ -389,6 +539,16 @@ impl TokenRefresher for OpenAiTokenClient {
#[async_trait]
impl AuthorizationCodeExchanger for OpenAiTokenClient {
async fn exchange_with_proxy(
&self,
grant: AuthorizationCodeGrant,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
self.with_proxy(proxy)
.map_err(|_| AuthorizationCodeExchangeError::Unavailable)?
.exchange_authorization_code(grant)
.await
}
async fn exchange_authorization_code(
&self,
grant: AuthorizationCodeGrant,
+9 -6
View File
@@ -153,12 +153,15 @@ pub async fn initialize(
.map_err(|_| OpenAiInitializeError::TokenClient)?,
);
let refresher: Arc<dyn TokenRefresher> = token_client.clone();
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client;
let credential_admin = Arc::new(CodexCredentialAdminService::new(
Arc::clone(&refresher),
Arc::clone(&leases),
Arc::clone(&runtime_policy),
));
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client.clone();
let credential_admin = Arc::new(
CodexCredentialAdminService::new(
Arc::clone(&refresher),
Arc::clone(&leases),
Arc::clone(&runtime_policy),
)
.with_personal_access_token_client(token_client),
);
let refresh = Arc::new(CodexCredentialRefreshService::new(
repository,
refresher,
@@ -124,7 +124,9 @@ impl CodexProvider {
crate::transport::request::scope_turn_metadata(metadata, lease.installation_id(), true)
});
let events = cold_json_response_stream(ColdJsonResponse {
client: self.client.clone(),
client: self.client.for_account(lease.account()).map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?,
response_origin: request.response_origin,
endpoint_path: request.endpoint_path,
body: request.body,
@@ -519,7 +519,9 @@ impl Provider for CodexProvider {
AttemptTransport::Default | AttemptTransport::Fallback => 0,
};
let events = cold_response_stream(ColdResponse {
client: self.client.clone(),
client: self.client.for_account(lease.account()).map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?,
response_origin: self.responses_url.clone(),
request: upstream_request,
upstream_model: upstream_model.clone(),
@@ -45,13 +45,20 @@ const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
pub(super) const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
"x-codex-ws-stream-request-start-ms";
type ReqwestClientCacheKey = Option<String>;
type ReqwestClientCacheKey = (Option<String>, String);
type ReqwestClientCache = Mutex<HashMap<ReqwestClientCacheKey, Client>>;
/// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。
pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
build_account_http_client("", None)
}
pub fn build_account_http_client(
account_id: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, CustomCaError> {
super::tls::ensure_rustls_provider();
let cache_key = custom_ca_env_cache_key();
let cache_key = (custom_ca_env_cache_key(), egress_key(account_id, proxy));
static CLIENTS: OnceLock<ReqwestClientCache> = OnceLock::new();
let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new()));
if let Some(client) = cache
@@ -62,7 +69,7 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
return Ok(client.clone());
}
let builder = Client::builder()
let mut builder = Client::builder()
.use_rustls_tls()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
@@ -73,13 +80,35 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
.http2_keep_alive_interval(Duration::from_secs(30))
.http2_keep_alive_timeout(Duration::from_secs(5))
.http2_keep_alive_while_idle(true);
if let Some(proxy) = proxy {
builder = builder.proxy(
reqwest::Proxy::all(proxy.expose_url())
.map_err(|_| CustomCaError::ProxyConfiguration)?,
);
}
let client = build_reqwest_client_with_custom_ca(builder)?;
let mut clients = cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if clients.len() >= 256 {
clients.clear();
}
Ok(clients.entry(cache_key).or_insert(client).clone())
}
fn egress_key(account_id: &str, proxy: Option<&gateway_core::account::OutboundProxy>) -> String {
use sha2::{Digest, Sha256};
let mut hash = Sha256::new();
hash.update(account_id.as_bytes());
hash.update([0]);
hash.update(
proxy
.map_or("direct", |proxy| proxy.expose_url())
.as_bytes(),
);
hex::encode(hash.finalize())
}
// ---------------------------------------------------------------------------
// 错误类型
// ---------------------------------------------------------------------------
@@ -599,11 +628,36 @@ pub struct CodexBackendJsonResponse {
#[derive(Clone)]
pub struct CodexBackendClient {
pub(super) client: Client,
pub(super) direct_client: Client,
pub(super) base_url: String,
pub(super) profile: CodexWireProfileState,
pub(super) websocket_pool: Option<Arc<CodexWebSocketPool>>,
pub(super) websocket_origin_breaker: WebSocketOriginBreaker,
pub(super) websocket_origin_key: String,
pub(super) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub(super) egress_key: String,
}
impl CodexBackendClient {
pub fn for_account(
&self,
account: &gateway_core::account::ProviderAccount,
) -> Result<Self, CodexClientError> {
let mut client = self.clone();
client.outbound_proxy = account.outbound_proxy().cloned();
client.egress_key = egress_key(account.id().as_str(), account.outbound_proxy());
client.websocket_origin_key = format!(
"{}:{}",
websocket_origin_key(&self.base_url),
client.egress_key
);
client.client = if account.outbound_proxy().is_some() {
build_account_http_client(account.id().as_str(), account.outbound_proxy())?
} else {
self.direct_client.clone()
};
Ok(client)
}
}
/// 已完成账号级 opening 准备、但尚未发送 payload 的 transport。
@@ -52,8 +52,11 @@ impl CodexBackendClient {
) -> Self {
let base_url = base_url.into().trim_end_matches('/').to_string();
Self {
direct_client: client.clone(),
client,
websocket_origin_key: websocket_origin_key(&base_url),
outbound_proxy: None,
egress_key: String::new(),
base_url,
profile,
websocket_pool: None,
@@ -234,13 +237,14 @@ impl CodexBackendClient {
let websocket_request = websocket_upstream_request(request);
let headers = self.request_headers_for_websocket_response(&websocket_request, context)?;
let websocket_create = CodexWebSocketConnection::responses_create_request(
let mut websocket_create = CodexWebSocketConnection::responses_create_request(
&self.base_url,
&generate_key(),
websocket_header_pairs(&headers),
&websocket_request,
)
.map_err(CodexClientError::WebSocketEncode)?;
websocket_create.connection.outbound_proxy = self.outbound_proxy.clone();
context.trace.cloned().unwrap_or_default().headers(
"upstream.request.headers",
serde_json::json!({"transport": "websocket", "phase": "prepared_opening"}),
@@ -469,6 +473,7 @@ impl CodexBackendClient {
.as_deref()
.or(request.previous_response_id())?;
let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id)
.with_egress_key(&self.egress_key)
.with_connection_profile(connection_profile);
if let Some(connection_id) = request.downstream_websocket_connection_id.as_deref() {
key = key.with_downstream_connection_id(connection_id);
@@ -33,6 +33,8 @@ pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE";
/// 自定义 CA 错误。
#[derive(Debug, Error)]
pub enum CustomCaError {
#[error("invalid account proxy configuration")]
ProxyConfiguration,
/// 读取 CA 证书文件失败。
#[error(
"Failed to read CA certificate file {} selected by {}: {source}. {hint}",
@@ -121,6 +123,17 @@ pub fn maybe_build_rustls_client_config_with_custom_ca() -> CustomCaResult<Optio
maybe_build_rustls_client_config_with_env(&ProcessEnv)
}
pub(crate) fn account_proxy_tls_config() -> CustomCaResult<Arc<ClientConfig>> {
if let Some(config) = maybe_build_rustls_client_config_with_custom_ca()? {
return Ok(config);
}
Ok(Arc::new(
ClientConfig::builder()
.with_root_certificates(native_root_store().map_err(CustomCaError::LoadNativeRoots)?)
.with_no_client_auth(),
))
}
fn build_reqwest_client_with_env(
env_source: &dyn EnvSource,
mut builder: reqwest::ClientBuilder,
@@ -5,7 +5,9 @@ use std::time::Duration;
use bytes::Bytes;
use gateway_protocol::openai::events;
use tokio::time::timeout;
use tokio_tungstenite::{Connector, connect_async_tls_with_config};
use tokio_tungstenite::{
Connector, MaybeTlsStream, client_async_tls_with_config, connect_async_tls_with_config,
};
use tungstenite::{
self, Message,
extensions::{ExtensionsConfig, compression::deflate::DeflateConfig},
@@ -60,7 +62,11 @@ impl CodexWebSocketConnection {
"sec-websocket-extensions".to_string(),
WEBSOCKET_EXTENSIONS.to_string(),
));
Self { endpoint, headers }
Self {
endpoint,
headers,
outbound_proxy: None,
}
}
/// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。
@@ -140,10 +146,39 @@ async fn connect_websocket(
)))
})?
.map(Connector::Rustls);
let result = timeout(
WEBSOCKET_CONNECT_TIMEOUT,
connect_async_tls_with_config(request, Some(websocket_config()), false, connector),
)
let result = timeout(WEBSOCKET_CONNECT_TIMEOUT, async {
// Preserve the native direct handshake; explicit egress never inherits a global proxy.
if connection.outbound_proxy.is_none()
&& matches!(
tungstenite::proxy::ProxyConfig::from_env(request.uri()),
Ok(None)
)
{
let (websocket, response) =
connect_async_tls_with_config(request, Some(websocket_config()), false, connector)
.await?;
return Ok((Box::new(websocket) as RawWsStream, response));
}
let stream = dial_account(connection).await?;
match stream {
MaybeTlsStream::Plain(tcp) => {
let (websocket, response) =
client_async_tls_with_config(request, tcp, Some(websocket_config()), connector)
.await?;
Ok((Box::new(websocket) as RawWsStream, response))
}
stream => {
let (websocket, response) = client_async_tls_with_config(
request,
stream,
Some(websocket_config()),
connector,
)
.await?;
Ok((Box::new(websocket) as RawWsStream, response))
}
}
})
.await
.map_err(|_| CodexWebSocketExchangeError::ConnectTimeout {
timeout: WEBSOCKET_CONNECT_TIMEOUT,
@@ -155,6 +190,89 @@ async fn connect_websocket(
}
}
async fn dial_account(
connection: &CodexWebSocketConnection,
) -> Result<MaybeTlsStream<tokio::net::TcpStream>, tungstenite::Error> {
use tokio::net::lookup_host;
use tungstenite::proxy::ProxyConfig;
let invalid =
|| tungstenite::Error::Io(std::io::Error::other("account WebSocket egress failed"));
let endpoint = url::Url::parse(connection.endpoint()).map_err(|_| invalid())?;
let host = endpoint.host_str().ok_or_else(invalid)?;
let dns_host = match endpoint.host() {
Some(url::Host::Ipv6(ip)) => ip.to_string(),
_ => host.to_owned(),
};
let port = endpoint.port_or_known_default().ok_or_else(invalid)?;
let Some(proxy) = connection.outbound_proxy.as_ref() else {
return Ok(MaybeTlsStream::Plain(connect_tcp(host, port).await?));
};
let mut proxy_url = url::Url::parse(proxy.expose_url()).map_err(|_| invalid())?;
let tls_proxy = proxy_url.scheme() == "https";
let proxy_port = proxy_url.port_or_known_default().ok_or_else(invalid)?;
if tls_proxy {
proxy_url.set_scheme("http").map_err(|_| invalid())?;
proxy_url
.set_port(Some(proxy_port))
.map_err(|_| invalid())?;
}
let config = ProxyConfig::parse(proxy_url.as_str()).map_err(|_| invalid())?;
let proxy_host = match proxy_url.host() {
Some(url::Host::Ipv6(ip)) => ip.to_string(),
_ => config.host.clone(),
};
let tcp = connect_tcp(&proxy_host, config.port).await?;
let stream = if tls_proxy {
let tls = tls::account_proxy_tls_config().map_err(|_| invalid())?;
let name = rustls_pki_types::ServerName::try_from(proxy_host).map_err(|_| invalid())?;
MaybeTlsStream::Rustls(
tokio_rustls::TlsConnector::from(tls)
.connect(name, tcp)
.await?,
)
} else {
MaybeTlsStream::Plain(tcp)
};
// The pinned tungstenite fork sends domains remotely for both SOCKS schemes.
let target = if proxy_url.scheme() == "socks5" {
lookup_host((dns_host.as_str(), port))
.await?
.next()
.ok_or_else(invalid)?
.ip()
.to_string()
} else if proxy_url.scheme() == "socks5h" {
dns_host
} else {
host.to_owned()
};
tokio_tungstenite::proxy::connect_via_proxy(stream, &config, &target, port)
.await
.map_err(|_| invalid())
}
async fn connect_tcp(host: &str, port: u16) -> Result<tokio::net::TcpStream, tungstenite::Error> {
use hyper_util::client::legacy::connect::HttpConnector;
use tower_service::Service;
// Reuse Hyper's DNS and Happy Eyeballs connector for explicit direct/proxy endpoints.
let authority = host.parse::<std::net::IpAddr>().map_or_else(
|_| format!("{host}:{port}"),
|ip| std::net::SocketAddr::new(ip, port).to_string(),
);
let uri = format!("http://{authority}")
.parse::<hyper::Uri>()
.map_err(|_| tungstenite::Error::Io(std::io::Error::other("invalid egress endpoint")))?;
let mut connector = HttpConnector::new();
connector
.call(uri)
.await
.map(|stream| stream.into_inner())
.map_err(|_| {
tungstenite::Error::Io(std::io::Error::other("account WebSocket connection failed"))
})
}
fn websocket_handshake_request(
connection: &CodexWebSocketConnection,
) -> Result<WsRequest, tungstenite::http::Error> {
@@ -9,6 +9,7 @@ use crate::transport::protocol::responses::{
/// WebSocket opening 描述。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexWebSocketConnection {
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub(super) endpoint: String,
pub(super) headers: Vec<(String, String)>,
}
@@ -16,7 +17,7 @@ pub struct CodexWebSocketConnection {
/// 已构造完成的 Responses WebSocket 请求描述。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexWebSocketRequest {
pub(super) connection: CodexWebSocketConnection,
pub(crate) connection: CodexWebSocketConnection,
pub(super) payload_text: String,
pub(super) continuation: WebSocketContinuationRequirement,
}
@@ -113,6 +114,7 @@ impl CodexWebSocketConnection {
Self {
endpoint: endpoint.into(),
headers,
outbound_proxy: None,
}
}
@@ -24,6 +24,7 @@ pub struct CodexWebSocketPoolKey {
conversation_id: String,
connection_profile: String,
downstream_connection_id: String,
egress_key: String,
}
impl CodexWebSocketPoolKey {
@@ -39,6 +40,7 @@ impl CodexWebSocketPoolKey {
conversation_id: conversation_id.into(),
connection_profile: String::new(),
downstream_connection_id: String::new(),
egress_key: String::new(),
}
}
@@ -61,6 +63,11 @@ impl CodexWebSocketPoolKey {
&self.account_id
}
pub(crate) fn with_egress_key(mut self, key: &str) -> Self {
self.egress_key = key.to_owned();
self
}
pub(crate) fn conversation_id_hash(&self) -> String {
short_sha256([self.conversation_id.as_str()])
}
@@ -72,6 +79,7 @@ impl CodexWebSocketPoolKey {
self.conversation_id.as_str(),
self.connection_profile.as_str(),
self.downstream_connection_id.as_str(),
self.egress_key.as_str(),
])
}
@@ -79,6 +87,7 @@ impl CodexWebSocketPoolKey {
self.base_url == other.base_url
&& self.account_id == other.account_id
&& self.conversation_id == other.conversation_id
&& self.egress_key == other.egress_key
}
}
@@ -19,20 +19,35 @@ use std::sync::{
};
use std::time::Duration;
use futures::{SinkExt, StreamExt};
use futures::{Sink, SinkExt, Stream, StreamExt};
use tokio::{
net::TcpStream,
sync::{mpsc, oneshot},
task::JoinHandle,
time::{Instant, MissedTickBehavior},
};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message};
use tokio_tungstenite::tungstenite::Message;
use uuid::Uuid;
use super::CodexWebSocketCloseError;
/// 底层 tungstenite WebSocket 流。
pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
pub(crate) trait WebSocketIo:
Stream<Item = Result<Message, tungstenite::Error>>
+ Sink<Message, Error = tungstenite::Error>
+ Send
+ Unpin
{
}
impl<T> WebSocketIo for T where
T: Stream<Item = Result<Message, tungstenite::Error>>
+ Sink<Message, Error = tungstenite::Error>
+ Send
+ Unpin
{
}
pub(crate) type RawWsStream = Box<dyn WebSocketIo>;
const PUMP_COMMAND_BUFFER: usize = 32;
@@ -655,9 +670,10 @@ pub(crate) fn transport_metric_reason(error: &tungstenite::Error) -> &'static st
std::io::ErrorKind::TimedOut => "transport_timeout",
_ => "transport_error",
},
// 未收到关闭握手只能证明连接异常结束,不能据此认定收到 TCP RST。
tungstenite::Error::Protocol(
tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
) => "tcp_reset",
) => "reset_without_closing_handshake",
_ => "transport_error",
}
}
@@ -443,6 +443,7 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
let error = bundle
.admin_provider()
.complete_authorization(CompleteAuthorization {
settings: None,
context: MutationContext {
actor: MutationActor::AdminSession {
admin_user_id: "different-owner".to_owned(),
@@ -647,6 +648,71 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export
assert!(exported_account.get("token").is_none());
}
#[tokio::test]
async fn openai_admin_projects_free_plan_from_cached_quota_when_account_claims_omit_it() {
let store = Arc::new(MemoryAccountStore::default());
let mut verified_account = profile("chatgpt-free-plan");
verified_account.plan_type = None;
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_free_plan".to_owned(),
name: "free plan".to_owned(),
secret: secret("free-plan-test-token"),
verified_account,
next_refresh_at: None,
enabled: true,
})
.await;
let account = store.account("acct_free_plan").expect("stored account");
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
json!({
"plan_type": "free",
"rate_limit": {
"allowed": true,
"primary_window": {
"used_percent": 0,
"limit_window_seconds": 2_592_000,
"reset_at": 1_900_000_000
}
}
})
.as_object()
.expect("quota object")
.clone(),
),
observed_at,
state: QuotaState::allowed(observed_at),
})
.await
.expect("persist existing quota");
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(store, Arc::new(TestOAuthPending::default())),
)
.await
.expect("OpenAI bundle");
let admin = bundle.admin_provider();
let quota = admin
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("read cached free quota");
assert_eq!(quota.plan_type.as_deref(), Some("free"));
assert_eq!(
admin.plan_type_display(quota.plan_type.as_deref().expect("plan")),
"Free"
);
}
#[tokio::test]
async fn openai_admin_provider_projects_official_codex_quota_and_independent_buckets_with_chinese_labels()
{
@@ -950,6 +1016,31 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
assert_eq!(refresh_error.kind(), ProviderAdminErrorKind::NotFound);
}
#[tokio::test]
async fn initialized_provider_reports_a_safe_pat_format_error_before_network_access() {
let config = valid_config();
let bundle = provider_openai::initialize(config.config.clone(), provider_ports())
.await
.expect("OpenAI bundle");
let error = bundle
.admin_provider()
.prepare_import(PrepareCredentialImport {
document: ProviderDocument::new(OpaqueProviderData::new(Map::from_iter([(
"accessToken".to_owned(),
json!("at-sensitive-token with whitespace"),
)]))),
})
.await
.expect_err("PAT format must be checked by the initialized provider");
assert_eq!(error.kind(), ProviderAdminErrorKind::Invalid);
assert_eq!(
error.public_message(),
Some("Codex PAT 格式无效:应为 at- 开头的完整令牌,不能包含空白或控制字符")
);
assert!(error.message().is_none());
assert!(!format!("{error:?}").contains("sensitive-token"));
}
#[tokio::test]
async fn openai_rotation_preserves_the_new_access_token_jwt_expiration() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1137,6 +1228,7 @@ fn provider_ports_with_catalog(
fn account_record(account: &ProviderAccount) -> AccountRecord {
let now = Utc::now();
AccountRecord {
outbound_proxy: None,
id: account.id().to_string(),
provider_kind: account.provider().clone(),
groups: Vec::new(),
@@ -1,16 +1,94 @@
use std::sync::Arc;
use std::{
num::NonZeroU32,
sync::Arc,
time::{Duration, SystemTime},
};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{TimeZone as _, Utc};
use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher};
use provider_openai::credential::token_client::{
OpenAiTokenClient, PersonalAccessTokenError, RefreshFailure, TokenClientConfig, TokenPair,
TokenRefresher, openai_token_client,
};
use provider_openai::credential::{CodexCredentialAdminService, CodexCredentialCodec};
use secrecy::ExposeSecret as _;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{header, method, path},
};
use crate::support::{TestLeaseCoordinator, runtime_policy};
struct UnusedRefresher;
#[tokio::test]
async fn sub2api_import_resolves_distinct_proxy_bindings_and_encodes_credentials() {
let service = CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
);
let prepared = service.prepare_import_document(serde_json::json!({"data": {
"accounts": [
{"name": "a", "platform": "openai", "type": "oauth", "proxy_key": "a", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-a"}}))}},
{"name": "b", "platform": "openai", "type": "oauth", "proxy_key": "b", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-b"}}))}},
{"name": "direct", "platform": "openai", "type": "oauth", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-c"}}))}}
],
"proxies": [
{"proxy_key": "a", "protocol": "http", "host": "127.0.0.1", "port": 18080, "username": "user@a", "password": "p:a/ss", "status": "active"},
{"proxy_key": "b", "protocol": "socks5", "host": "::1", "port": 1080, "status": "active", "fallback_mode": "none"}
]
}})).await.unwrap();
assert_eq!(prepared.accounts().len(), 3);
let first = prepared.accounts()[0].account.outbound_proxy().unwrap();
assert_eq!(first.endpoint(), "http://127.0.0.1:18080/");
assert!(first.expose_url().contains("user%40a:p%3Aa%2Fss@"));
assert_eq!(
prepared.accounts()[1]
.account
.outbound_proxy()
.unwrap()
.endpoint(),
"socks5h://[::1]:1080"
);
assert!(prepared.accounts()[2].account.outbound_proxy().is_none());
assert!(!format!("{prepared:?}").contains("p:a/ss"));
}
#[tokio::test]
async fn sub2api_invalid_proxy_bindings_are_rejected_before_any_token_refresh() {
let service = CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
);
let proxy = serde_json::json!({"proxy_key": "bound", "protocol": "http", "host": "127.0.0.1", "port": 18080, "status": "active"});
let mut cases = vec![
serde_json::json!([]),
serde_json::json!([proxy.clone(), proxy.clone()]),
];
for (field, value) in [
("status", serde_json::json!("inactive")),
("expires_at", serde_json::json!(2000000000)),
("fallback_mode", serde_json::json!("direct")),
("port", serde_json::json!(0)),
] {
let mut invalid = proxy.clone();
invalid[field] = value;
cases.push(serde_json::json!([invalid]));
}
for proxies in cases {
let result = service.prepare_import_document(serde_json::json!({
"accounts": [
{"platform": "openai", "type": "oauth", "credentials": {"refresh_token": "must-not-refresh"}},
{"platform": "openai", "type": "oauth", "proxy_key": "bound", "credentials": {"access_token": "at"}}
], "proxies": proxies
})).await;
assert!(result.is_err());
}
}
#[async_trait]
impl TokenRefresher for UnusedRefresher {
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
@@ -145,6 +223,390 @@ fn test_jwt(payload: serde_json::Value) -> String {
)
}
fn pat_service(server: &MockServer) -> CodexCredentialAdminService {
let client = openai_token_client(
TokenClientConfig {
client_id: "test-public-client".to_owned(),
token_endpoint: format!("{}/oauth/token", server.uri()),
},
provider_openai::OpenAiConfig::default().wire_profile_state(),
)
.expect("auth client");
CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
)
.with_personal_access_token_client(Arc::new(client))
}
fn pat_identity() -> serde_json::Value {
serde_json::json!({
"email": "pat@example.com",
"chatgpt_user_id": "pat-user",
"chatgpt_account_id": "pat-workspace",
"chatgpt_plan_type": "team",
"chatgpt_account_is_fedramp": false
})
}
#[tokio::test]
async fn pat_import_verifies_identity_and_becomes_schedulable_without_oauth_refresh_fields() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
.and(header("authorization", "Bearer at-test-token"))
.and(header("accept", "application/json"))
.respond_with(ResponseTemplate::new(200).set_body_json(pat_identity()))
.expect(1)
.mount(&server)
.await;
let service = pat_service(&server);
let prepared = service
.prepare_import_document(serde_json::json!({
"accounts": [{
"platform": "openai", "type": "codex", "name": "team PAT",
"access_token": " at-test-token ",
"refresh_token": "unrelated-refresh-token",
"id_token": test_jwt(serde_json::json!({
"https://api.openai.com/auth": {"chatgpt_user_id": "untrusted-jwt-user"}
})),
"account_id": "untrusted-document-account",
"email": "untrusted-document@example.com",
"planType": "untrusted-plan", "expires_at": "2000-01-01T00:00:00Z"
}]
}))
.await
.expect("PAT import without a document user ID");
let prepared = &prepared.accounts()[0];
let account = &prepared.account;
assert_eq!(account.upstream_user_id(), Some("pat-user"));
assert_eq!(account.upstream_account_id(), Some("pat-workspace"));
assert_eq!(account.email(), Some("pat@example.com"));
assert_eq!(account.plan_type(), Some("team"));
assert_eq!(account.name(), "team PAT");
assert_eq!(
account.credential_state(),
gateway_core::account::CredentialState::Ready
);
assert_eq!(
account.status_projection(SystemTime::now(), None).status,
gateway_core::account::AccountStatus::Normal
);
assert!(!account.has_refresh_token());
assert!(account.access_token_expires_at().is_none());
assert!(account.next_refresh_at().is_none());
let runtime = CodexCredentialCodec::decode(&prepared.credential).expect("stored PAT");
let secret = runtime
.authentication
.oauth()
.expect("shared Bearer transport");
assert_eq!(secret.access_token.expose_secret(), "at-test-token");
assert!(secret.refresh_token.is_none());
assert!(secret.id_token.is_none());
let now = SystemTime::now();
let refresh_query = gateway_core::account::ProviderRefreshQuery::new(
account.provider().clone(),
now + Duration::from_secs(86_400),
now + Duration::from_secs(86_400),
now,
Vec::new(),
NonZeroU32::new(1).expect("nonzero limit"),
);
assert!(!refresh_query.contains(account));
let requests = server.received_requests().await.expect("requests");
assert_eq!(requests[0].headers["originator"], "Codex Desktop");
assert!(!requests[0].headers.contains_key("chatgpt-account-id"));
}
#[tokio::test]
async fn pat_import_times_out_without_falling_back_to_document_identity() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(pat_identity())
.set_delay(Duration::from_secs(2)),
)
.mount(&server)
.await;
let client = OpenAiTokenClient::new(
reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.timeout(Duration::from_millis(100))
.build()
.expect("short-lived test client"),
TokenClientConfig {
client_id: "test-public-client".to_owned(),
token_endpoint: format!("{}/oauth/token", server.uri()),
},
provider_openai::OpenAiConfig::default().wire_profile_state(),
);
let service = CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
)
.with_personal_access_token_client(Arc::new(client));
let error = service
.prepare_import_document(serde_json::json!({
"accessToken": "at-timeout-token", "userId": "untrusted-user"
}))
.await
.expect_err("timeout must abort PAT import");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::Unavailable
)
);
}
#[tokio::test]
async fn pat_import_requires_a_validation_client() {
let service = CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
);
let error = service
.prepare_import_document(serde_json::json!({
"accessToken": "at-test-token", "userId": "untrusted-user"
}))
.await
.expect_err("PAT cannot be imported without validation");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::Unavailable
)
);
}
#[tokio::test]
async fn pat_auth_json_import_and_cpr_round_trip_preserve_a_token_without_email() {
use gateway_core::account::LoadedCredential;
use provider_openai::credential::{CodexCredentialAdmin, ExportManagedCodexCredential};
let server = MockServer::start().await;
let mut identity = pat_identity();
identity.as_object_mut().expect("object").remove("email");
Mock::given(method("GET"))
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
.respond_with(ResponseTemplate::new(200).set_body_json(identity))
.expect(2)
.mount(&server)
.await;
let service = pat_service(&server);
let imported = service
.prepare_import_document(serde_json::json!({
"auth_mode": "personalAccessToken", "personal_access_token": "at-test-token"
}))
.await
.expect("official PAT auth.json");
let prepared = imported.into_accounts().remove(0);
assert!(prepared.account.email().is_none());
let exported = CodexCredentialAdmin
.format_cpr_export(vec![ExportManagedCodexCredential {
current: LoadedCredential {
account: prepared.account,
credential: prepared.credential,
},
added_at: Utc::now(),
updated_at: Utc::now(),
}])
.expect("export PAT using the existing credential schema");
let imported = service
.prepare_import_document(serde_json::to_value(exported).expect("export JSON"))
.await
.expect("re-import verifies PAT again");
assert_eq!(
imported.accounts()[0].account.upstream_user_id(),
Some("pat-user")
);
assert!(!imported.accounts()[0].account.has_refresh_token());
}
#[tokio::test]
async fn pat_import_fails_closed_for_rejection_and_unavailable_upstream_without_leaking_secrets() {
for (status, expected) in [
(401, PersonalAccessTokenError::Rejected),
(403, PersonalAccessTokenError::Rejected),
(429, PersonalAccessTokenError::Unavailable),
(500, PersonalAccessTokenError::Unavailable),
] {
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(status).set_body_string("pat-response-secret-marker"),
)
.expect(1)
.mount(&server)
.await;
let error = pat_service(&server)
.prepare_import_document(serde_json::json!({
"accessToken": "at-request-secret-marker", "userId": "document-user"
}))
.await
.expect_err("no account may be prepared after failed validation");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(expected)
);
let message = format!("{error} {error:?}");
assert!(!message.contains("pat-response-secret-marker"));
assert!(!message.contains("at-request-secret-marker"));
}
}
#[tokio::test]
async fn pat_import_rejects_missing_or_invalid_whoami_identity_fields() {
for field in [
"chatgpt_user_id",
"chatgpt_account_id",
"chatgpt_plan_type",
"chatgpt_account_is_fedramp",
] {
for value in [
None,
Some(serde_json::json!("")),
Some(serde_json::json!(null)),
] {
let server = MockServer::start().await;
let mut identity = pat_identity();
if let Some(value) = value {
identity[field] = value;
} else {
identity.as_object_mut().expect("object").remove(field);
}
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(200).set_body_json(identity))
.expect(1)
.mount(&server)
.await;
let error = pat_service(&server)
.prepare_import_document(serde_json::json!({
"access_token": "at-test-token", "user_id": "not-a-fallback"
}))
.await
.expect_err("malformed whoami must not fall back to document identity");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::InvalidResponse
)
);
}
}
}
#[tokio::test]
async fn pat_import_rejects_invalid_json_and_oversized_chunked_success_bodies() {
for body in ["not-json-secret-marker".to_owned(), "x".repeat(70 * 1024)] {
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("transfer-encoding", "chunked")
.set_body_string(body),
)
.expect(1)
.mount(&server)
.await;
let error = pat_service(&server)
.prepare_import_document(serde_json::json!({"accessToken": "at-test-token"}))
.await
.expect_err("invalid body");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::InvalidResponse
)
);
assert!(!format!("{error:?}").contains("secret-marker"));
}
}
#[tokio::test]
async fn pat_import_does_not_follow_redirects_or_forward_bearer_to_another_endpoint() {
let server = MockServer::start().await;
let destination = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(302).insert_header("location", destination.uri()))
.expect(1)
.mount(&server)
.await;
let error = pat_service(&server)
.prepare_import_document(serde_json::json!({"accessToken": "at-test-token"}))
.await
.expect_err("redirects are not accepted");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::Unavailable
)
);
assert!(
destination
.received_requests()
.await
.expect("redirect target")
.is_empty()
);
}
#[tokio::test]
async fn pat_import_rejects_malformed_tokens_before_sending_a_request() {
let server = MockServer::start().await;
let service = pat_service(&server);
for token in ["at-", "at-has space", "at-has\nnewline", "at-has\0control"] {
let error = service
.prepare_import_document(serde_json::json!({"accessToken": token}))
.await
.expect_err("invalid PAT");
assert_eq!(
error,
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
PersonalAccessTokenError::InvalidToken
)
);
}
assert!(
server
.received_requests()
.await
.expect("requests")
.is_empty()
);
}
#[tokio::test]
async fn ordinary_jwt_import_does_not_call_pat_whoami_even_when_client_is_configured() {
let server = MockServer::start().await;
let jwt = test_jwt(
serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "jwt-user"}}),
);
let imported = pat_service(&server)
.prepare_import_document(serde_json::json!({
"accessToken": jwt, "refreshToken": "jwt-refresh", "userId": "untrusted"
}))
.await
.expect("existing JWT behavior");
assert_eq!(
imported.accounts()[0].account.upstream_user_id(),
Some("jwt-user")
);
assert!(imported.accounts()[0].account.has_refresh_token());
assert!(
server
.received_requests()
.await
.expect("requests")
.is_empty()
);
}
#[tokio::test]
async fn oauth_import_uses_official_chatgpt_access_token_claims() {
let access_token = test_jwt(serde_json::json!({
@@ -281,3 +743,41 @@ async fn oauth_import_rejects_legacy_bare_token_field() {
provider_openai::credential::CodexCredentialAdminError::InvalidCredential
);
}
#[tokio::test]
async fn pat_import_uses_account_proxy_without_direct_fallback() {
for status in [200, 503] {
let origin = MockServer::start().await;
let proxy = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
.and(header("authorization", "Bearer at-proxy-token"))
.respond_with(ResponseTemplate::new(status).set_body_json(pat_identity()))
.expect(1)
.mount(&proxy)
.await;
let imported = pat_service(&origin)
.prepare_import_document(serde_json::json!({
"accessToken": "at-proxy-token",
"outboundProxyUrl": proxy.uri(),
}))
.await;
if status == 200 {
let prepared = imported.expect("PAT validated through account proxy");
assert_eq!(
prepared.accounts()[0].account.upstream_user_id(),
Some("pat-user")
);
assert!(prepared.accounts()[0].account.outbound_proxy().is_some());
} else {
assert!(imported.is_err(), "proxy failure must abort PAT import");
}
assert!(
origin
.received_requests()
.await
.expect("origin requests")
.is_empty()
);
}
}
@@ -6,6 +6,46 @@ use secrecy::{ExposeSecret, SecretString};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[tokio::test]
async fn refresh_and_code_exchange_use_the_selected_account_proxy() {
let direct = MockServer::start().await;
let proxy_a = MockServer::start().await;
let proxy_b = MockServer::start().await;
for (proxy, access) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] {
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"access_token": access,
"refresh_token": "rotated-refresh",
"id_token": "header.e30.signature"
})))
.expect(2)
.mount(proxy)
.await;
}
let client = client(&direct);
for (proxy, expected) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] {
let proxy = gateway_core::account::OutboundProxy::parse(&proxy.uri()).unwrap();
let refreshed = client
.refresh_with_proxy("initial-refresh", Some(&proxy))
.await
.unwrap();
assert_eq!(refreshed.access_token.as_deref(), Some(expected));
let exchanged = client
.exchange_with_proxy(
AuthorizationCodeGrant {
code: SecretString::from("test-code"),
code_verifier: SecretString::from("test-verifier"),
},
Some(&proxy),
)
.await
.unwrap();
assert_eq!(exchanged.secret.access_token.expose_secret(), expected);
}
assert!(direct.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn refresh_tracks_the_online_profile_without_contaminating_code_exchange() {
let server = MockServer::start().await;
@@ -0,0 +1,336 @@
use super::*;
use gateway_core::{
account::{CredentialRevision, OutboundProxy, ProviderAccount, ProviderAccountId},
routing::ProviderKind,
};
fn account(id: &str, proxy: Option<&str>) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(format!("acct_{id}")).unwrap(),
ProviderKind::new("openai").unwrap(),
id.to_owned(),
Some(id.to_owned()),
"oauth".to_owned(),
CredentialRevision::new(1).unwrap(),
None,
)
.with_outbound_proxy(proxy.map(|value| OutboundProxy::parse(value).unwrap()))
}
async fn http_exit(label: &'static str) -> (String, tokio::task::JoinHandle<String>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request_with_body(&mut stream).await;
let body = format!(
"event: response.completed\ndata: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_{label}\",\"status\":\"completed\",\"output\":[],\"usage\":{{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}}}}\n\n"
);
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
let header_end = request
.windows(4)
.position(|bytes| bytes == b"\r\n\r\n")
.unwrap();
String::from_utf8(request[..header_end].to_vec()).unwrap()
});
(format!("http://{address}"), task)
}
#[tokio::test]
async fn http_sse_accounts_use_distinct_proxies_and_clearing_restores_direct_client() {
let (exit_a, a) = http_exit("exit_a").await;
let (exit_b, b) = http_exit("exit_b").await;
let (direct_url, direct) = http_exit("direct").await;
let base = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
&direct_url,
test_wire_profile(),
);
let client_a = base.for_account(&account("a", Some(&exit_a))).unwrap();
let client_b = base.for_account(&account("b", Some(&exit_b))).unwrap();
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = true;
let (response_a, response_b) = tokio::join!(
client_a.create_response(&request, request_context("proxy-a", Some("a"))),
client_b.create_response(&request, request_context("proxy-b", Some("b")))
);
assert!(response_a.unwrap().body.contains("resp_exit_a"));
assert!(response_b.unwrap().body.contains("resp_exit_b"));
let cleared = client_a.for_account(&account("a", None)).unwrap();
assert!(
cleared
.create_response(&request, request_context("direct", Some("a")))
.await
.unwrap()
.body
.contains("resp_direct")
);
for head in [a.await.unwrap(), b.await.unwrap()] {
assert!(head.starts_with(&format!("POST {direct_url}/codex/responses HTTP/1.1")));
}
assert!(
direct
.await
.unwrap()
.starts_with("POST /codex/responses HTTP/1.1")
);
}
async fn websocket_exit(label: &'static str) -> (String, tokio::task::JoinHandle<String>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let connect = read_http_request(&mut stream).await;
assert!(connect.starts_with("CONNECT upstream.invalid:80 HTTP/1.1"));
stream
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
.await
.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
assert!(matches!(websocket.next().await, Some(Ok(Message::Text(_)))));
websocket
.send(Message::Text(
completed_websocket_response(label, 2, 1).into(),
))
.await
.unwrap();
connect
});
(format!("http://user:pass@{address}"), task)
}
#[tokio::test]
async fn websocket_connect_and_pool_are_isolated_when_account_proxy_changes() {
let (exit_a, a) = websocket_exit("resp_exit_a").await;
let (exit_b, b) = websocket_exit("resp_exit_b").await;
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let base = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
"http://upstream.invalid",
test_wire_profile(),
)
.with_websocket_pool(pool);
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.set_previous_response_id(Some("resp_previous".to_owned()));
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
for (proxy, expected) in [(exit_a, "resp_exit_a"), (exit_b, "resp_exit_b")] {
let client = base
.for_account(&account("same-account", Some(&proxy)))
.unwrap();
let response = timeout(
Duration::from_secs(5),
client.create_response(&request, request_context("proxy-ws", Some("same-account"))),
)
.await
.unwrap()
.unwrap();
assert!(response.body.contains(expected));
}
for head in [a.await.unwrap(), b.await.unwrap()] {
assert_eq!(
read_header_value(&head, "Proxy-Authorization"),
Some("Basic dXNlcjpwYXNz")
);
}
}
#[tokio::test]
async fn unavailable_proxy_fails_without_direct_fallback() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let unreachable = format!("http://{}", listener.local_addr().unwrap());
drop(listener);
let direct = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{}", direct.local_addr().unwrap()),
test_wire_profile(),
);
let client = base.for_account(&account("a", Some(&unreachable))).unwrap();
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = true;
assert!(
client
.create_response(&request, request_context("blocked", Some("a")))
.await
.is_err()
);
assert!(
timeout(Duration::from_millis(30), direct.accept())
.await
.is_err()
);
}
async fn socks_exit(stream: &mut TcpStream) -> (u8, String, u16) {
assert_eq!(stream.read_u8().await.unwrap(), 5);
let count = stream.read_u8().await.unwrap();
let mut methods = vec![0; usize::from(count)];
stream.read_exact(&mut methods).await.unwrap();
assert!(methods.contains(&2));
stream.write_all(&[5, 2]).await.unwrap();
assert_eq!(stream.read_u8().await.unwrap(), 1);
let count = stream.read_u8().await.unwrap();
let mut user = vec![0; usize::from(count)];
stream.read_exact(&mut user).await.unwrap();
let count = stream.read_u8().await.unwrap();
let mut password = vec![0; usize::from(count)];
stream.read_exact(&mut password).await.unwrap();
assert_eq!(user, b"user@exit");
assert_eq!(password, b"pass:word");
stream.write_all(&[1, 0]).await.unwrap();
let mut header = [0; 4];
stream.read_exact(&mut header).await.unwrap();
assert_eq!(&header[..3], &[5, 1, 0]);
let host = match header[3] {
1 => {
let mut ip = [0; 4];
stream.read_exact(&mut ip).await.unwrap();
std::net::Ipv4Addr::from(ip).to_string()
}
4 => {
let mut ip = [0; 16];
stream.read_exact(&mut ip).await.unwrap();
std::net::Ipv6Addr::from(ip).to_string()
}
3 => {
let count = stream.read_u8().await.unwrap();
let mut host = vec![0; usize::from(count)];
stream.read_exact(&mut host).await.unwrap();
String::from_utf8(host).unwrap()
}
_ => panic!("unexpected SOCKS address type"),
};
let port = stream.read_u16().await.unwrap();
stream
.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 80])
.await
.unwrap();
(header[3], host, port)
}
#[tokio::test]
async fn socks_account_egress_preserves_dns_mode_auth_and_ipv6_for_sse_and_websocket() {
for websocket in [false, true] {
let mut targets = vec![
("socks5", "localhost"),
("socks5h", "upstream.invalid"),
("socks5h", "[::1]"),
];
// reqwest's local resolver does not accept IPv6 literals; use socks5h for those.
if websocket {
targets.push(("socks5", "[::1]"));
}
for (scheme, host) in targets {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = format!(
"{scheme}://user%40exit:pass%3Aword@{}",
listener.local_addr().unwrap()
);
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let target = socks_exit(&mut stream).await;
if websocket {
let mut stream = accept_codex_test_websocket(stream).await;
assert!(matches!(stream.next().await, Some(Ok(Message::Text(_)))));
stream
.send(Message::Text(
completed_websocket_response("resp_socks", 2, 1).into(),
))
.await
.unwrap();
} else {
read_http_request_with_body(&mut stream).await;
let body = format!(
"data: {}\n\n",
completed_websocket_response("resp_socks", 2, 1)
);
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
}
target
});
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{host}"),
test_wire_profile(),
)
.for_account(&account("socks", Some(&proxy)))
.unwrap();
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = !websocket;
if websocket {
request.set_previous_response_id(Some("resp_previous".to_owned()));
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
}
let response = timeout(
Duration::from_secs(5),
client.create_response(&request, request_context("socks", Some("socks"))),
)
.await
.unwrap()
.unwrap_or_else(|error| {
panic!("websocket={websocket}, scheme={scheme}, host={host}: {error}")
});
assert!(response.body.contains("resp_socks"));
let (address_type, target, port) = task.await.unwrap();
assert_eq!(port, 80);
match host {
"localhost" => {
assert!(matches!(address_type, 1 | 4));
assert!(target.parse::<std::net::IpAddr>().unwrap().is_loopback());
}
"[::1]" => assert_eq!((address_type, target.as_str()), (4, "::1")),
_ => assert_eq!((address_type, target.as_str()), (3, host)),
}
}
}
}
#[tokio::test]
async fn https_account_proxy_starts_tls_and_rejection_never_falls_back_to_direct() {
for websocket in [false, true] {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = format!("https://{}", listener.local_addr().unwrap());
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let mut header = [0; 5];
stream.read_exact(&mut header).await.unwrap();
assert_eq!(&header[..2], &[22, 3]);
let mut hello = vec![0; usize::from(u16::from_be_bytes([header[3], header[4]]))];
stream.read_exact(&mut hello).await.unwrap();
stream.write_all(&[21, 3, 3, 0, 2, 2, 40]).await.unwrap();
});
let direct = TcpListener::bind("127.0.0.1:0").await.unwrap();
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{}", direct.local_addr().unwrap()),
test_wire_profile(),
)
.for_account(&account("tls-proxy", Some(&proxy)))
.unwrap();
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = !websocket;
if websocket {
request.set_previous_response_id(Some("resp_previous".to_owned()));
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
}
assert!(
timeout(
Duration::from_secs(5),
client.create_response(&request, request_context("tls-proxy", Some("tls-proxy")))
)
.await
.unwrap()
.is_err()
);
timeout(Duration::from_secs(5), task)
.await
.unwrap()
.unwrap();
assert!(
timeout(Duration::from_millis(30), direct.accept())
.await
.is_err()
);
}
}
@@ -1000,7 +1000,7 @@ async fn connection_local_continuation_should_fail_after_its_live_socket_disappe
let observation = error
.connection_observation()
.expect("disappeared connection should retain its lifecycle observation");
assert_eq!(observation.exit_reason(), "tcp_reset");
assert_eq!(observation.exit_reason(), "reset_without_closing_handshake");
assert!(observation.age_ms() >= observation.idle_ms());
}
@@ -45,6 +45,7 @@ use tokio_tungstenite::{
},
};
mod account_proxy;
mod canonical;
mod catalog;
mod client;
+1 -1
View File
@@ -20,7 +20,7 @@ gateway-protocol = { path = "../../gateway-protocol" }
getrandom = "0.4.3"
jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs"] }
# 与官方 Grok CLI 一致:reqwest 不启用压缩特性,上游请求不带 accept-encoding。
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2", "socks"] }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
semver = "1.0.27"

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