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

feat(gateway): 流式执行支持 local tunnel 传输并改进 SSE 错误通知

Browse Source 
- 流式执行优先尝试 local tunnel 路径，不可用时降级到直连
- stream_pump 适配 Reqwest 和 LocalTunnel 双响应类型
- SSE passthrough 流中断时向下游发送 aether.error 终端事件
- 提取 frame 编码辅助函数消除重复代码
- 新增本地隧道流式场景的集成测试
This commit is contained in:
fawney19
2026-04-15 21:10:16 +08:00
parent c569081340
commit 704858390d
6 changed files with 1236 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

398
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -1,4 +1,4 @@
use std::collections::VecDeque;
use std::collections::{BTreeMap, VecDeque};
use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFrame, StreamFramePayload};
@@ -49,7 +49,8 @@ use crate::execution_runtime::submission::{
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::transport::{
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution,
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::{
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
@@ -96,6 +97,19 @@ fn record_stream_terminal_usage(
);
}
async fn execute_in_process_stream(
state: &AppState,
plan: &ExecutionPlan,
) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> {
if let Some(execution) = execute_stream_plan_via_local_tunnel(state, plan).await? {
return Ok(execution);
}
DirectSyncExecutionRuntime::new()
.execute_stream(plan.clone())
.await
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
pub(crate) async fn execute_execution_runtime_stream(
state: &AppState,
@@ -145,10 +159,7 @@ pub(crate) async fn execute_execution_runtime_stream(
.unwrap_or_else(|| "-".to_string());
#[cfg(not(test))]
{
let execution = match DirectSyncExecutionRuntime::new()
.execute_stream(plan.clone())
.await
{
let execution = match execute_in_process_stream(state, &plan).await {
Ok(execution) => execution,
Err(err) => {
info!(
@@ -204,10 +215,7 @@ pub(crate) async fn execute_execution_runtime_stream(
.execution_runtime_override_base_url()
.unwrap_or_default();
if remote_execution_runtime_base_url.trim().is_empty() {
let execution = match DirectSyncExecutionRuntime::new()
.execute_stream(plan.clone())
.await
{
let execution = match execute_in_process_stream(state, &plan).await {
Ok(execution) => execution,
Err(err) => {
info!(
@@ -355,6 +363,28 @@ fn decode_stream_data_chunk(
Ok(text.unwrap_or_default().as_bytes().to_vec())
}
fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
headers
.get("content-type")
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
}
fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Bytes, std::io::Error> {
let payload =
serde_json::to_string(&failure.body_json).map_err(|err| IoError::other(err.to_string()))?;
let mut event = String::from("event: aether.error\n");
for line in payload.lines() {
event.push_str("data: ");
event.push_str(line);
event.push('\n');
}
event.push('\n');
Ok(Bytes::from(event))
}
async fn next_stream_frame<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
@@ -1024,6 +1054,8 @@ async fn execute_stream_from_frame_stream(
let request_id_for_report = request_id.to_string();
let request_id_for_report_log = short_request_id(request_id);
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
let emit_passthrough_sse_terminal_error =
skip_direct_finalize_prefetch && response_headers_indicate_sse(&headers);
tokio::spawn(async move {
const MAX_STREAM_BODY_BUFFER_BYTES: usize = 256 * 1024; // 256KB
@@ -1330,6 +1362,43 @@ async fn execute_stream_from_frame_stream(
}
}
if !downstream_dropped && emit_passthrough_sse_terminal_error {
if let Some(failure) = terminal_failure.as_ref() {
match encode_terminal_sse_error_event(failure) {
Ok(error_event) => {
buffered_body.extend(error_event.iter().copied());
if buffered_body.len() > MAX_STREAM_BODY_BUFFER_BYTES {
let tail_start = buffered_body.len() - MAX_STREAM_BODY_BUFFER_BYTES;
buffered_body.drain(..tail_start);
client_body_truncated = true;
}
if tx.send(Ok(error_event)).await.is_err() {
warn!(
event_name = "stream_execution_downstream_terminal_error_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while sending terminal SSE error event"
);
downstream_dropped = true;
}
}
Err(err) => {
warn!(
event_name = "stream_execution_terminal_error_event_encode_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to encode terminal SSE error event"
);
}
}
}
}
drop(tx);
if downstream_dropped {
@@ -1500,7 +1569,41 @@ async fn execute_stream_from_frame_stream(
#[cfg(test)]
mod tests {
use super::should_skip_direct_finalize_prefetch;
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use axum::body::to_bytes;
use axum::extract::ws::Message;
use serde_json::{json, Value};
use tokio::sync::watch;
use super::{execute_execution_runtime_stream, should_skip_direct_finalize_prefetch};
use crate::control::GatewayControlDecision;
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
use crate::AppState;
fn test_decision() -> GatewayControlDecision {
GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
)
.with_execution_runtime_candidate(true)
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),
mode: Some("tunnel".into()),
node_id: Some("node-1".into()),
label: Some("relay-node".into()),
url: None,
extra: Some(json!({"tunnel_base_url": base_url})),
}
}
#[test]
fn skips_prefetch_for_same_format_passthrough_event_streams() {
@@ -1549,4 +1652,277 @@ mod tests {
true,
));
}
#[tokio::test]
async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data(
) {
let state = AppState::new().expect("app state should build");
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
901,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let plan = ExecutionPlan {
request_id: "req-client-stream-error-1".into(),
candidate_id: Some("cand-client-stream-error-1".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/chat".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({"stream": true})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-5".into()),
proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = test_decision();
let state_for_task = state.clone();
let plan_for_task = plan.clone();
let decision_for_task = decision.clone();
let execution_task = tokio::spawn(async move {
execute_execution_runtime_stream(
&state_for_task,
plan_for_task,
"trace-local-stream-client-error",
&decision_for_task,
"openai_chat_stream",
None,
Some(json!({
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
})),
)
.await
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = tunnel_protocol::FrameHeader::parse(&request_headers)
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body)
.expect("request body frame should parse");
assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
let response_meta = tunnel_protocol::ResponseMeta {
status: 200,
// Use a non-SSE content type so direct finalize prefetch stays enabled and the
// pre-body tunnel error is surfaced as a client-visible structured error response.
headers: vec![("content-type".to_string(), "application/json".to_string())],
};
let response_payload =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_HEADERS,
0,
&response_payload,
);
tunnel_app
.hub
.handle_proxy_frame(901, &mut response_headers_frame)
.await;
let original_error = "proxy disconnected before first upstream event";
let mut response_error_frame =
tunnel_protocol::encode_stream_error(request_header.stream_id, original_error);
tunnel_app
.hub
.handle_proxy_frame(901, &mut response_error_frame)
.await;
let response = execution_task
.await
.expect("execution task should complete")
.expect("execution should succeed")
.expect("execution should return a client response");
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_json: Value =
serde_json::from_slice(&body).expect("response body should decode as json");
let error_message = body_json
.get("error")
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.expect("response body should contain error.message");
assert_eq!(error_message, original_error);
assert!(
!error_message.contains("unexpected EOF during chunk size line"),
"client-facing response should preserve the original local tunnel error"
);
}
#[tokio::test]
async fn execute_execution_runtime_stream_emits_terminal_sse_error_event_after_body_started() {
let state = AppState::new().expect("app state should build");
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
902,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let plan = ExecutionPlan {
request_id: "req-client-stream-sse-error-1".into(),
candidate_id: Some("cand-client-stream-sse-error-1".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/chat".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({"stream": true})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-5".into()),
proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = test_decision();
let state_for_task = state.clone();
let plan_for_task = plan.clone();
let decision_for_task = decision.clone();
let execution_task = tokio::spawn(async move {
execute_execution_runtime_stream(
&state_for_task,
plan_for_task,
"trace-local-stream-sse-error",
&decision_for_task,
"openai_chat_stream",
None,
Some(json!({
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
})),
)
.await
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = tunnel_protocol::FrameHeader::parse(&request_headers)
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body)
.expect("request body frame should parse");
assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
let response_meta = tunnel_protocol::ResponseMeta {
status: 200,
headers: vec![("content-type".to_string(), "text/event-stream".to_string())],
};
let response_payload =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_HEADERS,
0,
&response_payload,
);
tunnel_app
.hub
.handle_proxy_frame(902, &mut response_headers_frame)
.await;
let mut response_body_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_BODY,
0,
b"data: hello\n\n",
);
tunnel_app
.hub
.handle_proxy_frame(902, &mut response_body_frame)
.await;
let response = execution_task
.await
.expect("execution task should complete")
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(
response
.headers()
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
let body_task = tokio::spawn(async move {
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
String::from_utf8(body.to_vec()).expect("response body should be utf8")
});
let original_error = "proxy disconnected while forwarding upstream body";
let mut response_error_frame =
tunnel_protocol::encode_stream_error(request_header.stream_id, original_error);
tunnel_app
.hub
.handle_proxy_frame(902, &mut response_error_frame)
.await;
let body = body_task.await.expect("body task should complete");
assert!(body.contains("data: hello\n\n"));
assert!(body.contains("event: aether.error\n"));
assert!(body.contains(original_error));
assert!(
!body.contains("unexpected EOF during chunk size line"),
"same-format SSE path should surface the original terminal error event"
);
}
}

394
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -11,6 +11,7 @@ use futures_util::{Stream, StreamExt};
use tracing::warn;
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::DirectUpstreamResponse;
use crate::execution_runtime::DirectUpstreamStreamExecution;
pub(crate) fn build_direct_execution_frame_stream(
@@ -26,14 +27,7 @@ pub(crate) fn build_direct_execution_frame_stream(
started_at,
} = execution;
let headers_frame = StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code,
headers,
},
};
match encode_stream_frame_ndjson(&headers_frame) {
match encode_headers_frame(status_code, headers) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -44,95 +38,109 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut upstream_bytes = 0u64;
let mut ttfb_ms = None;
let mut first_chunk_telemetry_emitted = false;
let mut bytes_stream = response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
match response {
DirectUpstreamResponse::Reqwest(response) => {
let mut bytes_stream = response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
}
Err(err) => {
let message = format_error_chain(&err);
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
}
if !first_chunk_telemetry_emitted {
let telemetry_frame = StreamFrame {
frame_type: StreamFrameType::Telemetry,
payload: StreamFramePayload::Telemetry {
telemetry: ExecutionTelemetry {
ttfb_ms,
elapsed_ms: ttfb_ms,
upstream_bytes: Some(upstream_bytes),
},
},
};
match encode_stream_frame_ndjson(&telemetry_frame) {
}
}
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
match response.next_chunk().await {
Ok(Some(chunk)) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
let frame = StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(base64::engine::general_purpose::STANDARD.encode(&chunk)),
text: None,
},
};
match encode_stream_frame_ndjson(&frame) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
Ok(None) => break,
Err(message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
}
Err(err) => {
let message = format_error_chain(&err);
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
"upstream body stream read error"
);
let frame = StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message,
upstream_status: Some(status_code),
retryable: false,
failover_recommended: false,
},
},
};
match encode_stream_frame_ndjson(&frame) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
}
}
let telemetry_frame = StreamFrame {
frame_type: StreamFrameType::Telemetry,
payload: StreamFramePayload::Telemetry {
telemetry: ExecutionTelemetry {
ttfb_ms,
elapsed_ms: Some(started_at.elapsed().as_millis() as u64),
upstream_bytes: Some(upstream_bytes),
},
},
};
match encode_stream_frame_ndjson(&telemetry_frame) {
match encode_telemetry_frame(
ttfb_ms,
Some(started_at.elapsed().as_millis() as u64),
upstream_bytes,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -146,6 +154,62 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
fn encode_headers_frame(
status_code: u16,
headers: std::collections::BTreeMap<String, String>,
) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code,
headers,
},
})
}
fn encode_telemetry_frame(
ttfb_ms: Option<u64>,
elapsed_ms: Option<u64>,
upstream_bytes: u64,
) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Telemetry,
payload: StreamFramePayload::Telemetry {
telemetry: ExecutionTelemetry {
ttfb_ms,
elapsed_ms,
upstream_bytes: Some(upstream_bytes),
},
},
})
}
fn encode_data_frame(chunk: &Bytes) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(base64::engine::general_purpose::STANDARD.encode(chunk)),
text: None,
},
})
}
fn encode_error_frame(status_code: u16, message: String) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message,
upstream_status: Some(status_code),
retryable: false,
failover_recommended: false,
},
},
})
}
fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
let mut message = err.to_string();
let mut source = err.source();
@@ -161,18 +225,36 @@ fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
mod tests {
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::extract::ws::Message;
use axum::routing::post;
use axum::{http::header, http::HeaderValue, Router};
use futures_util::StreamExt;
use serde_json::Value;
use tokio::sync::watch;
use super::build_direct_execution_frame_stream;
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
use crate::execution_runtime::transport::{
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
};
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
use crate::AppState;
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),
mode: Some("tunnel".into()),
node_id: Some("node-1".into()),
label: Some("relay-node".into()),
url: None,
extra: Some(serde_json::json!({"tunnel_base_url": base_url})),
}
}
#[tokio::test]
async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() {
@@ -358,4 +440,148 @@ mod tests {
"first telemetry frame should be emitted before the first data frame"
);
}
#[tokio::test]
async fn direct_execution_frame_stream_preserves_local_tunnel_stream_error_message() {
let state = AppState::new().expect("app state should build");
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
801,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let plan = ExecutionPlan {
request_id: "req-local-stream-error-1".into(),
candidate_id: Some("cand-local-stream-error-1".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/chat".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(serde_json::json!({"stream": true})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-5".into()),
proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let state_for_task = state.clone();
let plan_for_task = plan.clone();
let execution_task = tokio::spawn(async move {
execute_stream_plan_via_local_tunnel(&state_for_task, &plan_for_task).await
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = tunnel_protocol::FrameHeader::parse(&request_headers)
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body)
.expect("request body frame should parse");
assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
let response_meta = tunnel_protocol::ResponseMeta {
status: 200,
headers: vec![("content-type".to_string(), "text/event-stream".to_string())],
};
let response_payload =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_HEADERS,
0,
&response_payload,
);
tunnel_app
.hub
.handle_proxy_frame(801, &mut response_headers_frame)
.await;
let execution = execution_task
.await
.expect("execution task should complete")
.expect("local tunnel execution should resolve")
.expect("local tunnel execution should be available");
let frame_task = tokio::spawn(async move {
build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
.collect::<Vec<_>>()
.await
.into_iter()
.map(|bytes| String::from_utf8(bytes.to_vec()).expect("frame should be utf8"))
.collect::<Vec<_>>()
});
let mut response_body_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_BODY,
0,
b"data: hello\n\n",
);
tunnel_app
.hub
.handle_proxy_frame(801, &mut response_body_frame)
.await;
let original_error = "proxy disconnected while forwarding upstream body";
let mut response_error_frame =
tunnel_protocol::encode_stream_error(request_header.stream_id, original_error);
tunnel_app
.hub
.handle_proxy_frame(801, &mut response_error_frame)
.await;
let frames = frame_task.await.expect("frame task should complete");
let parsed_frames = frames
.iter()
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
.collect::<Vec<_>>();
assert!(
parsed_frames
.iter()
.any(|frame| { frame.get("type").and_then(Value::as_str) == Some("data") }),
"stream should contain at least one data frame before the error"
);
let error_message = parsed_frames
.iter()
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("error"))
.and_then(|frame| frame.get("payload"))
.and_then(|payload| payload.get("error"))
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.expect("error frame should include a message");
assert_eq!(error_message, original_error);
assert!(
!error_message.contains("unexpected EOF during chunk size line"),
"local tunnel path should preserve the original proxy error text"
);
}
}

