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

fix(ws): harden Responses continuation state

Browse Source
This commit is contained in:
ZheFox
2026-08-20 08:51:16 +08:00
parent bef282cfee
commit 654f798d25
45 changed files with 6403 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs
+132 -7
View File
@@ -13,7 +13,9 @@ use super::ownership::{
await_owned_responses_websocket_plan, begin_responses_websocket_turn_with_planned_lease,
spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision,
};
use super::request::{build_planning_parts, planned_response_create_event};
use super::request::{
build_planning_parts, planned_response_create_event, ResponsesLiteStaticConfig,
};
use super::state::BoundResponsesConnection;
use super::turn::{prepare_responses_websocket_turn_decision, ResponsesWebSocketTurnOutcome};
use super::upstream::{bind_responses_upstream, close_bound_upstream};
@@ -102,6 +104,14 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
context: &WebSocketRequestContext,
_previous_settled: PreviousAttemptSettled,
) -> bool {
// `LogicalTurn::client_event` is intentionally redacted before it is
// retained for replay. The binding, however, keeps the hash of the raw
// client-side Responses Lite tools/instructions so a later continuation
// can compare the client's plaintext configuration before redaction.
// Preserve that chain identity across this transparent rebind instead of
// replacing it with the hash `bind_responses_upstream` derives from the
// redacted replay event.
let responses_lite_static_config = bound.responses_lite_static_config.clone();
let Some(active) = bound.turn_state.logical_mut() else {
return false;
};
@@ -213,12 +223,14 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
);
return false;
}
let provider_event = match planned_response_create_event(&decision, &client_event).and_then(
|event| {
serde_json::from_str::<Value>(&event)
.map_err(|_| "response_create_serialization_failed")
},
) {
let provider_event = match planned_response_create_event(
&decision,
&normalization,
&client_event,
)
.and_then(|event| {
serde_json::from_str::<Value>(&event).map_err(|_| "response_create_serialization_failed")
}) {
Ok(event) => event,
Err(code) => {
planned_lease.release().await;
@@ -234,6 +246,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
return false;
}
};
let replacement_provider_store = provider_event.get("store") == Some(&Value::Bool(true));
let turn_decision = prepare_responses_websocket_turn_decision(
&decision,
turn_request_id,
@@ -309,12 +322,20 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
if let Some(mut previous_upstream) = bound.upstream.replace(replacement_upstream) {
close_upstream_socket(&mut previous_upstream, None).await;
}
// The replacement socket has no access to response IDs cached only on
// the exhausted physical connection. Continuation turns are never quota
// replayed, so clearing here cannot discard the active turn's parent.
bound.continuation_response_ids.clear();
let previous_key_id = bound.decision_template.key_id.clone();
bound.adapter = replacement.adapter;
bound.client_model = replacement.client_model;
bound.provider_model = replacement.provider_model;
bound.decision_template = replacement.decision_template;
bound.body_normalization = replacement.body_normalization;
bound.responses_lite_static_config = responses_lite_static_config_after_rebind(
responses_lite_static_config,
replacement.responses_lite_static_config,
);
bound.binding_identity = replacement.binding_identity;
// 同一个 logical turn 的下一个 attempt 就位。状态不符时把 attempt 交回
// drop guard 结算并让调用方走「透明重试失败」分支，不静默丢弃一条已经写了
@@ -323,6 +344,9 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
drop(orphan);
return false;
}
if let Some(logical) = bound.turn_state.logical_mut() {
logical.provider_store = replacement_provider_store;
}
bound.upstream_response_headers = replacement.upstream_response_headers;
bound.pending_adapter_drain = None;
debug!(
@@ -341,6 +365,13 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
true
}
fn responses_lite_static_config_after_rebind(
previous: Option<ResponsesLiteStaticConfig>,
replacement: Option<ResponsesLiteStaticConfig>,
) -> Option<ResponsesLiteStaticConfig> {
replacement.map(|replacement| previous.unwrap_or(replacement))
}
pub(super) fn is_usage_limit_error_event(event: &Value) -> bool {
let is_error = |value: &Value| {
value.get("type").and_then(Value::as_str) == Some("error")
@@ -375,3 +406,97 @@ pub(super) fn mark_active_response_retry_unsafe(
active.mark_retry_unsafe(reason);
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::responses_lite_static_config_after_rebind;
use crate::handlers::proxy::websocket::responses::request::{
prepare_responses_lite_continuation, ResponsesLiteStaticConfig,
};
#[test]
fn quota_rebind_preserves_the_raw_responses_lite_static_hash() {
let raw = json!({
"type": "response.create",
"model": "gpt-5.6-sol",
"instructions": "Contact [email protected] before using the tool",
"tools": [{
"type": "function",
"name": "lookup",
"description": "Look up [email protected]",
"parameters": {"type": "object", "properties": {}}
}],
"input": [{"role": "user", "content": "hello"}]
});
let redacted_replay = json!({
"type": "response.create",
"model": "gpt-5.6-sol",
"instructions": "Contact <EMAIL_2> before using the tool",
"tools": [{
"type": "function",
"name": "lookup",
"description": "Look up <EMAIL_1>",
"parameters": {"type": "object", "properties": {}}
}],
"input": [{"role": "user", "content": "hello"}]
});
let raw_static_config = ResponsesLiteStaticConfig::from_response_create(&raw);
let redacted_static_config =
ResponsesLiteStaticConfig::from_response_create(&redacted_replay);
assert_ne!(raw_static_config, redacted_static_config);
let rebound_static_config = responses_lite_static_config_after_rebind(
Some(raw_static_config.clone()),
Some(redacted_static_config),
)
.expect("the replacement still uses Responses Lite");
assert_eq!(rebound_static_config, raw_static_config);
let continuation = json!({
"type": "response.create",
"model": "gpt-5.6-sol",
"previous_response_id": "resp_after_quota_retry",
"instructions": "Contact [email protected] before using the tool",
"tools": [{
"type": "function",
"name": "lookup",
"description": "Look up [email protected]",
"parameters": {"type": "object", "properties": {}}
}],
"input": [{
"type": "function_call_output",
"call_id": "call_1",
"output": "ok"
}]
});
let prepared = prepare_responses_lite_continuation(&continuation, &rebound_static_config)
.expect("an unchanged plaintext continuation must survive a quota retry");
assert!(prepared.get("tools").is_none());
assert!(prepared.get("instructions").is_none());
}
#[test]
fn quota_rebind_still_tracks_the_replacement_contract() {
let raw = ResponsesLiteStaticConfig::from_response_create(&json!({
"type": "response.create",
"tools": [{"type": "function", "name": "lookup", "parameters": {}}]
}));
let replacement = ResponsesLiteStaticConfig::from_response_create(&json!({
"type": "response.create",
"instructions": "replacement"
}));
assert_eq!(
responses_lite_static_config_after_rebind(Some(raw), None),
None,
"a non-Lite replacement must clear the Lite chain marker"
);
assert_eq!(
responses_lite_static_config_after_rebind(None, Some(replacement.clone())),
Some(replacement),
"a newly selected Lite contract has no earlier raw hash to preserve"
);
}
}