Implement independent provider pool scheduling runtime

apps/aether-gateway/src/execution_runtime/mod.rs

@@ -6,6 +6,7 @@ use serde_json::{Map, Value};
 mod constants;
 mod fallback;
 pub(crate) mod ndjson;
+mod pool_feedback;
 #[cfg(test)]
 pub(crate) mod remote_compat;
 mod server;
@@ -28,6 +29,10 @@ pub(crate) use self::fallback::{
     should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
     LocalFailoverDecision,
 };
+pub(crate) use pool_feedback::{
+    record_pool_error_feedback, record_pool_stream_timeout_feedback,
+    record_stream_pool_success_feedback, record_sync_pool_success_feedback,
+};
 pub use server::{
     build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
     build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,

apps/aether-gateway/src/execution_runtime/pool_feedback.rs

@@ -0,0 +1,188 @@
+use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
+use aether_usage_runtime::{
+    build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome, TerminalUsageOutcome,
+};
+use serde_json::Value;
+use std::collections::BTreeMap;
+use tracing::warn;
+
+use crate::ai_pipeline::extract_pool_sticky_session_token;
+use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
+use crate::handlers::shared::provider_pool::{
+    record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
+    record_admin_provider_pool_success, AdminProviderPoolConfig,
+};
+use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
+use crate::AppState;
+
+struct PoolFeedbackContext {
+    runner: aether_data::redis::RedisKvRunner,
+    pool_config: AdminProviderPoolConfig,
+    sticky_session_token: Option<String>,
+}
+
+fn pool_feedback_request_body<'a>(
+    plan: &'a ExecutionPlan,
+    report_context: Option<&'a Value>,
+) -> Option<&'a Value> {
+    report_context
+        .and_then(Value::as_object)
+        .and_then(|object| object.get("original_request_body"))
+        .filter(|value| !value.is_null())
+        .or(plan.body.json_body.as_ref())
+}
+
+async fn resolve_pool_feedback_context(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+) -> Option<PoolFeedbackContext> {
+    let Some(runner) = state.redis_kv_runner() else {
+        return None;
+    };
+
+    let transport = match state
+        .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
+        .await
+    {
+        Ok(Some(transport)) => transport,
+        Ok(None) => return None,
+        Err(err) => {
+            warn!(
+                "gateway execution runtime pool feedback: failed to read transport snapshot for provider {} endpoint {} key {}: {:?}",
+                plan.provider_id, plan.endpoint_id, plan.key_id, err
+            );
+            return None;
+        }
+    };
+
+    let Some(pool_config) =
+        admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
+    else {
+        return None;
+    };
+
+    let sticky_session_token = pool_feedback_request_body(plan, report_context)
+        .and_then(extract_pool_sticky_session_token);
+
+    Some(PoolFeedbackContext {
+        runner,
+        pool_config,
+        sticky_session_token,
+    })
+}
+
+fn total_tokens_used(outcome: &TerminalUsageOutcome) -> u64 {
+    outcome
+        .standardized_usage
+        .as_ref()
+        .map(|usage| {
+            usage
+                .input_tokens
+                .saturating_add(usage.output_tokens)
+                .max(0) as u64
+        })
+        .unwrap_or(0)
+}
+
+fn resolve_ttfb_ms(telemetry: Option<&ExecutionTelemetry>) -> Option<u64> {
+    telemetry.and_then(|telemetry| telemetry.ttfb_ms.or(telemetry.elapsed_ms))
+}
+
+pub(crate) async fn record_sync_pool_success_feedback(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+    payload: &GatewaySyncReportRequest,
+) {
+    let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
+        return;
+    };
+
+    let usage_outcome = build_sync_terminal_usage_outcome(plan, report_context, payload);
+    record_admin_provider_pool_success(
+        &context.runner,
+        &plan.provider_id,
+        &plan.key_id,
+        &context.pool_config,
+        context.sticky_session_token.as_deref(),
+        total_tokens_used(&usage_outcome),
+        resolve_ttfb_ms(payload.telemetry.as_ref()),
+    )
+    .await;
+}
+
+pub(crate) async fn record_stream_pool_success_feedback(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+    payload: &GatewayStreamReportRequest,
+) {
+    let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
+        return;
+    };
+
+    let usage_outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
+    record_admin_provider_pool_success(
+        &context.runner,
+        &plan.provider_id,
+        &plan.key_id,
+        &context.pool_config,
+        context.sticky_session_token.as_deref(),
+        total_tokens_used(&usage_outcome),
+        resolve_ttfb_ms(payload.telemetry.as_ref()),
+    )
+    .await;
+}
+
+pub(crate) async fn record_pool_error_feedback(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+    status_code: u16,
+    headers: &BTreeMap<String, String>,
+    error_body: Option<&str>,
+) {
+    if status_code < 400 {
+        return;
+    }
+
+    let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
+        return;
+    };
+
+    if status_code == 401 {
+        let _ = state
+            .invalidate_local_oauth_refresh_entry(&plan.key_id)
+            .await;
+    }
+
+    record_admin_provider_pool_error(
+        &context.runner,
+        &plan.provider_id,
+        &plan.key_id,
+        &context.pool_config,
+        status_code,
+        error_body,
+        Some(headers),
+    )
+    .await;
+}
+
+pub(crate) async fn record_pool_stream_timeout_feedback(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+) {
+    let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
+        return;
+    };
+
+    record_admin_provider_pool_stream_timeout(
+        &context.runner,
+        &plan.provider_id,
+        &plan.key_id,
+        &context.pool_config,
+    )
+    .await;
+}

