Improve gateway scheduling and runtime admission

apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs

@@ -178,7 +178,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
 #[async_trait]
 impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        while let Some(attempt) = self.candidates.next_attempt().await? {
             match self.build_sync_attempt(attempt).await? {
                 Some(attempt) => return Ok(Some(attempt)),
                 None => continue,
@@ -206,7 +206,7 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSour
 #[async_trait]
 impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalStandardStreamAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        while let Some(attempt) = self.candidates.next_attempt().await? {
             match self.build_stream_attempt(attempt).await? {
                 Some(attempt) => return Ok(Some(attempt)),
                 None => continue,
@@ -340,7 +340,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
             .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )
@@ -390,7 +390,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
             .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )
@@ -449,7 +449,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
         return Ok(Vec::new());
     }
     let mut plans = Vec::new();
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )
@@ -524,7 +524,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
         return Ok(Vec::new());
     }
     let mut plans = Vec::new();
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs

@@ -12,6 +12,7 @@ use crate::ai_serving::planner::{
 use crate::ai_serving::transport::{
     resolve_transport_execution_timeouts, resolve_transport_profile,
 };
+use crate::stage_metrics::observe_gateway_stage_ms;
 use crate::{
     append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
     AiExecutionDecision, AppState, GatewayError,
@@ -40,6 +41,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
         candidate_id,
         ..
     } = attempt;
+    let payload_started_at = std::time::Instant::now();
     let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts(
         state,
         parts,
@@ -55,8 +57,16 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
     )
     .await?
     else {
+        observe_gateway_stage_ms(
+            "stream_candidate_payload_parts",
+            payload_started_at.elapsed().as_millis() as u64,
+        );
         return Ok(None);
     };
+    observe_gateway_stage_ms(
+        "stream_candidate_payload_parts",
+        payload_started_at.elapsed().as_millis() as u64,
+    );
     let candidate = &eligible.candidate;
 
     let prompt_cache_key = resolved
@@ -66,9 +76,14 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
         .map(str::trim)
         .filter(|value| !value.is_empty())
         .map(ToOwned::to_owned);
+    let proxy_started_at = std::time::Instant::now();
     let proxy = state
         .resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
         .await;
+    observe_gateway_stage_ms(
+        "stream_candidate_proxy",
+        proxy_started_at.elapsed().as_millis() as u64,
+    );
     let transport_profile = resolved
         .transport_profile
         .clone()
@@ -130,6 +145,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
         Some(body_json)
     };
     let effective_headers = input.effective_headers(&parts.headers);
+    let report_context_started_at = std::time::Instant::now();
     let report_context = append_local_failover_policy_to_value(
         append_execution_contract_fields_to_value(
             build_local_execution_report_context(LocalExecutionReportContextParts {
@@ -184,8 +200,13 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
         ),
         &transport,
     );
+    observe_gateway_stage_ms(
+        "stream_candidate_report_context",
+        report_context_started_at.elapsed().as_millis() as u64,
+    );
     let request_gzip = resolve_transport_request_gzip_policy(&transport);
 
+    let decision_started_at = std::time::Instant::now();
     let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
         decision_is_stream,
         decision_kind: decision_kind.to_string(),
@@ -222,5 +243,9 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
         auth_context: input.auth_context.clone(),
     });
     apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
+    observe_gateway_stage_ms(
+        "stream_candidate_decision_build",
+        decision_started_at.elapsed().as_millis() as u64,
+    );
     Ok(Some(decision))
 }

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs

@@ -54,6 +54,7 @@ use crate::ai_serving::{
     LocalResolvedOAuthRequestAuth,
 };
 use crate::ai_serving::{ConversionMode, ExecutionStrategy};
+use crate::stage_metrics::observe_gateway_stage_ms;
 use crate::{AppState, GatewayError};
 
 use super::support::{
@@ -102,6 +103,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
     report_kind: &str,
     upstream_is_stream: bool,
 ) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
+    let prepare_started_at = std::time::Instant::now();
     let planner_state = crate::ai_serving::PlannerAppState::new(state);
     let candidate = &eligible.candidate;
     let provider_api_format = eligible.provider_api_format.as_str();
@@ -109,6 +111,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
     let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
     let force_body_stream_field =
         endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
+    let model_directives_started_at = std::time::Instant::now();
     let enable_model_directives =
         crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
             state,
@@ -116,6 +119,11 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
             Some(&input.requested_model),
         )
         .await;
+    observe_gateway_stage_ms(
+        "openai_chat_payload_model_directives",
+        model_directives_started_at.elapsed().as_millis() as u64,
+    );
+    let redaction_started_at = std::time::Instant::now();
     let redaction = resolve_provider_chat_pii_redaction(
         state,
         parts,
@@ -125,6 +133,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
         candidate_id,
     )
     .await?;
+    observe_gateway_stage_ms(
+        "openai_chat_payload_redaction",
+        redaction_started_at.elapsed().as_millis() as u64,
+    );
     let body_json = redaction.body_json.as_ref();
     let effective_headers = input.effective_headers(&parts.headers);
     let is_grok = transport
