mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: 移除独立 hub/proxy/executor/gateway crate,统一为 gateway tunnel 架构
- 删除 aether-hub、aether-proxy 独立项目及其 Dockerfile/配置 - 删除 crates/aether-executor 和 crates/aether-gateway 全部模块 - 新增 apps/ 目录作为应用入口 - 将 hub 概念重构为 gateway tunnel transport - 将 executor 重构为 execution runtime - 新增 tunnel.rs 合约定义和 testkit tunnel/execution_runtime 模块 - 更新 Python 服务层和测试适配新架构命名
This commit is contained in:
3
apps/aether-gateway/src/execution_runtime/constants.rs
Normal file
3
apps/aether-gateway/src/execution_runtime/constants.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
pub(crate) const MAX_ERROR_BODY_BYTES: usize = 16_384;
|
||||
pub(crate) const MAX_STREAM_PREFETCH_FRAMES: usize = 5;
|
||||
pub(crate) const MAX_STREAM_PREFETCH_BYTES: usize = 16_384;
|
||||
273
apps/aether-gateway/src/execution_runtime/mod.rs
Normal file
273
apps/aether-gateway/src/execution_runtime/mod.rs
Normal file
@@ -0,0 +1,273 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
mod constants;
|
||||
pub(crate) mod ndjson;
|
||||
#[cfg(test)]
|
||||
pub(crate) mod remote_compat;
|
||||
mod server;
|
||||
pub(crate) mod stream;
|
||||
mod stream_pump;
|
||||
pub(crate) mod submission;
|
||||
pub(crate) mod sync;
|
||||
pub(crate) mod transport;
|
||||
|
||||
pub(crate) use self::constants::{
|
||||
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||
};
|
||||
pub use server::{
|
||||
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
||||
build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,
|
||||
serve_execution_runtime_unix,
|
||||
};
|
||||
pub(crate) use stream::{
|
||||
execute_execution_runtime_stream, maybe_execute_via_execution_runtime_stream,
|
||||
};
|
||||
pub(crate) use stream_pump::build_direct_execution_frame_stream;
|
||||
pub(crate) use sync::{
|
||||
execute_execution_runtime_sync, maybe_build_local_sync_finalize_response,
|
||||
maybe_build_local_video_error_response, maybe_build_local_video_success_outcome,
|
||||
maybe_execute_via_execution_runtime_sync, resolve_local_sync_error_background_report_kind,
|
||||
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
|
||||
};
|
||||
pub(crate) use transport::{
|
||||
execute_sync_plan as execute_execution_runtime_sync_plan, DirectSyncExecutionRuntime,
|
||||
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ExecutionStrategy {
|
||||
GatewayAffinityForward,
|
||||
RawPublicProxy,
|
||||
LocalSameFormat,
|
||||
LocalCrossFormat,
|
||||
}
|
||||
|
||||
impl ExecutionStrategy {
|
||||
pub(crate) const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::GatewayAffinityForward => "gateway_affinity_forward",
|
||||
Self::RawPublicProxy => "raw_public_proxy",
|
||||
Self::LocalSameFormat => "local_same_format",
|
||||
Self::LocalCrossFormat => "local_cross_format",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ConversionMode {
|
||||
None,
|
||||
RequestOnly,
|
||||
ResponseOnly,
|
||||
Bidirectional,
|
||||
}
|
||||
|
||||
impl ConversionMode {
|
||||
pub(crate) const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::None => "none",
|
||||
Self::RequestOnly => "request_only",
|
||||
Self::ResponseOnly => "response_only",
|
||||
Self::Bidirectional => "bidirectional",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub(crate) struct ClientIntent {
|
||||
pub(crate) client_contract: String,
|
||||
pub(crate) method: String,
|
||||
pub(crate) request_path: String,
|
||||
pub(crate) is_stream: bool,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) requested_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) original_request_headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) original_request_body: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub(crate) struct CompiledProviderRequest {
|
||||
pub(crate) execution_strategy: ExecutionStrategy,
|
||||
pub(crate) provider_contract: String,
|
||||
pub(crate) conversion_mode: ConversionMode,
|
||||
pub(crate) request_id: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) candidate_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) endpoint_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) key_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) mapped_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_request_headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_request_body_base64: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) upstream_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub(crate) struct ExecutionTerminalResult {
|
||||
pub(crate) status_code: u16,
|
||||
#[serde(default)]
|
||||
pub(crate) provider_headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_body_base64: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) telemetry: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_usage: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub(crate) struct FinalizedExecutionOutcome {
|
||||
pub(crate) report_kind: String,
|
||||
pub(crate) status_code: u16,
|
||||
pub(crate) terminal_state: FinalizedExecutionState,
|
||||
pub(crate) client_contract: String,
|
||||
pub(crate) provider_contract: String,
|
||||
pub(crate) execution_strategy: ExecutionStrategy,
|
||||
pub(crate) conversion_mode: ConversionMode,
|
||||
pub(crate) request_id: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) candidate_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) user_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) api_key_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) username: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) api_key_name: Option<String>,
|
||||
pub(crate) provider_name: String,
|
||||
pub(crate) model: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) target_model: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_endpoint_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_api_key_id: Option<String>,
|
||||
pub(crate) request_type: String,
|
||||
pub(crate) is_stream: bool,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) response_time_ms: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) first_byte_time_ms: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) request_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_request_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_request: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_response: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) client_response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) client_response: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) standardized_usage: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) request_metadata: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) audit_payload: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum FinalizedExecutionState {
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
pub(crate) fn append_execution_contract_fields(
|
||||
object: &mut Map<String, Value>,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
client_contract: &str,
|
||||
provider_contract: &str,
|
||||
) {
|
||||
object.insert(
|
||||
"execution_strategy".to_string(),
|
||||
Value::String(execution_strategy.as_str().to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"conversion_mode".to_string(),
|
||||
Value::String(conversion_mode.as_str().to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"client_contract".to_string(),
|
||||
Value::String(client_contract.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"provider_contract".to_string(),
|
||||
Value::String(provider_contract.to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn append_execution_contract_fields_to_value(
|
||||
value: Value,
|
||||
execution_strategy: ExecutionStrategy,
|
||||
conversion_mode: ConversionMode,
|
||||
client_contract: &str,
|
||||
provider_contract: &str,
|
||||
) -> Value {
|
||||
match value {
|
||||
Value::Object(mut object) => {
|
||||
append_execution_contract_fields(
|
||||
&mut object,
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
client_contract,
|
||||
provider_contract,
|
||||
);
|
||||
Value::Object(object)
|
||||
}
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn execution_contract_helper_appends_unified_fields() {
|
||||
let value = append_execution_contract_fields_to_value(
|
||||
json!({"provider_api_format": "gemini:chat"}),
|
||||
ExecutionStrategy::LocalCrossFormat,
|
||||
ConversionMode::Bidirectional,
|
||||
"openai:chat",
|
||||
"gemini:chat",
|
||||
);
|
||||
|
||||
assert_eq!(value["execution_strategy"], "local_cross_format");
|
||||
assert_eq!(value["conversion_mode"], "bidirectional");
|
||||
assert_eq!(value["client_contract"], "openai:chat");
|
||||
assert_eq!(value["provider_contract"], "gemini:chat");
|
||||
assert_eq!(value["provider_api_format"], "gemini:chat");
|
||||
}
|
||||
}
|
||||
41
apps/aether-gateway/src/execution_runtime/ndjson.rs
Normal file
41
apps/aether-gateway/src/execution_runtime/ndjson.rs
Normal file
@@ -0,0 +1,41 @@
|
||||
use std::io::Error as IoError;
|
||||
|
||||
use aether_contracts::StreamFrame;
|
||||
use axum::body::Bytes;
|
||||
|
||||
use crate::gateway::GatewayError;
|
||||
|
||||
pub(crate) fn encode_stream_frame_ndjson(frame: &StreamFrame) -> Result<Bytes, IoError> {
|
||||
let mut raw = serde_json::to_vec(frame).map_err(|err| IoError::other(err.to_string()))?;
|
||||
raw.push(b'\n');
|
||||
Ok(Bytes::from(raw))
|
||||
}
|
||||
|
||||
pub(crate) fn decode_stream_frame_ndjson(line: &[u8]) -> Result<StreamFrame, GatewayError> {
|
||||
serde_json::from_slice(line).map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{StreamFramePayload, StreamFrameType};
|
||||
|
||||
use super::{decode_stream_frame_ndjson, encode_stream_frame_ndjson};
|
||||
|
||||
#[test]
|
||||
fn ndjson_round_trip_preserves_frame() {
|
||||
let frame = aether_contracts::StreamFrame {
|
||||
frame_type: StreamFrameType::Headers,
|
||||
payload: StreamFramePayload::Headers {
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("content-type".into(), "text/event-stream".into())]),
|
||||
},
|
||||
};
|
||||
|
||||
let raw = encode_stream_frame_ndjson(&frame).expect("frame should encode");
|
||||
let decoded =
|
||||
decode_stream_frame_ndjson(raw.trim_ascii_end()).expect("frame should decode");
|
||||
assert_eq!(decoded, frame);
|
||||
}
|
||||
}
|
||||
83
apps/aether-gateway/src/execution_runtime/remote_compat.rs
Normal file
83
apps/aether-gateway/src/execution_runtime/remote_compat.rs
Normal file
@@ -0,0 +1,83 @@
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult};
|
||||
|
||||
use crate::gateway::constants::TRACE_ID_HEADER;
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
fn build_remote_execution_runtime_request(
|
||||
state: &AppState,
|
||||
remote_execution_runtime_base_url: &str,
|
||||
path: &str,
|
||||
trace_id: Option<&str>,
|
||||
plan: &ExecutionPlan,
|
||||
) -> reqwest::RequestBuilder {
|
||||
let mut request = state
|
||||
.client
|
||||
.post(format!("{remote_execution_runtime_base_url}{path}"))
|
||||
.json(plan);
|
||||
if let Some(trace_id) = trace_id.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
request = request.header(TRACE_ID_HEADER, trace_id);
|
||||
}
|
||||
request
|
||||
}
|
||||
|
||||
pub(crate) async fn post_sync_plan_to_remote_execution_runtime(
|
||||
state: &AppState,
|
||||
remote_execution_runtime_base_url: &str,
|
||||
trace_id: Option<&str>,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<reqwest::Response, GatewayError> {
|
||||
build_remote_execution_runtime_request(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
"/v1/execute/sync",
|
||||
trace_id,
|
||||
plan,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn post_stream_plan_to_remote_execution_runtime(
|
||||
state: &AppState,
|
||||
remote_execution_runtime_base_url: &str,
|
||||
trace_id: Option<&str>,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<reqwest::Response, GatewayError> {
|
||||
build_remote_execution_runtime_request(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
"/v1/execute/stream",
|
||||
trace_id,
|
||||
plan,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_sync_plan_via_remote_execution_runtime(
|
||||
state: &AppState,
|
||||
remote_execution_runtime_base_url: &str,
|
||||
trace_id: Option<&str>,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, GatewayError> {
|
||||
let response = post_sync_plan_to_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
trace_id,
|
||||
plan,
|
||||
)
|
||||
.await?;
|
||||
if response.status() != http::StatusCode::OK {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"execution runtime returned HTTP {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
response
|
||||
.json::<ExecutionResult>()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
602
apps/aether-gateway/src/execution_runtime/server.rs
Normal file
602
apps/aether-gateway/src/execution_runtime/server.rs
Normal file
@@ -0,0 +1,602 @@
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use aether_runtime::{
|
||||
maybe_hold_axum_response_permit, prometheus_response, service_up_sample, AdmissionPermit,
|
||||
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
|
||||
DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
|
||||
MetricSample,
|
||||
};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use serde_json::json;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::gateway::execution_runtime::{
|
||||
build_direct_execution_frame_stream, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
|
||||
};
|
||||
|
||||
const EXECUTION_RUNTIME_COMPONENT: &str = "aether-gateway-execution-runtime";
|
||||
const REQUEST_GATE_NAME: &str = "execution_runtime_requests";
|
||||
const DISTRIBUTED_REQUEST_GATE_NAME: &str = "execution_runtime_requests_distributed";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct ExecutionRuntimeAppState {
|
||||
execution_runtime: DirectSyncExecutionRuntime,
|
||||
request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
|
||||
}
|
||||
|
||||
impl ExecutionRuntimeAppState {
|
||||
fn with_request_concurrency_limit(limit: Option<usize>) -> Self {
|
||||
Self {
|
||||
execution_runtime: DirectSyncExecutionRuntime::new(),
|
||||
request_gate: limit
|
||||
.filter(|limit| *limit > 0)
|
||||
.map(|limit| Arc::new(ConcurrencyGate::new(REQUEST_GATE_NAME, limit))),
|
||||
distributed_request_gate: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
|
||||
self.distributed_request_gate = Some(Arc::new(gate));
|
||||
self
|
||||
}
|
||||
|
||||
fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
|
||||
self.request_gate.as_ref().map(|gate| gate.snapshot())
|
||||
}
|
||||
|
||||
async fn distributed_request_concurrency_snapshot(
|
||||
&self,
|
||||
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
|
||||
match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => gate.snapshot().await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn metric_samples(&self) -> Vec<MetricSample> {
|
||||
let mut samples = vec![service_up_sample(EXECUTION_RUNTIME_COMPONENT)];
|
||||
if let Some(snapshot) = self.request_concurrency_snapshot() {
|
||||
samples.extend(snapshot.to_metric_samples(REQUEST_GATE_NAME));
|
||||
}
|
||||
if let Some(gate) = self.distributed_request_gate.as_ref() {
|
||||
match gate.snapshot().await {
|
||||
Ok(snapshot) => {
|
||||
samples.extend(snapshot.to_metric_samples(DISTRIBUTED_REQUEST_GATE_NAME));
|
||||
}
|
||||
Err(_) => samples.push(
|
||||
MetricSample::new(
|
||||
"concurrency_unavailable",
|
||||
"Whether the distributed concurrency gate is currently unavailable.",
|
||||
MetricKind::Gauge,
|
||||
1,
|
||||
)
|
||||
.with_labels(vec![MetricLabel::new(
|
||||
"gate",
|
||||
DISTRIBUTED_REQUEST_GATE_NAME,
|
||||
)]),
|
||||
),
|
||||
}
|
||||
}
|
||||
samples
|
||||
}
|
||||
|
||||
async fn try_acquire_request_permit(
|
||||
&self,
|
||||
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
|
||||
let local = self
|
||||
.request_gate
|
||||
.as_ref()
|
||||
.map(|gate| gate.try_acquire())
|
||||
.transpose()
|
||||
.map_err(RequestAdmissionError::Local)?;
|
||||
let distributed = match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => Some(
|
||||
gate.try_acquire()
|
||||
.await
|
||||
.map_err(RequestAdmissionError::Distributed)?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
Ok(AdmissionPermit::from_parts(local, distributed))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_execution_runtime_router() -> Router {
|
||||
build_execution_runtime_router_with_request_concurrency_limit(None)
|
||||
}
|
||||
|
||||
pub fn build_execution_runtime_router_with_request_concurrency_limit(
|
||||
limit: Option<usize>,
|
||||
) -> Router {
|
||||
build_execution_runtime_router_with_request_gates(limit, None)
|
||||
}
|
||||
|
||||
pub fn build_execution_runtime_router_with_request_gates(
|
||||
limit: Option<usize>,
|
||||
distributed_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Router {
|
||||
let state = match distributed_gate {
|
||||
Some(gate) => ExecutionRuntimeAppState::with_request_concurrency_limit(limit)
|
||||
.with_distributed_request_gate(gate),
|
||||
None => ExecutionRuntimeAppState::with_request_concurrency_limit(limit),
|
||||
};
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/metrics", get(metrics))
|
||||
.route("/v1/execute/sync", post(execute_sync))
|
||||
.route("/v1/execute/stream", post(execute_stream))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
pub async fn serve_execution_runtime_tcp(
|
||||
bind: &str,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let listener = tokio::net::TcpListener::bind(bind).await?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_execution_runtime_router_with_request_gates(
|
||||
max_in_flight_requests,
|
||||
distributed_request_gate,
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn serve_execution_runtime_unix(
|
||||
socket_path: &Path,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if let Some(parent) = socket_path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
if socket_path.exists() {
|
||||
std::fs::remove_file(socket_path)?;
|
||||
}
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(socket_path)?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_execution_runtime_router_with_request_gates(
|
||||
max_in_flight_requests,
|
||||
distributed_request_gate,
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health(State(state): State<ExecutionRuntimeAppState>) -> impl IntoResponse {
|
||||
let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
let distributed_request_concurrency = state
|
||||
.distributed_request_concurrency_snapshot()
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"component": EXECUTION_RUNTIME_COMPONENT,
|
||||
"request_concurrency": request_concurrency,
|
||||
"distributed_request_concurrency": distributed_request_concurrency,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn metrics(State(state): State<ExecutionRuntimeAppState>) -> Response {
|
||||
prometheus_response(&state.metric_samples().await)
|
||||
}
|
||||
|
||||
async fn execute_sync(
|
||||
State(state): State<ExecutionRuntimeAppState>,
|
||||
request: Request,
|
||||
) -> Result<Response, ExecutionRuntimeAppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let result = state
|
||||
.execution_runtime
|
||||
.execute_sync(plan)
|
||||
.await
|
||||
.map_err(|err| ExecutionRuntimeAppError(ExecutionRuntimeServerError::Transport(err)))?;
|
||||
Ok(maybe_hold_axum_response_permit(
|
||||
Json(result).into_response(),
|
||||
request_permit,
|
||||
))
|
||||
}
|
||||
|
||||
async fn execute_stream(
|
||||
State(state): State<ExecutionRuntimeAppState>,
|
||||
request: Request,
|
||||
) -> Result<Response, ExecutionRuntimeAppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let execution = state
|
||||
.execution_runtime
|
||||
.execute_stream(plan)
|
||||
.await
|
||||
.map_err(|err| ExecutionRuntimeAppError(ExecutionRuntimeServerError::Transport(err)))?;
|
||||
|
||||
let mut response = Response::new(Body::from_stream(build_direct_execution_frame_stream(
|
||||
execution,
|
||||
)));
|
||||
*response.status_mut() = StatusCode::OK;
|
||||
response.headers_mut().insert(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
axum::http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
Ok(maybe_hold_axum_response_permit(response, request_permit))
|
||||
}
|
||||
|
||||
async fn acquire_request_permit(
|
||||
state: &ExecutionRuntimeAppState,
|
||||
) -> Result<Option<AdmissionPermit>, ExecutionRuntimeAppError> {
|
||||
match state.try_acquire_request_permit().await {
|
||||
Ok(permit) => Ok(permit),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { gate, limit }))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
|
||||
gate,
|
||||
limit,
|
||||
}))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
|
||||
gate,
|
||||
limit,
|
||||
..
|
||||
})) => Err(ExecutionRuntimeAppError(
|
||||
ExecutionRuntimeServerError::Overloaded { gate, limit },
|
||||
)),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => Err(
|
||||
ExecutionRuntimeAppError(ExecutionRuntimeServerError::RequestRead(format!(
|
||||
"execution runtime request concurrency gate {gate} is closed"
|
||||
))),
|
||||
),
|
||||
Err(RequestAdmissionError::Distributed(
|
||||
DistributedConcurrencyError::InvalidConfiguration(message),
|
||||
)) => Err(ExecutionRuntimeAppError(
|
||||
ExecutionRuntimeServerError::RequestRead(message),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum RequestAdmissionError {
|
||||
Local(ConcurrencyError),
|
||||
Distributed(DistributedConcurrencyError),
|
||||
}
|
||||
|
||||
async fn parse_request_json<T>(request: Request) -> Result<T, ExecutionRuntimeAppError>
|
||||
where
|
||||
T: serde::de::DeserializeOwned,
|
||||
{
|
||||
let body = to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
ExecutionRuntimeAppError(ExecutionRuntimeServerError::RequestRead(err.to_string()))
|
||||
})?;
|
||||
serde_json::from_slice(&body).map_err(|err| {
|
||||
ExecutionRuntimeAppError(ExecutionRuntimeServerError::InvalidRequestJson(err))
|
||||
})
|
||||
}
|
||||
|
||||
fn build_overloaded_response(message: &str) -> Response {
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"type": "overloaded",
|
||||
"message": message,
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
enum ExecutionRuntimeServerError {
|
||||
#[error("failed to read execution runtime request body: {0}")]
|
||||
RequestRead(String),
|
||||
#[error("execution runtime request body is not valid JSON: {0}")]
|
||||
InvalidRequestJson(serde_json::Error),
|
||||
#[error("execution runtime overloaded: gate {gate} saturated at {limit}")]
|
||||
Overloaded { gate: &'static str, limit: usize },
|
||||
#[error(transparent)]
|
||||
Transport(#[from] ExecutionRuntimeTransportError),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ExecutionRuntimeAppError(ExecutionRuntimeServerError);
|
||||
|
||||
impl IntoResponse for ExecutionRuntimeAppError {
|
||||
fn into_response(self) -> Response {
|
||||
let status_code = match self.0 {
|
||||
ExecutionRuntimeServerError::RequestRead(_)
|
||||
| ExecutionRuntimeServerError::InvalidRequestJson(_) => StatusCode::BAD_REQUEST,
|
||||
ExecutionRuntimeServerError::Overloaded { .. } => {
|
||||
return build_overloaded_response(&self.0.to_string());
|
||||
}
|
||||
ExecutionRuntimeServerError::Transport(
|
||||
ExecutionRuntimeTransportError::StreamUnsupported
|
||||
| ExecutionRuntimeTransportError::RequestBodyRequired
|
||||
| ExecutionRuntimeTransportError::BodyDecode(_)
|
||||
| ExecutionRuntimeTransportError::UnsupportedContentEncoding(_)
|
||||
| ExecutionRuntimeTransportError::ProxyUnsupported
|
||||
| ExecutionRuntimeTransportError::InvalidMethod(_)
|
||||
| ExecutionRuntimeTransportError::InvalidHeaderName(_)
|
||||
| ExecutionRuntimeTransportError::InvalidHeaderValue(_)
|
||||
| ExecutionRuntimeTransportError::InvalidProxy(_)
|
||||
| ExecutionRuntimeTransportError::BodyEncode(_),
|
||||
) => StatusCode::BAD_REQUEST,
|
||||
ExecutionRuntimeServerError::Transport(
|
||||
ExecutionRuntimeTransportError::ClientBuild(_)
|
||||
| ExecutionRuntimeTransportError::UpstreamRequest(_)
|
||||
| ExecutionRuntimeTransportError::RelayError(_)
|
||||
| ExecutionRuntimeTransportError::InvalidJson(_),
|
||||
) => StatusCode::BAD_GATEWAY,
|
||||
};
|
||||
|
||||
(
|
||||
status_code,
|
||||
Json(json!({
|
||||
"error": self.0.to_string(),
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_execution_runtime_router_with_request_concurrency_limit,
|
||||
build_execution_runtime_router_with_request_gates, DISTRIBUTED_REQUEST_GATE_NAME,
|
||||
};
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
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).await.expect("server should run");
|
||||
});
|
||||
(format!("http://{addr}"), handle)
|
||||
}
|
||||
|
||||
fn stream_plan(url: String) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "GET".into(),
|
||||
url,
|
||||
headers: std::collections::BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(30_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execution_runtime_rejects_second_in_flight_stream_request_with_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let runtime = build_execution_runtime_router_with_request_concurrency_limit(Some(1));
|
||||
let (runtime_url, runtime_handle) = start_server(runtime).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{runtime_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{runtime_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execution_runtime_rejects_second_in_flight_stream_request_with_distributed_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let distributed_gate = aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
DISTRIBUTED_REQUEST_GATE_NAME,
|
||||
1,
|
||||
);
|
||||
let runtime_a =
|
||||
build_execution_runtime_router_with_request_gates(None, Some(distributed_gate.clone()));
|
||||
let runtime_b =
|
||||
build_execution_runtime_router_with_request_gates(None, Some(distributed_gate));
|
||||
let (runtime_a_url, runtime_a_handle) = start_server(runtime_a).await;
|
||||
let (runtime_b_url, runtime_b_handle) = start_server(runtime_b).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{runtime_a_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{runtime_b_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
runtime_a_handle.abort();
|
||||
runtime_b_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execution_runtime_exposes_request_concurrency_metrics() {
|
||||
let runtime = build_execution_runtime_router_with_request_gates(
|
||||
Some(4),
|
||||
Some(aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
DISTRIBUTED_REQUEST_GATE_NAME,
|
||||
6,
|
||||
)),
|
||||
);
|
||||
let (runtime_url, runtime_handle) = start_server(runtime).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{runtime_url}/metrics"))
|
||||
.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("text/plain; version=0.0.4; charset=utf-8")
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("service_up{service=\"aether-gateway-execution-runtime\"} 1"));
|
||||
assert!(
|
||||
body.contains("concurrency_available_permits{gate=\"execution_runtime_requests\"} 4")
|
||||
);
|
||||
assert!(body.contains(
|
||||
"concurrency_available_permits{gate=\"execution_runtime_requests_distributed\"} 6"
|
||||
));
|
||||
|
||||
runtime_handle.abort();
|
||||
}
|
||||
}
|
||||
47
apps/aether-gateway/src/execution_runtime/stream.rs
Normal file
47
apps/aether-gateway/src/execution_runtime/stream.rs
Normal file
@@ -0,0 +1,47 @@
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::Response;
|
||||
|
||||
use crate::gateway::{AppState, GatewayControlDecision, GatewayError};
|
||||
|
||||
mod error;
|
||||
mod execution;
|
||||
|
||||
pub(crate) use execution::execute_execution_runtime_stream;
|
||||
|
||||
pub(crate) async fn maybe_execute_via_execution_runtime_stream(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
trace_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let Some(decision) = decision else {
|
||||
return Ok(None);
|
||||
};
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
let _ = state;
|
||||
if parts.method != http::Method::POST {
|
||||
return Ok(None);
|
||||
}
|
||||
return crate::gateway::ai_pipeline::planner::maybe_execute_stream_local_path(
|
||||
state, parts, body_bytes, trace_id, decision,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
#[cfg(test)]
|
||||
{
|
||||
if state
|
||||
.test_remote_execution_runtime_base_url()
|
||||
.unwrap_or_default()
|
||||
.is_empty()
|
||||
&& parts.method != http::Method::POST
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
crate::gateway::ai_pipeline::planner::maybe_execute_stream_local_path(
|
||||
state, parts, body_bytes, trace_id, decision,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
225
apps/aether-gateway/src/execution_runtime/stream/error.rs
Normal file
225
apps/aether-gateway/src/execution_runtime/stream/error.rs
Normal file
@@ -0,0 +1,225 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{StreamFrame, StreamFramePayload};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::json;
|
||||
use tokio_util::codec::{FramedRead, LinesCodec};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::gateway::execution_runtime::ndjson::decode_stream_frame_ndjson;
|
||||
use crate::gateway::execution_runtime::submission::{has_nested_error, strip_utf8_bom_and_ws};
|
||||
use crate::gateway::{build_client_response_from_parts, GatewayControlDecision, GatewayError};
|
||||
use crate::gateway::{
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) enum StreamPrefetchInspection {
|
||||
NeedMore,
|
||||
NonError,
|
||||
EmbeddedError(serde_json::Value),
|
||||
}
|
||||
|
||||
pub(super) fn decode_stream_error_body(
|
||||
headers: &BTreeMap<String, String>,
|
||||
error_body: &[u8],
|
||||
) -> (Option<serde_json::Value>, Option<String>) {
|
||||
if error_body.is_empty() {
|
||||
return (None, None);
|
||||
}
|
||||
|
||||
let content_type = headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.unwrap_or_default();
|
||||
let looks_json = content_type.contains("json") || content_type.ends_with("+json");
|
||||
if looks_json {
|
||||
if let Ok(json_body) = serde_json::from_slice::<serde_json::Value>(error_body) {
|
||||
return (Some(json_body), None);
|
||||
}
|
||||
}
|
||||
|
||||
(
|
||||
None,
|
||||
Some(base64::engine::general_purpose::STANDARD.encode(error_body)),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn inspect_prefetched_stream_body(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body: &[u8],
|
||||
) -> StreamPrefetchInspection {
|
||||
if body.is_empty() {
|
||||
return StreamPrefetchInspection::NeedMore;
|
||||
}
|
||||
|
||||
let stripped = strip_utf8_bom_and_ws(body);
|
||||
let content_type = headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.unwrap_or_default();
|
||||
let looks_json = content_type.contains("json") || content_type.ends_with("+json");
|
||||
if looks_json || stripped.starts_with(b"{") || stripped.starts_with(b"[") {
|
||||
if let Ok(json_body) = serde_json::from_slice::<serde_json::Value>(stripped) {
|
||||
return if has_nested_error(&json_body) {
|
||||
StreamPrefetchInspection::EmbeddedError(json_body)
|
||||
} else {
|
||||
StreamPrefetchInspection::NonError
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
let text = String::from_utf8_lossy(body);
|
||||
let mut saw_meaningful_line = false;
|
||||
for line in text.lines().take(MAX_STREAM_PREFETCH_FRAMES) {
|
||||
let line = line.trim_matches('\r').trim();
|
||||
if line.is_empty() || line.starts_with(':') || line.starts_with("event:") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let data_line = line.strip_prefix("data: ").unwrap_or(line).trim();
|
||||
if data_line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if data_line == "[DONE]" {
|
||||
return StreamPrefetchInspection::NonError;
|
||||
}
|
||||
|
||||
saw_meaningful_line = true;
|
||||
match serde_json::from_str::<serde_json::Value>(data_line) {
|
||||
Ok(json_body) => {
|
||||
return if has_nested_error(&json_body) {
|
||||
StreamPrefetchInspection::EmbeddedError(json_body)
|
||||
} else {
|
||||
StreamPrefetchInspection::NonError
|
||||
};
|
||||
}
|
||||
Err(_) => {
|
||||
if data_line.ends_with('}') || data_line.ends_with(']') {
|
||||
return StreamPrefetchInspection::NonError;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if saw_meaningful_line {
|
||||
StreamPrefetchInspection::NonError
|
||||
} else {
|
||||
StreamPrefetchInspection::NeedMore
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) 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, "execution runtime stream emitted error frame while collecting error body");
|
||||
break;
|
||||
}
|
||||
StreamFramePayload::Headers { .. } => {}
|
||||
}
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
pub(super) 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 = decode_stream_frame_ndjson(line.as_bytes())?;
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(super) fn build_execution_runtime_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),
|
||||
)
|
||||
}
|
||||
978
apps/aether-gateway/src/execution_runtime/stream/execution.rs
Normal file
978
apps/aether-gateway/src/execution_runtime/stream/execution.rs
Normal file
@@ -0,0 +1,978 @@
|
||||
use std::io::Error as IoError;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFramePayload};
|
||||
use async_stream::stream;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
use futures_util::stream::BoxStream;
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::codec::{FramedRead, LinesCodec};
|
||||
use tokio_util::io::StreamReader;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use super::error::{
|
||||
build_execution_runtime_error_response, collect_error_body, decode_stream_error_body,
|
||||
inspect_prefetched_stream_body, read_next_frame, StreamPrefetchInspection,
|
||||
};
|
||||
#[path = "execution_failures.rs"]
|
||||
mod execution_failures;
|
||||
use self::execution_failures::{
|
||||
build_stream_failure_from_execution_error, build_stream_failure_report,
|
||||
handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport,
|
||||
};
|
||||
use crate::gateway::ai_pipeline::runtime::{
|
||||
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
||||
};
|
||||
use crate::gateway::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
|
||||
use crate::gateway::execution_runtime::build_direct_execution_frame_stream;
|
||||
#[cfg(test)]
|
||||
use crate::gateway::execution_runtime::remote_compat::post_stream_plan_to_remote_execution_runtime;
|
||||
use crate::gateway::execution_runtime::submission::{
|
||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::gateway::execution_runtime::transport::{
|
||||
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution,
|
||||
};
|
||||
use crate::gateway::request_candidates::{
|
||||
current_unix_secs as current_request_candidate_unix_secs,
|
||||
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
|
||||
};
|
||||
use crate::gateway::scheduler::{
|
||||
resolve_core_stream_direct_finalize_report_kind,
|
||||
resolve_core_stream_error_finalize_report_kind, should_fallback_to_control_stream,
|
||||
should_retry_next_local_candidate_stream,
|
||||
};
|
||||
use crate::gateway::usage::submit_stream_report;
|
||||
use crate::gateway::{
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
maybe_build_stream_response_rewriter, AppState, GatewayControlDecision, GatewayError,
|
||||
GatewayStreamReportRequest, GatewaySyncReportRequest, MAX_STREAM_PREFETCH_BYTES,
|
||||
MAX_STREAM_PREFETCH_FRAMES,
|
||||
};
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
||||
pub(crate) async fn execute_execution_runtime_stream(
|
||||
state: &AppState,
|
||||
mut plan: ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
report_kind: Option<String>,
|
||||
mut report_context: Option<serde_json::Value>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
let execution = match DirectSyncExecutionRuntime::new()
|
||||
.execute_stream(plan.clone())
|
||||
.await
|
||||
{
|
||||
Ok(execution) => execution,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
error = %err,
|
||||
"gateway in-process stream execution unavailable"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
|
||||
return execute_stream_from_frame_stream(
|
||||
state,
|
||||
plan,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
report_kind,
|
||||
report_context,
|
||||
frame_stream,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
#[cfg(test)]
|
||||
{
|
||||
let remote_execution_runtime_base_url = state
|
||||
.test_remote_execution_runtime_base_url()
|
||||
.unwrap_or_default();
|
||||
if remote_execution_runtime_base_url.trim().is_empty() {
|
||||
let execution = match DirectSyncExecutionRuntime::new()
|
||||
.execute_stream(plan.clone())
|
||||
.await
|
||||
{
|
||||
Ok(execution) => execution,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
error = %err,
|
||||
"gateway in-process stream execution unavailable"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
|
||||
return execute_stream_from_frame_stream(
|
||||
state,
|
||||
plan,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
report_kind,
|
||||
report_context,
|
||||
frame_stream,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let response = match post_stream_plan_to_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
Some(trace_id),
|
||||
&plan,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway remote execution runtime stream unavailable"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
if response.status() != http::StatusCode::OK {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(response.status().as_u16()),
|
||||
Some("execution_runtime_http_error".to_string()),
|
||||
Some(format!(
|
||||
"execution runtime returned HTTP {}",
|
||||
response.status()
|
||||
)),
|
||||
None,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(plan.request_id.as_str()),
|
||||
plan.candidate_id.as_deref(),
|
||||
)?));
|
||||
}
|
||||
|
||||
let frame_stream = response
|
||||
.bytes_stream()
|
||||
.map_err(|err| IoError::other(err.to_string()))
|
||||
.boxed();
|
||||
return execute_stream_from_frame_stream(
|
||||
state,
|
||||
plan,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
report_kind,
|
||||
report_context,
|
||||
frame_stream,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_stream_from_frame_stream(
|
||||
state: &AppState,
|
||||
plan: ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
report_kind: Option<String>,
|
||||
report_context: Option<serde_json::Value>,
|
||||
frame_stream: BoxStream<'static, Result<Bytes, IoError>>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let request_id = plan.request_id.as_str();
|
||||
let candidate_id = plan.candidate_id.as_deref();
|
||||
let reader = StreamReader::new(frame_stream);
|
||||
let mut lines = FramedRead::new(reader, LinesCodec::new());
|
||||
|
||||
let first_frame = read_next_frame(&mut lines).await?.ok_or_else(|| {
|
||||
GatewayError::Internal("execution runtime stream ended before headers frame".to_string())
|
||||
})?;
|
||||
let StreamFramePayload::Headers {
|
||||
status_code,
|
||||
mut headers,
|
||||
} = first_frame.payload
|
||||
else {
|
||||
return Err(GatewayError::Internal(
|
||||
"execution runtime stream must start with headers frame".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
if should_retry_next_local_candidate_stream(plan_kind, report_context.as_ref(), status_code) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(status_code),
|
||||
Some("retryable_upstream_status".to_string()),
|
||||
Some(format!(
|
||||
"execution runtime stream returned retryable status {status_code}"
|
||||
)),
|
||||
None,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id,
|
||||
status_code,
|
||||
"gateway local stream decision retrying next candidate after retryable execution runtime status"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let stream_error_finalize_kind =
|
||||
resolve_core_stream_error_finalize_report_kind(plan_kind, status_code);
|
||||
|
||||
if should_fallback_to_control_stream(
|
||||
plan_kind,
|
||||
status_code,
|
||||
stream_error_finalize_kind.is_some(),
|
||||
) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(status_code),
|
||||
Some("control_fallback".to_string()),
|
||||
Some(format!(
|
||||
"stream decision fell back to control after status {status_code}"
|
||||
)),
|
||||
None,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if status_code >= 400 {
|
||||
let error_body = collect_error_body(&mut lines).await?;
|
||||
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
|
||||
let usage_report_kind = stream_error_finalize_kind
|
||||
.clone()
|
||||
.or_else(|| report_kind.clone())
|
||||
.unwrap_or_default();
|
||||
let usage_payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind: usage_report_kind,
|
||||
report_context: report_context.clone(),
|
||||
status_code,
|
||||
headers: headers.clone(),
|
||||
body_json: body_json.clone(),
|
||||
client_body_json: None,
|
||||
body_base64: body_base64.clone(),
|
||||
telemetry: None,
|
||||
};
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(status_code),
|
||||
Some("execution_runtime_stream_error".to_string()),
|
||||
Some(format!(
|
||||
"execution runtime stream returned error status {status_code}"
|
||||
)),
|
||||
None,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
if let Some(report_kind) = stream_error_finalize_kind {
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind,
|
||||
report_context,
|
||||
status_code,
|
||||
headers: headers.clone(),
|
||||
body_json,
|
||||
client_body_json: None,
|
||||
body_base64,
|
||||
telemetry: None,
|
||||
};
|
||||
let response =
|
||||
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload)
|
||||
.await?;
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_execution_runtime_error_response(
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
status_code,
|
||||
headers,
|
||||
error_body,
|
||||
)?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind);
|
||||
let normalized_stream_report_context =
|
||||
normalize_provider_private_report_context(report_context.as_ref());
|
||||
let mut private_stream_normalizer =
|
||||
maybe_build_provider_private_stream_normalizer(report_context.as_ref());
|
||||
let mut local_stream_rewriter =
|
||||
maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref());
|
||||
if private_stream_normalizer.is_some() || local_stream_rewriter.is_some() {
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "text/event-stream".to_string());
|
||||
}
|
||||
let mut prefetched_chunks: Vec<Bytes> = Vec::new();
|
||||
let mut provider_prefetched_body = Vec::new();
|
||||
let mut prefetched_body = Vec::new();
|
||||
let mut prefetched_inspection_body = Vec::new();
|
||||
let mut prefetched_telemetry: Option<ExecutionTelemetry> = None;
|
||||
let mut reached_eof = false;
|
||||
if let Some(ref report_kind) = direct_stream_finalize_kind {
|
||||
while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES
|
||||
&& prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES
|
||||
{
|
||||
let Some(frame) = (match read_next_frame(&mut lines).await {
|
||||
Ok(frame) => frame,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_stream_frame_decode_error",
|
||||
format!("failed to decode execution runtime stream frame: {err:?}"),
|
||||
502,
|
||||
);
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context.clone(),
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
&headers,
|
||||
prefetched_telemetry.clone(),
|
||||
&provider_prefetched_body,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}) else {
|
||||
reached_eof = true;
|
||||
break;
|
||||
};
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
let chunk = if let Some(chunk_b64) = chunk_b64 {
|
||||
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) {
|
||||
Ok(decoded) => decoded,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_stream_chunk_decode_error",
|
||||
format!(
|
||||
"failed to decode execution runtime stream chunk: {err}"
|
||||
),
|
||||
502,
|
||||
);
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context.clone(),
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
&headers,
|
||||
prefetched_telemetry.clone(),
|
||||
&prefetched_body,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else if let Some(text) = text {
|
||||
text.into_bytes()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
if chunk.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
provider_prefetched_body.extend_from_slice(&chunk);
|
||||
prefetched_inspection_body.extend_from_slice(&chunk);
|
||||
|
||||
let inspection =
|
||||
inspect_prefetched_stream_body(&headers, &prefetched_inspection_body);
|
||||
match inspection {
|
||||
StreamPrefetchInspection::EmbeddedError(body_json) => {
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind: report_kind.clone(),
|
||||
report_context: report_context.clone(),
|
||||
status_code,
|
||||
headers: headers.clone(),
|
||||
body_json: Some(body_json),
|
||||
client_body_json: None,
|
||||
body_base64: None,
|
||||
telemetry: prefetched_telemetry.clone(),
|
||||
};
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
&payload,
|
||||
)
|
||||
.await;
|
||||
let response = submit_local_core_error_or_sync_finalize(
|
||||
state, trace_id, decision, payload,
|
||||
)
|
||||
.await?;
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
StreamPrefetchInspection::NeedMore => {}
|
||||
StreamPrefetchInspection::NonError => {}
|
||||
}
|
||||
|
||||
let normalized_chunk = if let Some(normalizer) =
|
||||
private_stream_normalizer.as_mut()
|
||||
{
|
||||
match normalizer.push_chunk(&chunk) {
|
||||
Ok(normalized_chunk) => normalized_chunk,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_error",
|
||||
format!(
|
||||
"failed to normalize execution runtime stream chunk: {err:?}"
|
||||
),
|
||||
502,
|
||||
);
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context.clone(),
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
&headers,
|
||||
prefetched_telemetry.clone(),
|
||||
&provider_prefetched_body,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
chunk
|
||||
};
|
||||
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut() {
|
||||
match rewriter.push_chunk(&normalized_chunk) {
|
||||
Ok(rewritten_chunk) => rewritten_chunk,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_error",
|
||||
format!(
|
||||
"failed to rewrite execution runtime stream chunk: {err:?}"
|
||||
),
|
||||
502,
|
||||
);
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context.clone(),
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
&headers,
|
||||
prefetched_telemetry.clone(),
|
||||
&provider_prefetched_body,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
normalized_chunk
|
||||
};
|
||||
if !rewritten_chunk.is_empty() {
|
||||
prefetched_body.extend_from_slice(&rewritten_chunk);
|
||||
prefetched_chunks.push(Bytes::from(rewritten_chunk));
|
||||
}
|
||||
|
||||
if matches!(inspection, StreamPrefetchInspection::NonError) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
StreamFramePayload::Telemetry {
|
||||
telemetry: frame_telemetry,
|
||||
} => {
|
||||
prefetched_telemetry = Some(frame_telemetry);
|
||||
}
|
||||
StreamFramePayload::Eof { .. } => {
|
||||
reached_eof = true;
|
||||
break;
|
||||
}
|
||||
StreamFramePayload::Error { error } => {
|
||||
warn!(trace_id = %trace_id, error = %error.message, "execution runtime stream emitted error frame during prefetch");
|
||||
return handle_prefetch_stream_failure(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
report_context.clone(),
|
||||
request_id,
|
||||
candidate_id,
|
||||
report_kind,
|
||||
&headers,
|
||||
prefetched_telemetry.clone(),
|
||||
&provider_prefetched_body,
|
||||
build_stream_failure_from_execution_error(&error),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
StreamFramePayload::Headers { .. } => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_secs();
|
||||
state
|
||||
.usage_runtime
|
||||
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
|
||||
.await;
|
||||
state
|
||||
.usage_runtime
|
||||
.record_stream_started(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
status_code,
|
||||
&headers,
|
||||
prefetched_telemetry.as_ref(),
|
||||
)
|
||||
.await;
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Streaming,
|
||||
Some(status_code),
|
||||
None,
|
||||
None,
|
||||
prefetched_telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms),
|
||||
Some(candidate_started_unix_secs),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<Result<Bytes, IoError>>(16);
|
||||
let state_for_report = state.clone();
|
||||
let plan_for_report = plan.clone();
|
||||
let trace_id_owned = trace_id.to_string();
|
||||
let headers_for_report = headers.clone();
|
||||
let report_kind_owned = report_kind.clone();
|
||||
let report_context_owned = report_context.clone();
|
||||
let provider_prefetched_body_for_report = provider_prefetched_body.clone();
|
||||
let prefetched_body_for_report = prefetched_body.clone();
|
||||
let prefetched_chunks_for_body = prefetched_chunks.clone();
|
||||
let initial_telemetry = prefetched_telemetry.clone();
|
||||
let initial_reached_eof = reached_eof;
|
||||
let direct_stream_finalize_kind_owned = direct_stream_finalize_kind.clone();
|
||||
let candidate_started_unix_secs_for_report = candidate_started_unix_secs;
|
||||
tokio::spawn(async move {
|
||||
let mut provider_buffered_body = provider_prefetched_body_for_report;
|
||||
let mut buffered_body = prefetched_body_for_report;
|
||||
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
|
||||
let reached_eof = initial_reached_eof;
|
||||
let mut downstream_dropped = false;
|
||||
let mut terminal_failure: Option<StreamFailureReport> = None;
|
||||
|
||||
if !reached_eof {
|
||||
loop {
|
||||
let next_frame = match read_next_frame(&mut lines).await {
|
||||
Ok(frame) => frame,
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to decode execution runtime stream frame");
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"execution_runtime_stream_frame_decode_error",
|
||||
format!("failed to decode execution runtime stream frame: {err:?}"),
|
||||
502,
|
||||
));
|
||||
break;
|
||||
}
|
||||
};
|
||||
let Some(frame) = next_frame else {
|
||||
break;
|
||||
};
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
let chunk = if let Some(chunk_b64) = chunk_b64 {
|
||||
match base64::engine::general_purpose::STANDARD.decode(chunk_b64) {
|
||||
Ok(decoded) => decoded,
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = %err, "gateway failed to decode execution runtime chunk");
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"execution_runtime_stream_chunk_decode_error",
|
||||
format!("failed to decode execution runtime stream chunk: {err}"),
|
||||
502,
|
||||
));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else if let Some(text) = text {
|
||||
text.into_bytes()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
if chunk.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
provider_buffered_body.extend_from_slice(&chunk);
|
||||
let normalized_chunk = if let Some(normalizer) =
|
||||
private_stream_normalizer.as_mut()
|
||||
{
|
||||
match normalizer.push_chunk(&chunk) {
|
||||
Ok(normalized_chunk) => normalized_chunk,
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to normalize execution runtime stream chunk");
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_error",
|
||||
format!("failed to normalize execution runtime stream chunk: {err:?}"),
|
||||
502,
|
||||
));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
chunk
|
||||
};
|
||||
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
|
||||
{
|
||||
match rewriter.push_chunk(&normalized_chunk) {
|
||||
Ok(rewritten_chunk) => rewritten_chunk,
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to rewrite execution runtime stream chunk");
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_error",
|
||||
format!("failed to rewrite execution runtime stream chunk: {err:?}"),
|
||||
502,
|
||||
));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
normalized_chunk
|
||||
};
|
||||
|
||||
if rewritten_chunk.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
buffered_body.extend_from_slice(&rewritten_chunk);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway stream downstream dropped; stopping execution runtime stream forwarding"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
StreamFramePayload::Telemetry {
|
||||
telemetry: frame_telemetry,
|
||||
} => {
|
||||
telemetry = Some(frame_telemetry);
|
||||
}
|
||||
StreamFramePayload::Eof { .. } => {
|
||||
break;
|
||||
}
|
||||
StreamFramePayload::Error { error } => {
|
||||
warn!(trace_id = %trace_id_owned, error = %error.message, "execution runtime stream emitted error frame");
|
||||
terminal_failure = Some(build_stream_failure_from_execution_error(&error));
|
||||
break;
|
||||
}
|
||||
StreamFramePayload::Headers { .. } => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if downstream_dropped {
|
||||
debug!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped local stream flush after downstream disconnect"
|
||||
);
|
||||
} else {
|
||||
if let Some(normalizer) = private_stream_normalizer.as_mut() {
|
||||
match normalizer.finish() {
|
||||
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
|
||||
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
|
||||
{
|
||||
match rewriter.push_chunk(&normalized_chunk) {
|
||||
Ok(rewritten_chunk) => rewritten_chunk,
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to rewrite normalized private stream chunk during flush");
|
||||
terminal_failure.get_or_insert_with(|| {
|
||||
build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_flush_error",
|
||||
format!("failed to rewrite normalized private stream chunk during flush: {err:?}"),
|
||||
502,
|
||||
)
|
||||
});
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
normalized_chunk
|
||||
};
|
||||
if !rewritten_chunk.is_empty() {
|
||||
buffered_body.extend_from_slice(&rewritten_chunk);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway stream downstream dropped while flushing private stream normalization"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to flush private stream normalization");
|
||||
terminal_failure.get_or_insert_with(|| {
|
||||
build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_flush_error",
|
||||
format!("failed to flush private stream normalization: {err:?}"),
|
||||
502,
|
||||
)
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
if !downstream_dropped {
|
||||
if let Some(rewriter) = local_stream_rewriter.as_mut() {
|
||||
match rewriter.finish() {
|
||||
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
|
||||
buffered_body.extend_from_slice(&flushed_chunk);
|
||||
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
|
||||
warn!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway stream downstream dropped while flushing local stream rewrite"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(err) => {
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to flush local stream rewrite");
|
||||
terminal_failure.get_or_insert_with(|| {
|
||||
build_stream_failure_report(
|
||||
"execution_runtime_stream_rewrite_flush_error",
|
||||
format!("failed to flush local stream rewrite: {err:?}"),
|
||||
502,
|
||||
)
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
drop(tx);
|
||||
|
||||
if downstream_dropped {
|
||||
debug!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped stream report because downstream disconnected before completion"
|
||||
);
|
||||
state_for_report
|
||||
.usage_runtime
|
||||
.record_stream_terminal(
|
||||
state_for_report.data.as_ref(),
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&GatewayStreamReportRequest {
|
||||
trace_id: trace_id_owned.clone(),
|
||||
report_kind: report_kind_owned.clone().unwrap_or_default(),
|
||||
report_context: report_context_owned.clone(),
|
||||
status_code: 499,
|
||||
headers: headers_for_report.clone(),
|
||||
provider_body_base64: (!provider_buffered_body.is_empty()).then(|| {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.encode(&provider_buffered_body)
|
||||
}),
|
||||
client_body_base64: (!buffered_body.is_empty()).then(|| {
|
||||
base64::engine::general_purpose::STANDARD.encode(&buffered_body)
|
||||
}),
|
||||
telemetry: telemetry.clone(),
|
||||
},
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
record_local_request_candidate_status(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Cancelled,
|
||||
Some(499),
|
||||
Some("downstream_disconnect".to_string()),
|
||||
Some("client disconnected before stream completion".to_string()),
|
||||
telemetry.as_ref().and_then(|value| value.elapsed_ms),
|
||||
Some(candidate_started_unix_secs_for_report),
|
||||
Some(current_request_candidate_unix_secs()),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(failure) = terminal_failure {
|
||||
submit_midstream_stream_failure(
|
||||
&state_for_report,
|
||||
&trace_id_owned,
|
||||
&plan_for_report,
|
||||
direct_stream_finalize_kind_owned.as_deref(),
|
||||
report_context_owned.as_ref(),
|
||||
&headers_for_report,
|
||||
telemetry.clone(),
|
||||
&provider_buffered_body,
|
||||
candidate_started_unix_secs_for_report,
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
let usage_payload = GatewayStreamReportRequest {
|
||||
trace_id: trace_id_owned.clone(),
|
||||
report_kind: report_kind_owned.clone().unwrap_or_default(),
|
||||
report_context: report_context_owned.clone(),
|
||||
status_code,
|
||||
headers: headers_for_report.clone(),
|
||||
provider_body_base64: (!provider_buffered_body.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)),
|
||||
client_body_base64: (!buffered_body.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
|
||||
telemetry: telemetry.clone(),
|
||||
};
|
||||
state_for_report
|
||||
.usage_runtime
|
||||
.record_stream_terminal(
|
||||
state_for_report.data.as_ref(),
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&usage_payload,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
record_local_request_candidate_status(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Success,
|
||||
Some(status_code),
|
||||
None,
|
||||
None,
|
||||
telemetry.as_ref().and_then(|value| value.elapsed_ms),
|
||||
Some(candidate_started_unix_secs_for_report),
|
||||
Some(current_request_candidate_unix_secs()),
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Some(report_kind) = report_kind_owned {
|
||||
let mut report = usage_payload;
|
||||
report.report_kind = report_kind;
|
||||
if let Err(err) = submit_stream_report(&state_for_report, &trace_id_owned, report).await
|
||||
{
|
||||
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to submit stream execution report");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let body_stream = stream! {
|
||||
for chunk in prefetched_chunks_for_body {
|
||||
yield Ok(chunk);
|
||||
}
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
};
|
||||
|
||||
headers.insert(
|
||||
CONTROL_REQUEST_ID_HEADER.to_string(),
|
||||
request_id.to_string(),
|
||||
);
|
||||
|
||||
if let Some(candidate_id) = candidate_id
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
headers.insert(
|
||||
CONTROL_CANDIDATE_ID_HEADER.to_string(),
|
||||
candidate_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
Body::from_stream(body_stream),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
use serde_json::{Map, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::gateway::execution_runtime::submission::{
|
||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::gateway::request_candidates::{
|
||||
current_unix_secs as current_request_candidate_unix_secs,
|
||||
record_report_request_candidate_status,
|
||||
};
|
||||
use crate::gateway::usage::submit_sync_report;
|
||||
use crate::gateway::{
|
||||
attach_control_metadata_headers, AppState, GatewayControlDecision, GatewayError,
|
||||
GatewaySyncReportRequest,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct StreamFailureReport {
|
||||
pub(super) status_code: u16,
|
||||
pub(super) error_type: String,
|
||||
pub(super) error_message: String,
|
||||
pub(super) body_json: Value,
|
||||
}
|
||||
|
||||
pub(super) fn build_stream_failure_report(
|
||||
error_type: impl Into<String>,
|
||||
error_message: impl Into<String>,
|
||||
status_code: u16,
|
||||
) -> StreamFailureReport {
|
||||
let error_type = error_type.into();
|
||||
let error_message = error_message.into();
|
||||
StreamFailureReport {
|
||||
status_code,
|
||||
body_json: Value::Object(Map::from_iter([(
|
||||
"error".to_string(),
|
||||
Value::Object(Map::from_iter([
|
||||
("type".to_string(), Value::String(error_type.clone())),
|
||||
("message".to_string(), Value::String(error_message.clone())),
|
||||
("code".to_string(), Value::from(status_code)),
|
||||
])),
|
||||
)])),
|
||||
error_type,
|
||||
error_message,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn build_stream_failure_from_execution_error(
|
||||
error: &ExecutionError,
|
||||
) -> StreamFailureReport {
|
||||
let status_code = error.upstream_status.unwrap_or(502);
|
||||
let error_type = serde_json::to_value(&error.kind)
|
||||
.ok()
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "internal".to_string());
|
||||
let phase = serde_json::to_value(&error.phase).unwrap_or(Value::Null);
|
||||
let mut error_object = Map::from_iter([
|
||||
("type".to_string(), Value::String(error_type.clone())),
|
||||
("message".to_string(), Value::String(error.message.clone())),
|
||||
("code".to_string(), Value::from(status_code)),
|
||||
("phase".to_string(), phase),
|
||||
("retryable".to_string(), Value::Bool(error.retryable)),
|
||||
(
|
||||
"failover_recommended".to_string(),
|
||||
Value::Bool(error.failover_recommended),
|
||||
),
|
||||
]);
|
||||
if let Some(upstream_status) = error.upstream_status {
|
||||
error_object.insert("upstream_status".to_string(), Value::from(upstream_status));
|
||||
}
|
||||
|
||||
StreamFailureReport {
|
||||
status_code,
|
||||
error_type,
|
||||
error_message: error.message.trim().to_string(),
|
||||
body_json: Value::Object(Map::from_iter([(
|
||||
"error".to_string(),
|
||||
Value::Object(error_object),
|
||||
)])),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_stream_failure_sync_payload(
|
||||
trace_id: &str,
|
||||
report_kind: String,
|
||||
report_context: Option<Value>,
|
||||
headers: &std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
provider_buffered_body: &[u8],
|
||||
failure: &StreamFailureReport,
|
||||
) -> GatewaySyncReportRequest {
|
||||
let mut response_headers = headers.clone();
|
||||
response_headers.remove("content-encoding");
|
||||
response_headers.remove("content-length");
|
||||
response_headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
|
||||
GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind,
|
||||
report_context,
|
||||
status_code: failure.status_code,
|
||||
headers: response_headers,
|
||||
body_json: Some(failure.body_json.clone()),
|
||||
client_body_json: None,
|
||||
body_base64: (!provider_buffered_body.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(provider_buffered_body)),
|
||||
telemetry,
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_stream_sync_failure(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
failure: &StreamFailureReport,
|
||||
started_at_unix_secs: Option<u64>,
|
||||
) {
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_report_request_candidate_status(
|
||||
state,
|
||||
report_context,
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(failure.status_code),
|
||||
Some(failure.error_type.clone()),
|
||||
Some(failure.error_message.clone()),
|
||||
payload
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms),
|
||||
started_at_unix_secs.or(Some(terminal_unix_secs)),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal helper for prefetch error handling
|
||||
pub(super) async fn handle_prefetch_stream_failure(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<Value>,
|
||||
request_id: &str,
|
||||
candidate_id: Option<&str>,
|
||||
report_kind: &str,
|
||||
headers: &std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
buffered_body: &[u8],
|
||||
failure: StreamFailureReport,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let payload = build_stream_failure_sync_payload(
|
||||
trace_id,
|
||||
report_kind.to_string(),
|
||||
report_context.clone(),
|
||||
headers,
|
||||
telemetry,
|
||||
buffered_body,
|
||||
&failure,
|
||||
);
|
||||
record_stream_sync_failure(
|
||||
state,
|
||||
plan,
|
||||
report_context.as_ref(),
|
||||
&payload,
|
||||
&failure,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response =
|
||||
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
|
||||
Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?))
|
||||
}
|
||||
|
||||
pub(super) async fn submit_midstream_stream_failure(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
plan: &ExecutionPlan,
|
||||
direct_stream_finalize_kind: Option<&str>,
|
||||
report_context: Option<&Value>,
|
||||
headers: &std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
buffered_body: &[u8],
|
||||
started_at_unix_secs: u64,
|
||||
failure: StreamFailureReport,
|
||||
) {
|
||||
let Some(report_kind) =
|
||||
direct_stream_finalize_kind.and_then(resolve_core_error_background_report_kind)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let payload = build_stream_failure_sync_payload(
|
||||
trace_id,
|
||||
report_kind,
|
||||
report_context.cloned(),
|
||||
headers,
|
||||
telemetry,
|
||||
buffered_body,
|
||||
&failure,
|
||||
);
|
||||
record_stream_sync_failure(
|
||||
state,
|
||||
plan,
|
||||
report_context,
|
||||
&payload,
|
||||
&failure,
|
||||
Some(started_at_unix_secs),
|
||||
)
|
||||
.await;
|
||||
if let Err(err) = submit_sync_report(state, trace_id, payload).await {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway failed to submit sync execution report for terminal stream failure"
|
||||
);
|
||||
}
|
||||
}
|
||||
112
apps/aether-gateway/src/execution_runtime/stream_pump.rs
Normal file
112
apps/aether-gateway/src/execution_runtime/stream_pump.rs
Normal file
@@ -0,0 +1,112 @@
|
||||
use std::io::Error as IoError;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionTelemetry, StreamFrame,
|
||||
StreamFramePayload, StreamFrameType,
|
||||
};
|
||||
use async_stream::stream;
|
||||
use axum::body::Bytes;
|
||||
use base64::Engine as _;
|
||||
use futures_util::{Stream, StreamExt};
|
||||
|
||||
use crate::gateway::execution_runtime::ndjson::encode_stream_frame_ndjson;
|
||||
use crate::gateway::execution_runtime::DirectUpstreamStreamExecution;
|
||||
|
||||
pub(crate) fn build_direct_execution_frame_stream(
|
||||
execution: DirectUpstreamStreamExecution,
|
||||
) -> impl Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
||||
stream! {
|
||||
let DirectUpstreamStreamExecution {
|
||||
request_id: _,
|
||||
candidate_id: _,
|
||||
status_code,
|
||||
headers,
|
||||
response,
|
||||
started_at,
|
||||
} = execution;
|
||||
|
||||
let headers_frame = StreamFrame {
|
||||
frame_type: StreamFrameType::Headers,
|
||||
payload: StreamFramePayload::Headers {
|
||||
status_code,
|
||||
headers,
|
||||
},
|
||||
};
|
||||
match encode_stream_frame_ndjson(&headers_frame) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let mut upstream_bytes = 0u64;
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
while let Some(item) = bytes_stream.next().await {
|
||||
match item {
|
||||
Ok(chunk) => {
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
let frame = StreamFrame {
|
||||
frame_type: StreamFrameType::Data,
|
||||
payload: StreamFramePayload::Data {
|
||||
chunk_b64: Some(base64::engine::general_purpose::STANDARD.encode(&chunk)),
|
||||
text: None,
|
||||
},
|
||||
};
|
||||
match encode_stream_frame_ndjson(&frame) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
let frame = StreamFrame {
|
||||
frame_type: StreamFrameType::Error,
|
||||
payload: StreamFramePayload::Error {
|
||||
error: ExecutionError {
|
||||
kind: ExecutionErrorKind::Internal,
|
||||
phase: ExecutionPhase::StreamRead,
|
||||
message: err.to_string(),
|
||||
upstream_status: Some(status_code),
|
||||
retryable: false,
|
||||
failover_recommended: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
match encode_stream_frame_ndjson(&frame) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(encode_err) => {
|
||||
yield Err(encode_err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let telemetry_frame = StreamFrame {
|
||||
frame_type: StreamFrameType::Telemetry,
|
||||
payload: StreamFramePayload::Telemetry {
|
||||
telemetry: ExecutionTelemetry {
|
||||
ttfb_ms: None,
|
||||
elapsed_ms: Some(started_at.elapsed().as_millis() as u64),
|
||||
upstream_bytes: Some(upstream_bytes),
|
||||
},
|
||||
},
|
||||
};
|
||||
match encode_stream_frame_ndjson(&telemetry_frame) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
match encode_stream_frame_ndjson(&StreamFrame::eof()) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => yield Err(err),
|
||||
}
|
||||
}
|
||||
}
|
||||
466
apps/aether-gateway/src/execution_runtime/submission.rs
Normal file
466
apps/aether-gateway/src/execution_runtime/submission.rs
Normal file
@@ -0,0 +1,466 @@
|
||||
use crate::gateway::ai_pipeline::conversion::{
|
||||
build_core_error_body_for_client_format, core_error_background_report_kind,
|
||||
core_error_default_client_api_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
|
||||
};
|
||||
use crate::gateway::ai_pipeline::private_response::normalize_provider_private_response_value as unwrap_local_finalize_response_value;
|
||||
use crate::gateway::usage::spawn_sync_report;
|
||||
use crate::gateway::{
|
||||
build_client_response_from_parts, maybe_compile_sync_finalize_response, AppState,
|
||||
GatewayControlDecision, GatewayError, GatewaySyncReportRequest,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
use tracing::warn;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct LocalSyncErrorDetails {
|
||||
message: String,
|
||||
code: Option<String>,
|
||||
kind: LocalCoreSyncErrorKind,
|
||||
}
|
||||
|
||||
pub(super) fn maybe_build_local_core_error_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if !is_core_error_finalize_kind(payload.report_kind.as_str()) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(body_json) = payload.body_json.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut body_json = body_json.clone();
|
||||
if let Some(report_context) = payload.report_context.as_ref() {
|
||||
if let Some(unwrapped) =
|
||||
unwrap_local_finalize_response_value(body_json.clone(), report_context)?
|
||||
{
|
||||
body_json = unwrapped;
|
||||
}
|
||||
}
|
||||
|
||||
let Some(body_object) = body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !body_object.contains_key("error")
|
||||
&& !body_object
|
||||
.get("type")
|
||||
.and_then(|value| value.as_str())
|
||||
.is_some_and(|value| value == "error")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(response_body_json) = build_best_effort_local_core_error_body(payload, &body_json)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut response_headers = payload.headers.clone();
|
||||
response_headers.remove("content-encoding");
|
||||
response_headers.remove("content-length");
|
||||
response_headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
|
||||
let body_bytes = serde_json::to_vec(&response_body_json)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
resolve_local_sync_error_status_code(payload.status_code, &body_json),
|
||||
&response_headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
|
||||
fn maybe_resolve_local_sync_response_body_json(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
if let Some(client_body_json) = payload.client_body_json.clone() {
|
||||
return Ok(Some(client_body_json));
|
||||
}
|
||||
|
||||
let Some(mut body_json) = payload.body_json.clone() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Some(report_context) = payload.report_context.as_ref() {
|
||||
if let Some(unwrapped) =
|
||||
unwrap_local_finalize_response_value(body_json.clone(), report_context)?
|
||||
{
|
||||
body_json = unwrapped;
|
||||
}
|
||||
}
|
||||
|
||||
if is_core_error_finalize_kind(payload.report_kind.as_str()) {
|
||||
if let Some(converted) = build_best_effort_local_core_error_body(payload, &body_json)? {
|
||||
return Ok(Some(converted));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(body_json))
|
||||
}
|
||||
|
||||
fn build_local_sync_response_from_json(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_json: serde_json::Value,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let status_code = if is_core_error_finalize_kind(payload.report_kind.as_str())
|
||||
|| has_nested_error(&body_json)
|
||||
{
|
||||
resolve_local_sync_error_status_code(payload.status_code, &body_json)
|
||||
} else {
|
||||
payload.status_code
|
||||
};
|
||||
|
||||
let mut response_headers = payload.headers.clone();
|
||||
response_headers.remove("content-encoding");
|
||||
response_headers.remove("content-length");
|
||||
response_headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
|
||||
let body_bytes =
|
||||
serde_json::to_vec(&body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
build_client_response_from_parts(
|
||||
status_code,
|
||||
&response_headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
fn build_local_sync_response_from_bytes(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_bytes: Vec<u8>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let mut response_headers = payload.headers.clone();
|
||||
response_headers.remove("content-length");
|
||||
if body_bytes.is_empty() {
|
||||
response_headers.remove("content-encoding");
|
||||
}
|
||||
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
build_client_response_from_parts(
|
||||
payload.status_code,
|
||||
&response_headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
fn build_local_core_sync_finalize_fallback_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if let Some(body_json) = maybe_resolve_local_sync_response_body_json(payload)? {
|
||||
return build_local_sync_response_from_json(trace_id, decision, payload, body_json);
|
||||
}
|
||||
|
||||
if let Some(body_base64) = payload.body_base64.as_ref() {
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
return build_local_sync_response_from_bytes(trace_id, decision, payload, body_bytes);
|
||||
}
|
||||
|
||||
build_local_sync_response_from_bytes(trace_id, decision, payload, Vec::new())
|
||||
}
|
||||
|
||||
pub(crate) fn build_best_effort_local_core_error_body(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_json: &serde_json::Value,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let default_api_format = core_error_default_client_api_format(payload.report_kind.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let client_api_format = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("client_api_format"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or(default_api_format.as_str())
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let provider_api_format = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(|value| value.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| client_api_format.clone());
|
||||
|
||||
if client_api_format.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
if client_api_format == provider_api_format {
|
||||
return Ok(Some(body_json.clone()));
|
||||
}
|
||||
|
||||
let details = extract_local_sync_error_details(payload.status_code, body_json);
|
||||
Ok(build_core_error_body_for_client_format(
|
||||
&client_api_format,
|
||||
&details.message,
|
||||
details.code.as_deref(),
|
||||
details.kind,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_core_error_background_report_kind(report_kind: &str) -> Option<String> {
|
||||
core_error_background_report_kind(report_kind).map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn resolve_core_success_background_report_kind(report_kind: &str) -> Option<String> {
|
||||
crate::gateway::ai_pipeline::conversion::core_success_background_report_kind(report_kind)
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn resolve_local_sync_error_status_code(status_code: u16, body_json: &serde_json::Value) -> u16 {
|
||||
if (400..600).contains(&status_code) {
|
||||
return status_code;
|
||||
}
|
||||
|
||||
let Some(error_object) = body_json.get("error").and_then(|value| value.as_object()) else {
|
||||
return 400;
|
||||
};
|
||||
|
||||
for key in ["code", "status"] {
|
||||
let Some(value) = error_object.get(key) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(number) = value.as_u64() {
|
||||
if (400..600).contains(&number) {
|
||||
return number as u16;
|
||||
}
|
||||
}
|
||||
if let Some(text) = value.as_str() {
|
||||
if let Ok(number) = text.parse::<u16>() {
|
||||
if (400..600).contains(&number) {
|
||||
return number;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
400
|
||||
}
|
||||
|
||||
fn extract_local_sync_error_details(
|
||||
status_code: u16,
|
||||
body_json: &serde_json::Value,
|
||||
) -> LocalSyncErrorDetails {
|
||||
let resolved_status_code = resolve_local_sync_error_status_code(status_code, body_json);
|
||||
let body_object = body_json.as_object();
|
||||
let error_object = body_object
|
||||
.and_then(|object| object.get("error"))
|
||||
.and_then(|value| value.as_object());
|
||||
|
||||
let message = first_non_empty_error_text(
|
||||
error_object,
|
||||
body_object,
|
||||
&["message", "detail", "reason", "status", "type", "__type"],
|
||||
)
|
||||
.unwrap_or_else(|| format!("HTTP {resolved_status_code}"));
|
||||
let code = first_non_empty_error_text(error_object, body_object, &["code", "status"]);
|
||||
let raw_type = first_non_empty_error_text(error_object, body_object, &["type", "__type"]);
|
||||
let raw_status = first_non_empty_error_text(error_object, body_object, &["status"]);
|
||||
let kind = classify_local_sync_error_kind(
|
||||
resolved_status_code,
|
||||
raw_type.as_deref(),
|
||||
raw_status.as_deref(),
|
||||
code.as_deref(),
|
||||
message.as_str(),
|
||||
);
|
||||
|
||||
LocalSyncErrorDetails {
|
||||
message,
|
||||
code,
|
||||
kind,
|
||||
}
|
||||
}
|
||||
|
||||
fn first_non_empty_error_text(
|
||||
error_object: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
body_object: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
keys: &[&str],
|
||||
) -> Option<String> {
|
||||
for object in [error_object, body_object].into_iter().flatten() {
|
||||
for key in keys {
|
||||
let Some(value) = object.get(*key) else {
|
||||
continue;
|
||||
};
|
||||
match value {
|
||||
serde_json::Value::String(text) if !text.trim().is_empty() => {
|
||||
return Some(text.trim().to_string());
|
||||
}
|
||||
serde_json::Value::Number(number) => return Some(number.to_string()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn classify_local_sync_error_kind(
|
||||
status_code: u16,
|
||||
raw_type: Option<&str>,
|
||||
raw_status: Option<&str>,
|
||||
raw_code: Option<&str>,
|
||||
message: &str,
|
||||
) -> LocalCoreSyncErrorKind {
|
||||
let mut fingerprint = String::new();
|
||||
for segment in [raw_type, raw_status, raw_code, Some(message)] {
|
||||
if let Some(segment) = segment.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
if !fingerprint.is_empty() {
|
||||
fingerprint.push(' ');
|
||||
}
|
||||
fingerprint.push_str(&segment.to_ascii_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
if status_code == 429
|
||||
|| fingerprint.contains("rate_limit")
|
||||
|| fingerprint.contains("rate limited")
|
||||
|| fingerprint.contains("resource_exhausted")
|
||||
|| fingerprint.contains("throttl")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::RateLimit;
|
||||
}
|
||||
if fingerprint.contains("contextlength")
|
||||
|| fingerprint.contains("contentlengthexceeded")
|
||||
|| fingerprint.contains("context window")
|
||||
|| fingerprint.contains("context length")
|
||||
|| fingerprint.contains("max_tokens")
|
||||
|| (fingerprint.contains("context") && fingerprint.contains("token"))
|
||||
{
|
||||
return LocalCoreSyncErrorKind::ContextLengthExceeded;
|
||||
}
|
||||
if status_code == 401
|
||||
|| fingerprint.contains("unauth")
|
||||
|| fingerprint.contains("authentication")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::Authentication;
|
||||
}
|
||||
if status_code == 403 || fingerprint.contains("permission") || fingerprint.contains("forbidden")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::PermissionDenied;
|
||||
}
|
||||
if status_code == 404 || fingerprint.contains("not_found") || fingerprint.contains("not found")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::NotFound;
|
||||
}
|
||||
if status_code == 503 || fingerprint.contains("overload") || fingerprint.contains("unavailable")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::Overloaded;
|
||||
}
|
||||
if (500..600).contains(&status_code) {
|
||||
return LocalCoreSyncErrorKind::ServerError;
|
||||
}
|
||||
LocalCoreSyncErrorKind::InvalidRequest
|
||||
}
|
||||
|
||||
pub(crate) fn strip_utf8_bom_and_ws(mut body: &[u8]) -> &[u8] {
|
||||
loop {
|
||||
while let Some(first) = body.first() {
|
||||
if first.is_ascii_whitespace() {
|
||||
body = &body[1..];
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if body.starts_with(&[0xEF, 0xBB, 0xBF]) {
|
||||
body = &body[3..];
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
pub(crate) fn has_nested_error(value: &serde_json::Value) -> bool {
|
||||
let Some(object) = value.as_object() else {
|
||||
return false;
|
||||
};
|
||||
|
||||
if object.contains_key("error") {
|
||||
return true;
|
||||
}
|
||||
if object
|
||||
.get("type")
|
||||
.and_then(|value| value.as_str())
|
||||
.is_some_and(|value| value == "error")
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
object
|
||||
.get("chunks")
|
||||
.and_then(|value| value.as_array())
|
||||
.is_some_and(|chunks| {
|
||||
chunks.iter().any(|chunk| {
|
||||
chunk.as_object().is_some_and(|chunk_object| {
|
||||
chunk_object.contains_key("error")
|
||||
|| chunk_object
|
||||
.get("type")
|
||||
.and_then(|value| value.as_str())
|
||||
.is_some_and(|value| value == "error")
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn submit_local_core_error_or_sync_finalize(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: GatewaySyncReportRequest,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let response = if let Some(response) =
|
||||
maybe_compile_sync_finalize_response(trace_id, decision, &payload)?
|
||||
{
|
||||
response
|
||||
} else if let Some(response) =
|
||||
maybe_build_local_core_error_response(trace_id, decision, &payload)?
|
||||
{
|
||||
response
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
status_code = payload.status_code,
|
||||
client_api_format = payload.report_context.as_ref().and_then(|value| value.get("client_api_format")).and_then(|value| value.as_str()).unwrap_or(""),
|
||||
provider_api_format = payload.report_context.as_ref().and_then(|value| value.get("provider_api_format")).and_then(|value| value.as_str()).unwrap_or(""),
|
||||
envelope_name = payload.report_context.as_ref().and_then(|value| value.get("envelope_name")).and_then(|value| value.as_str()).unwrap_or(""),
|
||||
needs_conversion = payload.report_context.as_ref().and_then(|value| value.get("needs_conversion")).and_then(|value| value.as_bool()).unwrap_or(false),
|
||||
"gateway local core finalize fell back to raw response body"
|
||||
);
|
||||
build_local_core_sync_finalize_fallback_response(trace_id, decision, &payload)?
|
||||
};
|
||||
|
||||
if let Some(error_report_kind) =
|
||||
resolve_core_error_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
let mut report_payload = payload.clone();
|
||||
report_payload.report_kind = error_report_kind;
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway built local core finalize response without background error report mapping"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
53
apps/aether-gateway/src/execution_runtime/sync.rs
Normal file
53
apps/aether-gateway/src/execution_runtime/sync.rs
Normal file
@@ -0,0 +1,53 @@
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::Response;
|
||||
|
||||
use crate::gateway::{AppState, GatewayControlDecision, GatewayError};
|
||||
|
||||
mod execution;
|
||||
|
||||
pub(crate) use execution::execute_execution_runtime_sync;
|
||||
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use execution::{
|
||||
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
|
||||
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
|
||||
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_execute_via_execution_runtime_sync(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_bytes: &Bytes,
|
||||
trace_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let Some(decision) = decision else {
|
||||
return Ok(None);
|
||||
};
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
let _ = state;
|
||||
if parts.method != http::Method::POST {
|
||||
return Ok(None);
|
||||
}
|
||||
return crate::gateway::ai_pipeline::planner::maybe_execute_sync_local_path(
|
||||
state, parts, body_bytes, trace_id, decision,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
#[cfg(test)]
|
||||
{
|
||||
if state
|
||||
.test_remote_execution_runtime_base_url()
|
||||
.unwrap_or_default()
|
||||
.is_empty()
|
||||
&& parts.method != http::Method::POST
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
crate::gateway::ai_pipeline::planner::maybe_execute_sync_local_path(
|
||||
state, parts, body_bytes, trace_id, decision,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
617
apps/aether-gateway/src/execution_runtime/sync/execution.rs
Normal file
617
apps/aether-gateway/src/execution_runtime/sync/execution.rs
Normal file
@@ -0,0 +1,617 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTelemetry};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::gateway::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
|
||||
#[cfg(test)]
|
||||
use crate::gateway::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
|
||||
use crate::gateway::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
|
||||
use crate::gateway::execution_runtime::transport::DirectSyncExecutionRuntime;
|
||||
use crate::gateway::request_candidates::{
|
||||
current_unix_secs as current_request_candidate_unix_secs,
|
||||
ensure_execution_request_candidate_slot, execution_error_details,
|
||||
record_local_request_candidate_status,
|
||||
};
|
||||
use crate::gateway::scheduler::{
|
||||
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
||||
should_finalize_sync_response, should_retry_next_local_candidate_sync,
|
||||
};
|
||||
use crate::gateway::usage::{spawn_sync_report, submit_sync_report};
|
||||
use crate::gateway::video_tasks::VideoTaskSyncReportMode;
|
||||
use crate::gateway::{
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
maybe_build_sync_finalize_outcome, AppState, GatewayControlDecision, GatewayError,
|
||||
GatewaySyncReportRequest,
|
||||
};
|
||||
|
||||
#[path = "execution/policy.rs"]
|
||||
mod policy;
|
||||
#[path = "execution/response.rs"]
|
||||
mod response;
|
||||
|
||||
use policy::decode_execution_result_body;
|
||||
pub(crate) use response::{
|
||||
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
|
||||
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
|
||||
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
|
||||
};
|
||||
|
||||
struct ImplicitSyncFinalizeOutcome {
|
||||
payload: GatewaySyncReportRequest,
|
||||
outcome: crate::gateway::ai_pipeline::finalize::LocalCoreSyncFinalizeOutcome,
|
||||
}
|
||||
|
||||
async fn record_sync_terminal_usage(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) {
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
enum RemoteSyncFallbackOutcome {
|
||||
Executed(ExecutionResult),
|
||||
ClientResponse(Response<Body>),
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
||||
pub(crate) async fn execute_execution_runtime_sync(
|
||||
state: &AppState,
|
||||
request_path: &str,
|
||||
mut plan: ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
report_kind: Option<String>,
|
||||
mut report_context: Option<serde_json::Value>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
|
||||
let plan_request_id = plan.request_id.as_str();
|
||||
let plan_candidate_id = plan.candidate_id.as_deref();
|
||||
#[cfg(not(test))]
|
||||
let result = {
|
||||
match DirectSyncExecutionRuntime::new()
|
||||
.execute_sync(plan.clone())
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan_request_id,
|
||||
candidate_id = ?plan_candidate_id,
|
||||
error = %err,
|
||||
"gateway in-process sync execution unavailable"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
};
|
||||
#[cfg(test)]
|
||||
let result = {
|
||||
let remote_execution_runtime_base_url = state
|
||||
.test_remote_execution_runtime_base_url()
|
||||
.unwrap_or_default();
|
||||
if remote_execution_runtime_base_url.trim().is_empty() {
|
||||
match DirectSyncExecutionRuntime::new()
|
||||
.execute_sync(plan.clone())
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan_request_id,
|
||||
candidate_id = ?plan_candidate_id,
|
||||
error = %err,
|
||||
"gateway in-process sync execution unavailable"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let remote_outcome = execute_sync_via_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
trace_id,
|
||||
decision,
|
||||
&plan,
|
||||
plan_request_id,
|
||||
plan_candidate_id,
|
||||
report_context.as_ref(),
|
||||
)
|
||||
.await?;
|
||||
match remote_outcome {
|
||||
RemoteSyncFallbackOutcome::Executed(result) => result,
|
||||
RemoteSyncFallbackOutcome::ClientResponse(response) => return Ok(Some(response)),
|
||||
RemoteSyncFallbackOutcome::Unavailable => return Ok(None),
|
||||
}
|
||||
}
|
||||
};
|
||||
let result_body_json = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref());
|
||||
let (result_error_type, result_error_message) =
|
||||
execution_error_details(result.error.as_ref(), result_body_json);
|
||||
let result_latency_ms = result
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms);
|
||||
if should_retry_next_local_candidate_sync(plan_kind, report_context.as_ref(), &result) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(result.status_code),
|
||||
result_error_type.clone(),
|
||||
result_error_message.clone(),
|
||||
result_latency_ms,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan_request_id,
|
||||
status_code = result.status_code,
|
||||
"gateway local sync decision retrying next candidate after retryable execution runtime result"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
let request_id = (!result.request_id.trim().is_empty())
|
||||
.then_some(result.request_id.as_str())
|
||||
.or(Some(plan_request_id));
|
||||
let candidate_id = result.candidate_id.as_deref().or(plan_candidate_id);
|
||||
let mut headers = result.headers.clone();
|
||||
let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
|
||||
let has_body_bytes = body_base64.is_some();
|
||||
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
|
||||
let mapped_error_finalize_kind =
|
||||
resolve_core_sync_error_finalize_report_kind(plan_kind, &result, body_json.as_ref());
|
||||
let implicit_finalize = if explicit_finalize || mapped_error_finalize_kind.is_some() {
|
||||
None
|
||||
} else {
|
||||
maybe_build_implicit_sync_finalize_outcome(
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
report_context.clone(),
|
||||
result.status_code,
|
||||
headers.clone(),
|
||||
body_json.clone(),
|
||||
body_base64.clone(),
|
||||
result.telemetry.clone(),
|
||||
)?
|
||||
};
|
||||
let finalize_report_kind = if explicit_finalize {
|
||||
report_kind.clone()
|
||||
} else if let Some(implicit_finalize) = implicit_finalize.as_ref() {
|
||||
Some(implicit_finalize.payload.report_kind.clone())
|
||||
} else {
|
||||
mapped_error_finalize_kind.clone()
|
||||
};
|
||||
|
||||
if should_fallback_to_control_sync(
|
||||
plan_kind,
|
||||
&result,
|
||||
body_json.as_ref(),
|
||||
has_body_bytes,
|
||||
explicit_finalize || implicit_finalize.is_some(),
|
||||
mapped_error_finalize_kind.is_some(),
|
||||
) {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(result.status_code),
|
||||
result_error_type.clone(),
|
||||
result_error_message.clone(),
|
||||
result_latency_ms,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
state
|
||||
.usage_runtime
|
||||
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
|
||||
.await;
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
if result.status_code >= 400 {
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed
|
||||
} else {
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Success
|
||||
},
|
||||
Some(result.status_code),
|
||||
result_error_type.clone(),
|
||||
result_error_message.clone(),
|
||||
result_latency_ms,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
|
||||
let base_usage_payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind: finalize_report_kind
|
||||
.clone()
|
||||
.or_else(|| report_kind.clone())
|
||||
.unwrap_or_default(),
|
||||
report_context: report_context.clone(),
|
||||
status_code: result.status_code,
|
||||
headers: headers.clone(),
|
||||
body_json: body_json.clone(),
|
||||
client_body_json: None,
|
||||
body_base64: body_base64.clone(),
|
||||
telemetry: result.telemetry.clone(),
|
||||
};
|
||||
|
||||
if let Some(finalize_report_kind) = finalize_report_kind {
|
||||
if let Some(implicit_finalize) = implicit_finalize {
|
||||
let usage_payload = implicit_finalize
|
||||
.outcome
|
||||
.background_report
|
||||
.as_ref()
|
||||
.unwrap_or(&implicit_finalize.payload);
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), usage_payload).await;
|
||||
if let Some(report_payload) = implicit_finalize.outcome.background_report {
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %implicit_finalize.payload.report_kind,
|
||||
"gateway implicit local core finalize produced response without background success report mapping"
|
||||
);
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
implicit_finalize.outcome.response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind: finalize_report_kind,
|
||||
report_context,
|
||||
status_code: result.status_code,
|
||||
headers: headers.clone(),
|
||||
body_json: body_json.clone(),
|
||||
client_body_json: None,
|
||||
body_base64: body_base64.clone(),
|
||||
telemetry: result.telemetry.clone(),
|
||||
};
|
||||
if let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? {
|
||||
let usage_payload = outcome.background_report.as_ref().unwrap_or(&payload);
|
||||
record_sync_terminal_usage(
|
||||
state,
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
usage_payload,
|
||||
)
|
||||
.await;
|
||||
if let Some(report_payload) = outcome.background_report {
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway local core finalize produced response without background success report mapping"
|
||||
);
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
outcome.response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(outcome) = maybe_build_local_video_success_outcome(
|
||||
trace_id,
|
||||
decision,
|
||||
&payload,
|
||||
&state.video_tasks,
|
||||
&plan,
|
||||
)? {
|
||||
record_sync_terminal_usage(
|
||||
state,
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&outcome.report_payload,
|
||||
)
|
||||
.await;
|
||||
if let Some(snapshot) = outcome.local_task_snapshot.clone() {
|
||||
state.video_tasks.record_snapshot(snapshot.clone());
|
||||
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
||||
}
|
||||
match outcome.report_mode {
|
||||
VideoTaskSyncReportMode::InlineSync => {
|
||||
submit_sync_report(state, trace_id, outcome.report_payload).await?;
|
||||
}
|
||||
VideoTaskSyncReportMode::Background => {
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), outcome.report_payload);
|
||||
}
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
outcome.response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(response) =
|
||||
maybe_build_local_sync_finalize_response(trace_id, decision, &payload)?
|
||||
{
|
||||
let usage_payload = if let Some(success_report_kind) =
|
||||
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
let mut report_payload = payload.clone();
|
||||
report_payload.report_kind = success_report_kind;
|
||||
report_payload
|
||||
} else {
|
||||
payload.clone()
|
||||
};
|
||||
record_sync_terminal_usage(
|
||||
state,
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
state
|
||||
.video_tasks
|
||||
.apply_finalize_mutation(request_path, payload.report_kind.as_str());
|
||||
if let Some(snapshot) = state
|
||||
.video_tasks
|
||||
.snapshot_for_route(decision.route_family.as_deref(), request_path)
|
||||
{
|
||||
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
||||
}
|
||||
if let Some(success_report_kind) =
|
||||
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
let mut report_payload = usage_payload;
|
||||
report_payload.report_kind = success_report_kind;
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway local video finalize produced response without background success report mapping"
|
||||
);
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(response) =
|
||||
maybe_build_local_video_error_response(trace_id, decision, &payload)?
|
||||
{
|
||||
let usage_payload = if let Some(error_report_kind) =
|
||||
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
let mut report_payload = payload.clone();
|
||||
report_payload.report_kind = error_report_kind;
|
||||
report_payload
|
||||
} else {
|
||||
payload.clone()
|
||||
};
|
||||
record_sync_terminal_usage(
|
||||
state,
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
if let Some(error_report_kind) =
|
||||
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
let mut report_payload = usage_payload;
|
||||
report_payload.report_kind = error_report_kind;
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway local video finalize produced response without background error report mapping"
|
||||
);
|
||||
}
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload).await;
|
||||
let response =
|
||||
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &base_usage_payload).await;
|
||||
if let Some(report_kind) = report_kind {
|
||||
let report = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind,
|
||||
report_context,
|
||||
status_code: result.status_code,
|
||||
headers: headers.clone(),
|
||||
body_json: body_json.clone(),
|
||||
client_body_json: None,
|
||||
body_base64: body_base64.clone(),
|
||||
telemetry: result.telemetry.clone(),
|
||||
};
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report);
|
||||
}
|
||||
|
||||
let request_id_header: Option<&str> = request_id
|
||||
.map(str::trim)
|
||||
.filter(|value: &&str| !value.is_empty());
|
||||
if let Some(request_id) = request_id_header {
|
||||
headers.insert(
|
||||
CONTROL_REQUEST_ID_HEADER.to_string(),
|
||||
request_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
let candidate_id_header: Option<&str> = candidate_id
|
||||
.map(str::trim)
|
||||
.filter(|value: &&str| !value.is_empty());
|
||||
if let Some(candidate_id) = candidate_id_header {
|
||||
headers.insert(
|
||||
CONTROL_CANDIDATE_ID_HEADER.to_string(),
|
||||
candidate_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
result.status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
|
||||
fn resolve_implicit_sync_finalize_report_kind(plan_kind: &str) -> Option<&'static str> {
|
||||
match plan_kind {
|
||||
"openai_chat_sync" => Some("openai_chat_sync_finalize"),
|
||||
"claude_chat_sync" => Some("claude_chat_sync_finalize"),
|
||||
"gemini_chat_sync" => Some("gemini_chat_sync_finalize"),
|
||||
"openai_cli_sync" => Some("openai_cli_sync_finalize"),
|
||||
"openai_compact_sync" => Some("openai_compact_sync_finalize"),
|
||||
"claude_cli_sync" => Some("claude_cli_sync_finalize"),
|
||||
"gemini_cli_sync" => Some("gemini_cli_sync_finalize"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // mirrors sync execution context
|
||||
fn maybe_build_implicit_sync_finalize_outcome(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
report_context: Option<serde_json::Value>,
|
||||
status_code: u16,
|
||||
headers: BTreeMap<String, String>,
|
||||
body_json: Option<serde_json::Value>,
|
||||
body_base64: Option<String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
) -> Result<Option<ImplicitSyncFinalizeOutcome>, GatewayError> {
|
||||
if status_code >= 400 || body_json.is_some() || body_base64.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(report_kind) = resolve_implicit_sync_finalize_report_kind(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context,
|
||||
status_code,
|
||||
headers,
|
||||
body_json,
|
||||
client_body_json: None,
|
||||
body_base64,
|
||||
telemetry,
|
||||
};
|
||||
let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some(ImplicitSyncFinalizeOutcome { payload, outcome }))
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal helper mirroring execute path context
|
||||
#[cfg(test)]
|
||||
async fn execute_sync_via_remote_execution_runtime(
|
||||
state: &AppState,
|
||||
remote_execution_runtime_base_url: &str,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan: &ExecutionPlan,
|
||||
plan_request_id: &str,
|
||||
plan_candidate_id: Option<&str>,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
|
||||
let response = match post_sync_plan_to_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
Some(trace_id),
|
||||
plan,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
request_id = %plan_request_id,
|
||||
candidate_id = ?plan_candidate_id,
|
||||
error = ?err,
|
||||
"gateway remote execution runtime sync unavailable"
|
||||
);
|
||||
return Ok(RemoteSyncFallbackOutcome::Unavailable);
|
||||
}
|
||||
};
|
||||
|
||||
if response.status() != http::StatusCode::OK {
|
||||
let terminal_unix_secs = current_request_candidate_unix_secs();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
plan,
|
||||
report_context,
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed,
|
||||
Some(response.status().as_u16()),
|
||||
Some("execution_runtime_http_error".to_string()),
|
||||
Some(format!(
|
||||
"execution runtime returned HTTP {}",
|
||||
response.status()
|
||||
)),
|
||||
None,
|
||||
Some(terminal_unix_secs),
|
||||
Some(terminal_unix_secs),
|
||||
)
|
||||
.await;
|
||||
return Ok(RemoteSyncFallbackOutcome::ClientResponse(
|
||||
attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(plan_request_id),
|
||||
plan_candidate_id,
|
||||
)?,
|
||||
));
|
||||
}
|
||||
|
||||
response
|
||||
.json()
|
||||
.await
|
||||
.map(RemoteSyncFallbackOutcome::Executed)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ExecutionResult;
|
||||
use base64::Engine as _;
|
||||
|
||||
use crate::gateway::GatewayError;
|
||||
|
||||
type DecodedBody = (Vec<u8>, Option<serde_json::Value>, Option<String>);
|
||||
|
||||
pub(super) fn decode_execution_result_body(
|
||||
result: &ExecutionResult,
|
||||
headers: &mut BTreeMap<String, String>,
|
||||
) -> Result<DecodedBody, GatewayError> {
|
||||
let Some(body) = result.body.as_ref() else {
|
||||
return Ok((Vec::new(), None, None));
|
||||
};
|
||||
|
||||
if let Some(json_body) = body.json_body.clone() {
|
||||
headers
|
||||
.entry("content-type".to_string())
|
||||
.or_insert_with(|| "application/json".to_string());
|
||||
let bytes = serde_json::to_vec(&json_body)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
headers.insert("content-length".to_string(), bytes.len().to_string());
|
||||
return Ok((bytes, Some(json_body), None));
|
||||
}
|
||||
|
||||
if let Some(body_bytes_b64) = body.body_bytes_b64.clone() {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(&body_bytes_b64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
return Ok((bytes, None, Some(body_bytes_b64)));
|
||||
}
|
||||
|
||||
Ok((Vec::new(), None, None))
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::video_tasks::{LocalVideoTaskSnapshot, VideoTaskSyncReportMode};
|
||||
use crate::gateway::VideoTaskService;
|
||||
use crate::gateway::{
|
||||
build_client_response_from_parts, GatewayControlDecision, GatewayError,
|
||||
GatewaySyncReportRequest,
|
||||
};
|
||||
|
||||
pub(crate) struct LocalVideoSyncSuccessOutcome {
|
||||
pub(crate) response: Response<Body>,
|
||||
pub(crate) report_payload: GatewaySyncReportRequest,
|
||||
pub(crate) report_mode: VideoTaskSyncReportMode,
|
||||
pub(crate) local_task_snapshot: Option<LocalVideoTaskSnapshot>,
|
||||
}
|
||||
|
||||
fn cloned_report_context_object(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> serde_json::Map<String, serde_json::Value> {
|
||||
payload
|
||||
.report_context
|
||||
.clone()
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn build_local_video_success_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let body_bytes =
|
||||
serde_json::to_vec(body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
build_client_response_from_parts(
|
||||
http::StatusCode::OK.as_u16(),
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_video_success_outcome(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
video_tasks: &VideoTaskService,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<Option<LocalVideoSyncSuccessOutcome>, GatewayError> {
|
||||
if payload.status_code >= 400 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let provider_body = match payload
|
||||
.body_json
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
Some(value) => value,
|
||||
None => return Ok(None),
|
||||
};
|
||||
let mut report_context = cloned_report_context_object(payload);
|
||||
let Some(plan) = video_tasks.prepare_sync_success(
|
||||
payload.report_kind.as_str(),
|
||||
provider_body,
|
||||
&report_context,
|
||||
plan,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
plan.apply_to_report_context(&mut report_context);
|
||||
let client_body_json = plan.client_body_json();
|
||||
|
||||
let response = build_local_video_success_response(trace_id, decision, &client_body_json)?;
|
||||
let report_payload = GatewaySyncReportRequest {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: plan.success_report_kind().to_string(),
|
||||
report_context: Some(serde_json::Value::Object(report_context)),
|
||||
status_code: payload.status_code,
|
||||
headers: payload.headers.clone(),
|
||||
body_json: payload.body_json.clone(),
|
||||
client_body_json: Some(client_body_json),
|
||||
body_base64: None,
|
||||
telemetry: payload.telemetry.clone(),
|
||||
};
|
||||
|
||||
Ok(Some(LocalVideoSyncSuccessOutcome {
|
||||
response,
|
||||
report_payload,
|
||||
report_mode: plan.report_mode(),
|
||||
local_task_snapshot: matches!(plan.report_mode(), VideoTaskSyncReportMode::Background)
|
||||
.then(|| plan.to_snapshot()),
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let (status_code, body_json) = match payload.report_kind.as_str() {
|
||||
"openai_video_delete_sync_finalize" => {
|
||||
if payload.status_code >= 400 && payload.status_code != 404 {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(task_id) = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("task_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
(
|
||||
http::StatusCode::OK,
|
||||
json!({
|
||||
"id": task_id,
|
||||
"object": "video",
|
||||
"deleted": true,
|
||||
}),
|
||||
)
|
||||
}
|
||||
"openai_video_cancel_sync_finalize" | "gemini_video_cancel_sync_finalize" => {
|
||||
if payload.status_code >= 400 {
|
||||
return Ok(None);
|
||||
}
|
||||
(http::StatusCode::OK, json!({}))
|
||||
}
|
||||
_ => return Ok(None),
|
||||
};
|
||||
|
||||
let body_bytes =
|
||||
serde_json::to_vec(&body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
status_code.as_u16(),
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
|
||||
pub(crate) fn maybe_build_local_video_error_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if !matches!(
|
||||
payload.report_kind.as_str(),
|
||||
"openai_video_create_sync_finalize"
|
||||
| "openai_video_remix_sync_finalize"
|
||||
| "gemini_video_create_sync_finalize"
|
||||
| "openai_video_delete_sync_finalize"
|
||||
| "openai_video_cancel_sync_finalize"
|
||||
| "gemini_video_cancel_sync_finalize"
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if payload.status_code < 400 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_body = payload.body_json.clone().unwrap_or_else(|| json!({}));
|
||||
let body_bytes = serde_json::to_vec(&response_body)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
|
||||
let mut response_headers = payload.headers.clone();
|
||||
response_headers.remove("content-encoding");
|
||||
response_headers.remove("content-length");
|
||||
response_headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
payload.status_code,
|
||||
&response_headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_local_sync_success_background_report_kind(
|
||||
report_kind: &str,
|
||||
) -> Option<String> {
|
||||
let mapped = match report_kind {
|
||||
"openai_video_delete_sync_finalize" => "openai_video_delete_sync_success",
|
||||
"openai_video_cancel_sync_finalize" => "openai_video_cancel_sync_success",
|
||||
"gemini_video_cancel_sync_finalize" => "gemini_video_cancel_sync_success",
|
||||
_ => return None,
|
||||
};
|
||||
Some(mapped.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_local_sync_error_background_report_kind(report_kind: &str) -> Option<String> {
|
||||
let mapped = match report_kind {
|
||||
"openai_video_create_sync_finalize" => "openai_video_create_sync_error",
|
||||
"openai_video_remix_sync_finalize" => "openai_video_remix_sync_error",
|
||||
"gemini_video_create_sync_finalize" => "gemini_video_create_sync_error",
|
||||
"openai_video_delete_sync_finalize" => "openai_video_delete_sync_error",
|
||||
"openai_video_cancel_sync_finalize" => "openai_video_cancel_sync_error",
|
||||
"gemini_video_cancel_sync_finalize" => "gemini_video_cancel_sync_error",
|
||||
_ => return None,
|
||||
};
|
||||
Some(mapped.to_string())
|
||||
}
|
||||
557
apps/aether-gateway/src/execution_runtime/tests.rs
Normal file
557
apps/aether-gateway/src/execution_runtime/tests.rs
Normal file
@@ -0,0 +1,557 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use axum::http::Request;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::execution_runtime::submission::{
|
||||
build_best_effort_local_core_error_body, resolve_core_error_background_report_kind,
|
||||
resolve_core_success_background_report_kind,
|
||||
};
|
||||
use crate::gateway::ai_pipeline::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_openai_cli_stream_plan_from_decision, build_openai_cli_sync_plan_from_decision,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
build_standard_sync_plan_from_decision,
|
||||
};
|
||||
use crate::gateway::{
|
||||
resolve_local_sync_error_background_report_kind,
|
||||
resolve_local_sync_success_background_report_kind, GatewayControlSyncDecisionResponse,
|
||||
GatewaySyncReportRequest,
|
||||
};
|
||||
use crate::gateway::{should_bypass_intent_decision, should_bypass_intent_plan};
|
||||
|
||||
fn test_parts() -> http::request::Parts {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/chat/completions")
|
||||
.body(())
|
||||
.expect("request");
|
||||
let (parts, _) = request.into_parts();
|
||||
parts
|
||||
}
|
||||
|
||||
fn missing_exact_provider_request_payload(
|
||||
decision_kind: &str,
|
||||
) -> GatewayControlSyncDecisionResponse {
|
||||
GatewayControlSyncDecisionResponse {
|
||||
action: "execution_runtime".to_string(),
|
||||
decision_kind: Some(decision_kind.to_string()),
|
||||
execution_strategy: Some("local_same_format".to_string()),
|
||||
conversion_mode: Some("none".to_string()),
|
||||
request_id: Some("req_123".to_string()),
|
||||
candidate_id: Some("cand_123".to_string()),
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: Some("provider_id".to_string()),
|
||||
endpoint_id: Some("endpoint_id".to_string()),
|
||||
key_id: Some("key_id".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/messages".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer test".to_string()),
|
||||
provider_api_format: Some("anthropic".to_string()),
|
||||
client_api_format: Some("openai".to_string()),
|
||||
provider_contract: Some("anthropic".to_string()),
|
||||
client_contract: Some("openai".to_string()),
|
||||
model_name: Some("model".to_string()),
|
||||
mapped_model: Some("model".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: Default::default(),
|
||||
provider_request_headers: Default::default(),
|
||||
provider_request_body: None,
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn core_finalize_payload(
|
||||
report_kind: &str,
|
||||
client_api_format: &str,
|
||||
provider_api_format: &str,
|
||||
status_code: u16,
|
||||
body_json: serde_json::Value,
|
||||
) -> GatewaySyncReportRequest {
|
||||
GatewaySyncReportRequest {
|
||||
trace_id: "trace-core-error-123".to_string(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context: Some(json!({
|
||||
"client_api_format": client_api_format,
|
||||
"provider_api_format": provider_api_format,
|
||||
})),
|
||||
status_code,
|
||||
headers: Default::default(),
|
||||
body_json: Some(body_json),
|
||||
client_body_json: None,
|
||||
body_base64: None,
|
||||
telemetry: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_core_error_background_report_kind_maps_all_core_finalize_kinds() {
|
||||
let cases = [
|
||||
("openai_chat_sync_finalize", Some("openai_chat_sync_error")),
|
||||
("claude_chat_sync_finalize", Some("claude_chat_sync_error")),
|
||||
("gemini_chat_sync_finalize", Some("gemini_chat_sync_error")),
|
||||
("openai_cli_sync_finalize", Some("openai_cli_sync_error")),
|
||||
(
|
||||
"openai_compact_sync_finalize",
|
||||
Some("openai_compact_sync_error"),
|
||||
),
|
||||
("claude_cli_sync_finalize", Some("claude_cli_sync_error")),
|
||||
("gemini_cli_sync_finalize", Some("gemini_cli_sync_error")),
|
||||
("openai_video_create_sync_finalize", None),
|
||||
("gemini_video_cancel_sync_finalize", None),
|
||||
("unknown_finalize_kind", None),
|
||||
];
|
||||
|
||||
for (report_kind, expected) in cases {
|
||||
assert_eq!(
|
||||
resolve_core_error_background_report_kind(report_kind),
|
||||
expected.map(str::to_string),
|
||||
"unexpected mapping for {report_kind}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_core_success_background_report_kind_maps_all_core_finalize_kinds() {
|
||||
let cases = [
|
||||
(
|
||||
"openai_chat_sync_finalize",
|
||||
Some("openai_chat_sync_success"),
|
||||
),
|
||||
(
|
||||
"claude_chat_sync_finalize",
|
||||
Some("claude_chat_sync_success"),
|
||||
),
|
||||
(
|
||||
"gemini_chat_sync_finalize",
|
||||
Some("gemini_chat_sync_success"),
|
||||
),
|
||||
("openai_cli_sync_finalize", Some("openai_cli_sync_success")),
|
||||
(
|
||||
"openai_compact_sync_finalize",
|
||||
Some("openai_cli_sync_success"),
|
||||
),
|
||||
("claude_cli_sync_finalize", Some("claude_cli_sync_success")),
|
||||
("gemini_cli_sync_finalize", Some("gemini_cli_sync_success")),
|
||||
("openai_video_create_sync_finalize", None),
|
||||
("gemini_video_cancel_sync_finalize", None),
|
||||
("unknown_finalize_kind", None),
|
||||
];
|
||||
|
||||
for (report_kind, expected) in cases {
|
||||
assert_eq!(
|
||||
resolve_core_success_background_report_kind(report_kind),
|
||||
expected.map(str::to_string),
|
||||
"unexpected mapping for {report_kind}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_best_effort_local_core_error_body_converts_gemini_chat_error_to_openai_chat() {
|
||||
let payload = core_finalize_payload(
|
||||
"openai_chat_sync_finalize",
|
||||
"openai:chat",
|
||||
"gemini:chat",
|
||||
429,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "rate limited",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"code": 429
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let converted = build_best_effort_local_core_error_body(
|
||||
&payload,
|
||||
payload.body_json.as_ref().expect("body_json should exist"),
|
||||
)
|
||||
.expect("conversion should not error")
|
||||
.expect("conversion should produce a client error body");
|
||||
|
||||
assert_eq!(
|
||||
converted,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "rate limited",
|
||||
"type": "rate_limit_error",
|
||||
"code": "429"
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_best_effort_local_core_error_body_converts_claude_cli_error_to_openai_cli() {
|
||||
let payload = core_finalize_payload(
|
||||
"openai_cli_sync_finalize",
|
||||
"openai:cli",
|
||||
"claude:cli",
|
||||
401,
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "authentication_error",
|
||||
"message": "invalid auth token"
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let converted = build_best_effort_local_core_error_body(
|
||||
&payload,
|
||||
payload.body_json.as_ref().expect("body_json should exist"),
|
||||
)
|
||||
.expect("conversion should not error")
|
||||
.expect("conversion should produce a client error body");
|
||||
|
||||
assert_eq!(
|
||||
converted,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "invalid auth token",
|
||||
"type": "authentication_error"
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_local_sync_success_background_report_kind_maps_video_finalize_kinds() {
|
||||
let cases = [
|
||||
(
|
||||
"openai_video_delete_sync_finalize",
|
||||
Some("openai_video_delete_sync_success"),
|
||||
),
|
||||
(
|
||||
"openai_video_cancel_sync_finalize",
|
||||
Some("openai_video_cancel_sync_success"),
|
||||
),
|
||||
(
|
||||
"gemini_video_cancel_sync_finalize",
|
||||
Some("gemini_video_cancel_sync_success"),
|
||||
),
|
||||
("openai_video_create_sync_finalize", None),
|
||||
("unknown_finalize_kind", None),
|
||||
];
|
||||
|
||||
for (report_kind, expected) in cases {
|
||||
assert_eq!(
|
||||
resolve_local_sync_success_background_report_kind(report_kind),
|
||||
expected.map(str::to_string),
|
||||
"unexpected mapping for {report_kind}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_local_sync_error_background_report_kind_maps_video_finalize_kinds() {
|
||||
let cases = [
|
||||
(
|
||||
"openai_video_create_sync_finalize",
|
||||
Some("openai_video_create_sync_error"),
|
||||
),
|
||||
(
|
||||
"openai_video_remix_sync_finalize",
|
||||
Some("openai_video_remix_sync_error"),
|
||||
),
|
||||
(
|
||||
"gemini_video_create_sync_finalize",
|
||||
Some("gemini_video_create_sync_error"),
|
||||
),
|
||||
(
|
||||
"openai_video_delete_sync_finalize",
|
||||
Some("openai_video_delete_sync_error"),
|
||||
),
|
||||
(
|
||||
"openai_video_cancel_sync_finalize",
|
||||
Some("openai_video_cancel_sync_error"),
|
||||
),
|
||||
(
|
||||
"gemini_video_cancel_sync_finalize",
|
||||
Some("gemini_video_cancel_sync_error"),
|
||||
),
|
||||
("unknown_finalize_kind", None),
|
||||
];
|
||||
|
||||
for (report_kind, expected) in cases {
|
||||
assert_eq!(
|
||||
resolve_local_sync_error_background_report_kind(report_kind),
|
||||
expected.map(str::to_string),
|
||||
"unexpected mapping for {report_kind}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generic_decision_builders_require_exact_provider_request() {
|
||||
let parts = test_parts();
|
||||
let body_json = json!({"messages":[{"role":"user","content":"hi"}]});
|
||||
|
||||
assert!(build_openai_cli_stream_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("openai_cli_stream"),
|
||||
false,
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_openai_cli_stream_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("openai_compact_stream"),
|
||||
true,
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_standard_stream_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("claude_chat_stream"),
|
||||
false,
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_gemini_stream_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("gemini_chat_stream"),
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_openai_cli_sync_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("openai_cli_sync"),
|
||||
false,
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_standard_sync_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("claude_chat_sync"),
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
assert!(build_gemini_sync_plan_from_decision(
|
||||
&parts,
|
||||
&body_json,
|
||||
missing_exact_provider_request_payload("gemini_chat_sync"),
|
||||
)
|
||||
.expect("builder should not error")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passthrough_sync_plan_uses_provider_request_method_override() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/videos/task-123/cancel")
|
||||
.body(())
|
||||
.expect("request");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let mut payload = missing_exact_provider_request_payload("openai_video_cancel_sync");
|
||||
payload.provider_name = Some("openai".to_string());
|
||||
payload.provider_api_format = Some("openai:video".to_string());
|
||||
payload.client_api_format = Some("openai:video".to_string());
|
||||
payload.model_name = Some("sora-2".to_string());
|
||||
payload.upstream_url = Some("https://api.openai.example/v1/videos/ext-123".to_string());
|
||||
payload.provider_request_method = Some("DELETE".to_string());
|
||||
payload.provider_request_headers = [(
|
||||
"authorization".to_string(),
|
||||
"Bearer upstream-key".to_string(),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let plan_and_report = build_passthrough_sync_plan_from_decision(&parts, payload)
|
||||
.expect("builder should not error")
|
||||
.expect("plan should be built");
|
||||
|
||||
assert_eq!(plan_and_report.plan.method, "DELETE");
|
||||
assert_eq!(
|
||||
plan_and_report.plan.url,
|
||||
"https://api.openai.example/v1/videos/ext-123"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_cli_sync_plan_injects_auth_header_when_exact_headers_omit_it() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.body(())
|
||||
.expect("request");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let mut payload = missing_exact_provider_request_payload("openai_cli_sync");
|
||||
payload.provider_name = Some("openai".to_string());
|
||||
payload.provider_api_format = Some("openai:cli".to_string());
|
||||
payload.client_api_format = Some("openai:cli".to_string());
|
||||
payload.model_name = Some("gpt-5".to_string());
|
||||
payload.upstream_url = Some("https://chatgpt.com/backend-api/codex/responses".to_string());
|
||||
payload.provider_request_headers =
|
||||
[("content-type".to_string(), "application/json".to_string())]
|
||||
.into_iter()
|
||||
.collect();
|
||||
payload.provider_request_body = Some(json!({"model":"gpt-5"}));
|
||||
|
||||
let plan_and_report =
|
||||
build_openai_cli_sync_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("builder should not error")
|
||||
.expect("plan should be built");
|
||||
|
||||
assert_eq!(
|
||||
plan_and_report
|
||||
.plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.map(String::as_str),
|
||||
Some("Bearer test")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passthrough_sync_plan_uses_raw_body_bytes_when_decision_provides_base64_body() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/upload/v1beta/files?uploadType=resumable")
|
||||
.body(())
|
||||
.expect("request");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let mut payload = missing_exact_provider_request_payload("gemini_files_upload");
|
||||
payload.provider_name = Some("gemini".to_string());
|
||||
payload.provider_api_format = Some("gemini:files".to_string());
|
||||
payload.client_api_format = Some("gemini:files".to_string());
|
||||
payload.provider_request_method = Some("POST".to_string());
|
||||
payload.upstream_url = Some(
|
||||
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable"
|
||||
.to_string(),
|
||||
);
|
||||
payload.provider_request_headers = [
|
||||
(
|
||||
"content-type".to_string(),
|
||||
"application/octet-stream".to_string(),
|
||||
),
|
||||
("x-goog-api-key".to_string(), "upstream-key".to_string()),
|
||||
]
|
||||
.into_iter()
|
||||
.collect();
|
||||
payload.provider_request_body_base64 = Some("dXBsb2FkLWJ5dGVz".to_string());
|
||||
|
||||
let plan_and_report = build_passthrough_sync_plan_from_decision(&parts, payload)
|
||||
.expect("builder should not error")
|
||||
.expect("plan should be built");
|
||||
|
||||
assert_eq!(plan_and_report.plan.method, "POST");
|
||||
assert_eq!(
|
||||
plan_and_report.plan.url,
|
||||
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable"
|
||||
);
|
||||
assert_eq!(
|
||||
plan_and_report.plan.body.body_bytes_b64.as_deref(),
|
||||
Some("dXBsb2FkLWJ5dGVz")
|
||||
);
|
||||
assert_eq!(plan_and_report.plan.body.json_body, None);
|
||||
assert_eq!(
|
||||
plan_and_report
|
||||
.plan
|
||||
.headers
|
||||
.get("content-type")
|
||||
.map(String::as_str),
|
||||
Some("application/octet-stream")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_cli_stream_plan_injects_auth_header_when_exact_headers_omit_it() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.body(())
|
||||
.expect("request");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let mut payload = missing_exact_provider_request_payload("openai_cli_stream");
|
||||
payload.provider_name = Some("openai".to_string());
|
||||
payload.provider_api_format = Some("openai:cli".to_string());
|
||||
payload.client_api_format = Some("openai:cli".to_string());
|
||||
payload.model_name = Some("gpt-5".to_string());
|
||||
payload.upstream_url = Some("https://chatgpt.com/backend-api/codex/responses".to_string());
|
||||
payload.provider_request_headers =
|
||||
[("content-type".to_string(), "application/json".to_string())]
|
||||
.into_iter()
|
||||
.collect();
|
||||
payload.provider_request_body = Some(json!({"model":"gpt-5","stream":true}));
|
||||
|
||||
let plan_and_report =
|
||||
build_openai_cli_stream_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("builder should not error")
|
||||
.expect("plan should be built");
|
||||
|
||||
assert_eq!(
|
||||
plan_and_report
|
||||
.plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.map(String::as_str),
|
||||
Some("Bearer test")
|
||||
);
|
||||
assert_eq!(
|
||||
plan_and_report
|
||||
.plan
|
||||
.headers
|
||||
.get("accept")
|
||||
.map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bypasses_execution_runtime_for_codex_backendapi_variant() {
|
||||
let mut payload = missing_exact_provider_request_payload("openai_cli_stream");
|
||||
payload.provider_api_format = Some("openai:cli".to_string());
|
||||
payload.client_api_format = Some("openai:cli".to_string());
|
||||
payload.upstream_url = Some("https://chatgpt.com/backendapi/codex/responses".to_string());
|
||||
|
||||
assert!(should_bypass_intent_decision(&payload));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bypasses_execution_runtime_for_codex_plan_variant() {
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-123".to_string(),
|
||||
candidate_id: None,
|
||||
provider_name: Some("codex".to_string()),
|
||||
provider_id: "provider-123".to_string(),
|
||||
endpoint_id: "endpoint-123".to_string(),
|
||||
key_id: "key-123".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://chatgpt.com/backendapi/codex/responses".to_string(),
|
||||
headers: Default::default(),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model":"gpt-5.4"})),
|
||||
stream: true,
|
||||
client_api_format: "openai:cli".to_string(),
|
||||
provider_api_format: "openai:cli".to_string(),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
assert!(should_bypass_intent_plan(&plan));
|
||||
}
|
||||
937
apps/aether-gateway/src/execution_runtime/transport.rs
Normal file
937
apps/aether-gateway/src/execution_runtime/transport.rs
Normal file
@@ -0,0 +1,937 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::error::Error as _;
|
||||
use std::io::Write;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody,
|
||||
};
|
||||
use aether_http::{apply_http_client_config, HttpClientConfig};
|
||||
use base64::Engine as _;
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
|
||||
use reqwest::tls::Version;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
#[cfg(test)]
|
||||
use crate::gateway::execution_runtime::remote_compat::execute_sync_plan_via_remote_execution_runtime;
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
|
||||
const HUB_RELAY_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||
const DEFAULT_TUNNEL_BASE_URL: &str = "http://127.0.0.1:8084";
|
||||
const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay";
|
||||
const CLAUDE_CODE_TLS_PROFILE: &str = "claude_code_nodejs";
|
||||
|
||||
pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
|
||||
let mut kinds = Vec::new();
|
||||
if err.is_connect() {
|
||||
kinds.push("connect");
|
||||
}
|
||||
if err.is_timeout() {
|
||||
kinds.push("timeout");
|
||||
}
|
||||
if err.is_redirect() {
|
||||
kinds.push("redirect");
|
||||
}
|
||||
if err.is_body() {
|
||||
kinds.push("body");
|
||||
}
|
||||
if err.is_decode() {
|
||||
kinds.push("decode");
|
||||
}
|
||||
if err.is_request() {
|
||||
kinds.push("request");
|
||||
}
|
||||
|
||||
let mut detail = err.to_string();
|
||||
let mut source = err.source();
|
||||
while let Some(cause) = source {
|
||||
let cause_text = cause.to_string();
|
||||
if !cause_text.is_empty() && !detail.contains(&cause_text) {
|
||||
detail.push_str(": ");
|
||||
detail.push_str(&cause_text);
|
||||
}
|
||||
source = cause.source();
|
||||
}
|
||||
|
||||
if let Some(url) = err.url() {
|
||||
detail.push_str(" [url=");
|
||||
detail.push_str(url.as_str());
|
||||
detail.push(']');
|
||||
}
|
||||
if !kinds.is_empty() {
|
||||
detail.push_str(" [kind=");
|
||||
detail.push_str(&kinds.join(","));
|
||||
detail.push(']');
|
||||
}
|
||||
|
||||
detail
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(crate) enum ExecutionRuntimeTransportError {
|
||||
#[error("stream execution is not supported for this plan")]
|
||||
StreamUnsupported,
|
||||
#[error("request body must contain json_body or body_bytes_b64")]
|
||||
RequestBodyRequired,
|
||||
#[error("request body base64 is invalid: {0}")]
|
||||
BodyDecode(base64::DecodeError),
|
||||
#[error("request content-encoding is not supported: {0}")]
|
||||
UnsupportedContentEncoding(String),
|
||||
#[error("proxy execution is not supported")]
|
||||
ProxyUnsupported,
|
||||
#[error("invalid method: {0}")]
|
||||
InvalidMethod(#[from] http::method::InvalidMethod),
|
||||
#[error("invalid upstream header name: {0}")]
|
||||
InvalidHeaderName(String),
|
||||
#[error("invalid upstream header value for {0}")]
|
||||
InvalidHeaderValue(String),
|
||||
#[error("invalid proxy configuration: {0}")]
|
||||
InvalidProxy(reqwest::Error),
|
||||
#[error("failed to encode request body: {0}")]
|
||||
BodyEncode(serde_json::Error),
|
||||
#[error("failed to build HTTP client: {0}")]
|
||||
ClientBuild(reqwest::Error),
|
||||
#[error("failed to execute upstream request: {0}")]
|
||||
UpstreamRequest(String),
|
||||
#[error("hub relay request failed: {0}")]
|
||||
RelayError(String),
|
||||
#[error("upstream response is not valid JSON: {0}")]
|
||||
InvalidJson(serde_json::Error),
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct RelayRequestMeta {
|
||||
method: String,
|
||||
url: String,
|
||||
headers: BTreeMap<String, String>,
|
||||
timeout: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct DirectSyncExecutionRuntime;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct DirectUpstreamStreamExecution {
|
||||
pub(crate) request_id: String,
|
||||
pub(crate) candidate_id: Option<String>,
|
||||
pub(crate) status_code: u16,
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
pub(crate) response: reqwest::Response,
|
||||
pub(crate) started_at: Instant,
|
||||
}
|
||||
|
||||
impl DirectSyncExecutionRuntime {
|
||||
pub(crate) const fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_sync(
|
||||
&self,
|
||||
plan: ExecutionPlan,
|
||||
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
|
||||
let body_bytes = build_request_body(&plan)?;
|
||||
|
||||
let started_at = Instant::now();
|
||||
let response = send_request(&plan, body_bytes).await?;
|
||||
let status_code = response.status().as_u16();
|
||||
let headers = collect_response_headers(response.headers());
|
||||
let body_bytes = response.bytes().await.map_err(|err| {
|
||||
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
|
||||
})?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
|
||||
let body = if body_bytes.is_empty() {
|
||||
None
|
||||
} else if plan.stream {
|
||||
Some(ResponseBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
|
||||
})
|
||||
} else if response_body_is_json(&headers, &body_bytes) {
|
||||
let body_json: Value = serde_json::from_slice(&body_bytes)
|
||||
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
|
||||
Some(ResponseBody {
|
||||
json_body: Some(body_json),
|
||||
body_bytes_b64: None,
|
||||
})
|
||||
} else {
|
||||
Some(ResponseBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
|
||||
})
|
||||
};
|
||||
|
||||
Ok(ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: plan.candidate_id,
|
||||
status_code,
|
||||
headers,
|
||||
body,
|
||||
telemetry: Some(ExecutionTelemetry {
|
||||
ttfb_ms: None,
|
||||
elapsed_ms: Some(elapsed_ms),
|
||||
upstream_bytes: Some(upstream_bytes),
|
||||
}),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_stream(
|
||||
&self,
|
||||
plan: ExecutionPlan,
|
||||
) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> {
|
||||
if !plan.stream {
|
||||
return Err(ExecutionRuntimeTransportError::StreamUnsupported);
|
||||
}
|
||||
|
||||
let body_bytes = build_request_body(&plan)?;
|
||||
|
||||
let started_at = Instant::now();
|
||||
let response = send_request(&plan, body_bytes).await?;
|
||||
let status_code = response.status().as_u16();
|
||||
let headers = collect_response_headers(response.headers());
|
||||
|
||||
Ok(DirectUpstreamStreamExecution {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: plan.candidate_id,
|
||||
status_code,
|
||||
headers,
|
||||
response,
|
||||
started_at,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_sync_plan(
|
||||
state: &AppState,
|
||||
trace_id: Option<&str>,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<ExecutionResult, GatewayError> {
|
||||
#[cfg(test)]
|
||||
{
|
||||
let remote_execution_runtime_base_url = state
|
||||
.test_remote_execution_runtime_base_url()
|
||||
.unwrap_or_default();
|
||||
if !remote_execution_runtime_base_url.trim().is_empty() {
|
||||
return execute_sync_plan_via_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
trace_id,
|
||||
plan,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
let _ = state;
|
||||
let _ = trace_id;
|
||||
DirectSyncExecutionRuntime::new()
|
||||
.execute_sync(plan.clone())
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
async fn send_request(
|
||||
plan: &ExecutionPlan,
|
||||
body_bytes: Vec<u8>,
|
||||
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
|
||||
let method = plan.method.parse::<reqwest::Method>()?;
|
||||
let headers = build_request_headers(
|
||||
&plan.headers,
|
||||
plan.content_encoding.as_deref(),
|
||||
plan.body.body_bytes_b64.is_some(),
|
||||
)?;
|
||||
let total_timeout = plan
|
||||
.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| timeouts.total_ms)
|
||||
.map(Duration::from_millis);
|
||||
|
||||
if let Some(node_id) = resolve_tunnel_node_id(plan.proxy.as_ref()) {
|
||||
return send_via_tunnel_relay(plan, method, headers, body_bytes, &node_id, total_timeout)
|
||||
.await;
|
||||
}
|
||||
|
||||
let client = build_client(
|
||||
plan.timeouts.as_ref(),
|
||||
plan.proxy.as_ref(),
|
||||
plan.tls_profile.as_deref(),
|
||||
)?;
|
||||
let mut request = client.request(method, &plan.url);
|
||||
request = request.headers(headers).body(body_bytes);
|
||||
if let Some(timeout) = total_timeout {
|
||||
request = request.timeout(timeout);
|
||||
}
|
||||
request.send().await.map_err(|err| {
|
||||
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
|
||||
})
|
||||
}
|
||||
|
||||
async fn send_via_tunnel_relay(
|
||||
plan: &ExecutionPlan,
|
||||
method: reqwest::Method,
|
||||
headers: HeaderMap,
|
||||
body_bytes: Vec<u8>,
|
||||
node_id: &str,
|
||||
total_timeout: Option<Duration>,
|
||||
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
|
||||
let client = build_relay_client(plan.timeouts.as_ref())?;
|
||||
let relay_url = build_relay_url(plan.proxy.as_ref(), node_id);
|
||||
let envelope = build_relay_envelope(
|
||||
RelayRequestMeta {
|
||||
method: method.as_str().to_string(),
|
||||
url: plan.url.clone(),
|
||||
headers: header_map_to_string_map(&headers),
|
||||
timeout: resolve_relay_timeout_seconds(plan),
|
||||
},
|
||||
&body_bytes,
|
||||
)?;
|
||||
|
||||
let mut request = client
|
||||
.request(reqwest::Method::POST, relay_url)
|
||||
.header(reqwest::header::CONTENT_TYPE, HUB_RELAY_CONTENT_TYPE)
|
||||
.body(envelope);
|
||||
if let Some(timeout) = total_timeout {
|
||||
request = request.timeout(timeout);
|
||||
}
|
||||
|
||||
let response = request
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))?;
|
||||
|
||||
if let Some(kind) = response
|
||||
.headers()
|
||||
.get(HUB_RELAY_ERROR_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned)
|
||||
{
|
||||
let message = response
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|_| format!("hub relay error: {kind}"));
|
||||
return Err(ExecutionRuntimeTransportError::RelayError(message));
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
fn build_request_body(plan: &ExecutionPlan) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
let mut body_bytes = if let Some(json_body) = plan.body.json_body.clone() {
|
||||
serde_json::to_vec(&json_body).map_err(ExecutionRuntimeTransportError::BodyEncode)?
|
||||
} else if let Some(body_b64) = plan.body.body_bytes_b64.as_deref() {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(body_b64)
|
||||
.map_err(ExecutionRuntimeTransportError::BodyDecode)?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
if should_gzip_request_body(plan) && plan.body.json_body.is_some() {
|
||||
body_bytes = gzip_bytes(&body_bytes)?;
|
||||
}
|
||||
|
||||
Ok(body_bytes)
|
||||
}
|
||||
|
||||
fn should_gzip_request_body(plan: &ExecutionPlan) -> bool {
|
||||
matches!(
|
||||
normalize_content_encoding(plan.content_encoding.as_deref()).as_deref(),
|
||||
Some("gzip")
|
||||
)
|
||||
}
|
||||
|
||||
fn normalize_content_encoding(value: Option<&str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
}
|
||||
|
||||
fn gzip_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
encoder
|
||||
.write_all(body_bytes)
|
||||
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))?;
|
||||
encoder
|
||||
.finish()
|
||||
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))
|
||||
}
|
||||
|
||||
fn build_relay_client(
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
) -> Result<reqwest::Client, ExecutionRuntimeTransportError> {
|
||||
let builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
use_rustls_tls: false,
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder
|
||||
.build()
|
||||
.map_err(ExecutionRuntimeTransportError::ClientBuild)
|
||||
}
|
||||
|
||||
fn build_relay_envelope(
|
||||
meta: RelayRequestMeta,
|
||||
body_bytes: &[u8],
|
||||
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
let meta_bytes =
|
||||
serde_json::to_vec(&meta).map_err(ExecutionRuntimeTransportError::BodyEncode)?;
|
||||
let mut envelope = Vec::with_capacity(4 + meta_bytes.len() + body_bytes.len());
|
||||
envelope.extend_from_slice(&(meta_bytes.len() as u32).to_be_bytes());
|
||||
envelope.extend_from_slice(&meta_bytes);
|
||||
envelope.extend_from_slice(body_bytes);
|
||||
Ok(envelope)
|
||||
}
|
||||
|
||||
fn build_relay_url(proxy: Option<&ProxySnapshot>, node_id: &str) -> String {
|
||||
let base_url = proxy
|
||||
.and_then(resolve_tunnel_base_url_from_proxy)
|
||||
.or_else(|| std::env::var("AETHER_TUNNEL_BASE_URL").ok())
|
||||
.unwrap_or_else(|| DEFAULT_TUNNEL_BASE_URL.to_string());
|
||||
format!(
|
||||
"{}{}/{}",
|
||||
base_url.trim_end_matches('/'),
|
||||
TUNNEL_RELAY_PATH_PREFIX,
|
||||
node_id
|
||||
)
|
||||
}
|
||||
|
||||
fn resolve_tunnel_base_url_from_proxy(proxy: &ProxySnapshot) -> Option<String> {
|
||||
let extra = proxy.extra.as_ref()?;
|
||||
let value = extra.get("tunnel_base_url")?.as_str()?.trim();
|
||||
if !value.is_empty() {
|
||||
return Some(value.to_string());
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn resolve_relay_timeout_seconds(plan: &ExecutionPlan) -> u64 {
|
||||
let ms = plan
|
||||
.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| {
|
||||
timeouts
|
||||
.read_ms
|
||||
.or(timeouts.total_ms)
|
||||
.or(timeouts.connect_ms)
|
||||
})
|
||||
.unwrap_or(60_000);
|
||||
let secs = ms.div_ceil(1_000);
|
||||
secs.clamp(1, 300)
|
||||
}
|
||||
|
||||
fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
|
||||
let proxy = proxy?;
|
||||
if proxy.enabled == Some(false) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let proxy_mode = proxy
|
||||
.mode
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let node_id = proxy.node_id.as_deref().map(str::trim).unwrap_or_default();
|
||||
let has_node_id = !node_id.is_empty();
|
||||
let has_proxy_url = proxy
|
||||
.url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|url| !url.is_empty());
|
||||
|
||||
if has_node_id && (proxy_mode == "tunnel" || !has_proxy_url) {
|
||||
return Some(node_id.to_string());
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn build_client(
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
proxy: Option<&ProxySnapshot>,
|
||||
tls_profile: Option<&str>,
|
||||
) -> Result<reqwest::Client, ExecutionRuntimeTransportError> {
|
||||
let mut builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder = apply_tls_profile(builder, tls_profile);
|
||||
if let Some(proxy_url) = resolve_proxy_url(proxy)? {
|
||||
let proxy = reqwest::Proxy::all(&proxy_url)
|
||||
.map_err(ExecutionRuntimeTransportError::InvalidProxy)?;
|
||||
builder = builder.proxy(proxy);
|
||||
}
|
||||
builder
|
||||
.build()
|
||||
.map_err(ExecutionRuntimeTransportError::ClientBuild)
|
||||
}
|
||||
|
||||
fn apply_tls_profile(
|
||||
builder: reqwest::ClientBuilder,
|
||||
tls_profile: Option<&str>,
|
||||
) -> reqwest::ClientBuilder {
|
||||
let profile = normalize_tls_profile(tls_profile);
|
||||
if profile.is_none() {
|
||||
return builder;
|
||||
}
|
||||
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
|
||||
let tls_config = build_best_effort_tls_config();
|
||||
let builder = builder
|
||||
.use_preconfigured_tls(tls_config)
|
||||
.min_tls_version(Version::TLS_1_2)
|
||||
.max_tls_version(Version::TLS_1_3);
|
||||
|
||||
if profile.as_deref() == Some(CLAUDE_CODE_TLS_PROFILE) {
|
||||
return builder;
|
||||
}
|
||||
|
||||
builder
|
||||
}
|
||||
|
||||
fn normalize_tls_profile(tls_profile: Option<&str>) -> Option<String> {
|
||||
let profile = tls_profile
|
||||
.map(str::trim)
|
||||
.filter(|profile| !profile.is_empty())?
|
||||
.to_ascii_lowercase();
|
||||
Some(profile)
|
||||
}
|
||||
|
||||
fn build_best_effort_tls_config() -> rustls::ClientConfig {
|
||||
let root_store =
|
||||
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
let mut config = rustls::ClientConfig::builder_with_protocol_versions(&[
|
||||
&rustls::version::TLS13,
|
||||
&rustls::version::TLS12,
|
||||
])
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth();
|
||||
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
|
||||
config
|
||||
}
|
||||
|
||||
fn resolve_proxy_url(
|
||||
proxy: Option<&ProxySnapshot>,
|
||||
) -> Result<Option<String>, ExecutionRuntimeTransportError> {
|
||||
let Some(proxy) = proxy else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if proxy.enabled == Some(false) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Some(proxy_url) = proxy
|
||||
.url
|
||||
.as_ref()
|
||||
.map(|url| url.trim())
|
||||
.filter(|url| !url.is_empty())
|
||||
{
|
||||
return Ok(Some(proxy_url.to_string()));
|
||||
}
|
||||
|
||||
if proxy.node_id.is_some() || proxy.mode.as_deref() == Some("tunnel") {
|
||||
return Err(ExecutionRuntimeTransportError::ProxyUnsupported);
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn build_request_headers(
|
||||
headers: &BTreeMap<String, String>,
|
||||
content_encoding: Option<&str>,
|
||||
allow_passthrough_content_encoding: bool,
|
||||
) -> Result<HeaderMap, ExecutionRuntimeTransportError> {
|
||||
let mut out = HeaderMap::new();
|
||||
let normalized_content_encoding = normalize_content_encoding(content_encoding);
|
||||
if let Some(encoding) = normalized_content_encoding.as_deref() {
|
||||
if encoding != "gzip" && !allow_passthrough_content_encoding {
|
||||
return Err(ExecutionRuntimeTransportError::UnsupportedContentEncoding(
|
||||
encoding.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
for (key, value) in headers {
|
||||
let normalized_key = key.trim().to_ascii_lowercase();
|
||||
if is_hop_by_hop_header(&normalized_key) || normalized_key == "content-encoding" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let header_name = HeaderName::from_bytes(key.as_bytes())
|
||||
.map_err(|_| ExecutionRuntimeTransportError::InvalidHeaderName(key.clone()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|_| ExecutionRuntimeTransportError::InvalidHeaderValue(key.clone()))?;
|
||||
out.insert(header_name, header_value);
|
||||
}
|
||||
if let Some(encoding) = normalized_content_encoding {
|
||||
out.insert(
|
||||
reqwest::header::CONTENT_ENCODING,
|
||||
HeaderValue::from_str(&encoding).map_err(|_| {
|
||||
ExecutionRuntimeTransportError::InvalidHeaderValue("content-encoding".into())
|
||||
})?,
|
||||
);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn header_map_to_string_map(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
value
|
||||
.to_str()
|
||||
.ok()
|
||||
.map(|value| (name.as_str().to_string(), value.to_string()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn is_hop_by_hop_header(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
"host"
|
||||
| "content-length"
|
||||
| "connection"
|
||||
| "upgrade"
|
||||
| "keep-alive"
|
||||
| "proxy-authorization"
|
||||
| "proxy-connection"
|
||||
| "te"
|
||||
| "trailer"
|
||||
| "transfer-encoding"
|
||||
)
|
||||
}
|
||||
|
||||
fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
header_map_to_string_map(headers)
|
||||
}
|
||||
|
||||
fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
if headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.is_some_and(|value| value.contains("json"))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
serde_json::from_slice::<Value>(body_bytes).is_ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::io::Read;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::Path;
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use serde_json::json;
|
||||
|
||||
use super::DirectSyncExecutionRuntime;
|
||||
|
||||
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
|
||||
aether_contracts::ProxySnapshot {
|
||||
enabled: Some(true),
|
||||
mode: Some("tunnel".into()),
|
||||
node_id: Some("node-1".into()),
|
||||
label: Some("relay-node".into()),
|
||||
url: None,
|
||||
extra: Some(json!({"tunnel_base_url": base_url})),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_relay_envelope(body: &[u8]) -> (serde_json::Value, Vec<u8>) {
|
||||
assert!(
|
||||
body.len() >= 4,
|
||||
"relay body must contain meta length prefix"
|
||||
);
|
||||
let meta_len = u32::from_be_bytes([body[0], body[1], body[2], body[3]]) as usize;
|
||||
let meta_end = 4 + meta_len;
|
||||
let meta = serde_json::from_slice::<serde_json::Value>(&body[4..meta_end])
|
||||
.expect("relay meta should decode");
|
||||
(meta, body[meta_end..].to_vec())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_sync_execution_runtime_preserves_upstream_status_and_json_body() {
|
||||
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 app = Router::new().route(
|
||||
"/chat",
|
||||
post(|| async {
|
||||
(
|
||||
axum::http::StatusCode::TOO_MANY_REQUESTS,
|
||||
Json(json!({"error": {"message": "slow down"}})),
|
||||
)
|
||||
}),
|
||||
);
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("test server should run");
|
||||
});
|
||||
|
||||
let execution_runtime = DirectSyncExecutionRuntime::new();
|
||||
let result = execution_runtime
|
||||
.execute_sync(ExecutionPlan {
|
||||
request_id: "req-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: format!("http://{addr}/chat"),
|
||||
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("sync execution should succeed");
|
||||
|
||||
server.abort();
|
||||
|
||||
assert_eq!(result.status_code, 429);
|
||||
assert_eq!(
|
||||
result.body.and_then(|body| body.json_body),
|
||||
Some(json!({"error": {"message": "slow down"}}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_sync_execution_runtime_supports_tunnel_relay() {
|
||||
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 app = Router::new().route(
|
||||
"/api/internal/tunnel/relay/{node_id}",
|
||||
post(|Path(node_id): Path<String>, body: Bytes| async move {
|
||||
let (meta, request_body) = decode_relay_envelope(&body);
|
||||
assert_eq!(node_id, "node-1");
|
||||
assert_eq!(meta["method"], "POST");
|
||||
assert_eq!(meta["url"], "https://example.com/chat");
|
||||
let request_json: serde_json::Value =
|
||||
serde_json::from_slice(&request_body).expect("request body should be json");
|
||||
assert_eq!(request_json["model"], "gpt-4.1");
|
||||
(
|
||||
axum::http::StatusCode::OK,
|
||||
Json(json!({"tunnel": true, "node_id": node_id})),
|
||||
)
|
||||
}),
|
||||
);
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("relay test server should run");
|
||||
});
|
||||
|
||||
let execution_runtime = DirectSyncExecutionRuntime::new();
|
||||
let result = execution_runtime
|
||||
.execute_sync(ExecutionPlan {
|
||||
request_id: "req-1".into(),
|
||||
candidate_id: None,
|
||||
provider_name: None,
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://example.com/chat".into(),
|
||||
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: Some(tunnel_proxy_snapshot(format!("http://{addr}"))),
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("tunnel relay execution should succeed");
|
||||
|
||||
server.abort();
|
||||
|
||||
assert_eq!(result.status_code, 200);
|
||||
assert_eq!(
|
||||
result.body.and_then(|body| body.json_body),
|
||||
Some(json!({"tunnel": true, "node_id": "node-1"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_sync_execution_runtime_allows_tls_profile_best_effort() {
|
||||
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 app = Router::new().route(
|
||||
"/chat",
|
||||
post(|| async {
|
||||
(
|
||||
axum::http::StatusCode::OK,
|
||||
Json(json!({"tls_profile": true})),
|
||||
)
|
||||
}),
|
||||
);
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("test server should run");
|
||||
});
|
||||
|
||||
let execution_runtime = DirectSyncExecutionRuntime::new();
|
||||
let result = execution_runtime
|
||||
.execute_sync(ExecutionPlan {
|
||||
request_id: "req-tls-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("claude".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: format!("http://{addr}/chat"),
|
||||
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "claude-3.7-sonnet"})),
|
||||
stream: false,
|
||||
client_api_format: "claude:chat".into(),
|
||||
provider_api_format: "claude:chat".into(),
|
||||
model_name: Some("claude-3.7-sonnet".into()),
|
||||
proxy: None,
|
||||
tls_profile: Some("claude_code_nodejs".into()),
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("sync execution with tls profile should succeed");
|
||||
|
||||
server.abort();
|
||||
|
||||
assert_eq!(result.status_code, 200);
|
||||
assert_eq!(
|
||||
result.body.and_then(|body| body.json_body),
|
||||
Some(json!({"tls_profile": true}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_sync_execution_runtime_compresses_json_body_when_requested() {
|
||||
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 app = Router::new().route(
|
||||
"/chat",
|
||||
post(|headers: axum::http::HeaderMap, body: Bytes| async move {
|
||||
let header_encoding = headers
|
||||
.get(axum::http::header::CONTENT_ENCODING)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let mut decoder = flate2::read::GzDecoder::new(body.as_ref());
|
||||
let mut decoded = String::new();
|
||||
decoder
|
||||
.read_to_string(&mut decoded)
|
||||
.expect("gzip body should decode");
|
||||
let decoded_json: serde_json::Value =
|
||||
serde_json::from_str(&decoded).expect("decoded json should parse");
|
||||
(
|
||||
axum::http::StatusCode::OK,
|
||||
Json(json!({
|
||||
"content_encoding": header_encoding,
|
||||
"body": decoded_json,
|
||||
})),
|
||||
)
|
||||
}),
|
||||
);
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("test server should run");
|
||||
});
|
||||
|
||||
let execution_runtime = DirectSyncExecutionRuntime::new();
|
||||
let result = execution_runtime
|
||||
.execute_sync(ExecutionPlan {
|
||||
request_id: "req-gzip-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: format!("http://{addr}/chat"),
|
||||
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: Some("gzip".into()),
|
||||
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(5_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("gzip sync execution should succeed");
|
||||
|
||||
server.abort();
|
||||
|
||||
assert_eq!(result.status_code, 200);
|
||||
assert_eq!(
|
||||
result.body.and_then(|body| body.json_body),
|
||||
Some(json!({
|
||||
"content_encoding": "gzip",
|
||||
"body": {"model": "gpt-4.1"},
|
||||
}))
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user