From 7fc8bdbaba349af8192ed71e767ebb0d1069fb79 Mon Sep 17 00:00:00 2001 From: Michael Zeng Date: Thu, 14 May 2026 22:07:50 -0700 Subject: [PATCH] exec-server: use Codex auth for remote registration --- codex-rs/Cargo.lock | 2 + codex-rs/cli/src/main.rs | 47 ++++++++++++- codex-rs/exec-server/Cargo.toml | 2 + codex-rs/exec-server/README.md | 3 +- codex-rs/exec-server/src/lib.rs | 1 - codex-rs/exec-server/src/remote.rs | 103 +++++++++++++--------------- codex-rs/exec-server/tests/relay.rs | 24 ++++++- 7 files changed, 121 insertions(+), 61 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 50de8dc946..bee0a016b9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2747,6 +2747,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", + "codex-api", "codex-app-server-protocol", "codex-client", "codex-file-system", @@ -2758,6 +2759,7 @@ dependencies = [ "codex-utils-rustls-provider", "ctor 0.6.3", "futures", + "http 1.4.0", "pretty_assertions", "prost 0.14.3", "reqwest", diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 87ac5d6676..2e8d194fb6 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1388,7 +1388,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { root_remote_auth_token_env.as_deref(), "exec-server", )?; - run_exec_server_command(cmd, &arg0_paths).await?; + run_exec_server_command(cmd, &arg0_paths, &root_config_overrides, &interactive).await?; } Some(Subcommand::Features(FeaturesCli { sub })) => match sub { FeaturesSubcommand::List => { @@ -1485,6 +1485,8 @@ fn profile_v2_for_subcommand<'a>( async fn run_exec_server_command( cmd: ExecServerCommand, arg0_paths: &Arg0DispatchPaths, + root_config_overrides: &CliConfigOverrides, + interactive: &TuiCli, ) -> anyhow::Result<()> { let codex_self_exe = arg0_paths .codex_self_exe @@ -1498,8 +1500,10 @@ async fn run_exec_server_command( let executor_id = cmd .executor_id .ok_or_else(|| anyhow::anyhow!("--executor-id is required when --remote is set"))?; + let auth_provider = + load_exec_server_remote_auth_provider(root_config_overrides, interactive).await?; let mut remote_config = - codex_exec_server::RemoteExecutorConfig::new(base_url, executor_id)?; + codex_exec_server::RemoteExecutorConfig::new(base_url, executor_id, auth_provider)?; if let Some(name) = cmd.name { remote_config.name = name; } @@ -1515,6 +1519,45 @@ async fn run_exec_server_command( .map_err(anyhow::Error::from_boxed) } +async fn load_exec_server_remote_auth_provider( + root_config_overrides: &CliConfigOverrides, + interactive: &TuiCli, +) -> anyhow::Result { + let cli_kv_overrides = root_config_overrides + .parse_overrides() + .map_err(anyhow::Error::msg)?; + let config = ConfigBuilder::default() + .cli_overrides(cli_kv_overrides) + .harness_overrides(ConfigOverrides { + config_profile: interactive.config_profile.clone(), + ..Default::default() + }) + .build() + .await?; + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + + let auth = match auth_manager.auth().await { + Some(auth) => auth, + None => { + auth_manager.reload().await; + auth_manager.auth().await.ok_or_else(|| { + anyhow::anyhow!( + "remote exec-server registration requires ChatGPT authentication; run `codex login` first" + ) + })? + } + }; + + if !auth.is_chatgpt_auth() { + anyhow::bail!( + "remote exec-server registration requires ChatGPT authentication; API key and Agent Identity auth are not supported" + ); + } + + Ok(codex_model_provider::auth_provider_from_auth(&auth)) +} + async fn enable_feature_in_config(interactive: &TuiCli, feature: &str) -> anyhow::Result<()> { FeatureToggles::validate_feature(feature)?; let codex_home = find_codex_home()?; diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 09a9a71ea0..d842094a16 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -17,6 +17,7 @@ axum = { workspace = true, features = ["http1", "tokio", "ws"] } base64 = { workspace = true } bytes = { workspace = true } codex-app-server-protocol = { workspace = true } +codex-api = { workspace = true } codex-client = { workspace = true } codex-file-system = { workspace = true } codex-protocol = { workspace = true } @@ -51,6 +52,7 @@ uuid = { workspace = true, features = ["v4"] } anyhow = { workspace = true } codex-test-binary-support = { workspace = true } ctor = { workspace = true } +http = { workspace = true } pretty_assertions = { workspace = true } serial_test = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 1eaf6e69eb..614ec4dc87 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -26,7 +26,8 @@ The CLI entrypoint supports: Remote mode registers the local exec-server with the executor registry, then reconnects to the service-provided rendezvous websocket as the executor. -It requires a bearer token in `CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN`. +It uses the standard Codex ChatGPT sign-in state; run `codex login` first when +remote registration needs authentication. Wire framing: diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 872f16ce32..bd556638ea 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -91,7 +91,6 @@ pub use protocol::TerminateResponse; pub use protocol::WriteParams; pub use protocol::WriteResponse; pub use protocol::WriteStatus; -pub use remote::CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR; pub use remote::RemoteExecutorConfig; pub use remote::run_remote_executor; pub use runtime_paths::ExecServerRuntimePaths; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 32c0d5bc8e..2493135c9a 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -1,6 +1,6 @@ -use std::env; use std::time::Duration; +use codex_api::SharedAuthProvider; use reqwest::StatusCode; use serde::Deserialize; use tokio::time::sleep; @@ -14,15 +14,12 @@ use crate::ExecServerRuntimePaths; use crate::relay::run_multiplexed_executor; use crate::server::ConnectionProcessor; -pub const CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR: &str = - "CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN"; - const ERROR_BODY_PREVIEW_BYTES: usize = 4096; #[derive(Clone)] struct ExecutorRegistryClient { base_url: String, - bearer_token: String, + auth_provider: SharedAuthProvider, http: reqwest::Client, } @@ -30,17 +27,17 @@ impl std::fmt::Debug for ExecutorRegistryClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ExecutorRegistryClient") .field("base_url", &self.base_url) - .field("bearer_token", &"") + .field("auth_provider", &"") .finish_non_exhaustive() } } impl ExecutorRegistryClient { - fn new(base_url: String, bearer_token: String) -> Result { + fn new(base_url: String, auth_provider: SharedAuthProvider) -> Result { let base_url = normalize_base_url(base_url)?; Ok(Self { base_url, - bearer_token, + auth_provider, http: reqwest::Client::new(), }) } @@ -55,7 +52,7 @@ impl ExecutorRegistryClient { &self.base_url, &format!("/cloud/executor/{executor_id}/register"), )) - .bearer_auth(&self.bearer_token) + .headers(self.auth_provider.to_auth_headers()) .send() .await?; self.parse_json_response(response).await @@ -89,12 +86,12 @@ struct ExecutorRegistryExecutorRegistrationResponse { } /// Configuration for registering an exec-server for remote use. -#[derive(Clone, Eq, PartialEq)] +#[derive(Clone)] pub struct RemoteExecutorConfig { pub base_url: String, pub executor_id: String, pub name: String, - bearer_token: String, + auth_provider: SharedAuthProvider, } impl std::fmt::Debug for RemoteExecutorConfig { @@ -103,28 +100,23 @@ impl std::fmt::Debug for RemoteExecutorConfig { .field("base_url", &self.base_url) .field("executor_id", &self.executor_id) .field("name", &self.name) - .field("bearer_token", &"") + .field("auth_provider", &"") .finish() } } impl RemoteExecutorConfig { - pub fn new(base_url: String, executor_id: String) -> Result { - Self::with_bearer_token(base_url, executor_id, read_remote_bearer_token_from_env()?) - } - - pub fn with_bearer_token( + pub fn new( base_url: String, executor_id: String, - bearer_token: String, + auth_provider: SharedAuthProvider, ) -> Result { let executor_id = normalize_executor_id(executor_id)?; - let bearer_token = normalize_bearer_token(bearer_token)?; Ok(Self { base_url, executor_id, name: "codex-exec-server".to_string(), - bearer_token, + auth_provider, }) } } @@ -136,7 +128,8 @@ pub async fn run_remote_executor( runtime_paths: ExecServerRuntimePaths, ) -> Result<(), ExecServerError> { ensure_rustls_crypto_provider(); - let client = ExecutorRegistryClient::new(config.base_url.clone(), config.bearer_token.clone())?; + let client = + ExecutorRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; let processor = ConnectionProcessor::new(runtime_paths); let mut backoff = Duration::from_secs(1); @@ -162,32 +155,6 @@ pub async fn run_remote_executor( } } -fn read_remote_bearer_token_from_env() -> Result { - read_remote_bearer_token_from_env_with(|name| env::var(name)) -} - -fn read_remote_bearer_token_from_env_with(get_var: F) -> Result -where - F: FnOnce(&str) -> Result, -{ - let bearer_token = get_var(CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR).map_err(|_| { - ExecServerError::ExecutorRegistryAuth(format!( - "executor registry bearer token environment variable `{CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR}` is not set" - )) - })?; - normalize_bearer_token(bearer_token) -} - -fn normalize_bearer_token(bearer_token: String) -> Result { - let bearer_token = bearer_token.trim().to_string(); - if bearer_token.is_empty() { - return Err(ExecServerError::ExecutorRegistryAuth(format!( - "executor registry bearer token environment variable `{CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR}` is empty" - ))); - } - Ok(bearer_token) -} - fn normalize_executor_id(executor_id: String) -> Result { let executor_id = executor_id.trim().to_string(); if executor_id.is_empty() { @@ -274,6 +241,11 @@ fn preview_error_body(body: &str) -> Option { #[cfg(test)] mod tests { + use std::sync::Arc; + + use codex_api::AuthProvider; + use http::HeaderMap; + use http::HeaderValue; use pretty_assertions::assert_eq; use wiremock::Mock; use wiremock::MockServer; @@ -284,25 +256,46 @@ mod tests { use super::*; + #[derive(Debug)] + struct StaticRegistryAuthProvider; + + impl AuthProvider for StaticRegistryAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + let _ = headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer registry-token"), + ); + let _ = headers.insert( + "ChatGPT-Account-ID", + HeaderValue::from_static("workspace-123"), + ); + } + } + + fn static_registry_auth_provider() -> SharedAuthProvider { + Arc::new(StaticRegistryAuthProvider) + } + #[tokio::test] - async fn register_executor_posts_with_bearer_token_header() { + async fn register_executor_posts_with_auth_provider_headers() { let server = MockServer::start().await; - let config = RemoteExecutorConfig::with_bearer_token( + let config = RemoteExecutorConfig::new( server.uri(), "exec-requested".to_string(), - "registry-token".to_string(), + static_registry_auth_provider(), ) .expect("config"); Mock::given(method("POST")) .and(path("/cloud/executor/exec-requested/register")) .and(header("authorization", "Bearer registry-token")) + .and(header("chatgpt-account-id", "workspace-123")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "executor_id": "exec-1", "url": "wss://rendezvous.test/executor/exec-1?role=executor&sig=abc" }))) .mount(&server) .await; - let client = ExecutorRegistryClient::new(server.uri(), "registry-token".to_string()) + let client = ExecutorRegistryClient::new(server.uri(), static_registry_auth_provider()) .expect("client"); let response = client @@ -320,17 +313,17 @@ mod tests { } #[test] - fn debug_output_redacts_bearer_token() { - let config = RemoteExecutorConfig::with_bearer_token( + fn debug_output_redacts_auth_provider() { + let config = RemoteExecutorConfig::new( "https://registry.example".to_string(), "exec-1".to_string(), - "secret-token".to_string(), + static_registry_auth_provider(), ) .expect("config"); let debug = format!("{config:?}"); assert!(debug.contains("")); - assert!(!debug.contains("secret-token")); + assert!(!debug.contains("workspace-123")); } } diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index d228db9ee5..9f985e6982 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -10,6 +10,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use anyhow::bail; +use codex_api::AuthProvider; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; @@ -22,12 +23,15 @@ use codex_exec_server::InitializeResponse; use codex_exec_server::RemoteExecutorConfig; use futures::SinkExt; use futures::StreamExt; +use http::HeaderMap; +use http::HeaderValue; use pretty_assertions::assert_eq; use prost::Message as ProstMessage; use relay_proto::RelayData; use relay_proto::RelayMessageFrame; use relay_proto::RelayReset; use relay_proto::relay_message_frame; +use std::sync::Arc; use tokio::net::TcpListener; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; @@ -46,6 +50,22 @@ const REGISTRY_TOKEN: &str = "registry-token"; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const TEST_TIMEOUT: Duration = Duration::from_secs(5); +#[derive(Debug)] +struct StaticRegistryAuthProvider; + +impl AuthProvider for StaticRegistryAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + let _ = headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer registry-token"), + ); + } +} + +fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { + Arc::new(StaticRegistryAuthProvider) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -63,10 +83,10 @@ async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Res let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; - let config = RemoteExecutorConfig::with_bearer_token( + let config = RemoteExecutorConfig::new( registry.uri(), EXECUTOR_ID.to_string(), - REGISTRY_TOKEN.to_string(), + static_registry_auth_provider(), )?; let remote_executor = tokio::spawn(codex_exec_server::run_remote_executor( config,