From babfce7d9ca3cd68da39b37f6edcb0b2ea253aba Mon Sep 17 00:00:00 2001 From: bryanashley Date: Sun, 26 Apr 2026 19:45:22 -0700 Subject: [PATCH] Add opaque identity and make available to thread through to remote contracts --- codex-rs/app-server-client/src/lib.rs | 1 + .../app-server/src/codex_message_processor.rs | 3 + codex-rs/app-server/src/in_process.rs | 5 ++ codex-rs/app-server/src/lib.rs | 67 ++++++++++++++++- codex-rs/app-server/src/main.rs | 7 ++ codex-rs/app-server/src/message_processor.rs | 4 ++ .../src/message_processor/tracing_tests.rs | 1 + codex-rs/cli/src/main.rs | 26 ++++++- codex-rs/thread-store/src/lib.rs | 1 + .../thread-store/src/remote/list_threads.rs | 71 ++++++++++++++++++- codex-rs/thread-store/src/remote/mod.rs | 59 ++++++++++++++- 11 files changed, 237 insertions(+), 8 deletions(-) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index e1614c32db..be0383a96c 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -405,6 +405,7 @@ impl InProcessClientStartArgs { session_source: self.session_source, enable_codex_api_key_env: self.enable_codex_api_key_env, initialize, + identity_key: None, channel_capacity: self.channel_capacity, } } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c86d414359..ca0de74130 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/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, + pub(crate) identity_key: Option, } fn configured_thread_store(config: &Config) -> Arc { @@ -770,6 +772,7 @@ impl CodexMessageProcessor { config_manager, feedback, log_db, + identity_key: _identity_key, } = args; Self { auth_manager, diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index cc4e22e923..9b013d87e5 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/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, /// 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") diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 52b6de0e0c..3fffe398db 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/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, +} + +impl IdentityKey { + pub fn from_bytes(bytes: impl Into>) -> 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 { + self.bytes + } +} + +#[cfg(unix)] +fn os_string_to_bytes(value: OsString) -> Vec { + 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 { + 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, } 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); diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index 1cb4bd9a8e..1e9d7088e8 100644 --- a/codex-rs/app-server/src/main.rs +++ b/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, + /// 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; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 6cdb939364..0014fcf90f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/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, pub(crate) plugin_startup_tasks: crate::PluginStartupTasks, + pub(crate) identity_key: Option, } 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. diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs index 507d7865b9..455302fc21 100644 --- a/codex-rs/app-server/src/message_processor/tracing_tests.rs +++ b/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) } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 9f465521da..6c012e7fe7 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -35,6 +35,7 @@ use codex_tui::UpdateAction; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_cli::CliConfigOverrides; use owo_colors::OwoColorize; +use std::ffi::OsString; use std::io::IsTerminal; use std::path::PathBuf; use supports_color::Stream; @@ -440,6 +441,10 @@ struct AppServerCommand { #[command(flatten)] auth: codex_app_server::AppServerWebsocketAuthArgs, + + /// Opaque identity key forwarded to remote contract implementations. + #[arg(long = "identity-key", value_name = "KEY")] + identity_key: Option, } #[derive(Debug, Parser)] @@ -817,6 +822,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { listen, analytics_default_enabled, auth, + identity_key, } = app_server_cli; reject_remote_mode_for_app_server_subcommand( root_remote.as_deref(), @@ -827,7 +833,13 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { None => { let transport = listen; let auth = auth.try_into_settings()?; - codex_app_server::run_main_with_transport( + let identity_key = + identity_key.map(codex_app_server::IdentityKey::from_os_string); + let runtime_options = codex_app_server::AppServerRuntimeOptions { + identity_key, + ..Default::default() + }; + codex_app_server::run_main_with_transport_options( arg0_paths.clone(), root_config_overrides, codex_config::LoaderOverrides::default(), @@ -835,6 +847,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { transport, codex_protocol::protocol::SessionSource::VSCode, auth, + runtime_options, ) .await?; } @@ -2167,6 +2180,17 @@ mod tests { assert!(app_server.analytics_default_enabled); } + #[test] + fn app_server_identity_key_flag_parses() { + let app_server = app_server_from_args( + ["codex", "app-server", "--identity-key", "tenant-key-123"].as_ref(), + ); + assert_eq!( + app_server.identity_key.as_deref(), + Some(std::ffi::OsStr::new("tenant-key-123")) + ); + } + #[test] fn remote_flag_parses_for_interactive_root() { let cli = MultitoolCli::try_parse_from(["codex", "--remote", "ws://127.0.0.1:4500"]) diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 42b9297bca..ef58a9e628 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -22,6 +22,7 @@ pub use in_memory::InMemoryThreadStoreCalls; pub use live_thread::LiveThread; pub use live_thread::LiveThreadInitGuard; pub use local::LocalThreadStore; +pub use remote::IDENTITY_KEY_HEADER; pub use remote::RemoteThreadStore; pub use store::ThreadStore; pub use types::AppendThreadItemsParams; diff --git a/codex-rs/thread-store/src/remote/list_threads.rs b/codex-rs/thread-store/src/remote/list_threads.rs index 7fb0da6dc3..e5498e4123 100644 --- a/codex-rs/thread-store/src/remote/list_threads.rs +++ b/codex-rs/thread-store/src/remote/list_threads.rs @@ -66,6 +66,8 @@ pub(super) async fn list_threads( #[cfg(test)] mod tests { use std::path::PathBuf; + use std::sync::Arc; + use std::sync::Mutex; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; @@ -79,11 +81,14 @@ mod tests { use super::super::proto::thread_store_server; use super::super::proto::thread_store_server::ThreadStoreServer; use super::*; + use crate::IDENTITY_KEY_HEADER; use crate::ThreadSortKey; use crate::ThreadStore; #[derive(Default)] - struct TestServer; + struct TestServer { + captured_identity_key: Option>>>>, + } #[tonic::async_trait] impl thread_store_server::ThreadStore for TestServer { @@ -91,6 +96,16 @@ mod tests { &self, request: Request, ) -> Result, Status> { + if let Some(captured_identity_key) = &self.captured_identity_key { + let identity_key = request + .metadata() + .get_bin(IDENTITY_KEY_HEADER) + .and_then(|value| value.to_bytes().ok()) + .map(|value| value.to_vec()); + *captured_identity_key + .lock() + .expect("captured identity key mutex poisoned") = identity_key; + } let request = request.into_inner(); assert_eq!(request.page_size, 2); assert_eq!(request.cursor.as_deref(), Some("cursor-1")); @@ -170,7 +185,7 @@ mod tests { let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { Server::builder() - .add_service(ThreadStoreServer::new(TestServer)) + .add_service(ThreadStoreServer::new(TestServer::default())) .serve_with_incoming_shutdown( tokio_stream::wrappers::TcpListenerStream::new(listener), async { @@ -225,6 +240,58 @@ mod tests { server.await.expect("join server").expect("server"); } + #[tokio::test] + async fn list_threads_forwards_identity_key_as_metadata() { + let captured_identity_key = Arc::new(Mutex::new(None)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test server"); + let addr = listener.local_addr().expect("test server addr"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server_identity_key = captured_identity_key.clone(); + let server = tokio::spawn(async move { + Server::builder() + .add_service(ThreadStoreServer::new(TestServer { + captured_identity_key: Some(server_identity_key), + })) + .serve_with_incoming_shutdown( + tokio_stream::wrappers::TcpListenerStream::new(listener), + async { + let _ = shutdown_rx.await; + }, + ) + .await + }); + + let store = RemoteThreadStore::new_with_identity_key( + format!("http://{addr}"), + Some(b"tenant-key-\x00\xff".to_vec()), + ); + store + .list_threads(ListThreadsParams { + page_size: 2, + cursor: Some("cursor-1".to_string()), + sort_key: ThreadSortKey::UpdatedAt, + allowed_sources: vec![SessionSource::Cli], + model_providers: Some(vec!["openai".to_string()]), + archived: true, + search_term: Some("needle".to_string()), + }) + .await + .expect("list threads"); + + assert_eq!( + captured_identity_key + .lock() + .expect("captured identity key mutex poisoned") + .as_deref(), + Some(&b"tenant-key-\x00\xff"[..]) + ); + + let _ = shutdown_tx.send(()); + server.await.expect("join server").expect("server"); + } + #[test] fn stored_thread_proto_roundtrips_through_domain_type() { let thread = proto::StoredThread { diff --git a/codex-rs/thread-store/src/remote/mod.rs b/codex-rs/thread-store/src/remote/mod.rs index a3befac4e5..0d4e83e23a 100644 --- a/codex-rs/thread-store/src/remote/mod.rs +++ b/codex-rs/thread-store/src/remote/mod.rs @@ -20,10 +20,41 @@ use crate::ThreadStoreError; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; use proto::thread_store_client::ThreadStoreClient; +use tonic::codegen::InterceptedService; +use tonic::metadata::BinaryMetadataValue; +use tonic::service::Interceptor; +use tonic::transport::Channel; +use tonic::transport::Endpoint; #[path = "proto/codex.thread_store.v1.rs"] mod proto; +/// Metadata key used to forward the app-server's opaque identity key to remote contracts. +pub const IDENTITY_KEY_HEADER: &str = "x-codex-app-server-identity-key-bin"; + +#[derive(Clone, Debug)] +struct IdentityKeyInterceptor { + identity_key: Option>, +} + +impl Interceptor for IdentityKeyInterceptor { + fn call( + &mut self, + mut request: tonic::Request<()>, + ) -> Result, tonic::Status> { + if let Some(identity_key) = &self.identity_key { + request.metadata_mut().insert_bin( + IDENTITY_KEY_HEADER, + BinaryMetadataValue::from_bytes(identity_key), + ); + } + Ok(request) + } +} + +type RemoteThreadStoreClient = + ThreadStoreClient>; + /// gRPC-backed [`ThreadStore`] implementation for deployments whose durable thread data lives /// outside the app-server process. /// @@ -33,21 +64,43 @@ mod proto; #[derive(Clone, Debug)] pub struct RemoteThreadStore { endpoint: String, + identity_key: Option>, } impl RemoteThreadStore { pub fn new(endpoint: impl Into) -> Self { Self { endpoint: endpoint.into(), + identity_key: None, } } - async fn client(&self) -> ThreadStoreResult> { - ThreadStoreClient::connect(self.endpoint.clone()) + pub fn new_with_identity_key( + endpoint: impl Into, + identity_key: Option>, + ) -> Self { + Self { + endpoint: endpoint.into(), + identity_key, + } + } + + async fn client(&self) -> ThreadStoreResult { + let channel = Endpoint::new(self.endpoint.clone()) + .map_err(|err| ThreadStoreError::InvalidRequest { + message: format!("invalid remote thread store endpoint: {err}"), + })? + .connect() .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to connect to remote thread store: {err}"), - }) + })?; + Ok(ThreadStoreClient::with_interceptor( + channel, + IdentityKeyInterceptor { + identity_key: self.identity_key.clone(), + }, + )) } }