mirror of
https://github.com/openai/codex.git
synced 2026-11-06 20:20:34 +00:00
Bind runtime auth to the provider's selected credential source (#59823)
## Why Runtime authentication must use the credential selected by the model provider. Explicit API keys and command credentials need independent ownership so cached connections and in-flight work cannot cross credential sources. ## What changed - Register and reuse runtimes per credential source, supporting explicit provider keys and command credentials without falling back to saved login credentials. - Track `AuthOwner` by credential manager and owner generation across WebSocket state, compaction, and scoring. - Preserve cached identity and runtime registration state during credential-source outages or failed storage reloads. Require a successful auth reload before logout revokes credentials. - Reject requests whose production or staging environment differs from runtime registration, including model catalog requests. ## Testing Add regression coverage for credential selection and runtime reuse, owner isolation, source outages, failed reloads and logout recovery, and environment mismatches. GitOrigin-RevId: 7bab55e0ea502593393df10046309fd8502edbd0
This commit is contained in:
cooper-oai
@@ -3,6 +3,7 @@
|
||||
use crate::ActorInfo;
|
||||
use crate::RegistrationError;
|
||||
use crate::RuntimeAgentInstance;
|
||||
use crate::RuntimeAuthBinding;
|
||||
use crate::RuntimeAuthError;
|
||||
use codex_api::AuthError;
|
||||
use codex_api::AuthHeaderTelemetry;
|
||||
@@ -12,6 +13,7 @@ use codex_features::RuntimeAuthMode;
|
||||
use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_login::AuthChangeState;
|
||||
use codex_login::AuthOwner;
|
||||
use codex_login::CodexAuth;
|
||||
use http::HeaderMap;
|
||||
use http::HeaderValue;
|
||||
@@ -28,6 +30,8 @@ pub struct RuntimeAuth {
|
||||
pub(crate) auth: CodexAuth,
|
||||
pub(crate) http_client_factory: HttpClientFactory,
|
||||
pub(crate) change_state: AuthChangeState,
|
||||
pub(crate) owner: AuthOwner,
|
||||
pub(crate) binding: Arc<RuntimeAuthBinding>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RuntimeAuth {
|
||||
@@ -63,6 +67,14 @@ impl RuntimeAuth {
|
||||
self.change_state
|
||||
}
|
||||
|
||||
pub fn auth_owner(&self) -> &AuthOwner {
|
||||
&self.owner
|
||||
}
|
||||
|
||||
pub fn binding(&self) -> &Arc<RuntimeAuthBinding> {
|
||||
&self.binding
|
||||
}
|
||||
|
||||
/// Preserve the caller's transport settings and restrictions while binding requests to
|
||||
/// this credential's account. Independent managed policies must not replace each other.
|
||||
pub fn bind_http_client_factory(
|
||||
|
||||
@@ -0,0 +1,408 @@
|
||||
//! Register and refresh a runtime within one selected credential source.
|
||||
|
||||
use crate::AgentHostAuthorization;
|
||||
use crate::AuthenticatedActorAuthorization;
|
||||
use crate::PlatformResourceSelection;
|
||||
use crate::RegistrationError;
|
||||
use crate::RuntimeAgentInstance;
|
||||
use crate::RuntimeAuth;
|
||||
use crate::RuntimeRegistration;
|
||||
use crate::RuntimeRegistrationClient;
|
||||
use crate::RuntimeRegistrationMetadata;
|
||||
use codex_agent_identity::ChatGptEnvironment;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_http_client::NetworkPolicyDenied;
|
||||
use codex_login::AuthChangeState;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::AuthOwner;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::RefreshTokenError;
|
||||
use http::StatusCode;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use url::Url;
|
||||
|
||||
/// A selected credential source and its registration state, retained across request preparation.
|
||||
pub struct RuntimeAuthBinding {
|
||||
pub(crate) credentials: Arc<AuthManager>,
|
||||
pub(crate) registration: RegistrationTarget,
|
||||
pub(crate) metadata: RuntimeRegistrationMetadata,
|
||||
state: Mutex<Option<CachedRuntime>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum RegistrationTarget {
|
||||
Explicit {
|
||||
host: AgentHostAuthorization,
|
||||
url: Url,
|
||||
},
|
||||
ChatGpt {
|
||||
installation_id: String,
|
||||
base_url: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl RegistrationTarget {
|
||||
fn host(&self) -> Result<AgentHostAuthorization, RegistrationError> {
|
||||
match self {
|
||||
Self::Explicit { host, .. } => Ok(host.clone()),
|
||||
Self::ChatGpt {
|
||||
installation_id, ..
|
||||
} => Ok(AgentHostAuthorization::InstallationId {
|
||||
installation_id: installation_id
|
||||
.parse()
|
||||
.map_err(|_| RegistrationError::InvalidRequest)?,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn url(&self) -> Result<Url, RegistrationError> {
|
||||
match self {
|
||||
Self::Explicit { url, .. } => Ok(url.clone()),
|
||||
Self::ChatGpt { base_url, .. } => {
|
||||
let url = match ChatGptEnvironment::from_chatgpt_base_url(base_url)
|
||||
.map_err(|_| RegistrationError::InvalidEndpoint)?
|
||||
{
|
||||
ChatGptEnvironment::Production => crate::DEFAULT_REGISTRATION_URL,
|
||||
ChatGptEnvironment::Staging => "https://auth.api.openai.org/agent/register",
|
||||
};
|
||||
url.parse().map_err(|_| RegistrationError::InvalidEndpoint)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct CachedRuntime {
|
||||
source_generation: u64,
|
||||
selection: PlatformResourceSelection,
|
||||
state: RuntimeState,
|
||||
}
|
||||
|
||||
enum RuntimeState {
|
||||
Pending { registration: RuntimeRegistration },
|
||||
Ready { runtime: Arc<RuntimeAgentInstance> },
|
||||
Denied { error: RegistrationError },
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RuntimeAuthBinding {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("RuntimeAuthBinding").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl RuntimeAuthBinding {
|
||||
pub(crate) fn new(
|
||||
credentials: Arc<AuthManager>,
|
||||
registration: RegistrationTarget,
|
||||
metadata: RuntimeRegistrationMetadata,
|
||||
) -> Self {
|
||||
Self {
|
||||
credentials,
|
||||
registration,
|
||||
metadata,
|
||||
state: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn uses_credentials(&self, credentials: &Arc<AuthManager>) -> bool {
|
||||
Arc::ptr_eq(&self.credentials, credentials)
|
||||
}
|
||||
|
||||
pub fn auth_owner(&self) -> AuthOwner {
|
||||
self.credentials.auth_owner()
|
||||
}
|
||||
|
||||
pub fn registration_url(&self) -> Result<Url, RegistrationError> {
|
||||
self.registration.url()
|
||||
}
|
||||
|
||||
pub async fn auth(self: &Arc<Self>) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.auth_with_selection(&PlatformResourceSelection::default())
|
||||
.await
|
||||
}
|
||||
|
||||
/// Select auth and a runtime for the requested organization and project.
|
||||
pub async fn auth_with_selection(
|
||||
self: &Arc<Self>,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.select_auth(selection, /*expected_source*/ None).await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "runtime initialization serializes registration and refresh for one credential source"
|
||||
)]
|
||||
async fn select_auth(
|
||||
self: &Arc<Self>,
|
||||
selection: &PlatformResourceSelection,
|
||||
expected_source: Option<u64>,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
// This initializer's lock can span IO. No AuthManager lock is held across registration.
|
||||
let mut state = self.state.lock().await;
|
||||
self.check_source(expected_source)?;
|
||||
let (auth, http_client_factory) = self
|
||||
.credentials
|
||||
.auth_with_http_client_factory()
|
||||
.await
|
||||
.ok_or(RuntimeAuthError::Unavailable)?;
|
||||
let mut recovery = self.credentials.unauthorized_recovery();
|
||||
let mut result = self
|
||||
.auth_once(
|
||||
&mut state,
|
||||
auth,
|
||||
http_client_factory,
|
||||
selection,
|
||||
expected_source,
|
||||
)
|
||||
.await;
|
||||
while matches!(
|
||||
&result,
|
||||
Err(RuntimeAuthError::Registration(
|
||||
RegistrationError::Rejected {
|
||||
status: StatusCode::UNAUTHORIZED,
|
||||
..
|
||||
}
|
||||
))
|
||||
) && recovery.has_next()
|
||||
{
|
||||
self.check_source(expected_source)?;
|
||||
recovery.next().await.map_err(|error| match error {
|
||||
RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied),
|
||||
RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => {
|
||||
RuntimeAuthError::RefreshFailed
|
||||
}
|
||||
})?;
|
||||
let (auth, http_client_factory) = self
|
||||
.credentials
|
||||
.auth_with_http_client_factory()
|
||||
.await
|
||||
.ok_or(RuntimeAuthError::Unavailable)?;
|
||||
result = self
|
||||
.auth_once(
|
||||
&mut state,
|
||||
auth,
|
||||
http_client_factory,
|
||||
selection,
|
||||
expected_source,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn auth_once(
|
||||
self: &Arc<Self>,
|
||||
state: &mut Option<CachedRuntime>,
|
||||
auth: CodexAuth,
|
||||
http_client_factory: HttpClientFactory,
|
||||
selection: &PlatformResourceSelection,
|
||||
expected_source: Option<u64>,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
let actor_auth = actor_authorization(&auth, selection)?;
|
||||
let http_client_factory = http_client_factory.with_sensitive_header_protection();
|
||||
let change_state = self.auth_state(&auth)?;
|
||||
self.check_source(expected_source)?;
|
||||
// Registered runtimes can serve different projects. Pending and denied
|
||||
// registrations retain their original scope and idempotency key.
|
||||
if state.as_ref().is_some_and(|cached| {
|
||||
cached.source_generation != change_state.source_generation
|
||||
|| cached.selection.organization_id != selection.organization_id
|
||||
|| (!matches!(cached.state, RuntimeState::Ready { .. })
|
||||
&& cached.selection != *selection)
|
||||
}) {
|
||||
*state = None;
|
||||
}
|
||||
let cached = match state {
|
||||
Some(cached) => cached,
|
||||
None => state.insert(CachedRuntime {
|
||||
source_generation: change_state.source_generation,
|
||||
selection: selection.clone(),
|
||||
state: RuntimeState::Pending {
|
||||
registration: RuntimeRegistration::new(
|
||||
actor_auth.clone(),
|
||||
self.registration.host()?,
|
||||
self.metadata.clone(),
|
||||
)
|
||||
.map_err(|_| RuntimeAuthError::KeyGeneration)?,
|
||||
},
|
||||
}),
|
||||
};
|
||||
let runtime_state = &mut cached.state;
|
||||
let registration = match runtime_state {
|
||||
RuntimeState::Ready { runtime } => {
|
||||
return self.selected_auth(
|
||||
runtime.clone(),
|
||||
auth,
|
||||
http_client_factory,
|
||||
change_state,
|
||||
);
|
||||
}
|
||||
RuntimeState::Denied { error } => return Err(error.clone().into()),
|
||||
RuntimeState::Pending { registration } => {
|
||||
registration.request.actor_auth = actor_auth;
|
||||
registration
|
||||
}
|
||||
};
|
||||
let client =
|
||||
RuntimeRegistrationClient::new(http_client_factory.clone(), self.registration.url()?)?;
|
||||
let result = client.register(registration).await;
|
||||
if self.auth_state(&auth)?.source_generation != change_state.source_generation {
|
||||
// A late response must not publish an old identity after a credential switch.
|
||||
*state = None;
|
||||
return Err(RuntimeAuthError::Changed);
|
||||
}
|
||||
match result {
|
||||
Ok(info) => {
|
||||
let runtime = Arc::new(registration.finish(info));
|
||||
*runtime_state = RuntimeState::Ready {
|
||||
runtime: runtime.clone(),
|
||||
};
|
||||
self.selected_auth(runtime, auth, http_client_factory, change_state)
|
||||
}
|
||||
Err(
|
||||
error @ RegistrationError::Rejected {
|
||||
status: StatusCode::FORBIDDEN,
|
||||
..
|
||||
},
|
||||
) => {
|
||||
*runtime_state = RuntimeState::Denied {
|
||||
error: error.clone(),
|
||||
};
|
||||
Err(error.into())
|
||||
}
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn selected_auth(
|
||||
self: &Arc<Self>,
|
||||
runtime: Arc<RuntimeAgentInstance>,
|
||||
auth: CodexAuth,
|
||||
http_client_factory: HttpClientFactory,
|
||||
change_state: AuthChangeState,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
let owner = self
|
||||
.credentials
|
||||
.auth_owner_for(&auth)
|
||||
.ok_or(RuntimeAuthError::Changed)?;
|
||||
Ok(RuntimeAuth {
|
||||
runtime,
|
||||
auth,
|
||||
http_client_factory,
|
||||
change_state,
|
||||
owner,
|
||||
binding: Arc::clone(self),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve auth for work whose routing, cache key, or payload uses a captured credential.
|
||||
/// A same-source refresh is allowed; a different source must not take over that work.
|
||||
/// This binds credential identity, not when the caller captured it. Restoring the exact
|
||||
/// same credential can select a new runtime after logout; a different source cannot.
|
||||
pub async fn auth_for(
|
||||
self: &Arc<Self>,
|
||||
expected: &CodexAuth,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
let source = self
|
||||
.credentials
|
||||
.auth_change_state_for(expected)
|
||||
.ok_or(RuntimeAuthError::Changed)?;
|
||||
self.auth_for_source(source.source_generation, selection)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Select another request within a reusable client's original credential source.
|
||||
pub async fn auth_for_source(
|
||||
self: &Arc<Self>,
|
||||
source_generation: u64,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.select_auth(selection, Some(source_generation)).await
|
||||
}
|
||||
|
||||
fn check_source(&self, expected: Option<u64>) -> Result<(), RuntimeAuthError> {
|
||||
if expected.is_some_and(|expected| {
|
||||
self.credentials
|
||||
.auth_change_state_receiver()
|
||||
.borrow()
|
||||
.source_generation
|
||||
!= expected
|
||||
}) {
|
||||
return Err(RuntimeAuthError::Changed);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn auth_state(&self, auth: &CodexAuth) -> Result<AuthChangeState, RuntimeAuthError> {
|
||||
let state = self
|
||||
.credentials
|
||||
.auth_change_state_for(auth)
|
||||
.ok_or(RuntimeAuthError::Changed)?;
|
||||
if self.credentials.refresh_failure_for_auth(auth).is_some() {
|
||||
return Err(RuntimeAuthError::RefreshFailed);
|
||||
}
|
||||
Ok(state)
|
||||
}
|
||||
}
|
||||
|
||||
fn actor_authorization(
|
||||
auth: &CodexAuth,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<AuthenticatedActorAuthorization, RuntimeAuthError> {
|
||||
if !matches!(auth, CodexAuth::ApiKey(_)) && selection != &PlatformResourceSelection::default() {
|
||||
return Err(RuntimeAuthError::UnsupportedResourceSelection);
|
||||
}
|
||||
let credential = auth
|
||||
.get_token()
|
||||
.map_err(|_| RuntimeAuthError::Unsupported)?
|
||||
.into();
|
||||
Ok(match auth {
|
||||
CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey {
|
||||
credential,
|
||||
selection: selection.clone(),
|
||||
},
|
||||
CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => {
|
||||
AuthenticatedActorAuthorization::Oauth { credential }
|
||||
}
|
||||
CodexAuth::PersonalAccessToken(_) => {
|
||||
AuthenticatedActorAuthorization::PersonalAccessToken { credential }
|
||||
}
|
||||
CodexAuth::Headers(_)
|
||||
| CodexAuth::AgentIdentity(_)
|
||||
| CodexAuth::BedrockApiKey(_)
|
||||
| CodexAuth::BedrockAccessKeys(_) => return Err(RuntimeAuthError::Unsupported),
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum RuntimeAuthError {
|
||||
#[error("runtime registration requires an actor credential")]
|
||||
Unavailable,
|
||||
#[error(
|
||||
"this credential cannot register a runtime; sign in again with a supported OpenAI credential"
|
||||
)]
|
||||
Unsupported,
|
||||
#[error("organization and project selection requires an OpenAI API key")]
|
||||
UnsupportedResourceSelection,
|
||||
#[error("credentials changed during runtime registration; try again")]
|
||||
Changed,
|
||||
#[error("runtime registration could not refresh actor credentials")]
|
||||
RefreshFailed,
|
||||
#[error("failed to generate runtime signing key")]
|
||||
KeyGeneration,
|
||||
#[error(transparent)]
|
||||
Policy(#[from] NetworkPolicyDenied),
|
||||
#[error(transparent)]
|
||||
Registration(RegistrationError),
|
||||
}
|
||||
|
||||
impl From<RegistrationError> for RuntimeAuthError {
|
||||
fn from(error: RegistrationError) -> Self {
|
||||
match error {
|
||||
RegistrationError::Policy(denied) => Self::Policy(denied),
|
||||
error => Self::Registration(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
mod actor;
|
||||
mod auth;
|
||||
mod binding;
|
||||
mod bootstrap;
|
||||
mod manager;
|
||||
mod registration;
|
||||
@@ -20,7 +21,8 @@ pub use bootstrap::RuntimeRole;
|
||||
pub use codex_http_client::is_openai_api_url;
|
||||
|
||||
pub use auth::RuntimeAuth;
|
||||
pub use manager::RuntimeAuthError;
|
||||
pub use binding::RuntimeAuthBinding;
|
||||
pub use binding::RuntimeAuthError;
|
||||
pub use manager::RuntimeAuthManager;
|
||||
pub use registration::DEFAULT_REGISTRATION_URL;
|
||||
pub use registration::RegistrationError;
|
||||
|
||||
@@ -1,85 +1,28 @@
|
||||
//! Reuse one active runtime for the selected credential source and organization.
|
||||
//! Bind credential sources to shared runtime registrations.
|
||||
|
||||
use crate::AgentHostAuthorization;
|
||||
use crate::AuthenticatedActorAuthorization;
|
||||
use crate::PlatformResourceSelection;
|
||||
use crate::RegistrationError;
|
||||
use crate::RuntimeAgentInstance;
|
||||
use crate::RuntimeAuth;
|
||||
use crate::RuntimeRegistration;
|
||||
use crate::RuntimeRegistrationClient;
|
||||
use crate::RuntimeAuthBinding;
|
||||
use crate::RuntimeAuthError;
|
||||
use crate::RuntimeRegistrationMetadata;
|
||||
use codex_agent_identity::ChatGptEnvironment;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_http_client::NetworkPolicyDenied;
|
||||
use codex_login::AuthChangeState;
|
||||
use crate::binding::RegistrationTarget;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::AuthRuntimeConfig;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::RefreshTokenError;
|
||||
use http::StatusCode;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use url::Url;
|
||||
|
||||
/// Owned alongside an AuthManager, not created for every thread or request.
|
||||
/// Shares registration state across requests using the same credential source and context.
|
||||
pub struct RuntimeAuthManager {
|
||||
credentials: Arc<AuthManager>,
|
||||
registration: RegistrationTarget,
|
||||
metadata: RuntimeRegistrationMetadata,
|
||||
state: Mutex<Option<CachedRuntime>>,
|
||||
login: Arc<RuntimeAuthBinding>,
|
||||
registrations: std::sync::Mutex<Vec<CredentialRegistration>>,
|
||||
}
|
||||
|
||||
enum RegistrationTarget {
|
||||
Explicit {
|
||||
host: AgentHostAuthorization,
|
||||
url: Url,
|
||||
},
|
||||
ChatGpt {
|
||||
installation_id: String,
|
||||
base_url: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl RegistrationTarget {
|
||||
fn host(&self) -> Result<AgentHostAuthorization, RegistrationError> {
|
||||
match self {
|
||||
Self::Explicit { host, .. } => Ok(host.clone()),
|
||||
Self::ChatGpt {
|
||||
installation_id, ..
|
||||
} => Ok(AgentHostAuthorization::InstallationId {
|
||||
installation_id: installation_id
|
||||
.parse()
|
||||
.map_err(|_| RegistrationError::InvalidRequest)?,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn url(&self) -> Result<Url, RegistrationError> {
|
||||
match self {
|
||||
Self::Explicit { url, .. } => Ok(url.clone()),
|
||||
Self::ChatGpt { base_url, .. } => {
|
||||
let url = match ChatGptEnvironment::from_chatgpt_base_url(base_url)
|
||||
.map_err(|_| RegistrationError::InvalidEndpoint)?
|
||||
{
|
||||
ChatGptEnvironment::Production => crate::DEFAULT_REGISTRATION_URL,
|
||||
ChatGptEnvironment::Staging => "https://auth.api.openai.org/agent/register",
|
||||
};
|
||||
url.parse().map_err(|_| RegistrationError::InvalidEndpoint)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct CachedRuntime {
|
||||
source_generation: u64,
|
||||
selection: PlatformResourceSelection,
|
||||
state: RuntimeState,
|
||||
}
|
||||
|
||||
enum RuntimeState {
|
||||
Pending { registration: RuntimeRegistration },
|
||||
Ready { runtime: Arc<RuntimeAgentInstance> },
|
||||
Denied { error: RegistrationError },
|
||||
enum CredentialRegistration {
|
||||
Provider(Arc<RuntimeAuthBinding>),
|
||||
ApiKey(Arc<RuntimeAuthBinding>),
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RuntimeAuthManager {
|
||||
@@ -89,9 +32,8 @@ impl std::fmt::Debug for RuntimeAuthManager {
|
||||
}
|
||||
|
||||
impl RuntimeAuthManager {
|
||||
/// Check that a consumer uses this owner's credential provider, not an ambient fallback.
|
||||
pub fn uses_credentials(&self, credentials: &Arc<AuthManager>) -> bool {
|
||||
Arc::ptr_eq(&self.credentials, credentials)
|
||||
self.login.uses_credentials(credentials)
|
||||
}
|
||||
|
||||
pub fn new(
|
||||
@@ -100,15 +42,14 @@ impl RuntimeAuthManager {
|
||||
metadata: RuntimeRegistrationMetadata,
|
||||
registration_url: Url,
|
||||
) -> Self {
|
||||
Self {
|
||||
Self::from_registration(
|
||||
credentials,
|
||||
registration: RegistrationTarget::Explicit {
|
||||
RegistrationTarget::Explicit {
|
||||
host,
|
||||
url: registration_url,
|
||||
},
|
||||
metadata,
|
||||
state: Mutex::new(None),
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolve the ChatGPT registration origin only when a first-party request needs a runtime.
|
||||
@@ -118,287 +59,119 @@ impl RuntimeAuthManager {
|
||||
base_url: String,
|
||||
metadata: RuntimeRegistrationMetadata,
|
||||
) -> Self {
|
||||
Self {
|
||||
Self::from_registration(
|
||||
credentials,
|
||||
registration: RegistrationTarget::ChatGpt {
|
||||
RegistrationTarget::ChatGpt {
|
||||
installation_id,
|
||||
base_url,
|
||||
},
|
||||
metadata,
|
||||
state: Mutex::new(None),
|
||||
)
|
||||
}
|
||||
|
||||
fn from_registration(
|
||||
credentials: Arc<AuthManager>,
|
||||
registration: RegistrationTarget,
|
||||
metadata: RuntimeRegistrationMetadata,
|
||||
) -> Self {
|
||||
Self {
|
||||
login: Arc::new(RuntimeAuthBinding::new(credentials, registration, metadata)),
|
||||
registrations: std::sync::Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn auth(&self) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.auth_with_selection(&PlatformResourceSelection::default())
|
||||
.await
|
||||
pub fn registration_url(&self) -> Result<Url, RegistrationError> {
|
||||
self.login.registration_url()
|
||||
}
|
||||
|
||||
pub async fn auth(&self) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.login.auth().await
|
||||
}
|
||||
|
||||
/// Select auth and a runtime for the requested organization and project.
|
||||
pub async fn auth_with_selection(
|
||||
&self,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.select_auth(selection, /*expected_source*/ None).await
|
||||
self.login.auth_with_selection(selection).await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "runtime initialization serializes registration and refresh for one credential source"
|
||||
)]
|
||||
async fn select_auth(
|
||||
&self,
|
||||
selection: &PlatformResourceSelection,
|
||||
expected_source: Option<u64>,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
// This initializer's lock can span IO. No AuthManager lock is held across registration.
|
||||
let mut state = self.state.lock().await;
|
||||
self.check_source(expected_source)?;
|
||||
let (auth, http_client_factory) = self
|
||||
.credentials
|
||||
.auth_with_http_client_factory()
|
||||
.await
|
||||
.ok_or(RuntimeAuthError::Unavailable)?;
|
||||
let mut recovery = self.credentials.unauthorized_recovery();
|
||||
let mut result = self
|
||||
.auth_once(
|
||||
&mut state,
|
||||
auth,
|
||||
http_client_factory,
|
||||
selection,
|
||||
expected_source,
|
||||
)
|
||||
.await;
|
||||
while matches!(
|
||||
&result,
|
||||
Err(RuntimeAuthError::Registration(
|
||||
RegistrationError::Rejected {
|
||||
status: StatusCode::UNAUTHORIZED,
|
||||
..
|
||||
}
|
||||
))
|
||||
) && recovery.has_next()
|
||||
{
|
||||
self.check_source(expected_source)?;
|
||||
recovery.next().await.map_err(|error| match error {
|
||||
RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied),
|
||||
RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => {
|
||||
RuntimeAuthError::RefreshFailed
|
||||
}
|
||||
})?;
|
||||
let (auth, http_client_factory) = self
|
||||
.credentials
|
||||
.auth_with_http_client_factory()
|
||||
.await
|
||||
.ok_or(RuntimeAuthError::Unavailable)?;
|
||||
result = self
|
||||
.auth_once(
|
||||
&mut state,
|
||||
auth,
|
||||
http_client_factory,
|
||||
selection,
|
||||
expected_source,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn auth_once(
|
||||
&self,
|
||||
state: &mut Option<CachedRuntime>,
|
||||
auth: CodexAuth,
|
||||
http_client_factory: HttpClientFactory,
|
||||
selection: &PlatformResourceSelection,
|
||||
expected_source: Option<u64>,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
let actor_auth = actor_authorization(&auth, selection)?;
|
||||
let http_client_factory = http_client_factory.with_sensitive_header_protection();
|
||||
let change_state = self.auth_state(&auth)?;
|
||||
self.check_source(expected_source)?;
|
||||
// Registered runtimes can serve different projects. Pending and denied
|
||||
// registrations retain their original scope and idempotency key.
|
||||
if state.as_ref().is_some_and(|cached| {
|
||||
cached.source_generation != change_state.source_generation
|
||||
|| cached.selection.organization_id != selection.organization_id
|
||||
|| (!matches!(cached.state, RuntimeState::Ready { .. })
|
||||
&& cached.selection != *selection)
|
||||
}) {
|
||||
*state = None;
|
||||
}
|
||||
let cached = match state {
|
||||
Some(cached) => cached,
|
||||
None => state.insert(CachedRuntime {
|
||||
source_generation: change_state.source_generation,
|
||||
selection: selection.clone(),
|
||||
state: RuntimeState::Pending {
|
||||
registration: RuntimeRegistration::new(
|
||||
actor_auth.clone(),
|
||||
self.registration.host()?,
|
||||
self.metadata.clone(),
|
||||
)
|
||||
.map_err(|_| RuntimeAuthError::KeyGeneration)?,
|
||||
},
|
||||
}),
|
||||
};
|
||||
let runtime_state = &mut cached.state;
|
||||
let registration = match runtime_state {
|
||||
RuntimeState::Ready { runtime } => {
|
||||
return Ok(RuntimeAuth {
|
||||
runtime: runtime.clone(),
|
||||
auth,
|
||||
http_client_factory,
|
||||
change_state,
|
||||
});
|
||||
}
|
||||
RuntimeState::Denied { error } => return Err(error.clone().into()),
|
||||
RuntimeState::Pending { registration } => {
|
||||
registration.request.actor_auth = actor_auth;
|
||||
registration
|
||||
}
|
||||
};
|
||||
let client =
|
||||
RuntimeRegistrationClient::new(http_client_factory.clone(), self.registration.url()?)?;
|
||||
let result = client.register(registration).await;
|
||||
if self.auth_state(&auth)?.source_generation != change_state.source_generation {
|
||||
// A late response must not publish an old identity after a credential switch.
|
||||
*state = None;
|
||||
return Err(RuntimeAuthError::Changed);
|
||||
}
|
||||
match result {
|
||||
Ok(info) => {
|
||||
let runtime = Arc::new(registration.finish(info));
|
||||
*runtime_state = RuntimeState::Ready {
|
||||
runtime: runtime.clone(),
|
||||
};
|
||||
Ok(RuntimeAuth {
|
||||
runtime,
|
||||
auth,
|
||||
http_client_factory,
|
||||
change_state,
|
||||
})
|
||||
}
|
||||
Err(
|
||||
error @ RegistrationError::Rejected {
|
||||
status: StatusCode::FORBIDDEN,
|
||||
..
|
||||
},
|
||||
) => {
|
||||
*runtime_state = RuntimeState::Denied {
|
||||
error: error.clone(),
|
||||
};
|
||||
Err(error.into())
|
||||
}
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve auth for work whose routing, cache key, or payload uses a captured credential.
|
||||
/// A same-source refresh is allowed; a different source must not take over that work.
|
||||
/// This binds credential identity, not when the caller captured it. Restoring the exact
|
||||
/// same credential can select a new runtime after logout; a different source cannot.
|
||||
pub async fn auth_for(
|
||||
&self,
|
||||
expected: &CodexAuth,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
let source = self
|
||||
.credentials
|
||||
.auth_change_state_for(expected)
|
||||
.ok_or(RuntimeAuthError::Changed)?;
|
||||
self.auth_for_source(source.source_generation, selection)
|
||||
.await
|
||||
self.login.auth_for(expected, selection).await
|
||||
}
|
||||
|
||||
/// Select another request within a reusable client's original credential source.
|
||||
pub async fn auth_for_source(
|
||||
&self,
|
||||
source_generation: u64,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<RuntimeAuth, RuntimeAuthError> {
|
||||
self.select_auth(selection, Some(source_generation)).await
|
||||
self.login
|
||||
.auth_for_source(source_generation, selection)
|
||||
.await
|
||||
}
|
||||
|
||||
fn check_source(&self, expected: Option<u64>) -> Result<(), RuntimeAuthError> {
|
||||
if expected.is_some_and(|expected| {
|
||||
self.credentials
|
||||
.auth_change_state_receiver()
|
||||
.borrow()
|
||||
.source_generation
|
||||
!= expected
|
||||
/// Bind a login or command-backed credential without changing its refresh lifecycle.
|
||||
pub fn bind(&self, credentials: Arc<AuthManager>) -> Arc<RuntimeAuthBinding> {
|
||||
if self.login.uses_credentials(&credentials) {
|
||||
return Arc::clone(&self.login);
|
||||
}
|
||||
let mut registrations = self
|
||||
.registrations
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(binding) = registrations.iter().find_map(|entry| match entry {
|
||||
CredentialRegistration::Provider(binding) if binding.uses_credentials(&credentials) => {
|
||||
Some(binding)
|
||||
}
|
||||
_ => None,
|
||||
}) {
|
||||
return Err(RuntimeAuthError::Changed);
|
||||
return Arc::clone(binding);
|
||||
}
|
||||
Ok(())
|
||||
let binding = Arc::new(RuntimeAuthBinding::new(
|
||||
credentials,
|
||||
self.login.registration.clone(),
|
||||
self.login.metadata.clone(),
|
||||
));
|
||||
registrations.push(CredentialRegistration::Provider(Arc::clone(&binding)));
|
||||
binding
|
||||
}
|
||||
|
||||
fn auth_state(&self, auth: &CodexAuth) -> Result<AuthChangeState, RuntimeAuthError> {
|
||||
let state = self
|
||||
.credentials
|
||||
.auth_change_state_for(auth)
|
||||
.ok_or(RuntimeAuthError::Changed)?;
|
||||
if self.credentials.refresh_failure_for_auth(auth).is_some() {
|
||||
return Err(RuntimeAuthError::RefreshFailed);
|
||||
}
|
||||
Ok(state)
|
||||
}
|
||||
}
|
||||
|
||||
fn actor_authorization(
|
||||
auth: &CodexAuth,
|
||||
selection: &PlatformResourceSelection,
|
||||
) -> Result<AuthenticatedActorAuthorization, RuntimeAuthError> {
|
||||
if !matches!(auth, CodexAuth::ApiKey(_)) && selection != &PlatformResourceSelection::default() {
|
||||
return Err(RuntimeAuthError::UnsupportedResourceSelection);
|
||||
}
|
||||
let credential = auth
|
||||
.get_token()
|
||||
.map_err(|_| RuntimeAuthError::Unsupported)?
|
||||
.into();
|
||||
Ok(match auth {
|
||||
CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey {
|
||||
credential,
|
||||
selection: selection.clone(),
|
||||
},
|
||||
CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => {
|
||||
AuthenticatedActorAuthorization::Oauth { credential }
|
||||
}
|
||||
CodexAuth::PersonalAccessToken(_) => {
|
||||
AuthenticatedActorAuthorization::PersonalAccessToken { credential }
|
||||
}
|
||||
CodexAuth::Headers(_)
|
||||
| CodexAuth::AgentIdentity(_)
|
||||
| CodexAuth::BedrockApiKey(_)
|
||||
| CodexAuth::BedrockAccessKeys(_) => return Err(RuntimeAuthError::Unsupported),
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum RuntimeAuthError {
|
||||
#[error("runtime registration requires an actor credential")]
|
||||
Unavailable,
|
||||
#[error("this credential cannot register a runtime; use a supported OpenAI sign-in")]
|
||||
Unsupported,
|
||||
#[error("organization and project selection requires an OpenAI API key")]
|
||||
UnsupportedResourceSelection,
|
||||
#[error("credentials changed during runtime registration; try again")]
|
||||
Changed,
|
||||
#[error("runtime registration could not refresh actor credentials")]
|
||||
RefreshFailed,
|
||||
#[error("failed to generate runtime signing key")]
|
||||
KeyGeneration,
|
||||
#[error(transparent)]
|
||||
Policy(#[from] NetworkPolicyDenied),
|
||||
#[error(transparent)]
|
||||
Registration(RegistrationError),
|
||||
}
|
||||
|
||||
impl From<RegistrationError> for RuntimeAuthError {
|
||||
fn from(error: RegistrationError) -> Self {
|
||||
match error {
|
||||
RegistrationError::Policy(denied) => Self::Policy(denied),
|
||||
error => Self::Registration(error),
|
||||
/// An explicit key is independent of login storage and command output, even when bytes match.
|
||||
pub fn bind_api_key(
|
||||
&self,
|
||||
api_key: &str,
|
||||
runtime: &AuthRuntimeConfig,
|
||||
) -> Arc<RuntimeAuthBinding> {
|
||||
let mut registrations = self
|
||||
.registrations
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(binding) = registrations.iter().find_map(|entry| match entry {
|
||||
CredentialRegistration::ApiKey(binding)
|
||||
if binding
|
||||
.credentials
|
||||
.auth_cached()
|
||||
.as_ref()
|
||||
.and_then(CodexAuth::api_key)
|
||||
== Some(api_key)
|
||||
&& binding.credentials.runtime_config() == *runtime =>
|
||||
{
|
||||
Some(binding)
|
||||
}
|
||||
_ => None,
|
||||
}) {
|
||||
return Arc::clone(binding);
|
||||
}
|
||||
let binding = Arc::new(RuntimeAuthBinding::new(
|
||||
AuthManager::from_api_key(api_key, runtime.clone()),
|
||||
self.login.registration.clone(),
|
||||
self.login.metadata.clone(),
|
||||
));
|
||||
registrations.push(CredentialRegistration::ApiKey(Arc::clone(&binding)));
|
||||
binding
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ use codex_features::RuntimeAuthMode;
|
||||
use codex_http_client::DestinationPolicy;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_http_client::NetworkPolicyController;
|
||||
use codex_http_client::NetworkPolicyDenied;
|
||||
use codex_http_client::OutboundProxyPolicy;
|
||||
use codex_login::AuthCredentialsStoreMode;
|
||||
use codex_login::AuthKeyringBackendKind;
|
||||
@@ -34,13 +35,19 @@ use wiremock::matchers::body_partial_json;
|
||||
use wiremock::matchers::method;
|
||||
|
||||
struct Credentials {
|
||||
available: std::sync::atomic::AtomicBool,
|
||||
current: RwLock<CodexAuth>,
|
||||
refreshed: CodexAuth,
|
||||
}
|
||||
|
||||
impl ExternalAuth for Credentials {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async { Ok(self.current.read().unwrap().clone()) })
|
||||
Box::pin(async {
|
||||
if !self.available.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
return Err(std::io::Error::other("credential source unavailable"));
|
||||
}
|
||||
Ok(self.current.read().unwrap().clone())
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
@@ -58,6 +65,7 @@ async fn setup(
|
||||
) -> (Arc<RuntimeAuthManager>, Arc<Credentials>, Arc<AuthManager>) {
|
||||
let credentials = AuthManager::from_auth_for_testing(initial.clone());
|
||||
let source = Arc::new(Credentials {
|
||||
available: std::sync::atomic::AtomicBool::new(true),
|
||||
current: RwLock::new(initial),
|
||||
refreshed,
|
||||
});
|
||||
@@ -106,6 +114,97 @@ fn oauth(session: Option<&str>, token: &str) -> CodexAuth {
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_login_requires_sign_in_for_runtime_auth() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(registered)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let claims = json!({
|
||||
"exp": 4102444800_u64,
|
||||
"session_id": "session-a",
|
||||
"iss": "https://auth.openai.com",
|
||||
"https://api.openai.com/auth": {
|
||||
"chatgpt_user_id": "user-a", "chatgpt_account_id": "workspace-a",
|
||||
},
|
||||
});
|
||||
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string());
|
||||
let token = format!("e30.{payload}.c2ln");
|
||||
let stored: codex_login::AuthDotJson = serde_json::from_value(json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_identity": {
|
||||
"agent_runtime_id": "legacy-runtime",
|
||||
"agent_private_key": codex_agent_identity::generate_agent_key_material()
|
||||
.unwrap().private_key_pkcs8_base64,
|
||||
"account_id": "workspace-a", "chatgpt_user_id": "user-a", "plan_type": "pro",
|
||||
"chatgpt_account_is_fedramp": false, "task_id": "legacy-task",
|
||||
},
|
||||
"tokens": {
|
||||
"id_token": token, "access_token": token,
|
||||
"refresh_token": "refresh", "account_id": "workspace-a",
|
||||
},
|
||||
"last_refresh": "2099-01-01T00:00:00Z",
|
||||
}))
|
||||
.unwrap();
|
||||
codex_login::save_auth(
|
||||
home.path(),
|
||||
&stored,
|
||||
AuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::Direct,
|
||||
)
|
||||
.unwrap();
|
||||
let credentials = AuthManager::shared(
|
||||
home.path().to_path_buf(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*forced_chatgpt_workspace_id*/ None,
|
||||
/*chatgpt_base_url*/ None,
|
||||
AuthKeyringBackendKind::Direct,
|
||||
codex_login::test_support::transport_default_auth_route_config(),
|
||||
)
|
||||
.await;
|
||||
let legacy = credentials.auth_cached().unwrap();
|
||||
assert!(matches!(&legacy, CodexAuth::AgentIdentity(_)));
|
||||
let runtime = RuntimeAuthManager::new(
|
||||
credentials.clone(),
|
||||
AgentHostAuthorization::InstallationId {
|
||||
installation_id: Uuid::nil(),
|
||||
},
|
||||
RuntimeRegistrationMetadata {
|
||||
harness: Some("codex".into()),
|
||||
harness_version: None,
|
||||
role: RuntimeRole::OrchestratorAndExecutor,
|
||||
},
|
||||
format!("{}/agent/register", server.uri()).parse().unwrap(),
|
||||
);
|
||||
assert!(matches!(
|
||||
runtime.auth().await.unwrap_err(),
|
||||
RuntimeAuthError::Unsupported
|
||||
));
|
||||
assert!(matches!(
|
||||
runtime
|
||||
.auth_for(&legacy, &PlatformResourceSelection::default())
|
||||
.await
|
||||
.unwrap_err(),
|
||||
RuntimeAuthError::Unsupported,
|
||||
));
|
||||
|
||||
codex_login::login_with_api_key(
|
||||
home.path(),
|
||||
"replacement-key",
|
||||
AuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::Direct,
|
||||
)
|
||||
.unwrap();
|
||||
credentials.reload().await;
|
||||
assert_eq!(
|
||||
runtime.auth().await.unwrap().auth().api_key(),
|
||||
Some("replacement-key")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_requests_share_registration_and_replacement_preserves_in_flight_auth() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -276,6 +375,7 @@ async fn policy_recovery_preserves_pending_registration() {
|
||||
let auth = CodexAuth::from_api_key("key-a");
|
||||
credentials
|
||||
.set_external_auth(Arc::new(Credentials {
|
||||
available: std::sync::atomic::AtomicBool::new(true),
|
||||
current: RwLock::new(auth.clone()),
|
||||
refreshed: auth,
|
||||
}))
|
||||
@@ -708,3 +808,218 @@ async fn file_upload_runtime_auth_signs_control_requests_but_not_blob_uploads()
|
||||
assert_eq!(blobs[0].body, b"hello");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn source_outage_preserves_registered_pending_and_denied_runtimes() {
|
||||
for status in [200, 403, 503] {
|
||||
let server = MockServer::start().await;
|
||||
let initial = oauth(Some("login-a"), "initial");
|
||||
let initial_token = initial.get_token().unwrap();
|
||||
let attempts = std::sync::atomic::AtomicUsize::new(0);
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(move |request: &Request| {
|
||||
let attempt = attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
let body: Value = request.body_json().unwrap();
|
||||
if (status == 403 && body["actor_auth"]["credential"] == initial_token)
|
||||
|| (status == 503 && attempt == 0)
|
||||
{
|
||||
ResponseTemplate::new(status)
|
||||
} else {
|
||||
registered(request)
|
||||
}
|
||||
})
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (manager, source, credentials) = setup(&server, initial.clone(), initial).await;
|
||||
let first = manager.auth().await;
|
||||
assert_eq!(first.is_ok(), status == 200);
|
||||
let initial_source = *credentials.auth_change_state_receiver().borrow();
|
||||
source
|
||||
.available
|
||||
.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||
let cached = manager.auth().await;
|
||||
assert_eq!(cached.is_ok(), status != 403);
|
||||
assert_eq!(
|
||||
initial_source,
|
||||
*credentials.auth_change_state_receiver().borrow()
|
||||
);
|
||||
let attempts_before_replacement = if status == 503 { 2 } else { 1 };
|
||||
assert_eq!(
|
||||
server.received_requests().await.unwrap().len(),
|
||||
attempts_before_replacement
|
||||
);
|
||||
source
|
||||
.available
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
let recovered = manager.auth().await;
|
||||
assert_eq!(recovered.is_ok(), status != 403);
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), attempts_before_replacement);
|
||||
if let Ok(cached) = cached {
|
||||
assert_eq!(
|
||||
cached.runtime().id(),
|
||||
recovered.as_ref().unwrap().runtime().id()
|
||||
);
|
||||
}
|
||||
if status == 200 {
|
||||
assert_eq!(
|
||||
first.unwrap().runtime().id(),
|
||||
recovered.unwrap().runtime().id()
|
||||
);
|
||||
}
|
||||
if status == 503 {
|
||||
assert_eq!(
|
||||
requests[0].headers["idempotency-key"],
|
||||
requests[1].headers["idempotency-key"]
|
||||
);
|
||||
assert_eq!(requests[0].body, requests[1].body);
|
||||
}
|
||||
// Actual source replacement still retires the old key, including a denied registration.
|
||||
*source.current.write().unwrap() = oauth(Some("login-b"), "replacement");
|
||||
assert!(manager.auth().await.is_ok());
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), attempts_before_replacement + 1);
|
||||
assert_ne!(
|
||||
requests[0].headers["idempotency-key"],
|
||||
requests.last().unwrap().headers["idempotency-key"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_keys_reuse_their_owner_without_borrowing_a_command_source() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(registered)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (root, _, _) = setup(
|
||||
&server,
|
||||
CodexAuth::from_api_key("ambient"),
|
||||
CodexAuth::from_api_key("ambient"),
|
||||
)
|
||||
.await;
|
||||
let command_credentials =
|
||||
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("shared-key"));
|
||||
let command = Arc::new(Credentials {
|
||||
available: std::sync::atomic::AtomicBool::new(true),
|
||||
current: RwLock::new(CodexAuth::from_api_key("shared-key")),
|
||||
refreshed: CodexAuth::from_api_key("shared-key"),
|
||||
});
|
||||
command_credentials
|
||||
.set_external_auth(command.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let context = command_credentials.runtime_config();
|
||||
let command_runtime = root.bind(Arc::clone(&command_credentials));
|
||||
let original_command = command_runtime.auth().await.unwrap();
|
||||
command
|
||||
.available
|
||||
.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
let explicit = root
|
||||
.bind_api_key("shared-key", &context)
|
||||
.auth()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(explicit.auth().api_key(), Some("shared-key"));
|
||||
assert_ne!(explicit.runtime().id(), original_command.runtime().id());
|
||||
let explicit_id = explicit.runtime().id().to_string();
|
||||
drop(explicit);
|
||||
let reused = root
|
||||
.bind_api_key("shared-key", &context)
|
||||
.auth()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(reused.runtime().id(), explicit_id);
|
||||
assert_eq!(
|
||||
command_runtime.auth().await.unwrap().runtime().id(),
|
||||
original_command.runtime().id()
|
||||
);
|
||||
|
||||
command
|
||||
.available
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
let recovered = root.bind(command_credentials).auth().await.unwrap();
|
||||
assert_eq!(recovered.runtime().id(), original_command.runtime().id());
|
||||
*command.current.write().unwrap() = CodexAuth::from_api_key("replacement");
|
||||
assert_ne!(
|
||||
command_runtime.auth().await.unwrap().runtime().id(),
|
||||
original_command.runtime().id()
|
||||
);
|
||||
assert_eq!(
|
||||
root.bind_api_key("shared-key", &context)
|
||||
.auth()
|
||||
.await
|
||||
.unwrap()
|
||||
.runtime()
|
||||
.id(),
|
||||
explicit_id
|
||||
);
|
||||
assert_ne!(
|
||||
root.bind_api_key("other-key", &context)
|
||||
.auth()
|
||||
.await
|
||||
.unwrap()
|
||||
.runtime()
|
||||
.id(),
|
||||
explicit_id
|
||||
);
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_key_stays_independent_when_matching_saved_login_changes() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(registered)
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let credentials = AuthManager::from_auth_for_testing_with_home(
|
||||
CodexAuth::from_api_key("shared-key"),
|
||||
home.path().to_path_buf(),
|
||||
);
|
||||
let manager = Arc::new(RuntimeAuthManager::new(
|
||||
credentials.clone(),
|
||||
AgentHostAuthorization::InstallationId {
|
||||
installation_id: Uuid::nil(),
|
||||
},
|
||||
RuntimeRegistrationMetadata {
|
||||
harness: None,
|
||||
harness_version: None,
|
||||
role: RuntimeRole::Orchestrator,
|
||||
},
|
||||
format!("{}/agent/register", server.uri()).parse().unwrap(),
|
||||
));
|
||||
let explicit = manager.bind_api_key("shared-key", &credentials.runtime_config());
|
||||
let initial = explicit.auth().await.unwrap();
|
||||
codex_login::login_with_api_key(
|
||||
home.path(),
|
||||
"replacement",
|
||||
AuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::Direct,
|
||||
)
|
||||
.unwrap();
|
||||
credentials.reload().await;
|
||||
|
||||
let selected = explicit.auth().await.unwrap();
|
||||
assert_eq!(selected.auth().api_key(), Some("shared-key"));
|
||||
assert_eq!(selected.runtime().id(), initial.runtime().id());
|
||||
assert_eq!(
|
||||
manager.auth().await.unwrap().auth().api_key(),
|
||||
Some("replacement")
|
||||
);
|
||||
assert_eq!(
|
||||
manager
|
||||
.bind_api_key("shared-key", &credentials.runtime_config())
|
||||
.auth()
|
||||
.await
|
||||
.unwrap()
|
||||
.runtime()
|
||||
.id(),
|
||||
initial.runtime().id(),
|
||||
);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user