fix(core): attest websocket response handshakes

This commit is contained in:
Jiaming Zhang
2026-05-04 21:30:48 -07:00
parent 74cfe6f045
commit 698e5983b9
2 changed files with 27 additions and 2 deletions

View File

@@ -830,7 +830,9 @@ impl ModelClient {
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
let headers = self.build_websocket_headers(turn_state.as_ref(), turn_metadata_header);
let headers = self
.build_websocket_headers(&api_provider, turn_state.as_ref(), turn_metadata_header)
.await;
let websocket_telemetry = ModelClientSession::build_websocket_telemetry(
session_telemetry,
auth_context,
@@ -907,8 +909,9 @@ impl ModelClient {
///
/// Callers should pass the current turn-state lock when available so sticky-routing state is
/// replayed on reconnect within the same turn.
fn build_websocket_headers(
async fn build_websocket_headers(
&self,
provider: &codex_api::Provider,
turn_state: Option<&Arc<OnceLock<String>>>,
turn_metadata_header: Option<&str>,
) -> ApiHeaderMap {
@@ -924,6 +927,8 @@ impl ModelClient {
}
headers.extend(build_conversation_headers(Some(conversation_id)));
headers.extend(self.build_responses_identity_headers());
self.extend_attestation_header_for(&mut headers, provider, AttestationPurpose::Response)
.await;
headers.insert(
OPENAI_BETA_HEADER,
HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE),

View File

@@ -562,6 +562,26 @@ async fn responses_generate_fresh_attestation_headers_for_chatgpt_codex() {
assert_eq!(attestation_calls.load(Ordering::Relaxed), 2);
}
#[tokio::test]
async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() {
let provider = api_provider("https://chatgpt.com/backend-api/codex/");
let (model_client, attestation_calls) = model_client_with_counting_attestation();
let headers = model_client
.build_websocket_headers(
&provider, /*turn_state*/ None, /*turn_metadata_header*/ None,
)
.await;
assert_eq!(
headers
.get(crate::attestation::X_OAI_ATTESTATION_HEADER)
.and_then(|value| value.to_str().ok()),
Some("v1.header-1"),
);
assert_eq!(attestation_calls.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn compact_generate_fresh_attestation_headers_for_chatgpt_codex() {
let provider = api_provider("https://chatgpt.com/backend-api/codex/");