Made token refresh code resilient to missing id_token (#5782)
This PR does the following: 1. Changes `try_refresh_token` to handle the case where the endpoint returns a response without an `id_token`. The OpenID spec indicates that this field is optional and clients should not assume it's present. 2. Changes the `attempt_stream_responses` to propagate token refresh errors rather than silently ignoring them. 3. Fixes a typo in a couple of error messages (unrelated to the above, but something I noticed in passing) - "reconnect" should be spelled without a hyphen. This PR does not implement the additional suggestion from @pakrym-oai that we should sign out when receiving `refresh_token_expired` from the refresh endpoint. Leaving this as a follow-on because I'm undecided on whether this should be implemented in `try_refresh_token` or its callers.
This commit is contained in:
@@ -382,14 +382,16 @@ pub fn write_auth_json(auth_file: &Path, auth_dot_json: &AuthDotJson) -> std::io
|
|||||||
|
|
||||||
async fn update_tokens(
|
async fn update_tokens(
|
||||||
auth_file: &Path,
|
auth_file: &Path,
|
||||||
id_token: String,
|
id_token: Option<String>,
|
||||||
access_token: Option<String>,
|
access_token: Option<String>,
|
||||||
refresh_token: Option<String>,
|
refresh_token: Option<String>,
|
||||||
) -> std::io::Result<AuthDotJson> {
|
) -> std::io::Result<AuthDotJson> {
|
||||||
let mut auth_dot_json = try_read_auth_json(auth_file)?;
|
let mut auth_dot_json = try_read_auth_json(auth_file)?;
|
||||||
|
|
||||||
let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default);
|
let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default);
|
||||||
tokens.id_token = parse_id_token(&id_token).map_err(std::io::Error::other)?;
|
if let Some(id_token) = id_token {
|
||||||
|
tokens.id_token = parse_id_token(&id_token).map_err(std::io::Error::other)?;
|
||||||
|
}
|
||||||
if let Some(access_token) = access_token {
|
if let Some(access_token) = access_token {
|
||||||
tokens.access_token = access_token;
|
tokens.access_token = access_token;
|
||||||
}
|
}
|
||||||
@@ -445,7 +447,7 @@ struct RefreshRequest {
|
|||||||
|
|
||||||
#[derive(Deserialize, Clone)]
|
#[derive(Deserialize, Clone)]
|
||||||
struct RefreshResponse {
|
struct RefreshResponse {
|
||||||
id_token: String,
|
id_token: Option<String>,
|
||||||
access_token: Option<String>,
|
access_token: Option<String>,
|
||||||
refresh_token: Option<String>,
|
refresh_token: Option<String>,
|
||||||
}
|
}
|
||||||
@@ -511,6 +513,35 @@ mod tests {
|
|||||||
assert_eq!(auth_dot_json, same_auth_dot_json);
|
assert_eq!(auth_dot_json, same_auth_dot_json);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn refresh_without_id_token() {
|
||||||
|
let codex_home = tempdir().unwrap();
|
||||||
|
let fake_jwt = write_auth_file(
|
||||||
|
AuthFileParams {
|
||||||
|
openai_api_key: None,
|
||||||
|
chatgpt_plan_type: "pro".to_string(),
|
||||||
|
chatgpt_account_id: None,
|
||||||
|
},
|
||||||
|
codex_home.path(),
|
||||||
|
)
|
||||||
|
.expect("failed to write auth file");
|
||||||
|
|
||||||
|
let auth_file = super::get_auth_file(codex_home.path());
|
||||||
|
let updated = super::update_tokens(
|
||||||
|
auth_file.as_path(),
|
||||||
|
None,
|
||||||
|
Some("new-access-token".to_string()),
|
||||||
|
Some("new-refresh-token".to_string()),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("update_tokens should succeed");
|
||||||
|
|
||||||
|
let tokens = updated.tokens.expect("tokens should exist");
|
||||||
|
assert_eq!(tokens.id_token.raw_jwt, fake_jwt);
|
||||||
|
assert_eq!(tokens.access_token, "new-access-token");
|
||||||
|
assert_eq!(tokens.refresh_token, "new-refresh-token");
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn login_with_api_key_overwrites_existing_auth_json() {
|
fn login_with_api_key_overwrites_existing_auth_json() {
|
||||||
let dir = tempdir().unwrap();
|
let dir = tempdir().unwrap();
|
||||||
|
|||||||
@@ -389,9 +389,14 @@ impl ModelClient {
|
|||||||
|
|
||||||
if status == StatusCode::UNAUTHORIZED
|
if status == StatusCode::UNAUTHORIZED
|
||||||
&& let Some(manager) = auth_manager.as_ref()
|
&& let Some(manager) = auth_manager.as_ref()
|
||||||
&& manager.auth().is_some()
|
&& let Some(auth) = auth.as_ref()
|
||||||
|
&& auth.mode == AuthMode::ChatGPT
|
||||||
{
|
{
|
||||||
let _ = manager.refresh_token().await;
|
manager.refresh_token().await.map_err(|err| {
|
||||||
|
StreamAttemptError::Fatal(CodexErr::Fatal(format!(
|
||||||
|
"Failed to refresh ChatGPT credentials: {err}"
|
||||||
|
)))
|
||||||
|
})?;
|
||||||
}
|
}
|
||||||
|
|
||||||
// The OpenAI Responses endpoint returns structured JSON bodies even for 4xx/5xx
|
// The OpenAI Responses endpoint returns structured JSON bodies even for 4xx/5xx
|
||||||
|
|||||||
@@ -1911,7 +1911,7 @@ async fn run_turn(
|
|||||||
// at a seemingly frozen screen.
|
// at a seemingly frozen screen.
|
||||||
sess.notify_stream_error(
|
sess.notify_stream_error(
|
||||||
turn_context.as_ref(),
|
turn_context.as_ref(),
|
||||||
format!("Re-connecting... {retries}/{max_retries}"),
|
format!("Reconnecting... {retries}/{max_retries}"),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
|||||||
@@ -134,7 +134,7 @@ async fn run_compact_task_inner(
|
|||||||
let delay = backoff(retries);
|
let delay = backoff(retries);
|
||||||
sess.notify_stream_error(
|
sess.notify_stream_error(
|
||||||
turn_context.as_ref(),
|
turn_context.as_ref(),
|
||||||
format!("Re-connecting... {retries}/{max_retries}"),
|
format!("Reconnecting... {retries}/{max_retries}"),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
tokio::time::sleep(delay).await;
|
tokio::time::sleep(delay).await;
|
||||||
|
|||||||
@@ -2325,7 +2325,7 @@ fn plan_update_renders_history_cell() {
|
|||||||
fn stream_error_updates_status_indicator() {
|
fn stream_error_updates_status_indicator() {
|
||||||
let (mut chat, mut rx, _op_rx) = make_chatwidget_manual();
|
let (mut chat, mut rx, _op_rx) = make_chatwidget_manual();
|
||||||
chat.bottom_pane.set_task_running(true);
|
chat.bottom_pane.set_task_running(true);
|
||||||
let msg = "Re-connecting... 2/5";
|
let msg = "Reconnecting... 2/5";
|
||||||
chat.handle_codex_event(Event {
|
chat.handle_codex_event(Event {
|
||||||
id: "sub-1".into(),
|
id: "sub-1".into(),
|
||||||
msg: EventMsg::StreamError(StreamErrorEvent {
|
msg: EventMsg::StreamError(StreamErrorEvent {
|
||||||
|
|||||||
Reference in New Issue
Block a user