feat: 引入 Rust executor/gateway sidecar 及 Python 侧双后端适配

- 新增 Rust workspace crates: aether-contracts, aether-executor, aether-gateway
- aether-executor: 支持 Unix Socket/TCP 双传输模式,处理同步/流式上游请求
- aether-gateway: 作为本地主入口代理,集成 /api/internal/gateway/resolve 认证预解析
- Python 侧新增 ExecutionPlan 契约和 RustExecutorClient,各 handler 支持
  executor_backend=rust 时将可序列化请求转发给 Rust executor 执行
- 重构 dev.sh 支持 executor/gateway 进程编排与生命周期管理
- 新增 internal gateway 路由,提供 resolve/passthrough 端点
- handler 层(chat/cli/video/endpoint_checker 等)全面适配 Rust executor 回退逻辑
- pipeline 层支持 trusted auth context 跳过重复认证
- 新增 Rust CI workflow 及对应测试用例
This commit is contained in:
fawney19
2026-03-21 12:57:09 +08:00
parent 46737d32f8
commit d735b6316f
79 changed files with 19032 additions and 522 deletions

View File

@@ -0,0 +1,17 @@
pub(crate) const TRACE_ID_HEADER: &str = "x-trace-id";
pub(crate) const FORWARDED_HOST_HEADER: &str = "x-forwarded-host";
pub(crate) const FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
pub(crate) const FORWARDED_PROTO_HEADER: &str = "x-forwarded-proto";
pub(crate) const GATEWAY_HEADER: &str = "x-aether-gateway";
pub(crate) const CONTROL_ROUTE_CLASS_HEADER: &str = "x-aether-control-route-class";
pub(crate) const CONTROL_ROUTE_FAMILY_HEADER: &str = "x-aether-control-route-family";
pub(crate) const CONTROL_ROUTE_KIND_HEADER: &str = "x-aether-control-route-kind";
pub(crate) const CONTROL_EXECUTOR_HEADER: &str = "x-aether-control-executor-candidate";
pub(crate) const CONTROL_ENDPOINT_SIGNATURE_HEADER: &str = "x-aether-control-endpoint-signature";
pub(crate) const CONTROL_EXECUTED_HEADER: &str = "x-aether-control-executed";
pub(crate) const CONTROL_ACTION_HEADER: &str = "x-aether-control-action";
pub(crate) const CONTROL_ACTION_PROXY_PUBLIC: &str = "proxy_public";
pub(crate) const TRUSTED_AUTH_USER_ID_HEADER: &str = "x-aether-auth-user-id";
pub(crate) const TRUSTED_AUTH_API_KEY_ID_HEADER: &str = "x-aether-auth-api-key-id";
pub(crate) const TRUSTED_AUTH_BALANCE_HEADER: &str = "x-aether-auth-balance-remaining";
pub(crate) const TRUSTED_AUTH_ACCESS_ALLOWED_HEADER: &str = "x-aether-auth-access-allowed";

View File

@@ -0,0 +1,299 @@
use axum::body::{Body, Bytes};
use axum::http::{Response, StatusCode, Uri};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use serde::{Deserialize, Serialize};
use serde_json::json;
use crate::gateway::constants::*;
use crate::gateway::headers::{
collect_control_headers, header_equals, header_value_str, header_value_u64, is_json_request,
};
use crate::gateway::{build_client_response, AppState, GatewayError};
#[derive(Debug, Serialize)]
struct GatewayControlResolveRequest {
trace_id: String,
method: String,
path: String,
query_string: Option<String>,
headers: std::collections::BTreeMap<String, String>,
has_body: bool,
content_type: Option<String>,
content_length: Option<u64>,
}
#[derive(Debug, Serialize)]
struct GatewayControlExecuteRequest {
trace_id: String,
method: String,
path: String,
query_string: Option<String>,
headers: std::collections::BTreeMap<String, String>,
body_json: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
body_base64: Option<String>,
auth_context: Option<GatewayControlAuthContext>,
}
#[derive(Debug, Deserialize)]
struct GatewayControlResolveResponse {
action: String,
public_path: Option<String>,
public_query_string: Option<String>,
route_class: Option<String>,
route_family: Option<String>,
route_kind: Option<String>,
auth_endpoint_signature: Option<String>,
executor_candidate: Option<bool>,
auth_context: Option<GatewayControlAuthContext>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub(crate) struct GatewayControlAuthContext {
pub(crate) user_id: String,
pub(crate) api_key_id: String,
pub(crate) balance_remaining: Option<f64>,
pub(crate) access_allowed: bool,
}
#[derive(Debug, Clone)]
pub(crate) struct GatewayControlDecision {
pub(crate) public_path: String,
pub(crate) public_query_string: Option<String>,
pub(crate) route_class: Option<String>,
pub(crate) route_family: Option<String>,
pub(crate) route_kind: Option<String>,
pub(crate) auth_endpoint_signature: Option<String>,
pub(crate) executor_candidate: bool,
pub(crate) auth_context: Option<GatewayControlAuthContext>,
}
impl GatewayControlDecision {
pub(crate) fn proxy_path_and_query(&self) -> String {
if let Some(query) = self
.public_query_string
.as_deref()
.filter(|value| !value.is_empty())
{
format!("{}?{}", self.public_path, query)
} else {
self.public_path.clone()
}
}
}
pub(crate) async fn resolve_control_route(
state: &AppState,
method: &http::Method,
uri: &Uri,
headers: &http::HeaderMap,
trace_id: &str,
) -> Result<Option<GatewayControlDecision>, GatewayError> {
let Some(control_base_url) = state.control_base_url.as_deref() else {
return Ok(None);
};
let path = uri.path();
if !should_consult_control_api(path) {
return Ok(None);
}
let control_request = GatewayControlResolveRequest {
trace_id: trace_id.to_string(),
method: method.to_string(),
path: path.to_string(),
query_string: uri.query().map(ToOwned::to_owned),
headers: collect_control_headers(headers),
has_body: header_value_u64(headers, http::header::CONTENT_LENGTH.as_str()).unwrap_or(0) > 0
|| headers.contains_key(http::header::CONTENT_TYPE),
content_type: header_value_str(headers, http::header::CONTENT_TYPE.as_str()),
content_length: header_value_u64(headers, http::header::CONTENT_LENGTH.as_str()),
};
let response = state
.client
.post(format!("{control_base_url}/api/internal/gateway/resolve"))
.header(TRACE_ID_HEADER, trace_id)
.json(&control_request)
.send()
.await
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
let response = response
.error_for_status()
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
let payload: GatewayControlResolveResponse = response
.json()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if payload.action != "proxy_public" {
return Err(GatewayError::Internal(format!(
"unsupported gateway control action: {}",
payload.action
)));
}
Ok(Some(GatewayControlDecision {
public_path: payload.public_path.unwrap_or_else(|| path.to_string()),
public_query_string: payload
.public_query_string
.or_else(|| uri.query().map(ToOwned::to_owned)),
route_class: payload.route_class,
route_family: payload.route_family,
route_kind: payload.route_kind,
auth_endpoint_signature: payload.auth_endpoint_signature,
executor_candidate: payload.executor_candidate.unwrap_or(false),
auth_context: payload.auth_context,
}))
}
fn is_stream_route(path: &str) -> bool {
path.contains(":streamGenerateContent")
}
fn is_video_route(decision: &GatewayControlDecision) -> bool {
decision.route_kind.as_deref() == Some("video")
}
fn is_files_route(decision: &GatewayControlDecision) -> bool {
decision.route_kind.as_deref() == Some("files")
&& decision.route_family.as_deref() == Some("gemini")
}
pub(crate) async fn maybe_execute_via_control(
state: &AppState,
parts: &http::request::Parts,
body_bytes: Bytes,
trace_id: &str,
decision: Option<&GatewayControlDecision>,
) -> Result<Option<Response<Body>>, GatewayError> {
let Some(control_base_url) = state.control_base_url.as_deref() else {
return Ok(None);
};
let Some(decision) = decision else {
return Ok(None);
};
if !decision.executor_candidate {
return Ok(None);
}
if decision.route_class.as_deref() != Some("ai_public") {
return Ok(None);
}
let is_files_route = is_files_route(decision);
let is_video_route = is_video_route(decision);
if is_files_route || is_video_route {
if !matches!(
parts.method,
http::Method::GET | http::Method::POST | http::Method::DELETE
) {
return Ok(None);
}
} else if parts.method != http::Method::POST || !is_json_request(&parts.headers) {
return Ok(None);
}
let body_json = if is_json_request(&parts.headers) {
match serde_json::from_slice::<serde_json::Value>(&body_bytes) {
Ok(value) if value.is_object() => value,
_ if !is_files_route && !is_video_route => return Ok(None),
_ => json!({}),
}
} else if is_video_route && body_bytes.is_empty() {
json!({})
} else if is_files_route {
json!({})
} else {
return Ok(None);
};
let is_stream_request = if is_files_route || is_video_route {
false
} else {
is_stream_route(parts.uri.path())
|| body_json
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false)
};
let control_endpoint = if is_stream_request {
"execute-stream"
} else {
"execute-sync"
};
let request_payload = GatewayControlExecuteRequest {
trace_id: trace_id.to_string(),
method: parts.method.to_string(),
path: parts.uri.path().to_string(),
query_string: parts.uri.query().map(ToOwned::to_owned),
headers: collect_control_headers(&parts.headers),
body_json,
body_base64: if is_files_route && !body_bytes.is_empty() {
Some(BASE64_STANDARD.encode(&body_bytes))
} else {
None
},
auth_context: decision.auth_context.clone(),
};
let response = state
.client
.post(format!(
"{control_base_url}/api/internal/gateway/{control_endpoint}"
))
.header(TRACE_ID_HEADER, trace_id)
.json(&request_payload)
.send()
.await
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
if response.status() == StatusCode::CONFLICT
&& header_equals(response.headers(), CONTROL_ACTION_HEADER, "proxy_public")
{
return Ok(None);
}
if !header_equals(response.headers(), CONTROL_EXECUTED_HEADER, "true") {
return Ok(None);
}
Ok(Some(build_client_response(
response,
trace_id,
Some(decision),
)?))
}
fn should_consult_control_api(path: &str) -> bool {
matches!(
path,
"/v1/chat/completions" | "/v1/messages" | "/v1/responses" | "/v1/responses/compact"
) || path.starts_with("/v1/videos")
|| path == "/upload/v1beta/files"
|| path.starts_with("/v1beta/files")
|| is_gemini_models_route(path)
|| is_gemini_operation_route(path)
}
fn is_gemini_models_route(path: &str) -> bool {
(path.starts_with("/v1/models/") || path.starts_with("/v1beta/models/"))
&& (path.contains(":generateContent")
|| path.contains(":streamGenerateContent")
|| path.contains(":predictLongRunning"))
}
fn is_gemini_operation_route(path: &str) -> bool {
(path.starts_with("/v1beta/models/") && path.contains("/operations/"))
|| path == "/v1beta/operations"
|| path.starts_with("/v1beta/operations/")
}

