From 6bd17ba6ae6ca2788026bd789f28f38b754d7379 Mon Sep 17 00:00:00 2001 From: Won Park Date: Mon, 9 Nov 2026 19:38:16 +0000 Subject: [PATCH] Revalidate request auth ownership after budget checks (#60436) ## Why Account changes during asynchronous request setup can invalidate the credentials selected for a generation request. ## What changed - Restart HTTP request setup if the auth owner changes during the budget check. - Validate the WebSocket connection's retained auth owner with `AuthOwner::is_current()` after budget checks. ## Testing Add regression coverage for account switches during HTTP budget checks and external credential changes before WebSocket request setup, in both dual-header and assertion-only runtime auth modes. GitOrigin-RevId: 7a1e1475a40aa2d34cd94f60f729ad697f3f7028 --- codex-rs/core/src/client.rs | 15 ++- .../suite/client_websockets_runtime_auth.rs | 103 ++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 86310dd789..425967e392 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2069,6 +2069,14 @@ impl ResponsesHttpExecutor<'_> { let kind = codex_extension_api::ModelRequestKind::Generation; crate::model_request::check_budget(&self.client.request_contributors, &thread_id, kind) .await?; + // An account switch during the check requires fresh request setup. + if client_setup + .auth_owner + .as_ref() + .is_some_and(|owner| !owner.is_current()) + { + continue; + } let interceptors = crate::model_request::prepare( &self.client.request_contributors, &thread_id, @@ -2298,7 +2306,12 @@ impl ModelClientSession { ) .await?; // An account switch during the check requires fresh request setup and a new socket. - if self.websocket_session.auth_owner != self.client.auth_owner() { + if self + .websocket_session + .auth_owner + .as_ref() + .is_some_and(|owner| !owner.is_current()) + { continue; } } diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 2c3620480d..75630ab6e4 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -312,6 +312,66 @@ async fn http_requests_keep_matching_runtime_and_credentials_across_retries_and_ ); } +#[test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn http_reprepares_request_when_account_switches_during_budget_check(mode: RuntimeAuthMode) { + skip_if_no_network!(); + let server = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(registered) + .expect(2) + .mount(&server) + .await; + Mock::given(path("/v1/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(responses::sse(vec![ + ev_response_created("resp-1"), + ev_completed("resp-1"), + ])), + ) + .expect(1) + .mount(&server) + .await; + let mut provider = + ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); + provider.supports_websockets = false; + let budget = Arc::new(BudgetChecks::default()); + let (harness, credentials) = + runtime_harness(&server, provider, mode, vec![budget.clone()]).await; + let mut session = harness.client.new_session(); + let gate = Arc::new(Barrier::new(2)); + budget.0.lock().unwrap().push_back(Some(gate.clone())); + let generation = complete(&mut session, &harness); + tokio::pin!(generation); + tokio::select! { + result = &mut generation => panic!("generation finished before the budget check: {result:?}"), + _ = gate.wait() => {} + } + *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); + harness.auth_manager.reload().await; + let (result, _) = tokio::join!(generation, gate.wait()); + result.unwrap(); + + let requests = server.received_requests().await.unwrap(); + let responses: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/v1/responses") + .collect(); + assert_eq!(responses.len(), 1); + assertion_payload( + responses[0].headers["authorization"].to_str().unwrap(), + responses[0] + .headers + .get("x-openai-agent-assertion") + .map(|value| value.to_str().unwrap().to_string()), + mode, + "key-other-user", + ); +} + #[test_case(RuntimeAuthMode::DualHeader, false; "dual_header")] #[test_case(RuntimeAuthMode::AssertionOnly, false; "assertion_only")] #[test_case(RuntimeAuthMode::DualHeader, true; "dual_header_with_gateway")] @@ -469,6 +529,49 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( server.shutdown().await; } +#[test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_sends_after_external_auth_changes_during_request_setup(mode: RuntimeAuthMode) { + skip_if_no_network!(); + let registration = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .expect(2) + .mount(®istration) + .await; + let response = vec![ev_response_created("resp-1"), ev_completed("resp-1")]; + let server = start_websocket_server(vec![vec![response.clone()], vec![response]]).await; + let provider = websocket_provider(&server); + let budget = Arc::new(BudgetChecks::default()); + let (harness, credentials) = runtime_harness(®istration, provider, mode, vec![budget]).await; + let mut session = harness.client.new_session(); + complete(&mut session, &harness).await.unwrap(); + + *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); + complete(&mut session, &harness).await.unwrap(); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + for (handshake, credential) in handshakes.iter().zip(["key-a", "key-other-user"]) { + assertion_payload( + &handshake.header("authorization").unwrap(), + handshake.header("x-openai-agent-assertion"), + mode, + credential, + ); + } + assert_eq!( + server + .connections() + .iter() + .map(Vec::len) + .collect::>(), + vec![1, 1] + ); + server.shutdown().await; +} + #[test_case(RuntimeAuthMode::DualHeader; "dual_header")] #[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] #[tokio::test]