[codex] wire native HTTP state transport [ci changed_files]

This commit is contained in:
Cooper Gamble
2026-06-03 21:25:00 +00:00
parent 44afe34b77
commit b81167c920
30 changed files with 1401 additions and 81 deletions
+3
View File
@@ -2547,6 +2547,7 @@ dependencies = [
"codex-feedback",
"codex-git-utils",
"codex-hooks",
"codex-http-state",
"codex-install-context",
"codex-login",
"codex-mcp",
@@ -3321,6 +3322,7 @@ dependencies = [
"codex-aws-auth",
"codex-client",
"codex-feedback",
"codex-http-state",
"codex-login",
"codex-model-provider-info",
"codex-models-manager",
@@ -3330,6 +3332,7 @@ dependencies = [
"http 1.4.0",
"pretty_assertions",
"serde_json",
"tempfile",
"tokio",
"tracing",
"wiremock",
+22 -1
View File
@@ -41,6 +41,27 @@ pub trait AuthProvider: Send + Sync {
headers
}
/// Adds auth headers for a concrete HTTP request URL.
///
/// URL-sensitive providers can override this to scope request-specific
/// headers without exposing them through generic/non-HTTP auth helpers.
fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
self.add_auth_headers(headers);
}
/// Observes response headers for auth state that may need to rotate after a request.
///
/// Most providers do not need this. Providers with server-minted,
/// request-scoped state may validate the URL and selectively persist
/// response headers here.
fn observe_response_headers(
&self,
_request_url: &str,
_request_headers: &HeaderMap,
_response_headers: &HeaderMap,
) {
}
/// Applies auth to a complete outbound request and returns the request to send.
///
/// The input `request` is moved into this method. Implementations may mutate
@@ -54,7 +75,7 @@ pub trait AuthProvider: Send + Sync {
/// If this returns [`AuthError`], the request should not be sent.
async fn apply_auth(&self, request: Request) -> Result<Request, AuthError> {
let mut request = request;
self.add_auth_headers(&mut request.headers);
self.add_auth_headers_for_url(&request.url, &mut request.headers);
Ok(request)
}
}
+1
View File
@@ -6,6 +6,7 @@ pub(crate) mod realtime_call;
pub(crate) mod realtime_websocket;
pub(crate) mod responses;
pub(crate) mod responses_websocket;
mod responses_websocket_auth;
pub(crate) mod search;
mod session;
@@ -1,3 +1,4 @@
use crate::auth::SharedAuthProvider;
use crate::endpoint::realtime_websocket::methods_common::conversation_function_call_output_message;
use crate::endpoint::realtime_websocket::methods_common::conversation_item_create_message;
use crate::endpoint::realtime_websocket::methods_common::normalized_session_mode;
@@ -550,11 +551,21 @@ fn contains_transcript_entry(entries: &[RealtimeTranscriptEntry], role: &str, te
pub struct RealtimeWebsocketClient {
provider: Provider,
auth: Option<SharedAuthProvider>,
}
impl RealtimeWebsocketClient {
pub fn new(provider: Provider) -> Self {
Self { provider }
Self {
provider,
auth: None,
}
}
/// Applies URL-aware auth during handshakes and observes handshake responses.
pub fn with_auth(mut self, auth: SharedAuthProvider) -> Self {
self.auth = Some(auth);
self
}
pub async fn connect(
@@ -653,6 +664,10 @@ impl RealtimeWebsocketClient {
default_headers,
);
request.headers_mut().extend(headers);
if let Some(auth) = self.auth.as_ref() {
auth.add_auth_headers_for_url(ws_url.as_str(), request.headers_mut());
}
let request_headers = request.headers().clone();
info!("connecting realtime websocket: {ws_url}");
// Realtime websocket TLS should honor the same custom-CA env vars as the rest of Codex's
@@ -660,14 +675,39 @@ impl RealtimeWebsocketClient {
let connector = maybe_build_rustls_client_config_with_custom_ca()
.map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?
.map(tokio_tungstenite::Connector::Rustls);
let (stream, response) = tokio_tungstenite::connect_async_tls_with_config(
let connect_result = tokio_tungstenite::connect_async_tls_with_config(
request,
Some(websocket_config()),
false,
connector,
)
.await
.map_err(|err| ApiError::Stream(format!("failed to connect realtime websocket: {err}")))?;
.await;
let (stream, response) = match connect_result {
Ok((stream, response)) => {
if let Some(auth) = self.auth.as_ref() {
auth.observe_response_headers(
ws_url.as_str(),
&request_headers,
response.headers(),
);
}
(stream, response)
}
Err(err) => {
if let Some(auth) = self.auth.as_ref()
&& let WsError::Http(response) = &err
{
auth.observe_response_headers(
ws_url.as_str(),
&request_headers,
response.headers(),
);
}
return Err(ApiError::Stream(format!(
"failed to connect realtime websocket: {err}"
)));
}
};
info!(
ws_url = %ws_url,
status = %response.status(),
@@ -1,3 +1,4 @@
use super::responses_websocket_auth::WebsocketAuthContext;
use crate::auth::SharedAuthProvider;
use crate::common::ResponseEvent;
use crate::common::ResponseProcessedWsRequest;
@@ -167,6 +168,7 @@ pub struct ResponsesWebsocketConnection {
server_reasoning_included: bool,
models_etag: Option<String>,
server_model: Option<String>,
auth_context: WebsocketAuthContext,
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
}
@@ -178,6 +180,7 @@ impl std::fmt::Debug for ResponsesWebsocketConnection {
.field("server_reasoning_included", &self.server_reasoning_included)
.field("models_etag", &self.models_etag)
.field("server_model", &self.server_model)
.field("auth_context", &"<auth-context>")
.field("telemetry", &self.telemetry.as_ref().map(|_| "<telemetry>"))
.finish()
}
@@ -190,6 +193,7 @@ impl ResponsesWebsocketConnection {
server_reasoning_included: bool,
models_etag: Option<String>,
server_model: Option<String>,
auth_context: WebsocketAuthContext,
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
) -> Self {
Self {
@@ -198,6 +202,7 @@ impl ResponsesWebsocketConnection {
server_reasoning_included,
models_etag,
server_model,
auth_context,
telemetry,
}
}
@@ -258,6 +263,7 @@ impl ResponsesWebsocketConnection {
let server_reasoning_included = self.server_reasoning_included;
let models_etag = self.models_etag.clone();
let server_model = self.server_model.clone();
let auth_context = self.auth_context.clone();
let telemetry = self.telemetry.clone();
let request_body = serde_json::to_value(&request).map_err(|err| {
ApiError::Stream(format!("failed to encode websocket request: {err}"))
@@ -297,6 +303,7 @@ impl ResponsesWebsocketConnection {
tx_event.clone(),
request_body,
idle_timeout,
&auth_context,
telemetry,
connection_reused,
)
@@ -377,19 +384,35 @@ impl ResponsesWebsocketClient {
.provider
.websocket_url_for_path("responses")
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
let request_url = ws_url.to_string();
let mut headers =
merge_request_headers(&self.provider.headers, extra_headers, default_headers);
self.auth.add_auth_headers(&mut headers);
self.auth
.add_auth_headers_for_url(&request_url, &mut headers);
let auth_context =
WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());
let (stream, _status, server_reasoning_included, models_etag, server_model) =
connect_websocket(ws_url, headers, turn_state.clone()).await?;
let (
stream,
_status,
server_reasoning_included,
models_etag,
server_model,
response_headers,
) = connect_websocket(ws_url, headers, turn_state.clone())
.await
.inspect_err(|error| {
auth_context.observe_error_headers(error);
})?;
auth_context.observe_response_headers(&response_headers);
Ok(ResponsesWebsocketConnection::new(
stream,
self.provider.stream_idle_timeout,
server_reasoning_included,
models_etag,
server_model,
auth_context,
telemetry,
))
}
@@ -411,13 +434,22 @@ impl ResponsesWebsocketClient {
.provider
.websocket_url_for_path("responses")
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
let request_url = ws_url.to_string();
let mut headers =
merge_request_headers(&self.provider.headers, extra_headers, default_headers);
self.auth.add_auth_headers(&mut headers);
self.auth
.add_auth_headers_for_url(&request_url, &mut headers);
let auth_context =
WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());
let (mut stream, status, reasoning_included, models_etag, server_model) =
connect_websocket(ws_url.clone(), headers, /*turn_state*/ None).await?;
let (mut stream, status, reasoning_included, models_etag, server_model, response_headers) =
connect_websocket(ws_url.clone(), headers, /*turn_state*/ None)
.await
.inspect_err(|error| {
auth_context.observe_error_headers(error);
})?;
auth_context.observe_response_headers(&response_headers);
let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
.await
.ok()
@@ -472,7 +504,17 @@ async fn connect_websocket(
url: Url,
headers: HeaderMap,
turn_state: Option<Arc<OnceLock<String>>>,
) -> Result<(WsStream, StatusCode, bool, Option<String>, Option<String>), ApiError> {
) -> Result<
(
WsStream,
StatusCode,
bool,
Option<String>,
Option<String>,
HeaderMap,
),
ApiError,
> {
ensure_rustls_crypto_provider();
info!("connecting to websocket: {url}");
@@ -499,10 +541,7 @@ async fn connect_websocket(
let (stream, response) = match response {
Ok((stream, response)) => {
info!(
"successfully connected to websocket: {url}, headers: {:?}",
response.headers()
);
info!("successfully connected to websocket: {url}");
(stream, response)
}
Err(err) => {
@@ -536,6 +575,7 @@ async fn connect_websocket(
reasoning_included,
models_etag,
server_model,
response.headers().clone(),
))
}
@@ -630,12 +670,12 @@ fn map_wrapped_websocket_error_event(
Some(ApiError::Transport(TransportError::Http {
status,
url: None,
headers: headers.map(json_headers_to_http_headers),
headers: headers.as_ref().map(json_headers_to_http_headers),
body: Some(original_payload),
}))
}
fn json_headers_to_http_headers(headers: JsonMap<String, Value>) -> HeaderMap {
fn json_headers_to_http_headers(headers: &JsonMap<String, Value>) -> HeaderMap {
let mut mapped = HeaderMap::new();
for (name, value) in headers {
let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else {
@@ -649,9 +689,9 @@ fn json_headers_to_http_headers(headers: JsonMap<String, Value>) -> HeaderMap {
mapped
}
fn json_header_value(value: Value) -> Option<HeaderValue> {
fn json_header_value(value: &Value) -> Option<HeaderValue> {
let value = match value {
Value::String(value) => value,
Value::String(value) => value.clone(),
Value::Number(value) => value.to_string(),
Value::Bool(value) => value.to_string(),
_ => return None,
@@ -664,6 +704,7 @@ async fn run_websocket_response_stream(
tx_event: mpsc::Sender<std::result::Result<ResponseEvent, ApiError>>,
request_body: Value,
idle_timeout: Duration,
auth_context: &WebsocketAuthContext,
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
connection_reused: bool,
) -> Result<(), ApiError> {
@@ -702,12 +743,25 @@ async fn run_websocket_response_stream(
match message {
Message::Text(text) => {
trace!("websocket event: {text}");
if let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text)
&& let Some(error) =
let wrapped_error = parse_wrapped_websocket_error_event(&text);
if wrapped_error.is_some() {
trace!("websocket error event");
} else {
trace!("websocket event: {text}");
}
if let Some(wrapped_error) = wrapped_error {
if let Some(response_headers) = wrapped_error.headers.as_ref() {
auth_context.observe_response_headers(&json_headers_to_http_headers(
response_headers,
));
}
if let Some(error) =
map_wrapped_websocket_error_event(wrapped_error, text.to_string())
{
return Err(error);
{
return Err(error);
}
debug!("ignoring unmapped websocket error event");
continue;
}
let event = match serde_json::from_str::<ResponsesStreamEvent>(&text) {
@@ -0,0 +1,52 @@
use crate::auth::SharedAuthProvider;
use crate::error::ApiError;
use codex_client::TransportError;
use http::HeaderMap;
use std::sync::Arc;
use std::sync::Mutex;
#[derive(Clone)]
pub(super) struct WebsocketAuthContext {
auth: SharedAuthProvider,
request_url: String,
request_headers: Arc<Mutex<HeaderMap>>,
}
impl WebsocketAuthContext {
pub(super) fn new(
auth: SharedAuthProvider,
request_url: String,
request_headers: HeaderMap,
) -> Self {
Self {
auth,
request_url,
request_headers: Arc::new(Mutex::new(request_headers)),
}
}
pub(super) fn observe_response_headers(&self, response_headers: &HeaderMap) {
let mut request_headers = self
.request_headers
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
self.auth
.observe_response_headers(&self.request_url, &request_headers, response_headers);
self.auth
.add_auth_headers_for_url(&self.request_url, &mut request_headers);
}
pub(super) fn observe_error_headers(&self, error: &ApiError) {
if let ApiError::Transport(TransportError::Http {
headers: Some(response_headers),
..
}) = error
{
self.observe_response_headers(response_headers);
}
}
}
#[cfg(test)]
#[path = "responses_websocket_auth_tests.rs"]
mod tests;
@@ -0,0 +1,124 @@
use super::*;
use crate::auth::AuthProvider;
use http::HeaderValue;
use http::StatusCode;
use pretty_assertions::assert_eq;
use std::sync::Mutex as StdMutex;
const REQUEST_URL: &str = "wss://chatgpt.com/backend-api/codex/responses";
const SENT_STATE: &str = "sent-state";
const HANDSHAKE_STATE: &str = "handshake-state";
const WRAPPED_ERROR_STATE: &str = "wrapped-error-state";
struct RecordingAuthProvider {
state: StdMutex<String>,
observed_updates: StdMutex<Vec<(String, String, String)>>,
}
impl Default for RecordingAuthProvider {
fn default() -> Self {
Self {
state: StdMutex::new(SENT_STATE.to_string()),
observed_updates: StdMutex::new(Vec::new()),
}
}
}
impl AuthProvider for RecordingAuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
let state = self
.state
.lock()
.expect("state lock should not be poisoned");
headers.insert(
"x-test-state",
HeaderValue::from_str(&state).expect("state should be a valid header value"),
);
}
fn observe_response_headers(
&self,
request_url: &str,
request_headers: &HeaderMap,
response_headers: &HeaderMap,
) {
let Some(sent_state) = request_headers
.get("x-test-state")
.and_then(|value| value.to_str().ok())
else {
return;
};
let Some(update_state) = response_headers
.get("x-test-state-update")
.and_then(|value| value.to_str().ok())
else {
return;
};
self.observed_updates
.lock()
.expect("observed updates lock should not be poisoned")
.push((
request_url.to_string(),
sent_state.to_string(),
update_state.to_string(),
));
*self
.state
.lock()
.expect("state lock should not be poisoned") = update_state.to_string();
}
}
#[test]
fn observes_handshake_and_wrapped_error_headers() {
let auth = Arc::new(RecordingAuthProvider::default());
let mut request_headers = HeaderMap::new();
auth.add_auth_headers_for_url(REQUEST_URL, &mut request_headers);
assert_eq!(
request_headers
.get("x-test-state")
.and_then(|value| value.to_str().ok()),
Some(SENT_STATE)
);
let context = WebsocketAuthContext::new(auth.clone(), REQUEST_URL.to_string(), request_headers);
let mut handshake_headers = HeaderMap::new();
handshake_headers.insert(
"x-test-state-update",
HeaderValue::from_static(HANDSHAKE_STATE),
);
context.observe_response_headers(&handshake_headers);
let mut wrapped_error_headers = HeaderMap::new();
wrapped_error_headers.insert(
"x-test-state-update",
HeaderValue::from_static(WRAPPED_ERROR_STATE),
);
context.observe_error_headers(&ApiError::Transport(TransportError::Http {
status: StatusCode::BAD_REQUEST,
url: Some(REQUEST_URL.to_string()),
headers: Some(wrapped_error_headers),
body: None,
}));
assert_eq!(
*auth
.observed_updates
.lock()
.expect("observed updates lock should not be poisoned"),
vec![
(
REQUEST_URL.to_string(),
SENT_STATE.to_string(),
HANDSHAKE_STATE.to_string(),
),
(
REQUEST_URL.to_string(),
HANDSHAKE_STATE.to_string(),
WRAPPED_ERROR_STATE.to_string(),
),
]
);
}
+61 -2
View File
@@ -102,7 +102,16 @@ impl<T: HttpTransport> EndpointSession<T> {
let transport = &self.transport;
async move {
let req = auth.apply_auth(req).await.map_err(TransportError::from)?;
transport.execute(req).await
let request_url = req.url.clone();
let request_headers = req.headers.clone();
let response = transport.execute(req).await;
observe_auth_response_headers(
auth.as_ref(),
&request_url,
&request_headers,
&response,
);
response
}
},
)
@@ -143,7 +152,16 @@ impl<T: HttpTransport> EndpointSession<T> {
let transport = &self.transport;
async move {
let req = auth.apply_auth(req).await.map_err(TransportError::from)?;
transport.stream(req).await
let request_url = req.url.clone();
let request_headers = req.headers.clone();
let response = transport.stream(req).await;
observe_auth_response_headers(
auth.as_ref(),
&request_url,
&request_headers,
&response,
);
response
}
},
)
@@ -152,3 +170,44 @@ impl<T: HttpTransport> EndpointSession<T> {
Ok(stream)
}
}
fn observe_auth_response_headers<T>(
auth: &dyn crate::auth::AuthProvider,
request_url: &str,
request_headers: &HeaderMap,
response: &Result<T, TransportError>,
) where
T: ResponseHeaders,
{
match response {
Ok(response) => {
auth.observe_response_headers(request_url, request_headers, response.headers())
}
Err(TransportError::Http {
headers: Some(headers),
..
}) => auth.observe_response_headers(request_url, request_headers, headers),
Err(_) => {}
}
}
/// Exposes transport response headers so auth providers can observe request-scoped updates.
trait ResponseHeaders {
fn headers(&self) -> &HeaderMap;
}
impl ResponseHeaders for Response {
fn headers(&self) -> &HeaderMap {
&self.headers
}
}
impl ResponseHeaders for StreamResponse {
fn headers(&self) -> &HeaderMap {
&self.headers
}
}
#[cfg(test)]
#[path = "session_tests.rs"]
mod tests;
@@ -0,0 +1,174 @@
use super::*;
use crate::provider::RetryConfig;
use async_trait::async_trait;
use codex_client::Request;
use http::HeaderValue;
use http::StatusCode;
use pretty_assertions::assert_eq;
use std::sync::Mutex as StdMutex;
use std::time::Duration;
const RESPONSES_URL: &str = "https://chatgpt.com/backend-api/codex/responses";
#[derive(Default)]
struct RecordingAuthProvider {
observed_headers: StdMutex<Vec<(String, HeaderMap, HeaderMap)>>,
}
impl crate::auth::AuthProvider for RecordingAuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
headers.insert("x-test-state", HeaderValue::from_static("sent-state"));
}
fn observe_response_headers(
&self,
request_url: &str,
request_headers: &HeaderMap,
response_headers: &HeaderMap,
) {
self.observed_headers
.lock()
.expect("recording auth lock should not be poisoned")
.push((
request_url.to_string(),
request_headers.clone(),
response_headers.clone(),
));
}
}
#[derive(Clone)]
struct RejectingTransport {
requests: Arc<StdMutex<Vec<(String, HeaderMap)>>>,
response_headers: HeaderMap,
}
impl RejectingTransport {
fn new(response_headers: HeaderMap) -> Self {
Self {
requests: Arc::new(StdMutex::new(Vec::new())),
response_headers,
}
}
fn http_error(&self, request: Request) -> TransportError {
self.requests
.lock()
.expect("recording transport lock should not be poisoned")
.push((request.url.clone(), request.headers));
TransportError::Http {
status: StatusCode::UNAUTHORIZED,
url: Some(request.url),
headers: Some(self.response_headers.clone()),
body: None,
}
}
}
#[async_trait]
impl HttpTransport for RejectingTransport {
async fn execute(&self, request: Request) -> Result<Response, TransportError> {
Err(self.http_error(request))
}
async fn stream(&self, request: Request) -> Result<StreamResponse, TransportError> {
Err(self.http_error(request))
}
}
#[tokio::test]
async fn execute_attaches_url_scoped_auth_and_observes_http_error_headers() {
let (session, transport, auth, request_headers, response_headers) = test_session();
session
.execute(
Method::POST,
"responses",
HeaderMap::new(),
/*body*/ None,
)
.await
.expect_err("request should fail");
assert_request_and_response_headers(&transport, &auth, request_headers, response_headers);
}
#[tokio::test]
async fn stream_attaches_url_scoped_auth_and_observes_http_error_headers() {
let (session, transport, auth, request_headers, response_headers) = test_session();
let result = session
.stream_with(
Method::POST,
"responses",
HeaderMap::new(),
/*body*/ None,
|_| {},
)
.await;
assert!(result.is_err(), "request should fail");
assert_request_and_response_headers(&transport, &auth, request_headers, response_headers);
}
fn test_session() -> (
EndpointSession<RejectingTransport>,
RejectingTransport,
Arc<RecordingAuthProvider>,
HeaderMap,
HeaderMap,
) {
let mut request_headers = HeaderMap::new();
request_headers.insert("x-test-state", HeaderValue::from_static("sent-state"));
let mut response_headers = HeaderMap::new();
response_headers.insert(
"x-test-state-update",
HeaderValue::from_static("rotated-state"),
);
let transport = RejectingTransport::new(response_headers.clone());
let auth = Arc::new(RecordingAuthProvider::default());
let session = EndpointSession::new(transport.clone(), provider(), auth.clone());
(session, transport, auth, request_headers, response_headers)
}
fn provider() -> Provider {
Provider {
name: "test".to_string(),
base_url: "https://chatgpt.com/backend-api/codex".to_string(),
query_params: None,
headers: HeaderMap::new(),
retry: RetryConfig {
max_attempts: 1,
base_delay: Duration::from_millis(1),
retry_429: false,
retry_5xx: false,
retry_transport: false,
},
stream_idle_timeout: Duration::from_secs(1),
}
}
fn assert_request_and_response_headers(
transport: &RejectingTransport,
auth: &RecordingAuthProvider,
request_headers: HeaderMap,
response_headers: HeaderMap,
) {
assert_eq!(
transport
.requests
.lock()
.expect("recording transport lock should not be poisoned")
.as_slice(),
&[(RESPONSES_URL.to_string(), request_headers.clone())]
);
assert_eq!(
auth.observed_headers
.lock()
.expect("recording auth lock should not be poisoned")
.as_slice(),
&[(RESPONSES_URL.to_string(), request_headers, response_headers,)]
);
}
+99 -22
View File
@@ -129,7 +129,9 @@ pub async fn upload_local_file(
.unwrap_or("file")
.to_string();
let create_url = format!("{}/files", base_url.trim_end_matches('/'));
let create_response = authorized_request(auth, reqwest::Method::POST, &create_url)
let (create_request, create_request_headers) =
authorized_request(auth, reqwest::Method::POST, &create_url);
let create_response = create_request
.json(&serde_json::json!({
"file_name": file_name,
"file_size": metadata.len(),
@@ -141,6 +143,11 @@ pub async fn upload_local_file(
url: create_url.clone(),
source,
})?;
auth.observe_response_headers(
&create_url,
&create_request_headers,
create_response.headers(),
);
let create_status = create_response.status();
let create_body = create_response.text().await.unwrap_or_default();
if !create_status.is_success() {
@@ -191,7 +198,9 @@ pub async fn upload_local_file(
);
let finalize_started_at = Instant::now();
loop {
let finalize_response = authorized_request(auth, reqwest::Method::POST, &finalize_url)
let (finalize_request, finalize_request_headers) =
authorized_request(auth, reqwest::Method::POST, &finalize_url);
let finalize_response = finalize_request
.json(&serde_json::json!({}))
.send()
.await
@@ -199,6 +208,11 @@ pub async fn upload_local_file(
url: finalize_url.clone(),
source,
})?;
auth.observe_response_headers(
&finalize_url,
&finalize_request_headers,
finalize_response.headers(),
);
let finalize_status = finalize_response.status();
let finalize_body = finalize_response.text().await.unwrap_or_default();
if !finalize_status.is_success() {
@@ -255,15 +269,18 @@ fn authorized_request(
auth: &dyn AuthProvider,
method: reqwest::Method,
url: &str,
) -> reqwest::RequestBuilder {
) -> (reqwest::RequestBuilder, http::HeaderMap) {
let mut headers = http::HeaderMap::new();
auth.add_auth_headers(&mut headers);
auth.add_auth_headers_for_url(url, &mut headers);
let client = build_reqwest_client();
client
.request(method, url)
.timeout(OPENAI_FILE_REQUEST_TIMEOUT)
.headers(headers)
(
client
.request(method, url)
.timeout(OPENAI_FILE_REQUEST_TIMEOUT)
.headers(headers.clone()),
headers,
)
}
fn build_reqwest_client() -> reqwest::Client {
@@ -279,6 +296,7 @@ mod tests {
use pretty_assertions::assert_eq;
use reqwest::header::HeaderValue;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use tempfile::TempDir;
@@ -291,8 +309,16 @@ mod tests {
use wiremock::matchers::method;
use wiremock::matchers::path;
#[derive(Clone, Copy)]
struct ChatGptTestAuth;
const SENT_STATE: &str = "sent-state";
const CREATED_STATE: &str = "created-state";
const RETRY_STATE: &str = "retry-state";
const FINALIZED_STATE: &str = "finalized-state";
type ObservedUpdate = (String, Option<String>);
#[derive(Clone, Default)]
struct ChatGptTestAuth {
observed_updates: Arc<Mutex<Vec<ObservedUpdate>>>,
}
impl AuthProvider for ChatGptTestAuth {
fn add_auth_headers(&self, headers: &mut reqwest::header::HeaderMap) {
@@ -302,10 +328,33 @@ mod tests {
);
headers.insert("ChatGPT-Account-ID", HeaderValue::from_static("account_id"));
}
fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut http::HeaderMap) {
self.add_auth_headers(headers);
headers.insert("x-test-state", HeaderValue::from_static(SENT_STATE));
}
fn observe_response_headers(
&self,
request_url: &str,
_request_headers: &http::HeaderMap,
response_headers: &http::HeaderMap,
) {
self.observed_updates
.lock()
.expect("observed updates lock should not be poisoned")
.push((
request_url.to_string(),
response_headers
.get("x-test-state-update")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string),
));
}
}
fn chatgpt_auth() -> ChatGptTestAuth {
ChatGptTestAuth
ChatGptTestAuth::default()
}
fn base_url_for(server: &MockServer) -> String {
@@ -318,6 +367,7 @@ mod tests {
Mock::given(method("POST"))
.and(path("/backend-api/files"))
.and(header("chatgpt-account-id", "account_id"))
.and(header("x-test-state", SENT_STATE))
.and(body_json(serde_json::json!({
"file_name": "hello.txt",
"file_size": 5,
@@ -325,6 +375,7 @@ mod tests {
})))
.respond_with(
ResponseTemplate::new(200)
.insert_header("x-test-state-update", CREATED_STATE)
.set_body_json(serde_json::json!({"file_id": "file_123", "upload_url": format!("{}/upload/file_123", server.uri())})),
)
.mount(&server)
@@ -340,20 +391,25 @@ mod tests {
let download_url = format!("{}/download/file_123", server.uri());
Mock::given(method("POST"))
.and(path("/backend-api/files/file_123/uploaded"))
.and(header("x-test-state", SENT_STATE))
.respond_with(move |_request: &Request| {
if finalize_attempts_responder.fetch_add(1, Ordering::SeqCst) == 0 {
return ResponseTemplate::new(200).set_body_json(serde_json::json!({
"status": "retry"
}));
return ResponseTemplate::new(200)
.insert_header("x-test-state-update", RETRY_STATE)
.set_body_json(serde_json::json!({
"status": "retry"
}));
}
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"status": "success",
"download_url": download_url,
"file_name": "hello.txt",
"mime_type": "text/plain",
"file_size_bytes": 5
}))
ResponseTemplate::new(200)
.insert_header("x-test-state-update", FINALIZED_STATE)
.set_body_json(serde_json::json!({
"status": "success",
"download_url": download_url,
"file_name": "hello.txt",
"mime_type": "text/plain",
"file_size_bytes": 5
}))
})
.mount(&server)
.await;
@@ -363,7 +419,8 @@ mod tests {
let path = dir.path().join("hello.txt");
tokio::fs::write(&path, b"hello").await.expect("write file");
let uploaded = upload_local_file(&base_url, &chatgpt_auth(), &path)
let auth = chatgpt_auth();
let uploaded = upload_local_file(&base_url, &auth, &path)
.await
.expect("upload succeeds");
@@ -376,5 +433,25 @@ mod tests {
assert_eq!(uploaded.file_name, "hello.txt");
assert_eq!(uploaded.mime_type, Some("text/plain".to_string()));
assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2);
assert_eq!(
*auth
.observed_updates
.lock()
.expect("observed updates lock should not be poisoned"),
vec![
(
format!("{}/files", base_url.trim_end_matches('/')),
Some(CREATED_STATE.to_string()),
),
(
format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')),
Some(RETRY_STATE.to_string()),
),
(
format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')),
Some(FINALIZED_STATE.to_string()),
),
]
);
}
}
@@ -1,7 +1,11 @@
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_api::AuthProvider;
use codex_api::Provider;
use codex_api::RealtimeAudioFrame;
use codex_api::RealtimeEvent;
@@ -19,11 +23,15 @@ use codex_protocol::protocol::RealtimeVoice;
use futures::SinkExt;
use futures::StreamExt;
use http::HeaderMap;
use http::HeaderValue;
use serde_json::Value;
use serde_json::json;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::accept_hdr_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::handshake::server::Request;
use tokio_tungstenite::tungstenite::handshake::server::Response;
type RealtimeWsStream = tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>;
@@ -75,6 +83,37 @@ fn test_provider(base_url: String) -> Provider {
}
}
struct HandshakeAuthProvider {
observed_response: Arc<AtomicBool>,
}
impl AuthProvider for HandshakeAuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) {
assert!(request_url.starts_with("ws://127.0.0.1:"));
headers.insert("x-test-auth-state", HeaderValue::from_static("sent-state"));
}
fn observe_response_headers(
&self,
request_url: &str,
request_headers: &HeaderMap,
response_headers: &HeaderMap,
) {
assert!(request_url.starts_with("ws://127.0.0.1:"));
assert_eq!(
request_headers.get("x-test-auth-state"),
Some(&HeaderValue::from_static("sent-state")),
);
if response_headers.get("x-test-auth-state-update")
== Some(&HeaderValue::from_static("rotated-state"))
{
self.observed_response.store(true, Ordering::Relaxed);
}
}
}
#[tokio::test]
async fn realtime_ws_e2e_session_create_and_event_flow() {
let (addr, server) = spawn_realtime_ws_server(|mut ws: RealtimeWsStream| async move {
@@ -202,6 +241,66 @@ async fn realtime_ws_e2e_session_create_and_event_flow() {
server.await.expect("server task");
}
#[tokio::test]
async fn realtime_ws_e2e_applies_and_observes_url_aware_auth() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr");
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let mut ws = accept_hdr_async(stream, |request: &Request, mut response: Response| {
assert_eq!(
request.headers().get("x-test-auth-state"),
Some(&HeaderValue::from_static("sent-state")),
);
response.headers_mut().insert(
"x-test-auth-state-update",
HeaderValue::from_static("rotated-state"),
);
Ok(response)
})
.await
.expect("accept ws");
let first = ws
.next()
.await
.expect("first msg")
.expect("first msg ok")
.into_text()
.expect("text");
let first_json: Value = serde_json::from_str(&first).expect("json");
assert_eq!(first_json["type"], "session.update");
});
let observed_response = Arc::new(AtomicBool::new(false));
let client = RealtimeWebsocketClient::new(test_provider(format!("http://{addr}"))).with_auth(
Arc::new(HandshakeAuthProvider {
observed_response: Arc::clone(&observed_response),
}),
);
let connection = client
.connect(
RealtimeSessionConfig {
instructions: "backend prompt".to_string(),
model: Some("realtime-test-model".to_string()),
session_id: Some("conv_123".to_string()),
event_parser: RealtimeEventParser::V1,
session_mode: RealtimeSessionMode::Conversational,
output_modality: RealtimeOutputModality::Audio,
voice: RealtimeVoice::Cove,
},
HeaderMap::new(),
HeaderMap::new(),
)
.await
.expect("connect");
assert!(observed_response.load(Ordering::Relaxed));
connection.close().await.expect("close");
server.await.expect("server task");
}
#[tokio::test]
async fn realtime_ws_connect_webrtc_sideband_retries_join_until_server_is_available() {
let reserving_listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
@@ -10,6 +10,18 @@ pub fn is_allowed_chatgpt_host(host: &str) -> bool {
.any(|suffix| host.ends_with(suffix))
}
/// Returns whether `url` is an HTTPS or secure WebSocket ChatGPT URL that Codex
/// may treat as first-party traffic.
pub fn is_allowed_chatgpt_request_url(url: &str) -> bool {
let Ok(url) = reqwest::Url::parse(url) else {
return false;
};
if !matches!(url.scheme(), "https" | "wss") {
return false;
}
url.host_str().is_some_and(is_allowed_chatgpt_host)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -36,4 +48,25 @@ mod tests {
assert!(!is_allowed_chatgpt_host(host));
}
}
#[test]
fn recognizes_secure_chatgpt_request_urls() {
for url in [
"https://chatgpt.com/backend-api/codex/responses",
"https://preview.chatgpt.com/backend-api/codex/models",
"wss://chatgpt-staging.com/backend-api/codex/responses",
] {
assert!(is_allowed_chatgpt_request_url(url));
}
for url in [
"http://chatgpt.com/backend-api/codex/responses",
"ws://chatgpt.com/backend-api/codex/responses",
"https://api.openai.com/v1/responses",
"https://chatgpt.com.evil.example/backend-api",
"not a url",
] {
assert!(!is_allowed_chatgpt_request_url(url));
}
}
}
+46 -1
View File
@@ -1,3 +1,5 @@
use crate::INTEGRITY_STATE_HEADER_NAME;
use crate::INTEGRITY_STATE_UPDATE_HEADER_NAME;
use http::Error as HttpError;
use http::HeaderMap;
use http::HeaderName;
@@ -115,11 +117,12 @@ impl CodexRequestBuilder {
match self.builder.headers(headers).send().await {
Ok(response) => {
let headers = redacted_headers_for_logging(response.headers());
tracing::debug!(
method = %self.method,
url = %self.url,
status = %response.status(),
headers = ?response.headers(),
headers = ?headers,
version = ?response.version(),
"Request completed"
);
@@ -141,6 +144,19 @@ impl CodexRequestBuilder {
}
}
fn redacted_headers_for_logging(headers: &HeaderMap) -> HeaderMap {
let mut headers = headers.clone();
for name in [
INTEGRITY_STATE_HEADER_NAME,
INTEGRITY_STATE_UPDATE_HEADER_NAME,
] {
if headers.contains_key(name) {
headers.insert(name, HeaderValue::from_static("<redacted>"));
}
}
headers
}
struct HeaderMapInjector<'a>(&'a mut HeaderMap);
impl<'a> Injector for HeaderMapInjector<'a> {
@@ -174,10 +190,39 @@ mod tests {
use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::SdkTracerProvider;
use pretty_assertions::assert_eq;
use tracing::trace_span;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
#[test]
fn redacted_headers_for_logging_hides_integrity_state() {
let mut headers = HeaderMap::new();
headers.insert(
INTEGRITY_STATE_HEADER_NAME,
HeaderValue::from_static("ois1.request.nonce.ciphertext"),
);
headers.insert(
INTEGRITY_STATE_UPDATE_HEADER_NAME,
HeaderValue::from_static("ois1.response.nonce.ciphertext"),
);
headers.insert("x-request-id", HeaderValue::from_static("request-id"));
let redacted = redacted_headers_for_logging(&headers);
let mut expected = HeaderMap::new();
expected.insert(
INTEGRITY_STATE_HEADER_NAME,
HeaderValue::from_static("<redacted>"),
);
expected.insert(
INTEGRITY_STATE_UPDATE_HEADER_NAME,
HeaderValue::from_static("<redacted>"),
);
expected.insert("x-request-id", HeaderValue::from_static("request-id"));
assert_eq!(redacted, expected);
}
#[test]
fn inject_trace_headers_uses_current_span_context() {
global::set_text_map_propagator(TraceContextPropagator::new());
+4
View File
@@ -11,6 +11,7 @@ mod transport;
pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store;
pub use crate::chatgpt_hosts::is_allowed_chatgpt_host;
pub use crate::chatgpt_hosts::is_allowed_chatgpt_request_url;
pub use crate::custom_ca::BuildCustomCaTransportError;
/// Test-only subprocess hook for custom CA coverage.
///
@@ -40,3 +41,6 @@ pub use crate::transport::ByteStream;
pub use crate::transport::HttpTransport;
pub use crate::transport::ReqwestTransport;
pub use crate::transport::StreamResponse;
pub const INTEGRITY_STATE_HEADER_NAME: &str = "x-oai-is";
pub const INTEGRITY_STATE_UPDATE_HEADER_NAME: &str = "x-oai-is-update";
+1
View File
@@ -48,6 +48,7 @@ codex-shell-command = { workspace = true }
codex-execpolicy = { workspace = true }
codex-git-utils = { workspace = true }
codex-hooks = { workspace = true }
codex-http-state = { workspace = true }
codex-install-context = { workspace = true }
codex-network-proxy = { workspace = true }
codex-otel = { workspace = true }
+15
View File
@@ -307,6 +307,21 @@ impl AgentControl {
{
let client_metadata = match state.get_thread(*parent_thread_id).await {
Ok(parent_thread) => {
if let Some(surface) = parent_thread
.codex
.session
.services
.model_client
.http_state_surface()
{
new_thread
.thread
.codex
.session
.services
.model_client
.set_http_state_surface(surface);
}
parent_thread
.codex
.session
+46
View File
@@ -11,6 +11,7 @@ use crate::context::SubagentNotification;
use crate::init_state_db;
use assert_matches::assert_matches;
use codex_features::Feature;
use codex_http_state::HttpStateSurface;
use codex_login::CodexAuth;
use codex_protocol::AgentPath;
use codex_protocol::config_types::ModeKind;
@@ -1584,6 +1585,51 @@ async fn spawn_child_completion_notifies_parent_history() {
assert_eq!(wait_for_subagent_notification(&parent_thread).await, true);
}
#[tokio::test]
async fn spawn_thread_subagent_inherits_parent_http_state_surface() {
let harness = AgentControlHarness::new().await;
let (parent_thread_id, parent_thread) = harness.start_thread().await;
parent_thread
.set_app_server_client_info(
Some("codex_desktop".to_string()),
/*app_server_client_version*/ None,
/*mcp_elicitations_auto_deny*/ false,
)
.await
.expect("parent client info should update");
let child_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello child"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: Some("explorer".to_string()),
})),
)
.await
.expect("child spawn should succeed");
let child_thread = harness
.manager
.get_thread(child_thread_id)
.await
.expect("child thread should exist");
assert_eq!(
child_thread
.codex
.session
.services
.model_client
.http_state_surface(),
Some(HttpStateSurface::CodexDesktop)
);
}
#[tokio::test]
async fn multi_agent_v2_completion_ignores_dead_direct_parent() {
let harness = AgentControlHarness::new().await;
+102 -26
View File
@@ -62,6 +62,8 @@ use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_api::response_create_client_metadata;
use codex_app_server_protocol::AuthMode;
use codex_http_state::HttpStateContext;
use codex_http_state::HttpStateSurface;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::RefreshTokenError;
@@ -120,6 +122,7 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_model_provider::with_native_integrity_state;
#[cfg(test)]
use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS;
use codex_model_provider_info::ModelProviderInfo;
@@ -180,7 +183,9 @@ struct ModelClientState {
beta_features_header: Option<String>,
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
http_state: Option<HttpStateContext>,
disable_websockets: AtomicBool,
websocket_session_generation: AtomicU64,
cached_websocket_session: StdMutex<WebsocketSession>,
}
@@ -237,6 +242,7 @@ pub struct ModelClient {
/// contract and can cause routing bugs.
pub struct ModelClientSession {
client: ModelClient,
websocket_session_generation: u64,
websocket_session: WebsocketSession,
/// Turn state for sticky routing.
///
@@ -296,17 +302,7 @@ pub(crate) struct RealtimeWebrtcCallStart {
pub(crate) sdp: String,
pub(crate) call_id: String,
pub(crate) sideband_headers: ApiHeaderMap,
}
/// Reuses the API-auth material that created the WebRTC call for the sideband WebSocket join.
///
/// API-key sessions send that API bearer. ChatGPT-auth sessions send their bearer plus account id;
/// transceiver is responsible for accepting that same call-create identity on the direct
/// `api.openai.com` sideband path.
fn sideband_websocket_auth_headers(api_auth: &dyn AuthProvider) -> ApiHeaderMap {
let mut headers = ApiHeaderMap::new();
api_auth.add_auth_headers(&mut headers);
headers
pub(crate) sideband_auth: SharedAuthProvider,
}
impl ModelClient {
@@ -329,6 +325,12 @@ impl ModelClient {
beta_features_header: Option<String>,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
) -> Self {
let http_state = auth_manager.as_ref().map(|auth_manager| {
HttpStateContext::new(
auth_manager.codex_home().to_path_buf(),
http_state_surface_for_session_source(&session_source),
)
});
let model_provider = create_model_provider(provider_info, auth_manager);
let codex_api_key_env_enabled = model_provider
.auth_manager()
@@ -353,7 +355,9 @@ impl ModelClient {
beta_features_header,
include_attestation,
attestation_provider,
http_state,
disable_websockets: AtomicBool::new(false),
websocket_session_generation: AtomicU64::new(0),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
}),
prompt_cache_key_override: None,
@@ -379,9 +383,12 @@ impl ModelClient {
/// This constructor does not perform network I/O itself; the session opens a websocket lazily
/// when the first stream request is issued.
pub fn new_session(&self) -> ModelClientSession {
let (websocket_session_generation, websocket_session) =
self.take_cached_websocket_session();
ModelClientSession {
client: self.clone(),
websocket_session: self.take_cached_websocket_session(),
websocket_session_generation,
websocket_session,
turn_state: Arc::new(OnceLock::new()),
}
}
@@ -390,16 +397,45 @@ impl ModelClient {
self.state.provider.auth_manager()
}
pub(crate) fn set_app_server_client_name(&self, client_name: Option<&str>) {
let surface = client_name
.map(HttpStateSurface::from_app_server_client_name)
.unwrap_or(HttpStateSurface::CodexCli);
self.set_http_state_surface(surface);
}
pub(crate) fn http_state_surface(&self) -> Option<HttpStateSurface> {
self.state
.http_state
.as_ref()
.map(HttpStateContext::surface)
}
pub(crate) fn http_state_context(&self) -> Option<HttpStateContext> {
self.state.http_state.clone()
}
pub(crate) fn set_http_state_surface(&self, surface: HttpStateSurface) {
if self
.state
.http_state
.as_ref()
.is_some_and(|state| state.set_surface(surface))
{
self.invalidate_cached_websocket_session();
}
}
pub(crate) fn set_window_generation(&self, window_generation: u64) {
self.state
.window_generation
.store(window_generation, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
}
pub(crate) fn advance_window_generation(&self) {
self.state.window_generation.fetch_add(1, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
}
pub(crate) fn current_window_id(&self) -> String {
@@ -408,21 +444,45 @@ impl ModelClient {
format!("{thread_id}:{window_generation}")
}
fn take_cached_websocket_session(&self) -> WebsocketSession {
fn take_cached_websocket_session(&self) -> (u64, WebsocketSession) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::take(&mut *cached_websocket_session)
let generation = self
.state
.websocket_session_generation
.load(Ordering::Relaxed);
(generation, std::mem::take(&mut *cached_websocket_session))
}
fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) {
*self
fn invalidate_cached_websocket_session(&self) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session;
.unwrap_or_else(std::sync::PoisonError::into_inner);
self.state
.websocket_session_generation
.fetch_add(1, Ordering::Relaxed);
*cached_websocket_session = WebsocketSession::default();
}
fn store_cached_websocket_session(&self, generation: u64, websocket_session: WebsocketSession) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if self
.state
.websocket_session_generation
.load(Ordering::Relaxed)
== generation
{
*cached_websocket_session = websocket_session;
}
}
pub(crate) fn force_http_fallback(
@@ -442,7 +502,7 @@ impl ModelClient {
);
}
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
activated
}
@@ -554,10 +614,8 @@ impl ModelClient {
if let Some(header_value) = self.generate_attestation_header_for().await {
extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
let mut sideband_headers = extra_headers.clone();
sideband_headers.extend(sideband_websocket_auth_headers(
client_setup.api_auth.as_ref(),
));
let sideband_headers = extra_headers.clone();
let sideband_auth = Arc::clone(&client_setup.api_auth);
let transport = ReqwestTransport::new(build_reqwest_client());
let response =
ApiRealtimeCallClient::new(transport, client_setup.api_provider, client_setup.api_auth)
@@ -568,6 +626,7 @@ impl ModelClient {
sdp: response.sdp,
call_id: response.call_id,
sideband_headers,
sideband_auth,
})
}
@@ -812,7 +871,11 @@ impl ModelClient {
async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
let auth = self.state.provider.auth().await;
let api_provider = self.state.provider.api_provider().await?;
let api_auth = self.state.provider.api_auth().await?;
let api_auth = with_native_integrity_state(
self.state.provider.api_auth().await?,
auth.as_ref(),
self.state.http_state.clone(),
);
Ok(CurrentClientSetup {
auth,
api_provider,
@@ -953,7 +1016,7 @@ impl Drop for ModelClientSession {
fn drop(&mut self) {
let websocket_session = std::mem::take(&mut self.websocket_session);
self.client
.store_cached_websocket_session(websocket_session);
.store_cached_websocket_session(self.websocket_session_generation, websocket_session);
}
}
@@ -1736,6 +1799,19 @@ fn subagent_header_value(session_source: &SessionSource) -> Option<String> {
}
}
fn http_state_surface_for_session_source(session_source: &SessionSource) -> HttpStateSurface {
match session_source {
SessionSource::Cli => HttpStateSurface::CodexTui,
SessionSource::Exec => HttpStateSurface::CodexExec,
SessionSource::VSCode => HttpStateSurface::CodexVscode,
SessionSource::SubAgent(_)
| SessionSource::Internal(_)
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::Unknown => HttpStateSurface::CodexCli,
}
}
fn parent_thread_id_header_value(parent_thread_id: Option<ThreadId>) -> Option<String> {
parent_thread_id.map(|parent_thread_id| parent_thread_id.to_string())
}
+96 -1
View File
@@ -13,12 +13,15 @@ use crate::GenerateAttestationFuture;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_app_server_protocol::AuthMode;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_model_provider::BearerAuthProvider;
use codex_model_provider_info::CHATGPT_CODEX_BASE_URL;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_model_provider_info::built_in_model_providers;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_otel::SessionTelemetry;
use codex_protocol::SessionId;
@@ -67,11 +70,32 @@ fn test_model_client(session_source: SessionSource) -> ModelClient {
fn test_model_client_with_parent(
session_source: SessionSource,
parent_thread_id: Option<ThreadId>,
) -> ModelClient {
test_model_client_with_parent_and_auth(
session_source,
parent_thread_id,
/*auth_manager*/ None,
)
}
fn test_model_client_with_parent_and_auth(
session_source: SessionSource,
parent_thread_id: Option<ThreadId>,
auth_manager: Option<Arc<AuthManager>>,
) -> ModelClient {
let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
test_model_client_with_provider(session_source, parent_thread_id, auth_manager, provider)
}
fn test_model_client_with_provider(
session_source: SessionSource,
parent_thread_id: Option<ThreadId>,
auth_manager: Option<Arc<AuthManager>>,
provider: ModelProviderInfo,
) -> ModelClient {
let thread_id = ThreadId::new();
ModelClient::new(
/*auth_manager*/ None,
auth_manager,
thread_id.into(),
thread_id,
/*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(),
@@ -322,6 +346,77 @@ fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() {
);
}
#[tokio::test]
async fn current_client_setup_attaches_and_rotates_native_http_state() {
let codex_home = TempDir::new().expect("tempdir");
let store = HttpStateStore::new(codex_home.path().to_path_buf());
store
.set(HttpStateSurface::CodexTui, "old-state".to_string())
.expect("state should store");
let auth_manager = AuthManager::from_auth_for_testing_with_home(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
codex_home.path().to_path_buf(),
);
let client = test_model_client_with_provider(
SessionSource::Cli,
/*parent_thread_id*/ None,
Some(auth_manager),
built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(),
);
let setup = client
.current_client_setup()
.await
.expect("client setup should resolve");
let request_url = "https://chatgpt.com/backend-api/codex/responses";
let mut request_headers = http::HeaderMap::new();
setup
.api_auth
.add_auth_headers_for_url(request_url, &mut request_headers);
assert_eq!(
request_headers
.get("x-oai-is")
.and_then(|value| value.to_str().ok()),
Some("old-state"),
);
let mut response_headers = http::HeaderMap::new();
response_headers.insert(
"x-oai-is-update",
http::HeaderValue::from_static("new-state"),
);
setup
.api_auth
.observe_response_headers(request_url, &request_headers, &response_headers);
assert_eq!(
store
.get(HttpStateSurface::CodexTui)
.expect("state should load"),
Some("new-state".to_string()),
);
}
#[test]
fn surface_change_discards_checked_out_websocket_session() {
let codex_home = TempDir::new().expect("tempdir");
let auth_manager = AuthManager::from_auth_for_testing_with_home(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
codex_home.path().to_path_buf(),
);
let client = test_model_client_with_parent_and_auth(
SessionSource::Cli,
/*parent_thread_id*/ None,
Some(auth_manager),
);
let mut session = client.new_session();
session.websocket_session.last_response_from_untraced_warmup = true;
client.set_http_state_surface(HttpStateSurface::CodexDesktop);
drop(session);
let (_, cached_session) = client.take_cached_websocket_session();
assert!(!cached_session.last_response_from_untraced_warmup);
}
#[tokio::test]
async fn summarize_memories_returns_empty_for_empty_input() {
let client = test_model_client(SessionSource::Cli);
+7
View File
@@ -110,6 +110,13 @@ pub(crate) async fn run_codex_thread_interactive(
}))
.or_cancel(&cancel_token)
.await??;
if let Some(surface) = parent_session.services.model_client.http_state_surface() {
codex
.session
.services
.model_client
.set_http_state_surface(surface);
}
let thread_config = codex.thread_config_snapshot().await;
let client_metadata = parent_session.app_server_client_metadata().await;
emit_subagent_session_started(
+31
View File
@@ -2,6 +2,7 @@ use super::*;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX;
use async_channel::bounded;
use codex_http_state::HttpStateSurface;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::models::NetworkPermissions;
@@ -181,6 +182,36 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() {
assert!(matches!(result, Err(CodexErr::TurnAborted)));
}
#[tokio::test]
async fn run_codex_thread_interactive_inherits_parent_http_state_surface() {
let (parent_session, parent_ctx, _rx_events) =
crate::session::tests::make_session_and_context_with_rx().await;
parent_session
.services
.model_client
.set_http_state_surface(HttpStateSurface::CodexDesktop);
let cancel_token = CancellationToken::new();
let delegated = run_codex_thread_interactive(
parent_ctx.config.as_ref().clone(),
Arc::clone(&parent_session.services.auth_manager),
Arc::clone(&parent_session.services.models_manager),
parent_session,
parent_ctx,
cancel_token.clone(),
SubAgentSource::Review,
/*initial_history*/ None,
)
.await
.expect("delegate should spawn");
assert_eq!(
delegated.session.services.model_client.http_state_surface(),
Some(HttpStateSurface::CodexDesktop)
);
cancel_token.cancel();
}
#[tokio::test]
async fn handle_request_permissions_uses_tool_call_id_for_round_trip() {
let (parent_session, mut parent_ctx, rx_events) =
+22 -4
View File
@@ -13,6 +13,7 @@
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use codex_api::upload_local_file;
use codex_http_state::HttpStateContext;
use codex_login::CodexAuth;
use serde_json::Value as JsonValue;
@@ -33,15 +34,21 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
return Ok(Some(arguments_value));
};
let auth = sess.services.auth_manager.auth().await;
let http_state = sess.services.model_client.http_state_context();
let mut rewritten_arguments = arguments.clone();
for field_name in openai_file_input_params {
let Some(value) = arguments.get(field_name) else {
continue;
};
let Some(uploaded_value) =
rewrite_argument_value_for_openai_files(turn_context, auth.as_ref(), field_name, value)
.await?
let Some(uploaded_value) = rewrite_argument_value_for_openai_files(
turn_context,
auth.as_ref(),
http_state.clone(),
field_name,
value,
)
.await?
else {
continue;
};
@@ -58,6 +65,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
async fn rewrite_argument_value_for_openai_files(
turn_context: &TurnContext,
auth: Option<&CodexAuth>,
http_state: Option<HttpStateContext>,
field_name: &str,
value: &JsonValue,
) -> Result<Option<JsonValue>, String> {
@@ -66,6 +74,7 @@ async fn rewrite_argument_value_for_openai_files(
let rewritten = build_uploaded_local_argument_value(
turn_context,
auth,
http_state,
field_name,
/*index*/ None,
path_or_file_ref,
@@ -82,6 +91,7 @@ async fn rewrite_argument_value_for_openai_files(
let rewritten = build_uploaded_local_argument_value(
turn_context,
auth,
http_state.clone(),
field_name,
Some(index),
path_or_file_ref,
@@ -98,6 +108,7 @@ async fn rewrite_argument_value_for_openai_files(
async fn build_uploaded_local_argument_value(
turn_context: &TurnContext,
auth: Option<&CodexAuth>,
http_state: Option<HttpStateContext>,
field_name: &str,
index: Option<usize>,
file_path: &str,
@@ -114,7 +125,11 @@ async fn build_uploaded_local_argument_value(
"ChatGPT auth is required to upload local files for Codex Apps tools".to_string(),
);
}
let upload_auth = codex_model_provider::auth_provider_from_auth(auth);
let upload_auth = codex_model_provider::with_native_integrity_state(
codex_model_provider::auth_provider_from_auth(auth),
Some(auth),
http_state,
);
let uploaded = upload_local_file(
turn_context.config.chatgpt_base_url.trim_end_matches('/'),
upload_auth.as_ref(),
@@ -229,6 +244,7 @@ mod tests {
let rewritten = build_uploaded_local_argument_value(
&turn_context,
Some(&auth),
/*http_state*/ None,
"file",
/*index*/ None,
"file_report.csv",
@@ -312,6 +328,7 @@ mod tests {
let rewritten = rewrite_argument_value_for_openai_files(
&turn_context,
Some(&auth),
/*http_state*/ None,
"file",
&serde_json::json!("file_report.csv"),
)
@@ -429,6 +446,7 @@ mod tests {
let rewritten = rewrite_argument_value_for_openai_files(
&turn_context,
Some(&auth),
/*http_state*/ None,
"files",
&serde_json::json!(["one.csv", "two.csv"]),
)
@@ -19,6 +19,7 @@ use codex_api::RealtimeSessionMode;
use codex_api::RealtimeWebsocketClient;
use codex_api::RealtimeWebsocketEvents;
use codex_api::RealtimeWebsocketWriter;
use codex_api::SharedAuthProvider;
use codex_api::map_api_error;
use codex_app_server_protocol::AuthMode;
use codex_config::config_toml::RealtimeWsMode;
@@ -320,6 +321,7 @@ impl RealtimeConversationManager {
session_config,
call_id: call.call_id,
sideband_headers: call.sideband_headers,
sideband_auth: call.sideband_auth,
input_channels,
events_tx,
handoff_state: handoff.clone(),
@@ -1029,6 +1031,7 @@ struct RealtimeWebrtcSidebandInputTask {
session_config: RealtimeSessionConfig,
call_id: String,
sideband_headers: HeaderMap,
sideband_auth: SharedAuthProvider,
input_channels: RealtimeInputChannels,
events_tx: Sender<RealtimeEvent>,
handoff_state: RealtimeHandoffState,
@@ -1043,6 +1046,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J
session_config,
call_id,
sideband_headers,
sideband_auth,
input_channels,
events_tx,
handoff_state,
@@ -1057,6 +1061,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J
}
let connection = match client
.with_auth(sideband_auth)
.connect_webrtc_sideband(
session_config,
&call_id,
+4
View File
@@ -782,6 +782,10 @@ impl Codex {
app_server_client_version: Option<String>,
mcp_elicitations_auto_deny: bool,
) -> ConstraintResult<()> {
self.session
.services
.model_client
.set_app_server_client_name(app_server_client_name.as_deref());
self.session
.update_settings(SessionSettingsUpdate {
app_server_client_name,
+4
View File
@@ -241,6 +241,10 @@ impl HttpStateContext {
self.store.get(surface)
}
pub fn set_for_surface(&self, surface: HttpStateSurface, state: String) -> io::Result<()> {
self.store.set(surface, state)
}
pub fn compare_and_set_for_surface(
&self,
surface: HttpStateSurface,
+4
View File
@@ -1401,6 +1401,10 @@ impl AuthManager {
})
}
pub fn codex_home(&self) -> &Path {
&self.codex_home
}
/// Current cached auth (clone) without attempting a refresh.
pub fn auth_cached(&self) -> Option<CodexAuth> {
self.inner.read().ok().and_then(|c| c.auth.clone())
+2
View File
@@ -19,6 +19,7 @@ codex-agent-identity = { workspace = true }
codex-aws-auth = { workspace = true }
codex-client = { workspace = true }
codex-feedback = { workspace = true }
codex-http-state = { workspace = true }
codex-login = { workspace = true }
codex-model-provider-info = { workspace = true }
codex-models-manager = { workspace = true }
@@ -32,5 +33,6 @@ tracing = { workspace = true, features = ["log"] }
[dev-dependencies]
pretty_assertions = { workspace = true }
serde_json = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
wiremock = { workspace = true }
+19
View File
@@ -5,6 +5,7 @@ use codex_agent_identity::AgentTaskAuthorizationTarget;
use codex_agent_identity::authorization_header_for_agent_task;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_http_state::HttpStateContext;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_model_provider_info::ModelProviderInfo;
@@ -12,6 +13,7 @@ use http::HeaderMap;
use http::HeaderValue;
use crate::bearer_auth_provider::BearerAuthProvider;
use crate::bearer_auth_provider::NativeIntegrityAuthProvider;
#[derive(Clone, Debug)]
struct AgentIdentityAuthProvider {
@@ -119,6 +121,23 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider {
}
}
/// Decorates first-party ChatGPT auth with native integrity-state transport.
pub fn with_native_integrity_state(
auth_provider: SharedAuthProvider,
auth: Option<&CodexAuth>,
state: Option<HttpStateContext>,
) -> SharedAuthProvider {
let Some(state) = state else {
return auth_provider;
};
let Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_)) =
auth
else {
return auth_provider;
};
Arc::new(NativeIntegrityAuthProvider::new(auth_provider, state))
}
#[cfg(test)]
mod tests {
use codex_model_provider_info::WireApi;
@@ -1,7 +1,15 @@
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_client::INTEGRITY_STATE_HEADER_NAME;
use codex_client::INTEGRITY_STATE_UPDATE_HEADER_NAME;
use codex_client::is_allowed_chatgpt_request_url;
use codex_http_state::HttpStateContext;
use codex_http_state::HttpStateSurface;
use http::HeaderMap;
use http::HeaderValue;
const OAI_CLIENT_TYPE_HEADER_NAME: &str = "oai-client-type";
/// Bearer-token auth provider for OpenAI-compatible model-provider requests.
#[derive(Clone, Default)]
pub struct BearerAuthProvider {
@@ -46,10 +54,97 @@ impl AuthProvider for BearerAuthProvider {
}
}
#[derive(Clone)]
pub(crate) struct NativeIntegrityAuthProvider {
auth: SharedAuthProvider,
state: HttpStateContext,
surface: HttpStateSurface,
}
impl NativeIntegrityAuthProvider {
pub(crate) fn new(auth: SharedAuthProvider, state: HttpStateContext) -> Self {
let surface = state.surface();
Self {
auth,
state,
surface,
}
}
}
impl AuthProvider for NativeIntegrityAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
self.auth.add_auth_headers(headers);
}
fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) {
self.auth.add_auth_headers_for_url(request_url, headers);
if !is_allowed_chatgpt_request_url(request_url) {
return;
}
if let Ok(header) = HeaderValue::from_str(self.surface.as_str()) {
let _ = headers.insert(OAI_CLIENT_TYPE_HEADER_NAME, header);
}
match self.state.get_for_surface(self.surface) {
Ok(Some(state)) => {
if let Ok(header) = HeaderValue::from_str(&state) {
let _ = headers.insert(INTEGRITY_STATE_HEADER_NAME, header);
}
}
Ok(None) => {}
Err(error) => {
tracing::warn!("failed to load native integrity state: {error}");
}
}
}
fn observe_response_headers(
&self,
request_url: &str,
request_headers: &HeaderMap,
response_headers: &HeaderMap,
) {
self.auth
.observe_response_headers(request_url, request_headers, response_headers);
if !is_allowed_chatgpt_request_url(request_url) {
return;
}
let Some(next_state) = response_headers
.get(INTEGRITY_STATE_UPDATE_HEADER_NAME)
.and_then(|value| value.to_str().ok())
else {
return;
};
let expected_state = request_headers
.get(INTEGRITY_STATE_HEADER_NAME)
.and_then(|value| value.to_str().ok());
let result = match expected_state {
Some(expected_state) => self.state.compare_and_set_for_surface(
self.surface,
expected_state,
next_state.to_string(),
),
None => self
.state
.set_for_surface(self.surface, next_state.to_string())
.map(|()| true),
};
if let Err(error) = result {
tracing::warn!("failed to rotate native integrity state: {error}");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use codex_http_state::HttpStateStore;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use tempfile::TempDir;
#[test]
fn bearer_auth_provider_reports_when_auth_header_will_attach() {
@@ -107,4 +202,115 @@ mod tests {
Some("true")
);
}
#[test]
fn native_integrity_provider_scopes_and_rotates_surface_state() {
let codex_home = TempDir::new().expect("tempdir");
let context = HttpStateContext::new(
codex_home.path().to_path_buf(),
HttpStateSurface::CodexDesktop,
);
let store = HttpStateStore::new(codex_home.path().to_path_buf());
store
.set(
HttpStateSurface::CodexDesktop,
"ois1.initial.nonce.ciphertext".to_string(),
)
.expect("state should store");
let provider = NativeIntegrityAuthProvider::new(
Arc::new(BearerAuthProvider::new("access-token".to_string())),
context.clone(),
);
context.set_surface(HttpStateSurface::CodexCli);
store
.set(
HttpStateSurface::CodexCli,
"ois1.cli.nonce.ciphertext".to_string(),
)
.expect("CLI state should store");
let mut request_headers = HeaderMap::new();
provider.add_auth_headers_for_url(
"https://chatgpt.com/backend-api/codex/responses",
&mut request_headers,
);
assert_eq!(
request_headers
.get(INTEGRITY_STATE_HEADER_NAME)
.and_then(|value| value.to_str().ok()),
Some("ois1.initial.nonce.ciphertext")
);
assert_eq!(
request_headers
.get(OAI_CLIENT_TYPE_HEADER_NAME)
.and_then(|value| value.to_str().ok()),
Some("codex_desktop")
);
let mut response_headers = HeaderMap::new();
response_headers.insert(
INTEGRITY_STATE_UPDATE_HEADER_NAME,
HeaderValue::from_static("ois1.rotated.nonce.ciphertext"),
);
provider.observe_response_headers(
"https://chatgpt.com/backend-api/codex/responses",
&request_headers,
&response_headers,
);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load")
.expect("state should exist"),
"ois1.rotated.nonce.ciphertext"
);
assert_eq!(
store
.get(HttpStateSurface::CodexCli)
.expect("CLI state should load")
.expect("CLI state should exist"),
"ois1.cli.nonce.ciphertext"
);
let mut external_headers = HeaderMap::new();
provider.add_auth_headers_for_url(
"https://example.com/backend-api/codex/responses",
&mut external_headers,
);
assert!(!external_headers.contains_key(INTEGRITY_STATE_HEADER_NAME));
assert!(!external_headers.contains_key(OAI_CLIENT_TYPE_HEADER_NAME));
}
#[test]
fn native_integrity_provider_seeds_empty_surface_from_response() {
let codex_home = TempDir::new().expect("tempdir");
let context = HttpStateContext::new(
codex_home.path().to_path_buf(),
HttpStateSurface::CodexDesktop,
);
let store = HttpStateStore::new(codex_home.path().to_path_buf());
let provider = NativeIntegrityAuthProvider::new(
Arc::new(BearerAuthProvider::new("access-token".to_string())),
context,
);
let mut response_headers = HeaderMap::new();
response_headers.insert(
INTEGRITY_STATE_UPDATE_HEADER_NAME,
HeaderValue::from_static("ois1.seeded.nonce.ciphertext"),
);
provider.observe_response_headers(
"https://chatgpt.com/backend-api/codex/responses",
&HeaderMap::new(),
&response_headers,
);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load")
.expect("state should exist"),
"ois1.seeded.nonce.ciphertext"
);
}
}
+1
View File
@@ -6,6 +6,7 @@ mod provider;
pub use auth::auth_provider_from_auth;
pub use auth::unauthenticated_auth_provider;
pub use auth::with_native_integrity_state;
pub use bearer_auth_provider::BearerAuthProvider;
pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider;
pub use codex_protocol::account::ProviderAccount;