i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 02:47:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improve gateway transport and usage runtime

Browse Source
This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/ai_execute/sync/cli.rs
+65 -54
View File
@@ -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();