View File

@@ -0,0 +1,68 @@
use axum::body::Body;
use axum::http::{Response, StatusCode};
use axum::response::IntoResponse;
use axum::Json;
use serde_json::json;
use tracing::warn;
use crate::gateway::constants::*;
use crate::gateway::insert_header_if_missing;
#[derive(Debug)]
pub(crate) enum GatewayError {
UpstreamUnavailable { trace_id: String, message: String },
ControlUnavailable { trace_id: String, message: String },
Internal(String),
}
impl IntoResponse for GatewayError {
fn into_response(self) -> Response<Body> {
match self {
Self::UpstreamUnavailable { trace_id, message } => {
warn!(trace_id = %trace_id, error = %message, "gateway upstream unavailable");
let body = Json(json!({
"error": {
"message": "gateway upstream unavailable",
"trace_id": trace_id,
}
}));
let mut response = (StatusCode::BAD_GATEWAY, body).into_response();
let _ =
insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, &trace_id);
let _ = insert_header_if_missing(
response.headers_mut(),
GATEWAY_HEADER,
"rust-phase3b",
);
response
}
Self::ControlUnavailable { trace_id, message } => {
warn!(trace_id = %trace_id, error = %message, "gateway control unavailable");
let body = Json(json!({
"error": {
"message": "gateway control unavailable",
"trace_id": trace_id,
}
}));
let mut response = (StatusCode::BAD_GATEWAY, body).into_response();
let _ =
insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, &trace_id);
let _ = insert_header_if_missing(
response.headers_mut(),
GATEWAY_HEADER,
"rust-phase3b",
);
response
}
Self::Internal(message) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({
"error": {
"message": message,
}
})),
)
.into_response(),
}
}
}

View File

