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): 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!({

84
crates/aether-usage-runtime/src/write.rs
View File

@@ -158,6 +158,7 @@ pub struct StreamTerminalUsagePayloadSeed {
pub client_response: Option<Value>,
pub client_response_body_state: Option<UsageBodyCaptureState>,
pub standardized_usage: Option<StandardizedUsage>,
pub observed_stream_finish: Option<bool>,
pub capture_metadata: Option<Value>,
}
@@ -779,6 +780,10 @@ pub fn build_stream_terminal_usage_payload_seed(
let provider_response_headers = context_usage_value(context, "provider_response_headers")
.or_else(|| headers_to_json(&payload.headers));
let client_response_headers = headers_to_json(&payload.headers);
let observed_stream_finish = payload
.terminal_summary
.as_ref()
.map(|summary| summary.observed_finish);
StreamTerminalUsagePayloadSeed {
report_kind: payload.report_kind.clone(),
status_code: payload.status_code,
@@ -797,6 +802,7 @@ pub fn build_stream_terminal_usage_payload_seed(
.terminal_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.clone()),
observed_stream_finish,
capture_metadata: build_payload_body_capture_metadata(
payload.provider_body_base64.as_deref(),
payload.client_body_base64.as_deref(),
@@ -894,6 +900,7 @@ pub fn build_stream_terminal_usage_seed(
client_response,
client_response_body_state,
standardized_usage,
observed_stream_finish,
capture_metadata,
} = payload_seed;
let standardized_usage = standardized_usage.or_else(|| {
@@ -901,7 +908,16 @@ pub fn build_stream_terminal_usage_seed(
map_usage_from_response(response, context_seed.provider_contract.as_str())
})
});
let terminal_state = infer_stream_terminal_state(report_kind.as_str(), status_code, cancelled);
let missing_observed_finish = matches!(observed_stream_finish, Some(false))
&& !standardized_usage
.as_ref()
.is_some_and(StandardizedUsage::has_token_signal);
let terminal_state = infer_stream_terminal_state(
report_kind.as_str(),
status_code,
cancelled,
missing_observed_finish,
);
TerminalUsageSeed {
terminal_state,
@@ -970,10 +986,11 @@ fn infer_stream_terminal_state(
report_kind: &str,
status_code: u16,
cancelled: bool,
missing_observed_finish: bool,
) -> UsageTerminalState {
if cancelled || status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if !(200..300).contains(&status_code) {
} else if !(200..300).contains(&status_code) || missing_observed_finish {
UsageTerminalState::Failed
} else {
UsageTerminalState::Completed
@@ -3811,6 +3828,69 @@ mod tests {
assert!(event.data.client_response_body.is_none());
}
#[test]
fn stream_terminal_usage_marks_missing_observed_finish_as_failed() {
let plan = ExecutionPlan {
request_id: "req-stream-missing-finish-1".to_string(),
candidate_id: Some("cand-stream-missing-finish-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/responses".to_string(),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-missing-finish-1".to_string(),
report_kind: "openai_responses_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:responses",
"provider_api_format": "openai:responses"
})),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: Some(UsageBodyCaptureState::None),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
response_id: Some("resp_missing_finish".to_string()),
model: Some("gpt-5.5".to_string()),
observed_finish: false,
..ExecutionStreamTerminalSummary::default()
}),
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Failed);
assert_eq!(event.data.status_code, Some(200));
assert_eq!(
event.data.error_category.as_deref(),
Some("non_success_status")
);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
}
#[test]
fn completed_image_usage_estimates_request_tokens_when_provider_usage_is_missing() {
let plan = ExecutionPlan {