Improve gateway transport and usage runtime

This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
+183 -3
View File
@@ -9,7 +9,11 @@ use std::collections::BTreeMap;
use crate::ai_serving::api::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
supports_stream_execution_decision_kind, AiStreamAttempt, OPENAI_VIDEO_CONTENT_PLAN_KIND,
supports_stream_execution_decision_kind, AiStreamAttempt, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CLI_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
use crate::api::response::build_client_response_from_parts;
use crate::control::GatewayControlDecision;
@@ -34,24 +38,78 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
trace_id: &str,
decision: &GatewayControlDecision,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_kind_started_at = std::time::Instant::now();
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
observe_gateway_stage_ms(
"frontdoor_stream_plan_kind",
plan_kind_started_at.elapsed().as_millis() as u64,
);
return Ok(LocalExecutionRequestOutcome::NoPath);
};
observe_gateway_stage_ms(
"frontdoor_stream_plan_kind",
plan_kind_started_at.elapsed().as_millis() as u64,
);
let parse_started_at = std::time::Instant::now();
let Some((body_json, body_base64)) = parse_local_request_body(parts, body_bytes) else {
observe_gateway_stage_ms(
"frontdoor_stream_parse",
parse_started_at.elapsed().as_millis() as u64,
);
return Ok(LocalExecutionRequestOutcome::NoPath);
};
observe_gateway_stage_ms(
"frontdoor_stream_parse",
parse_started_at.elapsed().as_millis() as u64,
);
if !is_matching_stream_request(plan_kind, parts, &body_json, body_base64.as_deref()) {
let match_started_at = std::time::Instant::now();
let stream_matches =
is_matching_stream_request(plan_kind, parts, &body_json, body_base64.as_deref());
observe_gateway_stage_ms(
"frontdoor_stream_match",
match_started_at.elapsed().as_millis() as u64,
);
if !stream_matches {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
let bypass_started_at = std::time::Instant::now();
let bypass_cache_key =
build_direct_plan_bypass_cache_key(plan_kind, parts, body_bytes, decision);
if should_skip_direct_plan(state, &bypass_cache_key) {
let skip_direct_plan = should_skip_direct_plan(state, &bypass_cache_key);
observe_gateway_stage_ms(
"frontdoor_stream_bypass",
bypass_started_at.elapsed().as_millis() as u64,
);
if skip_direct_plan {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
if plan_kind == OPENAI_CHAT_STREAM_PLAN_KIND
&& supports_stream_execution_decision_kind(plan_kind)
&& decision.route_family.as_deref() == Some("openai")
{
let fast_path_started_at = std::time::Instant::now();
let outcome = execute_openai_chat_stream_fast_path(
state,
parts,
trace_id,
decision,
&body_json,
body_base64,
plan_kind,
bypass_cache_key,
)
.await;
observe_gateway_stage_ms(
"frontdoor_stream_fast_path_total",
fast_path_started_at.elapsed().as_millis() as u64,
);
return outcome;
}
let port = GatewayStreamExecutionPathPort {
state,
parts,
@@ -69,6 +127,64 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
))
}
#[allow(clippy::too_many_arguments)]
async fn execute_openai_chat_stream_fast_path(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
body_base64: Option<String>,
plan_kind: &str,
bypass_cache_key: String,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let started_at = std::time::Instant::now();
let local_outcome = maybe_execute_stream_via_local_decision(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?;
observe_gateway_stage_ms(
"stream_openai_chat_local_decision",
started_at.elapsed().as_millis() as u64,
);
match local_outcome {
LocalExecutionRequestOutcome::Responded(response) => {
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
LocalExecutionRequestOutcome::Exhausted(outcome) => {
return Ok(LocalExecutionRequestOutcome::Exhausted(outcome));
}
LocalExecutionRequestOutcome::NoPath => {}
}
if let Some(response) = maybe_execute_stream_via_remote_decision(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
{
return Ok(LocalExecutionRequestOutcome::Responded(response));
}
let fallback_started_at = std::time::Instant::now();
let fallback_outcome = maybe_execute_stream_via_plan_fallback(
state,
parts,
trace_id,
decision,
body_json,
body_base64,
plan_kind,
bypass_cache_key,
GatewayFallbackReason::RemoteDecisionMiss,
)
.await?;
observe_gateway_stage_ms(
"frontdoor_stream_fast_path",
fallback_started_at.elapsed().as_millis() as u64,
);
Ok(fallback_outcome)
}
struct GatewayStreamExecutionPathPort<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
@@ -91,6 +207,10 @@ impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> {
self.scheduler_supported
}
fn stream_execution_steps(&self) -> &'static [AiStreamExecutionStep] {
stream_execution_steps_for_plan_kind(self.plan_kind)
}
async fn execute_stream_step(
&self,
step: AiStreamExecutionStep,
@@ -194,6 +314,10 @@ impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> {
"stream_path_step",
step_started_at.elapsed().as_millis() as u64,
);
observe_gateway_stage_ms(
stream_path_stage_name(step),
step_started_at.elapsed().as_millis() as u64,
);
Ok(to_ai_serving_outcome(outcome))
}
@@ -217,6 +341,62 @@ impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> {
}
}
const STREAM_STEPS_VIDEO_CONTENT: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalVideoContent,
AiStreamExecutionStep::RemoteDecision,
];
const STREAM_STEPS_OPENAI_IMAGE: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalImage,
AiStreamExecutionStep::RemoteDecision,
];
const STREAM_STEPS_OPENAI_CHAT: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalOpenAiChat,
AiStreamExecutionStep::RemoteDecision,
];
const STREAM_STEPS_OPENAI_RESPONSES: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalOpenAiResponses,
AiStreamExecutionStep::RemoteDecision,
];
const STREAM_STEPS_STANDARD_TEXT: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalStandardFamily,
AiStreamExecutionStep::LocalSameFormatProvider,
AiStreamExecutionStep::RemoteDecision,
];
const STREAM_STEPS_GEMINI_FILES: &[AiStreamExecutionStep] = &[
AiStreamExecutionStep::LocalGeminiFiles,
AiStreamExecutionStep::RemoteDecision,
];
fn stream_execution_steps_for_plan_kind(plan_kind: &str) -> &'static [AiStreamExecutionStep] {
match plan_kind {
OPENAI_VIDEO_CONTENT_PLAN_KIND => STREAM_STEPS_VIDEO_CONTENT,
OPENAI_IMAGE_STREAM_PLAN_KIND => STREAM_STEPS_OPENAI_IMAGE,
OPENAI_CHAT_STREAM_PLAN_KIND => STREAM_STEPS_OPENAI_CHAT,
OPENAI_RESPONSES_STREAM_PLAN_KIND | OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND => {
STREAM_STEPS_OPENAI_RESPONSES
}
CLAUDE_CHAT_STREAM_PLAN_KIND
| CLAUDE_CLI_STREAM_PLAN_KIND
| GEMINI_CHAT_STREAM_PLAN_KIND
| GEMINI_CLI_STREAM_PLAN_KIND => STREAM_STEPS_STANDARD_TEXT,
GEMINI_FILES_DOWNLOAD_PLAN_KIND => STREAM_STEPS_GEMINI_FILES,
_ => aether_ai_serving::DEFAULT_STREAM_EXECUTION_STEPS,
}
}
fn stream_path_stage_name(step: AiStreamExecutionStep) -> &'static str {
match step {
AiStreamExecutionStep::LocalVideoContent => "stream_path_step_video_content",
AiStreamExecutionStep::LocalImage => "stream_path_step_image",
AiStreamExecutionStep::LocalOpenAiChat => "stream_path_step_openai_chat",
AiStreamExecutionStep::LocalOpenAiResponses => "stream_path_step_openai_responses",
AiStreamExecutionStep::LocalStandardFamily => "stream_path_step_standard_family",
AiStreamExecutionStep::LocalSameFormatProvider => "stream_path_step_same_format_provider",
AiStreamExecutionStep::LocalGeminiFiles => "stream_path_step_gemini_files",
AiStreamExecutionStep::RemoteDecision => "stream_path_step_remote_decision",
}
}
fn to_ai_serving_outcome(
outcome: LocalExecutionRequestOutcome,
) -> AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion> {