Compare commits

...
25 Commits
Author SHA1 Message Date
ZyyandCodex 68f05a4481 chore: bump version to 3.4.0
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 16:08:59 +08:00
ZyyandCodex bb4d178b5d docs(release): prepare v3.4.0 release notes
Document managed proxies, connection tests before saving, credential and
certificate fixes, account and interface improvements, and migration
behavior.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 16:08:26 +08:00
ZyyandCodex 73487fd98b fix(api-keys): trim trailing zeros from budget amounts
Display budget amounts with up to two decimal places without padding
trailing zeros. Keep exact amounts in the details popover and preserve
accounting precision.

Validation: frontend lint/type checks/build and browser checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 16:08:20 +08:00
ZyyandCodex 66fd09eba6 feat(proxies): test connections before saving and refine controls
Separate connection testing from saving in proxy forms. Probe unsaved URLs
with the shared certificate configuration and test concurrency limit,
without changing saved proxy settings or linked accounts.

Keep saved authentication when testing an unchanged connection, use
clearer placeholders, fix action column widths, and remove the redundant
list refresh button.

Validation: 5 proxy API tests, scoped strict Clippy, rustfmt, frontend
lint/type checks/build, and browser checks passed. Codex CLI requests
through real HTTP and SOCKS5 proxies completed with matching upstream
text while Clash TUN remained enabled.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 16:08:14 +08:00
ZyyandCodex bd152dec00 feat: integrate managed proxies and WebSocket fixes (#60)
Integrate managed account proxies and the non-streaming WebSocket fix.
Include shared CA trust, credential import protection, paginated linked
account management, and the reviewed interface refinements.

Preserve the original contribution and the separate proxy and UI commits.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 14:54:24 +08:00
ZyyandCodex 7090dca502 fix(ui): correct modal dragging and sidebar scrolling
Use positional offsets for modal dragging to prevent seams between table
cells. Scroll sidebar navigation independently while keeping the brand
and footer fixed, with an auto-hiding scrollbar at the outer edge.

Validation: frontend lint/type checks/build and browser checks for
modal dragging, sidebar overflow, alignment, and auto-hide passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 14:43:45 +08:00
ZyyandCodex 029f3124bf feat(proxies): paginate linked accounts and support removal
Load linked accounts on demand with shared plan metadata and guarded
removal to direct connections. Preserve credentials, scheduling settings,
and groups when clearing a proxy binding.

Align proxy forms and tables with shared components, integrate endpoint
details into outbound selectors, and consolidate probe results into the
latency column.

Validation: 11 proxy regressions, strict workspace Clippy, rustfmt,
frontend lint/type checks/build, and local browser checks passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-11 14:43:31 +08:00
ZyyandCodex 840707e5a1 fix(proxies): share CA trust and preserve credentials during import
Co-authored-by: Codex <noreply@openai.com>
2026-09-11 13:07:31 +08:00
Yanshan Chen 704b652f86 feat: add managed account proxies and fix non-streaming WebSocket requests 2026-09-11 11:55:09 +08:00
ZyyandCodex 2b98839c29 fix(admin): unify plan labels and refine budget popover
Normalize provider plan display names to PascalCase while preserving raw
plan values. Replace native budget tooltips with the shared borderless
popover, retaining exact amounts and daily/weekly reset information.
Reuse the dayjs date helper for explicit Beijing time formatting.

Validation: 371 Admin/API tests, strict Clippy, rustfmt, frontend lint and
type checks passed. Local frontend and backend health checks passed.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Co-authored-by: Codex <noreply@openai.com>
2026-09-10 15:33:43 +08:00
223 changed files with 12885 additions and 2248 deletions
+31 -9
View File
@@ -88,6 +88,13 @@ jobs:
fail("Release tag #{release_tag} does not match #{version_path} tag #{expected_tag}") unless release_tag == expected_tag
fail("#{version_path}: version #{version} is not a valid semver") unless version.match?(SEMVER)
notes_path = "release/notes.md"
fail("#{notes_path}: release notes are required") unless File.file?(notes_path)
notes = File.readlines(notes_path)
unless notes.first&.strip == "# #{release_tag}" && !notes.drop(1).join.strip.empty?
fail("#{notes_path}: expected non-empty notes starting with '# #{release_tag}'")
end
platforms = platforms_data["platforms"]
fail("#{platforms_path}: platforms must be a non-empty array") unless platforms.is_a?(Array) && !platforms.empty?
@@ -643,15 +650,10 @@ jobs:
cosign sign --yes "${images[@]}"
- name: Publish GitHub Release
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
with:
tag_name: ${{ needs.release-info.outputs.tag }}
generate_release_notes: true
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
body: |
- name: Prepare release notes
shell: bash
env:
INSTALLATION_NOTES: |
> [!WARNING]
> **Upgrading across major versions?** Online upgrades across major versions are not supported. Deploy the new major fresh with a new `.runtime/` data directory, then re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
@@ -672,6 +674,26 @@ jobs:
- [GitHub Repository](https://github.com/${{ github.repository }})
- [Installation Guide](https://github.com/${{ github.repository }}/blob/main/deploy/README.md)
run: |
python3 <<'PY'
import os
from pathlib import Path
notes = "\n".join(Path("release/notes.md").read_text().splitlines()[1:]).strip()
installation = os.environ["INSTALLATION_NOTES"].strip()
output = Path(os.environ["RUNNER_TEMP"]) / "release-notes.md"
output.write_text(f"{notes}\n\n{installation}\n")
PY
- name: Publish GitHub Release
uses: softprops/action-gh-release@efb35369e0ad2afab669f228072c1b0d510eae64 # v3.0.3
with:
tag_name: ${{ needs.release-info.outputs.tag }}
generate_release_notes: true
previous_tag: ${{ needs.release-info.outputs.previous_tag }}
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
body_path: ${{ runner.temp }}/release-notes.md
fail_on_unmatched_files: true
files: |
dist/release/*
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+12 -3
View File
@@ -10,8 +10,9 @@
[![CI](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/releases)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg?style=flat-square)](LICENSE)
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档)
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档) · [社区](#社区) · [许可证](#许可证)
</div>
@@ -20,8 +21,8 @@
## 功能
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看。
- **访问控制**:通过账号分组分配密钥权限,设置并发和请求频率。
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看;代理管理支持独立保存、连接测试和账号绑定,以及 sub2api OpenAI 账号代理导入。
- **访问控制**:通过账号分组分配密钥权限,设置日/周 USD 限额、并发和请求频率。
- **客户端接入**:支持 Codex CLI、桌面端与 Responses API 客户端,提供 HTTP 流式响应和 WebSocket。
- **请求诊断**:查看请求记录、用量、费用、延迟和错误详情。
- **部署运维**:Docker Compose 部署、S3/R2 数据库备份、同大版本在线更新与回滚。
@@ -122,3 +123,11 @@ curl http://127.0.0.1:8080/v1/models \
- [系统架构](docs/architecture.md)
- [管理端主题](docs/theme.md)
- [数据库迁移](backend/migrations/README.md)
## 社区
感谢 [LINUX DO](https://linux.do) 社区提供开放、友善的技术交流平台。
## 许可证
本项目基于 [Apache License 2.0](LICENSE) 开源。
+11 -7
View File
@@ -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
View File
@@ -15,7 +15,7 @@ resolver = "3"
[workspace.package]
version = "0.0.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
@@ -38,3 +38,7 @@ strip = "symbols"
lto = "thin"
codegen-units = 1
panic = "abort"
[patch.crates-io]
# Upstream SOCKS IPv6 fix; remove once a release includes hyperium/hyper-util#302.
hyper-util = { git = "https://github.com/hyperium/hyper-util", rev = "191b1acf0a6b18e807e4abf5a6d4a7c47979ecd6" }
+1
View File
@@ -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
+4 -1
View File
@@ -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",
+145
View File
@@ -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");
+1
View File
@@ -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"] }
+17 -2
View File
@@ -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(&quota),
),
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(&quota),
),
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(),
})
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-api"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
@@ -18,16 +18,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)),
))
}
+16 -2
View File
@@ -143,13 +143,27 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
/// Gateway 错误的 OpenAI HTTP 表达。
pub fn gateway_error_response(error: &GatewayError) -> Response {
let (status, default_type, default_code) = gateway_error_contract(error.kind());
openai_error_response(
let mut response = openai_error_response(
status,
error.client_message(),
error.client_error_type().unwrap_or(default_type),
error.client_error_code().unwrap_or(default_code),
)
.into_response()
.into_response();
if let Some(delay) = error.retry_after()
&& let Ok(value) = HeaderValue::from_str(
&delay
.as_secs()
.saturating_add(u64::from(delay.subsec_nanos() > 0))
.max(1)
.to_string(),
)
{
response
.headers_mut()
.insert(axum::http::header::RETRY_AFTER, value);
}
response
}
/// 在下游尚未提交时优先交付 Provider 的原始 HTTP 失败响应。
@@ -756,3 +756,60 @@ mod actions {
);
}
}
mod import_settings {
use gateway_api::admin::accounts::{AccountImportRequest, CompleteAccountAuthorizationRequest};
use serde_json::json;
#[test]
fn import_and_oauth_apply_the_same_settings_validation() {
for (field, value) in [
("concurrencyLimit", json!(0)),
("concurrencyLimit", json!(4_294_967_296_u64)),
("weight", json!(0)),
("weight", json!(101)),
("groupIds", json!(["invalid-group"])),
] {
let mut settings =
json!({"enabled": false, "concurrencyLimit": null, "weight": 1, "groupIds": []});
settings[field] = value;
let import: AccountImportRequest = serde_json::from_value(
json!({"provider": "openai", "data": {}, "settings": settings}),
)
.expect("import request");
let oauth: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({"provider": "xai", "flowId": "flow-test", "callbackUrl": "code", "settings": settings})).expect("OAuth request");
assert_eq!(
import.validate().expect_err("invalid settings").field(),
field
);
assert_eq!(
oauth.validate().expect_err("invalid settings").field(),
field
);
}
}
#[test]
fn import_settings_require_a_complete_explicit_configuration() {
let settings =
json!({"enabled": false, "concurrencyLimit": null, "weight": 100, "groupIds": []});
let request: AccountImportRequest =
serde_json::from_value(json!({"provider": "openai", "data": {}, "settings": settings}))
.expect("settings");
assert!(request.validate().is_ok());
for field in ["enabled", "concurrencyLimit", "weight", "groupIds"] {
let mut incomplete = settings.clone();
incomplete
.as_object_mut()
.expect("settings object")
.remove(field);
assert!(
serde_json::from_value::<AccountImportRequest>(
json!({"provider": "openai", "data": {}, "settings": incomplete})
)
.is_err(),
"missing {field}"
);
}
}
}
@@ -10,6 +10,46 @@ use serde_json::json;
use super::{AdminTestFixture, AdminTestState};
#[test]
fn budget_inputs_preserve_decimal_precision_and_omitted_updates() {
let payload = json!({"name": "budget", "groupIds": [], "maxConcurrency": 3,
"requestsPerMinute": 0, "dailyLimitUsd": "0.1234567891", "weeklyLimitUsd": "15"});
let command = serde_json::from_value::<CreateClientKeyRequest>(payload.clone())
.unwrap()
.into_command()
.unwrap();
assert_eq!(command.budget.daily_usd.canonical(), "0.1234567891");
assert_eq!(command.budget.weekly_usd.canonical(), "15");
assert_eq!(command.limits.max_concurrency, 3);
for field in ["dailyLimitUsd", "weeklyLimitUsd"] {
for invalid in ["-1", "NaN", "1e3", "10000000000", "0.00000000001", ""] {
let mut payload = payload.clone();
payload[field] = json!(invalid);
assert_eq!(
serde_json::from_value::<CreateClientKeyRequest>(payload)
.unwrap()
.into_command()
.unwrap_err()
.field(),
field
);
}
}
let mut update = payload;
update["id"] = json!("key_budget");
update.as_object_mut().unwrap().remove("dailyLimitUsd");
update["weeklyLimitUsd"] = json!("0");
let command = serde_json::from_value::<UpdateClientKeyRequest>(update)
.unwrap()
.into_command()
.unwrap();
assert_eq!(command.daily_limit_usd, None);
assert_eq!(
command.weekly_limit_usd,
Some(gateway_core::metering::Decimal::ZERO)
);
}
#[test]
fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
let unknown = serde_json::from_value::<ListClientKeysQuery>(json!({ "other": true }));
@@ -207,6 +247,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
.single()
.expect("valid time");
let view = ClientKeyView::from(gateway_admin::model::client_keys::ClientKeyRecord {
budget: Default::default(),
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
name: "visible".to_owned(),
label: None,
@@ -238,6 +279,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
assert!(list["items"][0].get("tokensPerMinute").is_none());
assert!(list["items"][0].get("policyJson").is_none());
assert_eq!(list["items"][0]["dailyLimitUsd"], "0");
assert_eq!(list["items"][0]["weeklyLimitUsd"], "0");
assert_eq!(list["items"][0]["dailyUsedUsd"], "0");
assert!(list["items"][0].get("unresolvedRequests").is_none());
assert_eq!(
DateTime::parse_from_rfc3339(
list["items"][0]["lastUsedAt"]
@@ -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 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-core"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
@@ -24,6 +24,7 @@ semver = "1.0.27"
thiserror = "2.0.18"
tracing = "0.1.44"
uuid = { version = "1.23.3", features = ["v7"] }
url = "2.5.8"
[lints]
workspace = true
@@ -2,6 +2,8 @@
mod error;
mod model;
mod proxy;
pub use proxy::{InvalidOutboundProxy, OutboundProxy};
pub mod scope;
mod selection;
mod store;
@@ -654,7 +654,7 @@ const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
}
}
/// 不含 secret 的账号持久事实。
/// 账号持久事实;代理认证信息只通过显式 secret accessor 读取。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
id: ProviderAccountId,
@@ -676,6 +676,7 @@ pub struct ProviderAccount {
access_token_expires_at: Option<SystemTime>,
next_refresh_at: Option<SystemTime>,
has_refresh_token: bool,
outbound_proxy: Option<super::OutboundProxy>,
}
impl ProviderAccount {
@@ -710,6 +711,7 @@ impl ProviderAccount {
access_token_expires_at,
next_refresh_at: None,
has_refresh_token: false,
outbound_proxy: None,
}
}
@@ -726,6 +728,17 @@ impl ProviderAccount {
self
}
#[must_use]
pub fn with_outbound_proxy(mut self, proxy: Option<super::OutboundProxy>) -> Self {
self.outbound_proxy = proxy;
self
}
#[must_use]
pub const fn outbound_proxy(&self) -> Option<&super::OutboundProxy> {
self.outbound_proxy.as_ref()
}
#[must_use]
pub fn with_account_facts(
mut self,
@@ -0,0 +1,49 @@
//! Explicit account egress. Credentials never appear in Debug or ordinary admin projections.
use std::fmt;
use url::Url;
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct OutboundProxy(Url);
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("invalid outbound proxy; expected http, https, socks5 or socks5h URL with a host and port")]
pub struct InvalidOutboundProxy;
impl OutboundProxy {
pub fn parse(value: &str) -> Result<Self, InvalidOutboundProxy> {
if value.len() > 4096 || value.chars().any(char::is_control) {
return Err(InvalidOutboundProxy);
}
let url = Url::parse(value).map_err(|_| InvalidOutboundProxy)?;
if !matches!(url.scheme(), "http" | "https" | "socks5" | "socks5h")
|| url.host_str().is_none()
|| url.port_or_known_default().is_none_or(|port| port == 0)
|| !matches!(url.path(), "" | "/")
|| url.query().is_some()
|| url.fragment().is_some()
{
return Err(InvalidOutboundProxy);
}
Ok(Self(url))
}
#[must_use]
pub fn expose_url(&self) -> &str {
self.0.as_str()
}
#[must_use]
pub fn endpoint(&self) -> String {
let mut url = self.0.clone();
let _ = url.set_username("");
let _ = url.set_password(None);
url.to_string()
}
}
impl fmt::Debug for OutboundProxy {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("OutboundProxy(<redacted>)")
}
}
@@ -0,0 +1,52 @@
//! 按 Key 持久化 USD 限额与已取得费用,不依赖尽力写入的请求观测。
use std::time::SystemTime;
use futures::future::BoxFuture;
use crate::{error::GatewayError, metering::Decimal, policy::ClientApiKeyId};
use super::ModelRequestId;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ClientBudgetLimits {
pub daily_usd: Decimal,
pub weekly_usd: Decimal,
}
impl ClientBudgetLimits {
#[must_use]
pub fn is_limited(self) -> bool {
self.daily_usd != Decimal::ZERO || self.weekly_usd != Decimal::ZERO
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ClientBudgetStatus {
pub limits: ClientBudgetLimits,
pub daily_used_usd: Decimal,
pub weekly_used_usd: Decimal,
pub daily_resets_at: Option<SystemTime>,
pub weekly_resets_at: Option<SystemTime>,
}
#[derive(Debug, Clone)]
pub struct ClientBudgetCharge {
pub key_id: ClientApiKeyId,
pub request_id: ModelRequestId,
/// 已取得的 USD 费用,包含重试;缺少费用的尝试按零累计。
pub amount_usd: Decimal,
pub completed_at: SystemTime,
}
#[derive(Debug, Clone, Copy, thiserror::Error)]
#[error("client budget store is unavailable")]
pub struct ClientBudgetError;
pub trait ClientBudgetPort: Send + Sync {
/// 原子检查当前限额与已用金额,不创建预扣费或待结算记录。
fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>>;
/// 按网关请求 ID 幂等累计已取得费用;写入失败由 Store 重试。
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>;
}
@@ -30,6 +30,7 @@ use crate::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
};
use crate::lifecycle::CancellationToken;
use crate::metering::Decimal;
use crate::operation::{Operation, ProviderSessionState};
use crate::routing::RoutingPlan;
use futures::future::Fuse;
@@ -166,6 +167,8 @@ where
request_id,
client_api_key_ref,
observation: ResponseObservation::new(timing_started_at),
budget_prior_attempts_usd: Decimal::ZERO,
budget_attempt_already_counted: false,
trace,
deadline,
deadline_timer: Delay::new(
@@ -197,6 +200,7 @@ where
delivery_pending: false,
upstream_complete: false,
finalized: false,
finalized_at: None,
image_generation_requested,
last_retryable_failure: None,
last_retryable_failure_events: Vec::new(),
@@ -253,6 +257,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
request_id: ModelRequestId,
client_api_key_ref: crate::policy::ClientApiKeyId,
observation: ResponseObservation,
budget_prior_attempts_usd: Decimal,
budget_attempt_already_counted: bool,
trace: TraceContext,
deadline: SystemTime,
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
@@ -288,6 +294,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
delivery_pending: bool,
upstream_complete: bool,
finalized: bool,
finalized_at: Option<SystemTime>,
image_generation_requested: bool,
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
last_retryable_failure: Option<ProviderError>,
@@ -475,6 +482,35 @@ where
self.finalized
}
#[must_use]
pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge {
// 超出数据库可表示范围时保留最大金额,避免溢出后误记为零。
let amount_usd = self
.budget_prior_attempts_usd
.checked_add(self.budget_attempt_usd())
.unwrap_or(Decimal::MAX);
super::budget::ClientBudgetCharge {
key_id: self.client_api_key_ref.clone(),
request_id: self.request_id.clone(),
amount_usd,
completed_at: self.finalized_at.unwrap_or_else(SystemTime::now),
}
}
fn budget_attempt_usd(&self) -> Decimal {
if self.budget_attempt_already_counted {
return Decimal::ZERO;
}
// Key 只累计已取得的 USD 费用;缺少费用的请求按零结算。
// 发送状态仍用于判断重放是否安全并保留在诊断中,不能据此推断存在欠费。
self.observation
.cost
.total()
.filter(|money| money.currency().as_str() == "USD")
.map(crate::metering::Money::amount)
.unwrap_or(Decimal::ZERO)
}
/// 返回截至当前已完成的实际上游调用结果。
///
/// 调用方可保存已消费下标;该切片在会话生命周期内只追加、不改写。
@@ -800,6 +836,7 @@ where
let mut events = std::mem::take(&mut self.last_retryable_failure_events);
if !events.is_empty() {
self.observe_atomic_terminal_events(&mut events);
self.budget_attempt_already_counted = true;
return Ok(Some(PullOutcome::TerminalFailure {
events,
error: last_failure,
@@ -1143,6 +1180,11 @@ where
"sendState": format!("{attempt_send_state:?}"),
}));
if retryable {
for event in &atomic_client_events {
for fact in event.canonical_facts() {
self.observation.observe_event(fact);
}
}
// 普通 clone 只保留稳定事实;原始 wire/HTTP response 由 request-local
// 所有权保留到下一次 attempt 成功,或最终空选路时返回客户端。
let persistence_error = self.request_persisted.then(|| error.clone());
@@ -1317,6 +1359,12 @@ where
}
fn reset_uncommitted_observations(&mut self) {
// 响应观测按尝试隔离;被丢弃尝试中已取得的费用仍计入 Key。
self.budget_prior_attempts_usd = self
.budget_prior_attempts_usd
.checked_add(self.budget_attempt_usd())
.unwrap_or(Decimal::MAX);
self.budget_attempt_already_counted = false;
self.observation.reset_for_attempt();
self.upstream_complete = false;
}
@@ -1326,6 +1374,7 @@ where
return Ok(());
}
let completed_at = SystemTime::now();
self.finalized_at = Some(completed_at);
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
@@ -1505,11 +1554,12 @@ where
"sendState": format!("{:?}", finalization.send_state), "attempts": self.attempts,
}),
);
let completed_at = SystemTime::now();
self.finalized_at = Some(completed_at);
if !self.request_persisted {
self.finalized = true;
return Ok(());
}
let completed_at = SystemTime::now();
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
@@ -14,6 +14,7 @@ use uuid::Uuid;
use crate::engine::admission::{
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
};
use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort};
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
NativeContinuationStoreErrorKind, PreviousResponseId,
@@ -256,6 +257,7 @@ pub struct DefaultExecutionService {
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
budget: Option<Arc<dyn ClientBudgetPort>>,
}
impl DefaultExecutionService {
@@ -283,9 +285,16 @@ impl DefaultExecutionService {
circuits,
continuation,
client_api_key_usage,
budget: None,
}
}
#[must_use]
pub fn with_budget(mut self, budget: Arc<dyn ClientBudgetPort>) -> Self {
self.budget = Some(budget);
self
}
async fn start_inner(&self, request: StartExecution) -> Result<StartedExecution, GatewayError> {
let StartExecution {
client,
@@ -325,9 +334,18 @@ impl DefaultExecutionService {
&self,
mut request: PendingStartExecution,
) -> Result<StartedExecution, GatewayError> {
request.client.policy.authorize().map_err(|_| {
GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled")
})?;
// 长连接每次执行都重新鉴权并冻结当前策略,确保限额和授权变更对新请求生效。
request.client = self
.authenticate(request.client.policy.plaintext_key().expose_for_auth())
.map_err(|error| match error {
ClientAuthenticationError::InvalidKey => {
GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid")
}
ClientAuthenticationError::SnapshotUnavailable => GatewayError::new(
GatewayErrorKind::Internal,
"runtime snapshot is unavailable",
),
})?;
let started_at = SystemTime::now();
let deadline_at = started_at
.checked_add(MODEL_REQUEST_DEADLINE)
@@ -504,6 +522,12 @@ impl DefaultExecutionService {
client_api_key_id: client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
if let Some(budget) = &self.budget
&& let Err(error) = budget.admit(client.policy.key_id().clone()).await
{
admission.release().await;
return Err(error);
}
let observation = plan
.candidates()
.first()
@@ -552,6 +576,18 @@ impl DefaultExecutionService {
{
Ok(core) => core,
Err(error) => {
if let Some(budget) = &self.budget {
settle_budget(
budget.as_ref(),
ClientBudgetCharge {
key_id: client.policy.key_id().clone(),
request_id: request_id.clone(),
amount_usd: crate::metering::Decimal::ZERO,
completed_at: SystemTime::now(),
},
)
.await;
}
admission.release().await;
return Err(gateway_error_from_engine(&error));
}
@@ -565,6 +601,7 @@ impl DefaultExecutionService {
admission,
Arc::clone(&self.circuits),
Arc::clone(&self.continuation),
self.budget.clone(),
)),
})
}
@@ -922,6 +959,12 @@ struct AdmissionLease {
model_request_id: ModelRequestId,
}
async fn settle_budget(port: &dyn ClientBudgetPort, charge: ClientBudgetCharge) {
if let Err(error) = port.settle(charge).await {
tracing::error!(%error, "Client budget settlement failed; storage will retry on the next request");
}
}
impl AdmissionLease {
async fn release(self) {
if let Err(error) = self
@@ -941,6 +984,7 @@ struct DefaultExecutionSession {
continuation: Arc<dyn NativeContinuationPort>,
observed_provider_outcomes: usize,
continuation_recorded: bool,
budget: Option<Arc<dyn ClientBudgetPort>>,
}
impl DefaultExecutionSession {
@@ -949,6 +993,7 @@ impl DefaultExecutionSession {
admission: AdmissionLease,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
budget: Option<Arc<dyn ClientBudgetPort>>,
) -> Self {
Self {
core: Some(core),
@@ -957,6 +1002,7 @@ impl DefaultExecutionSession {
continuation,
observed_provider_outcomes: 0,
continuation_recorded: false,
budget,
}
}
@@ -973,6 +1019,11 @@ impl DefaultExecutionSession {
.is_some_and(ResponseExecutionSession::is_finalized)
&& let Some(admission) = self.admission.take()
{
if let Some(budget) = self.budget.take()
&& let Some(core) = &self.core
{
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
admission.release().await;
}
}
@@ -1023,6 +1074,9 @@ impl DefaultExecutionSession {
.get(self.observed_provider_outcomes..)
.unwrap_or_default();
publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await;
if let Some(budget) = self.budget.take() {
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
if let Some(admission) = self.admission.take() {
admission.release().await;
}
@@ -1213,6 +1267,17 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError {
match error {
crate::validation::RoutingError::ModelNotFound {
model,
mapped_model,
} => GatewayError::new(
GatewayErrorKind::ModelNotFound,
if model == mapped_model {
"the requested model was not found in the provider catalogs available to this API key; check the model name"
} else {
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping"
},
),
crate::validation::RoutingError::NoCapableProvider { .. }
| crate::validation::RoutingError::NoCapableProviderEndpoint { .. }
| crate::validation::RoutingError::EmptyAccountScope => GatewayError::new(
@@ -1,6 +1,7 @@
//! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。
pub mod admission;
pub mod budget;
pub mod continuation;
pub mod coordinator;
pub mod execution;
+36
View File
@@ -975,6 +975,13 @@ pub struct GatewayError {
message: &'static str,
diagnostic: Option<Box<ProviderDiagnostic>>,
client_visible_upstream_error: Option<ClientVisibleUpstreamError>,
client_details: Option<Box<GatewayClientErrorDetails>>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct GatewayClientErrorDetails {
code: Option<&'static str>,
retry_after: Option<Duration>,
}
impl GatewayError {
@@ -986,9 +993,33 @@ impl GatewayError {
message,
diagnostic: None,
client_visible_upstream_error: None,
client_details: None,
}
}
#[must_use]
pub fn with_client_code(mut self, code: &'static str) -> Self {
self.client_details
.get_or_insert_with(Default::default)
.code = Some(code);
self
}
#[must_use]
pub fn with_retry_after(mut self, delay: Duration) -> Self {
self.client_details
.get_or_insert_with(Default::default)
.retry_after = Some(delay);
self
}
#[must_use]
pub fn retry_after(&self) -> Option<Duration> {
self.client_details
.as_ref()
.and_then(|details| details.retry_after)
}
/// 将 Provider 错误归一为客户端无关错误。
#[must_use]
pub fn from_provider(error: &ProviderError) -> Self {
@@ -1103,6 +1134,11 @@ impl GatewayError {
self.client_visible_upstream_error
.as_ref()
.and_then(ClientVisibleUpstreamError::code)
.or_else(|| {
self.client_details
.as_ref()
.and_then(|details| details.code)
})
}
}
+14 -2
View File
@@ -52,6 +52,7 @@ pub struct CoreStorePorts {
snapshots: Arc<dyn SnapshotStorePort>,
snapshot_subscriptions: Arc<dyn SnapshotSubscriptionPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
budget: Option<Arc<dyn engine::budget::ClientBudgetPort>>,
}
impl CoreStorePorts {
@@ -79,8 +80,15 @@ impl CoreStorePorts {
snapshots,
snapshot_subscriptions,
client_api_key_usage,
budget: None,
}
}
#[must_use]
pub fn with_budget(mut self, budget: Arc<dyn engine::budget::ClientBudgetPort>) -> Self {
self.budget = Some(budget);
self
}
}
pub struct CoreBundle {
@@ -145,7 +153,7 @@ pub async fn initialize(
Arc::clone(&ports.snapshot_subscriptions),
));
let worker_contributions = publisher.worker_contributions()?;
let service = Arc::new(DefaultExecutionService::new(
let mut service = DefaultExecutionService::new(
snapshots.clone(),
ports.execution,
providers,
@@ -153,7 +161,11 @@ pub async fn initialize(
ports.circuits,
ports.continuation,
ports.client_api_key_usage,
));
);
if let Some(budget) = ports.budget {
service = service.with_budget(budget);
}
let service = Arc::new(service);
let execution: Arc<dyn ExecutionService> = service.clone();
let account_probe: Arc<dyn AccountProbe> = service;
let snapshot_control: Arc<dyn SnapshotControl> = publisher;
@@ -14,6 +14,8 @@ pub struct Decimal(u128);
impl Decimal {
pub const ZERO: Self = Self(0);
/// 数据库可表示的最大非负金额。
pub const MAX: Self = Self(MAX_SCALED_DECIMAL);
/// 从按十位小数缩放的整数创建。
///
@@ -778,6 +778,19 @@ impl RuntimeSnapshot {
}
if candidates.is_empty() {
// 模型存在性必须包含被熔断的 Provider;目录未知时不能断言模型不存在。
let mut scoped_providers = providers.intersection(&self.providers).peekable();
if scoped_providers.peek().is_some()
&& scoped_providers.all(|provider| {
self.known_provider_catalogs.contains(provider)
&& !self.contains_public_model_for_provider(public_model, provider)
})
{
return Err(RoutingError::ModelNotFound {
model: public_model.as_str().to_owned(),
mapped_model: self.mapped_model(public_model.as_str()),
});
}
return Err(RoutingError::NoCapableProvider {
model: public_model.as_str().to_owned(),
});
@@ -85,6 +85,16 @@ pub enum RoutingError {
/// 实体 ID。
id: String,
},
/// 账号范围内的 Provider 目录均已知,且均没有请求或映射后的模型。
#[error(
"model `{model}` (mapped to `{mapped_model}`) was not found in the scoped provider catalogs"
)]
ModelNotFound {
/// 客户端提交的模型名称。
model: String,
/// 应用模型映射后的名称。
mapped_model: String,
},
/// 固定平台内没有可执行本次请求的 Provider。
#[error("no provider can execute model `{model}`")]
NoCapableProvider {
@@ -1,3 +1,4 @@
mod proxy;
mod selection;
mod selection_diagnostics;
@@ -0,0 +1,33 @@
use gateway_core::account::OutboundProxy;
#[test]
fn proxy_endpoints_support_explicit_schemes_and_redact_credentials() {
for scheme in ["http", "https", "socks5", "socks5h"] {
let proxy =
OutboundProxy::parse(&format!("{scheme}://user:p%40ss%3Aword@[::1]:1080")).unwrap();
assert!(proxy.expose_url().contains("p%40ss%3Aword"));
assert!(!proxy.endpoint().contains("user"));
assert!(!proxy.endpoint().contains("p%40ss"));
assert!(!format!("{proxy:?}").contains("user"));
assert!(proxy.endpoint().contains("[::1]:1080"));
}
}
#[test]
fn invalid_proxy_never_silently_becomes_direct() {
for url in [
"",
"localhost:1080",
"file:///etc/passwd",
"ftp://host:21",
"http://host:0",
"socks5://host",
"http://host:80/path",
"http://host:80?secret=1",
"http://host:80#fragment",
"http://user:secret@host:80\n",
] {
let error = OutboundProxy::parse(url).unwrap_err();
assert!(!error.to_string().contains("secret"));
}
}
@@ -2158,6 +2158,146 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() {
let state = store.state.lock().expect("store lock");
assert_eq!(state.finalizations[0].cost_source, CostSource::Unavailable);
assert_eq!(state.finalizations[0].cost_ticks, None);
assert_eq!(
session.budget_charge().amount_usd.scaled(),
999,
"missing final usage cannot erase a charged retry"
);
}
#[test]
fn key_budget_accumulates_only_observed_retry_costs() {
for (first_cost, first_send, expected) in [
(Some(999), UpstreamSendState::Sent, 1122),
(
Some(gateway_core::metering::Decimal::MAX.scaled()),
UpstreamSendState::Sent,
gateway_core::metering::Decimal::MAX.scaled(),
),
(None, UpstreamSendState::Sent, 123),
(None, UpstreamSendState::NotSent, 123),
] {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut first = Vec::new();
if first_send != UpstreamSendState::NotSent {
first.push(Ok(GatewayEvent::Started(ResponseMeta::new(
"resp_first",
"gpt-5",
))));
}
if let Some(cost) = first_cost {
first.push(Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
)));
}
first.push(Err(ProviderError::new(
ProviderErrorKind::Unavailable,
first_send,
)
.with_replay_safe()));
let (coordinator, _, _) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: first,
},
Script::Stream {
account_id: "acct_second",
items: vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"resp_done",
"gpt-5",
))),
Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(123).unwrap(),
)),
Ok(GatewayEvent::Completed(ResponseMeta::new(
"resp_done",
"gpt-5",
))),
],
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let charge = session.budget_charge();
assert_eq!(charge.amount_usd.scaled(), expected);
assert_eq!(
session.budget_charge().completed_at,
charge.completed_at,
"final settlement time stays stable"
);
}
}
#[test]
fn key_budget_settles_post_send_failures_from_observed_cost_without_replaying() {
for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] {
for cost in [None, Some(123)] {
for committed in [false, true] {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut items = vec![Ok(GatewayEvent::Started(ResponseMeta::new(
"response-interrupted",
"gpt-5",
)))];
if let Some(cost) = cost {
items.push(Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(cost).unwrap(),
)));
}
items.push(Err(ProviderError::new(
ProviderErrorKind::Transport,
send_state,
)));
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items,
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
if committed {
block_on(session.next_event()).unwrap().unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
loop {
match block_on(session.next_event()) {
Ok(Some(_)) => {}
Err(EngineError::Provider(_)) => break,
other => panic!("expected terminal transport error, got {other:?}"),
}
}
} else {
assert!(matches!(
block_on(session.collect_uncommitted()),
Err(EngineError::Provider(_))
));
}
assert!(session.is_finalized());
assert_eq!(
session.budget_charge().amount_usd.scaled(),
cost.unwrap_or(0),
"send={send_state:?}, committed={committed}, cost={cost:?}"
);
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
}
}
}
}
#[test]
@@ -2688,10 +2828,16 @@ fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once()
fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut failure = bare_atomic_response_failed("response-only-failed");
let mut events = failure.take_atomic_client_events();
events.push(ProviderEvent::canonical(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(500).unwrap(),
)));
let failure = failure.with_atomic_client_events(events);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_only",
items: vec![Err(bare_atomic_response_failed("response-only-failed"))],
items: vec![Err(failure)],
},
Script::Error(ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
@@ -2731,6 +2877,11 @@ fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire()
));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
assert_eq!(
session.budget_charge().amount_usd.scaled(),
500,
"restoring a discarded failure must not charge its usage twice"
);
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(state.intermediate_failures, 1);
@@ -2868,6 +3019,7 @@ fn ambiguous_send_state_stops_retry() {
let finalization = &state.finalizations[0];
assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
assert_eq!(finalization.send_state, UpstreamSendState::Ambiguous);
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
assert_eq!(finalization.attempt_count, 1);
assert!(!finalization.committed);
assert!(finalization.latency_ms.is_some());
@@ -3342,6 +3494,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
assert_eq!(state.commits, 0);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
assert!(!state.finalizations[0].committed);
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
}
#[test]
@@ -7,6 +7,231 @@ use std::{
time::{Duration, Instant, SystemTime},
};
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
use gateway_core::error::GatewayError;
#[derive(Default)]
struct Admissions {
active: Arc<AtomicBool>,
limits: Mutex<Vec<RateLimits>>,
}
impl ClientAdmissionPort for Admissions {
fn admit(
&self,
request: ClientAdmissionRequest,
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async move {
self.limits.lock().unwrap().push(request.limits);
assert!(!self.active.swap(true, Ordering::SeqCst));
Ok(ClientAdmissionDecision::Granted)
})
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
}
}
struct Budget {
reject: bool,
active: Arc<AtomicBool>,
charges: Mutex<Vec<ClientBudgetCharge>>,
}
impl ClientBudgetPort for Budget {
fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
Box::pin(async {
assert!(self.active.load(Ordering::SeqCst));
if self.reject {
Err(
GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted")
.with_client_code("key_daily_budget_exceeded"),
)
} else {
Ok(())
}
})
}
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
Box::pin(async move {
assert!(
self.active.load(Ordering::SeqCst),
"settle before releasing the concurrent request slot"
);
self.charges.lock().unwrap().push(charge);
Ok(())
})
}
}
fn service(admissions: Arc<Admissions>, budget: Arc<Budget>) -> DefaultExecutionService {
DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
.with_budget(budget)
}
fn request(service: &DefaultExecutionService, transport: ClientTransport) -> StartExecution {
StartExecution {
client: service.authenticate("sk_start_test").unwrap(),
public_model: PublicModelId::new("gpt-start").unwrap(),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport,
stream: transport != ClientTransport::HttpJson,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}
}
#[test]
fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
reject: true,
active: admissions.active.clone(),
charges: Mutex::default(),
});
let service = service(admissions.clone(), budget.clone());
for transport in [
ClientTransport::HttpJson,
ClientTransport::HttpSse,
ClientTransport::WebSocket,
] {
let result = block_on(service.start(request(&service, transport)));
assert!(
matches!(result, Err(error) if error.client_error_code() == Some("key_daily_budget_exceeded"))
);
assert!(!admissions.active.load(Ordering::SeqCst));
assert!(budget.charges.lock().unwrap().is_empty());
}
}
#[test]
fn reused_client_uses_updated_limits_for_each_execution() {
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
let admissions = Arc::new(Admissions::default());
let service = DefaultExecutionService::new(
snapshots.clone(),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service.authenticate("sk_start_test").unwrap();
let limited = RateLimits {
max_concurrency: 1,
requests_per_minute: 1,
};
for (revision, limits) in [(2, limited), (3, RateLimits::unlimited())] {
snapshots.publish(start_snapshot_with_policy(revision, true, limits));
let mut next = request(&service, ClientTransport::WebSocket);
next.client = client.clone();
let started = block_on(service.start(next)).expect("new execution");
block_on(started.session.detach_finalize());
}
assert_eq!(
*admissions.limits.lock().unwrap(),
vec![limited, RateLimits::unlimited()]
);
}
#[test]
fn reused_client_cannot_start_after_key_disable_or_snapshot_suspension() {
for suspend in [false, true] {
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
let admissions = Arc::new(Admissions::default());
let service = DefaultExecutionService::new(
snapshots.clone(),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let next = request(&service, ClientTransport::WebSocket);
if suspend {
snapshots.suspend();
} else {
snapshots.publish(start_snapshot_with_policy(
2,
false,
RateLimits::unlimited(),
));
}
let result = block_on(service.start(next));
let expected = if suspend {
GatewayErrorKind::Internal
} else {
GatewayErrorKind::Unauthorized
};
assert!(matches!(result, Err(error) if error.kind() == expected));
assert!(admissions.limits.lock().unwrap().is_empty());
}
}
#[test]
fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all_transports() {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
reject: false,
active: admissions.active.clone(),
charges: Mutex::default(),
});
let service = service(admissions.clone(), budget.clone());
for transport in [
ClientTransport::HttpJson,
ClientTransport::HttpSse,
ClientTransport::WebSocket,
] {
for detach in [true, false] {
let mut started = block_on(service.start(request(&service, transport))).unwrap();
assert!(admissions.active.load(Ordering::SeqCst));
if detach {
block_on(started.session.detach_finalize());
} else {
assert!(block_on(started.session.next_event()).is_err());
block_on(started.session.detach_finalize());
}
assert!(!admissions.active.load(Ordering::SeqCst));
}
}
let charges = budget.charges.lock().unwrap();
assert_eq!(
charges.len(),
6,
"one settlement per request, including detached finalized sessions"
);
assert!(
charges
.iter()
.all(|charge| charge.amount_usd == gateway_core::metering::Decimal::ZERO)
);
}
use async_trait::async_trait;
use bytes::Bytes;
use futures::{executor::block_on, future::BoxFuture};
@@ -495,7 +720,13 @@ fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
panic!("a known provider catalog must be authoritative for model availability");
};
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
assert_eq!(
(error.kind(), error.client_message()),
(
GatewayErrorKind::ModelNotFound,
"the requested model was not found in the provider catalogs available to this API key; check the model name",
)
);
}
#[test]
@@ -917,11 +1148,15 @@ fn client_snapshot() -> RuntimeSnapshot {
}
fn start_snapshot() -> RuntimeSnapshot {
start_snapshot_with_policy(1, true, RateLimits::unlimited())
}
fn start_snapshot_with_policy(revision: u64, enabled: bool, limits: RateLimits) -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
ConfigRevision::new(revision).expect("config revision"),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
std::num::NonZeroU32::new(1).expect("concurrency"),
@@ -937,8 +1172,8 @@ fn start_snapshot() -> RuntimeSnapshot {
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
account_scope(&provider, "acct_start"),
true,
RateLimits::unlimited(),
enabled,
limits,
)],
)
.expect("start snapshot")
+127 -13
View File
@@ -326,12 +326,123 @@ fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
)
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "gpt-5.4".to_owned(),
mapped_model: "gpt-5.5".to_owned(),
}
);
}
#[test]
fn missing_model_should_be_reported_when_all_scoped_catalogs_are_known() {
let snapshot = snapshot();
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect_err("neither provider publishes the requested model");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "missing-model".to_owned(),
mapped_model: "missing-model".to_owned(),
}
);
}
#[test]
fn missing_model_in_one_provider_should_still_route_to_another() {
let snapshot = snapshot();
let plan = snapshot
.plan(
&PublicModelId::new("grok-4.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("xAI publishes the model missing from OpenAI");
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
}
#[test]
fn blocked_model_provider_should_not_be_misreported_as_model_not_found() {
let snapshot = snapshot();
let error = snapshot
.plan(
&PublicModelId::new("gpt-5.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
..RoutingContext::default()
},
)
.expect_err("OpenAI is blocked and xAI does not publish this model");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn blocked_unknown_catalog_should_not_prove_model_absence() {
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![
ProviderKind::new("openai").expect("provider"),
ProviderKind::new("xai").expect("provider"),
],
vec![model("xai", "grok-4.5", capabilities())],
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory());
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
..RoutingContext::default()
},
)
.expect_err("the blocked provider's catalog cannot establish model absence");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn empty_account_scope_should_not_be_misreported_as_model_not_found() {
let snapshot = snapshot();
let empty_scope = Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::default()),
ClientRoutingScope::all_accounts(),
));
let error = snapshot
.plan(
&PublicModelId::new("missing-model").expect("model"),
&operation(),
empty_scope,
&RoutingContext::default(),
)
.expect_err("the API key has no accounts");
assert_eq!(error, gateway_core::error::RoutingError::EmptyAccountScope);
}
#[test]
fn provider_endpoint_plan_should_not_consult_or_publish_the_text_model_catalog() {
let snapshot = snapshot();
@@ -463,19 +574,22 @@ fn known_unsupported_operation_should_not_be_bypassed() {
)
.expect("snapshot");
assert!(
snapshot
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.is_err()
);
let error = snapshot
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
.expect_err("the model exists but cannot perform this operation");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
@@ -56,10 +56,13 @@ struct PublishingCatalogProvider {
queries: AtomicUsize,
}
struct UnavailableCatalogProvider;
enum TestCatalogProvider {
Unavailable,
Empty,
}
#[async_trait]
impl Provider for UnavailableCatalogProvider {
impl Provider for TestCatalogProvider {
fn name(&self) -> &'static str {
"alpha"
}
@@ -71,10 +74,13 @@ impl Provider for UnavailableCatalogProvider {
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
match self {
Self::Unavailable => Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Self::Empty => Ok(Vec::new()),
}
}
async fn execute(
@@ -145,7 +151,7 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
#[test]
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
let providers =
ProviderRegistry::new([Arc::new(UnavailableCatalogProvider) as Arc<dyn Provider>])
ProviderRegistry::new([Arc::new(TestCatalogProvider::Unavailable) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
@@ -162,6 +168,35 @@ fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
assert_eq!(snapshot.client_policies().count(), 1);
}
#[test]
fn known_empty_catalog_should_report_model_not_found() {
let providers =
ProviderRegistry::new([Arc::new(TestCatalogProvider::Empty) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let snapshot = block_on(compiler.compile()).expect("compile empty catalog");
let error = snapshot
.plan(
&PublicModelId::new("public-model").expect("model"),
&super::operation(),
snapshot.all_account_scope(),
&gateway_core::routing::RoutingContext {
required_provider: Some(ProviderKind::new("alpha").expect("provider")),
..gateway_core::routing::RoutingContext::default()
},
)
.expect_err("a successfully published empty catalog proves model absence");
assert_eq!(
error,
gateway_core::error::RoutingError::ModelNotFound {
model: "public-model".to_owned(),
mapped_model: "upstream-model".to_owned(),
},
);
}
#[test]
fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
let provider = Arc::new(PublishingCatalogProvider {
+1 -1
View File
@@ -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"] }
+12
View File
@@ -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"));
}
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-protocol"
version = "0.0.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
+1 -1
View File
@@ -2,7 +2,7 @@
name = "gateway-store"
version = "0.1.0"
edition = "2024"
license = "MIT"
license = "Apache-2.0"
publish = false
rust-version = "1.97"
autotests = false
+3 -1
View File
@@ -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;
}
+3 -1
View File
@@ -26,8 +26,9 @@ hex = "0.4.3"
aws-lc-rs = "=1.18.1"
hyper = { version = "=1.11.1", default-features = false }
hyper-rustls = "=0.27.7"
hyper-util = { version = "0.1.20", features = ["client-legacy", "http1", "tokio"] }
# 与官方 Codex Desktop 一致:reqwest 不启用压缩特性,上游请求不带 accept-encoding。
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2", "socks"] }
roxmltree = "0.21.1"
semver = "1.0.27"
# 与官方 rustls 路径一致:aws-lc 支持 P-521,优先提供 X25519MLKEM768。
@@ -44,6 +45,7 @@ thiserror = "2.0.18"
tokio = { version = "=1.53.1", default-features = false, features = ["fs", "macros", "rt", "sync", "time"] }
tokio-rustls = { version = "=0.26.4", default-features = false }
tokio-util = { version = "0.7.18", features = ["rt"] }
tower-service = "0.3.3"
# 官方 Codex WebSocket 需要上游尚未发布的 proxy 与 permessage-deflate 实现,固定到已审计提交。
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "0e5b2d73aa18dd9f0a50ee9ff199d5aef7594186", features = ["proxy", "rustls-tls-native-roots"] }
tracing = "0.1.44"
+64 -5
View File
@@ -166,6 +166,30 @@ impl ProviderAdmin for OpenAiAdminProvider {
&self.provider_kind
}
fn plan_type_display(&self, plan_type: &str) -> String {
// 与官方 Desktop 的套餐名称映射一致;原始子类型继续由 plan_type 保留。
match plan_type.to_ascii_lowercase().as_str() {
"free" | "free_workspace" | "guest" => "Free",
"go" => "Go",
"plus" => "Plus",
"pro" | "prolite" => "Pro",
"team" | "self_serve_business_prolite" | "self_serve_business_usage_based" => {
"Business"
}
"business"
| "ent26"
| "enterprise_cbp_automation"
| "enterprise_cbp_usage_based"
| "enterprise"
| "hc" => "Enterprise",
"edu" | "education" => "Edu",
"edu_plus" => "Edu Plus",
"edu_pro" => "Edu Pro",
_ => plan_type,
}
.to_owned()
}
async fn account_unavailable(&self, account_id: &ProviderAccountId) {
self.websocket_pool.evict_account(account_id.as_str()).await;
}
@@ -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