Add opaque identity and make available to thread through to remote contracts

codex-rs/app-server/src/codex_message_processor.rs

@@ -412,6 +412,7 @@ mod plugin_mcp_oauth;
 mod plugins;
 mod token_usage_replay;
 
+use crate::IdentityKey;
 use crate::filters::compute_source_filters;
 use crate::filters::source_kind_matches;
 use crate::thread_state::ThreadListenerCommand;
@@ -674,6 +675,7 @@ pub(crate) struct CodexMessageProcessorArgs {
     pub(crate) config_manager: ConfigManager,
     pub(crate) feedback: CodexFeedback,
     pub(crate) log_db: Option<LogDbLayer>,
+    pub(crate) identity_key: Option<IdentityKey>,
 }
 
 fn configured_thread_store(config: &Config) -> Arc<dyn ThreadStore> {
@@ -770,6 +772,7 @@ impl CodexMessageProcessor {
             config_manager,
             feedback,
             log_db,
+            identity_key: _identity_key,
         } = args;
         Self {
             auth_manager,

codex-rs/app-server/src/in_process.rs

@@ -51,6 +51,7 @@ use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use crate::config_manager::ConfigManager;
+use crate::IdentityKey;
 use crate::error_code::INTERNAL_ERROR_CODE;
 use crate::error_code::INVALID_REQUEST_ERROR_CODE;
 use crate::error_code::OVERLOADED_ERROR_CODE;
@@ -135,6 +136,8 @@ pub struct InProcessStartArgs {
     pub enable_codex_api_key_env: bool,
     /// Initialize params used for initial handshake.
     pub initialize: InitializeParams,
+    /// Opaque identity key forwarded to remote contract implementations.
+    pub identity_key: Option<IdentityKey>,
     /// Capacity used for all runtime queues (clamped to at least 1).
     pub channel_capacity: usize,
 }
@@ -416,6 +419,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
                 rpc_transport: AppServerRpcTransport::InProcess,
                 remote_control_handle: None,
                 plugin_startup_tasks: crate::PluginStartupTasks::Start,
+                identity_key: args.identity_key,
             }));
             let mut thread_created_rx = processor.thread_created_receiver();
             let session = Arc::new(ConnectionSessionState::new(ConnectionOrigin::InProcess));
@@ -763,6 +767,7 @@ mod tests {
                 },
                 capabilities: None,
             },
+            identity_key: None,
             channel_capacity,
         };
         start(args).await.expect("in-process runtime should start")

codex-rs/app-server/src/lib.rs

@@ -13,6 +13,7 @@ use codex_login::AuthManager;
 use codex_utils_cli::CliConfigOverrides;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::ffi::OsString;
 use std::io::ErrorKind;
 use std::io::Result as IoResult;
 use std::sync::Arc;
