openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor app server auth store injection

Browse Source
This commit is contained in:
starr-openai
2026-06-03 01:05:35 -07:00
parent 521829053d
commit cb357bdb6c
6 changed files with 192 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
codex-rs/app-server/src/in_process.rs
View File

@@ -85,6 +85,7 @@ use codex_core::resolve_installation_id;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_login::AuthManager;
use codex_login::AuthStores;
use codex_protocol::protocol::SessionSource;
pub use codex_rollout::StateDbHandle;
pub use codex_state::log_db::LogDbLayer;
@@ -348,8 +349,23 @@ impl InProcessClientHandle {
/// the handle, so callers receive a ready-to-use runtime. If initialize fails,
/// the runtime is shut down and an `InvalidData` error is returned.
pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
start_with_optional_auth_stores(args, /*auth_stores*/ None).await
}
/// Starts an in-process runtime backed by caller-provided auth stores.
pub async fn start_with_auth_stores(
args: InProcessStartArgs,
auth_stores: AuthStores,
) -> IoResult<InProcessClientHandle> {
start_with_optional_auth_stores(args, Some(auth_stores)).await
}
async fn start_with_optional_auth_stores(
args: InProcessStartArgs,
auth_stores: Option<AuthStores>,
) -> IoResult<InProcessClientHandle> {
let initialize = args.initialize.clone();
let client = start_uninitialized(args).await?;
let client = start_uninitialized(args, auth_stores).await?;
let initialize_response = client
.request(ClientRequest::Initialize {
@@ -369,7 +385,10 @@ pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle>
Ok(client)
}
async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
async fn start_uninitialized(
args: InProcessStartArgs,
auth_stores: Option<AuthStores>,
) -> IoResult<InProcessClientHandle> {
let channel_capacity = args.channel_capacity.max(1);
let installation_id = resolve_installation_id(&args.config.codex_home).await?;
let (client_tx, mut client_rx) = mpsc::channel::<InProcessClientMessage>(channel_capacity);
@@ -377,9 +396,22 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
let runtime_handle = tokio::spawn(async move {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
let auth_manager =
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env)
.await;
let auth_manager = match auth_stores {
Some(auth_stores) => {
AuthManager::shared_with_stores(
auth_stores,
args.enable_codex_api_key_env,
Some(args.config.chatgpt_base_url.clone()),
)
.await
}
None => {
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env)
.await
}
};
auth_manager
.set_forced_chatgpt_workspace_id(args.config.forced_chatgpt_workspace_id.clone());
let analytics_events_client =
analytics_events_client_from_config(Arc::clone(&auth_manager), args.config.as_ref());
let outgoing_message_sender = Arc::new(OutgoingMessageSender::new(

2
codex-rs/app-server/src/request_processors.rs
View File

@@ -341,9 +341,7 @@ use codex_login::CLIENT_ID;
use codex_login::CodexAuth;
use codex_login::ServerOptions as LoginServerOptions;
use codex_login::ShutdownHandle;
use codex_login::auth::login_with_chatgpt_auth_tokens;
use codex_login::complete_device_code_login;
use codex_login::login_with_api_key;
use codex_login::request_device_code;
use codex_login::run_login_server;
use codex_mcp::McpRuntimeContext;

21
codex-rs/app-server/src/request_processors/account_processor.rs
View File

@@ -273,11 +273,7 @@ impl AccountRequestProcessor {
}
}
match login_with_api_key(
&self.config.codex_home,
&params.api_key,
self.config.cli_auth_credentials_store_mode,
) {
match self.auth_manager.login_with_api_key(&params.api_key) {
Ok(()) => {
self.auth_manager.reload().await;
Ok(())
@@ -326,6 +322,7 @@ impl AccountRequestProcessor {
config.forced_chatgpt_workspace_id.clone(),
config.cli_auth_credentials_store_mode,
)
.with_configured_auth_store(self.auth_manager.configured_auth_store())
};
#[cfg(debug_assertions)]
let opts = {
@@ -578,13 +575,13 @@ impl AccountRequestProcessor {
)));
}
login_with_chatgpt_auth_tokens(
&self.config.codex_home,
&access_token,
&chatgpt_account_id,
chatgpt_plan_type.as_deref(),
)
.map_err(|err| internal_error(format!("failed to set external auth: {err}")))?;
self.auth_manager
.login_with_chatgpt_auth_tokens(
&access_token,
&chatgpt_account_id,
chatgpt_plan_type.as_deref(),
)
.map_err(|err| internal_error(format!("failed to set external auth: {err}")))?;
self.auth_manager.reload().await;
self.config_manager.replace_cloud_config_bundle_loader(
self.auth_manager.clone(),

5
codex-rs/login/src/auth/manager.rs
View File

@@ -1628,6 +1628,11 @@ impl AuthManager {
Arc::new(Self::new_with_stores(stores, enable_codex_api_key_env, chatgpt_base_url).await)
}
/// Returns the configured credential store for managed auth writers.
pub fn configured_auth_store(&self) -> Arc<dyn AuthCredentialStore> {
Arc::clone(&self.stores.configured)
}
/// Convenience constructor returning an `Arc` wrapper from resolved config.
pub async fn shared_from_config(
config: &impl AuthManagerConfig,

5
codex-rs/login/src/device_code_auth.rs
View File

@@ -210,13 +210,12 @@ pub async fn complete_device_code_login(
return Err(io::Error::new(io::ErrorKind::PermissionDenied, message));
}
crate::server::persist_tokens_async(
&opts.codex_home,
crate::server::persist_tokens_to_store_async(
opts.configured_auth_store(),
/*api_key*/ None,
tokens.id_token,
tokens.access_token,
tokens.refresh_token,
opts.cli_auth_credentials_store_mode,
)
.await
}

149
codex-rs/login/src/server.rs
View File

@@ -11,6 +11,7 @@
//! This module therefore keeps the user-facing error path and the structured-log path separate.
//! Returned `io::Error` values still carry the detail needed by CLI/browser callers, while
//! structured logs only emit explicitly reviewed fields plus redacted URL/error values.
use std::fmt::Debug;
use std::io::Cursor;
use std::io::Read;
use std::io::Write;
@@ -24,10 +25,10 @@ use std::sync::LazyLock;
use std::thread;
use std::time::Duration;
use crate::auth::AuthCredentialStore;
use crate::auth::AuthDotJson;
use crate::auth::load_auth_dot_json;
use crate::auth::AuthStores;
use crate::auth::revoke_auth_tokens;
use crate::auth::save_auth;
use crate::auth::should_revoke_auth_tokens;
use crate::default_client::originator;
use crate::pkce::PkceCodes;
@@ -61,7 +62,7 @@ static LOGIN_ERROR_PAGE_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
});
/// Options for launching the local login callback server.
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct ServerOptions {
pub codex_home: PathBuf,
pub client_id: String,
@@ -72,6 +73,7 @@ pub struct ServerOptions {
pub forced_chatgpt_workspace_id: Option<Vec<String>>,
pub codex_streamlined_login: bool,
pub cli_auth_credentials_store_mode: AuthCredentialsStoreMode,
configured_auth_store: Option<Arc<dyn AuthCredentialStore>>,
}
impl ServerOptions {
@@ -92,8 +94,54 @@ impl ServerOptions {
forced_chatgpt_workspace_id,
codex_streamlined_login: false,
cli_auth_credentials_store_mode,
configured_auth_store: None,
}
}
/// Overrides first-party ChatGPT credential persistence for this login flow.
pub fn with_configured_auth_store(
mut self,
configured_auth_store: Arc<dyn AuthCredentialStore>,
) -> Self {
self.configured_auth_store = Some(configured_auth_store);
self
}
pub(crate) fn configured_auth_store(&self) -> Arc<dyn AuthCredentialStore> {
self.configured_auth_store.clone().unwrap_or_else(|| {
AuthStores::local(
self.codex_home.clone(),
self.cli_auth_credentials_store_mode,
)
.configured
})
}
}
impl Debug for ServerOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ServerOptions")
.field("codex_home", &self.codex_home)
.field("client_id", &self.client_id)
.field("issuer", &self.issuer)
.field("port", &self.port)
.field("open_browser", &self.open_browser)
.field("force_state", &self.force_state)
.field(
"forced_chatgpt_workspace_id",
&self.forced_chatgpt_workspace_id,
)
.field("codex_streamlined_login", &self.codex_streamlined_login)
.field(
"cli_auth_credentials_store_mode",
&self.cli_auth_credentials_store_mode,
)
.field(
"has_configured_auth_store",
&self.configured_auth_store.is_some(),
)
.finish()
}
}
/// Handle for a running login callback server.
@@ -355,13 +403,12 @@ async fn process_request(
let api_key = obtain_api_key(&opts.issuer, &opts.client_id, &tokens.id_token)
.await
.ok();
if let Err(err) = persist_tokens_async(
&opts.codex_home,
if let Err(err) = persist_tokens_to_store_async(
opts.configured_auth_store(),
api_key.clone(),
tokens.id_token.clone(),
tokens.access_token.clone(),
tokens.refresh_token.clone(),
opts.cli_auth_credentials_store_mode,
)
.await
{
@@ -784,7 +831,7 @@ pub(crate) async fn exchange_code_for_tokens(
})
}
/// Persists exchanged credentials using the configured local auth store, then
/// Persists exchanged credentials using the configured auth store, then
/// best-effort revokes any superseded managed ChatGPT tokens.
pub(crate) async fn persist_tokens_async(
codex_home: &Path,
@@ -793,11 +840,27 @@ pub(crate) async fn persist_tokens_async(
access_token: String,
refresh_token: String,
auth_credentials_store_mode: AuthCredentialsStoreMode,
) -> io::Result<()> {
persist_tokens_to_store_async(
AuthStores::local(codex_home.to_path_buf(), auth_credentials_store_mode).configured,
api_key,
id_token,
access_token,
refresh_token,
)
.await
}
pub(crate) async fn persist_tokens_to_store_async(
configured_auth_store: Arc<dyn AuthCredentialStore>,
api_key: Option<String>,
id_token: String,
access_token: String,
refresh_token: String,
) -> io::Result<()> {
// Reuse existing synchronous logic but run it off the async runtime.
let codex_home = codex_home.to_path_buf();
let (previous_auth, auth) = tokio::task::spawn_blocking(move || {
let previous_auth = match load_auth_dot_json(&codex_home, auth_credentials_store_mode) {
let previous_auth = match configured_auth_store.load() {
Ok(auth) => auth,
Err(err) => {
warn!("failed to load previous auth before saving new login: {err}");
@@ -823,7 +886,7 @@ pub(crate) async fn persist_tokens_async(
last_refresh: Some(Utc::now()),
agent_identity: None,
};
save_auth(&codex_home, &auth, auth_credentials_store_mode)?;
configured_auth_store.save(&auth)?;
Ok::<_, io::Error>((previous_auth, auth))
})
.await
@@ -1156,6 +1219,8 @@ pub(crate) async fn obtain_api_key(
#[cfg(test)]
mod tests {
use std::ffi::OsString;
use std::sync::Arc;
use std::sync::Mutex;
use anyhow::Context;
use base64::Engine;
@@ -1170,6 +1235,7 @@ mod tests {
use wiremock::matchers::method;
use wiremock::matchers::path;
use crate::auth::AuthCredentialStore;
use crate::auth::AuthDotJson;
use crate::auth::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR;
use crate::auth::load_auth_dot_json;
@@ -1186,11 +1252,74 @@ mod tests {
use super::is_missing_codex_entitlement_error;
use super::parse_token_endpoint_error;
use super::persist_tokens_async;
use super::persist_tokens_to_store_async;
use super::redact_sensitive_query_value;
use super::redact_sensitive_url_parts;
use super::render_login_error_page;
use super::sanitize_url_for_logging;
#[derive(Debug, Default)]
struct FakeAuthCredentialStore {
auth: Mutex<Option<AuthDotJson>>,
}
impl AuthCredentialStore for FakeAuthCredentialStore {
fn load(&self) -> std::io::Result<Option<AuthDotJson>> {
Ok(self.auth.lock().expect("fake store lock").clone())
}
fn save(&self, auth: &AuthDotJson) -> std::io::Result<()> {
*self.auth.lock().expect("fake store lock") = Some(auth.clone());
Ok(())
}
fn delete(&self) -> std::io::Result<bool> {
Ok(self.auth.lock().expect("fake store lock").take().is_some())
}
}
#[tokio::test]
async fn configured_auth_store_overrides_local_login_persistence() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let configured = Arc::new(FakeAuthCredentialStore::default());
let opts = super::ServerOptions::new(
codex_home.path().to_path_buf(),
"client-id".to_string(),
/*forced_chatgpt_workspace_id*/ None,
AuthCredentialsStoreMode::File,
)
.with_configured_auth_store(configured.clone());
persist_tokens_to_store_async(
opts.configured_auth_store(),
/*api_key*/ None,
jwt_for_account("new-account"),
"new-access".to_string(),
"new-refresh".to_string(),
)
.await?;
assert_eq!(
load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?,
None
);
assert_eq!(
configured
.load()?
.context("configured auth should exist after login")?
.tokens
.context("configured tokens should be persisted")?,
TokenData {
id_token: parse_chatgpt_jwt_claims(&jwt_for_account("new-account"))
.expect("new JWT should parse"),
access_token: "new-access".to_string(),
refresh_token: "new-refresh".to_string(),
account_id: Some("new-account".to_string()),
}
);
Ok(())
}
#[serial_test::serial(logout_revoke)]
#[tokio::test]
async fn persist_tokens_async_revokes_previous_auth_without_failing_login() -> anyhow::Result<()>