@@ -234,7 +246,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
             redaction.redacted,
         );
 
-        return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
+        let result = Ok(Some(LocalOpenAiChatCandidatePayloadParts {
             client_api_format: "openai:chat".to_string(),
             auth_header: prepared_candidate.auth_header,
             auth_value: prepared_candidate.auth_value,
@@ -252,6 +264,11 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
             transport_profile,
             image_request_summary: None,
         }));
+        observe_gateway_stage_ms(
+            "openai_chat_payload_parts_prepare",
+            prepare_started_at.elapsed().as_millis() as u64,
+        );
+        return result;
     }
 
     if provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport) {
@@ -288,6 +305,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
             return Ok(None);
         };
 
+        let auth_prepare_started_at = std::time::Instant::now();
         let prepared_candidate = match prepare_header_authenticated_candidate(
             planner_state,
             transport,
@@ -316,7 +334,12 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
                 return Ok(None);
             }
         };
+        observe_gateway_stage_ms(
+            "openai_chat_payload_auth_prepare",
+            auth_prepare_started_at.elapsed().as_millis() as u64,
+        );
 
+        let body_build_started_at = std::time::Instant::now();
         let Some(mut provider_request_body) = build_local_openai_chat_request_body(
             body_json,
             &prepared_candidate.mapped_model,
@@ -343,6 +366,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
             .await;
             return Ok(None);
         };
+        observe_gateway_stage_ms(
+            "openai_chat_payload_body_build",
+            body_build_started_at.elapsed().as_millis() as u64,
+        );
         apply_deepseek_tool_call_thinking_compat(
             &mut provider_request_body,
             transport.provider.provider_type.as_str(),

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs

@@ -147,7 +147,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
     )
     .await;
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
             &attempt.eligible.transport,
             attempt.eligible.provider_api_format.as_str(),
@@ -199,7 +199,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
     )
     .await;
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
             &attempt.eligible.transport,
             attempt.eligible.provider_api_format.as_str(),

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs

@@ -18,6 +18,7 @@ use crate::ai_serving::planner::plan_builders::{
     build_openai_chat_stream_plan_from_decision, AiStreamAttempt,
 };
 use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
+use crate::stage_metrics::observe_gateway_stage_ms;
 
 pub(crate) struct LocalOpenAiChatStreamAttemptSource<'a> {
     state: &'a AppState,
@@ -93,10 +94,36 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
 #[async_trait]
 impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        loop {
+            let source_started_at = std::time::Instant::now();
+            let Some(attempt) = self.candidates.next_attempt().await? else {
+                observe_gateway_stage_ms(
+                    "stream_candidate_source_next",
+                    source_started_at.elapsed().as_millis() as u64,
+                );
+                break;
+            };
+            observe_gateway_stage_ms(
+                "stream_candidate_source_next",
+                source_started_at.elapsed().as_millis() as u64,
+            );
+
+            let plan_started_at = std::time::Instant::now();
             match self.build_stream_attempt(attempt).await? {
-                Some(attempt) => return Ok(Some(attempt)),
-                None => continue,
+                Some(attempt) => {
+                    observe_gateway_stage_ms(
+                        "stream_candidate_plan_build",
+                        plan_started_at.elapsed().as_millis() as u64,
+                    );
+                    return Ok(Some(attempt));
+                }
+                None => {
+                    observe_gateway_stage_ms(
+                        "stream_candidate_plan_build",
+                        plan_started_at.elapsed().as_millis() as u64,
+                    );
+                    continue;
+                }
             }
         }
         apply_local_runtime_candidate_terminal_reason(

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs

@@ -93,7 +93,7 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
 #[async_trait]
 impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiChatSyncAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        while let Some(attempt) = self.candidates.next_attempt().await? {
             match self.build_sync_attempt(attempt).await? {
                 Some(attempt) => return Ok(Some(attempt)),
                 None => continue,

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs

@@ -114,7 +114,7 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
     )
     .await?;
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )
@@ -153,7 +153,7 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
     )
     .await?;
 
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs

@@ -162,7 +162,7 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
 #[async_trait]
 impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        while let Some(attempt) = self.candidates.next_attempt().await? {
             match self.build_sync_attempt(attempt).await? {
                 Some(attempt) => return Ok(Some(attempt)),
                 None => continue,
@@ -190,7 +190,7 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAtte
 #[async_trait]
 impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiResponsesStreamAttemptSource<'_> {
     async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
-        while let Some(attempt) = self.candidates.next_attempt().await {
+        while let Some(attempt) = self.candidates.next_attempt().await? {
             match self.build_stream_attempt(attempt).await? {
                 Some(attempt) => return Ok(Some(attempt)),
                 None => continue,
@@ -331,7 +331,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
     }
 
     let mut plans = Vec::new();
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )
@@ -403,7 +403,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
     }
 
     let mut plans = Vec::new();
-    while let Some(attempt) = source.next_attempt().await {
+    while let Some(attempt) = source.next_attempt().await? {
         let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
             state, parts, trace_id, body_json, &input, attempt, spec,
         )