mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 16:05:27 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
90dd5b5a57 | ||
|
|
5ad9378b33 | ||
|
|
4a413f5bb7 | ||
|
|
0bab566d04 | ||
|
|
9b19cf083a | ||
|
|
8b22788746 | ||
|
|
e973fc516e | ||
|
|
12ecbf574a | ||
|
|
443bd6ad17 | ||
|
|
c795098a87 | ||
|
|
1f0009341d | ||
|
|
1359eaa894 | ||
|
|
74318af570 | ||
|
|
ddc10d7c1a | ||
|
|
6bc2189e6f | ||
|
|
b206ee8271 | ||
|
|
70ff923033 | ||
|
|
a2fbc073a4 | ||
|
|
9944041dc8 | ||
|
|
18a04f1213 | ||
|
|
35988d0eb7 |
@@ -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/*
|
||||
|
||||
@@ -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.
|
||||
@@ -10,8 +10,9 @@
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](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) 开源。
|
||||
|
||||
Generated
+9
-7
@@ -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
@@ -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(),
|
||||
})
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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,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;
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -65,7 +65,8 @@ pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
|
||||
let image_output_tokens =
|
||||
nested_number_field(body, &["tool_usage", "image_gen", "output_tokens"])
|
||||
.unwrap_or_default();
|
||||
let total_tokens = number_field(usage, "total_tokens").unwrap_or(input_tokens + output_tokens);
|
||||
let total_tokens = number_field(usage, "total_tokens")
|
||||
.unwrap_or_else(|| input_tokens.saturating_add(output_tokens));
|
||||
|
||||
let has_usage = [
|
||||
"input_tokens",
|
||||
@@ -98,7 +99,7 @@ pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
|
||||
})
|
||||
}
|
||||
|
||||
/// 校验计价所需的输入、输出和缓存读取字段,拒绝缺失、非法或不一致的数值。
|
||||
/// 校验计价所需的 token 字段,拒绝缺失、非法、溢出或不一致的数值。
|
||||
///
|
||||
/// 用量提取允许缺失字段补零;计价必须额外确认这些零值有明确的上游事实。
|
||||
#[must_use]
|
||||
@@ -122,11 +123,28 @@ pub fn billable_usage_is_complete(response: &Value, usage: TokenUsage) -> bool {
|
||||
Some(value) => value.as_u64(),
|
||||
None => Some(0),
|
||||
};
|
||||
let written = raw
|
||||
.pointer("/input_tokens_details/cache_write_tokens")
|
||||
.or_else(|| raw.pointer("/prompt_tokens_details/cache_write_tokens"))
|
||||
.or_else(|| raw.get("cache_write_tokens"));
|
||||
let written = match written {
|
||||
Some(value) => value.as_u64(),
|
||||
None => Some(0),
|
||||
};
|
||||
let total = usage.input_tokens.checked_add(usage.output_tokens);
|
||||
|
||||
input == Some(usage.input_tokens)
|
||||
&& output == Some(usage.output_tokens)
|
||||
&& cached == Some(usage.cached_tokens)
|
||||
&& usage.cached_tokens <= usage.input_tokens
|
||||
&& written == Some(usage.cache_write_tokens)
|
||||
&& usage
|
||||
.cached_tokens
|
||||
.checked_add(usage.cache_write_tokens)
|
||||
.is_some_and(|cached| cached <= usage.input_tokens)
|
||||
&& total == Some(usage.total_tokens)
|
||||
&& raw
|
||||
.get("total_tokens")
|
||||
.is_none_or(|value| value.as_u64() == total)
|
||||
}
|
||||
|
||||
/// 从完整 SSE 文本中提取最终可见用量。
|
||||
|
||||
@@ -5,6 +5,57 @@ use gateway_protocol::openai::events::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn billable_usage_should_validate_cache_writes_totals_and_overflow() {
|
||||
for (raw, expected) in [
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "input_tokens_details": {"cached_tokens": 20, "cache_write_tokens": 80}, "total_tokens": 110}),
|
||||
true,
|
||||
),
|
||||
(
|
||||
json!({"prompt_tokens": 100, "completion_tokens": 10, "prompt_tokens_details": {"cached_tokens": 20, "cache_write_tokens": 80}}),
|
||||
true,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "cached_tokens": 20, "cache_write_tokens": 90}),
|
||||
false,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "cache_write_tokens": -1}),
|
||||
false,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "input_tokens_details": {"cache_write_tokens": "1"}, "cache_write_tokens": 1}),
|
||||
false,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "cache_write_tokens": 1.5}),
|
||||
false,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "total_tokens": 111}),
|
||||
false,
|
||||
),
|
||||
(
|
||||
json!({"input_tokens": 100, "output_tokens": 10, "total_tokens": "110"}),
|
||||
false,
|
||||
),
|
||||
(json!({"input_tokens": u64::MAX, "output_tokens": 1}), false),
|
||||
(
|
||||
json!({"input_tokens": u64::MAX, "output_tokens": 0, "cached_tokens": u64::MAX, "cache_write_tokens": 1}),
|
||||
false,
|
||||
),
|
||||
] {
|
||||
let response = json!({"usage": raw});
|
||||
let usage = extract_usage(&response).expect("提取用量不得因溢出而崩溃");
|
||||
assert_eq!(
|
||||
billable_usage_is_complete(&response, usage),
|
||||
expected,
|
||||
"{response}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billable_usage_should_require_valid_input_output_and_cache_facts() {
|
||||
for (raw, expected) in [
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
@@ -55,10 +56,6 @@ pub(super) fn admin_account_store(pool: &PgPool) -> PgAdminAccountStore {
|
||||
|
||||
impl TestDatabase {
|
||||
pub(super) async fn create(label: &str) -> Option<Self> {
|
||||
Self::create_at(label, i64::MAX).await
|
||||
}
|
||||
|
||||
pub(super) async fn create_at(label: &str, version: i64) -> Option<Self> {
|
||||
let database_url = crate::support::test_env("CPR_TEST_DATABASE_URL")?;
|
||||
let schema = format!("cpr_store_{label}_{}", Uuid::new_v4().simple());
|
||||
let admin = PgPoolOptions::new()
|
||||
@@ -86,16 +83,10 @@ impl TestDatabase {
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect isolated test schema");
|
||||
sqlx::migrate::Migrator::with_migrations(
|
||||
TEST_MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.version <= version)
|
||||
.cloned()
|
||||
.collect(),
|
||||
)
|
||||
.run(&pool)
|
||||
.await
|
||||
.expect("apply test migrations");
|
||||
TEST_MIGRATOR
|
||||
.run(&pool)
|
||||
.await
|
||||
.expect("apply test migrations");
|
||||
Some(Self {
|
||||
admin,
|
||||
pool,
|
||||
@@ -240,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;
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use super::{TEST_MIGRATOR, TestDatabase};
|
||||
use super::TestDatabase;
|
||||
|
||||
async fn seed_request(pool: &sqlx::PgPool) {
|
||||
sqlx::query(
|
||||
@@ -101,8 +101,8 @@ async fn account_snapshots_are_required_before_live_foreign_keys_can_be_cleared(
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upgrade_normalizes_admin_identity_and_preserves_existing_event_rows() {
|
||||
let Some(db) = TestDatabase::create_at("audit_upgrade", 7).await else {
|
||||
async fn audit_requires_canonical_admin_identity_and_retains_it_after_admin_deletion() {
|
||||
let Some(db) = TestDatabase::create("audit_identity").await else {
|
||||
return;
|
||||
};
|
||||
sqlx::raw_sql(
|
||||
@@ -110,28 +110,12 @@ async fn upgrade_normalizes_admin_identity_and_preserves_existing_event_rows() {
|
||||
values ('admin_test', 'test_hash', now(), now());
|
||||
insert into admin_audit_events (id, actor_kind, actor_admin_user_id, actor_ref,
|
||||
action, entity_kind, entity_ref, created_at)
|
||||
values ('audit_before_upgrade', 'admin_session', 'admin_test', 'admin_test',
|
||||
'update', 'settings', '1', now());
|
||||
insert into ops_events (id, level, component, operation, failure_kind, message, created_at)
|
||||
values ('ops_before_upgrade', 'error', 'probe', 'probe', 'timeout', 'test', now());",
|
||||
values ('audit_identity', 'admin_session', 'admin_test', 'admin:admin_test',
|
||||
'update', 'settings', '1', now());",
|
||||
)
|
||||
.execute(&db.pool)
|
||||
.await
|
||||
.expect("seed data from pre-optimization schema");
|
||||
TEST_MIGRATOR
|
||||
.run(&db.pool)
|
||||
.await
|
||||
.expect("upgrade populated database");
|
||||
let identity: String = sqlx::query_scalar("select actor_ref from admin_audit_events")
|
||||
.fetch_one(&db.pool)
|
||||
.await
|
||||
.expect("read normalized identity");
|
||||
assert_eq!(identity, "admin:admin_test");
|
||||
let events: i64 = sqlx::query_scalar("select count(*) from ops_events")
|
||||
.fetch_one(&db.pool)
|
||||
.await
|
||||
.expect("read retained events");
|
||||
assert_eq!(events, 1);
|
||||
.expect("seed audit event with canonical admin identity");
|
||||
let error = sqlx::query("update admin_audit_events set actor_ref = 'admin_test'")
|
||||
.execute(&db.pool)
|
||||
.await
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -213,6 +213,7 @@ impl CodexWebSocketPoolSettings {
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct CodexQuotaSettings {
|
||||
/// 保留模型目录的刷新周期;额度独立每 30 秒检查周期复核和 reset 到期条件。
|
||||
pub refresh_interval_minutes: u64,
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -57,8 +57,8 @@ use snapshot::{
|
||||
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
pub(crate) const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
|
||||
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
|
||||
const EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL: Duration = Duration::from_secs(30 * 60);
|
||||
const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60);
|
||||
const EXHAUSTED_QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60);
|
||||
/// 首次 OAuth 异步观察失败时,由既有 quota worker 兜底重试的单轮上限。
|
||||
const INITIAL_QUOTA_SYNC_BATCH: usize = 100;
|
||||
// 5xx 上游拒绝的短退避重试预算;指数退避 1s/2s,吞掉瞬时抖动。
|
||||
@@ -68,8 +68,8 @@ const QUOTA_FETCH_5XX_BASE_DELAY: Duration = Duration::from_secs(1);
|
||||
/// OpenAI Provider 主动额度刷新的调度策略。
|
||||
///
|
||||
/// 正常账号依赖请求响应的被动额度同步;周期 worker 仅复核已耗尽账号。
|
||||
/// 该策略只控制 Provider 自己的 quota worker 频率,避免把 OpenAI
|
||||
/// 额度语义泄漏到公共调度层。
|
||||
/// 该策略保留模型目录的周期刷新频率;额度到期检查使用独立的短周期,
|
||||
/// 避免到达 reset 后还要等待完整的目录刷新周期。
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct CodexQuotaRefreshPolicy {
|
||||
interval: Duration,
|
||||
@@ -204,7 +204,12 @@ struct CodexQuotaSchedulingProjection {
|
||||
struct CodexQuotaProjectionState {
|
||||
next_version: u64,
|
||||
entries: BTreeMap<ProviderAccountId, CodexQuotaSchedulingEntry>,
|
||||
last_periodic_refresh_at: BTreeMap<ProviderAccountId, Instant>,
|
||||
last_periodic_refresh_at: BTreeMap<ProviderAccountId, CodexQuotaRefreshAttempt>,
|
||||
}
|
||||
|
||||
struct CodexQuotaRefreshAttempt {
|
||||
monotonic_at: Instant,
|
||||
wall_at: SystemTime,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
@@ -398,16 +403,27 @@ impl CodexQuotaSchedulingProjection {
|
||||
.last_periodic_refresh_at
|
||||
.retain(|account_id, _| candidate_ids.contains(account_id));
|
||||
|
||||
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只复核
|
||||
// 已耗尽且已到 reset/fallback 时刻的账号;重复复核至少间隔 30 分钟。
|
||||
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。
|
||||
// 已耗尽账号每 30 分钟复核,不能等待旧 reset:官方活动可能提前重置额度。
|
||||
// reset + 2 分钟可提前触发一次复核,给上游重置留出传播时间。
|
||||
let mut reserved = Vec::new();
|
||||
for account in candidates {
|
||||
if !periodic_quota_refresh_due(&state, account.id(), refreshed_at) {
|
||||
if !periodic_quota_refresh_due(
|
||||
&state,
|
||||
account.id(),
|
||||
account.quota().reset_at(),
|
||||
now,
|
||||
refreshed_at,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
state
|
||||
.last_periodic_refresh_at
|
||||
.insert(account.id().clone(), refreshed_at);
|
||||
state.last_periodic_refresh_at.insert(
|
||||
account.id().clone(),
|
||||
CodexQuotaRefreshAttempt {
|
||||
monotonic_at: refreshed_at,
|
||||
wall_at: now,
|
||||
},
|
||||
);
|
||||
reserved.push(account);
|
||||
}
|
||||
reserved
|
||||
@@ -415,24 +431,27 @@ impl CodexQuotaSchedulingProjection {
|
||||
}
|
||||
|
||||
fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option<ProviderAccount> {
|
||||
// reset 到期是“应刷新”的信号,不是“已恢复”的证据。
|
||||
(eligible_periodic_quota_refresh(&account, now)
|
||||
&& account
|
||||
.quota()
|
||||
.exhaustion_refresh_due(now, EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL))
|
||||
.then_some(account)
|
||||
(eligible_periodic_quota_refresh(&account, now) && account.quota().is_exhausted())
|
||||
.then_some(account)
|
||||
}
|
||||
|
||||
fn periodic_quota_refresh_due(
|
||||
state: &CodexQuotaProjectionState,
|
||||
account_id: &ProviderAccountId,
|
||||
now: Instant,
|
||||
reset_at: Option<SystemTime>,
|
||||
now: SystemTime,
|
||||
monotonic_now: Instant,
|
||||
) -> bool {
|
||||
state
|
||||
.last_periodic_refresh_at
|
||||
.get(account_id)
|
||||
.is_none_or(|last| {
|
||||
now.saturating_duration_since(*last) >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
|
||||
monotonic_now.saturating_duration_since(last.monotonic_at)
|
||||
>= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
|
||||
|| reset_at
|
||||
.and_then(|reset| reset.checked_add(EXHAUSTED_QUOTA_RESET_GRACE))
|
||||
// 已在该边界之后复核过时回到周期重试,避免过期 reset 每轮触发。
|
||||
.is_some_and(|due_at| last.wall_at < due_at && due_at <= now)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1206,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(),
|
||||
@@ -1229,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(),
|
||||
@@ -1310,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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -145,9 +145,9 @@ impl OpenAiResponseObservationState {
|
||||
}
|
||||
|
||||
pub(super) fn effective_service_tier(&self) -> Option<&str> {
|
||||
self.requested_service_tier
|
||||
self.upstream_service_tier
|
||||
.as_deref()
|
||||
.or(self.upstream_service_tier.as_deref())
|
||||
.or(self.requested_service_tier.as_deref())
|
||||
}
|
||||
|
||||
pub(super) fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
|
||||
|
||||
@@ -7,8 +7,10 @@ pub(super) const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60);
|
||||
pub(super) const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
|
||||
pub(super) const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60);
|
||||
pub(super) const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
|
||||
pub(super) const QUOTA_CHECK_INTERVAL: Duration = Duration::from_secs(30);
|
||||
pub(super) const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release";
|
||||
pub(super) const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag";
|
||||
pub(super) const MODEL_CATALOG_WORKER_OWNER: &str = "openai-model-catalog";
|
||||
|
||||
pub(crate) fn worker_contributions(
|
||||
refresh: Arc<CodexCredentialRefreshService>,
|
||||
@@ -20,6 +22,7 @@ pub(crate) fn worker_contributions(
|
||||
) -> Result<Vec<WorkerContribution>, WorkerDefinitionError> {
|
||||
let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, PROVIDER_NAME)?;
|
||||
let quota_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, PROVIDER_NAME)?;
|
||||
let catalog_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_CATALOG_WORKER_OWNER)?;
|
||||
let etag_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_ETAG_WORKER_OWNER)?;
|
||||
let desktop_release_id =
|
||||
WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?;
|
||||
@@ -34,9 +37,13 @@ pub(crate) fn worker_contributions(
|
||||
contributions.extend([
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
quota_id,
|
||||
QUOTA_CHECK_INTERVAL,
|
||||
Box::new(OpenAiQuotaTask { quota }),
|
||||
)?),
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
catalog_id,
|
||||
quota_refresh_policy.interval(),
|
||||
Box::new(OpenAiQuotaTask {
|
||||
quota,
|
||||
Box::new(OpenAiCatalogTask {
|
||||
catalog: Arc::clone(&catalog),
|
||||
}),
|
||||
)?),
|
||||
@@ -159,6 +166,9 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
|
||||
|
||||
pub(super) struct OpenAiQuotaTask {
|
||||
quota: Arc<CodexCredentialQuotaService>,
|
||||
}
|
||||
|
||||
pub(super) struct OpenAiCatalogTask {
|
||||
catalog: Arc<CodexCredentialCatalogService>,
|
||||
}
|
||||
|
||||
@@ -195,7 +205,6 @@ impl ScheduledTask for OpenAiQuotaTask {
|
||||
if context.cancellation().is_cancelled() {
|
||||
return Ok(());
|
||||
}
|
||||
let mut failures = false;
|
||||
match self.quota.synchronize().await {
|
||||
Ok(summary) if summary.has_operational_failures() => {
|
||||
tracing::warn!(
|
||||
@@ -209,30 +218,33 @@ impl ScheduledTask for OpenAiQuotaTask {
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(error) => {
|
||||
failures = true;
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
"OpenAI quota synchronization failed"
|
||||
);
|
||||
return Err(WorkerTaskError::safe("OpenAI quota synchronization failed"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ScheduledTask for OpenAiCatalogTask {
|
||||
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
if context.cancellation().is_cancelled() {
|
||||
return Ok(());
|
||||
}
|
||||
match self.catalog.refresh_catalogs().await {
|
||||
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {}
|
||||
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => Ok(()),
|
||||
Err(error) => {
|
||||
failures = true;
|
||||
tracing::warn!(error = %error, "OpenAI model catalog refresh failed");
|
||||
Err(WorkerTaskError::safe(
|
||||
"OpenAI model catalog synchronization failed",
|
||||
))
|
||||
}
|
||||
}
|
||||
if failures {
|
||||
Err(WorkerTaskError::safe(
|
||||
"OpenAI quota or catalog synchronization failed",
|
||||
))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -158,7 +158,7 @@ impl CodexCanonicalDecoder {
|
||||
self
|
||||
}
|
||||
|
||||
/// 请求显式指定的档位优先用于计费;响应档位只在请求未指定时兜底。
|
||||
/// 响应未报告实际档位时,才使用请求档位估算费用。
|
||||
#[must_use]
|
||||
pub fn with_requested_service_tier(mut self, service_tier: Option<&str>) -> Self {
|
||||
self.requested_service_tier = normalize_service_tier(service_tier);
|
||||
@@ -862,18 +862,19 @@ impl CodexCanonicalDecoder {
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.to_owned();
|
||||
let service_tier = self
|
||||
.requested_service_tier
|
||||
.response_service_tier
|
||||
.as_deref()
|
||||
.or(self.response_service_tier.as_deref());
|
||||
let web_search_calls = web_search_call_count(response);
|
||||
.or(self.requested_service_tier.as_deref());
|
||||
let tool_calls = billable_tool_calls(response);
|
||||
if let Some(breakdown) = usage
|
||||
.filter(|usage| billable_usage_is_complete(response, *usage))
|
||||
.and_then(|usage| {
|
||||
let web_search_calls = web_search_calls?;
|
||||
let (web_search_calls, file_search_calls) = tool_calls?;
|
||||
openai_billing_breakdown(
|
||||
&model,
|
||||
OpenAiBillingUsage::from(usage)
|
||||
.with_web_search_calls(web_search_calls, self.web_search_pricing),
|
||||
.with_web_search_calls(web_search_calls, self.web_search_pricing)
|
||||
.with_file_search_calls(file_search_calls),
|
||||
service_tier,
|
||||
)
|
||||
})
|
||||
@@ -1000,26 +1001,43 @@ fn core_usage(usage: TokenUsage) -> Usage {
|
||||
normalized
|
||||
}
|
||||
|
||||
fn web_search_call_count(response: &Value) -> Option<u64> {
|
||||
let mut calls = 0_u64;
|
||||
for item in response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("web_search_call"))
|
||||
{
|
||||
match item
|
||||
.get("action")
|
||||
.and_then(|action| action.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
Some("search") => calls = calls.checked_add(1)?,
|
||||
Some("open_page" | "find_in_page") => {}
|
||||
Some(_) | None => return None,
|
||||
fn billable_tool_calls(response: &Value) -> Option<(u64, u64)> {
|
||||
let mut web = 0_u64;
|
||||
let mut file = 0_u64;
|
||||
let items = match response.get("output") {
|
||||
Some(value) => value.as_array()?.as_slice(),
|
||||
None => &[],
|
||||
};
|
||||
for item in items {
|
||||
match item.get("type").and_then(Value::as_str)? {
|
||||
"web_search_call" => match item.pointer("/action/type").and_then(Value::as_str) {
|
||||
Some("search") => web = web.checked_add(1)?,
|
||||
Some("open_page" | "find_in_page") => {}
|
||||
_ => return None,
|
||||
},
|
||||
"file_search_call"
|
||||
if item.get("status").and_then(Value::as_str) == Some("completed") =>
|
||||
{
|
||||
file = file.checked_add(1)?;
|
||||
}
|
||||
// 这些输出没有独立的 OpenAI 工具调用费。
|
||||
"message"
|
||||
| "reasoning"
|
||||
| "function_call"
|
||||
| "custom_tool_call"
|
||||
| "computer_call"
|
||||
| "local_shell_call"
|
||||
| "apply_patch_call"
|
||||
| "mcp_call"
|
||||
| "mcp_list_tools"
|
||||
| "mcp_approval_request"
|
||||
| "compaction" => {}
|
||||
// 容器会话、生图及未知工具需要额外计费事实,不能把 token 小计
|
||||
// 当作完整费用。
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
Some(calls)
|
||||
Some((web, file))
|
||||
}
|
||||
|
||||
fn incomplete_finish_reason(response: &Value) -> FinishReason {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -16,6 +16,7 @@ use super::{
|
||||
const LONG_CONTEXT_THRESHOLD: u64 = 272_000;
|
||||
const WEB_SEARCH_CALL_TICKS: u128 = 100_000_000;
|
||||
const WEB_SEARCH_PREVIEW_NON_REASONING_CALL_TICKS: u128 = 250_000_000;
|
||||
const FILE_SEARCH_CALL_TICKS: u128 = 25_000_000;
|
||||
|
||||
/// OpenAI 公开 Token 价格计算所需的单次用量事实。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
@@ -27,6 +28,7 @@ pub struct OpenAiBillingUsage {
|
||||
image_input_tokens: u64,
|
||||
image_output_tokens: u64,
|
||||
web_search_calls: u64,
|
||||
file_search_calls: u64,
|
||||
web_search_pricing: Option<WebSearchPricing>,
|
||||
}
|
||||
|
||||
@@ -47,6 +49,7 @@ impl OpenAiBillingUsage {
|
||||
image_input_tokens: 0,
|
||||
image_output_tokens: 0,
|
||||
web_search_calls: 0,
|
||||
file_search_calls: 0,
|
||||
web_search_pricing: None,
|
||||
}
|
||||
}
|
||||
@@ -60,6 +63,11 @@ impl OpenAiBillingUsage {
|
||||
self.web_search_pricing = pricing;
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) const fn with_file_search_calls(mut self, calls: u64) -> Self {
|
||||
self.file_search_calls = calls;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl From<TokenUsage> for OpenAiBillingUsage {
|
||||
@@ -72,6 +80,7 @@ impl From<TokenUsage> for OpenAiBillingUsage {
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
web_search_calls: 0,
|
||||
file_search_calls: 0,
|
||||
web_search_pricing: None,
|
||||
}
|
||||
}
|
||||
@@ -125,6 +134,7 @@ struct ModelPricing {
|
||||
long_flex: TokenRates,
|
||||
long_fast: TokenRates,
|
||||
cache_write_percent: u32,
|
||||
unpriced_long_context: bool,
|
||||
}
|
||||
|
||||
impl ModelPricing {
|
||||
@@ -137,6 +147,7 @@ impl ModelPricing {
|
||||
long_flex: TokenRates::ZERO,
|
||||
long_fast: TokenRates::ZERO,
|
||||
cache_write_percent: 0,
|
||||
unpriced_long_context: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,7 +181,15 @@ impl ModelPricing {
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_unpriced_long_context(mut self) -> Self {
|
||||
self.unpriced_long_context = true;
|
||||
self
|
||||
}
|
||||
|
||||
fn rates(self, tier: PricingTier, long_context: bool) -> Option<TokenRates> {
|
||||
if long_context && self.unpriced_long_context {
|
||||
return None;
|
||||
}
|
||||
// Only models with a published long-context column switch at the threshold.
|
||||
// A dash in that column is not a second, unavailable price tier for models
|
||||
// whose entire supported context is covered by the short-context price.
|
||||
@@ -200,10 +219,12 @@ struct PricingRule {
|
||||
pricing: ModelPricing,
|
||||
}
|
||||
|
||||
// Source: https://developers.openai.com/api/docs/pricing (verified 2026-08-14).
|
||||
// 价格来源:https://developers.openai.com/api/docs/pricing,核验日期 2026-09-09。
|
||||
// 使用常规价,不采用 Sol 的临时 $4/$20 优惠;缓存、Flex、Fast 和长上下文
|
||||
// 档位也统一按常规价计算。
|
||||
const PRICING_RULES: &[PricingRule] = &[
|
||||
// Astra: https://developers.openai.com/api/docs/models/gpt-6-astra
|
||||
// Rates verified against the pricing table on 2026-09-05.
|
||||
// Astra:https://developers.openai.com/api/docs/models/gpt-6-astra
|
||||
// 已于 2026-09-09 对照官方价目表核验。
|
||||
PricingRule {
|
||||
model: "gpt-6-astra",
|
||||
pricing: ModelPricing::new(100_000, 500_000, 10_000)
|
||||
@@ -405,10 +426,60 @@ const PRICING_RULES: &[PricingRule] = &[
|
||||
model: "gpt-3.5-turbo",
|
||||
pricing: ModelPricing::new(5_000, 15_000, 0),
|
||||
},
|
||||
// 各变体独立采用已公开的价格与档位,不能继承父型号的所有档位。
|
||||
// Cyber 长上下文的官方来源存在差异,该区间暂不估价。
|
||||
PricingRule {
|
||||
model: "gpt-5.1-codex-mini",
|
||||
pricing: ModelPricing::new(2_500, 20_000, 250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.3-chat-latest",
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.6-cyber",
|
||||
pricing: ModelPricing::new(125_000, 750_000, 12_500)
|
||||
.with_cache_write(125)
|
||||
.with_unpriced_long_context(),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.5-cyber",
|
||||
pricing: ModelPricing::new(125_000, 750_000, 12_500).with_unpriced_long_context(),
|
||||
},
|
||||
PricingRule {
|
||||
model: "chat-latest",
|
||||
pricing: ModelPricing::new(50_000, 300_000, 5_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5-codex",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.1-codex",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.1-codex-max",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5-chat-latest",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.1-chat-latest",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.2-codex",
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.2-chat-latest",
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750),
|
||||
},
|
||||
];
|
||||
|
||||
const UNPRICED_MODELS: &[&str] = &["gpt-5.3-codex-spark"];
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct TokenAmounts {
|
||||
input_ticks: u128,
|
||||
@@ -439,13 +510,14 @@ fn openai_billing_breakdown_with_context(
|
||||
service_tier: Option<&str>,
|
||||
long_context: bool,
|
||||
) -> Option<CalculatedCostBreakdown> {
|
||||
if usage.cached_tokens > usage.input_tokens
|
||||
if usage.cached_tokens.checked_add(usage.cache_write_tokens)? > usage.input_tokens
|
||||
|| usage.image_input_tokens > 0
|
||||
|| usage.image_output_tokens > 0
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let web_search_ticks = web_search_amount_ticks(usage)?;
|
||||
let tool_ticks = web_search_amount_ticks(usage)?
|
||||
.checked_add(u128::from(usage.file_search_calls).checked_mul(FILE_SEARCH_CALL_TICKS)?)?;
|
||||
let pricing = model_pricing(model)?;
|
||||
let normalized_tier = normalize_service_tier(service_tier);
|
||||
let tier = pricing_tier(normalized_tier.as_deref())?;
|
||||
@@ -467,8 +539,8 @@ fn openai_billing_breakdown_with_context(
|
||||
usage.cached_tokens,
|
||||
usage.cache_write_tokens,
|
||||
)?;
|
||||
standard.total_ticks = standard.total_ticks.checked_add(web_search_ticks)?;
|
||||
selected.total_ticks = selected.total_ticks.checked_add(web_search_ticks)?;
|
||||
standard.total_ticks = standard.total_ticks.checked_add(tool_ticks)?;
|
||||
selected.total_ticks = selected.total_ticks.checked_add(tool_ticks)?;
|
||||
let multiplier_percent =
|
||||
effective_multiplier_percent(selected.total_ticks, standard.total_ticks)?;
|
||||
let cache_write_rate = cache_write_rate(selected_rates, pricing.cache_write_percent)?;
|
||||
@@ -503,7 +575,12 @@ pub(crate) fn image_calculated_cost(request_body: &[u8], usage: &Value) -> Optio
|
||||
let request = serde_json::from_slice::<ImageModel>(request_body).ok()?;
|
||||
if !matches!(
|
||||
request.model.as_str(),
|
||||
"gpt-image-2" | "gpt-image-2-2026-04-21"
|
||||
"gpt-image-2"
|
||||
| "gpt-image-2-2026-04-21"
|
||||
| "gpt-image-2.5-sunburst"
|
||||
| "gpt-image-2.5-sunburst-2026-09-08"
|
||||
| "gpt-image-2.5-flare"
|
||||
| "gpt-image-2.5-flare-2026-09-08"
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
@@ -542,8 +619,9 @@ pub(crate) fn image_calculated_cost(request_body: &[u8], usage: &Value) -> Optio
|
||||
(_, _, cached) if cached == input => (text_input, image_input),
|
||||
_ => return None,
|
||||
};
|
||||
// https://developers.openai.com/api/docs/pricing (2026-09-08)
|
||||
// USD / 1M tokens: text 5 / cached 1.25; image 8 / cached 2 / output 30.
|
||||
// 价格来源:https://developers.openai.com/api/docs/pricing,核验日期 2026-09-09。
|
||||
// GPT Image 2 和两个 2.5 型号按相同的 token 单价计费,不使用按张估价。
|
||||
// 每百万 token 的美元单价:文本 5 / 缓存 1.25;图片 8 / 缓存 2 / 输出 30。
|
||||
let text = token_amounts(
|
||||
TokenRates::new(50_000, 0, 12_500),
|
||||
0,
|
||||
@@ -565,16 +643,10 @@ pub(crate) fn image_calculated_cost(request_body: &[u8], usage: &Value) -> Optio
|
||||
|
||||
fn model_pricing(model: &str) -> Option<ModelPricing> {
|
||||
let normalized = normalize_model_name(model);
|
||||
if UNPRICED_MODELS
|
||||
.iter()
|
||||
.any(|rule| model_matches_rule(&normalized, rule))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let model = pricing_model_name(&normalized);
|
||||
PRICING_RULES
|
||||
.iter()
|
||||
.filter(|rule| model_matches_rule(&normalized, rule.model))
|
||||
.max_by_key(|rule| rule.model.len())
|
||||
.find(|rule| model == rule.model)
|
||||
.map(|rule| rule.pricing)
|
||||
}
|
||||
|
||||
@@ -613,16 +685,17 @@ fn fixed_block_web_search_model(model: &str) -> bool {
|
||||
|
||||
fn reasoning_model(model: &str) -> bool {
|
||||
let normalized = normalize_model_name(model);
|
||||
normalized.starts_with("gpt-5")
|
||||
|| model_matches_rule(&normalized, "gpt-6-astra")
|
||||
|| normalized.starts_with("o1")
|
||||
|| normalized.starts_with("o3")
|
||||
|| normalized.starts_with("o4")
|
||||
let model = pricing_model_name(&normalized);
|
||||
model.starts_with("gpt-5")
|
||||
|| model == "gpt-6-astra"
|
||||
|| model.starts_with("o1")
|
||||
|| model.starts_with("o3")
|
||||
|| model.starts_with("o4")
|
||||
}
|
||||
|
||||
fn pricing_tier(service_tier: Option<&str>) -> Option<PricingTier> {
|
||||
match service_tier {
|
||||
None | Some("auto" | "default" | "standard") => Some(PricingTier::Standard),
|
||||
None | Some("default" | "standard") => Some(PricingTier::Standard),
|
||||
Some("flex") => Some(PricingTier::Flex),
|
||||
Some("fast" | "priority") => Some(PricingTier::Fast),
|
||||
Some(_) => None,
|
||||
@@ -644,20 +717,23 @@ fn token_amounts(
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
) -> Option<TokenAmounts> {
|
||||
if cached_tokens.checked_add(cache_write_tokens)? > input_tokens {
|
||||
return None;
|
||||
}
|
||||
let billed_cache_read = if rates.cache_read_ticks > 0 {
|
||||
cached_tokens.min(input_tokens)
|
||||
cached_tokens
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let cache_write_rate = cache_write_rate(rates, cache_write_percent)?;
|
||||
let billed_cache_write = if cache_write_rate > 0 {
|
||||
cache_write_tokens.min(input_tokens.saturating_sub(billed_cache_read))
|
||||
cache_write_tokens
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let uncached_input = input_tokens
|
||||
.saturating_sub(billed_cache_read)
|
||||
.saturating_sub(billed_cache_write);
|
||||
.checked_sub(billed_cache_read)?
|
||||
.checked_sub(billed_cache_write)?;
|
||||
let input_ticks = u128::from(uncached_input).checked_mul(rates.input_ticks)?;
|
||||
let output_ticks = u128::from(output_tokens).checked_mul(rates.output_ticks)?;
|
||||
let cache_read_ticks = u128::from(billed_cache_read).checked_mul(rates.cache_read_ticks)?;
|
||||
@@ -703,13 +779,47 @@ fn normalize_model_name(model: &str) -> String {
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn model_matches_rule(model: &str, rule: &str) -> bool {
|
||||
if model == rule {
|
||||
return true;
|
||||
// 仅已核验的别名和快照可以共用价格;未知后缀、未来日期和微调模型
|
||||
// 不得直接继承父型号的价格。
|
||||
fn pricing_model_name(model: &str) -> &str {
|
||||
match model {
|
||||
"gpt-3.5-turbo-0125" => "gpt-3.5-turbo",
|
||||
"gpt-4-0314" => "gpt-4",
|
||||
"gpt-4-0613" => "gpt-4",
|
||||
"gpt-4-turbo-2024-04-09" => "gpt-4-turbo",
|
||||
"gpt-4.1-2025-04-14" => "gpt-4.1",
|
||||
"gpt-4.1-mini-2025-04-14" => "gpt-4.1-mini",
|
||||
"gpt-4.1-nano-2025-04-14" => "gpt-4.1-nano",
|
||||
"gpt-4o-2024-08-06" => "gpt-4o",
|
||||
"gpt-4o-2024-11-20" => "gpt-4o",
|
||||
"gpt-4o-mini-2024-07-18" => "gpt-4o-mini",
|
||||
"gpt-5-2025-08-07" => "gpt-5",
|
||||
"gpt-5-mini-2025-08-07" => "gpt-5-mini",
|
||||
"gpt-5-nano-2025-08-07" => "gpt-5-nano",
|
||||
"gpt-5-pro-2025-10-06" => "gpt-5-pro",
|
||||
"gpt-5.1-2025-11-13" => "gpt-5.1",
|
||||
"gpt-5.2-2025-12-11" => "gpt-5.2",
|
||||
"gpt-5.2-pro-2025-12-11" => "gpt-5.2-pro",
|
||||
"gpt-5.4-2026-03-05" => "gpt-5.4",
|
||||
"gpt-5.4-mini-2026-03-17" => "gpt-5.4-mini",
|
||||
"gpt-5.4-nano-2026-03-17" => "gpt-5.4-nano",
|
||||
"gpt-5.4-pro-2026-03-05" => "gpt-5.4-pro",
|
||||
"gpt-5.5-2026-04-23" => "gpt-5.5",
|
||||
"gpt-5.5-pro-2026-04-23" => "gpt-5.5-pro",
|
||||
"gpt-daybreak-blue-latest" => "gpt-5.6-sol",
|
||||
"gpt-daybreak-red-latest" => "gpt-5.6-cyber",
|
||||
"o1-2024-12-17" => "o1",
|
||||
"o1-pro-2025-03-19" => "o1-pro",
|
||||
"o3-2025-04-16" => "o3",
|
||||
"o3-mini-2025-01-31" => "o3-mini",
|
||||
"o3-pro-2025-06-10" => "o3-pro",
|
||||
"o4-mini-2025-04-16" => "o4-mini",
|
||||
_ => model,
|
||||
}
|
||||
model
|
||||
.strip_prefix(rule)
|
||||
.is_some_and(|suffix| matches!(suffix.as_bytes().first(), Some(b'-' | b'.' | b':')))
|
||||
}
|
||||
|
||||
fn model_matches_rule(model: &str, rule: &str) -> bool {
|
||||
pricing_model_name(model) == rule
|
||||
}
|
||||
|
||||
/// 规范化请求或响应携带的服务档位,供观测与计费共用。
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -36,7 +36,8 @@ const DEFAULT_MAX_CONNECTING: usize = 16;
|
||||
const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
|
||||
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
// 心跳也覆盖正在生成的连接,给短时链路停顿留出恢复余量。
|
||||
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
pub(crate) const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
|
||||
/// WebSocket 连接池。
|
||||
@@ -67,9 +68,9 @@ pub struct CodexWebSocketPoolConfig {
|
||||
pub max_connecting: usize,
|
||||
/// 后台维护间隔;`None` 表示不启动后台任务。
|
||||
pub maintenance_interval: Option<Duration>,
|
||||
/// idle socket 探活 ping 间隔;`None` 表示维护时只做过期清理。
|
||||
/// 池化连接的探活 ping 间隔(包括正在生成的连接);`None` 表示不主动 ping。
|
||||
pub ping_interval: Option<Duration>,
|
||||
/// 发送 ping 后等待上游响应的超时时间;零值表示不校验 Pong deadline。
|
||||
/// 发送 ping 后等待任意入站帧的超时时间;零值表示不校验 Pong deadline。
|
||||
pub ping_timeout: Duration,
|
||||
/// idle socket 无活动多久后视为失活。
|
||||
pub liveness_timeout: Option<Duration>,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user