@@ -102,6 +103,43 @@ pub use crate::transport::auth::WebsocketAuthCliMode;
 
 const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT";
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct IdentityKey {
+    bytes: Vec<u8>,
+}
+
+impl IdentityKey {
+    pub fn from_bytes(bytes: impl Into<Vec<u8>>) -> Self {
+        Self {
+            bytes: bytes.into(),
+        }
+    }
+
+    pub fn from_os_string(value: OsString) -> Self {
+        Self::from_bytes(os_string_to_bytes(value))
+    }
+
+    pub fn as_bytes(&self) -> &[u8] {
+        &self.bytes
+    }
+
+    pub fn into_bytes(self) -> Vec<u8> {
+        self.bytes
+    }
+}
+
+#[cfg(unix)]
+fn os_string_to_bytes(value: OsString) -> Vec<u8> {
+    use std::os::unix::ffi::OsStrExt;
+
+    value.as_os_str().as_bytes().to_vec()
+}
+
+#[cfg(not(unix))]
+fn os_string_to_bytes(value: OsString) -> Vec<u8> {
+    value.to_string_lossy().into_owned().into_bytes()
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 enum LogFormat {
     Default,
@@ -368,15 +406,17 @@ pub enum PluginStartupTasks {
     Skip,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AppServerRuntimeOptions {
     pub plugin_startup_tasks: PluginStartupTasks,
+    pub identity_key: Option<IdentityKey>,
 }
 
 impl Default for AppServerRuntimeOptions {
     fn default() -> Self {
         Self {
             plugin_startup_tasks: PluginStartupTasks::Start,
+            identity_key: None,
         }
     }
 }
@@ -414,6 +454,10 @@ pub async fn run_main_with_transport_options(
     auth: AppServerWebsocketAuthSettings,
     runtime_options: AppServerRuntimeOptions,
 ) -> IoResult<()> {
+    let AppServerRuntimeOptions {
+        plugin_startup_tasks,
+        identity_key,
+    } = runtime_options;
     let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::from_env(
         ExecServerRuntimePaths::from_optional_paths(
             arg0_paths.codex_self_exe.clone(),
@@ -726,7 +770,8 @@ pub async fn run_main_with_transport_options(
             auth_manager,
             rpc_transport: analytics_rpc_transport(&transport),
             remote_control_handle: Some(remote_control_handle),
-            plugin_startup_tasks: runtime_options.plugin_startup_tasks,
+            plugin_startup_tasks,
+            identity_key,
         }));
         let mut thread_created_rx = processor.thread_created_receiver();
         let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count();
@@ -968,9 +1013,27 @@ fn analytics_rpc_transport(transport: &AppServerTransport) -> AppServerRpcTransp
 
 #[cfg(test)]
 mod tests {
+    use super::IdentityKey;
     use super::LogFormat;
     use pretty_assertions::assert_eq;
 
+    #[test]
+    fn identity_key_preserves_opaque_bytes() {
+        let key = IdentityKey::from_bytes(b"tenant-key-\x00\xff".to_vec());
+        assert_eq!(key.as_bytes(), &b"tenant-key-\x00\xff"[..]);
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn identity_key_preserves_unix_argv_bytes() {
+        use std::ffi::OsString;
+        use std::os::unix::ffi::OsStringExt;
+
+        let key = IdentityKey::from_os_string(OsString::from_vec(b"tenant-key-\xff".to_vec()));
+
+        assert_eq!(key.as_bytes(), &b"tenant-key-\xff"[..]);
+    }
+
     #[test]
     fn log_format_from_env_value_matches_json_values_case_insensitively() {
         assert_eq!(LogFormat::from_env_value(Some("json")), LogFormat::Json);

codex-rs/app-server/src/main.rs

@@ -2,6 +2,7 @@ use clap::Parser;
 use codex_app_server::AppServerRuntimeOptions;
 use codex_app_server::AppServerTransport;
 use codex_app_server::AppServerWebsocketAuthArgs;
+use codex_app_server::IdentityKey;
 use codex_app_server::PluginStartupTasks;
 use codex_app_server::run_main_with_transport_options;
 use codex_arg0::Arg0DispatchPaths;
@@ -9,6 +10,7 @@ use codex_arg0::arg0_dispatch_or_else;
 use codex_config::LoaderOverrides;
 use codex_protocol::protocol::SessionSource;
 use codex_utils_cli::CliConfigOverrides;
+use std::ffi::OsString;
 use std::path::PathBuf;
 
 // Debug-only test hook: lets integration tests point the server at a temporary
@@ -39,6 +41,10 @@ struct AppServerArgs {
     #[command(flatten)]
     auth: AppServerWebsocketAuthArgs,
 
+    /// Opaque identity key forwarded to remote contract implementations.
+    #[arg(long = "identity-key", value_name = "KEY")]
+    identity_key: Option<OsString>,
+
     /// Hidden debug-only test hook used by integration tests that spawn the
     /// production app-server binary.
     #[cfg(debug_assertions)]
@@ -60,6 +66,7 @@ fn main() -> anyhow::Result<()> {
         let session_source = args.session_source;
         let auth = args.auth.try_into_settings()?;
         let mut runtime_options = AppServerRuntimeOptions::default();
+        runtime_options.identity_key = args.identity_key.map(IdentityKey::from_os_string);
         #[cfg(debug_assertions)]
         if args.disable_plugin_startup_tasks_for_tests {
             runtime_options.plugin_startup_tasks = PluginStartupTasks::Skip;

codex-rs/app-server/src/message_processor.rs

@@ -5,6 +5,7 @@ use std::sync::OnceLock;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 
+use crate::IdentityKey;
 use crate::codex_message_processor::CodexMessageProcessor;
 use crate::codex_message_processor::CodexMessageProcessorArgs;
 use crate::config_api::ConfigApi;
@@ -249,6 +250,7 @@ pub(crate) struct MessageProcessorArgs {
     pub(crate) rpc_transport: AppServerRpcTransport,
     pub(crate) remote_control_handle: Option<RemoteControlHandle>,
     pub(crate) plugin_startup_tasks: crate::PluginStartupTasks,
+    pub(crate) identity_key: Option<IdentityKey>,
 }
 
 impl MessageProcessor {
@@ -269,6 +271,7 @@ impl MessageProcessor {
             rpc_transport,
             remote_control_handle,
             plugin_startup_tasks,
+            identity_key,
         } = args;
         auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge {
             outgoing: outgoing.clone(),
@@ -304,6 +307,7 @@ impl MessageProcessor {
             config_manager: config_manager.clone(),
             feedback,
             log_db,
+            identity_key,
         });
         if matches!(plugin_startup_tasks, crate::PluginStartupTasks::Start) {
             // Keep plugin startup warmups aligned at app-server startup.

codex-rs/app-server/src/message_processor/tracing_tests.rs

@@ -289,6 +289,7 @@ fn build_test_processor(
         rpc_transport: AppServerRpcTransport::Stdio,
         remote_control_handle: None,
         plugin_startup_tasks: crate::PluginStartupTasks::Start,
+        identity_key: None,
     }));
     (processor, outgoing_rx)
 }