[codex] Use background task auth for additional backend calls (#18260)

## Summary Splits the larger PR4.1 background task auth rollout by moving additional backend/control-plane call sites into this downstream PR. This PR keeps callers on the same design as PR4.1: most code asks `AuthManager` for the default ChatGPT backend authorization header, and `AuthManager` decides bearer vs background AgentAssertion internally. Task-pinned inference auth remains separate because it needs the thread's registered task id. ## Stack - PR1: https://github.com/openai/codex/pull/17385 - add `features.use_agent_identity` - PR2: https://github.com/openai/codex/pull/17386 - register agent identities when enabled - PR3: https://github.com/openai/codex/pull/17387 - register agent tasks when enabled - PR3.1: https://github.com/openai/codex/pull/17978 - persist and prewarm registered tasks per thread - PR4: https://github.com/openai/codex/pull/17980 - use task-scoped `AgentAssertion` for downstream calls - PR4.1: https://github.com/openai/codex/pull/18094 - introduce AuthManager-owned background/control-plane `AgentAssertion` auth - PR4.2: this PR - use background task auth for additional backend/control-plane calls ## What Changed - pass full authorization header values through backend-client and cloud-tasks-client call paths where needed - move ChatGPT client, cloud requirements, cloud tasks, thread-manager, and models-manager background auth usage into this downstream slice - make app-server remote control enrollment/websocket auth ask `AuthManager` for the local backend authorization header instead of threading a background auth mode through transport options - keep the same feature-gated bearer fallback behavior from PR4.1 ## Validation - `just fmt` - `cargo check -p codex-core -p codex-login -p codex-analytics -p codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p codex-models-manager -p codex-chatgpt -p codex-model-provider -p codex-mcp -p codex-core-skills` - `cargo test -p codex-login agent_identity` - `cargo test -p codex-model-provider bearer_auth_provider` - `cargo test -p codex-core agent_assertion` - `cargo test -p codex-app-server remote_control` - `cargo test -p codex-cloud-requirements fetch_cloud_requirements` - `cargo test -p codex-models-manager manager::tests` - `cargo test -p codex-chatgpt` - `cargo test -p codex-cloud-tasks` - `just fix -p codex-core -p codex-login -p codex-analytics -p codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p codex-models-manager -p codex-chatgpt -p codex-model-provider -p codex-mcp -p codex-core-skills` - `just fix -p codex-app-server` - `git diff --check`
2026-07-01 00:31:56 +08:00 · 2026-04-20 07:24:29 -07:00
parent fa0e2ba87c
commit 19e2f21827
14 changed files with 364 additions and 98 deletions
@@ -2013,12 +2013,28 @@ impl CodexMessageProcessor {
            });
        }

-        let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth)
+        let authorization_header_value = self
+            .auth_manager
+            .chatgpt_authorization_header_for_auth(&auth)
+            .await;
+        let mut client = BackendClient::new(self.config.chatgpt_base_url.clone())
+            .map(|client| {
+                client.with_user_agent(codex_login::default_client::get_codex_user_agent())
+            })
            .map_err(|err| JSONRPCErrorError {
                code: INTERNAL_ERROR_CODE,
                message: format!("failed to construct backend client: {err}"),
                data: None,
            })?;
+        if let Some(authorization_header_value) = authorization_header_value {
+            client = client.with_authorization_header_value(authorization_header_value);
+        }
+        if let Some(account_id) = auth.get_account_id() {
+            client = client.with_chatgpt_account_id(account_id);
+        }
+        if auth.is_fedramp_account() {
+            client = client.with_fedramp_routing_header();
+        }

        let snapshots = client
            .get_rate_limits_many()
@@ -27,10 +27,11 @@ use crate::outgoing_message::QueuedOutgoingMessage;
 use crate::transport::CHANNEL_CAPACITY;
 use crate::transport::ConnectionState;
 use crate::transport::OutboundConnectionState;
+use crate::transport::RemoteControlStartOptions;
 use crate::transport::TransportEvent;
 use crate::transport::auth::policy_from_settings;
 use crate::transport::route_outgoing_envelope;
-use crate::transport::start_remote_control;
+use crate::transport::start_remote_control_with_options;
 use crate::transport::start_stdio_connection;
 use crate::transport::start_websocket_acceptor;
 use codex_analytics::AppServerRpcTransport;
@@ -575,16 +576,17 @@ pub async fn run_main_with_transport(
        ));
    }

