Use AgentAssertion downstream behind use_agent_identity

[adrian-openai]

Summary

Stack PR 4 of 4 for feature-gated agent identity support.

This PR uses the registered agent task to attach AgentAssertion authorization on downstream backend paths when features.use_agent_identity is enabled.

Stack

• PR1: Add use_agent_identity feature flag #17385 - add features.use_agent_identity
• PR2: Register agent identities behind use_agent_identity #17386 - register agent identities when enabled
• PR3: Register agent tasks behind use_agent_identity #17387 - register agent tasks when enabled
• PR4: Use AgentAssertion downstream behind use_agent_identity #17388 - this PR

Validation

Covered as part of the local stack validation pass:

• just fmt
• cargo test -p codex-core --lib agent_identity
• cargo test -p codex-core --lib agent_assertion
• cargo test -p codex-core --lib websocket_agent_task
• cargo test -p codex-api api_bridge
• cargo build -p codex-cli --bin codex

Notes

The full local app-server E2E path is still being debugged after PR creation. The current branch stack is directionally ready for review while that follow-up continues.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a9459f7442

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Re-register agent task when runtime IDs diverge

Handle runtime mismatches by refreshing the cached task instead of bailing. SessionState reuses persisted agent_task, but identities can be regenerated (e.g., auth/account change or missing secret on restart). This ensure! turns that recoverable state drift into a hard failure, causing every downstream request to fail before sending any API call.

Useful? React with 👍 / 👎.

[adrian-openai]

Moved/fixed in the replacement PR: #17807.

We changed this from a hard ensure! failure into a recoverable stale-task path. Assertion construction now reports AgentTaskRuntimeMismatch when the persisted task points at an old runtime, client setup maps that to CodexErr::AgentTaskStale, and the sampling loop clears that cached task, registers a fresh one once, and retries with a new task-scoped session.

Code pointers in #17807:

