mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-12 08:06:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3328a4356a | ||
|
|
77b0c29354 | ||
|
|
c872bc7881 | ||
|
|
487b2e87e2 | ||
|
|
eb8a01541a | ||
|
|
0acf114fa7 | ||
|
|
5d941008fc | ||
|
|
b44a790777 | ||
|
|
06e94ef69b | ||
|
|
0f2e0b18a9 | ||
|
|
71e26f44b7 | ||
|
|
8232e19bc0 | ||
|
|
dcc61915e9 | ||
|
|
47d08560a7 | ||
|
|
590c4a163b |
@@ -669,9 +669,6 @@ jobs:
|
||||
shell: bash
|
||||
env:
|
||||
INSTALLATION_NOTES: |
|
||||
> [!WARNING]
|
||||
> **跨大版本升级:** 不支持跨大版本在线升级。请使用全新的 `.runtime/` 数据目录重新部署,再重新导入或重新授权服务商账号,并重新导入客户端密钥。详见[升级说明](https://github.com/${{ github.repository }}/blob/${{ needs.release-info.outputs.tag }}/deploy/README.md#镜像升级与源码构建)。
|
||||
|
||||
## 📥 安装
|
||||
|
||||
**Docker 部署:**
|
||||
|
||||
@@ -79,6 +79,7 @@ pub struct AccountGroupListQuery {
|
||||
/// Complete account group summary.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupRecord {
|
||||
pub disable_fast: bool,
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
@@ -107,6 +108,7 @@ pub struct AccountGroupPage {
|
||||
/// Create an account group.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateAccountGroup {
|
||||
pub disable_fast: bool,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub color: AccountGroupColor,
|
||||
@@ -115,6 +117,7 @@ pub struct CreateAccountGroup {
|
||||
/// Store-ready create command with a generated stable ID.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct NewAccountGroup {
|
||||
pub disable_fast: bool,
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
@@ -124,6 +127,7 @@ pub struct NewAccountGroup {
|
||||
/// Update an account group's descriptive fields.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateAccountGroup {
|
||||
pub disable_fast: Option<bool>,
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
|
||||
@@ -544,6 +544,7 @@ pub struct UsageListRecord {
|
||||
pub provider_account_authentication_kind: Option<String>,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
@@ -602,6 +603,7 @@ pub struct UsageRecord {
|
||||
pub upstream_transport: Option<String>,
|
||||
pub http_version: Option<String>,
|
||||
pub websocket_pool: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
/// Provider 已筛选的专有观测 JSON;管理领域保持不透明。
|
||||
pub provider_metadata_json: Option<String>,
|
||||
|
||||
@@ -60,6 +60,8 @@ pub struct ProxyTestResult {
|
||||
pub success: bool,
|
||||
pub latency_ms: u64,
|
||||
pub exit_ip: Option<std::net::IpAddr>,
|
||||
pub exit_ipv4: Option<std::net::Ipv4Addr>,
|
||||
pub exit_ipv6: Option<std::net::Ipv6Addr>,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ pub use gateway_core::account::RotationStrategy;
|
||||
/// 完整运行设置事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RuntimeSettings {
|
||||
pub disable_fast: bool,
|
||||
pub config_revision: Revision,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
@@ -28,6 +29,7 @@ pub struct RuntimeSettings {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: RotationStrategy,
|
||||
pub min_codex_desktop_version: Option<String>,
|
||||
pub min_codex_cli_version: Option<String>,
|
||||
@@ -47,6 +49,7 @@ pub struct RuntimeSettings {
|
||||
/// 原子替换运行设置的命令。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ReplaceRuntimeSettings {
|
||||
pub disable_fast: Option<bool>,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
pub model_mappings: ModelMappings,
|
||||
@@ -57,6 +60,7 @@ pub struct ReplaceRuntimeSettings {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: RotationStrategy,
|
||||
pub min_codex_desktop_version: Option<String>,
|
||||
pub min_codex_cli_version: Option<String>,
|
||||
|
||||
@@ -144,6 +144,7 @@ impl AccountGroupService for DefaultAccountGroupService {
|
||||
id,
|
||||
name: command.name,
|
||||
description: command.description,
|
||||
disable_fast: command.disable_fast,
|
||||
color: command.color,
|
||||
},
|
||||
context,
|
||||
|
||||
@@ -191,9 +191,6 @@ async fn pending_authorization(
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if !record.last_test.is_some_and(|test| test.success) {
|
||||
return Err(AdminError::conflict("请先测试代理连接"));
|
||||
}
|
||||
proxy_id = Some(record.id);
|
||||
Some(record.proxy)
|
||||
}
|
||||
|
||||
@@ -115,6 +115,8 @@ impl SettingsService for DefaultSettingsService {
|
||||
|
||||
fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError> {
|
||||
let valid = command.request_location.validate().is_ok()
|
||||
&& command.responses_max_decompressed_body_bytes > 0
|
||||
&& isize::try_from(command.responses_max_decompressed_body_bytes).is_ok()
|
||||
&& command.refresh_margin_seconds > 0
|
||||
&& command.refresh_concurrency > 0
|
||||
&& command.max_concurrent_per_account > 0
|
||||
|
||||
@@ -202,6 +202,7 @@ impl AccountRuntimeStore for FakeRuntimeStore {
|
||||
fn group_record() -> AccountGroupRecord {
|
||||
let now = Utc::now();
|
||||
AccountGroupRecord {
|
||||
disable_fast: false,
|
||||
id: group_id(),
|
||||
name: "Primary".to_owned(),
|
||||
description: None,
|
||||
|
||||
@@ -973,6 +973,7 @@ struct StaticSettingsStore;
|
||||
impl SettingsStore for StaticSettingsStore {
|
||||
async fn load_runtime_settings(&self) -> AdminStoreResult<RuntimeSettings> {
|
||||
Ok(RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: revision(1),
|
||||
@@ -984,6 +985,7 @@ impl SettingsStore for StaticSettingsStore {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
min_codex_desktop_version: None,
|
||||
min_codex_cli_version: None,
|
||||
|
||||
@@ -25,6 +25,7 @@ use super::accounts::{FakeAccountStore, FakeProviderAdmin, account_record, event
|
||||
|
||||
fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> RuntimeSettings {
|
||||
RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: revision(1),
|
||||
@@ -36,6 +37,7 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: gateway_admin::model::settings::RotationStrategy::Smart,
|
||||
min_codex_desktop_version: None,
|
||||
min_codex_cli_version: None,
|
||||
|
||||
@@ -753,6 +753,7 @@ struct FixtureSettingsStore;
|
||||
impl SettingsStore for FixtureSettingsStore {
|
||||
async fn load_runtime_settings(&self) -> AdminStoreResult<RuntimeSettings> {
|
||||
Ok(RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: Revision::new(1).expect("revision"),
|
||||
@@ -764,6 +765,7 @@ impl SettingsStore for FixtureSettingsStore {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
min_codex_desktop_version: None,
|
||||
min_codex_cli_version: None,
|
||||
@@ -860,6 +862,7 @@ fn total_record(
|
||||
provider_account_authentication_kind: None,
|
||||
upstream_model_id: Some("gpt-5.5".to_owned()),
|
||||
upstream_transport: None,
|
||||
upstream_response_model: None,
|
||||
service_tier: None,
|
||||
input_tokens: Some(800),
|
||||
output_tokens: Some(200),
|
||||
|
||||
@@ -12,6 +12,7 @@ use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
pub(super) struct TestProxies {
|
||||
pub events: Option<super::accounts::EventLog>,
|
||||
pub accounts: Option<Vec<ProxyAccountRef>>,
|
||||
pub record: Option<ProxyRecord>,
|
||||
}
|
||||
|
||||
struct ImportGuard(super::accounts::EventLog);
|
||||
@@ -70,8 +71,11 @@ impl ProxyStore for TestProxies {
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn get(&self, _: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
Err(super::unavailable("proxy"))
|
||||
async fn get(&self, id: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
self.record
|
||||
.clone()
|
||||
.filter(|record| record.id == id)
|
||||
.ok_or_else(|| super::unavailable("proxy"))
|
||||
}
|
||||
async fn create(&self, _: NewProxy, _: &MutationContext) -> AdminStoreResult<ProxyMutation> {
|
||||
Err(super::unavailable("proxy"))
|
||||
@@ -105,6 +109,60 @@ impl ProxyProbe for TestProxies {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorization_uses_selected_proxy_regardless_of_probe_status() {
|
||||
use super::accounts::{FakeAccountStore, FakeProviderAdmin, context, events, recorded};
|
||||
use gateway_admin::model::provider_credentials::StartAuthorization;
|
||||
use std::sync::Arc;
|
||||
|
||||
for kind in ["openai", "xai"] {
|
||||
for probe_success in [None, Some(false), Some(true)] {
|
||||
let events = events();
|
||||
let now = chrono::Utc::now();
|
||||
let services = super::AdminHarness::new()
|
||||
.provider(FakeProviderAdmin::new(kind, events.clone()))
|
||||
.accounts(FakeAccountStore::new(kind, events.clone()))
|
||||
.proxies(Arc::new(TestProxies {
|
||||
record: Some(ProxyRecord {
|
||||
location: None,
|
||||
id: "proxy_oauth".to_owned(),
|
||||
name: "授权出口".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:8080").unwrap(),
|
||||
revision: Revision::new(1).unwrap(),
|
||||
account_count: 0,
|
||||
last_test_at: probe_success.map(|_| now),
|
||||
last_test: probe_success.map(|success| ProxyTestResult {
|
||||
success,
|
||||
latency_ms: 10,
|
||||
exit_ip: None,
|
||||
exit_ipv4: None,
|
||||
exit_ipv6: None,
|
||||
message: "出口探测结果".to_owned(),
|
||||
}),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
}),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let command = StartAuthorization {
|
||||
outbound_proxy: Some(AccountProxySelection::Saved("proxy_oauth".to_owned())),
|
||||
context: context("oauth-proxy-status"),
|
||||
name: "授权账号".to_owned(),
|
||||
reauthorization: None,
|
||||
};
|
||||
let result = if kind == "openai" {
|
||||
services.openai().start_authorization(command).await
|
||||
} else {
|
||||
services.xai().start_authorization(command).await
|
||||
};
|
||||
assert!(result.is_ok(), "{kind}, {probe_success:?}: {result:?}");
|
||||
assert_eq!(recorded(&events), ["provider.start_authorization"]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn linked_accounts_share_plan_resolution_and_only_read_cached_quota() {
|
||||
use super::accounts::{FakeProviderAdmin, events};
|
||||
|
||||
@@ -61,6 +61,7 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
|
||||
request_id: "request-settings".to_owned(),
|
||||
},
|
||||
ReplaceRuntimeSettings {
|
||||
disable_fast: None,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
model_mappings: Default::default(),
|
||||
@@ -71,6 +72,7 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
min_codex_desktop_version: None,
|
||||
min_codex_cli_version: None,
|
||||
|
||||
@@ -72,6 +72,8 @@ impl ListAccountGroupsQuery {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct CreateAccountGroupRequest {
|
||||
#[serde(default)]
|
||||
disable_fast: bool,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
color: String,
|
||||
@@ -80,6 +82,7 @@ struct CreateAccountGroupRequest {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct UpdateAccountGroupRequest {
|
||||
disable_fast: Option<bool>,
|
||||
id: String,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
@@ -95,6 +98,7 @@ struct AccountGroupIdRequest {
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupView {
|
||||
disable_fast: bool,
|
||||
id: String,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
@@ -138,6 +142,7 @@ impl From<AccountGroupRecord> for AccountGroupView {
|
||||
id: record.id.to_string(),
|
||||
name: record.name,
|
||||
description: record.description,
|
||||
disable_fast: record.disable_fast,
|
||||
color: record.color.as_str().to_owned(),
|
||||
enabled: record.enabled,
|
||||
member_count: record.member_count,
|
||||
@@ -273,6 +278,7 @@ where
|
||||
CreateAccountGroup {
|
||||
name: request.name,
|
||||
description: request.description,
|
||||
disable_fast: request.disable_fast,
|
||||
color: group_color(&request.color)?,
|
||||
},
|
||||
)
|
||||
@@ -300,6 +306,7 @@ where
|
||||
id: group_id(request.id)?,
|
||||
name: request.name,
|
||||
description: request.description,
|
||||
disable_fast: request.disable_fast,
|
||||
color: group_color(&request.color)?,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -267,6 +267,7 @@ pub(crate) fn usage_list_record_view(record: domain::UsageListRecord) -> UsageLi
|
||||
model,
|
||||
requested_model: record.requested_model_id,
|
||||
upstream_model: record.upstream_model_id,
|
||||
upstream_response_model: record.upstream_response_model,
|
||||
service_tier: record.service_tier,
|
||||
client_transport: record.client_transport,
|
||||
upstream_transport: record.upstream_transport,
|
||||
@@ -360,6 +361,7 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
model,
|
||||
requested_model: record.requested_model_id,
|
||||
upstream_model: record.upstream_model_id,
|
||||
upstream_response_model: record.upstream_response_model,
|
||||
service_tier: record.service_tier,
|
||||
status_code,
|
||||
client_transport: record.client_transport,
|
||||
|
||||
@@ -84,6 +84,7 @@ pub struct UsageListRecordView {
|
||||
pub model: Option<String>,
|
||||
pub requested_model: Option<String>,
|
||||
pub upstream_model: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub client_transport: String,
|
||||
pub upstream_transport: Option<String>,
|
||||
@@ -122,6 +123,7 @@ pub struct UsageRecordView {
|
||||
pub model: Option<String>,
|
||||
pub requested_model: Option<String>,
|
||||
pub upstream_model: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub client_transport: String,
|
||||
|
||||
@@ -92,6 +92,8 @@ struct ProxyTestView {
|
||||
success: bool,
|
||||
latency_ms: u64,
|
||||
exit_ip: Option<String>,
|
||||
exit_ipv4: Option<String>,
|
||||
exit_ipv6: Option<String>,
|
||||
message: String,
|
||||
}
|
||||
|
||||
@@ -101,6 +103,8 @@ impl From<ProxyTestResult> for ProxyTestView {
|
||||
success: result.success,
|
||||
latency_ms: result.latency_ms,
|
||||
exit_ip: result.exit_ip.map(|ip| ip.to_string()),
|
||||
exit_ipv4: result.exit_ipv4.map(|ip| ip.to_string()),
|
||||
exit_ipv6: result.exit_ipv6.map(|ip| ip.to_string()),
|
||||
message: result.message,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -34,6 +34,7 @@ pub type ModelMappings = BTreeMap<String, String>;
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RuntimeSettingsView {
|
||||
pub disable_fast: bool,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
pub model_mappings: ModelMappings,
|
||||
@@ -44,6 +45,7 @@ pub struct RuntimeSettingsView {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: String,
|
||||
pub min_codex_desktop_version: Option<String>,
|
||||
pub min_codex_cli_version: Option<String>,
|
||||
@@ -64,6 +66,7 @@ pub struct RuntimeSettingsView {
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct UpdateRuntimeSettingsRequest {
|
||||
pub disable_fast: Option<bool>,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
pub model_mappings: ModelMappings,
|
||||
@@ -74,6 +77,7 @@ pub struct UpdateRuntimeSettingsRequest {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: String,
|
||||
pub min_codex_desktop_version: Option<String>,
|
||||
pub min_codex_cli_version: Option<String>,
|
||||
@@ -104,6 +108,13 @@ impl UpdateRuntimeSettingsRequest {
|
||||
return Err(WireValidationError::new(field));
|
||||
}
|
||||
}
|
||||
if self.responses_max_decompressed_body_bytes == 0
|
||||
|| isize::try_from(self.responses_max_decompressed_body_bytes).is_err()
|
||||
{
|
||||
return Err(WireValidationError::new(
|
||||
"responsesMaxDecompressedBodyBytes",
|
||||
));
|
||||
}
|
||||
if !(1..=120).contains(&self.concurrency_wait_timeout_seconds) {
|
||||
return Err(WireValidationError::new("concurrencyWaitTimeoutSeconds"));
|
||||
}
|
||||
@@ -166,6 +177,7 @@ impl UpdateRuntimeSettingsRequest {
|
||||
fn into_command(self) -> Result<ReplaceRuntimeSettings, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(ReplaceRuntimeSettings {
|
||||
disable_fast: self.disable_fast,
|
||||
request_location_enabled: self.request_location_enabled,
|
||||
request_location: self
|
||||
.request_location
|
||||
@@ -181,6 +193,7 @@ impl UpdateRuntimeSettingsRequest {
|
||||
max_waiting_per_key: self.max_waiting_per_key,
|
||||
max_waiting_per_account: self.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: self.concurrency_wait_timeout_seconds,
|
||||
responses_max_decompressed_body_bytes: self.responses_max_decompressed_body_bytes,
|
||||
rotation_strategy: RotationStrategy::parse(&self.rotation_strategy)
|
||||
.ok_or_else(|| WireValidationError::new("rotationStrategy"))?,
|
||||
min_codex_desktop_version: self.min_codex_desktop_version,
|
||||
@@ -206,6 +219,7 @@ impl UpdateRuntimeSettingsRequest {
|
||||
impl From<RuntimeSettings> for RuntimeSettingsView {
|
||||
fn from(settings: RuntimeSettings) -> Self {
|
||||
Self {
|
||||
disable_fast: settings.disable_fast,
|
||||
request_location_enabled: settings.request_location_enabled,
|
||||
request_location: settings.request_location,
|
||||
model_mappings: wire_model_mappings(settings.model_mappings),
|
||||
@@ -216,6 +230,7 @@ impl From<RuntimeSettings> for RuntimeSettingsView {
|
||||
max_waiting_per_key: settings.max_waiting_per_key,
|
||||
max_waiting_per_account: settings.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
|
||||
responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes,
|
||||
rotation_strategy: settings.rotation_strategy.as_str().to_owned(),
|
||||
min_codex_desktop_version: settings.min_codex_desktop_version,
|
||||
min_codex_cli_version: settings.min_codex_cli_version,
|
||||
|
||||
@@ -143,7 +143,7 @@ pub(crate) fn client_access_error_response(error: ClientAccessError) -> Response
|
||||
}
|
||||
}
|
||||
|
||||
/// Desktop 优先于其内嵌的 CLI/Core 标记;未知客户端保持兼容并返回 `None`。
|
||||
/// Desktop 优先于内嵌 CLI/Core;未知客户端及未提供应用版本的手机远程客户端返回 `None`。
|
||||
#[must_use]
|
||||
pub fn identify_codex_client(headers: &HeaderMap) -> Option<IdentifiedCodexClient> {
|
||||
const MAXIMUM_HEADER_LENGTH: usize = 4096;
|
||||
@@ -154,11 +154,21 @@ pub fn identify_codex_client(headers: &HeaderMap) -> Option<IdentifiedCodexClien
|
||||
|| user_agent.is_some_and(|value| contains_ascii_case_insensitive(value, "Codex Desktop"));
|
||||
if is_desktop {
|
||||
let explicit_version = bounded_ascii_header(headers, "version", MAXIMUM_HEADER_LENGTH);
|
||||
let user_agent_version = user_agent.and_then(desktop_version_from_user_agent);
|
||||
// 手机远程初始化会覆盖进程级 UA 后缀,保留的 Desktop/Core 前缀不代表应用版本。
|
||||
// 仅缺失版本信息时免于门禁;已提供的版本即使非法,也继续按 Desktop 校验。
|
||||
if !headers.contains_key("version")
|
||||
&& user_agent_version.is_none()
|
||||
&& headers
|
||||
.get("user-agent")
|
||||
.is_some_and(|value| value.as_bytes().len() <= MAXIMUM_HEADER_LENGTH)
|
||||
&& user_agent.is_some_and(has_chatgpt_remote_user_agent_suffix)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let version = match explicit_version {
|
||||
Some(value) => CodexClientVersion::parse(value).ok(),
|
||||
None => user_agent
|
||||
.and_then(desktop_version_from_user_agent)
|
||||
.and_then(|value| CodexClientVersion::parse(value).ok()),
|
||||
None => user_agent_version.and_then(|value| CodexClientVersion::parse(value).ok()),
|
||||
};
|
||||
return Some(IdentifiedCodexClient {
|
||||
kind: CodexClientKind::Desktop,
|
||||
@@ -206,11 +216,34 @@ fn bounded_ascii_header<'a>(
|
||||
fn desktop_version_from_user_agent(user_agent: &str) -> Option<&str> {
|
||||
let marker = "Codex Desktop;";
|
||||
let start = find_ascii_case_insensitive(user_agent, marker)? + marker.len();
|
||||
let version = user_agent[start..]
|
||||
// 空候选也代表已提供应用版本,不能被手机远程的缺失版本规则放行。
|
||||
user_agent[start..]
|
||||
.trim_start()
|
||||
.split([')', ' ', ';', ','])
|
||||
.next()?;
|
||||
(!version.is_empty()).then_some(version)
|
||||
.next()
|
||||
}
|
||||
|
||||
fn has_chatgpt_remote_user_agent_suffix(user_agent: &str) -> bool {
|
||||
let Some((_, suffix)) = user_agent
|
||||
.strip_suffix(')')
|
||||
.and_then(|value| value.rsplit_once(" ("))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some((name, version)) = suffix.split_once("; ") else {
|
||||
return false;
|
||||
};
|
||||
let Some(platform) = name
|
||||
.strip_prefix("codex_chatgpt_")
|
||||
.and_then(|value| value.strip_suffix("_remote"))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
[platform, version].into_iter().all(|value| {
|
||||
!value.is_empty()
|
||||
&& !value.contains(['(', ')', ';'])
|
||||
&& !value.bytes().any(|byte| byte.is_ascii_whitespace())
|
||||
})
|
||||
}
|
||||
|
||||
fn contains_ascii_case_insensitive(haystack: &str, needle: &str) -> bool {
|
||||
|
||||
@@ -42,7 +42,10 @@ pub enum RequestDecodeError {
|
||||
MalformedJson,
|
||||
/// 请求体解压后超过允许上限。
|
||||
#[error("decompressed request body is too large")]
|
||||
DecompressedBodyTooLarge,
|
||||
DecompressedBodyTooLarge {
|
||||
/// 当前请求快照中的解压输出上限。
|
||||
limit_bytes: usize,
|
||||
},
|
||||
/// 请求使用了网关不支持的 `Content-Encoding`。
|
||||
#[error("request content-encoding `{encoding}` is not supported")]
|
||||
UnsupportedContentEncoding {
|
||||
@@ -108,9 +111,11 @@ impl RequestDecodeError {
|
||||
"Request body must be valid JSON.".to_owned(),
|
||||
None,
|
||||
),
|
||||
Self::DecompressedBodyTooLarge => (
|
||||
Self::DecompressedBodyTooLarge { limit_bytes } => (
|
||||
"request_too_large",
|
||||
"Decompressed request body exceeds the allowed size.".to_owned(),
|
||||
format!(
|
||||
"Decompressed request body exceeds the allowed size ({limit_bytes} bytes)."
|
||||
),
|
||||
None,
|
||||
),
|
||||
Self::UnsupportedContentEncoding { encoding } => (
|
||||
|
||||
@@ -52,7 +52,11 @@ pub(crate) async fn responses(
|
||||
Ok(client) => client,
|
||||
Err(error) => return client_access_error_response(error),
|
||||
};
|
||||
let decoded = match decode_request_with_headers(&body, &headers) {
|
||||
let decoded = match decode_request_with_headers(
|
||||
&body,
|
||||
&headers,
|
||||
client.snapshot().responses_max_decompressed_body_bytes(),
|
||||
) {
|
||||
Ok(decoded) => decoded,
|
||||
Err(error) => {
|
||||
return protocol_error_response(StatusCode::BAD_REQUEST, error.protocol_body());
|
||||
|
||||
@@ -327,13 +327,14 @@ impl fmt::Debug for DecodedResponsesRequest {
|
||||
pub fn decode_request_with_headers(
|
||||
body: &[u8],
|
||||
headers: &HeaderMap,
|
||||
max_decompressed_bytes: usize,
|
||||
) -> Result<DecodedResponsesRequest, RequestDecodeError> {
|
||||
let body = decompress_request_body(body, headers)?;
|
||||
let body = decompress_request_body(body, headers, max_decompressed_bytes)?;
|
||||
decode_request_inner(&body, &OpenAiRequestHeaders::from_headers(headers))
|
||||
}
|
||||
|
||||
/// 下游请求体解压后的最大字节数;防御解压炸弹。
|
||||
const MAX_DECOMPRESSED_REQUEST_BYTES: usize = 64 * 1024 * 1024;
|
||||
/// zstd 回溯窗口独立于输出上限,调整设置不能放大解码器内部窗口分配。
|
||||
const MAX_ZSTD_WINDOW_LOG: u32 = 26;
|
||||
|
||||
/// 按 `Content-Encoding` 解压下游请求体。
|
||||
///
|
||||
@@ -342,6 +343,7 @@ const MAX_DECOMPRESSED_REQUEST_BYTES: usize = 64 * 1024 * 1024;
|
||||
fn decompress_request_body<'a>(
|
||||
body: &'a [u8],
|
||||
headers: &HeaderMap,
|
||||
max_decompressed_bytes: usize,
|
||||
) -> Result<Cow<'a, [u8]>, RequestDecodeError> {
|
||||
let encoding = headers
|
||||
.get_all(axum::http::header::CONTENT_ENCODING)
|
||||
@@ -357,20 +359,20 @@ fn decompress_request_body<'a>(
|
||||
"" | "identity" => Ok(Cow::Borrowed(body)),
|
||||
"gzip" => {
|
||||
let decoder = flate2::read::MultiGzDecoder::new(body);
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
read_bounded(decoder, max_decompressed_bytes).map(Cow::Owned)
|
||||
}
|
||||
"deflate" => {
|
||||
let decoder = flate2::read::ZlibDecoder::new(body);
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
read_bounded(decoder, max_decompressed_bytes).map(Cow::Owned)
|
||||
}
|
||||
"zstd" => {
|
||||
let mut decoder = zstd::stream::read::Decoder::with_buffer(body)
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
// 输出有界之外,也限制压缩帧声明的回溯窗口,防止解码器内部过量分配。
|
||||
decoder
|
||||
.window_log_max(MAX_DECOMPRESSED_REQUEST_BYTES.ilog2())
|
||||
.window_log_max(MAX_ZSTD_WINDOW_LOG)
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
read_bounded(decoder, max_decompressed_bytes).map(Cow::Owned)
|
||||
}
|
||||
other => Err(RequestDecodeError::UnsupportedContentEncoding {
|
||||
encoding: other.to_owned(),
|
||||
@@ -378,12 +380,15 @@ fn decompress_request_body<'a>(
|
||||
}
|
||||
}
|
||||
|
||||
fn read_bounded<R: std::io::Read>(mut reader: R) -> Result<Vec<u8>, RequestDecodeError> {
|
||||
fn read_bounded<R: std::io::Read>(
|
||||
mut reader: R,
|
||||
max_decompressed_bytes: usize,
|
||||
) -> Result<Vec<u8>, RequestDecodeError> {
|
||||
let mut decoded = Vec::new();
|
||||
let mut chunk = [0; 8192];
|
||||
loop {
|
||||
let remaining = MAX_DECOMPRESSED_REQUEST_BYTES - decoded.len();
|
||||
let read_limit = chunk.len().min(remaining + 1);
|
||||
let remaining = max_decompressed_bytes - decoded.len();
|
||||
let read_limit = chunk.len().min(remaining.saturating_add(1));
|
||||
let read = reader
|
||||
.read(&mut chunk[..read_limit])
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
@@ -391,15 +396,19 @@ fn read_bounded<R: std::io::Read>(mut reader: R) -> Result<Vec<u8>, RequestDecod
|
||||
return Ok(decoded);
|
||||
}
|
||||
if read > remaining {
|
||||
return Err(RequestDecodeError::DecompressedBodyTooLarge);
|
||||
return Err(RequestDecodeError::DecompressedBodyTooLarge {
|
||||
limit_bytes: max_decompressed_bytes,
|
||||
});
|
||||
}
|
||||
// 默认 Vec 扩容和 read_to_end 的 EOF 探测可能突破输出上限;
|
||||
// 容量增长也受相同边界约束,越界探测只使用栈上分块缓冲区。
|
||||
let required = decoded.len() + read;
|
||||
if required > decoded.capacity() {
|
||||
let capacity = (decoded.capacity() * 2)
|
||||
let capacity = decoded
|
||||
.capacity()
|
||||
.saturating_mul(2)
|
||||
.max(required)
|
||||
.min(MAX_DECOMPRESSED_REQUEST_BYTES);
|
||||
.min(max_decompressed_bytes);
|
||||
decoded.reserve_exact(capacity - decoded.len());
|
||||
}
|
||||
decoded.extend_from_slice(&chunk[..read]);
|
||||
|
||||
@@ -272,3 +272,27 @@ async fn response_json(response: axum::response::Response) -> Value {
|
||||
.expect("account group response body");
|
||||
serde_json::from_slice(&body).expect("account group response JSON")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disable_fast_group_updates_preserve_omitted_values() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
for (value, expected) in [(Some(true), true), (None, true), (Some(false), false)] {
|
||||
let mut body = json!({"id": PRIMARY_GROUP_ID, "name":"Policy", "description":null,"color":"#F43F5ECC"});
|
||||
if let Some(value) = value {
|
||||
body["disableFast"] = json!(value);
|
||||
}
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/update",
|
||||
Some(body),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response_json(response).await["data"]["record"]["disableFast"],
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -402,6 +402,7 @@ impl SettingsStore for MemorySettingsStore {
|
||||
) -> AdminStoreResult<RuntimeSettings> {
|
||||
let mut settings = self.settings.lock().expect("settings");
|
||||
let updated = RuntimeSettings {
|
||||
disable_fast: command.disable_fast.unwrap_or(settings.disable_fast),
|
||||
request_location_enabled: command.request_location_enabled,
|
||||
request_location: command.request_location,
|
||||
config_revision: next_revision(settings.config_revision),
|
||||
@@ -413,6 +414,7 @@ impl SettingsStore for MemorySettingsStore {
|
||||
max_waiting_per_key: command.max_waiting_per_key,
|
||||
max_waiting_per_account: command.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds,
|
||||
responses_max_decompressed_body_bytes: command.responses_max_decompressed_body_bytes,
|
||||
rotation_strategy: command.rotation_strategy,
|
||||
min_codex_desktop_version: command.min_codex_desktop_version,
|
||||
min_codex_cli_version: command.min_codex_cli_version,
|
||||
@@ -491,6 +493,7 @@ impl MemoryAccountGroupStore {
|
||||
(
|
||||
primary_id.clone(),
|
||||
AccountGroupRecord {
|
||||
disable_fast: false,
|
||||
id: primary_id,
|
||||
name: "Alpha routing".to_owned(),
|
||||
description: Some("Primary traffic".to_owned()),
|
||||
@@ -512,6 +515,7 @@ impl MemoryAccountGroupStore {
|
||||
(
|
||||
secondary_id.clone(),
|
||||
AccountGroupRecord {
|
||||
disable_fast: false,
|
||||
id: secondary_id,
|
||||
name: "Beta routing".to_owned(),
|
||||
description: None,
|
||||
@@ -627,6 +631,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
let mut state = self.state.lock().expect("account groups");
|
||||
let now = Utc::now();
|
||||
let record = AccountGroupRecord {
|
||||
disable_fast: command.disable_fast,
|
||||
id: command.id.clone(),
|
||||
name: command.name,
|
||||
description: command.description,
|
||||
@@ -658,6 +663,9 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
record.name = command.name;
|
||||
record.description = command.description;
|
||||
record.color = command.color;
|
||||
if let Some(disable_fast) = command.disable_fast {
|
||||
record.disable_fast = disable_fast;
|
||||
}
|
||||
record.updated_at = Utc::now();
|
||||
mutation(&mut state, command.id, true)
|
||||
}
|
||||
@@ -1270,6 +1278,7 @@ fn test_runtime_settings() -> RuntimeSettings {
|
||||
),
|
||||
]);
|
||||
RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: Revision::new(7).expect("revision"),
|
||||
@@ -1281,6 +1290,7 @@ fn test_runtime_settings() -> RuntimeSettings {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
min_codex_desktop_version: None,
|
||||
min_codex_cli_version: None,
|
||||
|
||||
@@ -840,6 +840,7 @@ fn usage_record_with_account(
|
||||
upstream_transport: Some("http_sse".to_owned()),
|
||||
http_version: Some("h2".to_owned()),
|
||||
websocket_pool: None,
|
||||
upstream_response_model: None,
|
||||
service_tier: None,
|
||||
provider_metadata_json: None,
|
||||
attempt_count: 1,
|
||||
@@ -963,6 +964,7 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
|
||||
provider_account_authentication_kind: Some("oauth".to_owned()),
|
||||
upstream_model_id: Some("grok-4.5".to_owned()),
|
||||
upstream_transport: Some("http_sse".to_owned()),
|
||||
upstream_response_model: Some("grok-4.6-build".to_owned()),
|
||||
service_tier: Some("default".to_owned()),
|
||||
input_tokens: Some(1),
|
||||
output_tokens: Some(1),
|
||||
@@ -1075,6 +1077,7 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
|
||||
"imageOutputTokens": value["data"]["items"][0]["tokenDetails"]["imageOutputTokens"],
|
||||
"requestedModel": value["data"]["items"][0]["requestedModel"],
|
||||
"upstreamModel": value["data"]["items"][0]["upstreamModel"],
|
||||
"upstreamResponseModel": value["data"]["items"][0]["upstreamResponseModel"],
|
||||
"reasoningEffort": value["data"]["items"][0]["reasoningEffort"],
|
||||
"reasoningPreset": value["data"]["items"][0]["reasoningPreset"],
|
||||
"subagentKind": value["data"]["items"][0]["subagentKind"],
|
||||
@@ -1098,6 +1101,7 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
|
||||
"imageOutputTokens": 9,
|
||||
"requestedModel": "grok-4.5",
|
||||
"upstreamModel": "grok-4.5",
|
||||
"upstreamResponseModel": "grok-4.6-build",
|
||||
"reasoningEffort": "max",
|
||||
"reasoningPreset": "ultra",
|
||||
"subagentKind": "worker",
|
||||
|
||||
@@ -237,6 +237,8 @@ impl ProxyProbe for SuccessfulProbe {
|
||||
success: true,
|
||||
latency_ms: 15,
|
||||
exit_ip: Some("203.0.113.2".parse().unwrap()),
|
||||
exit_ipv4: Some("203.0.113.2".parse().unwrap()),
|
||||
exit_ipv6: None,
|
||||
message: "Connected".to_owned(),
|
||||
}
|
||||
}
|
||||
@@ -390,7 +392,14 @@ async fn proxy_probe_checks_unsaved_address_without_creating_or_changing_records
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
probed["data"],
|
||||
json!({"success": true, "latencyMs": 15, "exitIp": "203.0.113.2", "message": "Connected"})
|
||||
json!({
|
||||
"success": true,
|
||||
"latencyMs": 15,
|
||||
"exitIp": "203.0.113.2",
|
||||
"exitIpv4": "203.0.113.2",
|
||||
"exitIpv6": null,
|
||||
"message": "Connected"
|
||||
})
|
||||
);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["page"]["total"], 0);
|
||||
|
||||
@@ -37,6 +37,7 @@ async fn response_json(response: axum::response::Response) -> Value {
|
||||
|
||||
fn update_body() -> Value {
|
||||
json!({
|
||||
"disableFast": false,
|
||||
"requestLocationEnabled": false,
|
||||
"requestLocation": {"country":"US", "region":"Ohio", "city":"Piketon", "timezone":"America/New_York"},
|
||||
"modelMappings": {
|
||||
@@ -50,6 +51,7 @@ fn update_body() -> Value {
|
||||
"maxWaitingPerKey": 0,
|
||||
"maxWaitingPerAccount": 0,
|
||||
"concurrencyWaitTimeoutSeconds": 30,
|
||||
"responsesMaxDecompressedBodyBytes": 67108864,
|
||||
"rotationStrategy": "round_robin",
|
||||
"minCodexDesktopVersion": "26.825.6671",
|
||||
"minCodexCliVersion": "0.40.0",
|
||||
@@ -101,6 +103,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
use gateway_core::routing::{PublicModelId, UpstreamModelId};
|
||||
|
||||
let settings = RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: Revision::new(7).expect("revision"),
|
||||
@@ -121,6 +124,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::RoundRobin,
|
||||
min_codex_desktop_version: Some("26.825.6671".to_owned()),
|
||||
min_codex_cli_version: Some("0.40.0".to_owned()),
|
||||
@@ -144,7 +148,8 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
assert_eq!(
|
||||
value,
|
||||
json!({
|
||||
"requestLocationEnabled": false,
|
||||
"disableFast": false,
|
||||
"requestLocationEnabled": false,
|
||||
"requestLocation": {"country":"US", "region":"Ohio", "city":"Piketon", "timezone":"America/New_York"},
|
||||
"modelMappings": {
|
||||
"gpt-5.4": "gpt-5.5",
|
||||
@@ -157,6 +162,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
"maxWaitingPerKey": 0,
|
||||
"maxWaitingPerAccount": 0,
|
||||
"concurrencyWaitTimeoutSeconds": 30,
|
||||
"responsesMaxDecompressedBodyBytes": 67108864,
|
||||
"rotationStrategy": "round_robin",
|
||||
"minCodexDesktopVersion": "26.825.6671",
|
||||
"minCodexCliVersion": "0.40.0",
|
||||
@@ -196,6 +202,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
|
||||
.cloned()
|
||||
.collect();
|
||||
let settings = RuntimeSettings {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
config_revision: Revision::new(7).expect("revision"),
|
||||
@@ -217,6 +224,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: RotationStrategy::parse(&request.rotation_strategy)
|
||||
.expect("fixture rotation strategy"),
|
||||
min_codex_desktop_version: request.min_codex_desktop_version,
|
||||
@@ -575,3 +583,58 @@ async fn request_location_should_reject_invalid_or_missing_fields_without_replac
|
||||
original
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decompression_setting_should_reject_invalid_values() {
|
||||
let field = "responsesMaxDecompressedBodyBytes";
|
||||
for invalid in [
|
||||
json!(0),
|
||||
json!(-1),
|
||||
json!(1.5),
|
||||
json!(u64::MAX),
|
||||
Value::Null,
|
||||
] {
|
||||
let mut body = update_body();
|
||||
body[field] = invalid;
|
||||
if let Ok(request) = serde_json::from_value::<UpdateRuntimeSettingsRequest>(body) {
|
||||
assert_eq!(request.validate().unwrap_err().field(), field);
|
||||
}
|
||||
}
|
||||
let mut body = update_body();
|
||||
body.as_object_mut().unwrap().remove(field);
|
||||
assert!(serde_json::from_value::<UpdateRuntimeSettingsRequest>(body).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disable_fast_settings_updates_preserve_omitted_values() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let app = app(fixture.state());
|
||||
for (value, expected) in [(Some(true), true), (None, true), (Some(false), false)] {
|
||||
let mut body = update_body();
|
||||
if let Some(value) = value {
|
||||
body["disableFast"] = json!(value);
|
||||
} else {
|
||||
body.as_object_mut().unwrap().remove("disableFast");
|
||||
}
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(request(
|
||||
Method::POST,
|
||||
"/api/admin/settings/update",
|
||||
Some(body),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(request(Method::GET, "/api/admin/settings", None))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
response_json(response).await["data"]["disableFast"],
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,6 +100,7 @@ fn usage_record() -> UsageListRecord {
|
||||
provider_account_authentication_kind: Some("private-sentinel".to_owned()),
|
||||
upstream_model_id: Some("private-sentinel".to_owned()),
|
||||
upstream_transport: Some("websocket".to_owned()),
|
||||
upstream_response_model: Some("private-sentinel".to_owned()),
|
||||
service_tier: Some("private-sentinel".to_owned()),
|
||||
input_tokens: Some(1000),
|
||||
output_tokens: Some(100),
|
||||
|
||||
@@ -180,3 +180,121 @@ fn cli_user_agent_should_expose_semver_or_recognized_missing_version() {
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_remote_desktop_without_app_version_should_not_use_a_version_gate() {
|
||||
for name in [
|
||||
"codex_chatgpt_android_remote",
|
||||
"codex_chatgpt_ios_remote",
|
||||
"codex_chatgpt_future_os_remote",
|
||||
] {
|
||||
for remote_version in ["dev", "1.2.3"] {
|
||||
for (product, originator) in [
|
||||
("Codex Desktop", None),
|
||||
("codex_cli_rs", Some("Codex Desktop")),
|
||||
] {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"user-agent",
|
||||
HeaderValue::from_str(&format!(
|
||||
"{product}/0.154.0-alpha.6.2 (Windows 10.0.26200; x86_64) unknown ({name}; {remote_version})"
|
||||
))
|
||||
.unwrap(),
|
||||
);
|
||||
if let Some(originator) = originator {
|
||||
headers.insert("originator", HeaderValue::from_static(originator));
|
||||
}
|
||||
assert_eq!(identify_codex_client(&headers), None, "{headers:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_remote_marker_should_require_a_complete_exact_client_suffix() {
|
||||
for suffix in [
|
||||
"(codex_chatgpt_android_remote_extra; dev)",
|
||||
"(unofficial_codex_chatgpt_ios_remote; dev)",
|
||||
"(codex_chatgpt__remote; dev)",
|
||||
"(codex_chatgpt_remote; dev)",
|
||||
"(codex_chatgpt_future os_remote; dev)",
|
||||
"(codex_chatgpt_future;os_remote; dev)",
|
||||
"codex_chatgpt_android_remote; dev",
|
||||
"(codex_chatgpt_android_remote; dev",
|
||||
"(codex_chatgpt_android_remote; dev) trailing",
|
||||
"(codex_chatgpt_android_remote; dev) (another_client; dev)",
|
||||
"(codex_chatgpt_android_remote; )",
|
||||
"(another_client; codex_chatgpt_android_remote)",
|
||||
"(Codex Desktop; dev)",
|
||||
"(Codex Desktop; ) (codex_chatgpt_android_remote; dev)",
|
||||
"(Codex Desktop; invalid) (codex_chatgpt_ios_remote; dev)",
|
||||
"(Codex Desktop)",
|
||||
] {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"user-agent",
|
||||
HeaderValue::from_str(&format!("Codex Desktop/0.154.0-alpha.6.2 {suffix}")).unwrap(),
|
||||
);
|
||||
let client = identify_codex_client(&headers).expect("Desktop still requires a version");
|
||||
assert_eq!(client.kind(), CodexClientKind::Desktop, "{suffix}");
|
||||
assert!(client.version().is_none(), "{suffix}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_remote_desktop_should_preserve_explicit_version_validation() {
|
||||
for (value, expected) in [
|
||||
(
|
||||
HeaderValue::from_static("26.908.70816"),
|
||||
Some("26.908.70816"),
|
||||
),
|
||||
(HeaderValue::from_static("26.1.0"), Some("26.1.0")),
|
||||
(HeaderValue::from_static("dev"), None),
|
||||
(HeaderValue::from_static(""), None),
|
||||
(HeaderValue::from_bytes(&[0xff]).unwrap(), None),
|
||||
] {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("version", value);
|
||||
headers.insert(
|
||||
"user-agent",
|
||||
HeaderValue::from_static(
|
||||
"Codex Desktop/0.154.0-alpha.6.2 (codex_chatgpt_android_remote; dev)",
|
||||
),
|
||||
);
|
||||
let client = identify_codex_client(&headers).expect("explicit Desktop version");
|
||||
assert_eq!(client.kind(), CodexClientKind::Desktop);
|
||||
assert_eq!(
|
||||
client.version().map(ToString::to_string).as_deref(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_remote_marker_should_not_override_an_existing_desktop_version_suffix() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"user-agent",
|
||||
HeaderValue::from_static(
|
||||
"Codex Desktop/0.154.0-alpha.6.2 (Codex Desktop; 26.1.0) (codex_chatgpt_ios_remote; dev)",
|
||||
),
|
||||
);
|
||||
let client = identify_codex_client(&headers).expect("Desktop version remains available");
|
||||
assert_eq!(client.version().unwrap().to_string(), "26.1.0");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatgpt_remote_suffix_created_by_header_truncation_should_not_skip_the_gate() {
|
||||
let suffix = " (codex_chatgpt_android_remote; dev)";
|
||||
let mut user_agent = "Codex Desktop/0.154.0-alpha.6.2 ".to_owned();
|
||||
user_agent.push_str(&"x".repeat(4096 - user_agent.len() - suffix.len()));
|
||||
user_agent.push_str(suffix);
|
||||
user_agent.push_str(" (another_client; dev)");
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("user-agent", HeaderValue::from_str(&user_agent).unwrap());
|
||||
|
||||
let client =
|
||||
identify_codex_client(&headers).expect("incomplete headers still require a version");
|
||||
assert_eq!(client.kind(), CodexClientKind::Desktop);
|
||||
assert!(client.version().is_none());
|
||||
}
|
||||
|
||||
@@ -615,3 +615,33 @@ mod model_routing {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_recovery_http_response_uses_projection_without_replacing_upstream_facts() {
|
||||
use bytes::Bytes;
|
||||
use gateway_core::error::{ClientVisibleUpstreamResponse, OpaqueUpstreamValue};
|
||||
let body = Bytes::from_static(br#"{"error":{"code":"previous_response_not_found","type":"invalid_request_error","message":"Previous response was not found. Retrying the full request."}}"#);
|
||||
let provider = ProviderError::new(ProviderErrorKind::QuotaExhausted, UpstreamSendState::Sent)
|
||||
.with_status(429)
|
||||
.with_upstream_code(OpaqueUpstreamValue::new("usage_limit_reached"))
|
||||
.with_retry_after(std::time::Duration::from_secs(129_600))
|
||||
.with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
|
||||
400,
|
||||
Some(b"application/json".to_vec()),
|
||||
body.clone(),
|
||||
));
|
||||
let error = EngineError::Provider(provider);
|
||||
let response = gateway_api::openai::error::engine_error_response(&error);
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(response.headers()["content-type"], "application/json");
|
||||
assert!(!response.headers().contains_key("retry-after"));
|
||||
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), body);
|
||||
let EngineError::Provider(provider) = error else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(provider.upstream_status(), Some(429));
|
||||
assert_eq!(
|
||||
provider.upstream_code().unwrap().as_str(),
|
||||
"usage_limit_reached"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -135,8 +135,18 @@ pub(super) fn authenticated_client_for_provider(
|
||||
plaintext: &str,
|
||||
provider_name: &str,
|
||||
) -> AuthenticatedClient {
|
||||
authenticated_client_for_provider_with_limit(plaintext, provider_name, 64 * 1024 * 1024)
|
||||
}
|
||||
|
||||
pub(super) fn authenticated_client_for_provider_with_limit(
|
||||
plaintext: &str,
|
||||
provider_name: &str,
|
||||
bytes: usize,
|
||||
) -> AuthenticatedClient {
|
||||
let snapshot = snapshot(plaintext, provider_name)
|
||||
.with_responses_max_decompressed_body_bytes(std::num::NonZeroUsize::new(bytes).unwrap());
|
||||
let source = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot(plaintext, provider_name)),
|
||||
RuntimeSnapshotHandle::new(snapshot),
|
||||
Arc::new(UnusedExecutionStore),
|
||||
ProviderRegistry::default(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
|
||||
@@ -41,7 +41,9 @@ use serde_json::{Value, json};
|
||||
use gateway_api::openai::responses::{collect_execution_response, stream_execution_response};
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::openai::{api_router, authenticated_client_for_provider};
|
||||
use crate::openai::{
|
||||
api_router, authenticated_client_for_provider, authenticated_client_with_min_versions,
|
||||
};
|
||||
|
||||
#[derive(Default)]
|
||||
struct Trace {
|
||||
@@ -546,7 +548,7 @@ async fn invalid_compressed_http_requests_should_fail_before_execution() {
|
||||
super::request::encode_body("zstd", std::io::repeat(0).take(64 * 1024 * 1024 + 1))
|
||||
.into(),
|
||||
"request_too_large",
|
||||
"Decompressed request body exceeds the allowed size.",
|
||||
"Decompressed request body exceeds the allowed size (67108864 bytes).",
|
||||
),
|
||||
] {
|
||||
let headers = HeaderMap::from_iter([
|
||||
@@ -1968,6 +1970,112 @@ struct SessionExecution {
|
||||
session: Mutex<Option<Box<dyn ExecutionSession>>>,
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chatgpt_remote_responses_should_only_skip_missing_desktop_versions() {
|
||||
for name in [
|
||||
"codex_chatgpt_android_remote",
|
||||
"codex_chatgpt_ios_remote",
|
||||
"codex_chatgpt_future_os_remote",
|
||||
] {
|
||||
for (version, expected_code) in [
|
||||
(None, None),
|
||||
(Some("26.908.70816"), None),
|
||||
(Some("26.1.0"), Some("client_version_too_old")),
|
||||
(Some("dev"), Some("client_version_unavailable")),
|
||||
] {
|
||||
let trace = Arc::new(Trace::default());
|
||||
let execution = Arc::new(SessionExecution {
|
||||
client: authenticated_client_with_min_versions(
|
||||
"sk_remote_test",
|
||||
Some("26.908.70816"),
|
||||
Some("0.154.0"),
|
||||
),
|
||||
session: Mutex::new(Some(Box::new(FakeSession::streaming(
|
||||
Arc::clone(&trace),
|
||||
vec![
|
||||
NextStep::Event(delivery(
|
||||
started(),
|
||||
CommitRequirement::CommitBeforeDelivery,
|
||||
)),
|
||||
NextStep::Event(delivery(completed(), CommitRequirement::AlreadyCommitted)),
|
||||
NextStep::FinalizeSuccess,
|
||||
],
|
||||
)))),
|
||||
});
|
||||
let mut request = Request::post("/v1/responses")
|
||||
.header(AUTHORIZATION, "Bearer sk_remote_test")
|
||||
.header("user-agent", format!(
|
||||
"Codex Desktop/0.154.0-alpha.6.2 (Windows 10.0.26200; x86_64) unknown ({name}; dev)"
|
||||
));
|
||||
if let Some(version) = version {
|
||||
request = request.header("version", version);
|
||||
}
|
||||
let response = api_router(execution.clone())
|
||||
.await
|
||||
.oneshot(
|
||||
request
|
||||
.body(Body::from(
|
||||
json!({"model": "model-a", "input": "synthetic", "stream": true})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
if let Some(code) = expected_code {
|
||||
let value: Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(status, StatusCode::UPGRADE_REQUIRED, "{name} {version:?}");
|
||||
assert_eq!(value["error"]["code"], code);
|
||||
assert!(execution.session.lock().unwrap().is_some());
|
||||
} else {
|
||||
assert_eq!(status, StatusCode::OK, "{name} {version:?}");
|
||||
let body = String::from_utf8(body.to_vec()).unwrap();
|
||||
assert!(body.contains("event: response.completed\n"));
|
||||
assert!(body.ends_with("data: [DONE]\n\n"));
|
||||
assert!(trace.snapshot().contains(&"commit"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chatgpt_remote_responses_should_still_require_client_key_authentication() {
|
||||
let observed = Arc::new(Mutex::new(None));
|
||||
let execution = Arc::new(ContextCaptureExecution {
|
||||
observed: Arc::clone(&observed),
|
||||
client: authenticated_client_with_min_versions(
|
||||
"sk_context_test",
|
||||
Some("26.908.70816"),
|
||||
None,
|
||||
),
|
||||
});
|
||||
let router = api_router(execution).await;
|
||||
for authorization in [None, Some("Bearer sk_invalid")] {
|
||||
let mut request = Request::post("/v1/responses").header(
|
||||
"user-agent",
|
||||
"Codex Desktop/0.154.0-alpha.6.2 (codex_chatgpt_android_remote; dev)",
|
||||
);
|
||||
if let Some(authorization) = authorization {
|
||||
request = request.header(AUTHORIZATION, authorization);
|
||||
}
|
||||
let response = router
|
||||
.clone()
|
||||
.oneshot(
|
||||
request
|
||||
.body(Body::from(
|
||||
json!({"model": "model-a", "input": "synthetic"}).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert!(observed.lock().unwrap().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionService for SessionExecution {
|
||||
fn authenticate(&self, _: &str) -> Result<AuthenticatedClient, ClientAuthenticationError> {
|
||||
Ok(self.client.clone())
|
||||
@@ -2180,3 +2288,50 @@ async fn local_execution_failure_uses_the_existing_model_id() {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compressed_http_request_above_default_limit_should_reach_execution_after_increase() {
|
||||
let body = json!({"model": "model-a", "input": "x".repeat(64 * 1024 * 1024)}).to_string();
|
||||
let compressed = super::request::encode_body("zstd", body.as_bytes());
|
||||
for limit in [64 * 1024 * 1024, 128 * 1024 * 1024] {
|
||||
let observed = Arc::new(Mutex::new(None));
|
||||
let execution = Arc::new(ContextCaptureExecution {
|
||||
observed: Arc::clone(&observed),
|
||||
client: crate::openai::authenticated_client_for_provider_with_limit(
|
||||
"sk_context_test",
|
||||
"openai",
|
||||
limit,
|
||||
),
|
||||
});
|
||||
let response = api_router(execution)
|
||||
.await
|
||||
.oneshot(
|
||||
Request::post("/v1/responses")
|
||||
.header(AUTHORIZATION, "Bearer sk_context_test")
|
||||
.header(CONTENT_ENCODING, "zstd")
|
||||
.body(Body::from(compressed.clone()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
if limit == 64 * 1024 * 1024 {
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert!(observed.lock().unwrap().is_none());
|
||||
let error: Value =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), usize::MAX).await.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(error["error"]["code"], "request_too_large");
|
||||
} else {
|
||||
// 捕获执行服务刻意返回 500;断言完整输入到达执行层,而不是只验证路由状态码。
|
||||
let captured = observed
|
||||
.lock()
|
||||
.unwrap()
|
||||
.take()
|
||||
.expect("request reached execution");
|
||||
assert_eq!(
|
||||
captured.input.unwrap().as_str().unwrap().len(),
|
||||
64 * 1024 * 1024
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,7 +19,7 @@ use gateway_protocol::openai::sse::parse_sse_events;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
fn decode_request(body: &[u8]) -> Result<DecodedResponsesRequest, RequestDecodeError> {
|
||||
decode_request_with_headers(body, &HeaderMap::new())
|
||||
decode_request_with_headers(body, &HeaderMap::new(), 64 * 1024 * 1024)
|
||||
}
|
||||
|
||||
fn decode_response_create(
|
||||
@@ -205,8 +205,9 @@ fn decoder_should_preserve_connection_metadata_outside_the_openai_wire_body() {
|
||||
}
|
||||
});
|
||||
|
||||
let decoded = decode_request_with_headers(body.to_string().as_bytes(), &headers)
|
||||
.expect("request should decode");
|
||||
let decoded =
|
||||
decode_request_with_headers(body.to_string().as_bytes(), &headers, 64 * 1024 * 1024)
|
||||
.expect("request should decode");
|
||||
|
||||
assert_eq!(Value::Object(openai_wire_body(&decoded).clone()), body);
|
||||
assert_eq!(
|
||||
@@ -317,9 +318,12 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
HeaderValue::from_static("client-is-update"),
|
||||
);
|
||||
|
||||
let decoded =
|
||||
decode_request_with_headers(br#"{"model":"smart-code","input":"hello"}"#, &headers)
|
||||
.expect("opaque request headers should not affect body decoding");
|
||||
let decoded = decode_request_with_headers(
|
||||
br#"{"model":"smart-code","input":"hello"}"#,
|
||||
&headers,
|
||||
64 * 1024 * 1024,
|
||||
)
|
||||
.expect("opaque request headers should not affect body decoding");
|
||||
let entries = openai_protocol_context(&decoded)
|
||||
.get("opaque_request_headers")
|
||||
.and_then(Value::as_array)
|
||||
@@ -420,9 +424,12 @@ fn decoder_should_strip_http_transport_but_leave_source_headers_for_provider() {
|
||||
headers.insert("content-encoding", HeaderValue::from_static("identity"));
|
||||
headers.insert("x-openai-future-mode", HeaderValue::from_static("keep"));
|
||||
|
||||
let decoded =
|
||||
decode_request_with_headers(br#"{"model":"smart-code","input":"hello"}"#, &headers)
|
||||
.expect("decode request behind a reverse proxy");
|
||||
let decoded = decode_request_with_headers(
|
||||
br#"{"model":"smart-code","input":"hello"}"#,
|
||||
&headers,
|
||||
64 * 1024 * 1024,
|
||||
)
|
||||
.expect("decode request behind a reverse proxy");
|
||||
|
||||
assert_eq!(
|
||||
openai_protocol_context(&decoded).get("opaque_request_headers"),
|
||||
@@ -482,7 +489,8 @@ fn downstream_client_headers_should_remain_opaque_without_losing_session_semanti
|
||||
frame["type"] = json!("response.create");
|
||||
let opening = OpenAiRequestHeaders::from_headers(&headers);
|
||||
for decoded in [
|
||||
decode_request_with_headers(body.to_string().as_bytes(), &headers).unwrap(),
|
||||
decode_request_with_headers(body.to_string().as_bytes(), &headers, 64 * 1024 * 1024)
|
||||
.unwrap(),
|
||||
decode_response_create_with_context(&frame.to_string(), &opening).unwrap(),
|
||||
decode_response_create_with_context(&frame.to_string(), &opening).unwrap(),
|
||||
] {
|
||||
@@ -565,7 +573,7 @@ fn decoder_should_preserve_unknown_arbitrary_precision_numbers() {
|
||||
let number = "12345678901234567890123456789012345678901234567890";
|
||||
let body = format!(r#"{{"model":"smart-code","input":"hello","future_number":{number}}}"#);
|
||||
|
||||
let decoded = decode_request_with_headers(body.as_bytes(), &HeaderMap::new())
|
||||
let decoded = decode_request_with_headers(body.as_bytes(), &HeaderMap::new(), 64 * 1024 * 1024)
|
||||
.expect("opaque numeric field should decode");
|
||||
let encoded = serde_json::to_string(openai_wire_body(&decoded)).expect("encode request body");
|
||||
|
||||
|
||||
@@ -41,7 +41,9 @@ fn http_decode_should_preserve_compressed_request_fields() {
|
||||
let expected = serde_json::from_slice::<serde_json::Value>(REQUEST).expect("fixture JSON");
|
||||
for encoding in ENCODINGS {
|
||||
let compressed = encode_body(encoding, REQUEST);
|
||||
let decoded = decode_request_with_headers(&compressed, &headers(encoding)).expect(encoding);
|
||||
let decoded =
|
||||
decode_request_with_headers(&compressed, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect(encoding);
|
||||
assert_eq!(
|
||||
openai_wire_body(&decoded),
|
||||
expected.as_object().expect("object")
|
||||
@@ -52,10 +54,10 @@ fn http_decode_should_preserve_compressed_request_fields() {
|
||||
#[test]
|
||||
fn http_decode_should_accept_plain_identity_and_case_insensitive_encodings() {
|
||||
for request_headers in [HeaderMap::new(), headers(""), headers(" identity ")] {
|
||||
assert!(decode_request_with_headers(REQUEST, &request_headers).is_ok());
|
||||
assert!(decode_request_with_headers(REQUEST, &request_headers, 64 * 1024 * 1024).is_ok());
|
||||
}
|
||||
let compressed = encode_body("gzip", REQUEST);
|
||||
assert!(decode_request_with_headers(&compressed, &headers(" GZip ")).is_ok());
|
||||
assert!(decode_request_with_headers(&compressed, &headers(" GZip "), 64 * 1024 * 1024).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -63,7 +65,7 @@ fn http_decode_should_reject_unknown_and_stacked_encodings() {
|
||||
let mut repeated = headers("identity");
|
||||
repeated.append(CONTENT_ENCODING, HeaderValue::from_static("gzip"));
|
||||
for request_headers in [headers("br"), headers("gzip, zstd"), repeated] {
|
||||
let error = decode_request_with_headers(REQUEST, &request_headers)
|
||||
let error = decode_request_with_headers(REQUEST, &request_headers, 64 * 1024 * 1024)
|
||||
.expect_err("unsupported encoding");
|
||||
assert!(matches!(
|
||||
error,
|
||||
@@ -82,7 +84,8 @@ fn http_decode_should_reject_corrupted_and_truncated_compressed_bodies() {
|
||||
let mut truncated = encode_body(encoding, REQUEST);
|
||||
truncated.truncate(truncated.len() - 4);
|
||||
for body in [b"not-a-compressed-stream".as_slice(), truncated.as_slice()] {
|
||||
let error = decode_request_with_headers(body, &headers(encoding)).expect_err(encoding);
|
||||
let error = decode_request_with_headers(body, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect_err(encoding);
|
||||
assert_eq!(
|
||||
error.protocol_body().into_value(),
|
||||
json!({"error": {
|
||||
@@ -99,7 +102,9 @@ fn http_decode_should_read_all_gzip_members_and_zstd_frames() {
|
||||
for encoding in ["gzip", "zstd"] {
|
||||
let mut compressed = encode_body(encoding, &REQUEST[..12]);
|
||||
compressed.extend(encode_body(encoding, &REQUEST[12..]));
|
||||
let decoded = decode_request_with_headers(&compressed, &headers(encoding)).expect(encoding);
|
||||
let decoded =
|
||||
decode_request_with_headers(&compressed, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect(encoding);
|
||||
assert_eq!(openai_wire_body(&decoded)["input"], "hello");
|
||||
}
|
||||
}
|
||||
@@ -110,7 +115,8 @@ fn http_decode_should_not_ignore_invalid_trailing_members() {
|
||||
let mut compressed = encode_body(encoding, REQUEST);
|
||||
compressed.extend(encode_body(encoding, b"invalid trailing JSON".as_slice()));
|
||||
assert_eq!(
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding),
|
||||
decode_request_with_headers(&compressed, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect_err(encoding),
|
||||
RequestDecodeError::MalformedJson
|
||||
);
|
||||
}
|
||||
@@ -129,7 +135,8 @@ fn http_decode_should_reject_zstd_frames_with_an_excessive_window() {
|
||||
REQUEST
|
||||
);
|
||||
assert_eq!(
|
||||
decode_request_with_headers(&compressed, &headers("zstd")).expect_err("oversized window"),
|
||||
decode_request_with_headers(&compressed, &headers("zstd"), 128 * 1024 * 1024)
|
||||
.expect_err("oversized window"),
|
||||
RequestDecodeError::MalformedJson
|
||||
);
|
||||
}
|
||||
@@ -138,7 +145,8 @@ fn http_decode_should_reject_zstd_frames_with_an_excessive_window() {
|
||||
fn http_decode_should_accept_the_exact_decompressed_limit() {
|
||||
let padding = repeat(b' ').take(DECOMPRESSED_LIMIT - REQUEST.len() as u64);
|
||||
let compressed = encode_body("zstd", REQUEST.chain(padding));
|
||||
let decoded = decode_request_with_headers(&compressed, &headers("zstd")).expect("exact limit");
|
||||
let decoded = decode_request_with_headers(&compressed, &headers("zstd"), 64 * 1024 * 1024)
|
||||
.expect("exact limit");
|
||||
assert_eq!(decoded.metadata().requested_model(), "gpt-test");
|
||||
}
|
||||
|
||||
@@ -146,11 +154,13 @@ fn http_decode_should_accept_the_exact_decompressed_limit() {
|
||||
fn http_decode_should_reject_bodies_above_the_decompressed_limit() {
|
||||
for encoding in ENCODINGS {
|
||||
let compressed = encode_body(encoding, repeat(0).take(DECOMPRESSED_LIMIT + 1));
|
||||
let error =
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding);
|
||||
let error = decode_request_with_headers(&compressed, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect_err(encoding);
|
||||
assert_eq!(
|
||||
error,
|
||||
RequestDecodeError::DecompressedBodyTooLarge,
|
||||
RequestDecodeError::DecompressedBodyTooLarge {
|
||||
limit_bytes: 64 * 1024 * 1024
|
||||
},
|
||||
"{encoding}"
|
||||
);
|
||||
}
|
||||
@@ -163,12 +173,54 @@ fn http_decode_should_stop_at_limit_before_reading_invalid_trailing_frames() {
|
||||
// 这个断言验证越界时已经停止读取,而不只是最终返回了某个超限错误。
|
||||
let mut compressed = encode_body(encoding, repeat(0).take(DECOMPRESSED_LIMIT + 1));
|
||||
compressed.extend_from_slice(b"invalid trailing frame");
|
||||
let error =
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding);
|
||||
let error = decode_request_with_headers(&compressed, &headers(encoding), 64 * 1024 * 1024)
|
||||
.expect_err(encoding);
|
||||
assert_eq!(
|
||||
error,
|
||||
RequestDecodeError::DecompressedBodyTooLarge,
|
||||
RequestDecodeError::DecompressedBodyTooLarge {
|
||||
limit_bytes: 64 * 1024 * 1024
|
||||
},
|
||||
"{encoding}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_apply_configured_limits_to_every_encoding_and_frame() {
|
||||
for encoding in ENCODINGS {
|
||||
let compressed = encode_body(encoding, REQUEST);
|
||||
assert!(
|
||||
decode_request_with_headers(&compressed, &headers(encoding), REQUEST.len()).is_ok()
|
||||
);
|
||||
let error = decode_request_with_headers(&compressed, &headers(encoding), REQUEST.len() - 1)
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error,
|
||||
RequestDecodeError::DecompressedBodyTooLarge {
|
||||
limit_bytes: REQUEST.len() - 1
|
||||
}
|
||||
);
|
||||
assert!(
|
||||
error
|
||||
.protocol_body()
|
||||
.error
|
||||
.message
|
||||
.contains(&(REQUEST.len() - 1).to_string())
|
||||
);
|
||||
}
|
||||
for encoding in ["gzip", "zstd"] {
|
||||
let split = REQUEST.len() / 2;
|
||||
let mut compressed = encode_body(encoding, &REQUEST[..split]);
|
||||
compressed.extend(encode_body(encoding, &REQUEST[split..]));
|
||||
assert!(
|
||||
decode_request_with_headers(&compressed, &headers(encoding), REQUEST.len()).is_ok()
|
||||
);
|
||||
assert!(matches!(
|
||||
decode_request_with_headers(&compressed, &headers(encoding), REQUEST.len() - 1),
|
||||
Err(RequestDecodeError::DecompressedBodyTooLarge { .. })
|
||||
));
|
||||
}
|
||||
for request_headers in [HeaderMap::new(), headers("identity")] {
|
||||
assert!(decode_request_with_headers(REQUEST, &request_headers, 1).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -447,3 +447,29 @@ async fn deliverable_business_failures_are_not_rewritten_or_followed_by_a_second
|
||||
socket.close(None).await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_quota_recovery_delivers_client_projection_instead_of_upstream_status() {
|
||||
let detail = json!({"message":"Previous response was not found. Retrying the full request.","code":"previous_response_not_found","type":"invalid_request_error"});
|
||||
let provider = ProviderError::new(ProviderErrorKind::QuotaExhausted, UpstreamSendState::Sent)
|
||||
.with_status(429)
|
||||
.with_retry_after(Duration::from_secs(129_600))
|
||||
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
|
||||
detail["message"].as_str().unwrap(),
|
||||
Some("previous_response_not_found".to_owned()),
|
||||
Some("invalid_request_error".to_owned()),
|
||||
))
|
||||
.with_client_visible_upstream_response(
|
||||
ClientVisibleUpstreamResponse::new(
|
||||
400,
|
||||
Some(b"application/json".to_vec()),
|
||||
Bytes::from(json!({"error":detail}).to_string()),
|
||||
)
|
||||
.with_headers(vec![header("x-request-id", b"req-quota-rejected")]),
|
||||
);
|
||||
let error = initial_error(EngineError::Provider(provider), Vec::new()).await;
|
||||
assert_eq!(error["status"], 400);
|
||||
assert_eq!(error["error"], detail);
|
||||
assert_eq!(error["headers"]["x-request-id"], "req-quota-rejected");
|
||||
assert!(error["headers"].get("retry-after").is_none());
|
||||
}
|
||||
|
||||
@@ -247,17 +247,31 @@ impl ClientRoutingScope {
|
||||
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct FrozenAccountScope {
|
||||
disable_fast: bool,
|
||||
directory: Arc<RuntimeAccountDirectory>,
|
||||
client_scope: ClientRoutingScope,
|
||||
}
|
||||
|
||||
impl FrozenAccountScope {
|
||||
/// Key 绑定分组的冻结 Fast 限制,与账号成员资格无关。
|
||||
#[must_use]
|
||||
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
|
||||
self.disable_fast = disable_fast;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn disable_fast(&self) -> bool {
|
||||
self.disable_fast
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
directory: Arc<RuntimeAccountDirectory>,
|
||||
client_scope: ClientRoutingScope,
|
||||
) -> Self {
|
||||
Self {
|
||||
disable_fast: false,
|
||||
directory,
|
||||
client_scope,
|
||||
}
|
||||
|
||||
@@ -788,6 +788,7 @@ where
|
||||
}));
|
||||
let context = AttemptContext::new(
|
||||
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
|
||||
.with_disable_fast(self.plan.disable_fast())
|
||||
.with_request_location(self.plan.request_location().cloned())
|
||||
.with_concurrency_wait_budget(self.concurrency_wait_budget.clone())
|
||||
.with_timing_started_at(self.observation.timing_started_at)
|
||||
@@ -1498,6 +1499,12 @@ where
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
upstream_response_model: self
|
||||
.current
|
||||
.as_ref()
|
||||
.and_then(|current| current.response_observation.as_ref())
|
||||
.and_then(ProviderResponseObservation::upstream_response_model)
|
||||
.map(str::to_owned),
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: None,
|
||||
@@ -1659,6 +1666,12 @@ where
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
upstream_response_model: self
|
||||
.current
|
||||
.as_ref()
|
||||
.and_then(|current| current.response_observation.as_ref())
|
||||
.and_then(ProviderResponseObservation::upstream_response_model)
|
||||
.map(str::to_owned),
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: Some(finalization.error),
|
||||
|
||||
@@ -314,6 +314,7 @@ impl ContinuationAttempt {
|
||||
/// Provider 每次执行可见的 request-local context。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RequestAttemptContext {
|
||||
disable_fast: bool,
|
||||
request_location: Option<crate::account::RequestLocation>,
|
||||
request_id: ModelRequestId,
|
||||
client_api_key_ref: ClientApiKeyId,
|
||||
@@ -323,6 +324,12 @@ pub struct RequestAttemptContext {
|
||||
}
|
||||
|
||||
impl RequestAttemptContext {
|
||||
#[must_use]
|
||||
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
|
||||
self.disable_fast = disable_fast;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_request_location(
|
||||
mut self,
|
||||
@@ -337,6 +344,7 @@ impl RequestAttemptContext {
|
||||
Self {
|
||||
request_id,
|
||||
client_api_key_ref,
|
||||
disable_fast: false,
|
||||
request_location: None,
|
||||
timing_started_at: Instant::now(),
|
||||
trace: crate::diagnostics::TraceContext::default(),
|
||||
@@ -399,6 +407,11 @@ pub struct AttemptContext {
|
||||
}
|
||||
|
||||
impl AttemptContext {
|
||||
#[must_use]
|
||||
pub const fn disable_fast(&self) -> bool {
|
||||
self.request.disable_fast
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn request_location(&self) -> Option<&crate::account::RequestLocation> {
|
||||
self.request.request_location.as_ref()
|
||||
@@ -659,6 +672,7 @@ pub struct ModelRequestFinalization {
|
||||
pub http_version: Option<String>,
|
||||
pub websocket_pool: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
/// Provider 已筛选的专有观测 JSON;Core 不解释字段。
|
||||
pub provider_metadata_json: Option<String>,
|
||||
/// 请求全程的有界诊断快照,跨 Provider 与重试保留。
|
||||
|
||||
@@ -338,7 +338,10 @@ impl ProviderConnectionObservation {
|
||||
}
|
||||
}
|
||||
|
||||
/// 只供当前客户端请求使用的原始上游 HTTP 失败响应。
|
||||
/// 只供当前客户端请求使用的 HTTP 失败响应。
|
||||
///
|
||||
/// 通常保留原始上游响应;Provider 也可为明确的恢复流程准备客户端投影。
|
||||
/// 投影不能覆盖 [`ProviderError`] 中用于诊断和记账的真实上游事实。
|
||||
///
|
||||
/// 该值不属于稳定诊断事实,不能进入日志或持久化。它刻意不实现 [`Clone`];
|
||||
/// [`ProviderError`] 的普通 clone 会丢弃它,只有最终失败的原对象才能把响应交给
|
||||
@@ -352,7 +355,7 @@ pub struct ClientVisibleUpstreamResponse {
|
||||
}
|
||||
|
||||
impl ClientVisibleUpstreamResponse {
|
||||
/// 保存 transport 实际收到的状态码、Content-Type 原值和正文。
|
||||
/// 保存客户端响应的状态码、Content-Type 和正文。
|
||||
#[must_use]
|
||||
pub fn new(status: u16, content_type: Option<Vec<u8>>, body: Bytes) -> Self {
|
||||
Self {
|
||||
|
||||
@@ -530,6 +530,7 @@ pub struct ProviderResponseObservation {
|
||||
status_code: Option<u16>,
|
||||
request_id: Option<OpaqueUpstreamValue>,
|
||||
service_tier: Option<String>,
|
||||
upstream_response_model: Option<String>,
|
||||
timings: ProviderResponseTimings,
|
||||
client_headers: Vec<ProviderResponseHeader>,
|
||||
provider_metadata: Option<ProviderResponseMetadata>,
|
||||
@@ -545,6 +546,7 @@ impl ProviderResponseObservation {
|
||||
status_code: None,
|
||||
request_id: None,
|
||||
service_tier: None,
|
||||
upstream_response_model: None,
|
||||
timings: ProviderResponseTimings::default(),
|
||||
client_headers: Vec::new(),
|
||||
provider_metadata: None,
|
||||
@@ -600,6 +602,21 @@ impl ProviderResponseObservation {
|
||||
self
|
||||
}
|
||||
|
||||
/// 记录 Provider 从上游响应明确取得的模型,不改写发送模型或计费合同。
|
||||
#[must_use]
|
||||
pub fn with_upstream_response_model_if_valid(mut self, model: &str) -> Self {
|
||||
let model = model.trim();
|
||||
if validate_text(model, 256, false, None).is_ok() {
|
||||
self.upstream_response_model = Some(model.to_owned());
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn upstream_response_model(&self) -> Option<&str> {
|
||||
self.upstream_response_model.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_timings(mut self, timings: ProviderResponseTimings) -> Self {
|
||||
self.timings = timings;
|
||||
|
||||
@@ -536,6 +536,7 @@ impl ProviderCandidate {
|
||||
/// 一次请求冻结的 Provider 尝试顺序。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RoutingPlan {
|
||||
disable_fast: bool,
|
||||
request_location: Option<crate::account::RequestLocation>,
|
||||
config_revision: ConfigRevision,
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
@@ -546,6 +547,11 @@ pub struct RoutingPlan {
|
||||
}
|
||||
|
||||
impl RoutingPlan {
|
||||
#[must_use]
|
||||
pub const fn disable_fast(&self) -> bool {
|
||||
self.disable_fast
|
||||
}
|
||||
|
||||
/// 本次请求冻结的全局位置,重试时沿用同一份配置。
|
||||
#[must_use]
|
||||
pub const fn request_location(&self) -> Option<&crate::account::RequestLocation> {
|
||||
|
||||
@@ -28,12 +28,14 @@ const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4;
|
||||
/// Store 在一个一致性读取中提供的调度设置事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotSettingsFacts {
|
||||
disable_fast: bool,
|
||||
request_location_enabled: bool,
|
||||
request_location: crate::account::RequestLocation,
|
||||
max_concurrent_per_account: u32,
|
||||
max_waiting_per_key: u32,
|
||||
max_waiting_per_account: u32,
|
||||
concurrency_wait_timeout_seconds: u32,
|
||||
responses_max_decompressed_body_bytes: u64,
|
||||
request_interval_ms: u64,
|
||||
rotation_strategy: String,
|
||||
model_mappings: BTreeMap<String, String>,
|
||||
@@ -42,6 +44,18 @@ pub struct SnapshotSettingsFacts {
|
||||
}
|
||||
|
||||
impl SnapshotSettingsFacts {
|
||||
#[must_use]
|
||||
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
|
||||
self.disable_fast = disable_fast;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_responses_max_decompressed_body_bytes(mut self, bytes: u64) -> Self {
|
||||
self.responses_max_decompressed_body_bytes = bytes;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_request_location(
|
||||
mut self,
|
||||
@@ -76,12 +90,14 @@ impl SnapshotSettingsFacts {
|
||||
min_codex_cli_version: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
disable_fast: false,
|
||||
request_location_enabled: false,
|
||||
request_location: crate::account::RequestLocation::default(),
|
||||
max_concurrent_per_account,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
request_interval_ms,
|
||||
rotation_strategy: rotation_strategy.into(),
|
||||
model_mappings,
|
||||
@@ -120,15 +136,27 @@ impl SnapshotClientPolicyFacts {
|
||||
/// Store 读取到的账号分组事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotAccountGroupFacts {
|
||||
disable_fast: bool,
|
||||
id: AccountGroupId,
|
||||
name: String,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl SnapshotAccountGroupFacts {
|
||||
#[must_use]
|
||||
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
|
||||
self.disable_fast = disable_fast;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn new(id: AccountGroupId, name: String, enabled: bool) -> Self {
|
||||
Self { id, name, enabled }
|
||||
Self {
|
||||
id,
|
||||
name,
|
||||
enabled,
|
||||
disable_fast: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -391,6 +419,12 @@ async fn compile_runtime_snapshot(
|
||||
{
|
||||
return Err(RuntimeSnapshotCompileError::InvalidData);
|
||||
}
|
||||
let decompressed_body_limit =
|
||||
isize::try_from(facts.settings.responses_max_decompressed_body_bytes)
|
||||
.ok()
|
||||
.and_then(|bytes| usize::try_from(bytes).ok())
|
||||
.and_then(std::num::NonZeroUsize::new)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
let queue_timeout =
|
||||
Duration::from_secs(u64::from(facts.settings.concurrency_wait_timeout_seconds));
|
||||
let client_queue_policy = ConcurrencyQueuePolicy {
|
||||
@@ -428,6 +462,7 @@ async fn compile_runtime_snapshot(
|
||||
});
|
||||
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
|
||||
for policy in facts.client_policies {
|
||||
let mut disable_fast = false;
|
||||
let account_scope = if policy.group_ids.is_empty() {
|
||||
FrozenAccountScope::new(
|
||||
Arc::clone(&account_directory),
|
||||
@@ -444,6 +479,8 @@ async fn compile_runtime_snapshot(
|
||||
let group = groups
|
||||
.get(&group_id)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
// 禁用分组仅影响选号;Key 仍绑定其 Fast 限制。
|
||||
disable_fast |= group.disable_fast;
|
||||
bound_groups.push(RoutingGroupSnapshot::new(
|
||||
group.id.clone(),
|
||||
group.name.clone(),
|
||||
@@ -463,7 +500,7 @@ async fn compile_runtime_snapshot(
|
||||
client_policies.push(ClientPolicy::new(
|
||||
policy.key_id,
|
||||
policy.plaintext_key,
|
||||
Arc::new(account_scope),
|
||||
Arc::new(account_scope.with_disable_fast(disable_fast)),
|
||||
true,
|
||||
policy.limits,
|
||||
));
|
||||
@@ -491,7 +528,9 @@ async fn compile_runtime_snapshot(
|
||||
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)
|
||||
.map(|snapshot| {
|
||||
snapshot
|
||||
.with_disable_fast(facts.settings.disable_fast)
|
||||
.with_request_location(request_location)
|
||||
.with_responses_max_decompressed_body_bytes(decompressed_body_limit)
|
||||
.with_client_queue_policy(client_queue_policy)
|
||||
.with_model_mappings(model_mappings)
|
||||
.with_account_directory(account_directory)
|
||||
@@ -503,6 +542,8 @@ async fn compile_runtime_snapshot(
|
||||
/// 数据面使用的不可变配置快照。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RuntimeSnapshot {
|
||||
disable_fast: bool,
|
||||
responses_max_decompressed_body_bytes: std::num::NonZeroUsize,
|
||||
request_location: Option<crate::account::RequestLocation>,
|
||||
revision: ConfigRevision,
|
||||
client_queue_policy: ConcurrencyQueuePolicy,
|
||||
@@ -520,6 +561,26 @@ pub struct RuntimeSnapshot {
|
||||
}
|
||||
|
||||
impl RuntimeSnapshot {
|
||||
#[must_use]
|
||||
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
|
||||
self.disable_fast = disable_fast;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn responses_max_decompressed_body_bytes(&self) -> usize {
|
||||
self.responses_max_decompressed_body_bytes.get()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_responses_max_decompressed_body_bytes(
|
||||
mut self,
|
||||
bytes: std::num::NonZeroUsize,
|
||||
) -> Self {
|
||||
self.responses_max_decompressed_body_bytes = bytes;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_request_location(
|
||||
mut self,
|
||||
@@ -608,6 +669,9 @@ impl RuntimeSnapshot {
|
||||
client_policy_map.retain(|_, policy| policy.enabled());
|
||||
|
||||
Ok(Self {
|
||||
responses_max_decompressed_body_bytes: std::num::NonZeroUsize::new(64 * 1024 * 1024)
|
||||
.expect("positive default limit"),
|
||||
disable_fast: false,
|
||||
request_location: None,
|
||||
revision,
|
||||
account_selection_policy,
|
||||
@@ -919,6 +983,7 @@ impl RuntimeSnapshot {
|
||||
|
||||
Ok(RoutingPlan {
|
||||
config_revision: self.revision,
|
||||
disable_fast: self.disable_fast || account_scope.disable_fast(),
|
||||
request_location: self.request_location.clone(),
|
||||
account_selection_policy: self.account_selection_policy,
|
||||
operation: operation.kind(),
|
||||
@@ -961,6 +1026,7 @@ impl RuntimeSnapshot {
|
||||
};
|
||||
Ok(RoutingPlan {
|
||||
config_revision: self.revision,
|
||||
disable_fast: self.disable_fast || account_scope.disable_fast(),
|
||||
request_location: self.request_location.clone(),
|
||||
account_selection_policy: self.account_selection_policy,
|
||||
operation: operation.kind(),
|
||||
|
||||
@@ -73,6 +73,7 @@ struct FinalState {
|
||||
http_version: Option<String>,
|
||||
websocket_pool: Option<String>,
|
||||
service_tier: Option<String>,
|
||||
upstream_response_model: Option<String>,
|
||||
upstream_request_id: Option<String>,
|
||||
upstream_status_code: Option<u16>,
|
||||
transport_decision_wait_ms: Option<u64>,
|
||||
@@ -239,6 +240,7 @@ impl ExecutionStore for FakeStore {
|
||||
http_version: finalization.http_version,
|
||||
websocket_pool: finalization.websocket_pool,
|
||||
service_tier: finalization.service_tier,
|
||||
upstream_response_model: finalization.upstream_response_model,
|
||||
upstream_request_id: finalization.upstream_request_id,
|
||||
upstream_status_code: finalization.upstream_status_code,
|
||||
transport_decision_wait_ms: finalization.timings.transport_decision_wait_ms,
|
||||
@@ -592,6 +594,15 @@ fn plan_with_location(
|
||||
operation: &Operation,
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
request_location: gateway_core::account::RequestLocation,
|
||||
) -> RoutingPlan {
|
||||
plan_with_location_and_fast_policy(operation, account_selection_policy, request_location, false)
|
||||
}
|
||||
|
||||
fn plan_with_location_and_fast_policy(
|
||||
operation: &Operation,
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
request_location: gateway_core::account::RequestLocation,
|
||||
disable_fast: bool,
|
||||
) -> RoutingPlan {
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let public_model = PublicModelId::new("gpt-5").expect("public model");
|
||||
@@ -634,6 +645,7 @@ fn plan_with_location(
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("snapshot")
|
||||
.with_disable_fast(disable_fast)
|
||||
.with_request_location(Some(request_location))
|
||||
.with_account_directory(Arc::clone(&directory));
|
||||
snapshot
|
||||
@@ -1093,6 +1105,7 @@ fn response_observation_is_persisted_but_never_delivered() {
|
||||
.with_websocket_pool(WebSocketPoolKind::New)
|
||||
.with_status_code(200)
|
||||
.with_request_id(OpaqueUpstreamValue::new("upstream-observed"))
|
||||
.with_upstream_response_model_if_valid("gpt-returned")
|
||||
.try_with_service_tier("priority")
|
||||
.expect("service tier")
|
||||
.with_timings(ProviderResponseTimings {
|
||||
@@ -1132,6 +1145,10 @@ fn response_observation_is_persisted_but_never_delivered() {
|
||||
assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2"));
|
||||
assert_eq!(finalization.websocket_pool.as_deref(), Some("new"));
|
||||
assert_eq!(finalization.service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(
|
||||
finalization.upstream_response_model.as_deref(),
|
||||
Some("gpt-returned")
|
||||
);
|
||||
assert_eq!(
|
||||
finalization.upstream_request_id.as_deref(),
|
||||
Some("upstream-observed")
|
||||
@@ -1350,6 +1367,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
|
||||
.with_http_version(UpstreamHttpVersion::Http11)
|
||||
.with_status_code(503)
|
||||
.with_request_id(OpaqueUpstreamValue::new("discarded-request"))
|
||||
.with_upstream_response_model_if_valid("discarded-model")
|
||||
.with_timings(ProviderResponseTimings {
|
||||
transport_decision_wait_ms: Some(7),
|
||||
connect_ms: Some(11),
|
||||
@@ -1409,6 +1427,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
|
||||
|
||||
let state = store.state.lock().expect("store lock");
|
||||
let finalization = &state.finalizations[0];
|
||||
assert_eq!(finalization.upstream_response_model, None);
|
||||
let trace: Value =
|
||||
serde_json::from_str(finalization.diagnostic_trace_json.as_deref().unwrap()).unwrap();
|
||||
let events = trace["events"].as_array().unwrap();
|
||||
@@ -1710,61 +1729,72 @@ fn native_continuation_uses_provider_defined_owner_then_cross_account_recovery()
|
||||
|
||||
#[test]
|
||||
fn native_continuation_client_replay_required_is_terminal() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
let operation = Operation::Generate(generate.with_provider_session_state(
|
||||
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
|
||||
));
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_one",
|
||||
items: vec![Err(ProviderError::new(
|
||||
ProviderErrorKind::ContinuationRecoveryRequired,
|
||||
UpstreamSendState::NotSent,
|
||||
)
|
||||
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
|
||||
.with_continuation_recovery_disposition(
|
||||
ContinuationRecoveryDisposition::ClientReplayRequired,
|
||||
))],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_one",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-secret-id"),
|
||||
PreviousResponseId::new("provider-native-id"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
for (kind, send_state) in [
|
||||
(
|
||||
ProviderErrorKind::ContinuationRecoveryRequired,
|
||||
UpstreamSendState::NotSent,
|
||||
),
|
||||
(
|
||||
ProviderErrorKind::QuotaExhausted,
|
||||
UpstreamSendState::NotSent,
|
||||
),
|
||||
(ProviderErrorKind::QuotaExhausted, UpstreamSendState::Sent),
|
||||
] {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
let operation = Operation::Generate(generate.with_provider_session_state(
|
||||
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
|
||||
));
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_one",
|
||||
items: vec![Err(ProviderError::new(kind, send_state)
|
||||
.with_replay_safe()
|
||||
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
|
||||
.with_continuation_recovery_disposition(
|
||||
ContinuationRecoveryDisposition::ClientReplayRequired,
|
||||
))],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_one",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-secret-id"),
|
||||
PreviousResponseId::new("provider-native-id"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
let error = block_on(session.collect_uncommitted())
|
||||
.expect_err("client replay requirement must stop proxy recovery");
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
let error = block_on(session.collect_uncommitted())
|
||||
.expect_err("client replay requirement must stop proxy recovery");
|
||||
|
||||
assert!(matches!(error, EngineError::Provider(_)));
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert_eq!(
|
||||
contexts[0].continuation_attempt(),
|
||||
ContinuationAttempt::Native
|
||||
);
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.intermediate_failures, 0);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
|
||||
assert!(matches!(error, EngineError::Provider(_)));
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert_eq!(
|
||||
contexts[0].continuation_attempt(),
|
||||
ContinuationAttempt::Native
|
||||
);
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.intermediate_failures, 0);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
|
||||
assert_eq!(state.finalizations[0].send_state, send_state);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -4508,16 +4538,17 @@ fn deadline_before_first_event_records_no_provider_circuit_failure() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn global_request_location_should_reach_every_account_retry() {
|
||||
fn global_location_and_fast_policy_reach_every_account_retry() {
|
||||
let operation = generate_operation();
|
||||
let location = gateway_core::account::RequestLocation {
|
||||
timezone: "Asia/Tokyo".parse().unwrap(),
|
||||
..Default::default()
|
||||
};
|
||||
let route_plan = plan_with_location(
|
||||
let route_plan = plan_with_location_and_fast_policy(
|
||||
&operation,
|
||||
plan(&operation).account_selection_policy(),
|
||||
location.clone(),
|
||||
true,
|
||||
);
|
||||
let (coordinator, _, provider) = coordinator(vec![
|
||||
Script::Stream {
|
||||
@@ -4550,6 +4581,6 @@ fn global_request_location_should_reach_every_account_retry() {
|
||||
assert!(
|
||||
contexts
|
||||
.iter()
|
||||
.all(|context| context.request_location() == Some(&location))
|
||||
.all(|context| context.request_location() == Some(&location) && context.disable_fast())
|
||||
);
|
||||
}
|
||||
|
||||
@@ -483,7 +483,9 @@ fn detached_early_failure_resumes_cancelled_store_write_and_settles_once_for_all
|
||||
});
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ChargedProvider {
|
||||
policies: Mutex<Vec<bool>>,
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
@@ -510,8 +512,9 @@ impl Provider for ChargedProvider {
|
||||
async fn execute(
|
||||
&self,
|
||||
request: ProviderRequest,
|
||||
_: AttemptContext,
|
||||
context: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
self.policies.lock().unwrap().push(context.disable_fast());
|
||||
let candidate = request.candidate();
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
candidate.provider().clone(),
|
||||
@@ -552,7 +555,11 @@ fn charged_service(
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([Arc::new(ChargedProvider { fail }) as Arc<dyn Provider>]).unwrap(),
|
||||
ProviderRegistry::new([Arc::new(ChargedProvider {
|
||||
fail,
|
||||
..Default::default()
|
||||
}) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
@@ -2240,3 +2247,32 @@ fn cancelling_pending_admission_cleans_up_a_slot_acquired_before_the_reply() {
|
||||
drop(pending);
|
||||
assert!(admissions.active.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reused_websocket_client_gets_fast_policy_from_each_new_request_snapshot() {
|
||||
let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let provider = Arc::new(ChargedProvider::default());
|
||||
let service = DefaultExecutionService::new(
|
||||
snapshots.clone(),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let client = service.authenticate("sk_start_test").unwrap();
|
||||
for (revision, disable_fast) in [(2, true), (3, false)] {
|
||||
snapshots.publish(
|
||||
start_snapshot_with_policy(revision, true, RateLimits::unlimited())
|
||||
.with_disable_fast(disable_fast),
|
||||
);
|
||||
let mut next = request(&service, ClientTransport::WebSocket);
|
||||
next.client = client.clone();
|
||||
let mut started = block_on(service.start(next)).unwrap();
|
||||
block_on(started.session.collect_uncommitted()).unwrap();
|
||||
block_on(started.session.detach_finalize());
|
||||
}
|
||||
assert_eq!(*provider.policies.lock().unwrap(), vec![true, false]);
|
||||
}
|
||||
|
||||
@@ -407,3 +407,130 @@ fn global_request_location_should_be_frozen_when_snapshot_is_published() {
|
||||
"Asia/Tokyo"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decompression_setting_should_validate_and_remain_frozen_across_publication() {
|
||||
use gateway_core::runtime::RuntimeSnapshotHandle;
|
||||
let compile = |version, bytes| {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(version),
|
||||
revision(version),
|
||||
SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
.with_responses_max_decompressed_body_bytes(bytes),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
);
|
||||
block_on(compiler(Arc::new(TestSnapshotStore::new(Ok(facts)))).compile())
|
||||
};
|
||||
for invalid in [0, u64::MAX] {
|
||||
assert!(matches!(
|
||||
compile(1, invalid),
|
||||
Err(RuntimeSnapshotCompileError::InvalidData)
|
||||
));
|
||||
}
|
||||
let handle = RuntimeSnapshotHandle::new(compile(1, 64 * 1024 * 1024).unwrap());
|
||||
let frozen = handle.acquire().unwrap();
|
||||
handle.publish(compile(2, 128 * 1024 * 1024).unwrap());
|
||||
assert_eq!(
|
||||
handle
|
||||
.acquire()
|
||||
.unwrap()
|
||||
.responses_max_decompressed_body_bytes(),
|
||||
128 * 1024 * 1024
|
||||
);
|
||||
assert_eq!(
|
||||
frozen.responses_max_decompressed_body_bytes(),
|
||||
64 * 1024 * 1024
|
||||
);
|
||||
handle.publish(compile(3, 1024 * 1024).unwrap());
|
||||
assert_eq!(
|
||||
handle
|
||||
.acquire()
|
||||
.unwrap()
|
||||
.responses_max_decompressed_body_bytes(),
|
||||
1024 * 1024
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disable_fast_uses_bound_groups_and_global_policy_without_changing_account_scope() {
|
||||
use gateway_core::account::ProviderAccountId;
|
||||
use gateway_core::routing::AccountGroupId;
|
||||
for global in [false, true] {
|
||||
for group_enabled in [false, true] {
|
||||
for bound in [false, true] {
|
||||
let group_id = AccountGroupId::new("grp_00000000000000000000000000000001").unwrap();
|
||||
let open_group_id =
|
||||
AccountGroupId::new("grp_00000000000000000000000000000002").unwrap();
|
||||
let account_id = ProviderAccountId::new("acct_fast_policy").unwrap();
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(1),
|
||||
revision(1),
|
||||
SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
.with_disable_fast(global),
|
||||
vec![SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_fast_policy").unwrap(),
|
||||
PlaintextClientApiKey::new("sk_fast_policy").unwrap(),
|
||||
if bound {
|
||||
vec![group_id.clone(), open_group_id.clone()]
|
||||
} else {
|
||||
Vec::new()
|
||||
},
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
vec![
|
||||
SnapshotAccountGroupFacts::new(
|
||||
group_id.clone(),
|
||||
"Restricted".to_owned(),
|
||||
group_enabled,
|
||||
)
|
||||
.with_disable_fast(true),
|
||||
SnapshotAccountGroupFacts::new(
|
||||
open_group_id.clone(),
|
||||
"Open".to_owned(),
|
||||
true,
|
||||
),
|
||||
],
|
||||
vec![SnapshotProviderAccountFacts::new(
|
||||
account_id.clone(),
|
||||
"alpha",
|
||||
)],
|
||||
vec![
|
||||
SnapshotAccountGroupMemberFacts::new(group_id, account_id.clone()),
|
||||
SnapshotAccountGroupMemberFacts::new(open_group_id, account_id.clone()),
|
||||
],
|
||||
);
|
||||
let snapshot = block_on(
|
||||
RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts))),
|
||||
Arc::new(TestCatalog::Unavailable),
|
||||
)
|
||||
.compile(),
|
||||
)
|
||||
.unwrap();
|
||||
let scope = snapshot
|
||||
.client_policies()
|
||||
.next()
|
||||
.unwrap()
|
||||
.account_scope()
|
||||
.clone();
|
||||
assert!(scope.allows(&account_id));
|
||||
let plan = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("public-model").unwrap(),
|
||||
&super::operation(),
|
||||
scope,
|
||||
&Default::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
plan.disable_fast(),
|
||||
global || bound,
|
||||
"global={global}, enabled={group_enabled}, bound={bound}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 使用与 Provider 请求一致的显式代理协议,执行有超时和响应大小限制的出口测试。
|
||||
|
||||
use std::{
|
||||
net::IpAddr,
|
||||
net::{IpAddr, Ipv4Addr, Ipv6Addr},
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
@@ -11,8 +11,16 @@ use gateway_admin::{model::proxies::ProxyTestResult, ports::proxy::ProxyProbe};
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use serde::Deserialize;
|
||||
|
||||
enum ProbeStrategy {
|
||||
Single(String),
|
||||
Dual {
|
||||
ipv4_endpoint: String,
|
||||
ipv6_endpoint: String,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct HttpProxyProbe {
|
||||
endpoint: String,
|
||||
strategy: ProbeStrategy,
|
||||
build_client: Arc<ProxyClientBuilder>,
|
||||
}
|
||||
|
||||
@@ -21,8 +29,11 @@ type ProxyClientBuilder =
|
||||
|
||||
impl Default for HttpProxyProbe {
|
||||
fn default() -> Self {
|
||||
// 使用双栈端点,避免仅有 IPv6 出口的代理被 IPv4 专用检测服务误判为不可用。
|
||||
Self::new("https://api64.ipify.org?format=json")
|
||||
// 分别向 IPv4 和 IPv6 专用端点并发探测,以获取真实的双栈出口地址。
|
||||
Self::new_dual(
|
||||
"https://api.ipify.org?format=json",
|
||||
"https://api6.ipify.org?format=json",
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -30,7 +41,18 @@ impl HttpProxyProbe {
|
||||
#[must_use]
|
||||
pub fn new(endpoint: impl Into<String>) -> Self {
|
||||
Self {
|
||||
endpoint: endpoint.into(),
|
||||
strategy: ProbeStrategy::Single(endpoint.into()),
|
||||
build_client: Arc::new(|builder| builder.build().map_err(|_| "无法创建代理连接")),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn new_dual(ipv4_endpoint: impl Into<String>, ipv6_endpoint: impl Into<String>) -> Self {
|
||||
Self {
|
||||
strategy: ProbeStrategy::Dual {
|
||||
ipv4_endpoint: ipv4_endpoint.into(),
|
||||
ipv6_endpoint: ipv6_endpoint.into(),
|
||||
},
|
||||
build_client: Arc::new(|builder| builder.build().map_err(|_| "无法创建代理连接")),
|
||||
}
|
||||
}
|
||||
@@ -47,7 +69,11 @@ impl HttpProxyProbe {
|
||||
self
|
||||
}
|
||||
|
||||
async fn exit_ip(&self, proxy: &OutboundProxy) -> Result<IpAddr, &'static str> {
|
||||
async fn exit_ip_at(
|
||||
&self,
|
||||
proxy: &OutboundProxy,
|
||||
endpoint: &str,
|
||||
) -> Result<IpAddr, &'static str> {
|
||||
let proxy = reqwest::Proxy::all(proxy.expose_url()).map_err(|_| "代理地址不合法")?;
|
||||
let builder = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
@@ -56,7 +82,7 @@ impl HttpProxyProbe {
|
||||
.timeout(Duration::from_secs(12))
|
||||
.redirect(reqwest::redirect::Policy::none());
|
||||
let client = (self.build_client)(builder)?;
|
||||
let mut response = client.get(&self.endpoint).send().await.map_err(|error| {
|
||||
let mut response = client.get(endpoint).send().await.map_err(|error| {
|
||||
if error.is_timeout() {
|
||||
"代理连接超时"
|
||||
} else {
|
||||
@@ -94,13 +120,117 @@ impl HttpProxyProbe {
|
||||
impl ProxyProbe for HttpProxyProbe {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult {
|
||||
let started = Instant::now();
|
||||
let result = tokio::time::timeout(Duration::from_secs(15), self.exit_ip(proxy)).await;
|
||||
let result = result.unwrap_or(Err("代理连接超时"));
|
||||
ProxyTestResult {
|
||||
success: result.is_ok(),
|
||||
latency_ms: u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
|
||||
exit_ip: result.as_ref().ok().copied(),
|
||||
message: result.map_or_else(str::to_owned, |_| "连接成功".to_owned()),
|
||||
let timeout_limit = Duration::from_secs(15);
|
||||
|
||||
match &self.strategy {
|
||||
ProbeStrategy::Single(endpoint) => {
|
||||
let result =
|
||||
tokio::time::timeout(timeout_limit, self.exit_ip_at(proxy, endpoint)).await;
|
||||
let result = result.unwrap_or(Err("代理连接超时"));
|
||||
let latency_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
|
||||
match result {
|
||||
Ok(ip) => {
|
||||
let (exit_ipv4, exit_ipv6) = match ip {
|
||||
IpAddr::V4(v4) => (Some(v4), None),
|
||||
IpAddr::V6(v6) => (None, Some(v6)),
|
||||
};
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms,
|
||||
exit_ip: Some(ip),
|
||||
exit_ipv4,
|
||||
exit_ipv6,
|
||||
message: "连接成功".to_owned(),
|
||||
}
|
||||
}
|
||||
Err(err) => ProxyTestResult {
|
||||
success: false,
|
||||
latency_ms,
|
||||
exit_ip: None,
|
||||
exit_ipv4: None,
|
||||
exit_ipv6: None,
|
||||
message: err.to_owned(),
|
||||
},
|
||||
}
|
||||
}
|
||||
ProbeStrategy::Dual {
|
||||
ipv4_endpoint,
|
||||
ipv6_endpoint,
|
||||
} => {
|
||||
let probe_dual = async {
|
||||
tokio::join!(
|
||||
self.exit_ip_at(proxy, ipv4_endpoint),
|
||||
self.exit_ip_at(proxy, ipv6_endpoint),
|
||||
)
|
||||
};
|
||||
let result = tokio::time::timeout(timeout_limit, probe_dual).await;
|
||||
let latency_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
|
||||
match result {
|
||||
Ok((res_v4, res_v6)) => {
|
||||
let exit_ipv4: Option<Ipv4Addr> = match res_v4 {
|
||||
Ok(IpAddr::V4(v4)) => Some(v4),
|
||||
_ => None,
|
||||
};
|
||||
let exit_ipv6: Option<Ipv6Addr> = match res_v6 {
|
||||
Ok(IpAddr::V6(v6)) => Some(v6),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if exit_ipv4.is_some() && exit_ipv6.is_some() {
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms,
|
||||
exit_ip: exit_ipv4.map(IpAddr::V4),
|
||||
exit_ipv4,
|
||||
exit_ipv6,
|
||||
message: "连接成功(双栈可用)".to_owned(),
|
||||
}
|
||||
} else if let Some(v4) = exit_ipv4 {
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms,
|
||||
exit_ip: Some(IpAddr::V4(v4)),
|
||||
exit_ipv4: Some(v4),
|
||||
exit_ipv6: None,
|
||||
message: "连接成功(仅 IPv4)".to_owned(),
|
||||
}
|
||||
} else if let Some(v6) = exit_ipv6 {
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms,
|
||||
exit_ip: Some(IpAddr::V6(v6)),
|
||||
exit_ipv4: None,
|
||||
exit_ipv6: Some(v6),
|
||||
message: "连接成功(仅 IPv6)".to_owned(),
|
||||
}
|
||||
} else {
|
||||
let message = res_v4
|
||||
.err()
|
||||
.or(res_v6.err())
|
||||
.unwrap_or("代理连接失败")
|
||||
.to_owned();
|
||||
ProxyTestResult {
|
||||
success: false,
|
||||
latency_ms,
|
||||
exit_ip: None,
|
||||
exit_ipv4: None,
|
||||
exit_ipv6: None,
|
||||
message,
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(_) => ProxyTestResult {
|
||||
success: false,
|
||||
latency_ms,
|
||||
exit_ip: None,
|
||||
exit_ipv4: None,
|
||||
exit_ipv6: None,
|
||||
message: "代理连接超时".to_owned(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ use gateway_host::proxy_probe::HttpProxyProbe;
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{any, header},
|
||||
matchers::{any, header, path},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -89,3 +89,61 @@ async fn invalid_certificate_configuration_should_not_fall_back_or_expose_detail
|
||||
assert!(!result.success);
|
||||
assert!(!result.message.contains("private-certificate-path"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dual_stack_proxy_probe_reports_both_addresses_when_available() {
|
||||
let proxy_server = MockServer::start().await;
|
||||
Mock::given(path("/v4"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ip": "203.0.113.8"})))
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
|
||||
Mock::given(path("/v6"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ip": "2001:db8::8"})))
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
|
||||
let proxy = OutboundProxy::parse(&proxy_server.uri()).unwrap();
|
||||
let result = HttpProxyProbe::new_dual(
|
||||
format!("{}/v4", proxy_server.uri()),
|
||||
format!("{}/v6", proxy_server.uri()),
|
||||
)
|
||||
.test(&proxy)
|
||||
.await;
|
||||
|
||||
assert!(result.success);
|
||||
assert_eq!(result.exit_ipv4.unwrap().to_string(), "203.0.113.8");
|
||||
assert_eq!(result.exit_ipv6.unwrap().to_string(), "2001:db8::8");
|
||||
assert!(result.message.contains("双栈可用"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dual_stack_proxy_probe_reports_single_stack_when_only_one_succeeds() {
|
||||
let proxy_server = MockServer::start().await;
|
||||
Mock::given(path("/v4"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ip": "203.0.113.8"})))
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
|
||||
Mock::given(path("/v6"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.expect(1)
|
||||
.mount(&proxy_server)
|
||||
.await;
|
||||
|
||||
let proxy = OutboundProxy::parse(&proxy_server.uri()).unwrap();
|
||||
let result = HttpProxyProbe::new_dual(
|
||||
format!("{}/v4", proxy_server.uri()),
|
||||
format!("{}/v6", proxy_server.uri()),
|
||||
)
|
||||
.test(&proxy)
|
||||
.await;
|
||||
|
||||
assert!(result.success);
|
||||
assert_eq!(result.exit_ipv4.unwrap().to_string(), "203.0.113.8");
|
||||
assert!(result.exit_ipv6.is_none());
|
||||
assert!(result.message.contains("仅 IPv4"));
|
||||
}
|
||||
|
||||
@@ -5,6 +5,56 @@ use serde_json::Value;
|
||||
|
||||
use super::sse::{SseError, parse_sse_events};
|
||||
|
||||
/// 单次 Responses 流明确声明的模型;终态优先,缺失时不使用请求模型补齐。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ResponseModelObservation {
|
||||
model: Option<String>,
|
||||
terminal: bool,
|
||||
}
|
||||
|
||||
impl ResponseModelObservation {
|
||||
/// 从原始事件读取响应模型,避免工具转换或 canonical 兜底改写观测事实。
|
||||
pub fn observe(&mut self, event_type: Option<&str>, value: &Value) {
|
||||
let terminal = matches!(
|
||||
event_type,
|
||||
Some(
|
||||
"response.completed"
|
||||
| "response.done"
|
||||
| "response.incomplete"
|
||||
| "response.failed"
|
||||
| "response.cancelled"
|
||||
| "response.canceled"
|
||||
)
|
||||
);
|
||||
if self.terminal || (self.model.is_some() && !terminal) {
|
||||
return;
|
||||
}
|
||||
let Some(model) = value
|
||||
.pointer("/response/model")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(observed_model_name)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
self.model = Some(model.to_owned());
|
||||
self.terminal = terminal;
|
||||
}
|
||||
|
||||
/// 返回原始响应声明,未声明时保持未知。
|
||||
#[must_use]
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
self.model.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 规范化可展示的上游模型名,非法值不影响响应交付。
|
||||
#[must_use]
|
||||
pub fn observed_model_name(value: &str) -> Option<&str> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty() && value.len() <= 256 && !value.chars().any(char::is_control))
|
||||
.then_some(value)
|
||||
}
|
||||
|
||||
/// 从 Codex/OpenAI usage 结构中提取出的标准化 token 用量。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
|
||||
@@ -5,6 +5,50 @@ use gateway_protocol::openai::events::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn response_model_observation_keeps_first_declaration_until_terminal() {
|
||||
use gateway_protocol::openai::events::ResponseModelObservation;
|
||||
|
||||
let mut observed = ResponseModelObservation::default();
|
||||
observed.observe(
|
||||
Some("response.created"),
|
||||
&json!({"response": {"model": " first "}}),
|
||||
);
|
||||
observed.observe(
|
||||
Some("response.in_progress"),
|
||||
&json!({"response": {"model": "intermediate"}}),
|
||||
);
|
||||
assert_eq!(observed.model(), Some("first"));
|
||||
observed.observe(
|
||||
Some("response.completed"),
|
||||
&json!({"response": {"model": "final"}}),
|
||||
);
|
||||
observed.observe(
|
||||
Some("response.in_progress"),
|
||||
&json!({"response": {"model": "late"}}),
|
||||
);
|
||||
assert_eq!(observed.model(), Some("final"));
|
||||
assert_eq!(ResponseModelObservation::default().model(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_model_observation_ignores_missing_and_invalid_values() {
|
||||
use gateway_protocol::openai::events::ResponseModelObservation;
|
||||
|
||||
for response in [
|
||||
json!({}),
|
||||
json!({"model": null}),
|
||||
json!({"model": 42}),
|
||||
json!({"model": " "}),
|
||||
json!({"model": "a\nb"}),
|
||||
json!({"model": "a".repeat(257)}),
|
||||
] {
|
||||
let mut observed = ResponseModelObservation::default();
|
||||
observed.observe(Some("response.completed"), &json!({"response": response}));
|
||||
assert_eq!(observed.model(), None);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billable_usage_should_validate_cache_writes_totals_and_overflow() {
|
||||
for (raw, expected) in [
|
||||
|
||||
@@ -43,12 +43,15 @@ impl SettingsStore for AdminSettingsStoreAdapter {
|
||||
refresh_margin_seconds: command.refresh_margin_seconds,
|
||||
refresh_concurrency: command.refresh_concurrency,
|
||||
max_concurrent_per_account: command.max_concurrent_per_account,
|
||||
disable_fast: command.disable_fast,
|
||||
request_location_enabled: command.request_location_enabled,
|
||||
request_location: command.request_location,
|
||||
request_interval_ms: command.request_interval_ms,
|
||||
max_waiting_per_key: command.max_waiting_per_key,
|
||||
max_waiting_per_account: command.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds,
|
||||
responses_max_decompressed_body_bytes: command
|
||||
.responses_max_decompressed_body_bytes,
|
||||
rotation_strategy: command.rotation_strategy.as_str().to_owned(),
|
||||
model_mappings: store_model_mappings(command.model_mappings),
|
||||
min_codex_desktop_version: command.min_codex_desktop_version,
|
||||
@@ -71,6 +74,7 @@ impl SettingsStore for AdminSettingsStoreAdapter {
|
||||
"runtime_settings",
|
||||
"1",
|
||||
vec![
|
||||
"disable_fast".to_owned(),
|
||||
"request_location_enabled".to_owned(),
|
||||
"request_location_json".to_owned(),
|
||||
"model_mappings_json".to_owned(),
|
||||
@@ -81,6 +85,7 @@ impl SettingsStore for AdminSettingsStoreAdapter {
|
||||
"max_waiting_per_key".to_owned(),
|
||||
"max_waiting_per_account".to_owned(),
|
||||
"concurrency_wait_timeout_seconds".to_owned(),
|
||||
"responses_max_decompressed_body_bytes".to_owned(),
|
||||
"rotation_strategy".to_owned(),
|
||||
"min_codex_desktop_version".to_owned(),
|
||||
"min_codex_cli_version".to_owned(),
|
||||
@@ -180,6 +185,7 @@ pub(crate) fn admin_runtime_settings(
|
||||
.collect::<AdminStoreResult<ModelMappings>>()?;
|
||||
Ok(AdminRuntimeSettings {
|
||||
config_revision: admin_revision(settings.config_revision)?,
|
||||
disable_fast: settings.disable_fast,
|
||||
request_location_enabled: settings.request_location_enabled,
|
||||
request_location: settings.request_location,
|
||||
model_mappings,
|
||||
@@ -190,6 +196,7 @@ pub(crate) fn admin_runtime_settings(
|
||||
max_waiting_per_key: settings.max_waiting_per_key,
|
||||
max_waiting_per_account: settings.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
|
||||
responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes,
|
||||
rotation_strategy,
|
||||
min_codex_desktop_version: settings.min_codex_desktop_version,
|
||||
min_codex_cli_version: settings.min_codex_cli_version,
|
||||
|
||||
@@ -234,20 +234,25 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
"create",
|
||||
"account_group",
|
||||
id.as_str(),
|
||||
vec!["name".to_owned(), "description".to_owned()],
|
||||
vec![
|
||||
"name".to_owned(),
|
||||
"description".to_owned(),
|
||||
"disable_fast".to_owned(),
|
||||
],
|
||||
);
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
sqlx::query(
|
||||
"insert into account_groups
|
||||
(id, name, description, color, enabled, created_at, updated_at)
|
||||
values ($1, $2, $3, $4, true, now(), now())",
|
||||
(id, name, description, color, disable_fast, enabled, created_at, updated_at)
|
||||
values ($1, $2, $3, $4, $5, true, now(), now())",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.name)
|
||||
.bind(command.description)
|
||||
.bind(command.color.as_str())
|
||||
.bind(command.disable_fast)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
|
||||
@@ -274,20 +279,25 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
"update",
|
||||
"account_group",
|
||||
id.as_str(),
|
||||
vec!["name".to_owned(), "description".to_owned()],
|
||||
vec![
|
||||
"name".to_owned(),
|
||||
"description".to_owned(),
|
||||
"disable_fast".to_owned(),
|
||||
],
|
||||
);
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
let result = sqlx::query(
|
||||
"update account_groups
|
||||
set name = $2, description = $3, color = $4, updated_at = now()
|
||||
set name = $2, description = $3, color = $4, disable_fast = coalesce($5, disable_fast), updated_at = now()
|
||||
where id = $1",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.name)
|
||||
.bind(command.description)
|
||||
.bind(command.color.as_str())
|
||||
.bind(command.disable_fast)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
|
||||
@@ -386,7 +396,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
|
||||
fn group_select() -> QueryBuilder<Postgres> {
|
||||
QueryBuilder::new(
|
||||
"select g.id, g.name, g.description, g.color, g.enabled, g.created_at, g.updated_at,
|
||||
"select g.id, g.name, g.description, g.color, g.enabled, g.disable_fast, g.created_at, g.updated_at,
|
||||
coalesce(members.member_count, 0)::bigint as member_count,
|
||||
coalesce(keys.client_key_count, 0)::bigint as client_key_count,
|
||||
coalesce(members.provider_counts, '{}'::jsonb) as provider_counts
|
||||
@@ -463,6 +473,9 @@ fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord>
|
||||
let provider_counts = serde_json::from_value::<BTreeMap<String, u64>>(provider_counts)
|
||||
.map_err(|_| invalid("invalid provider counts"))?;
|
||||
Ok(AccountGroupRecord {
|
||||
disable_fast: row
|
||||
.try_get("disable_fast")
|
||||
.map_err(|_| invalid("invalid disable_fast"))?,
|
||||
id: AccountGroupId::new(
|
||||
row.try_get::<String, _>("id")
|
||||
.map_err(|_| invalid("invalid id"))?,
|
||||
|
||||
@@ -282,6 +282,7 @@ pub struct ModelRequestFinalization {
|
||||
pub http_version: Option<String>,
|
||||
pub websocket_pool: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub provider_metadata_json: Option<Value>,
|
||||
pub diagnostic_trace_json: Option<Value>,
|
||||
pub error_kind: Option<String>,
|
||||
@@ -347,6 +348,12 @@ impl ModelRequestFinalization {
|
||||
return Err(invalid("service tier is invalid"));
|
||||
}
|
||||
}
|
||||
if let Some(model) = self.upstream_response_model.as_deref() {
|
||||
require_nonempty(ENTITY, "upstream_response_model", model)?;
|
||||
if model.len() > 256 || model.chars().any(char::is_control) {
|
||||
return Err(invalid("upstream response model is invalid"));
|
||||
}
|
||||
}
|
||||
if self
|
||||
.provider_metadata_json
|
||||
.as_ref()
|
||||
@@ -735,7 +742,8 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
upstream_connection_id = $44,
|
||||
upstream_connection_exit_reason = $45,
|
||||
upstream_connection_age_ms = $46,
|
||||
upstream_connection_idle_ms = $47, diagnostic_trace_json = $48
|
||||
upstream_connection_idle_ms = $47, diagnostic_trace_json = $48,
|
||||
upstream_response_model = $49
|
||||
where id = $1 and outcome = 'running'
|
||||
returning id, client_api_key_ref, continuation_affinity_hash,
|
||||
continuation_requested, provider_kind, upstream_transport,
|
||||
@@ -927,6 +935,7 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
"upstream_connection_idle_ms",
|
||||
)?)
|
||||
.bind(finalization.diagnostic_trace_json.map(sqlx::types::Json))
|
||||
.bind(finalization.upstream_response_model)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("finalize model request"))?;
|
||||
@@ -1223,6 +1232,7 @@ impl ExecutionStore for PgExecutionStore {
|
||||
http_version: finalization.http_version,
|
||||
websocket_pool: finalization.websocket_pool,
|
||||
service_tier: finalization.service_tier,
|
||||
upstream_response_model: finalization.upstream_response_model,
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: finalization
|
||||
.diagnostic_trace_json
|
||||
|
||||
@@ -582,6 +582,7 @@ impl ExecutionObservationWrite {
|
||||
finalization.http_version.as_deref(),
|
||||
finalization.websocket_pool.as_deref(),
|
||||
finalization.service_tier.as_deref(),
|
||||
finalization.upstream_response_model.as_deref(),
|
||||
finalization.provider_metadata_json.as_deref(),
|
||||
finalization.diagnostic_trace_json.as_deref(),
|
||||
finalization.provider_error_code.as_deref(),
|
||||
|
||||
@@ -439,6 +439,7 @@ pub(crate) fn admin_usage_list_record(
|
||||
provider_account_authentication_kind: record.provider_account_authentication_kind,
|
||||
upstream_model_id: record.upstream_model_id,
|
||||
upstream_transport: record.upstream_transport,
|
||||
upstream_response_model: record.upstream_response_model,
|
||||
service_tier: record.service_tier,
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
@@ -515,6 +516,7 @@ pub(crate) fn admin_usage_record(
|
||||
upstream_transport: record.upstream_transport,
|
||||
http_version: record.http_version,
|
||||
websocket_pool: record.websocket_pool,
|
||||
upstream_response_model: record.upstream_response_model,
|
||||
service_tier: record.service_tier,
|
||||
provider_metadata_json: record.provider_metadata_json,
|
||||
attempt_count: record.attempt_count,
|
||||
@@ -764,6 +766,7 @@ pub(crate) fn usage_list_record_from_row(
|
||||
provider_account_authentication_kind: get(row, "provider_account_authentication_kind")?,
|
||||
upstream_model_id: get(row, "upstream_model_id")?,
|
||||
upstream_transport: get(row, "upstream_transport")?,
|
||||
upstream_response_model: get(row, "upstream_response_model")?,
|
||||
service_tier: get(row, "service_tier")?,
|
||||
input_tokens: optional_unsigned(row, "input_tokens")?,
|
||||
output_tokens: optional_unsigned(row, "output_tokens")?,
|
||||
@@ -825,6 +828,7 @@ pub(crate) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<
|
||||
upstream_transport: get(row, "upstream_transport")?,
|
||||
http_version: get(row, "http_version")?,
|
||||
websocket_pool: get(row, "websocket_pool")?,
|
||||
upstream_response_model: get(row, "upstream_response_model")?,
|
||||
service_tier: get(row, "service_tier")?,
|
||||
provider_metadata_json: get::<Option<serde_json::Value>>(row, "provider_observation_json")?
|
||||
.map(|value| serde_json::to_string(&value))
|
||||
|
||||
@@ -490,6 +490,7 @@ pub struct UsageListRecord {
|
||||
pub provider_account_authentication_kind: Option<String>,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
@@ -546,6 +547,7 @@ pub struct UsageRecord {
|
||||
pub upstream_transport: Option<String>,
|
||||
pub http_version: Option<String>,
|
||||
pub websocket_pool: Option<String>,
|
||||
pub upstream_response_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub provider_metadata_json: Option<String>,
|
||||
pub attempt_count: u32,
|
||||
|
||||
@@ -130,7 +130,7 @@ pub(crate) const USAGE_LIST_RECORD_SELECT: &str =
|
||||
mr.provider_account_email_snapshot as provider_account_email,
|
||||
mr.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
mr.upstream_model_id, mr.upstream_transport, mr.service_tier,
|
||||
mr.upstream_model_id, mr.upstream_transport, mr.upstream_response_model, mr.service_tier,
|
||||
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
|
||||
mr.reasoning_tokens, mr.image_input_tokens, mr.image_output_tokens,
|
||||
mr.total_tokens, mr.cost_source, mr.cost_amount::text, mr.cost_currency,
|
||||
@@ -154,7 +154,7 @@ pub(crate) const USAGE_RECORD_DETAIL_SELECT: &str =
|
||||
mr.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
mr.upstream_model_id, mr.upstream_transport, mr.http_version, mr.websocket_pool,
|
||||
mr.service_tier, mr.provider_observation_json, mr.diagnostic_trace_json,
|
||||
mr.service_tier, mr.upstream_response_model, mr.provider_observation_json, mr.diagnostic_trace_json,
|
||||
mr.attempt_count, mr.upstream_send_state, mr.downstream_committed_at,
|
||||
mr.outcome, mr.client_status_code, mr.upstream_status_code,
|
||||
mr.client_response_id, mr.upstream_request_id, mr.upstream_response_id,
|
||||
|
||||
@@ -96,6 +96,8 @@ fn invalid() -> StoreError {
|
||||
fn record(row: PgRow) -> StoreResult<ProxyRecord> {
|
||||
let success: Option<bool> = row.try_get("last_test_success").map_err(|_| invalid())?;
|
||||
let ip: Option<String> = row.try_get("last_test_ip").map_err(|_| invalid())?;
|
||||
let ipv4: Option<String> = row.try_get("last_test_ipv4").map_err(|_| invalid())?;
|
||||
let ipv6: Option<String> = row.try_get("last_test_ipv6").map_err(|_| invalid())?;
|
||||
let latency: Option<i64> = row.try_get("last_test_latency_ms").map_err(|_| invalid())?;
|
||||
Ok(ProxyRecord {
|
||||
location: location_from_row(&row)?,
|
||||
@@ -124,6 +126,12 @@ fn record(row: PgRow) -> StoreResult<ProxyRecord> {
|
||||
latency_ms: u64::try_from(latency.ok_or_else(invalid)?)
|
||||
.map_err(|_| invalid())?,
|
||||
exit_ip: ip.map(|ip| ip.parse().map_err(|_| invalid())).transpose()?,
|
||||
exit_ipv4: ipv4
|
||||
.map(|ip| ip.parse().map_err(|_| invalid()))
|
||||
.transpose()?,
|
||||
exit_ipv6: ipv6
|
||||
.map(|ip| ip.parse().map_err(|_| invalid()))
|
||||
.transpose()?,
|
||||
message: row
|
||||
.try_get::<Option<String>, _>("last_test_message")
|
||||
.map_err(|_| invalid())?
|
||||
@@ -222,8 +230,8 @@ pub(crate) async fn resolve_proxy_selection(
|
||||
Ok((Some(id), Some(proxy.clone())))
|
||||
}
|
||||
AccountProxySelection::Saved(id) => {
|
||||
let (value, tested) = sqlx::query_as::<_, (String, Option<bool>)>(
|
||||
"select proxy_url, last_test_success from outbound_proxies where id = $1 for share",
|
||||
let value: String = sqlx::query_scalar(
|
||||
"select proxy_url from outbound_proxies where id = $1 for share",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&mut **transaction)
|
||||
@@ -233,9 +241,6 @@ pub(crate) async fn resolve_proxy_selection(
|
||||
entity: ENTITY,
|
||||
id: id.clone(),
|
||||
})?;
|
||||
if tested != Some(true) {
|
||||
return Err(conflict(id));
|
||||
}
|
||||
Ok((
|
||||
Some(id.clone()),
|
||||
Some(OutboundProxy::parse(&value).map_err(|_| invalid())?),
|
||||
@@ -441,10 +446,18 @@ impl ProxyStore for PgProxyRepository {
|
||||
if !acquired {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let record = self.get(id).await?;
|
||||
if !record.last_test.is_some_and(|test| test.success) {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let record = match self.get(id).await {
|
||||
Ok(record) => record,
|
||||
Err(error) => {
|
||||
// 拒绝预留时先等待数据库释放锁,避免连接关闭尚未生效就误挡后续代理操作。
|
||||
sqlx::query("select pg_advisory_unlock_shared(hashtextextended($1, 739219))")
|
||||
.bind(id)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.map_err(|_| store_error(unavailable()))?;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
Ok(ProxyImportReservation {
|
||||
binding: ImportProxyBinding {
|
||||
id: record.id,
|
||||
@@ -591,6 +604,8 @@ impl ProxyStore for PgProxyRepository {
|
||||
last_test_success = case when $4 is not null and $4 <> proxy_url then null else last_test_success end,
|
||||
last_test_latency_ms = case when $4 is not null and $4 <> proxy_url then null else last_test_latency_ms end,
|
||||
last_test_ip = case when $4 is not null and $4 <> proxy_url then null else last_test_ip end,
|
||||
last_test_ipv4 = case when $4 is not null and $4 <> proxy_url then null else last_test_ipv4 end,
|
||||
last_test_ipv6 = case when $4 is not null and $4 <> proxy_url then null else last_test_ipv6 end,
|
||||
last_test_message = case when $4 is not null and $4 <> proxy_url then null else last_test_message end
|
||||
where id = $1 and revision = $2")
|
||||
.bind(&command.id).bind(i64::try_from(command.revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
@@ -697,10 +712,13 @@ impl ProxyStore for PgProxyRepository {
|
||||
exclude_active_imports(&mut transaction, id)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
let updated = sqlx::query("update outbound_proxies set last_test_at = now(), last_test_success = $3, last_test_latency_ms = $4, last_test_ip = $5, last_test_message = $6 where id = $1 and revision = $2")
|
||||
let updated = sqlx::query("update outbound_proxies set last_test_at = now(), last_test_success = $3, last_test_latency_ms = $4, last_test_ip = $5, last_test_ipv4 = $6, last_test_ipv6 = $7, last_test_message = $8 where id = $1 and revision = $2")
|
||||
.bind(id).bind(i64::try_from(revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.bind(result.success).bind(i64::try_from(result.latency_ms).map_err(|_| store_error(invalid()))?)
|
||||
.bind(result.exit_ip.map(|ip| ip.to_string())).bind(result.message)
|
||||
.bind(result.exit_ip.map(|ip| ip.to_string()))
|
||||
.bind(result.exit_ipv4.map(|ip| ip.to_string()))
|
||||
.bind(result.exit_ipv6.map(|ip| ip.to_string()))
|
||||
.bind(result.message)
|
||||
.execute(&mut *transaction).await.map_err(|_| store_error(unavailable()))?;
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(id)));
|
||||
|
||||
@@ -19,6 +19,7 @@ use crate::{Revision, StoreError, StoreResult, postgres_unavailable};
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct RuntimeSettings {
|
||||
pub disable_fast: bool,
|
||||
pub config_revision: Revision,
|
||||
pub admin_api_key: Option<String>,
|
||||
pub refresh_margin_seconds: u64,
|
||||
@@ -28,6 +29,7 @@ pub struct RuntimeSettings {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: String,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
@@ -63,6 +65,7 @@ impl fmt::Debug for RuntimeSettings {
|
||||
&self.max_concurrent_per_account,
|
||||
)
|
||||
.field("request_interval_ms", &self.request_interval_ms)
|
||||
.field("disable_fast", &self.disable_fast)
|
||||
.field("rotation_strategy", &self.rotation_strategy)
|
||||
.field("request_location_enabled", &self.request_location_enabled)
|
||||
.field("request_location", &self.request_location)
|
||||
@@ -107,6 +110,7 @@ impl fmt::Debug for RuntimeSettings {
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RuntimeSettingsUpdate {
|
||||
pub disable_fast: Option<bool>,
|
||||
pub admin_api_key: Option<String>,
|
||||
pub refresh_margin_seconds: u64,
|
||||
pub refresh_concurrency: u32,
|
||||
@@ -115,6 +119,7 @@ pub struct RuntimeSettingsUpdate {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: String,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
@@ -141,6 +146,7 @@ impl fmt::Debug for RuntimeSettingsUpdate {
|
||||
"admin_api_key",
|
||||
&self.admin_api_key.as_ref().map(|_| "[REDACTED]"),
|
||||
)
|
||||
.field("disable_fast", &self.disable_fast)
|
||||
.field("rotation_strategy", &self.rotation_strategy)
|
||||
.field("request_location_enabled", &self.request_location_enabled)
|
||||
.field("request_location", &self.request_location)
|
||||
@@ -152,6 +158,8 @@ impl fmt::Debug for RuntimeSettingsUpdate {
|
||||
impl RuntimeSettingsUpdate {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
if self.request_location.validate().is_err()
|
||||
|| self.responses_max_decompressed_body_bytes == 0
|
||||
|| isize::try_from(self.responses_max_decompressed_body_bytes).is_err()
|
||||
|| self.refresh_margin_seconds == 0
|
||||
|| self.refresh_concurrency == 0
|
||||
|| self.max_concurrent_per_account == 0
|
||||
@@ -225,11 +233,11 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
|
||||
|
||||
pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResult<RuntimeSettings> {
|
||||
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled,
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, disable_fast,
|
||||
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
|
||||
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
|
||||
audit_retention_days, min_codex_desktop_version,
|
||||
min_codex_cli_version, updated_at, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
|
||||
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
|
||||
account_auto_freeze_enabled, account_auto_freeze_threshold,
|
||||
account_auto_freeze_window_seconds, account_auto_freeze_duration_seconds,
|
||||
account_auto_freeze_probe_enabled, account_auto_freeze_probe_model,
|
||||
@@ -287,11 +295,11 @@ pub(crate) async fn load_runtime_settings_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<RuntimeSettings> {
|
||||
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled,
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, disable_fast,
|
||||
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
|
||||
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
|
||||
audit_retention_days, min_codex_desktop_version,
|
||||
min_codex_cli_version, updated_at, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
|
||||
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
|
||||
account_auto_freeze_enabled, account_auto_freeze_threshold,
|
||||
account_auto_freeze_window_seconds, account_auto_freeze_duration_seconds,
|
||||
account_auto_freeze_probe_enabled, account_auto_freeze_probe_model,
|
||||
@@ -342,6 +350,8 @@ pub(crate) async fn update_runtime_settings_in_transaction(
|
||||
account_auto_freeze_adaptive_concurrency = $22,
|
||||
request_location_json = $23,
|
||||
request_location_enabled = $24,
|
||||
responses_max_decompressed_body_bytes = $25,
|
||||
disable_fast = coalesce($26, disable_fast),
|
||||
updated_at = now()
|
||||
where id = 1
|
||||
returning config_revision",
|
||||
@@ -379,6 +389,11 @@ pub(crate) async fn update_runtime_settings_in_transaction(
|
||||
.map_err(|_| invalid_location())?,
|
||||
))
|
||||
.bind(update.request_location_enabled)
|
||||
.bind(
|
||||
i64::try_from(update.responses_max_decompressed_body_bytes)
|
||||
.map_err(|_| invalid_numeric())?,
|
||||
)
|
||||
.bind(update.disable_fast)
|
||||
.fetch_optional(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update runtime settings in transaction"))?
|
||||
@@ -428,6 +443,7 @@ pub(crate) async fn update_admin_api_key_in_transaction(
|
||||
|
||||
#[derive(sqlx::FromRow)]
|
||||
struct RuntimeSettingsRow {
|
||||
disable_fast: bool,
|
||||
config_revision: i64,
|
||||
admin_api_key: Option<String>,
|
||||
refresh_margin_seconds: i64,
|
||||
@@ -447,6 +463,7 @@ struct RuntimeSettingsRow {
|
||||
max_waiting_per_key: i64,
|
||||
max_waiting_per_account: i64,
|
||||
concurrency_wait_timeout_seconds: i64,
|
||||
responses_max_decompressed_body_bytes: i64,
|
||||
account_auto_freeze_enabled: bool,
|
||||
account_auto_freeze_threshold: i64,
|
||||
account_auto_freeze_window_seconds: i64,
|
||||
@@ -465,6 +482,7 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
|
||||
max_concurrent_per_account: to_u32(row.max_concurrent_per_account)?,
|
||||
request_interval_ms: to_u64(row.request_interval_ms)?,
|
||||
rotation_strategy: row.rotation_strategy,
|
||||
disable_fast: row.disable_fast,
|
||||
request_location_enabled: row.request_location_enabled,
|
||||
request_location: row
|
||||
.request_location_json
|
||||
@@ -481,6 +499,7 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
|
||||
max_waiting_per_key: to_u32(row.max_waiting_per_key)?,
|
||||
max_waiting_per_account: to_u32(row.max_waiting_per_account)?,
|
||||
concurrency_wait_timeout_seconds: to_u32(row.concurrency_wait_timeout_seconds)?,
|
||||
responses_max_decompressed_body_bytes: to_u64(row.responses_max_decompressed_body_bytes)?,
|
||||
account_auto_freeze_enabled: row.account_auto_freeze_enabled,
|
||||
account_auto_freeze_threshold: to_u32(row.account_auto_freeze_threshold)?,
|
||||
account_auto_freeze_window_seconds: to_u64(row.account_auto_freeze_window_seconds)?,
|
||||
|
||||
@@ -20,6 +20,7 @@ use super::ClientApiKeySnapshot;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotRuntimeSettings {
|
||||
pub disable_fast: bool,
|
||||
pub request_location_enabled: bool,
|
||||
pub request_location: gateway_core::account::RequestLocation,
|
||||
pub refresh_margin_seconds: u64,
|
||||
@@ -29,6 +30,7 @@ pub struct SnapshotRuntimeSettings {
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
pub responses_max_decompressed_body_bytes: u64,
|
||||
pub rotation_strategy: String,
|
||||
pub model_mappings: BTreeMap<String, String>,
|
||||
pub min_codex_desktop_version: Option<String>,
|
||||
@@ -48,6 +50,7 @@ pub struct RuntimeSnapshotData {
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotAccountGroupData {
|
||||
pub disable_fast: bool,
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub enabled: bool,
|
||||
@@ -154,6 +157,10 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
|
||||
data.settings.min_codex_desktop_version,
|
||||
data.settings.min_codex_cli_version,
|
||||
)
|
||||
.with_responses_max_decompressed_body_bytes(
|
||||
data.settings.responses_max_decompressed_body_bytes,
|
||||
)
|
||||
.with_disable_fast(data.settings.disable_fast)
|
||||
.with_request_location(
|
||||
data.settings.request_location,
|
||||
data.settings.request_location_enabled,
|
||||
@@ -178,7 +185,10 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
|
||||
let account_groups = data
|
||||
.account_groups
|
||||
.into_iter()
|
||||
.map(|group| SnapshotAccountGroupFacts::new(group.id, group.name, group.enabled))
|
||||
.map(|group| {
|
||||
SnapshotAccountGroupFacts::new(group.id, group.name, group.enabled)
|
||||
.with_disable_fast(group.disable_fast)
|
||||
})
|
||||
.collect();
|
||||
let provider_accounts = data
|
||||
.provider_accounts
|
||||
@@ -249,12 +259,14 @@ async fn load_settings(
|
||||
i64, i64, i64,
|
||||
sqlx::types::Json<gateway_core::account::RequestLocation>,
|
||||
bool,
|
||||
i64,
|
||||
bool,
|
||||
),
|
||||
>(
|
||||
"select config_revision, refresh_margin_seconds, refresh_concurrency,
|
||||
max_concurrent_per_account, request_interval_ms, rotation_strategy,
|
||||
model_mappings_json, min_codex_desktop_version,
|
||||
min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled
|
||||
min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled, responses_max_decompressed_body_bytes, disable_fast
|
||||
from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_optional(&mut **transaction)
|
||||
@@ -267,6 +279,8 @@ async fn load_settings(
|
||||
Ok((
|
||||
revision_from_i64(row.0)?,
|
||||
SnapshotRuntimeSettings {
|
||||
disable_fast: row.15,
|
||||
responses_max_decompressed_body_bytes: to_u64(row.14)?,
|
||||
request_location_enabled: row.13,
|
||||
request_location: row.12.0,
|
||||
refresh_margin_seconds: to_u64(row.1)?,
|
||||
@@ -309,15 +323,16 @@ async fn load_client_keys(
|
||||
async fn load_account_groups(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<Vec<SnapshotAccountGroupData>> {
|
||||
let rows = sqlx::query_as::<_, (String, String, bool)>(
|
||||
"select id, name, enabled from account_groups order by id",
|
||||
let rows = sqlx::query_as::<_, (String, String, bool, bool)>(
|
||||
"select id, name, enabled, disable_fast from account_groups order by id",
|
||||
)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load snapshot account groups"))?;
|
||||
rows.into_iter()
|
||||
.map(|(id, name, enabled)| {
|
||||
.map(|(id, name, enabled, disable_fast)| {
|
||||
Ok(SnapshotAccountGroupData {
|
||||
disable_fast,
|
||||
id: AccountGroupId::new(id).map_err(|_| invalid("invalid account group id"))?,
|
||||
name,
|
||||
enabled,
|
||||
|
||||
@@ -49,6 +49,7 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
disable_fast: false,
|
||||
id: mixed_group.clone(),
|
||||
name: "Mixed Production".to_owned(),
|
||||
description: Some("cross-provider".to_owned()),
|
||||
@@ -61,6 +62,7 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
disable_fast: false,
|
||||
id: empty_group.clone(),
|
||||
name: "Empty Pool".to_owned(),
|
||||
description: None,
|
||||
@@ -255,6 +257,7 @@ async fn group_costs_should_include_statusless_websocket_but_reject_statusless_h
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
disable_fast: false,
|
||||
id: group_id(EMPTY_GROUP),
|
||||
name: "Statusless Costs".to_owned(),
|
||||
description: None,
|
||||
@@ -430,3 +433,61 @@ async fn audit_count(pool: &sqlx::PgPool) -> u64 {
|
||||
.expect("count audit rows");
|
||||
u64::try_from(value).expect("non-negative audit count")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disable_fast_group_updates_preserve_omitted_values_and_publish_snapshot_facts() {
|
||||
use gateway_admin::model::account_groups::UpdateAccountGroup;
|
||||
use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository};
|
||||
let Some(database) = TestDatabase::create("disable_fast_group").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgAccountGroupRepository::new(database.pool.clone());
|
||||
let id = group_id(MIXED_GROUP);
|
||||
repository
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
id: id.clone(),
|
||||
name: "Fast policy".to_owned(),
|
||||
description: None,
|
||||
color: group_color("#2563EBFF"),
|
||||
disable_fast: true,
|
||||
},
|
||||
&context("create-fast"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for (value, expected) in [(None, true), (Some(false), false), (Some(true), true)] {
|
||||
let mutation = repository
|
||||
.update_account_group(
|
||||
UpdateAccountGroup {
|
||||
id: id.clone(),
|
||||
name: "Renamed policy".to_owned(),
|
||||
description: None,
|
||||
color: group_color("#2563EBFF"),
|
||||
disable_fast: value,
|
||||
},
|
||||
&context("update-fast"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(mutation.record.unwrap().disable_fast, expected);
|
||||
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
|
||||
.load_runtime_snapshot()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
snapshot.config_revision.get(),
|
||||
mutation.config_revision.get()
|
||||
);
|
||||
assert_eq!(
|
||||
snapshot
|
||||
.account_groups
|
||||
.iter()
|
||||
.find(|group| group.id == id)
|
||||
.unwrap()
|
||||
.disable_fast,
|
||||
expected
|
||||
);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -356,8 +356,10 @@ async fn core_adapter_should_persist_calculated_cost_exactly() {
|
||||
let repository = PgExecutionStore::new(database.pool.clone());
|
||||
|
||||
let mut finalization = successful_core_finalization("req_calculated_cost");
|
||||
finalization.downstream_committed_at = Some(std::time::SystemTime::now());
|
||||
finalization.websocket_pool = Some("reuse".to_owned());
|
||||
finalization.service_tier = Some("priority".to_owned());
|
||||
finalization.upstream_response_model = Some("grok-4.6-build".to_owned());
|
||||
finalization.cost = CalculatedCost::from_usd_ticks(12_345)
|
||||
.expect("calculated cost")
|
||||
.into_estimate();
|
||||
@@ -393,6 +395,37 @@ async fn core_adapter_should_persist_calculated_cost_exactly() {
|
||||
Some("priority".to_owned()),
|
||||
)
|
||||
);
|
||||
let observation = observability_repository(&database.pool);
|
||||
let detail = observation
|
||||
.usage_record_detail("req_calculated_cost")
|
||||
.await
|
||||
.expect("request detail");
|
||||
assert_eq!(
|
||||
detail.request.upstream_model_id.as_deref(),
|
||||
Some("grok-4.5")
|
||||
);
|
||||
assert_eq!(
|
||||
detail.request.upstream_response_model.as_deref(),
|
||||
Some("grok-4.6-build")
|
||||
);
|
||||
let page = observation
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range: ObservabilityRange::new(
|
||||
Utc::now() - Duration::hours(1),
|
||||
Utc::now() + Duration::hours(1),
|
||||
)
|
||||
.unwrap(),
|
||||
filter: UsageRecordFilter::default(),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).unwrap(),
|
||||
})
|
||||
.await
|
||||
.expect("request list");
|
||||
assert_eq!(page.items.len(), 1);
|
||||
assert_eq!(
|
||||
page.items[0].upstream_response_model.as_deref(),
|
||||
Some("grok-4.6-build")
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -494,6 +527,7 @@ fn successful_core_finalization(id: &str) -> CoreModelRequestFinalization {
|
||||
upstream_transport: Some("websocket".to_owned()),
|
||||
http_version: Some("HTTP/2".to_owned()),
|
||||
websocket_pool: None,
|
||||
upstream_response_model: None,
|
||||
service_tier: None,
|
||||
provider_metadata_json: None,
|
||||
error: None,
|
||||
@@ -1230,6 +1264,7 @@ pub(super) fn early_failure(request: &CoreNewModelRequest) -> CoreModelRequestFi
|
||||
upstream_transport: None,
|
||||
http_version: None,
|
||||
websocket_pool: None,
|
||||
upstream_response_model: None,
|
||||
service_tier: None,
|
||||
provider_metadata_json: None,
|
||||
diagnostic_trace_json: trace.snapshot().map(|value| value.to_string()),
|
||||
|
||||
@@ -2006,6 +2006,8 @@ async fn authorization_import_rejects_a_saved_proxy_changed_during_oauth() {
|
||||
success: true,
|
||||
latency_ms: 1,
|
||||
exit_ip: Some("203.0.113.5".parse().unwrap()),
|
||||
exit_ipv4: Some("203.0.113.5".parse().unwrap()),
|
||||
exit_ipv6: None,
|
||||
message: "Connected".to_owned(),
|
||||
};
|
||||
proxies
|
||||
|
||||
@@ -111,7 +111,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() {
|
||||
revision: stale_revision,
|
||||
name: saved.name.clone(),
|
||||
proxy: None,
|
||||
location: Some(Some(location.clone()))
|
||||
location: Some(Some(location.clone())),
|
||||
},
|
||||
&context()
|
||||
)
|
||||
@@ -161,6 +161,8 @@ fn success() -> ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms: 10,
|
||||
exit_ip: Some("203.0.113.5".parse().unwrap()),
|
||||
exit_ipv4: Some("203.0.113.5".parse().unwrap()),
|
||||
exit_ipv6: None,
|
||||
message: "Connected".to_owned(),
|
||||
}
|
||||
}
|
||||
@@ -473,6 +475,31 @@ async fn proxy_accounts_paginate_thousands_of_accounts_and_search_without_loadin
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejected_import_reservations_release_proxy_lock_before_returning() {
|
||||
let Some(database) = TestDatabase::create("rejected_proxy_import").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgProxyRepository::new(database.pool.clone());
|
||||
let id = "proxy_missing";
|
||||
|
||||
// 读取不存在的代理失败后,应立即允许其他事务取得同一资源的独占锁。
|
||||
for _ in 0..64 {
|
||||
let error = store.reserve_import(id).await.err().unwrap();
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::NotFound);
|
||||
let mut transaction = database.pool.begin().await.unwrap();
|
||||
let acquired: bool =
|
||||
sqlx::query_scalar("select pg_try_advisory_xact_lock(hashtextextended($1, 739219))")
|
||||
.bind(id)
|
||||
.fetch_one(&mut *transaction)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(acquired);
|
||||
transaction.commit().await.unwrap();
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are_committed() {
|
||||
use gateway_store::postgres::{
|
||||
@@ -496,9 +523,17 @@ async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert!(store.reserve_import(&saved.id).await.is_err());
|
||||
assert!(saved.last_test.is_none());
|
||||
store
|
||||
.record_test(&saved.id, saved.revision, success(), &context)
|
||||
.record_test(
|
||||
&saved.id,
|
||||
saved.revision,
|
||||
ProxyTestResult {
|
||||
success: false,
|
||||
..success()
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let reservation = store.reserve_import(&saved.id).await.unwrap();
|
||||
@@ -625,27 +660,33 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
.is_err()
|
||||
);
|
||||
let selection = AccountProxySelection::Saved(created.id.clone());
|
||||
assert!(
|
||||
admin
|
||||
.update_account(update("acct_one", selection.clone()), &context)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
admin
|
||||
.update_account(update("acct_one", selection.clone()), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let read = || async {
|
||||
sqlx::query_as::<_, (String, Option<String>, Option<String>, i64)>("select id, outbound_proxy_id, outbound_proxy_url, credential_revision from provider_accounts order by id")
|
||||
.fetch_all(&database.pool).await.unwrap()
|
||||
};
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.is_none() && row.2.is_none() && row.3 == 1)
|
||||
);
|
||||
let untested = read().await;
|
||||
assert_eq!(untested[0].1.as_deref(), Some(created.id.as_str()));
|
||||
assert_eq!(untested[0].2.as_deref(), Some(old_proxy.expose_url()));
|
||||
assert!(untested[1].1.is_none() && untested[1].2.is_none());
|
||||
assert!(untested.iter().all(|row| row.3 == 1));
|
||||
let tested = store
|
||||
.record_test(&created.id, created.revision, success(), &context)
|
||||
.record_test(
|
||||
&created.id,
|
||||
created.revision,
|
||||
ProxyTestResult {
|
||||
success: false,
|
||||
..success()
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(tested.last_test_at.is_some());
|
||||
assert!(!tested.last_test.unwrap().success);
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
admin
|
||||
.update_account(update(id, selection.clone()), &context)
|
||||
@@ -670,6 +711,10 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
&& row.3 == 1)
|
||||
);
|
||||
|
||||
store
|
||||
.record_test(&created.id, created.revision, success(), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let renamed = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
|
||||
@@ -9,6 +9,7 @@ use super::TestDatabase;
|
||||
|
||||
fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate {
|
||||
RuntimeSettingsUpdate {
|
||||
disable_fast: None,
|
||||
request_location_enabled: false,
|
||||
request_location: Default::default(),
|
||||
admin_api_key: None,
|
||||
@@ -19,6 +20,7 @@ fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate {
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
rotation_strategy: "smart".to_owned(),
|
||||
model_mappings: BTreeMap::from([
|
||||
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
|
||||
@@ -380,3 +382,77 @@ async fn auto_freeze_defaults_off_and_explicit_opt_in_round_trips() {
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decompression_setting_should_persist_and_reach_snapshot_facts() {
|
||||
use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository};
|
||||
let Some(database) = TestDatabase::create("decompression_settings").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
|
||||
let before = repository.load_runtime_settings().await.unwrap();
|
||||
assert_eq!(
|
||||
before.responses_max_decompressed_body_bytes,
|
||||
64 * 1024 * 1024
|
||||
);
|
||||
let mut update = settings_with_margin(3600);
|
||||
update.responses_max_decompressed_body_bytes = 128 * 1024 * 1024;
|
||||
repository.update_runtime_settings(update).await.unwrap();
|
||||
let reloaded = PgRuntimeSettingsRepository::new(database.pool.clone())
|
||||
.load_runtime_settings()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
reloaded.responses_max_decompressed_body_bytes,
|
||||
128 * 1024 * 1024
|
||||
);
|
||||
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
|
||||
.load_runtime_snapshot()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
snapshot.settings.responses_max_decompressed_body_bytes,
|
||||
reloaded.responses_max_decompressed_body_bytes
|
||||
);
|
||||
assert!(snapshot.config_revision > before.config_revision);
|
||||
for invalid in [0, u64::MAX] {
|
||||
let mut update = settings_with_margin(3600);
|
||||
update.responses_max_decompressed_body_bytes = invalid;
|
||||
assert!(repository.update_runtime_settings(update).await.is_err());
|
||||
assert_eq!(
|
||||
repository
|
||||
.load_runtime_settings()
|
||||
.await
|
||||
.unwrap()
|
||||
.config_revision,
|
||||
reloaded.config_revision
|
||||
);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disable_fast_persists_and_omitted_updates_preserve_the_restriction() {
|
||||
let Some(database) = TestDatabase::create("disable_fast_settings").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
|
||||
let initial = repository.load_runtime_settings().await.unwrap();
|
||||
assert!(!initial.disable_fast);
|
||||
for (value, expected) in [(Some(true), true), (None, true), (Some(false), false)] {
|
||||
let mut update = settings_with_margin(1_800);
|
||||
update.disable_fast = value;
|
||||
let revision = repository.update_runtime_settings(update).await.unwrap();
|
||||
let reloaded = repository.load_runtime_settings().await.unwrap();
|
||||
assert_eq!(reloaded.disable_fast, expected);
|
||||
assert_eq!(reloaded.config_revision, revision);
|
||||
use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository};
|
||||
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
|
||||
.load_runtime_snapshot()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(snapshot.settings.disable_fast, expected);
|
||||
assert_eq!(snapshot.config_revision, revision);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -465,9 +465,20 @@ impl CodexCredentialCatalogService {
|
||||
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
|
||||
let scope = CodexCatalogScope::for_account(&account)?;
|
||||
let mut candidates =
|
||||
catalog_candidates_by_scope(self.repository.list_for_provider().await?)?
|
||||
.remove(&scope)
|
||||
.unwrap_or_default();
|
||||
match catalog_candidates_by_scope(self.repository.list_for_provider().await?) {
|
||||
Ok(mut groups) => groups.remove(&scope).unwrap_or_default(),
|
||||
// 全部账号都被调度列表过滤时按空候选处理,让下方目标账号补回继续生效。
|
||||
Err(CodexCredentialCatalogError::NoEligibleCredential) => Vec::new(),
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
// 常规调度列表不含停用账号;管理端按账号查询模型要对停用账号返回真实上游
|
||||
// 结果,这里把不在候选里的目标账号本身补回,仍按优先顺序先试目标账号。
|
||||
if !candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.id() == account_id)
|
||||
{
|
||||
candidates.push(account.clone());
|
||||
}
|
||||
if candidates.is_empty() {
|
||||
return Err(CodexCredentialCatalogError::NoEligibleCredential);
|
||||
}
|
||||
|
||||
@@ -661,7 +661,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
}
|
||||
apply_failure(&failure_context, &active_account, &failure)
|
||||
.await;
|
||||
Err(failure.error)?;
|
||||
Err(quota_continuation_replay_error(
|
||||
failure.error,
|
||||
&request,
|
||||
ReplayBoundary::BeforeSemanticOutput,
|
||||
))?;
|
||||
return;
|
||||
}
|
||||
};
|
||||
@@ -730,12 +734,13 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let mut passive_quota_observation =
|
||||
OpenAiPassiveQuotaObservation::new(response.rate_limit_headers);
|
||||
let rate_limit_updates = response.rate_limit_updates;
|
||||
let turn_state_updates = response.turn_state_update;
|
||||
let response_metadata_updates = response.response_metadata_updates;
|
||||
// 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_reported_model(response.response_metadata.effective_model.as_deref())
|
||||
.with_requested_service_tier(request.service_tier())
|
||||
.with_request_tool_pricing(upstream_model.as_str(), request.tools())
|
||||
.with_raw_sse_passthrough();
|
||||
@@ -798,13 +803,14 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let update_headers = rate_limit_update_headers(&updates);
|
||||
observation_state.merge_rate_limit_headers(&update_headers)
|
||||
};
|
||||
let turn_state_merge = merge_turn_state_update(
|
||||
turn_state_updates.as_ref(),
|
||||
let metadata_merge = merge_response_metadata_updates(
|
||||
response_metadata_updates.as_ref(),
|
||||
&mut session_capture,
|
||||
&mut observation_state,
|
||||
&mut decoder,
|
||||
)
|
||||
.await;
|
||||
let observation_event = if rate_limits_changed || turn_state_merge.is_some() {
|
||||
let observation_event = if rate_limits_changed || metadata_merge.is_some() {
|
||||
observation_state.observation(None).map(ProviderEvent::observation)
|
||||
} else {
|
||||
None
|
||||
@@ -845,7 +851,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
}
|
||||
apply_failure(&failure_context, &active_account, &failure)
|
||||
.await;
|
||||
Err(failure.error)?;
|
||||
Err(quota_continuation_replay_error(
|
||||
failure.error,
|
||||
&request,
|
||||
ReplayBoundary::from_semantic_output(pre_commit_events.is_committed()),
|
||||
))?;
|
||||
return;
|
||||
}
|
||||
};
|
||||
@@ -857,13 +867,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
passive_quota_observation.observe(&updates);
|
||||
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
|
||||
};
|
||||
let turn_state_merge = merge_turn_state_update(
|
||||
turn_state_updates.as_ref(),
|
||||
&mut session_capture,
|
||||
&mut observation_state,
|
||||
)
|
||||
.await;
|
||||
let turn_state_changed = turn_state_merge.unwrap_or(false);
|
||||
let first_event_changed =
|
||||
observation_state.observe_stream_chunk(&chunk, output_started_at);
|
||||
let chunk_len = chunk.len();
|
||||
@@ -874,7 +877,17 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
(events, Some((error, semantic_output_seen)))
|
||||
}
|
||||
};
|
||||
let metadata_merge = merge_response_metadata_updates(
|
||||
response_metadata_updates.as_ref(),
|
||||
&mut session_capture,
|
||||
&mut observation_state,
|
||||
&mut decoder,
|
||||
)
|
||||
.await;
|
||||
let metadata_changed = metadata_merge.unwrap_or(false);
|
||||
pre_commit_events.observe_chunk(chunk_len);
|
||||
let response_model_changed = observation_state
|
||||
.observe_upstream_response_model(decoder.response_model());
|
||||
let service_tier_changed = observation_state
|
||||
.observe_upstream_service_tier(decoder.response_service_tier());
|
||||
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
|
||||
@@ -945,9 +958,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.await;
|
||||
}
|
||||
if (rate_limits_changed
|
||||
|| response_model_changed
|
||||
|| service_tier_changed
|
||||
|| timing_changed
|
||||
|| turn_state_changed
|
||||
|| metadata_changed
|
||||
|| terminal_changed
|
||||
|| (response_transport == CodexBackendTransport::WebSocket && terminal_failure.is_some()))
|
||||
&& let Some(observation) = observation_state.observation(
|
||||
@@ -968,7 +982,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.error
|
||||
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
|
||||
}
|
||||
Err(failure.error)?;
|
||||
Err(quota_continuation_replay_error(
|
||||
failure.error,
|
||||
&request,
|
||||
ReplayBoundary::from_semantic_output(failure_after_commit),
|
||||
))?;
|
||||
return;
|
||||
}
|
||||
let events = pre_commit_events.stage(events, timing_signals, completed);
|
||||
@@ -1011,6 +1029,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
)
|
||||
});
|
||||
let timing_signals = decoder.take_timing_signals();
|
||||
let response_model_changed = observation_state
|
||||
.observe_upstream_response_model(decoder.response_model());
|
||||
let service_tier_changed = observation_state
|
||||
.observe_upstream_service_tier(decoder.response_service_tier());
|
||||
let timing_changed = observation_state
|
||||
@@ -1034,10 +1054,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
apply_failure(&failure_context, &active_account, failure)
|
||||
.await;
|
||||
}
|
||||
let turn_state_changed = merge_turn_state_update(
|
||||
turn_state_updates.as_ref(),
|
||||
let metadata_changed = merge_response_metadata_updates(
|
||||
response_metadata_updates.as_ref(),
|
||||
&mut session_capture,
|
||||
&mut observation_state,
|
||||
&mut decoder,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
@@ -1067,10 +1088,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
.await;
|
||||
}
|
||||
if (service_tier_changed
|
||||
if (response_model_changed
|
||||
|| service_tier_changed
|
||||
|| timing_changed
|
||||
|| rate_limits_changed
|
||||
|| turn_state_changed
|
||||
|| metadata_changed
|
||||
|| terminal_changed
|
||||
|| (response_transport == CodexBackendTransport::WebSocket && terminal_failure.is_some()))
|
||||
&& let Some(observation) = observation_state.observation(
|
||||
@@ -1091,7 +1113,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.error
|
||||
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
|
||||
}
|
||||
Err(failure.error)?;
|
||||
Err(quota_continuation_replay_error(
|
||||
failure.error,
|
||||
&request,
|
||||
ReplayBoundary::from_semantic_output(failure_after_commit),
|
||||
))?;
|
||||
return;
|
||||
}
|
||||
let events = pre_commit_events.finish(events, timing_signals, completed);
|
||||
@@ -1101,15 +1127,30 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
})
|
||||
}
|
||||
|
||||
async fn merge_turn_state_update(
|
||||
updates: Option<&CodexTurnStateUpdate>,
|
||||
async fn merge_response_metadata_updates(
|
||||
updates: Option<&CodexResponseMetadataUpdates>,
|
||||
session_capture: &mut Option<OpenAiSessionCapture>,
|
||||
observation_state: &mut OpenAiResponseObservationState,
|
||||
decoder: &mut CodexCanonicalDecoder,
|
||||
) -> Option<bool> {
|
||||
let updates = updates?;
|
||||
let turn_state = updates.lock().await.take()?;
|
||||
if let Some(capture) = session_capture.as_mut() {
|
||||
capture.turn_state = Some(turn_state.clone());
|
||||
let mut pending = updates.lock().await;
|
||||
let turn_state = pending.turn_state.take();
|
||||
let reported_model = pending.reported_model.clone();
|
||||
drop(pending);
|
||||
if turn_state.is_none() && reported_model.is_none() {
|
||||
return None;
|
||||
}
|
||||
Some(observation_state.merge_client_header("x-codex-turn-state", &turn_state))
|
||||
let mut changed = false;
|
||||
if let Some(turn_state) = turn_state {
|
||||
if let Some(capture) = session_capture.as_mut() {
|
||||
capture.turn_state = Some(turn_state.clone());
|
||||
}
|
||||
changed |= observation_state.merge_client_header("x-codex-turn-state", &turn_state);
|
||||
}
|
||||
if let Some(model) = reported_model {
|
||||
decoder.observe_reported_model(&model);
|
||||
changed |= observation_state.observe_upstream_response_model(decoder.response_model());
|
||||
}
|
||||
Some(changed)
|
||||
}
|
||||
|
||||
@@ -677,11 +677,80 @@ pub(super) fn continuation_replay_required_error(reason: &'static str) -> Provid
|
||||
.with_diagnostic(ProviderDiagnostic::new(
|
||||
"previous response is unavailable in the selected upstream scope; client replay is required",
|
||||
))
|
||||
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
|
||||
.with_client_visible_upstream_error(continuation_replay_error_detail())
|
||||
}
|
||||
|
||||
/// 额度拒绝使当前账号无法继续原生历史,由客户端携带完整输入创建新链。
|
||||
/// 只投影客户端响应;真实上游分类、发送状态和错误正文继续用于隔离与记账。
|
||||
pub(super) fn quota_continuation_replay_error(
|
||||
mut error: ProviderError,
|
||||
request: &CodexResponsesRequest,
|
||||
replay_boundary: ReplayBoundary,
|
||||
) -> ProviderError {
|
||||
if request.previous_response_id().is_none()
|
||||
|| error.kind() != ProviderErrorKind::QuotaExhausted
|
||||
|| !error.replay_is_safe()
|
||||
|| error.send_state() == UpstreamSendState::Ambiguous
|
||||
|| !replay_boundary.permits_provider_proof()
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
let mut headers: Vec<ProviderResponseHeader> = error
|
||||
.client_visible_upstream_response()
|
||||
.map(|response| {
|
||||
response
|
||||
.headers()
|
||||
.iter()
|
||||
// 额度窗口的等待时间不适用于客户端重建历史。
|
||||
.filter(|header| !header.name().eq_ignore_ascii_case("retry-after"))
|
||||
.cloned()
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
// 流内拒绝没有原始 HTTP 失败响应,仍需把该次失败的关联 ID 交给 HTTP 客户端。
|
||||
if !headers.iter().any(|header| {
|
||||
header.name().eq_ignore_ascii_case("x-request-id")
|
||||
|| header.name().eq_ignore_ascii_case("x-oai-request-id")
|
||||
}) && let Some(request_id) = error.upstream_request_id()
|
||||
{
|
||||
headers.push(ProviderResponseHeader::new(
|
||||
"x-request-id",
|
||||
Bytes::copy_from_slice(request_id.as_str().as_bytes()),
|
||||
));
|
||||
}
|
||||
let detail = continuation_replay_error_detail();
|
||||
let body = json!({"error": {
|
||||
"message": detail.message(),
|
||||
"code": detail.code(),
|
||||
"type": detail.error_type(),
|
||||
}});
|
||||
// 未交付的 response.created 与额度错误帧必须一起丢弃,否则原错误帧会
|
||||
// 覆盖恢复投影。该拒绝发生在语义输出和 commit 之前,不能带走已交付事件。
|
||||
drop(error.take_atomic_client_events());
|
||||
error
|
||||
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
|
||||
.with_continuation_recovery_disposition(
|
||||
ContinuationRecoveryDisposition::ClientReplayRequired,
|
||||
)
|
||||
.with_continuation_unavailable_reason("quota_exhausted")
|
||||
.with_client_visible_upstream_error(detail)
|
||||
.with_client_visible_upstream_response(
|
||||
ClientVisibleUpstreamResponse::new(
|
||||
reqwest::StatusCode::BAD_REQUEST.as_u16(),
|
||||
Some(b"application/json".to_vec()),
|
||||
Bytes::from(body.to_string()),
|
||||
)
|
||||
.with_headers(headers),
|
||||
)
|
||||
}
|
||||
|
||||
fn continuation_replay_error_detail() -> ClientVisibleUpstreamError {
|
||||
ClientVisibleUpstreamError::new(
|
||||
PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
|
||||
Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
|
||||
Some("invalid_request_error".to_owned()),
|
||||
))
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
|
||||
|
||||
@@ -91,8 +91,8 @@ use crate::transport::{
|
||||
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
|
||||
CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport,
|
||||
CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata,
|
||||
CodexTransportMetrics, CodexTurnStateUpdate, CodexUpstreamDiagnostics, CodexWebSocketPool,
|
||||
endpoint_url,
|
||||
CodexResponseMetadataUpdates, CodexTransportMetrics, CodexUpstreamDiagnostics,
|
||||
CodexWebSocketPool, endpoint_url,
|
||||
};
|
||||
|
||||
mod execution;
|
||||
@@ -361,6 +361,8 @@ impl Provider for CodexProvider {
|
||||
let continuation_requested = generate.native_continuation_requested();
|
||||
let mut upstream_request = encode_generate_request(generate, upstream_model.as_str(), None)
|
||||
.map_err(map_request_error)?;
|
||||
// 编码已生成独立请求;HTTP、WS 与重试在头部和计量之前共用此策略。
|
||||
upstream_request.apply_fast_policy(context.disable_fast());
|
||||
if let Some(conversation_id) = previous_session
|
||||
.as_ref()
|
||||
.and_then(|state| state.conversation_id.as_ref())
|
||||
|
||||
@@ -33,6 +33,7 @@ pub(super) struct OpenAiResponseObservationState {
|
||||
stream: bool,
|
||||
requested_service_tier: Option<String>,
|
||||
upstream_service_tier: Option<String>,
|
||||
upstream_response_model: Option<String>,
|
||||
rate_limit_headers: Vec<(String, String)>,
|
||||
timings: ProviderResponseTimings,
|
||||
terminal: Option<OpenAiResponseTerminal>,
|
||||
@@ -79,6 +80,7 @@ impl OpenAiResponseObservationState {
|
||||
stream: request.stream(),
|
||||
requested_service_tier: normalize_service_tier(request.service_tier()),
|
||||
upstream_service_tier: None,
|
||||
upstream_response_model: response.response_metadata.effective_model.clone(),
|
||||
rate_limit_headers: selected_observation_headers(&response.rate_limit_headers),
|
||||
timings: openai_response_timings(
|
||||
&response.transport_metrics,
|
||||
@@ -121,9 +123,23 @@ impl OpenAiResponseObservationState {
|
||||
if let Some(service_tier) = &self.requested_service_tier {
|
||||
observation = observation.with_service_tier_if_valid(service_tier.clone());
|
||||
}
|
||||
if let Some(model) = &self.upstream_response_model {
|
||||
observation = observation.with_upstream_response_model_if_valid(model);
|
||||
}
|
||||
Some(observation)
|
||||
}
|
||||
|
||||
pub(super) fn observe_upstream_response_model(&mut self, model: Option<&str>) -> bool {
|
||||
let Some(model) = model else {
|
||||
return false;
|
||||
};
|
||||
if self.upstream_response_model.as_deref() == Some(model) {
|
||||
return false;
|
||||
}
|
||||
self.upstream_response_model = Some(model.to_owned());
|
||||
true
|
||||
}
|
||||
|
||||
pub(super) fn observe_stream_chunk(&mut self, chunk: &[u8], started_at: Instant) -> bool {
|
||||
if chunk.is_empty() {
|
||||
return false;
|
||||
|
||||
@@ -13,7 +13,10 @@ use gateway_core::event::{
|
||||
};
|
||||
use gateway_core::metering::Usage;
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_protocol::openai::events::{TokenUsage, billable_usage_is_complete, extract_usage};
|
||||
use gateway_protocol::openai::events::{
|
||||
ResponseModelObservation, TokenUsage, billable_usage_is_complete, extract_usage,
|
||||
observed_model_name,
|
||||
};
|
||||
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder, SseFrame, sse_frame_is_done};
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
@@ -34,7 +37,7 @@ const CONTENTS_PER_OUTPUT: u32 = 1_024;
|
||||
/// 因而未知或形状变化的 JSON event 只能放弃 canonical 投影,不能截断 wire 流。
|
||||
pub struct CodexCanonicalDecoder {
|
||||
decoder: SseEventDecoder,
|
||||
fallback_model: String,
|
||||
upstream_model: String,
|
||||
response_id: Option<String>,
|
||||
started: bool,
|
||||
completed: bool,
|
||||
@@ -46,6 +49,8 @@ pub struct CodexCanonicalDecoder {
|
||||
semantic_output_seen: bool,
|
||||
requested_service_tier: Option<String>,
|
||||
response_service_tier: Option<String>,
|
||||
response_model: ResponseModelObservation,
|
||||
reported_model: Option<String>,
|
||||
web_search_pricing: Option<WebSearchPricing>,
|
||||
timing_signals: ResponseEventSignals,
|
||||
raw_sse_passthrough: bool,
|
||||
@@ -128,10 +133,11 @@ impl CodexCanonicalFailure {
|
||||
}
|
||||
|
||||
impl CodexCanonicalDecoder {
|
||||
pub fn new(fallback_model: impl Into<String>) -> Self {
|
||||
/// 使用路由后最终发往上游的请求模型计价,并在响应缺少模型时用于 canonical 兜底。
|
||||
pub fn new(upstream_model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
decoder: SseEventDecoder::default(),
|
||||
fallback_model: fallback_model.into(),
|
||||
upstream_model: upstream_model.into(),
|
||||
response_id: None,
|
||||
started: false,
|
||||
completed: false,
|
||||
@@ -143,6 +149,8 @@ impl CodexCanonicalDecoder {
|
||||
semantic_output_seen: false,
|
||||
requested_service_tier: None,
|
||||
response_service_tier: None,
|
||||
response_model: ResponseModelObservation::default(),
|
||||
reported_model: None,
|
||||
web_search_pricing: None,
|
||||
timing_signals: ResponseEventSignals::default(),
|
||||
raw_sse_passthrough: false,
|
||||
@@ -223,6 +231,28 @@ impl CodexCanonicalDecoder {
|
||||
self.response_service_tier.as_deref()
|
||||
}
|
||||
|
||||
/// 真实 HTTP 响应头提供初始报告;流内请求级报告可覆盖它。
|
||||
#[must_use]
|
||||
pub fn with_reported_model(mut self, model: Option<&str>) -> Self {
|
||||
self.reported_model = model.and_then(observed_model_name).map(str::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
/// 接收 transport 已解析的内部 metadata 报告,内部帧无需交付客户端。
|
||||
pub(crate) fn observe_reported_model(&mut self, model: &str) {
|
||||
if let Some(model) = observed_model_name(model) {
|
||||
self.reported_model = Some(model.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
/// 官方服务端报告优先;缺少报告时仅使用正文明确声明,不使用请求兜底值。
|
||||
#[must_use]
|
||||
pub fn response_model(&self) -> Option<&str> {
|
||||
self.reported_model
|
||||
.as_deref()
|
||||
.or_else(|| self.response_model.model())
|
||||
}
|
||||
|
||||
fn decode(&mut self, events: Vec<SseEvent>) -> CodexCanonicalOutcome {
|
||||
let mut output = Vec::new();
|
||||
for event in events {
|
||||
@@ -288,6 +318,10 @@ impl CodexCanonicalDecoder {
|
||||
return Ok(());
|
||||
}
|
||||
self.observe_response_service_tier(&value);
|
||||
self.response_model.observe(event_type, &value);
|
||||
if let Some(model) = super::response_meta::reported_model_from_event(&value) {
|
||||
self.reported_model = Some(model.to_owned());
|
||||
}
|
||||
let signals = response_event_signals(event_type, &value);
|
||||
self.merge_timing_signals(signals);
|
||||
if matches!(event_type, Some("response.failed" | "error")) {
|
||||
@@ -452,7 +486,7 @@ impl CodexCanonicalDecoder {
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|model| !model.is_empty())
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.unwrap_or(&self.upstream_model)
|
||||
.to_owned();
|
||||
self.response_id = Some(response_id.clone());
|
||||
self.started = true;
|
||||
@@ -859,9 +893,9 @@ impl CodexCanonicalDecoder {
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|model| !model.is_empty())
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.unwrap_or(&self.upstream_model)
|
||||
.to_owned();
|
||||
// 与用量统计统一按最终发送档位估算,响应回显不改变本地计价口径。
|
||||
// 按最终发送的模型与档位估算,响应回显仅作观测,不改变本地计价口径。
|
||||
let service_tier = self.requested_service_tier.as_deref();
|
||||
let tool_calls = billable_tool_calls(response);
|
||||
if let Some(breakdown) = usage
|
||||
@@ -869,7 +903,7 @@ impl CodexCanonicalDecoder {
|
||||
.and_then(|usage| {
|
||||
let (web_search_calls, file_search_calls) = tool_calls?;
|
||||
openai_billing_breakdown(
|
||||
&model,
|
||||
&self.upstream_model,
|
||||
OpenAiBillingUsage::from(usage)
|
||||
.with_web_search_calls(web_search_calls, self.web_search_pricing)
|
||||
.with_file_search_calls(file_search_calls),
|
||||
|
||||
@@ -33,7 +33,7 @@ use super::response_meta::CodexResponseMetadata;
|
||||
use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key};
|
||||
use super::websocket::{
|
||||
CodexWebSocketExchangeError, CodexWebSocketPool, CodexWebSocketPoolKey,
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketRequest, CodexWebSocketResponseMetadataUpdates,
|
||||
PreparedWebSocket, WebSocketOriginBreaker, WebSocketPoolDecision,
|
||||
};
|
||||
|
||||
@@ -589,8 +589,8 @@ pub struct CodexTransportMetrics {
|
||||
/// 响应头之后在 live 流中采集的结构化限流更新。
|
||||
pub type CodexRateLimitUpdates = CodexWebSocketRateLimitUpdates;
|
||||
|
||||
/// 响应头之后在 live 流中采集的 turn state 更新。
|
||||
pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate;
|
||||
/// 响应头之后在 live 流中采集的请求级 metadata 更新。
|
||||
pub type CodexResponseMetadataUpdates = CodexWebSocketResponseMetadataUpdates;
|
||||
|
||||
/// Codex Responses 上游 live SSE 响应。
|
||||
pub struct CodexBackendStreamingResponse {
|
||||
@@ -608,8 +608,8 @@ pub struct CodexBackendStreamingResponse {
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
/// live stream 期间捕获的结构化限流更新。
|
||||
pub rate_limit_updates: Option<CodexRateLimitUpdates>,
|
||||
/// live stream 期间捕获的 turn-state 更新。
|
||||
pub turn_state_update: Option<CodexTurnStateUpdate>,
|
||||
/// live stream 期间捕获的请求级 metadata 更新。
|
||||
pub response_metadata_updates: Option<CodexResponseMetadataUpdates>,
|
||||
/// WebSocket 连接池决策。
|
||||
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// 上游诊断元数据。
|
||||
|
||||
@@ -197,7 +197,7 @@ impl CodexBackendClient {
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_updates: Some(rate_limit_updates),
|
||||
turn_state_update: None,
|
||||
response_metadata_updates: None,
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
response_metadata,
|
||||
@@ -464,7 +464,7 @@ impl CodexBackendClient {
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
rate_limit_updates: Some(exchange.rate_limit_updates),
|
||||
turn_state_update: Some(exchange.turn_state_update),
|
||||
response_metadata_updates: Some(exchange.response_metadata_updates),
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
response_metadata: exchange.response_metadata,
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
//! 下游客户端发送的非 Codex 模型请求参数兼容。
|
||||
//! 下游客户端模型请求正文的 Codex 协议兼容。
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// 仅移除已确认不适用于 Codex Responses 的顶层参数,不递归清洗业务正文。
|
||||
/// 补齐 Codex 请求缺省字段并移除已确认不适用的顶层参数,不递归清洗业务正文。
|
||||
///
|
||||
/// 兼容基准是 Codex Core/Desktop 的模型请求,不是公开 OpenAI Responses API。
|
||||
/// 未知字段继续透传,不能因官方请求结构中没有某个字段就将其列入过滤规则。
|
||||
pub(in crate::transport) fn strip_non_codex_request_fields(body: &mut Map<String, Value>) {
|
||||
pub(in crate::transport) fn normalize_codex_request_body(body: &mut Map<String, Value>) {
|
||||
// 官方 Core/Desktop 显式发送 store=false;仅为缺字段的下游请求补齐,保留显式值。
|
||||
body.entry("store").or_insert(Value::Bool(false));
|
||||
|
||||
for field in [
|
||||
// Pi 普通 Responses 适配将 maxTokens 映射为 max_output_tokens,
|
||||
// temperature 则原样写入;Pi 的 Codex 适配也可能发送 temperature。
|
||||
|
||||
@@ -4,5 +4,5 @@
|
||||
mod body;
|
||||
mod headers;
|
||||
|
||||
pub(super) use body::strip_non_codex_request_fields;
|
||||
pub(super) use body::normalize_codex_request_body;
|
||||
pub(super) use headers::is_non_codex_request_header;
|
||||
|
||||
@@ -35,7 +35,8 @@ pub use self::{
|
||||
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse,
|
||||
CodexBackendSseStream, CodexBackendStreamingResponse, CodexBackendTransport,
|
||||
CodexClientError, CodexClientResult, CodexRateLimitUpdates, CodexRequestContext,
|
||||
CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
|
||||
CodexResponseMetadataUpdates, CodexTransportDecision, CodexTransportMetrics,
|
||||
build_reqwest_client,
|
||||
},
|
||||
diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase},
|
||||
endpoints::{
|
||||
|
||||
@@ -611,6 +611,21 @@ impl CodexResponsesRequest {
|
||||
self.body.get("service_tier").and_then(Value::as_str)
|
||||
}
|
||||
|
||||
/// 只覆盖顶层 Fast 请求,显式使用官方标准档退出值。
|
||||
pub(crate) fn apply_fast_policy(&mut self, disable_fast: bool) {
|
||||
if disable_fast
|
||||
&& self.service_tier().is_some_and(|tier| {
|
||||
let tier = tier.trim();
|
||||
tier.eq_ignore_ascii_case("priority") || tier.eq_ignore_ascii_case("fast")
|
||||
})
|
||||
{
|
||||
self.body.insert(
|
||||
"service_tier".to_owned(),
|
||||
Value::String("default".to_owned()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 前一个 response ID。
|
||||
pub fn previous_response_id(&self) -> Option<&str> {
|
||||
self.body
|
||||
|
||||
@@ -14,7 +14,7 @@ use serde::Serialize as _;
|
||||
use serde_json::{Map, Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::transport::downstream::{is_non_codex_request_header, strip_non_codex_request_fields};
|
||||
use crate::transport::downstream::{is_non_codex_request_header, normalize_codex_request_body};
|
||||
use crate::transport::headers::is_managed_identity_header;
|
||||
use crate::transport::profile::CodexRequestLocation;
|
||||
use crate::transport::protocol::responses::{
|
||||
@@ -120,7 +120,7 @@ fn adapt_codex_responses_body(
|
||||
location: Option<&CodexRequestLocation>,
|
||||
) {
|
||||
body.insert("model".to_owned(), Value::String(upstream_model.to_owned()));
|
||||
strip_non_codex_request_fields(body);
|
||||
normalize_codex_request_body(body);
|
||||
if let Some(location) = location {
|
||||
align_structured_location_fields(body, Utc::now(), location);
|
||||
}
|
||||
|
||||
@@ -119,13 +119,37 @@ fn observe_typed_response_header(metadata: &mut CodexResponseMetadata, name: &st
|
||||
return;
|
||||
}
|
||||
match name.trim().to_ascii_lowercase().as_str() {
|
||||
"openai-model" => metadata.effective_model = Some(value.to_string()),
|
||||
"openai-model" | "x-openai-model" => {
|
||||
metadata.effective_model =
|
||||
gateway_protocol::openai::events::observed_model_name(value).map(str::to_owned)
|
||||
}
|
||||
"x-models-etag" => metadata.models_etag = Some(value.to_string()),
|
||||
"x-reasoning-included" => metadata.reasoning_included = true,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// 与官方 Codex 一致,优先读取 response.headers,再读取 WS metadata 的顶层 headers。
|
||||
pub(super) fn reported_model_from_event(value: &serde_json::Value) -> Option<&str> {
|
||||
[value.pointer("/response/headers"), value.get("headers")]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(serde_json::Value::as_object)
|
||||
.find_map(|headers| {
|
||||
headers.iter().find_map(|(name, value)| {
|
||||
if !name.eq_ignore_ascii_case("openai-model")
|
||||
&& !name.eq_ignore_ascii_case("x-openai-model")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
value
|
||||
.as_str()
|
||||
.or_else(|| value.as_array()?.first()?.as_str())
|
||||
.and_then(gateway_protocol::openai::events::observed_model_name)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn filter_client_headers(
|
||||
headers: impl IntoIterator<Item = (String, Bytes)>,
|
||||
) -> Vec<(String, Bytes)> {
|
||||
|
||||
@@ -37,8 +37,8 @@ pub struct CodexWebSocketStreamingExchange {
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
/// 上游内部 `codex.rate_limits` 事件里的结构化动态更新。
|
||||
pub rate_limit_updates: CodexWebSocketRateLimitUpdates,
|
||||
/// 上游内部 metadata 事件里的首个动态 turn state。
|
||||
pub turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
/// 上游内部 metadata 事件里的请求级动态更新。
|
||||
pub response_metadata_updates: CodexWebSocketResponseMetadataUpdates,
|
||||
/// WebSocket 连接池决策。
|
||||
pub pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// terminal completed 后该 socket 是否会保留 connection-local continuation。
|
||||
@@ -54,17 +54,31 @@ pub type CodexWebSocketSseStream =
|
||||
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
|
||||
/// live 流中的结构化限流动态更新。
|
||||
pub type CodexWebSocketRateLimitUpdates = Arc<Mutex<Vec<ParsedRateLimits>>>;
|
||||
/// live 流中的 turn state 动态更新。
|
||||
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
|
||||
/// 单次响应的动态 metadata,与连接池保存的握手快照隔离。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct CodexWebSocketResponseMetadataUpdate {
|
||||
/// 当前响应首次声明的会话续接状态。
|
||||
pub turn_state: Option<String>,
|
||||
/// 当前响应最新的服务端模型报告。
|
||||
pub reported_model: Option<String>,
|
||||
}
|
||||
|
||||
pub type CodexWebSocketResponseMetadataUpdates = Arc<Mutex<CodexWebSocketResponseMetadataUpdate>>;
|
||||
|
||||
pub(super) fn reusable_websocket_metadata(
|
||||
mut metadata: CodexWebSocketConnectionMetadata,
|
||||
) -> CodexWebSocketConnectionMetadata {
|
||||
metadata.rate_limit_headers.clear();
|
||||
metadata.turn_state = None;
|
||||
// 模型报告属于上一轮请求,池中连接不能把它带到下一轮缺失报告的响应。
|
||||
metadata.response_metadata.effective_model = None;
|
||||
metadata
|
||||
.response_metadata
|
||||
.client_headers
|
||||
.retain(|(name, _)| !name.eq_ignore_ascii_case("x-codex-turn-state"));
|
||||
.retain(|(name, _)| {
|
||||
!["x-codex-turn-state", "openai-model", "x-openai-model"]
|
||||
.iter()
|
||||
.any(|header| name.eq_ignore_ascii_case(header))
|
||||
});
|
||||
metadata
|
||||
}
|
||||
|
||||
@@ -63,6 +63,9 @@ pub(super) fn reduce_websocket_event(
|
||||
&mut metadata.response_metadata,
|
||||
websocket_metadata_headers(&value),
|
||||
);
|
||||
if let Some(model) = response_meta::reported_model_from_event(&value) {
|
||||
metadata.response_metadata.effective_model = Some(model.to_owned());
|
||||
}
|
||||
let turn_state_update = websocket_metadata_turn_state(&value).and_then(|turn_state| {
|
||||
if metadata.turn_state.is_some() {
|
||||
return None;
|
||||
|
||||
@@ -20,8 +20,9 @@ use super::super::{
|
||||
use super::io::{next_websocket_message, reused_stream_receive_error};
|
||||
use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event};
|
||||
use super::{
|
||||
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange,
|
||||
CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER, reusable_websocket_metadata,
|
||||
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates,
|
||||
CodexWebSocketResponseMetadataUpdate, CodexWebSocketResponseMetadataUpdates,
|
||||
CodexWebSocketStreamingExchange, WEBSOCKET_STREAM_BUFFER, reusable_websocket_metadata,
|
||||
};
|
||||
|
||||
pub(in crate::transport::websocket) struct WebSocketStreamPoolReturn {
|
||||
@@ -69,8 +70,11 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
let response_metadata = metadata.clone();
|
||||
let rate_limit_updates = Arc::new(Mutex::new(Vec::new()));
|
||||
let rate_limit_updates_for_task = Arc::clone(&rate_limit_updates);
|
||||
let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone()));
|
||||
let turn_state_update_for_task = Arc::clone(&turn_state_update);
|
||||
let response_metadata_updates = Arc::new(Mutex::new(CodexWebSocketResponseMetadataUpdate {
|
||||
turn_state: metadata.turn_state.clone(),
|
||||
reported_model: None,
|
||||
}));
|
||||
let response_metadata_updates_for_task = Arc::clone(&response_metadata_updates);
|
||||
let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER);
|
||||
let (task_tracker, shutdown) = pool_return
|
||||
.as_ref()
|
||||
@@ -89,7 +93,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
trace,
|
||||
shutdown,
|
||||
rate_limit_updates: rate_limit_updates_for_task,
|
||||
turn_state_update: turn_state_update_for_task,
|
||||
response_metadata_updates: response_metadata_updates_for_task,
|
||||
tx,
|
||||
})
|
||||
.await;
|
||||
@@ -111,7 +115,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
set_cookie_headers: response_metadata.set_cookie_headers,
|
||||
rate_limit_headers: response_metadata.rate_limit_headers,
|
||||
rate_limit_updates,
|
||||
turn_state_update,
|
||||
response_metadata_updates,
|
||||
pool_decision: None,
|
||||
connection_local_continuation: false,
|
||||
diagnostics: response_metadata.diagnostics,
|
||||
@@ -128,7 +132,7 @@ struct WebSocketStreamForwardState {
|
||||
stream_idle_timeout: Option<Duration>,
|
||||
shutdown: CancellationToken,
|
||||
rate_limit_updates: CodexWebSocketRateLimitUpdates,
|
||||
turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
response_metadata_updates: CodexWebSocketResponseMetadataUpdates,
|
||||
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
|
||||
}
|
||||
|
||||
@@ -142,7 +146,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
stream_idle_timeout,
|
||||
shutdown,
|
||||
rate_limit_updates,
|
||||
turn_state_update,
|
||||
response_metadata_updates,
|
||||
tx,
|
||||
} = state;
|
||||
let mut pool_return = pool_return;
|
||||
@@ -298,11 +302,14 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
last_event_type = Some(event_type);
|
||||
}
|
||||
if let Some(turn_state) = reduced.turn_state_update {
|
||||
let mut pending = turn_state_update.lock().await;
|
||||
if pending.is_none() {
|
||||
*pending = Some(turn_state);
|
||||
let mut pending = response_metadata_updates.lock().await;
|
||||
if pending.turn_state.is_none() {
|
||||
pending.turn_state = Some(turn_state);
|
||||
}
|
||||
}
|
||||
if let Some(model) = metadata.response_metadata.effective_model.as_ref() {
|
||||
response_metadata_updates.lock().await.reported_model = Some(model.clone());
|
||||
}
|
||||
let (frame, terminal) = match reduced.action {
|
||||
ExchangeAction::RateLimits(rate_limits) => {
|
||||
rate_limit_updates.lock().await.push(rate_limits);
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
use gateway_protocol::openai::events;
|
||||
use tokio::time::timeout;
|
||||
use tokio::{io::BufWriter, time::timeout};
|
||||
use tokio_tungstenite::{
|
||||
Connector, MaybeTlsStream, client_async_tls_with_config, connect_async_tls_with_config,
|
||||
};
|
||||
@@ -267,8 +267,12 @@ async fn dial_account(
|
||||
} else {
|
||||
host.to_owned()
|
||||
};
|
||||
tokio_tungstenite::proxy::connect_via_proxy(stream, &config, &target, port)
|
||||
// 依赖分两次写入 SOCKS 方法协商,再显式 flush;部分代理会提前拒绝半包。
|
||||
// 仅在代理握手期间合并写入,不缓冲读取,也不改变后续 TLS/WS 的发送方式。
|
||||
// 这只是兼容措施,不能保证网络层不再拆分 TCP 数据。
|
||||
tokio_tungstenite::proxy::connect_via_proxy(BufWriter::new(stream), &config, &target, port)
|
||||
.await
|
||||
.map(BufWriter::into_inner)
|
||||
.map_err(|_| invalid())
|
||||
}
|
||||
|
||||
|
||||
@@ -21,8 +21,8 @@ pub use self::{
|
||||
},
|
||||
error::{CodexWebSocketCloseError, CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
|
||||
exchange::{
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange,
|
||||
CodexWebSocketTurnStateUpdate,
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketResponseMetadataUpdates,
|
||||
CodexWebSocketSseStream, CodexWebSocketStreamingExchange,
|
||||
},
|
||||
handshake::responses_websocket_endpoint,
|
||||
model::{
|
||||
|
||||
@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use chrono::{TimeZone as _, Utc};
|
||||
use futures::executor::block_on;
|
||||
use gateway_core::account::{ProviderAccount, ProviderAccountId};
|
||||
use gateway_core::account::{ProviderAccount, ProviderAccountId, ProviderAccountStore};
|
||||
use gateway_core::provider_ports::ProviderCatalogCachePort;
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, FrozenAccountScope, ModelServiceTier, ProviderKind, RuntimeAccount,
|
||||
@@ -440,6 +440,39 @@ async fn missing_account_refresh_fails_before_network_io() {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_account_refresh_discovers_models_with_the_pinned_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account = seed_account(&store, "acct_disabled_models").await;
|
||||
store
|
||||
.set_enabled(account.id(), false)
|
||||
.await
|
||||
.expect("disable account");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/codex/models"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "application/json")
|
||||
.set_body_raw(
|
||||
br#"{"models":[{"slug":"gpt-5.4","display_name":"GPT-5.4"}]}"#.to_vec(),
|
||||
"application/json",
|
||||
),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = service_with_catalog_cache(&store, server.uri(), catalog_cache());
|
||||
|
||||
let catalog = service
|
||||
.refresh_account_catalog(account.id())
|
||||
.await
|
||||
.expect("disabled account refresh still discovers models");
|
||||
|
||||
assert_eq!(catalog.models(), ["gpt-5.4"]);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plan_catalog_cache_is_shared_and_manual_refresh_replaces_it() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
@@ -66,6 +66,152 @@ use crate::support::{
|
||||
};
|
||||
use crate::transport::accept_codex_test_websocket;
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_bill_sent_model_and_observe_unpriced_response_without_rewriting_it() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let server = MockServer::start().await;
|
||||
let body = concat!(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_model\",\"model\":\"gpt-created\"}}\n\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_model\",\"model\":\"gpt-6-sol\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(body),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
json!({"model":"gpt-5.6-sol","input":"hello"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
|
||||
));
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", operation),
|
||||
context("req_model", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut observed = None;
|
||||
let mut returned = None;
|
||||
let mut costs = Vec::new();
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("upstream response");
|
||||
for fact in event.canonical_facts() {
|
||||
if let GatewayEvent::CalculatedCost(cost) = fact {
|
||||
costs.push(cost.total().amount().scaled());
|
||||
}
|
||||
}
|
||||
if let Some(observation) = event.response_observation() {
|
||||
observed = observation.upstream_response_model().map(str::to_owned);
|
||||
}
|
||||
if let Some(wire) = event.wire_event()
|
||||
&& wire.data().get("type").and_then(Value::as_str) == Some("response.completed")
|
||||
{
|
||||
returned = wire
|
||||
.data()
|
||||
.pointer("/response/model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
}
|
||||
}
|
||||
let requests = server.received_requests().await.expect("upstream requests");
|
||||
assert_eq!(captured_request_body(&requests[0])["model"], "gpt-5.4");
|
||||
assert_eq!(costs, vec![3_437_500]);
|
||||
assert_eq!(observed.as_deref(), Some("gpt-6-sol"));
|
||||
assert_eq!(returned, observed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_bills_sent_model_and_observes_internal_model_report() {
|
||||
for metadata_type in ["response.metadata", "codex.response.metadata"] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
|
||||
let base_url = format!("http://{}", listener.local_addr().expect("address"));
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.expect("accept WebSocket");
|
||||
let mut websocket = accept_codex_test_websocket(socket).await;
|
||||
let request = websocket.next().await.expect("request").expect("frame");
|
||||
let request: Value =
|
||||
serde_json::from_str(request.to_text().expect("text")).expect("request JSON");
|
||||
assert_eq!(request["model"], "gpt-5.4");
|
||||
for event in [
|
||||
json!({"type":"response.created","response":{"id":"resp_model","model":"gpt-created"}}),
|
||||
json!({"type":metadata_type,"headers":{"X-OpenAI-Model":["gpt-server-report"]}}),
|
||||
json!({"type":"response.completed","response":{"id":"resp_model","model":"gpt-6-astra","status":"completed","output":[],"usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25},"total_tokens":110}}}),
|
||||
] {
|
||||
websocket
|
||||
.send(Message::Text(event.to_string().into()))
|
||||
.await
|
||||
.expect("response");
|
||||
}
|
||||
});
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.6-sol")),
|
||||
("input".to_owned(), json!("hello")),
|
||||
]),
|
||||
)
|
||||
.expect("payload")
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(true))]));
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
|
||||
),
|
||||
context("req_ws_model", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("provider stream");
|
||||
let mut observed = None;
|
||||
let mut returned = None;
|
||||
let mut costs = Vec::new();
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("provider event");
|
||||
for fact in event.canonical_facts() {
|
||||
if let GatewayEvent::CalculatedCost(cost) = fact {
|
||||
costs.push(cost.total().amount().scaled());
|
||||
}
|
||||
}
|
||||
if let Some(observation) = event.response_observation() {
|
||||
observed = observation.upstream_response_model().map(str::to_owned);
|
||||
}
|
||||
if let Some(wire) = event.wire_event()
|
||||
&& wire.data().get("type").and_then(Value::as_str) == Some("response.completed")
|
||||
{
|
||||
returned = wire
|
||||
.data()
|
||||
.pointer("/response/model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
}
|
||||
}
|
||||
server.await.expect("server task");
|
||||
assert_eq!(costs, vec![3_437_500]);
|
||||
assert_eq!(
|
||||
observed.as_deref(),
|
||||
Some("gpt-server-report"),
|
||||
"{metadata_type}"
|
||||
);
|
||||
assert_eq!(returned.as_deref(), Some("gpt-6-astra"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn selected_proxy_location_overrides_global_and_reloads_without_mutating_client_payload() {
|
||||
use gateway_core::account::{OutboundProxy, RequestLocation};
|
||||
@@ -573,11 +719,20 @@ fn global_request_location() -> gateway_core::account::RequestLocation {
|
||||
}
|
||||
|
||||
fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext {
|
||||
context_with_fast_policy(request_id, cancellation, false)
|
||||
}
|
||||
|
||||
fn context_with_fast_policy(
|
||||
request_id: &str,
|
||||
cancellation: CancellationToken,
|
||||
disable_fast: bool,
|
||||
) -> AttemptContext {
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new(request_id).expect("request id"),
|
||||
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
|
||||
)
|
||||
.with_disable_fast(disable_fast)
|
||||
.with_request_location(Some(global_request_location())),
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
@@ -5179,11 +5334,47 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
|
||||
#[tokio::test]
|
||||
async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_transports() {
|
||||
for use_websocket in [false, true] {
|
||||
for (requested, reported, expected_cost) in [
|
||||
(Some("priority"), Some("default"), 6_875_000),
|
||||
(Some("priority"), None, 6_875_000),
|
||||
(Some("default"), Some("priority"), 3_437_500),
|
||||
(None, Some("priority"), 3_437_500),
|
||||
for (disable_fast, requested, reported, expected_tier, expected_cost) in [
|
||||
(
|
||||
false,
|
||||
Some("priority"),
|
||||
Some("default"),
|
||||
Some("priority"),
|
||||
Some(6_875_000),
|
||||
),
|
||||
(
|
||||
false,
|
||||
Some("priority"),
|
||||
None,
|
||||
Some("priority"),
|
||||
Some(6_875_000),
|
||||
),
|
||||
(
|
||||
false,
|
||||
Some("default"),
|
||||
Some("priority"),
|
||||
Some("default"),
|
||||
Some(3_437_500),
|
||||
),
|
||||
(false, None, Some("priority"), None, Some(3_437_500)),
|
||||
(
|
||||
true,
|
||||
Some("priority"),
|
||||
Some("priority"),
|
||||
Some("default"),
|
||||
Some(3_437_500),
|
||||
),
|
||||
(true, Some("fast"), None, Some("default"), Some(3_437_500)),
|
||||
(
|
||||
true,
|
||||
Some("default"),
|
||||
None,
|
||||
Some("default"),
|
||||
Some(3_437_500),
|
||||
),
|
||||
(true, None, None, None, Some(3_437_500)),
|
||||
(true, Some("flex"), None, Some("flex"), Some(1_720_000)),
|
||||
(true, Some("ultrafast"), None, Some("ultrafast"), None),
|
||||
] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
@@ -5208,7 +5399,21 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
let base_url = format!("http://{}", listener.local_addr().expect("address"));
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.expect("accept WebSocket");
|
||||
let mut websocket = accept_codex_test_websocket(socket).await;
|
||||
let mut websocket = crate::transport::accept_codex_test_websocket_with(
|
||||
socket,
|
||||
|request, response| {
|
||||
let expected_hint = expected_tier.map_or_else(
|
||||
|| "model=gpt-5.4".to_owned(),
|
||||
|tier| format!("model=gpt-5.4;tier={tier}"),
|
||||
);
|
||||
assert_eq!(request.headers()["x-codex-routing-hint"], expected_hint);
|
||||
response.headers_mut().insert(
|
||||
"sec-websocket-extensions",
|
||||
"permessage-deflate".parse().unwrap(),
|
||||
);
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let request = websocket.next().await.expect("request").expect("frame");
|
||||
let request: Value = serde_json::from_str(request.to_text().expect("text"))
|
||||
.expect("request JSON");
|
||||
@@ -5237,23 +5442,30 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
let mut body = Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("hello")),
|
||||
(
|
||||
"metadata".to_owned(),
|
||||
json!({"service_tier":"priority","text":"fast priority"}),
|
||||
),
|
||||
]);
|
||||
if let Some(requested) = requested {
|
||||
body.insert("service_tier".to_owned(), json!(requested));
|
||||
}
|
||||
let payload = ProtocolPayload::json_object("openai", body)
|
||||
let payload = ProtocolPayload::json_object("openai", body.clone())
|
||||
.expect("payload")
|
||||
.with_context(Map::from_iter([(
|
||||
"use_websocket".to_owned(),
|
||||
json!(use_websocket),
|
||||
)]));
|
||||
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
|
||||
let original = operation.clone();
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
|
||||
planned_request("openai", operation.clone()),
|
||||
context_with_fast_policy(
|
||||
"req_service_tier",
|
||||
CancellationToken::new(),
|
||||
disable_fast,
|
||||
),
|
||||
context("req_service_tier", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("provider stream");
|
||||
@@ -5277,6 +5489,11 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
let outbound = if let Some(server) = http_server {
|
||||
server.verify().await;
|
||||
let requests = server.received_requests().await.expect("upstream requests");
|
||||
let expected_hint = expected_tier.map_or_else(
|
||||
|| "model=gpt-5.4".to_owned(),
|
||||
|tier| format!("model=gpt-5.4;tier={tier}"),
|
||||
);
|
||||
assert_eq!(requests[0].headers["x-codex-routing-hint"], expected_hint);
|
||||
captured_request_body(&requests[0])
|
||||
} else {
|
||||
websocket_server
|
||||
@@ -5286,13 +5503,15 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
};
|
||||
assert_eq!(
|
||||
outbound.get("service_tier").and_then(Value::as_str),
|
||||
requested
|
||||
expected_tier
|
||||
);
|
||||
assert_eq!(operation, original);
|
||||
assert_eq!(outbound["metadata"], body["metadata"]);
|
||||
assert!(!observations.is_empty());
|
||||
for observation in &observations {
|
||||
assert_eq!(
|
||||
observation.service_tier(),
|
||||
requested,
|
||||
expected_tier,
|
||||
"WebSocket={use_websocket}"
|
||||
);
|
||||
}
|
||||
@@ -5307,7 +5526,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
.expect("metadata JSON");
|
||||
assert_eq!(
|
||||
metadata.get("requestedServiceTier").and_then(Value::as_str),
|
||||
requested
|
||||
expected_tier
|
||||
);
|
||||
assert_eq!(
|
||||
metadata.get("upstreamServiceTier").and_then(Value::as_str),
|
||||
@@ -5315,7 +5534,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
|
||||
);
|
||||
assert_eq!(
|
||||
costs,
|
||||
vec![expected_cost],
|
||||
expected_cost.into_iter().collect::<Vec<_>>(),
|
||||
"WebSocket={use_websocket}, requested={requested:?}, reported={reported:?}"
|
||||
);
|
||||
assert_eq!(raw_response, response_frames.as_bytes());
|
||||
@@ -8366,3 +8585,595 @@ async fn api_key_websocket_uses_api_path_and_bearer_without_oauth_identity() {
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn quota_continuation_operation(use_websocket: bool) -> Operation {
|
||||
Operation::Generate(
|
||||
GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object("openai", json!({
|
||||
"model": "gpt-5.4", "input": "continue", "previous_response_id": "resp_previous",
|
||||
"session_id": "quota-replay", "thread_id": "turn",
|
||||
}).as_object().unwrap().clone())
|
||||
.unwrap()
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(use_websocket))])),
|
||||
)
|
||||
.with_provider_session_state(
|
||||
generate_with_persisted_session_context("acct_provider_contract", "conversation-quota", "quota-replay", "turn")
|
||||
.provider_session_state("openai").unwrap().clone(),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_continuation_opening_rejection_projects_replay_without_losing_upstream_facts() {
|
||||
for use_websocket in [false, true] {
|
||||
for (status, code, projects_replay) in [
|
||||
(429, "usage_limit_reached", true),
|
||||
(402, "insufficient_quota", true),
|
||||
(429, "quota_exceeded", false),
|
||||
(429, "rate_limit_exceeded", false),
|
||||
(429, "slow_down", false),
|
||||
(429, "unknown_error", false),
|
||||
] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let server = MockServer::start().await;
|
||||
let original =
|
||||
json!({"error": {"type": code, "code": code, "message": "upstream rejection"}});
|
||||
Mock::given(method(if use_websocket { "GET" } else { "POST" }))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(status)
|
||||
.insert_header("retry-after", "129600")
|
||||
.insert_header("x-request-id", "req-upstream-quota")
|
||||
.set_body_json(&original),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mut stream = provider_with_base_url_and_retry_budget(&store, server.uri(), 0)
|
||||
.execute(
|
||||
planned_request("openai", quota_continuation_operation(use_websocket)),
|
||||
context("req_quota_replay", CancellationToken::new())
|
||||
.with_continuation_attempt(ContinuationAttempt::Native),
|
||||
)
|
||||
.await
|
||||
.expect("prepare continuation");
|
||||
let mut error = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(event)) => assert!(!event.has_client_event()),
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("expected rejection"),
|
||||
}
|
||||
};
|
||||
assert_eq!(error.upstream_status(), Some(status));
|
||||
assert_eq!(error.retry_after(), Some(Duration::from_secs(129_600)));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(error.raw_upstream_error().unwrap().as_str())
|
||||
.unwrap(),
|
||||
original
|
||||
);
|
||||
assert!(error.take_atomic_client_events().is_empty());
|
||||
let response = error.client_visible_upstream_response().unwrap();
|
||||
let client: Value = serde_json::from_slice(response.body()).unwrap();
|
||||
if projects_replay {
|
||||
assert_eq!(error.upstream_code().map(|code| code.as_str()), Some(code));
|
||||
assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted);
|
||||
assert_eq!(
|
||||
error.continuation_recovery_disposition(),
|
||||
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
|
||||
);
|
||||
assert_eq!(response.status(), 400);
|
||||
assert_eq!(client["error"]["code"], "previous_response_not_found");
|
||||
assert!(
|
||||
!response
|
||||
.headers()
|
||||
.iter()
|
||||
.any(|header| header.name() == "retry-after")
|
||||
);
|
||||
assert!(
|
||||
response
|
||||
.headers()
|
||||
.iter()
|
||||
.any(|header| header.name() == "x-request-id"
|
||||
&& header.value().as_ref() == b"req-upstream-quota")
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_provider_contract")
|
||||
.unwrap()
|
||||
.quota()
|
||||
.access(),
|
||||
QuotaAccessState::Exhausted
|
||||
);
|
||||
} else {
|
||||
assert_ne!(
|
||||
error.continuation_recovery_disposition(),
|
||||
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
|
||||
);
|
||||
assert_eq!(response.status(), 429);
|
||||
assert_eq!(client, original);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_continuation_stream_rejection_only_projects_before_delivery() {
|
||||
for use_websocket in [false, true] {
|
||||
for semantic_output in [false, true] {
|
||||
for native_continuation in [false, true] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let mut events = vec![
|
||||
json!({"type":"response.created","response":{"id":"resp_quota","model":"gpt-5.4"}}),
|
||||
];
|
||||
if semantic_output {
|
||||
events.push(json!({"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}));
|
||||
}
|
||||
let original = json!({"type":"response.failed","response":{"id":"resp_quota","error":{"code":"usage_limit_reached","message":"You have reached your usage limit."}}});
|
||||
events.push(original.clone());
|
||||
let (base_url, _http_server, websocket_server) = if use_websocket {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut ws = accept_codex_test_websocket(socket).await;
|
||||
ws.next().await.unwrap().unwrap();
|
||||
for event in events {
|
||||
ws.send(Message::Text(event.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
(base_url, None, Some(server))
|
||||
} else {
|
||||
let server = MockServer::start().await;
|
||||
let frames = events
|
||||
.iter()
|
||||
.map(|event| {
|
||||
format!(
|
||||
"event: {}\ndata: {event}\n\n",
|
||||
event["type"].as_str().unwrap()
|
||||
)
|
||||
})
|
||||
.collect::<String>();
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("x-request-id", "req-quota-stream")
|
||||
.set_body_raw(frames, "text/event-stream"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
(server.uri(), Some(server), None)
|
||||
};
|
||||
let operation = if native_continuation {
|
||||
quota_continuation_operation(use_websocket)
|
||||
} else if use_websocket {
|
||||
generate_operation()
|
||||
} else {
|
||||
http_generate_operation()
|
||||
};
|
||||
let attempt = context("req_quota_stream", CancellationToken::new())
|
||||
.with_continuation_attempt(if native_continuation {
|
||||
ContinuationAttempt::Native
|
||||
} else {
|
||||
ContinuationAttempt::None
|
||||
});
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(planned_request("openai", operation), attempt)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut client_events = Vec::new();
|
||||
let mut error = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(event)) => {
|
||||
if event.has_client_event() {
|
||||
client_events.push(event);
|
||||
}
|
||||
}
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("expected quota failure"),
|
||||
}
|
||||
};
|
||||
assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted);
|
||||
assert_eq!(error.send_state(), UpstreamSendState::Sent);
|
||||
assert_eq!(error.replay_is_safe(), !semantic_output);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(error.raw_upstream_error().unwrap().as_str())
|
||||
.unwrap(),
|
||||
original
|
||||
);
|
||||
client_events.extend(error.take_atomic_client_events());
|
||||
if native_continuation && !semantic_output {
|
||||
assert!(
|
||||
client_events.is_empty(),
|
||||
"raw quota error must not override replay projection"
|
||||
);
|
||||
assert_eq!(
|
||||
error.continuation_recovery_disposition(),
|
||||
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
|
||||
);
|
||||
assert_eq!(
|
||||
error.client_visible_upstream_error().unwrap().code(),
|
||||
Some("previous_response_not_found")
|
||||
);
|
||||
if !use_websocket {
|
||||
assert_eq!(
|
||||
error.upstream_request_id().unwrap().as_str(),
|
||||
"req-quota-stream"
|
||||
);
|
||||
assert!(
|
||||
error
|
||||
.client_visible_upstream_response()
|
||||
.unwrap()
|
||||
.headers()
|
||||
.iter()
|
||||
.any(|header| header.name() == "x-request-id"
|
||||
&& header.value().as_ref() == b"req-quota-stream")
|
||||
);
|
||||
}
|
||||
} else {
|
||||
assert_ne!(
|
||||
error.continuation_recovery_disposition(),
|
||||
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
|
||||
);
|
||||
assert_eq!(
|
||||
client_events.last().unwrap().wire_event().unwrap().data(),
|
||||
&original
|
||||
);
|
||||
}
|
||||
if let Some(server) = websocket_server {
|
||||
server.await.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_continuation_after_structural_commit_preserves_original_failure() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let original = json!({"type":"response.failed","response":{"id":"resp_quota","error":{"code":"usage_limit_reached","message":"limit reached"}}});
|
||||
let (base_url, release, _, server) = paused_chunked_sse_server(
|
||||
"event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_quota\",\"model\":\"gpt-5.4\"}}\n\n".to_owned(),
|
||||
format!("event: response.failed\ndata: {original}\n\n"),
|
||||
).await;
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(
|
||||
planned_request("openai", quota_continuation_operation(false)),
|
||||
context("req_quota_commit", CancellationToken::new())
|
||||
.with_continuation_attempt(ContinuationAttempt::Native),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
timeout(Duration::from_secs(3), async {
|
||||
while let Some(event) = stream.next().await {
|
||||
if event.unwrap().has_client_event() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
panic!("expected grace commit");
|
||||
})
|
||||
.await
|
||||
.expect("bounded grace");
|
||||
release.send(()).unwrap();
|
||||
let mut delivered_failure = false;
|
||||
let error = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(event)) => {
|
||||
delivered_failure |= event
|
||||
.wire_event()
|
||||
.is_some_and(|wire| wire.data() == &original);
|
||||
}
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("expected failure"),
|
||||
}
|
||||
};
|
||||
assert!(delivered_failure);
|
||||
assert!(!error.replay_is_safe());
|
||||
assert_ne!(
|
||||
error.continuation_recovery_disposition(),
|
||||
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
|
||||
);
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_continuation_full_client_replay_selects_another_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let full_input = json!([
|
||||
{"role":"user","content":"hello"},
|
||||
{"role":"assistant","content":"first answer"},
|
||||
{"role":"user","content":"continue"},
|
||||
]);
|
||||
let expected_input = full_input.clone();
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut ws = accept_codex_test_websocket(socket).await;
|
||||
ws.next().await.unwrap().unwrap();
|
||||
for event in [
|
||||
json!({"type":"response.created","response":{"id":"resp_previous","model":"gpt-5.4"}}),
|
||||
json!({"type":"response.completed","response":{"id":"resp_previous","model":"gpt-5.4","status":"completed","output":[]}}),
|
||||
] {
|
||||
ws.send(Message::Text(event.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let delta: Value =
|
||||
serde_json::from_str(ws.next().await.unwrap().unwrap().to_text().unwrap()).unwrap();
|
||||
assert_eq!(delta["previous_response_id"], "resp_previous");
|
||||
assert_eq!(delta["input"], "continue");
|
||||
ws.send(Message::Text(json!({"type":"error","status":429,"error":{"type":"usage_limit_reached","code":"usage_limit_reached","message":"You have reached your usage limit."}}).to_string().into())).await.unwrap();
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut ws = accept_codex_test_websocket(socket).await;
|
||||
let replay: Value =
|
||||
serde_json::from_str(ws.next().await.unwrap().unwrap().to_text().unwrap()).unwrap();
|
||||
assert!(replay.get("previous_response_id").is_none());
|
||||
assert_eq!(replay["input"], expected_input);
|
||||
for event in [
|
||||
json!({"type":"response.created","response":{"id":"resp_replayed","model":"gpt-5.4"}}),
|
||||
json!({"type":"response.completed","response":{"id":"resp_replayed","model":"gpt-5.4","status":"completed","output":[]}}),
|
||||
] {
|
||||
ws.send(Message::Text(event.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
let provider = provider_with_base_url(&store, base_url);
|
||||
let first = Operation::Generate(generate_with_persisted_session_context(
|
||||
"acct_provider_contract",
|
||||
"conversation-quota",
|
||||
"quota-replay",
|
||||
"turn",
|
||||
));
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", first),
|
||||
context("req_quota_first", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
stream.metadata().provider_account_id().as_str(),
|
||||
"acct_provider_contract"
|
||||
);
|
||||
while let Some(event) = stream.next().await {
|
||||
event.unwrap();
|
||||
}
|
||||
drop(stream);
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", quota_continuation_operation(true)),
|
||||
context("req_quota_delta", CancellationToken::new())
|
||||
.with_continuation_attempt(ContinuationAttempt::Native),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut error = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(event)) => assert!(!event.has_client_event()),
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("expected quota rejection"),
|
||||
}
|
||||
};
|
||||
drop(stream);
|
||||
assert_eq!(
|
||||
error.client_visible_upstream_error().unwrap().code(),
|
||||
Some("previous_response_not_found")
|
||||
);
|
||||
assert!(error.take_atomic_client_events().is_empty());
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_provider_contract")
|
||||
.unwrap()
|
||||
.quota()
|
||||
.access(),
|
||||
QuotaAccessState::Exhausted
|
||||
);
|
||||
create_account(&store, "acct_affinity_switch_b").await;
|
||||
// 客户端重建完整历史,保留同一会话标识;选号必须跳过刚刚耗尽的原账号。
|
||||
let replay = Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object("openai", json!({"model":"gpt-5.4","session_id":"quota-replay","thread_id":"turn","input":full_input}).as_object().unwrap().clone()).unwrap(),
|
||||
));
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", replay),
|
||||
context("req_quota_full", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
stream.metadata().provider_account_id().as_str(),
|
||||
"acct_affinity_switch_b"
|
||||
);
|
||||
let mut completed = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
completed |= event
|
||||
.unwrap()
|
||||
.canonical_facts()
|
||||
.iter()
|
||||
.any(|fact| matches!(fact, GatewayEvent::Completed(_)));
|
||||
}
|
||||
assert!(completed);
|
||||
timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fast_policy_changes_preserve_the_websocket_continuation_and_meter_each_outbound_tier() {
|
||||
const ACCOUNT: &str = "acct_provider_contract";
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, ACCOUNT).await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut websocket =
|
||||
crate::transport::accept_codex_test_websocket_with(socket, |request, response| {
|
||||
// 提示头只在首次握手发送;后续档位由各自 response.create 正文指定。
|
||||
assert_eq!(
|
||||
request.headers()["x-codex-routing-hint"],
|
||||
"model=gpt-5.4;tier=priority"
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
"sec-websocket-extensions",
|
||||
"permessage-deflate".parse().unwrap(),
|
||||
);
|
||||
})
|
||||
.await;
|
||||
for (index, tier) in ["priority", "default", "priority"].into_iter().enumerate() {
|
||||
let message = websocket.next().await.unwrap().unwrap();
|
||||
let frame: Value = serde_json::from_str(message.to_text().unwrap()).unwrap();
|
||||
assert_eq!(frame["type"], "response.create");
|
||||
assert_eq!(frame["service_tier"], tier);
|
||||
assert_eq!(
|
||||
frame.get("previous_response_id").cloned(),
|
||||
index
|
||||
.checked_sub(1)
|
||||
.map(|previous| json!(format!("resp_fast_turn_{previous}")))
|
||||
);
|
||||
websocket.send(Message::Text(json!({
|
||||
"type":"response.created", "response":{"id":format!("resp_fast_turn_{index}"),"model":"gpt-5.4"}
|
||||
}).to_string().into())).await.unwrap();
|
||||
websocket.send(Message::Text(json!({
|
||||
"type":"response.completed", "response":{
|
||||
"id":format!("resp_fast_turn_{index}"), "model":"gpt-5.4", "status":"completed", "output":[],
|
||||
"service_tier":"priority", "usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25,"cache_write_tokens":0},"total_tokens":110}
|
||||
}
|
||||
}).to_string().into())).await.unwrap();
|
||||
}
|
||||
assert!(
|
||||
timeout(Duration::from_millis(100), listener.accept())
|
||||
.await
|
||||
.is_err(),
|
||||
"all three turns must use the same upstream connection"
|
||||
);
|
||||
});
|
||||
let provider = provider_with_base_url(&store, base_url);
|
||||
let mut session_state = ProviderSessionState::new(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("account_id".to_owned(), json!(ACCOUNT)),
|
||||
(
|
||||
"conversation_id".to_owned(),
|
||||
json!("fast-policy-continuation"),
|
||||
),
|
||||
("continuation_scope".to_owned(), json!("connection_local")),
|
||||
]),
|
||||
)
|
||||
.unwrap();
|
||||
for (index, disable_fast) in [false, true, false].into_iter().enumerate() {
|
||||
let previous = index
|
||||
.checked_sub(1)
|
||||
.map(|index| format!("resp_fast_turn_{index}"));
|
||||
let mut body =
|
||||
json!({"model":"gpt-5.4","input":"next turn","service_tier":"priority","store":false});
|
||||
if let Some(previous) = &previous {
|
||||
body["previous_response_id"] = json!(previous);
|
||||
}
|
||||
let operation = Operation::Generate(
|
||||
GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object("openai", body.as_object().unwrap().clone())
|
||||
.unwrap()
|
||||
// 与客户端 WebSocket 一致,首轮也禁止按快路径预算降级到 HTTP。
|
||||
.with_context(Map::from_iter([
|
||||
("use_websocket".to_owned(), json!(true)),
|
||||
(
|
||||
"downstream_websocket_connection_id".to_owned(),
|
||||
json!("ws_fast_policy_continuation"),
|
||||
),
|
||||
])),
|
||||
)
|
||||
.with_provider_session_state(session_state.clone()),
|
||||
);
|
||||
let account = ProviderAccountId::new(ACCOUNT).unwrap();
|
||||
let provider_kind = ProviderKind::new("openai").unwrap();
|
||||
let key = ClientApiKeyId::new("key_openai_contract").unwrap();
|
||||
let binding = previous.as_ref().map(|previous| {
|
||||
ContinuationBinding::Pinned(NativeContinuationPin::new(
|
||||
PreviousResponseId::new(previous),
|
||||
PreviousResponseId::new(previous),
|
||||
key.clone(),
|
||||
provider_kind.clone(),
|
||||
account.clone(),
|
||||
))
|
||||
});
|
||||
let context = AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new(format!("req_fast_turn_{index}")).unwrap(),
|
||||
key,
|
||||
)
|
||||
.with_disable_fast(disable_fast),
|
||||
NonZeroU32::new(1).unwrap(),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(
|
||||
BTreeSet::new(),
|
||||
None,
|
||||
Some(ProviderAccountStateOwner::new(provider_kind, account)),
|
||||
)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
binding,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.with_continuation_attempt(if previous.is_some() {
|
||||
ContinuationAttempt::Native
|
||||
} else {
|
||||
ContinuationAttempt::None
|
||||
});
|
||||
let mut stream = provider
|
||||
.execute(planned_request("openai", operation), context)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut costs = Vec::new();
|
||||
let mut pool = None;
|
||||
let mut observed_tier = None;
|
||||
while let Some(event) = timeout(Duration::from_secs(5), stream.next())
|
||||
.await
|
||||
.unwrap()
|
||||
{
|
||||
let event = event.unwrap();
|
||||
if let Some(update) = event.session_update() {
|
||||
session_state = update.clone();
|
||||
}
|
||||
if let Some(observation) = event.response_observation() {
|
||||
pool = observation.websocket_pool().or(pool);
|
||||
observed_tier = observation
|
||||
.service_tier()
|
||||
.map(str::to_owned)
|
||||
.or(observed_tier);
|
||||
}
|
||||
for fact in event.canonical_facts() {
|
||||
if let GatewayEvent::CalculatedCost(cost) = fact {
|
||||
costs.push(cost.total().amount().scaled());
|
||||
}
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
pool,
|
||||
Some(if index == 0 {
|
||||
WebSocketPoolKind::New
|
||||
} else {
|
||||
WebSocketPoolKind::Reuse
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
observed_tier.as_deref(),
|
||||
Some(if disable_fast { "default" } else { "priority" })
|
||||
);
|
||||
assert_eq!(
|
||||
costs,
|
||||
vec![if disable_fast { 3_437_500 } else { 6_875_000 }]
|
||||
);
|
||||
}
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
@@ -3,6 +3,12 @@ use gateway_core::{
|
||||
account::{CredentialRevision, OutboundProxy, ProviderAccount, ProviderAccountId},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
use std::{
|
||||
io::{self, Cursor, Read},
|
||||
pin::Pin,
|
||||
task::{Context, Poll},
|
||||
};
|
||||
use tokio::io::{AsyncRead, AsyncWrite, BufWriter, ReadBuf};
|
||||
|
||||
fn account(id: &str, proxy: Option<&str>) -> ProviderAccount {
|
||||
ProviderAccount::new(
|
||||
@@ -334,3 +340,115 @@ async fn https_account_proxy_starts_tls_and_rejection_never_falls_back_to_direct
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 以底层写入调用为边界模拟问题代理,避免将 TCP 单次读取误当成报文边界。
|
||||
struct GreetingSensitiveProxy {
|
||||
greeting: Vec<u8>,
|
||||
replies: Cursor<Vec<u8>>,
|
||||
writes: Vec<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl GreetingSensitiveProxy {
|
||||
fn new(authenticated: bool) -> Self {
|
||||
let (greeting, mut replies) = if authenticated {
|
||||
(vec![5, 2, 0, 2], vec![5, 2, 1, 0])
|
||||
} else {
|
||||
(vec![5, 1, 0], vec![5, 0])
|
||||
};
|
||||
replies.extend_from_slice(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 80]);
|
||||
Self {
|
||||
greeting,
|
||||
replies: Cursor::new(replies),
|
||||
writes: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncRead for GreetingSensitiveProxy {
|
||||
fn poll_read(
|
||||
mut self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
buf: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
if self.writes.first() != Some(&self.greeting) {
|
||||
return Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"proxy rejected partial method negotiation",
|
||||
)));
|
||||
}
|
||||
// 回复逐字节交付,验证客户端不依赖响应的读取边界。
|
||||
if buf.remaining() > 0 {
|
||||
let mut byte = [0];
|
||||
let count = Read::read(&mut self.replies, &mut byte)?;
|
||||
buf.put_slice(&byte[..count]);
|
||||
}
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for GreetingSensitiveProxy {
|
||||
fn poll_write(
|
||||
mut self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
self.writes.push(buf.to_vec());
|
||||
Poll::Ready(Ok(buf.len()))
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn socks_proxy_flush_contract_coalesces_greeting_without_losing_auth_or_connect() {
|
||||
use tokio_tungstenite::proxy::connect_via_proxy;
|
||||
use tungstenite::proxy::ProxyConfig;
|
||||
|
||||
// 直接验证锁定依赖与握手写缓冲的合同;生产接线由已有账号代理集成测试覆盖。
|
||||
for scheme in ["socks5", "socks5h"] {
|
||||
for authenticated in [false, true] {
|
||||
let authentication = if authenticated { "user:pass@" } else { "" };
|
||||
let config =
|
||||
ProxyConfig::parse(&format!("{scheme}://{authentication}127.0.0.1:1080")).unwrap();
|
||||
let raw = connect_via_proxy(
|
||||
GreetingSensitiveProxy::new(authenticated),
|
||||
&config,
|
||||
"upstream.invalid",
|
||||
443,
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(raw, Err(tungstenite::Error::Io(ref error))
|
||||
if error.kind() == io::ErrorKind::UnexpectedEof));
|
||||
|
||||
let buffered = connect_via_proxy(
|
||||
BufWriter::new(GreetingSensitiveProxy::new(authenticated)),
|
||||
&config,
|
||||
"upstream.invalid",
|
||||
443,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
buffered.buffer().is_empty(),
|
||||
"handshake must flush before unwrapping"
|
||||
);
|
||||
let stream = buffered.into_inner();
|
||||
let mut expected = vec![stream.greeting];
|
||||
if authenticated {
|
||||
expected.push(b"\x01\x04user\x04pass".to_vec());
|
||||
}
|
||||
expected.push(b"\x05\x01\x00\x03\x10upstream.invalid\x01\xbb".to_vec());
|
||||
assert_eq!(stream.writes, expected);
|
||||
assert_eq!(
|
||||
stream.replies.position(),
|
||||
stream.replies.get_ref().len() as u64
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,56 @@ use provider_openai::transport::canonical::{
|
||||
use provider_openai::transport::protocol::websocket::websocket_event_to_sse_frame;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn response_model_observation_prefers_reported_headers_without_changing_wire() {
|
||||
for websocket in [false, true] {
|
||||
let mut decoder = CodexCanonicalDecoder::new("requested")
|
||||
.with_reported_model(Some("opening-model"))
|
||||
.with_raw_sse_passthrough();
|
||||
for (value, expected) in [
|
||||
(
|
||||
json!({"type":"response.created","response":{"id":"resp_model","model":"body-model"}}),
|
||||
"opening-model",
|
||||
),
|
||||
(
|
||||
json!({"type":"codex.response.metadata","headers":{"X-OpenAI-Model":["metadata-model"]}}),
|
||||
"metadata-model",
|
||||
),
|
||||
(
|
||||
json!({"type":"response.completed","headers":{"openai-model":"top-level"},"response":{"id":"resp_model","model":"final-body-model","headers":{"OpenAI-Model":"terminal-model"}}}),
|
||||
"terminal-model",
|
||||
),
|
||||
] {
|
||||
let raw = value.to_string();
|
||||
let frame = if websocket {
|
||||
websocket_event_to_sse_frame(&raw).expect("WS event")
|
||||
} else {
|
||||
format!("data: {raw}\n\n")
|
||||
};
|
||||
let events = decoder
|
||||
.push(frame.as_bytes())
|
||||
.expect("decode model observation");
|
||||
assert_eq!(decoder.response_model(), Some(expected));
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|event| event.wire_event().is_some_and(|wire| wire.data() == &value))
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_model_observation_uses_explicit_body_and_never_request_fallback() {
|
||||
let mut decoder = CodexCanonicalDecoder::new("requested");
|
||||
decoder
|
||||
.push(b"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_model\"}}\n\n")
|
||||
.expect("response model event");
|
||||
assert_eq!(decoder.response_model(), None);
|
||||
decoder.push(b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_model\",\"model\":\"returned\"}}\n\n").expect("response model event");
|
||||
assert_eq!(decoder.response_model(), Some("returned"));
|
||||
}
|
||||
|
||||
const METADATA_PREFIX_FIXTURE: &str = include_str!("fixtures/metadata_only_prefix.sse");
|
||||
|
||||
#[test]
|
||||
@@ -199,6 +249,48 @@ fn decoder_should_preserve_long_event_metadata_without_raw_frame() {
|
||||
assert_eq!(wire.data()["marker"], serde_json::json!(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_bill_the_sent_model_independently_of_the_response_model() {
|
||||
for (sent, returned, expected_cost) in [
|
||||
("gpt-5.6-sol", Some("gpt-6-sol"), Some(6_875_000)),
|
||||
("gpt-5.6-sol", Some("gpt-6-astra"), Some(6_875_000)),
|
||||
("gpt-5.6-sol", Some("gpt-5.6-sol"), Some(6_875_000)),
|
||||
("gpt-5.6-sol", None, Some(6_875_000)),
|
||||
("gpt-6-sol", Some("gpt-5.6-sol"), None),
|
||||
] {
|
||||
let created =
|
||||
json!({"type":"response.created","response":{"id":"resp_model_cost","model":returned}});
|
||||
let completed = json!({
|
||||
"type":"response.completed",
|
||||
"response":{
|
||||
"id":"resp_model_cost","model":returned,"status":"completed","output":[],
|
||||
"usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25},"total_tokens":110}
|
||||
}
|
||||
});
|
||||
let body = format!("data: {created}\n\ndata: {completed}\n\n");
|
||||
let mut decoder = CodexCanonicalDecoder::new(sent).with_raw_sse_passthrough();
|
||||
let events = decoder.push(body.as_bytes()).expect("model cost response");
|
||||
let cost = canonical_facts(&events)
|
||||
.into_iter()
|
||||
.find_map(|event| match event {
|
||||
GatewayEvent::CalculatedCost(cost) => Some(cost.total().amount().scaled()),
|
||||
_ => None,
|
||||
});
|
||||
assert_eq!(cost, expected_cost, "{sent} -> {returned:?}");
|
||||
assert_eq!(decoder.response_model(), returned);
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::Completed(meta) if meta.model() == Some(returned.unwrap_or(sent))
|
||||
)));
|
||||
let raw: Vec<u8> = events
|
||||
.iter()
|
||||
.filter_map(|event| event.wire_event().and_then(|wire| wire.raw_sse_frame()))
|
||||
.flat_map(|frame| frame.iter().copied())
|
||||
.collect();
|
||||
assert_eq!(raw, body.as_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
|
||||
let body = concat!(
|
||||
@@ -207,7 +299,7 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
let events = CodexCanonicalDecoder::new("gpt-5.4")
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical priced response");
|
||||
|
||||
@@ -227,7 +319,7 @@ fn decoder_should_bill_requested_service_tier_despite_default_response() {
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
let mut decoder =
|
||||
CodexCanonicalDecoder::new("fallback").with_requested_service_tier(Some("priority"));
|
||||
CodexCanonicalDecoder::new("gpt-5.4").with_requested_service_tier(Some("priority"));
|
||||
let events = decoder
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical priority response");
|
||||
@@ -248,7 +340,7 @@ fn decoder_should_bill_standard_when_request_omits_service_tier() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
let mut decoder = CodexCanonicalDecoder::new("fallback");
|
||||
let mut decoder = CodexCanonicalDecoder::new("gpt-5.4");
|
||||
let events = decoder
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical response-tier priority response");
|
||||
@@ -270,7 +362,7 @@ fn decoder_should_add_standard_web_search_call_cost() {
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_search_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
let events = CodexCanonicalDecoder::new("gpt-5.4")
|
||||
.with_request_tool_pricing("gpt-5.4", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical web search response");
|
||||
@@ -291,7 +383,7 @@ fn decoder_should_use_non_reasoning_preview_web_search_price() {
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_preview_search_cost\",\"model\":\"gpt-4o\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search_preview" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
let events = CodexCanonicalDecoder::new("gpt-4o")
|
||||
.with_request_tool_pricing("gpt-4o", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical preview web search response");
|
||||
@@ -312,7 +404,7 @@ fn astra_decoder_should_use_reasoning_preview_web_search_price() {
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_preview_search_cost\",\"model\":\"gpt-6-astra\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search_preview" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
let events = CodexCanonicalDecoder::new("gpt-6-astra")
|
||||
.with_request_tool_pricing("gpt-6-astra", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical preview web search response");
|
||||
@@ -333,7 +425,7 @@ fn decoder_should_fail_closed_for_fixed_block_web_search_content() {
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fixed_search_cost\",\"model\":\"gpt-4o-mini\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
let events = CodexCanonicalDecoder::new("gpt-4o-mini")
|
||||
.with_request_tool_pricing("gpt-4o-mini", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical fixed-block search response");
|
||||
@@ -355,7 +447,7 @@ fn websocket_decoder_should_bill_requested_service_tier_despite_default_response
|
||||
r#"{"type":"response.completed","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","status":"completed","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25,"cache_write_tokens":0},"total_tokens":110}}}"#,
|
||||
)
|
||||
.expect("completed frame");
|
||||
let mut decoder = CodexCanonicalDecoder::new("fallback")
|
||||
let mut decoder = CodexCanonicalDecoder::new("gpt-5.4")
|
||||
.with_requested_service_tier(Some("priority"))
|
||||
.with_raw_sse_passthrough();
|
||||
let events = decoder
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use gateway_core::operation::{GenerateRequest, ProtocolPayload};
|
||||
use provider_openai::encode_generate_request;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::super::*;
|
||||
|
||||
#[test]
|
||||
fn encoder_should_adapt_pi_responses_parameters_without_losing_codex_fields() {
|
||||
@@ -43,13 +44,13 @@ fn encoder_should_adapt_pi_responses_parameters_without_losing_codex_fields() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_should_preserve_unknown_fields_and_nested_business_parameters() {
|
||||
let body = json!({
|
||||
fn encoder_should_default_missing_store_without_changing_unknown_or_nested_fields() {
|
||||
let mut body = json!({
|
||||
"model": "client-model",
|
||||
"input": "temperature and max_output_tokens are tool parameter names",
|
||||
"max_tokens": 256,
|
||||
"future_options": {"temperature": 0.7},
|
||||
"client_metadata": {"prompt_cache_retention": "business-value"},
|
||||
"client_metadata": {"prompt_cache_retention": "business-value", "store": true},
|
||||
"tools": [{
|
||||
"type": "function", "name": "configure_sampler",
|
||||
"parameters": {
|
||||
@@ -72,5 +73,121 @@ fn encoder_should_preserve_unknown_fields_and_nested_business_parameters() {
|
||||
)
|
||||
.expect("encode business parameters");
|
||||
|
||||
body["store"] = json!(false);
|
||||
assert_eq!(Value::Object(encoded.body().clone()), body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_should_preserve_explicit_store_values() {
|
||||
for store in [json!(false), json!(true), Value::Null, json!("false")] {
|
||||
let body = json!({"model": "gpt-test", "input": "hello", "store": store});
|
||||
let encoded = encode_downstream_request(body.clone());
|
||||
|
||||
assert_eq!(Value::Object(encoded.body().clone()), body);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_http_should_send_default_store_for_downstream_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind HTTP server");
|
||||
let address = listener.local_addr().expect("HTTP server address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.expect("accept HTTP client");
|
||||
let raw = read_http_request_with_body(&mut stream).await;
|
||||
write_completed_sse_response(&mut stream).await;
|
||||
let head_end = raw
|
||||
.windows(4)
|
||||
.position(|window| window == b"\r\n\r\n")
|
||||
.expect("HTTP head terminator")
|
||||
+ 4;
|
||||
let body = zstd::stream::decode_all(&raw[head_end..]).expect("decode request body");
|
||||
serde_json::from_slice::<Value>(&body).expect("HTTP request JSON")
|
||||
});
|
||||
let mut request = encode_downstream_request(json!({"model": "gpt-test", "input": "hello"}));
|
||||
request.force_http_sse = true;
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("HTTP client"),
|
||||
format!("http://{address}"),
|
||||
test_wire_profile(),
|
||||
);
|
||||
|
||||
client
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_default_store_http", Some("test-account")),
|
||||
)
|
||||
.await
|
||||
.expect("HTTP response");
|
||||
|
||||
assert_eq!(
|
||||
server.await.expect("HTTP server task"),
|
||||
json!({
|
||||
"model": "gpt-test", "input": "hello", "store": false, "stream": true
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_websocket_should_send_default_store_for_downstream_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind WebSocket server");
|
||||
let address = listener.local_addr().expect("WebSocket server address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.expect("accept WebSocket client");
|
||||
let mut websocket = accept_codex_test_websocket(stream).await;
|
||||
let message = websocket
|
||||
.next()
|
||||
.await
|
||||
.expect("response.create")
|
||||
.expect("valid frame");
|
||||
let body = serde_json::from_str::<Value>(message.to_text().expect("text frame"))
|
||||
.expect("WebSocket request JSON");
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_default_store", 1, 1).into(),
|
||||
))
|
||||
.await
|
||||
.expect("send terminal event");
|
||||
body
|
||||
});
|
||||
let mut request = encode_downstream_request(json!({"model": "gpt-test", "input": "hello"}));
|
||||
request.use_websocket = true;
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("HTTP client"),
|
||||
format!("http://{address}"),
|
||||
test_wire_profile(),
|
||||
)
|
||||
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
|
||||
|
||||
client
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_default_store_ws", Some("test-account")),
|
||||
)
|
||||
.await
|
||||
.expect("WebSocket response");
|
||||
|
||||
let body = server.await.expect("WebSocket server task");
|
||||
assert_eq!(body.get("store"), Some(&json!(false)));
|
||||
}
|
||||
|
||||
fn encode_downstream_request(body: Value) -> CodexResponsesRequest {
|
||||
let payload =
|
||||
ProtocolPayload::json_object("openai", body.as_object().expect("request object").clone())
|
||||
.expect("OpenAI payload");
|
||||
encode_generate_request(
|
||||
&GenerateRequest::from_protocol_payload(payload),
|
||||
"gpt-test",
|
||||
None,
|
||||
)
|
||||
.expect("encode downstream request")
|
||||
}
|
||||
|
||||
@@ -172,6 +172,7 @@ struct CollectedBackendResponse {
|
||||
rate_limit_headers: Vec<(String, String)>,
|
||||
websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
response_metadata: CodexResponseMetadata,
|
||||
reported_model: Option<String>,
|
||||
transport_metrics: CodexTransportMetrics,
|
||||
connection_local_continuation: bool,
|
||||
}
|
||||
@@ -246,7 +247,7 @@ async fn collect_backend_response(
|
||||
set_cookie_headers,
|
||||
mut rate_limit_headers,
|
||||
rate_limit_updates,
|
||||
turn_state_update,
|
||||
response_metadata_updates,
|
||||
websocket_pool_decision,
|
||||
diagnostics: _,
|
||||
response_metadata,
|
||||
@@ -268,8 +269,11 @@ async fn collect_backend_response(
|
||||
rate_limit_headers.extend(rate_limits_to_header_pairs(update));
|
||||
}
|
||||
}
|
||||
if let Some(update) = turn_state_update {
|
||||
turn_state = update.lock().await.clone().or(turn_state);
|
||||
let mut reported_model = response_metadata.effective_model.clone();
|
||||
if let Some(update) = response_metadata_updates {
|
||||
let update = update.lock().await;
|
||||
turn_state = update.turn_state.clone().or(turn_state);
|
||||
reported_model = update.reported_model.clone().or(reported_model);
|
||||
}
|
||||
let body = String::from_utf8_lossy(&body_bytes).into_owned();
|
||||
let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?;
|
||||
@@ -282,6 +286,7 @@ async fn collect_backend_response(
|
||||
rate_limit_headers,
|
||||
websocket_pool_decision,
|
||||
response_metadata,
|
||||
reported_model,
|
||||
transport_metrics,
|
||||
connection_local_continuation,
|
||||
})
|
||||
|
||||
@@ -118,7 +118,7 @@ fn encoder_should_preserve_openai_wire_fields_without_deriving_accountless_pool_
|
||||
|
||||
assert_eq!(encoded.body().get("model"), Some(&json!("gpt-routed")));
|
||||
assert!(encoded.body().get("stream").is_none());
|
||||
assert!(encoded.body().get("store").is_none());
|
||||
assert_eq!(encoded.body().get("store"), Some(&json!(false)));
|
||||
assert_eq!(encoded.body().get("tool_choice"), Some(&json!("auto")));
|
||||
assert_eq!(
|
||||
encoded.body().get("future_official_field"),
|
||||
|
||||
@@ -78,6 +78,7 @@ async fn reused_websocket_should_keep_response_metadata_scoped_to_each_exchange(
|
||||
.expect("response within timeout")
|
||||
.expect("response should complete on the same socket");
|
||||
assert_eq!(response.response_metadata, expected, "exchange {round}");
|
||||
assert_eq!(response.reported_model, Some(format!("model-{round}")));
|
||||
assert!(response.body.contains(&format!("etag-{round}")));
|
||||
assert_eq!(response.turn_state, Some(format!("turn-{round}")));
|
||||
if round > 0 {
|
||||
|
||||
@@ -931,7 +931,8 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": {
|
||||
"x-codex-turn-state": ["turn-from-metadata"]
|
||||
"x-codex-turn-state": ["turn-from-metadata"],
|
||||
"x-openai-model": ["gpt-internal-report"]
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
@@ -969,6 +970,10 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(response.turn_state.as_deref(), Some("turn-from-metadata"));
|
||||
assert_eq!(
|
||||
response.reported_model.as_deref(),
|
||||
Some("gpt-internal-report")
|
||||
);
|
||||
assert!(!response.body.contains("codex.rate_limits"));
|
||||
assert!(!response.body.contains("response.metadata"));
|
||||
assert!(response.body.contains("event: response.completed"));
|
||||
@@ -1048,7 +1053,7 @@ async fn reused_websocket_should_not_leak_turn_state_into_the_next_response() {
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "codex.response.metadata",
|
||||
"headers": {"x-codex-turn-state": "turn-from-first-response"}
|
||||
"headers": {"x-codex-turn-state": "turn-from-first-response", "openai-model": "first-response-model"}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
@@ -1104,6 +1109,11 @@ async fn reused_websocket_should_not_leak_turn_state_into_the_next_response() {
|
||||
.websocket_pool_decision
|
||||
.is_some_and(WebSocketPoolDecision::is_reuse)
|
||||
);
|
||||
assert_eq!(
|
||||
first.reported_model.as_deref(),
|
||||
Some("first-response-model")
|
||||
);
|
||||
assert_eq!(second.reported_model, None);
|
||||
assert_eq!(second.turn_state, None);
|
||||
assert!(
|
||||
second
|
||||
|
||||
@@ -929,16 +929,57 @@ impl GrokCredentialCatalogService {
|
||||
.list_loaded_for_provider()
|
||||
.await
|
||||
.map_err(|_| GrokCredentialCatalogError::Store)?;
|
||||
let mut groups = catalog_candidates_by_scope(candidates)?;
|
||||
let Some(scope) = groups.iter().find_map(|(scope, candidates)| {
|
||||
let mut groups = match catalog_candidates_by_scope(candidates) {
|
||||
Ok(groups) => groups,
|
||||
// 全部候选都被调度列表过滤时按空组处理,让下方目标账号补回继续生效。
|
||||
Err(GrokCredentialCatalogError::NoEligibleCredential) => BTreeMap::new(),
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
let scope = groups.iter().find_map(|(scope, candidates)| {
|
||||
candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.account.id() == account_id)
|
||||
.then(|| scope.clone())
|
||||
}) else {
|
||||
return Err(GrokCredentialCatalogError::NoEligibleCredential);
|
||||
});
|
||||
let mut pinned_candidate = None;
|
||||
let scope = match scope {
|
||||
Some(scope) => scope,
|
||||
None => {
|
||||
// 常规调度列表不含停用账号;管理端按账号查询模型要对停用账号返回
|
||||
// 真实上游结果。按同一 revision 加载账号与凭据,避免并发更新套餐时
|
||||
// 把新凭据的目录写入旧套餐 cache。
|
||||
let account = self
|
||||
.repository
|
||||
.account_by_id(account_id)
|
||||
.await
|
||||
.map_err(|_| GrokCredentialCatalogError::Store)?
|
||||
.ok_or(GrokCredentialCatalogError::NoEligibleCredential)?;
|
||||
let loaded = self
|
||||
.repository
|
||||
.load(account_id, account.revision())
|
||||
.await
|
||||
.map_err(|_| GrokCredentialCatalogError::Store)?;
|
||||
let scope = GrokCatalogScope::for_account(&loaded.account)
|
||||
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?;
|
||||
pinned_candidate = Some(loaded);
|
||||
scope
|
||||
}
|
||||
};
|
||||
let mut candidates = groups.remove(&scope).unwrap_or_default();
|
||||
if !candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.account.id() == account_id)
|
||||
{
|
||||
let loaded = match pinned_candidate {
|
||||
Some(loaded) => loaded,
|
||||
None => self
|
||||
.repository
|
||||
.load_current(account_id)
|
||||
.await
|
||||
.map_err(|_| GrokCredentialCatalogError::Store)?,
|
||||
};
|
||||
candidates.push(loaded);
|
||||
}
|
||||
candidates.sort_by(|left, right| {
|
||||
let left_preferred = left.account.id() == account_id;
|
||||
let right_preferred = right.account.id() == account_id;
|
||||
|
||||
@@ -252,6 +252,10 @@ pub(super) fn cold_compaction_http_sse_stream(
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(model) = canonical.response_model() {
|
||||
observation = observation.with_upstream_response_model_if_valid(model);
|
||||
yield ProviderEvent::observation(observation.clone());
|
||||
}
|
||||
if let Some(tier) = canonical.response_service_tier() {
|
||||
observation = observation.with_service_tier_if_valid(tier.to_owned());
|
||||
yield ProviderEvent::observation(observation);
|
||||
@@ -534,6 +538,14 @@ pub(super) fn cold_http_sse_stream(
|
||||
}),
|
||||
);
|
||||
}
|
||||
if let Some(model) = decoder.response_model()
|
||||
&& observation.upstream_response_model() != Some(model)
|
||||
{
|
||||
observation = observation.with_upstream_response_model_if_valid(model);
|
||||
yield ProviderEvent::observation(observation.clone().with_timings(
|
||||
ProviderResponseTimings { first_token_ms, ..base_timings },
|
||||
));
|
||||
}
|
||||
if let Some(tier) = decoder.response_service_tier()
|
||||
&& observation.service_tier() != Some(tier)
|
||||
{
|
||||
@@ -594,6 +606,14 @@ pub(super) fn cold_http_sse_stream(
|
||||
}),
|
||||
);
|
||||
}
|
||||
if let Some(model) = decoder.response_model()
|
||||
&& observation.upstream_response_model() != Some(model)
|
||||
{
|
||||
observation = observation.with_upstream_response_model_if_valid(model);
|
||||
yield ProviderEvent::observation(observation.clone().with_timings(
|
||||
ProviderResponseTimings { first_token_ms, ..base_timings },
|
||||
));
|
||||
}
|
||||
if let Some(tier) = decoder.response_service_tier()
|
||||
&& observation.service_tier() != Some(tier)
|
||||
{
|
||||
|
||||
@@ -14,7 +14,9 @@ use gateway_core::metering::{
|
||||
CurrencyCode, Decimal, Money, ProviderReportedCost, Usage,
|
||||
};
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_protocol::openai::events::{TokenUsage, billable_usage_is_complete, extract_usage};
|
||||
use gateway_protocol::openai::events::{
|
||||
ResponseModelObservation, TokenUsage, billable_usage_is_complete, extract_usage,
|
||||
};
|
||||
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder};
|
||||
use serde_json::Value;
|
||||
|
||||
@@ -167,7 +169,7 @@ fn usd_price_per_million(per_token_ticks: u128) -> Option<Money> {
|
||||
pub struct GrokCanonicalDecoder {
|
||||
decoder: SseEventDecoder,
|
||||
response_transform: GrokResponseTransform,
|
||||
fallback_model: String,
|
||||
upstream_model: String,
|
||||
response_id: Option<String>,
|
||||
started: bool,
|
||||
completed: bool,
|
||||
@@ -176,15 +178,17 @@ pub struct GrokCanonicalDecoder {
|
||||
usage_emitted: bool,
|
||||
output_start_seen: bool,
|
||||
response_service_tier: Option<String>,
|
||||
response_model: ResponseModelObservation,
|
||||
requires_provider_cost: bool,
|
||||
}
|
||||
|
||||
impl GrokCanonicalDecoder {
|
||||
pub fn new(fallback_model: impl Into<String>) -> Self {
|
||||
/// 使用路由后最终发往上游的请求模型计价,并在响应缺少模型时用于 canonical 兜底。
|
||||
pub fn new(upstream_model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
decoder: SseEventDecoder::default(),
|
||||
response_transform: GrokResponseTransform::default(),
|
||||
fallback_model: fallback_model.into(),
|
||||
upstream_model: upstream_model.into(),
|
||||
response_id: None,
|
||||
started: false,
|
||||
completed: false,
|
||||
@@ -193,6 +197,7 @@ impl GrokCanonicalDecoder {
|
||||
usage_emitted: false,
|
||||
output_start_seen: false,
|
||||
response_service_tier: None,
|
||||
response_model: ResponseModelObservation::default(),
|
||||
requires_provider_cost: false,
|
||||
}
|
||||
}
|
||||
@@ -200,7 +205,7 @@ impl GrokCanonicalDecoder {
|
||||
/// 创建在 canonical 与 wire 投影处理每个上游 event 前先还原请求级 tool 别名的
|
||||
/// 解码器。
|
||||
#[must_use]
|
||||
pub fn for_request(fallback_model: impl Into<String>, request: &GrokResponsesRequest) -> Self {
|
||||
pub fn for_request(upstream_model: impl Into<String>, request: &GrokResponsesRequest) -> Self {
|
||||
Self {
|
||||
response_transform: request.response_transform(),
|
||||
requires_provider_cost: request.body().get("tool_choice").and_then(Value::as_str)
|
||||
@@ -217,7 +222,7 @@ impl GrokCanonicalDecoder {
|
||||
)
|
||||
})
|
||||
}),
|
||||
..Self::new(fallback_model)
|
||||
..Self::new(upstream_model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -227,6 +232,12 @@ impl GrokCanonicalDecoder {
|
||||
self.response_service_tier.as_deref()
|
||||
}
|
||||
|
||||
/// 返回原始上游响应声明的模型,缺失时不使用请求模型补齐。
|
||||
#[must_use]
|
||||
pub fn response_model(&self) -> Option<&str> {
|
||||
self.response_model.model()
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<ProviderEvent>, ProviderError> {
|
||||
let events = self.decoder.push(chunk).map_err(protocol_error)?;
|
||||
self.decode(events)
|
||||
@@ -314,6 +325,7 @@ impl GrokCanonicalDecoder {
|
||||
continue;
|
||||
};
|
||||
let event_type = event_type.to_owned();
|
||||
self.response_model.observe(Some(&event_type), &value);
|
||||
// 工具转换可能隐藏注入的调用,计费事实必须从转换前的上游事件读取。
|
||||
if let Some(response) = value.get("response") {
|
||||
if let Some(tier) = response.get("service_tier").and_then(Value::as_str) {
|
||||
@@ -442,7 +454,7 @@ impl GrokCanonicalDecoder {
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|model| !model.is_empty())
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.unwrap_or(&self.upstream_model)
|
||||
.to_owned();
|
||||
self.response_id = Some(response_id.clone());
|
||||
self.started = true;
|
||||
@@ -649,13 +661,19 @@ impl GrokCanonicalDecoder {
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|model| !model.is_empty())
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.unwrap_or(&self.upstream_model)
|
||||
.to_owned();
|
||||
if let Some(cost) = provider_reported_cost(response)? {
|
||||
output.push(GatewayEvent::ProviderCost(cost));
|
||||
} else if !self.requires_provider_cost
|
||||
&& let Some(cost) = usage.and_then(|usage| {
|
||||
calculated_cost(response, &model, usage, self.response_service_tier())
|
||||
// 已报告金额仍优先;本地估价只使用实际发送模型,响应模型仅作观测。
|
||||
calculated_cost(
|
||||
response,
|
||||
&self.upstream_model,
|
||||
usage,
|
||||
self.response_service_tier(),
|
||||
)
|
||||
})
|
||||
{
|
||||
output.push(GatewayEvent::CalculatedCost(cost));
|
||||
|
||||
@@ -312,6 +312,27 @@ async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() {
|
||||
assert_eq!(cached.seed().models(), ["grok-4.5"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_account_refresh_discovers_models_with_the_pinned_account() {
|
||||
let (store, repository) =
|
||||
repository_with_accounts(&[("disabled-models", "subject-disabled-models")]).await;
|
||||
store
|
||||
.set_enabled(&account_id("disabled-models"), false)
|
||||
.await
|
||||
.expect("disable account");
|
||||
let service = crate::support::grok_catalog_service(
|
||||
repository,
|
||||
QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]),
|
||||
MemoryGrokCatalogCache::shared(),
|
||||
);
|
||||
|
||||
let refreshed = service
|
||||
.refresh_account_catalog(&account_id("disabled-models"))
|
||||
.await
|
||||
.expect("disabled account refresh still discovers models");
|
||||
assert_eq!(refreshed.seed().models(), ["grok-4.5"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn single_account_catalog_read_miss_does_not_call_upstream() {
|
||||
let (store, repository) =
|
||||
|
||||
@@ -10,6 +10,60 @@ use provider_xai::{
|
||||
grok_billing_breakdown_with_tier,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn decoder_should_bill_the_sent_model_independently_of_the_response_model() {
|
||||
for (sent, returned, reported_cost, expected_cost) in [
|
||||
("grok-4.5", Some("grok-future"), None, Some(2_175_000)),
|
||||
("grok-4.5", Some("grok-4.3"), None, Some(2_175_000)),
|
||||
("grok-4.5", None, None, Some(2_175_000)),
|
||||
("grok-future", Some("grok-4.5"), None, None),
|
||||
("grok-4.5", Some("grok-future"), Some(123), None),
|
||||
("grok-4.5", Some("grok-future"), Some(0), None),
|
||||
] {
|
||||
let created = serde_json::json!({"type":"response.created","response":{"id":"resp_model_cost","model":returned}});
|
||||
let mut completed = serde_json::json!({
|
||||
"type":"response.completed",
|
||||
"response":{
|
||||
"id":"resp_model_cost","model":returned,"status":"completed","output":[],
|
||||
"usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25},"total_tokens":110}
|
||||
}
|
||||
});
|
||||
if let Some(ticks) = reported_cost {
|
||||
completed["response"]["usage"]["cost_in_usd_ticks"] = serde_json::json!(ticks);
|
||||
}
|
||||
let body = format!("data: {created}\n\ndata: {completed}\n\n");
|
||||
let mut decoder = GrokCanonicalDecoder::new(sent);
|
||||
let events = decoder.push(body.as_bytes()).expect("model cost response");
|
||||
let facts: Vec<_> = events
|
||||
.into_iter()
|
||||
.flat_map(|event| event.into_parts().0)
|
||||
.collect();
|
||||
assert_eq!(
|
||||
calculated_cost_ticks(&facts),
|
||||
expected_cost,
|
||||
"{sent} -> {returned:?}"
|
||||
);
|
||||
assert_eq!(provider_cost_ticks(&facts), reported_cost);
|
||||
assert_eq!(decoder.response_model(), returned);
|
||||
assert!(facts.iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::Completed(meta) if meta.model() == Some(returned.unwrap_or(sent))
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_model_observation_preserves_raw_model_and_missing_values() {
|
||||
for model in [None, Some("grok-4.6-build")] {
|
||||
let mut decoder = GrokCanonicalDecoder::new("grok-4.6");
|
||||
let body = serde_json::json!({"type":"response.created","response":{"id":"resp_model","model":model}});
|
||||
decoder
|
||||
.push(format!("data: {body}\n\n").as_bytes())
|
||||
.expect("decode created");
|
||||
assert_eq!(decoder.response_model(), model);
|
||||
}
|
||||
}
|
||||
|
||||
fn terminal_cost_events(
|
||||
model: &str,
|
||||
input_tokens: u64,
|
||||
@@ -33,11 +87,11 @@ fn terminal_cost_events(
|
||||
cached_tokens = cached_tokens,
|
||||
provider_cost = provider_cost,
|
||||
);
|
||||
decode_canonical(body.as_bytes()).expect("canonical cost response")
|
||||
decode_canonical(body.as_bytes(), model).expect("canonical cost response")
|
||||
}
|
||||
|
||||
fn decode_canonical(body: &[u8]) -> Result<Vec<GatewayEvent>, ProviderError> {
|
||||
GrokCanonicalDecoder::new("fallback")
|
||||
fn decode_canonical(body: &[u8], upstream_model: &str) -> Result<Vec<GatewayEvent>, ProviderError> {
|
||||
GrokCanonicalDecoder::new(upstream_model)
|
||||
.push(body)
|
||||
.map(|events| {
|
||||
events
|
||||
@@ -97,7 +151,7 @@ fn decoder_should_normalize_text_usage_and_completion() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"model\":\"grok-code-test\",\"status\":\"completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":2,\"total_tokens\":5,\"cost_in_usd_ticks\":37756000}}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical response");
|
||||
let events = decode_canonical(body.as_bytes(), "fallback").expect("canonical response");
|
||||
|
||||
assert!(matches!(events[0], GatewayEvent::Started(_)));
|
||||
assert!(matches!(
|
||||
@@ -188,7 +242,8 @@ fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_usage\",\"model\":\"grok-4.5\",\"status\":\"completed\",\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"total_tokens\":17},\"tool_usage\":{\"image_gen\":{\"input_tokens\":31,\"output_tokens\":9}}}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical image usage response");
|
||||
let events =
|
||||
decode_canonical(body.as_bytes(), "fallback").expect("canonical image usage response");
|
||||
let image_usage = events.iter().find_map(|event| match event {
|
||||
GatewayEvent::Usage(usage) => Some((usage.image_input_tokens, usage.image_output_tokens)),
|
||||
_ => None,
|
||||
@@ -205,7 +260,7 @@ fn decoder_should_leave_cost_unavailable_when_upstream_omits_it() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_no_cost\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical response");
|
||||
let events = decode_canonical(body.as_bytes(), "fallback").expect("canonical response");
|
||||
|
||||
assert!(!events.iter().any(|event| matches!(
|
||||
event,
|
||||
@@ -420,7 +475,8 @@ fn decoder_should_leave_incomplete_usage_pricing_unavailable() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_partial\",\"model\":\"grok-4.5\",\"status\":\"completed\",\"usage\":{\"total_tokens\":3}}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical partial usage response");
|
||||
let events =
|
||||
decode_canonical(body.as_bytes(), "grok-4.5").expect("canonical partial usage response");
|
||||
|
||||
assert_eq!(calculated_cost_ticks(&events), None);
|
||||
}
|
||||
@@ -501,7 +557,7 @@ fn decoder_should_normalize_function_call_and_tool_finish() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_tool\",\"model\":\"grok-code-test\",\"status\":\"completed\"}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical function call");
|
||||
let events = decode_canonical(body.as_bytes(), "fallback").expect("canonical function call");
|
||||
|
||||
assert!(events.iter().any(|event| matches!(
|
||||
event,
|
||||
@@ -1076,7 +1132,7 @@ fn decoder_should_coalesce_reasoning_item_part_and_summary_index() {
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_reason\",\"status\":\"completed\"}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("canonical reasoning");
|
||||
let events = decode_canonical(body.as_bytes(), "fallback").expect("canonical reasoning");
|
||||
|
||||
assert_eq!(
|
||||
events
|
||||
@@ -1303,7 +1359,7 @@ fn decoder_should_preserve_incomplete_length_reason() {
|
||||
"event: response.incomplete\n",
|
||||
"data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_short\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n",
|
||||
);
|
||||
let events = decode_canonical(body.as_bytes()).expect("incomplete response");
|
||||
let events = decode_canonical(body.as_bytes(), "fallback").expect("incomplete response");
|
||||
|
||||
assert!(matches!(
|
||||
events.last(),
|
||||
|
||||
@@ -9,3 +9,7 @@ b7d3df7b48d835a00a3b74cbefc958a039d38fea5529cdad04bfbc703dca72aa 0001_initial.s
|
||||
161df9ae43da34d51b340bbf95d71c61a8c1a2aa142382f693a9d127c834630d 0009_account_model_access.sql
|
||||
3751a802e80b5d538abc5619a2397aa4bf13f52076068a489892f664907adf49 0010_account_auto_freeze.sql
|
||||
ae7f7acb3014f0a60d979dd9c04b3065322b206e42eb6256505baf78caf9dd22 0011_proxy_request_location.sql
|
||||
b20e425c0b062b72d35fdebe294e023eb17ceb44ac93083ff3d48400715bf430 0012_upstream_response_model.sql
|
||||
6e23fceef1096b50c1af3584cb1687eded2ad37e31e65391397ec984bf78c7e3 0013_responses_decompression_limit.sql
|
||||
12a540be31d65c4925cbe63dab42ea2d4ec2864fecf71a6dd4ee942e71726857 0014_disable_fast.sql
|
||||
9bd37bd314b6864816c43cb4fd1869fb415920d4f10f766976e344334a1f60fc 0015_proxy_dual_stack_exit.sql
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
-- 返回模型独立于发送模型,避免展示观测改写路由与统计口径。
|
||||
ALTER TABLE model_requests ADD COLUMN upstream_response_model text;
|
||||
|
||||
-- 历史记录只回填已经保存的明确报告,不从请求模型或映射模型推断。
|
||||
UPDATE model_requests
|
||||
SET upstream_response_model = btrim(provider_observation_json->>'upstreamReportedModel')
|
||||
WHERE provider_kind = 'openai'
|
||||
AND jsonb_typeof(provider_observation_json->'upstreamReportedModel') = 'string'
|
||||
AND octet_length(btrim(provider_observation_json->>'upstreamReportedModel')) BETWEEN 1 AND 256
|
||||
AND provider_observation_json->>'upstreamReportedModel' !~ '[[:cntrl:]]';
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user