-    let (remote_control_accept_handle, remote_control_handle) = start_remote_control(
-        config.chatgpt_base_url.clone(),
-        state_db.clone(),
-        auth_manager.clone(),
-        transport_event_tx.clone(),
-        transport_shutdown_token.clone(),
-        app_server_client_name_rx,
-        remote_control_enabled,
-    )
-    .await?;
+    let (remote_control_accept_handle, remote_control_handle) =
+        start_remote_control_with_options(RemoteControlStartOptions {
+            remote_control_url: config.chatgpt_base_url.clone(),
+            state_db: state_db.clone(),
+            auth_manager: auth_manager.clone(),
+            transport_event_tx: transport_event_tx.clone(),
+            shutdown_token: transport_shutdown_token.clone(),
+            app_server_client_name_rx,
+            initial_enabled: remote_control_enabled,
+        })
+        .await?;
    transport_accept_handles.push(remote_control_accept_handle);

    let outbound_handle = tokio::spawn(async move {
@@ -34,7 +34,8 @@ mod stdio;
 mod websocket;

 pub(crate) use remote_control::RemoteControlHandle;
-pub(crate) use remote_control::start_remote_control;
+pub(crate) use remote_control::RemoteControlStartOptions;
+pub(crate) use remote_control::start_remote_control_with_options;
 pub(crate) use stdio::start_stdio_connection;
 pub(crate) use websocket::start_websocket_acceptor;

@@ -17,6 +17,7 @@ const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096;
 const REQUEST_ID_HEADER: &str = "x-request-id";
 const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";
 const CF_RAY_HEADER: &str = "cf-ray";
+const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp";
 pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";

 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -29,8 +30,9 @@ pub(super) struct RemoteControlEnrollment {

 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) struct RemoteControlConnectionAuth {
-    pub(super) bearer_token: String,
+    pub(super) authorization_header_value: String,
    pub(super) account_id: String,
+    pub(super) is_fedramp_account: bool,
 }

 pub(super) async fn load_persisted_remote_control_enrollment(
@@ -199,12 +201,15 @@ pub(super) async fn enroll_remote_control_server(
        app_server_version: env!("CARGO_PKG_VERSION"),
    };
    let client = build_reqwest_client();
-    let http_request = client
+    let mut http_request = client
        .post(enroll_url)
        .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
-        .bearer_auth(&auth.bearer_token)
-        .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
-        .json(&request);
+        .header("authorization", &auth.authorization_header_value)
+        .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id);
+    if auth.is_fedramp_account {
+        http_request = http_request.header(REMOTE_CONTROL_FEDRAMP_HEADER, "true");
+    }
+    let http_request = http_request.json(&request);

    let response = http_request.send().await.map_err(|err| {
        io::Error::other(format!(
@@ -445,8 +450,9 @@ mod tests {
        let err = enroll_remote_control_server(
            &remote_control_target,
            &RemoteControlConnectionAuth {
-                bearer_token: "Access Token".to_string(),
+                authorization_header_value: "Bearer Access Token".to_string(),
                account_id: "account_id".to_string(),
+                is_fedramp_account: false,
            },
        )
        .await
@@ -4,6 +4,7 @@ mod protocol;
 mod websocket;

 use crate::transport::remote_control::websocket::RemoteControlWebsocket;
+use crate::transport::remote_control::websocket::RemoteControlWebsocketOptions;

 pub use self::protocol::ClientId;
 use self::protocol::ServerEvent;
@@ -44,6 +45,17 @@ impl RemoteControlHandle {
    }
 }

+pub(crate) struct RemoteControlStartOptions {
+    pub(crate) remote_control_url: String,
+    pub(crate) state_db: Option<Arc<StateRuntime>>,
+    pub(crate) auth_manager: Arc<AuthManager>,
+    pub(crate) transport_event_tx: mpsc::Sender<TransportEvent>,
+    pub(crate) shutdown_token: CancellationToken,
+    pub(crate) app_server_client_name_rx: Option<oneshot::Receiver<String>>,
+    pub(crate) initial_enabled: bool,
+}
+
+#[cfg(test)]
 pub(crate) async fn start_remote_control(
    remote_control_url: String,
    state_db: Option<Arc<StateRuntime>>,
@@ -53,15 +65,38 @@ pub(crate) async fn start_remote_control(
    app_server_client_name_rx: Option<oneshot::Receiver<String>>,
    initial_enabled: bool,
 ) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
+    start_remote_control_with_options(RemoteControlStartOptions {
+        remote_control_url,
+        state_db,
+        auth_manager,
+        transport_event_tx,
+        shutdown_token,
+        app_server_client_name_rx,
+        initial_enabled,
+    })
+    .await
+}
+
+pub(crate) async fn start_remote_control_with_options(
+    options: RemoteControlStartOptions,
+) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
+    let RemoteControlStartOptions {
+        remote_control_url,
+        state_db,
+        auth_manager,
+        transport_event_tx,
+        shutdown_token,
+        app_server_client_name_rx,
+        initial_enabled,
+    } = options;
    let remote_control_target = if initial_enabled {
        Some(normalize_remote_control_url(&remote_control_url)?)
    } else {
        None
    };
-
    let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
    let join_handle = tokio::spawn(async move {
-        RemoteControlWebsocket::new(
+        RemoteControlWebsocket::from_options(RemoteControlWebsocketOptions {
            remote_control_url,
            remote_control_target,
            state_db,
@@ -69,7 +104,7 @@ pub(crate) async fn start_remote_control(
            transport_event_tx,
            shutdown_token,
            enabled_rx,
-        )
+        })
        .run(app_server_client_name_rx)
        .await;
    });
@@ -49,6 +49,7 @@ use tracing::warn;

 pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2";
 pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
+const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp";
 const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor";
 const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration =
    std::time::Duration::from_secs(10);
@@ -128,6 +129,16 @@ pub(crate) struct RemoteControlWebsocket {
    enabled_rx: watch::Receiver<bool>,
 }

+pub(crate) struct RemoteControlWebsocketOptions {
+    pub(crate) remote_control_url: String,
+    pub(crate) remote_control_target: Option<RemoteControlTarget>,
+    pub(crate) state_db: Option<Arc<StateRuntime>>,
+    pub(crate) auth_manager: Arc<AuthManager>,
+    pub(crate) transport_event_tx: mpsc::Sender<TransportEvent>,
+    pub(crate) shutdown_token: CancellationToken,
+    pub(crate) enabled_rx: watch::Receiver<bool>,
+}
+
 enum ConnectOutcome {
    Connected(Box<WebSocketStream<MaybeTlsStream<TcpStream>>>),
    Disabled,
@@ -135,6 +146,7 @@ enum ConnectOutcome {
 }

 impl RemoteControlWebsocket {
+    #[cfg(test)]
    pub(crate) fn new(
        remote_control_url: String,
        remote_control_target: Option<RemoteControlTarget>,
@@ -144,6 +156,27 @@ impl RemoteControlWebsocket {
        shutdown_token: CancellationToken,
        enabled_rx: watch::Receiver<bool>,
    ) -> Self {
+        Self::from_options(RemoteControlWebsocketOptions {
+            remote_control_url,
+            remote_control_target,
+            state_db,
+            auth_manager,
+            transport_event_tx,
+            shutdown_token,
+            enabled_rx,
+        })
+    }
+
+    pub(crate) fn from_options(options: RemoteControlWebsocketOptions) -> Self {
+        let RemoteControlWebsocketOptions {
+            remote_control_url,
+            remote_control_target,
+            state_db,
+            auth_manager,
+            transport_event_tx,
+            shutdown_token,
+            enabled_rx,
+        } = options;
        let shutdown_token = shutdown_token.child_token();
        let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY);
        let client_tracker =
@@ -271,14 +304,16 @@ impl RemoteControlWebsocket {
                    }
                    return ConnectOutcome::Disabled;
                }
-                connect_result = connect_remote_control_websocket(
-                    &remote_control_target,
-                    self.state_db.as_deref(),
-                    &self.auth_manager,
-                    &mut self.auth_recovery,
-                    &mut self.enrollment,
-                    subscribe_cursor.as_deref(),
-                    app_server_client_name,
+                connect_result = connect_remote_control_websocket_with_options(
+                    ConnectRemoteControlWebsocketOptions {
+                        remote_control_target: &remote_control_target,
+                        state_db: self.state_db.as_deref(),
+                        auth_manager: &self.auth_manager,
+                        auth_recovery: &mut self.auth_recovery,
+                        enrollment: &mut self.enrollment,
+                        subscribe_cursor: subscribe_cursor.as_deref(),
+                        app_server_client_name,
+                    },
                ) => connect_result,
            };

@@ -668,12 +703,11 @@ fn build_remote_control_websocket_request(
        "x-codex-protocol-version",
        REMOTE_CONTROL_PROTOCOL_VERSION,
    )?;
-    set_remote_control_header(
-        headers,
-        "authorization",
-        &format!("Bearer {}", auth.bearer_token),
-    )?;
+    set_remote_control_header(headers, "authorization", &auth.authorization_header_value)?;
    set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?;
+    if auth.is_fedramp_account {
+        set_remote_control_header(headers, REMOTE_CONTROL_FEDRAMP_HEADER, "true")?;
+    }
    if let Some(subscribe_cursor) = subscribe_cursor {
        set_remote_control_header(
            headers,
@@ -718,8 +752,19 @@ pub(crate) async fn load_remote_control_auth(
        ));
    }

+    let authorization_header_value = auth_manager
+        .chatgpt_authorization_header_for_auth(&auth)
+        .await
+        .ok_or_else(|| {
+            io::Error::new(
+                ErrorKind::PermissionDenied,
+                "remote control requires ChatGPT authentication",
+            )
+        })?;
+
    Ok(RemoteControlConnectionAuth {
-        bearer_token: auth.get_token().map_err(io::Error::other)?,
+        authorization_header_value,
+        is_fedramp_account: auth.is_fedramp_account(),
        account_id: auth.get_account_id().ok_or_else(|| {
            io::Error::new(
                ErrorKind::WouldBlock,
@@ -729,6 +774,7 @@ pub(crate) async fn load_remote_control_auth(
    })
 }

+#[cfg(test)]
 pub(super) async fn connect_remote_control_websocket(
    remote_control_target: &RemoteControlTarget,
    state_db: Option<&StateRuntime>,
@@ -741,6 +787,44 @@ pub(super) async fn connect_remote_control_websocket(
    WebSocketStream<MaybeTlsStream<TcpStream>>,
    tungstenite::http::Response<()>,
 )> {
+    connect_remote_control_websocket_with_options(ConnectRemoteControlWebsocketOptions {
+        remote_control_target,
+        state_db,
+        auth_manager,
+        auth_recovery,
+        enrollment,
+        subscribe_cursor,
+        app_server_client_name,
+    })
+    .await
+}
+
+struct ConnectRemoteControlWebsocketOptions<'a> {
+    remote_control_target: &'a RemoteControlTarget,
+    state_db: Option<&'a StateRuntime>,
+    auth_manager: &'a Arc<AuthManager>,
+    auth_recovery: &'a mut UnauthorizedRecovery,
+    enrollment: &'a mut Option<RemoteControlEnrollment>,
+    subscribe_cursor: Option<&'a str>,
+    app_server_client_name: Option<&'a str>,
+}
+
+async fn connect_remote_control_websocket_with_options(
+    options: ConnectRemoteControlWebsocketOptions<'_>,
+) -> io::Result<(
+    WebSocketStream<MaybeTlsStream<TcpStream>>,
+    tungstenite::http::Response<()>,
+)> {
+    let ConnectRemoteControlWebsocketOptions {
+        remote_control_target,
+        state_db,
+        auth_manager,
+        auth_recovery,
+        enrollment,
+        subscribe_cursor,
+        app_server_client_name,
+    } = options;
+
    ensure_rustls_crypto_provider();

    let auth = load_remote_control_auth(auth_manager).await?;
@@ -1003,6 +1087,34 @@ mod tests {
        }
    }

+    #[test]
+    fn build_remote_control_websocket_request_includes_fedramp_header() {
+        let request = build_remote_control_websocket_request(
+            "ws://127.0.0.1/backend-api/wham/remote/control/server",
+            &RemoteControlEnrollment {
+                account_id: "account_id".to_string(),
+                environment_id: "env_test".to_string(),
+                server_id: "srv_e_test".to_string(),
+                server_name: "test-server".to_string(),
+            },
+            &RemoteControlConnectionAuth {
+                authorization_header_value: "AgentAssertion assertion".to_string(),
+                account_id: "account_id".to_string(),
+                is_fedramp_account: true,
+            },
+            /*subscribe_cursor*/ None,
+        )
+        .expect("request should build");
+
+        assert_eq!(
+            request
+                .headers()
+                .get(REMOTE_CONTROL_FEDRAMP_HEADER)
+                .and_then(|value| value.to_str().ok()),
+            Some("true")
+        );
+    }
+
    #[tokio::test]
    async fn connect_remote_control_websocket_includes_http_error_details() {
        let listener = TcpListener::bind("127.0.0.1:0")