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

fix(gateway): sanitize Claude thinking and handle missing stream finish

Browse Source
This commit is contained in:
zhefox
2026-05-21 12:30:42 +08:00
parent 923515ab28
commit e59e6c3797
5 changed files with 612 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
crates/aether-ai-formats/src/formats/openai/chat/stream.rs
View File

@@ -1291,6 +1291,8 @@ pub struct OpenAIChatClientEmitter {
model: Option<String>,
started: bool,
finished: bool,
next_tool_call_index: usize,
tool_call_index_by_canonical: BTreeMap<usize, usize>,
}
#[derive(Clone, Default)]
@@ -1357,6 +1359,17 @@ impl OpenAIChatClientEmitter {
)
}
fn chat_tool_call_index(&mut self, canonical_index: usize) -> usize {
if let Some(index) = self.tool_call_index_by_canonical.get(&canonical_index) {
return *index;
}
let index = self.next_tool_call_index;
self.next_tool_call_index += 1;
self.tool_call_index_by_canonical
.insert(canonical_index, index);
index
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
@@ -1465,6 +1478,7 @@ impl OpenAIChatClientEmitter {
name,
} => {
let mut out = self.ensure_started()?;
let chat_index = self.chat_tool_call_index(index);
out.extend(encode_json_sse(
None,
&build_openai_chat_chunk(
@@ -1474,7 +1488,7 @@ impl OpenAIChatClientEmitter {
self.model.as_deref().unwrap_or("unknown"),
String::new(),
Some(vec![json!({
"index": index,
"index": chat_index,
"id": call_id,
"type": "function",
"function": {
@@ -1489,6 +1503,7 @@ impl OpenAIChatClientEmitter {
}
CanonicalStreamEvent::ToolCallArgumentsDelta { index, arguments } => {
let mut out = self.ensure_started()?;
let chat_index = self.chat_tool_call_index(index);
out.extend(encode_json_sse(
None,
&json!({
@@ -1501,7 +1516,7 @@ impl OpenAIChatClientEmitter {
"index": 0,
"delta": {
"tool_calls": [{
"index": index,
"index": chat_index,
"function": {
"arguments": arguments,
}
@@ -2588,6 +2603,27 @@ mod tests {
parts
}
fn openai_chat_tool_call_indices(sse: &str) -> Vec<u64> {
let mut indices = Vec::new();
for payload in sse.lines().filter_map(|line| line.strip_prefix("data: ")) {
let Ok(value) = serde_json::from_str::<Value>(payload) else {
continue;
};
let Some(tool_calls) = value
.pointer("/choices/0/delta/tool_calls")
.and_then(Value::as_array)
else {
continue;
};
for tool_call in tool_calls {
if let Some(index) = tool_call.get("index").and_then(Value::as_u64) {
indices.push(index);
}
}
}
indices
}
#[test]
fn openai_chat_provider_state_emits_unknown_events_for_unrecognized_deltas() {
let mut state = OpenAIChatProviderState::default();
@@ -3250,6 +3286,50 @@ mod tests {
assert!(sse.contains("[Image]"));
}
#[test]
fn openai_chat_client_emitter_normalizes_sparse_tool_call_indices() {
let mut emitter = OpenAIChatClientEmitter::default();
let mut bytes = Vec::new();
for event in [
CanonicalStreamEvent::ToolCallStart {
index: 1,
call_id: "call_first".to_string(),
name: "first_tool".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 1,
arguments: "{\"first\":".to_string(),
},
CanonicalStreamEvent::ToolCallStart {
index: 3,
call_id: "call_second".to_string(),
name: "second_tool".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 3,
arguments: "{\"second\":true}".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 1,
arguments: "true}".to_string(),
},
] {
bytes.extend(
emitter
.emit(CanonicalStreamFrame {
id: "chatcmpl_sparse".to_string(),
model: "claude-opus-4-6".to_string(),
event,
})
.expect("tool event should encode"),
);
}
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert_eq!(openai_chat_tool_call_indices(&sse), vec![0, 0, 1, 1, 0]);
}
#[test]
fn openai_chat_client_emitter_emits_usage_only_final_chunk() {
let mut emitter = OpenAIChatClientEmitter::default();

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

@@ -165,14 +165,25 @@ fn inject_codex_default_variation_prompt(body_object: &mut serde_json::Map<Strin
);
}
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
let normalized = user_api_key_id.trim();
fn build_stable_codex_prompt_cache_key_from_seed(kind: &str, seed: &str) -> Option<String> {
let normalized = seed.trim();
if normalized.is_empty() {
return None;
}
let normalized_kind = kind
.trim()
.to_ascii_lowercase()
.chars()
.filter(|ch| ch.is_ascii_alphanumeric() || *ch == '_' || *ch == '-')
.collect::<String>();
let normalized_kind = if normalized_kind.is_empty() {
"seed".to_string()
} else {
normalized_kind
};
let namespace = format!(
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:user:{normalized}"
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:{normalized_kind}:{normalized}"
);
let mut hasher = Sha1::new();
hasher.update(UUID_NAMESPACE_OID_BYTES);
@@ -186,6 +197,48 @@ fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String>
Some(Uuid::from_bytes(bytes).to_string())
}
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
build_stable_codex_prompt_cache_key_from_seed("user", user_api_key_id)
}
fn extract_codex_prompt_cache_session_seed(provider_request_body: &Value) -> Option<String> {
fn non_empty_str(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn session_seed_from_metadata(metadata: &Value) -> Option<String> {
let object = metadata.as_object()?;
non_empty_str(object.get("session_id"))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.map(|value| format!("metadata:{value}"))
.or_else(|| {
let user_id = non_empty_str(object.get("user_id"))?;
serde_json::from_str::<Value>(user_id)
.ok()
.and_then(|decoded| {
non_empty_str(decoded.get("session_id"))
.or_else(|| non_empty_str(decoded.get("sessionId")))
.or_else(|| non_empty_str(decoded.get("conversation_id")))
.or_else(|| non_empty_str(decoded.get("conversationId")))
.map(|value| format!("metadata.user_id:{value}"))
})
})
}
let object = provider_request_body.as_object()?;
non_empty_str(object.get("session_id"))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.map(|value| format!("body:{value}"))
.or_else(|| object.get("metadata").and_then(session_seed_from_metadata))
}
fn build_short_codex_header_id(seed: &str) -> Option<String> {
let normalized = seed.trim();
if normalized.is_empty() {
@@ -271,11 +324,7 @@ fn maybe_inject_codex_prompt_cache_key(
return;
}
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
let existing = body_object
let existing = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
@@ -284,8 +333,14 @@ fn maybe_inject_codex_prompt_cache_key(
return;
}
let Some(prompt_cache_key) = user_api_key_id.and_then(build_stable_codex_prompt_cache_key)
else {
let prompt_cache_key = extract_codex_prompt_cache_session_seed(provider_request_body)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("session", &seed))
.or_else(|| user_api_key_id.and_then(build_stable_codex_prompt_cache_key));
let Some(prompt_cache_key) = prompt_cache_key else {
return;
};
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
@@ -431,6 +486,13 @@ pub fn apply_codex_openai_responses_special_body_edits(
return;
}
maybe_inject_codex_prompt_cache_key(
provider_request_body,
provider_type,
provider_api_format,
user_api_key_id,
);
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
@@ -478,13 +540,6 @@ pub fn apply_codex_openai_responses_special_body_edits(
apply_codex_openai_image_tool_overrides(body_object);
inject_codex_default_variation_prompt(body_object);
}
maybe_inject_codex_prompt_cache_key(
provider_request_body,
provider_type,
provider_api_format,
user_api_key_id,
);
}
pub fn apply_codex_openai_responses_chat_body_edits(
@@ -748,6 +803,57 @@ mod tests {
);
}
#[test]
fn codex_responses_body_edits_derive_prompt_cache_key_from_session_metadata() {
let mut body_a = json!({
"input": [{"role": "user", "content": "hello"}],
"model": "gpt-5.4",
"metadata": {
"user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-a\"}"
}
});
let mut body_b = json!({
"input": [{"role": "user", "content": "hello again"}],
"model": "gpt-5.4",
"metadata": {
"user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-b\"}"
}
});
let mut body_c = json!({
"input": [{"role": "user", "content": "hello"}],
"model": "gpt-5.4",
"metadata": {"session_id": "session-b"}
});
apply_codex_openai_responses_special_body_edits(
&mut body_a,
"codex",
"openai:responses",
None,
Some("key-123"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_b,
"codex",
"openai:responses",
None,
Some("different-key"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_c,
"codex",
"openai:responses",
None,
Some("key-123"),
);
assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]);
assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]);
assert!(body_a.get("metadata").is_none());
assert!(body_b.get("metadata").is_none());
assert!(body_c.get("metadata").is_none());
}
#[test]
fn compact_body_edits_strip_include_store_and_stream() {
let mut provider_request_body = json!({