Compare commits

...
33 Commits
Author SHA1 Message Date
zyycnandCodex 3110b71512 chore: bump version to 3.0.20
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 18:46:20 +08:00
zyycnandCodex eac1b2c39e fix(xai): classify upstream failures by scope
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 18:31:06 +08:00
zyycnandCodex 337bb80230 fix(openai): restore streaming failure fallback
Restore the v2.x stream commit policies so ordinary requests forward the first event immediately while continuation and same-batch failures remain behind the retry boundary. Preserve typed upstream failure batches through Core and both streaming adapters, with a 64 KiB prefetch limit and end-to-end regression coverage.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 17:29:15 +08:00
zyycnandCodex 82f2a0da06 fix(openai): preserve quota exhaustion state
Prefer explicit quota evidence over generic 429 feedback, keep exhausted and terminal accounts from being downgraded by temporary rate limits, and reconcile historical cooldown states from current quota observations.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 14:47:21 +08:00
zyycnandCodex f0d8dad71e chore: bump version to 3.0.19
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 13:05:04 +08:00
zyycnandCodex f5c1d773d7 fix(accounts): use current credential revisions
Load the current credential revision server-side for reauthorization and manual refresh so background token refreshes do not make admin mutations fail stale client fences.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 13:04:12 +08:00
zyycnandCodex 9148e6b11e chore: bump version to 3.0.18
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:48:21 +08:00
zyycnandCodex a7d8c3d3df fix(accounts): make fixed diagnostics resilient
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:45:35 +08:00
zyycnandCodex 997eb86aa0 fix(openai): refresh expired OAuth access tokens
Force refresh scheduling for OAuth accounts whose access token has already expired, even when a persisted next_refresh_at points to the future.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:41:44 +08:00
zyycnandCodex 3866acc269 chore: bump version to 3.0.17
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 00:06:15 +08:00
zyycnandCodex 38477f4328 fix(openai): normalize OAuth import credentials
Import OAuth accounts from access and refresh tokens only, then verify identity and observe the initial quota through official endpoints.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 23:54:59 +08:00
zyycnandCodex 5168f6a924 fix(openai): replace stale passive rate-limit snapshots
Treat each passive response-header observation as a complete rate-limit bucket snapshot, preventing stale secondary windows from surviving into a newer primary-only update.\n\nExtract duplicate OAuth callback parameter rejection to keep the release quality gate warning-free.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 23:12:17 +08:00
zyycnandCodex 3a407c635a chore: bump version to 3.0.16
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 22:39:31 +08:00
zyycnandCodex 6901e0ed45 fix(openai): accept standard OAuth callback parameters
Accept standard OAuth callback parameters while keeping one-shot flow and duplicate security parameter rejection.

Classify unavailable identity verification correctly and add redacted OAuth/import failure diagnostics.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 22:39:01 +08:00
zyycnandClaude Fable 5 9758c420fa refactor(availability): converge account availability on the gateway-core enum
- gateway-admin: drop the mirrored 7-variant enum; re-export the core type
- gateway-api: serialize wire availability via core as_str (bytes unchanged)
- gateway-store: drop ProviderAccountAvailability and four identity mappers;
  rows parse through core's codec with the same error string; redis state
  validation reuses core parse instead of a literal 7-string set
- providers/openai+xai: drop core<->admin identity mappers made moot by the
  admin re-export
- providers/xai: a single exhaustive From<GrokCredentialAvailability> impl
  is now the only Grok->core translation point

No behavior change: DB strings, wire JSON, and the AccountStatus fold
(including Invalid->Expired display) are byte-identical. 1244 tests green.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-27 20:40:47 +08:00
ZyyandCodex 6d2bf06ac7 chore: bump version to 3.0.15
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:29:36 +08:00
ZyyandCodex b8faaf250f fix(frontend): align dashboard account usage metrics
Add a metric layout variant for account usage windows so dashboard account rows align request bars with adjacent recent-usage text without changing the shared compact account-table presentation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:22:28 +08:00
ZyyandCodex 40137b1205 fix(backend): preserve safe upstream error details
Keep admin account connection tests routed through the provider probe even when the stored account is already marked unavailable so clicking test can refresh state from the upstream response.

Preserve provider-declared safe upstream details when OpenAI-compatible API errors collapse local provider unavailability to the stable no-provider contract.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:09:59 +08:00
ZyyandCodex c67cde4e90 chore: bump version to 3.0.14
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 16:28:29 +08:00
ZyyandCodex 2e9213a784 fix(backend): derive Codex metadata from provider catalogs
Share Codex Responses request semantics across OpenAI and xAI providers so reasoning effort, presets, request kind, subagent kind, and compaction observations stay consistent.

Read xAI reasoning efforts from official catalog feature options instead of model-name hardcoding, and make Codex service tiers data-driven through ModelPresentation rather than API adapter defaults.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 16:27:41 +08:00
ZyyandCodex 87496b5656 fix(frontend): polish dashboard layout
Keep dashboard header actions pinned to the action column, widen the xAI runtime field, and align active-account usage metrics without changing the shared usage window component.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 14:55:36 +08:00
ZyyandCodex c856e4306f fix(xai): mark rate limits as quota exhausted
Classify xAI 429 rate-limit responses as quota exhaustion so the scheduler removes the account from active rotation until the normal recovery flow makes it available again.

Cover the provider replay path and transport classification with tests.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 14:55:30 +08:00
ZyyandCodex 92090f28bf docs(backend): unify provider comments to Chinese
Convert the remaining English doc comments in the openai and xai provider
crates to Chinese, function-oriented wording so the whole backend follows
one convention. Explanatory narration is compressed into what the item does
and what its errors mean; code identifiers, route paths, rustdoc links and
`# Errors` headings stay as-is.

Also fixes two half-translated comments left in the websocket protocol
module and merges a duplicated doc block on the OIDC verifier constructor.
No code, signature or `#[error(...)]` string changes.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 13:13:44 +08:00
ZyyandCodex 12c2b612de chore(db): fold retention indexes into initial migration
Keep the database baseline as a single initial migration now that deployed instances have already been upgraded. The retention delete indexes move into 0001, the standalone 0002 file is removed, and the migration tests now expect one applied migration.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 12:05:49 +08:00
e8af0b9c66 chore(deps): Bump ed25519-dalek from 2.2.0 to 3.0.0 in /backend (#26)
* chore(deps): Bump ed25519-dalek from 2.2.0 to 3.0.0 in /backend

Bumps [ed25519-dalek](https://github.com/dalek-cryptography/curve25519-dalek) from 2.2.0 to 3.0.0.
- [Release notes](https://github.com/dalek-cryptography/curve25519-dalek/releases)
- [Changelog](https://github.com/dalek-cryptography/curve25519-dalek/blob/3.0.0/CHANGELOG.md)
- [Commits](https://github.com/dalek-cryptography/curve25519-dalek/compare/ed25519-2.2.0...3.0.0)

---
updated-dependencies:
- dependency-name: ed25519-dalek
  dependency-version: 3.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

* fix: enable ed25519 PKCS8 allocation support

ed25519-dalek 3 gates private-key PKCS#8 encoding behind the alloc feature. Enable it alongside pkcs8 so existing credential export tests continue to compile and exercise the same behavior.

Co-authored-by: Codex <noreply@openai.com>

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Zyy <No Email>
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 12:04:48 +08:00
80776e8e03 chore(deps): Bump jsonwebtoken from 10.3.0 to 11.0.0 in /backend (#24)
* chore(deps): Bump jsonwebtoken from 10.3.0 to 11.0.0 in /backend

Bumps [jsonwebtoken](https://github.com/Keats/jsonwebtoken) from 10.3.0 to 11.0.0.
- [Changelog](https://github.com/Keats/jsonwebtoken/blob/master/CHANGELOG.md)
- [Commits](https://github.com/Keats/jsonwebtoken/compare/v10.3.0...v11.0.0)

---
updated-dependencies:
- dependency-name: jsonwebtoken
  dependency-version: 11.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

* fix: adapt jsonwebtoken 11 header extras API

jsonwebtoken 11 wraps non-standard JWT header fields in Extras. Keep the existing strict header validation by checking the wrapped inner map instead of calling is_empty directly.

Co-authored-by: Codex <noreply@openai.com>

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Zyy <No Email>
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 11:54:06 +08:00
dependabot[bot] b296c795ac chore(deps): Bump the minor-and-patch group across 1 directory with 3 updates (#22)
Bumps the minor-and-patch group with 3 updates in the / directory: [actions/checkout](https://github.com/actions/checkout), [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) and [docker/login-action](https://github.com/docker/login-action).


Updates `actions/checkout` from 7.0.0 to 7.0.1
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0...3d3c42e5aac5ba805825da76410c181273ba90b1)

Updates `github/codeql-action/upload-sarif` from 4.37.1 to 4.37.3
- [Release notes](https://github.com/github/codeql-action/releases)
- [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md)
- [Commits](https://github.com/github/codeql-action/compare/7188fc363630916deb702c7fdcf4e481b751f97a...e4fba868fa4b1b91e1fdab776edc8cfbe6e9fb81)

Updates `docker/login-action` from 4.4.0 to 4.5.1
- [Release notes](https://github.com/docker/login-action/releases)
- [Commits](https://github.com/docker/login-action/compare/af1e73f918a031802d376d3c8bbc3fe56130a9b0...abd2ef45e78c5afb21d64d4ca52ee8550d9572c7)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: 7.0.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: docker/login-action
  dependency-version: 4.5.1
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: github/codeql-action/upload-sarif
  dependency-version: 4.37.3
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-27 11:41:57 +08:00
dependabot[bot] 9fc162a82f chore(deps): Bump base64 from 0.22.1 to 0.23.0 in /backend (#25)
Bumps [base64](https://github.com/marshallpierce/rust-base64) from 0.22.1 to 0.23.0.
- [Changelog](https://github.com/marshallpierce/rust-base64/blob/master/RELEASE-NOTES.md)
- [Commits](https://github.com/marshallpierce/rust-base64/compare/v0.22.1...v0.23.0)

---
updated-dependencies:
- dependency-name: base64
  dependency-version: 0.23.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-27 11:41:53 +08:00
dependabot[bot] 9532bd438b chore(deps): Bump getrandom from 0.3.4 to 0.4.3 in /backend (#27)
Bumps [getrandom](https://github.com/rust-random/getrandom) from 0.3.4 to 0.4.3.
- [Changelog](https://github.com/rust-random/getrandom/blob/master/CHANGELOG.md)
- [Commits](https://github.com/rust-random/getrandom/compare/v0.3.4...v0.4.3)

---
updated-dependencies:
- dependency-name: getrandom
  dependency-version: 0.4.3
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-27 11:41:49 +08:00
dependabot[bot] 73319f1943 chore(deps): Bump syn from 2.0.118 to 3.0.3 in /backend (#28)
Bumps [syn](https://github.com/dtolnay/syn) from 2.0.118 to 3.0.3.
- [Release notes](https://github.com/dtolnay/syn/releases)
- [Commits](https://github.com/dtolnay/syn/compare/2.0.118...3.0.3)

---
updated-dependencies:
- dependency-name: syn
  dependency-version: 3.0.3
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-27 11:41:45 +08:00
dependabot[bot] ee56d6a1d2 chore(deps): Bump the minor-and-patch group across 1 directory with 8 updates (#29)
Bumps the minor-and-patch group with 8 updates in the /frontend directory:

| Package | From | To |
| --- | --- | --- |
| [@fontsource-variable/inter](https://github.com/fontsource/font-files/tree/HEAD/fonts/variable/inter) | `5.2.8` | `5.3.0` |
| [@fontsource-variable/jetbrains-mono](https://github.com/fontsource/font-files/tree/HEAD/fonts/variable/jetbrains-mono) | `5.2.8` | `5.3.0` |
| [@lucide/vue](https://github.com/lucide-icons/lucide/tree/HEAD/packages/vue) | `1.26.0` | `1.27.0` |
| [es-toolkit](https://github.com/toss/es-toolkit) | `1.49.0` | `1.50.0` |
| [marked](https://github.com/markedjs/marked) | `18.0.6` | `18.0.7` |
| [@antfu/eslint-config](https://github.com/antfu/eslint-config) | `9.1.0` | `9.2.0` |
| [eslint](https://github.com/eslint/eslint) | `10.7.0` | `10.8.0` |
| [vue-tsc](https://github.com/vuejs/language-tools/tree/HEAD/packages/tsc) | `3.3.7` | `3.3.8` |



Updates `@fontsource-variable/inter` from 5.2.8 to 5.3.0
- [Changelog](https://github.com/fontsource/font-files/blob/main/CHANGELOG.md)
- [Commits](https://github.com/fontsource/font-files/commits/HEAD/fonts/variable/inter)

Updates `@fontsource-variable/jetbrains-mono` from 5.2.8 to 5.3.0
- [Changelog](https://github.com/fontsource/font-files/blob/main/CHANGELOG.md)
- [Commits](https://github.com/fontsource/font-files/commits/HEAD/fonts/variable/jetbrains-mono)

Updates `@lucide/vue` from 1.26.0 to 1.27.0
- [Release notes](https://github.com/lucide-icons/lucide/releases)
- [Commits](https://github.com/lucide-icons/lucide/commits/1.27.0/packages/vue)

Updates `es-toolkit` from 1.49.0 to 1.50.0
- [Release notes](https://github.com/toss/es-toolkit/releases)
- [Changelog](https://github.com/toss/es-toolkit/blob/main/CHANGELOG.md)
- [Commits](https://github.com/toss/es-toolkit/compare/v1.49.0...v1.50.0)

Updates `marked` from 18.0.6 to 18.0.7
- [Release notes](https://github.com/markedjs/marked/releases)
- [Commits](https://github.com/markedjs/marked/compare/v18.0.6...v18.0.7)

Updates `@antfu/eslint-config` from 9.1.0 to 9.2.0
- [Release notes](https://github.com/antfu/eslint-config/releases)
- [Commits](https://github.com/antfu/eslint-config/compare/v9.1.0...v9.2.0)

Updates `eslint` from 10.7.0 to 10.8.0
- [Release notes](https://github.com/eslint/eslint/releases)
- [Commits](https://github.com/eslint/eslint/compare/v10.7.0...v10.8.0)

Updates `vue-tsc` from 3.3.7 to 3.3.8
- [Release notes](https://github.com/vuejs/language-tools/releases)
- [Changelog](https://github.com/vuejs/language-tools/blob/master/CHANGELOG.md)
- [Commits](https://github.com/vuejs/language-tools/commits/v3.3.8/packages/tsc)

---
updated-dependencies:
- dependency-name: "@fontsource-variable/inter"
  dependency-version: 5.3.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: "@fontsource-variable/jetbrains-mono"
  dependency-version: 5.3.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: "@lucide/vue"
  dependency-version: 1.27.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: es-toolkit
  dependency-version: 1.50.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: marked
  dependency-version: 18.0.7
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: "@antfu/eslint-config"
  dependency-version: 9.2.0
  dependency-type: direct:development
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: eslint
  dependency-version: 10.8.0
  dependency-type: direct:development
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: vue-tsc
  dependency-version: 3.3.8
  dependency-type: direct:development
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-27 11:41:41 +08:00
ZyyandCodex 5d29ddfffc chore: bump version to 3.0.13
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 11:33:55 +08:00
ZyyandCodex 45903b684b fix(admin): remove config revision fences from API writes
Operator mutations no longer require callers to carry configRevision values through accounts, client keys, provider credentials, or settings. The store and admin use cases now apply current state directly, which removes stale-revision failures from normal create, update, and delete flows.

xAI Responses catalog entries now let the upstream validate feature support instead of rejecting requests locally when optional catalog metadata is absent, matching the Build routing behavior used by grok2api.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 11:30:59 +08:00
158 changed files with 9217 additions and 3832 deletions
+2 -2
View File
@@ -40,7 +40,7 @@ jobs:
} >> "$GITHUB_ENV"
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ inputs.ref }}
@@ -76,7 +76,7 @@ jobs:
${{ hashFiles('trivy-results.sarif') != '' &&
(github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository) }}
uses: github/codeql-action/upload-sarif@7188fc363630916deb702c7fdcf4e481b751f97a # v4.37.1
uses: github/codeql-action/upload-sarif@e4fba868fa4b1b91e1fdab776edc8cfbe6e9fb81 # v4.37.3
with:
sarif_file: trivy-results.sarif
category: runtime-container
+3 -3
View File
@@ -61,7 +61,7 @@ jobs:
--health-retries 10
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ inputs.ref }}
# PR 场景需要 base 提交做冻结清单的 append-only 比对。
@@ -152,7 +152,7 @@ jobs:
CI: 'true'
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ inputs.ref }}
@@ -191,7 +191,7 @@ jobs:
timeout-minutes: 5
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ inputs.ref }}
+1 -1
View File
@@ -25,7 +25,7 @@ jobs:
container: ${{ steps.filter.outputs.container }}
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
fetch-depth: 0
+8 -8
View File
@@ -41,7 +41,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref }}
fetch-depth: 0
@@ -177,7 +177,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ needs.release-info.outputs.git_sha }}
@@ -237,7 +237,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ needs.release-info.outputs.git_sha }}
@@ -289,7 +289,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ needs.release-info.outputs.git_sha }}
@@ -408,7 +408,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ needs.release-info.outputs.git_sha }}
@@ -454,7 +454,7 @@ jobs:
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
- name: Login to GHCR
uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0
uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7 # v4.5.1
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
@@ -508,7 +508,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ needs.release-info.outputs.git_sha }}
@@ -525,7 +525,7 @@ jobs:
uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2
- name: Login to GHCR
uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0
uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7 # v4.5.1
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
+2 -2
View File
@@ -23,7 +23,7 @@ jobs:
timeout-minutes: 20
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
@@ -47,7 +47,7 @@ jobs:
timeout-minutes: 15
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
+126 -59
View File
@@ -142,7 +142,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -208,7 +208,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
dependencies = [
"axum-core",
"base64",
"base64 0.22.1",
"bytes",
"form_urlencoded",
"futures-util",
@@ -275,6 +275,12 @@ version = "0.22.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
[[package]]
name = "base64"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9"
[[package]]
name = "base64ct"
version = "1.8.3"
@@ -444,7 +450,7 @@ dependencies = [
"provider-openai",
"provider-xai",
"serde",
"syn 2.0.118",
"syn 3.0.3",
"tempfile",
"thiserror",
"tokio",
@@ -669,7 +675,7 @@ dependencies = [
"aead",
"blake2",
"crypto_secretbox",
"curve25519-dalek",
"curve25519-dalek 4.1.3",
"salsa20",
"subtle",
"zeroize",
@@ -708,8 +714,23 @@ dependencies = [
"cfg-if",
"cpufeatures 0.2.17",
"curve25519-dalek-derive",
"digest 0.10.7",
"fiat-crypto",
"fiat-crypto 0.2.9",
"rustc_version",
"subtle",
"zeroize",
]
[[package]]
name = "curve25519-dalek"
version = "5.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b5eed333089e2e1c1ac8c6c0398e5e2497b4c9926ca6d0365ed1e099afa5bc23"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
"curve25519-dalek-derive",
"digest 0.11.3",
"fiat-crypto 0.3.0",
"rustc_version",
"subtle",
"zeroize",
@@ -723,7 +744,7 @@ checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -761,6 +782,16 @@ dependencies = [
"zeroize",
]
[[package]]
name = "der"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d"
dependencies = [
"const-oid 0.10.2",
"zeroize",
]
[[package]]
name = "deranged"
version = "0.5.8"
@@ -799,7 +830,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -829,34 +860,35 @@ version = "0.16.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca"
dependencies = [
"der",
"der 0.7.10",
"digest 0.10.7",
"elliptic-curve",
"rfc6979",
"signature",
"spki",
"signature 2.2.0",
"spki 0.7.3",
]
[[package]]
name = "ed25519"
version = "2.2.3"
version = "3.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53"
checksum = "29fcf32e6c73d1079f83ab4d782de2d81620346a5f38c6237a86a22f8368980a"
dependencies = [
"pkcs8",
"signature",
"pkcs8 0.11.0",
"signature 3.0.0",
]
[[package]]
name = "ed25519-dalek"
version = "2.2.0"
version = "3.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9"
checksum = "6ebaa1a2bf1290ab3bfe5a7b771d050ebffab2711c19a81691c683a5144a25de"
dependencies = [
"curve25519-dalek",
"curve25519-dalek 5.0.0",
"ed25519",
"serde",
"sha2 0.10.9",
"sha2 0.11.0",
"signature 3.0.0",
"subtle",
"zeroize",
]
@@ -883,7 +915,7 @@ dependencies = [
"generic-array",
"group",
"pem-rfc7468",
"pkcs8",
"pkcs8 0.10.2",
"rand_core 0.6.4",
"sec1",
"subtle",
@@ -968,6 +1000,12 @@ version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d"
[[package]]
name = "fiat-crypto"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "64cd1e32ddd350061ae6edb1b082d7c54915b5c672c389143b9a63403a109f24"
[[package]]
name = "filetime"
version = "0.2.29"
@@ -1100,7 +1138,7 @@ checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -1144,7 +1182,7 @@ version = "0.0.0"
dependencies = [
"argon2",
"async-trait",
"base64",
"base64 0.23.0",
"chrono",
"futures",
"gateway-core",
@@ -1156,6 +1194,7 @@ dependencies = [
"subtle",
"thiserror",
"tokio",
"tracing",
"uuid",
]
@@ -1165,7 +1204,7 @@ version = "0.1.0"
dependencies = [
"async-trait",
"axum",
"base64",
"base64 0.23.0",
"bytes",
"chrono",
"futures",
@@ -1176,6 +1215,7 @@ dependencies = [
"serde_json",
"thiserror",
"tokio",
"tokio-tungstenite 0.29.0",
"tower",
"tower-http 0.7.0",
"tracing",
@@ -1371,7 +1411,7 @@ version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"headers-core",
"http",
@@ -1538,7 +1578,7 @@ version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-util",
@@ -1736,19 +1776,20 @@ dependencies = [
[[package]]
name = "jsonwebtoken"
version = "10.3.0"
version = "11.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1"
checksum = "881733cbc631fc9e472e24447ce32a64bedf2da498d6d8570b08edc87de71f65"
dependencies = [
"aws-lc-rs",
"base64",
"base64 0.22.1",
"getrandom 0.2.17",
"js-sys",
"pem",
"serde",
"serde_json",
"signature",
"signature 2.2.0",
"simple_asn1",
"zeroize",
]
[[package]]
@@ -2015,7 +2056,7 @@ version = "3.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be"
dependencies = [
"base64",
"base64 0.22.1",
"serde_core",
]
@@ -2046,8 +2087,18 @@ version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
dependencies = [
"der",
"spki",
"der 0.7.10",
"spki 0.7.3",
]
[[package]]
name = "pkcs8"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "451913da69c775a56034ea8d9003d27ee8948e12443eae7c038ba100a4f21cb7"
dependencies = [
"der 0.8.1",
"spki 0.8.0",
]
[[package]]
@@ -2115,7 +2166,7 @@ version = "0.0.0"
dependencies = [
"async-stream",
"async-trait",
"base64",
"base64 0.23.0",
"bytes",
"chrono",
"cookie",
@@ -2125,7 +2176,7 @@ dependencies = [
"gateway-admin",
"gateway-core",
"gateway-protocol",
"getrandom 0.3.4",
"getrandom 0.4.3",
"hex",
"jsonwebtoken",
"reqwest",
@@ -2157,14 +2208,14 @@ version = "0.0.0"
dependencies = [
"async-stream",
"async-trait",
"base64",
"base64 0.23.0",
"bytes",
"chrono",
"futures",
"gateway-admin",
"gateway-core",
"gateway-protocol",
"getrandom 0.3.4",
"getrandom 0.4.3",
"jsonwebtoken",
"p256",
"reqwest",
@@ -2409,7 +2460,7 @@ version = "0.12.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"cookie",
"cookie_store",
@@ -2598,9 +2649,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc"
dependencies = [
"base16ct",
"der",
"der 0.7.10",
"generic-array",
"pkcs8",
"pkcs8 0.10.2",
"subtle",
"zeroize",
]
@@ -2796,6 +2847,12 @@ dependencies = [
"rand_core 0.6.4",
]
[[package]]
name = "signature"
version = "3.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5"
[[package]]
name = "simd-adler32"
version = "0.3.9"
@@ -2855,7 +2912,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
dependencies = [
"base64ct",
"der",
"der 0.7.10",
]
[[package]]
name = "spki"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d9efca8738c78ee9484207732f728b1ef517bbb1833d6fc0879ca898a522f6f"
dependencies = [
"base64ct",
"der 0.8.1",
]
[[package]]
@@ -2877,7 +2944,7 @@ version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"cfg-if",
"chrono",
@@ -2918,7 +2985,7 @@ dependencies = [
"quote",
"sqlx-core",
"sqlx-macros-core",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -2941,7 +3008,7 @@ dependencies = [
"sqlx-mysql",
"sqlx-postgres",
"sqlx-sqlite",
"syn 2.0.118",
"syn 2.0.119",
"thiserror",
"tokio",
"url",
@@ -2981,7 +3048,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e"
dependencies = [
"atoi",
"base64",
"base64 0.22.1",
"bitflags",
"byteorder",
"chrono",
@@ -3066,9 +3133,9 @@ checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a"
[[package]]
name = "syn"
version = "2.0.118"
version = "2.0.119"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422"
checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
dependencies = [
"proc-macro2",
"quote",
@@ -3103,7 +3170,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -3124,7 +3191,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.3.4",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
@@ -3238,7 +3305,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -3414,7 +3481,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -3691,7 +3758,7 @@ dependencies = [
"bumpalo",
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
"wasm-bindgen-shared",
]
@@ -3773,7 +3840,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -3784,7 +3851,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -3909,7 +3976,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031"
dependencies = [
"assert-json-diff",
"base64",
"base64 0.22.1",
"deadpool",
"futures",
"http",
@@ -3983,7 +4050,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
"synstructure",
]
@@ -4004,7 +4071,7 @@ checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -4024,7 +4091,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
"synstructure",
]
@@ -4045,7 +4112,7 @@ checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
@@ -4078,7 +4145,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
"syn 2.0.119",
]
[[package]]
+1 -1
View File
@@ -25,7 +25,7 @@ tokio = { version = "1.52.3", default-features = false, features = ["rt-multi-th
[dev-dependencies]
config = { version = "0.15.23", default-features = false, features = ["yaml"] }
syn = { version = "2.0.118", features = ["full"] }
syn = { version = "3.0.3", features = ["full"] }
tempfile = "3.27.0"
[lints]
+2 -1
View File
@@ -14,7 +14,7 @@ path = "tests/main.rs"
[dependencies]
argon2 = "0.5.3"
async-trait = "0.1.89"
base64 = "0.22.1"
base64 = "0.23.0"
chrono = "0.4.45"
futures = "0.3.32"
gateway-core = { path = "../gateway-core" }
@@ -24,6 +24,7 @@ serde = { version = "1.0.228", features = ["derive"] }
secrecy = { version = "0.10.3", features = ["serde"] }
subtle = "2.6.1"
thiserror = "2.0.18"
tracing = "0.1.44"
uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
@@ -9,17 +9,8 @@ use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision, observability::TimeRange};
/// Provider 账号当前可用性。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
/// Provider 账号当前可用性;权威定义与字符串编解码在 gateway-core,此处仅复用。
pub use gateway_core::engine::credential::AccountAvailability;
/// 管理页使用的归一化账号状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -189,7 +180,6 @@ pub struct AccountSummary {
/// 账号启停写入命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetAccountEnabled {
pub expected_config_revision: Revision,
pub account_id: String,
pub enabled: bool,
}
@@ -197,7 +187,6 @@ pub struct SetAccountEnabled {
/// 账号批量删除命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteAccounts {
pub expected_config_revision: Revision,
pub account_ids: Vec<String>,
}
@@ -141,7 +141,6 @@ impl fmt::Debug for ClientKeySecret {
/// API 提交的 Client Key 创建命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateClientKey {
pub expected_config_revision: Revision,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
@@ -151,7 +150,6 @@ pub struct CreateClientKey {
/// 管理用例生成 ID 与明文后的持久化命令。
#[derive(Clone, PartialEq, Eq)]
pub struct NewClientKey {
pub expected_config_revision: Revision,
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
@@ -175,7 +173,6 @@ impl fmt::Debug for NewClientKey {
/// 修改 Client Key 的公开策略字段。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientKey {
pub expected_config_revision: Revision,
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
@@ -186,7 +183,6 @@ pub struct UpdateClientKey {
/// 修改 Client Key 启用状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetClientKeyEnabled {
pub expected_config_revision: Revision,
pub id: ClientApiKeyId,
pub enabled: bool,
}
@@ -194,7 +190,6 @@ pub struct SetClientKeyEnabled {
/// 删除 Client Key。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteClientKey {
pub expected_config_revision: Revision,
pub id: ClientApiKeyId,
}
@@ -2,6 +2,7 @@
use std::fmt;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
@@ -9,7 +10,7 @@ use gateway_core::{
};
use super::{
MutationActor, MutationContext, PageSize, Revision,
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{AccountAvailability, AccountRecord, AccountStatus, AccountSummary, AccountUsage},
};
@@ -101,7 +102,6 @@ pub struct CredentialDetails {
/// Provider 正式文档批量导入命令。
pub struct ImportCredentials {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub document: ProviderDocument,
}
@@ -110,7 +110,6 @@ impl fmt::Debug for ImportCredentials {
formatter
.debug_struct("ImportCredentials")
.field("context", &self.context)
.field("expected_config_revision", &self.expected_config_revision)
.field("document", &self.document)
.finish()
}
@@ -169,17 +168,9 @@ pub struct PreparedCredentialImport {
/// Admin 交给 Store 的导入事务命令。
#[derive(Debug, Clone, PartialEq)]
pub struct CredentialImportCommit {
pub expected_config_revision: Revision,
pub prepared: PreparedCredentialImport,
}
/// 可选的重新授权目标。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReauthorizationTarget {
pub account_id: ProviderAccountId,
pub credential_revision: Revision,
}
/// OAuth pending owner 的中立身份;不编码具体 Provider 的 Redis key 或 JSON。
#[derive(Clone, PartialEq, Eq)]
pub enum AuthorizationOwner {
@@ -240,21 +231,18 @@ impl fmt::Debug for AuthorizationOwnerBinding {
}
/// OAuth 完成时应创建新账号还是 CAS 更新既有 credential。
///
/// 重新授权只绑定稳定的账号身份:credential revision 由服务端在临近写入时读取,
/// 长流程期间的后台刷新不得让恢复操作失效。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthorizationMutationTarget {
Create {
name: String,
},
Reauthorize {
account_id: ProviderAccountId,
expected_credential_revision: Revision,
},
Create { name: String },
Reauthorize { account_id: ProviderAccountId },
}
/// 必须完整进入 Provider opaque pending payload、并在 complete 后原样恢复的事务信封。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PendingAuthorizationMutation {
expected_config_revision: Revision,
provider_kind: ProviderKind,
target: AuthorizationMutationTarget,
owner_binding: AuthorizationOwnerBinding,
@@ -263,24 +251,17 @@ pub struct PendingAuthorizationMutation {
impl PendingAuthorizationMutation {
#[must_use]
pub const fn new(
expected_config_revision: Revision,
provider_kind: ProviderKind,
target: AuthorizationMutationTarget,
owner_binding: AuthorizationOwnerBinding,
) -> Self {
Self {
expected_config_revision,
provider_kind,
target,
owner_binding,
}
}
#[must_use]
pub const fn expected_config_revision(&self) -> Revision {
self.expected_config_revision
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
@@ -301,9 +282,8 @@ impl PendingAuthorizationMutation {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StartAuthorization {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub name: String,
pub reauthorization: Option<ReauthorizationTarget>,
pub reauthorization: Option<ProviderAccountId>,
}
/// OAuth 流程启动结果。
@@ -322,6 +302,20 @@ pub trait CredentialCommitGuard: Send + 'static {
fn finish(self: Box<Self>);
}
/// OAuth pending claim 在 Store 事务结果确定后执行的显式结算动作。
///
/// OAuth provider 必须在 credential 准备阶段持有 claim:Store 提交成功后消费 flow,提交或
/// 校验失败后释放 flow。异步结算不能依赖 `Drop`,否则请求提前返回会让同一授权流无谓地
/// 进入短暂的不可重试状态。
#[async_trait]
pub trait AuthorizationCommitGuard: Send + 'static {
/// Store 已提交 OAuth credential 后消费对应的一次性 flow。
async fn commit(self: Box<Self>) -> Result<(), AdminError>;
/// Store 未提交 OAuth credential 时释放 claim,允许同一 flow 重试。
async fn abort(self: Box<Self>) -> Result<(), AdminError>;
}
/// OAuth complete 后由 Provider 返回的准备结果;Store 仍是唯一提交者。
pub enum PreparedAuthorizationCredential {
Create(PreparedCredentialCreate),
@@ -332,6 +326,7 @@ pub enum PreparedAuthorizationCredential {
pub struct PreparedAuthorizationCommit {
pub pending: PendingAuthorizationMutation,
pub credential: PreparedAuthorizationCredential,
authorization_guard: Option<Box<dyn AuthorizationCommitGuard>>,
}
impl fmt::Debug for PreparedAuthorizationCommit {
@@ -358,9 +353,37 @@ pub struct AuthorizationCommit {
pub credential: AuthorizationCredentialCommit,
}
/// OAuth 准备结果拆解后的 Store 命令与两个结算 guard。
///
/// Store 成功时先结束 credential guard,再消费 OAuth claim;失败时丢弃 credential guard 并
/// 释放 OAuth claim,使同一授权流可以重试。
pub(crate) struct AuthorizationCommitSettlement {
pub(crate) command: AuthorizationCommit,
pub(crate) credential_guard: Option<Box<dyn CredentialCommitGuard>>,
pub(crate) authorization_guard: Option<Box<dyn AuthorizationCommitGuard>>,
}
impl PreparedAuthorizationCommit {
#[must_use]
pub fn into_commit(self) -> (AuthorizationCommit, Option<Box<dyn CredentialCommitGuard>>) {
pub fn new(
pending: PendingAuthorizationMutation,
credential: PreparedAuthorizationCredential,
) -> Self {
Self {
pending,
credential,
authorization_guard: None,
}
}
/// 让 OAuth claim 覆盖准备完成到 Store 事务结算之间的全部窗口。
#[must_use]
pub fn with_authorization_guard(mut self, guard: Box<dyn AuthorizationCommitGuard>) -> Self {
self.authorization_guard = Some(guard);
self
}
pub(crate) fn into_commit(self) -> AuthorizationCommitSettlement {
let (credential, guard) = match self.credential {
PreparedAuthorizationCredential::Create(credential) => {
(AuthorizationCredentialCommit::Create(credential), None)
@@ -373,13 +396,27 @@ impl PreparedAuthorizationCommit {
)
}
};
(
AuthorizationCommit {
AuthorizationCommitSettlement {
command: AuthorizationCommit {
pending: self.pending,
credential,
},
guard,
)
credential_guard: guard,
authorization_guard: self.authorization_guard,
}
}
pub(crate) async fn abort(self) -> Result<(), AdminError> {
let AuthorizationCommitSettlement {
credential_guard,
authorization_guard,
..
} = self.into_commit();
drop(credential_guard);
if let Some(guard) = authorization_guard {
guard.abort().await?;
}
Ok(())
}
}
@@ -406,7 +443,6 @@ impl fmt::Debug for CompleteAuthorization {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialMutation {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub account_id: ProviderAccountId,
}
@@ -422,7 +458,6 @@ pub struct CredentialMutationResult {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialDeletion {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub account_ids: Vec<ProviderAccountId>,
}
@@ -436,14 +471,12 @@ pub struct CredentialDeletionResult {
/// Provider-owned token 轮换命令。
pub struct RotateCredential {
pub mutation: CredentialMutation,
pub expected_credential_revision: Revision,
pub provider_material: ProviderDocument,
}
/// Provider 校验手工轮换材料时所需的非事务输入。
pub struct PrepareCredentialRotation {
pub account: AccountRecord,
pub expected_credential_revision: Revision,
pub provider_material: ProviderDocument,
}
@@ -452,10 +485,6 @@ impl fmt::Debug for PrepareCredentialRotation {
formatter
.debug_struct("PrepareCredentialRotation")
.field("account", &self.account)
.field(
"expected_credential_revision",
&self.expected_credential_revision,
)
.field("provider_material", &self.provider_material)
.finish()
}
@@ -520,7 +549,6 @@ impl fmt::Debug for PreparedCredentialRotation {
/// Admin 交给 Store 的轮换或 refresh 事务命令。
#[derive(Debug, Clone, PartialEq)]
pub struct CredentialRotationCommit {
pub expected_config_revision: Revision,
pub prepared: PreparedCredentialRotationFacts,
}
@@ -529,10 +557,6 @@ impl fmt::Debug for RotateCredential {
formatter
.debug_struct("RotateCredential")
.field("mutation", &self.mutation)
.field(
"expected_credential_revision",
&self.expected_credential_revision,
)
.field("provider_material", &self.provider_material)
.finish()
}
@@ -39,7 +39,6 @@ pub struct RuntimeSettings {
/// 原子替换运行设置的命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReplaceRuntimeSettings {
pub expected_config_revision: Revision,
pub model_mappings: ModelMappings,
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u32,
@@ -179,7 +179,7 @@ pub trait AuthStore: Send + Sync {
async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()>;
}
/// Client API Key 管理及其 revision CAS 写入。
/// Client API Key 管理写入。
#[async_trait]
pub trait ClientKeyStore: Send + Sync {
async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage>;
@@ -265,7 +265,7 @@ pub trait ObservabilityStore: Send + Sync {
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage>;
}
/// Runtime settings 与管理员 API Key 的 revision CAS 写入。
/// Runtime settings 与管理员 API Key 写入。
#[async_trait]
pub trait SettingsStore: Send + Sync {
async fn load_runtime_settings(&self) -> AdminStoreResult<RuntimeSettings>;
@@ -280,14 +280,12 @@ pub trait SettingsStore: Send + Sync {
async fn replace_admin_api_key(
&self,
expected_config_revision: Revision,
key: AdminApiKey,
context: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation>;
async fn delete_admin_api_key(
&self,
expected_config_revision: Revision,
context: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation>;
}
@@ -15,7 +15,7 @@ use gateway_core::{
use crate::{
model::{
AdminError, MutationContext, Revision,
AdminError, MutationContext,
accounts::{
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
AccountRecord, AccountStatus, AccountUsageWindowQuery,
@@ -54,7 +54,6 @@ pub trait AccountsService: Send + Sync {
async fn refresh(
&self,
context: &MutationContext,
expected_config_revision: Revision,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError>;
@@ -354,7 +353,6 @@ impl AccountsService for DefaultAccountsService {
async fn refresh(
&self,
context: &MutationContext,
expected_config_revision: Revision,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError> {
let (account, provider) = self.provider_for_account(&account_id).await?;
@@ -367,7 +365,6 @@ impl AccountsService for DefaultAccountsService {
validate_prepared_rotation(&account, &prepared, "provider credential refresh")?;
let result = commit_credential_refresh(
self.accounts.as_ref(),
expected_config_revision,
prepared,
context,
"provider credential refresh",
@@ -94,7 +94,6 @@ impl ClientKeyService for DefaultClientKeyService {
.store
.create_client_key(
NewClientKey {
expected_config_revision: command.expected_config_revision,
id,
name: command.name,
label: command.label,
@@ -90,17 +90,10 @@ async fn pending_authorization(
resource: &'static str,
) -> Result<PendingAuthorizationMutation, AdminError> {
let target = match &command.reauthorization {
Some(target) => {
let details =
required_credential(accounts, provider_kind, &target.account_id, resource).await?;
if details.credential.credential_revision != target.credential_revision {
return Err(AdminError::conflict(format!(
"{resource} reauthorization target is stale"
)));
}
Some(account_id) => {
required_credential(accounts, provider_kind, account_id, resource).await?;
AuthorizationMutationTarget::Reauthorize {
account_id: target.account_id.clone(),
expected_credential_revision: target.credential_revision,
account_id: account_id.clone(),
}
}
None => AuthorizationMutationTarget::Create {
@@ -108,7 +101,6 @@ async fn pending_authorization(
},
};
Ok(PendingAuthorizationMutation::new(
command.expected_config_revision,
provider_kind.clone(),
target,
crate::model::provider_credentials::AuthorizationOwnerBinding::from_context(
@@ -143,7 +135,6 @@ fn validate_prepared_rotation(
let facts = prepared.facts();
if facts.account_id.as_str() != account.id.as_str()
|| facts.provider_kind != account.provider_kind
|| facts.expected_credential_revision != account.credential_revision
{
return Err(AdminError::conflict(format!(
"{resource} prepared facts do not match the current credential"
@@ -152,18 +143,27 @@ fn validate_prepared_rotation(
Ok(())
}
fn validate_authorization_commit(
async fn validate_authorization_commit(
provider_kind: &ProviderKind,
context: &MutationContext,
prepared: &PreparedAuthorizationCommit,
prepared: PreparedAuthorizationCommit,
resource: &'static str,
) -> Result<(), AdminError> {
) -> Result<PreparedAuthorizationCommit, AdminError> {
if prepared.pending.provider_kind() != provider_kind
|| !prepared.pending.owner_binding().matches_context(context)
{
return Err(AdminError::conflict(format!(
let validation_error = AdminError::conflict(format!(
"{resource} pending authorization binding is invalid"
)));
));
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
let matches_target = match (prepared.pending.target(), &prepared.credential) {
(
@@ -171,25 +171,29 @@ fn validate_authorization_commit(
PreparedAuthorizationCredential::Create(credential),
) => credential.provider_kind == *provider_kind,
(
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
},
AuthorizationMutationTarget::Reauthorize { account_id },
PreparedAuthorizationCredential::Reauthorize(credential),
) => {
let facts = credential.facts();
facts.provider_kind == *provider_kind
&& facts.account_id == *account_id
&& facts.expected_credential_revision == *expected_credential_revision
facts.provider_kind == *provider_kind && facts.account_id == *account_id
}
_ => false,
};
if !matches_target {
return Err(AdminError::conflict(format!(
let validation_error = AdminError::conflict(format!(
"{resource} prepared credential does not match its pending target"
)));
));
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
Ok(())
Ok(prepared)
}
async fn commit_authorization(
@@ -198,16 +202,34 @@ async fn commit_authorization(
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (command, guard) = prepared.into_commit();
let crate::model::provider_credentials::AuthorizationCommitSettlement {
command,
credential_guard,
authorization_guard,
} = prepared.into_commit();
match accounts.commit_authorization(command, context).await {
Ok(result) => {
if let Some(guard) = guard {
if let Some(guard) = credential_guard {
guard.finish();
}
if let Some(guard) = authorization_guard {
guard.commit().await?;
}
Ok(result)
}
Err(error) => {
drop(guard);
drop(credential_guard);
if let Some(guard) = authorization_guard
&& let Err(settlement_error) = guard.abort().await
{
tracing::warn!(
resource,
store_error = %error,
settlement_error = %settlement_error,
"OAuth authorization claim release failed after Store commit failure"
);
return Err(settlement_error);
}
Err(map_store_error(error, resource))
}
}
@@ -215,20 +237,13 @@ async fn commit_authorization(
async fn commit_credential_rotation(
accounts: &dyn AccountStore,
expected_config_revision: crate::model::Revision,
prepared: PreparedCredentialRotation,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_rotation(
CredentialRotationCommit {
expected_config_revision,
prepared: facts,
},
context,
)
.commit_credential_rotation(CredentialRotationCommit { prepared: facts }, context)
.await
{
Ok(result) => {
@@ -244,20 +259,13 @@ async fn commit_credential_rotation(
async fn commit_credential_refresh(
accounts: &dyn AccountStore,
expected_config_revision: crate::model::Revision,
prepared: PreparedCredentialRotation,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_refresh(
CredentialRotationCommit {
expected_config_revision,
prepared: facts,
},
context,
)
.commit_credential_refresh(CredentialRotationCommit { prepared: facts }, context)
.await
{
Ok(result) => {
@@ -289,7 +297,6 @@ async fn set_credential_enabled(
let revision = accounts
.set_account_enabled(
SetAccountEnabled {
expected_config_revision: command.expected_config_revision,
account_id: account_id.as_str().to_owned(),
enabled,
},
@@ -320,7 +327,6 @@ async fn delete_credentials(
let revision = accounts
.delete_accounts(
DeleteAccounts {
expected_config_revision: command.expected_config_revision,
account_ids: account_ids
.iter()
.map(|account_id| account_id.as_str().to_owned())
@@ -13,7 +13,8 @@ use crate::{
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, ProviderQuotaRequest, RotateCredential, StartAuthorization,
PrepareCredentialRotation, PreparedAuthorizationCredential, ProviderQuotaRequest,
RotateCredential, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -83,6 +84,27 @@ impl DefaultOpenAiService {
snapshot,
}
}
async fn observe_initial_quotas(&self, account_ids: &[ProviderAccountId], request_id: &str) {
for account_id in account_ids {
if let Err(error) = self
.provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await
{
tracing::warn!(
request_id,
account_id = %account_id.as_str(),
quota_error = ?error.kind(),
"OpenAI initial quota observation failed"
);
}
}
}
}
#[async_trait]
@@ -112,7 +134,6 @@ impl OpenAiService for DefaultOpenAiService {
command: ImportCredentials,
) -> Result<CredentialImportResult, AdminError> {
let context = command.context;
let expected_config_revision = command.expected_config_revision;
let prepared = self
.provider
.prepare_import(PrepareCredentialImport {
@@ -127,27 +148,13 @@ impl OpenAiService for DefaultOpenAiService {
)?;
let result = self
.accounts
.commit_credential_import(
CredentialImportCommit {
expected_config_revision,
prepared,
},
&context,
)
.commit_credential_import(CredentialImportCommit { prepared }, &context)
.await
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
// 导入已经提交,额度属于可重建观察事实;单个账号查询失败由周期任务重试。
for account_id in &result.credential_ids {
let _ = self
.provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await;
}
// 导入完成后立即做一次观察;失败不回滚已经提交的 credential。
self.observe_initial_quotas(&result.credential_ids, &context.request_id)
.await;
Ok(result)
}
@@ -178,12 +185,17 @@ impl OpenAiService for DefaultOpenAiService {
.complete_authorization(command)
.await
.map_err(|error| map_provider_error(error, "OpenAI authorization"))?;
validate_authorization_commit(
let prepared = validate_authorization_commit(
self.provider.provider_kind(),
&context,
&prepared,
prepared,
"OpenAI authorization",
)?;
)
.await?;
let creates_account = matches!(
&prepared.credential,
PreparedAuthorizationCredential::Create(_)
);
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
@@ -192,6 +204,13 @@ impl OpenAiService for DefaultOpenAiService {
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
if creates_account {
self.observe_initial_quotas(
std::slice::from_ref(&result.account_id),
&context.request_id,
)
.await;
}
Ok(result)
}
@@ -200,7 +219,6 @@ impl OpenAiService for DefaultOpenAiService {
command: RotateCredential,
) -> Result<CredentialMutationResult, AdminError> {
let context = command.mutation.context;
let expected_config_revision = command.mutation.expected_config_revision;
let account_id = command.mutation.account_id;
let details = required_credential(
self.accounts.as_ref(),
@@ -209,17 +227,11 @@ impl OpenAiService for DefaultOpenAiService {
"OpenAI credential rotation",
)
.await?;
if details.credential.credential_revision != command.expected_credential_revision {
return Err(AdminError::conflict(
"OpenAI credential rotation revision is stale",
));
}
let account = details.credential;
let prepared = self
.provider
.prepare_rotation(PrepareCredentialRotation {
account: account.clone(),
expected_credential_revision: command.expected_credential_revision,
provider_material: command.provider_material,
})
.await
@@ -227,7 +239,6 @@ impl OpenAiService for DefaultOpenAiService {
validate_prepared_rotation(&account, &prepared, "OpenAI credential rotation")?;
let result = commit_credential_rotation(
self.accounts.as_ref(),
expected_config_revision,
prepared,
&context,
"OpenAI credential rotation",
@@ -86,13 +86,12 @@ impl SettingsService for DefaultSettingsService {
&self,
context: &MutationContext,
) -> Result<RegeneratedAdminApiKey, AdminError> {
let settings = self.load().await?;
let mut bytes = [0_u8; 32];
OsRng.fill_bytes(&mut bytes);
let key = AdminApiKey::new(format!("admin-{}", hex::encode(bytes)));
let mutation = self
.store
.replace_admin_api_key(settings.config_revision, key.clone(), context)
.replace_admin_api_key(key.clone(), context)
.await
.map_err(|error| map_store_error(error, "administrator API key"))?;
publish_committed(self.snapshot.as_ref(), mutation.config_revision).await?;
@@ -103,10 +102,9 @@ impl SettingsService for DefaultSettingsService {
&self,
context: &MutationContext,
) -> Result<AdminApiKeyMutation, AdminError> {
let settings = self.load().await?;
let mutation = self
.store
.delete_admin_api_key(settings.config_revision, context)
.delete_admin_api_key(context)
.await
.map_err(|error| map_store_error(error, "administrator API key"))?;
publish_committed(self.snapshot.as_ref(), mutation.config_revision).await?;
@@ -89,7 +89,6 @@ impl XaiService for DefaultXaiService {
command: ImportCredentials,
) -> Result<CredentialImportResult, AdminError> {
let context = command.context;
let expected_config_revision = command.expected_config_revision;
let prepared = self
.provider
.prepare_import(PrepareCredentialImport {
@@ -104,13 +103,7 @@ impl XaiService for DefaultXaiService {
)?;
let result = self
.accounts
.commit_credential_import(
CredentialImportCommit {
expected_config_revision,
prepared,
},
&context,
)
.commit_credential_import(CredentialImportCommit { prepared }, &context)
.await
.map_err(|error| map_store_error(error, "xAI credential import"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
@@ -155,12 +148,13 @@ impl XaiService for DefaultXaiService {
.complete_authorization(command)
.await
.map_err(|error| map_provider_error(error, "xAI authorization"))?;
validate_authorization_commit(
let prepared = validate_authorization_commit(
self.provider.provider_kind(),
&context,
&prepared,
prepared,
"xAI authorization",
)?;
)
.await?;
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
@@ -16,7 +16,7 @@ use gateway_core::{
use gateway_admin::{
AdminServices,
model::{
MutationContext, Revision,
AdminError, MutationContext, Revision,
accounts::{
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
@@ -24,12 +24,12 @@ use gateway_admin::{
},
observability::TimeRange,
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget,
AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit,
AuthorizationCommit, AuthorizationCommitGuard, AuthorizationCredentialCommit,
AuthorizationMutationTarget, AuthorizationStarted, CompleteAuthorization,
CredentialCommitGuard, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
@@ -58,11 +58,13 @@ pub(super) struct FakeProviderAdmin {
events: EventLog,
failure: Mutex<Option<ProviderAdminErrorKind>>,
quota_failure: Mutex<Option<ProviderAdminErrorKind>>,
pending: Mutex<Option<PendingAuthorizationMutation>>,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
retry_authorization_after_abort: Mutex<bool>,
export_inputs: Mutex<Vec<ProviderExportCredentialInput>>,
import_account_ids: Mutex<Vec<String>>,
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota: Mutex<ProviderQuota>,
current_credential_revision: Mutex<Revision>,
}
impl FakeProviderAdmin {
@@ -72,11 +74,13 @@ impl FakeProviderAdmin {
events,
failure: Mutex::new(None),
quota_failure: Mutex::new(None),
pending: Mutex::new(None),
pending: Arc::new(Mutex::new(None)),
retry_authorization_after_abort: Mutex::new(false),
export_inputs: Mutex::new(Vec::new()),
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
quota_requests: Mutex::new(Vec::new()),
quota: Mutex::new(empty_quota()),
current_credential_revision: Mutex::new(revision(1)),
})
}
@@ -113,6 +117,20 @@ impl FakeProviderAdmin {
self.pending.lock().expect("pending authorization").clone()
}
pub(super) fn retry_authorization_after_abort(&self) {
*self
.retry_authorization_after_abort
.lock()
.expect("authorization retry") = true;
}
pub(super) fn set_current_credential_revision(&self, revision: Revision) {
*self
.current_credential_revision
.lock()
.expect("current credential revision") = revision;
}
fn export_inputs(&self) -> Vec<ProviderExportCredentialInput> {
self.export_inputs.lock().expect("export inputs").clone()
}
@@ -133,7 +151,10 @@ impl FakeProviderAdmin {
PreparedCredentialRotationFacts {
account_id: ProviderAccountId::new(account.id.clone()).expect("account ID"),
provider_kind: account.provider_kind.clone(),
expected_credential_revision: account.credential_revision,
expected_credential_revision: *self
.current_credential_revision
.lock()
.expect("current credential revision"),
name: account.name.clone(),
email: account.email.clone(),
plan_type: account.plan_type.clone(),
@@ -246,23 +267,31 @@ impl ProviderAdmin for FakeProviderAdmin {
if !pending.owner_binding().matches_context(&command.context) {
return Err(ProviderAdminError::new(ProviderAdminErrorKind::NotFound));
}
let retry_pending = (*self
.retry_authorization_after_abort
.lock()
.expect("authorization retry"))
.then(|| pending.clone());
let credential = match pending.target() {
AuthorizationMutationTarget::Create { name } => {
PreparedAuthorizationCredential::Create(prepared_create(self.kind.clone(), name))
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
let mut account = account_record(self.kind.as_str());
account.id = account_id.as_str().to_owned();
account.credential_revision = *expected_credential_revision;
PreparedAuthorizationCredential::Reauthorize(self.prepared_rotation(&account))
}
};
Ok(PreparedAuthorizationCommit {
pending,
credential,
let prepared = PreparedAuthorizationCommit::new(pending, credential);
Ok(match retry_pending {
Some(pending) => {
prepared.with_authorization_guard(Box::new(RetryableAuthorizationGuard::new(
self.events.clone(),
Arc::clone(&self.pending),
pending,
)))
}
None => prepared,
})
}
@@ -342,18 +371,20 @@ pub(super) struct FakeAccountStore {
account: AccountRecord,
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
authorization_revision: Mutex<Option<Revision>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
}
impl FakeAccountStore {
pub(super) fn new(kind: &str, events: EventLog) -> Arc<Self> {
Self::with_account(account_record(kind), events)
}
fn with_account(account: AccountRecord, events: EventLog) -> Arc<Self> {
Arc::new(Self {
events,
account: account_record(kind),
account,
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
authorization_revision: Mutex::new(None),
quota_window_usage: Mutex::new(Vec::new()),
})
}
@@ -366,13 +397,6 @@ impl FakeAccountStore {
self.audit_requests.lock().expect("audit requests").clone()
}
pub(super) fn authorization_revision(&self) -> Option<Revision> {
*self
.authorization_revision
.lock()
.expect("authorization revision")
}
pub(super) fn set_quota_window_usage(&self, usage: Vec<AccountUsageWindowResult>) {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
@@ -519,10 +543,6 @@ impl AccountStore for FakeAccountStore {
) -> AdminStoreResult<CredentialMutationResult> {
self.record("store.commit_authorization");
self.record_context(context);
*self
.authorization_revision
.lock()
.expect("authorization revision") = Some(command.pending.expected_config_revision());
self.require_commit()?;
let (account_id, credential_revision) = match command.credential {
AuthorizationCredentialCommit::Create(credential) => (credential.account_id, None),
@@ -627,7 +647,6 @@ impl SettingsStore for StaticSettingsStore {
async fn replace_admin_api_key(
&self,
_: Revision,
_: AdminApiKey,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
@@ -636,7 +655,6 @@ impl SettingsStore for StaticSettingsStore {
async fn delete_admin_api_key(
&self,
_: Revision,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
Err(store_unavailable())
@@ -675,6 +693,51 @@ impl Drop for RecordingGuard {
}
}
struct RetryableAuthorizationGuard {
events: EventLog,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
authorization: PendingAuthorizationMutation,
}
impl RetryableAuthorizationGuard {
fn new(
events: EventLog,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
authorization: PendingAuthorizationMutation,
) -> Self {
Self {
events,
pending,
authorization,
}
}
}
#[async_trait]
impl AuthorizationCommitGuard for RetryableAuthorizationGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
self.events
.lock()
.expect("authorization guard events")
.push("authorization_guard.commit");
Ok(())
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
let Self {
events,
pending,
authorization,
} = *self;
*pending.lock().expect("pending authorization") = Some(authorization);
events
.lock()
.expect("authorization guard events")
.push("authorization_guard.abort");
Ok(())
}
}
#[test]
fn provider_registry_should_resolve_custom_kind_without_central_match() {
let provider = FakeProviderAdmin::new("custom-provider", events());
@@ -701,9 +764,11 @@ fn provider_registry_should_reject_duplicate_kind() {
}
#[tokio::test]
async fn connection_test_should_preserve_safe_provider_error_details() {
async fn connection_test_should_probe_unavailable_account() {
let provider = FakeProviderAdmin::new("xai", events());
let store = FakeAccountStore::new("xai", events());
let mut account = account_record("xai");
account.availability = AccountAvailability::QuotaExhausted;
let store = FakeAccountStore::with_account(account, events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
@@ -724,7 +789,7 @@ async fn connection_test_should_preserve_safe_provider_error_details() {
message,
provider_error_code: Some(code),
provider_error_type: Some(error_type),
..
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
}) if message == "included usage exhausted"
&& code == "usage_exhausted"
&& error_type == "invalid_request_error"
@@ -769,7 +834,6 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
.accounts()
.refresh(
&context("refresh-request"),
revision(1),
ProviderAccountId::new("acct_test").expect("account ID"),
)
.await
@@ -882,7 +946,6 @@ async fn accounts_refresh_provider_failure_should_not_call_store_commit() {
.accounts()
.refresh(
&context("refresh-provider-error"),
revision(1),
ProviderAccountId::new("acct_test").expect("account ID"),
)
.await
@@ -895,7 +958,7 @@ async fn accounts_refresh_provider_failure_should_not_call_store_commit() {
}
#[tokio::test]
async fn accounts_refresh_store_failure_should_drop_guard_after_cas_attempt() {
async fn accounts_refresh_store_failure_should_drop_guard_after_commit_attempt() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
@@ -906,11 +969,10 @@ async fn accounts_refresh_store_failure_should_drop_guard_after_cas_attempt() {
.accounts()
.refresh(
&context("refresh-store-error"),
revision(1),
ProviderAccountId::new("acct_test").expect("account ID"),
)
.await
.expect_err("Store CAS must fail");
.expect_err("Store commit must fail");
assert_eq!(
recorded(&events),
@@ -1039,7 +1101,7 @@ fn prepared_create_with_id(
fn rotation_result(command: CredentialRotationCommit) -> CredentialMutationResult {
CredentialMutationResult {
config_revision: revision(command.expected_config_revision.get() + 1),
config_revision: revision(2),
account_id: command.prepared.account_id,
credential_revision: Some(revision(
command.prepared.expected_credential_revision.get() + 1,
@@ -449,7 +449,6 @@ impl SettingsStore for UnavailableStore {
async fn replace_admin_api_key(
&self,
_: Revision,
_: AdminApiKey,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
@@ -458,7 +457,6 @@ impl SettingsStore for UnavailableStore {
async fn delete_admin_api_key(
&self,
_: Revision,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
Err(unavailable("admin API key"))
@@ -518,7 +518,6 @@ impl SettingsStore for FixtureSettingsStore {
async fn replace_admin_api_key(
&self,
_: Revision,
_: AdminApiKey,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
@@ -527,7 +526,6 @@ impl SettingsStore for FixtureSettingsStore {
async fn delete_admin_api_key(
&self,
_: Revision,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
Err(super::unavailable("admin API key"))
@@ -6,7 +6,7 @@ use gateway_admin::{
AdminServices,
model::provider_credentials::{
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
ImportCredentials, ProviderQuotaRequest, ReauthorizationTarget, StartAuthorization,
ImportCredentials, ProviderQuotaRequest, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -114,7 +114,6 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
.openai()
.import_document(ImportCredentials {
context: context("import-openai"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -151,7 +150,6 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
.openai()
.import_document(ImportCredentials {
context: context("import-openai-batch"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -193,7 +191,6 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
.openai()
.import_document(ImportCredentials {
context: context("import-openai-quota-failure"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -206,6 +203,133 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
assert_eq!(provider.quota_requests().len(), 1);
}
#[tokio::test]
async fn openai_authorization_create_should_observe_initial_quota() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider.clone(), store.clone()).await;
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-start-openai-create"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-start-openai-create"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
})
.await
.expect("complete authorization");
assert_eq!(
recorded(&events),
[
"provider.start_authorization",
"provider.complete_authorization",
"store.commit_authorization",
"provider.quota",
]
);
assert_eq!(
provider.quota_requests(),
[ProviderQuotaRequest {
account_id: ProviderAccountId::new("acct_prepared").expect("account ID"),
refresh: true,
rolling_usage: None,
}]
);
assert_eq!(store.audit_requests(), ["oauth-start-openai-create"]);
}
#[tokio::test]
async fn openai_authorization_create_should_remain_successful_when_initial_quota_fails() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider.fail_next_quota(ProviderAdminErrorKind::Unavailable);
let store = FakeAccountStore::new("openai", events);
let services = service(provider.clone(), store).await;
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-start-openai-quota-failure"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-start-openai-quota-failure"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
})
.await
.expect("quota observation must not fail authorization");
assert_eq!(result.account_id.as_str(), "acct_prepared");
assert_eq!(provider.quota_requests().len(), 1);
}
#[tokio::test]
async fn openai_authorization_store_failure_should_release_claim_for_retry() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider.retry_authorization_after_abort();
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = service(provider.clone(), store).await;
let command = || CompleteAuthorization {
context: context("oauth-store-retry"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
};
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-store-retry"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
services
.openai()
.complete_authorization(command())
.await
.expect_err("Store commit must fail");
services
.openai()
.complete_authorization(command())
.await
.expect("released claim must permit retry");
assert_eq!(
recorded(&events),
[
"provider.start_authorization",
"provider.complete_authorization",
"store.commit_authorization",
"authorization_guard.abort",
"provider.complete_authorization",
"store.commit_authorization",
"authorization_guard.commit",
"provider.quota",
]
);
assert!(provider.pending().is_none());
}
#[tokio::test]
async fn openai_import_provider_error_should_not_touch_store_transaction() {
let events = events();
@@ -218,7 +342,6 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
.openai()
.import_document(ImportCredentials {
context: context("import-openai-error"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -228,7 +351,7 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
}
#[tokio::test]
async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_through_commit() {
async fn openai_reauthorization_should_commit_after_credential_revision_advances() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
@@ -237,18 +360,13 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-start-openai"),
expected_config_revision: revision(7),
name: "reauthorize".to_owned(),
reauthorization: Some(ReauthorizationTarget {
account_id: ProviderAccountId::new("acct_test").expect("account ID"),
credential_revision: revision(1),
}),
reauthorization: Some(ProviderAccountId::new("acct_test").expect("account ID")),
})
.await
.expect("start reauthorization");
let pending = provider.pending().expect("pending envelope");
assert_eq!(pending.expected_config_revision(), revision(7));
assert_eq!(pending.provider_kind().as_str(), "openai");
assert_eq!(
pending.owner_binding().started_request_id(),
@@ -256,15 +374,14 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
);
assert!(matches!(
pending.target(),
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} if account_id.as_str() == "acct_test"
&& *expected_credential_revision == revision(1)
AuthorizationMutationTarget::Reauthorize { account_id }
if account_id.as_str() == "acct_test"
));
assert!(!format!("{pending:?}").contains("admin-test"));
services
provider.set_current_credential_revision(revision(2));
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-complete-openai"),
@@ -274,6 +391,8 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
.await
.expect("complete reauthorization");
assert_eq!(result.credential_revision, Some(revision(3)));
assert_eq!(
recorded(&events),
[
@@ -284,7 +403,6 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
"guard.finish",
]
);
assert_eq!(store.authorization_revision(), Some(revision(7)));
assert_eq!(store.audit_requests(), ["oauth-complete-openai"]);
}
@@ -299,7 +417,6 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-openai"),
expected_config_revision: revision(1),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
@@ -307,7 +424,6 @@ fn mutation(account_id: &str) -> CredentialMutation {
fn deletion(account_id: &str) -> CredentialDeletion {
CredentialDeletion {
context: context("request-openai"),
expected_config_revision: revision(1),
account_ids: vec![ProviderAccountId::new(account_id).expect("account ID")],
}
}
@@ -2,7 +2,7 @@ use async_trait::async_trait;
use gateway_admin::{
model::{
AdminErrorKind, MutationContext, Revision,
AdminErrorKind, MutationContext,
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
RuntimeSettings,
@@ -33,7 +33,6 @@ impl SettingsStore for UnusedSettingsStore {
async fn replace_admin_api_key(
&self,
_: Revision,
_: AdminApiKey,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
@@ -42,7 +41,6 @@ impl SettingsStore for UnusedSettingsStore {
async fn delete_admin_api_key(
&self,
_: Revision,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
Err(unused())
@@ -63,7 +61,6 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
request_id: "request-settings".to_owned(),
},
ReplaceRuntimeSettings {
expected_config_revision: Revision::new(1).expect("revision"),
model_mappings: Default::default(),
refresh_margin_seconds: 0,
refresh_concurrency: 1,
@@ -8,9 +8,7 @@ use gateway_admin::{
ports::provider::ProviderAdminErrorKind,
};
use super::accounts::{
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision,
};
use super::accounts::{FakeAccountStore, FakeProviderAdmin, context, document, events, recorded};
#[tokio::test]
async fn xai_disable_should_commit_then_notify_the_stateless_provider() {
@@ -47,7 +45,6 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
.xai()
.import_document(ImportCredentials {
context: context("import-xai"),
expected_config_revision: revision(3),
document: document(),
})
.await
@@ -84,7 +81,6 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
.xai()
.import_document(ImportCredentials {
context: context("import-xai-batch"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -119,7 +115,6 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
.xai()
.import_document(ImportCredentials {
context: context("import-xai-quota-failure"),
expected_config_revision: revision(1),
document: document(),
})
.await
@@ -140,7 +135,6 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-xai"),
expected_config_revision: revision(1),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
+4 -2
View File
@@ -13,7 +13,7 @@ path = "tests/main.rs"
[dependencies]
axum = { version = "0.8.9", default-features = false, features = ["json", "query", "ws"] }
base64 = "0.22.1"
base64 = "0.23.0"
bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde"] }
futures = "0.3.32"
@@ -31,7 +31,9 @@ uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
async-trait = "0.1.89"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] }
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread"] }
tokio-tungstenite = "0.29.0"
tower = { version = "0.5.3", features = ["util"] }
[lints]
+23 -127
View File
@@ -16,18 +16,18 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, FixedOffset, Utc};
use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision,
AdminError as AdminServiceError, AdminErrorKind, PageSize,
accounts::{
AccountAvailability, AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost,
AccountListQuery, AccountModelUsage, AccountSort, AccountSortField,
AccountStatus as DomainAccountStatus, AccountUsage, SortDirection,
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
AccountUsage, SortDirection,
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow,
ReauthorizationTarget, RotateCredential, StartAuthorization,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, RotateCredential,
StartAuthorization,
},
};
use gateway_core::{
@@ -161,7 +161,6 @@ fn parse_sort_direction(value: &str) -> Option<SortDirection> {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountPageData {
pub config_revision: u64,
pub items: Vec<AccountView>,
pub page: PageMeta,
pub summary: AccountSummaryView,
@@ -197,8 +196,6 @@ pub struct AccountView {
pub token_refreshing: bool,
pub availability: String,
pub enabled: bool,
pub credential_revision: u64,
pub state_revision: Option<u64>,
pub access_token_expires_at: Option<String>,
pub access_token_expires_at_display: Option<String>,
pub refresh_token_expires_at: Option<String>,
@@ -391,31 +388,21 @@ impl AccountActionRequest {
}
}
/// 手工 OAuth 刷新会变更持久 credential,因此必须携带配置 CAS revision。
/// 手工 OAuth 刷新会变更持久 credential。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountRefreshRequest {
pub account_id: String,
pub expected_config_revision: u64,
}
impl AccountRefreshRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.account_id, "accountId")?;
if self.expected_config_revision == 0 {
return Err(WireValidationError::new("expectedConfigRevision"));
}
Ok(())
require_account_id(&self.account_id, "accountId")
}
fn into_command(self) -> Result<(Revision, ProviderAccountId), WireValidationError> {
fn into_command(self) -> Result<ProviderAccountId, WireValidationError> {
self.validate()?;
Ok((
Revision::new(self.expected_config_revision)
.map_err(|_| WireValidationError::new("expectedConfigRevision"))?,
ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
))
ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId"))
}
}
@@ -460,7 +447,6 @@ pub struct AccountModelsData {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountRefreshData {
pub config_revision: u64,
pub account: AccountView,
}
@@ -579,14 +565,12 @@ impl AccountProvider {
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountImportRequest {
pub provider: String,
pub expected_config_revision: u64,
pub data: Value,
}
impl AccountImportRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
if !self.data.is_object()
|| serde_json::to_vec(&self.data)
.map_or(true, |encoded| encoded.len() > MAX_IMPORT_DATA_BYTES)
@@ -606,10 +590,6 @@ impl AccountImportRequest {
provider,
ImportCredentials {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
document: provider_document(self.data, "data")?,
},
))
@@ -620,28 +600,18 @@ impl AccountImportRequest {
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct StartAccountAuthorizationRequest {
pub provider: String,
pub expected_config_revision: u64,
pub name: String,
pub account_id: Option<String>,
pub expected_credential_revision: Option<u64>,
}
impl StartAccountAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_text(&self.name, MAX_NAME_BYTES, "name")?;
match (
self.account_id.as_deref(),
self.expected_credential_revision,
) {
(None, None) => Ok(()),
(Some(account_id), Some(credential_revision)) => {
require_account_id(account_id, "accountId")?;
require_positive(credential_revision, "expectedCredentialRevision")
}
_ => Err(WireValidationError::new("reauthorization")),
if let Some(account_id) = self.account_id.as_deref() {
require_account_id(account_id, "accountId")?;
}
Ok(())
}
fn into_command(
@@ -650,23 +620,15 @@ impl StartAccountAuthorizationRequest {
) -> Result<(AccountProvider, StartAuthorization), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
let reauthorization = match (self.account_id, self.expected_credential_revision) {
(None, None) => None,
(Some(account_id), Some(credential_revision)) => Some(ReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
credential_revision: revision(credential_revision, "expectedCredentialRevision")?,
}),
_ => return Err(WireValidationError::new("reauthorization")),
};
let reauthorization = self
.account_id
.map(ProviderAccountId::new)
.transpose()
.map_err(|_| WireValidationError::new("accountId"))?;
Ok((
provider,
StartAuthorization {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
name: self.name,
reauthorization,
},
@@ -721,14 +683,12 @@ impl CompleteAccountAuthorizationRequest {
pub struct AccountMutationRequest {
pub provider: String,
pub account_id: String,
pub expected_config_revision: u64,
}
impl AccountMutationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_account_id(&self.account_id, "accountId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")
require_account_id(&self.account_id, "accountId")
}
fn into_command(
@@ -741,10 +701,6 @@ impl AccountMutationRequest {
provider,
CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
@@ -757,7 +713,6 @@ impl AccountMutationRequest {
pub struct AccountDeletionRequest {
pub provider: String,
pub account_ids: Vec<String>,
pub expected_config_revision: u64,
}
impl AccountDeletionRequest {
@@ -773,7 +728,7 @@ impl AccountDeletionRequest {
return Err(WireValidationError::new("accountIds"));
}
}
require_positive(self.expected_config_revision, "expectedConfigRevision")
Ok(())
}
fn into_command(
@@ -786,10 +741,6 @@ impl AccountDeletionRequest {
provider,
CredentialDeletion {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_ids: self
.account_ids
.into_iter()
@@ -806,8 +757,6 @@ impl AccountDeletionRequest {
pub struct RotateAccountRequest {
pub provider: String,
pub account_id: String,
pub expected_config_revision: u64,
pub expected_credential_revision: u64,
pub access_token: String,
pub refresh_token: Option<String>,
}
@@ -818,11 +767,6 @@ impl RotateAccountRequest {
return Err(WireValidationError::new("provider"));
}
require_account_id(&self.account_id, "accountId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_positive(
self.expected_credential_revision,
"expectedCredentialRevision",
)?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
}
@@ -840,17 +784,9 @@ impl RotateAccountRequest {
Ok(RotateCredential {
mutation: CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
expected_credential_revision: revision(
self.expected_credential_revision,
"expectedCredentialRevision",
)?,
provider_material: ProviderDocument::new(OpaqueProviderData::new(material)),
})
}
@@ -859,7 +795,6 @@ impl RotateAccountRequest {
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountImportData {
pub config_revision: u64,
pub imported_count: usize,
pub account_ids: Vec<String>,
}
@@ -872,7 +807,6 @@ impl From<CredentialImportResult> for AccountImportData {
.map(|account_id| account_id.to_string())
.collect::<Vec<_>>();
Self {
config_revision: result.config_revision.get(),
imported_count: account_ids.len(),
account_ids,
}
@@ -900,18 +834,13 @@ impl From<AuthorizationStarted> for AccountAuthorizationData {
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountMutationData {
pub config_revision: u64,
pub account_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub credential_revision: Option<u64>,
}
impl From<CredentialMutationResult> for AccountMutationData {
fn from(result: CredentialMutationResult) -> Self {
Self {
config_revision: result.config_revision.get(),
account_id: result.account_id.to_string(),
credential_revision: result.credential_revision.map(Revision::get),
}
}
}
@@ -919,7 +848,6 @@ impl From<CredentialMutationResult> for AccountMutationData {
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountDeletionData {
pub config_revision: u64,
pub deleted_count: usize,
pub account_ids: Vec<String>,
}
@@ -932,7 +860,6 @@ impl From<CredentialDeletionResult> for AccountDeletionData {
.map(|account_id| account_id.to_string())
.collect::<Vec<_>>();
Self {
config_revision: result.config_revision.get(),
deleted_count: account_ids.len(),
account_ids,
}
@@ -1264,15 +1191,11 @@ async fn refresh_account<S>(
where
S: AdminSessionState + Send + Sync,
{
let (expected_config_revision, account_id) = request.into_command().map_err(map_wire_error)?;
let account_id = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.refresh(
&auth.context().mutation_context(),
expected_config_revision,
account_id,
)
.refresh(&auth.context().mutation_context(), account_id)
.await
.map_err(map_service_error)?;
let data = account_refresh_data(result, Utc::now());
@@ -1394,7 +1317,6 @@ fn account_page_data(
u32::try_from(result.total.div_ceil(u64::from(page_size))).unwrap_or(u32::MAX)
};
AccountPageData {
config_revision: result.config_revision.get(),
items: result
.items
.into_iter()
@@ -1412,7 +1334,6 @@ fn account_page_data(
fn account_refresh_data(result: AccountRefreshResult, now: DateTime<Utc>) -> AccountRefreshData {
AccountRefreshData {
config_revision: result.config_revision.get(),
account: account_view(result.account, now),
}
}
@@ -1460,10 +1381,8 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
status: status.clone(),
display_status: status,
token_refreshing: false,
availability: account_availability_name(account.availability).to_owned(),
availability: account.availability.as_str().to_owned(),
enabled: account.enabled,
credential_revision: account.credential_revision.get(),
state_revision: None,
access_token_expires_at: expires_at,
access_token_expires_at_display: account
.access_token_expires_at
@@ -1714,10 +1633,6 @@ fn empty_account_usage() -> AccountUsageView {
}
}
fn revision(value: u64, field: &'static str) -> Result<Revision, WireValidationError> {
Revision::new(value).map_err(|_| WireValidationError::new(field))
}
fn provider_document(
value: Value,
field: &'static str,
@@ -1751,13 +1666,6 @@ fn require_text(
Ok(())
}
fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
@@ -1805,18 +1713,6 @@ fn account_status_name(status: DomainAccountStatus) -> &'static str {
}
}
fn account_availability_name(availability: AccountAvailability) -> &'static str {
match availability {
AccountAvailability::Unknown => "unknown",
AccountAvailability::Ready => "ready",
AccountAvailability::Cooldown => "cooldown",
AccountAvailability::QuotaExhausted => "quota_exhausted",
AccountAvailability::Expired => "expired",
AccountAvailability::Banned => "banned",
AccountAvailability::Invalid => "invalid",
}
}
fn format_number(value: u64) -> String {
let text = value.to_string();
let mut output = String::with_capacity(text.len() + text.len() / 3);
@@ -4,14 +4,11 @@ use std::fmt;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
Revision,
client_keys::{
ClientKeyCursor, ClientKeyCursorValue as DomainCursorValue, ClientKeyListQuery,
ClientKeyMutation, ClientKeyPage, ClientKeyPageSize, ClientKeyRecord, ClientKeySecret,
ClientKeySort as DomainSort, ClientKeySortField as DomainSortField, CreateClientKey,
CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey,
},
use gateway_admin::model::client_keys::{
ClientKeyCursor, ClientKeyCursorValue as DomainCursorValue, ClientKeyListQuery,
ClientKeyMutation, ClientKeyPage, ClientKeyPageSize, ClientKeyRecord, ClientKeySecret,
ClientKeySort as DomainSort, ClientKeySortField as DomainSortField, CreateClientKey,
CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey,
};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
@@ -138,7 +135,6 @@ impl ClientKeySort {
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CreateClientKeyRequest {
expected_config_revision: u64,
name: String,
label: Option<String>,
provider_kind: String,
@@ -149,14 +145,12 @@ pub struct CreateClientKeyRequest {
impl CreateClientKeyRequest {
/// 校验 wire 边界并直接构造管理用例命令。
pub fn into_command(self) -> Result<CreateClientKey, WireValidationError> {
validate_revision(self.expected_config_revision)?;
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(CreateClientKey {
expected_config_revision: revision(self.expected_config_revision)?,
name: self.name,
label: self.label,
provider_kind: ProviderKind::new(self.provider_kind)
@@ -174,7 +168,6 @@ impl CreateClientKeyRequest {
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateClientKeyRequest {
id: String,
expected_config_revision: u64,
name: String,
label: Option<String>,
provider_kind: String,
@@ -186,14 +179,12 @@ impl UpdateClientKeyRequest {
/// 校验 wire 边界并直接构造管理用例命令。
pub fn into_command(self) -> Result<UpdateClientKey, WireValidationError> {
validate_required_text(&self.id, "id")?;
validate_revision(self.expected_config_revision)?;
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(UpdateClientKey {
expected_config_revision: revision(self.expected_config_revision)?,
id: client_key_id(self.id, "clientKeyMutationNotFound")?,
name: self.name,
label: self.label,
@@ -207,12 +198,11 @@ impl UpdateClientKeyRequest {
}
}
/// 只携带 ID 与配置 revision 的 Client Key mutation 请求。
/// 只携带 ID 的 Client Key mutation 请求。
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ClientKeyRevisionRequest {
pub struct ClientKeyMutationRequest {
id: String,
expected_config_revision: u64,
}
/// 读取一次完整 Key 的 ID query。
@@ -233,20 +223,15 @@ impl ClientKeyIdQuery {
}
}
impl ClientKeyRevisionRequest {
/// 校验请求并取出 ID 与 revision。
pub fn into_parts(self) -> Result<(String, u64), WireValidationError> {
impl ClientKeyMutationRequest {
/// 校验请求并取出 ID。
pub fn into_id(self) -> Result<String, WireValidationError> {
validate_required_text(&self.id, "id")?;
validate_revision(self.expected_config_revision)?;
Ok((self.id, self.expected_config_revision))
Ok(self.id)
}
fn into_command_parts(self) -> Result<(ClientApiKeyId, Revision), WireValidationError> {
let (id, expected_config_revision) = self.into_parts()?;
Ok((
client_key_id(id, "clientKeyMutationNotFound")?,
revision(expected_config_revision)?,
))
fn into_domain_id(self) -> Result<ClientApiKeyId, WireValidationError> {
client_key_id(self.into_id()?, "clientKeyMutationNotFound")
}
}
@@ -289,7 +274,6 @@ impl From<ClientKeyRecord> for ClientKeyView {
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientKeyListData {
config_revision: u64,
items: Vec<ClientKeyView>,
next_cursor: Option<String>,
total: u64,
@@ -298,14 +282,8 @@ pub struct ClientKeyListData {
impl ClientKeyListData {
/// 构造 Client Key 列表响应。
#[must_use]
pub fn new(
config_revision: u64,
items: Vec<ClientKeyView>,
next_cursor: Option<String>,
total: u64,
) -> Self {
pub fn new(items: Vec<ClientKeyView>, next_cursor: Option<String>, total: u64) -> Self {
Self {
config_revision,
items,
next_cursor,
total,
@@ -325,7 +303,6 @@ impl TryFrom<ClientKeyPage> for ClientKeyListData {
.map(encode_client_key_cursor)
.transpose()?;
Ok(Self::new(
page.config_revision.get(),
page.items.into_iter().map(Into::into).collect(),
next_cursor,
page.total,
@@ -337,7 +314,6 @@ impl TryFrom<ClientKeyPage> for ClientKeyListData {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CreatedClientKeyData {
config_revision: u64,
id: String,
prefix: String,
plaintext_key: String,
@@ -346,9 +322,8 @@ pub struct CreatedClientKeyData {
impl CreatedClientKeyData {
/// 构造一次性创建响应。
#[must_use]
pub fn new(config_revision: u64, id: String, prefix: String, plaintext_key: String) -> Self {
pub fn new(id: String, prefix: String, plaintext_key: String) -> Self {
Self {
config_revision,
id,
prefix,
plaintext_key,
@@ -359,7 +334,6 @@ impl CreatedClientKeyData {
impl From<CreatedClientKey> for CreatedClientKeyData {
fn from(created: CreatedClientKey) -> Self {
Self::new(
created.config_revision.get(),
created.secret.record.id.to_string(),
created.secret.record.prefix.clone(),
created.secret.expose_for_response().to_owned(),
@@ -371,7 +345,6 @@ impl fmt::Debug for CreatedClientKeyData {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CreatedClientKeyData")
.field("config_revision", &self.config_revision)
.field("id", &self.id)
.field("prefix", &self.prefix)
.field("plaintext_key", &"[REDACTED]")
@@ -417,24 +390,20 @@ impl fmt::Debug for RevealedClientKeyData {
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MutatedClientKeyData {
config_revision: u64,
id: String,
}
impl MutatedClientKeyData {
/// 构造 mutation 响应。
#[must_use]
pub fn new(config_revision: u64, id: String) -> Self {
Self {
config_revision,
id,
}
pub fn new(id: String) -> Self {
Self { id }
}
}
impl From<ClientKeyMutation> for MutatedClientKeyData {
fn from(mutation: ClientKeyMutation) -> Self {
Self::new(mutation.config_revision.get(), mutation.id.to_string())
Self::new(mutation.id.to_string())
}
}
@@ -582,18 +551,6 @@ fn wire_cursor(cursor: ClientKeyCursor) -> Result<ClientKeyCursorData, WireValid
Ok(cursor)
}
fn validate_revision(value: u64) -> Result<(), WireValidationError> {
if value == 0 {
return Err(WireValidationError::new("expectedConfigRevision"));
}
Ok(())
}
fn revision(value: u64) -> Result<Revision, WireValidationError> {
validate_revision(value)?;
Revision::new(value).map_err(|_| WireValidationError::new("expectedConfigRevision"))
}
fn client_key_id(
value: String,
field: &'static str,
@@ -768,23 +725,19 @@ where
async fn disable_client_key<S>(
auth: AdminAuth,
State(state): State<S>,
Json(payload): Json<ClientKeyRevisionRequest>,
Json(payload): Json<ClientKeyMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?;
let id = payload.into_domain_id().map_err(map_wire_error)?;
mutation_response(
state
.admin_services()
.client_keys()
.set_enabled(
&auth.context().mutation_context(),
SetClientKeyEnabled {
expected_config_revision,
id,
enabled: false,
},
SetClientKeyEnabled { id, enabled: false },
)
.await,
)
@@ -793,23 +746,19 @@ where
async fn enable_client_key<S>(
auth: AdminAuth,
State(state): State<S>,
Json(payload): Json<ClientKeyRevisionRequest>,
Json(payload): Json<ClientKeyMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?;
let id = payload.into_domain_id().map_err(map_wire_error)?;
mutation_response(
state
.admin_services()
.client_keys()
.set_enabled(
&auth.context().mutation_context(),
SetClientKeyEnabled {
expected_config_revision,
id,
enabled: true,
},
SetClientKeyEnabled { id, enabled: true },
)
.await,
)
@@ -818,23 +767,17 @@ where
async fn delete_client_key<S>(
auth: AdminAuth,
State(state): State<S>,
Json(payload): Json<ClientKeyRevisionRequest>,
Json(payload): Json<ClientKeyMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?;
let id = payload.into_domain_id().map_err(map_wire_error)?;
mutation_response(
state
.admin_services()
.client_keys()
.delete(
&auth.context().mutation_context(),
DeleteClientKey {
expected_config_revision,
id,
},
)
.delete(&auth.context().mutation_context(), DeleteClientKey { id })
.await,
)
}
@@ -854,9 +797,6 @@ fn map_wire_error(error: WireValidationError) -> AdminError {
"cursor" => AdminError::bad_request("Invalid client key cursor"),
"clientKeyRevealNotFound" => AdminError::not_found("Client API key was not found"),
"clientKeyMutationNotFound" => AdminError::not_found("client API key was not found"),
"expectedConfigRevision" => {
AdminError::bad_request("expectedConfigRevision must be positive")
}
_ => AdminError::bad_request("Invalid client API key request"),
}
}
@@ -10,12 +10,8 @@ use axum::{
routing::{get, post},
};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
Revision,
settings::{
ModelMappings as DomainModelMappings, ReplaceRuntimeSettings, RotationStrategy,
RuntimeSettings,
},
use gateway_admin::model::settings::{
ModelMappings as DomainModelMappings, ReplaceRuntimeSettings, RotationStrategy, RuntimeSettings,
};
use gateway_core::routing::{PublicModelId, UpstreamModelId};
use serde::{Deserialize, Serialize};
@@ -32,7 +28,6 @@ pub type ModelMappings = BTreeMap<String, String>;
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeSettingsView {
pub config_revision: u64,
pub model_mappings: ModelMappings,
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u64,
@@ -49,7 +44,6 @@ pub struct RuntimeSettingsView {
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateRuntimeSettingsRequest {
pub expected_config_revision: u64,
pub model_mappings: ModelMappings,
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u64,
@@ -64,9 +58,6 @@ pub struct UpdateRuntimeSettingsRequest {
impl UpdateRuntimeSettingsRequest {
/// 校验公共运行参数。
pub fn validate(&self) -> Result<(), WireValidationError> {
if self.expected_config_revision == 0 {
return Err(WireValidationError::new("expectedConfigRevision"));
}
validate_model_mappings(&self.model_mappings)?;
for (value, field) in [
(self.refresh_margin_seconds, "refreshMarginSeconds"),
@@ -93,8 +84,6 @@ impl UpdateRuntimeSettingsRequest {
fn into_command(self) -> Result<ReplaceRuntimeSettings, WireValidationError> {
self.validate()?;
Ok(ReplaceRuntimeSettings {
expected_config_revision: Revision::new(self.expected_config_revision)
.map_err(|_| WireValidationError::new("settingsRevisionOverflow"))?,
model_mappings: domain_model_mappings(self.model_mappings)?,
refresh_margin_seconds: self.refresh_margin_seconds,
refresh_concurrency: u32::try_from(self.refresh_concurrency)
@@ -117,7 +106,6 @@ impl UpdateRuntimeSettingsRequest {
impl From<RuntimeSettings> for RuntimeSettingsView {
fn from(settings: RuntimeSettings) -> Self {
Self {
config_revision: settings.config_revision.get(),
model_mappings: wire_model_mappings(settings.model_mappings),
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: u64::from(settings.refresh_concurrency),
@@ -358,7 +346,6 @@ fn valid_model_name(value: &str, max_len: usize) -> bool {
fn map_wire_error(error: WireValidationError) -> AdminError {
let message = match error.field() {
"settingsRevisionOverflow" => "expectedConfigRevision must be positive".to_owned(),
"settingsRefreshConcurrencyOverflow" => "Invalid refreshConcurrency".to_owned(),
"settingsMaxConcurrencyOverflow" => "Invalid maxConcurrentPerAccount".to_owned(),
"settingsUsageRetentionOverflow" => "Invalid usageRetentionDays".to_owned(),
@@ -75,10 +75,14 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
if error.kind() == ProviderErrorKind::Unavailable
&& error.send_state() == UpstreamSendState::NotSent =>
{
GatewayError::new(
let gateway = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream provider is currently available for this request",
)
);
match error.client_visible_upstream_error() {
Some(upstream) => gateway.with_client_visible_upstream_error(upstream.clone()),
None => gateway,
}
}
EngineError::Provider(error) => GatewayError::from_provider(error),
EngineError::Cancelled => {
@@ -189,6 +189,22 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
} else {
vec!["text"]
};
let additional_speed_tiers = presentation
.service_tiers()
.iter()
.filter_map(|tier| tier.speed_tier())
.collect::<Vec<_>>();
let service_tiers = presentation
.service_tiers()
.iter()
.map(|tier| {
json!({
"id": tier.id(),
"name": tier.name(),
"description": tier.description(),
})
})
.collect::<Vec<_>>();
json!({
"slug": id,
@@ -200,8 +216,8 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
"visibility": if presentation.hidden() { "hide" } else { "list" },
"supported_in_api": true,
"priority": index.saturating_add(1),
"additional_speed_tiers": [],
"service_tiers": [],
"additional_speed_tiers": additional_speed_tiers,
"service_tiers": service_tiers,
"availability_nux": Value::Null,
"upgrade": Value::Null,
"base_instructions": CODEX_BASE_INSTRUCTIONS,
@@ -105,12 +105,15 @@ mod actions {
use gateway_admin::model::{
Revision,
accounts::{AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus},
provider_credentials::{CredentialDeletionResult, CredentialImportResult},
provider_credentials::{
CredentialDeletionResult, CredentialImportResult, CredentialMutationResult,
},
};
use gateway_api::admin::accounts::{
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountRefreshRequest, AccountTestQuery,
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
AccountTestQuery, RotateAccountRequest, StartAccountAuthorizationRequest,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
@@ -147,7 +150,7 @@ mod actions {
}
#[test]
fn account_actions_should_require_frozen_account_ids_and_revision() {
fn account_actions_should_require_frozen_account_ids_and_reject_unknown_revision() {
let id: AccountIdQuery =
serde_json::from_value(json!({ "accountId": "acct_1" })).expect("decode ID query");
assert!(id.validate().is_ok());
@@ -155,13 +158,68 @@ mod actions {
serde_json::from_value(json!({ "accountId": "legacy-id" })).expect("decode action");
assert_eq!(action.validate().unwrap_err().field(), "accountId");
let refresh: AccountRefreshRequest = serde_json::from_value(json!({
"accountId": "acct_1",
"expectedConfigRevision": 0
"accountId": "acct_1"
}))
.expect("decode refresh");
assert!(refresh.validate().is_ok());
assert!(
serde_json::from_value::<AccountRefreshRequest>(json!({
"accountId": "acct_1",
"expectedConfigRevision": 0
}))
.is_err()
);
}
#[test]
fn credential_recovery_requests_should_not_accept_client_revision_fences() {
let authorization: StartAccountAuthorizationRequest = serde_json::from_value(json!({
"provider": "openai",
"name": "reauthorize",
"accountId": "acct_1"
}))
.expect("decode reauthorization");
assert!(authorization.validate().is_ok());
assert!(
serde_json::from_value::<StartAccountAuthorizationRequest>(json!({
"provider": "openai",
"name": "reauthorize",
"accountId": "acct_1",
"expectedCredentialRevision": 1
}))
.is_err()
);
let rotation: RotateAccountRequest = serde_json::from_value(json!({
"provider": "openai",
"accountId": "acct_1",
"accessToken": "header.payload.signature",
"refreshToken": "refresh-token"
}))
.expect("decode rotation");
assert!(rotation.validate().is_ok());
assert!(
serde_json::from_value::<RotateAccountRequest>(json!({
"provider": "openai",
"accountId": "acct_1",
"expectedCredentialRevision": 1,
"accessToken": "header.payload.signature"
}))
.is_err()
);
}
#[test]
fn credential_mutation_response_should_not_expose_internal_revision() {
let response = AccountMutationData::from(CredentialMutationResult {
config_revision: Revision::new(8).expect("config revision"),
account_id: ProviderAccountId::new("acct_1").expect("account ID"),
credential_revision: Some(Revision::new(9).expect("credential revision")),
});
assert_eq!(
refresh.validate().unwrap_err().field(),
"expectedConfigRevision"
serde_json::to_value(response).expect("serialize credential mutation"),
json!({ "accountId": "acct_1" })
);
}
@@ -169,7 +227,6 @@ mod actions {
fn account_import_should_use_provider_and_opaque_data_fields() {
let valid: AccountImportRequest = serde_json::from_value(json!({
"provider": "openai",
"expectedConfigRevision": 7,
"data": {
"providerOwnedUnknownField": {"nested": [1, 2, 3]},
"accounts": [{"credentials": {"access_token": "provider-validates-this"}}]
@@ -180,7 +237,6 @@ mod actions {
let invalid: AccountImportRequest = serde_json::from_value(json!({
"provider": "xai",
"expectedConfigRevision": 7,
"data": []
}))
.expect("decode invalid account import");
@@ -188,27 +244,32 @@ mod actions {
assert!(
serde_json::from_value::<AccountImportRequest>(json!({
"provider": "openai",
"expectedConfigRevision": 7,
"document": {}
}))
.is_err()
);
assert!(
serde_json::from_value::<AccountImportRequest>(json!({
"provider": "openai",
"expectedConfigRevision": 7,
"data": {}
}))
.is_err()
);
}
#[test]
fn account_deletion_should_validate_one_provider_batch_and_emit_account_ids() {
let request: AccountDeletionRequest = serde_json::from_value(json!({
"provider": "xai",
"accountIds": ["acct_1", "acct_2"],
"expectedConfigRevision": 8
"accountIds": ["acct_1", "acct_2"]
}))
.expect("decode account deletion");
assert!(request.validate().is_ok());
let duplicate: AccountDeletionRequest = serde_json::from_value(json!({
"provider": "xai",
"accountIds": ["acct_1", "acct_1"],
"expectedConfigRevision": 8
"accountIds": ["acct_1", "acct_1"]
}))
.expect("decode duplicate account deletion");
assert_eq!(duplicate.validate().unwrap_err().field(), "accountIds");
@@ -223,7 +284,6 @@ mod actions {
assert_eq!(
serde_json::to_value(response).expect("serialize account deletion"),
json!({
"configRevision": 9,
"deletedCount": 2,
"accountIds": ["acct_1", "acct_2"]
})
@@ -239,7 +299,6 @@ mod actions {
assert_eq!(
serde_json::to_value(response).expect("serialize account import"),
json!({
"configRevision": 8,
"importedCount": 1,
"accountIds": ["acct_imported"]
})
@@ -1,6 +1,6 @@
use chrono::{DateTime, TimeZone as _, Utc};
use gateway_api::admin::client_keys::{
self, ClientKeyCursorData, ClientKeyCursorValue, ClientKeyListData, ClientKeyRevisionRequest,
self, ClientKeyCursorData, ClientKeyCursorValue, ClientKeyListData, ClientKeyMutationRequest,
ClientKeySort, ClientKeySortDirection, ClientKeySortField, ClientKeyView,
CreateClientKeyRequest, CreatedClientKeyData, ListClientKeysQuery, MutatedClientKeyData,
RevealedClientKeyData, UpdateClientKeyRequest, decode_client_key_cursor,
@@ -101,9 +101,8 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
}
#[test]
fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields() {
fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
let valid = serde_json::from_value::<CreateClientKeyRequest>(json!({
"expectedConfigRevision": 7,
"name": "terminal key",
"label": "production",
"providerKind": "openai",
@@ -113,22 +112,11 @@ fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields(
.expect("deserialize create")
.into_command()
.expect("validate create");
assert_eq!(valid.expected_config_revision.get(), 7);
assert_eq!(valid.name, "terminal key");
for (payload, field) in [
(
json!({
"expectedConfigRevision": 0,
"name": "key",
"providerKind": "openai",
"maxConcurrency": 0,
"requestsPerMinute": 0
}),
"expectedConfigRevision",
),
(
json!({
"expectedConfigRevision": 1,
"name": " ",
"providerKind": "openai",
"maxConcurrency": 0,
@@ -138,7 +126,6 @@ fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields(
),
(
json!({
"expectedConfigRevision": 1,
"name": "key",
"providerKind": "openai",
"maxConcurrency": u64::MAX,
@@ -170,14 +157,23 @@ fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields(
}))
.is_err()
);
let revision = serde_json::from_value::<ClientKeyRevisionRequest>(json!({
"id": "key_1",
"expectedConfigRevision": 3
assert!(
serde_json::from_value::<CreateClientKeyRequest>(json!({
"expectedConfigRevision": 7,
"name": "terminal key",
"providerKind": "openai",
"maxConcurrency": 2,
"requestsPerMinute": 60
}))
.is_err()
);
let mutation_id = serde_json::from_value::<ClientKeyMutationRequest>(json!({
"id": "key_1"
}))
.expect("deserialize revision mutation")
.into_parts()
.expect("validate revision mutation");
assert_eq!(revision, ("key_1".to_owned(), 3));
.expect("deserialize mutation")
.into_id()
.expect("validate mutation");
assert_eq!(mutation_id, "key_1");
}
#[test]
@@ -225,9 +221,9 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
updated_at: created_at,
last_used_at: Some(created_at),
});
let list = serde_json::to_value(ClientKeyListData::new(9, vec![view], None, 1))
.expect("serialize list");
assert_eq!(list["configRevision"], 9);
let list =
serde_json::to_value(ClientKeyListData::new(vec![view], None, 1)).expect("serialize list");
assert!(list.get("configRevision").is_none());
assert_eq!(list["total"], 1);
assert_eq!(list["items"][0]["id"], "key_visible");
assert_eq!(list["items"][0]["providerKind"], "openai");
@@ -248,7 +244,6 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
let plaintext = format!("sk_{}", "a".repeat(43));
let created = CreatedClientKeyData::new(
10,
"key_created".to_owned(),
"sk_aaaaaaaaa".to_owned(),
plaintext.clone(),
@@ -263,9 +258,9 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
json!({ "id": "key_created", "plaintextKey": plaintext })
);
assert_eq!(
serde_json::to_value(MutatedClientKeyData::new(11, "key_created".to_owned()))
serde_json::to_value(MutatedClientKeyData::new("key_created".to_owned()))
.expect("serialize mutation"),
json!({ "configRevision": 11, "id": "key_created" })
json!({ "id": "key_created" })
);
}
@@ -279,8 +279,9 @@ impl SettingsStore for MemorySettingsStore {
command: ReplaceRuntimeSettings,
_: &MutationContext,
) -> AdminStoreResult<RuntimeSettings> {
let mut settings = self.settings.lock().expect("settings");
let updated = RuntimeSettings {
config_revision: next_revision(command.expected_config_revision),
config_revision: next_revision(settings.config_revision),
model_mappings: command.model_mappings,
refresh_margin_seconds: command.refresh_margin_seconds,
refresh_concurrency: command.refresh_concurrency,
@@ -292,31 +293,31 @@ impl SettingsStore for MemorySettingsStore {
audit_retention_days: command.audit_retention_days,
updated_at: Utc::now(),
};
*self.settings.lock().expect("settings") = updated.clone();
*settings = updated.clone();
Ok(updated)
}
async fn replace_admin_api_key(
&self,
expected_config_revision: Revision,
key: AdminApiKey,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
*self.api_key.lock().expect("API key") = Some(key);
let revision = self.settings.lock().expect("settings").config_revision;
Ok(AdminApiKeyMutation {
config_revision: next_revision(expected_config_revision),
config_revision: next_revision(revision),
exists: true,
})
}
async fn delete_admin_api_key(
&self,
expected_config_revision: Revision,
_: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
*self.api_key.lock().expect("API key") = None;
let revision = self.settings.lock().expect("settings").config_revision;
Ok(AdminApiKeyMutation {
config_revision: next_revision(expected_config_revision),
config_revision: next_revision(revision),
exists: false,
})
}
@@ -37,7 +37,6 @@ async fn response_json(response: axum::response::Response) -> Value {
fn update_body() -> Value {
json!({
"expectedConfigRevision": 7,
"modelMappings": {
"gpt-5.4": "gpt-5.5",
"grok-latest": "grok-4.5"
@@ -64,22 +63,11 @@ fn settings_request_should_reject_unknown_rotation_strategy() {
}
#[test]
fn settings_request_should_require_positive_revision() {
fn settings_request_should_reject_unknown_revision_field() {
let mut body = update_body();
body["expectedConfigRevision"] = json!(0);
let request: UpdateRuntimeSettingsRequest =
serde_json::from_value(body).expect("decode settings");
assert_eq!(
request.validate().unwrap_err().field(),
"expectedConfigRevision"
);
body["expectedConfigRevision"] = json!(7);
let mut missing = update_body();
missing
.as_object_mut()
.expect("settings object")
.remove("expectedConfigRevision");
assert!(serde_json::from_value::<UpdateRuntimeSettingsRequest>(missing).is_err());
assert!(serde_json::from_value::<UpdateRuntimeSettingsRequest>(body).is_err());
}
#[test]
@@ -124,7 +112,7 @@ async fn settings_post_should_replace_global_model_mappings() {
.expect("settings update response");
let data = response_json(response).await["data"].clone();
assert_eq!(data["configRevision"], 8);
assert!(data.get("configRevision").is_none());
assert_eq!(data["modelMappings"]["gpt-5.4"], "gpt-5.5");
assert_eq!(data["modelMappings"]["grok-latest"], "grok-4.5");
}
@@ -135,6 +135,36 @@ fn local_provider_capacity_exhaustion_should_map_to_service_unavailable() {
);
}
#[test]
fn local_provider_capacity_exhaustion_should_preserve_safe_provider_detail() {
let detail = ClientVisibleUpstreamError::new(
"no account is eligible for the requested model",
Some("no_eligible_account".to_owned()),
Some("account_unavailable_error".to_owned()),
)
.expect("safe provider detail");
let error = EngineError::Provider(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
.with_client_visible_upstream_error(detail),
);
let gateway = gateway_error_from_engine(&error);
assert_eq!(gateway.kind(), GatewayErrorKind::NoAvailableProvider);
assert_eq!(
gateway.client_message(),
"no account is eligible for the requested model"
);
assert_eq!(gateway.client_error_code(), Some("no_eligible_account"));
assert_eq!(
gateway.client_error_type(),
Some("account_unavailable_error")
);
assert_eq!(
gateway.safe_message(),
"no upstream provider is currently available for this request"
);
}
#[test]
fn no_eligible_provider_account_should_map_to_service_unavailable() {
let error = EngineError::Provider(ProviderError::new(
@@ -10,7 +10,9 @@ use gateway_core::engine::execution::{
StartedExecution,
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::routing::{ModelPresentation, PublicModelId, PublicModelProfile};
use gateway_core::routing::{
ModelPresentation, ModelServiceTier, PublicModelId, PublicModelProfile,
};
use tower::ServiceExt;
use super::{api_router, authenticated_client};
@@ -73,7 +75,15 @@ impl ExecutionService for ModelsExecution {
)
.with_context_window_tokens(Some(500_000))
.with_image_input(true)
.with_agent_tools(true, true),
.with_agent_tools(true, true)
.with_service_tiers(vec![
ModelServiceTier::new(
"priority",
"Fast",
"Route the request through Codex fast mode.",
)
.with_speed_tier("fast"),
]),
)]
}
@@ -145,7 +155,15 @@ async fn models_should_encode_provider_profiles_for_current_codex_clients() {
assert_eq!(model["default_reasoning_level"], "medium");
assert_eq!(model["context_window"], 500_000);
assert_eq!(model["apply_patch_tool_type"], "freeform");
assert_eq!(model["service_tiers"], serde_json::json!([]));
assert_eq!(model["additional_speed_tiers"], serde_json::json!(["fast"]));
assert_eq!(
model["service_tiers"],
serde_json::json!([{
"id": "priority",
"name": "Fast",
"description": "Route the request through Codex fast mode."
}])
);
assert_eq!(
model["supported_reasoning_levels"],
serde_json::json!([
@@ -1055,6 +1055,86 @@ async fn streaming_upstream_wire_failure_should_not_be_rewritten_as_a_gateway_er
);
}
#[tokio::test]
async fn atomic_uncommitted_upstream_failure_batch_should_be_forwarded_once() {
let trace = Arc::new(Trace::default());
let raw_created = Bytes::from_static(
b"event: response.created\r\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_atomic_api\",\"model\":\"public-model\",\"status\":\"in_progress\"}}\r\n\r\n",
);
let raw_failure = Bytes::from_static(
b"event: response.failed\r\ndata: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_atomic_api\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"atomic upstream marker\"}}}\r\n\r\n",
);
let started_wire = ProviderEvent::canonical_with_wire(
vec![started()],
ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {
"id": "resp_atomic_api",
"model": "public-model",
"status": "in_progress"
}
}),
raw_created.clone(),
None,
None,
)
.expect("valid upstream started wire"),
);
let failed_wire = ProviderEvent::wire(
ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": "resp_atomic_api",
"status": "failed",
"error": {
"code": "rate_limit_exceeded",
"message": "atomic upstream marker"
}
}
}),
raw_failure.clone(),
None,
None,
)
.expect("valid upstream failure wire"),
);
let batch = CoordinatedEvent::try_batch(
vec![started_wire, failed_wire],
CommitRequirement::CommitBeforeDelivery,
)
.expect("atomic failure delivery");
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::Event(batch),
NextStep::Error(EngineError::Provider(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
))),
],
);
let response = stream_execution_response(Box::new(session), None).await;
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read SSE body");
let body = String::from_utf8(body.to_vec()).expect("SSE is UTF-8");
assert_eq!(body.matches("atomic upstream marker").count(), 1);
assert_eq!(body.matches("event: response.failed").count(), 1);
assert!(body.contains(std::str::from_utf8(&raw_created).expect("created frame")));
assert!(body.contains(std::str::from_utf8(&raw_failure).expect("failure frame")));
assert!(!body.contains("upstream capacity is temporarily unavailable"));
assert!(body.ends_with("data: [DONE]\n\n"));
assert_eq!(trace.snapshot(), vec!["next_event", "commit", "next_error"]);
}
#[tokio::test]
async fn streaming_second_commit_request_should_fail_and_finalize_once() {
let trace = Arc::new(Trace::default());
@@ -1,14 +1,34 @@
use std::sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
};
use std::time::{Duration, SystemTime};
use axum::{
body::Body,
http::{Request, StatusCode, header::AUTHORIZATION},
};
use futures::{SinkExt, StreamExt, future::BoxFuture};
use gateway_api::openai::responses::ResponseCreateFrameError;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartedExecution,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState};
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind};
use gateway_core::event::{
GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta,
};
use gateway_core::operation::Operation;
use serde_json::json;
use gateway_core::routing::PublicModelId;
use serde_json::{Value, json};
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message as ClientMessage;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tower::ServiceExt;
use super::decode_response_create;
use crate::openai::{api_router, models::ModelsExecution};
use crate::openai::{api_router, authenticated_client, models::ModelsExecution};
#[test]
fn response_create_should_default_to_the_websocket_streaming_contract() {
@@ -148,11 +168,233 @@ fn response_create_should_reject_non_boolean_stream_without_disclosing_body_valu
assert!(!rendered.contains(opaque_stream_value));
}
// `serve_responses_websocket`/`forward_execution` 只有拿到 `hyper::upgrade::OnUpgrade`
// 的真实 HTTP/1.1 升级连接才会运行;本 crate 的依赖图按 axum `ws` feature 编译
// hyper(无 `http1`),测试进程内不存在能完成升级的服务端,`hyper::upgrade` 也不
// 提供公开构造器。集成测试因此止步于 upgrade 边界:会话循环内的事件转发、错误帧
// 与断连取消仍未被覆盖,需要 dev 依赖开启 `http1` 或 socket 泛化后才能落地。
#[derive(Default)]
struct AtomicFailureTrace {
next_calls: AtomicUsize,
committed: AtomicBool,
finalized: AtomicBool,
}
struct AtomicFailureSession {
trace: Arc<AtomicFailureTrace>,
}
impl ExecutionSession for AtomicFailureSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
match self.trace.next_calls.fetch_add(1, Ordering::AcqRel) {
0 => Ok(Some(atomic_failure_batch())),
1 => {
self.trace.finalized.store(true, Ordering::Release);
Err(EngineError::Provider(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
)))
}
_ => Ok(None),
}
})
}
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async { Err(EngineError::InvalidDeliveryState) })
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
&[]
}
fn commit_downstream(&mut self, _: Option<u16>) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
self.trace.committed.store(true, Ordering::Release);
Ok(())
})
}
fn record_client_status(&mut self, _: u16) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async { Ok(()) })
}
fn is_finalized(&self) -> bool {
self.trace.finalized.load(Ordering::Acquire)
}
fn cancel(&self) {}
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move {
self.trace.finalized.store(true, Ordering::Release);
})
}
}
struct AtomicFailureExecution {
client: AuthenticatedClient,
trace: Arc<AtomicFailureTrace>,
}
impl ExecutionService for AtomicFailureExecution {
fn authenticate(
&self,
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
(plaintext == "sk_ws_atomic")
.then(|| self.client.clone())
.ok_or(ClientAuthenticationError::InvalidKey)
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
vec![PublicModelId::new("model-a").expect("public model")]
}
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
model.as_str() == "model-a"
}
fn start(
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
.expect("request id"),
created_at: SystemTime::now(),
stream: request.metadata.stream,
session: Box::new(AtomicFailureSession {
trace: Arc::clone(&self.trace),
}),
})
})
}
}
fn atomic_failure_batch() -> CoordinatedEvent {
let started = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
"resp_ws_atomic",
"model-a",
))],
ProtocolWireEvent::json(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {
"id": "resp_ws_atomic",
"model": "model-a",
"status": "in_progress"
}
}),
)
.expect("created wire"),
);
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": "resp_ws_atomic",
"status": "failed",
"error": {
"code": "rate_limit_exceeded",
"message": "websocket atomic upstream marker"
}
}
}),
)
.expect("failed wire"),
);
CoordinatedEvent::try_batch(
vec![started, failed],
CommitRequirement::CommitBeforeDelivery,
)
.expect("atomic WebSocket batch")
}
#[tokio::test]
async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
let trace = Arc::new(AtomicFailureTrace::default());
let execution = Arc::new(AtomicFailureExecution {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_atomic".parse().expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
))
.await
.expect("send response.create");
let mut messages = Vec::<Value>::new();
loop {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("WebSocket response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
let value: Value = serde_json::from_str(&text).expect("JSON WebSocket event");
let failed = value.get("type").and_then(Value::as_str) == Some("response.failed");
messages.push(value);
if failed {
break;
}
}
assert_eq!(
messages
.iter()
.filter(|message| message.get("type").and_then(Value::as_str) == Some("response.failed"))
.count(),
1
);
let failed = messages
.iter()
.find(|message| message.get("type").and_then(Value::as_str) == Some("response.failed"))
.expect("upstream failed event");
assert_eq!(
failed
.pointer("/response/error/message")
.and_then(Value::as_str),
Some("websocket atomic upstream marker")
);
assert!(trace.committed.load(Ordering::Acquire));
assert_eq!(trace.next_calls.load(Ordering::Acquire), 2);
assert!(trace.finalized.load(Ordering::Acquire));
socket.close(None).await.expect("close WebSocket");
server.abort();
}
const TEST_WEBSOCKET_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAA==";
@@ -32,6 +32,21 @@ pub struct AttemptCoordinator<S: ?Sized> {
engine: Arc<GatewayEngine<S>>,
}
#[derive(Debug, Clone)]
enum AccountSelection {
Scheduled(Option<crate::engine::credential::ProviderAccountId>),
Diagnostic(crate::engine::credential::ProviderAccountId),
}
impl AccountSelection {
fn required_account(&self) -> Option<&crate::engine::credential::ProviderAccountId> {
match self {
Self::Scheduled(account) => account.as_ref(),
Self::Diagnostic(account) => Some(account),
}
}
}
impl<S: ?Sized> AttemptCoordinator<S>
where
S: ExecutionStore,
@@ -56,6 +71,49 @@ where
required_account: Option<crate::engine::credential::ProviderAccountId>,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
self.start_with_account_selection(
request,
operation,
plan,
AccountSelection::Scheduled(required_account),
continuation,
cancellation,
)
.await
}
/// 对固定账号执行诊断请求。
///
/// 仅跳过该账号的本地可用性投影;账号租约、并发和请求间隔仍必须满足。
pub async fn start_diagnostic(
&self,
request: NewModelRequest,
operation: Operation,
plan: RoutingPlan,
required_account: crate::engine::credential::ProviderAccountId,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
self.start_with_account_selection(
request,
operation,
plan,
AccountSelection::Diagnostic(required_account),
continuation,
cancellation,
)
.await
}
async fn start_with_account_selection(
&self,
request: NewModelRequest,
operation: Operation,
plan: RoutingPlan,
account_selection: AccountSelection,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
let request_id = request.id.clone();
let client_api_key_ref = request.client_api_key_ref.clone();
@@ -84,7 +142,7 @@ where
request_persisted: false,
operation,
plan,
required_account,
account_selection,
continuation,
continuation_attempt,
account_state_owner,
@@ -107,7 +165,9 @@ where
client_response_id: None,
upstream_response_id: None,
last_account_exhaustion: None,
last_account_exhaustion_events: Vec::new(),
provider_attempt_outcomes: Vec::new(),
pending_terminal_failure: None,
};
if session.cancellation.is_cancelled() {
@@ -158,7 +218,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
request_persisted: bool,
operation: Operation,
plan: RoutingPlan,
required_account: Option<crate::engine::credential::ProviderAccountId>,
account_selection: AccountSelection,
continuation: Option<ContinuationBinding>,
continuation_attempt: ContinuationAttempt,
account_state_owner: Option<ProviderAccountStateOwner>,
@@ -185,7 +245,16 @@ pub struct ResponseExecutionSession<S: ?Sized> {
client_response_id: Option<String>,
upstream_response_id: Option<String>,
last_account_exhaustion: Option<ProviderError>,
/// 最近一次账号耗尽在重试时被丢弃的原始失败批次;只用于后续空选路终态。
last_account_exhaustion_events: Vec<ProviderEvent>,
provider_attempt_outcomes: Vec<ProviderAttemptOutcome>,
/// 原子失败批次已交给协议层、但尚待下游提交后收敛的原 Provider 错误。
pending_terminal_failure: Option<PendingTerminalFailure>,
}
struct PendingTerminalFailure {
error: ProviderError,
send_state: UpstreamSendState,
}
impl<S: ?Sized> ResponseExecutionSession<S>
@@ -201,6 +270,12 @@ where
if self.delivery_pending {
return Err(EngineError::DownstreamCommitRequired);
}
if let Some(pending) = self.pending_terminal_failure.take() {
let engine_error = provider_engine_error(&pending.error);
self.finish_provider_error_with_send_state(&pending.error, pending.send_state)
.await?;
return Err(engine_error);
}
if self.finalized {
return Ok(None);
}
@@ -220,6 +295,21 @@ where
)));
}
PullOutcome::AttemptDiscarded => {}
PullOutcome::TerminalFailure {
events,
error,
send_state,
} => {
let requirement = if self.downstream_committed_at.is_some() {
CommitRequirement::AlreadyCommitted
} else {
self.delivery_pending = true;
CommitRequirement::CommitBeforeDelivery
};
self.pending_terminal_failure =
Some(PendingTerminalFailure { error, send_state });
return CoordinatedEvent::try_batch(events, requirement).map(Some);
}
PullOutcome::End => {
if self.downstream_committed_at.is_some() {
self.finish_success().await?;
@@ -248,6 +338,14 @@ where
match self.pull().await? {
PullOutcome::Event(event) => events.push(event),
PullOutcome::AttemptDiscarded => events.clear(),
PullOutcome::TerminalFailure {
error, send_state, ..
} => {
let engine_error = provider_engine_error(&error);
self.finish_provider_error_with_send_state(&error, send_state)
.await?;
return Err(engine_error);
}
PullOutcome::End => {
if events.is_empty() {
return Err(EngineError::InvalidDeliveryState);
@@ -376,13 +474,20 @@ where
return Ok(());
}
self.delivery_pending = false;
if let Some(pending) = self.pending_terminal_failure.take() {
return self
.finish_provider_error_with_send_state(&pending.error, pending.send_state)
.await;
}
self.finish_interruption(EngineError::Cancelled).await
}
async fn pull(&mut self) -> Result<PullOutcome, EngineError> {
loop {
if self.current.is_none() {
self.prepare_attempt().await?;
if self.current.is_none()
&& let Some(outcome) = self.prepare_attempt().await?
{
return Ok(outcome);
}
let boundary = {
@@ -412,10 +517,14 @@ where
if let Some(observation) = event.take_observation()
&& let Err(error) = self.observe_response(observation)
{
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
continue;
}
let mut identity_error = None;
for fact in event.canonical_facts() {
@@ -425,10 +534,14 @@ where
}
}
if let Some(error) = identity_error {
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
continue;
}
for fact in event.canonical_facts() {
self.observe_event(fact).await?;
@@ -442,8 +555,13 @@ where
return Ok(PullOutcome::Event(event));
}
PollBoundary::Item(Some(Err(error))) => {
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
}
PollBoundary::Item(None) => {
@@ -455,7 +573,7 @@ where
}
}
async fn prepare_attempt(&mut self) -> Result<(), EngineError> {
async fn prepare_attempt(&mut self) -> Result<Option<PullOutcome>, EngineError> {
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
@@ -483,27 +601,36 @@ where
.ok_or(EngineError::EmptyRoutingPlan)?;
// recovery 钉选在此被一次性消费,只绑定本次 replay attempt;
// 外部 required_account 每次 attempt 都重新生效。
let pinned_account = self
.recovery_account
.take()
.or_else(|| self.required_account.clone());
let pinned_account = match &self.account_selection {
AccountSelection::Diagnostic(account) => Some(account.clone()),
AccountSelection::Scheduled(_) => self
.recovery_account
.take()
.or_else(|| self.account_selection.required_account().cloned()),
};
let account_context = match &self.account_selection {
AccountSelection::Diagnostic(account) => AccountAttemptContext::diagnostic(
self.excluded_accounts.clone(),
account.clone(),
self.account_state_owner.clone(),
),
AccountSelection::Scheduled(_) => AccountAttemptContext::new(
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
),
}
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
.contains(account)
}));
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
.with_timing_started_at(self.timing_started_at),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
AccountAttemptContext::new(
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
)
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(
|account| {
self.credential_recovery_attempted_accounts
.contains(account)
},
)),
account_context,
self.continuation.clone(),
self.cancellation.clone(),
)
@@ -548,6 +675,16 @@ where
if error.kind() == ProviderErrorKind::NoEligibleAccount
&& let Some(exhaustion) = self.last_account_exhaustion.clone()
{
let send_state = self.current_send_state();
let mut events = std::mem::take(&mut self.last_account_exhaustion_events);
if !events.is_empty() {
self.observe_atomic_terminal_events(&mut events)?;
return Ok(Some(PullOutcome::TerminalFailure {
events,
error: exhaustion,
send_state,
}));
}
let engine_error = provider_engine_error(&exhaustion);
self.finish_provider_error(&exhaustion).await?;
return Err(engine_error);
@@ -672,10 +809,24 @@ where
send_observed: false,
response_observation: None,
});
Ok(())
Ok(None)
}
async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> {
self.observe_event_facts(event);
let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?;
if !current.send_observed {
self.engine
.store()
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
.await?;
current.send_observed = true;
self.send_state_watermark = UpstreamSendState::Sent;
}
Ok(())
}
fn observe_event_facts(&mut self, event: &GatewayEvent) {
let elapsed = elapsed_ms(self.timing_started_at);
observe_event_timing(&mut self.timings, event, elapsed);
if let GatewayEvent::Usage(observed) = event {
@@ -689,16 +840,6 @@ where
if let GatewayEvent::ProviderCost(observed) = event {
self.cost = observed.into_estimate();
}
let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?;
if !current.send_observed {
self.engine
.store()
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
.await?;
current.send_observed = true;
self.send_state_watermark = UpstreamSendState::Sent;
}
Ok(())
}
async fn observe_wire_event(&mut self) -> Result<(), EngineError> {
@@ -788,14 +929,22 @@ where
Ok(())
}
/// 返回 `true` 表示调用方必须丢弃本 attempt 已收集的未提交事件。
async fn handle_stream_error(&mut self, error: ProviderError) -> Result<bool, EngineError> {
/// 失败可重试时要求丢弃本 attempt;预算耗尽时可返回最后一批原始失败事件。
async fn handle_stream_error(
&mut self,
mut error: ProviderError,
) -> Result<StreamErrorOutcome, EngineError> {
// 原始 wire 只能活在本次决策栈内;后续 clone、attempt 记录与终态对象均只
// 接触已剥离的稳定错误字段。
let mut atomic_client_events = error.take_atomic_client_events();
let current = self.current.take().ok_or(EngineError::NoActiveAttempt)?;
if matches!(
let account_exhausted = matches!(
error.kind(),
ProviderErrorKind::RateLimited | ProviderErrorKind::QuotaExhausted
) {
);
if account_exhausted {
self.last_account_exhaustion = Some(error.clone());
self.last_account_exhaustion_events.clear();
}
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
// attempt_send_state 是本 attempt 自身的发送事实,驱动重试门;
@@ -818,7 +967,7 @@ where
attempt_send_state,
provider_proved_replay_safe,
);
let ordinary_retry = self.required_account.is_none()
let ordinary_retry = self.account_selection.required_account().is_none()
&& self.continuation_attempt == ContinuationAttempt::None
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
@@ -843,6 +992,9 @@ where
let retryable = continuation_retry || same_account_retry || ordinary_retry;
if retryable {
if account_exhausted {
self.last_account_exhaustion_events = atomic_client_events;
}
if same_account_retry {
if let Some(account) = current.metadata.provider_account_id() {
self.credential_recovery_attempted_accounts
@@ -879,16 +1031,42 @@ where
})
.await?;
self.reset_uncommitted_observations();
return Ok(true);
return Ok(StreamErrorOutcome::AttemptDiscarded);
}
let engine_error = provider_engine_error(&error);
self.current = Some(current);
if !atomic_client_events.is_empty() {
self.observe_atomic_terminal_events(&mut atomic_client_events)?;
return Ok(StreamErrorOutcome::TerminalFailure {
events: atomic_client_events,
error,
send_state,
});
}
self.finish_provider_error_with_send_state(&error, send_state)
.await?;
Err(engine_error)
}
fn observe_atomic_terminal_events(
&mut self,
events: &mut [ProviderEvent],
) -> Result<(), EngineError> {
for event in events {
if let Some(observation) = event.take_observation() {
self.observe_response(observation)
.map_err(EngineError::Provider)?;
}
for fact in event.canonical_facts() {
self.observe_response_identity(fact)
.map_err(EngineError::Provider)?;
self.observe_event_facts(fact);
}
}
Ok(())
}
fn prepare_continuation_retry(
&mut self,
current: &CurrentAttempt,
@@ -896,7 +1074,7 @@ where
send_state: UpstreamSendState,
provider_proved_replay_safe: bool,
) -> bool {
if self.required_account.is_some()
if self.account_selection.required_account().is_some()
|| self.continuation_attempt == ContinuationAttempt::None
|| self.downstream_committed_at.is_some()
|| self.delivery_pending
@@ -1238,9 +1416,40 @@ where
enum PullOutcome {
Event(ProviderEvent),
AttemptDiscarded,
TerminalFailure {
events: Vec<ProviderEvent>,
error: ProviderError,
send_state: UpstreamSendState,
},
End,
}
enum StreamErrorOutcome {
AttemptDiscarded,
TerminalFailure {
events: Vec<ProviderEvent>,
error: ProviderError,
send_state: UpstreamSendState,
},
}
impl StreamErrorOutcome {
fn into_pull_outcome(self) -> PullOutcome {
match self {
Self::AttemptDiscarded => PullOutcome::AttemptDiscarded,
Self::TerminalFailure {
events,
error,
send_state,
} => PullOutcome::TerminalFailure {
events,
error,
send_state,
},
}
}
}
enum PollBoundary {
Item(Option<Result<ProviderEvent, ProviderError>>),
Cancelled,
@@ -602,6 +602,15 @@ pub trait ProviderAccountStore: Send + Sync {
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError>;
/// 读取账号当前 credential 及其 revision,不做任何版本比对。
///
/// 管理写入必须在临近 CAS 时用它取 fence,而不是携带调用方持有的旧 revision:
/// 后台刷新随时会推进 revision,用陈旧快照会把正常的恢复操作误判为冲突。
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError>;
async fn compare_and_swap_credential(
&self,
update: CredentialCasUpdate,
@@ -901,6 +910,26 @@ pub struct AccountSelectionContext {
pub excluded_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
pub round_robin_cursor: u64,
pub availability: AccountAvailabilityPolicy,
}
/// 选择账号时是否信任本地可用性投影。
///
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
/// 请求间隔约束,且只能与固定账号约束组合使用。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AccountAvailabilityPolicy {
#[default]
Enforce,
BypassForDiagnostic,
}
impl AccountAvailabilityPolicy {
#[must_use]
pub const fn bypasses_local_availability(self) -> bool {
matches!(self, Self::BypassForDiagnostic)
}
}
/// 同一 target 内唯一的账号排序器。
@@ -1101,7 +1130,8 @@ fn smart_preference_should_escape(signals: &AccountRuntimeSignals) -> bool {
}
fn eligible(candidate: &AccountCandidate, context: &AccountSelectionContext) -> bool {
if !candidate.account.is_schedulable(context.now)
if (!context.availability.bypasses_local_availability()
&& !candidate.account.is_schedulable(context.now))
|| context.excluded_accounts.contains(candidate.account.id())
|| candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get()
{
@@ -477,11 +477,11 @@ impl DefaultExecutionService {
};
let mut session = match self
.probe_coordinator
.start(
.start_diagnostic(
new_request,
operation,
plan,
Some(account_id),
account_id,
None,
CancellationToken::new(),
)
@@ -525,7 +525,8 @@ impl DefaultExecutionService {
})
}
/// 探测失败先落库再转成客户端错误;转换会丢掉上游状态码与 provider code。
/// 探测失败先记录完整分类事实,再仅把 Provider 显式标记为客户端可见的
/// 结构化 message/code/type 带入协议错误。
async fn observe_probe_failure(
&self,
observed: &ProbeObservation,
@@ -232,6 +232,7 @@ pub struct AccountAttemptContext {
required_account: Option<ProviderAccountId>,
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
}
impl AccountAttemptContext {
@@ -246,6 +247,25 @@ impl AccountAttemptContext {
required_account,
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: false,
}
}
/// 为管理端对固定账号的上游诊断创建上下文。
///
/// 诊断只绕过本地的被动可用性投影,不能换号,并继续受账号租约保护。
#[must_use]
pub const fn diagnostic(
excluded_accounts: BTreeSet<ProviderAccountId>,
required_account: ProviderAccountId,
state_owner: Option<ProviderAccountStateOwner>,
) -> Self {
Self {
excluded_accounts,
required_account: Some(required_account),
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: true,
}
}
@@ -275,6 +295,12 @@ impl AccountAttemptContext {
pub const fn credential_recovery_attempted(&self) -> bool {
self.credential_recovery_attempted
}
/// 此尝试是否要由真实上游而非本地投影确认固定账号的可用性。
#[must_use]
pub const fn is_diagnostic_required_account(&self) -> bool {
self.diagnostic_required_account
}
}
/// 请求中 Provider 账号绑定状态的唯一归属。
@@ -474,6 +500,12 @@ impl AttemptContext {
self.account.credential_recovery_attempted()
}
/// 管理端固定账号诊断会跳过本地可用性投影,保留全部租约约束。
#[must_use]
pub const fn is_diagnostic_required_account(&self) -> bool {
self.account.is_diagnostic_required_account()
}
#[must_use]
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
self.continuation.as_ref()
+46 -1
View File
@@ -6,6 +6,7 @@ use std::time::Duration;
use thiserror::Error;
use crate::engine::UpstreamSendState;
use crate::event::ProviderEvent;
/// 应用层标识不满足约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
@@ -266,9 +267,13 @@ impl fmt::Debug for ClientVisibleUpstreamError {
/// 单次 Provider 调用的稳定错误。
///
/// 该类型不接收原始响应正文,也不会在 `Debug` 或 `Display` 中打印上游
/// 稳定诊断字段不接收原始响应正文,也不会在 `Debug` 或 `Display` 中打印上游
/// code/request ID/response ID。Adapter 若捕获到可能含 secret 的上下文,只能调用
/// [`ProviderError::redact_sensitive_context`] 丢弃正文并留下脱敏标记。
///
/// 唯一例外是 [`ProviderError::with_atomic_client_events`]:它只承载尚未交付客户端、
/// 必须与本错误一起完成重试判断的协议事件。Core 会在任何持久化或日志记录前取走
/// 这些事件,不能把它们当作诊断上下文使用。
#[derive(Clone)]
pub struct ProviderError {
kind: ProviderErrorKind,
@@ -284,8 +289,12 @@ pub struct ProviderError {
retry_same_account: bool,
sensitive_context_redacted: bool,
client_visible_upstream_error: Option<Box<ClientVisibleUpstreamError>>,
atomic_client_events: Option<Box<AtomicClientEvents>>,
}
#[derive(Clone)]
struct AtomicClientEvents(Vec<ProviderEvent>);
impl ProviderError {
/// 创建 Provider 错误。
#[must_use]
@@ -304,6 +313,7 @@ impl ProviderError {
retry_same_account: false,
sensitive_context_redacted: false,
client_visible_upstream_error: None,
atomic_client_events: None,
}
}
@@ -379,6 +389,17 @@ impl ProviderError {
self
}
/// 附加必须先由 Core 判断重试、再决定丢弃或交付的客户端协议事件。
///
/// Provider 只能放入当前 attempt 尚未交付的事件;调用方应设置明确的有界缓冲,
/// 并保证失败事件本身位于这批事件的末尾。
#[must_use]
pub fn with_atomic_client_events(mut self, events: Vec<ProviderEvent>) -> Self {
self.atomic_client_events =
(!events.is_empty()).then(|| Box::new(AtomicClientEvents(events)));
self
}
/// 丢弃敏感正文,只记录“发生过脱敏”这一事实。
#[must_use]
pub fn redact_sensitive_context(mut self, _sensitive_context: impl AsRef<str>) -> Self {
@@ -462,6 +483,23 @@ impl ProviderError {
pub fn client_visible_upstream_error(&self) -> Option<&ClientVisibleUpstreamError> {
self.client_visible_upstream_error.as_deref()
}
/// 取走只供本次重试/提交决策使用的原子客户端事件。
///
/// Core 必须在克隆错误、记录中间失败或构造终态前调用本方法,避免原始 wire
/// 进入诊断或持久化对象。
#[must_use]
pub fn take_atomic_client_events(&mut self) -> Vec<ProviderEvent> {
self.atomic_client_events
.take()
.map_or_else(Vec::new, |events| events.0)
}
/// 返回是否携带尚未提交的原子客户端事件;不暴露事件内容。
#[must_use]
pub fn has_atomic_client_events(&self) -> bool {
self.atomic_client_events.is_some()
}
}
impl fmt::Debug for ProviderError {
@@ -508,6 +546,13 @@ impl fmt::Debug for ProviderError {
.as_ref()
.map(|_| "<present>"),
)
.field(
"atomic_client_events",
&self
.atomic_client_events
.as_ref()
.map_or(0, |events| events.0.len()),
)
.finish()
}
}
@@ -13,7 +13,7 @@ use crate::engine::credential::{
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
};
use crate::error::{IdentifierError, validate_text};
use crate::routing::ProviderKind;
use crate::routing::{ProviderKind, UpstreamModelId};
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
@@ -527,6 +527,80 @@ impl ProviderCooldown {
}
}
/// 可丢失 cooldown 的细粒度作用域。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum ProviderCooldownScope {
/// 只阻止同一账号调用指定上游模型。
UpstreamModel(UpstreamModelId),
}
impl ProviderCooldownScope {
#[must_use]
pub const fn upstream_model(model: UpstreamModelId) -> Self {
Self::UpstreamModel(model)
}
#[must_use]
pub const fn kind(&self) -> &'static str {
match self {
Self::UpstreamModel(_) => "model",
}
}
#[must_use]
pub fn value(&self) -> &str {
match self {
Self::UpstreamModel(model) => model.as_str(),
}
}
}
/// 不进入账号持久状态的账号+作用域 cooldown。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderScopedCooldown {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
scope: ProviderCooldownScope,
until: SystemTime,
}
impl ProviderScopedCooldown {
#[must_use]
pub const fn new(
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
scope: ProviderCooldownScope,
until: SystemTime,
) -> Self {
Self {
account_id,
credential_revision,
scope,
until,
}
}
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
&self.account_id
}
#[must_use]
pub const fn credential_revision(&self) -> CredentialRevision {
self.credential_revision
}
#[must_use]
pub const fn scope(&self) -> &ProviderCooldownScope {
&self.scope
}
#[must_use]
pub const fn until(&self) -> SystemTime {
self.until
}
}
pub trait ProviderCooldownPort: Send + Sync {
fn put_if_later(
&self,
@@ -543,6 +617,24 @@ pub trait ProviderCooldownPort: Send + Sync {
account_id: &'a ProviderAccountId,
through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
fn put_scoped_if_later(
&self,
cooldown: ProviderScopedCooldown,
) -> BoxFuture<'_, Result<bool, ProviderStoreError>>;
fn read_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>>;
fn clear_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -778,35 +870,75 @@ pub enum OAuthPendingPutOutcome {
AlreadyExists,
}
/// OAuth 回调处理取得临时 flow 的独占权结果。
///
/// 与一次性消费不同,Provider 在上游换取令牌失败时可以释放 claim,让同一 flow
/// 使用新的回调地址重试;只有完整校验成功后才会消费 flow。
#[derive(Clone, PartialEq)]
pub enum OAuthPendingTakeOutcome {
Taken(OpaqueProviderData),
pub enum OAuthPendingClaimOutcome {
Claimed(OpaqueProviderData),
NotFound,
OwnerMismatch,
InProgress,
}
impl fmt::Debug for OAuthPendingTakeOutcome {
impl fmt::Debug for OAuthPendingClaimOutcome {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Taken(_) => formatter.write_str("Taken([PROVIDER-OWNED])"),
Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"),
Self::NotFound => formatter.write_str("NotFound"),
Self::OwnerMismatch => formatter.write_str("OwnerMismatch"),
Self::InProgress => formatter.write_str("InProgress"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OAuthPendingReleaseOutcome {
Released,
NotFound,
OwnerMismatch,
ClaimMismatch,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OAuthPendingConsumeOutcome {
Consumed,
NotFound,
OwnerMismatch,
ClaimMismatch,
}
pub trait OAuthPendingFlowPort: Send + Sync {
fn put_if_absent(
&self,
flow: NewOAuthPendingFlow,
) -> BoxFuture<'_, Result<OAuthPendingPutOutcome, ProviderStoreError>>;
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>>;
claim: &'a OAuthPendingBinding,
claim_ttl: Duration,
) -> BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>>;
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>>;
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>>;
}
/// Provider 只能按能力取用端口,无法取得 Redis client 或 repository 集合。
@@ -152,9 +152,61 @@ pub struct ModelPresentation {
search_tool: bool,
image_detail_original: bool,
verbosity: bool,
service_tiers: Vec<ModelServiceTier>,
hidden: bool,
}
/// Provider 声明给 Codex 客户端的服务档位。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelServiceTier {
id: String,
name: String,
description: String,
speed_tier: Option<String>,
}
impl ModelServiceTier {
#[must_use]
pub fn new(
id: impl Into<String>,
name: impl Into<String>,
description: impl Into<String>,
) -> Self {
Self {
id: id.into(),
name: name.into(),
description: description.into(),
speed_tier: None,
}
}
#[must_use]
pub fn with_speed_tier(mut self, speed_tier: impl Into<String>) -> Self {
self.speed_tier = Some(speed_tier.into());
self
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn description(&self) -> &str {
&self.description
}
#[must_use]
pub fn speed_tier(&self) -> Option<&str> {
self.speed_tier.as_deref()
}
}
impl ModelPresentation {
#[must_use]
pub fn new(display_name: Option<String>, description: Option<String>) -> Self {
@@ -213,6 +265,12 @@ impl ModelPresentation {
self
}
#[must_use]
pub fn with_service_tiers(mut self, service_tiers: Vec<ModelServiceTier>) -> Self {
self.service_tiers = service_tiers;
self
}
#[must_use]
pub const fn with_hidden(mut self, hidden: bool) -> Self {
self.hidden = hidden;
@@ -274,6 +332,11 @@ impl ModelPresentation {
self.verbosity
}
#[must_use]
pub fn service_tiers(&self) -> &[ModelServiceTier] {
&self.service_tiers
}
#[must_use]
pub const fn hidden(&self) -> bool {
self.hidden
@@ -399,6 +399,69 @@ fn complete_stream(total_tokens: Option<u64>) -> Vec<Result<GatewayEvent, Provid
events
}
fn atomic_response_failed(response_id: &str, marker: &str) -> ProviderError {
let started = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
response_id,
"gpt-5",
))],
ProtocolWireEvent::json(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {"id": response_id, "model": "gpt-5", "status": "in_progress"}
}),
)
.expect("started wire"),
);
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {"code": "rate_limit_exceeded", "message": marker}
}
}),
)
.expect("failed wire"),
);
ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_status(429)
.with_replay_safe()
.with_atomic_client_events(vec![started, failed])
}
fn bare_atomic_response_failed(response_id: &str) -> ProviderError {
let failed = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
response_id,
"fallback-model",
))],
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {"code": "rate_limit_exceeded", "message": "bare failure"}
}
}),
)
.expect("bare failed wire"),
);
ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_status(429)
.with_replay_safe()
.with_atomic_client_events(vec![failed])
}
fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, ProviderError>> {
let mut usage = Usage::new();
usage.input_tokens = Some(12);
@@ -1202,7 +1265,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
}
#[test]
fn required_account_reaches_provider_and_matching_metadata_succeeds() {
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
@@ -1210,11 +1273,11 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() {
items: complete_stream(None),
}]);
let required = ProviderAccountId::new("acct_required").expect("account id");
let mut session = block_on(coordinator.start(
let mut session = block_on(coordinator.start_diagnostic(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
Some(required.clone()),
required.clone(),
None,
CancellationToken::new(),
))
@@ -1225,6 +1288,7 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() {
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 1);
assert_eq!(contexts[0].required_account(), Some(&required));
assert!(contexts[0].is_diagnostic_required_account());
}
#[test]
@@ -1836,6 +1900,277 @@ fn explicit_429_after_structural_event_should_retry_before_commit() {
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn atomic_response_failed_should_rotate_account_before_stream_commit() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(atomic_response_failed(
"response-first-failed",
"first account failure",
))],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("second account should produce the first delivery")
.expect("started event");
let delivered = first.into_provider_events();
assert!(delivered.iter().flat_map(ProviderEvent::canonical_facts).any(
|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "response-1")
));
assert!(
!delivered
.iter()
.filter_map(ProviderEvent::wire_event)
.any(|wire| {
wire.data().pointer("/response/id").and_then(Value::as_str)
== Some("response-first-failed")
})
);
block_on(session.commit_downstream(Some(200))).expect("commit second account");
while block_on(session.next_event())
.expect("successful stream")
.is_some()
{}
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.commits, 1);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn bare_atomic_response_failed_should_rotate_before_the_first_delivery() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(bare_atomic_response_failed("response-bare-failed"))],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("second account should replace the bare failure")
.expect("second account started event");
assert!(first
.into_provider_events()
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "response-1")));
block_on(session.commit_downstream(Some(200))).expect("commit second account");
while block_on(session.next_event())
.expect("successful stream")
.is_some()
{}
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once() {
let operation = generate_operation();
let route_plan = plan(&operation);
let scripts = (1..=route_plan.max_attempts().get())
.map(|attempt| {
let response_id = format!("response-attempt-{attempt}-failed");
let marker = if attempt == route_plan.max_attempts().get() {
"last upstream failure".to_owned()
} else {
format!("discarded failure {attempt}")
};
Script::ObservedStream {
account_id: "acct_failed",
items: vec![Err(atomic_response_failed(&response_id, &marker))],
}
})
.collect::<Vec<_>>();
let expected_attempts = route_plan.max_attempts().get();
let expected_last_response = format!("response-attempt-{expected_attempts}-failed");
let (coordinator, store, provider) = coordinator(scripts);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("last failure remains deliverable")
.expect("terminal failure batch");
assert_eq!(
terminal.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
let events = terminal.into_provider_events();
let failures = events
.iter()
.filter_map(ProviderEvent::wire_event)
.filter(|wire| wire.event_type() == Some("response.failed"))
.collect::<Vec<_>>();
assert_eq!(failures.len(), 1);
assert_eq!(
failures[0]
.data()
.pointer("/response/id")
.and_then(Value::as_str),
Some(expected_last_response.as_str())
);
assert!(
failures[0]
.data()
.to_string()
.contains("last upstream failure")
);
block_on(session.commit_downstream(Some(200))).expect("commit original terminal failure");
let error = block_on(session.next_event()).expect_err("typed failure finalizes the request");
assert!(matches!(error, EngineError::Provider(_)));
assert!(session.is_finalized());
assert_eq!(
provider.contexts.lock().expect("contexts lock").len(),
usize::try_from(expected_attempts).expect("attempt count fits usize")
);
let state = store.state.lock().expect("store lock");
assert_eq!(
state.intermediate_failures,
usize::try_from(expected_attempts - 1).expect("attempt count fits usize")
);
assert_eq!(state.commits, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert!(state.finalizations[0].committed);
}
#[test]
fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_only",
items: vec![Err(bare_atomic_response_failed("response-only-failed"))],
},
Script::Error(ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
UpstreamSendState::NotSent,
)),
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("last account failure remains deliverable")
.expect("terminal failure batch");
let events = terminal.into_provider_events();
let failure = events
.iter()
.filter_map(ProviderEvent::wire_event)
.find(|wire| wire.event_type() == Some("response.failed"))
.expect("last upstream response.failed");
assert_eq!(
failure
.data()
.pointer("/response/id")
.and_then(Value::as_str),
Some("response-only-failed")
);
block_on(session.commit_downstream(Some(200))).expect("commit terminal failure");
assert!(matches!(
block_on(session.next_event()),
Err(EngineError::Provider(_))
));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert_eq!(
state.finalizations[0].client_response_id.as_deref(),
Some("response-only-failed")
);
assert_eq!(
state.finalizations[0].upstream_response_id.as_deref(),
Some("response-only-failed")
);
}
#[test]
fn atomic_response_failed_keeps_collect_uncommitted_error_semantics() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![Script::ObservedStream {
account_id: "acct_required",
items: vec![Err(atomic_response_failed(
"response-required-failed",
"required account failure",
))],
}]);
let required = ProviderAccountId::new("acct_required").expect("required account");
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
Some(required),
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted())
.expect_err("non-streaming collection must keep returning the typed failure");
assert!(matches!(error, EngineError::Provider(_)));
assert!(session.is_finalized());
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert_eq!(state.commits, 0);
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn native_continuation_explicit_429_is_not_retried() {
let continuation = NativeContinuationPin::new(
@@ -5,11 +5,11 @@ use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountCandidate, AccountFeedbackStats,
AccountQuotaSignals, AccountRuntimeSignals, AccountSelectionContext, AccountSelectionPolicy,
AccountSelector, CredentialCasUpdate, CredentialRevision, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
RotationStrategy,
AccountAttemptFeedback, AccountAvailability, AccountAvailabilityPolicy, AccountCandidate,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSelectionContext,
AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountUpdate, RotationStrategy,
};
use gateway_core::routing::ProviderKind;
@@ -51,6 +51,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
excluded_accounts: BTreeSet::new(),
preferred_account: None,
round_robin_cursor: 0,
availability: AccountAvailabilityPolicy::Enforce,
}
}
@@ -80,6 +81,55 @@ fn availability_should_round_trip_all_database_values() {
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_availability() {
let exhausted = AccountCandidate {
account: account("acct_exhausted").with_runtime_state(
true,
AccountAvailability::QuotaExhausted,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let disabled = AccountCandidate {
account: account("acct_disabled").with_runtime_state(
false,
AccountAvailability::QuotaExhausted,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let mut context = context(RotationStrategy::Sticky);
context.availability = AccountAvailabilityPolicy::BypassForDiagnostic;
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&exhausted), &context)
.map(|candidate| candidate.account.id()),
Some(exhausted.account.id())
);
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&disabled), &context)
.map(|candidate| candidate.account.id()),
Some(disabled.account.id())
);
}
#[test]
fn disabled_account_should_not_be_schedulable() {
let account =
@@ -2,6 +2,8 @@ use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{
ClientVisibleUpstreamError, GatewayError, ProviderError, ProviderErrorKind, SafeUpstreamValue,
};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent};
use serde_json::json;
#[test]
fn provider_error_debug_should_not_expose_sensitive_context() {
@@ -37,6 +39,26 @@ fn provider_error_replay_proof_should_be_explicit() {
assert!(error.replay_is_safe());
}
#[test]
fn provider_error_atomic_client_events_should_be_take_only_and_debug_redacted() {
let marker = "atomic-wire-must-not-enter-debug";
let event = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({"type": "response.failed", "message": marker}),
)
.expect("atomic wire"),
);
let mut error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_atomic_client_events(vec![event]);
assert!(error.has_atomic_client_events());
assert!(!format!("{error:?}").contains(marker));
assert_eq!(error.take_atomic_client_events().len(), 1);
assert!(!error.has_atomic_client_events());
}
#[test]
fn gateway_error_should_keep_client_visible_upstream_fields_out_of_safe_diagnostics() {
let message = "Your Codex quota is exhausted";
@@ -0,0 +1,126 @@
use serde_json::{Map, Value};
const MULTI_AGENT_MODE_OPEN_TAG: &str = "<multi_agent_mode>";
const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>";
const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active.";
/// 从 OpenAI Responses 请求中提取的稳定 Codex 请求语义。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexResponsesRequestSemantics {
/// 客户端原始选择的推理强度。
pub reasoning_effort: Option<String>,
/// Codex turn metadata 中的请求类型。
pub request_kind: Option<String>,
/// Codex turn metadata 中的子代理类型。
pub subagent_kind: Option<String>,
/// 网关侧派生出的推理预设。
pub reasoning_preset: Option<&'static str>,
/// 请求是否为 Codex 压缩请求。
pub compact: bool,
}
/// 从 Responses body 与非 wire 协议上下文中提取 Codex 语义。
#[must_use]
pub fn codex_responses_request_semantics(
body: &Map<String, Value>,
context: &Map<String, Value>,
) -> CodexResponsesRequestSemantics {
codex_responses_request_semantics_with_turn_metadata(
body,
non_empty_string(context.get("turn_metadata")),
)
}
/// 调用方已解析权威 turn metadata 时,从 Responses body 中提取 Codex 语义。
#[must_use]
pub fn codex_responses_request_semantics_with_turn_metadata(
body: &Map<String, Value>,
turn_metadata: Option<&str>,
) -> CodexResponsesRequestSemantics {
let reasoning_effort = body
.get("reasoning")
.and_then(Value::as_object)
.and_then(|reasoning| non_empty_string(reasoning.get("effort")))
.map(ToOwned::to_owned);
let parsed_turn_metadata = turn_metadata
.or_else(|| request_turn_metadata(body))
.and_then(|value| serde_json::from_str::<Value>(value).ok());
let request_kind = parsed_turn_metadata
.as_ref()
.and_then(|metadata| non_empty_string(metadata.get("request_kind")))
.map(ToOwned::to_owned);
let subagent_kind = parsed_turn_metadata
.as_ref()
.and_then(|metadata| non_empty_string(metadata.get("subagent_kind")))
.map(ToOwned::to_owned);
let proactive_multi_agent = latest_multi_agent_mode(input_items(body).iter())
.is_some_and(|mode| mode.starts_with(PROACTIVE_MULTI_AGENT_MODE_PREFIX));
let compact_trigger = input_items(body)
.iter()
.any(|item| non_empty_string(item.get("type")) == Some("compaction_trigger"));
let reasoning_preset = (subagent_kind.is_none()
&& reasoning_effort.as_deref() == Some("max")
&& proactive_multi_agent)
.then_some("ultra");
let compact = request_kind.as_deref() == Some("compaction") || compact_trigger;
CodexResponsesRequestSemantics {
reasoning_effort,
request_kind,
subagent_kind,
reasoning_preset,
compact,
}
}
fn request_turn_metadata(body: &Map<String, Value>) -> Option<&str> {
non_empty_string(body.get("turnMetadata"))
.or_else(|| non_empty_string(body.get("turn_metadata")))
.or_else(|| non_empty_string(body.get("x-codex-turn-metadata")))
.or_else(|| client_metadata_string(body, "x-codex-turn-metadata"))
.or_else(|| client_metadata_string(body, "turnMetadata"))
.or_else(|| client_metadata_string(body, "turn_metadata"))
}
fn client_metadata_string<'a>(body: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
body.get("client_metadata")
.and_then(Value::as_object)
.and_then(|metadata| non_empty_string(metadata.get(key)))
}
fn input_items(body: &Map<String, Value>) -> &[Value] {
body.get("input")
.and_then(Value::as_array)
.map_or(&[], Vec::as_slice)
}
fn latest_multi_agent_mode<'a>(
input: impl DoubleEndedIterator<Item = &'a Value>,
) -> Option<&'a str> {
input.rev().find_map(|item| {
if non_empty_string(item.get("role")) != Some("developer") {
return None;
}
item.get("content")?
.as_array()?
.iter()
.rev()
.filter_map(|content| non_empty_string(content.get("text")))
.find_map(multi_agent_mode_from_text)
})
}
fn multi_agent_mode_from_text(text: &str) -> Option<&str> {
let close = text.rfind(MULTI_AGENT_MODE_CLOSE_TAG)?;
let before_close = &text[..close];
let open = before_close.rfind(MULTI_AGENT_MODE_OPEN_TAG)?;
let body = &before_close[open + MULTI_AGENT_MODE_OPEN_TAG.len()..];
Some(body.trim())
}
fn non_empty_string(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
@@ -9,7 +9,14 @@ pub const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str =
/// Codex memory consolidation 请求标记。
pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request";
/// OpenAI Responses 请求中携带的 Codex 专属语义。
pub mod codex;
/// OpenAI/Codex 事件语义、用量与限流字段编解码。
pub mod events;
/// Server-Sent Events 帧的解析与编码。
pub mod sse;
pub use codex::{
CodexResponsesRequestSemantics, codex_responses_request_semantics,
codex_responses_request_semantics_with_turn_metadata,
};
@@ -0,0 +1,61 @@
use gateway_protocol::openai::codex_responses_request_semantics;
use serde_json::{Map, Value, json};
#[test]
fn codex_semantics_should_use_protocol_context_before_body_metadata() {
let body = json!({
"model": "grok-4.5",
"reasoning": {"effort": " xhigh "},
"client_metadata": {
"x-codex-turn-metadata": "{\"request_kind\":\"review\"}"
},
"input": [
{"type": "message", "role": "user", "content": "history"},
{"type": "compaction_trigger"}
]
});
let body = object(body);
let context = Map::from_iter([(
"turn_metadata".to_owned(),
json!("{\"request_kind\":\"compaction\",\"subagent_kind\":\"compact\"}"),
)]);
let semantics = codex_responses_request_semantics(&body, &context);
assert_eq!(semantics.reasoning_effort.as_deref(), Some("xhigh"));
assert_eq!(semantics.request_kind.as_deref(), Some("compaction"));
assert_eq!(semantics.subagent_kind.as_deref(), Some("compact"));
assert_eq!(semantics.reasoning_preset, None);
assert!(semantics.compact);
}
#[test]
fn codex_semantics_should_detect_ultra_preset_from_proactive_multi_agent() {
let body = object(json!({
"model": "gpt-test",
"reasoning": {"effort": "max"},
"input": [
{
"type": "message",
"role": "developer",
"content": [{
"type": "input_text",
"text": "<multi_agent_mode>Proactive multi-agent delegation is active.</multi_agent_mode>"
}]
}
]
}));
let semantics = codex_responses_request_semantics(&body, &Map::new());
assert_eq!(semantics.reasoning_effort.as_deref(), Some("max"));
assert_eq!(semantics.reasoning_preset, Some("ultra"));
assert!(!semantics.compact);
}
fn object(value: Value) -> Map<String, Value> {
let Value::Object(object) = value else {
panic!("test value must be an object");
};
object
}
@@ -1,2 +1,3 @@
mod codex;
mod events;
mod sse;
+3 -13
View File
@@ -582,7 +582,6 @@ impl SettingsStore for AdminSettingsStoreAdapter {
};
let snapshot = postgres::ControlPlaneRepository::replace_control_plane(
&self.control_plane,
store_revision(command.expected_config_revision)?,
replacement,
)
.await
@@ -592,32 +591,24 @@ impl SettingsStore for AdminSettingsStoreAdapter {
async fn replace_admin_api_key(
&self,
expected_config_revision: AdminRevision,
key: AdminApiKey,
context: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
self.replace_admin_api_key_value(
expected_config_revision,
Some(key.expose_for_auth().to_owned()),
context,
)
.await
self.replace_admin_api_key_value(Some(key.expose_for_auth().to_owned()), context)
.await
}
async fn delete_admin_api_key(
&self,
expected_config_revision: AdminRevision,
context: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
self.replace_admin_api_key_value(expected_config_revision, None, context)
.await
self.replace_admin_api_key_value(None, context).await
}
}
impl AdminSettingsStoreAdapter {
async fn replace_admin_api_key_value(
&self,
expected_config_revision: AdminRevision,
admin_api_key: Option<String>,
context: &MutationContext,
) -> AdminStoreResult<AdminApiKeyMutation> {
@@ -639,7 +630,6 @@ impl AdminSettingsStoreAdapter {
};
let snapshot = postgres::ControlPlaneRepository::replace_control_plane(
&self.control_plane,
store_revision(expected_config_revision)?,
postgres::ControlPlaneReplacement {
settings: update,
audit: mutation_audit(
@@ -36,7 +36,7 @@ use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
use crate::{
StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
postgres_unavailable, require_nonempty, store_revision,
postgres_unavailable, require_nonempty,
};
use super::{ControlPlaneRepository, PgControlPlaneRepository};
@@ -633,7 +633,6 @@ impl ClientKeyStore for PgAdminClientKeyStore {
let revision = self
.control_plane
.create_client_api_key(
store_revision(command.expected_config_revision)?,
NewClientApiKey {
id: id.as_str().to_owned(),
name: command.name,
@@ -676,7 +675,6 @@ impl ClientKeyStore for PgAdminClientKeyStore {
let revision = self
.control_plane
.update_client_api_key(
store_revision(command.expected_config_revision)?,
UpdateClientApiKeyDetails {
id: id.as_str().to_owned(),
name: command.name,
@@ -716,7 +714,6 @@ impl ClientKeyStore for PgAdminClientKeyStore {
let revision = self
.control_plane
.set_client_api_key_enabled(
store_revision(command.expected_config_revision)?,
id.as_str(),
command.enabled,
mutation_audit(
@@ -739,7 +736,6 @@ impl ClientKeyStore for PgAdminClientKeyStore {
) -> AdminStoreResult<gateway_admin::model::Revision> {
self.control_plane
.delete_client_api_key(
store_revision(command.expected_config_revision)?,
command.id.as_str(),
mutation_audit(
context,
@@ -74,27 +74,23 @@ pub trait ControlPlaneRepository: Send + Sync {
async fn replace_control_plane(
&self,
expected_revision: Revision,
replacement: ControlPlaneReplacement,
) -> StoreResult<ControlPlaneSnapshot>;
async fn create_client_api_key(
&self,
expected_revision: Revision,
key: NewClientApiKey,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn update_client_api_key(
&self,
expected_revision: Revision,
key: UpdateClientApiKeyDetails,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn set_client_api_key_enabled(
&self,
expected_revision: Revision,
id: &str,
enabled: bool,
audit: AdminAuditEvent,
@@ -102,7 +98,6 @@ pub trait ControlPlaneRepository: Send + Sync {
async fn delete_client_api_key(
&self,
expected_revision: Revision,
id: &str,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
@@ -142,7 +137,6 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
async fn replace_control_plane(
&self,
expected_revision: Revision,
replacement: ControlPlaneReplacement,
) -> StoreResult<ControlPlaneSnapshot> {
replacement.settings.validate()?;
@@ -152,12 +146,9 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
.await
.map_err(|_| postgres_unavailable("begin control plane replacement"))?;
let result = async {
let revision = update_runtime_settings_in_transaction(
&mut transaction,
expected_revision,
&replacement.settings,
)
.await?;
let revision =
update_runtime_settings_in_transaction(&mut transaction, &replacement.settings)
.await?;
append_admin_audit_event_in_transaction(&mut transaction, replacement.audit, revision)
.await?;
load_control_plane_in_transaction(&mut transaction).await
@@ -183,41 +174,29 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
async fn create_client_api_key(
&self,
expected_revision: Revision,
key: NewClientApiKey,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::CreateClientApiKey(key),
audit,
)
.await
self.apply_targeted_mutation(ControlPlaneMutation::CreateClientApiKey(key), audit)
.await
}
async fn update_client_api_key(
&self,
expected_revision: Revision,
key: UpdateClientApiKeyDetails,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::UpdateClientApiKey(key),
audit,
)
.await
self.apply_targeted_mutation(ControlPlaneMutation::UpdateClientApiKey(key), audit)
.await
}
async fn set_client_api_key_enabled(
&self,
expected_revision: Revision,
id: &str,
enabled: bool,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::SetClientApiKeyEnabled {
id: id.to_owned(),
enabled,
@@ -229,12 +208,10 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
async fn delete_client_api_key(
&self,
expected_revision: Revision,
id: &str,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::DeleteClientApiKey(id.to_owned()),
audit,
)
@@ -252,7 +229,6 @@ enum ControlPlaneMutation {
impl PgControlPlaneRepository {
async fn apply_targeted_mutation(
&self,
expected_revision: Revision,
mutation: ControlPlaneMutation,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
@@ -262,8 +238,7 @@ impl PgControlPlaneRepository {
.await
.map_err(|_| postgres_unavailable("begin targeted control plane mutation"))?;
let result = async {
let revision =
bump_config_revision_in_transaction(&mut transaction, expected_revision).await?;
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
match mutation {
ControlPlaneMutation::CreateClientApiKey(key) => {
insert_client_api_key_in_transaction(&mut transaction, &key).await?;
@@ -12,8 +12,7 @@ use gateway_admin::{
model::{
MutationContext, Revision as AdminRevision,
accounts::{
AccountAvailability as AdminAccountAvailability, AccountCost,
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountCost, AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
@@ -35,8 +34,8 @@ use gateway_admin::{
use sqlx::{PgPool, Postgres, Row, Transaction};
use gateway_core::engine::credential::{
AccountAvailability as CoreAccountAvailability, AccountStateChange, CredentialCasOutcome,
CredentialCasUpdate, CredentialRevision as CoreCredentialRevision, LoadedCredential,
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation,
@@ -121,43 +120,9 @@ impl ProviderAccountAdminScope {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderAccountAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
impl ProviderAccountAvailability {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Ready => "ready",
Self::Cooldown => "cooldown",
Self::QuotaExhausted => "quota_exhausted",
Self::Expired => "expired",
Self::Banned => "banned",
Self::Invalid => "invalid",
}
}
fn parse(value: &str) -> StoreResult<Self> {
match value {
"unknown" => Ok(Self::Unknown),
"ready" => Ok(Self::Ready),
"cooldown" => Ok(Self::Cooldown),
"quota_exhausted" => Ok(Self::QuotaExhausted),
"expired" => Ok(Self::Expired),
"banned" => Ok(Self::Banned),
"invalid" => Ok(Self::Invalid),
_ => Err(invalid("unknown availability value")),
}
}
/// 解析数据库 availability 列;字符串映射本体唯一归属 gateway-core。
fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -175,7 +140,7 @@ pub struct ProviderAccountSummary {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -220,7 +185,7 @@ pub struct NewProviderAccount {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -259,9 +224,7 @@ impl NewProviderAccount {
if !self.has_refresh_token && self.next_refresh_at.is_some() {
return Err(invalid("next_refresh_at requires a refresh token"));
}
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown_until must be present exactly for cooldown availability",
));
@@ -394,7 +357,7 @@ impl fmt::Debug for ProviderCredentialUpdate {
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderAccountObservation {
pub account_id: String,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub provider_quota_json: Option<JsonObject>,
@@ -406,7 +369,7 @@ pub struct ProviderAccountObservation {
pub struct ProviderAccountStateUpdate {
pub account_id: String,
pub expected_revision: Revision,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -415,9 +378,7 @@ pub struct ProviderAccountStateUpdate {
impl ProviderAccountStateUpdate {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
@@ -429,9 +390,7 @@ impl ProviderAccountStateUpdate {
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
@@ -485,25 +444,21 @@ pub trait ProviderAccountAdminRepository: Send + Sync {
async fn import_provider_accounts(
&self,
expected_config_revision: Revision,
command: ImportProviderAccounts,
) -> StoreResult<ProviderAccountAdminImport>;
async fn rotate_provider_account(
&self,
expected_config_revision: Revision,
command: RotateProviderAccount,
) -> StoreResult<ProviderAccountAdminRotation>;
async fn set_provider_account_enabled_admin(
&self,
expected_config_revision: Revision,
command: SetProviderAccountEnabled,
) -> StoreResult<Revision>;
async fn delete_provider_accounts_admin(
&self,
expected_config_revision: Revision,
command: DeleteProviderAccounts,
) -> StoreResult<Revision>;
}
@@ -767,7 +722,6 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
async fn import_provider_accounts(
&self,
expected_config_revision: Revision,
command: ImportProviderAccounts,
) -> StoreResult<ProviderAccountAdminImport> {
command.validate()?;
@@ -777,9 +731,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.await
.map_err(|_| postgres_unavailable("begin provider account admin import"))?;
let result = async {
let revision =
bump_config_revision_in_transaction(&mut transaction, expected_config_revision)
.await?;
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
let mut account_ids = Vec::with_capacity(command.accounts.len());
for account in &command.accounts {
account_ids
@@ -798,7 +750,6 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
async fn rotate_provider_account(
&self,
expected_config_revision: Revision,
command: RotateProviderAccount,
) -> StoreResult<ProviderAccountAdminRotation> {
command.scope.validate()?;
@@ -814,9 +765,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.await
.map_err(|_| postgres_unavailable("begin provider account admin rotation"))?;
let result = async {
let config_revision =
bump_config_revision_in_transaction(&mut transaction, expected_config_revision)
.await?;
let config_revision = bump_config_revision_in_transaction(&mut transaction).await?;
let credential_revision = rotate_provider_account_in_transaction(
&mut transaction,
&command.scope,
@@ -841,7 +790,6 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
async fn set_provider_account_enabled_admin(
&self,
expected_config_revision: Revision,
command: SetProviderAccountEnabled,
) -> StoreResult<Revision> {
command.scope.validate()?;
@@ -852,9 +800,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.await
.map_err(|_| postgres_unavailable("begin provider account admin state change"))?;
let result = async {
let revision =
bump_config_revision_in_transaction(&mut transaction, expected_config_revision)
.await?;
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
set_provider_account_enabled_in_transaction(
&mut transaction,
&command.scope,
@@ -872,7 +818,6 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
async fn delete_provider_accounts_admin(
&self,
expected_config_revision: Revision,
command: DeleteProviderAccounts,
) -> StoreResult<Revision> {
command.scope.validate()?;
@@ -883,9 +828,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.await
.map_err(|_| postgres_unavailable("begin provider account admin deletion"))?;
let result = async {
let revision =
bump_config_revision_in_transaction(&mut transaction, expected_config_revision)
.await?;
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
delete_provider_accounts_in_transaction(
&mut transaction,
&command.scope,
@@ -1032,7 +975,6 @@ impl PgAdminAccountStore {
async fn commit_prepared_import(
&self,
expected_config_revision: AdminRevision,
prepared: PreparedCredentialImport,
context: &MutationContext,
action: &str,
@@ -1046,22 +988,19 @@ impl PgAdminAccountStore {
.map_err(|error| admin_store_error(ENTITY, error))?;
let imported = self
.accounts
.import_provider_accounts(
store_revision(expected_config_revision)?,
ImportProviderAccounts {
scope: ProviderAccountAdminScope {
provider_kind: provider_kind.clone(),
},
accounts,
audit: mutation_audit(
context,
action,
"provider_account",
&provider_kind,
vec!["credentials".to_owned()],
),
.import_provider_accounts(ImportProviderAccounts {
scope: ProviderAccountAdminScope {
provider_kind: provider_kind.clone(),
},
)
accounts,
audit: mutation_audit(
context,
action,
"provider_account",
&provider_kind,
vec!["credentials".to_owned()],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
Ok(CredentialImportResult {
@@ -1083,7 +1022,6 @@ impl PgAdminAccountStore {
async fn commit_prepared_rotation(
&self,
expected_config_revision: AdminRevision,
prepared: PreparedCredentialRotationFacts,
context: &MutationContext,
action: &str,
@@ -1094,36 +1032,31 @@ impl PgAdminAccountStore {
};
let rotation = self
.accounts
.rotate_provider_account(
store_revision(expected_config_revision)?,
RotateProviderAccount {
scope,
profile: UpdateProviderAccount {
id: account_id.as_str().to_owned(),
name: prepared.name,
email: prepared.email,
plan_type: prepared.plan_type,
},
credential: ProviderCredentialUpdate {
account_id: account_id.as_str().to_owned(),
expected_revision: store_revision(prepared.expected_credential_revision)?,
provider_credentials_json: provider_document_json(
prepared.provider_material,
)
.map_err(|error| admin_store_error(ENTITY, error))?,
has_refresh_token: prepared.has_refresh_token,
access_token_expires_at: prepared.access_token_expires_at,
next_refresh_at: prepared.next_refresh_at,
},
audit: mutation_audit(
context,
action,
"provider_account",
account_id.as_str(),
vec!["credentials".to_owned()],
),
.rotate_provider_account(RotateProviderAccount {
scope,
profile: UpdateProviderAccount {
id: account_id.as_str().to_owned(),
name: prepared.name,
email: prepared.email,
plan_type: prepared.plan_type,
},
)
credential: ProviderCredentialUpdate {
account_id: account_id.as_str().to_owned(),
expected_revision: store_revision(prepared.expected_credential_revision)?,
provider_credentials_json: provider_document_json(prepared.provider_material)
.map_err(|error| admin_store_error(ENTITY, error))?,
has_refresh_token: prepared.has_refresh_token,
access_token_expires_at: prepared.access_token_expires_at,
next_refresh_at: prepared.next_refresh_at,
},
audit: mutation_audit(
context,
action,
"provider_account",
account_id.as_str(),
vec!["credentials".to_owned()],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
Ok(CredentialMutationResult {
@@ -1251,9 +1184,10 @@ impl AccountStore for PgAdminAccountStore {
.into_iter()
.filter(|account| account.provider_kind == provider_kind.as_str())
.filter(|account| {
query.availability.as_ref().is_none_or(|expected| {
expected.matches(admin_account_availability(account.availability))
})
query
.availability
.as_ref()
.is_none_or(|expected| expected.matches(account.availability))
})
.filter(|account| {
query
@@ -1383,13 +1317,8 @@ impl AccountStore for PgAdminAccountStore {
command: CredentialImportCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialImportResult> {
self.commit_prepared_import(
command.expected_config_revision,
command.prepared,
context,
"import_document",
)
.await
self.commit_prepared_import(command.prepared, context, "import_document")
.await
}
async fn commit_authorization(
@@ -1397,13 +1326,11 @@ impl AccountStore for PgAdminAccountStore {
command: AuthorizationCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult> {
let expected_config_revision = command.pending.expected_config_revision();
match command.credential {
AuthorizationCredentialCommit::Create(credential) => {
let account_id = credential.account_id.clone();
let result = self
.commit_prepared_import(
expected_config_revision,
PreparedCredentialImport {
provider_kind: credential.provider_kind.clone(),
credentials: vec![credential],
@@ -1431,13 +1358,8 @@ impl AccountStore for PgAdminAccountStore {
})
}
AuthorizationCredentialCommit::Reauthorize(prepared) => {
self.commit_prepared_rotation(
expected_config_revision,
prepared,
context,
"reauthorize",
)
.await
self.commit_prepared_rotation(prepared, context, "reauthorize")
.await
}
}
}
@@ -1447,13 +1369,8 @@ impl AccountStore for PgAdminAccountStore {
command: CredentialRotationCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult> {
self.commit_prepared_rotation(
command.expected_config_revision,
command.prepared,
context,
"rotate_credential",
)
.await
self.commit_prepared_rotation(command.prepared, context, "rotate_credential")
.await
}
async fn commit_credential_refresh(
@@ -1461,13 +1378,8 @@ impl AccountStore for PgAdminAccountStore {
command: CredentialRotationCommit,
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult> {
self.commit_prepared_rotation(
command.expected_config_revision,
command.prepared,
context,
"refresh_credential",
)
.await
self.commit_prepared_rotation(command.prepared, context, "refresh_credential")
.await
}
async fn set_account_enabled(
@@ -1478,21 +1390,18 @@ impl AccountStore for PgAdminAccountStore {
let scope = self.required_scope(&command.account_id).await?;
let action = if command.enabled { "enable" } else { "disable" };
self.accounts
.set_provider_account_enabled_admin(
store_revision(command.expected_config_revision)?,
SetProviderAccountEnabled {
scope,
account_id: command.account_id.clone(),
enabled: command.enabled,
audit: mutation_audit(
context,
action,
"provider_account",
&command.account_id,
vec!["enabled".to_owned()],
),
},
)
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
scope,
account_id: command.account_id.clone(),
enabled: command.enabled,
audit: mutation_audit(
context,
action,
"provider_account",
&command.account_id,
vec!["enabled".to_owned()],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)
@@ -1517,20 +1426,17 @@ impl AccountStore for PgAdminAccountStore {
"provider_accounts".to_owned()
};
self.accounts
.delete_provider_accounts_admin(
store_revision(command.expected_config_revision)?,
DeleteProviderAccounts {
scope,
account_ids: command.account_ids,
audit: mutation_audit(
context,
"delete",
"provider_account",
&audit_target,
Vec::new(),
),
},
)
.delete_provider_accounts_admin(DeleteProviderAccounts {
scope,
account_ids: command.account_ids,
audit: mutation_audit(
context,
"delete",
"provider_account",
&audit_target,
Vec::new(),
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)
@@ -1656,12 +1562,12 @@ fn admin_account_status(
) -> AdminAccountStatus {
if !account.enabled {
AdminAccountStatus::Disabled
} else if account.availability == ProviderAccountAvailability::Banned {
} else if account.availability == AccountAvailability::Banned {
AdminAccountStatus::Banned
} else if account.availability == ProviderAccountAvailability::QuotaExhausted {
} else if account.availability == AccountAvailability::QuotaExhausted {
AdminAccountStatus::QuotaExhausted
} else if account.availability == ProviderAccountAvailability::Expired
|| account.availability == ProviderAccountAvailability::Invalid
} else if account.availability == AccountAvailability::Expired
|| account.availability == AccountAvailability::Invalid
|| account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now)
@@ -1786,7 +1692,7 @@ fn admin_account_record(summary: ProviderAccountSummary) -> AdminStoreResult<Acc
access_token_expires_at: summary.access_token_expires_at,
next_refresh_at: summary.next_refresh_at,
enabled: summary.enabled,
availability: admin_account_availability(summary.availability),
availability: summary.availability,
availability_reason: summary.availability_reason,
cooldown_until: summary.cooldown_until,
availability_observed_at: summary.availability_observed_at,
@@ -1811,7 +1717,7 @@ fn prepared_account(credential: PreparedCredentialCreate) -> StoreResult<NewProv
access_token_expires_at: credential.access_token_expires_at,
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
availability: provider_account_availability(credential.availability),
availability: credential.availability,
availability_reason: credential.availability_reason,
cooldown_until: credential.cooldown_until,
availability_observed_at: credential.availability_observed_at,
@@ -1826,34 +1732,6 @@ fn provider_document_json(document: ProviderDocument) -> StoreResult<JsonObject>
)
}
const fn provider_account_availability(
availability: AdminAccountAvailability,
) -> ProviderAccountAvailability {
match availability {
AdminAccountAvailability::Unknown => ProviderAccountAvailability::Unknown,
AdminAccountAvailability::Ready => ProviderAccountAvailability::Ready,
AdminAccountAvailability::Cooldown => ProviderAccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => ProviderAccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => ProviderAccountAvailability::Expired,
AdminAccountAvailability::Banned => ProviderAccountAvailability::Banned,
AdminAccountAvailability::Invalid => ProviderAccountAvailability::Invalid,
}
}
const fn admin_account_availability(
availability: ProviderAccountAvailability,
) -> AdminAccountAvailability {
match availability {
ProviderAccountAvailability::Unknown => AdminAccountAvailability::Unknown,
ProviderAccountAvailability::Ready => AdminAccountAvailability::Ready,
ProviderAccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
ProviderAccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
ProviderAccountAvailability::Expired => AdminAccountAvailability::Expired,
ProviderAccountAvailability::Banned => AdminAccountAvailability::Banned,
ProviderAccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
fn admin_account_usage(usage: ProviderAccountUsageObservation) -> AdminStoreResult<AccountUsage> {
Ok(AccountUsage {
account_id: usage.account_id,
@@ -2268,7 +2146,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(DateTime::<Utc>::from),
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
availability: availability_from_core(account.account.availability()),
availability: account.account.availability(),
availability_reason: None,
cooldown_until: account.account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: Utc::now(),
@@ -2313,15 +2191,23 @@ impl ProviderAccountStore for PgProviderAccountRepository {
&self,
account: &CoreProviderAccountId,
expected_revision: CoreCredentialRevision,
) -> Result<LoadedCredential, CoreStoreError> {
let loaded = self.load_current_credential(account).await?;
if loaded.account.revision().get() != expected_revision.get() {
return Err(CoreStoreError::new(CoreStoreErrorKind::Conflict));
}
Ok(loaded)
}
async fn load_current_credential(
&self,
account: &CoreProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
let record = self
.load_provider_account(account.as_str())
.await
.map_err(core_store_error)?
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
if record.summary.credential_revision.get() != expected_revision.get() {
return Err(CoreStoreError::new(CoreStoreErrorKind::Conflict));
}
Ok(LoadedCredential {
account: core_account_from_summary(record.summary)?,
credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()),
@@ -2470,7 +2356,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
account_id: change.account_id.as_str().to_owned(),
expected_revision: Revision::new(change.expected_revision.get())
.map_err(core_store_error)?,
availability: availability_from_core(change.availability),
availability: change.availability,
availability_reason: change.reason,
cooldown_until: change.cooldown_until.map(DateTime::<Utc>::from),
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
@@ -2581,7 +2467,7 @@ fn core_account_from_summary(
)
.with_runtime_state(
summary.enabled,
availability_to_core(summary.availability),
summary.availability,
summary.cooldown_until.map(Into::into),
)
.with_refresh_schedule(
@@ -2590,30 +2476,6 @@ fn core_account_from_summary(
))
}
const fn availability_to_core(value: ProviderAccountAvailability) -> CoreAccountAvailability {
match value {
ProviderAccountAvailability::Unknown => CoreAccountAvailability::Unknown,
ProviderAccountAvailability::Ready => CoreAccountAvailability::Ready,
ProviderAccountAvailability::Cooldown => CoreAccountAvailability::Cooldown,
ProviderAccountAvailability::QuotaExhausted => CoreAccountAvailability::QuotaExhausted,
ProviderAccountAvailability::Expired => CoreAccountAvailability::Expired,
ProviderAccountAvailability::Banned => CoreAccountAvailability::Banned,
ProviderAccountAvailability::Invalid => CoreAccountAvailability::Invalid,
}
}
const fn availability_from_core(value: CoreAccountAvailability) -> ProviderAccountAvailability {
match value {
CoreAccountAvailability::Unknown => ProviderAccountAvailability::Unknown,
CoreAccountAvailability::Ready => ProviderAccountAvailability::Ready,
CoreAccountAvailability::Cooldown => ProviderAccountAvailability::Cooldown,
CoreAccountAvailability::QuotaExhausted => ProviderAccountAvailability::QuotaExhausted,
CoreAccountAvailability::Expired => ProviderAccountAvailability::Expired,
CoreAccountAvailability::Banned => ProviderAccountAvailability::Banned,
CoreAccountAvailability::Invalid => ProviderAccountAvailability::Invalid,
}
}
fn core_store_error(error: StoreError) -> CoreStoreError {
let kind = match error {
StoreError::Unavailable { .. } => CoreStoreErrorKind::Unavailable,
@@ -2654,7 +2516,7 @@ fn account_summary_from_row(row: sqlx::postgres::PgRow) -> StoreResult<ProviderA
access_token_expires_at: get(&row, "access_token_expires_at")?,
next_refresh_at: get(&row, "next_refresh_at")?,
enabled: get(&row, "enabled")?,
availability: ProviderAccountAvailability::parse(&availability)?,
availability: parse_availability(&availability)?,
availability_reason: get(&row, "availability_reason")?,
cooldown_until: get(&row, "cooldown_until")?,
availability_observed_at: get(&row, "availability_observed_at")?,
@@ -12,7 +12,7 @@ use gateway_core::provider_ports::{
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use crate::{ConflictKind, Revision, StoreError, StoreResult, postgres_unavailable};
use crate::{Revision, StoreError, StoreResult, postgres_unavailable};
#[derive(Clone, PartialEq, Eq)]
pub struct RuntimeSettings {
@@ -111,11 +111,8 @@ impl RuntimeSettingsUpdate {
pub trait RuntimeSettingsRepository: Send + Sync {
async fn load_runtime_settings(&self) -> StoreResult<RuntimeSettings>;
async fn update_runtime_settings(
&self,
expected_revision: Revision,
update: RuntimeSettingsUpdate,
) -> StoreResult<Revision>;
async fn update_runtime_settings(&self, update: RuntimeSettingsUpdate)
-> StoreResult<Revision>;
}
#[derive(Clone)]
@@ -169,7 +166,6 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
async fn update_runtime_settings(
&self,
expected_revision: Revision,
update: RuntimeSettingsUpdate,
) -> StoreResult<Revision> {
let mut transaction = self
@@ -177,9 +173,7 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
.begin()
.await
.map_err(|_| postgres_unavailable("begin runtime settings update"))?;
let revision =
update_runtime_settings_in_transaction(&mut transaction, expected_revision, &update)
.await?;
let revision = update_runtime_settings_in_transaction(&mut transaction, &update).await?;
transaction
.commit()
.await
@@ -228,28 +222,26 @@ pub(crate) async fn load_runtime_settings_in_transaction(
pub(crate) async fn update_runtime_settings_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
expected_revision: Revision,
update: &RuntimeSettingsUpdate,
) -> StoreResult<Revision> {
update.validate()?;
let next = sqlx::query_scalar::<_, i64>(
"update runtime_settings
set config_revision = config_revision + 1,
admin_api_key = $2,
refresh_margin_seconds = $3,
refresh_concurrency = $4,
max_concurrent_per_account = $5,
request_interval_ms = $6,
rotation_strategy = $7,
model_mappings_json = $8,
usage_retention_days = $9,
ops_event_retention_days = $10,
audit_retention_days = $11,
updated_at = now()
where id = 1 and config_revision = $1
returning config_revision",
admin_api_key = $1,
refresh_margin_seconds = $2,
refresh_concurrency = $3,
max_concurrent_per_account = $4,
request_interval_ms = $5,
rotation_strategy = $6,
model_mappings_json = $7,
usage_retention_days = $8,
ops_event_retention_days = $9,
audit_retention_days = $10,
updated_at = now()
where id = 1
returning config_revision",
)
.bind(i64::try_from(expected_revision.get()).map_err(|_| invalid_numeric())?)
.bind(update.admin_api_key.as_deref())
.bind(i64::try_from(update.refresh_margin_seconds).map_err(|_| invalid_numeric())?)
.bind(i64::from(update.refresh_concurrency))
@@ -263,32 +255,28 @@ pub(crate) async fn update_runtime_settings_in_transaction(
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update runtime settings in transaction"))?
.ok_or_else(|| StoreError::Conflict {
.ok_or_else(|| StoreError::NotFound {
entity: "runtime settings",
id: "1".to_owned(),
kind: ConflictKind::StaleRevision,
})?;
Revision::new(u64::try_from(next).map_err(|_| invalid_numeric())?)
}
pub(crate) async fn bump_config_revision_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
expected_revision: Revision,
) -> StoreResult<Revision> {
let next = sqlx::query_scalar::<_, i64>(
"update runtime_settings
set config_revision = config_revision + 1, updated_at = now()
where id = 1 and config_revision = $1
where id = 1
returning config_revision",
)
.bind(i64::try_from(expected_revision.get()).map_err(|_| invalid_numeric())?)
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("bump config revision in transaction"))?
.ok_or_else(|| StoreError::Conflict {
.ok_or_else(|| StoreError::NotFound {
entity: "runtime settings",
id: "1".to_owned(),
kind: ConflictKind::StaleRevision,
})?;
Revision::new(u64::try_from(next).map_err(|_| invalid_numeric())?)
}
@@ -13,7 +13,8 @@ use gateway_core::{
},
error::StoreError as CoreStoreError,
provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderStoreError, ProviderStoreErrorKind,
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown,
ProviderStoreError, ProviderStoreErrorKind,
},
routing::ProviderKind,
};
@@ -104,24 +105,35 @@ impl RedisCredentialCooldownRepository {
let fingerprint = resource_fingerprint("credential cooldown", provider_account_id)?;
Ok(format!("{}:account:{fingerprint}:cooldown", self.namespace))
}
}
#[async_trait]
impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
require_nonempty(
"credential cooldown",
"provider_account_id",
&cooldown.provider_account_id,
)?;
let until_ms = cooldown.cooldown_until.timestamp_millis();
fn scoped_key(
&self,
provider_account_id: &str,
scope: &ProviderCooldownScope,
) -> StoreResult<String> {
let account_fingerprint = resource_fingerprint("credential cooldown", provider_account_id)?;
let scope_fingerprint = resource_fingerprint("credential cooldown scope", scope.value())?;
Ok(format!(
"{}:account:{account_fingerprint}:cooldown:{}:{scope_fingerprint}",
self.namespace,
scope.kind(),
))
}
async fn cache_at_key(
&self,
key: String,
credential_revision: Revision,
cooldown_until: DateTime<Utc>,
) -> StoreResult<bool> {
let until_ms = cooldown_until.timestamp_millis();
if until_ms <= 0 {
return Err(invalid("cooldown expiry must be positive"));
}
let mut connection = self.connection.clone();
let written = Script::new(WRITE_SCRIPT)
.key(self.key(&cooldown.provider_account_id)?)
.arg(cooldown.credential_revision.get())
.key(key)
.arg(credential_revision.get())
.arg(until_ms)
.invoke_async::<i64>(&mut connection)
.await
@@ -129,18 +141,10 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
Ok(written == 1)
}
async fn read_credential_cooldown(
&self,
provider_account_id: &str,
) -> StoreResult<Option<CredentialCooldown>> {
require_nonempty(
"credential cooldown",
"provider_account_id",
provider_account_id,
)?;
async fn read_at_key(&self, key: String) -> StoreResult<Option<(Revision, DateTime<Utc>)>> {
let mut connection = self.connection.clone();
let (present, revision, until_ms): (i64, String, String) = Script::new(READ_SCRIPT)
.key(self.key(provider_account_id)?)
.key(key)
.invoke_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read credential cooldown"))?;
@@ -155,11 +159,59 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
.map_err(|_| invalid("cached cooldown expiry is invalid"))?;
let cooldown_until = DateTime::from_timestamp_millis(until_ms)
.ok_or_else(|| invalid("cached cooldown expiry is invalid"))?;
Ok(Some(CredentialCooldown {
provider_account_id: provider_account_id.to_owned(),
credential_revision: Revision::new(revision)?,
cooldown_until,
}))
Ok(Some((Revision::new(revision)?, cooldown_until)))
}
async fn invalidate_at_key(
&self,
key: String,
through_revision: Revision,
) -> StoreResult<bool> {
let mut connection = self.connection.clone();
let removed = Script::new(INVALIDATE_SCRIPT)
.key(key)
.arg(through_revision.get())
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
Ok(removed == 1)
}
}
#[async_trait]
impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
require_nonempty(
"credential cooldown",
"provider_account_id",
&cooldown.provider_account_id,
)?;
self.cache_at_key(
self.key(&cooldown.provider_account_id)?,
cooldown.credential_revision,
cooldown.cooldown_until,
)
.await
}
async fn read_credential_cooldown(
&self,
provider_account_id: &str,
) -> StoreResult<Option<CredentialCooldown>> {
require_nonempty(
"credential cooldown",
"provider_account_id",
provider_account_id,
)?;
self.read_at_key(self.key(provider_account_id)?)
.await
.map(|value| {
value.map(|(credential_revision, cooldown_until)| CredentialCooldown {
provider_account_id: provider_account_id.to_owned(),
credential_revision,
cooldown_until,
})
})
}
async fn invalidate_credential_cooldown(
@@ -172,14 +224,8 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
"provider_account_id",
provider_account_id,
)?;
let mut connection = self.connection.clone();
let removed = Script::new(INVALIDATE_SCRIPT)
.key(self.key(provider_account_id)?)
.arg(through_revision.get())
.invoke_async::<i64>(&mut connection)
self.invalidate_at_key(self.key(provider_account_id)?, through_revision)
.await
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
Ok(removed == 1)
}
}
@@ -241,6 +287,69 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
.map_err(|_| provider_unavailable("clear credential cooldown"))
})
}
fn put_scoped_if_later(
&self,
cooldown: ProviderScopedCooldown,
) -> futures::future::BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let revision = Revision::new(cooldown.credential_revision().get())
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?;
self.cache_at_key(
self.scoped_key(cooldown.account_id().as_str(), cooldown.scope())
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
revision,
cooldown.until().into(),
)
.await
.map_err(|_| provider_unavailable("cache scoped credential cooldown"))
})
}
fn read_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>>
{
Box::pin(async move {
self.read_at_key(
self.scoped_key(account_id.as_str(), scope)
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
)
.await
.map_err(|_| provider_unavailable("read scoped credential cooldown"))?
.map(|(revision, until)| {
Ok(ProviderScopedCooldown::new(
account_id.clone(),
CredentialRevision::new(revision.get())
.map_err(|_| provider_invalid("decode scoped credential cooldown"))?,
scope.clone(),
until.into(),
))
})
.transpose()
})
}
fn clear_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
through_revision: CredentialRevision,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let revision = Revision::new(through_revision.get())
.map_err(|_| provider_invalid("encode scoped cooldown revision"))?;
self.invalidate_at_key(
self.scoped_key(account_id.as_str(), scope)
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
revision,
)
.await
.map_err(|_| provider_unavailable("clear scoped credential cooldown"))
})
}
}
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
@@ -348,6 +457,13 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore {
.await
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
self.authoritative.load_current_credential(account).await
}
async fn compare_and_swap_credential(
&self,
update: CredentialCasUpdate,
@@ -55,10 +55,7 @@ impl CredentialStateCache {
"provider_account_id",
&self.provider_account_id,
)?;
if !matches!(
self.availability.as_str(),
"unknown" | "ready" | "cooldown" | "quota_exhausted" | "expired" | "banned" | "invalid"
) {
if AccountAvailability::parse(&self.availability).is_none() {
return Err(invalid("availability is not registered"));
}
Ok(())
@@ -60,7 +60,7 @@ end
return 1
"#;
/// Redis-backed, best-effort native response continuation affinity.
/// 基于 Redis 的 best-effort 原生响应 continuation 亲和。
#[derive(Clone)]
pub struct RedisNativeContinuationRepository {
connection: ConnectionManager,
@@ -3,8 +3,9 @@
use std::time::Duration;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderStoreError, ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome, ProviderStoreError,
ProviderStoreErrorKind,
};
use gateway_core::routing::ProviderKind;
use redis::{Script, aio::ConnectionManager};
@@ -31,7 +32,7 @@ redis.call('PEXPIREAT', KEYS[1], expires_at_ms)
return 1
"#;
const TAKE_SCRIPT: &str = r#"
const CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
@@ -51,14 +52,82 @@ end
if not equal_bytes(owner, ARGV[1]) then
return {-1, ''}
end
local clock = redis.call('TIME')
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
local existing_claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if existing_claim ~= false then
local claim_expires_at = redis.call('HGET', KEYS[1], 'claim_expires_at_epoch_millis')
if claim_expires_at ~= false and tonumber(claim_expires_at) > now_ms then
return {-2, ''}
end
redis.call('HDEL', KEYS[1], 'claim_fingerprint', 'claim_expires_at_epoch_millis')
end
local payload = redis.call('HGET', KEYS[1], 'provider_payload')
if payload == false then
return redis.error_reply('OAuth pending flow payload is missing')
end
redis.call('DEL', KEYS[1])
redis.call('HSET', KEYS[1],
'claim_fingerprint', ARGV[2],
'claim_expires_at_epoch_millis', now_ms + tonumber(ARGV[3]))
return {1, payload}
"#;
const RELEASE_CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
for index = 1, length do
local left_byte = string.byte(left, index) or 0
local right_byte = string.byte(right, index) or 0
if left_byte ~= right_byte then
different = different + 1
end
end
return different == 0
end
local owner = redis.call('HGET', KEYS[1], 'owner_fingerprint')
if owner == false then
return 0
end
if not equal_bytes(owner, ARGV[1]) then
return -1
end
local claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if claim == false or not equal_bytes(claim, ARGV[2]) then
return -2
end
redis.call('HDEL', KEYS[1], 'claim_fingerprint', 'claim_expires_at_epoch_millis')
return 1
"#;
const CONSUME_CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
for index = 1, length do
local left_byte = string.byte(left, index) or 0
local right_byte = string.byte(right, index) or 0
if left_byte ~= right_byte then
different = different + 1
end
end
return different == 0
end
local owner = redis.call('HGET', KEYS[1], 'owner_fingerprint')
if owner == false then
return 0
end
if not equal_bytes(owner, ARGV[1]) then
return -1
end
local claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if claim == false or not equal_bytes(claim, ARGV[2]) then
return -2
end
redis.call('DEL', KEYS[1])
return 1
"#;
/// OAuth pending flow 的 Redis 原子持久化能力。
#[derive(Clone)]
pub struct RedisOAuthPendingFlowRepository {
@@ -118,39 +187,115 @@ impl OAuthPendingFlowPort for RedisOAuthPendingFlowRepository {
})
}
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>> {
claim: &'a OAuthPendingBinding,
claim_ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let claim_ttl = ttl_millis(claim_ttl)
.map_err(|_| provider_invalid("encode OAuth pending claim TTL"))?;
let mut connection = self.connection.clone();
let (status, payload): (i64, Vec<u8>) = Script::new(TAKE_SCRIPT)
let (status, payload): (i64, Vec<u8>) = Script::new(CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.arg(claim_ttl)
.invoke_async(&mut connection)
.await
.map_err(|_| provider_unavailable("take OAuth pending flow"))?;
.map_err(|_| provider_unavailable("claim OAuth pending flow"))?;
match status {
0 => Ok(OAuthPendingTakeOutcome::NotFound),
-1 => Ok(OAuthPendingTakeOutcome::OwnerMismatch),
0 => Ok(OAuthPendingClaimOutcome::NotFound),
-1 => Ok(OAuthPendingClaimOutcome::OwnerMismatch),
-2 => Ok(OAuthPendingClaimOutcome::InProgress),
1 => {
let value: serde_json::Value = serde_json::from_slice(&payload)
.map_err(|_| provider_invalid("decode OAuth pending payload"))?;
let serde_json::Value::Object(fields) = value else {
return Err(provider_invalid("decode OAuth pending payload"));
};
Ok(OAuthPendingTakeOutcome::Taken(
Ok(OAuthPendingClaimOutcome::Claimed(
gateway_core::engine::credential::OpaqueProviderData::new(fields),
))
}
_ => Err(provider_invalid("decode OAuth pending outcome")),
_ => Err(provider_invalid("decode OAuth pending claim outcome")),
}
})
}
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>>
{
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let mut connection = self.connection.clone();
let status = Script::new(RELEASE_CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| provider_unavailable("release OAuth pending claim"))?;
Ok(match status {
0 => OAuthPendingReleaseOutcome::NotFound,
-1 => OAuthPendingReleaseOutcome::OwnerMismatch,
-2 => OAuthPendingReleaseOutcome::ClaimMismatch,
1 => OAuthPendingReleaseOutcome::Released,
_ => {
return Err(provider_invalid(
"decode OAuth pending claim release outcome",
));
}
})
})
}
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>>
{
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let mut connection = self.connection.clone();
let status = Script::new(CONSUME_CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| provider_unavailable("consume OAuth pending claim"))?;
Ok(match status {
0 => OAuthPendingConsumeOutcome::NotFound,
-1 => OAuthPendingConsumeOutcome::OwnerMismatch,
-2 => OAuthPendingConsumeOutcome::ClaimMismatch,
1 => OAuthPendingConsumeOutcome::Consumed,
_ => {
return Err(provider_invalid(
"decode OAuth pending claim consume outcome",
));
}
})
})
}
}
fn provider_scoped_fingerprint(
@@ -58,12 +58,6 @@ impl TestDatabase {
.execute(&mut *migration)
.await
.expect("apply terminal migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0002_retention_delete_indexes.sql"
))
.execute(&mut *migration)
.await
.expect("apply retention index migration");
migration.commit().await.expect("commit terminal migration");
Some(Self {
admin,
@@ -86,7 +80,7 @@ impl TestDatabase {
}
#[tokio::test]
async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
async fn connect_and_migrate_should_apply_initial_migration_once_and_reopen_cleanly() {
let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else {
return;
};
@@ -110,7 +104,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
.to_string();
let first = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default())
.await
.expect("apply 0001 through production migrator");
.expect("apply initial migration through production migrator");
let first_table_count = sqlx::query_scalar::<_, i64>(
"select count(*) from information_schema.tables where table_schema = 'public'",
)
@@ -137,7 +131,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (8, 2));
assert_eq!((first_table_count, migration_count), (8, 1));
}
#[test]
@@ -3,7 +3,7 @@ use std::time::{Duration, SystemTime};
use chrono::{TimeDelta, Utc};
use gateway_admin::{
model::{
MutationActor, MutationContext, PageSize, Revision as AdminRevision,
MutationActor, MutationContext, PageSize,
accounts::{
AccountListQuery, AccountSort, AccountSortField, AccountStatus,
AccountUsageWindowQuery, DeleteAccounts, SetAccountEnabled as AdminSetAccountEnabled,
@@ -14,12 +14,12 @@ use gateway_admin::{
CredentialAvailabilityFilter, CredentialListQuery, CredentialListWindow,
},
},
ports::store::{AccountStore, AdminStoreErrorKind},
ports::store::AccountStore,
};
use gateway_core::engine::credential::{
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, OpaqueProviderData,
PlaintextCredential, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
QuotaObservation, QuotaWriteOutcome,
AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
};
use gateway_core::routing::ProviderKind;
use gateway_store::{
@@ -27,9 +27,9 @@ use gateway_store::{
postgres::{
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountAvailability,
ProviderAccountObservation, ProviderAccountRepository, ProviderCredentialUpdate,
RotateProviderAccount, SetProviderAccountEnabled, UpdateProviderAccount,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountObservation,
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
SetProviderAccountEnabled, UpdateProviderAccount,
},
};
use serde_json::json;
@@ -123,15 +123,15 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
let mut beta = account("acct_beta", "user-beta");
beta.provider_kind = "xai".to_owned();
beta.email = Some("beta@example.invalid".to_owned());
beta.availability = ProviderAccountAvailability::Banned;
beta.availability = AccountAvailability::Banned;
let mut charlie = account("acct_charlie", "user-charlie");
charlie.email = Some("charlie@example.invalid".to_owned());
let mut invalid = account("acct_invalid", "user-invalid");
invalid.email = Some("invalid@example.invalid".to_owned());
invalid.availability = ProviderAccountAvailability::Invalid;
invalid.availability = AccountAvailability::Invalid;
let mut quota_exhausted = account("acct_quota_exhausted", "user-quota-exhausted");
quota_exhausted.email = Some("quota-exhausted@example.invalid".to_owned());
quota_exhausted.availability = ProviderAccountAvailability::QuotaExhausted;
quota_exhausted.availability = AccountAvailability::QuotaExhausted;
for account in [alpha, beta, charlie, invalid, quota_exhausted] {
repository
.insert_provider_account(account)
@@ -292,10 +292,10 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
);
credential.provider_kind = "xai".to_owned();
credential.availability = match index {
0 => ProviderAccountAvailability::Expired,
1 => ProviderAccountAvailability::Banned,
2 => ProviderAccountAvailability::Invalid,
_ => ProviderAccountAvailability::Ready,
0 => AccountAvailability::Expired,
1 => AccountAvailability::Banned,
2 => AccountAvailability::Invalid,
_ => AccountAvailability::Ready,
};
repository
.insert_provider_account(credential)
@@ -525,7 +525,6 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
let revision = store
.set_account_enabled(
AdminSetAccountEnabled {
expected_config_revision: AdminRevision::new(1).expect("initial revision"),
account_id: "acct_terminal_mutation".to_owned(),
enabled: false,
},
@@ -542,32 +541,9 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
.expect("read disabled state");
assert!(!enabled);
let error = store
.delete_accounts(
DeleteAccounts {
expected_config_revision: AdminRevision::new(1).expect("stale revision"),
account_ids: vec!["acct_terminal_mutation".to_owned()],
},
&context,
)
.await
.expect_err("stale delete must roll back");
assert_eq!(error.kind(), AdminStoreErrorKind::StaleRevision);
assert_eq!(current_revision(&database.pool).await, 2);
assert_eq!(
account_count(&database.pool, "acct_terminal_mutation").await,
1
);
let audit_count: i64 = sqlx::query_scalar("select count(*) from admin_audit_events")
.fetch_one(&database.pool)
.await
.expect("count audit rows after rollback");
assert_eq!(audit_count, 1);
let revision = store
.delete_accounts(
DeleteAccounts {
expected_config_revision: revision,
account_ids: vec!["acct_terminal_mutation".to_owned()],
},
&context,
@@ -616,7 +592,6 @@ async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
let revision = PgAdminAccountStore::new(database.pool.clone())
.delete_accounts(
DeleteAccounts {
expected_config_revision: AdminRevision::new(1).expect("initial revision"),
account_ids: vec![
"acct_enabled_delete_a".to_owned(),
"acct_enabled_delete_b".to_owned(),
@@ -649,18 +624,14 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
provider_kind: "openai".to_owned(),
};
let imported = repository
.import_provider_accounts(
Revision::new(1).expect("initial revision"),
ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_admin_upsert", "user-admin-upsert")],
audit: audit("audit_admin_upsert_create", "import", "acct_admin_upsert"),
},
)
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_admin_upsert", "user-admin-upsert")],
audit: audit("audit_admin_upsert_create", "import", "acct_admin_upsert"),
})
.await
.expect("create imported account");
assert_eq!(imported.account_ids, ["acct_admin_upsert"]);
let revision = imported.config_revision;
sqlx::query(
"update provider_accounts
set provider_quota_json = '{}'::jsonb, quota_observed_at = now(), updated_at = now()
@@ -673,17 +644,14 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
let mut updated = account("acct_admin_reimport", "user-admin-upsert");
updated.name = "updated import".to_owned();
updated.provider_credentials_json = credential_json("updated-import-secret");
updated.availability = ProviderAccountAvailability::Banned;
updated.availability = AccountAvailability::Banned;
updated.availability_reason = Some("upstream_account_deactivated".to_owned());
let imported = repository
.import_provider_accounts(
revision,
ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![updated],
audit: audit("audit_admin_upsert_update", "import", "acct_admin_upsert"),
},
)
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![updated],
audit: audit("audit_admin_upsert_update", "import", "acct_admin_upsert"),
})
.await
.expect("update the same imported identity");
assert_eq!(imported.account_ids, ["acct_admin_upsert"]);
@@ -715,14 +683,11 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
assert_eq!(row.5.as_deref(), Some("upstream_account_deactivated"));
repository
.import_provider_accounts(
revision,
ImportProviderAccounts {
scope,
accounts: vec![account("acct_admin_upsert", "another-user")],
audit: audit("audit_admin_upsert_rebind", "import", "acct_admin_upsert"),
},
)
.import_provider_accounts(ImportProviderAccounts {
scope,
accounts: vec![account("acct_admin_upsert", "another-user")],
audit: audit("audit_admin_upsert_rebind", "import", "acct_admin_upsert"),
})
.await
.expect_err("an existing account ID must not be rebound");
assert_eq!(current_revision(&database.pool).await, 3);
@@ -751,7 +716,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
access_token_expires_at: Some(Utc::now() + TimeDelta::minutes(5)),
next_refresh_at: None,
enabled: true,
availability: ProviderAccountAvailability::Ready,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
@@ -804,6 +769,25 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
assert_eq!(row.3["access_token"], "after-secret");
assert_eq!(row.4, 2);
assert!(
repository
.load_credential(
&account_id,
CredentialRevision::new(1).expect("stale credential revision"),
)
.await
.is_err()
);
let current = repository
.load_current_credential(&account_id)
.await
.expect("load current credential without a caller revision");
assert_eq!(current.account.revision().get(), 2);
assert_eq!(
current.credential.expose_to_provider()["access_token"],
"after-secret"
);
let stale = CredentialCasUpdate::new(
account_id.clone(),
CredentialRevision::new(1).expect("stale credential revision"),
@@ -856,21 +840,17 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
};
let mut cooldown_account = account("acct_admin_b", "user-admin-b");
cooldown_account.availability = ProviderAccountAvailability::Cooldown;
cooldown_account.availability = AccountAvailability::Cooldown;
cooldown_account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(10));
let imported = repository
.import_provider_accounts(
Revision::new(1).expect("initial revision"),
ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_admin_a", "user-admin-a"), cooldown_account],
audit: audit("audit_account_batch", "import_batch", "provider_accounts"),
},
)
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_admin_a", "user-admin-a"), cooldown_account],
audit: audit("audit_account_batch", "import_batch", "provider_accounts"),
})
.await
.expect("import provider accounts");
let revision = imported.config_revision;
assert_eq!(revision.get(), 2);
assert_eq!(imported.config_revision.get(), 2);
let ready: (bool, String, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
"select enabled, availability, cooldown_until from provider_accounts where id = $1",
)
@@ -891,21 +871,18 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
assert!(cooldown.2.is_some());
repository
.import_provider_accounts(
revision,
ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![
account("acct_admin_transient", "user-admin-transient"),
account("acct_admin_a", "user-admin-duplicate"),
],
audit: audit(
"audit_account_failed_batch",
"import_batch",
"provider_accounts",
),
},
)
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![
account("acct_admin_transient", "user-admin-transient"),
account("acct_admin_a", "user-admin-duplicate"),
],
audit: audit(
"audit_account_failed_batch",
"import_batch",
"provider_accounts",
),
})
.await
.expect_err("duplicate batch row must roll back the entire import");
assert_eq!(current_revision(&database.pool).await, 2);
@@ -918,15 +895,12 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
provider_kind: "xai".to_owned(),
};
let wrong_scope_error = repository
.rotate_provider_account(
revision,
RotateProviderAccount {
scope: wrong_scope,
profile: profile("acct_admin_a", "wrong scope"),
credential: credential_update("acct_admin_a", 1, "wrong-scope-secret"),
audit: audit("audit_wrong_scope", "rotate", "acct_admin_a"),
},
)
.rotate_provider_account(RotateProviderAccount {
scope: wrong_scope,
profile: profile("acct_admin_a", "wrong scope"),
credential: credential_update("acct_admin_a", 1, "wrong-scope-secret"),
audit: audit("audit_wrong_scope", "rotate", "acct_admin_a"),
})
.await
.expect_err("Provider endpoint must not rotate another Provider account");
assert!(matches!(
@@ -939,32 +913,27 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
assert_eq!(current_revision(&database.pool).await, 2);
let rotation = repository
.rotate_provider_account(
revision,
RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "rotated account"),
credential: credential_update("acct_admin_a", 1, "rotated-secret"),
audit: audit("audit_account_rotate", "rotate", "acct_admin_a"),
},
)
.rotate_provider_account(RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "rotated account"),
credential: credential_update("acct_admin_a", 1, "rotated-secret"),
audit: audit("audit_account_rotate", "rotate", "acct_admin_a"),
})
.await
.expect("rotate provider account");
assert_eq!(rotation.config_revision.get(), 3);
assert_eq!(rotation.credential_revision.get(), 2);
let revision = repository
.set_provider_account_enabled_admin(
rotation.config_revision,
SetProviderAccountEnabled {
scope: scope.clone(),
account_id: "acct_admin_a".to_owned(),
enabled: false,
audit: audit("audit_account_disable", "disable", "acct_admin_a"),
},
)
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
scope: scope.clone(),
account_id: "acct_admin_a".to_owned(),
enabled: false,
audit: audit("audit_account_disable", "disable", "acct_admin_a"),
})
.await
.expect("disable provider account");
assert_eq!(revision.get(), 4);
let exports = repository
.export_provider_accounts(
scope.clone(),
@@ -977,14 +946,11 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
assert!(!format!("{exports:?}").contains("rotated-secret"));
let revision = repository
.delete_provider_accounts_admin(
revision,
DeleteProviderAccounts {
scope,
account_ids: vec!["acct_admin_a".to_owned(), "acct_admin_b".to_owned()],
audit: audit("audit_account_delete", "delete", "provider_accounts"),
},
)
.delete_provider_accounts_admin(DeleteProviderAccounts {
scope,
account_ids: vec!["acct_admin_a".to_owned(), "acct_admin_b".to_owned()],
audit: audit("audit_account_delete", "delete", "provider_accounts"),
})
.await
.expect("delete selected accounts regardless of enabled state");
assert_eq!(revision.get(), 5);
@@ -1014,7 +980,7 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
next_refresh_at: None,
enabled: true,
availability: ProviderAccountAvailability::Ready,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
@@ -1147,7 +1113,7 @@ fn provider_credentials_are_redacted_from_debug() {
fn cooldown_observation_requires_expiry() {
let observation = ProviderAccountObservation {
account_id: "account-1".to_owned(),
availability: ProviderAccountAvailability::Cooldown,
availability: AccountAvailability::Cooldown,
availability_reason: None,
cooldown_until: None,
provider_quota_json: None,
@@ -13,6 +13,7 @@ use gateway_core::{
QuotaObservation, QuotaWriteOutcome,
},
error::{StoreError as CoreStoreError, StoreErrorKind},
provider_ports::{ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown},
routing::ProviderKind,
};
use gateway_store::{
@@ -73,6 +74,13 @@ impl ProviderAccountStore for CooldownTestAccountStore {
unreachable!("cooldown tests do not load credentials")
}
async fn load_current_credential(
&self,
_account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
unreachable!("cooldown tests do not load credentials")
}
async fn compare_and_swap_credential(
&self,
_update: CredentialCasUpdate,
@@ -523,6 +531,109 @@ async fn credential_cooldown_read_removes_expired_grace_key() {
assert!(namespace_keys(&mut connection, &namespace).await.is_empty());
}
#[tokio::test]
async fn scoped_cooldown_isolated_by_model_and_revision_fenced() {
let Some((repository, mut connection, namespace)) = repository().await else {
return;
};
let account_id = ProviderAccountId::new("acct_scoped_cooldown").expect("valid account ID");
let revision = CredentialRevision::new(2).expect("positive revision");
let model_a = ProviderCooldownScope::upstream_model(
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
);
let model_b = ProviderCooldownScope::upstream_model(
gateway_core::routing::UpstreamModelId::new("grok-4.6").expect("model"),
);
let until_a: SystemTime = millisecond_precision(Utc::now() + Duration::seconds(30)).into();
let until_b: SystemTime = millisecond_precision(Utc::now() + Duration::seconds(60)).into();
assert!(
repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
revision,
model_a.clone(),
until_a,
))
.await
.expect("cache model A cooldown")
);
assert!(
!repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
CredentialRevision::new(1).expect("stale revision"),
model_a.clone(),
until_b,
))
.await
.expect("reject stale model A cooldown")
);
assert!(
repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
revision,
model_b.clone(),
until_b,
))
.await
.expect("cache model B cooldown")
);
assert_eq!(
repository
.read_scoped(&account_id, &model_a)
.await
.expect("read model A cooldown")
.expect("model A cooldown")
.until(),
until_a
);
assert_eq!(
repository
.read_scoped(&account_id, &model_b)
.await
.expect("read model B cooldown")
.expect("model B cooldown")
.until(),
until_b
);
assert!(
!repository
.clear_scoped(
&account_id,
&model_a,
CredentialRevision::new(1).expect("stale revision"),
)
.await
.expect("fence stale model A invalidation")
);
assert!(
repository
.clear_scoped(&account_id, &model_a, revision)
.await
.expect("clear model A cooldown")
);
assert!(
repository
.read_scoped(&account_id, &model_a)
.await
.expect("read cleared model A cooldown")
.is_none()
);
assert!(
repository
.read_scoped(&account_id, &model_b)
.await
.expect("read retained model B cooldown")
.is_some()
);
let keys = namespace_keys(&mut connection, &namespace).await;
assert_eq!(keys.len(), 1);
assert!(!keys[0].contains(account_id.as_str()));
assert!(!keys[0].contains(model_b.value()));
}
fn cooldown(account_id: &str, revision: u64, seconds: i64) -> CredentialCooldown {
CredentialCooldown {
provider_account_id: account_id.to_owned(),
@@ -3,8 +3,8 @@ use std::{sync::Arc, time::Duration};
use gateway_core::{
engine::credential::OpaqueProviderData,
provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome,
OAuthPendingConsumeOutcome, OAuthPendingFlowPort, OAuthPendingPutOutcome,
},
routing::ProviderKind,
};
@@ -23,6 +23,7 @@ async fn codex_pending_flow_wrong_owner_does_not_consume_it() {
OAuthPendingBinding::try_new(format!("flow-{}", Uuid::new_v4())).expect("flow binding");
let owner = OAuthPendingBinding::try_new("owner-a").expect("owner binding");
let other_owner = OAuthPendingBinding::try_new("owner-b").expect("other owner binding");
let claim = OAuthPendingBinding::try_new("claim-a").expect("claim binding");
let mut payload = serde_json::Map::new();
payload.insert(
"state".to_owned(),
@@ -77,24 +78,37 @@ async fn codex_pending_flow_wrong_owner_does_not_consume_it() {
assert!(ttl > 0 && ttl <= 60_000);
assert_eq!(
repository
.take_if_owner(&provider, &flow, &other_owner)
.claim_if_owner(
&provider,
&flow,
&other_owner,
&claim,
Duration::from_secs(30),
)
.await
.expect("reject wrong owner"),
OAuthPendingTakeOutcome::OwnerMismatch
OAuthPendingClaimOutcome::OwnerMismatch
);
assert!(matches!(
repository
.take_if_owner(&provider, &flow, &owner)
.claim_if_owner(&provider, &flow, &owner, &claim, Duration::from_secs(30),)
.await
.expect("consume right owner"),
OAuthPendingTakeOutcome::Taken(_)
.expect("claim right owner"),
OAuthPendingClaimOutcome::Claimed(_)
));
assert_eq!(
repository
.take_if_owner(&provider, &flow, &owner)
.consume_claim(&provider, &flow, &owner, &claim)
.await
.expect("consume right owner claim"),
OAuthPendingConsumeOutcome::Consumed
);
assert_eq!(
repository
.claim_if_owner(&provider, &flow, &owner, &claim, Duration::from_secs(30),)
.await
.expect("flow is one shot"),
OAuthPendingTakeOutcome::NotFound
OAuthPendingClaimOutcome::NotFound
);
}
+4 -4
View File
@@ -10,19 +10,19 @@ autotests = false
[dependencies]
async-stream = "0.3.6"
async-trait = "0.1.89"
base64 = "0.22.1"
base64 = "0.23.0"
bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["clock", "serde"] }
cookie = "0.18.1"
crypto_box = { version = "0.9.1", features = ["seal"] }
ed25519-dalek = { version = "2.2.0", features = ["pkcs8"] }
ed25519-dalek = { version = "3.0.0", features = ["alloc", "pkcs8"] }
futures = "0.3.32"
gateway-admin = { path = "../../gateway-admin" }
gateway-core = { path = "../../gateway-core" }
gateway-protocol = { path = "../../gateway-protocol" }
getrandom = "0.3.4"
getrandom = "0.4.3"
hex = "0.4.3"
jsonwebtoken = { version = "=10.3.0", default-features = false, features = ["aws_lc_rs", "use_pem"] }
jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs", "use_pem"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
roxmltree = "0.21.1"
semver = "1.0.27"
+215 -134
View File
@@ -5,9 +5,7 @@ use std::time::SystemTime;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_admin::model::accounts::{
AccountAvailability as AdminAccountAvailability, AccountRecord,
};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireAttribute,
DashboardWireProfile, DashboardWireTarget, DecimalAmount, DesktopReleaseStatus,
@@ -27,15 +25,15 @@ use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use secrecy::{ExposeSecret as _, SecretString};
@@ -50,10 +48,11 @@ use crate::credential::{
CodexAccountQuotaSnapshot, CodexCredentialAdmin, CodexCredentialAdminError,
CodexCredentialAdminService, CodexCredentialCatalogError, CodexCredentialCatalogService,
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexPendingAuthorization,
CodexQuotaWindowKind, CodexQuotaWindowRole, CompleteCodexOAuthAuthorization,
CompletedCodexOAuthCredential, ExportManagedCodexCredential, RotateManagedCodexCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization, refresh_time,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CodexQuotaWindowKind, CodexQuotaWindowRole,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
RotateManagedCodexCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
refresh_time,
};
use crate::transport::CodexWebSocketPool;
use crate::transport::openai_billing_breakdown;
@@ -62,7 +61,7 @@ use crate::transport::profile::{
};
const PROVIDER_NAME: &str = "openai";
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 2;
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 3;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
@@ -147,7 +146,7 @@ impl OpenAiAdminProvider {
}
let current = self
.accounts
.load_credential(existing.id(), existing.revision())
.load_current_credential(existing.id())
.await
.map_err(map_store_error)?;
incoming.credential = CodexCredentialCodec::preserve_installation_id(
@@ -253,11 +252,22 @@ impl ProviderAdmin for OpenAiAdminProvider {
command.document.into_provider_data().into_inner(),
))
.await
.inspect_err(|error| {
log_import_failure("prepare_document", credential_admin_error_code(*error));
})
.map_err(map_credential_admin_error)?;
let observed_at = Utc::now();
let mut credentials = Vec::with_capacity(prepared.accounts().len());
for account in prepared.into_accounts() {
let account = self.preserve_existing_installation_id(account).await?;
let account = self
.preserve_existing_installation_id(account)
.await
.inspect_err(|error| {
log_import_failure(
"preserve_installation_id",
provider_admin_error_code(error.kind()),
);
})?;
credentials.push(prepared_create(account, observed_at)?);
}
Ok(PreparedCredentialImport {
@@ -289,6 +299,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
&self,
command: CompleteAuthorization,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let request_id = command.context.request_id.clone();
let binding = AuthorizationOwnerBinding::from_context(&command.context);
let completed = self
.oauth
@@ -298,30 +309,46 @@ impl ProviderAdmin for OpenAiAdminProvider {
callback_url: SecretString::from(command.callback_url),
})
.await
.inspect_err(|error| {
tracing::warn!(
request_id = %request_id,
oauth_stage = "complete_authorization",
oauth_error = oauth_error_code(error),
"OpenAI OAuth authorization completion failed"
);
})
.map_err(map_oauth_error)?;
let credential = match completed.credential {
CompletedCodexOAuthCredential::Create(credential) => {
PreparedAuthorizationCredential::Create(prepared_create(credential, Utc::now())?)
let (mutation, completed_credential, authorization_guard) = completed.into_parts();
let credential = match (mutation.target(), completed_credential) {
(
AuthorizationMutationTarget::Create { .. },
CompletedCodexOAuthCredential::Create(credential),
) => {
prepared_create(credential, Utc::now()).map(PreparedAuthorizationCredential::Create)
}
CompletedCodexOAuthCredential::Reauthorize(credential) => {
let provider_kind = match completed.mutation.target() {
AuthorizationMutationTarget::Reauthorize { .. } => {
completed.mutation.provider_kind().clone()
}
AuthorizationMutationTarget::Create { .. } => {
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
}
};
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
credential,
provider_kind,
)?)
(
AuthorizationMutationTarget::Reauthorize { .. },
CompletedCodexOAuthCredential::Reauthorize(credential),
) => prepared_rotation(credential, mutation.provider_kind().clone())
.map(PreparedAuthorizationCredential::Reauthorize),
_ => Err(provider_admin_error(ProviderAdminErrorKind::Internal)),
};
let credential = match credential {
Ok(credential) => credential,
Err(error) => {
if let Err(settlement_error) = authorization_guard.abort().await {
tracing::warn!(
request_id = %request_id,
settlement_error = %settlement_error,
"OpenAI OAuth claim release failed after credential preparation"
);
return Err(provider_admin_error(ProviderAdminErrorKind::Unavailable));
}
return Err(error);
}
};
Ok(PreparedAuthorizationCommit {
pending: completed.mutation,
credential,
})
Ok(PreparedAuthorizationCommit::new(mutation, credential)
.with_authorization_guard(authorization_guard))
}
async fn prepare_rotation(
@@ -329,16 +356,11 @@ impl ProviderAdmin for OpenAiAdminProvider {
command: PrepareCredentialRotation,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
validate_account_record(&command.account, &self.provider_kind)?;
if command.account.credential_revision != command.expected_credential_revision {
return Err(provider_admin_error(ProviderAdminErrorKind::Conflict));
}
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let expected_revision = CredentialRevision::new(command.expected_credential_revision.get())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_credential(&account_id, expected_revision)
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
@@ -398,11 +420,17 @@ impl ProviderAdmin for OpenAiAdminProvider {
validate_account_record(&command.account, &self.provider_kind)?;
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let revision = CredentialRevision::new(command.account.credential_revision.get())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
return Err(provider_admin_error(ProviderAdminErrorKind::Conflict));
}
let prepared = self
.credentials
.manual_refresh(account_id, revision)
.manual_refresh(current)
.await
.map_err(map_credential_admin_error)?;
prepared_rotation(prepared, command.account.provider_kind)
@@ -519,7 +547,7 @@ fn prepared_create(
account,
credential,
} = prepared;
let availability = admin_availability(account.availability());
let availability = account.availability();
Ok(PreparedCredentialCreate {
account_id: account.id().clone(),
provider_kind: account.provider().clone(),
@@ -639,7 +667,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
)
.with_runtime_state(
account.enabled,
core_availability(account.availability),
account.availability,
account.cooldown_until.map(SystemTime::from),
)
.with_refresh_schedule(
@@ -690,44 +718,9 @@ fn dashboard_desktop_release(
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
&& account.email() == record.email.as_deref()
&& account.upstream_user_id() == record.upstream_user_id
&& account.upstream_account_id() == record.upstream_account_id.as_deref()
&& account.plan_type() == record.plan_type.as_deref()
&& account.authentication_kind() == record.authentication_kind
&& account.enabled() == record.enabled
&& admin_availability(account.availability()) == record.availability
&& account.cooldown_until().map(DateTime::<Utc>::from) == record.cooldown_until
&& account.access_token_expires_at().map(DateTime::<Utc>::from)
== record.access_token_expires_at
&& account.next_refresh_at().map(DateTime::<Utc>::from) == record.next_refresh_at
&& account.has_refresh_token() == record.has_refresh_token
}
const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability {
match value {
AccountAvailability::Unknown => AdminAccountAvailability::Unknown,
AccountAvailability::Ready => AdminAccountAvailability::Ready,
AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
AccountAvailability::Expired => AdminAccountAvailability::Expired,
AccountAvailability::Banned => AdminAccountAvailability::Banned,
AccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability {
match value {
AdminAccountAvailability::Unknown => AccountAvailability::Unknown,
AdminAccountAvailability::Ready => AccountAvailability::Ready,
AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => AccountAvailability::Expired,
AdminAccountAvailability::Banned => AccountAvailability::Banned,
AdminAccountAvailability::Invalid => AccountAvailability::Invalid,
}
}
fn empty_quota() -> ProviderQuota {
@@ -903,21 +896,86 @@ impl CodexOAuthPendingStore for OpenAiOAuthPendingStore {
}
}
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError> {
claim_ref: &str,
claim_ttl: std::time::Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
let outcome = self
.port
.claim_if_owner(&self.provider_kind, &flow, &owner, &claim, claim_ttl)
.await
.map_err(map_pending_store_error)?;
match outcome {
OAuthPendingClaimOutcome::Claimed(payload) => match decode_pending(payload) {
Ok(pending) => Ok(CodexOAuthPendingClaimOutcome::Claimed(Box::new(pending))),
Err(error) => match self
.port
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingReleaseOutcome::Released => Err(error),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => {
Err(CodexOAuthPendingStoreError::Unavailable)
}
},
},
OAuthPendingClaimOutcome::NotFound | OAuthPendingClaimOutcome::OwnerMismatch => {
Ok(CodexOAuthPendingClaimOutcome::NotFound)
}
OAuthPendingClaimOutcome::InProgress => Ok(CodexOAuthPendingClaimOutcome::InProgress),
}
}
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.port
.take_if_owner(&self.provider_kind, &flow, &owner)
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingTakeOutcome::Taken(payload) => decode_pending(payload).map(Some),
OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => Ok(None),
OAuthPendingReleaseOutcome::Released => Ok(true),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => Ok(false),
}
}
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.port
.consume_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingConsumeOutcome::Consumed => Ok(true),
OAuthPendingConsumeOutcome::NotFound
| OAuthPendingConsumeOutcome::OwnerMismatch
| OAuthPendingConsumeOutcome::ClaimMismatch => Ok(false),
}
}
}
@@ -934,7 +992,6 @@ struct PendingDocument {
nonce: String,
code_verifier: String,
reauthorization_account_id: Option<String>,
reauthorization_credential_revision: Option<u64>,
mutation: PendingMutationDocument,
}
@@ -942,7 +999,6 @@ struct PendingDocument {
#[serde(deny_unknown_fields)]
struct PendingMutationDocument {
schema_version: u64,
expected_config_revision: u64,
provider_kind: String,
target: PendingTargetDocument,
owner: PendingOwnerDocument,
@@ -952,13 +1008,8 @@ struct PendingMutationDocument {
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
name: String,
},
Reauthorize {
account_id: String,
expected_credential_revision: u64,
},
Create { name: String },
Reauthorize { account_id: String },
}
#[derive(Deserialize)]
@@ -1002,15 +1053,9 @@ fn encode_pending(pending: &CodexPendingAuthorization) -> Map<String, Value> {
);
document.insert(
"reauthorization_account_id".to_owned(),
pending.reauthorization().map_or(Value::Null, |target| {
Value::String(target.account_id().to_string())
}),
);
document.insert(
"reauthorization_credential_revision".to_owned(),
pending.reauthorization().map_or(Value::Null, |target| {
Value::Number(Number::from(target.credential_revision().get()))
}),
pending
.reauthorization()
.map_or(Value::Null, |target| Value::String(target.to_string())),
);
document.insert(
"mutation".to_owned(),
@@ -1025,10 +1070,6 @@ fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map<String, Value
"schema_version".to_owned(),
Value::Number(Number::from(PENDING_DOCUMENT_SCHEMA_VERSION)),
);
document.insert(
"expected_config_revision".to_owned(),
Value::Number(Number::from(mutation.expected_config_revision().get())),
);
document.insert(
"provider_kind".to_owned(),
Value::String(mutation.provider_kind().as_str().to_owned()),
@@ -1055,19 +1096,12 @@ fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
);
document.insert(
"expected_credential_revision".to_owned(),
Value::Number(Number::from(expected_credential_revision.get())),
);
}
}
document
@@ -1108,7 +1142,6 @@ fn decode_pending(
nonce: SecretString::from(document.nonce),
code_verifier: SecretString::from(document.code_verifier),
reauthorization_account_id: document.reauthorization_account_id,
reauthorization_credential_revision: document.reauthorization_credential_revision,
mutation: decode_mutation(document.mutation)?,
})
}
@@ -1119,21 +1152,16 @@ fn decode_mutation(
if document.schema_version != PENDING_DOCUMENT_SCHEMA_VERSION {
return Err(CodexOAuthPendingStoreError::InvalidValue);
}
let expected_config_revision = Revision::new(document.expected_config_revision)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let provider_kind = ProviderKind::new(document.provider_kind)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let target = match document.target {
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
expected_credential_revision: Revision::new(expected_credential_revision)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
},
PendingTargetDocument::Reauthorize { account_id } => {
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
}
}
};
let actor = match document.owner {
PendingOwnerDocument::AdminSession { admin_user_id } => {
@@ -1147,7 +1175,6 @@ fn decode_mutation(
request_id: document.started_request_id,
};
Ok(PendingAuthorizationMutation::new(
expected_config_revision,
provider_kind,
target,
AuthorizationOwnerBinding::from_context(&context),
@@ -1237,21 +1264,59 @@ fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdmin
| Error::AccountBanned
| Error::IdentityRejected => ProviderAdminErrorKind::Invalid,
Error::NotFound => ProviderAdminErrorKind::NotFound,
Error::RevisionConflict | Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
ProviderAdminErrorKind::Conflict
}
Error::StoreUnavailable | Error::RefreshUnavailable | Error::IdentityUnavailable => {
Error::RefreshUnavailable | Error::IdentityUnavailable => {
ProviderAdminErrorKind::Unavailable
}
})
}
const fn credential_admin_error_code(error: CodexCredentialAdminError) -> &'static str {
match error {
CodexCredentialAdminError::InvalidInput => "invalid_input",
CodexCredentialAdminError::IdentityMismatch => "identity_mismatch",
CodexCredentialAdminError::InvalidCredential => "invalid_credential",
CodexCredentialAdminError::NotFound => "not_found",
CodexCredentialAdminError::MissingRefreshToken => "missing_refresh_token",
CodexCredentialAdminError::RefreshLeaseUnavailable => "refresh_lease_unavailable",
CodexCredentialAdminError::RefreshRejected => "refresh_rejected",
CodexCredentialAdminError::AccountBanned => "account_banned",
CodexCredentialAdminError::RefreshUnavailable => "refresh_unavailable",
CodexCredentialAdminError::RefreshAmbiguous => "refresh_ambiguous",
CodexCredentialAdminError::IdentityRejected => "identity_rejected",
CodexCredentialAdminError::IdentityUnavailable => "identity_unavailable",
}
}
const fn provider_admin_error_code(kind: ProviderAdminErrorKind) -> &'static str {
match kind {
ProviderAdminErrorKind::Invalid => "invalid",
ProviderAdminErrorKind::Unsupported => "unsupported",
ProviderAdminErrorKind::NotFound => "not_found",
ProviderAdminErrorKind::Conflict => "conflict",
ProviderAdminErrorKind::Unavailable => "unavailable",
ProviderAdminErrorKind::Internal => "internal",
}
}
fn log_import_failure(stage: &'static str, error: &'static str) {
tracing::warn!(
import_stage = stage,
import_error = error,
"OpenAI credential import preparation failed"
);
}
fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError {
use CodexOAuthAdminError as Error;
provider_admin_error(match error {
Error::InvalidInput | Error::UpstreamRejected | Error::Credential => {
ProviderAdminErrorKind::Invalid
}
Error::InvalidInput
| Error::CallbackRejected
| Error::TokenRejected
| Error::IdentityRejected
| Error::Credential => ProviderAdminErrorKind::Invalid,
Error::NotFound | Error::FlowExpired => ProviderAdminErrorKind::NotFound,
Error::Conflict | Error::Ambiguous => ProviderAdminErrorKind::Conflict,
Error::UpstreamUnavailable | Error::StorageUnavailable => {
@@ -1260,6 +1325,22 @@ fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError {
})
}
const fn oauth_error_code(error: &CodexOAuthAdminError) -> &'static str {
match error {
CodexOAuthAdminError::InvalidInput => "invalid_input",
CodexOAuthAdminError::NotFound => "not_found",
CodexOAuthAdminError::Conflict => "conflict",
CodexOAuthAdminError::FlowExpired => "flow_expired",
CodexOAuthAdminError::CallbackRejected => "callback_rejected",
CodexOAuthAdminError::TokenRejected => "token_rejected",
CodexOAuthAdminError::IdentityRejected => "identity_rejected",
CodexOAuthAdminError::UpstreamUnavailable => "upstream_unavailable",
CodexOAuthAdminError::Ambiguous => "ambiguous",
CodexOAuthAdminError::StorageUnavailable => "storage_unavailable",
CodexOAuthAdminError::Credential => "credential",
}
}
fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
use CodexCredentialQuotaError as Error;
provider_admin_error(match error {
@@ -13,7 +13,6 @@ use gateway_core::engine::credential::{
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
@@ -28,7 +27,6 @@ use thiserror::Error;
use super::identity::{
CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError,
};
use super::repository::CodexCredentialRepository;
use super::security::CodexCredentialCodec;
use super::token_client::{RefreshFailure, TokenRefresher};
use super::types::{
@@ -42,32 +40,6 @@ const CODEX_CREDENTIAL_SCHEMA_VERSION: u32 = 1;
const MAX_BATCH: usize = 200;
const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1024 * 1024;
const CPR_CONTAINER_KEYS: &[&str] = &["sourceFormat", "source_format", "accounts"];
const CPR_ACCOUNT_KEYS: &[&str] = &[
"id",
"email",
"accountId",
"account_id",
"userId",
"user_id",
"label",
"planType",
"plan_type",
"token",
"at",
"accessToken",
"access_token",
"refreshToken",
"refresh_token",
"accessTokenExpiresAt",
"access_token_expires_at",
"status",
"addedAt",
"added_at",
"updatedAt",
"updated_at",
];
pub struct ImportCodexOAuthCredential {
pub account_id: String,
pub name: String,
@@ -133,7 +105,6 @@ enum ParsedCodexAuthentication {
OAuth {
access_token: Option<String>,
refresh_token: Option<String>,
id_token: Option<String>,
},
AgentIdentity {
runtime_id: String,
@@ -148,7 +119,6 @@ impl fmt::Debug for ParsedCodexAuthentication {
Self::OAuth {
access_token,
refresh_token,
id_token,
} => formatter
.debug_struct("OAuth")
.field("access_token", &access_token.as_ref().map(|_| "<redacted>"))
@@ -156,7 +126,6 @@ impl fmt::Debug for ParsedCodexAuthentication {
"refresh_token",
&refresh_token.as_ref().map(|_| "<redacted>"),
)
.field("id_token", &id_token.as_ref().map(|_| "<redacted>"))
.finish(),
Self::AgentIdentity {
runtime_id: _,
@@ -379,10 +348,6 @@ pub enum CodexCredentialAdminError {
InvalidCredential,
#[error("Codex account was not found")]
NotFound,
#[error("Codex credential revision is stale")]
RevisionConflict,
#[error("Codex account store is unavailable")]
StoreUnavailable,
#[error("Codex account has no refresh token")]
MissingRefreshToken,
#[error("Codex refresh lease is unavailable")]
@@ -580,9 +545,8 @@ impl CodexCredentialAdmin {
}
}
/// 有状态的 Codex 手工刷新边界;只读取 Store 并准备 CAS,不自行持久化。
/// 有状态的 Codex 手工刷新边界;消费调用方刚读取的当前 credential 并准备 CAS。
pub struct CodexCredentialAdminService {
repository: CodexCredentialRepository,
refresher: Arc<dyn TokenRefresher>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
leases: Arc<dyn ProviderLeasePort>,
@@ -593,7 +557,6 @@ impl fmt::Debug for CodexCredentialAdminService {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexCredentialAdminService")
.field("repository", &"CodexCredentialRepository")
.field("refresher", &"TokenRefresher")
.field("verifier", &"CodexAccountIdentityVerifier")
.field("leases", &"ProviderLeasePort")
@@ -604,14 +567,12 @@ impl fmt::Debug for CodexCredentialAdminService {
impl CodexCredentialAdminService {
pub fn new(
repository: CodexCredentialRepository,
refresher: Arc<dyn TokenRefresher>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
leases: Arc<dyn ProviderLeasePort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
) -> Self {
Self {
repository,
refresher,
verifier,
leases,
@@ -622,19 +583,11 @@ impl CodexCredentialAdminService {
/// 官方 RT exchange + AT 身份验证;结果由 App 在同一 revision/audit 事务中提交。
pub async fn manual_refresh(
&self,
account_id: ProviderAccountId,
expected_revision: CredentialRevision,
current: LoadedCredential,
) -> Result<PreparedCodexCredentialRotation, CodexCredentialAdminError> {
let current = self
.repository
.store()
.load_credential(&account_id, expected_revision)
.await
.map_err(map_admin_store_error)?;
if current.account.provider().as_str() != PROVIDER_NAME
|| current.account.id() != &account_id
|| current.account.revision() != expected_revision
{
let account_id = current.account.id().clone();
let expected_revision = current.account.revision();
if current.account.provider().as_str() != PROVIDER_NAME {
return Err(CodexCredentialAdminError::NotFound);
}
let runtime = CodexCredentialCodec::decode(&current.credential)
@@ -720,7 +673,7 @@ impl CodexCredentialAdminService {
Ok(prepared)
}
/// 按文档结构严格判别账号布局,并归一到唯一 `NewProviderAccount` 写入路径。
/// 从 OAuth token 对补全账号事实,并归一到唯一 `NewProviderAccount` 写入路径。
pub async fn prepare_import_document(
&self,
payload: Value,
@@ -741,19 +694,18 @@ impl CodexCredentialAdminService {
let mut accounts = Vec::with_capacity(candidates.len());
let mut policy = None;
for candidate in candidates {
let account_id = candidate
.id
.clone()
.filter(|id| ProviderAccountId::new(id.clone()).is_ok())
.unwrap_or_else(|| format!("acct_{}", uuid::Uuid::now_v7().simple()));
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (enabled, availability) = import_runtime_state(candidate.status.as_deref())?;
let mut prepared = match &candidate.authentication {
ParsedCodexAuthentication::OAuth { .. } => {
let (secret, verified_account) =
self.verify_import_candidate(&candidate).await?;
let prepared = match &candidate.authentication {
ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
} => {
let account_id = format!("acct_{}", uuid::Uuid::now_v7().simple());
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (secret, verified_account) = self
.verify_import_tokens(access_token.clone(), refresh_token.clone())
.await?;
if !upstream_identities.insert((
verified_account.chatgpt_user_id.clone(),
verified_account.chatgpt_account_id.clone(),
@@ -780,19 +732,29 @@ impl CodexCredentialAdminService {
)?;
CodexCredentialAdmin.prepare_import(ImportCodexOAuthCredential {
account_id,
name: candidate
.name
name: verified_account
.email
.clone()
.filter(|name| !name.trim().is_empty())
.unwrap_or_else(|| "Codex OAuth".to_owned()),
secret,
verified_account,
next_refresh_at,
enabled,
enabled: true,
})?
}
ParsedCodexAuthentication::AgentIdentity { .. } => {
let account =
let account_id = candidate
.id
.clone()
.filter(|id| ProviderAccountId::new(id.clone()).is_ok())
.unwrap_or_else(|| format!("acct_{}", uuid::Uuid::now_v7().simple()));
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (enabled, availability) =
import_runtime_state(candidate.status.as_deref())?;
let mut account =
self.prepare_agent_identity_import(&candidate, account_id, enabled)?;
if !upstream_identities.insert((
account.account.upstream_user_id().to_owned(),
@@ -804,29 +766,23 @@ impl CodexCredentialAdminService {
)) {
return Err(CodexCredentialAdminError::InvalidInput);
}
account.account =
account
.account
.with_runtime_state(enabled, availability, None);
account
}
};
prepared.account = prepared
.account
.with_runtime_state(enabled, availability, None);
accounts.push(prepared);
}
Ok(PreparedCodexAccountImport { accounts })
}
async fn verify_import_candidate(
async fn verify_import_tokens(
&self,
candidate: &ParsedCodexImportAccount,
access_token: Option<String>,
refresh_token: Option<String>,
) -> Result<(CodexOAuthSecret, CodexAccountProfile), CodexCredentialAdminError> {
let ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token,
} = &candidate.authentication
else {
return Err(CodexCredentialAdminError::InvalidCredential);
};
let access_token = access_token
.as_deref()
.map(normalize_bearer)
@@ -835,26 +791,14 @@ impl CodexCredentialAdminService {
.as_deref()
.map(normalize_bearer)
.filter(|token| !token.is_empty());
let id_token = id_token
.as_deref()
.map(normalize_bearer)
.filter(|token| !token.is_empty());
let expectation = CodexIdentityExpectation::imported(
candidate.chatgpt_account_id.clone(),
candidate.chatgpt_user_id.clone(),
)
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
if let Some(access_token) = access_token {
let secret = CodexOAuthSecret {
access_token: SecretString::from(access_token),
refresh_token: refresh_token.clone().map(SecretString::from),
id_token: id_token.clone().map(SecretString::from),
id_token: None,
};
match self.verifier.verify(&secret, &expectation).await {
Ok(verification) => {
let profile = verification.into_complete().map_err(map_identity_error)?;
return Ok((secret, profile));
}
match self.verifier.verify_import(&secret).await {
Ok(profile) => return Ok((secret, profile)),
Err(CodexIdentityVerificationError::Unavailable) => {
return Err(CodexCredentialAdminError::IdentityUnavailable);
}
@@ -873,17 +817,19 @@ impl CodexCredentialAdminService {
if tokens.access_token.trim().is_empty() || tokens.expires_in.is_zero() {
return Err(CodexCredentialAdminError::InvalidCredential);
}
// 导入时已经消费来源 RT;只有上游明确轮换返回的新 RT 才能继续保存。
let secret = CodexOAuthSecret {
access_token: SecretString::from(tokens.access_token),
refresh_token: tokens.refresh_token.map(SecretString::from),
id_token: id_token.map(SecretString::from),
// RT 未轮换时仍保留导入提供的 RT,保证本次补全不会丢失后续刷新能力。
refresh_token: tokens
.refresh_token
.map(SecretString::from)
.or_else(|| Some(SecretString::from(refresh_token))),
id_token: None,
};
let profile = self
.verifier
.verify(&secret, &expectation)
.verify_import(&secret)
.await
.and_then(super::identity::CodexIdentityVerification::into_complete)
.map_err(map_identity_error)?;
Ok((secret, profile))
}
@@ -1026,17 +972,6 @@ fn china_rfc3339(value: DateTime<Utc>) -> String {
.to_rfc3339()
}
fn map_admin_store_error(error: gateway_core::error::StoreError) -> CodexCredentialAdminError {
match error.kind() {
StoreErrorKind::Conflict => CodexCredentialAdminError::RevisionConflict,
StoreErrorKind::Unavailable => CodexCredentialAdminError::StoreUnavailable,
StoreErrorKind::InvalidState | StoreErrorKind::InvalidData => {
CodexCredentialAdminError::NotFound
}
_ => CodexCredentialAdminError::StoreUnavailable,
}
}
const fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
match error {
RefreshFailure::InvalidGrant => CodexCredentialAdminError::RefreshRejected,
@@ -1058,67 +993,29 @@ const fn map_identity_error(error: CodexIdentityVerificationError) -> CodexCrede
fn parse_import_document(
payload: &Value,
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let shape = if looks_like_agent_identity_document(payload) {
ImportDocumentShape::AgentIdentity
} else if looks_like_credential_bundle(payload) {
ImportDocumentShape::CredentialBundle
} else if looks_like_auth_document(payload) {
ImportDocumentShape::AuthDocument
} else {
ImportDocumentShape::Native
};
if shape == ImportDocumentShape::Native
&& payload.get("accounts").is_some()
&& payload.as_object().is_none_or(|object| {
object
.keys()
.any(|key| !CPR_CONTAINER_KEYS.contains(&key.as_str()))
})
{
return Err(CodexCredentialAdminError::InvalidInput);
}
let values = import_account_values(payload)?;
let mut accounts = Vec::new();
for value in values {
// CPR 导出允许 OAuth 与 Agent Identity 共存;按账号选择 Provider
// 解析器,避免一个 Agent Identity 条目把整份文档误判成单一格式。
let account_shape = if looks_like_agent_identity_account(value) {
ImportDocumentShape::AgentIdentity
} else if shape == ImportDocumentShape::AgentIdentity {
if value.get("credentials").is_some() || value.get("tokens").is_some() {
ImportDocumentShape::CredentialBundle
} else if value
.as_object()
.is_some_and(auth_document_provider_is_openai)
{
ImportDocumentShape::AuthDocument
} else {
ImportDocumentShape::Native
}
} else {
shape
};
let parsed = match account_shape {
ImportDocumentShape::Native => Some(parse_native_account(value)?),
ImportDocumentShape::AgentIdentity => Some(parse_agent_identity_account(value)?),
ImportDocumentShape::CredentialBundle => parse_credential_bundle_account(value)?,
ImportDocumentShape::AuthDocument => parse_auth_document_account(value)?,
};
if let Some(parsed) = parsed {
accounts.push(parsed);
for value in import_account_values(payload)? {
if looks_like_agent_identity_account(value) {
accounts.push(parse_agent_identity_account(value)?);
continue;
}
if !is_openai_oauth_candidate(value) {
continue;
}
accounts.push(ParsedCodexImportAccount {
id: None,
name: None,
email: None,
plan_type: None,
chatgpt_account_id: None,
chatgpt_user_id: None,
authentication: parse_oauth_import_tokens(value)?,
status: None,
});
}
Ok(accounts)
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum ImportDocumentShape {
Native,
AgentIdentity,
CredentialBundle,
AuthDocument,
}
fn parse_agent_identity_account(
value: &Value,
) -> Result<ParsedCodexImportAccount, CodexCredentialAdminError> {
@@ -1189,274 +1086,44 @@ fn import_account_values(payload: &Value) -> Result<Vec<&Value>, CodexCredential
Ok(vec![payload])
}
fn parse_native_account(
fn parse_oauth_import_tokens(
value: &Value,
) -> Result<ParsedCodexImportAccount, CodexCredentialAdminError> {
if value.get("accounts").is_some() {
return Err(CodexCredentialAdminError::InvalidInput);
) -> Result<ParsedCodexAuthentication, CodexCredentialAdminError> {
let mut access_token = None;
let mut refresh_token = None;
let mut pending = vec![value];
while let Some(current) = pending.pop() {
match current {
Value::Object(object) => {
for (key, value) in object {
let token = match key.as_str() {
"access_token" | "accessToken" | "at" | "token" => &mut access_token,
"refresh_token" | "refreshToken" | "rt" => &mut refresh_token,
_ => {
pending.push(value);
continue;
}
};
let value = value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or(CodexCredentialAdminError::InvalidCredential)?;
if token.replace(value.to_owned()).is_some() {
return Err(CodexCredentialAdminError::InvalidInput);
}
}
}
Value::Array(values) => pending.extend(values),
_ => {}
}
}
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if account
.keys()
.any(|key| !CPR_ACCOUNT_KEYS.contains(&key.as_str()))
{
return Err(CodexCredentialAdminError::InvalidInput);
}
let access_token = first_string(value, &["token", "at", "accessToken", "access_token"]);
let refresh_token = first_string(value, &["refreshToken", "refresh_token"]);
if access_token.is_none() && refresh_token.is_none() {
return Err(CodexCredentialAdminError::InvalidCredential);
}
Ok(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_string(value, &["label", "email"]),
email: first_string(value, &["email"]),
plan_type: first_string(value, &["planType", "plan_type"]),
chatgpt_account_id: first_string(value, &["accountId", "account_id"]),
chatgpt_user_id: first_string(value, &["userId", "user_id"]),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: None,
},
status: first_string(value, &["status"]),
})
}
fn parse_credential_bundle_account(
value: &Value,
) -> Result<Option<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if let Some(credentials) = value.get("credentials") {
if account
.get("platform")
.and_then(Value::as_str)
.is_some_and(|value| !value.eq_ignore_ascii_case("openai"))
|| account
.get("type")
.and_then(Value::as_str)
.is_some_and(|value| !value.eq_ignore_ascii_case("oauth"))
{
return Ok(None);
}
let access_token = first_path_string(
credentials,
&[
&["access_token"],
&["accessToken"],
&["at"],
&["token", "access_token"],
&["token", "accessToken"],
&["token", "at"],
],
);
let refresh_token = first_path_string(
credentials,
&[
&["refresh_token"],
&["refreshToken"],
&["rt"],
&["token", "refresh_token"],
&["token", "refreshToken"],
&["token", "rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
return Ok(Some(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_string(value, &["label", "name"]),
email: first_path_string(credentials, &[&["email"], &["user", "email"]]),
plan_type: first_path_string(
credentials,
&[&["plan_type"], &["planType"], &["plan", "type"]],
),
chatgpt_account_id: first_path_string(
credentials,
&[
&["chatgpt_account_id"],
&["chatgptAccountId"],
&["account_id"],
&["accountId"],
&["account", "id"],
&["account", "account_id"],
&["account", "chatgpt_account_id"],
],
),
chatgpt_user_id: first_path_string(
credentials,
&[
&["chatgpt_user_id"],
&["chatgptUserId"],
&["user_id"],
&["userId"],
&["user", "id"],
&["account", "user_id"],
],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_path_string(credentials, &[&["id_token"], &["idToken"]]),
},
status: first_string(value, &["status"]),
}));
}
let access_token = first_path_string(
value,
&[
&["tokens", "access_token"],
&["tokens", "accessToken"],
&["tokens", "at"],
&["access_token"],
&["accessToken"],
&["token"],
&["at"],
],
);
let refresh_token = first_path_string(
value,
&[
&["tokens", "refresh_token"],
&["tokens", "refreshToken"],
&["tokens", "rt"],
&["refresh_token"],
&["refreshToken"],
&["rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
Ok(Some(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_path_string(value, &[&["label"], &["name"], &["user", "name"]]),
email: first_path_string(value, &[&["email"], &["user", "email"]]),
plan_type: first_path_string(value, &[&["plan_type"], &["planType"], &["plan", "type"]]),
chatgpt_account_id: first_path_string(
value,
&[
&["chatgpt_account_id"],
&["chatgptAccountId"],
&["account_id"],
&["accountId"],
&["account", "id"],
&["account", "account_id"],
&["account", "chatgpt_account_id"],
],
),
chatgpt_user_id: first_path_string(
value,
&[
&["chatgpt_user_id"],
&["chatgptUserId"],
&["user_id"],
&["userId"],
&["user", "id"],
&["account", "user_id"],
],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_path_string(
value,
&[&["tokens", "id_token"], &["id_token"], &["idToken"]],
),
},
status: first_string(value, &["status"]),
}))
}
fn parse_auth_document_account(
value: &Value,
) -> Result<Option<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if !auth_document_provider_is_openai(account) {
return Ok(None);
}
let access_token = first_path_string(
value,
&[
&["access_token"],
&["accessToken"],
&["at"],
&["token", "access_token"],
&["token", "accessToken"],
&["token", "at"],
&["token"],
],
);
let refresh_token = first_path_string(
value,
&[
&["refresh_token"],
&["refreshToken"],
&["rt"],
&["token", "refresh_token"],
&["token", "refreshToken"],
&["token", "rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
Ok(Some(ParsedCodexImportAccount {
id: None,
name: first_string(value, &["label", "name", "email"]),
email: first_string(value, &["email"]),
plan_type: first_string(value, &["plan_type", "planType"]),
chatgpt_account_id: first_string(
value,
&[
"chatgpt_account_id",
"chatgptAccountId",
"account_id",
"accountId",
],
),
chatgpt_user_id: first_string(
value,
&["chatgpt_user_id", "chatgptUserId", "user_id", "userId"],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_string(value, &["id_token", "idToken"]),
},
status: account
.get("disabled")
.and_then(Value::as_bool)
.filter(|disabled| *disabled)
.map(|_| "disabled".to_owned()),
}))
}
fn looks_like_credential_bundle(value: &Value) -> bool {
value.get("exported_at").is_some()
|| value.get("proxies").is_some()
|| import_account_values(value).is_ok_and(|accounts| {
accounts.iter().any(|account| {
account.get("credentials").is_some()
|| account.get("tokens").is_some()
|| account.get("cachedQuota").is_some()
|| account.get("cached_quota").is_some()
})
})
}
fn looks_like_agent_identity_document(value: &Value) -> bool {
import_account_values(value).is_ok_and(|accounts| {
accounts
.iter()
.any(|account| looks_like_agent_identity_account(account))
Ok(ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
})
}
@@ -1471,26 +1138,24 @@ fn looks_like_agent_identity_account(value: &Value) -> bool {
|| first_string(identity, &["agent_runtime_id", "agentRuntimeId"]).is_some()
}
fn looks_like_auth_document(value: &Value) -> bool {
import_account_values(value).is_ok_and(|accounts| {
accounts.iter().any(|account| {
account
.as_object()
.is_some_and(auth_document_provider_is_openai)
})
})
}
fn auth_document_provider_is_openai(account: &serde_json::Map<String, Value>) -> bool {
account
.get("type")
fn is_openai_oauth_candidate(value: &Value) -> bool {
let Some(account) = value.as_object() else {
return false;
};
if let Some(provider) = account
.get("platform")
.or_else(|| account.get("provider"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|provider| {
provider.eq_ignore_ascii_case("openai")
// CLIProxyAPI 的 OpenAI Codex OAuth 文件使用 `type: codex`。
|| provider.eq_ignore_ascii_case("codex")
{
return provider.eq_ignore_ascii_case("openai") || provider.eq_ignore_ascii_case("codex");
}
account
.get("type")
.and_then(Value::as_str)
.is_none_or(|kind| {
kind.eq_ignore_ascii_case("openai")
|| kind.eq_ignore_ascii_case("codex")
|| kind.eq_ignore_ascii_case("oauth")
})
}
@@ -1502,23 +1167,6 @@ fn first_string(value: &Value, keys: &[&str]) -> Option<String> {
.map(str::to_owned)
}
fn first_path_string(value: &Value, paths: &[&[&str]]) -> Option<String> {
paths
.iter()
.find_map(|path| value_at_path(value, path).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned)
}
fn value_at_path<'a>(value: &'a Value, path: &[&str]) -> Option<&'a Value> {
let mut current = value;
for key in path {
current = current.get(*key)?;
}
Some(current)
}
fn normalize_bearer(value: &str) -> String {
let value = value.trim();
value
@@ -347,7 +347,7 @@ impl CodexAgentIdentityTaskService {
self.load_current(account_id).await
}
async fn load_current(
pub(crate) async fn load_current(
&self,
account_id: &ProviderAccountId,
) -> Result<PreparedCodexRuntimeCredential, CodexAgentIdentityError> {
@@ -116,6 +116,24 @@ impl CodexSignedIdentity {
pub const fn access_token_expires_at(&self) -> DateTime<Utc> {
self.access_token_expires_at
}
fn into_import_profile(self) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
let chatgpt_account_id = self
.claimed_account_id
.ok_or(CodexIdentityVerificationError::Rejected)?;
let chatgpt_user_id = self
.claimed_user_id
.ok_or(CodexIdentityVerificationError::Rejected)?;
Ok(CodexAccountProfile {
email: self.email,
oauth_subject: self.oauth_subject,
poid: self.poid,
chatgpt_account_id,
chatgpt_user_id,
plan_type: self.plan_type,
access_token_expires_at: Some(self.access_token_expires_at),
})
}
}
/// 导入文档或现有 credential 对认证结果施加的一致性约束。
@@ -256,6 +274,19 @@ pub trait CodexAccountIdentityVerifier: Send + Sync {
expectation: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError>;
/// 为无状态导入补全 OAuth 账号事实。
///
/// 导入不信任导出文件中的身份投影;实现必须只使用已验签的 token
/// 与官方认证接口返回的事实生成可持久化 profile。
async fn verify_import(
&self,
secret: &CodexOAuthSecret,
) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
self.verify(secret, &CodexIdentityExpectation::default())
.await?
.into_complete()
}
async fn verify_authorization(
&self,
secret: &CodexOAuthSecret,
@@ -580,6 +611,31 @@ impl CodexAccountIdentityVerifier for CodexAccountIdentityService {
self.complete(secret, signed, expectation).await
}
async fn verify_import(
&self,
secret: &CodexOAuthSecret,
) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
let signed = self.signed.verify_access(secret).await?;
let expectation = CodexIdentityExpectation::imported(
signed.claimed_account_id.clone(),
signed.claimed_user_id.clone(),
)?;
let fallback = signed.clone().into_import_profile().ok();
match self.accounts.fetch(secret, &expectation).await {
Ok(account) => match complete_profile(signed, account, &expectation) {
Ok(profile) => Ok(profile),
Err(_) => fallback.ok_or(CodexIdentityVerificationError::Rejected),
},
// `/wham/usage` 是导入时的补全来源,不是 token 有效性的裁决者。AT
// 已完成官方验签且带有完整账号声明时,接口拒绝、暂时不可用或返回默认
// workspace 的冲突投影,都不能阻止恢复该 OAuth 账号。
Err(
CodexIdentityVerificationError::Rejected
| CodexIdentityVerificationError::Unavailable,
) => fallback.ok_or(CodexIdentityVerificationError::Rejected),
}
}
async fn verify_authorization(
&self,
secret: &CodexOAuthSecret,
@@ -962,7 +1018,7 @@ fn valid_signed_token_header(header: &jsonwebtoken::Header) -> bool {
&& header.zip.is_none()
&& header.url.is_none()
&& header.nonce.is_none()
&& header.extras.is_empty()
&& header.extras.inner().is_empty()
}
fn constant_time_equal(left: &[u8], right: &[u8]) -> bool {
@@ -1,4 +1,4 @@
//! Codex credential domain exports。
//! Codex credential 领域导出。
mod admin;
mod affinity;
@@ -47,7 +47,7 @@ pub use identity::{
};
pub use oauth::{
CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthAuthorizationStarted,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthReauthorizationTarget,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthAuthorization,
CompletedCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
};
@@ -1,16 +1,19 @@
//! Codex Authorization Code + PKCE/OIDC 管理流。
use std::fmt;
use std::sync::Arc;
use std::{fmt, sync::Arc, time::Duration};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwner, PendingAuthorizationMutation,
use gateway_admin::model::{
AdminError,
provider_credentials::{
AuthorizationCommitGuard, AuthorizationMutationTarget, AuthorizationOwner,
PendingAuthorizationMutation,
},
};
use gateway_core::engine::credential::{
CredentialRevision, NewProviderAccount, ProviderAccountId, ProviderAccountStore,
NewProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::provider_ports::ProviderRuntimePolicyPort;
use secrecy::{ExposeSecret, SecretString};
@@ -35,6 +38,7 @@ use super::token_client::{
const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
const AUTHORIZATION_SCOPE: &str = "openid profile email offline_access";
const AUTHORIZATION_TTL: TimeDelta = TimeDelta::minutes(10);
const AUTHORIZATION_CLAIM_TTL: Duration = Duration::from_secs(90);
const MAX_CALLBACK_BYTES: usize = 16 * 1024;
const MAX_TEXT_BYTES: usize = 512;
@@ -43,24 +47,6 @@ pub struct StartCodexOAuthAuthorization {
pub mutation: PendingAuthorizationMutation,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodexOAuthReauthorizationTarget {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
}
impl CodexOAuthReauthorizationTarget {
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
&self.account_id
}
#[must_use]
pub const fn credential_revision(&self) -> CredentialRevision {
self.credential_revision
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct CodexOAuthAuthorizationStarted {
pub flow_id: String,
@@ -89,6 +75,7 @@ pub struct CompleteCodexOAuthAuthorization {
pub struct CompletedCodexOAuthAuthorization<T> {
pub mutation: PendingAuthorizationMutation,
pub credential: T,
authorization_guard: Box<dyn AuthorizationCommitGuard>,
}
/// OAuth exchange 后唯一的 credential preparation 结果。
@@ -116,6 +103,18 @@ impl<T> fmt::Debug for CompletedCodexOAuthAuthorization<T> {
}
}
impl<T> CompletedCodexOAuthAuthorization<T> {
pub(crate) fn into_parts(
self,
) -> (
PendingAuthorizationMutation,
T,
Box<dyn AuthorizationCommitGuard>,
) {
(self.mutation, self.credential, self.authorization_guard)
}
}
impl fmt::Debug for CompleteCodexOAuthAuthorization {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -136,7 +135,7 @@ pub struct CodexPendingAuthorization {
state: SecretString,
nonce: SecretString,
code_verifier: SecretString,
reauthorization: Option<CodexOAuthReauthorizationTarget>,
reauthorization: Option<ProviderAccountId>,
mutation: PendingAuthorizationMutation,
}
@@ -150,7 +149,6 @@ pub struct StoredCodexPendingAuthorization {
pub nonce: SecretString,
pub code_verifier: SecretString,
pub reauthorization_account_id: Option<String>,
pub reauthorization_credential_revision: Option<u64>,
pub mutation: PendingAuthorizationMutation,
}
@@ -158,19 +156,11 @@ impl CodexPendingAuthorization {
pub fn from_stored(
input: StoredCodexPendingAuthorization,
) -> Result<Self, CodexOAuthPendingStoreError> {
let reauthorization = match (
input.reauthorization_account_id,
input.reauthorization_credential_revision,
) {
(None, None) => None,
(Some(account_id), Some(revision)) => Some(CodexOAuthReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
credential_revision: CredentialRevision::new(revision)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
}),
_ => return Err(CodexOAuthPendingStoreError::InvalidValue),
};
let reauthorization = input
.reauthorization_account_id
.map(ProviderAccountId::new)
.transpose()
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let pending = Self {
flow_id: input.flow_id,
owner_ref: input.owner_ref,
@@ -239,7 +229,7 @@ impl CodexPendingAuthorization {
}
#[must_use]
pub const fn reauthorization(&self) -> Option<&CodexOAuthReauthorizationTarget> {
pub const fn reauthorization(&self) -> Option<&ProviderAccountId> {
self.reauthorization.as_ref()
}
@@ -284,11 +274,33 @@ pub trait CodexOAuthPendingStore: Send + Sync {
pending: &CodexPendingAuthorization,
) -> Result<(), CodexOAuthPendingStoreError>;
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError>;
claim_ref: &str,
claim_ttl: Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError>;
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError>;
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError>;
}
pub enum CodexOAuthPendingClaimOutcome {
Claimed(Box<CodexPendingAuthorization>),
NotFound,
InProgress,
}
#[derive(Debug, PartialEq, Eq, thiserror::Error)]
@@ -301,8 +313,12 @@ pub enum CodexOAuthAdminError {
Conflict,
#[error("Codex OAuth flow expired")]
FlowExpired,
#[error("Codex OAuth upstream rejected the operation")]
UpstreamRejected,
#[error("Codex OAuth callback was rejected")]
CallbackRejected,
#[error("Codex OAuth token exchange was rejected")]
TokenRejected,
#[error("Codex OAuth identity verification was rejected")]
IdentityRejected,
#[error("Codex OAuth upstream is unavailable")]
UpstreamUnavailable,
#[error("Codex OAuth exchange send state is ambiguous")]
@@ -336,6 +352,46 @@ pub struct CodexOAuthAdminService {
oauth_client_id: String,
}
struct CodexOAuthAuthorizationCommitGuard {
pending: Arc<dyn CodexOAuthPendingStore>,
owner_ref: String,
flow_id: String,
claim_ref: String,
}
#[async_trait]
impl AuthorizationCommitGuard for CodexOAuthAuthorizationCommitGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
let consumed = self
.pending
.consume_claim(&self.owner_ref, &self.flow_id, &self.claim_ref)
.await
.map_err(map_authorization_settlement_error)?;
if consumed {
Ok(())
} else {
Err(AdminError::conflict(
"OpenAI OAuth pending claim is no longer current",
))
}
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
let released = self
.pending
.release_claim(&self.owner_ref, &self.flow_id, &self.claim_ref)
.await
.map_err(map_authorization_settlement_error)?;
if released {
Ok(())
} else {
Err(AdminError::unavailable(
"OpenAI OAuth pending claim could not be released",
))
}
}
}
impl CodexOAuthAdminService {
#[must_use]
pub fn new(
@@ -372,15 +428,10 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
let (name, reauthorization) = match command.mutation.target() {
AuthorizationMutationTarget::Create { name } => (name.clone(), None),
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
let expected_revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| CodexOAuthAdminError::InvalidInput)?;
AuthorizationMutationTarget::Reauthorize { account_id } => {
let current = self
.store
.load_credential(account_id, expected_revision)
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider().as_str() != "openai"
@@ -388,13 +439,7 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
{
return Err(CodexOAuthAdminError::NotFound);
}
(
current.account.name().to_owned(),
Some(CodexOAuthReauthorizationTarget {
account_id: account_id.clone(),
credential_revision: expected_revision,
}),
)
(current.account.name().to_owned(), Some(account_id.clone()))
}
};
self.start_pending(command.mutation, name, reauthorization)
@@ -412,12 +457,43 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
{
return Err(CodexOAuthAdminError::InvalidInput);
}
let (pending, mut secret, id_token) = self.exchange_pending(command).await?;
let claim_ref = random_secret()?;
let pending = self.claim_pending(&command, &claim_ref).await?;
let completed = self
.complete_claimed_authorization(pending, command.callback_url.expose_secret())
.await;
match completed {
Ok((mutation, credential)) => Ok(CompletedCodexOAuthAuthorization {
mutation,
credential,
authorization_guard: Box::new(CodexOAuthAuthorizationCommitGuard {
pending: Arc::clone(&self.pending),
owner_ref: command.owner_ref,
flow_id: command.flow_id,
claim_ref,
}),
}),
Err(error) => {
self.release_claim(&command, &claim_ref).await?;
Err(error)
}
}
}
}
impl CodexOAuthAdminService {
async fn complete_claimed_authorization(
&self,
pending: CodexPendingAuthorization,
callback_url: &str,
) -> Result<(PendingAuthorizationMutation, CompletedCodexOAuthCredential), CodexOAuthAdminError>
{
let (mut secret, id_token) = self.exchange_pending(&pending, callback_url).await?;
let mutation = pending.mutation.clone();
let current = if let Some(target) = pending.reauthorization() {
let current = self
.store
.load_credential(target.account_id(), target.credential_revision())
.load_current_credential(target)
.await
.map_err(map_store_error)?;
Some(current)
@@ -483,10 +559,7 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
.map_err(map_admin_error)?,
)
};
Ok(CompletedCodexOAuthAuthorization {
mutation,
credential,
})
Ok((mutation, credential))
}
}
@@ -495,7 +568,7 @@ impl CodexOAuthAdminService {
&self,
mutation: PendingAuthorizationMutation,
name: String,
reauthorization: Option<CodexOAuthReauthorizationTarget>,
reauthorization: Option<ProviderAccountId>,
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
if !valid_text(&name) {
return Err(CodexOAuthAdminError::InvalidInput);
@@ -514,11 +587,7 @@ impl CodexOAuthAdminService {
state: SecretString::from(random_secret()?),
nonce: SecretString::from(random_secret()?),
code_verifier: SecretString::from(random_secret()?),
reauthorization_account_id: reauthorization
.as_ref()
.map(|target| target.account_id().to_string()),
reauthorization_credential_revision: reauthorization
.map(|target| target.credential_revision().get()),
reauthorization_account_id: reauthorization.as_ref().map(ToString::to_string),
mutation,
})
.map_err(map_pending_error)?;
@@ -533,38 +602,70 @@ impl CodexOAuthAdminService {
})
}
async fn exchange_pending(
async fn claim_pending(
&self,
command: CompleteCodexOAuthAuthorization,
) -> Result<
(
CodexPendingAuthorization,
super::types::CodexOAuthSecret,
SecretString,
),
CodexOAuthAdminError,
> {
command: &CompleteCodexOAuthAuthorization,
claim_ref: &str,
) -> Result<CodexPendingAuthorization, CodexOAuthAdminError> {
if !valid_text(&command.owner_ref)
|| !valid_text(&command.flow_id)
|| !valid_secret(claim_ref)
|| command.callback_url.expose_secret().len() > MAX_CALLBACK_BYTES
{
return Err(CodexOAuthAdminError::InvalidInput);
}
let pending = self
let pending = match self
.pending
.take(&command.owner_ref, &command.flow_id)
.claim(
&command.owner_ref,
&command.flow_id,
claim_ref,
AUTHORIZATION_CLAIM_TTL,
)
.await
.map_err(map_pending_error)?
.ok_or(CodexOAuthAdminError::NotFound)?;
{
CodexOAuthPendingClaimOutcome::Claimed(pending) => *pending,
CodexOAuthPendingClaimOutcome::NotFound => return Err(CodexOAuthAdminError::NotFound),
CodexOAuthPendingClaimOutcome::InProgress => {
return Err(CodexOAuthAdminError::Conflict);
}
};
if pending.expires_at <= Utc::now() {
self.release_claim(command, claim_ref).await?;
return Err(CodexOAuthAdminError::FlowExpired);
}
let (code, callback_state) = callback_parts(command.callback_url.expose_secret())?;
Ok(pending)
}
async fn release_claim(
&self,
command: &CompleteCodexOAuthAuthorization,
claim_ref: &str,
) -> Result<(), CodexOAuthAdminError> {
if self
.pending
.release_claim(&command.owner_ref, &command.flow_id, claim_ref)
.await
.map_err(map_pending_error)?
{
Ok(())
} else {
Err(CodexOAuthAdminError::StorageUnavailable)
}
}
async fn exchange_pending(
&self,
pending: &CodexPendingAuthorization,
callback_url: &str,
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
let (code, callback_state) = callback_parts(callback_url)?;
if !constant_time_equal(
pending.state.expose_secret().as_bytes(),
callback_state.expose_secret().as_bytes(),
) {
return Err(CodexOAuthAdminError::UpstreamRejected);
return Err(CodexOAuthAdminError::CallbackRejected);
}
let tokens = self
.exchanger
@@ -574,7 +675,7 @@ impl CodexOAuthAdminService {
})
.await
.map_err(map_exchange_error)?;
Ok((pending, tokens.secret, tokens.id_token))
Ok((tokens.secret, tokens.id_token))
}
}
@@ -622,32 +723,44 @@ fn authorization_url(
}
fn callback_parts(value: &str) -> Result<(SecretString, SecretString), CodexOAuthAdminError> {
let url = Url::parse(value).map_err(|_| CodexOAuthAdminError::UpstreamRejected)?;
let url = Url::parse(value).map_err(|_| CodexOAuthAdminError::CallbackRejected)?;
if url.scheme() != "http"
|| url.host_str() != Some("localhost")
|| url.port() != Some(1455)
|| url.path() != "/auth/callback"
|| url.fragment().is_some()
{
return Err(CodexOAuthAdminError::UpstreamRejected);
return Err(CodexOAuthAdminError::CallbackRejected);
}
let mut code = None;
let mut state = None;
for (key, value) in url.query_pairs() {
match key.as_ref() {
"code" if code.is_none() => code = Some(value.into_owned()),
"state" if state.is_none() => state = Some(value.into_owned()),
_ => return Err(CodexOAuthAdminError::UpstreamRejected),
"code" => set_unique_callback_parameter(&mut code, value.into_owned())?,
"state" => set_unique_callback_parameter(&mut state, value.into_owned())?,
// OAuth 服务可能附加 `scope`、`iss` 等标准参数。
// 固定回调地址与唯一的 `code`/`state` 才是安全边界,无关参数不参与校验。
_ => {}
}
}
let code = code.filter(|value| valid_secret(value));
let state = state.filter(|value| valid_secret(value));
match (code, state) {
(Some(code), Some(state)) => Ok((SecretString::from(code), SecretString::from(state))),
_ => Err(CodexOAuthAdminError::UpstreamRejected),
_ => Err(CodexOAuthAdminError::CallbackRejected),
}
}
fn set_unique_callback_parameter(
target: &mut Option<String>,
value: String,
) -> Result<(), CodexOAuthAdminError> {
if target.replace(value).is_some() {
return Err(CodexOAuthAdminError::CallbackRejected);
}
Ok(())
}
fn random_secret() -> Result<String, CodexOAuthAdminError> {
let mut bytes = [0_u8; 32];
getrandom::fill(&mut bytes).map_err(|_| CodexOAuthAdminError::StorageUnavailable)?;
@@ -688,15 +801,8 @@ fn pending_mutation_matches(pending: &CodexPendingAuthorization) -> bool {
}
match (pending.mutation.target(), pending.reauthorization.as_ref()) {
(AuthorizationMutationTarget::Create { name }, None) => name == &pending.name,
(
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
},
Some(target),
) => {
account_id == target.account_id()
&& expected_credential_revision.get() == target.credential_revision().get()
(AuthorizationMutationTarget::Reauthorize { account_id }, Some(target)) => {
account_id == target
}
_ => false,
}
@@ -710,16 +816,33 @@ fn map_pending_error(error: CodexOAuthPendingStoreError) -> CodexOAuthAdminError
}
}
fn map_authorization_settlement_error(error: CodexOAuthPendingStoreError) -> AdminError {
match error {
CodexOAuthPendingStoreError::InvalidValue => {
AdminError::internal("OpenAI OAuth pending claim is invalid")
}
CodexOAuthPendingStoreError::Conflict => {
AdminError::conflict("OpenAI OAuth pending claim conflicts with current state")
}
CodexOAuthPendingStoreError::Unavailable => {
AdminError::unavailable("OpenAI OAuth pending claim store is unavailable")
}
}
}
fn map_exchange_error(error: AuthorizationCodeExchangeError) -> CodexOAuthAdminError {
match error {
AuthorizationCodeExchangeError::Rejected => CodexOAuthAdminError::UpstreamRejected,
AuthorizationCodeExchangeError::Rejected => CodexOAuthAdminError::TokenRejected,
AuthorizationCodeExchangeError::Unavailable => CodexOAuthAdminError::UpstreamUnavailable,
AuthorizationCodeExchangeError::Ambiguous => CodexOAuthAdminError::Ambiguous,
}
}
fn map_identity_error(_: CodexIdentityVerificationError) -> CodexOAuthAdminError {
CodexOAuthAdminError::UpstreamRejected
fn map_identity_error(error: CodexIdentityVerificationError) -> CodexOAuthAdminError {
match error {
CodexIdentityVerificationError::Rejected => CodexOAuthAdminError::IdentityRejected,
CodexIdentityVerificationError::Unavailable => CodexOAuthAdminError::UpstreamUnavailable,
}
}
fn map_admin_error(_: CodexCredentialAdminError) -> CodexOAuthAdminError {
@@ -30,6 +30,9 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
/// OpenAI Provider 主动额度刷新的调度策略。
///
@@ -175,6 +178,14 @@ impl CodexQuotaFact {
}
}
fn quota_fact_is_active_exhaustion(fact: CodexQuotaFact, now: SystemTime) -> bool {
fact.exhausted()
&& fact
.resets_at()
.map(SystemTime::from)
.is_none_or(|reset_at| reset_at > now)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CodexQuotaSyncSummary {
pub updated: u64,
@@ -235,6 +246,8 @@ struct CodexQuotaProjectionState {
next_version: u64,
entries: BTreeMap<ProviderAccountId, CodexQuotaSchedulingEntry>,
last_periodic_refresh_at: BTreeMap<ProviderAccountId, Instant>,
last_exhausted_reset_recheck: BTreeMap<ProviderAccountId, CodexQuotaResetMarker>,
ready_refresh_cursor: Option<ProviderAccountId>,
}
#[derive(Debug, Clone, Copy)]
@@ -245,6 +258,24 @@ struct CodexQuotaSchedulingEntry {
signals: Option<AccountQuotaSignals>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CodexQuotaResetMarker {
revision: CredentialRevision,
reset_at: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodexQuotaRefreshKind {
Exhausted { reset_at: Option<SystemTime> },
Ready,
}
#[derive(Clone)]
struct CodexQuotaRefreshCandidate {
account: ProviderAccount,
kind: CodexQuotaRefreshKind,
}
#[derive(Clone)]
struct CodexQuotaHydrationTarget {
account: ProviderAccount,
@@ -395,17 +426,18 @@ impl CodexQuotaSchedulingProjection {
fn reserve_periodic_refreshes(
&self,
accounts: Vec<ProviderAccount>,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Vec<ProviderAccount> {
let candidates = accounts
.into_iter()
.filter(|account| eligible_quota_sync_account(account, now))
.filter_map(|account| quota_refresh_candidate(account, observations, now))
.collect::<Vec<_>>();
let candidate_ids = candidates
.iter()
.map(|account| account.id().clone())
.map(|candidate| candidate.account.id().clone())
.collect::<BTreeSet<_>>();
let now = Instant::now();
let refreshed_at = Instant::now();
let mut state = self
.state
.write()
@@ -413,26 +445,185 @@ impl CodexQuotaSchedulingProjection {
state
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
candidates
.into_iter()
.filter(|account| {
let due = state
.last_periodic_refresh_at
.get(account.id())
.is_none_or(|last| {
now.saturating_duration_since(*last) >= QUOTA_PERIODIC_REFRESH_MIN_INTERVAL
});
if due {
state
.last_periodic_refresh_at
.insert(account.id().clone(), now);
state
.last_exhausted_reset_recheck
.retain(|account_id, _| candidate_ids.contains(account_id));
let mut reserved = Vec::new();
let mut ready = Vec::new();
for candidate in candidates {
match candidate.kind {
CodexQuotaRefreshKind::Exhausted { reset_at } => {
if reserve_exhausted_quota_refresh(
&mut state,
&candidate.account,
reset_at,
now,
refreshed_at,
) {
reserved.push(candidate.account);
}
}
due
})
.collect()
CodexQuotaRefreshKind::Ready => ready.push(candidate.account),
}
}
ready.sort_by(|left, right| left.id().cmp(right.id()));
let ready = ready
.into_iter()
.filter(|account| periodic_quota_refresh_due(&state, account.id(), refreshed_at))
.collect::<Vec<_>>();
let ready_len = ready.len();
if ready_len == 0 {
return reserved;
}
let start = state
.ready_refresh_cursor
.as_ref()
.and_then(|cursor| ready.iter().position(|account| account.id() > cursor))
.unwrap_or(0);
let ready = ready
.into_iter()
.cycle()
.skip(start)
.take(ready_len.min(QUOTA_READY_REFRESH_BATCH_SIZE))
.collect::<Vec<_>>();
if let Some(last) = ready.last() {
state.ready_refresh_cursor = Some(last.id().clone());
}
for account in &ready {
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
}
reserved.extend(ready);
reserved
}
}
fn quota_refresh_candidate(
account: ProviderAccount,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Option<CodexQuotaRefreshCandidate> {
if !eligible_periodic_quota_refresh(&account, now) {
return None;
}
let observation = observations.and_then(|observations| observations.get(account.id()));
let snapshot = quota_snapshot_for_account(&account, observation);
let kind = match account.availability() {
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
},
AccountAvailability::Cooldown
if snapshot
.as_ref()
.is_some_and(|snapshot| snapshot.fact().exhausted()) =>
{
CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
}
}
AccountAvailability::Ready
if observations.is_some() && quota_observation_is_stale(&account, observation, now) =>
{
CodexQuotaRefreshKind::Ready
}
_ => return None,
};
Some(CodexQuotaRefreshCandidate { account, kind })
}
fn quota_snapshot_for_account(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
) -> Option<CodexAccountQuotaSnapshot> {
observation
.filter(|observation| observation.expected_revision == account.revision())
.and_then(quota_snapshot_from_observation)
}
fn quota_reset_at_from_snapshot(
snapshot: Option<&CodexAccountQuotaSnapshot>,
) -> Option<SystemTime> {
snapshot
.and_then(|snapshot| snapshot.fact().resets_at())
.map(SystemTime::from)
}
fn quota_observation_confirms_active_exhaustion(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
now: SystemTime,
) -> bool {
quota_snapshot_for_account(account, observation)
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))
}
fn quota_observation_is_stale(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
now: SystemTime,
) -> bool {
let Some(observation) = observation.filter(|observation| {
observation.expected_revision == account.revision()
&& quota_snapshot_from_observation(observation).is_some()
}) else {
return true;
};
observation
.observed_at
.and_then(|observed_at| now.duration_since(observed_at).ok())
.is_none_or(|age| age >= QUOTA_READY_REFRESH_MAX_AGE)
}
fn reserve_exhausted_quota_refresh(
state: &mut CodexQuotaProjectionState,
account: &ProviderAccount,
reset_at: Option<SystemTime>,
now: SystemTime,
refreshed_at: Instant,
) -> bool {
let reset_recheck_due = reset_at.is_some_and(|reset_at| {
reset_at <= now
&& state
.last_exhausted_reset_recheck
.get(account.id())
.is_none_or(|marker| {
marker.revision != account.revision() || marker.reset_at != reset_at
})
});
if !reset_recheck_due && !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
if let Some(reset_at) = reset_at.filter(|reset_at| *reset_at <= now) {
state.last_exhausted_reset_recheck.insert(
account.id().clone(),
CodexQuotaResetMarker {
revision: account.revision(),
reset_at,
},
);
}
true
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
now: Instant,
) -> bool {
state
.last_periodic_refresh_at
.get(account_id)
.is_none_or(|last| {
now.saturating_duration_since(*last) >= QUOTA_PERIODIC_REFRESH_MIN_INTERVAL
})
}
fn insert_projection_entry(
state: &mut CodexQuotaProjectionState,
account_id: ProviderAccountId,
@@ -516,10 +707,31 @@ impl CodexCredentialQuotaService {
}
pub async fn synchronize(&self) -> Result<CodexQuotaSyncSummary, CodexCredentialQuotaError> {
let accounts = self.repository.list_for_provider().await?;
let mut accounts = self.repository.list_for_provider().await?;
let mut summary = CodexQuotaSyncSummary::default();
let now = SystemTime::now();
let accounts = self.scheduling.reserve_periodic_refreshes(accounts, now);
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
.collect::<Vec<_>>();
let observations = match self.store.get_quotas(&account_ids).await {
Ok(observations) => Some(
observations
.into_iter()
.map(|observation| (observation.account_id.clone(), observation))
.collect::<BTreeMap<_, _>>(),
),
Err(error) => {
// 持久观察不可用时不盲扫 ready;保留已耗尽账号的既有恢复轮询。
tracing::warn!(error = %error, "OpenAI quota refresh skipped ready-account scan");
None
}
};
self.reconcile_persisted_exhaustion(&mut accounts, observations.as_ref(), now)
.await;
let accounts =
self.scheduling
.reserve_periodic_refreshes(accounts, observations.as_ref(), now);
if accounts.is_empty() {
return Ok(summary);
}
@@ -553,6 +765,9 @@ impl CodexCredentialQuotaService {
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
let (availability, reason, cooldown) = classify_error(&error, observed_at);
let preserve_exhaustion = account.availability()
== AccountAvailability::QuotaExhausted
&& availability == Some(AccountAvailability::Cooldown);
match availability {
Some(AccountAvailability::Invalid | AccountAvailability::Banned) => {
summary.invalid += 1;
@@ -564,6 +779,9 @@ impl CodexCredentialQuotaService {
continue;
}
}
if preserve_exhaustion {
continue;
}
let _ = self
.repository
.apply_state(
@@ -580,6 +798,63 @@ impl CodexCredentialQuotaService {
Ok(summary)
}
async fn reconcile_persisted_exhaustion(
&self,
accounts: &mut [ProviderAccount],
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
observed_at: SystemTime,
) {
let Some(observations) = observations else {
return;
};
for account in accounts {
if account.availability() != AccountAvailability::Cooldown
|| !quota_observation_confirms_active_exhaustion(
account,
observations.get(account.id()),
observed_at,
)
{
continue;
}
let Ok(Some(current)) = self.store.get_account(account.id()).await else {
continue;
};
if current.revision() != account.revision()
|| current.availability() != AccountAvailability::Cooldown
{
continue;
}
if !matches!(
self.has_active_exhaustion(&current, observed_at).await,
Ok(true)
) {
continue;
}
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
observed_at,
)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
"OpenAI persisted quota exhaustion reconciliation failed"
);
continue;
}
let enabled = current.enabled();
*account =
current.with_runtime_state(enabled, AccountAvailability::QuotaExhausted, None);
}
}
/// 解析并 revision-fenced 落库单账号的 Provider quota JSON。
async fn apply_fetched_quota(
&self,
@@ -717,6 +992,30 @@ impl CodexCredentialQuotaService {
.await?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialQuotaError::NotFound)?;
self.read_snapshot_for(&account).await
}
pub(super) async fn has_active_exhaustion(
&self,
account: &ProviderAccount,
now: SystemTime,
) -> Result<bool, CodexCredentialQuotaError> {
match self.read_snapshot_for(account).await {
Ok(snapshot) => Ok(snapshot
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))),
Err(CodexCredentialQuotaError::RevisionConflict) => Ok(false),
Err(error) => Err(error),
}
}
async fn read_snapshot_for(
&self,
account: &ProviderAccount,
) -> Result<Option<CodexAccountQuotaSnapshot>, CodexCredentialQuotaError> {
if account.provider().as_str() != "openai" {
return Err(CodexCredentialQuotaError::NotFound);
}
let account_id = account.id();
let Some(observation) = self
.store
.get_quotas(std::slice::from_ref(account_id))
@@ -918,12 +1217,11 @@ async fn fetch_usage_once(
.map_err(CodexQuotaFetchAttemptError::Upstream)
}
fn eligible_quota_sync_account(account: &ProviderAccount, now: SystemTime) -> bool {
fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled()
&& account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
&& account.availability() == AccountAvailability::QuotaExhausted
}
fn merge_passive_quota(
@@ -973,13 +1271,11 @@ fn merge_passive_quota(
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (limit_id, details) in &rate_limits.limits {
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
if Some(limit_id.as_str()) == active_limit {
let mut value = quota
.remove("rate_limit")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_details(&mut value, details);
quota.insert("rate_limit".to_owned(), Value::Object(value));
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
continue;
}
let index = additional.iter().position(|item| {
@@ -1000,12 +1296,7 @@ fn merge_passive_quota(
if let Some(name) = details.limit_name.as_ref() {
item.insert("limit_name".to_owned(), Value::String(name.clone()));
}
let mut value = item
.remove("rate_limit")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_details(&mut value, details);
item.insert("rate_limit".to_owned(), Value::Object(value));
item.insert("rate_limit".to_owned(), Value::Object(rate_limit));
match index {
Some(index) => additional[index] = Value::Object(item),
None => additional.push(Value::Object(item)),
@@ -1018,12 +1309,14 @@ fn merge_passive_quota(
quota
}
fn merge_rate_limit_details(target: &mut Map<String, Value>, details: &RateLimitDetails) {
fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option<Map<String, Value>> {
// 响应头的窗口属于同一次上游观测;跨响应合并会制造不存在的额度窗口。
let mut snapshot = Map::new();
if let Some(allowed) = details.allowed {
target.insert("allowed".to_owned(), Value::Bool(allowed));
snapshot.insert("allowed".to_owned(), Value::Bool(allowed));
}
if let Some(limit_reached) = details.limit_reached {
target.insert("limit_reached".to_owned(), Value::Bool(limit_reached));
snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached));
}
for (field, window) in [
("primary_window", details.primary),
@@ -1032,28 +1325,29 @@ fn merge_rate_limit_details(target: &mut Map<String, Value>, details: &RateLimit
let Some(window) = window else {
continue;
};
let mut value = target
.remove(field)
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_window(&mut value, window);
target.insert(field.to_owned(), Value::Object(value));
snapshot.insert(
field.to_owned(),
Value::Object(passive_rate_limit_window(window)),
);
}
(!snapshot.is_empty()).then_some(snapshot)
}
fn merge_rate_limit_window(target: &mut Map<String, Value>, window: RateLimitWindow) {
fn passive_rate_limit_window(window: RateLimitWindow) -> Map<String, Value> {
let mut snapshot = Map::new();
if let Some(number) = serde_json::Number::from_f64(window.used_percent) {
target.insert("used_percent".to_owned(), Value::Number(number));
snapshot.insert("used_percent".to_owned(), Value::Number(number));
}
if let Some(seconds) = window
.window_minutes
.and_then(|minutes| minutes.checked_mul(60))
{
target.insert("limit_window_seconds".to_owned(), Value::from(seconds));
snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds));
}
if let Some(reset_at) = window.reset_at {
target.insert("reset_at".to_owned(), Value::from(reset_at));
snapshot.insert("reset_at".to_owned(), Value::from(reset_at));
}
snapshot
}
fn quota_success_availability(
@@ -28,6 +28,16 @@ const MAX_REFRESH_BATCH: u32 = 1_000;
/// 连续失败计数窗口;每次瞬态失败刷新该 TTL,静默满窗后计数过期归零。
const REFRESH_BACKOFF_WINDOW: Duration = Duration::from_secs(30 * 60);
fn refresh_due_at(account: &ProviderAccount, now: SystemTime) -> Option<SystemTime> {
if let Some(expires_at) = account.access_token_expires_at()
&& expires_at <= now
{
// AT 已过期时不能继续等待持久退避;后续额度/请求探测都必须先用 RT 续出可用 AT。
return Some(expires_at);
}
account.next_refresh_at().filter(|next| *next <= now)
}
pub struct DueCodexCredential {
pub account: ProviderAccount,
pub secret: CodexOAuthSecret,
@@ -248,9 +258,9 @@ impl CodexCredentialRefreshService {
)
&& (account.availability() != AccountAvailability::Cooldown
|| account.cooldown_until().is_some_and(|until| until <= now))
&& account.next_refresh_at().is_some_and(|next| next <= now)
&& refresh_due_at(account, now).is_some()
});
accounts.sort_by_key(|account| (account.next_refresh_at(), account.id().clone()));
accounts.sort_by_key(|account| (refresh_due_at(account, now), account.id().clone()));
accounts.truncate(MAX_REFRESH_BATCH as usize);
let mut due = Vec::with_capacity(accounts.len());
let mut failures = Vec::new();
@@ -6,8 +6,9 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use gateway_core::engine::credential::{
AccountAvailability, AccountCandidate, AccountFeedbackStats, AccountRuntimeSignals,
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountRuntimeSignals, AccountSelectionContext, AccountSelector, ProviderAccount,
ProviderAccountId,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt};
use gateway_core::provider_ports::{
@@ -162,22 +163,26 @@ impl CodexCredentialSelector {
request: &SelectCodexCredential<'_>,
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let diagnostic = request.attempt.is_diagnostic_required_account();
let accounts = self.repository.list_for_provider().await?;
let accounts = accounts
.into_iter()
.filter(|account| {
account.provider() == &self.provider_kind
&& !matches!(
self.catalog.observed_model_support(
account.id(),
account.revision(),
request.upstream_model,
),
Ok(Some(false))
)
&& (diagnostic
|| !matches!(
self.catalog.observed_model_support(
account.id(),
account.revision(),
request.upstream_model,
),
Ok(Some(false))
))
})
.collect::<Vec<_>>();
self.quota.prepare_scheduling(&accounts).await;
if !diagnostic {
self.quota.prepare_scheduling(&accounts).await;
}
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
@@ -223,23 +228,27 @@ impl CodexCredentialSelector {
.collect::<Vec<_>>()
});
let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates);
let affinity_account = match request.session_affinity_key {
Some(key) => match self.lookup_session_affinity(key).await {
SessionAffinityLookup::Bound(account_id)
if scheduling_candidates.iter().any(|candidate| {
candidate.account.id() == &account_id
&& candidate.account.is_schedulable(SystemTime::now())
}) =>
{
Some(account_id)
}
SessionAffinityLookup::Bound(_) => {
self.clear_session_affinity(key).await;
None
}
SessionAffinityLookup::Missing | SessionAffinityLookup::Unavailable => None,
},
None => None,
let affinity_account = if diagnostic {
None
} else {
match request.session_affinity_key {
Some(key) => match self.lookup_session_affinity(key).await {
SessionAffinityLookup::Bound(account_id)
if scheduling_candidates.iter().any(|candidate| {
candidate.account.id() == &account_id
&& candidate.account.is_schedulable(SystemTime::now())
}) =>
{
Some(account_id)
}
SessionAffinityLookup::Bound(_) => {
self.clear_session_affinity(key).await;
None
}
SessionAffinityLookup::Missing | SessionAffinityLookup::Unavailable => None,
},
None => None,
}
};
let cyber_policy_scope = self
.prepare_cyber_policy_scope(cyber_policy_session_key)
@@ -290,6 +299,11 @@ impl CodexCredentialSelector {
excluded_accounts: excluded.clone(),
preferred_account: preferred.clone(),
round_robin_cursor,
availability: if diagnostic {
AccountAvailabilityPolicy::BypassForDiagnostic
} else {
AccountAvailabilityPolicy::Enforce
},
};
let Some(selected) = AccountSelector.select(scheduling_candidates, &context) else {
return match shortest_retry {
@@ -304,6 +318,7 @@ impl CodexCredentialSelector {
.find(|candidate| candidate.account.id() == selected.account.id())
.map(|candidate| candidate.account.clone())
.ok_or(CredentialSelectionError::InvalidCredential)?;
let allows_account_state_mutation = !diagnostic || account.enabled();
let policy = request.attempt.account_selection_policy();
match self
.leases
@@ -327,11 +342,12 @@ impl CodexCredentialSelector {
let (account, runtime) = if account.authentication_kind()
== CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY
{
let prepared = self
.agent_identity
.prepare(&account)
.await
.map_err(|_| CredentialSelectionError::InvalidCredential)?;
let prepared = if allows_account_state_mutation {
self.agent_identity.prepare(&account).await
} else {
self.agent_identity.load_current(account.id()).await
}
.map_err(|_| CredentialSelectionError::InvalidCredential)?;
(prepared.account, prepared.credential)
} else {
let runtime = self.repository.load_runtime_credential(&account).await?;
@@ -359,6 +375,7 @@ impl CodexCredentialSelector {
authentication: runtime.authentication,
cookies,
cyber_policy_scope,
allows_account_state_mutation,
_guard: guard,
});
}
@@ -535,13 +552,9 @@ impl CodexCredentialSelector {
CookieRecovery::None,
),
CodexAccountFailure::RateLimited { retry_after } => {
let until = now.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
(
AccountAvailability::QuotaExhausted,
Some("rate_limited".to_owned()),
until,
CookieRecovery::None,
)
return self
.record_rate_limit_failure(account, retry_after, now)
.await;
}
CodexAccountFailure::CloudflareChallenge { retry_after } => {
let delay = self.cloudflare_challenge_delay(account.id(), now, retry_after);
@@ -579,11 +592,55 @@ impl CodexCredentialSelector {
Ok(())
}
async fn record_rate_limit_failure(
&self,
account: &ProviderAccount,
retry_after: Option<Duration>,
observed_at: SystemTime,
) -> Result<(), CredentialSelectionError> {
let current = self.current_account(account.id()).await?;
if current.revision() != account.revision() {
return Err(CredentialRepositoryError::RevisionConflict.into());
}
if matches!(
current.availability(),
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::QuotaExhausted
) {
return Ok(());
}
let quota_exhausted = self
.quota
.has_active_exhaustion(&current, observed_at)
.await
.map_err(|_| CredentialSelectionError::Store)?;
let (availability, reason, cooldown_until) = if quota_exhausted {
(
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
)
} else {
(
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)),
)
};
self.repository
.apply_state(&current, availability, reason, cooldown_until, observed_at)
.await?;
Ok(())
}
pub async fn record_success(
&self,
account: &ProviderAccount,
session_affinity_key: Option<&ProviderSessionAffinityKey>,
) {
self.restore_recoverable_account_state(account).await;
self.risk_recovery
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
@@ -626,6 +683,42 @@ impl CodexCredentialSelector {
}
}
async fn restore_recoverable_account_state(&self, account: &ProviderAccount) {
let Ok(current) = self.current_account(account.id()).await else {
return;
};
if current.provider() != &self.provider_kind
|| current.revision() != account.revision()
|| !current.enabled()
|| !matches!(
current.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
return;
}
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::Ready,
None,
None,
SystemTime::now(),
)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
"OpenAI account state recovery after successful upstream response failed"
);
}
}
pub async fn current_account(
&self,
account_id: &ProviderAccountId,
@@ -801,6 +894,7 @@ pub struct CodexCredentialLease {
cookies: Vec<RuntimeCodexCookie>,
installation_id: String,
cyber_policy_scope: Option<CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -834,6 +928,12 @@ impl CodexCredentialLease {
pub(crate) const fn cyber_policy_scope(&self) -> Option<&CodexCyberPolicyScope> {
self.cyber_policy_scope.as_ref()
}
/// 禁用账号的管理端诊断必须只返回真实上游结果,不能回写账号侧状态。
#[must_use]
pub(crate) const fn allows_account_state_mutation(&self) -> bool {
self.allows_account_state_mutation
}
}
impl fmt::Debug for CodexCredentialLease {
@@ -160,7 +160,6 @@ pub async fn initialize(
let refresher: Arc<dyn TokenRefresher> = token_client.clone();
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client;
let credential_admin = Arc::new(CodexCredentialAdminService::new(
repository.clone(),
Arc::clone(&refresher),
Arc::clone(&identity),
Arc::clone(&leases),
+292 -131
View File
@@ -79,6 +79,8 @@ const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const WEBSOCKET_TRANSPORT: &str = "websocket";
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 与 2.x 预取器一致:提交边界前最多保留 64 KiB 原始上游 chunk。
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api";
pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api";
@@ -177,14 +179,9 @@ impl Provider for CodexProvider {
input.extend(encoded.input().iter().cloned());
encoded.set_input(input);
}
let reasoning_effort = encoded
.reasoning()
.and_then(|reasoning| reasoning.get("effort"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned);
(encoded.semantics(), reasoning_effort)
let semantics = encoded.semantics();
let reasoning_effort = semantics.reasoning_effort.clone();
(semantics, reasoning_effort)
})
.unwrap_or_default();
ProviderRequestObservation {
@@ -381,6 +378,7 @@ impl Provider for CodexProvider {
response_store,
continuation_scope: None,
});
let allows_account_state_mutation = lease.allows_account_state_mutation();
let events = cold_response_stream(ColdResponse {
client: self.client.clone(),
response_origin: self.responses_url.clone(),
@@ -396,8 +394,12 @@ impl Provider for CodexProvider {
session_affinity_key,
session_capture,
});
Ok(ProviderStream::new(metadata, events, lease)
.with_account_feedback(Arc::clone(&self.account_feedback)))
let stream = ProviderStream::new(metadata, events, lease);
Ok(if allows_account_state_mutation {
stream.with_account_feedback(Arc::clone(&self.account_feedback))
} else {
stream
})
}
}
@@ -662,6 +664,15 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
} = response;
Box::pin(async_stream::try_stream! {
let cyber_policy_scope = lease.cyber_policy_scope().cloned();
let allows_account_state_mutation = lease.allows_account_state_mutation();
let failure_context = OpenAiFailureContext {
client: &client,
selector: &selector,
quota: &quota,
response_origin: &response_origin,
cyber_policy_scope: cyber_policy_scope.as_ref(),
allows_account_state_mutation,
};
let mut active_account = lease.account().clone();
let cookie_header = build_cookie_header(lease.cookies())?;
let mut authorization = lease
@@ -675,6 +686,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
})?;
let request_id = context.request_id().as_str().to_owned();
let cancellation = context.cancellation().clone();
let stream_commit_policy =
StreamCommitPolicy::for_continuation(context.continuation_attempt());
let mut response = create_response_attempt(
&client,
&request,
@@ -691,7 +704,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
&cancellation,
)
.await;
if let Err(CodexHandshakeAttemptError::Client(error)) = &response {
if allows_account_state_mutation
&& let Err(CodexHandshakeAttemptError::Client(error)) = &response
{
match agent_identity
.recover_after_rejected_task(
active_account.id(),
@@ -740,15 +755,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if let Some(observation) = failure.observation.take() {
yield ProviderEvent::observation(observation);
}
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -759,15 +766,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
ProviderErrorKind::Protocol,
UpstreamSendState::Sent,
));
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -791,7 +790,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
context.attempt_index().get(),
);
yield ProviderEvent::observation(observation_state.observation()?);
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
}
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
{
@@ -800,7 +801,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
"OpenAI model ETag observation was rejected"
);
}
if !response.set_cookie_headers.is_empty()
if allows_account_state_mutation
&& !response.set_cookie_headers.is_empty()
&& let Ok(outcome) = selector
.capture_response_cookies(
&active_account,
@@ -826,6 +828,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
let mut decoder =
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
let mut pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
@@ -852,20 +855,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
if observation_state.merge_rate_limit_headers(&updates) {
yield ProviderEvent::observation(observation_state.observation()?);
}
}
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -876,7 +873,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
};
if let (Some(capture), Some(updates)) =
@@ -888,56 +887,48 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, context.timing_started_at());
let (mut events, terminal_failure) = match decoder.push(&chunk) {
let chunk_len = chunk.len();
let (mut events, canonical_failure) = match decoder.push(&chunk) {
CodexCanonicalOutcome::Events(events) => (events, None),
CodexCanonicalOutcome::Failed(failure) => {
let (events, error, semantic_output_seen) = failure.into_parts();
let failure = map_canonical_error(
(events, Some((error, semantic_output_seen)))
}
};
pre_commit_events.observe_chunk(chunk_len);
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
map_canonical_error(
error,
&failure_diagnostics,
&failure_set_cookie_headers,
&failure_rate_limit_headers,
ReplayBoundary::from_semantic_output(semantic_output_seen),
);
(events, Some(failure))
}
};
ReplayBoundary::from_semantic_output(
semantic_output_seen || pre_commit_events.is_committed(),
),
),
atomic_upstream_failure,
)
});
let timing_signals = decoder.take_timing_signals();
let timing_changed = first_event_changed
|| observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
.observe_timing_signals(timing_signals, context.timing_started_at());
let completed = events
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
@@ -945,56 +936,78 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
return;
}
for event in events {
yield event;
}
if let Some(failure) = terminal_failure {
Err(failure.error)?;
return;
}
if completed {
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 完成事件一旦交给下游,Core 可以立刻停止轮询 Provider stream;
// 在此之前持久化亲和关系,保证成功请求不会因流被提前 drop 而丢失绑定。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
timing_signals.semantic_output || pre_commit_events.is_committed();
if failure_after_commit {
for event in pre_commit_events.commit(events) {
yield event;
}
} else if atomic_upstream_failure {
failure.error = failure
.error
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
}
Err(failure.error)?;
return;
}
let events = pre_commit_events.stage(events, timing_signals, completed)?;
for event in events {
yield event;
}
if completed {
return;
}
}
let (mut events, terminal_failure) = match decoder.finish() {
let (mut events, canonical_failure) = match decoder.finish() {
CodexCanonicalOutcome::Events(events) => (events, None),
CodexCanonicalOutcome::Failed(failure) => {
let (events, error, semantic_output_seen) = failure.into_parts();
let failure = map_canonical_error(
(events, Some((error, semantic_output_seen)))
}
};
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
map_canonical_error(
error,
&failure_diagnostics,
&failure_set_cookie_headers,
&failure_rate_limit_headers,
ReplayBoundary::from_semantic_output(semantic_output_seen),
);
(events, Some(failure))
}
};
ReplayBoundary::from_semantic_output(
semantic_output_seen || pre_commit_events.is_committed(),
),
),
atomic_upstream_failure,
)
});
let timing_signals = decoder.take_timing_signals();
let timing_changed = observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
.observe_timing_signals(timing_signals, context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
};
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let (Some(capture), Some(updates)) =
@@ -1006,15 +1019,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
@@ -1028,17 +1033,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
for event in events {
yield event;
}
if let Some(failure) = terminal_failure {
Err(failure.error)?;
return;
}
if completed {
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
@@ -1046,6 +1042,28 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
timing_signals.semantic_output || pre_commit_events.is_committed();
if failure_after_commit {
for event in pre_commit_events.commit(events) {
yield event;
}
} else if atomic_upstream_failure {
failure.error = failure
.error
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
}
Err(failure.error)?;
return;
}
let events = pre_commit_events.finish(events, timing_signals, completed)?;
for event in events {
yield event;
}
})
}
@@ -1766,6 +1784,128 @@ impl MappedProviderFailure {
}
}
/// 与 2.x 一致的流提交策略:普通请求优先首事件,续写/重放优先可恢复性。
#[derive(Clone, Copy)]
enum StreamCommitPolicy {
FirstForwardableEvent,
UntilOutputOrTerminal,
}
impl StreamCommitPolicy {
const fn for_continuation(attempt: ContinuationAttempt) -> Self {
match attempt {
ContinuationAttempt::None => Self::FirstForwardableEvent,
ContinuationAttempt::Native
| ContinuationAttempt::ReplayOwner
| ContinuationAttempt::ReplayAny => Self::UntilOutputOrTerminal,
}
}
fn boundary_reached(
self,
events: &[ProviderEvent],
timing_signals: ResponseEventSignals,
completed: bool,
) -> bool {
match self {
Self::FirstForwardableEvent => events.iter().any(ProviderEvent::has_client_event),
Self::UntilOutputOrTerminal => timing_signals.semantic_output || completed,
}
}
}
/// 提交边界前的上游事件预取。
///
/// 原始 chunk 计数而不是重编码后的 event 大小,保持 2.x 的 64 KiB 资源边界;一旦
/// 提交,后续事件不再具备无痕重放资格。
struct PreCommitClientEvents {
policy: StreamCommitPolicy,
pending: Vec<ProviderEvent>,
prefetched_bytes: usize,
committed: bool,
}
impl PreCommitClientEvents {
const fn new(policy: StreamCommitPolicy) -> Self {
Self {
policy,
pending: Vec::new(),
prefetched_bytes: 0,
committed: false,
}
}
fn observe_chunk(&mut self, bytes: usize) {
if !self.committed {
self.prefetched_bytes = self.prefetched_bytes.saturating_add(bytes);
}
}
fn stage(
&mut self,
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
if self.committed {
return Ok(incoming);
}
let boundary_reached = self
.policy
.boundary_reached(&incoming, timing_signals, completed);
self.pending.extend(incoming);
if boundary_reached {
return Ok(self.commit_pending());
}
if self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
return Err(prefetch_protocol_error());
}
Ok(Vec::new())
}
fn finish(
&mut self,
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
let events = self.stage(incoming, timing_signals, completed)?;
if self.committed {
Ok(events)
} else {
Err(prefetch_protocol_error())
}
}
fn commit(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
if self.committed {
return incoming;
}
self.pending.extend(incoming);
self.commit_pending()
}
fn take_for_failure(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
self.pending.extend(incoming);
self.prefetched_bytes = 0;
std::mem::take(&mut self.pending)
}
const fn is_committed(&self) -> bool {
self.committed
}
fn commit_pending(&mut self) -> Vec<ProviderEvent> {
self.committed = true;
self.prefetched_bytes = 0;
std::mem::take(&mut self.pending)
}
}
fn prefetch_protocol_error() -> ProviderError {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
}
#[derive(Clone, Copy)]
enum ReplayBoundary {
BeforeSemanticOutput,
@@ -1786,24 +1926,40 @@ impl ReplayBoundary {
}
}
struct OpenAiFailureContext<'a> {
client: &'a CodexBackendClient,
selector: &'a CodexCredentialSelector,
quota: &'a CodexCredentialQuotaService,
response_origin: &'a Url,
cyber_policy_scope: Option<&'a CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
}
async fn apply_failure(
client: &CodexBackendClient,
selector: &CodexCredentialSelector,
quota: &CodexCredentialQuotaService,
context: &OpenAiFailureContext<'_>,
account: &ProviderAccount,
response_origin: &Url,
cyber_policy_scope: Option<&CodexCyberPolicyScope>,
failure: &MappedProviderFailure,
) {
synchronize_passive_quota(quota, account, &failure.rate_limit_headers).await;
if !context.allows_account_state_mutation {
return;
}
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
if failure.cyber_policy_failure {
selector
.record_cyber_policy_failure(cyber_policy_scope, account)
context
.selector
.record_cyber_policy_failure(context.cyber_policy_scope, account)
.await;
}
if let Some(account_failure) = failure.account_failure {
client.evict_websocket_account(account.id().as_str()).await;
if let Err(error) = selector.record_failure(account, account_failure).await {
context
.client
.evict_websocket_account(account.id().as_str())
.await;
if let Err(error) = context
.selector
.record_failure(account, account_failure)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
@@ -1813,8 +1969,13 @@ async fn apply_failure(
}
if failure.capture_response_cookies
&& !failure.set_cookie_headers.is_empty()
&& let Err(error) = selector
.capture_response_cookies(account, response_origin, &failure.set_cookie_headers)
&& let Err(error) = context
.selector
.capture_response_cookies(
account,
context.response_origin,
&failure.set_cookie_headers,
)
.await
{
tracing::warn!(
@@ -273,8 +273,20 @@ impl CodexCanonicalDecoder {
self.merge_timing_signals(signals);
if matches!(event_type, Some("response.failed" | "error")) {
let failure = ResponsesSseFailure::from_event(event_type.unwrap_or_default(), &value);
let mut canonical = Vec::new();
if !self.started {
// 失败首帧可能是唯一携带 response_id 的上游事实。身份投影只做
// 最佳努力:缺少 response/id 时仍必须保留原 typed failure。
let _ = self.start(&value, &mut canonical);
}
if let Some(wire) = Self::wire_for_event(event, value, raw_sse_frame) {
output.push(ProviderEvent::wire(wire));
output.push(if canonical.is_empty() {
ProviderEvent::wire(wire)
} else {
ProviderEvent::canonical_with_wire(canonical, wire)
});
} else {
output.extend(canonical.into_iter().map(ProviderEvent::canonical));
}
return Err(CodexCanonicalError::Upstream(failure));
}
@@ -35,7 +35,7 @@ use super::websocket::{
};
// ---------------------------------------------------------------------------
// Constants
// 常量
// ---------------------------------------------------------------------------
const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024;
@@ -82,7 +82,7 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
}
// ---------------------------------------------------------------------------
// Error types
// 错误类型
// ---------------------------------------------------------------------------
/// Codex 上游 HTTP 客户端错误。
@@ -245,7 +245,7 @@ pub type CodexBackendSseStream =
Pin<Box<dyn Stream<Item = CodexClientResult<Bytes>> + Send + 'static>>;
// ---------------------------------------------------------------------------
// Request context & response types
// 请求上下文与响应类型
// ---------------------------------------------------------------------------
/// 单次 Codex 上游请求的上下文。
@@ -253,7 +253,7 @@ pub type CodexBackendSseStream =
pub struct CodexRequestContext<'a> {
/// Provider 已构造并脱敏持有的完整 Authorization 值。
pub authorization: &'a str,
/// ChatGPT account id。
/// ChatGPT 账号 ID。
pub account_id: Option<&'a str>,
/// 代理请求 ID。
pub request_id: &'a str,
@@ -261,11 +261,11 @@ pub struct CodexRequestContext<'a> {
pub turn_state: Option<&'a str>,
/// 客户端 turn metadata;其中 installation ID 已按当前账号处理。
pub turn_metadata: Option<&'a str>,
/// x-codex-beta-features。
/// x-codex-beta-features 头值。
pub beta_features: Option<&'a str>,
/// x-responsesapi-include-timing-metrics。
/// x-responsesapi-include-timing-metrics 头值。
pub include_timing_metrics: Option<&'a str>,
/// version。
/// 客户端版本号。
pub version: Option<&'a str>,
/// 客户端 window ID。
pub codex_window_id: Option<&'a str>,
@@ -335,9 +335,9 @@ impl fmt::Debug for CodexRequestContext<'_> {
/// Codex Responses 实际使用的上游传输。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodexBackendTransport {
/// HTTP SSE transport.
/// HTTP SSE 传输。
HttpSse,
/// WebSocket transport.
/// WebSocket 传输。
WebSocket,
}
@@ -382,10 +382,10 @@ pub struct CodexTransportMetrics {
pub http_version: Option<String>,
}
/// Live stream rate-limit updates captured after the response headers.
/// 响应头之后在 live 流中采集的限流更新。
pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates;
/// Live stream turn-state updates captured after the response headers.
/// 响应头之后在 live 流中采集的 turn state 更新。
pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate;
/// Codex Responses 上游 live SSE 响应。
@@ -449,7 +449,7 @@ pub(super) struct PreparedWebSocketRoute {
}
// ---------------------------------------------------------------------------
// Response helpers
// 响应辅助函数
// ---------------------------------------------------------------------------
pub(super) fn log_websocket_pool_decision(
@@ -587,7 +587,7 @@ pub(super) async fn read_capped_error_body(
}
// ---------------------------------------------------------------------------
// Request helpers
// 请求辅助函数
// ---------------------------------------------------------------------------
pub(super) fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponsesRequest {
@@ -649,7 +649,7 @@ pub(super) fn openai_subagent_from_metadata(client_metadata: Option<&Value>) ->
}
// ---------------------------------------------------------------------------
// Error conversion helpers
// 错误转换辅助函数
// ---------------------------------------------------------------------------
pub(super) fn websocket_exchange_error_to_client_error(
@@ -1,4 +1,4 @@
//! Upstream response diagnostics captured at the transport boundary.
//! 在 transport 边界采集的上游响应诊断信息。
use super::protocol::responses::ResponsesSseFailure;
use base64::{Engine as _, engine::general_purpose::STANDARD};
@@ -415,15 +415,6 @@ fn classify_upstream_failure(
if status == Some(StatusCode::NOT_FOUND) && body.trim().is_empty() {
return CodexFailureCategory::CloudflarePathBlocked;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_rate_limit_signal)
|| is_rate_limit_message(&message)
|| is_rate_limit_message(&body)
|| status == Some(StatusCode::TOO_MANY_REQUESTS)
{
return CodexFailureCategory::RateLimited;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_quota_signal)
@@ -433,6 +424,15 @@ fn classify_upstream_failure(
{
return CodexFailureCategory::QuotaExhausted;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_rate_limit_signal)
|| is_rate_limit_message(&message)
|| is_rate_limit_message(&body)
|| status == Some(StatusCode::TOO_MANY_REQUESTS)
{
return CodexFailureCategory::RateLimited;
}
if is_upstream_overload(&code) || is_upstream_overload(&message) {
return CodexFailureCategory::Unavailable;
}
@@ -1,13 +1,12 @@
use std::fmt;
use gateway_protocol::openai::events;
use gateway_protocol::openai::{
CodexResponsesRequestSemantics as CodexRequestSemantics,
codex_responses_request_semantics_with_turn_metadata, events,
};
use serde::Serialize;
use serde_json::{Map, Value};
const MULTI_AGENT_MODE_OPEN_TAG: &str = "<multi_agent_mode>";
const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>";
const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active.";
/// Codex Responses 上游请求体。
///
/// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object,逐字段(含顺序、
@@ -60,7 +59,7 @@ pub struct CodexResponsesRequest {
pub responses_lite: Option<String>,
/// Memory consolidation 请求语义;HTTP 与 WebSocket opening 均使用 header。
pub memgen_request: Option<String>,
/// codex window id。
/// Codex window ID。
pub codex_window_id: Option<String>,
/// 父线程 ID。
pub parent_thread_id: Option<String>,
@@ -93,19 +92,6 @@ pub enum PreviousResponseScope {
ExternalUnknown,
}
/// Responses 请求携带的 Codex 运行语义。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexRequestSemantics {
/// Codex turn metadata 中的请求类型。
pub request_kind: Option<String>,
/// Codex turn metadata 中的子代理类型。
pub subagent_kind: Option<String>,
/// Codex 客户端选择的推理预设。
pub reasoning_preset: Option<&'static str>,
/// 请求是否为远端压缩。
pub compact: bool,
}
impl Serialize for CodexResponsesRequest {
/// 上游 body 序列化即原始 `body` map(HTTP SSE 直发;WebSocket 在外层前置 `type`)。
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
@@ -608,34 +594,12 @@ impl CodexResponsesRequest {
/// 提取 Codex 请求类型、子代理类型、推理预设与压缩语义。
pub fn semantics(&self) -> CodexRequestSemantics {
let turn_metadata = self.turn_metadata.as_deref().or_else(|| {
self.client_metadata()?
.get("x-codex-turn-metadata")?
.as_str()
});
let effort = self
.reasoning()
.and_then(|value| value.get("effort"))
.and_then(Value::as_str);
let proactive_multi_agent = self
.latest_multi_agent_mode()
.is_some_and(|mode| mode.starts_with(PROACTIVE_MULTI_AGENT_MODE_PREFIX));
let compact_trigger = self
.input()
.iter()
.any(|item| item.get("type").and_then(Value::as_str) == Some("compaction_trigger"));
codex_request_semantics_from_parts(
turn_metadata,
effort,
proactive_multi_agent,
compact_trigger,
codex_responses_request_semantics_with_turn_metadata(
self.body(),
self.turn_metadata.as_deref(),
)
}
fn latest_multi_agent_mode(&self) -> Option<&str> {
latest_multi_agent_mode(self.input().iter())
}
/// 设置 / 合并 client metadata。
pub fn set_client_metadata(&mut self, client_metadata: Option<Value>) {
match client_metadata {
@@ -664,61 +628,3 @@ impl CodexResponsesRequest {
}
}
}
fn non_empty_owned_string(value: &str) -> Option<String> {
let value = value.trim();
(!value.is_empty()).then(|| value.to_string())
}
pub(crate) fn codex_request_semantics_from_parts(
turn_metadata: Option<&str>,
effort: Option<&str>,
proactive_multi_agent: bool,
compact_trigger: bool,
) -> CodexRequestSemantics {
let turn_metadata = turn_metadata.and_then(|value| serde_json::from_str::<Value>(value).ok());
let request_kind = turn_metadata
.as_ref()
.and_then(|metadata| metadata.get("request_kind"))
.and_then(Value::as_str)
.and_then(non_empty_owned_string);
let subagent_kind = turn_metadata
.as_ref()
.and_then(|metadata| metadata.get("subagent_kind"))
.and_then(Value::as_str)
.and_then(non_empty_owned_string);
let reasoning_preset =
(subagent_kind.is_none() && effort == Some("max") && proactive_multi_agent)
.then_some("ultra");
let compact = request_kind.as_deref() == Some("compaction") || compact_trigger;
CodexRequestSemantics {
request_kind,
subagent_kind,
reasoning_preset,
compact,
}
}
fn latest_multi_agent_mode<'a>(
input: impl DoubleEndedIterator<Item = &'a Value>,
) -> Option<&'a str> {
input.rev().find_map(|item| {
if item.get("role").and_then(Value::as_str) != Some("developer") {
return None;
}
item.get("content")?
.as_array()?
.iter()
.rev()
.filter_map(|content| content.get("text").and_then(Value::as_str))
.find_map(multi_agent_mode_from_text)
})
}
fn multi_agent_mode_from_text(text: &str) -> Option<&str> {
let close = text.rfind(MULTI_AGENT_MODE_CLOSE_TAG)?;
let before_close = &text[..close];
let open = before_close.rfind(MULTI_AGENT_MODE_OPEN_TAG)?;
let body = &before_close[open + MULTI_AGENT_MODE_OPEN_TAG.len()..];
Some(body.trim())
}
@@ -13,7 +13,7 @@ pub struct OpeningAuditSnapshot {
pub request_line: String,
/// 请求头顺序。
pub header_order: Vec<String>,
/// 红action后的请求头。
/// 脱敏后的请求头。
pub headers: Vec<OpeningAuditHeader>,
}
@@ -31,11 +31,11 @@ pub struct OpeningAuditHeader {
pub struct PayloadAuditSnapshot {
/// 按构造顺序记录的顶层字段。
pub top_level_keys: Vec<String>,
/// 红action后的 JSON payload。
/// 脱敏后的 JSON payload。
pub body: Value,
}
/// WebSocket audit artifact.
/// WebSocket 审计产物。
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct WebSocketAuditArtifact {
/// 实际选择的传输模式。
@@ -10,7 +10,7 @@ use crate::transport::diagnostics::CodexUpstreamSendPhase;
use super::PreviousResponseUnavailableReason;
/// Responses WebSocket exchange error.
/// Responses WebSocket 交互错误。
#[derive(Debug, Error)]
pub enum CodexWebSocketExchangeError {
/// opening request 无法构造。
@@ -98,7 +98,7 @@ pub enum CodexWebSocketExchangeError {
/// WebSocket 上游错误帧载荷。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexWebSocketUpstreamError {
/// HTTP-style upstream status code.
/// 上游返回的 HTTP 风格状态码。
pub status_code: u16,
/// 推导出的重试秒数。
pub retry_after_seconds: Option<u64>,
@@ -24,11 +24,11 @@ const WEBSOCKET_STREAM_BUFFER: usize = 16;
pub(super) use self::stream::{WebSocketStreamPoolReturn, stream_websocket_response};
/// Responses WebSocket live SSE exchange result.
/// Responses WebSocket live SSE 交互结果。
pub struct CodexWebSocketStreamingExchange {
/// 关联请求日志与底层 pump 生命周期日志的连接标识。
pub(crate) websocket_connection_id: Uuid,
/// Live SSE bytes converted from WebSocket events.
/// 由 WebSocket 事件转换出的 live SSE 字节流。
pub body: CodexWebSocketSseStream,
/// 上游 metadata 帧中的最新 turn state。
pub turn_state: Option<String>,
@@ -50,12 +50,12 @@ pub struct CodexWebSocketStreamingExchange {
pub response_metadata: CodexResponseMetadata,
}
/// Responses WebSocket live SSE byte stream.
/// Responses WebSocket live SSE 字节流。
pub type CodexWebSocketSseStream =
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
/// WebSocket live stream rate-limit header updates.
/// live 流中的限流头动态更新。
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
/// WebSocket live stream turn-state update.
/// live 流中的 turn state 动态更新。
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
pub(super) fn reusable_websocket_metadata(
@@ -1,4 +1,4 @@
//! Codex HTTP Responses WebSocket transport。
//! Codex Responses WebSocket 传输。
mod audit;
mod breaker;
@@ -13,7 +13,7 @@ pub struct CodexWebSocketConnection {
pub(super) headers: Vec<(String, String)>,
}
/// Prepared Responses WebSocket request descriptor.
/// 已构造完成的 Responses WebSocket 请求描述。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexWebSocketRequest {
pub(super) connection: CodexWebSocketConnection,
+202 -20
View File
@@ -33,10 +33,12 @@ use gateway_core::operation::OperationKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingTakeOutcome,
NewOAuthPendingFlow, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCredentialState, ProviderCredentialStatePort, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts,
ProviderCooldownScope, ProviderCredentialState, ProviderCredentialStatePort,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError,
ProviderStorePorts,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
@@ -329,7 +331,6 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
let started = bundle
.admin_provider()
.start_authorization(PendingAuthorizationMutation::new(
Revision::new(7).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Create {
name: "OAuth account".to_owned(),
@@ -340,17 +341,28 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
.expect("start authorization");
{
let values = pending.values.lock().expect("OAuth pending");
let (_, payload, _) = values.values().next().expect("stored pending");
let (_, payload, _, _) = values.values().next().expect("stored pending");
let mutation = payload
.expose_to_provider()
.get("mutation")
.and_then(Value::as_object)
.expect("pending mutation");
assert!(
payload
.expose_to_provider()
.get("reauthorization_credential_revision")
.is_none()
);
assert_eq!(
mutation.get("schema_version").and_then(Value::as_u64),
Some(3)
);
assert!(mutation.get("expected_config_revision").is_none());
assert!(
mutation
.get("expected_config_revision")
.and_then(Value::as_u64),
Some(7)
.get("target")
.and_then(Value::as_object)
.is_some_and(|target| target.get("expected_credential_revision").is_none())
);
assert_eq!(
mutation.get("started_request_id").and_then(Value::as_str),
@@ -375,6 +387,81 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
assert_eq!(pending.values.lock().expect("OAuth pending").len(), 1);
}
#[tokio::test]
async fn openai_reauthorization_pending_payload_contains_only_stable_account_identity() {
let accounts = Arc::new(MemoryAccountStore::default());
accounts
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_pending_reauth".to_owned(),
name: "pending reauthorization".to_owned(),
secret: secret("pending-reauth-access"),
verified_account: profile("chatgpt-pending-reauth"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let pending = Arc::new(TestOAuthPending::default());
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(accounts, Arc::clone(&pending)),
)
.await
.expect("OpenAI bundle");
let context = MutationContext {
actor: MutationActor::AdminApiKey,
request_id: "request-pending-reauth".to_owned(),
};
bundle
.admin_provider()
.start_authorization(PendingAuthorizationMutation::new(
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new("acct_pending_reauth").expect("account id"),
},
AuthorizationOwnerBinding::from_context(&context),
))
.await
.expect("start reauthorization");
let values = pending.values.lock().expect("OAuth pending");
let (_, payload, _, _) = values.values().next().expect("stored pending");
let document = payload.expose_to_provider();
let mutation = document
.get("mutation")
.and_then(Value::as_object)
.expect("pending mutation");
let target = mutation
.get("target")
.and_then(Value::as_object)
.expect("pending target");
assert_eq!(
mutation.get("schema_version").and_then(Value::as_u64),
Some(3)
);
assert_eq!(
document
.get("reauthorization_account_id")
.and_then(Value::as_str),
Some("acct_pending_reauth")
);
assert!(
document
.get("reauthorization_credential_revision")
.is_none()
);
assert_eq!(
target.get("kind").and_then(Value::as_str),
Some("reauthorize")
);
assert_eq!(
target.get("account_id").and_then(Value::as_str),
Some("acct_pending_reauth")
);
assert!(target.get("expected_credential_revision").is_none());
}
#[tokio::test]
async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export() {
let store = Arc::new(MemoryAccountStore::default());
@@ -573,10 +660,20 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
.await
.expect_err("invalid import");
assert_eq!(import_error.kind(), ProviderAdminErrorKind::Invalid);
let mut stale_record = record.clone();
stale_record.name = "stale name".to_owned();
stale_record.email = None;
stale_record.plan_type = None;
stale_record.credential_revision = Revision::new(99).expect("stale revision");
stale_record.has_refresh_token = false;
stale_record.access_token_expires_at = None;
stale_record.next_refresh_at = None;
stale_record.enabled = false;
stale_record.availability = AdminAccountAvailability::Banned;
stale_record.cooldown_until = Some(Utc::now() + chrono::Duration::minutes(5));
let rotation_error = admin
.prepare_rotation(PrepareCredentialRotation {
account: record.clone(),
expected_credential_revision: record.credential_revision,
account: stale_record,
provider_material: ProviderDocument::new(OpaqueProviderData::new(Map::new())),
})
.await
@@ -805,6 +902,30 @@ impl ProviderCooldownPort for TestCooldown {
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn put_scoped_if_later(
&self,
_cooldown: ProviderScopedCooldown,
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn read_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>> {
Box::pin(async { Ok(None) })
}
fn clear_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
_through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
}
struct TestRuntimePolicy;
@@ -828,7 +949,7 @@ struct TestOAuthPending {
}
type PendingKey = (String, String);
type PendingValue = (String, OpaqueProviderData, SystemTime);
type PendingValue = (String, OpaqueProviderData, SystemTime, Option<String>);
impl OAuthPendingFlowPort for TestOAuthPending {
fn put_if_absent(
@@ -850,36 +971,97 @@ impl OAuthPendingFlowPort for TestOAuthPending {
flow.owner().expose_to_store().to_owned(),
flow.payload().clone(),
SystemTime::now() + flow.ttl(),
None,
),
);
Ok(OAuthPendingPutOutcome::Stored)
})
}
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>> {
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
_claim_ttl: Duration,
) -> BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, expires_at)) = values.get(&key) else {
return Ok(OAuthPendingTakeOutcome::NotFound);
let Some((stored_owner, payload, expires_at, stored_claim)) = values.get_mut(&key)
else {
return Ok(OAuthPendingClaimOutcome::NotFound);
};
if expires_at <= &SystemTime::now() {
if *expires_at <= SystemTime::now() {
values.remove(&key);
return Ok(OAuthPendingTakeOutcome::NotFound);
return Ok(OAuthPendingClaimOutcome::NotFound);
}
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingTakeOutcome::OwnerMismatch);
return Ok(OAuthPendingClaimOutcome::OwnerMismatch);
}
let (_, payload, _) = values.remove(&key).expect("checked pending");
Ok(OAuthPendingTakeOutcome::Taken(payload))
if stored_claim.is_some() {
return Ok(OAuthPendingClaimOutcome::InProgress);
}
*stored_claim = Some(claim.expose_to_store().to_owned());
Ok(OAuthPendingClaimOutcome::Claimed(payload.clone()))
})
}
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get_mut(&key) else {
return Ok(OAuthPendingReleaseOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingReleaseOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingReleaseOutcome::ClaimMismatch);
}
*stored_claim = None;
Ok(OAuthPendingReleaseOutcome::Released)
})
}
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get(&key) else {
return Ok(OAuthPendingConsumeOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingConsumeOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingConsumeOutcome::ClaimMismatch);
}
values.remove(&key);
Ok(OAuthPendingConsumeOutcome::Consumed)
})
}
}
@@ -11,7 +11,7 @@ use ed25519_dalek::pkcs8::EncodePrivateKey as _;
use futures::future::BoxFuture;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, PlaintextCredential,
ProviderAccount, ProviderAccountId,
ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest,
@@ -456,7 +456,6 @@ async fn manual_refresh_fixture(
available,
});
let service = CodexCredentialAdminService::new(
store.repository(),
refresher.clone(),
Arc::new(ManualVerifier),
leases.clone(),
@@ -478,11 +477,12 @@ async fn manual_refresh_prepares_revision_fenced_rotation_without_store_mutation
let (store, refresher, leases, service) =
manual_refresh_fixture(Ok(refreshed_tokens()), true).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let current = store
.load_current_credential(&account_id)
.await
.expect("current credential");
let prepared = service
.manual_refresh(
account_id.clone(),
CredentialRevision::new(1).expect("revision"),
)
.manual_refresh(current)
.await
.expect("manual refresh");
@@ -528,24 +528,16 @@ async fn manual_refresh_prepares_revision_fenced_rotation_without_store_mutation
}
#[tokio::test]
async fn manual_refresh_stale_revision_and_missing_lease_fail_before_exchange() {
let (_, stale_refresher, _, stale_service) =
manual_refresh_fixture(Ok(refreshed_tokens()), true).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let stale = stale_service
.manual_refresh(
account_id.clone(),
CredentialRevision::new(2).expect("revision"),
)
.await
.expect_err("stale revision");
assert_eq!(stale, CodexCredentialAdminError::RevisionConflict);
assert!(stale_refresher.seen.lock().expect("seen tokens").is_empty());
let (_, unavailable_refresher, _, unavailable_service) =
async fn manual_refresh_missing_lease_fails_before_exchange() {
let (store, unavailable_refresher, _, unavailable_service) =
manual_refresh_fixture(Ok(refreshed_tokens()), false).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let current = store
.load_current_credential(&account_id)
.await
.expect("current credential");
let unavailable = unavailable_service
.manual_refresh(account_id, CredentialRevision::new(1).expect("revision"))
.manual_refresh(current)
.await
.expect_err("lease unavailable");
assert_eq!(
@@ -606,9 +598,7 @@ impl CodexAccountIdentityVerifier for ImportVerifier {
pub(super) fn import_service(
refresher: Arc<ManualRefresher>,
) -> provider_openai::credential::CodexCredentialAdminService {
let store = Arc::new(MemoryAccountStore::default());
CodexCredentialAdminService::new(
store.repository(),
refresher,
Arc::new(ImportVerifier),
Arc::new(ManualLeases {
@@ -628,7 +618,7 @@ pub(super) fn unused_import_refresher() -> Arc<ManualRefresher> {
}
#[tokio::test]
async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
async fn oauth_import_uses_only_tokens_and_verified_identity() {
let refresher = unused_import_refresher();
let prepared = import_service(refresher.clone())
.prepare_import_document(serde_json::json!({
@@ -645,15 +635,16 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
"accessTokenExpiresAt": "2100-01-01T00:00:00+00:00",
"status": "disabled",
"addedAt": "2026-07-18T10:47:01+08:00",
"updatedAt": "2026-07-19T11:00:00+08:00"
"updatedAt": "2026-07-19T11:00:00+08:00",
"unrelated": { "value": true }
}]
}))
.await
.expect("CPR import");
assert_eq!(prepared.accounts().len(), 1);
let account = &prepared.accounts()[0];
assert_eq!(account.account.id().as_str(), "acct_cpr_import");
assert!(!account.account.enabled());
assert!(account.account.id().as_str().starts_with("acct_"));
assert!(account.account.enabled());
assert_eq!(account.account.upstream_account_id(), Some("chatgpt-cpr"));
let runtime = CodexCredentialCodec::decode(&account.credential).expect("credential");
assert_eq!(
@@ -666,33 +657,6 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
"token-cpr"
);
assert!(refresher.seen.lock().expect("seen tokens").is_empty());
let error = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{"token": "token-cpr", "unexpected": true}]
}))
.await
.expect_err("unknown CPR account key");
assert_eq!(error, CodexCredentialAdminError::InvalidInput);
}
#[tokio::test]
async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_projections() {
let error = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_verified_identity",
"accountId": "user-chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"token": "token-cpr"
}]
}))
.await
.expect_err("stale document identity must not override authenticated identity");
assert_eq!(error, CodexCredentialAdminError::IdentityRejected);
}
#[tokio::test]
@@ -744,7 +708,7 @@ async fn cpr_batch_allows_distinct_users_in_the_same_workspace() {
}
#[tokio::test]
async fn credential_bundle_and_auth_document_normalize_to_the_same_core_accounts() {
async fn oauth_import_extracts_tokens_without_parsing_source_projections() {
let bundle = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"exported_at": "2026-07-03T15:46:38.717Z",
@@ -754,7 +718,7 @@ async fn credential_bundle_and_auth_document_normalize_to_the_same_core_accounts
"platform": "openai",
"type": "oauth",
"credentials": {
"at": "token-bundle",
"access_token": "token-bundle",
"refresh_token": "refresh-bundle",
"chatgpt_account_id": "chatgpt-bundle",
"chatgpt_user_id": "user-chatgpt-bundle"
@@ -832,7 +796,6 @@ async fn real_cpr_fixture_import_contract() {
&std::fs::read(fixture).expect("read CPR fixture"),
)
.expect("parse CPR fixture");
let store = Arc::new(MemoryAccountStore::default());
let refresher = Arc::new(
provider_openai::credential::token_client::official_openai_token_client()
.expect("official token client"),
@@ -860,7 +823,6 @@ async fn real_cpr_fixture_import_contract() {
);
let verifier = Arc::new(CodexAccountIdentityService::new(signed, accounts));
let service = CodexCredentialAdminService::new(
store.repository(),
refresher,
verifier,
Arc::new(ManualLeases {
@@ -7,7 +7,8 @@ use chrono::Utc;
use futures::executor::block_on;
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountSelectionPolicy,
ProviderAccountId, RotationStrategy,
AccountStateChange, OpaqueProviderData, ProviderAccountId, ProviderAccountStore as _,
QuotaObservation, RotationStrategy,
};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
@@ -23,6 +24,7 @@ use provider_openai::credential::{
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use secrecy::ExposeSecret;
use serde_json::json;
use url::Url;
use crate::support::{
@@ -603,7 +605,7 @@ fn credential_expired_failure_marks_unified_account_expired() {
}
#[test]
fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -629,7 +631,7 @@ fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
.expect("record rate-limit failure");
let account = store.account("acct_primary").expect("account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
@@ -637,6 +639,171 @@ fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
);
}
#[test]
fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist exhausted quota");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[test]
fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() - 1;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist expired quota exhaustion");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
);
}
#[test]
fn rate_limited_failure_preserves_terminal_account_state() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
block_on(store.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Invalid,
reason: Some("identity_verification_required".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("mark account invalid");
let current = store.account("acct_primary").expect("invalid account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&current,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Invalid
);
}
#[test]
fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let first = store.account("acct_first").expect("first account");
let second = store.account("acct_second").expect("second account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(async {
let (first_result, second_result) = futures::join!(
selector.record_failure(
&first,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
),
selector.record_failure(
&second,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
),
);
first_result.expect("record first account failure");
second_result.expect("record second account failure");
});
assert_eq!(
[
store
.account("acct_first")
.expect("persisted first account")
.availability(),
store
.account("acct_second")
.expect("persisted second account")
.availability(),
],
[AccountAvailability::Cooldown; 2]
);
}
#[test]
fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip_it() {
let store = Arc::new(MemoryAccountStore::default());
@@ -264,6 +264,55 @@ async fn account_identity_rejects_imported_workspace_conflict() {
assert_eq!(error, CodexIdentityVerificationError::Rejected);
}
#[tokio::test]
async fn import_identity_uses_signed_claims_when_usage_rejects_the_request() {
let (signed, _) = verifier();
let service = CodexAccountIdentityService::new(
Arc::new(signed),
Arc::new(StaticAccountSource {
result: Err(CodexIdentityVerificationError::Rejected),
}),
);
let profile = service
.verify_import(&secret(
&claims(get_current_timestamp() + 3_600),
Some("refresh-token"),
))
.await
.expect("signed import identity");
assert_eq!(profile.chatgpt_account_id, "account-signed");
assert_eq!(profile.chatgpt_user_id, "user-signed");
}
#[tokio::test]
async fn import_identity_keeps_signed_workspace_when_usage_returns_a_different_workspace() {
let (signed, _) = verifier();
let service = CodexAccountIdentityService::new(
Arc::new(signed),
Arc::new(StaticAccountSource {
result: Ok(CodexAuthenticatedAccount {
chatgpt_account_id: "account-default".to_owned(),
chatgpt_user_id: "user-default".to_owned(),
email: None,
plan_type: None,
}),
}),
);
let profile = service
.verify_import(&secret(
&claims(get_current_timestamp() + 3_600),
Some("refresh-token"),
))
.await
.expect("signed import identity");
assert_eq!(profile.chatgpt_account_id, "account-signed");
assert_eq!(profile.chatgpt_user_id, "user-signed");
}
#[tokio::test]
async fn account_identity_keeps_usage_unavailable_out_of_ready_profile() {
let (signed, _) = verifier();
@@ -5,7 +5,7 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::model::{MutationActor, MutationContext};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::routing::ProviderKind;
use provider_openai::credential::token_client::{
@@ -15,10 +15,10 @@ use provider_openai::credential::token_client::{
use provider_openai::credential::{
CodexAccountIdentityVerifier, CodexCredentialAdmin, CodexIdentityExpectation,
CodexIdentityVerification, CodexIdentityVerificationError, CodexOAuthAdmin,
CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthPendingStore,
CodexOAuthPendingStoreError, CodexOAuthSecret, CodexPendingAuthorization,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ImportCodexOAuthCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthPendingClaimOutcome,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthSecret,
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential,
ImportCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
};
use secrecy::{ExposeSecret, SecretString};
use sha2::{Digest as _, Sha256};
@@ -28,7 +28,24 @@ use crate::support::{MemoryAccountStore, profile, runtime_policy, secret};
#[derive(Default)]
struct PendingStore {
value: Mutex<Option<StoredCodexPendingAuthorization>>,
value: Mutex<Option<(StoredCodexPendingAuthorization, Option<String>)>>,
}
fn duplicate_pending(
pending: &StoredCodexPendingAuthorization,
) -> Result<CodexPendingAuthorization, CodexOAuthPendingStoreError> {
CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization {
flow_id: pending.flow_id.clone(),
owner_ref: pending.owner_ref.clone(),
started_request_ref: pending.started_request_ref.clone(),
name: pending.name.clone(),
expires_at: pending.expires_at,
state: pending.state.clone(),
nonce: pending.nonce.clone(),
code_verifier: pending.code_verifier.clone(),
reauthorization_account_id: pending.reauthorization_account_id.clone(),
mutation: pending.mutation.clone(),
})
}
#[async_trait]
@@ -41,42 +58,85 @@ impl CodexOAuthPendingStore for PendingStore {
if value.is_some() {
return Err(CodexOAuthPendingStoreError::Conflict);
}
*value = Some(StoredCodexPendingAuthorization {
flow_id: pending.flow_id().to_owned(),
owner_ref: pending.owner_ref().to_owned(),
started_request_ref: pending.started_request_ref().to_owned(),
name: pending.name().to_owned(),
expires_at: pending.expires_at(),
state: pending.state().clone(),
nonce: pending.nonce().clone(),
code_verifier: pending.code_verifier().clone(),
reauthorization_account_id: pending
.reauthorization()
.map(|target| target.account_id().to_string()),
reauthorization_credential_revision: pending
.reauthorization()
.map(|target| target.credential_revision().get()),
mutation: pending.mutation().clone(),
});
*value = Some((
StoredCodexPendingAuthorization {
flow_id: pending.flow_id().to_owned(),
owner_ref: pending.owner_ref().to_owned(),
started_request_ref: pending.started_request_ref().to_owned(),
name: pending.name().to_owned(),
expires_at: pending.expires_at(),
state: pending.state().clone(),
nonce: pending.nonce().clone(),
code_verifier: pending.code_verifier().clone(),
reauthorization_account_id: pending.reauthorization().map(ToString::to_string),
mutation: pending.mutation().clone(),
},
None,
));
Ok(())
}
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError> {
claim_ref: &str,
_claim_ttl: std::time::Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
if value
.as_ref()
.is_some_and(|pending| pending.owner_ref == owner_ref && pending.flow_id == flow_id)
{
return value
.take()
.map(CodexPendingAuthorization::from_stored)
.transpose();
let Some((pending, claim)) = value.as_mut() else {
return Ok(CodexOAuthPendingClaimOutcome::NotFound);
};
if pending.owner_ref != owner_ref || pending.flow_id != flow_id {
return Ok(CodexOAuthPendingClaimOutcome::NotFound);
}
Ok(None)
if claim.is_some() {
return Ok(CodexOAuthPendingClaimOutcome::InProgress);
}
*claim = Some(claim_ref.to_owned());
duplicate_pending(pending)
.map(Box::new)
.map(CodexOAuthPendingClaimOutcome::Claimed)
}
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
let Some((pending, claim)) = value.as_mut() else {
return Ok(false);
};
if pending.owner_ref != owner_ref
|| pending.flow_id != flow_id
|| claim.as_deref() != Some(claim_ref)
{
return Ok(false);
}
*claim = None;
Ok(true)
}
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
let Some((pending, claim)) = value.as_ref() else {
return Ok(false);
};
if pending.owner_ref != owner_ref
|| pending.flow_id != flow_id
|| claim.as_deref() != Some(claim_ref)
{
return Ok(false);
}
*value = None;
Ok(true)
}
}
@@ -172,7 +232,6 @@ fn owner_context(request_id: &str) -> MutationContext {
fn create_mutation(request_id: &str) -> PendingAuthorizationMutation {
let context = owner_context(request_id);
PendingAuthorizationMutation::new(
Revision::new(1).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Create {
name: "OAuth Account".to_owned(),
@@ -184,11 +243,9 @@ fn create_mutation(request_id: &str) -> PendingAuthorizationMutation {
fn reauthorization_mutation(request_id: &str, account_id: &str) -> PendingAuthorizationMutation {
let context = owner_context(request_id);
PendingAuthorizationMutation::new(
Revision::new(1).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id).expect("account id"),
expected_credential_revision: Revision::new(1).expect("revision"),
},
AuthorizationOwnerBinding::from_context(&context),
)
@@ -268,7 +325,130 @@ async fn completion_returns_verified_core_account_without_writing_store() {
}
#[tokio::test]
async fn callback_state_mismatch_fails_closed_and_consumes_flow() {
async fn completion_accepts_unrelated_oauth_callback_parameters() {
let service = service();
let started = started(&service).await;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let prepared = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}&scope=openid%20profile&iss=https%3A%2F%2Fauth.openai.com"
)),
})
.await
.expect("complete OAuth with standard callback parameters");
assert!(matches!(
prepared.credential,
CompletedCodexOAuthCredential::Create(_)
));
}
#[tokio::test]
async fn completion_rejects_repeated_code_callback_parameter() {
let service = service();
let started = started(&service).await;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&code=another-code&state={state}"
)),
})
.await
.expect_err("repeated code");
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
#[tokio::test]
async fn completion_rejects_repeated_state_callback_parameter() {
let service = service();
let started = started(&service).await;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}&state=another-state-value"
)),
})
.await
.expect_err("repeated state");
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
struct UnavailableVerifier;
#[async_trait]
impl CodexAccountIdentityVerifier for UnavailableVerifier {
async fn verify(
&self,
_: &CodexOAuthSecret,
_: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
Err(CodexIdentityVerificationError::Unavailable)
}
async fn verify_authorization(
&self,
_: &CodexOAuthSecret,
_: &SecretString,
_: &SecretString,
_: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
Err(CodexIdentityVerificationError::Unavailable)
}
}
#[tokio::test]
async fn completion_classifies_unavailable_identity_verification_as_upstream_unavailable() {
let store = Arc::new(MemoryAccountStore::default());
let service = service_with_store(store, Arc::new(UnavailableVerifier));
let started = started(&service).await;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}"
)),
})
.await
.expect_err("identity verification unavailable");
assert_eq!(error, CodexOAuthAdminError::UpstreamUnavailable);
}
#[tokio::test]
async fn callback_state_mismatch_releases_flow_for_retry() {
let service = service();
let started = started(&service).await;
let flow_id = started.flow_id;
@@ -284,14 +464,48 @@ async fn callback_state_mismatch_fails_closed_and_consumes_flow() {
.complete_authorization(command())
.await
.expect_err("bad state"),
CodexOAuthAdminError::UpstreamRejected
CodexOAuthAdminError::CallbackRejected
);
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("one shot"),
CodexOAuthAdminError::NotFound
.expect_err("retryable callback rejection"),
CodexOAuthAdminError::CallbackRejected
);
}
#[tokio::test]
async fn token_exchange_rejection_releases_flow_for_retry() {
let service = service();
let started = started(&service).await;
let flow_id = started.flow_id;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let command = || CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: flow_id.clone(),
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=rejected-authorization-code&state={state}"
)),
};
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("token exchange rejection"),
CodexOAuthAdminError::TokenRejected
);
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("retryable token exchange rejection"),
CodexOAuthAdminError::TokenRejected
);
}
@@ -309,7 +523,7 @@ fn oauth_command_debug_redacts_owner_flow_and_callback() {
}
#[tokio::test]
async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotation() {
async fn reauthorization_uses_revision_advanced_after_start_for_prepared_rotation() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
@@ -337,6 +551,8 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
})
.await
.expect("start reauthorization");
let advanced_revision = store.advance_credential_revision("acct_oauth_reauth").await;
assert_eq!(advanced_revision.get(), 2);
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
@@ -360,7 +576,7 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
prepared.credential.account_id().as_str(),
"acct_oauth_reauth"
);
assert_eq!(prepared.credential.expected_revision().get(), 1);
assert_eq!(prepared.credential.expected_revision().get(), 2);
let runtime =
provider_openai::credential::CodexCredentialCodec::decode(prepared.credential.credential())
.expect("prepared credential");
@@ -380,7 +596,7 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
.expect("unchanged account")
.revision()
.get(),
1
2
);
}
@@ -361,6 +361,73 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
);
}
#[tokio::test]
async fn passive_rate_limit_headers_replace_stale_secondary_window() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_quota_snapshot").await;
let account = store
.account("acct_passive_quota_snapshot")
.expect("account");
let baseline = json!({
"rate_limit": {
"primary_window": {
"used_percent": 91,
"reset_at": 1_900_000_000,
"limit_window_seconds": 2_592_000
},
"secondary_window": {
"used_percent": 43,
"reset_at": 1_900_000_000,
"limit_window_seconds": 2_592_000
}
}
});
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
baseline.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist baseline quota");
let service = quota_service(&store);
let headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "92".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"43200".to_owned(),
),
(
"x-codex-primary-reset-at".to_owned(),
"1900000000".to_owned(),
),
];
assert!(
service
.synchronize_passive_headers(&account, &headers)
.await
.expect("passive quota")
);
let snapshot = service
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
let monthly = snapshot
.windows()
.iter()
.filter(|window| window.kind() == CodexQuotaWindowKind::Monthly)
.collect::<Vec<_>>();
assert_eq!(monthly.len(), 1);
assert_eq!(monthly[0].used_percent(), Some(92.0));
}
#[tokio::test]
async fn synchronize_without_accounts_is_a_noop_before_network_io() {
let store = Arc::new(MemoryAccountStore::default());
@@ -379,18 +446,248 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() {
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ready_accounts() {
async fn periodic_quota_synchronization_skips_fresh_ready_accounts() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_ready").await;
let account = store
.account("acct_periodic_ready")
.expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist fresh quota");
let summary = blocked_network_quota_service(&store)
.synchronize()
.await
.expect("ready account must not query upstream");
.expect("fresh ready account must not query upstream");
assert_eq!(summary, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_quota() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_stuck_cooldown").await;
let account = store
.account("acct_periodic_stuck_cooldown")
.expect("created account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account cooldown");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist exhausted quota");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(429))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_limit() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_exhausted").await;
let account = store
.account("acct_periodic_exhausted")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(429))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_exhausted")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhausted_quota() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_temporary_cooldown").await;
let account = store
.account("acct_periodic_temporary_cooldown")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account cooldown");
let summary = blocked_network_quota_service(&store)
.synchronize()
.await
.expect("ordinary cooldown must not query upstream");
assert_eq!(summary, CodexQuotaSyncSummary::default());
assert_eq!(
store
.account("acct_periodic_temporary_cooldown")
.expect("cooldown account")
.availability(),
AccountAvailability::Cooldown
);
}
#[tokio::test]
async fn periodic_quota_synchronization_limits_stale_ready_account_scans() {
let store = Arc::new(MemoryAccountStore::default());
for index in 0..9 {
let account_id = format!("acct_periodic_stale_ready_{index}");
create_account(&store, &account_id).await;
let account = store.account(&account_id).expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::UNIX_EPOCH),
})
.await
.expect("persist stale quota");
}
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let summary = service.synchronize().await.expect("quota synchronization");
assert_eq!(summary.exhausted, 8);
assert_eq!(
store
.account("acct_periodic_stale_ready_0")
.expect("scanned account")
.availability(),
AccountAvailability::QuotaExhausted
);
assert_eq!(
store
.account("acct_periodic_stale_ready_8")
.expect("deferred account")
.availability(),
AccountAvailability::Ready
);
}
#[tokio::test]
async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
let store = Arc::new(MemoryAccountStore::default());
@@ -444,6 +741,90 @@ async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
assert_eq!(summary, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_its_reset_is_due() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_reset_due").await;
let account = store
.account("acct_periodic_reset_due")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist future reset");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let first = service.synchronize().await.expect("initial refresh");
assert_eq!(first.transient, 1);
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": 1}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist due reset");
let reset_recheck = service.synchronize().await.expect("reset recheck");
let repeated = service
.synchronize()
.await
.expect("throttled reset recheck");
assert_eq!(reset_recheck.transient, 1);
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
let store = Arc::new(MemoryAccountStore::default());
@@ -596,6 +596,75 @@ async fn transient_refresh_defers_without_invalidating_account() {
assert_eq!(refresher.seen.lock().expect("seen tokens lock").len(), 1);
}
#[tokio::test]
async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
let store = Arc::new(MemoryAccountStore::default());
seed_due_account(&store).await;
let original = store.account("acct_refresh").expect("seeded account");
let mut expired_profile = profile("chatgpt-acct_refresh");
expired_profile.access_token_expires_at =
Some(chrono::Utc::now() - chrono::Duration::minutes(5));
store
.repository()
.rotate_oauth_secret(RotateCodexCredential {
account_id: "acct_refresh".to_owned(),
expected_credential_revision: original.revision().get(),
secret: secret("old-access"),
verified_account: expired_profile,
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::hours(12)),
})
.await
.expect("seed expired access token with deferred refresh");
let deferred = store.account("acct_refresh").expect("deferred account");
store
.repository()
.apply_state(
&deferred,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
.expect("mark account quota exhausted");
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
let service = CodexCredentialRefreshService::new(
store.repository(),
refresher.clone(),
Arc::new(VerifiedIdentity),
Arc::new(RefreshLeases {
available: true,
requests: Mutex::new(Vec::new()),
}),
Arc::new(CountingCredentialState::default()),
runtime_policy(),
);
let outcomes = service.refresh_due().await.expect("refresh due");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed {
account_id,
credential_revision: 3,
}] if account_id == "acct_refresh"
));
assert_eq!(
refresher.seen.lock().expect("seen tokens lock").as_slice(),
["rt-old-access"]
);
let refreshed = store.account("acct_refresh").expect("refreshed account");
assert_eq!(
refreshed.availability(),
AccountAvailability::QuotaExhausted
);
assert!(
refreshed
.access_token_expires_at()
.is_some_and(|expires_at| expires_at > SystemTime::now())
);
}
#[tokio::test]
async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() {
let (store, refresher, service) = setup(Ok(success_tokens()), false).await;
@@ -4,13 +4,17 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{AccountFeedbackStats, ProviderAccountId};
use futures::StreamExt;
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccountId,
};
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId,
AccountAttemptContext, AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId,
RequestAttemptContext, UpstreamSendState,
};
use gateway_core::error::ProviderErrorKind;
use gateway_core::event::GatewayEvent;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
@@ -25,6 +29,12 @@ use provider_openai::transport::CodexWebSocketPool;
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use provider_openai::{CodexProvider, OFFICIAL_CODEX_BASE_URL};
use serde_json::{Map, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
sync::oneshot,
time::timeout,
};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
@@ -122,6 +132,10 @@ fn provider_with_affinity_and_base_url(
}
async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
create_account_with_enabled(store, id, true).await;
}
async fn create_account_with_enabled(store: &Arc<MemoryAccountStore>, id: &str, enabled: bool) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: id.to_owned(),
@@ -129,7 +143,7 @@ async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
secret: secret(&format!("at-{id}")),
verified_account: profile(&format!("chatgpt-{id}")),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
enabled,
})
.await;
}
@@ -147,6 +161,19 @@ fn generate_operation() -> Operation {
))
}
fn http_generate_operation() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest {
let provider = ProviderKind::new(provider_name).expect("provider");
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
@@ -185,6 +212,97 @@ fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext
)
}
fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::diagnostic(
BTreeSet::new(),
ProviderAccountId::new(account_id).expect("account id"),
None,
),
None,
CancellationToken::new(),
)
}
async fn paused_chunked_sse_server(
first_chunk: String,
second_chunk: String,
) -> (
String,
oneshot::Sender<()>,
oneshot::Receiver<()>,
tokio::task::JoinHandle<()>,
) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind chunked SSE listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let (release_sender, release_receiver) = oneshot::channel();
let (first_chunk_sender, first_chunk_sent) = oneshot::channel();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept chunked SSE request");
read_http_request(&mut stream).await;
stream
.write_all(
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n",
)
.await
.expect("write chunked SSE headers");
write_http_chunk(&mut stream, &first_chunk).await;
let _ = first_chunk_sender.send(());
let _ = release_receiver.await;
if !second_chunk.is_empty() {
write_http_chunk(&mut stream, &second_chunk).await;
}
stream
.write_all(b"0\r\n\r\n")
.await
.expect("terminate chunked SSE response");
});
(base_url, release_sender, first_chunk_sent, server)
}
async fn read_http_request(stream: &mut TcpStream) {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let read = stream.read(&mut buffer).await.expect("read HTTP request");
if read == 0 {
return;
}
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
return;
}
}
}
async fn write_http_chunk(stream: &mut TcpStream, body: &str) {
stream
.write_all(format!("{:X}\r\n", body.len()).as_bytes())
.await
.expect("write HTTP chunk size");
stream
.write_all(body.as_bytes())
.await
.expect("write HTTP chunk body");
stream
.write_all(b"\r\n")
.await
.expect("terminate HTTP chunk");
stream.flush().await.expect("flush HTTP chunk");
}
#[tokio::test]
async fn openai_provider_rejects_a_foreign_provider_candidate_before_account_selection() {
let store = Arc::new(MemoryAccountStore::default());
@@ -271,6 +389,25 @@ async fn generate_without_an_eligible_openai_account_fails_before_network_io() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn disabled_account_is_excluded_from_normal_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
create_account_with_enabled(&store, "acct_disabled_scheduling", false).await;
let result = provider(&store)
.execute(
planned_request("openai", generate_operation()),
context("req_disabled_scheduling", CancellationToken::new()),
)
.await;
let Err(error) = result else {
panic!("disabled account must not be scheduled normally")
};
assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount);
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn opaque_provider_options_do_not_change_openai_account_selection() {
let store = Arc::new(MemoryAccountStore::default());
@@ -334,6 +471,478 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key()
assert!(keys[0].bytes().all(|byte| byte.is_ascii_hexdigit()));
}
#[tokio::test]
async fn completed_response_persists_session_affinity_before_stream_consumer_stops_polling() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_completed_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
("prompt_cache_key".to_owned(), json!("affinity-key")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
);
let mut stream =
provider_with_affinity_and_base_url(&store, Arc::clone(&affinity), server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_completed_affinity", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)))
{
break;
}
}
drop(stream);
assert_eq!(affinity.binding_count(), 1);
}
#[tokio::test]
async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_atomic_failure";
create_account(&store, account_id).await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_atomic_failure\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":17,\"response\":{\"id\":\"resp_atomic_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"atomic upstream failure\"}}}\n\n"
)
.to_owned(),
String::new(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_atomic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let mut failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("response.failed must produce a typed failure"),
}
};
assert_eq!(visible_before_error, 0);
assert_eq!(failure.kind(), ProviderErrorKind::RateLimited);
assert_eq!(failure.send_state(), UpstreamSendState::Sent);
assert_eq!(failure.upstream_status(), Some(429));
assert!(failure.replay_is_safe());
let events = failure.take_atomic_client_events();
assert_eq!(
events
.iter()
.filter_map(|event| event.wire_event()?.event_type())
.collect::<Vec<_>>(),
vec!["response.created", "response.failed"]
);
let account = store.account(account_id).expect("rate-limited account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first_event_latency").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_first_event\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
)
.to_owned(),
concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_first_event\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"later failure\"}}}\n\n"
)
.to_owned(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_first_event_latency", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let first_event = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("first chunk must not wait for the withheld failure")
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
break next;
}
};
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
release.send(()).expect("release second upstream chunk");
let mut later_wire_types = Vec::new();
let failure = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("later failure must arrive after release")
.expect("provider stream must return typed failure");
match next {
Ok(event) => {
if let Some(event_type) = event.wire_event().and_then(|wire| wire.event_type()) {
later_wire_types.push(event_type.to_owned());
}
}
Err(error) => break error,
}
};
assert_eq!(later_wire_types, vec!["response.failed"]);
assert!(!failure.replay_is_safe());
assert!(!failure.has_atomic_client_events());
server.await.expect("chunked SSE server");
}
#[tokio::test]
async fn continuation_should_hold_created_until_semantic_output() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_continuation_prefetch").await;
let (base_url, release, first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_continuation\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
)
.to_owned(),
concat!(
"event: response.content_part.added\n",
"data: {\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n",
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_continuation\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)
.to_owned(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_continuation_prefetch", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::Native),
)
.await
.expect("prepare provider stream");
let mut first_chunk_sent = Box::pin(first_chunk_sent);
loop {
tokio::select! {
sent = &mut first_chunk_sent => {
sent.expect("first upstream chunk");
break;
}
next = stream.next() => {
let event = next
.expect("provider stream must stay open")
.expect("provider event");
assert!(!event.has_client_event(), "created was exposed before the first chunk barrier");
}
}
}
let blocked = timeout(Duration::from_millis(100), async {
loop {
let next = stream
.next()
.await
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
return next;
}
}
})
.await;
assert!(blocked.is_err());
release.send(()).expect("release semantic output chunk");
let first_event = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("semantic output must release the prefetched batch")
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
break next;
}
};
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
while let Some(event) = stream.next().await {
event.expect("continuation response must complete");
}
server.await.expect("chunked SSE server");
}
#[tokio::test]
async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_bare_atomic_failure").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_bare_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"bare failure\"}}}\n\n"
)
.to_owned(),
String::new(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_bare_atomic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let mut failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("bare response.failed must produce a typed failure"),
}
};
assert_eq!(visible_before_error, 0);
assert!(failure.replay_is_safe());
let events = failure.take_atomic_client_events();
assert_eq!(events.len(), 1);
assert_eq!(
events[0].wire_event().and_then(|wire| wire.event_type()),
Some("response.failed")
);
assert!(events[0]
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "resp_bare_failure")));
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn continuation_prefetch_over_64_kib_without_a_boundary_should_fail_as_protocol_error() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_prefetch_limit").await;
let padding = "x".repeat(64 * 1024);
let body = format!(
"event: response.created\ndata: {}\n\n",
json!({
"type": "response.created",
"response": {
"id": "resp_prefetch_limit",
"model": "gpt-5.4",
"status": "in_progress",
"padding": padding,
}
})
);
assert!(body.len() > 64 * 1024);
let (base_url, release, _first_chunk_sent, server) =
paused_chunked_sse_server(body, String::new()).await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_prefetch_limit", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::Native),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("oversized prefetch must fail"),
}
};
assert_eq!(visible_before_error, 0);
assert_eq!(failure.kind(), ProviderErrorKind::Protocol);
assert!(!failure.has_atomic_client_events());
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn response_failed_after_semantic_output_is_exposed_and_not_replay_safe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_semantic_failure").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_semantic_failure\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.content_part.added\n",
"data: {\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n",
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_semantic_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"failure after output\"}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_semantic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut wire_types = Vec::new();
let failure = loop {
match stream.next().await {
Some(Ok(event)) => {
if let Some(event_type) = event.wire_event().and_then(|wire| wire.event_type()) {
wire_types.push(event_type.to_owned());
}
}
Some(Err(error)) => break error,
None => panic!("response.failed must produce a typed failure"),
}
};
assert_eq!(
wire_types,
vec![
"response.created",
"response.content_part.added",
"response.output_text.delta",
"response.failed"
]
);
assert!(!failure.replay_is_safe());
assert!(!failure.has_atomic_client_events());
}
#[tokio::test]
async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_disabled_diagnostic";
create_account_with_enabled(&store, account_id, false).await;
let account = store.account(account_id).expect("disabled test account");
store
.repository()
.apply_state(
&account,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
.expect("seed quota-exhausted state");
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.insert_header("x-codex-active-limit", "codex")
.insert_header("x-codex-primary-used-percent", "100")
.insert_header("x-codex-primary-window-minutes", "300")
.insert_header("x-codex-primary-reset-at", "1900000000")
.insert_header("x-codex-limit-reached", "true")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_disabled_diagnostic\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_disabled_diagnostic\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream =
provider_with_affinity_and_base_url(&store, Arc::clone(&affinity), server.uri())
.execute(
planned_request("openai", http_generate_operation()),
diagnostic_context("req_disabled_diagnostic", account_id),
)
.await
.expect("disabled diagnostic should prepare a fixed-account stream");
let mut completed = false;
while let Some(event) = stream.next().await {
let event = event.expect("disabled diagnostic upstream response");
completed |= event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)));
}
assert!(completed);
let account = store
.account(account_id)
.expect("disabled test account after test");
assert!(!account.enabled());
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(!store.has_quota(account_id));
assert_eq!(affinity.binding_count(), 0);
}
#[test]
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
let store = Arc::new(MemoryAccountStore::default());
@@ -67,9 +67,53 @@ impl MemoryAccountStore {
.map(|stored| stored.account.clone())
}
pub(crate) async fn advance_credential_revision(&self, id: &str) -> CredentialRevision {
let account_id = ProviderAccountId::new(id).expect("test account id");
let current = self
.load_current_credential(&account_id)
.await
.expect("current test credential");
let account = current.account;
let expected_revision = account.revision();
let update = CredentialCasUpdate::new(
account_id.clone(),
expected_revision,
ProviderAccountUpdate {
account_id,
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
plan_type: account.plan_type().map(str::to_owned),
},
current.credential,
account.has_refresh_token(),
account.access_token_expires_at(),
account.next_refresh_at(),
)
.expect("valid test credential update");
match self
.compare_and_swap_credential(update)
.await
.expect("advance test credential revision")
{
CredentialCasOutcome::Updated(revision) => revision,
CredentialCasOutcome::Conflict => panic!("test credential revision conflicted"),
}
}
pub(crate) fn quota_reads(&self) -> usize {
self.quota_reads.load(Ordering::SeqCst)
}
pub(crate) fn has_quota(&self, id: &str) -> bool {
let Some(id) = ProviderAccountId::new(id).ok() else {
return false;
};
self.accounts
.lock()
.expect("account store lock")
.get(&id)
.is_some_and(|stored| stored.quota.is_some())
}
}
#[async_trait]
@@ -130,14 +174,22 @@ impl ProviderAccountStore for MemoryAccountStore {
&self,
account: &ProviderAccountId,
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError> {
let loaded = self.load_current_credential(account).await?;
if loaded.account.revision() != expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
Ok(loaded)
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError> {
let accounts = self.accounts.lock().expect("account store lock");
let stored = accounts
.get(account)
.ok_or_else(|| store_error(StoreErrorKind::InvalidData))?;
if stored.account.revision() != expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
Ok(LoadedCredential {
account: stored.account.clone(),
credential: stored.credential.clone(),
@@ -440,6 +492,10 @@ impl MemorySessionAffinity {
pub(crate) fn lookup_keys(&self) -> Vec<String> {
self.lookups.lock().expect("session affinity lock").clone()
}
pub(crate) fn binding_count(&self) -> usize {
self.bindings.lock().expect("session affinity lock").len()
}
}
impl ProviderSessionAffinityPort for MemorySessionAffinity {
@@ -528,6 +528,40 @@ fn raw_sse_failure_should_keep_the_original_frame_before_reporting_the_typed_fai
assert!(!format!("{failure:?}").contains("raw failure marker"));
}
#[test]
fn bare_response_failed_should_project_started_identity_with_the_fallback_model() {
let raw = concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_bare_failed\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"bare failure\"}}}\n\n",
);
let failure = CodexCanonicalDecoder::new("fallback-model")
.push(raw.as_bytes())
.expect_err("bare response.failed remains a typed lifecycle failure");
assert!(matches!(
canonical_facts(failure.events()).as_slice(),
[GatewayEvent::Started(metadata)]
if metadata.response_id() == "resp_bare_failed"
&& metadata.model() == "fallback-model"
));
}
#[test]
fn response_failed_without_identity_should_keep_the_upstream_typed_failure() {
let raw = concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"missing identity\"}}}\n\n",
);
let failure = CodexCanonicalDecoder::new("fallback")
.push(raw.as_bytes())
.expect_err("missing identity must not replace the upstream lifecycle failure");
assert!(matches!(failure.error(), CodexCanonicalError::Upstream(_)));
assert!(canonical_facts(failure.events()).is_empty());
}
#[test]
fn decoder_should_preserve_same_chunk_output_before_typed_failure() {
let marker = "same-chunk-secret-marker";
@@ -143,3 +143,19 @@ fn structured_expiry_code_classifies_without_an_auth_status() {
CodexFailureCategory::CredentialExpired
);
}
#[test]
fn structured_quota_signal_takes_precedence_over_http_429() {
assert_eq!(
classify(
429,
r#"{"error":{"code":"insufficient_quota","message":"rate limit reached"}}"#
),
CodexFailureCategory::QuotaExhausted
);
}
#[test]
fn bare_http_429_remains_a_temporary_rate_limit() {
assert_eq!(classify(429, ""), CodexFailureCategory::RateLimited);
}
+3 -3
View File
@@ -10,15 +10,15 @@ autotests = false
[dependencies]
async-stream = "0.3.6"
async-trait = "0.1.89"
base64 = "0.22.1"
base64 = "0.23.0"
bytes = "1.9.0"
chrono = { version = "0.4.45", features = ["clock", "serde"] }
futures = "0.3.32"
gateway-admin = { path = "../../gateway-admin" }
gateway-core = { path = "../../gateway-core" }
gateway-protocol = { path = "../../gateway-protocol" }
getrandom = "0.3.4"
jsonwebtoken = { version = "=10.3.0", default-features = false, features = ["aws_lc_rs"] }
getrandom = "0.4.3"
jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["preserve_order"] }
+275 -178
View File
@@ -2,22 +2,20 @@
use std::collections::BTreeSet;
use std::sync::Arc;
use std::time::SystemTime;
use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::model::accounts::{
AccountAvailability as AdminAccountAvailability, AccountRecord,
};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireAttribute,
DashboardWireProfile, DashboardWireTarget, DecimalAmount, DesktopReleaseStatus,
ProviderBillingInput,
};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwner, AuthorizationOwnerBinding,
AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard,
AuthorizationCommitGuard, AuthorizationMutationTarget, AuthorizationOwner,
AuthorizationOwnerBinding, AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedAuthorizationCredential,
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
@@ -25,20 +23,19 @@ use gateway_admin::model::provider_credentials::{
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::model::{AdminError, MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use serde::Deserialize;
@@ -61,10 +58,78 @@ use crate::credential::{
use crate::transport::profile::{GrokCliReleaseSnapshot, GrokCliReleaseStatus};
use crate::transport::{GROK_CLI_BASE_URL, XAI_PROVIDER_NAME, grok_billing_breakdown};
const PENDING_SCHEMA_VERSION: u64 = 2;
const PENDING_SCHEMA_VERSION: u64 = 3;
const PENDING_TTL: TimeDelta = TimeDelta::minutes(30);
const PENDING_CLAIM_TTL: Duration = Duration::from_secs(90);
const MAX_PENDING_TEXT_BYTES: usize = 512;
struct ClaimedXaiAuthorization {
authorization: StoredXaiAuthorization,
owner_ref: String,
claim_ref: String,
}
struct XaiAuthorizationCommitGuard {
pending: Arc<dyn OAuthPendingFlowPort>,
provider_kind: ProviderKind,
flow: OAuthPendingBinding,
owner: OAuthPendingBinding,
claim: OAuthPendingBinding,
}
impl XaiAuthorizationCommitGuard {
fn new(
pending: Arc<dyn OAuthPendingFlowPort>,
provider_kind: ProviderKind,
flow_id: &str,
owner_ref: &str,
claim_ref: &str,
) -> Result<Self, ProviderAdminError> {
Ok(Self {
pending,
provider_kind,
flow: binding(flow_id)?,
owner: binding(owner_ref)?,
claim: binding(claim_ref)?,
})
}
}
#[async_trait]
impl AuthorizationCommitGuard for XaiAuthorizationCommitGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
match self
.pending
.consume_claim(&self.provider_kind, &self.flow, &self.owner, &self.claim)
.await
.map_err(map_pending_claim_settlement_error)?
{
OAuthPendingConsumeOutcome::Consumed => Ok(()),
OAuthPendingConsumeOutcome::NotFound
| OAuthPendingConsumeOutcome::OwnerMismatch
| OAuthPendingConsumeOutcome::ClaimMismatch => Err(AdminError::conflict(
"xAI OAuth pending claim is no longer current",
)),
}
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
match self
.pending
.release_claim(&self.provider_kind, &self.flow, &self.owner, &self.claim)
.await
.map_err(map_pending_claim_settlement_error)?
{
OAuthPendingReleaseOutcome::Released => Ok(()),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => Err(AdminError::unavailable(
"xAI OAuth pending claim could not be released",
)),
}
}
}
pub(crate) struct XaiAdminProvider {
provider_kind: ProviderKind,
wire_profile: XaiWireProfileState,
@@ -175,31 +240,81 @@ impl XaiAdminProvider {
}
}
async fn take_pending(
async fn claim_pending(
&self,
context: &MutationContext,
flow_id: &str,
) -> Result<StoredXaiAuthorization, ProviderAdminError> {
) -> Result<ClaimedXaiAuthorization, ProviderAdminError> {
let flow = binding(flow_id)?;
let owner_binding = AuthorizationOwnerBinding::from_context(context);
let owner_ref = owner_ref(owner_binding.owner());
let owner = binding(&owner_ref)?;
let claim_ref = random_flow_id()?;
let claim = binding(&claim_ref)?;
match self
.pending
.take_if_owner(&self.provider_kind, &flow, &owner)
.claim_if_owner(
&self.provider_kind,
&flow,
&owner,
&claim,
PENDING_CLAIM_TTL,
)
.await
.map_err(map_provider_store_error)?
{
OAuthPendingTakeOutcome::Taken(payload) => {
let stored = decode_pending(payload, &self.oauth_config)?;
if stored.flow_id != flow_id || stored.owner_ref != owner_ref {
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
Ok(stored)
OAuthPendingClaimOutcome::Claimed(payload) => {
let stored = match decode_pending(payload, &self.oauth_config) {
Ok(stored) if stored.flow_id == flow_id && stored.owner_ref == owner_ref => {
stored
}
Ok(_) => {
self.release_pending_claim(flow_id, &owner_ref, &claim_ref)
.await?;
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
Err(error) => {
self.release_pending_claim(flow_id, &owner_ref, &claim_ref)
.await?;
return Err(error);
}
};
Ok(ClaimedXaiAuthorization {
authorization: stored,
owner_ref,
claim_ref,
})
}
OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => {
OAuthPendingClaimOutcome::NotFound | OAuthPendingClaimOutcome::OwnerMismatch => {
Err(provider_error(ProviderAdminErrorKind::NotFound))
}
OAuthPendingClaimOutcome::InProgress => {
Err(provider_error(ProviderAdminErrorKind::Conflict))
}
}
}
async fn release_pending_claim(
&self,
flow_id: &str,
owner_ref: &str,
claim_ref: &str,
) -> Result<(), ProviderAdminError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.pending
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_provider_store_error)?
{
OAuthPendingReleaseOutcome::Released => Ok(()),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => {
Err(provider_error(ProviderAdminErrorKind::Unavailable))
}
}
}
}
@@ -405,16 +520,10 @@ impl ProviderAdmin for XaiAdminProvider {
if pending.provider_kind() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
if let AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} = pending.target()
{
let revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
if let AuthorizationMutationTarget::Reauthorize { account_id } = pending.target() {
let current = self
.accounts
.load_credential(account_id, revision)
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
@@ -436,76 +545,38 @@ impl ProviderAdmin for XaiAdminProvider {
&self,
command: CompleteAuthorization,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let stored = self
.take_pending(&command.context, &command.flow_id)
let ClaimedXaiAuthorization {
authorization: stored,
owner_ref,
claim_ref,
} = self
.claim_pending(&command.context, &command.flow_id)
.await?;
let authorization = PendingAuthorization::from_server_state(
&self.oauth_config,
&SecretValue::new(stored.server_state),
)
.map_err(map_oauth_error)?;
let callback = callback(&command.callback_url)?;
let grant = authorization
.accept_callback(callback)
.map_err(map_oauth_error)?;
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
let tokens = self
.oauth
.exchange_authorization_code(&discovery, grant)
.await
.map_err(map_oauth_error)?;
let policy = self
.runtime_policy
.load_refresh_policy()
.await
.map_err(map_provider_store_error)?;
let credential = match stored.mutation.target() {
AuthorizationMutationTarget::Create { name } => {
let prepared = GrokCredentialAdmin
.prepare_verified_account(
&VerifiedGrokAccount {
account_id: ProviderAccountId::new(format!(
"acct_{}",
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
name: name.clone(),
email: None,
upstream_account_id: None,
plan_type: None,
tokens,
enabled: true,
},
policy,
)
.map_err(map_repository_error)?;
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
let revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?;
let current = self
.accounts
.load_credential(account_id, revision)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
let prepared = verified_rotation(current, tokens, policy)?;
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
prepared,
self.provider_kind.clone(),
)?)
let authorization_guard = match XaiAuthorizationCommitGuard::new(
Arc::clone(&self.pending),
self.provider_kind.clone(),
&command.flow_id,
&owner_ref,
&claim_ref,
) {
Ok(guard) => guard,
Err(error) => {
self.release_pending_claim(&command.flow_id, &owner_ref, &claim_ref)
.await?;
return Err(error);
}
};
Ok(PreparedAuthorizationCommit {
pending: stored.mutation,
credential,
})
let completed = self
.complete_claimed_authorization(stored, &command.callback_url)
.await;
match completed {
Ok(completed) => Ok(completed.with_authorization_guard(Box::new(authorization_guard))),
Err(error) => {
self.release_pending_claim(&command.flow_id, &owner_ref, &claim_ref)
.await?;
Err(error)
}
}
}
async fn prepare_rotation(
@@ -513,16 +584,11 @@ impl ProviderAdmin for XaiAdminProvider {
command: PrepareCredentialRotation,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
validate_account_record(&command.account, &self.provider_kind)?;
if command.account.credential_revision != command.expected_credential_revision {
return Err(provider_error(ProviderAdminErrorKind::Conflict));
}
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let revision = CredentialRevision::new(command.expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_credential(&account_id, revision)
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
@@ -576,15 +642,17 @@ impl ProviderAdmin for XaiAdminProvider {
validate_account_record(&command.account, &self.provider_kind)?;
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let current = self.account(&account_id).await?;
if !account_matches_record(&current, &command.account) {
let current = self
.accounts
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
return Err(provider_error(ProviderAdminErrorKind::Conflict));
}
let revision = CredentialRevision::new(command.account.credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let prepared = self
.refresh
.prepare_manual_refresh(&account_id, revision)
.prepare_manual_refresh(current)
.await
.map_err(map_refresh_error)?;
prepared_rotation(prepared, command.account.provider_kind)
@@ -692,6 +760,77 @@ impl ProviderAdmin for XaiAdminProvider {
}
}
impl XaiAdminProvider {
async fn complete_claimed_authorization(
&self,
stored: StoredXaiAuthorization,
callback_url: &str,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let authorization = PendingAuthorization::from_server_state(
&self.oauth_config,
&SecretValue::new(stored.server_state),
)
.map_err(map_oauth_error)?;
let callback = callback(callback_url)?;
let grant = authorization
.accept_callback(callback)
.map_err(map_oauth_error)?;
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
let tokens = self
.oauth
.exchange_authorization_code(&discovery, grant)
.await
.map_err(map_oauth_error)?;
let policy = self
.runtime_policy
.load_refresh_policy()
.await
.map_err(map_provider_store_error)?;
let credential = match stored.mutation.target() {
AuthorizationMutationTarget::Create { name } => {
let prepared = GrokCredentialAdmin
.prepare_verified_account(
&VerifiedGrokAccount {
account_id: ProviderAccountId::new(format!(
"acct_{}",
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
name: name.clone(),
email: None,
upstream_account_id: None,
plan_type: None,
tokens,
enabled: true,
},
policy,
)
.map_err(map_repository_error)?;
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
}
AuthorizationMutationTarget::Reauthorize { account_id } => {
let current = self
.accounts
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
let prepared = verified_rotation(current, tokens, policy)?;
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
prepared,
self.provider_kind.clone(),
)?)
}
};
Ok(PreparedAuthorizationCommit::new(
stored.mutation,
credential,
))
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RotationDocument {
@@ -773,7 +912,7 @@ fn prepared_create(
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
next_refresh_at: account.next_refresh_at().map(Into::into),
enabled: account.enabled(),
availability: admin_availability(account.availability()),
availability: account.availability(),
availability_reason: None,
cooldown_until: account.cooldown_until().map(Into::into),
availability_observed_at: observed_at,
@@ -880,7 +1019,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
)
.with_runtime_state(
account.enabled,
core_availability(account.availability),
account.availability,
account.cooldown_until.map(Into::into),
)
.with_refresh_schedule(
@@ -892,43 +1031,9 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
&& account.email() == record.email.as_deref()
&& account.upstream_user_id() == record.upstream_user_id
&& account.upstream_account_id() == record.upstream_account_id.as_deref()
&& account.plan_type() == record.plan_type.as_deref()
&& account.enabled() == record.enabled
&& admin_availability(account.availability()) == record.availability
&& account.cooldown_until().map(DateTime::<Utc>::from) == record.cooldown_until
&& account.access_token_expires_at().map(DateTime::<Utc>::from)
== record.access_token_expires_at
&& account.next_refresh_at().map(DateTime::<Utc>::from) == record.next_refresh_at
&& account.has_refresh_token() == record.has_refresh_token
}
const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability {
match value {
AccountAvailability::Unknown => AdminAccountAvailability::Unknown,
AccountAvailability::Ready => AdminAccountAvailability::Ready,
AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
AccountAvailability::Expired => AdminAccountAvailability::Expired,
AccountAvailability::Banned => AdminAccountAvailability::Banned,
AccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability {
match value {
AdminAccountAvailability::Unknown => AccountAvailability::Unknown,
AdminAccountAvailability::Ready => AccountAvailability::Ready,
AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => AccountAvailability::Expired,
AdminAccountAvailability::Banned => AccountAvailability::Banned,
AdminAccountAvailability::Invalid => AccountAvailability::Invalid,
}
&& account.authentication_kind() == record.authentication_kind
}
fn project_quota(
@@ -1071,7 +1176,6 @@ struct PendingDocument {
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PendingMutationDocument {
expected_config_revision: u64,
provider_kind: String,
target: PendingTargetDocument,
owner: PendingOwnerDocument,
@@ -1081,13 +1185,8 @@ struct PendingMutationDocument {
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
name: String,
},
Reauthorize {
account_id: String,
expected_credential_revision: u64,
},
Create { name: String },
Reauthorize { account_id: String },
}
#[derive(Deserialize)]
@@ -1126,10 +1225,6 @@ fn encode_pending(pending: &StoredXaiAuthorization) -> Map<String, Value> {
fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map<String, Value> {
let mut document = Map::new();
document.insert(
"expected_config_revision".to_owned(),
Value::Number(Number::from(mutation.expected_config_revision().get())),
);
document.insert(
"provider_kind".to_owned(),
Value::String(mutation.provider_kind().as_str().to_owned()),
@@ -1156,19 +1251,12 @@ fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
);
document.insert(
"expected_credential_revision".to_owned(),
Value::Number(Number::from(expected_credential_revision.get())),
);
}
}
document
@@ -1234,15 +1322,12 @@ fn decode_mutation(
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let target = match document.target {
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
expected_credential_revision: Revision::new(expected_credential_revision)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
},
PendingTargetDocument::Reauthorize { account_id } => {
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
}
}
};
let actor = match document.owner {
PendingOwnerDocument::AdminSession { admin_user_id } => {
@@ -1256,8 +1341,6 @@ fn decode_mutation(
request_id: document.started_request_id,
};
Ok(PendingAuthorizationMutation::new(
Revision::new(document.expected_config_revision)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
provider_kind,
target,
AuthorizationOwnerBinding::from_context(&context),
@@ -1376,6 +1459,20 @@ fn map_provider_store_error(error: ProviderStoreError) -> ProviderAdminError {
})
}
fn map_pending_claim_settlement_error(error: ProviderStoreError) -> AdminError {
match error.kind() {
ProviderStoreErrorKind::InvalidData => {
AdminError::internal("xAI OAuth pending claim is invalid")
}
ProviderStoreErrorKind::Conflict => {
AdminError::conflict("xAI OAuth pending claim conflicts with current state")
}
ProviderStoreErrorKind::Unavailable => {
AdminError::unavailable("xAI OAuth pending claim store is unavailable")
}
}
}
fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError {
provider_error(match error.kind() {
StoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict,

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