mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8b22788746 | ||
|
|
e973fc516e | ||
|
|
12ecbf574a | ||
|
|
443bd6ad17 | ||
|
|
c795098a87 | ||
|
|
1f0009341d | ||
|
|
1359eaa894 | ||
|
|
74318af570 | ||
|
|
ddc10d7c1a | ||
|
|
6bc2189e6f |
@@ -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 启用状态。
|
||||
@@ -217,3 +223,19 @@ pub struct ClientKeyMutation {
|
||||
pub record: Option<ClientKeyRecord>,
|
||||
pub id: ClientApiKeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UnresolvedClientCharge {
|
||||
pub request_id: String,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub state: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReconcileClientCharge {
|
||||
pub key_id: ClientApiKeyId,
|
||||
pub request_id: String,
|
||||
pub amount_usd: gateway_core::metering::Decimal,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -220,6 +220,29 @@ pub trait AuthStore: Send + Sync {
|
||||
/// Client API Key 管理写入。
|
||||
#[async_trait]
|
||||
pub trait ClientKeyStore: Send + Sync {
|
||||
async fn unresolved_charges(
|
||||
&self,
|
||||
_id: &gateway_core::policy::ClientApiKeyId,
|
||||
) -> AdminStoreResult<Vec<crate::model::client_keys::UnresolvedClientCharge>> {
|
||||
Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
"client budget",
|
||||
"unsupported",
|
||||
))
|
||||
}
|
||||
|
||||
async fn reconcile_charge(
|
||||
&self,
|
||||
_command: crate::model::client_keys::ReconcileClientCharge,
|
||||
_context: &MutationContext,
|
||||
) -> AdminStoreResult<()> {
|
||||
Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
"client budget",
|
||||
"unsupported",
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage>;
|
||||
|
||||
async fn reveal_client_key(
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -26,6 +26,19 @@ use super::{map_store_error, publish_committed};
|
||||
/// API 消费的 Client Key 管理服务。
|
||||
#[async_trait]
|
||||
pub trait ClientKeyService: Send + Sync {
|
||||
async fn unresolved_charges(
|
||||
&self,
|
||||
_id: &ClientApiKeyId,
|
||||
) -> Result<Vec<crate::model::client_keys::UnresolvedClientCharge>, AdminError> {
|
||||
Err(AdminError::unavailable("client budget unavailable"))
|
||||
}
|
||||
async fn reconcile_charge(
|
||||
&self,
|
||||
_context: &MutationContext,
|
||||
_command: crate::model::client_keys::ReconcileClientCharge,
|
||||
) -> Result<(), AdminError> {
|
||||
Err(AdminError::unavailable("client budget unavailable"))
|
||||
}
|
||||
async fn list(&self, query: ClientKeyListQuery) -> Result<ClientKeyPage, AdminError>;
|
||||
async fn reveal(&self, id: &ClientApiKeyId) -> Result<ClientKeySecret, AdminError>;
|
||||
async fn create(
|
||||
@@ -64,6 +77,34 @@ impl DefaultClientKeyService {
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyService for DefaultClientKeyService {
|
||||
async fn unresolved_charges(
|
||||
&self,
|
||||
id: &ClientApiKeyId,
|
||||
) -> Result<Vec<crate::model::client_keys::UnresolvedClientCharge>, AdminError> {
|
||||
self.store
|
||||
.unresolved_charges(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "client budget"))
|
||||
}
|
||||
|
||||
async fn reconcile_charge(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: crate::model::client_keys::ReconcileClientCharge,
|
||||
) -> Result<(), AdminError> {
|
||||
if command.request_id.is_empty()
|
||||
|| command.request_id.len() > 128
|
||||
|| command.reason.trim().is_empty()
|
||||
|| command.reason.len() > 1024
|
||||
|| command.reason.chars().any(char::is_control)
|
||||
{
|
||||
return Err(AdminError::invalid("invalid charge reconciliation"));
|
||||
}
|
||||
self.store
|
||||
.reconcile_charge(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "client budget"))
|
||||
}
|
||||
async fn list(&self, query: ClientKeyListQuery) -> Result<ClientKeyPage, AdminError> {
|
||||
validate_cursor(&query)?;
|
||||
self.store
|
||||
@@ -99,6 +140,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,13 @@ 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>>,
|
||||
unresolved_requests: u64,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
@@ -292,6 +321,13 @@ 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),
|
||||
unresolved_requests: record.budget.unresolved_requests,
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
last_used_at: record.last_used_at,
|
||||
@@ -643,6 +679,14 @@ where
|
||||
{
|
||||
Router::new()
|
||||
.route("/api/admin/client-keys", get(list_client_keys::<S>))
|
||||
.route(
|
||||
"/api/admin/client-keys/unresolved-charges",
|
||||
get(unresolved_client_charges::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/client-keys/reconcile-charge",
|
||||
post(reconcile_client_charge::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/client-keys/create",
|
||||
post(create_client_key::<S>),
|
||||
@@ -666,6 +710,66 @@ where
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ReconcileClientChargeRequest {
|
||||
id: String,
|
||||
request_id: String,
|
||||
amount_usd: String,
|
||||
reason: String,
|
||||
}
|
||||
|
||||
async fn unresolved_client_charges<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<ClientKeyIdQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let id = query.into_domain_id().map_err(map_wire_error)?;
|
||||
let items = state.admin_services().client_keys().unresolved_charges(&id).await.map_err(map_service_error)?
|
||||
.into_iter().map(|event| serde_json::json!({
|
||||
"requestId": event.request_id, "startedAt": event.started_at.to_rfc3339(),
|
||||
"completedAt": event.completed_at.map(|time| time.to_rfc3339()), "state": event.state,
|
||||
})).collect::<Vec<_>>();
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"items": items})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn reconcile_client_charge<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(payload): AdminJson<ReconcileClientChargeRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let key_id = client_key_id(payload.id, "id").map_err(map_wire_error)?;
|
||||
let amount_usd = payload
|
||||
.amount_usd
|
||||
.parse::<Decimal>()
|
||||
.map_err(|_| map_wire_error(WireValidationError::new("amountUsd")))?;
|
||||
let command = gateway_admin::model::client_keys::ReconcileClientCharge {
|
||||
key_id,
|
||||
request_id: payload.request_id,
|
||||
amount_usd,
|
||||
reason: payload.reason,
|
||||
};
|
||||
state
|
||||
.admin_services()
|
||||
.client_keys()
|
||||
.reconcile_charge(&auth.context().mutation_context(), command)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"reconciled": true})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_client_keys<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
|
||||
@@ -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,84 @@ 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)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn budget_reconciliation_routes_require_an_admin_session() {
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use tower::ServiceExt as _;
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
let router = client_keys::router::<AdminTestState>().with_state(fixture.state());
|
||||
for (method, path, body) in [
|
||||
(
|
||||
"GET",
|
||||
"/api/admin/client-keys/unresolved-charges?id=key_1",
|
||||
"",
|
||||
),
|
||||
(
|
||||
"POST",
|
||||
"/api/admin/client-keys/reconcile-charge",
|
||||
r#"{"id":"key_1","requestId":"req_charge","amountUsd":"0","reason":"verified"}"#,
|
||||
),
|
||||
] {
|
||||
let response = router
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(method)
|
||||
.uri(path)
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_auth_budget")
|
||||
.body(Body::from(body))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
|
||||
let unknown = serde_json::from_value::<ListClientKeysQuery>(json!({ "other": true }));
|
||||
@@ -207,6 +285,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 +317,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_eq!(list["items"][0]["unresolvedRequests"], 0);
|
||||
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,59 @@
|
||||
//! Durable per-key USD budgets, independent of best-effort request observations.
|
||||
|
||||
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>,
|
||||
pub unresolved_requests: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientBudgetAdmission {
|
||||
pub key_id: ClientApiKeyId,
|
||||
pub request_id: ModelRequestId,
|
||||
pub deadline_at: SystemTime,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientBudgetCharge {
|
||||
pub request_id: ModelRequestId,
|
||||
/// None means the request may have incurred an unknown charge, never zero.
|
||||
pub amount_usd: Option<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 {
|
||||
/// Persist a pending charge before any upstream work; read the current policy atomically.
|
||||
fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>>;
|
||||
|
||||
/// Idempotent by gateway request ID. Failed writes leave the durable pending charge intact.
|
||||
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,9 @@ where
|
||||
request_id,
|
||||
client_api_key_ref,
|
||||
observation: ResponseObservation::new(timing_started_at),
|
||||
budget_prior_attempts_usd: Some(Decimal::ZERO),
|
||||
budget_attempt_send_state: UpstreamSendState::NotSent,
|
||||
budget_attempt_already_counted: false,
|
||||
trace,
|
||||
deadline,
|
||||
deadline_timer: Delay::new(
|
||||
@@ -197,6 +201,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 +258,9 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
request_id: ModelRequestId,
|
||||
client_api_key_ref: crate::policy::ClientApiKeyId,
|
||||
observation: ResponseObservation,
|
||||
budget_prior_attempts_usd: Option<Decimal>,
|
||||
budget_attempt_send_state: UpstreamSendState,
|
||||
budget_attempt_already_counted: bool,
|
||||
trace: TraceContext,
|
||||
deadline: SystemTime,
|
||||
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
|
||||
@@ -288,6 +296,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 +484,34 @@ where
|
||||
self.finalized
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge {
|
||||
let amount_usd = self
|
||||
.budget_prior_attempts_usd
|
||||
.zip(self.budget_attempt_usd())
|
||||
.and_then(|(prior, current)| prior.checked_add(current));
|
||||
super::budget::ClientBudgetCharge {
|
||||
request_id: self.request_id.clone(),
|
||||
amount_usd,
|
||||
completed_at: self.finalized_at.unwrap_or_else(SystemTime::now),
|
||||
}
|
||||
}
|
||||
|
||||
fn budget_attempt_usd(&self) -> Option<Decimal> {
|
||||
if self.budget_attempt_already_counted {
|
||||
return Some(Decimal::ZERO);
|
||||
}
|
||||
self.observation
|
||||
.cost
|
||||
.total()
|
||||
.filter(|money| money.currency().as_str() == "USD")
|
||||
.map(crate::metering::Money::amount)
|
||||
.or_else(|| {
|
||||
(self.budget_attempt_send_state == UpstreamSendState::NotSent)
|
||||
.then_some(Decimal::ZERO)
|
||||
})
|
||||
}
|
||||
|
||||
/// 返回截至当前已完成的实际上游调用结果。
|
||||
///
|
||||
/// 调用方可保存已消费下标;该切片在会话生命周期内只追加、不改写。
|
||||
@@ -749,6 +786,8 @@ where
|
||||
}),
|
||||
);
|
||||
let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone());
|
||||
// Cancellation while a provider is opening cannot prove the request was never sent.
|
||||
self.budget_attempt_send_state = UpstreamSendState::Ambiguous;
|
||||
let stream = match poll_provider(
|
||||
provider,
|
||||
provider_request,
|
||||
@@ -772,6 +811,7 @@ where
|
||||
ProviderBoundary::Result(result) => match *result {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
self.budget_attempt_send_state = error.send_state();
|
||||
record_trace_error(&attempt_trace, &error);
|
||||
if self.prepare_unavailable_native_continuation_replay(&error) {
|
||||
return Ok(Some(PullOutcome::AttemptDiscarded));
|
||||
@@ -800,6 +840,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,
|
||||
@@ -1027,6 +1068,7 @@ where
|
||||
.await;
|
||||
}
|
||||
current.send_observed = true;
|
||||
self.budget_attempt_send_state = UpstreamSendState::Sent;
|
||||
self.send_state_watermark = UpstreamSendState::Sent;
|
||||
}
|
||||
}
|
||||
@@ -1066,6 +1108,7 @@ where
|
||||
error.send_state()
|
||||
};
|
||||
let send_state = self.raise_send_watermark(attempt_send_state);
|
||||
self.budget_attempt_send_state = attempt_send_state;
|
||||
if self.request_persisted {
|
||||
best_effort_store_write(
|
||||
"mark_send_state",
|
||||
@@ -1143,6 +1186,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 +1365,13 @@ where
|
||||
}
|
||||
|
||||
fn reset_uncommitted_observations(&mut self) {
|
||||
// Response facts are attempt-local; the key must also pay for discarded attempts.
|
||||
self.budget_prior_attempts_usd = self
|
||||
.budget_prior_attempts_usd
|
||||
.zip(self.budget_attempt_usd())
|
||||
.and_then(|(prior, current)| prior.checked_add(current));
|
||||
self.budget_attempt_send_state = UpstreamSendState::NotSent;
|
||||
self.budget_attempt_already_counted = false;
|
||||
self.observation.reset_for_attempt();
|
||||
self.upstream_complete = false;
|
||||
}
|
||||
@@ -1326,6 +1381,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 +1561,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::{ClientBudgetAdmission, 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,18 @@ 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(ClientBudgetAdmission {
|
||||
key_id: client.policy.key_id().clone(),
|
||||
request_id: request_id.clone(),
|
||||
deadline_at,
|
||||
})
|
||||
.await
|
||||
{
|
||||
admission.release().await;
|
||||
return Err(error);
|
||||
}
|
||||
let observation = plan
|
||||
.candidates()
|
||||
.first()
|
||||
@@ -552,6 +582,17 @@ impl DefaultExecutionService {
|
||||
{
|
||||
Ok(core) => core,
|
||||
Err(error) => {
|
||||
if let Some(budget) = &self.budget {
|
||||
settle_budget(
|
||||
budget.as_ref(),
|
||||
ClientBudgetCharge {
|
||||
request_id: request_id.clone(),
|
||||
amount_usd: Some(crate::metering::Decimal::ZERO),
|
||||
completed_at: SystemTime::now(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
admission.release().await;
|
||||
return Err(gateway_error_from_engine(&error));
|
||||
}
|
||||
@@ -565,6 +606,7 @@ impl DefaultExecutionService {
|
||||
admission,
|
||||
Arc::clone(&self.circuits),
|
||||
Arc::clone(&self.continuation),
|
||||
self.budget.clone(),
|
||||
)),
|
||||
})
|
||||
}
|
||||
@@ -922,6 +964,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 pending; durable admission requires reconciliation");
|
||||
}
|
||||
}
|
||||
|
||||
impl AdmissionLease {
|
||||
async fn release(self) {
|
||||
if let Err(error) = self
|
||||
@@ -941,6 +989,7 @@ struct DefaultExecutionSession {
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
observed_provider_outcomes: usize,
|
||||
continuation_recorded: bool,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
impl DefaultExecutionSession {
|
||||
@@ -949,6 +998,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 +1007,7 @@ impl DefaultExecutionSession {
|
||||
continuation,
|
||||
observed_provider_outcomes: 0,
|
||||
continuation_recorded: false,
|
||||
budget,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -973,6 +1024,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 +1079,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 +1272,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;
|
||||
|
||||
@@ -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,80 @@ 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,
|
||||
None,
|
||||
"unknown final usage cannot erase a charged retry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_budget_accumulates_retries_and_preserves_unknown_attempts() {
|
||||
for (first_cost, first_send, expected) in [
|
||||
(Some(999), UpstreamSendState::Sent, Some(1122)),
|
||||
(None, UpstreamSendState::Sent, None),
|
||||
(None, UpstreamSendState::NotSent, Some(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.map(|amount| amount.scaled()), expected);
|
||||
assert_eq!(
|
||||
session.budget_charge().completed_at,
|
||||
charge.completed_at,
|
||||
"final settlement time stays stable"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2688,10 +2762,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 +2811,14 @@ 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
|
||||
.map(|amount| amount.scaled()),
|
||||
Some(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);
|
||||
|
||||
@@ -7,6 +7,233 @@ use std::{
|
||||
time::{Duration, Instant, SystemTime},
|
||||
};
|
||||
|
||||
use gateway_core::engine::budget::{
|
||||
ClientBudgetAdmission, 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, _: ClientBudgetAdmission) -> 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 == Some(gateway_core::metering::Decimal::ZERO))
|
||||
);
|
||||
}
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::{executor::block_on, future::BoxFuture};
|
||||
@@ -495,7 +722,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 +1150,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 +1174,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
|
||||
|
||||
@@ -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,412 @@
|
||||
//! Durable admission and idempotent USD settlement, serialized per client key.
|
||||
|
||||
use std::{collections::BTreeMap, sync::Mutex, time::Duration};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::{
|
||||
engine::budget::{
|
||||
ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits,
|
||||
ClientBudgetPort, ClientBudgetStatus,
|
||||
},
|
||||
error::{GatewayError, GatewayErrorKind},
|
||||
metering::Decimal,
|
||||
};
|
||||
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, request: ClientBudgetAdmission) -> Result<(), GatewayError> {
|
||||
// Preserve exact charges across a transient outage. On process loss the durable
|
||||
// pending event still blocks a limited key after its deadline, requiring reconciliation.
|
||||
let retries = self
|
||||
.retry
|
||||
.lock()
|
||||
.map_err(|_| unavailable())?
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if !retries.is_empty() {
|
||||
let ids = retries
|
||||
.iter()
|
||||
.map(|charge| charge.request_id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
let owned = sqlx::query_scalar::<_, String>(
|
||||
"select request_id from client_key_charge_events
|
||||
where client_api_key_id = $1 and request_id = any($2)",
|
||||
)
|
||||
.bind(request.key_id.as_str())
|
||||
.bind(ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
for charge in retries {
|
||||
if owned.iter().any(|id| id == charge.request_id.as_str()) {
|
||||
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(request.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, request.key_id.as_str(), now)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
if limits.is_limited() {
|
||||
let unresolved = sqlx::query_scalar::<_, bool>("select exists(select 1 from client_key_charge_events
|
||||
where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= $2)))")
|
||||
.bind(request.key_id.as_str()).bind(now).fetch_one(&mut *tx).await.map_err(|_| unavailable())?;
|
||||
if unresolved {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::RateLimited,
|
||||
"client API key has unresolved usage; administrator reconciliation is required",
|
||||
)
|
||||
.with_client_code("key_budget_unresolved"));
|
||||
}
|
||||
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(request.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));
|
||||
}
|
||||
}
|
||||
sqlx::query(
|
||||
"insert into client_key_charge_events (request_id, client_api_key_id, deadline_at)
|
||||
values ($1, $2, $3)",
|
||||
)
|
||||
.bind(request.request_id.as_str())
|
||||
.bind(request.key_id.as_str())
|
||||
.bind(DateTime::<Utc>::from(request.deadline_at))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
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)?;
|
||||
// Match admission's key -> window -> event lock order, including concurrent settlements.
|
||||
let key = sqlx::query_scalar::<_, String>(
|
||||
"select k.id from client_api_keys k
|
||||
join client_key_charge_events e on e.client_api_key_id = k.id
|
||||
where e.request_id = $1 for update of k",
|
||||
)
|
||||
.bind(charge.request_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| ClientBudgetError)?;
|
||||
let Some(key) = key else { return Ok(()) }; // Deleting a key also deletes its ledger.
|
||||
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?;
|
||||
let changed = sqlx::query(
|
||||
"update client_key_charge_events
|
||||
set state = $2, amount_usd = $3::text::numeric, completed_at = $4
|
||||
where request_id = $1 and state <> 'settled'",
|
||||
)
|
||||
.bind(charge.request_id.as_str())
|
||||
.bind(if charge.amount_usd.is_some() {
|
||||
"settled"
|
||||
} else {
|
||||
"unknown"
|
||||
})
|
||||
.bind(charge.amount_usd.map(|amount| amount.canonical()))
|
||||
.bind(DateTime::<Utc>::from(charge.completed_at))
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
if changed == 1
|
||||
&& let Some(amount) = charge.amount_usd
|
||||
{
|
||||
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(amount.canonical()).bind(DateTime::<Utc>::from(charge.completed_at))
|
||||
.execute(&mut **tx).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for PgClientBudgetStore {
|
||||
fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async move { self.admit_inner(request).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,
|
||||
(select count(*) from client_key_charge_events e where e.client_api_key_id = k.id
|
||||
and (e.state = 'unknown' or (e.state = 'pending' and e.deadline_at <= now()))) as unresolved
|
||||
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),
|
||||
unresolved_requests: row.get::<i64, _>("unresolved").unsigned_abs(),
|
||||
},
|
||||
);
|
||||
}
|
||||
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")
|
||||
}
|
||||
|
||||
pub(super) async fn unresolved_charges(
|
||||
pool: &PgPool,
|
||||
key: &gateway_core::policy::ClientApiKeyId,
|
||||
) -> gateway_admin::ports::store::AdminStoreResult<
|
||||
Vec<gateway_admin::model::client_keys::UnresolvedClientCharge>,
|
||||
> {
|
||||
let rows = sqlx::query("select request_id, started_at, completed_at, state from client_key_charge_events
|
||||
where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= now()))
|
||||
order by started_at, request_id limit 200")
|
||||
.bind(key.as_str()).fetch_all(pool).await.map_err(|_| budget_admin_unavailable())?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(
|
||||
|row| gateway_admin::model::client_keys::UnresolvedClientCharge {
|
||||
request_id: row.get("request_id"),
|
||||
started_at: row.get("started_at"),
|
||||
completed_at: row.get("completed_at"),
|
||||
state: row.get("state"),
|
||||
},
|
||||
)
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub(super) async fn reconcile_charge(
|
||||
pool: &PgPool,
|
||||
command: gateway_admin::model::client_keys::ReconcileClientCharge,
|
||||
context: &gateway_admin::model::MutationContext,
|
||||
revision: gateway_admin::model::Revision,
|
||||
) -> gateway_admin::ports::store::AdminStoreResult<()> {
|
||||
use gateway_admin::ports::store::{AdminStoreError, AdminStoreErrorKind};
|
||||
let conflict = || {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Conflict,
|
||||
"client budget",
|
||||
"charge is missing, active, or already settled with another amount",
|
||||
)
|
||||
};
|
||||
let mut tx = pool.begin().await.map_err(|_| budget_admin_unavailable())?;
|
||||
sqlx::query("select id from client_api_keys where id = $1 for update")
|
||||
.bind(command.key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| budget_admin_unavailable())?
|
||||
.ok_or_else(conflict)?;
|
||||
let row = sqlx::query("select state, amount_usd::text, coalesce(completed_at, deadline_at) as completed_at, deadline_at <= now() as expired
|
||||
from client_key_charge_events where request_id = $1 and client_api_key_id = $2 for update")
|
||||
.bind(&command.request_id).bind(command.key_id.as_str()).fetch_optional(&mut *tx).await.map_err(|_| budget_admin_unavailable())?.ok_or_else(conflict)?;
|
||||
let state: &str = row.get("state");
|
||||
if state == "settled" {
|
||||
return if row.get::<String, _>("amount_usd").parse::<Decimal>().ok()
|
||||
== Some(command.amount_usd)
|
||||
{
|
||||
Ok(())
|
||||
} else {
|
||||
Err(conflict())
|
||||
};
|
||||
}
|
||||
if state == "pending" && !row.get::<bool, _>("expired") {
|
||||
return Err(conflict());
|
||||
}
|
||||
let charge = ClientBudgetCharge {
|
||||
request_id: gateway_core::engine::ModelRequestId::new(command.request_id.clone())
|
||||
.map_err(|_| conflict())?,
|
||||
amount_usd: Some(command.amount_usd),
|
||||
completed_at: row.get::<DateTime<Utc>, _>("completed_at").into(),
|
||||
};
|
||||
settle_in_transaction(&mut tx, command.key_id.as_str(), &charge)
|
||||
.await
|
||||
.map_err(|_| budget_admin_unavailable())?;
|
||||
sqlx::query("update client_key_charge_events set reconciled_at = now(), reconciliation_reason = $2 where request_id = $1")
|
||||
.bind(&command.request_id).bind(command.reason).execute(&mut *tx).await.map_err(|_| budget_admin_unavailable())?;
|
||||
let audit = crate::mutation_audit(
|
||||
context,
|
||||
"reconcile",
|
||||
"client_key_charge",
|
||||
&command.request_id,
|
||||
vec!["amount_usd".to_owned(), "reconciliation_reason".to_owned()],
|
||||
);
|
||||
super::admin_security_audit::append_admin_audit_event_in_transaction(
|
||||
&mut tx,
|
||||
audit,
|
||||
crate::store_revision(revision)?,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| budget_admin_unavailable())?;
|
||||
tx.commit().await.map_err(|_| budget_admin_unavailable())
|
||||
}
|
||||
|
||||
fn budget_admin_unavailable() -> gateway_admin::ports::store::AdminStoreError {
|
||||
gateway_admin::ports::store::AdminStoreError::new(
|
||||
gateway_admin::ports::store::AdminStoreErrorKind::Unavailable,
|
||||
"client budget",
|
||||
"storage 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(
|
||||
@@ -662,6 +678,22 @@ impl PgAdminClientKeyStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
async fn unresolved_charges(
|
||||
&self,
|
||||
id: &ClientApiKeyId,
|
||||
) -> AdminStoreResult<Vec<gateway_admin::model::client_keys::UnresolvedClientCharge>> {
|
||||
self.required_record(id).await?;
|
||||
super::client_budgets::unresolved_charges(&self.keys.pool, id).await
|
||||
}
|
||||
|
||||
async fn reconcile_charge(
|
||||
&self,
|
||||
command: gateway_admin::model::client_keys::ReconcileClientCharge,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<()> {
|
||||
let revision = self.revision().await?;
|
||||
super::client_budgets::reconcile_charge(&self.keys.pool, command, context, revision).await
|
||||
}
|
||||
async fn list_client_keys(
|
||||
&self,
|
||||
query: AdminClientKeyListQuery,
|
||||
@@ -719,6 +751,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 +768,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 +801,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 +815,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 +981,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 +996,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 +1005,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 +1022,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 +1032,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 +1188,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,383 @@
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext,
|
||||
client_keys::{ReconcileClientCharge, UpdateClientKey},
|
||||
},
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
ModelRequestId,
|
||||
budget::{ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus},
|
||||
},
|
||||
error::GatewayErrorKind,
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
ClientApiKeyRepository as _, PgAdminClientKeyStore, PgClientApiKeyRepository,
|
||||
PgClientBudgetStore,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
fn admission(key: &str, request: &str) -> ClientBudgetAdmission {
|
||||
ClientBudgetAdmission {
|
||||
key_id: ClientApiKeyId::new(key).unwrap(),
|
||||
request_id: ModelRequestId::new(format!("req_{request}")).unwrap(),
|
||||
deadline_at: SystemTime::now() + Duration::from_secs(600),
|
||||
}
|
||||
}
|
||||
|
||||
fn charge(request: &str, amount: Option<&str>) -> ClientBudgetCharge {
|
||||
ClientBudgetCharge {
|
||||
request_id: ModelRequestId::new(format!("req_{request}")).unwrap(),
|
||||
amount_usd: amount.map(|value| value.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: "reconcile-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 index in 0..3 {
|
||||
let id = format!("{prefix}-{index}");
|
||||
// Admitted work can finish above the threshold; no estimated cost reservation.
|
||||
first.admit(admission(key, &id)).await.unwrap();
|
||||
}
|
||||
for index in 0..3 {
|
||||
let id = format!("{prefix}-{index}");
|
||||
let (a, b) = tokio::join!(
|
||||
first.settle(charge(&id, Some(amount))),
|
||||
second.settle(charge(&id, Some(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(admission("day", "d-rejected"))
|
||||
.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(admission("week", "w-rejected"))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
week_error.client_error_code(),
|
||||
Some("key_weekly_budget_exceeded")
|
||||
);
|
||||
let rejected_count: i64 = sqlx::query_scalar(
|
||||
"select count(*) from client_key_charge_events where request_id like '%rejected'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rejected_count, 0);
|
||||
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(admission("key", "first")).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("first", Some("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(admission("key", "second")).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("second", Some("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(admission("key", "after-reset")).await.unwrap();
|
||||
let old = ClientBudgetCharge {
|
||||
completed_at: (day_start - chrono::Duration::seconds(1)).into(),
|
||||
..charge("after-reset", Some("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 unknown_and_orphan_charges_block_limited_keys_until_audited_reconciliation() {
|
||||
let Some(database) = TestDatabase::create("budgets_unknown").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "1", "5").await;
|
||||
seed(&database, "other", "0", "0").await;
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
let admin = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
store.admit(admission("key", "unknown")).await.unwrap();
|
||||
store.settle(charge("unknown", None)).await.unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(admission("key", "blocked"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_budget_unresolved")
|
||||
);
|
||||
let key_id = ClientApiKeyId::new("key").unwrap();
|
||||
assert_eq!(admin.unresolved_charges(&key_id).await.unwrap().len(), 1);
|
||||
assert_eq!(status(&database, "key").await.unresolved_requests, 1);
|
||||
let reconcile = |key: &str, request: &str, amount: &str| ReconcileClientCharge {
|
||||
key_id: ClientApiKeyId::new(key).unwrap(),
|
||||
request_id: format!("req_{request}"),
|
||||
amount_usd: amount.parse().unwrap(),
|
||||
reason: "verified against upstream usage".to_owned(),
|
||||
};
|
||||
assert!(
|
||||
admin
|
||||
.reconcile_charge(reconcile("other", "unknown", "0.4"), &context())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
admin
|
||||
.reconcile_charge(reconcile("key", "unknown", "0.4"), &context())
|
||||
.await
|
||||
.unwrap();
|
||||
admin
|
||||
.reconcile_charge(reconcile("key", "unknown", "0.4"), &context())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
admin
|
||||
.reconcile_charge(reconcile("key", "unknown", "0.5"), &context())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0.4"
|
||||
);
|
||||
let audits: i64 = sqlx::query_scalar(
|
||||
"select count(*) from admin_audit_events where entity_kind = 'client_key_charge'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(audits, 1);
|
||||
store.admit(admission("key", "orphan")).await.unwrap();
|
||||
assert!(
|
||||
admin
|
||||
.reconcile_charge(reconcile("key", "orphan", "0"), &context())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
sqlx::query("update client_key_charge_events set deadline_at = now() - interval '1 second' where request_id = 'req_orphan'")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
let restarted = PgClientBudgetStore::new(database.pool.clone());
|
||||
assert_eq!(
|
||||
restarted
|
||||
.admit(admission("key", "restart-blocked"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_budget_unresolved")
|
||||
);
|
||||
admin
|
||||
.reconcile_charge(reconcile("key", "orphan", "0"), &context())
|
||||
.await
|
||||
.unwrap();
|
||||
restarted
|
||||
.admit(admission("key", "recovered"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(status(&database, "key").await.unresolved_requests, 0);
|
||||
store
|
||||
.admit(admission("other", "unlimited-unknown"))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.settle(charge("unlimited-unknown", None))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.admit(admission("other", "unlimited-next"))
|
||||
.await
|
||||
.unwrap();
|
||||
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(admission("key", "unlimited")).await.unwrap();
|
||||
store
|
||||
.settle(charge("unlimited", Some("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(admission("key", "blocked"))
|
||||
.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(admission("key", "allowed")).await.unwrap();
|
||||
sqlx::query("update client_api_keys set enabled = false where id = 'key'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(admission("key", "disabled"))
|
||||
.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(admission("key", "deleted"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
GatewayErrorKind::Unauthorized
|
||||
);
|
||||
store.settle(charge("allowed", Some("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(admission("key", "offline"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_budget_unavailable")
|
||||
);
|
||||
assert!(store.settle(charge("offline", Some("1"))).await.is_err());
|
||||
database.close().await;
|
||||
}
|
||||
@@ -24,6 +24,7 @@ use super::TestDatabase;
|
||||
#[test]
|
||||
fn client_key_requires_the_frozen_plaintext_format() {
|
||||
let key = NewClientApiKey {
|
||||
budget: Default::default(),
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
@@ -382,6 +383,7 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
|
||||
fn client_key_debug_redacts_plaintext() {
|
||||
let secret = format!("sk_{}", "s".repeat(43));
|
||||
let key = NewClientApiKey {
|
||||
budget: Default::default(),
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
|
||||
@@ -14,6 +14,7 @@ mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod execution_buffer;
|
||||
@@ -230,6 +231,8 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
"backup_settings",
|
||||
"client_api_key_groups",
|
||||
"client_api_keys",
|
||||
"client_key_budget_windows",
|
||||
"client_key_charge_events",
|
||||
"model_requests",
|
||||
"ops_events",
|
||||
"provider_accounts",
|
||||
|
||||
@@ -1124,6 +1124,7 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
let result = store
|
||||
.update_account(
|
||||
UpdateAccount {
|
||||
outbound_proxy: None,
|
||||
account_id: "acct_terminal_mutation".to_owned(),
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
@@ -1184,6 +1185,78 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_proxy_edits_preserve_credentials_and_clear_egress_without_audit_secrets() {
|
||||
let Some(database) = TestDatabase::create("account_proxy").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let mut seed = account("acct_proxy", "proxy-user");
|
||||
seed.outbound_proxy = Some(
|
||||
gateway_core::account::OutboundProxy::parse("http://initial:secret@127.0.0.1:18080")
|
||||
.unwrap(),
|
||||
);
|
||||
repository.insert_provider_account(seed).await.unwrap();
|
||||
let store = admin_account_store(&database.pool);
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "proxy-edit".to_owned(),
|
||||
};
|
||||
let command = UpdateAccount {
|
||||
account_id: "acct_proxy".to_owned(),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::account::AccountWeight::DEFAULT,
|
||||
group_ids: vec![],
|
||||
outbound_proxy: None,
|
||||
};
|
||||
store
|
||||
.update_account(command.clone(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let read = || async {
|
||||
sqlx::query_as::<_, (Option<String>, i64)>("select outbound_proxy_url, credential_revision from provider_accounts where id = 'acct_proxy'").fetch_one(&database.pool).await.unwrap()
|
||||
};
|
||||
assert_eq!(
|
||||
read().await,
|
||||
(Some("http://initial:secret@127.0.0.1:18080/".to_owned()), 1)
|
||||
);
|
||||
store
|
||||
.update_account(
|
||||
UpdateAccount {
|
||||
outbound_proxy: Some(Some(
|
||||
gateway_core::account::OutboundProxy::parse(
|
||||
"socks5h://next:new-secret@127.0.0.1:1080",
|
||||
)
|
||||
.unwrap(),
|
||||
)),
|
||||
..command.clone()
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(read().await.1, 1);
|
||||
store
|
||||
.update_account(
|
||||
UpdateAccount {
|
||||
outbound_proxy: Some(None),
|
||||
..command
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(read().await, (None, 1));
|
||||
let audits: Vec<serde_json::Value> =
|
||||
sqlx::query_scalar("select to_jsonb(a) from admin_audit_events a")
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!serde_json::to_string(&audits).unwrap().contains("secret"));
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_recovery_resets_status_facts_without_changing_credentials_or_scheduling() {
|
||||
const GROUP_ID: &str = "grp_00000000000000000000000000000070";
|
||||
@@ -1342,6 +1415,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever
|
||||
let result = store
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
outbound_proxy: None,
|
||||
account_ids: account_ids.clone(),
|
||||
enabled: false,
|
||||
concurrency_limit: gateway_core::account::AccountConcurrencyLimit::new(7),
|
||||
@@ -1377,6 +1451,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever
|
||||
store
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
outbound_proxy: None,
|
||||
account_ids: account_ids.clone(),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
@@ -1470,6 +1545,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
};
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_admin_upsert", "user-admin-upsert")],
|
||||
audit: audit("audit_admin_upsert_create", "import", "acct_admin_upsert"),
|
||||
@@ -1492,6 +1568,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
updated.credential_state = CredentialState::Banned;
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![updated],
|
||||
audit: audit("audit_admin_upsert_update", "import", "acct_admin_upsert"),
|
||||
@@ -1526,6 +1603,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope,
|
||||
accounts: vec![account("acct_admin_upsert", "another-user")],
|
||||
audit: audit("audit_admin_upsert_rebind", "import", "acct_admin_upsert"),
|
||||
@@ -1563,6 +1641,12 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
|
||||
let result = admin_account_store(&database.pool)
|
||||
.commit_authorization(
|
||||
AuthorizationCommit {
|
||||
settings: Some(gateway_admin::model::accounts::AccountImportSettings {
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::account::AccountWeight::new(9).expect("weight"),
|
||||
group_ids: Vec::new(),
|
||||
}),
|
||||
pending: PendingAuthorizationMutation::new(
|
||||
provider_kind.clone(),
|
||||
AuthorizationMutationTarget::Create {
|
||||
@@ -1571,6 +1655,7 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
|
||||
AuthorizationOwnerBinding::from_context(&context),
|
||||
),
|
||||
credential: AuthorizationCredentialCommit::Create(PreparedCredentialCreate {
|
||||
outbound_proxy: None,
|
||||
account_id: ProviderAccountId::new("acct_authorization_candidate")
|
||||
.expect("candidate account ID"),
|
||||
provider_kind,
|
||||
@@ -1603,6 +1688,8 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
|
||||
),
|
||||
("acct_authorization_existing", Some(2)),
|
||||
);
|
||||
let settings: (bool, Option<i64>, i16) = sqlx::query_as("select enabled, concurrency_limit, weight from provider_accounts where id = 'acct_authorization_existing'").fetch_one(&database.pool).await.expect("OAuth settings");
|
||||
assert_eq!(settings, (false, None, 9));
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -1614,6 +1701,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
repository
|
||||
.insert_provider_account(NewProviderAccount {
|
||||
outbound_proxy: None,
|
||||
id: "acct_core_refresh".to_owned(),
|
||||
provider_kind: "xai".to_owned(),
|
||||
name: "before refresh".to_owned(),
|
||||
@@ -1795,6 +1883,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![
|
||||
account("acct_admin_a", "user-admin-a"),
|
||||
@@ -1822,6 +1911,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![
|
||||
account("acct_admin_transient", "user-admin-transient"),
|
||||
@@ -1947,6 +2037,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
let revision = repository
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
outbound_proxy: None,
|
||||
account_ids: vec!["acct_admin_a".to_owned()],
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
@@ -2009,6 +2100,7 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
|
||||
};
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_grouped_import", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_import", "import", "acct_grouped_import"),
|
||||
@@ -2042,6 +2134,7 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
|
||||
.expect("set account scheduling before reimport");
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_reimport_candidate", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_reimport", "import", "acct_grouped_import"),
|
||||
@@ -2109,6 +2202,7 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
|
||||
};
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_rotation_quota", "user-rotation-quota")],
|
||||
audit: audit(
|
||||
@@ -2250,6 +2344,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
disabled.credential_state = CredentialState::Expired;
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![disabled],
|
||||
audit: audit(
|
||||
@@ -2305,6 +2400,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
|
||||
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
|
||||
NewProviderAccount {
|
||||
outbound_proxy: None,
|
||||
id: id.to_owned(),
|
||||
provider_kind: "openai".to_owned(),
|
||||
name: id.to_owned(),
|
||||
@@ -2473,3 +2569,167 @@ fn imported_ready_account_passes_validation() {
|
||||
let imported = account("acct_ready_import", "user-ready-import");
|
||||
assert!(imported.validate().is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_edit_preserves_an_inflight_token_refresh() {
|
||||
let Some(database) = TestDatabase::create("review_proxy_refresh").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let id = ProviderAccountId::new("acct_review_proxy").unwrap();
|
||||
let mut seed = account(id.as_str(), "review-proxy-user");
|
||||
seed.has_refresh_token = true;
|
||||
repository.insert_provider_account(seed).await.unwrap();
|
||||
let loaded = repository.load_current_credential(&id).await.unwrap();
|
||||
let refreshed = CredentialCasUpdate::new(
|
||||
id.clone(),
|
||||
loaded.account.revision(),
|
||||
ProviderAccountUpdate {
|
||||
account_id: id.clone(),
|
||||
name: loaded.account.name().to_owned(),
|
||||
email: loaded.account.email().map(str::to_owned),
|
||||
plan_type: loaded.account.plan_type().map(str::to_owned),
|
||||
},
|
||||
plaintext_credential("refreshed-marker"),
|
||||
true,
|
||||
Some(SystemTime::now() + Duration::from_secs(3600)),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
admin_account_store(&database.pool)
|
||||
.update_account(
|
||||
UpdateAccount {
|
||||
account_id: id.as_str().to_owned(),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::account::AccountWeight::DEFAULT,
|
||||
group_ids: vec![],
|
||||
outbound_proxy: Some(Some(
|
||||
gateway_core::account::OutboundProxy::parse("http://127.0.0.1:18080").unwrap(),
|
||||
)),
|
||||
},
|
||||
&MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "review-proxy-change".to_owned(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let result = repository
|
||||
.compare_and_swap_credential(refreshed)
|
||||
.await
|
||||
.unwrap();
|
||||
let saved = repository.load_current_credential(&id).await.unwrap();
|
||||
let saved_new_token =
|
||||
saved.credential.expose_to_provider()["access_token"] == "refreshed-marker";
|
||||
database.close().await;
|
||||
assert!(
|
||||
matches!(result, CredentialCasOutcome::Updated(_)),
|
||||
"Changing only egress must preserve a refresh of unchanged OAuth credentials"
|
||||
);
|
||||
assert!(saved_new_token);
|
||||
assert_eq!(
|
||||
saved.account.outbound_proxy().unwrap().expose_url(),
|
||||
"http://127.0.0.1:18080/"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_import_settings_apply_atomically_to_new_and_existing_identities() {
|
||||
use gateway_admin::model::accounts::AccountImportSettings;
|
||||
use gateway_core::account::{AccountConcurrencyLimit, AccountWeight};
|
||||
const GROUP_ID: &str = "grp_00000000000000000000000000000091";
|
||||
let Some(database) = TestDatabase::create("import_settings").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
repository
|
||||
.insert_provider_account(account("acct_existing_settings", "existing-settings-user"))
|
||||
.await
|
||||
.expect("seed account");
|
||||
sqlx::query(
|
||||
"insert into account_groups (id, name, color, created_at, updated_at) values ($1, 'Import settings', '#2563EBFF', now(), now())",
|
||||
)
|
||||
.bind(GROUP_ID)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed group");
|
||||
let settings = AccountImportSettings {
|
||||
enabled: false,
|
||||
concurrency_limit: Some(AccountConcurrencyLimit::new(3).expect("concurrency")),
|
||||
weight: AccountWeight::new(7).expect("weight"),
|
||||
group_ids: vec![AccountGroupId::new(GROUP_ID).expect("group ID")],
|
||||
};
|
||||
let result = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: Some(settings.clone()),
|
||||
scope: ProviderAccountAdminScope {
|
||||
provider_kind: "openai".to_owned(),
|
||||
},
|
||||
accounts: vec![
|
||||
account("acct_new_settings", "new-settings-user"),
|
||||
account("acct_existing_candidate", "existing-settings-user"),
|
||||
account("acct_duplicate_candidate", "existing-settings-user"),
|
||||
],
|
||||
audit: audit("audit_import_settings", "import", "provider_accounts"),
|
||||
})
|
||||
.await
|
||||
.expect("import with settings");
|
||||
assert_eq!(result.config_revision.get(), 2);
|
||||
for id in ["acct_new_settings", "acct_existing_settings"] {
|
||||
let row: (bool, Option<i64>, i16) = sqlx::query_as(
|
||||
"select enabled, concurrency_limit, weight from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("saved settings");
|
||||
assert_eq!(row, (false, Some(3), 7));
|
||||
assert_eq!(account_group_ids(&database.pool, id).await, [GROUP_ID]);
|
||||
}
|
||||
let before = repository
|
||||
.load_provider_account("acct_existing_settings")
|
||||
.await
|
||||
.expect("existing account");
|
||||
let failed = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: Some(AccountImportSettings {
|
||||
group_ids: vec![
|
||||
AccountGroupId::new("grp_00000000000000000000000000000092")
|
||||
.expect("missing group"),
|
||||
],
|
||||
..settings
|
||||
}),
|
||||
scope: ProviderAccountAdminScope {
|
||||
provider_kind: "openai".to_owned(),
|
||||
},
|
||||
accounts: vec![
|
||||
account("acct_rollback_settings", "rollback-settings-user"),
|
||||
account("acct_rollback_candidate", "existing-settings-user"),
|
||||
],
|
||||
audit: audit(
|
||||
"audit_import_settings_failed",
|
||||
"import",
|
||||
"provider_accounts",
|
||||
),
|
||||
})
|
||||
.await;
|
||||
assert!(failed.is_err());
|
||||
assert_eq!(current_revision(&database.pool).await, 2);
|
||||
assert_eq!(
|
||||
account_count(&database.pool, "acct_rollback_settings").await,
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.load_provider_account("acct_existing_settings")
|
||||
.await
|
||||
.expect("unchanged account"),
|
||||
before
|
||||
);
|
||||
assert_eq!(
|
||||
account_group_ids(&database.pool, "acct_existing_settings").await,
|
||||
[GROUP_ID]
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -26,7 +26,9 @@ use thiserror::Error;
|
||||
|
||||
use super::recovery_log::{CodexOAuthRecoveryOperation, record_oauth_recovery};
|
||||
use super::security::CodexCredentialCodec;
|
||||
use super::token_client::{RefreshFailure, TokenRefresher};
|
||||
use super::token_client::{
|
||||
OpenAiTokenClient, PersonalAccessTokenError, RefreshFailure, TokenRefresher,
|
||||
};
|
||||
use super::types::{
|
||||
CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile, CodexCredentialData,
|
||||
CodexCredentialPrincipal, CodexOAuthMetadata, CodexOAuthSecret, parse_access_token_expiration,
|
||||
@@ -48,8 +50,8 @@ pub struct ImportCodexOAuthCredential {
|
||||
|
||||
/// OAuth credential 的最小创建输入。
|
||||
///
|
||||
/// metadata 来自 ID token/access token 的官方本地 payload 解析;不包含签名、
|
||||
/// issuer、audience、nonce 或 token 验证。
|
||||
/// OAuth metadata 来自 ID token/access token 的本地 payload 解析;PAT metadata
|
||||
/// 来自 auth whoami 验证。两者均不从导入文档信任身份字段。
|
||||
pub(crate) struct UnresolvedCodexOAuthCredential {
|
||||
pub(crate) account_id: String,
|
||||
pub(crate) name: String,
|
||||
@@ -122,6 +124,7 @@ struct ParsedCodexImportAccount {
|
||||
name: Option<String>,
|
||||
email: Option<String>,
|
||||
authentication: ParsedCodexAuthentication,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
struct ParsedCodexAuthentication {
|
||||
@@ -226,6 +229,8 @@ struct CodexCprExportCommon {
|
||||
status: &'static str,
|
||||
added_at: String,
|
||||
updated_at: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
outbound_proxy_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
@@ -321,6 +326,8 @@ impl fmt::Debug for PreparedCodexCredentialRotationGuard {
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum CodexCredentialAdminError {
|
||||
#[error(transparent)]
|
||||
PersonalAccessToken(#[from] PersonalAccessTokenError),
|
||||
#[error("Codex account input is invalid")]
|
||||
InvalidInput,
|
||||
#[error("Codex credential JSON is invalid")]
|
||||
@@ -348,7 +355,8 @@ impl CodexCredentialAdminError {
|
||||
Self::RefreshRejected { message }
|
||||
| Self::AccountBanned { message }
|
||||
| Self::RefreshAmbiguous { message } => message.as_deref(),
|
||||
Self::InvalidInput
|
||||
Self::PersonalAccessToken(_)
|
||||
| Self::InvalidInput
|
||||
| Self::InvalidCredential
|
||||
| Self::NotFound
|
||||
| Self::MissingRefreshToken
|
||||
@@ -414,7 +422,7 @@ impl CodexCredentialAdmin {
|
||||
|
||||
/// 为已取得 OAuth access token、但资料尚未补全的账号创建最小记录。
|
||||
///
|
||||
/// ID token payload 只用于本地字段投影;创建路径绝不发起 usage/profile 请求。
|
||||
/// 接受调用方已取得的身份投影;本方法不发起 usage/profile 请求。
|
||||
pub(crate) fn prepare_unresolved_oauth(
|
||||
&self,
|
||||
input: UnresolvedCodexOAuthCredential,
|
||||
@@ -502,6 +510,9 @@ impl CodexCredentialAdmin {
|
||||
status: cpr_status(&account),
|
||||
added_at: china_rfc3339(item.added_at),
|
||||
updated_at: china_rfc3339(item.updated_at),
|
||||
outbound_proxy_url: account
|
||||
.outbound_proxy()
|
||||
.map(|proxy| proxy.expose_url().to_owned()),
|
||||
};
|
||||
let exported = match data {
|
||||
CodexCredentialData::OAuth(data) => {
|
||||
@@ -662,6 +673,7 @@ impl CodexCredentialAdmin {
|
||||
/// 有状态的 Codex 手工刷新边界;消费调用方刚读取的当前 credential 并准备 CAS。
|
||||
pub struct CodexCredentialAdminService {
|
||||
refresher: Arc<dyn TokenRefresher>,
|
||||
personal_access_token_client: Option<Arc<OpenAiTokenClient>>,
|
||||
leases: Arc<dyn ProviderLeasePort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
}
|
||||
@@ -671,6 +683,10 @@ impl fmt::Debug for CodexCredentialAdminService {
|
||||
formatter
|
||||
.debug_struct("CodexCredentialAdminService")
|
||||
.field("refresher", &"TokenRefresher")
|
||||
.field(
|
||||
"personal_access_token_client",
|
||||
&self.personal_access_token_client.is_some(),
|
||||
)
|
||||
.field("leases", &"ProviderLeasePort")
|
||||
.field("runtime_policy", &"ProviderRuntimePolicyPort")
|
||||
.finish()
|
||||
@@ -685,11 +701,19 @@ impl CodexCredentialAdminService {
|
||||
) -> Self {
|
||||
Self {
|
||||
refresher,
|
||||
personal_access_token_client: None,
|
||||
leases,
|
||||
runtime_policy,
|
||||
}
|
||||
}
|
||||
|
||||
/// 在生产组装时复用 OAuth auth client,为 at- 导入启用上游身份验证。
|
||||
#[must_use]
|
||||
pub fn with_personal_access_token_client(mut self, client: Arc<OpenAiTokenClient>) -> Self {
|
||||
self.personal_access_token_client = Some(client);
|
||||
self
|
||||
}
|
||||
|
||||
/// 官方 RT exchange;结果由 App 在同一 revision/audit 事务中提交。
|
||||
pub async fn manual_refresh(
|
||||
&self,
|
||||
@@ -743,7 +767,10 @@ impl CodexCredentialAdminService {
|
||||
};
|
||||
let tokens = self
|
||||
.refresher
|
||||
.refresh(refresh_token.expose_secret())
|
||||
.refresh_with_proxy(
|
||||
refresh_token.expose_secret(),
|
||||
current.account.outbound_proxy(),
|
||||
)
|
||||
.await
|
||||
.inspect_err(|error| log_manual_refresh_failure(&account_id, error))
|
||||
.map_err(map_refresh_failure)?;
|
||||
@@ -793,6 +820,7 @@ impl CodexCredentialAdminService {
|
||||
///
|
||||
/// OAuth 导入先取得 access token(直接提供或 RT exchange),再按官方
|
||||
/// `parse_chatgpt_jwt_claims` 从 ID token/access token 本地投影账号资料。
|
||||
/// at- PAT 使用 whoami 取得身份,并丢弃不适用的 RT、ID token 和刷新计划。
|
||||
pub async fn prepare_import_document(
|
||||
&self,
|
||||
payload: Value,
|
||||
@@ -813,34 +841,56 @@ impl CodexCredentialAdminService {
|
||||
let account_id = format!("acct_{}", uuid::Uuid::now_v7().simple());
|
||||
let typed_account_id = ProviderAccountId::new(account_id.clone())
|
||||
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
|
||||
let (secret, access_token_expires_at) = self
|
||||
let (mut secret, mut access_token_expires_at) = self
|
||||
.resolve_import_tokens(
|
||||
&typed_account_id,
|
||||
candidate.authentication.access_token.clone(),
|
||||
candidate.authentication.refresh_token.clone(),
|
||||
candidate.authentication.id_token.clone(),
|
||||
candidate.outbound_proxy.as_ref(),
|
||||
)
|
||||
.await?;
|
||||
// 与官方 `parse_chatgpt_jwt_claims` 一致:先读 ID token,缺失字段
|
||||
// 再由 access token 补齐。解析失败不阻塞 token 入库。
|
||||
let id_metadata = secret
|
||||
.id_token
|
||||
.as_ref()
|
||||
.and_then(|token| parse_chatgpt_jwt_claims(token.expose_secret()).ok())
|
||||
.unwrap_or_default();
|
||||
let access_metadata =
|
||||
parse_chatgpt_jwt_claims(secret.access_token.expose_secret()).unwrap_or_default();
|
||||
let metadata = CodexOAuthMetadata {
|
||||
email: id_metadata.email.or(access_metadata.email),
|
||||
chatgpt_plan_type: id_metadata
|
||||
.chatgpt_plan_type
|
||||
.or(access_metadata.chatgpt_plan_type),
|
||||
chatgpt_user_id: id_metadata
|
||||
.chatgpt_user_id
|
||||
.or(access_metadata.chatgpt_user_id),
|
||||
chatgpt_account_id: id_metadata
|
||||
.chatgpt_account_id
|
||||
.or(access_metadata.chatgpt_account_id),
|
||||
let metadata = if secret
|
||||
.access_token
|
||||
.expose_secret()
|
||||
.trim()
|
||||
.starts_with("at-")
|
||||
{
|
||||
let token = secret.access_token.expose_secret().trim();
|
||||
let metadata = self
|
||||
.personal_access_token_client
|
||||
.as_ref()
|
||||
.ok_or(PersonalAccessTokenError::Unavailable)?
|
||||
.personal_access_token_metadata(token, candidate.outbound_proxy.as_ref())
|
||||
.await?;
|
||||
secret = CodexOAuthSecret {
|
||||
access_token: SecretString::from(token),
|
||||
refresh_token: None,
|
||||
id_token: None,
|
||||
};
|
||||
access_token_expires_at = None;
|
||||
metadata
|
||||
} else {
|
||||
// 普通 OAuth 仍按 ID token、access token 的顺序投影,不调用 whoami。
|
||||
let id_metadata = secret
|
||||
.id_token
|
||||
.as_ref()
|
||||
.and_then(|token| parse_chatgpt_jwt_claims(token.expose_secret()).ok())
|
||||
.unwrap_or_default();
|
||||
let access_metadata = parse_chatgpt_jwt_claims(secret.access_token.expose_secret())
|
||||
.unwrap_or_default();
|
||||
CodexOAuthMetadata {
|
||||
email: id_metadata.email.or(access_metadata.email),
|
||||
chatgpt_plan_type: id_metadata
|
||||
.chatgpt_plan_type
|
||||
.or(access_metadata.chatgpt_plan_type),
|
||||
chatgpt_user_id: id_metadata
|
||||
.chatgpt_user_id
|
||||
.or(access_metadata.chatgpt_user_id),
|
||||
chatgpt_account_id: id_metadata
|
||||
.chatgpt_account_id
|
||||
.or(access_metadata.chatgpt_account_id),
|
||||
}
|
||||
};
|
||||
let prepared =
|
||||
CodexCredentialAdmin.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
|
||||
@@ -859,7 +909,12 @@ impl CodexCredentialAdminService {
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})?;
|
||||
accounts.push(prepared);
|
||||
accounts.push(NewProviderAccount {
|
||||
account: prepared
|
||||
.account
|
||||
.with_outbound_proxy(candidate.outbound_proxy),
|
||||
credential: prepared.credential,
|
||||
});
|
||||
}
|
||||
Ok(PreparedCodexAccountImport { accounts })
|
||||
}
|
||||
@@ -870,6 +925,7 @@ impl CodexCredentialAdminService {
|
||||
access_token: Option<String>,
|
||||
refresh_token: Option<String>,
|
||||
id_token: Option<String>,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<(CodexOAuthSecret, Option<DateTime<Utc>>), CodexCredentialAdminError> {
|
||||
let id_token = id_token.map(SecretString::from);
|
||||
if let Some(access_token) = access_token {
|
||||
@@ -889,7 +945,7 @@ impl CodexCredentialAdminService {
|
||||
let refresh_token = refresh_token.ok_or(CodexCredentialAdminError::InvalidCredential)?;
|
||||
let tokens = self
|
||||
.refresher
|
||||
.refresh(&refresh_token)
|
||||
.refresh_with_proxy(&refresh_token, proxy)
|
||||
.await
|
||||
.map_err(map_refresh_failure)?;
|
||||
let access_token = tokens
|
||||
@@ -986,12 +1042,18 @@ fn log_manual_refresh_failure(account_id: &ProviderAccountId, error: &RefreshFai
|
||||
fn parse_import_document(
|
||||
payload: &Value,
|
||||
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
|
||||
let payload = payload
|
||||
.get("data")
|
||||
.filter(|data| data.get("accounts").is_some())
|
||||
.unwrap_or(payload);
|
||||
let mut accounts = Vec::new();
|
||||
for value in import_account_values(payload)? {
|
||||
if !is_openai_oauth_candidate(value) {
|
||||
continue;
|
||||
}
|
||||
accounts.push(parse_oauth_import_account(value)?);
|
||||
let mut account = parse_oauth_import_account(value)?;
|
||||
account.outbound_proxy = import_proxy(payload, value)?;
|
||||
accounts.push(account);
|
||||
}
|
||||
Ok(accounts)
|
||||
}
|
||||
@@ -1004,9 +1066,89 @@ fn parse_oauth_import_account(
|
||||
name: first_string(value, &["name", "label"]),
|
||||
email: first_string(value, &["email"]).or_else(|| first_string(credentials, &["email"])),
|
||||
authentication: parse_oauth_import_tokens(value)?,
|
||||
outbound_proxy: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn import_proxy(
|
||||
payload: &Value,
|
||||
account: &Value,
|
||||
) -> Result<Option<gateway_core::account::OutboundProxy>, CodexCredentialAdminError> {
|
||||
let invalid = || CodexCredentialAdminError::InvalidInput;
|
||||
if let Some(value) = account
|
||||
.get("outbound_proxy_url")
|
||||
.or_else(|| account.get("outboundProxyUrl"))
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
let text = value.as_str().ok_or_else(invalid)?;
|
||||
if text.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
return gateway_core::account::OutboundProxy::parse(text)
|
||||
.map(Some)
|
||||
.map_err(|_| invalid());
|
||||
}
|
||||
let Some(key) = account.get("proxy_key").filter(|value| !value.is_null()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let key = key
|
||||
.as_str()
|
||||
.filter(|key| !key.is_empty())
|
||||
.ok_or_else(invalid)?;
|
||||
let proxies = payload
|
||||
.get("proxies")
|
||||
.and_then(Value::as_array)
|
||||
.ok_or_else(invalid)?;
|
||||
let mut matches = proxies
|
||||
.iter()
|
||||
.filter(|proxy| proxy.get("proxy_key").and_then(Value::as_str) == Some(key));
|
||||
let proxy = matches.next().ok_or_else(invalid)?;
|
||||
if matches.next().is_some()
|
||||
|| proxy
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value != "active")
|
||||
|| proxy
|
||||
.get("expires_at")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
|| proxy
|
||||
.get("fallback_mode")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| !value.is_empty() && value != "none")
|
||||
{
|
||||
return Err(invalid());
|
||||
}
|
||||
let text = |field| proxy.get(field).and_then(Value::as_str).ok_or_else(invalid);
|
||||
let scheme = text("protocol")?;
|
||||
// sub2api resolves SOCKS targets at the proxy.
|
||||
let scheme = if scheme == "socks5" {
|
||||
"socks5h"
|
||||
} else {
|
||||
scheme
|
||||
};
|
||||
let mut url = url::Url::parse(&format!("{scheme}://localhost")).map_err(|_| invalid())?;
|
||||
let host = text("host")?;
|
||||
if let Ok(address) = host.parse::<std::net::IpAddr>() {
|
||||
url.set_ip_host(address).map_err(|_| invalid())?;
|
||||
} else {
|
||||
url.set_host(Some(host)).map_err(|_| invalid())?;
|
||||
}
|
||||
let port = proxy
|
||||
.get("port")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|port| u16::try_from(port).ok())
|
||||
.filter(|port| *port != 0)
|
||||
.ok_or_else(invalid)?;
|
||||
url.set_port(Some(port)).map_err(|_| invalid())?;
|
||||
url.set_username(proxy.get("username").and_then(Value::as_str).unwrap_or(""))
|
||||
.map_err(|_| invalid())?;
|
||||
url.set_password(proxy.get("password").and_then(Value::as_str))
|
||||
.map_err(|_| invalid())?;
|
||||
gateway_core::account::OutboundProxy::parse(url.as_str())
|
||||
.map(Some)
|
||||
.map_err(|_| invalid())
|
||||
}
|
||||
|
||||
fn import_account_values(payload: &Value) -> Result<Vec<&Value>, CodexCredentialAdminError> {
|
||||
if let Some(accounts) = payload.get("accounts") {
|
||||
return accounts
|
||||
@@ -1032,7 +1174,10 @@ fn parse_oauth_import_tokens(
|
||||
Value::Object(object) => {
|
||||
for (key, value) in object {
|
||||
let token = match key.as_str() {
|
||||
"accessToken" | "access_token" => &mut access_token,
|
||||
"accessToken"
|
||||
| "access_token"
|
||||
| "personal_access_token"
|
||||
| "personalAccessToken" => &mut access_token,
|
||||
"refreshToken" | "refresh_token" => &mut refresh_token,
|
||||
"idToken" | "id_token" => &mut id_token,
|
||||
_ => {
|
||||
|
||||
@@ -436,6 +436,10 @@ impl CodexCredentialCatalogService {
|
||||
.authorization_header()
|
||||
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
|
||||
let result = client
|
||||
.for_account(account)
|
||||
.map_err(|error| CodexCredentialCatalogError::Upstream {
|
||||
detail: error.to_string(),
|
||||
})?
|
||||
.fetch_models_with_context(CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
account.upstream_account_id(),
|
||||
|
||||
@@ -524,7 +524,16 @@ impl CodexOAuthAdminService {
|
||||
.transpose()?
|
||||
.flatten();
|
||||
let (mut secret, id_token) = self
|
||||
.exchange_pending(&pending, callback_url, fallback_refresh_token)
|
||||
.exchange_pending(
|
||||
&pending,
|
||||
callback_url,
|
||||
fallback_refresh_token,
|
||||
current
|
||||
.as_ref()
|
||||
.map_or(pending.mutation.outbound_proxy(), |current| {
|
||||
current.account.outbound_proxy()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let mutation = pending.mutation.clone();
|
||||
let access_token_expires_at =
|
||||
@@ -545,20 +554,25 @@ impl CodexOAuthAdminService {
|
||||
)
|
||||
} else {
|
||||
let account_id = format!("acct_{}", Uuid::now_v7().simple());
|
||||
CompletedCodexOAuthCredential::Create(
|
||||
self.credentials
|
||||
.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
|
||||
account_id,
|
||||
name: pending.name,
|
||||
installation_id: pending.installation_id,
|
||||
secret,
|
||||
metadata,
|
||||
access_token_expires_at,
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.map_err(map_admin_error)?,
|
||||
)
|
||||
let prepared = self
|
||||
.credentials
|
||||
.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
|
||||
account_id,
|
||||
name: pending.name,
|
||||
installation_id: pending.installation_id,
|
||||
secret,
|
||||
metadata,
|
||||
access_token_expires_at,
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.map_err(map_admin_error)?;
|
||||
CompletedCodexOAuthCredential::Create(gateway_core::account::NewProviderAccount {
|
||||
account: prepared
|
||||
.account
|
||||
.with_outbound_proxy(mutation.outbound_proxy().cloned()),
|
||||
credential: prepared.credential,
|
||||
})
|
||||
};
|
||||
Ok((mutation, credential))
|
||||
}
|
||||
@@ -666,6 +680,7 @@ impl CodexOAuthAdminService {
|
||||
pending: &CodexPendingAuthorization,
|
||||
callback_url: &str,
|
||||
fallback_refresh_token: Option<SecretString>,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
|
||||
let (code, callback_state) = callback_parts(callback_url)?;
|
||||
if !constant_time_equal(
|
||||
@@ -676,10 +691,13 @@ impl CodexOAuthAdminService {
|
||||
}
|
||||
let tokens = self
|
||||
.exchanger
|
||||
.exchange_authorization_code(AuthorizationCodeGrant {
|
||||
code,
|
||||
code_verifier: pending.code_verifier.clone(),
|
||||
})
|
||||
.exchange_with_proxy(
|
||||
AuthorizationCodeGrant {
|
||||
code,
|
||||
code_verifier: pending.code_verifier.clone(),
|
||||
},
|
||||
proxy,
|
||||
)
|
||||
.await
|
||||
.map_err(map_exchange_error)?;
|
||||
let mut secret = tokens.secret;
|
||||
|
||||
@@ -133,13 +133,16 @@ impl CodexCredentialProfileService {
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CodexProfileStatistics, CodexProfileStatisticsError> {
|
||||
let (authorization, upstream_account_id) = self.account_authentication(account_id).await?;
|
||||
let (authorization, upstream_account_id, account) =
|
||||
self.account_authentication(account_id).await?;
|
||||
let request_id = format!("profile_statistics_{}", Uuid::now_v7().simple());
|
||||
let statistics = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
)
|
||||
.for_account(&account)
|
||||
.map_err(map_client_error)?
|
||||
.fetch_profile_statistics(CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
upstream_account_id.as_deref(),
|
||||
@@ -156,7 +159,14 @@ impl CodexCredentialProfileService {
|
||||
async fn account_authentication(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<(SecretString, Option<String>), CodexProfileStatisticsError> {
|
||||
) -> Result<
|
||||
(
|
||||
SecretString,
|
||||
Option<String>,
|
||||
gateway_core::account::ProviderAccount,
|
||||
),
|
||||
CodexProfileStatisticsError,
|
||||
> {
|
||||
let account = self
|
||||
.repository
|
||||
.store()
|
||||
@@ -187,6 +197,7 @@ impl CodexCredentialProfileService {
|
||||
Ok((
|
||||
authorization,
|
||||
account.upstream_account_id().map(str::to_owned),
|
||||
account,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -202,9 +213,19 @@ impl CodexCredentialProfileService {
|
||||
.image_url
|
||||
.ok_or(CodexProfileAvatarError::Missing)?,
|
||||
};
|
||||
let (authorization, upstream_account_id) = self.account_authentication(account_id).await?;
|
||||
let (authorization, upstream_account_id, account) =
|
||||
self.account_authentication(account_id).await?;
|
||||
let http = if account.outbound_proxy().is_some() {
|
||||
crate::transport::client::build_account_http_client(
|
||||
account.id().as_str(),
|
||||
account.outbound_proxy(),
|
||||
)
|
||||
.map_err(|_| CodexProfileAvatarError::TransportUnavailable)?
|
||||
} else {
|
||||
self.http.clone()
|
||||
};
|
||||
fetch_profile_avatar(
|
||||
&self.http,
|
||||
&http,
|
||||
&self.base_url,
|
||||
&self.profile.snapshot(),
|
||||
&source,
|
||||
|
||||
@@ -1225,6 +1225,8 @@ async fn list_reset_credits_once(
|
||||
.authorization_header()
|
||||
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
|
||||
client
|
||||
.for_account(&prepared.account)
|
||||
.map_err(ResetCreditAttemptError::Upstream)?
|
||||
.list_rate_limit_reset_credits(CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
@@ -1248,6 +1250,8 @@ async fn consume_reset_credit_once(
|
||||
.authorization_header()
|
||||
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
|
||||
client
|
||||
.for_account(&prepared.account)
|
||||
.map_err(ResetCreditAttemptError::Upstream)?
|
||||
.consume_rate_limit_reset_credit(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
@@ -1329,6 +1333,8 @@ async fn fetch_usage_once(
|
||||
.map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?;
|
||||
let request_id = format!("quota_{}", Uuid::now_v7().simple());
|
||||
client
|
||||
.for_account(&prepared.account)
|
||||
.map_err(CodexQuotaFetchAttemptError::Upstream)?
|
||||
.fetch_usage(CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
|
||||
@@ -113,6 +113,7 @@ pub struct CodexAccountQuotaSnapshot {
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
observed_at: SystemTime,
|
||||
plan_type: Option<String>,
|
||||
fact: CodexQuotaFact,
|
||||
quota: QuotaState,
|
||||
windows: Vec<CodexQuotaWindow>,
|
||||
@@ -135,6 +136,12 @@ impl CodexAccountQuotaSnapshot {
|
||||
self.observed_at
|
||||
}
|
||||
|
||||
/// 上游额度响应明确返回的套餐;缺失时不推断账号级别。
|
||||
#[must_use]
|
||||
pub fn plan_type(&self) -> Option<&str> {
|
||||
self.plan_type.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn fact(&self) -> CodexQuotaFact {
|
||||
self.fact
|
||||
@@ -339,6 +346,10 @@ pub(crate) fn parse_account_quota_snapshot(
|
||||
account_id,
|
||||
credential_revision,
|
||||
observed_at,
|
||||
plan_type: object
|
||||
.get("plan_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned),
|
||||
fact,
|
||||
quota,
|
||||
windows,
|
||||
|
||||
@@ -288,7 +288,10 @@ impl CodexCredentialRefreshService {
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?;
|
||||
let refresh_result = self.refresher.refresh(refresh_token.expose_secret()).await;
|
||||
let refresh_result = self
|
||||
.refresher
|
||||
.refresh_with_proxy(refresh_token.expose_secret(), due.account.outbound_proxy())
|
||||
.await;
|
||||
if recovery_window_exhausted && let Err(failure) = &refresh_result {
|
||||
let message = failure.message().map(str::to_owned);
|
||||
let upstream = failure.upstream();
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
//! OpenAI token 续期 Reqwest 适配器。
|
||||
//! OpenAI OAuth token exchange 与 Codex PAT 验证的 Reqwest 适配器。
|
||||
|
||||
use super::types::CodexOAuthMetadata;
|
||||
|
||||
use crate::transport::{
|
||||
headers::build_codex_profile_headers,
|
||||
@@ -22,6 +24,33 @@ pub const OFFICIAL_CODEX_OAUTH_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
|
||||
pub const OFFICIAL_CODEX_TOKEN_ENDPOINT: &str = "https://auth.openai.com/oauth/token";
|
||||
/// Codex Desktop loopback callback;管理员复制完整回调 URL 交回固定 complete API。
|
||||
pub const OFFICIAL_CODEX_REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
|
||||
const PERSONAL_ACCESS_TOKEN_WHOAMI_PATH: &str = "/api/accounts/v1/user-auth-credential/whoami";
|
||||
|
||||
/// PAT 验证失败;不保留令牌、响应体或可能包含秘密的底层 HTTP 错误。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum PersonalAccessTokenError {
|
||||
#[error("Codex PAT must be a non-empty at- token without whitespace or control characters")]
|
||||
InvalidToken,
|
||||
#[error(
|
||||
"Codex PAT was rejected by OpenAI; it may be invalid, expired, revoked, or lack permission"
|
||||
)]
|
||||
Rejected,
|
||||
#[error("Codex PAT validation is unavailable; retry later")]
|
||||
Unavailable,
|
||||
#[error("OpenAI returned an invalid Codex PAT identity response")]
|
||||
InvalidResponse,
|
||||
}
|
||||
|
||||
// 与 Codex personal_access_token.rs 一致:email 可缺失,其余身份字段必填。
|
||||
#[derive(Deserialize)]
|
||||
struct PersonalAccessTokenResponse {
|
||||
email: Option<String>,
|
||||
chatgpt_user_id: String,
|
||||
chatgpt_account_id: String,
|
||||
chatgpt_plan_type: String,
|
||||
#[serde(rename = "chatgpt_account_is_fedramp")]
|
||||
_chatgpt_account_is_fedramp: bool,
|
||||
}
|
||||
|
||||
/// Token 刷新成功后得到的认证材料。
|
||||
#[derive(Clone)]
|
||||
@@ -178,6 +207,22 @@ impl fmt::Debug for RefreshUpstreamFailure {
|
||||
#[async_trait]
|
||||
pub trait TokenRefresher: Send + Sync + 'static {
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
|
||||
async fn refresh_with_proxy(
|
||||
&self,
|
||||
refresh_token: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<TokenPair, RefreshFailure> {
|
||||
if proxy.is_some() {
|
||||
return Err(proxy_refresh_failure());
|
||||
}
|
||||
self.refresh(refresh_token).await
|
||||
}
|
||||
}
|
||||
|
||||
fn proxy_refresh_failure() -> RefreshFailure {
|
||||
RefreshFailure::RetryableTransport {
|
||||
message: "account OAuth egress unavailable".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Authorization Code + PKCE 的一次性 grant。
|
||||
@@ -232,6 +277,16 @@ pub trait AuthorizationCodeExchanger: Send + Sync + 'static {
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError>;
|
||||
async fn exchange_with_proxy(
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
|
||||
if proxy.is_some() {
|
||||
return Err(AuthorizationCodeExchangeError::Unavailable);
|
||||
}
|
||||
self.exchange_authorization_code(grant).await
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI token 续期客户端配置。
|
||||
@@ -257,6 +312,24 @@ pub struct OpenAiTokenClient {
|
||||
pub struct TokenClientBuildError;
|
||||
|
||||
impl OpenAiTokenClient {
|
||||
fn with_proxy(
|
||||
&self,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<Self, TokenClientBuildError> {
|
||||
let Some(proxy) = proxy else {
|
||||
return Ok(self.clone());
|
||||
};
|
||||
let builder = Client::builder()
|
||||
.use_rustls_tls()
|
||||
.no_proxy()
|
||||
.proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| TokenClientBuildError)?)
|
||||
.redirect(Policy::none())
|
||||
.connect_timeout(TOKEN_CONNECT_TIMEOUT)
|
||||
.timeout(TOKEN_REQUEST_TIMEOUT);
|
||||
let client =
|
||||
build_reqwest_client_with_custom_ca(builder).map_err(|_| TokenClientBuildError)?;
|
||||
Ok(Self::new(client, self.config.clone(), self.profile.clone()))
|
||||
}
|
||||
/// 共享运行时画像;刷新时取快照,授权码交换仍使用 raw auth 请求。
|
||||
pub fn new(client: Client, config: TokenClientConfig, profile: CodexWireProfileState) -> Self {
|
||||
Self {
|
||||
@@ -265,6 +338,73 @@ impl OpenAiTokenClient {
|
||||
profile,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn personal_access_token_metadata(
|
||||
&self,
|
||||
access_token: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<CodexOAuthMetadata, PersonalAccessTokenError> {
|
||||
if !access_token.starts_with("at-")
|
||||
|| access_token.len() <= 3
|
||||
|| access_token.len() > MAX_OAUTH_RESPONSE_BYTES
|
||||
|| access_token
|
||||
.chars()
|
||||
.any(|ch| ch.is_whitespace() || ch.is_control())
|
||||
{
|
||||
return Err(PersonalAccessTokenError::InvalidToken);
|
||||
}
|
||||
// 复用固定 auth origin 及其 TLS/超时/禁止重定向策略;导入文档不能指定验证地址。
|
||||
let mut endpoint = reqwest::Url::parse(&self.config.token_endpoint)
|
||||
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
|
||||
endpoint.set_path(PERSONAL_ACCESS_TOKEN_WHOAMI_PATH);
|
||||
endpoint.set_query(None);
|
||||
endpoint.set_fragment(None);
|
||||
let headers = build_codex_profile_headers(&self.profile.snapshot())
|
||||
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
|
||||
let client = self
|
||||
.with_proxy(proxy)
|
||||
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
|
||||
let response = client
|
||||
.client
|
||||
.get(endpoint)
|
||||
.headers(headers)
|
||||
.bearer_auth(access_token)
|
||||
.header(reqwest::header::ACCEPT, "application/json")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| PersonalAccessTokenError::Unavailable)?;
|
||||
match response.status() {
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
|
||||
return Err(PersonalAccessTokenError::Rejected);
|
||||
}
|
||||
status if !status.is_success() => return Err(PersonalAccessTokenError::Unavailable),
|
||||
_ => {}
|
||||
}
|
||||
let (_, body) = read_bounded_response(response)
|
||||
.await
|
||||
.map_err(|_| PersonalAccessTokenError::InvalidResponse)?;
|
||||
let identity: PersonalAccessTokenResponse =
|
||||
serde_json::from_slice(&body).map_err(|_| PersonalAccessTokenError::InvalidResponse)?;
|
||||
if [
|
||||
&identity.chatgpt_user_id,
|
||||
&identity.chatgpt_account_id,
|
||||
&identity.chatgpt_plan_type,
|
||||
]
|
||||
.iter()
|
||||
.any(|value| value.trim().is_empty() || value.chars().any(char::is_control))
|
||||
{
|
||||
return Err(PersonalAccessTokenError::InvalidResponse);
|
||||
}
|
||||
Ok(CodexOAuthMetadata {
|
||||
email: identity
|
||||
.email
|
||||
.map(|value| value.trim().to_owned())
|
||||
.filter(|value| !value.is_empty()),
|
||||
chatgpt_user_id: Some(identity.chatgpt_user_id.trim().to_owned()),
|
||||
chatgpt_account_id: Some(identity.chatgpt_account_id.trim().to_owned()),
|
||||
chatgpt_plan_type: Some(identity.chatgpt_plan_type.trim().to_owned()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建禁止 redirect 且无自动重试的 Codex token client。
|
||||
@@ -357,6 +497,16 @@ struct AuthorizationCodeResponse {
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for OpenAiTokenClient {
|
||||
async fn refresh_with_proxy(
|
||||
&self,
|
||||
refresh_token: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<TokenPair, RefreshFailure> {
|
||||
self.with_proxy(proxy)
|
||||
.map_err(|_| proxy_refresh_failure())?
|
||||
.refresh(refresh_token)
|
||||
.await
|
||||
}
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
let headers = build_codex_profile_headers(&self.profile.snapshot()).map_err(|_| {
|
||||
RefreshFailure::Transport {
|
||||
@@ -389,6 +539,16 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
|
||||
#[async_trait]
|
||||
impl AuthorizationCodeExchanger for OpenAiTokenClient {
|
||||
async fn exchange_with_proxy(
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
|
||||
self.with_proxy(proxy)
|
||||
.map_err(|_| AuthorizationCodeExchangeError::Unavailable)?
|
||||
.exchange_authorization_code(grant)
|
||||
.await
|
||||
}
|
||||
async fn exchange_authorization_code(
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
|
||||
@@ -153,12 +153,15 @@ pub async fn initialize(
|
||||
.map_err(|_| OpenAiInitializeError::TokenClient)?,
|
||||
);
|
||||
let refresher: Arc<dyn TokenRefresher> = token_client.clone();
|
||||
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client;
|
||||
let credential_admin = Arc::new(CodexCredentialAdminService::new(
|
||||
Arc::clone(&refresher),
|
||||
Arc::clone(&leases),
|
||||
Arc::clone(&runtime_policy),
|
||||
));
|
||||
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client.clone();
|
||||
let credential_admin = Arc::new(
|
||||
CodexCredentialAdminService::new(
|
||||
Arc::clone(&refresher),
|
||||
Arc::clone(&leases),
|
||||
Arc::clone(&runtime_policy),
|
||||
)
|
||||
.with_personal_access_token_client(token_client),
|
||||
);
|
||||
let refresh = Arc::new(CodexCredentialRefreshService::new(
|
||||
repository,
|
||||
refresher,
|
||||
|
||||
@@ -124,7 +124,9 @@ impl CodexProvider {
|
||||
crate::transport::request::scope_turn_metadata(metadata, lease.installation_id(), true)
|
||||
});
|
||||
let events = cold_json_response_stream(ColdJsonResponse {
|
||||
client: self.client.clone(),
|
||||
client: self.client.for_account(lease.account()).map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
})?,
|
||||
response_origin: request.response_origin,
|
||||
endpoint_path: request.endpoint_path,
|
||||
body: request.body,
|
||||
|
||||
@@ -519,7 +519,9 @@ impl Provider for CodexProvider {
|
||||
AttemptTransport::Default | AttemptTransport::Fallback => 0,
|
||||
};
|
||||
let events = cold_response_stream(ColdResponse {
|
||||
client: self.client.clone(),
|
||||
client: self.client.for_account(lease.account()).map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
})?,
|
||||
response_origin: self.responses_url.clone(),
|
||||
request: upstream_request,
|
||||
upstream_model: upstream_model.clone(),
|
||||
|
||||
@@ -45,13 +45,20 @@ const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
pub(super) const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
|
||||
"x-codex-ws-stream-request-start-ms";
|
||||
type ReqwestClientCacheKey = Option<String>;
|
||||
type ReqwestClientCacheKey = (Option<String>, String);
|
||||
type ReqwestClientCache = Mutex<HashMap<ReqwestClientCacheKey, Client>>;
|
||||
|
||||
/// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。
|
||||
pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
|
||||
build_account_http_client("", None)
|
||||
}
|
||||
|
||||
pub fn build_account_http_client(
|
||||
account_id: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<Client, CustomCaError> {
|
||||
super::tls::ensure_rustls_provider();
|
||||
let cache_key = custom_ca_env_cache_key();
|
||||
let cache_key = (custom_ca_env_cache_key(), egress_key(account_id, proxy));
|
||||
static CLIENTS: OnceLock<ReqwestClientCache> = OnceLock::new();
|
||||
let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new()));
|
||||
if let Some(client) = cache
|
||||
@@ -62,7 +69,7 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
|
||||
return Ok(client.clone());
|
||||
}
|
||||
|
||||
let builder = Client::builder()
|
||||
let mut builder = Client::builder()
|
||||
.use_rustls_tls()
|
||||
.no_proxy()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
@@ -73,13 +80,35 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
|
||||
.http2_keep_alive_interval(Duration::from_secs(30))
|
||||
.http2_keep_alive_timeout(Duration::from_secs(5))
|
||||
.http2_keep_alive_while_idle(true);
|
||||
if let Some(proxy) = proxy {
|
||||
builder = builder.proxy(
|
||||
reqwest::Proxy::all(proxy.expose_url())
|
||||
.map_err(|_| CustomCaError::ProxyConfiguration)?,
|
||||
);
|
||||
}
|
||||
let client = build_reqwest_client_with_custom_ca(builder)?;
|
||||
let mut clients = cache
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if clients.len() >= 256 {
|
||||
clients.clear();
|
||||
}
|
||||
Ok(clients.entry(cache_key).or_insert(client).clone())
|
||||
}
|
||||
|
||||
fn egress_key(account_id: &str, proxy: Option<&gateway_core::account::OutboundProxy>) -> String {
|
||||
use sha2::{Digest, Sha256};
|
||||
let mut hash = Sha256::new();
|
||||
hash.update(account_id.as_bytes());
|
||||
hash.update([0]);
|
||||
hash.update(
|
||||
proxy
|
||||
.map_or("direct", |proxy| proxy.expose_url())
|
||||
.as_bytes(),
|
||||
);
|
||||
hex::encode(hash.finalize())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 错误类型
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -599,11 +628,36 @@ pub struct CodexBackendJsonResponse {
|
||||
#[derive(Clone)]
|
||||
pub struct CodexBackendClient {
|
||||
pub(super) client: Client,
|
||||
pub(super) direct_client: Client,
|
||||
pub(super) base_url: String,
|
||||
pub(super) profile: CodexWireProfileState,
|
||||
pub(super) websocket_pool: Option<Arc<CodexWebSocketPool>>,
|
||||
pub(super) websocket_origin_breaker: WebSocketOriginBreaker,
|
||||
pub(super) websocket_origin_key: String,
|
||||
pub(super) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub(super) egress_key: String,
|
||||
}
|
||||
|
||||
impl CodexBackendClient {
|
||||
pub fn for_account(
|
||||
&self,
|
||||
account: &gateway_core::account::ProviderAccount,
|
||||
) -> Result<Self, CodexClientError> {
|
||||
let mut client = self.clone();
|
||||
client.outbound_proxy = account.outbound_proxy().cloned();
|
||||
client.egress_key = egress_key(account.id().as_str(), account.outbound_proxy());
|
||||
client.websocket_origin_key = format!(
|
||||
"{}:{}",
|
||||
websocket_origin_key(&self.base_url),
|
||||
client.egress_key
|
||||
);
|
||||
client.client = if account.outbound_proxy().is_some() {
|
||||
build_account_http_client(account.id().as_str(), account.outbound_proxy())?
|
||||
} else {
|
||||
self.direct_client.clone()
|
||||
};
|
||||
Ok(client)
|
||||
}
|
||||
}
|
||||
|
||||
/// 已完成账号级 opening 准备、但尚未发送 payload 的 transport。
|
||||
|
||||
@@ -52,8 +52,11 @@ impl CodexBackendClient {
|
||||
) -> Self {
|
||||
let base_url = base_url.into().trim_end_matches('/').to_string();
|
||||
Self {
|
||||
direct_client: client.clone(),
|
||||
client,
|
||||
websocket_origin_key: websocket_origin_key(&base_url),
|
||||
outbound_proxy: None,
|
||||
egress_key: String::new(),
|
||||
base_url,
|
||||
profile,
|
||||
websocket_pool: None,
|
||||
@@ -234,13 +237,14 @@ impl CodexBackendClient {
|
||||
|
||||
let websocket_request = websocket_upstream_request(request);
|
||||
let headers = self.request_headers_for_websocket_response(&websocket_request, context)?;
|
||||
let websocket_create = CodexWebSocketConnection::responses_create_request(
|
||||
let mut websocket_create = CodexWebSocketConnection::responses_create_request(
|
||||
&self.base_url,
|
||||
&generate_key(),
|
||||
websocket_header_pairs(&headers),
|
||||
&websocket_request,
|
||||
)
|
||||
.map_err(CodexClientError::WebSocketEncode)?;
|
||||
websocket_create.connection.outbound_proxy = self.outbound_proxy.clone();
|
||||
context.trace.cloned().unwrap_or_default().headers(
|
||||
"upstream.request.headers",
|
||||
serde_json::json!({"transport": "websocket", "phase": "prepared_opening"}),
|
||||
@@ -469,6 +473,7 @@ impl CodexBackendClient {
|
||||
.as_deref()
|
||||
.or(request.previous_response_id())?;
|
||||
let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id)
|
||||
.with_egress_key(&self.egress_key)
|
||||
.with_connection_profile(connection_profile);
|
||||
if let Some(connection_id) = request.downstream_websocket_connection_id.as_deref() {
|
||||
key = key.with_downstream_connection_id(connection_id);
|
||||
|
||||
@@ -33,6 +33,8 @@ pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE";
|
||||
/// 自定义 CA 错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum CustomCaError {
|
||||
#[error("invalid account proxy configuration")]
|
||||
ProxyConfiguration,
|
||||
/// 读取 CA 证书文件失败。
|
||||
#[error(
|
||||
"Failed to read CA certificate file {} selected by {}: {source}. {hint}",
|
||||
@@ -121,6 +123,17 @@ pub fn maybe_build_rustls_client_config_with_custom_ca() -> CustomCaResult<Optio
|
||||
maybe_build_rustls_client_config_with_env(&ProcessEnv)
|
||||
}
|
||||
|
||||
pub(crate) fn account_proxy_tls_config() -> CustomCaResult<Arc<ClientConfig>> {
|
||||
if let Some(config) = maybe_build_rustls_client_config_with_custom_ca()? {
|
||||
return Ok(config);
|
||||
}
|
||||
Ok(Arc::new(
|
||||
ClientConfig::builder()
|
||||
.with_root_certificates(native_root_store().map_err(CustomCaError::LoadNativeRoots)?)
|
||||
.with_no_client_auth(),
|
||||
))
|
||||
}
|
||||
|
||||
fn build_reqwest_client_with_env(
|
||||
env_source: &dyn EnvSource,
|
||||
mut builder: reqwest::ClientBuilder,
|
||||
|
||||
@@ -5,7 +5,9 @@ use std::time::Duration;
|
||||
use bytes::Bytes;
|
||||
use gateway_protocol::openai::events;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::{Connector, connect_async_tls_with_config};
|
||||
use tokio_tungstenite::{
|
||||
Connector, MaybeTlsStream, client_async_tls_with_config, connect_async_tls_with_config,
|
||||
};
|
||||
use tungstenite::{
|
||||
self, Message,
|
||||
extensions::{ExtensionsConfig, compression::deflate::DeflateConfig},
|
||||
@@ -60,7 +62,11 @@ impl CodexWebSocketConnection {
|
||||
"sec-websocket-extensions".to_string(),
|
||||
WEBSOCKET_EXTENSIONS.to_string(),
|
||||
));
|
||||
Self { endpoint, headers }
|
||||
Self {
|
||||
endpoint,
|
||||
headers,
|
||||
outbound_proxy: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。
|
||||
@@ -140,10 +146,39 @@ async fn connect_websocket(
|
||||
)))
|
||||
})?
|
||||
.map(Connector::Rustls);
|
||||
let result = timeout(
|
||||
WEBSOCKET_CONNECT_TIMEOUT,
|
||||
connect_async_tls_with_config(request, Some(websocket_config()), false, connector),
|
||||
)
|
||||
let result = timeout(WEBSOCKET_CONNECT_TIMEOUT, async {
|
||||
// Preserve the native direct handshake; explicit egress never inherits a global proxy.
|
||||
if connection.outbound_proxy.is_none()
|
||||
&& matches!(
|
||||
tungstenite::proxy::ProxyConfig::from_env(request.uri()),
|
||||
Ok(None)
|
||||
)
|
||||
{
|
||||
let (websocket, response) =
|
||||
connect_async_tls_with_config(request, Some(websocket_config()), false, connector)
|
||||
.await?;
|
||||
return Ok((Box::new(websocket) as RawWsStream, response));
|
||||
}
|
||||
let stream = dial_account(connection).await?;
|
||||
match stream {
|
||||
MaybeTlsStream::Plain(tcp) => {
|
||||
let (websocket, response) =
|
||||
client_async_tls_with_config(request, tcp, Some(websocket_config()), connector)
|
||||
.await?;
|
||||
Ok((Box::new(websocket) as RawWsStream, response))
|
||||
}
|
||||
stream => {
|
||||
let (websocket, response) = client_async_tls_with_config(
|
||||
request,
|
||||
stream,
|
||||
Some(websocket_config()),
|
||||
connector,
|
||||
)
|
||||
.await?;
|
||||
Ok((Box::new(websocket) as RawWsStream, response))
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::ConnectTimeout {
|
||||
timeout: WEBSOCKET_CONNECT_TIMEOUT,
|
||||
@@ -155,6 +190,89 @@ async fn connect_websocket(
|
||||
}
|
||||
}
|
||||
|
||||
async fn dial_account(
|
||||
connection: &CodexWebSocketConnection,
|
||||
) -> Result<MaybeTlsStream<tokio::net::TcpStream>, tungstenite::Error> {
|
||||
use tokio::net::lookup_host;
|
||||
use tungstenite::proxy::ProxyConfig;
|
||||
let invalid =
|
||||
|| tungstenite::Error::Io(std::io::Error::other("account WebSocket egress failed"));
|
||||
let endpoint = url::Url::parse(connection.endpoint()).map_err(|_| invalid())?;
|
||||
let host = endpoint.host_str().ok_or_else(invalid)?;
|
||||
let dns_host = match endpoint.host() {
|
||||
Some(url::Host::Ipv6(ip)) => ip.to_string(),
|
||||
_ => host.to_owned(),
|
||||
};
|
||||
let port = endpoint.port_or_known_default().ok_or_else(invalid)?;
|
||||
let Some(proxy) = connection.outbound_proxy.as_ref() else {
|
||||
return Ok(MaybeTlsStream::Plain(connect_tcp(host, port).await?));
|
||||
};
|
||||
let mut proxy_url = url::Url::parse(proxy.expose_url()).map_err(|_| invalid())?;
|
||||
let tls_proxy = proxy_url.scheme() == "https";
|
||||
let proxy_port = proxy_url.port_or_known_default().ok_or_else(invalid)?;
|
||||
if tls_proxy {
|
||||
proxy_url.set_scheme("http").map_err(|_| invalid())?;
|
||||
proxy_url
|
||||
.set_port(Some(proxy_port))
|
||||
.map_err(|_| invalid())?;
|
||||
}
|
||||
let config = ProxyConfig::parse(proxy_url.as_str()).map_err(|_| invalid())?;
|
||||
let proxy_host = match proxy_url.host() {
|
||||
Some(url::Host::Ipv6(ip)) => ip.to_string(),
|
||||
_ => config.host.clone(),
|
||||
};
|
||||
let tcp = connect_tcp(&proxy_host, config.port).await?;
|
||||
let stream = if tls_proxy {
|
||||
let tls = tls::account_proxy_tls_config().map_err(|_| invalid())?;
|
||||
let name = rustls_pki_types::ServerName::try_from(proxy_host).map_err(|_| invalid())?;
|
||||
MaybeTlsStream::Rustls(
|
||||
tokio_rustls::TlsConnector::from(tls)
|
||||
.connect(name, tcp)
|
||||
.await?,
|
||||
)
|
||||
} else {
|
||||
MaybeTlsStream::Plain(tcp)
|
||||
};
|
||||
// The pinned tungstenite fork sends domains remotely for both SOCKS schemes.
|
||||
let target = if proxy_url.scheme() == "socks5" {
|
||||
lookup_host((dns_host.as_str(), port))
|
||||
.await?
|
||||
.next()
|
||||
.ok_or_else(invalid)?
|
||||
.ip()
|
||||
.to_string()
|
||||
} else if proxy_url.scheme() == "socks5h" {
|
||||
dns_host
|
||||
} else {
|
||||
host.to_owned()
|
||||
};
|
||||
tokio_tungstenite::proxy::connect_via_proxy(stream, &config, &target, port)
|
||||
.await
|
||||
.map_err(|_| invalid())
|
||||
}
|
||||
|
||||
async fn connect_tcp(host: &str, port: u16) -> Result<tokio::net::TcpStream, tungstenite::Error> {
|
||||
use hyper_util::client::legacy::connect::HttpConnector;
|
||||
use tower_service::Service;
|
||||
|
||||
// Reuse Hyper's DNS and Happy Eyeballs connector for explicit direct/proxy endpoints.
|
||||
let authority = host.parse::<std::net::IpAddr>().map_or_else(
|
||||
|_| format!("{host}:{port}"),
|
||||
|ip| std::net::SocketAddr::new(ip, port).to_string(),
|
||||
);
|
||||
let uri = format!("http://{authority}")
|
||||
.parse::<hyper::Uri>()
|
||||
.map_err(|_| tungstenite::Error::Io(std::io::Error::other("invalid egress endpoint")))?;
|
||||
let mut connector = HttpConnector::new();
|
||||
connector
|
||||
.call(uri)
|
||||
.await
|
||||
.map(|stream| stream.into_inner())
|
||||
.map_err(|_| {
|
||||
tungstenite::Error::Io(std::io::Error::other("account WebSocket connection failed"))
|
||||
})
|
||||
}
|
||||
|
||||
fn websocket_handshake_request(
|
||||
connection: &CodexWebSocketConnection,
|
||||
) -> Result<WsRequest, tungstenite::http::Error> {
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::transport::protocol::responses::{
|
||||
/// WebSocket opening 描述。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWebSocketConnection {
|
||||
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub(super) endpoint: String,
|
||||
pub(super) headers: Vec<(String, String)>,
|
||||
}
|
||||
@@ -16,7 +17,7 @@ pub struct CodexWebSocketConnection {
|
||||
/// 已构造完成的 Responses WebSocket 请求描述。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWebSocketRequest {
|
||||
pub(super) connection: CodexWebSocketConnection,
|
||||
pub(crate) connection: CodexWebSocketConnection,
|
||||
pub(super) payload_text: String,
|
||||
pub(super) continuation: WebSocketContinuationRequirement,
|
||||
}
|
||||
@@ -113,6 +114,7 @@ impl CodexWebSocketConnection {
|
||||
Self {
|
||||
endpoint: endpoint.into(),
|
||||
headers,
|
||||
outbound_proxy: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -24,6 +24,7 @@ pub struct CodexWebSocketPoolKey {
|
||||
conversation_id: String,
|
||||
connection_profile: String,
|
||||
downstream_connection_id: String,
|
||||
egress_key: String,
|
||||
}
|
||||
|
||||
impl CodexWebSocketPoolKey {
|
||||
@@ -39,6 +40,7 @@ impl CodexWebSocketPoolKey {
|
||||
conversation_id: conversation_id.into(),
|
||||
connection_profile: String::new(),
|
||||
downstream_connection_id: String::new(),
|
||||
egress_key: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,6 +63,11 @@ impl CodexWebSocketPoolKey {
|
||||
&self.account_id
|
||||
}
|
||||
|
||||
pub(crate) fn with_egress_key(mut self, key: &str) -> Self {
|
||||
self.egress_key = key.to_owned();
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn conversation_id_hash(&self) -> String {
|
||||
short_sha256([self.conversation_id.as_str()])
|
||||
}
|
||||
@@ -72,6 +79,7 @@ impl CodexWebSocketPoolKey {
|
||||
self.conversation_id.as_str(),
|
||||
self.connection_profile.as_str(),
|
||||
self.downstream_connection_id.as_str(),
|
||||
self.egress_key.as_str(),
|
||||
])
|
||||
}
|
||||
|
||||
@@ -79,6 +87,7 @@ impl CodexWebSocketPoolKey {
|
||||
self.base_url == other.base_url
|
||||
&& self.account_id == other.account_id
|
||||
&& self.conversation_id == other.conversation_id
|
||||
&& self.egress_key == other.egress_key
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,20 +19,35 @@ use std::sync::{
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use futures::{Sink, SinkExt, Stream, StreamExt};
|
||||
use tokio::{
|
||||
net::TcpStream,
|
||||
sync::{mpsc, oneshot},
|
||||
task::JoinHandle,
|
||||
time::{Instant, MissedTickBehavior},
|
||||
};
|
||||
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::CodexWebSocketCloseError;
|
||||
|
||||
/// 底层 tungstenite WebSocket 流。
|
||||
pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
pub(crate) trait WebSocketIo:
|
||||
Stream<Item = Result<Message, tungstenite::Error>>
|
||||
+ Sink<Message, Error = tungstenite::Error>
|
||||
+ Send
|
||||
+ Unpin
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> WebSocketIo for T where
|
||||
T: Stream<Item = Result<Message, tungstenite::Error>>
|
||||
+ Sink<Message, Error = tungstenite::Error>
|
||||
+ Send
|
||||
+ Unpin
|
||||
{
|
||||
}
|
||||
|
||||
pub(crate) type RawWsStream = Box<dyn WebSocketIo>;
|
||||
|
||||
const PUMP_COMMAND_BUFFER: usize = 32;
|
||||
|
||||
|
||||
@@ -443,6 +443,7 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
|
||||
let error = bundle
|
||||
.admin_provider()
|
||||
.complete_authorization(CompleteAuthorization {
|
||||
settings: None,
|
||||
context: MutationContext {
|
||||
actor: MutationActor::AdminSession {
|
||||
admin_user_id: "different-owner".to_owned(),
|
||||
@@ -647,6 +648,71 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export
|
||||
assert!(exported_account.get("token").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_admin_projects_free_plan_from_cached_quota_when_account_claims_omit_it() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let mut verified_account = profile("chatgpt-free-plan");
|
||||
verified_account.plan_type = None;
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: "acct_free_plan".to_owned(),
|
||||
name: "free plan".to_owned(),
|
||||
secret: secret("free-plan-test-token"),
|
||||
verified_account,
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
let account = store.account("acct_free_plan").expect("stored account");
|
||||
let observed_at = SystemTime::now();
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: OpaqueProviderData::new(
|
||||
json!({
|
||||
"plan_type": "free",
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"primary_window": {
|
||||
"used_percent": 0,
|
||||
"limit_window_seconds": 2_592_000,
|
||||
"reset_at": 1_900_000_000
|
||||
}
|
||||
}
|
||||
})
|
||||
.as_object()
|
||||
.expect("quota object")
|
||||
.clone(),
|
||||
),
|
||||
observed_at,
|
||||
state: QuotaState::allowed(observed_at),
|
||||
})
|
||||
.await
|
||||
.expect("persist existing quota");
|
||||
let config = valid_config();
|
||||
let bundle = provider_openai::initialize(
|
||||
config.config.clone(),
|
||||
provider_ports_with(store, Arc::new(TestOAuthPending::default())),
|
||||
)
|
||||
.await
|
||||
.expect("OpenAI bundle");
|
||||
let admin = bundle.admin_provider();
|
||||
let quota = admin
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: account.id().clone(),
|
||||
refresh: false,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
.expect("read cached free quota");
|
||||
assert_eq!(quota.plan_type.as_deref(), Some("free"));
|
||||
assert_eq!(
|
||||
admin.plan_type_display(quota.plan_type.as_deref().expect("plan")),
|
||||
"Free"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_admin_provider_projects_official_codex_quota_and_independent_buckets_with_chinese_labels()
|
||||
{
|
||||
@@ -950,6 +1016,31 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
|
||||
assert_eq!(refresh_error.kind(), ProviderAdminErrorKind::NotFound);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialized_provider_reports_a_safe_pat_format_error_before_network_access() {
|
||||
let config = valid_config();
|
||||
let bundle = provider_openai::initialize(config.config.clone(), provider_ports())
|
||||
.await
|
||||
.expect("OpenAI bundle");
|
||||
let error = bundle
|
||||
.admin_provider()
|
||||
.prepare_import(PrepareCredentialImport {
|
||||
document: ProviderDocument::new(OpaqueProviderData::new(Map::from_iter([(
|
||||
"accessToken".to_owned(),
|
||||
json!("at-sensitive-token with whitespace"),
|
||||
)]))),
|
||||
})
|
||||
.await
|
||||
.expect_err("PAT format must be checked by the initialized provider");
|
||||
assert_eq!(error.kind(), ProviderAdminErrorKind::Invalid);
|
||||
assert_eq!(
|
||||
error.public_message(),
|
||||
Some("Codex PAT 格式无效:应为 at- 开头的完整令牌,不能包含空白或控制字符")
|
||||
);
|
||||
assert!(error.message().is_none());
|
||||
assert!(!format!("{error:?}").contains("sensitive-token"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_rotation_preserves_the_new_access_token_jwt_expiration() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1137,6 +1228,7 @@ fn provider_ports_with_catalog(
|
||||
fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
let now = Utc::now();
|
||||
AccountRecord {
|
||||
outbound_proxy: None,
|
||||
id: account.id().to_string(),
|
||||
provider_kind: account.provider().clone(),
|
||||
groups: Vec::new(),
|
||||
|
||||
@@ -1,16 +1,94 @@
|
||||
use std::sync::Arc;
|
||||
use std::{
|
||||
num::NonZeroU32,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{TimeZone as _, Utc};
|
||||
use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use provider_openai::credential::token_client::{
|
||||
OpenAiTokenClient, PersonalAccessTokenError, RefreshFailure, TokenClientConfig, TokenPair,
|
||||
TokenRefresher, openai_token_client,
|
||||
};
|
||||
use provider_openai::credential::{CodexCredentialAdminService, CodexCredentialCodec};
|
||||
use secrecy::ExposeSecret as _;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{header, method, path},
|
||||
};
|
||||
|
||||
use crate::support::{TestLeaseCoordinator, runtime_policy};
|
||||
|
||||
struct UnusedRefresher;
|
||||
|
||||
#[tokio::test]
|
||||
async fn sub2api_import_resolves_distinct_proxy_bindings_and_encodes_credentials() {
|
||||
let service = CodexCredentialAdminService::new(
|
||||
Arc::new(UnusedRefresher),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
let prepared = service.prepare_import_document(serde_json::json!({"data": {
|
||||
"accounts": [
|
||||
{"name": "a", "platform": "openai", "type": "oauth", "proxy_key": "a", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-a"}}))}},
|
||||
{"name": "b", "platform": "openai", "type": "oauth", "proxy_key": "b", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-b"}}))}},
|
||||
{"name": "direct", "platform": "openai", "type": "oauth", "credentials": {"access_token": test_jwt(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-c"}}))}}
|
||||
],
|
||||
"proxies": [
|
||||
{"proxy_key": "a", "protocol": "http", "host": "127.0.0.1", "port": 18080, "username": "user@a", "password": "p:a/ss", "status": "active"},
|
||||
{"proxy_key": "b", "protocol": "socks5", "host": "::1", "port": 1080, "status": "active", "fallback_mode": "none"}
|
||||
]
|
||||
}})).await.unwrap();
|
||||
assert_eq!(prepared.accounts().len(), 3);
|
||||
let first = prepared.accounts()[0].account.outbound_proxy().unwrap();
|
||||
assert_eq!(first.endpoint(), "http://127.0.0.1:18080/");
|
||||
assert!(first.expose_url().contains("user%40a:p%3Aa%2Fss@"));
|
||||
assert_eq!(
|
||||
prepared.accounts()[1]
|
||||
.account
|
||||
.outbound_proxy()
|
||||
.unwrap()
|
||||
.endpoint(),
|
||||
"socks5h://[::1]:1080"
|
||||
);
|
||||
assert!(prepared.accounts()[2].account.outbound_proxy().is_none());
|
||||
assert!(!format!("{prepared:?}").contains("p:a/ss"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sub2api_invalid_proxy_bindings_are_rejected_before_any_token_refresh() {
|
||||
let service = CodexCredentialAdminService::new(
|
||||
Arc::new(UnusedRefresher),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
let proxy = serde_json::json!({"proxy_key": "bound", "protocol": "http", "host": "127.0.0.1", "port": 18080, "status": "active"});
|
||||
let mut cases = vec![
|
||||
serde_json::json!([]),
|
||||
serde_json::json!([proxy.clone(), proxy.clone()]),
|
||||
];
|
||||
for (field, value) in [
|
||||
("status", serde_json::json!("inactive")),
|
||||
("expires_at", serde_json::json!(2000000000)),
|
||||
("fallback_mode", serde_json::json!("direct")),
|
||||
("port", serde_json::json!(0)),
|
||||
] {
|
||||
let mut invalid = proxy.clone();
|
||||
invalid[field] = value;
|
||||
cases.push(serde_json::json!([invalid]));
|
||||
}
|
||||
for proxies in cases {
|
||||
let result = service.prepare_import_document(serde_json::json!({
|
||||
"accounts": [
|
||||
{"platform": "openai", "type": "oauth", "credentials": {"refresh_token": "must-not-refresh"}},
|
||||
{"platform": "openai", "type": "oauth", "proxy_key": "bound", "credentials": {"access_token": "at"}}
|
||||
], "proxies": proxies
|
||||
})).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for UnusedRefresher {
|
||||
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
@@ -145,6 +223,390 @@ fn test_jwt(payload: serde_json::Value) -> String {
|
||||
)
|
||||
}
|
||||
|
||||
fn pat_service(server: &MockServer) -> CodexCredentialAdminService {
|
||||
let client = openai_token_client(
|
||||
TokenClientConfig {
|
||||
client_id: "test-public-client".to_owned(),
|
||||
token_endpoint: format!("{}/oauth/token", server.uri()),
|
||||
},
|
||||
provider_openai::OpenAiConfig::default().wire_profile_state(),
|
||||
)
|
||||
.expect("auth client");
|
||||
CodexCredentialAdminService::new(
|
||||
Arc::new(UnusedRefresher),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
runtime_policy(),
|
||||
)
|
||||
.with_personal_access_token_client(Arc::new(client))
|
||||
}
|
||||
|
||||
fn pat_identity() -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"email": "pat@example.com",
|
||||
"chatgpt_user_id": "pat-user",
|
||||
"chatgpt_account_id": "pat-workspace",
|
||||
"chatgpt_plan_type": "team",
|
||||
"chatgpt_account_is_fedramp": false
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_verifies_identity_and_becomes_schedulable_without_oauth_refresh_fields() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
|
||||
.and(header("authorization", "Bearer at-test-token"))
|
||||
.and(header("accept", "application/json"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(pat_identity()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = pat_service(&server);
|
||||
let prepared = service
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accounts": [{
|
||||
"platform": "openai", "type": "codex", "name": "team PAT",
|
||||
"access_token": " at-test-token ",
|
||||
"refresh_token": "unrelated-refresh-token",
|
||||
"id_token": test_jwt(serde_json::json!({
|
||||
"https://api.openai.com/auth": {"chatgpt_user_id": "untrusted-jwt-user"}
|
||||
})),
|
||||
"account_id": "untrusted-document-account",
|
||||
"email": "untrusted-document@example.com",
|
||||
"planType": "untrusted-plan", "expires_at": "2000-01-01T00:00:00Z"
|
||||
}]
|
||||
}))
|
||||
.await
|
||||
.expect("PAT import without a document user ID");
|
||||
let prepared = &prepared.accounts()[0];
|
||||
let account = &prepared.account;
|
||||
assert_eq!(account.upstream_user_id(), Some("pat-user"));
|
||||
assert_eq!(account.upstream_account_id(), Some("pat-workspace"));
|
||||
assert_eq!(account.email(), Some("pat@example.com"));
|
||||
assert_eq!(account.plan_type(), Some("team"));
|
||||
assert_eq!(account.name(), "team PAT");
|
||||
assert_eq!(
|
||||
account.credential_state(),
|
||||
gateway_core::account::CredentialState::Ready
|
||||
);
|
||||
assert_eq!(
|
||||
account.status_projection(SystemTime::now(), None).status,
|
||||
gateway_core::account::AccountStatus::Normal
|
||||
);
|
||||
assert!(!account.has_refresh_token());
|
||||
assert!(account.access_token_expires_at().is_none());
|
||||
assert!(account.next_refresh_at().is_none());
|
||||
let runtime = CodexCredentialCodec::decode(&prepared.credential).expect("stored PAT");
|
||||
let secret = runtime
|
||||
.authentication
|
||||
.oauth()
|
||||
.expect("shared Bearer transport");
|
||||
assert_eq!(secret.access_token.expose_secret(), "at-test-token");
|
||||
assert!(secret.refresh_token.is_none());
|
||||
assert!(secret.id_token.is_none());
|
||||
let now = SystemTime::now();
|
||||
let refresh_query = gateway_core::account::ProviderRefreshQuery::new(
|
||||
account.provider().clone(),
|
||||
now + Duration::from_secs(86_400),
|
||||
now + Duration::from_secs(86_400),
|
||||
now,
|
||||
Vec::new(),
|
||||
NonZeroU32::new(1).expect("nonzero limit"),
|
||||
);
|
||||
assert!(!refresh_query.contains(account));
|
||||
let requests = server.received_requests().await.expect("requests");
|
||||
assert_eq!(requests[0].headers["originator"], "Codex Desktop");
|
||||
assert!(!requests[0].headers.contains_key("chatgpt-account-id"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_times_out_without_falling_back_to_document_identity() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(pat_identity())
|
||||
.set_delay(Duration::from_secs(2)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = OpenAiTokenClient::new(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.timeout(Duration::from_millis(100))
|
||||
.build()
|
||||
.expect("short-lived test client"),
|
||||
TokenClientConfig {
|
||||
client_id: "test-public-client".to_owned(),
|
||||
token_endpoint: format!("{}/oauth/token", server.uri()),
|
||||
},
|
||||
provider_openai::OpenAiConfig::default().wire_profile_state(),
|
||||
);
|
||||
let service = CodexCredentialAdminService::new(
|
||||
Arc::new(UnusedRefresher),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
runtime_policy(),
|
||||
)
|
||||
.with_personal_access_token_client(Arc::new(client));
|
||||
let error = service
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accessToken": "at-timeout-token", "userId": "untrusted-user"
|
||||
}))
|
||||
.await
|
||||
.expect_err("timeout must abort PAT import");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::Unavailable
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_requires_a_validation_client() {
|
||||
let service = CodexCredentialAdminService::new(
|
||||
Arc::new(UnusedRefresher),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
let error = service
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accessToken": "at-test-token", "userId": "untrusted-user"
|
||||
}))
|
||||
.await
|
||||
.expect_err("PAT cannot be imported without validation");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::Unavailable
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_auth_json_import_and_cpr_round_trip_preserve_a_token_without_email() {
|
||||
use gateway_core::account::LoadedCredential;
|
||||
use provider_openai::credential::{CodexCredentialAdmin, ExportManagedCodexCredential};
|
||||
|
||||
let server = MockServer::start().await;
|
||||
let mut identity = pat_identity();
|
||||
identity.as_object_mut().expect("object").remove("email");
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(identity))
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = pat_service(&server);
|
||||
let imported = service
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"auth_mode": "personalAccessToken", "personal_access_token": "at-test-token"
|
||||
}))
|
||||
.await
|
||||
.expect("official PAT auth.json");
|
||||
let prepared = imported.into_accounts().remove(0);
|
||||
assert!(prepared.account.email().is_none());
|
||||
let exported = CodexCredentialAdmin
|
||||
.format_cpr_export(vec![ExportManagedCodexCredential {
|
||||
current: LoadedCredential {
|
||||
account: prepared.account,
|
||||
credential: prepared.credential,
|
||||
},
|
||||
added_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
}])
|
||||
.expect("export PAT using the existing credential schema");
|
||||
let imported = service
|
||||
.prepare_import_document(serde_json::to_value(exported).expect("export JSON"))
|
||||
.await
|
||||
.expect("re-import verifies PAT again");
|
||||
assert_eq!(
|
||||
imported.accounts()[0].account.upstream_user_id(),
|
||||
Some("pat-user")
|
||||
);
|
||||
assert!(!imported.accounts()[0].account.has_refresh_token());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_fails_closed_for_rejection_and_unavailable_upstream_without_leaking_secrets() {
|
||||
for (status, expected) in [
|
||||
(401, PersonalAccessTokenError::Rejected),
|
||||
(403, PersonalAccessTokenError::Rejected),
|
||||
(429, PersonalAccessTokenError::Unavailable),
|
||||
(500, PersonalAccessTokenError::Unavailable),
|
||||
] {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(status).set_body_string("pat-response-secret-marker"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let error = pat_service(&server)
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accessToken": "at-request-secret-marker", "userId": "document-user"
|
||||
}))
|
||||
.await
|
||||
.expect_err("no account may be prepared after failed validation");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(expected)
|
||||
);
|
||||
let message = format!("{error} {error:?}");
|
||||
assert!(!message.contains("pat-response-secret-marker"));
|
||||
assert!(!message.contains("at-request-secret-marker"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_rejects_missing_or_invalid_whoami_identity_fields() {
|
||||
for field in [
|
||||
"chatgpt_user_id",
|
||||
"chatgpt_account_id",
|
||||
"chatgpt_plan_type",
|
||||
"chatgpt_account_is_fedramp",
|
||||
] {
|
||||
for value in [
|
||||
None,
|
||||
Some(serde_json::json!("")),
|
||||
Some(serde_json::json!(null)),
|
||||
] {
|
||||
let server = MockServer::start().await;
|
||||
let mut identity = pat_identity();
|
||||
if let Some(value) = value {
|
||||
identity[field] = value;
|
||||
} else {
|
||||
identity.as_object_mut().expect("object").remove(field);
|
||||
}
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(identity))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let error = pat_service(&server)
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"access_token": "at-test-token", "user_id": "not-a-fallback"
|
||||
}))
|
||||
.await
|
||||
.expect_err("malformed whoami must not fall back to document identity");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::InvalidResponse
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_rejects_invalid_json_and_oversized_chunked_success_bodies() {
|
||||
for body in ["not-json-secret-marker".to_owned(), "x".repeat(70 * 1024)] {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("transfer-encoding", "chunked")
|
||||
.set_body_string(body),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let error = pat_service(&server)
|
||||
.prepare_import_document(serde_json::json!({"accessToken": "at-test-token"}))
|
||||
.await
|
||||
.expect_err("invalid body");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::InvalidResponse
|
||||
)
|
||||
);
|
||||
assert!(!format!("{error:?}").contains("secret-marker"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_does_not_follow_redirects_or_forward_bearer_to_another_endpoint() {
|
||||
let server = MockServer::start().await;
|
||||
let destination = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(302).insert_header("location", destination.uri()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let error = pat_service(&server)
|
||||
.prepare_import_document(serde_json::json!({"accessToken": "at-test-token"}))
|
||||
.await
|
||||
.expect_err("redirects are not accepted");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::Unavailable
|
||||
)
|
||||
);
|
||||
assert!(
|
||||
destination
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("redirect target")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_rejects_malformed_tokens_before_sending_a_request() {
|
||||
let server = MockServer::start().await;
|
||||
let service = pat_service(&server);
|
||||
for token in ["at-", "at-has space", "at-has\nnewline", "at-has\0control"] {
|
||||
let error = service
|
||||
.prepare_import_document(serde_json::json!({"accessToken": token}))
|
||||
.await
|
||||
.expect_err("invalid PAT");
|
||||
assert_eq!(
|
||||
error,
|
||||
provider_openai::credential::CodexCredentialAdminError::PersonalAccessToken(
|
||||
PersonalAccessTokenError::InvalidToken
|
||||
)
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_jwt_import_does_not_call_pat_whoami_even_when_client_is_configured() {
|
||||
let server = MockServer::start().await;
|
||||
let jwt = test_jwt(
|
||||
serde_json::json!({"https://api.openai.com/auth": {"chatgpt_user_id": "jwt-user"}}),
|
||||
);
|
||||
let imported = pat_service(&server)
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accessToken": jwt, "refreshToken": "jwt-refresh", "userId": "untrusted"
|
||||
}))
|
||||
.await
|
||||
.expect("existing JWT behavior");
|
||||
assert_eq!(
|
||||
imported.accounts()[0].account.upstream_user_id(),
|
||||
Some("jwt-user")
|
||||
);
|
||||
assert!(imported.accounts()[0].account.has_refresh_token());
|
||||
assert!(
|
||||
server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oauth_import_uses_official_chatgpt_access_token_claims() {
|
||||
let access_token = test_jwt(serde_json::json!({
|
||||
@@ -281,3 +743,41 @@ async fn oauth_import_rejects_legacy_bare_token_field() {
|
||||
provider_openai::credential::CodexCredentialAdminError::InvalidCredential
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pat_import_uses_account_proxy_without_direct_fallback() {
|
||||
for status in [200, 503] {
|
||||
let origin = MockServer::start().await;
|
||||
let proxy = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/accounts/v1/user-auth-credential/whoami"))
|
||||
.and(header("authorization", "Bearer at-proxy-token"))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(pat_identity()))
|
||||
.expect(1)
|
||||
.mount(&proxy)
|
||||
.await;
|
||||
let imported = pat_service(&origin)
|
||||
.prepare_import_document(serde_json::json!({
|
||||
"accessToken": "at-proxy-token",
|
||||
"outboundProxyUrl": proxy.uri(),
|
||||
}))
|
||||
.await;
|
||||
if status == 200 {
|
||||
let prepared = imported.expect("PAT validated through account proxy");
|
||||
assert_eq!(
|
||||
prepared.accounts()[0].account.upstream_user_id(),
|
||||
Some("pat-user")
|
||||
);
|
||||
assert!(prepared.accounts()[0].account.outbound_proxy().is_some());
|
||||
} else {
|
||||
assert!(imported.is_err(), "proxy failure must abort PAT import");
|
||||
}
|
||||
assert!(
|
||||
origin
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("origin requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,46 @@ use secrecy::{ExposeSecret, SecretString};
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_and_code_exchange_use_the_selected_account_proxy() {
|
||||
let direct = MockServer::start().await;
|
||||
let proxy_a = MockServer::start().await;
|
||||
let proxy_b = MockServer::start().await;
|
||||
for (proxy, access) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] {
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": access,
|
||||
"refresh_token": "rotated-refresh",
|
||||
"id_token": "header.e30.signature"
|
||||
})))
|
||||
.expect(2)
|
||||
.mount(proxy)
|
||||
.await;
|
||||
}
|
||||
let client = client(&direct);
|
||||
for (proxy, expected) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] {
|
||||
let proxy = gateway_core::account::OutboundProxy::parse(&proxy.uri()).unwrap();
|
||||
let refreshed = client
|
||||
.refresh_with_proxy("initial-refresh", Some(&proxy))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(refreshed.access_token.as_deref(), Some(expected));
|
||||
let exchanged = client
|
||||
.exchange_with_proxy(
|
||||
AuthorizationCodeGrant {
|
||||
code: SecretString::from("test-code"),
|
||||
code_verifier: SecretString::from("test-verifier"),
|
||||
},
|
||||
Some(&proxy),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(exchanged.secret.access_token.expose_secret(), expected);
|
||||
}
|
||||
assert!(direct.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_tracks_the_online_profile_without_contaminating_code_exchange() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
@@ -0,0 +1,336 @@
|
||||
use super::*;
|
||||
use gateway_core::{
|
||||
account::{CredentialRevision, OutboundProxy, ProviderAccount, ProviderAccountId},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
|
||||
fn account(id: &str, proxy: Option<&str>) -> ProviderAccount {
|
||||
ProviderAccount::new(
|
||||
ProviderAccountId::new(format!("acct_{id}")).unwrap(),
|
||||
ProviderKind::new("openai").unwrap(),
|
||||
id.to_owned(),
|
||||
Some(id.to_owned()),
|
||||
"oauth".to_owned(),
|
||||
CredentialRevision::new(1).unwrap(),
|
||||
None,
|
||||
)
|
||||
.with_outbound_proxy(proxy.map(|value| OutboundProxy::parse(value).unwrap()))
|
||||
}
|
||||
|
||||
async fn http_exit(label: &'static str) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request_with_body(&mut stream).await;
|
||||
let body = format!(
|
||||
"event: response.completed\ndata: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_{label}\",\"status\":\"completed\",\"output\":[],\"usage\":{{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}}}}\n\n"
|
||||
);
|
||||
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
let header_end = request
|
||||
.windows(4)
|
||||
.position(|bytes| bytes == b"\r\n\r\n")
|
||||
.unwrap();
|
||||
String::from_utf8(request[..header_end].to_vec()).unwrap()
|
||||
});
|
||||
(format!("http://{address}"), task)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_sse_accounts_use_distinct_proxies_and_clearing_restores_direct_client() {
|
||||
let (exit_a, a) = http_exit("exit_a").await;
|
||||
let (exit_b, b) = http_exit("exit_b").await;
|
||||
let (direct_url, direct) = http_exit("direct").await;
|
||||
let base = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
&direct_url,
|
||||
test_wire_profile(),
|
||||
);
|
||||
let client_a = base.for_account(&account("a", Some(&exit_a))).unwrap();
|
||||
let client_b = base.for_account(&account("b", Some(&exit_b))).unwrap();
|
||||
let mut request = codex_request("gpt-5.5", "", Vec::new());
|
||||
request.force_http_sse = true;
|
||||
let (response_a, response_b) = tokio::join!(
|
||||
client_a.create_response(&request, request_context("proxy-a", Some("a"))),
|
||||
client_b.create_response(&request, request_context("proxy-b", Some("b")))
|
||||
);
|
||||
assert!(response_a.unwrap().body.contains("resp_exit_a"));
|
||||
assert!(response_b.unwrap().body.contains("resp_exit_b"));
|
||||
let cleared = client_a.for_account(&account("a", None)).unwrap();
|
||||
assert!(
|
||||
cleared
|
||||
.create_response(&request, request_context("direct", Some("a")))
|
||||
.await
|
||||
.unwrap()
|
||||
.body
|
||||
.contains("resp_direct")
|
||||
);
|
||||
for head in [a.await.unwrap(), b.await.unwrap()] {
|
||||
assert!(head.starts_with(&format!("POST {direct_url}/codex/responses HTTP/1.1")));
|
||||
}
|
||||
assert!(
|
||||
direct
|
||||
.await
|
||||
.unwrap()
|
||||
.starts_with("POST /codex/responses HTTP/1.1")
|
||||
);
|
||||
}
|
||||
|
||||
async fn websocket_exit(label: &'static str) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let connect = read_http_request(&mut stream).await;
|
||||
assert!(connect.starts_with("CONNECT upstream.invalid:80 HTTP/1.1"));
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let mut websocket = accept_codex_test_websocket(stream).await;
|
||||
assert!(matches!(websocket.next().await, Some(Ok(Message::Text(_)))));
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response(label, 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
connect
|
||||
});
|
||||
(format!("http://user:pass@{address}"), task)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_connect_and_pool_are_isolated_when_account_proxy_changes() {
|
||||
let (exit_a, a) = websocket_exit("resp_exit_a").await;
|
||||
let (exit_b, b) = websocket_exit("resp_exit_b").await;
|
||||
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
|
||||
let base = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
"http://upstream.invalid",
|
||||
test_wire_profile(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let mut request = codex_request("gpt-5.5", "", Vec::new());
|
||||
request.set_previous_response_id(Some("resp_previous".to_owned()));
|
||||
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
|
||||
for (proxy, expected) in [(exit_a, "resp_exit_a"), (exit_b, "resp_exit_b")] {
|
||||
let client = base
|
||||
.for_account(&account("same-account", Some(&proxy)))
|
||||
.unwrap();
|
||||
let response = timeout(
|
||||
Duration::from_secs(5),
|
||||
client.create_response(&request, request_context("proxy-ws", Some("same-account"))),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(response.body.contains(expected));
|
||||
}
|
||||
for head in [a.await.unwrap(), b.await.unwrap()] {
|
||||
assert_eq!(
|
||||
read_header_value(&head, "Proxy-Authorization"),
|
||||
Some("Basic dXNlcjpwYXNz")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unavailable_proxy_fails_without_direct_fallback() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let unreachable = format!("http://{}", listener.local_addr().unwrap());
|
||||
drop(listener);
|
||||
let direct = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{}", direct.local_addr().unwrap()),
|
||||
test_wire_profile(),
|
||||
);
|
||||
let client = base.for_account(&account("a", Some(&unreachable))).unwrap();
|
||||
let mut request = codex_request("gpt-5.5", "", Vec::new());
|
||||
request.force_http_sse = true;
|
||||
assert!(
|
||||
client
|
||||
.create_response(&request, request_context("blocked", Some("a")))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
timeout(Duration::from_millis(30), direct.accept())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
async fn socks_exit(stream: &mut TcpStream) -> (u8, String, u16) {
|
||||
assert_eq!(stream.read_u8().await.unwrap(), 5);
|
||||
let count = stream.read_u8().await.unwrap();
|
||||
let mut methods = vec![0; usize::from(count)];
|
||||
stream.read_exact(&mut methods).await.unwrap();
|
||||
assert!(methods.contains(&2));
|
||||
stream.write_all(&[5, 2]).await.unwrap();
|
||||
assert_eq!(stream.read_u8().await.unwrap(), 1);
|
||||
let count = stream.read_u8().await.unwrap();
|
||||
let mut user = vec![0; usize::from(count)];
|
||||
stream.read_exact(&mut user).await.unwrap();
|
||||
let count = stream.read_u8().await.unwrap();
|
||||
let mut password = vec![0; usize::from(count)];
|
||||
stream.read_exact(&mut password).await.unwrap();
|
||||
assert_eq!(user, b"user@exit");
|
||||
assert_eq!(password, b"pass:word");
|
||||
stream.write_all(&[1, 0]).await.unwrap();
|
||||
let mut header = [0; 4];
|
||||
stream.read_exact(&mut header).await.unwrap();
|
||||
assert_eq!(&header[..3], &[5, 1, 0]);
|
||||
let host = match header[3] {
|
||||
1 => {
|
||||
let mut ip = [0; 4];
|
||||
stream.read_exact(&mut ip).await.unwrap();
|
||||
std::net::Ipv4Addr::from(ip).to_string()
|
||||
}
|
||||
4 => {
|
||||
let mut ip = [0; 16];
|
||||
stream.read_exact(&mut ip).await.unwrap();
|
||||
std::net::Ipv6Addr::from(ip).to_string()
|
||||
}
|
||||
3 => {
|
||||
let count = stream.read_u8().await.unwrap();
|
||||
let mut host = vec![0; usize::from(count)];
|
||||
stream.read_exact(&mut host).await.unwrap();
|
||||
String::from_utf8(host).unwrap()
|
||||
}
|
||||
_ => panic!("unexpected SOCKS address type"),
|
||||
};
|
||||
let port = stream.read_u16().await.unwrap();
|
||||
stream
|
||||
.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 80])
|
||||
.await
|
||||
.unwrap();
|
||||
(header[3], host, port)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn socks_account_egress_preserves_dns_mode_auth_and_ipv6_for_sse_and_websocket() {
|
||||
for websocket in [false, true] {
|
||||
let mut targets = vec![
|
||||
("socks5", "localhost"),
|
||||
("socks5h", "upstream.invalid"),
|
||||
("socks5h", "[::1]"),
|
||||
];
|
||||
// reqwest's local resolver does not accept IPv6 literals; use socks5h for those.
|
||||
if websocket {
|
||||
targets.push(("socks5", "[::1]"));
|
||||
}
|
||||
for (scheme, host) in targets {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let proxy = format!(
|
||||
"{scheme}://user%40exit:pass%3Aword@{}",
|
||||
listener.local_addr().unwrap()
|
||||
);
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let target = socks_exit(&mut stream).await;
|
||||
if websocket {
|
||||
let mut stream = accept_codex_test_websocket(stream).await;
|
||||
assert!(matches!(stream.next().await, Some(Ok(Message::Text(_)))));
|
||||
stream
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_socks", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
read_http_request_with_body(&mut stream).await;
|
||||
let body = format!(
|
||||
"data: {}\n\n",
|
||||
completed_websocket_response("resp_socks", 2, 1)
|
||||
);
|
||||
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
}
|
||||
target
|
||||
});
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{host}"),
|
||||
test_wire_profile(),
|
||||
)
|
||||
.for_account(&account("socks", Some(&proxy)))
|
||||
.unwrap();
|
||||
let mut request = codex_request("gpt-5.5", "", Vec::new());
|
||||
request.force_http_sse = !websocket;
|
||||
if websocket {
|
||||
request.set_previous_response_id(Some("resp_previous".to_owned()));
|
||||
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
|
||||
}
|
||||
let response = timeout(
|
||||
Duration::from_secs(5),
|
||||
client.create_response(&request, request_context("socks", Some("socks"))),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap_or_else(|error| {
|
||||
panic!("websocket={websocket}, scheme={scheme}, host={host}: {error}")
|
||||
});
|
||||
assert!(response.body.contains("resp_socks"));
|
||||
let (address_type, target, port) = task.await.unwrap();
|
||||
assert_eq!(port, 80);
|
||||
match host {
|
||||
"localhost" => {
|
||||
assert!(matches!(address_type, 1 | 4));
|
||||
assert!(target.parse::<std::net::IpAddr>().unwrap().is_loopback());
|
||||
}
|
||||
"[::1]" => assert_eq!((address_type, target.as_str()), (4, "::1")),
|
||||
_ => assert_eq!((address_type, target.as_str()), (3, host)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn https_account_proxy_starts_tls_and_rejection_never_falls_back_to_direct() {
|
||||
for websocket in [false, true] {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let proxy = format!("https://{}", listener.local_addr().unwrap());
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let mut header = [0; 5];
|
||||
stream.read_exact(&mut header).await.unwrap();
|
||||
assert_eq!(&header[..2], &[22, 3]);
|
||||
let mut hello = vec![0; usize::from(u16::from_be_bytes([header[3], header[4]]))];
|
||||
stream.read_exact(&mut hello).await.unwrap();
|
||||
stream.write_all(&[21, 3, 3, 0, 2, 2, 40]).await.unwrap();
|
||||
});
|
||||
let direct = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{}", direct.local_addr().unwrap()),
|
||||
test_wire_profile(),
|
||||
)
|
||||
.for_account(&account("tls-proxy", Some(&proxy)))
|
||||
.unwrap();
|
||||
let mut request = codex_request("gpt-5.5", "", Vec::new());
|
||||
request.force_http_sse = !websocket;
|
||||
if websocket {
|
||||
request.set_previous_response_id(Some("resp_previous".to_owned()));
|
||||
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
|
||||
}
|
||||
assert!(
|
||||
timeout(
|
||||
Duration::from_secs(5),
|
||||
client.create_response(&request, request_context("tls-proxy", Some("tls-proxy")))
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
timeout(Duration::from_secs(5), task)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(
|
||||
timeout(Duration::from_millis(30), direct.accept())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,7 @@ use tokio_tungstenite::{
|
||||
},
|
||||
};
|
||||
|
||||
mod account_proxy;
|
||||
mod canonical;
|
||||
mod catalog;
|
||||
mod client;
|
||||
|
||||
@@ -20,7 +20,7 @@ gateway-protocol = { path = "../../gateway-protocol" }
|
||||
getrandom = "0.4.3"
|
||||
jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs"] }
|
||||
# 与官方 Grok CLI 一致:reqwest 不启用压缩特性,上游请求不带 accept-encoding。
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2"] }
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2", "socks"] }
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
|
||||
semver = "1.0.27"
|
||||
|
||||
@@ -431,7 +431,6 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
|
||||
let document = GrokOAuthImportDocument::parse_json(&document)
|
||||
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
|
||||
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let mut subjects = BTreeSet::new();
|
||||
let mut credentials = Vec::new();
|
||||
// 逐条目独立成败:verify 会真实轮换 RT,后续条目的任何失败都不得
|
||||
@@ -444,11 +443,16 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
for entry in document.into_entries() {
|
||||
let name = entry.name().to_owned();
|
||||
let email = entry.email().map(str::to_owned);
|
||||
let tokens = match self
|
||||
.oauth
|
||||
.verify_imported_credential(&discovery, entry.into_candidate())
|
||||
.await
|
||||
{
|
||||
let outbound_proxy = entry.outbound_proxy().cloned();
|
||||
let oauth = self.oauth.with_outbound_proxy(outbound_proxy.clone());
|
||||
let verified = async {
|
||||
let discovery = oauth.discover().await?;
|
||||
oauth
|
||||
.verify_imported_credential(&discovery, entry.into_candidate())
|
||||
.await
|
||||
}
|
||||
.await;
|
||||
let tokens = match verified {
|
||||
Ok(tokens) => tokens,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
@@ -493,6 +497,10 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let prepared = NewProviderAccount {
|
||||
account: prepared.account.with_outbound_proxy(outbound_proxy),
|
||||
credential: prepared.credential,
|
||||
};
|
||||
match prepared_create(prepared, Utc::now()) {
|
||||
Ok(credential) => credentials.push(credential),
|
||||
Err(error) => {
|
||||
@@ -533,7 +541,12 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
return Err(provider_error(ProviderAdminErrorKind::NotFound));
|
||||
}
|
||||
}
|
||||
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let discovery = self
|
||||
.oauth
|
||||
.with_outbound_proxy(pending.outbound_proxy().cloned())
|
||||
.discover()
|
||||
.await
|
||||
.map_err(map_oauth_error)?;
|
||||
let redirect = RedirectUriAllowlist::new([crate::OFFICIAL_REDIRECT_URI])
|
||||
.and_then(|allowlist| allowlist.authorize(crate::OFFICIAL_REDIRECT_URI))
|
||||
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?;
|
||||
@@ -623,9 +636,11 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
rotation.expires_at,
|
||||
),
|
||||
);
|
||||
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let tokens = self
|
||||
let oauth = self
|
||||
.oauth
|
||||
.with_outbound_proxy(current.account.outbound_proxy().cloned());
|
||||
let discovery = oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let tokens = oauth
|
||||
.verify_imported_credential(&discovery, candidate)
|
||||
.await
|
||||
.map_err(|error| map_failure_class(error.class()))?;
|
||||
@@ -775,9 +790,28 @@ impl XaiAdminProvider {
|
||||
let grant = authorization
|
||||
.accept_authorization_input(callback_url)
|
||||
.map_err(map_oauth_error)?;
|
||||
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let tokens = self
|
||||
.oauth
|
||||
let current = match stored.mutation.target() {
|
||||
AuthorizationMutationTarget::Create { .. } => None,
|
||||
AuthorizationMutationTarget::Reauthorize { account_id } => {
|
||||
let current = self
|
||||
.accounts
|
||||
.load_current_credential(account_id)
|
||||
.await
|
||||
.map_err(map_store_error)?;
|
||||
if current.account.provider() != &self.provider_kind {
|
||||
return Err(provider_error(ProviderAdminErrorKind::NotFound));
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
};
|
||||
let proxy = current
|
||||
.as_ref()
|
||||
.map_or(stored.mutation.outbound_proxy(), |current| {
|
||||
current.account.outbound_proxy()
|
||||
});
|
||||
let oauth = self.oauth.with_outbound_proxy(proxy.cloned());
|
||||
let discovery = oauth.discover().await.map_err(map_oauth_error)?;
|
||||
let tokens = oauth
|
||||
.exchange_authorization_code(&discovery, grant)
|
||||
.await
|
||||
.map_err(map_oauth_error)?;
|
||||
@@ -798,17 +832,17 @@ impl XaiAdminProvider {
|
||||
enabled: true,
|
||||
})
|
||||
.map_err(map_repository_error)?;
|
||||
let prepared = NewProviderAccount {
|
||||
account: prepared
|
||||
.account
|
||||
.with_outbound_proxy(stored.mutation.outbound_proxy().cloned()),
|
||||
credential: prepared.credential,
|
||||
};
|
||||
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
|
||||
}
|
||||
AuthorizationMutationTarget::Reauthorize { account_id } => {
|
||||
let current = self
|
||||
.accounts
|
||||
.load_current_credential(account_id)
|
||||
.await
|
||||
.map_err(map_store_error)?;
|
||||
if current.account.provider() != &self.provider_kind {
|
||||
return Err(provider_error(ProviderAdminErrorKind::NotFound));
|
||||
}
|
||||
AuthorizationMutationTarget::Reauthorize { .. } => {
|
||||
let current =
|
||||
current.ok_or_else(|| provider_error(ProviderAdminErrorKind::Internal))?;
|
||||
let prepared = verified_rotation(current, tokens)?;
|
||||
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
|
||||
prepared,
|
||||
@@ -898,6 +932,7 @@ fn prepared_create(
|
||||
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
|
||||
next_refresh_at: account.next_refresh_at().map(Into::into),
|
||||
enabled: account.enabled(),
|
||||
outbound_proxy: account.outbound_proxy().cloned(),
|
||||
credential_state: account.credential_state(),
|
||||
credential_observed_at: observed_at,
|
||||
})
|
||||
@@ -998,6 +1033,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(),
|
||||
@@ -1032,6 +1068,7 @@ fn project_quota(
|
||||
) -> ProviderQuota {
|
||||
let Some(snapshot) = snapshot else {
|
||||
return ProviderQuota {
|
||||
plan_type: None,
|
||||
observed_at: None,
|
||||
refresh_token_expires_at,
|
||||
windows: Vec::new(),
|
||||
@@ -1103,6 +1140,7 @@ fn project_quota(
|
||||
}
|
||||
};
|
||||
ProviderQuota {
|
||||
plan_type: None,
|
||||
observed_at: Some(snapshot.observed_at()),
|
||||
refresh_token_expires_at,
|
||||
windows: vec![window],
|
||||
|
||||
@@ -1162,7 +1162,8 @@ async fn fetch_candidate_catalog(
|
||||
.map(|value| SecretValue::new(value.to_owned())),
|
||||
wire_profile,
|
||||
)
|
||||
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?;
|
||||
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?
|
||||
.with_outbound_proxy(candidate.account.outbound_proxy().cloned());
|
||||
let snapshot = client
|
||||
.fetch(&session)
|
||||
.await
|
||||
@@ -1214,6 +1215,7 @@ fn billing_session(
|
||||
.map(|value| SecretValue::new(value.to_owned())),
|
||||
wire_profile.clone(),
|
||||
)
|
||||
.map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned()))
|
||||
.map_err(|_| GrokQuotaError::InvalidData)
|
||||
}
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user