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,13 @@
[package]
name = "aether-contracts"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared contracts for Python and Rust Aether components"
[dependencies]
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true

View File

@@ -0,0 +1,41 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionErrorKind {
ConnectTimeout,
FirstByteTimeout,
ReadTimeout,
Upstream4xx,
Upstream5xx,
TlsError,
ProxyError,
ProtocolError,
Cancelled,
Internal,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionPhase {
Connect,
Handshake,
Write,
FirstByte,
StreamRead,
Decode,
Finalize,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExecutionError {
pub kind: ExecutionErrorKind,
pub phase: ExecutionPhase,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub upstream_status: Option<u16>,
#[serde(default)]
pub retryable: bool,
#[serde(default)]
pub failover_recommended: bool,
}

View File

@@ -0,0 +1,58 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{ExecutionError, ExecutionTelemetry};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum StreamFrameType {
Headers,
Data,
Error,
Telemetry,
Eof,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum StreamFramePayload {
Headers {
status_code: u16,
#[serde(default)]
headers: BTreeMap<String, String>,
},
Data {
#[serde(default, skip_serializing_if = "Option::is_none")]
chunk_b64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
text: Option<String>,
},
Error {
error: ExecutionError,
},
Telemetry {
telemetry: ExecutionTelemetry,
},
Eof {
#[serde(default, skip_serializing_if = "Option::is_none")]
summary: Option<Value>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct StreamFrame {
#[serde(rename = "type")]
pub frame_type: StreamFrameType,
pub payload: StreamFramePayload,
}
impl StreamFrame {
pub fn eof() -> Self {
Self {
frame_type: StreamFrameType::Eof,
payload: StreamFramePayload::Eof { summary: None },
}
}
}

View File

@@ -0,0 +1,9 @@
mod error;
mod frame;
mod plan;
mod result;
pub use error::{ExecutionError, ExecutionErrorKind, ExecutionPhase};
pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType};
pub use plan::{ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody};
pub use result::{ExecutionResult, ExecutionTelemetry, ResponseBody};

View File

@@ -0,0 +1,196 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct ExecutionTimeouts {
#[serde(skip_serializing_if = "Option::is_none")]
pub connect_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub read_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub first_byte_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub write_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pool_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total_ms: Option<u64>,
}
impl Default for ExecutionTimeouts {
fn default() -> Self {
Self {
connect_ms: None,
read_ms: None,
first_byte_ms: None,
write_ms: None,
pool_ms: None,
total_ms: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RequestBody {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub json_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body_bytes_b64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body_ref: Option<String>,
}
impl RequestBody {
pub fn from_json(json_body: Value) -> Self {
Self {
json_body: Some(json_body),
body_bytes_b64: None,
body_ref: None,
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct ProxySnapshot {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mode: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub node_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
#[serde(default, alias = "proxy_url", skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub extra: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ExecutionPlan {
pub request_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_name: Option<String>,
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub method: String,
#[serde(alias = "upstream_url")]
pub url: String,
#[serde(default)]
pub headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content_encoding: Option<String>,
pub body: RequestBody,
#[serde(default)]
pub stream: bool,
pub client_api_format: String,
pub provider_api_format: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub proxy: Option<ProxySnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tls_profile: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timeouts: Option<ExecutionTimeouts>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn serializes_plan_with_json_body() {
let plan = ExecutionPlan {
request_id: "req_123".into(),
candidate_id: Some("cand_123".into()),
provider_name: Some("openai".into()),
provider_id: "prov_123".into(),
endpoint_id: "ep_123".into(),
key_id: "key_123".into(),
method: "POST".into(),
url: "https://example.com/v1/chat/completions".into(),
headers: BTreeMap::from([("authorization".into(), "Bearer test".into())]),
content_type: Some("application/json".into()),
content_encoding: Some("gzip".into()),
body: RequestBody::from_json(serde_json::json!({"model":"gpt-test"})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-test".into()),
proxy: None,
tls_profile: Some("chrome".into()),
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(3_600_000),
first_byte_ms: Some(30_000),
..ExecutionTimeouts::default()
}),
};
let raw = serde_json::to_value(&plan).expect("plan should serialize");
assert_eq!(raw["body"]["json_body"]["model"], "gpt-test");
assert_eq!(raw["content_encoding"], "gzip");
assert_eq!(raw["stream"], true);
}
#[test]
fn deserializes_python_control_plane_plan_shape() {
let raw = serde_json::json!({
"request_id": "req-1",
"candidate_id": null,
"provider_name": "openai",
"provider_id": "prov-1",
"endpoint_id": "ep-1",
"key_id": "key-1",
"method": "POST",
"url": "https://example.com/v1/chat/completions",
"headers": {"content-type": "application/json"},
"content_encoding": "gzip",
"body": {"json_body": {"model": "gpt-4.1"}},
"stream": false,
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
"model_name": "gpt-4.1",
"proxy": {
"enabled": true,
"mode": "direct",
"label": "no-proxy",
"url": "http://proxy.internal"
},
"timeouts": {
"connect_ms": 10000,
"read_ms": 30000,
"write_ms": 30000,
"pool_ms": 10000,
"total_ms": 300000
}
});
let plan: ExecutionPlan =
serde_json::from_value(raw).expect("python payload should deserialize");
assert_eq!(plan.url, "https://example.com/v1/chat/completions");
assert_eq!(plan.candidate_id, None);
assert_eq!(plan.provider_name.as_deref(), Some("openai"));
assert_eq!(plan.model_name.as_deref(), Some("gpt-4.1"));
assert_eq!(plan.content_encoding.as_deref(), Some("gzip"));
assert_eq!(
plan.proxy.as_ref().and_then(|proxy| proxy.url.as_deref()),
Some("http://proxy.internal")
);
assert_eq!(
plan.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms),
Some(300_000)
);
}
}

View File

@@ -0,0 +1,40 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::ExecutionError;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExecutionTelemetry {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ttfb_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub elapsed_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub upstream_bytes: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ResponseBody {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub json_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body_bytes_b64: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ExecutionResult {
pub request_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_id: Option<String>,
pub status_code: u16,
#[serde(default)]
pub headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body: Option<ResponseBody>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub telemetry: Option<ExecutionTelemetry>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<ExecutionError>,
}

View File

@@ -0,0 +1,27 @@
[package]
name = "aether-executor"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Rust executor scaffold for Aether request execution"
[dependencies]
aether-contracts.workspace = true
async-stream.workspace = true
axum = { version = "0.8" }
base64.workspace = true
bytes.workspace = true
clap = { version = "4", features = ["derive", "env"] }
flate2.workspace = true
futures-util.workspace = true
http.workspace = true
reqwest.workspace = true
rustls.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
webpki-roots.workspace = true

View File

@@ -0,0 +1,57 @@
use std::path::PathBuf;
use aether_contracts::{ExecutionPlan, ExecutionResult, StreamFrame};
use crate::{ExecutorClientError, TransportMode};
#[derive(Debug, Clone)]
pub struct ExecutorClientConfig {
pub transport: TransportMode,
pub endpoint: Option<String>,
pub socket_path: Option<PathBuf>,
}
impl ExecutorClientConfig {
pub fn unix_socket(socket_path: impl Into<PathBuf>) -> Self {
Self {
transport: TransportMode::UnixSocketHttp,
endpoint: None,
socket_path: Some(socket_path.into()),
}
}
}
#[derive(Debug, Clone)]
pub struct ExecutorClient {
config: ExecutorClientConfig,
}
impl ExecutorClient {
pub fn new(config: ExecutorClientConfig) -> Self {
Self { config }
}
pub fn config(&self) -> &ExecutorClientConfig {
&self.config
}
pub async fn execute(
&self,
_plan: &ExecutionPlan,
) -> Result<ExecutionResult, ExecutorClientError> {
if self.config.endpoint.is_none() && self.config.socket_path.is_none() {
return Err(ExecutorClientError::MissingEndpoint);
}
Err(ExecutorClientError::Unimplemented)
}
pub async fn execute_stream(
&self,
_plan: &ExecutionPlan,
) -> Result<Vec<StreamFrame>, ExecutorClientError> {
if self.config.endpoint.is_none() && self.config.socket_path.is_none() {
return Err(ExecutorClientError::MissingEndpoint);
}
Err(ExecutorClientError::Unimplemented)
}
}

View File

@@ -0,0 +1,48 @@
use http::method::InvalidMethod;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum ExecutorClientError {
#[error("executor endpoint is not configured")]
MissingEndpoint,
#[error("executor request is not implemented yet")]
Unimplemented,
#[error("failed to encode NDJSON frame: {0}")]
Encode(#[from] serde_json::Error),
}
#[derive(Debug, Error)]
pub enum ExecutorServiceError {
#[error("stream execution is not implemented yet")]
StreamUnsupported,
#[error("request body must contain json_body or body_bytes_b64")]
RequestBodyRequired,
#[error("request body base64 is invalid: {0}")]
BodyDecode(base64::DecodeError),
#[error("request content-encoding is not supported: {0}")]
UnsupportedContentEncoding(String),
#[error("proxy execution is not implemented yet")]
ProxyUnsupported,
#[error("tls profile overrides are not implemented yet")]
TlsProfileUnsupported,
#[error("tunnel delegate execution is not implemented yet")]
DelegateUnsupported,
#[error("invalid method: {0}")]
InvalidMethod(#[from] InvalidMethod),
#[error("invalid upstream header name: {0}")]
InvalidHeaderName(String),
#[error("invalid upstream header value for {0}")]
InvalidHeaderValue(String),
#[error("invalid proxy configuration: {0}")]
InvalidProxy(reqwest::Error),
#[error("failed to encode request body: {0}")]
BodyEncode(serde_json::Error),
#[error("failed to build HTTP client: {0}")]
ClientBuild(reqwest::Error),
#[error("failed to execute upstream request: {0}")]
UpstreamRequest(reqwest::Error),
#[error("hub relay request failed: {0}")]
RelayError(String),
#[error("upstream response is not valid JSON: {0}")]
InvalidJson(serde_json::Error),
}

View File

@@ -0,0 +1,12 @@
mod client;
mod error;
mod ndjson;
pub mod server;
mod service;
mod transport;
pub use client::{ExecutorClient, ExecutorClientConfig};
pub use error::{ExecutorClientError, ExecutorServiceError};
pub use ndjson::{decode_frame, encode_frame};
pub use service::{SyncExecutor, UpstreamStreamExecution};
pub use transport::TransportMode;

View File

@@ -0,0 +1,52 @@
use std::path::PathBuf;
use clap::Parser;
use tracing::info;
use aether_executor::server;
#[derive(Parser, Debug)]
#[command(name = "aether-executor", about = "Internal Rust executor for Aether")]
struct Args {
#[arg(long, env = "AETHER_EXECUTOR_TRANSPORT", default_value = "unix_socket")]
transport: String,
#[arg(long, env = "AETHER_EXECUTOR_BIND", default_value = "127.0.0.1:5219")]
bind: String,
#[arg(
long,
env = "AETHER_EXECUTOR_UNIX_SOCKET",
default_value = "/tmp/aether-executor.sock"
)]
unix_socket: PathBuf,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let _ = rustls::crypto::ring::default_provider().install_default();
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "aether_executor=info".into()),
)
.init();
let args = Args::parse();
match args.transport.trim().to_ascii_lowercase().as_str() {
"unix_socket" | "unix" | "uds" => {
info!(socket = %args.unix_socket.display(), "aether-executor started");
server::serve_unix(&args.unix_socket).await?;
}
"tcp" => {
info!(bind = %args.bind, "aether-executor started");
server::serve_tcp(&args.bind).await?;
}
other => {
return Err(format!("unsupported executor transport: {other}").into());
}
}
Ok(())
}

View File

@@ -0,0 +1,38 @@
use aether_contracts::StreamFrame;
use bytes::Bytes;
use crate::ExecutorClientError;
pub fn encode_frame(frame: &StreamFrame) -> Result<Bytes, ExecutorClientError> {
let mut raw = serde_json::to_vec(frame)?;
raw.push(b'\n');
Ok(Bytes::from(raw))
}
pub fn decode_frame(line: &[u8]) -> Result<StreamFrame, ExecutorClientError> {
Ok(serde_json::from_slice(line)?)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use aether_contracts::{StreamFrame, StreamFramePayload, StreamFrameType};
use super::{decode_frame, encode_frame};
#[test]
fn ndjson_round_trip_preserves_frame() {
let frame = StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([("content-type".into(), "text/event-stream".into())]),
},
};
let raw = encode_frame(&frame).expect("frame should encode");
let decoded = decode_frame(raw.trim_ascii_end()).expect("frame should decode");
assert_eq!(decoded, frame);
}
}

View File

@@ -0,0 +1,187 @@
use std::convert::Infallible;
use std::path::Path;
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
};
use async_stream::stream;
use axum::body::Body;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{Json, Router};
use base64::Engine as _;
use bytes::Bytes;
use futures_util::StreamExt;
use serde_json::json;
use crate::{encode_frame, ExecutorServiceError, SyncExecutor};
#[derive(Debug, Clone, Default)]
pub struct AppState {
executor: SyncExecutor,
}
pub fn build_router() -> Router {
Router::new()
.route("/health", get(health))
.route("/v1/execute/sync", post(execute_sync))
.route("/v1/execute/stream", post(execute_stream))
.with_state(AppState {
executor: SyncExecutor::new(),
})
}
pub async fn serve_tcp(bind: &str) -> Result<(), Box<dyn std::error::Error>> {
let listener = tokio::net::TcpListener::bind(bind).await?;
axum::serve(listener, build_router()).await?;
Ok(())
}
pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
if let Some(parent) = socket_path.parent() {
std::fs::create_dir_all(parent)?;
}
if socket_path.exists() {
std::fs::remove_file(socket_path)?;
}
let listener = tokio::net::UnixListener::bind(socket_path)?;
axum::serve(listener, build_router()).await?;
Ok(())
}
async fn health() -> impl IntoResponse {
Json(json!({"status": "ok"}))
}
async fn execute_sync(
State(state): State<AppState>,
Json(plan): Json<ExecutionPlan>,
) -> Result<Json<ExecutionResult>, AppError> {
state
.executor
.execute_sync(plan)
.await
.map(Json)
.map_err(AppError)
}
async fn execute_stream(
State(state): State<AppState>,
Json(plan): Json<ExecutionPlan>,
) -> Result<Response, AppError> {
let execution = state
.executor
.execute_stream(plan)
.await
.map_err(AppError)?;
let status_code = execution.status_code;
let response_headers = execution.headers.clone();
let started_at = execution.started_at;
let upstream_response = execution.response;
let body_stream = stream! {
let headers_frame = StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code,
headers: response_headers,
},
};
yield Ok::<Bytes, Infallible>(encode_frame(&headers_frame).expect("headers frame should encode"));
let mut upstream_bytes = 0u64;
let mut bytes_stream = upstream_response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
Ok(chunk) => {
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,
},
};
yield Ok::<Bytes, Infallible>(encode_frame(&frame).expect("data frame should encode"));
}
Err(err) => {
let frame = StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message: err.to_string(),
upstream_status: Some(status_code),
retryable: false,
failover_recommended: false,
},
},
};
yield Ok::<Bytes, Infallible>(encode_frame(&frame).expect("error frame should encode"));
break;
}
}
}
let telemetry_frame = StreamFrame {
frame_type: StreamFrameType::Telemetry,
payload: StreamFramePayload::Telemetry {
telemetry: ExecutionTelemetry {
ttfb_ms: None,
elapsed_ms: Some(started_at.elapsed().as_millis() as u64),
upstream_bytes: Some(upstream_bytes),
},
},
};
yield Ok::<Bytes, Infallible>(encode_frame(&telemetry_frame).expect("telemetry frame should encode"));
yield Ok::<Bytes, Infallible>(encode_frame(&StreamFrame::eof()).expect("eof frame should encode"));
};
let mut response = Response::new(Body::from_stream(body_stream));
*response.status_mut() = StatusCode::OK;
response.headers_mut().insert(
axum::http::header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("application/x-ndjson"),
);
Ok(response)
}
#[derive(Debug)]
struct AppError(ExecutorServiceError);
impl IntoResponse for AppError {
fn into_response(self) -> Response {
let status_code = match self.0 {
ExecutorServiceError::StreamUnsupported
| ExecutorServiceError::RequestBodyRequired
| ExecutorServiceError::BodyDecode(_)
| ExecutorServiceError::UnsupportedContentEncoding(_)
| ExecutorServiceError::ProxyUnsupported
| ExecutorServiceError::TlsProfileUnsupported
| ExecutorServiceError::DelegateUnsupported
| ExecutorServiceError::InvalidMethod(_)
| ExecutorServiceError::InvalidHeaderName(_)
| ExecutorServiceError::InvalidHeaderValue(_)
| ExecutorServiceError::InvalidProxy(_)
| ExecutorServiceError::BodyEncode(_) => StatusCode::BAD_REQUEST,
ExecutorServiceError::ClientBuild(_)
| ExecutorServiceError::UpstreamRequest(_)
| ExecutorServiceError::RelayError(_)
| ExecutorServiceError::InvalidJson(_) => StatusCode::BAD_GATEWAY,
};
(
status_code,
Json(json!({
"error": self.0.to_string(),
})),
)
.into_response()
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,5 @@
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransportMode {
UnixSocketHttp,
TcpHttp,
}

View File

@@ -0,0 +1,26 @@
[package]
name = "aether-gateway"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Rust ingress gateway for Aether phase 3a transparent proxy"
[dependencies]
aether-contracts.workspace = true
async-stream.workspace = true
axum = { version = "0.8" }
base64.workspace = true
bytes.workspace = true
clap = { version = "4", features = ["derive", "env"] }
futures-util.workspace = true
http.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
tokio-util.workspace = true
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
uuid.workspace = true

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)
}