Fix login for internal employees (#2528)

This PR:
- fixes for internal employee because we currently want to prefer SIWC
for them.
- fixes retrying forever on unauthorized access. we need to break
eventually on max retries.
This commit is contained in:
Ahmed Ibrahim
2025-08-20 14:05:20 -07:00
committed by GitHub
parent 0d12380c3b
commit c579ae41ae
3 changed files with 25 additions and 9 deletions

View File

@@ -248,12 +248,10 @@ impl ModelClient {
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<u64>().ok());
if status == StatusCode::UNAUTHORIZED {
if let Some(a) = auth.as_ref() {
let _ = a.refresh_token().await;
}
// Retry immediately with refreshed credentials.
continue;
if status == StatusCode::UNAUTHORIZED
&& let Some(a) = auth.as_ref()
{
let _ = a.refresh_token().await;
}
// The OpenAI Responses endpoint returns structured JSON bodies even for 4xx/5xx
@@ -263,7 +261,10 @@ impl ModelClient {
// exact error message (e.g. "Unknown parameter: 'input[0].metadata'"). The body is
// small and this branch only runs on error paths so the extra allocation is
// negligible.
if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) {
if !(status == StatusCode::TOO_MANY_REQUESTS
|| status == StatusCode::UNAUTHORIZED
|| status.is_server_error())
{
// Surface the error body to callers. Use `unwrap_or_default` per Clippy.
let body = res.text().await.unwrap_or_default();
return Err(CodexErr::UnexpectedStatus(status, body));

View File

@@ -242,7 +242,7 @@ fn load_auth(
// "refreshable" even if we are using the API key for auth?
match &tokens {
Some(tokens) => {
if tokens.should_use_api_key(preferred_auth_method) {
if tokens.should_use_api_key(preferred_auth_method, tokens.is_openai_email()) {
return Ok(Some(CodexAuth::from_api_key(api_key)));
} else {
// Ignore the API key and fall through to ChatGPT auth.

View File

@@ -25,16 +25,31 @@ pub struct TokenData {
impl TokenData {
/// Returns true if this is a plan that should use the traditional
/// "metered" billing via an API key.
pub(crate) fn should_use_api_key(&self, preferred_auth_method: AuthMode) -> bool {
pub(crate) fn should_use_api_key(
&self,
preferred_auth_method: AuthMode,
is_openai_email: bool,
) -> bool {
if preferred_auth_method == AuthMode::ApiKey {
return true;
}
// If the email is an OpenAI email, use AuthMode::ChatGPT unless preferred_auth_method is AuthMode::ApiKey.
if is_openai_email {
return false;
}
self.id_token
.chatgpt_plan_type
.as_ref()
.is_none_or(|plan| plan.is_plan_that_should_use_api_key())
}
pub fn is_openai_email(&self) -> bool {
self.id_token
.email
.as_deref()
.is_some_and(|email| email.trim().to_ascii_lowercase().ends_with("@openai.com"))
}
}
/// Flat subset of useful claims in id_token from auth.json.