mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-09 00:05:47 +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 |
@@ -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
+1
-3
@@ -1189,7 +1189,6 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
"serde",
|
||||
@@ -1225,7 +1224,6 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-appender",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
@@ -2151,7 +2149,6 @@ dependencies = [
|
||||
"url",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2161,6 +2158,7 @@ dependencies = [
|
||||
"async-stream",
|
||||
"async-trait",
|
||||
"base64",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"gateway-admin",
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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<()>;
|
||||
}
|
||||
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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, PublicModelProfile};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::ApiState;
|
||||
@@ -19,29 +20,48 @@ 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<_>>();
|
||||
let profiles = service.public_model_profiles(&client);
|
||||
|
||||
if !profiles.is_empty() {
|
||||
let models = profiles
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, profile)| codex_model_json(profile, index))
|
||||
.collect::<Vec<_>>();
|
||||
return (StatusCode::OK, Json(json!({ "models": models }))).into_response();
|
||||
}
|
||||
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
@@ -209,6 +229,7 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
|
||||
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",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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(),
|
||||
})
|
||||
|
||||
@@ -146,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,
|
||||
@@ -157,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),
|
||||
@@ -239,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);
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
|
||||
@@ -15,13 +15,13 @@ 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,
|
||||
@@ -67,7 +67,9 @@ impl ExecutionService for ModelsExecution {
|
||||
)
|
||||
.with_reasoning(
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high"].map(str::to_owned).to_vec(),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
)
|
||||
.with_context_window_tokens(Some(500_000))
|
||||
.with_image_input(true)
|
||||
@@ -133,11 +135,77 @@ async fn models_should_encode_provider_profiles_for_current_codex_clients() {
|
||||
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]
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -73,6 +74,12 @@ where
|
||||
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,6 +93,7 @@ where
|
||||
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,
|
||||
@@ -143,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,
|
||||
@@ -156,6 +166,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
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,
|
||||
@@ -374,6 +386,7 @@ where
|
||||
&mut current.stream,
|
||||
self.cancellation.clone(),
|
||||
self.deadline,
|
||||
&mut self.deadline_timer,
|
||||
)
|
||||
.await
|
||||
};
|
||||
@@ -614,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,
|
||||
@@ -642,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);
|
||||
@@ -676,6 +685,7 @@ where
|
||||
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
|
||||
.await?;
|
||||
current.send_observed = true;
|
||||
self.send_state_watermark = UpstreamSendState::Sent;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -688,6 +698,7 @@ where
|
||||
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
|
||||
.await?;
|
||||
current.send_observed = true;
|
||||
self.send_state_watermark = UpstreamSendState::Sent;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -776,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)
|
||||
@@ -789,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();
|
||||
@@ -804,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
|
||||
@@ -990,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() {
|
||||
@@ -1042,6 +1058,7 @@ where
|
||||
} else {
|
||||
UpstreamSendState::Ambiguous
|
||||
};
|
||||
let send_state = self.raise_send_watermark(send_state);
|
||||
if self.attempts > 0 {
|
||||
self.engine
|
||||
.store()
|
||||
@@ -1162,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)
|
||||
@@ -1172,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) {
|
||||
@@ -1225,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),
|
||||
}
|
||||
}
|
||||
@@ -1300,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,7 +22,8 @@ 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};
|
||||
@@ -31,6 +32,7 @@ use crate::policy::{ClientApiKeyId, ClientPolicy};
|
||||
use crate::routing::snapshot::RuntimeSnapshotHandle;
|
||||
use crate::routing::{
|
||||
ProviderKind, PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
};
|
||||
|
||||
const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60);
|
||||
@@ -187,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>,
|
||||
@@ -205,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());
|
||||
@@ -212,6 +217,7 @@ impl DefaultExecutionService {
|
||||
snapshots,
|
||||
coordinator: Arc::new(AttemptCoordinator::new(engine)),
|
||||
probe_coordinator: Arc::new(AttemptCoordinator::new(probe_engine)),
|
||||
observations,
|
||||
providers,
|
||||
admissions,
|
||||
circuits,
|
||||
@@ -417,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,
|
||||
@@ -464,7 +475,7 @@ impl DefaultExecutionService {
|
||||
started_at,
|
||||
deadline_at,
|
||||
};
|
||||
let mut session = self
|
||||
let mut session = match self
|
||||
.probe_coordinator
|
||||
.start(
|
||||
new_request,
|
||||
@@ -475,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()
|
||||
@@ -498,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;
|
||||
@@ -868,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,
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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,6 +474,19 @@ 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 只保存可丢失的调度截止时间,不进入账号持久状态。
|
||||
@@ -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,
|
||||
|
||||
@@ -1620,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(())
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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,6 +227,11 @@ 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
|
||||
|
||||
@@ -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();
|
||||
@@ -397,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"))?;
|
||||
|
||||
@@ -759,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,
|
||||
@@ -1651,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,
|
||||
@@ -1659,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()
|
||||
@@ -1709,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,
|
||||
@@ -1720,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
|
||||
@@ -2249,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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,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]
|
||||
|
||||
@@ -174,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));
|
||||
@@ -217,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,
|
||||
);
|
||||
|
||||
@@ -235,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
|
||||
@@ -442,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));
|
||||
@@ -501,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
|
||||
);
|
||||
|
||||
@@ -592,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(
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -11,6 +11,7 @@ autotests = false
|
||||
async-stream = "0.3.6"
|
||||
async-trait = "0.1.89"
|
||||
base64 = "0.22.1"
|
||||
bytes = "1.9.0"
|
||||
chrono = { version = "0.4.45", features = ["clock", "serde"] }
|
||||
futures = "0.3.32"
|
||||
gateway-admin = { path = "../../gateway-admin" }
|
||||
|
||||
@@ -316,36 +316,81 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
.map_err(map_provider_store_error)?;
|
||||
let mut subjects = BTreeSet::new();
|
||||
let mut credentials = Vec::new();
|
||||
// 逐条目独立成败:verify 会真实轮换 RT,后续条目的任何失败都不得
|
||||
// 丢弃已轮换成功的凭据(旧 RT 已被上游作废,丢弃即不可逆报废)。
|
||||
// 失败条目跳过;仅当没有任何条目成功时返回首个失败。响应只携带
|
||||
// 成功计数,每个被跳过的条目在此逐条留日志——尤其 verify 成功后
|
||||
// 才失败/被去重的条目,其源文件里的 RT 已被轮换作废,运维只能从
|
||||
// 日志得知需要重新授权。
|
||||
let mut first_failure: Option<ProviderAdminError> = None;
|
||||
for entry in document.into_entries() {
|
||||
let name = entry.name().to_owned();
|
||||
let email = entry.email().map(str::to_owned);
|
||||
let tokens = self
|
||||
let tokens = match self
|
||||
.oauth
|
||||
.verify_imported_credential(&discovery, entry.into_candidate())
|
||||
.await
|
||||
.map_err(|error| map_failure_class(error.class()))?;
|
||||
{
|
||||
Ok(tokens) => tokens,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
target: "xai_admin",
|
||||
account = name.as_str(),
|
||||
class = ?error.class(),
|
||||
"导入条目验证失败,已跳过"
|
||||
);
|
||||
first_failure.get_or_insert_with(|| map_failure_class(error.class()));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
// 同一账号的另一份 grant:账号已由更早条目入册,本条目跳过;
|
||||
// 本条目若经历了 refresh,其源文件 RT 已作废。
|
||||
if !subjects.insert(tokens.evidence().subject().to_owned()) {
|
||||
return Err(provider_error(ProviderAdminErrorKind::Invalid));
|
||||
tracing::warn!(
|
||||
target: "xai_admin",
|
||||
account = name.as_str(),
|
||||
"导入条目与更早条目属同一账号,已跳过;该条目的 refresh token 可能已被轮换作废"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let prepared = GrokCredentialAdmin
|
||||
.prepare_verified_account(
|
||||
&VerifiedGrokAccount {
|
||||
account_id: ProviderAccountId::new(format!(
|
||||
"acct_{}",
|
||||
Uuid::now_v7().simple()
|
||||
))
|
||||
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
|
||||
name,
|
||||
email,
|
||||
upstream_account_id: None,
|
||||
plan_type: None,
|
||||
tokens,
|
||||
enabled: true,
|
||||
},
|
||||
policy,
|
||||
)
|
||||
.map_err(map_repository_error)?;
|
||||
credentials.push(prepared_create(prepared, Utc::now())?);
|
||||
let account = VerifiedGrokAccount {
|
||||
account_id: ProviderAccountId::new(format!("acct_{}", Uuid::now_v7().simple()))
|
||||
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
|
||||
name,
|
||||
email,
|
||||
upstream_account_id: None,
|
||||
plan_type: None,
|
||||
tokens,
|
||||
enabled: true,
|
||||
};
|
||||
let prepared = match GrokCredentialAdmin.prepare_verified_account(&account, policy) {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
target: "xai_admin",
|
||||
account = account.name.as_str(),
|
||||
"已验证条目本地入册失败,已跳过;其 refresh token 已被轮换,需重新授权该账号"
|
||||
);
|
||||
first_failure.get_or_insert_with(|| map_repository_error(error));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match prepared_create(prepared, Utc::now()) {
|
||||
Ok(credential) => credentials.push(credential),
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
target: "xai_admin",
|
||||
account = account.name.as_str(),
|
||||
"已验证条目序列化失败,已跳过;其 refresh token 已被轮换,需重新授权该账号"
|
||||
);
|
||||
first_failure.get_or_insert(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
if credentials.is_empty() {
|
||||
return Err(
|
||||
first_failure.unwrap_or_else(|| provider_error(ProviderAdminErrorKind::Invalid))
|
||||
);
|
||||
}
|
||||
Ok(PreparedCredentialImport {
|
||||
provider_kind: self.provider_kind.clone(),
|
||||
|
||||
@@ -172,6 +172,12 @@ impl GrokOAuthImportDocument {
|
||||
entries.push(entry);
|
||||
}
|
||||
}
|
||||
// 同一 refresh token 的重复条目只保留首个:RT exchange 会作废旧
|
||||
// grant,重复条目逐条交换会让先完成的轮换被后续条目作废。
|
||||
let mut seen_refresh_tokens = HashSet::new();
|
||||
entries.retain(|entry| {
|
||||
seen_refresh_tokens.insert(entry.candidate.refresh_token.expose().to_owned())
|
||||
});
|
||||
if entries.is_empty() {
|
||||
return Err(GrokOAuthImportError::InvalidField("document"));
|
||||
}
|
||||
|
||||
@@ -12,10 +12,10 @@ use gateway_core::engine::credential::{
|
||||
AccountAvailability, CredentialRevision, ProviderAccountId,
|
||||
};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCooldown, ProviderCooldownPort, ProviderLeaseAcquisition, ProviderLeasePort,
|
||||
ProviderLeaseRequest, ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest,
|
||||
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
|
||||
provider_refresh_retry_at,
|
||||
ProviderCooldown, ProviderCooldownPort, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshCapacityRequest,
|
||||
ProviderRefreshLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
|
||||
ProviderStoreError, provider_refresh_backoff_at,
|
||||
};
|
||||
|
||||
use super::catalog::GrokCredentialCatalogService;
|
||||
@@ -36,7 +36,8 @@ const MAX_REFRESH_BATCH: u32 = 100;
|
||||
const MAX_REFRESH_EXCLUSIONS: usize = 400;
|
||||
const MAX_SECRET_BYTES: usize = 64 * 1_024;
|
||||
const DISCOVERY_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(5);
|
||||
const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60);
|
||||
/// 连续失败计数窗口;静默满窗后计数过期归零。
|
||||
const REFRESH_BACKOFF_WINDOW: Duration = Duration::from_secs(30 * 60);
|
||||
const REFRESH_AMBIGUITY_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
/// 一个到期且已按 revision 读取明文 RT 的 xAI account。
|
||||
@@ -277,6 +278,7 @@ pub struct GrokCredentialRefreshService {
|
||||
catalog: Arc<GrokCredentialCatalogService>,
|
||||
leases: Arc<dyn ProviderLeasePort>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
}
|
||||
|
||||
@@ -287,6 +289,7 @@ impl GrokCredentialRefreshService {
|
||||
catalog: Arc<GrokCredentialCatalogService>,
|
||||
leases: Arc<dyn ProviderLeasePort>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
) -> Self {
|
||||
Self {
|
||||
@@ -295,6 +298,7 @@ impl GrokCredentialRefreshService {
|
||||
catalog,
|
||||
leases,
|
||||
cooldowns,
|
||||
credential_state,
|
||||
runtime_policy,
|
||||
}
|
||||
}
|
||||
@@ -339,17 +343,6 @@ impl GrokCredentialRefreshService {
|
||||
self.refresh_one_with_policy(credential, policy).await
|
||||
}
|
||||
|
||||
pub async fn prepare_cycle_if_due(&self) -> Result<bool, GrokCredentialRefreshError> {
|
||||
if !self.repository.has_due_refresh().await? {
|
||||
return Ok(false);
|
||||
}
|
||||
self.refresher
|
||||
.prepare_cycle()
|
||||
.await
|
||||
.map_err(|_| GrokCredentialRefreshError::Preparation)?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// 手工刷新一个指定 revision;只返回 Provider 验证后的 CAS command,不写 Store。
|
||||
pub async fn prepare_manual_refresh(
|
||||
&self,
|
||||
@@ -650,6 +643,10 @@ impl GrokCredentialRefreshService {
|
||||
{
|
||||
let _ = self.catalog.cache_seed(&account_id, seed).await;
|
||||
}
|
||||
let _ = self
|
||||
.credential_state
|
||||
.clear_refresh_backoff(&account_id)
|
||||
.await;
|
||||
Ok(GrokCredentialRefreshOutcome::Refreshed {
|
||||
account_id,
|
||||
credential_revision: record.credential_revision,
|
||||
@@ -729,10 +726,15 @@ impl GrokCredentialRefreshService {
|
||||
credential: &DueGrokCredential,
|
||||
reason: &'static str,
|
||||
) -> Result<bool, GrokCredentialRefreshError> {
|
||||
let retry_at = provider_refresh_retry_at(
|
||||
let attempt = self
|
||||
.credential_state
|
||||
.record_refresh_backoff(&credential.account_id, REFRESH_BACKOFF_WINDOW)
|
||||
.await
|
||||
.unwrap_or(1);
|
||||
let retry_at = provider_refresh_backoff_at(
|
||||
&credential.account_id,
|
||||
SystemTime::now(),
|
||||
REFRESH_RETRY_DELAY,
|
||||
attempt,
|
||||
reason,
|
||||
)?;
|
||||
match self
|
||||
@@ -781,22 +783,6 @@ impl GrokCredentialRecovery for GrokCredentialRefreshService {
|
||||
}
|
||||
|
||||
impl GrokCredentialRepository {
|
||||
pub async fn has_due_refresh(&self) -> Result<bool, GrokCredentialRepositoryError> {
|
||||
Ok(self
|
||||
.list_all_accounts()
|
||||
.await?
|
||||
.iter()
|
||||
.any(|account| account_due(account, SystemTime::now())))
|
||||
}
|
||||
|
||||
pub async fn list_due_refresh(
|
||||
&self,
|
||||
) -> Result<Vec<DueGrokCredential>, GrokCredentialRepositoryError> {
|
||||
self.list_due_refresh_batch_excluding(&[])
|
||||
.await
|
||||
.map(|batch| batch.credentials)
|
||||
}
|
||||
|
||||
async fn list_due_refresh_batch_excluding(
|
||||
&self,
|
||||
excluded_account_ids: &[ProviderAccountId],
|
||||
|
||||
@@ -102,7 +102,7 @@ impl GrokAccountSessionSelector {
|
||||
let runtime_cooldowns = self.runtime_cooldowns(&catalog_eligible).await;
|
||||
let mut candidates = catalog_eligible
|
||||
.into_iter()
|
||||
.filter(|account| !runtime_cooldowns.contains(account.id()))
|
||||
.filter(|account| !runtime_cooldowns.contains_key(account.id()))
|
||||
.map(|account| {
|
||||
let health = self
|
||||
.account_feedback
|
||||
@@ -197,18 +197,30 @@ impl GrokAccountSessionSelector {
|
||||
}
|
||||
|
||||
if capacity_denied {
|
||||
Err(GrokSessionSelectorError::CapacityUnavailable { retry_after })
|
||||
} else {
|
||||
Err(GrokSessionSelectorError::NoEligibleSession)
|
||||
return Err(GrokSessionSelectorError::CapacityUnavailable { retry_after });
|
||||
}
|
||||
// 钉死账号时只看该账号,否则取最早退出 cooldown 的候选。
|
||||
let cooled_until = match request.required_account() {
|
||||
Some(required) => runtime_cooldowns.get(required).copied(),
|
||||
None => runtime_cooldowns.values().copied().min(),
|
||||
};
|
||||
if let Some(until) = cooled_until {
|
||||
return Err(GrokSessionSelectorError::AccountCoolingDown {
|
||||
retry_after: until
|
||||
.duration_since(SystemTime::now())
|
||||
.ok()
|
||||
.filter(|remaining| !remaining.is_zero()),
|
||||
});
|
||||
}
|
||||
Err(GrokSessionSelectorError::NoEligibleSession)
|
||||
}
|
||||
|
||||
async fn runtime_cooldowns(
|
||||
&self,
|
||||
accounts: &[ProviderAccount],
|
||||
) -> std::collections::BTreeSet<ProviderAccountId> {
|
||||
) -> std::collections::BTreeMap<ProviderAccountId, SystemTime> {
|
||||
let now = SystemTime::now();
|
||||
let mut cooled = std::collections::BTreeSet::new();
|
||||
let mut cooled = std::collections::BTreeMap::new();
|
||||
for account in accounts {
|
||||
let Ok(Some(cooldown)) = self.cooldowns.read(account.id()).await else {
|
||||
continue;
|
||||
@@ -220,7 +232,7 @@ impl GrokAccountSessionSelector {
|
||||
continue;
|
||||
}
|
||||
if cooldown.until() > now {
|
||||
cooled.insert(account.id().clone());
|
||||
cooled.insert(account.id().clone(), cooldown.until());
|
||||
}
|
||||
}
|
||||
cooled
|
||||
|
||||
@@ -57,24 +57,25 @@ pub use transport::{
|
||||
GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture,
|
||||
GrokBillingTransportResponse, GrokCanonicalDecoder, GrokCatalogApiBackend,
|
||||
GrokCatalogCapabilities, GrokCatalogCapabilityEvidence, GrokCatalogLimits, GrokCatalogMetadata,
|
||||
GrokCatalogModel, GrokClientIdentity, GrokCompactionDecodeError, GrokCompactionRequest,
|
||||
GrokCompactionSummaryDecoder, GrokCredentialFailure, GrokCredentialFeedbackFuture,
|
||||
GrokDnsResolutionError, GrokDnsResolutionPlan, GrokDnsResolutionPolicy, GrokEndpointPolicy,
|
||||
GrokHeader, GrokHeaderValue, GrokInferenceChunkStream, GrokInferenceClientCacheStatus,
|
||||
GrokInferenceDnsObservation, GrokInferenceDnsSource, GrokInferenceRequest,
|
||||
GrokInferenceResponse, GrokInferenceTransport, GrokInferenceTransportError,
|
||||
GrokInferenceTransportErrorKind, GrokInferenceTransportFuture, GrokInferenceTransportMetrics,
|
||||
GrokModelCatalogClient, GrokModelCatalogError, GrokModelCatalogRequest,
|
||||
GrokModelCatalogSession, GrokModelCatalogSessionError, GrokModelCatalogSnapshot,
|
||||
GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind,
|
||||
GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokProviderConfigError,
|
||||
GrokProviderTransport, GrokRequestEncodeError, GrokReqwestTransportBuildError,
|
||||
GrokResponsesRequest, GrokSessionAffinityKey, GrokSessionBinding, GrokSessionDataError,
|
||||
GrokSessionLeaseGuard, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError,
|
||||
GrokSessionSelectorFuture, MAX_GROK_BILLING_BYTES, MAX_GROK_MODEL_CATALOG_BYTES,
|
||||
OfficialGrokEndpointPolicy, ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport,
|
||||
ReqwestOAuthTransport, SelectedGrokSession, XAI_PROVIDER_NAME, build_grok_headers,
|
||||
grok_billing_breakdown, parse_grok_billing, parse_grok_model_catalog,
|
||||
GrokCatalogModel, GrokCatalogReasoningEffort, GrokClientIdentity, GrokCompactionDecodeError,
|
||||
GrokCompactionRequest, GrokCompactionSummaryDecoder, GrokCredentialFailure,
|
||||
GrokCredentialFeedbackFuture, GrokDnsResolutionError, GrokDnsResolutionPlan,
|
||||
GrokDnsResolutionPolicy, GrokEndpointPolicy, GrokHeader, GrokHeaderValue,
|
||||
GrokInferenceChunkStream, GrokInferenceClientCacheStatus, GrokInferenceDnsObservation,
|
||||
GrokInferenceDnsSource, GrokInferenceRequest, GrokInferenceResponse, GrokInferenceTransport,
|
||||
GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportFuture,
|
||||
GrokInferenceTransportMetrics, GrokModelCatalogClient, GrokModelCatalogError,
|
||||
GrokModelCatalogRequest, GrokModelCatalogSession, GrokModelCatalogSessionError,
|
||||
GrokModelCatalogSnapshot, GrokModelCatalogTransport, GrokModelCatalogTransportError,
|
||||
GrokModelCatalogTransportErrorKind, GrokModelCatalogTransportFuture,
|
||||
GrokModelCatalogTransportResponse, GrokProviderConfigError, GrokProviderTransport,
|
||||
GrokRequestEncodeError, GrokReqwestTransportBuildError, GrokResponsesRequest,
|
||||
GrokSessionAffinityKey, GrokSessionBinding, GrokSessionDataError, GrokSessionLeaseGuard,
|
||||
GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, GrokSessionSelectorFuture,
|
||||
MAX_GROK_BILLING_BYTES, MAX_GROK_MODEL_CATALOG_BYTES, OfficialGrokEndpointPolicy,
|
||||
ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport,
|
||||
SelectedGrokSession, XAI_PROVIDER_NAME, build_grok_headers, grok_billing_breakdown,
|
||||
parse_grok_billing, parse_grok_model_catalog,
|
||||
};
|
||||
|
||||
/// xAI 初始化后交给组装根的最小能力集。
|
||||
@@ -167,6 +168,7 @@ pub async fn initialize(
|
||||
Arc::clone(&catalog),
|
||||
Arc::clone(&leases),
|
||||
cooldowns,
|
||||
ports.credential_state(),
|
||||
Arc::clone(&runtime_policy),
|
||||
));
|
||||
let credential_recovery: Arc<dyn GrokCredentialRecovery> = refresh.clone();
|
||||
|
||||
@@ -430,7 +430,9 @@ fn default_grok_model_presentation() -> ModelPresentation {
|
||||
)
|
||||
.with_reasoning(
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high"].map(str::to_owned).to_vec(),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
)
|
||||
.with_context_window_tokens(Some(500_000))
|
||||
.with_image_input(true)
|
||||
@@ -444,17 +446,35 @@ fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation {
|
||||
DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest"
|
||||
);
|
||||
let reasoning_evidence = model.capabilities().reasoning_effort();
|
||||
let catalog_reasoning_efforts = model
|
||||
.capabilities()
|
||||
.reasoning_efforts()
|
||||
.iter()
|
||||
.map(|effort| effort.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let catalog_default_reasoning = model
|
||||
.capabilities()
|
||||
.default_reasoning_effort()
|
||||
.map(|effort| effort.as_str().to_owned());
|
||||
let (default_reasoning, reasoning_efforts) = match reasoning_evidence {
|
||||
GrokCatalogCapabilityEvidence::DeclaredNative if !catalog_reasoning_efforts.is_empty() => (
|
||||
catalog_default_reasoning.or_else(|| catalog_reasoning_efforts.first().cloned()),
|
||||
catalog_reasoning_efforts,
|
||||
),
|
||||
GrokCatalogCapabilityEvidence::DeclaredNative => (
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high"].map(str::to_owned).to_vec(),
|
||||
catalog_default_reasoning.or_else(|| Some("medium".to_owned())),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
),
|
||||
GrokCatalogCapabilityEvidence::DeclaredUnsupported => {
|
||||
(Some("none".to_owned()), vec!["none".to_owned()])
|
||||
}
|
||||
GrokCatalogCapabilityEvidence::Unknown if known_grok_4_5 => (
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high"].map(str::to_owned).to_vec(),
|
||||
catalog_default_reasoning.or_else(|| Some("medium".to_owned())),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
),
|
||||
GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()),
|
||||
};
|
||||
@@ -969,7 +989,7 @@ async fn next_grok_chunk(
|
||||
selector: &dyn GrokSessionSelector,
|
||||
session: &SelectedGrokSession,
|
||||
context: &AttemptContext,
|
||||
) -> Result<Option<Vec<u8>>, ProviderError> {
|
||||
) -> Result<Option<bytes::Bytes>, ProviderError> {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
return Err(provider_error(
|
||||
ProviderErrorKind::Timeout,
|
||||
@@ -1572,21 +1592,55 @@ fn map_request_error(error: GrokRequestEncodeError) -> ProviderError {
|
||||
}
|
||||
}
|
||||
|
||||
/// 将选择阶段失败映射为带结构化 code 与 retry_after 的 Provider 错误。
|
||||
fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError {
|
||||
match error {
|
||||
GrokSessionSelectorError::CapacityUnavailable { retry_after } => {
|
||||
let error = provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent);
|
||||
match retry_after {
|
||||
Some(retry_after) => error.with_retry_after(retry_after),
|
||||
None => error,
|
||||
}
|
||||
}
|
||||
GrokSessionSelectorError::NoEligibleSession | GrokSessionSelectorError::Unavailable => {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
}
|
||||
let (retry_after, message, code) = match error {
|
||||
GrokSessionSelectorError::AccountCoolingDown { retry_after } => (
|
||||
retry_after,
|
||||
cooling_down_message(retry_after),
|
||||
"account_cooling_down",
|
||||
),
|
||||
GrokSessionSelectorError::CapacityUnavailable { retry_after } => (
|
||||
retry_after,
|
||||
"account is at its concurrency or request-interval limit".to_owned(),
|
||||
"account_capacity_busy",
|
||||
),
|
||||
GrokSessionSelectorError::NoEligibleSession => (
|
||||
None,
|
||||
"no account is eligible for the requested model".to_owned(),
|
||||
"no_eligible_account",
|
||||
),
|
||||
GrokSessionSelectorError::Unavailable => (
|
||||
None,
|
||||
"account scheduling state is temporarily unreadable".to_owned(),
|
||||
"account_selector_unavailable",
|
||||
),
|
||||
GrokSessionSelectorError::InvalidSession => {
|
||||
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
|
||||
return provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent);
|
||||
}
|
||||
};
|
||||
let error = provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent);
|
||||
let error = match retry_after {
|
||||
Some(retry_after) => error.with_retry_after(retry_after),
|
||||
None => error,
|
||||
};
|
||||
match ClientVisibleUpstreamError::new(
|
||||
message,
|
||||
Some(code.to_owned()),
|
||||
Some("account_unavailable_error".to_owned()),
|
||||
) {
|
||||
Ok(detail) => error.with_client_visible_upstream_error(detail),
|
||||
Err(_) => error,
|
||||
}
|
||||
}
|
||||
|
||||
fn cooling_down_message(retry_after: Option<Duration>) -> String {
|
||||
match retry_after {
|
||||
Some(retry_after) => format!(
|
||||
"account is cooling down after an interrupted upstream stream; retry in {}s",
|
||||
retry_after.as_secs().saturating_add(1)
|
||||
),
|
||||
None => "account is cooling down after an interrupted upstream stream".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,9 @@ use gateway_core::accounting::{
|
||||
CurrencyCode, Decimal, Money, ProviderReportedCost, Usage,
|
||||
};
|
||||
use gateway_core::engine::UpstreamSendState;
|
||||
use gateway_core::error::{ProviderError, ProviderErrorKind, SafeUpstreamValue};
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamError, ProviderError, ProviderErrorKind, SafeUpstreamValue,
|
||||
};
|
||||
use gateway_core::event::{
|
||||
ContentItem, ContentKind, FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent,
|
||||
ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta,
|
||||
@@ -211,10 +213,13 @@ impl GrokCanonicalDecoder {
|
||||
continue;
|
||||
};
|
||||
let event_type = event_type.to_owned();
|
||||
let transformed = self
|
||||
// 无法转换的上游帧(多为畸形帧)直接丢弃,保持已开始的客户端流不中断。
|
||||
let Ok(transformed) = self
|
||||
.response_transform
|
||||
.rewrite_stream_event(&event_type, value)
|
||||
.map_err(protocol_error)?;
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
for (index, transformed) in transformed.into_iter().enumerate() {
|
||||
let transformed_type = transformed.event_type().to_owned();
|
||||
if !client_visible_event(&transformed_type) {
|
||||
@@ -222,7 +227,23 @@ impl GrokCanonicalDecoder {
|
||||
}
|
||||
let value = transformed.into_value();
|
||||
let mut canonical = Vec::new();
|
||||
self.decode_event(&transformed_type, &value, &mut canonical)?;
|
||||
// 终态事件(completed/incomplete)fail-closed:用量/计费校验失败即断流。
|
||||
// 其余内容事件容忍字段校验失败——正常上游变体(空 delta、重复 index、
|
||||
// 缺字段等)不打断已开始的客户端流:跳过 canonical 提取、wire 原样转发。
|
||||
// 真·上游错误(response.failed/error)为非 Protocol 类别,按终态传播。
|
||||
let terminal_event = matches!(
|
||||
transformed_type.as_str(),
|
||||
"response.completed" | "response.incomplete"
|
||||
);
|
||||
match self.decode_event(&transformed_type, &value, &mut canonical) {
|
||||
Ok(()) => {}
|
||||
Err(error)
|
||||
if !terminal_event && error.kind() == ProviderErrorKind::Protocol =>
|
||||
{
|
||||
canonical.clear();
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
let wire_event = if transformed_type == event_type {
|
||||
event.event.clone()
|
||||
} else {
|
||||
@@ -762,11 +783,7 @@ fn incomplete_finish_reason(response: &Value) -> FinishReason {
|
||||
}
|
||||
|
||||
fn upstream_event_error(value: &Value) -> ProviderError {
|
||||
let code = value
|
||||
.pointer("/error/code")
|
||||
.or_else(|| value.pointer("/response/error/code"))
|
||||
.or_else(|| value.get("code"))
|
||||
.and_then(Value::as_str);
|
||||
let code = upstream_error_field(value, "code");
|
||||
let kind = match code {
|
||||
Some("invalid_request" | "invalid_prompt") => ProviderErrorKind::InvalidRequest,
|
||||
Some("unsupported" | "unsupported_feature") => ProviderErrorKind::Unsupported,
|
||||
@@ -781,9 +798,75 @@ fn upstream_event_error(value: &Value) -> ProviderError {
|
||||
if let Some(code) = code.and_then(|code| SafeUpstreamValue::new(code.to_owned()).ok()) {
|
||||
error = error.with_upstream_code(code);
|
||||
}
|
||||
// 结构化 message/code/type 供原客户端展示与重试分类;message 先脱去
|
||||
// 账号指纹(上游限流文案内嵌 team UUID),非结构化正文不透出。
|
||||
if let Some(message) = upstream_error_field(value, "message")
|
||||
&& let Ok(client_error) = ClientVisibleUpstreamError::new(
|
||||
scrub_account_fingerprints(message),
|
||||
code.map(str::to_owned),
|
||||
upstream_error_field(value, "type").map(str::to_owned),
|
||||
)
|
||||
{
|
||||
error = error.with_client_visible_upstream_error(client_error);
|
||||
}
|
||||
error
|
||||
}
|
||||
|
||||
fn upstream_error_field<'a>(value: &'a Value, field: &str) -> Option<&'a str> {
|
||||
value
|
||||
.pointer(&format!("/error/{field}"))
|
||||
.or_else(|| value.pointer(&format!("/response/error/{field}")))
|
||||
.or_else(|| value.get(field))
|
||||
.and_then(Value::as_str)
|
||||
}
|
||||
|
||||
const UUID_TEXT_LEN: usize = 36;
|
||||
|
||||
/// 把 message 中的 UUID 替换为占位符,控制字符归一为空格,
|
||||
/// 使文本满足客户端可见错误的约束且不携带可定位账号的标识。
|
||||
fn scrub_account_fingerprints(message: &str) -> String {
|
||||
let bytes = message.as_bytes();
|
||||
let mut scrubbed = String::with_capacity(message.len());
|
||||
let mut index = 0;
|
||||
while index < message.len() {
|
||||
if uuid_at(bytes, index) {
|
||||
scrubbed.push_str("[redacted]");
|
||||
index += UUID_TEXT_LEN;
|
||||
continue;
|
||||
}
|
||||
let character = message[index..].chars().next().expect("char boundary");
|
||||
scrubbed.push(if character.is_control() {
|
||||
' '
|
||||
} else {
|
||||
character
|
||||
});
|
||||
index += character.len_utf8();
|
||||
}
|
||||
scrubbed
|
||||
}
|
||||
|
||||
fn uuid_at(bytes: &[u8], index: usize) -> bool {
|
||||
let Some(candidate) = bytes.get(index..index + UUID_TEXT_LEN) else {
|
||||
return false;
|
||||
};
|
||||
if index > 0 && bytes[index - 1].is_ascii_alphanumeric() {
|
||||
return false;
|
||||
}
|
||||
if bytes
|
||||
.get(index + UUID_TEXT_LEN)
|
||||
.is_some_and(u8::is_ascii_alphanumeric)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
candidate.iter().enumerate().all(|(offset, byte)| {
|
||||
if matches!(offset, 8 | 13 | 18 | 23) {
|
||||
*byte == b'-'
|
||||
} else {
|
||||
byte.is_ascii_hexdigit()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn protocol_error(_error: impl std::fmt::Debug) -> ProviderError {
|
||||
protocol_error_marker()
|
||||
}
|
||||
|
||||
@@ -589,11 +589,46 @@ pub enum GrokCatalogApiBackend {
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// Grok 官方目录允许声明的 reasoning effort。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum GrokCatalogReasoningEffort {
|
||||
None,
|
||||
Minimal,
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
Xhigh,
|
||||
Max,
|
||||
/// 上游新增、当前 adapter 尚未识别的 effort;不进入能力投影。
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl GrokCatalogReasoningEffort {
|
||||
/// 返回 Responses API 使用的规范值。
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::None => "none",
|
||||
Self::Minimal => "minimal",
|
||||
Self::Low => "low",
|
||||
Self::Medium => "medium",
|
||||
Self::High => "high",
|
||||
Self::Xhigh => "xhigh",
|
||||
Self::Max => "max",
|
||||
Self::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Grok 目录中允许进入控制面的能力证据。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct GrokCatalogCapabilities {
|
||||
responses_api: GrokCatalogCapabilityEvidence,
|
||||
reasoning_effort: GrokCatalogCapabilityEvidence,
|
||||
reasoning_efforts: Vec<GrokCatalogReasoningEffort>,
|
||||
default_reasoning_effort: Option<GrokCatalogReasoningEffort>,
|
||||
backend_search: GrokCatalogCapabilityEvidence,
|
||||
streaming_tool_calls: GrokCatalogCapabilityEvidence,
|
||||
api_backend: Option<GrokCatalogApiBackend>,
|
||||
@@ -612,6 +647,18 @@ impl GrokCatalogCapabilities {
|
||||
self.reasoning_effort
|
||||
}
|
||||
|
||||
/// 返回上游声明的可选 reasoning effort,顺序与目录一致。
|
||||
#[must_use]
|
||||
pub fn reasoning_efforts(&self) -> &[GrokCatalogReasoningEffort] {
|
||||
&self.reasoning_efforts
|
||||
}
|
||||
|
||||
/// 返回上游声明或菜单推导出的默认 reasoning effort。
|
||||
#[must_use]
|
||||
pub const fn default_reasoning_effort(&self) -> Option<GrokCatalogReasoningEffort> {
|
||||
self.default_reasoning_effort
|
||||
}
|
||||
|
||||
/// 返回后端 search 支持证据。
|
||||
#[must_use]
|
||||
pub const fn backend_search(&self) -> GrokCatalogCapabilityEvidence {
|
||||
@@ -850,6 +897,10 @@ struct GrokModelWire {
|
||||
alias = "supports_reasoning_effort"
|
||||
)]
|
||||
supports_reasoning_effort: Option<bool>,
|
||||
#[serde(rename = "reasoningEffort", alias = "reasoning_effort")]
|
||||
reasoning_effort: Option<GrokCatalogReasoningEffort>,
|
||||
#[serde(default, rename = "reasoningEfforts", alias = "reasoning_efforts")]
|
||||
reasoning_efforts: Vec<GrokReasoningEffortOptionWire>,
|
||||
#[serde(rename = "supportsBackendSearch", alias = "supports_backend_search")]
|
||||
supports_backend_search: Option<bool>,
|
||||
#[serde(rename = "streamToolCalls", alias = "stream_tool_calls")]
|
||||
@@ -857,6 +908,26 @@ struct GrokModelWire {
|
||||
hidden: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum GrokReasoningEffortOptionWire {
|
||||
Bare(GrokCatalogReasoningEffort),
|
||||
Detailed {
|
||||
value: GrokCatalogReasoningEffort,
|
||||
#[serde(default)]
|
||||
default: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl GrokReasoningEffortOptionWire {
|
||||
fn into_parts(self) -> (GrokCatalogReasoningEffort, bool) {
|
||||
match self {
|
||||
Self::Bare(value) => (value, false),
|
||||
Self::Detailed { value, default } => (value, default),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 解析唯一受支持的 Grok 官方完整快照。
|
||||
///
|
||||
/// # Errors
|
||||
@@ -922,15 +993,41 @@ fn normalize_model(wire: GrokModelWire) -> Result<GrokCatalogModel, GrokModelCat
|
||||
.api_backend
|
||||
.filter(|backend| *backend != GrokCatalogApiBackend::Unknown);
|
||||
let responses_api = responses_evidence(wire.supported_in_api, api_backend);
|
||||
let mut menu_default = None;
|
||||
// 未识别的 effort 值不进入能力投影,也不作为默认值候选。
|
||||
let reasoning_efforts = wire
|
||||
.reasoning_efforts
|
||||
.into_iter()
|
||||
.filter_map(|option| {
|
||||
let (effort, is_default) = option.into_parts();
|
||||
if effort == GrokCatalogReasoningEffort::Unknown {
|
||||
return None;
|
||||
}
|
||||
if is_default && menu_default.is_none() {
|
||||
menu_default = Some(effort);
|
||||
}
|
||||
Some(effort)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let default_reasoning_effort = wire
|
||||
.reasoning_effort
|
||||
.filter(|effort| *effort != GrokCatalogReasoningEffort::Unknown)
|
||||
.or(menu_default)
|
||||
.or_else(|| reasoning_efforts.first().copied());
|
||||
let reasoning_effort = if reasoning_efforts.is_empty() {
|
||||
GrokCatalogCapabilityEvidence::from_wire(wire.supports_reasoning_effort)
|
||||
} else {
|
||||
GrokCatalogCapabilityEvidence::DeclaredNative
|
||||
};
|
||||
|
||||
Ok(GrokCatalogModel {
|
||||
request_model,
|
||||
display_name: wire.name,
|
||||
capabilities: GrokCatalogCapabilities {
|
||||
responses_api,
|
||||
reasoning_effort: GrokCatalogCapabilityEvidence::from_wire(
|
||||
wire.supports_reasoning_effort,
|
||||
),
|
||||
reasoning_effort,
|
||||
reasoning_efforts,
|
||||
default_reasoning_effort,
|
||||
backend_search: GrokCatalogCapabilityEvidence::from_wire(wire.supports_backend_search),
|
||||
streaming_tool_calls: GrokCatalogCapabilityEvidence::from_wire(wire.stream_tool_calls),
|
||||
api_backend,
|
||||
|
||||
@@ -21,10 +21,10 @@ pub use catalog::{
|
||||
GrokBillingRequest, GrokBillingSnapshot, GrokBillingTransport, GrokBillingTransportError,
|
||||
GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse,
|
||||
GrokCatalogApiBackend, GrokCatalogCapabilities, GrokCatalogCapabilityEvidence,
|
||||
GrokCatalogLimits, GrokCatalogMetadata, GrokCatalogModel, GrokModelCatalogClient,
|
||||
GrokModelCatalogError, GrokModelCatalogRequest, GrokModelCatalogSession,
|
||||
GrokModelCatalogSessionError, GrokModelCatalogSnapshot, GrokModelCatalogTransport,
|
||||
GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind,
|
||||
GrokCatalogLimits, GrokCatalogMetadata, GrokCatalogModel, GrokCatalogReasoningEffort,
|
||||
GrokModelCatalogClient, GrokModelCatalogError, GrokModelCatalogRequest,
|
||||
GrokModelCatalogSession, GrokModelCatalogSessionError, GrokModelCatalogSnapshot,
|
||||
GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind,
|
||||
GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, MAX_GROK_BILLING_BYTES,
|
||||
MAX_GROK_MODEL_CATALOG_BYTES, parse_grok_billing, parse_grok_model_catalog,
|
||||
};
|
||||
|
||||
@@ -409,7 +409,8 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport {
|
||||
.is_some_and(|total| total <= MAX_INFERENCE_BODY_BYTES) =>
|
||||
{
|
||||
*observed += chunk.len();
|
||||
Ok(chunk.to_vec())
|
||||
// Bytes 直传,长流不再逐 chunk 复制到 Vec。
|
||||
Ok(chunk)
|
||||
}
|
||||
Ok(_) => Err(GrokInferenceTransportError::new(
|
||||
GrokInferenceTransportErrorKind::Protocol,
|
||||
|
||||
@@ -311,6 +311,12 @@ pub enum GrokSessionSelectorError {
|
||||
/// Optional selector-derived delay.
|
||||
retry_after: Option<Duration>,
|
||||
},
|
||||
/// Every eligible session is inside its runtime cooldown window.
|
||||
#[error("Grok Build account is cooling down")]
|
||||
AccountCoolingDown {
|
||||
/// Remaining cooldown of the earliest recovering session.
|
||||
retry_after: Option<Duration>,
|
||||
},
|
||||
/// Session metadata or Provider-owned plaintext secret is invalid.
|
||||
#[error("Grok Build session data is invalid")]
|
||||
InvalidSession,
|
||||
|
||||
@@ -77,7 +77,7 @@ impl fmt::Debug for GrokInferenceRequest {
|
||||
|
||||
/// Stream of raw SSE byte chunks returned after the POST is accepted.
|
||||
pub type GrokInferenceChunkStream =
|
||||
Pin<Box<dyn Stream<Item = Result<Vec<u8>, GrokInferenceTransportError>> + Send + 'static>>;
|
||||
Pin<Box<dyn Stream<Item = Result<bytes::Bytes, GrokInferenceTransportError>> + Send + 'static>>;
|
||||
|
||||
/// Whether the account-isolated inference client was already cached.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
|
||||
@@ -775,6 +775,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;
|
||||
|
||||
@@ -283,7 +283,7 @@ fn minimal_oauth_credential_should_not_require_a_source_format() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oauth_account_document_should_not_require_unique_source_names() {
|
||||
fn duplicate_refresh_token_entries_should_collapse_to_the_first() {
|
||||
let mut document: serde_json::Value =
|
||||
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
|
||||
let duplicate = document["accounts"][0].clone();
|
||||
@@ -292,10 +292,29 @@ fn oauth_account_document_should_not_require_unique_source_names() {
|
||||
.expect("accounts array")
|
||||
.push(duplicate);
|
||||
|
||||
// 同一 refresh token 重复出现时解析层去重,避免逐条 RT exchange
|
||||
// 让先完成的轮换被后续重复条目作废。
|
||||
let parsed = GrokOAuthImportDocument::parse_json(
|
||||
&serde_json::to_vec(&document).expect("serialize fixture"),
|
||||
)
|
||||
.expect("upstream identity verification owns duplicate detection");
|
||||
.expect("duplicate refresh tokens collapse at parse");
|
||||
assert_eq!(parsed.into_entries().len(), 1);
|
||||
|
||||
// 不同 refresh token 的同名条目不受影响,重复检测仍归上游身份验证。
|
||||
let mut document: serde_json::Value =
|
||||
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
|
||||
let mut distinct = document["accounts"][0].clone();
|
||||
distinct["credentials"]["refresh_token"] =
|
||||
serde_json::Value::String("fixture-refresh-token-second".to_owned());
|
||||
document["accounts"]
|
||||
.as_array_mut()
|
||||
.expect("accounts array")
|
||||
.push(distinct);
|
||||
|
||||
let parsed = GrokOAuthImportDocument::parse_json(
|
||||
&serde_json::to_vec(&document).expect("serialize fixture"),
|
||||
)
|
||||
.expect("distinct refresh tokens keep their entries");
|
||||
|
||||
assert_eq!(parsed.into_entries().len(), 2);
|
||||
}
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::collections::{BTreeMap, VecDeque};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAvailability, CredentialCasOutcome, CredentialRevision, ProviderAccountStore,
|
||||
AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
|
||||
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
|
||||
};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError,
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError,
|
||||
};
|
||||
use provider_xai::{
|
||||
GrokCredentialCatalogCache, GrokCredentialRecovery, GrokCredentialRecoveryOutcome,
|
||||
@@ -141,6 +143,86 @@ async fn fixture(
|
||||
fixture_many([input], responses, lease_available).await
|
||||
}
|
||||
|
||||
/// 真实累加连续失败计数的测试 double,用于断言退避的指数增长与成功清零。
|
||||
#[derive(Default)]
|
||||
struct CountingCredentialState {
|
||||
counts: Mutex<BTreeMap<String, u32>>,
|
||||
cleared: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl CountingCredentialState {
|
||||
fn count(&self, account_id: &ProviderAccountId) -> u32 {
|
||||
self.counts
|
||||
.lock()
|
||||
.expect("backoff counts lock")
|
||||
.get(account_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn was_cleared(&self, account_id: &ProviderAccountId) -> bool {
|
||||
self.cleared
|
||||
.lock()
|
||||
.expect("cleared accounts lock")
|
||||
.iter()
|
||||
.any(|id| id == account_id.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCredentialStatePort for CountingCredentialState {
|
||||
fn replace(
|
||||
&self,
|
||||
_state: ProviderCredentialState,
|
||||
) -> futures::future::BoxFuture<'_, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> futures::future::BoxFuture<'a, Result<Option<ProviderCredentialState>, ProviderStoreError>>
|
||||
{
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
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 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,
|
||||
) -> futures::future::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 fixture_many(
|
||||
inputs: impl IntoIterator<Item = provider_xai::CreateGrokCredential>,
|
||||
responses: impl IntoIterator<Item = Result<GrokRefreshTokens, GrokRefreshFailure>>,
|
||||
@@ -150,6 +232,26 @@ async fn fixture_many(
|
||||
GrokCredentialRepository,
|
||||
Arc<QueueRefresher>,
|
||||
GrokCredentialRefreshService,
|
||||
) {
|
||||
fixture_many_with_state(
|
||||
inputs,
|
||||
responses,
|
||||
lease_available,
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn fixture_many_with_state(
|
||||
inputs: impl IntoIterator<Item = provider_xai::CreateGrokCredential>,
|
||||
responses: impl IntoIterator<Item = Result<GrokRefreshTokens, GrokRefreshFailure>>,
|
||||
lease_available: bool,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
) -> (
|
||||
Arc<MemoryProviderAccountStore>,
|
||||
GrokCredentialRepository,
|
||||
Arc<QueueRefresher>,
|
||||
GrokCredentialRefreshService,
|
||||
) {
|
||||
let store = MemoryProviderAccountStore::shared();
|
||||
let account_store: Arc<dyn ProviderAccountStore> = store.clone();
|
||||
@@ -176,11 +278,40 @@ async fn fixture_many(
|
||||
catalog,
|
||||
leases,
|
||||
cooldowns,
|
||||
credential_state,
|
||||
runtime_policy(),
|
||||
);
|
||||
(store, repository, refresher, service)
|
||||
}
|
||||
|
||||
/// 用 store 的 CAS 把 next_refresh_at 复位到过去,使已退避到未来的账号再次到期。
|
||||
async fn force_due(store: &MemoryProviderAccountStore, id: &ProviderAccountId) {
|
||||
let account = store.account(id).expect("account to reset");
|
||||
let credential = store.credential(id).expect("credential to reset");
|
||||
let update = CredentialCasUpdate::new(
|
||||
id.clone(),
|
||||
account.revision(),
|
||||
ProviderAccountUpdate {
|
||||
account_id: id.clone(),
|
||||
name: account.name().to_owned(),
|
||||
email: account.email().map(str::to_owned),
|
||||
plan_type: account.plan_type().map(str::to_owned),
|
||||
},
|
||||
credential,
|
||||
true,
|
||||
account.access_token_expires_at(),
|
||||
Some(SystemTime::now() - Duration::from_secs(1)),
|
||||
)
|
||||
.expect("valid due-time reset");
|
||||
assert!(matches!(
|
||||
store
|
||||
.compare_and_swap_credential(update)
|
||||
.await
|
||||
.expect("reset next_refresh_at to the past"),
|
||||
CredentialCasOutcome::Updated(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_refresh_rotates_plaintext_tokens_once() {
|
||||
let input = due_input("success");
|
||||
@@ -423,6 +554,95 @@ async fn pre_send_transient_failure_applies_bounded_cooldown() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_backoff_grows_exponentially_across_attempts() {
|
||||
let input = due_input("backoff-growth");
|
||||
let id = input.account_id.clone();
|
||||
let counting = Arc::new(CountingCredentialState::default());
|
||||
let (store, _, _, service) = fixture_many_with_state(
|
||||
[input],
|
||||
[
|
||||
Err(GrokRefreshFailure::Transient),
|
||||
Err(GrokRefreshFailure::Transient),
|
||||
],
|
||||
true,
|
||||
counting.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let before_first = SystemTime::now();
|
||||
let first = service.refresh_due().await.expect("first refresh cycle");
|
||||
assert!(matches!(
|
||||
first.as_slice(),
|
||||
[GrokCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let first_delay = store
|
||||
.account(&id)
|
||||
.expect("account after first defer")
|
||||
.next_refresh_at()
|
||||
.expect("first retry scheduled")
|
||||
.duration_since(before_first)
|
||||
.expect("first delay is in the future");
|
||||
assert_eq!(counting.count(&id), 1);
|
||||
|
||||
force_due(&store, &id).await;
|
||||
|
||||
let before_second = SystemTime::now();
|
||||
let second = service.refresh_due().await.expect("second refresh cycle");
|
||||
assert!(matches!(
|
||||
second.as_slice(),
|
||||
[GrokCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
let second_delay = store
|
||||
.account(&id)
|
||||
.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!(counting.count(&id), 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 input = due_input("backoff-clear");
|
||||
let id = input.account_id.clone();
|
||||
let counting = Arc::new(CountingCredentialState::default());
|
||||
let (store, _, _, service) = fixture_many_with_state(
|
||||
[input],
|
||||
[
|
||||
Err(GrokRefreshFailure::Transient),
|
||||
Ok(success_tokens(Some("cleared-refresh"))),
|
||||
],
|
||||
true,
|
||||
counting.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let first = service.refresh_due().await.expect("first refresh cycle");
|
||||
assert!(matches!(
|
||||
first.as_slice(),
|
||||
[GrokCredentialRefreshOutcome::Transient { .. }]
|
||||
));
|
||||
assert_eq!(counting.count(&id), 1);
|
||||
|
||||
force_due(&store, &id).await;
|
||||
|
||||
let second = service.refresh_due().await.expect("second refresh cycle");
|
||||
assert!(matches!(
|
||||
second.as_slice(),
|
||||
[GrokCredentialRefreshOutcome::Refreshed { .. }]
|
||||
));
|
||||
assert!(counting.was_cleared(&id));
|
||||
assert_eq!(counting.count(&id), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_refresh_lifetime_is_rejected_without_cas_write() {
|
||||
let input = due_input("bad-lifetime");
|
||||
|
||||
@@ -411,13 +411,16 @@ async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting
|
||||
.cooldown(session.account_id())
|
||||
.is_some_and(|cooldown| cooldown.until() > SystemTime::now())
|
||||
);
|
||||
assert!(matches!(
|
||||
fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await,
|
||||
Err(GrokSessionSelectorError::NoEligibleSession)
|
||||
));
|
||||
let retry_after = match fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
{
|
||||
Err(GrokSessionSelectorError::AccountCoolingDown { retry_after }) => retry_after,
|
||||
other => panic!("expected AccountCoolingDown, got {other:?}"),
|
||||
};
|
||||
let retry_after = retry_after.expect("retry_after");
|
||||
assert!(retry_after <= Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -325,20 +325,20 @@ impl GrokInferenceTransport for StubInferenceTransport {
|
||||
Box::pin(async move {
|
||||
match mode {
|
||||
InferenceMode::Success => Ok(GrokInferenceResponse::new(
|
||||
Box::pin(stream::iter([Ok(SUCCESS_SSE.to_vec())])),
|
||||
Box::pin(stream::iter([Ok(bytes::Bytes::from_static(SUCCESS_SSE))])),
|
||||
UpstreamHttpVersion::Http2,
|
||||
200,
|
||||
None,
|
||||
)),
|
||||
InferenceMode::SuccessWithMetrics(metrics) => Ok(GrokInferenceResponse::new(
|
||||
Box::pin(stream::iter([Ok(SUCCESS_SSE.to_vec())])),
|
||||
Box::pin(stream::iter([Ok(bytes::Bytes::from_static(SUCCESS_SSE))])),
|
||||
UpstreamHttpVersion::Http2,
|
||||
200,
|
||||
None,
|
||||
)
|
||||
.with_transport_metrics(metrics)),
|
||||
InferenceMode::SuccessBody(body) => Ok(GrokInferenceResponse::new(
|
||||
Box::pin(stream::iter([Ok(body)])),
|
||||
Box::pin(stream::iter([Ok(bytes::Bytes::from(body))])),
|
||||
UpstreamHttpVersion::Http2,
|
||||
200,
|
||||
None,
|
||||
@@ -1330,6 +1330,60 @@ async fn selector_capacity_failure_occurs_before_visible_upstream_send() {
|
||||
assert_eq!(transport.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn selection_failures_carry_a_distinguishable_client_error_code() {
|
||||
let cases = [
|
||||
(
|
||||
GrokSessionSelectorError::AccountCoolingDown {
|
||||
retry_after: Some(Duration::from_secs(12)),
|
||||
},
|
||||
"account_cooling_down",
|
||||
),
|
||||
(
|
||||
GrokSessionSelectorError::CapacityUnavailable {
|
||||
retry_after: Some(Duration::from_millis(20)),
|
||||
},
|
||||
"account_capacity_busy",
|
||||
),
|
||||
(
|
||||
GrokSessionSelectorError::NoEligibleSession,
|
||||
"no_eligible_account",
|
||||
),
|
||||
(
|
||||
GrokSessionSelectorError::Unavailable,
|
||||
"account_selector_unavailable",
|
||||
),
|
||||
];
|
||||
for (failure, expected_code) in cases {
|
||||
let transport = StubInferenceTransport::success();
|
||||
let provider = provider(StubSelector::failing(failure), transport.clone()).await;
|
||||
let error = match provider
|
||||
.execute(
|
||||
provider_request("xai"),
|
||||
context(CancellationToken::new(), None),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("selection must fail"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let detail = error
|
||||
.client_visible_upstream_error()
|
||||
.expect("client visible detail");
|
||||
assert_eq!(error.kind(), ProviderErrorKind::Unavailable);
|
||||
assert_eq!(detail.code(), Some(expected_code));
|
||||
assert_eq!(detail.error_type(), Some("account_unavailable_error"));
|
||||
assert_eq!(transport.calls.load(Ordering::SeqCst), 0);
|
||||
if expected_code == "account_cooling_down" {
|
||||
assert!(
|
||||
detail.message().contains("retry in 13s"),
|
||||
"{}",
|
||||
detail.message()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_previous_response_is_rejected_before_selection() {
|
||||
let selector = StubSelector::success();
|
||||
@@ -1827,7 +1881,7 @@ async fn provider_compiles_realtime_catalog_capabilities() {
|
||||
assert_eq!(presentation.display_name(), Some("Grok 4.5"));
|
||||
assert_eq!(
|
||||
presentation.supported_reasoning_efforts(),
|
||||
["low", "medium", "high"]
|
||||
["low", "medium", "high", "xhigh"]
|
||||
);
|
||||
assert_eq!(presentation.context_window_tokens(), Some(1_000_000));
|
||||
assert!(presentation.agent_tools());
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user