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

feat(gateway): Codex/OpenAI Responses WebSocket 代理模式

Browse Source 
在 /v1/responses 上支持 WebSocket 升级，把客户端帧中继到上游 Codex /
OpenAI Responses WebSocket 端点，同时保持既有的路由、鉴权、配额与用量
语义：

- 路由与准入：control/route/ai.rs 识别 WebSocket 升级请求；
  websocket/ingress.rs 复用 API Key 鉴权、IP 规则与并发许可，并引入
  独立的 WebSocket 连接许可
- 中继：websocket/responses/* 按 connection / session / turn 分层，
  帧解析归一化、socket 写入有界、continuation 保持调度亲和性
- 配额：orchestration/codex_quota_breaker.rs 在账号配额耗尽时熔断并
  自动恢复，不再直接断开客户端连接
- 用量：每个 turn 的终态用量落库，request_metadata 记录
  websocket_mode / websocket_transport，管理端与 usage 视图暴露
  is_websocket
- 管理端：provider 可配置 Responses WebSocket 开关
This commit is contained in:
AAEE86
2026-08-17 14:50:33 +08:00
committed by ZheFox
parent 9a0d346ff3
commit 71b54070e8
72 changed files with 10441 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/websocket/responses/binding.rs
+413
View File
@@ -0,0 +1,413 @@
//! Identity of the physical upstream connection backing a Responses session.
//!
//! A Responses continuation carries state that lives on one provider socket.
//! Comparing only the selected key is therefore not sufficient: transport
//! settings, stable account headers, and the protocol adapter can all change
//! the connection that would receive the next event. Rotating bearer values
//! are intentionally excluded because they do not change an already-upgraded
//! socket's physical binding.
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use aether_contracts::{ProxySnapshot, ResolvedTransportProfile};
use sha2::{Digest, Sha256};
use super::adapter::ResponsesWebSocketProtocolAdapter;
use crate::ai_serving::AiExecutionDecision;
use crate::handlers::proxy::websocket::transport::{
websocket_handshake_headers, websocket_upstream_url,
};
use crate::orchestration::ResponsesWebSocketAdapter;
/// Stable, comparable identity for the actual WebSocket connection target.
///
/// The identity deliberately owns the normalized handshake values rather than
/// retaining a reference to the planner decision. A later re-plan can then
/// be compared without accidentally ignoring a field that changes the
/// physical connection.
#[derive(Clone, PartialEq)]
pub(super) struct UpstreamBindingIdentity {
adapter_kind: ResponsesWebSocketAdapter,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
upstream_url: String,
handshake_headers: BTreeMap<String, String>,
/// Authentication values are not part of a stable key binding when the
/// planner has already supplied a key identity. If that identity is
/// unavailable, retain only a one-way fingerprint so two accounts cannot
/// accidentally share a continuation socket.
auth_fingerprint: Option<[u8; 32]>,
proxy: Option<ProxySnapshot>,
transport_profile: Option<ResolvedTransportProfile>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum UpstreamBindingIdentityError {
MissingUpstreamUrl,
InvalidUpstreamUrl,
InvalidHandshakeHeaders,
}
impl UpstreamBindingIdentity {
/// Builds an identity from the same normalized URL and headers used by
/// the WebSocket transport client.
pub(super) fn from_decision(
adapter: &'static dyn ResponsesWebSocketProtocolAdapter,
decision: &AiExecutionDecision,
) -> Result<Self, UpstreamBindingIdentityError> {
let raw_url = decision
.upstream_url
.as_deref()
.filter(|value| !value.trim().is_empty())
.ok_or(UpstreamBindingIdentityError::MissingUpstreamUrl)?;
let upstream_url = websocket_upstream_url(raw_url, "invalid")
.map_err(|_| UpstreamBindingIdentityError::InvalidUpstreamUrl)?
.to_string();
let headers = websocket_handshake_headers(&decision.provider_request_headers, "invalid")
.map_err(|_| UpstreamBindingIdentityError::InvalidHandshakeHeaders)?;
let authentication_header_names = authentication_header_names(decision);
let mut handshake_headers = BTreeMap::new();
let mut authentication_headers = BTreeMap::new();
for (name, value) in &headers {
let name = name.as_str().to_ascii_lowercase();
let value = value
.to_str()
.map_err(|_| UpstreamBindingIdentityError::InvalidHandshakeHeaders)?;
if authentication_header_names.contains(name.as_str()) {
authentication_headers.insert(name, value.to_string());
} else {
handshake_headers.insert(name, value.to_string());
}
}
let auth_fingerprint = decision
.key_id
.is_none()
.then(|| fingerprint_headers(&authentication_headers))
.filter(|_| !authentication_headers.is_empty());
Ok(Self {
adapter_kind: adapter.kind(),
provider_id: decision.provider_id.clone(),
endpoint_id: decision.endpoint_id.clone(),
key_id: decision.key_id.clone(),
upstream_url,
handshake_headers,
auth_fingerprint,
proxy: effective_proxy_snapshot(decision.proxy.as_ref()),
transport_profile: decision.transport_profile.clone(),
})
}
}
/// Header names that carry credentials in the provider handshake. The
/// planner's explicit `auth_header` extends this list for provider-specific
/// schemes; unknown headers remain part of the stable handshake identity.
fn authentication_header_names(decision: &AiExecutionDecision) -> BTreeSet<String> {
let mut names = BTreeSet::from([
"authorization".to_string(),
"proxy-authorization".to_string(),
"x-api-key".to_string(),
"api-key".to_string(),
"x-goog-api-key".to_string(),
"x-azure-api-key".to_string(),
]);
if let Some(name) = decision
.auth_header
.as_deref()
.map(str::trim)
.filter(|name| !name.is_empty())
{
names.insert(name.to_ascii_lowercase());
}
names
}
fn fingerprint_headers(headers: &BTreeMap<String, String>) -> [u8; 32] {
let mut hasher = Sha256::new();
for (name, value) in headers {
hasher.update((name.len() as u64).to_be_bytes());
hasher.update(name.as_bytes());
hasher.update((value.len() as u64).to_be_bytes());
hasher.update(value.as_bytes());
}
hasher.finalize().into()
}
/// Normalize only values that are provably direct transport. Keep node/tunnel
/// fields even though the current WebSocket builder rejects those proxies: a
/// re-plan must not accidentally reuse an already-bound direct socket for a
/// decision that selected a different proxy topology.
fn effective_proxy_snapshot(proxy: Option<&ProxySnapshot>) -> Option<ProxySnapshot> {
let proxy = proxy?;
if proxy.enabled == Some(false) {
return None;
}
let mut normalized = proxy.clone();
normalized.url = normalized
.url
.take()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
normalized.mode = normalized
.mode
.take()
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
normalized.node_id = normalized
.node_id
.take()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
normalized.label = normalized
.label
.take()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
let has_effective_proxy = normalized.url.is_some()
|| normalized.node_id.is_some()
|| normalized.mode.is_some()
|| normalized.extra.is_some();
has_effective_proxy.then_some(normalized)
}
impl fmt::Debug for UpstreamBindingIdentity {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("UpstreamBindingIdentity")
.field("adapter_kind", &self.adapter_kind)
.field("provider_id", &self.provider_id)
.field("endpoint_id", &self.endpoint_id)
.field("key_id", &self.key_id)
.field("upstream_url", &self.upstream_url)
.field(
"handshake_header_names",
&self.handshake_headers.keys().collect::<Vec<_>>(),
)
.field("proxy_configured", &self.proxy.is_some())
.field(
"transport_profile_id",
&self
.transport_profile
.as_ref()
.map(|profile| profile.profile_id.as_str()),
)
.finish()
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::{UpstreamBindingIdentity, UpstreamBindingIdentityError};
use crate::ai_serving::AiExecutionDecision;
use crate::handlers::proxy::websocket::responses::adapter::resolve_responses_websocket_adapter;
use crate::orchestration::ResponsesWebSocketAdapter;
fn decision() -> AiExecutionDecision {
AiExecutionDecision {
action: "execute".to_string(),
decision_kind: None,
execution_strategy: None,
conversion_mode: None,
request_id: Some("request-1".to_string()),
candidate_id: Some("candidate-1".to_string()),
provider_name: Some("provider".to_string()),
provider_type: Some("openai".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
upstream_base_url: Some("https://api.example.test".to_string()),
upstream_url: Some("https://api.example.test/v1/responses".to_string()),
provider_request_method: Some("POST".to_string()),
auth_header: Some("authorization".to_string()),
auth_value: Some("Bearer secret".to_string()),
provider_api_format: Some("openai:responses".to_string()),
client_api_format: Some("openai:responses".to_string()),
provider_contract: None,
client_contract: None,
model_name: Some("gpt-5.6-sol".to_string()),
mapped_model: None,
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: BTreeMap::from([
("Authorization".to_string(), "Bearer secret".to_string()),
("X-Client".to_string(), "aether".to_string()),
("Connection".to_string(), "keep-alive".to_string()),
]),
provider_request_body: Some(json!({"model": "gpt-5.6-sol"})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
upstream_is_stream: true,
report_kind: None,
report_context: None,
auth_context: None,
}
}
#[test]
fn identity_normalizes_url_and_hop_by_hop_headers() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let identity = UpstreamBindingIdentity::from_decision(adapter, &decision()).unwrap();
assert_eq!(identity.upstream_url, "wss://api.example.test/v1/responses");
assert_eq!(
identity.handshake_headers,
BTreeMap::from([("x-client".to_string(), "aether".to_string())])
);
}
#[test]
fn identity_changes_when_physical_binding_changes() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let base = decision();
let identity = UpstreamBindingIdentity::from_decision(adapter, &base).unwrap();
let codex_adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex);
assert_ne!(
identity,
UpstreamBindingIdentity::from_decision(codex_adapter, &base).unwrap()
);
for mutate in [
|decision: &mut AiExecutionDecision| {
decision.key_id = Some("key-2".to_string());
},
|decision: &mut AiExecutionDecision| {
decision.upstream_url = Some("https://other.example.test/v1/responses".to_string());
},
|decision: &mut AiExecutionDecision| {
decision
.provider_request_headers
.insert("X-Client".to_string(), "other".to_string());
},
|decision: &mut AiExecutionDecision| {
decision.proxy = Some(aether_contracts::ProxySnapshot {
enabled: Some(true),
url: Some("http://proxy.example.test:8080".to_string()),
..Default::default()
});
},
|decision: &mut AiExecutionDecision| {
decision.transport_profile = Some(aether_contracts::ResolvedTransportProfile {
profile_id: "chrome136".to_string(),
..Default::default()
});
},
] {
let mut changed = base.clone();
mutate(&mut changed);
let changed_identity =
UpstreamBindingIdentity::from_decision(adapter, &changed).unwrap();
assert_ne!(identity, changed_identity);
}
let mut rotated = base.clone();
rotated
.provider_request_headers
.insert("Authorization".to_string(), "Bearer rotated".to_string());
assert_eq!(
identity,
UpstreamBindingIdentity::from_decision(adapter, &rotated).unwrap()
);
}
#[test]
fn stable_key_identity_ignores_custom_auth_value_rotation() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let mut base = decision();
base.auth_header = Some("X-Provider-Token".to_string());
base.provider_request_headers.remove("Authorization");
base.provider_request_headers.insert(
"X-Provider-Token".to_string(),
"provider-token-1".to_string(),
);
let identity = UpstreamBindingIdentity::from_decision(adapter, &base).unwrap();
assert!(!identity.handshake_headers.contains_key("x-provider-token"));
assert!(identity.auth_fingerprint.is_none());
let mut rotated = base;
rotated.provider_request_headers.insert(
"X-Provider-Token".to_string(),
"provider-token-2".to_string(),
);
assert_eq!(
identity,
UpstreamBindingIdentity::from_decision(adapter, &rotated).unwrap()
);
}
#[test]
fn missing_key_identity_fingerprints_authentication_values() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let mut first = decision();
first.key_id = None;
let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap();
assert!(first_identity.auth_fingerprint.is_some());
let mut same_account_rotation = first.clone();
same_account_rotation.provider_request_headers.insert(
"Authorization".to_string(),
"Bearer different-account-or-token".to_string(),
);
let changed_identity =
UpstreamBindingIdentity::from_decision(adapter, &same_account_rotation).unwrap();
assert_ne!(first_identity, changed_identity);
let mut non_auth_change = first;
non_auth_change
.provider_request_headers
.insert("X-Client".to_string(), "other-client".to_string());
assert_ne!(
first_identity,
UpstreamBindingIdentity::from_decision(adapter, &non_auth_change).unwrap()
);
}
#[test]
fn disabled_proxy_is_equivalent_to_direct_transport() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let direct = decision();
let direct_identity = UpstreamBindingIdentity::from_decision(adapter, &direct).unwrap();
let mut explicitly_disabled = direct;
explicitly_disabled.proxy = Some(aether_contracts::ProxySnapshot {
enabled: Some(false),
url: Some("http://ignored.example.test:8080".to_string()),
..Default::default()
});
assert_eq!(
direct_identity,
UpstreamBindingIdentity::from_decision(adapter, &explicitly_disabled).unwrap()
);
}
#[test]
fn identity_rejects_missing_or_invalid_connection_fields() {
let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard);
let mut missing = decision();
missing.upstream_url = None;
assert_eq!(
UpstreamBindingIdentity::from_decision(adapter, &missing),
Err(UpstreamBindingIdentityError::MissingUpstreamUrl)
);
let mut invalid = decision();
invalid.upstream_url = Some("file:///tmp/responses".to_string());
assert_eq!(
UpstreamBindingIdentity::from_decision(adapter, &invalid),
Err(UpstreamBindingIdentityError::InvalidUpstreamUrl)
);
}
}