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
This commit is contained in:
jif
2026-11-06 13:54:03 +00:00
committed by copyberry
parent 2210380232
commit 2c502d8b3d
3 changed files with 144 additions and 34 deletions
@@ -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<AuthManager>),
RuntimeLogin(Arc<AuthManager>, u64),
}
impl From<&str> for DecisionsCredential {
fn from(key: &str) -> Self {
Self::ApiKey(key.to_owned())
}
}
impl DecisionsCredential {
fn source_generation(&self) -> Option<u64> {
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<u64>) -> Option<Cow<'_, str>> {
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<VecDeque<AbortHandle>>,
@@ -102,10 +144,10 @@ impl DecisionsTask {
impl DecisionsSampler {
pub(super) fn new(
client: HttpClient,
api_key: String,
api_key: DecisionsCredential,
url: String,
) -> Result<Self, DecisionsError> {
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<Self>,
body: Result<Value, DecisionsError>,
started: Instant,
source_generation: Option<u64>,
) -> 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<u64>,
) -> 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()
@@ -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<AuthManager>,
) -> Result<String, DecisionsError> {
select_api_key(provider_id, base_url, /*context*/ None, manager).and_then(current_key)
}
fn current_key(credential: DecisionsCredential) -> Result<String, DecisionsError> {
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::<Vec<_>>();
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]
@@ -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<AuthManager>,
) -> Option<DecisionsSampler> {
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::<ProviderAuthContext>().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<String, DecisionsError> {
context: Option<&ProviderAuthContext>,
fallback: &Arc<AuthManager>,
) -> Result<DecisionsCredential, DecisionsError> {
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(_)) => {