i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-11 03:39:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix rust ci failures

Browse Source
This commit is contained in:
elky
2026-06-01 02:42:10 +08:00
parent 0f4141ef3f
commit 6c2a57f237
3 changed files with 173 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/usage/direct.rs
+148 -123
View File
@@ -19,6 +19,27 @@ fn large_request_body(stream: bool) -> String {
.expect("request body should encode") .expect("request body should encode")
} }
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 usage direct test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[tokio::test] #[tokio::test]
async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() { async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -131,74 +152,76 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled()
upstream_handle.abort(); upstream_handle.abort();
} }
#[tokio::test] #[test]
async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() { fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); run_async_test_on_large_stack("pending-usage-sync-before-runtime-result", async move {
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let execution_request_started = Arc::new(tokio::sync::Notify::new()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let allow_execution_response = Arc::new(tokio::sync::Notify::new()); let execution_request_started = Arc::new(tokio::sync::Notify::new());
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
let upstream = Router::new().route( let upstream = Router::new().route(
"/api/internal/gateway/report-sync", "/api/internal/gateway/report-sync",
any(|_request: Request| async move { Json(json!({"ok": true})) }), any(|_request: Request| async move { Json(json!({"ok": true})) }),
); );
let execution_runtime = Router::new().route( let execution_runtime = Router::new().route(
"/v1/execute/sync", "/v1/execute/sync",
any({ any({
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
move |_request: Request| {
let execution_request_started = Arc::clone(&execution_request_started); let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response); let allow_execution_response = Arc::clone(&allow_execution_response);
async move { move |_request: Request| {
execution_request_started.notify_one(); let execution_request_started = Arc::clone(&execution_request_started);
allow_execution_response.notified().await; let allow_execution_response = Arc::clone(&allow_execution_response);
Json(json!({ async move {
"request_id": "req-usage-sync-pending-123", execution_request_started.notify_one();
"status_code": 200, allow_execution_response.notified().await;
"headers": { Json(json!({
"content-type": "application/json" "request_id": "req-usage-sync-pending-123",
}, "status_code": 200,
"body": { "headers": {
"json_body": { "content-type": "application/json"
"id": "chatcmpl-usage-sync-pending-123", },
"usage": { "body": {
"input_tokens": 3, "json_body": {
"output_tokens": 5, "id": "chatcmpl-usage-sync-pending-123",
"total_tokens": 8 "usage": {
"input_tokens": 3,
"output_tokens": 5,
"total_tokens": 8
}
} }
},
"telemetry": {
"elapsed_ms": 45
} }
}, }))
"telemetry": { }
"elapsed_ms": 45
}
}))
} }
} }),
}), );
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-usage-sync-pending")), Some(hash_api_key("sk-client-openai-usage-sync-pending")),
sample_local_openai_auth_snapshot( sample_local_openai_auth_snapshot(
"api-key-usage-sync-pending-123", "api-key-usage-sync-pending-123",
"user-usage-sync-pending-123", "user-usage-sync-pending-123",
), ),
)])); )]));
let candidate_selection_repository = let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(), sample_local_openai_candidate_row(),
])); ]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()], vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()], vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()], vec![sample_local_openai_key()],
)); ));
let (upstream_url, upstream_handle) = start_server(upstream).await; let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let (execution_runtime_url, execution_runtime_handle) =
let gateway_state = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url) build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests( .with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
@@ -214,74 +237,76 @@ async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arri
enabled: true, enabled: true,
..UsageRuntimeConfig::default() ..UsageRuntimeConfig::default()
}); });
let gateway = build_router_with_state(gateway_state); let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await; let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_task = tokio::spawn({ let request_task = tokio::spawn({
let gateway_url = gateway_url.clone(); let gateway_url = gateway_url.clone();
async move { async move {
let response = reqwest::Client::new() let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions")) .post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json") .header(http::header::CONTENT_TYPE, "application/json")
.header( .header(
http::header::AUTHORIZATION, http::header::AUTHORIZATION,
"Bearer sk-client-openai-usage-sync-pending", "Bearer sk-client-openai-usage-sync-pending",
) )
.header(TRACE_ID_HEADER, "req-usage-sync-pending-123") .header(TRACE_ID_HEADER, "req-usage-sync-pending-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}") .body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send() .send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("body should read");
(status, body)
}
});
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await .await
.expect("request should succeed"); .expect("usage lookup should succeed");
let status = response.status(); if pending
let body = response.text().await.expect("body should read"); .as_ref()
(status, body) .is_some_and(|stored| stored.status == "pending")
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
} }
let pending =
pending.expect("pending usage should be recorded before sync result resolves");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.response_time_ms, Some(45));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}); });
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if pending
.as_ref()
.is_some_and(|stored| stored.status == "pending")
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let pending = pending.expect("pending usage should be recorded before sync result resolves");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.response_time_ms, Some(45));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
} }
#[tokio::test] #[tokio::test]
apps/aether-gateway/src/tunnel/embedded/hub.rs
+3 -5
View File
@@ -44,8 +44,7 @@ static DRAIN_DEADLINE_MS: LazyLock<u64> = LazyLock::new(|| {
static STREAM_MIN_WINDOW_UPDATE_BYTES: LazyLock<u32> = LazyLock::new(|| { static STREAM_MIN_WINDOW_UPDATE_BYTES: LazyLock<u32> = LazyLock::new(|| {
STREAM_INITIAL_WINDOW_BYTES STREAM_INITIAL_WINDOW_BYTES
.saturating_div(4) .saturating_div(4)
.max(1) .clamp(1, 1024 * 1024)
.min(1024 * 1024)
}); });
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -376,9 +375,8 @@ impl ProxyConn {
} }
pub fn update_protocol_version(&self, protocol_version: u8) { pub fn update_protocol_version(&self, protocol_version: u8) {
let negotiated = protocol_version let negotiated =
.max(1) protocol_version.clamp(1, aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION);
.min(aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION);
self.protocol_version.store(negotiated, Ordering::Relaxed); self.protocol_version.store(negotiated, Ordering::Relaxed);
} }
apps/aether-tunnel/src/tunnel/stream_handler.rs
+22 -14
View File
@@ -1394,17 +1394,14 @@ pub async fn handle_stream(
server.active_connections.fetch_add(1, Ordering::Release); server.active_connections.fetch_add(1, Ordering::Release);
let connect_elapsed = handle_stream_inner( let stream_io = StreamIo {
&state,
&server,
stream_id,
meta,
body_rx, body_rx,
&frame_tx, frame_tx: &frame_tx,
response_window.as_ref(), response_window: response_window.as_ref(),
permit, admission_permit: permit,
) };
.await;
let connect_elapsed = handle_stream_inner(&state, &server, stream_id, meta, stream_io).await;
server.active_connections.fetch_sub(1, Ordering::Release); server.active_connections.fetch_sub(1, Ordering::Release);
if let Some(d) = connect_elapsed { if let Some(d) = connect_elapsed {
@@ -1459,16 +1456,27 @@ async fn send_frame(tx: &FrameSender, frame: TunnelFrame) -> bool {
/// Returns the connection-establishment duration (DNS + TCP/TLS + TTFB) if the /// Returns the connection-establishment duration (DNS + TCP/TLS + TTFB) if the
/// upstream request succeeded, or `None` if the request never reached the /// upstream request succeeded, or `None` if the request never reached the
/// response-headers stage. /// response-headers stage.
struct StreamIo<'a> {
body_rx: mpsc::Receiver<TunnelFrame>,
frame_tx: &'a FrameSender,
response_window: &'a StreamSendWindow,
admission_permit: Option<AdmissionPermit>,
}
async fn handle_stream_inner( async fn handle_stream_inner(
state: &AppState, state: &AppState,
server: &ServerContext, server: &ServerContext,
stream_id: u32, stream_id: u32,
meta: RequestMeta, meta: RequestMeta,
body_rx: mpsc::Receiver<TunnelFrame>, stream_io: StreamIo<'_>,
frame_tx: &FrameSender,
response_window: &StreamSendWindow,
mut admission_permit: Option<AdmissionPermit>,
) -> Option<Duration> { ) -> Option<Duration> {
let StreamIo {
body_rx,
frame_tx,
response_window,
mut admission_permit,
} = stream_io;
let mut current_method: hyper::Method = parse_request_method(&meta.method); let mut current_method: hyper::Method = parse_request_method(&meta.method);
let mut current_url = match url::Url::parse(&meta.url) { let mut current_url = match url::Url::parse(&meta.url) {
Ok(u) => u, Ok(u) => u,