openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: refresh codex auth providers per request

Browse Source
This commit is contained in:
Edward Frazer
2026-04-22 13:50:58 -07:00
parent 27d0647682
commit b0017b5333
25 changed files with 307 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+3
View File
@@ -1829,6 +1829,7 @@ dependencies = [
"chrono",
"clap",
"codex-analytics",
"codex-api",
"codex-app-server-protocol",
"codex-arg0",
"codex-backend-client",
@@ -2829,6 +2830,7 @@ version = "0.0.0"
dependencies = [
"anyhow",
"async-channel",
"codex-api",
"codex-async-utils",
"codex-config",
"codex-exec-server",
@@ -3135,6 +3137,7 @@ version = "0.0.0"
dependencies = [
"anyhow",
"axum",
"codex-api",
"codex-client",
"codex-config",
"codex-exec-server",
codex-rs/app-server/Cargo.toml
+1
View File
@@ -30,6 +30,7 @@ axum = { workspace = true, default-features = false, features = [
"ws",
] }
codex-analytics = { workspace = true }
codex-api = { workspace = true }
codex-arg0 = { workspace = true }
codex-cloud-requirements = { workspace = true }
codex-config = { workspace = true }
codex-rs/app-server/src/transport/remote_control/enroll.rs
+6 -4
View File
@@ -2,6 +2,7 @@ use super::protocol::EnrollRemoteServerRequest;
use super::protocol::EnrollRemoteServerResponse;
use super::protocol::RemoteControlTarget;
use axum::http::HeaderMap;
use codex_api::SharedAuthProvider;
use codex_login::default_client::build_reqwest_client;
use codex_state::RemoteControlEnrollmentRecord;
use codex_state::StateRuntime;
@@ -27,9 +28,8 @@ pub(super) struct RemoteControlEnrollment {
pub(super) server_name: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct RemoteControlConnectionAuth {
pub(super) auth_headers: HeaderMap,
pub(super) auth_provider: SharedAuthProvider,
pub(super) account_id: String,
}
@@ -199,10 +199,12 @@ pub(super) async fn enroll_remote_control_server(
app_server_version: env!("CARGO_PKG_VERSION"),
};
let client = build_reqwest_client();
let mut auth_headers = HeaderMap::new();
auth.auth_provider.add_auth_headers(&mut auth_headers);
let http_request = client
.post(enroll_url)
.timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
.headers(auth.auth_headers.clone())
.headers(auth_headers)
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
.json(&request);
@@ -445,7 +447,7 @@ mod tests {
let err = enroll_remote_control_server(
&remote_control_target,
&RemoteControlConnectionAuth {
auth_headers: HeaderMap::new(),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
account_id: "account_id".to_string(),
},
)
codex-rs/app-server/src/transport/remote_control/websocket.rs
+4 -2
View File
@@ -680,7 +680,9 @@ fn build_remote_control_websocket_request(
"x-codex-protocol-version",
REMOTE_CONTROL_PROTOCOL_VERSION,
)?;
headers.extend(auth.auth_headers.clone());
let mut auth_headers = tungstenite::http::HeaderMap::new();
auth.auth_provider.add_auth_headers(&mut auth_headers);
headers.extend(auth_headers);
set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?;
if let Some(subscribe_cursor) = subscribe_cursor {
set_remote_control_header(
@@ -727,7 +729,7 @@ pub(crate) async fn load_remote_control_auth(
}
Ok(RemoteControlConnectionAuth {
auth_headers: codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers(),
auth_provider: codex_model_provider::auth_provider_from_auth(&auth),
account_id: auth.get_account_id().ok_or_else(|| {
io::Error::new(
ErrorKind::WouldBlock,
codex-rs/backend-client/src/client.rs
+26 -26
View File
@@ -5,7 +5,7 @@ use crate::types::RateLimitReachedKind as BackendRateLimitReachedKind;
use crate::types::RateLimitStatusPayload;
use crate::types::TurnAttemptsSiblingTurnsResponse;
use anyhow::Result;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_client::build_reqwest_client_with_custom_ca;
use codex_client::with_chatgpt_cloudflare_cookie_store;
use codex_login::CodexAuth;
@@ -16,7 +16,6 @@ use codex_protocol::protocol::RateLimitReachedType;
use codex_protocol::protocol::RateLimitSnapshot;
use codex_protocol::protocol::RateLimitWindow;
use reqwest::StatusCode;
use reqwest::header::AUTHORIZATION;
use reqwest::header::CONTENT_TYPE;
use reqwest::header::HeaderMap;
use reqwest::header::HeaderName;
@@ -114,18 +113,33 @@ impl PathStyle {
}
}
#[derive(Clone, Debug)]
#[derive(Clone)]
pub struct Client {
base_url: String,
http: reqwest::Client,
bearer_token: Option<String>,
auth_headers: HeaderMap,
auth_provider: SharedAuthProvider,
user_agent: Option<HeaderValue>,
chatgpt_account_id: Option<String>,
chatgpt_account_is_fedramp: bool,
path_style: PathStyle,
}
impl fmt::Debug for Client {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Client")
.field("base_url", &self.base_url)
.field("auth_provider", &"<provider>")
.field("user_agent", &self.user_agent)
.field("chatgpt_account_id", &self.chatgpt_account_id)
.field(
"chatgpt_account_is_fedramp",
&self.chatgpt_account_is_fedramp,
)
.field("path_style", &self.path_style)
.finish_non_exhaustive()
}
}
impl Client {
pub fn new(base_url: impl Into<String>) -> Result<Self> {
let mut base_url = base_url.into();
@@ -147,8 +161,7 @@ impl Client {
Ok(Self {
base_url,
http,
bearer_token: None,
auth_headers: HeaderMap::new(),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -159,16 +172,11 @@ impl Client {
pub fn from_auth(base_url: impl Into<String>, auth: &CodexAuth) -> Result<Self> {
Ok(Self::new(base_url)?
.with_user_agent(get_codex_user_agent())
.with_auth_provider(codex_model_provider::auth_provider_from_auth(auth).as_ref()))
.with_auth_provider(codex_model_provider::auth_provider_from_auth(auth)))
}
pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self {
auth.add_auth_headers(&mut self.auth_headers);
self
}
pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
self.bearer_token = Some(token.into());
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self {
self.auth_provider = auth;
self
}
@@ -201,13 +209,7 @@ 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);
}
}
h.extend(self.auth_headers.clone());
self.auth_provider.add_auth_headers(&mut h);
if let Some(acc) = &self.chatgpt_account_id
&& let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id")
&& let Ok(hv) = HeaderValue::from_str(acc)
@@ -820,8 +822,7 @@ mod tests {
let codex_client = Client {
base_url: "https://example.test".to_string(),
http: reqwest::Client::new(),
bearer_token: None,
auth_headers: HeaderMap::new(),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -835,8 +836,7 @@ mod tests {
let chatgpt_client = Client {
base_url: "https://chatgpt.com/backend-api".to_string(),
http: reqwest::Client::new(),
bearer_token: None,
auth_headers: HeaderMap::new(),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
codex-rs/cli/src/mcp_cmd.rs
+6 -2
View File
@@ -486,8 +486,12 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
let mut entries: Vec<_> = mcp_servers.iter().collect();
entries.sort_by(|(a, _), (b, _)| a.cmp(b));
let auth_statuses =
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await;
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
/*auth*/ None,
)
.await;
if list_args.json {
let json_entries: Vec<_> = entries
codex-rs/cloud-tasks-client/src/http.rs
+2 -7
View File
@@ -14,7 +14,7 @@ use crate::api::TaskText;
use chrono::DateTime;
use chrono::Utc;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_backend_client as backend;
use codex_backend_client::CodeTaskDetailsResponseExt;
use codex_git_utils::ApplyGitRequest;
@@ -33,17 +33,12 @@ impl HttpClient {
Ok(Self { base_url, backend })
}
pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
self.backend = self.backend.clone().with_bearer_token(token);
self
}
pub fn with_user_agent(mut self, ua: impl Into<String>) -> Self {
self.backend = self.backend.clone().with_user_agent(ua);
self
}
pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self {
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self {
self.backend = self.backend.clone().with_auth_provider(auth);
self
}
codex-rs/cloud-tasks/src/lib.rs
+2 -2
View File
@@ -68,7 +68,7 @@ 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_manager = util::load_auth_manager(Some(base_url.clone())).await;
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
@@ -95,7 +95,7 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result<BackendContext>
}
let auth_provider = codex_model_provider::auth_provider_from_auth(&auth);
http = http.with_auth_provider(auth_provider.as_ref());
http = http.with_auth_provider(auth_provider);
if let Some(acc) = auth.get_account_id() {
append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}"));
}
codex-rs/cloud-tasks/src/util.rs
+3 -3
View File
@@ -41,14 +41,14 @@ pub fn normalize_base_url(input: &str) -> String {
base_url
}
pub async fn load_auth_manager() -> Option<AuthManager> {
pub async fn load_auth_manager(chatgpt_base_url: Option<String>) -> Option<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,
Some(config.chatgpt_base_url),
chatgpt_base_url.or(Some(config.chatgpt_base_url)),
))
}
@@ -65,7 +65,7 @@ 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
if let Some(am) = load_auth_manager(/*chatgpt_base_url*/ None).await
&& let Some(auth) = am.auth().await
&& auth.uses_codex_backend()
{
codex-rs/codex-mcp/Cargo.toml
+1
View File
@@ -15,6 +15,7 @@ workspace = true
anyhow = { workspace = true }
async-channel = { workspace = true }
codex-async-utils = { workspace = true }
codex-api = { workspace = true }
codex-config = { workspace = true }
codex-exec-server = { workspace = true }
codex-login = { workspace = true }
codex-rs/codex-mcp/src/mcp/auth.rs
+27 -9
View File
@@ -1,7 +1,10 @@
use std::collections::HashMap;
use anyhow::Result;
use codex_config::McpServerConfig;
use codex_config::McpServerTransportConfig;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_login::CodexAuth;
use codex_protocol::protocol::McpAuthStatus;
use codex_rmcp_client::OAuthProviderError;
use codex_rmcp_client::determine_streamable_http_auth_status;
@@ -9,8 +12,7 @@ use codex_rmcp_client::discover_streamable_http_oauth;
use futures::future::join_all;
use tracing::warn;
use codex_config::McpServerConfig;
use codex_config::McpServerTransportConfig;
use super::CODEX_APPS_MCP_SERVER_NAME;
#[derive(Debug, Clone)]
pub struct McpOAuthLoginConfig {
@@ -126,6 +128,7 @@ pub struct McpAuthStatusEntry {
pub async fn compute_auth_statuses<'a, I>(
servers: I,
store_mode: OAuthCredentialsStoreMode,
auth: Option<&CodexAuth>,
) -> HashMap<String, McpAuthStatusEntry>
where
I: IntoIterator<Item = (&'a String, &'a McpServerConfig)>,
@@ -133,14 +136,24 @@ where
let futures = servers.into_iter().map(|(name, config)| {
let name = name.clone();
let config = config.clone();
async move {
let auth_status = match compute_auth_status(&name, &config, store_mode).await {
Ok(status) => status,
Err(error) => {
warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
McpAuthStatus::Unsupported
let has_runtime_auth = name == CODEX_APPS_MCP_SERVER_NAME
&& auth.is_some_and(CodexAuth::uses_codex_backend)
&& matches!(
&config.transport,
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var: None,
..
}
};
);
async move {
let auth_status =
match compute_auth_status(&name, &config, store_mode, has_runtime_auth).await {
Ok(status) => status,
Err(error) => {
warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
McpAuthStatus::Unsupported
}
};
let entry = McpAuthStatusEntry {
config,
auth_status,
@@ -156,11 +169,16 @@ async fn compute_auth_status(
server_name: &str,
config: &McpServerConfig,
store_mode: OAuthCredentialsStoreMode,
has_runtime_auth: bool,
) -> Result<McpAuthStatus> {
if !config.enabled {
return Ok(McpAuthStatus::Unsupported);
}
if has_runtime_auth {
return Ok(McpAuthStatus::BearerToken);
}
match &config.transport {
McpServerTransportConfig::Stdio { .. } => Ok(McpAuthStatus::Unsupported),
McpServerTransportConfig::StreamableHttp {
codex-rs/codex-mcp/src/mcp/mod.rs
+25 -32
View File
@@ -205,22 +205,6 @@ fn codex_apps_mcp_bearer_token_env_var() -> Option<String> {
}
}
fn codex_apps_mcp_http_headers(auth: Option<&CodexAuth>) -> Option<HashMap<String, String>> {
let auth = auth.filter(|auth| auth.uses_codex_backend())?;
let headers = codex_model_provider::auth_provider_from_auth(auth).to_auth_headers();
let headers: HashMap<_, _> = headers
.iter()
.filter_map(|(name, value)| {
Some((name.as_str().to_string(), value.to_str().ok()?.to_string()))
})
.collect();
if headers.is_empty() {
None
} else {
Some(headers)
}
}
fn normalize_codex_apps_base_url(base_url: &str) -> String {
let mut base_url = base_url.trim_end_matches('/').to_string();
if (base_url.starts_with("https://chatgpt.com")
@@ -247,20 +231,14 @@ pub(crate) fn codex_apps_mcp_url(config: &McpConfig) -> String {
codex_apps_mcp_url_for_base_url(&config.chatgpt_base_url)
}
fn codex_apps_mcp_server_config(config: &McpConfig, auth: Option<&CodexAuth>) -> McpServerConfig {
let bearer_token_env_var = codex_apps_mcp_bearer_token_env_var();
let http_headers = if bearer_token_env_var.is_some() {
None
} else {
codex_apps_mcp_http_headers(auth)
};
fn codex_apps_mcp_server_config(config: &McpConfig) -> McpServerConfig {
let url = codex_apps_mcp_url(config);
McpServerConfig {
transport: McpServerTransportConfig::StreamableHttp {
url,
bearer_token_env_var,
http_headers,
bearer_token_env_var: codex_apps_mcp_bearer_token_env_var(),
http_headers: None,
env_http_headers: None,
},
experimental_environment: None,
@@ -287,7 +265,7 @@ pub fn with_codex_apps_mcp(
if config.apps_enabled && auth.is_some_and(CodexAuth::uses_codex_backend) {
servers.insert(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
codex_apps_mcp_server_config(config, auth),
codex_apps_mcp_server_config(config),
);
} else {
servers.remove(CODEX_APPS_MCP_SERVER_NAME);
@@ -320,8 +298,12 @@ pub async fn read_mcp_resource(
) -> anyhow::Result<ReadResourceResult> {
let mut mcp_servers = effective_mcp_servers(config, auth);
mcp_servers.retain(|name, _| name == server);
let auth_statuses =
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await;
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
auth,
)
.await;
let (tx_event, rx_event) = unbounded();
drop(rx_event);
let (manager, cancel_token) = McpConnectionManager::new(
@@ -336,6 +318,7 @@ pub async fn read_mcp_resource(
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance(config),
auth,
)
.await;
@@ -386,8 +369,12 @@ pub async fn collect_mcp_snapshot_with_detail(
};
}
let auth_status_entries =
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await;
let auth_status_entries = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
auth,
)
.await;
let (tx_event, rx_event) = unbounded();
drop(rx_event);
@@ -404,6 +391,7 @@ pub async fn collect_mcp_snapshot_with_detail(
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
auth,
)
.await;
@@ -461,8 +449,12 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
};
}
let auth_status_entries =
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await;
let auth_status_entries = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
auth,
)
.await;
let (tx_event, rx_event) = unbounded();
drop(rx_event);
@@ -479,6 +471,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
auth,
)
.await;
codex-rs/codex-mcp/src/mcp_connection_manager.rs
+23
View File
@@ -32,6 +32,7 @@ use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use async_channel::Sender;
use codex_api::SharedAuthProvider;
use codex_async_utils::CancelErr;
use codex_async_utils::OrCancelExt;
use codex_config::Constrained;
@@ -485,6 +486,7 @@ impl AsyncManagedClient {
codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
tool_plugin_provenance: Arc<ToolPluginProvenance>,
runtime_environment: McpRuntimeEnvironment,
runtime_auth_provider: Option<SharedAuthProvider>,
) -> Self {
let tool_filter = ToolFilter::from_config(&config);
let startup_snapshot = load_startup_cached_codex_apps_tools_snapshot(
@@ -507,6 +509,7 @@ impl AsyncManagedClient {
config.clone(),
store_mode,
runtime_environment,
runtime_auth_provider,
)
.await?,
);
@@ -744,6 +747,7 @@ impl McpConnectionManager {
codex_home: PathBuf,
codex_apps_tools_cache_key: CodexAppsToolsCacheKey,
tool_plugin_provenance: ToolPluginProvenance,
auth: Option<&CodexAuth>,
) -> (Self, CancellationToken) {
let cancel_token = CancellationToken::new();
let mut clients = HashMap::new();
@@ -753,6 +757,9 @@ impl McpConnectionManager {
ElicitationRequestManager::new(approval_policy.value(), initial_sandbox_policy);
let tool_plugin_provenance = Arc::new(tool_plugin_provenance);
let startup_submit_id = submit_id.clone();
let codex_apps_auth_provider = auth
.filter(|auth| auth.uses_codex_backend())
.map(codex_model_provider::auth_provider_from_auth);
let mcp_servers = mcp_servers.clone();
for (server_name, cfg) in mcp_servers.into_iter().filter(|(_, cfg)| cfg.enabled) {
if let Some(origin) = transport_origin(&cfg.transport) {
@@ -776,6 +783,19 @@ impl McpConnectionManager {
} else {
None
};
let uses_env_bearer_token = match &cfg.transport {
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var,
..
} => bearer_token_env_var.is_some(),
McpServerTransportConfig::Stdio { .. } => false,
};
let runtime_auth_provider =
if server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token {
codex_apps_auth_provider.clone()
} else {
None
};
let async_managed_client = AsyncManagedClient::new(
server_name.clone(),
cfg,
@@ -786,6 +806,7 @@ impl McpConnectionManager {
codex_apps_tools_cache_context,
Arc::clone(&tool_plugin_provenance),
runtime_environment.clone(),
runtime_auth_provider,
);
clients.insert(server_name.clone(), async_managed_client.clone());
let tx_event = tx_event.clone();
@@ -1519,6 +1540,7 @@ async fn make_rmcp_client(
config: McpServerConfig,
store_mode: OAuthCredentialsStoreMode,
runtime_environment: McpRuntimeEnvironment,
runtime_auth_provider: Option<SharedAuthProvider>,
) -> Result<RmcpClient, StartupOutcomeError> {
let McpServerConfig {
transport,
@@ -1607,6 +1629,7 @@ async fn make_rmcp_client(
http_headers,
env_http_headers,
store_mode,
runtime_auth_provider,
)
.await
.map_err(StartupOutcomeError::from)
codex-rs/core/src/connectors.rs
+15 -9
View File
@@ -8,6 +8,7 @@ use std::time::Instant;
use anyhow::Context;
use async_channel::unbounded;
use codex_api::SharedAuthProvider;
pub use codex_app_server_protocol::AppBranding;
pub use codex_app_server_protocol::AppInfo;
pub use codex_app_server_protocol::AppMetadata;
@@ -252,8 +253,12 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
});
}
let auth_status_entries =
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await;
let auth_status_entries = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
auth.as_ref(),
)
.await;
let (tx_event, rx_event) = unbounded();
drop(rx_event);
@@ -274,6 +279,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
config.codex_home.to_path_buf(),
codex_apps_tools_cache_key(auth.as_ref()),
ToolPluginProvenance::default(),
auth.as_ref(),
)
.await;
@@ -440,7 +446,7 @@ async fn list_directory_connectors_for_tool_suggest_with_auth(
Some(account_id) if !account_id.is_empty() => account_id,
_ => return Ok(Vec::new()),
};
let auth_headers = codex_model_provider::auth_provider_from_auth(auth).to_auth_headers();
let auth_provider = codex_model_provider::auth_provider_from_auth(auth);
let is_workspace_account = auth.is_workspace_account();
let cache_key = AllConnectorsCacheKey::new(
config.chatgpt_base_url.clone(),
@@ -454,12 +460,12 @@ async fn list_directory_connectors_for_tool_suggest_with_auth(
is_workspace_account,
/*force_refetch*/ false,
|path| {
let auth_headers = auth_headers.clone();
let auth_provider = auth_provider.clone();
async move {
chatgpt_get_request_with_auth_headers::<DirectoryListResponse>(
chatgpt_get_request_with_auth_provider::<DirectoryListResponse>(
config,
path,
auth_headers,
auth_provider,
)
.await
}
@@ -468,16 +474,16 @@ async fn list_directory_connectors_for_tool_suggest_with_auth(
.await
}
async fn chatgpt_get_request_with_auth_headers<T: DeserializeOwned>(
async fn chatgpt_get_request_with_auth_provider<T: DeserializeOwned>(
config: &Config,
path: String,
auth_headers: http::HeaderMap,
auth_provider: SharedAuthProvider,
) -> anyhow::Result<T> {
let client = create_client();
let url = format!("{}{}", config.chatgpt_base_url, path);
let response = client
.get(&url)
.headers(auth_headers)
.headers(auth_provider.to_auth_headers())
.header("Content-Type", "application/json")
.timeout(DIRECTORY_CONNECTORS_TIMEOUT)
.send()
codex-rs/core/src/session/handlers.rs
+6 -1
View File
@@ -489,7 +489,12 @@ pub async fn list_mcp_tools(sess: &Session, config: &Arc<Config>, sub_id: String
.await;
let snapshot = collect_mcp_snapshot_from_manager(
&mcp_connection_manager,
compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await,
compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
auth.as_ref(),
)
.await,
)
.await;
let event = Event {
codex-rs/core/src/session/mcp.rs
+3 -1
View File
@@ -219,7 +219,8 @@ impl Session {
.tool_plugin_provenance(config.as_ref())
.await;
let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config);
let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode).await;
let auth_statuses =
compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await;
{
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
guard.cancel();
@@ -243,6 +244,7 @@ impl Session {
config.codex_home.to_path_buf(),
codex_apps_tools_cache_key(auth.as_ref()),
tool_plugin_provenance,
auth.as_ref(),
)
.await;
{
codex-rs/core/src/session/session.rs
+2
View File
@@ -350,6 +350,7 @@ impl Session {
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
config_for_mcp.mcp_oauth_credentials_store_mode,
auth.as_ref(),
)
.await;
(auth, mcp_servers, auth_statuses)
@@ -787,6 +788,7 @@ impl Session {
config.codex_home.to_path_buf(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
auth,
)
.instrument(info_span!(
"session_init.mcp_manager_init",
codex-rs/model-provider/src/auth.rs
+13 -21
View File
@@ -55,31 +55,19 @@ impl AuthProvider for AgentIdentityAuthProvider {
}
}
#[derive(Clone, Debug)]
struct CodexBearerAuthProvider {
auth: CodexAuth,
}
impl AuthProvider for CodexBearerAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
BearerAuthProvider {
token: self.auth.get_token().ok(),
account_id: self.auth.get_account_id(),
is_fedramp_account: self.auth.is_fedramp_account(),
}
.add_auth_headers(headers);
}
}
// Some providers are meant to send no auth headers. Examples include local OSS
// providers and custom test providers with `requires_openai_auth = false`.
#[derive(Clone, Debug)]
struct NoAuthProvider;
struct UnauthenticatedAuthProvider;
impl AuthProvider for NoAuthProvider {
impl AuthProvider for UnauthenticatedAuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
}
pub fn unauthenticated_auth_provider() -> SharedAuthProvider {
Arc::new(UnauthenticatedAuthProvider)
}
/// Returns the provider-scoped auth manager when this provider uses command-backed auth.
///
/// Providers without custom auth continue using the caller-supplied base manager, when present.
@@ -103,7 +91,7 @@ pub(crate) fn resolve_provider_auth(
Ok(match auth {
Some(auth) => auth_provider_from_auth(auth),
None => Arc::new(NoAuthProvider),
None => unauthenticated_auth_provider(),
})
}
@@ -128,7 +116,11 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider {
Arc::new(AgentIdentityAuthProvider { auth: auth.clone() })
}
CodexAuth::ApiKey(_) | CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => {
Arc::new(CodexBearerAuthProvider { auth: auth.clone() })
Arc::new(BearerAuthProvider {
token: auth.get_token().ok(),
account_id: auth.get_account_id(),
is_fedramp_account: auth.is_fedramp_account(),
})
}
}
}
@@ -141,7 +133,7 @@ mod tests {
use super::*;
#[test]
fn no_auth_provider_adds_no_headers() {
fn unauthenticated_auth_provider_adds_no_headers() {
let provider =
create_oss_provider_with_base_url("http://localhost:11434/v1", WireApi::Responses);
let auth = resolve_provider_auth(/*auth*/ None, &provider).expect("auth should resolve");
codex-rs/model-provider/src/lib.rs
+1
View File
@@ -4,6 +4,7 @@ mod bearer_auth_provider;
mod provider;
pub use auth::auth_provider_from_auth;
pub use auth::unauthenticated_auth_provider;
pub use bearer_auth_provider::BearerAuthProvider;
pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider;
pub use provider::ModelProvider;
codex-rs/protocol/src/account.rs
+12 -3
View File
@@ -37,9 +37,15 @@ impl PlanType {
}
pub fn is_workspace_account(self) -> bool {
self.is_team_like()
|| self.is_business_like()
|| matches!(self, Self::Enterprise | Self::Edu)
matches!(
self,
Self::Team
| Self::SelfServeBusinessUsageBased
| Self::Business
| Self::EnterpriseCbpUsageBased
| Self::Enterprise
| Self::Edu
)
}
}
@@ -89,7 +95,10 @@ mod tests {
assert_eq!(PlanType::Business.is_business_like(), true);
assert_eq!(PlanType::EnterpriseCbpUsageBased.is_business_like(), true);
assert_eq!(PlanType::Team.is_business_like(), false);
}
#[test]
fn workspace_account_helper_includes_usage_based_workspace_plans() {
assert_eq!(PlanType::Team.is_workspace_account(), true);
assert_eq!(
PlanType::SelfServeBusinessUsageBased.is_workspace_account(),
codex-rs/protocol/src/auth.rs
+9 -11
View File
@@ -50,14 +50,6 @@ pub enum KnownPlan {
}
impl KnownPlan {
pub fn is_team_like(self) -> bool {
matches!(self, Self::Team | Self::SelfServeBusinessUsageBased)
}
pub fn is_business_like(self) -> bool {
matches!(self, Self::Business | Self::EnterpriseCbpUsageBased)
}
pub fn display_name(self) -> &'static str {
match self {
Self::Free => "Free",
@@ -91,9 +83,15 @@ impl KnownPlan {
}
pub fn is_workspace_account(self) -> bool {
self.is_team_like()
|| self.is_business_like()
|| matches!(self, Self::Enterprise | Self::Edu)
matches!(
self,
Self::Team
| Self::SelfServeBusinessUsageBased
| Self::Business
| Self::EnterpriseCbpUsageBased
| Self::Enterprise
| Self::Edu
)
}
}
codex-rs/protocol/src/error.rs
+28 -32
View File
@@ -472,44 +472,40 @@ impl std::fmt::Display for UsageLimitReachedError {
}
let message = match self.plan_type.as_ref() {
Some(PlanType::Known(plan)) => match plan {
KnownPlan::Plus => format!(
"You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits{}",
retry_suffix_after_or(self.resets_at.as_ref())
),
plan if plan.is_team_like() || plan.is_business_like() => {
format!(
"You've hit your usage limit. To get more access now, send a request to your admin{}",
retry_suffix_after_or(self.resets_at.as_ref())
)
}
KnownPlan::Free | KnownPlan::Go => {
format!(
"You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}",
retry_suffix_after_or(self.resets_at.as_ref())
)
}
KnownPlan::Pro | KnownPlan::ProLite => format!(
"You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}",
retry_suffix_after_or(self.resets_at.as_ref())
),
KnownPlan::Enterprise | KnownPlan::Edu => format!(
"You've hit your usage limit.{}",
retry_suffix(self.resets_at.as_ref())
),
Some(PlanType::Known(KnownPlan::Plus)) => format!(
"You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits{}",
retry_suffix_after_or(self.resets_at.as_ref())
),
Some(PlanType::Known(
KnownPlan::Team
| KnownPlan::SelfServeBusinessUsageBased
| KnownPlan::Business
| KnownPlan::EnterpriseCbpUsageBased => {
unreachable!("team-like and business-like plans are handled above")
}
},
Some(PlanType::Unknown(_)) | None => {
| KnownPlan::EnterpriseCbpUsageBased,
)) => {
format!(
"You've hit your usage limit.{}",
retry_suffix(self.resets_at.as_ref())
"You've hit your usage limit. To get more access now, send a request to your admin{}",
retry_suffix_after_or(self.resets_at.as_ref())
)
}
Some(PlanType::Known(KnownPlan::Free)) | Some(PlanType::Known(KnownPlan::Go)) => {
format!(
"You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}",
retry_suffix_after_or(self.resets_at.as_ref())
)
}
Some(PlanType::Known(KnownPlan::Pro | KnownPlan::ProLite)) => format!(
"You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}",
retry_suffix_after_or(self.resets_at.as_ref())
),
Some(PlanType::Known(KnownPlan::Enterprise))
| Some(PlanType::Known(KnownPlan::Edu)) => format!(
"You've hit your usage limit.{}",
retry_suffix(self.resets_at.as_ref())
),
Some(PlanType::Unknown(_)) | None => format!(
"You've hit your usage limit.{}",
retry_suffix(self.resets_at.as_ref())
),
};
write!(f, "{message}")
codex-rs/rmcp-client/Cargo.toml
+1
View File
@@ -13,6 +13,7 @@ axum = { workspace = true, default-features = false, features = [
"http1",
"tokio",
] }
codex-api = { workspace = true }
codex-client = { workspace = true }
codex-config = { workspace = true }
codex-exec-server = { workspace = true }
codex-rs/rmcp-client/src/rmcp_client.rs
+87 -16
View File
@@ -12,6 +12,7 @@ use std::time::Instant;
use anyhow::Result;
use anyhow::anyhow;
use codex_api::SharedAuthProvider;
use codex_client::build_reqwest_client_with_custom_ca;
use codex_config::types::McpServerEnvVar;
use futures::FutureExt;
@@ -88,11 +89,15 @@ const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192;
#[derive(Clone)]
struct StreamableHttpResponseClient {
inner: reqwest::Client,
auth_provider: Option<SharedAuthProvider>,
}
impl StreamableHttpResponseClient {
fn new(inner: reqwest::Client) -> Self {
Self { inner }
fn new(inner: reqwest::Client, auth_provider: Option<SharedAuthProvider>) -> Self {
Self {
inner,
auth_provider,
}
}
fn reqwest_error(
@@ -129,6 +134,9 @@ impl StreamableHttpClient for StreamableHttpResponseClient {
.inner
.post(uri.as_ref())
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
if let Some(auth_provider) = &self.auth_provider {
request = request.headers(auth_provider.to_auth_headers());
}
if let Some(auth_header) = auth_token {
request = request.bearer_auth(auth_header);
}
@@ -339,6 +347,7 @@ enum TransportRecipe {
http_headers: Option<HashMap<String, String>>,
env_http_headers: Option<HashMap<String, String>>,
store_mode: OAuthCredentialsStoreMode,
auth_provider: Option<SharedAuthProvider>,
},
}
@@ -536,6 +545,7 @@ impl RmcpClient {
http_headers: Option<HashMap<String, String>>,
env_http_headers: Option<HashMap<String, String>>,
store_mode: OAuthCredentialsStoreMode,
auth_provider: Option<SharedAuthProvider>,
) -> Result<Self> {
let transport_recipe = TransportRecipe::StreamableHttp {
server_name: server_name.to_string(),
@@ -544,6 +554,7 @@ impl RmcpClient {
http_headers,
env_http_headers,
store_mode,
auth_provider,
};
let transport = Self::create_pending_transport(&transport_recipe).await?;
Ok(Self {
@@ -895,22 +906,25 @@ impl RmcpClient {
http_headers,
env_http_headers,
store_mode,
auth_provider,
} => {
let default_headers =
build_default_headers(http_headers.clone(), env_http_headers.clone())?;
let initial_oauth_tokens =
if bearer_token.is_none() && !default_headers.contains_key(AUTHORIZATION) {
match load_oauth_tokens(server_name, url, *store_mode) {
Ok(tokens) => tokens,
Err(err) => {
warn!("failed to read tokens for server `{server_name}`: {err}");
None
}
let initial_oauth_tokens = if bearer_token.is_none()
&& auth_provider.is_none()
&& !default_headers.contains_key(AUTHORIZATION)
{
match load_oauth_tokens(server_name, url, *store_mode) {
Ok(tokens) => tokens,
Err(err) => {
warn!("failed to read tokens for server `{server_name}`: {err}");
None
}
} else {
None
};
}
} else {
None
};
if let Some(initial_tokens) = initial_oauth_tokens.clone() {
match create_oauth_transport_and_runtime(
@@ -947,7 +961,7 @@ impl RmcpClient {
.auth_header(access_token);
let http_client = build_http_client(&default_headers)?;
let transport = StreamableHttpClientTransport::with_client(
StreamableHttpResponseClient::new(http_client),
StreamableHttpResponseClient::new(http_client, None),
http_config,
);
Ok(PendingTransport::StreamableHttp { transport })
@@ -964,7 +978,7 @@ impl RmcpClient {
let http_client = build_http_client(&default_headers)?;
let transport = StreamableHttpClientTransport::with_client(
StreamableHttpResponseClient::new(http_client),
StreamableHttpResponseClient::new(http_client, auth_provider.clone()),
http_config,
);
Ok(PendingTransport::StreamableHttp { transport })
@@ -1178,7 +1192,10 @@ async fn create_oauth_transport_and_runtime(
}
};
let auth_client = AuthClient::new(StreamableHttpResponseClient::new(http_client), manager);
let auth_client = AuthClient::new(
StreamableHttpResponseClient::new(http_client, None),
manager,
);
let auth_manager = auth_client.auth_manager.clone();
let transport = StreamableHttpClientTransport::with_client(
@@ -1199,13 +1216,67 @@ async fn create_oauth_transport_and_runtime(
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_api::AuthProvider;
use pretty_assertions::assert_eq;
use reqwest::header::HeaderMap;
use rmcp::model::ClientJsonRpcMessage;
use rmcp::model::ClientNotification;
use rmcp::model::InitializedNotification;
use tokio::time;
use super::*;
#[derive(Debug)]
struct CountingAuthProvider {
calls: Arc<AtomicUsize>,
}
impl AuthProvider for CountingAuthProvider {
fn add_auth_headers(&self, _headers: &mut HeaderMap) {
self.calls.fetch_add(1, Ordering::SeqCst);
}
}
#[tokio::test]
async fn streamable_http_client_applies_runtime_auth_per_request() {
let calls = Arc::new(AtomicUsize::new(0));
let client = StreamableHttpResponseClient::new(
reqwest::Client::new(),
Some(Arc::new(CountingAuthProvider {
calls: Arc::clone(&calls),
})),
);
let message = || {
ClientJsonRpcMessage::notification(ClientNotification::InitializedNotification(
InitializedNotification::default(),
))
};
let _ = client
.post_message(
Arc::from("http://127.0.0.1:1/mcp"),
message(),
/*session_id*/ None,
/*auth_token*/ None,
)
.await;
let _ = client
.post_message(
Arc::from("http://127.0.0.1:1/mcp"),
message(),
/*session_id*/ None,
/*auth_token*/ None,
)
.await;
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn active_time_timeout_pauses_while_elicitation_is_pending() {
let pause_state = ElicitationPauseState::new();
codex-rs/rmcp-client/tests/streamable_http_recovery.rs
+1
View File
@@ -77,6 +77,7 @@ async fn create_client(base_url: &str) -> anyhow::Result<RmcpClient> {
/*http_headers*/ None,
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::File,
/*auth_provider*/ None,
)
.await?;