307
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -9,6 +9,7 @@ use aether_contracts::{
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
use aether_http::{apply_http_client_config, HttpClientConfig};
use axum::body::Bytes;
use base64::Engine as _;
use flate2::read::{DeflateDecoder, GzDecoder};
use flate2::write::GzEncoder;
@@ -29,6 +30,7 @@ use crate::execution_runtime::remote_compat::execute_sync_plan_via_remote_execut
use crate::frontdoor_loop_guard::{
configured_gateway_frontdoor_base_url, gateway_frontdoor_self_loop_guard_error,
};
use crate::tunnel::{self, tunnel_protocol};
use crate::{AppState, GatewayError};
const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
@@ -135,13 +137,17 @@ struct ExecutionTransportControls {
http1_only: bool,
}
#[derive(Debug)]
pub(crate) enum DirectUpstreamResponse {
Reqwest(reqwest::Response),
LocalTunnel(tunnel::DirectRelayResponse),
}
pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) request_id: String,
pub(crate) candidate_id: Option<String>,
pub(crate) status_code: u16,
pub(crate) headers: BTreeMap<String, String>,
pub(crate) response: reqwest::Response,
pub(crate) response: DirectUpstreamResponse,
pub(crate) started_at: Instant,
}
@@ -225,7 +231,7 @@ impl DirectSyncExecutionRuntime {
candidate_id: plan.candidate_id,
status_code,
headers,
response,
response: DirectUpstreamResponse::Reqwest(response),
started_at,
})
}
@@ -252,6 +258,12 @@ pub(crate) async fn execute_sync_plan(
}
}
if resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).is_some() {
return execute_sync_plan_via_local_tunnel(state, plan)
.await
.map_err(|err| GatewayError::Internal(err.to_string()));
}
let _ = state;
let _ = trace_id;
DirectSyncExecutionRuntime::new()
@@ -260,6 +272,146 @@ pub(crate) async fn execute_sync_plan(
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn execute_stream_plan_via_local_tunnel(
state: &AppState,
plan: &ExecutionPlan,
) -> Result<Option<DirectUpstreamStreamExecution>, ExecutionRuntimeTransportError> {
let Some(node_id) = resolve_local_tunnel_node_id(state, plan.proxy.as_ref()) else {
return Ok(None);
};
if let Some(detail) = gateway_frontdoor_self_loop_guard_error(plan.url.as_str()) {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(detail));
}
let body_bytes = build_request_body(plan)?;
let transport_controls = resolve_execution_transport_controls(&plan.headers);
let headers = build_request_headers(
&plan.headers,
plan.content_encoding.as_deref(),
plan.body.body_bytes_b64.is_some(),
)?;
let headers = append_execution_loop_guard_header(headers);
let started_at = Instant::now();
let response = state
.tunnel
.open_direct_relay_stream(
&node_id,
build_direct_tunnel_request_meta(plan, &headers, transport_controls),
Bytes::from(body_bytes),
)
.await
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
Ok(Some(DirectUpstreamStreamExecution {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response: DirectUpstreamResponse::LocalTunnel(response),
started_at,
}))
}
async fn execute_sync_plan_via_local_tunnel(
state: &AppState,
plan: &ExecutionPlan,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
let node_id = resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).ok_or_else(|| {
ExecutionRuntimeTransportError::RelayError("local tunnel node unavailable".to_string())
})?;
if let Some(detail) = gateway_frontdoor_self_loop_guard_error(plan.url.as_str()) {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(detail));
}
let body_bytes = build_request_body(plan)?;
let transport_controls = resolve_execution_transport_controls(&plan.headers);
let headers = build_request_headers(
&plan.headers,
plan.content_encoding.as_deref(),
plan.body.body_bytes_b64.is_some(),
)?;
let headers = append_execution_loop_guard_header(headers);
let started_at = Instant::now();
let mut response = state
.tunnel
.open_direct_relay_stream(
&node_id,
build_direct_tunnel_request_meta(plan, &headers, transport_controls),
Bytes::from(body_bytes),
)
.await
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
let mut body_bytes = Vec::new();
while let Some(chunk) = response
.next_chunk()
.await
.map_err(ExecutionRuntimeTransportError::UpstreamRequest)?
{
body_bytes.extend_from_slice(&chunk);
}
let decoded_body_bytes =
decode_response_body_bytes(&headers, &body_bytes).unwrap_or_else(|| body_bytes.clone());
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let upstream_bytes = body_bytes.len() as u64;
let body = if body_bytes.is_empty() {
None
} else if plan.stream {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
} else if response_body_is_json(&headers, &decoded_body_bytes) {
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
})
} else {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
};
Ok(ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: Some(ttfb_ms),
elapsed_ms: Some(elapsed_ms),
upstream_bytes: Some(upstream_bytes),
}),
error: None,
})
}
fn build_direct_tunnel_request_meta(
plan: &ExecutionPlan,
headers: &HeaderMap,
transport_controls: ExecutionTransportControls,
) -> tunnel_protocol::RequestMeta {
tunnel_protocol::RequestMeta {
method: plan.method.clone(),
url: plan.url.clone(),
headers: header_map_to_string_map(headers).into_iter().collect(),
timeout: resolve_relay_timeout_seconds(plan),
follow_redirects: transport_controls.follow_redirects,
http1_only: transport_controls.http1_only,
}
}
async fn send_request(
plan: &ExecutionPlan,
body_bytes: Vec<u8>,
@@ -526,6 +678,11 @@ fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
None
}
fn resolve_local_tunnel_node_id(state: &AppState, proxy: Option<&ProxySnapshot>) -> Option<String> {
let node_id = resolve_tunnel_node_id(proxy)?;
state.tunnel.has_local_proxy(&node_id).then_some(node_id)
}
fn build_client(
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
proxy: Option<&ProxySnapshot>,
@@ -739,6 +896,13 @@ fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
header_map_to_string_map(headers)
}
fn collect_tunnel_response_headers(headers: &[(String, String)]) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| (name.to_ascii_lowercase(), value.clone()))
.collect()
}
fn decode_response_body_bytes(
headers: &BTreeMap<String, String>,
body_bytes: &[u8],
@@ -782,22 +946,29 @@ fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8])
mod tests {
use std::collections::BTreeMap;
use std::io::Read;
use std::sync::Arc;
use aether_contracts::{
ExecutionPlan, ExecutionTimeouts, RequestBody, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
use axum::body::Bytes;
use axum::extract::ws::Message;
use axum::extract::Path;
use axum::routing::post;
use axum::{Json, Router};
use serde_json::json;
use tokio::sync::watch;
use super::{build_client, DirectSyncExecutionRuntime, ExecutionTransportControls};
use super::{
build_client, execute_sync_plan, DirectSyncExecutionRuntime, ExecutionTransportControls,
};
use crate::frontdoor_loop_guard::{
frontdoor_self_loop_public_ai_path, gateway_frontdoor_self_loop_guard_error_with_port,
gateway_frontdoor_self_loop_guard_matches_with_port,
};
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
use crate::AppState;
#[test]
fn gateway_frontdoor_self_loop_guard_matches_loopback_public_ai_route() {
@@ -1018,6 +1189,134 @@ mod tests {
);
}
#[tokio::test]
async fn execute_sync_plan_prefers_local_tunnel_stream_over_http_relay_loopback() {
let state = AppState::new().expect("app state should build");
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
701,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let plan = ExecutionPlan {
request_id: "req-local-tunnel-1".into(),
candidate_id: Some("cand-local-tunnel-1".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/chat".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
stream: false,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-4.1".into()),
proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let state_for_task = state.clone();
let plan_for_task = plan.clone();
let execution_task = tokio::spawn(async move {
execute_sync_plan(&state_for_task, Some("trace-local-tunnel"), &plan_for_task).await
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = tunnel_protocol::FrameHeader::parse(&request_headers)
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let request_meta_payload =
tunnel_protocol::decode_payload(&request_headers, &request_header)
.expect("request meta payload should decode");
let request_meta =
serde_json::from_slice::<tunnel_protocol::RequestMeta>(&request_meta_payload)
.expect("request meta should decode");
assert_eq!(request_meta.method, "POST");
assert_eq!(request_meta.url, "https://example.com/chat");
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body)
.expect("request body frame should parse");
assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
let request_body_payload =
tunnel_protocol::decode_payload(&request_body, &request_body_header)
.expect("request body payload should decode");
let request_json = serde_json::from_slice::<serde_json::Value>(&request_body_payload)
.expect("request body should decode");
assert_eq!(request_json["model"], "gpt-4.1");
let response_meta = tunnel_protocol::ResponseMeta {
status: 200,
headers: vec![("content-type".to_string(), "application/json".to_string())],
};
let response_payload =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_HEADERS,
0,
&response_payload,
);
tunnel_app
.hub
.handle_proxy_frame(701, &mut response_headers_frame)
.await;
let mut response_body_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_BODY,
0,
br#"{"local_tunnel":true}"#,
);
tunnel_app
.hub
.handle_proxy_frame(701, &mut response_body_frame)
.await;
let mut response_end_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::STREAM_END,
0,
&[],
);
tunnel_app
.hub
.handle_proxy_frame(701, &mut response_end_frame)
.await;
let result = execution_task
.await
.expect("execution task should complete")
.expect("local tunnel execution should succeed");
assert_eq!(result.status_code, 200);
assert_eq!(
result.body.and_then(|body| body.json_body),
Some(json!({"local_tunnel": true}))
);
}
#[tokio::test]
async fn direct_sync_execution_runtime_disables_redirects_by_default() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")

225
apps/aether-gateway/src/tunnel/embedded/local_relay.rs
View File

@@ -11,9 +11,11 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
use axum::response::IntoResponse;
use bytes::BytesMut;
use futures_util::StreamExt;
use tokio::sync::mpsc;
use tracing::warn;
use crate::api::response::apply_streaming_response_headers;
use crate::headers::should_skip_response_header;
use crate::maintenance::record_proxy_upgrade_traffic_success;
use super::hub::{LocalBodyEvent, LocalStream};
@@ -38,6 +40,114 @@ impl Drop for StreamGuard {
}
}
pub(crate) struct DirectRelayResponse {
status: u16,
headers: Vec<(String, String)>,
body_rx: mpsc::Receiver<LocalBodyEvent>,
request_guard: StreamGuard,
_request_permit: Option<AdmissionPermit>,
}
impl DirectRelayResponse {
pub(crate) fn status(&self) -> u16 {
self.status
}
pub(crate) fn headers(&self) -> &[(String, String)] {
&self.headers
}
pub(crate) async fn next_chunk(&mut self) -> Result<Option<Bytes>, String> {
let event = self.body_rx.recv().await;
match event {
Some(LocalBodyEvent::Chunk(chunk)) => Ok(Some(chunk)),
Some(LocalBodyEvent::End) | None => {
self.request_guard.finished = true;
Ok(None)
}
Some(LocalBodyEvent::Error(error)) => {
self.request_guard.finished = true;
Err(error)
}
}
}
}
pub(crate) async fn open_direct_relay_stream(
state: &AppState,
node_id: &str,
meta: protocol::RequestMeta,
body: Bytes,
) -> Result<DirectRelayResponse, String> {
let request_permit = state
.try_acquire_request_permit()
.await
.map_err(map_request_admission_error)?;
let stream = state
.hub
.open_local_stream(node_id, &meta)
.map_err(|error| format!("connect: {error}"))?;
if let Err(error) = state.hub.push_local_request_body(stream.id, body, true) {
state.hub.cancel_local_stream(stream.id, &error);
return Err(format!("connect: {error}"));
}
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
let response_head = match stream.wait_headers(wait_timeout).await {
Ok(response) => response,
Err(error) => {
state.hub.cancel_local_stream(stream.id, &error);
return Err(format!("timeout: {error}"));
}
};
if let Err(error) = record_proxy_upgrade_traffic_success(state.data.as_ref(), node_id).await {
warn!(
node_id = %node_id,
error = %error,
"failed to record proxy upgrade traffic confirmation"
);
}
let Some(body_rx) = stream.take_body_receiver() else {
state
.hub
.cancel_local_stream(stream.id, "missing relay response body receiver");
return Err("relay: missing relay response body receiver".to_string());
};
Ok(DirectRelayResponse {
status: response_head.status,
headers: response_head.headers,
body_rx,
request_guard: StreamGuard {
hub: state.hub.clone(),
stream_id: stream.id,
finished: false,
},
_request_permit: request_permit,
})
}
fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
match error {
super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Saturated {
..
})
| super::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::Saturated { .. },
)
| super::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::Unavailable { .. },
) => "overloaded: hub relay overloaded".to_string(),
super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Closed {
..
}) => "overloaded: hub relay gate closed".to_string(),
super::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::InvalidConfiguration(_),
) => "overloaded: hub relay distributed gate invalid".to_string(),
}
}
pub async fn relay_request(
Path(node_id): Path<String>,
State(state): State<AppState>,
@@ -319,6 +429,9 @@ fn try_decode_envelope_meta(
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
for (name, value) in headers {
if should_skip_local_relay_response_header(name) {
continue;
}
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
continue;
};
@@ -329,6 +442,10 @@ fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
}
}
fn should_skip_local_relay_response_header(name: &str) -> bool {
should_skip_response_header(name) || name.eq_ignore_ascii_case("content-length")
}
fn tunnel_error_response(status: StatusCode, kind: &str, message: &str) -> Response<Body> {
let mut builder = Response::builder().status(status);
if let Some(headers) = builder.headers_mut() {
@@ -609,4 +726,112 @@ mod tests {
assert!(tracked_nodes[0]["version_confirmed_at_unix_secs"].is_u64());
assert!(tracked_nodes[0]["traffic_confirmed_at_unix_secs"].is_u64());
}
#[tokio::test]
async fn relay_strips_hop_by_hop_and_stale_length_headers_from_proxy_response() {
let state = test_app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
state.hub.register_proxy(Arc::new(ProxyConn::new(
501,
"node-123".to_string(),
"Node 123".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let meta = protocol::RequestMeta {
method: "GET".to_string(),
url: "https://example.com/headers".to_string(),
headers: HashMap::new(),
timeout: 30,
follow_redirects: None,
http1_only: false,
};
let request = Request::builder()
.body(Body::from(encode_relay_envelope(&meta, &[])))
.expect("request should build");
let relay_state = state.clone();
let relay_task = tokio::spawn(async move {
relay_request(
Path("node-123".to_string()),
State(relay_state),
ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 4242))),
request,
)
.await
.into_response()
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = protocol::FrameHeader::parse(&request_headers)
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, protocol::REQUEST_HEADERS);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header =
protocol::FrameHeader::parse(&request_body).expect("request body frame should parse");
assert_eq!(request_body_header.msg_type, protocol::REQUEST_BODY);
let response_meta = protocol::ResponseMeta {
status: 200,
headers: vec![
("content-length".to_string(), "999".to_string()),
("transfer-encoding".to_string(), "chunked".to_string()),
("connection".to_string(), "keep-alive".to_string()),
("content-type".to_string(), "text/plain".to_string()),
(
"x-proxy-timing".to_string(),
"{\"mode\":\"tunnel\"}".to_string(),
),
],
};
let response_payload =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = protocol::encode_frame(
request_header.stream_id,
protocol::RESPONSE_HEADERS,
0,
&response_payload,
);
state
.hub
.handle_proxy_frame(501, &mut response_headers_frame)
.await;
let mut response_end_frame =
protocol::encode_frame(request_header.stream_id, protocol::STREAM_END, 0, &[]);
state
.hub
.handle_proxy_frame(501, &mut response_end_frame)
.await;
let response = relay_task.await.expect("relay task should complete");
assert_eq!(response.status(), StatusCode::OK);
assert!(response.headers().get("content-length").is_none());
assert!(response.headers().get("transfer-encoding").is_none());
assert!(response.headers().get("connection").is_none());
assert_eq!(
response
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok()),
Some("text/plain")
);
assert_eq!(
response
.headers()
.get("x-proxy-timing")
.and_then(|value| value.to_str().ok()),
Some("{\"mode\":\"tunnel\"}")
);
}
}

1
apps/aether-gateway/src/tunnel/embedded/mod.rs
View File

@@ -24,6 +24,7 @@ use crate::data::GatewayDataState;
pub use control_plane::ControlPlaneClient;
pub use hub::{ConnConfig, HubRouter, LocalBodyEvent, ProxyConn};
pub use local_relay::relay_request;
pub(crate) use local_relay::{open_direct_relay_stream, DirectRelayResponse};
#[derive(Clone)]
pub struct AppState {

10
apps/aether-gateway/src/tunnel/mod.rs
View File

@@ -33,6 +33,7 @@ use super::error::GatewayError;
use super::headers::{extract_or_generate_trace_id, should_skip_request_header};
use super::AppState;
pub(crate) use embedded::DirectRelayResponse;
pub(crate) use embedded::ProxyConn as TunnelProxyConn;
pub use embedded::{
build_router_with_state as build_tunnel_runtime_router_with_state, protocol as tunnel_protocol,
@@ -434,6 +435,15 @@ impl EmbeddedTunnelState {
self.inner.hub.has_local_proxy(node_id)
}
pub(crate) async fn open_direct_relay_stream(
&self,
node_id: &str,
meta: tunnel_protocol::RequestMeta,
body: Bytes,
) -> Result<DirectRelayResponse, String> {
embedded::open_direct_relay_stream(&self.inner, node_id, meta, body).await
}
pub(crate) fn request_close_all_proxies(&self) -> usize {
self.inner.hub.request_close_all_proxies()
}