From 6eb51a438ac95769e2d9d04df048e72b9b2626e0 Mon Sep 17 00:00:00 2001 From: alexsong-oai Date: Fri, 6 Nov 2026 02:54:40 +0000 Subject: [PATCH] Retry gateway credential exports after unauthorized responses (#59816) ## Why Cached gateway command credentials need to be refreshed when the gateway rejects them with a `401`. ## What changed Clear the cached `credential_export` token after an inference HTTP `401` or a WebSocket upgrade rejected with `401`, then rerun the command for one retry. Preserve subsequent primary authentication recovery if the retry is also unauthorized, and document the behavior. ## Testing Add integration coverage for using a new exported token, reporting an invalid export, recovering primary authentication after the gateway retry, and stopping after a second unauthorized WebSocket upgrade. GitOrigin-RevId: 3397476dbfd3bec1bd44f90657c4f9e17bb7df67 --- codex-rs/app-server/README.md | 4 +- .../core/tests/suite/client_websockets.rs | 76 ++++++++++++++++--- codex-rs/core/tests/suite/external_auth.rs | 62 +++++++++++++-- codex-rs/core/tests/suite/gateway_auth.rs | 41 +++++++++- codex-rs/model-provider/src/gateway_auth.rs | 11 +++ codex-rs/model-provider/src/gateway_export.rs | 4 + codex-rs/model-provider/src/provider.rs | 11 +++ 7 files changed, 188 insertions(+), 21 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 57d4742b64..fbea4730d4 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -638,7 +638,9 @@ cookie. Choose either `credential_export` or OAuth settings such as `authorization_url`, `token_url`, and `client_id`. If the command fails, the request returns the export error. Use direct inference and catalog URLs because command credentials require same-origin catalog -access and reject HTTP redirects. +access and reject HTTP redirects. After an inference HTTP 401 or a rejected +WebSocket upgrade with status 401, Codex clears the cached export and reruns +the command once for the retry. Clients that omit the capability or set it to `false`, including the TUI, retain automatic browser authorization after initialization. Startup credential reads diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 560d5a504c..620bd2f231 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2828,6 +2828,56 @@ fn websocket_provider_with_connect_timeout( } } +#[tokio::test] +async fn command_gateway_stops_after_a_second_unauthorized_websocket_upgrade() { + skip_if_no_network!(); + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::header("upgrade", "websocket")) + .and(wiremock::matchers::header("cookie", "gateway=exported")) + .respond_with(wiremock::ResponseTemplate::new(401)) + .expect(2) + .mount(&server) + .await; + let command = if cfg!(windows) { + "command = 'cmd.exe', args = ['/D', '/C', 'echo', 'exported']" + } else { + "command = 'echo', args = ['exported']" + }; + let provider = toml::from_str(&format!( + r#" +name = "gateway" +base_url = "{}/v1" +supports_websockets = true +[gateway_oauth] +delivery = {{ kind = "cookie", name = "gateway" }} +credential_export = {{ {command} }} +"#, + server.uri() + )) + .unwrap(); + let harness = websocket_harness_with_provider_options_and_auth( + provider, + /*runtime_metrics_enabled*/ false, + /*concurrent_reasoning_summaries_enabled*/ false, + /*enabled_features*/ &[], + Some(CodexAuth::from_api_key("primary")), + ) + .await; + harness + .client + .new_session() + .preconnect_websocket( + &harness.model_info, + /*service_tier*/ None, + &harness.session_telemetry, + &websocket_connection_metadata(&harness), + ) + .await + .expect_err("a rejected replay ends recovery"); + server.verify().await; +} + async fn websocket_harness(server: &WebSocketTestServer) -> WebsocketTestHarness { websocket_harness_with_runtime_metrics(server, /*runtime_metrics_enabled*/ false).await } @@ -2976,17 +3026,21 @@ async fn websocket_harness_with_provider_factory( .with_metrics(metrics); let effort = None; let summary = ReasoningSummary::Auto; - let _gateway = provider.gateway_oauth.as_ref().map(|_| { - codex_model_provider::test_support::seed_gateway_auth( - &provider, - &auth_manager, - serde_json::json!({ - "access_token": "gateway-token", - "refresh_token": "refresh-token", - "expires_at": i64::MAX, - }), - ) - }); + let _gateway = provider + .gateway_oauth + .as_ref() + .filter(|config| !config.uses_credential_export()) + .map(|_| { + codex_model_provider::test_support::seed_gateway_auth( + &provider, + &auth_manager, + serde_json::json!({ + "access_token": "gateway-token", + "refresh_token": "refresh-token", + "expires_at": i64::MAX, + }), + ) + }); let client = ModelClient::new( provider_factory(provider.clone(), client_auth_manager), AgentIdentityAuthPolicy::JwtOnly, diff --git a/codex-rs/core/tests/suite/external_auth.rs b/codex-rs/core/tests/suite/external_auth.rs index 5a77f6e992..b63d09aaa4 100644 --- a/codex-rs/core/tests/suite/external_auth.rs +++ b/codex-rs/core/tests/suite/external_auth.rs @@ -758,14 +758,7 @@ async fn external_auth_401_retry_uses_refreshed_chatgpt_headers() -> anyhow::Res let mut builder = test_codex().with_auth(CodexAuth::from_api_key("seed")); let test = builder.build_with_auto_env(&server).await?; - let external_auth = Arc::new(ScriptedExternalAuth { - current: Mutex::new(external_chatgpt_auth(INITIAL_ACCESS_TOKEN)?), - refreshed: external_chatgpt_auth(REFRESHED_ACCESS_TOKEN)?, - }); - test.thread_manager - .auth_manager() - .set_external_auth(external_auth.clone()) - .await?; + install_scripted_external_auth(&test).await?; test.submit_turn("hello").await?; @@ -790,3 +783,56 @@ async fn external_auth_401_retry_uses_refreshed_chatgpt_headers() -> anyhow::Res ); Ok(()) } + +async fn install_scripted_external_auth( + test: &core_test_support::test_codex::TestCodex, +) -> anyhow::Result<()> { + let external_auth = Arc::new(ScriptedExternalAuth { + current: Mutex::new(external_chatgpt_auth(INITIAL_ACCESS_TOKEN)?), + refreshed: external_chatgpt_auth(REFRESHED_ACCESS_TOKEN)?, + }); + test.thread_manager + .auth_manager() + .set_external_auth(external_auth) + .await?; + Ok(()) +} + +#[tokio::test] +async fn command_gateway_retry_then_recovers_refreshable_primary_auth() -> anyhow::Result<()> { + let server = start_mock_server().await; + let source = tempfile::TempDir::new()?; + let token_path = source.path().join("gateway-token"); + std::fs::write(&token_path, "gateway-token")?; + let test = super::gateway_auth::command_gateway(&server, &token_path).await?; + install_scripted_external_auth(&test).await?; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(header("cookie", "gateway=gateway-token")) + .and(header( + "authorization", + format!("Bearer {INITIAL_ACCESS_TOKEN}"), + )) + .respond_with(ResponseTemplate::new(401)) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(header("cookie", "gateway=gateway-token")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ev_completed("resp-1")])), + ) + .expect(1) + .mount(&server) + .await; + test.submit_turn("hello").await?; + server.verify().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/gateway_auth.rs b/codex-rs/core/tests/suite/gateway_auth.rs index 577cb63d57..f42d8eb522 100644 --- a/codex-rs/core/tests/suite/gateway_auth.rs +++ b/codex-rs/core/tests/suite/gateway_auth.rs @@ -41,7 +41,10 @@ fn gateway_export_settings(token_path: &std::path::Path) -> Result { )) } -async fn command_gateway(server: &MockServer, token_path: &std::path::Path) -> Result { +pub(super) async fn command_gateway( + server: &MockServer, + token_path: &std::path::Path, +) -> Result { let provider: ModelProviderInfo = toml::from_str(&format!( r#" name = "gateway" @@ -60,6 +63,42 @@ requires_openai_auth = true .await } +#[test_case::test_case("second", None; "uses_new_export")] +#[test_case::test_case("", Some("Gateway credential export must print one access token"); "reports_invalid_export")] +#[tokio::test] +async fn command_gateway_reexports_after_401( + next_export: &'static str, + expected_error: Option<&'static str>, +) -> Result<()> { + let server = MockServer::start().await; + let source = TempDir::new()?; + let token_path = source.path().join("access-token"); + std::fs::write(&token_path, "first")?; + let test = command_gateway(&server, &token_path).await?; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(wiremock::matchers::header("cookie", "gateway=first")) + .respond_with(move |_: &wiremock::Request| { + std::fs::write(&token_path, next_export).expect("change the export for the retry"); + ResponseTemplate::new(401) + }) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(wiremock::matchers::header("cookie", "gateway=second")) + .respond_with(responses::sse_response(responses::sse(vec![ + responses::ev_completed("response"), + ]))) + .expect(u64::from(expected_error.is_none())) + .mount(&server) + .await; + assert_eq!(run_gateway_turn(&test).await?.as_deref(), expected_error); + server.verify().await; + Ok(()) +} + #[tokio::test] async fn command_gateway_authenticates_without_an_oauth_grant_and_caches_the_token() -> Result<()> { let server = MockServer::start().await; diff --git a/codex-rs/model-provider/src/gateway_auth.rs b/codex-rs/model-provider/src/gateway_auth.rs index 79c2661180..81ce4e0243 100644 --- a/codex-rs/model-provider/src/gateway_auth.rs +++ b/codex-rs/model-provider/src/gateway_auth.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use crate::ProviderUnauthorizedRecovery; use crate::gateway_export::GatewayExport; use codex_login::AuthManager; use codex_login::GatewayAuthError; @@ -46,6 +47,16 @@ impl GatewayAuthManager { Ok(Self(source)) } + pub(crate) async fn prepare_retry(&self) -> ProviderUnauthorizedRecovery { + match &self.0 { + GatewayCredentialSource::Export(export) => { + export.invalidate().await; + ProviderUnauthorizedRecovery::Recovered + } + GatewayCredentialSource::OAuth(_) => ProviderUnauthorizedRecovery::NotConfigured, + } + } + /// The session used for explicit sign-in and OAuth status, when configured. pub fn oauth_session(&self) -> Option> { match &self.0 { diff --git a/codex-rs/model-provider/src/gateway_export.rs b/codex-rs/model-provider/src/gateway_export.rs index 48132abe64..e062c1246f 100644 --- a/codex-rs/model-provider/src/gateway_export.rs +++ b/codex-rs/model-provider/src/gateway_export.rs @@ -30,6 +30,10 @@ impl GatewayExport { } } + pub(crate) async fn invalidate(&self) { + *self.cached.lock().await = None; + } + #[expect( clippy::await_holding_invalid_type, reason = "concurrent cache misses must share a single credential export" diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index c562d6e40d..daf14c8ddd 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -561,6 +561,17 @@ impl ModelProvider for ConfiguredModelProvider { .map_err(std::io::Error::other) } + fn recover_from_unauthorized( + &self, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + Box::pin(async { + let Some(gateway) = self.gateway_auth_manager()? else { + return Ok(ProviderUnauthorizedRecovery::NotConfigured); + }; + Ok(gateway.prepare_retry().await) + }) + } + fn supports_attestation(&self) -> bool { self.auth_manager .as_ref()