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
@@ -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| {