mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
68f05a4481 | ||
|
|
bb4d178b5d | ||
|
|
73487fd98b | ||
|
|
66fd09eba6 | ||
|
|
bd152dec00 | ||
|
|
7090dca502 | ||
|
|
029f3124bf | ||
|
|
840707e5a1 | ||
|
|
704b652f86 | ||
|
|
2b98839c29 | ||
|
|
90dd5b5a57 | ||
|
|
5ad9378b33 | ||
|
|
4a413f5bb7 | ||
|
|
0bab566d04 | ||
|
|
9b19cf083a | ||
|
|
8b22788746 | ||
|
|
e973fc516e | ||
|
|
12ecbf574a | ||
|
|
443bd6ad17 | ||
|
|
c795098a87 | ||
|
|
1f0009341d | ||
|
|
1359eaa894 | ||
|
|
74318af570 | ||
|
|
ddc10d7c1a | ||
|
|
6bc2189e6f |
@@ -88,6 +88,13 @@ jobs:
|
||||
fail("Release tag #{release_tag} does not match #{version_path} tag #{expected_tag}") unless release_tag == expected_tag
|
||||
fail("#{version_path}: version #{version} is not a valid semver") unless version.match?(SEMVER)
|
||||
|
||||
notes_path = "release/notes.md"
|
||||
fail("#{notes_path}: release notes are required") unless File.file?(notes_path)
|
||||
notes = File.readlines(notes_path)
|
||||
unless notes.first&.strip == "# #{release_tag}" && !notes.drop(1).join.strip.empty?
|
||||
fail("#{notes_path}: expected non-empty notes starting with '# #{release_tag}'")
|
||||
end
|
||||
|
||||
platforms = platforms_data["platforms"]
|
||||
fail("#{platforms_path}: platforms must be a non-empty array") unless platforms.is_a?(Array) && !platforms.empty?
|
||||
|
||||
@@ -643,15 +650,10 @@ jobs:
|
||||
|
||||
cosign sign --yes "${images[@]}"
|
||||
|
||||
- name: Publish GitHub Release
|
||||
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
|
||||
with:
|
||||
tag_name: ${{ needs.release-info.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
|
||||
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
|
||||
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
|
||||
body: |
|
||||
- name: Prepare release notes
|
||||
shell: bash
|
||||
env:
|
||||
INSTALLATION_NOTES: |
|
||||
> [!WARNING]
|
||||
> **Upgrading across major versions?** Online upgrades across major versions are not supported. Deploy the new major fresh with a new `.runtime/` data directory, then re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
|
||||
|
||||
@@ -672,6 +674,26 @@ jobs:
|
||||
|
||||
- [GitHub Repository](https://github.com/${{ github.repository }})
|
||||
- [Installation Guide](https://github.com/${{ github.repository }}/blob/main/deploy/README.md)
|
||||
run: |
|
||||
python3 <<'PY'
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
notes = "\n".join(Path("release/notes.md").read_text().splitlines()[1:]).strip()
|
||||
installation = os.environ["INSTALLATION_NOTES"].strip()
|
||||
output = Path(os.environ["RUNNER_TEMP"]) / "release-notes.md"
|
||||
output.write_text(f"{notes}\n\n{installation}\n")
|
||||
PY
|
||||
|
||||
- name: Publish GitHub Release
|
||||
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
|
||||
with:
|
||||
tag_name: ${{ needs.release-info.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
|
||||
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
|
||||
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
|
||||
body_path: ${{ runner.temp }}/release-notes.md
|
||||
fail_on_unmatched_files: true
|
||||
files: |
|
||||
dist/release/*
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -10,8 +10,9 @@
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](LICENSE)
|
||||
|
||||
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档)
|
||||
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档) · [社区](#社区) · [许可证](#许可证)
|
||||
|
||||
</div>
|
||||
|
||||
@@ -20,8 +21,8 @@
|
||||
|
||||
## 功能
|
||||
|
||||
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看。
|
||||
- **访问控制**:通过账号分组分配密钥权限,设置并发和请求频率。
|
||||
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看;代理管理支持独立保存、连接测试和账号绑定,以及 sub2api OpenAI 账号代理导入。
|
||||
- **访问控制**:通过账号分组分配密钥权限,设置日/周 USD 限额、并发和请求频率。
|
||||
- **客户端接入**:支持 Codex CLI、桌面端与 Responses API 客户端,提供 HTTP 流式响应和 WebSocket。
|
||||
- **请求诊断**:查看请求记录、用量、费用、延迟和错误详情。
|
||||
- **部署运维**:Docker Compose 部署、S3/R2 数据库备份、同大版本在线更新与回滚。
|
||||
@@ -122,3 +123,11 @@ curl http://127.0.0.1:8080/v1/models \
|
||||
- [系统架构](docs/architecture.md)
|
||||
- [管理端主题](docs/theme.md)
|
||||
- [数据库迁移](backend/migrations/README.md)
|
||||
|
||||
## 社区
|
||||
|
||||
感谢 [LINUX DO](https://linux.do) 社区提供开放、友善的技术交流平台。
|
||||
|
||||
## 许可证
|
||||
|
||||
本项目基于 [Apache License 2.0](LICENSE) 开源。
|
||||
|
||||
Generated
+11
-7
@@ -744,6 +744,7 @@ dependencies = [
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1191,7 +1192,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]]
|
||||
@@ -1433,6 +1434,7 @@ dependencies = [
|
||||
"cron",
|
||||
"futures",
|
||||
"gateway-core",
|
||||
"heck",
|
||||
"hex",
|
||||
"rand_core 0.6.4",
|
||||
"secrecy",
|
||||
@@ -1488,6 +1490,7 @@ dependencies = [
|
||||
"sha2 0.11.0",
|
||||
"thiserror",
|
||||
"tracing",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
@@ -1877,10 +1880,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 +2612,7 @@ dependencies = [
|
||||
"hex",
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
"hyper-util",
|
||||
"reqwest",
|
||||
"roxmltree",
|
||||
"rustls",
|
||||
@@ -2627,6 +2630,7 @@ dependencies = [
|
||||
"tokio-rustls",
|
||||
"tokio-tungstenite 0.28.0",
|
||||
"tokio-util",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"tungstenite 0.27.0",
|
||||
@@ -2735,7 +2739,7 @@ dependencies = [
|
||||
"once_cell",
|
||||
"socket2",
|
||||
"tracing",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3005,7 +3009,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3707,7 +3711,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" }
|
||||
|
||||
@@ -28,6 +28,7 @@ config = { version = "0.15.23", default-features = false, features = ["yaml"] }
|
||||
serde_json = "1"
|
||||
syn = { version = "3.0.3", features = ["full", "visit"] }
|
||||
tempfile = "3.27.0"
|
||||
tokio-rustls = { version = "=0.26.4", default-features = false }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -78,7 +78,10 @@ pub async fn run() -> Result<(), BootstrapError> {
|
||||
store.admin_ports(),
|
||||
vec![openai.admin_provider(), xai.admin_provider()],
|
||||
core.snapshot_control(),
|
||||
core.account_probe(),
|
||||
(
|
||||
core.account_probe(),
|
||||
host.proxy_probe(provider_openai::build_reqwest_client_with_custom_ca),
|
||||
),
|
||||
host.client_distribution_resolver(),
|
||||
host.system_operations(),
|
||||
)
|
||||
|
||||
@@ -182,6 +182,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[
|
||||
&[
|
||||
"client_distribution",
|
||||
"config",
|
||||
"proxy_probe",
|
||||
"serve",
|
||||
"system_update",
|
||||
"workers",
|
||||
|
||||
@@ -9,6 +9,151 @@ const REDIS_PASSWORD: &str = "222222222222222222222222222222222222222222222222";
|
||||
const ADMIN_PASSWORD: &str = "test-admin-password";
|
||||
const TOPOLOGY_CHILD_ENV: &str = "CPR_TEST_TOPOLOGY_CHILD";
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_should_use_provider_custom_ca_for_https_proxies() {
|
||||
use gateway_admin::ports::proxy::ProxyProbe;
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use gateway_host::proxy_probe::HttpProxyProbe;
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
use tokio_rustls::{
|
||||
TlsAcceptor,
|
||||
rustls::{
|
||||
ServerConfig,
|
||||
pki_types::{CertificateDer, PrivateKeyDer, pem::PemObject as _},
|
||||
},
|
||||
};
|
||||
|
||||
const CHILD_ENV: &str = "CPR_TEST_PROXY_TLS_DIRECTORY";
|
||||
let Ok(directory) = std::env::var(CHILD_ENV) else {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
generate_proxy_test_certificates(directory.path());
|
||||
// 使用子进程隔离环境变量,避免并行测试读取到临时 CA 配置。
|
||||
for ca_env in ["CODEX_CA_CERTIFICATE", "SSL_CERT_FILE"] {
|
||||
let mut child = Command::new(std::env::current_exe().unwrap());
|
||||
child
|
||||
.args([
|
||||
"--exact",
|
||||
"bootstrap::proxy_probe_should_use_provider_custom_ca_for_https_proxies",
|
||||
"--nocapture",
|
||||
])
|
||||
.env(CHILD_ENV, directory.path())
|
||||
.env_remove("CODEX_CA_CERTIFICATE")
|
||||
.env_remove("SSL_CERT_FILE")
|
||||
.env(ca_env, directory.path().join("ca.pem"));
|
||||
if ca_env == "CODEX_CA_CERTIFICATE" {
|
||||
child.env(
|
||||
"SSL_CERT_FILE",
|
||||
directory.path().join("missing-fallback.pem"),
|
||||
);
|
||||
}
|
||||
let output = child.output().unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{ca_env}: {}\n{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
}
|
||||
return;
|
||||
};
|
||||
provider_openai::ensure_rustls_provider();
|
||||
let directory = std::path::Path::new(&directory);
|
||||
let certificate = CertificateDer::from_pem_file(directory.join("server.pem")).unwrap();
|
||||
let key = PrivateKeyDer::from_pem_file(directory.join("server.key")).unwrap();
|
||||
let acceptor = TlsAcceptor::from(Arc::new(
|
||||
ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(vec![certificate], key)
|
||||
.unwrap(),
|
||||
));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let proxy =
|
||||
OutboundProxy::parse(&format!("https://{}", listener.local_addr().unwrap())).unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut stream = acceptor.accept(socket).await.unwrap();
|
||||
let mut request = Vec::new();
|
||||
while !request.ends_with(b"\r\n\r\n") {
|
||||
assert!(request.len() < 8192);
|
||||
request.push(stream.read_u8().await.unwrap());
|
||||
}
|
||||
assert!(request.starts_with(b"GET http://unresolvable.invalid/ip "));
|
||||
let body = "{\"ip\":\"203.0.113.8\"}";
|
||||
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
stream.shutdown().await.unwrap();
|
||||
});
|
||||
let probe = HttpProxyProbe::new("http://unresolvable.invalid/ip")
|
||||
.with_client_builder(provider_openai::build_reqwest_client_with_custom_ca);
|
||||
let result = probe.test(&proxy).await;
|
||||
assert!(result.success, "{}", result.message);
|
||||
assert_eq!(result.exit_ip.unwrap().to_string(), "203.0.113.8");
|
||||
tokio::time::timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn generate_proxy_test_certificates(directory: &std::path::Path) {
|
||||
let openssl = |args: &[&str]| {
|
||||
let output = Command::new("openssl")
|
||||
.args(args)
|
||||
.current_dir(directory)
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
};
|
||||
openssl(&[
|
||||
"req",
|
||||
"-x509",
|
||||
"-newkey",
|
||||
"rsa:2048",
|
||||
"-nodes",
|
||||
"-keyout",
|
||||
"ca.key",
|
||||
"-out",
|
||||
"ca.pem",
|
||||
"-days",
|
||||
"1",
|
||||
"-subj",
|
||||
"/CN=Proxy Test CA",
|
||||
]);
|
||||
openssl(&[
|
||||
"req",
|
||||
"-newkey",
|
||||
"rsa:2048",
|
||||
"-nodes",
|
||||
"-keyout",
|
||||
"server.key",
|
||||
"-out",
|
||||
"server.csr",
|
||||
"-subj",
|
||||
"/CN=localhost",
|
||||
]);
|
||||
fs::write(directory.join("extensions"), "subjectAltName=IP:127.0.0.1\nbasicConstraints=critical,CA:FALSE\nkeyUsage=digitalSignature,keyEncipherment\nextendedKeyUsage=serverAuth\n").unwrap();
|
||||
openssl(&[
|
||||
"x509",
|
||||
"-req",
|
||||
"-in",
|
||||
"server.csr",
|
||||
"-CA",
|
||||
"ca.pem",
|
||||
"-CAkey",
|
||||
"ca.key",
|
||||
"-CAcreateserial",
|
||||
"-out",
|
||||
"server.pem",
|
||||
"-days",
|
||||
"1",
|
||||
"-extfile",
|
||||
"extensions",
|
||||
]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_load_complete_terminal_example() {
|
||||
parse_config(&valid_config()).expect("terminal config example");
|
||||
|
||||
@@ -19,6 +19,7 @@ bytes = "1.11.1"
|
||||
chrono = "0.4.45"
|
||||
futures = "0.3.32"
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
heck = "0.5.0"
|
||||
hex = "0.4.3"
|
||||
rand_core = { version = "0.6.4", features = ["getrandom"] }
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
|
||||
@@ -25,7 +25,7 @@ pub use use_case::{
|
||||
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
|
||||
backup::BackupService, client_distribution::ClientDistributionService,
|
||||
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
|
||||
settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
proxies::ProxiesService, settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
};
|
||||
|
||||
use model::{AdminError, AdminErrorKind};
|
||||
@@ -149,6 +149,7 @@ pub enum AdminConfigError {
|
||||
/// 字段全部私有;调用方经 accessor 直接调用能力,不需要命名内部 `use_case` 模块。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminServices {
|
||||
proxies: Arc<dyn ProxiesService>,
|
||||
auth: Arc<dyn AuthService>,
|
||||
accounts: Arc<dyn AccountsService>,
|
||||
account_groups: Arc<dyn AccountGroupService>,
|
||||
@@ -163,6 +164,11 @@ pub struct AdminServices {
|
||||
}
|
||||
|
||||
impl AdminServices {
|
||||
#[must_use]
|
||||
pub fn proxies(&self) -> &dyn ProxiesService {
|
||||
self.proxies.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auth(&self) -> &dyn AuthService {
|
||||
self.auth.as_ref()
|
||||
@@ -247,10 +253,11 @@ pub async fn initialize(
|
||||
store: AdminStorePorts,
|
||||
providers: Vec<Arc<dyn ProviderAdmin>>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
probes: (Arc<dyn AccountProbe>, Arc<dyn ports::proxy::ProxyProbe>),
|
||||
client_distribution: Arc<dyn ClientDistributionResolver>,
|
||||
system: Arc<dyn SystemOperations>,
|
||||
) -> Result<AdminBundle, AdminError> {
|
||||
let (probe, proxy_probe) = probes;
|
||||
config
|
||||
.resolve_and_validate(Path::new("."))
|
||||
.map_err(|error| AdminError::invalid(error.to_string()))?;
|
||||
@@ -290,6 +297,12 @@ pub async fn initialize(
|
||||
backup_ports.object_store(),
|
||||
);
|
||||
let services = AdminServices {
|
||||
proxies: Arc::new(use_case::proxies::DefaultProxiesService::new(
|
||||
store.proxies(),
|
||||
proxy_probe,
|
||||
snapshot.clone(),
|
||||
registry.clone(),
|
||||
)),
|
||||
auth,
|
||||
accounts,
|
||||
account_groups: Arc::new(DefaultAccountGroupService::new(
|
||||
@@ -316,11 +329,13 @@ pub async fn initialize(
|
||||
openai: Arc::new(DefaultOpenAiService::new(
|
||||
openai,
|
||||
store.accounts(),
|
||||
store.proxies(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
xai: Arc::new(DefaultXaiService::new(
|
||||
xai,
|
||||
store.accounts(),
|
||||
store.proxies(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
backups,
|
||||
|
||||
@@ -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<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// 账号更新结果。
|
||||
@@ -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<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// 批量账号更新结果。
|
||||
|
||||
@@ -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 启用状态。
|
||||
|
||||
@@ -10,6 +10,7 @@ pub mod client_distribution;
|
||||
pub mod client_keys;
|
||||
pub mod observability;
|
||||
pub mod provider_credentials;
|
||||
pub mod proxies;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
|
||||
|
||||
@@ -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,8 @@ pub struct CredentialDetails {
|
||||
|
||||
/// Provider 正式文档批量导入命令。
|
||||
pub struct ImportCredentials {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub settings: Option<AccountImportSettings>,
|
||||
pub context: MutationContext,
|
||||
pub document: ProviderDocument,
|
||||
}
|
||||
@@ -129,6 +133,7 @@ pub struct CredentialImportResult {
|
||||
|
||||
/// Provider 解析导入文档时只接收不透明文档,不接触 revision 或审计上下文。
|
||||
pub struct PrepareCredentialImport {
|
||||
pub default_outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub document: ProviderDocument,
|
||||
}
|
||||
|
||||
@@ -144,6 +149,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 +177,8 @@ pub struct PreparedCredentialImport {
|
||||
/// Admin 交给 Store 的导入事务命令。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CredentialImportCommit {
|
||||
pub outbound_proxy: Option<super::proxies::ImportProxyBinding>,
|
||||
pub settings: Option<AccountImportSettings>,
|
||||
pub prepared: PreparedCredentialImport,
|
||||
}
|
||||
|
||||
@@ -249,6 +257,8 @@ pub struct PendingAuthorizationMutation {
|
||||
provider_kind: ProviderKind,
|
||||
target: AuthorizationMutationTarget,
|
||||
owner_binding: AuthorizationOwnerBinding,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
outbound_proxy_id: Option<String>,
|
||||
}
|
||||
|
||||
impl PendingAuthorizationMutation {
|
||||
@@ -262,6 +272,8 @@ impl PendingAuthorizationMutation {
|
||||
provider_kind,
|
||||
target,
|
||||
owner_binding,
|
||||
outbound_proxy: None,
|
||||
outbound_proxy_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -279,6 +291,29 @@ 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()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_outbound_proxy_id(mut self, id: Option<String>) -> Self {
|
||||
self.outbound_proxy_id = id;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn outbound_proxy_id(&self) -> Option<&str> {
|
||||
self.outbound_proxy_id.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
impl PendingAuthorizationMutation {
|
||||
@@ -300,7 +335,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 +346,17 @@ 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()),
|
||||
);
|
||||
}
|
||||
if let Some(id) = &self.outbound_proxy_id {
|
||||
document.insert("outbound_proxy_id".to_owned(), Value::String(id.clone()));
|
||||
}
|
||||
document
|
||||
}
|
||||
|
||||
/// 恢复 adapter 已确认版本为 v1 的中立事务字段。
|
||||
@@ -341,6 +386,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 +399,9 @@ impl PendingAuthorizationMutation {
|
||||
owner,
|
||||
started_request_id: document.started_request_id,
|
||||
},
|
||||
))
|
||||
)
|
||||
.with_outbound_proxy(proxy)
|
||||
.with_outbound_proxy_id(document.outbound_proxy_id))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -359,6 +412,10 @@ struct StoredAuthorizationMutationV1 {
|
||||
target: StoredAuthorizationTargetV1,
|
||||
owner: StoredAuthorizationOwnerV1,
|
||||
started_request_id: String,
|
||||
#[serde(default)]
|
||||
outbound_proxy_url: Option<String>,
|
||||
#[serde(default)]
|
||||
outbound_proxy_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -382,6 +439,7 @@ pub struct StartAuthorization {
|
||||
pub context: MutationContext,
|
||||
pub name: String,
|
||||
pub reauthorization: Option<ProviderAccountId>,
|
||||
pub outbound_proxy: Option<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// OAuth 流程启动结果。
|
||||
@@ -447,6 +505,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 +540,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 +558,7 @@ impl PreparedAuthorizationCommit {
|
||||
};
|
||||
AuthorizationCommitSettlement {
|
||||
command: AuthorizationCommit {
|
||||
settings,
|
||||
pending: self.pending,
|
||||
credential,
|
||||
},
|
||||
@@ -509,7 +572,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 +584,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 +782,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 +792,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 +914,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 +1100,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,
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
//! 可复用的账号出口配置,以及脱敏后的连通性测试结果。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::account::OutboundProxy;
|
||||
|
||||
use super::{PageSize, Revision, account_groups::AccountGroupRef};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AccountProxySelection {
|
||||
Direct,
|
||||
Url(OutboundProxy),
|
||||
Saved(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ImportProxyBinding {
|
||||
pub id: String,
|
||||
pub proxy: OutboundProxy,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyListQuery {
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub search: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProxyAccountRef {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub email: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub authentication_kind: String,
|
||||
pub plan_type: Option<String>,
|
||||
pub plan_type_display: Option<String>,
|
||||
pub groups: Vec<AccountGroupRef>,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// 按代理查询关联账号,分页与搜索均在存储层执行。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyAccountListQuery {
|
||||
pub proxy_id: String,
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub search: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyAccountPage {
|
||||
pub items: Vec<ProxyAccountRef>,
|
||||
pub total: u64,
|
||||
pub page: u32,
|
||||
pub page_size: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProxyTestResult {
|
||||
pub success: bool,
|
||||
pub latency_ms: u64,
|
||||
pub exit_ip: Option<std::net::IpAddr>,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub proxy: OutboundProxy,
|
||||
pub revision: Revision,
|
||||
pub account_count: u64,
|
||||
pub last_test_at: Option<DateTime<Utc>>,
|
||||
pub last_test: Option<ProxyTestResult>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyPage {
|
||||
pub items: Vec<ProxyRecord>,
|
||||
pub total: u64,
|
||||
pub page: u32,
|
||||
pub page_size: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NewProxy {
|
||||
pub name: String,
|
||||
pub proxy: OutboundProxy,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpdateProxy {
|
||||
pub id: String,
|
||||
pub revision: Revision,
|
||||
pub name: String,
|
||||
pub proxy: Option<OutboundProxy>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyMutation {
|
||||
pub config_revision: Revision,
|
||||
pub record: ProxyRecord,
|
||||
}
|
||||
@@ -3,5 +3,6 @@
|
||||
pub mod backup;
|
||||
pub mod client_distribution;
|
||||
pub mod provider;
|
||||
pub mod proxy;
|
||||
pub mod store;
|
||||
pub mod system;
|
||||
|
||||
@@ -8,6 +8,7 @@ use gateway_core::{
|
||||
operation::Operation,
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
};
|
||||
use heck::ToUpperCamelCase;
|
||||
|
||||
use crate::model::observability::{
|
||||
CalculatedBillingBreakdown, DashboardWireProfile, ProviderBillingInput,
|
||||
@@ -18,7 +19,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 +41,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 +56,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 +67,7 @@ impl ProviderAdminError {
|
||||
Self {
|
||||
kind,
|
||||
message: None,
|
||||
public_message: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +77,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 +108,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 +252,37 @@ impl ProviderAdminRegistry {
|
||||
.ok_or_else(|| ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 账号目录与关联列表共用套餐补全和展示规则,已知账号套餐优先于额度快照。
|
||||
pub(crate) fn resolve_account_plan(
|
||||
&self,
|
||||
provider_kind: &str,
|
||||
plan_type: &mut Option<String>,
|
||||
quota: Option<&ProviderQuota>,
|
||||
) -> Option<String> {
|
||||
if let Some(quota) = quota {
|
||||
quota.fill_missing_plan_type(plan_type);
|
||||
}
|
||||
self.plan_type_display(provider_kind, plan_type.as_deref())
|
||||
}
|
||||
|
||||
/// 账号页和 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));
|
||||
let display = provider.map_or_else(
|
||||
|| plan_type.to_owned(),
|
||||
|provider| provider.plan_type_display(plan_type),
|
||||
);
|
||||
// 所有 Provider 及未知套餐统一排版,保留 Plus 后缀的含义。
|
||||
Some(display.replace('+', " Plus ").to_upper_camel_case())
|
||||
}
|
||||
|
||||
/// 返回所有已注册 Provider 的 Dashboard 上游身份画像。
|
||||
pub fn dashboard_wire_profiles(&self) -> Vec<DashboardWireProfile> {
|
||||
self.providers
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
|
||||
use super::store::AdminStoreResult;
|
||||
use crate::model::{
|
||||
MutationContext, Revision,
|
||||
proxies::{
|
||||
ImportProxyBinding, NewProxy, ProxyAccountListQuery, ProxyAccountPage, ProxyListQuery,
|
||||
ProxyMutation, ProxyPage, ProxyRecord, ProxyTestResult, UpdateProxy,
|
||||
},
|
||||
};
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyStore: Send + Sync {
|
||||
/// 在凭据交换到提交期间保护选定代理的连接配置和测试结果。
|
||||
async fn reserve_import(&self, id: &str) -> AdminStoreResult<ProxyImportReservation>;
|
||||
async fn list(&self, query: ProxyListQuery) -> AdminStoreResult<ProxyPage>;
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage>;
|
||||
async fn get(&self, id: &str) -> AdminStoreResult<ProxyRecord>;
|
||||
/// 仅在账号仍绑定指定代理时解除关联,并清除账号保存的连接地址。
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &ProviderAccountId,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision>;
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation>;
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation>;
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision>;
|
||||
async fn record_test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
result: ProxyTestResult,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord>;
|
||||
}
|
||||
|
||||
/// 离开作用域时释放保护,错误返回和请求取消也遵循相同规则。
|
||||
pub trait ProxyImportGuard: Send + Sync {}
|
||||
|
||||
pub struct ProxyImportReservation {
|
||||
pub binding: ImportProxyBinding,
|
||||
pub guard: Box<dyn ProxyImportGuard>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyProbe: Send + Sync {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult;
|
||||
}
|
||||
@@ -382,6 +382,7 @@ pub struct AdminAccountStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
proxies: Arc<dyn super::proxy::ProxyStore>,
|
||||
}
|
||||
|
||||
impl AdminAccountStorePorts {
|
||||
@@ -390,11 +391,13 @@ impl AdminAccountStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
proxies: Arc<dyn super::proxy::ProxyStore>,
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
runtime,
|
||||
groups,
|
||||
proxies,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -447,6 +450,11 @@ impl AdminStorePorts {
|
||||
self.accounts.groups.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn proxies(&self) -> Arc<dyn super::proxy::ProxyStore> {
|
||||
self.accounts.proxies.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auth(&self) -> Arc<dyn AuthStore> {
|
||||
self.auth.clone()
|
||||
|
||||
@@ -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
|
||||
@@ -281,6 +281,11 @@ impl DefaultAccountsService {
|
||||
.await?;
|
||||
let usage = quota.representative_window_usage().cloned();
|
||||
Ok(AccountDirectoryItem {
|
||||
plan_type_display: self.providers.resolve_account_plan(
|
||||
stored.account.provider_kind.as_str(),
|
||||
&mut stored.account.plan_type,
|
||||
Some("a),
|
||||
),
|
||||
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();
|
||||
AccountDirectoryItem {
|
||||
plan_type_display: self.providers.resolve_account_plan(
|
||||
item.account.provider_kind.as_str(),
|
||||
&mut item.account.plan_type,
|
||||
Some("a),
|
||||
),
|
||||
usage,
|
||||
account: item.account,
|
||||
projection: item.projection,
|
||||
@@ -735,6 +745,7 @@ fn map_reset_credits_error_after_refresh(
|
||||
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
|
||||
fn empty_quota() -> ProviderQuota {
|
||||
ProviderQuota {
|
||||
plan_type: None,
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: Vec::new(),
|
||||
|
||||
@@ -99,6 +99,7 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
label: command.label,
|
||||
group_ids: command.group_ids,
|
||||
limits: command.limits,
|
||||
budget: command.budget,
|
||||
plaintext: plaintext.clone(),
|
||||
},
|
||||
context,
|
||||
|
||||
@@ -8,6 +8,7 @@ pub mod client_distribution;
|
||||
pub mod client_keys;
|
||||
pub mod observability;
|
||||
pub mod openai;
|
||||
pub mod proxies;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
pub mod xai;
|
||||
@@ -83,7 +84,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(
|
||||
@@ -109,22 +110,58 @@ async fn required_credential(
|
||||
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
|
||||
}
|
||||
|
||||
async fn import_proxy_binding(
|
||||
proxies: &dyn crate::ports::proxy::ProxyStore,
|
||||
id: Option<&str>,
|
||||
) -> Result<Option<crate::ports::proxy::ProxyImportReservation>, AdminError> {
|
||||
let Some(id) = id else { return Ok(None) };
|
||||
let reservation = proxies
|
||||
.reserve_import(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
Ok(Some(reservation))
|
||||
}
|
||||
|
||||
async fn pending_authorization(
|
||||
accounts: &dyn AccountStore,
|
||||
proxies: &dyn crate::ports::proxy::ProxyStore,
|
||||
provider_kind: &ProviderKind,
|
||||
command: &StartAuthorization,
|
||||
resource: &'static str,
|
||||
) -> Result<PendingAuthorizationMutation, AdminError> {
|
||||
let proxy;
|
||||
let mut proxy_id = None;
|
||||
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(),
|
||||
}
|
||||
}
|
||||
None => AuthorizationMutationTarget::Create {
|
||||
name: command.name.clone(),
|
||||
},
|
||||
None => {
|
||||
use crate::model::proxies::AccountProxySelection;
|
||||
proxy = match &command.outbound_proxy {
|
||||
Some(AccountProxySelection::Url(value)) => Some(value.clone()),
|
||||
Some(AccountProxySelection::Saved(id)) => {
|
||||
let record = proxies
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if !record.last_test.is_some_and(|test| test.success) {
|
||||
return Err(AdminError::conflict("请先测试代理连接"));
|
||||
}
|
||||
proxy_id = Some(record.id);
|
||||
Some(record.proxy)
|
||||
}
|
||||
Some(AccountProxySelection::Direct) | None => None,
|
||||
};
|
||||
AuthorizationMutationTarget::Create {
|
||||
name: command.name.clone(),
|
||||
}
|
||||
}
|
||||
};
|
||||
Ok(PendingAuthorizationMutation::new(
|
||||
provider_kind.clone(),
|
||||
@@ -132,7 +169,9 @@ async fn pending_authorization(
|
||||
crate::model::provider_credentials::AuthorizationOwnerBinding::from_context(
|
||||
&command.context,
|
||||
),
|
||||
))
|
||||
)
|
||||
.with_outbound_proxy(proxy)
|
||||
.with_outbound_proxy_id(proxy_id))
|
||||
}
|
||||
|
||||
fn validate_prepared_import(
|
||||
@@ -217,14 +256,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();
|
||||
|
||||
@@ -58,6 +58,7 @@ pub trait OpenAiService: Send + Sync {
|
||||
pub(crate) struct DefaultOpenAiService {
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
@@ -66,11 +67,13 @@ impl DefaultOpenAiService {
|
||||
pub(crate) fn new(
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
accounts,
|
||||
proxies,
|
||||
snapshot,
|
||||
}
|
||||
}
|
||||
@@ -128,9 +131,20 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
command: ImportCredentials,
|
||||
) -> Result<CredentialImportResult, AdminError> {
|
||||
let context = command.context;
|
||||
let proxy_reservation = super::import_proxy_binding(
|
||||
self.proxies.as_ref(),
|
||||
command.outbound_proxy_id.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
let outbound_proxy = proxy_reservation
|
||||
.as_ref()
|
||||
.map(|reservation| reservation.binding.clone());
|
||||
let prepared = self
|
||||
.provider
|
||||
.prepare_import(PrepareCredentialImport {
|
||||
default_outbound_proxy: outbound_proxy
|
||||
.as_ref()
|
||||
.map(|binding| binding.proxy.clone()),
|
||||
document: command.document,
|
||||
})
|
||||
.await
|
||||
@@ -142,9 +156,17 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
)?;
|
||||
let result = self
|
||||
.accounts
|
||||
.commit_credential_import(CredentialImportCommit { prepared }, &context)
|
||||
.commit_credential_import(
|
||||
CredentialImportCommit {
|
||||
outbound_proxy,
|
||||
prepared,
|
||||
settings: command.settings,
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
|
||||
drop(proxy_reservation);
|
||||
self.provider
|
||||
.account_facts_changed(&result.credential_ids)
|
||||
.await;
|
||||
@@ -160,6 +182,7 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
) -> Result<AuthorizationStarted, AdminError> {
|
||||
let pending = pending_authorization(
|
||||
self.accounts.as_ref(),
|
||||
self.proxies.as_ref(),
|
||||
self.provider.provider_kind(),
|
||||
&command,
|
||||
"OpenAI credential",
|
||||
@@ -173,9 +196,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 +215,7 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
let result = commit_authorization(
|
||||
self.accounts.as_ref(),
|
||||
prepared,
|
||||
settings,
|
||||
&context,
|
||||
"OpenAI authorization",
|
||||
)
|
||||
|
||||
@@ -0,0 +1,261 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::{StreamExt, stream};
|
||||
use gateway_core::{
|
||||
account::{OutboundProxy, ProviderAccountId},
|
||||
routing::ProviderKind,
|
||||
runtime::SnapshotControl,
|
||||
};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use super::{map_store_error, publish_committed};
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind, MutationContext, Revision,
|
||||
provider_credentials::{ProviderQuotaRequest, explicit_plan_type},
|
||||
proxies::*,
|
||||
},
|
||||
ports::{
|
||||
provider::ProviderAdminRegistry,
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::AdminStoreErrorKind,
|
||||
},
|
||||
};
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxiesService: Send + Sync {
|
||||
async fn list(&self, query: ProxyListQuery) -> Result<ProxyPage, AdminError>;
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> Result<ProxyAccountPage, AdminError>;
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &str,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError>;
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError>;
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError>;
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError>;
|
||||
async fn test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyRecord, AdminError>;
|
||||
/// 探测未保存的连接地址,不写入代理记录或修改账号绑定。
|
||||
async fn probe(&self, proxy: &OutboundProxy) -> Result<ProxyTestResult, AdminError>;
|
||||
}
|
||||
|
||||
pub(crate) struct DefaultProxiesService {
|
||||
store: Arc<dyn ProxyStore>,
|
||||
probe: Arc<dyn ProxyProbe>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
providers: ProviderAdminRegistry,
|
||||
test_slots: Semaphore,
|
||||
}
|
||||
|
||||
impl DefaultProxiesService {
|
||||
pub(crate) fn new(
|
||||
store: Arc<dyn ProxyStore>,
|
||||
probe: Arc<dyn ProxyProbe>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
providers: ProviderAdminRegistry,
|
||||
) -> Self {
|
||||
Self {
|
||||
store,
|
||||
probe,
|
||||
snapshot,
|
||||
providers,
|
||||
test_slots: Semaphore::new(4),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxiesService for DefaultProxiesService {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &str,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError> {
|
||||
if proxy_id.is_empty() || proxy_id.len() > 128 || proxy_id.chars().any(char::is_control) {
|
||||
return Err(AdminError::invalid("代理 ID 不合法"));
|
||||
}
|
||||
let account_id = ProviderAccountId::new(account_id.to_owned())
|
||||
.map_err(|_| AdminError::invalid("账号 ID 不合法"))?;
|
||||
let revision = self
|
||||
.store
|
||||
.remove_account(proxy_id, &account_id, context)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
if error.kind() == AdminStoreErrorKind::Conflict {
|
||||
AdminError::conflict("账号的代理绑定已变化,请刷新后重试")
|
||||
} else {
|
||||
map_store_error(error, "proxy account")
|
||||
}
|
||||
})?;
|
||||
publish_committed(self.snapshot.as_ref(), revision).await?;
|
||||
Ok(revision)
|
||||
}
|
||||
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> Result<ProxyAccountPage, AdminError> {
|
||||
if query.proxy_id.is_empty()
|
||||
|| query.proxy_id.len() > 128
|
||||
|| query.proxy_id.chars().any(char::is_control)
|
||||
|| query.page == 0
|
||||
|| query.search.len() > 256
|
||||
|| query.search.chars().any(char::is_control)
|
||||
{
|
||||
return Err(AdminError::invalid("关联账号查询参数不合法"));
|
||||
}
|
||||
let mut page = self
|
||||
.store
|
||||
.list_accounts(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
page.items = stream::iter(page.items)
|
||||
.map(|mut account| async {
|
||||
// 与账号目录一致,缺失套餐时读取已有额度快照;限制并发且不触发上游刷新。
|
||||
let mut cached_quota = None;
|
||||
if explicit_plan_type(account.plan_type.as_deref()).is_none()
|
||||
&& let Ok(kind) = ProviderKind::new(account.provider_kind.clone())
|
||||
&& let Ok(provider) = self.providers.require(&kind)
|
||||
&& let Ok(account_id) = ProviderAccountId::new(account.id.clone())
|
||||
&& let Ok(quota) = provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id,
|
||||
refresh: false,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
cached_quota = Some(quota);
|
||||
}
|
||||
account.plan_type_display = self.providers.resolve_account_plan(
|
||||
&account.provider_kind,
|
||||
&mut account.plan_type,
|
||||
cached_quota.as_ref(),
|
||||
);
|
||||
account
|
||||
})
|
||||
.buffered(8)
|
||||
.collect()
|
||||
.await;
|
||||
Ok(page)
|
||||
}
|
||||
|
||||
async fn list(&self, query: ProxyListQuery) -> Result<ProxyPage, AdminError> {
|
||||
if query.page == 0 || query.search.len() > 256 || query.search.chars().any(char::is_control)
|
||||
{
|
||||
return Err(AdminError::invalid("代理查询参数不合法"));
|
||||
}
|
||||
self.store
|
||||
.list(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))
|
||||
}
|
||||
|
||||
async fn create(
|
||||
&self,
|
||||
mut command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError> {
|
||||
command.name = validate_name(&command.name)?;
|
||||
let result = self
|
||||
.store
|
||||
.create(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
mut command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError> {
|
||||
command.name = validate_name(&command.name)?;
|
||||
let result = self
|
||||
.store
|
||||
.update(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError> {
|
||||
let result = self
|
||||
.store
|
||||
.delete(id, revision, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn probe(&self, proxy: &OutboundProxy) -> Result<ProxyTestResult, AdminError> {
|
||||
let _permit = self.test_slots.try_acquire().map_err(|_| {
|
||||
AdminError::new(AdminErrorKind::RateLimited, "代理测试繁忙,请稍后重试")
|
||||
})?;
|
||||
Ok(self.probe.test(proxy).await)
|
||||
}
|
||||
|
||||
async fn test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyRecord, AdminError> {
|
||||
let _permit = self.test_slots.try_acquire().map_err(|_| {
|
||||
AdminError::new(AdminErrorKind::RateLimited, "代理测试繁忙,请稍后重试")
|
||||
})?;
|
||||
let record = self
|
||||
.store
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if record.revision != revision {
|
||||
return Err(AdminError::conflict("代理已被修改,请刷新后重新测试"));
|
||||
}
|
||||
let result = self.probe.test(&record.proxy).await;
|
||||
self.store
|
||||
.record_test(id, revision, result, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_name(value: &str) -> Result<String, AdminError> {
|
||||
let value = value.trim();
|
||||
if value.is_empty() || value.chars().count() > 100 || value.chars().any(char::is_control) {
|
||||
return Err(AdminError::invalid("代理名称需要 1 至 100 个字符"));
|
||||
}
|
||||
Ok(value.to_owned())
|
||||
}
|
||||
@@ -49,6 +49,7 @@ pub trait XaiService: Send + Sync {
|
||||
pub(crate) struct DefaultXaiService {
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
@@ -57,11 +58,13 @@ impl DefaultXaiService {
|
||||
pub(crate) fn new(
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
accounts,
|
||||
proxies,
|
||||
snapshot,
|
||||
}
|
||||
}
|
||||
@@ -81,9 +84,20 @@ impl XaiService for DefaultXaiService {
|
||||
command: ImportCredentials,
|
||||
) -> Result<CredentialImportResult, AdminError> {
|
||||
let context = command.context;
|
||||
let proxy_reservation = super::import_proxy_binding(
|
||||
self.proxies.as_ref(),
|
||||
command.outbound_proxy_id.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
let outbound_proxy = proxy_reservation
|
||||
.as_ref()
|
||||
.map(|reservation| reservation.binding.clone());
|
||||
let prepared = self
|
||||
.provider
|
||||
.prepare_import(PrepareCredentialImport {
|
||||
default_outbound_proxy: outbound_proxy
|
||||
.as_ref()
|
||||
.map(|binding| binding.proxy.clone()),
|
||||
document: command.document,
|
||||
})
|
||||
.await
|
||||
@@ -95,9 +109,17 @@ impl XaiService for DefaultXaiService {
|
||||
)?;
|
||||
let result = self
|
||||
.accounts
|
||||
.commit_credential_import(CredentialImportCommit { prepared }, &context)
|
||||
.commit_credential_import(
|
||||
CredentialImportCommit {
|
||||
outbound_proxy,
|
||||
prepared,
|
||||
settings: command.settings,
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "xAI credential import"))?;
|
||||
drop(proxy_reservation);
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
// 导入已经提交,额度属于可重建观察事实;单个账号查询失败由周期任务重试。
|
||||
for account_id in &result.credential_ids {
|
||||
@@ -119,6 +141,7 @@ impl XaiService for DefaultXaiService {
|
||||
) -> Result<AuthorizationStarted, AdminError> {
|
||||
let pending = pending_authorization(
|
||||
self.accounts.as_ref(),
|
||||
self.proxies.as_ref(),
|
||||
self.provider.provider_kind(),
|
||||
&command,
|
||||
"xAI credential",
|
||||
@@ -132,9 +155,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 +174,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(),
|
||||
@@ -96,7 +97,21 @@ fn pending_mutation_v1_should_round_trip_all_targets_and_owners() {
|
||||
PendingAuthorizationMutation::from_storage_v1(Value::Object(
|
||||
expected.to_storage_v1()
|
||||
)),
|
||||
Ok(expected),
|
||||
Ok(expected.clone()),
|
||||
);
|
||||
let with_proxy = expected
|
||||
.with_outbound_proxy(Some(
|
||||
gateway_core::account::OutboundProxy::parse(
|
||||
"http://user:secret@proxy.example:8080",
|
||||
)
|
||||
.unwrap(),
|
||||
))
|
||||
.with_outbound_proxy_id(Some("proxy_saved".to_owned()));
|
||||
assert_eq!(
|
||||
PendingAuthorizationMutation::from_storage_v1(Value::Object(
|
||||
with_proxy.to_storage_v1()
|
||||
)),
|
||||
Ok(with_proxy)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,91 @@ 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("OpenaiDisplaySelfServeBusinessProlite"),
|
||||
)
|
||||
);
|
||||
|
||||
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, expected_display) in [
|
||||
(None, Some("free"), Some("free"), Some("OpenaiDisplayFree")),
|
||||
(
|
||||
Some(" "),
|
||||
Some("free"),
|
||||
Some("free"),
|
||||
Some("OpenaiDisplayFree"),
|
||||
),
|
||||
(
|
||||
Some("unknown"),
|
||||
Some("free"),
|
||||
Some("free"),
|
||||
Some("OpenaiDisplayFree"),
|
||||
),
|
||||
(
|
||||
Some("self_serve_business_prolite"),
|
||||
Some("team"),
|
||||
Some("self_serve_business_prolite"),
|
||||
Some("OpenaiDisplaySelfServeBusinessProlite"),
|
||||
),
|
||||
(None, 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.as_deref(), expected_display);
|
||||
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 +1476,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 +1608,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 +1656,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 +1705,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 +1772,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 +2027,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 +2062,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 +2104,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 +2231,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"),
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ mod backup;
|
||||
mod client_keys;
|
||||
mod observability;
|
||||
mod openai;
|
||||
mod proxies;
|
||||
mod settings;
|
||||
mod system;
|
||||
mod xai;
|
||||
@@ -82,6 +83,7 @@ pub(super) struct AdminHarness {
|
||||
default_password: String,
|
||||
session_ttl_minutes: u64,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn gateway_admin::ports::proxy::ProxyStore>,
|
||||
account_runtime: Arc<dyn AccountRuntimeStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
@@ -101,6 +103,7 @@ impl AdminHarness {
|
||||
default_password: "strong-test-password".to_owned(),
|
||||
session_ttl_minutes: 60,
|
||||
accounts: unavailable.clone(),
|
||||
proxies: Arc::new(proxies::TestProxies::default()),
|
||||
account_runtime: unavailable.clone(),
|
||||
account_groups: Arc::new(UnavailableAccountGroupStore),
|
||||
auth: Arc::new(BootstrapAuthStore::default()),
|
||||
@@ -179,6 +182,14 @@ impl AdminHarness {
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn proxies(
|
||||
mut self,
|
||||
proxies: Arc<dyn gateway_admin::ports::proxy::ProxyStore>,
|
||||
) -> Self {
|
||||
self.proxies = proxies;
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn system(mut self, system: Arc<dyn SystemOperations>) -> Self {
|
||||
self.system = system;
|
||||
self
|
||||
@@ -196,6 +207,7 @@ impl AdminHarness {
|
||||
self.accounts,
|
||||
self.account_runtime,
|
||||
self.account_groups,
|
||||
self.proxies,
|
||||
),
|
||||
self.auth,
|
||||
self.client_keys,
|
||||
@@ -205,7 +217,7 @@ impl AdminHarness {
|
||||
),
|
||||
self.providers,
|
||||
Arc::new(NoopSnapshot),
|
||||
self.probe,
|
||||
(self.probe, Arc::new(proxies::TestProxies::default())),
|
||||
Arc::new(NoopClientDistribution),
|
||||
self.system,
|
||||
)
|
||||
|
||||
@@ -50,6 +50,8 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
|
||||
services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: Some(super::accounts::import_settings()),
|
||||
context: context("import-openai"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -75,9 +77,52 @@ 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 {
|
||||
outbound_proxy_id: None,
|
||||
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 +134,8 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
|
||||
let result = services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-batch"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -132,6 +179,8 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
let result = services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-quota-failure"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -157,6 +206,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 +216,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 +244,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 +262,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 +272,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 +295,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 +304,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 +352,8 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
|
||||
services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-error"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -311,6 +372,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 +398,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 +450,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());
|
||||
}
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
use async_trait::async_trait;
|
||||
use gateway_admin::{
|
||||
model::{MutationContext, Revision, proxies::*},
|
||||
ports::{
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::AdminStoreResult,
|
||||
},
|
||||
};
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct TestProxies {
|
||||
pub events: Option<super::accounts::EventLog>,
|
||||
pub accounts: Option<Vec<ProxyAccountRef>>,
|
||||
}
|
||||
|
||||
struct ImportGuard(super::accounts::EventLog);
|
||||
|
||||
impl gateway_admin::ports::proxy::ProxyImportGuard for ImportGuard {}
|
||||
|
||||
impl Drop for ImportGuard {
|
||||
fn drop(&mut self) {
|
||||
self.0.lock().unwrap().push("proxy.release");
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyStore for TestProxies {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
_: &str,
|
||||
_: &ProviderAccountId,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
|
||||
async fn reserve_import(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> AdminStoreResult<gateway_admin::ports::proxy::ProxyImportReservation> {
|
||||
let events = self
|
||||
.events
|
||||
.as_ref()
|
||||
.ok_or_else(|| super::unavailable("proxy"))?;
|
||||
events.lock().unwrap().push("proxy.reserve");
|
||||
Ok(gateway_admin::ports::proxy::ProxyImportReservation {
|
||||
binding: ImportProxyBinding {
|
||||
id: id.to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:8080").unwrap(),
|
||||
},
|
||||
guard: Box::new(ImportGuard(events.clone())),
|
||||
})
|
||||
}
|
||||
async fn list(&self, _: ProxyListQuery) -> AdminStoreResult<ProxyPage> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage> {
|
||||
let items = self
|
||||
.accounts
|
||||
.clone()
|
||||
.ok_or_else(|| super::unavailable("proxy"))?;
|
||||
Ok(ProxyAccountPage {
|
||||
total: items.len() as u64,
|
||||
items,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn get(&self, _: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn create(&self, _: NewProxy, _: &MutationContext) -> AdminStoreResult<ProxyMutation> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn update(&self, _: UpdateProxy, _: &MutationContext) -> AdminStoreResult<ProxyMutation> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn delete(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn record_test(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: ProxyTestResult,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyProbe for TestProxies {
|
||||
async fn test(&self, _: &OutboundProxy) -> ProxyTestResult {
|
||||
panic!("unexpected proxy probe")
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn linked_accounts_share_plan_resolution_and_only_read_cached_quota() {
|
||||
use super::accounts::{FakeProviderAdmin, events};
|
||||
use gateway_admin::model::{PageSize, provider_credentials::ProviderQuota};
|
||||
use std::sync::Arc;
|
||||
|
||||
for (stored, cached, expected) in [
|
||||
(None, Some("free"), Some("free")),
|
||||
(Some("unknown"), Some("free"), Some("free")),
|
||||
(Some("plus"), Some("free"), Some("plus")),
|
||||
(None, None, None),
|
||||
] {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
plan_type: cached.map(str::to_owned),
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![],
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
});
|
||||
let services = super::AdminHarness::new()
|
||||
.provider(provider.clone())
|
||||
.proxies(Arc::new(TestProxies {
|
||||
accounts: Some(vec![ProxyAccountRef {
|
||||
id: "acct_plan".to_owned(),
|
||||
name: "套餐测试".to_owned(),
|
||||
email: None,
|
||||
provider_kind: "openai".to_owned(),
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
plan_type: stored.map(str::to_owned),
|
||||
plan_type_display: None,
|
||||
groups: vec![],
|
||||
enabled: true,
|
||||
}]),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let result = services
|
||||
.proxies()
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: "proxy_plan".to_owned(),
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: String::new(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.items[0].plan_type.as_deref(), expected);
|
||||
assert_eq!(
|
||||
result.items[0].plan_type_display.as_deref(),
|
||||
match expected {
|
||||
Some("free") => Some("OpenaiDisplayFree"),
|
||||
Some("plus") => Some("OpenaiDisplayPlus"),
|
||||
_ => None,
|
||||
}
|
||||
);
|
||||
let requests = provider.quota_requests();
|
||||
assert_eq!(requests.len(), usize::from(stored != Some("plus")));
|
||||
assert!(requests.iter().all(|request| !request.refresh));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_import_keeps_proxy_reserved_until_commit_and_releases_on_errors() {
|
||||
use super::accounts::{
|
||||
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded,
|
||||
};
|
||||
use gateway_admin::model::provider_credentials::ImportCredentials;
|
||||
use gateway_admin::ports::provider::ProviderAdminErrorKind;
|
||||
use std::sync::Arc;
|
||||
|
||||
for kind in ["openai", "xai"] {
|
||||
for failure in [None, Some("prepare"), Some("commit")] {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new(kind, events.clone());
|
||||
let store = FakeAccountStore::new(kind, events.clone());
|
||||
if failure == Some("prepare") {
|
||||
provider.fail_next(ProviderAdminErrorKind::Unavailable);
|
||||
}
|
||||
if failure == Some("commit") {
|
||||
store.fail_next_commit();
|
||||
}
|
||||
let services = super::AdminHarness::new()
|
||||
.provider(provider)
|
||||
.accounts(store)
|
||||
.proxies(Arc::new(TestProxies {
|
||||
events: Some(events.clone()),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let command = ImportCredentials {
|
||||
outbound_proxy_id: Some("proxy_import".to_owned()),
|
||||
settings: None,
|
||||
context: context("reserved-import"),
|
||||
document: document(),
|
||||
};
|
||||
let result = if kind == "openai" {
|
||||
services.openai().import_document(command).await
|
||||
} else {
|
||||
services.xai().import_document(command).await
|
||||
};
|
||||
assert_eq!(result.is_err(), failure.is_some());
|
||||
let events = recorded(&events);
|
||||
let expected = if failure == Some("prepare") {
|
||||
vec!["proxy.reserve", "provider.prepare_import", "proxy.release"]
|
||||
} else {
|
||||
vec![
|
||||
"proxy.reserve",
|
||||
"provider.prepare_import",
|
||||
"store.commit_import",
|
||||
"proxy.release",
|
||||
]
|
||||
};
|
||||
assert_eq!(&events[..expected.len()], expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,8 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
|
||||
services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: Some(super::accounts::import_settings()),
|
||||
context: context("import-xai"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -42,6 +44,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 +62,8 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
|
||||
services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-xai-batch"),
|
||||
document: document(),
|
||||
})
|
||||
@@ -90,6 +98,8 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
let result = services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
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,55 @@ 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 outbound_proxy_id: Option<String>,
|
||||
pub settings: Option<AccountImportSettingsRequest>,
|
||||
pub provider: String,
|
||||
pub data: Value,
|
||||
}
|
||||
|
||||
impl AccountImportRequest {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
if let Some(id) = &self.outbound_proxy_id {
|
||||
require_wire_id(id, "outboundProxyId")?;
|
||||
}
|
||||
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 +86,11 @@ impl AccountImportRequest {
|
||||
Ok((
|
||||
provider,
|
||||
ImportCredentials {
|
||||
outbound_proxy_id: self.outbound_proxy_id,
|
||||
settings: self
|
||||
.settings
|
||||
.map(AccountImportSettingsRequest::into_settings)
|
||||
.transpose()?,
|
||||
context,
|
||||
document: provider_document(self.data, "data")?,
|
||||
},
|
||||
@@ -57,9 +101,11 @@ impl AccountImportRequest {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct StartAccountAuthorizationRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub account_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
|
||||
}
|
||||
|
||||
impl StartAccountAuthorizationRequest {
|
||||
@@ -89,6 +135,10 @@ impl StartAccountAuthorizationRequest {
|
||||
context,
|
||||
name: self.name,
|
||||
reauthorization,
|
||||
outbound_proxy: super::wire::proxy_selection(
|
||||
self.outbound_proxy_id,
|
||||
self.outbound_proxy_url,
|
||||
)?,
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -97,6 +147,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 +155,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 +182,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 +197,8 @@ impl CompleteAccountAuthorizationRequest {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct UpdateAccountRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
|
||||
pub account_id: String,
|
||||
pub enabled: bool,
|
||||
#[serde(deserialize_with = "deserialize_required_nullable")]
|
||||
@@ -159,6 +219,10 @@ impl UpdateAccountRequest {
|
||||
pub(super) fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(UpdateAccount {
|
||||
outbound_proxy: super::wire::proxy_selection(
|
||||
self.outbound_proxy_id,
|
||||
self.outbound_proxy_url,
|
||||
)?,
|
||||
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,48 @@
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn proxy_selection(
|
||||
id: Option<String>,
|
||||
url: Option<AccountProxyUpdate>,
|
||||
) -> Result<Option<gateway_admin::model::proxies::AccountProxySelection>, WireValidationError> {
|
||||
use gateway_admin::model::proxies::AccountProxySelection;
|
||||
match (id, url) {
|
||||
(Some(_), Some(_)) => Err(WireValidationError::new("outboundProxyId")),
|
||||
(Some(id), None) if id.is_empty() => Ok(Some(AccountProxySelection::Direct)),
|
||||
(Some(id), None) => {
|
||||
if id.len() > 128
|
||||
|| !id
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
|
||||
{
|
||||
return Err(WireValidationError::new("outboundProxyId"));
|
||||
}
|
||||
Ok(Some(AccountProxySelection::Saved(id)))
|
||||
}
|
||||
(None, Some(value)) => Ok(Some(
|
||||
value
|
||||
.0
|
||||
.map_or(AccountProxySelection::Direct, AccountProxySelection::Url),
|
||||
)),
|
||||
(None, None) => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号列表查询参数。
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -19,6 +61,8 @@ pub struct ListQuery {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct BatchUpdateAccountsRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<AccountProxyUpdate>,
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
#[serde(deserialize_with = "deserialize_required_nullable")]
|
||||
@@ -68,6 +112,7 @@ impl BatchUpdateAccountsRequest {
|
||||
pub(super) fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(BatchUpdateAccounts {
|
||||
outbound_proxy: proxy_selection(self.outbound_proxy_id, self.outbound_proxy_url)?,
|
||||
account_ids: self.account_ids,
|
||||
enabled: self.enabled,
|
||||
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
|
||||
@@ -192,6 +237,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 +248,8 @@ pub struct AccountView {
|
||||
pub user_id: Option<String>,
|
||||
pub label: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
/// 后端生成的套餐展示名称;缺失套餐时为“未知套餐”。
|
||||
pub plan_type_display: String,
|
||||
pub authentication_kind: String,
|
||||
pub has_refresh_token: bool,
|
||||
pub status: String,
|
||||
|
||||
@@ -11,6 +11,8 @@ use gateway_admin::model::client_keys::{
|
||||
CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::budget::ClientBudgetLimits,
|
||||
metering::Decimal,
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
@@ -33,6 +35,15 @@ const MAX_CURSOR_BYTES: usize = 512;
|
||||
const MAX_SEARCH_BYTES: usize = 256;
|
||||
const DEFAULT_PAGE_SIZE: u16 = 50;
|
||||
|
||||
fn parse_budget(
|
||||
value: Option<String>,
|
||||
field: &'static str,
|
||||
) -> Result<Option<Decimal>, WireValidationError> {
|
||||
value
|
||||
.map(|value| value.parse().map_err(|_| WireValidationError::new(field)))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Client Key 列表查询。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -140,6 +151,8 @@ pub struct CreateClientKeyRequest {
|
||||
group_ids: Vec<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
daily_limit_usd: Option<String>,
|
||||
weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
impl CreateClientKeyRequest {
|
||||
@@ -154,6 +167,11 @@ impl CreateClientKeyRequest {
|
||||
name: self.name,
|
||||
label: self.label,
|
||||
group_ids,
|
||||
budget: ClientBudgetLimits {
|
||||
daily_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?.unwrap_or_default(),
|
||||
weekly_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")?
|
||||
.unwrap_or_default(),
|
||||
},
|
||||
limits: RateLimits {
|
||||
max_concurrency: self.max_concurrency,
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
@@ -172,6 +190,8 @@ pub struct UpdateClientKeyRequest {
|
||||
group_ids: Vec<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
daily_limit_usd: Option<String>,
|
||||
weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
impl UpdateClientKeyRequest {
|
||||
@@ -188,6 +208,8 @@ impl UpdateClientKeyRequest {
|
||||
name: self.name,
|
||||
label: self.label,
|
||||
group_ids,
|
||||
daily_limit_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?,
|
||||
weekly_limit_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")?,
|
||||
limits: RateLimits {
|
||||
max_concurrency: self.max_concurrency,
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
@@ -247,6 +269,12 @@ pub struct ClientKeyView {
|
||||
enabled: bool,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
daily_limit_usd: String,
|
||||
weekly_limit_usd: String,
|
||||
daily_used_usd: String,
|
||||
weekly_used_usd: String,
|
||||
daily_resets_at: Option<DateTime<Utc>>,
|
||||
weekly_resets_at: Option<DateTime<Utc>>,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
@@ -292,6 +320,12 @@ impl From<ClientKeyRecord> for ClientKeyView {
|
||||
enabled: record.enabled,
|
||||
max_concurrency: record.limits.max_concurrency,
|
||||
requests_per_minute: record.limits.requests_per_minute,
|
||||
daily_limit_usd: record.budget.limits.daily_usd.canonical(),
|
||||
weekly_limit_usd: record.budget.limits.weekly_usd.canonical(),
|
||||
daily_used_usd: record.budget.daily_used_usd.canonical(),
|
||||
weekly_used_usd: record.budget.weekly_used_usd.canonical(),
|
||||
daily_resets_at: record.budget.daily_resets_at.map(DateTime::from),
|
||||
weekly_resets_at: record.budget.weekly_resets_at.map(DateTime::from),
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
last_used_at: record.last_used_at,
|
||||
|
||||
@@ -16,6 +16,7 @@ pub mod client_keys;
|
||||
mod extract;
|
||||
pub mod observability;
|
||||
pub mod presenter;
|
||||
pub mod proxies;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
pub mod wire;
|
||||
@@ -34,6 +35,7 @@ where
|
||||
{
|
||||
Router::new()
|
||||
.merge(account_groups::router::<S>())
|
||||
.merge(proxies::router::<S>())
|
||||
.merge(accounts::router::<S>())
|
||||
.merge(auth::router::<S>())
|
||||
.merge(backups::router::<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>,
|
||||
|
||||
@@ -0,0 +1,452 @@
|
||||
use axum::{
|
||||
Router,
|
||||
extract::State,
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use gateway_admin::model::{
|
||||
PageSize, Revision,
|
||||
proxies::{
|
||||
NewProxy, ProxyAccountListQuery, ProxyListQuery, ProxyMutation, ProxyRecord,
|
||||
ProxyTestResult, UpdateProxy,
|
||||
},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{
|
||||
AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse, AdminSessionState,
|
||||
PageMeta,
|
||||
accounts::{AccountGroupRefView, AccountProxyUpdate},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ListQuery {
|
||||
page: Option<u32>,
|
||||
page_size: Option<u16>,
|
||||
search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct AccountsQuery {
|
||||
proxy_id: String,
|
||||
page: Option<u32>,
|
||||
page_size: Option<u16>,
|
||||
search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct RemoveAccountRequest {
|
||||
proxy_id: String,
|
||||
account_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct CreateRequest {
|
||||
name: String,
|
||||
proxy_url: AccountProxyUpdate,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct UpdateRequest {
|
||||
id: String,
|
||||
revision: u64,
|
||||
name: String,
|
||||
proxy_url: Option<AccountProxyUpdate>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct IdRequest {
|
||||
id: String,
|
||||
revision: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ProbeRequest {
|
||||
proxy_url: AccountProxyUpdate,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyTestView {
|
||||
success: bool,
|
||||
latency_ms: u64,
|
||||
exit_ip: Option<String>,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl From<ProxyTestResult> for ProxyTestView {
|
||||
fn from(result: ProxyTestResult) -> Self {
|
||||
Self {
|
||||
success: result.success,
|
||||
latency_ms: result.latency_ms,
|
||||
exit_ip: result.exit_ip.map(|ip| ip.to_string()),
|
||||
message: result.message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyAccountView {
|
||||
id: String,
|
||||
name: String,
|
||||
email: Option<String>,
|
||||
provider: String,
|
||||
authentication_kind: String,
|
||||
plan_type: Option<String>,
|
||||
plan_type_display: Option<String>,
|
||||
groups: Vec<AccountGroupRefView>,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyView {
|
||||
id: String,
|
||||
name: String,
|
||||
endpoint: String,
|
||||
has_authentication: bool,
|
||||
revision: u64,
|
||||
account_count: u64,
|
||||
last_test_at: Option<String>,
|
||||
last_test: Option<ProxyTestView>,
|
||||
created_at: String,
|
||||
updated_at: String,
|
||||
}
|
||||
|
||||
impl From<ProxyRecord> for ProxyView {
|
||||
fn from(record: ProxyRecord) -> Self {
|
||||
let endpoint = record.proxy.endpoint();
|
||||
Self {
|
||||
id: record.id,
|
||||
name: record.name,
|
||||
has_authentication: record.proxy.expose_url() != endpoint,
|
||||
endpoint,
|
||||
revision: record.revision.get(),
|
||||
account_count: record.account_count,
|
||||
last_test_at: record.last_test_at.map(|at| at.to_rfc3339()),
|
||||
last_test: record.last_test.map(Into::into),
|
||||
created_at: record.created_at.to_rfc3339(),
|
||||
updated_at: record.updated_at.to_rfc3339(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ProxyPageView {
|
||||
items: Vec<ProxyView>,
|
||||
page: PageMeta,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ProxyAccountPageView {
|
||||
items: Vec<ProxyAccountView>,
|
||||
page: PageMeta,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct MutationView {
|
||||
record: ProxyView,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<ProxyMutation> for MutationView {
|
||||
fn from(value: ProxyMutation) -> Self {
|
||||
Self {
|
||||
record: value.record.into(),
|
||||
config_revision: value.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
S: AdminSessionState + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Router::new()
|
||||
.route("/api/admin/proxies", get(list::<S>))
|
||||
.route("/api/admin/proxies/accounts", get(list_accounts::<S>))
|
||||
.route(
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
post(remove_account::<S>),
|
||||
)
|
||||
.route("/api/admin/proxies/create", post(create::<S>))
|
||||
.route("/api/admin/proxies/update", post(update::<S>))
|
||||
.route("/api/admin/proxies/delete", post(delete::<S>))
|
||||
.route("/api/admin/proxies/test", post(test::<S>))
|
||||
.route("/api/admin/proxies/probe", post(probe::<S>))
|
||||
}
|
||||
|
||||
fn revision(value: u64) -> Result<Revision, AdminError> {
|
||||
if value > i64::MAX as u64 {
|
||||
return Err(AdminError::bad_request("代理版本不合法"));
|
||||
}
|
||||
Revision::new(value).map_err(|_| AdminError::bad_request("代理版本不合法"))
|
||||
}
|
||||
|
||||
fn map_error(error: gateway_admin::model::AdminError) -> AdminError {
|
||||
super::wire::map_admin_service_error(error)
|
||||
}
|
||||
|
||||
async fn list<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<ListQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.list(ProxyListQuery {
|
||||
page: query.page.unwrap_or(1),
|
||||
page_size: PageSize::new(query.page_size.unwrap_or(20))
|
||||
.map_err(|_| AdminError::bad_request("分页大小不合法"))?,
|
||||
search: query.search.unwrap_or_default().trim().to_owned(),
|
||||
})
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
let total_pages = result.total.div_ceil(u64::from(result.page_size));
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyPageView {
|
||||
items: result.items.into_iter().map(Into::into).collect(),
|
||||
page: PageMeta::new(
|
||||
result.page,
|
||||
u32::from(result.page_size),
|
||||
result.total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_accounts<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<AccountsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: query.proxy_id,
|
||||
page: query.page.unwrap_or(1),
|
||||
page_size: PageSize::new(query.page_size.unwrap_or(20))
|
||||
.map_err(|_| AdminError::bad_request("分页大小不合法"))?,
|
||||
search: query.search.unwrap_or_default().trim().to_owned(),
|
||||
})
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
let total_pages = result.total.div_ceil(u64::from(result.page_size));
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyAccountPageView {
|
||||
items: result
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|account| ProxyAccountView {
|
||||
id: account.id,
|
||||
name: account.name,
|
||||
email: account.email,
|
||||
provider: account.provider_kind,
|
||||
authentication_kind: account.authentication_kind,
|
||||
plan_type: account.plan_type,
|
||||
plan_type_display: account.plan_type_display,
|
||||
groups: account
|
||||
.groups
|
||||
.into_iter()
|
||||
.map(|group| AccountGroupRefView {
|
||||
id: group.id.to_string(),
|
||||
name: group.name,
|
||||
color: group.color.as_str().to_owned(),
|
||||
enabled: group.enabled,
|
||||
})
|
||||
.collect(),
|
||||
enabled: account.enabled,
|
||||
})
|
||||
.collect(),
|
||||
page: PageMeta::new(
|
||||
result.page,
|
||||
u32::from(result.page_size),
|
||||
result.total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn create<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<CreateRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.create(
|
||||
NewProxy {
|
||||
name: request.name,
|
||||
proxy,
|
||||
},
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::CREATED,
|
||||
AdminEnvelope::ok(MutationView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn remove_account<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<RemoveAccountRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let revision = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.remove_account(
|
||||
&request.proxy_id,
|
||||
&request.account_id,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"configRevision": revision.get()})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn update<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<UpdateRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.map(|value| {
|
||||
value
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))
|
||||
})
|
||||
.transpose()?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.update(
|
||||
UpdateProxy {
|
||||
id: request.id,
|
||||
revision: revision(request.revision)?,
|
||||
name: request.name,
|
||||
proxy,
|
||||
},
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(MutationView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn delete<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<IdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.delete(
|
||||
&request.id,
|
||||
revision(request.revision)?,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"configRevision": result.get()})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn probe<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<ProbeRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.probe(&proxy)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyTestView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn test<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<IdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.test(
|
||||
&request.id,
|
||||
revision(request.revision)?,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyView::from(result)),
|
||||
))
|
||||
}
|
||||
@@ -143,13 +143,27 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
|
||||
/// Gateway 错误的 OpenAI HTTP 表达。
|
||||
pub fn gateway_error_response(error: &GatewayError) -> Response {
|
||||
let (status, default_type, default_code) = gateway_error_contract(error.kind());
|
||||
openai_error_response(
|
||||
let mut response = openai_error_response(
|
||||
status,
|
||||
error.client_message(),
|
||||
error.client_error_type().unwrap_or(default_type),
|
||||
error.client_error_code().unwrap_or(default_code),
|
||||
)
|
||||
.into_response()
|
||||
.into_response();
|
||||
if let Some(delay) = error.retry_after()
|
||||
&& let Ok(value) = HeaderValue::from_str(
|
||||
&delay
|
||||
.as_secs()
|
||||
.saturating_add(u64::from(delay.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
.to_string(),
|
||||
)
|
||||
{
|
||||
response
|
||||
.headers_mut()
|
||||
.insert(axum::http::header::RETRY_AFTER, value);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
/// 在下游尚未提交时优先交付 Provider 的原始 HTTP 失败响应。
|
||||
|
||||
@@ -756,3 +756,60 @@ mod actions {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
mod import_settings {
|
||||
use gateway_api::admin::accounts::{AccountImportRequest, CompleteAccountAuthorizationRequest};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn import_and_oauth_apply_the_same_settings_validation() {
|
||||
for (field, value) in [
|
||||
("concurrencyLimit", json!(0)),
|
||||
("concurrencyLimit", json!(4_294_967_296_u64)),
|
||||
("weight", json!(0)),
|
||||
("weight", json!(101)),
|
||||
("groupIds", json!(["invalid-group"])),
|
||||
] {
|
||||
let mut settings =
|
||||
json!({"enabled": false, "concurrencyLimit": null, "weight": 1, "groupIds": []});
|
||||
settings[field] = value;
|
||||
let import: AccountImportRequest = serde_json::from_value(
|
||||
json!({"provider": "openai", "data": {}, "settings": settings}),
|
||||
)
|
||||
.expect("import request");
|
||||
let oauth: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({"provider": "xai", "flowId": "flow-test", "callbackUrl": "code", "settings": settings})).expect("OAuth request");
|
||||
assert_eq!(
|
||||
import.validate().expect_err("invalid settings").field(),
|
||||
field
|
||||
);
|
||||
assert_eq!(
|
||||
oauth.validate().expect_err("invalid settings").field(),
|
||||
field
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn import_settings_require_a_complete_explicit_configuration() {
|
||||
let settings =
|
||||
json!({"enabled": false, "concurrencyLimit": null, "weight": 100, "groupIds": []});
|
||||
let request: AccountImportRequest =
|
||||
serde_json::from_value(json!({"provider": "openai", "data": {}, "settings": settings}))
|
||||
.expect("settings");
|
||||
assert!(request.validate().is_ok());
|
||||
for field in ["enabled", "concurrencyLimit", "weight", "groupIds"] {
|
||||
let mut incomplete = settings.clone();
|
||||
incomplete
|
||||
.as_object_mut()
|
||||
.expect("settings object")
|
||||
.remove(field);
|
||||
assert!(
|
||||
serde_json::from_value::<AccountImportRequest>(
|
||||
json!({"provider": "openai", "data": {}, "settings": incomplete})
|
||||
)
|
||||
.is_err(),
|
||||
"missing {field}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,46 @@ use serde_json::json;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState};
|
||||
|
||||
#[test]
|
||||
fn budget_inputs_preserve_decimal_precision_and_omitted_updates() {
|
||||
let payload = json!({"name": "budget", "groupIds": [], "maxConcurrency": 3,
|
||||
"requestsPerMinute": 0, "dailyLimitUsd": "0.1234567891", "weeklyLimitUsd": "15"});
|
||||
let command = serde_json::from_value::<CreateClientKeyRequest>(payload.clone())
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap();
|
||||
assert_eq!(command.budget.daily_usd.canonical(), "0.1234567891");
|
||||
assert_eq!(command.budget.weekly_usd.canonical(), "15");
|
||||
assert_eq!(command.limits.max_concurrency, 3);
|
||||
for field in ["dailyLimitUsd", "weeklyLimitUsd"] {
|
||||
for invalid in ["-1", "NaN", "1e3", "10000000000", "0.00000000001", ""] {
|
||||
let mut payload = payload.clone();
|
||||
payload[field] = json!(invalid);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<CreateClientKeyRequest>(payload)
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap_err()
|
||||
.field(),
|
||||
field
|
||||
);
|
||||
}
|
||||
}
|
||||
let mut update = payload;
|
||||
update["id"] = json!("key_budget");
|
||||
update.as_object_mut().unwrap().remove("dailyLimitUsd");
|
||||
update["weeklyLimitUsd"] = json!("0");
|
||||
let command = serde_json::from_value::<UpdateClientKeyRequest>(update)
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap();
|
||||
assert_eq!(command.daily_limit_usd, None);
|
||||
assert_eq!(
|
||||
command.weekly_limit_usd,
|
||||
Some(gateway_core::metering::Decimal::ZERO)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
|
||||
let unknown = serde_json::from_value::<ListClientKeysQuery>(json!({ "other": true }));
|
||||
@@ -207,6 +247,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
|
||||
.single()
|
||||
.expect("valid time");
|
||||
let view = ClientKeyView::from(gateway_admin::model::client_keys::ClientKeyRecord {
|
||||
budget: Default::default(),
|
||||
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
|
||||
name: "visible".to_owned(),
|
||||
label: None,
|
||||
@@ -238,6 +279,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
|
||||
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
|
||||
assert!(list["items"][0].get("tokensPerMinute").is_none());
|
||||
assert!(list["items"][0].get("policyJson").is_none());
|
||||
assert_eq!(list["items"][0]["dailyLimitUsd"], "0");
|
||||
assert_eq!(list["items"][0]["weeklyLimitUsd"], "0");
|
||||
assert_eq!(list["items"][0]["dailyUsedUsd"], "0");
|
||||
assert!(list["items"][0].get("unresolvedRequests").is_none());
|
||||
assert_eq!(
|
||||
DateTime::parse_from_rfc3339(
|
||||
list["items"][0]["lastUsedAt"]
|
||||
|
||||
@@ -82,6 +82,7 @@ mod auth;
|
||||
mod client_keys;
|
||||
mod errors;
|
||||
mod observability;
|
||||
mod proxies;
|
||||
mod settings;
|
||||
mod system;
|
||||
mod wire;
|
||||
@@ -124,7 +125,12 @@ impl AdminTestFixture {
|
||||
dashboard_summary_range: Arc::clone(&dashboard_summary_range),
|
||||
});
|
||||
let stores = AdminStorePorts::new(
|
||||
AdminAccountStorePorts::new(unused.clone(), unused.clone(), account_groups.clone()),
|
||||
AdminAccountStorePorts::new(
|
||||
unused.clone(),
|
||||
unused.clone(),
|
||||
account_groups.clone(),
|
||||
Arc::new(proxies::MemoryProxies::default()),
|
||||
),
|
||||
auth.clone(),
|
||||
client_keys.clone(),
|
||||
unused,
|
||||
@@ -144,7 +150,7 @@ impl AdminTestFixture {
|
||||
stores,
|
||||
providers,
|
||||
Arc::new(NoopSnapshot),
|
||||
Arc::new(NoopProbe),
|
||||
(Arc::new(NoopProbe), Arc::new(proxies::SuccessfulProbe)),
|
||||
Arc::new(StaticClientDistribution),
|
||||
system,
|
||||
)
|
||||
@@ -640,6 +646,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 {
|
||||
|
||||
@@ -0,0 +1,520 @@
|
||||
use std::sync::Mutex;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use gateway_admin::{
|
||||
model::{MutationContext, Revision, proxies::*},
|
||||
ports::{
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
|
||||
},
|
||||
};
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
use serde_json::{Value, json};
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState};
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct MemoryProxies(Mutex<Option<ProxyRecord>>);
|
||||
|
||||
fn missing() -> AdminStoreError {
|
||||
AdminStoreError::new(AdminStoreErrorKind::NotFound, "proxy", "missing proxy")
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyStore for MemoryProxies {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &ProviderAccountId,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
if proxy_id != "proxy_test" || account_id.as_str() != "acct_linked" {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Conflict,
|
||||
"proxy account",
|
||||
"changed binding",
|
||||
));
|
||||
}
|
||||
Ok(Revision::new(4).unwrap())
|
||||
}
|
||||
|
||||
async fn reserve_import(
|
||||
&self,
|
||||
_: &str,
|
||||
) -> AdminStoreResult<gateway_admin::ports::proxy::ProxyImportReservation> {
|
||||
Err(missing())
|
||||
}
|
||||
async fn list(&self, query: ProxyListQuery) -> AdminStoreResult<ProxyPage> {
|
||||
let items: Vec<_> = self.0.lock().unwrap().iter().cloned().collect();
|
||||
Ok(ProxyPage {
|
||||
total: u64::try_from(items.len()).unwrap(),
|
||||
items,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn get(&self, _: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
self.0.lock().unwrap().clone().ok_or_else(missing)
|
||||
}
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage> {
|
||||
if !self
|
||||
.0
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.is_some_and(|record| record.id == query.proxy_id)
|
||||
{
|
||||
return Err(missing());
|
||||
}
|
||||
Ok(ProxyAccountPage {
|
||||
items: vec![ProxyAccountRef {
|
||||
id: "acct_linked".to_owned(),
|
||||
name: "工作账号".to_owned(),
|
||||
email: Some("work@example.invalid".to_owned()),
|
||||
provider_kind: "openai".to_owned(),
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
plan_type: Some("plus".to_owned()),
|
||||
plan_type_display: None,
|
||||
groups: vec![],
|
||||
enabled: true,
|
||||
}],
|
||||
total: 21,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let record = ProxyRecord {
|
||||
id: "proxy_test".to_owned(),
|
||||
name: command.name,
|
||||
proxy: command.proxy,
|
||||
revision: Revision::new(1).unwrap(),
|
||||
account_count: 0,
|
||||
last_test: None,
|
||||
last_test_at: None,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
};
|
||||
*self.0.lock().unwrap() = Some(record.clone());
|
||||
Ok(ProxyMutation {
|
||||
config_revision: record.revision,
|
||||
record,
|
||||
})
|
||||
}
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let mut stored = self.0.lock().unwrap();
|
||||
let record = stored.as_mut().ok_or_else(missing)?;
|
||||
record.name = command.name;
|
||||
if let Some(proxy) = command.proxy {
|
||||
record.proxy = proxy;
|
||||
}
|
||||
record.revision = Revision::new(record.revision.get() + 1).unwrap();
|
||||
Ok(ProxyMutation {
|
||||
config_revision: record.revision,
|
||||
record: record.clone(),
|
||||
})
|
||||
}
|
||||
async fn delete(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
self.0.lock().unwrap().take().ok_or_else(missing)?;
|
||||
Ok(Revision::new(3).unwrap())
|
||||
}
|
||||
async fn record_test(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
result: ProxyTestResult,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord> {
|
||||
let mut stored = self.0.lock().unwrap();
|
||||
let record = stored.as_mut().ok_or_else(missing)?;
|
||||
record.last_test = Some(result);
|
||||
record.last_test_at = Some(Utc::now());
|
||||
Ok(record.clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct SuccessfulProbe;
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyProbe for SuccessfulProbe {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult {
|
||||
assert_eq!(
|
||||
proxy.expose_url(),
|
||||
"http://test-user:private-password@proxy.example:8080/"
|
||||
);
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms: 15,
|
||||
exit_ip: Some("203.0.113.2".parse().unwrap()),
|
||||
message: "Connected".to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn request(
|
||||
fixture: &AdminTestFixture,
|
||||
path: &str,
|
||||
body: Option<Value>,
|
||||
authenticated: bool,
|
||||
) -> (StatusCode, Value) {
|
||||
let mut builder = Request::builder()
|
||||
.header("x-request-id", "req_proxy_tests")
|
||||
.uri(path)
|
||||
.method(if body.is_some() { "POST" } else { "GET" });
|
||||
if authenticated {
|
||||
builder = builder.header(header::COOKIE, "cpr_admin_session=valid-session");
|
||||
}
|
||||
let body = body.map_or_else(Body::empty, |body| Body::from(body.to_string()));
|
||||
let response = gateway_api::admin::proxies::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(
|
||||
builder
|
||||
.header(header::CONTENT_TYPE, "application/json")
|
||||
.body(body)
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
|
||||
let text = std::str::from_utf8(&body).unwrap();
|
||||
assert!(!text.contains("private-password"));
|
||||
assert!(!text.contains("test-user"));
|
||||
(status, serde_json::from_slice(&body).unwrap())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_routes_save_reload_test_rename_and_delete_without_exposing_credentials() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let (status, created) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({
|
||||
"name": " Office ", "proxyUrl": "http://test-user:private-password@proxy.example:8080"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED);
|
||||
assert_eq!(created["data"]["record"]["name"], "Office");
|
||||
assert_eq!(created["data"]["record"]["hasAuthentication"], true);
|
||||
assert_eq!(
|
||||
created["data"]["record"]["endpoint"],
|
||||
"http://proxy.example:8080/"
|
||||
);
|
||||
assert!(created["data"]["record"].get("proxyUrl").is_none());
|
||||
assert_eq!(created["data"]["record"]["accountCount"], 0);
|
||||
assert!(created["data"]["record"].get("accounts").is_none());
|
||||
|
||||
let (status, linked) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=proxy_test&page=2&pageSize=20&search=work",
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
linked["data"]["page"],
|
||||
json!({"page":2,"pageSize":20,"total":21,"totalPages":2})
|
||||
);
|
||||
assert_eq!(
|
||||
linked["data"]["items"],
|
||||
json!([{
|
||||
"id":"acct_linked", "name":"工作账号", "email":"work@example.invalid",
|
||||
"provider":"openai", "enabled":true, "authenticationKind":"oauth",
|
||||
"planType":"plus", "planTypeDisplay":"Plus", "groups":[]
|
||||
}])
|
||||
);
|
||||
|
||||
let (status, tested) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 1})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(tested["data"]["lastTest"]["exitIp"], "203.0.113.2");
|
||||
let (_, listed) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies?page=1&pageSize=20",
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(listed["data"]["items"][0], tested["data"]);
|
||||
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/update",
|
||||
Some(json!({"id": "proxy_test", "revision": 1, "name": "Renamed"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 1})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CONFLICT);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 2})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/delete",
|
||||
Some(json!({"id": "proxy_test", "revision": 2})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["page"]["total"], 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_checks_unsaved_address_without_creating_or_changing_records() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let draft = json!({
|
||||
"proxyUrl": "http://test-user:private-password@proxy.example:8080"
|
||||
});
|
||||
let (status, probed) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/probe",
|
||||
Some(draft.clone()),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
probed["data"],
|
||||
json!({"success": true, "latencyMs": 15, "exitIp": "203.0.113.2", "message": "Connected"})
|
||||
);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["page"]["total"], 0);
|
||||
|
||||
let (_, created) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({"name": "Saved", "proxyUrl": "http://saved.example:8080"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
let (status, _) = request(&fixture, "/api/admin/proxies/probe", Some(draft), true).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["items"][0], created["data"]["record"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_rejects_empty_or_invalid_addresses() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (body, expected) in [
|
||||
(json!({"proxyUrl": ""}), StatusCode::BAD_REQUEST),
|
||||
(json!({"proxyUrl": null}), StatusCode::UNPROCESSABLE_ENTITY),
|
||||
(json!({}), StatusCode::UNPROCESSABLE_ENTITY),
|
||||
(
|
||||
json!({"proxyUrl": "ftp://test-user:private-password@proxy.example"}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
(
|
||||
json!({"proxyUrl": "http://proxy.example:8080", "name": "Not saved"}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies/probe", Some(body), true)
|
||||
.await
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_routes_require_auth_and_reject_invalid_input() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies", None, false).await.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=proxy_test",
|
||||
None,
|
||||
false
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
for query in [
|
||||
"",
|
||||
"proxyId=",
|
||||
"proxyId=proxy_test&page=0",
|
||||
"proxyId=proxy_test&pageSize=0",
|
||||
"proxyId=proxy_test&pageSize=201",
|
||||
"proxyId=proxy_test&search=%00",
|
||||
"proxyId=proxy_test&unknown=1",
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
&format!("/api/admin/proxies/accounts?{query}"),
|
||||
None,
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=missing",
|
||||
None,
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::NOT_FOUND
|
||||
);
|
||||
for action in [
|
||||
"create",
|
||||
"update",
|
||||
"delete",
|
||||
"test",
|
||||
"probe",
|
||||
"accounts/remove",
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
&format!("/api/admin/proxies/{action}"),
|
||||
Some(json!({})),
|
||||
false
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
}
|
||||
for query in ["page=0", "pageSize=0", "pageSize=201", "unknown=1"] {
|
||||
assert_eq!(
|
||||
request(&fixture, &format!("/api/admin/proxies?{query}"), None, true)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
for body in [
|
||||
json!({"name":" ","proxyUrl":"http://proxy.example:8080"}),
|
||||
json!({"name":"Invalid","proxyUrl":""}),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies/create", Some(body), true)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({"name":"Invalid","proxyUrl":"ftp://test-user:private-password@proxy.example"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_account_removal_validates_binding_and_input() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (body, expected) in [
|
||||
(
|
||||
json!({"proxyId":"", "accountId":"acct_linked"}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":""}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":"invalid"}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":"acct_linked", "enabled":false}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(body),
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
}
|
||||
let (status, result) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(json!({"proxyId":"proxy_test", "accountId":"acct_linked"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(result["data"], json!({"configRevision":4}));
|
||||
let (status, result) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(json!({"proxyId":"proxy_previous", "accountId":"acct_linked"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CONFLICT);
|
||||
assert_eq!(result["message"], "账号的代理绑定已变化,请刷新后重试");
|
||||
}
|
||||
@@ -50,6 +50,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
|
||||
"src/admin/observability/routes.rs",
|
||||
"src/admin/observability/wire.rs",
|
||||
"src/admin/presenter.rs",
|
||||
"src/admin/proxies.rs",
|
||||
"src/admin/settings.rs",
|
||||
"src/admin/system.rs",
|
||||
"src/admin/wire.rs",
|
||||
@@ -96,6 +97,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/admin/observability/mod.rs",
|
||||
"tests/admin/observability/query.rs",
|
||||
"tests/admin/observability/response.rs",
|
||||
"tests/admin/proxies.rs",
|
||||
"tests/admin/settings.rs",
|
||||
"tests/admin/system.rs",
|
||||
"tests/admin/wire.rs",
|
||||
|
||||
@@ -11,6 +11,21 @@ use gateway_api::openai::error::{
|
||||
openai_error_response,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn key_budget_errors_preserve_limit_code_and_retry_after() {
|
||||
for code in ["key_daily_budget_exceeded", "key_weekly_budget_exceeded"] {
|
||||
let error = GatewayError::new(GatewayErrorKind::RateLimited, "key budget exhausted")
|
||||
.with_client_code(code)
|
||||
.with_retry_after(std::time::Duration::from_millis(1501));
|
||||
let response = gateway_error_response(&error);
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(response.headers()["retry-after"], "2");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap();
|
||||
assert_eq!(body["error"]["code"], code);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_request_error_should_map_to_openai_bad_request() {
|
||||
assert_eq!(
|
||||
@@ -433,3 +448,151 @@ async fn continuation_recovery_should_preserve_the_official_retry_signal() {
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
mod model_routing {
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::execution::{
|
||||
DefaultExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
|
||||
};
|
||||
use gateway_core::engine::provider::ProviderRegistry;
|
||||
use gateway_core::routing::{ProviderKind, RuntimeSnapshot};
|
||||
use gateway_core::runtime::RuntimeSnapshotHandle;
|
||||
use serde_json::{Value, json};
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::openai::{
|
||||
IgnoredClientApiKeyUsage, UnusedAdmissions, UnusedContinuation, UnusedExecutionStore,
|
||||
};
|
||||
|
||||
struct TestCircuits {
|
||||
blocked: bool,
|
||||
}
|
||||
|
||||
impl ProviderCircuitPort for TestCircuits {
|
||||
fn decision<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
|
||||
Box::pin(async move {
|
||||
Ok(if self.blocked {
|
||||
ProviderCircuitDecision::BlockedUntil(
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
)
|
||||
} else {
|
||||
ProviderCircuitDecision::Allow
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn observe_failure<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
|
||||
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
|
||||
}
|
||||
|
||||
fn observe_success<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
|
||||
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
|
||||
}
|
||||
}
|
||||
|
||||
async fn request_model(
|
||||
snapshot: RuntimeSnapshot,
|
||||
model: &str,
|
||||
stream: bool,
|
||||
blocked: bool,
|
||||
) -> (StatusCode, Value) {
|
||||
let execution = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot),
|
||||
Arc::new(UnusedExecutionStore),
|
||||
ProviderRegistry::default(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(TestCircuits { blocked }),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(IgnoredClientApiKeyUsage),
|
||||
);
|
||||
let response = crate::openai::api_router(Arc::new(execution))
|
||||
.await
|
||||
.oneshot(
|
||||
Request::post("/v1/responses")
|
||||
.header("authorization", "Bearer sk_model_routing")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({"model": model, "input": "hello", "stream": stream}).to_string(),
|
||||
))
|
||||
.expect("request"),
|
||||
)
|
||||
.await
|
||||
.expect("response");
|
||||
let status = response.status();
|
||||
let body = to_bytes(response.into_body(), 16 * 1024)
|
||||
.await
|
||||
.expect("error body");
|
||||
(status, serde_json::from_slice(&body).expect("error JSON"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_models_should_return_404_with_a_specific_explanation() {
|
||||
for (model, mappings, message) in [
|
||||
(
|
||||
"missing-model",
|
||||
BTreeMap::new(),
|
||||
"the requested model was not found in the provider catalogs available to this API key; check the model name",
|
||||
),
|
||||
(
|
||||
"model-a",
|
||||
BTreeMap::from([("model-a".to_owned(), "missing-target".to_owned())]),
|
||||
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping",
|
||||
),
|
||||
] {
|
||||
for stream in [false, true] {
|
||||
let snapshot = crate::openai::snapshot("sk_model_routing", "openai")
|
||||
.with_model_mappings(mappings.clone());
|
||||
let response = request_model(snapshot, model, stream, false).await;
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
json!({"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "model_not_found",
|
||||
"message": message,
|
||||
}}),
|
||||
),
|
||||
"model={model}, stream={stream}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_existing_model_with_a_blocked_provider_should_remain_service_unavailable() {
|
||||
let response = request_model(
|
||||
crate::openai::snapshot("sk_model_routing", "openai"),
|
||||
"model-a",
|
||||
true,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
(response.0, response.1["error"]["code"].as_str()),
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Some("no_available_provider")
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
name = "gateway-core"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
license = "MIT"
|
||||
license = "Apache-2.0"
|
||||
publish = false
|
||||
rust-version = "1.97"
|
||||
autotests = false
|
||||
@@ -24,6 +24,7 @@ semver = "1.0.27"
|
||||
thiserror = "2.0.18"
|
||||
tracing = "0.1.44"
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
url = "2.5.8"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
mod error;
|
||||
mod model;
|
||||
mod proxy;
|
||||
pub use proxy::{InvalidOutboundProxy, OutboundProxy};
|
||||
pub mod scope;
|
||||
mod selection;
|
||||
mod store;
|
||||
|
||||
@@ -654,7 +654,7 @@ const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
|
||||
}
|
||||
}
|
||||
|
||||
/// 不含 secret 的账号持久事实。
|
||||
/// 账号持久事实;代理认证信息只通过显式 secret accessor 读取。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderAccount {
|
||||
id: ProviderAccountId,
|
||||
@@ -676,6 +676,7 @@ pub struct ProviderAccount {
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
has_refresh_token: bool,
|
||||
outbound_proxy: Option<super::OutboundProxy>,
|
||||
}
|
||||
|
||||
impl ProviderAccount {
|
||||
@@ -710,6 +711,7 @@ impl ProviderAccount {
|
||||
access_token_expires_at,
|
||||
next_refresh_at: None,
|
||||
has_refresh_token: false,
|
||||
outbound_proxy: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -726,6 +728,17 @@ impl ProviderAccount {
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_outbound_proxy(mut self, proxy: Option<super::OutboundProxy>) -> Self {
|
||||
self.outbound_proxy = proxy;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn outbound_proxy(&self) -> Option<&super::OutboundProxy> {
|
||||
self.outbound_proxy.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_account_facts(
|
||||
mut self,
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
//! Explicit account egress. Credentials never appear in Debug or ordinary admin projections.
|
||||
|
||||
use std::fmt;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub struct OutboundProxy(Url);
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("invalid outbound proxy; expected http, https, socks5 or socks5h URL with a host and port")]
|
||||
pub struct InvalidOutboundProxy;
|
||||
|
||||
impl OutboundProxy {
|
||||
pub fn parse(value: &str) -> Result<Self, InvalidOutboundProxy> {
|
||||
if value.len() > 4096 || value.chars().any(char::is_control) {
|
||||
return Err(InvalidOutboundProxy);
|
||||
}
|
||||
let url = Url::parse(value).map_err(|_| InvalidOutboundProxy)?;
|
||||
if !matches!(url.scheme(), "http" | "https" | "socks5" | "socks5h")
|
||||
|| url.host_str().is_none()
|
||||
|| url.port_or_known_default().is_none_or(|port| port == 0)
|
||||
|| !matches!(url.path(), "" | "/")
|
||||
|| url.query().is_some()
|
||||
|| url.fragment().is_some()
|
||||
{
|
||||
return Err(InvalidOutboundProxy);
|
||||
}
|
||||
Ok(Self(url))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn expose_url(&self) -> &str {
|
||||
self.0.as_str()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn endpoint(&self) -> String {
|
||||
let mut url = self.0.clone();
|
||||
let _ = url.set_username("");
|
||||
let _ = url.set_password(None);
|
||||
url.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for OutboundProxy {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("OutboundProxy(<redacted>)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
//! 按 Key 持久化 USD 限额与已取得费用,不依赖尽力写入的请求观测。
|
||||
|
||||
use std::time::SystemTime;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::{error::GatewayError, metering::Decimal, policy::ClientApiKeyId};
|
||||
|
||||
use super::ModelRequestId;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct ClientBudgetLimits {
|
||||
pub daily_usd: Decimal,
|
||||
pub weekly_usd: Decimal,
|
||||
}
|
||||
|
||||
impl ClientBudgetLimits {
|
||||
#[must_use]
|
||||
pub fn is_limited(self) -> bool {
|
||||
self.daily_usd != Decimal::ZERO || self.weekly_usd != Decimal::ZERO
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ClientBudgetStatus {
|
||||
pub limits: ClientBudgetLimits,
|
||||
pub daily_used_usd: Decimal,
|
||||
pub weekly_used_usd: Decimal,
|
||||
pub daily_resets_at: Option<SystemTime>,
|
||||
pub weekly_resets_at: Option<SystemTime>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientBudgetCharge {
|
||||
pub key_id: ClientApiKeyId,
|
||||
pub request_id: ModelRequestId,
|
||||
/// 已取得的 USD 费用,包含重试;缺少费用的尝试按零累计。
|
||||
pub amount_usd: Decimal,
|
||||
pub completed_at: SystemTime,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, thiserror::Error)]
|
||||
#[error("client budget store is unavailable")]
|
||||
pub struct ClientBudgetError;
|
||||
|
||||
pub trait ClientBudgetPort: Send + Sync {
|
||||
/// 原子检查当前限额与已用金额,不创建预扣费或待结算记录。
|
||||
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>>;
|
||||
|
||||
/// 按网关请求 ID 幂等累计已取得费用;写入失败由 Store 重试。
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>;
|
||||
}
|
||||
@@ -30,6 +30,7 @@ use crate::event::{
|
||||
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
|
||||
};
|
||||
use crate::lifecycle::CancellationToken;
|
||||
use crate::metering::Decimal;
|
||||
use crate::operation::{Operation, ProviderSessionState};
|
||||
use crate::routing::RoutingPlan;
|
||||
use futures::future::Fuse;
|
||||
@@ -166,6 +167,8 @@ where
|
||||
request_id,
|
||||
client_api_key_ref,
|
||||
observation: ResponseObservation::new(timing_started_at),
|
||||
budget_prior_attempts_usd: Decimal::ZERO,
|
||||
budget_attempt_already_counted: false,
|
||||
trace,
|
||||
deadline,
|
||||
deadline_timer: Delay::new(
|
||||
@@ -197,6 +200,7 @@ where
|
||||
delivery_pending: false,
|
||||
upstream_complete: false,
|
||||
finalized: false,
|
||||
finalized_at: None,
|
||||
image_generation_requested,
|
||||
last_retryable_failure: None,
|
||||
last_retryable_failure_events: Vec::new(),
|
||||
@@ -253,6 +257,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
request_id: ModelRequestId,
|
||||
client_api_key_ref: crate::policy::ClientApiKeyId,
|
||||
observation: ResponseObservation,
|
||||
budget_prior_attempts_usd: Decimal,
|
||||
budget_attempt_already_counted: bool,
|
||||
trace: TraceContext,
|
||||
deadline: SystemTime,
|
||||
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
|
||||
@@ -288,6 +294,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
delivery_pending: bool,
|
||||
upstream_complete: bool,
|
||||
finalized: bool,
|
||||
finalized_at: Option<SystemTime>,
|
||||
image_generation_requested: bool,
|
||||
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
|
||||
last_retryable_failure: Option<ProviderError>,
|
||||
@@ -475,6 +482,35 @@ where
|
||||
self.finalized
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge {
|
||||
// 超出数据库可表示范围时保留最大金额,避免溢出后误记为零。
|
||||
let amount_usd = self
|
||||
.budget_prior_attempts_usd
|
||||
.checked_add(self.budget_attempt_usd())
|
||||
.unwrap_or(Decimal::MAX);
|
||||
super::budget::ClientBudgetCharge {
|
||||
key_id: self.client_api_key_ref.clone(),
|
||||
request_id: self.request_id.clone(),
|
||||
amount_usd,
|
||||
completed_at: self.finalized_at.unwrap_or_else(SystemTime::now),
|
||||
}
|
||||
}
|
||||
|
||||
fn budget_attempt_usd(&self) -> Decimal {
|
||||
if self.budget_attempt_already_counted {
|
||||
return Decimal::ZERO;
|
||||
}
|
||||
// Key 只累计已取得的 USD 费用;缺少费用的请求按零结算。
|
||||
// 发送状态仍用于判断重放是否安全并保留在诊断中,不能据此推断存在欠费。
|
||||
self.observation
|
||||
.cost
|
||||
.total()
|
||||
.filter(|money| money.currency().as_str() == "USD")
|
||||
.map(crate::metering::Money::amount)
|
||||
.unwrap_or(Decimal::ZERO)
|
||||
}
|
||||
|
||||
/// 返回截至当前已完成的实际上游调用结果。
|
||||
///
|
||||
/// 调用方可保存已消费下标;该切片在会话生命周期内只追加、不改写。
|
||||
@@ -800,6 +836,7 @@ where
|
||||
let mut events = std::mem::take(&mut self.last_retryable_failure_events);
|
||||
if !events.is_empty() {
|
||||
self.observe_atomic_terminal_events(&mut events);
|
||||
self.budget_attempt_already_counted = true;
|
||||
return Ok(Some(PullOutcome::TerminalFailure {
|
||||
events,
|
||||
error: last_failure,
|
||||
@@ -1143,6 +1180,11 @@ where
|
||||
"sendState": format!("{attempt_send_state:?}"),
|
||||
}));
|
||||
if retryable {
|
||||
for event in &atomic_client_events {
|
||||
for fact in event.canonical_facts() {
|
||||
self.observation.observe_event(fact);
|
||||
}
|
||||
}
|
||||
// 普通 clone 只保留稳定事实;原始 wire/HTTP response 由 request-local
|
||||
// 所有权保留到下一次 attempt 成功,或最终空选路时返回客户端。
|
||||
let persistence_error = self.request_persisted.then(|| error.clone());
|
||||
@@ -1317,6 +1359,12 @@ where
|
||||
}
|
||||
|
||||
fn reset_uncommitted_observations(&mut self) {
|
||||
// 响应观测按尝试隔离;被丢弃尝试中已取得的费用仍计入 Key。
|
||||
self.budget_prior_attempts_usd = self
|
||||
.budget_prior_attempts_usd
|
||||
.checked_add(self.budget_attempt_usd())
|
||||
.unwrap_or(Decimal::MAX);
|
||||
self.budget_attempt_already_counted = false;
|
||||
self.observation.reset_for_attempt();
|
||||
self.upstream_complete = false;
|
||||
}
|
||||
@@ -1326,6 +1374,7 @@ where
|
||||
return Ok(());
|
||||
}
|
||||
let completed_at = SystemTime::now();
|
||||
self.finalized_at = Some(completed_at);
|
||||
self.observation.finish();
|
||||
let upstream_request_id = self.current.as_ref().and_then(|current| {
|
||||
current
|
||||
@@ -1505,11 +1554,12 @@ where
|
||||
"sendState": format!("{:?}", finalization.send_state), "attempts": self.attempts,
|
||||
}),
|
||||
);
|
||||
let completed_at = SystemTime::now();
|
||||
self.finalized_at = Some(completed_at);
|
||||
if !self.request_persisted {
|
||||
self.finalized = true;
|
||||
return Ok(());
|
||||
}
|
||||
let completed_at = SystemTime::now();
|
||||
self.observation.finish();
|
||||
let upstream_request_id = self.current.as_ref().and_then(|current| {
|
||||
current
|
||||
|
||||
@@ -14,6 +14,7 @@ use uuid::Uuid;
|
||||
use crate::engine::admission::{
|
||||
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
|
||||
};
|
||||
use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort};
|
||||
use crate::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
|
||||
NativeContinuationStoreErrorKind, PreviousResponseId,
|
||||
@@ -256,6 +257,7 @@ pub struct DefaultExecutionService {
|
||||
circuits: Arc<dyn ProviderCircuitPort>,
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
impl DefaultExecutionService {
|
||||
@@ -283,9 +285,16 @@ impl DefaultExecutionService {
|
||||
circuits,
|
||||
continuation,
|
||||
client_api_key_usage,
|
||||
budget: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_budget(mut self, budget: Arc<dyn ClientBudgetPort>) -> Self {
|
||||
self.budget = Some(budget);
|
||||
self
|
||||
}
|
||||
|
||||
async fn start_inner(&self, request: StartExecution) -> Result<StartedExecution, GatewayError> {
|
||||
let StartExecution {
|
||||
client,
|
||||
@@ -325,9 +334,18 @@ impl DefaultExecutionService {
|
||||
&self,
|
||||
mut request: PendingStartExecution,
|
||||
) -> Result<StartedExecution, GatewayError> {
|
||||
request.client.policy.authorize().map_err(|_| {
|
||||
GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled")
|
||||
})?;
|
||||
// 长连接每次执行都重新鉴权并冻结当前策略,确保限额和授权变更对新请求生效。
|
||||
request.client = self
|
||||
.authenticate(request.client.policy.plaintext_key().expose_for_auth())
|
||||
.map_err(|error| match error {
|
||||
ClientAuthenticationError::InvalidKey => {
|
||||
GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid")
|
||||
}
|
||||
ClientAuthenticationError::SnapshotUnavailable => GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"runtime snapshot is unavailable",
|
||||
),
|
||||
})?;
|
||||
let started_at = SystemTime::now();
|
||||
let deadline_at = started_at
|
||||
.checked_add(MODEL_REQUEST_DEADLINE)
|
||||
@@ -504,6 +522,12 @@ impl DefaultExecutionService {
|
||||
client_api_key_id: client.policy.key_id().clone(),
|
||||
model_request_id: request_id.clone(),
|
||||
};
|
||||
if let Some(budget) = &self.budget
|
||||
&& let Err(error) = budget.admit(client.policy.key_id().clone()).await
|
||||
{
|
||||
admission.release().await;
|
||||
return Err(error);
|
||||
}
|
||||
let observation = plan
|
||||
.candidates()
|
||||
.first()
|
||||
@@ -552,6 +576,18 @@ impl DefaultExecutionService {
|
||||
{
|
||||
Ok(core) => core,
|
||||
Err(error) => {
|
||||
if let Some(budget) = &self.budget {
|
||||
settle_budget(
|
||||
budget.as_ref(),
|
||||
ClientBudgetCharge {
|
||||
key_id: client.policy.key_id().clone(),
|
||||
request_id: request_id.clone(),
|
||||
amount_usd: crate::metering::Decimal::ZERO,
|
||||
completed_at: SystemTime::now(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
admission.release().await;
|
||||
return Err(gateway_error_from_engine(&error));
|
||||
}
|
||||
@@ -565,6 +601,7 @@ impl DefaultExecutionService {
|
||||
admission,
|
||||
Arc::clone(&self.circuits),
|
||||
Arc::clone(&self.continuation),
|
||||
self.budget.clone(),
|
||||
)),
|
||||
})
|
||||
}
|
||||
@@ -922,6 +959,12 @@ struct AdmissionLease {
|
||||
model_request_id: ModelRequestId,
|
||||
}
|
||||
|
||||
async fn settle_budget(port: &dyn ClientBudgetPort, charge: ClientBudgetCharge) {
|
||||
if let Err(error) = port.settle(charge).await {
|
||||
tracing::error!(%error, "Client budget settlement failed; storage will retry on the next request");
|
||||
}
|
||||
}
|
||||
|
||||
impl AdmissionLease {
|
||||
async fn release(self) {
|
||||
if let Err(error) = self
|
||||
@@ -941,6 +984,7 @@ struct DefaultExecutionSession {
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
observed_provider_outcomes: usize,
|
||||
continuation_recorded: bool,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
impl DefaultExecutionSession {
|
||||
@@ -949,6 +993,7 @@ impl DefaultExecutionSession {
|
||||
admission: AdmissionLease,
|
||||
circuits: Arc<dyn ProviderCircuitPort>,
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
core: Some(core),
|
||||
@@ -957,6 +1002,7 @@ impl DefaultExecutionSession {
|
||||
continuation,
|
||||
observed_provider_outcomes: 0,
|
||||
continuation_recorded: false,
|
||||
budget,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -973,6 +1019,11 @@ impl DefaultExecutionSession {
|
||||
.is_some_and(ResponseExecutionSession::is_finalized)
|
||||
&& let Some(admission) = self.admission.take()
|
||||
{
|
||||
if let Some(budget) = self.budget.take()
|
||||
&& let Some(core) = &self.core
|
||||
{
|
||||
settle_budget(budget.as_ref(), core.budget_charge()).await;
|
||||
}
|
||||
admission.release().await;
|
||||
}
|
||||
}
|
||||
@@ -1023,6 +1074,9 @@ impl DefaultExecutionSession {
|
||||
.get(self.observed_provider_outcomes..)
|
||||
.unwrap_or_default();
|
||||
publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await;
|
||||
if let Some(budget) = self.budget.take() {
|
||||
settle_budget(budget.as_ref(), core.budget_charge()).await;
|
||||
}
|
||||
if let Some(admission) = self.admission.take() {
|
||||
admission.release().await;
|
||||
}
|
||||
@@ -1213,6 +1267,17 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
|
||||
|
||||
fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError {
|
||||
match error {
|
||||
crate::validation::RoutingError::ModelNotFound {
|
||||
model,
|
||||
mapped_model,
|
||||
} => GatewayError::new(
|
||||
GatewayErrorKind::ModelNotFound,
|
||||
if model == mapped_model {
|
||||
"the requested model was not found in the provider catalogs available to this API key; check the model name"
|
||||
} else {
|
||||
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping"
|
||||
},
|
||||
),
|
||||
crate::validation::RoutingError::NoCapableProvider { .. }
|
||||
| crate::validation::RoutingError::NoCapableProviderEndpoint { .. }
|
||||
| crate::validation::RoutingError::EmptyAccountScope => GatewayError::new(
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。
|
||||
|
||||
pub mod admission;
|
||||
pub mod budget;
|
||||
pub mod continuation;
|
||||
pub mod coordinator;
|
||||
pub mod execution;
|
||||
|
||||
@@ -975,6 +975,13 @@ pub struct GatewayError {
|
||||
message: &'static str,
|
||||
diagnostic: Option<Box<ProviderDiagnostic>>,
|
||||
client_visible_upstream_error: Option<ClientVisibleUpstreamError>,
|
||||
client_details: Option<Box<GatewayClientErrorDetails>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct GatewayClientErrorDetails {
|
||||
code: Option<&'static str>,
|
||||
retry_after: Option<Duration>,
|
||||
}
|
||||
|
||||
impl GatewayError {
|
||||
@@ -986,9 +993,33 @@ impl GatewayError {
|
||||
message,
|
||||
diagnostic: None,
|
||||
client_visible_upstream_error: None,
|
||||
client_details: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_client_code(mut self, code: &'static str) -> Self {
|
||||
self.client_details
|
||||
.get_or_insert_with(Default::default)
|
||||
.code = Some(code);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_retry_after(mut self, delay: Duration) -> Self {
|
||||
self.client_details
|
||||
.get_or_insert_with(Default::default)
|
||||
.retry_after = Some(delay);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn retry_after(&self) -> Option<Duration> {
|
||||
self.client_details
|
||||
.as_ref()
|
||||
.and_then(|details| details.retry_after)
|
||||
}
|
||||
|
||||
/// 将 Provider 错误归一为客户端无关错误。
|
||||
#[must_use]
|
||||
pub fn from_provider(error: &ProviderError) -> Self {
|
||||
@@ -1103,6 +1134,11 @@ impl GatewayError {
|
||||
self.client_visible_upstream_error
|
||||
.as_ref()
|
||||
.and_then(ClientVisibleUpstreamError::code)
|
||||
.or_else(|| {
|
||||
self.client_details
|
||||
.as_ref()
|
||||
.and_then(|details| details.code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -52,6 +52,7 @@ pub struct CoreStorePorts {
|
||||
snapshots: Arc<dyn SnapshotStorePort>,
|
||||
snapshot_subscriptions: Arc<dyn SnapshotSubscriptionPort>,
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
budget: Option<Arc<dyn engine::budget::ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
impl CoreStorePorts {
|
||||
@@ -79,8 +80,15 @@ impl CoreStorePorts {
|
||||
snapshots,
|
||||
snapshot_subscriptions,
|
||||
client_api_key_usage,
|
||||
budget: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_budget(mut self, budget: Arc<dyn engine::budget::ClientBudgetPort>) -> Self {
|
||||
self.budget = Some(budget);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CoreBundle {
|
||||
@@ -145,7 +153,7 @@ pub async fn initialize(
|
||||
Arc::clone(&ports.snapshot_subscriptions),
|
||||
));
|
||||
let worker_contributions = publisher.worker_contributions()?;
|
||||
let service = Arc::new(DefaultExecutionService::new(
|
||||
let mut service = DefaultExecutionService::new(
|
||||
snapshots.clone(),
|
||||
ports.execution,
|
||||
providers,
|
||||
@@ -153,7 +161,11 @@ pub async fn initialize(
|
||||
ports.circuits,
|
||||
ports.continuation,
|
||||
ports.client_api_key_usage,
|
||||
));
|
||||
);
|
||||
if let Some(budget) = ports.budget {
|
||||
service = service.with_budget(budget);
|
||||
}
|
||||
let service = Arc::new(service);
|
||||
let execution: Arc<dyn ExecutionService> = service.clone();
|
||||
let account_probe: Arc<dyn AccountProbe> = service;
|
||||
let snapshot_control: Arc<dyn SnapshotControl> = publisher;
|
||||
|
||||
@@ -14,6 +14,8 @@ pub struct Decimal(u128);
|
||||
|
||||
impl Decimal {
|
||||
pub const ZERO: Self = Self(0);
|
||||
/// 数据库可表示的最大非负金额。
|
||||
pub const MAX: Self = Self(MAX_SCALED_DECIMAL);
|
||||
|
||||
/// 从按十位小数缩放的整数创建。
|
||||
///
|
||||
|
||||
@@ -778,6 +778,19 @@ impl RuntimeSnapshot {
|
||||
}
|
||||
|
||||
if candidates.is_empty() {
|
||||
// 模型存在性必须包含被熔断的 Provider;目录未知时不能断言模型不存在。
|
||||
let mut scoped_providers = providers.intersection(&self.providers).peekable();
|
||||
if scoped_providers.peek().is_some()
|
||||
&& scoped_providers.all(|provider| {
|
||||
self.known_provider_catalogs.contains(provider)
|
||||
&& !self.contains_public_model_for_provider(public_model, provider)
|
||||
})
|
||||
{
|
||||
return Err(RoutingError::ModelNotFound {
|
||||
model: public_model.as_str().to_owned(),
|
||||
mapped_model: self.mapped_model(public_model.as_str()),
|
||||
});
|
||||
}
|
||||
return Err(RoutingError::NoCapableProvider {
|
||||
model: public_model.as_str().to_owned(),
|
||||
});
|
||||
|
||||
@@ -85,6 +85,16 @@ pub enum RoutingError {
|
||||
/// 实体 ID。
|
||||
id: String,
|
||||
},
|
||||
/// 账号范围内的 Provider 目录均已知,且均没有请求或映射后的模型。
|
||||
#[error(
|
||||
"model `{model}` (mapped to `{mapped_model}`) was not found in the scoped provider catalogs"
|
||||
)]
|
||||
ModelNotFound {
|
||||
/// 客户端提交的模型名称。
|
||||
model: String,
|
||||
/// 应用模型映射后的名称。
|
||||
mapped_model: String,
|
||||
},
|
||||
/// 固定平台内没有可执行本次请求的 Provider。
|
||||
#[error("no provider can execute model `{model}`")]
|
||||
NoCapableProvider {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod proxy;
|
||||
mod selection;
|
||||
mod selection_diagnostics;
|
||||
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
use gateway_core::account::OutboundProxy;
|
||||
|
||||
#[test]
|
||||
fn proxy_endpoints_support_explicit_schemes_and_redact_credentials() {
|
||||
for scheme in ["http", "https", "socks5", "socks5h"] {
|
||||
let proxy =
|
||||
OutboundProxy::parse(&format!("{scheme}://user:p%40ss%3Aword@[::1]:1080")).unwrap();
|
||||
assert!(proxy.expose_url().contains("p%40ss%3Aword"));
|
||||
assert!(!proxy.endpoint().contains("user"));
|
||||
assert!(!proxy.endpoint().contains("p%40ss"));
|
||||
assert!(!format!("{proxy:?}").contains("user"));
|
||||
assert!(proxy.endpoint().contains("[::1]:1080"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_proxy_never_silently_becomes_direct() {
|
||||
for url in [
|
||||
"",
|
||||
"localhost:1080",
|
||||
"file:///etc/passwd",
|
||||
"ftp://host:21",
|
||||
"http://host:0",
|
||||
"socks5://host",
|
||||
"http://host:80/path",
|
||||
"http://host:80?secret=1",
|
||||
"http://host:80#fragment",
|
||||
"http://user:secret@host:80\n",
|
||||
] {
|
||||
let error = OutboundProxy::parse(url).unwrap_err();
|
||||
assert!(!error.to_string().contains("secret"));
|
||||
}
|
||||
}
|
||||
@@ -2158,6 +2158,146 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() {
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.finalizations[0].cost_source, CostSource::Unavailable);
|
||||
assert_eq!(state.finalizations[0].cost_ticks, None);
|
||||
assert_eq!(
|
||||
session.budget_charge().amount_usd.scaled(),
|
||||
999,
|
||||
"missing final usage cannot erase a charged retry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_budget_accumulates_only_observed_retry_costs() {
|
||||
for (first_cost, first_send, expected) in [
|
||||
(Some(999), UpstreamSendState::Sent, 1122),
|
||||
(
|
||||
Some(gateway_core::metering::Decimal::MAX.scaled()),
|
||||
UpstreamSendState::Sent,
|
||||
gateway_core::metering::Decimal::MAX.scaled(),
|
||||
),
|
||||
(None, UpstreamSendState::Sent, 123),
|
||||
(None, UpstreamSendState::NotSent, 123),
|
||||
] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let mut first = Vec::new();
|
||||
if first_send != UpstreamSendState::NotSent {
|
||||
first.push(Ok(GatewayEvent::Started(ResponseMeta::new(
|
||||
"resp_first",
|
||||
"gpt-5",
|
||||
))));
|
||||
}
|
||||
if let Some(cost) = first_cost {
|
||||
first.push(Ok(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
|
||||
)));
|
||||
}
|
||||
first.push(Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
first_send,
|
||||
)
|
||||
.with_replay_safe()));
|
||||
let (coordinator, _, _) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: first,
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_second",
|
||||
items: vec![
|
||||
Ok(GatewayEvent::Started(ResponseMeta::new(
|
||||
"resp_done",
|
||||
"gpt-5",
|
||||
))),
|
||||
Ok(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(123).unwrap(),
|
||||
)),
|
||||
Ok(GatewayEvent::Completed(ResponseMeta::new(
|
||||
"resp_done",
|
||||
"gpt-5",
|
||||
))),
|
||||
],
|
||||
},
|
||||
]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
block_on(session.collect_uncommitted()).unwrap();
|
||||
block_on(session.commit_downstream(Some(200))).unwrap();
|
||||
let charge = session.budget_charge();
|
||||
assert_eq!(charge.amount_usd.scaled(), expected);
|
||||
assert_eq!(
|
||||
session.budget_charge().completed_at,
|
||||
charge.completed_at,
|
||||
"final settlement time stays stable"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_budget_settles_post_send_failures_from_observed_cost_without_replaying() {
|
||||
for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] {
|
||||
for cost in [None, Some(123)] {
|
||||
for committed in [false, true] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let mut items = vec![Ok(GatewayEvent::Started(ResponseMeta::new(
|
||||
"response-interrupted",
|
||||
"gpt-5",
|
||||
)))];
|
||||
if let Some(cost) = cost {
|
||||
items.push(Ok(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
|
||||
)));
|
||||
}
|
||||
items.push(Err(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
send_state,
|
||||
)));
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items,
|
||||
}]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
if committed {
|
||||
block_on(session.next_event()).unwrap().unwrap();
|
||||
block_on(session.commit_downstream(Some(200))).unwrap();
|
||||
loop {
|
||||
match block_on(session.next_event()) {
|
||||
Ok(Some(_)) => {}
|
||||
Err(EngineError::Provider(_)) => break,
|
||||
other => panic!("expected terminal transport error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
} else {
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Provider(_))
|
||||
));
|
||||
}
|
||||
assert!(session.is_finalized());
|
||||
assert_eq!(
|
||||
session.budget_charge().amount_usd.scaled(),
|
||||
cost.unwrap_or(0),
|
||||
"send={send_state:?}, committed={committed}, cost={cost:?}"
|
||||
);
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2688,10 +2828,16 @@ fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once()
|
||||
fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let mut failure = bare_atomic_response_failed("response-only-failed");
|
||||
let mut events = failure.take_atomic_client_events();
|
||||
events.push(ProviderEvent::canonical(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(500).unwrap(),
|
||||
)));
|
||||
let failure = failure.with_atomic_client_events(events);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::ObservedStream {
|
||||
account_id: "acct_only",
|
||||
items: vec![Err(bare_atomic_response_failed("response-only-failed"))],
|
||||
items: vec![Err(failure)],
|
||||
},
|
||||
Script::Error(ProviderError::new(
|
||||
ProviderErrorKind::NoEligibleAccount,
|
||||
@@ -2731,6 +2877,11 @@ fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire()
|
||||
));
|
||||
|
||||
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
|
||||
assert_eq!(
|
||||
session.budget_charge().amount_usd.scaled(),
|
||||
500,
|
||||
"restoring a discarded failure must not charge its usage twice"
|
||||
);
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.attempts.len(), 1);
|
||||
assert_eq!(state.intermediate_failures, 1);
|
||||
@@ -2868,6 +3019,7 @@ fn ambiguous_send_state_stops_retry() {
|
||||
let finalization = &state.finalizations[0];
|
||||
assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
|
||||
assert_eq!(finalization.send_state, UpstreamSendState::Ambiguous);
|
||||
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
|
||||
assert_eq!(finalization.attempt_count, 1);
|
||||
assert!(!finalization.committed);
|
||||
assert!(finalization.latency_ms.is_some());
|
||||
@@ -3342,6 +3494,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
|
||||
assert_eq!(state.commits, 0);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
|
||||
assert!(!state.finalizations[0].committed);
|
||||
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -7,6 +7,231 @@ use std::{
|
||||
time::{Duration, Instant, SystemTime},
|
||||
};
|
||||
|
||||
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
|
||||
use gateway_core::error::GatewayError;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Admissions {
|
||||
active: Arc<AtomicBool>,
|
||||
limits: Mutex<Vec<RateLimits>>,
|
||||
}
|
||||
|
||||
impl ClientAdmissionPort for Admissions {
|
||||
fn admit(
|
||||
&self,
|
||||
request: ClientAdmissionRequest,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
|
||||
Box::pin(async move {
|
||||
self.limits.lock().unwrap().push(request.limits);
|
||||
assert!(!self.active.swap(true, Ordering::SeqCst));
|
||||
Ok(ClientAdmissionDecision::Granted)
|
||||
})
|
||||
}
|
||||
fn release<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
|
||||
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
|
||||
}
|
||||
fn restore(
|
||||
&self,
|
||||
_: ClientAdmissionRecovery,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
|
||||
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
|
||||
}
|
||||
}
|
||||
|
||||
struct Budget {
|
||||
reject: bool,
|
||||
active: Arc<AtomicBool>,
|
||||
charges: Mutex<Vec<ClientBudgetCharge>>,
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for Budget {
|
||||
fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async {
|
||||
assert!(self.active.load(Ordering::SeqCst));
|
||||
if self.reject {
|
||||
Err(
|
||||
GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted")
|
||||
.with_client_code("key_daily_budget_exceeded"),
|
||||
)
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
})
|
||||
}
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
assert!(
|
||||
self.active.load(Ordering::SeqCst),
|
||||
"settle before releasing the concurrent request slot"
|
||||
);
|
||||
self.charges.lock().unwrap().push(charge);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn service(admissions: Arc<Admissions>, budget: Arc<Budget>) -> DefaultExecutionService {
|
||||
DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
.with_budget(budget)
|
||||
}
|
||||
|
||||
fn request(service: &DefaultExecutionService, transport: ClientTransport) -> StartExecution {
|
||||
StartExecution {
|
||||
client: service.authenticate("sk_start_test").unwrap(),
|
||||
public_model: PublicModelId::new("gpt-start").unwrap(),
|
||||
operation: start_operation(),
|
||||
metadata: ExecutionRequestMetadata {
|
||||
protocol: "openai".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
transport,
|
||||
stream: transport != ClientTransport::HttpJson,
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
previous_response_id: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let budget = Arc::new(Budget {
|
||||
reject: true,
|
||||
active: admissions.active.clone(),
|
||||
charges: Mutex::default(),
|
||||
});
|
||||
let service = service(admissions.clone(), budget.clone());
|
||||
for transport in [
|
||||
ClientTransport::HttpJson,
|
||||
ClientTransport::HttpSse,
|
||||
ClientTransport::WebSocket,
|
||||
] {
|
||||
let result = block_on(service.start(request(&service, transport)));
|
||||
assert!(
|
||||
matches!(result, Err(error) if error.client_error_code() == Some("key_daily_budget_exceeded"))
|
||||
);
|
||||
assert!(!admissions.active.load(Ordering::SeqCst));
|
||||
assert!(budget.charges.lock().unwrap().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reused_client_uses_updated_limits_for_each_execution() {
|
||||
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let service = DefaultExecutionService::new(
|
||||
snapshots.clone(),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
admissions.clone(),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let client = service.authenticate("sk_start_test").unwrap();
|
||||
let limited = RateLimits {
|
||||
max_concurrency: 1,
|
||||
requests_per_minute: 1,
|
||||
};
|
||||
for (revision, limits) in [(2, limited), (3, RateLimits::unlimited())] {
|
||||
snapshots.publish(start_snapshot_with_policy(revision, true, limits));
|
||||
let mut next = request(&service, ClientTransport::WebSocket);
|
||||
next.client = client.clone();
|
||||
let started = block_on(service.start(next)).expect("new execution");
|
||||
block_on(started.session.detach_finalize());
|
||||
}
|
||||
assert_eq!(
|
||||
*admissions.limits.lock().unwrap(),
|
||||
vec![limited, RateLimits::unlimited()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reused_client_cannot_start_after_key_disable_or_snapshot_suspension() {
|
||||
for suspend in [false, true] {
|
||||
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let service = DefaultExecutionService::new(
|
||||
snapshots.clone(),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
admissions.clone(),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let next = request(&service, ClientTransport::WebSocket);
|
||||
if suspend {
|
||||
snapshots.suspend();
|
||||
} else {
|
||||
snapshots.publish(start_snapshot_with_policy(
|
||||
2,
|
||||
false,
|
||||
RateLimits::unlimited(),
|
||||
));
|
||||
}
|
||||
let result = block_on(service.start(next));
|
||||
let expected = if suspend {
|
||||
GatewayErrorKind::Internal
|
||||
} else {
|
||||
GatewayErrorKind::Unauthorized
|
||||
};
|
||||
assert!(matches!(result, Err(error) if error.kind() == expected));
|
||||
assert!(admissions.limits.lock().unwrap().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all_transports() {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let budget = Arc::new(Budget {
|
||||
reject: false,
|
||||
active: admissions.active.clone(),
|
||||
charges: Mutex::default(),
|
||||
});
|
||||
let service = service(admissions.clone(), budget.clone());
|
||||
for transport in [
|
||||
ClientTransport::HttpJson,
|
||||
ClientTransport::HttpSse,
|
||||
ClientTransport::WebSocket,
|
||||
] {
|
||||
for detach in [true, false] {
|
||||
let mut started = block_on(service.start(request(&service, transport))).unwrap();
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
if detach {
|
||||
block_on(started.session.detach_finalize());
|
||||
} else {
|
||||
assert!(block_on(started.session.next_event()).is_err());
|
||||
block_on(started.session.detach_finalize());
|
||||
}
|
||||
assert!(!admissions.active.load(Ordering::SeqCst));
|
||||
}
|
||||
}
|
||||
let charges = budget.charges.lock().unwrap();
|
||||
assert_eq!(
|
||||
charges.len(),
|
||||
6,
|
||||
"one settlement per request, including detached finalized sessions"
|
||||
);
|
||||
assert!(
|
||||
charges
|
||||
.iter()
|
||||
.all(|charge| charge.amount_usd == gateway_core::metering::Decimal::ZERO)
|
||||
);
|
||||
}
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::{executor::block_on, future::BoxFuture};
|
||||
@@ -495,7 +720,13 @@ fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
|
||||
panic!("a known provider catalog must be authoritative for model availability");
|
||||
};
|
||||
|
||||
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
|
||||
assert_eq!(
|
||||
(error.kind(), error.client_message()),
|
||||
(
|
||||
GatewayErrorKind::ModelNotFound,
|
||||
"the requested model was not found in the provider catalogs available to this API key; check the model name",
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -917,11 +1148,15 @@ fn client_snapshot() -> RuntimeSnapshot {
|
||||
}
|
||||
|
||||
fn start_snapshot() -> RuntimeSnapshot {
|
||||
start_snapshot_with_policy(1, true, RateLimits::unlimited())
|
||||
}
|
||||
|
||||
fn start_snapshot_with_policy(revision: u64, enabled: bool, limits: RateLimits) -> RuntimeSnapshot {
|
||||
let provider = ProviderKind::new("openai").expect("provider kind");
|
||||
let capabilities =
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
ConfigRevision::new(revision).expect("config revision"),
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
std::num::NonZeroU32::new(1).expect("concurrency"),
|
||||
@@ -937,8 +1172,8 @@ fn start_snapshot() -> RuntimeSnapshot {
|
||||
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
|
||||
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
|
||||
account_scope(&provider, "acct_start"),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
enabled,
|
||||
limits,
|
||||
)],
|
||||
)
|
||||
.expect("start snapshot")
|
||||
|
||||
@@ -326,12 +326,123 @@ fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
|
||||
)
|
||||
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::ModelNotFound {
|
||||
model: "gpt-5.4".to_owned(),
|
||||
mapped_model: "gpt-5.5".to_owned(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_model_should_be_reported_when_all_scoped_catalogs_are_known() {
|
||||
let snapshot = snapshot();
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("missing-model").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect_err("neither provider publishes the requested model");
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::ModelNotFound {
|
||||
model: "missing-model".to_owned(),
|
||||
mapped_model: "missing-model".to_owned(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_model_in_one_provider_should_still_route_to_another() {
|
||||
let snapshot = snapshot();
|
||||
let plan = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("grok-4.5").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect("xAI publishes the model missing from OpenAI");
|
||||
|
||||
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocked_model_provider_should_not_be_misreported_as_model_not_found() {
|
||||
let snapshot = snapshot();
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-5.5").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect_err("OpenAI is blocked and xAI does not publish this model");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::NoCapableProvider { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocked_unknown_catalog_should_not_prove_model_absence() {
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("revision"),
|
||||
scheduling(),
|
||||
vec![
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
],
|
||||
vec![model("xai", "grok-4.5", capabilities())],
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("snapshot")
|
||||
.with_account_directory(account_directory());
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("missing-model").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect_err("the blocked provider's catalog cannot establish model absence");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::NoCapableProvider { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_account_scope_should_not_be_misreported_as_model_not_found() {
|
||||
let snapshot = snapshot();
|
||||
let empty_scope = Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::default()),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
));
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("missing-model").expect("model"),
|
||||
&operation(),
|
||||
empty_scope,
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect_err("the API key has no accounts");
|
||||
|
||||
assert_eq!(error, gateway_core::error::RoutingError::EmptyAccountScope);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_endpoint_plan_should_not_consult_or_publish_the_text_model_catalog() {
|
||||
let snapshot = snapshot();
|
||||
@@ -463,19 +574,22 @@ fn known_unsupported_operation_should_not_be_bypassed() {
|
||||
)
|
||||
.expect("snapshot");
|
||||
|
||||
assert!(
|
||||
snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-known-unsupported").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-known-unsupported").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect_err("the model exists but cannot perform this operation");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::NoCapableProvider { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -56,10 +56,13 @@ struct PublishingCatalogProvider {
|
||||
queries: AtomicUsize,
|
||||
}
|
||||
|
||||
struct UnavailableCatalogProvider;
|
||||
enum TestCatalogProvider {
|
||||
Unavailable,
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for UnavailableCatalogProvider {
|
||||
impl Provider for TestCatalogProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"alpha"
|
||||
}
|
||||
@@ -71,10 +74,13 @@ impl Provider for UnavailableCatalogProvider {
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
))
|
||||
match self {
|
||||
Self::Unavailable => Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
)),
|
||||
Self::Empty => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
@@ -145,7 +151,7 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
|
||||
#[test]
|
||||
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
|
||||
let providers =
|
||||
ProviderRegistry::new([Arc::new(UnavailableCatalogProvider) as Arc<dyn Provider>])
|
||||
ProviderRegistry::new([Arc::new(TestCatalogProvider::Unavailable) as Arc<dyn Provider>])
|
||||
.expect("provider registry");
|
||||
let compiler =
|
||||
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
|
||||
@@ -162,6 +168,35 @@ fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
|
||||
assert_eq!(snapshot.client_policies().count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn known_empty_catalog_should_report_model_not_found() {
|
||||
let providers =
|
||||
ProviderRegistry::new([Arc::new(TestCatalogProvider::Empty) as Arc<dyn Provider>])
|
||||
.expect("provider registry");
|
||||
let compiler =
|
||||
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
|
||||
let snapshot = block_on(compiler.compile()).expect("compile empty catalog");
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("public-model").expect("model"),
|
||||
&super::operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&gateway_core::routing::RoutingContext {
|
||||
required_provider: Some(ProviderKind::new("alpha").expect("provider")),
|
||||
..gateway_core::routing::RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect_err("a successfully published empty catalog proves model absence");
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::ModelNotFound {
|
||||
model: "public-model".to_owned(),
|
||||
mapped_model: "upstream-model".to_owned(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
|
||||
let provider = Arc::new(PublishingCatalogProvider {
|
||||
|
||||
@@ -17,7 +17,7 @@ futures = "0.3.32"
|
||||
gateway-admin = { path = "../gateway-admin" }
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
hex = "0.4.3"
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "rustls-tls-native-roots"] }
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "rustls-tls-native-roots", "socks"] }
|
||||
scraper = "0.27.0"
|
||||
semver = "1.0.27"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
pub mod client_distribution;
|
||||
pub mod config;
|
||||
mod logging;
|
||||
pub mod proxy_probe;
|
||||
pub mod serve;
|
||||
pub mod system_update;
|
||||
pub mod workers;
|
||||
@@ -81,6 +82,17 @@ impl HostBundle {
|
||||
self.client_distribution.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn proxy_probe<E>(
|
||||
&self,
|
||||
build_client: impl Fn(reqwest::ClientBuilder) -> Result<reqwest::Client, E>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> Arc<dyn gateway_admin::ports::proxy::ProxyProbe> {
|
||||
Arc::new(proxy_probe::HttpProxyProbe::default().with_client_builder(build_client))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn connection_lifecycle(&self) -> Arc<dyn ConnectionLifecycle> {
|
||||
self.connections.clone()
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
//! 使用与 Provider 请求一致的显式代理协议,执行有超时和响应大小限制的出口测试。
|
||||
|
||||
use std::{
|
||||
net::IpAddr,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_admin::{model::proxies::ProxyTestResult, ports::proxy::ProxyProbe};
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use serde::Deserialize;
|
||||
|
||||
pub struct HttpProxyProbe {
|
||||
endpoint: String,
|
||||
build_client: Arc<ProxyClientBuilder>,
|
||||
}
|
||||
|
||||
type ProxyClientBuilder =
|
||||
dyn Fn(reqwest::ClientBuilder) -> Result<reqwest::Client, &'static str> + Send + Sync;
|
||||
|
||||
impl Default for HttpProxyProbe {
|
||||
fn default() -> Self {
|
||||
Self::new("https://api.ipify.org?format=json")
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpProxyProbe {
|
||||
#[must_use]
|
||||
pub fn new(endpoint: impl Into<String>) -> Self {
|
||||
Self {
|
||||
endpoint: endpoint.into(),
|
||||
build_client: Arc::new(|builder| builder.build().map_err(|_| "无法创建代理连接")),
|
||||
}
|
||||
}
|
||||
|
||||
/// 由组合根注入与 Provider 请求一致的证书信任策略。
|
||||
#[must_use]
|
||||
pub fn with_client_builder<E>(
|
||||
mut self,
|
||||
build: impl Fn(reqwest::ClientBuilder) -> Result<reqwest::Client, E> + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
self.build_client = Arc::new(move |builder| {
|
||||
build(builder).map_err(|_| "无法创建代理连接,请检查证书信任配置")
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
async fn exit_ip(&self, proxy: &OutboundProxy) -> Result<IpAddr, &'static str> {
|
||||
let proxy = reqwest::Proxy::all(proxy.expose_url()).map_err(|_| "代理地址不合法")?;
|
||||
let builder = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.proxy(proxy)
|
||||
.connect_timeout(Duration::from_secs(5))
|
||||
.timeout(Duration::from_secs(12))
|
||||
.redirect(reqwest::redirect::Policy::none());
|
||||
let client = (self.build_client)(builder)?;
|
||||
let mut response = client.get(&self.endpoint).send().await.map_err(|error| {
|
||||
if error.is_timeout() {
|
||||
"代理连接超时"
|
||||
} else {
|
||||
"代理连接失败,请检查地址、认证和网络"
|
||||
}
|
||||
})?;
|
||||
if !response.status().is_success() {
|
||||
return Err(
|
||||
if response.status() == reqwest::StatusCode::PROXY_AUTHENTICATION_REQUIRED {
|
||||
"代理认证失败"
|
||||
} else {
|
||||
"出口检测服务返回错误状态"
|
||||
},
|
||||
);
|
||||
}
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| "出口检测响应读取失败")?
|
||||
{
|
||||
if body.len() + chunk.len() > 1024 {
|
||||
return Err("出口检测响应过大");
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct Response {
|
||||
ip: IpAddr,
|
||||
}
|
||||
serde_json::from_slice::<Response>(&body)
|
||||
.map(|response| response.ip)
|
||||
.map_err(|_| "出口检测响应不合法")
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyProbe for HttpProxyProbe {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult {
|
||||
let started = Instant::now();
|
||||
let result = tokio::time::timeout(Duration::from_secs(15), self.exit_ip(proxy)).await;
|
||||
let result = result.unwrap_or(Err("代理连接超时"));
|
||||
ProxyTestResult {
|
||||
success: result.is_ok(),
|
||||
latency_ms: u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
|
||||
exit_ip: result.as_ref().ok().copied(),
|
||||
message: result.map_or_else(str::to_owned, |_| "连接成功".to_owned()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
mod client_distribution;
|
||||
mod config;
|
||||
mod logging;
|
||||
mod proxy_probe;
|
||||
mod serve;
|
||||
mod system_update;
|
||||
mod workers;
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
use gateway_admin::ports::proxy::ProxyProbe;
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use gateway_host::proxy_probe::HttpProxyProbe;
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{any, header},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_sends_authentication_through_explicit_proxy() {
|
||||
let proxy_server = MockServer::start().await;
|
||||
Mock::given(header("proxy-authorization", "Basic dXNlcjpwYXNzd29yZA=="))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ip": "203.0.113.8"})))
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
let proxy =
|
||||
OutboundProxy::parse(&format!("http://user:password@{}", proxy_server.address())).unwrap();
|
||||
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
|
||||
.test(&proxy)
|
||||
.await;
|
||||
assert!(result.success);
|
||||
assert_eq!(result.exit_ip.unwrap().to_string(), "203.0.113.8");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_rejects_auth_errors_redirects_and_invalid_or_oversized_responses() {
|
||||
for response in [
|
||||
ResponseTemplate::new(407),
|
||||
ResponseTemplate::new(302).insert_header("Location", "http://127.0.0.1/"),
|
||||
ResponseTemplate::new(200).set_body_json(json!({"ip": "not-an-ip"})),
|
||||
ResponseTemplate::new(200).set_body_string("a".repeat(1025)),
|
||||
] {
|
||||
let proxy_server = MockServer::start().await;
|
||||
Mock::given(any())
|
||||
.respond_with(response)
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
|
||||
.test(&OutboundProxy::parse(&proxy_server.uri()).unwrap())
|
||||
.await;
|
||||
assert!(!result.success);
|
||||
assert!(result.exit_ip.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unavailable_proxy_never_falls_back_to_direct_connection() {
|
||||
let target = MockServer::start().await;
|
||||
Mock::given(any())
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ip":"203.0.113.8"})))
|
||||
.expect(0)
|
||||
.mount(&target)
|
||||
.await;
|
||||
let unused = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let proxy = OutboundProxy::parse(&format!("http://{}", unused.local_addr().unwrap())).unwrap();
|
||||
drop(unused);
|
||||
let result = HttpProxyProbe::new(target.uri()).test(&proxy).await;
|
||||
assert!(!result.success);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_certificate_configuration_should_not_fall_back_or_expose_details() {
|
||||
let target = MockServer::start().await;
|
||||
Mock::given(any())
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(0)
|
||||
.mount(&target)
|
||||
.await;
|
||||
let result = HttpProxyProbe::new(target.uri())
|
||||
.with_client_builder(|_| Err("private-certificate-path"))
|
||||
.test(&OutboundProxy::parse(&target.uri()).unwrap())
|
||||
.await;
|
||||
assert!(!result.success);
|
||||
assert!(!result.message.contains("private-certificate-path"));
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -106,6 +106,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
)),
|
||||
admin_account_runtime,
|
||||
Arc::new(postgres::PgAccountGroupRepository::new(pool.clone())),
|
||||
Arc::new(postgres::PgProxyRepository::new(pool.clone())),
|
||||
),
|
||||
Arc::new(AdminAuthStoreAdapter {
|
||||
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
|
||||
@@ -171,7 +172,8 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
)?),
|
||||
),
|
||||
Arc::new(client_key_usage),
|
||||
);
|
||||
)
|
||||
.with_budget(Arc::new(postgres::PgClientBudgetStore::new(pool.clone())));
|
||||
|
||||
let provider_ports = ProviderStorePorts::new(
|
||||
account_store,
|
||||
|
||||
@@ -167,7 +167,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
"select membership.account_group_id,
|
||||
account.id, account.provider_kind, account.name, account.email,
|
||||
account.upstream_user_id, account.upstream_account_id, account.plan_type,
|
||||
account.authentication_kind, account.credential_revision,
|
||||
account.authentication_kind, account.credential_revision, account.outbound_proxy_url,
|
||||
account.has_refresh_token, account.access_token_expires_at,
|
||||
account.next_refresh_at, account.enabled, account.concurrency_limit,
|
||||
account.weight, account.credential_state, account.quota_access_state,
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
//! 按 Key 串行检查限额,并幂等累计已取得的 USD 费用。
|
||||
|
||||
use std::{collections::BTreeMap, sync::Mutex, time::Duration};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::{
|
||||
engine::budget::{
|
||||
ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, ClientBudgetPort,
|
||||
ClientBudgetStatus,
|
||||
},
|
||||
error::{GatewayError, GatewayErrorKind},
|
||||
metering::Decimal,
|
||||
policy::ClientApiKeyId,
|
||||
};
|
||||
use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use crate::{StoreResult, postgres_unavailable};
|
||||
|
||||
pub struct PgClientBudgetStore {
|
||||
pool: PgPool,
|
||||
retry: Mutex<BTreeMap<String, ClientBudgetCharge>>,
|
||||
}
|
||||
|
||||
impl PgClientBudgetStore {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
retry: Mutex::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn admit_inner(&self, key_id: ClientApiKeyId) -> Result<(), GatewayError> {
|
||||
// 短暂存储故障后按原金额重试;进程退出丢失的费用不转成人工核账或阻断 Key。
|
||||
let retries = self
|
||||
.retry
|
||||
.lock()
|
||||
.map_err(|_| unavailable())?
|
||||
.values()
|
||||
.filter(|charge| charge.key_id == key_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
for charge in retries {
|
||||
self.settle(charge).await.map_err(|_| unavailable())?;
|
||||
}
|
||||
let mut tx = self.pool.begin().await.map_err(|_| unavailable())?;
|
||||
let row = sqlx::query(
|
||||
"select daily_limit_usd::text, weekly_limit_usd::text, enabled
|
||||
from client_api_keys where id = $1 for update",
|
||||
)
|
||||
.bind(key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| unavailable())?
|
||||
.ok_or_else(|| {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Unauthorized,
|
||||
"client API key no longer exists",
|
||||
)
|
||||
})?;
|
||||
if !row.get::<bool, _>("enabled") {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"client API key is disabled",
|
||||
));
|
||||
}
|
||||
let limits = ClientBudgetLimits {
|
||||
daily_usd: row
|
||||
.get::<String, _>("daily_limit_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?,
|
||||
weekly_usd: row
|
||||
.get::<String, _>("weekly_limit_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?,
|
||||
};
|
||||
let now = Utc::now();
|
||||
advance_windows(&mut tx, key_id.as_str(), now)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
if limits.is_limited() {
|
||||
let window = sqlx::query(
|
||||
"select daily_used_usd::text, weekly_used_usd::text, daily_end, weekly_end
|
||||
from client_key_budget_windows where client_api_key_id = $1",
|
||||
)
|
||||
.bind(key_id.as_str())
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
let daily: Decimal = window
|
||||
.get::<String, _>("daily_used_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?;
|
||||
let weekly: Decimal = window
|
||||
.get::<String, _>("weekly_used_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?;
|
||||
let daily_exceeded = limits.daily_usd != Decimal::ZERO && daily >= limits.daily_usd;
|
||||
let weekly_exceeded = limits.weekly_usd != Decimal::ZERO && weekly >= limits.weekly_usd;
|
||||
if daily_exceeded || weekly_exceeded {
|
||||
let daily_end: DateTime<Utc> = window.get("daily_end");
|
||||
let weekly_end: DateTime<Utc> = window.get("weekly_end");
|
||||
let reset = if weekly_exceeded {
|
||||
weekly_end
|
||||
} else {
|
||||
daily_end
|
||||
};
|
||||
let retry = (reset - now).to_std().unwrap_or(Duration::from_secs(1));
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::RateLimited,
|
||||
"client API key budget is exhausted",
|
||||
)
|
||||
.with_client_code(if weekly_exceeded {
|
||||
"key_weekly_budget_exceeded"
|
||||
} else {
|
||||
"key_daily_budget_exceeded"
|
||||
})
|
||||
.with_retry_after(retry));
|
||||
}
|
||||
}
|
||||
tx.commit().await.map_err(|_| unavailable())
|
||||
}
|
||||
|
||||
async fn settle_inner(&self, charge: &ClientBudgetCharge) -> Result<(), ClientBudgetError> {
|
||||
let mut tx = self.pool.begin().await.map_err(|_| ClientBudgetError)?;
|
||||
// 与准入统一先锁 Key,再写窗口和费用,串行化同一 Key 的并发结算。
|
||||
let key = sqlx::query_scalar::<_, String>(
|
||||
"select id from client_api_keys where id = $1 for update",
|
||||
)
|
||||
.bind(charge.key_id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| ClientBudgetError)?;
|
||||
let Some(key) = key else { return Ok(()) }; // 删除 Key 时也会删除其费用记录。
|
||||
settle_in_transaction(&mut tx, &key, charge)
|
||||
.await
|
||||
.map_err(|_| ClientBudgetError)?;
|
||||
tx.commit().await.map_err(|_| ClientBudgetError)
|
||||
}
|
||||
}
|
||||
|
||||
async fn settle_in_transaction(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
key: &str,
|
||||
charge: &ClientBudgetCharge,
|
||||
) -> Result<(), sqlx::Error> {
|
||||
advance_windows(tx, key, Utc::now()).await?;
|
||||
// 仅在请求结束时写入费用;请求 ID 冲突时不重复累计。
|
||||
let changed = sqlx::query(
|
||||
"insert into client_key_charge_events (request_id, client_api_key_id, amount_usd, completed_at)
|
||||
values ($1, $2, $3::text::numeric, $4)
|
||||
on conflict (request_id) do nothing",
|
||||
)
|
||||
.bind(charge.request_id.as_str())
|
||||
.bind(key)
|
||||
.bind(charge.amount_usd.canonical())
|
||||
.bind(DateTime::<Utc>::from(charge.completed_at))
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
if changed == 1 {
|
||||
sqlx::query("update client_key_budget_windows set
|
||||
daily_used_usd = daily_used_usd + case when $3 >= daily_start and $3 < daily_end then $2::text::numeric else 0 end,
|
||||
weekly_used_usd = weekly_used_usd + case when $3 >= weekly_start and $3 < weekly_end then $2::text::numeric else 0 end
|
||||
where client_api_key_id = $1")
|
||||
.bind(key).bind(charge.amount_usd.canonical()).bind(DateTime::<Utc>::from(charge.completed_at))
|
||||
.execute(&mut **tx).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for PgClientBudgetStore {
|
||||
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async move { self.admit_inner(key_id).await })
|
||||
}
|
||||
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
let result = self.settle_inner(&charge).await;
|
||||
let mut retry = self.retry.lock().map_err(|_| ClientBudgetError)?;
|
||||
if result.is_err() {
|
||||
retry.insert(charge.request_id.as_str().to_owned(), charge);
|
||||
} else {
|
||||
retry.remove(charge.request_id.as_str());
|
||||
}
|
||||
result
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn advance_windows(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
key: &str,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<(), sqlx::Error> {
|
||||
sqlx::query("insert into client_key_budget_windows
|
||||
(client_api_key_id, daily_start, daily_end, weekly_start, weekly_end)
|
||||
select $1, day, day + interval '24 hours', day, day + interval '168 hours'
|
||||
from (select date_trunc('day', $2::timestamptz at time zone 'Asia/Shanghai') at time zone 'Asia/Shanghai' as day) d
|
||||
on conflict (client_api_key_id) do update set
|
||||
daily_start = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_start else client_key_budget_windows.daily_start end,
|
||||
daily_end = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_end else client_key_budget_windows.daily_end end,
|
||||
daily_used_usd = case when client_key_budget_windows.daily_end <= $2 then 0 else client_key_budget_windows.daily_used_usd end,
|
||||
weekly_start = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_start else client_key_budget_windows.weekly_start end,
|
||||
weekly_end = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_end else client_key_budget_windows.weekly_end end,
|
||||
weekly_used_usd = case when client_key_budget_windows.weekly_end <= $2 then 0 else client_key_budget_windows.weekly_used_usd end")
|
||||
.bind(key).bind(now).execute(&mut **tx).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn load_client_key_budgets(
|
||||
pool: &PgPool,
|
||||
records: &mut [super::ClientApiKeyRecord],
|
||||
) -> StoreResult<()> {
|
||||
if records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let ids = records
|
||||
.iter()
|
||||
.map(|record| record.id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(
|
||||
"select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text,
|
||||
(case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used,
|
||||
(case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used,
|
||||
case when w.daily_end > now() then w.daily_end end as daily_end,
|
||||
case when w.weekly_end > now() then w.weekly_end end as weekly_end
|
||||
from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id
|
||||
where k.id = any($1)",
|
||||
)
|
||||
.bind(ids)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load client budgets"))?;
|
||||
let mut budgets = BTreeMap::new();
|
||||
for row in rows {
|
||||
let parse = |field| -> StoreResult<Decimal> {
|
||||
row.get::<String, _>(field)
|
||||
.parse()
|
||||
.map_err(|_| postgres_unavailable("decode client budget"))
|
||||
};
|
||||
budgets.insert(
|
||||
row.get::<String, _>("id"),
|
||||
ClientBudgetStatus {
|
||||
limits: ClientBudgetLimits {
|
||||
daily_usd: parse("daily_limit_usd")?,
|
||||
weekly_usd: parse("weekly_limit_usd")?,
|
||||
},
|
||||
daily_used_usd: parse("daily_used")?,
|
||||
weekly_used_usd: parse("weekly_used")?,
|
||||
daily_resets_at: row
|
||||
.get::<Option<DateTime<Utc>>, _>("daily_end")
|
||||
.map(Into::into),
|
||||
weekly_resets_at: row
|
||||
.get::<Option<DateTime<Utc>>, _>("weekly_end")
|
||||
.map(Into::into),
|
||||
},
|
||||
);
|
||||
}
|
||||
for record in records {
|
||||
record.budget = budgets
|
||||
.remove(&record.id)
|
||||
.ok_or_else(|| postgres_unavailable("load client budget policy"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unavailable() -> GatewayError {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::ProviderInfrastructureUnavailable,
|
||||
"client budget service is temporarily unavailable",
|
||||
)
|
||||
.with_client_code("key_budget_unavailable")
|
||||
}
|
||||
@@ -27,8 +27,10 @@ use gateway_admin::{
|
||||
ports::store::{AdminStoreResult, ClientKeyStore},
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
|
||||
engine::execution::ClientApiKeyUsageSink,
|
||||
lifecycle::CancellationToken,
|
||||
metering::Decimal,
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
task::{DaemonTask, WorkerTaskError},
|
||||
@@ -114,6 +116,7 @@ pub struct ClientApiKeyRecord {
|
||||
pub enabled: bool,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
pub budget: ClientBudgetStatus,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
@@ -271,6 +274,7 @@ pub struct NewClientApiKey {
|
||||
pub key: String,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
pub budget: ClientBudgetLimits,
|
||||
}
|
||||
|
||||
impl fmt::Debug for NewClientApiKey {
|
||||
@@ -301,6 +305,8 @@ pub struct UpdateClientApiKeyDetails {
|
||||
pub group_ids: Vec<String>,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
pub daily_limit_usd: Option<Decimal>,
|
||||
pub weekly_limit_usd: Option<Decimal>,
|
||||
}
|
||||
|
||||
impl UpdateClientApiKeyDetails {
|
||||
@@ -383,6 +389,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
None
|
||||
};
|
||||
load_client_key_memberships(&self.pool, &mut items).await?;
|
||||
super::client_budgets::load_client_key_budgets(&self.pool, &mut items).await?;
|
||||
Ok(ClientApiKeyPage {
|
||||
items,
|
||||
total,
|
||||
@@ -406,7 +413,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
|
||||
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
sqlx::query(
|
||||
let record = sqlx::query(
|
||||
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
|
||||
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
|
||||
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
|
||||
@@ -445,7 +452,16 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
.map_err(|_| postgres_unavailable("get client API key"))?
|
||||
.as_ref()
|
||||
.map(client_record_from_row)
|
||||
.transpose()
|
||||
.transpose()?;
|
||||
let Some(mut record) = record else {
|
||||
return Ok(None);
|
||||
};
|
||||
super::client_budgets::load_client_key_budgets(
|
||||
&self.pool,
|
||||
std::slice::from_mut(&mut record),
|
||||
)
|
||||
.await?;
|
||||
Ok(Some(record))
|
||||
}
|
||||
|
||||
async fn touch_client_api_keys(
|
||||
@@ -719,6 +735,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
.map(|id| id.as_str().to_owned())
|
||||
.collect(),
|
||||
key: command.plaintext,
|
||||
budget: command.budget,
|
||||
max_concurrency: command.limits.max_concurrency,
|
||||
requests_per_minute: command.limits.requests_per_minute,
|
||||
},
|
||||
@@ -735,6 +752,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
"enabled",
|
||||
"max_concurrency",
|
||||
"requests_per_minute",
|
||||
"daily_limit_usd",
|
||||
"weekly_limit_usd",
|
||||
]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
@@ -766,6 +785,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
.collect(),
|
||||
max_concurrency: command.limits.max_concurrency,
|
||||
requests_per_minute: command.limits.requests_per_minute,
|
||||
daily_limit_usd: command.daily_limit_usd,
|
||||
weekly_limit_usd: command.weekly_limit_usd,
|
||||
},
|
||||
mutation_audit(
|
||||
context,
|
||||
@@ -778,6 +799,8 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
"group_ids",
|
||||
"max_concurrency",
|
||||
"requests_per_minute",
|
||||
"daily_limit_usd",
|
||||
"weekly_limit_usd",
|
||||
]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
@@ -942,6 +965,7 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<Admin
|
||||
max_concurrency: record.max_concurrency,
|
||||
requests_per_minute: record.requests_per_minute,
|
||||
},
|
||||
budget: record.budget,
|
||||
last_used_at: record.last_used_at,
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
@@ -956,8 +980,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, key, enabled, max_concurrency, requests_per_minute,
|
||||
last_used_at, created_at, updated_at
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
|
||||
last_used_at, created_at, updated_at, daily_limit_usd, weekly_limit_usd
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric)",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
@@ -965,6 +989,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
.bind(&key.key)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.bind(key.budget.daily_usd.canonical())
|
||||
.bind(key.budget.weekly_usd.canonical())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
|
||||
@@ -980,7 +1006,9 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
let result = sqlx::query(
|
||||
"update client_api_keys
|
||||
set name = $2, label = $3, max_concurrency = $4,
|
||||
requests_per_minute = $5, updated_at = now()
|
||||
requests_per_minute = $5, updated_at = now(),
|
||||
daily_limit_usd = coalesce($6::text::numeric, daily_limit_usd),
|
||||
weekly_limit_usd = coalesce($7::text::numeric, weekly_limit_usd)
|
||||
where id = $1",
|
||||
)
|
||||
.bind(&key.id)
|
||||
@@ -988,6 +1016,8 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
.bind(&key.label)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.bind(key.daily_limit_usd.map(|amount| amount.canonical()))
|
||||
.bind(key.weekly_limit_usd.map(|amount| amount.canonical()))
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
|
||||
@@ -1142,6 +1172,7 @@ fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiK
|
||||
row.try_get("requests_per_minute")
|
||||
.map_err(|_| invalid("invalid requests per minute"))?,
|
||||
)?,
|
||||
budget: ClientBudgetStatus::default(),
|
||||
last_used_at: row
|
||||
.try_get("last_used_at")
|
||||
.map_err(|_| invalid("invalid last used at"))?,
|
||||
|
||||
@@ -15,12 +15,14 @@ mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod execution_buffer;
|
||||
mod observability;
|
||||
mod ops_events;
|
||||
mod provider_accounts;
|
||||
mod proxies;
|
||||
mod retention;
|
||||
mod runtime_settings;
|
||||
mod snapshot;
|
||||
@@ -30,12 +32,14 @@ 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::*;
|
||||
pub use observability::*;
|
||||
pub use ops_events::*;
|
||||
pub use provider_accounts::*;
|
||||
pub use proxies::PgProxyRepository;
|
||||
pub use retention::*;
|
||||
pub use runtime_settings::*;
|
||||
pub use snapshot::*;
|
||||
|
||||
@@ -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,8 +213,10 @@ impl PgAdminAccountStore {
|
||||
async fn commit_prepared_import(
|
||||
&self,
|
||||
prepared: PreparedCredentialImport,
|
||||
settings: Option<AccountImportSettings>,
|
||||
context: &MutationContext,
|
||||
action: &str,
|
||||
outbound_proxy: Option<gateway_admin::model::proxies::ImportProxyBinding>,
|
||||
) -> AdminStoreResult<CredentialImportResult> {
|
||||
let provider_kind = prepared.provider_kind.as_str().to_owned();
|
||||
let accounts = prepared
|
||||
@@ -223,9 +225,16 @@ 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,
|
||||
outbound_proxy,
|
||||
scope: ProviderAccountAdminScope {
|
||||
provider_kind: provider_kind.clone(),
|
||||
},
|
||||
@@ -235,7 +244,7 @@ impl PgAdminAccountStore {
|
||||
action,
|
||||
"provider_account",
|
||||
&provider_kind,
|
||||
vec!["credentials".to_owned()],
|
||||
changed_fields,
|
||||
),
|
||||
})
|
||||
.await
|
||||
@@ -615,8 +624,14 @@ 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",
|
||||
command.outbound_proxy,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn commit_authorization(
|
||||
@@ -635,8 +650,19 @@ impl AccountStore for PgAdminAccountStore {
|
||||
provider_kind: credential.provider_kind.clone(),
|
||||
credentials: vec![credential],
|
||||
},
|
||||
command.settings,
|
||||
context,
|
||||
"authorize",
|
||||
command
|
||||
.pending
|
||||
.outbound_proxy_id()
|
||||
.zip(command.pending.outbound_proxy())
|
||||
.map(
|
||||
|(id, proxy)| gateway_admin::model::proxies::ImportProxyBinding {
|
||||
id: id.to_owned(),
|
||||
proxy: proxy.clone(),
|
||||
},
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
let [account_id]: [CoreProviderAccountId; 1] =
|
||||
@@ -666,6 +692,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 +735,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 +752,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 +828,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 +845,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,
|
||||
@@ -133,15 +133,28 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
} else {
|
||||
CredentialState::Unknown
|
||||
};
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("begin provider account insert"))?;
|
||||
let proxy_id = match account.outbound_proxy.as_ref() {
|
||||
Some(proxy) => Some(
|
||||
super::super::proxies::ensure_proxy_for_url(&mut transaction, proxy, None)
|
||||
.await?
|
||||
.0,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
sqlx::query(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
outbound_proxy_url, outbound_proxy_id, 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, $19, $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,9 +175,15 @@ 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)
|
||||
.execute(&self.pool)
|
||||
.bind(account.outbound_proxy.as_ref().map(|proxy| proxy.expose_url()))
|
||||
.bind(proxy_id)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert provider account"))?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("commit provider account insert"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -442,11 +461,50 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
|
||||
.map_err(|_| postgres_unavailable("begin provider account admin import"))?;
|
||||
let result = async {
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
|
||||
if let Some(binding) = &command.outbound_proxy {
|
||||
let (_, current) = super::super::proxies::resolve_proxy_selection(
|
||||
&mut transaction,
|
||||
&gateway_admin::model::proxies::AccountProxySelection::Saved(
|
||||
binding.id.clone(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
if current.as_ref() != Some(&binding.proxy) {
|
||||
return Err(StoreError::Conflict {
|
||||
entity: "outbound proxy",
|
||||
id: binding.id.clone(),
|
||||
kind: ConflictKind::StaleRevision,
|
||||
});
|
||||
}
|
||||
}
|
||||
let mut account_ids = Vec::with_capacity(command.accounts.len());
|
||||
for account in &command.accounts {
|
||||
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 +582,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
|
||||
command.enabled,
|
||||
command.concurrency_limit,
|
||||
command.weight,
|
||||
command.outbound_proxy.as_ref(),
|
||||
)
|
||||
.await?;
|
||||
replace_account_group_assignments_in_transaction(
|
||||
@@ -676,15 +735,23 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
} else {
|
||||
CredentialState::Unknown
|
||||
};
|
||||
let proxy_id = match account.outbound_proxy.as_ref() {
|
||||
Some(proxy) => Some(
|
||||
super::super::proxies::ensure_proxy_for_url(transaction, proxy, None)
|
||||
.await?
|
||||
.0,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
let imported_id = sqlx::query_scalar::<_, String>(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
outbound_proxy_url, outbound_proxy_id, 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, $19, $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 +764,8 @@ 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),
|
||||
outbound_proxy_id = coalesce(excluded.outbound_proxy_id, provider_accounts.outbound_proxy_id),
|
||||
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 +801,8 @@ 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()))
|
||||
.bind(proxy_id)
|
||||
.fetch_optional(&mut **transaction)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
@@ -838,10 +909,19 @@ pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
|
||||
enabled: bool,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
outbound_proxy: Option<&gateway_admin::model::proxies::AccountProxySelection>,
|
||||
) -> StoreResult<()> {
|
||||
let (proxy_id, proxy) = match outbound_proxy {
|
||||
Some(selection) => {
|
||||
super::super::proxies::resolve_proxy_selection(transaction, selection).await?
|
||||
}
|
||||
None => (None, None),
|
||||
};
|
||||
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,
|
||||
outbound_proxy_id = case when $5 then $7 else outbound_proxy_id end
|
||||
where id = any($1::text[])
|
||||
returning id",
|
||||
)
|
||||
@@ -849,6 +929,9 @@ 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(proxy.as_ref().map(gateway_core::account::OutboundProxy::expose_url))
|
||||
.bind(proxy_id)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))?
|
||||
@@ -937,7 +1020,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,8 @@ impl NewProviderAccount {
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ImportProviderAccounts {
|
||||
pub settings: Option<AccountImportSettings>,
|
||||
pub outbound_proxy: Option<gateway_admin::model::proxies::ImportProxyBinding>,
|
||||
pub scope: ProviderAccountAdminScope,
|
||||
pub accounts: Vec<NewProviderAccount>,
|
||||
pub audit: AdminAuditEvent,
|
||||
@@ -261,6 +265,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 +314,7 @@ impl fmt::Debug for RotateProviderAccount {
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchUpdateProviderAccountsAdmin {
|
||||
pub outbound_proxy: Option<gateway_admin::model::proxies::AccountProxySelection>,
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
@@ -372,7 +380,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 +388,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 +398,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 +473,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 +532,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")?,
|
||||
|
||||
@@ -0,0 +1,678 @@
|
||||
//! 命名代理持有认证信息,账号保留解析后的 URL 供 Provider 传输使用。
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext, Revision as AdminRevision,
|
||||
account_groups::{AccountGroupColor, AccountGroupRef},
|
||||
proxies::*,
|
||||
},
|
||||
ports::{
|
||||
proxy::{ProxyImportGuard, ProxyImportReservation, ProxyStore},
|
||||
store::{AdminStoreError, AdminStoreResult},
|
||||
},
|
||||
};
|
||||
use gateway_core::{
|
||||
account::{OutboundProxy, ProviderAccountId},
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
use sqlx::{PgConnection, PgPool, Postgres, QueryBuilder, Row as _, Transaction, postgres::PgRow};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use super::{append_admin_audit_event_in_transaction, bump_config_revision_in_transaction};
|
||||
use crate::{
|
||||
ConflictKind, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
|
||||
mutation_audit, postgres_unavailable,
|
||||
};
|
||||
|
||||
const ENTITY: &str = "outbound proxy";
|
||||
const SELECT: &str = "select p.*,
|
||||
(select count(*) from provider_accounts a where a.outbound_proxy_id = p.id) as account_count
|
||||
from outbound_proxies p";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PgProxyRepository {
|
||||
pool: PgPool,
|
||||
import_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl PgProxyRepository {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
import_slots: Arc::new(Semaphore::new(4)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct PgProxyImportGuard {
|
||||
// 独立会话持有咨询锁,避免空闲事务超时,也不给提交事务占用的连接池制造死锁。
|
||||
// 每个仓储最多增加四个会话;取消请求或进程退出时,关闭连接会自动释放锁。
|
||||
_connection: PgConnection,
|
||||
_slot: OwnedSemaphorePermit,
|
||||
}
|
||||
|
||||
impl ProxyImportGuard for PgProxyImportGuard {}
|
||||
|
||||
async fn exclude_active_imports(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
id: &str,
|
||||
) -> StoreResult<()> {
|
||||
let acquired: bool =
|
||||
sqlx::query_scalar("select pg_try_advisory_xact_lock(hashtextextended($1, 739219))")
|
||||
.bind(id)
|
||||
.fetch_one(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
if !acquired {
|
||||
return Err(conflict(id));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn store_error(error: StoreError) -> AdminStoreError {
|
||||
admin_store_error(ENTITY, error)
|
||||
}
|
||||
fn unavailable() -> StoreError {
|
||||
postgres_unavailable("outbound proxy operation")
|
||||
}
|
||||
fn conflict(id: &str) -> StoreError {
|
||||
StoreError::Conflict {
|
||||
entity: ENTITY,
|
||||
id: id.to_owned(),
|
||||
kind: ConflictKind::InvalidTransition,
|
||||
}
|
||||
}
|
||||
fn invalid() -> StoreError {
|
||||
StoreError::InvalidData {
|
||||
entity: ENTITY,
|
||||
message: "invalid proxy record".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn record(row: PgRow) -> StoreResult<ProxyRecord> {
|
||||
let success: Option<bool> = row.try_get("last_test_success").map_err(|_| invalid())?;
|
||||
let ip: Option<String> = row.try_get("last_test_ip").map_err(|_| invalid())?;
|
||||
let latency: Option<i64> = row.try_get("last_test_latency_ms").map_err(|_| invalid())?;
|
||||
Ok(ProxyRecord {
|
||||
id: row.try_get("id").map_err(|_| invalid())?,
|
||||
name: row.try_get("name").map_err(|_| invalid())?,
|
||||
proxy: OutboundProxy::parse(
|
||||
&row.try_get::<String, _>("proxy_url")
|
||||
.map_err(|_| invalid())?,
|
||||
)
|
||||
.map_err(|_| invalid())?,
|
||||
revision: AdminRevision::new(
|
||||
u64::try_from(row.try_get::<i64, _>("revision").map_err(|_| invalid())?)
|
||||
.map_err(|_| invalid())?,
|
||||
)
|
||||
.map_err(|_| invalid())?,
|
||||
account_count: u64::try_from(
|
||||
row.try_get::<i64, _>("account_count")
|
||||
.map_err(|_| invalid())?,
|
||||
)
|
||||
.map_err(|_| invalid())?,
|
||||
last_test_at: row.try_get("last_test_at").map_err(|_| invalid())?,
|
||||
last_test: success
|
||||
.map(|success| -> StoreResult<_> {
|
||||
Ok(ProxyTestResult {
|
||||
success,
|
||||
latency_ms: u64::try_from(latency.ok_or_else(invalid)?)
|
||||
.map_err(|_| invalid())?,
|
||||
exit_ip: ip.map(|ip| ip.parse().map_err(|_| invalid())).transpose()?,
|
||||
message: row
|
||||
.try_get::<Option<String>, _>("last_test_message")
|
||||
.map_err(|_| invalid())?
|
||||
.unwrap_or_default(),
|
||||
})
|
||||
})
|
||||
.transpose()?,
|
||||
created_at: row.try_get("created_at").map_err(|_| invalid())?,
|
||||
updated_at: row.try_get("updated_at").map_err(|_| invalid())?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn lock_url(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
proxy: &OutboundProxy,
|
||||
) -> StoreResult<()> {
|
||||
sqlx::query("select pg_advisory_xact_lock(hashtextextended($1, 739218))")
|
||||
.bind(proxy.expose_url())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 导入和旧版 URL 写入在账号事务内登记到共享代理目录。
|
||||
pub(crate) async fn ensure_proxy_for_url(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
proxy: &OutboundProxy,
|
||||
name: Option<&str>,
|
||||
) -> StoreResult<(String, bool)> {
|
||||
lock_url(transaction, proxy).await?;
|
||||
if let Some(id) = sqlx::query_scalar::<_, String>("select id from outbound_proxies where proxy_url = $1 order by created_at, id limit 1 for share")
|
||||
.bind(proxy.expose_url()).fetch_optional(&mut **transaction).await.map_err(|_| unavailable())? {
|
||||
return Ok((id, false));
|
||||
}
|
||||
let id = format!("proxy_{}", uuid::Uuid::now_v7().simple());
|
||||
let generated_name = format!("Imported proxy {}", &id[id.len() - 8..]);
|
||||
sqlx::query("insert into outbound_proxies (id, name, proxy_url) values ($1, $2, $3)")
|
||||
.bind(&id)
|
||||
.bind(name.unwrap_or(&generated_name))
|
||||
.bind(proxy.expose_url())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
Ok((id, true))
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_proxy_selection(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
selection: &AccountProxySelection,
|
||||
) -> StoreResult<(Option<String>, Option<OutboundProxy>)> {
|
||||
match selection {
|
||||
AccountProxySelection::Direct => Ok((None, None)),
|
||||
AccountProxySelection::Url(proxy) => {
|
||||
let (id, _) = ensure_proxy_for_url(transaction, proxy, None).await?;
|
||||
Ok((Some(id), Some(proxy.clone())))
|
||||
}
|
||||
AccountProxySelection::Saved(id) => {
|
||||
let (value, tested) = sqlx::query_as::<_, (String, Option<bool>)>(
|
||||
"select proxy_url, last_test_success from outbound_proxies where id = $1 for share",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable())?
|
||||
.ok_or_else(|| StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: id.clone(),
|
||||
})?;
|
||||
if tested != Some(true) {
|
||||
return Err(conflict(id));
|
||||
}
|
||||
Ok((
|
||||
Some(id.clone()),
|
||||
Some(OutboundProxy::parse(&value).map_err(|_| invalid())?),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn audit(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
context: &MutationContext,
|
||||
action: &str,
|
||||
id: &str,
|
||||
fields: &[&str],
|
||||
revision: Revision,
|
||||
) -> StoreResult<()> {
|
||||
append_admin_audit_event_in_transaction(
|
||||
transaction,
|
||||
mutation_audit(
|
||||
context,
|
||||
action,
|
||||
"outbound_proxy",
|
||||
id,
|
||||
fields.iter().map(|value| (*value).to_owned()).collect(),
|
||||
),
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyStore for PgProxyRepository {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &ProviderAccountId,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AdminRevision> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
// 在同一条更新中校验绑定,避免旧弹窗清除账号后来选择的其他代理。
|
||||
let updated = sqlx::query(
|
||||
"update provider_accounts
|
||||
set outbound_proxy_id = null, outbound_proxy_url = null,
|
||||
updated_at = greatest(now(), updated_at)
|
||||
where id = $1 and outbound_proxy_id = $2",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
.bind(proxy_id)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(proxy_id)));
|
||||
}
|
||||
append_admin_audit_event_in_transaction(
|
||||
&mut transaction,
|
||||
mutation_audit(
|
||||
context,
|
||||
"update",
|
||||
"provider_account",
|
||||
account_id.as_str(),
|
||||
vec!["outbound_proxy".to_owned()],
|
||||
),
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
admin_revision(revision)
|
||||
}
|
||||
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage> {
|
||||
if query.page == 0 {
|
||||
return Err(store_error(invalid()));
|
||||
}
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
// 数量和当前页共用只读快照,避免绑定变化使同一次响应的分页事实不一致。
|
||||
sqlx::query("set transaction isolation level repeatable read, read only")
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let total: i64 = sqlx::query_scalar(
|
||||
"select count(a.id) from outbound_proxies p
|
||||
left join provider_accounts a on a.outbound_proxy_id = p.id
|
||||
and (strpos(lower(a.name), lower($2)) > 0
|
||||
or strpos(lower(coalesce(a.email, '')), lower($2)) > 0)
|
||||
where p.id = $1 group by p.id",
|
||||
)
|
||||
.bind(&query.proxy_id)
|
||||
.bind(&query.search)
|
||||
.fetch_optional(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?
|
||||
.ok_or_else(|| {
|
||||
store_error(StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: query.proxy_id.clone(),
|
||||
})
|
||||
})?;
|
||||
let mut items = sqlx::query(
|
||||
"select id, name, email, provider_kind, authentication_kind, plan_type, enabled from provider_accounts
|
||||
where outbound_proxy_id = $1 and (strpos(lower(name), lower($2)) > 0
|
||||
or strpos(lower(coalesce(email, '')), lower($2)) > 0)
|
||||
order by name, id limit $3 offset $4",
|
||||
)
|
||||
.bind(&query.proxy_id)
|
||||
.bind(&query.search)
|
||||
.bind(i64::from(query.page_size.get()))
|
||||
.bind(i64::from(query.page - 1) * i64::from(query.page_size.get()))
|
||||
.fetch_all(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(ProxyAccountRef {
|
||||
id: row.try_get("id").map_err(|_| invalid())?,
|
||||
name: row.try_get("name").map_err(|_| invalid())?,
|
||||
email: row.try_get("email").map_err(|_| invalid())?,
|
||||
provider_kind: row.try_get("provider_kind").map_err(|_| invalid())?,
|
||||
authentication_kind: row.try_get("authentication_kind").map_err(|_| invalid())?,
|
||||
plan_type: row.try_get("plan_type").map_err(|_| invalid())?,
|
||||
plan_type_display: None,
|
||||
groups: Vec::new(),
|
||||
enabled: row.try_get("enabled").map_err(|_| invalid())?,
|
||||
})
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(store_error)?;
|
||||
if !items.is_empty() {
|
||||
// 仅批量加载当前页的分组,并沿用同一快照,避免逐账号查询。
|
||||
let account_ids: Vec<&str> = items.iter().map(|account| account.id.as_str()).collect();
|
||||
let rows = sqlx::query_as::<_, (String, String, String, String, bool)>(
|
||||
"select m.provider_account_id, g.id, g.name, g.color, g.enabled
|
||||
from account_group_accounts m join account_groups g on g.id = m.account_group_id
|
||||
where m.provider_account_id = any($1::text[])
|
||||
order by m.provider_account_id, g.id",
|
||||
)
|
||||
.bind(&account_ids)
|
||||
.fetch_all(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let mut groups = BTreeMap::<String, Vec<AccountGroupRef>>::new();
|
||||
for (account_id, group_id, name, color, enabled) in rows {
|
||||
groups.entry(account_id).or_default().push(AccountGroupRef {
|
||||
id: AccountGroupId::new(group_id).map_err(|_| store_error(invalid()))?,
|
||||
name,
|
||||
color: AccountGroupColor::parse(&color)
|
||||
.ok_or_else(|| store_error(invalid()))?,
|
||||
enabled,
|
||||
});
|
||||
}
|
||||
for account in &mut items {
|
||||
account.groups = groups.remove(&account.id).unwrap_or_default();
|
||||
}
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
Ok(ProxyAccountPage {
|
||||
items,
|
||||
total: u64::try_from(total).map_err(|_| store_error(invalid()))?,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn reserve_import(&self, id: &str) -> AdminStoreResult<ProxyImportReservation> {
|
||||
let slot = self
|
||||
.import_slots
|
||||
.clone()
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| store_error(conflict(id)))?;
|
||||
let mut connection = self
|
||||
.pool
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?
|
||||
.detach();
|
||||
let acquired: bool =
|
||||
sqlx::query_scalar("select pg_try_advisory_lock_shared(hashtextextended($1, 739219))")
|
||||
.bind(id)
|
||||
.fetch_one(&mut connection)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
if !acquired {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let record = self.get(id).await?;
|
||||
if !record.last_test.is_some_and(|test| test.success) {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
Ok(ProxyImportReservation {
|
||||
binding: ImportProxyBinding {
|
||||
id: record.id,
|
||||
proxy: record.proxy,
|
||||
},
|
||||
guard: Box::new(PgProxyImportGuard {
|
||||
_connection: connection,
|
||||
_slot: slot,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list(&self, query: ProxyListQuery) -> AdminStoreResult<ProxyPage> {
|
||||
if query.page == 0 {
|
||||
return Err(store_error(invalid()));
|
||||
}
|
||||
let total: i64 = sqlx::query_scalar(
|
||||
"select count(*) from outbound_proxies where strpos(lower(name), lower($1)) > 0",
|
||||
)
|
||||
.bind(&query.search)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let mut builder = QueryBuilder::<Postgres>::new(SELECT);
|
||||
builder
|
||||
.push(" where strpos(lower(p.name), lower(")
|
||||
.push_bind(&query.search)
|
||||
.push(")) > 0 order by p.created_at desc, p.id limit ")
|
||||
.push_bind(i64::from(query.page_size.get()))
|
||||
.push(" offset ")
|
||||
.push_bind(i64::from(query.page - 1) * i64::from(query.page_size.get()));
|
||||
let items = builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?
|
||||
.into_iter()
|
||||
.map(record)
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(store_error)?;
|
||||
Ok(ProxyPage {
|
||||
items,
|
||||
total: u64::try_from(total).map_err(|_| store_error(invalid()))?,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn get(&self, id: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(SELECT);
|
||||
builder.push(" where p.id = ").push_bind(id);
|
||||
let row = builder
|
||||
.build()
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?
|
||||
.ok_or_else(|| {
|
||||
store_error(StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: id.to_owned(),
|
||||
})
|
||||
})?;
|
||||
record(row).map_err(store_error)
|
||||
}
|
||||
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let (id, created) =
|
||||
ensure_proxy_for_url(&mut transaction, &command.proxy, Some(&command.name))
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
if !created {
|
||||
return Err(store_error(conflict(&id)));
|
||||
}
|
||||
audit(
|
||||
&mut transaction,
|
||||
context,
|
||||
"create",
|
||||
&id,
|
||||
&["name", "proxy_url"],
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
Ok(ProxyMutation {
|
||||
config_revision: admin_revision(revision)?,
|
||||
record: self.get(&id).await?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
exclude_active_imports(&mut transaction, &command.id)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
if let Some(proxy) = &command.proxy {
|
||||
lock_url(&mut transaction, proxy)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let duplicate: bool = sqlx::query_scalar(
|
||||
"select exists(select 1 from outbound_proxies where proxy_url = $1 and id <> $2)",
|
||||
)
|
||||
.bind(proxy.expose_url())
|
||||
.bind(&command.id)
|
||||
.fetch_one(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
if duplicate {
|
||||
return Err(store_error(conflict(&command.id)));
|
||||
}
|
||||
}
|
||||
let changed = sqlx::query(
|
||||
"update outbound_proxies set name = $3, proxy_url = coalesce($4, proxy_url), revision = revision + 1, updated_at = now(),
|
||||
last_test_at = case when $4 is not null and $4 <> proxy_url then null else last_test_at end,
|
||||
last_test_success = case when $4 is not null and $4 <> proxy_url then null else last_test_success end,
|
||||
last_test_latency_ms = case when $4 is not null and $4 <> proxy_url then null else last_test_latency_ms end,
|
||||
last_test_ip = case when $4 is not null and $4 <> proxy_url then null else last_test_ip end,
|
||||
last_test_message = case when $4 is not null and $4 <> proxy_url then null else last_test_message end
|
||||
where id = $1 and revision = $2")
|
||||
.bind(&command.id).bind(i64::try_from(command.revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.bind(&command.name).bind(command.proxy.as_ref().map(OutboundProxy::expose_url))
|
||||
.execute(&mut *transaction).await.map_err(|_| store_error(unavailable()))?;
|
||||
if changed.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(&command.id)));
|
||||
}
|
||||
sqlx::query("update provider_accounts a set outbound_proxy_url = p.proxy_url, updated_at = greatest(now(), a.updated_at) from outbound_proxies p where p.id = $1 and a.outbound_proxy_id = p.id and a.outbound_proxy_url is distinct from p.proxy_url")
|
||||
.bind(&command.id).execute(&mut *transaction).await.map_err(|_| store_error(unavailable()))?;
|
||||
audit(
|
||||
&mut transaction,
|
||||
context,
|
||||
"update",
|
||||
&command.id,
|
||||
&["name", "proxy_url"],
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
Ok(ProxyMutation {
|
||||
config_revision: admin_revision(revision)?,
|
||||
record: self.get(&command.id).await?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: AdminRevision,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AdminRevision> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
exclude_active_imports(&mut transaction, id)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let config_revision = bump_config_revision_in_transaction(&mut transaction)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let result = sqlx::query("delete from outbound_proxies where id = $1 and revision = $2")
|
||||
.bind(id)
|
||||
.bind(i64::try_from(revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
// PostgreSQL 18 为 RESTRICT 返回不同于普通外键违规的错误码。
|
||||
if error.as_database_error().is_some_and(|error| {
|
||||
error.is_foreign_key_violation() || error.code().as_deref() == Some("23001")
|
||||
}) {
|
||||
store_error(conflict(id))
|
||||
} else {
|
||||
store_error(unavailable())
|
||||
}
|
||||
})?;
|
||||
if result.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
audit(
|
||||
&mut transaction,
|
||||
context,
|
||||
"delete",
|
||||
id,
|
||||
&[],
|
||||
config_revision,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
admin_revision(config_revision)
|
||||
}
|
||||
|
||||
async fn record_test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: AdminRevision,
|
||||
result: ProxyTestResult,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
exclude_active_imports(&mut transaction, id)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let updated = sqlx::query("update outbound_proxies set last_test_at = now(), last_test_success = $3, last_test_latency_ms = $4, last_test_ip = $5, last_test_message = $6 where id = $1 and revision = $2")
|
||||
.bind(id).bind(i64::try_from(revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.bind(result.success).bind(i64::try_from(result.latency_ms).map_err(|_| store_error(invalid()))?)
|
||||
.bind(result.exit_ip.map(|ip| ip.to_string())).bind(result.message)
|
||||
.execute(&mut *transaction).await.map_err(|_| store_error(unavailable()))?;
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let current: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
let current = Revision::new(u64::try_from(current).map_err(|_| store_error(invalid()))?)
|
||||
.map_err(store_error)?;
|
||||
audit(
|
||||
&mut transaction,
|
||||
context,
|
||||
"test",
|
||||
id,
|
||||
&["last_test"],
|
||||
current,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
self.get(id).await
|
||||
}
|
||||
}
|
||||
@@ -168,6 +168,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
let (scope_revision, widened) = keys
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: None,
|
||||
id: client_key_id("key_group_one"),
|
||||
name: "key_group_one".to_owned(),
|
||||
label: None,
|
||||
@@ -191,6 +193,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
let (restricted_revision, restricted) = keys
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: None,
|
||||
id: client_key_id("key_group_one"),
|
||||
name: "key_group_one".to_owned(),
|
||||
label: None,
|
||||
@@ -311,6 +315,7 @@ async fn group_costs_should_include_statusless_websocket_but_reject_statusless_h
|
||||
fn new_key(id: &str, group_ids: Vec<AccountGroupId>) -> NewClientKey {
|
||||
let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k'));
|
||||
NewClientKey {
|
||||
budget: Default::default(),
|
||||
id: client_key_id(id),
|
||||
name: id.to_owned(),
|
||||
label: None,
|
||||
|
||||
@@ -0,0 +1,337 @@
|
||||
use std::time::SystemTime;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_admin::{
|
||||
model::{MutationActor, MutationContext, client_keys::UpdateClientKey},
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
ModelRequestId,
|
||||
budget::{ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus},
|
||||
},
|
||||
error::GatewayErrorKind,
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
ClientApiKeyRepository as _, PgAdminClientKeyStore, PgClientApiKeyRepository,
|
||||
PgClientBudgetStore,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
fn key_id(key: &str) -> ClientApiKeyId {
|
||||
ClientApiKeyId::new(key).unwrap()
|
||||
}
|
||||
|
||||
fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge {
|
||||
ClientBudgetCharge {
|
||||
key_id: key_id(key),
|
||||
request_id: ModelRequestId::new(format!("req_{request}")).unwrap(),
|
||||
amount_usd: amount.parse().unwrap(),
|
||||
completed_at: SystemTime::now(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, key, daily_limit_usd, weekly_limit_usd, created_at, updated_at)
|
||||
values ($1, $1, $2, $3::text::numeric, $4::text::numeric, now(), now())",
|
||||
)
|
||||
.bind(key)
|
||||
.bind(format!("sk_{key:a<43}"))
|
||||
.bind(daily)
|
||||
.bind(weekly)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus {
|
||||
PgClientApiKeyRepository::new(database.pool.clone())
|
||||
.get_client_api_key(key)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.budget
|
||||
}
|
||||
|
||||
fn context() -> MutationContext {
|
||||
MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "budget-test".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_instances() {
|
||||
let Some(database) = TestDatabase::create("budgets_exact").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "day", "0.3", "2").await;
|
||||
seed(&database, "week", "2", "0.2").await;
|
||||
let first = PgClientBudgetStore::new(database.pool.clone());
|
||||
let second = PgClientBudgetStore::new(database.pool.clone());
|
||||
for (key, prefix, amount) in [("day", "d", "0.1"), ("week", "w", "0.1")] {
|
||||
for _ in 0..3 {
|
||||
// 已准入请求可完成并超过限额,不预占估算费用。
|
||||
first.admit(key_id(key)).await.unwrap();
|
||||
}
|
||||
for index in 0..3 {
|
||||
let id = format!("{prefix}-{index}");
|
||||
let (a, b) = tokio::join!(
|
||||
first.settle(charge(key, &id, amount)),
|
||||
second.settle(charge(key, &id, amount))
|
||||
);
|
||||
a.unwrap();
|
||||
b.unwrap();
|
||||
}
|
||||
}
|
||||
let day = status(&database, "day").await;
|
||||
assert_eq!(day.daily_used_usd.canonical(), "0.3");
|
||||
assert_eq!(day.weekly_used_usd.canonical(), "0.3");
|
||||
let day_error = second.admit(key_id("day")).await.unwrap_err();
|
||||
assert_eq!(day_error.kind(), GatewayErrorKind::RateLimited);
|
||||
assert_eq!(
|
||||
day_error.client_error_code(),
|
||||
Some("key_daily_budget_exceeded")
|
||||
);
|
||||
assert!(day_error.retry_after().is_some());
|
||||
let week_error = first.admit(key_id("week")).await.unwrap_err();
|
||||
assert_eq!(
|
||||
week_error.client_error_code(),
|
||||
Some("key_weekly_budget_exceeded")
|
||||
);
|
||||
let event_count: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(event_count, 6, "rejected admission must not create charges");
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlement() {
|
||||
let Some(database) = TestDatabase::create("budgets_windows").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "1", "2").await;
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
let (day_start, day_end, week_end): (DateTime<Utc>, DateTime<Utc>, DateTime<Utc>) =
|
||||
sqlx::query_as("select daily_start, daily_end, weekly_end from client_key_budget_windows where client_api_key_id = 'key'")
|
||||
.fetch_one(&database.pool).await.unwrap();
|
||||
assert_eq!(day_start.timestamp().rem_euclid(86400), 16 * 3600);
|
||||
assert_eq!((day_end - day_start).num_hours(), 24);
|
||||
assert_eq!((week_end - day_start).num_hours(), 168);
|
||||
store.settle(charge("key", "first", "1")).await.unwrap();
|
||||
sqlx::query("update client_key_budget_windows set daily_start = daily_start - interval '1 day', daily_end = daily_start")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0"
|
||||
);
|
||||
assert_eq!(
|
||||
status(&database, "key").await.weekly_used_usd.canonical(),
|
||||
"1"
|
||||
);
|
||||
store.settle(charge("key", "second", "1")).await.unwrap();
|
||||
sqlx::query("update client_key_budget_windows set daily_end = now() - interval '1 second', weekly_end = now() - interval '1 second'")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
let virtual_reset = status(&database, "key").await;
|
||||
assert_eq!(virtual_reset.daily_used_usd.canonical(), "0");
|
||||
assert_eq!(virtual_reset.weekly_used_usd.canonical(), "0");
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
let old = ClientBudgetCharge {
|
||||
completed_at: (day_start - chrono::Duration::seconds(1)).into(),
|
||||
..charge("key", "after-reset", "0.9")
|
||||
};
|
||||
store.settle(old).await.unwrap();
|
||||
let reset = status(&database, "key").await;
|
||||
assert_eq!(reset.daily_used_usd.canonical(), "0");
|
||||
assert_eq!(reset.weekly_used_usd.canonical(), "0");
|
||||
assert!(reset.weekly_resets_at.is_some());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_cost_and_interrupted_requests_never_block_limited_keys() {
|
||||
let Some(database) = TestDatabase::create("budgets_zero").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "1", "5").await;
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
store.settle(charge("key", "no-cost", "0")).await.unwrap();
|
||||
// 模拟准入后进程退出,没有费用可结算;重启后仍应允许同一 Key 使用。
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
let restarted = PgClientBudgetStore::new(database.pool.clone());
|
||||
restarted.admit(key_id("key")).await.unwrap();
|
||||
let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(events, 1, "admission must not leave pending charges");
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0"
|
||||
);
|
||||
restarted
|
||||
.settle(charge("key", "known", "0.4"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0.4"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn transient_settlement_failure_rolls_back_and_retries_exact_cost_before_admission() {
|
||||
let Some(database) = TestDatabase::create("budgets_retry").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "1", "5").await;
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
// 在费用事件插入后让窗口写入失败,验证整个事务回滚。
|
||||
sqlx::raw_sql(
|
||||
"create function reject_budget_update() returns trigger language plpgsql as $$
|
||||
begin
|
||||
if new.daily_used_usd > 0 then raise exception 'temporary test failure'; end if;
|
||||
return new;
|
||||
end $$;
|
||||
create trigger reject_budget_update before update on client_key_budget_windows
|
||||
for each row execute function reject_budget_update();",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(store.settle(charge("key", "retry", "1.25")).await.is_err());
|
||||
let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(events, 0);
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0"
|
||||
);
|
||||
sqlx::query("drop trigger reject_budget_update on client_key_budget_windows")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = store.admit(key_id("key")).await.unwrap_err();
|
||||
assert_eq!(error.client_error_code(), Some("key_daily_budget_exceeded"));
|
||||
store.settle(charge("key", "retry", "1.25")).await.unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"1.25"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() {
|
||||
let Some(database) = TestDatabase::create("budgets_policy").await else {
|
||||
return;
|
||||
};
|
||||
seed(&database, "key", "0", "0").await;
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
let admin = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
store
|
||||
.settle(charge("key", "unlimited", "2.75"))
|
||||
.await
|
||||
.unwrap();
|
||||
let update = UpdateClientKey {
|
||||
id: ClientApiKeyId::new("key").unwrap(),
|
||||
name: "key".to_owned(),
|
||||
label: None,
|
||||
group_ids: vec![],
|
||||
limits: RateLimits {
|
||||
max_concurrency: 3,
|
||||
requests_per_minute: 0,
|
||||
},
|
||||
daily_limit_usd: Some("2".parse().unwrap()),
|
||||
weekly_limit_usd: Some("10".parse().unwrap()),
|
||||
};
|
||||
admin
|
||||
.update_client_key(update.clone(), &context())
|
||||
.await
|
||||
.unwrap();
|
||||
admin
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: None,
|
||||
..update.clone()
|
||||
},
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let current = status(&database, "key").await;
|
||||
assert_eq!(current.limits.daily_usd.canonical(), "2");
|
||||
assert_eq!(current.limits.weekly_usd.canonical(), "10");
|
||||
assert_eq!(current.daily_used_usd.canonical(), "2.75");
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(key_id("key"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_daily_budget_exceeded")
|
||||
);
|
||||
admin
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
daily_limit_usd: Some("0".parse().unwrap()),
|
||||
weekly_limit_usd: None,
|
||||
..update
|
||||
},
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store.admit(key_id("key")).await.unwrap();
|
||||
sqlx::query("update client_api_keys set enabled = false where id = 'key'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.admit(key_id("key")).await.unwrap_err().kind(),
|
||||
GatewayErrorKind::PolicyDenied
|
||||
);
|
||||
sqlx::query("delete from client_api_keys where id = 'key'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.admit(key_id("key")).await.unwrap_err().kind(),
|
||||
GatewayErrorKind::Unauthorized
|
||||
);
|
||||
store.settle(charge("key", "allowed", "1")).await.unwrap();
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn budget_database_outage_fails_closed() {
|
||||
let Some(database) = TestDatabase::create("budgets_outage").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgClientBudgetStore::new(database.pool.clone());
|
||||
database.pool.close().await;
|
||||
assert_eq!(
|
||||
store
|
||||
.admit(key_id("key"))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.client_error_code(),
|
||||
Some("key_budget_unavailable")
|
||||
);
|
||||
assert!(store.settle(charge("key", "offline", "1")).await.is_err());
|
||||
database.close().await;
|
||||
}
|
||||
@@ -24,6 +24,7 @@ use super::TestDatabase;
|
||||
#[test]
|
||||
fn client_key_requires_the_frozen_plaintext_format() {
|
||||
let key = NewClientApiKey {
|
||||
budget: Default::default(),
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
@@ -382,6 +383,7 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
|
||||
fn client_key_debug_redacts_plaintext() {
|
||||
let secret = format!("sk_{}", "s".repeat(43));
|
||||
let key = NewClientApiKey {
|
||||
budget: Default::default(),
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
|
||||
@@ -14,6 +14,7 @@ mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod execution_buffer;
|
||||
@@ -21,6 +22,7 @@ mod health;
|
||||
mod observability;
|
||||
mod ops_events;
|
||||
mod provider_accounts;
|
||||
mod proxies;
|
||||
mod query_budget;
|
||||
mod retention;
|
||||
mod runtime_settings;
|
||||
@@ -230,8 +232,11 @@ 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",
|
||||
"outbound_proxies",
|
||||
"provider_accounts",
|
||||
"runtime_settings",
|
||||
]
|
||||
|
||||
@@ -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(gateway_admin::model::proxies::AccountProxySelection::Url(
|
||||
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(gateway_admin::model::proxies::AccountProxySelection::Direct),
|
||||
..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,8 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
};
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: 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 +1569,8 @@ 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,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![updated],
|
||||
audit: audit("audit_admin_upsert_update", "import", "acct_admin_upsert"),
|
||||
@@ -1526,6 +1605,8 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope,
|
||||
accounts: vec![account("acct_admin_upsert", "another-user")],
|
||||
audit: audit("audit_admin_upsert_rebind", "import", "acct_admin_upsert"),
|
||||
@@ -1563,6 +1644,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 +1658,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 +1691,113 @@ 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;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorization_import_rejects_a_saved_proxy_changed_during_oauth() {
|
||||
use gateway_admin::{
|
||||
model::proxies::{NewProxy, ProxyTestResult, UpdateProxy},
|
||||
ports::proxy::ProxyStore,
|
||||
};
|
||||
let Some(database) = TestDatabase::create("oauth_proxy_guard").await else {
|
||||
return;
|
||||
};
|
||||
let proxies = gateway_store::postgres::PgProxyRepository::new(database.pool.clone());
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "oauth-proxy".to_owned(),
|
||||
};
|
||||
let original = gateway_core::account::OutboundProxy::parse("http://127.0.0.1:8080").unwrap();
|
||||
let saved = proxies
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "OAuth".to_owned(),
|
||||
proxy: original.clone(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
let success = ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms: 1,
|
||||
exit_ip: Some("203.0.113.5".parse().unwrap()),
|
||||
message: "Connected".to_owned(),
|
||||
};
|
||||
proxies
|
||||
.record_test(&saved.id, saved.revision, success.clone(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let replacement = gateway_core::account::OutboundProxy::parse("http://127.0.0.1:9090").unwrap();
|
||||
let edited = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name,
|
||||
proxy: Some(replacement.clone()),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
proxies
|
||||
.record_test(&edited.id, edited.revision, success, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let mut candidate = account("acct_oauth_proxy", "oauth-user");
|
||||
candidate.outbound_proxy = Some(original);
|
||||
let command = |candidate: NewProviderAccount| ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: Some(gateway_admin::model::proxies::ImportProxyBinding {
|
||||
id: saved.id.clone(),
|
||||
proxy: candidate.outbound_proxy.clone().unwrap(),
|
||||
}),
|
||||
scope: ProviderAccountAdminScope {
|
||||
provider_kind: "openai".to_owned(),
|
||||
},
|
||||
accounts: vec![candidate],
|
||||
audit: audit("audit_oauth_proxy", "authorize", "acct_oauth_proxy"),
|
||||
};
|
||||
assert!(
|
||||
repository
|
||||
.import_provider_accounts(command(candidate.clone()))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(account_count(&database.pool, "acct_oauth_proxy").await, 0);
|
||||
candidate.outbound_proxy = Some(replacement);
|
||||
repository
|
||||
.import_provider_accounts(command(candidate))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
proxies
|
||||
.list_accounts(gateway_admin::model::proxies::ProxyAccountListQuery {
|
||||
proxy_id: saved.id.clone(),
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).unwrap(),
|
||||
search: String::new(),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.items[0]
|
||||
.id,
|
||||
"acct_oauth_proxy"
|
||||
);
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("select count(*) from outbound_proxies")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -1614,6 +1809,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 +1991,8 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![
|
||||
account("acct_admin_a", "user-admin-a"),
|
||||
@@ -1822,6 +2020,8 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![
|
||||
account("acct_admin_transient", "user-admin-transient"),
|
||||
@@ -1947,6 +2147,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 +2210,8 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
|
||||
};
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_grouped_import", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_import", "import", "acct_grouped_import"),
|
||||
@@ -2042,6 +2245,8 @@ async fn provider_account_import_and_reauthorization_preserve_existing_membershi
|
||||
.expect("set account scheduling before reimport");
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_reimport_candidate", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_reimport", "import", "acct_grouped_import"),
|
||||
@@ -2109,6 +2314,8 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
|
||||
};
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_rotation_quota", "user-rotation-quota")],
|
||||
audit: audit(
|
||||
@@ -2250,6 +2457,8 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
disabled.credential_state = CredentialState::Expired;
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: None,
|
||||
scope: scope.clone(),
|
||||
accounts: vec![disabled],
|
||||
audit: audit(
|
||||
@@ -2303,8 +2512,9 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
|
||||
pub(super) 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(),
|
||||
@@ -2362,7 +2572,7 @@ fn plaintext_credential(marker: &str) -> PlaintextCredential {
|
||||
)
|
||||
}
|
||||
|
||||
fn audit(id: &str, action: &str, entity_ref: &str) -> AdminAuditEvent {
|
||||
pub(super) fn audit(id: &str, action: &str, entity_ref: &str) -> AdminAuditEvent {
|
||||
AdminAuditEvent {
|
||||
id: id.to_owned(),
|
||||
actor_kind: AdminAuditActorKind::System,
|
||||
@@ -2473,3 +2683,169 @@ 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(gateway_admin::model::proxies::AccountProxySelection::Url(
|
||||
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()),
|
||||
outbound_proxy: None,
|
||||
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 {
|
||||
outbound_proxy: None,
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,707 @@
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext, PageSize,
|
||||
accounts::{BatchUpdateAccounts, UpdateAccount},
|
||||
proxies::*,
|
||||
},
|
||||
ports::{
|
||||
proxy::ProxyStore,
|
||||
store::{AccountStore, AdminStoreErrorKind},
|
||||
},
|
||||
};
|
||||
use gateway_core::account::{AccountWeight, OutboundProxy, ProviderAccountId};
|
||||
use gateway_store::postgres::{
|
||||
PgProviderAccountRepository, PgProxyRepository, ProviderAccountRepository,
|
||||
};
|
||||
|
||||
use super::{TestDatabase, admin_account_store, provider_accounts::account};
|
||||
|
||||
fn context() -> MutationContext {
|
||||
MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "managed-proxy-test".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn success() -> ProxyTestResult {
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms: 10,
|
||||
exit_ip: Some("203.0.113.5".parse().unwrap()),
|
||||
message: "Connected".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn update(account_id: &str, selection: AccountProxySelection) -> UpdateAccount {
|
||||
UpdateAccount {
|
||||
account_id: account_id.to_owned(),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: AccountWeight::DEFAULT,
|
||||
group_ids: vec![],
|
||||
outbound_proxy: Some(selection),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_account_removal_preserves_settings_and_rejects_changed_bindings() {
|
||||
let Some(database) = TestDatabase::create("proxy_account_remove").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let context = context();
|
||||
let saved = store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "解绑测试".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://user:secret@127.0.0.1:17890").unwrap(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
sqlx::query(
|
||||
"insert into provider_accounts
|
||||
(id, provider_kind, name, email, authentication_kind, provider_credentials_json,
|
||||
has_refresh_token, enabled, concurrency_limit, weight, plan_type,
|
||||
credential_observed_at, created_at, updated_at, outbound_proxy_id, outbound_proxy_url)
|
||||
values ('acct_remove', 'openai', '解绑账号', 'remove@example.invalid', 'oauth',
|
||||
'{\"access_token\":\"preserve-test-secret\"}'::jsonb, false, false, 3, 7, 'plus',
|
||||
now(), now(), now(), $1, $2)",
|
||||
)
|
||||
.bind(&saved.id)
|
||||
.bind(saved.proxy.expose_url())
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_groups (id, name, color, enabled, created_at, updated_at)
|
||||
values ('grp_00000000000000000000000000000001', '保留分组', '#123456FF', true, now(), now())",
|
||||
).execute(&database.pool).await.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_group_accounts (account_group_id, provider_account_id, created_at)
|
||||
values ('grp_00000000000000000000000000000001', 'acct_remove', now())",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let account_id = ProviderAccountId::new("acct_remove").unwrap();
|
||||
let snapshot = "select to_jsonb(a) - 'updated_at' - 'outbound_proxy_id' - 'outbound_proxy_url'
|
||||
from provider_accounts a where id = 'acct_remove'";
|
||||
let before: serde_json::Value = sqlx::query_scalar(snapshot)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let audit_count: i64 = sqlx::query_scalar("select count(*) from admin_audit_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.remove_account("proxy_previous", &account_id, &context)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert_eq!(store.get(&saved.id).await.unwrap().account_count, 1);
|
||||
let unchanged: (Option<String>, Option<String>) = sqlx::query_as(
|
||||
"select outbound_proxy_id, outbound_proxy_url from provider_accounts where id = 'acct_remove'",
|
||||
).fetch_one(&database.pool).await.unwrap();
|
||||
assert_eq!(
|
||||
unchanged,
|
||||
(
|
||||
Some(saved.id.clone()),
|
||||
Some(saved.proxy.expose_url().to_owned())
|
||||
)
|
||||
);
|
||||
let current_revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(current_revision, revision);
|
||||
|
||||
let committed = store
|
||||
.remove_account(&saved.id, &account_id, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(committed.get(), u64::try_from(revision + 1).unwrap());
|
||||
let after: serde_json::Value = sqlx::query_scalar(snapshot)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after, before);
|
||||
let direct: (Option<String>, Option<String>) = sqlx::query_as(
|
||||
"select outbound_proxy_id, outbound_proxy_url from provider_accounts where id = 'acct_remove'",
|
||||
).fetch_one(&database.pool).await.unwrap();
|
||||
assert_eq!(direct, (None, None));
|
||||
let group_count: i64 = sqlx::query_scalar(
|
||||
"select count(*) from account_group_accounts where provider_account_id = 'acct_remove'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(group_count, 1);
|
||||
assert_eq!(store.get(&saved.id).await.unwrap().account_count, 0);
|
||||
assert_eq!(
|
||||
store
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: saved.id.clone(),
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: String::new(),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.total,
|
||||
0
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.remove_account(&saved.id, &account_id, &context)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
let current_audits: i64 = sqlx::query_scalar("select count(*) from admin_audit_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(current_audits, audit_count + 1);
|
||||
let current_revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(u64::try_from(current_revision).unwrap(), committed.get());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_accounts_paginate_thousands_of_accounts_and_search_without_loading_the_catalog() {
|
||||
let Some(database) = TestDatabase::create("proxy_account_pagination").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let context = context();
|
||||
let saved = store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "分页测试".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:17890").unwrap(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
let empty = store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "无关联账号".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:17891").unwrap(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
// 相同名称验证稳定排序;账号仅写入隔离测试数据库,不接触真实账号池。
|
||||
sqlx::query(
|
||||
"insert into provider_accounts
|
||||
(id, provider_kind, name, email, authentication_kind, provider_credentials_json,
|
||||
has_refresh_token, enabled, credential_observed_at, created_at, updated_at,
|
||||
outbound_proxy_id, outbound_proxy_url)
|
||||
select 'acct_page_' || lpad(n::text, 4, '0'),
|
||||
case when n % 2 = 0 then 'xai' else 'openai' end, '共享账号',
|
||||
'page_' || lpad(n::text, 4, '0') || '@example.invalid', 'oauth',
|
||||
'{\"access_token\":\"page-test-secret\"}'::jsonb, false, n % 2 = 0,
|
||||
now(), now(), now(), $1, $2
|
||||
from generate_series(1, 1005) as n",
|
||||
)
|
||||
.bind(&saved.id)
|
||||
.bind(saved.proxy.expose_url())
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("update provider_accounts set plan_type = 'plus' where id = 'acct_page_0001'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_groups (id, name, color, enabled, created_at, updated_at)
|
||||
values ('grp_00000000000000000000000000000001', '工作组', '#123456FF', false, now(), now()),
|
||||
('grp_00000000000000000000000000000002', '次页分组', '#654321FF', true, now(), now())",
|
||||
).execute(&database.pool).await.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_group_accounts (account_group_id, provider_account_id, created_at)
|
||||
values ('grp_00000000000000000000000000000001', 'acct_page_0001', now()),
|
||||
('grp_00000000000000000000000000000002', 'acct_page_0021', now())",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
PgProviderAccountRepository::new(database.pool.clone())
|
||||
.insert_provider_account(account("acct_unrelated", "unrelated-user"))
|
||||
.await
|
||||
.unwrap();
|
||||
let query = |page, search: &str| ProxyAccountListQuery {
|
||||
proxy_id: saved.id.clone(),
|
||||
page,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: search.to_owned(),
|
||||
};
|
||||
let first = store.list_accounts(query(1, "")).await.unwrap();
|
||||
assert_eq!((first.total, first.items.len()), (1005, 20));
|
||||
assert_eq!(first.items[0].id, "acct_page_0001");
|
||||
assert_eq!(
|
||||
first.items[0].email.as_deref(),
|
||||
Some("page_0001@example.invalid")
|
||||
);
|
||||
assert_eq!(first.items[0].provider_kind, "openai");
|
||||
assert_eq!(first.items[0].authentication_kind, "oauth");
|
||||
assert_eq!(first.items[0].plan_type.as_deref(), Some("plus"));
|
||||
assert_eq!(first.items[0].groups.len(), 1);
|
||||
assert_eq!(first.items[0].groups[0].name, "工作组");
|
||||
assert_eq!(first.items[0].groups[0].color.as_str(), "#123456FF");
|
||||
assert!(!first.items[0].groups[0].enabled);
|
||||
assert!(first.items[1].groups.is_empty());
|
||||
assert!(!first.items[0].enabled);
|
||||
let second = store.list_accounts(query(2, "")).await.unwrap();
|
||||
assert_eq!(second.items[0].id, "acct_page_0021");
|
||||
assert_eq!(second.items[0].groups[0].name, "次页分组");
|
||||
assert!(
|
||||
first
|
||||
.items
|
||||
.iter()
|
||||
.all(|left| second.items.iter().all(|right| left.id != right.id))
|
||||
);
|
||||
let last = store.list_accounts(query(51, "")).await.unwrap();
|
||||
assert_eq!((last.total, last.items.len()), (1005, 5));
|
||||
assert_eq!(last.items[4].id, "acct_page_1005");
|
||||
let filtered = store.list_accounts(query(1, "PAGE_004")).await.unwrap();
|
||||
assert_eq!((filtered.total, filtered.items.len()), (10, 10));
|
||||
assert_eq!(filtered.items[0].id, "acct_page_0040");
|
||||
assert_eq!(store.list_accounts(query(1, "%")).await.unwrap().total, 0);
|
||||
assert!(
|
||||
store
|
||||
.list_accounts(query(52, ""))
|
||||
.await
|
||||
.unwrap()
|
||||
.items
|
||||
.is_empty()
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: empty.id,
|
||||
..query(1, "")
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.total,
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: "missing".to_owned(),
|
||||
..query(1, "")
|
||||
})
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::NotFound
|
||||
);
|
||||
let catalog = store
|
||||
.list(ProxyListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: "分页测试".to_owned(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(catalog.items[0].account_count, 1005);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are_committed() {
|
||||
use gateway_store::postgres::{
|
||||
ImportProviderAccounts, ProviderAccountAdminRepository, ProviderAccountAdminScope,
|
||||
};
|
||||
let Some(database) = TestDatabase::create("proxy_import_reservation").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let other_process = PgProxyRepository::new(database.pool.clone());
|
||||
let context = context();
|
||||
let saved = store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "导入出口".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:8080").unwrap(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert!(store.reserve_import(&saved.id).await.is_err());
|
||||
store
|
||||
.record_test(&saved.id, saved.revision, success(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let reservation = store.reserve_import(&saved.id).await.unwrap();
|
||||
let replacement = UpdateProxy {
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name,
|
||||
proxy: Some(OutboundProxy::parse("http://127.0.0.1:9090").unwrap()),
|
||||
};
|
||||
assert!(
|
||||
other_process
|
||||
.update(replacement.clone(), &context)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
other_process
|
||||
.delete(&saved.id, saved.revision, &context)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
other_process
|
||||
.record_test(
|
||||
&saved.id,
|
||||
saved.revision,
|
||||
ProxyTestResult {
|
||||
success: false,
|
||||
..success()
|
||||
},
|
||||
&context
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
|
||||
// 模拟上游已轮换的新凭据;持有保护时,另一个连接仍能完成导入事务。
|
||||
let mut candidate = account("acct_reserved_import", "reserved-import-user");
|
||||
candidate.outbound_proxy = Some(reservation.binding.proxy.clone());
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let audit =
|
||||
super::provider_accounts::audit("audit_reserved_import", "import", "acct_reserved_import");
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
repository.import_provider_accounts(ImportProviderAccounts {
|
||||
settings: None,
|
||||
outbound_proxy: Some(reservation.binding.clone()),
|
||||
scope: ProviderAccountAdminScope {
|
||||
provider_kind: "openai".to_owned(),
|
||||
},
|
||||
accounts: vec![candidate],
|
||||
audit,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
drop(reservation);
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
loop {
|
||||
if other_process
|
||||
.update(replacement.clone(), &context)
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
repository
|
||||
.load_provider_account("acct_reserved_import")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_tests() {
|
||||
let Some(database) = TestDatabase::create("managed_proxies").await else {
|
||||
return;
|
||||
};
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
accounts
|
||||
.insert_provider_account(account(id, id))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let admin = admin_account_store(&database.pool);
|
||||
let context = context();
|
||||
let old_proxy = OutboundProxy::parse("http://user:secret@127.0.0.1:8080").unwrap();
|
||||
let created = store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "Office".to_owned(),
|
||||
proxy: old_proxy.clone(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert_eq!(store.get(&created.id).await.unwrap().proxy, old_proxy);
|
||||
assert!(
|
||||
store
|
||||
.create(
|
||||
NewProxy {
|
||||
name: "Duplicate".to_owned(),
|
||||
proxy: old_proxy.clone()
|
||||
},
|
||||
&context
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let selection = AccountProxySelection::Saved(created.id.clone());
|
||||
assert!(
|
||||
admin
|
||||
.update_account(update("acct_one", selection.clone()), &context)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let read = || async {
|
||||
sqlx::query_as::<_, (String, Option<String>, Option<String>, i64)>("select id, outbound_proxy_id, outbound_proxy_url, credential_revision from provider_accounts order by id")
|
||||
.fetch_all(&database.pool).await.unwrap()
|
||||
};
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.is_none() && row.2.is_none() && row.3 == 1)
|
||||
);
|
||||
let tested = store
|
||||
.record_test(&created.id, created.revision, success(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(tested.last_test_at.is_some());
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
admin
|
||||
.update_account(update(id, selection.clone()), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(store.get(&created.id).await.unwrap().account_count, 2);
|
||||
assert_eq!(
|
||||
store
|
||||
.delete(&created.id, created.revision, &context)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.as_deref() == Some(created.id.as_str())
|
||||
&& row.2.as_deref() == Some(old_proxy.expose_url())
|
||||
&& row.3 == 1)
|
||||
);
|
||||
|
||||
let renamed = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
id: created.id.clone(),
|
||||
revision: created.revision,
|
||||
name: "Renamed".to_owned(),
|
||||
proxy: None,
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert_eq!(renamed.proxy, old_proxy);
|
||||
assert_eq!(renamed.last_test, Some(success()));
|
||||
assert!(read().await.iter().all(|row| row.3 == 1));
|
||||
let new_proxy = OutboundProxy::parse("socks5h://next:new-secret@127.0.0.1:1080").unwrap();
|
||||
let edited = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
id: created.id.clone(),
|
||||
revision: renamed.revision,
|
||||
name: renamed.name,
|
||||
proxy: Some(new_proxy.clone()),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert!(edited.last_test.is_none());
|
||||
assert!(edited.last_test_at.is_none());
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.2.as_deref() == Some(new_proxy.expose_url()) && row.3 == 1)
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.record_test(&created.id, created.revision, success(), &context)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(store.get(&created.id).await.unwrap().last_test.is_none());
|
||||
store
|
||||
.record_test(&created.id, edited.revision, success(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
admin
|
||||
.update_account(update(id, AccountProxySelection::Direct), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.is_none() && row.2.is_none() && row.3 == 1)
|
||||
);
|
||||
store
|
||||
.delete(&created.id, edited.revision, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(store.get(&created.id).await.is_err());
|
||||
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 legacy_urls_join_one_catalog_entry_and_invalid_batch_rolls_back() {
|
||||
let Some(database) = TestDatabase::create("proxy_catalog_legacy").await else {
|
||||
return;
|
||||
};
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let proxy = OutboundProxy::parse("http://user:secret@127.0.0.1:8080").unwrap();
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
let mut seed = account(id, id);
|
||||
seed.outbound_proxy = Some(proxy.clone());
|
||||
accounts.insert_provider_account(seed).await.unwrap();
|
||||
}
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let page = store
|
||||
.list(ProxyListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: String::new(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].account_count, 2);
|
||||
let admin = admin_account_store(&database.pool);
|
||||
assert!(
|
||||
admin
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
account_ids: vec!["acct_one".to_owned(), "acct_missing".to_owned()],
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: AccountWeight::DEFAULT,
|
||||
group_ids: vec![],
|
||||
outbound_proxy: Some(AccountProxySelection::Url(
|
||||
OutboundProxy::parse("http://127.0.0.1:9090").unwrap()
|
||||
)),
|
||||
},
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("select count(*) from outbound_proxies")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
sqlx::query_as::<_, (bool, String)>(
|
||||
"select enabled, outbound_proxy_url from provider_accounts where id = 'acct_one'"
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
(true, proxy.expose_url().to_owned())
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn migration_backfills_shared_proxies_without_changing_credentials() {
|
||||
let Some(database) = TestDatabase::create("proxy_backfill").await else {
|
||||
return;
|
||||
};
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
for id in ["acct_one", "acct_two", "acct_direct"] {
|
||||
accounts
|
||||
.insert_provider_account(account(id, id))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
sqlx::raw_sql("alter table provider_accounts drop column outbound_proxy_id; drop table outbound_proxies;
|
||||
update provider_accounts set outbound_proxy_url = 'http://user:secret@127.0.0.1:8080/' where id <> 'acct_direct';")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../migrations/0005_managed_outbound_proxies.sql"
|
||||
))
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("select count(*) from outbound_proxies")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
assert_eq!(sqlx::query_scalar::<_, i64>("select count(*) from provider_accounts a join outbound_proxies p on a.outbound_proxy_id = p.id where a.outbound_proxy_url = p.proxy_url and a.credential_revision = 1").fetch_one(&database.pool).await.unwrap(), 2);
|
||||
assert!(
|
||||
sqlx::query_scalar::<_, Option<String>>(
|
||||
"select outbound_proxy_id from provider_accounts where id = 'acct_direct'"
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
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;
|
||||
}
|
||||
@@ -267,9 +291,10 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
) -> Result<PreparedCredentialImport, ProviderAdminError> {
|
||||
let prepared = self
|
||||
.credentials
|
||||
.prepare_import_document(Value::Object(
|
||||
command.document.into_provider_data().into_inner(),
|
||||
))
|
||||
.prepare_import_document_with_proxy(
|
||||
Value::Object(command.document.into_provider_data().into_inner()),
|
||||
command.default_outbound_proxy.as_ref(),
|
||||
)
|
||||
.await
|
||||
.inspect_err(|error| {
|
||||
log_import_failure("prepare_document", credential_admin_error_code(error));
|
||||
@@ -595,6 +620,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 +724,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 +792,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 +853,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 +1311,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 +1342,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 +1361,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",
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user