mirror of
https://github.com/openai/codex.git
synced 2026-11-06 20:20:34 +00:00
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:
@@ -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(_)) => {
|
||||
|
||||
Reference in New Issue
Block a user