Fix Codex image progress heartbeat merge regressions

apps/aether-gateway/src/tests/ai_execute/stream/decision.rs

@@ -928,6 +928,428 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
     upstream_handle.abort();
 }
 
+#[tokio::test]
+async fn gateway_executes_openai_chat_stream_via_local_cross_format_gemini_candidate_with_stream_path_rewrite(
+) {
+    #[derive(Debug, Clone)]
+    struct SeenExecutionRuntimeStreamRequest {
+        trace_id: String,
+        url: String,
+        provider_model: String,
+        auth_header_value: String,
+        accept: String,
+        endpoint_tag: String,
+    }
+
+    fn hash_api_key(value: &str) -> String {
+        let mut hasher = Sha256::new();
+        hasher.update(value.as_bytes());
+        format!("{:x}", hasher.finalize())
+    }
+
+    fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
+        StoredAuthApiKeySnapshot::new(
+            user_id.to_string(),
+            "alice".to_string(),
+            Some("alice@example.com".to_string()),
+            "user".to_string(),
+            "local".to_string(),
+            true,
+            false,
+            Some(serde_json::json!(["openai", "gemini"])),
+            Some(serde_json::json!(["openai:chat"])),
+            Some(serde_json::json!(["gpt-5"])),
+            api_key_id.to_string(),
+            Some("default".to_string()),
+            true,
+            false,
+            false,
+            Some(60),
+            Some(5),
+            Some(4_102_444_800),
+            Some(serde_json::json!(["openai", "gemini"])),
+            Some(serde_json::json!(["openai:chat"])),
+            Some(serde_json::json!(["gpt-5"])),
+        )
+        .expect("auth snapshot should build")
+    }
+
+    fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
+        StoredMinimalCandidateSelectionRow {
+            provider_id: "provider-openai-chat-gemini-stream-local-1".to_string(),
+            provider_name: "gemini".to_string(),
+            provider_type: "custom".to_string(),
+            provider_priority: 10,
+            provider_is_active: true,
+            endpoint_id: "endpoint-openai-chat-gemini-stream-local-1".to_string(),
+            endpoint_api_format: "gemini:generate_content".to_string(),
+            endpoint_api_family: Some("gemini".to_string()),
+            endpoint_kind: Some("chat".to_string()),
+            endpoint_is_active: true,
+            key_id: "key-openai-chat-gemini-stream-local-1".to_string(),
+            key_name: "prod".to_string(),
+            key_auth_type: "api_key".to_string(),
+            key_is_active: true,
+            key_api_formats: Some(vec!["gemini:generate_content".to_string()]),
+            key_allowed_models: None,
+            key_capabilities: None,
+            key_internal_priority: 5,
+            key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
+            model_id: "model-openai-chat-gemini-stream-local-1".to_string(),
+            global_model_id: "global-model-openai-chat-gemini-stream-local-1".to_string(),
+            global_model_name: "gpt-5".to_string(),
+            global_model_mappings: None,
+            global_model_supports_streaming: Some(true),
+            model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
+            model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
+                name: "gemini-2.5-pro-upstream".to_string(),
+                priority: 1,
+                api_formats: Some(vec!["gemini:generate_content".to_string()]),
+                endpoint_ids: None,
+            }]),
+            model_supports_streaming: Some(true),
+            model_is_active: true,
+            model_is_available: true,
+        }
+    }
+
+    fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
+        StoredProviderCatalogProvider::new(
+            "provider-openai-chat-gemini-stream-local-1".to_string(),
+            "gemini".to_string(),
+            Some("https://example.com".to_string()),
+            "custom".to_string(),
+        )
+        .expect("provider should build")
+        .with_transport_fields(
+            true,
+            false,
+            true,
+            None,
+            Some(2),
+            None,
+            Some(20.0),
+            None,
+            None,
+        )
+    }
+
+    fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
+        StoredProviderCatalogEndpoint::new(
+            "endpoint-openai-chat-gemini-stream-local-1".to_string(),
+            "provider-openai-chat-gemini-stream-local-1".to_string(),
+            "gemini:generate_content".to_string(),
+            Some("gemini".to_string()),
+            Some("chat".to_string()),
+            true,
+        )
+        .expect("endpoint should build")
+        .with_transport_fields(
+            "https://generativelanguage.googleapis.com".to_string(),
+            Some(serde_json::json!([
+                {"action":"set","key":"x-endpoint-tag","value":"openai-chat-gemini-cross-format-stream"}
+            ])),
+            None,
+            Some(2),
+            Some("/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent".to_string()),
+            None,
+            None,
+            None,
+        )
+        .expect("endpoint transport should build")
+    }
+
+    fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
+        StoredProviderCatalogKey::new(
+            "key-openai-chat-gemini-stream-local-1".to_string(),
+            "provider-openai-chat-gemini-stream-local-1".to_string(),
+            "prod".to_string(),
+            "api_key".to_string(),
+            None,
+            true,
+        )
+        .expect("key should build")
+        .with_transport_fields(
+            Some(serde_json::json!(["gemini:generate_content"])),
+            encrypt_python_fernet_plaintext(
+                DEVELOPMENT_ENCRYPTION_KEY,
+                "sk-upstream-openai-chat-gemini-stream",
+            )
+            .expect("api key should encrypt"),
+            None,
+            None,
+            Some(serde_json::json!({"gemini:generate_content": 1})),
+            None,
+            None,
+            None,
+            None,
+        )
+        .expect("key transport should build")
+    }
+
+    let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
+    let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
+    let seen_report = Arc::new(Mutex::new(false));
+    let seen_report_clone = Arc::clone(&seen_report);
+    let decision_hits = Arc::new(Mutex::new(0usize));
+    let decision_hits_clone = Arc::clone(&decision_hits);
+    let plan_hits = Arc::new(Mutex::new(0usize));
+    let plan_hits_clone = Arc::clone(&plan_hits);
+    let public_hits = Arc::new(Mutex::new(0usize));
+    let public_hits_clone = Arc::clone(&public_hits);
+    let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
+
+    let upstream = Router::new()
+        .route(
+            "/api/internal/gateway/resolve",
+            any(|_request: Request| async move {
+                Json(json!({
+                    "action": "proxy_public",
+                    "route_class": "ai_public",
+                    "route_family": "openai",
+                    "route_kind": "chat",
+                    "auth_endpoint_signature": "openai:chat",
+                    "execution_runtime_candidate": true,
+                    "auth_context": {
+                        "user_id": "user-openai-chat-gemini-stream-local-1",
+                        "api_key_id": "api-key-openai-chat-gemini-stream-local-1",
+                        "access_allowed": true
+                    },
+                    "public_path": "/v1/chat/completions"
+                }))
+            }),
+        )
+        .route(
+            "/api/internal/gateway/decision-stream",
+            any(move |_request: Request| {
+                let decision_hits_inner = Arc::clone(&decision_hits_clone);
+                async move {
+                    *decision_hits_inner.lock().expect("mutex should lock") += 1;
+                    Json(json!({"action": "proxy_public"}))
+                }
+            }),
+        )
+        .route(
+            "/api/internal/gateway/plan-stream",
+            any(move |_request: Request| {
+                let plan_hits_inner = Arc::clone(&plan_hits_clone);
+                async move {
+                    *plan_hits_inner.lock().expect("mutex should lock") += 1;
+                    Json(json!({"action": "proxy_public"}))
+                }
+            }),
+        )
+        .route(
+            "/api/internal/gateway/report-stream",
+            any(move |request: Request| {
+                let seen_report_inner = Arc::clone(&seen_report_clone);
+                async move {
+                    let (_parts, body) = request.into_parts();
+                    let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
+                    *seen_report_inner.lock().expect("mutex should lock") = true;
+                    Json(json!({"ok": true}))
+                }
+            }),
+        )
+        .route(
+            "/v1/chat/completions",
+            any(move |_request: Request| {
+                let public_hits_inner = Arc::clone(&public_hits_clone);
+                async move {
+                    *public_hits_inner.lock().expect("mutex should lock") += 1;
+                    (StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
+                }
+            }),
+        );
+
+    let execution_runtime = Router::new().route(
+        "/v1/execute/stream",
+        any(move |request: Request| {
+            let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
+            async move {
+                let (parts, body) = request.into_parts();
+                let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
+                let payload: serde_json::Value =
+                    serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
+                *seen_execution_runtime_inner
+                    .lock()
+                    .expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest {
+                    trace_id: parts
+                        .headers
+                        .get(TRACE_ID_HEADER)
+                        .and_then(|value| value.to_str().ok())
+                        .unwrap_or_default()
+                        .to_string(),
+                    url: payload
+                        .get("url")
+                        .and_then(|value| value.as_str())
+                        .unwrap_or_default()
+                        .to_string(),
+                    provider_model: payload
+                        .get("body")
+                        .and_then(|value| value.get("json_body"))
+                        .and_then(|value| value.get("model"))
+                        .and_then(|value| value.as_str())
+                        .unwrap_or_default()
+                        .to_string(),
+                    auth_header_value: payload
+                        .get("headers")
+                        .and_then(|value| value.get("x-goog-api-key"))
+                        .and_then(|value| value.as_str())
+                        .unwrap_or_default()
+                        .to_string(),
+                    accept: payload
+                        .get("headers")
+                        .and_then(|value| value.get("accept"))
+                        .and_then(|value| value.as_str())
+                        .unwrap_or_default()
+                        .to_string(),
+                    endpoint_tag: payload
+                        .get("headers")
+                        .and_then(|value| value.get("x-endpoint-tag"))
+                        .and_then(|value| value.as_str())
+                        .unwrap_or_default()
+                        .to_string(),
+                });
+                let frames = concat!(
+                    "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
+                    "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello \\\"}],\\\"role\\\":\\\"model\\\"},\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\"}\\n\\n\"}}\n",
+                    "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"responseId\\\":\\\"resp_openai_chat_gemini_stream_123\\\",\\\"candidates\\\":[{\\\"content\\\":{\\\"parts\\\":[{\\\"text\\\":\\\"Hello Gemini stream\\\"}],\\\"role\\\":\\\"model\\\"},\\\"finishReason\\\":\\\"STOP\\\",\\\"index\\\":0}],\\\"modelVersion\\\":\\\"gemini-2.5-pro-upstream\\\",\\\"usageMetadata\\\":{\\\"promptTokenCount\\\":1,\\\"candidatesTokenCount\\\":2,\\\"totalTokenCount\\\":3}}\\n\\n\"}}\n",
+                    "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n",
+                    "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
+                );
+                let mut response = Response::builder()
+                    .status(StatusCode::OK)
+                    .body(Body::from(frames))
+                    .expect("response should build");
+                response.headers_mut().insert(
+                    http::header::CONTENT_TYPE,
+                    HeaderValue::from_static("application/x-ndjson"),
+                );
+                response
+            }
+        }),
+    );
+
+    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
+        Some(hash_api_key("sk-client-openai-chat-gemini-stream")),
+        sample_auth_snapshot(
+            "api-key-openai-chat-gemini-stream-local-1",
+            "user-openai-chat-gemini-stream-local-1",
+        ),
+    )]));
+    let candidate_selection_repository =
+        Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+            sample_candidate_row(),
+        ]));
+    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![sample_provider_catalog_provider()],
+        vec![sample_provider_catalog_endpoint()],
+        vec![sample_provider_catalog_key()],
+    ));
+
+    let (upstream_url, upstream_handle) = start_server(upstream).await;
+    let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
+    let gateway_state =
+        build_state_with_execution_runtime_override(execution_runtime_url.clone())
+    .with_data_state_for_tests(
+        crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
+            auth_repository,
+            candidate_selection_repository,
+            provider_catalog_repository,
+            Arc::clone(&request_candidate_repository),
+            DEVELOPMENT_ENCRYPTION_KEY,
+        )
+        .with_system_config_values_for_tests(vec![(
+            "provider_priority_mode".to_string(),
+            json!("global_key"),
+        )]),
+    );
+    let gateway = build_router_with_state(gateway_state);
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!("{gateway_url}/v1/chat/completions"))
+        .header(http::header::CONTENT_TYPE, "application/json")
+        .header(
+            http::header::AUTHORIZATION,
+            "Bearer sk-client-openai-chat-gemini-stream",
+        )
+        .header(TRACE_ID_HEADER, "trace-openai-chat-gemini-stream-123")
+        .body(
+            "{\"model\":\"gpt-5\",\"messages\":[{\"role\":\"system\",\"content\":\"You are terse.\"},{\"role\":\"user\",\"content\":\"Say hello\"}],\"stream\":true}",
+        )
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response
+            .headers()
+            .get(EXECUTION_PATH_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
+    );
+    let response_text = response.text().await.expect("body should read");
+    assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
+    assert!(response_text.contains("data: [DONE]"));
+
+    let seen_execution_runtime_request = seen_execution_runtime
+        .lock()
+        .expect("mutex should lock")
+        .clone()
+        .expect("execution runtime stream should be captured");
+    assert_eq!(
+        seen_execution_runtime_request.trace_id,
+        "trace-openai-chat-gemini-stream-123"
+    );
+    assert_eq!(
+        seen_execution_runtime_request.url,
+        "https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-2.5-pro-upstream:streamGenerateContent?alt=sse"
+    );
+    assert_eq!(
+        seen_execution_runtime_request.provider_model,
+        "gemini-2.5-pro-upstream"
+    );
+    assert_eq!(
+        seen_execution_runtime_request.auth_header_value,
+        "sk-upstream-openai-chat-gemini-stream"
+    );
+    assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
+    assert_eq!(
+        seen_execution_runtime_request.endpoint_tag,
+        "openai-chat-gemini-cross-format-stream"
+    );
+
+    let stored_candidates = request_candidate_repository
+        .list_by_request_id("trace-openai-chat-gemini-stream-123")
+        .await
+        .expect("request candidate trace should read");
+    assert_eq!(stored_candidates.len(), 1);
+    assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
+    let extra_data = stored_candidates[0]
+        .extra_data
+        .as_ref()
+        .expect("request candidate extra_data should exist");
+    assert_eq!(extra_data["client_api_format"], "openai:chat");
+    assert_eq!(extra_data["provider_api_format"], "gemini:generate_content");
+
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    assert!(
+        !*seen_report.lock().expect("mutex should lock"),
+        "report-stream should stay local when request candidate persistence is available"
+    );
+
+    assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
+    assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
+    assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+    upstream_handle.abort();
+}
+
 #[tokio::test]
 async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision_gate_with_local_stream_decision(
 ) {

apps/aether-gateway/src/tests/files/registry_cleanup.rs

@@ -72,7 +72,5 @@ async fn gateway_background_gemini_file_mapping_cleanup_deletes_expired_entries(
         "active mapping should remain"
     );
 
-    for handle in background_tasks {
-        handle.abort();
-    }
+    background_tasks.shutdown().await;
 }

apps/aether-gateway/src/tests/video/registry_poller.rs

@@ -231,9 +231,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
         }]
     );
 
-    for handle in background_tasks {
-        handle.abort();
-    }
+    background_tasks.shutdown().await;
     execution_runtime_handle.abort();
 }
 
@@ -321,8 +319,6 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
         }]
     );
 
-    for handle in background_tasks {
-        handle.abort();
-    }
+    background_tasks.shutdown().await;
     upstream_handle.abort();
 }