mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
Merge 658d2977df into sapling-pr-archive-bolinfest
This commit is contained in:
Michael Bolin
@@ -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<OutgoingMessageSender>,
|
||||
@@ -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(());
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<CodexAuth>,
|
||||
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<CodexAuth>,
|
||||
runtime_context: McpRuntimeContext,
|
||||
codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache,
|
||||
) -> Result<ListMcpServerStatusResponse, JSONRPCErrorError> {
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -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<ItemStartedNotification> {
|
||||
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<ItemCompletedNotification> {
|
||||
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"
|
||||
"#
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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<F, E>(run: F) -> Result<(), E>
|
||||
where
|
||||
F: Future<Output = Result<(), E>>,
|
||||
{
|
||||
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<ShutdownSignal> {
|
||||
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<ShutdownSignal> {
|
||||
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))
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<u8>,
|
||||
}
|
||||
|
||||
impl EncodedFrame {
|
||||
pub fn encode<T>(message: &T) -> io::Result<Self>
|
||||
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<R> {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String>,
|
||||
pub(crate) chatgpt_user_id: Option<String>,
|
||||
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<ToolInfo>),
|
||||
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<Vec<ToolInfo>> {
|
||||
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<McpServerInfo> {
|
||||
load_cached_codex_apps_server_info(cache_context?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn read_cached_codex_apps_tools(
|
||||
cache_context: &CodexAppsToolsCacheContext,
|
||||
) -> Option<Vec<ToolInfo>> {
|
||||
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<McpServerInfo> {
|
||||
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<ToolInfo>,
|
||||
}
|
||||
|
||||
#[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}")
|
||||
}
|
||||
|
||||
@@ -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<String>,
|
||||
pub(crate) chatgpt_user_id: Option<String>,
|
||||
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<Mutex<HashMap<CodexAppsToolsCacheIdentity, Arc<CodexAppsToolsCacheEntry>>>>,
|
||||
}
|
||||
|
||||
/// 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<CodexAppsToolsCacheEntry>,
|
||||
}
|
||||
|
||||
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<Vec<ToolInfo>> {
|
||||
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<ToolInfo>,
|
||||
) -> Vec<ToolInfo> {
|
||||
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<ToolInfo>) {
|
||||
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<Vec<ToolInfo>>,
|
||||
next_fetch_generation: AtomicU64,
|
||||
last_accepted_generation: Mutex<u64>,
|
||||
}
|
||||
|
||||
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<McpServerInfo> {
|
||||
load_cached_codex_apps_server_info(cache_context)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn read_cached_codex_apps_tools(
|
||||
cache_context: &CodexAppsToolsCacheContext,
|
||||
) -> Option<Vec<ToolInfo>> {
|
||||
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<Vec<ToolInfo>> {
|
||||
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<McpServerInfo> {
|
||||
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<u8>,
|
||||
) -> 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<ToolInfo>,
|
||||
}
|
||||
|
||||
#[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<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
||||
mutex
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "codex_apps_cache_tests.rs"]
|
||||
mod tests;
|
||||
@@ -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<ToolName> {
|
||||
tools
|
||||
.iter()
|
||||
.map(ToolInfo::canonical_tool_name)
|
||||
.collect::<HashSet<_>>()
|
||||
}
|
||||
|
||||
#[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<_>>(),
|
||||
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::<HashSet<_>>().len(),
|
||||
cache_paths.len()
|
||||
);
|
||||
}
|
||||
@@ -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<ToolInfo> {
|
||||
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<Vec<ToolInfo>> {
|
||||
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<CodexAppsToolsCacheContext> {
|
||||
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<CodexAppsToolsCacheContext> {
|
||||
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<Event>,
|
||||
|
||||
@@ -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<DuplexStream>> {
|
||||
async {
|
||||
let (client_stream, _server_stream) = tokio::io::duplex(1);
|
||||
Ok(client_stream)
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> 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::<Result<ManagedClient, StartupOutcomeError>>(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<ToolName> {
|
||||
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<_>>(),
|
||||
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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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::<Result<ManagedClient, StartupOutcomeError>>(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::<Result<ManagedClient, StartupOutcomeError>>()
|
||||
.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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<ReadResourceResult> {
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<ToolInfo> {
|
||||
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<BoxFuture<'static, Result<ManagedClient, StartupOutcomeError>>>,
|
||||
pub(crate) is_codex_apps_mcp_server: bool,
|
||||
pub(crate) cached_tool_info_snapshot: Option<Vec<ToolInfo>>,
|
||||
pub(crate) cached_server_info: Option<McpServerInfo>,
|
||||
pub(crate) codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
|
||||
pub(crate) tool_filter: ToolFilter,
|
||||
pub(crate) startup_complete: Arc<AtomicBool>,
|
||||
pub(crate) tool_plugin_provenance: Arc<ToolPluginProvenance>,
|
||||
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<Vec<ToolInfo>> {
|
||||
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<Vec<ToolInfo>> {
|
||||
// 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<Vec<ToolInfo>> {
|
||||
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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<PluginsManager>,
|
||||
extensions: Arc<ExtensionRegistry<Config>>,
|
||||
codex_apps_tools_cache: CodexAppsToolsCache,
|
||||
}
|
||||
|
||||
impl McpManager {
|
||||
pub fn new(plugins_manager: Arc<PluginsManager>) -> 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 {
|
||||
|
||||
@@ -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| {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Self, ExecServerError> {
|
||||
@@ -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,
|
||||
),
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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, ExecServerError> {
|
||||
Self::new_with_telemetry(base_url, auth_provider, ExecServerTelemetry::default())
|
||||
}
|
||||
|
||||
fn new_with_telemetry(
|
||||
base_url: String,
|
||||
auth_provider: SharedAuthProvider,
|
||||
telemetry: ExecServerTelemetry,
|
||||
) -> Result<Self, ExecServerError> {
|
||||
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<EnvironmentRegistryRegistrationResponse, ExecServerError> {
|
||||
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<EnvironmentRegistryRegistrationResponse, ExecServerError> {
|
||||
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<Self, ExecServerError> {
|
||||
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<MaybeTlsStream<tokio::net::TcpStream>>,
|
||||
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<String, ExecServerError> {
|
||||
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");
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
}
|
||||
Reference in New Issue
Block a user