Mark sync usage active earlier

2026-10-07 01:47:47 +08:00 · 2026-07-01 02:21:20 +08:00
parent 6ed2f9bd0a
commit 7c32be46ca
10 changed files with 626 additions and 32 deletions
@@ -78,7 +78,8 @@ use crate::orchestration::{
 use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
 use crate::request_candidate_runtime::{
    ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
-    record_local_request_candidate_status,
+    record_local_request_candidate_status, record_local_request_candidate_status_snapshot,
+    snapshot_local_request_candidate_status,
 };
 use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context;
 use crate::usage::{spawn_sync_report, submit_sync_report};
@@ -292,6 +293,82 @@ struct ImplicitSyncFinalizeOutcome {
    outcome: LocalCoreSyncFinalizeOutcome,
 }

+fn spawn_sync_candidate_status_update(
+    state: AppState,
+    snapshot: crate::request_candidate_runtime::LocalRequestCandidateStatusSnapshot,
+    status_update: SchedulerRequestCandidateStatusUpdate,
+) {
+    tokio::spawn(async move {
+        record_local_request_candidate_status_snapshot(&state, &snapshot, status_update).await;
+    });
+}
+
+fn record_sync_response_started(
+    state: &AppState,
+    lifecycle_seed: aether_usage_runtime::LifecycleUsageSeed,
+    request_candidate_status_snapshot: Option<
+        crate::request_candidate_runtime::LocalRequestCandidateStatusSnapshot,
+    >,
+    candidate_started_unix_ms: u64,
+    status_code: u16,
+    ttfb_ms: u64,
+) {
+    state.usage_runtime.record_stream_started_immediate_async(
+        state.data.as_ref(),
+        lifecycle_seed,
+        status_code,
+        Some(ExecutionTelemetry {
+            ttfb_ms: Some(ttfb_ms),
+            elapsed_ms: Some(ttfb_ms),
+            upstream_bytes: None,
+        }),
+    );
+
+    if let Some(snapshot) = request_candidate_status_snapshot {
+        spawn_sync_candidate_status_update(
+            state.clone(),
+            snapshot,
+            SchedulerRequestCandidateStatusUpdate {
+                status: RequestCandidateStatus::Streaming,
+                status_code: Some(status_code),
+                error_type: None,
+                error_message: None,
+                latency_ms: Some(ttfb_ms),
+                started_at_unix_ms: Some(candidate_started_unix_ms),
+                finished_at_unix_ms: None,
+            },
+        );
+    }
+}
+
+fn record_sync_execution_active(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    report_context: Option<&Value>,
+    candidate_started_unix_ms: u64,
+) {
+    let lifecycle_seed = build_lifecycle_usage_seed(plan, report_context);
+    state
+        .usage_runtime
+        .record_sync_active_immediate_async(state.data.as_ref(), lifecycle_seed);
+
+    if let Some(snapshot) = snapshot_local_request_candidate_status(plan, report_context) {
+        spawn_sync_candidate_status_update(
+            state.clone(),
+            snapshot,
+            SchedulerRequestCandidateStatusUpdate {
+                status: RequestCandidateStatus::Streaming,
+                status_code: None,
+                error_type: None,
+                error_message: None,
+                latency_ms: None,
+                started_at_unix_ms: Some(candidate_started_unix_ms),
+                finished_at_unix_ms: None,
+            },
+        );
+    }
+}
+
 fn record_sync_terminal_usage(
    state: &AppState,
    plan: &ExecutionPlan,
@@ -1019,6 +1096,7 @@ async fn execute_direct_sync_runtime_candidate(
    report_context: Option<&Value>,
    trace_id: &str,
    plan_kind: &str,
+    candidate_started_unix_ms: u64,
    request_id_for_log: &str,
    candidate_id: Option<&str>,
    provider_name: &str,
@@ -1035,8 +1113,21 @@ async fn execute_direct_sync_runtime_candidate(
        return Ok(result);
    }
    if !should_track_openai_image_sync_upstream_sse(plan_kind, plan, report_context) {
+        let state_for_response_started = state.clone();
+        let response_started_lifecycle_seed = build_lifecycle_usage_seed(plan, report_context);
+        let response_started_candidate_snapshot =
+            snapshot_local_request_candidate_status(plan, report_context);
        return DirectSyncExecutionRuntime::new()
-            .execute_sync(plan)
+            .execute_sync_with_response_started(plan, move |event| {
+                record_sync_response_started(
+                    &state_for_response_started,
+                    response_started_lifecycle_seed,
+                    response_started_candidate_snapshot,
+                    candidate_started_unix_ms,
+                    event.status_code,
+                    event.ttfb_ms,
+                )
+            })
            .await
            .map_err(SyncExecutionFailure::from_transport);
    }
@@ -1573,6 +1664,12 @@ async fn execute_execution_runtime_sync_impl(
        key_id.as_str(),
    )
    .await;
+    record_sync_execution_active(
+        state,
+        &plan,
+        report_context.as_ref(),
+        candidate_started_unix_secs,
+    );
    #[cfg(not(test))]
    let mut result = {
        match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
@@ -1588,6 +1685,7 @@ async fn execute_execution_runtime_sync_impl(
                        report_context.as_ref(),
                        trace_id,
                        plan_kind,
+                        candidate_started_unix_secs,
                        plan_request_id_for_log.as_str(),
                        plan_candidate_id.as_deref(),
                        provider_name.as_str(),
@@ -1768,6 +1866,7 @@ async fn execute_execution_runtime_sync_impl(
                        report_context.as_ref(),
                        trace_id,
                        plan_kind,
+                        candidate_started_unix_secs,
                        plan_request_id_for_log.as_str(),
                        plan_candidate_id.as_deref(),
                        provider_name.as_str(),
@@ -2690,6 +2789,7 @@ mod tests {
    use serde_json::json;
    use std::collections::BTreeMap;
    use std::sync::Arc;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    fn test_openai_image_plan(stream: bool) -> ExecutionPlan {
        ExecutionPlan {
@@ -2723,6 +2823,17 @@ mod tests {
        plan
    }

+    fn test_decision() -> GatewayControlDecision {
+        GatewayControlDecision::synthetic(
+            "/v1/chat/completions",
+            Some("ai_public".to_string()),
+            Some("openai".to_string()),
+            Some("chat".to_string()),
+            Some("openai:chat".to_string()),
+        )
+        .with_execution_runtime_candidate(true)
+    }
+
    fn test_kiro_sync_plan() -> ExecutionPlan {
        ExecutionPlan {
            request_id: "req-kiro-sync-cache-1".to_string(),
@@ -2917,6 +3028,293 @@ mod tests {
        );
    }

+    #[tokio::test]
+    async fn sync_direct_response_start_marks_usage_and_candidate_active_before_body_finishes() {
+        let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
+        let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
+        let state = AppState::new()
+            .expect("gateway 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 listener = crate::test_support::bind_loopback_listener()
+            .await
+            .expect("listener should bind");
+        let addr = listener.local_addr().expect("local addr should resolve");
+        let (headers_tx, headers_rx) = tokio::sync::oneshot::channel();
+        let (body_tx, body_rx) = tokio::sync::oneshot::channel();
+        let server = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.expect("client should connect");
+            let mut request = [0_u8; 4096];
+            let _ = socket
+                .read(&mut request)
+                .await
+                .expect("request should read");
+            socket
+                .write_all(
+                    b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n",
+                )
+                .await
+                .expect("headers should write");
+            socket.flush().await.expect("headers should flush");
+            let _ = headers_tx.send(());
+            let _ = body_rx.await;
+            socket
+                .write_all(br#"{"ok":true}"#)
+                .await
+                .expect("body should write");
+        });
+
+        let mut plan = test_gemini_chat_plan();
+        plan.request_id = "sync-response-start-active-request".to_string();
+        plan.candidate_id = Some("sync-response-start-active-candidate".to_string());
+        plan.url = format!("http://{addr}/chat");
+        plan.provider_api_format = "openai:chat".to_string();
+        plan.model_name = Some("gpt-5".to_string());
+        plan.body = aether_contracts::RequestBody::from_json(json!({
+            "model": "gpt-5",
+            "messages": [{"role": "user", "content": "slow body"}],
+        }));
+        let report_context = Some(json!({
+            "candidate_index": 0,
+            "retry_index": 0,
+            "user_id": "user-active",
+            "api_key_id": "api-key-active",
+            "client_api_format": "openai:chat",
+            "provider_api_format": "openai:chat",
+            "request_path": "/v1/chat/completions",
+            "request_path_and_query": "/v1/chat/completions",
+            "upstream_url": plan.url.clone(),
+            "mapped_model": "gpt-5",
+        }));
+        let started_at = current_request_candidate_unix_ms();
+        state
+            .usage_runtime
+            .record_pending_direct(
+                state.data.as_ref(),
+                build_lifecycle_usage_seed(&plan, report_context.as_ref()),
+            )
+            .await;
+        record_local_request_candidate_status(
+            &state,
+            &plan,
+            report_context.as_ref(),
+            SchedulerRequestCandidateStatusUpdate {
+                status: RequestCandidateStatus::Pending,
+                status_code: None,
+                error_type: None,
+                error_message: None,
+                latency_ms: None,
+                started_at_unix_ms: Some(started_at),
+                finished_at_unix_ms: None,
+            },
+        )
+        .await;
+
+        let state_for_exec = state.clone();
+        let plan_for_exec = plan.clone();
+        let report_context_for_exec = report_context.clone();
+        let exec = tokio::spawn(async move {
+            execute_direct_sync_runtime_candidate(
+                &state_for_exec,
+                &plan_for_exec,
+                report_context_for_exec.as_ref(),
+                "trace-response-start-active",
+                "openai_chat_sync",
+                started_at,
+                "sync-response-start-active-request",
+                plan_for_exec.candidate_id.as_deref(),
+                "openai",
+                "endpoint-1",
+                "key-1",
+                "gpt-5",
+                "0",
+                None,
+            )
+            .await
+        });
+
+        headers_rx.await.expect("headers should be written");
+        let mut active_usage = None;
+        for _ in 0..50 {
+            if let Some(usage) = usage_repository
+                .find_by_request_id("sync-response-start-active-request")
+                .await
+                .expect("usage should read")
+            {
+                if usage.status == "streaming" {
+                    active_usage = Some(usage);
+                    break;
+                }
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+        let active_usage = active_usage.expect("usage should become active before body finishes");
+        assert_eq!(active_usage.status_code, Some(200));
+        assert!(active_usage.first_byte_time_ms.is_some());
+        assert!(active_usage.response_time_ms.is_some());
+
+        let stored_candidates = request_candidate_repository
+            .list_by_request_id("sync-response-start-active-request")
+            .await
+            .expect("candidate should read");
+        let active_candidate = stored_candidates
+            .iter()
+            .find(|candidate| candidate.id == "sync-response-start-active-candidate")
+            .expect("candidate should exist");
+        assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
+        assert_eq!(active_candidate.status_code, Some(200));
+
+        let _ = body_tx.send(());
+        let result = tokio::time::timeout(Duration::from_secs(2), exec)
+            .await
+            .expect("sync execution should finish")
+            .expect("sync execution task should not panic")
+            .expect("sync execution should succeed");
+        assert_eq!(result.status_code, 200);
+        server.abort();
+    }
+
+    #[tokio::test]
+    async fn sync_execution_active_marks_usage_before_response_headers() {
+        let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
+        let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
+        let state = AppState::new()
+            .expect("gateway 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 listener = crate::test_support::bind_loopback_listener()
+            .await
+            .expect("listener should bind");
+        let addr = listener.local_addr().expect("local addr should resolve");
+        let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel();
+        let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
+        let server = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.expect("client should connect");
+            let mut request = [0_u8; 4096];
+            let _ = socket
+                .read(&mut request)
+                .await
+                .expect("request should read");
+            let _ = request_seen_tx.send(());
+            let _ = finish_rx.await;
+            socket
+                .write_all(
+                    b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n{\"ok\":true}",
+                )
+                .await
+                .expect("response should write");
+        });
+
+        let mut plan = test_gemini_chat_plan();
+        plan.request_id = "sync-active-before-headers-request".to_string();
+        plan.candidate_id = Some("sync-active-before-headers-candidate".to_string());
+        plan.url = format!("http://{addr}/chat");
+        plan.provider_name = Some("OpenAI".to_string());
+        plan.provider_api_format = "openai:chat".to_string();
+        plan.model_name = Some("gpt-5".to_string());
+        plan.body = aether_contracts::RequestBody::from_json(json!({
+            "model": "gpt-5",
+            "messages": [{"role": "user", "content": "slow headers"}],
+        }));
+        let report_context = Some(json!({
+            "candidate_index": 0,
+            "retry_index": 0,
+            "user_id": "user-active-before-headers",
+            "api_key_id": "api-key-active-before-headers",
+            "candidate_id": "sync-active-before-headers-candidate",
+            "provider_id": "provider-1",
+            "endpoint_id": "endpoint-1",
+            "key_id": "key-1",
+            "provider_name": "OpenAI",
+            "client_api_format": "openai:chat",
+            "provider_api_format": "openai:chat",
+            "request_path": "/v1/chat/completions",
+            "request_path_and_query": "/v1/chat/completions",
+            "upstream_url": plan.url.clone(),
+            "mapped_model": "gpt-5",
+        }));
+        let state_for_exec = state.clone();
+        let plan_for_exec = plan.clone();
+        let report_context_for_exec = report_context.clone();
+        let exec = tokio::spawn(async move {
+            execute_execution_runtime_sync(
+                &state_for_exec,
+                "/v1/chat/completions",
+                plan_for_exec,
+                "trace-active-before-headers",
+                &test_decision(),
+                "openai_chat_sync",
+                Some("openai_chat_sync".to_string()),
+                report_context_for_exec,
+            )
+            .await
+        });
+
+        request_seen_rx
+            .await
+            .expect("upstream request should be observed");
+        let mut active_usage = None;
+        for _ in 0..50 {
+            if let Some(usage) = usage_repository
+                .find_by_request_id("sync-active-before-headers-request")
+                .await
+                .expect("usage should read")
+            {
+                if usage.status == "streaming" {
+                    active_usage = Some(usage);
+                    break;
+                }
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+        let active_usage =
+            active_usage.expect("usage should become active before upstream headers");
+        assert_eq!(active_usage.status_code, None);
+        assert_eq!(active_usage.first_byte_time_ms, None);
+        assert_eq!(active_usage.response_time_ms, None);
+
+        let stored_candidates = request_candidate_repository
+            .list_by_request_id("sync-active-before-headers-request")
+            .await
+            .expect("candidate should read");
+        let active_candidate = stored_candidates
+            .iter()
+            .find(|candidate| candidate.id == "sync-active-before-headers-candidate")
+            .expect("candidate should exist");
+        assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
+        assert_eq!(active_candidate.status_code, None);
+        assert!(active_candidate.started_at_unix_ms.is_some());
+        assert!(active_candidate.finished_at_unix_ms.is_none());
+
+        let _ = finish_tx.send(());
+        let response = tokio::time::timeout(Duration::from_secs(2), exec)
+            .await
+            .expect("sync execution should finish")
+            .expect("sync execution task should not panic")
+            .expect("sync execution should succeed")
+            .expect("sync execution should produce a response");
+        assert_eq!(response.status(), StatusCode::OK);
+        server.abort();
+    }
+
    #[test]
    fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body() {
        let plan = test_kiro_sync_plan();
@@ -608,6 +608,12 @@ pub(crate) struct DirectUpstreamStreamExecution {
    pub(crate) upstream_target_permit: Option<UpstreamTargetAdmissionPermit>,
 }

+#[derive(Debug, Clone)]
+pub(crate) struct DirectSyncResponseStarted {
+    pub(crate) status_code: u16,
+    pub(crate) ttfb_ms: u64,
+}
+
 impl DirectSyncExecutionRuntime {
    pub(crate) const fn new() -> Self {
        Self
@@ -617,6 +623,17 @@ impl DirectSyncExecutionRuntime {
        &self,
        plan: &ExecutionPlan,
    ) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
+        self.execute_sync_with_response_started(plan, |_| {}).await
+    }
+
+    pub(crate) async fn execute_sync_with_response_started<F>(
+        &self,
+        plan: &ExecutionPlan,
+        on_response_started: F,
+    ) -> Result<ExecutionResult, ExecutionRuntimeTransportError>
+    where
+        F: FnOnce(DirectSyncResponseStarted),
+    {
        let body_bytes = build_request_body(plan)?;

        let started_at = Instant::now();
@@ -625,6 +642,10 @@ impl DirectSyncExecutionRuntime {
            let ttfb_ms = started_at.elapsed().as_millis() as u64;
            let status_code = response.status_code();
            let headers = response.headers();
+            on_response_started(DirectSyncResponseStarted {
+                status_code,
+                ttfb_ms,
+            });
            let (body_bytes, stream_ttfb_ms) =
                response.bytes_with_stream_timeout(plan, started_at).await?;
            let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
@@ -160,14 +160,14 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled_i
 }

 #[test]
-fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() {
+fn gateway_records_active_usage_before_execution_runtime_sync_result_arrives() {
    run_async_test_on_large_stack(
-        "gateway_records_pending_usage_before_execution_runtime_sync_result_arrives",
-        gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl(),
+        "gateway_records_active_usage_before_execution_runtime_sync_result_arrives",
+        gateway_records_active_usage_before_execution_runtime_sync_result_arrives_impl(),
    );
 }

-async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl() {
+async fn gateway_records_active_usage_before_execution_runtime_sync_result_arrives_impl() {
    let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
    let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
    let execution_request_started = Arc::new(tokio::sync::Notify::new());
@@ -275,24 +275,26 @@ async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arri

    execution_request_started.notified().await;

-    let mut pending = None;
+    let mut active = None;
    for _ in 0..50 {
-        pending = usage_repository
+        active = usage_repository
            .find_by_request_id("req-usage-sync-pending-123")
            .await
            .expect("usage lookup should succeed");
-        if pending
+        if active
            .as_ref()
-            .is_some_and(|stored| stored.status == "pending")
+            .is_some_and(|stored| stored.status == "streaming")
        {
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
-    let pending = pending.expect("pending usage should be recorded before sync result resolves");
-    assert_eq!(pending.status, "pending");
-    assert_eq!(pending.billing_status, "pending");
-    assert_eq!(pending.response_time_ms, None);
+    let active = active.expect("active usage should be recorded before sync result resolves");
+    assert_eq!(active.status, "streaming");
+    assert_eq!(active.billing_status, "pending");
+    assert_eq!(active.status_code, None);
+    assert_eq!(active.first_byte_time_ms, None);
+    assert_eq!(active.response_time_ms, None);

    allow_execution_response.notify_one();

@@ -320,14 +322,14 @@ async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arri
 }

 #[test]
-fn gateway_keeps_pending_sync_usage_lightweight_for_large_request_body() {
+fn gateway_keeps_active_sync_usage_lightweight_for_large_request_body() {
    run_async_test_on_large_stack(
-        "gateway_keeps_pending_sync_usage_lightweight_for_large_request_body",
-        gateway_keeps_pending_sync_usage_lightweight_for_large_request_body_impl(),
+        "gateway_keeps_active_sync_usage_lightweight_for_large_request_body",
+        gateway_keeps_active_sync_usage_lightweight_for_large_request_body_impl(),
    );
 }

-async fn gateway_keeps_pending_sync_usage_lightweight_for_large_request_body_impl() {
+async fn gateway_keeps_active_sync_usage_lightweight_for_large_request_body_impl() {
    let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
    let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
    let execution_request_started = Arc::new(tokio::sync::Notify::new());
@@ -434,28 +436,28 @@ async fn gateway_keeps_pending_sync_usage_lightweight_for_large_request_body_imp

    execution_request_started.notified().await;

-    let mut pending = None;
+    let mut active = None;
    for _ in 0..50 {
-        pending = usage_repository
+        active = usage_repository
            .find_by_request_id("req-usage-sync-large-pending-123")
            .await
            .expect("usage lookup should succeed");
-        if pending
+        if active
            .as_ref()
-            .is_some_and(|stored| stored.status == "pending")
+            .is_some_and(|stored| stored.status == "streaming")
        {
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
-    let pending = pending.expect("pending usage should be recorded before sync result resolves");
-    assert_eq!(pending.status, "pending");
-    assert!(pending.request_headers.is_none());
-    assert!(pending.request_body.is_none());
-    assert!(pending.provider_request_headers.is_none());
-    assert!(pending.provider_request_body.is_none());
-    assert!(pending.response_headers.is_none());
-    assert!(pending.client_response_headers.is_none());
+    let active = active.expect("active usage should be recorded before sync result resolves");
+    assert_eq!(active.status, "streaming");
+    assert!(active.request_headers.is_none());
+    assert!(active.request_body.is_none());
+    assert!(active.provider_request_headers.is_none());
+    assert!(active.provider_request_body.is_none());
+    assert!(active.response_headers.is_none());
+    assert!(active.client_response_headers.is_none());

    allow_execution_response.notify_one();