feat(gateway): configure cyber policy failover

2026-09-12 22:20:19 +08:00 · 2026-07-19 23:27:19 +08:00
parent e0dbb233f7
commit f8778c4a23
11 changed files with 429 additions and 53 deletions
@@ -940,7 +940,7 @@ mod tests {
                continue_status_codes: [409, 429].into_iter().collect(),
                success_failover_patterns: Vec::new(),
                error_stop_patterns: Vec::new(),
-                stop_cyber_policy_errors: false,
+                stop_cyber_policy_errors: true,
                retry_client_errors_by_default: true,
            }
        );
@@ -98,11 +98,12 @@ use crate::execution_runtime::{
 use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
 use crate::log_ids::short_request_id;
 use crate::orchestration::{
-    apply_local_execution_effect, build_local_error_flow_metadata, trace_upstream_response_body,
-    with_error_flow_report_context, with_upstream_response_report_context,
-    LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
-    LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
-    LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
+    apply_local_execution_effect, build_local_error_flow_metadata, cyber_continue_failover_enabled,
+    trace_upstream_response_body, with_error_flow_report_context,
+    with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
+    LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
+    LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
+    LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
 };
 use crate::provider_pool_demand::{
    acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
@@ -1121,6 +1122,7 @@ fn should_use_direct_sse_passthrough(
        plan.client_api_format.as_str(),
        false,
        false,
+        false,
    )
 }

@@ -4074,11 +4076,16 @@ fn should_skip_direct_finalize_prefetch(
    client_api_format: &str,
    has_private_stream_normalizer: bool,
    has_local_stream_rewriter: bool,
+    force_prefetch: bool,
 ) -> bool {
    if direct_stream_finalize_kind.is_none() {
        return false;
    }

+    if force_prefetch {
+        return false;
+    }
+
    let content_type = content_type
        .map(str::trim)
        .filter(|value| !value.is_empty())
@@ -4103,6 +4110,36 @@ fn should_skip_direct_finalize_prefetch(
    !(content_type.contains("json") || content_type.ends_with("+json"))
 }

+fn prefetched_openai_responses_body_has_output_boundary(body: &[u8]) -> bool {
+    let Ok(text) = std::str::from_utf8(body) else {
+        return true;
+    };
+    for line in text.lines() {
+        let Some(data) = line.trim().strip_prefix("data:").map(str::trim) else {
+            continue;
+        };
+        if data.is_empty() {
+            continue;
+        }
+        if data == "[DONE]" {
+            return true;
+        }
+        let Ok(event) = serde_json::from_str::<Value>(data) else {
+            continue;
+        };
+        let event_type = event.get("type").and_then(Value::as_str).map(str::trim);
+        if !event_type.is_some_and(|event_type| {
+            matches!(
+                event_type,
+                "response.created" | "response.in_progress" | "response.queued"
+            )
+        }) {
+            return true;
+        }
+    }
+    false
+}
+
 fn should_probe_success_failover_before_stream(headers: &BTreeMap<String, String>) -> bool {
    let content_type = headers
        .get("content-type")
@@ -4577,6 +4614,9 @@ async fn execute_stream_from_frame_stream(
        headers.insert("content-type".to_string(), "text/event-stream".to_string());
    }
    let upstream_content_type = upstream_headers.get("content-type").map(String::as_str);
+    let prefetch_for_cyber_failover =
+        is_openai_responses_family_format(plan.provider_api_format.as_str())
+            && cyber_continue_failover_enabled(state).await;
    let skip_direct_finalize_prefetch = should_skip_direct_finalize_prefetch(
        direct_stream_finalize_kind.as_deref(),
        upstream_content_type,
@@ -4584,6 +4624,7 @@ async fn execute_stream_from_frame_stream(
        plan.client_api_format.as_str(),
        private_stream_normalizer.is_some(),
        local_stream_rewriter.is_some(),
+        prefetch_for_cyber_failover,
    );
    let limit_direct_finalize_prefetch =
        should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some());
@@ -4930,7 +4971,12 @@ async fn execute_stream_from_frame_stream(
                        prefetched_chunks.push(Bytes::from(rewritten_chunk));
                    }

-                    if matches!(inspection, StreamPrefetchInspection::NonError) {
+                    if matches!(inspection, StreamPrefetchInspection::NonError)
+                        && (!prefetch_for_cyber_failover
+                            || prefetched_openai_responses_body_has_output_boundary(
+                                &prefetched_inspection_body,
+                            ))
+                    {
                        break;
                    }
                }
@@ -6190,10 +6236,10 @@ mod tests {
        ensure_stream_terminal_summary_for_missing_observed_finish,
        execute_execution_runtime_stream, execute_stream_from_frame_stream,
        maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
-        parse_direct_passthrough_mode, should_limit_direct_finalize_prefetch,
-        should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
-        stream_chunk_contains_sse_done, stream_requires_observed_terminal_event,
-        stream_terminal_summary_missing_observed_finish,
+        parse_direct_passthrough_mode, prefetched_openai_responses_body_has_output_boundary,
+        should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream,
+        should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done,
+        stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish,
        stream_terminal_summary_missing_observed_finish_with_requirement,
        stream_terminal_summary_represents_failure_with_requirement,
        ClientVisibleStreamCompletionTracker, DirectPassthroughMode, ProviderStreamErrorInspection,
@@ -6205,6 +6251,20 @@ mod tests {

    fn provider_catalog_stop_429_for_plan(
        plan: &ExecutionPlan,
+    ) -> InMemoryProviderCatalogReadRepository {
+        provider_catalog_for_plan(
+            plan,
+            Some(json!({
+                "failover_rules": {
+                    "stop_status_codes": [429]
+                }
+            })),
+        )
+    }
+
+    fn provider_catalog_for_plan(
+        plan: &ExecutionPlan,
+        provider_config: Option<Value>,
    ) -> InMemoryProviderCatalogReadRepository {
        let provider_type = plan.provider_name.as_deref().unwrap_or("custom");
        let provider = StoredProviderCatalogProvider::new(
@@ -6223,11 +6283,7 @@ mod tests {
            None,
            None,
            None,
-            Some(json!({
-                "failover_rules": {
-                    "stop_status_codes": [429]
-                }
-            })),
+            provider_config,
        );
        let endpoint = StoredProviderCatalogEndpoint::new(
            plan.endpoint_id.clone(),
@@ -6274,6 +6330,118 @@ mod tests {
        InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key])
    }

+    fn codex_cyber_policy_plan(request_id: &str) -> ExecutionPlan {
+        ExecutionPlan {
+            request_id: request_id.to_string(),
+            candidate_id: Some(format!("candidate-{request_id}")),
+            provider_name: Some("codex".to_string()),
+            provider_id: format!("provider-{request_id}"),
+            endpoint_id: format!("endpoint-{request_id}"),
+            key_id: format!("key-{request_id}"),
+            method: "POST".to_string(),
+            url: "https://chatgpt.com/backend-api/codex/responses".to_string(),
+            headers: BTreeMap::from([
+                ("content-type".to_string(), "application/json".to_string()),
+                ("accept".to_string(), "text/event-stream".to_string()),
+            ]),
+            content_type: Some("application/json".to_string()),
+            content_encoding: None,
+            body: RequestBody::from_json(json!({
+                "model": "gpt-5.5",
+                "input": [],
+                "stream": true
+            })),
+            stream: true,
+            client_api_format: "openai:responses".to_string(),
+            provider_api_format: "openai:responses".to_string(),
+            model_name: Some("gpt-5.5".to_string()),
+            proxy: None,
+            transport_profile: None,
+            timeouts: None,
+        }
+    }
+
+    async fn execute_prefetched_codex_cyber_policy_failure(
+        continue_failover: bool,
+    ) -> Option<axum::http::Response<Body>> {
+        let request_id = if continue_failover {
+            "req-cyber-policy-retry"
+        } else {
+            "req-cyber-policy-stop"
+        };
+        let plan = codex_cyber_policy_plan(request_id);
+        let provider_catalog = provider_catalog_for_plan(&plan, None);
+        let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
+            Arc::new(provider_catalog),
+            "development-key",
+        );
+        let data_state = if continue_failover {
+            data_state.with_system_config_values_for_tests([(
+                crate::orchestration::CYBER_CONTINUE_FAILOVER_CONFIG_KEY.to_string(),
+                json!(true),
+            )])
+        } else {
+            data_state
+        };
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data_state);
+        let upstream_setup = "event: response.created\ndata: {\"type\":\"response.created\"}\n\n";
+        let upstream_error = "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"type\":\"invalid_request\",\"message\":\"cyber policy rejected the request\",\"code\":\"cyber_policy_violation\",\"param\":\"input\"}}}\n\n";
+        let frame_stream = stream! {
+            yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
+                frame_type: StreamFrameType::Headers,
+                payload: StreamFramePayload::Headers {
+                    status_code: 200,
+                    headers: BTreeMap::from([(
+                        "content-type".to_string(),
+                        "text/event-stream".to_string(),
+                    )]),
+                },
+            }));
+            yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
+                frame_type: StreamFrameType::Data,
+                payload: StreamFramePayload::Data {
+                    chunk_b64: None,
+                    text: Some(upstream_setup.to_string()),
+                },
+            }));
+            yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
+                frame_type: StreamFrameType::Data,
+                payload: StreamFramePayload::Data {
+                    chunk_b64: None,
+                    text: Some(upstream_error.to_string()),
+                },
+            }));
+            yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame::eof()));
+        }
+        .boxed();
+
+        execute_stream_from_frame_stream(
+            &state,
+            plan,
+            &format!("trace-{request_id}"),
+            &test_decision(),
+            "openai_responses_stream",
+            Some("openai_responses_stream_success".to_string()),
+            Some(json!({
+                "request_id": request_id,
+                "candidate_id": format!("candidate-{request_id}"),
+                "candidate_index": 0,
+                "retry_index": 0,
+                "provider_api_format": "openai:responses",
+                "client_api_format": "openai:responses"
+            })),
+            crate::clock::current_unix_ms(),
+            Instant::now(),
+            RequestStageTrace::from_env(),
+            frame_stream,
+            None,
+        )
+        .await
+        .expect("execution should succeed")
+    }
+
    fn test_decision() -> GatewayControlDecision {
        GatewayControlDecision::synthetic(
            "/v1/chat/completions",
@@ -6365,6 +6533,25 @@ mod tests {
        assert_eq!(detected.pointer("/error/param"), Some(&json!("input")));
    }

+    #[tokio::test]
+    async fn prefetched_codex_cyber_policy_violation_stops_failover_by_default() {
+        let response = execute_prefetched_codex_cyber_policy_failure(false)
+            .await
+            .expect("default Codex cyber policy handling should return the provider error");
+
+        assert_eq!(response.status().as_u16(), 200);
+    }
+
+    #[tokio::test]
+    async fn prefetched_codex_cyber_policy_violation_retries_when_system_setting_is_enabled() {
+        assert!(
+            execute_prefetched_codex_cyber_policy_failure(true)
+                .await
+                .is_none(),
+            "enabling cyber failover should retry the next candidate"
+        );
+    }
+
    fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
        aether_contracts::ProxySnapshot {
            enabled: Some(true),
@@ -7258,6 +7445,7 @@ mod tests {
            "claude:messages",
            false,
            false,
+            false,
        ));
    }

@@ -7270,6 +7458,7 @@ mod tests {
            "claude:messages",
            false,
            false,
+            false,
        ));
    }

