openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-03 19:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[codex-login] persist attributable refresh HTTP state [ci changed_files]

Browse Source
This commit is contained in:
Cooper Gamble
2026-06-03 09:54:41 +00:00
parent 1c5ddd9234
commit 37a29a6195
11 changed files with 287 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-transport/src/transport/remote_control/mod.rs
+1 -1
View File
@@ -263,7 +263,7 @@ async fn refresh_pairing_enrollment(
if err.kind() != io::ErrorKind::PermissionDenied {
return handle_pairing_refresh_error(current_enrollment, enrollment, err);
}
let mut auth_recovery = auth_manager.unauthorized_recovery();
let mut auth_recovery = websocket::remote_control_auth_recovery(auth_manager);
let mut auth_change_rx = auth_manager.auth_change_receiver();
if !websocket::recover_remote_control_auth(&mut auth_recovery, &mut auth_change_rx).await {
return Err(err);
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+9 -3
View File
@@ -75,6 +75,12 @@ const REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT: std::time::Duration =
const REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT: std::time::Duration =
std::time::Duration::from_secs(5);
pub(super) fn remote_control_auth_recovery(
auth_manager: &Arc<AuthManager>,
) -> UnauthorizedRecovery {
auth_manager.unauthorized_recovery_for_surface(HttpStateSurface::CodexRemoteControl)
}
struct BoundedOutboundBuffer {
buffer_by_stream: HashMap<(ClientId, StreamId), VecDeque<ServerEnvelope>>,
used_tx: watch::Sender<usize>,
@@ -394,7 +400,7 @@ impl RemoteControlWebsocket {
&shutdown_token,
);
let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
let auth_recovery = auth_manager.unauthorized_recovery();
let auth_recovery = remote_control_auth_recovery(&auth_manager);
let auth_change_rx = auth_manager.auth_change_receiver();
Self {
@@ -628,7 +634,7 @@ impl RemoteControlWebsocket {
return ConnectOutcome::Disabled;
}
self.reconnect_attempt = 0;
self.auth_recovery = self.auth_manager.unauthorized_recovery();
self.auth_recovery = remote_control_auth_recovery(&self.auth_manager);
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Connected);
let enrollment = self.enrollment.as_ref();
@@ -692,7 +698,7 @@ impl RemoteControlWebsocket {
if changed.is_err() {
return ConnectOutcome::Shutdown;
}
self.auth_recovery = self.auth_manager.unauthorized_recovery();
self.auth_recovery = remote_control_auth_recovery(&self.auth_manager);
self.reconnect_attempt = 0;
info!("retrying app-server remote control websocket after auth changed");
}
codex-rs/app-server/src/message_processor.rs
+6 -2
View File
@@ -1297,10 +1297,14 @@ impl MessageProcessor {
self.account_processor.cancel_login_account(params).await
}
ClientRequest::GetAccount { params, .. } => {
self.account_processor.get_account(params).await
self.account_processor
.get_account(params, app_server_client_name.as_deref())
.await
}
ClientRequest::GetAuthStatus { params, .. } => {
self.account_processor.get_auth_status(params).await
self.account_processor
.get_auth_status(params, app_server_client_name.as_deref())
.await
}
ClientRequest::GetAccountRateLimits { .. } => {
self.account_processor.get_account_rate_limits().await
codex-rs/app-server/src/request_processors/account_processor.rs
+21 -6
View File
@@ -116,8 +116,9 @@ impl AccountRequestProcessor {
pub(crate) async fn get_account(
&self,
params: GetAccountParams,
app_server_client_name: Option<&str>,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
self.get_account_response(params)
self.get_account_response(params, Self::http_state_surface(app_server_client_name))
.await
.map(|response| Some(response.into()))
}
@@ -125,8 +126,9 @@ impl AccountRequestProcessor {
pub(crate) async fn get_auth_status(
&self,
params: GetAuthStatusParams,
app_server_client_name: Option<&str>,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
self.get_auth_status_response(params)
self.get_auth_status_response(params, Self::http_state_surface(app_server_client_name))
.await
.map(|response| Some(response.into()))
}
@@ -750,11 +752,20 @@ impl AccountRequestProcessor {
Ok(())
}
async fn refresh_token_if_requested(&self, do_refresh: bool) -> RefreshTokenRequestOutcome {
async fn refresh_token_if_requested(
&self,
do_refresh: bool,
http_state_surface: HttpStateSurface,
) -> RefreshTokenRequestOutcome {
if self.auth_manager.is_external_chatgpt_auth_active() {
return RefreshTokenRequestOutcome::NotAttemptedOrSucceeded;
}
if do_refresh && let Err(err) = self.auth_manager.refresh_token().await {
if do_refresh
&& let Err(err) = self
.auth_manager
.refresh_token_for_surface(http_state_surface)
.await
{
let failed_reason = err.failed_reason();
if failed_reason.is_none() {
tracing::warn!("failed to refresh token while getting account: {err}");
@@ -768,11 +779,13 @@ impl AccountRequestProcessor {
async fn get_auth_status_response(
&self,
params: GetAuthStatusParams,
http_state_surface: HttpStateSurface,
) -> Result<GetAuthStatusResponse, JSONRPCErrorError> {
let include_token = params.include_token.unwrap_or(false);
let do_refresh = params.refresh_token.unwrap_or(false);
self.refresh_token_if_requested(do_refresh).await;
self.refresh_token_if_requested(do_refresh, http_state_surface)
.await;
// Determine whether auth is required based on the active model provider.
// If a custom provider is configured with `requires_openai_auth == false`,
@@ -834,10 +847,12 @@ impl AccountRequestProcessor {
async fn get_account_response(
&self,
params: GetAccountParams,
http_state_surface: HttpStateSurface,
) -> Result<GetAccountResponse, JSONRPCErrorError> {
let do_refresh = params.refresh_token;
self.refresh_token_if_requested(do_refresh).await;
self.refresh_token_if_requested(do_refresh, http_state_surface)
.await;
let provider = create_model_provider(
self.config.model_provider.clone(),
codex-rs/app-server/tests/suite/v2/account.rs
+72
View File
@@ -1774,6 +1774,78 @@ async fn get_account_with_chatgpt() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn get_auth_status_refresh_persists_state_for_initialized_surface() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("stale-access-token")
.refresh_token("stale-refresh-token")
.account_id(WORKSPACE_ID_STALE)
.email("user@example.com")
.plan_type("pro"),
AuthCredentialsStoreMode::File,
)?;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": "new-access-token",
"refresh_token": "new-refresh-token",
"state": "new-state"
})))
.expect(1)
.mount(&server)
.await;
let refresh_url = format!("{}/oauth/token", server.uri());
let mut mcp = TestAppServer::new_with_env(
codex_home.path(),
&[
("OPENAI_API_KEY", None),
(
REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR,
Some(refresh_url.as_str()),
),
],
)
.await?;
let initialized = timeout(
DEFAULT_READ_TIMEOUT,
mcp.initialize_with_client_info(client_info("codex_desktop_ssh")),
)
.await??;
assert!(matches!(initialized, JSONRPCMessage::Response(_)));
let request_id = mcp
.send_get_auth_status_request(GetAuthStatusParams {
include_token: Some(true),
refresh_token: Some(true),
})
.await?;
let resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let _: GetAuthStatusResponse = to_response(resp)?;
let http_state = HttpStateStore::new(codex_home.path().to_path_buf());
assert_eq!(
http_state.get(HttpStateSurface::CodexDesktopSsh)?,
Some("new-state".to_string()),
);
let requests = server.received_requests().await.unwrap_or_default();
assert_eq!(requests.len(), 1);
let body: serde_json::Value = serde_json::from_slice(&requests[0].body)?;
assert_eq!(body["surface"], "codex_desktop_ssh");
server.verify().await;
Ok(())
}
#[tokio::test]
async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<()> {
let codex_home = TempDir::new()?;
codex-rs/core/src/client.rs
+16 -7
View File
@@ -1316,9 +1316,14 @@ impl ModelClientSession {
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
let auth_manager = self.client.state.provider.auth_manager();
let mut auth_recovery = auth_manager
.as_ref()
.map(AuthManager::unauthorized_recovery);
let http_state_surface = self.client.http_state_surface();
let mut auth_recovery = auth_manager.as_ref().map(|auth_manager| {
if let Some(http_state_surface) = http_state_surface {
auth_manager.unauthorized_recovery_for_surface(http_state_surface)
} else {
auth_manager.unauthorized_recovery()
}
});
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
@@ -1431,10 +1436,14 @@ impl ModelClientSession {
inference_trace: &InferenceTraceContext,
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.provider.auth_manager();
let mut auth_recovery = auth_manager
.as_ref()
.map(AuthManager::unauthorized_recovery);
let http_state_surface = self.client.http_state_surface();
let mut auth_recovery = auth_manager.as_ref().map(|auth_manager| {
if let Some(http_state_surface) = http_state_surface {
auth_manager.unauthorized_recovery_for_surface(http_state_surface)
} else {
auth_manager.unauthorized_recovery()
}
});
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
codex-rs/login/src/auth/auth_tests.rs
+2
View File
@@ -312,6 +312,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
step: UnauthorizedRecoveryStep::Reload,
expected_account_id: None,
mode: UnauthorizedRecoveryMode::Managed,
http_state_surface: None,
};
assert_eq!(managed.mode_name(), "managed");
assert_eq!(managed.step_name(), "reload");
@@ -321,6 +322,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
step: UnauthorizedRecoveryStep::ExternalRefresh,
expected_account_id: None,
mode: UnauthorizedRecoveryMode::External,
http_state_surface: None,
};
assert_eq!(external.mode_name(), "external");
assert_eq!(external.step_name(), "external_refresh");
codex-rs/login/src/auth/manager.rs
+68 -9
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;
@@ -824,11 +825,13 @@ fn persist_tokens(
async fn request_chatgpt_token_refresh(
refresh_token: String,
client: &CodexHttpClient,
http_state_surface: Option<HttpStateSurface>,
) -> Result<RefreshResponse, RefreshTokenError> {
let refresh_request = RefreshRequest {
client_id: CLIENT_ID,
grant_type: "refresh_token",
refresh_token,
surface: http_state_surface.map(HttpStateSurface::as_str),
};
let endpoint = refresh_token_endpoint();
@@ -924,6 +927,8 @@ struct RefreshRequest {
client_id: &'static str,
grant_type: &'static str,
refresh_token: String,
#[serde(skip_serializing_if = "Option::is_none")]
surface: Option<&'static str>,
}
#[derive(Deserialize, Clone)]
@@ -931,6 +936,7 @@ struct RefreshResponse {
id_token: Option<String>,
access_token: Option<String>,
refresh_token: Option<String>,
state: Option<String>,
}
// Shared constant for token refresh (client id used for oauth token refresh flow)
@@ -1086,6 +1092,7 @@ pub struct UnauthorizedRecovery {
step: UnauthorizedRecoveryStep,
expected_account_id: Option<String>,
mode: UnauthorizedRecoveryMode,
http_state_surface: Option<HttpStateSurface>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -1100,7 +1107,7 @@ impl UnauthorizedRecoveryStepResult {
}
impl UnauthorizedRecovery {
fn new(manager: Arc<AuthManager>) -> Self {
fn new(manager: Arc<AuthManager>, http_state_surface: Option<HttpStateSurface>) -> Self {
let cached_auth = manager.auth_cached();
let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id);
let mode = if manager.has_external_api_key_auth()
@@ -1121,6 +1128,7 @@ impl UnauthorizedRecovery {
step,
expected_account_id,
mode,
http_state_surface,
}
}
@@ -1227,7 +1235,9 @@ impl UnauthorizedRecovery {
}
}
UnauthorizedRecoveryStep::RefreshToken => {
self.manager.refresh_token_from_authority().await?;
self.manager
.refresh_token_from_authority_with_surface(self.http_state_surface)
.await?;
self.step = UnauthorizedRecoveryStep::Done;
return Ok(UnauthorizedRecoveryStepResult {
auth_state_changed: Some(true),
@@ -1649,7 +1659,14 @@ impl AuthManager {
}
pub fn unauthorized_recovery(self: &Arc<Self>) -> UnauthorizedRecovery {
UnauthorizedRecovery::new(Arc::clone(self))
UnauthorizedRecovery::new(Arc::clone(self), None)
}
pub fn unauthorized_recovery_for_surface(
self: &Arc<Self>,
http_state_surface: HttpStateSurface,
) -> UnauthorizedRecovery {
UnauthorizedRecovery::new(Arc::clone(self), Some(http_state_surface))
}
fn external_auth(&self) -> Option<Arc<dyn ExternalAuth>> {
@@ -1692,6 +1709,21 @@ impl AuthManager {
/// can assume that some other instance already refreshed it. If the persisted
/// token is the same as the cached, then ask the token authority to refresh.
pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> {
self.refresh_token_with_surface(None).await
}
pub async fn refresh_token_for_surface(
&self,
http_state_surface: HttpStateSurface,
) -> Result<(), RefreshTokenError> {
self.refresh_token_with_surface(Some(http_state_surface))
.await
}
async fn refresh_token_with_surface(
&self,
http_state_surface: Option<HttpStateSurface>,
) -> Result<(), RefreshTokenError> {
let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| {
RefreshTokenError::Permanent(RefreshTokenFailedError::new(
RefreshTokenFailedReason::Other,
@@ -1717,7 +1749,10 @@ impl AuthManager {
tracing::info!("Skipping token refresh because auth changed after guarded reload.");
Ok(())
}
ReloadOutcome::ReloadedNoChange => self.refresh_token_from_authority_impl().await,
ReloadOutcome::ReloadedNoChange => {
self.refresh_token_from_authority_impl(http_state_surface)
.await
}
ReloadOutcome::Skipped => {
Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new(
RefreshTokenFailedReason::Other,
@@ -1732,16 +1767,27 @@ impl AuthManager {
/// observe refreshed token. If the token refresh fails, returns the error to
/// the caller.
pub async fn refresh_token_from_authority(&self) -> Result<(), RefreshTokenError> {
self.refresh_token_from_authority_with_surface(None).await
}
async fn refresh_token_from_authority_with_surface(
&self,
http_state_surface: Option<HttpStateSurface>,
) -> Result<(), RefreshTokenError> {
let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| {
RefreshTokenError::Permanent(RefreshTokenFailedError::new(
RefreshTokenFailedReason::Other,
REFRESH_TOKEN_UNKNOWN_MESSAGE.to_string(),
))
})?;
self.refresh_token_from_authority_impl().await
self.refresh_token_from_authority_impl(http_state_surface)
.await
}
async fn refresh_token_from_authority_impl(&self) -> Result<(), RefreshTokenError> {
async fn refresh_token_from_authority_impl(
&self,
http_state_surface: Option<HttpStateSurface>,
) -> Result<(), RefreshTokenError> {
tracing::info!("Refreshing token");
let auth = match self.auth_cached() {
@@ -1764,8 +1810,12 @@ impl AuthManager {
"Token data is not available.",
))
})?;
self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token)
.await
self.refresh_and_persist_chatgpt_token(
&chatgpt_auth,
token_data.refresh_token,
http_state_surface,
)
.await
}
CodexAuth::ApiKey(_) | CodexAuth::AgentIdentity(_) => Ok(()),
};
@@ -1905,8 +1955,10 @@ impl AuthManager {
&self,
auth: &ChatgptAuth,
refresh_token: String,
http_state_surface: Option<HttpStateSurface>,
) -> Result<(), RefreshTokenError> {
let refresh_response = request_chatgpt_token_refresh(refresh_token, auth.client()).await?;
let refresh_response =
request_chatgpt_token_refresh(refresh_token, auth.client(), http_state_surface).await?;
persist_tokens(
auth.storage(),
@@ -1916,6 +1968,13 @@ impl AuthManager {
)
.map_err(RefreshTokenError::from)?;
self.reload().await;
if let Some(http_state_surface) = http_state_surface {
crate::http_state::replace_after_refresh(
&self.codex_home,
http_state_surface,
refresh_response.state,
);
}
Ok(())
}
codex-rs/login/src/http_state.rs
+13
View File
@@ -18,6 +18,19 @@ pub(crate) fn replace_after_login(
}
}
pub(crate) fn replace_after_refresh(
codex_home: &Path,
surface: HttpStateSurface,
state: Option<String>,
) {
let Some(state) = state else {
return;
};
if let Err(err) = HttpStateStore::new(codex_home.to_path_buf()).set(surface, state) {
warn!(%surface, "failed to reset HTTP state after refresh: {err}");
}
}
pub(crate) fn clear_all(codex_home: &Path) {
let store = HttpStateStore::new(codex_home.to_path_buf());
for surface in HttpStateSurface::ALL {
codex-rs/login/src/http_state_tests.rs
+29
View File
@@ -36,3 +36,32 @@ fn replace_after_login_sets_and_clears_surface_state() {
None,
);
}
#[test]
fn replace_after_refresh_preserves_state_without_replacement_and_sets_returned_state() {
let codex_home = TempDir::new().expect("tempdir");
let store = HttpStateStore::new(codex_home.path().to_path_buf());
store
.set(HttpStateSurface::CodexDesktop, "old-state".to_string())
.expect("state should store");
replace_after_refresh(codex_home.path(), HttpStateSurface::CodexDesktop, None);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load"),
Some("old-state".to_string()),
);
replace_after_refresh(
codex_home.path(),
HttpStateSurface::CodexDesktop,
Some("new-state".to_string()),
);
assert_eq!(
store
.get(HttpStateSurface::CodexDesktop)
.expect("state should load"),
Some("new-state".to_string()),
);
}
codex-rs/login/tests/suite/auth_refresh.rs
+50
View File
@@ -5,6 +5,8 @@ use chrono::Duration;
use chrono::Utc;
use codex_app_server_protocol::AuthMode;
use codex_config::types::AuthCredentialsStoreMode;
use codex_http_state::HttpStateStore;
use codex_http_state::HttpStateSurface;
use codex_login::AuthDotJson;
use codex_login::AuthManager;
use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
@@ -94,6 +96,54 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> {
Ok(())
}
#[serial_test::serial(auth_refresh)]
#[tokio::test]
async fn refresh_token_for_surface_sends_surface_and_persists_returned_state() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": "new-access-token",
"refresh_token": "new-refresh-token",
"state": "new-state"
})))
.expect(1)
.mount(&server)
.await;
let ctx = RefreshTokenTestContext::new(&server).await?;
let initial_auth = AuthDotJson {
auth_mode: Some(AuthMode::Chatgpt),
openai_api_key: None,
tokens: Some(build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN)),
last_refresh: Some(Utc::now() - Duration::days(1)),
agent_identity: None,
};
ctx.write_auth(&initial_auth).await?;
let http_state = HttpStateStore::new(ctx.codex_home.path().to_path_buf());
http_state.set(HttpStateSurface::CodexDesktop, "old-state".to_string())?;
ctx.auth_manager
.refresh_token_for_surface(HttpStateSurface::CodexDesktop)
.await
.context("refresh should succeed")?;
assert_eq!(
http_state.get(HttpStateSurface::CodexDesktop)?,
Some("new-state".to_string()),
);
let requests = server.received_requests().await.unwrap_or_default();
assert_eq!(requests.len(), 1);
let body: serde_json::Value = serde_json::from_slice(&requests[0].body)?;
assert_eq!(body["surface"], "codex_desktop");
server.verify().await;
Ok(())
}
#[serial_test::serial(auth_refresh)]
#[tokio::test]
async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> {