mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -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")]
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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?
|
||||
|
||||
Reference in New Issue
Block a user