@@ -0,0 +1,396 @@
use std::collections::BTreeMap;
use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, StreamFrame, StreamFramePayload};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::http::Response;
use base64::Engine as _;
use futures_util::{StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::io::StreamReader;
use tracing::warn;
use crate::gateway::constants::*;
use crate::gateway::headers::{collect_control_headers, header_equals};
use crate::gateway::{
build_client_response, build_client_response_from_parts, AppState, GatewayControlAuthContext,
GatewayControlDecision, GatewayError,
};
const GEMINI_FILES_DOWNLOAD_PLAN_KIND: &str = "gemini_files_download";
const OPENAI_VIDEO_CONTENT_PLAN_KIND: &str = "openai_video_content";
const EXECUTOR_STREAM_ACTION: &str = "executor_stream";
const MAX_ERROR_BODY_BYTES: usize = 16_384;
#[derive(Debug, Serialize)]
struct GatewayControlPlanRequest {
trace_id: String,
method: String,
path: String,
query_string: Option<String>,
headers: BTreeMap<String, String>,
body_json: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
body_base64: Option<String>,
auth_context: Option<GatewayControlAuthContext>,
}
#[derive(Debug, Deserialize)]
struct GatewayControlPlanResponse {
action: String,
#[serde(default)]
plan_kind: Option<String>,
#[serde(default)]
plan: Option<ExecutionPlan>,
}
pub(crate) async fn maybe_execute_via_executor_stream(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: Option<&GatewayControlDecision>,
) -> Result<Option<Response<Body>>, GatewayError> {
let Some(control_base_url) = state.control_base_url.as_deref() else {
return Ok(None);
};
let Some(executor_base_url) = state.executor_base_url.as_deref() else {
return Ok(None);
};
let Some(decision) = decision else {
return Ok(None);
};
let Some(plan_kind) = resolve_direct_executor_stream_plan_kind(parts, decision) else {
return Ok(None);
};
let request_payload = GatewayControlPlanRequest {
trace_id: trace_id.to_string(),
method: parts.method.to_string(),
path: parts.uri.path().to_string(),
query_string: parts.uri.query().map(ToOwned::to_owned),
headers: collect_control_headers(&parts.headers),
body_json: json!({}),
body_base64: None,
auth_context: decision.auth_context.clone(),
};
let response = state
.client
.post(format!(
"{control_base_url}/api/internal/gateway/plan-stream"
))
.header(TRACE_ID_HEADER, trace_id)
.json(&request_payload)
.send()
.await
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
if response.status() == http::StatusCode::CONFLICT
&& header_equals(
response.headers(),
CONTROL_ACTION_HEADER,
CONTROL_ACTION_PROXY_PUBLIC,
)
{
return Ok(None);
}
if header_equals(response.headers(), CONTROL_EXECUTED_HEADER, "true")
&& response.status() != http::StatusCode::OK
{
return Ok(Some(build_client_response(
response,
trace_id,
Some(decision),
)?));
}
let response = response
.error_for_status()
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
let payload: GatewayControlPlanResponse = response
.json()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if payload.action != EXECUTOR_STREAM_ACTION {
return Ok(None);
}
if payload.plan_kind.as_deref() != Some(plan_kind) {
return Ok(None);
}
let Some(plan) = payload.plan else {
return Err(GatewayError::Internal(
"gateway plan response missing execution plan".to_string(),
));
};
execute_executor_stream(
state,
executor_base_url,
plan,
trace_id,
decision,
plan_kind,
)
.await
}
fn resolve_direct_executor_stream_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
if parts.method != http::Method::GET || decision.route_class.as_deref() != Some("ai_public") {
return None;
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("files")
&& parts.uri.path().ends_with(":download")
{
return Some(GEMINI_FILES_DOWNLOAD_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("video")
&& parts.uri.path().ends_with("/content")
{
return Some(OPENAI_VIDEO_CONTENT_PLAN_KIND);
}
None
}
async fn execute_executor_stream(
state: &AppState,
executor_base_url: &str,
plan: ExecutionPlan,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<Option<Response<Body>>, GatewayError> {
let response = match state
.client
.post(format!("{executor_base_url}/v1/execute/stream"))
.header(TRACE_ID_HEADER, trace_id)
.json(&plan)
.send()
.await
{
Ok(response) => response,
Err(err) => {
warn!(trace_id = %trace_id, error = %err, "gateway direct executor stream unavailable");
return Ok(None);
}
};
if response.status() != http::StatusCode::OK {
return Ok(Some(build_client_response(
response,
trace_id,
Some(decision),
)?));
}
let stream = response
.bytes_stream()
.map_err(|err| IoError::other(err.to_string()));
let reader = StreamReader::new(stream);
let mut lines = FramedRead::new(reader, LinesCodec::new());
let first_frame = read_next_frame(&mut lines).await?.ok_or_else(|| {
GatewayError::Internal("executor stream ended before headers frame".to_string())
})?;
let StreamFramePayload::Headers {
status_code,
headers,
} = first_frame.payload
else {
return Err(GatewayError::Internal(
"executor stream must start with headers frame".to_string(),
));
};
if status_code >= 400 {
let error_body = collect_error_body(&mut lines).await?;
return Ok(Some(build_executor_error_response(
trace_id,
decision,
plan_kind,
status_code,
headers,
error_body,
)?));
}
let trace_id_owned = trace_id.to_string();
let body_stream = stream! {
loop {
let next_frame = match read_next_frame(&mut lines).await {
Ok(frame) => frame,
Err(err) => {
warn!(trace_id = %trace_id_owned, error = %format!("{err:?}"), "gateway failed to decode executor stream frame");
break;
}
};
let Some(frame) = next_frame else {
break;
};
match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => {
if let Some(chunk_b64) = chunk_b64 {
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) {
Ok(decoded) => yield Ok::<Bytes, IoError>(Bytes::from(decoded)),
Err(err) => {
warn!(trace_id = %trace_id_owned, error = %err, "gateway failed to decode executor chunk");
break;
}
}
} else if let Some(text) = text {
yield Ok::<Bytes, IoError>(Bytes::from(text.into_bytes()));
}
}
StreamFramePayload::Telemetry { .. } => {}
StreamFramePayload::Eof { .. } => break,
StreamFramePayload::Error { error } => {
warn!(trace_id = %trace_id_owned, error = %error.message, "executor stream emitted error frame");
break;
}
StreamFramePayload::Headers { .. } => {}
}
}
};
Ok(Some(build_client_response_from_parts(
status_code,
&headers,
Body::from_stream(body_stream),
trace_id,
Some(decision),
)?))
}
async fn collect_error_body<R>(
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Vec<u8>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
let mut body = Vec::new();
while let Some(frame) = read_next_frame(lines).await? {
match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => {
let chunk = if let Some(chunk_b64) = chunk_b64 {
base64::engine::general_purpose::STANDARD
.decode(chunk_b64)
.map_err(|err| GatewayError::Internal(err.to_string()))?
} else {
text.unwrap_or_default().into_bytes()
};
body.extend_from_slice(&chunk);
if body.len() >= MAX_ERROR_BODY_BYTES {
body.truncate(MAX_ERROR_BODY_BYTES);
break;
}
}
StreamFramePayload::Telemetry { .. } => {}
StreamFramePayload::Eof { .. } => break,
StreamFramePayload::Error { error } => {
warn!(error = %error.message, "executor stream emitted error frame while collecting error body");
break;
}
StreamFramePayload::Headers { .. } => {}
}
}
Ok(body)
}
async fn read_next_frame<R>(
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Option<StreamFrame>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
while let Some(line) = lines.next().await {
let line = line.map_err(|err| GatewayError::Internal(err.to_string()))?;
if line.trim().is_empty() {
continue;
}
let frame: StreamFrame =
serde_json::from_str(&line).map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(frame));
}
Ok(None)
}
fn build_executor_error_response(
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
status_code: u16,
headers: BTreeMap<String, String>,
error_body: Vec<u8>,
) -> Result<Response<Body>, GatewayError> {
let content_type = headers
.get("content-type")
.map(|value| value.to_ascii_lowercase())
.unwrap_or_default();
if plan_kind == GEMINI_FILES_DOWNLOAD_PLAN_KIND && !content_type.starts_with("application/json")
{
let wrapped = serde_json::to_vec(&json!({
"error": String::from_utf8_lossy(&error_body).to_string(),
}))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let wrapped_headers =
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
return build_client_response_from_parts(
status_code,
&wrapped_headers,
Body::from(wrapped),
trace_id,
Some(decision),
);
}
if plan_kind == OPENAI_VIDEO_CONTENT_PLAN_KIND && !content_type.starts_with("application/json")
{
let wrapped = serde_json::to_vec(&json!({
"error": {
"type": "upstream_error",
"message": "Video not available",
}
}))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let wrapped_headers =
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
return build_client_response_from_parts(
status_code,
&wrapped_headers,
Body::from(wrapped),
trace_id,
Some(decision),
);
}
build_client_response_from_parts(
status_code,
&headers,
Body::from(error_body),
trace_id,
Some(decision),
)
}

View File

