diff --git a/codex-rs/app-server/src/current_time.rs b/codex-rs/app-server/src/current_time.rs index 205068be2b..ae873bfc63 100644 --- a/codex-rs/app-server/src/current_time.rs +++ b/codex-rs/app-server/src/current_time.rs @@ -23,6 +23,7 @@ use crate::outgoing_message::OutgoingMessageSender; use crate::thread_state::ThreadStateManager; const CURRENT_TIME_REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const CURRENT_TIME_POLL_INTERVAL: Duration = Duration::from_secs(1); pub(crate) fn app_server_time_provider( outgoing: Arc, @@ -51,10 +52,32 @@ impl TimeProvider for AppServerTimeProvider { }) } - fn sleep(&self, _thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + fn sleep(&self, thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + let outgoing = self.outgoing.clone(); + let thread_state_manager = self.thread_state_manager.clone(); Box::pin(async move { - tokio::time::sleep(duration).await; - Ok(()) + let outgoing = outgoing + .upgrade() + .context("app-server current-time provider is unavailable")?; + let started_at = + request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await?; + let wake_at = started_at + .checked_add_signed( + chrono::Duration::from_std(duration) + .context("external sleep duration is outside the supported range")?, + ) + .context("external sleep deadline is outside the supported range")?; + + loop { + tokio::time::sleep(CURRENT_TIME_POLL_INTERVAL).await; + if request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await? + >= wake_at + { + return Ok(()); + } + } }) } } diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 18b49ede3b..a3d2191782 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -249,15 +249,12 @@ impl McpRequestProcessor { } None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; + let mcp_manager = self.thread_manager.mcp_manager(); let mcp_config = match thread { Some(thread) => thread.runtime_mcp_config(&config).await, - None => { - self.thread_manager - .mcp_manager() - .runtime_config(&config) - .await - } + None => mcp_manager.runtime_config(&config).await, }; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd @@ -274,6 +271,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; }); @@ -287,6 +285,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -294,6 +293,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; outgoing.send_result(request_id, result).await; @@ -305,6 +305,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, @@ -316,6 +317,7 @@ impl McpRequestProcessor { auth.as_ref(), request_id, runtime_context, + codex_apps_tools_cache, detail, ) .await; @@ -406,11 +408,9 @@ impl McpRequestProcessor { } let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let mcp_config = self - .thread_manager - .mcp_manager() - .runtime_config(&config) - .await; + let mcp_manager = self.thread_manager.mcp_manager(); + let mcp_config = mcp_manager.runtime_config(&config).await; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected @@ -425,6 +425,7 @@ impl McpRequestProcessor { &mcp_config, auth.as_ref(), runtime_context, + codex_apps_tools_cache, &server, &uri, ) diff --git a/codex-rs/app-server/tests/suite/v2/sleep.rs b/codex-rs/app-server/tests/suite/v2/sleep.rs index 872983f76d..8d88816fde 100644 --- a/codex-rs/app-server/tests/suite/v2/sleep.rs +++ b/codex-rs/app-server/tests/suite/v2/sleep.rs @@ -1,11 +1,12 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::to_response; +use codex_app_server_protocol::CurrentTimeReadResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -19,12 +20,16 @@ use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; +#[cfg(windows)] +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(25); +#[cfg(not(windows))] const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +const CURRENT_TIME_AT: i64 = 1_781_717_655; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sleep_emits_started_and_completed_items() -> Result<()> { +async fn external_sleep_polls_current_time_and_emits_items() -> Result<()> { const CALL_ID: &str = "sleep-1"; - const DURATION_MS: u64 = 1; + const DURATION_MS: u64 = 2_000; let server = responses::start_mock_server().await; responses::mount_sse_sequence( @@ -85,38 +90,19 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { .await??; let TurnStartResponse { turn, .. } = to_response(turn_start_response)?; - let (started, completed) = timeout(DEFAULT_READ_TIMEOUT, async { - let mut started = None; - let mut completed = None; - while started.is_none() || completed.is_none() { - let JSONRPCMessage::Notification(notification) = mcp.read_next_message().await? else { - continue; - }; - match notification.method.as_str() { - "item/started" => { - let payload: ItemStartedNotification = - serde_json::from_value(notification.params.expect("item/started params"))?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - started = Some(payload); - } - } - "item/completed" => { - let payload: ItemCompletedNotification = serde_json::from_value( - notification.params.expect("item/completed params"), - )?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - completed = Some(payload); - } - } - _ => {} - } - } - Ok::<_, anyhow::Error>(( - started.expect("sleep started"), - completed.expect("sleep completed"), - )) - }) - .await??; + // Read once for the initial reminder, then once to establish the sleep deadline. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + let started = wait_for_sleep_started(&mut mcp, CALL_ID).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + + // The first poll remains below the deadline, so the provider must request time again. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 1).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; + + let completed = wait_for_sleep_completed(&mut mcp, CALL_ID).await?; + + // The next inference boundary reads the same external clock after the sleep completes. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -150,6 +136,64 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { Ok(()) } +async fn wait_for_sleep_started( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/started"), + ) + .await??; + let started: ItemStartedNotification = + serde_json::from_value(notification.params.expect("item/started params"))?; + if matches!(&started.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(started); + } + } +} + +async fn wait_for_sleep_completed( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/completed"), + ) + .await??; + let completed: ItemCompletedNotification = + serde_json::from_value(notification.params.expect("item/completed params"))?; + if matches!(&completed.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(completed); + } + } +} + +async fn respond_to_current_time_read( + mcp: &mut TestAppServer, + thread_id: &str, + current_time_at: i64, +) -> Result<()> { + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::CurrentTimeRead { request_id, params } = request else { + panic!("expected CurrentTimeRead request, got: {request:?}"); + }; + assert_eq!(params.thread_id, thread_id); + mcp.send_response( + request_id, + serde_json::to_value(CurrentTimeReadResponse { current_time_at })?, + ) + .await?; + Ok(()) +} + fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { std::fs::write( codex_home.join("config.toml"), @@ -170,6 +214,7 @@ stream_max_retries = 0 [features.current_time_reminder] enabled = true sleep_tool = true +clock_source = "external" "# ), ) diff --git a/codex-rs/cli/src/exec_server_telemetry.rs b/codex-rs/cli/src/exec_server_telemetry.rs index 1645bd5744..b5af769e6a 100644 --- a/codex-rs/cli/src/exec_server_telemetry.rs +++ b/codex-rs/cli/src/exec_server_telemetry.rs @@ -1,3 +1,5 @@ +use std::future::Future; + use tracing_subscriber::EnvFilter; use tracing_subscriber::prelude::*; @@ -43,6 +45,70 @@ pub(crate) fn init( (otel, telemetry) } +pub(crate) async fn run_until_shutdown(run: F) -> Result<(), E> +where + F: Future>, +{ + let shutdown_signal = match shutdown_signal() { + Ok(signal) => Some(signal), + Err(error) => { + eprintln!("Could not listen for exec-server shutdown signal: {error}"); + None + } + }; + tokio::pin!(run); + + if let Some(shutdown_signal) = shutdown_signal { + tokio::select! { + result = &mut run => result, + signal = wait_for_shutdown_signal(shutdown_signal) => { + match signal { + Ok(()) => Ok(()), + Err(error) => { + eprintln!("Could not listen for exec-server shutdown signal: {error}"); + run.await + } + } + } + } + } else { + run.await + } +} + +#[cfg(unix)] +struct ShutdownSignal { + terminate: tokio::signal::unix::Signal, +} + +#[cfg(unix)] +fn shutdown_signal() -> std::io::Result { + Ok(ShutdownSignal { + terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?, + }) +} + +#[cfg(unix)] +async fn wait_for_shutdown_signal(mut shutdown_signal: ShutdownSignal) -> std::io::Result<()> { + tokio::select! { + result = tokio::signal::ctrl_c() => result, + _ = shutdown_signal.terminate.recv() => Ok(()), + } +} + +#[cfg(not(unix))] +struct ShutdownSignal; + +#[cfg(not(unix))] +fn shutdown_signal() -> std::io::Result { + Ok(ShutdownSignal) +} + +#[cfg(not(unix))] +async fn wait_for_shutdown_signal(_: ShutdownSignal) -> std::io::Result<()> { + tokio::signal::ctrl_c().await +} + fn stderr_env_filter() -> EnvFilter { EnvFilter::try_from_default_env() .or_else(|_| EnvFilter::try_new(DEFAULT_LOG_FILTER)) diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index d7dbff47ec..c766baa048 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1706,7 +1706,10 @@ async fn run_exec_server_command( remote_config.name = name; } let remote_config = remote_config.with_telemetry(telemetry); - codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; + exec_server_telemetry::run_until_shutdown(async move { + codex_exec_server::run_remote_environment(remote_config, runtime_paths).await + }) + .await?; Ok(()) } else { let config_result = load_exec_server_config(root_config_overrides, strict_config).await; @@ -1718,11 +1721,12 @@ async fn run_exec_server_command( let (_otel, telemetry) = exec_server_telemetry::init(config.as_ref()); let listen_url = cmd .listen - .as_deref() - .unwrap_or(codex_exec_server::DEFAULT_LISTEN_URL); - codex_exec_server::run_main_with_telemetry(listen_url, runtime_paths, telemetry) - .await - .map_err(anyhow::Error::from_boxed) + .unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string()); + exec_server_telemetry::run_until_shutdown(async move { + codex_exec_server::run_main_with_telemetry(&listen_url, runtime_paths, telemetry).await + }) + .await + .map_err(anyhow::Error::from_boxed) } } diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index c54e9c40b2..aabe9c0453 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -1,6 +1,20 @@ +#[cfg(unix)] +use std::io::BufRead as _; +#[cfg(unix)] +use std::io::BufReader as StdBufReader; +#[cfg(unix)] +use std::io::Read as _; +#[cfg(unix)] +use std::io::Write as _; +#[cfg(unix)] +use std::net::TcpStream; use std::path::Path; use std::process::Stdio; +#[cfg(unix)] +use std::thread; use std::time::Duration; +#[cfg(unix)] +use std::time::Instant; use anyhow::Result; use predicates::prelude::PredicateBooleanExt; @@ -220,6 +234,54 @@ async fn send_json_line( Ok(()) } +#[cfg(unix)] +#[test] +fn local_exec_server_exits_successfully_on_sigterm() -> Result<()> { + let codex_home = TempDir::new()?; + let mut child = std::process::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?) + .env("CODEX_HOME", codex_home.path()) + .args(["exec-server", "--listen", "ws://127.0.0.1:0"]) + .stdout(Stdio::piped()) + .spawn()?; + let mut listen_url = String::new(); + StdBufReader::new(child.stdout.take().expect("child stdout")).read_line(&mut listen_url)?; + assert!(listen_url.starts_with("ws://127.0.0.1:"), "{listen_url}"); + + let listen_addr = listen_url + .trim() + .strip_prefix("ws://") + .expect("listen URL should use ws://") + .parse()?; + let deadline = Instant::now() + Duration::from_secs(5); + let mut ready = false; + while let Some(remaining) = deadline.checked_duration_since(Instant::now()) { + if let Ok(mut stream) = + TcpStream::connect_timeout(&listen_addr, remaining.min(Duration::from_millis(100))) + { + let _ = stream.set_read_timeout(Some(Duration::from_secs(1))); + let request = + format!("GET /readyz HTTP/1.1\r\nHost: {listen_addr}\r\nConnection: close\r\n\r\n"); + let mut response = String::new(); + if stream.write_all(request.as_bytes()).is_ok() + && stream.read_to_string(&mut response).is_ok() + && response.starts_with("HTTP/1.1 200") + { + ready = true; + break; + } + } + thread::sleep(Duration::from_millis(10)); + } + assert!(ready, "exec-server did not become ready at {listen_url}"); + + // SAFETY: `child.id()` is the live process spawned above. + let result = unsafe { libc::kill(child.id() as libc::pid_t, libc::SIGTERM) }; + assert_eq!(result, 0); + let status = child.wait()?; + assert!(status.success(), "{status}"); + Ok(()) +} + async fn wait_for_response( stdout: &mut (impl tokio::io::AsyncBufRead + Unpin), expected_id: i64, diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs index 6bf83d367d..ec07b92d41 100644 --- a/codex-rs/code-mode-protocol/src/host/codec.rs +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -11,6 +11,36 @@ use tokio::io::AsyncWriteExt; /// Maximum JSON payload size accepted for one IPC frame. pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024; +/// A serialized IPC frame that has already passed the payload size limit. +#[derive(Clone, Debug)] +pub struct EncodedFrame { + payload: Vec, +} + +impl EncodedFrame { + pub fn encode(message: &T) -> io::Result + where + T: Serialize, + { + let payload = serde_json::to_vec(message).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to encode code-mode IPC frame: {err}"), + ) + })?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", + payload.len() + ), + )); + } + Ok(Self { payload }) + } +} + /// Decodes JSON messages prefixed by a four-byte little-endian payload length. pub struct FramedReader { reader: R, @@ -72,22 +102,12 @@ where where T: Serialize, { - let payload = serde_json::to_vec(message).map_err(|err| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("failed to encode code-mode IPC frame: {err}"), - ) - })?; - if payload.len() > MAX_FRAME_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!( - "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", - payload.len() - ), - )); - } - let length = u32::try_from(payload.len()).map_err(|_| { + self.write_frame(&EncodedFrame::encode(message)?).await + } + + /// Writes and flushes a frame encoded before it entered an I/O queue. + pub async fn write_frame(&mut self, frame: &EncodedFrame) -> io::Result<()> { + let length = u32::try_from(frame.payload.len()).map_err(|_| { io::Error::new( io::ErrorKind::InvalidData, "code-mode IPC frame length exceeds u32", @@ -95,7 +115,7 @@ where })?; self.writer.write_all(&length.to_le_bytes()).await?; - self.writer.write_all(&payload).await?; + self.writer.write_all(&frame.payload).await?; self.writer.flush().await } } diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 3aa2a18307..d56e56e1fa 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -365,6 +365,16 @@ fn client_to_host_v1_variants_are_pinned() { }), ); } + + assert_wire_round_trip( + ClientToHost::CancelRequest { + id: request_id(/*value*/ 9), + }, + json!({ + "type": "operation/cancel", + "id": 9, + }), + ); } #[test] diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index 78f0961d51..0e83c866e4 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -132,6 +132,8 @@ pub enum ClientToHost { ClientHello(ClientHello), #[serde(rename = "operation/request")] Request { id: RequestId, request: HostRequest }, + #[serde(rename = "operation/cancel")] + CancelRequest { id: RequestId }, #[serde(rename = "delegate/response")] DelegateResponse { id: DelegateRequestId, diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index 7e78768ab6..3583ccb3e9 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -11,6 +11,7 @@ mod message; mod payload; mod types; +pub use codec::EncodedFrame; pub use codec::FramedReader; pub use codec::FramedWriter; pub use codec::MAX_FRAME_BYTES; diff --git a/codex-rs/codex-mcp/src/codex_apps.rs b/codex-rs/codex-mcp/src/codex_apps.rs index c1b5fd5fa5..a33f735366 100644 --- a/codex-rs/codex-mcp/src/codex_apps.rs +++ b/codex-rs/codex-mcp/src/codex_apps.rs @@ -1,70 +1,9 @@ //! Codex Apps support for the host-owned apps MCP server. //! -//! This module owns the pieces that are unique to ChatGPT-hosted app -//! connectors: cache scoping by authenticated user, disk cache reads/writes, -//! connector allow-list filtering, and the normalization that turns app +//! This module owns the normalization that turns ChatGPT-hosted app //! connector/tool metadata into model-visible MCP callable names. -use std::path::PathBuf; -use std::time::Instant; - -use crate::runtime::emit_duration; -use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; -use crate::tools::ToolInfo; -use anyhow::Context; -use codex_login::CodexAuth; -use codex_protocol::mcp::McpServerInfo; use codex_utils_plugins::mcp_connector::sanitize_name; -use serde::Deserialize; -use serde::Serialize; -use sha1::Digest; -use sha1::Sha1; -use tracing::instrument; - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CodexAppsToolsCacheKey { - pub(crate) account_id: Option, - pub(crate) chatgpt_user_id: Option, - pub(crate) is_workspace_account: bool, -} - -pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { - CodexAppsToolsCacheKey { - account_id: auth.and_then(CodexAuth::get_account_id), - chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), - is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), - } -} - -#[derive(Clone)] -pub(crate) struct CodexAppsToolsCacheContext { - pub(crate) codex_home: PathBuf, - pub(crate) user_key: CodexAppsToolsCacheKey, -} - -impl CodexAppsToolsCacheContext { - pub(crate) fn tools_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) - } - - pub(crate) fn server_info_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) - } - - fn cache_path_in(&self, cache_dir: &str) -> PathBuf { - let user_key_json = serde_json::to_string(&self.user_key).unwrap_or_default(); - let user_key_hash = sha1_hex(&user_key_json); - self.codex_home - .join(cache_dir) - .join(format!("{user_key_hash}.json")) - } -} - -pub(crate) enum CachedCodexAppsToolsLoad { - Hit(Vec), - Missing, - Invalid, -} pub(crate) fn normalize_codex_apps_tool_title(connector_name: Option<&str>, value: &str) -> String { let Some(connector_name) = connector_name @@ -124,152 +63,3 @@ pub(crate) fn normalize_codex_apps_callable_namespace( server_name.to_string() } } - -pub(crate) fn write_codex_apps_tools_cache( - cache_context: Option<&CodexAppsToolsCacheContext>, - server_info: &McpServerInfo, - tools: &[ToolInfo], -) { - if let Some(cache_context) = cache_context { - let cache_write_start = Instant::now(); - write_cached_codex_apps_tools(cache_context, tools); - if let Err(err) = write_cached_codex_apps_server_info(cache_context, server_info) { - tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); - } - emit_duration( - MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, - cache_write_start.elapsed(), - &[], - ); - } -} - -pub(crate) fn load_startup_cached_codex_apps_tools_snapshot( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option> { - let cache_context = cache_context?; - - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -pub(crate) fn load_startup_cached_codex_apps_server_info( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option { - load_cached_codex_apps_server_info(cache_context?) -} - -#[cfg(test)] -pub(crate) fn read_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> Option> { - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> CachedCodexAppsToolsLoad { - let cache_path = cache_context.tools_cache_path(); - let bytes = match std::fs::read(cache_path) { - Ok(bytes) => bytes, - Err(err) if err.kind() == std::io::ErrorKind::NotFound => { - return CachedCodexAppsToolsLoad::Missing; - } - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - let cache: CodexAppsToolsDiskCache = match serde_json::from_slice(&bytes) { - Ok(cache) => cache, - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - if cache.schema_version != CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION { - return CachedCodexAppsToolsLoad::Invalid; - } - CachedCodexAppsToolsLoad::Hit(cache.tools) -} - -pub(crate) fn write_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, - tools: &[ToolInfo], -) { - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() - && std::fs::create_dir_all(parent).is_err() - { - return; - } - let Ok(bytes) = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { - schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - tools: tools.to_vec(), - }) else { - return; - }; - let _ = std::fs::write(cache_path, bytes); -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, -) -> Option { - let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; - let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; - (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) - .then_some(cache.server_info) -} - -fn write_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, - server_info: &McpServerInfo, -) -> anyhow::Result<()> { - let cache_path = cache_context.server_info_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).with_context(|| { - format!( - "failed to create Codex Apps server info cache directory `{}`", - parent.display() - ) - })?; - } - let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { - schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, - server_info: server_info.clone(), - }) - .context("failed to serialize Codex Apps server info cache")?; - std::fs::write(&cache_path, bytes).with_context(|| { - format!( - "failed to write Codex Apps server info cache `{}`", - cache_path.display() - ) - })?; - Ok(()) -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsToolsDiskCache { - schema_version: u8, - tools: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsServerInfoDiskCache { - schema_version: u8, - server_info: McpServerInfo, -} - -const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; -pub(crate) const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; - -const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; -const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; - -fn sha1_hex(s: &str) -> String { - let mut hasher = Sha1::new(); - hasher.update(s.as_bytes()); - let sha1 = hasher.finalize(); - format!("{sha1:x}") -} diff --git a/codex-rs/codex-mcp/src/codex_apps_cache.rs b/codex-rs/codex-mcp/src/codex_apps_cache.rs new file mode 100644 index 0000000000..ea39444628 --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache.rs @@ -0,0 +1,378 @@ +//! Shared raw tool cache for the host-owned Codex Apps MCP server. +//! +//! Cache entries are process-local live state scoped by the active Codex auth +//! key. Disk is best-effort cold-start persistence; entries do not reread disk +//! after creation. + +use std::collections::HashMap; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Context; +use arc_swap::ArcSwapOption; +use codex_login::CodexAuth; +use codex_protocol::mcp::McpServerInfo; +use serde::Deserialize; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; +use tracing::instrument; + +use crate::runtime::emit_duration; +use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; +use crate::tools::ToolInfo; + +const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms"; + +/// The CodexAuth bits that identify a Codex Apps catalog. +/// +/// Debug bearer-token overrides bypass the shared cache, so shared entries only +/// need the CodexAuth-backed identity. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct CodexAppsToolsCacheKey { + pub(crate) account_id: Option, + pub(crate) chatgpt_user_id: Option, + pub(crate) is_workspace_account: bool, +} + +/// Builds the CodexAuth-backed Codex Apps cache key. +pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { + account_id: auth.and_then(CodexAuth::get_account_id), + chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), + is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), + } +} + +/// Process-scoped registry for shared Codex Apps raw tool snapshots. +/// +/// Two clients share an entry only when they would read the same Codex Apps +/// catalog. New entries may seed from disk; live entries read from memory only. +#[derive(Clone, Default)] +pub struct CodexAppsToolsCache { + entries: Arc>>>, +} + +/// Handle to one shared Codex Apps tools cache entry. +/// +/// The connection manager creates this from the auth key, then tool +/// reads and refreshes for that managed client use the same entry. +#[derive(Clone)] +pub(crate) struct CodexAppsToolsCacheContext { + entry: Arc, +} + +impl CodexAppsToolsCacheContext { + pub(crate) fn tools_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) + } + + pub(crate) fn server_info_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) + } + + pub(crate) fn current_tools(&self) -> Option> { + self.entry + .current_tools + .load_full() + .map(|tools| tools.as_ref().clone()) + } + + pub(crate) fn has_current_tools(&self) -> bool { + self.entry.current_tools.load_full().is_some() + } + + pub(crate) fn begin_fetch( + &self, + source: CodexAppsToolsFetchSource, + ) -> CodexAppsToolsFetchTicket { + CodexAppsToolsFetchTicket { + generation: self + .entry + .next_fetch_generation + .fetch_add(1, Ordering::Relaxed) + + 1, + source, + } + } + + pub(crate) fn publish_if_newest_accepted( + &self, + ticket: CodexAppsToolsFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + ) -> Vec { + let publish_start = Instant::now(); + let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation); + if ticket.generation <= *last_accepted_generation { + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "stale")], + ); + return self.current_tools().unwrap_or(tools); + } + + *last_accepted_generation = ticket.generation; + self.entry + .current_tools + .store(Some(Arc::new(tools.clone()))); + persist_codex_apps_cache(self, server_info, &tools); + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "published")], + ); + tools + } + + #[cfg(test)] + pub(crate) fn store_current_tools_for_test(&self, tools: Vec) { + self.entry.current_tools.store(Some(Arc::new(tools))); + } +} + +impl CodexAppsToolsCache { + pub(crate) fn context( + &self, + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, + ) -> CodexAppsToolsCacheContext { + let identity = CodexAppsToolsCacheIdentity { + codex_home, + auth_key, + }; + let mut entries = lock_unpoisoned(&self.entries); + let entry = entries + .entry(identity.clone()) + .or_insert_with(|| Arc::new(CodexAppsToolsCacheEntry::new(identity))) + .clone(); + CodexAppsToolsCacheContext { entry } + } +} + +#[derive(Debug, Clone, Copy)] +pub(crate) enum CodexAppsToolsFetchSource { + Startup, + HardRefresh, +} + +impl CodexAppsToolsFetchSource { + fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::HardRefresh => "hard_refresh", + } + } +} + +pub(crate) struct CodexAppsToolsFetchTicket { + generation: u64, + source: CodexAppsToolsFetchSource, +} + +struct CodexAppsToolsCacheEntry { + identity: CodexAppsToolsCacheIdentity, + current_tools: ArcSwapOption>, + next_fetch_generation: AtomicU64, + last_accepted_generation: Mutex, +} + +impl CodexAppsToolsCacheEntry { + fn new(identity: CodexAppsToolsCacheIdentity) -> Self { + let current_tools = load_cached_codex_apps_tools_for_identity(&identity).map(Arc::new); + Self { + identity, + current_tools: ArcSwapOption::from(current_tools), + next_fetch_generation: AtomicU64::new(0), + last_accepted_generation: Mutex::new(0), + } + } +} + +/// Everything that decides whether two Codex Apps clients can share tools. +/// +/// The auth key says whose catalog we are reading. `codex_home` keeps the +/// persisted cache under the right home directory. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CodexAppsToolsCacheIdentity { + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, +} + +impl CodexAppsToolsCacheIdentity { + fn cache_path_in(&self, cache_dir: &str) -> PathBuf { + // `codex_home` is already the parent directory. Keep it out of the + // filename hash so non-UTF-8 Unix paths cannot collapse distinct auth + // keys onto the same disk cache file. + let identity_json = serde_json::to_string(&self.auth_key).unwrap_or_default(); + let identity_hash = sha1_hex(&identity_json); + self.codex_home + .join(cache_dir) + .join(format!("{identity_hash}.json")) + } +} + +#[cfg(test)] +fn write_cached_codex_apps_tools_for_test( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + cache_context + .entry + .current_tools + .store(Some(Arc::new(tools.to_vec()))); + persist_codex_apps_cache(cache_context, server_info, tools); +} + +pub(crate) fn load_startup_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + load_cached_codex_apps_server_info(cache_context) +} + +#[cfg(test)] +fn read_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, +) -> Option> { + load_cached_codex_apps_tools_for_identity(&cache_context.entry.identity) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_tools_for_identity( + identity: &CodexAppsToolsCacheIdentity, +) -> Option> { + let cache_path = identity.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR); + let bytes = std::fs::read(cache_path).ok()?; + let cache: CodexAppsToolsDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION).then_some(cache.tools) +} + +fn write_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, + tools: &[ToolInfo], +) -> anyhow::Result<()> { + let cache_path = cache_context.tools_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { + schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + tools: tools.to_vec(), + }) + .context("failed to serialize Codex Apps tools cache")?; + write_codex_apps_cache_file(&cache_path, "tools", bytes) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; + let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) + .then_some(cache.server_info) +} + +fn write_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, +) -> anyhow::Result<()> { + let cache_path = cache_context.server_info_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { + schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, + server_info: server_info.clone(), + }) + .context("failed to serialize Codex Apps server info cache")?; + write_codex_apps_cache_file(&cache_path, "server info", bytes) +} + +fn write_codex_apps_cache_file( + cache_path: &Path, + cache_name: &str, + bytes: Vec, +) -> anyhow::Result<()> { + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create Codex Apps {cache_name} cache directory `{}`", + parent.display() + ) + })?; + } + std::fs::write(cache_path, bytes).with_context(|| { + format!( + "failed to write Codex Apps {cache_name} cache `{}`", + cache_path.display() + ) + })?; + Ok(()) +} + +fn persist_codex_apps_cache( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + let cache_write_start = Instant::now(); + let tools_result = write_cached_codex_apps_tools(cache_context, tools); + if let Err(err) = &tools_result { + tracing::warn!("failed to write Codex Apps tools cache: {err:#}"); + } + let server_info_result = write_cached_codex_apps_server_info(cache_context, server_info); + if let Err(err) = &server_info_result { + tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); + } + let status = if tools_result.is_ok() && server_info_result.is_ok() { + "success" + } else { + "failure" + }; + emit_duration( + MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, + cache_write_start.elapsed(), + &[("status", status)], + ); +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsToolsDiskCache { + schema_version: u8, + tools: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsServerInfoDiskCache { + schema_version: u8, + server_info: McpServerInfo, +} + +const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; +const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; + +const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; +const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; + +fn sha1_hex(s: &str) -> String { + let mut hasher = Sha1::new(); + hasher.update(s.as_bytes()); + let sha1 = hasher.finalize(); + format!("{sha1:x}") +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +#[cfg(test)] +#[path = "codex_apps_cache_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs new file mode 100644 index 0000000000..46ad311a1b --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs @@ -0,0 +1,464 @@ +use super::*; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; +use crate::tools::ToolInfo; +use codex_protocol::ToolName; +use codex_protocol::mcp::McpServerInfo; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Tool; +use std::collections::HashSet; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::tempdir; + +fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: server_name.to_string(), + namespace_description: None, + tool: Tool::new( + tool_name.to_string(), + format!("Test tool: {tool_name}"), + Arc::new(JsonObject::default()), + ), + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } +} + +fn create_test_tool_with_connector( + server_name: &str, + tool_name: &str, + connector_id: &str, + connector_name: Option<&str>, +) -> ToolInfo { + let mut tool = create_test_tool(server_name, tool_name); + tool.connector_id = Some(connector_id.to_string()); + tool.connector_name = connector_name.map(ToOwned::to_owned); + tool +} + +fn create_codex_apps_tools_cache_context( + codex_home: PathBuf, + account_id: Option<&str>, + chatgpt_user_id: Option<&str>, +) -> CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( + codex_home, + CodexAppsToolsCacheKey { + account_id: account_id.map(ToOwned::to_owned), + chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), + is_workspace_account: false, + }, + ) +} + +fn create_test_server_info(title: &str) -> McpServerInfo { + McpServerInfo { + name: "codex-apps".to_string(), + title: Some(title.to_string()), + version: "1.0.0".to_string(), + description: None, + icons: None, + website_url: None, + } +} + +fn model_tool_names(tools: &[ToolInfo]) -> HashSet { + tools + .iter() + .map(ToolInfo::canonical_tool_name) + .collect::>() +} + +#[test] +fn codex_apps_tools_cache_is_overwritten_by_last_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_1).expect("write first cache"); + let cached_gateway_1 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); + assert_eq!(cached_gateway_1[0].callable_name, "one"); + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_2).expect("write second cache"); + let cached_gateway_2 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); + assert_eq!(cached_gateway_2[0].callable_name, "two"); +} + +#[test] +fn codex_apps_tools_cache_is_scoped_per_user() { + let codex_home = tempdir().expect("tempdir"); + let cache_context_user_1 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_context_user_2 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-two"), + Some("user-two"), + ); + let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1) + .expect("write user one cache"); + write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2) + .expect("write user two cache"); + + let read_user_1 = + read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); + let read_user_2 = + read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); + + assert_eq!(read_user_1[0].callable_name, "one"); + assert_eq!(read_user_2[0].callable_name, "two"); + assert_ne!( + cache_context_user_1.tools_cache_path(), + cache_context_user_2.tools_cache_path(), + "each user should get an isolated cache file" + ); +} + +#[test] +fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools = vec![ + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "formerly_blocked_tool", + "connector_2b0a9009c9c64bf9933a3dae3f2b1254", + Some("Formerly Blocked"), + ), + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_tool", + "calendar", + Some("Calendar"), + ), + ]; + + write_cached_codex_apps_tools(&cache_context, &tools).expect("write cache"); + let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); + + assert_eq!( + cached + .iter() + .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) + .collect::>(), + vec![ + ( + "formerly_blocked_tool", + Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") + ), + ("calendar_tool", Some("calendar")), + ] + ); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + std::fs::write(cache_path, b"{not json").expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_from_disk_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )]; + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test(&writer_cache_context, &server_info, &cached_tools); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("expected startup snapshot to load from cache"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, Some(server_info)); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = writer_cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("legacy startup snapshot should remain available"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, None); +} + +#[test] +fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test( + &cache_context, + &server_info, + &[create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )], + ); + + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write legacy tools cache"); + let startup_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + assert_eq!( + load_startup_cached_codex_apps_server_info(&startup_cache_context), + Some(server_info) + ); + assert!(startup_cache_context.current_tools().is_none()); +} + +#[test] +fn codex_apps_tools_cache_context_does_not_reread_disk_after_creation() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "cached")]; + write_cached_codex_apps_tools(&writer_cache_context, &cached_tools).expect("write cache"); + let reader_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let updated_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")]; + write_cached_codex_apps_tools(&writer_cache_context, &updated_tools).expect("rewrite cache"); + + assert_eq!( + reader_cache_context + .current_tools() + .expect("in-memory tools")[0] + .callable_name, + "cached" + ); + assert_eq!( + read_cached_codex_apps_tools(&writer_cache_context).expect("disk tools")[0].callable_name, + "updated" + ); +} + +#[test] +fn codex_apps_tools_cache_publishes_newest_shared_snapshot() { + let codex_home = tempdir().expect("tempdir"); + let cache = CodexAppsToolsCache::default(); + let cache_context_1 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let cache_context_2 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let older_ticket = cache_context_1.begin_fetch(CodexAppsToolsFetchSource::Startup); + let newer_ticket = cache_context_2.begin_fetch(CodexAppsToolsFetchSource::HardRefresh); + let server_info = create_test_server_info("Codex Apps"); + let newer_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")]; + let older_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")]; + + let published_tools = + cache_context_2.publish_if_newest_accepted(newer_ticket, &server_info, newer_tools); + assert_eq!( + model_tool_names(&published_tools), + model_tool_names( + &cache_context_1 + .current_tools() + .expect("new snapshot should publish") + ) + ); + let current_tools = + cache_context_1.publish_if_newest_accepted(older_ticket, &server_info, older_tools); + + assert_eq!(current_tools[0].callable_name, "newer"); + assert_eq!( + cache_context_2.current_tools().expect("shared snapshot")[0].callable_name, + "newer" + ); + assert_eq!( + read_cached_codex_apps_tools(&cache_context_1).expect("persisted snapshot")[0] + .callable_name, + "newer" + ); +} + +#[test] +fn codex_apps_tools_cache_keeps_live_publish_when_disk_persistence_fails() { + let codex_home = tempdir().expect("tempdir"); + let codex_home_file = codex_home.path().join("not-a-directory"); + std::fs::write(&codex_home_file, b"occupied").expect("create codex home file"); + let cache_context = CodexAppsToolsCache::default().context( + codex_home_file, + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "live")]; + let published_tools = cache_context.publish_if_newest_accepted( + cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools.clone(), + ); + + assert_eq!(model_tool_names(&published_tools), model_tool_names(&tools)); + assert_eq!( + model_tool_names(&cache_context.current_tools().expect("live snapshot")), + model_tool_names(&tools) + ); +} + +#[cfg(unix)] +#[test] +fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { + let codex_home = PathBuf::from(std::ffi::OsString::from_vec( + b"/tmp/codex-home-\xff".to_vec(), + )); + let cache = CodexAppsToolsCache::default(); + let user_one_context = cache.context( + codex_home.clone(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let user_two_context = cache.context( + codex_home, + CodexAppsToolsCacheKey { + account_id: Some("account-two".to_string()), + chatgpt_user_id: Some("user-two".to_string()), + is_workspace_account: false, + }, + ); + let cache_paths = [ + user_one_context.tools_cache_path(), + user_two_context.tools_cache_path(), + ]; + + assert_eq!( + cache_paths.iter().collect::>().len(), + cache_paths.len() + ); +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 55c72b0467..0e81549556 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,7 +7,6 @@ //! `codex-core`. use std::collections::HashMap; -use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::Ordering; @@ -15,9 +14,9 @@ use std::time::Duration; use std::time::Instant; use crate::McpAuthStatusEntry; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::CodexAppsToolsCacheKey; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheKey; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -131,6 +130,7 @@ impl McpConnectionManager { initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, + codex_apps_tools_cache: CodexAppsToolsCache, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, @@ -174,13 +174,32 @@ impl McpConnectionManager { }, ) .await; - let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { - codex_apps_tools_cache_context(&codex_home, &codex_apps_tools_cache_key) - } else { - regular_mcp_tools_cache_context() - }; + let configured_config = server.configured_config(); + // For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug + // override: it supplies runtime auth but bypasses the shared tools + // cache. + let uses_env_bearer_token = + configured_config.is_some_and(|config| match &config.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }); + let shares_codex_apps_tools_cache = + should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); + let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { + codex_apps_tools_cache + .context(codex_home.clone(), codex_apps_tools_cache_key.clone()) + }); + // If Codex Apps has an env bearer token, that is its auth path. Do + // not also attach the ambient CodexAuth provider. let runtime_auth_provider = - chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()); + if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token { + None + } else { + chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()) + }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), server, @@ -445,13 +464,13 @@ impl McpConnectionManager { pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); for (server_name, managed_client) in &self.clients { - let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some(); + let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete .load(std::sync::atomic::Ordering::Acquire); trace!( server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete, "waiting for MCP server tools while building tool list" ); @@ -460,7 +479,7 @@ impl McpConnectionManager { .instrument(trace_span!( "list_tools_for_server", server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete )) .await @@ -483,9 +502,9 @@ impl McpConnectionManager { /// Force-refresh codex apps tools by bypassing the in-process cache. /// - /// On success, the refreshed tools replace the cache contents and the - /// latest filtered tools are returned directly to the caller. On - /// failure, the existing cache remains unchanged. + /// On success, the refreshed tools replace shared cache contents when the + /// cache is enabled and the latest filtered tools are returned directly to + /// the caller. On failure, existing shared cache contents remain unchanged. pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let managed_client = self .clients @@ -497,6 +516,10 @@ impl McpConnectionManager { let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = managed_client + .codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh)); let tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, @@ -514,11 +537,16 @@ impl McpConnectionManager { &[], ); - write_codex_apps_tools_cache( - managed_client.codex_apps_tools_cache_context.as_ref(), - &managed_client.server_info, - &tools, - ); + let tools = + match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context + .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), @@ -849,22 +877,6 @@ impl Drop for McpConnectionManager { } } -/// Creates the per-user tools cache context used only by the Codex Apps server. -fn codex_apps_tools_cache_context( - codex_home: &Path, - codex_apps_tools_cache_key: &CodexAppsToolsCacheKey, -) -> Option { - Some(CodexAppsToolsCacheContext { - codex_home: codex_home.to_path_buf(), - user_key: codex_apps_tools_cache_key.clone(), - }) -} - -/// Keeps regular MCP servers isolated from the Codex Apps tools cache. -fn regular_mcp_tools_cache_context() -> Option { - None -} - /// Makes ChatGPT authentication available to servers that explicitly opt in. /// The HTTP transport applies it only when no configured authorization resolves. fn chatgpt_auth_provider_for_server( @@ -880,6 +892,10 @@ fn chatgpt_auth_provider_for_server( chatgpt_auth_provider } +fn should_share_codex_apps_tools_cache(server_name: &str, uses_env_bearer_token: bool) -> bool { + server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token +} + async fn emit_update( submit_id: &str, tx_event: &Sender, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d66474c7bb..05caf0489c 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,11 +1,6 @@ use super::*; -use crate::codex_apps::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; -use crate::codex_apps::read_cached_codex_apps_tools; -use crate::codex_apps::write_cached_codex_apps_tools; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; use crate::declared_openai_file_input_param_names; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::elicitation_is_rejected_by_policy; @@ -31,7 +26,10 @@ use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpAuthStatus; +use codex_rmcp_client::InProcessTransportFactory; +use codex_rmcp_client::RmcpClient; use futures::FutureExt; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; @@ -41,8 +39,10 @@ use rmcp::model::Meta; use rmcp::model::NumberOrString; use rmcp::model::Tool; use std::collections::HashSet; +use std::io; use std::sync::Arc; use tempfile::tempdir; +use tokio::io::DuplexStream; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -63,31 +63,19 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { } } -fn create_test_tool_with_connector( - server_name: &str, - tool_name: &str, - connector_id: &str, - connector_name: Option<&str>, -) -> ToolInfo { - let mut tool = create_test_tool(server_name, tool_name); - tool.connector_id = Some(connector_id.to_string()); - tool.connector_name = connector_name.map(ToOwned::to_owned); - tool -} - fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, chatgpt_user_id: Option<&str>, ) -> CodexAppsToolsCacheContext { - CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( codex_home, - user_key: CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { account_id: account_id.map(ToOwned::to_owned), chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), is_workspace_account: false, }, - } + ) } fn create_test_server_info(title: &str) -> McpServerInfo { @@ -101,6 +89,51 @@ fn create_test_server_info(title: &str) -> McpServerInfo { } } +struct TestInProcessTransportFactory; + +impl InProcessTransportFactory for TestInProcessTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + async { + let (client_stream, _server_stream) = tokio::io::duplex(1); + Ok(client_stream) + } + .boxed() + } +} + +async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { + let tool_filter = ToolFilter::default(); + let managed_client = ManagedClient { + client: Arc::new( + RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) + .await + .expect("create in-process RMCP client"), + ), + server_info: create_test_server_info("Ready"), + tools, + tool_filter: tool_filter.clone(), + tool_timeout: None, + server_instructions: None, + server_supports_sandbox_state_meta_capability: false, + codex_apps_tools_cache_context: None, + }; + + AsyncManagedClient { + client: futures::future::ready::>(Ok( + managed_client, + )) + .boxed() + .shared(), + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + } +} + fn model_tool_names(tools: &[ToolInfo]) -> HashSet { tools .iter() @@ -554,235 +587,25 @@ fn filter_tools_applies_per_server_filters() { } #[test] -fn codex_apps_tools_cache_is_overwritten_by_last_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_1); - let cached_gateway_1 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); - assert_eq!(cached_gateway_1[0].callable_name, "one"); - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_2); - let cached_gateway_2 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); - assert_eq!(cached_gateway_2[0].callable_name, "two"); -} - -#[test] -fn codex_apps_tools_cache_is_scoped_per_user() { - let codex_home = tempdir().expect("tempdir"); - let cache_context_user_1 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_context_user_2 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-two"), - Some("user-two"), - ); - let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1); - write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2); - - let read_user_1 = - read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); - let read_user_2 = - read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); - - assert_eq!(read_user_1[0].callable_name, "one"); - assert_eq!(read_user_2[0].callable_name, "two"); - assert_ne!( - cache_context_user_1.tools_cache_path(), - cache_context_user_2.tools_cache_path(), - "each user should get an isolated cache file" - ); -} - -#[test] -fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools = vec![ - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "formerly_blocked_tool", - "connector_2b0a9009c9c64bf9933a3dae3f2b1254", - Some("Formerly Blocked"), - ), - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_tool", - "calendar", - Some("Calendar"), - ), - ]; - - write_cached_codex_apps_tools(&cache_context, &tools); - let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); - - assert_eq!( - cached - .iter() - .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) - .collect::>(), - vec![ - ( - "formerly_blocked_tool", - Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") - ), - ("calendar_tool", Some("calendar")), - ] - ); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - std::fs::write(cache_path, b"{not json").expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_from_disk_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cached_tools = vec![create_test_tool( +fn codex_apps_env_bearer_token_bypasses_shared_tools_cache() { + assert!(!should_share_codex_apps_tools_cache( CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )]; - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache(Some(&cache_context), &server_info, &cached_tools); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("expected startup snapshot to load from cache"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, Some(server_info)); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("legacy startup snapshot should remain available"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, None); -} - -#[test] -fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache( - Some(&cache_context), - &server_info, - &[create_test_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )], - ); - - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write legacy tools cache"); - - assert_eq!( - load_startup_cached_codex_apps_server_info(Some(&cache_context)), - Some(server_info) - ); - assert!(load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)).is_none()); + /*uses_env_bearer_token*/ true, + )); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -798,8 +621,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -813,7 +637,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); } @@ -836,8 +660,9 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -858,10 +683,8 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { #[tokio::test] async fn list_all_tools_accepts_canonical_namespaced_tool_names() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -869,18 +692,7 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { &permission_profile, /*prefix_mcp_tool_names*/ false, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -902,10 +714,8 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -913,18 +723,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -945,7 +744,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { } #[tokio::test] -async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() { +async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -961,8 +760,9 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -998,8 +798,9 @@ async fn shutdown_cancels_pending_tool_listing() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, @@ -1043,8 +844,9 @@ async fn shutdown_continues_after_caller_is_aborted() { AsyncManagedClient { client: blocking_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1071,7 +873,14 @@ async fn shutdown_continues_after_caller_is_aborted() { } #[tokio::test] -async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() { +async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(Vec::new()); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -1087,8 +896,9 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1097,16 +907,22 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() let timeout_result = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; - let tools = timeout_result.expect("cache-hit startup snapshot should not block"); + let tools = timeout_result.expect("shared empty cache should not block"); assert!(tools.is_empty()); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let server_info = create_test_server_info("Codex Apps"); let failed_client = futures::future::ready::>(Err( StartupOutcomeError::Failed { @@ -1128,8 +944,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails AsyncManagedClient { client: failed_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1143,7 +960,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); assert_eq!( @@ -1156,12 +973,10 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails } #[tokio::test] -async fn list_all_tools_adds_server_metadata_to_cached_tools() { +async fn list_all_tools_adds_server_metadata_to_tools() { let server_name = "docs"; - let startup_tools = vec![create_test_tool(server_name, "search")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool(server_name, "search")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -1182,18 +997,9 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() { tool_approval_modes: HashMap::new(), }, ); - manager.clients.insert( - server_name.to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager + .clients + .insert(server_name.to_string(), managed_client); let tools = manager.list_all_tools().await; assert_eq!(tools.len(), 1); @@ -1340,6 +1146,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { PathBuf::from("/tmp"), ), codex_home.path().to_path_buf(), + CodexAppsToolsCache::default(), CodexAppsToolsCacheKey { account_id: None, chatgpt_user_id: None, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 98efdc1b56..e611623d44 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -35,8 +35,9 @@ pub use auth_elicitation::auth_elicitation_id; pub use auth_elicitation::build_auth_elicitation; pub use auth_elicitation::build_auth_elicitation_plan; pub use auth_elicitation::connector_auth_failure_from_tool_result; -pub use codex_apps::CodexAppsToolsCacheKey; -pub use codex_apps::codex_apps_tools_cache_key; +pub use codex_apps_cache::CodexAppsToolsCache; +pub use codex_apps_cache::CodexAppsToolsCacheKey; +pub use codex_apps_cache::codex_apps_tools_cache_key; pub use mcp::codex_apps_mcp_server_config; pub use mcp::configured_mcp_servers; pub use mcp::effective_mcp_servers; @@ -75,6 +76,7 @@ pub use tools::declared_openai_file_input_param_names; pub(crate) mod auth_elicitation; mod catalog; pub(crate) mod codex_apps; +pub(crate) mod codex_apps_cache; pub(crate) mod connection_manager; pub(crate) mod elicitation; pub(crate) mod mcp; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 4550afa799..7532c6466b 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -44,7 +44,8 @@ use serde_json::Value; use tokio_util::sync::CancellationToken; use crate::ResolvedMcpCatalog; -use crate::codex_apps::codex_apps_tools_cache_key; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::codex_apps_tools_cache_key; use crate::connection_manager::McpConnectionManager; use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; @@ -301,6 +302,7 @@ pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, server: &str, uri: &str, ) -> anyhow::Result { @@ -329,6 +331,7 @@ pub async fn read_mcp_resource( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), @@ -361,6 +364,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( auth: Option<&CodexAuth>, submit_id: String, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -403,6 +407,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 1e8f5a331d..73ad33df7f 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -2,7 +2,7 @@ //! //! This module owns startup of individual RMCP clients: building the transport, //! initializing the server, listing raw tools, applying per-server tool filters, -//! and exposing cached startup snapshots while a client is still connecting. +//! and exposing cached Codex Apps tools while a client is still connecting. //! Higher-level aggregation and resource/tool APIs live in //! [`crate::connection_manager`]. @@ -17,15 +17,12 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; -use crate::codex_apps::CachedCodexAppsToolsLoad; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_cached_codex_apps_tools; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; use crate::codex_apps::normalize_codex_apps_tool_title; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; +use crate::codex_apps_cache::load_startup_cached_codex_apps_server_info; use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; @@ -105,9 +102,10 @@ pub(crate) struct ManagedClient { impl ManagedClient { fn listed_tools(&self) -> Vec { let total_start = Instant::now(); - if let Some(cache_context) = self.codex_apps_tools_cache_context.as_ref() - && let CachedCodexAppsToolsLoad::Hit(tools) = - load_cached_codex_apps_tools(cache_context) + if let Some(tools) = self + .codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, @@ -133,8 +131,9 @@ impl ManagedClient { pub(crate) struct AsyncManagedClient { pub(crate) client: Shared>>, pub(crate) is_codex_apps_mcp_server: bool, - pub(crate) cached_tool_info_snapshot: Option>, pub(crate) cached_server_info: Option, + pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) tool_plugin_provenance: Arc, pub(crate) cancel_token: CancellationToken, @@ -165,19 +164,15 @@ impl AsyncManagedClient { .configured_config() .map(ToolFilter::from_config) .unwrap_or_default(); - let (cached_tool_info_snapshot, cached_server_info) = if is_codex_apps_mcp_server { - ( - load_startup_cached_codex_apps_tools_snapshot( - codex_apps_tools_cache_context.as_ref(), - ), - load_startup_cached_codex_apps_server_info(codex_apps_tools_cache_context.as_ref()), - ) + let cached_server_info = if is_codex_apps_mcp_server { + codex_apps_tools_cache_context + .as_ref() + .and_then(load_startup_cached_codex_apps_server_info) } else { - (None, None) + None }; - let cached_tool_info_snapshot = - cached_tool_info_snapshot.map(|tools| filter_tools(tools, &tool_filter)); - let startup_tool_filter = tool_filter; + let startup_tool_filter = tool_filter.clone(); + let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone(); let startup_complete = Arc::new(AtomicBool::new(false)); let startup_complete_for_fut = Arc::clone(&startup_complete); let cancel_token_for_fut = cancel_token.clone(); @@ -214,7 +209,7 @@ impl AsyncManagedClient { tool_filter: startup_tool_filter, tx_event, elicitation_requests, - codex_apps_tools_cache_context, + codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -232,7 +227,10 @@ impl AsyncManagedClient { outcome }; let client = fut.in_current_span().boxed().shared(); - if cached_tool_info_snapshot.is_some() { + if codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + { let startup_task = client.clone(); tokio::spawn(async move { let _ = startup_task.await; @@ -242,8 +240,9 @@ impl AsyncManagedClient { Self { client, is_codex_apps_mcp_server, - cached_tool_info_snapshot, cached_server_info, + codex_apps_tools_cache_context, + tool_filter, startup_complete, tool_plugin_provenance, cancel_token, @@ -265,15 +264,29 @@ impl AsyncManagedClient { } } + pub(crate) fn has_cached_tools(&self) -> bool { + self.codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + } + + fn cached_tools(&self) -> Option> { + self.codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) + .map(|tools| filter_tools(tools, &self.tool_filter)) + } + pub(crate) async fn listed_tools(&self) -> Option> { // Keep cache payloads raw; plugin provenance is resolved per-session at read time. - let tools = if let Some(startup_tools) = self.cached_tool_info_snapshot_while_initializing() + let tools = if !self.startup_complete.load(Ordering::Acquire) + && let Some(startup_tools) = self.cached_tools() { Some(startup_tools) } else { match self.client().await { Ok(client) => Some(client.listed_tools()), - Err(_) => self.cached_tool_info_snapshot.clone(), + Err(_) => self.cached_tools(), } }?; Some(if self.is_codex_apps_mcp_server { @@ -282,13 +295,6 @@ impl AsyncManagedClient { prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) }) } - - fn cached_tool_info_snapshot_while_initializing(&self) -> Option> { - if !self.startup_complete.load(Ordering::Acquire) { - return self.cached_tool_info_snapshot.clone(); - } - None - } } #[derive(Debug, Clone, thiserror::Error)] @@ -570,6 +576,9 @@ async fn start_server_task( .is_some(); let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); let tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, @@ -585,21 +594,20 @@ async fn start_server_task( &[], ); let server_info = mcp_server_info_from_implementation(initialize_result.server_info); - let codex_apps_tools_cache_context = if is_codex_apps_mcp_server { - write_codex_apps_tools_cache( - codex_apps_tools_cache_context.as_ref(), - &server_info, - &tools, - ); + let tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { + (Some(cache_context), Some(fetch_ticket)) => { + cache_context.publish_if_newest_accepted(fetch_ticket, &server_info, tools) + } + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; + if is_codex_apps_mcp_server { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), &[("cache", "miss")], ); - codex_apps_tools_cache_context - } else { - None - }; + } let tools = filter_tools(tools, &tool_filter); let managed = ManagedClient { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 76b4a84830..97f24e043f 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -275,6 +275,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( // one exists, but do not reintroduce the old hidden-local fallback. runtime_context, config.codex_home.to_path_buf(), + mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 90e6b01e75..658c5471ee 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -10,6 +10,7 @@ use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::CodexAppsToolsCache; use codex_mcp::EffectiveMcpServer; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; @@ -38,14 +39,15 @@ enum OrderedMcpOverlay { pub struct McpManager { plugins_manager: Arc, extensions: Arc>, + codex_apps_tools_cache: CodexAppsToolsCache, } impl McpManager { pub fn new(plugins_manager: Arc) -> Self { - Self { + Self::new_with_extensions( plugins_manager, - extensions: codex_extension_api::empty_extension_registry(), - } + codex_extension_api::empty_extension_registry(), + ) } /// Creates a manager that resolves host-installed MCP contributions. @@ -56,9 +58,14 @@ impl McpManager { Self { plugins_manager, extensions, + codex_apps_tools_cache: CodexAppsToolsCache::default(), } } + pub fn codex_apps_tools_cache(&self) -> CodexAppsToolsCache { + self.codex_apps_tools_cache.clone() + } + /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index df2a375e35..a270af7851 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -1471,9 +1471,15 @@ pub(crate) async fn lookup_mcp_tool_metadata( { Some(connectors) => Some(connectors), None => { - connectors::list_accessible_connectors_from_mcp_tools(turn_context.config.as_ref()) - .await - .ok() + connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + turn_context.config.as_ref(), + /*force_refetch*/ false, + sess.services.turn_environments.environment_manager(), + Arc::clone(&sess.services.mcp_manager), + ) + .await + .ok() + .map(|status| status.connectors) } }; connectors.and_then(|connectors| { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 387987c9f6..41701d557c 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1400,6 +1400,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: }, ), turn_context.config.codex_home.to_path_buf(), + session.services.mcp_manager.codex_apps_tools_cache(), codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c060ef090b..249a6994b1 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -311,6 +311,7 @@ impl Session { turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + self.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 96dffed888..57a66b9986 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1193,6 +1193,7 @@ impl Session { session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + sess.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names(), client_elicitation_capability, diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index c6cff4ed30..752bf27037 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -8,6 +8,7 @@ use tokio::process::Command; use tokio::time::timeout; use tokio_tungstenite::connect_async; use tokio_tungstenite::connect_async_with_config; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::warn; @@ -30,6 +31,7 @@ use crate::noise_relay::NoiseHarnessConnectionArgs; use crate::noise_relay::noise_harness_connection_from_websocket; use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; +use crate::trace_context::current_trace_context_headers; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; @@ -216,6 +218,14 @@ impl ExecServerClient { /// only ciphertext after that. Environment-managed connections use a /// retained [`NoiseRendezvousConnectProvider`] so recovery can fetch a fresh /// bundle for each reconnect. + #[tracing::instrument( + name = "codex.exec_server.remote.harness.connect", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.connect", + ) + )] pub async fn connect_noise_rendezvous( args: NoiseRendezvousConnectArgs, ) -> Result { @@ -249,10 +259,20 @@ impl ExecServerClient { .next() .unwrap_or(websocket_url.as_str()) .to_string(); + let mut request = websocket_url + .as_str() + .into_client_request() + .map_err(|source| ExecServerError::WebSocketConnect { + url: diagnostic_url.clone(), + source, + })?; + request + .headers_mut() + .extend(current_trace_context_headers()); let (stream, _) = timeout( connect_timeout, connect_async_with_config( - websocket_url.as_str(), + request, Some(noise_relay_websocket_config()), /*disable_nagle*/ false, ), diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 87593374cb..0068c2cc1a 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -28,6 +28,7 @@ mod runtime_paths; mod sandboxed_file_system; mod server; mod telemetry; +mod trace_context; use codex_exec_server_protocol as protocol; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 6609b6cd92..f5e233ce9d 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -1,5 +1,6 @@ use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; @@ -10,7 +11,10 @@ use http::HeaderValue; use reqwest::StatusCode; use serde::Deserialize; use tokio::time::sleep; +use tokio_tungstenite::MaybeTlsStream; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async_with_config; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::info; use tracing::warn; @@ -35,6 +39,7 @@ use crate::noise_relay::noise_relay_websocket_config; use crate::relay::HarnessKeyValidator; use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; +use crate::trace_context::current_trace_context_headers; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; @@ -45,6 +50,7 @@ struct EnvironmentRegistryClient { auth_provider: SharedAuthProvider, http: reqwest::Client, connect_timeout: Duration, + telemetry: ExecServerTelemetry, } impl std::fmt::Debug for EnvironmentRegistryClient { @@ -57,7 +63,16 @@ impl std::fmt::Debug for EnvironmentRegistryClient { } impl EnvironmentRegistryClient { + #[cfg(test)] fn new(base_url: String, auth_provider: SharedAuthProvider) -> Result { + Self::new_with_telemetry(base_url, auth_provider, ExecServerTelemetry::default()) + } + + fn new_with_telemetry( + base_url: String, + auth_provider: SharedAuthProvider, + telemetry: ExecServerTelemetry, + ) -> Result { let base_url = normalize_base_url(base_url)?; Ok(Self { base_url, @@ -66,15 +81,41 @@ impl EnvironmentRegistryClient { .redirect(reqwest::redirect::Policy::none()) .build()?, connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + telemetry, }) } /// Register the executor public key and obtain the rendezvous allocation. /// The returned registration ID is included in each stream's Noise prologue. + #[tracing::instrument( + name = "codex.exec_server.remote.register", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.register", + result = tracing::field::Empty, + ) + )] async fn register_environment( &self, environment_id: &str, executor_public_key: &NoiseChannelPublicKey, + ) -> Result { + let started_at = Instant::now(); + let response = self + .register_environment_inner(environment_id, executor_public_key) + .await; + let result = if response.is_ok() { "success" } else { "error" }; + tracing::Span::current().record("result", result); + self.telemetry + .remote_registration_completed(result, started_at.elapsed()); + response + } + + async fn register_environment_inner( + &self, + environment_id: &str, + executor_public_key: &NoiseChannelPublicKey, ) -> Result { let response = self .http @@ -83,6 +124,7 @@ impl EnvironmentRegistryClient { &format!("/cloud/environment/{environment_id}/register"), )) .headers(self.auth_provider.to_auth_headers()) + .headers(current_trace_context_headers()) .json(&EnvironmentRegistryRegistrationRequest { security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), executor_public_key: executor_public_key.clone(), @@ -269,7 +311,11 @@ impl NoiseRendezvousEnvironmentConfig { ) -> Result { let environment_id = normalize_environment_id(environment_id)?; let auth_provider = static_bearer_auth_provider(bearer_token, chatgpt_account_id)?; - let client = EnvironmentRegistryClient::new(base_url, auth_provider)?; + let client = EnvironmentRegistryClient::new_with_telemetry( + base_url, + auth_provider, + ExecServerTelemetry::default(), + )?; Ok(Self { provider: Arc::new(EnvironmentRegistryNoiseConnectProvider { client, @@ -416,18 +462,16 @@ impl RemoteEnvironmentConfig { /// reconnects. The registration and rendezvous URL are also reused until /// rendezvous rejects the URL, at which point the next attempt registers again. /// The websocket carries cleartext routing metadata and encrypted payloads. -#[tracing::instrument( - name = "codex.exec_server", - skip_all, - fields(otel.kind = "internal") -)] pub async fn run_remote_environment( config: RemoteEnvironmentConfig, runtime_paths: ExecServerRuntimePaths, ) -> Result<(), ExecServerError> { ensure_rustls_crypto_provider(); - let client = - EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; + let client = EnvironmentRegistryClient::new_with_telemetry( + config.base_url.clone(), + config.auth_provider.clone(), + config.telemetry.clone(), + )?; let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, config.telemetry.clone()); let identity = NoiseChannelIdentity::generate().map_err(|error| { @@ -439,14 +483,8 @@ pub async fn run_remote_environment( .await?; loop { - match connect_async_with_config( - response.url.as_str(), - Some(noise_relay_websocket_config()), - /*disable_nagle*/ false, - ) - .await - { - Ok((websocket, _)) => { + match connect_rendezvous(&response.url, &config.telemetry).await { + Ok(websocket) => { backoff = Duration::from_secs(1); let executor_registration_id = response.executor_registration_id.clone(); info!( @@ -467,6 +505,7 @@ pub async fn run_remote_environment( }, ) .await; + config.telemetry.remote_reconnect("disconnected"); } Err(error) => { let registration_rejected = matches!( @@ -482,9 +521,12 @@ pub async fn run_remote_environment( ); debug!(error = %error, "Noise executor rendezvous connection error"); if registration_rejected { + config.telemetry.remote_reconnect("registration_rejected"); response = client .register_environment(&config.environment_id, &identity.public_key()) .await?; + } else { + config.telemetry.remote_reconnect("connect_failed"); } } } @@ -494,6 +536,43 @@ pub async fn run_remote_environment( } } +#[tracing::instrument( + name = "codex.exec_server.remote.rendezvous.connect", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.rendezvous.connect", + result = tracing::field::Empty, + ) +)] +async fn connect_rendezvous( + url: &str, + telemetry: &ExecServerTelemetry, +) -> Result< + WebSocketStream>, + tokio_tungstenite::tungstenite::Error, +> { + let started_at = Instant::now(); + let result = async { + let mut request = url.into_client_request()?; + request + .headers_mut() + .extend(current_trace_context_headers()); + connect_async_with_config( + request, + Some(noise_relay_websocket_config()), + /*disable_nagle*/ false, + ) + .await + .map(|(websocket, _)| websocket) + } + .await; + let result_name = if result.is_ok() { "success" } else { "error" }; + tracing::Span::current().record("result", result_name); + telemetry.remote_rendezvous_completed(result_name, started_at.elapsed()); + result +} + fn normalize_environment_id(environment_id: String) -> Result { let environment_id = environment_id.trim().to_string(); if environment_id.is_empty() { @@ -585,12 +664,17 @@ mod tests { use codex_api::AuthProvider; use http::HeaderMap; use http::HeaderValue; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; + use tracing::Instrument; + use tracing_subscriber::prelude::*; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::body_partial_json; use wiremock::matchers::header; + use wiremock::matchers::header_regex; use wiremock::matchers::method; use wiremock::matchers::path; @@ -616,8 +700,14 @@ mod tests { Arc::new(StaticRegistryAuthProvider) } - #[tokio::test] + #[tokio::test(flavor = "current_thread")] async fn register_environment_posts_with_auth_provider_headers() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); let server = MockServer::start().await; let executor_public_key = NoiseChannelIdentity::generate() .expect("identity") @@ -626,6 +716,10 @@ mod tests { .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) + .and(header_regex( + "traceparent", + "^00-[0-9a-f]{32}-[0-9a-f]{16}-0[01]$", + )) .and(body_partial_json(serde_json::json!({ "security_profile": NOISE_RELAY_SECURITY_PROFILE, "executor_public_key": executor_public_key.clone(), @@ -643,6 +737,7 @@ mod tests { let response = client .register_environment("environment-requested", &executor_public_key) + .instrument(tracing::info_span!("remote-operation")) .await .expect("register environment"); diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index 4624a2ff2c..650faacb2c 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/codex-rs/exec-server/src/telemetry.rs @@ -21,6 +21,28 @@ const PROCESSES_FINISHED_TOTAL_DESCRIPTION: &str = "Total number of finished exec-server processes."; const PROCESS_DURATION_METRIC: &str = "exec_server_process_duration_seconds"; const PROCESS_DURATION_DESCRIPTION: &str = "Duration of exec-server processes in seconds."; +const REMOTE_REGISTRATION_METRICS: OperationMetrics = OperationMetrics { + total_name: "exec_server_remote_registration_total", + total_description: "Total number of remote exec-server registration attempts.", + duration_name: "exec_server_remote_registration_duration_seconds", + duration_description: "Duration of remote exec-server registration attempts in seconds.", +}; +const REMOTE_RENDEZVOUS_METRICS: OperationMetrics = OperationMetrics { + total_name: "exec_server_remote_rendezvous_connect_total", + total_description: "Total number of remote exec-server rendezvous connection attempts.", + duration_name: "exec_server_remote_rendezvous_connect_duration_seconds", + duration_description: "Duration of remote exec-server rendezvous connection attempts in seconds.", +}; +const REMOTE_RECONNECTS_TOTAL_METRIC: &str = "exec_server_remote_reconnects_total"; +const REMOTE_RECONNECTS_TOTAL_DESCRIPTION: &str = "Total number of remote exec-server reconnects."; + +#[derive(Clone, Copy)] +struct OperationMetrics { + total_name: &'static str, + total_description: &'static str, + duration_name: &'static str, + duration_description: &'static str, +} #[derive(Clone, Copy)] pub(crate) enum ConnectionTransport { @@ -123,6 +145,24 @@ impl ExecServerTelemetry { }); } + pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) { + self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration); + } + + pub(crate) fn remote_rendezvous_completed(&self, result: &'static str, duration: Duration) { + self.record_operation(REMOTE_RENDEZVOUS_METRICS, result, duration); + } + + pub(crate) fn remote_reconnect(&self, reason: &'static str) { + self.with_inner(|inner| { + inner.counter( + REMOTE_RECONNECTS_TOTAL_METRIC, + REMOTE_RECONNECTS_TOTAL_DESCRIPTION, + &[("reason", reason)], + ); + }); + } + pub(crate) fn process_started(&self) -> ProcessMetricGuard { self.with_inner(|inner| { inner.adjust_process_count(/*delta*/ 1); @@ -162,6 +202,24 @@ impl ExecServerTelemetry { emit(inner); } } + + fn record_operation( + &self, + metrics: OperationMetrics, + result: &'static str, + duration: Duration, + ) { + self.with_inner(|inner| { + let tags = [("result", result)]; + inner.counter(metrics.total_name, metrics.total_description, &tags); + inner.duration( + metrics.duration_name, + metrics.duration_description, + duration, + &tags, + ); + }); + } } impl Drop for ConnectionMetricGuard { diff --git a/codex-rs/exec-server/src/trace_context.rs b/codex-rs/exec-server/src/trace_context.rs new file mode 100644 index 0000000000..70a8f3c61e --- /dev/null +++ b/codex-rs/exec-server/src/trace_context.rs @@ -0,0 +1,24 @@ +use reqwest::header::HeaderMap; +use reqwest::header::HeaderValue; + +pub(crate) fn current_trace_context_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + let Some(trace) = codex_otel::current_span_w3c_trace_context() else { + return headers; + }; + if let Some(traceparent) = trace.traceparent + && let Ok(value) = HeaderValue::try_from(traceparent) + { + headers.insert("traceparent", value); + } + if let Some(tracestate) = trace.tracestate + && let Ok(value) = HeaderValue::try_from(tracestate) + { + headers.insert("tracestate", value); + } + headers +} + +#[cfg(test)] +#[path = "trace_context_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/trace_context_tests.rs b/codex-rs/exec-server/src/trace_context_tests.rs new file mode 100644 index 0000000000..92a820afc4 --- /dev/null +++ b/codex-rs/exec-server/src/trace_context_tests.rs @@ -0,0 +1,27 @@ +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::SdkTracerProvider; +use tracing_subscriber::prelude::*; + +use super::current_trace_context_headers; + +#[test] +fn creates_traceparent_header_from_current_span() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); + let span = tracing::info_span!("outbound-request"); + let _entered = span.enter(); + + let headers = current_trace_context_headers(); + + let traceparent = headers + .get("traceparent") + .expect("traceparent header") + .to_str() + .expect("valid traceparent header"); + assert!(traceparent.starts_with("00-")); + assert_eq!(traceparent.len(), 55); +}