• mismatch detection:

  codex/codex-rs/core/src/agent_identity/assertion.rs

  Lines 15 to 52 in a87e178

  	#[derive(Debug, Error)]
  	#[error(
  	"agent task runtime {agent_runtime_id} does not match stored agent identity {stored_agent_runtime_id}"
  	)]
  	pub(crate) struct AgentTaskRuntimeMismatch {
  	pub(crate) agent_runtime_id: String,
  	pub(crate) task_id: String,
  	pub(crate) stored_agent_runtime_id: String,
  	}
  	#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
  	pub(crate) struct AgentAssertionEnvelope {
  	pub(crate) agent_runtime_id: String,
  	pub(crate) task_id: String,
  	pub(crate) timestamp: String,
  	pub(crate) signature: String,
  	}
  	impl AgentIdentityManager {
  	pub(crate) async fn authorization_header_for_task(
  	&self,
  	agent_task: &RegisteredAgentTask,
  	) -> Result<Option<String>> {
  	if !self.feature_enabled {
  	return Ok(None);
  	}
  	let Some(stored_identity) = self.ensure_registered_identity().await? else {
  	return Ok(None);
  	};
  	if stored_identity.agent_runtime_id != agent_task.agent_runtime_id {
  	return Err(AgentTaskRuntimeMismatch {
  	agent_runtime_id: agent_task.agent_runtime_id.clone(),
  	task_id: agent_task.task_id.clone(),
  	stored_agent_runtime_id: stored_identity.agent_runtime_id,
  	}
  	.into());
  	}

• refresh/retry path:

  codex/codex-rs/core/src/codex.rs

  Lines 7072 to 7089 in a87e178

  	Err(CodexErr::AgentTaskStale) => {
  	if stale_agent_task_refreshed {
  	return Err(CodexErr::AgentTaskStale);
  	}
  	stale_agent_task_refreshed = true;
  	let stale_agent_task = client_session.agent_task().cloned();
  	client_session.disable_cached_websocket_session_on_drop();
  	if let Some(stale_agent_task) = stale_agent_task.as_ref() {
  	sess.clear_cached_agent_task(stale_agent_task).await;
  	}
  	match sess.ensure_agent_task_registered().await {
  	Ok(Some(agent_task)) => {
  	*client_session = sess
  	.services
  	.model_client
  	.new_session_with_agent_task(Some(agent_task));
  	retries = 0;
  	continue;

• regression coverage:

  codex/codex-rs/core/src/client_tests.rs

  Lines 397 to 424 in a87e178

  	#[tokio::test]
  	async fn responses_http_reports_stale_agent_task_when_identity_changed() {
  	let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
  	let (_codex_home, client, mut agent_task, _stored_identity) =
  	model_client_with_agent_task(provider).await;
  	agent_task.agent_runtime_id = "agent-stale".to_string();
  	let model_info = test_model_info();
  	let session_telemetry = test_session_telemetry();
  	let mut client_session = client.new_session_with_agent_task(Some(agent_task));
  	let error = match client_session
  	.stream(
  	&test_prompt("hello"),
  	&model_info,
  	&session_telemetry,
  	/*effort*/ None,
  	ReasoningSummary::Auto,
  	/*service_tier*/ None,
  	/*turn_metadata_header*/ None,
  	)
  	.await
  	{
  	Ok(_) => panic!("stale task should be reported before sending a request"),
  	Err(error) => error,
  	};
  	assert!(matches!(error, CodexErr::AgentTaskStale));
  	}

[shijie-oai]

Just to clarify - it is plumbed but not fully at this point?

[adrian-openai]

I'm not sure what happened, but somehow this PR was merged. I'll respond anyway, and we can figure out what's going on, I'll re-open this on another PR, I think.

[adrian-openai]

Ah ok, my codex fucked up, and github closed this automatically. Will reopen here: #17807

[shijie-oai]

No worries - let me know if you want to continue the conversation here or there.

[adrian-openai]

Yep, that read is right: the path was plumbed, but we are intentionally not making API-key auth work with AgentAssertion in this PR.

In the replacement PR (#17807), the current scope is explicit: AgentAssertion is only for ChatGPT-backed Codex sessions, and API-key sessions continue using the API key until the registration service supports API-key identity binding. The downstream request setup still switches to AgentAssertion when a task exists for a ChatGPT-auth session.

Code pointers in #17807:

• ChatGPT-only scope comment:

  codex/codex-rs/core/src/agent_identity.rs

  Lines 454 to 459 in a87e178

  	fn from_auth(auth: &CodexAuth, forced_workspace_id: Option<String>) -> Option<Self> {
  	// AgentAssertion is currently supported only for ChatGPT-backed Codex sessions. API-key
  	// sessions keep using their API key until the registration service supports API-key
  	// identity binding.
  	if !auth.is_chatgpt_auth() {
  	return None;

• downstream AgentAssertion header selection:

  codex/codex-rs/core/src/client.rs

  Lines 751 to 775 in a87e178

  	let api_auth = match (agent_task, self.state.agent_identity_manager.as_ref()) {
  	(Some(agent_task), Some(agent_identity_manager)) => {
  	if let Some(authorization_header_value) = agent_identity_manager
  	.authorization_header_for_task(agent_task)
  	.await
  	.map_err(|err| {
  	if let Some(mismatch) = err.downcast_ref::<AgentTaskRuntimeMismatch>() {
  	debug!(
  	agent_runtime_id = %mismatch.agent_runtime_id,
  	task_id = %mismatch.task_id,
  	stored_agent_runtime_id = %mismatch.stored_agent_runtime_id,
  	"agent task no longer matches stored identity"
  	);
  	return CodexErr::AgentTaskStale;
  	}
  	CodexErr::Stream(
  	format!("failed to build agent assertion authorization: {err}"),
  	None,
  	)
  	})?
  	{
  	debug!(
  	agent_runtime_id = %agent_task.agent_runtime_id,
  	task_id = %agent_task.task_id,
  	"using agent assertion authorization for downstream request"

• API-key task registration remains skipped:

  codex/codex-rs/core/src/agent_identity/task_registration.rs

  Lines 193 to 204 in a87e178

  	#[tokio::test]
  	async fn register_task_skips_for_api_key_auth() {
  	let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key"));
  	let manager = AgentIdentityManager::new_for_tests(
  	auth_manager,
  	/*feature_enabled*/ true,
  	"https://chatgpt.com/backend-api/".to_string(),
  	SessionSource::Cli,
  	);
  	assert_eq!(manager.register_task().await.unwrap(), None);
  	}

[adrian-openai]

Let me know if these followups answer your questions @shijie-oai

[shijie-oai]

let me play it back a little bit - cache_websocket_session_on_drop is a bit that indicates if a websocket session can be cached and reused on drop.

[adrian-openai]

Exactly, and the replacement PR (#17807) makes that bit a little less load-bearing for auth correctness.

The cache now stores the RegisteredAgentTask alongside the WebSocket session. cache_websocket_session_on_drop only controls whether this session should be put back in the cache at all; actual reuse is allowed only when the cached connection's task matches the next session's task. That means same-thread/same-task WebSockets can be reused, while taskless vs task-scoped or different-task auth contexts are dropped instead of mixed.

Code pointers in #17807:

• task-keyed cache/reuse logic:

  codex/codex-rs/core/src/client.rs

  Lines 383 to 459 in a87e178

  	pub(crate) fn new_session_with_agent_task(
  	&self,
  	agent_task: Option<RegisteredAgentTask>,
  	) -> ModelClientSession {
  	// WebSocket auth is bound to the task that opened the connection. Reuse only when the
  	// cached connection was created for the same task, and drop mismatched taskless/task-scoped
  	// sessions rather than mixing auth contexts.
  	let websocket_session = self.take_cached_websocket_session(agent_task.as_ref());
  	ModelClientSession {
  	client: self.clone(),
  	websocket_session,
  	agent_task,
  	cache_websocket_session_on_drop: true,
  	turn_state: Arc::new(OnceLock::new()),
  	}
  	}
  	pub(crate) fn auth_manager(&self) -> Option<Arc<AuthManager>> {
  	self.state.auth_manager.clone()
  	}
  	pub(crate) fn set_window_generation(&self, window_generation: u64) {
  	self.state
  	.window_generation
  	.store(window_generation, Ordering::Relaxed);
  	self.clear_cached_websocket_session();
  	}
  	pub(crate) fn advance_window_generation(&self) {
  	self.state.window_generation.fetch_add(1, Ordering::Relaxed);
  	self.clear_cached_websocket_session();
  	}
  	fn current_window_id(&self) -> String {
  	let conversation_id = self.state.conversation_id;
  	let window_generation = self.state.window_generation.load(Ordering::Relaxed);
  	format!("{conversation_id}:{window_generation}")
  	}
  	fn take_cached_websocket_session(
  	&self,
  	agent_task: Option<&RegisteredAgentTask>,
  	) -> WebsocketSession {
  	let mut cached_websocket_session = self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner);
  	if cached_websocket_session.agent_task.as_ref() == agent_task {
  	return std::mem::take(&mut *cached_websocket_session).websocket_session;
  	}
  	*cached_websocket_session = CachedWebsocketSession::default();
  	WebsocketSession::default()
  	}
  	fn store_cached_websocket_session(
  	&self,
  	agent_task: Option<RegisteredAgentTask>,
  	websocket_session: WebsocketSession,
  	) {
  	*self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession {
  	agent_task,
  	websocket_session,
  	};
  	}
  	fn clear_cached_websocket_session(&self) {
  	*self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession::default();

• cache-on-drop stores the task key:

  codex/codex-rs/core/src/client.rs

  Lines 917 to 923 in a87e178

  	impl Drop for ModelClientSession {
  	fn drop(&mut self) {
  	let websocket_session = std::mem::take(&mut self.websocket_session);
  	if self.cache_websocket_session_on_drop {
  	self.client
  	.store_cached_websocket_session(self.agent_task.clone(), websocket_session);
  	}

• reuse coverage for the same task:

  codex/codex-rs/core/src/client_tests.rs

  Lines 503 to 578 in a87e178

  	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
  	async fn websocket_agent_task_reuses_cached_connection_for_same_task() {
  	core_test_support::skip_if_no_network!();
  	let server = responses::start_websocket_server(vec![vec![
  	vec![
  	responses::ev_response_created("resp-1"),
  	responses::ev_completed("resp-1"),
  	],
  	vec![
  	responses::ev_response_created("resp-2"),
  	responses::ev_completed("resp-2"),
  	],
  	]])
  	.await;
  	let mut provider =
  	create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses);
  	provider.supports_websockets = true;
  	provider.websocket_connect_timeout_ms = Some(5_000);
  	let (_codex_home, client, agent_task, stored_identity) =
  	model_client_with_agent_task(provider).await;
  	let model_info = test_model_info();
  	let session_telemetry = test_session_telemetry();
  	let prompt = test_prompt("hello");
  	{
  	let mut first_session = client.new_session_with_agent_task(Some(agent_task.clone()));
  	let mut stream = first_session
  	.stream(
  	&prompt,
  	&model_info,
  	&session_telemetry,
  	/*effort*/ None,
  	ReasoningSummary::Auto,
  	/*service_tier*/ None,
  	/*turn_metadata_header*/ None,
  	)
  	.await
  	.expect("first agent task stream should succeed");
  	drain_stream_to_completion(&mut stream)
  	.await
  	.expect("first agent task websocket stream should complete");
  	}
  	let mut second_session = client.new_session_with_agent_task(Some(agent_task.clone()));
  	let mut stream = second_session
  	.stream(
  	&prompt,
  	&model_info,
  	&session_telemetry,
  	/*effort*/ None,
  	ReasoningSummary::Auto,
  	/*service_tier*/ None,
  	/*turn_metadata_header*/ None,
  	)
  	.await
  	.expect("second agent task stream should succeed");
  	drain_stream_to_completion(&mut stream)
  	.await
  	.expect("second agent task websocket stream should complete");
  	let handshakes = server.handshakes();
  	assert_eq!(handshakes.len(), 1);
  	let agent_authorization = handshakes[0]
  	.header("authorization")
  	.expect("agent handshake should include authorization");
  	assert_agent_assertion_header(
  	&agent_authorization,
  	&stored_identity,
  	&agent_task.agent_runtime_id,
  	&agent_task.task_id,
  	);
  	assert_eq!(server.single_connection().len(), 2);
  	server.shutdown().await;
  	}

[shijie-oai]

For my personal understanding of this code block - why would an agent task associated with a thread not able to use the prewarmed client session of that thread? Would that be a possible scenario?

[shijie-oai]

Chatted with Codex a bit locally and correct me if I am wrong - in this case an initial prewarmed thread is created without agent_task (none) and if we want to reuse it, it has to have a matching task. The effectively meant that the initial prewarmed thread on thread creation is always discarded and we do not get the benefit of it?

[adrian-openai]

Yep, your follow-up was the key point here: a taskless startup-prewarmed WebSocket is not useful once the first real turn needs task-scoped AgentAssertion auth, because we would have to discard that taskless connection.

In the replacement PR (#17807), we stopped scheduling startup prewarm for ChatGPT/AgentIdentity sessions. Instead, task registration starts once a user turn is accepted and overlaps with the rest of turn prep, then we await it before the model request. That preserves the "do not create a task at startup" behavior while still warming up as late/early as core can observe. We did not hook literal keypress/typing events in this PR, since that would require app/TUI input plumbing rather than the core turn path.

Code pointers in #17807:

• skip taskless startup prewarm for AgentIdentity sessions:

  codex/codex-rs/core/src/session_startup_prewarm.rs

  Lines 159 to 168 in a87e178

  	impl Session {
  	pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>, base_instructions: String) {
  	if self.services.agent_identity_manager.is_enabled()
  	&& self.services.auth_manager.auth_mode() != Some(AuthMode::ApiKey)
  	{
  	info!(
  	"skipping startup websocket prewarm because agent identity requires task-scoped auth"
  	);
  	return;
  	}

• lazy task registration overlapped with turn prep:

  codex/codex-rs/core/src/codex.rs

  Lines 6357 to 6394 in a87e178

  	let agent_task_registration = if sess.services.agent_identity_manager.is_enabled() {
  	let sess = Arc::clone(&sess);
  	Some(tokio::spawn(async move {
  	sess.ensure_agent_task_registered().await
  	}))
  	} else {
  	None
  	};
  	sess.services
  	.analytics_events_client
  	.track_app_mentioned(tracking.clone(), mentioned_app_invocations);
  	for plugin in mentioned_plugin_metadata {
  	sess.services
  	.analytics_events_client
  	.track_plugin_used(tracking.clone(), plugin);
  	}
  	sess.merge_connector_selection(explicitly_enabled_connectors.clone())
  	.await;
  	record_additional_contexts(&sess, &turn_context, additional_contexts).await;
  	if !input.is_empty() {
  	// Track the previous-turn baseline from the regular user-turn path only so
  	// standalone tasks (compact/shell/review/undo) cannot suppress future
  	// model/realtime injections.
  	sess.set_previous_turn_settings(Some(PreviousTurnSettings {
  	model: turn_context.model_info.slug.clone(),
  	realtime_active: Some(turn_context.realtime_active),
  	}))
  	.await;
  	}
  	let agent_task_result = match agent_task_registration {
  	Some(registration) => registration.await.unwrap_or_else(|error| {
  	Err(anyhow::anyhow!(
  	"agent task registration task failed: {error}"
  	))
  	}),
  	None => sess.ensure_agent_task_registered().await,
  	};
  	let agent_task = match agent_task_result {

• task-aware WebSocket reuse after that:

  codex/codex-rs/core/src/client.rs

  Lines 383 to 459 in a87e178

  	pub(crate) fn new_session_with_agent_task(
  	&self,
  	agent_task: Option<RegisteredAgentTask>,
  	) -> ModelClientSession {
  	// WebSocket auth is bound to the task that opened the connection. Reuse only when the
  	// cached connection was created for the same task, and drop mismatched taskless/task-scoped
  	// sessions rather than mixing auth contexts.
  	let websocket_session = self.take_cached_websocket_session(agent_task.as_ref());
  	ModelClientSession {
  	client: self.clone(),
  	websocket_session,
  	agent_task,
  	cache_websocket_session_on_drop: true,
  	turn_state: Arc::new(OnceLock::new()),
  	}
  	}
  	pub(crate) fn auth_manager(&self) -> Option<Arc<AuthManager>> {
  	self.state.auth_manager.clone()
  	}
  	pub(crate) fn set_window_generation(&self, window_generation: u64) {
  	self.state
  	.window_generation
  	.store(window_generation, Ordering::Relaxed);
  	self.clear_cached_websocket_session();
  	}
  	pub(crate) fn advance_window_generation(&self) {
  	self.state.window_generation.fetch_add(1, Ordering::Relaxed);
  	self.clear_cached_websocket_session();
  	}
  	fn current_window_id(&self) -> String {
  	let conversation_id = self.state.conversation_id;
  	let window_generation = self.state.window_generation.load(Ordering::Relaxed);
  	format!("{conversation_id}:{window_generation}")
  	}
  	fn take_cached_websocket_session(
  	&self,
  	agent_task: Option<&RegisteredAgentTask>,
  	) -> WebsocketSession {
  	let mut cached_websocket_session = self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner);
  	if cached_websocket_session.agent_task.as_ref() == agent_task {
  	return std::mem::take(&mut *cached_websocket_session).websocket_session;
  	}
  	*cached_websocket_session = CachedWebsocketSession::default();
  	WebsocketSession::default()
  	}
  	fn store_cached_websocket_session(
  	&self,
  	agent_task: Option<RegisteredAgentTask>,
  	websocket_session: WebsocketSession,
  	) {
  	*self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession {
  	agent_task,
  	websocket_session,
  	};
  	}
  	fn clear_cached_websocket_session(&self) {
  	*self
  	.state
  	.cached_websocket_session
  	.lock()
  	.unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession::default();

[shijie-oai]

Sorry I am re-reading those and I am not 100% sure no longer leveraging prewarming for chatgpt auth like method is ideal. Correct me if I am wrong but in long term agent identity would be applicable to all user instance and if that is the case we take away all prewarming capability essentially and lose the performance benefit (we are still keeping it for the api auth path)