//! 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, credentials, and the protocol adapter can //! all change the connection that would receive the next event. Ordinary Codex //! OAuth access-token refreshes retain the credential generation and therefore //! do not unnecessarily replace an already-upgraded socket. 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, endpoint_id: Option, key_id: Option, upstream_url: String, handshake_headers: BTreeMap, /// One-way identity for the credential generation used by this socket. /// /// A provider key id identifies a catalog row, not the secret currently /// stored in that row. Codex decisions carry a server-owned credential /// generation which is stable across access-token refreshes but rotates /// when the account/static/refresh credential is replaced. Other /// decisions conservatively fingerprint the effective authentication /// handshake values. credential_fingerprint: [u8; 32], proxy: Option, transport_profile: Option, } #[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 { 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 adapter_kind = adapter.kind(); 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 if !is_turn_scoped_handshake_header(adapter_kind, name.as_str()) { handshake_headers.insert(name, value.to_string()); } } let credential_fingerprint = credential_binding_fingerprint(decision, &authentication_headers); Ok(Self { adapter_kind, provider_id: decision.provider_id.clone(), endpoint_id: decision.endpoint_id.clone(), key_id: decision.key_id.clone(), upstream_url, handshake_headers, credential_fingerprint, proxy: effective_proxy_snapshot(decision.proxy.as_ref()), transport_profile: decision.transport_profile.clone(), }) } /// Returns only safe field/header names for operator diagnostics. Header /// values and credential fingerprints must never reach logs. pub(super) fn changed_field_names(&self, other: &Self) -> Vec { let mut changed = Vec::new(); if self.adapter_kind != other.adapter_kind { changed.push("adapter_kind".to_string()); } if self.provider_id != other.provider_id { changed.push("provider_id".to_string()); } if self.endpoint_id != other.endpoint_id { changed.push("endpoint_id".to_string()); } if self.key_id != other.key_id { changed.push("key_id".to_string()); } if self.upstream_url != other.upstream_url { changed.push("upstream_url".to_string()); } for name in self .handshake_headers .keys() .chain(other.handshake_headers.keys()) .collect::>() { if self.handshake_headers.get(name) != other.handshake_headers.get(name) { changed.push(format!("handshake_header:{name}")); } } if self.credential_fingerprint != other.credential_fingerprint { changed.push("credential_fingerprint".to_string()); } if self.proxy != other.proxy { changed.push("proxy".to_string()); } if self.transport_profile != other.transport_profile { changed.push("transport_profile".to_string()); } changed } } /// `x-codex-turn-metadata` describes one logical response turn. Fingerprint /// convergence intentionally rewrites its `turn_id` and timestamp for every /// `response.create`, including a continuation on an already-upgraded socket. /// It therefore cannot identify the physical handshake that owns a /// `previous_response_id`; the per-turn copy in `client_metadata` still travels /// in the response.create body. fn is_turn_scoped_handshake_header(adapter_kind: ResponsesWebSocketAdapter, name: &str) -> bool { adapter_kind == ResponsesWebSocketAdapter::Codex && name.eq_ignore_ascii_case("x-codex-turn-metadata") } /// 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 { 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) -> [u8; 32] { let mut hasher = Sha256::new(); hasher.update(b"aether-responses-websocket-auth-headers-v1"); 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() } /// Returns the non-secret credential identity represented by a planner /// decision. The generation is emitted by Aether's trusted Codex planner from /// provider-key metadata; it is not sourced from the downstream request. fn credential_binding_fingerprint( decision: &AiExecutionDecision, authentication_headers: &BTreeMap, ) -> [u8; 32] { if decision .provider_type .as_deref() .is_some_and(|provider_type| provider_type.trim().eq_ignore_ascii_case("codex")) { if let Some(generation) = decision .report_context .as_ref() .and_then(|context| context.get("codex_credential_generation")) .and_then(serde_json::Value::as_str) .map(str::trim) .filter(|generation| !generation.is_empty()) { let mut hasher = Sha256::new(); hasher.update(b"aether-responses-websocket-codex-credential-generation-v1"); hasher.update((generation.len() as u64).to_be_bytes()); hasher.update(generation.as_bytes()); // Only a planner-owned Codex bearer access token is expected to // rotate without changing credential generation. Compare the // effective handshake value with the decision's original auth // value: auth-config/routing/header overrides change only the // former and therefore must force a rebind. let stable_authentication_headers = authentication_headers .iter() .filter(|(name, value)| { !is_planner_owned_codex_bearer(decision, name.as_str(), value.as_str()) }) .map(|(name, value)| (name.clone(), value.clone())) .collect::>(); hasher.update(fingerprint_headers(&stable_authentication_headers)); return hasher.finalize().into(); } } // Fail closed when no trusted generation is available. Rebinding after an // access-token change is preferable to sending a continuation over a // socket authenticated with a credential that may have been replaced. fingerprint_headers(authentication_headers) } fn is_planner_owned_codex_bearer( decision: &AiExecutionDecision, name: &str, effective_value: &str, ) -> bool { name.eq_ignore_ascii_case("authorization") && decision .auth_header .as_deref() .is_some_and(|header| header.eq_ignore_ascii_case(name)) && decision.auth_value.as_deref() == Some(effective_value) && effective_value .get(.."bearer ".len()) .is_some_and(|scheme| scheme.eq_ignore_ascii_case("bearer ")) } /// 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 { 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::>(), ) .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 static_secret_rotated = base.clone(); static_secret_rotated .provider_request_headers .insert("Authorization".to_string(), "Bearer rotated".to_string()); assert_ne!( identity, UpstreamBindingIdentity::from_decision(adapter, &static_secret_rotated).unwrap() ); } #[test] fn stable_key_identity_rejects_custom_static_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")); let mut rotated = base; rotated.provider_request_headers.insert( "X-Provider-Token".to_string(), "provider-token-2".to_string(), ); assert_ne!( identity, UpstreamBindingIdentity::from_decision(adapter, &rotated).unwrap() ); } #[test] fn codex_access_token_refresh_reuses_the_same_credential_generation() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); let mut access_token_refreshed = first; access_token_refreshed.auth_value = Some("Bearer refreshed-access-token".to_string()); access_token_refreshed.provider_request_headers.insert( "Authorization".to_string(), "Bearer refreshed-access-token".to_string(), ); assert_eq!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &access_token_refreshed).unwrap() ); } #[test] fn codex_turn_metadata_changes_do_not_change_the_physical_binding() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); first.provider_request_headers.insert( "x-codex-turn-metadata".to_string(), json!({ "session_id": "session-1", "thread_id": "thread-1", "turn_id": "turn-1", "turn_started_at_unix_ms": 1 }) .to_string(), ); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); assert!(!first_identity .handshake_headers .contains_key("x-codex-turn-metadata")); let mut continuation = first; continuation.provider_request_headers.insert( "x-codex-turn-metadata".to_string(), json!({ "session_id": "session-1", "thread_id": "thread-1", "turn_id": "turn-2", "turn_started_at_unix_ms": 2 }) .to_string(), ); assert_eq!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &continuation).unwrap() ); } #[test] fn only_codex_turn_metadata_is_excluded_from_binding_headers() { let codex_adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); first.provider_request_headers.insert( "x-codex-turn-metadata".to_string(), r#"{"turn_id":"turn-1"}"#.to_string(), ); let first_identity = UpstreamBindingIdentity::from_decision(codex_adapter, &first).unwrap(); let mut changed = first; changed .provider_request_headers .insert("x-codex-window-id".to_string(), "window-2".to_string()); let changed_identity = UpstreamBindingIdentity::from_decision(codex_adapter, &changed).unwrap(); assert_ne!(first_identity, changed_identity); assert_eq!( first_identity.changed_field_names(&changed_identity), vec!["handshake_header:x-codex-window-id".to_string()] ); let standard_adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard); let standard_identity = UpstreamBindingIdentity::from_decision(standard_adapter, &changed).unwrap(); assert!(standard_identity .handshake_headers .contains_key("x-codex-turn-metadata")); } #[test] fn codex_authorization_override_changes_binding_with_the_same_generation() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); // The planner-owned auth value remains unchanged while an effective // auth-config/header override replaces the actual handshake value. first.provider_request_headers.insert( "Authorization".to_string(), "Bearer endpoint-override".to_string(), ); assert_ne!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &first).unwrap() ); } #[test] fn codex_credential_replacement_changes_binding_for_the_same_key_id() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); let mut replaced = first; replaced.provider_request_headers.insert( "Authorization".to_string(), "Bearer replacement-access-token".to_string(), ); replaced.report_context = Some(json!({ "codex_credential_generation": "credential-generation-2" })); assert_ne!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &replaced).unwrap() ); } #[test] fn codex_custom_auth_rotation_changes_binding_with_the_same_generation() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); first.auth_header = Some("X-Provider-Token".to_string()); first.provider_request_headers.insert( "X-Provider-Token".to_string(), "provider-token-1".to_string(), ); first.report_context = Some(json!({ "codex_credential_generation": "credential-generation-1" })); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); first.provider_request_headers.insert( "X-Provider-Token".to_string(), "provider-token-2".to_string(), ); assert_ne!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &first).unwrap() ); } #[test] fn missing_codex_credential_generation_fails_closed_on_auth_rotation() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); first.provider_request_headers.insert( "Authorization".to_string(), "Bearer possibly-replaced-credential".to_string(), ); assert_ne!( first_identity, UpstreamBindingIdentity::from_decision(adapter, &first).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) ); } }