mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
fix(codex): serve versioned dynamic model catalogs
This commit is contained in:
@@ -339,7 +339,7 @@ fn apply_codex_oauth_fingerprint_convergence_to_decision(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
crate::provider_transport::apply_codex_oauth_fingerprint_convergence(
|
||||
crate::ai_serving::transport::apply_codex_oauth_fingerprint_convergence(
|
||||
transport,
|
||||
provider_api_format,
|
||||
input.original_client_session_id.as_deref(),
|
||||
|
||||
@@ -82,10 +82,10 @@ pub(crate) use aether_ai_formats::api::{
|
||||
parse_codex_auth_identity, parse_direct_request_body, parse_model_directive,
|
||||
parse_model_directive_with_suffixes, parse_openai_stop_sequences,
|
||||
parse_openai_tool_result_content, prepare_local_success_response_parts,
|
||||
prepare_local_success_response_parts_owned, project_codex_openai_image_api_request_body,
|
||||
project_openai_image_api_request_body, provider_adaptation_allows_sync_finalize_envelope,
|
||||
provider_adaptation_anchor_api_format, provider_adaptation_descriptor_for_envelope,
|
||||
provider_adaptation_descriptor_for_provider_type,
|
||||
prepare_local_success_response_parts_owned, project_codex_catalog_model_card,
|
||||
project_codex_openai_image_api_request_body, project_openai_image_api_request_body,
|
||||
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format,
|
||||
provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type,
|
||||
provider_adaptation_requires_eventstream_accept,
|
||||
provider_adaptation_should_unwrap_stream_envelope,
|
||||
provider_private_response_allows_sync_finalize, record_converted_response_history,
|
||||
@@ -175,6 +175,7 @@ pub(crate) use aether_ai_formats::{
|
||||
is_rerank_api_format, openai_responses_request_operation,
|
||||
openai_responses_synthetic_reasoning_item_id,
|
||||
strip_incompatible_openai_responses_reasoning_items, ApiOperation, ClientSurface,
|
||||
CODEX_CLIENT_VERSION,
|
||||
};
|
||||
|
||||
pub(crate) fn plan_kind_matches_api_operation(
|
||||
|
||||
@@ -59,8 +59,9 @@ pub(crate) mod windsurf {
|
||||
}
|
||||
|
||||
pub(crate) use aether_provider_transport::{
|
||||
append_transport_diagnostics_to_value, apply_local_auth_config_header_overrides,
|
||||
apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
|
||||
append_transport_diagnostics_to_value, apply_codex_oauth_fingerprint_convergence,
|
||||
apply_local_auth_config_header_overrides, apply_local_body_rules,
|
||||
apply_local_body_rules_with_request_headers, apply_local_header_rules,
|
||||
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
|
||||
apply_standard_provider_request_body_rules_with_request_headers,
|
||||
apply_transport_request_body_semantics, body_rules_are_locally_supported,
|
||||
|
||||
@@ -85,29 +85,17 @@ async fn read_last_backup_slot(app: &AppState) -> Result<Option<String>, Gateway
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::backup::schedule::{BackupSchedule, BackupScheduleUnit};
|
||||
use crate::task_runtime::{task_definition, TASK_KEY_SYSTEM_S3_BACKUP};
|
||||
|
||||
#[test]
|
||||
fn backup_worker_skips_already_recorded_slot() {
|
||||
let schedule = BackupSchedule {
|
||||
unit: BackupScheduleUnit::Days,
|
||||
interval: 1,
|
||||
minute: 0,
|
||||
hour: 3,
|
||||
weekday: 1,
|
||||
month_day: 1,
|
||||
};
|
||||
let now = chrono::DateTime::parse_from_rfc3339("2026-05-24T03:00:30+08:00")
|
||||
.unwrap()
|
||||
.with_timezone(&chrono::Utc);
|
||||
let slot = schedule.due_slot(now).expect("slot should be due");
|
||||
let slot = "days:2026-05-23T19:00:00Z";
|
||||
|
||||
assert!(super::should_start_scheduled_backup(
|
||||
Some("days:2026-05-22T19:00:00Z"),
|
||||
&slot
|
||||
slot
|
||||
));
|
||||
assert!(!super::should_start_scheduled_backup(Some(&slot), &slot));
|
||||
assert!(!super::should_start_scheduled_backup(Some(slot), slot));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
-10
@@ -1274,7 +1274,6 @@ mod tests {
|
||||
let first_cache = Arc::clone(&cache);
|
||||
let first_key = key.clone();
|
||||
let first_calls = Arc::clone(&calls);
|
||||
let first_started = Instant::now();
|
||||
let first = tokio::spawn(async move {
|
||||
first_cache
|
||||
.get_or_load_once_stale_while_refreshing::<(), _, _>(
|
||||
@@ -1292,7 +1291,6 @@ mod tests {
|
||||
});
|
||||
|
||||
let follower_cache = Arc::clone(&cache);
|
||||
let follower_started = Instant::now();
|
||||
let follower_calls = Arc::clone(&calls);
|
||||
let follower = tokio::spawn(async move {
|
||||
follower_cache
|
||||
@@ -1310,15 +1308,7 @@ mod tests {
|
||||
});
|
||||
|
||||
assert_eq!(first.await.unwrap().unwrap(), Some(1));
|
||||
assert!(
|
||||
first_started.elapsed() < Duration::from_millis(80),
|
||||
"stale value should not wait for request-path refresh"
|
||||
);
|
||||
assert_eq!(follower.await.unwrap().unwrap(), Some(1));
|
||||
assert!(
|
||||
follower_started.elapsed() < Duration::from_millis(80),
|
||||
"follower should return stale value without waiting for refresh"
|
||||
);
|
||||
assert_eq!(calls.load(Ordering::Acquire), 0);
|
||||
}
|
||||
|
||||
|
||||
@@ -247,8 +247,7 @@ pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::Hea
|
||||
|
||||
let has_codex_client_version = uri.path() == "/v1/models"
|
||||
&& uri.query().is_some_and(|query| {
|
||||
url::form_urlencoded::parse(query.as_bytes())
|
||||
.any(|(key, value)| key == "client_version" && !value.trim().is_empty())
|
||||
url::form_urlencoded::parse(query.as_bytes()).any(|(key, _)| key == "client_version")
|
||||
});
|
||||
if has_codex_client_version {
|
||||
return "openai:responses".to_string();
|
||||
|
||||
@@ -39,7 +39,7 @@ fn classifies_codex_models_list_with_responses_auth_signature() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_codex_client_version_keeps_standard_openai_models_signature() {
|
||||
fn empty_codex_client_version_uses_responses_signature_for_bounded_fallback() {
|
||||
let headers = headers(&[("authorization", "Bearer sk-test")]);
|
||||
let uri: Uri = "/v1/models?client_version="
|
||||
.parse()
|
||||
@@ -49,7 +49,7 @@ fn empty_codex_client_version_keeps_standard_openai_models_signature() {
|
||||
|
||||
assert_eq!(
|
||||
decision.auth_endpoint_signature.as_deref(),
|
||||
Some("openai:chat")
|
||||
Some("openai:responses")
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -282,6 +282,16 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids_strong(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_keys_by_ids_strong(key_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
|
||||
@@ -334,6 +334,13 @@ impl ProviderCatalogReadRepository for CachedProviderCatalogReadRepository {
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids_strong(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
self.inner.list_keys_by_ids_strong(key_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
@@ -538,6 +545,7 @@ fn normalize_ids(ids: &[String]) -> Vec<String> {
|
||||
mod tests {
|
||||
use super::*;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||
|
||||
fn cache() -> CachedProviderCatalogReadRepository {
|
||||
CachedProviderCatalogReadRepository::new(Arc::new(
|
||||
@@ -555,6 +563,55 @@ mod tests {
|
||||
.expect("provider should be valid")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_catalog_strong_key_read_bypasses_fresh_cached_generation() {
|
||||
let old_metadata = serde_json::json!({
|
||||
"codex": {"credential_generation": "old"}
|
||||
});
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should be valid");
|
||||
key.upstream_metadata = Some(old_metadata.clone());
|
||||
let inner = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider("provider-1")],
|
||||
Vec::new(),
|
||||
vec![key],
|
||||
));
|
||||
let cache = CachedProviderCatalogReadRepository::new(inner.clone());
|
||||
let key_ids = vec!["key-1".to_string()];
|
||||
|
||||
let first = cache
|
||||
.list_keys_by_ids(&key_ids)
|
||||
.await
|
||||
.expect("initial key read should succeed");
|
||||
assert_eq!(first[0].upstream_metadata.as_ref(), Some(&old_metadata));
|
||||
|
||||
let new_metadata = serde_json::json!({
|
||||
"codex": {"credential_generation": "new"}
|
||||
});
|
||||
assert!(inner
|
||||
.upsert_key_upstream_metadata_namespace("key-1", "codex", &new_metadata["codex"], None,)
|
||||
.await
|
||||
.expect("inner metadata update should succeed"));
|
||||
|
||||
let cached = cache
|
||||
.list_keys_by_ids(&key_ids)
|
||||
.await
|
||||
.expect("cached key read should succeed");
|
||||
assert_eq!(cached[0].upstream_metadata.as_ref(), Some(&old_metadata));
|
||||
let strong = cache
|
||||
.list_keys_by_ids_strong(&key_ids)
|
||||
.await
|
||||
.expect("strong key read should succeed");
|
||||
assert_eq!(strong[0].upstream_metadata.as_ref(), Some(&new_metadata));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_catalog_follower_observes_completion_before_first_poll() {
|
||||
let cache = cache();
|
||||
|
||||
@@ -55,8 +55,9 @@ use crate::execution_runtime::submission::{
|
||||
resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::execution_runtime::transport::{
|
||||
append_upstream_response_body_chunk, build_execution_response_body, build_request_body,
|
||||
collect_response_headers, decode_response_body_bytes, execution_response_body_mode,
|
||||
append_upstream_response_body_chunk_with_limit, build_execution_response_body,
|
||||
build_request_body, collect_response_headers, decode_response_body_bytes_with_limit,
|
||||
execution_plan_response_body_limit_bytes, execution_response_body_mode,
|
||||
format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error,
|
||||
response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
|
||||
ExecutionRuntimeTransportError,
|
||||
@@ -1491,6 +1492,7 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
||||
) -> Result<ExecutionResult, SyncExecutionFailure> {
|
||||
let request_body = build_request_body(plan).map_err(SyncExecutionFailure::from_transport)?;
|
||||
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
|
||||
let started_at = Instant::now();
|
||||
let mut progress =
|
||||
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
|
||||
@@ -1529,8 +1531,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
@@ -1547,8 +1553,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
)),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
@@ -1565,8 +1575,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
@@ -1575,8 +1589,9 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
}
|
||||
}
|
||||
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let decoded_body_bytes =
|
||||
decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
progress.finish(status_code, elapsed_ms).await;
|
||||
|
||||
@@ -21,6 +21,7 @@ use aether_http::{apply_http_client_config, HttpClientConfig};
|
||||
use aether_runtime::{MetricKind, MetricSample};
|
||||
use axum::body::Bytes;
|
||||
use base64::Engine as _;
|
||||
use brotli::Decompressor as BrotliDecoder;
|
||||
use flate2::read::{DeflateDecoder, GzDecoder};
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
@@ -62,6 +63,10 @@ const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000;
|
||||
const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000;
|
||||
const DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS: u64 = 1_200_000;
|
||||
const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1;
|
||||
const EXECUTION_RESPONSE_BODY_LIMIT_HEADER: &str = "x-aether-execution-response-body-limit-bytes";
|
||||
const DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024;
|
||||
const MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024;
|
||||
const MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024;
|
||||
const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS";
|
||||
const DIRECT_REQWEST_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS";
|
||||
const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str =
|
||||
@@ -607,6 +612,56 @@ impl std::fmt::Display for UpstreamResponseBodyPhase {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn with_upstream_response_body_limit(
|
||||
plan: &ExecutionPlan,
|
||||
limit_bytes: usize,
|
||||
) -> ExecutionPlan {
|
||||
let mut bounded_plan = plan.clone();
|
||||
bounded_plan
|
||||
.headers
|
||||
.retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER));
|
||||
bounded_plan.headers.insert(
|
||||
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(),
|
||||
normalize_scoped_response_body_limit(limit_bytes)
|
||||
.unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
|
||||
.to_string(),
|
||||
);
|
||||
bounded_plan
|
||||
}
|
||||
|
||||
pub(crate) fn execution_plan_response_body_limit_bytes(plan: &ExecutionPlan) -> usize {
|
||||
effective_response_body_limit_bytes(
|
||||
execution_transport_header_value(&plan.headers, EXECUTION_RESPONSE_BODY_LIMIT_HEADER),
|
||||
crate::headers::max_internal_buffered_body_bytes(),
|
||||
)
|
||||
}
|
||||
|
||||
fn effective_response_body_limit_bytes(
|
||||
raw_scoped_limit: Option<&str>,
|
||||
global_limit: usize,
|
||||
) -> usize {
|
||||
let Some(raw_scoped_limit) = raw_scoped_limit else {
|
||||
return global_limit;
|
||||
};
|
||||
parse_scoped_response_body_limit(raw_scoped_limit)
|
||||
.unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
|
||||
.min(global_limit)
|
||||
}
|
||||
|
||||
fn parse_scoped_response_body_limit(value: &str) -> Option<usize> {
|
||||
let raw_limit = value.trim().parse::<u64>().ok()?;
|
||||
usize::try_from(raw_limit)
|
||||
.ok()
|
||||
.and_then(normalize_scoped_response_body_limit)
|
||||
}
|
||||
|
||||
fn normalize_scoped_response_body_limit(limit_bytes: usize) -> Option<usize> {
|
||||
(limit_bytes > 0).then_some(limit_bytes.clamp(
|
||||
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn append_upstream_response_body_chunk(
|
||||
body: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
@@ -618,7 +673,7 @@ pub(crate) fn append_upstream_response_body_chunk(
|
||||
)
|
||||
}
|
||||
|
||||
fn append_upstream_response_body_chunk_with_limit(
|
||||
pub(crate) fn append_upstream_response_body_chunk_with_limit(
|
||||
body: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
limit_bytes: usize,
|
||||
@@ -724,6 +779,7 @@ impl DirectSyncExecutionRuntime {
|
||||
F: FnOnce(DirectSyncResponseStarted),
|
||||
{
|
||||
let body_bytes = build_request_body(plan)?;
|
||||
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
|
||||
|
||||
let started_at = Instant::now();
|
||||
let request_started_at_unix_ms = crate::clock::current_unix_ms();
|
||||
@@ -744,9 +800,14 @@ impl DirectSyncExecutionRuntime {
|
||||
ttfb_ms,
|
||||
response_observation: response_observation.clone(),
|
||||
});
|
||||
let (body_bytes, stream_ttfb_ms) =
|
||||
response.bytes_with_stream_timeout(plan, started_at).await?;
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
|
||||
let (body_bytes, stream_ttfb_ms) = response
|
||||
.bytes_with_stream_timeout(plan, started_at, response_body_limit_bytes)
|
||||
.await?;
|
||||
let decoded_body_bytes = decode_response_body_bytes_with_limit(
|
||||
&headers,
|
||||
&body_bytes,
|
||||
response_body_limit_bytes,
|
||||
)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
|
||||
@@ -1058,6 +1119,7 @@ async fn execute_sync_plan_via_local_tunnel_inner(
|
||||
}
|
||||
|
||||
let body_bytes = build_request_body(plan)?;
|
||||
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
|
||||
let transport_controls = resolve_execution_transport_controls(&plan.headers);
|
||||
let headers = build_request_headers(
|
||||
&plan.headers,
|
||||
@@ -1113,8 +1175,10 @@ async fn execute_sync_plan_via_local_tunnel_inner(
|
||||
);
|
||||
let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-");
|
||||
let (body_bytes, stream_ttfb_ms) =
|
||||
collect_local_tunnel_response_body(response, plan, started_at).await?;
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
|
||||
collect_local_tunnel_response_body(response, plan, started_at, response_body_limit_bytes)
|
||||
.await?;
|
||||
let decoded_body_bytes =
|
||||
decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
if status_code >= 400 {
|
||||
@@ -1179,6 +1243,7 @@ async fn collect_local_tunnel_response_body(
|
||||
mut response: tunnel::DirectRelayResponse,
|
||||
plan: &ExecutionPlan,
|
||||
started_at: Instant,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<(Vec<u8>, Option<u64>), ExecutionRuntimeTransportError> {
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut first_byte_ms = None;
|
||||
@@ -1201,7 +1266,11 @@ async fn collect_local_tunnel_response_body(
|
||||
if plan.stream && first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok((body_bytes, first_byte_ms))
|
||||
@@ -1362,20 +1431,28 @@ impl DirectHttpResponse {
|
||||
}
|
||||
|
||||
pub(crate) async fn bytes(self) -> Result<Bytes, ExecutionRuntimeTransportError> {
|
||||
self.bytes_with_limit(crate::headers::max_internal_buffered_body_bytes())
|
||||
.await
|
||||
}
|
||||
|
||||
async fn bytes_with_limit(
|
||||
self,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<Bytes, ExecutionRuntimeTransportError> {
|
||||
let started_at = Instant::now();
|
||||
match self {
|
||||
DirectHttpResponse::Reqwest(response) => {
|
||||
collect_reqwest_stream_body(response, started_at, None)
|
||||
collect_reqwest_stream_body(response, started_at, None, response_body_limit_bytes)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
DirectHttpResponse::HyperH2c(response) => {
|
||||
collect_hyper_stream_body(response, started_at, None)
|
||||
collect_hyper_stream_body(response, started_at, None, response_body_limit_bytes)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
DirectHttpResponse::BrowserWreq(response) => {
|
||||
collect_wreq_stream_body(response, started_at, None)
|
||||
collect_wreq_stream_body(response, started_at, None, response_body_limit_bytes)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
@@ -1386,21 +1463,43 @@ impl DirectHttpResponse {
|
||||
self,
|
||||
plan: &ExecutionPlan,
|
||||
started_at: Instant,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
|
||||
if !plan.stream {
|
||||
return self.bytes().await.map(|bytes| (bytes, None));
|
||||
return self
|
||||
.bytes_with_limit(response_body_limit_bytes)
|
||||
.await
|
||||
.map(|bytes| (bytes, None));
|
||||
}
|
||||
|
||||
let first_byte_timeout = resolve_stream_first_byte_timeout(plan);
|
||||
match self {
|
||||
DirectHttpResponse::Reqwest(response) => {
|
||||
collect_reqwest_stream_body(response, started_at, first_byte_timeout).await
|
||||
collect_reqwest_stream_body(
|
||||
response,
|
||||
started_at,
|
||||
first_byte_timeout,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.await
|
||||
}
|
||||
DirectHttpResponse::HyperH2c(response) => {
|
||||
collect_hyper_stream_body(response, started_at, first_byte_timeout).await
|
||||
collect_hyper_stream_body(
|
||||
response,
|
||||
started_at,
|
||||
first_byte_timeout,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.await
|
||||
}
|
||||
DirectHttpResponse::BrowserWreq(response) => {
|
||||
collect_wreq_stream_body(response, started_at, first_byte_timeout).await
|
||||
collect_wreq_stream_body(
|
||||
response,
|
||||
started_at,
|
||||
first_byte_timeout,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1446,6 +1545,7 @@ async fn collect_reqwest_stream_body(
|
||||
response: reqwest::Response,
|
||||
started_at: Instant,
|
||||
first_byte_timeout: Option<Duration>,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut body_bytes = Vec::new();
|
||||
@@ -1466,7 +1566,11 @@ async fn collect_reqwest_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -1476,6 +1580,7 @@ async fn collect_hyper_stream_body(
|
||||
response: hyper::Response<HyperIncomingBody>,
|
||||
started_at: Instant,
|
||||
first_byte_timeout: Option<Duration>,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
|
||||
let mut stream = response.into_body().into_data_stream();
|
||||
let mut body_bytes = Vec::new();
|
||||
@@ -1496,7 +1601,11 @@ async fn collect_hyper_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -1506,6 +1615,7 @@ async fn collect_wreq_stream_body(
|
||||
response: wreq::Response,
|
||||
started_at: Instant,
|
||||
first_byte_timeout: Option<Duration>,
|
||||
response_body_limit_bytes: usize,
|
||||
) -> Result<(Bytes, Option<u64>), ExecutionRuntimeTransportError> {
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut body_bytes = Vec::new();
|
||||
@@ -1526,7 +1636,11 @@ async fn collect_wreq_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
response_body_limit_bytes,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -2378,10 +2492,31 @@ async fn send_via_tunnel_relay(
|
||||
error_kind = %kind,
|
||||
"gateway execution runtime tunnel relay returned relay error"
|
||||
);
|
||||
let message = response
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|_| format!("hub relay error: {kind}"));
|
||||
let response_headers = collect_response_headers(response.headers());
|
||||
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
|
||||
let (wire_body, _) =
|
||||
collect_reqwest_stream_body(response, Instant::now(), None, response_body_limit_bytes)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ExecutionRuntimeTransportError::RelayError(format!(
|
||||
"hub relay error: {kind}: bounded error body read failed: {error}"
|
||||
))
|
||||
})?;
|
||||
let decoded_body = decode_response_body_bytes_with_limit(
|
||||
&response_headers,
|
||||
&wire_body,
|
||||
response_body_limit_bytes,
|
||||
)
|
||||
.map_err(|error| {
|
||||
ExecutionRuntimeTransportError::RelayError(format!(
|
||||
"hub relay error: {kind}: bounded error body decode failed: {error}"
|
||||
))
|
||||
})?;
|
||||
let message = if decoded_body.is_empty() {
|
||||
format!("hub relay error: {kind}")
|
||||
} else {
|
||||
String::from_utf8_lossy(decoded_body.as_ref()).into_owned()
|
||||
};
|
||||
return Err(ExecutionRuntimeTransportError::RelayError(message));
|
||||
}
|
||||
|
||||
@@ -3903,6 +4038,7 @@ pub(crate) fn build_request_headers(
|
||||
|| normalized_key == EXECUTION_REQUEST_HTTP1_ONLY_HEADER
|
||||
|| normalized_key == EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER
|
||||
|| normalized_key == EXECUTION_RESPONSE_BODY_MODE_HEADER
|
||||
|| normalized_key == EXECUTION_RESPONSE_BODY_LIMIT_HEADER
|
||||
{
|
||||
continue;
|
||||
}
|
||||
@@ -4051,7 +4187,7 @@ pub(crate) fn decode_response_body_bytes<'a>(
|
||||
)
|
||||
}
|
||||
|
||||
fn decode_response_body_bytes_with_limit<'a>(
|
||||
pub(crate) fn decode_response_body_bytes_with_limit<'a>(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &'a [u8],
|
||||
limit_bytes: usize,
|
||||
@@ -4073,6 +4209,11 @@ fn decode_response_body_bytes_with_limit<'a>(
|
||||
read_upstream_response_decoder_with_limit("deflate", &mut decoder, limit_bytes)
|
||||
.map(Cow::Owned)
|
||||
}
|
||||
Some("br") => {
|
||||
let mut decoder = BrotliDecoder::new(body_bytes, 4_096);
|
||||
read_upstream_response_decoder_with_limit("br", &mut decoder, limit_bytes)
|
||||
.map(Cow::Owned)
|
||||
}
|
||||
_ => Ok(Cow::Borrowed(body_bytes)),
|
||||
}
|
||||
}
|
||||
@@ -4192,12 +4333,16 @@ mod tests {
|
||||
use super::{
|
||||
append_upstream_response_body_chunk_with_limit, build_browser_wreq_client, build_client,
|
||||
build_direct_tunnel_request_meta, build_execution_response_body, build_request_headers,
|
||||
decode_response_body_bytes_with_limit, execute_sync_plan, execution_response_body_mode,
|
||||
decode_response_body_bytes_with_limit, effective_response_body_limit_bytes,
|
||||
execute_sync_plan, execution_plan_response_body_limit_bytes, execution_response_body_mode,
|
||||
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
|
||||
record_manual_proxy_request_success, record_manual_proxy_stream_error,
|
||||
resolve_execution_transport_controls, resolve_non_stream_total_timeout,
|
||||
resolve_stream_first_byte_timeout, response_body_is_json, DirectSyncExecutionRuntime,
|
||||
resolve_stream_first_byte_timeout, response_body_is_json,
|
||||
with_upstream_response_body_limit, DirectSyncExecutionRuntime,
|
||||
ExecutionRuntimeTransportError, ExecutionTransportControls, UpstreamResponseBodyPhase,
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES, EXECUTION_RESPONSE_BODY_LIMIT_HEADER,
|
||||
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES, MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
};
|
||||
use crate::constants::{
|
||||
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
|
||||
@@ -4252,6 +4397,162 @@ mod tests {
|
||||
assert!(!materialized.contains_key("x-aether-future-control"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_response_body_limit_injection_preserves_transport_profile_and_extra() {
|
||||
let mut plan = tunnel_timeout_plan(false);
|
||||
let original_profile = ResolvedTransportProfile {
|
||||
profile_id: "existing-profile".into(),
|
||||
backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(),
|
||||
http_mode: TRANSPORT_HTTP_MODE_HTTP1_ONLY.into(),
|
||||
pool_scope: "provider".into(),
|
||||
header_fingerprint: Some(json!({"user_agent": "existing"})),
|
||||
extra: Some(json!({"existing": {"nested": true}})),
|
||||
};
|
||||
plan.transport_profile = Some(original_profile.clone());
|
||||
|
||||
let bounded_plan =
|
||||
with_upstream_response_body_limit(&plan, DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES);
|
||||
|
||||
assert_eq!(plan.transport_profile, Some(original_profile.clone()));
|
||||
assert_eq!(bounded_plan.transport_profile, Some(original_profile));
|
||||
assert_eq!(
|
||||
bounded_plan
|
||||
.headers
|
||||
.get(EXECUTION_RESPONSE_BODY_LIMIT_HEADER)
|
||||
.and_then(|value| value.parse::<usize>().ok()),
|
||||
Some(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES)
|
||||
);
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&bounded_plan),
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
|
||||
);
|
||||
|
||||
let unprofiled_plan = tunnel_timeout_plan(false);
|
||||
let bounded_unprofiled_plan = with_upstream_response_body_limit(
|
||||
&unprofiled_plan,
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
);
|
||||
assert!(unprofiled_plan.transport_profile.is_none());
|
||||
assert!(bounded_unprofiled_plan.transport_profile.is_none());
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&bounded_unprofiled_plan),
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
|
||||
);
|
||||
|
||||
let mut shadowed_plan = tunnel_timeout_plan(false);
|
||||
shadowed_plan.headers.insert(
|
||||
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_ascii_uppercase(),
|
||||
"65536".to_string(),
|
||||
);
|
||||
let bounded_shadowed_plan = with_upstream_response_body_limit(
|
||||
&shadowed_plan,
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
);
|
||||
assert_eq!(
|
||||
bounded_shadowed_plan
|
||||
.headers
|
||||
.keys()
|
||||
.filter(|name| name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_response_body_limit_parsing_rejects_invalid_values_and_clamps_bounds() {
|
||||
let scoped_plan = |raw_limit: &str| {
|
||||
let mut plan = tunnel_timeout_plan(false);
|
||||
plan.headers.insert(
|
||||
EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(),
|
||||
raw_limit.to_string(),
|
||||
);
|
||||
plan
|
||||
};
|
||||
|
||||
for invalid in ["0", "-1", "1.5", "", "invalid"] {
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&scoped_plan(invalid)),
|
||||
DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&scoped_plan("1")),
|
||||
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES
|
||||
);
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&scoped_plan(
|
||||
&(MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES as u64 + 1).to_string()
|
||||
)),
|
||||
MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES
|
||||
);
|
||||
assert_eq!(
|
||||
execution_plan_response_body_limit_bytes(&scoped_plan("1048576")),
|
||||
1_048_576
|
||||
);
|
||||
assert_eq!(
|
||||
effective_response_body_limit_bytes(
|
||||
Some(&(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES * 2).to_string()),
|
||||
1024 * 1024,
|
||||
),
|
||||
1024 * 1024,
|
||||
"a scoped limit must never raise the operator's global cap"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_response_body_wire_limit_rejects_overflow() {
|
||||
let bounded_plan = with_upstream_response_body_limit(
|
||||
&tunnel_timeout_plan(false),
|
||||
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
);
|
||||
let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan);
|
||||
let mut body = vec![b'x'; limit_bytes];
|
||||
|
||||
let error =
|
||||
append_upstream_response_body_chunk_with_limit(&mut body, b"overflow", limit_bytes)
|
||||
.expect_err("wire body above the plan-scoped limit should fail");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Wire,
|
||||
limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_response_body_limit_rejects_gzip_bomb_after_wire_check() {
|
||||
let bounded_plan = with_upstream_response_body_limit(
|
||||
&tunnel_timeout_plan(false),
|
||||
MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
);
|
||||
let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan);
|
||||
let payload = vec![b'x'; limit_bytes + 1];
|
||||
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
|
||||
encoder
|
||||
.write_all(&payload)
|
||||
.expect("gzip payload should encode");
|
||||
let encoded = encoder.finish().expect("gzip payload should finish");
|
||||
assert!(encoded.len() < limit_bytes);
|
||||
|
||||
let mut wire_body = Vec::new();
|
||||
append_upstream_response_body_chunk_with_limit(&mut wire_body, &encoded, limit_bytes)
|
||||
.expect("compressed wire body should fit within the plan-scoped limit");
|
||||
let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]);
|
||||
|
||||
let error = decode_response_body_bytes_with_limit(&headers, &wire_body, limit_bytes)
|
||||
.expect_err("decoded body above the plan-scoped limit should fail");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Decoded,
|
||||
limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upstream_response_wire_limit_allows_exact_body_and_rejects_next_byte() {
|
||||
let mut body = Vec::new();
|
||||
|
||||
@@ -1241,7 +1241,7 @@ async fn refresh_codex_provider_quota_locally_with_reset_fence(
|
||||
continue;
|
||||
}
|
||||
let persisted_key = state
|
||||
.read_provider_catalog_keys_by_ids(&[key.id.clone()])
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key.id))
|
||||
.await?
|
||||
.into_iter()
|
||||
.next();
|
||||
|
||||
@@ -22,6 +22,8 @@ pub(crate) use self::system_modules_helpers::{
|
||||
serialize_public_capability, supported_capability_names, PUBLIC_CAPABILITY_DEFINITIONS,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use self::support::filter_eligible_model_rows;
|
||||
pub(crate) use self::support::{
|
||||
build_api_key_install_session_response, build_proxy_node_install_session_response,
|
||||
build_unhandled_public_support_response, matches_model_mapping_for_models,
|
||||
|
||||
@@ -52,6 +52,8 @@ mod support_user_me;
|
||||
mod support_wallet;
|
||||
|
||||
pub(crate) use self::support_announcements::maybe_build_local_admin_announcements_response;
|
||||
#[cfg(test)]
|
||||
pub(crate) use self::support_models::filter_eligible_model_rows;
|
||||
pub(crate) use self::support_models::matches_model_mapping_for_models;
|
||||
|
||||
use self::support_announcements::{
|
||||
|
||||
@@ -10,6 +10,8 @@ mod models_route;
|
||||
mod models_shared;
|
||||
|
||||
pub(crate) use self::models_responses::build_models_auth_error_response;
|
||||
#[cfg(test)]
|
||||
pub(crate) use self::models_shared::filter_eligible_model_rows;
|
||||
pub(crate) use self::models_shared::{matches_model_mapping_for_models, models_api_format};
|
||||
|
||||
pub(super) async fn maybe_build_local_models_response(
|
||||
|
||||
@@ -102,8 +102,15 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Bod
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn build_codex_models_list_response(models: Vec<serde_json::Value>) -> Response<Body> {
|
||||
Json(json!({ "models": models })).into_response()
|
||||
pub(super) fn build_codex_models_list_response(
|
||||
models: Vec<serde_json::Value>,
|
||||
etag: Option<&str>,
|
||||
) -> Response<Body> {
|
||||
let mut response = Json(json!({ "models": models })).into_response();
|
||||
if let Some(etag) = etag.and_then(|value| http::HeaderValue::from_str(value).ok()) {
|
||||
response.headers_mut().insert(http::header::ETAG, etag);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
pub(super) fn build_openai_models_list_response(
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::BTreeSet;
|
||||
use std::fmt::Debug;
|
||||
use std::future::Future;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
@@ -25,8 +25,18 @@ use super::{query_param_value, AppState, GatewayPublicRequestContext};
|
||||
#[cfg(not(test))]
|
||||
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
#[cfg(test)]
|
||||
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50);
|
||||
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(1);
|
||||
const CODEX_MODELS_QUERY_API_FORMATS: &[&str] = &["openai:responses"];
|
||||
const CODEX_MODELS_MAX_RESPONSE_MODELS: usize = 512;
|
||||
const CODEX_MODELS_MAX_RESPONSE_JSON_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
fn codex_projected_catalog_fits_response_limits(cards: &[Value]) -> bool {
|
||||
if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS {
|
||||
return false;
|
||||
}
|
||||
serde_json::to_vec(&serde_json::json!({ "models": cards }))
|
||||
.is_ok_and(|body| body.len() <= CODEX_MODELS_MAX_RESPONSE_JSON_BYTES)
|
||||
}
|
||||
|
||||
async fn await_models_route_read<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T>
|
||||
where
|
||||
@@ -115,124 +125,130 @@ fn is_codex_provider_row(row: &StoredMinimalCandidateSelectionRow) -> bool {
|
||||
row.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
fn codex_model_card_is_complete(card: &serde_json::Map<String, Value>) -> bool {
|
||||
card.get("slug").and_then(Value::as_str).is_some()
|
||||
&& card.get("display_name").and_then(Value::as_str).is_some()
|
||||
&& card
|
||||
.get("supported_reasoning_levels")
|
||||
.and_then(Value::as_array)
|
||||
.is_some()
|
||||
&& card.get("shell_type").and_then(Value::as_str).is_some()
|
||||
&& card.get("visibility").and_then(Value::as_str).is_some()
|
||||
&& card
|
||||
.get("supported_in_api")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some()
|
||||
&& card.get("priority").and_then(Value::as_i64).is_some()
|
||||
&& card
|
||||
.get("base_instructions")
|
||||
.and_then(Value::as_str)
|
||||
.is_some()
|
||||
&& card
|
||||
.get("supports_reasoning_summary_parameter")
|
||||
.is_none_or(Value::is_boolean)
|
||||
&& card
|
||||
.get("support_verbosity")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some()
|
||||
&& card
|
||||
.get("truncation_policy")
|
||||
.and_then(Value::as_object)
|
||||
.is_some()
|
||||
&& card
|
||||
.get("supports_parallel_tool_calls")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some()
|
||||
&& card
|
||||
.get("experimental_supported_tools")
|
||||
.and_then(Value::as_array)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn project_codex_model_card(
|
||||
cached_models: &[Value],
|
||||
source_model: &str,
|
||||
global_model: &str,
|
||||
) -> Option<Value> {
|
||||
let mut card = cached_models
|
||||
.iter()
|
||||
.find(|model| {
|
||||
model.get("id").and_then(Value::as_str) == Some(source_model)
|
||||
|| model.get("slug").and_then(Value::as_str) == Some(source_model)
|
||||
})?
|
||||
.as_object()?
|
||||
.clone();
|
||||
if !codex_model_card_is_complete(&card) {
|
||||
return None;
|
||||
}
|
||||
|
||||
card.remove("id");
|
||||
card.remove("api_formats");
|
||||
card.insert("slug".to_string(), Value::String(global_model.to_string()));
|
||||
Some(Value::Object(card))
|
||||
}
|
||||
|
||||
async fn load_codex_model_cards(
|
||||
state: &AppState,
|
||||
rows: &[StoredMinimalCandidateSelectionRow],
|
||||
) -> Vec<Value> {
|
||||
let cache_keys = rows
|
||||
targets: &[crate::model_fetch::CodexCatalogTarget],
|
||||
client_version: crate::model_fetch::NormalizedCodexClientVersion,
|
||||
) -> (Vec<Value>, Option<String>) {
|
||||
let catalogs = crate::model_fetch::load_codex_catalogs(state, targets, &client_version).await;
|
||||
for target in catalogs.stale_targets() {
|
||||
let state = state.clone();
|
||||
let target = target.clone();
|
||||
let client_version = client_version.clone();
|
||||
tokio::spawn(async move {
|
||||
crate::model_fetch::refresh_codex_catalog_target(&state, &target, &client_version)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
if !catalogs.is_complete() {
|
||||
warn!(
|
||||
event_name = "codex_catalog_aggregate_incomplete",
|
||||
client_version = %client_version.as_str(),
|
||||
target_count = targets.len(),
|
||||
"Codex catalog aggregation was incomplete; returning an empty remote catalog so the client can use its bundled fallback"
|
||||
);
|
||||
return (Vec::new(), None);
|
||||
}
|
||||
let mut seen_global_models = BTreeSet::new();
|
||||
let possible_inference_catalogs = rows
|
||||
.iter()
|
||||
.filter(|row| is_codex_provider_row(row))
|
||||
.map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id))
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
let cached_values = await_models_route_read(
|
||||
"codex_models_cache",
|
||||
state.runtime_state.kv_get_many(&cache_keys),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let cached_models_by_key = cache_keys
|
||||
.into_iter()
|
||||
.zip(cached_values)
|
||||
.filter_map(|(key, raw)| {
|
||||
let models = serde_json::from_str::<Vec<Value>>(raw.as_deref()?).ok()?;
|
||||
Some((key, models))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
let mut seen_global_models = BTreeSet::new();
|
||||
.map(|row| (row.provider_id.clone(), row.key_id.clone()))
|
||||
.collect::<BTreeSet<_>>();
|
||||
let expected_global_models = rows
|
||||
.iter()
|
||||
.filter(|row| is_codex_provider_row(row))
|
||||
.map(|row| row.global_model_name.clone())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let mut cards = Vec::new();
|
||||
for row in rows.iter().filter(|row| is_codex_provider_row(row)) {
|
||||
if seen_global_models.contains(&row.global_model_name) {
|
||||
continue;
|
||||
}
|
||||
let cache_key = format!("upstream_models:{}:{}", row.provider_id, row.key_id);
|
||||
let Some(cached_models) = cached_models_by_key.get(&cache_key) else {
|
||||
let Some(snapshot) = catalogs.snapshot(&row.provider_id, &row.key_id) else {
|
||||
continue;
|
||||
};
|
||||
let source_model =
|
||||
aether_scheduler_core::select_provider_model_name(row, "openai:responses");
|
||||
let Some(card) = project_codex_model_card(
|
||||
cached_models,
|
||||
let Some(card) = crate::ai_serving::project_codex_catalog_model_card(
|
||||
&snapshot.models,
|
||||
source_model.as_str(),
|
||||
row.global_model_name.as_str(),
|
||||
) else {
|
||||
warn!(
|
||||
event_name = "codex_catalog_authorized_model_missing",
|
||||
provider_id = %row.provider_id,
|
||||
key_id = %row.key_id,
|
||||
client_version = %client_version.as_str(),
|
||||
source_model = %source_model,
|
||||
global_model = %row.global_model_name,
|
||||
"authorized Codex model was not present in this upstream catalog mapping"
|
||||
);
|
||||
continue;
|
||||
};
|
||||
seen_global_models.insert(row.global_model_name.clone());
|
||||
cards.push(card);
|
||||
if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS {
|
||||
warn!(
|
||||
event_name = "codex_catalog_aggregate_model_limit",
|
||||
client_version = %client_version.as_str(),
|
||||
model_count = cards.len(),
|
||||
limit = CODEX_MODELS_MAX_RESPONSE_MODELS,
|
||||
"Codex projected catalog exceeded the aggregate model limit; returning an empty remote catalog"
|
||||
);
|
||||
return (Vec::new(), None);
|
||||
}
|
||||
}
|
||||
cards
|
||||
let missing_model_count = expected_global_models
|
||||
.difference(&seen_global_models)
|
||||
.count();
|
||||
if missing_model_count > 0 {
|
||||
warn!(
|
||||
event_name = "codex_catalog_authorized_models_incomplete",
|
||||
client_version = %client_version.as_str(),
|
||||
expected_model_count = expected_global_models.len(),
|
||||
projected_model_count = cards.len(),
|
||||
missing_model_count,
|
||||
"Codex upstream catalogs omitted authorized mappings; returning an empty remote catalog so the client can use its bundled fallback"
|
||||
);
|
||||
return (Vec::new(), None);
|
||||
}
|
||||
if !codex_projected_catalog_fits_response_limits(&cards) {
|
||||
warn!(
|
||||
event_name = "codex_catalog_aggregate_body_limit",
|
||||
client_version = %client_version.as_str(),
|
||||
model_count = cards.len(),
|
||||
limit_bytes = CODEX_MODELS_MAX_RESPONSE_JSON_BYTES,
|
||||
"Codex projected catalog exceeded the aggregate response body limit; returning an empty remote catalog"
|
||||
);
|
||||
return (Vec::new(), None);
|
||||
}
|
||||
if cards.is_empty() {
|
||||
return (cards, None);
|
||||
}
|
||||
let etag = if possible_inference_catalogs.len() == 1 {
|
||||
possible_inference_catalogs
|
||||
.iter()
|
||||
.next()
|
||||
.and_then(|(provider_id, key_id)| catalogs.snapshot(provider_id, key_id))
|
||||
.and_then(|snapshot| snapshot.etag.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
(cards, etag)
|
||||
}
|
||||
|
||||
struct ModelRowsForClientFormat {
|
||||
rows: Vec<StoredMinimalCandidateSelectionRow>,
|
||||
codex_catalog_targets: Vec<crate::model_fetch::CodexCatalogTarget>,
|
||||
}
|
||||
|
||||
async fn list_model_rows_for_client_format(
|
||||
state: &AppState,
|
||||
api_format: &str,
|
||||
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
|
||||
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
|
||||
) -> Option<ModelRowsForClientFormat> {
|
||||
let mut collected = Vec::new();
|
||||
let query_api_formats = if is_codex_models_api_format(api_format) {
|
||||
CODEX_MODELS_QUERY_API_FORMATS
|
||||
@@ -254,9 +270,16 @@ async fn list_model_rows_for_client_format(
|
||||
}
|
||||
if is_codex_models_api_format(api_format) {
|
||||
collected.retain(is_codex_provider_row);
|
||||
Some(sort_model_rows(collected))
|
||||
let codex_catalog_targets = crate::model_fetch::codex_catalog_targets(&collected);
|
||||
Some(ModelRowsForClientFormat {
|
||||
rows: sort_model_rows(collected),
|
||||
codex_catalog_targets,
|
||||
})
|
||||
} else {
|
||||
Some(sort_and_dedup_model_rows(collected))
|
||||
Some(ModelRowsForClientFormat {
|
||||
rows: sort_and_dedup_model_rows(collected),
|
||||
codex_catalog_targets: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -296,6 +319,9 @@ pub(super) async fn maybe_build_local_models_route_response(
|
||||
}
|
||||
|
||||
let auth_context = decision.auth_context.as_ref()?;
|
||||
if !auth_context.access_allowed || auth_context.local_rejection.is_some() {
|
||||
return Some(build_models_auth_error_response(api_format));
|
||||
}
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
@@ -318,11 +344,23 @@ pub(super) async fn maybe_build_local_models_route_response(
|
||||
))
|
||||
}
|
||||
};
|
||||
let auth_snapshot = auth_snapshot.as_ref();
|
||||
let Some(auth_snapshot) = auth_snapshot.as_ref() else {
|
||||
warn!(
|
||||
event_name = "models_route_auth_snapshot_missing",
|
||||
user_id = %auth_context.user_id,
|
||||
api_key_id = %auth_context.api_key_id,
|
||||
"gateway models route rejected a request whose authenticated API key snapshot disappeared"
|
||||
);
|
||||
return Some(build_models_auth_error_response(api_format));
|
||||
};
|
||||
if !auth_snapshot.currently_usable {
|
||||
return Some(build_models_auth_error_response(api_format));
|
||||
}
|
||||
let auth_snapshot = Some(auth_snapshot);
|
||||
|
||||
match decision.route_kind.as_deref() {
|
||||
Some("list") => {
|
||||
let rows =
|
||||
let listed =
|
||||
match list_model_rows_for_client_format(state, api_format, auth_snapshot).await {
|
||||
Some(rows) => rows,
|
||||
None => {
|
||||
@@ -332,12 +370,34 @@ pub(super) async fn maybe_build_local_models_route_response(
|
||||
))
|
||||
}
|
||||
};
|
||||
let rows = listed.rows;
|
||||
if rows.is_empty() {
|
||||
return Some(build_empty_models_list_response(api_format));
|
||||
}
|
||||
if is_codex_models_api_format(api_format) {
|
||||
let models = load_codex_model_cards(state, &rows).await;
|
||||
return Some(build_codex_models_list_response(models));
|
||||
let raw_client_version = query_param_value(
|
||||
request_context.request_query_string.as_deref(),
|
||||
"client_version",
|
||||
);
|
||||
let client_version = crate::model_fetch::normalize_codex_client_version(
|
||||
raw_client_version.as_deref(),
|
||||
);
|
||||
if client_version.used_fallback() {
|
||||
warn!(
|
||||
event_name = "codex_catalog_invalid_client_version",
|
||||
raw_length = raw_client_version.as_ref().map_or(0, String::len),
|
||||
fallback_version = %client_version.as_str(),
|
||||
"invalid Codex client_version used the bounded fallback version"
|
||||
);
|
||||
}
|
||||
let (models, etag) = load_codex_model_cards(
|
||||
state,
|
||||
&rows,
|
||||
&listed.codex_catalog_targets,
|
||||
client_version,
|
||||
)
|
||||
.await;
|
||||
return Some(build_codex_models_list_response(models, etag.as_deref()));
|
||||
}
|
||||
let response = match api_format {
|
||||
"claude:messages" => {
|
||||
@@ -410,3 +470,32 @@ pub(super) async fn maybe_build_local_models_route_response(
|
||||
_ => Some(build_models_auth_error_response(api_format)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
codex_projected_catalog_fits_response_limits, CODEX_MODELS_MAX_RESPONSE_JSON_BYTES,
|
||||
CODEX_MODELS_MAX_RESPONSE_MODELS,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn projected_codex_catalog_enforces_aggregate_count_and_body_limits() {
|
||||
assert!(codex_projected_catalog_fits_response_limits(&[json!({
|
||||
"slug": "gpt-future-dynamic",
|
||||
"model_messages": {"instructions_template": "opaque"}
|
||||
})]));
|
||||
|
||||
let too_many = (0..=CODEX_MODELS_MAX_RESPONSE_MODELS)
|
||||
.map(|index| json!({"slug": format!("gpt-future-{index}")}))
|
||||
.collect::<Vec<_>>();
|
||||
assert!(!codex_projected_catalog_fits_response_limits(&too_many));
|
||||
|
||||
let oversized = vec![json!({
|
||||
"slug": "gpt-future-oversized",
|
||||
"future_capability": "x".repeat(CODEX_MODELS_MAX_RESPONSE_JSON_BYTES)
|
||||
})];
|
||||
assert!(!codex_projected_catalog_fits_response_limits(&oversized));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -194,7 +194,7 @@ fn row_exposes_global_model_for_models(
|
||||
false
|
||||
}
|
||||
|
||||
pub(super) fn filter_eligible_model_rows(
|
||||
pub(crate) fn filter_eligible_model_rows(
|
||||
rows: Vec<StoredMinimalCandidateSelectionRow>,
|
||||
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
|
||||
api_format: &str,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,8 +1,15 @@
|
||||
mod catalog;
|
||||
mod runtime;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub(crate) use aether_model_fetch::ModelFetchRunSummary;
|
||||
pub(crate) use catalog::{
|
||||
codex_catalog_credential_scope_from_stored_key, codex_catalog_targets, load_codex_catalogs,
|
||||
normalize_codex_client_version, read_recent_codex_catalog_client_version,
|
||||
refresh_codex_catalog_target, CodexCatalogLoad, CodexCatalogRuntime, CodexCatalogTarget,
|
||||
NormalizedCodexClientVersion,
|
||||
};
|
||||
pub(crate) use runtime::state::ModelFetchRuntimeState;
|
||||
pub(crate) use runtime::{
|
||||
perform_model_fetch_for_key, perform_model_fetch_for_keys, perform_model_fetch_once,
|
||||
|
||||
@@ -7,7 +7,7 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_model_fetch::{
|
||||
apply_model_filters, fetch_models_from_transports, json_string_list,
|
||||
apply_model_filters, fetch_models_from_transports_for_client_version, json_string_list,
|
||||
model_catalog_upstream_metadata, model_fetch_interval_minutes,
|
||||
model_fetch_startup_delay_seconds, model_fetch_startup_enabled, preset_models_for_provider,
|
||||
selected_models_fetch_endpoints, sync_provider_model_whitelist_associations,
|
||||
@@ -315,7 +315,25 @@ async fn fetch_and_persist_key_models(
|
||||
return Ok(KeyFetchDisposition::Skipped);
|
||||
}
|
||||
|
||||
let result = match fetch_models_from_transports(state, &transports).await {
|
||||
let codex_client_version = if target
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
{
|
||||
state
|
||||
.read_recent_codex_catalog_client_version(&target.provider.id, &target.key.id)
|
||||
.await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let result = match fetch_models_from_transports_for_client_version(
|
||||
state,
|
||||
&transports,
|
||||
codex_client_version.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
persist_key_fetch_failure(state, &target.key, now_unix_secs, err.clone()).await?;
|
||||
|
||||
@@ -43,6 +43,14 @@ pub(crate) trait ModelFetchRuntimeState:
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, GatewayError>;
|
||||
|
||||
async fn read_recent_codex_catalog_client_version(
|
||||
&self,
|
||||
_provider_id: &str,
|
||||
_key_id: &str,
|
||||
) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn update_provider_catalog_key_model_fetch_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
|
||||
@@ -3838,7 +3838,9 @@ mod tests {
|
||||
.expect("overflow should resume after repository recovery")
|
||||
.expect("overflow task should complete");
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while runtime.metrics.pending_current.load(Ordering::Acquire) != 0 {
|
||||
while runtime.metrics.pending_current.load(Ordering::Acquire) != 0
|
||||
|| runtime.priority_admission.available_permits() != 2
|
||||
{
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
|
||||
@@ -480,6 +480,16 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids_strong(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
|
||||
self.data
|
||||
.list_provider_catalog_keys_by_ids_strong(key_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_page(
|
||||
&self,
|
||||
query: &provider_catalog::ProviderCatalogKeyListQuery,
|
||||
|
||||
@@ -25,7 +25,7 @@ use tracing::{debug, warn};
|
||||
|
||||
use super::{AppState, GatewayError};
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::model_fetch::ModelFetchRuntimeState;
|
||||
use crate::model_fetch::{CodexCatalogRuntime, ModelFetchRuntimeState};
|
||||
use crate::provider_transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
|
||||
use crate::request_candidate_runtime::{
|
||||
RequestCandidateRuntimeCapabilityReader, RequestCandidateRuntimeReader,
|
||||
@@ -34,6 +34,8 @@ use crate::request_candidate_runtime::{
|
||||
use crate::scheduler::state::SchedulerRuntimeState;
|
||||
use crate::{execution_runtime, provider_transport};
|
||||
|
||||
const MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
impl AppState {
|
||||
pub(crate) async fn hydrate_antigravity_project_metadata_for_transport(
|
||||
&self,
|
||||
@@ -274,12 +276,90 @@ impl ModelFetchTransportRuntime for AppState {
|
||||
&self,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, String> {
|
||||
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan)
|
||||
let bounded_plan = execution_runtime::transport::with_upstream_response_body_limit(
|
||||
plan,
|
||||
MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES,
|
||||
);
|
||||
execution_runtime::execute_execution_runtime_sync_plan(self, None, &bounded_plan)
|
||||
.await
|
||||
.map_err(GatewayError::into_message)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CodexCatalogRuntime for AppState {
|
||||
fn codex_catalog_runtime_state(&self) -> &aether_runtime_state::RuntimeState {
|
||||
self.runtime_state.as_ref()
|
||||
}
|
||||
|
||||
async fn read_codex_catalog_transport_snapshot(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
|
||||
self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
|
||||
.await
|
||||
.map_err(GatewayError::into_message)
|
||||
}
|
||||
|
||||
async fn read_codex_catalog_credential_scope_strong(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
key_id: &str,
|
||||
) -> Result<Option<String>, String> {
|
||||
let Some(key) = self
|
||||
.list_provider_catalog_keys_by_ids_strong(&[key_id.to_string()])
|
||||
.await
|
||||
.map_err(GatewayError::into_message)?
|
||||
.into_iter()
|
||||
.find(|key| key.id == key_id && key.provider_id == provider_id && key.is_active)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Some(scope) =
|
||||
crate::model_fetch::codex_catalog_credential_scope_from_stored_key(&key, None, None)
|
||||
{
|
||||
return Ok(Some(scope));
|
||||
}
|
||||
|
||||
let decrypted_auth_config = match key.encrypted_auth_config.as_deref() {
|
||||
Some(ciphertext) => Some(
|
||||
crate::handlers::shared::decrypt_catalog_secret_with_fallbacks(
|
||||
self.encryption_key(),
|
||||
ciphertext,
|
||||
)
|
||||
.ok_or_else(|| {
|
||||
"Codex catalog auth config could not be verified for credential fencing"
|
||||
.to_string()
|
||||
})?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
let decrypted_api_key = match key.encrypted_api_key.as_deref() {
|
||||
Some(ciphertext) => Some(
|
||||
crate::handlers::shared::decrypt_catalog_secret_with_fallbacks(
|
||||
self.encryption_key(),
|
||||
ciphertext,
|
||||
)
|
||||
.ok_or_else(|| {
|
||||
"Codex catalog API key could not be verified for credential fencing".to_string()
|
||||
})?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
|
||||
Ok(
|
||||
crate::model_fetch::codex_catalog_credential_scope_from_stored_key(
|
||||
&key,
|
||||
decrypted_auth_config.as_deref(),
|
||||
decrypted_api_key.as_deref(),
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ModelFetchRuntimeState for AppState {
|
||||
fn has_provider_catalog_data_reader(&self) -> bool {
|
||||
@@ -320,6 +400,29 @@ impl ModelFetchRuntimeState for AppState {
|
||||
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan).await
|
||||
}
|
||||
|
||||
async fn read_recent_codex_catalog_client_version(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
key_id: &str,
|
||||
) -> Option<String> {
|
||||
let credential_scope =
|
||||
<AppState as CodexCatalogRuntime>::read_codex_catalog_credential_scope_strong(
|
||||
self,
|
||||
provider_id,
|
||||
key_id,
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()?;
|
||||
crate::model_fetch::read_recent_codex_catalog_client_version(
|
||||
self.runtime_state.as_ref(),
|
||||
provider_id,
|
||||
key_id,
|
||||
&credential_scope,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_provider_catalog_key_model_fetch_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
|
||||
@@ -8266,24 +8266,21 @@ async fn gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over
|
||||
"Bearer cached-codex-access-token"
|
||||
);
|
||||
|
||||
let mut updated_key = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-codex-oauth-stale-cache".to_string()])
|
||||
.await
|
||||
.expect("keys should list")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("key should exist");
|
||||
updated_key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#,
|
||||
let fresh_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#,
|
||||
)
|
||||
.expect("updated auth config ciphertext should build");
|
||||
assert!(provider_catalog_repository
|
||||
.update_key_oauth_runtime_state(
|
||||
"key-codex-oauth-stale-cache",
|
||||
None,
|
||||
None,
|
||||
Some(&fresh_auth_config),
|
||||
Some(4_102_444_810),
|
||||
)
|
||||
.expect("updated auth config ciphertext should build"),
|
||||
);
|
||||
provider_catalog_repository
|
||||
.update_key(&updated_key)
|
||||
.await
|
||||
.expect("key should update");
|
||||
.expect("OAuth runtime state should update"));
|
||||
|
||||
let gateway = build_router_with_state(app_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
@@ -7,11 +7,16 @@ use super::{
|
||||
use crate::image_capabilities::openai_image_gateway_max_generation_count;
|
||||
use crate::tests::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
|
||||
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
to_bytes, wait_until, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_AI_PUBLIC,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
};
|
||||
use aether_contracts::{ExecutionResult, ExecutionTelemetry, ResponseBody};
|
||||
use aether_crypto::encrypt_python_fernet_plaintext;
|
||||
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::DataLayerError;
|
||||
use aether_data_contracts::repository::auth::AuthApiKeyWriteRepository;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
@@ -20,13 +25,21 @@ use aether_data_contracts::repository::candidate_selection::{
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
StoredAdminGlobalModel, UpdateAdminGlobalModelRecord,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use axum::response::IntoResponse;
|
||||
use std::collections::HashMap;
|
||||
use std::future::pending;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
|
||||
fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
fn codex_models_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
allowed_models: &[&str],
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
@@ -37,7 +50,7 @@ fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySna
|
||||
false,
|
||||
Some(json!(["codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
Some(json!(["frontier-sol", "broken-luna"])),
|
||||
Some(json!(allowed_models)),
|
||||
api_key_id.to_string(),
|
||||
Some("codex-models".to_string()),
|
||||
true,
|
||||
@@ -48,7 +61,7 @@ fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySna
|
||||
Some(4_102_444_800),
|
||||
Some(json!(["codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
Some(json!(["frontier-sol", "broken-luna"])),
|
||||
Some(json!(allowed_models)),
|
||||
)
|
||||
.expect("Codex models auth snapshot should build")
|
||||
}
|
||||
@@ -104,6 +117,7 @@ fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value {
|
||||
"upgrade": null,
|
||||
"base_instructions": "Use the current Codex instructions.",
|
||||
"model_messages": null,
|
||||
"available_in_plans": ["plus", "pro"],
|
||||
"support_verbosity": true,
|
||||
"default_verbosity": "low",
|
||||
"apply_patch_tool_type": "freeform",
|
||||
@@ -115,6 +129,105 @@ fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value {
|
||||
})
|
||||
}
|
||||
|
||||
fn codex_catalog_provider(provider_id: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
provider_id.to_string(),
|
||||
"codex".to_string(),
|
||||
Some("https://chatgpt.com".to_string()),
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("Codex provider should build")
|
||||
}
|
||||
|
||||
fn codex_catalog_endpoint(provider_id: &str, endpoint_id: &str) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
endpoint_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"openai:responses".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("responses".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("Codex endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://chatgpt.example/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("Codex endpoint transport should build")
|
||||
}
|
||||
|
||||
fn codex_catalog_key(
|
||||
provider_id: &str,
|
||||
key_id: &str,
|
||||
allowed_models: &[&str],
|
||||
) -> StoredProviderCatalogKey {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
key_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"manual".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("Codex key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["openai:responses"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "oauth-upstream-secret")
|
||||
.expect("Codex test token should encrypt"),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(json!(allowed_models)),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("Codex key transport should build");
|
||||
key.auto_fetch_models = false;
|
||||
key.locked_models = Some(json!(["manual-locked-model"]));
|
||||
key.model_include_patterns = Some(json!(["gpt-future-*"]));
|
||||
key.model_exclude_patterns = Some(json!(["gpt-future-denied"]));
|
||||
key
|
||||
}
|
||||
|
||||
fn codex_catalog_execution_result(
|
||||
plan: &aether_contracts::ExecutionPlan,
|
||||
status_code: u16,
|
||||
body: serde_json::Value,
|
||||
etag: Option<&str>,
|
||||
) -> ExecutionResult {
|
||||
let mut headers = std::collections::BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]);
|
||||
if let Some(etag) = etag {
|
||||
headers.insert("ETag".to_string(), etag.to_string());
|
||||
}
|
||||
ExecutionResult {
|
||||
request_id: plan.request_id.clone(),
|
||||
candidate_id: plan.candidate_id.clone(),
|
||||
status_code,
|
||||
headers,
|
||||
response_observation: None,
|
||||
body: Some(ResponseBody {
|
||||
json_body: Some(body),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: Some(ExecutionTelemetry {
|
||||
ttfb_ms: Some(1),
|
||||
elapsed_ms: Some(2),
|
||||
upstream_bytes: None,
|
||||
}),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
|
||||
match status {
|
||||
VideoTaskStatus::Pending => "Pending",
|
||||
@@ -449,102 +562,618 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
|
||||
fallback_probe_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_serves_codex_model_cards_for_versioned_models_requests() {
|
||||
let codex_row =
|
||||
sample_codex_models_candidate_row("provider-codex-models", "frontier-sol", "gpt-5.6-sol");
|
||||
let incomplete_codex_row = sample_codex_models_candidate_row(
|
||||
"provider-codex-incomplete",
|
||||
"broken-luna",
|
||||
"gpt-5.6-luna",
|
||||
);
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
codex_row.clone(),
|
||||
incomplete_codex_row.clone(),
|
||||
sample_models_candidate_row(
|
||||
"provider-openai-responses",
|
||||
"openai",
|
||||
"openai:responses",
|
||||
"custom-responses-model",
|
||||
20,
|
||||
#[test]
|
||||
fn gateway_versioned_models_fail_closed_when_cached_auth_becomes_unusable_or_missing() {
|
||||
std::thread::Builder::new()
|
||||
.name("codex-model-catalog-auth-race".to_string())
|
||||
.stack_size(16 * 1024 * 1024)
|
||||
.spawn(|| {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("Codex auth race test runtime should build")
|
||||
.block_on(run_versioned_models_auth_race_scenario());
|
||||
})
|
||||
.expect("Codex auth race test thread should spawn")
|
||||
.join()
|
||||
.expect("Codex auth race test thread should finish");
|
||||
}
|
||||
|
||||
async fn run_versioned_models_auth_race_scenario() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-codex-models-auth-race")),
|
||||
codex_models_snapshot(
|
||||
"key-codex-models-auth-race",
|
||||
"user-codex-models-auth-race",
|
||||
&["future-alias"],
|
||||
),
|
||||
)]));
|
||||
let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(
|
||||
Vec::new(),
|
||||
));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
|
||||
candidate_repository,
|
||||
auth_repository.clone(),
|
||||
),
|
||||
),
|
||||
]));
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
let models_url = format!("{gateway_url}/v1/models?client_version=0.145.2");
|
||||
|
||||
let warm_response = client
|
||||
.get(&models_url)
|
||||
.header("authorization", "Bearer sk-codex-models-auth-race")
|
||||
.send()
|
||||
.await
|
||||
.expect("initial versioned models request should succeed");
|
||||
assert_eq!(warm_response.status(), StatusCode::OK);
|
||||
|
||||
assert!(auth_repository
|
||||
.set_user_api_key_locked(
|
||||
"user-codex-models-auth-race",
|
||||
"key-codex-models-auth-race",
|
||||
true,
|
||||
)
|
||||
.await
|
||||
.expect("locking the API key should succeed"));
|
||||
let locked_response = client
|
||||
.get(&models_url)
|
||||
.header("authorization", "Bearer sk-codex-models-auth-race")
|
||||
.send()
|
||||
.await
|
||||
.expect("locked versioned models request should complete");
|
||||
assert_eq!(locked_response.status(), StatusCode::UNAUTHORIZED);
|
||||
|
||||
assert!(auth_repository
|
||||
.delete_user_api_key("user-codex-models-auth-race", "key-codex-models-auth-race",)
|
||||
.await
|
||||
.expect("deleting the API key should succeed"));
|
||||
let missing_response = client
|
||||
.get(&models_url)
|
||||
.header("authorization", "Bearer sk-codex-models-auth-race")
|
||||
.send()
|
||||
.await
|
||||
.expect("missing versioned models request should complete");
|
||||
assert_eq!(missing_response.status(), StatusCode::UNAUTHORIZED);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_serves_codex_model_cards_for_versioned_models_requests() {
|
||||
std::thread::Builder::new()
|
||||
.name("codex-model-catalog-frontdoor".to_string())
|
||||
.stack_size(16 * 1024 * 1024)
|
||||
.spawn(|| {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("Codex frontdoor test runtime should build")
|
||||
.block_on(run_versioned_codex_model_cards_frontdoor_scenario());
|
||||
})
|
||||
.expect("Codex frontdoor test thread should spawn")
|
||||
.join()
|
||||
.expect("Codex frontdoor test thread should finish");
|
||||
}
|
||||
|
||||
async fn run_versioned_codex_model_cards_frontdoor_scenario() {
|
||||
const PROVIDER_ID: &str = "provider-codex-models";
|
||||
const CATALOG_KEY_ID: &str = "key-provider-codex-models";
|
||||
const CATALOG_ENDPOINT_ID: &str = "endpoint-provider-codex-models";
|
||||
const SOURCE_MODELS: &[&str] = &[
|
||||
"gpt-future-dynamic",
|
||||
"gpt-future-legacy",
|
||||
"gpt-future-second",
|
||||
"gpt-hidden-direct",
|
||||
];
|
||||
const GLOBAL_MODELS: &[&str] = &["future-alias", "legacy-alias"];
|
||||
|
||||
let mut codex_rows = vec![
|
||||
sample_codex_models_candidate_row(PROVIDER_ID, "future-alias", "gpt-future-dynamic"),
|
||||
sample_codex_models_candidate_row(PROVIDER_ID, "legacy-alias", "gpt-future-legacy"),
|
||||
sample_codex_models_candidate_row(PROVIDER_ID, "second-alias", "gpt-future-second"),
|
||||
sample_codex_models_candidate_row(PROVIDER_ID, "hidden-alias", "gpt-hidden-direct"),
|
||||
];
|
||||
for row in &mut codex_rows {
|
||||
row.key_allowed_models = Some(
|
||||
SOURCE_MODELS
|
||||
.iter()
|
||||
.map(|value| value.to_string())
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
let mut all_candidate_rows = codex_rows.clone();
|
||||
all_candidate_rows.push(sample_models_candidate_row(
|
||||
"provider-openai-responses",
|
||||
"openai",
|
||||
"openai:responses",
|
||||
"custom-responses-model",
|
||||
20,
|
||||
));
|
||||
let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(
|
||||
all_candidate_rows,
|
||||
));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![
|
||||
(
|
||||
Some(hash_api_key("sk-codex-models")),
|
||||
codex_models_snapshot("key-codex-models", "user-codex-models"),
|
||||
codex_models_snapshot("key-codex-models", "user-codex-models", GLOBAL_MODELS),
|
||||
),
|
||||
(
|
||||
Some(hash_api_key("sk-standard-models")),
|
||||
unrestricted_models_snapshot("key-standard-models", "user-standard-models"),
|
||||
),
|
||||
(
|
||||
Some(hash_api_key("sk-codex-legacy-only")),
|
||||
codex_models_snapshot(
|
||||
"key-codex-legacy-only",
|
||||
"user-codex-legacy-only",
|
||||
&["legacy-alias"],
|
||||
),
|
||||
),
|
||||
(
|
||||
Some(hash_api_key("sk-codex-hidden-mixed")),
|
||||
codex_models_snapshot(
|
||||
"key-codex-hidden-mixed",
|
||||
"user-codex-hidden-mixed",
|
||||
&["future-alias", "hidden-alias"],
|
||||
),
|
||||
),
|
||||
(
|
||||
Some(hash_api_key("sk-codex-second-mixed")),
|
||||
codex_models_snapshot(
|
||||
"key-codex-second-mixed",
|
||||
"user-codex-second-mixed",
|
||||
&["future-alias", "second-alias"],
|
||||
),
|
||||
),
|
||||
]));
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
let original_catalog_key = codex_catalog_key(PROVIDER_ID, CATALOG_KEY_ID, SOURCE_MODELS);
|
||||
assert!(!original_catalog_key.auto_fetch_models);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_catalog_provider(PROVIDER_ID)],
|
||||
vec![codex_catalog_endpoint(PROVIDER_ID, CATALOG_ENDPOINT_ID)],
|
||||
vec![original_catalog_key.clone()],
|
||||
));
|
||||
let rows_before = candidate_repository
|
||||
.list_for_exact_api_format("openai:responses")
|
||||
.await
|
||||
.expect("candidate rows should load before request");
|
||||
|
||||
let catalog_generation = Arc::new(AtomicUsize::new(0));
|
||||
let catalog_hits = Arc::new(AtomicUsize::new(0));
|
||||
let captured_plans = Arc::new(Mutex::new(Vec::<(String, Option<String>)>::new()));
|
||||
let generation_for_runtime = Arc::clone(&catalog_generation);
|
||||
let hits_for_runtime = Arc::clone(&catalog_hits);
|
||||
let plans_for_runtime = Arc::clone(&captured_plans);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let generation_for_request = Arc::clone(&generation_for_runtime);
|
||||
let hits_for_request = Arc::clone(&hits_for_runtime);
|
||||
let plans_for_request = Arc::clone(&plans_for_runtime);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX)
|
||||
.await
|
||||
.expect("execution runtime request body should read");
|
||||
let plan: aether_contracts::ExecutionPlan =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime plan should parse");
|
||||
plans_for_request
|
||||
.lock()
|
||||
.expect("plans mutex")
|
||||
.push((plan.url.clone(), plan.headers.get("user-agent").cloned()));
|
||||
if plan.url.contains("/models?") {
|
||||
hits_for_request.fetch_add(1, Ordering::SeqCst);
|
||||
let generation = generation_for_request.load(Ordering::SeqCst);
|
||||
if generation == 1 {
|
||||
return Json(codex_catalog_execution_result(
|
||||
&plan,
|
||||
503,
|
||||
json!({"error":{"message":"temporary catalog outage"}}),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
let mut current = complete_codex_model_card("gpt-future-dynamic");
|
||||
let current_object = current.as_object_mut().expect("current card object");
|
||||
current_object.remove("base_instructions");
|
||||
current_object.insert(
|
||||
"model_messages".to_string(),
|
||||
json!({"instructions_template":"Use future dynamic instructions."}),
|
||||
);
|
||||
current_object.insert(
|
||||
"future_capability".to_string(),
|
||||
json!({"mode":"opaque-current"}),
|
||||
);
|
||||
|
||||
let mut legacy = complete_codex_model_card("gpt-future-legacy");
|
||||
legacy["base_instructions"] = json!("Use legacy future instructions.");
|
||||
legacy["future_capability"] = json!({"mode":"opaque-legacy"});
|
||||
|
||||
let mut models = vec![current, legacy];
|
||||
if generation >= 2 {
|
||||
let mut second = complete_codex_model_card("gpt-future-second");
|
||||
second
|
||||
.as_object_mut()
|
||||
.expect("second card object")
|
||||
.remove("base_instructions");
|
||||
second["model_messages"] =
|
||||
json!({"instructions_template":"Use second future instructions."});
|
||||
second["future_capability"] = json!({"mode":"added-without-code-change"});
|
||||
models.push(second);
|
||||
models.push(complete_codex_model_card("gpt-future-unmapped"));
|
||||
}
|
||||
let etag = if generation >= 2 {
|
||||
"\"catalog-etag-v2\""
|
||||
} else {
|
||||
"\"catalog-etag-v1\""
|
||||
};
|
||||
return Json(codex_catalog_execution_result(
|
||||
&plan,
|
||||
200,
|
||||
json!({"models": models, "future_top_level": true}),
|
||||
Some(etag),
|
||||
));
|
||||
}
|
||||
|
||||
Json(codex_catalog_execution_result(
|
||||
&plan,
|
||||
200,
|
||||
json!({
|
||||
"id": "resp-future-dynamic",
|
||||
"object": "response",
|
||||
"model": "gpt-future-dynamic",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}),
|
||||
Some("\"catalog-etag-v2\""),
|
||||
))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
|
||||
candidate_repository,
|
||||
candidate_repository.clone(),
|
||||
auth_repository,
|
||||
),
|
||||
)
|
||||
.attach_provider_catalog_repository_for_tests(provider_catalog_repository.clone())
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
);
|
||||
state
|
||||
.runtime_kv_setex(
|
||||
&format!(
|
||||
"upstream_models:{}:{}",
|
||||
codex_row.provider_id, codex_row.key_id
|
||||
),
|
||||
&serde_json::to_string(&vec![complete_codex_model_card("gpt-5.6-sol")])
|
||||
.expect("model cache should serialize"),
|
||||
60,
|
||||
|
||||
let configured_rows = state
|
||||
.list_minimal_candidate_selection_rows_for_api_format("openai:responses")
|
||||
.await
|
||||
.expect("configured Codex candidate rows should be readable");
|
||||
assert_eq!(
|
||||
configured_rows
|
||||
.iter()
|
||||
.filter(|row| row.provider_type == "codex")
|
||||
.count(),
|
||||
codex_rows.len()
|
||||
);
|
||||
let resolved_auth =
|
||||
aether_data_contracts::repository::auth::ResolvedAuthApiKeySnapshot::from_stored(
|
||||
codex_models_snapshot("key-codex-models", "user-codex-models", GLOBAL_MODELS),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
);
|
||||
let eligible_rows = crate::handlers::public::filter_eligible_model_rows(
|
||||
configured_rows.clone(),
|
||||
Some(&resolved_auth),
|
||||
"openai:responses",
|
||||
);
|
||||
assert_eq!(
|
||||
eligible_rows
|
||||
.iter()
|
||||
.filter(|row| row.provider_type == "codex")
|
||||
.count(),
|
||||
GLOBAL_MODELS.len(),
|
||||
"Codex fixture rows must survive the same provider/model/key authorization filters as the route"
|
||||
);
|
||||
let actual_auth = state
|
||||
.data
|
||||
.read_auth_api_key_snapshot(
|
||||
"user-codex-models",
|
||||
"key-codex-models",
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
)
|
||||
.await
|
||||
.expect("model cache should seed");
|
||||
state
|
||||
.runtime_kv_setex(
|
||||
&format!(
|
||||
"upstream_models:{}:{}",
|
||||
incomplete_codex_row.provider_id, incomplete_codex_row.key_id
|
||||
),
|
||||
&serde_json::to_string(&vec![json!({
|
||||
"id": "gpt-5.6-luna",
|
||||
"slug": "gpt-5.6-luna",
|
||||
"display_name": "GPT-5.6-Luna"
|
||||
})])
|
||||
.expect("incomplete model cache should serialize"),
|
||||
60,
|
||||
.expect("Codex auth snapshot read should succeed")
|
||||
.expect("Codex auth snapshot should exist");
|
||||
assert_eq!(
|
||||
crate::handlers::public::filter_eligible_model_rows(
|
||||
configured_rows.clone(),
|
||||
Some(&actual_auth),
|
||||
"openai:responses",
|
||||
)
|
||||
.iter()
|
||||
.filter(|row| row.provider_type == "codex")
|
||||
.count(),
|
||||
GLOBAL_MODELS.len(),
|
||||
"stored Codex auth snapshot must preserve every authorized manual mapping"
|
||||
);
|
||||
assert!(
|
||||
<AppState as crate::model_fetch::CodexCatalogRuntime>::read_codex_catalog_transport_snapshot(
|
||||
&state,
|
||||
PROVIDER_ID,
|
||||
CATALOG_ENDPOINT_ID,
|
||||
CATALOG_KEY_ID,
|
||||
)
|
||||
.await
|
||||
.expect("incomplete model cache should seed");
|
||||
.expect("Codex catalog transport lookup should succeed")
|
||||
.is_some(),
|
||||
"Codex catalog transport must be available even when auto_fetch_models is disabled"
|
||||
);
|
||||
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let codex_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.144.1"))
|
||||
.get(format!(
|
||||
"{gateway_url}/v1/models?client_version=0.145.2-beta.7%2Bdesktop.9"
|
||||
))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("Codex models request should succeed");
|
||||
assert_eq!(codex_response.status(), StatusCode::OK);
|
||||
let codex_etag = codex_response
|
||||
.headers()
|
||||
.get(http::header::ETAG)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
let codex_payload: serde_json::Value = codex_response
|
||||
.json()
|
||||
.await
|
||||
.expect("Codex models body should parse");
|
||||
assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(codex_payload["models"][0]["slug"], "frontier-sol");
|
||||
assert_eq!(
|
||||
codex_payload["models"][0]["supported_reasoning_levels"][5]["effort"],
|
||||
"ultra"
|
||||
catalog_hits.load(Ordering::SeqCst),
|
||||
1,
|
||||
"cold Codex request must reach the upstream catalog once; payload={codex_payload}"
|
||||
);
|
||||
assert_eq!(codex_etag.as_deref(), Some("\"catalog-etag-v1\""));
|
||||
assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(2));
|
||||
let current_card = codex_payload["models"]
|
||||
.as_array()
|
||||
.and_then(|models| models.iter().find(|model| model["slug"] == "future-alias"))
|
||||
.expect("current model card should be projected");
|
||||
assert_eq!(
|
||||
current_card["model_messages"]["instructions_template"],
|
||||
"Use future dynamic instructions."
|
||||
);
|
||||
assert_eq!(
|
||||
codex_payload["models"][0]["future_capability"],
|
||||
json!({"enabled": true})
|
||||
current_card["future_capability"],
|
||||
json!({"mode":"opaque-current"})
|
||||
);
|
||||
assert!(codex_payload["models"][0].get("id").is_none());
|
||||
assert!(codex_payload["models"][0].get("api_formats").is_none());
|
||||
assert_eq!(current_card["available_in_plans"], json!(["plus", "pro"]));
|
||||
assert!(current_card.get("base_instructions").is_none());
|
||||
assert!(current_card.get("id").is_none());
|
||||
assert!(current_card.get("api_formats").is_none());
|
||||
let legacy_card = codex_payload["models"]
|
||||
.as_array()
|
||||
.and_then(|models| models.iter().find(|model| model["slug"] == "legacy-alias"))
|
||||
.expect("legacy model card should be projected");
|
||||
assert_eq!(
|
||||
legacy_card["base_instructions"],
|
||||
"Use legacy future instructions."
|
||||
);
|
||||
assert_eq!(
|
||||
legacy_card["future_capability"],
|
||||
json!({"mode":"opaque-legacy"})
|
||||
);
|
||||
assert!(codex_payload["models"]
|
||||
.as_array()
|
||||
.is_some_and(
|
||||
|models| models.iter().all(|model| model["slug"] != "hidden-alias"
|
||||
&& model["slug"] != "second-alias"
|
||||
&& model["slug"] != "gpt-future-unmapped")
|
||||
));
|
||||
assert!(codex_payload.get("object").is_none());
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let restricted_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-legacy-only")
|
||||
.send()
|
||||
.await
|
||||
.expect("restricted Codex models request should succeed");
|
||||
assert_eq!(restricted_response.status(), StatusCode::OK);
|
||||
let restricted_payload: serde_json::Value = restricted_response
|
||||
.json()
|
||||
.await
|
||||
.expect("restricted Codex models body should parse");
|
||||
assert_eq!(
|
||||
restricted_payload["models"].as_array().map(|models| models
|
||||
.iter()
|
||||
.map(|model| model["slug"].as_str().unwrap_or_default())
|
||||
.collect::<Vec<_>>()),
|
||||
Some(vec!["legacy-alias"])
|
||||
);
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let incomplete_authorized_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-hidden-mixed")
|
||||
.send()
|
||||
.await
|
||||
.expect("incomplete authorized Codex catalog request should succeed");
|
||||
assert_eq!(incomplete_authorized_response.status(), StatusCode::OK);
|
||||
assert!(incomplete_authorized_response
|
||||
.headers()
|
||||
.get(http::header::ETAG)
|
||||
.is_none());
|
||||
let incomplete_authorized_payload: serde_json::Value = incomplete_authorized_response
|
||||
.json()
|
||||
.await
|
||||
.expect("incomplete authorized Codex body should parse");
|
||||
assert_eq!(
|
||||
incomplete_authorized_payload["models"]
|
||||
.as_array()
|
||||
.map(Vec::len),
|
||||
Some(0),
|
||||
"a partial non-empty remote catalog would hide the client's bundled fallback models"
|
||||
);
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let pending_second_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-second-mixed")
|
||||
.send()
|
||||
.await
|
||||
.expect("not-yet-published authorized model request should succeed");
|
||||
assert_eq!(pending_second_response.status(), StatusCode::OK);
|
||||
let pending_second_payload: serde_json::Value = pending_second_response
|
||||
.json()
|
||||
.await
|
||||
.expect("not-yet-published authorized model body should parse");
|
||||
assert_eq!(
|
||||
pending_second_payload["models"].as_array().map(Vec::len),
|
||||
Some(0)
|
||||
);
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let captured_catalog_plan = captured_plans
|
||||
.lock()
|
||||
.expect("plans mutex")
|
||||
.iter()
|
||||
.find(|(url, _)| url.contains("/models?"))
|
||||
.cloned()
|
||||
.expect("catalog execution plan should be captured");
|
||||
assert_eq!(
|
||||
captured_catalog_plan.0,
|
||||
"https://chatgpt.example/backend-api/codex/models?client_version=0.145.2"
|
||||
);
|
||||
assert_eq!(
|
||||
captured_catalog_plan.1.as_deref(),
|
||||
Some("codex_cli_rs/0.145.2")
|
||||
);
|
||||
|
||||
let fresh_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("fresh Codex models request should succeed");
|
||||
assert_eq!(fresh_response.status(), StatusCode::OK);
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1_100)).await;
|
||||
catalog_generation.store(1, Ordering::SeqCst);
|
||||
let stale_started = std::time::Instant::now();
|
||||
let stale_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("stale Codex models request should succeed");
|
||||
assert_eq!(stale_response.status(), StatusCode::OK);
|
||||
assert!(stale_started.elapsed() < std::time::Duration::from_millis(400));
|
||||
let stale_payload: serde_json::Value = stale_response
|
||||
.json()
|
||||
.await
|
||||
.expect("stale body should parse");
|
||||
assert!(stale_payload["models"]
|
||||
.as_array()
|
||||
.is_some_and(|models| models.iter().any(|model| model["slug"] == "future-alias")));
|
||||
wait_until(1_000, || catalog_hits.load(Ordering::SeqCst) >= 2).await;
|
||||
|
||||
let failed_refresh_lkg_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("failed refresh should keep serving LKG");
|
||||
let failed_refresh_lkg_payload: serde_json::Value = failed_refresh_lkg_response
|
||||
.json()
|
||||
.await
|
||||
.expect("failed refresh LKG body should parse");
|
||||
assert!(failed_refresh_lkg_payload["models"]
|
||||
.as_array()
|
||||
.is_some_and(|models| models.iter().any(|model| model["slug"] == "future-alias")));
|
||||
assert_eq!(catalog_hits.load(Ordering::SeqCst), 2);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(150)).await;
|
||||
catalog_generation.store(2, Ordering::SeqCst);
|
||||
let refresh_trigger = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("recovered refresh trigger should succeed");
|
||||
assert_eq!(refresh_trigger.status(), StatusCode::OK);
|
||||
wait_until(1_000, || catalog_hits.load(Ordering::SeqCst) >= 3).await;
|
||||
|
||||
let updated_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.145.2"))
|
||||
.header("authorization", "Bearer sk-codex-second-mixed")
|
||||
.send()
|
||||
.await
|
||||
.expect("updated catalog request should succeed");
|
||||
assert_eq!(
|
||||
updated_response
|
||||
.headers()
|
||||
.get(http::header::ETAG)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("\"catalog-etag-v2\"")
|
||||
);
|
||||
let updated_payload: serde_json::Value = updated_response
|
||||
.json()
|
||||
.await
|
||||
.expect("updated body should parse");
|
||||
assert!(updated_payload["models"]
|
||||
.as_array()
|
||||
.is_some_and(|models| models.iter().any(|model| {
|
||||
model["slug"] == "second-alias"
|
||||
&& model["future_capability"] == json!({"mode":"added-without-code-change"})
|
||||
})));
|
||||
assert!(updated_payload["models"]
|
||||
.as_array()
|
||||
.is_some_and(|models| models.iter().all(|model| {
|
||||
model["slug"] != "hidden-alias" && model["slug"] != "gpt-future-unmapped"
|
||||
})));
|
||||
|
||||
let stored_keys = provider_catalog_repository
|
||||
.list_keys_by_ids(&[CATALOG_KEY_ID.to_string()])
|
||||
.await
|
||||
.expect("catalog key should remain readable");
|
||||
assert_eq!(stored_keys.as_slice(), &[original_catalog_key]);
|
||||
let rows_after = candidate_repository
|
||||
.list_for_exact_api_format("openai:responses")
|
||||
.await
|
||||
.expect("candidate rows should load after request");
|
||||
assert_eq!(rows_after, rows_before);
|
||||
|
||||
let inference_response = client
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.body(r#"{"model":"future-alias","input":"hello","store":false}"#)
|
||||
.send()
|
||||
.await
|
||||
.expect("manual alias inference should succeed");
|
||||
assert_eq!(inference_response.status(), StatusCode::OK);
|
||||
let inference_payload: serde_json::Value = inference_response
|
||||
.json()
|
||||
.await
|
||||
.expect("inference body should parse");
|
||||
assert_eq!(inference_payload["model"], "gpt-future-dynamic");
|
||||
|
||||
let standard_response = client
|
||||
.get(format!("{gateway_url}/v1/models"))
|
||||
@@ -562,6 +1191,7 @@ async fn gateway_serves_codex_model_cards_for_versioned_models_requests() {
|
||||
assert!(standard_payload.get("models").is_none());
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -695,7 +1325,7 @@ async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() {
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_millis(500))
|
||||
.timeout(std::time::Duration::from_secs(3))
|
||||
.build()
|
||||
.expect("client should build")
|
||||
.get(format!("{gateway_url}/v1/models"))
|
||||
@@ -755,7 +1385,7 @@ async fn gateway_returns_not_found_for_openai_model_detail_when_candidate_rows_s
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_millis(500))
|
||||
.timeout(std::time::Duration::from_secs(3))
|
||||
.build()
|
||||
.expect("client should build")
|
||||
.get(format!("{gateway_url}/v1/models/gpt-stalled"))
|
||||
|
||||
Reference in New Issue
Block a user