Improve gateway transport and usage runtime

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

@@ -15,6 +15,7 @@ use base64::Engine as _;
 use futures_util::stream::{self, BoxStream};
 use futures_util::StreamExt;
 use http::{HeaderMap, HeaderName, HeaderValue};
+use http_body_util::BodyExt;
 use regex::{Captures, Regex};
 use serde_json::{json, Map, Value};
 use uuid::Uuid;
@@ -29,10 +30,10 @@ use crate::clock::current_unix_secs;
 use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
 use crate::execution_runtime::transport::{
     build_browser_wreq_client, build_request_body, build_request_headers,
-    decode_response_body_bytes, format_upstream_request_error, format_wreq_upstream_request_error,
-    resolve_stream_first_byte_timeout, send_request, stream_first_byte_timeout_message,
-    with_non_stream_total_timeout, DirectHttpResponse, ExecutionRuntimeTransportError,
-    ExecutionTransportControls,
+    decode_response_body_bytes, format_hyper_error_chain, format_upstream_request_error,
+    format_wreq_upstream_request_error, resolve_stream_first_byte_timeout, send_request,
+    stream_first_byte_timeout_message, with_non_stream_total_timeout, DirectHttpResponse,
+    ExecutionRuntimeTransportError, ExecutionTransportControls,
 };
 
 const GROK_INTERNAL_HEADER: &str = "x-aether-grok-runtime";
@@ -503,6 +504,15 @@ async fn collect_grok_response_stream(
                 collect_grok_response_chunk(status_code, upstream_bytes, raw_body, adapter, &chunk);
             }
         }
+        DirectHttpResponse::HyperH2c(response) => {
+            let mut stream = response.into_body().into_data_stream();
+            while let Some(chunk) = stream.next().await {
+                let chunk = chunk.map_err(|err| {
+                    ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(&err))
+                })?;
+                collect_grok_response_chunk(status_code, upstream_bytes, raw_body, adapter, &chunk);
+            }
+        }
         DirectHttpResponse::BrowserWreq(response) => {
             let mut stream = response.bytes_stream();
             while let Some(chunk) = stream.next().await {
@@ -547,6 +557,16 @@ fn grok_response_body_stream(response: DirectHttpResponse) -> GrokUpstreamBodySt
                 })
             })
             .boxed(),
+        DirectHttpResponse::HyperH2c(response) => response
+            .into_body()
+            .into_data_stream()
+            .map(|chunk| {
+                chunk.map_err(|err| {
+                    ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(&err))
+                        .to_string()
+                })
+            })
+            .boxed(),
         DirectHttpResponse::BrowserWreq(response) => response
             .bytes_stream()
             .map(|chunk| {

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

@@ -48,6 +48,15 @@ pub use server::{
     build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,
     serve_execution_runtime_unix,
 };
+pub use transport::DirectH2cSenderPrewarmReport;
+
+pub async fn prewarm_direct_h2c_sender_cache_from_env_for_startup(
+) -> Result<Option<DirectH2cSenderPrewarmReport>, String> {
+    transport::prewarm_direct_h2c_sender_cache_from_env()
+        .await
+        .map_err(|err| err.to_string())
+}
+
 pub(crate) use stream::execute_execution_runtime_stream;
 pub(crate) use stream_pump::build_direct_execution_frame_stream;
 pub(crate) use sync::{

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

@@ -11,6 +11,7 @@ use async_stream::stream;
 use axum::body::Bytes;
 use base64::Engine as _;
 use futures_util::{Stream, StreamExt};
+use http_body_util::BodyExt;
 use serde_json::Value;
 use tracing::warn;
 
@@ -20,7 +21,8 @@ use crate::ai_serving::api::{
 };
 use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
 use crate::execution_runtime::transport::{
-    format_wreq_upstream_request_error, stream_first_byte_timeout_message, DirectUpstreamResponse,
+    format_hyper_error_chain, format_wreq_upstream_request_error,
+    stream_first_byte_timeout_message, DirectUpstreamResponse,
 };
 use crate::execution_runtime::DirectUpstreamStreamExecution;
 use crate::GatewayError;
@@ -273,6 +275,88 @@ pub(crate) fn build_direct_execution_frame_stream(
                     }
                 }
             }
+            DirectUpstreamResponse::HyperH2c(response) => {
+                let mut bytes_stream = response.into_body().into_data_stream();
+                loop {
+                    let item = if ttfb_ms.is_none() {
+                        match await_stream_first_byte(
+                            bytes_stream.next(),
+                            started_at,
+                            stream_first_byte_timeout,
+                        )
+                        .await
+                        {
+                            Ok(item) => item,
+                            Err(timeout) => {
+                                match encode_first_byte_timeout_frame(timeout) {
+                                    Ok(frame) => yield Ok(frame),
+                                    Err(err) => {
+                                        yield Err(err);
+                                        return;
+                                    }
+                                }
+                                break;
+                            }
+                        }
+                    } else {
+                        bytes_stream.next().await
+                    };
+                    let Some(item) = item else {
+                        break;
+                    };
+                    match item {
+                        Ok(chunk) => {
+                            if ttfb_ms.is_none() {
+                                ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
+                            }
+                            if !first_chunk_telemetry_emitted {
+                                match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
+                                    Ok(frame) => yield Ok(frame),
+                                    Err(err) => {
+                                        yield Err(err);
+                                        return;
+                                    }
+                                }
+                                first_chunk_telemetry_emitted = true;
+                            }
+                            upstream_bytes += chunk.len() as u64;
+                            observe_stream_chunk(
+                                &mut stream_terminal_observer,
+                                &normalized_observer_context,
+                                private_stream_normalizer.as_mut(),
+                                &mut observer_buffered,
+                                chunk.as_ref(),
+                            );
+                            match encode_data_frame(&chunk) {
+                                Ok(frame) => yield Ok(frame),
+                                Err(err) => {
+                                    yield Err(err);
+                                    return;
+                                }
+                            }
+                        }
+                        Err(err) => {
+                            let message = format_hyper_error_chain(&err);
+                            warn!(
+                                event_name = "stream_pump_body_read_error",
+                                log_type = "ops",
+                                status_code,
+                                upstream_bytes,
+                                error = %message,
+                                "upstream body stream read error"
+                            );
+                            match encode_error_frame(status_code, message) {
+                                Ok(frame) => yield Ok(frame),
+                                Err(encode_err) => {
+                                    yield Err(encode_err);
+                                    return;
+                                }
+                            }
+                            break;
+                        }
+                    }
+                }
+            }
             DirectUpstreamResponse::BrowserWreq(response) => {
                 let mut bytes_stream = response.bytes_stream();
                 loop {
@@ -662,6 +746,60 @@ async fn buffer_non_sse_upstream_body(
                 }
             }
         }
