i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +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()
{
tokio::select! {
biased;
_ = tx.closed(), if client_visible_stream_completed => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
_ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms
@@ -2797,7 +2802,14 @@ async fn execute_stream_from_frame_stream(
}
}
} 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 {
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!(
seen_execution_runtime_request.prompt_cache_key,
"b6741389-8b9e-5c00-bef6-fbce92aee45a"
"bc749eb7-a9e2-5793-8d14-abd659c700b0"
);
assert_eq!(
seen_execution_runtime_request.session_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
seen_execution_runtime_request.conversation_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
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());
}
fn maybe_inject_codex_prompt_cache_key(
provider_request_body: &mut Value,
fn codex_prompt_cache_key_to_insert(
provider_request_body: &Value,
provider_type: &str,
provider_api_format: &str,
user_api_key_id: Option<&str>,
) {
) -> Option<String> {
if !is_codex_openai_responses_request(provider_type, provider_api_format) {
return;
return None;
}
let existing = provider_request_body
@@ -548,10 +548,10 @@ fn maybe_inject_codex_prompt_cache_key(
.map(str::trim)
.unwrap_or_default();
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))
.or_else(|| {
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)
.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 {
return;
};
@@ -712,7 +718,7 @@ pub fn apply_codex_openai_responses_special_body_edits(
return;
}
maybe_inject_codex_prompt_cache_key(
let prompt_cache_key = codex_prompt_cache_key_to_insert(
provider_request_body,
provider_type,
provider_api_format,
@@ -766,6 +772,8 @@ pub fn apply_codex_openai_responses_special_body_edits(
apply_codex_openai_image_tool_overrides(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(
@@ -790,6 +798,9 @@ pub fn apply_codex_openai_responses_chat_body_edits(
return;
};
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(