agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: use run agent task auth for inference (#19051)

Browse Source 
## Stack

This is PR 3 of the simplified HAI single-run-task stack:

- [#19047](https://github.com/openai/codex/pull/19047) Agent Identity
assertion and task-registration primitives, including the shared
run-task helper used by existing Agent Identity JWT auth.
- [#19049](https://github.com/openai/codex/pull/19049)
Disabled-by-default ChatGPT auth opt-in that provisions/reuses persisted
Agent Identity runtime auth and its single run task.
- [#19051](https://github.com/openai/codex/pull/19051) Run-scoped
provider auth that uses one backend-owned task id for first-party
inference and compaction requests.

[#19054](https://github.com/openai/codex/pull/19054) collapsed out of
the active stack because the simplified design no longer needs a
separate background/control-plane task helper.

## Summary

This PR moves Agent Identity usage into provider auth resolution. That
keeps `AgentAssertion` auth tied to first-party OpenAI provider requests
instead of applying a late session-wide override that could affect
local, custom, Bedrock, API-key, or external-bearer providers.

What changed:

- adds a small `ProviderAuthScope` struct carrying the run auth policy
and session source needed by provider-scoped auth resolution
- lets `Session` opt the existing `ModelClient` into `ChatGptAuth`
policy when `use_agent_identity` is enabled, without adding a second
model-client constructor
- resolves Agent Identity only for first-party OpenAI provider auth
paths
- uses the persisted run task id from the `AgentIdentityAuth` record to
build `AgentAssertion` auth for Responses requests
- routes shared request setup through scoped provider auth so unary
compact requests use the same run-task assertion path as inference turns
- keeps local/custom/Bedrock/env-key/external-bearer provider auth
unchanged
- lets missing run-task state surface through the existing model-request
error path instead of silently falling back to bearer auth

This PR intentionally does not create thread-scoped, target-scoped, or
background-scoped task identities. The run task is the only task Codex
registers in this POC shape.

## Testing

- `just test -p codex-model-provider`
- `just test -p codex-core client::tests::provider_auth_scope_uses`
- `just test -p codex-core remote_compact_uses_agent_identity_assertion`
This commit is contained in:
Adrian
2026-06-24 22:31:41 -07:00
committed by GitHub
parent f66d793a2d
commit 51864b0b4b
22 changed files with 1020 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/auth.rs
+6
View File
@@ -67,6 +67,12 @@ pub type AuthProviderFuture<'a> =
/// Shared auth handle passed through API clients.
pub type SharedAuthProvider = Arc<dyn AuthProvider>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentIdentityTelemetry {
pub agent_id: String,
pub task_id: String,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct AuthHeaderTelemetry {
pub attached: bool,
codex-rs/codex-api/src/lib.rs
+1
View File
@@ -19,6 +19,7 @@ pub use codex_client::ReqwestTransport;
pub use codex_client::TransportError;
pub use crate::api_bridge::map_api_error;
pub use crate::auth::AgentIdentityTelemetry;
pub use crate::auth::AuthError;
pub use crate::auth::AuthHeaderTelemetry;
pub use crate::auth::AuthProvider;
codex-rs/core/src/client.rs
+42 -4
View File
@@ -30,6 +30,7 @@ use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::AuthProvider;
use codex_api::CompactClient as ApiCompactClient;
@@ -115,8 +116,11 @@ use crate::responses_metadata::subagent_header_value;
use crate::util::emit_feedback_auth_recovery_tags;
use codex_feedback::FeedbackRequestTags;
use codex_feedback::emit_feedback_request_tags_with_auth_env;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_login::auth_env_telemetry::AuthEnvTelemetry;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_model_provider::AgentIdentitySessionFallback;
use codex_model_provider::ProviderAuthScope;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
#[cfg(test)]
@@ -202,6 +206,7 @@ struct ModelClientState {
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
disable_websockets: AtomicBool,
agent_identity_session_fallback: AgentIdentitySessionFallback,
cached_websocket_session: StdMutex<WebsocketSession>,
}
@@ -213,6 +218,7 @@ struct CurrentClientSetup {
auth: Option<CodexAuth>,
api_provider: ApiProvider,
api_auth: SharedAuthProvider,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
}
#[derive(Clone, Copy)]
@@ -240,6 +246,7 @@ impl RequestRouteTelemetry {
#[derive(Debug, Clone)]
pub struct ModelClient {
state: Arc<ModelClientState>,
agent_identity_policy: AgentIdentityAuthPolicy,
prompt_cache_key_override: Option<String>,
}
@@ -392,6 +399,7 @@ impl ModelClient {
/// are passed to [`ModelClientSession::stream`] (and other turn-scoped methods) explicitly.
pub fn new(
auth_manager: Option<Arc<AuthManager>>,
agent_identity_policy: AgentIdentityAuthPolicy,
thread_id: ThreadId,
provider_info: ModelProviderInfo,
session_source: SessionSource,
@@ -426,8 +434,10 @@ impl ModelClient {
include_attestation,
attestation_provider,
disable_websockets: AtomicBool::new(false),
agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
}),
agent_identity_policy,
prompt_cache_key_override: None,
}
}
@@ -528,6 +538,7 @@ impl ModelClient {
AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
),
RequestRouteTelemetry::for_endpoint(RESPONSES_COMPACT_ENDPOINT),
@@ -660,6 +671,7 @@ impl ModelClient {
AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
),
RequestRouteTelemetry::for_endpoint(MEMORIES_SUMMARIZE_ENDPOINT),
@@ -909,14 +921,27 @@ impl ModelClient {
async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
let auth = self.state.provider.auth().await;
let api_provider = self.state.provider.api_provider().await?;
let api_auth = self.state.provider.api_auth().await?;
let resolved_auth = self
.state
.provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: self.agent_identity_policy,
session_source: self.state.session_source.clone(),
agent_identity_session_fallback: self.state.agent_identity_session_fallback.clone(),
})
.await?;
Ok(CurrentClientSetup {
auth,
api_provider,
api_auth,
api_auth: resolved_auth.auth,
agent_identity_telemetry: resolved_auth.agent_identity_telemetry,
})
}
pub(crate) async fn prewarm_auth(&self) -> Result<()> {
self.current_client_setup().await.map(|_| ())
}
/// Opens a websocket connection using the same header and telemetry wiring as normal turns.
///
/// Both startup prewarm and in-turn `needs_new` reconnects call this path so handshake
@@ -934,7 +959,7 @@ impl ModelClient {
let headers = self.build_websocket_headers(responses_metadata).await;
let websocket_telemetry = ModelClientSession::build_websocket_telemetry(
session_telemetry,
auth_context,
auth_context.clone(),
request_route_telemetry,
self.state.auth_env_telemetry.clone(),
);
@@ -976,6 +1001,7 @@ impl ModelClient {
response_debug.cf_ray.as_deref(),
response_debug.auth_error.as_deref(),
response_debug.auth_error_code.as_deref(),
auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
@@ -1205,6 +1231,7 @@ impl ModelClientSession {
let auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
);
let connection = self
@@ -1343,6 +1370,7 @@ impl ModelClientSession {
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
pending_retry,
);
let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry(
@@ -1470,6 +1498,7 @@ impl ModelClientSession {
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
pending_retry,
);
let request = self.client.build_responses_request(
@@ -2035,11 +2064,12 @@ impl PendingUnauthorizedRetry {
}
}
#[derive(Clone, Copy, Debug, Default)]
#[derive(Clone, Debug, Default)]
struct AuthRequestTelemetryContext {
auth_mode: Option<&'static str>,
auth_header_attached: bool,
auth_header_name: Option<&'static str>,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
retry_after_unauthorized: bool,
recovery_mode: Option<&'static str>,
recovery_phase: Option<&'static str>,
@@ -2049,6 +2079,7 @@ impl AuthRequestTelemetryContext {
fn new(
auth_mode: Option<AuthMode>,
api_auth: &dyn AuthProvider,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
retry: PendingUnauthorizedRetry,
) -> Self {
let auth_telemetry = auth_header_telemetry(api_auth);
@@ -2062,11 +2093,16 @@ impl AuthRequestTelemetryContext {
}),
auth_header_attached: auth_telemetry.attached,
auth_header_name: auth_telemetry.name,
agent_identity_telemetry,
retry_after_unauthorized: retry.retry_after_unauthorized,
recovery_mode: retry.recovery_mode,
recovery_phase: retry.recovery_phase,
}
}
fn agent_identity_telemetry(&self) -> Option<&AgentIdentityTelemetry> {
self.agent_identity_telemetry.as_ref()
}
}
struct WebsocketConnectParams<'a> {
@@ -2253,6 +2289,7 @@ impl RequestTelemetry for ApiTelemetry {
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
self.auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
@@ -2307,6 +2344,7 @@ impl WebsocketTelemetry for ApiTelemetry {
duration,
error_message.as_deref(),
connection_reused,
self.auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
codex-rs/core/src/client_tests.rs
+193 -1
View File
@@ -1,6 +1,8 @@
use super::AuthRequestTelemetryContext;
use super::CompactConversationRequestSettings;
use super::ModelClient;
use super::PendingUnauthorizedRetry;
use super::Prompt;
use super::UnauthorizedRecoveryExecution;
use super::X_CODEX_INSTALLATION_ID_HEADER;
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
@@ -13,11 +15,15 @@ use crate::GenerateAttestationFuture;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::test_support::TestCodexResponsesRequestKind;
use crate::test_support::responses_metadata as test_responses_metadata;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_api::TransportError;
use codex_login::AuthCredentialsStoreMode;
use codex_login::AuthKeyringBackendKind;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider::BearerAuthProvider;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
@@ -28,6 +34,7 @@ use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_otel::SessionTelemetry;
use codex_protocol::ThreadId;
use codex_protocol::auth::AuthMode;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
@@ -35,6 +42,7 @@ use codex_protocol::openai_models::ReasoningEffort;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_rollout_trace::CompactionTraceContext;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::InferenceTraceAttempt;
use codex_rollout_trace::InferenceTraceContext;
@@ -65,14 +73,27 @@ use tracing_subscriber::layer::Context as LayerContext;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::util::SubscriberInitExt;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
const TEST_CHATGPT_ID_TOKEN: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7ImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzNDUiLCJ1c2VyX2lkIjoidXNlci0xMjM0NSIsImNoYXRncHRfcGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjb3VudC0xMjMifX0.c2ln";
const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
fn test_model_client(session_source: SessionSource) -> ModelClient {
test_model_client_with_thread_id(ThreadId::new(), session_source)
}
fn test_model_client_with_thread_id(
thread_id: ThreadId,
session_source: SessionSource,
) -> ModelClient {
let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
let thread_id = ThreadId::new();
ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider,
session_source,
@@ -86,6 +107,115 @@ fn test_model_client(session_source: SessionSource) -> ModelClient {
)
}
#[tokio::test]
async fn compact_uses_bearer_after_agent_identity_session_fallback() -> anyhow::Result<()> {
let server = MockServer::start().await;
let registration_count = Arc::new(AtomicUsize::new(0));
let response_count = Arc::clone(&registration_count);
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.respond_with(move |_request: &wiremock::Request| {
response_count.fetch_add(1, Ordering::SeqCst);
ResponseTemplate::new(/*status*/ 503)
})
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/responses/compact"))
.respond_with(ResponseTemplate::new(/*status*/ 200).set_body_json(json!({
"output": []
})))
.expect(/*requests*/ 1)
.mount(&server)
.await;
let codex_home = TempDir::new()?;
let auth_manager = chatgpt_auth_manager(&codex_home, server.uri()).await;
let mut provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
provider.base_url = Some(format!("{}/v1", server.uri()));
provider.supports_websockets = false;
let thread_id = ThreadId::new();
let client = ModelClient::new(
Some(auth_manager),
AgentIdentityAuthPolicy::ChatGptAuth,
thread_id,
provider,
SessionSource::Cli,
"test_originator".to_string(),
/*model_verbosity*/ None,
/*enable_request_compression*/ false,
/*include_timing_metrics*/ false,
/*beta_features_header*/ None,
/*item_ids_enabled*/ false,
/*attestation_provider*/ None,
);
let prompt = Prompt {
input: vec![ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "please compact".to_string(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
}],
base_instructions: BaseInstructions {
text: "base instructions".to_string(),
},
..Default::default()
};
let responses_metadata = test_responses_metadata_for_client(
&client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
);
let output = client
.compact_conversation_history(
&prompt,
&test_model_info(),
/*turn_state*/ None,
CompactConversationRequestSettings {
effort: None,
summary: codex_protocol::config_types::ReasoningSummary::None,
service_tier: None,
},
&test_session_telemetry(),
&CompactionTraceContext::disabled(),
&responses_metadata,
)
.await?;
assert!(output.is_empty());
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
let requests = server
.received_requests()
.await
.expect("server should record requests");
let compact_request = requests
.iter()
.find(|request| request.url.path() == "/v1/responses/compact")
.expect("compact request should be captured");
assert_eq!(
compact_request
.headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok()),
Some("Bearer test-access-token")
);
assert_eq!(
compact_request
.headers
.get("ChatGPT-Account-ID")
.and_then(|value| value.to_str().ok()),
Some("account-123")
);
Ok(())
}
fn test_model_provider() -> SharedModelProvider {
test_model_client(SessionSource::Cli).state.provider.clone()
}
@@ -169,6 +299,45 @@ fn ultra_reasoning_uses_max_for_requests() {
);
}
fn write_chatgpt_auth_json(codex_home: &std::path::Path) {
let auth_json = json!({
"tokens": {
"id_token": TEST_CHATGPT_ID_TOKEN,
"access_token": "test-access-token",
"refresh_token": "test-refresh-token",
"account_id": "account-123"
},
"last_refresh": "2099-01-01T00:00:00Z"
});
std::fs::write(
codex_home.join("auth.json"),
serde_json::to_string_pretty(&auth_json).expect("serialize auth.json"),
)
.expect("write auth.json");
}
async fn chatgpt_auth_manager(
codex_home: &TempDir,
agent_identity_authapi_base_url: String,
) -> Arc<AuthManager> {
write_chatgpt_auth_json(codex_home.path());
let auth_manager = AuthManager::shared(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
/*auth_route_config*/ None,
)
.await;
let auth = auth_manager.auth().await.expect("auth should load");
AuthManager::from_auth_for_testing_with_agent_identity_authapi_base_url(
auth,
agent_identity_authapi_base_url,
)
}
#[derive(Default)]
struct TagCollectorVisitor {
tags: BTreeMap<String, String>,
@@ -572,6 +741,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
let auth_context = AuthRequestTelemetryContext::new(
Some(AuthMode::Chatgpt),
&BearerAuthProvider::for_test(Some("access-token"), Some("workspace-123")),
/*agent_identity_telemetry*/ None,
PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution {
mode: "managed",
phase: "refresh_token",
@@ -586,6 +756,27 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
assert_eq!(auth_context.recovery_phase, Some("refresh_token"));
}
#[test]
fn auth_request_telemetry_context_tracks_agent_identity_ids() {
let auth_context = AuthRequestTelemetryContext::new(
Some(AuthMode::Chatgpt),
&BearerAuthProvider::for_test(/*token*/ None, /*account_id*/ None),
Some(AgentIdentityTelemetry {
agent_id: "agent-runtime-context".to_string(),
task_id: "task-run-context".to_string(),
}),
PendingUnauthorizedRetry::default(),
);
assert_eq!(
auth_context.agent_identity_telemetry(),
Some(&AgentIdentityTelemetry {
agent_id: "agent-runtime-context".to_string(),
task_id: "task-run-context".to_string(),
})
);
}
fn model_client_with_counting_attestation(
include_attestation: bool,
) -> (ModelClient, Arc<AtomicUsize>) {
@@ -623,6 +814,7 @@ fn model_client_with_counting_attestation(
};
let model_client = ModelClient::new(
auth_manager,
AgentIdentityAuthPolicy::JwtOnly,
ThreadId::new(),
provider,
SessionSource::Exec,
codex-rs/core/src/session/session.rs
+6
View File
@@ -8,6 +8,7 @@ use crate::shell_snapshot::ShellSnapshot;
use crate::skills::SkillError;
use crate::state::ActiveTurn;
use codex_extension_api::ExtensionDataInit;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_protocol::SessionId;
use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE;
use codex_protocol::config_types::ServiceTier;
@@ -1057,6 +1058,11 @@ impl Session {
time_provider,
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
if config.features.enabled(Feature::UseAgentIdentity) {
AgentIdentityAuthPolicy::ChatGptAuth
} else {
AgentIdentityAuthPolicy::JwtOnly
},
thread_id,
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
codex-rs/core/src/session/tests.rs
+4
View File
@@ -30,6 +30,7 @@ use core_test_support::test_codex::local_selections;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider_info::ModelProviderInfo;
use codex_models_manager::bundled_models_response;
use codex_models_manager::model_info;
@@ -469,6 +470,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession {
.expect("test thread id should be valid");
crate::client::ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
ModelProviderInfo::create_openai_provider(/* base_url */ /*base_url*/ None),
codex_protocol::protocol::SessionSource::Exec,
@@ -5415,6 +5417,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
time_provider: Arc::new(crate::current_time::SystemTimeProvider),
model_client: ModelClient::new(
Some(auth_manager.clone()),
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
@@ -7492,6 +7495,7 @@ where
time_provider: Arc::new(crate::current_time::SystemTimeProvider),
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
codex-rs/core/src/session_startup_prewarm.rs
+8
View File
@@ -183,6 +183,14 @@ impl SessionStartupPrewarmHandle {
impl Session {
pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>, base_instructions: String) {
if !self.services.model_client.responses_websocket_enabled() {
// Without websocket prewarm, resolve auth once so Agent Identity bootstrap can
// register or engage this session's bearer fallback before the first user request.
let model_client = self.services.model_client.clone();
tokio::spawn(async move {
if let Err(err) = model_client.prewarm_auth().await {
warn!("startup auth prewarm failed: {err:#}");
}
});
return;
}
codex-rs/core/tests/responses_headers.rs
+4
View File
@@ -5,6 +5,7 @@ use codex_core::ModelClient;
use codex_core::Prompt;
use codex_core::ResponseEvent;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_otel::SessionTelemetry;
@@ -120,6 +121,7 @@ async fn responses_stream_includes_subagent_header_on_review() {
let client = ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider.clone(),
session_source.clone(),
@@ -253,6 +255,7 @@ async fn responses_stream_includes_subagent_header_on_other() {
let client = ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider.clone(),
session_source.clone(),
@@ -372,6 +375,7 @@ async fn responses_respects_model_info_overrides_from_config() {
let client = ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider.clone(),
session_source.clone(),
codex-rs/core/tests/suite/client.rs
+3
View File
@@ -12,6 +12,7 @@ use codex_features::Feature;
use codex_login::AuthKeyringBackendKind;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_login::default_client::originator;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
@@ -1220,6 +1221,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
"unused-api-key",
))),
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider,
SessionSource::Exec,
@@ -2833,6 +2835,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
let client = ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider.clone(),
SessionSource::Exec,
codex-rs/core/tests/suite/client_websockets.rs
+2
View File
@@ -9,6 +9,7 @@ use codex_core::ResponseEvent;
use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_otel::MetricsClient;
@@ -2189,6 +2190,7 @@ async fn websocket_harness_with_provider_options(
let summary = ReasoningSummary::Auto;
let client = ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider.clone(),
SessionSource::Exec,
codex-rs/core/tests/suite/compact_remote.rs
+78
View File
@@ -5,6 +5,9 @@ use anyhow::Result;
use codex_core::compact::SUMMARY_PREFIX;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuth;
use codex_login::auth::AgentIdentityAuthRecord;
use codex_protocol::account::PlanType as AccountPlanType;
use codex_protocol::config_types::ServiceTier;
use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec;
@@ -133,6 +136,8 @@ const PRETURN_CONTEXT_DIFF_CWD: &str = "/tmp/PRETURN_CONTEXT_DIFF_CWD";
const DUMMY_FUNCTION_NAME: &str = "test_tool";
const TURN_STATE_HEADER: &str = "x-codex-turn-state";
const REMOTE_COMPACT_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30);
const TEST_AGENT_IDENTITY_PRIVATE_KEY: &str =
"MC4CAQAwBQYDK2VwBCIEIJ7kFBaOujmoz1gvBNEC+BeM2IX87FFB0xmISOZ/XO0c";
fn summary_with_prefix(summary: &str) -> String {
format!("{SUMMARY_PREFIX}\n{summary}")
@@ -532,6 +537,79 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_compact_uses_agent_identity_assertion() -> Result<()> {
skip_if_no_network!(Ok(()));
let harness = TestCodexHarness::with_builder(
test_codex().with_auth(CodexAuth::AgentIdentity(
AgentIdentityAuth::from_record(
AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-compact".to_string(),
agent_private_key: TEST_AGENT_IDENTITY_PRIVATE_KEY.to_string(),
account_id: "account-compact".to_string(),
chatgpt_user_id: "user-compact".to_string(),
email: Some("agent@example.com".to_string()),
plan_type: AccountPlanType::Plus,
chatgpt_account_is_fedramp: false,
task_id: Some("task-compact".to_string()),
},
"https://auth.openai.com/api/accounts",
/*auth_route_config*/ None,
)
.await?,
)),
)
.await?;
let codex = harness.test().codex.clone();
let _responses_mock = responses::mount_sse_once(
harness.server(),
responses::sse(vec![
responses::ev_assistant_message("m1", "REMOTE_REPLY"),
responses::ev_completed("resp-1"),
]),
)
.await;
let compact_mock = responses::mount_compact_json_once(
harness.server(),
serde_json::json!({ "output": compacted_summary_only_output("COMPACTED") }),
)
.await;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello remote compact".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
wait_for_turn_complete(&codex).await;
codex.submit(Op::Compact).await?;
wait_for_turn_complete(&codex).await;
let compact_request = compact_mock.single_request();
assert_eq!(compact_request.path(), "/v1/responses/compact");
assert!(
compact_request
.header("authorization")
.is_some_and(|value| value.starts_with("AgentAssertion ")),
"compact request should use task-scoped AgentAssertion auth"
);
assert_eq!(
compact_request.header("chatgpt-account-id").as_deref(),
Some("account-compact")
);
Ok(())
}
async fn assert_remote_manual_compact_request_parity(
auth: CodexAuth,
configured_service_tier: Option<ServiceTier>,
codex-rs/login/src/auth/agent_identity.rs
+9 -8
View File
@@ -46,7 +46,7 @@ pub(super) fn require_agent_identity_authapi_base_url(
})
}
#[derive(Debug, Error)]
#[derive(Clone, Debug, Error)]
pub enum AgentIdentityAuthError {
#[error(
"agent identity bootstrap unavailable after {attempts} attempts during {operation}: {message}"
@@ -59,13 +59,14 @@ pub enum AgentIdentityAuthError {
}
impl AgentIdentityAuthError {
pub fn is_bootstrap_unavailable(error: &std::io::Error) -> bool {
matches!(
error
.get_ref()
.and_then(|source| source.downcast_ref::<Self>()),
Some(Self::BootstrapUnavailable { .. })
)
pub(super) fn bootstrap_unavailable(error: &std::io::Error) -> Option<&Self> {
match error
.get_ref()
.and_then(|source| source.downcast_ref::<Self>())
{
Some(error @ Self::BootstrapUnavailable { .. }) => Some(error),
None => None,
}
}
}
codex-rs/login/src/auth/auth_tests.rs
+4 -4
View File
@@ -584,7 +584,7 @@ async fn chatgpt_auth_registration_retry_exhaustion_is_fallback_eligible() -> an
.await
.expect_err("retry exhaustion should return an error");
assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_some());
assert!(
auth.stored_managed_chatgpt_agent_identity_record("account-123")
.is_none()
@@ -648,7 +648,7 @@ async fn chatgpt_auth_task_registration_retry_exhaustion_is_fallback_eligible()
.await
.expect_err("task retry exhaustion should return an error");
assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_some());
record.task_id = None;
assert_eq!(
auth.stored_managed_chatgpt_agent_identity_record("account-123"),
@@ -701,7 +701,7 @@ async fn chatgpt_auth_non_retryable_registration_error_is_hard_failure() -> anyh
.await
.expect_err("hard registration failure should return an error");
assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_none());
assert!(
auth.stored_managed_chatgpt_agent_identity_record("account-123")
.is_none()
@@ -742,7 +742,7 @@ async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> an
.await
.expect_err("agent identity jwt task retry exhaustion should fail");
assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_none());
Ok(())
}
codex-rs/login/src/auth/manager.rs
+96 -2
View File
@@ -15,6 +15,8 @@ use std::sync::Mutex;
use std::sync::RwLock;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use tokio::sync::Semaphore;
use tokio::sync::watch;
use tracing::instrument;
@@ -83,6 +85,63 @@ pub enum AgentIdentityAuthPolicy {
ChatGptAuth,
}
const AGENT_IDENTITY_BOOTSTRAP_FAILURE_COOLDOWN: Duration = Duration::from_secs(60 * 60);
#[derive(Debug)]
struct CachedAgentIdentityBootstrapFailure {
account_id: String,
authapi_base_url: String,
retry_at: Instant,
error: AgentIdentityAuthError,
}
#[derive(Debug, Default)]
struct AgentIdentityBootstrapCooldown {
failure: Option<CachedAgentIdentityBootstrapFailure>,
}
impl AgentIdentityBootstrapCooldown {
fn error_for(
&mut self,
account_id: &str,
authapi_base_url: &str,
now: Instant,
) -> Option<AgentIdentityAuthError> {
let error = self
.failure
.as_ref()
.filter(|failure| {
failure.account_id == account_id
&& failure.authapi_base_url == authapi_base_url
&& failure.retry_at > now
})
.map(|failure| failure.error.clone());
if error.is_none() {
self.clear();
}
error
}
fn record_failure(
&mut self,
account_id: String,
authapi_base_url: String,
error: AgentIdentityAuthError,
now: Instant,
) {
self.failure = Some(CachedAgentIdentityBootstrapFailure {
account_id,
authapi_base_url,
retry_at: now + AGENT_IDENTITY_BOOTSTRAP_FAILURE_COOLDOWN,
error,
});
}
fn clear(&mut self) {
self.failure = None;
}
}
impl PartialEq for CodexAuth {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
@@ -1742,6 +1801,7 @@ pub struct AuthManager {
agent_identity_authapi_base_url: Option<String>,
refresh_lock: Semaphore,
agent_identity_lock: Semaphore,
agent_identity_bootstrap_cooldown: Mutex<AgentIdentityBootstrapCooldown>,
external_auth: RwLock<Option<Arc<dyn ExternalAuth>>>,
auth_route_config: Option<AuthRouteConfig>,
}
@@ -1843,6 +1903,7 @@ impl AuthManager {
agent_identity_authapi_base_url,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
agent_identity_bootstrap_cooldown: Mutex::default(),
external_auth: RwLock::new(None),
auth_route_config,
}
@@ -1868,6 +1929,7 @@ impl AuthManager {
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
agent_identity_bootstrap_cooldown: Mutex::default(),
external_auth: RwLock::new(None),
auth_route_config: None,
})
@@ -1892,6 +1954,7 @@ impl AuthManager {
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
agent_identity_bootstrap_cooldown: Mutex::default(),
external_auth: RwLock::new(None),
auth_route_config: None,
})
@@ -1924,6 +1987,7 @@ impl AuthManager {
),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
agent_identity_bootstrap_cooldown: Mutex::default(),
external_auth: RwLock::new(None),
auth_route_config: None,
})
@@ -1946,6 +2010,7 @@ impl AuthManager {
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
agent_identity_bootstrap_cooldown: Mutex::default(),
external_auth: RwLock::new(Some(
Arc::new(BearerTokenRefresher::new(config)) as Arc<dyn ExternalAuth>
)),
@@ -2006,15 +2071,44 @@ impl AuthManager {
.acquire()
.await
.map_err(std::io::Error::other)?;
return auth
let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id();
let cooldown_key = ManagedChatGptAgentIdentityBinding::from_auth(
&auth,
forced_chatgpt_workspace_id.clone(),
)
.and_then(|binding| {
self.agent_identity_authapi_base_url
.as_ref()
.map(|base_url| (binding.account_id, base_url.clone()))
});
if let Some((account_id, authapi_base_url)) = cooldown_key.as_ref()
&& let Ok(mut cooldown) = self.agent_identity_bootstrap_cooldown.lock()
&& let Some(error) =
cooldown.error_for(account_id, authapi_base_url, Instant::now())
{
tracing::warn!("agent identity bootstrap retry suppressed during shared cooldown");
return Err(std::io::Error::other(error));
}
let result = auth
.agent_identity_auth(
policy,
self.agent_identity_authapi_base_url.as_deref(),
self.forced_chatgpt_workspace_id(),
forced_chatgpt_workspace_id,
self.auth_route_config.as_ref(),
session_source,
)
.await;
if let Ok(mut cooldown) = self.agent_identity_bootstrap_cooldown.lock() {
if let (Err(err), Some((account_id, authapi_base_url))) = (&result, cooldown_key)
&& let Some(error) = AgentIdentityAuthError::bootstrap_unavailable(err).cloned()
{
cooldown.record_failure(account_id, authapi_base_url, error, Instant::now());
} else {
cooldown.clear();
}
}
return result;
}
auth.agent_identity_auth(
policy,
codex-rs/memories/write/src/runtime.rs
+2
View File
@@ -12,6 +12,7 @@ use codex_core::resolve_installation_id;
use codex_features::Feature;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_login::default_client::originator;
use codex_model_provider::ModelProvider;
@@ -250,6 +251,7 @@ impl MemoryStartupContext {
let session_id_string = session_id.to_string();
let model_client = ModelClient::new(
Some(Arc::clone(&self.auth_manager)),
AgentIdentityAuthPolicy::JwtOnly,
self.thread_id,
config.model_provider.clone(),
session_source.clone(),
codex-rs/model-provider/src/auth.rs
+420 -1
View File
@@ -1,13 +1,20 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::authorization_header_for_agent_task;
use codex_api::AgentIdentityTelemetry;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuth;
use codex_login::auth::AgentIdentityAuthError;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider_info::ModelProviderInfo;
use codex_protocol::error::CodexErr;
use codex_protocol::protocol::SessionSource;
use http::HeaderMap;
use http::HeaderValue;
@@ -16,9 +23,62 @@ use crate::bearer_auth_provider::BearerAuthProvider;
const BEDROCK_API_KEY_UNSUPPORTED_MESSAGE: &str =
"Bedrock API key auth is only supported by the Amazon Bedrock model provider";
#[derive(Clone, Debug)]
pub struct ProviderAuthScope {
pub agent_identity_policy: AgentIdentityAuthPolicy,
pub session_source: SessionSource,
pub agent_identity_session_fallback: AgentIdentitySessionFallback,
}
#[derive(Clone, Debug, Default)]
pub struct AgentIdentitySessionFallback {
engaged: Arc<AtomicBool>,
}
impl AgentIdentitySessionFallback {
pub fn is_engaged(&self) -> bool {
self.engaged.load(Ordering::Relaxed)
}
fn engage(&self) -> bool {
!self.engaged.swap(true, Ordering::Relaxed)
}
}
/// Provider auth resolved for a request, plus metadata describing the effective auth.
#[derive(Clone)]
pub struct ResolvedProviderAuth {
pub auth: SharedAuthProvider,
pub agent_identity_telemetry: Option<AgentIdentityTelemetry>,
}
impl ResolvedProviderAuth {
pub(crate) fn new(auth: SharedAuthProvider) -> Self {
Self {
auth,
agent_identity_telemetry: None,
}
}
fn for_agent_identity(auth: AgentIdentityAuth) -> Self {
let agent_identity_telemetry = agent_identity_telemetry(&auth);
Self {
auth: Arc::new(AgentIdentityAuthProvider { auth }),
agent_identity_telemetry: Some(agent_identity_telemetry),
}
}
}
pub(crate) fn agent_identity_telemetry(auth: &AgentIdentityAuth) -> AgentIdentityTelemetry {
AgentIdentityTelemetry {
agent_id: auth.record().agent_runtime_id.clone(),
task_id: auth.run_task_id().to_string(),
}
}
#[derive(Clone, Debug)]
struct AgentIdentityAuthProvider {
auth: codex_login::auth::AgentIdentityAuth,
auth: AgentIdentityAuth,
}
impl AuthProvider for AgentIdentityAuthProvider {
@@ -95,6 +155,74 @@ pub(crate) fn resolve_provider_auth(
})
}
pub(crate) async fn resolve_provider_auth_for_scope(
auth_manager: Option<Arc<AuthManager>>,
auth: Option<&CodexAuth>,
provider: &ModelProviderInfo,
scope: ProviderAuthScope,
) -> codex_protocol::error::Result<ResolvedProviderAuth> {
let ProviderAuthScope {
agent_identity_policy,
session_source,
agent_identity_session_fallback,
} = scope;
if let Some(CodexAuth::AgentIdentity(agent_identity_auth)) = auth {
return Ok(ResolvedProviderAuth::for_agent_identity(
agent_identity_auth.clone(),
));
}
if !should_bootstrap_chatgpt_agent_identity(agent_identity_policy, auth)
|| agent_identity_session_fallback.is_engaged()
{
return resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new);
}
let Some(auth_manager) = auth_manager else {
return resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new);
};
match auth_manager
.agent_identity_auth(agent_identity_policy, session_source)
.await
{
Ok(Some(agent_identity_auth)) => Ok(ResolvedProviderAuth::for_agent_identity(
agent_identity_auth,
)),
Ok(None) => resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new),
Err(err) => {
if let Some(AgentIdentityAuthError::BootstrapUnavailable {
operation,
attempts,
message,
}) = err
.get_ref()
.and_then(|source| source.downcast_ref::<AgentIdentityAuthError>())
{
let newly_engaged = agent_identity_session_fallback.engage();
tracing::warn!(
operation,
attempts = *attempts,
error = %message,
newly_engaged,
"agent identity bootstrap unavailable; using ChatGPT bearer auth for this session"
);
resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new)
} else {
Err(err.into())
}
}
}
}
fn should_bootstrap_chatgpt_agent_identity(
agent_identity_policy: AgentIdentityAuthPolicy,
auth: Option<&CodexAuth>,
) -> bool {
agent_identity_policy == AgentIdentityAuthPolicy::ChatGptAuth
&& matches!(auth, Some(CodexAuth::Chatgpt(_)))
}
fn bearer_auth_for_provider(
provider: &ModelProviderInfo,
) -> codex_protocol::error::Result<Option<BearerAuthProvider>> {
@@ -129,13 +257,128 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider {
#[cfg(test)]
mod tests {
use codex_agent_identity::generate_agent_key_material;
use codex_login::AuthCredentialsStoreMode;
use codex_login::AuthKeyringBackendKind;
use codex_login::auth::AgentIdentityAuthRecord;
use codex_login::auth::BedrockApiKeyAuth;
use codex_model_provider_info::WireApi;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_protocol::account::PlanType;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::*;
static NEXT_CODEX_HOME_ID: AtomicUsize = AtomicUsize::new(0);
const TEST_CHATGPT_ID_TOKEN: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7ImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzNDUiLCJ1c2VyX2lkIjoidXNlci0xMjM0NSIsImNoYXRncHRfcGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjb3VudC0xMjMifX0.c2ln";
async fn agent_identity_auth(chatgpt_account_is_fedramp: bool) -> AgentIdentityAuth {
let key_material = generate_agent_key_material().expect("generate key material");
AgentIdentityAuth::from_record(
AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-1".to_string(),
agent_private_key: key_material.private_key_pkcs8_base64,
account_id: "account-1".to_string(),
chatgpt_user_id: "user-1".to_string(),
email: Some("agent@example.com".to_string()),
plan_type: PlanType::Plus,
chatgpt_account_is_fedramp,
task_id: Some("task-run-1".to_string()),
},
"https://auth.openai.com/api/accounts",
/*auth_route_config*/ None,
)
.await
.expect("agent identity auth record should include task id")
}
fn provider_auth_scope(
policy: AgentIdentityAuthPolicy,
fallback: AgentIdentitySessionFallback,
) -> ProviderAuthScope {
ProviderAuthScope {
agent_identity_policy: policy,
session_source: SessionSource::Cli,
agent_identity_session_fallback: fallback,
}
}
fn test_codex_home() -> PathBuf {
let id = NEXT_CODEX_HOME_ID.fetch_add(1, Ordering::Relaxed);
let path = std::env::temp_dir().join(format!(
"codex-model-provider-agent-identity-{pid}-{id}",
pid = std::process::id()
));
let _ = std::fs::remove_dir_all(&path);
std::fs::create_dir_all(&path).expect("create temp codex home");
path
}
fn write_chatgpt_auth_json(codex_home: &Path) {
let auth_json = json!({
"tokens": {
"id_token": TEST_CHATGPT_ID_TOKEN,
"access_token": "test-access-token",
"refresh_token": "test-refresh-token",
"account_id": "account-123"
},
"last_refresh": "2099-01-01T00:00:00Z"
});
std::fs::write(
codex_home.join("auth.json"),
serde_json::to_string_pretty(&auth_json).expect("serialize auth.json"),
)
.expect("write auth.json");
}
async fn chatgpt_auth_manager(
agent_identity_authapi_base_url: String,
) -> (PathBuf, Arc<AuthManager>, CodexAuth) {
let codex_home = test_codex_home();
write_chatgpt_auth_json(&codex_home);
let auth_manager = AuthManager::shared(
codex_home.clone(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
/*auth_route_config*/ None,
)
.await;
let auth = auth_manager.auth().await.expect("auth should load");
let auth_manager = AuthManager::from_auth_for_testing_with_agent_identity_authapi_base_url(
auth.clone(),
agent_identity_authapi_base_url,
);
(codex_home, auth_manager, auth)
}
async fn mount_transient_agent_registration(
server: &MockServer,
status: u16,
registration_count: Arc<AtomicUsize>,
) {
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.respond_with(move |_request: &wiremock::Request| {
registration_count.fetch_add(1, Ordering::SeqCst);
ResponseTemplate::new(status)
})
.mount(server)
.await;
}
#[test]
fn unauthenticated_auth_provider_adds_no_headers() {
let provider =
@@ -161,4 +404,180 @@ mod tests {
Ok(_) => panic!("Bedrock API key auth should be rejected"),
}
}
#[tokio::test]
async fn first_party_run_scope_uses_agent_assertion_and_exposes_telemetry() {
let auth = CodexAuth::AgentIdentity(
agent_identity_auth(/*chatgpt_account_is_fedramp*/ false).await,
);
let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
let auth = resolve_provider_auth_for_scope(
/*auth_manager*/ None,
Some(&auth),
&provider,
provider_auth_scope(
AgentIdentityAuthPolicy::JwtOnly,
AgentIdentitySessionFallback::default(),
),
)
.await
.expect("auth should resolve");
assert_eq!(
auth.agent_identity_telemetry,
Some(AgentIdentityTelemetry {
agent_id: "agent-runtime-1".to_string(),
task_id: "task-run-1".to_string(),
})
);
let headers = auth.auth.to_auth_headers();
assert!(
headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.starts_with("AgentAssertion "))
);
}
#[tokio::test]
async fn agent_identity_auth_provider_preserves_account_routing_headers() {
let auth = agent_identity_auth(/*chatgpt_account_is_fedramp*/ true).await;
let provider = auth_provider_from_auth(&CodexAuth::AgentIdentity(auth));
let headers = provider.to_auth_headers();
assert!(
headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.starts_with("AgentAssertion "))
);
assert_eq!(
headers
.get("ChatGPT-Account-ID")
.and_then(|value| value.to_str().ok()),
Some("account-1")
);
assert_eq!(
headers
.get("X-OpenAI-Fedramp")
.and_then(|value| value.to_str().ok()),
Some("true")
);
}
#[tokio::test]
async fn chatgpt_bootstrap_unavailable_uses_session_bearer_fallback() {
let server = MockServer::start().await;
let registration_count = Arc::new(AtomicUsize::new(0));
mount_transient_agent_registration(
&server,
/*status*/ 503,
Arc::clone(&registration_count),
)
.await;
let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await;
let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
let fallback = AgentIdentitySessionFallback::default();
let provider_auth = resolve_provider_auth_for_scope(
Some(auth_manager),
Some(&auth),
&provider,
provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback.clone()),
)
.await
.expect("fallback should resolve bearer auth");
let headers = provider_auth.auth.to_auth_headers();
assert_eq!(
headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok()),
Some("Bearer test-access-token")
);
assert_eq!(
headers
.get("ChatGPT-Account-ID")
.and_then(|value| value.to_str().ok()),
Some("account-123")
);
assert!(fallback.is_engaged());
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn chatgpt_session_fallback_skips_later_agent_identity_bootstrap() {
let server = MockServer::start().await;
let registration_count = Arc::new(AtomicUsize::new(0));
mount_transient_agent_registration(
&server,
/*status*/ 503,
Arc::clone(&registration_count),
)
.await;
let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await;
let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
let fallback = AgentIdentitySessionFallback::default();
resolve_provider_auth_for_scope(
Some(Arc::clone(&auth_manager)),
Some(&auth),
&provider,
provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback.clone()),
)
.await
.expect("first fallback should resolve bearer auth");
resolve_provider_auth_for_scope(
Some(auth_manager),
Some(&auth),
&provider,
provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback),
)
.await
.expect("second fallback should resolve bearer auth");
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn chatgpt_sessions_share_bootstrap_failure_cooldown() {
let server = MockServer::start().await;
let registration_count = Arc::new(AtomicUsize::new(0));
mount_transient_agent_registration(
&server,
/*status*/ 503,
Arc::clone(&registration_count),
)
.await;
let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await;
let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
let first_fallback = AgentIdentitySessionFallback::default();
let second_fallback = AgentIdentitySessionFallback::default();
resolve_provider_auth_for_scope(
Some(Arc::clone(&auth_manager)),
Some(&auth),
&provider,
provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, first_fallback.clone()),
)
.await
.expect("first session fallback should resolve bearer auth");
resolve_provider_auth_for_scope(
Some(auth_manager),
Some(&auth),
&provider,
provider_auth_scope(
AgentIdentityAuthPolicy::ChatGptAuth,
second_fallback.clone(),
),
)
.await
.expect("second session fallback should resolve bearer auth");
assert!(first_fallback.is_engaged());
assert!(second_fallback.is_engaged());
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
}
}
codex-rs/model-provider/src/lib.rs
+3
View File
@@ -4,6 +4,9 @@ mod bearer_auth_provider;
mod models_endpoint;
mod provider;
pub use auth::AgentIdentitySessionFallback;
pub use auth::ProviderAuthScope;
pub use auth::ResolvedProviderAuth;
pub use auth::auth_provider_from_auth;
pub use auth::unauthenticated_auth_provider;
pub use bearer_auth_provider::BearerAuthProvider;
codex-rs/model-provider/src/models_endpoint.rs
+13
View File
@@ -1,6 +1,7 @@
use std::sync::Arc;
use std::time::Duration;
use codex_api::AgentIdentityTelemetry;
use codex_api::ModelsClient;
use codex_api::RequestTelemetry;
use codex_api::ReqwestTransport;
@@ -26,6 +27,7 @@ use codex_response_debug_context::telemetry_transport_error_message;
use http::HeaderMap;
use tokio::time::timeout;
use crate::auth::agent_identity_telemetry;
use crate::auth::resolve_provider_auth;
const MODELS_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
@@ -75,10 +77,16 @@ impl OpenAiModelsEndpoint {
let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?;
let transport = ReqwestTransport::new(build_reqwest_client());
let auth_telemetry = auth_header_telemetry(api_auth.as_ref());
let agent_identity_telemetry = if let Some(CodexAuth::AgentIdentity(auth)) = auth.as_ref() {
Some(agent_identity_telemetry(auth))
} else {
None
};
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()),
auth_header_attached: auth_telemetry.attached,
auth_header_name: auth_telemetry.name,
agent_identity_telemetry,
auth_env: self.auth_env(),
});
let client = ModelsClient::new(transport, api_provider, api_auth)
@@ -124,6 +132,7 @@ struct ModelsRequestTelemetry {
auth_mode: Option<String>,
auth_header_attached: bool,
auth_header_name: Option<&'static str>,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
auth_env: AuthEnvTelemetry,
}
@@ -164,6 +173,8 @@ impl RequestTelemetry for ModelsRequestTelemetry {
auth.error = response_debug.auth_error.as_deref(),
auth.error_code = response_debug.auth_error_code.as_deref(),
auth.mode = self.auth_mode.as_deref(),
auth.agent_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.agent_id.as_str()),
auth.task_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.task_id.as_str()),
);
tracing::event!(
target: "codex_otel.trace_safe",
@@ -188,6 +199,8 @@ impl RequestTelemetry for ModelsRequestTelemetry {
auth.error = response_debug.auth_error.as_deref(),
auth.error_code = response_debug.auth_error_code.as_deref(),
auth.mode = self.auth_mode.as_deref(),
auth.agent_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.agent_id.as_str()),
auth.task_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.task_id.as_str()),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
codex-rs/model-provider/src/provider.rs
+49
View File
@@ -18,8 +18,11 @@ use codex_protocol::error::CodexErr;
use codex_protocol::openai_models::ModelsResponse;
use crate::amazon_bedrock::AmazonBedrockModelProvider;
use crate::auth::ProviderAuthScope;
use crate::auth::ResolvedProviderAuth;
use crate::auth::auth_manager_for_provider;
use crate::auth::resolve_provider_auth;
use crate::auth::resolve_provider_auth_for_scope;
use crate::models_endpoint::OpenAiModelsEndpoint;
/// Optional provider-backed features that Codex may expose at runtime.
@@ -175,6 +178,21 @@ pub trait ModelProvider: fmt::Debug + Send + Sync {
})
}
/// Returns request credentials, optionally scoped to a Codex session task.
fn api_auth_for_scope(
&self,
scope: ProviderAuthScope,
) -> ModelProviderFuture<'_, codex_protocol::error::Result<ResolvedProviderAuth>> {
Box::pin(async move {
if !provider_uses_first_party_auth_path(self.info()) {
return self.api_auth().await.map(ResolvedProviderAuth::new);
}
let auth = self.auth().await;
resolve_provider_auth_for_scope(self.auth_manager(), auth.as_ref(), self.info(), scope)
.await
})
}
/// Creates the model manager implementation appropriate for this provider.
fn models_manager(
&self,
@@ -188,6 +206,14 @@ pub type ModelProviderFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a
/// Shared runtime model provider handle.
pub type SharedModelProvider = Arc<dyn ModelProvider>;
fn provider_uses_first_party_auth_path(provider: &ModelProviderInfo) -> bool {
provider.requires_openai_auth
&& provider.env_key.is_none()
&& provider.experimental_bearer_token.is_none()
&& provider.auth.is_none()
&& provider.aws.is_none()
}
/// Creates the default runtime model provider for configured provider metadata.
pub fn create_model_provider(
provider_info: ModelProviderInfo,
@@ -310,14 +336,17 @@ impl ModelProvider for ConfiguredModelProvider {
mod tests {
use std::num::NonZeroU64;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_login::auth::BedrockApiKeyAuth;
use codex_model_provider_info::ModelProviderAwsAuthInfo;
use codex_model_provider_info::WireApi;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_models_manager::manager::RefreshStrategy;
use codex_protocol::account::PlanType;
use codex_protocol::config_types::ModelProviderAuthInfo;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::protocol::SessionSource;
use pretty_assertions::assert_eq;
use serde_json::json;
use wiremock::Mock;
@@ -328,6 +357,7 @@ mod tests {
use wiremock::matchers::path;
use super::*;
use crate::auth::AgentIdentitySessionFallback;
fn provider_info_with_command_auth() -> ModelProviderInfo {
ModelProviderInfo {
@@ -406,6 +436,25 @@ mod tests {
})
}
#[tokio::test]
async fn scoped_auth_ignores_scope_for_non_openai_provider() {
let provider = create_model_provider(
create_oss_provider_with_base_url("http://localhost:11434/v1", WireApi::Responses),
/*auth_manager*/ None,
);
let auth = provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly,
session_source: SessionSource::Cli,
agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
})
.await
.expect("auth should resolve");
assert!(auth.auth.to_auth_headers().is_empty());
}
#[test]
fn configured_provider_uses_default_capabilities() {
let provider = create_model_provider(
codex-rs/otel/src/events/session_telemetry.rs
+11
View File
@@ -32,6 +32,7 @@ use crate::metrics::runtime_metrics::RuntimeMetricsSummary;
use crate::metrics::timer::Timer;
use crate::provider::OtelProvider;
use crate::sanitize_metric_tag_value;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_protocol::ThreadId;
@@ -516,6 +517,7 @@ impl SessionTelemetry {
/*cf_ray*/ None,
/*auth_error*/ None,
/*auth_error_code*/ None,
/*agent_identity_telemetry*/ None,
);
response
@@ -538,6 +540,7 @@ impl SessionTelemetry {
cf_ray: Option<&str>,
auth_error: Option<&str>,
auth_error_code: Option<&str>,
agent_identity_telemetry: Option<&AgentIdentityTelemetry>,
) {
let success = status.is_some_and(|code| (200..=299).contains(&code)) && error.is_none();
let success_str = if success { "true" } else { "false" };
@@ -578,6 +581,8 @@ impl SessionTelemetry {
auth.cf_ray = cf_ray,
auth.error = auth_error,
auth.error_code = auth_error_code,
auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()),
auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()),
},
log: {},
trace: {},
@@ -601,6 +606,7 @@ impl SessionTelemetry {
cf_ray: Option<&str>,
auth_error: Option<&str>,
auth_error_code: Option<&str>,
agent_identity_telemetry: Option<&AgentIdentityTelemetry>,
) {
let success = error.is_none()
&& status
@@ -632,6 +638,8 @@ impl SessionTelemetry {
auth.cf_ray = cf_ray,
auth.error = auth_error,
auth.error_code = auth_error_code,
auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()),
auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()),
},
log: {},
trace: {},
@@ -643,6 +651,7 @@ impl SessionTelemetry {
duration: Duration,
error: Option<&str>,
connection_reused: bool,
agent_identity_telemetry: Option<&AgentIdentityTelemetry>,
) {
let success_str = if error.is_none() { "true" } else { "false" };
self.counter(
@@ -669,6 +678,8 @@ impl SessionTelemetry {
auth.env_provider_key_present = self.metadata.auth_env.provider_env_key_present,
auth.env_refresh_token_url_override_present = self.metadata.auth_env.refresh_token_url_override_present,
auth.connection_reused = connection_reused,
auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()),
auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()),
},
log: {},
trace: {},
codex-rs/otel/tests/suite/otel_export_routing_policy.rs
+64
View File
@@ -1,3 +1,4 @@
use codex_api::AgentIdentityTelemetry;
use codex_otel::AuthEnvTelemetryMetadata;
use codex_otel::OtelProvider;
use codex_otel::SessionTelemetry;
@@ -511,6 +512,10 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() {
SandboxPolicy::DangerFullAccess,
Vec::new(),
);
let agent_identity_telemetry = AgentIdentityTelemetry {
agent_id: "agent-runtime-otel".to_string(),
task_id: "task-run-otel".to_string(),
};
manager.record_api_request(
/*attempt*/ 1,
Some(401),
@@ -526,6 +531,7 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() {
Some("ray-401"),
Some("missing_authorization_header"),
Some("token_expired"),
Some(&agent_identity_telemetry),
);
});
@@ -599,6 +605,14 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() {
.map(String::as_str),
Some("true")
);
assert_eq!(
request_log_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-otel")
);
assert_eq!(
request_log_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-otel")
);
let spans = span_exporter.get_finished_spans().expect("span export");
let conversation_trace_event =
@@ -641,6 +655,14 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() {
.map(String::as_str),
Some("true")
);
assert_eq!(
request_trace_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-otel")
);
assert_eq!(
request_trace_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-otel")
);
}
#[test]
@@ -685,6 +707,10 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() {
.with_auth_env(auth_env_metadata());
let root_span = tracing::info_span!("root");
let _root_guard = root_span.enter();
let agent_identity_telemetry = AgentIdentityTelemetry {
agent_id: "agent-runtime-ws".to_string(),
task_id: "task-run-ws".to_string(),
};
manager.record_websocket_connect(
std::time::Duration::from_millis(17),
Some(401),
@@ -700,6 +726,7 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() {
Some("ray-ws-401"),
Some("missing_authorization_header"),
Some("token_expired"),
Some(&agent_identity_telemetry),
);
});
@@ -741,6 +768,14 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() {
.map(String::as_str),
Some("configured")
);
assert_eq!(
connect_log_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-ws")
);
assert_eq!(
connect_log_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-ws")
);
let spans = span_exporter.get_finished_spans().expect("span export");
let connect_trace_event =
@@ -758,6 +793,14 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() {
.map(String::as_str),
Some("true")
);
assert_eq!(
connect_trace_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-ws")
);
assert_eq!(
connect_trace_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-ws")
);
}
#[test]
@@ -802,10 +845,15 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability()
.with_auth_env(auth_env_metadata());
let root_span = tracing::info_span!("root");
let _root_guard = root_span.enter();
let agent_identity_telemetry = AgentIdentityTelemetry {
agent_id: "agent-runtime-ws-request".to_string(),
task_id: "task-run-ws-request".to_string(),
};
manager.record_websocket_request(
std::time::Duration::from_millis(23),
Some("stream error"),
/*connection_reused*/ true,
Some(&agent_identity_telemetry),
);
});
@@ -831,6 +879,14 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability()
.map(String::as_str),
Some("true")
);
assert_eq!(
request_log_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-ws-request")
);
assert_eq!(
request_log_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-ws-request")
);
let spans = span_exporter.get_finished_spans().expect("span export");
let request_trace_event =
@@ -848,4 +904,12 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability()
.map(String::as_str),
Some("true")
);
assert_eq!(
request_trace_attrs.get("auth.agent_id").map(String::as_str),
Some("agent-runtime-ws-request")
);
assert_eq!(
request_trace_attrs.get("auth.task_id").map(String::as_str),
Some("task-run-ws-request")
);
}
codex-rs/otel/tests/suite/runtime_summary.rs
+2
View File
@@ -61,11 +61,13 @@ fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<(
/*cf_ray*/ None,
/*auth_error*/ None,
/*auth_error_code*/ None,
/*agent_identity_telemetry*/ None,
);
manager.record_websocket_request(
Duration::from_millis(400),
/*error*/ None,
/*connection_reused*/ false,
/*agent_identity_telemetry*/ None,
);
let sse_response: std::result::Result<
Option<std::result::Result<StreamEvent, eventsource_stream::EventStreamError<&str>>>,