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

2026-10-06 17:37:47 +08:00 · 2026-05-18 01:32:45 +08:00
parent c024c782e4
commit 691ccaaa04
4 changed files with 511 additions and 6 deletions
@@ -991,6 +991,12 @@ fn build_sse_body_stream(
    }
 }

+fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
+    std::str::from_utf8(chunk)
+        .ok()
+        .is_some_and(|text| text.lines().any(|line| line.trim() == "data: [DONE]"))
+}
+
 async fn next_stream_frame<R>(
    buffered_frames: &mut VecDeque<StreamFrame>,
    lines: &mut FramedRead<R, LinesCodec>,
@@ -2026,6 +2032,8 @@ async fn execute_stream_from_frame_stream(
            max_stream_body_buffer_bytes,
            &mut client_body_truncated,
        );
+        let mut client_visible_stream_completed =
+            stream_chunk_contains_sse_done(&prefetched_body_for_report);
        let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
        let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
        let reached_eof = initial_reached_eof;
@@ -2452,6 +2460,8 @@ async fn execute_stream_from_frame_stream(
                        );
                        let rewritten_chunk_len =
                            u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
+                        let chunk_completed_stream =
+                            stream_chunk_contains_sse_done(&rewritten_chunk);
                        if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
                            warn!(
                                event_name = "stream_execution_downstream_disconnected",
@@ -2464,6 +2474,7 @@ async fn execute_stream_from_frame_stream(
                            downstream_dropped = true;
                            break;
                        } else {
+                            client_visible_stream_completed |= chunk_completed_stream;
                            client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
                            last_client_chunk_elapsed_ms.store(
                                stream_started_at_for_report
@@ -2578,6 +2589,8 @@ async fn execute_stream_from_frame_stream(
                            );
                            let rewritten_chunk_len =
                                u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
+                            let chunk_completed_stream =
+                                stream_chunk_contains_sse_done(&rewritten_chunk);
                            if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
                                warn!(
                                    event_name = "stream_execution_downstream_flush_disconnected",
@@ -2589,6 +2602,7 @@ async fn execute_stream_from_frame_stream(
                                );
                                downstream_dropped = true;
                            } else {
+                                client_visible_stream_completed |= chunk_completed_stream;
                                client_stream_bytes
                                    .fetch_add(rewritten_chunk_len, Ordering::Relaxed);
                                last_client_chunk_elapsed_ms.store(
@@ -2635,6 +2649,8 @@ async fn execute_stream_from_frame_stream(
                            );
                            let flushed_chunk_len =
                                u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
+                            let chunk_completed_stream =
+                                stream_chunk_contains_sse_done(&flushed_chunk);
                            if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
                                warn!(
                                    event_name = "stream_execution_downstream_rewrite_flush_disconnected",
@@ -2646,6 +2662,7 @@ async fn execute_stream_from_frame_stream(
                                );
                                downstream_dropped = true;
                            } else {
+                                client_visible_stream_completed |= chunk_completed_stream;
                                client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
                                last_client_chunk_elapsed_ms.store(
                                    stream_started_at_for_report
@@ -2755,6 +2772,18 @@ async fn execute_stream_from_frame_stream(
            ),
        );

+        if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() {
+            debug!(
+                event_name = "execution_runtime_stream_downstream_closed_after_done",
+                log_type = "debug",
+                trace_id = %trace_id_owned,
+                request_id = %request_id_for_report_log,
+                candidate_id = ?candidate_id_for_report.as_deref(),
+                "gateway treats downstream close after client-visible SSE DONE as completed"
+            );
+            downstream_dropped = false;
+        }
+
        if downstream_dropped {
            debug!(
                event_name = "execution_runtime_stream_report_skipped",
@@ -3417,6 +3446,142 @@ mod tests {
        );
    }

+    #[tokio::test]
+    async fn image_stream_downstream_close_after_done_is_recorded_success() {
+        let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
+        let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(
+                crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
+                    Arc::clone(&request_candidate_repository),
+                    Arc::clone(&usage_repository),
+                ),
+            )
+            .with_usage_runtime_for_tests(UsageRuntimeConfig {
+                enabled: true,
+                ..UsageRuntimeConfig::default()
+            });
+        let plan = ExecutionPlan {
+            request_id: "req-image-done-close-success".into(),
+            candidate_id: Some("cand-image-done-close-success".into()),
+            provider_name: Some("openai".into()),
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/v1/images/generations".into(),
+            headers: BTreeMap::from([("accept".into(), "text/event-stream".into())]),
+            content_type: Some("application/json".into()),
+            content_encoding: None,
+            body: RequestBody::from_json(json!({
+                "model": "gpt-image-2",
+                "prompt": "draw a small image",
+                "stream": true
+            })),
+            stream: true,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:image".into(),
+            model_name: Some("gpt-image-2".into()),
+            proxy: None,
+            transport_profile: None,
+            timeouts: None,
+        };
+        let frame_stream = stream! {
+            yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
+                b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
+            ));
+            yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
+                b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_1\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\nevent: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_1\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":null}}\\n\\n\"}}\n",
+            ));
+            std::future::pending::<()>().await;
+        }
+        .boxed();
+
+        let response = execute_stream_from_frame_stream(
+            &state,
+            plan,
+            "trace-image-done-close-success",
+            &test_decision(),
+            "openai_chat_stream",
+            Some("openai_chat_stream_success".to_string()),
+            Some(json!({
+                "request_id": "req-image-done-close-success",
+                "candidate_id": "cand-image-done-close-success",
+                "candidate_index": 0,
+                "retry_index": 0,
+                "provider_api_format": "openai:image",
+                "client_api_format": "openai:chat",
+                "image_request": {
+                    "size": "1024x1024",
+                    "quality": "medium"
+                }
+            })),
+            crate::clock::current_unix_ms(),
+            Instant::now(),
+            frame_stream,
+            None,
+        )
+        .await
+        .expect("execution should succeed")
+        .expect("execution should return a client response");
+
+        let mut body_stream = response.into_body().into_data_stream();
+        let mut body = Vec::new();
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while !String::from_utf8_lossy(&body).contains("data: [DONE]") {
+                let chunk = body_stream
+                    .next()
+                    .await
+                    .expect("body should yield until done")
+                    .expect("chunk should be ok");
+                body.extend_from_slice(&chunk);
+            }
+        })
+        .await
+        .expect("final DONE should arrive");
+        drop(body_stream);
+
+        let candidates = tokio::time::timeout(Duration::from_secs(1), async {
+            loop {
+                let candidates = request_candidate_repository
+                    .list_by_request_id("req-image-done-close-success")
+                    .await
+                    .expect("request candidates should read");
+                if candidates
+                    .first()
+                    .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success)
+                {
+                    break candidates;
+                }
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .expect("candidate should be marked success");
+        assert_eq!(candidates[0].status_code, Some(200));
+
+        let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
+            loop {
+                let usage = usage_repository
+                    .find_by_request_id("req-image-done-close-success")
+                    .await
+                    .expect("usage should read");
+                if usage
+                    .as_ref()
+                    .is_some_and(|usage| usage.status == "completed")
+                {
+                    break usage.expect("completed usage should exist");
+                }
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .expect("usage should be marked completed");
+        assert_eq!(stored_usage.status_code, Some(200));
+        assert!(stored_usage.total_tokens > 0);
+    }
+
    #[tokio::test]
    async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
    ) {