Improve request trace upstream diagnostics

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

@@ -1,23 +1,17 @@
 use std::collections::BTreeMap;
 
 use aether_contracts::{StreamFrame, StreamFramePayload};
-use axum::body::Body;
-use axum::http::Response;
+use axum::http::StatusCode;
 use base64::Engine as _;
 use futures_util::StreamExt;
-use serde_json::json;
+use serde_json::{json, Map, Value};
 use tokio_util::codec::{FramedRead, LinesCodec};
 use tracing::warn;
 
-use crate::api::response::build_client_response_from_parts;
-use crate::control::GatewayControlDecision;
 use crate::execution_runtime::ndjson::decode_stream_frame_ndjson;
 use crate::execution_runtime::submission::{has_nested_error, strip_utf8_bom_and_ws};
 use crate::GatewayError;
-use crate::{
-    GEMINI_FILES_DOWNLOAD_PLAN_KIND, MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES,
-    OPENAI_VIDEO_CONTENT_PLAN_KIND,
-};
+use crate::{MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_FRAMES};
 
 #[derive(Debug)]
 pub(super) enum StreamPrefetchInspection {
@@ -51,6 +45,88 @@ pub(super) fn decode_stream_error_body(
     )
 }
 
+fn header_value_case_insensitive<'a>(
+    headers: &'a BTreeMap<String, String>,
+    name: &str,
+) -> Option<&'a str> {
+    headers
+        .iter()
+        .find(|(key, _)| key.eq_ignore_ascii_case(name))
+        .map(|(_, value)| value.as_str())
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+}
+
+fn remove_header_case_insensitive(headers: &mut BTreeMap<String, String>, name: &str) {
+    let keys = headers
+        .keys()
+        .filter(|key| key.eq_ignore_ascii_case(name))
+        .cloned()
+        .collect::<Vec<_>>();
+    for key in keys {
+        headers.remove(&key);
+    }
+}
+
+pub(super) fn should_synthesize_non_success_stream_error_body(
+    status_code: u16,
+    error_body: &[u8],
+) -> bool {
+    !(200..300).contains(&status_code)
+        && ((300..400).contains(&status_code) || error_body.is_empty())
+}
+
+pub(super) fn build_synthetic_non_success_stream_error_body(
+    status_code: u16,
+    headers: &BTreeMap<String, String>,
+) -> Value {
+    let mut error = Map::from_iter([
+        (
+            "type".to_string(),
+            Value::String("execution_runtime_non_success_status".to_string()),
+        ),
+        (
+            "message".to_string(),
+            Value::String(format!(
+                "execution runtime stream returned non-success status {status_code}"
+            )),
+        ),
+        ("code".to_string(), Value::from(status_code)),
+        ("upstream_status".to_string(), Value::from(status_code)),
+    ]);
+    if let Some(location) = header_value_case_insensitive(headers, "location") {
+        error.insert("location".to_string(), Value::String(location.to_string()));
+    }
+
+    Value::Object(Map::from_iter([(
+        "error".to_string(),
+        Value::Object(error),
+    )]))
+}
+
+pub(super) fn synthetic_error_response_headers(
+    mut headers: BTreeMap<String, String>,
+) -> BTreeMap<String, String> {
+    remove_header_case_insensitive(&mut headers, "content-encoding");
+    remove_header_case_insensitive(&mut headers, "content-length");
+    remove_header_case_insensitive(&mut headers, "content-type");
+    remove_header_case_insensitive(&mut headers, "location");
+    headers.insert("content-type".to_string(), "application/json".to_string());
+    headers
+}
+
+fn client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
+    if (300..400).contains(&status_code) || status_code < 200 {
+        StatusCode::BAD_GATEWAY.as_u16()
+    } else {
+        status_code
+    }
+}
+
+pub(super) fn stream_client_error_status_code_for_upstream_status(status_code: u16) -> u16 {
+    client_error_status_code_for_upstream_status(status_code)
+}
+
 pub(super) fn inspect_prefetched_stream_body(
     headers: &BTreeMap<String, String>,
     body: &[u8],
@@ -166,62 +242,3 @@ where
     }
     Ok(None)
 }
-
-pub(super) fn build_execution_runtime_error_response(
-    trace_id: &str,
-    decision: &GatewayControlDecision,
-    plan_kind: &str,
-    status_code: u16,
-    headers: BTreeMap<String, String>,
-    error_body: Vec<u8>,
-) -> Result<Response<Body>, GatewayError> {
-    let content_type = headers
-        .get("content-type")
-        .map(|value| value.to_ascii_lowercase())
-        .unwrap_or_default();
-
-    if plan_kind == GEMINI_FILES_DOWNLOAD_PLAN_KIND && !content_type.starts_with("application/json")
-    {
-        let wrapped = serde_json::to_vec(&json!({
-            "error": String::from_utf8_lossy(&error_body).to_string(),
-        }))
-        .map_err(|err| GatewayError::Internal(err.to_string()))?;
-        let wrapped_headers =
-            BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
-        return build_client_response_from_parts(
-            status_code,
-            &wrapped_headers,
-            Body::from(wrapped),
-            trace_id,
-            Some(decision),
-        );
-    }
-
-    if plan_kind == OPENAI_VIDEO_CONTENT_PLAN_KIND && !content_type.starts_with("application/json")
-    {
-        let wrapped = serde_json::to_vec(&json!({
-            "error": {
-                "type": "upstream_error",
-                "message": "Video not available",
-            }
-        }))
-        .map_err(|err| GatewayError::Internal(err.to_string()))?;
-        let wrapped_headers =
-            BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
-        return build_client_response_from_parts(
-            status_code,
-            &wrapped_headers,
-            Body::from(wrapped),
-            trace_id,
-            Some(decision),
-        );
-    }
-
-    build_client_response_from_parts(
-        status_code,
-        &headers,
-        Body::from(error_body),
-        trace_id,
-        Some(decision),
-    )
-}