feat: 实现 executor sync 执行路径及 plan-sync/report-sync 端点

Rust gateway 新增同步执行模式,支持 AI 请求(OpenAI/Claude/Gemini chat/cli)
及 Gemini Files CRUD 通过 executor 直接同步执行,执行完成后通过 report-sync
端点向 control 回报结果。Python 侧新增 /plan-sync 和 /report-sync 内部端点,
构建同步执行计划并处理结果上报。包含完整测试覆盖。
This commit is contained in:
fawney19
2026-03-21 14:07:45 +08:00
parent d735b6316f
commit 53bb23b510
9 changed files with 4926 additions and 25 deletions

View File

@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct ExecutionTimeouts {
#[serde(skip_serializing_if = "Option::is_none")]
@@ -20,19 +20,6 @@ pub struct ExecutionTimeouts {
pub total_ms: Option<u64>,
}
impl Default for ExecutionTimeouts {
fn default() -> Self {
Self {
connect_ms: None,
read_ms: None,
first_byte_ms: None,
write_ms: None,
pool_ms: None,
total_ms: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RequestBody {
#[serde(default, skip_serializing_if = "Option::is_none")]

View File

@@ -206,9 +206,7 @@ pub(crate) async fn maybe_execute_via_control(
_ if !is_files_route && !is_video_route => return Ok(None),
_ => json!({}),
}
} else if is_video_route && body_bytes.is_empty() {
json!({})
} else if is_files_route {
} else if (is_video_route && body_bytes.is_empty()) || is_files_route {
json!({})
} else {
return Ok(None);

View File

@@ -1,7 +1,9 @@
use std::collections::BTreeMap;
use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, StreamFrame, StreamFramePayload};
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, StreamFrame, StreamFramePayload,
};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::http::Response;
@@ -14,14 +16,26 @@ use tokio_util::io::StreamReader;
use tracing::warn;
use crate::gateway::constants::*;
use crate::gateway::headers::{collect_control_headers, header_equals};
use crate::gateway::headers::{collect_control_headers, header_equals, is_json_request};
use crate::gateway::{
build_client_response, build_client_response_from_parts, AppState, GatewayControlAuthContext,
GatewayControlDecision, GatewayError,
};
const GEMINI_FILES_GET_PLAN_KIND: &str = "gemini_files_get";
const GEMINI_FILES_LIST_PLAN_KIND: &str = "gemini_files_list";
const GEMINI_FILES_UPLOAD_PLAN_KIND: &str = "gemini_files_upload";
const GEMINI_FILES_DELETE_PLAN_KIND: &str = "gemini_files_delete";
const GEMINI_FILES_DOWNLOAD_PLAN_KIND: &str = "gemini_files_download";
const OPENAI_VIDEO_CONTENT_PLAN_KIND: &str = "openai_video_content";
const OPENAI_CHAT_SYNC_PLAN_KIND: &str = "openai_chat_sync";
const OPENAI_CLI_SYNC_PLAN_KIND: &str = "openai_cli_sync";
const OPENAI_COMPACT_SYNC_PLAN_KIND: &str = "openai_compact_sync";
const CLAUDE_CHAT_SYNC_PLAN_KIND: &str = "claude_chat_sync";
const GEMINI_CHAT_SYNC_PLAN_KIND: &str = "gemini_chat_sync";
const CLAUDE_CLI_SYNC_PLAN_KIND: &str = "claude_cli_sync";
const GEMINI_CLI_SYNC_PLAN_KIND: &str = "gemini_cli_sync";
const EXECUTOR_SYNC_ACTION: &str = "executor_sync";
const EXECUTOR_STREAM_ACTION: &str = "executor_stream";
const MAX_ERROR_BODY_BYTES: usize = 16_384;
@@ -45,6 +59,147 @@ struct GatewayControlPlanResponse {
plan_kind: Option<String>,
#[serde(default)]
plan: Option<ExecutionPlan>,
#[serde(default)]
report_kind: Option<String>,
#[serde(default)]
report_context: Option<serde_json::Value>,
}
#[derive(Debug, Serialize)]
struct GatewaySyncReportRequest {
trace_id: String,
report_kind: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
report_context: Option<serde_json::Value>,
status_code: u16,
headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
body_json: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
telemetry: Option<ExecutionTelemetry>,
}
pub(crate) async fn maybe_execute_via_executor_sync(
state: &AppState,
parts: &http::request::Parts,
body_bytes: &Bytes,
trace_id: &str,
decision: Option<&GatewayControlDecision>,
) -> Result<Option<Response<Body>>, GatewayError> {
let Some(control_base_url) = state.control_base_url.as_deref() else {
return Ok(None);
};
let Some(executor_base_url) = state.executor_base_url.as_deref() else {
return Ok(None);
};
let Some(decision) = decision else {
return Ok(None);
};
let Some(plan_kind) = resolve_direct_executor_sync_plan_kind(parts, decision) else {
return Ok(None);
};
let (body_json, body_base64) = if is_json_request(&parts.headers) {
if body_bytes.is_empty() {
(json!({}), None)
} else {
match serde_json::from_slice::<serde_json::Value>(body_bytes) {
Ok(value) => (value, None),
Err(_) => return Ok(None),
}
}
} else {
(
json!({}),
(!body_bytes.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(body_bytes)),
)
};
let request_payload = GatewayControlPlanRequest {
trace_id: trace_id.to_string(),
method: parts.method.to_string(),
path: parts.uri.path().to_string(),
query_string: parts.uri.query().map(ToOwned::to_owned),
headers: collect_control_headers(&parts.headers),
body_json,
body_base64,
auth_context: decision.auth_context.clone(),
};
let response = state
.client
.post(format!("{control_base_url}/api/internal/gateway/plan-sync"))
.header(TRACE_ID_HEADER, trace_id)
.json(&request_payload)
.send()
.await
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
if response.status() == http::StatusCode::CONFLICT
&& header_equals(
response.headers(),
CONTROL_ACTION_HEADER,
CONTROL_ACTION_PROXY_PUBLIC,
)
{
return Ok(None);
}
if header_equals(response.headers(), CONTROL_EXECUTED_HEADER, "true")
&& response.status() != http::StatusCode::OK
{
return Ok(Some(build_client_response(
response,
trace_id,
Some(decision),
)?));
}
let response = response
.error_for_status()
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
let payload: GatewayControlPlanResponse = response
.json()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if payload.action != EXECUTOR_SYNC_ACTION {
return Ok(None);
}
if payload.plan_kind.as_deref() != Some(plan_kind) {
return Ok(None);
}
let Some(plan) = payload.plan else {
return Err(GatewayError::Internal(
"gateway sync plan response missing execution plan".to_string(),
));
};
execute_executor_sync(
state,
control_base_url,
executor_base_url,
plan,
trace_id,
decision,
plan_kind,
payload.report_kind,
payload.report_context,
)
.await
}
pub(crate) async fn maybe_execute_via_executor_stream(
@@ -174,6 +329,96 @@ fn resolve_direct_executor_stream_plan_kind(
None
}
fn resolve_direct_executor_sync_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
if decision.route_class.as_deref() != Some("ai_public") {
return None;
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/chat/completions"
{
return Some(OPENAI_CHAT_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses"
{
return Some(OPENAI_CLI_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("compact")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/responses/compact"
{
return Some(OPENAI_COMPACT_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some(CLAUDE_CHAT_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path() == "/v1/messages"
{
return Some(CLAUDE_CLI_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("chat")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":generateContent")
{
return Some(GEMINI_CHAT_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("cli")
&& parts.method == http::Method::POST
&& parts.uri.path().ends_with(":generateContent")
{
return Some(GEMINI_CLI_SYNC_PLAN_KIND);
}
if decision.route_family.as_deref() == Some("gemini")
&& decision.route_kind.as_deref() == Some("files")
{
if parts.method == http::Method::POST && parts.uri.path() == "/upload/v1beta/files" {
return Some(GEMINI_FILES_UPLOAD_PLAN_KIND);
}
if parts.method == http::Method::GET && parts.uri.path() == "/v1beta/files" {
return Some(GEMINI_FILES_LIST_PLAN_KIND);
}
if parts.method == http::Method::GET
&& parts.uri.path().starts_with("/v1beta/files/")
&& !parts.uri.path().ends_with(":download")
{
return Some(GEMINI_FILES_GET_PLAN_KIND);
}
if parts.method == http::Method::DELETE
&& parts.uri.path().starts_with("/v1beta/files/")
&& !parts.uri.path().ends_with(":download")
{
return Some(GEMINI_FILES_DELETE_PLAN_KIND);
}
}
None
}
async fn execute_executor_stream(
state: &AppState,
executor_base_url: &str,
@@ -242,7 +487,7 @@ async fn execute_executor_stream(
let next_frame = match read_next_frame(&mut lines).await {
Ok(frame) => frame,
Err(err) => {
warn!(trace_id = %trace_id_owned, error = %format!("{err:?}"), "gateway failed to decode executor stream frame");
warn!(trace_id = %trace_id_owned, error = ?err, "gateway failed to decode executor stream frame");
break;
}
};
@@ -283,6 +528,165 @@ async fn execute_executor_stream(
)?))
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
async fn execute_executor_sync(
state: &AppState,
control_base_url: &str,
executor_base_url: &str,
plan: ExecutionPlan,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
report_kind: Option<String>,
report_context: Option<serde_json::Value>,
) -> Result<Option<Response<Body>>, GatewayError> {
let response = match state
.client
.post(format!("{executor_base_url}/v1/execute/sync"))
.header(TRACE_ID_HEADER, trace_id)
.json(&plan)
.send()
.await
{
Ok(response) => response,
Err(err) => {
warn!(trace_id = %trace_id, error = %err, "gateway direct executor sync unavailable");
return Ok(None);
}
};
if response.status() != http::StatusCode::OK {
return Ok(Some(build_client_response(
response,
trace_id,
Some(decision),
)?));
}
let result: ExecutionResult = response
.json()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let mut headers = result.headers.clone();
let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
if should_fallback_to_control_sync(plan_kind, &result, body_json.as_ref()) {
return Ok(None);
}
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(),
body_base64: body_base64.clone(),
telemetry: result.telemetry.clone(),
};
if let Err(err) = submit_sync_report(state, control_base_url, trace_id, report).await {
warn!(trace_id = %trace_id, error = ?err, "gateway failed to submit sync execution report");
}
}
Ok(Some(build_client_response_from_parts(
result.status_code,
&headers,
Body::from(body_bytes),
trace_id,
Some(decision),
)?))
}
fn should_fallback_to_control_sync(
plan_kind: &str,
result: &ExecutionResult,
body_json: Option<&serde_json::Value>,
) -> bool {
if !matches!(
plan_kind,
OPENAI_CHAT_SYNC_PLAN_KIND
| OPENAI_CLI_SYNC_PLAN_KIND
| OPENAI_COMPACT_SYNC_PLAN_KIND
| CLAUDE_CHAT_SYNC_PLAN_KIND
| GEMINI_CHAT_SYNC_PLAN_KIND
| CLAUDE_CLI_SYNC_PLAN_KIND
| GEMINI_CLI_SYNC_PLAN_KIND
) {
return false;
}
if result.status_code >= 400 {
return true;
}
let Some(body_json) = body_json else {
return true;
};
body_json.get("error").is_some()
}
type DecodedBody = (Vec<u8>, Option<serde_json::Value>, Option<String>);
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))
}
async fn submit_sync_report(
state: &AppState,
control_base_url: &str,
trace_id: &str,
payload: GatewaySyncReportRequest,
) -> Result<(), GatewayError> {
let response = state
.client
.post(format!(
"{control_base_url}/api/internal/gateway/report-sync"
))
.header(TRACE_ID_HEADER, trace_id)
.json(&payload)
.send()
.await
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
response
.error_for_status()
.map_err(|err| GatewayError::ControlUnavailable {
trace_id: trace_id.to_string(),
message: err.to_string(),
})?;
Ok(())
}
async fn collect_error_body<R>(
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Vec<u8>, GatewayError>

View File

@@ -22,7 +22,7 @@ pub(crate) use control::{
GatewayControlDecision,
};
pub(crate) use error::GatewayError;
pub(crate) use executor::maybe_execute_via_executor_stream;
pub(crate) use executor::{maybe_execute_via_executor_stream, maybe_execute_via_executor_sync};
use handlers::{health, proxy_request};
pub(crate) use response::{build_client_response, build_client_response_from_parts};

File diff suppressed because it is too large Load Diff

View File

@@ -535,3 +535,807 @@ async fn gateway_executes_gemini_files_download_via_executor_stream_plan() {
executor_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_files_get_via_executor_sync_plan() {
#[derive(Debug, Clone)]
struct SeenPlanSyncRequest {
trace_id: String,
method: String,
path: String,
query_string: String,
user_id: String,
}
#[derive(Debug, Clone)]
struct SeenExecutorSyncRequest {
trace_id: String,
method: String,
url: String,
stream: bool,
client_api_format: String,
}
#[derive(Debug, Clone)]
struct SeenReportSyncRequest {
trace_id: String,
report_kind: String,
status_code: u64,
file_key_id: String,
user_id: String,
file_name: String,
}
let seen_plan = Arc::new(Mutex::new(None::<SeenPlanSyncRequest>));
let seen_plan_clone = Arc::clone(&seen_plan);
let seen_executor = Arc::new(Mutex::new(None::<SeenExecutorSyncRequest>));
let seen_executor_clone = Arc::clone(&seen_executor);
let seen_report = Arc::new(Mutex::new(None::<SeenReportSyncRequest>));
let seen_report_clone = Arc::clone(&seen_report);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-sync-123",
"api_key_id": "key-files-sync-123",
"access_allowed": true
},
"public_path": "/v1beta/files/files/abc-123"
}))
}),
)
.route(
"/api/internal/gateway/plan-sync",
any(move |request: Request| {
let seen_plan_inner = Arc::clone(&seen_plan_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("plan payload should parse");
*seen_plan_inner.lock().expect("mutex should lock") =
Some(SeenPlanSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
query_string: payload
.get("query_string")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("auth_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"action": "executor_sync",
"plan_kind": "gemini_files_get",
"plan": {
"request_id": "req-files-sync-123",
"provider_id": "provider-files-sync-123",
"endpoint_id": "endpoint-files-sync-123",
"key_id": "file-key-sync-123",
"provider_name": "gemini",
"method": "GET",
"url": "https://files.example/v1beta/files/files/abc-123?view=FULL",
"headers": {
"authorization": "Bearer upstream-key"
},
"body": {},
"stream": false,
"client_api_format": "gemini:files",
"provider_api_format": "gemini:files",
"model_name": "gemini-files"
},
"report_kind": "gemini_files_store_mapping",
"report_context": {
"file_key_id": "file-key-sync-123",
"user_id": "user-files-sync-123"
}
}))
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |request: Request| {
let seen_report_inner = Arc::clone(&seen_report_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("report payload should parse");
*seen_report_inner.lock().expect("mutex should lock") =
Some(SeenReportSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
report_kind: payload
.get("report_kind")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
status_code: payload
.get("status_code")
.and_then(|value| value.as_u64())
.unwrap_or_default(),
file_key_id: payload
.get("report_context")
.and_then(|value| value.get("file_key_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("report_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
file_name: payload
.get("body_json")
.and_then(|value| value.get("name"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({"ok": true}))
}
}),
)
.route(
"/v1beta/files/files/abc-123",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let executor = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_executor_inner = Arc::clone(&seen_executor_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("executor payload should parse");
*seen_executor_inner.lock().expect("mutex should lock") =
Some(SeenExecutorSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
stream: payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(true),
client_api_format: payload
.get("client_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "req-files-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"name": "files/abc-123",
"displayName": "ABC"
}
}
}))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (executor_url, executor_handle) = start_server(executor).await;
let gateway =
build_router_with_endpoints(upstream_url.clone(), Some(upstream_url), Some(executor_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/v1beta/files/files/abc-123?view=FULL"
))
.header(TRACE_ID_HEADER, "trace-files-sync-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("json should parse"),
json!({
"name": "files/abc-123",
"displayName": "ABC"
})
);
let seen_plan_request = seen_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("plan-sync should be captured");
assert_eq!(seen_plan_request.trace_id, "trace-files-sync-123");
assert_eq!(seen_plan_request.method, "GET");
assert_eq!(seen_plan_request.path, "/v1beta/files/files/abc-123");
assert_eq!(seen_plan_request.query_string, "view=FULL");
assert_eq!(seen_plan_request.user_id, "user-files-sync-123");
let seen_executor_request = seen_executor
.lock()
.expect("mutex should lock")
.clone()
.expect("executor sync should be captured");
assert_eq!(seen_executor_request.trace_id, "trace-files-sync-123");
assert_eq!(seen_executor_request.method, "GET");
assert_eq!(
seen_executor_request.url,
"https://files.example/v1beta/files/files/abc-123?view=FULL"
);
assert!(!seen_executor_request.stream);
assert_eq!(seen_executor_request.client_api_format, "gemini:files");
let seen_report_request = seen_report
.lock()
.expect("mutex should lock")
.clone()
.expect("report-sync should be captured");
assert_eq!(seen_report_request.trace_id, "trace-files-sync-123");
assert_eq!(
seen_report_request.report_kind,
"gemini_files_store_mapping"
);
assert_eq!(seen_report_request.status_code, 200);
assert_eq!(seen_report_request.file_key_id, "file-key-sync-123");
assert_eq!(seen_report_request.user_id, "user-files-sync-123");
assert_eq!(seen_report_request.file_name, "files/abc-123");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
executor_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_files_upload_via_executor_sync_plan() {
#[derive(Debug, Clone)]
struct SeenPlanSyncRequest {
method: String,
path: String,
body_base64: String,
}
#[derive(Debug, Clone)]
struct SeenExecutorSyncRequest {
method: String,
url: String,
body_bytes_b64: String,
}
#[derive(Debug, Clone)]
struct SeenReportSyncRequest {
report_kind: String,
file_key_id: String,
user_id: String,
file_name: String,
}
let seen_plan = Arc::new(Mutex::new(None::<SeenPlanSyncRequest>));
let seen_plan_clone = Arc::clone(&seen_plan);
let seen_executor = Arc::new(Mutex::new(None::<SeenExecutorSyncRequest>));
let seen_executor_clone = Arc::clone(&seen_executor);
let seen_report = Arc::new(Mutex::new(None::<SeenReportSyncRequest>));
let seen_report_clone = Arc::clone(&seen_report);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-upload-123",
"api_key_id": "key-files-upload-123",
"access_allowed": true
},
"public_path": "/upload/v1beta/files"
}))
}),
)
.route(
"/api/internal/gateway/plan-sync",
any(move |request: Request| {
let seen_plan_inner = Arc::clone(&seen_plan_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("plan payload should parse");
*seen_plan_inner.lock().expect("mutex should lock") =
Some(SeenPlanSyncRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
path: payload
.get("path")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_base64: payload
.get("body_base64")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"action": "executor_sync",
"plan_kind": "gemini_files_upload",
"plan": {
"request_id": "req-files-upload-123",
"provider_id": "provider-files-upload-123",
"endpoint_id": "endpoint-files-upload-123",
"key_id": "file-key-upload-123",
"provider_name": "gemini",
"method": "POST",
"url": "https://files.example/upload/v1beta/files?uploadType=resumable",
"headers": {
"content-type": "application/octet-stream",
"x-goog-api-key": "upstream-key"
},
"body": {
"body_bytes_b64": "dXBsb2FkLWJ5dGVz"
},
"stream": false,
"client_api_format": "gemini:files",
"provider_api_format": "gemini:files",
"model_name": "gemini-files"
},
"report_kind": "gemini_files_store_mapping",
"report_context": {
"file_key_id": "file-key-upload-123",
"user_id": "user-files-upload-123"
}
}))
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |request: Request| {
let seen_report_inner = Arc::clone(&seen_report_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("report payload should parse");
*seen_report_inner.lock().expect("mutex should lock") =
Some(SeenReportSyncRequest {
report_kind: payload
.get("report_kind")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
file_key_id: payload
.get("report_context")
.and_then(|value| value.get("file_key_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_id: payload
.get("report_context")
.and_then(|value| value.get("user_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
file_name: payload
.get("body_json")
.and_then(|value| value.get("file"))
.and_then(|value| value.get("name"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({"ok": true}))
}
}),
)
.route(
"/upload/v1beta/files",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let executor = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_executor_inner = Arc::clone(&seen_executor_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("executor payload should parse");
*seen_executor_inner.lock().expect("mutex should lock") =
Some(SeenExecutorSyncRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_bytes_b64: payload
.get("body")
.and_then(|value| value.get("body_bytes_b64"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "req-files-upload-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"file": {
"name": "files/uploaded-123"
}
}
}
}))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (executor_url, executor_handle) = start_server(executor).await;
let gateway =
build_router_with_endpoints(upstream_url.clone(), Some(upstream_url), Some(executor_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/upload/v1beta/files?uploadType=resumable"
))
.header(http::header::CONTENT_TYPE, "application/octet-stream")
.header(TRACE_ID_HEADER, "trace-files-upload-123")
.body("upload-bytes")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("json should parse"),
json!({"file": {"name": "files/uploaded-123"}})
);
let seen_plan_request = seen_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("plan-sync should be captured");
assert_eq!(seen_plan_request.method, "POST");
assert_eq!(seen_plan_request.path, "/upload/v1beta/files");
assert_eq!(
BASE64_STANDARD
.decode(seen_plan_request.body_base64)
.expect("body should decode"),
b"upload-bytes"
);
let seen_executor_request = seen_executor
.lock()
.expect("mutex should lock")
.clone()
.expect("executor sync should be captured");
assert_eq!(seen_executor_request.method, "POST");
assert_eq!(
seen_executor_request.url,
"https://files.example/upload/v1beta/files?uploadType=resumable"
);
assert_eq!(
BASE64_STANDARD
.decode(seen_executor_request.body_bytes_b64)
.expect("executor body should decode"),
b"upload-bytes"
);
let seen_report_request = seen_report
.lock()
.expect("mutex should lock")
.clone()
.expect("report-sync should be captured");
assert_eq!(
seen_report_request.report_kind,
"gemini_files_store_mapping"
);
assert_eq!(seen_report_request.file_key_id, "file-key-upload-123");
assert_eq!(seen_report_request.user_id, "user-files-upload-123");
assert_eq!(seen_report_request.file_name, "files/uploaded-123");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
executor_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_executes_gemini_files_delete_via_executor_sync_plan() {
#[derive(Debug, Clone)]
struct SeenExecutorSyncRequest {
method: String,
url: String,
}
#[derive(Debug, Clone)]
struct SeenReportSyncRequest {
report_kind: String,
file_name: String,
}
let seen_executor = Arc::new(Mutex::new(None::<SeenExecutorSyncRequest>));
let seen_executor_clone = Arc::clone(&seen_executor);
let seen_report = Arc::new(Mutex::new(None::<SeenReportSyncRequest>));
let seen_report_clone = Arc::clone(&seen_report);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "gemini",
"route_kind": "files",
"auth_endpoint_signature": "gemini:chat",
"executor_candidate": true,
"auth_context": {
"user_id": "user-files-delete-123",
"api_key_id": "key-files-delete-123",
"access_allowed": true
},
"public_path": "/v1beta/files/files/abc-123"
}))
}),
)
.route(
"/api/internal/gateway/plan-sync",
any(|_request: Request| async move {
Json(json!({
"action": "executor_sync",
"plan_kind": "gemini_files_delete",
"plan": {
"request_id": "req-files-delete-123",
"provider_id": "provider-files-delete-123",
"endpoint_id": "endpoint-files-delete-123",
"key_id": "file-key-delete-123",
"provider_name": "gemini",
"method": "DELETE",
"url": "https://files.example/v1beta/files/files/abc-123",
"headers": {
"x-goog-api-key": "upstream-key"
},
"body": {},
"stream": false,
"client_api_format": "gemini:files",
"provider_api_format": "gemini:files",
"model_name": "gemini-files"
},
"report_kind": "gemini_files_delete_mapping",
"report_context": {
"file_name": "files/abc-123"
}
}))
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |request: Request| {
let seen_report_inner = Arc::clone(&seen_report_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("report payload should parse");
*seen_report_inner.lock().expect("mutex should lock") =
Some(SeenReportSyncRequest {
report_kind: payload
.get("report_kind")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
file_name: payload
.get("report_context")
.and_then(|value| value.get("file_name"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({"ok": true}))
}
}),
)
.route(
"/v1beta/files/files/abc-123",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let executor = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_executor_inner = Arc::clone(&seen_executor_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("executor payload should parse");
*seen_executor_inner.lock().expect("mutex should lock") =
Some(SeenExecutorSyncRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "req-files-delete-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (executor_url, executor_handle) = start_server(executor).await;
let gateway =
build_router_with_endpoints(upstream_url.clone(), Some(upstream_url), Some(executor_url))
.expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!("{gateway_url}/v1beta/files/files/abc-123"))
.header(TRACE_ID_HEADER, "trace-files-delete-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("json should parse"),
json!({})
);
let seen_executor_request = seen_executor
.lock()
.expect("mutex should lock")
.clone()
.expect("executor sync should be captured");
assert_eq!(seen_executor_request.method, "DELETE");
assert_eq!(
seen_executor_request.url,
"https://files.example/v1beta/files/files/abc-123"
);
let seen_report_request = seen_report
.lock()
.expect("mutex should lock")
.clone()
.expect("report-sync should be captured");
assert_eq!(
seen_report_request.report_kind,
"gemini_files_delete_mapping"
);
assert_eq!(seen_report_request.file_name, "files/abc-123");
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
executor_handle.abort();
upstream_handle.abort();
}

View File

@@ -15,7 +15,7 @@ use crate::gateway::headers::{
};
use crate::gateway::{
build_client_response, maybe_execute_via_control, maybe_execute_via_executor_stream,
resolve_control_route, AppState, GatewayError,
maybe_execute_via_executor_sync, resolve_control_route, AppState, GatewayError,
};
pub(crate) async fn health(State(state): State<AppState>) -> impl IntoResponse {
@@ -132,6 +132,17 @@ pub(crate) async fn proxy_request(
let buffered_body = to_bytes(body, usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if let Some(executor_response) = maybe_execute_via_executor_sync(
&state,
&parts,
&buffered_body,
&trace_id,
control_decision.as_ref(),
)
.await?
{
return Ok(executor_response);
}
if let Some(executor_response) =
maybe_execute_via_executor_stream(&state, &parts, &trace_id, control_decision.as_ref())
.await?