@@ -0,0 +1,149 @@
#[path = "constants.rs"]
mod constants;
#[path = "control.rs"]
mod control;
#[path = "error.rs"]
mod error;
#[path = "executor.rs"]
mod executor;
#[path = "handlers.rs"]
mod handlers;
#[path = "headers.rs"]
mod headers;
#[path = "response.rs"]
mod response;
use axum::http::header::{HeaderName, HeaderValue};
use axum::routing::{any, get};
use axum::Router;
pub(crate) use control::{
maybe_execute_via_control, resolve_control_route, GatewayControlAuthContext,
GatewayControlDecision,
};
pub(crate) use error::GatewayError;
pub(crate) use executor::maybe_execute_via_executor_stream;
use handlers::{health, proxy_request};
pub(crate) use response::{build_client_response, build_client_response_from_parts};
#[derive(Debug, Clone)]
pub struct AppState {
upstream_base_url: String,
control_base_url: Option<String>,
executor_base_url: Option<String>,
client: reqwest::Client,
}
impl AppState {
pub fn new(
upstream_base_url: impl Into<String>,
control_base_url: Option<String>,
) -> Result<Self, reqwest::Error> {
Self::new_with_executor(upstream_base_url, control_base_url, None)
}
pub fn new_with_executor(
upstream_base_url: impl Into<String>,
control_base_url: Option<String>,
executor_base_url: Option<String>,
) -> Result<Self, reqwest::Error> {
let client = reqwest::Client::builder()
.http2_adaptive_window(true)
.connect_timeout(std::time::Duration::from_secs(10))
.build()?;
Ok(Self {
upstream_base_url: normalize_upstream_base_url(upstream_base_url.into()),
control_base_url: control_base_url
.map(normalize_upstream_base_url)
.filter(|value| !value.is_empty()),
executor_base_url: executor_base_url
.map(normalize_upstream_base_url)
.filter(|value| !value.is_empty()),
client,
})
}
}
pub fn build_router(upstream_base_url: impl Into<String>) -> Result<Router, reqwest::Error> {
build_router_with_control(upstream_base_url, None)
}
pub fn build_router_with_control(
upstream_base_url: impl Into<String>,
control_base_url: Option<String>,
) -> Result<Router, reqwest::Error> {
Ok(build_router_with_state(AppState::new(
upstream_base_url,
control_base_url,
)?))
}
pub fn build_router_with_endpoints(
upstream_base_url: impl Into<String>,
control_base_url: Option<String>,
executor_base_url: Option<String>,
) -> Result<Router, reqwest::Error> {
Ok(build_router_with_state(AppState::new_with_executor(
upstream_base_url,
control_base_url,
executor_base_url,
)?))
}
pub fn build_router_with_state(state: AppState) -> Router {
Router::new()
.route("/_gateway/health", get(health))
.route("/", any(proxy_request))
.route("/{*path}", any(proxy_request))
.with_state(state)
}
pub async fn serve_tcp(
bind: &str,
upstream_base_url: &str,
control_base_url: Option<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
serve_tcp_with_endpoints(bind, upstream_base_url, control_base_url, None).await
}
pub async fn serve_tcp_with_endpoints(
bind: &str,
upstream_base_url: &str,
control_base_url: Option<&str>,
executor_base_url: Option<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
let listener = tokio::net::TcpListener::bind(bind).await?;
let router = build_router_with_endpoints(
upstream_base_url.to_string(),
control_base_url.map(ToOwned::to_owned),
executor_base_url.map(ToOwned::to_owned),
)?;
axum::serve(
listener,
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.await?;
Ok(())
}
fn normalize_upstream_base_url(upstream_base_url: String) -> String {
upstream_base_url.trim_end_matches('/').to_string()
}
fn insert_header_if_missing(
headers: &mut http::HeaderMap,
key: &'static str,
value: &str,
) -> Result<(), GatewayError> {
if headers.contains_key(key) {
return Ok(());
}
let name = HeaderName::from_static(key);
let value =
HeaderValue::from_str(value).map_err(|err| GatewayError::Internal(err.to_string()))?;
headers.insert(name, value);
Ok(())
}
#[cfg(test)]
mod tests;

View File

@@ -0,0 +1,290 @@
use super::*;
#[tokio::test]
async fn gateway_executes_sync_ai_route_via_control_execute_endpoint() {
#[derive(Debug, Clone)]
struct SeenExecuteSyncRequest {
trace_id: String,
path: String,
model: String,
user_id: String,
}
let seen_execute = Arc::new(Mutex::new(None::<SeenExecuteSyncRequest>));
let seen_execute_clone = Arc::clone(&seen_execute);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "chat",
"auth_endpoint_signature": "openai:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-sync-123",
"api_key_id": "key-sync-123",
"balance_remaining": 12.5,
"access_allowed": true
},
"public_path": "/v1/chat/completions"
}))
}),
)
.route(
"/api/internal/gateway/execute-sync",
any(move |request: Request| {
let seen_execute_inner = Arc::clone(&seen_execute_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execute payload should parse");
*seen_execute_inner.lock().expect("mutex should lock") =
Some(SeenExecuteSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
model: payload
.get("body_json")
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("auth_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
let mut response = Response::builder()
.status(StatusCode::CREATED)
.body(Body::from("{\"ok\":true}"))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response.headers_mut().insert(
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
HeaderValue::from_static("true"),
);
response
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("ai_public")
);
assert_eq!(
response
.headers()
.get(GATEWAY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("rust-phase3b")
);
assert_eq!(
response.text().await.expect("body should read"),
"{\"ok\":true}"
);
let seen_execute_request = seen_execute
.lock()
.expect("mutex should lock")
.clone()
.expect("execute-sync should be captured");
assert_eq!(seen_execute_request.trace_id, "trace-sync-123");
assert_eq!(seen_execute_request.path, "/v1/chat/completions");
assert_eq!(seen_execute_request.model, "gpt-5");
assert_eq!(seen_execute_request.user_id, "user-sync-123");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_stream_ai_route_via_control_stream_endpoint() {
#[derive(Debug, Clone)]
struct SeenExecuteStreamRequest {
trace_id: String,
path: String,
stream: bool,
}
let seen_execute = Arc::new(Mutex::new(None::<SeenExecuteStreamRequest>));
let seen_execute_clone = Arc::clone(&seen_execute);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "chat",
"auth_endpoint_signature": "openai:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-stream-123",
"api_key_id": "key-stream-123",
"balance_remaining": 8.0,
"access_allowed": true
},
"public_path": "/v1/chat/completions"
}))
}),
)
.route(
"/api/internal/gateway/execute-stream",
any(move |request: Request| {
let seen_execute_inner = Arc::clone(&seen_execute_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execute payload should parse");
*seen_execute_inner.lock().expect("mutex should lock") =
Some(SeenExecuteStreamRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
stream: payload
.get("body_json")
.and_then(|value| value.get("stream"))
.and_then(|value| value.as_bool())
.unwrap_or(false),
});
let stream = futures_util::stream::iter([
Ok::<_, Infallible>(Bytes::from_static(b"data: one\n\n")),
Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n")),
]);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from_stream(stream))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream"),
);
response.headers_mut().insert(
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
HeaderValue::from_static("true"),
);
response
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-stream-123")
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("ai_public")
);
assert_eq!(
response
.headers()
.get(GATEWAY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("rust-phase3b")
);
assert_eq!(
response.text().await.expect("body should read"),
"data: one\n\ndata: [DONE]\n\n"
);
let seen_execute_request = seen_execute
.lock()
.expect("mutex should lock")
.clone()
.expect("execute-stream should be captured");
assert_eq!(seen_execute_request.trace_id, "trace-stream-123");
assert_eq!(seen_execute_request.path, "/v1/chat/completions");
assert!(seen_execute_request.stream);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,222 @@
use super::*;
#[tokio::test]
async fn gateway_consults_control_api_for_ai_routes_and_propagates_decision_headers() {
#[derive(Debug, Clone)]
struct SeenControlRequest {
method: String,
path: String,
query_string: String,
trace_id: String,
}
#[derive(Debug, Clone)]
struct SeenPublicRequest {
control_route_class: String,
control_route_family: String,
control_route_kind: String,
control_executor_candidate: String,
control_endpoint_signature: String,
trusted_user_id: String,
trusted_api_key_id: String,
trusted_balance_remaining: String,
trusted_access_allowed: String,
trace_id: String,
}
let seen_control = Arc::new(Mutex::new(None::<SeenControlRequest>));
let seen_control_clone = Arc::clone(&seen_control);
let seen_public = Arc::new(Mutex::new(None::<SeenPublicRequest>));
let seen_public_clone = Arc::clone(&seen_public);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(move |request: Request| {
let seen_control_inner = Arc::clone(&seen_control_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("control payload should parse");
*seen_control_inner.lock().expect("mutex should lock") =
Some(SeenControlRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
query_string: payload
.get("query_string")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "chat",
"auth_endpoint_signature": "openai:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-123",
"api_key_id": "key-123",
"balance_remaining": 42.5,
"access_allowed": true
},
"public_path": "/v1/chat/completions",
"public_query_string": "stream=true"
}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |request: Request| {
let seen_public_inner = Arc::clone(&seen_public_clone);
async move {
*seen_public_inner.lock().expect("mutex should lock") =
Some(SeenPublicRequest {
control_route_class: request
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
control_route_family: request
.headers()
.get(CONTROL_ROUTE_FAMILY_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
control_route_kind: request
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
control_executor_candidate: request
.headers()
.get(CONTROL_EXECUTOR_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
control_endpoint_signature: request
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
trusted_user_id: request
.headers()
.get(TRUSTED_AUTH_USER_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
trusted_api_key_id: request
.headers()
.get(TRUSTED_AUTH_API_KEY_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
trusted_balance_remaining: request
.headers()
.get(TRUSTED_AUTH_BALANCE_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
trusted_access_allowed: request
.headers()
.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
trace_id: request
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
(
StatusCode::OK,
[(GATEWAY_HEADER, "python-upstream")],
Body::from("proxied"),
)
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions?stream=true"))
.header(TRACE_ID_HEADER, "trace-control-123")
.body("{\"hello\":\"world\"}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("ai_public")
);
assert_eq!(
response
.headers()
.get(CONTROL_EXECUTOR_HEADER)
.and_then(|value| value.to_str().ok()),
Some("true")
);
let seen_control_request = seen_control
.lock()
.expect("mutex should lock")
.clone()
.expect("control request should be captured");
assert_eq!(seen_control_request.method, "POST");
assert_eq!(seen_control_request.path, "/v1/chat/completions");
assert_eq!(seen_control_request.query_string, "stream=true");
assert_eq!(seen_control_request.trace_id, "trace-control-123");
let seen_public_request = seen_public
.lock()
.expect("mutex should lock")
.clone()
.expect("public request should be captured");
assert_eq!(seen_public_request.control_route_class, "ai_public");
assert_eq!(seen_public_request.control_route_family, "openai");
assert_eq!(seen_public_request.control_route_kind, "chat");
assert_eq!(seen_public_request.control_executor_candidate, "true");
assert_eq!(
seen_public_request.control_endpoint_signature,
"openai:chat"
);
assert_eq!(seen_public_request.trusted_user_id, "user-123");
assert_eq!(seen_public_request.trusted_api_key_id, "key-123");
assert_eq!(seen_public_request.trusted_balance_remaining, "42.5");
assert_eq!(seen_public_request.trusted_access_allowed, "true");
assert_eq!(seen_public_request.trace_id, "trace-control-123");
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,537 @@
use aether_contracts::{StreamFrame, StreamFramePayload, StreamFrameType};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use super::*;
#[tokio::test]
async fn gateway_executes_gemini_files_download_via_control_sync_endpoint() {
#[derive(Debug, Clone)]
struct SeenExecuteFilesRequest {
trace_id: String,
method: String,
path: String,
body_base64: Option<String>,
}
let seen_execute = Arc::new(Mutex::new(None::<SeenExecuteFilesRequest>));
let seen_execute_clone = Arc::clone(&seen_execute);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-123",
"api_key_id": "key-files-123",
"access_allowed": true
},
"public_path": "/v1beta/files/file-123:download"
}))
}),
)
.route(
"/api/internal/gateway/execute-sync",
any(move |request: Request| {
let seen_execute_inner = Arc::clone(&seen_execute_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execute payload should parse");
*seen_execute_inner.lock().expect("mutex should lock") =
Some(SeenExecuteFilesRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_base64: payload
.get("body_base64")
.and_then(|value| value.as_str())
.map(ToOwned::to_owned),
});
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from("file-bytes"))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/octet-stream"),
);
response.headers_mut().insert(
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
HeaderValue::from_static("true"),
);
response
}
}),
)
.route(
"/v1beta/files/file-123:download",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/files/file-123:download?alt=media"
))
.header(TRACE_ID_HEADER, "trace-files-download-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
"file-bytes"
);
let seen_execute_request = seen_execute
.lock()
.expect("mutex should lock")
.clone()
.expect("execute-sync should be captured");
assert_eq!(seen_execute_request.trace_id, "trace-files-download-123");
assert_eq!(seen_execute_request.method, "GET");
assert_eq!(seen_execute_request.path, "/v1beta/files/file-123:download");
assert!(seen_execute_request.body_base64.is_none());
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_files_upload_via_control_sync_endpoint() {
#[derive(Debug, Clone)]
struct SeenExecuteFilesUploadRequest {
path: String,
body_base64: String,
content_type: String,
}
let seen_execute = Arc::new(Mutex::new(None::<SeenExecuteFilesUploadRequest>));
let seen_execute_clone = Arc::clone(&seen_execute);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-456",
"api_key_id": "key-files-456",
"access_allowed": true
},
"public_path": "/upload/v1beta/files"
}))
}),
)
.route(
"/api/internal/gateway/execute-sync",
any(move |request: Request| {
let seen_execute_inner = Arc::clone(&seen_execute_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execute payload should parse");
*seen_execute_inner.lock().expect("mutex should lock") =
Some(SeenExecuteFilesUploadRequest {
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_base64: payload
.get("body_base64")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
content_type: payload
.get("headers")
.and_then(|value| value.get("content-type"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
let mut response = Response::builder()
.status(StatusCode::CREATED)
.body(Body::from("{\"uploaded\":true}"))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response.headers_mut().insert(
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
HeaderValue::from_static("true"),
);
response
}
}),
)
.route(
"/upload/v1beta/files",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/upload/v1beta/files?uploadType=resumable"
))
.header(http::header::CONTENT_TYPE, "application/octet-stream")
.body("upload-body-bytes")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response.text().await.expect("body should read"),
"{\"uploaded\":true}"
);
let seen_execute_request = seen_execute
.lock()
.expect("mutex should lock")
.clone()
.expect("execute-sync should be captured");
assert_eq!(seen_execute_request.path, "/upload/v1beta/files");
assert_eq!(
BASE64_STANDARD
.decode(seen_execute_request.body_base64)
.expect("body should decode"),
b"upload-body-bytes"
);
assert_eq!(
seen_execute_request.content_type,
"application/octet-stream"
);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_files_download_via_executor_stream_plan() {
#[derive(Debug, Clone)]
struct SeenPlanStreamRequest {
trace_id: String,
method: String,
path: String,
query_string: String,
user_id: String,
}
#[derive(Debug, Clone)]
struct SeenExecutorStreamRequest {
trace_id: String,
method: String,
url: String,
stream: bool,
client_api_format: String,
}
let seen_plan = Arc::new(Mutex::new(None::<SeenPlanStreamRequest>));
let seen_plan_clone = Arc::clone(&seen_plan);
let seen_executor = Arc::new(Mutex::new(None::<SeenExecutorStreamRequest>));
let seen_executor_clone = Arc::clone(&seen_executor);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-direct-123",
"api_key_id": "key-files-direct-123",
"access_allowed": true
},
"public_path": "/v1beta/files/file-123:download"
}))
}),
)
.route(
"/api/internal/gateway/plan-stream",
any(move |request: Request| {
let seen_plan_inner = Arc::clone(&seen_plan_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("plan payload should parse");
*seen_plan_inner.lock().expect("mutex should lock") =
Some(SeenPlanStreamRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
query_string: payload
.get("query_string")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("auth_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"action": "executor_stream",
"plan_kind": "gemini_files_download",
"plan": {
"request_id": "req-files-direct-123",
"provider_id": "provider-files-direct-123",
"endpoint_id": "endpoint-files-direct-123",
"key_id": "key-files-direct-123",
"provider_name": "gemini",
"method": "GET",
"url": "https://files.example/v1beta/files/file-123:download?alt=media",
"headers": {
"authorization": "Bearer upstream-key"
},
"body": {},
"stream": true,
"client_api_format": "gemini:files",
"provider_api_format": "gemini:files",
"model_name": "gemini-files"
}
}))
}
}),
)
.route(
"/v1beta/files/file-123:download",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let executor = Router::new().route(
"/v1/execute/stream",
any(move |request: Request| {
let seen_executor_inner = Arc::clone(&seen_executor_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("executor payload should parse");
*seen_executor_inner.lock().expect("mutex should lock") =
Some(SeenExecutorStreamRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
stream: payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false),
client_api_format: payload
.get("client_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
let frames = [
StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: std::collections::BTreeMap::from([(
"content-type".to_string(),
"application/octet-stream".to_string(),
)]),
},
},
StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(BASE64_STANDARD.encode(b"file-direct-")),
text: None,
},
},
StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(BASE64_STANDARD.encode(b"bytes")),
text: None,
},
},
StreamFrame::eof(),
];
let body = frames
.into_iter()
.map(|frame| serde_json::to_string(&frame).expect("frame should serialize"))
.collect::<Vec<_>>()
.join("\n")
+ "\n";
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(body))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (executor_url, executor_handle) = start_server(executor).await;
let gateway =
build_router_with_endpoints(upstream_url.clone(), Some(upstream_url), Some(executor_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/files/file-123:download?alt=media"
))
.header(TRACE_ID_HEADER, "trace-files-direct-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/octet-stream")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("ai_public")
);
assert_eq!(
response.bytes().await.expect("body should read"),
Bytes::from_static(b"file-direct-bytes")
);
let seen_plan_request = seen_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("plan-stream should be captured");
assert_eq!(seen_plan_request.trace_id, "trace-files-direct-123");
assert_eq!(seen_plan_request.method, "GET");
assert_eq!(seen_plan_request.path, "/v1beta/files/file-123:download");
assert_eq!(seen_plan_request.query_string, "alt=media");
assert_eq!(seen_plan_request.user_id, "user-files-direct-123");
let seen_executor_request = seen_executor
.lock()
.expect("mutex should lock")
.clone()
.expect("executor stream should be captured");
assert_eq!(seen_executor_request.trace_id, "trace-files-direct-123");
assert_eq!(seen_executor_request.method, "GET");
assert_eq!(
seen_executor_request.url,
"https://files.example/v1beta/files/file-123:download?alt=media"
);
assert!(seen_executor_request.stream);
assert_eq!(seen_executor_request.client_api_format, "gemini:files");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
executor_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,35 @@
pub(super) use std::convert::Infallible;
pub(super) use std::sync::{Arc, Mutex};
pub(super) use axum::body::{to_bytes, Body, Bytes};
pub(super) use axum::response::Response;
pub(super) use axum::routing::any;
pub(super) use axum::{extract::Request, Json, Router};
pub(super) use http::header::{HeaderName, HeaderValue};
pub(super) use http::StatusCode;
pub(super) use serde_json::json;
mod ai_execute;
mod control;
mod files;
mod proxy;
mod video;
pub(super) use super::constants::*;
pub(super) use super::{build_router, build_router_with_control, build_router_with_endpoints};
pub(super) async fn start_server(app: Router) -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let handle = tokio::spawn(async move {
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.await
.expect("server should run");
});
(format!("http://{addr}"), handle)
}

