From 2c502d8b3d382dd609179164f31396d95deee55c Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 6 Nov 2026 13:38:14 +0000 Subject: [PATCH] Bind Guardian Decisions requests to their credential source (#59914) ## Why Guardian Decisions retained a copied login key after the saved login changed or was removed. Queued comparisons need to respect credential changes, while runtime-bound threads must remain tied to their startup credential source. ## What changed - Select login credentials from the thread's `ProviderAuthContext`, falling back to the supplied auth manager when no context exists. - Resolve saved login keys before sending requests and reject queued comparisons if their credential source has changed. Pin runtime-bound credentials to the source selected at startup. - Revalidate implicit keys before sending, rejecting missing, padded, or FedRAMP keys. Preserve explicit key overrides and the environment-key fallback. ## Testing Extend credential lookup and queued-request tests to cover reloaded logins, runtime source pinning, rejection of stale queued comparisons, replacement-key authorization, invalid keys, and logout. GitOrigin-RevId: 1d975613040e08c3825ba6355a166cf0fa812cad --- .../guardian-v2/src/async_scorer/decisions.rs | 70 ++++++++++++++++--- .../src/async_scorer/decisions_tests.rs | 69 ++++++++++++++---- .../guardian-v2/src/async_scorer/startup.rs | 39 ++++++++--- 3 files changed, 144 insertions(+), 34 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/decisions.rs b/codex-rs/ext/guardian-v2/src/async_scorer/decisions.rs index d39203478e..e68ccfc710 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/decisions.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/decisions.rs @@ -3,16 +3,19 @@ //! Comparisons omit reasoning and compaction. //! Other unsupported evidence is rejected intact. //! Errors never contain credentials or wire bodies. +//! Queued comparisons retain their credential source; runtime-bound threads pin it at startup. //! Evidence preserves message roles, content order, and boundaries. //! The caller records measurements after baseline publication; dropping its task aborts Decisions work. use codex_history::ResponseItemEnvelope; use codex_http_client::HttpClient; +use codex_login::AuthManager; use codex_protocol::models::ContentItem; use codex_protocol::models::ImageReference; use codex_protocol::models::ResponseItem; use serde_json::Value; use serde_json::json; +use std::borrow::Cow; use std::collections::VecDeque; use std::sync::Arc; use std::sync::Mutex; @@ -62,10 +65,49 @@ pub(super) enum DecisionsError { InvalidResponse, } +pub(super) enum DecisionsCredential { + ApiKey(String), + SavedLogin(Arc), + RuntimeLogin(Arc, u64), +} + +impl From<&str> for DecisionsCredential { + fn from(key: &str) -> Self { + Self::ApiKey(key.to_owned()) + } +} + +impl DecisionsCredential { + fn source_generation(&self) -> Option { + match self { + Self::ApiKey(_) => None, + Self::RuntimeLogin(_, source) => Some(*source), + Self::SavedLogin(manager) => { + let changes = manager.auth_change_state_receiver(); + Some(changes.borrow().source_generation) + } + } + } + + fn resolve(&self, source_generation: Option) -> Option> { + match self { + Self::ApiKey(key) => Some(Cow::Borrowed(key)), + Self::SavedLogin(manager) | Self::RuntimeLogin(manager, _) => { + let auth = manager.auth_cached()?; + let source = manager.auth_change_state_for(&auth)?.source_generation; + let key = auth.api_key()?; + (Some(source) == source_generation + && super::startup::valid_implicit_decisions_key(key)) + .then(|| Cow::Owned(key.to_owned())) + } + } + } +} + pub(super) struct DecisionsSampler { client: HttpClient, // No Debug implementation: credentials must not reach diagnostics. - api_key: String, + api_key: DecisionsCredential, url: String, slots: Semaphore, active_requests: Mutex>, @@ -102,10 +144,10 @@ impl DecisionsTask { impl DecisionsSampler { pub(super) fn new( client: HttpClient, - api_key: String, + api_key: DecisionsCredential, url: String, ) -> Result { - if api_key.trim().is_empty() { + if matches!(&api_key, DecisionsCredential::ApiKey(key) if key.trim().is_empty()) { return Err(DecisionsError::Credentials); } Ok(Self { @@ -123,6 +165,7 @@ impl DecisionsSampler { input: &[ResponseItemEnvelope], ) -> DecisionsTask { let started = Instant::now(); + let source_generation = self.api_key.source_generation(); let body = (|| { let [tools, instructions, evidence @ ..] = input else { return Err(DecisionsError::UnexpectedRequestSetup); @@ -145,13 +188,14 @@ impl DecisionsSampler { } request_body(rubric, evidence.iter().map(|item| &item.item)) })(); - self.spawn_request(body, started) + self.spawn_request(body, started, source_generation) } fn spawn_request( self: &Arc, body: Result, started: Instant, + source_generation: Option, ) -> DecisionsTask { let body = match body { Ok(body) => body, @@ -177,7 +221,7 @@ impl DecisionsSampler { let Ok(_permit) = sampler.slots.acquire().await else { return (Err(DecisionsError::Transport), started.elapsed()); }; - let result = tokio::time::timeout(DEADLINE, sampler.request(body)) + let result = tokio::time::timeout(DEADLINE, sampler.request(body, source_generation)) .await .unwrap_or(Err(DecisionsError::Timeout)); (result, started.elapsed()) @@ -189,12 +233,16 @@ impl DecisionsSampler { } } - async fn request(&self, body: Value) -> Result<&'static str, DecisionsError> { - let request = self - .client - .post(&self.url) - .bearer_auth(&self.api_key) - .json(&body); + async fn request( + &self, + body: Value, + source_generation: Option, + ) -> Result<&'static str, DecisionsError> { + let key = self + .api_key + .resolve(source_generation) + .ok_or(DecisionsError::Credentials)?; + let request = self.client.post(&self.url).bearer_auth(key).json(&body); drop(body); let mut response = request .send() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs index 512a0ffd9c..ca910ef77f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs @@ -1,14 +1,16 @@ -use super::super::startup::decisions_api_key; +use super::super::startup::decisions_api_key as select_api_key; use super::*; use crate::async_scorer::sampler::LunaSampler; use codex_context_fragments::AnnotatedContent; use codex_context_fragments::ContextualUserFragment; use codex_context_fragments::RenderedFragment; +use codex_features::Feature::RuntimeAgentInstances; use codex_guardian_context::TrustedTool; use codex_guardian_context::TrustedToolSource; use codex_http_client::HttpClientBuilder; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_model_provider::ProviderAuthContext; use codex_protocol::models::ContentItemKind; use codex_protocol::models::ContentItemMetadata; use codex_protocol_test_support::ContentItemTestExt; @@ -97,6 +99,21 @@ fn adapter_preserves_trusted_tool_authority_and_scope() { ); } +fn decisions_api_key( + provider_id: &str, + base_url: Option<&str>, + manager: &Arc, +) -> Result { + select_api_key(provider_id, base_url, /*context*/ None, manager).and_then(current_key) +} + +fn current_key(credential: DecisionsCredential) -> Result { + credential + .resolve(credential.source_generation()) + .map(Cow::into_owned) + .ok_or(DecisionsError::Credentials) +} + #[tokio::test] async fn decisions_key_lookup_honors_override_or_reloaded_login() { let home = tempfile::tempdir().unwrap(); @@ -104,6 +121,18 @@ async fn decisions_key_lookup_honors_override_or_reloaded_login() { codex_login::CodexAuth::from_api_key("old-account-key"), home.path().to_path_buf(), ); + let mut config = core_test_support::load_default_config_for_test(&home).await; + config.features.enable(RuntimeAgentInstances).unwrap(); + let runtime = + codex_core::build_runtime_auth_manager(&config, manager.clone(), "installation").unwrap(); + let mut context = ProviderAuthContext { + credentials: manager.clone(), + runtime_auth: Some((runtime, config.runtime_auth_mode)), + }; + let harness = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("harness-key")); + let pinned = select_api_key("openai", /*base_url*/ None, Some(&context), &harness); + context.runtime_auth = None; + let live = select_api_key("openai", /*base_url*/ None, Some(&context), &harness); codex_login::login_with_api_key( home.path(), "synthetic-key", @@ -112,16 +141,13 @@ async fn decisions_key_lookup_honors_override_or_reloaded_login() { ) .unwrap(); manager.reload().await; - let default = - super::super::startup::decisions_api_key("openai", /*base_url*/ None, &manager); - let custom_url = super::super::startup::decisions_api_key( - "openai", - Some("https://gateway.example/v1"), - &manager, - ); + let default = live.and_then(current_key); + let custom_url = decisions_api_key("openai", Some("https://gateway.example/v1"), &manager); let expected = match std::env::var("CODEX_GUARDIAN_DECISIONS_API_KEY") { Ok(key) => (Ok(key.clone()), Ok(key)), Err(std::env::VarError::NotPresent) => { + let pinned = pinned.unwrap(); + assert!(pinned.resolve(pinned.source_generation()).is_none()); (Ok("synthetic-key".into()), Err(DecisionsError::Credentials)) } Err(std::env::VarError::NotUnicode(_)) => ( @@ -292,22 +318,27 @@ async fn http_contract_and_untrusted_response_validation() { .mount(&server) .await; Arc::get_mut(&mut sampler).unwrap().url = format!("{}/denied", server.uri()); - assert_eq!(sampler.request(body).await, Err(DecisionsError::Http(403))); + let result = sampler.request(body, /*source_generation*/ None).await; + assert_eq!(result, Err(DecisionsError::Http(403))); } #[tokio::test] async fn decisions_admits_newest_request_by_cancelling_oldest() { core_test_support::skip_if_no_network!(); + let home = tempfile::tempdir().unwrap(); + let initial = CodexAuth::from_api_key("old-key"); + let manager = AuthManager::from_auth_for_testing_with_home(initial, home.path().into()); let server = MockServer::start().await; Mock::given(method("POST")) + .and(header("authorization", "Bearer replacement-key")) .respond_with(ResponseTemplate::new(/*s*/ 403)) - .expect(MAX_CONCURRENT_REQUESTS as u64) + .expect((MAX_CONCURRENT_REQUESTS - 1) as u64) .mount(&server) .await; let sampler = Arc::new( DecisionsSampler::new( HttpClientBuilder::new().build_direct().unwrap(), - "synthetic-key".into(), + DecisionsCredential::SavedLogin(Arc::clone(&manager)), server.uri(), ) .unwrap(), @@ -330,19 +361,33 @@ async fn decisions_admits_newest_request_by_cancelling_oldest() { .await .unwrap(); let oldest = sampler.start_conversation_comparison(&input); - let mut remaining = (1..MAX_CONCURRENT_REQUESTS) + let stale = sampler.start_conversation_comparison(&input); + let store = codex_login::AuthCredentialsStoreMode::File; + let keyring = codex_login::AuthKeyringBackendKind::default(); + let login = |key| codex_login::login_with_api_key(home.path(), key, store, keyring).unwrap(); + login("replacement-key"); + manager.reload().await; + let mut remaining = (2..MAX_CONCURRENT_REQUESTS) .map(|_| sampler.start_conversation_comparison(&input)) .collect::>(); remaining.push(sampler.start_conversation_comparison(&input)); assert!(oldest.finish().await.unwrap_err().is_cancelled()); assert!(server.received_requests().await.unwrap().is_empty()); drop(permits); + let rejected = Err(DecisionsError::Credentials); + assert_eq!(stale.finish().await.unwrap().0, rejected); for task in remaining { assert_eq!( task.finish().await.unwrap().0, Err(DecisionsError::Http(403)) ); } + let compare = || sampler.start_conversation_comparison(&input); + login(" sk-proj-fedramp-synthetic"); + manager.reload().await; + assert_eq!(compare().finish().await.unwrap().0, rejected); + manager.logout().await.unwrap(); + assert_eq!(compare().finish().await.unwrap().0, rejected); } #[test] diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs index 9f90e8c14f..3e6c30cf92 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs @@ -11,8 +11,9 @@ use codex_extension_api::ThreadStartInput; use codex_features::Feature; use codex_login::AgentIdentityAuthPolicy; use codex_login::AuthManager; -use codex_login::CodexAuth; +use codex_model_provider::ProviderAuthContext; +use super::decisions::DecisionsCredential; use super::decisions::DecisionsError; use super::decisions::DecisionsSampler; use super::decisions::URL; @@ -22,12 +23,13 @@ use super::sampler::MODEL; // Decisions setup failures are telemetry only; they never affect the baseline sampler. pub(super) fn decisions_sampler( input: &ThreadStartInput<'_, Config>, - auth_manager: &AuthManager, + auth_manager: &Arc, ) -> Option { use codex_http_client::ClientRouteClass; let result = decisions_api_key( &input.config.model_provider_id, input.config.model_provider.base_url.as_deref(), + input.thread_store.get::().as_deref(), auth_manager, ) .and_then(|key| { @@ -56,18 +58,33 @@ pub(super) fn decisions_sampler( pub(super) fn decisions_api_key( provider_id: &str, base_url: Option<&str>, - auth_manager: &AuthManager, -) -> Result { + context: Option<&ProviderAuthContext>, + fallback: &Arc, +) -> Result { match std::env::var("CODEX_GUARDIAN_DECISIONS_API_KEY") { - Ok(key) => Ok(key), + Ok(key) => Ok(DecisionsCredential::ApiKey(key)), Err(std::env::VarError::NotPresent) if provider_id == "openai" && base_url.is_none() => { - auth_manager - .auth_cached() - .as_ref() - .and_then(CodexAuth::api_key) - .map(str::to_owned) - .or_else(codex_login::read_openai_api_key_from_env) + let manager = context.map_or(fallback, |context| &context.credentials); + if let Some(auth) = manager.auth_cached() + && let Some(key) = auth.api_key() + { + if !valid_implicit_decisions_key(key) { + return Err(DecisionsError::Credentials); + } + if context.is_some_and(|context| context.runtime_auth.is_some()) { + let state = manager + .auth_change_state_for(&auth) + .ok_or(DecisionsError::Credentials)?; + return Ok(DecisionsCredential::RuntimeLogin( + Arc::clone(manager), + state.source_generation, + )); + } + return Ok(DecisionsCredential::SavedLogin(Arc::clone(manager))); + } + codex_login::read_openai_api_key_from_env() .filter(|key| valid_implicit_decisions_key(key)) + .map(DecisionsCredential::ApiKey) .ok_or(DecisionsError::Credentials) } Err(std::env::VarError::NotPresent | std::env::VarError::NotUnicode(_)) => {