From 76ceaddb29444388fbbbae07c46f7e8849f9658b Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Tue, 18 Aug 2026 15:54:38 +0000 Subject: [PATCH] Reconnect Guardian sampling WebSockets after auth changes (#39220) ## Why Guardian sampling WebSockets authenticate when the connection is opened. Reusing a pooled connection after credentials change can therefore keep using the old authorization. ## What changed - Track authentication changes on each pooled sampling connection. - Discard stale connections after an auth change and reject connections whose authentication changes while the handshake is in progress. ## Testing - Verify the sampler reconnects with the refreshed bearer token instead of reusing its existing connection. - Verify the installed Guardian extension reconnects after an external auth refresh. GitOrigin-RevId: 6d2e7df776fd21c78be0928f71162f5419a8b8f0 --- .../ext/guardian-v2/src/extension_tests.rs | 122 ++++++++++++++++++ codex-rs/ext/guardian-v2/src/sampler.rs | 18 ++- codex-rs/ext/guardian-v2/src/sampler_tests.rs | 37 +++++- 3 files changed, 171 insertions(+), 6 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 2b367ab762..57d0821aac 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -19,6 +19,9 @@ use codex_features::Feature; use codex_history::RolloutItem; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::ResponseItemId; use codex_protocol::models::ContentItem; @@ -55,6 +58,125 @@ const TEST_GUARDIAN_POLICY: &str = const TEST_CATALOG_GUARDIAN_POLICY: &str = "Require review before sending organization data to third-party services."; +struct RefreshableAuth(std::sync::Mutex<&'static str>); + +impl ExternalAuth for RefreshableAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(CodexAuth::from_api_key(*self.0.lock().expect("auth"))) }) + } + + fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + *self.0.lock().expect("auth") = "refreshed"; + self.resolve() + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> { + skip_if_no_network!(Ok(())); + + let thread_server = responses::start_mock_server().await; + let test = test_codex().build_with_auto_env(&thread_server).await?; + let events = vec![ + ev_assistant_message("sample", r#"{"scores":{"action_risk":0.25}}"#), + ev_completed("response-1"), + ]; + // Keep the sampled connection open for another request so only auth + // invalidation, not a server close, forces the next handshake. + let server = responses::start_websocket_server(vec![ + Vec::new(), + vec![events.clone(), events.clone()], + vec![events], + ]) + .await; + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("original")); + auth_manager + .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new("original")))) + .await?; + let mut config = test.config.clone(); + config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!( + "http://{}/v1", + server.uri().trim_start_matches("ws://") + ))); + config.features.enable(Feature::GuardianV2)?; + let mut builder = ExtensionRegistryBuilder::new(); + crate::install( + &mut builder, + auth_manager.clone(), + Arc::downgrade(&test.thread_manager), + ); + let registry = builder.build(); + let session_store = ExtensionData::new("session-1"); + let thread_store = test.codex.thread_extension_data(); + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store, + }) + .await; + let progress = thread_store + .get::() + .expect("Guardian v2 should initialize"); + let turn_store = ExtensionData::new("turn-1"); + let tool_name = ToolName::plain("read_file"); + let payload = ToolPayload::Function { + arguments: r#"{"path":"README.md"}"#.to_owned(), + }; + + for (call_index, call_id) in [(1, "call-1"), (2, "call-2")] { + if call_index == 2 { + auth_manager.refresh_token_from_authority().await?; + } + registry.tool_lifecycle_contributors()[0] + .on_tool_start(ToolStartInput { + session_store: &session_store, + thread_store, + turn_store: &turn_store, + turn_id: "turn-1", + call_id, + tool_name: &tool_name, + payload: &payload, + conversation_history: Arc::new(TestConversationHistory(Vec::new())), + source: ToolCallSource::Direct, + }) + .await; + tokio::time::timeout(Duration::from_secs(5), async { + while progress.latest_scored_tool_call.load(Ordering::Acquire) < call_index { + tokio::task::yield_now().await; + } + }) + .await?; + } + + assert_eq!( + server + .handshakes() + .iter() + .map(|handshake| handshake.header("authorization")) + .collect::>(), + vec![ + Some("Bearer original".to_owned()), + Some("Bearer original".to_owned()), + Some("Bearer refreshed".to_owned()), + ] + ); + assert_eq!( + server + .connections() + .iter() + .map(Vec::len) + .collect::>(), + vec![0, 1, 1] + ); + Ok(()) +} + struct TestConversationHistory(Vec); impl ConversationHistorySnapshot for TestConversationHistory { diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index 97ccf41623..1953a770f7 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -118,6 +118,7 @@ pub enum LunaSamplerError { struct PooledConnection { connection: ResponsesWebsocketConnection, connected_at: Instant, + auth_changes: Option>, } struct ConnectionLease { @@ -178,6 +179,8 @@ impl LunaSampler { .api_provider() .await .map_err(LunaSamplerError::Provider)?; + let auth_manager = self.config.provider.auth_manager(); + let auth_changes = auth_manager.map(|manager| manager.auth_change_receiver()); let auth = self .config .provider @@ -244,10 +247,19 @@ impl LunaSampler { .await .map_err(|_| LunaSamplerError::ConnectionTimeout)? .map_err(LunaSamplerError::Api)?; + if auth_changes + .as_ref() + .is_some_and(|auth| auth.has_changed().unwrap_or(true)) + { + return Err(LunaSamplerError::Api(ApiError::Stream( + "authentication changed while connecting".into(), + ))); + } Ok(PooledConnection { connection, connected_at: Instant::now(), + auth_changes, }) } @@ -264,7 +276,11 @@ impl LunaSampler { .pop(); match idle { Some(connection) - if connection.connected_at.elapsed() < MAX_WEBSOCKET_AGE + if connection + .auth_changes + .as_ref() + .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) + && connection.connected_at.elapsed() < MAX_WEBSOCKET_AGE && !connection.connection.is_closed().await => { break connection; diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index 3a73acdbc8..08fb38c719 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -4,6 +4,9 @@ use codex_http_client::OutboundProxyPolicy; use codex_login::AgentIdentityAuthPolicy; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::ResponseItemId; @@ -92,6 +95,17 @@ fn sample_request(turn_id: &str) -> LunaSamplingRequest { } } +struct RefreshableAuth(std::sync::Mutex<&'static str>); +impl ExternalAuth for RefreshableAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(CodexAuth::from_api_key(*self.0.lock().expect("auth"))) }) + } + fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + *self.0.lock().expect("auth") = "refreshed"; + self.resolve() + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requests() -> Result<()> { @@ -111,13 +125,19 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ ], ]; let idle_server = responses::start_websocket_server(vec![scripted_requests.clone()]).await; + let refreshed = + responses::start_websocket_server(vec![vec![scripted_requests[1].clone()]]).await; let server = responses::start_websocket_server(vec![scripted_requests]).await; - let base_url = proxy_websocket_servers(&[&idle_server, &server]).await?; + let base_url = proxy_websocket_servers(&[&idle_server, &server, &refreshed]).await?; + let manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-api-key")); + manager + .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new( + "test-api-key", + )))) + .await?; let provider = create_model_provider( ModelProviderInfo::create_openai_provider(Some(base_url)), - Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( - "test-api-key", - ))), + Some(manager.clone()), ); let sampler = LunaSampler::connect(LunaSamplerConfig { @@ -193,6 +213,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ } }) .await?; + manager.refresh_token_from_authority().await?; let second = sampler .sample(LunaSamplingRequest { instructions: "Return a risk score.".to_owned(), @@ -208,7 +229,13 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ assert_eq!(first, r#"{"score":0.25}"#); assert_eq!(second, r#"{"score":0.75}"#); - let requests = server.single_connection(); + let mut requests = server.single_connection(); + assert_eq!(requests.len(), 1); + assert_eq!( + refreshed.single_handshake().header("authorization"), + Some("Bearer refreshed".to_owned()) + ); + requests.extend(refreshed.single_connection()); assert_eq!(requests.len(), 2); assert_eq!( requests[0].body_json()["input"][2]["content"],