Add Codex image progress heartbeat

This commit is contained in:
Entropy.Xu
2026-05-09 01:21:26 +08:00
parent 9a84a6ff6c
commit 3b4f27f767
44 changed files with 2676 additions and 208 deletions

View File

@@ -20,6 +20,7 @@ use aether_data_contracts::repository::provider_catalog::{
use sha2::{Digest, Sha256};
use crate::data::GatewayDataState;
use crate::tests::next_non_keepalive_chunk;
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
@@ -370,11 +371,7 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects() {
assert_eq!(response.status(), StatusCode::OK);
let mut response = response;
let first_chunk = response
.chunk()
.await
.expect("first chunk should read")
.expect("first chunk should exist");
let first_chunk = next_non_keepalive_chunk(&mut response).await;
assert_eq!(
first_chunk,
Bytes::from_static(b"data: {\"id\":\"chatcmpl-first\"}\n\n")

View File

@@ -18,9 +18,10 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
VideoTaskTruthSourceMode,
};
mod control_execute;

View File

@@ -1,8 +1,8 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response,
Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER,
TRACE_ID_HEADER,
strip_sse_keepalive_comments, to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue,
Json, Mutex, Request, Response, Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -391,7 +391,7 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-123\"}\n\ndata: [DONE]\n\n"
);
@@ -1352,7 +1352,7 @@ async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-custom-path-123\"}\n\ndata: [DONE]\n\n"
);
@@ -1863,7 +1863,7 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-stream-failover-123\"}\n\ndata: [DONE]\n\n"
);

View File

@@ -18,9 +18,10 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
VideoTaskTruthSourceMode,
};
mod decision;

View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
strip_sse_keepalive_comments, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible,
Json, Mutex, Request, Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -469,7 +469,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_codex_cli_stream_local_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
);

View File

@@ -18,9 +18,9 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
mod compact;

View File

@@ -1,7 +1,8 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router,
StatusCode, TRACE_ID_HEADER,
any, build_router_with_state, build_state_with_execution_runtime_override, json,
next_non_keepalive_chunk, start_server, strip_sse_keepalive_comments, to_bytes, Arc, Body,
Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode,
TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -952,11 +953,12 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
tokio::time::timeout(std::time::Duration::from_millis(100), response.chunk())
.await
.expect("same-format passthrough should yield first chunk before eof")
.expect("first chunk should read")
.expect("first chunk should exist"),
tokio::time::timeout(
std::time::Duration::from_millis(100),
next_non_keepalive_chunk(&mut response),
)
.await
.expect("same-format passthrough should yield first chunk before eof"),
Bytes::from_static(b"event: message_start\ndata: {\"type\":\"message_start\"}\n\n")
);
assert_eq!(
@@ -1470,7 +1472,7 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
);
@@ -1924,7 +1926,7 @@ async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
);

View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
@@ -404,7 +404,7 @@ async fn gateway_executes_gemini_chat_stream_via_local_decision_gate_with_local_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);

View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
@@ -381,7 +381,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_with_local_s
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -872,7 +872,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -1339,7 +1339,7 @@ async fn gateway_executes_vertex_ai_gemini_cli_stream_via_local_decision_gate_wi
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -1849,7 +1849,8 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let response_text = response.text().await.expect("body should read");
let response_text =
strip_sse_keepalive_comments(&response.text().await.expect("body should read"));
let payload = response_text
.trim()
.strip_prefix("data: ")

View File

@@ -18,9 +18,9 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{