agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Use background task auth for additional backend calls (#18260)

Browse Source 
## Summary

Splits the larger PR4.1 background task auth rollout by moving
additional backend/control-plane call sites into this downstream PR.

This PR keeps callers on the same design as PR4.1: most code asks
`AuthManager` for the default ChatGPT backend authorization header, and
`AuthManager` decides bearer vs background AgentAssertion internally.
Task-pinned inference auth remains separate because it needs the
thread's registered task id.

## Stack

- PR1: https://github.com/openai/codex/pull/17385 - add
`features.use_agent_identity`
- PR2: https://github.com/openai/codex/pull/17386 - register agent
identities when enabled
- PR3: https://github.com/openai/codex/pull/17387 - register agent tasks
when enabled
- PR3.1: https://github.com/openai/codex/pull/17978 - persist and
prewarm registered tasks per thread
- PR4: https://github.com/openai/codex/pull/17980 - use task-scoped
`AgentAssertion` for downstream calls
- PR4.1: https://github.com/openai/codex/pull/18094 - introduce
AuthManager-owned background/control-plane `AgentAssertion` auth
- PR4.2: this PR - use background task auth for additional
backend/control-plane calls

## What Changed

- pass full authorization header values through backend-client and
cloud-tasks-client call paths where needed
- move ChatGPT client, cloud requirements, cloud tasks, thread-manager,
and models-manager background auth usage into this downstream slice
- make app-server remote control enrollment/websocket auth ask
`AuthManager` for the local backend authorization header instead of
threading a background auth mode through transport options
- keep the same feature-gated bearer fallback behavior from PR4.1

## Validation

- `just fmt`
- `cargo check -p codex-core -p codex-login -p codex-analytics -p
codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p
codex-models-manager -p codex-chatgpt -p codex-model-provider -p
codex-mcp -p codex-core-skills`
- `cargo test -p codex-login agent_identity`
- `cargo test -p codex-model-provider bearer_auth_provider`
- `cargo test -p codex-core agent_assertion`
- `cargo test -p codex-app-server remote_control`
- `cargo test -p codex-cloud-requirements fetch_cloud_requirements`
- `cargo test -p codex-models-manager manager::tests`
- `cargo test -p codex-chatgpt`
- `cargo test -p codex-cloud-tasks`
- `just fix -p codex-core -p codex-login -p codex-analytics -p
codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p
codex-models-manager -p codex-chatgpt -p codex-model-provider -p
codex-mcp -p codex-core-skills`
- `just fix -p codex-app-server`
- `git diff --check`
This commit is contained in:
Adrian
2026-04-20 07:24:29 -07:00
committed by GitHub
Unverified
parent fa0e2ba87c
commit 19e2f21827
14 changed files with 364 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/codex_message_processor.rs
+17 -1
View File
@@ -2013,12 +2013,28 @@ impl CodexMessageProcessor {
});
}
let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth)
let authorization_header_value = self
.auth_manager
.chatgpt_authorization_header_for_auth(&auth)
.await;
let mut client = BackendClient::new(self.config.chatgpt_base_url.clone())
.map(|client| {
client.with_user_agent(codex_login::default_client::get_codex_user_agent())
})
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to construct backend client: {err}"),
data: None,
})?;
if let Some(authorization_header_value) = authorization_header_value {
client = client.with_authorization_header_value(authorization_header_value);
}
if let Some(account_id) = auth.get_account_id() {
client = client.with_chatgpt_account_id(account_id);
}
if auth.is_fedramp_account() {
client = client.with_fedramp_routing_header();
}
let snapshots = client
.get_rate_limits_many()
codex-rs/app-server/src/lib.rs
+13 -11
View File
@@ -27,10 +27,11 @@ use crate::outgoing_message::QueuedOutgoingMessage;
use crate::transport::CHANNEL_CAPACITY;
use crate::transport::ConnectionState;
use crate::transport::OutboundConnectionState;
use crate::transport::RemoteControlStartOptions;
use crate::transport::TransportEvent;
use crate::transport::auth::policy_from_settings;
use crate::transport::route_outgoing_envelope;
use crate::transport::start_remote_control;
use crate::transport::start_remote_control_with_options;
use crate::transport::start_stdio_connection;
use crate::transport::start_websocket_acceptor;
use codex_analytics::AppServerRpcTransport;
@@ -575,16 +576,17 @@ pub async fn run_main_with_transport(
));
}
let (remote_control_accept_handle, remote_control_handle) = start_remote_control(
config.chatgpt_base_url.clone(),
state_db.clone(),
auth_manager.clone(),
transport_event_tx.clone(),
transport_shutdown_token.clone(),
app_server_client_name_rx,
remote_control_enabled,
)
.await?;
let (remote_control_accept_handle, remote_control_handle) =
start_remote_control_with_options(RemoteControlStartOptions {
remote_control_url: config.chatgpt_base_url.clone(),
state_db: state_db.clone(),
auth_manager: auth_manager.clone(),
transport_event_tx: transport_event_tx.clone(),
shutdown_token: transport_shutdown_token.clone(),
app_server_client_name_rx,
initial_enabled: remote_control_enabled,
})
.await?;
transport_accept_handles.push(remote_control_accept_handle);
let outbound_handle = tokio::spawn(async move {
codex-rs/app-server/src/transport/mod.rs
+2 -1
View File
@@ -34,7 +34,8 @@ mod stdio;
mod websocket;
pub(crate) use remote_control::RemoteControlHandle;
pub(crate) use remote_control::start_remote_control;
pub(crate) use remote_control::RemoteControlStartOptions;
pub(crate) use remote_control::start_remote_control_with_options;
pub(crate) use stdio::start_stdio_connection;
pub(crate) use websocket::start_websocket_acceptor;
codex-rs/app-server/src/transport/remote_control/enroll.rs
+12 -6
View File
@@ -17,6 +17,7 @@ const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096;
const REQUEST_ID_HEADER: &str = "x-request-id";
const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";
const CF_RAY_HEADER: &str = "cf-ray";
const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp";
pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -29,8 +30,9 @@ pub(super) struct RemoteControlEnrollment {
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct RemoteControlConnectionAuth {
pub(super) bearer_token: String,
pub(super) authorization_header_value: String,
pub(super) account_id: String,
pub(super) is_fedramp_account: bool,
}
pub(super) async fn load_persisted_remote_control_enrollment(
@@ -199,12 +201,15 @@ pub(super) async fn enroll_remote_control_server(
app_server_version: env!("CARGO_PKG_VERSION"),
};
let client = build_reqwest_client();
let http_request = client
let mut http_request = client
.post(enroll_url)
.timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
.bearer_auth(&auth.bearer_token)
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
.json(&request);
.header("authorization", &auth.authorization_header_value)
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id);
if auth.is_fedramp_account {
http_request = http_request.header(REMOTE_CONTROL_FEDRAMP_HEADER, "true");
}
let http_request = http_request.json(&request);
let response = http_request.send().await.map_err(|err| {
io::Error::other(format!(
@@ -445,8 +450,9 @@ mod tests {
let err = enroll_remote_control_server(
&remote_control_target,
&RemoteControlConnectionAuth {
bearer_token: "Access Token".to_string(),
authorization_header_value: "Bearer Access Token".to_string(),
account_id: "account_id".to_string(),
is_fedramp_account: false,
},
)
.await
codex-rs/app-server/src/transport/remote_control/mod.rs
+38 -3
View File
@@ -4,6 +4,7 @@ mod protocol;
mod websocket;
use crate::transport::remote_control::websocket::RemoteControlWebsocket;
use crate::transport::remote_control::websocket::RemoteControlWebsocketOptions;
pub use self::protocol::ClientId;
use self::protocol::ServerEvent;
@@ -44,6 +45,17 @@ impl RemoteControlHandle {
}
}
pub(crate) struct RemoteControlStartOptions {
pub(crate) remote_control_url: String,
pub(crate) state_db: Option<Arc<StateRuntime>>,
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) transport_event_tx: mpsc::Sender<TransportEvent>,
pub(crate) shutdown_token: CancellationToken,
pub(crate) app_server_client_name_rx: Option<oneshot::Receiver<String>>,
pub(crate) initial_enabled: bool,
}
#[cfg(test)]
pub(crate) async fn start_remote_control(
remote_control_url: String,
state_db: Option<Arc<StateRuntime>>,
@@ -53,15 +65,38 @@ pub(crate) async fn start_remote_control(
app_server_client_name_rx: Option<oneshot::Receiver<String>>,
initial_enabled: bool,
) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
start_remote_control_with_options(RemoteControlStartOptions {
remote_control_url,
state_db,
auth_manager,
transport_event_tx,
shutdown_token,
app_server_client_name_rx,
initial_enabled,
})
.await
}
pub(crate) async fn start_remote_control_with_options(
options: RemoteControlStartOptions,
) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
let RemoteControlStartOptions {
remote_control_url,
state_db,
auth_manager,
transport_event_tx,
shutdown_token,
app_server_client_name_rx,
initial_enabled,
} = options;
let remote_control_target = if initial_enabled {
Some(normalize_remote_control_url(&remote_control_url)?)
} else {
None
};
let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
let join_handle = tokio::spawn(async move {
RemoteControlWebsocket::new(
RemoteControlWebsocket::from_options(RemoteControlWebsocketOptions {
remote_control_url,
remote_control_target,
state_db,
@@ -69,7 +104,7 @@ pub(crate) async fn start_remote_control(
transport_event_tx,
shutdown_token,
enabled_rx,
)
})
.run(app_server_client_name_rx)
.await;
});
codex-rs/app-server/src/transport/remote_control/websocket.rs
+126 -14
View File
@@ -49,6 +49,7 @@ use tracing::warn;
pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2";
pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp";
const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor";
const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration =
std::time::Duration::from_secs(10);
@@ -128,6 +129,16 @@ pub(crate) struct RemoteControlWebsocket {
enabled_rx: watch::Receiver<bool>,
}
pub(crate) struct RemoteControlWebsocketOptions {
pub(crate) remote_control_url: String,
pub(crate) remote_control_target: Option<RemoteControlTarget>,
pub(crate) state_db: Option<Arc<StateRuntime>>,
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) transport_event_tx: mpsc::Sender<TransportEvent>,
pub(crate) shutdown_token: CancellationToken,
pub(crate) enabled_rx: watch::Receiver<bool>,
}
enum ConnectOutcome {
Connected(Box<WebSocketStream<MaybeTlsStream<TcpStream>>>),
Disabled,
@@ -135,6 +146,7 @@ enum ConnectOutcome {
}
impl RemoteControlWebsocket {
#[cfg(test)]
pub(crate) fn new(
remote_control_url: String,
remote_control_target: Option<RemoteControlTarget>,
@@ -144,6 +156,27 @@ impl RemoteControlWebsocket {
shutdown_token: CancellationToken,
enabled_rx: watch::Receiver<bool>,
) -> Self {
Self::from_options(RemoteControlWebsocketOptions {
remote_control_url,
remote_control_target,
state_db,
auth_manager,
transport_event_tx,
shutdown_token,
enabled_rx,
})
}
pub(crate) fn from_options(options: RemoteControlWebsocketOptions) -> Self {
let RemoteControlWebsocketOptions {
remote_control_url,
remote_control_target,
state_db,
auth_manager,
transport_event_tx,
shutdown_token,
enabled_rx,
} = options;
let shutdown_token = shutdown_token.child_token();
let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY);
let client_tracker =
@@ -271,14 +304,16 @@ impl RemoteControlWebsocket {
}
return ConnectOutcome::Disabled;
}
connect_result = connect_remote_control_websocket(
&remote_control_target,
self.state_db.as_deref(),
&self.auth_manager,
&mut self.auth_recovery,
&mut self.enrollment,
subscribe_cursor.as_deref(),
app_server_client_name,
connect_result = connect_remote_control_websocket_with_options(
ConnectRemoteControlWebsocketOptions {
remote_control_target: &remote_control_target,
state_db: self.state_db.as_deref(),
auth_manager: &self.auth_manager,
auth_recovery: &mut self.auth_recovery,
enrollment: &mut self.enrollment,
subscribe_cursor: subscribe_cursor.as_deref(),
app_server_client_name,
},
) => connect_result,
};
@@ -668,12 +703,11 @@ fn build_remote_control_websocket_request(
"x-codex-protocol-version",
REMOTE_CONTROL_PROTOCOL_VERSION,
)?;
set_remote_control_header(
headers,
"authorization",
&format!("Bearer {}", auth.bearer_token),
)?;
set_remote_control_header(headers, "authorization", &auth.authorization_header_value)?;
set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?;
if auth.is_fedramp_account {
set_remote_control_header(headers, REMOTE_CONTROL_FEDRAMP_HEADER, "true")?;
}
if let Some(subscribe_cursor) = subscribe_cursor {
set_remote_control_header(
headers,
@@ -718,8 +752,19 @@ pub(crate) async fn load_remote_control_auth(
));
}
let authorization_header_value = auth_manager
.chatgpt_authorization_header_for_auth(&auth)
.await
.ok_or_else(|| {
io::Error::new(
ErrorKind::PermissionDenied,
"remote control requires ChatGPT authentication",
)
})?;
Ok(RemoteControlConnectionAuth {
bearer_token: auth.get_token().map_err(io::Error::other)?,
authorization_header_value,
is_fedramp_account: auth.is_fedramp_account(),
account_id: auth.get_account_id().ok_or_else(|| {
io::Error::new(
ErrorKind::WouldBlock,
@@ -729,6 +774,7 @@ pub(crate) async fn load_remote_control_auth(
})
}
#[cfg(test)]
pub(super) async fn connect_remote_control_websocket(
remote_control_target: &RemoteControlTarget,
state_db: Option<&StateRuntime>,
@@ -741,6 +787,44 @@ pub(super) async fn connect_remote_control_websocket(
WebSocketStream<MaybeTlsStream<TcpStream>>,
tungstenite::http::Response<()>,
)> {
connect_remote_control_websocket_with_options(ConnectRemoteControlWebsocketOptions {
remote_control_target,
state_db,
auth_manager,
auth_recovery,
enrollment,
subscribe_cursor,
app_server_client_name,
})
.await
}
struct ConnectRemoteControlWebsocketOptions<'a> {
remote_control_target: &'a RemoteControlTarget,
state_db: Option<&'a StateRuntime>,
auth_manager: &'a Arc<AuthManager>,
auth_recovery: &'a mut UnauthorizedRecovery,
enrollment: &'a mut Option<RemoteControlEnrollment>,
subscribe_cursor: Option<&'a str>,
app_server_client_name: Option<&'a str>,
}
async fn connect_remote_control_websocket_with_options(
options: ConnectRemoteControlWebsocketOptions<'_>,
) -> io::Result<(
WebSocketStream<MaybeTlsStream<TcpStream>>,
tungstenite::http::Response<()>,
)> {
let ConnectRemoteControlWebsocketOptions {
remote_control_target,
state_db,
auth_manager,
auth_recovery,
enrollment,
subscribe_cursor,
app_server_client_name,
} = options;
ensure_rustls_crypto_provider();
let auth = load_remote_control_auth(auth_manager).await?;
@@ -1003,6 +1087,34 @@ mod tests {
}
}
#[test]
fn build_remote_control_websocket_request_includes_fedramp_header() {
let request = build_remote_control_websocket_request(
"ws://127.0.0.1/backend-api/wham/remote/control/server",
&RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_test".to_string(),
server_id: "srv_e_test".to_string(),
server_name: "test-server".to_string(),
},
&RemoteControlConnectionAuth {
authorization_header_value: "AgentAssertion assertion".to_string(),
account_id: "account_id".to_string(),
is_fedramp_account: true,
},
/*subscribe_cursor*/ None,
)
.expect("request should build");
assert_eq!(
request
.headers()
.get(REMOTE_CONTROL_FEDRAMP_HEADER)
.and_then(|value| value.to_str().ok()),
Some("true")
);
}
#[tokio::test]
async fn connect_remote_control_websocket_includes_http_error_details() {
let listener = TcpListener::bind("127.0.0.1:0")
codex-rs/backend-client/src/client.rs
+14 -10
View File
@@ -116,7 +116,7 @@ impl PathStyle {
pub struct Client {
base_url: String,
http: reqwest::Client,
bearer_token: Option<String>,
authorization_header_value: Option<String>,
user_agent: Option<HeaderValue>,
chatgpt_account_id: Option<String>,
chatgpt_account_is_fedramp: bool,
@@ -142,7 +142,7 @@ impl Client {
Ok(Self {
base_url,
http,
bearer_token: None,
authorization_header_value: None,
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -165,7 +165,12 @@ impl Client {
}
pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
self.bearer_token = Some(token.into());
self.authorization_header_value = Some(format!("Bearer {}", token.into()));
self
}
pub fn with_authorization_header_value(mut self, value: impl Into<String>) -> Self {
self.authorization_header_value = Some(value.into());
self
}
@@ -198,11 +203,10 @@ impl Client {
} else {
h.insert(USER_AGENT, HeaderValue::from_static("codex-cli"));
}
if let Some(token) = &self.bearer_token {
let value = format!("Bearer {token}");
if let Ok(hv) = HeaderValue::from_str(&value) {
h.insert(AUTHORIZATION, hv);
}
if let Some(value) = &self.authorization_header_value
&& let Ok(hv) = HeaderValue::from_str(value)
{
h.insert(AUTHORIZATION, hv);
}
if let Some(acc) = &self.chatgpt_account_id
&& let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id")
@@ -816,7 +820,7 @@ mod tests {
let codex_client = Client {
base_url: "https://example.test".to_string(),
http: reqwest::Client::new(),
bearer_token: None,
authorization_header_value: None,
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -830,7 +834,7 @@ mod tests {
let chatgpt_client = Client {
base_url: "https://chatgpt.com/backend-api".to_string(),
http: reqwest::Client::new(),
bearer_token: None,
authorization_header_value: None,
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
codex-rs/chatgpt/src/chatgpt_client.rs
+20 -3
View File
@@ -1,4 +1,5 @@
use codex_core::config::Config;
use codex_login::AuthManager;
use codex_login::default_client::create_client;
use crate::chatgpt_token::get_chatgpt_token_data;
@@ -31,16 +32,32 @@ pub(crate) async fn chatgpt_get_request_with_timeout<T: DeserializeOwned>(
let token =
get_chatgpt_token_data().ok_or_else(|| anyhow::anyhow!("ChatGPT token not available"))?;
let auth_manager =
AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false);
let auth = auth_manager.auth().await;
let is_fedramp_account = auth
.as_ref()
.is_some_and(codex_login::CodexAuth::is_fedramp_account);
let authorization_header_value = match auth.as_ref() {
Some(auth) if auth.is_chatgpt_auth() => auth_manager
.chatgpt_authorization_header_for_auth(auth)
.await
.unwrap_or_else(|| format!("Bearer {}", token.access_token)),
_ => format!("Bearer {}", token.access_token),
};
let account_id = token.account_id.ok_or_else(|| {
anyhow::anyhow!("ChatGPT account ID not available, please re-run `codex login`")
});
})?;
let mut request = client
.get(&url)
.bearer_auth(&token.access_token)
.header("chatgpt-account-id", account_id?)
.header("authorization", authorization_header_value)
.header("chatgpt-account-id", account_id)
.header("Content-Type", "application/json");
if is_fedramp_account {
request = request.header("X-OpenAI-Fedramp", "true");
}
if let Some(timeout) = timeout {
request = request.timeout(timeout);
codex-rs/cloud-requirements/src/lib.rs
+28 -5
View File
@@ -195,11 +195,15 @@ trait RequirementsFetcher: Send + Sync {
struct BackendRequirementsFetcher {
base_url: String,
auth_manager: Arc<AuthManager>,
}
impl BackendRequirementsFetcher {
fn new(base_url: String) -> Self {
Self { base_url }
fn new(auth_manager: Arc<AuthManager>, base_url: String) -> Self {
Self {
base_url,
auth_manager,
}
}
}
@@ -209,7 +213,14 @@ impl RequirementsFetcher for BackendRequirementsFetcher {
&self,
auth: &CodexAuth,
) -> Result<Option<String>, FetchAttemptError> {
let client = BackendClient::from_auth(self.base_url.clone(), auth)
let authorization_header_value = self
.auth_manager
.chatgpt_authorization_header_for_auth(auth)
.await;
let mut client = BackendClient::new(self.base_url.clone())
.map(|client| {
client.with_user_agent(codex_login::default_client::get_codex_user_agent())
})
.inspect_err(|err| {
tracing::warn!(
error = %err,
@@ -217,6 +228,15 @@ impl RequirementsFetcher for BackendRequirementsFetcher {
);
})
.map_err(|_| FetchAttemptError::Retryable(RetryableFailureKind::BackendClientInit))?;
if let Some(authorization_header_value) = authorization_header_value {
client = client.with_authorization_header_value(authorization_header_value);
}
if let Some(account_id) = auth.get_account_id() {
client = client.with_chatgpt_account_id(account_id);
}
if auth.is_fedramp_account() {
client = client.with_fedramp_routing_header();
}
let response = client
.get_config_requirements_file()
@@ -693,8 +713,11 @@ pub fn cloud_requirements_loader(
codex_home: PathBuf,
) -> CloudRequirementsLoader {
let service = CloudRequirementsService::new(
auth_manager,
Arc::new(BackendRequirementsFetcher::new(chatgpt_base_url)),
auth_manager.clone(),
Arc::new(BackendRequirementsFetcher::new(
auth_manager,
chatgpt_base_url,
)),
codex_home,
CLOUD_REQUIREMENTS_TIMEOUT,
);
codex-rs/cloud-tasks-client/src/http.rs
+10
View File
@@ -37,6 +37,11 @@ impl HttpClient {
self
}
pub fn with_authorization_header_value(mut self, value: impl Into<String>) -> Self {
self.backend = self.backend.clone().with_authorization_header_value(value);
self
}
pub fn with_user_agent(mut self, ua: impl Into<String>) -> Self {
self.backend = self.backend.clone().with_user_agent(ua);
self
@@ -47,6 +52,11 @@ impl HttpClient {
self
}
pub fn with_fedramp_routing_header(mut self) -> Self {
self.backend = self.backend.clone().with_fedramp_routing_header();
self
}
fn tasks_api(&self) -> api::Tasks<'_> {
api::Tasks::new(self)
}
codex-rs/cloud-tasks/src/lib.rs
+27 -25
View File
@@ -68,43 +68,45 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result<BackendContext>
};
append_error_log(format!("startup: base_url={base_url} path_style={style}"));
let auth_manager = util::load_auth_manager().await;
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
let Some(auth_manager) = util::load_auth_manager(Some(base_url.clone())).await else {
eprintln!(
"Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'."
);
std::process::exit(1);
};
let auth = match auth {
Some(auth) => auth,
None => {
eprintln!(
"Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'."
);
std::process::exit(1);
}
let Some(auth) = auth_manager.auth().await else {
eprintln!(
"Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'."
);
std::process::exit(1);
};
if let Some(acc) = auth.get_account_id() {
append_error_log(format!("auth: mode=ChatGPT account_id={acc}"));
}
let token = match auth.get_token() {
Ok(t) if !t.is_empty() => t,
_ => {
eprintln!(
"Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'."
);
std::process::exit(1);
}
let authorization_header_value = auth_manager
.chatgpt_authorization_header_for_auth(&auth)
.await;
let Some(authorization_header_value) = authorization_header_value else {
eprintln!(
"Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'."
);
std::process::exit(1);
};
http = http.with_bearer_token(token.clone());
if let Some(acc) = auth
.get_account_id()
.or_else(|| util::extract_chatgpt_account_id(&token))
{
http = http.with_authorization_header_value(authorization_header_value);
if let Some(acc) = auth.get_account_id().or_else(|| {
auth.get_token()
.ok()
.and_then(|token| util::extract_chatgpt_account_id(&token))
}) {
append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}"));
http = http.with_chatgpt_account_id(acc);
}
if auth.is_fedramp_account() {
http = http.with_fedramp_routing_header();
}
Ok(BackendContext {
backend: Arc::new(http),
codex-rs/cloud-tasks/src/util.rs
+30 -18
View File
@@ -3,6 +3,7 @@ use chrono::DateTime;
use chrono::Local;
use chrono::Utc;
use reqwest::header::HeaderMap;
use std::sync::Arc;
use codex_core::config::Config;
use codex_login::AuthManager;
@@ -59,18 +60,18 @@ pub fn extract_chatgpt_account_id(token: &str) -> Option<String> {
.map(str::to_string)
}
pub async fn load_auth_manager() -> Option<AuthManager> {
pub async fn load_auth_manager(chatgpt_base_url: Option<String>) -> Option<Arc<AuthManager>> {
// TODO: pass in cli overrides once cloud tasks properly support them.
let config = Config::load_with_cli_overrides(Vec::new()).await.ok()?;
Some(AuthManager::new(
config.codex_home.to_path_buf(),
/*enable_codex_api_key_env*/ false,
config.cli_auth_credentials_store_mode,
))
let auth_manager =
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false);
if let Some(chatgpt_base_url) = chatgpt_base_url {
auth_manager.set_chatgpt_backend_base_url(Some(chatgpt_base_url));
}
Some(auth_manager)
}
/// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`,
/// and optional `ChatGPT-Account-Id`.
/// Build headers for ChatGPT-backed requests.
pub async fn build_chatgpt_headers() -> HeaderMap {
use reqwest::header::AUTHORIZATION;
use reqwest::header::HeaderName;
@@ -84,23 +85,34 @@ pub async fn build_chatgpt_headers() -> HeaderMap {
USER_AGENT,
HeaderValue::from_str(&ua).unwrap_or(HeaderValue::from_static("codex-cli")),
);
if let Some(am) = load_auth_manager().await
&& let Some(auth) = am.auth().await
&& let Ok(tok) = auth.get_token()
&& !tok.is_empty()
let base_url = normalize_base_url(
&std::env::var("CODEX_CLOUD_TASKS_BASE_URL")
.unwrap_or_else(|_| "https://chatgpt.com/backend-api".to_string()),
);
if let Some(auth_manager) = load_auth_manager(Some(base_url)).await
&& let Some(auth) = auth_manager.auth().await
{
let v = format!("Bearer {tok}");
if let Ok(hv) = HeaderValue::from_str(&v) {
if let Some(authorization_header_value) = auth_manager
.chatgpt_authorization_header_for_auth(&auth)
.await
&& let Ok(hv) = HeaderValue::from_str(&authorization_header_value)
{
headers.insert(AUTHORIZATION, hv);
}
if let Some(acc) = auth
.get_account_id()
.or_else(|| extract_chatgpt_account_id(&tok))
&& let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id")
if let Some(acc) = auth.get_account_id().or_else(|| {
auth.get_token()
.ok()
.and_then(|token| extract_chatgpt_account_id(&token))
}) && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id")
&& let Ok(hv) = HeaderValue::from_str(&acc)
{
headers.insert(name, hv);
}
if auth.is_fedramp_account()
&& let Ok(name) = HeaderName::from_bytes(b"X-OpenAI-Fedramp")
{
headers.insert(name, HeaderValue::from_static("true"));
}
}
headers
}
codex-rs/login/src/auth/manager.rs
+5
View File
@@ -1581,6 +1581,11 @@ impl AuthManager {
}
}
pub fn set_chatgpt_backend_base_url(&self, chatgpt_base_url: Option<String>) {
let (_, auth_mode) = self.chatgpt_backend_auth_config();
self.set_chatgpt_backend_auth_config(chatgpt_base_url, auth_mode);
}
fn chatgpt_backend_auth_config(&self) -> (Option<String>, BackgroundAgentTaskAuthMode) {
let chatgpt_base_url = self
.chatgpt_base_url
codex-rs/models-manager/src/manager.rs
+22 -1
View File
@@ -17,6 +17,7 @@ use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::collect_auth_env_telemetry;
use codex_login::default_client::build_reqwest_client;
use codex_model_provider::AuthorizationHeaderAuthProvider;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo;
@@ -453,7 +454,23 @@ impl ModelsManager {
let auth = self.provider.auth().await;
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
let api_provider = self.provider.api_provider().await?;
let api_auth = self.provider.api_auth().await?;
let mut api_auth = self.provider.api_auth().await?;
if let Some(auth_manager) = auth_manager.as_ref()
&& let Some(auth) = auth.as_ref().filter(|auth| auth.is_chatgpt_auth())
&& provider_uses_codex_login_auth(self.provider.info())
&& let Some(authorization_header_value) = auth_manager
.chatgpt_authorization_header_for_auth(auth)
.await
{
let mut auth_provider = AuthorizationHeaderAuthProvider::new(
Some(authorization_header_value),
auth.get_account_id(),
);
if auth.is_fedramp_account() {
auth_provider = auth_provider.with_fedramp_routing_header();
}
api_auth = Arc::new(auth_provider);
}
let auth_env = collect_auth_env_telemetry(self.provider.info(), codex_api_key_env_enabled);
let transport = ReqwestTransport::new(build_reqwest_client());
let auth_telemetry = auth_header_telemetry(api_auth.as_ref());
@@ -601,6 +618,10 @@ impl ModelsManager {
}
}
fn provider_uses_codex_login_auth(provider: &ModelProviderInfo) -> bool {
provider.env_key.is_none() && provider.experimental_bearer_token.is_none()
}
#[cfg(test)]
#[path = "manager_tests.rs"]
mod tests;