force reauth on workspace-restricted ChatGPT 401

## Summary

- detect workspace-restricted ChatGPT 401 responses before normal auth recovery
- clear active Codex-backend auth state and return an Unauthorized reauth error
- leave generic 401s, malformed bodies, and API-key auth on existing paths

## Testing

- `cargo test -p codex-core workspace_restricted --lib`
- `cargo test -p codex-core generic_or_malformed_401_uses_existing_recovery_path --lib`
- `cargo fmt`
- `git diff --check`
This commit is contained in:
radai
2026-06-18 12:27:40 -07:00
parent 66f0220c56
commit 3ea7067d07
3 changed files with 237 additions and 0 deletions

View File

@@ -140,6 +140,8 @@ pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request";
pub const X_OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str =
"x-responsesapi-include-timing-metrics";
pub const CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE: &str = "chatgpt_ip_workspace_restricted";
const CHATGPT_IP_WORKSPACE_RESTRICTED_MESSAGE: &str = "Your ChatGPT session is no longer authorized for this network or workspace. Please sign in again and choose an allowed workspace.";
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
"x-codex-ws-stream-request-start-ms";
const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str =
@@ -2051,6 +2053,53 @@ async fn handle_unauthorized(
provider: &SharedModelProvider,
) -> Result<UnauthorizedRecoveryExecution> {
let debug = extract_response_debug_context(&transport);
if is_chatgpt_ip_workspace_restricted_unauthorized(&transport)
&& let Some(recovery) = auth_recovery
&& recovery.current_auth_uses_codex_backend()
{
let mode = recovery.mode_name();
let phase = recovery.step_name();
let auth_error_code = Some(
debug
.auth_error_code
.as_deref()
.unwrap_or(CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE),
);
let auth_state_changed = match recovery.force_logout_due_to_server_auth_rejection().await {
Ok(changed) => Some(changed),
Err(err) => {
warn!("failed to clear auth after server auth rejection: {err}");
None
}
};
session_telemetry.record_auth_recovery(
mode,
phase,
"forced_logout",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
auth_error_code,
Some(CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE),
auth_state_changed,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"forced_logout",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
auth_error_code,
);
return Err(CodexErr::RefreshTokenFailed(
codex_protocol::auth::RefreshTokenFailedError::new(
codex_protocol::auth::RefreshTokenFailedReason::Other,
CHATGPT_IP_WORKSPACE_RESTRICTED_MESSAGE,
),
));
}
if let Some(recovery) = auth_recovery
&& recovery.has_next()
{
@@ -2161,6 +2210,29 @@ async fn handle_unauthorized(
Err(provider.map_api_error(ApiError::Transport(transport)))
}
fn is_chatgpt_ip_workspace_restricted_unauthorized(transport: &TransportError) -> bool {
let TransportError::Http { status, body, .. } = transport else {
return false;
};
if *status != StatusCode::UNAUTHORIZED {
return false;
}
body.as_deref()
.is_some_and(|body| body_has_error_code(body, CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE))
}
fn body_has_error_code(body: &str, expected_code: &str) -> bool {
let Ok(value) = serde_json::from_str::<serde_json::Value>(body) else {
return false;
};
value
.get("error")
.and_then(|error| error.get("code"))
.and_then(serde_json::Value::as_str)
.or_else(|| value.get("code").and_then(serde_json::Value::as_str))
== Some(expected_code)
}
fn api_error_http_status(error: &ApiError) -> Option<u16> {
match error {
ApiError::Transport(TransportError::Http { status, .. }) => Some(status.as_u16()),

View File

@@ -17,8 +17,14 @@ use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_api::TransportError;
use codex_app_server_protocol::AuthMode;
use codex_config::types::AuthCredentialsStoreMode;
use codex_login::AuthDotJson;
use codex_login::AuthKeyringBackendKind;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::TokenData;
use codex_login::load_auth_dot_json;
use codex_login::save_auth;
use codex_model_provider::BearerAuthProvider;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
@@ -28,6 +34,7 @@ use codex_model_provider_info::WireApi;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_otel::SessionTelemetry;
use codex_protocol::ThreadId;
use codex_protocol::error::CodexErr;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
@@ -153,6 +160,32 @@ fn test_session_telemetry() -> SessionTelemetry {
)
}
fn test_auth_dot_json() -> AuthDotJson {
AuthDotJson {
auth_mode: Some(AuthMode::Chatgpt),
openai_api_key: None,
tokens: Some(TokenData {
id_token: Default::default(),
access_token: "access-token".to_string(),
refresh_token: "refresh-token".to_string(),
account_id: Some("account_id".to_string()),
}),
last_refresh: Some(chrono::Utc::now()),
agent_identity: None,
personal_access_token: None,
bedrock_api_key: None,
}
}
fn unauthorized_transport(body: &str) -> TransportError {
TransportError::Http {
status: http::StatusCode::UNAUTHORIZED,
url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()),
headers: None,
body: Some(body.to_string()),
}
}
#[derive(Default)]
struct TagCollectorVisitor {
tags: BTreeMap<String, String>,
@@ -566,6 +599,110 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
assert_eq!(auth_context.recovery_phase, Some("refresh_token"));
}
#[tokio::test]
async fn workspace_restricted_401_forces_chatgpt_logout() -> anyhow::Result<()> {
let codex_home = TempDir::new()?;
save_auth(
codex_home.path(),
&test_auth_dot_json(),
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
let auth_manager = AuthManager::from_auth_for_testing_with_home(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
codex_home.path().to_path_buf(),
);
assert!(auth_manager.auth_cached().is_some());
let mut recovery = Some(auth_manager.unauthorized_recovery());
let err = super::handle_unauthorized(
unauthorized_transport(r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#),
&mut recovery,
&test_session_telemetry(),
)
.await
.expect_err("matching 401 should force reauth without retry");
match err {
CodexErr::RefreshTokenFailed(failed) => {
assert_eq!(
failed.message,
"Your ChatGPT session is no longer authorized for this network or workspace. Please sign in again and choose an allowed workspace."
);
}
other => panic!("expected RefreshTokenFailed, got {other:?}"),
}
assert!(auth_manager.auth_cached().is_none());
assert!(
load_auth_dot_json(
codex_home.path(),
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?
.is_none()
);
assert!(!recovery.as_ref().expect("recovery").has_next());
Ok(())
}
#[tokio::test]
async fn workspace_restricted_401_does_not_clear_api_key_auth() {
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("sk-test"));
let mut recovery = Some(auth_manager.unauthorized_recovery());
let err = super::handle_unauthorized(
unauthorized_transport(r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#),
&mut recovery,
&test_session_telemetry(),
)
.await
.expect_err("api key auth should not use ChatGPT forced logout");
assert!(matches!(err, CodexErr::UnexpectedStatus(_)));
assert!(
auth_manager
.auth_cached()
.is_some_and(|auth| auth.is_api_key_auth())
);
}
#[tokio::test]
async fn generic_or_malformed_401_uses_existing_recovery_path() {
for body in [
r#"{"error":{"code":"token_expired"}}"#,
"not json",
r#"{"error":{"message":"missing code"}}"#,
] {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let mut recovery = Some(auth_manager.unauthorized_recovery());
let err = super::handle_unauthorized(
unauthorized_transport(body),
&mut recovery,
&test_session_telemetry(),
)
.await
.expect_err("non-matching 401 should fall back to normal recovery");
assert!(matches!(err, CodexErr::RefreshTokenFailed(_)));
assert!(auth_manager.auth_cached().is_some());
}
}
#[test]
fn workspace_restricted_error_code_parser_accepts_supported_shapes() {
for body in [
r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#,
r#"{"code":"chatgpt_ip_workspace_restricted"}"#,
] {
assert!(super::is_chatgpt_ip_workspace_restricted_unauthorized(
&unauthorized_transport(body)
));
}
}
fn model_client_with_counting_attestation(
include_attestation: bool,
) -> (ModelClient, Arc<AtomicUsize>) {

View File

@@ -1661,6 +1661,17 @@ impl UnauthorizedRecovery {
}
}
pub fn current_auth_uses_codex_backend(&self) -> bool {
self.manager.current_auth_uses_codex_backend()
}
pub async fn force_logout_due_to_server_auth_rejection(&mut self) -> std::io::Result<bool> {
self.step = UnauthorizedRecoveryStep::Done;
self.manager
.force_logout_due_to_server_auth_rejection()
.await
}
pub async fn next(&mut self) -> Result<UnauthorizedRecoveryStepResult, RefreshTokenError> {
if !self.has_next() {
return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new(
@@ -2401,6 +2412,23 @@ impl AuthManager {
Ok(result)
}
pub async fn force_logout_due_to_server_auth_rejection(&self) -> std::io::Result<bool> {
if !self.current_auth_uses_codex_backend() {
return Ok(false);
}
let removal_result = logout_all_stores(
&self.codex_home,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
);
let had_external_auth = self.has_external_auth();
self.clear_external_auth();
let cache_changed = self.set_cached_auth(None);
let removed = removal_result?;
Ok(removed || had_external_auth || cache_changed)
}
pub fn get_api_auth_mode(&self) -> Option<ApiAuthMode> {
if self.has_external_api_key_auth() {
return Some(ApiAuthMode::ApiKey);