View File

@@ -0,0 +1,163 @@
use super::*;
#[tokio::test]
async fn gateway_proxies_method_path_body_and_generates_trace_id() {
#[derive(Debug, Clone)]
struct SeenRequest {
method: String,
path: String,
trace_id: String,
host: String,
forwarded_for: String,
body: String,
}
let seen = Arc::new(Mutex::new(None::<SeenRequest>));
let seen_clone = Arc::clone(&seen);
let upstream = Router::new()
.route("/", any(|| async { StatusCode::OK }))
.route(
"/{*path}",
any(move |request: Request| {
let seen_inner = Arc::clone(&seen_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
*seen_inner.lock().expect("mutex should lock") = Some(SeenRequest {
method: parts.method.to_string(),
path: parts
.uri
.path_and_query()
.map(|value| value.as_str())
.unwrap_or("/")
.to_string(),
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
host: parts
.headers
.get(http::header::HOST)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
forwarded_for: parts
.headers
.get(FORWARDED_FOR_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"),
});
(
StatusCode::CREATED,
[(GATEWAY_HEADER, "python-upstream")],
Body::from("proxied"),
)
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router(upstream_url).expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let response = client
.post(format!("{gateway_url}/v1/chat/completions?stream=true"))
.header(http::header::HOST, "api.example.com")
.body("{\"hello\":\"world\"}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response
.headers()
.get(GATEWAY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("python-upstream")
);
let response_trace_id = response
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.expect("response trace id should exist")
.to_string();
assert_eq!(response.text().await.expect("body should read"), "proxied");
let seen_request = seen
.lock()
.expect("mutex should lock")
.clone()
.expect("upstream request should be captured");
assert_eq!(seen_request.method, "POST");
assert_eq!(seen_request.path, "/v1/chat/completions?stream=true");
assert_eq!(seen_request.body, "{\"hello\":\"world\"}");
assert_eq!(seen_request.host, "api.example.com");
assert_eq!(seen_request.forwarded_for, "127.0.0.1");
assert_eq!(seen_request.trace_id, response_trace_id);
assert!(!seen_request.trace_id.is_empty());
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_preserves_existing_trace_id_and_streams_response() {
let upstream = Router::new().route(
"/{*path}",
any(|request: Request| async move {
let incoming_trace_id = request
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let stream = futures_util::stream::iter([
Ok::<_, Infallible>(Bytes::from_static(b"chunk-1")),
Ok::<_, Infallible>(Bytes::from_static(b"chunk-2")),
]);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from_stream(stream))
.expect("response should build");
response.headers_mut().insert(
HeaderName::from_static(TRACE_ID_HEADER),
HeaderValue::from_str(&incoming_trace_id).expect("trace id header"),
);
response
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router(upstream_url).expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/messages"))
.header(TRACE_ID_HEADER, "trace-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok()),
Some("trace-123")
);
assert_eq!(
response.bytes().await.expect("bytes should read"),
Bytes::from_static(b"chunk-1chunk-2")
);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,357 @@
use aether_contracts::{StreamFrame, StreamFramePayload, StreamFrameType};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use super::*;
#[tokio::test]
async fn gateway_executes_video_get_route_via_control_sync_endpoint() {
#[derive(Debug, Clone)]
struct SeenExecuteVideoRequest {
method: String,
path: String,
body_json: serde_json::Value,
}
let seen_execute = Arc::new(Mutex::new(None::<SeenExecuteVideoRequest>));
let seen_execute_clone = Arc::clone(&seen_execute);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "video",
"auth_endpoint_signature": "openai:video",
"executor_candidate": true,
"auth_context": {
"user_id": "user-video-123",
"api_key_id": "key-video-123",
"access_allowed": true
},
"public_path": "/v1/videos/task-123"
}))
}),
)
.route(
"/api/internal/gateway/execute-sync",
any(move |request: Request| {
let seen_execute_inner = Arc::clone(&seen_execute_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execute payload should parse");
*seen_execute_inner.lock().expect("mutex should lock") =
Some(SeenExecuteVideoRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_json: payload
.get("body_json")
.cloned()
.unwrap_or_else(|| json!({})),
});
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from("{\"status\":\"queued\"}"))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response.headers_mut().insert(
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
HeaderValue::from_static("true"),
);
response
}
}),
)
.route(
"/v1/videos/task-123",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_control(upstream_url.clone(), Some(upstream_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/videos/task-123"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
"{\"status\":\"queued\"}"
);
let seen_execute_request = seen_execute
.lock()
.expect("mutex should lock")
.clone()
.expect("execute-sync should be captured");
assert_eq!(seen_execute_request.method, "GET");
assert_eq!(seen_execute_request.path, "/v1/videos/task-123");
assert_eq!(seen_execute_request.body_json, json!({}));
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_openai_video_content_via_executor_stream_plan() {
#[derive(Debug, Clone)]
struct SeenPlanStreamRequest {
trace_id: String,
path: String,
user_id: String,
}
#[derive(Debug, Clone)]
struct SeenExecutorStreamRequest {
method: String,
url: String,
provider_api_format: String,
}
let seen_plan = Arc::new(Mutex::new(None::<SeenPlanStreamRequest>));
let seen_plan_clone = Arc::clone(&seen_plan);
let seen_executor = Arc::new(Mutex::new(None::<SeenExecutorStreamRequest>));
let seen_executor_clone = Arc::clone(&seen_executor);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "video",
"auth_endpoint_signature": "openai:video",
"executor_candidate": true,
"auth_context": {
"user_id": "user-video-direct-123",
"api_key_id": "key-video-direct-123",
"access_allowed": true
},
"public_path": "/v1/videos/task-123/content"
}))
}),
)
.route(
"/api/internal/gateway/plan-stream",
any(move |request: Request| {
let seen_plan_inner = Arc::clone(&seen_plan_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("plan payload should parse");
*seen_plan_inner.lock().expect("mutex should lock") =
Some(SeenPlanStreamRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("auth_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"action": "executor_stream",
"plan_kind": "openai_video_content",
"plan": {
"request_id": "req-video-direct-123",
"provider_id": "provider-video-direct-123",
"endpoint_id": "endpoint-video-direct-123",
"key_id": "key-video-direct-123",
"provider_name": "openai",
"method": "GET",
"url": "https://cdn.example.com/video.mp4",
"headers": {},
"body": {},
"stream": true,
"client_api_format": "openai:video",
"provider_api_format": "openai:video",
"model_name": "sora-2"
}
}))
}
}),
)
.route(
"/v1/videos/task-123/content",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let executor = Router::new().route(
"/v1/execute/stream",
any(move |request: Request| {
let seen_executor_inner = Arc::clone(&seen_executor_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("executor payload should parse");
*seen_executor_inner.lock().expect("mutex should lock") =
Some(SeenExecutorStreamRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
provider_api_format: payload
.get("provider_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
let frames = [
StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: std::collections::BTreeMap::from([(
"content-type".to_string(),
"video/mp4".to_string(),
)]),
},
},
StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(BASE64_STANDARD.encode(b"openai-")),
text: None,
},
},
StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: Some(BASE64_STANDARD.encode(b"video")),
text: None,
},
},
StreamFrame::eof(),
];
let body = frames
.into_iter()
.map(|frame| serde_json::to_string(&frame).expect("frame should serialize"))
.collect::<Vec<_>>()
.join("\n")
+ "\n";
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(body))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (executor_url, executor_handle) = start_server(executor).await;
let gateway =
build_router_with_endpoints(upstream_url.clone(), Some(upstream_url), Some(executor_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1/videos/task-123/content?variant=video"
))
.header(TRACE_ID_HEADER, "trace-video-direct-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("video/mp4")
);
assert_eq!(
response.bytes().await.expect("body should read"),
Bytes::from_static(b"openai-video")
);
let seen_plan_request = seen_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("plan-stream should be captured");
assert_eq!(seen_plan_request.trace_id, "trace-video-direct-123");
assert_eq!(seen_plan_request.path, "/v1/videos/task-123/content");
assert_eq!(seen_plan_request.user_id, "user-video-direct-123");
let seen_executor_request = seen_executor
.lock()
.expect("mutex should lock")
.clone()
.expect("executor stream should be captured");
assert_eq!(seen_executor_request.method, "GET");
assert_eq!(
seen_executor_request.url,
"https://cdn.example.com/video.mp4"
);
assert_eq!(seen_executor_request.provider_api_format, "openai:video");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
executor_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,193 @@
use std::time::Instant;
use axum::body::{to_bytes, Body};
use axum::extract::{ConnectInfo, Request, State};
use axum::http::Response;
use axum::response::IntoResponse;
use axum::Json;
use futures_util::TryStreamExt;
use serde_json::json;
use tracing::info;
use crate::gateway::constants::*;
use crate::gateway::headers::{
extract_or_generate_trace_id, header_value_str, should_skip_request_header,
};
use crate::gateway::{
build_client_response, maybe_execute_via_control, maybe_execute_via_executor_stream,
resolve_control_route, AppState, GatewayError,
};
pub(crate) async fn health(State(state): State<AppState>) -> impl IntoResponse {
Json(json!({
"status": "ok",
"component": "aether-gateway",
"control_api_enabled": state.control_base_url.is_some(),
}))
}
pub(crate) async fn proxy_request(
State(state): State<AppState>,
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,
request: Request,
) -> Result<Response<Body>, GatewayError> {
let started_at = Instant::now();
let (parts, body) = request.into_parts();
let method = parts.method.clone();
let path_and_query = parts
.uri
.path_and_query()
.map(|value| value.as_str())
.unwrap_or("/");
let host_header = header_value_str(&parts.headers, http::header::HOST.as_str());
let trace_id = extract_or_generate_trace_id(&parts.headers);
let control_decision =
resolve_control_route(&state, &method, &parts.uri, &parts.headers, &trace_id).await?;
let upstream_path_and_query = control_decision
.as_ref()
.map(|decision| decision.proxy_path_and_query())
.unwrap_or_else(|| path_and_query.to_string());
let target_url = format!("{}{}", state.upstream_base_url, upstream_path_and_query);
let should_try_control_execute = control_decision
.as_ref()
.map(|decision| {
decision.executor_candidate && decision.route_class.as_deref() == Some("ai_public")
})
.unwrap_or(false);
let mut upstream_request = state.client.request(method.clone(), &target_url);
for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) {
continue;
}
upstream_request = upstream_request.header(name, value);
}
if let Some(host) = host_header.as_deref() {
if !parts.headers.contains_key(FORWARDED_HOST_HEADER) {
upstream_request = upstream_request.header(FORWARDED_HOST_HEADER, host);
}
}
if !parts.headers.contains_key(FORWARDED_FOR_HEADER) {
upstream_request =
upstream_request.header(FORWARDED_FOR_HEADER, remote_addr.ip().to_string());
}
if !parts.headers.contains_key(FORWARDED_PROTO_HEADER) {
upstream_request = upstream_request.header(FORWARDED_PROTO_HEADER, "http");
}
if !parts.headers.contains_key(TRACE_ID_HEADER) {
upstream_request = upstream_request.header(TRACE_ID_HEADER, &trace_id);
}
if let Some(decision) = control_decision.as_ref() {
upstream_request = upstream_request
.header(
CONTROL_ROUTE_CLASS_HEADER,
decision.route_class.as_deref().unwrap_or("passthrough"),
)
.header(
CONTROL_EXECUTOR_HEADER,
if decision.executor_candidate {
"true"
} else {
"false"
},
);
if let Some(route_family) = decision.route_family.as_deref() {
upstream_request = upstream_request.header(CONTROL_ROUTE_FAMILY_HEADER, route_family);
}
if let Some(route_kind) = decision.route_kind.as_deref() {
upstream_request = upstream_request.header(CONTROL_ROUTE_KIND_HEADER, route_kind);
}
if let Some(endpoint_signature) = decision.auth_endpoint_signature.as_deref() {
upstream_request =
upstream_request.header(CONTROL_ENDPOINT_SIGNATURE_HEADER, endpoint_signature);
}
if let Some(auth_context) = decision.auth_context.as_ref() {
upstream_request = upstream_request
.header(TRUSTED_AUTH_USER_ID_HEADER, &auth_context.user_id)
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, &auth_context.api_key_id)
.header(
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER,
if auth_context.access_allowed {
"true"
} else {
"false"
},
);
if let Some(balance_remaining) = auth_context.balance_remaining {
upstream_request = upstream_request
.header(TRUSTED_AUTH_BALANCE_HEADER, balance_remaining.to_string());
}
}
}
upstream_request = upstream_request.header(GATEWAY_HEADER, "rust-phase3b");
let upstream_response = if should_try_control_execute {
let buffered_body = to_bytes(body, usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Some(executor_response) =
maybe_execute_via_executor_stream(&state, &parts, &trace_id, control_decision.as_ref())
.await?
{
return Ok(executor_response);
}
if let Some(control_response) = maybe_execute_via_control(
&state,
&parts,
buffered_body.clone(),
&trace_id,
control_decision.as_ref(),
)
.await?
{
return Ok(control_response);
}
upstream_request
.body(buffered_body)
.send()
.await
.map_err(|err| GatewayError::UpstreamUnavailable {
trace_id: trace_id.clone(),
message: err.to_string(),
})?
} else {
let request_body_stream = body
.into_data_stream()
.map_err(|err| std::io::Error::other(err.to_string()));
upstream_request
.body(reqwest::Body::wrap_stream(request_body_stream))
.send()
.await
.map_err(|err| GatewayError::UpstreamUnavailable {
trace_id: trace_id.clone(),
message: err.to_string(),
})?
};
let response = build_client_response(upstream_response, &trace_id, control_decision.as_ref())?;
let response_status = response.status();
let elapsed_ms = started_at.elapsed().as_millis() as u64;
info!(
trace_id = %trace_id,
remote_addr = %remote_addr,
method = %method,
path = %path_and_query,
route_class = control_decision
.as_ref()
.and_then(|decision| decision.route_class.as_deref())
.unwrap_or("passthrough"),
status = response_status.as_u16(),
elapsed_ms,
"gateway proxied request"
);
Ok(response)
}