@@ -7282,6 +7471,7 @@ mod tests {
            "claude:messages",
            false,
            false,
+            false,
        ));
    }

@@ -7294,6 +7484,30 @@ mod tests {
            "claude:messages",
            false,
            true,
+            false,
+        ));
+    }
+
+    #[test]
+    fn cyber_failover_setting_forces_prefetch_for_event_streams() {
+        assert!(!should_skip_direct_finalize_prefetch(
+            Some("openai_responses_sync_finalize"),
+            Some("text/event-stream"),
+            "openai:responses",
+            "openai:responses",
+            false,
+            false,
+            true,
+        ));
+    }
+
+    #[test]
+    fn cyber_prefetch_waits_through_response_setup_until_output() {
+        assert!(!prefetched_openai_responses_body_has_output_boundary(
+            b"event: response.created\ndata: {\"type\":\"response.created\"}\n\n"
+        ));
+        assert!(prefetched_openai_responses_body_has_output_boundary(
+            b"event: response.created\ndata: {\"type\":\"response.created\"}\n\nevent: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n"
        ));
    }

@@ -21,8 +21,9 @@ use crate::log_ids::short_request_id;
 use crate::orchestration::{
    apply_local_execution_effect, resolve_local_failover_analysis_for_attempt,
    with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
-    LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
-    LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
+    LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis,
+    LocalFailoverDecision, LocalHealthFailureEffect, LocalOAuthInvalidationEffect,
+    LocalPoolErrorEffect,
 };
 use crate::request_candidate_runtime::record_report_request_candidate_status;
 use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context;
