Merge 4707af72db into sapling-pr-archive-bolinfest

This commit is contained in:
Michael Bolin
2026-04-27 21:29:55 -07:00
committed by GitHub
7 changed files with 48 additions and 191 deletions
+2 -2
View File
@@ -148,7 +148,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools(
let auth = auth_manager.auth().await;
if !config
.features
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
{
return Some(Vec::new());
}
@@ -220,7 +220,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
let auth = auth_manager.auth().await;
if !config
.features
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
{
return Ok(AccessibleConnectorsStatus {
connectors: Vec::new(),
+20 -40
View File
@@ -1,9 +1,6 @@
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::register_agent_task;
use codex_protocol::account::PlanType as AccountPlanType;
use tokio::sync::OnceCell;
use crate::default_client::build_reqwest_client;
@@ -11,50 +8,33 @@ use super::storage::AgentIdentityAuthRecord;
const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
#[derive(Debug)]
#[derive(Clone, Debug)]
pub struct AgentIdentityAuth {
record: AgentIdentityAuthRecord,
process_task_id: Arc<OnceCell<String>>,
}
impl Clone for AgentIdentityAuth {
fn clone(&self) -> Self {
Self {
record: self.record.clone(),
process_task_id: Arc::clone(&self.process_task_id),
}
}
process_task_id: String,
}
impl AgentIdentityAuth {
pub fn new(record: AgentIdentityAuthRecord) -> Self {
Self {
pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
let process_task_id = register_agent_task(
&build_reqwest_client(),
AGENT_IDENTITY_AUTHAPI_BASE_URL,
key(&record),
)
.await
.map_err(std::io::Error::other)?;
Ok(Self {
record,
process_task_id: Arc::new(OnceCell::new()),
}
process_task_id,
})
}
pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.record
}
pub fn process_task_id(&self) -> Option<&str> {
self.process_task_id.get().map(String::as_str)
}
pub async fn ensure_runtime(&self) -> std::io::Result<()> {
self.process_task_id
.get_or_try_init(|| async {
register_agent_task(
&build_reqwest_client(),
AGENT_IDENTITY_AUTHAPI_BASE_URL,
self.key(),
)
.await
.map_err(std::io::Error::other)
})
.await
.map(|_| ())
pub fn process_task_id(&self) -> &str {
&self.process_task_id
}
pub fn account_id(&self) -> &str {
@@ -76,11 +56,11 @@ impl AgentIdentityAuth {
pub fn is_fedramp_account(&self) -> bool {
self.record.chatgpt_account_is_fedramp
}
}
fn key(&self) -> AgentIdentityKey<'_> {
AgentIdentityKey {
agent_runtime_id: &self.record.agent_runtime_id,
private_key_pkcs8_base64: &self.record.agent_private_key,
}
fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
}
}
+3 -28
View File
@@ -625,33 +625,6 @@ impl Drop for EnvVarGuard {
}
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn load_auth_reads_agent_identity_from_env() {
let codex_home = tempdir().unwrap();
let expected_record = agent_identity_record("account-123");
let agent_identity = fake_agent_identity_jwt(&expected_record).expect("fake agent identity");
let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
)
.await
.expect("env auth should load")
.expect("env auth should be present");
let CodexAuth::AgentIdentity(agent_identity) = auth else {
panic!("env auth should load as agent identity");
};
assert_eq!(agent_identity.record(), &expected_record);
assert!(
!get_auth_file(codex_home.path()).exists(),
"env auth should not write auth.json"
);
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn load_auth_keeps_codex_api_key_env_precedence() {
@@ -805,9 +778,11 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
}
fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
let key_material =
codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(),
agent_private_key: key_material.private_key_pkcs8_base64,
account_id: account_id.to_string(),
chatgpt_user_id: "user-id".to_string(),
email: "user@example.com".to_string(),
+7 -20
View File
@@ -212,7 +212,7 @@ impl CodexAuth {
"agent identity auth is missing an agent identity token.",
));
};
return Self::from_agent_identity_jwt(&agent_identity);
return Self::from_agent_identity_jwt(&agent_identity).await;
}
let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
@@ -246,9 +246,9 @@ impl CodexAuth {
.await
}
pub fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
Ok(Self::AgentIdentity(AgentIdentityAuth::new(record)))
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
}
pub fn auth_mode(&self) -> AuthMode {
@@ -322,16 +322,6 @@ impl CodexAuth {
}
}
pub async fn initialize_runtime(
&self,
_chatgpt_base_url: Option<String>,
) -> std::io::Result<()> {
match self {
Self::AgentIdentity(auth) => auth.ensure_runtime().await,
Self::ApiKey(_) | Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => Ok(()),
}
}
/// Returns `None` if Codex backend auth does not expose an account id.
pub fn get_account_id(&self) -> Option<String> {
match self {
@@ -749,7 +739,9 @@ async fn load_auth(
}
if let Some(agent_identity) = read_codex_agent_identity_from_env() {
return CodexAuth::from_agent_identity_jwt(&agent_identity).map(Some);
return CodexAuth::from_agent_identity_jwt(&agent_identity)
.await
.map(Some);
}
// Fall back to the configured persistent store (file/keyring/auto) for managed auth.
@@ -1400,12 +1392,7 @@ impl AuthManager {
tracing::error!("Failed to refresh token: {}", err);
return Some(auth);
}
let auth = self.auth_cached()?;
if let Err(err) = auth.initialize_runtime(self.chatgpt_base_url.clone()).await {
tracing::error!("Failed to initialize auth runtime: {err}");
return None;
}
Some(auth)
self.auth_cached()
}
/// Force a reload of the auth information from auth.json. Returns
+11 -17
View File
@@ -21,23 +21,17 @@ struct AgentIdentityAuthProvider {
impl AuthProvider for AgentIdentityAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
let record = self.auth.record();
let header_value = self
.auth
.process_task_id()
.ok_or_else(|| std::io::Error::other("agent identity process task is not initialized"))
.and_then(|task_id| {
authorization_header_for_agent_task(
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
},
AgentTaskAuthorizationTarget {
agent_runtime_id: &record.agent_runtime_id,
task_id,
},
)
.map_err(std::io::Error::other)
});
let header_value = authorization_header_for_agent_task(
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
},
AgentTaskAuthorizationTarget {
agent_runtime_id: &record.agent_runtime_id,
task_id: self.auth.process_task_id(),
},
)
.map_err(std::io::Error::other);
if let Ok(header_value) = header_value
&& let Ok(header) = HeaderValue::from_str(&header_value)
@@ -9,9 +9,6 @@ use codex_login::ExternalAuth;
use codex_login::ExternalAuthRefreshContext;
use codex_login::ExternalAuthTokens;
use codex_login::TokenData;
use codex_login::auth::AgentIdentityAuth;
use codex_login::auth::AgentIdentityAuthRecord;
use codex_protocol::account::PlanType;
use codex_protocol::openai_models::ModelsResponse;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -237,18 +234,6 @@ c2ln",
.expect("auth should be present")
}
fn agent_identity_auth_for_tests() -> CodexAuth {
CodexAuth::AgentIdentity(AgentIdentityAuth::new(AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "agent-private-key".to_string(),
account_id: "account-id".to_string(),
chatgpt_user_id: "chatgpt-user-id".to_string(),
email: "agent@example.com".to_string(),
plan_type: PlanType::Pro,
chatgpt_account_is_fedramp: false,
}))
}
#[tokio::test]
async fn get_model_info_tracks_fallback_usage() {
let codex_home = tempdir().expect("temp dir");
@@ -713,43 +698,6 @@ async fn refresh_available_models_fetches_with_chatgpt_auth_tokens() {
);
}
#[tokio::test]
async fn refresh_available_models_fetches_with_agent_identity() {
let dynamic_slug = "dynamic-model-only-for-test-agent-identity";
let codex_home = tempdir().expect("temp dir");
let endpoint = TestModelsEndpoint::new(vec![vec![remote_model(
dynamic_slug,
"Agent Identity",
/*priority*/ 1,
)]]);
let manager = openai_manager_for_tests_with_auth(
codex_home.path().to_path_buf(),
endpoint.clone(),
Some(AuthManager::from_auth_for_testing(
agent_identity_auth_for_tests(),
)),
);
manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("refresh should fetch with agent identity");
assert!(
manager
.get_remote_models()
.await
.iter()
.any(|candidate| candidate.slug == dynamic_slug),
"remote refresh should include models fetched with agent identity"
);
assert_eq!(
endpoint.fetch_count(),
1,
"endpoint should fetch models with agent identity"
);
}
#[test]
fn build_available_models_picks_default_after_hiding_hidden_models() {
let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() });
@@ -768,35 +716,6 @@ fn build_available_models_picks_default_after_hiding_hidden_models() {
assert_eq!(available, vec![expected_hidden, expected_visible]);
}
#[tokio::test]
async fn static_manager_treats_agent_identity_as_backend_auth_for_filtering() {
let chatgpt_only_model = {
let mut model = remote_model("chatgpt-only", "ChatGPT Only", /*priority*/ 0);
model.supported_in_api = false;
model
};
let api_model = remote_model("api-model", "API Model", /*priority*/ 1);
let manager = StaticModelsManager::new(
Some(AuthManager::from_auth_for_testing(
agent_identity_auth_for_tests(),
)),
ModelsResponse {
models: vec![chatgpt_only_model, api_model],
},
CollaborationModesConfig::default(),
);
let agent_identity_models = manager.list_models(RefreshStrategy::Online).await;
assert_eq!(
agent_identity_models
.iter()
.map(|model| model.model.as_str())
.collect::<Vec<_>>(),
vec!["chatgpt-only", "api-model"]
);
}
#[tokio::test]
async fn static_manager_reads_latest_auth_mode() {
let auth_manager =
+5 -3
View File
@@ -196,6 +196,8 @@ impl SlashCommand {
| SlashCommand::Mcp
| SlashCommand::Apps
| SlashCommand::Plugins
| SlashCommand::Title
| SlashCommand::Statusline
| SlashCommand::AutoReview
| SlashCommand::Feedback
| SlashCommand::Quit
@@ -207,9 +209,7 @@ impl SlashCommand {
SlashCommand::Settings => true,
SlashCommand::Collab => true,
SlashCommand::Agent | SlashCommand::MultiAgents => true,
SlashCommand::Statusline => false,
SlashCommand::Theme => false,
SlashCommand::Title => false,
}
}
@@ -249,8 +249,10 @@ mod tests {
}
#[test]
fn goal_command_is_available_during_task() {
fn certain_commands_are_available_during_task() {
assert!(SlashCommand::Goal.available_during_task());
assert!(SlashCommand::Title.available_during_task());
assert!(SlashCommand::Statusline.available_during_task());
}
#[test]