mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0af8610a48 | ||
|
|
8226089ade | ||
|
|
1bd40ffa67 | ||
|
|
34a6ccce43 | ||
|
|
c275119f12 | ||
|
|
9322b7aeb8 | ||
|
|
c6444b2bf1 | ||
|
|
55075677eb | ||
|
|
733aa5f88a | ||
|
|
d704910857 | ||
|
|
ad6c21392f | ||
|
|
55e87c3b16 | ||
|
|
0c4da6c838 | ||
|
|
225c1166c6 | ||
|
|
97b91d3409 | ||
|
|
c953f6f44d | ||
|
|
d6b9fedbc7 | ||
|
|
720eb0c227 | ||
|
|
495c19b5e4 | ||
|
|
237ba96233 | ||
|
|
ae46d89de6 | ||
|
|
69f73e8c5d | ||
|
|
245f87ef35 | ||
|
|
81e51db0b6 | ||
|
|
4634496638 | ||
|
|
4e2f0a9346 | ||
|
|
304b3d4772 | ||
|
|
74629f3ef2 | ||
|
|
c9d48188af | ||
|
|
7e7b88a59e | ||
|
|
5c66dc198d | ||
|
|
47efce7a90 | ||
|
|
ea4d3f0ba7 | ||
|
|
2d03e27dcf | ||
|
|
cf7ff80f71 | ||
|
|
93ab75e42b | ||
|
|
449dc7da79 | ||
|
|
3b129a483a | ||
|
|
3e0d60bb1f | ||
|
|
1d3645fc26 | ||
|
|
18bff86214 | ||
|
|
114084107c | ||
|
|
d60f7238da | ||
|
|
ea684dceeb | ||
|
|
dfd496b41a | ||
|
|
6125b10276 | ||
|
|
071271f003 | ||
|
|
13a9df2e07 | ||
|
|
37a5a6702b | ||
|
|
afb68e97d6 | ||
|
|
93b9741af2 | ||
|
|
013d592454 | ||
|
|
0f56e9b80c | ||
|
|
6ec5ecb9ac | ||
|
|
460d439351 | ||
|
|
f39406dd80 | ||
|
|
9de8e75238 | ||
|
|
fee34b6bd1 | ||
|
|
cbe8699322 | ||
|
|
066e5551f8 | ||
|
|
1d73b736a9 | ||
|
|
8ae6635b9c | ||
|
|
574f0d580e | ||
|
|
0879515ee8 | ||
|
|
be40783454 |
@@ -64,6 +64,41 @@ jobs:
|
||||
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
with:
|
||||
ref: ${{ inputs.ref }}
|
||||
# PR 场景需要 base 提交做冻结清单的 append-only 比对。
|
||||
fetch-depth: ${{ github.event_name == 'pull_request' && 0 || 1 }}
|
||||
|
||||
# 已应用迁移的字节被 sqlx 启动校验钉死:改动已发布迁移会让所有
|
||||
# 存量部署重启即失败。冻结清单保证旧文件不被改、新文件必须入册。
|
||||
- name: Migration freeze check
|
||||
working-directory: backend/migrations
|
||||
shell: bash
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
sha256sum --check --strict .frozen-sha256
|
||||
for file in *.sql; do
|
||||
if ! awk '{print $2}' .frozen-sha256 | grep -Fxq "$file"; then
|
||||
echo "::error file=backend/migrations/$file::migration $file is not frozen; append its sha256 to backend/migrations/.frozen-sha256"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
# 清单自证不足以防「改迁移 + 同 PR 刷新清单行」:校验对新哈希通过。
|
||||
# 对 PR,断言 .frozen-sha256 相对 base 只增不改——base 里的每一行
|
||||
# 必须原样保留(改哈希=删行,删迁移=删行,都会被挡)。
|
||||
if [ -n "${BASE_SHA:-}" ]; then
|
||||
base_manifest="$(git show "${BASE_SHA}:backend/migrations/.frozen-sha256" 2>/dev/null || true)"
|
||||
if [ -n "$base_manifest" ]; then
|
||||
while IFS= read -r line; do
|
||||
[ -z "$line" ] && continue
|
||||
if ! grep -Fxq "$line" .frozen-sha256; then
|
||||
file="$(printf '%s' "$line" | awk '{print $2}')"
|
||||
echo "::error file=backend/migrations/.frozen-sha256::frozen migration $file was modified or removed; released migrations are append-only — add a new numbered migration instead"
|
||||
exit 1
|
||||
fi
|
||||
done <<< "$base_manifest"
|
||||
fi
|
||||
fi
|
||||
|
||||
- name: Setup Rust
|
||||
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
|
||||
|
||||
@@ -629,6 +629,9 @@ jobs:
|
||||
prerelease: ${{ needs.release-info.outputs.prerelease == 'true' }}
|
||||
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
|
||||
body: |
|
||||
> [!WARNING]
|
||||
> **Upgrading across major versions?** Database schema and configuration are only compatible within a major version, and the one-click self-update refuses cross-major targets (releases before 3.0 lack this gate — do not use their built-in update to reach a newer major). Deploy the new major fresh (new `.runtime/` data directory) and re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
|
||||
|
||||
## 📥 Installation
|
||||
|
||||
**Docker:**
|
||||
|
||||
@@ -26,7 +26,7 @@
|
||||
| 协议 | OpenAI Responses JSON、SSE、WebSocket 与模型目录 |
|
||||
| Provider | 固定 OpenAI、xAI 两个 Provider;各自管理 OAuth、账号、额度与模型目录 |
|
||||
| 路由 | Client Key 绑定 Provider、全局模型映射、会话亲和与同 Provider 安全 fallback |
|
||||
| 透明边界 | OpenAI 保持 Responses 语义透明;xAI 的协议适配只驻留在 `providers/xai` |
|
||||
| 透明边界 | OpenAI SSE 与 WebSocket 均原样转发上游字节;xAI 在 `providers/xai` 内做 Grok 与 Responses 协议转换 |
|
||||
| 延续 | OpenAI native/replay continuation;xAI 使用客户端提交的完整历史 |
|
||||
| 管理 | Client Key、账号、模型目录、设置、观测、系统更新与 OAuth |
|
||||
| 计量 | 模型请求 Token、费用、延迟、账号与 Provider 归因 |
|
||||
@@ -91,7 +91,7 @@ curl -i http://127.0.0.1:8080/healthz
|
||||
2. 在账号页选择对应 Provider:OpenAI 支持 OAuth、账号文件与 Agent Identity;xAI 支持 OAuth
|
||||
和兼容账号文件导入。
|
||||
3. 按需配置客户端模型到上游模型的全局映射;未命中映射的模型名原样透传。
|
||||
4. 创建 `sk_...` Client Key,绑定 `openai` 或 `xai`,并设置模型范围、速率与并发限制。
|
||||
4. 创建 `sk_...` Client Key,绑定 `openai` 或 `xai`,并设置速率与并发限制。
|
||||
|
||||
> [!IMPORTANT]
|
||||
> xAI 使用 OAuth session,不支持把 xAI API Key 作为上游 credential。
|
||||
@@ -158,9 +158,15 @@ docker compose -f deploy/compose.yaml up -d
|
||||
```
|
||||
|
||||
> [!IMPORTANT]
|
||||
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点和自更新状态。删除该目录会永久清除
|
||||
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点密钥和自更新状态。删除该目录会永久清除
|
||||
> 运行状态。
|
||||
|
||||
> [!WARNING]
|
||||
> 跨大版本不是原地升级:数据库 schema 与配置只在同一大版本内兼容,内置一键自更新
|
||||
> 会拒绝跨大版本目标(3.0 之前的旧版没有这道门禁,不要用旧版界面的一键更新跨大版本)。
|
||||
> 升级到新的大版本需要全新部署(全新 `.runtime/` 数据目录),并重新导入或重新授权
|
||||
> Provider 账号与客户端 Key。
|
||||
|
||||
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
|
||||
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
|
||||
|
||||
|
||||
Generated
+2
-3
@@ -1152,6 +1152,7 @@ dependencies = [
|
||||
"rand_core 0.6.4",
|
||||
"secrecy",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"subtle",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
@@ -1188,7 +1189,6 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
"serde",
|
||||
@@ -1224,7 +1224,6 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-appender",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
@@ -2150,7 +2149,6 @@ dependencies = [
|
||||
"url",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2160,6 +2158,7 @@ dependencies = [
|
||||
"async-stream",
|
||||
"async-trait",
|
||||
"base64",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"gateway-admin",
|
||||
|
||||
@@ -27,6 +27,7 @@ thiserror = "2.0.18"
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = { version = "1.0.150", features = ["preserve_order"] }
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] }
|
||||
|
||||
[lints]
|
||||
|
||||
@@ -29,7 +29,6 @@ pub enum AccountStatus {
|
||||
QuotaExhausted,
|
||||
Disabled,
|
||||
Banned,
|
||||
Attention,
|
||||
}
|
||||
|
||||
/// 账号列表排序字段。
|
||||
@@ -184,7 +183,7 @@ pub struct AccountSummary {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub attention: u64,
|
||||
pub unavailable: u64,
|
||||
}
|
||||
|
||||
/// 账号启停写入命令。
|
||||
@@ -222,6 +221,8 @@ pub enum AccountConnectionTestEvent {
|
||||
},
|
||||
Failed {
|
||||
message: String,
|
||||
provider_error_code: Option<String>,
|
||||
provider_error_type: Option<String>,
|
||||
account_status: AccountStatus,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -39,7 +39,6 @@ impl AdminRequestContext {
|
||||
pub struct LoginCommand {
|
||||
pub username: Option<String>,
|
||||
pub password: String,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for LoginCommand {
|
||||
@@ -48,7 +47,6 @@ impl std::fmt::Debug for LoginCommand {
|
||||
.debug_struct("LoginCommand")
|
||||
.field("username", &self.username)
|
||||
.field("password", &"[REDACTED]")
|
||||
.field("source", &self.source)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -65,8 +63,6 @@ pub struct LoginResult {
|
||||
pub enum LoginError {
|
||||
#[error("invalid administrator credentials")]
|
||||
InvalidCredentials,
|
||||
#[error("administrator login is throttled")]
|
||||
Throttled,
|
||||
#[error("administrator authentication is unavailable")]
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
@@ -718,7 +718,8 @@ pub struct OpsError {
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub failure_kind: String,
|
||||
pub status_code: Option<u16>,
|
||||
pub client_status_code: Option<u16>,
|
||||
pub upstream_status_code: Option<u16>,
|
||||
pub provider_error_code: Option<String>,
|
||||
pub client_response_id: Option<String>,
|
||||
pub upstream_request_id: Option<String>,
|
||||
|
||||
@@ -156,7 +156,7 @@ pub trait AccountStore: Send + Sync {
|
||||
) -> AdminStoreResult<()>;
|
||||
}
|
||||
|
||||
/// 管理员密码、会话、登录限流和安全审计。
|
||||
/// 管理员密码、会话和安全审计。
|
||||
#[async_trait]
|
||||
pub trait AuthStore: Send + Sync {
|
||||
async fn load_password_hash(&self, admin_user_id: &str) -> AdminStoreResult<Option<String>>;
|
||||
@@ -176,22 +176,6 @@ pub trait AuthStore: Send + Sync {
|
||||
|
||||
async fn delete_session(&self, session_id: &str) -> AdminStoreResult<Option<AdminSession>>;
|
||||
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> AdminStoreResult<bool>;
|
||||
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> AdminStoreResult<bool>;
|
||||
|
||||
async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()>;
|
||||
|
||||
async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()>;
|
||||
}
|
||||
|
||||
|
||||
@@ -10,7 +10,6 @@ use gateway_core::{
|
||||
credential::ProviderAccountId,
|
||||
probe::{AccountProbe, AccountProbeRequest},
|
||||
},
|
||||
error::GatewayErrorKind,
|
||||
routing::{ProviderKind, UpstreamModelId, snapshot::SnapshotControl},
|
||||
};
|
||||
|
||||
@@ -454,12 +453,9 @@ impl AccountsService for DefaultAccountsService {
|
||||
}))
|
||||
.collect(),
|
||||
Err(error) => vec![AccountConnectionTestEvent::Failed {
|
||||
message: match error.kind() {
|
||||
GatewayErrorKind::InvalidRequest
|
||||
| GatewayErrorKind::Unsupported
|
||||
| GatewayErrorKind::ModelNotFound => error.safe_message().to_owned(),
|
||||
_ => "Provider connection test failed".to_owned(),
|
||||
},
|
||||
message: error.client_message().to_owned(),
|
||||
provider_error_code: error.client_error_code().map(ToOwned::to_owned),
|
||||
provider_error_type: error.client_error_type().map(ToOwned::to_owned),
|
||||
account_status: status,
|
||||
}],
|
||||
}
|
||||
@@ -485,8 +481,7 @@ fn account_status(account: &AccountRecord, now: chrono::DateTime<Utc>) -> Accoun
|
||||
{
|
||||
AccountStatus::Expired
|
||||
}
|
||||
crate::model::accounts::AccountAvailability::Ready => AccountStatus::Active,
|
||||
_ => AccountStatus::Attention,
|
||||
_ => AccountStatus::Active,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,9 +22,6 @@ use crate::{
|
||||
|
||||
use super::map_store_error;
|
||||
|
||||
const LOGIN_FAILURE_LIMIT: u32 = 5;
|
||||
const LOGIN_FAILURE_WINDOW_SECONDS: u64 = 15 * 60;
|
||||
|
||||
/// API 鉴权与管理员登录消费的统一服务。
|
||||
#[async_trait]
|
||||
pub trait AuthService: Send + Sync {
|
||||
@@ -63,19 +60,6 @@ impl DefaultAuthService {
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_failed_login(&self, source: &str) -> Result<LoginError, LoginError> {
|
||||
let throttled = self
|
||||
.store
|
||||
.record_login_failure(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS)
|
||||
.await
|
||||
.map_err(|_| LoginError::Unavailable)?;
|
||||
Ok(if throttled {
|
||||
LoginError::Throttled
|
||||
} else {
|
||||
LoginError::InvalidCredentials
|
||||
})
|
||||
}
|
||||
|
||||
fn auth_audit(&self, action: &str, occurred_at: chrono::DateTime<Utc>) -> AdminAuditEvent {
|
||||
AdminAuditEvent {
|
||||
id: format!("audit_{}", Uuid::now_v7().simple()),
|
||||
@@ -137,22 +121,13 @@ impl AuthService for DefaultAuthService {
|
||||
}
|
||||
|
||||
async fn login(&self, command: LoginCommand) -> Result<LoginResult, LoginError> {
|
||||
let source = normalized_login_source(&command.source);
|
||||
if self
|
||||
.store
|
||||
.login_source_is_throttled(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS)
|
||||
.await
|
||||
.map_err(|_| LoginError::Unavailable)?
|
||||
{
|
||||
return Err(LoginError::Throttled);
|
||||
}
|
||||
if command
|
||||
.username
|
||||
.as_deref()
|
||||
.unwrap_or(&self.default_admin_user_id)
|
||||
!= self.default_admin_user_id
|
||||
{
|
||||
return Err(self.record_failed_login(source).await?);
|
||||
return Err(LoginError::InvalidCredentials);
|
||||
}
|
||||
let hash = self
|
||||
.store
|
||||
@@ -161,13 +136,9 @@ impl AuthService for DefaultAuthService {
|
||||
.map_err(|_| LoginError::Unavailable)?
|
||||
.ok_or(LoginError::InvalidCredentials)?;
|
||||
if !verify_admin_password(&command.password, &hash).map_err(|_| LoginError::Unavailable)? {
|
||||
return Err(self.record_failed_login(source).await?);
|
||||
return Err(LoginError::InvalidCredentials);
|
||||
}
|
||||
|
||||
self.store
|
||||
.clear_login_failures(source)
|
||||
.await
|
||||
.map_err(|_| LoginError::Unavailable)?;
|
||||
let session_id = random_session_token();
|
||||
let expires_at = Utc::now() + self.session_ttl;
|
||||
self.store
|
||||
@@ -246,8 +217,3 @@ fn valid_admin_api_key_shape(value: &str) -> bool {
|
||||
&& value.starts_with("admin-")
|
||||
&& value[6..].bytes().all(|byte| byte.is_ascii_hexdigit())
|
||||
}
|
||||
|
||||
fn normalized_login_source(source: &str) -> &str {
|
||||
let source = source.trim();
|
||||
if source.is_empty() { "unknown" } else { source }
|
||||
}
|
||||
|
||||
@@ -87,7 +87,6 @@ fn login_command_debug_should_redact_password() {
|
||||
let command = LoginCommand {
|
||||
username: None,
|
||||
password: "secret-password".to_owned(),
|
||||
source: "127.0.0.1".to_owned(),
|
||||
};
|
||||
assert!(!format!("{command:?}").contains("secret-password"));
|
||||
}
|
||||
|
||||
@@ -2,13 +2,14 @@ use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{TimeDelta, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use futures::{StreamExt as _, future::BoxFuture};
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
credential::{OpaqueProviderData, ProviderAccountId},
|
||||
probe::{AccountProbe, AccountProbeRequest, AccountProbeResult},
|
||||
},
|
||||
error::GatewayError,
|
||||
error::{ClientVisibleUpstreamError, GatewayError, GatewayErrorKind},
|
||||
operation::{GenerateRequest, Operation, ProtocolPayload},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
|
||||
@@ -17,9 +18,9 @@ use gateway_admin::{
|
||||
model::{
|
||||
MutationContext, Revision,
|
||||
accounts::{
|
||||
AccountAvailability, AccountListQuery, AccountPage, AccountRecord, AccountSummary,
|
||||
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, DeleteAccounts,
|
||||
SetAccountEnabled,
|
||||
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
|
||||
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
@@ -48,6 +49,7 @@ use gateway_admin::{
|
||||
},
|
||||
},
|
||||
};
|
||||
use serde_json::{Map, json};
|
||||
|
||||
pub(super) type EventLog = Arc<Mutex<Vec<&'static str>>>;
|
||||
|
||||
@@ -159,10 +161,22 @@ impl ProviderAdmin for FakeProviderAdmin {
|
||||
|
||||
fn connection_test_operation(
|
||||
&self,
|
||||
_: &gateway_core::routing::UpstreamModelId,
|
||||
_: &str,
|
||||
model: &gateway_core::routing::UpstreamModelId,
|
||||
input: &str,
|
||||
) -> Result<gateway_core::operation::Operation, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!(model.as_str())),
|
||||
("input".to_owned(), json!(input)),
|
||||
("stream".to_owned(), json!(true)),
|
||||
("store".to_owned(), json!(false)),
|
||||
]),
|
||||
)
|
||||
.map_err(|_| ProviderAdminError::new(ProviderAdminErrorKind::Invalid))?;
|
||||
Ok(Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
payload,
|
||||
)))
|
||||
}
|
||||
|
||||
fn dashboard_wire_profile(
|
||||
@@ -395,7 +409,7 @@ impl AccountStore for FakeAccountStore {
|
||||
total: 1,
|
||||
active: 1,
|
||||
quota_exhausted: 0,
|
||||
attention: 0,
|
||||
unavailable: 0,
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -686,6 +700,37 @@ fn provider_registry_should_reject_duplicate_kind() {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_test_should_preserve_safe_provider_error_details() {
|
||||
let provider = FakeProviderAdmin::new("xai", events());
|
||||
let store = FakeAccountStore::new("xai", events());
|
||||
let services =
|
||||
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
|
||||
|
||||
let events = services
|
||||
.accounts()
|
||||
.test_connection(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
|
||||
)
|
||||
.await
|
||||
.expect("connection test stream")
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
events.last(),
|
||||
Some(AccountConnectionTestEvent::Failed {
|
||||
message,
|
||||
provider_error_code: Some(code),
|
||||
provider_error_type: Some(error_type),
|
||||
..
|
||||
}) if message == "included usage exhausted"
|
||||
&& code == "usage_exhausted"
|
||||
&& error_type == "invalid_request_error"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
@@ -1005,12 +1050,20 @@ fn rotation_result(command: CredentialRotationCommit) -> CredentialMutationResul
|
||||
async fn accounts_service(
|
||||
provider: Arc<FakeProviderAdmin>,
|
||||
store: Arc<FakeAccountStore>,
|
||||
) -> AdminServices {
|
||||
accounts_service_with_probe(provider, store, Arc::new(SuccessfulAccountProbe)).await
|
||||
}
|
||||
|
||||
async fn accounts_service_with_probe(
|
||||
provider: Arc<FakeProviderAdmin>,
|
||||
store: Arc<FakeAccountStore>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
) -> AdminServices {
|
||||
super::AdminHarness::new()
|
||||
.accounts(store)
|
||||
.settings(Arc::new(StaticSettingsStore))
|
||||
.provider(provider)
|
||||
.probe(Arc::new(SuccessfulAccountProbe))
|
||||
.probe(probe)
|
||||
.build()
|
||||
.await
|
||||
}
|
||||
@@ -1030,6 +1083,29 @@ impl AccountProbe for SuccessfulAccountProbe {
|
||||
}
|
||||
}
|
||||
|
||||
struct FailingAccountProbe;
|
||||
|
||||
impl AccountProbe for FailingAccountProbe {
|
||||
fn probe(
|
||||
&self,
|
||||
_: AccountProbeRequest,
|
||||
) -> BoxFuture<'_, Result<AccountProbeResult, GatewayError>> {
|
||||
Box::pin(async {
|
||||
let detail = ClientVisibleUpstreamError::new(
|
||||
"included usage exhausted",
|
||||
Some("usage_exhausted".to_owned()),
|
||||
Some("invalid_request_error".to_owned()),
|
||||
)
|
||||
.expect("safe provider error");
|
||||
Err(GatewayError::new(
|
||||
GatewayErrorKind::RateLimited,
|
||||
"upstream capacity is temporarily unavailable",
|
||||
)
|
||||
.with_client_visible_upstream_error(detail))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn store_unavailable() -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
|
||||
@@ -15,7 +15,6 @@ struct MemoryAuthStore {
|
||||
password_hash: Mutex<Option<String>>,
|
||||
sessions: Mutex<BTreeMap<String, AdminSession>>,
|
||||
audits: Mutex<Vec<AdminAuditEvent>>,
|
||||
failures: Mutex<u32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -66,31 +65,6 @@ impl AuthStore for MemoryAuthStore {
|
||||
Ok(self.sessions.lock().expect("sessions").remove(session_id))
|
||||
}
|
||||
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
_: &str,
|
||||
failure_limit: u32,
|
||||
_: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
Ok(*self.failures.lock().expect("failures") >= failure_limit)
|
||||
}
|
||||
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
_: &str,
|
||||
failure_limit: u32,
|
||||
_: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
let mut failures = self.failures.lock().expect("failures");
|
||||
*failures += 1;
|
||||
Ok(*failures >= failure_limit)
|
||||
}
|
||||
|
||||
async fn clear_login_failures(&self, _: &str) -> AdminStoreResult<()> {
|
||||
*self.failures.lock().expect("failures") = 0;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()> {
|
||||
self.audits.lock().expect("audits").push(event);
|
||||
Ok(())
|
||||
@@ -107,7 +81,6 @@ async fn successful_login_should_create_expiring_session_and_audit() {
|
||||
.login(LoginCommand {
|
||||
username: Some("admin".to_owned()),
|
||||
password: "strong-test-password".to_owned(),
|
||||
source: "127.0.0.1".to_owned(),
|
||||
})
|
||||
.await
|
||||
.expect("login");
|
||||
@@ -142,7 +115,6 @@ async fn repeated_default_initialization_should_not_replace_password() {
|
||||
.login(LoginCommand {
|
||||
username: Some("admin".to_owned()),
|
||||
password: "first-strong-password".to_owned(),
|
||||
source: "127.0.0.1".to_owned(),
|
||||
})
|
||||
.await
|
||||
.is_ok()
|
||||
|
||||
@@ -210,18 +210,6 @@ impl AuthStore for BootstrapAuthStore {
|
||||
Err(unavailable("admin session"))
|
||||
}
|
||||
|
||||
async fn login_source_is_throttled(&self, _: &str, _: u32, _: u64) -> AdminStoreResult<bool> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
async fn record_login_failure(&self, _: &str, _: u32, _: u64) -> AdminStoreResult<bool> {
|
||||
Err(unavailable("login failure"))
|
||||
}
|
||||
|
||||
async fn clear_login_failures(&self, _: &str) -> AdminStoreResult<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn append_audit_event(&self, _: AdminAuditEvent) -> AdminStoreResult<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -174,7 +174,7 @@ pub struct AccountSummaryView {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub attention: u64,
|
||||
pub unavailable: u64,
|
||||
}
|
||||
|
||||
/// 一条安全账号视图。
|
||||
@@ -543,10 +543,14 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
|
||||
}),
|
||||
DomainConnectionTestEvent::Failed {
|
||||
message,
|
||||
provider_error_code,
|
||||
provider_error_type,
|
||||
account_status,
|
||||
} => serde_json::json!({
|
||||
"type": "error",
|
||||
"error": message,
|
||||
"providerErrorCode": provider_error_code,
|
||||
"providerErrorType": provider_error_type,
|
||||
"accountStatus": account_status_name(account_status)
|
||||
}),
|
||||
};
|
||||
@@ -1401,7 +1405,7 @@ fn account_page_data(
|
||||
total: result.summary.total,
|
||||
active: result.summary.active,
|
||||
quota_exhausted: result.summary.quota_exhausted,
|
||||
attention: result.summary.attention,
|
||||
unavailable: result.summary.unavailable,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1798,7 +1802,6 @@ fn account_status_name(status: DomainAccountStatus) -> &'static str {
|
||||
DomainAccountStatus::QuotaExhausted => "quota_exhausted",
|
||||
DomainAccountStatus::Disabled => "disabled",
|
||||
DomainAccountStatus::Banned => "banned",
|
||||
DomainAccountStatus::Attention => "attention",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use std::fmt;
|
||||
|
||||
use axum::{
|
||||
Extension, Json, Router,
|
||||
Json, Router,
|
||||
extract::{FromRequestParts, State},
|
||||
http::{HeaderMap, HeaderValue, StatusCode, header::SET_COOKIE, request::Parts},
|
||||
response::{IntoResponse, Response},
|
||||
@@ -14,6 +14,7 @@ use gateway_admin::{
|
||||
model::auth::{AdminPrincipal, AdminRequestContext, LoginCommand, LoginError},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tower_http::request_id::RequestId;
|
||||
|
||||
use super::{AdminEnvelope, AdminError, AdminResponse};
|
||||
|
||||
@@ -51,22 +52,32 @@ where
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
|
||||
let principal = require_admin_auth(state, &parts.headers).await?;
|
||||
let request_id = parts
|
||||
.headers
|
||||
.get(REQUEST_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
let request_id = admin_request_id(parts)
|
||||
.ok_or_else(|| AdminError::internal("Missing admin request context"))?;
|
||||
Ok(Self {
|
||||
context: AdminRequestContext {
|
||||
principal,
|
||||
request_id: request_id.to_owned(),
|
||||
request_id,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// request-id 层按配置的 header 名注入,同时写入与名字无关的扩展;
|
||||
/// 优先读扩展,使自定义 header 名不会让管理请求失去请求上下文。
|
||||
/// header 回退覆盖未装配该层的调用方。
|
||||
fn admin_request_id(parts: &Parts) -> Option<String> {
|
||||
parts
|
||||
.extensions
|
||||
.get::<RequestId>()
|
||||
.map(RequestId::header_value)
|
||||
.or_else(|| parts.headers.get(REQUEST_ID_HEADER))
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub async fn require_admin_session<S>(state: &S, headers: &HeaderMap) -> Result<String, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
@@ -185,22 +196,6 @@ impl Default for AdminLogoutData {
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求上下文解析出的客户端登录来源。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminLoginSource(String);
|
||||
|
||||
impl AdminLoginSource {
|
||||
#[must_use]
|
||||
pub fn new(value: impl Into<String>) -> Self {
|
||||
Self(value.into())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造固定 GET/POST 管理员认证路由。
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
@@ -214,24 +209,16 @@ where
|
||||
|
||||
async fn login<S>(
|
||||
State(state): State<S>,
|
||||
source: Option<Extension<AdminLoginSource>>,
|
||||
Json(payload): Json<AdminLoginRequest>,
|
||||
) -> Result<Response, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let source = source
|
||||
.as_ref()
|
||||
.map_or("unknown", |Extension(source)| source.as_str());
|
||||
let (username, password) = payload.into_parts();
|
||||
let session = state
|
||||
.admin_services()
|
||||
.auth()
|
||||
.login(LoginCommand {
|
||||
username,
|
||||
password,
|
||||
source: source.to_owned(),
|
||||
})
|
||||
.login(LoginCommand { username, password })
|
||||
.await
|
||||
.map_err(map_login_error)?;
|
||||
let mut response = AdminResponse::new(
|
||||
@@ -300,7 +287,6 @@ fn admin_session_cookie(headers: &HeaderMap) -> Option<String> {
|
||||
fn map_login_error(error: LoginError) -> AdminError {
|
||||
match error {
|
||||
LoginError::InvalidCredentials => AdminError::invalid_admin_credentials(),
|
||||
LoginError::Throttled => AdminError::too_many_login_attempts(),
|
||||
LoginError::Unavailable => AdminError::internal("Failed to create admin session"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -801,8 +801,8 @@ pub struct DashboardCardsView {
|
||||
pub struct DashboardCredentialsCardView {
|
||||
pub total: String,
|
||||
pub total_value: u64,
|
||||
pub enabled: String,
|
||||
pub enabled_value: u64,
|
||||
pub available: String,
|
||||
pub available_value: u64,
|
||||
pub unavailable: String,
|
||||
pub unavailable_value: u64,
|
||||
}
|
||||
@@ -1235,7 +1235,6 @@ pub struct OpsErrorView {
|
||||
pub account_id: Option<String>,
|
||||
pub route: String,
|
||||
pub model: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub client_status_code: Option<i64>,
|
||||
pub upstream_status_code: Option<i64>,
|
||||
pub transport: Option<String>,
|
||||
@@ -1954,7 +1953,6 @@ fn usage_page_view(
|
||||
}
|
||||
|
||||
fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
|
||||
let status = error.status_code.map(i64::from);
|
||||
OpsErrorView {
|
||||
id: error.event_id,
|
||||
request_id: error.request_id,
|
||||
@@ -1964,9 +1962,8 @@ fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
|
||||
account_id: error.provider_account_ref,
|
||||
route: error.operation,
|
||||
model: error.upstream_model_id,
|
||||
status_code: status,
|
||||
client_status_code: None,
|
||||
upstream_status_code: None,
|
||||
client_status_code: error.client_status_code.map(i64::from),
|
||||
upstream_status_code: error.upstream_status_code.map(i64::from),
|
||||
transport: error.upstream_transport,
|
||||
attempt_index: error.attempt_index,
|
||||
failure_class: error.failure_kind,
|
||||
@@ -2290,8 +2287,8 @@ fn dashboard_view(result: domain::DashboardResult, kind: TrendKind) -> Dashboard
|
||||
credentials: DashboardCredentialsCardView {
|
||||
total: format_compact_number(provider_accounts.total),
|
||||
total_value: provider_accounts.total,
|
||||
enabled: format_compact_number(provider_accounts.enabled),
|
||||
enabled_value: provider_accounts.enabled,
|
||||
available: format_compact_number(provider_accounts.active),
|
||||
available_value: provider_accounts.active,
|
||||
unavailable: format_compact_number(provider_accounts.unavailable),
|
||||
unavailable_value: provider_accounts.unavailable,
|
||||
},
|
||||
|
||||
@@ -201,11 +201,13 @@ struct RestartAcceptedView {
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SystemUpdateEventView {
|
||||
id: String,
|
||||
operation_id: Option<String>,
|
||||
level: &'static str,
|
||||
step: Option<String>,
|
||||
message: String,
|
||||
terminal: bool,
|
||||
progress_percent: Option<u8>,
|
||||
at: String,
|
||||
}
|
||||
|
||||
@@ -403,11 +405,13 @@ where
|
||||
impl From<SystemUpdateEvent> for SystemUpdateEventView {
|
||||
fn from(event: SystemUpdateEvent) -> Self {
|
||||
Self {
|
||||
id: event.id,
|
||||
operation_id: event.operation_id,
|
||||
level: update_event_level_name(event.level),
|
||||
step: event.step,
|
||||
message: event.message,
|
||||
terminal: event.terminal,
|
||||
progress_percent: event.progress_percent,
|
||||
at: event.occurred_at.to_rfc3339(),
|
||||
}
|
||||
}
|
||||
@@ -416,11 +420,13 @@ impl From<SystemUpdateEvent> for SystemUpdateEventView {
|
||||
impl SystemUpdateEventView {
|
||||
fn into_json(self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"id": self.id,
|
||||
"operationId": self.operation_id,
|
||||
"level": self.level,
|
||||
"step": self.step,
|
||||
"message": self.message,
|
||||
"terminal": self.terminal,
|
||||
"progressPercent": self.progress_percent,
|
||||
"at": self.at,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ use gateway_core::engine::execution::ExecutionService;
|
||||
use gateway_core::health::{HealthProbe, WorkerHealthSource};
|
||||
use gateway_core::lifecycle::ConnectionLifecycle;
|
||||
use serde::Deserialize;
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
|
||||
use tower_http::services::{ServeDir, ServeFile};
|
||||
use tower_http::timeout::TimeoutLayer;
|
||||
@@ -146,11 +146,17 @@ pub fn initialize(
|
||||
.map_err(|_| ApiError::Config(ApiConfigError::InvalidCorsOrigin))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
// credentials 模式禁止通配 allow_headers;显式列出鉴权与内容协商头。
|
||||
router = router.layer(
|
||||
CorsLayer::new()
|
||||
.allow_origin(origins)
|
||||
.allow_methods([Method::GET, Method::POST])
|
||||
.allow_headers(Any)
|
||||
.allow_headers([
|
||||
axum::http::header::AUTHORIZATION,
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
HeaderName::from_static("x-api-key"),
|
||||
request_id_header.clone(),
|
||||
])
|
||||
.allow_credentials(true),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2,11 +2,12 @@
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
extract::{Path, Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use gateway_core::routing::PublicModelId;
|
||||
use gateway_core::routing::{PublicModelId, PublicModelProfile};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::ApiState;
|
||||
@@ -17,20 +18,50 @@ use super::{
|
||||
};
|
||||
|
||||
const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000;
|
||||
const CODEX_BASE_INSTRUCTIONS: &str = "You are Codex, a coding agent. Follow the user's instructions and use the available tools to complete software engineering tasks. Inspect relevant files before editing, preserve unrelated changes, and verify the result.";
|
||||
|
||||
/// `GET /v1/models`。
|
||||
pub(crate) async fn models(State(state): State<ApiState>, headers: HeaderMap) -> Response {
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ModelsQuery {
|
||||
client_version: Option<String>,
|
||||
}
|
||||
|
||||
/// `GET /v1/models`。Codex 携带 `client_version` 时返回其专用目录合同。
|
||||
pub(crate) async fn models(
|
||||
State(state): State<ApiState>,
|
||||
Query(query): Query<ModelsQuery>,
|
||||
headers: HeaderMap,
|
||||
) -> Response {
|
||||
let service = state.openai();
|
||||
let client = match authenticate_client(service, &headers) {
|
||||
Ok(client) => client,
|
||||
Err(error) => return authentication_error_response(error),
|
||||
};
|
||||
|
||||
if query
|
||||
.client_version
|
||||
.as_deref()
|
||||
.is_some_and(|version| !version.trim().is_empty())
|
||||
{
|
||||
let models = service
|
||||
.public_model_profiles(&client)
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, profile)| codex_model_json(profile, index))
|
||||
.collect::<Vec<_>>();
|
||||
return (
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"models": models,
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
let data = service
|
||||
.public_models(&client)
|
||||
.into_iter()
|
||||
.map(|model| openai_model_json(&model))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
@@ -134,3 +165,76 @@ fn catalog_model_json(id: &str) -> Value {
|
||||
"source": "gateway",
|
||||
})
|
||||
}
|
||||
|
||||
fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
|
||||
let id = profile.model().as_str();
|
||||
let presentation = profile.presentation();
|
||||
let reasoning_levels = presentation
|
||||
.supported_reasoning_efforts()
|
||||
.iter()
|
||||
.map(|effort| {
|
||||
json!({
|
||||
"effort": effort,
|
||||
"description": reasoning_effort_description(effort),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let reasoning_supported = presentation
|
||||
.supported_reasoning_efforts()
|
||||
.iter()
|
||||
.any(|effort| effort != "none");
|
||||
let context_window = presentation.context_window_tokens();
|
||||
let input_modalities = if presentation.image_input() {
|
||||
vec!["text", "image"]
|
||||
} else {
|
||||
vec!["text"]
|
||||
};
|
||||
|
||||
json!({
|
||||
"slug": id,
|
||||
"display_name": presentation.display_name().unwrap_or(id),
|
||||
"description": presentation.description(),
|
||||
"default_reasoning_level": presentation.default_reasoning_effort(),
|
||||
"supported_reasoning_levels": reasoning_levels,
|
||||
"shell_type": "shell_command",
|
||||
"visibility": if presentation.hidden() { "hide" } else { "list" },
|
||||
"supported_in_api": true,
|
||||
"priority": index.saturating_add(1),
|
||||
"additional_speed_tiers": [],
|
||||
"service_tiers": [],
|
||||
"availability_nux": Value::Null,
|
||||
"upgrade": Value::Null,
|
||||
"base_instructions": CODEX_BASE_INSTRUCTIONS,
|
||||
"model_messages": Value::Null,
|
||||
"include_skills_usage_instructions": false,
|
||||
"supports_reasoning_summary_parameter": reasoning_supported,
|
||||
"default_reasoning_summary": "auto",
|
||||
"support_verbosity": presentation.verbosity(),
|
||||
"default_verbosity": Value::Null,
|
||||
"apply_patch_tool_type": presentation.agent_tools().then_some("freeform"),
|
||||
"web_search_tool_type": "text",
|
||||
"truncation_policy": { "mode": "tokens", "limit": 10_000 },
|
||||
"supports_parallel_tool_calls": presentation.parallel_tool_calls(),
|
||||
"supports_image_detail_original": presentation.image_detail_original(),
|
||||
"context_window": context_window,
|
||||
"max_context_window": context_window,
|
||||
"effective_context_window_percent": 95,
|
||||
"experimental_supported_tools": [],
|
||||
"input_modalities": input_modalities,
|
||||
"supports_search_tool": presentation.search_tool(),
|
||||
"use_responses_lite": false,
|
||||
})
|
||||
}
|
||||
|
||||
fn reasoning_effort_description(effort: &str) -> &'static str {
|
||||
match effort {
|
||||
"none" => "No reasoning",
|
||||
"minimal" => "Minimal reasoning",
|
||||
"low" => "Fast responses with lighter reasoning",
|
||||
"medium" => "Balances speed and reasoning depth for everyday tasks",
|
||||
"high" => "Greater reasoning depth for complex problems",
|
||||
"xhigh" => "Extra high reasoning depth for complex problems",
|
||||
"max" => "Maximum reasoning depth for the hardest problems",
|
||||
_ => "Provider-supported reasoning level",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -310,9 +310,18 @@ pub(super) fn decode_request_inner(
|
||||
) -> Result<DecodedResponsesRequest, RequestDecodeError> {
|
||||
let value =
|
||||
serde_json::from_slice::<Value>(body).map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
let Value::Object(mut object) = value else {
|
||||
let Value::Object(object) = value else {
|
||||
return Err(RequestDecodeError::ExpectedObject);
|
||||
};
|
||||
decode_request_object(object, review, request_headers)
|
||||
}
|
||||
|
||||
/// 解码已解析的顶层 object;WebSocket 帧解码复用此内核,避免序列化往返。
|
||||
pub(super) fn decode_request_object(
|
||||
mut object: Map<String, Value>,
|
||||
review: bool,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
) -> Result<DecodedResponsesRequest, RequestDecodeError> {
|
||||
// `use_websocket` 只影响本地 transport。无论值是否可识别,都不得进入上游 body。
|
||||
let use_websocket = object
|
||||
.remove("use_websocket")
|
||||
|
||||
@@ -549,9 +549,7 @@ fn decode_response_create_inner(
|
||||
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
|
||||
return Err(ResponseCreateFrameError::StreamingRequired);
|
||||
}
|
||||
let encoded = serde_json::to_vec(&Value::Object(body))
|
||||
.map_err(|_| ResponseCreateFrameError::InvalidJson)?;
|
||||
super::request::decode_request_inner(&encoded, false, request_headers)
|
||||
super::request::decode_request_object(body, false, request_headers)
|
||||
.map_err(ResponseCreateFrameError::Request)
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ use gateway_core::engine::execution::{
|
||||
};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle};
|
||||
use gateway_core::routing::PublicModelId;
|
||||
use gateway_core::routing::{PublicModelId, PublicModelProfile};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::auth::ClientApiKeyAuthError;
|
||||
@@ -51,6 +51,13 @@ impl OpenAiService {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn public_model_profiles(
|
||||
&self,
|
||||
client: &AuthenticatedClient,
|
||||
) -> Vec<PublicModelProfile> {
|
||||
self.execution.public_model_profiles(client)
|
||||
}
|
||||
|
||||
pub(crate) fn contains_public_model(
|
||||
&self,
|
||||
client: &AuthenticatedClient,
|
||||
|
||||
@@ -276,6 +276,8 @@ mod actions {
|
||||
},
|
||||
DomainConnectionTestEvent::Failed {
|
||||
message: "upstream unavailable".to_owned(),
|
||||
provider_error_code: Some("usage_exhausted".to_owned()),
|
||||
provider_error_type: Some("invalid_request_error".to_owned()),
|
||||
account_status: AccountStatus::Active,
|
||||
},
|
||||
]
|
||||
@@ -302,7 +304,13 @@ mod actions {
|
||||
}),
|
||||
json!({ "type": "content", "text": "OK" }),
|
||||
json!({ "type": "test_complete", "success": true, "accountStatus": "active" }),
|
||||
json!({ "type": "error", "error": "upstream unavailable", "accountStatus": "active" }),
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": "upstream unavailable",
|
||||
"providerErrorCode": "usage_exhausted",
|
||||
"providerErrorType": "invalid_request_error",
|
||||
"accountStatus": "active"
|
||||
}),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -51,7 +51,6 @@ async fn default_auth_service_should_initialize_login_validate_and_logout() {
|
||||
let service = fixture.services.auth();
|
||||
let session = service
|
||||
.login(LoginCommand {
|
||||
source: "127.0.0.1".to_owned(),
|
||||
username: Some("admin_1".to_owned()),
|
||||
password: "strong-admin-password".to_owned(),
|
||||
})
|
||||
@@ -84,42 +83,6 @@ async fn default_auth_service_should_initialize_login_validate_and_logout() {
|
||||
assert_eq!(fixture.auth.audit_count(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_auth_service_should_throttle_repeated_invalid_identity() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
for attempt in 1..=5 {
|
||||
let error = fixture
|
||||
.services
|
||||
.auth()
|
||||
.login(LoginCommand {
|
||||
source: "shared-source".to_owned(),
|
||||
username: Some("wrong-user".to_owned()),
|
||||
password: "wrong-password".to_owned(),
|
||||
})
|
||||
.await
|
||||
.expect_err("invalid login");
|
||||
if attempt < 5 {
|
||||
assert_eq!(error, LoginError::InvalidCredentials);
|
||||
} else {
|
||||
assert_eq!(error, LoginError::Throttled);
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
fixture
|
||||
.services
|
||||
.auth()
|
||||
.login(LoginCommand {
|
||||
source: "shared-source".to_owned(),
|
||||
username: None,
|
||||
password: "anything".to_owned(),
|
||||
})
|
||||
.await
|
||||
.expect_err("source remains throttled"),
|
||||
LoginError::Throttled
|
||||
);
|
||||
assert_eq!(fixture.auth.session_count(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_auth_service_should_verify_only_full_plaintext_admin_key() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
@@ -162,7 +125,6 @@ async fn audit_failure_should_revoke_new_session_before_returning_it() {
|
||||
.services
|
||||
.auth()
|
||||
.login(LoginCommand {
|
||||
source: "source".to_owned(),
|
||||
username: None,
|
||||
password: "strong-admin-password".to_owned(),
|
||||
})
|
||||
|
||||
@@ -84,6 +84,10 @@ pub(super) struct AdminTestFixture {
|
||||
|
||||
impl AdminTestFixture {
|
||||
pub async fn new() -> Self {
|
||||
Self::with_system(Arc::new(UnusedSystem)).await
|
||||
}
|
||||
|
||||
pub async fn with_system(system: Arc<dyn SystemOperations>) -> Self {
|
||||
let api_key = Arc::new(Mutex::new(None));
|
||||
let auth = Arc::new(MemoryAuthStore::new(api_key.clone()));
|
||||
let settings = Arc::new(MemorySettingsStore::new(api_key));
|
||||
@@ -113,7 +117,7 @@ impl AdminTestFixture {
|
||||
providers,
|
||||
Arc::new(NoopSnapshot),
|
||||
Arc::new(NoopProbe),
|
||||
Arc::new(UnusedSystem),
|
||||
system,
|
||||
)
|
||||
.await
|
||||
.expect("initialize test admin services");
|
||||
@@ -142,7 +146,6 @@ impl AdminSessionState for AdminTestState {
|
||||
pub(super) struct MemoryAuthStore {
|
||||
password_hash: Mutex<Option<String>>,
|
||||
sessions: Mutex<BTreeMap<String, AdminSession>>,
|
||||
failures: Mutex<BTreeMap<String, u32>>,
|
||||
audits: Mutex<Vec<AdminAuditEvent>>,
|
||||
api_key: Arc<Mutex<Option<AdminApiKey>>>,
|
||||
fail_audit: AtomicBool,
|
||||
@@ -153,7 +156,6 @@ impl MemoryAuthStore {
|
||||
Self {
|
||||
password_hash: Mutex::new(None),
|
||||
sessions: Mutex::new(BTreeMap::new()),
|
||||
failures: Mutex::new(BTreeMap::new()),
|
||||
audits: Mutex::new(Vec::new()),
|
||||
api_key,
|
||||
fail_audit: AtomicBool::new(false),
|
||||
@@ -235,37 +237,6 @@ impl AuthStore for MemoryAuthStore {
|
||||
Ok(self.sessions.lock().expect("sessions").remove(session_id))
|
||||
}
|
||||
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
_: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
Ok(self
|
||||
.failures
|
||||
.lock()
|
||||
.expect("failures")
|
||||
.get(source)
|
||||
.is_some_and(|count| *count >= failure_limit))
|
||||
}
|
||||
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
_: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
let mut failures = self.failures.lock().expect("failures");
|
||||
let count = failures.entry(source.to_owned()).or_default();
|
||||
*count = count.saturating_add(1);
|
||||
Ok(*count >= failure_limit)
|
||||
}
|
||||
|
||||
async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()> {
|
||||
self.failures.lock().expect("failures").remove(source);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()> {
|
||||
if self.fail_audit.load(Ordering::SeqCst) {
|
||||
return Err(unavailable("auth audit"));
|
||||
|
||||
@@ -213,3 +213,26 @@ async fn settings_should_accept_admin_api_key_header() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_should_accept_a_configured_request_id_header_name() {
|
||||
use axum::http::HeaderName;
|
||||
use tower_http::request_id::{MakeRequestUuid, SetRequestIdLayer};
|
||||
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
// 部署把 api.request_id_header 改名后,注入的 header 不再叫 x-request-id;
|
||||
// 管理请求仍须拿到请求上下文,而不是退化为 500。
|
||||
let custom = HeaderName::from_static("x-trace-id");
|
||||
let app = app(fixture.state()).layer(SetRequestIdLayer::new(custom, MakeRequestUuid));
|
||||
let unlabelled = Request::builder()
|
||||
.method(Method::GET)
|
||||
.uri("/api/admin/settings")
|
||||
.header(header::COOKIE, "cpr_admin_session=valid-session")
|
||||
.body(Body::empty())
|
||||
.expect("build settings request");
|
||||
|
||||
let response = app.oneshot(unlabelled).await.expect("settings response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,24 @@
|
||||
use gateway_api::admin::system::{UpdateDetailQuery, UpdateRequest};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, header},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use futures::stream;
|
||||
use gateway_admin::{
|
||||
model::system::{
|
||||
SystemOperationAccepted, SystemUpdateDetail, SystemUpdateEvent, SystemUpdateEventLevel,
|
||||
SystemUpdateStatus, SystemVersion,
|
||||
},
|
||||
ports::system::{SystemOperationError, SystemOperations, SystemUpdateEventStream},
|
||||
};
|
||||
use gateway_api::admin::system::{self, UpdateDetailQuery, UpdateRequest};
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState, unavailable_system};
|
||||
|
||||
#[test]
|
||||
fn update_detail_query_should_default_to_no_refresh_and_reject_unknown_fields() {
|
||||
@@ -16,3 +36,84 @@ fn update_request_should_preserve_target_version_for_domain_validation() {
|
||||
serde_json::from_str::<UpdateRequest>(r#"{"targetVersion":"v0.2.0","extra":1}"#).is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_event_stream_should_preserve_event_id_and_download_progress_percent() {
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem)).await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let response = app(fixture.state())
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/api/admin/system/update/events")
|
||||
.header(header::COOKIE, "cpr_admin_session=valid-session")
|
||||
.header("x-request-id", "req_system_update_events")
|
||||
.body(Body::empty())
|
||||
.expect("update event request"),
|
||||
)
|
||||
.await
|
||||
.expect("update event response");
|
||||
let body = to_bytes(response.into_body(), 1024 * 1024)
|
||||
.await
|
||||
.expect("update event body");
|
||||
let body = String::from_utf8(body.to_vec()).expect("UTF-8 event stream");
|
||||
|
||||
assert!(
|
||||
body.contains(r#""id":"update-progress-1""#),
|
||||
"unexpected update event stream: {body}"
|
||||
);
|
||||
assert!(
|
||||
body.contains(r#""progressPercent":40"#),
|
||||
"unexpected update event stream: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
fn app(state: AdminTestState) -> Router {
|
||||
system::router::<AdminTestState>().with_state(state)
|
||||
}
|
||||
|
||||
struct ProgressSystem;
|
||||
|
||||
#[async_trait]
|
||||
impl SystemOperations for ProgressSystem {
|
||||
async fn version(&self) -> Result<SystemVersion, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
|
||||
async fn update_detail(&self, _: bool) -> Result<SystemUpdateDetail, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
|
||||
fn update_events(&self) -> SystemUpdateEventStream {
|
||||
Box::pin(stream::once(async {
|
||||
SystemUpdateEvent {
|
||||
id: "update-progress-1".to_owned(),
|
||||
operation_id: Some("update-1".to_owned()),
|
||||
level: SystemUpdateEventLevel::Info,
|
||||
step: Some("download".to_owned()),
|
||||
message: "已下载 4.0 MiB / 10.0 MiB (40%)".to_owned(),
|
||||
terminal: true,
|
||||
progress_percent: Some(40),
|
||||
occurred_at: Utc::now(),
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
async fn perform_update(
|
||||
&self,
|
||||
_: Option<String>,
|
||||
) -> Result<SystemOperationAccepted, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
|
||||
async fn update_status(&self) -> Result<SystemUpdateStatus, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
|
||||
async fn rollback(&self) -> Result<SystemOperationAccepted, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
|
||||
async fn restart(&self) -> Result<SystemOperationAccepted, SystemOperationError> {
|
||||
Err(unavailable_system())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,11 +39,18 @@ use gateway_core::routing::{
|
||||
};
|
||||
|
||||
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
|
||||
api_router_with_origins(execution, Vec::new()).await
|
||||
}
|
||||
|
||||
pub(super) async fn api_router_with_origins(
|
||||
execution: Arc<dyn ExecutionService>,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
) -> axum::Router {
|
||||
let admin = crate::admin::AdminTestFixture::new().await;
|
||||
gateway_api::initialize(
|
||||
gateway_api::ApiConfig {
|
||||
asset_directory: std::env::temp_dir(),
|
||||
cors_allowed_origins: Vec::new(),
|
||||
cors_allowed_origins,
|
||||
request_timeout_seconds: None,
|
||||
request_id_header: "x-request-id".to_owned(),
|
||||
},
|
||||
|
||||
@@ -10,19 +10,28 @@ use gateway_core::engine::execution::{
|
||||
StartedExecution,
|
||||
};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::routing::PublicModelId;
|
||||
use gateway_core::routing::{ModelPresentation, PublicModelId, PublicModelProfile};
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::{api_router, authenticated_client};
|
||||
|
||||
struct ModelsExecution {
|
||||
pub(super) struct ModelsExecution {
|
||||
client: AuthenticatedClient,
|
||||
profiles: bool,
|
||||
}
|
||||
|
||||
impl ModelsExecution {
|
||||
fn new() -> Arc<Self> {
|
||||
pub(super) fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
profiles: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn with_profiles() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
profiles: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -46,6 +55,28 @@ impl ExecutionService for ModelsExecution {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, _: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
if !self.profiles {
|
||||
return Vec::new();
|
||||
}
|
||||
vec![PublicModelProfile::new(
|
||||
PublicModelId::new("grok-4.5").expect("model"),
|
||||
ModelPresentation::new(
|
||||
Some("Grok 4.5".to_owned()),
|
||||
Some("xAI Grok 4.5 frontier model.".to_owned()),
|
||||
)
|
||||
.with_reasoning(
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
)
|
||||
.with_context_window_tokens(Some(500_000))
|
||||
.with_image_input(true)
|
||||
.with_agent_tools(true, true),
|
||||
)]
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
matches!(model.as_str(), "model-a" | "model-b")
|
||||
}
|
||||
@@ -91,6 +122,92 @@ async fn models_should_encode_the_service_visible_catalog() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_should_encode_provider_profiles_for_current_codex_clients() {
|
||||
let response = api_router(ModelsExecution::with_profiles())
|
||||
.await
|
||||
.oneshot(authorized_request("/v1/models?client_version=0.145.0"))
|
||||
.await
|
||||
.expect("list Codex models response");
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("read Codex models body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("models JSON");
|
||||
let model = &value["models"][0];
|
||||
|
||||
assert_eq!(
|
||||
value
|
||||
.as_object()
|
||||
.map(|catalog| catalog.keys().map(String::as_str).collect::<Vec<_>>()),
|
||||
Some(vec!["models"])
|
||||
);
|
||||
assert_eq!(model["slug"], "grok-4.5");
|
||||
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["supported_reasoning_levels"],
|
||||
serde_json::json!([
|
||||
{
|
||||
"effort": "low",
|
||||
"description": "Fast responses with lighter reasoning"
|
||||
},
|
||||
{
|
||||
"effort": "medium",
|
||||
"description": "Balances speed and reasoning depth for everyday tasks"
|
||||
},
|
||||
{
|
||||
"effort": "high",
|
||||
"description": "Greater reasoning depth for complex problems"
|
||||
},
|
||||
{
|
||||
"effort": "xhigh",
|
||||
"description": "Extra high reasoning depth for complex problems"
|
||||
}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_should_keep_the_codex_contract_when_profiles_are_empty() {
|
||||
let response = api_router(ModelsExecution::new())
|
||||
.await
|
||||
.oneshot(authorized_request("/v1/models?client_version=0.145.0"))
|
||||
.await
|
||||
.expect("list empty Codex models response");
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("read empty Codex models body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("Codex models JSON");
|
||||
|
||||
assert_eq!(value, serde_json::json!({ "models": [] }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_with_provider_profiles_should_keep_the_openai_list_contract() {
|
||||
let response = api_router(ModelsExecution::with_profiles())
|
||||
.await
|
||||
.oneshot(authorized_request("/v1/models"))
|
||||
.await
|
||||
.expect("list compatible models response");
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("read compatible models body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("models JSON");
|
||||
|
||||
assert_eq!(
|
||||
value,
|
||||
serde_json::json!({
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"id":"model-a","object":"model","created":1700000000_i64,"owned_by":"gateway"},
|
||||
{"id":"model-b","object":"model","created":1700000000_i64,"owned_by":"gateway"}
|
||||
]
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_detail_should_keep_the_official_path_id_contract() {
|
||||
let response = api_router(ModelsExecution::new())
|
||||
|
||||
@@ -2,12 +2,15 @@ use axum::{
|
||||
Router,
|
||||
body::{Body, Bytes},
|
||||
extract::DefaultBodyLimit,
|
||||
http::{Request, StatusCode},
|
||||
http::{Method, Request, StatusCode},
|
||||
routing::post,
|
||||
};
|
||||
use gateway_api::openai::router::MAX_CLIENT_REQUEST_BODY_BYTES;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::api_router_with_origins;
|
||||
use super::models::ModelsExecution;
|
||||
|
||||
fn body_limit_app() -> Router {
|
||||
Router::new()
|
||||
.route(
|
||||
@@ -50,3 +53,50 @@ async fn responses_body_limit_should_reject_one_byte_over_sixteen_mibibytes() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn configured_cors_origin_assembles_router_and_answers_preflight() {
|
||||
let router = api_router_with_origins(
|
||||
ModelsExecution::new(),
|
||||
vec!["https://app.example.com".into()],
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = router
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(Method::OPTIONS)
|
||||
.uri("/v1/models")
|
||||
.header("origin", "https://app.example.com")
|
||||
.header("access-control-request-method", "GET")
|
||||
.header("access-control-request-headers", "authorization")
|
||||
.body(Body::empty())
|
||||
.expect("build preflight request"),
|
||||
)
|
||||
.await
|
||||
.expect("route preflight request");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get("access-control-allow-origin")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("https://app.example.com")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get("access-control-allow-credentials")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("true")
|
||||
);
|
||||
let allow_headers = response
|
||||
.headers()
|
||||
.get("access-control-allow-headers")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.expect("allow-headers present");
|
||||
assert!(allow_headers.contains("authorization"));
|
||||
assert!(allow_headers.contains("x-api-key"));
|
||||
assert!(allow_headers.contains("x-request-id"));
|
||||
}
|
||||
|
||||
@@ -14,7 +14,6 @@ path = "tests/main.rs"
|
||||
[dependencies]
|
||||
async-trait = "0.1.89"
|
||||
bytes = "1.11.1"
|
||||
chrono = "0.4.45"
|
||||
futures = "0.3.32"
|
||||
futures-timer = "3.0.3"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 唯一的账号重试、发送与下游 commit barrier owner。
|
||||
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
use std::collections::BTreeSet;
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
@@ -23,6 +23,7 @@ use crate::event::{
|
||||
};
|
||||
use crate::operation::{Operation, ProviderSessionState, RetrySafety};
|
||||
use crate::routing::RoutingPlan;
|
||||
use futures::future::Fuse;
|
||||
use futures::{FutureExt, StreamExt, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
@@ -67,18 +68,18 @@ where
|
||||
let continuation_attempt =
|
||||
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
|
||||
let image_generation_requested = operation.image_generation_requested();
|
||||
let commit_policy = if continuation_attempt == ContinuationAttempt::None {
|
||||
StreamCommitPolicy::FirstCanonicalEvent
|
||||
} else {
|
||||
StreamCommitPolicy::UntilOutputOrTerminal
|
||||
};
|
||||
|
||||
let mut session = ResponseExecutionSession {
|
||||
engine: Arc::clone(&self.engine),
|
||||
request_id,
|
||||
client_api_key_ref,
|
||||
timing_started_at,
|
||||
deadline,
|
||||
deadline_timer: Delay::new(
|
||||
deadline
|
||||
.duration_since(SystemTime::now())
|
||||
.unwrap_or(Duration::ZERO),
|
||||
)
|
||||
.fuse(),
|
||||
pending_request: Some(request),
|
||||
request_persisted: false,
|
||||
operation,
|
||||
@@ -86,13 +87,13 @@ where
|
||||
required_account,
|
||||
continuation,
|
||||
continuation_attempt,
|
||||
commit_policy,
|
||||
account_state_owner,
|
||||
cancellation,
|
||||
attempts: 0,
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
credential_recovery_attempted_accounts: BTreeSet::new(),
|
||||
current: None,
|
||||
send_state_watermark: UpstreamSendState::NotSent,
|
||||
downstream_committed_at: None,
|
||||
client_status_code: None,
|
||||
delivery_pending: false,
|
||||
@@ -106,7 +107,6 @@ where
|
||||
upstream_response_id: None,
|
||||
last_account_exhaustion: None,
|
||||
provider_attempt_outcomes: Vec::new(),
|
||||
before_commit: VecDeque::new(),
|
||||
};
|
||||
|
||||
if session.cancellation.is_cancelled() {
|
||||
@@ -141,12 +141,6 @@ struct FailureFinalization {
|
||||
provider_response_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum StreamCommitPolicy {
|
||||
FirstCanonicalEvent,
|
||||
UntilOutputOrTerminal,
|
||||
}
|
||||
|
||||
/// API 可逐事件消费的 Core 执行会话。
|
||||
///
|
||||
/// API 只能提交下游 delivery 边界;账号重试、断流终结与
|
||||
@@ -157,6 +151,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
client_api_key_ref: crate::policy::ClientApiKeyId,
|
||||
timing_started_at: Instant,
|
||||
deadline: SystemTime,
|
||||
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
|
||||
deadline_timer: Fuse<Delay>,
|
||||
pending_request: Option<NewModelRequest>,
|
||||
request_persisted: bool,
|
||||
operation: Operation,
|
||||
@@ -164,13 +160,14 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
required_account: Option<crate::engine::credential::ProviderAccountId>,
|
||||
continuation: Option<ContinuationBinding>,
|
||||
continuation_attempt: ContinuationAttempt,
|
||||
commit_policy: StreamCommitPolicy,
|
||||
account_state_owner: Option<ProviderAccountStateOwner>,
|
||||
cancellation: CancellationToken,
|
||||
attempts: u32,
|
||||
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
|
||||
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
|
||||
current: Option<CurrentAttempt>,
|
||||
/// 请求级发送状态水位;跨 attempt 单调不降,终态写回不得低于此档。
|
||||
send_state_watermark: UpstreamSendState,
|
||||
downstream_committed_at: Option<SystemTime>,
|
||||
client_status_code: Option<u16>,
|
||||
delivery_pending: bool,
|
||||
@@ -184,7 +181,6 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
upstream_response_id: Option<String>,
|
||||
last_account_exhaustion: Option<ProviderError>,
|
||||
provider_attempt_outcomes: Vec<ProviderAttemptOutcome>,
|
||||
before_commit: VecDeque<ProviderEvent>,
|
||||
}
|
||||
|
||||
impl<S: ?Sized> ResponseExecutionSession<S>
|
||||
@@ -203,15 +199,6 @@ where
|
||||
if self.finalized {
|
||||
return Ok(None);
|
||||
}
|
||||
if self.downstream_committed_at.is_some()
|
||||
&& let Some(event) = self.before_commit.pop_front()
|
||||
{
|
||||
return Ok(Some(CoordinatedEvent::single(
|
||||
event,
|
||||
CommitRequirement::AlreadyCommitted,
|
||||
)));
|
||||
}
|
||||
|
||||
loop {
|
||||
match self.pull().await? {
|
||||
PullOutcome::Event(event) => {
|
||||
@@ -221,25 +208,13 @@ where
|
||||
CommitRequirement::AlreadyCommitted,
|
||||
)));
|
||||
}
|
||||
let commit_significant = match self.commit_policy {
|
||||
StreamCommitPolicy::FirstCanonicalEvent => event.has_canonical_facts(),
|
||||
StreamCommitPolicy::UntilOutputOrTerminal => event.is_commit_significant(),
|
||||
};
|
||||
self.before_commit.push_back(event);
|
||||
if !commit_significant {
|
||||
continue;
|
||||
}
|
||||
self.delivery_pending = true;
|
||||
let events = self.before_commit.drain(..).collect();
|
||||
return CoordinatedEvent::try_batch(
|
||||
events,
|
||||
return Ok(Some(CoordinatedEvent::single(
|
||||
event,
|
||||
CommitRequirement::CommitBeforeDelivery,
|
||||
)
|
||||
.map(Some);
|
||||
}
|
||||
PullOutcome::AttemptDiscarded => {
|
||||
self.before_commit.clear();
|
||||
)));
|
||||
}
|
||||
PullOutcome::AttemptDiscarded => {}
|
||||
PullOutcome::End => {
|
||||
if self.downstream_committed_at.is_some() {
|
||||
self.finish_success().await?;
|
||||
@@ -411,6 +386,7 @@ where
|
||||
&mut current.stream,
|
||||
self.cancellation.clone(),
|
||||
self.deadline,
|
||||
&mut self.deadline_timer,
|
||||
)
|
||||
.await
|
||||
};
|
||||
@@ -651,21 +627,30 @@ where
|
||||
account.clone(),
|
||||
));
|
||||
}
|
||||
self.persist_request().await?;
|
||||
self.engine
|
||||
.store()
|
||||
.record_attempt(AttemptRecord {
|
||||
request_id: self.request_id.clone(),
|
||||
attempt_count: next_attempt,
|
||||
trigger,
|
||||
provider_kind: metadata.provider().clone(),
|
||||
provider_account_id: metadata.provider_account_id().cloned(),
|
||||
provider_account_ref: metadata.provider_account_id().cloned(),
|
||||
upstream_model_id: metadata.upstream_model().clone(),
|
||||
upstream_transport: metadata.transport().as_str().to_owned(),
|
||||
http_version: None,
|
||||
})
|
||||
.await?;
|
||||
let attempt_record = AttemptRecord {
|
||||
request_id: self.request_id.clone(),
|
||||
attempt_count: next_attempt,
|
||||
trigger,
|
||||
provider_kind: metadata.provider().clone(),
|
||||
provider_account_id: metadata.provider_account_id().cloned(),
|
||||
provider_account_ref: metadata.provider_account_id().cloned(),
|
||||
upstream_model_id: metadata.upstream_model().clone(),
|
||||
upstream_transport: metadata.transport().as_str().to_owned(),
|
||||
http_version: None,
|
||||
};
|
||||
if self.request_persisted {
|
||||
self.engine.store().record_attempt(attempt_record).await?;
|
||||
} else {
|
||||
let request = self
|
||||
.pending_request
|
||||
.take()
|
||||
.ok_or(EngineError::InvalidDeliveryState)?;
|
||||
self.engine
|
||||
.store()
|
||||
.create_model_request_with_attempt(request, attempt_record)
|
||||
.await?;
|
||||
self.request_persisted = true;
|
||||
}
|
||||
self.attempts = next_attempt.get();
|
||||
self.current = Some(CurrentAttempt {
|
||||
stream,
|
||||
@@ -679,19 +664,6 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn persist_request(&mut self) -> Result<(), EngineError> {
|
||||
if self.request_persisted {
|
||||
return Ok(());
|
||||
}
|
||||
let request = self
|
||||
.pending_request
|
||||
.take()
|
||||
.ok_or(EngineError::InvalidDeliveryState)?;
|
||||
self.engine.store().create_model_request(request).await?;
|
||||
self.request_persisted = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> {
|
||||
let elapsed = elapsed_ms(self.timing_started_at);
|
||||
observe_event_timing(&mut self.timings, event, elapsed);
|
||||
@@ -713,6 +685,7 @@ where
|
||||
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
|
||||
.await?;
|
||||
current.send_observed = true;
|
||||
self.send_state_watermark = UpstreamSendState::Sent;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -725,6 +698,7 @@ where
|
||||
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
|
||||
.await?;
|
||||
current.send_observed = true;
|
||||
self.send_state_watermark = UpstreamSendState::Sent;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -813,11 +787,15 @@ where
|
||||
self.last_account_exhaustion = Some(error.clone());
|
||||
}
|
||||
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
|
||||
let send_state = if current.send_observed {
|
||||
// attempt_send_state 是本 attempt 自身的发送事实,驱动重试门;
|
||||
// 持久化与终态用请求级水位,二者不可混用(水位会把早先 attempt 的
|
||||
// sent 传染给本 attempt,从而错误放行/拦截重试)。
|
||||
let attempt_send_state = if current.send_observed {
|
||||
UpstreamSendState::Sent
|
||||
} else {
|
||||
error.send_state()
|
||||
};
|
||||
let send_state = self.raise_send_watermark(attempt_send_state);
|
||||
self.engine
|
||||
.store()
|
||||
.mark_send_state(&self.request_id, send_state)
|
||||
@@ -826,14 +804,14 @@ where
|
||||
let continuation_retry = self.prepare_continuation_retry(
|
||||
¤t,
|
||||
&error,
|
||||
send_state,
|
||||
attempt_send_state,
|
||||
provider_proved_replay_safe,
|
||||
);
|
||||
let ordinary_retry = self.required_account.is_none()
|
||||
&& self.continuation_attempt == ContinuationAttempt::None
|
||||
&& self.downstream_committed_at.is_none()
|
||||
&& !self.delivery_pending
|
||||
&& send_state != UpstreamSendState::Ambiguous
|
||||
&& attempt_send_state != UpstreamSendState::Ambiguous
|
||||
&& (self.operation.retry_safety() == RetrySafety::Idempotent
|
||||
|| provider_proved_replay_safe)
|
||||
&& self.attempts < self.plan.max_attempts().get();
|
||||
@@ -841,7 +819,7 @@ where
|
||||
&& provider_proved_replay_safe
|
||||
&& self.downstream_committed_at.is_none()
|
||||
&& !self.delivery_pending
|
||||
&& send_state != UpstreamSendState::Ambiguous
|
||||
&& attempt_send_state != UpstreamSendState::Ambiguous
|
||||
&& self.attempts < self.plan.max_attempts().get()
|
||||
&& current
|
||||
.metadata
|
||||
@@ -953,7 +931,6 @@ where
|
||||
self.timings.first_text_ms = None;
|
||||
self.timings.first_token_ms = None;
|
||||
self.upstream_complete = false;
|
||||
self.before_commit.clear();
|
||||
}
|
||||
|
||||
async fn finish_success(&mut self) -> Result<(), EngineError> {
|
||||
@@ -1028,6 +1005,7 @@ where
|
||||
error: &ProviderError,
|
||||
send_state: UpstreamSendState,
|
||||
) -> Result<(), EngineError> {
|
||||
let send_state = self.raise_send_watermark(send_state);
|
||||
let outcome = if error.kind() == ProviderErrorKind::Cancelled {
|
||||
ExecutionOutcome::Cancelled
|
||||
} else if self.downstream_committed_at.is_some() {
|
||||
@@ -1080,6 +1058,7 @@ where
|
||||
} else {
|
||||
UpstreamSendState::Ambiguous
|
||||
};
|
||||
let send_state = self.raise_send_watermark(send_state);
|
||||
if self.attempts > 0 {
|
||||
self.engine
|
||||
.store()
|
||||
@@ -1200,9 +1179,7 @@ where
|
||||
}
|
||||
|
||||
fn current_send_state(&self) -> UpstreamSendState {
|
||||
if self.attempts == 0 {
|
||||
UpstreamSendState::NotSent
|
||||
} else if self
|
||||
let observed = if self
|
||||
.current
|
||||
.as_ref()
|
||||
.is_some_and(|current| current.send_observed)
|
||||
@@ -1210,7 +1187,15 @@ where
|
||||
UpstreamSendState::Sent
|
||||
} else {
|
||||
UpstreamSendState::NotSent
|
||||
}
|
||||
};
|
||||
escalate_send_state(self.send_state_watermark, observed)
|
||||
}
|
||||
|
||||
/// 抬升并返回请求级发送水位;attempt 间切换(`current` 被取走)后,
|
||||
/// 后续终态沿用已达到的最高档,不会把已落库的 `sent` 写回 `not_sent`。
|
||||
fn raise_send_watermark(&mut self, observed: UpstreamSendState) -> UpstreamSendState {
|
||||
self.send_state_watermark = escalate_send_state(self.send_state_watermark, observed);
|
||||
self.send_state_watermark
|
||||
}
|
||||
|
||||
fn record_current_provider_success(&mut self) {
|
||||
@@ -1263,17 +1248,17 @@ async fn poll_stream_item(
|
||||
stream: &mut ProviderStream,
|
||||
cancellation: CancellationToken,
|
||||
deadline: SystemTime,
|
||||
mut deadline_timer: &mut Fuse<Delay>,
|
||||
) -> PollBoundary {
|
||||
let Ok(remaining) = deadline.duration_since(SystemTime::now()) else {
|
||||
if SystemTime::now() >= deadline {
|
||||
return PollBoundary::Deadline;
|
||||
};
|
||||
}
|
||||
let next = stream.next().fuse();
|
||||
let cancelled = cancellation.cancelled().fuse();
|
||||
let timeout = Delay::new(remaining).fuse();
|
||||
pin_mut!(next, cancelled, timeout);
|
||||
pin_mut!(next, cancelled);
|
||||
select_biased! {
|
||||
_ = cancelled => PollBoundary::Cancelled,
|
||||
_ = timeout => PollBoundary::Deadline,
|
||||
_ = deadline_timer => PollBoundary::Deadline,
|
||||
item = next => PollBoundary::Item(item),
|
||||
}
|
||||
}
|
||||
@@ -1338,6 +1323,18 @@ fn provider_proved_replay_safe(error: &ProviderError) -> bool {
|
||||
|| (error.send_state() != UpstreamSendState::Ambiguous && error.replay_is_safe())
|
||||
}
|
||||
|
||||
/// 发送状态合并档位:`Sent` > `Ambiguous` > `NotSent`。
|
||||
/// 只要任一 attempt 达到过高档,请求整体就不允许回落到低档。
|
||||
const fn escalate_send_state(a: UpstreamSendState, b: UpstreamSendState) -> UpstreamSendState {
|
||||
match (a, b) {
|
||||
(UpstreamSendState::Sent, _) | (_, UpstreamSendState::Sent) => UpstreamSendState::Sent,
|
||||
(UpstreamSendState::Ambiguous, _) | (_, UpstreamSendState::Ambiguous) => {
|
||||
UpstreamSendState::Ambiguous
|
||||
}
|
||||
(UpstreamSendState::NotSent, UpstreamSendState::NotSent) => UpstreamSendState::NotSent,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_event_timing(timings: &mut ModelRequestTimings, event: &GatewayEvent, elapsed_ms: u64) {
|
||||
timings.first_event_ms.get_or_insert(elapsed_ms);
|
||||
match event {
|
||||
|
||||
@@ -22,14 +22,18 @@ use crate::engine::provider::ProviderRegistry;
|
||||
use crate::engine::{
|
||||
AttemptCoordinator, AttemptRecord, CancellationToken, CommitRequirement, CoordinatedEvent,
|
||||
EngineError, ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization,
|
||||
ModelRequestId, NewModelRequest, ProviderAttemptOutcome, RecoveryReport, UpstreamSendState,
|
||||
ModelRequestId, NewModelRequest, ProbeFailure, ProviderAccountId, ProviderAttemptOutcome,
|
||||
RecoveryReport, UpstreamSendState,
|
||||
};
|
||||
use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind, StoreError};
|
||||
use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader};
|
||||
use crate::operation::{Operation, ProviderSessionState};
|
||||
use crate::policy::{ClientApiKeyId, ClientPolicy};
|
||||
use crate::routing::snapshot::RuntimeSnapshotHandle;
|
||||
use crate::routing::{ProviderKind, PublicModelId, RoutingContext, RuntimeSnapshot};
|
||||
use crate::routing::{
|
||||
ProviderKind, PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
};
|
||||
|
||||
const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60);
|
||||
const CONTINUATION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
@@ -139,6 +143,9 @@ pub trait ExecutionService: Send + Sync {
|
||||
plaintext: &str,
|
||||
) -> Result<AuthenticatedClient, ClientAuthenticationError>;
|
||||
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId>;
|
||||
fn public_model_profiles(&self, _client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
Vec::new()
|
||||
}
|
||||
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool;
|
||||
fn start(
|
||||
&self,
|
||||
@@ -182,6 +189,8 @@ pub struct DefaultExecutionService {
|
||||
snapshots: RuntimeSnapshotHandle,
|
||||
coordinator: Arc<AttemptCoordinator<dyn ExecutionStore>>,
|
||||
probe_coordinator: Arc<AttemptCoordinator<dyn ExecutionStore>>,
|
||||
/// probe 自身走 transient store,探测失败仍写入持久 store 的 ops_events。
|
||||
observations: Arc<dyn ExecutionStore>,
|
||||
providers: ProviderRegistry,
|
||||
admissions: Arc<dyn ClientAdmissionPort>,
|
||||
circuits: Arc<dyn ProviderCircuitPort>,
|
||||
@@ -200,6 +209,7 @@ impl DefaultExecutionService {
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
) -> Self {
|
||||
let observations = Arc::clone(&execution);
|
||||
let engine = GatewayEngine::<dyn ExecutionStore>::new(execution, providers.clone());
|
||||
let transient: Arc<dyn ExecutionStore> = Arc::new(TransientExecutionStore);
|
||||
let probe_engine = GatewayEngine::<dyn ExecutionStore>::new(transient, providers.clone());
|
||||
@@ -207,6 +217,7 @@ impl DefaultExecutionService {
|
||||
snapshots,
|
||||
coordinator: Arc::new(AttemptCoordinator::new(engine)),
|
||||
probe_coordinator: Arc::new(AttemptCoordinator::new(probe_engine)),
|
||||
observations,
|
||||
providers,
|
||||
admissions,
|
||||
circuits,
|
||||
@@ -412,6 +423,11 @@ impl DefaultExecutionService {
|
||||
upstream_model,
|
||||
operation,
|
||||
} = request;
|
||||
let observed = ProbeObservation {
|
||||
provider_kind: provider_kind.clone(),
|
||||
account_id: account_id.clone(),
|
||||
upstream_model: upstream_model.clone(),
|
||||
};
|
||||
let snapshot = self.snapshots.acquire().map_err(|_| {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
@@ -459,7 +475,7 @@ impl DefaultExecutionService {
|
||||
started_at,
|
||||
deadline_at,
|
||||
};
|
||||
let mut session = self
|
||||
let mut session = match self
|
||||
.probe_coordinator
|
||||
.start(
|
||||
new_request,
|
||||
@@ -470,18 +486,33 @@ impl DefaultExecutionService {
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.map_err(|error| gateway_error_from_engine(&error))?;
|
||||
{
|
||||
Ok(session) => session,
|
||||
Err(error) => {
|
||||
return Err(self
|
||||
.observe_probe_failure(&observed, started_at, &error)
|
||||
.await);
|
||||
}
|
||||
};
|
||||
let events = session.collect_uncommitted().await;
|
||||
publish_provider_attempt_outcomes(
|
||||
self.circuits.as_ref(),
|
||||
session.provider_attempt_outcomes(),
|
||||
)
|
||||
.await;
|
||||
let events = events.map_err(|error| gateway_error_from_engine(&error))?;
|
||||
session
|
||||
.commit_downstream(Some(200))
|
||||
.await
|
||||
.map_err(|error| gateway_error_from_engine(&error))?;
|
||||
let events = match events {
|
||||
Ok(events) => events,
|
||||
Err(error) => {
|
||||
return Err(self
|
||||
.observe_probe_failure(&observed, started_at, &error)
|
||||
.await);
|
||||
}
|
||||
};
|
||||
if let Err(error) = session.commit_downstream(Some(200)).await {
|
||||
return Err(self
|
||||
.observe_probe_failure(&observed, started_at, &error)
|
||||
.await);
|
||||
}
|
||||
Ok(AccountProbeResult {
|
||||
text: events
|
||||
.into_iter()
|
||||
@@ -493,6 +524,34 @@ impl DefaultExecutionService {
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 探测失败先落库再转成客户端错误;转换会丢掉上游状态码与 provider code。
|
||||
async fn observe_probe_failure(
|
||||
&self,
|
||||
observed: &ProbeObservation,
|
||||
started_at: SystemTime,
|
||||
error: &EngineError,
|
||||
) -> GatewayError {
|
||||
if let EngineError::Provider(provider_error) = error {
|
||||
let _ = self
|
||||
.observations
|
||||
.record_probe_failure(ProbeFailure {
|
||||
provider_kind: observed.provider_kind.clone(),
|
||||
account_id: observed.account_id.clone(),
|
||||
upstream_model_id: observed.upstream_model.clone(),
|
||||
error: provider_error.clone(),
|
||||
latency: started_at.elapsed().unwrap_or_default(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
gateway_error_from_engine(error)
|
||||
}
|
||||
}
|
||||
|
||||
struct ProbeObservation {
|
||||
provider_kind: ProviderKind,
|
||||
account_id: ProviderAccountId,
|
||||
upstream_model: UpstreamModelId,
|
||||
}
|
||||
|
||||
struct TransientExecutionStore;
|
||||
@@ -568,6 +627,12 @@ impl ExecutionService for DefaultExecutionService {
|
||||
.public_models_for_provider(client.policy.provider_kind())
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
client
|
||||
.snapshot
|
||||
.public_model_profiles_for_provider(client.policy.provider_kind())
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
client
|
||||
.snapshot
|
||||
@@ -857,7 +922,7 @@ fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeCont
|
||||
{
|
||||
return;
|
||||
}
|
||||
*operation = operation.clone().with_provider_session_state(state.clone());
|
||||
operation.set_provider_session_state(state.clone());
|
||||
}
|
||||
|
||||
pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
|
||||
|
||||
@@ -543,6 +543,16 @@ pub struct IntermediateFailure {
|
||||
pub latency: Duration,
|
||||
}
|
||||
|
||||
/// 不属于任何 `model_requests` 行的管理端账号探测失败。
|
||||
#[derive(Debug)]
|
||||
pub struct ProbeFailure {
|
||||
pub provider_kind: ProviderKind,
|
||||
pub account_id: ProviderAccountId,
|
||||
pub upstream_model_id: UpstreamModelId,
|
||||
pub error: ProviderError,
|
||||
pub latency: Duration,
|
||||
}
|
||||
|
||||
/// `model_requests` 可用的毫秒级阶段耗时。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ModelRequestTimings {
|
||||
@@ -596,6 +606,19 @@ pub struct RecoveryReport {
|
||||
pub trait ExecutionStore: Send + Sync {
|
||||
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError>;
|
||||
async fn record_attempt(&self, attempt: AttemptRecord) -> Result<(), StoreError>;
|
||||
/// 请求插入与首次 attempt 合并持久化;两者写同一行,支持合并写的
|
||||
/// store 可覆写为单次往返,缩短首 token 前的关键路径。
|
||||
///
|
||||
/// 覆写实现应保证原子性:默认的两步实现里第二步失败会留下一条
|
||||
/// `running` 请求行,由 deadline 回收器收敛。
|
||||
async fn create_model_request_with_attempt(
|
||||
&self,
|
||||
request: NewModelRequest,
|
||||
attempt: AttemptRecord,
|
||||
) -> Result<(), StoreError> {
|
||||
self.create_model_request(request).await?;
|
||||
self.record_attempt(attempt).await
|
||||
}
|
||||
async fn mark_send_state(
|
||||
&self,
|
||||
request_id: &ModelRequestId,
|
||||
@@ -616,6 +639,10 @@ pub trait ExecutionStore: Send + Sync {
|
||||
&self,
|
||||
failure: IntermediateFailure,
|
||||
) -> Result<(), StoreError>;
|
||||
/// 记录不挂在 `model_requests` 上的账号探测失败;默认丢弃。
|
||||
async fn record_probe_failure(&self, _failure: ProbeFailure) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn finalize_model_request(
|
||||
&self,
|
||||
finalization: ModelRequestFinalization,
|
||||
|
||||
@@ -20,7 +20,9 @@ use crate::error::{
|
||||
};
|
||||
use crate::event::{EventSequenceValidator, ProviderEvent};
|
||||
use crate::operation::Operation;
|
||||
use crate::routing::{ModelCapabilities, ProviderCandidate, ProviderKind, UpstreamModelId};
|
||||
use crate::routing::{
|
||||
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId,
|
||||
};
|
||||
|
||||
/// Box 只出现在 Provider Registry 的统一 event envelope 边界。
|
||||
pub type EventStream =
|
||||
@@ -433,6 +435,7 @@ pub struct ProviderRequestObservation {
|
||||
pub struct ProviderModelCapabilities {
|
||||
upstream_model: UpstreamModelId,
|
||||
capabilities: ModelCapabilities,
|
||||
presentation: Option<ModelPresentation>,
|
||||
}
|
||||
|
||||
/// Provider 实时目录成功发布后的进程内单调代次。
|
||||
@@ -459,9 +462,16 @@ impl ProviderModelCapabilities {
|
||||
Self {
|
||||
upstream_model,
|
||||
capabilities,
|
||||
presentation: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
|
||||
self.presentation = Some(presentation);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn upstream_model(&self) -> &UpstreamModelId {
|
||||
&self.upstream_model
|
||||
@@ -471,6 +481,11 @@ impl ProviderModelCapabilities {
|
||||
pub const fn capabilities(&self) -> &ModelCapabilities {
|
||||
&self.capabilities
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn presentation(&self) -> Option<&ModelPresentation> {
|
||||
self.presentation.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderRequest {
|
||||
|
||||
@@ -715,12 +715,6 @@ impl ProviderEvent {
|
||||
&self.canonical
|
||||
}
|
||||
|
||||
/// 返回 Core 可改写客户端响应身份的全部 facts。
|
||||
#[must_use]
|
||||
pub fn canonical_facts_mut(&mut self) -> &mut [GatewayEvent] {
|
||||
&mut self.canonical
|
||||
}
|
||||
|
||||
/// 返回协议原生表达。
|
||||
#[must_use]
|
||||
pub fn wire_event(&self) -> Option<&ProtocolWireEvent> {
|
||||
@@ -756,18 +750,6 @@ impl ProviderEvent {
|
||||
pub fn has_client_event(&self) -> bool {
|
||||
!self.canonical.is_empty() || self.wire.is_some()
|
||||
}
|
||||
|
||||
/// 返回是否包含足以冻结下游 commit barrier 的 canonical fact。
|
||||
///
|
||||
/// 单独的 wire event 不能冻结边界:协议层在看到 `Started`
|
||||
/// 携带的网关响应身份前无法安全编码它。Core 会保留这些
|
||||
/// 事件,直到首个 canonical 可见增量或终态再一并交付。
|
||||
#[must_use]
|
||||
pub fn is_commit_significant(&self) -> bool {
|
||||
self.canonical
|
||||
.iter()
|
||||
.any(GatewayEvent::is_commit_significant)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<GatewayEvent> for ProviderEvent {
|
||||
@@ -776,23 +758,6 @@ impl From<GatewayEvent> for ProviderEvent {
|
||||
}
|
||||
}
|
||||
|
||||
impl GatewayEvent {
|
||||
/// 返回该 fact 是否足以冻结下游 commit barrier。
|
||||
///
|
||||
/// 生命周期、结构与结算 facts 可以在换号前安全丢弃;首个可见增量或
|
||||
/// Provider 正常终态才会使客户端输出不可撤回。
|
||||
#[must_use]
|
||||
pub const fn is_commit_significant(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::TextDelta(_)
|
||||
| Self::ReasoningDelta(_)
|
||||
| Self::ToolCallDelta(_)
|
||||
| Self::Completed(_)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonical event 顺序错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum EventSequenceError {
|
||||
|
||||
@@ -269,10 +269,15 @@ impl GenerateRequest {
|
||||
/// 附着同一客户端连接上一轮由 Provider 返回的不透明状态。
|
||||
#[must_use]
|
||||
pub fn with_provider_session_state(mut self, state: ProviderSessionState) -> Self {
|
||||
Arc::make_mut(&mut self.payload).provider_session_state = Some(state);
|
||||
self.set_provider_session_state(state);
|
||||
self
|
||||
}
|
||||
|
||||
/// 原地附着会话状态;payload 独占时不复制正文。
|
||||
pub fn set_provider_session_state(&mut self, state: ProviderSessionState) {
|
||||
Arc::make_mut(&mut self.payload).provider_session_state = Some(state);
|
||||
}
|
||||
|
||||
/// 返回最大输出 token 数。
|
||||
#[must_use]
|
||||
pub fn max_output_tokens(&self) -> Option<u64> {
|
||||
@@ -597,6 +602,13 @@ impl Operation {
|
||||
}
|
||||
}
|
||||
|
||||
/// 原地附着 Provider 私有状态;payload 独占时不复制正文。
|
||||
pub fn set_provider_session_state(&mut self, state: ProviderSessionState) {
|
||||
if let Self::Generate(request) = self {
|
||||
request.set_provider_session_state(state);
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回稳定 operation 分类。
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> OperationKind {
|
||||
|
||||
@@ -474,9 +474,22 @@ pub trait ProviderCredentialStatePort: Send + Sync {
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
|
||||
|
||||
/// 记录一次瞬态刷新失败并返回窗口内累计失败次数;每次失败刷新 TTL。
|
||||
fn record_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
window: Duration,
|
||||
) -> BoxFuture<'a, Result<u32, ProviderStoreError>>;
|
||||
|
||||
/// 凭据完整成功轮换后清零失败计数,退避窗口重新从 base 起步。
|
||||
fn clear_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<(), ProviderStoreError>>;
|
||||
}
|
||||
|
||||
/// 临时 cooldown 只保存调度截止时间;原因事实仍由 PostgreSQL 账号状态持有。
|
||||
/// 临时 cooldown 只保存可丢失的调度截止时间,不进入账号持久状态。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderCooldown {
|
||||
account_id: ProviderAccountId,
|
||||
@@ -592,6 +605,33 @@ impl ProviderRefreshPolicy {
|
||||
}
|
||||
}
|
||||
|
||||
/// 指数退避基准延迟;attempt=1 即为该值。
|
||||
const REFRESH_BACKOFF_BASE_DELAY: Duration = Duration::from_secs(5);
|
||||
/// 每多一次连续失败,基准延迟乘以该因子。
|
||||
const REFRESH_BACKOFF_FACTOR: u32 = 3;
|
||||
/// 退避延迟上限,避免连续失败时无限增长。
|
||||
const REFRESH_BACKOFF_CAP: Duration = Duration::from_secs(300);
|
||||
|
||||
/// 基于连续失败计数的指数退避重试时刻;复用 `provider_refresh_retry_at` 的稳定扰动。
|
||||
///
|
||||
/// `attempt` 为窗口内累计失败次数(0 与 1 等价,均取基准延迟)。延迟按
|
||||
/// `base * factor^(attempt-1)` 增长并封顶到 `REFRESH_BACKOFF_CAP`。
|
||||
pub fn provider_refresh_backoff_at(
|
||||
account_id: &ProviderAccountId,
|
||||
observed_at: SystemTime,
|
||||
attempt: u32,
|
||||
reason: &'static str,
|
||||
) -> Result<SystemTime, ProviderStoreError> {
|
||||
let exponent = attempt.saturating_sub(1);
|
||||
let multiplier = REFRESH_BACKOFF_FACTOR.saturating_pow(exponent);
|
||||
let scaled_seconds = REFRESH_BACKOFF_BASE_DELAY
|
||||
.as_secs()
|
||||
.saturating_mul(u64::from(multiplier))
|
||||
.min(REFRESH_BACKOFF_CAP.as_secs());
|
||||
let base_delay = Duration::from_secs(scaled_seconds);
|
||||
provider_refresh_retry_at(account_id, observed_at, base_delay, reason)
|
||||
}
|
||||
|
||||
/// 临时失败后的持久重试时刻;稳定扰动避免多实例同频重试。
|
||||
pub fn provider_refresh_retry_at(
|
||||
account_id: &ProviderAccountId,
|
||||
|
||||
@@ -136,6 +136,177 @@ pub struct ModelCapabilities {
|
||||
upstream_validates_features: bool,
|
||||
}
|
||||
|
||||
/// Provider 为客户端模型目录提供的展示与交互能力。
|
||||
///
|
||||
/// 该值不参与路由,也不把任一 Provider 的 wire 类型带入 Core。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ModelPresentation {
|
||||
display_name: Option<String>,
|
||||
description: Option<String>,
|
||||
default_reasoning_effort: Option<String>,
|
||||
supported_reasoning_efforts: Vec<String>,
|
||||
context_window_tokens: Option<u64>,
|
||||
image_input: bool,
|
||||
agent_tools: bool,
|
||||
parallel_tool_calls: bool,
|
||||
search_tool: bool,
|
||||
image_detail_original: bool,
|
||||
verbosity: bool,
|
||||
hidden: bool,
|
||||
}
|
||||
|
||||
impl ModelPresentation {
|
||||
#[must_use]
|
||||
pub fn new(display_name: Option<String>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
display_name,
|
||||
description,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_reasoning(
|
||||
mut self,
|
||||
default_effort: Option<String>,
|
||||
supported_efforts: Vec<String>,
|
||||
) -> Self {
|
||||
self.default_reasoning_effort = default_effort;
|
||||
self.supported_reasoning_efforts = supported_efforts;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_context_window_tokens(mut self, context_window_tokens: Option<u64>) -> Self {
|
||||
self.context_window_tokens = context_window_tokens;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_image_input(mut self, image_input: bool) -> Self {
|
||||
self.image_input = image_input;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_agent_tools(mut self, agent_tools: bool, parallel_tool_calls: bool) -> Self {
|
||||
self.agent_tools = agent_tools;
|
||||
self.parallel_tool_calls = parallel_tool_calls;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_search_tool(mut self, search_tool: bool) -> Self {
|
||||
self.search_tool = search_tool;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_image_detail_original(mut self, image_detail_original: bool) -> Self {
|
||||
self.image_detail_original = image_detail_original;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_verbosity(mut self, verbosity: bool) -> Self {
|
||||
self.verbosity = verbosity;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_hidden(mut self, hidden: bool) -> Self {
|
||||
self.hidden = hidden;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn display_name(&self) -> Option<&str> {
|
||||
self.display_name.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn description(&self) -> Option<&str> {
|
||||
self.description.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn default_reasoning_effort(&self) -> Option<&str> {
|
||||
self.default_reasoning_effort.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn supported_reasoning_efforts(&self) -> &[String] {
|
||||
&self.supported_reasoning_efforts
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn context_window_tokens(&self) -> Option<u64> {
|
||||
self.context_window_tokens
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn image_input(&self) -> bool {
|
||||
self.image_input
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn agent_tools(&self) -> bool {
|
||||
self.agent_tools
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn parallel_tool_calls(&self) -> bool {
|
||||
self.parallel_tool_calls
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn search_tool(&self) -> bool {
|
||||
self.search_tool
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn image_detail_original(&self) -> bool {
|
||||
self.image_detail_original
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn verbosity(&self) -> bool {
|
||||
self.verbosity
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn hidden(&self) -> bool {
|
||||
self.hidden
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个公开模型及其 Provider 编译后的客户端画像。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PublicModelProfile {
|
||||
model: PublicModelId,
|
||||
presentation: ModelPresentation,
|
||||
}
|
||||
|
||||
impl PublicModelProfile {
|
||||
#[must_use]
|
||||
pub const fn new(model: PublicModelId, presentation: ModelPresentation) -> Self {
|
||||
Self {
|
||||
model,
|
||||
presentation,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn model(&self) -> &PublicModelId {
|
||||
&self.model
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn presentation(&self) -> &ModelPresentation {
|
||||
&self.presentation
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelCapabilities {
|
||||
#[must_use]
|
||||
pub fn new(operations: BTreeSet<OperationKind>, max_output_tokens: Option<u64>) -> Self {
|
||||
@@ -205,6 +376,7 @@ pub struct ProviderModel {
|
||||
provider: ProviderKind,
|
||||
upstream_model: UpstreamModelId,
|
||||
capabilities: ModelCapabilities,
|
||||
presentation: Option<ModelPresentation>,
|
||||
}
|
||||
|
||||
impl ProviderModel {
|
||||
@@ -218,9 +390,16 @@ impl ProviderModel {
|
||||
provider,
|
||||
upstream_model,
|
||||
capabilities,
|
||||
presentation: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
|
||||
self.presentation = Some(presentation);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider(&self) -> &ProviderKind {
|
||||
&self.provider
|
||||
@@ -230,6 +409,11 @@ impl ProviderModel {
|
||||
pub const fn upstream_model(&self) -> &UpstreamModelId {
|
||||
&self.upstream_model
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn presentation(&self) -> Option<&ModelPresentation> {
|
||||
self.presentation.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 本次请求选择 Provider 时使用的动态过滤事实。
|
||||
|
||||
@@ -231,11 +231,15 @@ async fn compile_runtime_snapshot(
|
||||
continue;
|
||||
};
|
||||
provider_models.extend(models.into_iter().map(|model| {
|
||||
ProviderModel::new(
|
||||
let compiled = ProviderModel::new(
|
||||
provider.clone(),
|
||||
model.upstream_model().clone(),
|
||||
model.capabilities().clone(),
|
||||
)
|
||||
);
|
||||
match model.presentation().cloned() {
|
||||
Some(presentation) => compiled.with_presentation(presentation),
|
||||
None => compiled,
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -287,6 +291,8 @@ pub struct RuntimeSnapshot {
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
providers: Arc<BTreeSet<ProviderKind>>,
|
||||
provider_models: Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>>,
|
||||
provider_model_presentations:
|
||||
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
|
||||
model_mappings: Arc<BTreeMap<String, String>>,
|
||||
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
|
||||
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
|
||||
@@ -313,17 +319,24 @@ impl RuntimeSnapshot {
|
||||
|
||||
let mut model_map =
|
||||
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
|
||||
let mut presentation_map =
|
||||
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>::new();
|
||||
for model in provider_models {
|
||||
if !provider_set.contains(model.provider()) {
|
||||
let ProviderModel {
|
||||
provider,
|
||||
upstream_model,
|
||||
capabilities,
|
||||
presentation,
|
||||
} = model;
|
||||
if !provider_set.contains(&provider) {
|
||||
return Err(RoutingError::NotFound {
|
||||
entity: "provider",
|
||||
id: model.provider().to_string(),
|
||||
id: provider.to_string(),
|
||||
});
|
||||
}
|
||||
let models = model_map.entry(model.provider).or_default();
|
||||
let upstream_model = model.upstream_model;
|
||||
let models = model_map.entry(provider.clone()).or_default();
|
||||
if models
|
||||
.insert(upstream_model.clone(), model.capabilities)
|
||||
.insert(upstream_model.clone(), capabilities)
|
||||
.is_some()
|
||||
{
|
||||
return Err(RoutingError::DuplicateEntity {
|
||||
@@ -331,6 +344,12 @@ impl RuntimeSnapshot {
|
||||
id: upstream_model.to_string(),
|
||||
});
|
||||
}
|
||||
if let Some(presentation) = presentation {
|
||||
presentation_map
|
||||
.entry(provider)
|
||||
.or_default()
|
||||
.insert(upstream_model, presentation);
|
||||
}
|
||||
}
|
||||
|
||||
let mut client_policy_map = BTreeMap::new();
|
||||
@@ -350,6 +369,7 @@ impl RuntimeSnapshot {
|
||||
account_selection_policy,
|
||||
providers: Arc::new(provider_set),
|
||||
provider_models: Arc::new(model_map),
|
||||
provider_model_presentations: Arc::new(presentation_map),
|
||||
model_mappings: Arc::new(BTreeMap::new()),
|
||||
provider_catalog_generations: Arc::new(BTreeMap::new()),
|
||||
client_policies: Arc::new(client_policy_map),
|
||||
@@ -402,6 +422,38 @@ impl RuntimeSnapshot {
|
||||
models.into_iter().collect()
|
||||
}
|
||||
|
||||
/// 返回 Provider 已明确声明画像的公开模型;没有画像时不猜测 Provider 语义。
|
||||
#[must_use]
|
||||
pub fn public_model_profiles_for_provider(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Vec<super::PublicModelProfile> {
|
||||
let Some(presentations) = self.provider_model_presentations.get(provider) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut profiles = BTreeMap::new();
|
||||
for (model, presentation) in presentations {
|
||||
if let Ok(public_model) = PublicModelId::new(model.as_str().to_owned()) {
|
||||
profiles.insert(public_model, presentation.clone());
|
||||
}
|
||||
}
|
||||
for alias in self.model_mappings.keys() {
|
||||
let target = self.mapped_model(alias);
|
||||
let Some(presentation) = presentations.iter().find_map(|(model, presentation)| {
|
||||
(model.as_str() == target).then_some(presentation)
|
||||
}) else {
|
||||
continue;
|
||||
};
|
||||
if let Ok(public_model) = PublicModelId::new(alias.clone()) {
|
||||
profiles.insert(public_model, presentation.clone());
|
||||
}
|
||||
}
|
||||
profiles
|
||||
.into_iter()
|
||||
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn public_models(&self) -> Vec<PublicModelId> {
|
||||
self.providers
|
||||
|
||||
@@ -1005,6 +1005,51 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streaming_native_continuation_delivers_started_event_before_later_output() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, _) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_one",
|
||||
items: complete_stream(None),
|
||||
}]);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response").expect("previous response"),
|
||||
SafeUpstreamValue::new("upstream-response").expect("upstream response"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
|
||||
let first = block_on(session.next_event())
|
||||
.expect("first event")
|
||||
.expect("started event");
|
||||
assert_eq!(
|
||||
first.commit_requirement(),
|
||||
CommitRequirement::CommitBeforeDelivery
|
||||
);
|
||||
let events = first.into_provider_events();
|
||||
assert!(matches!(
|
||||
events.as_slice(),
|
||||
[event]
|
||||
if matches!(event.canonical_facts(), [GatewayEvent::Started(_)])
|
||||
));
|
||||
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit started event");
|
||||
while block_on(session.next_event())
|
||||
.expect("remaining event")
|
||||
.is_some()
|
||||
{}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
@@ -1575,6 +1620,8 @@ fn rate_limited_account_exhaustion_survives_a_later_empty_selection() {
|
||||
assert_eq!(state.attempts.len(), 1);
|
||||
assert_eq!(state.intermediate_failures, 1);
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
// attempt-1 已把 sent 落库;attempt-2 空选路终态不得降级回 not_sent。
|
||||
assert_eq!(state.finalizations[0].send_state, UpstreamSendState::Sent);
|
||||
assert_eq!(state.finalizations[0].upstream_status_code, Some(429));
|
||||
assert_eq!(
|
||||
state.finalizations[0].provider_error_code.as_deref(),
|
||||
|
||||
@@ -24,12 +24,15 @@ use gateway_core::engine::execution::{
|
||||
ProviderCircuitError, ProviderCircuitPort, provider_failure_affects_circuit,
|
||||
};
|
||||
use gateway_core::engine::probe::{AccountProbe, AccountProbeRequest};
|
||||
use gateway_core::engine::provider::ProviderRegistry;
|
||||
use gateway_core::engine::{
|
||||
AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId,
|
||||
NewModelRequest, RecoveryReport, UpstreamSendState,
|
||||
use gateway_core::engine::provider::{
|
||||
Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry,
|
||||
ProviderRequest, ProviderStream,
|
||||
};
|
||||
use gateway_core::error::{GatewayErrorKind, ProviderErrorKind, StoreError};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
|
||||
ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport, UpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::{GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError};
|
||||
use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload};
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
|
||||
@@ -78,6 +81,64 @@ fn account_probe_should_not_write_to_the_persistent_execution_store() {
|
||||
assert!(!store.touched.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn probe_failures_should_be_observable_without_a_model_request_row() {
|
||||
let store = Arc::new(TrackingExecutionStore::default());
|
||||
let providers = ProviderRegistry::new([Arc::new(FailingProvider) as Arc<dyn Provider>])
|
||||
.expect("provider registry");
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(probe_snapshot()),
|
||||
store.clone(),
|
||||
providers,
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
|
||||
let error = block_on(service.probe(AccountProbeRequest {
|
||||
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
|
||||
operation: probe_operation(),
|
||||
}))
|
||||
.expect_err("the provider rejects every probe");
|
||||
|
||||
assert_eq!(error.kind(), GatewayErrorKind::UpstreamUnavailable);
|
||||
assert!(!store.touched.load(Ordering::SeqCst));
|
||||
assert_eq!(store.probe_failures(), vec!["transport".to_owned()]);
|
||||
}
|
||||
|
||||
struct FailingProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for FailingProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"openai"
|
||||
}
|
||||
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::default()
|
||||
}
|
||||
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_request: ProviderRequest,
|
||||
_context: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::NotSent,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn successful_authentication_should_record_client_key_usage() {
|
||||
let usage = Arc::new(RecordingClientApiKeyUsage::default());
|
||||
@@ -121,12 +182,20 @@ impl ClientApiKeyUsageSink for RecordingClientApiKeyUsage {
|
||||
#[derive(Default)]
|
||||
struct TrackingExecutionStore {
|
||||
touched: AtomicBool,
|
||||
probe_failures: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl TrackingExecutionStore {
|
||||
fn touch(&self) {
|
||||
self.touched.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
fn probe_failures(&self) -> Vec<String> {
|
||||
self.probe_failures
|
||||
.lock()
|
||||
.expect("probe failures lock")
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -170,6 +239,14 @@ impl ExecutionStore for TrackingExecutionStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), StoreError> {
|
||||
self.probe_failures
|
||||
.lock()
|
||||
.expect("probe failures lock")
|
||||
.push(failure.error.kind().as_str().to_owned());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> {
|
||||
self.touch();
|
||||
Ok(())
|
||||
|
||||
@@ -24,8 +24,8 @@ use gateway_core::routing::snapshot::{
|
||||
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, PublicModelId, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderKind, PublicModelId,
|
||||
RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use gateway_core::task::WorkerKind;
|
||||
|
||||
@@ -135,13 +135,19 @@ impl Provider for PublishingCatalogProvider {
|
||||
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
self.generation.store(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(vec![ProviderModelCapabilities::new(
|
||||
UpstreamModelId::new("upstream-model").expect("model"),
|
||||
ModelCapabilities::new(
|
||||
std::collections::BTreeSet::from([OperationKind::Generate]),
|
||||
Ok(vec![
|
||||
ProviderModelCapabilities::new(
|
||||
UpstreamModelId::new("upstream-model").expect("model"),
|
||||
ModelCapabilities::new(
|
||||
std::collections::BTreeSet::from([OperationKind::Generate]),
|
||||
None,
|
||||
),
|
||||
)
|
||||
.with_presentation(ModelPresentation::new(
|
||||
Some("Upstream Model".to_owned()),
|
||||
None,
|
||||
),
|
||||
)])
|
||||
)),
|
||||
])
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
@@ -215,6 +221,15 @@ fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
|
||||
.map(|generation| generation.get()),
|
||||
Some(1),
|
||||
);
|
||||
let profiles =
|
||||
snapshot.public_model_profiles_for_provider(&ProviderKind::new("alpha").expect("provider"));
|
||||
assert_eq!(
|
||||
profiles
|
||||
.iter()
|
||||
.map(|profile| profile.model().as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["public-model", "upstream-model"],
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -28,7 +28,6 @@ tokio = { version = "1.52.3", default-features = false, features = ["macros", "r
|
||||
tracing = "0.1.44"
|
||||
tracing-appender = "0.2.5"
|
||||
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
|
||||
[dev-dependencies]
|
||||
filetime = "0.2.27"
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! HTTP 监听、OS signal、原子连接注册与优雅 drain。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::Router;
|
||||
@@ -121,16 +121,44 @@ pub(crate) async fn serve_router(
|
||||
wait_for_shutdown(&shutdown_cancellation).await;
|
||||
shutdown_connections.begin_draining();
|
||||
};
|
||||
axum::serve(
|
||||
// axum 的优雅关闭会无限等待存量连接结束,而 SSE 等长响应体不观察
|
||||
// 取消信号:慢客户端可以把关闭拖住直到被 SIGKILL。整个 drain(axum
|
||||
// 优雅关闭 + 游离 guard 等待)共享同一个从取消信号起算的绝对截止点,
|
||||
// 逾期即放弃等待,存量连接随进程退出终止。
|
||||
let drain_deadline_at = Arc::new(OnceLock::new());
|
||||
let drain_deadline = {
|
||||
let cancellation = cancellation.clone();
|
||||
let deadline_at = Arc::clone(&drain_deadline_at);
|
||||
async move {
|
||||
cancellation.cancelled().await;
|
||||
let deadline = tokio::time::Instant::now() + drain_timeout;
|
||||
let _ = deadline_at.set(deadline);
|
||||
tokio::time::sleep_until(deadline).await;
|
||||
}
|
||||
};
|
||||
let serve = axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(shutdown)
|
||||
.await
|
||||
.map_err(ServeError::Serve)?;
|
||||
|
||||
connections.begin_draining();
|
||||
connections.wait_until_idle(drain_timeout).await;
|
||||
.with_graceful_shutdown(shutdown);
|
||||
tokio::select! {
|
||||
biased;
|
||||
result = serve => {
|
||||
result.map_err(ServeError::Serve)?;
|
||||
connections.begin_draining();
|
||||
let remaining = drain_deadline_at.get().map_or(drain_timeout, |deadline| {
|
||||
deadline.saturating_duration_since(tokio::time::Instant::now())
|
||||
});
|
||||
connections.wait_until_idle(remaining).await;
|
||||
}
|
||||
() = drain_deadline => {
|
||||
tracing::warn!(
|
||||
target: "gateway_shutdown",
|
||||
timeout_secs = drain_timeout.as_secs(),
|
||||
"优雅 drain 超时,放弃等待存量连接"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ pub(crate) async fn download_file(
|
||||
.map_err(|error| super::internal(format!("failed to create download file: {error}")))?;
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut downloaded = 0_u64;
|
||||
let mut last_percent = 0_u8;
|
||||
let mut next_progress = 10_u8;
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|error| upstream(format!("download stream failed: {error}")))?;
|
||||
downloaded = downloaded.saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX));
|
||||
@@ -59,9 +59,17 @@ pub(crate) async fn download_file(
|
||||
file.write_all(&chunk)
|
||||
.map_err(|error| super::internal(format!("failed to write download: {error}")))?;
|
||||
let percent = download_percent(downloaded, expected_size);
|
||||
if percent > last_percent {
|
||||
last_percent = percent;
|
||||
events.emit_progress(operation_id, "downloading release archive", percent);
|
||||
if percent >= next_progress || downloaded == expected_size {
|
||||
events.emit_progress(
|
||||
operation_id,
|
||||
&format!(
|
||||
"已下载 {} / {} ({percent}%)",
|
||||
format_bytes(downloaded),
|
||||
format_bytes(expected_size)
|
||||
),
|
||||
percent,
|
||||
);
|
||||
next_progress = percent.saturating_add(10);
|
||||
}
|
||||
}
|
||||
file.sync_all()
|
||||
@@ -155,3 +163,15 @@ fn download_percent(downloaded: u64, total: u64) -> u8 {
|
||||
)
|
||||
.unwrap_or(100)
|
||||
}
|
||||
|
||||
pub(crate) fn format_bytes(bytes: u64) -> String {
|
||||
const KIB: u64 = 1024;
|
||||
const MIB: u64 = KIB * 1024;
|
||||
if bytes >= MIB {
|
||||
return format!("{:.1} MiB", bytes as f64 / MIB as f64);
|
||||
}
|
||||
if bytes >= KIB {
|
||||
return format!("{:.1} KiB", bytes as f64 / KIB as f64);
|
||||
}
|
||||
format!("{bytes} B")
|
||||
}
|
||||
|
||||
@@ -33,6 +33,7 @@ use self::download::{MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, verify
|
||||
use self::process::{environment_value, spawn_replacement};
|
||||
use self::release::{
|
||||
ReleaseCache, confirmed_target, detail_from_release, fetch_latest, select_archive,
|
||||
versions_cross_major,
|
||||
};
|
||||
use self::state::{
|
||||
OperationFileLock, UpdateTempDir, finish, operation_id, read_status, set_running,
|
||||
@@ -226,14 +227,26 @@ impl ProcessSystemOperations {
|
||||
return Err(conflict(reason));
|
||||
}
|
||||
let target = confirmed_target(target_version)?;
|
||||
if versions_cross_major(&self.config.version, &target).unwrap_or(false) {
|
||||
let reason = "跨大版本升级不提供自动更新,请按发布说明手动迁移";
|
||||
self.events.error_terminal(None, Some("preflight"), reason);
|
||||
return Err(conflict(reason));
|
||||
}
|
||||
let repository = self
|
||||
.config
|
||||
.update_repository
|
||||
.as_deref()
|
||||
.ok_or_else(|| conflict("update repository is not configured"))?;
|
||||
self.events
|
||||
.info(None, Some("release"), "fetching latest release");
|
||||
let release = fetch_latest(&self.config.github_api_base, repository).await?;
|
||||
.info(None, Some("release"), "正在获取最新 Release 信息");
|
||||
let release = match fetch_latest(&self.config.github_api_base, repository).await {
|
||||
Ok(release) => release,
|
||||
Err(error) => {
|
||||
self.events
|
||||
.error_terminal(None, Some("release"), error.to_string());
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let detail = detail_from_release(&self.config, &release);
|
||||
if detail.latest_version != target {
|
||||
self.events.warning_terminal(
|
||||
@@ -258,12 +271,17 @@ impl ProcessSystemOperations {
|
||||
Some(&target),
|
||||
&self.config.version,
|
||||
)?;
|
||||
self.events.info(
|
||||
Some(&operation_id),
|
||||
Some("prepare"),
|
||||
format!("准备更新到 v{target}"),
|
||||
);
|
||||
let result = self.install_release(&release, &target, &operation_id).await;
|
||||
match &result {
|
||||
Ok(()) => self.events.success_terminal(
|
||||
Some(&operation_id),
|
||||
Some("done"),
|
||||
"release files replaced",
|
||||
"更新文件已替换,等待服务重启生效",
|
||||
),
|
||||
Err(error) => {
|
||||
self.events
|
||||
@@ -294,7 +312,23 @@ impl ProcessSystemOperations {
|
||||
version: &str,
|
||||
operation_id: &str,
|
||||
) -> Result<(), OperationError> {
|
||||
self.events.info(
|
||||
Some(operation_id),
|
||||
Some("asset"),
|
||||
"正在选择匹配当前平台的更新包",
|
||||
);
|
||||
let archive = select_archive(release, version)?;
|
||||
self.events.info(
|
||||
Some(operation_id),
|
||||
Some("asset"),
|
||||
format!(
|
||||
"已选择更新包 {} ({})",
|
||||
archive.name,
|
||||
download::format_bytes(archive.size)
|
||||
),
|
||||
);
|
||||
self.events
|
||||
.info(Some(operation_id), Some("verify"), "正在校验更新资源");
|
||||
if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE {
|
||||
return Err(invalid("release archive size is invalid"));
|
||||
}
|
||||
@@ -306,12 +340,16 @@ impl ProcessSystemOperations {
|
||||
if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE {
|
||||
return Err(invalid("release checksum size is invalid"));
|
||||
}
|
||||
self.events
|
||||
.info(Some(operation_id), Some("prepare"), "正在创建临时更新目录");
|
||||
fs::create_dir_all(&self.config.update_temp_dir)
|
||||
.map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?;
|
||||
let temp_root = fs::canonicalize(&self.config.update_temp_dir)
|
||||
.map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?;
|
||||
let temp = UpdateTempDir::create(&temp_root)?;
|
||||
let archive_path = temp.path().join(&archive.name);
|
||||
self.events
|
||||
.info(Some(operation_id), Some("download"), "开始下载更新包");
|
||||
download_file(
|
||||
&archive.browser_download_url,
|
||||
&archive_path,
|
||||
@@ -321,6 +359,10 @@ impl ProcessSystemOperations {
|
||||
&self.events,
|
||||
)
|
||||
.await?;
|
||||
self.events
|
||||
.success(Some(operation_id), Some("download"), "更新包下载完成");
|
||||
self.events
|
||||
.info(Some(operation_id), Some("checksum"), "正在校验 checksum");
|
||||
verify_checksum(
|
||||
&archive_path,
|
||||
&archive.name,
|
||||
@@ -329,12 +371,23 @@ impl ProcessSystemOperations {
|
||||
&self.config.github_api_base,
|
||||
)
|
||||
.await?;
|
||||
self.events
|
||||
.success(Some(operation_id), Some("checksum"), "checksum 校验通过");
|
||||
self.events
|
||||
.info(Some(operation_id), Some("extract"), "正在解压更新包");
|
||||
let extracted = extract_release(&archive_path, temp.path())?;
|
||||
self.events
|
||||
.success(Some(operation_id), Some("extract"), "更新包解压完成");
|
||||
self.events
|
||||
.info(Some(operation_id), Some("replace"), "正在替换应用文件");
|
||||
replace_release_files(
|
||||
&self.config.executable_path()?,
|
||||
&self.config.web_dist_dir,
|
||||
extracted,
|
||||
)
|
||||
)?;
|
||||
self.events
|
||||
.success(Some(operation_id), Some("replace"), "应用文件替换完成");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -515,6 +568,17 @@ impl UpdateEvents {
|
||||
);
|
||||
}
|
||||
|
||||
fn success(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into<String>) {
|
||||
self.emit(
|
||||
SystemUpdateEventLevel::Success,
|
||||
operation_id,
|
||||
step,
|
||||
message,
|
||||
None,
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
fn warning_terminal(
|
||||
&self,
|
||||
operation_id: Option<&str>,
|
||||
|
||||
@@ -121,12 +121,13 @@ pub(crate) fn detail_from_release(
|
||||
release: &GitHubRelease,
|
||||
) -> SystemUpdateDetail {
|
||||
let latest = normalize_version(&release.tag_name);
|
||||
let available = (config.update_channel != "stable" || !release.prerelease)
|
||||
let newer = (config.update_channel != "stable" || !release.prerelease)
|
||||
&& version_is_newer(&config.version, &latest).unwrap_or(false);
|
||||
let cross_major = versions_cross_major(&config.version, &latest).unwrap_or(false);
|
||||
SystemUpdateDetail {
|
||||
current_version: config.version.clone(),
|
||||
latest_version: latest,
|
||||
has_update: available,
|
||||
latest_version: latest.clone(),
|
||||
has_update: newer && !cross_major,
|
||||
deployment_mode: config.deployment_mode.clone(),
|
||||
build_type: config.build_type.clone(),
|
||||
release_url: release.html_url.clone(),
|
||||
@@ -134,7 +135,9 @@ pub(crate) fn detail_from_release(
|
||||
cached: false,
|
||||
update_supported: config.update_support_error().is_none(),
|
||||
unsupported_reason: config.update_support_error(),
|
||||
warning: None,
|
||||
warning: (newer && cross_major).then(|| {
|
||||
format!("检测到新的大版本 v{latest}:跨大版本升级不提供自动更新,请按发布说明手动迁移")
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -265,6 +268,13 @@ fn version_is_newer(current: &str, latest: &str) -> Option<bool> {
|
||||
Some(latest > current)
|
||||
}
|
||||
|
||||
/// 自动更新只在同 major 内提供;跨大版本的配置与数据迁移契约不保证兼容。
|
||||
pub(crate) fn versions_cross_major(current: &str, latest: &str) -> Option<bool> {
|
||||
let current = semver::Version::parse(&normalize_version(current)).ok()?;
|
||||
let latest = semver::Version::parse(&normalize_version(latest)).ok()?;
|
||||
Some(current.major != latest.major)
|
||||
}
|
||||
|
||||
fn repository_character(character: char) -> bool {
|
||||
character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.')
|
||||
}
|
||||
|
||||
@@ -108,22 +108,19 @@ impl WorkerSupervisor {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 先协作取消;超时后 abort 不合作任务,lease 依赖 TTL 兜底。
|
||||
/// 先协作取消;所有任务共享同一绝对截止点,逾期任务被 abort,
|
||||
/// lease 依赖 TTL 兜底。每个句柄恰好 join 一次:已在截止点前退出的
|
||||
/// 不会被二次 poll,逾期的在 abort 后等待其真正终止。
|
||||
pub async fn shutdown(&self, timeout: Duration) {
|
||||
self.cancellation.cancel();
|
||||
let mut handles = std::mem::take(&mut *lock_unpoisoned(&self.handles));
|
||||
let joined = tokio::time::timeout(timeout, async {
|
||||
for (_, handle) in &mut handles {
|
||||
let _ = handle.await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.is_ok();
|
||||
if !joined {
|
||||
for (_, handle) in &handles {
|
||||
let handles = std::mem::take(&mut *lock_unpoisoned(&self.handles));
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
for (_, mut handle) in handles {
|
||||
if tokio::time::timeout_at(deadline, &mut handle)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
handle.abort();
|
||||
}
|
||||
for (_, handle) in handles {
|
||||
let _ = handle.await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,7 +18,8 @@ use tar::{Builder, EntryType, Header};
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
const TARGET_VERSION: &str = "9.9.9";
|
||||
const TARGET_VERSION: &str = "1.9.9";
|
||||
const CROSS_MAJOR_VERSION: &str = "2.0.0";
|
||||
|
||||
#[tokio::test]
|
||||
async fn restart_should_not_shutdown_when_replacement_spawn_fails() {
|
||||
@@ -174,6 +175,43 @@ async fn update_detail_should_use_cached_release_when_not_refreshed() {
|
||||
assert!(service.update_detail(false).await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_detail_should_withhold_cross_major_release() {
|
||||
let server = MockServer::start().await;
|
||||
let fixture = Fixture::new();
|
||||
fixture
|
||||
.mount_release_once(&server, CROSS_MAJOR_VERSION)
|
||||
.await;
|
||||
|
||||
let detail = fixture
|
||||
.service(&server)
|
||||
.update_detail(true)
|
||||
.await
|
||||
.expect("detail");
|
||||
assert_eq!(detail.latest_version, CROSS_MAJOR_VERSION);
|
||||
assert!(!detail.has_update);
|
||||
assert!(detail.warning.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_should_reject_cross_major_target_before_fetching_release() {
|
||||
let fixture = Fixture::new();
|
||||
let service = ProcessSystemOperations::new(
|
||||
CancellationToken::new(),
|
||||
fixture.config("http://127.0.0.1:1/repos"),
|
||||
);
|
||||
|
||||
let error = service
|
||||
.perform_update(Some(CROSS_MAJOR_VERSION.to_owned()))
|
||||
.await
|
||||
.expect_err("cross-major target must be rejected");
|
||||
assert_eq!(error.kind(), SystemOperationErrorKind::Conflict);
|
||||
assert_eq!(
|
||||
fs::read(fixture.executable()).expect("binary"),
|
||||
b"old-binary"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_events_should_close_after_terminal_update_log() {
|
||||
let fixture = Fixture::new();
|
||||
@@ -204,6 +242,68 @@ async fn update_events_should_open_authenticated_sse_stream() {
|
||||
assert!(stream.next().await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_events_should_preserve_complete_release_stage_sequence() {
|
||||
let server = MockServer::start().await;
|
||||
let fixture = Fixture::new();
|
||||
fixture
|
||||
.mount_release(
|
||||
&server,
|
||||
TARGET_VERSION,
|
||||
ArchiveKind::Safe,
|
||||
ChecksumKind::Valid,
|
||||
)
|
||||
.await;
|
||||
let service = fixture.service(&server);
|
||||
let events = service.update_events();
|
||||
service
|
||||
.perform_update(Some(TARGET_VERSION.to_owned()))
|
||||
.await
|
||||
.expect("update");
|
||||
let steps = events
|
||||
.map(|event| event.step.expect("release update step"))
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
steps,
|
||||
[
|
||||
"release", "prepare", "asset", "asset", "verify", "prepare", "download", "download",
|
||||
"download", "checksum", "checksum", "extract", "extract", "replace", "replace", "done",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_events_should_report_download_bytes_and_percent() {
|
||||
let server = MockServer::start().await;
|
||||
let fixture = Fixture::new();
|
||||
fixture
|
||||
.mount_release(
|
||||
&server,
|
||||
TARGET_VERSION,
|
||||
ArchiveKind::Safe,
|
||||
ChecksumKind::Valid,
|
||||
)
|
||||
.await;
|
||||
let service = fixture.service(&server);
|
||||
let events = service.update_events();
|
||||
service
|
||||
.perform_update(Some(TARGET_VERSION.to_owned()))
|
||||
.await
|
||||
.expect("update");
|
||||
let progress = events
|
||||
.filter_map(
|
||||
|event| async move { event.progress_percent.map(|value| (value, event.message)) },
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
assert!(progress.iter().any(|(percent, message)| {
|
||||
*percent == 100 && message.starts_with("已下载 ") && message.ends_with(" (100%)")
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_should_fail_when_release_checksum_is_missing() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -335,7 +435,7 @@ async fn update_should_reject_when_confirmed_target_differs_from_remote_latest()
|
||||
let server = MockServer::start().await;
|
||||
let fixture = Fixture::new();
|
||||
fixture
|
||||
.mount_release(&server, "9.9.8", ArchiveKind::Safe, ChecksumKind::Valid)
|
||||
.mount_release(&server, "1.9.8", ArchiveKind::Safe, ChecksumKind::Valid)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
|
||||
@@ -626,6 +626,55 @@ async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_timeout_must_not_rejoin_workers_that_already_exited() {
|
||||
// 挂死任务放在 join 顺序最末的 kind(Retention),保证超时触发时
|
||||
// 排在它前面的合作任务句柄都已被 join 过一次。
|
||||
let hung = HungTask {
|
||||
entered: Arc::new(AtomicBool::new(false)),
|
||||
notification: Arc::new(Notify::new()),
|
||||
};
|
||||
let cooperative = CancellationTask {
|
||||
entered: Arc::new(AtomicUsize::new(0)),
|
||||
observed: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let mut plan = ACTIVE_KINDS
|
||||
.into_iter()
|
||||
.map(|kind| {
|
||||
if kind == WorkerKind::Retention {
|
||||
registration(kind, kind.as_str(), hung.clone(), test_schedule(), true)
|
||||
} else {
|
||||
registration(
|
||||
kind,
|
||||
kind.as_str(),
|
||||
cooperative.clone(),
|
||||
test_schedule(),
|
||||
true,
|
||||
)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
plan.push(WorkerContribution::Disabled {
|
||||
kind: WorkerKind::OpsFlush,
|
||||
reason: WorkerDisabledReason::NoBufferedOpsEvents,
|
||||
});
|
||||
let leases = Arc::new(FakeLeasePort::default());
|
||||
let supervisor = supervisor();
|
||||
supervisor
|
||||
.start(plan, leases.clone())
|
||||
.expect("complete plan");
|
||||
hung.notification.notified().await;
|
||||
wait_until(|| cooperative.entered.load(Ordering::SeqCst) == ACTIVE_KINDS.len() - 1).await;
|
||||
|
||||
supervisor.shutdown(Duration::from_millis(100)).await;
|
||||
|
||||
assert_eq!(
|
||||
cooperative.observed.load(Ordering::SeqCst),
|
||||
ACTIVE_KINDS.len() - 1
|
||||
);
|
||||
assert!(lock_unpoisoned(&leases.active_guards).is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn running_worker_becomes_unhealthy_when_last_success_is_stale() {
|
||||
let task = SucceedsThenHangsTask {
|
||||
|
||||
@@ -72,26 +72,50 @@ pub const DONE_SSE_FRAME: &str = "data: [DONE]\n\n";
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SseEventDecoder {
|
||||
pending: Vec<u8>,
|
||||
/// `pending` 中已确认不含帧分隔符的前缀长度。新数据到达时从该边界
|
||||
/// 回退分隔符最大跨界宽度继续扫描,大帧跨多个 chunk 不会被反复重扫。
|
||||
scanned: usize,
|
||||
}
|
||||
|
||||
impl SseEventDecoder {
|
||||
/// 返回从 `consumed` 起的下一个完整帧结束位置(含分隔符),并推进扫描边界。
|
||||
fn next_frame_end(&mut self, consumed: usize) -> Option<usize> {
|
||||
// `\r\n\r\n` 分隔符可能跨 chunk 边界,扫描起点回退 3 字节。
|
||||
let from = self.scanned.saturating_sub(3).max(consumed);
|
||||
match sse_frame_separator_bytes(&self.pending[from..]) {
|
||||
Some((position, separator_len)) => {
|
||||
let end = from + position + separator_len;
|
||||
self.scanned = end;
|
||||
Some(end)
|
||||
}
|
||||
None => {
|
||||
self.scanned = self.pending.len();
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn commit_consumed(&mut self, consumed: usize) {
|
||||
if consumed != 0 {
|
||||
self.pending.drain(..consumed);
|
||||
self.scanned = self.scanned.saturating_sub(consumed);
|
||||
}
|
||||
}
|
||||
|
||||
/// 追加一个字节块并返回其中已经完整的事件。
|
||||
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<SseEvent>, SseError> {
|
||||
self.pending.extend_from_slice(chunk);
|
||||
let mut events = Vec::new();
|
||||
let mut consumed = 0usize;
|
||||
|
||||
while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) {
|
||||
let end = consumed.saturating_add(frame_len);
|
||||
while let Some(end) = self.next_frame_end(consumed) {
|
||||
let frame = std::str::from_utf8(&self.pending[consumed..end])
|
||||
.map_err(|error| SseError::ParseError(error.to_string()))?;
|
||||
events.extend(parse_sse_events(frame)?);
|
||||
consumed = end;
|
||||
}
|
||||
|
||||
if consumed != 0 {
|
||||
self.pending.drain(..consumed);
|
||||
}
|
||||
self.commit_consumed(consumed);
|
||||
if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES {
|
||||
return Err(SseError::BufferExceeded {
|
||||
max_bytes: MAX_SSE_EVENT_BUFFER_BYTES,
|
||||
@@ -110,8 +134,7 @@ impl SseEventDecoder {
|
||||
let mut frames = Vec::new();
|
||||
let mut consumed = 0usize;
|
||||
|
||||
while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) {
|
||||
let end = consumed.saturating_add(frame_len);
|
||||
while let Some(end) = self.next_frame_end(consumed) {
|
||||
let raw = self.pending[consumed..end].to_vec();
|
||||
let events = std::str::from_utf8(&raw)
|
||||
.ok()
|
||||
@@ -121,9 +144,7 @@ impl SseEventDecoder {
|
||||
consumed = end;
|
||||
}
|
||||
|
||||
if consumed != 0 {
|
||||
self.pending.drain(..consumed);
|
||||
}
|
||||
self.commit_consumed(consumed);
|
||||
if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES {
|
||||
return Err(SseError::BufferExceeded {
|
||||
max_bytes: MAX_SSE_EVENT_BUFFER_BYTES,
|
||||
@@ -134,6 +155,7 @@ impl SseEventDecoder {
|
||||
|
||||
/// 流结束时解析尚未带空行分隔符的最后一帧。
|
||||
pub fn finish(&mut self) -> Result<Vec<SseEvent>, SseError> {
|
||||
self.scanned = 0;
|
||||
if self.pending.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -147,6 +169,7 @@ impl SseEventDecoder {
|
||||
///
|
||||
/// 与 [`Self::push_frames`] 一致,旁路解析失败不会丢弃原始帧。
|
||||
pub fn finish_frames(&mut self) -> Result<Vec<SseFrame>, SseError> {
|
||||
self.scanned = 0;
|
||||
if self.pending.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -235,9 +258,6 @@ pub enum SseError {
|
||||
/// 解析错误。
|
||||
#[error("SSE parse error: {0}")]
|
||||
ParseError(String),
|
||||
/// 用量提取错误。
|
||||
#[error("Usage extraction error: {0}")]
|
||||
UsageExtraction(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
|
||||
@@ -107,6 +107,29 @@ fn incremental_decoder_should_decode_across_arbitrary_chunk_boundaries() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_decoder_should_find_separators_split_at_every_byte_boundary() {
|
||||
// CRLF 分隔符按单字节到达时必须仍被识别;帧内容跨任意多次 push 不丢失。
|
||||
let body = b"event: one\r\ndata: first\r\n\r\nevent: two\ndata: second\n\ndata: third\r\n\r\n";
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
let mut events = Vec::new();
|
||||
for byte in body {
|
||||
events.extend(decoder.push(&[*byte]).expect("single-byte chunk"));
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
events
|
||||
.iter()
|
||||
.map(|event| (event.event.as_deref(), event.data.as_str()))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
(Some("one"), "first"),
|
||||
(Some("two"), "second"),
|
||||
(None, "third"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_frame_decoder_should_preserve_exact_frame_across_chunk_boundaries() {
|
||||
let raw = b": keep-alive\r\nid: evt_1\r\nevent: response.future\r\nretry: 250\r\ndata: { \"type\": \"response.future\", \"opaque\": true }\r\n\r\n";
|
||||
|
||||
@@ -86,6 +86,37 @@ pub type StoreResult<T> = Result<T, StoreError>;
|
||||
pub struct StoreConfig {
|
||||
database: StoreConnectionConfig,
|
||||
redis: StoreConnectionConfig,
|
||||
#[serde(default)]
|
||||
pool: StorePoolConfig,
|
||||
}
|
||||
|
||||
/// PostgreSQL 连接池预算;acquire 超时决定池耗尽时快速失败而非排队积压。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
|
||||
#[serde(deny_unknown_fields, default)]
|
||||
pub struct StorePoolConfig {
|
||||
pub max_connections: u32,
|
||||
pub acquire_timeout_seconds: u64,
|
||||
}
|
||||
|
||||
impl Default for StorePoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_connections: 10,
|
||||
acquire_timeout_seconds: 5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StorePoolConfig {
|
||||
fn validate(&self) -> StoreResult<()> {
|
||||
if self.max_connections == 0 || self.acquire_timeout_seconds == 0 {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "store config",
|
||||
message: "pool limits must be positive".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl StoreConfig {
|
||||
@@ -104,6 +135,7 @@ impl StoreConfig {
|
||||
}
|
||||
self.database.validate("database")?;
|
||||
self.redis.validate("redis")?;
|
||||
self.pool.validate()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -227,7 +259,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
const REDIS_NAMESPACE: &str = "codex-proxy-rs";
|
||||
|
||||
config.resolve_and_validate(std::path::Path::new("."))?;
|
||||
let pool = postgres::connect_and_migrate(&config.database_url()?).await?;
|
||||
let pool = postgres::connect_and_migrate(&config.database_url()?, config.pool).await?;
|
||||
let redis_client = ::redis::Client::open(config.redis_url()?)
|
||||
.map_err(|_| redis_unavailable("create Redis client"))?;
|
||||
let redis_connection = redis_client
|
||||
@@ -823,44 +855,6 @@ impl AuthStore for AdminAuthStoreAdapter {
|
||||
.map_err(|error| admin_store_error("admin session", error))
|
||||
}
|
||||
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
redis::AdminAuthStateRepository::login_source_is_throttled(
|
||||
&self.state,
|
||||
source,
|
||||
failure_limit,
|
||||
window_seconds,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| admin_store_error("admin login throttle", error))
|
||||
}
|
||||
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> AdminStoreResult<bool> {
|
||||
redis::AdminAuthStateRepository::record_login_failure(
|
||||
&self.state,
|
||||
source,
|
||||
failure_limit,
|
||||
window_seconds,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| admin_store_error("admin login throttle", error))
|
||||
}
|
||||
|
||||
async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()> {
|
||||
redis::AdminAuthStateRepository::clear_login_failures(&self.state, source)
|
||||
.await
|
||||
.map_err(|error| admin_store_error("admin login throttle", error))
|
||||
}
|
||||
|
||||
async fn append_audit_event(&self, event: AdminAuditModel) -> AdminStoreResult<()> {
|
||||
let config_revision = event
|
||||
.config_revision
|
||||
|
||||
@@ -310,17 +310,12 @@ impl UpdateClientApiKeyDetails {
|
||||
|
||||
#[async_trait]
|
||||
pub trait ClientApiKeyRepository: Send + Sync {
|
||||
async fn authenticate_client_api_key(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> StoreResult<Option<ClientApiKeySecret>>;
|
||||
async fn list_client_api_keys(
|
||||
&self,
|
||||
query: ClientApiKeyListQuery,
|
||||
) -> StoreResult<ClientApiKeyPage>;
|
||||
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>>;
|
||||
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>>;
|
||||
async fn insert_client_api_key(&self, key: NewClientApiKey) -> StoreResult<()>;
|
||||
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool>;
|
||||
async fn touch_client_api_keys(
|
||||
&self,
|
||||
@@ -342,22 +337,6 @@ impl PgClientApiKeyRepository {
|
||||
|
||||
#[async_trait]
|
||||
impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
async fn authenticate_client_api_key(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> StoreResult<Option<ClientApiKeySecret>> {
|
||||
validate_key(key)?;
|
||||
let row = sqlx::query_as::<_, (String, String, String, bool, i64, i64)>(
|
||||
"select id, key, provider_kind, enabled, max_concurrency, requests_per_minute
|
||||
from client_api_keys where key = $1",
|
||||
)
|
||||
.bind(key)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("authenticate client API key"))?;
|
||||
row.map(client_secret_from_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_client_api_keys(
|
||||
&self,
|
||||
query: ClientApiKeyListQuery,
|
||||
@@ -432,27 +411,6 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn insert_client_api_key(&self, key: NewClientApiKey) -> StoreResult<()> {
|
||||
key.validate()?;
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
last_used_at, created_at, updated_at
|
||||
) values ($1, $2, $3, $4, $5, true, $6, $7, null, now(), now())",
|
||||
)
|
||||
.bind(key.id)
|
||||
.bind(key.name)
|
||||
.bind(key.label)
|
||||
.bind(key.provider_kind)
|
||||
.bind(key.key)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "id", &key.id)?;
|
||||
require_nonempty(ENTITY, "name", &key.name)?;
|
||||
|
||||
@@ -10,7 +10,7 @@ use gateway_core::accounting::{CostSource as CoreCostSource, Usage as CoreUsage}
|
||||
use gateway_core::engine::{
|
||||
AttemptRecord as CoreAttemptRecord, ExecutionStore, IntermediateFailure,
|
||||
ModelRequestFinalization as CoreModelRequestFinalization, ModelRequestId,
|
||||
NewModelRequest as CoreNewModelRequest, RecoveryReport as CoreRecoveryReport,
|
||||
NewModelRequest as CoreNewModelRequest, ProbeFailure, RecoveryReport as CoreRecoveryReport,
|
||||
UpstreamSendState as CoreUpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
@@ -295,6 +295,12 @@ pub struct ModelRequestRecoveryReport {
|
||||
#[async_trait]
|
||||
pub trait ModelRequestRepository: Send + Sync {
|
||||
async fn insert_model_request(&self, request: NewModelRequest) -> StoreResult<()>;
|
||||
/// 请求行与首次 attempt 列一次插入;`attempt.attempt_count` 必须为 1。
|
||||
async fn insert_model_request_with_first_attempt(
|
||||
&self,
|
||||
request: NewModelRequest,
|
||||
attempt: ModelRequestAttemptStart,
|
||||
) -> StoreResult<()>;
|
||||
async fn begin_model_request_attempt(
|
||||
&self,
|
||||
attempt: ModelRequestAttemptStart,
|
||||
@@ -383,11 +389,71 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_model_request_with_first_attempt(
|
||||
&self,
|
||||
request: NewModelRequest,
|
||||
attempt: ModelRequestAttemptStart,
|
||||
) -> StoreResult<()> {
|
||||
request.validate()?;
|
||||
attempt.validate()?;
|
||||
if attempt.attempt_count != 1 || attempt.model_request_id != request.id {
|
||||
return Err(invalid("first attempt must target the inserted request"));
|
||||
}
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
|
||||
operation, endpoint, client_transport, requested_model_id,
|
||||
client_ip, user_agent, reasoning_effort,
|
||||
reasoning_preset, request_kind, subagent_kind, compact,
|
||||
image_generation_requested, started_at, deadline_at,
|
||||
provider_kind, provider_account_id, provider_account_ref,
|
||||
upstream_model_id, upstream_transport, http_version,
|
||||
attempt_count, upstream_send_state
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
|
||||
$13, $14, $15, $16, $17, $18, $19,
|
||||
$20, $21, $22, $23, $24, $25, 1, 'not_sent'
|
||||
)",
|
||||
)
|
||||
.bind(request.id)
|
||||
.bind(request.client_api_key_id)
|
||||
.bind(request.client_api_key_ref)
|
||||
.bind(to_i64(request.config_revision, "config_revision")?)
|
||||
.bind(request.protocol)
|
||||
.bind(request.operation)
|
||||
.bind(request.endpoint)
|
||||
.bind(request.client_transport)
|
||||
.bind(request.requested_model_id)
|
||||
.bind(request.client_ip)
|
||||
.bind(request.user_agent)
|
||||
.bind(request.reasoning_effort)
|
||||
.bind(request.reasoning_preset)
|
||||
.bind(request.request_kind)
|
||||
.bind(request.subagent_kind)
|
||||
.bind(request.compact)
|
||||
.bind(request.image_generation_requested)
|
||||
.bind(request.started_at)
|
||||
.bind(request.deadline_at)
|
||||
.bind(attempt.provider_kind)
|
||||
.bind(attempt.provider_account_id)
|
||||
.bind(attempt.provider_account_ref)
|
||||
.bind(attempt.upstream_model_id)
|
||||
.bind(attempt.upstream_transport)
|
||||
.bind(attempt.http_version)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert model request with first attempt"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn begin_model_request_attempt(
|
||||
&self,
|
||||
attempt: ModelRequestAttemptStart,
|
||||
) -> StoreResult<u32> {
|
||||
attempt.validate()?;
|
||||
// upstream_send_state 是请求级单调水位(sent > ambiguous > not_sent),
|
||||
// 由 mark/finalize 抬升;开启新 attempt 不得把已持久化的 sent 重置回
|
||||
// not_sent——崩溃恢复的终态写回会原样继承本列。
|
||||
let count = sqlx::query_scalar::<_, i32>(
|
||||
"update model_requests
|
||||
set provider_kind = $2,
|
||||
@@ -396,8 +462,7 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
upstream_model_id = $5,
|
||||
upstream_transport = $6,
|
||||
http_version = $7,
|
||||
attempt_count = $8,
|
||||
upstream_send_state = 'not_sent'
|
||||
attempt_count = $8
|
||||
where id = $1 and outcome = 'running' and downstream_committed_at is null
|
||||
and $8 = attempt_count + 1
|
||||
returning attempt_count",
|
||||
@@ -632,53 +697,15 @@ impl ExecutionStore for PgExecutionStore {
|
||||
&self,
|
||||
request: CoreNewModelRequest,
|
||||
) -> Result<(), CoreStoreError> {
|
||||
self.insert_model_request(NewModelRequest {
|
||||
id: request.id.as_str().to_owned(),
|
||||
client_api_key_id: request
|
||||
.client_api_key_id
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
client_api_key_ref: request.client_api_key_ref.as_str().to_owned(),
|
||||
config_revision: request.config_revision.get(),
|
||||
protocol: request.protocol,
|
||||
operation: request.operation.as_str().to_owned(),
|
||||
endpoint: request.endpoint,
|
||||
client_transport: request.client_transport,
|
||||
requested_model_id: request.requested_model.as_str().to_owned(),
|
||||
client_ip: request.client_ip.map(|address| address.to_string()),
|
||||
user_agent: request.user_agent,
|
||||
reasoning_effort: request.reasoning_effort,
|
||||
reasoning_preset: request.reasoning_preset,
|
||||
request_kind: request.request_kind,
|
||||
subagent_kind: request.subagent_kind,
|
||||
compact: request.compact,
|
||||
image_generation_requested: request.image_generation_requested,
|
||||
started_at: DateTime::<Utc>::from(request.started_at),
|
||||
deadline_at: DateTime::<Utc>::from(request.deadline_at),
|
||||
})
|
||||
.await
|
||||
.map_err(core_store_error)
|
||||
self.insert_model_request(new_model_request_row(request))
|
||||
.await
|
||||
.map_err(core_store_error)
|
||||
}
|
||||
|
||||
async fn record_attempt(&self, attempt: CoreAttemptRecord) -> Result<(), CoreStoreError> {
|
||||
let expected_count = attempt.attempt_count.get();
|
||||
let persisted = self
|
||||
.begin_model_request_attempt(ModelRequestAttemptStart {
|
||||
model_request_id: attempt.request_id.as_str().to_owned(),
|
||||
attempt_count: expected_count,
|
||||
provider_kind: attempt.provider_kind.as_str().to_owned(),
|
||||
provider_account_id: attempt
|
||||
.provider_account_id
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
provider_account_ref: attempt
|
||||
.provider_account_ref
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
upstream_model_id: attempt.upstream_model_id.as_str().to_owned(),
|
||||
upstream_transport: attempt.upstream_transport,
|
||||
http_version: attempt.http_version,
|
||||
})
|
||||
.begin_model_request_attempt(attempt_start_row(attempt))
|
||||
.await
|
||||
.map_err(core_store_error)?;
|
||||
if persisted == expected_count {
|
||||
@@ -688,6 +715,24 @@ impl ExecutionStore for PgExecutionStore {
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_model_request_with_attempt(
|
||||
&self,
|
||||
request: CoreNewModelRequest,
|
||||
attempt: CoreAttemptRecord,
|
||||
) -> Result<(), CoreStoreError> {
|
||||
// 合并写只对首次 attempt 成立(插入即带 attempt 列);其余走基础两步。
|
||||
if attempt.attempt_count.get() != 1 {
|
||||
self.create_model_request(request).await?;
|
||||
return self.record_attempt(attempt).await;
|
||||
}
|
||||
self.insert_model_request_with_first_attempt(
|
||||
new_model_request_row(request),
|
||||
attempt_start_row(attempt),
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)
|
||||
}
|
||||
|
||||
async fn mark_send_state(
|
||||
&self,
|
||||
request_id: &ModelRequestId,
|
||||
@@ -776,6 +821,56 @@ impl ExecutionStore for PgExecutionStore {
|
||||
.map_err(core_store_error)
|
||||
}
|
||||
|
||||
async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), CoreStoreError> {
|
||||
let error = failure.error;
|
||||
let retry_after_ms = error
|
||||
.retry_after()
|
||||
.map(|duration| u64::try_from(duration.as_millis()))
|
||||
.transpose()
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
|
||||
let latency_ms = u64::try_from(failure.latency.as_millis())
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
|
||||
let provider_error_code = error.upstream_code().map(|code| code.as_str().to_owned());
|
||||
tracing::warn!(
|
||||
target: "gateway_probe",
|
||||
provider_kind = failure.provider_kind.as_str(),
|
||||
account_id = failure.account_id.as_str(),
|
||||
upstream_model = failure.upstream_model_id.as_str(),
|
||||
failure_kind = error.kind().as_str(),
|
||||
send_state = ?error.send_state(),
|
||||
upstream_status = ?error.upstream_status(),
|
||||
provider_error_code = ?provider_error_code,
|
||||
latency_ms,
|
||||
"账号连接测试失败"
|
||||
);
|
||||
super::OpsEventRepository::append_ops_event(
|
||||
&super::PgOpsEventRepository::new(self.pool.clone()),
|
||||
super::OpsEvent {
|
||||
id: Uuid::now_v7().to_string(),
|
||||
model_request_id: None,
|
||||
attempt_index: None,
|
||||
level: super::OpsEventLevel::Warning,
|
||||
component: "account_probe".to_owned(),
|
||||
operation: "connection_test".to_owned(),
|
||||
provider_kind: Some(failure.provider_kind.as_str().to_owned()),
|
||||
provider_account_id: Some(failure.account_id.as_str().to_owned()),
|
||||
provider_account_ref: Some(failure.account_id.as_str().to_owned()),
|
||||
upstream_model_id: Some(failure.upstream_model_id.as_str().to_owned()),
|
||||
failure_kind: error.kind().as_str().to_owned(),
|
||||
status_code: error.upstream_status(),
|
||||
provider_error_code,
|
||||
retry_after_ms,
|
||||
upstream_request_id: error.upstream_request_id().map(|id| id.as_str().to_owned()),
|
||||
latency_ms: Some(latency_ms),
|
||||
message: "account connection test failed".to_owned(),
|
||||
occurrence_count: 1,
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)
|
||||
}
|
||||
|
||||
async fn finalize_model_request(
|
||||
&self,
|
||||
finalization: CoreModelRequestFinalization,
|
||||
@@ -937,6 +1032,52 @@ fn core_store_error(error: StoreError) -> CoreStoreError {
|
||||
CoreStoreError::new(kind)
|
||||
}
|
||||
|
||||
fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
|
||||
NewModelRequest {
|
||||
id: request.id.as_str().to_owned(),
|
||||
client_api_key_id: request
|
||||
.client_api_key_id
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
client_api_key_ref: request.client_api_key_ref.as_str().to_owned(),
|
||||
config_revision: request.config_revision.get(),
|
||||
protocol: request.protocol,
|
||||
operation: request.operation.as_str().to_owned(),
|
||||
endpoint: request.endpoint,
|
||||
client_transport: request.client_transport,
|
||||
requested_model_id: request.requested_model.as_str().to_owned(),
|
||||
client_ip: request.client_ip.map(|address| address.to_string()),
|
||||
user_agent: request.user_agent,
|
||||
reasoning_effort: request.reasoning_effort,
|
||||
reasoning_preset: request.reasoning_preset,
|
||||
request_kind: request.request_kind,
|
||||
subagent_kind: request.subagent_kind,
|
||||
compact: request.compact,
|
||||
image_generation_requested: request.image_generation_requested,
|
||||
started_at: DateTime::<Utc>::from(request.started_at),
|
||||
deadline_at: DateTime::<Utc>::from(request.deadline_at),
|
||||
}
|
||||
}
|
||||
|
||||
fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
|
||||
ModelRequestAttemptStart {
|
||||
model_request_id: attempt.request_id.as_str().to_owned(),
|
||||
attempt_count: attempt.attempt_count.get(),
|
||||
provider_kind: attempt.provider_kind.as_str().to_owned(),
|
||||
provider_account_id: attempt
|
||||
.provider_account_id
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
provider_account_ref: attempt
|
||||
.provider_account_ref
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
upstream_model_id: attempt.upstream_model_id.as_str().to_owned(),
|
||||
upstream_transport: attempt.upstream_transport,
|
||||
http_version: attempt.http_version,
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_i64(value: Option<u64>, field: &'static str) -> StoreResult<Option<i64>> {
|
||||
value.map(|value| to_i64(value, field)).transpose()
|
||||
}
|
||||
|
||||
@@ -1,111 +0,0 @@
|
||||
//! 不含正文或会话状态的模型请求历史查询。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::PgPool;
|
||||
|
||||
use crate::{StoreResult, postgres_unavailable, require_nonempty};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ModelRequestHistoryRecord {
|
||||
pub id: String,
|
||||
pub client_api_key_ref: String,
|
||||
pub requested_model_id: String,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub outcome: String,
|
||||
pub client_response_id: Option<String>,
|
||||
pub upstream_response_id: Option<String>,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ModelRequestHistoryRepository: Send + Sync {
|
||||
async fn find_model_request_history(
|
||||
&self,
|
||||
model_request_id: &str,
|
||||
) -> StoreResult<Option<ModelRequestHistoryRecord>>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PgHistoryRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl PgHistoryRepository {
|
||||
#[must_use]
|
||||
pub const fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelRequestHistoryRepository for PgHistoryRepository {
|
||||
async fn find_model_request_history(
|
||||
&self,
|
||||
model_request_id: &str,
|
||||
) -> StoreResult<Option<ModelRequestHistoryRecord>> {
|
||||
require_nonempty(
|
||||
"model request history",
|
||||
"model_request_id",
|
||||
model_request_id,
|
||||
)?;
|
||||
fetch_history(&self.pool, "id", model_request_id).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_history(
|
||||
pool: &PgPool,
|
||||
lookup: &'static str,
|
||||
value: &str,
|
||||
) -> StoreResult<Option<ModelRequestHistoryRecord>> {
|
||||
let row = match lookup {
|
||||
"id" => {
|
||||
sqlx::query_as::<_, HistoryRow>(
|
||||
"select id, client_api_key_ref, requested_model_id, provider_kind,
|
||||
provider_account_ref, upstream_model_id, outcome,
|
||||
client_response_id, upstream_response_id,
|
||||
started_at, completed_at
|
||||
from model_requests where id = $1",
|
||||
)
|
||||
.bind(value)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
}
|
||||
_ => return Ok(None),
|
||||
}
|
||||
.map_err(|_| postgres_unavailable("read model request history"))?;
|
||||
Ok(row.map(history_from_row))
|
||||
}
|
||||
|
||||
type HistoryRow = (
|
||||
String,
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
DateTime<Utc>,
|
||||
Option<DateTime<Utc>>,
|
||||
);
|
||||
|
||||
fn history_from_row(row: HistoryRow) -> ModelRequestHistoryRecord {
|
||||
ModelRequestHistoryRecord {
|
||||
id: row.0,
|
||||
client_api_key_ref: row.1,
|
||||
requested_model_id: row.2,
|
||||
provider_kind: row.3,
|
||||
provider_account_ref: row.4,
|
||||
upstream_model_id: row.5,
|
||||
outcome: row.6,
|
||||
client_response_id: row.7,
|
||||
upstream_response_id: row.8,
|
||||
started_at: row.9,
|
||||
completed_at: row.10,
|
||||
}
|
||||
}
|
||||
@@ -3,13 +3,14 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, postgres::PgPoolOptions};
|
||||
|
||||
use crate::{Revision, StoreBackend, StoreError, StoreResult, postgres_unavailable};
|
||||
use crate::{
|
||||
Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
|
||||
};
|
||||
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod history;
|
||||
mod observability;
|
||||
mod ops_events;
|
||||
mod provider_accounts;
|
||||
@@ -21,7 +22,6 @@ pub use admin_security_audit::*;
|
||||
pub use admission_recovery::*;
|
||||
pub use client_keys::*;
|
||||
pub use execution::*;
|
||||
pub use history::*;
|
||||
pub use observability::*;
|
||||
pub use ops_events::*;
|
||||
pub use provider_accounts::*;
|
||||
@@ -31,12 +31,19 @@ pub use snapshot::*;
|
||||
|
||||
static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations");
|
||||
|
||||
/// 建立 PostgreSQL pool 并只执行冻结的 `0001_initial.sql` migration 集。
|
||||
pub async fn connect_and_migrate(database_url: &str) -> StoreResult<PgPool> {
|
||||
/// 建立 PostgreSQL pool 并只执行冻结的 migration 集。
|
||||
pub async fn connect_and_migrate(
|
||||
database_url: &str,
|
||||
pool_config: StorePoolConfig,
|
||||
) -> StoreResult<PgPool> {
|
||||
if database_url.trim().is_empty() {
|
||||
return Err(postgres_unavailable("connect PostgreSQL"));
|
||||
}
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(pool_config.max_connections)
|
||||
.acquire_timeout(std::time::Duration::from_secs(
|
||||
pool_config.acquire_timeout_seconds,
|
||||
))
|
||||
.connect(database_url)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("connect PostgreSQL"))?;
|
||||
|
||||
@@ -23,7 +23,8 @@ const MAX_PAGE_SIZE: u16 = 100;
|
||||
const MAX_FILTER_BYTES: usize = 256;
|
||||
const MAX_SEARCH_BYTES: usize = 512;
|
||||
const MAX_ACCOUNT_IDS: usize = 200;
|
||||
const DASHBOARD_ACCOUNT_LIMIT: u16 = 50;
|
||||
/// 概览卡只展示最近使用的四个账号;完整账号用量由账号管理页单独查询。
|
||||
const DASHBOARD_ACCOUNT_LIMIT: u16 = 4;
|
||||
const DIAGNOSTIC_LIMIT: i64 = 100;
|
||||
const ACCOUNT_USAGE_TIMELINE_HOURS: i64 = 24;
|
||||
|
||||
@@ -660,7 +661,8 @@ pub struct OpsErrorRecord {
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub failure_kind: String,
|
||||
pub status_code: Option<u16>,
|
||||
pub client_status_code: Option<u16>,
|
||||
pub upstream_status_code: Option<u16>,
|
||||
pub provider_error_code: Option<String>,
|
||||
pub client_response_id: Option<String>,
|
||||
pub upstream_request_id: Option<String>,
|
||||
@@ -757,35 +759,35 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
range: ObservabilityRange,
|
||||
) -> StoreResult<DashboardObservation> {
|
||||
let filter = UsageRecordFilter::default();
|
||||
let requests = request_metrics(&self.pool, range, &filter).await?;
|
||||
let attempts = attempt_metrics(&self.pool, range, &filter).await?;
|
||||
let provider_accounts = provider_account_metrics(&self.pool, range.end).await?;
|
||||
let trend = request_metric_series(&self.pool, range, &filter).await?;
|
||||
let account_usage_range = ObservabilityRange::new(
|
||||
range.end - TimeDelta::hours(ACCOUNT_USAGE_TIMELINE_HOURS),
|
||||
range.end,
|
||||
)?;
|
||||
let account_usage = provider_account_usage(
|
||||
&self.pool,
|
||||
let account_usage_query =
|
||||
ProviderAccountUsageQuery::recent(account_usage_range, DASHBOARD_ACCOUNT_LIMIT)?
|
||||
.with_hourly_request_buckets()?,
|
||||
)
|
||||
.await?;
|
||||
let recent_requests = list_usage_records(
|
||||
&self.pool,
|
||||
UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
outcome: Some("succeeded".to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1)?,
|
||||
page_size: ObservabilityPageSize::new(10)?,
|
||||
.with_hourly_request_buckets()?;
|
||||
let recent_query = UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
outcome: Some("succeeded".to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
)
|
||||
.await?
|
||||
.items;
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1)?,
|
||||
page_size: ObservabilityPageSize::new(10)?,
|
||||
};
|
||||
// 分两批并发:页面延迟从六路之和降为两批各自最慢者之和,同时单次
|
||||
// 渲染最多占用 3 个池连接,不与数据面写路径争抢整个连接池。
|
||||
let (requests, attempts, provider_accounts) = futures::try_join!(
|
||||
request_metrics(&self.pool, range, &filter),
|
||||
attempt_metrics(&self.pool, range, &filter),
|
||||
provider_account_metrics(&self.pool, range.end),
|
||||
)?;
|
||||
let (trend, account_usage, recent_requests) = futures::try_join!(
|
||||
request_metric_series(&self.pool, range, &filter),
|
||||
provider_account_usage(&self.pool, account_usage_query),
|
||||
async { Ok(list_usage_records(&self.pool, recent_query).await?.items) },
|
||||
)?;
|
||||
Ok(DashboardObservation {
|
||||
range,
|
||||
requests,
|
||||
@@ -1626,7 +1628,8 @@ fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::OpsError {
|
||||
upstream_model_id: error.upstream_model_id,
|
||||
upstream_transport: error.upstream_transport,
|
||||
failure_kind: error.failure_kind,
|
||||
status_code: error.status_code,
|
||||
client_status_code: error.client_status_code,
|
||||
upstream_status_code: error.upstream_status_code,
|
||||
provider_error_code: error.provider_error_code,
|
||||
client_response_id: error.client_response_id,
|
||||
upstream_request_id: error.upstream_request_id,
|
||||
@@ -1648,7 +1651,10 @@ async fn request_metrics(
|
||||
filter: &UsageRecordFilter,
|
||||
) -> StoreResult<RequestMetrics> {
|
||||
filter.validate()?;
|
||||
let mut query = QueryBuilder::<Postgres>::new(
|
||||
// 结果计数覆盖范围内全部请求(成功率/失败率的分母分子);用量、缓存、
|
||||
// 延迟与成本聚合仅统计用量事实(已完整交付客户端的成功响应)。
|
||||
let fact = completed_usage_fact_predicate("mr");
|
||||
let mut query = QueryBuilder::<Postgres>::new(format!(
|
||||
"select count(*)::bigint as request_count,
|
||||
count(*) filter (where outcome = 'succeeded')::bigint as success_count,
|
||||
count(*) filter (where outcome = 'failed')::bigint as failure_count,
|
||||
@@ -1656,40 +1662,43 @@ async fn request_metrics(
|
||||
count(*) filter (where outcome = 'incomplete')::bigint as incomplete_count,
|
||||
count(*) filter (where client_status_code between 400 and 499)::bigint
|
||||
as caller_error_count,
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
|
||||
coalesce(sum(reasoning_tokens), 0)::bigint as reasoning_tokens,
|
||||
coalesce(sum(total_tokens), 0)::bigint as total_tokens,
|
||||
coalesce(sum(first_token_ms), 0)::bigint as first_token_latency_sum,
|
||||
count(first_token_ms)::bigint as first_token_latency_count,
|
||||
coalesce(sum(latency_ms), 0)::bigint as latency_sum,
|
||||
count(latency_ms)::bigint as latency_count,
|
||||
max(latency_ms)::bigint as max_latency_ms,
|
||||
min(latency_ms)::bigint as min_latency_ms,
|
||||
count(*) filter (where input_tokens is not null)::bigint
|
||||
coalesce(sum(input_tokens) filter (where {fact}), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens) filter (where {fact}), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens) filter (where {fact}), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens) filter (where {fact}), 0)::bigint
|
||||
as cache_write_tokens,
|
||||
coalesce(sum(reasoning_tokens) filter (where {fact}), 0)::bigint
|
||||
as reasoning_tokens,
|
||||
coalesce(sum(total_tokens) filter (where {fact}), 0)::bigint as total_tokens,
|
||||
coalesce(sum(first_token_ms) filter (where {fact}), 0)::bigint
|
||||
as first_token_latency_sum,
|
||||
count(first_token_ms) filter (where {fact})::bigint as first_token_latency_count,
|
||||
coalesce(sum(latency_ms) filter (where {fact}), 0)::bigint as latency_sum,
|
||||
count(latency_ms) filter (where {fact})::bigint as latency_count,
|
||||
max(latency_ms) filter (where {fact})::bigint as max_latency_ms,
|
||||
min(latency_ms) filter (where {fact})::bigint as min_latency_ms,
|
||||
count(*) filter (where {fact} and input_tokens is not null)::bigint
|
||||
as cache_eligible_request_count,
|
||||
count(*) filter (where input_tokens is not null and cached_tokens > 0)::bigint
|
||||
count(*) filter (where {fact} and input_tokens is not null
|
||||
and cached_tokens > 0)::bigint
|
||||
as cache_hit_request_count,
|
||||
percentile_cont(0.50) within group (order by latency_ms)
|
||||
as latency_p50_ms,
|
||||
filter (where {fact}) as latency_p50_ms,
|
||||
percentile_cont(0.95) within group (order by latency_ms)
|
||||
as latency_p95_ms,
|
||||
filter (where {fact}) as latency_p95_ms,
|
||||
percentile_cont(0.99) within group (order by latency_ms)
|
||||
as latency_p99_ms,
|
||||
filter (where {fact}) as latency_p99_ms,
|
||||
percentile_cont(0.50) within group (order by first_token_ms)
|
||||
as first_token_p50_ms,
|
||||
filter (where {fact}) as first_token_p50_ms,
|
||||
percentile_cont(0.95) within group (order by first_token_ms)
|
||||
as first_token_p95_ms,
|
||||
filter (where {fact}) as first_token_p95_ms,
|
||||
percentile_cont(0.99) within group (order by first_token_ms)
|
||||
as first_token_p99_ms
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
filter (where {fact}) as first_token_p99_ms
|
||||
from model_requests mr where mr.started_at >= "
|
||||
));
|
||||
query.push_bind(range.start);
|
||||
query.push(" and mr.started_at < ");
|
||||
query.push_bind(range.end);
|
||||
push_completed_usage_fact_filter(&mut query, "mr");
|
||||
push_usage_filter(&mut query, filter, "mr");
|
||||
let row = query
|
||||
.build()
|
||||
@@ -1706,9 +1715,12 @@ async fn request_metric_series(
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
filter.validate()?;
|
||||
let granularity = granularity_for(range);
|
||||
// 与 request_metrics 同一契约:结果计数覆盖全部请求,用量/延迟/成本
|
||||
// 聚合仅统计用量事实。
|
||||
let fact = completed_usage_fact_predicate("mr");
|
||||
let mut query = QueryBuilder::<Postgres>::new("select date_bin(");
|
||||
query.push_bind(granularity.sql_interval());
|
||||
query.push(
|
||||
query.push(format!(
|
||||
"::interval, mr.started_at, timestamptz '1970-01-01 00:00:00+00') as bucket_start,
|
||||
count(*)::bigint as request_count,
|
||||
count(*) filter (where outcome = 'succeeded')::bigint as success_count,
|
||||
@@ -1717,46 +1729,49 @@ async fn request_metric_series(
|
||||
count(*) filter (where outcome = 'incomplete')::bigint as incomplete_count,
|
||||
count(*) filter (where client_status_code between 400 and 499)::bigint
|
||||
as caller_error_count,
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
|
||||
coalesce(sum(reasoning_tokens), 0)::bigint as reasoning_tokens,
|
||||
coalesce(sum(total_tokens), 0)::bigint as total_tokens,
|
||||
coalesce(sum(first_token_ms), 0)::bigint as first_token_latency_sum,
|
||||
count(first_token_ms)::bigint as first_token_latency_count,
|
||||
coalesce(sum(latency_ms), 0)::bigint as latency_sum,
|
||||
count(latency_ms)::bigint as latency_count,
|
||||
max(latency_ms)::bigint as max_latency_ms,
|
||||
min(latency_ms)::bigint as min_latency_ms,
|
||||
count(*) filter (where input_tokens is not null)::bigint
|
||||
coalesce(sum(input_tokens) filter (where {fact}), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens) filter (where {fact}), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens) filter (where {fact}), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens) filter (where {fact}), 0)::bigint
|
||||
as cache_write_tokens,
|
||||
coalesce(sum(reasoning_tokens) filter (where {fact}), 0)::bigint
|
||||
as reasoning_tokens,
|
||||
coalesce(sum(total_tokens) filter (where {fact}), 0)::bigint as total_tokens,
|
||||
coalesce(sum(first_token_ms) filter (where {fact}), 0)::bigint
|
||||
as first_token_latency_sum,
|
||||
count(first_token_ms) filter (where {fact})::bigint as first_token_latency_count,
|
||||
coalesce(sum(latency_ms) filter (where {fact}), 0)::bigint as latency_sum,
|
||||
count(latency_ms) filter (where {fact})::bigint as latency_count,
|
||||
max(latency_ms) filter (where {fact})::bigint as max_latency_ms,
|
||||
min(latency_ms) filter (where {fact})::bigint as min_latency_ms,
|
||||
count(*) filter (where {fact} and input_tokens is not null)::bigint
|
||||
as cache_eligible_request_count,
|
||||
count(*) filter (where input_tokens is not null and cached_tokens > 0)::bigint
|
||||
count(*) filter (where {fact} and input_tokens is not null
|
||||
and cached_tokens > 0)::bigint
|
||||
as cache_hit_request_count,
|
||||
percentile_cont(0.50) within group (order by latency_ms)
|
||||
as latency_p50_ms,
|
||||
filter (where {fact}) as latency_p50_ms,
|
||||
percentile_cont(0.95) within group (order by latency_ms)
|
||||
as latency_p95_ms,
|
||||
filter (where {fact}) as latency_p95_ms,
|
||||
percentile_cont(0.99) within group (order by latency_ms)
|
||||
as latency_p99_ms,
|
||||
filter (where {fact}) as latency_p99_ms,
|
||||
percentile_cont(0.50) within group (order by first_token_ms)
|
||||
as first_token_p50_ms,
|
||||
filter (where {fact}) as first_token_p50_ms,
|
||||
percentile_cont(0.95) within group (order by first_token_ms)
|
||||
as first_token_p95_ms,
|
||||
filter (where {fact}) as first_token_p95_ms,
|
||||
percentile_cont(0.99) within group (order by first_token_ms)
|
||||
as first_token_p99_ms,
|
||||
count(*) filter (where cost_source = 'provider_reported')::bigint
|
||||
filter (where {fact}) as first_token_p99_ms,
|
||||
count(*) filter (where {fact} and cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(*) filter (where cost_source = 'calculated')::bigint
|
||||
count(*) filter (where {fact} and cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(*) filter (where cost_source = 'unavailable')::bigint
|
||||
count(*) filter (where {fact} and cost_source = 'unavailable')::bigint
|
||||
as unavailable_count
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
from model_requests mr where mr.started_at >= "
|
||||
));
|
||||
query.push_bind(range.start);
|
||||
query.push(" and mr.started_at < ");
|
||||
query.push_bind(range.end);
|
||||
push_completed_usage_fact_filter(&mut query, "mr");
|
||||
push_usage_filter(&mut query, filter, "mr");
|
||||
query.push(" group by bucket_start order by bucket_start");
|
||||
let rows = query
|
||||
@@ -1977,8 +1992,12 @@ async fn provider_account_metrics(
|
||||
let row = sqlx::query(
|
||||
"select count(*)::bigint as total,
|
||||
count(*) filter (where enabled)::bigint as enabled,
|
||||
count(*) filter (where not enabled or availability <> 'ready')::bigint
|
||||
as unavailable,
|
||||
count(*) filter (
|
||||
where not (
|
||||
enabled and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
)
|
||||
)::bigint as unavailable,
|
||||
count(*) filter (
|
||||
where enabled and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
@@ -2242,8 +2261,14 @@ fn push_usage_filter(query: &mut QueryBuilder<Postgres>, filter: &UsageRecordFil
|
||||
/// `model_requests` 同时承担执行审计:包括上游已发送、但尚未收到首个事件就断开的
|
||||
/// WebSocket 请求。这些失败仍必须留给 Ops Errors 和调度健康度分析,不能混入用量、
|
||||
/// 成本、账号使用次数或请求明细。
|
||||
fn completed_usage_fact_predicate(alias: &str) -> String {
|
||||
format!(
|
||||
"{alias}.outcome = 'succeeded' and {alias}.downstream_committed_at is not null and {alias}.client_status_code between 200 and 399"
|
||||
)
|
||||
}
|
||||
|
||||
fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres>, alias: &str) {
|
||||
query.push(format!(" and {alias}.outcome = 'succeeded' and {alias}.downstream_committed_at is not null and {alias}.client_status_code between 200 and 399"));
|
||||
query.push(format!(" and {}", completed_usage_fact_predicate(alias)));
|
||||
}
|
||||
|
||||
async fn provider_account_usage(
|
||||
@@ -2921,7 +2946,7 @@ const OPS_ERRORS_CTE: &str = "with errors as (
|
||||
mr.provider_kind,
|
||||
mr.provider_account_ref, mr.upstream_model_id, mr.upstream_transport,
|
||||
coalesce(mr.error_kind, 'failed') as failure_kind,
|
||||
coalesce(mr.upstream_status_code, mr.client_status_code) as status_code,
|
||||
mr.client_status_code, mr.upstream_status_code,
|
||||
mr.provider_error_code, mr.client_response_id, mr.upstream_request_id,
|
||||
mr.latency_ms, coalesce(mr.error_message, mr.error_kind, 'request failed') as message,
|
||||
1::integer as occurrence_count,
|
||||
@@ -2932,7 +2957,8 @@ const OPS_ERRORS_CTE: &str = "with errors as (
|
||||
select 'ops_event'::text, oe.id, oe.model_request_id, oe.attempt_index,
|
||||
mr.client_api_key_ref, oe.component, oe.operation,
|
||||
oe.provider_kind, oe.provider_account_ref,
|
||||
oe.upstream_model_id, null::text, oe.failure_kind, oe.status_code,
|
||||
oe.upstream_model_id, null::text, oe.failure_kind,
|
||||
null::integer as client_status_code, oe.status_code as upstream_status_code,
|
||||
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
|
||||
oe.latency_ms, oe.message, oe.occurrence_count, oe.created_at,
|
||||
'ops_event:' || oe.id
|
||||
@@ -2947,7 +2973,8 @@ async fn list_ops_errors(pool: &PgPool, query: OpsErrorQuery) -> StoreResult<Ops
|
||||
" select source, event_id, request_id, attempt_index, client_api_key_ref,
|
||||
component, operation, provider_kind, provider_account_ref,
|
||||
upstream_model_id, upstream_transport,
|
||||
failure_kind, status_code, provider_error_code, client_response_id,
|
||||
failure_kind, client_status_code, upstream_status_code,
|
||||
provider_error_code, client_response_id,
|
||||
upstream_request_id, latency_ms, message, occurrence_count, occurred_at,
|
||||
stable_sort_id
|
||||
from errors e where e.occurred_at >= ",
|
||||
@@ -3046,7 +3073,7 @@ fn push_ops_filter(statement: &mut QueryBuilder<Postgres>, filter: &OpsErrorFilt
|
||||
statement.push_bind(i32::try_from(index).unwrap_or(i32::MAX));
|
||||
}
|
||||
if let Some(status) = filter.status_code {
|
||||
statement.push(" and e.status_code = ");
|
||||
statement.push(" and e.upstream_status_code = ");
|
||||
statement.push_bind(i32::from(status));
|
||||
}
|
||||
if let Some(search) = &filter.search {
|
||||
@@ -3147,7 +3174,8 @@ fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<OpsErrorRecord
|
||||
upstream_model_id: get(row, "upstream_model_id")?,
|
||||
upstream_transport: get(row, "upstream_transport")?,
|
||||
failure_kind: get(row, "failure_kind")?,
|
||||
status_code: optional_status(row, "status_code")?,
|
||||
client_status_code: optional_status(row, "client_status_code")?,
|
||||
upstream_status_code: optional_status(row, "upstream_status_code")?,
|
||||
provider_error_code: get(row, "provider_error_code")?,
|
||||
client_response_id: get(row, "client_response_id")?,
|
||||
upstream_request_id: get(row, "upstream_request_id")?,
|
||||
|
||||
@@ -460,10 +460,6 @@ pub trait ProviderAccountRepository: Send + Sync {
|
||||
&self,
|
||||
update: ProviderCredentialUpdate,
|
||||
) -> StoreResult<Revision>;
|
||||
async fn update_provider_account_observation(
|
||||
&self,
|
||||
observation: ProviderAccountObservation,
|
||||
) -> StoreResult<bool>;
|
||||
async fn apply_provider_account_state(
|
||||
&self,
|
||||
update: ProviderAccountStateUpdate,
|
||||
@@ -654,36 +650,6 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
Revision::new(to_u64(next)?)
|
||||
}
|
||||
|
||||
async fn update_provider_account_observation(
|
||||
&self,
|
||||
observation: ProviderAccountObservation,
|
||||
) -> StoreResult<bool> {
|
||||
observation.validate()?;
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts
|
||||
set availability = $2, availability_reason = $3, cooldown_until = $4,
|
||||
provider_quota_json = $5, availability_observed_at = $6,
|
||||
quota_observed_at = $7,
|
||||
updated_at = greatest(now(), $6, coalesce($7, $6))
|
||||
where id = $1",
|
||||
)
|
||||
.bind(observation.account_id)
|
||||
.bind(observation.availability.as_str())
|
||||
.bind(observation.availability_reason)
|
||||
.bind(observation.cooldown_until)
|
||||
.bind(
|
||||
observation
|
||||
.provider_quota_json
|
||||
.map(|quota| quota.as_value()),
|
||||
)
|
||||
.bind(observation.availability_observed_at)
|
||||
.bind(observation.quota_observed_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update provider account observation"))?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn apply_provider_account_state(
|
||||
&self,
|
||||
update: ProviderAccountStateUpdate,
|
||||
@@ -1700,10 +1666,8 @@ fn admin_account_status(
|
||||
.is_some_and(|expires_at| expires_at <= now)
|
||||
{
|
||||
AdminAccountStatus::Expired
|
||||
} else if account.availability == ProviderAccountAvailability::Ready {
|
||||
AdminAccountStatus::Active
|
||||
} else {
|
||||
AdminAccountStatus::Attention
|
||||
AdminAccountStatus::Active
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1728,7 +1692,7 @@ fn admin_account_summary(
|
||||
.count(),
|
||||
)
|
||||
.unwrap_or(u64::MAX);
|
||||
let attention = u64::try_from(
|
||||
let unavailable = u64::try_from(
|
||||
accounts
|
||||
.iter()
|
||||
.filter(|account| {
|
||||
@@ -1746,7 +1710,7 @@ fn admin_account_summary(
|
||||
total,
|
||||
active,
|
||||
quota_exhausted,
|
||||
attention,
|
||||
unavailable,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1797,7 +1761,6 @@ const fn admin_account_status_name(status: AdminAccountStatus) -> &'static str {
|
||||
AdminAccountStatus::QuotaExhausted => "quota_exhausted",
|
||||
AdminAccountStatus::Disabled => "disabled",
|
||||
AdminAccountStatus::Banned => "banned",
|
||||
AdminAccountStatus::Attention => "attention",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -78,47 +78,49 @@ impl RetentionRepository for PgRetentionRepository {
|
||||
});
|
||||
}
|
||||
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("begin retention transaction"))?;
|
||||
let model_requests = sqlx::query(
|
||||
// 分批删除,每批独立提交:清理是幂等的,无需跨表原子性;
|
||||
// 有界批量避免百万行单事务的长锁与 WAL 峰值。
|
||||
let model_requests = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from model_requests
|
||||
where outcome <> 'running'
|
||||
and completed_at < $1 - ($2 * interval '1 day')",
|
||||
where ctid in (
|
||||
select ctid from model_requests
|
||||
where outcome <> 'running'
|
||||
and completed_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.usage_retention_days,
|
||||
"delete expired model requests",
|
||||
)
|
||||
.bind(now)
|
||||
.bind(i64::from(settings.usage_retention_days))
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("delete expired model requests"))?
|
||||
.rows_affected();
|
||||
let ops_events = sqlx::query(
|
||||
.await?;
|
||||
let ops_events = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from ops_events
|
||||
where model_request_id is null
|
||||
and created_at < $1 - ($2 * interval '1 day')",
|
||||
where ctid in (
|
||||
select ctid from ops_events
|
||||
where model_request_id is null
|
||||
and created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.ops_event_retention_days,
|
||||
"delete expired ops events",
|
||||
)
|
||||
.bind(now)
|
||||
.bind(i64::from(settings.ops_event_retention_days))
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("delete expired ops events"))?
|
||||
.rows_affected();
|
||||
let admin_audit_events = sqlx::query(
|
||||
.await?;
|
||||
let admin_audit_events = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from admin_audit_events
|
||||
where created_at < $1 - ($2 * interval '1 day')",
|
||||
where ctid in (
|
||||
select ctid from admin_audit_events
|
||||
where created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.audit_retention_days,
|
||||
"delete expired admin audit events",
|
||||
)
|
||||
.bind(now)
|
||||
.bind(i64::from(settings.audit_retention_days))
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("delete expired admin audit events"))?
|
||||
.rows_affected();
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("commit retention transaction"))?;
|
||||
.await?;
|
||||
Ok(RetentionReport {
|
||||
model_requests,
|
||||
ops_events,
|
||||
@@ -127,6 +129,32 @@ impl RetentionRepository for PgRetentionRepository {
|
||||
}
|
||||
}
|
||||
|
||||
const RETENTION_DELETE_BATCH_ROWS: i64 = 10_000;
|
||||
|
||||
async fn purge_in_batches(
|
||||
pool: &PgPool,
|
||||
delete_sql: &'static str,
|
||||
now: DateTime<Utc>,
|
||||
retention_days: u32,
|
||||
label: &'static str,
|
||||
) -> StoreResult<u64> {
|
||||
let mut total = 0u64;
|
||||
loop {
|
||||
let deleted = sqlx::query(delete_sql)
|
||||
.bind(now)
|
||||
.bind(i64::from(retention_days))
|
||||
.bind(RETENTION_DELETE_BATCH_ROWS)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable(label))?
|
||||
.rows_affected();
|
||||
total += deleted;
|
||||
if deleted < RETENTION_DELETE_BATCH_ROWS.unsigned_abs() {
|
||||
return Ok(total);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn to_u32(value: i64) -> StoreResult<u32> {
|
||||
u32::try_from(value).map_err(|_| StoreError::InvalidData {
|
||||
entity: "retention settings",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! PostgreSQL 账号 cooldown 事实的可丢失 Redis 热缓存。
|
||||
//! 请求调度用的可丢失账号 cooldown Redis 存储。
|
||||
|
||||
use std::{sync::Arc, time::SystemTime};
|
||||
|
||||
@@ -243,7 +243,7 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
}
|
||||
}
|
||||
|
||||
/// PostgreSQL 账号事实到可丢失 cooldown cache 的 write-through adapter。
|
||||
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
|
||||
pub struct CooldownCachingProviderAccountStore {
|
||||
authoritative: Arc<dyn ProviderAccountStore>,
|
||||
cooldowns: Arc<dyn CredentialCooldownRepository>,
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! Provider account availability 的可重建 Redis cache fence。
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use redis::{Script, aio::ConnectionManager};
|
||||
@@ -28,6 +30,12 @@ redis.call('PEXPIRE', KEYS[1], ARGV[5])
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const RECORD_REFRESH_BACKOFF_SCRIPT: &str = r#"
|
||||
local n = redis.call('INCR', KEYS[1])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[1])
|
||||
return n
|
||||
"#;
|
||||
|
||||
const MAX_PROVIDER_CATALOG_BYTES: usize = 1024 * 1024;
|
||||
const CREDENTIAL_STATE_TTL_SECONDS: u64 = 24 * 60 * 60;
|
||||
|
||||
@@ -132,6 +140,14 @@ impl RedisCredentialStateRepository {
|
||||
Ok(format!("{}:account:{fingerprint}:state", self.namespace))
|
||||
}
|
||||
|
||||
fn backoff_key(&self, provider_account_id: &str) -> StoreResult<String> {
|
||||
let fingerprint = resource_fingerprint("credential refresh backoff", provider_account_id)?;
|
||||
Ok(format!(
|
||||
"{}:account:{fingerprint}:refresh-backoff",
|
||||
self.namespace
|
||||
))
|
||||
}
|
||||
|
||||
fn catalog_key(&self, key: &RedisProviderCatalogCacheKey) -> StoreResult<String> {
|
||||
key.validate()?;
|
||||
let fingerprint = resource_fingerprint("provider catalog cache", &key.catalog_scope)?;
|
||||
@@ -340,6 +356,48 @@ impl ProviderCredentialStatePort for RedisCredentialStateRepository {
|
||||
.map_err(|_| provider_unavailable("clear credential state"))
|
||||
})
|
||||
}
|
||||
|
||||
fn record_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
window: Duration,
|
||||
) -> futures::future::BoxFuture<'a, Result<u32, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let window_millis = u64::try_from(window.as_millis())
|
||||
.ok()
|
||||
.filter(|millis| *millis > 0)
|
||||
.ok_or_else(|| provider_invalid("record refresh backoff"))?;
|
||||
let key = self
|
||||
.backoff_key(account_id.as_str())
|
||||
.map_err(|_| provider_unavailable("record refresh backoff"))?;
|
||||
let mut connection = self.connection.clone();
|
||||
let count = Script::new(RECORD_REFRESH_BACKOFF_SCRIPT)
|
||||
.key(key)
|
||||
.arg(window_millis)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("record refresh backoff"))?;
|
||||
Ok(u32::try_from(count).unwrap_or(u32::MAX))
|
||||
})
|
||||
}
|
||||
|
||||
fn clear_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> futures::future::BoxFuture<'a, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let key = self
|
||||
.backoff_key(account_id.as_str())
|
||||
.map_err(|_| provider_unavailable("clear refresh backoff"))?;
|
||||
let mut connection = self.connection.clone();
|
||||
redis::cmd("DEL")
|
||||
.arg(key)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("clear refresh backoff"))?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCatalogCachePort for RedisCredentialStateRepository {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 可丢失、可从 PostgreSQL 或 Provider 重建的 Redis 协调状态。
|
||||
|
||||
use chrono::{DateTime, SecondsFormat, Utc};
|
||||
use redis::{Script, aio::ConnectionManager};
|
||||
use redis::aio::ConnectionManager;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
@@ -33,17 +33,6 @@ pub use runtime_change::*;
|
||||
use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty};
|
||||
|
||||
pub(crate) const MAX_REDIS_EXACT_INTEGER: u64 = (1_u64 << 53) - 1;
|
||||
const RECORD_ADMIN_LOGIN_FAILURE_SCRIPT: &str = r#"
|
||||
local count = redis.call('INCR', KEYS[1])
|
||||
local ttl = redis.call('PTTL', KEYS[1])
|
||||
if count == 1 or ttl < 0 then
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[2])
|
||||
end
|
||||
if count >= tonumber(ARGV[1]) then
|
||||
return 1
|
||||
end
|
||||
return 0
|
||||
"#;
|
||||
|
||||
/// Redis 中可丢失的管理员会话事实;认证秘密不属于该结构。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -71,7 +60,7 @@ impl AdminSessionRecord {
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理员会话与登录失败窗口的 Redis 基础设施端口。
|
||||
/// 管理员会话的 Redis 基础设施端口。
|
||||
#[async_trait]
|
||||
pub trait AdminAuthStateRepository: Send + Sync {
|
||||
async fn load_admin_session(&self, session_id: &str)
|
||||
@@ -85,22 +74,9 @@ pub trait AdminAuthStateRepository: Send + Sync {
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> StoreResult<Option<AdminSessionRecord>>;
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> StoreResult<bool>;
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> StoreResult<bool>;
|
||||
async fn clear_login_failures(&self, source: &str) -> StoreResult<()>;
|
||||
}
|
||||
|
||||
/// Redis 管理员会话与登录失败窗口 adapter。
|
||||
/// Redis 管理员会话 adapter。
|
||||
#[derive(Clone)]
|
||||
pub struct RedisAdminAuthStateRepository {
|
||||
connection: ConnectionManager,
|
||||
@@ -119,11 +95,6 @@ impl RedisAdminAuthStateRepository {
|
||||
let fingerprint = resource_fingerprint("admin session", session_id)?;
|
||||
Ok(format!("{}:session:{{{fingerprint}}}", self.namespace))
|
||||
}
|
||||
|
||||
fn failure_key(&self, source: &str) -> StoreResult<String> {
|
||||
let fingerprint = resource_fingerprint("admin login failure", source)?;
|
||||
Ok(format!("{}:failure:{{{fingerprint}}}", self.namespace))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -179,54 +150,6 @@ impl AdminAuthStateRepository for RedisAdminAuthStateRepository {
|
||||
.map(|value| decode_admin_session(&value))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn login_source_is_throttled(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> StoreResult<bool> {
|
||||
validate_failure_policy(failure_limit, window_seconds)?;
|
||||
let key = self.failure_key(source)?;
|
||||
let mut connection = self.connection.clone();
|
||||
let failures = redis::cmd("GET")
|
||||
.arg(key)
|
||||
.query_async::<Option<u64>>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("load admin login failure window"))?
|
||||
.unwrap_or_default();
|
||||
Ok(failures >= u64::from(failure_limit))
|
||||
}
|
||||
|
||||
async fn record_login_failure(
|
||||
&self,
|
||||
source: &str,
|
||||
failure_limit: u32,
|
||||
window_seconds: u64,
|
||||
) -> StoreResult<bool> {
|
||||
let window_millis = validate_failure_policy(failure_limit, window_seconds)?;
|
||||
let key = self.failure_key(source)?;
|
||||
let mut connection = self.connection.clone();
|
||||
let throttled = Script::new(RECORD_ADMIN_LOGIN_FAILURE_SCRIPT)
|
||||
.key(key)
|
||||
.arg(failure_limit)
|
||||
.arg(window_millis)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("record admin login failure"))?;
|
||||
Ok(throttled == 1)
|
||||
}
|
||||
|
||||
async fn clear_login_failures(&self, source: &str) -> StoreResult<()> {
|
||||
let key = self.failure_key(source)?;
|
||||
let mut connection = self.connection.clone();
|
||||
redis::cmd("DEL")
|
||||
.arg(key)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("clear admin login failures"))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
@@ -259,17 +182,6 @@ fn decode_admin_session(value: &str) -> StoreResult<AdminSessionRecord> {
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_failure_policy(failure_limit: u32, window_seconds: u64) -> StoreResult<u64> {
|
||||
if failure_limit == 0 {
|
||||
return Err(admin_auth_invalid("login failure limit must be positive"));
|
||||
}
|
||||
let window_millis = window_seconds
|
||||
.checked_mul(1_000)
|
||||
.filter(|value| *value > 0 && *value <= MAX_REDIS_EXACT_INTEGER)
|
||||
.ok_or_else(|| admin_auth_invalid("login failure window is outside the supported range"))?;
|
||||
Ok(window_millis)
|
||||
}
|
||||
|
||||
fn admin_auth_invalid(message: &str) -> StoreError {
|
||||
StoreError::InvalidData {
|
||||
entity: "admin authentication state",
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
mod postgres;
|
||||
mod redis;
|
||||
mod support;
|
||||
|
||||
@@ -4,7 +4,9 @@ use gateway_core::engine::{
|
||||
ExecutionOutcome, ExecutionStore, ModelRequestFinalization as CoreModelRequestFinalization,
|
||||
ModelRequestId, ModelRequestTimings as CoreModelRequestTimings, UpstreamSendState,
|
||||
};
|
||||
use gateway_store::postgres::{ModelRequestRepository, NewModelRequest, PgExecutionStore};
|
||||
use gateway_store::postgres::{
|
||||
ModelRequestAttemptStart, ModelRequestRepository, NewModelRequest, PgExecutionStore,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
@@ -41,6 +43,116 @@ fn model_request_rejects_mismatched_client_key_live_id() {
|
||||
assert!(request.validate().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn merged_first_attempt_insert_should_match_sequential_semantics() {
|
||||
let Some(database) = TestDatabase::create("execution_merged_insert").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgExecutionStore::new(database.pool.clone());
|
||||
let started_at = Utc::now();
|
||||
sqlx::query(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind,
|
||||
provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, enabled, availability,
|
||||
availability_observed_at, created_at, updated_at
|
||||
) values (
|
||||
'acct_merged', 'openai', 'merged', null, 'user-merged', null, null, 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1
|
||||
)",
|
||||
)
|
||||
.bind(started_at)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed provider account");
|
||||
let request = NewModelRequest {
|
||||
id: "req_merged".to_owned(),
|
||||
client_api_key_id: None,
|
||||
client_api_key_ref: "key_merged".to_owned(),
|
||||
config_revision: 1,
|
||||
protocol: "openai".to_owned(),
|
||||
operation: "responses".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
reasoning_preset: None,
|
||||
request_kind: None,
|
||||
subagent_kind: None,
|
||||
compact: false,
|
||||
image_generation_requested: false,
|
||||
started_at,
|
||||
deadline_at: started_at + Duration::seconds(30),
|
||||
};
|
||||
let attempt = ModelRequestAttemptStart {
|
||||
model_request_id: "req_merged".to_owned(),
|
||||
attempt_count: 1,
|
||||
provider_kind: "openai".to_owned(),
|
||||
provider_account_id: Some("acct_merged".to_owned()),
|
||||
provider_account_ref: Some("acct_merged".to_owned()),
|
||||
upstream_model_id: "gpt-test".to_owned(),
|
||||
upstream_transport: "http_sse".to_owned(),
|
||||
http_version: None,
|
||||
};
|
||||
|
||||
repository
|
||||
.insert_model_request_with_first_attempt(request, attempt)
|
||||
.await
|
||||
.expect("merged insert");
|
||||
|
||||
let (attempt_count, send_state, provider_kind, outcome): (i32, String, String, String) =
|
||||
sqlx::query_as(
|
||||
"select attempt_count, upstream_send_state, provider_kind, outcome
|
||||
from model_requests where id = 'req_merged'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load merged request");
|
||||
assert_eq!(
|
||||
(
|
||||
attempt_count,
|
||||
send_state.as_str(),
|
||||
provider_kind.as_str(),
|
||||
outcome.as_str()
|
||||
),
|
||||
(1, "not_sent", "openai", "running")
|
||||
);
|
||||
|
||||
// 后续 attempt 沿用常规 CAS 递增路径;已持久化的 sent 水位不被重试重置。
|
||||
repository
|
||||
.mark_upstream_send_state(
|
||||
"req_merged",
|
||||
gateway_store::postgres::UpstreamSendState::Sent,
|
||||
)
|
||||
.await
|
||||
.expect("mark sent before retry");
|
||||
let second = repository
|
||||
.begin_model_request_attempt(ModelRequestAttemptStart {
|
||||
model_request_id: "req_merged".to_owned(),
|
||||
attempt_count: 2,
|
||||
provider_kind: "openai".to_owned(),
|
||||
provider_account_id: Some("acct_merged".to_owned()),
|
||||
provider_account_ref: Some("acct_merged".to_owned()),
|
||||
upstream_model_id: "gpt-test".to_owned(),
|
||||
upstream_transport: "http_sse".to_owned(),
|
||||
http_version: None,
|
||||
})
|
||||
.await
|
||||
.expect("second attempt");
|
||||
assert_eq!(second, 2);
|
||||
let send_state: String = sqlx::query_scalar(
|
||||
"select upstream_send_state from model_requests where id = 'req_merged'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load send state after retry");
|
||||
assert_eq!(send_state, "sent");
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn downstream_commit_should_atomically_record_http_status_once() {
|
||||
let Some(database) = TestDatabase::create("execution_commit_status").await else {
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
use chrono::Utc;
|
||||
use gateway_store::postgres::ModelRequestHistoryRecord;
|
||||
|
||||
#[test]
|
||||
fn history_contains_no_payload_field() {
|
||||
let history = ModelRequestHistoryRecord {
|
||||
id: "request-1".to_owned(),
|
||||
client_api_key_ref: "key-1".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
provider_kind: None,
|
||||
provider_account_ref: None,
|
||||
upstream_model_id: None,
|
||||
outcome: "failed".to_owned(),
|
||||
client_response_id: None,
|
||||
upstream_response_id: None,
|
||||
started_at: Utc::now(),
|
||||
completed_at: None,
|
||||
};
|
||||
|
||||
assert_eq!(history.id, "request-1");
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::{env, str::FromStr};
|
||||
use std::str::FromStr;
|
||||
|
||||
use gateway_store::postgres::connect_and_migrate;
|
||||
use sqlx::{
|
||||
@@ -11,7 +11,6 @@ mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod history;
|
||||
mod observability;
|
||||
mod ops_events;
|
||||
mod provider_accounts;
|
||||
@@ -27,7 +26,7 @@ pub(super) struct TestDatabase {
|
||||
|
||||
impl TestDatabase {
|
||||
pub(super) async fn create(label: &str) -> Option<Self> {
|
||||
let database_url = env::var("CPR_TEST_DATABASE_URL").ok()?;
|
||||
let database_url = crate::support::test_env("CPR_TEST_DATABASE_URL")?;
|
||||
let schema = format!("cpr_store_{label}_{}", Uuid::new_v4().simple());
|
||||
let admin = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
@@ -59,6 +58,12 @@ 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,
|
||||
@@ -82,7 +87,7 @@ impl TestDatabase {
|
||||
|
||||
#[tokio::test]
|
||||
async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
|
||||
let Ok(database_url) = env::var("CPR_TEST_DATABASE_URL") else {
|
||||
let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else {
|
||||
return;
|
||||
};
|
||||
let database = format!("cpr_store_migrator_{}", Uuid::new_v4().simple());
|
||||
@@ -103,7 +108,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
|
||||
.database(&database)
|
||||
.to_url_lossy()
|
||||
.to_string();
|
||||
let first = connect_and_migrate(&isolated_url)
|
||||
let first = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default())
|
||||
.await
|
||||
.expect("apply 0001 through production migrator");
|
||||
let first_table_count = sqlx::query_scalar::<_, i64>(
|
||||
@@ -114,7 +119,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
|
||||
.expect("count migrated tables");
|
||||
first.close().await;
|
||||
|
||||
let second = connect_and_migrate(&isolated_url)
|
||||
let second = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default())
|
||||
.await
|
||||
.expect("reopen database through production migrator");
|
||||
let migration_count =
|
||||
@@ -132,7 +137,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, 1));
|
||||
assert_eq!((first_table_count, migration_count), (8, 2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -62,6 +62,57 @@ fn postgres_admin_observability_adapter_implements_terminal_port() {
|
||||
assert_port::<PgAdminObservabilityStore>();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_account_metrics_should_partition_available_and_unavailable_accounts() {
|
||||
let Some(database) = TestDatabase::create("dashboard_account_metrics").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
sqlx::query(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, upstream_user_id, authentication_kind,
|
||||
provider_credentials_json, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, enabled, availability,
|
||||
availability_observed_at, created_at, updated_at
|
||||
) values
|
||||
('acct_available', 'openai', 'available', 'user-available', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1),
|
||||
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', $1, $1, $1),
|
||||
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', $1, $1, $1),
|
||||
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', $1, $1, $1)",
|
||||
)
|
||||
.bind(now)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed account metric states");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("dashboard range");
|
||||
|
||||
let metrics = repository
|
||||
.dashboard_summary(range)
|
||||
.await
|
||||
.expect("dashboard summary")
|
||||
.provider_accounts;
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
metrics.total,
|
||||
metrics.enabled,
|
||||
metrics.active,
|
||||
metrics.unavailable,
|
||||
metrics.expired,
|
||||
metrics.disabled,
|
||||
metrics.banned,
|
||||
),
|
||||
(4, 3, 1, 3, 1, 1, 1),
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
let Some(database) = TestDatabase::create("calculated_usage_billing_facts").await else {
|
||||
@@ -123,7 +174,9 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.await
|
||||
.expect("admin dashboard summary");
|
||||
assert_eq!(dashboard.range, range);
|
||||
assert_eq!(dashboard.requests.request_count, 1);
|
||||
assert_eq!(dashboard.requests.request_count, 3);
|
||||
assert_eq!(dashboard.requests.success_count, 2);
|
||||
assert_eq!(dashboard.requests.failure_count, 1);
|
||||
assert_eq!(dashboard.requests.first_token_latency_sum_ms, 120);
|
||||
assert_eq!(dashboard.requests.latency_sum_ms, 900);
|
||||
assert_eq!(dashboard.requests.min_latency_ms, Some(900));
|
||||
@@ -166,6 +219,13 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
3,
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard_trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.failure_count)
|
||||
.sum::<u64>(),
|
||||
1,
|
||||
);
|
||||
|
||||
@@ -184,7 +244,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
1,
|
||||
2,
|
||||
);
|
||||
assert_eq!(
|
||||
trend
|
||||
@@ -391,7 +451,10 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.dashboard_summary(range)
|
||||
.await
|
||||
.expect("dashboard summary");
|
||||
assert_eq!(dashboard.requests.request_count, 1);
|
||||
assert_eq!(dashboard.requests.request_count, 3);
|
||||
assert_eq!(dashboard.requests.success_count, 2);
|
||||
assert_eq!(dashboard.requests.failure_count, 1);
|
||||
assert_eq!(dashboard.requests.caller_error_count, 0);
|
||||
assert_eq!(dashboard.requests.cache_eligible_request_count, 1);
|
||||
assert_eq!(dashboard.requests.cache_hit_request_count, 1);
|
||||
assert_eq!(dashboard.requests.cache_hit_request_rate(), Some(1.0));
|
||||
@@ -450,6 +513,14 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
3
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.failure_count)
|
||||
.sum::<u64>(),
|
||||
1
|
||||
);
|
||||
|
||||
@@ -541,7 +612,8 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("usage summary");
|
||||
assert_eq!(overview.attempts.attempt_count, 4);
|
||||
assert_eq!(overview.attempts.failure_count, 2);
|
||||
assert_eq!(overview.requests.request_count, 1);
|
||||
assert_eq!(overview.requests.request_count, 3);
|
||||
assert_eq!(overview.requests.failure_count, 1);
|
||||
|
||||
let succeeded = repository
|
||||
.usage_summary(
|
||||
@@ -589,6 +661,20 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.await
|
||||
.expect("ops errors");
|
||||
assert_eq!(errors.total, 2);
|
||||
let request_error = errors
|
||||
.items
|
||||
.iter()
|
||||
.find(|error| error.source == "model_request")
|
||||
.expect("request error");
|
||||
assert_eq!(request_error.client_status_code, Some(502));
|
||||
assert_eq!(request_error.upstream_status_code, Some(429));
|
||||
let attempt_error = errors
|
||||
.items
|
||||
.iter()
|
||||
.find(|error| error.source == "ops_event")
|
||||
.expect("attempt error");
|
||||
assert_eq!(attempt_error.client_status_code, None);
|
||||
assert_eq!(attempt_error.upstream_status_code, Some(429));
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -126,13 +126,13 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
beta.availability = ProviderAccountAvailability::Banned;
|
||||
let mut charlie = account("acct_charlie", "user-charlie");
|
||||
charlie.email = Some("charlie@example.invalid".to_owned());
|
||||
let mut attention = account("acct_attention", "user-attention");
|
||||
attention.email = Some("attention@example.invalid".to_owned());
|
||||
attention.availability = ProviderAccountAvailability::Invalid;
|
||||
let mut invalid = account("acct_invalid", "user-invalid");
|
||||
invalid.email = Some("invalid@example.invalid".to_owned());
|
||||
invalid.availability = ProviderAccountAvailability::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;
|
||||
for account in [alpha, beta, charlie, attention, quota_exhausted] {
|
||||
for account in [alpha, beta, charlie, invalid, quota_exhausted] {
|
||||
repository
|
||||
.insert_provider_account(account)
|
||||
.await
|
||||
@@ -214,9 +214,8 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
assert_eq!(usage_page.config_revision.get(), 1);
|
||||
assert_eq!(usage_page.total, 5);
|
||||
assert_eq!(usage_page.summary.total, 5);
|
||||
assert_eq!(usage_page.summary.active, 2);
|
||||
assert_eq!(usage_page.summary.active, 3);
|
||||
assert_eq!(usage_page.summary.quota_exhausted, 1);
|
||||
assert_eq!(usage_page.summary.attention, 1);
|
||||
assert_eq!(
|
||||
usage_page
|
||||
.items
|
||||
@@ -277,19 +276,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
.await
|
||||
.expect("filter banned accounts");
|
||||
assert_eq!(banned.items[0].id, "acct_beta");
|
||||
let attention = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::Attention),
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("filter attention accounts");
|
||||
assert_eq!(attention.items[0].id, "acct_attention");
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
|
||||
@@ -283,7 +283,7 @@ fn running_request(
|
||||
}
|
||||
|
||||
async fn repository() -> Option<(RedisClientAdmissionRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -536,7 +536,7 @@ fn millisecond_precision(value: DateTime<Utc>) -> DateTime<Utc> {
|
||||
}
|
||||
|
||||
async fn repository() -> Option<(RedisCredentialCooldownRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -404,7 +404,7 @@ fn acquired(
|
||||
}
|
||||
|
||||
async fn repository() -> Option<(RedisCredentialLeaseRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -32,25 +32,14 @@ fn admin_auth_state_rejects_invalid_ttl_boundaries() {
|
||||
.block_on(repository.store_admin_session("expired-session", &session))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.block_on(repository.record_login_failure("source", 0, 60))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.block_on(repository.record_login_failure("source", 5, 0))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_state_keeps_fixed_ttl_and_atomic_failure_window() {
|
||||
async fn admin_auth_state_keeps_fixed_ttl_and_opaque_keys() {
|
||||
let Some((repository, mut connection, namespace)) = admin_auth_repository().await else {
|
||||
return;
|
||||
};
|
||||
let session_id = "session-real-secret-id";
|
||||
let source = "203.0.113.7";
|
||||
let password = "must-never-enter-redis";
|
||||
let admin_api_key = "admin-must-never-enter-redis";
|
||||
let session = AdminSessionRecord {
|
||||
@@ -70,29 +59,14 @@ async fn admin_auth_state_keeps_fixed_ttl_and_atomic_failure_window() {
|
||||
Some(session.clone())
|
||||
);
|
||||
|
||||
for attempt in 1..=5 {
|
||||
let throttled = repository
|
||||
.record_login_failure(source, 5, 60)
|
||||
.await
|
||||
.expect("record failure");
|
||||
assert_eq!(throttled, attempt == 5);
|
||||
}
|
||||
assert!(
|
||||
repository
|
||||
.login_source_is_throttled(source, 5, 60)
|
||||
.await
|
||||
.expect("read failure window")
|
||||
);
|
||||
|
||||
let keys = redis::cmd("KEYS")
|
||||
.arg(format!("{namespace}:*"))
|
||||
.query_async::<Vec<String>>(&mut connection)
|
||||
.await
|
||||
.expect("list isolated test keys");
|
||||
assert_eq!(keys.len(), 2);
|
||||
assert_eq!(keys.len(), 1);
|
||||
for key in &keys {
|
||||
assert!(!key.contains(session_id));
|
||||
assert!(!key.contains(source));
|
||||
let ttl = redis::cmd("PTTL")
|
||||
.arg(key)
|
||||
.query_async::<i64>(&mut connection)
|
||||
@@ -105,21 +79,10 @@ async fn admin_auth_state_keeps_fixed_ttl_and_atomic_failure_window() {
|
||||
.await
|
||||
.expect("read isolated test value");
|
||||
assert!(!value.contains(session_id));
|
||||
assert!(!value.contains(source));
|
||||
assert!(!value.contains(password));
|
||||
assert!(!value.contains(admin_api_key));
|
||||
}
|
||||
|
||||
repository
|
||||
.clear_login_failures(source)
|
||||
.await
|
||||
.expect("clear failure window");
|
||||
assert!(
|
||||
!repository
|
||||
.login_source_is_throttled(source, 5, 60)
|
||||
.await
|
||||
.expect("read cleared failure window")
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.delete_admin_session(session_id)
|
||||
@@ -138,7 +101,7 @@ async fn admin_auth_state_keeps_fixed_ttl_and_atomic_failure_window() {
|
||||
|
||||
async fn admin_auth_repository()
|
||||
-> Option<(RedisAdminAuthStateRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -91,7 +91,7 @@ async fn native_continuation_rejects_provider_state_for_a_different_provider() {
|
||||
}
|
||||
|
||||
async fn repository() -> Option<(RedisNativeContinuationRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -132,7 +132,7 @@ async fn repository() -> Option<(
|
||||
ConnectionManager,
|
||||
String,
|
||||
)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).ok()?;
|
||||
let connection: ConnectionManager = client.get_connection_manager().await.ok()?;
|
||||
let namespace = format!("gateway-store-test-{}", Uuid::new_v4());
|
||||
|
||||
@@ -157,7 +157,7 @@ async fn repository_with_policy(
|
||||
threshold: u32,
|
||||
open_duration: Duration,
|
||||
) -> Option<(RedisProviderCircuitRepository, ConnectionManager, String)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -113,7 +113,7 @@ async fn affinity_repository() -> Option<(
|
||||
ConnectionManager,
|
||||
String,
|
||||
)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -92,7 +92,7 @@ async fn repository() -> Option<(
|
||||
ConnectionManager,
|
||||
String,
|
||||
)> {
|
||||
let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?;
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
|
||||
@@ -15,8 +15,8 @@ use uuid::Uuid;
|
||||
#[tokio::test]
|
||||
async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() {
|
||||
let (Some(database_url), Some(redis_url)) = (
|
||||
std::env::var("CPR_TEST_DATABASE_URL").ok(),
|
||||
std::env::var("CPR_TEST_REDIS_URL").ok(),
|
||||
crate::support::test_env("CPR_TEST_DATABASE_URL"),
|
||||
crate::support::test_env("CPR_TEST_REDIS_URL"),
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
/// 读取集成测试基础设施地址(PostgreSQL / Redis)。
|
||||
///
|
||||
/// 本地未配置时返回 `None`,测试跳过;CI 中缺失则直接失败——
|
||||
/// 静默跳过产生的绿灯与真实回归不可区分。
|
||||
pub(crate) fn test_env(name: &str) -> Option<String> {
|
||||
let value = std::env::var(name).ok();
|
||||
assert!(
|
||||
value.is_some() || !running_in_ci(),
|
||||
"{name} must be set in CI: store integration tests are not allowed to skip silently"
|
||||
);
|
||||
value
|
||||
}
|
||||
|
||||
/// 按 is-ci 约定判定 CI:`CI` 存在且非 falsey(`0`/`false`/空)。
|
||||
/// 开发者用 `CI=false` 显式声明「非 CI」时不应把跳过变成 panic。
|
||||
fn running_in_ci() -> bool {
|
||||
std::env::var("CI")
|
||||
.ok()
|
||||
.is_some_and(|value| !matches!(value.as_str(), "" | "0" | "false"))
|
||||
}
|
||||
@@ -43,7 +43,6 @@ tracing = "0.1.44"
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f", features = ["deflate", "proxy"] }
|
||||
url = "2.5.8"
|
||||
uuid = { version = "1.23.3", features = ["v4", "v7"] }
|
||||
zeroize = { version = "1.8.2", features = ["derive"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.27.0"
|
||||
|
||||
@@ -658,6 +658,8 @@ impl CodexCredentialQuotaService {
|
||||
if outcome == QuotaWriteOutcome::Conflict {
|
||||
return Ok(false);
|
||||
}
|
||||
// 可用性写回必须基于最新账号状态:revision fence 只覆盖凭据轮换,
|
||||
// 不覆盖其他路径的仅状态变化(Cooldown/QuotaExhausted 转换)。
|
||||
let Some(current) = self.store.get_account(account.id()).await? else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
@@ -8,9 +8,9 @@ use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt as _, stream};
|
||||
use gateway_core::engine::credential::{AccountAvailability, ProviderAccount, ProviderAccountId};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
ProviderCredentialStatePort, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
|
||||
ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_retry_at,
|
||||
ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_backoff_at,
|
||||
};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use thiserror::Error;
|
||||
@@ -25,7 +25,8 @@ use super::types::{CodexCredentialPrincipal, CodexOAuthSecret, RotateCodexCreden
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
const MAX_REFRESH_BATCH: u32 = 1_000;
|
||||
const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60);
|
||||
/// 连续失败计数窗口;每次瞬态失败刷新该 TTL,静默满窗后计数过期归零。
|
||||
const REFRESH_BACKOFF_WINDOW: Duration = Duration::from_secs(30 * 60);
|
||||
|
||||
pub struct DueCodexCredential {
|
||||
pub account: ProviderAccount,
|
||||
@@ -59,9 +60,6 @@ pub enum CodexCredentialRefreshOutcome {
|
||||
Transient {
|
||||
account_id: String,
|
||||
},
|
||||
Ambiguous {
|
||||
account_id: String,
|
||||
},
|
||||
LeaseUnavailable {
|
||||
account_id: String,
|
||||
},
|
||||
@@ -88,6 +86,7 @@ pub struct CodexCredentialRefreshService {
|
||||
refresher: Arc<dyn TokenRefresher>,
|
||||
identity: Arc<dyn CodexAccountIdentityVerifier>,
|
||||
leases: Arc<dyn ProviderLeasePort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
}
|
||||
|
||||
@@ -97,6 +96,7 @@ impl CodexCredentialRefreshService {
|
||||
refresher: Arc<dyn TokenRefresher>,
|
||||
identity: Arc<dyn CodexAccountIdentityVerifier>,
|
||||
leases: Arc<dyn ProviderLeasePort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
) -> Self {
|
||||
Self {
|
||||
@@ -104,6 +104,7 @@ impl CodexCredentialRefreshService {
|
||||
refresher,
|
||||
identity,
|
||||
leases,
|
||||
credential_state,
|
||||
runtime_policy,
|
||||
}
|
||||
}
|
||||
@@ -206,13 +207,16 @@ impl CodexCredentialRefreshService {
|
||||
}
|
||||
}
|
||||
Err(RefreshFailure::Transport) => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
AccountAvailability::Invalid,
|
||||
"refresh_ambiguous",
|
||||
CodexCredentialRefreshOutcome::Ambiguous { account_id },
|
||||
)
|
||||
.await
|
||||
// 上游瞬态(429/5xx/超时/畸形响应等,发送状态未知)保留现有凭据、
|
||||
// 推进退避重试;仅 invalid_grant/banned 才终态失效账号。
|
||||
if self
|
||||
.defer_refresh(&due.account, "transport-ambiguous")
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient { account_id })
|
||||
} else {
|
||||
Ok(CodexCredentialRefreshOutcome::Stale { account_id })
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -326,25 +330,34 @@ impl CodexCredentialRefreshService {
|
||||
.await;
|
||||
}
|
||||
Err(CodexIdentityVerificationError::Unavailable) => {
|
||||
return self
|
||||
.persist_terminal(
|
||||
&due.account,
|
||||
AccountAvailability::Invalid,
|
||||
"refreshed_identity_unavailable",
|
||||
CodexCredentialRefreshOutcome::Ambiguous {
|
||||
account_id: due.account.id().to_string(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
// JWKS/签名边界短暂不可用属瞬态:保留现有凭据、推进退避重试,
|
||||
// 不终态失效账号。
|
||||
return if self
|
||||
.defer_refresh(&due.account, "identity-unavailable")
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient {
|
||||
account_id: due.account.id().to_string(),
|
||||
})
|
||||
} else {
|
||||
Ok(CodexCredentialRefreshOutcome::Stale {
|
||||
account_id: due.account.id().to_string(),
|
||||
})
|
||||
};
|
||||
}
|
||||
};
|
||||
let profile = match verification {
|
||||
CodexIdentityVerification::Complete(profile) => profile,
|
||||
CodexIdentityVerification::SignedOnly(signed) => {
|
||||
let retry_at = provider_refresh_retry_at(
|
||||
let attempt = self
|
||||
.credential_state
|
||||
.record_refresh_backoff(due.account.id(), REFRESH_BACKOFF_WINDOW)
|
||||
.await
|
||||
.unwrap_or(1);
|
||||
let retry_at = provider_refresh_backoff_at(
|
||||
due.account.id(),
|
||||
SystemTime::now(),
|
||||
REFRESH_RETRY_DELAY,
|
||||
attempt,
|
||||
"identity-completion",
|
||||
)?;
|
||||
match self
|
||||
@@ -388,10 +401,17 @@ impl CodexCredentialRefreshService {
|
||||
})
|
||||
.await;
|
||||
match result {
|
||||
Ok(revision) => Ok(CodexCredentialRefreshOutcome::Refreshed {
|
||||
account_id: due.account.id().to_string(),
|
||||
credential_revision: revision.get(),
|
||||
}),
|
||||
Ok(revision) => {
|
||||
// 凭据完整成功轮换后清零连续失败计数,退避窗口重新从 base 起步。
|
||||
let _ = self
|
||||
.credential_state
|
||||
.clear_refresh_backoff(due.account.id())
|
||||
.await;
|
||||
Ok(CodexCredentialRefreshOutcome::Refreshed {
|
||||
account_id: due.account.id().to_string(),
|
||||
credential_revision: revision.get(),
|
||||
})
|
||||
}
|
||||
Err(CredentialRepositoryError::RevisionConflict) => {
|
||||
Ok(CodexCredentialRefreshOutcome::Stale {
|
||||
account_id: due.account.id().to_string(),
|
||||
@@ -443,12 +463,13 @@ impl CodexCredentialRefreshService {
|
||||
account: &ProviderAccount,
|
||||
reason: &'static str,
|
||||
) -> Result<bool, CodexCredentialRefreshError> {
|
||||
let retry_at = provider_refresh_retry_at(
|
||||
account.id(),
|
||||
SystemTime::now(),
|
||||
REFRESH_RETRY_DELAY,
|
||||
reason,
|
||||
)?;
|
||||
let attempt = self
|
||||
.credential_state
|
||||
.record_refresh_backoff(account.id(), REFRESH_BACKOFF_WINDOW)
|
||||
.await
|
||||
.unwrap_or(1);
|
||||
let retry_at =
|
||||
provider_refresh_backoff_at(account.id(), SystemTime::now(), attempt, reason)?;
|
||||
match self.repository.defer_refresh(account, retry_at).await {
|
||||
Ok(_) => Ok(true),
|
||||
Err(CredentialRepositoryError::RevisionConflict) => Ok(false),
|
||||
|
||||
@@ -307,7 +307,13 @@ fn parse_token_pair(body: &[u8]) -> Result<TokenPair, ()> {
|
||||
})
|
||||
}
|
||||
|
||||
fn classify_refresh_failure(_status: StatusCode, body: &[u8]) -> RefreshFailure {
|
||||
fn classify_refresh_failure(status: StatusCode, body: &[u8]) -> RefreshFailure {
|
||||
// 5xx 与 429 按状态码判为瞬态:其正文(CDN/网关页等)不是权威 OAuth 错误响应,
|
||||
// 不据此终态失效账号。
|
||||
if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS {
|
||||
return RefreshFailure::Transport;
|
||||
}
|
||||
// 4xx 正文才是权威 OAuth 错误响应,据其区分真过期/封禁与其余瞬态。
|
||||
let lower = String::from_utf8_lossy(body).to_ascii_lowercase();
|
||||
if lower.contains("account has been deactivated") {
|
||||
return RefreshFailure::Banned;
|
||||
|
||||
@@ -64,6 +64,7 @@ pub async fn initialize(
|
||||
let session_exclusions = ports.session_exclusions();
|
||||
let account_feedback = ports.account_feedback();
|
||||
let runtime_policy = ports.runtime_policy();
|
||||
let credential_state = ports.credential_state();
|
||||
let profile = config.wire_profile_state();
|
||||
let session_identity = config
|
||||
.session_identity()
|
||||
@@ -170,6 +171,7 @@ pub async fn initialize(
|
||||
refresher,
|
||||
Arc::clone(&identity),
|
||||
Arc::clone(&leases),
|
||||
credential_state,
|
||||
Arc::clone(&runtime_policy),
|
||||
));
|
||||
let pending = Arc::new(OpenAiOAuthPendingStore::new(
|
||||
|
||||
@@ -818,10 +818,12 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let failure_rate_limit_headers = response.rate_limit_headers.clone();
|
||||
let rate_limit_updates = response.rate_limit_header_updates;
|
||||
let turn_state_updates = response.turn_state_update;
|
||||
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str());
|
||||
if response.transport == CodexBackendTransport::HttpSse {
|
||||
decoder = decoder.with_raw_sse_passthrough();
|
||||
}
|
||||
// OpenAI 线路为透明代理:HTTP SSE 与 WebSocket 两条上游均启用 raw 透传,
|
||||
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
|
||||
// WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON
|
||||
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
|
||||
let mut decoder =
|
||||
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
|
||||
loop {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
|
||||
@@ -877,7 +879,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
};
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
&& let Some(turn_state) = updates.lock().await.clone()
|
||||
// take:仅在上游写入过新 turn_state 时更新 capture,避免逐 chunk 克隆。
|
||||
&& let Some(turn_state) = updates.lock().await.take()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
@@ -994,7 +997,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
}
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
&& let Some(turn_state) = updates.lock().await.clone()
|
||||
&& let Some(turn_state) = updates.lock().await.take()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
@@ -2194,7 +2197,6 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
|
||||
matches!(
|
||||
outcome,
|
||||
CodexCredentialRefreshOutcome::Transient { .. }
|
||||
| CodexCredentialRefreshOutcome::Ambiguous { .. }
|
||||
| CodexCredentialRefreshOutcome::Failed { .. }
|
||||
)
|
||||
})
|
||||
|
||||
@@ -31,7 +31,7 @@ use crate::transport::{
|
||||
CodexWebSocketConnection, CodexWebSocketPool, CodexWebSocketPoolKey,
|
||||
DEFAULT_INITIAL_EVENT_TIMEOUT, WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker,
|
||||
WebSocketPoolDecision, execute_prepared_response_create_request_stream,
|
||||
post_send_ambiguous, prepare_response_create_request_with_pool,
|
||||
post_send_ambiguous, prepare_response_create_request_with_pool, websocket_audit_dir,
|
||||
write_websocket_audit_artifact_from_env,
|
||||
},
|
||||
};
|
||||
@@ -176,13 +176,17 @@ impl CodexBackendClient {
|
||||
&websocket_request,
|
||||
)
|
||||
.map_err(CodexClientError::WebSocketEncode)?;
|
||||
let artifact = websocket_audit_artifact_from_attempt(
|
||||
&websocket_request,
|
||||
websocket_create.connection().opening_audit_snapshot(),
|
||||
websocket_payload_audit_snapshot(&websocket_request),
|
||||
);
|
||||
if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await {
|
||||
tracing::warn!(error = %error, "Failed to write Codex WebSocket audit artifact");
|
||||
// 审计未启用时跳过 artifact 构造:payload 快照会深拷贝整个请求 body,
|
||||
// 且位于首字节前的关键路径上。
|
||||
if websocket_audit_dir().is_some() {
|
||||
let artifact = websocket_audit_artifact_from_attempt(
|
||||
&websocket_request,
|
||||
websocket_create.connection().opening_audit_snapshot(),
|
||||
websocket_payload_audit_snapshot(&websocket_request),
|
||||
);
|
||||
if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await {
|
||||
tracing::warn!(error = %error, "Failed to write Codex WebSocket audit artifact");
|
||||
}
|
||||
}
|
||||
let connection_profile = websocket_connection_profile(&headers);
|
||||
let pool_key =
|
||||
|
||||
@@ -366,21 +366,6 @@ fn is_tool_execution_event(event_type: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
/// 从 Codex SSE 收集出的非流式 Responses 结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum CollectedResponse {
|
||||
/// 收集到 `response.completed`。
|
||||
Completed(Value),
|
||||
/// 收集到 `response.incomplete`。
|
||||
Incomplete(Value),
|
||||
/// 收集到 `response.failed` 或 `error`。
|
||||
Failed(ResponsesSseFailure),
|
||||
/// SSE 未包含成功终止响应。
|
||||
MissingCompleted,
|
||||
/// 完成响应为空。
|
||||
Empty,
|
||||
}
|
||||
|
||||
/// Codex Responses SSE 失败事件。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct ResponsesSseFailure {
|
||||
|
||||
@@ -50,43 +50,32 @@ pub struct WebSocketAuditArtifact {
|
||||
|
||||
/// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。
|
||||
pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
|
||||
let event = websocket_event_type(raw)?;
|
||||
if is_internal_websocket_event(&event) {
|
||||
return None;
|
||||
}
|
||||
if websocket_event_should_skip(raw) {
|
||||
return None;
|
||||
}
|
||||
Some(encode_sse_event(&event, raw))
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
websocket_event_frame(&value, raw)
|
||||
}
|
||||
|
||||
pub(crate) fn websocket_event_type(raw: &str) -> Option<String> {
|
||||
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
|
||||
value
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string)
|
||||
})
|
||||
/// 同上,复用已解析的事件 JSON;`data` 段逐字节内嵌 `raw` 原文。
|
||||
///
|
||||
/// 仅剥离传输层内部帧:`codex.rate_limits` 与承载 turn_state 的
|
||||
/// `response.metadata` 不下发客户端;业务事件一律原样转发,不做 schema 校验。
|
||||
pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option<String> {
|
||||
let event = websocket_event_type(value)?;
|
||||
if is_internal_websocket_event(event) || event == "response.metadata" {
|
||||
return None;
|
||||
}
|
||||
Some(encode_sse_event(event, raw))
|
||||
}
|
||||
|
||||
pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> {
|
||||
value.get("type").and_then(Value::as_str)
|
||||
}
|
||||
|
||||
fn is_internal_websocket_event(event: &str) -> bool {
|
||||
event == "codex.rate_limits"
|
||||
}
|
||||
|
||||
/// 判断 WebSocket 事件是否应在转发前剥离。
|
||||
///
|
||||
/// 仅剥离传输层内部帧:`response.metadata` 承载 turn_state,由上层提取转存到
|
||||
/// 会话状态,不下发客户端。业务事件一律原样转发,不做 schema 校验。
|
||||
/// (`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。)
|
||||
fn websocket_event_should_skip(raw: &str) -> bool {
|
||||
websocket_event_type(raw).as_deref() == Some("response.metadata")
|
||||
}
|
||||
|
||||
/// 提取 `response.metadata` 帧中的字符串响应头。
|
||||
pub fn websocket_metadata_headers(raw: &str) -> Vec<(String, String)> {
|
||||
let Ok(value) = serde_json::from_str::<Value>(raw) else {
|
||||
return Vec::new();
|
||||
};
|
||||
pub fn websocket_metadata_headers(value: &Value) -> Vec<(String, String)> {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.metadata") {
|
||||
return Vec::new();
|
||||
}
|
||||
@@ -100,9 +89,8 @@ pub fn websocket_metadata_headers(raw: &str) -> Vec<(String, String)> {
|
||||
}
|
||||
|
||||
/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。
|
||||
pub fn websocket_metadata_turn_state(raw: &str) -> Option<String> {
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
websocket_metadata_headers(raw)
|
||||
pub fn websocket_metadata_turn_state(value: &Value) -> Option<String> {
|
||||
websocket_metadata_headers(value)
|
||||
.into_iter()
|
||||
.find_map(|(name, value)| {
|
||||
name.eq_ignore_ascii_case("x-codex-turn-state")
|
||||
@@ -128,8 +116,7 @@ fn json_value_as_string(value: &Value) -> Option<String> {
|
||||
/// 从 `response.completed` 旁路提取 response ID,供连接池记录续接能力。
|
||||
///
|
||||
/// 该值不参与客户端 wire 的可交付性判断;无法读取时只是不记录连接内续接状态。
|
||||
pub fn websocket_response_completed_id(raw: &str) -> Option<String> {
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.completed") {
|
||||
return None;
|
||||
}
|
||||
@@ -179,11 +166,39 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val
|
||||
Value::Object(payload)
|
||||
}
|
||||
|
||||
/// 借用原始 body 序列化 `response.create` 帧,不复制字段。
|
||||
///
|
||||
/// 输出与序列化 [`websocket_response_create_payload`] 的合并 Map 逐字节一致:
|
||||
/// body 自带 `type` 键时以 body 值置于首位,其余字段按插入顺序原样透传。
|
||||
struct ResponseCreateFrame<'a> {
|
||||
body: &'a Map<String, Value>,
|
||||
}
|
||||
|
||||
impl serde::Serialize for ResponseCreateFrame<'_> {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
use serde::ser::SerializeMap as _;
|
||||
let mut map = serializer.serialize_map(Some(self.body.len() + 1))?;
|
||||
match self.body.get("type") {
|
||||
Some(body_type) => map.serialize_entry("type", body_type)?,
|
||||
None => map.serialize_entry("type", "response.create")?,
|
||||
}
|
||||
for (key, value) in self.body {
|
||||
if key == "type" {
|
||||
continue;
|
||||
}
|
||||
map.serialize_entry(key, value)?;
|
||||
}
|
||||
map.end()
|
||||
}
|
||||
}
|
||||
|
||||
/// 生成 Responses WebSocket `response.create` 文本帧内容。
|
||||
pub fn websocket_response_create_payload_text(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Result<String, serde_json::Error> {
|
||||
serde_json::to_string(&websocket_response_create_payload(request))
|
||||
serde_json::to_string(&ResponseCreateFrame {
|
||||
body: request.body(),
|
||||
})
|
||||
}
|
||||
|
||||
fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use std::{
|
||||
io,
|
||||
path::{Path, PathBuf},
|
||||
sync::OnceLock,
|
||||
};
|
||||
|
||||
use chrono::Utc;
|
||||
@@ -19,6 +20,19 @@ const REDACTED_HEADER_VALUE: &str = "<redacted>";
|
||||
/// WebSocket audit artifact 输出目录环境变量。
|
||||
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
|
||||
|
||||
static WS_AUDIT_DIR: OnceLock<Option<PathBuf>> = OnceLock::new();
|
||||
|
||||
/// 返回进程级审计输出目录;未配置时调用方应整体跳过 artifact 构造。
|
||||
pub fn websocket_audit_dir() -> Option<&'static Path> {
|
||||
WS_AUDIT_DIR
|
||||
.get_or_init(|| {
|
||||
std::env::var_os(WS_AUDIT_DIR_ENV)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(PathBuf::from)
|
||||
})
|
||||
.as_deref()
|
||||
}
|
||||
|
||||
/// 显式写入 WebSocket audit artifact。
|
||||
pub async fn write_websocket_audit_artifact_for_dir(
|
||||
dir: Option<&Path>,
|
||||
@@ -39,14 +53,7 @@ pub async fn write_websocket_audit_artifact_for_dir(
|
||||
pub async fn write_websocket_audit_artifact_from_env(
|
||||
artifact: &WebSocketAuditArtifact,
|
||||
) -> io::Result<Option<PathBuf>> {
|
||||
let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(PathBuf::from)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await
|
||||
write_websocket_audit_artifact_for_dir(websocket_audit_dir(), artifact).await
|
||||
}
|
||||
|
||||
impl CodexWebSocketConnection {
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
//! WebSocket aggregate/stream 共用事件归约器。
|
||||
|
||||
use gateway_protocol::openai::events;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::transport::protocol::websocket::{
|
||||
websocket_event_to_sse_frame, websocket_event_type, websocket_metadata_headers,
|
||||
websocket_event_frame, websocket_event_type, websocket_metadata_headers,
|
||||
websocket_metadata_turn_state, websocket_response_completed_id,
|
||||
};
|
||||
use crate::transport::response_meta;
|
||||
@@ -33,40 +34,41 @@ pub(super) fn reduce_websocket_event(
|
||||
metadata: &mut CodexWebSocketConnectionMetadata,
|
||||
continuation: &mut WebSocketContinuationState,
|
||||
) -> Result<ExchangeAction, CodexWebSocketExchangeError> {
|
||||
if let Some(headers) = websocket_rate_limit_event_headers(raw) {
|
||||
// 每帧只解析一次 JSON,后续提取全部复用同一 Value;
|
||||
// 不可解析的帧不承载可路由的事件类型,忽略。
|
||||
let Ok(value) = serde_json::from_str::<Value>(raw) else {
|
||||
return Ok(ExchangeAction::Ignore);
|
||||
};
|
||||
if let Some(parsed) = events::parse_rate_limits_event(&value) {
|
||||
let headers = events::rate_limits_to_header_pairs(&parsed);
|
||||
metadata.rate_limit_headers.extend(headers.iter().cloned());
|
||||
return Ok(ExchangeAction::RateLimits(headers));
|
||||
}
|
||||
|
||||
response_meta::merge_response_metadata(
|
||||
&mut metadata.response_metadata,
|
||||
websocket_metadata_headers(raw),
|
||||
websocket_metadata_headers(&value),
|
||||
);
|
||||
if let Some(turn_state) = websocket_metadata_turn_state(raw) {
|
||||
if let Some(turn_state) = websocket_metadata_turn_state(&value) {
|
||||
metadata.turn_state = Some(turn_state.clone());
|
||||
return Ok(ExchangeAction::TurnState(turn_state));
|
||||
}
|
||||
|
||||
let event = websocket_event_type(raw);
|
||||
if event.as_deref() == Some("response.completed")
|
||||
&& let Some(response_id) = websocket_response_completed_id(raw)
|
||||
let event = websocket_event_type(&value);
|
||||
if event == Some("response.completed")
|
||||
&& let Some(response_id) = websocket_response_completed_id(&value)
|
||||
{
|
||||
continuation.record_completed(response_id);
|
||||
}
|
||||
|
||||
let terminal = match event.as_deref() {
|
||||
let terminal = match event {
|
||||
Some("response.completed") => Some(WebSocketTerminalKind::Completed),
|
||||
Some("response.incomplete") => Some(WebSocketTerminalKind::Incomplete),
|
||||
Some("response.failed" | "error") => Some(WebSocketTerminalKind::Failed),
|
||||
_ => None,
|
||||
};
|
||||
Ok(match websocket_event_to_sse_frame(raw) {
|
||||
Ok(match websocket_event_frame(&value, raw) {
|
||||
Some(frame) => ExchangeAction::Forward { frame, terminal },
|
||||
None => ExchangeAction::Ignore,
|
||||
})
|
||||
}
|
||||
|
||||
fn websocket_rate_limit_event_headers(raw: &str) -> Option<Vec<(String, String)>> {
|
||||
events::parse_rate_limits_event_raw(raw)
|
||||
.map(|parsed| events::rate_limits_to_header_pairs(&parsed))
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@ mod pump;
|
||||
|
||||
pub use self::{
|
||||
audit::{
|
||||
WS_AUDIT_DIR_ENV, write_websocket_audit_artifact_for_dir,
|
||||
WS_AUDIT_DIR_ENV, websocket_audit_dir, write_websocket_audit_artifact_for_dir,
|
||||
write_websocket_audit_artifact_from_env,
|
||||
},
|
||||
breaker::{
|
||||
|
||||
@@ -764,6 +764,21 @@ impl ProviderCredentialStatePort for TestCredentialState {
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn record_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_window: Duration,
|
||||
) -> BoxFuture<'a, Result<u32, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(1) })
|
||||
}
|
||||
|
||||
fn clear_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
struct TestCooldown;
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::credential::AccountAvailability;
|
||||
use gateway_core::engine::credential::{AccountAvailability, ProviderAccountId};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest,
|
||||
ProviderStoreError,
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest, ProviderStoreError,
|
||||
};
|
||||
use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use provider_openai::credential::{
|
||||
@@ -214,6 +214,85 @@ impl ProviderLeasePort for RefreshLeases {
|
||||
}
|
||||
}
|
||||
|
||||
/// 真实累加连续失败计数的测试 double,用于断言退避的指数增长与成功清零。
|
||||
#[derive(Default)]
|
||||
struct CountingCredentialState {
|
||||
counts: Mutex<HashMap<String, u32>>,
|
||||
cleared: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl CountingCredentialState {
|
||||
fn count(&self, account_id: &str) -> u32 {
|
||||
self.counts
|
||||
.lock()
|
||||
.expect("backoff counts lock")
|
||||
.get(account_id)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn was_cleared(&self, account_id: &str) -> bool {
|
||||
self.cleared
|
||||
.lock()
|
||||
.expect("cleared accounts lock")
|
||||
.iter()
|
||||
.any(|id| id == account_id)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCredentialStatePort for CountingCredentialState {
|
||||
fn replace(
|
||||
&self,
|
||||
_state: ProviderCredentialState,
|
||||
) -> BoxFuture<'_, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderCredentialState>, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn record_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
_window: Duration,
|
||||
) -> BoxFuture<'a, Result<u32, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut counts = self.counts.lock().expect("backoff counts lock");
|
||||
let entry = counts.entry(account_id.as_str().to_owned()).or_insert(0);
|
||||
*entry += 1;
|
||||
Ok(*entry)
|
||||
})
|
||||
}
|
||||
|
||||
fn clear_refresh_backoff<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
self.counts
|
||||
.lock()
|
||||
.expect("backoff counts lock")
|
||||
.remove(account_id.as_str());
|
||||
self.cleared
|
||||
.lock()
|
||||
.expect("cleared accounts lock")
|
||||
.push(account_id.as_str().to_owned());
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn setup(
|
||||
outcome: Result<TokenPair, RefreshFailure>,
|
||||
lease_available: bool,
|
||||
@@ -242,6 +321,7 @@ async fn setup(
|
||||
available: lease_available,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
(store, refresher, service)
|
||||
@@ -307,6 +387,7 @@ async fn refreshed_identity_rejection_revision_fences_account_as_invalid() {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
@@ -336,6 +417,7 @@ async fn unavailable_signature_verification_persists_refresh_backoff() {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
let before = SystemTime::now();
|
||||
@@ -344,15 +426,16 @@ async fn unavailable_signature_verification_persists_refresh_backoff() {
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Ambiguous { .. }]
|
||||
[CodexCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let account = store.account("acct_refresh").expect("deferred account");
|
||||
assert_eq!(account.revision().get(), 1);
|
||||
assert_eq!(account.availability(), AccountAvailability::Invalid);
|
||||
assert_eq!(account.revision().get(), 2);
|
||||
// JWKS/签名边界短暂不可用是瞬态:账号保持可用、不被永久失效,仅推进退避重试。
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert!(
|
||||
account
|
||||
.next_refresh_at()
|
||||
.is_some_and(|retry| retry < before)
|
||||
.is_some_and(|retry| retry > before)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -382,6 +465,7 @@ async fn unavailable_usage_preserves_rotated_tokens_and_persists_backoff() {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
let before = SystemTime::now();
|
||||
@@ -445,6 +529,7 @@ async fn permanent_refresh_failure_cannot_overwrite_a_newer_credential_revision(
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
@@ -489,16 +574,18 @@ async fn proven_pre_send_transport_failure_uses_short_cooldown() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ambiguous_refresh_does_not_mutate_tokens_or_account_state() {
|
||||
async fn transient_refresh_defers_without_invalidating_account() {
|
||||
let (store, refresher, service) = setup(Err(RefreshFailure::Transport), true).await;
|
||||
let outcomes = service.refresh_due().await.expect("refresh due");
|
||||
// 上游瞬态(含可能已发出的 ambiguous)不再永久失效账号,改为退避重试。
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Ambiguous { .. }]
|
||||
[CodexCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let account = store.account("acct_refresh").expect("account");
|
||||
assert_eq!(account.revision().get(), 1);
|
||||
assert_eq!(account.availability(), AccountAvailability::Invalid);
|
||||
assert_eq!(account.revision().get(), 2);
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
// 已退避到未来,本轮不再 due,且未重复兑换 token。
|
||||
assert!(
|
||||
service
|
||||
.refresh_due()
|
||||
@@ -572,6 +659,7 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
@@ -596,3 +684,128 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() {
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
async fn seed_due_account(store: &MemoryAccountStore) {
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: "acct_refresh".to_owned(),
|
||||
name: "refresh".to_owned(),
|
||||
secret: secret("old-access"),
|
||||
verified_account: profile("chatgpt-acct_refresh"),
|
||||
next_refresh_at: Some(chrono::Utc::now() - chrono::Duration::seconds(1)),
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_backoff_grows_exponentially_across_attempts() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed_due_account(&store).await;
|
||||
let credential_state = Arc::new(CountingCredentialState::default());
|
||||
let state_port: Arc<dyn ProviderCredentialStatePort> = credential_state.clone();
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
Arc::new(Refresher::scripted([
|
||||
Err(RefreshFailure::Transport),
|
||||
Err(RefreshFailure::Transport),
|
||||
])),
|
||||
Arc::new(VerifiedIdentity),
|
||||
Arc::new(RefreshLeases {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
state_port,
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
let before_first = SystemTime::now();
|
||||
let first = service.refresh_due().await.expect("first refresh cycle");
|
||||
assert!(matches!(
|
||||
first.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let first_account = store
|
||||
.account("acct_refresh")
|
||||
.expect("account after first defer");
|
||||
let first_delay = first_account
|
||||
.next_refresh_at()
|
||||
.expect("first retry scheduled")
|
||||
.duration_since(before_first)
|
||||
.expect("first delay is in the future");
|
||||
assert_eq!(credential_state.count("acct_refresh"), 1);
|
||||
|
||||
// 已退避到未来的账号真实调度不会再次到期,这里手动复位以触发第二次退避。
|
||||
store
|
||||
.repository()
|
||||
.defer_refresh(&first_account, SystemTime::now() - Duration::from_secs(1))
|
||||
.await
|
||||
.expect("reset next_refresh_at to the past");
|
||||
|
||||
let before_second = SystemTime::now();
|
||||
let second = service.refresh_due().await.expect("second refresh cycle");
|
||||
assert!(matches!(
|
||||
second.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let second_delay = store
|
||||
.account("acct_refresh")
|
||||
.expect("account after second defer")
|
||||
.next_refresh_at()
|
||||
.expect("second retry scheduled")
|
||||
.duration_since(before_second)
|
||||
.expect("second delay is in the future");
|
||||
assert_eq!(credential_state.count("acct_refresh"), 2);
|
||||
|
||||
// base=5s、factor=3:第二次(attempt=2)应比第一次(attempt=1)显著更久。
|
||||
assert!(
|
||||
second_delay > first_delay * 2,
|
||||
"second backoff {second_delay:?} should grow well beyond first {first_delay:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_refresh_clears_backoff_counter() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed_due_account(&store).await;
|
||||
let credential_state = Arc::new(CountingCredentialState::default());
|
||||
let state_port: Arc<dyn ProviderCredentialStatePort> = credential_state.clone();
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
Arc::new(Refresher::scripted([
|
||||
Err(RefreshFailure::Transport),
|
||||
Ok(success_tokens()),
|
||||
])),
|
||||
Arc::new(VerifiedIdentity),
|
||||
Arc::new(RefreshLeases {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
state_port,
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
// 第一次瞬态失败 → 计数累加到 1。
|
||||
let first = service.refresh_due().await.expect("first refresh cycle");
|
||||
assert!(matches!(
|
||||
first.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
assert_eq!(credential_state.count("acct_refresh"), 1);
|
||||
|
||||
// 复位到期后成功刷新,退避计数被清零。
|
||||
let deferred = store.account("acct_refresh").expect("deferred account");
|
||||
store
|
||||
.repository()
|
||||
.defer_refresh(&deferred, SystemTime::now() - Duration::from_secs(1))
|
||||
.await
|
||||
.expect("reset next_refresh_at to the past");
|
||||
|
||||
let second = service.refresh_due().await.expect("second refresh cycle");
|
||||
assert!(matches!(
|
||||
second.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Refreshed { .. }]
|
||||
));
|
||||
assert!(credential_state.was_cleared("acct_refresh"));
|
||||
assert_eq!(credential_state.count("acct_refresh"), 0);
|
||||
}
|
||||
|
||||
@@ -221,6 +221,24 @@ async fn token_revoked_text_without_invalid_grant_should_remain_temporary() {
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_error_or_rate_limit_must_stay_transient_even_with_oauth_error_body() {
|
||||
// 5xx/429(含 CDN/网关页恰好嵌入 invalid_grant 字样)按状态码判瞬态,
|
||||
// 不因正文子串把账号永久终态——正文只在 4xx 时才是权威 OAuth 错误。
|
||||
for status in [500, 502, 503, 429] {
|
||||
let failure = refresh_failure(
|
||||
status,
|
||||
r#"{"error":"invalid_grant","error_description":"refresh_token_expired"}"#,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
failure,
|
||||
RefreshFailure::Transport,
|
||||
"status {status} must classify as transient"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh_failure(status: u16, body: &str) -> RefreshFailure {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
|
||||
@@ -3,6 +3,7 @@ use gateway_core::event::{ContentKind, FinishReason, GatewayEvent, ProviderEvent
|
||||
use provider_openai::transport::canonical::{
|
||||
CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalFailure, CodexCanonicalOutcome,
|
||||
};
|
||||
use provider_openai::transport::protocol::websocket::websocket_event_to_sse_frame;
|
||||
|
||||
const METADATA_PREFIX_FIXTURE: &str = include_str!(
|
||||
"../../../../gateway-api/tests/openai/responses/fixtures/http_sse/metadata_only_prefix.sse"
|
||||
@@ -63,6 +64,53 @@ fn raw_sse_passthrough_should_keep_original_bytes_alongside_canonical_facts() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_raw_passthrough_should_preserve_upstream_number_bytes() {
|
||||
// OpenAI 线路透明代理:WebSocket 上游帧经 reducer→SSE→decoder(raw 透传) 后,
|
||||
// data 段必须与上游原文逐字节一致,不得经 serde 往返改写数值/精度。
|
||||
// `1e3` 若被重序列化会变成 `1000.0`——用它作为字节改写的探针。
|
||||
let upstream = r#"{"type":"response.created","response":{"id":"resp_ws_raw","model":"gpt-test","x_precision":1e3}}"#;
|
||||
let frame = websocket_event_to_sse_frame(upstream)
|
||||
.expect("client-visible WS event yields an SSE frame");
|
||||
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_raw_sse_passthrough()
|
||||
.push(frame.as_bytes())
|
||||
.expect("WS-derived SSE frame should remain deliverable");
|
||||
let wire = events
|
||||
.iter()
|
||||
.filter_map(ProviderEvent::wire_event)
|
||||
.find(|wire| wire.has_json_data())
|
||||
.expect("response.created wire");
|
||||
let raw = wire
|
||||
.raw_sse_frame()
|
||||
.map(AsRef::as_ref)
|
||||
.expect("raw frame bytes");
|
||||
let raw = std::str::from_utf8(raw).expect("utf8 raw frame");
|
||||
|
||||
assert!(
|
||||
raw.contains("1e3"),
|
||||
"raw frame must keep upstream `1e3` verbatim: {raw}"
|
||||
);
|
||||
assert!(
|
||||
!raw.contains("1000.0"),
|
||||
"raw frame must not re-serialize the number via serde: {raw}"
|
||||
);
|
||||
assert!(matches!(
|
||||
wire_for_response_created(&events),
|
||||
[GatewayEvent::Started(metadata)] if metadata.response_id() == "resp_ws_raw"
|
||||
));
|
||||
}
|
||||
|
||||
fn wire_for_response_created(events: &[ProviderEvent]) -> &[GatewayEvent] {
|
||||
events
|
||||
.iter()
|
||||
.find(|event| {
|
||||
event.wire_event().and_then(|wire| wire.event_type()) == Some("response.created")
|
||||
})
|
||||
.map_or(&[], ProviderEvent::canonical_facts)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_sse_passthrough_should_forward_unparseable_frames_without_failure() {
|
||||
let raw = b"retry: later\ndata: opaque\n\n";
|
||||
|
||||
@@ -90,6 +90,32 @@ fn websocket_response_create_payload_should_keep_explicit_empty_instructions() {
|
||||
assert_eq!(value["instructions"], "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_response_create_payload_text_should_match_merged_map_serialization() {
|
||||
// body 自带 type 键与未知字段:文本帧必须与「合并 Map 再序列化」逐字节一致。
|
||||
let request = CodexResponsesRequest::from_body(
|
||||
json!({
|
||||
"model": "gpt-test",
|
||||
"type": "client-supplied",
|
||||
"x_unknown": {"nested": [1, 2.5, "3"]},
|
||||
"input": []
|
||||
})
|
||||
.as_object()
|
||||
.expect("request object")
|
||||
.clone(),
|
||||
);
|
||||
|
||||
let merged = serde_json::to_string(
|
||||
&provider_openai::transport::protocol::websocket::websocket_response_create_payload(
|
||||
&request,
|
||||
),
|
||||
)
|
||||
.expect("merged payload");
|
||||
let streamed = websocket_response_create_payload_text(&request).expect("streamed payload");
|
||||
|
||||
assert_eq!(streamed, merged);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_websocket_audit_should_forbid_http_fallback() {
|
||||
let request = CodexResponsesRequest::from_body(
|
||||
@@ -138,8 +164,7 @@ fn websocket_metadata_turn_state_should_accept_case_insensitive_header() {
|
||||
let event = json!({
|
||||
"type": "response.metadata",
|
||||
"headers": {"X-Codex-Turn-State": ["turn-from-metadata"]}
|
||||
})
|
||||
.to_string();
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
websocket_metadata_turn_state(&event).as_deref(),
|
||||
@@ -155,16 +180,14 @@ fn websocket_completed_id_should_read_the_id_without_validating_the_response_sha
|
||||
"id": "resp_valid",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
});
|
||||
let invalid_usage = json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_invalid",
|
||||
"usage": {"input_tokens": "bad", "output_tokens": 1, "total_tokens": 1}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
websocket_response_completed_id(&valid),
|
||||
@@ -226,16 +249,14 @@ fn websocket_completed_id_should_leave_unreadable_completion_untracked() {
|
||||
"response": {
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
});
|
||||
let incomplete_usage = json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_incomplete_usage",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
});
|
||||
|
||||
assert_eq!(websocket_response_completed_id(&missing_id), None);
|
||||
assert_eq!(
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user