fix(usage): track OpenAI image SSE completion and usage estimates

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

@@ -3,19 +3,20 @@ use std::io::Error as IoError;
 use std::time::Instant;
 
 use aether_contracts::{
-    ExecutionPlan, ExecutionResult, ExecutionTelemetry, RequestBody, ResponseBody, StreamFrame,
-    StreamFramePayload, StreamFrameType, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
-    EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
+    ExecutionPlan, ExecutionResult, ExecutionStreamTerminalSummary, ExecutionTelemetry,
+    RequestBody, ResponseBody, StreamFrame, StreamFramePayload, StreamFrameType,
+    EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
 };
 use axum::body::Bytes;
 use base64::Engine as _;
 use chrono::{FixedOffset, Utc};
 use futures_util::stream::{self, BoxStream};
 use futures_util::StreamExt;
-use serde_json::{json, Value};
+use serde_json::{json, Map, Value};
 use tracing::debug;
 use uuid::Uuid;
 
+use crate::ai_serving::api::StreamingStandardTerminalObserver;
 use crate::clock::current_unix_secs;
 use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
 use crate::execution_runtime::transport::{
@@ -109,7 +110,7 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
         Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
     };
     Ok(Some(ChatGptWebImageStream {
-        frame_stream: execution_result_frame_stream(&result),
+        frame_stream: execution_result_frame_stream(plan, &result, report_context),
         report_context: report_context.cloned(),
     }))
 }
@@ -1479,9 +1480,12 @@ fn bytes_execution_result(
 }
 
 fn execution_result_frame_stream(
+    plan: &ExecutionPlan,
     result: &ExecutionResult,
+    report_context: Option<&Value>,
 ) -> BoxStream<'static, Result<Bytes, IoError>> {
     let body = execution_result_body_bytes_lossy(result);
+    let terminal_summary = chatgpt_web_stream_terminal_summary(plan, result, report_context, &body);
     let mut frames = vec![
         StreamFrame {
             frame_type: StreamFrameType::Headers,
@@ -1520,7 +1524,7 @@ fn execution_result_frame_stream(
             }),
         },
     });
-    frames.push(StreamFrame::eof());
+    frames.push(StreamFrame::eof_with_summary(terminal_summary));
     stream::iter(
         frames
             .into_iter()
@@ -1529,6 +1533,81 @@ fn execution_result_frame_stream(
     .boxed()
 }
 
+fn chatgpt_web_stream_terminal_summary(
+    plan: &ExecutionPlan,
+    result: &ExecutionResult,
+    report_context: Option<&Value>,
+    body: &[u8],
+) -> Option<ExecutionStreamTerminalSummary> {
+    if !(200..300).contains(&result.status_code) || body.is_empty() {
+        return None;
+    }
+
+    let observer_context = chatgpt_web_stream_observer_context(plan, report_context);
+    let mut observer = StreamingStandardTerminalObserver::default();
+    let mut line_start = 0usize;
+    for (index, byte) in body.iter().enumerate() {
+        if *byte != b'\n' {
+            continue;
+        }
+        observer
+            .push_line(&observer_context, body[line_start..=index].to_vec())
+            .ok()?;
+        line_start = index.saturating_add(1);
+    }
+    if line_start < body.len() {
+        observer
+            .push_line(&observer_context, body[line_start..].to_vec())
+            .ok()?;
+    }
+    observer.finish(&observer_context).ok().flatten()
+}
+
+fn chatgpt_web_stream_observer_context(
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+) -> Value {
+    let mut context = report_context
+        .cloned()
+        .filter(Value::is_object)
+        .unwrap_or_else(|| json!({}));
+    let object = context
+        .as_object_mut()
+        .expect("observer context should be an object");
+    object
+        .entry("provider_api_format".to_string())
+        .or_insert_with(|| Value::String(plan.provider_api_format.clone()));
+    object
+        .entry("client_api_format".to_string())
+        .or_insert_with(|| Value::String(plan.client_api_format.clone()));
+    object
+        .entry("model".to_string())
+        .or_insert_with(|| Value::String(plan.model_name.clone().unwrap_or_default()));
+    if !object.contains_key("image_request") {
+        if let Some(image_request) = chatgpt_web_image_request_context(plan) {
+            object.insert("image_request".to_string(), image_request);
+        }
+    }
+    context
+}
+
+fn chatgpt_web_image_request_context(plan: &ExecutionPlan) -> Option<Value> {
+    let body = plan.body.json_body.as_ref()?.as_object()?;
+    let mut image_request = Map::new();
+    image_request.insert(
+        "operation".to_string(),
+        Value::String("generate".to_string()),
+    );
+    for key in ["model", "size", "quality", "output_format"] {
+        if let Some(value) = body.get(key).and_then(Value::as_str).map(str::trim) {
+            if !value.is_empty() {
+                image_request.insert(key.to_string(), Value::String(value.to_string()));
+            }
+        }
+    }
+    Some(Value::Object(image_request))
+}
+
 fn telemetry(started_at: Instant, upstream_bytes: u64) -> ExecutionTelemetry {
     let elapsed_ms = started_at.elapsed().as_millis() as u64;
     ExecutionTelemetry {
@@ -2225,6 +2304,31 @@ data: [DONE]
         assert!(decoded_data.contains("\"width\":2"));
         assert!(decoded_data.contains("\"height\":3"));
         assert!(text.contains("\"type\":\"eof\""));
+        let eof_frame = text
+            .lines()
+            .filter_map(|line| serde_json::from_str::<Value>(line).ok())
+            .find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
+            .expect("eof frame should exist");
+        assert_eq!(
+            eof_frame
+                .get("payload")
+                .and_then(|payload| payload.get("summary"))
+                .and_then(|summary| summary.get("standardized_usage"))
+                .and_then(|usage| usage.get("dimensions"))
+                .and_then(|dimensions| dimensions.get("image_count"))
+                .and_then(Value::as_u64),
+            Some(1)
+        );
+        assert_eq!(
+            eof_frame
+                .get("payload")
+                .and_then(|payload| payload.get("summary"))
+                .and_then(|summary| summary.get("standardized_usage"))
+                .and_then(|usage| usage.get("dimensions"))
+                .and_then(|dimensions| dimensions.get("image_size"))
+                .and_then(Value::as_str),
+            Some("1024x1024")
+        );
 
         handle.abort();
     }