fix(ci): stabilize cross-platform workflow checks

This commit is contained in:
zhefox
2026-07-29 21:55:43 +08:00
parent 37df5b93b1
commit f8000012f7
40 changed files with 251 additions and 139 deletions
+1 -1
View File
@@ -591,7 +591,7 @@ jobs:
RUSTC_WRAPPER: sccache RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true" SCCACHE_GHA_ENABLED: "true"
AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test
run: cargo test -p aether-data-mysql mysql_usage_write_repository_upserts_when_url_is_set --lib -- --nocapture run: cargo test -p aether-data-mysql mysql_usage_write_repository_upserts_and_flushes_counters_when_url_is_set --lib -- --nocapture
- name: Run MySQL usage read smoke test - name: Run MySQL usage read smoke test
env: env:
Generated
+17
View File
@@ -1114,6 +1114,19 @@ dependencies = [
"zeroize", "zeroize",
] ]
[[package]]
name = "bigdecimal"
version = "0.4.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695"
dependencies = [
"autocfg",
"libm",
"num-bigint",
"num-integer",
"num-traits",
]
[[package]] [[package]]
name = "bindgen" name = "bindgen"
version = "0.72.1" version = "0.72.1"
@@ -4447,6 +4460,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6"
dependencies = [ dependencies = [
"base64 0.22.1", "base64 0.22.1",
"bigdecimal",
"bytes", "bytes",
"chrono", "chrono",
"crc", "crc",
@@ -4523,6 +4537,7 @@ checksum = "aa003f0038df784eb8fecbbac13affe3da23b45194bd57dba231c8f48199c526"
dependencies = [ dependencies = [
"atoi", "atoi",
"base64 0.22.1", "base64 0.22.1",
"bigdecimal",
"bitflags 2.11.0", "bitflags 2.11.0",
"byteorder", "byteorder",
"bytes", "bytes",
@@ -4566,6 +4581,7 @@ checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46"
dependencies = [ dependencies = [
"atoi", "atoi",
"base64 0.22.1", "base64 0.22.1",
"bigdecimal",
"bitflags 2.11.0", "bitflags 2.11.0",
"byteorder", "byteorder",
"chrono", "chrono",
@@ -4583,6 +4599,7 @@ dependencies = [
"log", "log",
"md-5", "md-5",
"memchr", "memchr",
"num-bigint",
"once_cell", "once_cell",
"rand 0.8.5", "rand 0.8.5",
"serde", "serde",
@@ -2307,8 +2307,8 @@ mod tests {
transport.provider.provider_type = "antigravity".to_string(); transport.provider.provider_type = "antigravity".to_string();
transport.endpoint.base_url = "https://antigravity.googleapis.com".to_string(); transport.endpoint.base_url = "https://antigravity.googleapis.com".to_string();
transport.endpoint.custom_path = None; transport.endpoint.custom_path = None;
transport.key.auth_type = "oauth".to_string(); transport.key.auth_type = "bearer".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string(); transport.key.decrypted_api_key = "imported-antigravity-token".to_string();
transport.key.upstream_metadata = None; transport.key.upstream_metadata = None;
transport.key.decrypted_auth_config = Some( transport.key.decrypted_auth_config = Some(
json!({ json!({
@@ -8,9 +8,9 @@ const ANTHROPIC_PRECOMMIT_MAX_WAIT: Duration = Duration::from_millis(750);
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum StreamCommitPolicy { pub(super) enum StreamCommitPolicy {
OnResponseHeaders, ResponseHeaders,
OnFirstClassifiedBody, FirstClassifiedBody,
OnFirstAnthropicSemanticEvent { FirstAnthropicSemanticEvent {
max_bytes: usize, max_bytes: usize,
max_wait: Duration, max_wait: Duration,
}, },
@@ -28,11 +28,11 @@ impl StreamCommitPolicy {
force_prefetch: bool, force_prefetch: bool,
) -> Self { ) -> Self {
if !has_direct_finalize { if !has_direct_finalize {
return Self::OnFirstClassifiedBody; return Self::FirstClassifiedBody;
} }
if force_prefetch { if force_prefetch {
return Self::OnFirstClassifiedBody; return Self::FirstClassifiedBody;
} }
let content_type = content_type let content_type = content_type
@@ -46,50 +46,50 @@ impl StreamCommitPolicy {
&& !has_private_stream_normalizer && !has_private_stream_normalizer
&& !has_local_stream_rewriter && !has_local_stream_rewriter
{ {
return Self::OnFirstAnthropicSemanticEvent { return Self::FirstAnthropicSemanticEvent {
max_bytes: MAX_STREAM_PREFETCH_BYTES, max_bytes: MAX_STREAM_PREFETCH_BYTES,
max_wait: ANTHROPIC_PRECOMMIT_MAX_WAIT, max_wait: ANTHROPIC_PRECOMMIT_MAX_WAIT,
}; };
} }
return Self::OnResponseHeaders; return Self::ResponseHeaders;
} }
if has_private_stream_normalizer || has_local_stream_rewriter { if has_private_stream_normalizer || has_local_stream_rewriter {
return Self::OnFirstClassifiedBody; return Self::FirstClassifiedBody;
} }
if !provider_api_format.eq_ignore_ascii_case(client_api_format) { if !provider_api_format.eq_ignore_ascii_case(client_api_format) {
return Self::OnFirstClassifiedBody; return Self::FirstClassifiedBody;
} }
if content_type.is_empty() { if content_type.is_empty() {
return Self::OnResponseHeaders; return Self::ResponseHeaders;
} }
if content_type.contains("json") || content_type.ends_with("+json") { if content_type.contains("json") || content_type.ends_with("+json") {
Self::OnFirstClassifiedBody Self::FirstClassifiedBody
} else { } else {
Self::OnResponseHeaders Self::ResponseHeaders
} }
} }
pub(super) const fn commits_on_response_headers(self) -> bool { pub(super) const fn commits_on_response_headers(self) -> bool {
matches!(self, Self::OnResponseHeaders) matches!(self, Self::ResponseHeaders)
} }
pub(super) const fn requires_bounded_frame_wait(self) -> bool { pub(super) const fn requires_bounded_frame_wait(self) -> bool {
matches!(self, Self::OnFirstAnthropicSemanticEvent { .. }) matches!(self, Self::FirstAnthropicSemanticEvent { .. })
} }
pub(super) const fn max_precommit_wait(self) -> Option<Duration> { pub(super) const fn max_precommit_wait(self) -> Option<Duration> {
match self { match self {
Self::OnFirstAnthropicSemanticEvent { max_wait, .. } => Some(max_wait), Self::FirstAnthropicSemanticEvent { max_wait, .. } => Some(max_wait),
Self::OnResponseHeaders | Self::OnFirstClassifiedBody => None, Self::ResponseHeaders | Self::FirstClassifiedBody => None,
} }
} }
pub(super) const fn is_native_anthropic(self) -> bool { pub(super) const fn is_native_anthropic(self) -> bool {
matches!(self, Self::OnFirstAnthropicSemanticEvent { .. }) matches!(self, Self::FirstAnthropicSemanticEvent { .. })
} }
} }
@@ -143,8 +143,7 @@ impl StreamCommitGate {
return StreamPrecommitObservation::Commit; return StreamPrecommitObservation::Commit;
} }
let StreamCommitPolicy::OnFirstAnthropicSemanticEvent { max_bytes, .. } = self.policy let StreamCommitPolicy::FirstAnthropicSemanticEvent { max_bytes, .. } = self.policy else {
else {
return StreamPrecommitObservation::Pending; return StreamPrecommitObservation::Pending;
}; };
@@ -341,7 +340,7 @@ mod tests {
}; };
fn native_anthropic_policy() -> StreamCommitPolicy { fn native_anthropic_policy() -> StreamCommitPolicy {
StreamCommitPolicy::OnFirstAnthropicSemanticEvent { StreamCommitPolicy::FirstAnthropicSemanticEvent {
max_bytes: 16_384, max_bytes: 16_384,
max_wait: Duration::from_millis(750), max_wait: Duration::from_millis(750),
} }
@@ -480,7 +479,7 @@ mod tests {
#[test] #[test]
fn transport_fragment_count_does_not_commit_an_incomplete_anthropic_error() { fn transport_fragment_count_does_not_commit_an_incomplete_anthropic_error() {
let policy = StreamCommitPolicy::OnFirstAnthropicSemanticEvent { let policy = StreamCommitPolicy::FirstAnthropicSemanticEvent {
max_bytes: 1024, max_bytes: 1024,
max_wait: Duration::from_millis(750), max_wait: Duration::from_millis(750),
}; };
@@ -4041,8 +4041,8 @@ async fn execute_execution_runtime_stream_inner(
frame_stream, frame_stream,
stream_precommit_committed, stream_precommit_committed,
provider_pool_in_flight_guard.take(), provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(), retry_scope_out,
retry_fallback_out.as_deref_mut(), retry_fallback_out,
) )
.await; .await;
} }
@@ -573,7 +573,7 @@ impl ProviderTransferStateTracker {
now: Instant, now: Instant,
) -> Option<ProviderTransferLimitReached> { ) -> Option<ProviderTransferLimitReached> {
if self.exhausted_provider_ids.contains(&plan.provider_id) { if self.exhausted_provider_ids.contains(&plan.provider_id) {
return Some(self.reached_snapshot(plan.provider_id.as_str(), now, false, false)?); return self.reached_snapshot(plan.provider_id.as_str(), now, false, false);
} }
let state = self.by_provider.get(&plan.provider_id)?; let state = self.by_provider.get(&plan.provider_id)?;
@@ -7,7 +7,5 @@ mod task;
pub(super) use orchestration::handle_admin_provider_oauth_batch_import; pub(super) use orchestration::handle_admin_provider_oauth_batch_import;
pub(super) use parse::build_admin_provider_oauth_batch_task_state; pub(super) use parse::build_admin_provider_oauth_batch_task_state;
pub(super) use task::{ pub(super) use task::handle_admin_provider_oauth_start_agent_identity_import_task;
handle_admin_provider_oauth_start_agent_identity_import_task, pub(super) use task::handle_admin_provider_oauth_start_batch_import_task;
handle_admin_provider_oauth_start_batch_import_task,
};
@@ -1611,7 +1611,9 @@ async fn proxy_request_inner(
let buffered_body = buffered_body let buffered_body = buffered_body
.as_ref() .as_ref()
.expect("execution runtime/control auth gate should have buffered request body"); .expect("execution runtime/control auth gate should have buffered request body");
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body); let stream_request = control_decision.is_some_and(|decision| {
owner_forward_request_is_stream(&parts, decision, buffered_body)
});
let mut local_execution_exhaustion = None; let mut local_execution_exhaustion = None;
if stream_request { if stream_request {
let execute_stream_started_at = Instant::now(); let execute_stream_started_at = Instant::now();
@@ -3915,7 +3915,9 @@ mod tests {
repository.failing.store(false, Ordering::Release); repository.failing.store(false, Ordering::Release);
tokio::time::timeout(Duration::from_secs(2), async { 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.normal_admission.available_permits() != 2
{
tokio::task::yield_now().await; tokio::task::yield_now().await;
} }
}) })
@@ -383,7 +383,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
.header(TRACE_ID_HEADER, "trace-search-1") .header(TRACE_ID_HEADER, "trace-search-1")
.json(&json!({ .json(&json!({
"id": "session-search-1", "id": "session-search-1",
"model": "gpt-5.6-sol-ultra-fast", "model": "gpt-5.6-sol-ultra",
"reasoning": {"effort": "low", "summary": "auto"}, "reasoning": {"effort": "low", "summary": "auto"},
"input": "find current OpenAI documentation", "input": "find current OpenAI documentation",
"commands": { "commands": {
@@ -2133,7 +2133,7 @@ fn admin_provider_oauth_dispatch_batch_mod_stays_thin() {
"pub(super) fn apply_admin_provider_oauth_batch_import_hints(", "pub(super) fn apply_admin_provider_oauth_batch_import_hints(",
"pub(super) async fn extract_admin_provider_oauth_batch_error_detail(", "pub(super) async fn extract_admin_provider_oauth_batch_error_detail(",
"pub(super) fn build_admin_provider_oauth_batch_import_response(", "pub(super) fn build_admin_provider_oauth_batch_import_response(",
"pub(super) fn build_admin_provider_oauth_batch_task_state(", "pub(in super::super) fn build_admin_provider_oauth_batch_task_state(",
] { ] {
assert!( assert!(
batch_parse.contains(pattern), batch_parse.contains(pattern),
@@ -4994,6 +4994,8 @@ fn retired_api_format_occurrences_are_whitelisted() {
"crates/aether-data/runtime/src/lifecycle/migrate/tests.rs", "crates/aether-data/runtime/src/lifecycle/migrate/tests.rs",
"crates/aether-usage/runtime/src/report.rs", "crates/aether-usage/runtime/src/report.rs",
"frontend/src/api/endpoints/types/__tests__/api-format.spec.ts", "frontend/src/api/endpoints/types/__tests__/api-format.spec.ts",
"frontend/src/views/admin/module-management/modelDirectivesConfig.ts",
"frontend/src/views/admin/module-management/__tests__/modelDirectivesConfig.spec.ts",
]; ];
let allowed = allowed_paths let allowed = allowed_paths
.into_iter() .into_iter()
@@ -1750,7 +1750,8 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() {
} }
let mysql_facade = let mysql_facade =
read_workspace_file("crates/aether-data/runtime/src/repository/usage/mysql.rs"); read_workspace_file("crates/aether-data/runtime/src/repository/usage/mysql.rs");
assert!(mysql_facade.contains("aether_data_mysql::MysqlUsageStorage")); assert!(mysql_facade.contains("use aether_data_mysql::"));
assert!(mysql_facade.contains("MysqlUsageStorage"));
for forbidden in ["sqlx::query", "FROM `usage`", "INSERT INTO `usage`"] { for forbidden in ["sqlx::query", "FROM `usage`", "INSERT INTO `usage`"] {
assert!( assert!(
!mysql_facade.contains(forbidden), !mysql_facade.contains(forbidden),
@@ -332,7 +332,8 @@ fn testkit_gateway_harness_is_opt_in() {
"aether-gateway = { workspace = true, features = [\"testkit\"], optional = true }" "aether-gateway = { workspace = true, features = [\"testkit\"], optional = true }"
)); ));
let testkit_lib = read_workspace_file("crates/aether-testing/testkit/src/lib.rs"); let testkit_lib =
read_workspace_file("crates/aether-testing/testkit/src/lib.rs").replace("\r\n", "\n");
for module in ["execution_runtime", "gateway", "tunnel"] { for module in ["execution_runtime", "gateway", "tunnel"] {
assert!( assert!(
testkit_lib.contains(&format!("#[cfg(feature = \"gateway\")]\nmod {module};")), testkit_lib.contains(&format!("#[cfg(feature = \"gateway\")]\nmod {module};")),
@@ -2828,12 +2828,14 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
.expect("Agent Identity auth config should encrypt"), .expect("Agent Identity auth config should encrypt"),
); );
let mut codex_provider = sample_provider("provider-codex", "codex", 30);
codex_provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(SummaryNullingProviderCatalogReadRepository::seed( let provider_catalog_repository = Arc::new(SummaryNullingProviderCatalogReadRepository::seed(
vec![ vec![
sample_provider("provider-openai", "openai", 10), sample_provider("provider-openai", "openai", 10),
sample_provider("provider-claude", "claude", 20) sample_provider("provider-claude", "claude", 20)
.with_transport_fields(false, false, true, None, None, None, None, None, None), .with_transport_fields(false, false, true, None, None, None, None, None, None),
sample_provider("provider-codex", "codex", 30), codex_provider,
], ],
vec![ vec![
sample_endpoint( sample_endpoint(
@@ -474,12 +474,19 @@ async fn gateway_auto_removes_codex_key_when_quota_proves_oauth_invalid() {
"stale-access-token", "stale-access-token",
); );
key.auth_type = "oauth".to_string(); key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1); key.expires_at_unix_secs = Some(4_102_444_800);
key.oauth_invalid_at_unix_secs = Some(1); key.oauth_invalid_at_unix_secs = Some(1);
key.oauth_invalid_reason = Some( key.oauth_invalid_reason = Some(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销,请重新登录授权" "[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销,请重新登录授权"
.to_string(), .to_string(),
); );
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"invalid-refresh-token","expires_at":4102444800}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider], vec![provider],
@@ -521,7 +528,7 @@ async fn gateway_auto_removes_codex_key_when_quota_proves_oauth_invalid() {
let payload: serde_json::Value = response.json().await.expect("json body should parse"); let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 0); assert_eq!(payload["success"], 0);
assert_eq!(payload["failed"], 1); assert_eq!(payload["failed"], 1);
assert_eq!(payload["auto_removed"], 1); assert_eq!(payload["auto_removed"], 1, "payload={payload}");
assert_eq!(payload["results"][0]["status"], "auth_invalid"); assert_eq!(payload["results"][0]["status"], "auth_invalid");
assert_eq!(payload["results"][0]["auto_removed"], true); assert_eq!(payload["results"][0]["auto_removed"], true);
@@ -152,25 +152,33 @@ fn sample_kiro_device_access_token_without_email() -> String {
} }
fn sample_codex_access_token_with_profile_email(email: &str, account_id: &str) -> String { fn sample_codex_access_token_with_profile_email(email: &str, account_id: &str) -> String {
sample_codex_access_token_with_profile_email_and_expiry(email, account_id, Some(2_000_000_000))
}
fn sample_codex_access_token_with_profile_email_and_expiry(
email: &str,
account_id: &str,
expires_at: Option<u64>,
) -> String {
use base64::Engine as _; use base64::Engine as _;
let header = let header =
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode( let mut claims = json!({
json!({ "iss": "https://auth.openai.com",
"iss": "https://auth.openai.com", "aud": ["https://api.openai.com/v1"],
"aud": ["https://api.openai.com/v1"], "https://api.openai.com/profile": {
"exp": 2_000_000_000u64, "email": email,
"https://api.openai.com/profile": { "email_verified": true,
"email": email, },
"email_verified": true, "https://api.openai.com/auth": {
}, "chatgpt_account_id": account_id,
"https://api.openai.com/auth": { },
"chatgpt_account_id": account_id, });
}, if let Some(expires_at) = expires_at {
}) claims["exp"] = json!(expires_at);
.to_string(), }
); let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string());
format!("{header}.{payload}.sig") format!("{header}.{payload}.sig")
} }
@@ -4352,6 +4360,11 @@ fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_
} }
async fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp_impl() { async fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp_impl() {
let access_token = sample_codex_access_token_with_profile_email_and_expiry(
"opaque@example.com",
"acct-opaque-123",
None,
);
let mut provider = sample_provider("provider-codex", "codex", 10); let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string(); provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint( let endpoint = sample_endpoint(
@@ -4387,7 +4400,7 @@ async fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_h
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({ .json(&json!({
"access_token": "opaque-codex-access-token", "access_token": access_token,
"expiresAt": 2_100_000_000u64, "expiresAt": 2_100_000_000u64,
"name": "temporary-codex-opaque-access-token", "name": "temporary-codex-opaque-access-token",
})) }))
@@ -7763,14 +7776,14 @@ async fn gateway_consecutive_manual_oauth_refresh_uses_rotated_refresh_token_imp
} }
#[test] #[test]
fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait() { fn gateway_concurrent_manual_oauth_refresh_reuses_winner_after_lock_wait() {
run_admin_oauth_test( run_admin_oauth_test(
"gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait", "gateway_concurrent_manual_oauth_refresh_reuses_winner_after_lock_wait",
gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait_impl, gateway_concurrent_manual_oauth_refresh_reuses_winner_after_lock_wait_impl,
); );
} }
async fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait_impl() { async fn gateway_concurrent_manual_oauth_refresh_reuses_winner_after_lock_wait_impl() {
let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new())); let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new()));
let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies); let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies);
let execution_runtime = Router::new().route( let execution_runtime = Router::new().route(
@@ -7963,17 +7976,12 @@ async fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_afte
.lock() .lock()
.expect("mutex should lock") .expect("mutex should lock")
.clone(); .clone();
assert_eq!(bodies.len(), 2); assert_eq!(bodies.len(), 1);
assert!( assert!(
bodies[0].contains("refresh_token=old-codex-refresh-token"), bodies[0].contains("refresh_token=old-codex-refresh-token"),
"unexpected first refresh body: {}", "unexpected first refresh body: {}",
bodies[0] bodies[0]
); );
assert!(
bodies[1].contains("refresh_token=rotated-codex-refresh-token"),
"unexpected second refresh body: {}",
bodies[1]
);
let stored_key = provider_catalog_repository let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-concurrent-refresh".to_string()]) .list_keys_by_ids(&["key-codex-oauth-concurrent-refresh".to_string()])
@@ -7992,10 +8000,7 @@ async fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_afte
.expect("auth config should decrypt"); .expect("auth config should decrypt");
let auth_config: serde_json::Value = let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse"); serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!( assert_eq!(auth_config["refresh_token"], "rotated-codex-refresh-token");
auth_config["refresh_token"],
"rotated-codex-refresh-token-2"
);
gateway_handle.abort(); gateway_handle.abort();
execution_runtime_handle.abort(); execution_runtime_handle.abort();
@@ -4,6 +4,7 @@ use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionResult, ResponseBody}; use aether_contracts::{ExecutionPlan, ExecutionResult, ResponseBody};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY; use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository; use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use base64::Engine as _;
use http::StatusCode; use http::StatusCode;
use serde_json::json; use serde_json::json;
@@ -634,7 +635,17 @@ fn assert_native_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
Some("gemini-embedding-2-preview") Some("gemini-embedding-2-preview")
); );
assert!(!plan.stream); assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body"); assert!(plan.body.json_body.is_none());
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(
plan.body
.body_bytes_b64
.as_deref()
.expect("original request body bytes"),
)
.expect("request body should decode");
let body: serde_json::Value =
serde_json::from_slice(&body_bytes).expect("request body should parse");
assert_eq!(body["content"]["parts"][0]["text"], "hello"); assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("input").is_none()); assert!(body.get("input").is_none());
assert!(body.get("messages").is_none()); assert!(body.get("messages").is_none());
+1
View File
@@ -1203,6 +1203,7 @@ mod tests {
#[test] #[test]
fn to_config_enables_pretty_file_logging_with_defaults() { fn to_config_enables_pretty_file_logging_with_defaults() {
let mut app = sample_app(); let mut app = sample_app();
set_global_field(&mut app, "install_service", "false");
set_global_field(&mut app, "save_logs_to_file", "true"); set_global_field(&mut app, "save_logs_to_file", "true");
let cfg = app.to_config().expect("config should serialize"); let cfg = app.to_config().expect("config should serialize");
@@ -374,6 +374,9 @@ pub fn apply_model_directive_overrides_from_model(
) -> Option<ModelDirective> { ) -> Option<ModelDirective> {
let directive = parse_model_directive(source_model)?; let directive = parse_model_directive(source_model)?;
let mut patched_body = provider_request_body.clone(); let mut patched_body = provider_request_body.clone();
let is_openai_search =
crate::normalize_api_format_alias(provider_api_format) == "openai:search";
let mut applied_override = false;
for override_item in &directive.overrides { for override_item in &directive.overrides {
match override_item { match override_item {
ModelOverride::ReasoningEffort(effort) => { ModelOverride::ReasoningEffort(effort) => {
@@ -384,6 +387,7 @@ pub fn apply_model_directive_overrides_from_model(
&directive.base_model, &directive.base_model,
*effort, *effort,
)?; )?;
applied_override = true;
} }
ModelOverride::CodexReasoningPreset(preset) => { ModelOverride::CodexReasoningPreset(preset) => {
apply_codex_reasoning_preset_override( apply_codex_reasoning_preset_override(
@@ -391,12 +395,20 @@ pub fn apply_model_directive_overrides_from_model(
provider_api_format, provider_api_format,
*preset, *preset,
)?; )?;
applied_override = true;
} }
ModelOverride::ServiceTier(tier) => { ModelOverride::ServiceTier(tier) => {
if is_openai_search {
continue;
}
apply_service_tier_override(&mut patched_body, provider_api_format, *tier)?; apply_service_tier_override(&mut patched_body, provider_api_format, *tier)?;
applied_override = true;
} }
} }
} }
if !applied_override {
return None;
}
*provider_request_body = patched_body; *provider_request_body = patched_body;
Some(directive) Some(directive)
} }
@@ -54,7 +54,7 @@ fn mysql_dashboard_reads_imported_daily_aggregates() {
#[test] #[test]
fn mysql_usage_stat_rebuilds_aggregate_in_sql() { fn mysql_usage_stat_rebuilds_aggregate_in_sql() {
let source = include_str!("../usage.rs"); let source = include_str!("../usage.rs").replace("\r\n", "\n");
assert!(source.contains("UPDATE api_keys\nJOIN (")); assert!(source.contains("UPDATE api_keys\nJOIN ("));
assert!(source.contains("AND status NOT IN ('pending', 'streaming')")); assert!(source.contains("AND status NOT IN ('pending', 'streaming')"));
assert!(source.contains("MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR")); assert!(source.contains("MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR"));
@@ -171,7 +171,7 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
let page = repository let page = repository
.list_admin_wallets(&AdminWalletListQuery { .list_admin_wallets(&AdminWalletListQuery {
status: Some("active".to_string()), status: Some("wallet-read-smoke".to_string()),
owner_type: Some("user".to_string()), owner_type: Some("user".to_string()),
limit: 1, limit: 1,
offset: 1, offset: 1,
@@ -186,7 +186,7 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
let unknown_owner = repository let unknown_owner = repository
.list_admin_wallets(&AdminWalletListQuery { .list_admin_wallets(&AdminWalletListQuery {
status: Some("active".to_string()), status: Some("wallet-read-smoke".to_string()),
owner_type: Some("unknown".to_string()), owner_type: Some("unknown".to_string()),
limit: 10, limit: 10,
offset: 0, offset: 0,
@@ -430,12 +430,12 @@ VALUES ('api-key-1', 'user-1', 'wallet-test-api-key-hash-1', 'Standalone Key', 1
sqlx::query( sqlx::query(
r#" r#"
INSERT INTO wallets ( INSERT INTO wallets (
id, user_id, api_key_id, balance, gift_balance, total_recharged, total_consumed, id, user_id, api_key_id, balance, gift_balance, status, total_recharged,
total_refunded, total_adjusted, created_at, updated_at total_consumed, total_refunded, total_adjusted, created_at, updated_at
) VALUES ) VALUES
('wallet-1', 'user-1', NULL, 10.0, 2.0, 20.0, 4.0, 1.0, 3.0, 1, 2), ('wallet-1', 'user-1', NULL, 10.0, 2.0, 'wallet-read-smoke', 20.0, 4.0, 1.0, 3.0, 1, 2),
('wallet-2', 'user-2', NULL, 4.0, 1.0, 5.0, 2.0, 0.0, 0.0, 1, 3), ('wallet-2', 'user-2', NULL, 4.0, 1.0, 'wallet-read-smoke', 5.0, 2.0, 0.0, 0.0, 1, 3),
('wallet-api-1', NULL, 'api-key-1', 7.0, 0.0, 7.0, 0.0, 0.0, 0.0, 1, 4) ('wallet-api-1', NULL, 'api-key-1', 7.0, 0.0, 'wallet-read-smoke', 7.0, 0.0, 0.0, 0.0, 1, 4)
"#, "#,
) )
.execute(pool) .execute(pool)
@@ -1152,9 +1152,11 @@ fn sort_candidate_selection_rows(
include_global_model: bool, include_global_model: bool,
) { ) {
rows.sort_by(|left, right| { rows.sort_by(|left, right| {
let global_model_order = include_global_model let global_model_order = if include_global_model {
.then(|| left.global_model_name.cmp(&right.global_model_name)) left.global_model_name.cmp(&right.global_model_name)
.unwrap_or(std::cmp::Ordering::Equal); } else {
std::cmp::Ordering::Equal
};
global_model_order global_model_order
.then(left.provider_priority.cmp(&right.provider_priority)) .then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority)) .then(left.key_internal_priority.cmp(&right.key_internal_priority))
@@ -3329,7 +3329,7 @@ mod tests {
assert!(sql.contains("concurrent_limit")); assert!(sql.contains("concurrent_limit"));
} }
let source = include_str!("provider_catalog.rs"); let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
assert!(source.contains("concurrent_limit,")); assert!(source.contains("concurrent_limit,"));
assert!(source.contains("concurrent_limit = $13")); assert!(source.contains("concurrent_limit = $13"));
assert!(source.contains(".bind(key.concurrent_limit)")); assert!(source.contains(".bind(key.concurrent_limit)"));
@@ -3365,7 +3365,7 @@ mod tests {
assert!(sql.contains("allow_auth_channel_mismatch_formats")); assert!(sql.contains("allow_auth_channel_mismatch_formats"));
} }
let source = include_str!("provider_catalog.rs"); let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
assert!( assert!(
source source
.matches( .matches(
@@ -3381,7 +3381,7 @@ mod tests {
#[test] #[test]
fn provider_api_keys_create_key_insert_placeholders_match_bind_order() { fn provider_api_keys_create_key_insert_placeholders_match_bind_order() {
let source = include_str!("provider_catalog.rs"); let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
assert!(source.contains( assert!(source.contains(
" $24,\n $25,\n $26,\n CASE\n WHEN $27::double precision IS NULL THEN NULL" " $24,\n $25,\n $26,\n CASE\n WHEN $27::double precision IS NULL THEN NULL"
)); ));
@@ -25,6 +25,10 @@ use aether_data_contracts::repository::usage::{
UsageProviderPerformanceQuery, UsageTimeSeriesGranularity, UsageWriteRepository, UsageProviderPerformanceQuery, UsageTimeSeriesGranularity, UsageWriteRepository,
}; };
fn normalize_newlines(value: &str) -> String {
value.replace("\r\n", "\n")
}
fn fast_clear_usage_record( fn fast_clear_usage_record(
request_id: &str, request_id: &str,
provider_name: &str, provider_name: &str,
@@ -2352,7 +2356,7 @@ fn usage_counter_pending_health_does_not_scan_processed_history() {
#[test] #[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() { fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
let sql = super::REBUILD_API_KEY_USAGE_STATS_SQL; let sql = normalize_newlines(super::REBUILD_API_KEY_USAGE_STATS_SQL);
assert!(sql.contains("COUNT(*)::BIGINT")); assert!(sql.contains("COUNT(*)::BIGINT"));
assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),")); assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),"));
assert!(!sql.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)")); assert!(!sql.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
@@ -2362,7 +2366,7 @@ fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
#[test] #[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() { fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
let sql = super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL; let sql = normalize_newlines(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL);
assert!(sql.contains("COUNT(*)::BIGINT")); assert!(sql.contains("COUNT(*)::BIGINT"));
assert!(sql.contains("NULLIF(BTRIM(error_message), '') IS NULL")); assert!(sql.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),")); assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),"));
@@ -2574,7 +2578,7 @@ fn usage_sql_aggregate_reads_use_materialized_total_tokens_only_when_available()
projections projections
} }
let source = include_str!("mod.rs"); let source = normalize_newlines(include_str!("mod.rs"));
let breakdown = source let breakdown = source
.split("async fn summarize_usage_breakdown_from_daily_aggregates") .split("async fn summarize_usage_breakdown_from_daily_aggregates")
.nth(1) .nth(1)
@@ -2647,7 +2651,7 @@ fn usage_sql_aggregate_usage_audits_supports_daily_model_and_provider_aggregates
#[test] #[test]
fn usage_sql_provider_aggregation_excludes_unknown_provider_labels() { fn usage_sql_provider_aggregation_excludes_unknown_provider_labels() {
let source = include_str!("mod.rs"); let source = normalize_newlines(include_str!("mod.rs"));
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_FILTER_SQL")); assert!(source.contains("const USAGE_PROVIDER_IDENTITY_FILTER_SQL"));
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_SOURCE_SQL")); assert!(source.contains("const USAGE_PROVIDER_IDENTITY_SOURCE_SQL"));
assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_id, '')) <> ''"#)); assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_id, '')) <> ''"#));
@@ -3280,9 +3284,10 @@ fn usage_sql_casts_json_payload_bind_parameters_explicitly() {
#[test] #[test]
fn usage_sql_insert_values_aligns_request_metadata_and_timestamps() { fn usage_sql_insert_values_aligns_request_metadata_and_timestamps() {
assert!(super::UPSERT_SQL.contains("\n $51::json,\n $52,\n $53::json,\n CASE")); let sql = normalize_newlines(super::UPSERT_SQL);
assert!(super::UPSERT_SQL.contains("WHEN $54 IS NULL THEN NULL")); assert!(sql.contains("\n $51::json,\n $52,\n $53::json,\n CASE"));
assert!(super::UPSERT_SQL.contains("TO_TIMESTAMP($55::double precision)")); assert!(sql.contains("WHEN $54 IS NULL THEN NULL"));
assert!(sql.contains("TO_TIMESTAMP($55::double precision)"));
} }
#[test] #[test]
@@ -3529,7 +3534,7 @@ fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() {
#[test] #[test]
fn first_byte_upsert_sql_is_single_row_guarded_and_preserves_existing_metadata() { fn first_byte_upsert_sql_is_single_row_guarded_and_preserves_existing_metadata() {
let sql = super::UPSERT_FIRST_BYTE_SQL; let sql = normalize_newlines(super::UPSERT_FIRST_BYTE_SQL);
assert_eq!(sql.matches("INSERT INTO").count(), 1); assert_eq!(sql.matches("INSERT INTO").count(), 1);
assert!(!sql.contains("usage_http_audits")); assert!(!sql.contains("usage_http_audits"));
assert!(!sql.contains("usage_routing_snapshots")); assert!(!sql.contains("usage_routing_snapshots"));
@@ -1027,9 +1027,11 @@ fn sort_candidate_selection_rows(
include_global_model: bool, include_global_model: bool,
) { ) {
rows.sort_by(|left, right| { rows.sort_by(|left, right| {
let global_model_order = include_global_model let global_model_order = if include_global_model {
.then(|| left.global_model_name.cmp(&right.global_model_name)) left.global_model_name.cmp(&right.global_model_name)
.unwrap_or(std::cmp::Ordering::Equal); } else {
std::cmp::Ordering::Equal
};
global_model_order global_model_order
.then(left.provider_priority.cmp(&right.provider_priority)) .then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority)) .then(left.key_internal_priority.cmp(&right.key_internal_priority))
+1 -1
View File
@@ -30,7 +30,7 @@ flate2.workspace = true
serde.workspace = true serde.workspace = true
serde_json.workspace = true serde_json.workspace = true
sha2.workspace = true sha2.workspace = true
sqlx = { workspace = true, features = ["migrate", "macros"] } sqlx = { workspace = true, features = ["bigdecimal", "migrate", "macros"] }
thiserror.workspace = true thiserror.workspace = true
tokio.workspace = true tokio.workspace = true
tracing.workspace = true tracing.workspace = true
@@ -28,6 +28,9 @@ BEGIN
END $$; END $$;
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx
ON public.routing_groups USING btree (is_system_default, enabled); ON public.routing_groups USING btree (is_system_default, enabled);
CREATE UNIQUE INDEX IF NOT EXISTS routing_groups_one_system_default_key
ON public.routing_groups (is_system_default)
WHERE is_system_default = TRUE;
CREATE TABLE IF NOT EXISTS public.routing_group_bindings ( CREATE TABLE IF NOT EXISTS public.routing_group_bindings (
id character varying(64) NOT NULL, id character varying(64) NOT NULL,
@@ -53,6 +56,9 @@ CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx
ON public.routing_group_bindings USING btree (group_id); ON public.routing_group_bindings USING btree (group_id);
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx
ON public.routing_group_bindings USING btree (subject_type, subject_id); ON public.routing_group_bindings USING btree (subject_type, subject_id);
CREATE UNIQUE INDEX IF NOT EXISTS routing_group_bindings_subject_default_key
ON public.routing_group_bindings (subject_type, subject_id)
WHERE is_default = TRUE;
CREATE TABLE IF NOT EXISTS public.routing_group_versions ( CREATE TABLE IF NOT EXISTS public.routing_group_versions (
id character varying(64) NOT NULL, id character varying(64) NOT NULL,
@@ -651,7 +651,7 @@ SELECT SHA2(UUID(), 256), usage.api_key_id, ?, COUNT(*),
COALESCE(SUM(usage.cache_creation_input_tokens), 0), COALESCE(SUM(usage.cache_creation_input_tokens), 0),
COALESCE(SUM(usage.cache_read_input_tokens), 0), COALESCE(SUM(usage.cache_read_input_tokens), 0),
COALESCE(SUM(usage.total_cost_usd), 0), MAX(api_keys.name), ?, ? COALESCE(SUM(usage.total_cost_usd), 0), MAX(api_keys.name), ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN api_keys ON api_keys.id = usage.api_key_id LEFT JOIN api_keys ON api_keys.id = usage.api_key_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND usage.api_key_id IS NOT NULL AND usage.api_key_id <> '' AND usage.api_key_id IS NOT NULL AND usage.api_key_id <> ''
@@ -736,7 +736,7 @@ SELECT SHA2(UUID(), 256), usage.user_id, ?, COUNT(*),
COALESCE(SUM(usage.cache_creation_input_tokens), 0), COALESCE(SUM(usage.cache_creation_input_tokens), 0),
COALESCE(SUM(usage.cache_read_input_tokens), 0), COALESCE(SUM(usage.cache_read_input_tokens), 0),
COALESCE(SUM(usage.total_cost_usd), 0), MAX(users.username), ?, ? COALESCE(SUM(usage.total_cost_usd), 0), MAX(users.username), ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN users ON users.id = usage.user_id LEFT JOIN users ON users.id = usage.user_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND usage.user_id IS NOT NULL AND usage.user_id <> ''
@@ -143,7 +143,7 @@ JOIN (
MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs, MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs,
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms,
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
) AS aggregated ) AS aggregated
@@ -207,7 +207,7 @@ JOIN (
COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens,
MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs,
MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE}
@@ -266,7 +266,7 @@ JOIN (
SELECT {select_dimensions}, SELECT {select_dimensions},
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms,
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples
FROM `usage` AS usage FROM `usage` AS `usage`
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? AND {AGGREGATABLE} WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? AND {AGGREGATABLE}
GROUP BY {group_by} GROUP BY {group_by}
) AS aggregated ON {join} ) AS aggregated ON {join}
@@ -391,7 +391,7 @@ JOIN (
COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens,
MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs,
MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
) AS aggregated ) AS aggregated
@@ -461,7 +461,7 @@ JOIN (
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples,
COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens, COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens,
COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens
FROM `usage` AS usage FROM `usage` AS `usage`
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND {AGGREGATABLE} AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} AND usage.model IS NOT NULL AND usage.model <> ''
GROUP BY usage.model GROUP BY usage.model
@@ -503,7 +503,7 @@ JOIN (
COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens,
MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs,
MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE}
@@ -581,7 +581,7 @@ SELECT SHA2(UUID(), 256), usage.user_id,
COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0),
?, ? ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN users ON users.id = usage.user_id LEFT JOIN users ON users.id = usage.user_id
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
@@ -637,7 +637,7 @@ SELECT SHA2(UUID(), 256), ?, usage.model, usage.provider_name, COUNT(*),
COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0),
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE}
@@ -669,7 +669,7 @@ SELECT SHA2(UUID(), 256), usage.user_id, MAX(COALESCE(usage.username, users.user
COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0),
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN users ON users.id = usage.user_id LEFT JOIN users ON users.id = usage.user_id
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ?
@@ -834,7 +834,7 @@ SELECT SHA2(UUID(), 256), {user_select}?, {dimension_select_sql}
COALESCE(settlement.input_price_per_1m, usage.input_price_per_1m, 0) COALESCE(settlement.input_price_per_1m, usage.input_price_per_1m, 0)
* GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) / 1000000.0 * GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) / 1000000.0
), 0), ?, ? ), 0), ?, ?
FROM `usage` AS usage FROM `usage` AS `usage`
LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id
{user_join} {user_join}
WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? {user_filter} WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? {user_filter}
@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres. // Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str = pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql")); include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260720000000; pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260727000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#" const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT SELECT COUNT(*)::BIGINT
@@ -536,6 +536,7 @@ const REQUEST_BODY_DETAIL_TABLES: &[&str] = &["usage_body_blobs"];
#[cfg(all(feature = "postgres", feature = "sqlite"))] #[cfg(all(feature = "postgres", feature = "sqlite"))]
const LIFECYCLE_TABLES: &[&str] = &["_sqlx_migrations", "schema_backfills"]; const LIFECYCLE_TABLES: &[&str] = &["_sqlx_migrations", "schema_backfills"];
#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))]
fn import_column_stores_timestamp(column_name: &str) -> bool { fn import_column_stores_timestamp(column_name: &str) -> bool {
column_name.ends_with("_at") column_name.ends_with("_at")
|| column_name.ends_with("_unix_secs") || column_name.ends_with("_unix_secs")
@@ -547,6 +548,7 @@ fn import_column_stores_timestamp(column_name: &str) -> bool {
) )
} }
#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))]
fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool { fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool {
if !column_name.ends_with("_unix_ms") { if !column_name.ends_with("_unix_ms") {
return false; return false;
@@ -562,6 +564,7 @@ fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool {
!(relation_name == "usage" && column_name == "created_at_unix_ms") !(relation_name == "usage" && column_name == "created_at_unix_ms")
} }
#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))]
fn normalize_imported_integer_timestamp( fn normalize_imported_integer_timestamp(
driver_name: &str, driver_name: &str,
table_name: &str, table_name: &str,
@@ -597,6 +600,7 @@ fn normalize_imported_integer_timestamp(
Ok(Some(timestamp)) Ok(Some(timestamp))
} }
#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))]
fn parse_imported_datetime(value: &str) -> Option<chrono::DateTime<chrono::Utc>> { fn parse_imported_datetime(value: &str) -> Option<chrono::DateTime<chrono::Utc>> {
let value = value.trim(); let value = value.trim();
if let Ok(datetime) = chrono::DateTime::parse_from_rfc3339(value) { if let Ok(datetime) = chrono::DateTime::parse_from_rfc3339(value) {
@@ -476,9 +476,9 @@ async fn load_mysql_import_columns(
let rows = sqlx::query( let rows = sqlx::query(
r#" r#"
SELECT SELECT
COLUMN_NAME AS column_name, CAST(COLUMN_NAME AS CHAR) AS column_name,
DATA_TYPE AS data_type, CAST(DATA_TYPE AS CHAR) AS data_type,
COLUMN_KEY AS column_key, CAST(COLUMN_KEY AS CHAR) AS column_key,
ORDINAL_POSITION AS ordinal_position ORDINAL_POSITION AS ordinal_position
FROM information_schema.columns FROM information_schema.columns
WHERE table_schema = DATABASE() WHERE table_schema = DATABASE()
@@ -506,7 +506,7 @@ WHERE table_schema = DATABASE()
.eq_ignore_ascii_case("PRI") .eq_ignore_ascii_case("PRI")
{ {
primary_key.insert( primary_key.insert(
row.try_get::<i64, _>("ordinal_position").map_sql_err()?, row.try_get::<u32, _>("ordinal_position").map_sql_err()?,
name, name,
); );
} }
@@ -662,7 +662,9 @@ fn mysql_value_to_json(row: &sqlx::mysql::MySqlRow, index: usize) -> Result<Valu
}) })
} }
"DECIMAL" | "NEWDECIMAL" => Ok(Value::String( "DECIMAL" | "NEWDECIMAL" => Ok(Value::String(
row.try_get::<String, _>(index).map_sql_err()?, row.try_get::<sqlx::types::BigDecimal, _>(index)
.map_sql_err()?
.to_string(),
)), )),
"VARCHAR" | "VAR_STRING" | "STRING" | "TEXT" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" "VARCHAR" | "VAR_STRING" | "STRING" | "TEXT" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT"
| "JSON" | "ENUM" | "SET" | "DATE" | "DATETIME" | "TIMESTAMP" | "TIME" => Ok( | "JSON" | "ENUM" | "SET" | "DATE" | "DATETIME" | "TIMESTAMP" | "TIME" => Ok(
@@ -488,6 +488,10 @@ fn normalize_postgres_import_value(
if is_postgres_boolean_column(target_column) { if is_postgres_boolean_column(target_column) {
return normalize_postgres_boolean_value(column_name, value); return normalize_postgres_boolean_value(column_name, value);
} }
if is_postgres_integer_column(target_column) && import_column_stores_timestamp(column_name) {
return normalize_imported_integer_timestamp("postgres", table_name, column_name, value)
.map(|timestamp| timestamp.map_or(Value::Null, |value| Value::Number(value.into())));
}
if is_postgres_timestamp_column(target_column) { if is_postgres_timestamp_column(target_column) {
return normalize_postgres_timestamp_value(table_name, column_name, value); return normalize_postgres_timestamp_value(table_name, column_name, value);
} }
@@ -512,6 +516,13 @@ pub(super) fn is_postgres_timestamp_column(target_column: &PostgresImportColumn)
) || matches!(target_column.udt_name.as_str(), "timestamptz" | "timestamp") ) || matches!(target_column.udt_name.as_str(), "timestamptz" | "timestamp")
} }
fn is_postgres_integer_column(target_column: &PostgresImportColumn) -> bool {
matches!(
target_column.data_type.as_str(),
"smallint" | "integer" | "bigint"
) || matches!(target_column.udt_name.as_str(), "int2" | "int4" | "int8")
}
fn is_postgres_json_column(target_column: &PostgresImportColumn) -> bool { fn is_postgres_json_column(target_column: &PostgresImportColumn) -> bool {
matches!(target_column.data_type.as_str(), "json" | "jsonb") matches!(target_column.data_type.as_str(), "json" | "jsonb")
|| matches!(target_column.udt_name.as_str(), "json" | "jsonb") || matches!(target_column.udt_name.as_str(), "json" | "jsonb")
@@ -655,20 +666,6 @@ fn postgres_billing_table_name(
} }
} }
#[cfg(test)]
mod billing_table_tests {
use super::postgres_billing_table_name;
#[test]
fn settlement_snapshot_import_uses_request_id_conflict_key() {
assert_eq!(
postgres_billing_table_name("usage_settlement_snapshots")
.expect("settlement snapshot table should be supported"),
("public.usage_settlement_snapshots", "request_id")
);
}
}
async fn import_postgres_wallet_row( async fn import_postgres_wallet_row(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
row: &ExportRow, row: &ExportRow,
@@ -721,3 +718,17 @@ fn postgres_wallet_table_name(
)) ))
}) })
} }
#[cfg(test)]
mod billing_table_tests {
use super::postgres_billing_table_name;
#[test]
fn settlement_snapshot_import_uses_request_id_conflict_key() {
assert_eq!(
postgres_billing_table_name("usage_settlement_snapshots")
.expect("settlement snapshot table should be supported"),
("public.usage_settlement_snapshots", "request_id")
);
}
}
@@ -290,6 +290,23 @@ fn cross_driver_timestamp_normalization_preserves_usage_second_contract() {
normalized["created_at_unix_ms"], normalized["created_at_unix_ms"],
json!("2023-11-14T22:13:20+00:00") json!("2023-11-14T22:13:20+00:00")
); );
let target_columns = BTreeMap::from([(
"created_at_unix_ms".to_string(),
postgres_column("bigint", "int8"),
)]);
let row = ExportRow {
id: "usage-1".to_string(),
payload: json!({ "created_at_unix_ms": "1970-01-01T00:00:01.234900Z" }),
};
let normalized = normalize_postgres_import_payload(
"public.usage",
ExportDomain::Usage,
&row,
&target_columns,
)
.expect("postgres integer usage timestamp should normalize");
assert_eq!(normalized["created_at_unix_ms"], json!(1));
} }
#[test] #[test]
@@ -409,6 +409,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260718000000, 20260718000000,
20260718010000, 20260718010000,
20260720000000, 20260720000000,
20260727000000,
] ]
); );
} }
@@ -1023,6 +1024,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260725010000, 20260725010000,
20260725020000, 20260725020000,
20260725030000, 20260725030000,
20260727000000,
] ]
); );
assert_eq!( assert_eq!(
@@ -1055,6 +1057,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260725020000, 20260725020000,
20260725030000, 20260725030000,
20260725040000, 20260725040000,
20260727000000,
] ]
); );
} }
@@ -664,7 +664,7 @@ mod tests {
let provider_id = format!( let provider_id = format!(
"vertex-{}-{}", "vertex-{}-{}",
auth_type, auth_type,
api_format.replace(':', "-").replace('_', "-") api_format.replace([':', '_'], "-")
); );
let mut row = sample_row(&provider_id, api_format, "vertex-model", 10); let mut row = sample_row(&provider_id, api_format, "vertex-model", 10);
row.provider_type = "vertex_ai".to_string(); row.provider_type = "vertex_ai".to_string();
+1 -1
View File
@@ -786,7 +786,7 @@ fn assert_file_contents(path: PathBuf, expected: &str) -> Result<(), SchemaError
path: path.clone(), path: path.clone(),
source, source,
})?; })?;
if actual == expected { if actual.replace("\r\n", "\n") == expected {
Ok(()) Ok(())
} else { } else {
Err(SchemaError::Validation(format!( Err(SchemaError::Validation(format!(
@@ -541,10 +541,11 @@ fn build_configured_claude_count_tokens_url(
.map(str::trim) .map(str::trim)
.filter(|value| !value.is_empty()); .filter(|value| !value.is_empty());
let path = path?; let path = path?;
let path = path let path = if path.starts_with('/') {
.starts_with('/') path.to_string()
.then(|| path.to_string()) } else {
.unwrap_or_else(|| format!("/{path}")); format!("/{path}")
};
build_passthrough_path_url( build_passthrough_path_url(
&transport.endpoint.base_url, &transport.endpoint.base_url,
path.as_str(), path.as_str(),
@@ -1018,7 +1018,7 @@ mod tests {
.expect("headers should build") .expect("headers should build")
}; };
assert!( assert!(
build_headers(compat_behavior, None).get("x-app").is_none(), !build_headers(compat_behavior, None).contains_key("x-app"),
"compatibility profile must not impersonate the Claude Code transport" "compatibility profile must not impersonate the Claude Code transport"
); );
assert_eq!( assert_eq!(
@@ -1029,12 +1029,11 @@ mod tests {
"Claude Code transport headers must survive a transparent body profile" "Claude Code transport headers must survive a transparent body profile"
); );
assert!( assert!(
build_headers( !build_headers(
native_behavior, native_behavior,
Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), Some(aether_ai_formats::ApiOperation::ClaudeCountTokens),
) )
.get("anthropic-beta") .contains_key("anthropic-beta"),
.is_none(),
"native transparent token counting must not inject compatibility betas" "native transparent token counting must not inject compatibility betas"
); );
assert!( assert!(