mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
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:
26
crates/aether-gateway/Cargo.toml
Normal file
26
crates/aether-gateway/Cargo.toml
Normal 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
|
||||
17
crates/aether-gateway/src/constants.rs
Normal file
17
crates/aether-gateway/src/constants.rs
Normal 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";
|
||||
299
crates/aether-gateway/src/control.rs
Normal file
299
crates/aether-gateway/src/control.rs
Normal 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/")
|
||||
}
|
||||
68
crates/aether-gateway/src/error.rs
Normal file
68
crates/aether-gateway/src/error.rs
Normal 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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
396
crates/aether-gateway/src/executor.rs
Normal file
396
crates/aether-gateway/src/executor.rs
Normal 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),
|
||||
)
|
||||
}
|
||||
149
crates/aether-gateway/src/gateway.rs
Normal file
149
crates/aether-gateway/src/gateway.rs
Normal 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;
|
||||
290
crates/aether-gateway/src/gateway/tests/ai_execute.rs
Normal file
290
crates/aether-gateway/src/gateway/tests/ai_execute.rs
Normal 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();
|
||||
}
|
||||
222
crates/aether-gateway/src/gateway/tests/control.rs
Normal file
222
crates/aether-gateway/src/gateway/tests/control.rs
Normal 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();
|
||||
}
|
||||
537
crates/aether-gateway/src/gateway/tests/files.rs
Normal file
537
crates/aether-gateway/src/gateway/tests/files.rs
Normal 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();
|
||||
}
|
||||
35
crates/aether-gateway/src/gateway/tests/mod.rs
Normal file
35
crates/aether-gateway/src/gateway/tests/mod.rs
Normal 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)
|
||||
}
|
||||
163
crates/aether-gateway/src/gateway/tests/proxy.rs
Normal file
163
crates/aether-gateway/src/gateway/tests/proxy.rs
Normal 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();
|
||||
}
|
||||
357
crates/aether-gateway/src/gateway/tests/video.rs
Normal file
357
crates/aether-gateway/src/gateway/tests/video.rs
Normal 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();
|
||||
}
|
||||
193
crates/aether-gateway/src/handlers.rs
Normal file
193
crates/aether-gateway/src/handlers.rs
Normal 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)
|
||||
}
|
||||
83
crates/aether-gateway/src/headers.rs
Normal file
83
crates/aether-gateway/src/headers.rs
Normal 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)
|
||||
}
|
||||
6
crates/aether-gateway/src/lib.rs
Normal file
6
crates/aether-gateway/src/lib.rs
Normal 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,
|
||||
};
|
||||
62
crates/aether-gateway/src/main.rs
Normal file
62
crates/aether-gateway/src/main.rs
Normal 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(())
|
||||
}
|
||||
92
crates/aether-gateway/src/response.rs
Normal file
92
crates/aether-gateway/src/response.rs
Normal 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)
|
||||
}
|
||||
Reference in New Issue
Block a user