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

Merge PR #669: align GPT-5.6 and Codex request protocols

Browse Source
This commit is contained in:
elky
2026-07-12 21:50:20 +08:00
parent 7f61bb43c7 25c49dd804
commit e58621a735
313 changed files with 29823 additions and 5178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tunnel/embedded/local_relay.rs
+40 -43
View File
@@ -2,7 +2,10 @@ use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER;
use aether_contracts::tunnel::{
resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta,
TUNNEL_RELAY_FORWARDED_BY_HEADER,
};
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use async_stream::stream;
use axum::body::{Body, Bytes};
@@ -23,9 +26,6 @@ use super::protocol;
use super::AppState;
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
const MAX_RELAY_META_LEN: usize = 256 * 1024;
const MIN_RELAY_TIMEOUT_MS: u64 = 1;
const MAX_RELAY_TIMEOUT_MS: u64 = 300_000;
struct StreamGuard {
hub: std::sync::Arc<super::hub::HubRouter>,
@@ -156,15 +156,7 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
}
fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
let timeout_ms = if meta.stream {
meta.stream_first_byte_timeout_ms
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
};
Duration::from_millis(timeout_ms.clamp(MIN_RELAY_TIMEOUT_MS, MAX_RELAY_TIMEOUT_MS))
Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms)
}
pub async fn relay_request(
@@ -252,15 +244,17 @@ pub async fn relay_request(
if stream.is_none() {
envelope_buf.extend_from_slice(&chunk);
let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) {
Ok(result) => result,
Err(error) => {
return release_permit_response(
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
request_permit,
);
}
}) else {
let Some((parsed_meta, body_offset)) =
(match try_decode_tunnel_relay_request_meta(&envelope_buf) {
Ok(result) => result,
Err(error) => {
return release_permit_response(
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
request_permit,
);
}
})
else {
continue;
};
@@ -427,27 +421,6 @@ fn release_permit_response(
response
}
fn try_decode_envelope_meta(
buffer: &BytesMut,
) -> Result<Option<(protocol::RequestMeta, usize)>, String> {
if buffer.len() < 4 {
return Ok(None);
}
let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
if meta_len > MAX_RELAY_META_LEN {
return Err("relay metadata too large".to_string());
}
let meta_end = 4usize
.checked_add(meta_len)
.ok_or_else(|| "relay envelope length overflow".to_string())?;
if buffer.len() < meta_end {
return Ok(None);
}
let meta = serde_json::from_slice::<protocol::RequestMeta>(&buffer[4..meta_end])
.map_err(|e| format!("invalid relay metadata: {e}"))?;
Ok(Some((meta, meta_end)))
}
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
for (name, value) in headers {
if should_skip_local_relay_response_header(name) {
@@ -542,6 +515,30 @@ mod tests {
assert_eq!(relay_header_timeout(&meta), Duration::from_secs(7));
}
#[test]
fn relay_header_timeout_keeps_the_protocol_maximum_for_non_stream_requests() {
let meta = protocol::RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "POST".to_string(),
url: "https://example.com/responses".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS),
stream_first_byte_timeout_ms: None,
timeout: 60,
follow_redirects: None,
http1_only: false,
transport_profile: None,
};
assert_eq!(
relay_header_timeout(&meta),
Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
);
}
fn sample_connected_proxy_node(node_id: &str) -> StoredProxyNode {
StoredProxyNode::new(
node_id.to_string(),
apps/aether-gateway/src/tunnel/mod.rs
+125 -16
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};
use aether_contracts::tunnel::{
resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta, RequestMeta,
TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use aether_data::repository::proxy_nodes::{
@@ -21,6 +22,7 @@ use axum::extract::ws::WebSocketUpgrade;
use axum::extract::{ConnectInfo, Path, Request, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
use bytes::BytesMut;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::json;
@@ -61,6 +63,22 @@ const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
const TUNNEL_RELAY_BASE_URL_ENV: &str = "AETHER_TUNNEL_RELAY_BASE_URL";
const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
pub(crate) async fn send_owner_forward_request(
request: reqwest::RequestBuilder,
first_byte_timeout: Option<Duration>,
) -> Result<reqwest::Response, String> {
match first_byte_timeout {
Some(timeout) => match tokio::time::timeout(timeout, request.send()).await {
Ok(result) => result.map_err(|error| error.to_string()),
Err(_) => Err(format!(
"owner gateway first byte timeout after {} ms",
timeout.as_millis()
)),
},
None => request.send().await.map_err(|error| error.to_string()),
}
}
#[derive(Debug, Deserialize)]
struct InternalTunnelHeartbeatRequest {
node_id: String,
@@ -776,8 +794,29 @@ async fn forward_relay_request_to_owner(
);
}
let limit_exceeded = Arc::new(AtomicBool::new(false));
let prepared_body =
match prepare_owner_relay_request_body(body, body_limit, Arc::clone(&limit_exceeded)).await
{
Ok(prepared_body) => prepared_body,
Err(_) if limit_exceeded.load(Ordering::SeqCst) => {
return build_local_http_error_response(
trace_id,
None,
StatusCode::PAYLOAD_TOO_LARGE,
&format!("tunnel relay body exceeds {body_limit} bytes"),
);
}
Err(error) => {
return build_local_http_error_response(
trace_id,
None,
StatusCode::BAD_REQUEST,
&error,
);
}
};
let mut upstream_request = state.client.post(owner_url);
let mut upstream_request = state.owner_forward_client.post(owner_url);
for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) || name == http::header::HOST {
continue;
@@ -793,18 +832,23 @@ async fn forward_relay_request_to_owner(
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
owner.gateway_instance_id.as_str(),
);
let resolved_timeouts = resolve_tunnel_request_timeouts(&prepared_body.meta);
if let Some(timeout_ms) = resolved_timeouts.response_body_ms {
upstream_request = upstream_request.timeout(Duration::from_millis(timeout_ms));
}
if !parts.headers.contains_key(TRACE_ID_HEADER) {
upstream_request = upstream_request.header(TRACE_ID_HEADER, trace_id);
}
let upstream_response = match upstream_request
.body(build_owner_relay_request_body(
body,
body_limit,
Arc::clone(&limit_exceeded),
))
.send()
.await
let first_byte_timeout = prepared_body
.meta
.stream
.then_some(Duration::from_millis(resolved_timeouts.first_byte_ms));
let upstream_response = match send_owner_forward_request(
upstream_request.body(prepared_body.body),
first_byte_timeout,
)
.await
{
Ok(response) => response,
Err(err) if limit_exceeded.load(Ordering::SeqCst) => {
@@ -862,14 +906,53 @@ fn request_content_length_exceeds_limit(headers: &HeaderMap, body_limit: usize)
.is_some_and(|value| value > body_limit)
}
fn build_owner_relay_request_body(
struct PreparedOwnerRelayRequestBody {
body: reqwest::Body,
meta: RequestMeta,
}
async fn prepare_owner_relay_request_body(
body: Body,
body_limit: usize,
limit_exceeded: Arc<AtomicBool>,
) -> reqwest::Body {
) -> Result<PreparedOwnerRelayRequestBody, String> {
let mut body_stream = body.into_data_stream();
reqwest::Body::wrap_stream(stream! {
let mut forwarded = 0usize;
let mut buffered_chunks = Vec::new();
let mut meta_buffer = BytesMut::new();
let mut forwarded = 0usize;
let mut meta = None;
while meta.is_none() {
let Some(next_chunk) = body_stream.next().await else {
return Err("incomplete tunnel relay metadata".to_string());
};
match next_chunk {
Ok(chunk) => {
let next_forwarded = forwarded.saturating_add(chunk.len());
if next_forwarded > body_limit {
limit_exceeded.store(true, Ordering::SeqCst);
return Err(format!("tunnel relay body exceeds {body_limit} bytes"));
}
forwarded = next_forwarded;
meta_buffer.extend_from_slice(&chunk);
buffered_chunks.push(chunk);
match try_decode_tunnel_relay_request_meta(&meta_buffer) {
Ok(Some((parsed, _))) => meta = Some(parsed),
Ok(None) => {}
Err(error) => return Err(error),
}
}
Err(error) => {
return Err(format!("tunnel relay body read failed: {error}"));
}
}
}
let meta = meta.ok_or_else(|| "incomplete tunnel relay metadata".to_string())?;
let forwarded_body = reqwest::Body::wrap_stream(stream! {
for chunk in buffered_chunks {
yield Ok::<Bytes, io::Error>(chunk);
}
while let Some(next_chunk) = body_stream.next().await {
match next_chunk {
Ok(chunk) => {
@@ -884,12 +967,17 @@ fn build_owner_relay_request_body(
}
yield Ok::<Bytes, io::Error>(chunk);
}
Err(err) => {
yield Err::<Bytes, io::Error>(io::Error::other(err));
Err(error) => {
yield Err::<Bytes, io::Error>(io::Error::other(error));
break;
}
}
}
});
Ok(PreparedOwnerRelayRequestBody {
body: forwarded_body,
meta,
})
}
@@ -1047,12 +1135,15 @@ fn parse_embedded_tunnel_heartbeat_request(
mod tests {
use super::{
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status, current_unix_secs,
tunnel_attachment_key, GatewayDataState, TunnelAttachmentDirectory, TunnelAttachmentRecord,
prepare_owner_relay_request_body, tunnel_attachment_key, GatewayDataState,
TunnelAttachmentDirectory, TunnelAttachmentRecord,
};
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
};
use axum::body::Body;
use serde_json::json;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
@@ -1092,6 +1183,24 @@ mod tests {
)
}
#[tokio::test]
async fn owner_relay_body_preparation_rejects_invalid_metadata() {
let mut envelope = Vec::new();
envelope.extend_from_slice(&1u32.to_be_bytes());
envelope.push(b'{');
let error = prepare_owner_relay_request_body(
Body::from(envelope),
1024,
Arc::new(AtomicBool::new(false)),
)
.await
.err()
.expect("invalid metadata should fail");
assert!(error.contains("invalid relay metadata"));
}
#[tokio::test]
async fn embedded_tunnel_heartbeat_updates_proxy_node_repository() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(