[codex-login] reset per-surface HTTP state on auth lifecycle [ci changed_files]

This commit is contained in:
Cooper Gamble
2026-06-03 09:54:36 +00:00
parent b6fc0b7788
commit 236fc646dc
13 changed files with 360 additions and 56 deletions
+1
View File
@@ -3130,6 +3130,7 @@ dependencies = [
"codex-app-server-protocol",
"codex-client",
"codex-config",
"codex-http-state",
"codex-keyring-store",
"codex-model-provider-info",
"codex-otel",
+6 -2
View File
@@ -1281,12 +1281,16 @@ impl MessageProcessor {
}
ClientRequest::LoginAccount { params, .. } => {
self.account_processor
.login_account(request_id.clone(), params)
.login_account(
request_id.clone(),
params,
app_server_client_name.as_deref(),
)
.await
}
ClientRequest::LogoutAccount { .. } => {
self.account_processor
.logout_account(request_id.clone())
.logout_account(request_id.clone(), app_server_client_name.as_deref())
.await
}
ClientRequest::CancelLoginAccount { params, .. } => {
@@ -341,9 +341,9 @@ 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::auth::login_with_chatgpt_auth_tokens_for_surface;
use codex_login::complete_device_code_login;
use codex_login::login_with_api_key;
use codex_login::login_with_api_key_for_surface;
use codex_login::request_device_code;
use codex_login::run_login_server;
use codex_mcp::McpRuntimeContext;
@@ -1,4 +1,5 @@
use super::*;
use codex_http_state::HttpStateSurface;
// Duration before a browser ChatGPT login attempt is abandoned.
const LOGIN_CHATGPT_TIMEOUT: Duration = Duration::from_secs(10 * 60);
@@ -85,15 +86,25 @@ impl AccountRequestProcessor {
&self,
request_id: ConnectionRequestId,
params: LoginAccountParams,
app_server_client_name: Option<&str>,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
self.login_v2(request_id, params).await.map(|()| None)
self.login_v2(
request_id,
params,
Self::http_state_surface(app_server_client_name),
)
.await
.map(|()| None)
}
pub(crate) async fn logout_account(
&self,
request_id: ConnectionRequestId,
app_server_client_name: Option<&str>,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
self.logout_v2(request_id).await.map(|()| None)
self.logout_v2(request_id, Self::http_state_surface(app_server_client_name))
.await
.map(|()| None)
}
pub(crate) async fn cancel_login_account(
@@ -151,6 +162,10 @@ impl AccountRequestProcessor {
self.auth_manager.clear_external_auth();
}
fn http_state_surface(app_server_client_name: Option<&str>) -> HttpStateSurface {
HttpStateSurface::from_app_server_client_name(app_server_client_name.unwrap_or_default())
}
fn current_account_updated_notification(&self) -> AccountUpdatedNotification {
let auth = self.auth_manager.auth_cached();
AccountUpdatedNotification {
@@ -210,20 +225,26 @@ impl AccountRequestProcessor {
&self,
request_id: ConnectionRequestId,
params: LoginAccountParams,
http_state_surface: HttpStateSurface,
) -> Result<(), JSONRPCErrorError> {
match params {
LoginAccountParams::ApiKey { api_key } => {
self.login_api_key_v2(request_id, LoginApiKeyParams { api_key })
.await;
self.login_api_key_v2(
request_id,
LoginApiKeyParams { api_key },
http_state_surface,
)
.await;
}
LoginAccountParams::Chatgpt {
codex_streamlined_login,
} => {
self.login_chatgpt_v2(request_id, codex_streamlined_login)
self.login_chatgpt_v2(request_id, codex_streamlined_login, http_state_surface)
.await;
}
LoginAccountParams::ChatgptDeviceCode => {
self.login_chatgpt_device_code_v2(request_id).await;
self.login_chatgpt_device_code_v2(request_id, http_state_surface)
.await;
}
LoginAccountParams::ChatgptAuthTokens {
access_token,
@@ -235,6 +256,7 @@ impl AccountRequestProcessor {
access_token,
chatgpt_account_id,
chatgpt_plan_type,
http_state_surface,
)
.await;
}
@@ -251,6 +273,7 @@ impl AccountRequestProcessor {
async fn login_api_key_common(
&self,
params: &LoginApiKeyParams,
http_state_surface: HttpStateSurface,
) -> std::result::Result<(), JSONRPCErrorError> {
if self.auth_manager.is_external_chatgpt_auth_active() {
return Err(self.external_auth_active_error());
@@ -273,10 +296,11 @@ impl AccountRequestProcessor {
}
}
match login_with_api_key(
match login_with_api_key_for_surface(
&self.config.codex_home,
&params.api_key,
self.config.cli_auth_credentials_store_mode,
http_state_surface,
) {
Ok(()) => {
self.auth_manager.reload().await;
@@ -286,9 +310,14 @@ impl AccountRequestProcessor {
}
}
async fn login_api_key_v2(&self, request_id: ConnectionRequestId, params: LoginApiKeyParams) {
async fn login_api_key_v2(
&self,
request_id: ConnectionRequestId,
params: LoginApiKeyParams,
http_state_surface: HttpStateSurface,
) {
let result = self
.login_api_key_common(&params)
.login_api_key_common(&params, http_state_surface)
.await
.map(|()| LoginAccountResponse::ApiKey {});
let logged_in = result.is_ok();
@@ -304,6 +333,7 @@ impl AccountRequestProcessor {
async fn login_chatgpt_common(
&self,
codex_streamlined_login: bool,
http_state_surface: HttpStateSurface,
) -> std::result::Result<LoginServerOptions, JSONRPCErrorError> {
let config = self.config.as_ref();
@@ -320,6 +350,7 @@ impl AccountRequestProcessor {
let opts = LoginServerOptions {
open_browser: false,
codex_streamlined_login,
http_state_surface,
..LoginServerOptions::new(
config.codex_home.to_path_buf(),
CLIENT_ID.to_string(),
@@ -354,16 +385,22 @@ impl AccountRequestProcessor {
&self,
request_id: ConnectionRequestId,
codex_streamlined_login: bool,
http_state_surface: HttpStateSurface,
) {
let result = self.login_chatgpt_response(codex_streamlined_login).await;
let result = self
.login_chatgpt_response(codex_streamlined_login, http_state_surface)
.await;
self.outgoing.send_result(request_id, result).await;
}
async fn login_chatgpt_response(
&self,
codex_streamlined_login: bool,
http_state_surface: HttpStateSurface,
) -> Result<LoginAccountResponse, JSONRPCErrorError> {
let opts = self.login_chatgpt_common(codex_streamlined_login).await?;
let opts = self
.login_chatgpt_common(codex_streamlined_login, http_state_surface)
.await?;
let server = run_login_server(opts)
.map_err(|err| internal_error(format!("failed to start login server: {err}")))?;
let login_id = Uuid::new_v4();
@@ -426,16 +463,23 @@ impl AccountRequestProcessor {
})
}
async fn login_chatgpt_device_code_v2(&self, request_id: ConnectionRequestId) {
let result = self.login_chatgpt_device_code_response().await;
async fn login_chatgpt_device_code_v2(
&self,
request_id: ConnectionRequestId,
http_state_surface: HttpStateSurface,
) {
let result = self
.login_chatgpt_device_code_response(http_state_surface)
.await;
self.outgoing.send_result(request_id, result).await;
}
async fn login_chatgpt_device_code_response(
&self,
http_state_surface: HttpStateSurface,
) -> Result<LoginAccountResponse, JSONRPCErrorError> {
let opts = self
.login_chatgpt_common(/*codex_streamlined_login*/ false)
.login_chatgpt_common(/*codex_streamlined_login*/ false, http_state_surface)
.await?;
let device_code = request_device_code(&opts)
.await
@@ -534,9 +578,15 @@ impl AccountRequestProcessor {
access_token: String,
chatgpt_account_id: String,
chatgpt_plan_type: Option<String>,
http_state_surface: HttpStateSurface,
) {
let result = self
.login_chatgpt_auth_tokens_response(access_token, chatgpt_account_id, chatgpt_plan_type)
.login_chatgpt_auth_tokens_response(
access_token,
chatgpt_account_id,
chatgpt_plan_type,
http_state_surface,
)
.await;
let logged_in = result.is_ok();
self.outgoing.send_result(request_id, result).await;
@@ -552,6 +602,7 @@ impl AccountRequestProcessor {
access_token: String,
chatgpt_account_id: String,
chatgpt_plan_type: Option<String>,
http_state_surface: HttpStateSurface,
) -> Result<LoginAccountResponse, JSONRPCErrorError> {
if matches!(
self.config.forced_login_method,
@@ -578,11 +629,12 @@ impl AccountRequestProcessor {
)));
}
login_with_chatgpt_auth_tokens(
login_with_chatgpt_auth_tokens_for_surface(
&self.config.codex_home,
&access_token,
&chatgpt_account_id,
chatgpt_plan_type.as_deref(),
http_state_surface,
)
.map_err(|err| internal_error(format!("failed to set external auth: {err}")))?;
self.auth_manager.reload().await;
@@ -667,7 +719,10 @@ impl AccountRequestProcessor {
}
}
async fn logout_common(&self) -> std::result::Result<Option<AuthMode>, JSONRPCErrorError> {
async fn logout_common(
&self,
http_state_surface: HttpStateSurface,
) -> std::result::Result<Option<AuthMode>, JSONRPCErrorError> {
// Cancel any active login attempt.
{
let mut guard = self.active_login.lock().await;
@@ -676,7 +731,11 @@ impl AccountRequestProcessor {
}
}
match self.auth_manager.logout_with_revoke().await {
match self
.auth_manager
.logout_with_revoke_for_surface(http_state_surface)
.await
{
Ok(_) => {}
Err(err) => {
return Err(internal_error(format!("logout failed: {err}")));
@@ -698,8 +757,12 @@ impl AccountRequestProcessor {
.map(CodexAuth::api_auth_mode))
}
async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> {
let result = self.logout_common().await;
async fn logout_v2(
&self,
request_id: ConnectionRequestId,
http_state_surface: HttpStateSurface,
) -> Result<(), JSONRPCErrorError> {
let result = self.logout_common(http_state_surface).await;
let account_updated =
result
.as_ref()
+37 -2
View File
@@ -17,12 +17,14 @@ use codex_app_server_protocol::CancelLoginAccountResponse;
use codex_app_server_protocol::CancelLoginAccountStatus;
use codex_app_server_protocol::ChatgptAuthTokensRefreshReason;
use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::GetAccountParams;
use codex_app_server_protocol::GetAccountResponse;
use codex_app_server_protocol::GetAuthStatusParams;
use codex_app_server_protocol::GetAuthStatusResponse;
use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::LoginAccountResponse;
@@ -33,6 +35,8 @@ use codex_app_server_protocol::ServerRequest;
use codex_app_server_protocol::TurnCompletedNotification;
use codex_app_server_protocol::TurnStatus;
use codex_config::types::AuthCredentialsStoreMode;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
use codex_login::login_with_api_key;
use codex_protocol::account::PlanType as AccountPlanType;
@@ -61,6 +65,7 @@ const WORKSPACE_ID_INITIAL: &str = "123e4567-e89b-42d3-a456-426614174011";
const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012";
const WORKSPACE_ID_DEVICE: &str = "123e4567-e89b-42d3-a456-426614174013";
const WORKSPACE_ID_STALE: &str = "123e4567-e89b-42d3-a456-426614174014";
const INTEGRITY_STATE: &str = "ois1.header.nonce.ciphertext";
// Helper to create a minimal config.toml for the app server
#[derive(Default)]
@@ -74,6 +79,14 @@ struct CreateConfigTomlParams {
extra_provider_config: Option<String>,
}
fn client_info(name: &str) -> ClientInfo {
ClientInfo {
name: name.to_string(),
title: None,
version: "1.0.0".to_string(),
}
}
fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std::io::Result<()> {
let config_toml = codex_home.join("config.toml");
let base_url = params
@@ -199,11 +212,21 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> {
"sk-test-key",
AuthCredentialsStoreMode::File,
)?;
let http_state = HttpStateStore::new(codex_home.path().to_path_buf());
http_state.set(
HttpStateSurface::CodexDesktopSsh,
INTEGRITY_STATE.to_string(),
)?;
assert!(codex_home.path().join("auth.json").exists());
let mut mcp =
TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let initialized = timeout(
DEFAULT_READ_TIMEOUT,
mcp.initialize_with_client_info(client_info("codex_desktop_ssh")),
)
.await??;
assert!(matches!(initialized, JSONRPCMessage::Response(_)));
let id = mcp.send_logout_account_request().await?;
let resp: JSONRPCResponse = timeout(
@@ -232,6 +255,7 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> {
!codex_home.path().join("auth.json").exists(),
"auth.json should be deleted"
);
assert_eq!(http_state.get(HttpStateSurface::CodexDesktopSsh)?, None);
let get_id = mcp
.send_get_account_request(GetAccountParams {
@@ -913,8 +937,18 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
let http_state = HttpStateStore::new(codex_home.path().to_path_buf());
http_state.set(
HttpStateSurface::CodexDesktopSsh,
INTEGRITY_STATE.to_string(),
)?;
let mut mcp = TestAppServer::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let initialized = timeout(
DEFAULT_READ_TIMEOUT,
mcp.initialize_with_client_info(client_info("codex_desktop_ssh")),
)
.await??;
assert!(matches!(initialized, JSONRPCMessage::Response(_)));
let req_id = mcp
.send_login_account_api_key_request("sk-test-key")
@@ -953,6 +987,7 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> {
pretty_assertions::assert_eq!(payload.plan_type, None);
assert!(codex_home.path().join("auth.json").exists());
assert_eq!(http_state.get(HttpStateSurface::CodexDesktopSsh)?, None);
Ok(())
}
+1
View File
@@ -15,6 +15,7 @@ codex-agent-identity = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-client = { workspace = true }
codex-config = { workspace = true }
codex-http-state = { workspace = true }
codex-keyring-store = { workspace = true }
codex-model-provider-info = { workspace = true }
codex-otel = { workspace = true }
+93 -8
View File
@@ -21,6 +21,7 @@ use codex_agent_identity::decode_agent_identity_jwt;
use codex_agent_identity::fetch_agent_identity_jwks;
use codex_app_server_protocol::AuthMode;
use codex_app_server_protocol::AuthMode as ApiAuthMode;
use codex_http_state::HttpStateSurface;
use codex_protocol::config_types::ForcedLoginMethod;
use codex_protocol::config_types::ModelProviderAuthInfo;
@@ -515,6 +516,19 @@ pub fn logout(
pub async fn logout_with_revoke(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
) -> std::io::Result<bool> {
logout_with_revoke_for_surface(
codex_home,
auth_credentials_store_mode,
HttpStateSurface::CodexCli,
)
.await
}
pub async fn logout_with_revoke_for_surface(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
http_state_surface: HttpStateSurface,
) -> std::io::Result<bool> {
AuthManager::new(
codex_home.to_path_buf(),
@@ -523,7 +537,7 @@ pub async fn logout_with_revoke(
/*chatgpt_base_url*/ None,
)
.await
.logout_with_revoke()
.logout_with_revoke_for_surface(http_state_surface)
.await
}
@@ -532,6 +546,20 @@ pub fn login_with_api_key(
codex_home: &Path,
api_key: &str,
auth_credentials_store_mode: AuthCredentialsStoreMode,
) -> std::io::Result<()> {
login_with_api_key_for_surface(
codex_home,
api_key,
auth_credentials_store_mode,
HttpStateSurface::CodexCli,
)
}
pub fn login_with_api_key_for_surface(
codex_home: &Path,
api_key: &str,
auth_credentials_store_mode: AuthCredentialsStoreMode,
http_state_surface: HttpStateSurface,
) -> std::io::Result<()> {
let auth_dot_json = AuthDotJson {
auth_mode: Some(ApiAuthMode::ApiKey),
@@ -540,7 +568,9 @@ pub fn login_with_api_key(
last_refresh: None,
agent_identity: None,
};
save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode)
save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode)?;
crate::http_state::clear(codex_home, http_state_surface);
Ok(())
}
/// Writes an `auth.json` that contains only the access token.
@@ -549,6 +579,23 @@ pub async fn login_with_access_token(
access_token: &str,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
) -> std::io::Result<()> {
login_with_access_token_for_surface(
codex_home,
access_token,
auth_credentials_store_mode,
chatgpt_base_url,
HttpStateSurface::CodexCli,
)
.await
}
pub async fn login_with_access_token_for_surface(
codex_home: &Path,
access_token: &str,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
http_state_surface: HttpStateSurface,
) -> std::io::Result<()> {
let base_url = chatgpt_base_url
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
@@ -562,7 +609,9 @@ pub async fn login_with_access_token(
last_refresh: None,
agent_identity: Some(access_token.to_string()),
};
save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode)
save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode)?;
crate::http_state::clear(codex_home, http_state_surface);
Ok(())
}
/// Writes an in-memory auth payload for externally managed ChatGPT tokens.
@@ -571,6 +620,22 @@ pub fn login_with_chatgpt_auth_tokens(
access_token: &str,
chatgpt_account_id: &str,
chatgpt_plan_type: Option<&str>,
) -> std::io::Result<()> {
login_with_chatgpt_auth_tokens_for_surface(
codex_home,
access_token,
chatgpt_account_id,
chatgpt_plan_type,
HttpStateSurface::CodexCli,
)
}
pub fn login_with_chatgpt_auth_tokens_for_surface(
codex_home: &Path,
access_token: &str,
chatgpt_account_id: &str,
chatgpt_plan_type: Option<&str>,
http_state_surface: HttpStateSurface,
) -> std::io::Result<()> {
let auth_dot_json = AuthDotJson::from_external_access_token(
access_token,
@@ -581,7 +646,9 @@ pub fn login_with_chatgpt_auth_tokens(
codex_home,
&auth_dot_json,
AuthCredentialsStoreMode::Ephemeral,
)
)?;
crate::http_state::clear(codex_home, http_state_surface);
Ok(())
}
/// Persist the provided auth payload using the specified backend.
@@ -711,6 +778,7 @@ fn logout_with_message(
// External auth tokens live in the ephemeral store, but persistent auth may still exist
// from earlier logins. Clear both so a forced logout truly removes all active auth.
let removal_result = logout_all_stores(codex_home, auth_credentials_store_mode);
crate::http_state::clear(codex_home, HttpStateSurface::CodexCli);
let error_message = match removal_result {
Ok(_) => message,
Err(err) => format!("{message}. Failed to remove auth.json: {err}"),
@@ -1773,23 +1841,40 @@ impl AuthManager {
/// reloads the in‑memory auth cache so callers immediately observe the
/// unauthenticated state.
pub async fn logout(&self) -> std::io::Result<bool> {
let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?;
self.logout_for_surface(HttpStateSurface::CodexCli).await
}
pub async fn logout_for_surface(
&self,
http_state_surface: HttpStateSurface,
) -> std::io::Result<bool> {
let result = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode);
crate::http_state::clear(&self.codex_home, http_state_surface);
// Always reload to clear any cached auth (even if file absent).
self.reload().await;
Ok(removed)
result
}
pub async fn logout_with_revoke(&self) -> std::io::Result<bool> {
self.logout_with_revoke_for_surface(HttpStateSurface::CodexCli)
.await
}
pub async fn logout_with_revoke_for_surface(
&self,
http_state_surface: HttpStateSurface,
) -> std::io::Result<bool> {
let auth_dot_json = self
.auth_cached()
.and_then(|auth| auth.get_current_auth_json());
if let Err(err) = revoke_auth_tokens(auth_dot_json.as_ref()).await {
tracing::warn!("failed to revoke auth tokens during logout: {err}");
}
let result = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?;
let result = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode);
crate::http_state::clear(&self.codex_home, http_state_surface);
// Always reload to clear any cached auth (even if file absent).
self.reload().await;
Ok(result)
result
}
pub fn get_api_auth_mode(&self) -> Option<ApiAuthMode> {
+3 -3
View File
@@ -196,6 +196,7 @@ pub async fn complete_device_code_login(
let tokens = crate::server::exchange_code_for_tokens(
base_url,
&opts.client_id,
opts.http_state_surface,
&redirect_uri,
&pkce,
&code_resp.authorization_code,
@@ -213,9 +214,8 @@ pub async fn complete_device_code_login(
crate::server::persist_tokens_async(
&opts.codex_home,
/*api_key*/ None,
tokens.id_token,
tokens.access_token,
tokens.refresh_token,
&tokens,
opts.http_state_surface,
opts.cli_auth_credentials_store_mode,
)
.await
+58
View File
@@ -0,0 +1,58 @@
use std::path::Path;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use tracing::warn;
pub(crate) fn replace_after_login(
codex_home: &Path,
surface: HttpStateSurface,
state: Option<String>,
) {
let store = HttpStateStore::new(codex_home.to_path_buf());
let result = match state {
Some(state) => store.set(surface, state),
None => store.clear(surface),
};
if let Err(err) = result {
warn!(%surface, "failed to reset HTTP state after login: {err}");
}
}
pub(crate) fn clear(codex_home: &Path, surface: HttpStateSurface) {
if let Err(err) = HttpStateStore::new(codex_home.to_path_buf()).clear(surface) {
warn!(%surface, "failed to clear HTTP state: {err}");
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn replace_after_login_sets_and_clears_surface_state() {
let codex_home = TempDir::new().expect("tempdir");
let store = HttpStateStore::new(codex_home.path().to_path_buf());
replace_after_login(
codex_home.path(),
HttpStateSurface::CodexDesktop,
Some("minted-state".to_string()),
);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load"),
Some("minted-state".to_string()),
);
replace_after_login(codex_home.path(), HttpStateSurface::CodexDesktop, None);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load"),
None,
);
}
}
+2
View File
@@ -3,6 +3,7 @@ pub mod auth_env_telemetry;
pub mod token_data;
mod device_code_auth;
mod http_state;
mod pkce;
mod server;
@@ -40,6 +41,7 @@ pub use auth::enforce_login_restrictions;
pub use auth::load_auth_dot_json;
pub use auth::login_with_access_token;
pub use auth::login_with_api_key;
pub use auth::login_with_api_key_for_surface;
pub use auth::logout;
pub use auth::logout_with_revoke;
pub use auth::read_codex_access_token_from_env;
+49 -18
View File
@@ -39,6 +39,7 @@ use chrono::Utc;
use codex_app_server_protocol::AuthMode;
use codex_client::build_reqwest_client_with_custom_ca;
use codex_config::types::AuthCredentialsStoreMode;
use codex_http_state::HttpStateSurface;
use codex_utils_template::Template;
use rand::RngCore;
use serde_json::Value as JsonValue;
@@ -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,
pub http_state_surface: HttpStateSurface,
}
impl ServerOptions {
@@ -92,6 +94,7 @@ impl ServerOptions {
forced_chatgpt_workspace_id,
codex_streamlined_login: false,
cli_auth_credentials_store_mode,
http_state_surface: HttpStateSurface::CodexCli,
}
}
}
@@ -335,8 +338,15 @@ async fn process_request(
}
};
match exchange_code_for_tokens(&opts.issuer, &opts.client_id, redirect_uri, pkce, &code)
.await
match exchange_code_for_tokens(
&opts.issuer,
&opts.client_id,
opts.http_state_surface,
redirect_uri,
pkce,
&code,
)
.await
{
Ok(tokens) => {
if let Err(message) = ensure_workspace_allowed(
@@ -358,9 +368,8 @@ async fn process_request(
if let Err(err) = persist_tokens_async(
&opts.codex_home,
api_key.clone(),
tokens.id_token.clone(),
tokens.access_token.clone(),
tokens.refresh_token.clone(),
&tokens,
opts.http_state_surface,
opts.cli_auth_credentials_store_mode,
)
.await
@@ -606,6 +615,7 @@ pub(crate) struct ExchangedTokens {
pub id_token: String,
pub access_token: String,
pub refresh_token: String,
pub state: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -714,6 +724,7 @@ fn sanitize_url_for_logging(url: &str) -> String {
pub(crate) async fn exchange_code_for_tokens(
issuer: &str,
client_id: &str,
http_state_surface: HttpStateSurface,
redirect_uri: &str,
pkce: &PkceCodes,
code: &str,
@@ -723,6 +734,7 @@ pub(crate) async fn exchange_code_for_tokens(
id_token: String,
access_token: String,
refresh_token: String,
state: Option<String>,
}
let client = build_reqwest_client_with_custom_ca(reqwest::Client::builder())?;
@@ -737,11 +749,12 @@ pub(crate) async fn exchange_code_for_tokens(
.post(token_endpoint)
.header("Content-Type", "application/x-www-form-urlencoded")
.body(format!(
"grant_type=authorization_code&code={}&redirect_uri={}&client_id={}&code_verifier={}",
"grant_type=authorization_code&code={}&redirect_uri={}&client_id={}&code_verifier={}&surface={}",
urlencoding::encode(code),
urlencoding::encode(redirect_uri),
urlencoding::encode(client_id),
urlencoding::encode(&pkce.code_verifier)
urlencoding::encode(&pkce.code_verifier),
urlencoding::encode(http_state_surface.as_str()),
))
.send()
.await;
@@ -781,6 +794,7 @@ pub(crate) async fn exchange_code_for_tokens(
id_token: tokens.id_token,
access_token: tokens.access_token,
refresh_token: tokens.refresh_token,
state: tokens.state,
})
}
@@ -789,15 +803,20 @@ pub(crate) async fn exchange_code_for_tokens(
pub(crate) async fn persist_tokens_async(
codex_home: &Path,
api_key: Option<String>,
id_token: String,
access_token: String,
refresh_token: String,
tokens: &ExchangedTokens,
http_state_surface: HttpStateSurface,
auth_credentials_store_mode: AuthCredentialsStoreMode,
) -> io::Result<()> {
// Reuse existing synchronous logic but run it off the async runtime.
let codex_home = codex_home.to_path_buf();
let auth_codex_home = codex_home.clone();
let id_token = tokens.id_token.clone();
let access_token = tokens.access_token.clone();
let refresh_token = tokens.refresh_token.clone();
let state = tokens.state.clone();
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 load_auth_dot_json(&auth_codex_home, auth_credentials_store_mode)
{
Ok(auth) => auth,
Err(err) => {
warn!("failed to load previous auth before saving new login: {err}");
@@ -823,12 +842,14 @@ pub(crate) async fn persist_tokens_async(
last_refresh: Some(Utc::now()),
agent_identity: None,
};
save_auth(&codex_home, &auth, auth_credentials_store_mode)?;
save_auth(&auth_codex_home, &auth, auth_credentials_store_mode)?;
Ok::<_, io::Error>((previous_auth, auth))
})
.await
.map_err(|e| io::Error::other(format!("persist task failed: {e}")))??;
crate::http_state::replace_after_login(&codex_home, http_state_surface, state);
if should_revoke_auth_tokens(previous_auth.as_ref(), &auth)
&& let Err(err) = revoke_auth_tokens(previous_auth.as_ref()).await
{
@@ -1176,10 +1197,12 @@ mod tests {
use crate::auth::save_auth;
use crate::token_data::TokenData;
use crate::token_data::parse_chatgpt_jwt_claims;
use codex_http_state::HttpStateSurface;
use core_test_support::skip_if_no_network;
use pretty_assertions::assert_eq;
use super::DEFAULT_ISSUER;
use super::ExchangedTokens;
use super::TokenEndpointErrorDetail;
use super::compose_success_url;
use super::html_escape;
@@ -1223,9 +1246,13 @@ mod tests {
persist_tokens_async(
codex_home.path(),
/*api_key*/ None,
jwt_for_account("new-account"),
"new-access".to_string(),
"new-refresh".to_string(),
&ExchangedTokens {
id_token: jwt_for_account("new-account"),
access_token: "new-access".to_string(),
refresh_token: "new-refresh".to_string(),
state: None,
},
HttpStateSurface::CodexCli,
AuthCredentialsStoreMode::File,
)
.await?;
@@ -1283,9 +1310,13 @@ mod tests {
persist_tokens_async(
codex_home.path(),
/*api_key*/ None,
jwt_for_account("new-account"),
"new-access".to_string(),
"shared-refresh".to_string(),
&ExchangedTokens {
id_token: jwt_for_account("new-account"),
access_token: "new-access".to_string(),
refresh_token: "shared-refresh".to_string(),
state: None,
},
HttpStateSurface::CodexCli,
AuthCredentialsStoreMode::File,
)
.await?;
@@ -4,6 +4,8 @@ use anyhow::Context;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use codex_config::types::AuthCredentialsStoreMode;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use codex_login::ServerOptions;
use codex_login::auth::load_auth_dot_json;
use codex_login::run_device_code_login;
@@ -16,6 +18,7 @@ use wiremock::Mock;
use wiremock::MockServer;
use wiremock::Request;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_string_contains;
use wiremock::matchers::method;
use wiremock::matchers::path;
@@ -91,10 +94,12 @@ async fn mock_poll_token_single(server: &MockServer, endpoint: &str, response: R
async fn mock_oauth_token_single(server: &MockServer, jwt: String) {
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("surface=codex_cli"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"id_token": jwt.clone(),
"access_token": "access-token-123",
"refresh_token": "refresh-token-123"
"refresh_token": "refresh-token-123",
"state": "ois1.header.nonce.ciphertext"
})))
.mount(server)
.await;
@@ -156,6 +161,10 @@ async fn device_code_login_integration_succeeds() -> anyhow::Result<()> {
assert_eq!(tokens.refresh_token, "refresh-token-123");
assert_eq!(tokens.id_token.raw_jwt, jwt);
assert_eq!(tokens.account_id.as_deref(), Some(WORKSPACE_ID_ALLOWED));
assert_eq!(
HttpStateStore::new(codex_home.path().to_path_buf()).get(HttpStateSurface::CodexCli)?,
Some("ois1.header.nonce.ciphertext".to_string())
);
Ok(())
}
@@ -9,6 +9,8 @@ use std::time::Duration;
use anyhow::Result;
use base64::Engine;
use codex_config::types::AuthCredentialsStoreMode;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use codex_login::ServerOptions;
use codex_login::run_login_server;
use core_test_support::skip_if_no_network;
@@ -69,6 +71,7 @@ fn start_mock_issuer(chatgpt_account_id: &str) -> (SocketAddr, thread::JoinHandl
"id_token": id_token,
"access_token": "access-123",
"refresh_token": "refresh-123",
"state": "ois1.header.nonce.ciphertext",
});
let data = serde_json::to_vec(&tokens).unwrap();
let mut resp = tiny_http::Response::from_data(data);
@@ -128,6 +131,7 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
force_state: Some(state),
forced_chatgpt_workspace_id: Some(vec![chatgpt_account_id.to_string()]),
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexDesktop,
};
let server = run_login_server(opts)?;
assert!(
@@ -160,6 +164,10 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
assert_eq!(json["tokens"]["access_token"], "access-123");
assert_eq!(json["tokens"]["refresh_token"], "refresh-123");
assert_eq!(json["tokens"]["account_id"], chatgpt_account_id);
assert_eq!(
HttpStateStore::new(codex_home.clone()).get(HttpStateSurface::CodexDesktop)?,
Some("ois1.header.nonce.ciphertext".to_string())
);
// Stop mock issuer
drop(issuer_handle);
@@ -190,6 +198,7 @@ async fn creates_missing_codex_home_dir() -> Result<()> {
force_state: Some(state),
forced_chatgpt_workspace_id: None,
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let server = run_login_server(opts)?;
let login_port = server.actual_port;
@@ -233,6 +242,7 @@ async fn login_server_includes_forced_workspaces_as_one_query_param() -> Result<
WORKSPACE_ID_SECOND_ALLOWED.to_string(),
]),
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let server = run_login_server(opts)?;
let auth_url = Url::parse(&server.auth_url)?;
@@ -271,6 +281,7 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> {
force_state: Some(state.clone()),
forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let server = run_login_server(opts)?;
assert!(
@@ -331,6 +342,7 @@ async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error()
force_state: Some(state.clone()),
forced_chatgpt_workspace_id: None,
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let server = run_login_server(opts)?;
let login_port = server.actual_port;
@@ -399,6 +411,7 @@ async fn oauth_access_denied_unknown_reason_uses_generic_error_page() -> Result<
force_state: Some(state.clone()),
forced_chatgpt_workspace_id: None,
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let server = run_login_server(opts)?;
let login_port = server.actual_port;
@@ -544,6 +557,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> {
force_state: Some("cancel_state".to_string()),
forced_chatgpt_workspace_id: None,
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let first_server = run_login_server(first_opts)?;
@@ -565,6 +579,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> {
force_state: Some("cancel_state_2".to_string()),
forced_chatgpt_workspace_id: None,
codex_streamlined_login: false,
http_state_surface: HttpStateSurface::CodexCli,
};
let second_server = run_login_server(second_opts)?;