mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
Improve gateway transport and usage runtime
This commit is contained in:
@@ -1,11 +1,9 @@
|
||||
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,
|
||||
to_bytes, wait_until, Arc, Body, Json, Mutex, Request, Router, StatusCode,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
use crate::constants::{
|
||||
EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||
};
|
||||
use crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
@@ -780,6 +778,7 @@ async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_respons
|
||||
"auth_context": {
|
||||
"user_id": "user-openai-cli-local-limit-123",
|
||||
"api_key_id": "key-openai-cli-local-limit-123",
|
||||
"api_key_concurrent_limit": 1,
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/responses"
|
||||
@@ -969,15 +968,15 @@ async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_respons
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync() {
|
||||
fn gateway_executes_openai_responses_sync_after_api_key_concurrency_wait_budget_elapses() {
|
||||
run_cli_sync_test(
|
||||
"gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync",
|
||||
gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl,
|
||||
"gateway_executes_openai_responses_sync_after_api_key_concurrency_wait_budget_elapses",
|
||||
gateway_executes_openai_responses_sync_after_api_key_concurrency_wait_budget_elapses_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl(
|
||||
) {
|
||||
async fn gateway_executes_openai_responses_sync_after_api_key_concurrency_wait_budget_elapses_impl()
|
||||
{
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
@@ -1126,36 +1125,7 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let now_unix_ms = chrono::Utc::now().timestamp_millis().max(0);
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
aether_data_contracts::repository::candidates::StoredRequestCandidate::new(
|
||||
"cand-pending-openai-cli-local-timeout-1".to_string(),
|
||||
"req-inflight-openai-cli-local-timeout-1".to_string(),
|
||||
Some("user-openai-cli-local-timeout-123".to_string()),
|
||||
Some("key-openai-cli-local-timeout-123".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-openai-cli-local-timeout-1".to_string()),
|
||||
Some("endpoint-openai-cli-local-timeout-1".to_string()),
|
||||
Some("key-openai-cli-local-timeout-1".to_string()),
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
now_unix_ms,
|
||||
Some(now_unix_ms),
|
||||
None,
|
||||
)
|
||||
.expect("pending candidate should build"),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
@@ -1171,6 +1141,7 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
"auth_context": {
|
||||
"user_id": "user-openai-cli-local-timeout-123",
|
||||
"api_key_id": "key-openai-cli-local-timeout-123",
|
||||
"api_key_concurrent_limit": 1,
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/responses"
|
||||
@@ -1204,6 +1175,7 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
*execution_runtime_hits_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") += 1;
|
||||
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-cli-local-timeout-123",
|
||||
"status_code": 200,
|
||||
@@ -1265,8 +1237,51 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_client = client.clone();
|
||||
let first_gateway_url = gateway_url.clone();
|
||||
let first_request = tokio::spawn(async move {
|
||||
first_client
|
||||
.post(format!("{first_gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-cli-local-timeout",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-openai-cli-local-timeout-inflight-123",
|
||||
)
|
||||
.body("{\"model\":\"gpt-5\",\"input\":\"first\",\"store\":false}")
|
||||
.send()
|
||||
.await
|
||||
.expect("inflight request should complete")
|
||||
});
|
||||
wait_until(1_000, || {
|
||||
*execution_runtime_hits.lock().expect("mutex should lock") >= 1
|
||||
})
|
||||
.await;
|
||||
let pending_deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(1_000);
|
||||
loop {
|
||||
let inflight_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-cli-local-timeout-inflight-123")
|
||||
.await
|
||||
.expect("inflight request candidate trace should read");
|
||||
if inflight_candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
|
||||
{
|
||||
break;
|
||||
}
|
||||
assert!(
|
||||
tokio::time::Instant::now() < pending_deadline,
|
||||
"inflight request candidate did not become pending"
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let started_at = std::time::Instant::now();
|
||||
let response = reqwest::Client::new()
|
||||
let response = client
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
@@ -1281,44 +1296,40 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena
|
||||
|
||||
assert!(
|
||||
started_at.elapsed() >= std::time::Duration::from_millis(100),
|
||||
"request should wait for the bounded concurrency window before failing"
|
||||
"request should wait for the bounded concurrency window before retrying"
|
||||
);
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED)
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("auth_api_key_concurrency_limit_reached")
|
||||
None
|
||||
);
|
||||
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())
|
||||
);
|
||||
assert_eq!(payload["model"], "gpt-5-upstream");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-cli-local-timeout-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
|
||||
assert_eq!(
|
||||
stored_candidates[0].skip_reason.as_deref(),
|
||||
Some("auth_api_key_concurrency_limit_reached")
|
||||
);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
assert_eq!(stored_candidates[0].skip_reason.as_deref(), None);
|
||||
assert_eq!(
|
||||
*execution_runtime_hits.lock().expect("mutex should lock"),
|
||||
0
|
||||
2
|
||||
);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
let first_response = first_request.await.expect("inflight request should join");
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
|
||||
@@ -43,8 +43,15 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
|
||||
#[test]
|
||||
fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
|
||||
run_image_sync_test(
|
||||
"gateway_converts_openai_image_sync_to_gemini_image_provider",
|
||||
gateway_converts_openai_image_sync_to_gemini_image_provider_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_converts_openai_image_sync_to_gemini_image_provider_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -765,8 +772,15 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh() {
|
||||
#[test]
|
||||
fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh() {
|
||||
run_image_sync_test(
|
||||
"gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh",
|
||||
gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
@@ -1271,8 +1285,15 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
refresh_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_plans_chatgpt_web_image_sync_with_internal_web_executor_url() {
|
||||
#[test]
|
||||
fn gateway_plans_chatgpt_web_image_sync_with_internal_web_executor_url() {
|
||||
run_image_sync_test(
|
||||
"gateway_plans_chatgpt_web_image_sync_with_internal_web_executor_url",
|
||||
gateway_plans_chatgpt_web_image_sync_with_internal_web_executor_url_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_plans_chatgpt_web_image_sync_with_internal_web_executor_url_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
|
||||
Reference in New Issue
Block a user