View File

@@ -0,0 +1,83 @@
use std::collections::BTreeMap;
use crate::gateway::constants::*;
use uuid::Uuid;
pub(crate) fn extract_or_generate_trace_id(headers: &http::HeaderMap) -> String {
header_value_str(headers, TRACE_ID_HEADER).unwrap_or_else(|| Uuid::new_v4().to_string())
}
pub(crate) fn header_value_str(headers: &http::HeaderMap, key: &str) -> Option<String> {
headers
.get(key)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub(crate) fn header_value_u64(headers: &http::HeaderMap, key: &str) -> Option<u64> {
header_value_str(headers, key).and_then(|value| value.parse::<u64>().ok())
}
pub(crate) fn should_skip_request_header(name: &str) -> bool {
matches!(
name.to_ascii_lowercase().as_str(),
"connection"
| "keep-alive"
| "proxy-authenticate"
| "proxy-authorization"
| "proxy-connection"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
)
}
pub(crate) fn should_skip_response_header(name: &str) -> bool {
matches!(
name.to_ascii_lowercase().as_str(),
"connection"
| "keep-alive"
| "proxy-authenticate"
| "proxy-authorization"
| "proxy-connection"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
| "x-aether-control-executed"
| "x-aether-control-action"
)
}
pub(crate) fn collect_control_headers(headers: &http::HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_ascii_lowercase(), value.trim().to_string()))
})
.collect()
}
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
header_value_str(headers, http::header::CONTENT_TYPE.as_str())
.map(|value| value.to_ascii_lowercase().contains("application/json"))
.unwrap_or(false)
}
pub(crate) fn header_equals(
headers: &reqwest::header::HeaderMap,
key: &'static str,
expected: &str,
) -> bool {
headers
.get(key)
.and_then(|value| value.to_str().ok())
.map(|value| value.eq_ignore_ascii_case(expected))
.unwrap_or(false)
}

