i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): stop stream polling on downstream disconnect and preserve Codex cache keys

Browse Source
This commit is contained in:
zhefox
2026-05-21 22:26:14 +08:00
parent 77c2d91eb0
commit 68d8f86dc6
3 changed files with 35 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -2746,6 +2746,11 @@ async fn execute_stream_from_frame_stream(
image_stream_total_timeout.as_mut() image_stream_total_timeout.as_mut()
{ {
tokio::select! { tokio::select! {
biased;
_ = tx.closed(), if client_visible_stream_completed => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result, result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
_ = timeout_sleep.as_mut() => { _ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms let timeout_ms = openai_image_stream_total_timeout_ms
@@ -2797,7 +2802,14 @@ async fn execute_stream_from_frame_stream(
} }
} }
} else { } else {
next_stream_frame(&mut buffered_frames, &mut lines).await tokio::select! {
biased;
_ = tx.closed(), if client_visible_stream_completed => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
}
}; };
let next_frame = match next_frame_result { let next_frame = match next_frame_result {
Ok(frame) => frame, Ok(frame) => frame,

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

@@ -884,15 +884,15 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
); );
assert_eq!( assert_eq!(
seen_execution_runtime_request.prompt_cache_key, seen_execution_runtime_request.prompt_cache_key,
"b6741389-8b9e-5c00-bef6-fbce92aee45a" "bc749eb7-a9e2-5793-8d14-abd659c700b0"
); );
assert_eq!( assert_eq!(
seen_execution_runtime_request.session_id, seen_execution_runtime_request.session_id,
"9fa08f4f14ccba13" "d1e9b802644e1f52"
); );
assert_eq!( assert_eq!(
seen_execution_runtime_request.conversation_id, seen_execution_runtime_request.conversation_id,
"9fa08f4f14ccba13" "d1e9b802644e1f52"
); );
assert_eq!( assert_eq!(
seen_execution_runtime_request.instructions, seen_execution_runtime_request.instructions,

27
crates/aether-ai-formats/src/formats/openai/responses/codex.rs
View File

@@ -532,14 +532,14 @@ fn maybe_insert_default_codex_header(
provider_request_headers.insert(header_name.to_string(), header_value.to_string()); provider_request_headers.insert(header_name.to_string(), header_value.to_string());
} }
fn maybe_inject_codex_prompt_cache_key( fn codex_prompt_cache_key_to_insert(
provider_request_body: &mut Value, provider_request_body: &Value,
provider_type: &str, provider_type: &str,
provider_api_format: &str, provider_api_format: &str,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
) { ) -> Option<String> {
if !is_codex_openai_responses_request(provider_type, provider_api_format) { if !is_codex_openai_responses_request(provider_type, provider_api_format) {
return; return None;
} }
let existing = provider_request_body let existing = provider_request_body
@@ -548,10 +548,10 @@ fn maybe_inject_codex_prompt_cache_key(
.map(str::trim) .map(str::trim)
.unwrap_or_default(); .unwrap_or_default();
if !existing.is_empty() { if !existing.is_empty() {
return; return None;
} }
let prompt_cache_key = extract_codex_prompt_cache_session_seed(provider_request_body) extract_codex_prompt_cache_session_seed(provider_request_body)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("session", &seed)) .and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("session", &seed))
.or_else(|| { .or_else(|| {
extract_codex_prompt_cache_control_seed(provider_request_body) extract_codex_prompt_cache_control_seed(provider_request_body)
@@ -561,7 +561,13 @@ fn maybe_inject_codex_prompt_cache_key(
extract_codex_stable_request_prompt_cache_seed(provider_request_body, user_api_key_id) extract_codex_stable_request_prompt_cache_seed(provider_request_body, user_api_key_id)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("request", &seed)) .and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("request", &seed))
}) })
.or_else(|| user_api_key_id.and_then(build_stable_codex_prompt_cache_key)); .or_else(|| user_api_key_id.and_then(build_stable_codex_prompt_cache_key))
}
fn insert_codex_prompt_cache_key(
provider_request_body: &mut Value,
prompt_cache_key: Option<String>,
) {
let Some(prompt_cache_key) = prompt_cache_key else { let Some(prompt_cache_key) = prompt_cache_key else {
return; return;
}; };
@@ -712,7 +718,7 @@ pub fn apply_codex_openai_responses_special_body_edits(
return; return;
} }
maybe_inject_codex_prompt_cache_key( let prompt_cache_key = codex_prompt_cache_key_to_insert(
provider_request_body, provider_request_body,
provider_type, provider_type,
provider_api_format, provider_api_format,
@@ -766,6 +772,8 @@ pub fn apply_codex_openai_responses_special_body_edits(
apply_codex_openai_image_tool_overrides(body_object); apply_codex_openai_image_tool_overrides(body_object);
inject_codex_default_variation_prompt(body_object); inject_codex_default_variation_prompt(body_object);
} }
insert_codex_prompt_cache_key(provider_request_body, prompt_cache_key);
} }
pub fn apply_codex_openai_responses_chat_body_edits( pub fn apply_codex_openai_responses_chat_body_edits(
@@ -790,6 +798,9 @@ pub fn apply_codex_openai_responses_chat_body_edits(
return; return;
}; };
ensure_codex_chat_reasoning_defaults(body_object, provider_api_format, body_rules); ensure_codex_chat_reasoning_defaults(body_object, provider_api_format, body_rules);
if let Some(prompt_cache_key) = body_object.remove("prompt_cache_key") {
body_object.insert("prompt_cache_key".to_string(), prompt_cache_key);
}
} }
pub fn apply_codex_openai_responses_special_headers( pub fn apply_codex_openai_responses_special_headers(