@@ -53,6 +54,12 @@ struct StreamFailureBodyFields<'a> {
    extra_error_fields: &'a Map<String, Value>,
 }

+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum StreamFailureHandling {
+    Terminal,
+    HonorLocalFailover,
+}
+
 impl StreamFailureReport {
    fn into_body_jsons(self) -> (Value, Option<Value>) {
        let Self {
@@ -257,7 +264,8 @@ async fn record_stream_sync_failure(
    report_context: Option<&Value>,
    payload: &GatewaySyncReportRequest,
    started_at_unix_ms: Option<u64>,
-) {
+    handling: StreamFailureHandling,
+) -> LocalFailoverAnalysis {
    let error_type = stream_failure_body_field(payload, "type").unwrap_or("internal");
    let error_message = stream_failure_body_field(payload, "message").unwrap_or_default();
    let error_body = payload
@@ -346,16 +354,22 @@ async fn record_stream_sync_failure(
        }),
    )
    .await;
-    let report_context_with_diagnostics =
-        attach_current_request_diagnostics_to_report_context(report_context);
-    let context_seed = build_terminal_usage_context_seed(
-        plan,
-        report_context_with_diagnostics.as_ref().or(report_context),
+    let retrying_next_candidate = matches!(
+        failure_analysis.decision,
+        LocalFailoverDecision::RetryNextCandidate
    );
-    let payload_seed = build_sync_terminal_usage_payload_seed(payload);
-    state
-        .usage_runtime
-        .record_sync_terminal(state.data.as_ref(), context_seed, payload_seed);
+    if !matches!(handling, StreamFailureHandling::HonorLocalFailover) || !retrying_next_candidate {
+        let report_context_with_diagnostics =
+            attach_current_request_diagnostics_to_report_context(report_context);
+        let context_seed = build_terminal_usage_context_seed(
+            plan,
+            report_context_with_diagnostics.as_ref().or(report_context),
+        );
+        let payload_seed = build_sync_terminal_usage_payload_seed(payload);
+        state
+            .usage_runtime
+            .record_sync_terminal(state.data.as_ref(), context_seed, payload_seed);
+    }
    let terminal_unix_secs = current_request_candidate_unix_ms();
    record_report_request_candidate_status(
        state,
@@ -374,6 +388,7 @@ async fn record_stream_sync_failure(
        },
    )
    .await;
+    failure_analysis
 }

 #[allow(clippy::too_many_arguments)] // internal helper for prefetch error handling
@@ -408,7 +423,31 @@ pub(super) async fn handle_prefetch_provider_private_stream_error(
            .then(|| base64::engine::general_purpose::STANDARD.encode(buffered_body)),
        telemetry,
    };
-    record_stream_sync_failure(state, plan, payload.report_context.as_ref(), &payload, None).await;
+    let failure_analysis = record_stream_sync_failure(
+        state,
+        plan,
+        payload.report_context.as_ref(),
+        &payload,
+        None,
+        StreamFailureHandling::HonorLocalFailover,
+    )
+    .await;
+    if matches!(
+        failure_analysis.decision,
+        LocalFailoverDecision::RetryNextCandidate
+    ) {
+        warn!(
+            event_name = "local_stream_candidate_retry_scheduled",
+            log_type = "event",
+            trace_id = %trace_id,
+            request_id = %request_id,
+            candidate_id = ?candidate_id,
+            status_code,
+            failover_classification = failure_analysis.classification.as_str(),
+            "gateway local stream decision retrying next candidate after prefetched provider error"
+        );
+        return Ok(None);
+    }

    let response =
        submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
@@ -443,7 +482,15 @@ pub(super) async fn handle_prefetch_stream_failure(
        buffered_body,
        failure,
    );
-    record_stream_sync_failure(state, plan, payload.report_context.as_ref(), &payload, None).await;
+    record_stream_sync_failure(
+        state,
+        plan,
+        payload.report_context.as_ref(),
+        &payload,
+        None,
+        StreamFailureHandling::Terminal,
+    )
+    .await;

    let response =
        submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
@@ -487,6 +534,7 @@ pub(super) async fn submit_midstream_stream_failure(
        payload.report_context.as_ref(),
        &payload,
        Some(started_at_unix_ms),
+        StreamFailureHandling::Terminal,
    )
    .await;
    if let Err(err) = submit_sync_report(state, payload).await {