View File

@@ -0,0 +1,6 @@
mod gateway;
pub use gateway::{
build_router, build_router_with_control, build_router_with_endpoints, build_router_with_state,
serve_tcp, serve_tcp_with_endpoints, AppState,
};

View File

@@ -0,0 +1,62 @@
use clap::Parser;
use tracing::info;
use aether_gateway::{serve_tcp, serve_tcp_with_endpoints};
#[derive(Parser, Debug)]
#[command(
name = "aether-gateway",
about = "Phase 3a Rust ingress gateway for Aether"
)]
struct Args {
#[arg(long, env = "AETHER_GATEWAY_BIND", default_value = "0.0.0.0:8084")]
bind: String,
#[arg(
long,
env = "AETHER_GATEWAY_UPSTREAM",
default_value = "http://127.0.0.1:18084"
)]
upstream: String,
#[arg(long, env = "AETHER_GATEWAY_CONTROL_URL")]
control_url: Option<String>,
#[arg(long, env = "AETHER_GATEWAY_EXECUTOR_URL")]
executor_url: Option<String>,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "aether_gateway=info".into()),
)
.init();
let args = Args::parse();
let control_url = args
.control_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let executor_url = args
.executor_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
info!(
bind = %args.bind,
upstream = %args.upstream,
control_url = control_url.unwrap_or("-"),
executor_url = executor_url.unwrap_or("-"),
"aether-gateway started"
);
if executor_url.is_some() {
serve_tcp_with_endpoints(&args.bind, &args.upstream, control_url, executor_url).await?;
} else {
serve_tcp(&args.bind, &args.upstream, control_url).await?;
}
Ok(())
}