+        DirectUpstreamResponse::HyperH2c(response) => {
+            let mut bytes_stream = response.into_body().into_data_stream();
+            loop {
+                let item = if ttfb_ms.is_none() {
+                    match await_stream_first_byte(
+                        bytes_stream.next(),
+                        started_at,
+                        stream_first_byte_timeout,
+                    )
+                    .await
+                    {
+                        Ok(item) => item,
+                        Err(timeout) => {
+                            return Err(BufferedUpstreamBodyError {
+                                message: stream_first_byte_timeout_message(timeout),
+                                ttfb_ms,
+                                upstream_bytes,
+                                first_byte_timeout: Some(timeout),
+                            });
+                        }
+                    }
+                } else {
+                    bytes_stream.next().await
+                };
+                let Some(item) = item else {
+                    break;
+                };
+                match item {
+                    Ok(chunk) => {
+                        if ttfb_ms.is_none() {
+                            ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
+                        }
+                        upstream_bytes += chunk.len() as u64;
+                        body_bytes.extend_from_slice(&chunk);
+                    }
+                    Err(err) => {
+                        let message = format_hyper_error_chain(&err);
+                        warn!(
+                            event_name = "stream_pump_body_read_error",
+                            log_type = "ops",
+                            upstream_bytes,
+                            error = %message,
+                            "upstream body stream read error"
+                        );
+                        return Err(BufferedUpstreamBodyError {
+                            message,
+                            ttfb_ms,
+                            upstream_bytes,
+                            first_byte_timeout: None,
+                        });
+                    }
+                }
+            }
+        }
         DirectUpstreamResponse::BrowserWreq(response) => {
             let mut bytes_stream = response.bytes_stream();
             loop {

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

@@ -22,6 +22,7 @@ use axum::body::{to_bytes, Body, Bytes};
 use axum::http::header::{CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE};
 use axum::http::{HeaderName, HeaderValue, Response, StatusCode};
 use futures_util::StreamExt;
+use http_body_util::BodyExt;
 use serde_json::{json, Value};
 use tokio::sync::mpsc;
 use tokio::sync::Mutex;
@@ -55,9 +56,9 @@ use crate::execution_runtime::submission::{
 };
 use crate::execution_runtime::transport::{
     build_execution_response_body, build_request_body, collect_response_headers,
-    decode_response_body_bytes, format_upstream_request_error, format_wreq_upstream_request_error,
-    response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
-    ExecutionRuntimeTransportError,
+    decode_response_body_bytes, format_hyper_error_chain, format_upstream_request_error,
+    format_wreq_upstream_request_error, response_body_is_json, send_request, DirectHttpResponse,
+    DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
 };
 use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
 use crate::execution_runtime::{
@@ -1165,6 +1166,23 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
                 body_bytes.extend_from_slice(&chunk);
             }
         }
+        DirectHttpResponse::HyperH2c(response) => {
+            let mut upstream_stream = response.into_body().into_data_stream();
+            while let Some(chunk) = upstream_stream.next().await {
+                let chunk = chunk.map_err(|err| {
+                    SyncExecutionFailure::from_transport(
+                        ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(
+                            &err,
+                        )),
+                    )
+                })?;
+                let elapsed_ms = started_at.elapsed().as_millis() as u64;
+                progress
+                    .observe_chunk(&chunk, status_code, elapsed_ms)
+                    .await;
+                body_bytes.extend_from_slice(&chunk);
+            }
+        }
         DirectHttpResponse::BrowserWreq(response) => {
             let mut upstream_stream = response.bytes_stream();
             while let Some(chunk) = upstream_stream.next().await {
@@ -1522,7 +1540,8 @@ async fn execute_execution_runtime_sync_impl(
     let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
     state
         .usage_runtime
-        .record_pending(state.data.as_ref(), lifecycle_seed);
+        .record_pending_direct(state.data.as_ref(), lifecycle_seed)
+        .await;
     record_local_request_candidate_status(
         state,
         &plan,