apps/aether-gateway/src/execution_runtime/stream/execution.rs

@@ -53,7 +53,8 @@ use crate::execution_runtime::transport::{
     DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
 };
 use crate::execution_runtime::{
-    local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
+    local_failover_response_text, record_pool_error_feedback, record_stream_pool_success_feedback,
+    resolve_core_stream_direct_finalize_report_kind,
     resolve_core_stream_error_finalize_report_kind,
     resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
     should_retry_next_local_candidate_stream, LocalFailoverDecision,
@@ -565,6 +566,15 @@ async fn execute_stream_from_frame_stream(
         let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
         let error_response_text =
             local_failover_response_text(body_json.as_ref(), &error_body, None);
+        record_pool_error_feedback(
+            state,
+            &plan,
+            report_context.as_ref(),
+            status_code,
+            &headers,
+            error_response_text.as_deref(),
+        )
+        .await;
         let failover_decision = resolve_local_candidate_failover_decision_stream(
             state,
             &plan,
@@ -1492,6 +1502,13 @@ async fn execute_stream_from_frame_stream(
             }),
             telemetry: telemetry.clone(),
         };
+        record_stream_pool_success_feedback(
+            &state_for_report,
+            &plan_for_report,
+            report_context_owned.as_ref(),
+            &usage_payload,
+        )
+        .await;
         record_stream_terminal_usage(
             &state_for_report,
             &plan_for_report,

apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs

@@ -16,6 +16,7 @@ use crate::control::GatewayControlDecision;
 use crate::execution_runtime::submission::{
     resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
 };
+use crate::execution_runtime::{record_pool_error_feedback, record_pool_stream_timeout_feedback};
 use crate::log_ids::short_request_id;
 use crate::request_candidate_runtime::record_report_request_candidate_status;
 use crate::usage::submit_sync_report;
@@ -122,6 +123,22 @@ async fn record_stream_sync_failure(
     failure: &StreamFailureReport,
     started_at_unix_ms: Option<u64>,
 ) {
+    if matches!(
+        failure.error_type.as_str(),
+        "first_byte_timeout" | "read_timeout"
+    ) {
+        record_pool_stream_timeout_feedback(state, plan, report_context).await;
+    }
+    let error_body = serde_json::to_string(&failure.body_json).ok();
+    record_pool_error_feedback(
+        state,
+        plan,
+        report_context,
+        failure.status_code,
+        &payload.headers,
+        error_body.as_deref(),
+    )
+    .await;
     let context_seed = build_terminal_usage_context_seed(plan, report_context);
     let payload_seed = build_sync_terminal_usage_payload_seed(payload);
     state

apps/aether-gateway/src/execution_runtime/sync/execution.rs

@@ -30,9 +30,10 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
 use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
 use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
 use crate::execution_runtime::{
-    local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
-    should_fallback_to_control_sync, should_finalize_sync_response,
-    should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_sync,
+    local_failover_response_text, record_pool_error_feedback, record_sync_pool_success_feedback,
+    resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
+    should_finalize_sync_response, should_retry_next_local_candidate_sync,
+    should_stop_local_candidate_failover_sync,
 };
 use crate::log_ids::short_request_id;
 use crate::request_candidate_runtime::{
@@ -255,6 +256,17 @@ pub(crate) async fn execute_execution_runtime_sync(
         local_failover_response_text.as_deref(),
     )
     .await;
+    if result.status_code >= 400 {
+        record_pool_error_feedback(
+            state,
+            &plan,
+            report_context.as_ref(),
+            result.status_code,
+            &headers,
+            local_failover_response_text.as_deref(),
+        )
+        .await;
+    }
     if should_retry_next_local_candidate_sync(
         state,
         &plan,
@@ -393,6 +405,15 @@ pub(crate) async fn execute_execution_runtime_sync(
         body_base64: body_base64.clone(),
         telemetry: result.telemetry.clone(),
     };
+    if result.status_code < 400 {
+        record_sync_pool_success_feedback(
+            state,
+            &plan,
+            report_context.as_ref(),
+            &base_usage_payload,
+        )
+        .await;
+    }
 
     if let Some(finalize_report_kind) = finalize_report_kind {
         if let Some(implicit_finalize) = implicit_finalize {