View File

@@ -0,0 +1,92 @@
use std::collections::BTreeMap;
use axum::body::Body;
use axum::http::header::{HeaderName, HeaderValue};
use axum::http::Response;
use crate::gateway::constants::*;
use crate::gateway::headers::should_skip_response_header;
use crate::gateway::{insert_header_if_missing, GatewayControlDecision, GatewayError};
pub(crate) fn build_client_response(
upstream_response: reqwest::Response,
trace_id: &str,
control_decision: Option<&GatewayControlDecision>,
) -> Result<Response<Body>, GatewayError> {
let status = upstream_response.status();
let upstream_headers = upstream_response
.headers()
.iter()
.map(|(name, value)| {
(
name.as_str().to_string(),
value.to_str().unwrap_or_default().to_string(),
)
})
.collect::<BTreeMap<_, _>>();
let upstream_stream = upstream_response.bytes_stream();
build_client_response_from_parts(
status.as_u16(),
&upstream_headers,
Body::from_stream(upstream_stream),
trace_id,
control_decision,
)
}
pub(crate) fn build_client_response_from_parts(
status_code: u16,
upstream_headers: &BTreeMap<String, String>,
body: Body,
trace_id: &str,
control_decision: Option<&GatewayControlDecision>,
) -> Result<Response<Body>, GatewayError> {
let mut response = Response::builder()
.status(status_code)
.body(body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
for (name, value) in upstream_headers {
if should_skip_response_header(name.as_str()) {
continue;
}
let header_name = HeaderName::from_bytes(name.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let header_value =
HeaderValue::from_str(value).map_err(|err| GatewayError::Internal(err.to_string()))?;
response.headers_mut().insert(header_name, header_value);
}
insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, trace_id)?;
insert_header_if_missing(response.headers_mut(), GATEWAY_HEADER, "rust-phase3b")?;
if let Some(decision) = control_decision {
insert_header_if_missing(
response.headers_mut(),
CONTROL_ROUTE_CLASS_HEADER,
decision.route_class.as_deref().unwrap_or("passthrough"),
)?;
insert_header_if_missing(
response.headers_mut(),
CONTROL_EXECUTOR_HEADER,
if decision.executor_candidate {
"true"
} else {
"false"
},
)?;
if let Some(route_family) = decision.route_family.as_deref() {
insert_header_if_missing(
response.headers_mut(),
CONTROL_ROUTE_FAMILY_HEADER,
route_family,
)?;
}
if let Some(route_kind) = decision.route_kind.as_deref() {
insert_header_if_missing(
response.headers_mut(),
CONTROL_ROUTE_KIND_HEADER,
route_kind,
)?;
}
}
Ok(response)
}