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

fix(ws): harden Responses connection lifecycle

Browse Source 
Revalidate control policy per turn, isolate downstream credentials, and make planner/turn ownership cancellation-safe.

Preserve opaque protocol events, align configurable timeout semantics, and extend end-to-end security and settlement coverage.
This commit is contained in:
ZheFox
2026-08-17 18:50:29 +08:00
parent 4a0775c4ea
commit c8118edf36
42 changed files with 4017 additions and 1276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/standard/mod.rs
+1
View File
@@ -50,6 +50,7 @@ pub(crate) use self::openai::{
maybe_build_sync_local_openai_responses_decision_payload, parse_openai_stop_sequences,
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
value_as_u64, ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
ResponsesWebSocketPinnedCandidate,
};
pub(crate) use crate::ai_serving::normalize_standard_request_to_openai_chat_request;
pub(crate) use crate::ai_serving::{
apps/aether-gateway/src/ai_serving/planner/standard/openai/mod.rs
+1 -1
View File
@@ -26,5 +26,5 @@ pub(crate) use responses::{
maybe_build_responses_websocket_decision,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, ResponsesWebSocketBodyNormalization,
ResponsesWebSocketDecision,
ResponsesWebSocketDecision, ResponsesWebSocketPinnedCandidate,
};
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs
+4 -2
View File
@@ -9,7 +9,9 @@ pub(super) use self::payload::maybe_build_local_openai_responses_decision_payloa
pub(super) use self::support::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
resolve_local_openai_responses_decision_input,
resolve_local_openai_responses_decision_input_with_snapshot,
LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesCandidateAttemptSource,
LocalOpenAiResponsesDecisionInput,
};
pub(super) use crate::ai_serving::LocalOpenAiResponsesSpec;
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs
+40 -11
View File
@@ -22,6 +22,7 @@ use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::planner::decision_input::{
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
resolve_local_authenticated_decision_input_with_snapshot,
};
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
@@ -32,7 +33,8 @@ use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
openai_responses_request_operation, resolve_local_decision_execution_runtime_auth_context,
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayControlDecision,
PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
@@ -51,6 +53,21 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
resolve_local_openai_responses_decision_input_with_snapshot(
state, parts, trace_id, decision, body_json, plan_kind, None,
)
.await
}
pub(crate) async fn resolve_local_openai_responses_decision_input_with_snapshot(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
auth_snapshot_override: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
warn!(
@@ -87,16 +104,28 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
return Ok(None);
};
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
{
let resolved_input = match if let Some(auth_snapshot) = auth_snapshot_override {
resolve_local_authenticated_decision_input_with_snapshot(
state,
auth_context.clone(),
Some(auth_snapshot.clone()),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
} else {
resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None,
&decision.model_directive_policy,
)
.await
} {
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => {
warn!(
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
+126 -23
View File
@@ -12,6 +12,49 @@ use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_runtime_state::RuntimeLockLease;
use std::collections::BTreeSet;
/// Releases a scheduler pool-key lease if WebSocket planning is cancelled
/// after candidate selection but before ownership reaches the turn lifecycle.
struct ResponsesWebSocketPlanningLeaseGuard {
state: AppState,
lease: Option<RuntimeLockLease>,
}
impl ResponsesWebSocketPlanningLeaseGuard {
fn new(state: &AppState, lease: Option<&RuntimeLockLease>) -> Self {
Self {
state: state.clone(),
lease: lease.cloned(),
}
}
async fn release(mut self) {
// Keep the lease armed across the await. If the owner task is aborted
// or reaches its hard deadline while the runtime backend is stalled,
// Drop can still hand cleanup to a detached owner.
if release_responses_websocket_planning_lease(&self.state, self.lease.as_ref()).await {
self.lease = None;
}
}
fn disarm(&mut self) {
self.lease = None;
}
}
impl Drop for ResponsesWebSocketPlanningLeaseGuard {
fn drop(&mut self) {
let Some(lease) = self.lease.take() else {
return;
};
let state = self.state.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = release_responses_websocket_planning_lease(&state, Some(&lease)).await;
});
}
}
}
mod decision;
mod plans;
@@ -19,6 +62,7 @@ use self::decision::{
build_local_openai_responses_candidate_attempt_source,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input,
resolve_local_openai_responses_decision_input_with_snapshot,
};
use self::plans::{
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
@@ -187,6 +231,45 @@ pub(crate) struct ResponsesWebSocketDecision {
pub(crate) normalization: ResponsesWebSocketBodyNormalization,
}
/// The scheduler identity a continuation is allowed to reuse.
///
/// A `previous_response_id` chain cannot move to another provider connection,
/// but it still has to pass the current scheduler runtime checks on every
/// turn. The planner uses this identity as a filter rather than selecting an
/// arbitrary eligible replacement.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ResponsesWebSocketPinnedCandidate {
provider_id: String,
endpoint_id: String,
key_id: String,
}
impl ResponsesWebSocketPinnedCandidate {
pub(crate) fn from_decision(decision: &AiExecutionDecision) -> Option<Self> {
Some(Self {
provider_id: non_empty_decision_identity(decision.provider_id.as_deref())?,
endpoint_id: non_empty_decision_identity(decision.endpoint_id.as_deref())?,
key_id: non_empty_decision_identity(decision.key_id.as_deref())?,
})
}
fn matches(
&self,
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
candidate.provider_id == self.provider_id
&& candidate.endpoint_id == self.endpoint_id
&& candidate.key_id == self.key_id
}
}
fn non_empty_decision_identity(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
}
/// Everything needed to re-run provider-body normalization for the candidate a
/// socket is already bound to.
///
@@ -321,21 +404,24 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
auth_snapshot: Option<&crate::ai_serving::GatewayAuthApiKeySnapshot>,
body_json: &serde_json::Value,
excluded_key_ids: Option<&BTreeSet<String>>,
excluded_codex_account_ids: Option<&BTreeSet<String>>,
pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>,
) -> Result<Option<ResponsesWebSocketDecision>, GatewayError> {
let Some(spec) = resolve_stream_spec(crate::ai_serving::OPENAI_RESPONSES_STREAM_PLAN_KIND)
else {
return Ok(None);
};
let Some(input) = resolve_local_openai_responses_decision_input(
let Some(input) = resolve_local_openai_responses_decision_input_with_snapshot(
state,
parts,
trace_id,
decision,
body_json,
spec.decision_kind,
auth_snapshot,
)
.await?
else {
@@ -348,18 +434,28 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
.await?;
while let Some(attempt) = source.next_attempt().await? {
let pool_key_lease = attempt.eligible.orchestration.pool_key_lease.clone();
// `next_attempt` may return with a distributed pool-key lease. Arm a
// guard before the first await so owner-task timeout/cancellation
// cannot strand that lease until its server-side TTL expires.
let mut planning_lease = ResponsesWebSocketPlanningLeaseGuard::new(
state,
attempt.eligible.orchestration.pool_key_lease.as_ref(),
);
if pinned_candidate.is_some_and(|pinned| !pinned.matches(&attempt.eligible.candidate)) {
planning_lease.release().await;
continue;
}
if excluded_key_ids
.is_some_and(|key_ids| key_ids.contains(attempt.eligible.candidate.key_id.as_str()))
{
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
continue;
}
let Some(adapter) = responses_websocket_adapter(
&attempt.eligible.transport.provider.provider_type,
attempt.eligible.transport.provider.config.as_ref(),
) else {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
continue;
};
// Captured before `attempt` is consumed so a later continuation turn can
@@ -373,11 +469,11 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
{
Ok(Some(payload)) => payload,
Ok(None) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
continue;
}
Err(error) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
return Err(error);
}
};
@@ -393,7 +489,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
.is_some_and(|account_ids| account_ids.contains(account_id))
})
{
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
continue;
}
match codex_quota_breaker_blocks_candidate(
@@ -405,7 +501,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
.await
{
Ok(true) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
continue;
}
Ok(false) => {}
@@ -454,13 +550,17 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
.flatten(),
mapped_model,
};
return Ok(Some(ResponsesWebSocketDecision {
let decision = ResponsesWebSocketDecision {
execution: payload,
adapter,
normalization,
}));
};
// The decision report context now carries the lease identity. The
// WebSocket ownership layer takes over before any further await.
planning_lease.disarm();
return Ok(Some(decision));
}
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
planning_lease.release().await;
}
Ok(None)
@@ -469,20 +569,23 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
async fn release_responses_websocket_planning_lease(
state: &AppState,
lease: Option<&RuntimeLockLease>,
) {
) -> bool {
let Some(lease) = lease else {
return;
return true;
};
if let Err(error) =
crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease(
state.runtime_state.as_ref(),
lease,
)
.await
match crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease(
state.runtime_state.as_ref(),
lease,
)
.await
{
tracing::warn!(
error = ?error,
"gateway Responses WebSocket planner failed to release an unused pool key lease"
);
Ok(_) => true,
Err(error) => {
tracing::warn!(
error = ?error,
"gateway Responses WebSocket planner failed to release an unused pool key lease"
);
false
}
}
}