[codex] wire native integrity state surfaces [ci changed_files]

This commit is contained in:
Cooper Gamble
2026-06-03 09:54:34 +00:00
parent b2df930c51
commit 5b6b61613b
13 changed files with 355 additions and 1 deletions
+2
View File
@@ -2529,6 +2529,7 @@ dependencies = [
"codex-app-server-protocol",
"codex-apply-patch",
"codex-async-utils",
"codex-client",
"codex-code-mode",
"codex-config",
"codex-connectors",
@@ -3322,6 +3323,7 @@ dependencies = [
"http 1.4.0",
"pretty_assertions",
"serde_json",
"tempfile",
"tokio",
"tracing",
"wiremock",
+1
View File
@@ -26,6 +26,7 @@ chrono = { workspace = true, features = ["serde"] }
clap = { workspace = true, features = ["derive"] }
codex-analytics = { workspace = true }
codex-api = { workspace = true }
codex-client = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-apply-patch = { workspace = true }
codex-async-utils = { workspace = true }
+15
View File
@@ -307,6 +307,21 @@ impl AgentControl {
{
let client_metadata = match state.get_thread(*parent_thread_id).await {
Ok(parent_thread) => {
if let Some(surface) = parent_thread
.codex
.session
.services
.model_client
.native_integrity_surface()
{
new_thread
.thread
.codex
.session
.services
.model_client
.set_native_integrity_surface(surface);
}
parent_thread
.codex
.session
+46
View File
@@ -10,6 +10,7 @@ use crate::context::ContextualUserFragment;
use crate::context::SubagentNotification;
use crate::init_state_db;
use assert_matches::assert_matches;
use codex_client::NativeIntegritySurface;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_protocol::AgentPath;
@@ -1584,6 +1585,51 @@ async fn spawn_child_completion_notifies_parent_history() {
assert_eq!(wait_for_subagent_notification(&parent_thread).await, true);
}
#[tokio::test]
async fn spawn_thread_subagent_inherits_parent_native_integrity_surface() {
let harness = AgentControlHarness::new().await;
let (parent_thread_id, parent_thread) = harness.start_thread().await;
parent_thread
.set_app_server_client_info(
Some("codex_desktop".to_string()),
/*app_server_client_version*/ None,
/*mcp_elicitations_auto_deny*/ false,
)
.await
.expect("parent client info should update");
let child_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello child"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: Some("explorer".to_string()),
})),
)
.await
.expect("child spawn should succeed");
let child_thread = harness
.manager
.get_thread(child_thread_id)
.await
.expect("child thread should exist");
assert_eq!(
child_thread
.codex
.session
.services
.model_client
.native_integrity_surface(),
Some(NativeIntegritySurface::CodexDesktop)
);
}
#[tokio::test]
async fn multi_agent_v2_completion_ignores_dead_direct_parent() {
let harness = AgentControlHarness::new().await;
+56 -1
View File
@@ -62,6 +62,8 @@ use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_api::response_create_client_metadata;
use codex_app_server_protocol::AuthMode;
use codex_client::NativeIntegrityStateContext;
use codex_client::NativeIntegritySurface;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::RefreshTokenError;
@@ -120,6 +122,7 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_model_provider::with_native_integrity_state;
#[cfg(test)]
use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS;
use codex_model_provider_info::ModelProviderInfo;
@@ -180,6 +183,7 @@ struct ModelClientState {
beta_features_header: Option<String>,
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
native_integrity_state: Option<NativeIntegrityStateContext>,
disable_websockets: AtomicBool,
cached_websocket_session: StdMutex<WebsocketSession>,
}
@@ -329,6 +333,12 @@ impl ModelClient {
beta_features_header: Option<String>,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
) -> Self {
let native_integrity_state = auth_manager.as_ref().map(|auth_manager| {
NativeIntegrityStateContext::new(
auth_manager.codex_home().to_path_buf(),
native_integrity_surface_for_session_source(&session_source),
)
});
let model_provider = create_model_provider(provider_info, auth_manager);
let codex_api_key_env_enabled = model_provider
.auth_manager()
@@ -353,6 +363,7 @@ impl ModelClient {
beta_features_header,
include_attestation,
attestation_provider,
native_integrity_state,
disable_websockets: AtomicBool::new(false),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
}),
@@ -390,6 +401,31 @@ impl ModelClient {
self.state.provider.auth_manager()
}
pub(crate) fn set_app_server_client_name(&self, client_name: Option<&str>) {
let surface = client_name
.map(NativeIntegritySurface::from_app_server_client_name)
.unwrap_or(NativeIntegritySurface::CodexCli);
self.set_native_integrity_surface(surface);
}
pub(crate) fn native_integrity_surface(&self) -> Option<NativeIntegritySurface> {
self.state
.native_integrity_state
.as_ref()
.map(NativeIntegrityStateContext::surface)
}
pub(crate) fn set_native_integrity_surface(&self, surface: NativeIntegritySurface) {
if self
.state
.native_integrity_state
.as_ref()
.is_some_and(|state| state.set_surface(surface))
{
self.store_cached_websocket_session(WebsocketSession::default());
}
}
pub(crate) fn set_window_generation(&self, window_generation: u64) {
self.state
.window_generation
@@ -812,7 +848,11 @@ 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 api_auth = with_native_integrity_state(
self.state.provider.api_auth().await?,
auth.as_ref(),
self.state.native_integrity_state.clone(),
);
Ok(CurrentClientSetup {
auth,
api_provider,
@@ -1736,6 +1776,21 @@ fn subagent_header_value(session_source: &SessionSource) -> Option<String> {
}
}
fn native_integrity_surface_for_session_source(
session_source: &SessionSource,
) -> NativeIntegritySurface {
match session_source {
SessionSource::Cli => NativeIntegritySurface::CodexTui,
SessionSource::Exec => NativeIntegritySurface::CodexExec,
SessionSource::VSCode => NativeIntegritySurface::CodexVscode,
SessionSource::SubAgent(_)
| SessionSource::Internal(_)
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::Unknown => NativeIntegritySurface::CodexCli,
}
}
fn parent_thread_id_header_value(parent_thread_id: Option<ThreadId>) -> Option<String> {
parent_thread_id.map(|parent_thread_id| parent_thread_id.to_string())
}
+11
View File
@@ -110,6 +110,17 @@ pub(crate) async fn run_codex_thread_interactive(
}))
.or_cancel(&cancel_token)
.await??;
if let Some(surface) = parent_session
.services
.model_client
.native_integrity_surface()
{
codex
.session
.services
.model_client
.set_native_integrity_surface(surface);
}
let thread_config = codex.thread_config_snapshot().await;
let client_metadata = parent_session.app_server_client_metadata().await;
emit_subagent_session_started(
+35
View File
@@ -2,6 +2,7 @@ use super::*;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX;
use async_channel::bounded;
use codex_client::NativeIntegritySurface;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::models::NetworkPermissions;
@@ -181,6 +182,40 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() {
assert!(matches!(result, Err(CodexErr::TurnAborted)));
}
#[tokio::test]
async fn run_codex_thread_interactive_inherits_parent_native_integrity_surface() {
let (parent_session, parent_ctx, _rx_events) =
crate::session::tests::make_session_and_context_with_rx().await;
parent_session
.services
.model_client
.set_native_integrity_surface(NativeIntegritySurface::CodexDesktop);
let cancel_token = CancellationToken::new();
let delegated = run_codex_thread_interactive(
parent_ctx.config.as_ref().clone(),
Arc::clone(&parent_session.services.auth_manager),
Arc::clone(&parent_session.services.models_manager),
parent_session,
parent_ctx,
cancel_token.clone(),
SubAgentSource::Review,
/*initial_history*/ None,
)
.await
.expect("delegate should spawn");
assert_eq!(
delegated
.session
.services
.model_client
.native_integrity_surface(),
Some(NativeIntegritySurface::CodexDesktop)
);
cancel_token.cancel();
}
#[tokio::test]
async fn handle_request_permissions_uses_tool_call_id_for_round_trip() {
let (parent_session, parent_ctx, rx_events) =
+4
View File
@@ -781,6 +781,10 @@ impl Codex {
app_server_client_version: Option<String>,
mcp_elicitations_auto_deny: bool,
) -> ConstraintResult<()> {
self.session
.services
.model_client
.set_app_server_client_name(app_server_client_name.as_deref());
self.session
.update_settings(SessionSettingsUpdate {
app_server_client_name,
+4
View File
@@ -1401,6 +1401,10 @@ impl AuthManager {
})
}
pub fn codex_home(&self) -> &Path {
&self.codex_home
}
/// Current cached auth (clone) without attempting a refresh.
pub fn auth_cached(&self) -> Option<CodexAuth> {
self.inner.read().ok().and_then(|c| c.auth.clone())
+1
View File
@@ -32,5 +32,6 @@ tracing = { workspace = true, features = ["log"] }
[dev-dependencies]
pretty_assertions = { workspace = true }
serde_json = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
wiremock = { workspace = true }
+17
View File
@@ -5,6 +5,7 @@ use codex_agent_identity::AgentTaskAuthorizationTarget;
use codex_agent_identity::authorization_header_for_agent_task;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_client::NativeIntegrityStateContext;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_model_provider_info::ModelProviderInfo;
@@ -12,6 +13,7 @@ use http::HeaderMap;
use http::HeaderValue;
use crate::bearer_auth_provider::BearerAuthProvider;
use crate::bearer_auth_provider::NativeIntegrityAuthProvider;
#[derive(Clone, Debug)]
struct AgentIdentityAuthProvider {
@@ -119,6 +121,21 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider {
}
}
/// Decorates first-party ChatGPT auth with native integrity-state transport.
pub fn with_native_integrity_state(
auth_provider: SharedAuthProvider,
auth: Option<&CodexAuth>,
state: Option<NativeIntegrityStateContext>,
) -> SharedAuthProvider {
let Some(state) = state else {
return auth_provider;
};
let Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_)) = auth else {
return auth_provider;
};
Arc::new(NativeIntegrityAuthProvider::new(auth_provider, state))
}
#[cfg(test)]
mod tests {
use codex_model_provider_info::WireApi;
@@ -1,4 +1,10 @@
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_client::INTEGRITY_STATE_HEADER_NAME;
use codex_client::INTEGRITY_STATE_UPDATE_HEADER_NAME;
use codex_client::NativeIntegrityStateContext;
use codex_client::NativeIntegritySurface;
use codex_client::is_allowed_chatgpt_request_url;
use http::HeaderMap;
use http::HeaderValue;
@@ -46,10 +52,91 @@ impl AuthProvider for BearerAuthProvider {
}
}
#[derive(Clone)]
pub(crate) struct NativeIntegrityAuthProvider {
auth: SharedAuthProvider,
state: NativeIntegrityStateContext,
surface: NativeIntegritySurface,
}
impl NativeIntegrityAuthProvider {
pub(crate) fn new(auth: SharedAuthProvider, state: NativeIntegrityStateContext) -> Self {
let surface = state.surface();
Self {
auth,
state,
surface,
}
}
}
impl AuthProvider for NativeIntegrityAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
self.auth.add_auth_headers(headers);
}
fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) {
self.auth.add_auth_headers_for_url(request_url, headers);
if !is_allowed_chatgpt_request_url(request_url) {
return;
}
match self.state.load_for_surface(self.surface) {
Ok(Some(state_file)) => {
if let Ok(header) = HeaderValue::from_str(&state_file.state) {
let _ = headers.insert(INTEGRITY_STATE_HEADER_NAME, header);
}
}
Ok(None) => {}
Err(error) => {
tracing::warn!("failed to load native integrity state: {error}");
}
}
}
fn observe_response_headers(
&self,
request_url: &str,
request_headers: &HeaderMap,
response_headers: &HeaderMap,
) {
self.auth
.observe_response_headers(request_url, request_headers, response_headers);
if !is_allowed_chatgpt_request_url(request_url) {
return;
}
let Some(expected_state) = request_headers
.get(INTEGRITY_STATE_HEADER_NAME)
.and_then(|value| value.to_str().ok())
else {
return;
};
let Some(next_state) = response_headers
.get(INTEGRITY_STATE_UPDATE_HEADER_NAME)
.and_then(|value| value.to_str().ok())
else {
return;
};
if let Err(error) = self.state.compare_and_store_for_surface(
self.surface,
expected_state,
next_state.to_string(),
) {
tracing::warn!("failed to rotate native integrity state: {error}");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use codex_client::NativeIntegrityStateFile;
use codex_client::NativeIntegrityStateStore;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use tempfile::TempDir;
#[test]
fn bearer_auth_provider_reports_when_auth_header_will_attach() {
@@ -107,4 +194,79 @@ mod tests {
Some("true")
);
}
#[test]
fn native_integrity_provider_scopes_and_rotates_surface_state() {
let codex_home = TempDir::new().expect("tempdir");
let context = NativeIntegrityStateContext::new(
codex_home.path().to_path_buf(),
NativeIntegritySurface::CodexDesktop,
);
let store = NativeIntegrityStateStore::new(codex_home.path().to_path_buf());
store
.replace(
NativeIntegritySurface::CodexDesktop,
"ois1.initial.nonce.ciphertext".to_string(),
)
.expect("state should store");
let provider = NativeIntegrityAuthProvider::new(
Arc::new(BearerAuthProvider::new("access-token".to_string())),
context.clone(),
);
context.set_surface(NativeIntegritySurface::CodexCli);
store
.replace(
NativeIntegritySurface::CodexCli,
"ois1.cli.nonce.ciphertext".to_string(),
)
.expect("CLI state should store");
let mut request_headers = HeaderMap::new();
provider.add_auth_headers_for_url(
"https://chatgpt.com/backend-api/codex/responses",
&mut request_headers,
);
assert_eq!(
request_headers
.get(INTEGRITY_STATE_HEADER_NAME)
.and_then(|value| value.to_str().ok()),
Some("ois1.initial.nonce.ciphertext")
);
let mut response_headers = HeaderMap::new();
response_headers.insert(
INTEGRITY_STATE_UPDATE_HEADER_NAME,
HeaderValue::from_static("ois1.rotated.nonce.ciphertext"),
);
provider.observe_response_headers(
"https://chatgpt.com/backend-api/codex/responses",
&request_headers,
&response_headers,
);
assert_eq!(
store
.load(NativeIntegritySurface::CodexDesktop)
.expect("state should load")
.expect("state should exist"),
NativeIntegrityStateFile {
state: "ois1.rotated.nonce.ciphertext".to_string(),
}
);
assert_eq!(
store
.load(NativeIntegritySurface::CodexCli)
.expect("CLI state should load")
.expect("CLI state should exist"),
NativeIntegrityStateFile {
state: "ois1.cli.nonce.ciphertext".to_string(),
}
);
let mut external_headers = HeaderMap::new();
provider.add_auth_headers_for_url(
"https://example.com/backend-api/codex/responses",
&mut external_headers,
);
assert!(!external_headers.contains_key(INTEGRITY_STATE_HEADER_NAME));
}
}
+1
View File
@@ -6,6 +6,7 @@ mod provider;
pub use auth::auth_provider_from_auth;
pub use auth::unauthenticated_auth_provider;
pub use auth::with_native_integrity_state;
pub use bearer_auth_provider::BearerAuthProvider;
pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider;
pub use codex_protocol::account::ProviderAccount;