Merge origin/main into fix/gemini-cli-v1internal

This commit is contained in:
Mas0nShi
2026-05-28 11:58:00 +08:00
410 changed files with 38026 additions and 6621 deletions
@@ -1481,7 +1481,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
false,
Some(serde_json::json!(["gemini", "antigravity"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-cli"])),
Some(serde_json::json!(["gemini-cli", "gemini-3.1-flash-lite"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
@@ -1492,12 +1492,23 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
Some(4_102_444_800),
Some(serde_json::json!(["gemini", "antigravity"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-cli"])),
Some(serde_json::json!(["gemini-cli", "gemini-3.1-flash-lite"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
sample_candidate_row_for("gemini-cli", "1")
}
fn sample_native_antigravity_candidate_row() -> StoredMinimalCandidateSelectionRow {
sample_candidate_row_for("gemini-3.1-flash-lite", "native-1")
}
fn sample_candidate_row_for(
global_model_name: &str,
row_suffix: &str,
) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-antigravity-cli-oauth-stream-local-1".to_string(),
provider_name: "antigravity".to_string(),
@@ -1518,9 +1529,11 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
model_id: "model-antigravity-cli-oauth-stream-local-1".to_string(),
global_model_id: "global-model-antigravity-cli-oauth-stream-local-1".to_string(),
global_model_name: "gemini-cli".to_string(),
model_id: format!("model-antigravity-cli-oauth-stream-local-{row_suffix}"),
global_model_id: format!(
"global-model-antigravity-cli-oauth-stream-local-{row_suffix}"
),
global_model_name: global_model_name.to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "claude-sonnet-4-5".to_string(),
@@ -1840,6 +1853,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
sample_native_antigravity_candidate_row(),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
@@ -1858,17 +1872,26 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
.with_token_url_for_tests("antigravity", format!("{refresh_url}/oauth/token")),
),
]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
let data_state =
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
)
.with_system_config_values_for_tests([(
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY.to_string(),
json!({
"enabled": true,
"auth_user_id": "user-antigravity-cli-oauth-stream-local-1",
"auth_api_key_id": "api-key-antigravity-cli-oauth-stream-local-1",
"allow_unverified_google_bearer": true
}),
)]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(data_state)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -1990,6 +2013,380 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let inbound_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1internal:streamGenerateContent?alt=sse"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header("authorization", "Bearer google-antigravity-access-token")
.header("x-api-key", client_api_key)
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-inbound-stream-456",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-456",
"model": "gemini-cli",
"userAgent": "antigravity",
"requestType": "checkpoint",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "checkpoint context"}]
}],
"generationConfig": {
"temperature": 0.4,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("inbound antigravity request should succeed");
let inbound_status = inbound_response.status();
let inbound_miss_reason = inbound_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let inbound_response_body = inbound_response.text().await.expect("body should read");
assert_eq!(
inbound_status,
StatusCode::OK,
"unexpected inbound antigravity response body: {inbound_response_body}; miss_reason={inbound_miss_reason}"
);
let inbound_response_text = strip_sse_keepalive_comments(&inbound_response_body);
let inbound_payload = inbound_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let inbound_response_json: serde_json::Value =
serde_json::from_str(inbound_payload).expect("stream payload should parse");
assert_eq!(
inbound_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
inbound_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_inbound_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("inbound execution runtime stream should be captured");
assert_eq!(
seen_inbound_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-inbound-stream-456"
);
assert_eq!(
seen_inbound_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_inbound_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_inbound_execution_runtime_request.project,
"project-antigravity-stream-local-1"
);
assert_eq!(
seen_inbound_execution_runtime_request.request_id,
"client-v1internal-request-456"
);
assert_eq!(
seen_inbound_execution_runtime_request.model,
"claude-sonnet-4-5"
);
assert_eq!(
seen_inbound_execution_runtime_request.user_agent,
"antigravity"
);
assert_eq!(
seen_inbound_execution_runtime_request.request_type,
"checkpoint"
);
assert_eq!(seen_inbound_execution_runtime_request.contents_len, 1);
assert!((seen_inbound_execution_runtime_request.exact_temperature - 0.4).abs() < f64::EPSILON);
assert!(!seen_inbound_execution_runtime_request.request_has_model);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let bearer_only_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1internal:streamGenerateContent?alt=sse"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header("authorization", "Bearer google-antigravity-access-token")
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-bearer-only-stream-789",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-789",
"model": "gemini-cli",
"userAgent": "antigravity",
"requestType": "agent",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "bearer-only request"}]
}],
"generationConfig": {
"temperature": 0.5,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("bearer-only antigravity request should succeed");
let bearer_only_status = bearer_only_response.status();
let bearer_only_miss_reason = bearer_only_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let bearer_only_response_body = bearer_only_response.text().await.expect("body should read");
assert_eq!(
bearer_only_status,
StatusCode::OK,
"unexpected bearer-only antigravity response body: {bearer_only_response_body}; miss_reason={bearer_only_miss_reason}"
);
let bearer_only_response_text = strip_sse_keepalive_comments(&bearer_only_response_body);
let bearer_only_payload = bearer_only_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let bearer_only_response_json: serde_json::Value =
serde_json::from_str(bearer_only_payload).expect("stream payload should parse");
assert_eq!(
bearer_only_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
bearer_only_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_bearer_only_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("bearer-only inbound execution runtime stream should be captured");
assert_eq!(
seen_bearer_only_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-bearer-only-stream-789"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.request_id,
"client-v1internal-request-789"
);
assert_eq!(
seen_bearer_only_execution_runtime_request.request_type,
"agent"
);
assert_eq!(seen_bearer_only_execution_runtime_request.contents_len, 1);
assert!(
(seen_bearer_only_execution_runtime_request.exact_temperature - 0.5).abs() < f64::EPSILON
);
assert!(!seen_bearer_only_execution_runtime_request.request_has_model);
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let native_model_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1internal:streamGenerateContent?alt=sse"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header("authorization", "Bearer google-antigravity-access-token")
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.header(
TRACE_ID_HEADER,
"trace-antigravity-v1internal-native-model-stream-790",
)
.json(&json!({
"project": "client-side-project-should-not-leak",
"requestId": "client-v1internal-request-790",
"model": "gemini-3.1-flash-lite",
"userAgent": "antigravity",
"requestType": "agent",
"request": {
"contents": [{
"role": "user",
"parts": [{"text": "native antigravity model request"}]
}],
"generationConfig": {
"temperature": 0.6,
"thinkingConfig": {
"includeThoughts": true
}
},
"toolConfig": {
"functionCallingConfig": {
"mode": "NONE"
}
}
}
}))
.send()
.await
.expect("native-model antigravity request should succeed");
let native_model_status = native_model_response.status();
let native_model_miss_reason = native_model_response
.headers()
.get(crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-")
.to_string();
let native_model_response_body = native_model_response
.text()
.await
.expect("body should read");
assert_eq!(
native_model_status,
StatusCode::OK,
"unexpected native-model antigravity response body: {native_model_response_body}; miss_reason={native_model_miss_reason}"
);
let native_model_response_text = strip_sse_keepalive_comments(&native_model_response_body);
let native_model_payload = native_model_response_text
.trim()
.strip_prefix("data: ")
.expect("response should start with sse data prefix");
let native_model_response_json: serde_json::Value =
serde_json::from_str(native_model_payload).expect("stream payload should parse");
assert_eq!(
native_model_response_json["responseId"],
"resp_antigravity_cli_local_stream_123"
);
assert_eq!(
native_model_response_json["response"]["candidates"][0]["content"]["parts"][0]["text"],
"Hello Antigravity Stream"
);
let seen_native_model_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("native-model inbound execution runtime stream should be captured");
assert_eq!(
seen_native_model_execution_runtime_request.trace_id,
"trace-antigravity-v1internal-native-model-stream-790"
);
assert_eq!(
seen_native_model_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
seen_native_model_execution_runtime_request.authorization,
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(
seen_native_model_execution_runtime_request.model,
"claude-sonnet-4-5"
);
assert_eq!(
seen_native_model_execution_runtime_request.request_id,
"client-v1internal-request-790"
);
assert!(
(seen_native_model_execution_runtime_request.exact_temperature - 0.6).abs() < f64::EPSILON
);
assert!(!seen_native_model_execution_runtime_request.request_has_model);
if std::env::var("AETHER_REAL_AGY_CLI_SMOKE").ok().as_deref() == Some("1") {
*seen_execution_runtime.lock().expect("mutex should lock") = None;
let log_path = std::env::var("AETHER_REAL_AGY_CLI_LOG")
.unwrap_or_else(|_| "/tmp/aether-real-agy-cli-smoke.log".to_string());
let workdir = std::env::var("AETHER_REAL_AGY_CLI_WORKDIR")
.unwrap_or_else(|_| "/tmp/aether-real-agy-cli-work".to_string());
std::fs::create_dir_all(&workdir).expect("agy smoke workdir should create");
let gateway_url_for_agy = gateway_url.clone();
let log_path_for_agy = log_path.clone();
let workdir_for_agy = workdir.clone();
let output = tokio::task::spawn_blocking(move || {
std::process::Command::new("agy")
.arg("--log-file")
.arg(&log_path_for_agy)
.arg("-p")
.arg("Reply with AETHER_CLOSED_LOOP_OK only.")
.arg("--print-timeout")
.arg("45s")
.env("AGY_CLI_DISABLE_AUTO_UPDATE", "true")
.env("CLOUD_CODE_URL", &gateway_url_for_agy)
.current_dir(&workdir_for_agy)
.output()
})
.await
.expect("agy smoke blocking task should join")
.expect("agy smoke process should spawn");
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
let agy_log = std::fs::read_to_string(&log_path).unwrap_or_default();
let seen_agy_execution_runtime_snapshot = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone();
assert!(
output.status.success(),
"agy smoke failed: status={:?}\nseen_execution_runtime={seen_agy_execution_runtime_snapshot:?}\nstdout={stdout}\nstderr={stderr}\nlog={agy_log}",
output.status
);
assert!(
stdout.contains("Hello Antigravity Stream")
|| stdout.contains("AETHER_CLOSED_LOOP_OK"),
"agy smoke stdout did not contain the local runtime response: stdout={stdout}\nstderr={stderr}\nlog={agy_log}"
);
let seen_agy_execution_runtime_request = seen_agy_execution_runtime_snapshot
.expect("real agy smoke should reach execution runtime");
assert_eq!(
seen_agy_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
}
let inbound_stored_candidates = request_candidate_repository
.list_by_request_id("trace-antigravity-v1internal-inbound-stream-456")
.await
.expect("inbound request candidate trace should read");
assert_eq!(inbound_stored_candidates.len(), 1);
assert_eq!(
inbound_stored_candidates[0].status,
RequestCandidateStatus::Success
);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
!*seen_report.lock().expect("mutex should lock"),
@@ -1249,12 +1249,12 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok()),
Some("api_key_concurrency_limit_reached")
Some("auth_api_key_concurrency_limit_reached")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["error"]["message"],
serde_json::Value::String("当前 API Key 并发请求数已达上限,请稍后重试".to_string())
serde_json::Value::String("当前调用方 API Key 并发请求数已达上限,请稍后重试".to_string())
);
let stored_candidates = request_candidate_repository
@@ -1265,7 +1265,7 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored_candidates[0].skip_reason.as_deref(),
Some("api_key_concurrency_limit_reached")
Some("auth_api_key_concurrency_limit_reached")
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
@@ -19,6 +19,30 @@ use aether_data_contracts::repository::provider_catalog::{
use base64::Engine as _;
use sha2::{Digest, Sha256};
const IMAGE_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_image_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(IMAGE_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("image sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[tokio::test]
async fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
#[derive(Debug, Clone)]
@@ -365,8 +389,15 @@ async fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_converts_gemini_image_sync_to_openai_image_provider() {
#[test]
fn gateway_converts_gemini_image_sync_to_openai_image_provider() {
run_image_sync_test(
"gateway_converts_gemini_image_sync_to_openai_image_provider",
gateway_converts_gemini_image_sync_to_openai_image_provider_impl,
);
}
async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
@@ -28,3 +28,4 @@ mod claude;
mod cli;
mod gemini;
mod image;
mod pii_redaction_formats;
@@ -0,0 +1,787 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode,
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use sha2::{Digest, Sha256};
const ORIGINAL_EMAIL: &str = "[email protected]";
fn run_async_test_on_large_stack<F>(name: &'static str, future: F)
where
F: std::future::Future<Output = ()> + Send + 'static,
{
let handle = std::thread::Builder::new()
.name(name.to_string())
.stack_size(16 * 1024 * 1024)
.spawn(move || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime should build")
.block_on(future);
})
.expect("large-stack pii redaction format test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum StandardFormat {
OpenAiChat,
OpenAiResponses,
ClaudeMessages,
}
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
body: serde_json::Value,
headers: serde_json::Value,
url: String,
}
struct RedactionFormatCase {
test_id: &'static str,
trace_id: &'static str,
client_format: StandardFormat,
provider_format: StandardFormat,
}
impl StandardFormat {
fn api_format(self) -> &'static str {
match self {
Self::OpenAiChat => "openai:chat",
Self::OpenAiResponses => "openai:responses",
Self::ClaudeMessages => "claude:messages",
}
}
fn provider_name(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "openai",
Self::ClaudeMessages => "claude",
}
}
fn endpoint_kind(self) -> &'static str {
match self {
Self::OpenAiChat | Self::ClaudeMessages => "chat",
Self::OpenAiResponses => "cli",
}
}
fn client_path(self) -> &'static str {
match self {
Self::OpenAiChat => "/v1/chat/completions",
Self::OpenAiResponses => "/v1/responses",
Self::ClaudeMessages => "/v1/messages",
}
}
fn upstream_base_url(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "https://api.openai.example",
Self::ClaudeMessages => "https://api.anthropic.example",
}
}
fn upstream_path(self) -> &'static str {
match self {
Self::OpenAiChat => "/custom/v1/chat/completions",
Self::OpenAiResponses => "/custom/v1/responses",
Self::ClaudeMessages => "/custom/v1/messages",
}
}
fn client_model(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "gpt-5",
Self::ClaudeMessages => "claude-sonnet-4-5",
}
}
fn provider_model(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "gpt-5-upstream",
Self::ClaudeMessages => "claude-sonnet-4-5-upstream",
}
}
fn provider_auth_type(self) -> &'static str {
match self {
Self::OpenAiChat | Self::OpenAiResponses => "api_key",
Self::ClaudeMessages => "api_key",
}
}
fn client_request_body(self) -> serde_json::Value {
match self {
Self::OpenAiChat => json!({
"model": self.client_model(),
"messages": [
{"role": "system", "content": "Keep answers short."},
{"role": "user", "content": format!("Please contact {ORIGINAL_EMAIL}")}
]
}),
Self::OpenAiResponses => json!({
"model": self.client_model(),
"instructions": format!("Never expose {ORIGINAL_EMAIL}."),
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": format!("Send a status update to {ORIGINAL_EMAIL}")
}]
}],
"store": false
}),
Self::ClaudeMessages => json!({
"model": self.client_model(),
"system": format!("The private contact is {ORIGINAL_EMAIL}."),
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": format!("Draft a reply for {ORIGINAL_EMAIL}")
}]
}],
"max_tokens": 64
}),
}
}
fn execution_runtime_response_body(self, sentinel: &str) -> serde_json::Value {
let restored_text = format!("restored {sentinel}");
match self {
Self::OpenAiChat => json!({
"id": "chatcmpl-redaction-format",
"object": "chat.completion",
"model": self.provider_model(),
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": restored_text},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
}),
Self::OpenAiResponses => json!({
"id": "resp-redaction-format",
"object": "response",
"status": "completed",
"model": self.provider_model(),
"output": [{
"type": "message",
"id": "resp-redaction-format-msg",
"role": "assistant",
"status": "completed",
"content": [{
"type": "output_text",
"text": restored_text,
"annotations": []
}]
}],
"usage": {
"input_tokens": 2,
"output_tokens": 3,
"total_tokens": 5
}
}),
Self::ClaudeMessages => json!({
"id": "msg_redaction_format",
"type": "message",
"model": self.provider_model(),
"role": "assistant",
"content": [{"type": "text", "text": restored_text}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 2,
"output_tokens": 3
}
}),
}
}
}
#[tokio::test]
async fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-responses-pii-redaction-same-format",
trace_id: "trace-openai-responses-pii-redaction-same-format",
client_format: StandardFormat::OpenAiResponses,
provider_format: StandardFormat::OpenAiResponses,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses);
assert!(seen.body.get("input").is_some());
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
}
#[test]
fn ai_execute_claude_messages_pii_redaction_round_trip_same_format() {
run_async_test_on_large_stack(
"ai_execute_claude_messages_pii_redaction_round_trip_same_format",
async {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "claude-messages-pii-redaction-same-format",
trace_id: "trace-claude-messages-pii-redaction-same-format",
client_format: StandardFormat::ClaudeMessages,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_restored_response(&response_json, StandardFormat::ClaudeMessages);
},
);
}
#[tokio::test]
async fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-chat-pii-redaction-before-claude-conversion",
trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion",
client_format: StandardFormat::OpenAiChat,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::ClaudeMessages.provider_model()
);
assert_restored_response(&response_json, StandardFormat::OpenAiChat);
}
#[tokio::test]
async fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "openai-responses-pii-redaction-before-claude-conversion",
trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion",
client_format: StandardFormat::OpenAiResponses,
provider_format: StandardFormat::ClaudeMessages,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::ClaudeMessages.provider_model()
);
assert_restored_response(&response_json, StandardFormat::OpenAiResponses);
}
#[test]
fn ai_execute_claude_messages_pii_redaction_before_openai_chat_conversion() {
run_async_test_on_large_stack(
"ai_execute_claude_messages_pii_redaction_before_openai_chat_conversion",
async {
let (response_json, seen) = run_redaction_format_case(RedactionFormatCase {
test_id: "claude-messages-pii-redaction-before-openai-chat-conversion",
trace_id: "trace-claude-messages-pii-redaction-before-openai-chat-conversion",
client_format: StandardFormat::ClaudeMessages,
provider_format: StandardFormat::OpenAiChat,
})
.await;
assert_provider_request_redacted(&seen, StandardFormat::OpenAiChat);
assert!(seen.body.get("messages").is_some());
assert_eq!(
seen.body["model"],
StandardFormat::OpenAiChat.provider_model()
);
assert_restored_response(&response_json, StandardFormat::ClaudeMessages);
},
);
}
async fn run_redaction_format_case(
case: RedactionFormatCase,
) -> (serde_json::Value, SeenExecutionRuntimeSyncRequest) {
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let provider_format = case.provider_format;
let trace_id = case.trace_id.to_string();
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
let trace_id = trace_id.clone();
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let provider_body = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.expect("json body should exist");
let provider_body_text =
serde_json::to_string(&provider_body).expect("json body should serialize");
let email_sentinel = collect_sentinels(&provider_body_text, "EMAIL")
.into_iter()
.next()
.expect("email sentinel should exist in provider body");
let headers = payload.get("headers").cloned().unwrap_or_else(|| json!({}));
let url = payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string();
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
body: provider_body,
headers,
url,
});
Json(json!({
"request_id": trace_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": provider_format.execution_runtime_response_body(&email_sentinel)
},
"telemetry": {
"elapsed_ms": 19
}
}))
}
}),
);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let auth_repository = auth_repository(&case);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(&case),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider(&case)],
vec![endpoint(&case)],
vec![key(&case)],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(redaction_config()),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let mut request = client
.post(format!("{gateway_url}{}", case.client_format.client_path()))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::ACCEPT_ENCODING, "gzip")
.header(TRACE_ID_HEADER, case.trace_id)
.body(case.client_format.client_request_body().to_string());
request = match case.client_format {
StandardFormat::OpenAiChat | StandardFormat::OpenAiResponses => request.header(
http::header::AUTHORIZATION,
format!("Bearer {}", client_api_key(&case)),
),
StandardFormat::ClaudeMessages => request
.header("x-api-key", client_api_key(&case))
.header("anthropic-version", "2023-06-01"),
};
let response = request.send().await.expect("request should succeed");
let status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let response_text = response.text().await.expect("response body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert_eq!(
execution_path.as_deref(),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
);
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("response body should parse");
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id(case.trace_id)
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
let seen = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
gateway_handle.abort();
execution_runtime_handle.abort();
(response_json, seen)
}
fn auth_repository(case: &RedactionFormatCase) -> Arc<InMemoryAuthApiKeySnapshotRepository> {
let snapshot = auth_snapshot(case);
let key_hash = hash_api_key(&client_api_key(case));
Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(key_hash.clone()),
snapshot.clone(),
)])
.with_export_records(vec![auth_export_record(
&snapshot,
key_hash,
Some(json!({
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": true
}
})),
)]),
)
}
fn auth_snapshot(case: &RedactionFormatCase) -> StoredAuthApiKeySnapshot {
let allowed_providers = unique_json_array([
case.client_format.provider_name(),
case.provider_format.provider_name(),
]);
StoredAuthApiKeySnapshot::new(
format!("user-{}", case.test_id),
"alice".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(allowed_providers.clone()),
Some(json!([case.client_format.api_format()])),
Some(json!([case.client_format.client_model()])),
format!("api-key-{}", case.test_id),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(allowed_providers),
Some(json!([case.client_format.api_format()])),
Some(json!([case.client_format.client_model()])),
)
.expect("auth snapshot should build")
}
fn auth_export_record(
snapshot: &StoredAuthApiKeySnapshot,
key_hash: String,
feature_settings: Option<serde_json::Value>,
) -> StoredAuthApiKeyExportRecord {
StoredAuthApiKeyExportRecord::new(
snapshot.user_id.clone(),
snapshot.api_key_id.clone(),
key_hash,
None,
snapshot.api_key_name.clone(),
snapshot
.api_key_allowed_providers
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot
.api_key_allowed_api_formats
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot
.api_key_allowed_models
.as_ref()
.map(|value| serde_json::json!(value)),
snapshot.api_key_rate_limit,
snapshot.api_key_concurrent_limit,
None,
snapshot.api_key_is_active,
snapshot
.api_key_expires_at_unix_secs
.map(|value| value as i64),
false,
0,
0,
0.0,
snapshot.api_key_is_standalone,
)
.expect("auth api key export record should build")
.with_feature_settings(feature_settings)
}
fn candidate_row(case: &RedactionFormatCase) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{}", case.test_id),
provider_name: case.provider_format.provider_name().to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{}", case.test_id),
endpoint_api_format: case.provider_format.api_format().to_string(),
endpoint_api_family: Some(case.provider_format.provider_name().to_string()),
endpoint_kind: Some(case.provider_format.endpoint_kind().to_string()),
endpoint_is_active: true,
key_id: format!("key-{}", case.test_id),
key_name: "prod".to_string(),
key_auth_type: case.provider_format.provider_auth_type().to_string(),
key_is_active: true,
key_api_formats: Some(vec![case.provider_format.api_format().to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(json!({case.provider_format.api_format(): 1})),
model_id: format!("model-{}", case.test_id),
global_model_id: format!("global-model-{}", case.test_id),
global_model_name: case.client_format.client_model().to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: case.provider_format.provider_model().to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: case.provider_format.provider_model().to_string(),
priority: 1,
api_formats: Some(vec![case.provider_format.api_format().to_string()]),
endpoint_ids: Some(vec![format!("endpoint-{}", case.test_id)]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn provider(case: &RedactionFormatCase) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{}", case.test_id),
case.provider_format.provider_name().to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
case.client_format != case.provider_format,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn endpoint(case: &RedactionFormatCase) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
format!("endpoint-{}", case.test_id),
format!("provider-{}", case.test_id),
case.provider_format.api_format().to_string(),
Some(case.provider_format.provider_name().to_string()),
Some(case.provider_format.endpoint_kind().to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
case.provider_format.upstream_base_url().to_string(),
None,
None,
Some(2),
Some(case.provider_format.upstream_path().to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn key(case: &RedactionFormatCase) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
format!("key-{}", case.test_id),
format!("provider-{}", case.test_id),
"prod".to_string(),
case.provider_format.provider_auth_type().to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!([case.provider_format.api_format()])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&format!("sk-upstream-{}", case.test_id),
)
.expect("api key should encrypt"),
None,
None,
Some(json!({case.provider_format.api_format(): 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn redaction_config() -> Vec<(String, serde_json::Value)> {
vec![
("module.chat_pii_redaction.enabled".to_string(), json!(true)),
(
"module.chat_pii_redaction.rules".to_string(),
json!([{
"id": "email",
"name": "邮箱",
"pattern": r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}",
"enabled": true,
"features": {"validator": "email"},
"system": true
}]),
),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
]
}
fn assert_provider_request_redacted(
seen: &SeenExecutionRuntimeSyncRequest,
provider_format: StandardFormat,
) {
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(
provider_body_text.contains("<AETHER:EMAIL:"),
"provider body was not redacted: {provider_body_text}"
);
assert!(
!provider_body_text.contains(ORIGINAL_EMAIL),
"provider body leaked original email: {provider_body_text}"
);
assert_eq!(seen.headers["accept-encoding"], "identity");
assert!(
seen.url.ends_with(provider_format.upstream_path()),
"unexpected provider url {}",
seen.url
);
}
fn assert_restored_response(response_json: &serde_json::Value, client_format: StandardFormat) {
match client_format {
StandardFormat::OpenAiChat => assert!(response_json["choices"].is_array()),
StandardFormat::ClaudeMessages => assert_eq!(response_json["type"], "message"),
StandardFormat::OpenAiResponses => {}
}
let mut strings = Vec::new();
collect_json_strings(response_json, &mut strings);
assert!(
strings.iter().any(|value| value.contains(ORIGINAL_EMAIL)),
"client response did not restore original email: {response_json}"
);
assert!(
strings.iter().all(|value| !value.contains("<AETHER:")),
"client response still contains redaction sentinel: {response_json}"
);
}
fn collect_json_strings<'a>(value: &'a serde_json::Value, strings: &mut Vec<&'a str>) {
match value {
serde_json::Value::String(value) => strings.push(value),
serde_json::Value::Array(items) => {
for item in items {
collect_json_strings(item, strings);
}
}
serde_json::Value::Object(map) => {
for value in map.values() {
collect_json_strings(value, strings);
}
}
_ => {}
}
}
fn collect_sentinels(text: &str, kind: &str) -> Vec<String> {
let prefix = format!("<AETHER:{kind}:");
let mut sentinels = Vec::new();
let mut offset = 0;
while let Some(relative_start) = text[offset..].find(&prefix) {
let start = offset + relative_start;
let Some(relative_end) = text[start..].find('>') else {
break;
};
let end = start + relative_end + 1;
sentinels.push(text[start..end].to_string());
offset = end;
}
sentinels
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn client_api_key(case: &RedactionFormatCase) -> String {
format!("sk-client-{}", case.test_id)
}
fn unique_json_array(values: [&str; 2]) -> serde_json::Value {
let mut unique = Vec::new();
for value in values {
if !unique.contains(&value) {
unique.push(value);
}
}
json!(unique)
}
@@ -24,6 +24,42 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
"Dockerfile.app.local should pass explicit build version pattern {pattern}"
);
}
let deploy = read_workspace_file("deploy.sh");
for pattern in [
"detect_build_version()",
"git describe --tags --always --dirty",
"AETHER_BUILD_VERSION=\"${AETHER_BUILD_VERSION:-$(detect_build_version)}\"",
"--build-arg \"AETHER_BUILD_VERSION=$AETHER_BUILD_VERSION\"",
">>> AETHER_BUILD_VERSION",
] {
assert!(
deploy.contains(pattern),
"deploy.sh should pass deterministic local build version pattern {pattern}"
);
}
let vite_config = read_workspace_file("frontend/vite.config.ts");
for pattern in [
"process.env.AETHER_BUILD_VERSION",
"process.env.AETHER_VERSION",
] {
assert!(
vite_config.contains(pattern),
"frontend/vite.config.ts should consume local build version pattern {pattern}"
);
}
let core_api = read_workspace_file("apps/aether-gateway/src/api/core.rs");
for pattern in [
"option_env!(\"AETHER_BUILD_VERSION\")",
"\"version\": current_gateway_version()",
] {
assert!(
core_api.contains(pattern),
"api/core.rs should expose build version pattern {pattern}"
);
}
}
#[test]
@@ -371,7 +371,7 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"candidate/mod.rs should not own the core selectable-candidate collector anymore"
);
assert!(
!candidate_mod.contains("auth_api_key_concurrency_limit_reached"),
!candidate_mod.contains("auth_api_key_concurrency_limit_reached("),
"candidate/mod.rs should not own the core auth api key concurrency helper anymore"
);
assert!(
@@ -463,7 +463,7 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"read_provider_key_rpm_states(",
"candidate_is_selectable_with_runtime_state",
"collect_selectable_candidates_from_keys",
"auth_api_key_concurrency_limit_reached",
"auth_api_key_concurrency_limit_reached(",
"build_provider_concurrent_limit_map(",
"reorder_candidates_by_scheduler_health",
] {
@@ -487,7 +487,7 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"candidate/runtime.rs should keep affinity out of runtime eligibility checks"
);
assert!(
runtime.contains("auth_api_key_concurrency_limit_reached"),
runtime.contains("auth_api_key_concurrency_limit_reached("),
"candidate/runtime.rs should depend on core auth api key concurrency helper"
);
assert!(
@@ -7,8 +7,8 @@ use aether_crypto::{
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::DataLayerError;
use axum::body::Body;
@@ -107,6 +107,15 @@ impl ProviderCatalogReadRepository for SummaryNullingProviderCatalogReadReposito
Ok(keys)
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
self.inner
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -232,8 +241,8 @@ async fn gateway_provider_keys_expose_circuit_breaker_and_recover_clears_it() {
"open": true,
"open_at": "2026-03-26T12:00:00+00:00",
"reason": "consecutive_failures_8",
"next_probe_at": "2026-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 1774526460u64,
"next_probe_at": "2099-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 4078209660u64,
"probe_interval_minutes": 1,
"max_probe_interval_minutes": 32,
"half_open_until": null,
@@ -1533,8 +1533,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_gemini_cli_with_trus
assert_eq!(
seen_request.request_body,
Some(json!({
"project": "gemini-cli-project-1",
"userAgent": "GeminiCLI/0.1.5 (Windows; AMD64)"
"project": "gemini-cli-project-1"
}))
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -346,7 +346,7 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at": "2026-03-26T12:05:00+00:00",
"next_probe_at": "2099-03-26T12:05:00+00:00",
"half_open_until": null,
"half_open_successes": 1,
"half_open_failures": 0
@@ -399,7 +399,7 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
payload["circuit_breaker_open_at"],
"2026-03-26T12:01:00+00:00"
);
assert_eq!(payload["next_probe_at"], "2026-03-26T12:05:00+00:00");
assert_eq!(payload["next_probe_at"], "2099-03-26T12:05:00+00:00");
assert_eq!(payload["half_open_successes"], 1);
assert_eq!(payload["half_open_failures"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -408,6 +408,85 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_admin_key_health_summary_treats_expired_unix_circuit_as_closed() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/health/key/key-openai",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![sample_endpoint(
"endpoint-openai",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![
sample_key("key-openai", "provider-openai", "openai:chat", "sk-test")
.with_health_fields(
Some(json!({"openai:chat": {
"health_score": 0.7,
"consecutive_failures": 2
}})),
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 1u64
}})),
),
],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/health/key/key-openai"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "body={body}");
let payload: serde_json::Value = serde_json::from_str(&body).expect("json body should parse");
let circuit = &payload["health_by_format"]["openai:chat"]["circuit_breaker"];
assert_eq!(payload["any_circuit_open"], false);
assert_eq!(circuit["open"], false);
assert_eq!(circuit["state"], "closed");
assert_eq!(payload["key_health_score"], 0.7);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_recovers_admin_key_health_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -442,7 +521,7 @@ async fn gateway_recovers_admin_key_health_locally_with_trusted_admin_principal(
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at": "2026-03-26T12:05:00+00:00",
"next_probe_at": "2099-03-26T12:05:00+00:00",
"half_open_until": null,
"half_open_successes": 0,
"half_open_failures": 1
@@ -818,6 +897,15 @@ async fn gateway_handles_admin_modules_status_locally_with_trusted_admin_princip
);
assert_eq!(payload["bark_push"]["display_name"], "Bark 推送");
assert_eq!(payload["bark_push"]["admin_route"], "/admin/modules/bark");
assert_eq!(payload["s3_backup"]["display_name"], "S3 备份");
assert_eq!(
payload["s3_backup"]["admin_route"],
"/admin/modules/s3-backup"
);
assert_eq!(
payload["s3_backup"]["admin_menu_group"],
serde_json::Value::Null
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -789,7 +789,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
"openai:chat": {"health_score": 0.66}
}));
primary_key.circuit_breaker_by_format = Some(json!({
"openai:chat": {"open": true, "next_probe_at": "2026-03-27T15:00:00Z"}
"openai:chat": {"open": true, "next_probe_at": "2099-03-27T15:00:00Z"}
}));
let mut mapped_key = sample_key(
@@ -907,7 +907,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
openai_keys[0]["circuit_breaker_formats"],
json!(["openai:chat"])
);
assert_eq!(openai_keys[0]["next_probe_at"], "2026-03-27T15:00:00Z");
assert_eq!(openai_keys[0]["next_probe_at"], "2099-03-27T15:00:00Z");
let alt_endpoints = providers[1]["endpoints"]
.as_array()
@@ -2040,7 +2040,8 @@ async fn gateway_handles_admin_monitoring_resilience_circuit_history_locally_wit
"openai:chat": {
"open": true,
"open_at": "2026-03-30T12:00:00+00:00",
"next_probe_at": "2026-03-30T12:05:00+00:00",
"next_probe_at": "2099-03-30T12:05:00+00:00",
"recovery_seconds": 300,
"reason": "错误率过高"
}
})),
@@ -225,7 +225,7 @@ fn windsurf_register_user_execution_result(request_id: &str) -> serde_json::Valu
},
"body": {
"json_body": {
"apiKey": "devin-session-token$registered",
"sessionToken": "devin-session-token$registered",
"name": "Windsurf User",
"apiServerUrl": "https://server.codeium.com"
}
@@ -431,7 +431,7 @@ async fn gateway_rejects_generic_oauth_start_for_windsurf_provider() {
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_callback_token() {
async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
@@ -521,7 +521,7 @@ async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_callback_
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?token=firebase-id-token&state=session-windsurf&provider=google"
"token": "ott$browser-token"
})),
)
.await;
@@ -590,14 +590,22 @@ async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_callback_
.expect("register plan should execute");
assert_eq!(register_plan.method, "POST");
assert_eq!(
register_plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("firebase_id_token"))
.and_then(serde_json::Value::as_str),
Some("firebase-id-token")
register_plan.content_type.as_deref(),
Some("application/proto")
);
assert!(register_plan.body.json_body.is_none());
let encoded_body = register_plan
.body
.body_bytes_b64
.as_deref()
.expect("register body should be bytes");
use base64::Engine as _;
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(encoded_body)
.expect("register body should decode");
let mut expected_body = vec![0x0a, "ott$browser-token".len() as u8];
expected_body.extend_from_slice(b"ott$browser-token");
assert_eq!(body_bytes, expected_body);
assert_eq!(
register_plan
.proxy
@@ -661,7 +669,7 @@ async fn gateway_rejects_windsurf_callback_state_mismatch_and_missing_token() {
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?token=firebase-id-token&state=wrong-state"
"callback_url": "https://windsurf.com/show-auth-token?token=ott$wrong-state&state=wrong-state"
})),
)
.await;
@@ -2874,15 +2882,10 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["expires_at"], 4_102_444_800u64);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["account_state_recheck_attempted"], true);
let account_state_recheck_error = payload["account_state_recheck_error"]
.as_str()
.expect("account_state_recheck_error should be string when recheck is attempted");
assert!(
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error == "wham/usage API 返回状态码 403"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
assert_eq!(payload["account_state_recheck_attempted"], false);
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
@@ -6144,10 +6147,10 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
StatusCode::UNAUTHORIZED,
Json(json!({
"error": {
"message": "Your refresh token has already been used to generate a new access token. Please try signing in again.",
"message": "Could not validate your refresh token. Please try signing in again.",
"type": "invalid_request_error",
"param": serde_json::Value::Null,
"code": "refresh_token_reused"
"code": "refresh_token_expired"
}
})),
)
@@ -6178,19 +6181,18 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-expired",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
"expired-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
r#"{"provider_type":"codex","refresh_token":"expired-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
@@ -6200,7 +6202,6 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
vec![endpoint],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
@@ -6240,6 +6241,7 @@ async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_ex
.await
.expect("refresh payload should parse");
assert_eq!(refresh_payload["status"], json!("auto_removed"));
assert_eq!(refresh_payload["message"], json!("已自动删除"));
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
@@ -3526,6 +3526,14 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
);
banned_key.name = "banned".to_string();
banned_key.oauth_invalid_reason = Some("account_banned".to_string());
let mut oauth_expired_key = sample_key(
"key-openai-oauth-expired",
"provider-openai",
"openai:chat",
"sk-oauth-expired",
);
oauth_expired_key.name = "oauth-expired".to_string();
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
let mut healthy_key = sample_key(
"key-openai-healthy",
"provider-openai",
@@ -3537,7 +3545,7 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![banned_key, healthy_key],
vec![banned_key, oauth_expired_key, healthy_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -3574,8 +3582,13 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("remaining keys should load");
assert_eq!(remaining_keys.len(), 1);
assert_eq!(remaining_keys[0].id, "key-openai-healthy");
assert_eq!(remaining_keys.len(), 2);
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-expired"));
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-healthy"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -32,6 +32,30 @@ use crate::constants::{
};
use crate::data::{GatewayDataConfig, GatewayDataState};
const SUB2API_PROVIDER_OPS_BALANCE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_sub2api_provider_ops_balance_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(SUB2API_PROVIDER_OPS_BALANCE_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("sub2api provider ops balance test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
match ManagedRedisServer::start().await {
Ok(server) => Some(server),
@@ -4328,8 +4352,15 @@ async fn gateway_handles_admin_provider_ops_batch_balance_with_pending_cache_hit
ops_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_rotation() {
#[test]
fn gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_rotation() {
run_sub2api_provider_ops_balance_test(
"gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_rotation",
gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_rotation_impl,
);
}
async fn gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_rotation_impl() {
let ops = Router::new()
.route(
"/api/v1/auth/refresh",
@@ -4533,8 +4564,15 @@ async fn gateway_handles_admin_provider_ops_sub2api_balance_with_refresh_token_r
ops_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_when_base_url_has_path(
#[test]
fn gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_when_base_url_has_path() {
run_sub2api_provider_ops_balance_test(
"gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_when_base_url_has_path",
gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_when_base_url_has_path_impl,
);
}
async fn gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_when_base_url_has_path_impl(
) {
let nested_refresh_hits = Arc::new(Mutex::new(0usize));
let nested_refresh_hits_clone = Arc::clone(&nested_refresh_hits);
@@ -4742,8 +4780,15 @@ async fn gateway_handles_admin_provider_ops_sub2api_balance_against_site_root_wh
ops_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_ops_sub2api_balance_with_session_login() {
#[test]
fn gateway_handles_admin_provider_ops_sub2api_balance_with_session_login() {
run_sub2api_provider_ops_balance_test(
"gateway_handles_admin_provider_ops_sub2api_balance_with_session_login",
gateway_handles_admin_provider_ops_sub2api_balance_with_session_login_impl,
);
}
async fn gateway_handles_admin_provider_ops_sub2api_balance_with_session_login_impl() {
let ops = Router::new()
.route(
"/api/v1/auth/login",
@@ -36,6 +36,8 @@ use crate::constants::{
};
use crate::data::GatewayDataState;
static SYSTEM_UPDATE_TEST_MUTEX: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
#[tokio::test]
async fn gateway_handles_admin_system_version_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -158,6 +160,269 @@ async fn gateway_handles_admin_system_check_update_locally_with_bearer_admin_ses
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_system_update_capability_locally() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/update-capability",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/system/update-capability"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["supported"].is_boolean());
assert!(payload["build_type"].is_string());
assert!(payload["task_status"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_prepares_admin_system_update_locally() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/prepare-update",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/prepare-update"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["detail"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_system_apply_update_without_prepared_version() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/apply-update",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/apply-update"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["detail"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_system_rollback_without_previous_release() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/rollback",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/rollback"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::PRECONDITION_REQUIRED);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["detail"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_system_releases_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/releases",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/system/releases"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["current_version"]
.as_str()
.is_some_and(|value| !value.is_empty()));
assert!(payload["releases"].is_array());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_system_apply_update_with_nonexistent_version() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/apply-update",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/apply-update"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.header("content-type", "application/json")
.body(r#"{"version":"v99.99.99"}"#)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["detail"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_system_update_status_locally() {
let _lock = SYSTEM_UPDATE_TEST_MUTEX.lock().await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/system/update-status"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["phase"].is_string());
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_system_aws_regions_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -470,11 +735,11 @@ async fn gateway_handles_admin_system_config_export_locally_with_trusted_admin_p
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-5",
"gpt-5",
"GPT 5",
)])
.with_admin_global_models(vec![{
let mut model = sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5");
model.usage_count = 7;
model
}])
.with_admin_provider_models(vec![sample_admin_provider_model(
"model-gpt-5",
&provider_id,
@@ -537,9 +802,10 @@ async fn gateway_handles_admin_system_config_export_locally_with_trusted_admin_p
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["version"], "2.2");
assert_eq!(payload["version"], "2.3");
assert!(payload["exported_at"].as_str().is_some());
assert_eq!(payload["global_models"][0]["name"], "gpt-5");
assert_eq!(payload["global_models"][0]["usage_count"], json!(7));
assert_eq!(payload["providers"][0]["name"], "openai");
assert_eq!(
payload["providers"][0]["config"]["provider_ops"]["connector"]["credentials"]
@@ -760,7 +1026,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["version"], "1.4");
assert_eq!(payload["version"], "1.5");
assert!(payload["exported_at"].as_str().is_some());
assert_eq!(payload["user_groups"][0]["name"], "Restricted GPT");
assert!(payload["user_groups"][0].get("priority").is_none());
@@ -778,6 +1044,9 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["group_names"],
json!(["Restricted GPT"])
);
assert_eq!(payload["users"][0]["id"], json!("user-1"));
assert_eq!(payload["users"][0]["request_count"], json!(0));
assert_eq!(payload["users"][0]["total_tokens"], json!(0));
assert_eq!(payload["users"][0]["wallet"]["balance"], json!(12.5));
assert_eq!(
payload["users"][0]["wallet"]["recharge_balance"],
@@ -798,6 +1067,10 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["api_keys"][0]["is_standalone"],
json!(false)
);
assert_eq!(
payload["users"][0]["api_keys"][0]["api_key_id"],
json!("key-user-1")
);
assert_eq!(
payload["users"][0]["api_keys"][0]["total_tokens"],
json!(420)
@@ -806,12 +1079,22 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["standalone_keys"][0]["key"],
json!("ak-standalone-live-1")
);
assert_eq!(
payload["standalone_keys"][0]["api_key_id"],
json!("key-standalone-1")
);
assert_eq!(payload["standalone_keys"][0]["total_tokens"], json!(84));
assert_eq!(
payload["standalone_keys"][0]["wallet"]["unlimited"],
json!(true)
);
assert_eq!(payload["standalone_keys"][0].get("is_standalone"), None,);
assert_eq!(payload["usage_aggregates"]["stats_daily"], json!([]));
assert_eq!(payload["usage_aggregates"]["stats_user_daily"], json!([]));
assert_eq!(
payload["usage_aggregates"]["stats_daily_api_key"],
json!([])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -4,7 +4,9 @@ use aether_contracts::ExecutionPlan;
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
AuthModuleReadRepository, InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
};
@@ -27,7 +29,7 @@ use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::{json, Value};
use super::super::helpers::{sample_endpoint, sample_key, sample_provider};
use super::super::helpers::{hash_api_key, sample_endpoint, sample_key, sample_provider};
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, start_server, AppState,
};
@@ -79,6 +81,7 @@ fn sample_system_import_payload() -> Value {
"global_models": [{
"name": "gpt-5",
"display_name": "GPT 5",
"usage_count": 123,
"default_price_per_request": 0.03,
"default_tiered_pricing": {
"tiers": [{
@@ -122,6 +125,14 @@ fn sample_system_import_payload() -> Value {
"name": "primary",
"api_formats": ["openai:chat"],
"auth_type": "api_key",
"auth_type_by_format": {
"openai:chat": "api_key",
"openai:video": "bearer"
},
"allow_auth_channel_mismatch_formats": [
"openai:chat",
"openai:video"
],
"api_key": "sk-import-123",
"internal_priority": 5,
"is_active": true
@@ -335,6 +346,7 @@ async fn gateway_imports_admin_system_config_locally_and_persists_data() {
.expect("global models should load");
assert_eq!(global_models.items.len(), 1);
assert_eq!(global_models.items[0].name, "gpt-5");
assert_eq!(global_models.items[0].usage_count, 123);
let providers = provider_catalog_repository
.list_providers(false)
@@ -371,6 +383,15 @@ async fn gateway_imports_admin_system_config_locally_and_persists_data() {
.expect("api key should decrypt"),
"sk-import-123"
);
assert_eq!(keys[0].api_formats, Some(json!(["openai:chat"])));
assert_eq!(
keys[0].auth_type_by_format,
Some(json!({ "openai:chat": "api_key" }))
);
assert_eq!(
keys[0].allow_auth_channel_mismatch_formats,
Some(json!(["openai:chat"]))
);
let provider_models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
@@ -817,7 +838,7 @@ async fn gateway_rejects_unknown_admin_system_config_import_versions() {
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for version in ["1.9", "2.3"] {
for version in ["1.9", "2.4"] {
let response = client
.post(format!("{gateway_url}/api/admin/system/config/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
@@ -840,7 +861,7 @@ async fn gateway_rejects_unknown_admin_system_config_import_versions() {
.as_str()
.expect("detail should be a string");
assert!(detail.contains(&format!("不支持的配置版本: {version}")));
assert!(detail.contains("支持的版本: 2.0, 2.1, 2.2"));
assert!(detail.contains("支持的版本: 2.0, 2.1, 2.2, 2.3"));
}
gateway_handle.abort();
@@ -1052,6 +1073,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.expect("user api keys should load");
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_requests, 12);
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(user_api_keys[0].total_cost_usd, 1.25);
assert_eq!(
@@ -1079,6 +1101,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
standalone_keys[0].name.as_deref(),
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_requests, 3);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(standalone_keys[0].total_cost_usd, 0.75);
assert_eq!(
@@ -1119,6 +1142,211 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
let _ = upstream_url;
}
#[tokio::test]
async fn gateway_overwrites_existing_admin_system_user_key_usage_totals() {
let user_key_hash = hash_api_key("sk-existing-user-key");
let standalone_key_hash = hash_api_key("sk-existing-standalone-key");
let existing_user = StoredUserAuthRecord::new(
"user-existing".to_string(),
Some("[email protected]".to_string()),
true,
"existing".to_string(),
Some("existing-hash".to_string()),
"user".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(chrono::Utc::now()),
Some(chrono::Utc::now()),
)
.expect("existing user should build");
let user_key_snapshot = StoredAuthApiKeySnapshot::new(
"user-existing".to_string(),
"existing".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-user-existing".to_string(),
Some("Existing User Key".to_string()),
true,
false,
false,
Some(10),
None,
None,
None,
None,
None,
)
.expect("user key snapshot should build");
let standalone_key_snapshot = StoredAuthApiKeySnapshot::new(
"admin-user-123".to_string(),
"admin".to_string(),
Some("[email protected]".to_string()),
"admin".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-standalone-existing".to_string(),
Some("Existing Standalone Key".to_string()),
true,
false,
true,
Some(20),
None,
None,
None,
None,
None,
)
.expect("standalone key snapshot should build");
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![
(Some(user_key_hash.clone()), user_key_snapshot),
(Some(standalone_key_hash.clone()), standalone_key_snapshot),
])
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-existing".to_string(),
"key-user-existing".to_string(),
user_key_hash.clone(),
None,
Some("Existing User Key".to_string()),
None,
None,
None,
Some(10),
None,
None,
true,
None,
false,
1,
2,
0.03,
false,
)
.expect("existing user key export should build"),
StoredAuthApiKeyExportRecord::new(
"admin-user-123".to_string(),
"key-standalone-existing".to_string(),
standalone_key_hash.clone(),
None,
Some("Existing Standalone Key".to_string()),
None,
None,
None,
Some(20),
None,
None,
true,
None,
false,
4,
5,
0.06,
true,
)
.expect("existing standalone key export should build"),
]),
);
let user_repository =
Arc::new(aether_data::repository::users::InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
.with_user_reader(user_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123"), existing_user])
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/users/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"version": "1.4",
"merge_mode": "overwrite",
"users": [{
"id": "source-user-existing",
"email": "[email protected]",
"username": "existing",
"password_hash": "existing-hash",
"role": "user",
"is_active": true,
"api_keys": [{
"api_key_id": "source-user-key",
"key_hash": user_key_hash,
"name": "Imported User Key",
"is_active": true,
"total_requests": 222,
"total_tokens": 3333,
"total_cost_usd": 4.56
}]
}],
"standalone_keys": [{
"api_key_id": "source-standalone-key",
"key_hash": standalone_key_hash,
"name": "Imported Standalone Key",
"is_active": true,
"total_requests": 444,
"total_tokens": 5555,
"total_cost_usd": 6.78
}]
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["stats"]["users"]["updated"], json!(1));
assert_eq!(payload["stats"]["api_keys"]["updated"], json!(1));
assert_eq!(payload["stats"]["standalone_keys"]["updated"], json!(1));
let updated_records = state
.list_auth_api_key_export_records_by_ids(&[
"key-user-existing".to_string(),
"key-standalone-existing".to_string(),
])
.await
.expect("api key export records should load");
let user_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-user-existing")
.expect("updated user key should exist");
assert_eq!(user_key.total_requests, 222);
assert_eq!(user_key.total_tokens, 3333);
assert_eq!(user_key.total_cost_usd, 4.56);
let standalone_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-standalone-existing")
.expect("updated standalone key should exist");
assert_eq!(standalone_key.total_requests, 444);
assert_eq!(standalone_key.total_tokens, 5555);
assert_eq!(standalone_key.total_cost_usd, 6.78);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_system_config_fixture_v22() {
let gateway = build_router_with_state(
@@ -1671,33 +1899,20 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
}
#[tokio::test]
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_and_forces_refresh(
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_without_refresh(
) {
#[derive(Debug, Clone)]
struct SeenRefreshRequest {
content_type: String,
body: String,
}
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_refresh = Arc::new(Mutex::new(false));
let seen_refresh_clone = Arc::clone(&seen_refresh);
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
post(move |_headers: HeaderMap, _body: Bytes| {
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
*seen_refresh_inner.lock().expect("mutex should lock") = true;
axum::Json(json!({
"access_token": "oauth-access-token-refreshed",
"refresh_token": "oauth-refresh-token-refreshed",
@@ -1792,21 +2007,8 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let seen_refresh = seen_refresh
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(
seen_refresh.content_type,
"application/x-www-form-urlencoded"
);
assert!(seen_refresh.body.contains("grant_type=refresh_token"));
assert!(seen_refresh
.body
.contains("refresh_token=oauth-refresh-token-new"));
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 0);
assert!(!*seen_refresh.lock().expect("mutex should lock"));
let providers = provider_catalog_repository
.list_providers(false)
@@ -1823,7 +2025,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
assert_eq!(key.name, "oauth-primary");
assert_eq!(key.oauth_invalid_at_unix_secs, None);
assert_eq!(key.oauth_invalid_reason, None);
assert!(key.expires_at_unix_secs.is_some());
assert_eq!(key.expires_at_unix_secs, None);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -1832,7 +2034,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("api key should be present"),
)
.expect("oauth access token should decrypt"),
"oauth-access-token-refreshed"
"oauth-access-token-new"
);
let auth_config = decrypt_python_fernet_ciphertext(
@@ -1845,15 +2047,12 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
let auth_config: Value =
serde_json::from_str(&auth_config).expect("oauth auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"oauth-refresh-token-refreshed"
);
assert_eq!(auth_config["refresh_token"], "oauth-refresh-token-new");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["token_type"], "Bearer");
assert_eq!(auth_config["expires_at"].as_u64(), key.expires_at_unix_secs);
assert!(auth_config.get("token_type").is_none());
assert!(auth_config.get("expires_at").is_none());
gateway_handle.abort();
refresh_handle.abort();
@@ -15,7 +15,7 @@ use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use axum::body::Body;
use axum::routing::{any, delete, get, patch, post, put};
use axum::{extract::Request, Router};
use chrono::Utc;
use chrono::{TimeZone, Utc};
use http::StatusCode;
use serde_json::json;
@@ -37,6 +37,16 @@ fn sample_admin_user_with_role(
role: &str,
email: &str,
username: &str,
) -> StoredUserAuthRecord {
sample_admin_user_with_role_and_created_at(user_id, role, email, username, Utc::now())
}
fn sample_admin_user_with_role_and_created_at(
user_id: &str,
role: &str,
email: &str,
username: &str,
created_at: chrono::DateTime<Utc>,
) -> StoredUserAuthRecord {
StoredUserAuthRecord::new(
user_id.to_string(),
@@ -51,7 +61,7 @@ fn sample_admin_user_with_role(
Some(json!(["gpt-4.1"])),
true,
false,
Some(Utc::now()),
Some(created_at),
Some(Utc::now()),
)
.expect("user should build")
@@ -206,6 +216,113 @@ fn sample_admin_api_key_snapshot(user_id: &str, api_key_id: &str) -> StoredAuthA
.expect("api key snapshot should build")
}
#[tokio::test]
async fn gateway_sorts_admin_users_by_created_at() {
let oldest = Utc
.with_ymd_and_hms(2026, 1, 10, 0, 0, 0)
.single()
.expect("valid timestamp");
let middle = Utc
.with_ymd_and_hms(2026, 2, 10, 0, 0, 0)
.single()
.expect("valid timestamp");
let newest = Utc
.with_ymd_and_hms(2026, 3, 10, 0, 0, 0)
.single()
.expect("valid timestamp");
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![
sample_admin_user_with_role_and_created_at(
"user-old",
"user",
"[email protected]",
"old",
oldest,
),
sample_admin_user_with_role_and_created_at(
"user-middle",
"user",
"[email protected]",
"middle",
middle,
),
sample_admin_user_with_role_and_created_at(
"user-new",
"user",
"[email protected]",
"new",
newest,
),
])
.with_export_users(vec![
sample_admin_export_user_with("user", true, "user-old", "[email protected]", "old"),
sample_admin_export_user_with(
"user",
true,
"user-middle",
"[email protected]",
"middle",
),
sample_admin_export_user_with("user", true, "user-new", "[email protected]", "new"),
]),
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_user_reader_for_tests(
user_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let desc_response = client
.get(format!(
"{gateway_url}/api/admin/users?skip=0&limit=10&sort_by=created_at&sort_order=desc"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(desc_response.status(), StatusCode::OK);
let desc_payload: serde_json::Value = desc_response.json().await.expect("json should parse");
let desc_ids = desc_payload["items"]
.as_array()
.expect("items should be array")
.iter()
.map(|item| item["id"].as_str().expect("id should be string"))
.collect::<Vec<_>>();
assert_eq!(desc_ids, vec!["user-new", "user-middle", "user-old"]);
let asc_response = client
.get(format!(
"{gateway_url}/api/admin/users?skip=0&limit=10&sort_by=created_at&sort_order=asc"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(asc_response.status(), StatusCode::OK);
let asc_payload: serde_json::Value = asc_response.json().await.expect("json should parse");
let asc_ids = asc_payload["items"]
.as_array()
.expect("items should be array")
.iter()
.map(|item| item["id"].as_str().expect("id should be string"))
.collect::<Vec<_>>();
assert_eq!(asc_ids, vec!["user-old", "user-middle", "user-new"]);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_users_root_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -1698,6 +1815,83 @@ async fn gateway_returns_conflict_for_admin_lock_user_api_key_when_writer_unavai
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_allows_admin_update_user_to_clear_explicit_groups() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().fallback(any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![sample_admin_user("user-1")])
.with_export_users(vec![sample_admin_export_user("user-1")]),
);
let default_group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "GPT Adapt".to_string(),
description: None,
priority: 0,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("default group should create")
.expect("default group should exist");
user_repository
.add_user_to_group(&default_group.id, "user-1")
.await
.expect("default membership should create");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_user_reader_for_tests(user_repository.clone())
.with_system_config_values_for_tests(vec![(
crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY.to_string(),
json!(default_group.id),
)]),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!("{gateway_url}/api/admin/users/user-1"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({ "group_ids": [] }))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["groups"], json!([]));
assert!(user_repository
.list_user_groups_for_user("user-1")
.await
.expect("memberships should load")
.is_empty());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_conflict_for_admin_update_user_when_writer_unavailable() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -10,7 +10,15 @@ use crate::tests::{
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
};
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
};
use async_trait::async_trait;
use axum::response::IntoResponse;
use std::future::pending;
fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
match status {
@@ -118,6 +126,63 @@ fn sample_gemini_video_task(
}
}
struct PendingMinimalCandidateSelectionReadRepository;
impl PendingMinimalCandidateSelectionReadRepository {
async fn pending_rows(
&self,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
pending::<Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>().await
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for PendingMinimalCandidateSelectionReadRepository {
async fn list_for_exact_api_format(
&self,
_api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
async fn list_for_exact_api_format_and_global_model(
&self,
_api_format: &str,
_global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
async fn list_for_exact_api_format_and_requested_model(
&self,
_api_format: &str,
_requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
async fn list_for_exact_api_format_and_requested_model_page(
&self,
_query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
async fn list_pool_key_rows_for_group(
&self,
_query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
_query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.pending_rows().await
}
}
#[tokio::test]
async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
@@ -175,6 +240,119 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-openai-models-stalled")),
unrestricted_models_snapshot("key-stalled", "user-stalled"),
)]));
let candidate_repository = Arc::new(PendingMinimalCandidateSelectionReadRepository);
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
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,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::builder()
.timeout(std::time::Duration::from_millis(500))
.build()
.expect("client should build")
.get(format!("{gateway_url}/v1/models"))
.header("authorization", "Bearer sk-openai-models-stalled")
.send()
.await
.expect("request should return before client timeout");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(
payload["data"]
.as_array()
.expect("data should be an array")
.len(),
0
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_returns_not_found_for_openai_model_detail_when_candidate_rows_stall() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-openai-model-detail-stalled")),
unrestricted_models_snapshot("key-detail-stalled", "user-detail-stalled"),
)]));
let candidate_repository = Arc::new(PendingMinimalCandidateSelectionReadRepository);
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
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,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::builder()
.timeout(std::time::Duration::from_millis(500))
.build()
.expect("client should build")
.get(format!("{gateway_url}/v1/models/gpt-stalled"))
.header("authorization", "Bearer sk-openai-model-detail-stalled")
.send()
.await
.expect("request should return before client timeout");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["code"], "model_not_found");
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_openai_models_with_cross_format_candidates_without_hitting_fallback_probe(
) {
@@ -522,6 +700,174 @@ async fn gateway_rejects_invalid_claude_count_tokens_payload_without_hitting_fal
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
}
}),
);
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let user_settings = json!({
"preferredModelId": "gemini-3.1-flash-lite",
"theme": "dark"
});
let requests = vec![
(
"/v1internal:loadCodeAssist",
json!({"metadata": {"ideType": "ANTIGRAVITY_CLI"}}),
),
(
"/v1internal:fetchAvailableModels",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:fetchUserInfo",
json!({"project": "aether-antigravity-local"}),
),
(
"/v1internal:fetchAdminControls",
json!({"project": "aether-antigravity-local"}),
),
("/v1internal:listExperiments", json!({})),
(
"/v1internal:recordCodeAssistMetrics",
json!({
"project": "aether-antigravity-local",
"requestId": "opaque-request-id",
"metrics": []
}),
),
(
"/v1internal:setUserSettings",
json!({"userSettings": user_settings.clone()}),
),
];
for (path, request_body) in requests {
let response = client
.post(format!("{gateway_url}{path}"))
.header("authorization", "Bearer ant-access-token")
.header("user-agent", "antigravity/cli/1.0.2 linux/arm64")
.json(&request_body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK, "path {path}");
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC),
"path {path}"
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
match path {
"/v1internal:loadCodeAssist" => {
assert_eq!(
payload["cloudaicompanionProject"],
"aether-antigravity-local"
);
assert_eq!(payload["currentTier"]["id"], "free-tier");
assert_eq!(payload["currentTier"]["name"], "Antigravity");
assert_eq!(payload["paidTier"]["id"], "g1-pro-tier");
assert_eq!(payload["gcpManaged"], false);
assert_eq!(payload["allowedTiers"][0]["id"], "free-tier");
assert_eq!(payload["allowedTiers"][0]["isDefault"], true);
assert_eq!(payload["allowedTiers"][1]["id"], "standard-tier");
assert_eq!(
payload["upgradeSubscriptionUri"],
"https://codeassist.google.com/upgrade"
);
}
"/v1internal:fetchAvailableModels" => {
assert_eq!(payload["defaultAgentModelId"], "gemini-3.1-flash-lite");
assert_eq!(
payload["tieredModelIds"]["flash"],
json!(["gemini-3-flash-agent"])
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["displayName"],
"Gemini 3.5 Flash Low"
);
assert_eq!(
payload["models"]["gemini-3.5-flash-low"]["apiProvider"],
"API_PROVIDER_GOOGLE_GEMINI"
);
assert_eq!(
payload["models"]["gemini-2.5-flash-lite"]["model"],
"MODEL_GOOGLE_GEMINI_2_5_FLASH_LITE"
);
assert_eq!(
payload["agentModelSorts"][0]["groups"][0]["modelIds"],
json!([
"gemini-3.1-flash-lite",
"gemini-3-flash-agent",
"gemini-3.1-pro-low",
"gemini-3.5-flash-low"
])
);
assert_eq!(payload["deprecatedModelIds"], json!({}));
assert_eq!(payload["commandModelIds"], json!(["gemini-3-flash"]));
assert_eq!(
payload["imageGenerationModelIds"],
json!(["gemini-3.1-flash-image"])
);
assert_eq!(payload["mqueryModelIds"], json!(["gemini-3.1-flash-lite"]));
assert_eq!(
payload["webSearchModelIds"],
json!(["gemini-3.1-flash-lite"])
);
assert_eq!(
payload["commitMessageModelIds"],
json!(["gemini-3.1-flash-lite"])
);
}
"/v1internal:fetchUserInfo" => {
assert_eq!(payload["regionCode"], "US");
assert_eq!(
payload["userSettings"]["preferredModelId"],
"gemini-3.1-flash-lite"
);
}
"/v1internal:fetchAdminControls" => {
assert_eq!(payload, json!({}));
}
"/v1internal:listExperiments" => {
assert_eq!(payload["experimentIds"], json!([]));
assert_eq!(payload["flags"], json!([]));
}
"/v1internal:recordCodeAssistMetrics" => {
assert_eq!(payload, json!({}));
}
"/v1internal:setUserSettings" => {
assert_eq!(payload["userSettings"], user_settings);
}
other => panic!("unexpected path {other}"),
}
}
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
@@ -145,6 +145,13 @@ async fn gateway_serves_frontend_routes_and_assets_without_shadowing_public_api(
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("cache-control")
.and_then(|value| value.to_str().ok()),
Some("no-store, no-cache, must-revalidate")
);
let content_type = response
.headers()
.get("content-type")
@@ -161,6 +168,13 @@ async fn gateway_serves_frontend_routes_and_assets_without_shadowing_public_api(
.await
.expect("spa request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("cache-control")
.and_then(|value| value.to_str().ok()),
Some("no-store, no-cache, must-revalidate")
);
let body = response.text().await.expect("spa body should be readable");
assert!(body.contains("Aether Frontend"));
@@ -168,6 +168,15 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() {
assert!(owned_routes
.iter()
.any(|value| value == "/v1beta/files/{path...}"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:loadCodeAssist"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:fetchAvailableModels"));
assert!(owned_routes
.iter()
.any(|value| value == "/v1internal:streamGenerateContent"));
assert_eq!(
payload["rust_frontdoor"]["internal_gateway"]["status"],
"rust_native_control_plane"
@@ -485,7 +485,7 @@ async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -539,7 +539,8 @@ async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.total_tokens, 6);
assert_eq!(stored.first_byte_time_ms, Some(19));
assert!(stored.first_byte_time_ms.is_some());
assert!(stored.response_time_ms >= stored.first_byte_time_ms);
assert_eq!(stored.is_stream, true);
gateway_handle.abort();
@@ -572,7 +573,7 @@ async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_a
allow_execution_response.notified().await;
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-pending-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-pending-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -688,7 +689,8 @@ async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_a
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.first_byte_time_ms, Some(19));
assert!(stored.first_byte_time_ms.is_some());
assert!(stored.response_time_ms >= stored.first_byte_time_ms);
gateway_handle.abort();
execution_runtime_handle.abort();
+4 -3
View File
@@ -1213,7 +1213,7 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-local-report-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-local-report-stream-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -1286,7 +1286,7 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read"));
assert_eq!(
body_text,
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}],\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"
);
let stored_usage = wait_for_usage_status(
@@ -1297,7 +1297,8 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
.await;
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 6);
assert_eq!(stored_usage.first_byte_time_ms, Some(11));
assert!(stored_usage.first_byte_time_ms.is_some());
assert!(stored_usage.response_time_ms >= stored_usage.first_byte_time_ms);
assert!(stored_usage.is_stream);
let stored_candidates = request_candidate_repository
+18 -4
View File
@@ -637,7 +637,20 @@ fn assert_usage_and_pricing(
Some(expected_response_time_ms)
);
}
assert_eq!(stored_usage.first_byte_time_ms, expected_ttfb_ms);
if expected_ttfb_ms.is_some() {
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("stream usage should record first visible text time");
assert!(
stored_usage
.response_time_ms
.is_some_and(|response_time_ms| response_time_ms >= first_byte_time_ms),
"stream first_byte_time_ms should not exceed response_time_ms: first_byte={first_byte_time_ms}, response={:?}",
stored_usage.response_time_ms
);
} else {
assert_eq!(stored_usage.first_byte_time_ms, None);
}
assert_eq!(
stored_usage.settlement_input_price_per_1m(),
Some(INPUT_PRICE_PER_1M)
@@ -893,7 +906,7 @@ async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl(
cache_read_tokens: 40,
};
let stream_body = [
"data: {\"id\":\"chatcmpl-openai-usage-pricing-stream-123\",\"choices\":[]}\n\n",
"data: {\"id\":\"chatcmpl-openai-usage-pricing-stream-123\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n",
"data: [DONE]\n\n",
];
let frames = build_stream_frames(
@@ -1084,7 +1097,7 @@ async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_im
cache_read_tokens: 10,
};
let stream_body = [
"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-claude-usage-pricing-stream-123\",\"type\":\"message\",\"model\":\"claude-sonnet-4-5-upstream\",\"role\":\"assistant\",\"content\":[]}}\n\n",
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n",
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
];
let frames = build_stream_frames(
@@ -1273,7 +1286,8 @@ async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 30,
};
let stream_body = ["data: {\"candidates\":[]}\n\n"];
let stream_body =
["data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}]}\n\n"];
let frames = build_stream_frames(
&stream_body,
standardized_usage_json(expected),
@@ -177,7 +177,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
assert!(!background_tasks.is_empty(), "poller task should spawn");
let stored = {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(500);
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let stored = repository
.find(VideoTaskLookupKey::Id("task-local-123"))
@@ -189,7 +189,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}
assert!(
tokio::time::Instant::now() < deadline,
"poller did not refresh task within 500ms"
"poller did not refresh task within 2s"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
@@ -283,7 +283,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
assert!(!background_tasks.is_empty(), "poller task should spawn");
let stored = {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(500);
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let stored = repository
.find(VideoTaskLookupKey::Id("task-local-123"))
@@ -295,7 +295,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}
assert!(
tokio::time::Instant::now() < deadline,
"poller did not refresh task within 500ms"
"poller did not refresh task within 2s"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}