From 5579792b3b881b44ba42139f9fcd191c4a944128 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 25 Jun 2026 02:02:07 -0700 Subject: [PATCH 01/33] Parallelize environment skill loading (#29990) ## Why Avoid a request waterfall for loading lots of skills at once by hiding latency in concurrent tasks. ## What changed Poll the per-skill parse futures concurrently with an order-preserving stream capped at 64 in-flight loads. Results retain discovery order, and the existing filtering, warnings, and final catalog sorting are unchanged. --- .../core-skills/src/loader/environment.rs | 42 ++++++++++++------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/codex-rs/core-skills/src/loader/environment.rs b/codex-rs/core-skills/src/loader/environment.rs index 5dc29ebfb2..93ab14e9aa 100644 --- a/codex-rs/core-skills/src/loader/environment.rs +++ b/codex-rs/core-skills/src/loader/environment.rs @@ -10,6 +10,7 @@ use codex_utils_path_uri::PathUri; use codex_utils_plugins::DISCOVERABLE_PLUGIN_MANIFEST_PATHS; use codex_utils_plugins::plugin_namespace_for_root_uri; use codex_utils_plugins::plugin_namespace_for_skill_uri; +use futures::StreamExt; use futures::future::join_all; use crate::model::SkillDependencies; @@ -31,6 +32,7 @@ use super::sanitize_single_line; use super::validate_len; const MAX_SKILLS_ENTRIES_PER_ROOT: usize = 20_000; +const MAX_CONCURRENT_SKILL_LOADS: usize = 64; /// URI-native metadata for one skill owned by an execution environment. #[derive(Clone, Debug, PartialEq, Eq)] @@ -220,24 +222,32 @@ pub async fn load_environment_skills_from_root( .filter_map(|(plugin_root, namespace)| namespace.map(|namespace| (plugin_root, namespace))) .collect::>(); - for path in discovery.skill_files { - let mut ancestor = path.parent(); - let plugin_namespace = loop { - let Some(current) = ancestor else { - break None; + // Remote executors can multiplex these independent per-skill reads, so polling a bounded + // number together allows the I/O for each skill and its metadata to happen concurrently. + let skill_results = futures::stream::iter(discovery.skill_files) + .map(|path| { + let mut ancestor = path.parent(); + let plugin_namespace = loop { + let Some(current) = ancestor else { + break None; + }; + if let Some(namespace) = plugin_namespaces.get(¤t) { + break Some(namespace.as_str()); + } + ancestor = current.parent(); }; - if let Some(namespace) = plugin_namespaces.get(¤t) { - break Some(namespace.as_str()); + async move { + let result = + EnvironmentSkillMetadata::parse(file_system, &path, plugin_namespace).await; + (path, result) } - ancestor = current.parent(); - }; - match EnvironmentSkillMetadata::parse( - file_system, - &path, - /*plugin_namespace*/ plugin_namespace, - ) - .await - { + }) + .buffered(MAX_CONCURRENT_SKILL_LOADS) + .collect::>() + .await; + + for (path, result) in skill_results { + match result { Ok(skill) if skill.matches_product_restriction(restriction_product) => { outcome.skills.push(skill); } From 6368937939dceb07b4a3c47c4448027d0d1a85a6 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 25 Jun 2026 10:10:36 +0100 Subject: [PATCH 02/33] Support HTTP MCP servers from selected executor plugins (#28522) ## Why Selected executor plugins can declare both stdio and Streamable HTTP MCP servers, but only stdio registrations were retained. That silently drops part of the plugin's tool surface and prevents HTTP traffic from using the owning executor's network. ## What changed - retain selected-plugin Streamable HTTP MCP declarations alongside stdio declarations - route their HTTP clients through the owning executor environment - preserve local auth-header environment references while rejecting them for executor-hosted declarations - cover thread isolation, refresh, and an executor-only HTTP route end to end --- codex-rs/app-server/README.md | 2 +- .../app-server/tests/suite/v2/executor_mcp.rs | 143 ++++++++++++++++-- codex-rs/codex-mcp/src/plugin_config.rs | 31 +++- codex-rs/codex-mcp/src/plugin_config_tests.rs | 92 +++++++++++ .../ext/mcp/src/executor_plugin/provider.rs | 19 +-- .../mcp/src/executor_plugin/provider_tests.rs | 80 ++++++---- 6 files changed, 308 insertions(+), 59 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 246b2bb0c0..981e8556dd 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -137,7 +137,7 @@ Example with notification opt-out: ## API Overview -- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are also started in that environment; HTTP MCP declarations remain inactive. +- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are started in that environment, and HTTP MCP connections use that environment's HTTP client. - `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`. - `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. - `thread/start`, `thread/resume`, and `thread/fork` responses include the legacy `sandbox` compatibility projection. `instructionSources` lists loaded instruction files using each source environment's native absolute path syntax, including files loaded from remote environments. Experimental clients can read `runtimeWorkspaceRoots` for the thread-scoped runtime roots and `activePermissionProfile` for the named or implicit built-in profile identity/provenance when known. Their deprecated experimental `multiAgentMode` field, and the corresponding thread setting, always report `explicitRequestOnly`; Ultra reasoning effort is the source of proactive multi-agent behavior. diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 299ca81aa5..55e5f2dd4b 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -2,6 +2,7 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml; +use axum::Router; use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; @@ -17,13 +18,32 @@ use codex_utils_path_uri::PathUri; use core_test_support::responses; use core_test_support::stdio_server_bin; use pretty_assertions::assert_eq; +use rmcp::handler::server::ServerHandler; +use rmcp::model::CallToolRequestParams; +use rmcp::model::CallToolResult; +use rmcp::model::JsonObject; +use rmcp::model::ListToolsResult; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerInfo; +use rmcp::model::Tool; +use rmcp::model::ToolAnnotations; +use rmcp::service::RequestContext; +use rmcp::service::RoleServer; +use rmcp::transport::StreamableHttpServerConfig; +use rmcp::transport::StreamableHttpService; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use serde_json::json; +use std::borrow::Cow; use std::collections::BTreeMap; +use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; +use tokio::net::TcpListener; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(20); +const EXECUTOR_HTTP_MCP_URL: &str = "http://executor-only.invalid/mcp"; +const HTTP_MCP_SERVER_NAME: &str = "executor_http"; const MCP_SERVER_NAME: &str = "executor_demo"; const EXECUTOR_ENV_NAME: &str = "MCP_EXECUTOR_MARKER"; const EXECUTOR_ENV_VALUE: &str = "executor-only"; @@ -32,8 +52,19 @@ const REFRESH_PROBE_SERVER_NAME: &str = "refresh_probe"; const TOOL_CALL_ID: &str = "executor-mcp-call"; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn selected_executor_plugin_exposes_its_stdio_mcp_only_to_that_thread() -> Result<()> { +async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Result<()> { let responses_server = responses::start_mock_server().await; + let http_listener = TcpListener::bind("127.0.0.1:0").await?; + let http_addr = http_listener.local_addr()?; + let http_mcp_service = StreamableHttpService::new( + || Ok(ExecutorHttpMcpServer), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default().with_allowed_hosts(["executor-only.invalid"]), + ); + let http_router = Router::new().nest_service("/mcp", http_mcp_service); + let http_server_handle = tokio::spawn(async move { + let _ = axum::serve(http_listener, http_router).await; + }); let codex_home = TempDir::new()?; write_mock_responses_config_toml( codex_home.path(), @@ -44,6 +75,12 @@ async fn selected_executor_plugin_exposes_its_stdio_mcp_only_to_that_thread() -> "mock_provider", "compact", )?; + let codex_bin = toml::Value::String( + codex_utils_cargo_bin::cargo_bin("codex")? + .to_string_lossy() + .into_owned(), + ); + let http_proxy = toml::Value::String(format!("http://{http_addr}")); std::fs::write( codex_home.path().join("environments.toml"), format!( @@ -52,16 +89,12 @@ include_local = true [[environments]] id = "{EXECUTOR_ID}" -program = {} +program = {codex_bin} args = ["exec-server", "--listen", "stdio"] [environments.env] {EXECUTOR_ENV_NAME} = "{EXECUTOR_ENV_VALUE}" -"#, - toml::Value::String( - codex_utils_cargo_bin::cargo_bin("codex")? - .to_string_lossy() - .into_owned() - ) +HTTP_PROXY = {http_proxy} +"# ), )?; @@ -79,6 +112,11 @@ args = ["exec-server", "--listen", "stdio"] "command": stdio_server_bin()?, "env_vars": [EXECUTOR_ENV_NAME], "startup_timeout_sec": 10, + }, + (HTTP_MCP_SERVER_NAME): { + "url": EXECUTOR_HTTP_MCP_URL, + "environment_id": "local", + "startup_timeout_sec": 10, } } }))?, @@ -178,6 +216,26 @@ startup_timeout_sec = 10 assert!(output.contains("ECHOING: hello from executor")); assert!(output.contains(EXECUTOR_ENV_VALUE)); + let request_id = app_server + .send_mcp_server_tool_call_request(McpServerToolCallParams { + thread_id: selected_thread.clone(), + server: HTTP_MCP_SERVER_NAME.to_string(), + tool: "echo".to_string(), + arguments: Some(json!({"message": "hello over executor HTTP"})), + meta: None, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: McpServerToolCallResponse = to_response(response)?; + assert_eq!( + response.structured_content, + Some(json!({"echo": "ECHOING: hello over executor HTTP"})) + ); + let request_id = app_server .send_mcp_server_tool_call_request(McpServerToolCallParams { thread_id: selected_thread.clone(), @@ -200,25 +258,84 @@ startup_timeout_sec = 10 Some(json!("ECHOING: refresh applied")) ); + let selected_server_names = mcp_server_names(&mut app_server, selected_thread).await?; assert!( - mcp_server_names(&mut app_server, selected_thread) - .await? + selected_server_names .iter() .any(|name| name == MCP_SERVER_NAME) ); + assert!( + selected_server_names + .iter() + .any(|name| name == HTTP_MCP_SERVER_NAME) + ); let unselected_thread = start_thread(&mut app_server, /*selected_capability_roots*/ None).await?; + let unselected_server_names = mcp_server_names(&mut app_server, unselected_thread).await?; assert!( - mcp_server_names(&mut app_server, unselected_thread) - .await? + unselected_server_names .iter() - .all(|name| name != MCP_SERVER_NAME) + .all(|name| { name != MCP_SERVER_NAME && name != HTTP_MCP_SERVER_NAME }) ); + http_server_handle.abort(); + let _ = http_server_handle.await; + Ok(()) } +#[derive(Clone, Copy)] +struct ExecutorHttpMcpServer; + +impl ServerHandler for ExecutorHttpMcpServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let input_schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "properties": {"message": {"type": "string"}}, + "required": ["message"], + "additionalProperties": false + })) + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; + let mut tool = Tool::new( + Cow::Borrowed("echo"), + Cow::Borrowed("Echo a message."), + Arc::new(input_schema), + ); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + + Ok(ListToolsResult { + tools: vec![tool], + next_cursor: None, + meta: None, + }) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + let message = request + .arguments + .as_ref() + .and_then(|arguments| arguments.get("message")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + Ok(CallToolResult::structured(json!({ + "echo": format!("ECHOING: {message}") + }))) + } +} + async fn mcp_server_names( app_server: &mut TestAppServer, thread_id: String, diff --git a/codex-rs/codex-mcp/src/plugin_config.rs b/codex-rs/codex-mcp/src/plugin_config.rs index 968d6a64df..4f74a97d16 100644 --- a/codex-rs/codex-mcp/src/plugin_config.rs +++ b/codex-rs/codex-mcp/src/plugin_config.rs @@ -173,8 +173,35 @@ fn environment_cwd( fn bind_environment_env_vars(config: &mut McpServerConfig) -> Result<(), String> { let is_local_environment = config.is_local_environment(); - let McpServerTransportConfig::Stdio { env_vars, .. } = &mut config.transport else { - return Ok(()); + let env_vars = match &mut config.transport { + McpServerTransportConfig::Stdio { env_vars, .. } => env_vars, + // Never resolve executor-owned environment references in the host process. + // Remove this rejection once the owning executor resolves these fields. + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + env_http_headers, + .. + } => { + if is_local_environment { + return Ok(()); + } + if bearer_token_env_var.is_some() { + return Err( + "`bearer_token_env_var` requires executor-side environment resolution for an executor-owned HTTP MCP" + .to_string(), + ); + } + if env_http_headers + .as_ref() + .is_some_and(|headers| !headers.is_empty()) + { + return Err( + "`env_http_headers` requires executor-side environment resolution for an executor-owned HTTP MCP" + .to_string(), + ); + } + return Ok(()); + } }; for env_var in env_vars { match env_var { diff --git a/codex-rs/codex-mcp/src/plugin_config_tests.rs b/codex-rs/codex-mcp/src/plugin_config_tests.rs index 576bd29618..c2f1f71a54 100644 --- a/codex-rs/codex-mcp/src/plugin_config_tests.rs +++ b/codex-rs/codex-mcp/src/plugin_config_tests.rs @@ -275,6 +275,98 @@ fn environment_placement_rejects_orchestrator_env_vars() { ); } +#[test] +fn remote_environment_placement_rejects_http_env_references() { + let plugin_root = plugin_root(); + let outcome = parse_executor_plugin_mcp_config( + &plugin_root_uri(&plugin_root), + r#"{ + "bearer": { + "url": "https://example.com/bearer", + "bearer_token_env_var": "TOKEN" + }, + "headers": { + "url": "https://example.com/headers", + "env_http_headers": {"Authorization": "TOKEN"} + } + }"#, + "executor-1", + ) + .expect("parse plugin MCP config"); + + assert_eq!( + outcome, + PluginMcpConfigParseOutcome { + servers: BTreeMap::new(), + errors: vec![ + PluginMcpServerParseError { + name: "bearer".to_string(), + message: "`bearer_token_env_var` requires executor-side environment resolution for an executor-owned HTTP MCP" + .to_string(), + }, + PluginMcpServerParseError { + name: "headers".to_string(), + message: "`env_http_headers` requires executor-side environment resolution for an executor-owned HTTP MCP" + .to_string(), + }, + ], + } + ); +} + +#[test] +fn local_environment_placement_preserves_http_env_references() { + let plugin_root = plugin_root(); + let outcome = parse_executor_plugin_mcp_config( + &plugin_root_uri(&plugin_root), + r#"{ + "demo": { + "url": "https://example.com/mcp", + "bearer_token_env_var": "TOKEN", + "env_http_headers": {"X-Account": "ACCOUNT_ID"} + } + }"#, + DEFAULT_MCP_SERVER_ENVIRONMENT_ID, + ) + .expect("parse plugin MCP config"); + + assert_eq!( + outcome, + PluginMcpConfigParseOutcome { + servers: BTreeMap::from([( + "demo".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::StreamableHttp { + url: "https://example.com/mcp".to_string(), + bearer_token_env_var: Some("TOKEN".to_string()), + http_headers: None, + env_http_headers: Some(HashMap::from([( + "X-Account".to_string(), + "ACCOUNT_ID".to_string(), + )])), + }, + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + )]), + errors: Vec::new(), + } + ); +} + #[test] fn local_environment_placement_preserves_local_env_vars() { let plugin_root = plugin_root(); diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider.rs b/codex-rs/ext/mcp/src/executor_plugin/provider.rs index 9b8a9878e3..a6972df096 100644 --- a/codex-rs/ext/mcp/src/executor_plugin/provider.rs +++ b/codex-rs/ext/mcp/src/executor_plugin/provider.rs @@ -1,5 +1,4 @@ use codex_config::McpServerConfig; -use codex_config::McpServerTransportConfig; use codex_core_plugins::ResolvedExecutorPlugin; use codex_exec_server::ExecutorFileSystem; use codex_mcp::parse_executor_plugin_mcp_config; @@ -48,7 +47,7 @@ pub(super) enum ExecutorPluginMcpProviderError { } impl ExecutorPluginMcpProvider { - /// Returns stdio servers declared by `plugin`, bound to its environment. + /// Returns MCP servers declared by `plugin`, bound to its environment. pub(super) async fn load( &self, plugin: &ResolvedExecutorPlugin, @@ -131,21 +130,7 @@ async fn load_from_file_system( ); } - Ok(parsed - .servers - .into_iter() - .filter_map(|(name, config)| match &config.transport { - McpServerTransportConfig::Stdio { .. } => Some((name, config)), - McpServerTransportConfig::StreamableHttp { .. } => { - tracing::warn!( - plugin = plugin_id, - server = name, - "ignoring HTTP MCP server from executor plugin" - ); - None - } - }) - .collect()) + Ok(parsed.servers.into_iter().collect()) } #[cfg(test)] diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs index 22e5d1bec5..bc59e104a0 100644 --- a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs +++ b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs @@ -163,33 +163,61 @@ async fn reads_declared_config_only_through_executor_file_system() { assert_eq!( servers, - vec![( - "demo".to_string(), - McpServerConfig { - auth: Default::default(), - transport: McpServerTransportConfig::Stdio { - command: "demo-mcp".to_string(), - args: Vec::new(), - env: None, - env_vars: Vec::new(), - cwd: Some(LegacyAppPathString::from_path(plugin_root.as_path())), + vec![ + ( + "demo".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::Stdio { + command: "demo-mcp".to_string(), + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: Some(LegacyAppPathString::from_path(plugin_root.as_path())), + }, + environment_id: "executor-test".to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), }, - environment_id: "executor-test".to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - disabled_reason: None, - startup_timeout_sec: None, - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }, - )] + ), + ( + "hosted".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::StreamableHttp { + url: "https://example.com/mcp".to_string(), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + environment_id: "executor-test".to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + ), + ] ); assert_eq!(reads(&file_system), vec![config_path]); } From b215961a56b2553a4612a22a812747403c08a58b Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 25 Jun 2026 10:31:17 +0100 Subject: [PATCH 03/33] Support OAuth for HTTP MCP servers from selected executor plugins (#28529) ## Why #28522 routes selected-plugin HTTP MCP traffic through the owning executor, but OAuth bootstrap and refresh still used host-local clients. Executor-only servers therefore cannot complete discovery or login through the same network boundary as the MCP connection. ## What changed - adapt `codex_exec_server::HttpClient` to RMCP 1.8's `OAuthHttpClient` contract - let RMCP own discovery, dynamic registration, PKCE, token exchange, and refresh - route auth status, persisted-token startup, and app-server login through the server runtime while preserving the existing local discovery path - add optional `threadId` to `mcpServer/oauth/login` and echo it in the completion notification - implement RMCP's redirect policy and 1 MiB OAuth response limit over executor HTTP - cover selected-thread OAuth discovery and login through an executor-only route Depends on #28522. --- codex-rs/Cargo.lock | 1 - .../schema/json/ClientRequest.json | 6 + .../schema/json/ServerNotification.json | 6 + .../codex_app_server_protocol.schemas.json | 12 ++ .../codex_app_server_protocol.v2.schemas.json | 12 ++ ...ServerOauthLoginCompletedNotification.json | 6 + .../json/v2/McpServerOauthLoginParams.json | 6 + ...cpServerOauthLoginCompletedNotification.ts | 2 +- .../v2/McpServerOauthLoginParams.ts | 2 +- .../src/protocol/common.rs | 1 + .../src/protocol/v2/mcp.rs | 3 + codex-rs/app-server/README.md | 4 +- codex-rs/app-server/src/request_processors.rs | 4 +- .../src/request_processors/mcp_processor.rs | 48 +++++-- .../src/request_processors/plugins.rs | 1 + .../app-server/tests/suite/v2/executor_mcp.rs | 90 ++++++++++++- codex-rs/cli/src/mcp_cmd.rs | 6 + codex-rs/codex-mcp/src/lib.rs | 2 + codex-rs/codex-mcp/src/mcp/auth.rs | 126 ++++++++++++++---- codex-rs/codex-mcp/src/mcp/mod.rs | 4 + codex-rs/codex-mcp/src/runtime.rs | 16 +++ codex-rs/core/src/connectors.rs | 5 +- codex-rs/core/src/session/mcp.rs | 15 ++- codex-rs/core/src/session/session.rs | 5 + codex-rs/core/tests/suite/rmcp_client.rs | 2 + codex-rs/exec-server-protocol/src/protocol.rs | 14 ++ .../src/client/reqwest_http_client.rs | 21 ++- codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/server/handler.rs | 2 +- codex-rs/exec-server/tests/http_client.rs | 25 ++++ codex-rs/exec-server/tests/http_request.rs | 62 +++++++++ codex-rs/rmcp-client/Cargo.toml | 1 - codex-rs/rmcp-client/src/auth_status.rs | 124 ++++++++++++++++- .../rmcp-client/src/http_client_adapter.rs | 4 + codex-rs/rmcp-client/src/lib.rs | 4 + codex-rs/rmcp-client/src/oauth_http_client.rs | 116 ++++++++++++++++ .../rmcp-client/src/perform_oauth_login.rs | 83 +++++++++--- codex-rs/rmcp-client/src/rmcp_client.rs | 17 +-- 38 files changed, 763 insertions(+), 96 deletions(-) create mode 100644 codex-rs/rmcp-client/src/oauth_http_client.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2c50f268f0..bc5953803e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3757,7 +3757,6 @@ dependencies = [ "base64 0.22.1", "bytes", "codex-api", - "codex-client", "codex-config", "codex-exec-server", "codex-keyring-store", diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 68c242504f..0027988c42 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1619,6 +1619,12 @@ "null" ] }, + "threadId": { + "type": [ + "string", + "null" + ] + }, "timeoutSecs": { "format": "int64", "type": [ diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index a3d07c24a0..18435f6f2b 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -2370,6 +2370,12 @@ }, "success": { "type": "boolean" + }, + "threadId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 4d417ffb2f..33981b5b5f 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -12026,6 +12026,12 @@ }, "success": { "type": "boolean" + }, + "threadId": { + "type": [ + "string", + "null" + ] } }, "required": [ @@ -12050,6 +12056,12 @@ "null" ] }, + "threadId": { + "type": [ + "string", + "null" + ] + }, "timeoutSecs": { "format": "int64", "type": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 0df923a0f4..5a7cf80e83 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -8430,6 +8430,12 @@ }, "success": { "type": "boolean" + }, + "threadId": { + "type": [ + "string", + "null" + ] } }, "required": [ @@ -8454,6 +8460,12 @@ "null" ] }, + "threadId": { + "type": [ + "string", + "null" + ] + }, "timeoutSecs": { "format": "int64", "type": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json index 35efd2baf2..6204ff672b 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json @@ -12,6 +12,12 @@ }, "success": { "type": "boolean" + }, + "threadId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginParams.json b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginParams.json index 4370f444b9..de66dcefd6 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginParams.json @@ -13,6 +13,12 @@ "null" ] }, + "threadId": { + "type": [ + "string", + "null" + ] + }, "timeoutSecs": { "format": "int64", "type": [ diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts index 592860ae39..cfa6603094 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts @@ -2,4 +2,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type McpServerOauthLoginCompletedNotification = { name: string, success: boolean, error?: string, }; +export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null, success: boolean, error?: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginParams.ts index a61c304609..ff088b8d9f 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginParams.ts @@ -2,4 +2,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type McpServerOauthLoginParams = { name: string, scopes?: Array | null, timeoutSecs?: bigint | null, }; +export type McpServerOauthLoginParams = { name: string, threadId?: string | null, scopes?: Array | null, timeoutSecs?: bigint | null, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 51a37fbaa9..3120007e53 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1913,6 +1913,7 @@ mod tests { request_id: request_id(), params: v2::McpServerOauthLoginParams { name: "server-a".to_string(), + thread_id: None, scopes: None, timeout_secs: None, }, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs index cdfeb524f5..ca0e124d59 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs @@ -186,6 +186,8 @@ pub struct McpServerRefreshResponse {} #[ts(export_to = "v2/")] pub struct McpServerOauthLoginParams { pub name: String, + #[ts(optional = nullable)] + pub thread_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional = nullable)] pub scopes: Option>, @@ -215,6 +217,7 @@ pub struct McpToolCallProgressNotification { #[ts(export_to = "v2/")] pub struct McpServerOauthLoginCompletedNotification { pub name: String, + pub thread_id: Option, pub success: bool, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 981e8556dd..397fa3d58f 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -229,7 +229,7 @@ Example with notification opt-out: - `skills/config/write` — write user-level skill config by name or absolute path. - `plugin/install` — install a plugin from a discovered marketplace entry, rejecting marketplace entries marked unavailable for install, install MCPs if any, and return the effective plugin auth policy plus any apps that still need auth (**under development; do not call from production clients yet**). - `plugin/uninstall` — uninstall a local plugin by `pluginId` in `@` form by removing its cached files and clearing its user-level config entry, or uninstall a remote ChatGPT plugin by backend `pluginId` by forwarding the uninstall to the ChatGPT plugin backend and removing any downloaded remote-plugin cache (**under development; do not call from production clients yet**). -- `mcpServer/oauth/login` — start an OAuth login for a configured MCP server; returns an `authorization_url` and later emits `mcpServer/oauthLogin/completed` once the browser flow finishes. +- `mcpServer/oauth/login` — start an OAuth login for a configured MCP server; pass `threadId` to resolve servers from that thread's selected plugins and executor, and receive an `authorization_url` followed by `mcpServer/oauthLogin/completed` once the browser flow finishes. - `tool/requestUserInput` — prompt the user with 1–3 short questions for a tool call and return their answers (experimental). - `config/mcpServer/reload` — reload MCP server config from disk and queue a refresh for loaded threads (applied on each thread's next active turn); returns `{}`. Use this after editing `config.toml` without restarting the server. - `mcpServerStatus/list` — enumerate configured MCP servers with their tools, auth status, server info, plus resources/resource templates for `full` detail; supports optional `threadId` and cursor+limit pagination. If `threadId` is omitted, the server reads from the latest global config directly. If `detail` is omitted, the server defaults to `full`. @@ -1934,7 +1934,7 @@ Codex supports these authentication modes. The current mode is surfaced in `acco - `account/workspaceMessages/read` — fetch active workspace messages, including workspace notification headlines when available. - `account/rateLimits/updated` (notify) — emitted whenever a user's ChatGPT rate limits change. This is a sparse rolling update; merge available values into the most recent `account/rateLimits/read` response or refetch that snapshot. - `account/sendAddCreditsNudgeEmail` — ask ChatGPT to email the workspace owner about depleted credits or a reached usage limit. -- `mcpServer/oauthLogin/completed` (notify) — emitted after a `mcpServer/oauth/login` flow finishes for a server; payload includes `{ name, success, error? }`. +- `mcpServer/oauthLogin/completed` (notify) — emitted after a `mcpServer/oauth/login` flow finishes for a server; payload includes `{ name, threadId, success, error? }`. - `mcpServer/startupStatus/updated` (notify) — emitted when a configured MCP server's startup status changes; payload includes `{ threadId, name, status, error }`, where `threadId` is the owning thread when startup is thread-scoped and `null` when it is app-scoped, and `status` is `starting`, `ready`, `failed`, or `cancelled`. ### 1) Check auth state diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index d73a41706c..9b56957a53 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -380,7 +380,7 @@ use codex_mcp::McpRuntimeContext; use codex_mcp::McpServerStatusSnapshot; use codex_mcp::McpSnapshotDetail; use codex_mcp::collect_mcp_server_status_snapshot_with_detail; -use codex_mcp::discover_supported_scopes; +use codex_mcp::discover_supported_scopes_with_http_client; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; use codex_memories_write::clear_memory_roots_contents; @@ -428,7 +428,7 @@ use codex_protocol::protocol::USER_MESSAGE_BEGIN; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; use codex_protocol::user_input::UserInput as CoreInputItem; -use codex_rmcp_client::perform_oauth_login_return_url; +use codex_rmcp_client::perform_oauth_login_return_url_with_http_client; use codex_rollout::is_persisted_rollout_item; use codex_rollout::state_db::StateDbHandle; use codex_rollout::state_db::reconcile_rollout; diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 08d464598b..18b49ede3b 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -113,19 +113,37 @@ impl McpRequestProcessor { &self, params: McpServerOauthLoginParams, ) -> Result { - let config = self.load_latest_config(/*fallback_cwd*/ None).await?; let McpServerOauthLoginParams { name, + thread_id, scopes, timeout_secs, } = params; let auth = self.auth_manager.auth().await; - let effective_servers = self - .thread_manager - .mcp_manager() - .effective_servers(&config, auth.as_ref()) - .await; + let (config, mcp_config) = match thread_id.as_deref() { + Some(thread_id) => { + let (_, thread) = self.load_thread(thread_id).await?; + let thread_config = thread.config().await; + let config = self + .config_manager + .load_latest_config_for_thread(thread_config.as_ref()) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + let mcp_config = thread.runtime_mcp_config(&config).await; + (config, mcp_config) + } + None => { + let config = self.load_latest_config(/*fallback_cwd*/ None).await?; + let mcp_config = self + .thread_manager + .mcp_manager() + .runtime_config(&config) + .await; + (config, mcp_config) + } + }; + let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); let Some(server) = effective_servers .get(&name) .and_then(codex_mcp::EffectiveMcpServer::configured_config) @@ -149,15 +167,26 @@ impl McpRequestProcessor { } }; + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + let http_client = runtime_context + .resolve_http_client(&name, server) + .map_err(|err| { + internal_error(format!("failed to resolve MCP server runtime: {err}")) + })?; + let discovered_scopes = if scopes.is_none() && server.scopes.is_none() { - discover_supported_scopes(&server.transport).await + discover_supported_scopes_with_http_client(&server.transport, Arc::clone(&http_client)) + .await } else { None }; let resolved_scopes = resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes); - let handle = perform_oauth_login_return_url( + let handle = perform_oauth_login_return_url_with_http_client( &name, &url, config.mcp_oauth_credentials_store_mode, @@ -170,11 +199,13 @@ impl McpRequestProcessor { timeout_secs, config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), + http_client, ) .await .map_err(|err| internal_error(format!("failed to login to MCP server '{name}': {err}")))?; let authorization_url = handle.authorization_url().to_string(); let notification_name = name.clone(); + let notification_thread_id = thread_id; let outgoing = Arc::clone(&self.outgoing); tokio::spawn(async move { @@ -186,6 +217,7 @@ impl McpRequestProcessor { let notification = ServerNotification::McpServerOauthLoginCompleted( McpServerOauthLoginCompletedNotification { name: notification_name, + thread_id: notification_thread_id, success, error, }, diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 4aa8352d15..d123f0379c 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1883,6 +1883,7 @@ impl PluginRequestProcessor { let notification = ServerNotification::McpServerOauthLoginCompleted( McpServerOauthLoginCompletedNotification { name: notification_name, + thread_id: None, success, error, }, diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 55e5f2dd4b..f93af6821a 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -2,10 +2,14 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml; +use axum::Json; use axum::Router; +use axum::routing::get; use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; +use codex_app_server_protocol::McpServerOauthLoginCompletedNotification; +use codex_app_server_protocol::McpServerOauthLoginResponse; use codex_app_server_protocol::McpServerToolCallParams; use codex_app_server_protocol::McpServerToolCallResponse; use codex_app_server_protocol::RequestId; @@ -45,6 +49,8 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(20); const EXECUTOR_HTTP_MCP_URL: &str = "http://executor-only.invalid/mcp"; const HTTP_MCP_SERVER_NAME: &str = "executor_http"; const MCP_SERVER_NAME: &str = "executor_demo"; +const OAUTH_MCP_SERVER_NAME: &str = "executor_oauth"; +const EXECUTOR_OAUTH_MCP_URL: &str = "http://oauth-only.invalid/oauth-mcp"; const EXECUTOR_ENV_NAME: &str = "MCP_EXECUTOR_MARKER"; const EXECUTOR_ENV_VALUE: &str = "executor-only"; const EXECUTOR_ID: &str = "executor-1"; @@ -56,12 +62,35 @@ async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Resu let responses_server = responses::start_mock_server().await; let http_listener = TcpListener::bind("127.0.0.1:0").await?; let http_addr = http_listener.local_addr()?; + let http_server_config = StreamableHttpServerConfig::default() + .with_allowed_hosts(["executor-only.invalid", "oauth-only.invalid"]); let http_mcp_service = StreamableHttpService::new( || Ok(ExecutorHttpMcpServer), Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default().with_allowed_hosts(["executor-only.invalid"]), + http_server_config.clone(), ); - let http_router = Router::new().nest_service("/mcp", http_mcp_service); + let oauth_mcp_service = StreamableHttpService::new( + || Ok(ExecutorHttpMcpServer), + Arc::new(LocalSessionManager::default()), + http_server_config, + ); + let oauth_metadata = json!({ + "authorization_endpoint": "https://oauth-only.invalid/authorize", + "token_endpoint": "https://oauth-only.invalid/token", + "scopes_supported": ["read", "write"], + "response_types_supported": ["code"], + "code_challenge_methods_supported": ["S256"], + }); + let http_router = Router::new() + .route( + "/.well-known/oauth-authorization-server/oauth-mcp", + get(move || { + let metadata = oauth_metadata.clone(); + async move { Json(metadata) } + }), + ) + .nest_service("/mcp", http_mcp_service) + .nest_service("/oauth-mcp", oauth_mcp_service); let http_server_handle = tokio::spawn(async move { let _ = axum::serve(http_listener, http_router).await; }); @@ -117,6 +146,12 @@ HTTP_PROXY = {http_proxy} "url": EXECUTOR_HTTP_MCP_URL, "environment_id": "local", "startup_timeout_sec": 10, + }, + (OAUTH_MCP_SERVER_NAME): { + "url": EXECUTOR_OAUTH_MCP_URL, + "environment_id": "local", + "oauth": {"clientId": "executor-oauth-client"}, + "startup_timeout_sec": 10, } } }))?, @@ -158,6 +193,44 @@ startup_timeout_sec = 10 ) .await??; + let request_id = app_server + .send_raw_request( + "mcpServer/oauth/login", + Some(json!({ + "name": OAUTH_MCP_SERVER_NAME, + "threadId": selected_thread.clone(), + "timeoutSecs": 1, + })), + ) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: McpServerOauthLoginResponse = to_response(response)?; + assert!( + response + .authorization_url + .starts_with("https://oauth-only.invalid/authorize?") + ); + assert!( + response + .authorization_url + .contains("client_id=executor-oauth-client") + ); + let notification = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_notification_message("mcpServer/oauthLogin/completed"), + ) + .await??; + let completed: McpServerOauthLoginCompletedNotification = + serde_json::from_value(notification.params.expect("notification params"))?; + assert_eq!( + completed.thread_id.as_deref(), + Some(selected_thread.as_str()) + ); + let namespace = format!("mcp__{MCP_SERVER_NAME}"); let response_mock = responses::mount_sse_sequence( &responses_server, @@ -269,15 +342,18 @@ startup_timeout_sec = 10 .iter() .any(|name| name == HTTP_MCP_SERVER_NAME) ); + assert!( + selected_server_names + .iter() + .any(|name| name == OAUTH_MCP_SERVER_NAME) + ); let unselected_thread = start_thread(&mut app_server, /*selected_capability_roots*/ None).await?; let unselected_server_names = mcp_server_names(&mut app_server, unselected_thread).await?; - assert!( - unselected_server_names - .iter() - .all(|name| { name != MCP_SERVER_NAME && name != HTTP_MCP_SERVER_NAME }) - ); + assert!(unselected_server_names.iter().all(|name| { + name != MCP_SERVER_NAME && name != HTTP_MCP_SERVER_NAME && name != OAUTH_MCP_SERVER_NAME + })); http_server_handle.abort(); let _ = http_server_handle.await; diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 15f0b0bb4d..586f9c184b 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -18,8 +18,10 @@ use codex_core::config::edit::ConfigEditsBuilder; use codex_core::config::find_codex_home; use codex_core::config::load_global_mcp_servers; use codex_core_plugins::PluginsManager; +use codex_exec_server::EnvironmentManager; use codex_login::AuthManager; use codex_mcp::McpOAuthLoginSupport; +use codex_mcp::McpRuntimeContext; use codex_mcp::ResolvedMcpOAuthScopes; use codex_mcp::compute_auth_statuses; use codex_mcp::discover_supported_scopes; @@ -558,6 +560,10 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), auth.as_ref(), + &McpRuntimeContext::new( + Arc::new(EnvironmentManager::without_environments()), + config.cwd.to_path_buf(), + ), ) .await; diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index b4d397eac9..98efdc1b56 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -61,7 +61,9 @@ pub use mcp::McpOAuthScopesSource; pub use mcp::ResolvedMcpOAuthScopes; pub use mcp::compute_auth_statuses; pub use mcp::discover_supported_scopes; +pub use mcp::discover_supported_scopes_with_http_client; pub use mcp::oauth_login_support; +pub use mcp::oauth_login_support_with_http_client; pub use mcp::resolve_oauth_scopes; pub use mcp::should_retry_without_scopes; diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index d19766af1f..111f04163e 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::sync::Arc; use anyhow::Result; use codex_config::McpServerAuth; @@ -6,15 +7,19 @@ use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::HttpClient; use codex_login::CodexAuth; use codex_protocol::protocol::McpAuthStatus; use codex_rmcp_client::OAuthProviderError; use codex_rmcp_client::determine_streamable_http_auth_status; +use codex_rmcp_client::determine_streamable_http_auth_status_with_http_client; use codex_rmcp_client::discover_streamable_http_oauth; +use codex_rmcp_client::discover_streamable_http_oauth_with_http_client; use futures::FutureExt; use futures::future::join_all; use tracing::warn; +use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; #[derive(Debug, Clone)] @@ -53,6 +58,50 @@ pub struct McpAuthStatusEntry { } pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport { + let Some(mut config) = oauth_login_candidate(transport) else { + return McpOAuthLoginSupport::Unsupported; + }; + match discover_streamable_http_oauth( + &config.url, + config.http_headers.clone(), + config.env_http_headers.clone(), + ) + .await + { + Ok(Some(discovery)) => { + config.discovered_scopes = discovery.scopes_supported; + McpOAuthLoginSupport::Supported(config) + } + Ok(None) => McpOAuthLoginSupport::Unsupported, + Err(err) => McpOAuthLoginSupport::Unknown(err), + } +} + +pub async fn oauth_login_support_with_http_client( + transport: &McpServerTransportConfig, + http_client: Arc, +) -> McpOAuthLoginSupport { + let Some(mut config) = oauth_login_candidate(transport) else { + return McpOAuthLoginSupport::Unsupported; + }; + match discover_streamable_http_oauth_with_http_client( + &config.url, + config.http_headers.clone(), + config.env_http_headers.clone(), + http_client, + ) + .await + { + Ok(Some(discovery)) => { + config.discovered_scopes = discovery.scopes_supported; + McpOAuthLoginSupport::Supported(config) + } + Ok(None) => McpOAuthLoginSupport::Unsupported, + Err(err) => McpOAuthLoginSupport::Unknown(err), + } +} + +fn oauth_login_candidate(transport: &McpServerTransportConfig) -> Option { let McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, @@ -60,24 +109,17 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu env_http_headers, } = transport else { - return McpOAuthLoginSupport::Unsupported; + return None; }; - if bearer_token_env_var.is_some() { - return McpOAuthLoginSupport::Unsupported; - } - - match discover_streamable_http_oauth(url, http_headers.clone(), env_http_headers.clone()).await - { - Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(McpOAuthLoginConfig { - url: url.clone(), - http_headers: http_headers.clone(), - env_http_headers: env_http_headers.clone(), - discovered_scopes: discovery.scopes_supported, - }), - Ok(None) => McpOAuthLoginSupport::Unsupported, - Err(err) => McpOAuthLoginSupport::Unknown(err), + return None; } + Some(McpOAuthLoginConfig { + url: url.clone(), + http_headers: http_headers.clone(), + env_http_headers: env_http_headers.clone(), + discovered_scopes: None, + }) } pub async fn discover_supported_scopes( @@ -89,6 +131,16 @@ pub async fn discover_supported_scopes( } } +pub async fn discover_supported_scopes_with_http_client( + transport: &McpServerTransportConfig, + http_client: Arc, +) -> Option> { + match oauth_login_support_with_http_client(transport, http_client).await { + McpOAuthLoginSupport::Supported(config) => config.discovered_scopes, + McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None, + } +} + pub fn resolve_oauth_scopes( explicit_scopes: Option>, configured_scopes: Option>, @@ -133,6 +185,7 @@ pub async fn compute_auth_statuses<'a, I>( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, auth: Option<&CodexAuth>, + runtime_context: &McpRuntimeContext, ) -> HashMap where I: IntoIterator, @@ -140,6 +193,7 @@ where let futures = servers.into_iter().map(|(name, server)| { let name = name.clone(); let config = server.configured_config().cloned(); + let runtime_context = runtime_context.clone(); let has_runtime_auth = config .as_ref() .is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt)) @@ -162,6 +216,7 @@ where store_mode, keyring_backend_kind, has_runtime_auth, + &runtime_context, ) .await { @@ -193,6 +248,7 @@ async fn compute_auth_status( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, has_runtime_auth: bool, + runtime_context: &McpRuntimeContext, ) -> Result { if !config.enabled { return Ok(McpAuthStatus::Unsupported); @@ -210,17 +266,35 @@ async fn compute_auth_status( http_headers, env_http_headers, } => { - determine_streamable_http_auth_status( - server_name, - url, - bearer_token_env_var.as_deref(), - http_headers.clone(), - env_http_headers.clone(), - store_mode, - keyring_backend_kind, - ) - .boxed() - .await + if config.is_local_environment() { + determine_streamable_http_auth_status( + server_name, + url, + bearer_token_env_var.as_deref(), + http_headers.clone(), + env_http_headers.clone(), + store_mode, + keyring_backend_kind, + ) + .boxed() + .await + } else { + let http_client = runtime_context + .resolve_http_client(server_name, config) + .map_err(anyhow::Error::msg)?; + determine_streamable_http_auth_status_with_http_client( + server_name, + url, + bearer_token_env_var.as_deref(), + http_headers.clone(), + env_http_headers.clone(), + store_mode, + keyring_backend_kind, + http_client, + ) + .boxed() + .await + } } } } diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index f2c8f1d020..4550afa799 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -5,7 +5,9 @@ pub use auth::McpOAuthScopesSource; pub use auth::ResolvedMcpOAuthScopes; pub use auth::compute_auth_statuses; pub use auth::discover_supported_scopes; +pub use auth::discover_supported_scopes_with_http_client; pub use auth::oauth_login_support; +pub use auth::oauth_login_support_with_http_client; pub use auth::resolve_oauth_scopes; pub use auth::should_retry_without_scopes; @@ -309,6 +311,7 @@ pub async fn read_mcp_resource( config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind, auth, + &runtime_context, ) .await; let (tx_event, rx_event) = unbounded(); @@ -378,6 +381,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind, auth, + &runtime_context, ) .await; diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index be9d50d3f4..b9b769ffb4 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -11,6 +11,8 @@ use std::time::Duration; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use codex_protocol::models::PermissionProfile; use codex_utils_path_uri::PathUri; use serde::Deserialize; @@ -81,6 +83,20 @@ impl McpRuntimeContext { config.environment_id )) } + + /// Resolves the HTTP capability owned by the server's configured environment. + pub fn resolve_http_client( + &self, + server_name: &str, + config: &codex_config::McpServerConfig, + ) -> Result, String> { + Ok(self + .resolve_server_environment(server_name, config)? + .map_or_else( + || Arc::new(ReqwestHttpClient) as Arc, + |environment| environment.get_http_client(), + )) + } } pub(crate) fn emit_duration(metric: &str, duration: Duration, tags: &[(&str, &str)]) { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 79978da998..76b4a84830 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -246,11 +246,14 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( }); } + let runtime_context = + McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf()); let auth_status_entries = compute_auth_statuses( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), auth.as_ref(), + &runtime_context, ) .await; @@ -270,7 +273,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( PermissionProfile::default(), // Connector discovery is threadless. Use an actually configured env if // one exists, but do not reintroduce the old hidden-local fallback. - McpRuntimeContext::new(environment_manager, config.cwd.to_path_buf()), + runtime_context, config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 5a115b917f..c060ef090b 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -271,13 +271,6 @@ impl Session { let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let mcp_servers = effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); - let auth_statuses = compute_auth_statuses( - mcp_servers.iter(), - store_mode, - keyring_backend_kind, - auth.as_ref(), - ) - .await; let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. @@ -291,6 +284,14 @@ impl Session { turn_context.cwd.to_path_buf() }); let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd); + let auth_statuses = compute_auth_statuses( + mcp_servers.iter(), + store_mode, + keyring_backend_kind, + auth.as_ref(), + &mcp_runtime_context, + ) + .await; let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 5988834ee5..e43a26bbe5 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -648,6 +648,10 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; + let mcp_runtime_context_for_auth = McpRuntimeContext::new( + Arc::clone(&environment_manager), + session_configuration.cwd().to_path_buf(), + ); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp @@ -660,6 +664,7 @@ impl Session { config_for_mcp.mcp_oauth_credentials_store_mode, config_for_mcp.auth_keyring_backend_kind(), auth.as_ref(), + &mcp_runtime_context_for_auth, ) .await; (auth, mcp_servers, auth_statuses, tool_plugin_provenance) diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index a4d56f8cca..fc1a7df18b 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -24,6 +24,7 @@ use codex_config::types::McpServerTransportConfig; use codex_core::config::Config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::Environment; +use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRequestParams; use codex_login::CodexAuth; use codex_mcp::MCP_SANDBOX_STATE_META_CAPABILITY; @@ -2934,6 +2935,7 @@ async fn wait_for_remote_streamable_http_server( headers: Vec::new(), body: None, timeout_ms: Some(remaining.as_millis().clamp(1, 1_000) as u64), + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: false, }; diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index 133654dc0a..491fc8cf97 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -419,6 +419,17 @@ pub struct HttpHeader { pub value: String, } +/// Redirect behavior for an executor-side HTTP request. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum HttpRedirectPolicy { + /// Follow redirects using the HTTP client's normal limits. + #[default] + Follow, + /// Return the redirect response without following its location. + Stop, +} + /// Executor-side HTTP request envelope. /// /// This intentionally stays transport-shaped rather than MCP-shaped so callers @@ -443,6 +454,9 @@ pub struct HttpRequestParams { /// millisecond deadline. #[serde(default, skip_serializing_if = "Option::is_none")] pub timeout_ms: Option, + /// Whether the executor should follow HTTP redirects. + #[serde(default)] + pub redirect_policy: HttpRedirectPolicy, /// Caller-chosen stream id for `http/request/bodyDelta` notifications. /// /// The id must remain unique on a connection until the terminal body delta diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index 9de35315e1..2fcb7c7e6c 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -24,6 +24,7 @@ use super::response_body_stream::send_body_delta; use crate::HttpClient; use crate::client::ExecServerError; use crate::protocol::HttpHeader; +use crate::protocol::HttpRedirectPolicy; use crate::protocol::HttpRequestBodyDeltaNotification; use crate::protocol::HttpRequestParams; use crate::protocol::HttpRequestResponse; @@ -50,13 +51,20 @@ pub(crate) struct ReqwestHttpRequestRunner { } impl ReqwestHttpClient { - fn build_client(timeout_ms: Option) -> Result { + fn build_client( + timeout_ms: Option, + redirect_policy: HttpRedirectPolicy, + ) -> Result { let builder = match timeout_ms { None => reqwest::Client::builder(), Some(timeout_ms) => { reqwest::Client::builder().timeout(Duration::from_millis(timeout_ms)) } }; + let builder = match redirect_policy { + HttpRedirectPolicy::Follow => builder, + HttpRedirectPolicy::Stop => builder.redirect(reqwest::redirect::Policy::none()), + }; build_reqwest_client_with_custom_ca(builder) .map_err(|error| ExecServerError::HttpRequest(error.to_string())) } @@ -68,7 +76,7 @@ impl HttpClient for ReqwestHttpClient { params: HttpRequestParams, ) -> BoxFuture<'_, Result> { async move { - let runner = ReqwestHttpRequestRunner::new(params.timeout_ms) + let runner = ReqwestHttpRequestRunner::new(params.timeout_ms, params.redirect_policy) .map_err(|error| ExecServerError::HttpRequest(error.message))?; let (response, _) = runner .run(HttpRequestParams { @@ -87,7 +95,7 @@ impl HttpClient for ReqwestHttpClient { params: HttpRequestParams, ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { async move { - let runner = ReqwestHttpRequestRunner::new(params.timeout_ms) + let runner = ReqwestHttpRequestRunner::new(params.timeout_ms, params.redirect_policy) .map_err(|error| ExecServerError::HttpRequest(error.message))?; let (response, pending_stream) = runner .run(HttpRequestParams { @@ -111,8 +119,11 @@ impl HttpClient for ReqwestHttpClient { } impl ReqwestHttpRequestRunner { - pub(crate) fn new(timeout_ms: Option) -> Result { - let client = ReqwestHttpClient::build_client(timeout_ms) + pub(crate) fn new( + timeout_ms: Option, + redirect_policy: HttpRedirectPolicy, + ) -> Result { + let client = ReqwestHttpClient::build_client(timeout_ms, redirect_policy) .map_err(|error| internal_error(error.to_string()))?; Ok(Self { client }) } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 69d4c07dfb..5b2bfa7e55 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -123,6 +123,7 @@ pub use protocol::FsWalkResponse; pub use protocol::FsWriteFileParams; pub use protocol::FsWriteFileResponse; pub use protocol::HttpHeader; +pub use protocol::HttpRedirectPolicy; pub use protocol::HttpRequestBodyDeltaNotification; pub use protocol::HttpRequestParams; pub use protocol::HttpRequestResponse; diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index 561cf0ff63..02858ae887 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -211,7 +211,7 @@ impl ExecServerHandler { if stream_response { self.reserve_http_body_stream(&http_request_id).await?; } - let response = ReqwestHttpRequestRunner::new(params.timeout_ms)? + let response = ReqwestHttpRequestRunner::new(params.timeout_ms, params.redirect_policy)? .run(params) .await; if response.is_err() && stream_response { diff --git a/codex-rs/exec-server/tests/http_client.rs b/codex-rs/exec-server/tests/http_client.rs index 248564b493..b3f0182c63 100644 --- a/codex-rs/exec-server/tests/http_client.rs +++ b/codex-rs/exec-server/tests/http_client.rs @@ -6,6 +6,7 @@ use anyhow::Result; use anyhow::bail; use codex_exec_server::ExecServerClient; use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRequestBodyDeltaNotification; use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpRequestResponse; @@ -63,6 +64,7 @@ async fn http_request_forces_buffered_request_params() -> Result<()> { headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "ignored-stream-id".to_string(), stream_response: false, } @@ -91,6 +93,7 @@ async fn http_request_forces_buffered_request_params() -> Result<()> { headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "ignored-stream-id".to_string(), stream_response: true, }), @@ -131,6 +134,7 @@ async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_delta }], body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -186,6 +190,7 @@ async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_delta headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-2".to_string(), stream_response: true, } @@ -215,6 +220,7 @@ async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_delta }], body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -253,6 +259,7 @@ async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_delta headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -289,6 +296,7 @@ async fn http_response_body_stream_drops_queued_terminal_before_next_generated_i headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -322,6 +330,7 @@ async fn http_response_body_stream_drops_queued_terminal_before_next_generated_i headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-2".to_string(), stream_response: true, } @@ -348,6 +357,7 @@ async fn http_response_body_stream_drops_queued_terminal_before_next_generated_i headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -372,6 +382,7 @@ async fn http_response_body_stream_drops_queued_terminal_before_next_generated_i headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }; @@ -411,6 +422,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -430,6 +442,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-2".to_string(), stream_response: true, } @@ -474,6 +487,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }) @@ -495,6 +509,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -541,6 +556,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -580,6 +596,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-2".to_string(), stream_response: true, } @@ -615,6 +632,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -647,6 +665,7 @@ async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -691,6 +710,7 @@ async fn http_response_body_stream_fails_when_transport_disconnects() -> Result< headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -717,6 +737,7 @@ async fn http_response_body_stream_fails_when_transport_disconnects() -> Result< headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -759,6 +780,7 @@ async fn http_response_body_stream_reports_disconnect_when_queue_is_full() -> Re headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -796,6 +818,7 @@ async fn http_response_body_stream_reports_disconnect_when_queue_is_full() -> Re headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), @@ -852,6 +875,7 @@ async fn http_response_body_stream_reports_backpressure_truncation() -> Result<( headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "http-1".to_string(), stream_response: true, } @@ -894,6 +918,7 @@ async fn http_response_body_stream_reports_backpressure_truncation() -> Result<( headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "caller-stream-id".to_string(), stream_response: false, }), diff --git a/codex-rs/exec-server/tests/http_request.rs b/codex-rs/exec-server/tests/http_request.rs index a4734bf92d..1edacc4b13 100644 --- a/codex-rs/exec-server/tests/http_request.rs +++ b/codex-rs/exec-server/tests/http_request.rs @@ -7,6 +7,7 @@ use std::io::ErrorKind; use std::time::Duration; use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRequestBodyDeltaNotification; use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpRequestResponse; @@ -63,6 +64,7 @@ async fn exec_server_http_request_buffers_response_body() -> anyhow::Result<()> }], body: Some(b"request-body".to_vec().into()), timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: false, })?, @@ -108,6 +110,61 @@ async fn exec_server_http_request_buffers_response_body() -> anyhow::Result<()> Ok(()) } +/// What this tests: OAuth callers can inspect redirect responses without the +/// executor following the Location header. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_server_http_request_can_stop_at_redirects() -> anyhow::Result<()> { + let mut server = exec_server().await?; + initialize_exec_server(&mut server).await?; + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let base_url = format!("http://{}", listener.local_addr()?); + let http_request_id = server + .send_request( + "http/request", + serde_json::to_value(HttpRequestParams { + method: "GET".to_string(), + url: format!("{base_url}/redirect"), + headers: Vec::new(), + body: None, + timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Stop, + request_id: "redirect-request".to_string(), + stream_response: false, + })?, + ) + .await?; + + let captured = accept_http_request(&listener).await?; + assert_eq!(captured.request_line, "GET /redirect HTTP/1.1"); + respond_with_status_and_headers( + captured.stream, + "302 Found", + &[("location", &format!("{base_url}/final"))], + b"redirect", + ) + .await?; + + let response: HttpRequestResponse = wait_for_response(&mut server, http_request_id).await?; + assert_eq!( + ( + response.status, + response_header(&response.headers, "location"), + response.body.into_inner(), + ), + (302, Some(format!("{base_url}/final")), b"redirect".to_vec(),) + ); + assert!( + timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err(), + "redirect target should not be requested" + ); + + server.shutdown().await?; + Ok(()) +} + /// What this tests: a real exec-server websocket `http/request` can return /// response headers immediately and stream the response body as ordered /// `http/request/bodyDelta` notifications. @@ -132,6 +189,7 @@ async fn exec_server_http_request_streams_response_body_notifications() -> anyho }], body: None, timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Follow, request_id: "stream-1".to_string(), stream_response: true, })?, @@ -218,6 +276,7 @@ async fn exec_server_http_request_rejects_duplicate_stream_request_ids() -> anyh headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "stream-dup".to_string(), stream_response: true, })?, @@ -241,6 +300,7 @@ async fn exec_server_http_request_rejects_duplicate_stream_request_ids() -> anyh headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "stream-dup".to_string(), stream_response: true, })?, @@ -295,6 +355,7 @@ async fn exec_server_http_request_honors_optional_timeout() -> anyhow::Result<() headers: Vec::new(), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: false, })?, @@ -320,6 +381,7 @@ async fn exec_server_http_request_honors_optional_timeout() -> anyhow::Result<() headers: Vec::new(), body: None, timeout_ms: Some(10), + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: false, })?, diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 44f3eb1397..517f7270bb 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -16,7 +16,6 @@ axum = { workspace = true, default-features = false, features = [ ] } base64 = { workspace = true } codex-api = { workspace = true } -codex-client = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } codex-keyring-store = { workspace = true } diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index b2bd834d59..5b8378e67a 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -1,7 +1,9 @@ use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use anyhow::Result; +use codex_exec_server::HttpClient; use codex_protocol::protocol::McpAuthStatus; use futures::FutureExt; use reqwest::Client; @@ -13,6 +15,7 @@ use tracing::debug; use crate::oauth::StoredOAuthTokenStatus; use crate::oauth::oauth_token_status; +use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; @@ -25,6 +28,11 @@ pub struct StreamableHttpOAuthDiscovery { pub scopes_supported: Option>, } +enum AuthStatusCheck { + Complete(McpAuthStatus), + Discover(HeaderMap), +} + /// Determine the authentication status for a streamable HTTP MCP server. pub async fn determine_streamable_http_auth_status( server_name: &str, @@ -35,24 +43,100 @@ pub async fn determine_streamable_http_auth_status( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { + let default_headers = match auth_status_before_discovery( + server_name, + url, + bearer_token_env_var, + http_headers, + env_http_headers, + store_mode, + keyring_backend_kind, + )? { + AuthStatusCheck::Complete(status) => return Ok(status), + AuthStatusCheck::Discover(default_headers) => default_headers, + }; + + determine_auth_status_from_discovery( + server_name, + url, + discover_streamable_http_oauth_with_headers(url, &default_headers).await, + ) +} + +/// Determine authentication status while routing OAuth discovery through the +/// provided HTTP client. +#[allow(clippy::too_many_arguments)] +pub async fn determine_streamable_http_auth_status_with_http_client( + server_name: &str, + url: &str, + bearer_token_env_var: Option<&str>, + http_headers: Option>, + env_http_headers: Option>, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + http_client: Arc, +) -> Result { + let default_headers = match auth_status_before_discovery( + server_name, + url, + bearer_token_env_var, + http_headers, + env_http_headers, + store_mode, + keyring_backend_kind, + )? { + AuthStatusCheck::Complete(status) => return Ok(status), + AuthStatusCheck::Discover(default_headers) => default_headers, + }; + determine_auth_status_from_discovery( + server_name, + url, + discover_streamable_http_oauth_with_headers_and_http_client( + url, + default_headers, + http_client, + ) + .await, + ) +} + +fn auth_status_before_discovery( + server_name: &str, + url: &str, + bearer_token_env_var: Option<&str>, + http_headers: Option>, + env_http_headers: Option>, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result { if bearer_token_env_var.is_some() { - return Ok(McpAuthStatus::BearerToken); + return Ok(AuthStatusCheck::Complete(McpAuthStatus::BearerToken)); } let default_headers = build_default_headers(http_headers, env_http_headers)?; if default_headers.contains_key(AUTHORIZATION) { - return Ok(McpAuthStatus::BearerToken); + return Ok(AuthStatusCheck::Complete(McpAuthStatus::BearerToken)); } match oauth_token_status(server_name, url, store_mode, keyring_backend_kind)? { - StoredOAuthTokenStatus::Usable => return Ok(McpAuthStatus::OAuth), + StoredOAuthTokenStatus::Usable => { + return Ok(AuthStatusCheck::Complete(McpAuthStatus::OAuth)); + } StoredOAuthTokenStatus::AuthorizationRequired => { - return Ok(McpAuthStatus::NotLoggedIn); + return Ok(AuthStatusCheck::Complete(McpAuthStatus::NotLoggedIn)); } StoredOAuthTokenStatus::Missing => {} } - match discover_streamable_http_oauth_with_headers(url, &default_headers).await { + Ok(AuthStatusCheck::Discover(default_headers)) +} + +fn determine_auth_status_from_discovery( + server_name: &str, + url: &str, + discovery: Result>, +) -> Result { + match discovery { Ok(Some(_)) => Ok(McpAuthStatus::NotLoggedIn), Ok(None) => Ok(McpAuthStatus::Unsupported), Err(error) => { @@ -82,6 +166,17 @@ pub async fn discover_streamable_http_oauth( discover_streamable_http_oauth_with_headers(url, &default_headers).await } +pub async fn discover_streamable_http_oauth_with_http_client( + url: &str, + http_headers: Option>, + env_http_headers: Option>, + http_client: Arc, +) -> Result> { + let default_headers = build_default_headers(http_headers, env_http_headers)?; + discover_streamable_http_oauth_with_headers_and_http_client(url, default_headers, http_client) + .await +} + async fn discover_streamable_http_oauth_with_headers( url: &str, default_headers: &HeaderMap, @@ -92,6 +187,25 @@ async fn discover_streamable_http_oauth_with_headers( let client = apply_default_headers(builder, default_headers).build()?; let mut authorization_manager = AuthorizationManager::new(url).await?; authorization_manager.with_client(client)?; + discover_streamable_http_oauth_with_manager(&authorization_manager).await +} + +async fn discover_streamable_http_oauth_with_headers_and_http_client( + url: &str, + default_headers: HeaderMap, + http_client: Arc, +) -> Result> { + let authorization_manager = AuthorizationManager::new_with_oauth_http_client( + url, + Arc::new(OAuthHttpClientAdapter::new(http_client, default_headers)), + ) + .await?; + discover_streamable_http_oauth_with_manager(&authorization_manager).await +} + +async fn discover_streamable_http_oauth_with_manager( + authorization_manager: &AuthorizationManager, +) -> Result> { match authorization_manager.discover_metadata().boxed().await { Ok(metadata) => Ok(Some(StreamableHttpOAuthDiscovery { scopes_supported: normalize_scopes(metadata.scopes_supported), diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index a9cdc7b88f..19befb6235 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -16,6 +16,7 @@ use codex_api::SharedAuthProvider; use codex_exec_server::ExecServerError; use codex_exec_server::HttpClient; use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpResponseBodyStream; use futures::StreamExt; @@ -135,6 +136,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { headers: protocol_headers(&headers), body: Some(body.into()), timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: true, }) @@ -261,6 +263,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { headers: protocol_headers(&headers), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: false, }) @@ -330,6 +333,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { headers: protocol_headers(&headers), body: None, timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, request_id: "buffered-request".to_string(), stream_response: true, }) diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1ee18c753..e1041fdafb 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod http_client_adapter; mod in_process_transport; mod logging_client_handler; mod oauth; +mod oauth_http_client; mod perform_oauth_login; mod program_resolver; mod rmcp_client; @@ -13,7 +14,9 @@ mod utils; pub use auth_status::StreamableHttpOAuthDiscovery; pub use auth_status::determine_streamable_http_auth_status; +pub use auth_status::determine_streamable_http_auth_status_with_http_client; pub use auth_status::discover_streamable_http_oauth; +pub use auth_status::discover_streamable_http_oauth_with_http_client; pub use auth_status::supports_oauth_login; pub use codex_protocol::protocol::McpAuthStatus; pub use in_process_transport::InProcessTransportFactory; @@ -26,6 +29,7 @@ pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; pub use perform_oauth_login::perform_oauth_login; pub use perform_oauth_login::perform_oauth_login_return_url; +pub use perform_oauth_login::perform_oauth_login_return_url_with_http_client; pub use perform_oauth_login::perform_oauth_login_silent; pub use rmcp::model::ElicitationAction; pub use rmcp_client::Elicitation; diff --git a/codex-rs/rmcp-client/src/oauth_http_client.rs b/codex-rs/rmcp-client/src/oauth_http_client.rs new file mode 100644 index 0000000000..44cf63cb7c --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_http_client.rs @@ -0,0 +1,116 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use oauth2::HttpRequest; +use oauth2::HttpResponse; +use reqwest::header::HeaderMap; +use rmcp::transport::auth::OAuthHttpClient; +use rmcp::transport::auth::OAuthHttpClientError; +use rmcp::transport::auth::OAuthHttpClientFuture; +use rmcp::transport::auth::OAuthHttpRedirectPolicy; +use rmcp::transport::auth::OAuthHttpRequest; + +const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024; +static NEXT_OAUTH_REQUEST_ID: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone)] +pub(crate) struct OAuthHttpClientAdapter { + http_client: Arc, + default_headers: HeaderMap, +} + +impl OAuthHttpClientAdapter { + pub(crate) fn new(http_client: Arc, default_headers: HeaderMap) -> Self { + Self { + http_client, + default_headers, + } + } + + async fn execute_request( + &self, + request: HttpRequest, + redirect_policy: OAuthHttpRedirectPolicy, + timeout: Option, + ) -> Result { + let redirect_policy = match redirect_policy { + OAuthHttpRedirectPolicy::Follow => HttpRedirectPolicy::Follow, + OAuthHttpRedirectPolicy::Stop => HttpRedirectPolicy::Stop, + _ => { + return Err(OAuthHttpClientError::new( + "unsupported OAuth HTTP redirect policy", + )); + } + }; + let (parts, body) = request.into_parts(); + let mut headers = self.default_headers.clone(); + for name in parts.headers.keys() { + headers.remove(name); + } + headers.extend(parts.headers); + let headers = headers + .iter() + .map(|(name, value)| { + Ok(HttpHeader { + name: name.as_str().to_string(), + value: value + .to_str() + .map_err(|error| OAuthHttpClientError::new(error.to_string()))? + .to_string(), + }) + }) + .collect::, OAuthHttpClientError>>()?; + let timeout_ms = timeout.map(|timeout| { + u64::try_from(timeout.as_millis()) + .unwrap_or(u64::MAX) + .max(1) + }); + let request_id = NEXT_OAUTH_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let (response, mut body_stream) = self + .http_client + .http_request_stream(HttpRequestParams { + method: parts.method.to_string(), + url: parts.uri.to_string(), + headers, + body: (!body.is_empty()).then_some(body.into()), + timeout_ms, + redirect_policy, + request_id: format!("oauth-request-{request_id}"), + stream_response: true, + }) + .await + .map_err(|error| OAuthHttpClientError::new(error.to_string()))?; + let mut body = Vec::new(); + while let Some(chunk) = body_stream + .recv() + .await + .map_err(|error| OAuthHttpClientError::new(error.to_string()))? + { + if chunk.len() > MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES - body.len() { + return Err(OAuthHttpClientError::new(format!( + "OAuth HTTP response body exceeds {MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES} bytes" + ))); + } + body.extend_from_slice(&chunk); + } + let mut builder = oauth2::http::Response::builder().status(response.status); + for header in response.headers { + builder = builder.header(header.name, header.value); + } + builder + .body(body) + .map_err(|error| OAuthHttpClientError::new(error.to_string())) + } +} + +impl OAuthHttpClient for OAuthHttpClientAdapter { + fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + Box::pin(self.execute_request(request.request, request.redirect_policy, request.timeout)) + } +} diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 8c72f320bc..7a9d242bc4 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -9,11 +9,13 @@ use anyhow::anyhow; use anyhow::bail; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use reqwest::ClientBuilder; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use reqwest::Url; use rmcp::transport::AuthorizationManager; use rmcp::transport::AuthorizationSession; use rmcp::transport::auth::OAuthClientConfig; +use rmcp::transport::auth::OAuthHttpClient; use rmcp::transport::auth::OAuthState; use sha2::Digest; use sha2::Sha256; @@ -26,15 +28,16 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; +use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::save_oauth_tokens; -use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; -struct OauthHeaders { +struct OAuthHttpContext { http_headers: Option>, env_http_headers: Option>, + http_client: Arc, } struct CallbackServerGuard { @@ -154,16 +157,17 @@ async fn perform_oauth_login_with_browser_output( callback_url: Option<&str>, emit_browser_url: bool, ) -> Result<()> { - let headers = OauthHeaders { + let http_context = OAuthHttpContext { http_headers, env_http_headers, + http_client: Arc::new(ReqwestHttpClient), }; OauthLoginFlow::new( server_name, server_url, store_mode, keyring_backend_kind, - headers, + http_context, scopes, oauth_client_id, oauth_resource, @@ -192,16 +196,51 @@ pub async fn perform_oauth_login_return_url( callback_port: Option, callback_url: Option<&str>, ) -> Result { - let headers = OauthHeaders { + perform_oauth_login_return_url_with_http_client( + server_name, + server_url, + store_mode, + keyring_backend_kind, http_headers, env_http_headers, + scopes, + oauth_client_id, + oauth_resource, + timeout_secs, + callback_port, + callback_url, + Arc::new(ReqwestHttpClient), + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_return_url_with_http_client( + server_name: &str, + server_url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + http_headers: Option>, + env_http_headers: Option>, + scopes: &[String], + oauth_client_id: Option<&str>, + oauth_resource: Option<&str>, + timeout_secs: Option, + callback_port: Option, + callback_url: Option<&str>, + http_client: Arc, +) -> Result { + let http_context = OAuthHttpContext { + http_headers, + env_http_headers, + http_client, }; let flow = OauthLoginFlow::new( server_name, server_url, store_mode, keyring_backend_kind, - headers, + http_context, scopes, oauth_client_id, oauth_resource, @@ -456,7 +495,7 @@ impl OauthLoginFlow { server_url: &str, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, - headers: OauthHeaders, + http_context: OAuthHttpContext, scopes: &[String], oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -487,17 +526,18 @@ impl OauthLoginFlow { let (tx, rx) = oneshot::channel(); spawn_callback_server(server, tx, callback_path); - let OauthHeaders { + let OAuthHttpContext { http_headers, env_http_headers, - } = headers; + http_client, + } = http_context; let default_headers = build_default_headers(http_headers, env_http_headers)?; - let http_client = apply_default_headers(ClientBuilder::new(), &default_headers).build()?; + let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new(http_client, default_headers)); let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); let oauth_state = start_authorization( server_url, - http_client, + oauth_http_client, &scope_refs, &redirect_uri, oauth_client_id, @@ -620,22 +660,23 @@ impl OauthLoginFlow { async fn start_authorization( server_url: &str, - http_client: reqwest::Client, + http_client: Arc, scopes: &[&str], redirect_uri: &str, oauth_client_id: Option<&str>, ) -> Result { let Some(oauth_client_id) = oauth_client_id.filter(|client_id| !client_id.trim().is_empty()) else { - let mut oauth_state = OAuthState::new(server_url, Some(http_client)).await?; + let mut oauth_state = + OAuthState::new_with_oauth_http_client(server_url, http_client).await?; oauth_state .start_authorization(scopes, redirect_uri, Some("Codex")) .await?; return Ok(oauth_state); }; - let mut auth_manager = AuthorizationManager::new(server_url).await?; - auth_manager.with_client(http_client)?; + let mut auth_manager = + AuthorizationManager::new_with_oauth_http_client(server_url, http_client).await?; let metadata = auth_manager.discover_metadata().await?; auth_manager.set_metadata(metadata); auth_manager.configure_client( @@ -668,15 +709,20 @@ fn append_query_param(url: &str, key: &str, value: Option<&str>) -> String { #[cfg(test)] mod tests { + use std::sync::Arc; + use axum::Json; use axum::Router; use axum::routing::get; + use codex_exec_server::ReqwestHttpClient; use pretty_assertions::assert_eq; use reqwest::Url; + use reqwest::header::HeaderMap; use serde_json::json; use tokio::net::TcpListener; use super::CallbackOutcome; + use super::OAuthHttpClientAdapter; use super::OAuthProviderError; use super::append_callback_id_to_redirect_uri; use super::append_query_param; @@ -727,7 +773,10 @@ mod tests { let base_url = spawn_oauth_metadata_server().await; let oauth_state = start_authorization( &format!("{base_url}/mcp"), - reqwest::Client::new(), + Arc::new(OAuthHttpClientAdapter::new( + Arc::new(ReqwestHttpClient), + HeaderMap::new(), + )), &[], "http://127.0.0.1/callback", Some("eci-prd-pub-codex-123"), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 93abf7d114..c6527990fa 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -11,7 +11,6 @@ use std::time::Instant; use anyhow::Result; use anyhow::anyhow; use codex_api::SharedAuthProvider; -use codex_client::maybe_build_rustls_client_config_with_custom_ca; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerEnvVar; use codex_exec_server::HttpClient; @@ -68,11 +67,11 @@ use crate::in_process_transport::InProcessTransportFactory; use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; +use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; use crate::stdio_server_launcher::StdioServerTransport; -use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; @@ -1166,16 +1165,12 @@ async fn create_oauth_transport_and_runtime( StreamableHttpClientTransport>, OAuthPersistor, )> { - let mut builder = apply_default_headers(reqwest::Client::builder(), &default_headers); - if let Some(tls_config) = maybe_build_rustls_client_config_with_custom_ca()? { - builder = builder.tls_backend_preconfigured(tls_config.as_ref().clone()); - } - let oauth_metadata_client = builder.build()?; - // TODO(aibrahim): teach OAuth bootstrap and refresh to use the same - // shared HTTP client abstraction instead of always creating the local - // reqwest metadata client here. + let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( + http_client.clone(), + default_headers.clone(), + )); let mut oauth_state = - OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?; + OAuthState::new_with_oauth_http_client(url.to_string(), oauth_http_client).await?; oauth_state .set_credentials( From c38b2e9ba69cb57d197c6e5ba78b5e52ae0870f9 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 25 Jun 2026 10:45:20 +0100 Subject: [PATCH 04/33] Test executor-routed MCP OAuth token exchange (#29656) ## Why #28529 proves OAuth discovery uses the selected executor, but its end-to-end test stops before the callback and token exchange. ## What changed - add an executor-only mock token endpoint - complete the OAuth callback using the authorization URL's `state` and `redirect_uri` - assert the PKCE token exchange reaches the executor-only endpoint - assert the completion notification reports the selected thread and succeeds Depends on #28529. --- .../app-server/tests/suite/v2/executor_mcp.rs | 59 +++++++++++++++++-- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index f93af6821a..67d97bc248 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -4,7 +4,9 @@ use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml; use axum::Json; use axum::Router; +use axum::body::Bytes; use axum::routing::get; +use axum::routing::post; use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; @@ -43,6 +45,7 @@ use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; use tokio::net::TcpListener; +use tokio::sync::mpsc; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(20); @@ -74,9 +77,10 @@ async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Resu Arc::new(LocalSessionManager::default()), http_server_config, ); + let (token_request_tx, mut token_request_rx) = mpsc::unbounded_channel(); let oauth_metadata = json!({ "authorization_endpoint": "https://oauth-only.invalid/authorize", - "token_endpoint": "https://oauth-only.invalid/token", + "token_endpoint": "http://oauth-only.invalid/token", "scopes_supported": ["read", "write"], "response_types_supported": ["code"], "code_challenge_methods_supported": ["S256"], @@ -89,6 +93,21 @@ async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Resu async move { Json(metadata) } }), ) + .route( + "/token", + post(move |body: Bytes| { + let token_request_tx = token_request_tx.clone(); + async move { + let _ = token_request_tx.send(String::from_utf8_lossy(&body).into_owned()); + Json(json!({ + "access_token": "executor-access-token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "executor-refresh-token", + })) + } + }), + ) .nest_service("/mcp", http_mcp_service) .nest_service("/oauth-mcp", oauth_mcp_service); let http_server_handle = tokio::spawn(async move { @@ -199,7 +218,7 @@ startup_timeout_sec = 10 Some(json!({ "name": OAUTH_MCP_SERVER_NAME, "threadId": selected_thread.clone(), - "timeoutSecs": 1, + "timeoutSecs": 10, })), ) .await?; @@ -219,6 +238,33 @@ startup_timeout_sec = 10 .authorization_url .contains("client_id=executor-oauth-client") ); + let authorization_url = reqwest::Url::parse(&response.authorization_url)?; + let state = authorization_url + .query_pairs() + .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) + .expect("authorization URL should include state"); + let redirect_uri = authorization_url + .query_pairs() + .find_map(|(key, value)| (key == "redirect_uri").then(|| value.into_owned())) + .expect("authorization URL should include redirect_uri"); + let mut callback_url = reqwest::Url::parse(&redirect_uri)?; + callback_url + .query_pairs_mut() + .append_pair("code", "executor-test-code") + .append_pair("state", &state); + reqwest::Client::builder() + .no_proxy() + .build()? + .get(callback_url) + .send() + .await? + .error_for_status()?; + let token_request = timeout(DEFAULT_READ_TIMEOUT, token_request_rx.recv()) + .await? + .expect("executor token endpoint should receive a request"); + assert!(token_request.contains("grant_type=authorization_code")); + assert!(token_request.contains("code=executor-test-code")); + assert!(token_request.contains("code_verifier=")); let notification = timeout( DEFAULT_READ_TIMEOUT, app_server.read_stream_until_notification_message("mcpServer/oauthLogin/completed"), @@ -227,8 +273,13 @@ startup_timeout_sec = 10 let completed: McpServerOauthLoginCompletedNotification = serde_json::from_value(notification.params.expect("notification params"))?; assert_eq!( - completed.thread_id.as_deref(), - Some(selected_thread.as_str()) + completed, + McpServerOauthLoginCompletedNotification { + name: OAUTH_MCP_SERVER_NAME.to_string(), + thread_id: Some(selected_thread.clone()), + success: true, + error: None, + } ); let namespace = format!("mcp__{MCP_SERVER_NAME}"); From 268328001f899ec02cebab238fdeafda052525b6 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 25 Jun 2026 10:15:46 -0700 Subject: [PATCH 05/33] chore(app-server): mark thread/rollback as deprecated (#29928) We will drop support for this in the near future due to the complexity it introduces. --- .../schema/json/ClientRequest.json | 1 + .../codex_app_server_protocol.schemas.json | 1 + .../codex_app_server_protocol.v2.schemas.json | 1 + .../schema/json/v2/ThreadRollbackParams.json | 1 + .../typescript/v2/ThreadRollbackParams.ts | 3 +++ .../src/protocol/v2/thread.rs | 1 + codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/thread_processor.rs | 16 +++++++++++++++ .../tests/suite/v2/thread_rollback.rs | 20 +++++++++++++++++++ 10 files changed, 46 insertions(+), 1 deletion(-) diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 0027988c42..acfffc370f 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -4174,6 +4174,7 @@ "type": "object" }, "ThreadRollbackParams": { + "description": "DEPRECATED: `thread/rollback` will be removed soon.", "properties": { "numTurns": { "description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 33981b5b5f..8625c08da4 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -18997,6 +18997,7 @@ }, "ThreadRollbackParams": { "$schema": "http://json-schema.org/draft-07/schema#", + "description": "DEPRECATED: `thread/rollback` will be removed soon.", "properties": { "numTurns": { "description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 5a7cf80e83..94fb8a9890 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -16776,6 +16776,7 @@ }, "ThreadRollbackParams": { "$schema": "http://json-schema.org/draft-07/schema#", + "description": "DEPRECATED: `thread/rollback` will be removed soon.", "properties": { "numTurns": { "description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.", diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackParams.json index cb3ba0db39..aa52fbd598 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackParams.json @@ -1,5 +1,6 @@ { "$schema": "http://json-schema.org/draft-07/schema#", + "description": "DEPRECATED: `thread/rollback` will be removed soon.", "properties": { "numTurns": { "description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadRollbackParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadRollbackParams.ts index 1c938e3bfd..af416d1872 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadRollbackParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadRollbackParams.ts @@ -2,6 +2,9 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +/** + * DEPRECATED: `thread/rollback` will be removed soon. + */ export type ThreadRollbackParams = { threadId: string, /** * The number of turns to drop from the end of the thread. Must be >= 1. diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index c330c7eefd..0be7776279 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -1025,6 +1025,7 @@ pub struct ThreadBackgroundTerminalsTerminateResponse { #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] +/// DEPRECATED: `thread/rollback` will be removed soon. pub struct ThreadRollbackParams { pub thread_id: String, /// The number of turns to drop from the end of the thread. Must be >= 1. diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 397fa3d58f..4c88d4b69d 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -167,7 +167,7 @@ Example with notification opt-out: - `thread/backgroundTerminals/clean` — terminate all running background terminals for a thread (experimental; requires `capabilities.experimentalApi`); returns `{}` when the cleanup request is accepted. - `thread/backgroundTerminals/list` — list running background terminals for a loaded thread (experimental; requires `capabilities.experimentalApi`); returns `data` with the running terminal ids. - `thread/backgroundTerminals/terminate` — terminate one running background terminal by app-server `processId` (experimental; requires `capabilities.experimentalApi`); returns whether a process was terminated. -- `thread/rollback` — drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success. +- `thread/rollback` — deprecated and will be removed soon. Drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success. - `turn/start` — add user input to a thread and begin Codex generation; responds with the initial `turn` object and streams `turn/started`, `item/*`, and `turn/completed` notifications. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. Prefer experimental `permissions` profile selection by id for permission overrides; the legacy `sandboxPolicy` field is still accepted but cannot be combined with `permissions`. For `collaborationMode`, `settings.developer_instructions: null` means "use built-in instructions for the selected mode". Deprecated experimental `multiAgentMode` is ignored; Ultra reasoning effort selects proactive behavior. - `thread/inject_items` — append raw Responses API items to a loaded thread’s model-visible history without starting a user turn; returns `{}` on success. - `turn/steer` — add user input to an already in-flight regular turn without starting a new turn; returns the active `turnId` that accepted the input. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Review and manual compaction turns reject `turn/steer`. diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 9b56957a53..57a65fd313 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -56,6 +56,7 @@ use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditParams; use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditResponse; use codex_app_server_protocol::ConversationGitInfo; use codex_app_server_protocol::ConversationSummary; +use codex_app_server_protocol::DeprecationNoticeNotification; use codex_app_server_protocol::DynamicToolFunctionSpec; use codex_app_server_protocol::DynamicToolNamespaceTool; use codex_app_server_protocol::DynamicToolSpec; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 7242e816be..ab4e241867 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -8,6 +8,8 @@ use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; +const THREAD_ROLLBACK_DEPRECATION_SUMMARY: &str = + "thread/rollback is deprecated and will be removed soon"; struct ThreadListFilters { model_providers: Option>, @@ -633,11 +635,25 @@ impl ThreadRequestProcessor { request_id: &ConnectionRequestId, params: ThreadRollbackParams, ) -> Result, JSONRPCErrorError> { + self.send_thread_rollback_deprecation_notice(request_id.connection_id) + .await; self.thread_rollback_inner(request_id, params) .await .map(|()| None) } + async fn send_thread_rollback_deprecation_notice(&self, connection_id: ConnectionId) { + self.outgoing + .send_server_notification_to_connections( + &[connection_id], + ServerNotification::DeprecationNotice(DeprecationNoticeNotification { + summary: THREAD_ROLLBACK_DEPRECATION_SUMMARY.to_string(), + details: None, + }), + ) + .await; + } + pub(crate) async fn thread_list( &self, params: ThreadListParams, diff --git a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs index 9f0df0ff2d..9481f69232 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs @@ -3,6 +3,8 @@ use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; +use codex_app_server_protocol::DeprecationNoticeNotification; +use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; @@ -97,6 +99,7 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() mcp.read_stream_until_notification_message("turn/completed"), ) .await??; + mcp.clear_message_buffer(); // Roll back the last turn. let rollback_id = mcp @@ -105,6 +108,23 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() num_turns: 1, }) .await?; + let deprecation_notice = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??; + let JSONRPCMessage::Notification(deprecation_notice) = deprecation_notice else { + panic!("thread/rollback should emit deprecationNotice before its response"); + }; + assert_eq!(deprecation_notice.method, "deprecationNotice"); + let deprecation_notice: DeprecationNoticeNotification = serde_json::from_value( + deprecation_notice + .params + .expect("deprecationNotice params should be present"), + )?; + assert_eq!( + deprecation_notice, + DeprecationNoticeNotification { + summary: "thread/rollback is deprecated and will be removed soon".to_string(), + details: None, + } + ); let rollback_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_response_message(RequestId::Integer(rollback_id)), From 8f02973d2541fb98d9a626f07d8f3661cb2a9e71 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 25 Jun 2026 18:49:43 +0100 Subject: [PATCH 06/33] Persist selected capability roots and resolve availability per model step (#29856) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why `selectedCapabilityRoots` is durable thread intent: “use this capability root from environment `worker`.” The important product assumption is: > One environment ID always names the same logical executor and stable contents. `worker` does not silently change from executor A to an unrelated executor B. The process-local connection handle for `worker` can still be replaced while Codex is running, though, for example when `environment/add` registers a fresh handle for the same logical environment. The thread should persist only the stable selection. Each model step should pair that selection with the exact ready handle captured for that step. ## The boundary ```text persisted thread intent plugin@1 -> environment "worker" | | capture the current step v model-step view unavailable, or plugin@1 + worker's exact captured ready handle ``` The environment ID is the stable identity and cache key. The `Arc` is only a process-local handle retained so consumers of one model step use the same captured environment. It is never persisted and it does not imply different environment contents. ## What changes ### Persist the stable selection Selected roots are written into `SessionMeta` and restored with the thread. Forked subagents inherit the same selections, including bounded-history forks. Only stable data is persisted: root ID, environment ID, and root path. ### Capture readiness together with the exact handle The environment snapshot records: ```rust environment_id -> Some(Arc) // ready in this step environment_id -> None // still starting in this step ``` This prevents readiness and execution from coming from different registry snapshots. For example: ```text step snapshot: worker -> handle A, ready environment/add: worker -> fresh handle B for the same logical environment current step: plugin@1 still uses captured handle A ``` Without carrying handle A in the snapshot, the resolver could combine “A was ready” with handle B and treat B as ready before it had finished starting. This does not change cache invalidation. Stable capability metadata remains identified by environment ID and capability root. Replacing a process-local handle under the same stable environment ID does not invalidate or rediscover that metadata. ### Resolve availability per model step - A ready captured environment produces resolved roots using its captured handle. - A starting, missing, or failed environment is omitted from that step. - A selected lazy environment that is outside the turn's captured environment set is asked to start, and a later step can observe it as ready. - No capability files are scanned here. Transient transport disconnects remain the remote client's reconnect concern. This PR models initial attachment/readiness; it does not add live socket-connectivity state. ## Example ```text thread selection: plugin@1 -> environment "worker" step 1: worker is starting -> plugin@1 unavailable step 2: worker is ready -> plugin@1 resolves through worker's captured handle step 3: fresh local handle -> current step remains pinned; a later step captures its own view ``` Temporary unavailability does not discard the durable selection. Later PRs can retain stable metadata caches while projecting only currently available capabilities into model-visible World State. ## Compatibility The app-server request shape does not change. Older rollouts without `selected_capability_roots` deserialize to an empty list. ## Stack 1. **This PR:** persist stable selected roots and resolve them through an exact model-step handle. 2. #29960: cache stable skill metadata and project available skills into World State. 3. #29946: cache stable plugin declarations and manage the separate live MCP runtime. --- .../external_agent_session_import.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 2 + .../tests/suite/conversation_summary.rs | 1 + .../tests/suite/v2/remote_thread_store.rs | 1 + .../app-server/tests/suite/v2/thread_read.rs | 1 + .../tests/suite/v2/thread_resume.rs | 1 + .../tests/suite/v2/thread_unarchive.rs | 1 + codex-rs/core/src/agent/control/spawn.rs | 14 +++ codex-rs/core/src/agent/control_tests.rs | 39 +++++- codex-rs/core/src/environment_selection.rs | 19 +++ .../core/src/personality_migration_tests.rs | 1 + codex-rs/core/src/session/mod.rs | 10 ++ codex-rs/core/src/session/session.rs | 7 ++ codex-rs/core/src/session/step_context.rs | 5 + codex-rs/core/src/session/tests.rs | 5 + codex-rs/core/src/session/world_state.rs | 4 + codex-rs/core/src/state/service.rs | 4 + codex-rs/core/src/thread_manager.rs | 3 +- codex-rs/core/src/thread_manager_tests.rs | 67 ++++++++++ codex-rs/core/src/tools/spec_plan_tests.rs | 1 + .../core/tests/suite/personality_migration.rs | 2 + codex-rs/core/tests/suite/sqlite_state.rs | 1 + codex-rs/exec-server/src/environment.rs | 21 +++- codex-rs/exec-server/src/lib.rs | 2 + .../exec-server/src/resolved_capability.rs | 114 ++++++++++++++++++ .../tests/selected_capability_roots.rs | 69 +++++++++++ codex-rs/protocol/src/protocol.rs | 11 ++ codex-rs/rollout/src/compression_tests.rs | 1 + codex-rs/rollout/src/metadata_tests.rs | 3 + codex-rs/rollout/src/recorder.rs | 19 +++ codex-rs/rollout/src/recorder_tests.rs | 1 + codex-rs/rollout/src/session_index_tests.rs | 1 + codex-rs/rollout/src/state_db_tests.rs | 1 + codex-rs/rollout/src/tests.rs | 1 + codex-rs/state/src/extract.rs | 2 + codex-rs/state/src/runtime/threads.rs | 2 + codex-rs/thread-store/src/in_memory.rs | 2 + .../thread-store/src/local/create_thread.rs | 1 + codex-rs/thread-store/src/local/mod.rs | 1 + codex-rs/thread-store/src/types.rs | 4 + 40 files changed, 443 insertions(+), 3 deletions(-) create mode 100644 codex-rs/exec-server/src/resolved_capability.rs create mode 100644 codex-rs/exec-server/tests/selected_capability_roots.rs diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index 77b9bfcbf3..bc0a64ef72 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -215,6 +215,7 @@ impl ExternalAgentSessionImporter { .unwrap_or_else(|| model_info.get_model_instructions(config.personality)), }, dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: Some(MultiAgentVersion::V1), initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index ec01dacc47..8d2567c010 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -200,6 +200,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( model_provider: model_provider.map(str::to_string), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -288,6 +289,7 @@ pub fn create_fake_rollout_with_text_elements( model_provider: model_provider.map(str::to_string), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 7f964d3147..63ad473261 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -131,6 +131,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 8fdb38e5ea..7e0db1b5b3 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -158,6 +158,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 6b427b0d97..9f67761757 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1369,6 +1369,7 @@ async fn seed_pathless_store_thread( originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index ce95f4810c..533087a11e 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2070,6 +2070,7 @@ stream_max_retries = 0 model_provider: Some("mock_provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index f647dc2cfd..231ede47b9 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -218,6 +218,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 6919d8c627..c4c40577ca 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,5 +1,6 @@ use super::residency::is_v2_resident_session_source; use super::*; +use codex_extension_api::ExtensionDataInit; const AGENT_NAMES: &str = include_str!("../agent_names.txt"); @@ -433,6 +434,16 @@ impl AgentControl { )) })?; + let selected_capability_roots = parent_history + .items + .iter() + .find_map(|item| { + let RolloutItem::SessionMeta(meta_line) = item else { + return None; + }; + Some(meta_line.meta.selected_capability_roots.clone()) + }) + .unwrap_or_default(); let mut forked_rollout_items = parent_history.items; if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { forked_rollout_items = @@ -504,6 +515,8 @@ impl AgentControl { { forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message)); } + let mut thread_extension_init = ExtensionDataInit::new(); + thread_extension_init.insert(selected_capability_roots); state .fork_thread_with_source( @@ -517,6 +530,7 @@ impl AgentControl { inherited_environments, inherited_exec_policy, options.environments.clone(), + thread_extension_init, ) .await } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index e00a8624f4..de8754ad92 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -17,6 +17,8 @@ use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_protocol::AgentPath; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ModeKind; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; @@ -38,6 +40,7 @@ use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; use codex_thread_store::ThreadStore; +use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::time::Duration; @@ -1446,7 +1449,33 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { #[tokio::test] async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_limit() { let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; + let selected_capability_roots = vec![SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "build".to_string(), + path: PathUri::parse("file:///plugins/demo").expect("plugin root URI"), + }, + }]; + let mut thread_extension_init = ExtensionDataInit::new(); + thread_extension_init.insert(selected_capability_roots.clone()); + let parent = harness + .manager + .start_thread_with_options(StartThreadOptions { + config: harness.config.clone(), + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + thread_extension_init, + supports_openai_form_elicitation: false, + }) + .await + .expect("start parent thread"); + let parent_thread_id = parent.thread_id; + let parent_thread = parent.thread; let startup_turn_context = parent_thread.codex.session.new_default_turn().await; parent_thread .codex @@ -1525,6 +1554,14 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li !history_contains_text(history.raw_items(), "parent startup developer context"), "bounded fork should drop parent startup context even when fewer turns exist than requested" ); + assert_eq!( + &child_thread + .codex + .session + .services + .selected_capability_roots, + &selected_capability_roots + ); assert!( child_thread .codex diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index fa86dded55..ef6a3ec5ea 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::collections::HashSet; use std::fmt; use std::sync::Arc; @@ -227,6 +228,24 @@ pub(crate) struct TurnEnvironmentSnapshot { } impl TurnEnvironmentSnapshot { + /// Maps each captured environment to its exact ready handle, or `None` when it was starting. + pub(crate) fn captured_environments(&self) -> HashMap>> { + self.turn_environments + .iter() + .map(|environment| { + ( + environment.environment_id.clone(), + Some(Arc::clone(&environment.environment)), + ) + }) + .chain( + self.starting + .iter() + .map(|environment| (environment.selection.environment_id.clone(), None)), + ) + .collect() + } + pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { self.turn_environments.first() } diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index dd0da327c4..396d065644 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -59,6 +59,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index be12e61e08..94bef646aa 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2828,9 +2828,19 @@ impl Session { .await; } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; + let selected_capability_roots = self + .services + .turn_environments + .environment_manager() + .resolve_selected_capability_roots( + &self.services.selected_capability_roots, + &environments.captured_environments(), + ) + .await; Arc::new(StepContext::new( turn_context, environments, + selected_capability_roots, loaded_agents_md, )) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e43a26bbe5..96dffed888 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -10,6 +10,7 @@ use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; use codex_login::auth::AgentIdentityAuthPolicy; use codex_protocol::SessionId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; use codex_protocol::permissions::FileSystemPath; @@ -555,6 +556,10 @@ impl Session { config.current_time_reminder.as_ref(), external_time_provider, )?; + let selected_capability_roots = thread_extension_init + .get::>() + .map(|roots| roots.as_ref().clone()) + .unwrap_or_else(|| initial_history.get_selected_capability_roots()); let mcp_thread_init = thread_extension_init.clone(); let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), @@ -584,6 +589,7 @@ impl Session { text: session_configuration.base_instructions.clone(), }, dynamic_tools: session_configuration.dynamic_tools.clone(), + selected_capability_roots: selected_capability_roots.clone(), multi_agent_version: initial_multi_agent_version, initial_window_id: initial_auto_compact_window_ids .window_id @@ -1041,6 +1047,7 @@ impl Session { // TODO(jif): extract session to share between sub-agents session_extension_data, thread_extension_data, + selected_capability_roots, mcp_thread_init, supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new( supports_openai_form_elicitation, diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index bf672db6e7..e463015459 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -3,12 +3,15 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::turn_context::TurnContext; +use codex_exec_server::ResolvedSelectedCapabilityRoot; /// Request-scoped state that may change between model sampling requests. #[derive(Debug)] pub(crate) struct StepContext { pub(crate) turn: Arc, pub(crate) environments: TurnEnvironmentSnapshot, + /// Capability roots bound to ready environments in this exact step. + pub(crate) selected_capability_roots: Vec, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -17,11 +20,13 @@ impl StepContext { pub(crate) fn new( turn: Arc, environments: TurnEnvironmentSnapshot, + selected_capability_roots: Vec, loaded_agents_md: Option>, ) -> Self { Self { turn, environments, + selected_capability_roots, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 4553d54a0c..525fbc3b99 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -190,6 +190,7 @@ impl StepContext { Arc::new(Self::new( turn, environments, + Vec::new(), /*loaded_agents_md*/ None, )) } @@ -4043,6 +4044,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { @@ -5404,6 +5406,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, @@ -6919,6 +6922,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { @@ -7479,6 +7483,7 @@ where agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index e19e8c2ff4..cedeaab25e 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -10,6 +10,10 @@ impl Session { step_context: &StepContext, ) -> WorldState { let turn_context = step_context.turn.as_ref(); + tracing::trace!( + selected_capability_root_count = step_context.selected_capability_roots.len(), + "building step world state" + ); let environment_subagents = if turn_context.config.include_environment_context { self.services .agent_control diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index cfed7d90d0..3721cb1848 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -33,6 +33,7 @@ use codex_login::AuthManager; use codex_mcp::McpConnectionManager; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_rollout::state_db::StateDbHandle; use codex_rollout_trace::ThreadTraceContext; use codex_thread_store::LiveThread; @@ -72,6 +73,9 @@ pub(crate) struct SessionServices { pub(crate) session_extension_data: ExtensionData, pub(crate) thread_extension_data: ExtensionData, pub(crate) supports_openai_form_elicitation: AtomicBool, + /// Raw capability selections for this thread. Each model step resolves them against its + /// current executor environments before using them. + pub(crate) selected_capability_roots: Vec, pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: AgentControl, pub(crate) network_proxy: ArcSwapOption, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 98e13c5d67..b2add06024 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1402,6 +1402,7 @@ impl ThreadManagerState { inherited_environments: Option, inherited_exec_policy: Option>, environments: Option>, + thread_extension_init: ExtensionDataInit, ) -> CodexResult { let environments = environments.unwrap_or_else(|| { default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) @@ -1421,7 +1422,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, - /*thread_extension_init*/ ExtensionDataInit::default(), + thread_extension_init, /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index f1acc2d979..0d407e19c5 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -20,6 +20,8 @@ use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::ResumedHistory; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; @@ -610,6 +612,71 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() ); } +#[tokio::test] +async fn selected_capability_roots_round_trip_through_fork() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + let selected_roots = vec![SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "build".to_string(), + path: PathUri::parse("file:///plugins/demo").expect("plugin root URI"), + }, + }]; + let inherited = manager + .start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::Forked(vec![RolloutItem::SessionMeta( + SessionMetaLine { + meta: SessionMeta { + selected_capability_roots: selected_roots.clone(), + ..SessionMeta::default() + }, + git: None, + }, + )]), + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, + }) + .await + .expect("start inherited fork"); + inherited.thread.ensure_rollout_materialized().await; + inherited + .thread + .flush_rollout() + .await + .expect("flush inherited fork"); + let inherited_history = RolloutRecorder::get_rollout_history( + &inherited + .thread + .rollout_path() + .expect("inherited fork rollout path"), + ) + .await + .expect("read inherited fork rollout"); + + assert_eq!( + inherited_history.get_selected_capability_roots(), + selected_roots + ); +} + #[tokio::test] async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let temp_dir = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index c422cc749e..6fc5adcb11 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -689,6 +689,7 @@ async fn environment_tools_follow_the_step_context() { let step_context = Arc::new(StepContext::new( Arc::new(turn), environments, + Vec::new(), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index ca1d148728..3dd0f3c016 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -75,6 +75,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -127,6 +128,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 5eea5a09e7..1c2f08262a 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -372,6 +372,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 664ef730d0..66f7381103 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -54,7 +54,7 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str = #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, - environments: RwLock>>, + pub(super) environments: RwLock>>, local_environment: Option>, local_runtime_paths: Option, } @@ -574,6 +574,25 @@ impl Environment { } } + /// Starts the initial connection after an environment is actually selected for use. + pub(crate) fn start_connecting_for_use(environment: &Arc) { + if environment.remote_client.is_none() { + return; + } + let mut startup_task = environment + .startup_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if startup_task.is_none() { + let environment = Arc::clone(environment); + *startup_task = Some(AbortOnDropHandle::new(tokio::spawn(async move { + if let Err(error) = environment.wait_until_ready().await { + tracing::debug!(%error, "exec-server environment startup failed"); + } + }))); + } + } + /// Returns whether initial startup has either succeeded or permanently failed. pub fn startup_finished(&self) -> bool { self.remote_client diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 5b2bfa7e55..92c7f85edf 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -22,6 +22,7 @@ mod relay_proto; mod remote; mod remote_file_system; mod remote_process; +mod resolved_capability; mod rpc; mod runtime_paths; mod sandboxed_file_system; @@ -143,6 +144,7 @@ pub use protocol::WriteResponse; pub use protocol::WriteStatus; pub use remote::RemoteEnvironmentConfig; pub use remote::run_remote_environment; +pub use resolved_capability::ResolvedSelectedCapabilityRoot; pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; diff --git a/codex-rs/exec-server/src/resolved_capability.rs b/codex-rs/exec-server/src/resolved_capability.rs new file mode 100644 index 0000000000..2902cf9eef --- /dev/null +++ b/codex-rs/exec-server/src/resolved_capability.rs @@ -0,0 +1,114 @@ +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; + +use crate::Environment; +use crate::EnvironmentManager; +use crate::ExecutorFileSystem; + +/// A selected capability root paired with its currently ready environment handle. +/// +/// Environment IDs have stable identity and contents. This process-local value must not be +/// persisted: it only keeps the current connection handle alive while one model step uses the +/// stable environment. +#[derive(Clone)] +pub struct ResolvedSelectedCapabilityRoot { + selected_root: SelectedCapabilityRoot, + environment: Arc, +} + +impl ResolvedSelectedCapabilityRoot { + pub fn selected_root(&self) -> &SelectedCapabilityRoot { + &self.selected_root + } + + pub fn environment(&self) -> &Arc { + &self.environment + } + + pub fn file_system(&self) -> Arc { + self.environment.get_filesystem() + } +} + +impl EnvironmentManager { + /// Resolves selected roots whose stable environments are ready for the current model step. + /// + /// Environment identity comes from the selected root's stable environment ID. A ready + /// environment captured for the step carries its exact process-local handle so readiness and + /// execution cannot come from different registry snapshots. Missing, starting, or failed + /// environments are omitted. A lazy environment is started for a later step. + pub async fn resolve_selected_capability_roots( + &self, + selected_roots: &[SelectedCapabilityRoot], + captured_environments: &HashMap>>, + ) -> Vec { + let candidates = { + let environments = self + .environments + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + selected_roots + .iter() + .filter_map(|selected_root| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + let (environment, already_ready) = + match captured_environments.get(environment_id) { + Some(Some(environment)) => (Arc::clone(environment), true), + Some(None) => return None, + None => (Arc::clone(environments.get(environment_id)?), false), + }; + Some(( + ResolvedSelectedCapabilityRoot { + selected_root: selected_root.clone(), + environment, + }, + already_ready, + )) + }) + .collect::>() + }; + + let mut readiness = HashMap::new(); + for (candidate, already_ready) in &candidates { + let CapabilityRootLocation::Environment { environment_id, .. } = + &candidate.selected_root().location; + if readiness.contains_key(environment_id) { + continue; + } + let environment = candidate.environment(); + let ready = if *already_ready { + true + } else if environment.startup_finished() { + environment.wait_until_ready().await.is_ok() + } else { + Environment::start_connecting_for_use(environment); + false + }; + readiness.insert(environment_id.clone(), ready); + } + + candidates + .into_iter() + .map(|(candidate, _)| candidate) + .filter(|candidate| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &candidate.selected_root().location; + readiness.get(environment_id).copied().unwrap_or(false) + }) + .collect() + } +} + +impl fmt::Debug for ResolvedSelectedCapabilityRoot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ResolvedSelectedCapabilityRoot") + .field("selected_root", &self.selected_root) + .finish_non_exhaustive() + } +} diff --git a/codex-rs/exec-server/tests/selected_capability_roots.rs b/codex-rs/exec-server/tests/selected_capability_roots.rs new file mode 100644 index 0000000000..639d652ca8 --- /dev/null +++ b/codex-rs/exec-server/tests/selected_capability_roots.rs @@ -0,0 +1,69 @@ +#![cfg(unix)] + +mod common; + +use std::collections::HashMap; +use std::sync::Arc; + +use codex_exec_server::EnvironmentManager; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_utils_path_uri::PathUri; +use common::exec_server::exec_server; +use pretty_assertions::assert_eq; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_capability_roots_use_captured_handle_after_replacement() -> anyhow::Result<()> { + let mut executor = exec_server().await?; + let manager = EnvironmentManager::without_environments(); + let selected_root = SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "tools".to_string(), + path: PathUri::parse("file:///plugins/demo")?, + }, + }; + + manager.upsert_environment( + "tools".to_string(), + executor.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + let environment_a = manager + .get_environment("tools") + .expect("executor A should be registered"); + environment_a.wait_until_ready().await?; + + let unavailable = manager + .resolve_selected_capability_roots( + std::slice::from_ref(&selected_root), + &HashMap::from([("tools".to_string(), None)]), + ) + .await; + assert!(unavailable.is_empty()); + + let captured_environments = + HashMap::from([("tools".to_string(), Some(Arc::clone(&environment_a)))]); + // Replace only the process-local handle; the stable environment ID and executor stay the same. + manager.upsert_environment( + "tools".to_string(), + executor.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + + let available = manager + .resolve_selected_capability_roots( + std::slice::from_ref(&selected_root), + &captured_environments, + ) + .await; + let [resolved] = available.as_slice() else { + anyhow::bail!("selected root should resolve through its stable environment"); + }; + + assert_eq!(resolved.selected_root(), &selected_root); + assert!(Arc::ptr_eq(resolved.environment(), &environment_a)); + + executor.shutdown().await?; + Ok(()) +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 1affdf0d20..42b01b16f3 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -19,6 +19,7 @@ use crate::AgentPath; use crate::SessionId; use crate::ThreadId; use crate::approvals::ElicitationRequestEvent; +use crate::capabilities::SelectedCapabilityRoot; use crate::config_types::ApprovalsReviewer; use crate::config_types::CollaborationMode; use crate::config_types::ModeKind; @@ -2570,6 +2571,12 @@ impl InitialHistory { } } + pub fn get_selected_capability_roots(&self) -> Vec { + self.get_session_meta() + .map(|meta| meta.selected_capability_roots.clone()) + .unwrap_or_default() + } + pub fn get_multi_agent_version(&self) -> Option { match self { InitialHistory::New | InitialHistory::Cleared => None, @@ -3003,6 +3010,9 @@ pub struct SessionMeta { skip_serializing_if = "Option::is_none" )] pub dynamic_tools: Option>, + /// Capability roots selected for this thread by the hosting platform. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub selected_capability_roots: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub memory_mode: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -3032,6 +3042,7 @@ impl Default for SessionMeta { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index a4db9a965c..733e6cc5fb 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -472,6 +472,7 @@ fn write_rollout(path: &std::path::Path, thread_id: ThreadId, message: &str) -> model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 0744e7780e..c0aea248c5 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -49,6 +49,7 @@ async fn extract_metadata_from_rollout_uses_session_meta() { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -106,6 +107,7 @@ async fn extract_metadata_from_rollout_returns_latest_memory_mode() { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -375,6 +377,7 @@ fn write_rollout_in_sessions_with_cwd( model_provider: Some("test-provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 3908604c77..58d5c76662 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -12,6 +12,7 @@ use std::sync::Mutex; use chrono::SecondsFormat; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; use serde_json::Value; @@ -91,6 +92,7 @@ pub enum RolloutRecorderParams { originator: String, base_instructions: BaseInstructions, dynamic_tools: Vec, + selected_capability_roots: Vec, multi_agent_version: Option, initial_window_id: Option, }, @@ -182,6 +184,7 @@ impl RolloutRecorderParams { originator, base_instructions, dynamic_tools, + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: None, } @@ -194,6 +197,20 @@ impl RolloutRecorderParams { self } + pub fn with_selected_capability_roots( + mut self, + selected_capability_roots: Vec, + ) -> Self { + if let Self::Create { + selected_capability_roots: roots, + .. + } = &mut self + { + *roots = selected_capability_roots; + } + self + } + pub fn with_multi_agent_version( mut self, multi_agent_version: Option, @@ -732,6 +749,7 @@ impl RolloutRecorder { originator, base_instructions, dynamic_tools, + selected_capability_roots, multi_agent_version, initial_window_id, } => { @@ -769,6 +787,7 @@ impl RolloutRecorder { } else { Some(dynamic_tools) }, + selected_capability_roots, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), multi_agent_version, context_window: initial_window_id.map(SessionContextWindow::new), diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 4b72ff7673..71ec252c23 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -101,6 +101,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index 7f57c2ba08..be39fb06ab 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -41,6 +41,7 @@ fn write_rollout_with_metadata(path: &Path, thread_id: ThreadId) -> std::io::Res model_provider: Some("test-provider".into()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 5712b65cee..bb341aa85a 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -174,6 +174,7 @@ fn write_rollout_with_user_message( model_provider: Some("test-provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index e6a3c728ef..901461cf56 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -1288,6 +1288,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> { model_provider: Some("test-provider".into()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index e5ae07839e..81de0cc39f 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -343,6 +343,7 @@ mod tests { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -535,6 +536,7 @@ mod tests { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index d1c2c90762..e5ab6fb31d 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -2148,6 +2148,7 @@ mod tests { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: Some("polluted".to_string()), multi_agent_version: None, context_window: None, @@ -2211,6 +2212,7 @@ mod tests { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index dfed2ffac2..d64510f223 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -126,6 +126,7 @@ mod tests { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { @@ -281,6 +282,7 @@ impl InMemoryThreadStore { model_provider: Some(params.metadata.model_provider.clone()), base_instructions: Some(params.base_instructions.clone()), dynamic_tools: (!params.dynamic_tools.is_empty()).then(|| params.dynamic_tools.clone()), + selected_capability_roots: params.selected_capability_roots.clone(), memory_mode: matches!(params.metadata.memory_mode, ThreadMemoryMode::Disabled) .then_some("disabled".to_string()), multi_agent_version: params.multi_agent_version, diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index eba6a5a6a9..451931b4dc 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -38,6 +38,7 @@ pub(super) async fn create_thread( params.dynamic_tools, ) .with_session_id(params.session_id) + .with_selected_capability_roots(params.selected_capability_roots) .with_multi_agent_version(params.multi_agent_version) .with_initial_window_id(params.initial_window_id), ) diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 51e622fc5b..717ae09b69 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1133,6 +1133,7 @@ mod tests { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: thread_metadata(), diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 1f761cea93..f96a4df31b 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -5,6 +5,7 @@ use chrono::DateTime; use chrono::Utc; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; use codex_protocol::models::PermissionProfile; @@ -85,6 +86,9 @@ pub struct CreateThreadParams { pub base_instructions: BaseInstructions, /// Dynamic tools available to the thread at startup. pub dynamic_tools: Vec, + /// Environment-qualified capability roots selected for this thread. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub selected_capability_roots: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, /// Initial context-window identity captured when the thread was created. From 2dec46e30ae81aa4844cc7a5474ff5c95c6998c0 Mon Sep 17 00:00:00 2001 From: richardopenai Date: Thu, 25 Jun 2026 11:02:11 -0700 Subject: [PATCH 07/33] [codex] Record exec-server lifecycle metrics (#27467) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Record bounded connection, request, and process lifecycle metrics. - Report active gauges from callbacks on every collection, including delta exports. - Serialize active-count updates so concurrent starts and finishes cannot publish stale values. - Serialize process exit, explicit termination, and shutdown through the process registry so exactly one completion result wins. - Keep the implementation small with single-owner RAII guards and one real OTLP/HTTP integration test using the existing `wiremock` dependency. ## Root cause Process exit and session shutdown previously used cloned completion state. That avoided duplicate emission, but it duplicated lifecycle ownership and made the ordering harder to reason about. The process registry mutex already defines the lifecycle ordering, so the final implementation stores the metric guard and termination flag directly on the process entry. Whichever path claims the entry first owns the completion result. Production metric export uses delta temporality. Event-only synchronous gauge recordings disappear after the next collection when no count changes, so active counts now use observable callbacks that report current state on every collection. The cleanup also removes the constant `result="accepted"` connection tag, redundant route and response assertions, a custom HTTP collector, and fallback initialization machinery that did not add behavior. ## Stack Review and land this stack in order: 1. #27466 — trace exec-server JSON-RPC requests 2. #27467 — record bounded connection, request, and process lifecycle metrics **(this PR)** 3. #27470 — observe remote registration and Noise rendezvous lifecycle ## Validation - `just test -p codex-exec-server --lib` (158 passed) - `just test -p codex-cli --test exec_server` (3 passed) - `just test -p codex-otel observable_gauge_is_collected_on_every_delta_snapshot` (1 passed) - `CARGO_BUILD_JOBS=1 just fix -p codex-otel -p codex-exec-server` - `just fmt` - `git diff --check` --- codex-rs/cli/src/exec_server_telemetry.rs | 19 +- codex-rs/cli/src/main.rs | 11 +- codex-rs/cli/tests/exec_server.rs | 233 +++++++++++++++ codex-rs/exec-server/src/lib.rs | 3 + codex-rs/exec-server/src/local_process.rs | 155 +++++++++- .../src/noise_relay/executor_stream.rs | 5 +- codex-rs/exec-server/src/remote.rs | 11 +- codex-rs/exec-server/src/rpc.rs | 6 +- codex-rs/exec-server/src/server.rs | 60 +++- .../exec-server/src/server/handler/tests.rs | 8 +- .../exec-server/src/server/process_handler.rs | 4 +- codex-rs/exec-server/src/server/processor.rs | 68 ++++- .../src/server/session_registry.rs | 15 +- codex-rs/exec-server/src/server/transport.rs | 38 ++- .../exec-server/src/server/transport_tests.rs | 1 + codex-rs/exec-server/src/telemetry.rs | 282 ++++++++++++++++++ codex-rs/otel/src/metrics/client.rs | 30 ++ codex-rs/otel/tests/suite/snapshot.rs | 32 ++ 18 files changed, 924 insertions(+), 57 deletions(-) create mode 100644 codex-rs/exec-server/src/telemetry.rs diff --git a/codex-rs/cli/src/exec_server_telemetry.rs b/codex-rs/cli/src/exec_server_telemetry.rs index 4d688a0159..1645bd5744 100644 --- a/codex-rs/cli/src/exec_server_telemetry.rs +++ b/codex-rs/cli/src/exec_server_telemetry.rs @@ -7,7 +7,7 @@ const OTEL_SERVICE_NAME: &str = "codex-exec-server"; pub(crate) fn init( config: Option<&codex_core::config::Config>, -) -> Result> { +) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) { let fmt_layer = tracing_subscriber::fmt::layer() .with_writer(std::io::stderr) .with_filter(stderr_env_filter()); @@ -17,21 +17,30 @@ pub(crate) fn init( env!("CARGO_PKG_VERSION"), Some(OTEL_SERVICE_NAME), DEFAULT_ANALYTICS_ENABLED, - ), - None => Ok(None), + ) + .unwrap_or_else(|error| { + eprintln!("Could not create otel exporter: {error}"); + None + }), + None => None, }; - let provider = otel.as_ref().ok().and_then(Option::as_ref); + let provider = otel.as_ref(); codex_core::otel_init::record_process_start(provider, OTEL_SERVICE_NAME); let otel_logger_layer = provider.and_then(|otel| otel.logger_layer()); let otel_tracing_layer = provider.and_then(|otel| otel.tracing_layer()); + let telemetry = provider + .and_then(|otel| otel.metrics()) + .cloned() + .map(codex_exec_server::ExecServerTelemetry::new) + .unwrap_or_default(); let _ = tracing_subscriber::registry() .with(fmt_layer) .with(otel_tracing_layer) .with(otel_logger_layer) .try_init(); tracing::callsite::rebuild_interest_cache(); - otel + (otel, telemetry) } fn stderr_env_filter() -> EnvFilter { diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 959ec3a3b3..79af9831df 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1693,9 +1693,7 @@ async fn run_exec_server_command( .environment_id .ok_or_else(|| anyhow::anyhow!("--environment-id is required when --remote is set"))?; let config = load_exec_server_config(root_config_overrides, strict_config).await?; - let _otel = exec_server_telemetry::init(Some(&config)) - .inspect_err(|err| eprintln!("Could not create otel exporter: {err}")) - .ok(); + let (_otel, telemetry) = exec_server_telemetry::init(Some(&config)); let auth_provider = load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth) .await?; @@ -1707,6 +1705,7 @@ async fn run_exec_server_command( if let Some(name) = cmd.name { remote_config.name = name; } + let remote_config = remote_config.with_telemetry(telemetry); codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; Ok(()) } else { @@ -1716,14 +1715,12 @@ async fn run_exec_server_command( } else { config_result.ok() }; - let _otel = exec_server_telemetry::init(config.as_ref()) - .inspect_err(|err| eprintln!("Could not create otel exporter: {err}")) - .ok(); + 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(listen_url, runtime_paths) + codex_exec_server::run_main_with_telemetry(listen_url, runtime_paths, telemetry) .await .map_err(anyhow::Error::from_boxed) } diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index 5e412a9de1..c54e9c40b2 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -1,9 +1,20 @@ use std::path::Path; +use std::process::Stdio; +use std::time::Duration; use anyhow::Result; use predicates::prelude::PredicateBooleanExt; use predicates::str::contains; use tempfile::TempDir; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; fn codex_command(codex_home: &Path) -> Result { let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); @@ -48,3 +59,225 @@ fn local_exec_server_ignores_invalid_config_without_strict_config() -> Result<() Ok(()) } + +#[tokio::test] +async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> { + let collector = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/metrics")) + .respond_with(ResponseTemplate::new(202)) + .mount(&collector) + .await; + let codex_home = TempDir::new()?; + let base_url = collector.uri(); + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +[analytics] +enabled = true + +[otel] +environment = "test" +metrics_exporter = {{ otlp-http = {{ endpoint = "{base_url}/v1/metrics", protocol = "json" }} }} +"# + ), + )?; + + let cwd = url::Url::from_directory_path(std::env::current_dir()?) + .map_err(|()| anyhow::anyhow!("could not convert cwd to file URL"))?; + #[cfg(windows)] + let argv = vec!["ping.exe", "-n", "61", "127.0.0.1"]; + #[cfg(not(windows))] + let argv = vec!["/bin/sleep", "60"]; + let codex_bin = codex_utils_cargo_bin::cargo_bin("codex")?; + let codex_home = codex_home.path().to_path_buf(); + let subprocess = async move { + let mut command = tokio::process::Command::new(codex_bin); + command + .env("CODEX_HOME", codex_home) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("no_proxy", "127.0.0.1,localhost") + .args(["exec-server", "--listen", "stdio"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .kill_on_drop(true); + let mut child = command.spawn()?; + let mut stdin = child + .stdin + .take() + .ok_or_else(|| anyhow::anyhow!("exec-server stdin was not piped"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("exec-server stdout was not piped"))?; + let mut stdout = BufReader::new(stdout); + send_json_line( + &mut stdin, + &serde_json::json!({ + "id": 1, + "method": "initialize", + "params": {"clientName": "otel-test", "resumeSessionId": null} + }), + ) + .await?; + wait_for_response(&mut stdout, /*expected_id*/ 1).await?; + send_json_line( + &mut stdin, + &serde_json::json!({"method": "initialized", "params": {}}), + ) + .await?; + send_json_line( + &mut stdin, + &serde_json::json!({ + "id": 2, + "method": "process/start", + "params": { + "processId": "otel-process", + "argv": argv, + "cwd": cwd, + "env": {}, + "tty": false, + "pipeStdin": false, + "arg0": null + } + }), + ) + .await?; + wait_for_response(&mut stdout, /*expected_id*/ 2).await?; + drop(stdin); + let mut remaining_stdout = String::new(); + stdout.read_to_string(&mut remaining_stdout).await?; + let status = child.wait().await?; + anyhow::ensure!( + status.success(), + "exec-server exited with {status}; remaining stdout: {remaining_stdout}" + ); + Ok::<(), anyhow::Error>(()) + }; + let subprocess_result = tokio::time::timeout(Duration::from_secs(30), subprocess) + .await + .map_err(|_| anyhow::anyhow!("exec-server subprocess timed out"))?; + subprocess_result?; + + let requests = collector + .received_requests() + .await + .ok_or_else(|| anyhow::anyhow!("failed to read OTLP collector requests"))?; + let metrics = requests + .iter() + .filter(|request| request.url.path() == "/v1/metrics") + .map(|request| serde_json::from_slice::(&request.body)) + .collect::>>()?; + assert_metric_point( + &metrics, + "exec_server_connections_active", + &[("transport", "stdio")], + Some(0), + ); + assert_metric_point( + &metrics, + "exec_server_connections_total", + &[("transport", "stdio")], + Some(1), + ); + assert_metric_point( + &metrics, + "exec_server_requests_total", + &[("method", "process/start"), ("result", "success")], + Some(1), + ); + assert_metric_point(&metrics, "exec_server_processes_active", &[], Some(0)); + assert_metric_point( + &metrics, + "exec_server_processes_finished_total", + &[("result", "terminated")], + Some(1), + ); + assert_metric_point( + &metrics, + "exec_server_request_duration_seconds", + &[("method", "process/start"), ("result", "success")], + /*value*/ None, + ); + assert_metric_point( + &metrics, + "exec_server_process_duration_seconds", + &[("result", "terminated")], + /*value*/ None, + ); + Ok(()) +} + +async fn send_json_line( + stdin: &mut (impl tokio::io::AsyncWrite + Unpin), + message: &serde_json::Value, +) -> Result<()> { + let mut encoded = serde_json::to_vec(message)?; + encoded.push(b'\n'); + stdin.write_all(&encoded).await?; + stdin.flush().await?; + Ok(()) +} + +async fn wait_for_response( + stdout: &mut (impl tokio::io::AsyncBufRead + Unpin), + expected_id: i64, +) -> Result<()> { + loop { + let mut line = String::new(); + if stdout.read_line(&mut line).await? == 0 { + anyhow::bail!("exec-server stdout closed before response {expected_id}"); + } + let message: serde_json::Value = serde_json::from_str(&line)?; + if message["id"].as_i64() == Some(expected_id) { + anyhow::ensure!( + message.get("error").is_none(), + "exec-server request {expected_id} failed: {message}" + ); + return Ok(()); + } + } +} + +fn assert_metric_point( + payloads: &[serde_json::Value], + name: &str, + attributes: &[(&str, &str)], + value: Option, +) { + let found = payloads + .iter() + .flat_map(|payload| payload["resourceMetrics"].as_array().into_iter().flatten()) + .flat_map(|resource| resource["scopeMetrics"].as_array().into_iter().flatten()) + .flat_map(|scope| scope["metrics"].as_array().into_iter().flatten()) + .filter(|metric| metric["name"].as_str() == Some(name)) + .flat_map(|metric| { + ["gauge", "sum", "histogram"] + .into_iter() + .find_map(|kind| metric[kind]["dataPoints"].as_array()) + .into_iter() + .flatten() + }) + .any(|point| { + let actual_attributes = point["attributes"] + .as_array() + .map(Vec::as_slice) + .unwrap_or_default(); + let attributes_match = actual_attributes.len() == attributes.len() + && attributes.iter().all(|(expected_key, expected_value)| { + actual_attributes.iter().any(|actual| { + actual["key"].as_str() == Some(*expected_key) + && actual["value"]["stringValue"].as_str() == Some(*expected_value) + }) + }); + let actual_value = point["asInt"] + .as_i64() + .or_else(|| point["asInt"].as_str()?.parse().ok()); + attributes_match && value.is_none_or(|expected| actual_value == Some(expected)) + }); + assert!( + found, + "metric {name} with attributes {attributes:?} and value {value:?} missing" + ); +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 92c7f85edf..87593374cb 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -27,6 +27,7 @@ mod rpc; mod runtime_paths; mod sandboxed_file_system; mod server; +mod telemetry; use codex_exec_server_protocol as protocol; @@ -149,3 +150,5 @@ pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; pub use server::run_main; +pub use server::run_main_with_telemetry; +pub use telemetry::ExecServerTelemetry; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 973019808a..7335bd1eed 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -59,6 +59,8 @@ use crate::rpc::RpcServerOutboundMessage; use crate::rpc::internal_error; use crate::rpc::invalid_params; use crate::rpc::invalid_request; +use crate::telemetry::ExecServerTelemetry; +use crate::telemetry::ProcessMetricGuard; const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; const NOTIFICATION_CHANNEL_CAPACITY: usize = 256; @@ -91,6 +93,8 @@ struct RunningProcess { output_notify: Arc, open_streams: usize, closed: bool, + metrics: Option, + termination_requested: bool, sandbox: SandboxType, sandbox_denied: bool, } @@ -136,6 +140,7 @@ enum ProcessEntry { struct Inner { notifications: std::sync::RwLock>, processes: Mutex>, + telemetry: ExecServerTelemetry, } #[derive(Clone)] @@ -166,24 +171,31 @@ impl LocalProcess { let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(NOTIFICATION_CHANNEL_CAPACITY); tokio::spawn(async move { while outgoing_rx.recv().await.is_some() {} }); - Self::with_runtime_paths(RpcNotificationSender::new(outgoing_tx), runtime_paths) + Self::with_runtime_paths( + RpcNotificationSender::new(outgoing_tx), + ExecServerTelemetry::default(), + runtime_paths, + ) } pub(crate) fn new( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: ExecServerRuntimePaths, ) -> Self { - Self::with_runtime_paths(notifications, Some(runtime_paths)) + Self::with_runtime_paths(notifications, telemetry, Some(runtime_paths)) } fn with_runtime_paths( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: Option, ) -> Self { Self { inner: Arc::new(Inner { notifications: std::sync::RwLock::new(Some(notifications)), processes: Mutex::new(HashMap::new()), + telemetry, }), runtime_paths, } @@ -200,7 +212,10 @@ impl LocalProcess { }) .collect::>() }; - for process in remaining { + for mut process in remaining { + if let Some(metrics) = process.metrics.take() { + metrics.finish("terminated"); + } process.session.terminate(); } } @@ -319,12 +334,13 @@ impl LocalProcess { output_notify: Arc::clone(&output_notify), open_streams: 2, closed: false, + metrics: Some(self.inner.telemetry.process_started()), + termination_requested: false, sandbox: prepared.sandbox, sandbox_denied: false, })), ); } - tokio::spawn(stream_output( process_id.clone(), if params.tty { @@ -536,11 +552,12 @@ impl LocalProcess { ) -> Result { let running = { let mut process_map = self.inner.processes.lock().await; - match process_map.get(¶ms.process_id) { + match process_map.get_mut(¶ms.process_id) { Some(ProcessEntry::Running(process)) => { if process.exit_code.is_some() { return Ok(TerminateResponse { running: false }); } + process.termination_requested = true; process.session.terminate(); true } @@ -807,11 +824,23 @@ async fn watch_exit( ) { let exit_code = exit_rx.await.unwrap_or(-1); let sandboxed = { - let processes = inner.processes.lock().await; - matches!( - processes.get(&process_id), - Some(ProcessEntry::Running(process)) if process.sandbox != SandboxType::None - ) + let mut processes = inner.processes.lock().await; + match processes.get_mut(&process_id) { + Some(ProcessEntry::Running(process)) => { + let sandboxed = process.sandbox != SandboxType::None; + if let Some(metrics) = process.metrics.take() { + metrics.finish(if process.termination_requested { + "terminated" + } else if exit_code == 0 { + "success" + } else { + "error" + }); + } + sandboxed + } + Some(ProcessEntry::Starting(_)) | None => false, + } }; if sandboxed { let _ = tokio::time::timeout(Duration::from_millis(20), output_notify.notified()).await; @@ -945,9 +974,13 @@ fn notification_sender(inner: &Inner) -> Option { #[cfg(test)] mod tests { use super::*; + use codex_otel::MetricsConfig; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_utils_path_uri::PathUri; use codex_utils_pty::ProcessDriver; + use opentelemetry_sdk::metrics::InMemoryMetricExporter; + use opentelemetry_sdk::metrics::data::AggregatedMetrics; + use opentelemetry_sdk::metrics::data::MetricData; use pretty_assertions::assert_eq; use tokio::sync::oneshot; use tokio::time::timeout; @@ -969,6 +1002,62 @@ mod tests { } } + fn telemetry_backend() -> ( + LocalProcess, + codex_otel::MetricsClient, + InMemoryMetricExporter, + ) { + let exporter = InMemoryMetricExporter::default(); + let metrics = codex_otel::MetricsClient::new(MetricsConfig::in_memory( + "test", + "codex-exec-server", + env!("CARGO_PKG_VERSION"), + exporter.clone(), + )) + .expect("metrics"); + let telemetry = ExecServerTelemetry::new(metrics.clone()); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(NOTIFICATION_CHANNEL_CAPACITY); + tokio::spawn(async move { while outgoing_rx.recv().await.is_some() {} }); + ( + LocalProcess::with_runtime_paths( + RpcNotificationSender::new(outgoing_tx), + telemetry, + /*runtime_paths*/ None, + ), + metrics, + exporter, + ) + } + + fn assert_finished_process_result( + metrics: codex_otel::MetricsClient, + exporter: &InMemoryMetricExporter, + expected: &str, + ) { + metrics.shutdown().expect("shutdown metrics"); + let resource_metrics = exporter + .get_finished_metrics() + .expect("finished metrics") + .into_iter() + .last() + .expect("metrics export"); + let finished_processes = resource_metrics + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == "exec_server_processes_finished_total") + .expect("finished process metric"); + let AggregatedMetrics::U64(MetricData::Sum(sum)) = finished_processes.data() else { + panic!("finished process metric should be a u64 sum"); + }; + let results = sum + .data_points() + .flat_map(opentelemetry_sdk::metrics::data::SumDataPoint::attributes) + .filter(|attribute| attribute.key.as_str() == "result") + .map(|attribute| attribute.value.as_str().into_owned()) + .collect::>(); + assert_eq!(results, vec![expected.to_string()]); + } + #[tokio::test] async fn start_process_rejects_non_native_cwd_before_launch() { #[cfg(unix)] @@ -1028,6 +1117,50 @@ mod tests { assert_eq!(child_env(¶ms), expected); } + #[tokio::test] + async fn exit_before_shutdown_records_success() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "exit-before-shutdown").await; + + process.exit(/*exit_code*/ 0); + let _ = read_process_until_change(&backend, &process.process_id, /*after_seq*/ None).await; + backend.shutdown().await; + + assert_finished_process_result(metrics, &exporter, "success"); + } + + #[tokio::test] + async fn termination_request_before_exit_records_terminated() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "terminate-before-exit").await; + + assert_eq!( + backend + .terminate_process(TerminateParams { + process_id: process.process_id.clone(), + }) + .await + .expect("terminate process"), + TerminateResponse { running: true }, + ); + process.exit(/*exit_code*/ 0); + let _ = read_process_until_change(&backend, &process.process_id, /*after_seq*/ None).await; + backend.shutdown().await; + + assert_finished_process_result(metrics, &exporter, "terminated"); + } + + #[tokio::test] + async fn shutdown_before_exit_records_terminated() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "shutdown-before-exit").await; + + backend.shutdown().await; + process.exit(/*exit_code*/ 0); + + assert_finished_process_result(metrics, &exporter, "terminated"); + } + #[tokio::test] async fn exited_process_retains_late_output_past_retention() { let backend = LocalProcess::default(); @@ -1170,6 +1303,8 @@ mod tests { output_notify: Arc::clone(&output_notify), open_streams: 2, closed: false, + metrics: Some(backend.inner.telemetry.process_started()), + termination_requested: false, sandbox: SandboxType::None, sandbox_denied: false, })), diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 57a05579dc..89f33d938a 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -27,6 +27,7 @@ use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; use crate::server::ConnectionProcessor; +use crate::telemetry::ConnectionTransport; /// Identifies one completed virtual-stream instance. /// @@ -169,7 +170,9 @@ pub(crate) fn spawn_noise_virtual_stream( transport: JsonRpcTransport::Plain, }; tokio::spawn(async move { - processor.run_connection(connection).await; + processor + .run_connection(connection, ConnectionTransport::Relay) + .await; let _ = processor_closed_stream_tx .send(ClosedNoiseVirtualStream { stream_id: processor_stream_id, diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 199e787542..6609b6cd92 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -25,6 +25,7 @@ use crate::EnvironmentRegistryRegistrationRequest; use crate::EnvironmentRegistryRegistrationResponse; use crate::ExecServerError; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; use crate::NoiseRendezvousConnectBundle; @@ -373,6 +374,7 @@ pub struct RemoteEnvironmentConfig { pub environment_id: String, pub name: String, auth_provider: SharedAuthProvider, + telemetry: ExecServerTelemetry, } impl std::fmt::Debug for RemoteEnvironmentConfig { @@ -398,8 +400,14 @@ impl RemoteEnvironmentConfig { environment_id, name: "codex-exec-server".to_string(), auth_provider, + telemetry: ExecServerTelemetry::default(), }) } + + pub fn with_telemetry(mut self, telemetry: ExecServerTelemetry) -> Self { + self.telemetry = telemetry; + self + } } /// Register an exec-server for remote use and serve requests over Noise. @@ -420,7 +428,8 @@ pub async fn run_remote_environment( ensure_rustls_crypto_provider(); let client = EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; - let processor = ConnectionProcessor::new(runtime_paths); + let processor = + ConnectionProcessor::new_with_telemetry(runtime_paths, config.telemetry.clone()); let identity = NoiseChannelIdentity::generate().map_err(|error| { ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}")) })?; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 4613c41f63..a941dce41f 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -212,8 +212,10 @@ where ); } - pub(crate) fn request_route(&self, method: &str) -> Option<&RequestRoute> { - self.request_routes.get(method) + pub(crate) fn request_route(&self, method: &str) -> Option<(&'static str, &RequestRoute)> { + self.request_routes + .get_key_value(method) + .map(|(&method, route)| (method, route)) } pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index d443b0ed1f..7988e0870b 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -12,15 +12,71 @@ pub use transport::DEFAULT_LISTEN_URL; pub use transport::ExecServerListenUrlParseError; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; + +pub async fn run_main( + listen_url: &str, + runtime_paths: ExecServerRuntimePaths, +) -> Result<(), Box> { + run_main_with_telemetry(listen_url, runtime_paths, ExecServerTelemetry::default()).await +} #[tracing::instrument( name = "codex.exec_server", skip_all, fields(otel.kind = "internal") )] -pub async fn run_main( +pub async fn run_main_with_telemetry( listen_url: &str, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { - transport::run_transport(listen_url, runtime_paths).await + transport::run_transport(listen_url, runtime_paths, telemetry).await +} + +#[cfg(test)] +mod tests { + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::InMemorySpanExporter; + use opentelemetry_sdk::trace::SdkTracerProvider; + use tracing::instrument::WithSubscriber; + use tracing_subscriber::prelude::*; + + use super::run_main_with_telemetry; + use crate::ExecServerRuntimePaths; + use crate::ExecServerTelemetry; + + #[tokio::test] + async fn telemetry_entrypoint_emits_root_span() { + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("exec-server-test"))); + + async { + tracing::callsite::rebuild_interest_cache(); + run_main_with_telemetry( + "invalid", + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current executable"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"), + ExecServerTelemetry::default(), + ) + .await + .expect_err("invalid listen URL should fail"); + } + .with_subscriber(subscriber) + .await; + + provider.force_flush().expect("flush traces"); + let spans = exporter.get_finished_spans().expect("span export"); + assert!( + spans.iter().any(|span| span.name == "codex.exec_server"), + "root exec-server span missing: {spans:?}" + ); + } } diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 2cea99b214..47416039c6 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -82,7 +82,7 @@ fn test_runtime_paths() -> ExecServerRuntimePaths { async fn initialized_handler() -> Arc { let (outgoing_tx, _outgoing_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let handler = Arc::new(ExecServerHandler::new( registry, RpcNotificationSender::new(outgoing_tx), @@ -160,7 +160,7 @@ async fn terminate_reports_false_after_process_exit() { #[tokio::test] async fn long_poll_read_fails_after_session_resume() { let (first_tx, _first_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), RpcNotificationSender::new(first_tx), @@ -233,7 +233,7 @@ async fn long_poll_read_fails_after_session_resume() { #[tokio::test] async fn active_session_resume_is_rejected() { let (first_tx, _first_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), RpcNotificationSender::new(first_tx), @@ -277,7 +277,7 @@ async fn active_session_resume_is_rejected() { async fn output_and_exit_are_retained_after_notification_receiver_closes() { let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let handler = Arc::new(ExecServerHandler::new( - SessionRegistry::new(), + SessionRegistry::new(crate::ExecServerTelemetry::default()), RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), )); diff --git a/codex-rs/exec-server/src/server/process_handler.rs b/codex-rs/exec-server/src/server/process_handler.rs index 9ff6690411..4d933afbe1 100644 --- a/codex-rs/exec-server/src/server/process_handler.rs +++ b/codex-rs/exec-server/src/server/process_handler.rs @@ -13,6 +13,7 @@ use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::rpc::RpcNotificationSender; +use crate::telemetry::ExecServerTelemetry; #[derive(Clone)] pub(crate) struct ProcessHandler { @@ -22,10 +23,11 @@ pub(crate) struct ProcessHandler { impl ProcessHandler { pub(crate) fn new( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: ExecServerRuntimePaths, ) -> Self { Self { - process: LocalProcess::new(notifications, runtime_paths), + process: LocalProcess::new(notifications, telemetry, runtime_paths), } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 3eedfb2f20..63f6bd45d9 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::Instant; use tokio::sync::mpsc; use tracing::Instrument; @@ -17,36 +18,61 @@ use crate::rpc::method_not_found; use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; +use crate::telemetry::ConnectionTransport; +use crate::telemetry::ExecServerTelemetry; #[derive(Clone)] pub(crate) struct ConnectionProcessor { session_registry: Arc, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, } impl ConnectionProcessor { + #[cfg(test)] pub(crate) fn new(runtime_paths: ExecServerRuntimePaths) -> Self { + Self::new_with_telemetry(runtime_paths, ExecServerTelemetry::default()) + } + + pub(crate) fn new_with_telemetry( + runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, + ) -> Self { Self { - session_registry: SessionRegistry::new(), + session_registry: SessionRegistry::new(telemetry.clone()), runtime_paths, + telemetry, } } - pub(crate) async fn run_connection(&self, connection: JsonRpcConnection) { + pub(crate) async fn run_connection( + &self, + connection: JsonRpcConnection, + transport: ConnectionTransport, + ) { run_connection( connection, Arc::clone(&self.session_registry), self.runtime_paths.clone(), + self.telemetry.clone(), + transport, ) .await; } + + pub(crate) async fn shutdown(&self) { + self.session_registry.shutdown().await; + } } async fn run_connection( connection: JsonRpcConnection, session_registry: Arc, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, + transport: ConnectionTransport, ) { + let _connection_metrics = telemetry.connection_started(transport); let router = Arc::new(build_router()); let JsonRpcConnection { outgoing_tx: json_outgoing_tx, @@ -101,12 +127,18 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_exec_server_protocol::JSONRPCMessage::Request(request) => { - if let Some(route) = router.request_route(request.method.as_str()) { - let request_span = request_span(request.method.as_str(), &request); + let request_started_at = Instant::now(); + if let Some((method, route)) = router.request_route(request.method.as_str()) { + let request_span = request_span(method, &request); let message = tokio::select! { message = route(Arc::clone(&handler), request).instrument(request_span.clone()) => message, _ = disconnected_rx.changed() => { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); debug!("exec-server transport disconnected while handling request"); break; } @@ -116,11 +148,19 @@ async fn run_connection( && outgoing_tx.send(message).await.is_err() { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); break; } request_span.record("result", result); + telemetry.request_completed(method, result, request_started_at.elapsed()); + drop(request_span); } else { - let request_span = request_span("unknown", &request); + let method = "unknown"; + let request_span = request_span(method, &request); if outgoing_tx .send(RpcServerOutboundMessage::Error { request_id: request.id, @@ -133,9 +173,15 @@ async fn run_connection( .is_err() { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); break; } request_span.record("result", "error"); + telemetry.request_completed(method, "error", request_started_at.elapsed()); } } codex_exec_server_protocol::JSONRPCMessage::Notification(notification) => { @@ -330,7 +376,7 @@ mod tests { #[tokio::test] async fn connection_accepts_pipelined_scalar_requests() { - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let (mut writer, mut lines, task) = spawn_test_connection(registry, "pipelined-scalar"); send_request( @@ -362,7 +408,7 @@ mod tests { #[tokio::test] async fn transport_disconnect_detaches_session_during_in_flight_read() { - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let (mut first_writer, mut first_lines, first_task) = spawn_test_connection(Arc::clone(®istry), "first"); @@ -458,7 +504,13 @@ mod tests { let (server_writer, client_reader) = duplex(1 << 20); let connection = JsonRpcConnection::from_stdio(server_reader, server_writer, label.to_string()); - let task = tokio::spawn(run_connection(connection, registry, test_runtime_paths())); + let task = tokio::spawn(run_connection( + connection, + registry, + test_runtime_paths(), + crate::ExecServerTelemetry::default(), + crate::telemetry::ConnectionTransport::Stdio, + )); (client_writer, BufReader::new(client_reader).lines(), task) } diff --git a/codex-rs/exec-server/src/server/session_registry.rs b/codex-rs/exec-server/src/server/session_registry.rs index 4d2ef7f893..89a861ab7b 100644 --- a/codex-rs/exec-server/src/server/session_registry.rs +++ b/codex-rs/exec-server/src/server/session_registry.rs @@ -12,6 +12,7 @@ use crate::rpc::RpcNotificationSender; use crate::rpc::invalid_request; use crate::rpc::session_already_attached; use crate::server::process_handler::ProcessHandler; +use crate::telemetry::ExecServerTelemetry; #[cfg(test)] const DETACHED_SESSION_TTL: Duration = Duration::from_millis(200); @@ -20,6 +21,7 @@ const DETACHED_SESSION_TTL: Duration = Duration::from_secs(30); pub(crate) struct SessionRegistry { sessions: Mutex>>, + telemetry: ExecServerTelemetry, } struct SessionEntry { @@ -51,9 +53,10 @@ pub(crate) struct SessionHandle { } impl SessionRegistry { - pub(crate) fn new() -> Arc { + pub(crate) fn new(telemetry: ExecServerTelemetry) -> Arc { Arc::new(Self { sessions: Mutex::new(HashMap::new()), + telemetry, }) } @@ -97,7 +100,7 @@ impl SessionRegistry { let session_id = Uuid::new_v4().to_string(); let entry = Arc::new(SessionEntry::new( session_id.clone(), - ProcessHandler::new(notifications, runtime_paths), + ProcessHandler::new(notifications, self.telemetry.clone(), runtime_paths), connection_id, )); sessions.insert(session_id, Arc::clone(&entry)); @@ -119,6 +122,13 @@ impl SessionRegistry { }) } + pub(crate) async fn shutdown(&self) { + let sessions = std::mem::take(&mut *self.sessions.lock().await); + for entry in sessions.into_values() { + entry.process.shutdown().await; + } + } + async fn expire_if_detached(&self, session_id: String, connection_id: ConnectionId) { tokio::time::sleep(DETACHED_SESSION_TTL).await; @@ -143,6 +153,7 @@ impl Default for SessionRegistry { fn default() -> Self { Self { sessions: Mutex::new(HashMap::new()), + telemetry: ExecServerTelemetry::default(), } } } diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs index 57d7048e2b..2affab0070 100644 --- a/codex-rs/exec-server/src/server/transport.rs +++ b/codex-rs/exec-server/src/server/transport.rs @@ -22,8 +22,10 @@ use tracing::info; use tracing::warn; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; use crate::connection::JsonRpcConnection; use crate::server::processor::ConnectionProcessor; +use crate::telemetry::ConnectionTransport; pub const DEFAULT_LISTEN_URL: &str = "ws://127.0.0.1:0"; @@ -80,49 +82,54 @@ pub(crate) fn parse_listen_url( pub(crate) async fn run_transport( listen_url: &str, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { match parse_listen_url(listen_url)? { ExecServerListenTransport::WebSocket(bind_address) => { - run_websocket_listener(bind_address, runtime_paths).await + run_websocket_listener(bind_address, runtime_paths, telemetry).await } - ExecServerListenTransport::Stdio => run_stdio_connection(runtime_paths).await, + ExecServerListenTransport::Stdio => run_stdio_connection(runtime_paths, telemetry).await, } } async fn run_stdio_connection( runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { - run_stdio_connection_with_io(io::stdin(), io::stdout(), runtime_paths).await + run_stdio_connection_with_io(io::stdin(), io::stdout(), runtime_paths, telemetry).await } async fn run_stdio_connection_with_io( reader: R, writer: W, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, { - let processor = ConnectionProcessor::new(runtime_paths); + let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, telemetry); tracing::info!("codex-exec-server listening on stdio"); processor - .run_connection(JsonRpcConnection::from_stdio( - reader, - writer, - "exec-server stdio".to_string(), - )) + .run_connection( + JsonRpcConnection::from_stdio(reader, writer, "exec-server stdio".to_string()), + ConnectionTransport::Stdio, + ) .await; + // Stdio serves exactly one connection, so detached sessions cannot be resumed. + processor.shutdown().await; Ok(()) } async fn run_websocket_listener( bind_address: SocketAddr, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; - let processor = ConnectionProcessor::new(runtime_paths); + let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, telemetry); info!("codex-exec-server listening on ws://{local_addr}"); println!("ws://{local_addr}"); std::io::stdout().flush()?; @@ -174,10 +181,13 @@ async fn websocket_upgrade_handler( websocket.on_upgrade(move |stream| async move { state .processor - .run_connection(JsonRpcConnection::from_axum_websocket( - stream, - format!("exec-server websocket {peer_addr}"), - )) + .run_connection( + JsonRpcConnection::from_axum_websocket( + stream, + format!("exec-server websocket {peer_addr}"), + ), + ConnectionTransport::WebSocket, + ) .await; }) } diff --git a/codex-rs/exec-server/src/server/transport_tests.rs b/codex-rs/exec-server/src/server/transport_tests.rs index a8db4250bf..3016effa6c 100644 --- a/codex-rs/exec-server/src/server/transport_tests.rs +++ b/codex-rs/exec-server/src/server/transport_tests.rs @@ -61,6 +61,7 @@ async fn stdio_listen_transport_serves_initialize() { server_reader, server_writer, test_runtime_paths(), + crate::ExecServerTelemetry::default(), )); let mut client_lines = BufReader::new(client_reader).lines(); diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs new file mode 100644 index 0000000000..4624a2ff2c --- /dev/null +++ b/codex-rs/exec-server/src/telemetry.rs @@ -0,0 +1,282 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; + +use codex_otel::MetricsClient; +use tracing::warn; + +const CONNECTIONS_ACTIVE_METRIC: &str = "exec_server_connections_active"; +const CONNECTIONS_ACTIVE_DESCRIPTION: &str = "Number of active exec-server connections."; +const CONNECTIONS_TOTAL_METRIC: &str = "exec_server_connections_total"; +const CONNECTIONS_TOTAL_DESCRIPTION: &str = "Total number of accepted exec-server connections."; +const REQUESTS_TOTAL_METRIC: &str = "exec_server_requests_total"; +const REQUESTS_TOTAL_DESCRIPTION: &str = "Total number of exec-server requests."; +const REQUEST_DURATION_METRIC: &str = "exec_server_request_duration_seconds"; +const REQUEST_DURATION_DESCRIPTION: &str = "Duration of exec-server requests in seconds."; +const PROCESSES_ACTIVE_METRIC: &str = "exec_server_processes_active"; +const PROCESSES_ACTIVE_DESCRIPTION: &str = "Number of active exec-server processes."; +const PROCESSES_FINISHED_TOTAL_METRIC: &str = "exec_server_processes_finished_total"; +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."; + +#[derive(Clone, Copy)] +pub(crate) enum ConnectionTransport { + Relay, + Stdio, + WebSocket, +} + +impl ConnectionTransport { + fn metric_tag(self) -> &'static str { + match self { + Self::Relay => "relay", + Self::Stdio => "stdio", + Self::WebSocket => "websocket", + } + } +} + +#[derive(Clone, Default)] +pub struct ExecServerTelemetry { + inner: Option>, +} + +struct ExecServerTelemetryInner { + metrics: MetricsClient, + active: Arc>, +} + +#[derive(Default)] +struct ActiveCounts { + relay_connections: i64, + stdio_connections: i64, + websocket_connections: i64, + processes: i64, +} + +impl ActiveCounts { + fn connections(&self, transport: ConnectionTransport) -> i64 { + match transport { + ConnectionTransport::Relay => self.relay_connections, + ConnectionTransport::Stdio => self.stdio_connections, + ConnectionTransport::WebSocket => self.websocket_connections, + } + } +} + +pub(crate) struct ConnectionMetricGuard { + telemetry: ExecServerTelemetry, + transport: ConnectionTransport, +} + +pub(crate) struct ProcessMetricGuard { + telemetry: ExecServerTelemetry, + started_at: Instant, + result: &'static str, +} + +impl ExecServerTelemetry { + pub fn new(metrics: MetricsClient) -> Self { + let active = Arc::new(Mutex::new(ActiveCounts::default())); + register_active_gauges(&metrics, &active); + Self { + inner: Some(Arc::new(ExecServerTelemetryInner { metrics, active })), + } + } + + pub(crate) fn connection_started( + &self, + transport: ConnectionTransport, + ) -> ConnectionMetricGuard { + self.with_inner(|inner| { + inner.adjust_connection_count(transport, /*delta*/ 1); + inner.counter( + CONNECTIONS_TOTAL_METRIC, + CONNECTIONS_TOTAL_DESCRIPTION, + &[("transport", transport.metric_tag())], + ); + }); + ConnectionMetricGuard { + telemetry: self.clone(), + transport, + } + } + + pub(crate) fn request_completed( + &self, + method: &'static str, + result: &'static str, + duration: Duration, + ) { + self.with_inner(|inner| { + let tags = [("method", method), ("result", result)]; + inner.counter(REQUESTS_TOTAL_METRIC, REQUESTS_TOTAL_DESCRIPTION, &tags); + inner.duration( + REQUEST_DURATION_METRIC, + REQUEST_DURATION_DESCRIPTION, + duration, + &tags, + ); + }); + } + + pub(crate) fn process_started(&self) -> ProcessMetricGuard { + self.with_inner(|inner| { + inner.adjust_process_count(/*delta*/ 1); + }); + ProcessMetricGuard { + telemetry: self.clone(), + started_at: Instant::now(), + result: "unknown", + } + } + + fn process_finished(&self, result: &'static str, duration: Duration) { + self.with_inner(|inner| { + inner.adjust_process_count(/*delta*/ -1); + inner.counter( + PROCESSES_FINISHED_TOTAL_METRIC, + PROCESSES_FINISHED_TOTAL_DESCRIPTION, + &[("result", result)], + ); + inner.duration( + PROCESS_DURATION_METRIC, + PROCESS_DURATION_DESCRIPTION, + duration, + &[("result", result)], + ); + }); + } + + fn connection_finished(&self, transport: ConnectionTransport) { + self.with_inner(|inner| { + inner.adjust_connection_count(transport, /*delta*/ -1); + }); + } + + fn with_inner(&self, emit: impl FnOnce(&ExecServerTelemetryInner)) { + if let Some(inner) = &self.inner { + emit(inner); + } + } +} + +impl Drop for ConnectionMetricGuard { + fn drop(&mut self) { + self.telemetry.connection_finished(self.transport); + } +} + +impl ProcessMetricGuard { + pub(crate) fn finish(mut self, result: &'static str) { + self.result = result; + } +} + +impl Drop for ProcessMetricGuard { + fn drop(&mut self) { + self.telemetry + .process_finished(self.result, self.started_at.elapsed()); + } +} + +impl ExecServerTelemetryInner { + fn active_counts(&self) -> std::sync::MutexGuard<'_, ActiveCounts> { + // These are independent integer counts, so a panic cannot leave a cross-field invariant + // half-updated. Recovering a poisoned lock preserves the last completed count update. + self.active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + fn adjust_connection_count(&self, transport: ConnectionTransport, delta: i64) { + let mut active = self.active_counts(); + let count = match transport { + ConnectionTransport::Relay => &mut active.relay_connections, + ConnectionTransport::Stdio => &mut active.stdio_connections, + ConnectionTransport::WebSocket => &mut active.websocket_connections, + }; + *count += delta; + } + + fn adjust_process_count(&self, delta: i64) { + let mut active = self.active_counts(); + active.processes += delta; + } + + fn counter(&self, name: &str, description: &str, tags: &[(&str, &str)]) { + if self + .metrics + .counter_with_description(name, description, /*inc*/ 1, tags) + .is_err() + { + warn!(metric = name, "failed to emit exec-server counter"); + } + } + + fn duration(&self, name: &str, description: &str, duration: Duration, tags: &[(&str, &str)]) { + if self + .metrics + .record_duration_seconds_with_description(name, description, duration, tags) + .is_err() + { + warn!(metric = name, "failed to emit exec-server duration"); + } + } +} + +fn register_active_gauges(metrics: &MetricsClient, active: &Arc>) { + for transport in [ + ConnectionTransport::Relay, + ConnectionTransport::Stdio, + ConnectionTransport::WebSocket, + ] { + register_active_gauge( + metrics, + active, + CONNECTIONS_ACTIVE_METRIC, + CONNECTIONS_ACTIVE_DESCRIPTION, + &[("transport", transport.metric_tag())], + move |active| active.connections(transport), + ); + } + + register_active_gauge( + metrics, + active, + PROCESSES_ACTIVE_METRIC, + PROCESSES_ACTIVE_DESCRIPTION, + &[], + |active| active.processes, + ); +} + +fn register_active_gauge( + metrics: &MetricsClient, + active: &Arc>, + name: &str, + description: &str, + tags: &[(&str, &str)], + read: impl Fn(&ActiveCounts) -> i64 + Send + Sync + 'static, +) { + let active = Arc::clone(active); + if metrics + .register_observable_gauge_with_description( + name, + description, + move || { + let active = active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + read(&active) + }, + tags, + ) + .is_err() + { + warn!(metric = name, "failed to register exec-server gauge"); + } +} diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index 73fc3c6613..251fc74d77 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -188,6 +188,24 @@ impl MetricsClientInner { Ok(()) } + fn register_observable_gauge( + &self, + name: &str, + description: &str, + observe: impl Fn() -> i64 + Send + Sync + 'static, + tags: &[(&str, &str)], + ) -> Result<()> { + validate_metric_name(name)?; + let attributes = self.attributes(tags)?; + let _gauge = self + .meter + .i64_observable_gauge(name.to_string()) + .with_description(description.to_string()) + .with_callback(move |observer| observer.observe(observe(), &attributes)) + .build(); + Ok(()) + } + fn duration_histogram( &self, name: &str, @@ -354,6 +372,18 @@ impl MetricsClient { self.0.gauge(name, Some(description), value, tags) } + /// Register a gauge callback that reports the current value on every collection. + pub fn register_observable_gauge_with_description( + &self, + name: &str, + description: &str, + observe: impl Fn() -> i64 + Send + Sync + 'static, + tags: &[(&str, &str)], + ) -> Result<()> { + self.0 + .register_observable_gauge(name, description, observe, tags) + } + /// Record a duration in milliseconds using a histogram. pub fn record_duration( &self, diff --git a/codex-rs/otel/tests/suite/snapshot.rs b/codex-rs/otel/tests/suite/snapshot.rs index e631100004..ffc2458a04 100644 --- a/codex-rs/otel/tests/suite/snapshot.rs +++ b/codex-rs/otel/tests/suite/snapshot.rs @@ -62,6 +62,38 @@ fn snapshot_collects_metrics_without_shutdown() -> Result<()> { Ok(()) } +#[test] +fn observable_gauge_is_collected_on_every_delta_snapshot() -> Result<()> { + let exporter = InMemoryMetricExporter::default(); + let config = MetricsConfig::in_memory("test", "codex-cli", env!("CARGO_PKG_VERSION"), exporter) + .with_runtime_reader(); + let metrics = MetricsClient::new(config)?; + metrics.register_observable_gauge_with_description( + "codex.active", + "Number of active operations.", + || 1, + &[("component", "test")], + )?; + + for snapshot in [metrics.snapshot()?, metrics.snapshot()?] { + let gauge = find_metric(&snapshot, "codex.active").expect("gauge metric missing"); + let point = match gauge.data() { + AggregatedMetrics::I64(MetricData::Gauge(gauge)) => { + gauge.data_points().next().expect("gauge point") + } + _ => panic!("unexpected gauge metric data type"), + }; + assert_eq!(point.value(), 1); + assert_eq!( + attributes_to_map(point.attributes()), + BTreeMap::from([("component".to_string(), "test".to_string())]) + ); + } + + metrics.shutdown()?; + Ok(()) +} + #[test] fn manager_snapshot_metrics_collects_without_shutdown() -> Result<()> { let exporter = InMemoryMetricExporter::default(); From 6d9dbacf1af69fc9b633132e34d63be202556a24 Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Thu, 25 Jun 2026 11:24:34 -0700 Subject: [PATCH 08/33] feat: add provider-aware model fallback to thread start (#29942) ## Why Helper threads such as task title generation can request a model ID that is valid for the default OpenAI provider but unavailable from the active provider. With Amazon Bedrock, `gpt-5.4-mini` is rejected while the provider static catalog exposes Bedrock model IDs such as `openai.gpt-5.5` and `openai.gpt-5.4`. This causes repeated background 404s and can surface a misleading turn error even when the main turn succeeds. Clients need an explicit way to ask app-server to resolve an unavailable helper model to the active provider default. That fallback must remain limited to providers with an authoritative static catalog so custom or dynamically discovered model IDs are not rewritten based on an incomplete catalog. Fixes #28741. ## What changed - Add the experimental `allowProviderModelFallback` option to `thread/start`, defaulting to `false` to preserve existing behavior. - Thread the option through thread creation and model selection. - When enabled for a static model manager, preserve requested models present in the catalog and replace unavailable models with the provider default. - Continue preserving explicit model IDs for dynamic model managers without fetching a catalog solely to validate them. - Document the new `thread/start` behavior in the app-server API overview. ## Test Temporary test-client harness: ``` ThreadStartParams { model: Some("gpt-5.4-mini".to_string()), allow_provider_model_fallback: true, ..Default::default() } ``` Command: ``` CODEX_HOME=/tmp/codex-bedrock-thread-start-home \ CODEX_E2E_BEDROCK_THREAD_START_ONLY=1 \ ./target/debug/codex-app-server-test-client \ --codex-bin ./target/debug/codex \ -c 'model_provider="amazon-bedrock"' \ send-message-v2 --experimental-api ignored ``` Relevant output: ``` > "method": "thread/start", > "params": { > "model": "gpt-5.4-mini", > "modelProvider": null, > "allowProviderModelFallback": true, > ... > } < "result": { < "model": "openai.gpt-5.5", < "modelProvider": "amazon-bedrock", < ... < } ``` --- .../src/protocol/v2/thread.rs | 5 + codex-rs/app-server/README.md | 2 +- .../external_agent_session_import.rs | 6 +- .../request_processors/thread_processor.rs | 4 + .../app-server/tests/suite/v2/skills_list.rs | 1 + .../app-server/tests/suite/v2/thread_start.rs | 109 ++++++++++++++++++ codex-rs/core/src/agent/control_tests.rs | 2 + codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 19 ++- .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/thread_manager.rs | 11 ++ codex-rs/core/src/thread_manager_tests.rs | 6 + codex-rs/core/tests/common/test_codex.rs | 1 + codex-rs/core/tests/suite/agents_md.rs | 2 + codex-rs/core/tests/suite/remote_models.rs | 12 +- .../tests/suite/subagent_notifications.rs | 1 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/models-manager/src/manager.rs | 51 +++++++- codex-rs/models-manager/src/manager_tests.rs | 99 ++++++++++++++++ 19 files changed, 327 insertions(+), 7 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 0be7776279..e712ab0ad3 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -57,6 +57,11 @@ pub struct ThreadStartParams { pub model: Option, #[ts(optional = nullable)] pub model_provider: Option, + /// Allow a provider with an authoritative static model catalog to replace an unavailable + /// requested model with its default. + #[experimental("thread/start.allowProviderModelFallback")] + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub allow_provider_model_fallback: bool, #[serde( default, deserialize_with = "crate::protocol::serde_helpers::deserialize_double_option", diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 4c88d4b69d..1a3684a024 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -137,7 +137,7 @@ Example with notification opt-out: ## API Overview -- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are started in that environment, and HTTP MCP connections use that environment's HTTP client. +- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `allowProviderModelFallback` lets providers backed by an authoritative static model catalog replace an unavailable requested `model` with the catalog default; dynamic or cached catalogs preserve the requested model. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are started in that environment, and HTTP MCP connections use that environment's HTTP client. - `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`. - `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. - `thread/start`, `thread/resume`, and `thread/fork` responses include the legacy `sandbox` compatibility projection. `instructionSources` lists loaded instruction files using each source environment's native absolute path syntax, including files loaded from remote environments. Experimental clients can read `runtimeWorkspaceRoots` for the thread-scoped runtime roots and `activePermissionProfile` for the named or implicit built-in profile identity/provenance when known. Their deprecated experimental `multiAgentMode` field, and the corresponding thread setting, always report `explicitRequestOnly`; Ultra reasoning effort is the source of proactive multi-agent behavior. diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index bc0a64ef72..15b6e3a43a 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -184,7 +184,11 @@ impl ExternalAgentSessionImporter { .map_err(|err| format!("failed to load imported session config: {err}"))?; let models_manager = self.thread_manager.get_models_manager(); let model = models_manager - .get_default_model(&config.model, RefreshStrategy::Offline) + .get_default_model( + &config.model, + /*allow_provider_model_fallback*/ false, + RefreshStrategy::Offline, + ) .await; let model_info = models_manager .get_model_info(model.as_str(), &config.to_models_manager_config()) diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index ab4e241867..7c6b0aad74 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -903,6 +903,7 @@ impl ThreadRequestProcessor { let ThreadStartParams { model, model_provider, + allow_provider_model_fallback, service_tier, cwd, runtime_workspace_roots, @@ -979,6 +980,7 @@ impl ThreadRequestProcessor { thread_source.map(Into::into), environment_selections, service_name, + allow_provider_model_fallback, experimental_raw_events, request_trace, ) @@ -1053,6 +1055,7 @@ impl ThreadRequestProcessor { thread_source: Option, environments: Option>, service_name: Option, + allow_provider_model_fallback: bool, experimental_raw_events: bool, request_trace: Option, ) -> Result<(), JSONRPCErrorError> { @@ -1160,6 +1163,7 @@ impl ThreadRequestProcessor { .thread_manager .start_thread_with_options(StartThreadOptions { config, + allow_provider_model_fallback, initial_history: match session_start_source .unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup) { diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index 075cc007ac..64b7fe0504 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -882,6 +882,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( .send_thread_start_request(ThreadStartParams { model: None, model_provider: None, + allow_provider_model_fallback: false, service_tier: None, cwd: None, runtime_workspace_roots: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index 39f8935a93..93dcb47ede 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -54,6 +54,115 @@ use super::analytics::wait_for_analytics_payload; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const INVALID_REQUEST_ERROR_CODE: i64 = -32600; +async fn start_thread_with_model( + mcp: &mut TestAppServer, + model: &str, + allow_provider_model_fallback: bool, +) -> Result { + let request_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + model: Some(model.to_string()), + allow_provider_model_fallback, + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +#[tokio::test] +async fn thread_start_provider_model_fallback_applies_to_configured_model() -> Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + r#"model_provider = "amazon-bedrock" +model = "gpt-5.4-mini" +"#, + )?; + let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + allow_provider_model_fallback: true, + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: ThreadStartResponse = to_response(response)?; + + assert_eq!(response.model, "openai.gpt-5.5"); + Ok(()) +} + +#[tokio::test] +async fn thread_start_provider_model_fallback_uses_bedrock_static_catalog() -> Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + r#"model_provider = "amazon-bedrock" +"#, + )?; + let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let unsupported_with_fallback = start_thread_with_model( + &mut mcp, + "gpt-5.4-mini", + /*allow_provider_model_fallback*/ true, + ) + .await?; + let supported_with_fallback = start_thread_with_model( + &mut mcp, + "openai.gpt-5.4", + /*allow_provider_model_fallback*/ true, + ) + .await?; + let unsupported_without_fallback = start_thread_with_model( + &mut mcp, + "gpt-5.4-mini", + /*allow_provider_model_fallback*/ false, + ) + .await?; + + assert_eq!( + vec![ + unsupported_with_fallback.model, + supported_with_fallback.model, + unsupported_without_fallback.model, + ], + vec!["openai.gpt-5.5", "openai.gpt-5.4", "gpt-5.4-mini"] + ); + Ok(()) +} + +#[tokio::test] +async fn thread_start_provider_model_fallback_ignores_dynamic_catalog() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?; + let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let response = start_thread_with_model( + &mut mcp, + "unlisted-dynamic-model", + /*allow_provider_model_fallback*/ true, + ) + .await?; + + assert_eq!(response.model, "unlisted-dynamic-model"); + Ok(()) +} + #[tokio::test] async fn thread_start_creates_thread_and_emits_started() -> Result<()> { // Provide a mock server and config so model wiring is valid. diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index de8754ad92..acb7af3779 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1462,6 +1462,7 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li .manager .start_thread_with_options(StartThreadOptions { config: harness.config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -2250,6 +2251,7 @@ async fn spawn_thread_subagents_persist_parent_originator_across_new_and_truncat .manager .start_thread_with_options(StartThreadOptions { config: harness.config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 6fbd195fce..21e8de83ca 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -86,6 +86,7 @@ pub(crate) async fn run_codex_thread_interactive( }; let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { config, + allow_provider_model_fallback: false, user_instructions, installation_id: parent_session.installation_id.clone(), auth_manager, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 94bef646aa..442722003b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -414,6 +414,7 @@ pub struct CodexSpawnOk { pub(crate) struct CodexSpawnArgs { pub(crate) config: Config, + pub(crate) allow_provider_model_fallback: bool, pub(crate) user_instructions: LoadedUserInstructions, pub(crate) installation_id: String, pub(crate) auth_manager: Arc, @@ -503,6 +504,7 @@ impl Codex { async fn spawn_internal(args: CodexSpawnArgs) -> CodexResult { let CodexSpawnArgs { mut config, + allow_provider_model_fallback, user_instructions, installation_id, auth_manager, @@ -576,8 +578,23 @@ impl Codex { let _ = models_manager.list_models(refresh_strategy).await; } let model = models_manager - .get_default_model(&config.model, refresh_strategy) + .get_default_model( + &config.model, + allow_provider_model_fallback, + refresh_strategy, + ) .await; + if allow_provider_model_fallback + && let Some(requested_model) = config.model.as_ref() + && model != *requested_model + { + info!( + model_provider = %config.model_provider_id, + requested_model, + fallback_model = %model, + "replaced unavailable requested model with provider default" + ); + } // Resolve base instructions for the session. Priority order: // 1. config.base_instructions override diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 64c4ba29f2..4f88c9ed66 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -718,6 +718,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { config, + allow_provider_model_fallback: false, user_instructions: Default::default(), installation_id: "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b2add06024..358eae3527 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -182,6 +182,7 @@ pub struct ThreadManager { pub struct StartThreadOptions { pub config: Config, + pub allow_provider_model_fallback: bool, pub initial_history: InitialHistory, pub session_source: Option, pub thread_source: Option, @@ -632,6 +633,7 @@ impl ThreadManager { ); Box::pin(self.start_thread_with_options(StartThreadOptions { config, + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -668,6 +670,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( options.config, options.initial_history, + options.allow_provider_model_fallback, Arc::clone(&self.state.auth_manager), agent_control, session_source, @@ -763,6 +766,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( config, initial_history, + /*allow_provider_model_fallback*/ false, auth_manager, agent_control, session_source, @@ -832,6 +836,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( config, initial_history, + /*allow_provider_model_fallback*/ false, auth_manager, agent_control, session_source, @@ -1332,6 +1337,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, InitialHistory::New, + /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, session_source, @@ -1370,6 +1376,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, session_source, @@ -1410,6 +1417,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, session_source, @@ -1451,6 +1459,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*allow_provider_model_fallback*/ false, auth_manager, agent_control, self.session_source.clone(), @@ -1475,6 +1484,7 @@ impl ThreadManagerState { &self, config: Config, initial_history: InitialHistory, + allow_provider_model_fallback: bool, auth_manager: Arc, agent_control: AgentControl, session_source: SessionSource, @@ -1541,6 +1551,7 @@ impl ThreadManagerState { codex, thread_id, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { config, + allow_provider_model_fallback, user_instructions, installation_id: self.installation_id.clone(), auth_manager, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 0d407e19c5..34df66cb67 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -400,6 +400,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { let thread = manager .start_thread_with_options(StartThreadOptions { config, + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: Some(SessionSource::Internal( InternalSessionSource::MemoryConsolidation, @@ -542,6 +543,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() let first_thread = manager .start_thread_with_options(StartThreadOptions { config: config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -557,6 +559,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() let second_thread = manager .start_thread_with_options(StartThreadOptions { config: config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -636,6 +639,7 @@ async fn selected_capability_roots_round_trip_through_fork() { let inherited = manager .start_thread_with_options(StartThreadOptions { config, + allow_provider_model_fallback: false, initial_history: InitialHistory::Forked(vec![RolloutItem::SessionMeta( SessionMetaLine { meta: SessionMeta { @@ -714,6 +718,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let source = manager .start_thread_with_options(StartThreadOptions { config: source_config, + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -997,6 +1002,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { let source = manager .start_thread_with_options(StartThreadOptions { config: config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: Some(ThreadSource::User), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index eb4fa3b39f..6b11267774 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -658,6 +658,7 @@ impl TestCodexBuilder { Box::pin( thread_manager.start_thread_with_options(StartThreadOptions { config: config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 605e0ada7d..4e0ca794e3 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -486,6 +486,7 @@ async fn loads_user_instructions_without_a_primary_environment() -> Result<()> { .thread_manager .start_thread_with_options(StartThreadOptions { config: test.config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -691,6 +692,7 @@ async fn multi_environment_thread_loads_every_project_and_keeps_creation_snapsho .thread_manager .start_thread_with_options(StartThreadOptions { config: test.config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: None, thread_source: None, diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index b63aff5400..63d891bada 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -1059,7 +1059,11 @@ async fn remote_models_request_times_out_after_5s() -> Result<()> { let start = Instant::now(); let model = timeout( Duration::from_secs(7), - manager.get_default_model(&None, RefreshStrategy::OnlineIfUncached), + manager.get_default_model( + &None, + /*allow_provider_model_fallback*/ false, + RefreshStrategy::OnlineIfUncached, + ), ) .await; let elapsed = start.elapsed(); @@ -1127,7 +1131,11 @@ async fn remote_models_hide_picker_only_models() -> Result<()> { ); let selected = manager - .get_default_model(&None, RefreshStrategy::OnlineIfUncached) + .get_default_model( + &None, + /*allow_provider_model_fallback*/ false, + RefreshStrategy::OnlineIfUncached, + ) .await; assert_eq!(selected, bundled_default_model_slug()); diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index aad8fc494b..62a367a229 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -761,6 +761,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() .thread_manager .start_thread_with_options(StartThreadOptions { config: test.config.clone(), + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: Some(SessionSource::SubAgent(SubAgentSource::Review)), thread_source: None, diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 0ee0313b58..73744e0a89 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -329,6 +329,7 @@ impl MemoryStartupContext { .thread_manager .start_thread_with_options(StartThreadOptions { config, + allow_provider_model_fallback: false, initial_history: InitialHistory::New, session_source: Some(SessionSource::Internal( InternalSessionSource::MemoryConsolidation, diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 7a68c254fe..67b939051b 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -144,11 +144,13 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { // todo(aibrahim): should be visible to core only and sent on session_configured event /// Get the model identifier to use, refreshing according to the specified strategy. /// - /// If `model` is provided, returns it directly. Otherwise selects the default based on - /// auth mode and available models. + /// If `model` is provided, preserves it unless the implementation supports and the policy + /// allows provider fallback. Otherwise selects the default based on auth mode and available + /// models. fn get_default_model<'a>( &'a self, model: &'a Option, + allow_provider_model_fallback: bool, refresh_strategy: RefreshStrategy, ) -> ModelsManagerFuture<'a, String> { Box::pin( @@ -161,6 +163,7 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { .instrument(tracing::info_span!( "get_default_model", model.provided = model.is_some(), + allow_provider_model_fallback, refresh_strategy = %refresh_strategy )), ) @@ -408,6 +411,39 @@ impl OpenAiModelsManager { } impl ModelsManager for StaticModelsManager { + fn get_default_model<'a>( + &'a self, + model: &'a Option, + allow_provider_model_fallback: bool, + refresh_strategy: RefreshStrategy, + ) -> ModelsManagerFuture<'a, String> { + Box::pin( + async move { + let available_models = self.list_models(refresh_strategy).await; + let requested_model = model.as_deref(); + + if allow_provider_model_fallback { + if requested_model_is_available(requested_model, &available_models) + && let Some(requested_model) = requested_model + { + return requested_model.to_string(); + } + return default_model_from_available(available_models); + } + + model + .clone() + .unwrap_or_else(|| default_model_from_available(available_models)) + } + .instrument(tracing::info_span!( + "get_default_model", + model.provided = model.is_some(), + allow_provider_model_fallback, + refresh_strategy = %refresh_strategy + )), + ) + } + fn raw_model_catalog( &self, _refresh_strategy: RefreshStrategy, @@ -453,6 +489,17 @@ fn default_model_from_available(available: Vec) -> String { .unwrap_or_default() } +fn requested_model_is_available( + requested_model: Option<&str>, + available_models: &[ModelPreset], +) -> bool { + requested_model.is_some_and(|requested_model| { + available_models + .iter() + .any(|available_model| available_model.model == requested_model) + }) +} + fn find_model_by_longest_prefix(model: &str, candidates: &[ModelInfo]) -> Option { let mut best: Option = None; for candidate in candidates { diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index bd692baad1..e7d1691a3a 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -239,6 +239,105 @@ c2ln", .expect("auth should be present") } +#[tokio::test] +async fn static_manager_preserves_supported_requested_model_when_fallback_is_allowed() { + let manager = static_manager_for_tests(ModelsResponse { + models: vec![ + remote_model("provider-default", "Default", /*priority*/ 0), + remote_model("provider-supported", "Supported", /*priority*/ 1), + ], + }); + let requested_model = Some("provider-supported".to_string()); + + let model = manager + .get_default_model( + &requested_model, + /*allow_provider_model_fallback*/ true, + RefreshStrategy::Offline, + ) + .await; + + assert_eq!(model, "provider-supported"); +} + +#[tokio::test] +async fn static_manager_falls_back_from_unsupported_requested_model_when_allowed() { + let manager = static_manager_for_tests(ModelsResponse { + models: vec![ + remote_model("provider-default", "Default", /*priority*/ 0), + remote_model("provider-supported", "Supported", /*priority*/ 1), + ], + }); + let requested_model = Some("unsupported".to_string()); + + let model = manager + .get_default_model( + &requested_model, + /*allow_provider_model_fallback*/ true, + RefreshStrategy::Offline, + ) + .await; + + assert_eq!(model, "provider-default"); +} + +#[tokio::test] +async fn static_manager_preserves_unsupported_requested_model_when_fallback_is_disabled() { + let manager = static_manager_for_tests(ModelsResponse { + models: vec![remote_model( + "provider-default", + "Default", + /*priority*/ 0, + )], + }); + let requested_model = Some("unsupported".to_string()); + + let model = manager + .get_default_model( + &requested_model, + /*allow_provider_model_fallback*/ false, + RefreshStrategy::Offline, + ) + .await; + + assert_eq!(model, "unsupported"); +} + +#[tokio::test] +async fn static_manager_uses_empty_default_when_fallback_is_allowed_and_catalog_is_empty() { + let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() }); + let requested_model = Some("unsupported".to_string()); + + let model = manager + .get_default_model( + &requested_model, + /*allow_provider_model_fallback*/ true, + RefreshStrategy::Offline, + ) + .await; + + assert_eq!(model, ""); +} + +#[tokio::test] +async fn dynamic_manager_preserves_requested_model_when_fallback_is_allowed() { + let codex_home = tempdir().expect("temp dir"); + let endpoint = TestModelsEndpoint::new(Vec::new()); + let manager = openai_manager_for_tests(codex_home.path().to_path_buf(), endpoint.clone()); + let requested_model = Some("unsupported".to_string()); + + let model = manager + .get_default_model( + &requested_model, + /*allow_provider_model_fallback*/ true, + RefreshStrategy::Online, + ) + .await; + + assert_eq!(model, "unsupported"); + assert_eq!(endpoint.fetch_count(), 0); +} + #[tokio::test] async fn get_model_info_tracks_fallback_usage() { let codex_home = tempdir().expect("temp dir"); From 31b99f65cf4b3eaf9180f625b48b5614806cd121 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 25 Jun 2026 11:25:19 -0700 Subject: [PATCH 09/33] cli: rename sandbox permission profile flag (#30095) ## Why `codex sandbox` accepts a single named permissions profile, so the existing plural `--permissions-profile` spelling is misleading. The canonical flag and its help text should use the singular form without breaking scripts that already use the old spelling. ## What changed - Make `--permission-profile` the canonical flag for all sandbox backends. - Keep `--permissions-profile` as a hidden backwards-compatible alias. - Cover the canonical spelling, legacy alias, and help visibility with regression tests. ## Testing Ran `just c sandbox --help` and verified I saw: ```shell -P, --permission-profile Named permissions profile to apply from the active configuration stack ``` --- codex-rs/cli/src/lib.rs | 21 ++++++++++++-- codex-rs/cli/src/main.rs | 31 ++++++++++++++++++++- codex-rs/cli/tests/sandbox_network_proxy.rs | 2 +- 3 files changed, 49 insertions(+), 5 deletions(-) diff --git a/codex-rs/cli/src/lib.rs b/codex-rs/cli/src/lib.rs index 1c0a2e7e02..c008567b96 100644 --- a/codex-rs/cli/src/lib.rs +++ b/codex-rs/cli/src/lib.rs @@ -53,7 +53,12 @@ pub struct SeatbeltCommand { pub sandbox_state: SandboxStateArgs, /// Named permissions profile to apply from the active configuration stack. - #[arg(long = "permissions-profile", short = 'P', value_name = "NAME")] + #[arg( + long = "permission-profile", + alias = "permissions-profile", + short = 'P', + value_name = "NAME" + )] pub permissions_profile: Option, /// Layer $CODEX_HOME/.config.toml on top of the base user config. @@ -104,7 +109,12 @@ pub struct LandlockCommand { pub sandbox_state: SandboxStateArgs, /// Named permissions profile to apply from the active configuration stack. - #[arg(long = "permissions-profile", short = 'P', value_name = "NAME")] + #[arg( + long = "permission-profile", + alias = "permissions-profile", + short = 'P', + value_name = "NAME" + )] pub permissions_profile: Option, /// Layer $CODEX_HOME/.config.toml on top of the base user config. @@ -142,7 +152,12 @@ pub struct WindowsCommand { pub sandbox_state: SandboxStateArgs, /// Named permissions profile to apply from the active configuration stack. - #[arg(long = "permissions-profile", short = 'P', value_name = "NAME")] + #[arg( + long = "permission-profile", + alias = "permissions-profile", + short = 'P', + value_name = "NAME" + )] pub permissions_profile: Option, /// Layer $CODEX_HOME/.config.toml on top of the base user config. diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 79af9831df..d7dbff47ec 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -2962,7 +2962,28 @@ mod tests { #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] #[test] - fn sandbox_parses_permissions_profile() { + fn sandbox_parses_permission_profile() { + let cli = MultitoolCli::try_parse_from([ + "codex", + "sandbox", + "--permission-profile", + ":workspace", + "--", + "echo", + ]) + .expect("parse"); + + let Some(Subcommand::Sandbox(command)) = cli.subcommand else { + panic!("expected sandbox command"); + }; + + assert_eq!(command.permissions_profile.as_deref(), Some(":workspace")); + assert_eq!(command.command, vec!["echo"]); + } + + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] + #[test] + fn sandbox_parses_legacy_permissions_profile_alias() { let cli = MultitoolCli::try_parse_from([ "codex", "sandbox", @@ -2981,6 +3002,14 @@ mod tests { assert_eq!(command.command, vec!["echo"]); } + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] + #[test] + fn sandbox_help_only_shows_singular_permission_profile() { + let help = help_from_args(&["codex", "sandbox", "--help"]); + assert!(help.contains("--permission-profile"), "{help}"); + assert!(!help.contains("--permissions-profile"), "{help}"); + } + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] #[test] fn sandbox_parses_permissions_profile_short_alias() { diff --git a/codex-rs/cli/tests/sandbox_network_proxy.rs b/codex-rs/cli/tests/sandbox_network_proxy.rs index 22bdcce49d..9ec1014688 100644 --- a/codex-rs/cli/tests/sandbox_network_proxy.rs +++ b/codex-rs/cli/tests/sandbox_network_proxy.rs @@ -35,7 +35,7 @@ mode = "full" .env("CODEX_HOME", codex_home.path()) .args([ "sandbox", - "--permissions-profile", + "--permission-profile", "network-test", "--", "curl", From cc78903379773233264ae9e3ccd3740439d70b51 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 25 Jun 2026 11:30:53 -0700 Subject: [PATCH 10/33] [codex] current time reminder interval to be set to 0 (#30029) A zero interval lets callers request a reminder at every otherwise-eligible inference boundary. ## Validation - just test -p codex-core load_config_resolves_current_time_reminder --- codex-rs/core/config.schema.json | 2 +- codex-rs/core/src/config/config_tests.rs | 24 +----------- codex-rs/core/src/config/mod.rs | 6 --- codex-rs/core/src/session/time_reminder.rs | 1 + .../core/tests/suite/current_time_reminder.rs | 39 +++++++++++++++++++ codex-rs/features/src/feature_configs.rs | 1 - 6 files changed, 43 insertions(+), 30 deletions(-) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 92a45d1978..bfdd179d78 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -815,7 +815,7 @@ }, "reminder_interval_seconds": { "format": "uint64", - "minimum": 1.0, + "minimum": 0.0, "type": "integer" }, "sleep_tool": { diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index a178dfae04..17779117f6 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -625,12 +625,12 @@ current_time_reminder = true r#" [features.current_time_reminder] enabled = true -reminder_interval_seconds = 4 +reminder_interval_seconds = 0 clock_source = "external" sleep_tool = true "#, CurrentTimeReminderConfig { - reminder_interval_seconds: 4, + reminder_interval_seconds: 0, clock_source: CurrentTimeSource::External, sleep_tool: true, }, @@ -643,26 +643,6 @@ sleep_tool = true Ok(()) } -#[tokio::test] -async fn load_config_rejects_zero_current_time_reminder_interval() -> std::io::Result<()> { - let error = load_current_time_reminder_config( - r#" -[features.current_time_reminder] -enabled = true -reminder_interval_seconds = 0 -"#, - ) - .await - .expect_err("zero reminder interval should be rejected"); - - assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); - assert_eq!( - error.to_string(), - "features.current_time_reminder.reminder_interval_seconds must be positive" - ); - Ok(()) -} - async fn load_current_time_reminder_config(config_toml: &str) -> std::io::Result { let codex_home = tempdir()?; let config_toml = toml::from_str(config_toml).expect("TOML should deserialize"); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 78188d92e2..58ae94906e 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -2685,12 +2685,6 @@ fn resolve_current_time_reminder_config( let reminder_interval_seconds = base .and_then(|config| config.reminder_interval_seconds) .unwrap_or(default.reminder_interval_seconds); - if reminder_interval_seconds == 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "features.current_time_reminder.reminder_interval_seconds must be positive", - )); - } Ok(Some(CurrentTimeReminderConfig { reminder_interval_seconds, diff --git a/codex-rs/core/src/session/time_reminder.rs b/codex-rs/core/src/session/time_reminder.rs index 85831d2947..8f612e8fcb 100644 --- a/codex-rs/core/src/session/time_reminder.rs +++ b/codex-rs/core/src/session/time_reminder.rs @@ -21,6 +21,7 @@ impl CurrentTimeReminderState { interval_seconds: u64, ) -> bool { let reminder_is_due = self.last_window_id.as_deref() != Some(window_id) + || interval_seconds == 0 || self.last_delivery_time.is_none_or(|last_delivery_time| { current_time .signed_duration_since(last_delivery_time) diff --git a/codex-rs/core/tests/suite/current_time_reminder.rs b/codex-rs/core/tests/suite/current_time_reminder.rs index 8ebf5d6b16..c5ebd0ed13 100644 --- a/codex-rs/core/tests/suite/current_time_reminder.rs +++ b/codex-rs/core/tests/suite/current_time_reminder.rs @@ -39,6 +39,7 @@ use pretty_assertions::assert_eq; use serde_json::json; const FIRST_REMINDER: &str = "It is 2026-06-17 17:34:15 UTC."; +const EARLIER_REMINDER: &str = "It is 2026-06-17 17:33:15 UTC."; const SECOND_REMINDER: &str = "It is 2026-06-17 17:35:15 UTC."; const THIRD_REMINDER: &str = "It is 2026-06-17 17:36:15 UTC."; const FIRST_TIME_UNIX_SECONDS: i64 = 1_781_717_655; @@ -162,6 +163,44 @@ async fn current_time_reminders_follow_time_interval_and_persist_in_history() -> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_current_time_reminder_interval_delivers_when_time_moves_backward() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let time_provider = Arc::new(TestTimeProvider::default()); + let test = test_codex() + .with_config(|config| { + enable_current_time_reminder(config, /*interval*/ 0, CurrentTimeSource::External) + }) + .with_external_time_provider(time_provider.clone()) + .build(&server) + .await?; + + test.submit_turn("first turn").await?; + time_provider + .current_time + .store(FIRST_TIME_UNIX_SECONDS - 60, Ordering::Relaxed); + test.submit_turn("second turn").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(current_time_reminders(&requests[0]), vec![FIRST_REMINDER]); + assert_eq!( + current_time_reminders(&requests[1]), + vec![FIRST_REMINDER, EARLIER_REMINDER] + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn system_time_source_adds_current_time_reminder() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 2b638fce6e..24036eee58 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -147,7 +147,6 @@ pub struct CurrentTimeReminderConfigToml { #[serde(skip_serializing_if = "Option::is_none")] pub enabled: Option, #[serde(skip_serializing_if = "Option::is_none")] - #[schemars(range(min = 1))] pub reminder_interval_seconds: Option, #[serde(skip_serializing_if = "Option::is_none")] pub clock_source: Option, From c65cfeab1478c4c8356b42aca3d06c92b2730334 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 25 Jun 2026 12:00:23 -0700 Subject: [PATCH 11/33] core: expose permission profile to shell tools (#29941) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## tl;dr Inject a `CODEX_PERMISSION_PROFILE` environment variable with the name of the current permission profile when invoking a shell tool. ## Why Shell tool owners may need to launch nested commands under the same named permission profile, including through `codex sandbox -P PROFILE --include-managed-config`. Until now, child processes could observe sandbox and network metadata but could not identify the active named permission profile. The `--include-managed-config` flag is essential when a helper reconstructs the sandbox from a profile name: it ensures the nested sandbox also loads managed enterprise requirements. Without it, using the inherited profile could unintentionally create a sandbox that does not enforce the organization's managed restrictions. The new environment value is intentionally informational and **must not be treated as trusted input**. Any process in the ancestry can overwrite an environment variable, so a consumer that passes this value to `codex sandbox -P` must first validate it against the profiles that helper is authorized to use. ## Example Use Case Suppose an organization provides a trusted `remote-bash` wrapper that lets Codex run a command on an approved build host. The local shell command uses the named `:workspace` permission profile: ```toml default_permissions = ":workspace" ``` The command exposed to the model is a small zsh wrapper. It deliberately delegates with `exec`, preserving the original arguments and process environment: ```zsh #!/usr/bin/env zsh exec /opt/codex-tools/remote_bash.py "$@" ``` The model invokes the public wrapper, not its Python implementation: ```sh /opt/codex-tools/remote-bash \ --host builder.example.com \ -- printf '%s' 'hello world' ``` Only the inner implementation is authorized to escape the local sandbox: ```starlark prefix_rule( pattern=["/opt/codex-tools/remote_bash.py"], decision="allow", ) ``` With zsh-fork, execution begins with `remote-bash` inside the `:workspace` sandbox. When the wrapper calls `exec`, the exact prefix rule matches `remote_bash.py`, so that inner script is restarted unsandboxed. The escalated process inherits: ```text CODEX_PERMISSION_PROFILE=:workspace ``` Inheritance does not make the value trustworthy. `remote_bash.py` independently allowlists both the remote host and the permission profile before using either value. In particular, a forged value such as `:danger-full-access` is rejected before it can reach `codex sandbox -P`: ```python import argparse import os import shlex import sys ALLOWED_HOSTS = {"builder.example.com"} ALLOWED_PROFILES = {":workspace"} parser = argparse.ArgumentParser() parser.add_argument("--host", required=True) separator = sys.argv.index("--") args = parser.parse_args(sys.argv[1:separator]) command = sys.argv[separator + 1:] if args.host not in ALLOWED_HOSTS: parser.error("host is not allowlisted") if not command: parser.error("the remote command must not be empty") profile = os.environ.get("CODEX_PERMISSION_PROFILE") if not profile: raise SystemExit("CODEX_PERMISSION_PROFILE must not be empty") if profile not in ALLOWED_PROFILES: raise SystemExit("CODEX_PERMISSION_PROFILE is not allowlisted") remote_command = shlex.join(command) sandbox_command = shlex.join([ "codex", "sandbox", "-P", profile, "--include-managed-config", "--", "bash", "-lc", remote_command, ]) print(shlex.join(["ssh", args.host, sandbox_command])) ``` This builds each command layer as an argument vector and uses `shlex.join()` at the boundary, rather than interpolating untrusted shell text. After validation and parsing, the nested command has this structure: ```text ssh argv: ["ssh", "builder.example.com", SANDBOX_COMMAND] SANDBOX_COMMAND argv: ["codex", "sandbox", "-P", ":workspace", "--include-managed-config", "--", "bash", "-lc", "printf %s 'hello world'"] bash -lc payload argv: ["printf", "%s", "hello world"] ``` A production implementation could execute that SSH command. The integration fixture prints it and parses the result back into arguments, verifying the complete flow: ```text model invokes outer wrapper -> zsh-fork starts wrapper under :workspace -> wrapper execs allowlisted Python script -> prefix rule restarts Python script unsandboxed -> Python script inherits CODEX_PERMISSION_PROFILE=:workspace -> Python script verifies :workspace is allowlisted -> remote command runs codex sandbox -P :workspace with --include-managed-config -> nested sandbox honors managed enterprise requirements ``` This gives the trusted helper access to resources outside the local sandbox—such as SSH credentials—while ensuring that it can select only an explicitly authorized profile and that work on the remote host remains subject to the organization's managed requirements. ## What changed - Inject `CODEX_PERMISSION_PROFILE` after shell environment policy evaluation so the active profile wins over inherited or configured stale values. - Apply the variable to both `shell_command` and unified `exec_command`, including local, zsh-fork, and remote exec-server paths. - Remove stale values when the session has no active named profile. - Preserve the current profile value when loading a shell snapshot so a parent snapshot cannot restore an older profile. ## Testing - Added classic-shell integration coverage proving an exact prefix rule can run a `require_escalated` script outside the `:workspace` sandbox while preserving `CODEX_PERMISSION_PROFILE=:workspace`. - Added zsh-fork integration coverage in which the model invokes an outer zsh wrapper, an inner allowlisted `remote_bash.py` runs unsandboxed, and its printed SSH command reconstructs the inherited `:workspace` sandbox with `--include-managed-config` while preserving every argument after `--`. - The example helper treats `CODEX_PERMISSION_PROFILE` as untrusted and validates it against `ALLOWED_PROFILES` before constructing the nested command. - Assert that the reconstructed sandbox command includes `--include-managed-config` so nested use of the inherited profile cannot bypass managed enterprise requirements. - Added coverage for overriding and removing stale profile values. - Verified `shell_command` receives the selected active profile. - Added shell snapshot coverage using `printenv CODEX_PERMISSION_PROFILE`. --- codex-rs/core/src/exec_env.rs | 26 ++ codex-rs/core/src/exec_env_tests.rs | 53 +++ .../src/tools/handlers/shell/shell_command.rs | 13 +- .../core/src/tools/handlers/shell_tests.rs | 24 +- codex-rs/core/src/tools/runtimes/mod.rs | 18 +- codex-rs/core/src/tools/runtimes/mod_tests.rs | 72 ++++ .../core/src/unified_exec/process_manager.rs | 25 +- .../src/unified_exec/process_manager_tests.rs | 37 ++ codex-rs/core/tests/common/zsh_fork.rs | 35 +- codex-rs/core/tests/suite/approvals.rs | 370 ++++++++++++++++++ 10 files changed, 642 insertions(+), 31 deletions(-) diff --git a/codex-rs/core/src/exec_env.rs b/codex-rs/core/src/exec_env.rs index 938667b12e..f33061698b 100644 --- a/codex-rs/core/src/exec_env.rs +++ b/codex-rs/core/src/exec_env.rs @@ -2,11 +2,16 @@ use codex_protocol::ThreadId; #[cfg(test)] use codex_protocol::config_types::EnvironmentVariablePattern; use codex_protocol::config_types::ShellEnvironmentPolicy; +use codex_protocol::models::ActivePermissionProfile; use codex_protocol::shell_environment; use std::collections::HashMap; pub use codex_protocol::shell_environment::CODEX_THREAD_ID_ENV_VAR; +/// Informational name of the active permission profile. Child processes can +/// overwrite this value, so it must not be treated as proof of enforcement. +pub const CODEX_PERMISSION_PROFILE_ENV_VAR: &str = "CODEX_PERMISSION_PROFILE"; + /// Construct an environment map based on the rules in the specified policy. The /// resulting map can be passed directly to `Command::envs()` after calling /// `env_clear()` to ensure no unintended variables are leaked to the spawned @@ -25,6 +30,27 @@ pub fn create_env( shell_environment::create_env(policy, thread_id.as_deref()) } +/// Injects the selected named permission profile into a shell tool's environment. +/// +/// This is applied after the shell environment policy so the runtime-selected +/// profile wins over inherited or configured values. +pub(crate) fn inject_permission_profile_env( + env: &mut HashMap, + active_permission_profile: Option<&ActivePermissionProfile>, +) { + if cfg!(windows) { + env.retain(|key, _| !key.eq_ignore_ascii_case(CODEX_PERMISSION_PROFILE_ENV_VAR)); + } else { + env.remove(CODEX_PERMISSION_PROFILE_ENV_VAR); + } + if let Some(active_permission_profile) = active_permission_profile { + env.insert( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + active_permission_profile.id.clone(), + ); + } +} + #[cfg(all(test, target_os = "windows"))] fn create_env_from_vars( vars: I, diff --git a/codex-rs/core/src/exec_env_tests.rs b/codex-rs/core/src/exec_env_tests.rs index 725edd8cc5..73c0944c14 100644 --- a/codex-rs/core/src/exec_env_tests.rs +++ b/codex-rs/core/src/exec_env_tests.rs @@ -10,6 +10,59 @@ fn make_vars(pairs: &[(&str, &str)]) -> Vec<(String, String)> { .collect() } +#[test] +fn inject_permission_profile_env_overrides_policy_value() { + let mut env = HashMap::from([( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "stale-profile".to_string(), + )]); + + inject_permission_profile_env( + &mut env, + Some(&ActivePermissionProfile::new("current-profile")), + ); + + assert_eq!( + env.get(CODEX_PERMISSION_PROFILE_ENV_VAR) + .map(String::as_str), + Some("current-profile") + ); +} + +#[test] +fn inject_permission_profile_env_removes_stale_value_without_active_profile() { + let mut env = HashMap::from([( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "stale-profile".to_string(), + )]); + + inject_permission_profile_env(&mut env, /*active_permission_profile*/ None); + + assert_eq!(env.get(CODEX_PERMISSION_PROFILE_ENV_VAR), None); +} + +#[cfg(target_os = "windows")] +#[test] +fn inject_permission_profile_env_replaces_differently_cased_windows_key() { + let mut env = HashMap::from([( + "codex_permission_profile".to_string(), + "stale-profile".to_string(), + )]); + + inject_permission_profile_env( + &mut env, + Some(&ActivePermissionProfile::new("current-profile")), + ); + + assert_eq!( + env, + HashMap::from([( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "current-profile".to_string(), + )]) + ); +} + #[test] fn test_core_inherit_defaults_keep_sensitive_vars() { let vars = make_vars(&[ diff --git a/codex-rs/core/src/tools/handlers/shell/shell_command.rs b/codex-rs/core/src/tools/handlers/shell/shell_command.rs index 1faa89b31f..62bcee8bdf 100644 --- a/codex-rs/core/src/tools/handlers/shell/shell_command.rs +++ b/codex-rs/core/src/tools/handlers/shell/shell_command.rs @@ -6,6 +6,7 @@ use codex_utils_absolute_path::AbsolutePathBuf; use crate::exec::ExecCapturePolicy; use crate::exec::ExecParams; use crate::exec_env::create_env; +use crate::exec_env::inject_permission_profile_env; use crate::function_tool::FunctionCallError; use crate::maybe_emit_implicit_skill_invocation; use crate::session::turn_context::TurnContext; @@ -99,15 +100,19 @@ impl ShellCommandHandler { let use_login_shell = Self::resolve_use_login_shell(params.login, allow_login_shell)?; let command = Self::base_command(shell, ¶ms.command, use_login_shell); + let mut env = create_env( + &turn_context.config.permissions.shell_environment_policy, + Some(session.thread_id), + ); + let active_permission_profile = turn_context.config.permissions.active_permission_profile(); + inject_permission_profile_env(&mut env, active_permission_profile.as_ref()); + Ok(ExecParams { command, cwd, expiration: params.timeout_ms.into(), capture_policy: ExecCapturePolicy::ShellTool, - env: create_env( - &turn_context.config.permissions.shell_environment_policy, - Some(session.thread_id), - ), + env, network: turn_context.network.clone(), network_environment_id: Some(turn_environment.environment_id.clone()), sandbox_permissions: params.sandbox_permissions.unwrap_or_default(), diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index ded7733022..ebfbf9127c 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -1,10 +1,14 @@ use std::path::PathBuf; use std::sync::Arc; +use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::ShellCommandToolCallParams; use pretty_assertions::assert_eq; +use crate::config::PermissionProfileSnapshot; +use crate::exec_env::CODEX_PERMISSION_PROFILE_ENV_VAR; use crate::exec_env::create_env; +use crate::exec_env::inject_permission_profile_env; use crate::sandboxing::SandboxPermissions; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; @@ -71,7 +75,15 @@ fn assert_safe(shell: &Shell, command: &str) { #[tokio::test] async fn shell_command_handler_to_exec_params_uses_selected_environment() { - let (session, turn_context) = make_session_and_context().await; + let (session, mut turn_context) = make_session_and_context().await; + let permission_profile = turn_context.config.permissions.permission_profile().clone(); + Arc::make_mut(&mut turn_context.config) + .permissions + .set_permission_profile_from_session_snapshot(PermissionProfileSnapshot::active( + permission_profile, + ActivePermissionProfile::new("test-profile"), + )) + .expect("set active permission profile"); let command = "echo hello".to_string(); let workdir = Some("subdir".to_string()); @@ -99,10 +111,12 @@ async fn shell_command_handler_to_exec_params_uses_selected_environment() { PathUri::from_abs_path(&selected_cwd), Some(selected_shell), ); - let expected_env = create_env( + let mut expected_env = create_env( &turn_context.config.permissions.shell_environment_policy, Some(session.thread_id), ); + let active_permission_profile = turn_context.config.permissions.active_permission_profile(); + inject_permission_profile_env(&mut expected_env, active_permission_profile.as_ref()); let params = ShellCommandToolCallParams { command, @@ -129,6 +143,12 @@ async fn shell_command_handler_to_exec_params_uses_selected_environment() { assert_eq!(exec_params.command, expected_command); assert_eq!(exec_params.cwd, expected_cwd); assert_eq!(exec_params.env, expected_env); + assert_eq!( + exec_params.env.get(CODEX_PERMISSION_PROFILE_ENV_VAR), + active_permission_profile + .as_ref() + .map(|profile| &profile.id) + ); assert_eq!(exec_params.network, turn_context.network); assert_eq!( exec_params.network_environment_id.as_deref(), diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index db55cc2bbc..1a83ef5657 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -4,6 +4,7 @@ Module: runtimes Concrete ToolRuntime implementations for specific tools. Each runtime stays small and focused and reuses the orchestrator for approvals + sandbox + retry. */ +use crate::exec_env::CODEX_PERMISSION_PROFILE_ENV_VAR; use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; use crate::sandboxing::SandboxPermissions; use crate::shell::Shell; @@ -287,10 +288,14 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot( .map(|arg| format!(" '{}'", shell_single_quote(arg))) .collect::(); let mut override_env = explicit_env_overrides.clone(); - if let Some(thread_id) = env.get(CODEX_THREAD_ID_ENV_VAR) { - override_env.insert(CODEX_THREAD_ID_ENV_VAR.to_string(), thread_id.clone()); + for key in [CODEX_THREAD_ID_ENV_VAR, CODEX_PERMISSION_PROFILE_ENV_VAR] { + if let Some(value) = env.get(key) { + override_env.insert(key.to_string(), value.clone()); + } } - let (override_captures, override_exports) = build_override_exports(&override_env); + // Do not let a snapshot resurrect a stale profile when no named profile is active. + let (override_captures, override_exports) = + build_override_exports(&override_env, &[CODEX_PERMISSION_PROFILE_ENV_VAR]); let (proxy_captures, proxy_exports) = build_proxy_env_exports(); let runtime_path_prepend_exports = runtime_path_prepends.shell_exports_after_snapshot(explicit_env_overrides); @@ -313,13 +318,18 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot( vec![shell_path.to_string(), "-c".to_string(), rewritten_script] } -fn build_override_exports(explicit_env_overrides: &HashMap) -> (String, String) { +fn build_override_exports( + explicit_env_overrides: &HashMap, + restore_even_when_absent: &[&str], +) -> (String, String) { let mut keys = explicit_env_overrides .keys() .map(String::as_str) + .chain(restore_even_when_absent.iter().copied()) .filter(|key| is_valid_shell_variable_name(key)) .collect::>(); keys.sort_unstable(); + keys.dedup(); build_override_exports_for_keys("__CODEX_SNAPSHOT_OVERRIDE", &keys) } diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index 9b485dbe58..5c2614c9c0 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -534,6 +534,78 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_codex_thread_id_from_env() { assert_eq!(String::from_utf8_lossy(&output.stdout), "nested-thread"); } +#[test] +fn maybe_wrap_shell_lc_with_snapshot_restores_permission_profile_from_env() { + let dir = tempdir().expect("create temp dir"); + let snapshot_path = dir.path().join("snapshot.sh"); + std::fs::write( + &snapshot_path, + "# Snapshot file\nexport CODEX_PERMISSION_PROFILE='parent-profile'\n", + ) + .expect("write snapshot"); + let (session_shell, shell_snapshot) = + shell_with_snapshot(ShellType::Bash, "/bin/bash", snapshot_path.abs()); + let command = vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "printenv CODEX_PERMISSION_PROFILE".to_string(), + ]; + let env = HashMap::from([( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "current-profile".to_string(), + )]); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + Some(&shell_snapshot), + &HashMap::new(), + &env, + &RuntimePathPrepends::default(), + ); + let output = Command::new(&rewritten[0]) + .args(&rewritten[1..]) + .env(CODEX_PERMISSION_PROFILE_ENV_VAR, "current-profile") + .output() + .expect("run rewritten command"); + + assert!(output.status.success(), "command failed: {output:?}"); + assert_eq!(String::from_utf8_lossy(&output.stdout), "current-profile\n"); +} + +#[test] +fn maybe_wrap_shell_lc_with_snapshot_unsets_absent_permission_profile() { + let dir = tempdir().expect("create temp dir"); + let snapshot_path = dir.path().join("snapshot.sh"); + std::fs::write( + &snapshot_path, + "# Snapshot file\nexport CODEX_PERMISSION_PROFILE='stale-profile'\n", + ) + .expect("write snapshot"); + let (session_shell, shell_snapshot) = + shell_with_snapshot(ShellType::Bash, "/bin/bash", snapshot_path.abs()); + let command = vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "printenv CODEX_PERMISSION_PROFILE".to_string(), + ]; + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + Some(&shell_snapshot), + &HashMap::new(), + &HashMap::new(), + &RuntimePathPrepends::default(), + ); + let output = Command::new(&rewritten[0]) + .args(&rewritten[1..]) + .env_remove(CODEX_PERMISSION_PROFILE_ENV_VAR) + .output() + .expect("run rewritten command"); + + assert_eq!(output.status.code(), Some(1)); + assert_eq!(output.stdout, b""); +} + #[test] fn maybe_wrap_shell_lc_with_snapshot_restores_proxy_env_from_process_env() { let dir = tempdir().expect("create temp dir"); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 5e89445765..ef22918cd2 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -13,8 +13,10 @@ use tokio_util::sync::CancellationToken; use uuid::Uuid; use crate::codex_thread::BackgroundTerminalInfo; +use crate::exec_env::CODEX_PERMISSION_PROFILE_ENV_VAR; use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; use crate::exec_env::create_env; +use crate::exec_env::inject_permission_profile_env; use crate::exec_policy::ExecApprovalRequest; use crate::sandboxing::ExecOptions; use crate::sandboxing::ExecRequest; @@ -109,15 +111,19 @@ fn apply_unified_exec_env(mut env: HashMap) -> HashMap codex_exec_server::ExecEnvPolicy { + let mut exclude = policy + .exclude + .iter() + .map(std::string::ToString::to_string) + .collect::>(); + exclude.push(CODEX_PERMISSION_PROFILE_ENV_VAR.to_string()); + let mut r#set = policy.r#set.clone(); + r#set.retain(|key, _| !key.eq_ignore_ascii_case(CODEX_PERMISSION_PROFILE_ENV_VAR)); codex_exec_server::ExecEnvPolicy { inherit: policy.inherit.clone(), ignore_default_excludes: policy.ignore_default_excludes, - exclude: policy - .exclude - .iter() - .map(std::string::ToString::to_string) - .collect(), - r#set: policy.r#set.clone(), + exclude, + r#set, include_only: policy .include_only .iter() @@ -132,7 +138,10 @@ fn env_overlay_for_exec_server( ) -> HashMap { request_env .iter() - .filter(|(key, value)| local_policy_env.get(*key) != Some(*value)) + .filter(|(key, value)| { + key.as_str() == CODEX_PERMISSION_PROFILE_ENV_VAR + || local_policy_env.get(*key) != Some(*value) + }) .map(|(key, value)| (key.clone(), value.clone())) .collect() } @@ -1110,6 +1119,8 @@ impl UnifiedExecProcessManager { CODEX_THREAD_ID_ENV_VAR.to_string(), context.session.thread_id.to_string(), ); + let active_permission_profile = context.turn.config.permissions.active_permission_profile(); + inject_permission_profile_env(&mut env, active_permission_profile.as_ref()); let env = apply_unified_exec_env(env); let exec_server_env_config = ExecServerEnvConfig { policy: exec_env_policy_from_shell_policy( diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 93ffd0cedb..096cbe0a13 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -42,12 +42,20 @@ fn env_overlay_for_exec_server_keeps_runtime_changes_only() { ("HOME".to_string(), "/client-home".to_string()), ("PATH".to_string(), "/client-path".to_string()), ("SHELL_SET".to_string(), "policy".to_string()), + ( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "current-profile".to_string(), + ), ]); let request_env = HashMap::from([ ("HOME".to_string(), "/client-home".to_string()), ("PATH".to_string(), "/sandbox-path".to_string()), ("SHELL_SET".to_string(), "policy".to_string()), ("CODEX_THREAD_ID".to_string(), "thread-1".to_string()), + ( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "current-profile".to_string(), + ), ( "CODEX_SANDBOX_NETWORK_DISABLED".to_string(), "1".to_string(), @@ -59,6 +67,10 @@ fn env_overlay_for_exec_server_keeps_runtime_changes_only() { HashMap::from([ ("PATH".to_string(), "/sandbox-path".to_string()), ("CODEX_THREAD_ID".to_string(), "thread-1".to_string()), + ( + CODEX_PERMISSION_PROFILE_ENV_VAR.to_string(), + "current-profile".to_string(), + ), ( "CODEX_SANDBOX_NETWORK_DISABLED".to_string(), "1".to_string() @@ -67,6 +79,31 @@ fn env_overlay_for_exec_server_keeps_runtime_changes_only() { ); } +#[test] +fn exec_env_policy_excludes_runtime_permission_profile() { + let policy = ShellEnvironmentPolicy { + r#set: HashMap::from([ + ( + "codex_permission_profile".to_string(), + "stale-profile".to_string(), + ), + ("KEEP".to_string(), "value".to_string()), + ]), + ..Default::default() + }; + + assert_eq!( + exec_env_policy_from_shell_policy(&policy), + codex_exec_server::ExecEnvPolicy { + inherit: policy.inherit, + ignore_default_excludes: policy.ignore_default_excludes, + exclude: vec![CODEX_PERMISSION_PROFILE_ENV_VAR.to_string()], + r#set: HashMap::from([("KEEP".to_string(), "value".to_string())]), + include_only: Vec::new(), + } + ); +} + #[test] fn exec_server_params_use_path_uri_and_env_policy_overlay_contract() { let cwd: codex_utils_absolute_path::AbsolutePathBuf = std::env::current_dir() diff --git a/codex-rs/core/tests/common/zsh_fork.rs b/codex-rs/core/tests/common/zsh_fork.rs index 5abfc753b9..8736da72c3 100644 --- a/codex-rs/core/tests/common/zsh_fork.rs +++ b/codex-rs/core/tests/common/zsh_fork.rs @@ -10,6 +10,7 @@ use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; use crate::test_codex::TestCodex; +use crate::test_codex::TestCodexBuilder; use crate::test_codex::test_codex; #[derive(Clone)] @@ -19,12 +20,7 @@ pub struct ZshForkRuntime { } impl ZshForkRuntime { - fn apply_to_config( - &self, - config: &mut Config, - approval_policy: AskForApproval, - permission_profile: PermissionProfile, - ) { + fn apply_to_config(&self, config: &mut Config, approval_policy: AskForApproval) { config .features .enable(Feature::ShellTool) @@ -37,10 +33,6 @@ impl ZshForkRuntime { config.main_execve_wrapper_exe = Some(self.main_execve_wrapper_exe.clone()); config.permissions.allow_login_shell = false; config.permissions.approval_policy = Constrained::allow_any(approval_policy); - config - .permissions - .set_permission_profile(permission_profile) - .expect("set permission profile"); } } @@ -86,10 +78,13 @@ pub async fn build_zsh_fork_test( where F: FnOnce(&Path) + Send + 'static, { - let mut builder = test_codex() + let mut builder = zsh_fork_test_builder(runtime, approval_policy) .with_pre_build_hook(pre_build_hook) .with_config(move |config| { - runtime.apply_to_config(config, approval_policy, permission_profile); + config + .permissions + .set_permission_profile(permission_profile) + .expect("set permission profile"); }); builder.build(server).await } @@ -104,10 +99,13 @@ pub async fn build_unified_exec_zsh_fork_test( where F: FnOnce(&Path) + Send + 'static, { - let mut builder = test_codex() + let mut builder = zsh_fork_test_builder(runtime, approval_policy) .with_pre_build_hook(pre_build_hook) .with_config(move |config| { - runtime.apply_to_config(config, approval_policy, permission_profile); + config + .permissions + .set_permission_profile(permission_profile) + .expect("set permission profile"); config.use_experimental_unified_exec_tool = true; config .features @@ -121,6 +119,15 @@ where builder.build(server).await } +pub fn zsh_fork_test_builder( + runtime: ZshForkRuntime, + approval_policy: AskForApproval, +) -> TestCodexBuilder { + test_codex().with_config(move |config| { + runtime.apply_to_config(config, approval_policy); + }) +} + fn find_test_zsh_path() -> Result> { let repo_root = codex_utils_cargo_bin::repo_root()?; let dotslash_zsh = repo_root.join("codex-rs/app-server/tests/suite/zsh"); diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index f0ab48c5e1..54205f7e8b 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -10,6 +10,7 @@ use codex_features::Feature; use codex_protocol::approvals::NetworkApprovalProtocol; use codex_protocol::approvals::NetworkPolicyAmendment; use codex_protocol::approvals::NetworkPolicyRuleAction; +use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -47,6 +48,7 @@ use core_test_support::wait_for_event_with_timeout; use core_test_support::zsh_fork::build_zsh_fork_test; use core_test_support::zsh_fork::restrictive_workspace_write_profile; use core_test_support::zsh_fork::zsh_fork_runtime; +use core_test_support::zsh_fork::zsh_fork_test_builder; use pretty_assertions::assert_eq; use regex_lite::Regex; use serde_json::Value; @@ -131,6 +133,8 @@ enum ActionKind { const DEFAULT_UNIFIED_EXEC_JUSTIFICATION: &str = "Requires escalated permissions to bypass the sandbox in tests."; +const WORKSPACE_PERMISSION_PROFILE_CONFIG: &str = r#"default_permissions = ":workspace" +"#; impl ActionKind { fn policy_src(&self) -> Option<&'static str> { @@ -640,6 +644,7 @@ enum ScenarioGroup { UnifiedExec, } +#[derive(Debug, Eq, PartialEq)] struct CommandResult { exit_code: Option, stdout: String, @@ -683,6 +688,52 @@ async fn submit_turn( Ok(()) } +async fn submit_turn_preserving_active_permission_profile( + test: &TestCodex, + prompt: &str, + approval_policy: AskForApproval, +) -> Result<()> { + let session_model = test.session_configured.model.clone(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: prompt.into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + approval_policy: Some(approval_policy), + approvals_reviewer: Some(ApprovalsReviewer::User), + collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { + mode: codex_protocol::config_types::ModeKind::Default, + settings: codex_protocol::config_types::Settings { + model: session_model, + reasoning_effort: None, + developer_instructions: None, + }, + }), + ..Default::default() + }, + }) + .await?; + + Ok(()) +} + +fn assert_active_workspace_permission_profile(test: &TestCodex) { + assert_eq!( + test.session_configured + .active_permission_profile + .as_ref() + .map(|profile| profile.id.as_str()), + Some(BUILT_IN_PERMISSION_PROFILE_WORKSPACE) + ); +} + fn parse_result(item: &Value) -> CommandResult { let output_str = item .get("output") @@ -2696,6 +2747,325 @@ async fn matched_prefix_rule_runs_unsandboxed_under_zsh_fork() -> Result<()> { Ok(()) } +/// Verifies that an allowlisted script retains the originating named profile +/// when its shell tool call requests escalated, unsandboxed execution. +/// +/// Tool owners use this pattern when a trusted wrapper must run outside the +/// current sandbox, but then needs to launch child commands back inside the +/// same sandbox with `codex sandbox -P`. The nested invocation must also pass +/// `--include-managed-config` so it continues to honor enterprise requirements. +/// The test proves both halves of that contract: the wrapper writes outside the +/// `:workspace` sandbox, while its inherited profile name remains `:workspace`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg(unix)] +async fn allowed_escalated_shell_command_inherits_active_permission_profile() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + fs::write( + home.path().join("config.toml"), + WORKSPACE_PERMISSION_PROFILE_CONFIG, + )?; + + let script_dir = tempfile::tempdir_in(std::env::current_dir()?)?; + let script_path = script_dir.path().join("print-permission-profile.sh"); + let outside_path = script_dir.path().join("unsandboxed-marker"); + fs::write( + &script_path, + format!( + r#"#!/bin/sh +# Print the inherited profile so the test can verify that it reached this script. +printenv CODEX_PERMISSION_PROFILE +touch {outside_path:?} +"# + ), + )?; + + let rules_dir = home.path().join("rules"); + fs::create_dir_all(&rules_dir)?; + let script_pattern = serde_json::to_string(&script_path.to_string_lossy())?; + fs::write( + rules_dir.join("default.rules"), + format!(r#"prefix_rule(pattern=["/bin/sh", {script_pattern}], decision="allow")"#), + )?; + + let approval_policy = AskForApproval::OnRequest; + let mut builder = test_codex().with_home(home).with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(approval_policy); + }); + let test = builder.build(&server).await?; + assert!(!outside_path.starts_with(test.config.cwd.as_path())); + assert_active_workspace_permission_profile(&test); + + let call_id = "allowed-escalated-shell-inherits-permission-profile"; + let command = format!("/bin/sh {script_path:?}"); + let event = shell_event( + call_id, + &command, + /*timeout_ms*/ 5_000, + SandboxPermissions::RequireEscalated, + )?; + let _ = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-escalated-profile-1"), + event, + ev_completed("resp-escalated-profile-1"), + ]), + ) + .await; + let results = mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-escalated-profile-1", "done"), + ev_completed("resp-escalated-profile-2"), + ]), + ) + .await; + + submit_turn_preserving_active_permission_profile( + &test, + "run the allowed script with escalated permissions", + approval_policy, + ) + .await?; + + wait_for_completion_without_approval(&test).await; + + let result = parse_result(&results.single_request().function_call_output(call_id)); + assert_eq!( + result, + CommandResult { + exit_code: Some(0), + stdout: format!("{BUILT_IN_PERMISSION_PROFILE_WORKSPACE}\n"), + }, + "the unsandboxed script should inherit CODEX_PERMISSION_PROFILE from the shell command" + ); + assert!( + outside_path.exists(), + "allowed escalated script should run outside the :workspace sandbox" + ); + + Ok(()) +} + +/// Verifies that zsh-fork applies an inner script's allow rule even when the +/// model invokes an outer wrapper, and that the escalated script retains the +/// named profile needed to reconstruct the original sandbox remotely without +/// dropping managed enterprise requirements. The script treats the inherited +/// environment value as untrusted and accepts only explicitly allowlisted +/// profile names before passing one to `codex sandbox -P`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg(unix)] +async fn zsh_fork_inner_allowed_script_inherits_active_permission_profile() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("zsh-fork remote sandbox wrapper test")? else { + return Ok(()); + }; + + const HOST: &str = "builder.example.com"; + let approval_policy = AskForApproval::OnRequest; + let script_dir = tempfile::tempdir_in(std::env::current_dir()?)?; + let wrapper_path = script_dir.path().join("remote-bash"); + let remote_bash_path = script_dir.path().join("remote_bash.py"); + let outside_path = script_dir.path().join("remote-bash-unsandboxed-marker"); + let outside_path_literal = serde_json::to_string(&outside_path.to_string_lossy())?; + fs::write( + &remote_bash_path, + format!( + r#"#!/usr/bin/env python3 +import argparse +import os +from pathlib import Path +import re +import shlex +import sys + +ALLOWED_HOSTS = ("builder.example.com",) +ALLOWED_PROFILES = (":workspace",) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Print an ssh command that recreates the current Codex sandbox remotely." + ) + parser.add_argument("--host", required=True) + try: + separator = sys.argv.index("--") + except ValueError: + parser.error("the remote command must follow --") + args = parser.parse_args(sys.argv[1:separator]) + args.command = sys.argv[separator + 1:] + if not args.command: + parser.error("the remote command must not be empty") + if not re.fullmatch(r"[a-z0-9.-]+", args.host) or args.host not in ALLOWED_HOSTS: + parser.error("host is not allowlisted") + return args + + +def main(): + args = parse_args() + profile_name = os.environ.get("CODEX_PERMISSION_PROFILE") + if not profile_name: + raise SystemExit("CODEX_PERMISSION_PROFILE must not be empty") + if profile_name not in ALLOWED_PROFILES: + raise SystemExit("CODEX_PERMISSION_PROFILE is not allowlisted") + + shell_command = shlex.join(args.command) + sandbox_command = shlex.join( + [ + "codex", + "sandbox", + "-P", + profile_name, + "--include-managed-config", + "--", + "bash", + "-lc", + shell_command, + ] + ) + print(shlex.join(["ssh", args.host, sandbox_command])) + + # Test-only proof that this inner script was allowed to run unsandboxed. + Path({outside_path_literal}).write_text("unsandboxed", encoding="utf-8") + + +if __name__ == "__main__": + main() +"# + ), + )?; + let remote_bash_exec = shlex::try_join([remote_bash_path.to_string_lossy().as_ref()])?; + fs::write( + &wrapper_path, + format!( + r#"#!/usr/bin/env zsh +exec {remote_bash_exec} "$@" +"# + ), + )?; + for path in [&wrapper_path, &remote_bash_path] { + let mut permissions = fs::metadata(path)?.permissions(); + permissions.set_mode(0o755); + fs::set_permissions(path, permissions)?; + } + + let remote_bash_pattern = serde_json::to_string(&remote_bash_path.to_string_lossy())?; + let rules = format!(r#"prefix_rule(pattern=[{remote_bash_pattern}], decision="allow")"#); + let server = start_mock_server().await; + let mut builder = + zsh_fork_test_builder(runtime, approval_policy).with_pre_build_hook(move |home| { + fs::write( + home.join("config.toml"), + WORKSPACE_PERMISSION_PROFILE_CONFIG, + ) + .expect("write config"); + let rules_dir = home.join("rules"); + fs::create_dir_all(&rules_dir).expect("create rules dir"); + fs::write(rules_dir.join("default.rules"), rules).expect("write rules"); + }); + let test = builder.build(&server).await?; + assert!(!outside_path.starts_with(test.config.cwd.as_path())); + assert_active_workspace_permission_profile(&test); + + let command = shlex::try_join([ + wrapper_path.to_string_lossy().as_ref(), + "--host", + HOST, + "--", + "printf", + "%s", + "hello world", + ])?; + let call_id = "zsh-fork-remote-sandbox-wrapper"; + let event = shell_event( + call_id, + &command, + /*timeout_ms*/ 30_000, + SandboxPermissions::UseDefault, + )?; + let _ = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-zsh-fork-remote-wrapper-1"), + event, + ev_completed("resp-zsh-fork-remote-wrapper-1"), + ]), + ) + .await; + let results = mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-zsh-fork-remote-wrapper-1", "done"), + ev_completed("resp-zsh-fork-remote-wrapper-2"), + ]), + ) + .await; + + submit_turn_preserving_active_permission_profile( + &test, + "run the remote sandbox wrapper", + approval_policy, + ) + .await?; + wait_for_completion_without_approval(&test).await; + + let result = parse_result(&results.single_request().function_call_output(call_id)); + assert_eq!( + result.exit_code.unwrap_or(0), + 0, + "the inner remote_bash.py script should run successfully: {}", + result.stdout + ); + let ssh_argv = shlex::split(result.stdout.trim()).context("parse printed ssh command")?; + assert_eq!( + ssh_argv.len(), + 3, + "expected ssh HOST LONG_COMMAND, got: {}", + result.stdout + ); + assert_eq!( + ssh_argv[..2], + ["ssh", HOST], + "remote_bash.py should target only the allowlisted host" + ); + let sandbox_argv = shlex::split(&ssh_argv[2]).context("parse remote sandbox command")?; + assert_eq!( + sandbox_argv.len(), + 9, + "expected codex sandbox ... bash -lc CMD" + ); + assert_eq!( + sandbox_argv[..8], + [ + "codex", + "sandbox", + "-P", + BUILT_IN_PERMISSION_PROFILE_WORKSPACE, + "--include-managed-config", + "--", + "bash", + "-lc", + ], + "remote_bash.py should use the allowlisted inherited profile and managed configuration to reconstruct the Codex sandbox" + ); + let command_argv = shlex::split(&sandbox_argv[8]).context("parse remote bash command")?; + assert_eq!( + command_argv, + ["printf", "%s", "hello world"], + "remote_bash.py should preserve every argument after --" + ); + assert!( + outside_path.exists(), + "the inner allowlisted script should run outside the :workspace sandbox" + ); + + Ok(()) +} + #[tokio::test(flavor = "current_thread")] #[cfg(unix)] async fn invalid_requested_prefix_rule_falls_back_for_compound_command() -> Result<()> { From e8d4a1a411bbab9be85e704ae3847915916bf347 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 25 Jun 2026 12:06:43 -0700 Subject: [PATCH 12/33] [codex] add current time reminder delivery mode config (#30031) ```python delivery_mode = "any_inference" # default delivery_mode = "after_user_or_tool_output" # new mode ``` ## Validation - just test -p codex-core load_config_resolves_current_time_reminder - just test -p codex-core lock_contains_prompts_and_materializes_features --- codex-rs/core/config.schema.json | 22 ++++++++++++++++++++++ codex-rs/core/src/config/config_tests.rs | 2 ++ codex-rs/core/src/config/mod.rs | 6 ++++++ codex-rs/core/src/session/config_lock.rs | 1 + codex-rs/features/src/feature_configs.rs | 13 +++++++++++++ codex-rs/features/src/lib.rs | 1 + 6 files changed, 45 insertions(+) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index bfdd179d78..512c050a20 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -810,6 +810,9 @@ "clock_source": { "$ref": "#/definitions/CurrentTimeSource" }, + "delivery_mode": { + "$ref": "#/definitions/CurrentTimeReminderDeliveryMode" + }, "enabled": { "type": "boolean" }, @@ -825,6 +828,25 @@ }, "type": "object" }, + "CurrentTimeReminderDeliveryMode": { + "description": "Which inference boundaries may receive current-time reminders.", + "oneOf": [ + { + "description": "Allow a reminder before any inference request once the interval is due.", + "enum": [ + "any_inference" + ], + "type": "string" + }, + { + "description": "Allow reminders after user input or tool output; new context windows still force one.", + "enum": [ + "after_user_or_tool_output" + ], + "type": "string" + } + ] + }, "CurrentTimeSource": { "enum": [ "system", diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 17779117f6..237e67d6c0 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -627,11 +627,13 @@ current_time_reminder = true enabled = true reminder_interval_seconds = 0 clock_source = "external" +delivery_mode = "after_user_or_tool_output" sleep_tool = true "#, CurrentTimeReminderConfig { reminder_interval_seconds: 0, clock_source: CurrentTimeSource::External, + delivery_mode: CurrentTimeReminderDeliveryMode::AfterUserOrToolOutput, sleep_tool: true, }, ), diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 58ae94906e..897a31c8db 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -59,6 +59,7 @@ use codex_exec_server::ExecutorFileSystem; use codex_exec_server::LOCAL_FS; use codex_features::CodeModeConfigToml; use codex_features::CurrentTimeReminderConfigToml; +use codex_features::CurrentTimeReminderDeliveryMode; use codex_features::CurrentTimeSource; use codex_features::Feature; use codex_features::FeatureConfigSource; @@ -1120,6 +1121,7 @@ pub struct RolloutBudgetConfig { pub struct CurrentTimeReminderConfig { pub reminder_interval_seconds: u64, pub clock_source: CurrentTimeSource, + pub delivery_mode: CurrentTimeReminderDeliveryMode, /// Whether to expose the input-interruptible `clock.sleep` tool. pub sleep_tool: bool, } @@ -1129,6 +1131,7 @@ impl Default for CurrentTimeReminderConfig { Self { reminder_interval_seconds: 1, clock_source: CurrentTimeSource::System, + delivery_mode: CurrentTimeReminderDeliveryMode::AnyInference, sleep_tool: false, } } @@ -2691,6 +2694,9 @@ fn resolve_current_time_reminder_config( clock_source: base .and_then(|config| config.clock_source) .unwrap_or(default.clock_source), + delivery_mode: base + .and_then(|config| config.delivery_mode) + .unwrap_or(default.delivery_mode), sleep_tool: base .and_then(|config| config.sleep_tool) .unwrap_or(default.sleep_tool), diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index b1adf94466..606d9d2be0 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -361,6 +361,7 @@ mod tests { enabled: Some(true), reminder_interval_seconds: Some(1), clock_source: Some(codex_features::CurrentTimeSource::System), + delivery_mode: Some(codex_features::CurrentTimeReminderDeliveryMode::AnyInference), sleep_tool: Some(false), })) ); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 24036eee58..9fc5df2d59 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -141,6 +141,17 @@ pub enum CurrentTimeSource { External, } +/// Which inference boundaries may receive current-time reminders. +#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum CurrentTimeReminderDeliveryMode { + /// Allow a reminder before any inference request once the interval is due. + #[default] + AnyInference, + /// Allow reminders after user input or tool output; new context windows still force one. + AfterUserOrToolOutput, +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct CurrentTimeReminderConfigToml { @@ -150,6 +161,8 @@ pub struct CurrentTimeReminderConfigToml { pub reminder_interval_seconds: Option, #[serde(skip_serializing_if = "Option::is_none")] pub clock_source: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub delivery_mode: Option, /// Expose the input-interruptible `clock.sleep` tool. #[serde(skip_serializing_if = "Option::is_none")] pub sleep_tool: Option, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 151e13f615..3bce1a7d12 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -18,6 +18,7 @@ mod feature_configs; mod legacy; pub use feature_configs::CodeModeConfigToml; pub use feature_configs::CurrentTimeReminderConfigToml; +pub use feature_configs::CurrentTimeReminderDeliveryMode; pub use feature_configs::CurrentTimeSource; pub use feature_configs::MultiAgentV2ConfigToml; pub use feature_configs::NetworkProxyConfigToml; From 964b138c3db16088a3bff5cc599bfd31428d03e9 Mon Sep 17 00:00:00 2001 From: richardopenai Date: Thu, 25 Jun 2026 12:25:04 -0700 Subject: [PATCH 13/33] [codex] Retry temporarily offline exec-server recovery (#30098) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - retry ERS `409 environment_offline` responses inside the existing exec-server recovery loop - keep all other registry conflicts terminal - add focused coverage for both cases ## Root cause When an exec server disconnects and reconnects, the client already starts recovery and calls ERS `/connect`. During the transient executor presence gap, ERS can return `409 environment_offline`. The retry classifier treated every 409 as terminal, so the first response aborted the existing 25-second recovery window before the executor came back online. That then caused active processes to be marked lost. This change classifies only the structured `environment_offline` conflict as retryable. Recovery continues with the existing bounded deadline, exponential backoff, and jitter. ## Validation - `just test -p codex-exec-server client::recovery::tests` — 4 passed - `just fix -p codex-exec-server` — passed - `just fmt` — passed - Full `just test -p codex-exec-server` reached unrelated macOS filesystem-sandbox integration failures because nested `/usr/bin/sandbox-exec` is denied in this environment (`sandbox_apply: Operation not permitted`). --- codex-rs/exec-server/src/client_recovery.rs | 7 +++++- .../exec-server/src/client_recovery_tests.rs | 24 +++++++++++++++---- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index f58e365728..61cee009cd 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -530,10 +530,15 @@ fn is_retryable_registry_error(error: &ExecServerError) -> bool { if error.is_connect() || error.is_timeout() ) || matches!( error, - ExecServerError::EnvironmentRegistryHttp { status, .. } + ExecServerError::EnvironmentRegistryHttp { status, code, .. } if status.is_server_error() || *status == reqwest::StatusCode::REQUEST_TIMEOUT || *status == reqwest::StatusCode::TOO_MANY_REQUESTS + // TODO: Replace this coarse retry with an explicit registry/presence + // recovery FSM so `environment_offline` is retried only while the + // executor is expected to reconnect. + || (*status == reqwest::StatusCode::CONFLICT + && code.as_deref() == Some("environment_offline")) ) } diff --git a/codex-rs/exec-server/src/client_recovery_tests.rs b/codex-rs/exec-server/src/client_recovery_tests.rs index a6ff91fc9c..b54b80e3c5 100644 --- a/codex-rs/exec-server/src/client_recovery_tests.rs +++ b/codex-rs/exec-server/src/client_recovery_tests.rs @@ -2,10 +2,10 @@ use std::time::Duration; use super::*; -fn registry_error() -> ExecServerError { +fn registry_error(status: reqwest::StatusCode, code: Option<&str>) -> ExecServerError { ExecServerError::EnvironmentRegistryHttp { - status: reqwest::StatusCode::TOO_MANY_REQUESTS, - code: None, + status, + code: code.map(str::to_string), message: "registry unavailable".to_string(), } } @@ -33,8 +33,24 @@ fn registry_recovery_retry_delay_exponentially_backs_off_and_caps() { #[test] fn recovery_retries_transient_registry_errors() { - let error = registry_error(); + let error = registry_error(reqwest::StatusCode::TOO_MANY_REQUESTS, /*code*/ None); assert!(is_retryable_registry_error(&error)); assert!(is_retryable_recovery_error(&error)); } + +#[test] +fn recovery_retries_environment_offline_conflicts() { + let error = registry_error(reqwest::StatusCode::CONFLICT, Some("environment_offline")); + + assert!(is_retryable_registry_error(&error)); + assert!(is_retryable_recovery_error(&error)); +} + +#[test] +fn recovery_does_not_retry_other_registry_conflicts() { + let error = registry_error(reqwest::StatusCode::CONFLICT, Some("registration_conflict")); + + assert!(!is_retryable_registry_error(&error)); + assert!(!is_retryable_recovery_error(&error)); +} From adccb464d075a39d5450d6bc879e3bb6c97ce14b Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 25 Jun 2026 12:28:50 -0700 Subject: [PATCH 14/33] [codex] impl delivery_mode: current time reminders on response boundaries (#30033) ## Summary - track user-like input and tool-output boundaries in current-time reminder state - gate reminder injection when delivery_mode is after_user_or_tool_output - preserve interval debounce and forced reminders after context-window changes ## Why Training can request reminders only after user or tool-output items while keeping the existing canonical pre-inference history-injection path. ## Validation - just test -p codex-core current_time_reminders_can_follow_only_user_or_tool_outputs - just test -p codex-core current_time_reminders_follow_time_interval_and_persist_in_history - just test -p codex-core current_time_reminder_is_refreshed_after_compaction - just fix -p codex-core --- codex-rs/core/src/session/mod.rs | 2 + codex-rs/core/src/session/time_reminder.rs | 33 ++++++++- .../core/tests/suite/current_time_reminder.rs | 69 +++++++++++++++++++ 3 files changed, 103 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 442722003b..6fe61222cf 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2775,6 +2775,7 @@ impl Session { let items = items.as_ref(); { let mut state = self.state.lock().await; + state.current_time_reminder.note_recorded_items(items); state.record_items( items.iter(), turn_context.model_info.truncation_policy.into(), @@ -2877,6 +2878,7 @@ impl Session { let response_item = items[0].clone(); { let mut state = self.state.lock().await; + state.current_time_reminder.note_recorded_items(items); state.record_items( items.iter(), turn_context.model_info.truncation_policy.into(), diff --git a/codex-rs/core/src/session/time_reminder.rs b/codex-rs/core/src/session/time_reminder.rs index 8f612e8fcb..398ae58114 100644 --- a/codex-rs/core/src/session/time_reminder.rs +++ b/codex-rs/core/src/session/time_reminder.rs @@ -1,26 +1,56 @@ use chrono::DateTime; use chrono::Utc; +use codex_features::CurrentTimeReminderDeliveryMode; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; +use codex_protocol::models::ResponseItem; use super::session::Session; use super::turn_context::TurnContext; use crate::context::ContextualUserFragment; +use crate::context_manager::is_user_turn_boundary; #[derive(Default)] pub(crate) struct CurrentTimeReminderState { last_delivery_time: Option>, last_window_id: Option, + pending_user_or_tool_output_boundary: bool, } impl CurrentTimeReminderState { + pub(super) fn note_recorded_items(&mut self, items: &[ResponseItem]) { + if items.iter().any(|item| { + is_user_turn_boundary(item) + || matches!( + item, + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::ToolSearchOutput { .. } + ) + }) { + self.pending_user_or_tool_output_boundary = true; + } + } + fn take_reminder_due( &mut self, window_id: &str, current_time: DateTime, interval_seconds: u64, + delivery_mode: CurrentTimeReminderDeliveryMode, ) -> bool { - let reminder_is_due = self.last_window_id.as_deref() != Some(window_id) + let is_new_window = self.last_window_id.as_deref() != Some(window_id); + // Consume the boundary for this inference even if the interval suppresses delivery. + let follows_user_or_tool_output = + std::mem::take(&mut self.pending_user_or_tool_output_boundary); + if delivery_mode == CurrentTimeReminderDeliveryMode::AfterUserOrToolOutput + && !is_new_window + && !follows_user_or_tool_output + { + return false; + } + + let reminder_is_due = is_new_window || interval_seconds == 0 || self.last_delivery_time.is_none_or(|last_delivery_time| { current_time @@ -60,6 +90,7 @@ pub(super) async fn maybe_record_current_time_reminder( window_id, current_time, config.reminder_interval_seconds, + config.delivery_mode, ) }; if !reminder_is_due { diff --git a/codex-rs/core/tests/suite/current_time_reminder.rs b/codex-rs/core/tests/suite/current_time_reminder.rs index c5ebd0ed13..9a9440fb33 100644 --- a/codex-rs/core/tests/suite/current_time_reminder.rs +++ b/codex-rs/core/tests/suite/current_time_reminder.rs @@ -12,6 +12,7 @@ use codex_core::SleepFuture; use codex_core::TimeFuture; use codex_core::TimeProvider; use codex_core::config::CurrentTimeReminderConfig; +use codex_features::CurrentTimeReminderDeliveryMode; use codex_features::CurrentTimeSource; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; @@ -201,6 +202,69 @@ async fn zero_current_time_reminder_interval_delivers_when_time_moves_backward() Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn current_time_reminders_can_follow_only_user_or_tool_outputs() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let tool_args = json!({ + "command": "echo current time", + "timeout_ms": 1_000, + }); + let mut continue_response = ev_completed("resp-2"); + // Ask for another inference without recording a new user message or tool output. + continue_response["response"]["end_turn"] = json!(false); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + "current-time-tool-call", + "shell_command", + &serde_json::to_string(&tool_args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "continue"), + continue_response, + ]), + sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + enable_current_time_reminder(config, /*interval*/ 0, CurrentTimeSource::External); + config + .current_time_reminder + .as_mut() + .expect("current-time reminder should be configured") + .delivery_mode = CurrentTimeReminderDeliveryMode::AfterUserOrToolOutput; + }) + .with_external_time_provider(Arc::new(TestTimeProvider::default())) + .build(&server) + .await?; + + test.submit_turn_with_permission_profile("first turn", PermissionProfile::Disabled) + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert_eq!(current_time_reminders(&requests[0]), vec![FIRST_REMINDER]); + assert_eq!( + current_time_reminders(&requests[1]), + vec![FIRST_REMINDER, SECOND_REMINDER] + ); + assert_eq!( + current_time_reminders(&requests[2]), + vec![FIRST_REMINDER, SECOND_REMINDER] + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn system_time_source_adds_current_time_reminder() -> Result<()> { skip_if_no_network!(Ok(())); @@ -260,6 +324,11 @@ async fn current_time_reminder_is_refreshed_after_compaction() -> Result<()> { /*interval*/ 3_000, CurrentTimeSource::External, ); + config + .current_time_reminder + .as_mut() + .expect("current-time reminder should be configured") + .delivery_mode = CurrentTimeReminderDeliveryMode::AfterUserOrToolOutput; }) .with_external_time_provider(Arc::new(TestTimeProvider::default())) .build(&server) From 3b78f58fb2e672fc1e5f808963dc65544f571392 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 13:26:47 -0700 Subject: [PATCH 15/33] [codex] extend code-mode host IPC transport (#30108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - add an `EncodedFrame` type so IPC payloads are serialized and size-checked before entering bounded queues - add the V1 `operation/cancel` client-to-host message - pin the new wire shape with protocol tests ## Why The process-owned code-mode host needs bounded, pre-encoded outbound messages and a best-effort cancellation signal. Keeping these wire primitives in a protocol-only change lets their compatibility contract be reviewed independently from either endpoint. ## Stack This is **1 of 4** in the process-owned code-mode session stack. The next PR targets this branch. ## Validation - `just test -p codex-code-mode-protocol` — 22 passed - `just fix -p codex-code-mode-protocol` - `just fmt` --- codex-rs/code-mode-protocol/src/host/codec.rs | 54 +++++++++++++------ .../code-mode-protocol/src/host/host_tests.rs | 10 ++++ .../code-mode-protocol/src/host/message.rs | 2 + codex-rs/code-mode-protocol/src/host/mod.rs | 1 + 4 files changed, 50 insertions(+), 17 deletions(-) diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs index 6bf83d367d..ec07b92d41 100644 --- a/codex-rs/code-mode-protocol/src/host/codec.rs +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -11,6 +11,36 @@ use tokio::io::AsyncWriteExt; /// Maximum JSON payload size accepted for one IPC frame. pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024; +/// A serialized IPC frame that has already passed the payload size limit. +#[derive(Clone, Debug)] +pub struct EncodedFrame { + payload: Vec, +} + +impl EncodedFrame { + pub fn encode(message: &T) -> io::Result + where + T: Serialize, + { + let payload = serde_json::to_vec(message).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to encode code-mode IPC frame: {err}"), + ) + })?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", + payload.len() + ), + )); + } + Ok(Self { payload }) + } +} + /// Decodes JSON messages prefixed by a four-byte little-endian payload length. pub struct FramedReader { reader: R, @@ -72,22 +102,12 @@ where where T: Serialize, { - let payload = serde_json::to_vec(message).map_err(|err| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("failed to encode code-mode IPC frame: {err}"), - ) - })?; - if payload.len() > MAX_FRAME_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!( - "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", - payload.len() - ), - )); - } - let length = u32::try_from(payload.len()).map_err(|_| { + self.write_frame(&EncodedFrame::encode(message)?).await + } + + /// Writes and flushes a frame encoded before it entered an I/O queue. + pub async fn write_frame(&mut self, frame: &EncodedFrame) -> io::Result<()> { + let length = u32::try_from(frame.payload.len()).map_err(|_| { io::Error::new( io::ErrorKind::InvalidData, "code-mode IPC frame length exceeds u32", @@ -95,7 +115,7 @@ where })?; self.writer.write_all(&length.to_le_bytes()).await?; - self.writer.write_all(&payload).await?; + self.writer.write_all(&frame.payload).await?; self.writer.flush().await } } diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 3aa2a18307..d56e56e1fa 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -365,6 +365,16 @@ fn client_to_host_v1_variants_are_pinned() { }), ); } + + assert_wire_round_trip( + ClientToHost::CancelRequest { + id: request_id(/*value*/ 9), + }, + json!({ + "type": "operation/cancel", + "id": 9, + }), + ); } #[test] diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index 78f0961d51..0e83c866e4 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -132,6 +132,8 @@ pub enum ClientToHost { ClientHello(ClientHello), #[serde(rename = "operation/request")] Request { id: RequestId, request: HostRequest }, + #[serde(rename = "operation/cancel")] + CancelRequest { id: RequestId }, #[serde(rename = "delegate/response")] DelegateResponse { id: DelegateRequestId, diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index 7e78768ab6..3583ccb3e9 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -11,6 +11,7 @@ mod message; mod payload; mod types; +pub use codec::EncodedFrame; pub use codec::FramedReader; pub use codec::FramedWriter; pub use codec::MAX_FRAME_BYTES; From 3b22498f69112eb1f7dd3061e278e344626e834c Mon Sep 17 00:00:00 2001 From: richardopenai Date: Thu, 25 Jun 2026 13:42:40 -0700 Subject: [PATCH 16/33] [codex] Observe remote exec-server lifecycle (#27470) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Record bounded duration and outcome metrics for remote environment registration and Noise rendezvous connection attempts. - Count reconnects by bounded reason: disconnect, connection failure, or rejected registration. - Trace registration at the owning client boundary without exporting raw environment or registration identifiers. - Replace the stale pre-Noise WebSocket observability design with the current remote transport model. ## Stack Review and land this stack in order: 1. #27466 — trace exec-server JSON-RPC requests 2. #27467 — record bounded connection, request, and process lifecycle metrics 3. #27470 — observe remote registration and Noise rendezvous lifecycle **(this PR)** ## Validation - `just test -p codex-exec-server --lib` (149 passed) - `just test -p codex-cli --test exec_server` (4 passed) - `just argument-comment-lint` - `just bazel-lock-check` - `just fix -p codex-exec-server -p codex-cli` - `just fmt` --- codex-rs/cli/src/exec_server_telemetry.rs | 66 +++++++++ codex-rs/cli/src/main.rs | 16 ++- codex-rs/cli/tests/exec_server.rs | 62 +++++++++ codex-rs/exec-server/src/client_transport.rs | 22 ++- codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/remote.rs | 129 +++++++++++++++--- codex-rs/exec-server/src/telemetry.rs | 58 ++++++++ codex-rs/exec-server/src/trace_context.rs | 24 ++++ .../exec-server/src/trace_context_tests.rs | 27 ++++ 9 files changed, 381 insertions(+), 24 deletions(-) create mode 100644 codex-rs/exec-server/src/trace_context.rs create mode 100644 codex-rs/exec-server/src/trace_context_tests.rs diff --git a/codex-rs/cli/src/exec_server_telemetry.rs b/codex-rs/cli/src/exec_server_telemetry.rs index 1645bd5744..b5af769e6a 100644 --- a/codex-rs/cli/src/exec_server_telemetry.rs +++ b/codex-rs/cli/src/exec_server_telemetry.rs @@ -1,3 +1,5 @@ +use std::future::Future; + use tracing_subscriber::EnvFilter; use tracing_subscriber::prelude::*; @@ -43,6 +45,70 @@ pub(crate) fn init( (otel, telemetry) } +pub(crate) async fn run_until_shutdown(run: F) -> Result<(), E> +where + F: Future>, +{ + let shutdown_signal = match shutdown_signal() { + Ok(signal) => Some(signal), + Err(error) => { + eprintln!("Could not listen for exec-server shutdown signal: {error}"); + None + } + }; + tokio::pin!(run); + + if let Some(shutdown_signal) = shutdown_signal { + tokio::select! { + result = &mut run => result, + signal = wait_for_shutdown_signal(shutdown_signal) => { + match signal { + Ok(()) => Ok(()), + Err(error) => { + eprintln!("Could not listen for exec-server shutdown signal: {error}"); + run.await + } + } + } + } + } else { + run.await + } +} + +#[cfg(unix)] +struct ShutdownSignal { + terminate: tokio::signal::unix::Signal, +} + +#[cfg(unix)] +fn shutdown_signal() -> std::io::Result { + Ok(ShutdownSignal { + terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?, + }) +} + +#[cfg(unix)] +async fn wait_for_shutdown_signal(mut shutdown_signal: ShutdownSignal) -> std::io::Result<()> { + tokio::select! { + result = tokio::signal::ctrl_c() => result, + _ = shutdown_signal.terminate.recv() => Ok(()), + } +} + +#[cfg(not(unix))] +struct ShutdownSignal; + +#[cfg(not(unix))] +fn shutdown_signal() -> std::io::Result { + Ok(ShutdownSignal) +} + +#[cfg(not(unix))] +async fn wait_for_shutdown_signal(_: ShutdownSignal) -> std::io::Result<()> { + tokio::signal::ctrl_c().await +} + fn stderr_env_filter() -> EnvFilter { EnvFilter::try_from_default_env() .or_else(|_| EnvFilter::try_new(DEFAULT_LOG_FILTER)) diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index d7dbff47ec..c766baa048 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1706,7 +1706,10 @@ async fn run_exec_server_command( remote_config.name = name; } let remote_config = remote_config.with_telemetry(telemetry); - codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; + exec_server_telemetry::run_until_shutdown(async move { + codex_exec_server::run_remote_environment(remote_config, runtime_paths).await + }) + .await?; Ok(()) } else { let config_result = load_exec_server_config(root_config_overrides, strict_config).await; @@ -1718,11 +1721,12 @@ async fn run_exec_server_command( let (_otel, telemetry) = exec_server_telemetry::init(config.as_ref()); let listen_url = cmd .listen - .as_deref() - .unwrap_or(codex_exec_server::DEFAULT_LISTEN_URL); - codex_exec_server::run_main_with_telemetry(listen_url, runtime_paths, telemetry) - .await - .map_err(anyhow::Error::from_boxed) + .unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string()); + exec_server_telemetry::run_until_shutdown(async move { + codex_exec_server::run_main_with_telemetry(&listen_url, runtime_paths, telemetry).await + }) + .await + .map_err(anyhow::Error::from_boxed) } } diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index c54e9c40b2..aabe9c0453 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -1,6 +1,20 @@ +#[cfg(unix)] +use std::io::BufRead as _; +#[cfg(unix)] +use std::io::BufReader as StdBufReader; +#[cfg(unix)] +use std::io::Read as _; +#[cfg(unix)] +use std::io::Write as _; +#[cfg(unix)] +use std::net::TcpStream; use std::path::Path; use std::process::Stdio; +#[cfg(unix)] +use std::thread; use std::time::Duration; +#[cfg(unix)] +use std::time::Instant; use anyhow::Result; use predicates::prelude::PredicateBooleanExt; @@ -220,6 +234,54 @@ async fn send_json_line( Ok(()) } +#[cfg(unix)] +#[test] +fn local_exec_server_exits_successfully_on_sigterm() -> Result<()> { + let codex_home = TempDir::new()?; + let mut child = std::process::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?) + .env("CODEX_HOME", codex_home.path()) + .args(["exec-server", "--listen", "ws://127.0.0.1:0"]) + .stdout(Stdio::piped()) + .spawn()?; + let mut listen_url = String::new(); + StdBufReader::new(child.stdout.take().expect("child stdout")).read_line(&mut listen_url)?; + assert!(listen_url.starts_with("ws://127.0.0.1:"), "{listen_url}"); + + let listen_addr = listen_url + .trim() + .strip_prefix("ws://") + .expect("listen URL should use ws://") + .parse()?; + let deadline = Instant::now() + Duration::from_secs(5); + let mut ready = false; + while let Some(remaining) = deadline.checked_duration_since(Instant::now()) { + if let Ok(mut stream) = + TcpStream::connect_timeout(&listen_addr, remaining.min(Duration::from_millis(100))) + { + let _ = stream.set_read_timeout(Some(Duration::from_secs(1))); + let request = + format!("GET /readyz HTTP/1.1\r\nHost: {listen_addr}\r\nConnection: close\r\n\r\n"); + let mut response = String::new(); + if stream.write_all(request.as_bytes()).is_ok() + && stream.read_to_string(&mut response).is_ok() + && response.starts_with("HTTP/1.1 200") + { + ready = true; + break; + } + } + thread::sleep(Duration::from_millis(10)); + } + assert!(ready, "exec-server did not become ready at {listen_url}"); + + // SAFETY: `child.id()` is the live process spawned above. + let result = unsafe { libc::kill(child.id() as libc::pid_t, libc::SIGTERM) }; + assert_eq!(result, 0); + let status = child.wait()?; + assert!(status.success(), "{status}"); + Ok(()) +} + async fn wait_for_response( stdout: &mut (impl tokio::io::AsyncBufRead + Unpin), expected_id: i64, diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index c6cff4ed30..752bf27037 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -8,6 +8,7 @@ use tokio::process::Command; use tokio::time::timeout; use tokio_tungstenite::connect_async; use tokio_tungstenite::connect_async_with_config; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::warn; @@ -30,6 +31,7 @@ use crate::noise_relay::NoiseHarnessConnectionArgs; use crate::noise_relay::noise_harness_connection_from_websocket; use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; +use crate::trace_context::current_trace_context_headers; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; @@ -216,6 +218,14 @@ impl ExecServerClient { /// only ciphertext after that. Environment-managed connections use a /// retained [`NoiseRendezvousConnectProvider`] so recovery can fetch a fresh /// bundle for each reconnect. + #[tracing::instrument( + name = "codex.exec_server.remote.harness.connect", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.connect", + ) + )] pub async fn connect_noise_rendezvous( args: NoiseRendezvousConnectArgs, ) -> Result { @@ -249,10 +259,20 @@ impl ExecServerClient { .next() .unwrap_or(websocket_url.as_str()) .to_string(); + let mut request = websocket_url + .as_str() + .into_client_request() + .map_err(|source| ExecServerError::WebSocketConnect { + url: diagnostic_url.clone(), + source, + })?; + request + .headers_mut() + .extend(current_trace_context_headers()); let (stream, _) = timeout( connect_timeout, connect_async_with_config( - websocket_url.as_str(), + request, Some(noise_relay_websocket_config()), /*disable_nagle*/ false, ), diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 87593374cb..0068c2cc1a 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -28,6 +28,7 @@ mod runtime_paths; mod sandboxed_file_system; mod server; mod telemetry; +mod trace_context; use codex_exec_server_protocol as protocol; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 6609b6cd92..f5e233ce9d 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -1,5 +1,6 @@ use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; @@ -10,7 +11,10 @@ use http::HeaderValue; use reqwest::StatusCode; use serde::Deserialize; use tokio::time::sleep; +use tokio_tungstenite::MaybeTlsStream; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async_with_config; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::info; use tracing::warn; @@ -35,6 +39,7 @@ use crate::noise_relay::noise_relay_websocket_config; use crate::relay::HarnessKeyValidator; use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; +use crate::trace_context::current_trace_context_headers; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; @@ -45,6 +50,7 @@ struct EnvironmentRegistryClient { auth_provider: SharedAuthProvider, http: reqwest::Client, connect_timeout: Duration, + telemetry: ExecServerTelemetry, } impl std::fmt::Debug for EnvironmentRegistryClient { @@ -57,7 +63,16 @@ impl std::fmt::Debug for EnvironmentRegistryClient { } impl EnvironmentRegistryClient { + #[cfg(test)] fn new(base_url: String, auth_provider: SharedAuthProvider) -> Result { + Self::new_with_telemetry(base_url, auth_provider, ExecServerTelemetry::default()) + } + + fn new_with_telemetry( + base_url: String, + auth_provider: SharedAuthProvider, + telemetry: ExecServerTelemetry, + ) -> Result { let base_url = normalize_base_url(base_url)?; Ok(Self { base_url, @@ -66,15 +81,41 @@ impl EnvironmentRegistryClient { .redirect(reqwest::redirect::Policy::none()) .build()?, connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + telemetry, }) } /// Register the executor public key and obtain the rendezvous allocation. /// The returned registration ID is included in each stream's Noise prologue. + #[tracing::instrument( + name = "codex.exec_server.remote.register", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.register", + result = tracing::field::Empty, + ) + )] async fn register_environment( &self, environment_id: &str, executor_public_key: &NoiseChannelPublicKey, + ) -> Result { + let started_at = Instant::now(); + let response = self + .register_environment_inner(environment_id, executor_public_key) + .await; + let result = if response.is_ok() { "success" } else { "error" }; + tracing::Span::current().record("result", result); + self.telemetry + .remote_registration_completed(result, started_at.elapsed()); + response + } + + async fn register_environment_inner( + &self, + environment_id: &str, + executor_public_key: &NoiseChannelPublicKey, ) -> Result { let response = self .http @@ -83,6 +124,7 @@ impl EnvironmentRegistryClient { &format!("/cloud/environment/{environment_id}/register"), )) .headers(self.auth_provider.to_auth_headers()) + .headers(current_trace_context_headers()) .json(&EnvironmentRegistryRegistrationRequest { security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), executor_public_key: executor_public_key.clone(), @@ -269,7 +311,11 @@ impl NoiseRendezvousEnvironmentConfig { ) -> Result { let environment_id = normalize_environment_id(environment_id)?; let auth_provider = static_bearer_auth_provider(bearer_token, chatgpt_account_id)?; - let client = EnvironmentRegistryClient::new(base_url, auth_provider)?; + let client = EnvironmentRegistryClient::new_with_telemetry( + base_url, + auth_provider, + ExecServerTelemetry::default(), + )?; Ok(Self { provider: Arc::new(EnvironmentRegistryNoiseConnectProvider { client, @@ -416,18 +462,16 @@ impl RemoteEnvironmentConfig { /// reconnects. The registration and rendezvous URL are also reused until /// rendezvous rejects the URL, at which point the next attempt registers again. /// The websocket carries cleartext routing metadata and encrypted payloads. -#[tracing::instrument( - name = "codex.exec_server", - skip_all, - fields(otel.kind = "internal") -)] pub async fn run_remote_environment( config: RemoteEnvironmentConfig, runtime_paths: ExecServerRuntimePaths, ) -> Result<(), ExecServerError> { ensure_rustls_crypto_provider(); - let client = - EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; + let client = EnvironmentRegistryClient::new_with_telemetry( + config.base_url.clone(), + config.auth_provider.clone(), + config.telemetry.clone(), + )?; let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, config.telemetry.clone()); let identity = NoiseChannelIdentity::generate().map_err(|error| { @@ -439,14 +483,8 @@ pub async fn run_remote_environment( .await?; loop { - match connect_async_with_config( - response.url.as_str(), - Some(noise_relay_websocket_config()), - /*disable_nagle*/ false, - ) - .await - { - Ok((websocket, _)) => { + match connect_rendezvous(&response.url, &config.telemetry).await { + Ok(websocket) => { backoff = Duration::from_secs(1); let executor_registration_id = response.executor_registration_id.clone(); info!( @@ -467,6 +505,7 @@ pub async fn run_remote_environment( }, ) .await; + config.telemetry.remote_reconnect("disconnected"); } Err(error) => { let registration_rejected = matches!( @@ -482,9 +521,12 @@ pub async fn run_remote_environment( ); debug!(error = %error, "Noise executor rendezvous connection error"); if registration_rejected { + config.telemetry.remote_reconnect("registration_rejected"); response = client .register_environment(&config.environment_id, &identity.public_key()) .await?; + } else { + config.telemetry.remote_reconnect("connect_failed"); } } } @@ -494,6 +536,43 @@ pub async fn run_remote_environment( } } +#[tracing::instrument( + name = "codex.exec_server.remote.rendezvous.connect", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.rendezvous.connect", + result = tracing::field::Empty, + ) +)] +async fn connect_rendezvous( + url: &str, + telemetry: &ExecServerTelemetry, +) -> Result< + WebSocketStream>, + tokio_tungstenite::tungstenite::Error, +> { + let started_at = Instant::now(); + let result = async { + let mut request = url.into_client_request()?; + request + .headers_mut() + .extend(current_trace_context_headers()); + connect_async_with_config( + request, + Some(noise_relay_websocket_config()), + /*disable_nagle*/ false, + ) + .await + .map(|(websocket, _)| websocket) + } + .await; + let result_name = if result.is_ok() { "success" } else { "error" }; + tracing::Span::current().record("result", result_name); + telemetry.remote_rendezvous_completed(result_name, started_at.elapsed()); + result +} + fn normalize_environment_id(environment_id: String) -> Result { let environment_id = environment_id.trim().to_string(); if environment_id.is_empty() { @@ -585,12 +664,17 @@ mod tests { use codex_api::AuthProvider; use http::HeaderMap; use http::HeaderValue; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; + use tracing::Instrument; + use tracing_subscriber::prelude::*; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::body_partial_json; use wiremock::matchers::header; + use wiremock::matchers::header_regex; use wiremock::matchers::method; use wiremock::matchers::path; @@ -616,8 +700,14 @@ mod tests { Arc::new(StaticRegistryAuthProvider) } - #[tokio::test] + #[tokio::test(flavor = "current_thread")] async fn register_environment_posts_with_auth_provider_headers() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); let server = MockServer::start().await; let executor_public_key = NoiseChannelIdentity::generate() .expect("identity") @@ -626,6 +716,10 @@ mod tests { .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) + .and(header_regex( + "traceparent", + "^00-[0-9a-f]{32}-[0-9a-f]{16}-0[01]$", + )) .and(body_partial_json(serde_json::json!({ "security_profile": NOISE_RELAY_SECURITY_PROFILE, "executor_public_key": executor_public_key.clone(), @@ -643,6 +737,7 @@ mod tests { let response = client .register_environment("environment-requested", &executor_public_key) + .instrument(tracing::info_span!("remote-operation")) .await .expect("register environment"); diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index 4624a2ff2c..650faacb2c 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/codex-rs/exec-server/src/telemetry.rs @@ -21,6 +21,28 @@ const PROCESSES_FINISHED_TOTAL_DESCRIPTION: &str = "Total number of finished exec-server processes."; const PROCESS_DURATION_METRIC: &str = "exec_server_process_duration_seconds"; const PROCESS_DURATION_DESCRIPTION: &str = "Duration of exec-server processes in seconds."; +const REMOTE_REGISTRATION_METRICS: OperationMetrics = OperationMetrics { + total_name: "exec_server_remote_registration_total", + total_description: "Total number of remote exec-server registration attempts.", + duration_name: "exec_server_remote_registration_duration_seconds", + duration_description: "Duration of remote exec-server registration attempts in seconds.", +}; +const REMOTE_RENDEZVOUS_METRICS: OperationMetrics = OperationMetrics { + total_name: "exec_server_remote_rendezvous_connect_total", + total_description: "Total number of remote exec-server rendezvous connection attempts.", + duration_name: "exec_server_remote_rendezvous_connect_duration_seconds", + duration_description: "Duration of remote exec-server rendezvous connection attempts in seconds.", +}; +const REMOTE_RECONNECTS_TOTAL_METRIC: &str = "exec_server_remote_reconnects_total"; +const REMOTE_RECONNECTS_TOTAL_DESCRIPTION: &str = "Total number of remote exec-server reconnects."; + +#[derive(Clone, Copy)] +struct OperationMetrics { + total_name: &'static str, + total_description: &'static str, + duration_name: &'static str, + duration_description: &'static str, +} #[derive(Clone, Copy)] pub(crate) enum ConnectionTransport { @@ -123,6 +145,24 @@ impl ExecServerTelemetry { }); } + pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) { + self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration); + } + + pub(crate) fn remote_rendezvous_completed(&self, result: &'static str, duration: Duration) { + self.record_operation(REMOTE_RENDEZVOUS_METRICS, result, duration); + } + + pub(crate) fn remote_reconnect(&self, reason: &'static str) { + self.with_inner(|inner| { + inner.counter( + REMOTE_RECONNECTS_TOTAL_METRIC, + REMOTE_RECONNECTS_TOTAL_DESCRIPTION, + &[("reason", reason)], + ); + }); + } + pub(crate) fn process_started(&self) -> ProcessMetricGuard { self.with_inner(|inner| { inner.adjust_process_count(/*delta*/ 1); @@ -162,6 +202,24 @@ impl ExecServerTelemetry { emit(inner); } } + + fn record_operation( + &self, + metrics: OperationMetrics, + result: &'static str, + duration: Duration, + ) { + self.with_inner(|inner| { + let tags = [("result", result)]; + inner.counter(metrics.total_name, metrics.total_description, &tags); + inner.duration( + metrics.duration_name, + metrics.duration_description, + duration, + &tags, + ); + }); + } } impl Drop for ConnectionMetricGuard { diff --git a/codex-rs/exec-server/src/trace_context.rs b/codex-rs/exec-server/src/trace_context.rs new file mode 100644 index 0000000000..70a8f3c61e --- /dev/null +++ b/codex-rs/exec-server/src/trace_context.rs @@ -0,0 +1,24 @@ +use reqwest::header::HeaderMap; +use reqwest::header::HeaderValue; + +pub(crate) fn current_trace_context_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + let Some(trace) = codex_otel::current_span_w3c_trace_context() else { + return headers; + }; + if let Some(traceparent) = trace.traceparent + && let Ok(value) = HeaderValue::try_from(traceparent) + { + headers.insert("traceparent", value); + } + if let Some(tracestate) = trace.tracestate + && let Ok(value) = HeaderValue::try_from(tracestate) + { + headers.insert("tracestate", value); + } + headers +} + +#[cfg(test)] +#[path = "trace_context_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/trace_context_tests.rs b/codex-rs/exec-server/src/trace_context_tests.rs new file mode 100644 index 0000000000..92a820afc4 --- /dev/null +++ b/codex-rs/exec-server/src/trace_context_tests.rs @@ -0,0 +1,27 @@ +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::SdkTracerProvider; +use tracing_subscriber::prelude::*; + +use super::current_trace_context_headers; + +#[test] +fn creates_traceparent_header_from_current_span() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); + let span = tracing::info_span!("outbound-request"); + let _entered = span.enter(); + + let headers = current_trace_context_headers(); + + let traceparent = headers + .get("traceparent") + .expect("traceparent header") + .to_str() + .expect("valid traceparent header"); + assert!(traceparent.starts_with("00-")); + assert_eq!(traceparent.len(), 55); +} From 62c7f506d9a8892724ee025da04bfef74f092ccb Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 25 Jun 2026 13:46:42 -0700 Subject: [PATCH 17/33] [codex] poll external clock during sleep (#30113) ## Summary - make the external app-server time provider establish sleep deadlines using `currentTime/read` - poll the external clock once per second and complete `clock.sleep` when the deadline is reached - keep the system-clock timer and existing steer/agent-message interruption behavior unchanged ## Why This lets training control `clock.sleep` through its existing external simulated clock without adding separate sleep/wake protocol methods. ## Testing - `just fmt` - `just test -p codex-app-server external_sleep_polls_current_time_and_emits_items` --- codex-rs/app-server/src/current_time.rs | 29 ++++- codex-rs/app-server/tests/suite/v2/sleep.rs | 115 ++++++++++++++------ 2 files changed, 106 insertions(+), 38 deletions(-) diff --git a/codex-rs/app-server/src/current_time.rs b/codex-rs/app-server/src/current_time.rs index 205068be2b..ae873bfc63 100644 --- a/codex-rs/app-server/src/current_time.rs +++ b/codex-rs/app-server/src/current_time.rs @@ -23,6 +23,7 @@ use crate::outgoing_message::OutgoingMessageSender; use crate::thread_state::ThreadStateManager; const CURRENT_TIME_REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const CURRENT_TIME_POLL_INTERVAL: Duration = Duration::from_secs(1); pub(crate) fn app_server_time_provider( outgoing: Arc, @@ -51,10 +52,32 @@ impl TimeProvider for AppServerTimeProvider { }) } - fn sleep(&self, _thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + fn sleep(&self, thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + let outgoing = self.outgoing.clone(); + let thread_state_manager = self.thread_state_manager.clone(); Box::pin(async move { - tokio::time::sleep(duration).await; - Ok(()) + let outgoing = outgoing + .upgrade() + .context("app-server current-time provider is unavailable")?; + let started_at = + request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await?; + let wake_at = started_at + .checked_add_signed( + chrono::Duration::from_std(duration) + .context("external sleep duration is outside the supported range")?, + ) + .context("external sleep deadline is outside the supported range")?; + + loop { + tokio::time::sleep(CURRENT_TIME_POLL_INTERVAL).await; + if request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await? + >= wake_at + { + return Ok(()); + } + } }) } } diff --git a/codex-rs/app-server/tests/suite/v2/sleep.rs b/codex-rs/app-server/tests/suite/v2/sleep.rs index 872983f76d..8d88816fde 100644 --- a/codex-rs/app-server/tests/suite/v2/sleep.rs +++ b/codex-rs/app-server/tests/suite/v2/sleep.rs @@ -1,11 +1,12 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::to_response; +use codex_app_server_protocol::CurrentTimeReadResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -19,12 +20,16 @@ use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; +#[cfg(windows)] +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(25); +#[cfg(not(windows))] const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +const CURRENT_TIME_AT: i64 = 1_781_717_655; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sleep_emits_started_and_completed_items() -> Result<()> { +async fn external_sleep_polls_current_time_and_emits_items() -> Result<()> { const CALL_ID: &str = "sleep-1"; - const DURATION_MS: u64 = 1; + const DURATION_MS: u64 = 2_000; let server = responses::start_mock_server().await; responses::mount_sse_sequence( @@ -85,38 +90,19 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { .await??; let TurnStartResponse { turn, .. } = to_response(turn_start_response)?; - let (started, completed) = timeout(DEFAULT_READ_TIMEOUT, async { - let mut started = None; - let mut completed = None; - while started.is_none() || completed.is_none() { - let JSONRPCMessage::Notification(notification) = mcp.read_next_message().await? else { - continue; - }; - match notification.method.as_str() { - "item/started" => { - let payload: ItemStartedNotification = - serde_json::from_value(notification.params.expect("item/started params"))?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - started = Some(payload); - } - } - "item/completed" => { - let payload: ItemCompletedNotification = serde_json::from_value( - notification.params.expect("item/completed params"), - )?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - completed = Some(payload); - } - } - _ => {} - } - } - Ok::<_, anyhow::Error>(( - started.expect("sleep started"), - completed.expect("sleep completed"), - )) - }) - .await??; + // Read once for the initial reminder, then once to establish the sleep deadline. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + let started = wait_for_sleep_started(&mut mcp, CALL_ID).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + + // The first poll remains below the deadline, so the provider must request time again. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 1).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; + + let completed = wait_for_sleep_completed(&mut mcp, CALL_ID).await?; + + // The next inference boundary reads the same external clock after the sleep completes. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -150,6 +136,64 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { Ok(()) } +async fn wait_for_sleep_started( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/started"), + ) + .await??; + let started: ItemStartedNotification = + serde_json::from_value(notification.params.expect("item/started params"))?; + if matches!(&started.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(started); + } + } +} + +async fn wait_for_sleep_completed( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/completed"), + ) + .await??; + let completed: ItemCompletedNotification = + serde_json::from_value(notification.params.expect("item/completed params"))?; + if matches!(&completed.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(completed); + } + } +} + +async fn respond_to_current_time_read( + mcp: &mut TestAppServer, + thread_id: &str, + current_time_at: i64, +) -> Result<()> { + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::CurrentTimeRead { request_id, params } = request else { + panic!("expected CurrentTimeRead request, got: {request:?}"); + }; + assert_eq!(params.thread_id, thread_id); + mcp.send_response( + request_id, + serde_json::to_value(CurrentTimeReadResponse { current_time_at })?, + ) + .await?; + Ok(()) +} + fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { std::fs::write( codex_home.join("config.toml"), @@ -170,6 +214,7 @@ stream_max_retries = 0 [features.current_time_reminder] enabled = true sleep_tool = true +clock_source = "external" "# ), ) From 703793c22efc1e47de51703150f5cf16976ff456 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 25 Jun 2026 13:54:48 -0700 Subject: [PATCH 18/33] feat(core, mcp): cache codex_apps tools in memory (#29003) ## Description This makes Codex Apps tool reads use a shared in-memory snapshot instead of rereading the disk cache every time `list_all_tools()` runs. Disk still seeds the cache on startup and gets updated after successful fetches, but it is no longer the live read path. The core change is that `McpManager` now owns a process-scoped `CodexAppsToolsCache`. Codex threads in the same app-server process now share this Codex Apps in-memory tools snapshot. The snapshot is keyed by the Codex home plus the Codex Apps identity: the active Codex auth user/workspace and the effective Codex Apps MCP source config. There's already code to hard-refresh the cache, so we respect it in this PR. ## Local benchmark I ran a local steady-state microbenchmark of the exact repeated Codex Apps cached-tools read this PR removes, using the same real local cache payload in both trees: `3,678,138` bytes and `381` tools. The cache file was already warm in the OS page cache, so this measures same-process reread/deserialization work rather than cold-disk latency or full turn latency. Each run is 25 iterations (mimicking a turn that makes 25 inference calls). | Version | Run 1 | Run 2 | Avg | |---|---:|---:|---:| | `origin/main` disk read + JSON deserialize + `filter_tools` | `50.755 ms` | `52.894 ms` | `51.825 ms` | | This branch in-memory `current_tools` + `filter_tools` | `0.740 ms` | `0.778 ms` | `0.759 ms` | That removes about `51 ms` from each repeated Codex Apps cached-tools read on this machine, roughly `68x` faster for that subpath. It is useful evidence for the hot path this PR changes, but not a claim that every production turn gets `51 ms` faster; end-to-end impact also depends on the rest of `list_all_tools()` and tool-payload construction. This is on my M2 Max macbook, so with a slower disk this would be much worse (and indeed we did see this really blew up turn runtime with a slow disk). --- .../src/request_processors/mcp_processor.rs | 23 +- codex-rs/codex-mcp/src/codex_apps.rs | 212 +------- codex-rs/codex-mcp/src/codex_apps_cache.rs | 378 ++++++++++++++ .../codex-mcp/src/codex_apps_cache_tests.rs | 464 ++++++++++++++++++ codex-rs/codex-mcp/src/connection_manager.rs | 90 ++-- .../codex-mcp/src/connection_manager_tests.rs | 425 +++++----------- codex-rs/codex-mcp/src/lib.rs | 6 +- codex-rs/codex-mcp/src/mcp/mod.rs | 7 +- codex-rs/codex-mcp/src/rmcp_client.rs | 96 ++-- codex-rs/core/src/connectors.rs | 1 + codex-rs/core/src/mcp.rs | 13 +- codex-rs/core/src/mcp_tool_call.rs | 12 +- codex-rs/core/src/mcp_tool_call_tests.rs | 1 + codex-rs/core/src/session/mcp.rs | 1 + codex-rs/core/src/session/session.rs | 1 + 15 files changed, 1109 insertions(+), 621 deletions(-) create mode 100644 codex-rs/codex-mcp/src/codex_apps_cache.rs create mode 100644 codex-rs/codex-mcp/src/codex_apps_cache_tests.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 18b49ede3b..a3d2191782 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -249,15 +249,12 @@ impl McpRequestProcessor { } None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; + let mcp_manager = self.thread_manager.mcp_manager(); let mcp_config = match thread { Some(thread) => thread.runtime_mcp_config(&config).await, - None => { - self.thread_manager - .mcp_manager() - .runtime_config(&config) - .await - } + None => mcp_manager.runtime_config(&config).await, }; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd @@ -274,6 +271,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; }); @@ -287,6 +285,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -294,6 +293,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; outgoing.send_result(request_id, result).await; @@ -305,6 +305,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, @@ -316,6 +317,7 @@ impl McpRequestProcessor { auth.as_ref(), request_id, runtime_context, + codex_apps_tools_cache, detail, ) .await; @@ -406,11 +408,9 @@ impl McpRequestProcessor { } let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let mcp_config = self - .thread_manager - .mcp_manager() - .runtime_config(&config) - .await; + let mcp_manager = self.thread_manager.mcp_manager(); + let mcp_config = mcp_manager.runtime_config(&config).await; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected @@ -425,6 +425,7 @@ impl McpRequestProcessor { &mcp_config, auth.as_ref(), runtime_context, + codex_apps_tools_cache, &server, &uri, ) diff --git a/codex-rs/codex-mcp/src/codex_apps.rs b/codex-rs/codex-mcp/src/codex_apps.rs index c1b5fd5fa5..a33f735366 100644 --- a/codex-rs/codex-mcp/src/codex_apps.rs +++ b/codex-rs/codex-mcp/src/codex_apps.rs @@ -1,70 +1,9 @@ //! Codex Apps support for the host-owned apps MCP server. //! -//! This module owns the pieces that are unique to ChatGPT-hosted app -//! connectors: cache scoping by authenticated user, disk cache reads/writes, -//! connector allow-list filtering, and the normalization that turns app +//! This module owns the normalization that turns ChatGPT-hosted app //! connector/tool metadata into model-visible MCP callable names. -use std::path::PathBuf; -use std::time::Instant; - -use crate::runtime::emit_duration; -use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; -use crate::tools::ToolInfo; -use anyhow::Context; -use codex_login::CodexAuth; -use codex_protocol::mcp::McpServerInfo; use codex_utils_plugins::mcp_connector::sanitize_name; -use serde::Deserialize; -use serde::Serialize; -use sha1::Digest; -use sha1::Sha1; -use tracing::instrument; - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CodexAppsToolsCacheKey { - pub(crate) account_id: Option, - pub(crate) chatgpt_user_id: Option, - pub(crate) is_workspace_account: bool, -} - -pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { - CodexAppsToolsCacheKey { - account_id: auth.and_then(CodexAuth::get_account_id), - chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), - is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), - } -} - -#[derive(Clone)] -pub(crate) struct CodexAppsToolsCacheContext { - pub(crate) codex_home: PathBuf, - pub(crate) user_key: CodexAppsToolsCacheKey, -} - -impl CodexAppsToolsCacheContext { - pub(crate) fn tools_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) - } - - pub(crate) fn server_info_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) - } - - fn cache_path_in(&self, cache_dir: &str) -> PathBuf { - let user_key_json = serde_json::to_string(&self.user_key).unwrap_or_default(); - let user_key_hash = sha1_hex(&user_key_json); - self.codex_home - .join(cache_dir) - .join(format!("{user_key_hash}.json")) - } -} - -pub(crate) enum CachedCodexAppsToolsLoad { - Hit(Vec), - Missing, - Invalid, -} pub(crate) fn normalize_codex_apps_tool_title(connector_name: Option<&str>, value: &str) -> String { let Some(connector_name) = connector_name @@ -124,152 +63,3 @@ pub(crate) fn normalize_codex_apps_callable_namespace( server_name.to_string() } } - -pub(crate) fn write_codex_apps_tools_cache( - cache_context: Option<&CodexAppsToolsCacheContext>, - server_info: &McpServerInfo, - tools: &[ToolInfo], -) { - if let Some(cache_context) = cache_context { - let cache_write_start = Instant::now(); - write_cached_codex_apps_tools(cache_context, tools); - if let Err(err) = write_cached_codex_apps_server_info(cache_context, server_info) { - tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); - } - emit_duration( - MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, - cache_write_start.elapsed(), - &[], - ); - } -} - -pub(crate) fn load_startup_cached_codex_apps_tools_snapshot( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option> { - let cache_context = cache_context?; - - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -pub(crate) fn load_startup_cached_codex_apps_server_info( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option { - load_cached_codex_apps_server_info(cache_context?) -} - -#[cfg(test)] -pub(crate) fn read_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> Option> { - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> CachedCodexAppsToolsLoad { - let cache_path = cache_context.tools_cache_path(); - let bytes = match std::fs::read(cache_path) { - Ok(bytes) => bytes, - Err(err) if err.kind() == std::io::ErrorKind::NotFound => { - return CachedCodexAppsToolsLoad::Missing; - } - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - let cache: CodexAppsToolsDiskCache = match serde_json::from_slice(&bytes) { - Ok(cache) => cache, - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - if cache.schema_version != CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION { - return CachedCodexAppsToolsLoad::Invalid; - } - CachedCodexAppsToolsLoad::Hit(cache.tools) -} - -pub(crate) fn write_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, - tools: &[ToolInfo], -) { - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() - && std::fs::create_dir_all(parent).is_err() - { - return; - } - let Ok(bytes) = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { - schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - tools: tools.to_vec(), - }) else { - return; - }; - let _ = std::fs::write(cache_path, bytes); -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, -) -> Option { - let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; - let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; - (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) - .then_some(cache.server_info) -} - -fn write_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, - server_info: &McpServerInfo, -) -> anyhow::Result<()> { - let cache_path = cache_context.server_info_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).with_context(|| { - format!( - "failed to create Codex Apps server info cache directory `{}`", - parent.display() - ) - })?; - } - let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { - schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, - server_info: server_info.clone(), - }) - .context("failed to serialize Codex Apps server info cache")?; - std::fs::write(&cache_path, bytes).with_context(|| { - format!( - "failed to write Codex Apps server info cache `{}`", - cache_path.display() - ) - })?; - Ok(()) -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsToolsDiskCache { - schema_version: u8, - tools: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsServerInfoDiskCache { - schema_version: u8, - server_info: McpServerInfo, -} - -const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; -pub(crate) const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; - -const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; -const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; - -fn sha1_hex(s: &str) -> String { - let mut hasher = Sha1::new(); - hasher.update(s.as_bytes()); - let sha1 = hasher.finalize(); - format!("{sha1:x}") -} diff --git a/codex-rs/codex-mcp/src/codex_apps_cache.rs b/codex-rs/codex-mcp/src/codex_apps_cache.rs new file mode 100644 index 0000000000..ea39444628 --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache.rs @@ -0,0 +1,378 @@ +//! Shared raw tool cache for the host-owned Codex Apps MCP server. +//! +//! Cache entries are process-local live state scoped by the active Codex auth +//! key. Disk is best-effort cold-start persistence; entries do not reread disk +//! after creation. + +use std::collections::HashMap; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Context; +use arc_swap::ArcSwapOption; +use codex_login::CodexAuth; +use codex_protocol::mcp::McpServerInfo; +use serde::Deserialize; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; +use tracing::instrument; + +use crate::runtime::emit_duration; +use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; +use crate::tools::ToolInfo; + +const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms"; + +/// The CodexAuth bits that identify a Codex Apps catalog. +/// +/// Debug bearer-token overrides bypass the shared cache, so shared entries only +/// need the CodexAuth-backed identity. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct CodexAppsToolsCacheKey { + pub(crate) account_id: Option, + pub(crate) chatgpt_user_id: Option, + pub(crate) is_workspace_account: bool, +} + +/// Builds the CodexAuth-backed Codex Apps cache key. +pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { + account_id: auth.and_then(CodexAuth::get_account_id), + chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), + is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), + } +} + +/// Process-scoped registry for shared Codex Apps raw tool snapshots. +/// +/// Two clients share an entry only when they would read the same Codex Apps +/// catalog. New entries may seed from disk; live entries read from memory only. +#[derive(Clone, Default)] +pub struct CodexAppsToolsCache { + entries: Arc>>>, +} + +/// Handle to one shared Codex Apps tools cache entry. +/// +/// The connection manager creates this from the auth key, then tool +/// reads and refreshes for that managed client use the same entry. +#[derive(Clone)] +pub(crate) struct CodexAppsToolsCacheContext { + entry: Arc, +} + +impl CodexAppsToolsCacheContext { + pub(crate) fn tools_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) + } + + pub(crate) fn server_info_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) + } + + pub(crate) fn current_tools(&self) -> Option> { + self.entry + .current_tools + .load_full() + .map(|tools| tools.as_ref().clone()) + } + + pub(crate) fn has_current_tools(&self) -> bool { + self.entry.current_tools.load_full().is_some() + } + + pub(crate) fn begin_fetch( + &self, + source: CodexAppsToolsFetchSource, + ) -> CodexAppsToolsFetchTicket { + CodexAppsToolsFetchTicket { + generation: self + .entry + .next_fetch_generation + .fetch_add(1, Ordering::Relaxed) + + 1, + source, + } + } + + pub(crate) fn publish_if_newest_accepted( + &self, + ticket: CodexAppsToolsFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + ) -> Vec { + let publish_start = Instant::now(); + let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation); + if ticket.generation <= *last_accepted_generation { + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "stale")], + ); + return self.current_tools().unwrap_or(tools); + } + + *last_accepted_generation = ticket.generation; + self.entry + .current_tools + .store(Some(Arc::new(tools.clone()))); + persist_codex_apps_cache(self, server_info, &tools); + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "published")], + ); + tools + } + + #[cfg(test)] + pub(crate) fn store_current_tools_for_test(&self, tools: Vec) { + self.entry.current_tools.store(Some(Arc::new(tools))); + } +} + +impl CodexAppsToolsCache { + pub(crate) fn context( + &self, + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, + ) -> CodexAppsToolsCacheContext { + let identity = CodexAppsToolsCacheIdentity { + codex_home, + auth_key, + }; + let mut entries = lock_unpoisoned(&self.entries); + let entry = entries + .entry(identity.clone()) + .or_insert_with(|| Arc::new(CodexAppsToolsCacheEntry::new(identity))) + .clone(); + CodexAppsToolsCacheContext { entry } + } +} + +#[derive(Debug, Clone, Copy)] +pub(crate) enum CodexAppsToolsFetchSource { + Startup, + HardRefresh, +} + +impl CodexAppsToolsFetchSource { + fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::HardRefresh => "hard_refresh", + } + } +} + +pub(crate) struct CodexAppsToolsFetchTicket { + generation: u64, + source: CodexAppsToolsFetchSource, +} + +struct CodexAppsToolsCacheEntry { + identity: CodexAppsToolsCacheIdentity, + current_tools: ArcSwapOption>, + next_fetch_generation: AtomicU64, + last_accepted_generation: Mutex, +} + +impl CodexAppsToolsCacheEntry { + fn new(identity: CodexAppsToolsCacheIdentity) -> Self { + let current_tools = load_cached_codex_apps_tools_for_identity(&identity).map(Arc::new); + Self { + identity, + current_tools: ArcSwapOption::from(current_tools), + next_fetch_generation: AtomicU64::new(0), + last_accepted_generation: Mutex::new(0), + } + } +} + +/// Everything that decides whether two Codex Apps clients can share tools. +/// +/// The auth key says whose catalog we are reading. `codex_home` keeps the +/// persisted cache under the right home directory. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CodexAppsToolsCacheIdentity { + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, +} + +impl CodexAppsToolsCacheIdentity { + fn cache_path_in(&self, cache_dir: &str) -> PathBuf { + // `codex_home` is already the parent directory. Keep it out of the + // filename hash so non-UTF-8 Unix paths cannot collapse distinct auth + // keys onto the same disk cache file. + let identity_json = serde_json::to_string(&self.auth_key).unwrap_or_default(); + let identity_hash = sha1_hex(&identity_json); + self.codex_home + .join(cache_dir) + .join(format!("{identity_hash}.json")) + } +} + +#[cfg(test)] +fn write_cached_codex_apps_tools_for_test( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + cache_context + .entry + .current_tools + .store(Some(Arc::new(tools.to_vec()))); + persist_codex_apps_cache(cache_context, server_info, tools); +} + +pub(crate) fn load_startup_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + load_cached_codex_apps_server_info(cache_context) +} + +#[cfg(test)] +fn read_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, +) -> Option> { + load_cached_codex_apps_tools_for_identity(&cache_context.entry.identity) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_tools_for_identity( + identity: &CodexAppsToolsCacheIdentity, +) -> Option> { + let cache_path = identity.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR); + let bytes = std::fs::read(cache_path).ok()?; + let cache: CodexAppsToolsDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION).then_some(cache.tools) +} + +fn write_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, + tools: &[ToolInfo], +) -> anyhow::Result<()> { + let cache_path = cache_context.tools_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { + schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + tools: tools.to_vec(), + }) + .context("failed to serialize Codex Apps tools cache")?; + write_codex_apps_cache_file(&cache_path, "tools", bytes) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; + let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) + .then_some(cache.server_info) +} + +fn write_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, +) -> anyhow::Result<()> { + let cache_path = cache_context.server_info_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { + schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, + server_info: server_info.clone(), + }) + .context("failed to serialize Codex Apps server info cache")?; + write_codex_apps_cache_file(&cache_path, "server info", bytes) +} + +fn write_codex_apps_cache_file( + cache_path: &Path, + cache_name: &str, + bytes: Vec, +) -> anyhow::Result<()> { + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create Codex Apps {cache_name} cache directory `{}`", + parent.display() + ) + })?; + } + std::fs::write(cache_path, bytes).with_context(|| { + format!( + "failed to write Codex Apps {cache_name} cache `{}`", + cache_path.display() + ) + })?; + Ok(()) +} + +fn persist_codex_apps_cache( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + let cache_write_start = Instant::now(); + let tools_result = write_cached_codex_apps_tools(cache_context, tools); + if let Err(err) = &tools_result { + tracing::warn!("failed to write Codex Apps tools cache: {err:#}"); + } + let server_info_result = write_cached_codex_apps_server_info(cache_context, server_info); + if let Err(err) = &server_info_result { + tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); + } + let status = if tools_result.is_ok() && server_info_result.is_ok() { + "success" + } else { + "failure" + }; + emit_duration( + MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, + cache_write_start.elapsed(), + &[("status", status)], + ); +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsToolsDiskCache { + schema_version: u8, + tools: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsServerInfoDiskCache { + schema_version: u8, + server_info: McpServerInfo, +} + +const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; +const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; + +const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; +const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; + +fn sha1_hex(s: &str) -> String { + let mut hasher = Sha1::new(); + hasher.update(s.as_bytes()); + let sha1 = hasher.finalize(); + format!("{sha1:x}") +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +#[cfg(test)] +#[path = "codex_apps_cache_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs new file mode 100644 index 0000000000..46ad311a1b --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs @@ -0,0 +1,464 @@ +use super::*; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; +use crate::tools::ToolInfo; +use codex_protocol::ToolName; +use codex_protocol::mcp::McpServerInfo; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Tool; +use std::collections::HashSet; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::tempdir; + +fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: server_name.to_string(), + namespace_description: None, + tool: Tool::new( + tool_name.to_string(), + format!("Test tool: {tool_name}"), + Arc::new(JsonObject::default()), + ), + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } +} + +fn create_test_tool_with_connector( + server_name: &str, + tool_name: &str, + connector_id: &str, + connector_name: Option<&str>, +) -> ToolInfo { + let mut tool = create_test_tool(server_name, tool_name); + tool.connector_id = Some(connector_id.to_string()); + tool.connector_name = connector_name.map(ToOwned::to_owned); + tool +} + +fn create_codex_apps_tools_cache_context( + codex_home: PathBuf, + account_id: Option<&str>, + chatgpt_user_id: Option<&str>, +) -> CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( + codex_home, + CodexAppsToolsCacheKey { + account_id: account_id.map(ToOwned::to_owned), + chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), + is_workspace_account: false, + }, + ) +} + +fn create_test_server_info(title: &str) -> McpServerInfo { + McpServerInfo { + name: "codex-apps".to_string(), + title: Some(title.to_string()), + version: "1.0.0".to_string(), + description: None, + icons: None, + website_url: None, + } +} + +fn model_tool_names(tools: &[ToolInfo]) -> HashSet { + tools + .iter() + .map(ToolInfo::canonical_tool_name) + .collect::>() +} + +#[test] +fn codex_apps_tools_cache_is_overwritten_by_last_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_1).expect("write first cache"); + let cached_gateway_1 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); + assert_eq!(cached_gateway_1[0].callable_name, "one"); + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_2).expect("write second cache"); + let cached_gateway_2 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); + assert_eq!(cached_gateway_2[0].callable_name, "two"); +} + +#[test] +fn codex_apps_tools_cache_is_scoped_per_user() { + let codex_home = tempdir().expect("tempdir"); + let cache_context_user_1 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_context_user_2 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-two"), + Some("user-two"), + ); + let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1) + .expect("write user one cache"); + write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2) + .expect("write user two cache"); + + let read_user_1 = + read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); + let read_user_2 = + read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); + + assert_eq!(read_user_1[0].callable_name, "one"); + assert_eq!(read_user_2[0].callable_name, "two"); + assert_ne!( + cache_context_user_1.tools_cache_path(), + cache_context_user_2.tools_cache_path(), + "each user should get an isolated cache file" + ); +} + +#[test] +fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools = vec![ + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "formerly_blocked_tool", + "connector_2b0a9009c9c64bf9933a3dae3f2b1254", + Some("Formerly Blocked"), + ), + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_tool", + "calendar", + Some("Calendar"), + ), + ]; + + write_cached_codex_apps_tools(&cache_context, &tools).expect("write cache"); + let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); + + assert_eq!( + cached + .iter() + .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) + .collect::>(), + vec![ + ( + "formerly_blocked_tool", + Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") + ), + ("calendar_tool", Some("calendar")), + ] + ); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + std::fs::write(cache_path, b"{not json").expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_from_disk_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )]; + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test(&writer_cache_context, &server_info, &cached_tools); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("expected startup snapshot to load from cache"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, Some(server_info)); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = writer_cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("legacy startup snapshot should remain available"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, None); +} + +#[test] +fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test( + &cache_context, + &server_info, + &[create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )], + ); + + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write legacy tools cache"); + let startup_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + assert_eq!( + load_startup_cached_codex_apps_server_info(&startup_cache_context), + Some(server_info) + ); + assert!(startup_cache_context.current_tools().is_none()); +} + +#[test] +fn codex_apps_tools_cache_context_does_not_reread_disk_after_creation() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "cached")]; + write_cached_codex_apps_tools(&writer_cache_context, &cached_tools).expect("write cache"); + let reader_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let updated_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")]; + write_cached_codex_apps_tools(&writer_cache_context, &updated_tools).expect("rewrite cache"); + + assert_eq!( + reader_cache_context + .current_tools() + .expect("in-memory tools")[0] + .callable_name, + "cached" + ); + assert_eq!( + read_cached_codex_apps_tools(&writer_cache_context).expect("disk tools")[0].callable_name, + "updated" + ); +} + +#[test] +fn codex_apps_tools_cache_publishes_newest_shared_snapshot() { + let codex_home = tempdir().expect("tempdir"); + let cache = CodexAppsToolsCache::default(); + let cache_context_1 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let cache_context_2 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let older_ticket = cache_context_1.begin_fetch(CodexAppsToolsFetchSource::Startup); + let newer_ticket = cache_context_2.begin_fetch(CodexAppsToolsFetchSource::HardRefresh); + let server_info = create_test_server_info("Codex Apps"); + let newer_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")]; + let older_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")]; + + let published_tools = + cache_context_2.publish_if_newest_accepted(newer_ticket, &server_info, newer_tools); + assert_eq!( + model_tool_names(&published_tools), + model_tool_names( + &cache_context_1 + .current_tools() + .expect("new snapshot should publish") + ) + ); + let current_tools = + cache_context_1.publish_if_newest_accepted(older_ticket, &server_info, older_tools); + + assert_eq!(current_tools[0].callable_name, "newer"); + assert_eq!( + cache_context_2.current_tools().expect("shared snapshot")[0].callable_name, + "newer" + ); + assert_eq!( + read_cached_codex_apps_tools(&cache_context_1).expect("persisted snapshot")[0] + .callable_name, + "newer" + ); +} + +#[test] +fn codex_apps_tools_cache_keeps_live_publish_when_disk_persistence_fails() { + let codex_home = tempdir().expect("tempdir"); + let codex_home_file = codex_home.path().join("not-a-directory"); + std::fs::write(&codex_home_file, b"occupied").expect("create codex home file"); + let cache_context = CodexAppsToolsCache::default().context( + codex_home_file, + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "live")]; + let published_tools = cache_context.publish_if_newest_accepted( + cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools.clone(), + ); + + assert_eq!(model_tool_names(&published_tools), model_tool_names(&tools)); + assert_eq!( + model_tool_names(&cache_context.current_tools().expect("live snapshot")), + model_tool_names(&tools) + ); +} + +#[cfg(unix)] +#[test] +fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { + let codex_home = PathBuf::from(std::ffi::OsString::from_vec( + b"/tmp/codex-home-\xff".to_vec(), + )); + let cache = CodexAppsToolsCache::default(); + let user_one_context = cache.context( + codex_home.clone(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let user_two_context = cache.context( + codex_home, + CodexAppsToolsCacheKey { + account_id: Some("account-two".to_string()), + chatgpt_user_id: Some("user-two".to_string()), + is_workspace_account: false, + }, + ); + let cache_paths = [ + user_one_context.tools_cache_path(), + user_two_context.tools_cache_path(), + ]; + + assert_eq!( + cache_paths.iter().collect::>().len(), + cache_paths.len() + ); +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 55c72b0467..0e81549556 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,7 +7,6 @@ //! `codex-core`. use std::collections::HashMap; -use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::Ordering; @@ -15,9 +14,9 @@ use std::time::Duration; use std::time::Instant; use crate::McpAuthStatusEntry; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::CodexAppsToolsCacheKey; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheKey; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -131,6 +130,7 @@ impl McpConnectionManager { initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, + codex_apps_tools_cache: CodexAppsToolsCache, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, @@ -174,13 +174,32 @@ impl McpConnectionManager { }, ) .await; - let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { - codex_apps_tools_cache_context(&codex_home, &codex_apps_tools_cache_key) - } else { - regular_mcp_tools_cache_context() - }; + let configured_config = server.configured_config(); + // For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug + // override: it supplies runtime auth but bypasses the shared tools + // cache. + let uses_env_bearer_token = + configured_config.is_some_and(|config| match &config.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }); + let shares_codex_apps_tools_cache = + should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); + let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { + codex_apps_tools_cache + .context(codex_home.clone(), codex_apps_tools_cache_key.clone()) + }); + // If Codex Apps has an env bearer token, that is its auth path. Do + // not also attach the ambient CodexAuth provider. let runtime_auth_provider = - chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()); + if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token { + None + } else { + chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()) + }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), server, @@ -445,13 +464,13 @@ impl McpConnectionManager { pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); for (server_name, managed_client) in &self.clients { - let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some(); + let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete .load(std::sync::atomic::Ordering::Acquire); trace!( server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete, "waiting for MCP server tools while building tool list" ); @@ -460,7 +479,7 @@ impl McpConnectionManager { .instrument(trace_span!( "list_tools_for_server", server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete )) .await @@ -483,9 +502,9 @@ impl McpConnectionManager { /// Force-refresh codex apps tools by bypassing the in-process cache. /// - /// On success, the refreshed tools replace the cache contents and the - /// latest filtered tools are returned directly to the caller. On - /// failure, the existing cache remains unchanged. + /// On success, the refreshed tools replace shared cache contents when the + /// cache is enabled and the latest filtered tools are returned directly to + /// the caller. On failure, existing shared cache contents remain unchanged. pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let managed_client = self .clients @@ -497,6 +516,10 @@ impl McpConnectionManager { let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = managed_client + .codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh)); let tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, @@ -514,11 +537,16 @@ impl McpConnectionManager { &[], ); - write_codex_apps_tools_cache( - managed_client.codex_apps_tools_cache_context.as_ref(), - &managed_client.server_info, - &tools, - ); + let tools = + match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context + .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), @@ -849,22 +877,6 @@ impl Drop for McpConnectionManager { } } -/// Creates the per-user tools cache context used only by the Codex Apps server. -fn codex_apps_tools_cache_context( - codex_home: &Path, - codex_apps_tools_cache_key: &CodexAppsToolsCacheKey, -) -> Option { - Some(CodexAppsToolsCacheContext { - codex_home: codex_home.to_path_buf(), - user_key: codex_apps_tools_cache_key.clone(), - }) -} - -/// Keeps regular MCP servers isolated from the Codex Apps tools cache. -fn regular_mcp_tools_cache_context() -> Option { - None -} - /// Makes ChatGPT authentication available to servers that explicitly opt in. /// The HTTP transport applies it only when no configured authorization resolves. fn chatgpt_auth_provider_for_server( @@ -880,6 +892,10 @@ fn chatgpt_auth_provider_for_server( chatgpt_auth_provider } +fn should_share_codex_apps_tools_cache(server_name: &str, uses_env_bearer_token: bool) -> bool { + server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token +} + async fn emit_update( submit_id: &str, tx_event: &Sender, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d66474c7bb..05caf0489c 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,11 +1,6 @@ use super::*; -use crate::codex_apps::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; -use crate::codex_apps::read_cached_codex_apps_tools; -use crate::codex_apps::write_cached_codex_apps_tools; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; use crate::declared_openai_file_input_param_names; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::elicitation_is_rejected_by_policy; @@ -31,7 +26,10 @@ use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpAuthStatus; +use codex_rmcp_client::InProcessTransportFactory; +use codex_rmcp_client::RmcpClient; use futures::FutureExt; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; @@ -41,8 +39,10 @@ use rmcp::model::Meta; use rmcp::model::NumberOrString; use rmcp::model::Tool; use std::collections::HashSet; +use std::io; use std::sync::Arc; use tempfile::tempdir; +use tokio::io::DuplexStream; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -63,31 +63,19 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { } } -fn create_test_tool_with_connector( - server_name: &str, - tool_name: &str, - connector_id: &str, - connector_name: Option<&str>, -) -> ToolInfo { - let mut tool = create_test_tool(server_name, tool_name); - tool.connector_id = Some(connector_id.to_string()); - tool.connector_name = connector_name.map(ToOwned::to_owned); - tool -} - fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, chatgpt_user_id: Option<&str>, ) -> CodexAppsToolsCacheContext { - CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( codex_home, - user_key: CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { account_id: account_id.map(ToOwned::to_owned), chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), is_workspace_account: false, }, - } + ) } fn create_test_server_info(title: &str) -> McpServerInfo { @@ -101,6 +89,51 @@ fn create_test_server_info(title: &str) -> McpServerInfo { } } +struct TestInProcessTransportFactory; + +impl InProcessTransportFactory for TestInProcessTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + async { + let (client_stream, _server_stream) = tokio::io::duplex(1); + Ok(client_stream) + } + .boxed() + } +} + +async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { + let tool_filter = ToolFilter::default(); + let managed_client = ManagedClient { + client: Arc::new( + RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) + .await + .expect("create in-process RMCP client"), + ), + server_info: create_test_server_info("Ready"), + tools, + tool_filter: tool_filter.clone(), + tool_timeout: None, + server_instructions: None, + server_supports_sandbox_state_meta_capability: false, + codex_apps_tools_cache_context: None, + }; + + AsyncManagedClient { + client: futures::future::ready::>(Ok( + managed_client, + )) + .boxed() + .shared(), + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + } +} + fn model_tool_names(tools: &[ToolInfo]) -> HashSet { tools .iter() @@ -554,235 +587,25 @@ fn filter_tools_applies_per_server_filters() { } #[test] -fn codex_apps_tools_cache_is_overwritten_by_last_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_1); - let cached_gateway_1 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); - assert_eq!(cached_gateway_1[0].callable_name, "one"); - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_2); - let cached_gateway_2 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); - assert_eq!(cached_gateway_2[0].callable_name, "two"); -} - -#[test] -fn codex_apps_tools_cache_is_scoped_per_user() { - let codex_home = tempdir().expect("tempdir"); - let cache_context_user_1 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_context_user_2 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-two"), - Some("user-two"), - ); - let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1); - write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2); - - let read_user_1 = - read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); - let read_user_2 = - read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); - - assert_eq!(read_user_1[0].callable_name, "one"); - assert_eq!(read_user_2[0].callable_name, "two"); - assert_ne!( - cache_context_user_1.tools_cache_path(), - cache_context_user_2.tools_cache_path(), - "each user should get an isolated cache file" - ); -} - -#[test] -fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools = vec![ - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "formerly_blocked_tool", - "connector_2b0a9009c9c64bf9933a3dae3f2b1254", - Some("Formerly Blocked"), - ), - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_tool", - "calendar", - Some("Calendar"), - ), - ]; - - write_cached_codex_apps_tools(&cache_context, &tools); - let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); - - assert_eq!( - cached - .iter() - .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) - .collect::>(), - vec![ - ( - "formerly_blocked_tool", - Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") - ), - ("calendar_tool", Some("calendar")), - ] - ); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - std::fs::write(cache_path, b"{not json").expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_from_disk_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cached_tools = vec![create_test_tool( +fn codex_apps_env_bearer_token_bypasses_shared_tools_cache() { + assert!(!should_share_codex_apps_tools_cache( CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )]; - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache(Some(&cache_context), &server_info, &cached_tools); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("expected startup snapshot to load from cache"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, Some(server_info)); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("legacy startup snapshot should remain available"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, None); -} - -#[test] -fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache( - Some(&cache_context), - &server_info, - &[create_test_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )], - ); - - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write legacy tools cache"); - - assert_eq!( - load_startup_cached_codex_apps_server_info(Some(&cache_context)), - Some(server_info) - ); - assert!(load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)).is_none()); + /*uses_env_bearer_token*/ true, + )); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -798,8 +621,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -813,7 +637,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); } @@ -836,8 +660,9 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -858,10 +683,8 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { #[tokio::test] async fn list_all_tools_accepts_canonical_namespaced_tool_names() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -869,18 +692,7 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { &permission_profile, /*prefix_mcp_tool_names*/ false, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -902,10 +714,8 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -913,18 +723,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -945,7 +744,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { } #[tokio::test] -async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() { +async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -961,8 +760,9 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -998,8 +798,9 @@ async fn shutdown_cancels_pending_tool_listing() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, @@ -1043,8 +844,9 @@ async fn shutdown_continues_after_caller_is_aborted() { AsyncManagedClient { client: blocking_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1071,7 +873,14 @@ async fn shutdown_continues_after_caller_is_aborted() { } #[tokio::test] -async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() { +async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(Vec::new()); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -1087,8 +896,9 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1097,16 +907,22 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() let timeout_result = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; - let tools = timeout_result.expect("cache-hit startup snapshot should not block"); + let tools = timeout_result.expect("shared empty cache should not block"); assert!(tools.is_empty()); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let server_info = create_test_server_info("Codex Apps"); let failed_client = futures::future::ready::>(Err( StartupOutcomeError::Failed { @@ -1128,8 +944,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails AsyncManagedClient { client: failed_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1143,7 +960,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); assert_eq!( @@ -1156,12 +973,10 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails } #[tokio::test] -async fn list_all_tools_adds_server_metadata_to_cached_tools() { +async fn list_all_tools_adds_server_metadata_to_tools() { let server_name = "docs"; - let startup_tools = vec![create_test_tool(server_name, "search")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool(server_name, "search")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -1182,18 +997,9 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() { tool_approval_modes: HashMap::new(), }, ); - manager.clients.insert( - server_name.to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager + .clients + .insert(server_name.to_string(), managed_client); let tools = manager.list_all_tools().await; assert_eq!(tools.len(), 1); @@ -1340,6 +1146,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { PathBuf::from("/tmp"), ), codex_home.path().to_path_buf(), + CodexAppsToolsCache::default(), CodexAppsToolsCacheKey { account_id: None, chatgpt_user_id: None, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 98efdc1b56..e611623d44 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -35,8 +35,9 @@ pub use auth_elicitation::auth_elicitation_id; pub use auth_elicitation::build_auth_elicitation; pub use auth_elicitation::build_auth_elicitation_plan; pub use auth_elicitation::connector_auth_failure_from_tool_result; -pub use codex_apps::CodexAppsToolsCacheKey; -pub use codex_apps::codex_apps_tools_cache_key; +pub use codex_apps_cache::CodexAppsToolsCache; +pub use codex_apps_cache::CodexAppsToolsCacheKey; +pub use codex_apps_cache::codex_apps_tools_cache_key; pub use mcp::codex_apps_mcp_server_config; pub use mcp::configured_mcp_servers; pub use mcp::effective_mcp_servers; @@ -75,6 +76,7 @@ pub use tools::declared_openai_file_input_param_names; pub(crate) mod auth_elicitation; mod catalog; pub(crate) mod codex_apps; +pub(crate) mod codex_apps_cache; pub(crate) mod connection_manager; pub(crate) mod elicitation; pub(crate) mod mcp; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 4550afa799..7532c6466b 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -44,7 +44,8 @@ use serde_json::Value; use tokio_util::sync::CancellationToken; use crate::ResolvedMcpCatalog; -use crate::codex_apps::codex_apps_tools_cache_key; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::codex_apps_tools_cache_key; use crate::connection_manager::McpConnectionManager; use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; @@ -301,6 +302,7 @@ pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, server: &str, uri: &str, ) -> anyhow::Result { @@ -329,6 +331,7 @@ pub async fn read_mcp_resource( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), @@ -361,6 +364,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( auth: Option<&CodexAuth>, submit_id: String, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -403,6 +407,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 1e8f5a331d..73ad33df7f 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -2,7 +2,7 @@ //! //! This module owns startup of individual RMCP clients: building the transport, //! initializing the server, listing raw tools, applying per-server tool filters, -//! and exposing cached startup snapshots while a client is still connecting. +//! and exposing cached Codex Apps tools while a client is still connecting. //! Higher-level aggregation and resource/tool APIs live in //! [`crate::connection_manager`]. @@ -17,15 +17,12 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; -use crate::codex_apps::CachedCodexAppsToolsLoad; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_cached_codex_apps_tools; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; use crate::codex_apps::normalize_codex_apps_tool_title; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; +use crate::codex_apps_cache::load_startup_cached_codex_apps_server_info; use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; @@ -105,9 +102,10 @@ pub(crate) struct ManagedClient { impl ManagedClient { fn listed_tools(&self) -> Vec { let total_start = Instant::now(); - if let Some(cache_context) = self.codex_apps_tools_cache_context.as_ref() - && let CachedCodexAppsToolsLoad::Hit(tools) = - load_cached_codex_apps_tools(cache_context) + if let Some(tools) = self + .codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, @@ -133,8 +131,9 @@ impl ManagedClient { pub(crate) struct AsyncManagedClient { pub(crate) client: Shared>>, pub(crate) is_codex_apps_mcp_server: bool, - pub(crate) cached_tool_info_snapshot: Option>, pub(crate) cached_server_info: Option, + pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) tool_plugin_provenance: Arc, pub(crate) cancel_token: CancellationToken, @@ -165,19 +164,15 @@ impl AsyncManagedClient { .configured_config() .map(ToolFilter::from_config) .unwrap_or_default(); - let (cached_tool_info_snapshot, cached_server_info) = if is_codex_apps_mcp_server { - ( - load_startup_cached_codex_apps_tools_snapshot( - codex_apps_tools_cache_context.as_ref(), - ), - load_startup_cached_codex_apps_server_info(codex_apps_tools_cache_context.as_ref()), - ) + let cached_server_info = if is_codex_apps_mcp_server { + codex_apps_tools_cache_context + .as_ref() + .and_then(load_startup_cached_codex_apps_server_info) } else { - (None, None) + None }; - let cached_tool_info_snapshot = - cached_tool_info_snapshot.map(|tools| filter_tools(tools, &tool_filter)); - let startup_tool_filter = tool_filter; + let startup_tool_filter = tool_filter.clone(); + let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone(); let startup_complete = Arc::new(AtomicBool::new(false)); let startup_complete_for_fut = Arc::clone(&startup_complete); let cancel_token_for_fut = cancel_token.clone(); @@ -214,7 +209,7 @@ impl AsyncManagedClient { tool_filter: startup_tool_filter, tx_event, elicitation_requests, - codex_apps_tools_cache_context, + codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -232,7 +227,10 @@ impl AsyncManagedClient { outcome }; let client = fut.in_current_span().boxed().shared(); - if cached_tool_info_snapshot.is_some() { + if codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + { let startup_task = client.clone(); tokio::spawn(async move { let _ = startup_task.await; @@ -242,8 +240,9 @@ impl AsyncManagedClient { Self { client, is_codex_apps_mcp_server, - cached_tool_info_snapshot, cached_server_info, + codex_apps_tools_cache_context, + tool_filter, startup_complete, tool_plugin_provenance, cancel_token, @@ -265,15 +264,29 @@ impl AsyncManagedClient { } } + pub(crate) fn has_cached_tools(&self) -> bool { + self.codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + } + + fn cached_tools(&self) -> Option> { + self.codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) + .map(|tools| filter_tools(tools, &self.tool_filter)) + } + pub(crate) async fn listed_tools(&self) -> Option> { // Keep cache payloads raw; plugin provenance is resolved per-session at read time. - let tools = if let Some(startup_tools) = self.cached_tool_info_snapshot_while_initializing() + let tools = if !self.startup_complete.load(Ordering::Acquire) + && let Some(startup_tools) = self.cached_tools() { Some(startup_tools) } else { match self.client().await { Ok(client) => Some(client.listed_tools()), - Err(_) => self.cached_tool_info_snapshot.clone(), + Err(_) => self.cached_tools(), } }?; Some(if self.is_codex_apps_mcp_server { @@ -282,13 +295,6 @@ impl AsyncManagedClient { prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) }) } - - fn cached_tool_info_snapshot_while_initializing(&self) -> Option> { - if !self.startup_complete.load(Ordering::Acquire) { - return self.cached_tool_info_snapshot.clone(); - } - None - } } #[derive(Debug, Clone, thiserror::Error)] @@ -570,6 +576,9 @@ async fn start_server_task( .is_some(); let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); let tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, @@ -585,21 +594,20 @@ async fn start_server_task( &[], ); let server_info = mcp_server_info_from_implementation(initialize_result.server_info); - let codex_apps_tools_cache_context = if is_codex_apps_mcp_server { - write_codex_apps_tools_cache( - codex_apps_tools_cache_context.as_ref(), - &server_info, - &tools, - ); + let tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { + (Some(cache_context), Some(fetch_ticket)) => { + cache_context.publish_if_newest_accepted(fetch_ticket, &server_info, tools) + } + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; + if is_codex_apps_mcp_server { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), &[("cache", "miss")], ); - codex_apps_tools_cache_context - } else { - None - }; + } let tools = filter_tools(tools, &tool_filter); let managed = ManagedClient { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 76b4a84830..97f24e043f 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -275,6 +275,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( // one exists, but do not reintroduce the old hidden-local fallback. runtime_context, config.codex_home.to_path_buf(), + mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 90e6b01e75..658c5471ee 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -10,6 +10,7 @@ use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::CodexAppsToolsCache; use codex_mcp::EffectiveMcpServer; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; @@ -38,14 +39,15 @@ enum OrderedMcpOverlay { pub struct McpManager { plugins_manager: Arc, extensions: Arc>, + codex_apps_tools_cache: CodexAppsToolsCache, } impl McpManager { pub fn new(plugins_manager: Arc) -> Self { - Self { + Self::new_with_extensions( plugins_manager, - extensions: codex_extension_api::empty_extension_registry(), - } + codex_extension_api::empty_extension_registry(), + ) } /// Creates a manager that resolves host-installed MCP contributions. @@ -56,9 +58,14 @@ impl McpManager { Self { plugins_manager, extensions, + codex_apps_tools_cache: CodexAppsToolsCache::default(), } } + pub fn codex_apps_tools_cache(&self) -> CodexAppsToolsCache { + self.codex_apps_tools_cache.clone() + } + /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index df2a375e35..a270af7851 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -1471,9 +1471,15 @@ pub(crate) async fn lookup_mcp_tool_metadata( { Some(connectors) => Some(connectors), None => { - connectors::list_accessible_connectors_from_mcp_tools(turn_context.config.as_ref()) - .await - .ok() + connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + turn_context.config.as_ref(), + /*force_refetch*/ false, + sess.services.turn_environments.environment_manager(), + Arc::clone(&sess.services.mcp_manager), + ) + .await + .ok() + .map(|status| status.connectors) } }; connectors.and_then(|connectors| { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 387987c9f6..41701d557c 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1400,6 +1400,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: }, ), turn_context.config.codex_home.to_path_buf(), + session.services.mcp_manager.codex_apps_tools_cache(), codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c060ef090b..249a6994b1 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -311,6 +311,7 @@ impl Session { turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + self.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 96dffed888..57a66b9986 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1193,6 +1193,7 @@ impl Session { session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + sess.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names(), client_elicitation_capability, From 891f1f4c8584a082fc4658cabd48f1a8b01354e0 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 25 Jun 2026 13:56:08 -0700 Subject: [PATCH 19/33] release: publish standalone zsh artifacts (#30114) ## Why The patched zsh artifacts rarely change, but `.github/workflows/rust-release-zsh.yml` currently runs as part of every Rust release. Rebuilding the same four binaries for each Codex version wastes release capacity and ties an independently versioned runtime dependency to the main release cadence. This establishes the producer side of a build-once flow. The existing Rust release workflow remains unchanged until the first standalone artifact release has been published and the checked-in DotSlash manifests can be updated with its URLs and checksums. ## What changed - Run the zsh release workflow for protected `codex-zsh-vX.Y.Z` tags instead of as a reusable workflow. - Validate the semantic release tag before starting the platform builds. - Publish the four zsh archives to a GitHub prerelease so the release never becomes the repository latest release. - Publish the generated `codex-zsh` DotSlash manifest alongside the archives. - Document how to publish the next artifact version after changing the pinned zsh commit or patch. ## Tag protection An active repository tag ruleset named `codex-zsh-v*.*.*` targets `refs/tags/codex-zsh-v*.*.*`. It restricts tag creation, updates, deletion, and non-fast-forward changes; requires linear history; and limits bypass to the configured repository role. This was verified with: ```shell gh api repos/openai/codex/rulesets/18140982 ``` The response reported `"enforcement":"active"`, the expected tag condition, and the `creation`, `update`, `deletion`, `non_fast_forward`, and `required_linear_history` rules. ## Rollout After this lands, publish the first `codex-zsh-vX.Y.Z` release. A follow-up can then update the checked-in DotSlash manifests and remove the zsh rebuild from `.github/workflows/rust-release.yml`. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/30114). * #30116 * __->__ #30114 --- .github/dotslash-zsh-config.json | 4 ++ .github/workflows/rust-release-zsh.yml | 81 +++++++++++++++++++++++++- codex-rs/shell-escalation/README.md | 5 ++ 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/.github/dotslash-zsh-config.json b/.github/dotslash-zsh-config.json index 37285f19e2..2c2b11eec5 100644 --- a/.github/dotslash-zsh-config.json +++ b/.github/dotslash-zsh-config.json @@ -5,21 +5,25 @@ "macos-aarch64": { "name": "codex-zsh-aarch64-apple-darwin.tar.gz", "format": "tar.gz", + "hash": "sha256", "path": "codex-zsh/bin/zsh" }, "macos-x86_64": { "name": "codex-zsh-x86_64-apple-darwin.tar.gz", "format": "tar.gz", + "hash": "sha256", "path": "codex-zsh/bin/zsh" }, "linux-x86_64": { "name": "codex-zsh-x86_64-unknown-linux-musl.tar.gz", "format": "tar.gz", + "hash": "sha256", "path": "codex-zsh/bin/zsh" }, "linux-aarch64": { "name": "codex-zsh-aarch64-unknown-linux-musl.tar.gz", "format": "tar.gz", + "hash": "sha256", "path": "codex-zsh/bin/zsh" } } diff --git a/.github/workflows/rust-release-zsh.yml b/.github/workflows/rust-release-zsh.yml index b55d2e714b..e814ec32e1 100644 --- a/.github/workflows/rust-release-zsh.yml +++ b/.github/workflows/rust-release-zsh.yml @@ -1,15 +1,56 @@ name: rust-release-zsh on: - workflow_call: + push: + tags: + - "codex-zsh-v*.*.*" env: ZSH_COMMIT: 77045ef899e53b9598bebc5a41db93a548a40ca6 ZSH_PATCH: codex-rs/shell-escalation/patches/zsh-exec-wrapper.patch +concurrency: + group: ${{ github.workflow }}::${{ github.ref_name }} + cancel-in-progress: false + jobs: + metadata: + runs-on: ubuntu-latest + outputs: + release_tag: ${{ steps.release_tag.outputs.release_tag }} + + steps: + - name: Validate release tag + id: release_tag + env: + RELEASE_TAG: ${{ github.ref_name }} + shell: bash + run: | + set -euo pipefail + + if [[ ! "${RELEASE_TAG}" =~ ^codex-zsh-v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then + echo "Tag ${RELEASE_TAG} does not match codex-zsh-vX.Y.Z." >&2 + exit 1 + fi + + echo "release_tag=${RELEASE_TAG}" >> "${GITHUB_OUTPUT}" + + - name: Ensure release does not exist + env: + GH_TOKEN: ${{ github.token }} + RELEASE_TAG: ${{ steps.release_tag.outputs.release_tag }} + shell: bash + run: | + set -euo pipefail + + if gh release view "${RELEASE_TAG}" --repo "${GITHUB_REPOSITORY}" > /dev/null 2>&1; then + echo "Release ${RELEASE_TAG} already exists; publish changed artifacts under a new tag." >&2 + exit 1 + fi + linux: name: Build zsh (Linux) - ${{ matrix.variant }} - ${{ matrix.target }} + needs: metadata runs-on: ${{ matrix.runner }} timeout-minutes: 30 container: @@ -62,6 +103,7 @@ jobs: darwin: name: Build zsh (macOS) - ${{ matrix.variant }} - ${{ matrix.target }} + needs: metadata runs-on: ${{ matrix.runner }} timeout-minutes: 30 @@ -101,3 +143,40 @@ jobs: with: name: codex-zsh-${{ matrix.target }} path: dist/zsh/${{ matrix.target }}/* + + publish-release: + needs: + - metadata + - linux + - darwin + runs-on: ubuntu-latest + permissions: + contents: write + actions: read + + steps: + - name: Checkout repository + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + path: dist + + - name: Create GitHub Release + uses: softprops/action-gh-release@153bb8e04406b158c6c84fc1615b65b24149a1fe # v2.6.1 + with: + tag_name: ${{ needs.metadata.outputs.release_tag }} + name: ${{ needs.metadata.outputs.release_tag }} + files: dist/** + # Keep zsh artifact releases out of Codex's normal "latest release" channel. + prerelease: true + + - name: Publish DotSlash manifest + uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + tag: ${{ needs.metadata.outputs.release_tag }} + config: .github/dotslash-zsh-config.json diff --git a/codex-rs/shell-escalation/README.md b/codex-rs/shell-escalation/README.md index e4d3fecd66..69cd038220 100644 --- a/codex-rs/shell-escalation/README.md +++ b/codex-rs/shell-escalation/README.md @@ -27,3 +27,8 @@ git apply /path/to/patches/zsh-exec-wrapper.patch ./configure make -j"$(nproc)" ``` + +Release artifacts are built by `.github/workflows/rust-release-zsh.yml` when a +`codex-zsh-vX.Y.Z` tag is pushed. When the zsh commit or patch changes, publish +the next version tag and update the checked-in DotSlash manifests to use the new +release. From e23e7cbe4634bf1eed1584414e6bbb8a2be349f6 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 25 Jun 2026 14:05:49 -0700 Subject: [PATCH 20/33] release: consume standalone zsh artifacts (#30116) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Once #30114 publishes zsh independently, regular Rust releases should reuse that protected, versioned artifact set instead of rebuilding identical zsh binaries for every Codex version. Keeping the zsh release tag explicit in the workflow also makes future artifact upgrades deliberate and easy to review. This PR assumes the first standalone artifact release will be published as `codex-zsh-v0.1.0` before this change lands. ## What changed - Added `CODEX_ZSH_RELEASE_TAG` near the top of `.github/workflows/rust-release.yml`, initially pinned to `codex-zsh-v0.1.0`. - Download the standalone release’s generated `codex-zsh` DotSlash manifest before assembling Linux and macOS Codex packages. - Added a `--zsh-manifest` package-builder override so release packaging fetches the matching target archive and verifies the size and SHA-256 digest recorded in that manifest. - Removed the reusable zsh build job from regular Rust releases. - Stopped copying zsh archives into each Rust release and stopped regenerating a zsh DotSlash manifest there. Windows packaging remains unchanged because the patched zsh resource is only shipped for supported Unix targets. ## Testing - Added package-helper coverage that supplies a standalone manifest override and verifies the extracted zsh bytes. - Ran the `scripts/codex_package` unit test suite. - Validated `.github/scripts/build-codex-package-archive.sh` with `bash -n`. --- .../scripts/build-codex-package-archive.sh | 5 ++ .github/workflows/rust-release.yml | 44 +++++++------ scripts/codex_package/README.md | 4 +- scripts/codex_package/cli.py | 10 ++- scripts/codex_package/test_zsh.py | 63 +++++++++++++++++++ scripts/codex_package/zsh.py | 7 ++- 6 files changed, 111 insertions(+), 22 deletions(-) create mode 100644 scripts/codex_package/test_zsh.py diff --git a/.github/scripts/build-codex-package-archive.sh b/.github/scripts/build-codex-package-archive.sh index 80da4cf20c..229213cc53 100644 --- a/.github/scripts/build-codex-package-archive.sh +++ b/.github/scripts/build-codex-package-archive.sh @@ -9,6 +9,7 @@ Usage: build-codex-package-archive.sh \ --entrypoint-dir \ --archive-dir \ [--bwrap-bin ] \ + [--zsh-manifest ] \ [--codex-command-runner-bin ] \ [--codex-windows-sandbox-setup-bin ] \ [--target-suffixed-entrypoint] @@ -48,6 +49,10 @@ while [[ $# -gt 0 ]]; do bwrap_bin_provided="true" shift 2 ;; + --zsh-manifest) + resource_args+=(--zsh-manifest "${2:?--zsh-manifest requires a value}") + shift 2 + ;; --codex-command-runner-bin) resource_args+=( --codex-command-runner-bin diff --git a/.github/workflows/rust-release.yml b/.github/workflows/rust-release.yml index a2646e3eae..49d24abc18 100644 --- a/.github/workflows/rust-release.yml +++ b/.github/workflows/rust-release.yml @@ -14,6 +14,9 @@ on: tags: - "rust-v*.*.*" +env: + CODEX_ZSH_RELEASE_TAG: codex-zsh-v0.1.0 + concurrency: group: ${{ github.workflow }} cancel-in-progress: true @@ -354,6 +357,15 @@ jobs: cp target/${{ matrix.target }}/release/codex-${{ matrix.target }}.dmg "$dest/codex-${{ matrix.target }}.dmg" fi + - name: Download packaged zsh manifest + if: ${{ runner.os != 'macOS' }} + shell: bash + run: | + set -euo pipefail + curl -fsSL \ + "https://github.com/${GITHUB_REPOSITORY}/releases/download/${CODEX_ZSH_RELEASE_TAG}/codex-zsh" \ + -o "${RUNNER_TEMP}/codex-zsh" + - name: Build Codex package archive if: ${{ runner.os != 'macOS' }} shell: bash @@ -366,7 +378,8 @@ jobs: --target "$TARGET" \ --bundle "$BUNDLE" \ --entrypoint-dir "target/${TARGET}/release" \ - --archive-dir "dist/${TARGET}" + --archive-dir "dist/${TARGET}" \ + --zsh-manifest "${RUNNER_TEMP}/codex-zsh" - name: Build Python runtime wheel if: ${{ matrix.bundle == 'primary' && runner.os != 'macOS' }} @@ -697,6 +710,14 @@ jobs: cp "target/${{ matrix.target }}/release/${binary}" "$dest/${binary}-${{ matrix.target }}" done + - name: Download packaged zsh manifest + shell: bash + run: | + set -euo pipefail + curl -fsSL \ + "https://github.com/${GITHUB_REPOSITORY}/releases/download/${CODEX_ZSH_RELEASE_TAG}/codex-zsh" \ + -o "${RUNNER_TEMP}/codex-zsh" + - name: Build Codex package archive shell: bash env: @@ -708,7 +729,8 @@ jobs: --target "$TARGET" \ --bundle "$BUNDLE" \ --entrypoint-dir "target/${TARGET}/release" \ - --archive-dir "dist/${TARGET}" + --archive-dir "dist/${TARGET}" \ + --zsh-manifest "${RUNNER_TEMP}/codex-zsh" - name: Build Python runtime wheel if: ${{ matrix.bundle == 'primary' }} @@ -1043,11 +1065,6 @@ jobs: with: publish: true - zsh-release-assets: - name: zsh release assets - needs: tag-check - uses: ./.github/workflows/rust-release-zsh.yml - release: needs: - tag-check @@ -1055,7 +1072,6 @@ jobs: - finalize-macos - build-windows - argument-comment-lint-release-assets - - zsh-release-assets if: >- ${{ always() && @@ -1063,8 +1079,7 @@ jobs: needs.build.result == 'success' && needs.finalize-macos.result == 'success' && needs.build-windows.result == 'success' && - needs.argument-comment-lint-release-assets.result == 'success' && - needs.zsh-release-assets.result == 'success' + needs.argument-comment-lint-release-assets.result == 'success' }} name: release runs-on: ubuntu-latest @@ -1113,7 +1128,7 @@ jobs: uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: path: dist - pattern: "{*-symbols,argument-comment-lint-*,codex-zsh-*,python-runtime-wheel-*}" + pattern: "{*-symbols,argument-comment-lint-*,python-runtime-wheel-*}" - name: List run: ls -R dist/ @@ -1236,13 +1251,6 @@ jobs: tag: ${{ github.ref_name }} config: .github/dotslash-config.json - - uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2 - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - with: - tag: ${{ github.ref_name }} - config: .github/dotslash-zsh-config.json - - uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/scripts/codex_package/README.md b/scripts/codex_package/README.md index 323a3ce5ba..b76c18c3c2 100644 --- a/scripts/codex_package/README.md +++ b/scripts/codex_package/README.md @@ -73,4 +73,6 @@ The patched zsh fork used by `shell_zsh_fork` is fetched from the DotSlash manifest at `scripts/codex_package/codex-zsh` when the selected target has a matching prebuilt artifact. Downloaded archives are cached under `$TMPDIR/codex-package/-zsh` and installed at -`codex-resources/zsh/bin/zsh`. +`codex-resources/zsh/bin/zsh`. Pass `--zsh-manifest` to use a different +DotSlash manifest, such as the manifest published with a standalone zsh +artifact release. diff --git a/scripts/codex_package/cli.py b/scripts/codex_package/cli.py index b7d919e4ff..1dd276827e 100644 --- a/scripts/codex_package/cli.py +++ b/scripts/codex_package/cli.py @@ -90,6 +90,14 @@ def parse_args() -> argparse.Namespace: "targets, bwrap is built with Cargo." ), ) + parser.add_argument( + "--zsh-manifest", + type=Path, + help=( + "Optional DotSlash manifest for the patched zsh fork instead of " + "scripts/codex_package/codex-zsh." + ), + ) parser.add_argument( "--codex-command-runner-bin", type=Path, @@ -160,7 +168,7 @@ def main() -> int: inputs = PackageInputs( entrypoint_bin=source_outputs.entrypoint_bin, rg_bin=resolve_rg_bin(spec, args.rg_bin), - zsh_bin=resolve_zsh_bin(spec), + zsh_bin=resolve_zsh_bin(spec, args.zsh_manifest), bwrap_bin=source_outputs.bwrap_bin, codex_command_runner_bin=source_outputs.codex_command_runner_bin, codex_windows_sandbox_setup_bin=source_outputs.codex_windows_sandbox_setup_bin, diff --git a/scripts/codex_package/test_zsh.py b/scripts/codex_package/test_zsh.py new file mode 100644 index 0000000000..b4f536cfc2 --- /dev/null +++ b/scripts/codex_package/test_zsh.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 + +import hashlib +import json +from pathlib import Path +import sys +import tarfile +import tempfile +import unittest +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from codex_package.targets import TARGET_SPECS +from codex_package.zsh import resolve_zsh_bin + + +class ResolveZshBinTest(unittest.TestCase): + def test_uses_manifest_override(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + archive = root / "codex-zsh.tar.gz" + source = root / "zsh" + source.write_bytes(b"standalone zsh") + with tarfile.open(archive, "w:gz") as tar: + tar.add(source, arcname="codex-zsh/bin/zsh") + + manifest = root / "codex-zsh" + manifest.write_text( + json.dumps( + { + "platforms": { + "linux-x86_64": { + "size": archive.stat().st_size, + "hash": "sha256", + "digest": hashlib.sha256( + archive.read_bytes() + ).hexdigest(), + "format": "tar.gz", + "path": "codex-zsh/bin/zsh", + "providers": [{"url": archive.as_uri()}], + } + } + } + ), + encoding="utf-8", + ) + + with patch( + "codex_package.dotslash.default_cache_root", + return_value=root / "cache", + ): + zsh_bin = resolve_zsh_bin( + TARGET_SPECS["x86_64-unknown-linux-musl"], manifest + ) + + self.assertIsNotNone(zsh_bin) + assert zsh_bin is not None + self.assertEqual(zsh_bin.read_bytes(), b"standalone zsh") + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/codex_package/zsh.py b/scripts/codex_package/zsh.py index f4ee7e19dd..4ec2db502c 100644 --- a/scripts/codex_package/zsh.py +++ b/scripts/codex_package/zsh.py @@ -11,10 +11,13 @@ ZSH_MANIFEST = REPO_ROOT / "scripts" / "codex_package" / "codex-zsh" ZSH_RESOURCE_PATH = Path("zsh") / "bin" / "zsh" -def resolve_zsh_bin(spec: TargetSpec) -> Path | None: +def resolve_zsh_bin( + spec: TargetSpec, + manifest_path: Path | None = None, +) -> Path | None: return fetch_dotslash_executable( spec, - manifest_path=ZSH_MANIFEST, + manifest_path=manifest_path or ZSH_MANIFEST, artifact_label="codex-zsh", cache_key=f"{spec.target}-zsh", dest_name="zsh", From db541f45536f18741b50c74b05904a2c5f628c96 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Thu, 25 Jun 2026 22:15:50 +0100 Subject: [PATCH 21/33] [codex] Add managed MCP server matchers (#29648) ## Summary This PR extends the existing managed `mcp_servers` identity requirement so that one name-qualified rule can use either: - the released exact command or URL identity; - an exact stdio executable with an exact-length, ordered argument matcher list; or - a direct MCP URL matcher. Matcher-based rules stay under the released `identity` key and use the same `McpServerRequirement` abstraction and `mcp_servers.` namespace. ## Behavior Policy activation and name qualification are unchanged: - If `mcp_servers` is absent, ordinary configured MCP servers remain unrestricted. - If `mcp_servers` is present, a server needs a matching same-name requirement. - `mcp_servers = {}` continues to deny every configured MCP server. - Existing exact identity requirements keep their released semantics. Plugin-bundled MCP servers use the same requirement shapes under `plugins..mcp_servers.`. Top-level non-empty rules continue to govern only ordinary configured servers; plugin rules remain explicitly plugin-scoped. The existing globally empty `mcp_servers = {}` plugin kill switch is preserved. Requirements layers continue to use the existing regular TOML merge behavior. Atomic replacement of named MCP requirements is intentionally out of scope here and is tracked independently in #30118. ## Requirement contract The released exact identity contract remains valid: ```toml [mcp_servers.docs.identity] command = "codex-mcp" [mcp_servers.remote.identity] url = "https://example.com/mcp" ``` Command identities continue to check only `command`; they do not inspect arguments, `cwd`, `env`, or `env_vars`. A command matcher uses an exact executable plus an exact-length, ordered argument list. Each argument position supports `exact`, `prefix`, or full-value `regex` matching: ```toml [mcp_servers.internal_mcp_proxy.identity] command = { executable = "company-cli", args = [ { match = "exact", value = "mcp" }, { match = "exact", value = "proxy" }, { match = "exact", value = "--server" }, { match = "regex", expression = '^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$' }, ] } ``` Direct streamable HTTP MCP definitions can use the same value matcher types through `identity.url`: ```toml [mcp_servers.internal_http.identity] url = { match = "regex", expression = '^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?:/.*)?$', } ``` Plugin-bundled MCP matchers use the same contract inside the plugin-qualified allowlist: ```toml [plugins."sample@test".mcp_servers.internal_mcp_proxy.identity] command = { executable = "company-cli", args = [ { match = "exact", value = "mcp" }, { match = "exact", value = "proxy" }, ] } ``` Regexes are validated while managed requirements are loaded, and regex matching must cover the complete value. Command matchers constrain only the executable and arguments. ## Why Enterprise administrators need to allow MCP servers by executable and positional-argument shape, including fixed arguments plus constrained values such as internal MCP URLs passed to a proxy. ## Validation - `just fmt` - `git diff --check` - `just test -p codex-config` (198 passed) - `just test -p codex-core mcp_servers_by_matchers --lib` (2 passed) --- codex-rs/Cargo.lock | 1 + codex-rs/config/Cargo.toml | 1 + codex-rs/config/src/config_requirements.rs | 197 ++++++++++++++-- codex-rs/config/src/constraint.rs | 9 + codex-rs/config/src/lib.rs | 7 +- codex-rs/config/src/mcp_requirements.rs | 163 +++++++++++++ codex-rs/config/src/mcp_requirements_tests.rs | 217 ++++++++++++++++++ codex-rs/core/src/config/config_tests.rs | 202 +++++++++++++++- codex-rs/core/src/config/mod.rs | 26 +-- codex-rs/tui/src/debug_config.rs | 4 +- 10 files changed, 771 insertions(+), 56 deletions(-) create mode 100644 codex-rs/config/src/mcp_requirements.rs create mode 100644 codex-rs/config/src/mcp_requirements_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index bc5953803e..fbae5fd049 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2557,6 +2557,7 @@ dependencies = [ "multimap", "pretty_assertions", "prost 0.14.3", + "regex-lite", "schemars 0.8.22", "serde", "serde_ignored", diff --git a/codex-rs/config/Cargo.toml b/codex-rs/config/Cargo.toml index 509ffd189a..4fe529d743 100644 --- a/codex-rs/config/Cargo.toml +++ b/codex-rs/config/Cargo.toml @@ -30,6 +30,7 @@ gethostname = { workspace = true } indexmap = { workspace = true, features = ["serde"] } multimap = { workspace = true } prost = "0.14.3" +regex-lite = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_ignored = { workspace = true } diff --git a/codex-rs/config/src/config_requirements.rs b/codex-rs/config/src/config_requirements.rs index e757c2149c..199ca4eafa 100644 --- a/codex-rs/config/src/config_requirements.rs +++ b/codex-rs/config/src/config_requirements.rs @@ -19,6 +19,7 @@ use super::requirements_exec_policy::RequirementsExecPolicyToml; use crate::Constrained; use crate::ConstraintError; use crate::ManagedHooksRequirementsToml; +use crate::mcp_requirements::McpServerRequirement; use crate::mcp_types::AppToolApproval; use crate::permissions_toml::PermissionProfileToml; use crate::types::WindowsSandboxModeToml; @@ -217,18 +218,6 @@ impl ConfigRequirements { } } -#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] -#[serde(untagged)] -pub enum McpServerIdentity { - Command { command: String }, - Url { url: String }, -} - -#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] -pub struct McpServerRequirement { - pub identity: McpServerIdentity, -} - #[derive(Deserialize, Debug, Clone, Default, PartialEq, Eq)] pub struct PluginRequirementsToml { pub mcp_servers: Option>, @@ -1201,6 +1190,25 @@ impl ConfigRequirementsToml { } } +fn validate_mcp_server_requirements( + requirements: &BTreeMap, + source: &RequirementSource, + plugin_name: Option<&str>, +) -> Result<(), ConstraintError> { + for (server_name, requirement) in requirements { + requirement + .validate() + .map_err(|reason| ConstraintError::McpServerRequirementParse { + server_name: plugin_name + .map(|plugin_name| format!("{plugin_name}/{server_name}")) + .unwrap_or_else(|| server_name.clone()), + requirement_source: source.clone(), + reason, + })?; + } + Ok(()) +} + impl TryFrom for ConfigRequirements { type Error = ConstraintError; @@ -1233,6 +1241,25 @@ impl TryFrom for ConfigRequirements { guardian_policy_config, } = toml; + if let Some(requirements) = &mcp_servers { + validate_mcp_server_requirements( + &requirements.value, + &requirements.source, + /*plugin_name*/ None, + )?; + } + if let Some(plugin_requirements) = &plugins { + for (plugin_name, plugin) in &plugin_requirements.value { + if let Some(requirements) = &plugin.mcp_servers { + validate_mcp_server_requirements( + requirements, + &plugin_requirements.source, + Some(plugin_name), + )?; + } + } + } + let approval_policy = match allowed_approval_policies { Some(Sourced { value: policies, @@ -1542,6 +1569,9 @@ pub fn sandbox_mode_requirement_for_permission_profile( mod tests { use super::*; use crate::HookEventsToml; + use crate::McpServerCommandMatcher; + use crate::McpServerIdentity; + use crate::McpServerValueMatcher; use anyhow::Result; use codex_execpolicy::Decision; use codex_execpolicy::Evaluation; @@ -3454,6 +3484,9 @@ command = "python3 /enterprise/hooks/pre.py" #[test] fn deserialize_mcp_server_requirements() -> Result<()> { let toml_str = r#" + [mcp_servers.docs] + description = "ignored legacy field" + [mcp_servers.docs.identity] command = "codex-mcp" @@ -3469,7 +3502,7 @@ command = "python3 /enterprise/hooks/pre.py" BTreeMap::from([ ( "docs".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "codex-mcp".to_string(), }, @@ -3477,7 +3510,7 @@ command = "python3 /enterprise/hooks/pre.py" ), ( "remote".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Url { url: "https://example.com/mcp".to_string(), }, @@ -3490,6 +3523,74 @@ command = "python3 /enterprise/hooks/pre.py" Ok(()) } + #[test] + fn deserialize_mcp_server_matcher_requirements() -> Result<()> { + let toml_str = r#" + [mcp_servers.internal_mcp_proxy.identity] + command = { executable = "company-cli", args = [ + { match = "exact", value = "mcp" }, + { match = "exact", value = "proxy" }, + { match = "exact", value = "--server" }, + { match = "regex", expression = '^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$' }, + ] } + "#; + let requirements: ConfigRequirements = + with_unknown_source(from_str(toml_str)?).try_into()?; + + assert_eq!( + requirements.mcp_servers, + Some(Sourced::new( + BTreeMap::from([( + "internal_mcp_proxy".to_string(), + McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![ + McpServerValueMatcher::Exact { + value: "mcp".to_string(), + }, + McpServerValueMatcher::Exact { + value: "proxy".to_string(), + }, + McpServerValueMatcher::Exact { + value: "--server".to_string(), + }, + McpServerValueMatcher::Regex { + expression: r"^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$" + .to_string(), + }, + ], + }), + )]), + RequirementSource::Unknown, + )) + ); + Ok(()) + } + + #[test] + fn invalid_mcp_server_requirement_regex_reports_the_server_name_and_source() -> Result<()> { + let toml_str = r#" + [mcp_servers.broken_rule.identity] + url = { match = "regex", expression = "[" } + "#; + + let err = ConfigRequirements::try_from(with_unknown_source(from_str(toml_str)?)) + .expect_err("invalid matcher regex should fail requirements normalization"); + let ConstraintError::McpServerRequirementParse { + server_name, + requirement_source, + reason, + } = err + else { + panic!("unexpected error: {err:?}"); + }; + + assert_eq!(server_name, "broken_rule"); + assert_eq!(requirement_source, RequirementSource::Unknown); + assert!(reason.contains("invalid regex `[`"), "{reason}"); + Ok(()) + } + #[test] fn deserialize_plugin_mcp_server_requirements() -> Result<()> { let toml_str = r#" @@ -3511,7 +3612,7 @@ command = "python3 /enterprise/hooks/pre.py" PluginRequirementsToml { mcp_servers: Some(BTreeMap::from([( "remote".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Url { url: "https://example.com/mcp".to_string(), }, @@ -3524,7 +3625,7 @@ command = "python3 /enterprise/hooks/pre.py" PluginRequirementsToml { mcp_servers: Some(BTreeMap::from([( "sample".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "sample-mcp".to_string(), }, @@ -3539,6 +3640,70 @@ command = "python3 /enterprise/hooks/pre.py" Ok(()) } + #[test] + fn deserialize_plugin_mcp_server_matcher_requirement() -> Result<()> { + let toml_str = r#" + [plugins."sample@test".mcp_servers.internal_proxy.identity] + command = { executable = "company-cli", args = [ + { match = "exact", value = "mcp" }, + { match = "regex", expression = '^https://[a-z]+\.example\.com$' }, + ] } + "#; + let requirements: ConfigRequirements = + with_unknown_source(from_str(toml_str)?).try_into()?; + + assert_eq!( + requirements.plugins, + Some(Sourced::new( + BTreeMap::from([( + "sample@test".to_string(), + PluginRequirementsToml { + mcp_servers: Some(BTreeMap::from([( + "internal_proxy".to_string(), + McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![ + McpServerValueMatcher::Exact { + value: "mcp".to_string(), + }, + McpServerValueMatcher::Regex { + expression: r"^https://[a-z]+\.example\.com$".to_string(), + }, + ], + }), + )])), + }, + )]), + RequirementSource::Unknown, + )) + ); + Ok(()) + } + + #[test] + fn invalid_plugin_mcp_server_regex_reports_plugin_and_server_name() -> Result<()> { + let toml_str = r#" + [plugins."sample@test".mcp_servers.broken_rule.identity] + url = { match = "regex", expression = "[" } + "#; + + let err = ConfigRequirements::try_from(with_unknown_source(from_str(toml_str)?)) + .expect_err("invalid plugin MCP regex should fail requirements normalization"); + let ConstraintError::McpServerRequirementParse { + server_name, + requirement_source, + reason, + } = err + else { + panic!("unexpected error: {err:?}"); + }; + + assert_eq!(server_name, "sample@test/broken_rule"); + assert_eq!(requirement_source, RequirementSource::Unknown); + assert!(reason.contains("invalid regex `[`"), "{reason}"); + Ok(()) + } + #[test] fn deserialize_exec_policy_requirements() -> Result<()> { let toml_str = r#" diff --git a/codex-rs/config/src/constraint.rs b/codex-rs/config/src/constraint.rs index 64b604cbc2..8628f3909c 100644 --- a/codex-rs/config/src/constraint.rs +++ b/codex-rs/config/src/constraint.rs @@ -24,6 +24,15 @@ pub enum ConstraintError { requirement_source: RequirementSource, reason: String, }, + + #[error( + "invalid requirement for MCP server `{server_name}` (set by {requirement_source}): {reason}" + )] + McpServerRequirementParse { + server_name: String, + requirement_source: RequirementSource, + reason: String, + }, } impl ConstraintError { diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index afda2cf9f3..61e6b6966a 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -12,6 +12,7 @@ mod key_aliases; pub mod loader; mod marketplace_edit; mod mcp_edit; +mod mcp_requirements; mod mcp_types; mod merge; mod overrides; @@ -66,8 +67,6 @@ pub use config_requirements::FilesystemDenyReadPattern; pub use config_requirements::MarketplaceAllowedSourceKind; pub use config_requirements::MarketplaceAllowedSourceToml; pub use config_requirements::MarketplaceRequirementsToml; -pub use config_requirements::McpServerIdentity; -pub use config_requirements::McpServerRequirement; pub use config_requirements::NetworkConstraints; pub use config_requirements::NetworkDomainPermissionToml; pub use config_requirements::NetworkDomainPermissionsToml; @@ -113,6 +112,10 @@ pub use marketplace_edit::remove_user_marketplace; pub use marketplace_edit::remove_user_marketplace_config; pub use mcp_edit::ConfigEditsBuilder; pub use mcp_edit::load_global_mcp_servers; +pub use mcp_requirements::McpServerCommandMatcher; +pub use mcp_requirements::McpServerIdentity; +pub use mcp_requirements::McpServerRequirement; +pub use mcp_requirements::McpServerValueMatcher; pub use mcp_types::AppToolApproval; pub use mcp_types::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; pub use mcp_types::McpServerAuth; diff --git a/codex-rs/config/src/mcp_requirements.rs b/codex-rs/config/src/mcp_requirements.rs new file mode 100644 index 0000000000..b8f9a12d69 --- /dev/null +++ b/codex-rs/config/src/mcp_requirements.rs @@ -0,0 +1,163 @@ +use crate::mcp_types::McpServerConfig; +use crate::mcp_types::McpServerTransportConfig; +use regex_lite::Regex; +use serde::Deserialize; + +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(untagged)] +pub enum McpServerIdentity { + Command { command: String }, + Url { url: String }, +} + +/// String matching operations available to managed MCP server matchers. +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(tag = "match", rename_all = "snake_case", deny_unknown_fields)] +pub enum McpServerValueMatcher { + Exact { value: String }, + Prefix { value: String }, + Regex { expression: String }, +} + +impl McpServerValueMatcher { + fn compile_full_regex(expression: &str) -> Result { + Regex::new(&format!(r"\A(?:{expression})\z")).map_err(|err| { + format!("regex `{expression}` cannot be used for full-value matching: {err}") + }) + } + + fn validate(&self) -> Result<(), String> { + let Self::Regex { expression } = self else { + return Ok(()); + }; + + Regex::new(expression).map_err(|err| format!("invalid regex `{expression}`: {err}"))?; + Self::compile_full_regex(expression).map(|_| ()) + } + + fn matches(&self, candidate: &str) -> bool { + match self { + Self::Exact { value } => candidate == value, + Self::Prefix { value } => candidate.starts_with(value), + Self::Regex { expression } => Self::compile_full_regex(expression) + .ok() + .is_some_and(|regex| regex.is_match(candidate)), + } + } +} + +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct McpServerCommandMatcher { + pub executable: String, + pub args: Vec, +} + +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +struct RawMcpServerCommandIdentity { + command: McpServerCommandMatcher, +} + +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +struct RawMcpServerUrlIdentity { + url: McpServerValueMatcher, +} + +/// A requirement for one named MCP server. +/// +/// The `Identity` variant preserves the released exact-match contract. The +/// command and URL variants are the normalized matcher-based forms accepted +/// under the `identity` key. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum McpServerRequirement { + Identity { identity: McpServerIdentity }, + Command(McpServerCommandMatcher), + Url(McpServerValueMatcher), +} + +#[derive(Deserialize)] +struct RawMcpServerRequirement { + identity: RawMcpServerIdentity, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum RawMcpServerIdentity { + Exact(McpServerIdentity), + Command(RawMcpServerCommandIdentity), + Url(RawMcpServerUrlIdentity), +} + +impl<'de> Deserialize<'de> for McpServerRequirement { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let RawMcpServerRequirement { identity } = + RawMcpServerRequirement::deserialize(deserializer)?; + match identity { + RawMcpServerIdentity::Exact(identity) => Ok(Self::Identity { identity }), + RawMcpServerIdentity::Command(matcher) => Ok(Self::Command(matcher.command)), + RawMcpServerIdentity::Url(matcher) => Ok(Self::Url(matcher.url)), + } + } +} + +impl McpServerRequirement { + pub(crate) fn validate(&self) -> Result<(), String> { + match self { + Self::Identity { .. } => Ok(()), + Self::Command(matcher) => { + for (index, arg) in matcher.args.iter().enumerate() { + arg.validate().map_err(|err| { + format!("invalid argument matcher at index {index}: {err}") + })?; + } + Ok(()) + } + Self::Url(matcher) => matcher.validate(), + } + } + + pub fn matches(&self, server: &McpServerConfig) -> bool { + match (self, &server.transport) { + ( + Self::Identity { + identity: + McpServerIdentity::Command { + command: want_command, + }, + }, + McpServerTransportConfig::Stdio { + command: got_command, + .. + }, + ) => got_command == want_command, + ( + Self::Identity { + identity: McpServerIdentity::Url { url: want_url }, + }, + McpServerTransportConfig::StreamableHttp { url: got_url, .. }, + ) => got_url == want_url, + (Self::Command(matcher), McpServerTransportConfig::Stdio { command, args, .. }) => { + matcher.executable == *command + && matcher.args.len() == args.len() + && matcher + .args + .iter() + .zip(args) + .all(|(matcher, arg)| matcher.matches(arg)) + } + (Self::Url(matcher), McpServerTransportConfig::StreamableHttp { url, .. }) => { + matcher.matches(url) + } + _ => false, + } + } +} + +#[cfg(test)] +#[path = "mcp_requirements_tests.rs"] +mod tests; diff --git a/codex-rs/config/src/mcp_requirements_tests.rs b/codex-rs/config/src/mcp_requirements_tests.rs new file mode 100644 index 0000000000..b854873568 --- /dev/null +++ b/codex-rs/config/src/mcp_requirements_tests.rs @@ -0,0 +1,217 @@ +use super::*; +use crate::mcp_types::McpServerConfig; +use pretty_assertions::assert_eq; +use std::collections::HashMap; + +fn stdio_server(command: &str, args: &[&str]) -> McpServerConfig { + McpServerConfig { + transport: McpServerTransportConfig::Stdio { + command: command.to_string(), + args: args.iter().map(ToString::to_string).collect(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + environment_id: crate::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + } +} + +#[test] +fn command_matcher_matches_exact_positional_arguments() { + let requirement = McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![ + McpServerValueMatcher::Exact { + value: "mcp".to_string(), + }, + McpServerValueMatcher::Regex { + expression: r"https://[a-z]+\.example\.com".to_string(), + }, + ], + }); + + assert!(requirement.matches(&stdio_server( + "company-cli", + &["mcp", "https://pricing.example.com"] + ))); + assert!(!requirement.matches(&stdio_server( + "company-cli", + &["https://pricing.example.com", "mcp"] + ))); + assert!(!requirement.matches(&stdio_server( + "company-cli", + &["mcp", "https://pricing.example.com", "--verbose"] + ))); + assert!(!requirement.matches(&stdio_server( + "/usr/local/bin/company-cli", + &["mcp", "https://pricing.example.com"] + ))); +} + +#[test] +fn regex_matcher_requires_a_full_value_match() { + let matcher = McpServerValueMatcher::Regex { + expression: "mcp".to_string(), + }; + + assert!(matcher.matches("mcp")); + assert!(!matcher.matches("mcp-proxy")); + assert!(!matcher.matches("prefix-mcp")); +} + +#[test] +fn regex_matcher_allows_a_later_alternative_to_match_the_full_value() { + let matcher = McpServerValueMatcher::Regex { + expression: r"https://api\.example\.com|https://api\.example\.com/mcp".to_string(), + }; + + assert!(matcher.matches("https://api.example.com/mcp")); +} + +#[test] +fn regex_matcher_validation_rejects_expression_that_cannot_be_wrapped() { + let matcher = McpServerValueMatcher::Regex { + expression: "(?x)mcp # trailing comment".to_string(), + }; + + let err = matcher + .validate() + .expect_err("expression should not be valid for full-value matching"); + assert!( + err.contains("cannot be used for full-value matching"), + "{err}" + ); +} + +#[test] +fn legacy_command_identity_keeps_ignoring_arguments() { + let requirement: McpServerRequirement = toml::from_str( + r#" +[identity] +command = "company-cli" +"#, + ) + .expect("legacy command identity"); + + assert!(requirement.matches(&stdio_server( + "company-cli", + &["any", "arguments", "remain", "allowed"] + ))); + assert!(!requirement.matches(&stdio_server("different-cli", &[]))); +} + +#[test] +fn requirement_deserializes_command_and_url_matcher_shapes() { + let command: McpServerRequirement = toml::from_str( + r#" +[identity] +command = { executable = "company-cli", args = [ + { match = "exact", value = "mcp" }, + { match = "regex", expression = '^https://[a-z]+\.example\.com$' }, +] } +"#, + ) + .expect("command matcher"); + let url: McpServerRequirement = toml::from_str( + r#" +[identity] +url = { match = "prefix", value = "https://mcp.example.com/" } +"#, + ) + .expect("URL matcher"); + + assert_eq!( + command, + McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![ + McpServerValueMatcher::Exact { + value: "mcp".to_string(), + }, + McpServerValueMatcher::Regex { + expression: r"^https://[a-z]+\.example\.com$".to_string(), + }, + ], + }) + ); + assert_eq!( + url, + McpServerRequirement::Url(McpServerValueMatcher::Prefix { + value: "https://mcp.example.com/".to_string(), + }) + ); +} + +#[test] +fn requirement_rejects_matchers_outside_identity() { + for contents in [ + r#" +command = "company-cli" +"#, + r#" +command = { executable = "company-cli", args = [] } +"#, + r#" +url = { match = "prefix", value = "https://mcp.example.com/" } +"#, + ] { + let err = toml::from_str::(contents) + .expect_err("MCP server requirements should use the identity key"); + assert!( + err.to_string().contains("missing field `identity`"), + "{err}" + ); + } +} + +#[test] +fn matcher_identity_rejects_unknown_fields() { + for contents in [ + r#" +[identity] +unknown = "value" +command = { executable = "company-cli", args = [] } +"#, + r#" +[identity] +command = { executable = "company-cli", args = [], unknown = "value" } +"#, + ] { + toml::from_str::(contents) + .expect_err("matcher identities should reject unknown fields"); + } +} + +#[test] +fn identity_requirement_keeps_ignoring_unrelated_sibling_fields() { + let requirement: McpServerRequirement = toml::from_str( + r#" +unrelated = "ignored" +[identity] +command = "company-cli" +"#, + ) + .expect("legacy identity with unrelated sibling field"); + + assert_eq!( + requirement, + McpServerRequirement::Identity { + identity: McpServerIdentity::Command { + command: "company-cli".to_string(), + }, + } + ); +} diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 237e67d6c0..6f7114928e 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -6,6 +6,10 @@ use codex_config::CONFIG_TOML_FILE; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerSource; use codex_config::ConfigLayerStack; +use codex_config::McpServerCommandMatcher; +use codex_config::McpServerIdentity; +use codex_config::McpServerRequirement; +use codex_config::McpServerValueMatcher; use codex_config::ProfileV2Name; use codex_config::RequirementSource; use codex_config::Sourced; @@ -113,11 +117,15 @@ use std::time::Duration; use tempfile::TempDir; fn stdio_mcp(command: &str) -> McpServerConfig { + stdio_mcp_with_args(command, &[]) +} + +fn stdio_mcp_with_args(command: &str, args: &[&str]) -> McpServerConfig { McpServerConfig { auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: command.to_string(), - args: Vec::new(), + args: args.iter().map(ToString::to_string).collect(), env: None, env_vars: Vec::new(), cwd: None, @@ -4128,7 +4136,7 @@ fn filter_mcp_servers_by_allowlist_enforces_identity_rules() { BTreeMap::from([ ( MISMATCHED_URL_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Url { url: "https://example.com/other".to_string(), }, @@ -4136,7 +4144,7 @@ fn filter_mcp_servers_by_allowlist_enforces_identity_rules() { ), ( MISMATCHED_COMMAND_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "other-cmd".to_string(), }, @@ -4144,7 +4152,7 @@ fn filter_mcp_servers_by_allowlist_enforces_identity_rules() { ), ( MATCHED_URL_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Url { url: GOOD_URL.to_string(), }, @@ -4152,7 +4160,7 @@ fn filter_mcp_servers_by_allowlist_enforces_identity_rules() { ), ( MATCHED_COMMAND_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: GOOD_CMD.to_string(), }, @@ -4209,6 +4217,117 @@ fn filter_mcp_servers_by_allowlist_allows_all_when_unset() { ); } +#[test] +fn filter_mcp_servers_by_matchers_enforces_command_and_positional_args() { + let mut servers = HashMap::from([ + ( + "internal_mcp_proxy".to_string(), + stdio_mcp_with_args( + "company-cli", + &[ + "mcp", + "proxy", + "--server", + "https://pricing.mcp.internal.example.com", + ], + ), + ), + ( + "unlisted".to_string(), + stdio_mcp_with_args( + "company-cli", + &[ + "mcp", + "proxy", + "--server", + "https://pricing.mcp.internal.example.com", + ], + ), + ), + ( + "wrong-order".to_string(), + stdio_mcp_with_args( + "company-cli", + &[ + "proxy", + "mcp", + "--server", + "https://pricing.mcp.internal.example.com", + ], + ), + ), + ( + "trailing-arg".to_string(), + stdio_mcp_with_args( + "company-cli", + &[ + "mcp", + "proxy", + "--server", + "https://pricing.mcp.internal.example.com", + "--verbose", + ], + ), + ), + ( + "wrong-host".to_string(), + stdio_mcp_with_args( + "company-cli", + &["mcp", "proxy", "--server", "https://mcp.example.com"], + ), + ), + ]); + let source = RequirementSource::LegacyManagedConfigTomlFromMdm; + let requirement = McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![ + McpServerValueMatcher::Exact { + value: "mcp".to_string(), + }, + McpServerValueMatcher::Exact { + value: "proxy".to_string(), + }, + McpServerValueMatcher::Exact { + value: "--server".to_string(), + }, + McpServerValueMatcher::Regex { + expression: + r"^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$" + .to_string(), + }, + ], + }); + let requirements = Sourced::new( + BTreeMap::from([ + ("internal_mcp_proxy".to_string(), requirement.clone()), + ("wrong-order".to_string(), requirement.clone()), + ("trailing-arg".to_string(), requirement.clone()), + ("wrong-host".to_string(), requirement), + ]), + source.clone(), + ); + + filter_mcp_servers_by_requirements(&mut servers, Some(&requirements)); + + let reason = Some(McpServerDisabledReason::Requirements { source }); + assert_eq!( + servers + .iter() + .map(|(name, server)| ( + name.clone(), + (server.enabled, server.disabled_reason.clone()) + )) + .collect::)>>(), + HashMap::from([ + ("internal_mcp_proxy".to_string(), (true, None)), + ("unlisted".to_string(), (false, reason.clone())), + ("wrong-order".to_string(), (false, reason.clone())), + ("trailing-arg".to_string(), (false, reason.clone())), + ("wrong-host".to_string(), (false, reason)), + ]) + ); +} + #[test] fn filter_mcp_servers_by_allowlist_blocks_all_when_empty() { let mut servers = HashMap::from([ @@ -4259,7 +4378,7 @@ fn filter_plugin_mcp_servers_by_allowlist_enforces_plugin_and_identity_rules() { mcp_servers: Some(BTreeMap::from([ ( MATCHED_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: GOOD_CMD.to_string(), }, @@ -4267,7 +4386,7 @@ fn filter_plugin_mcp_servers_by_allowlist_enforces_plugin_and_identity_rules() { ), ( MISMATCHED_SERVER.to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: GOOD_CMD.to_string(), }, @@ -4308,7 +4427,7 @@ fn filter_plugin_mcp_servers_by_allowlist_blocks_unlisted_plugin() { codex_config::PluginRequirementsToml { mcp_servers: Some(BTreeMap::from([( "server-a".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "cmd-a".to_string(), }, @@ -4339,6 +4458,65 @@ fn filter_plugin_mcp_servers_by_allowlist_blocks_unlisted_plugin() { ); } +#[test] +fn filter_plugin_mcp_servers_by_matchers_enforces_name_and_invocation() { + const MATCHED_SERVER: &str = "matched"; + const MISMATCHED_SERVER: &str = "mismatched"; + const UNLISTED_SERVER: &str = "unlisted"; + + let mut servers = HashMap::from([ + ( + MATCHED_SERVER.to_string(), + stdio_mcp_with_args("company-cli", &["approved"]), + ), + ( + MISMATCHED_SERVER.to_string(), + stdio_mcp_with_args("company-cli", &["rejected"]), + ), + ( + UNLISTED_SERVER.to_string(), + stdio_mcp_with_args("company-cli", &["approved"]), + ), + ]); + let source = RequirementSource::LegacyManagedConfigTomlFromMdm; + let requirement = McpServerRequirement::Command(McpServerCommandMatcher { + executable: "company-cli".to_string(), + args: vec![McpServerValueMatcher::Exact { + value: "approved".to_string(), + }], + }); + let requirements = Sourced::new( + BTreeMap::from([( + "sample@test".to_string(), + codex_config::PluginRequirementsToml { + mcp_servers: Some(BTreeMap::from([ + (MATCHED_SERVER.to_string(), requirement.clone()), + (MISMATCHED_SERVER.to_string(), requirement), + ])), + }, + )]), + source.clone(), + ); + + filter_plugin_mcp_servers_by_requirements("sample@test", &mut servers, Some(&requirements)); + + let reason = Some(McpServerDisabledReason::Requirements { source }); + assert_eq!( + servers + .iter() + .map(|(name, server)| ( + name.clone(), + (server.enabled, server.disabled_reason.clone()) + )) + .collect::)>>(), + HashMap::from([ + (MATCHED_SERVER.to_string(), (true, None)), + (MISMATCHED_SERVER.to_string(), (false, reason.clone())), + (UNLISTED_SERVER.to_string(), (false, reason)), + ]) + ); +} + #[tokio::test] async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io::Result<()> { let codex_home = TempDir::new()?; @@ -4348,7 +4526,7 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: let mcp_requirements = BTreeMap::from([ ( "session_overrides_user".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "session-command".to_string(), }, @@ -4356,7 +4534,7 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: ), ( "managed_overrides_session".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "managed-command".to_string(), }, @@ -4364,7 +4542,7 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: ), ( "fresh_global".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "fresh-global-command".to_string(), }, @@ -4372,7 +4550,7 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: ), ( "fresh_project".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "fresh-project-command".to_string(), }, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 897a31c8db..d90d010102 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -14,7 +14,6 @@ use codex_config::ConfigRequirements; use codex_config::ConfigRequirementsToml; use codex_config::ConstrainedWithSource; use codex_config::FeatureRequirementsToml; -use codex_config::McpServerIdentity; use codex_config::McpServerRequirement; use codex_config::PluginRequirementsToml; use codex_config::ProfileV2Name; @@ -40,7 +39,6 @@ use codex_config::types::AuthKeyringBackendKind; use codex_config::types::History; use codex_config::types::McpServerConfig; use codex_config::types::McpServerDisabledReason; -use codex_config::types::McpServerTransportConfig; use codex_config::types::MemoriesConfig; use codex_config::types::ModelAvailabilityNuxConfig; use codex_config::types::Notice; @@ -1890,7 +1888,7 @@ fn filter_mcp_servers_by_requirements( let allowed = allowlist .value .get(name) - .is_some_and(|requirement| mcp_server_matches_requirement(requirement, server)); + .is_some_and(|requirement| requirement.matches(server)); if allowed { server.disabled_reason = None; } else { @@ -1919,7 +1917,7 @@ fn filter_plugin_mcp_servers_by_requirements( for (name, server) in mcp_servers.iter_mut() { let allowed = plugin_mcp_requirements .and_then(|mcp_requirements| mcp_requirements.get(name)) - .is_some_and(|requirement| mcp_server_matches_requirement(requirement, server)); + .is_some_and(|requirement| requirement.matches(server)); if allowed { server.disabled_reason = None; } else { @@ -1982,26 +1980,6 @@ where Ok(false) } -fn mcp_server_matches_requirement( - requirement: &McpServerRequirement, - server: &McpServerConfig, -) -> bool { - match &requirement.identity { - McpServerIdentity::Command { - command: want_command, - } => matches!( - &server.transport, - McpServerTransportConfig::Stdio { command: got_command, .. } - if got_command == want_command - ), - McpServerIdentity::Url { url: want_url } => matches!( - &server.transport, - McpServerTransportConfig::StreamableHttp { url: got_url, .. } - if got_url == want_url - ), - } -} - pub async fn load_global_mcp_servers( codex_home: &Path, ) -> std::io::Result> { diff --git a/codex-rs/tui/src/debug_config.rs b/codex-rs/tui/src/debug_config.rs index ed9aeffddb..7a35cc4506 100644 --- a/codex-rs/tui/src/debug_config.rs +++ b/codex-rs/tui/src/debug_config.rs @@ -699,7 +699,7 @@ mod tests { mcp_servers: Some(Sourced::new( BTreeMap::from([( "docs".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "codex-mcp".to_string(), }, @@ -778,7 +778,7 @@ mod tests { hooks: None, mcp_servers: Some(BTreeMap::from([( "docs".to_string(), - McpServerRequirement { + McpServerRequirement::Identity { identity: McpServerIdentity::Command { command: "codex-mcp".to_string(), }, From c9e6d9783dfe99a819f00318f753b6f5cbcc38e8 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 25 Jun 2026 22:23:51 +0100 Subject: [PATCH 22/33] Let extensions contribute World State sections (#30100) ## Why #29856 already owns the durable thread intent and exact environment binding. This PR adds only the small missing extension boundary: an extension can contribute one named World State section, while core still owns persistence, diffing, and model-visible fragment types. This lets skills stay in the skills extension instead of moving their runtime into core. ## Shape ```text extension-owned state | | contribute section id + JSON snapshot + renderer v core World State | | compare with the previous snapshot v no message, or one incremental model-visible update ``` The extension API is deliberately small: ```rust fn contribute_world_state(...) -> Vec ``` Core adapts the rendered result to `ContextualUserFragment`, records the snapshot, and keeps the existing compaction/resume behavior. ## What changes - Adds extension-owned World State section contributions. - Calls those contributors from the existing per-step World State builder. - Restores durable selected capability roots into extension thread state on resume. - Keeps the actual model-context fragment and rollout machinery in core. ## What does not change - No skill or MCP implementation moves out of its extension. - No new file watcher, generation, or RPC. - No generic migration of existing World State sections. - No change to the stable environment-ID assumption from #29856. ## Example ```text step 1 snapshot: skills = [] step 2 snapshot: skills = [executor-demo:deploy] core asks the skills extension to render only that change. ``` ## Stack 1. **This PR:** let extensions contribute World State sections. 2. Project executor skills through the skills extension. 3. Pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario. --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/context/world_state/mod.rs | 61 +++++++++ .../context/world_state/world_state_tests.rs | 39 ++++++ codex-rs/core/src/session/session.rs | 17 ++- codex-rs/core/src/session/world_state.rs | 17 +++ codex-rs/ext/extension-api/Cargo.toml | 1 + .../ext/extension-api/src/contributors.rs | 16 +++ .../src/contributors/world_state.rs | 122 ++++++++++++++++++ codex-rs/ext/extension-api/src/lib.rs | 4 + 9 files changed, 273 insertions(+), 5 deletions(-) create mode 100644 codex-rs/ext/extension-api/src/contributors/world_state.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index fbae5fd049..3837df62ea 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3010,6 +3010,7 @@ dependencies = [ "codex-tools", "codex-utils-absolute-path", "pretty_assertions", + "serde_json", "tokio", ] diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 50b9967096..94960a8322 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -2,6 +2,9 @@ mod agents_md; mod environment; use crate::context::ContextualUserFragment; +use codex_extension_api::PreviousWorldStateSection; +use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::WorldStateSectionContribution; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use indexmap::IndexMap; @@ -83,6 +86,54 @@ impl ErasedWorldStateSection for S { } } +struct ExtensionWorldStateSection(WorldStateSectionContribution); + +impl ErasedWorldStateSection for ExtensionWorldStateSection { + fn snapshot(&self) -> Option { + let mut snapshot = self.0.snapshot().clone(); + remove_null_object_fields(&mut snapshot); + (!snapshot.is_null()).then_some(snapshot) + } + + fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { + self.0.matches_legacy_fragment(role, text) + } + + fn render_diff( + &self, + previous: PreviousSectionState<'_, Value>, + ) -> Option> { + let previous = match previous { + PreviousSectionState::Absent => PreviousWorldStateSection::Absent, + PreviousSectionState::Unknown => PreviousWorldStateSection::Unknown, + PreviousSectionState::Known(previous) => PreviousWorldStateSection::Known(previous), + }; + self.0 + .render_diff(previous) + .map(|fragment| Box::new(WorldStateContextFragment(fragment)) as _) + } +} + +struct WorldStateContextFragment(RenderedWorldStateFragment); + +impl ContextualUserFragment for WorldStateContextFragment { + fn role(&self) -> &'static str { + self.0.role() + } + + fn markers(&self) -> (&'static str, &'static str) { + self.0.markers() + } + + fn body(&self) -> String { + self.0.body().to_string() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } +} + /// What is known about a section's previously model-visible state. pub(crate) enum PreviousSectionState<'a, T> { /// No persisted snapshot or matching fragment exists in retained history. @@ -169,6 +220,16 @@ impl WorldState { self.sections.insert(id, Box::new(section)); } + pub(crate) fn add_extension_section(&mut self, section: WorldStateSectionContribution) { + let id = section.id(); + assert!( + !self.sections.contains_key(id), + "duplicate world-state section ID: {id}" + ); + self.sections + .insert(id, Box::new(ExtensionWorldStateSection(section))); + } + pub(crate) fn snapshot(&self) -> WorldStateSnapshot { WorldStateSnapshot { sections: self diff --git a/codex-rs/core/src/context/world_state/world_state_tests.rs b/codex-rs/core/src/context/world_state/world_state_tests.rs index d17304bb7b..7ff6f73285 100644 --- a/codex-rs/core/src/context/world_state/world_state_tests.rs +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -110,6 +110,45 @@ fn render_diff_restores_the_typed_section_snapshot() { ); } +#[test] +fn extension_owned_section_uses_its_snapshot_and_renderer() { + let mut world_state = WorldState::default(); + world_state.add_extension_section(WorldStateSectionContribution::new( + "extension_test", + json!({"value": "after", "optional": null}), + |previous| match previous { + PreviousWorldStateSection::Known(previous) + if previous == &json!({"value": "before"}) => + { + Some(RenderedWorldStateFragment::new( + "developer", + ("", ""), + "after", + )) + } + PreviousWorldStateSection::Absent + | PreviousWorldStateSection::Unknown + | PreviousWorldStateSection::Known(_) => None, + }, + )); + let previous = WorldStateSnapshot { + sections: BTreeMap::from([("extension_test".to_string(), json!({"value": "before"}))]), + }; + + let rendered = world_state.render_diff(&previous); + + assert_eq!( + serde_json::to_value(world_state.snapshot()).expect("serialize world-state snapshot"), + json!({"extension_test": {"value": "after"}}) + ); + assert_eq!(rendered.len(), 1); + assert_eq!(rendered[0].role(), "developer"); + assert_eq!( + rendered[0].render(), + "after" + ); +} + #[test] fn unreadable_section_snapshot_is_treated_as_unknown() { let mut current = WorldState::default(); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 57a66b9986..f385a8c983 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -490,7 +490,7 @@ impl Session { plugins_manager: Arc, mcp_manager: Arc, extensions: Arc>, - thread_extension_init: ExtensionDataInit, + mut thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, agent_control: AgentControl, environment_manager: Arc, @@ -556,10 +556,17 @@ impl Session { config.current_time_reminder.as_ref(), external_time_provider, )?; - let selected_capability_roots = thread_extension_init - .get::>() - .map(|roots| roots.as_ref().clone()) - .unwrap_or_else(|| initial_history.get_selected_capability_roots()); + let selected_capability_roots = + match thread_extension_init.get::>() { + Some(roots) => roots.as_ref().clone(), + None => { + let roots = initial_history.get_selected_capability_roots(); + if !roots.is_empty() { + thread_extension_init.insert(roots.clone()); + } + roots + } + }; let mcp_thread_init = thread_extension_init.clone(); let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index cedeaab25e..2541972252 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -3,6 +3,7 @@ use super::step_context::StepContext; use crate::context::world_state::AgentsMdState; use crate::context::world_state::EnvironmentsState; use crate::context::world_state::WorldState; +use codex_extension_api::WorldStateContributionInput; impl Session { pub(crate) async fn build_world_state_for_step( @@ -34,6 +35,22 @@ impl Session { .with_subagents(environment_subagents), ); } + let environments = step_context.environments.to_selections(); + for contributor in self.services.extensions.context_contributors() { + for section in contributor + .contribute_world_state(WorldStateContributionInput { + thread_id: self.thread_id(), + turn_id: turn_context.sub_id.as_str(), + environments: &environments, + session_store: &self.services.session_extension_data, + thread_store: &self.services.thread_extension_data, + turn_store: turn_context.extension_data.as_ref(), + }) + .await + { + world_state.add_extension_section(section); + } + } world_state } } diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 1a5a9742fb..ce441576ec 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -19,6 +19,7 @@ codex-context-fragments = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } +serde_json = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 7d16efd3d4..96f1ee0b48 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -18,6 +18,7 @@ mod thread_lifecycle; mod tool_lifecycle; mod turn_input; mod turn_lifecycle; +mod world_state; pub use context::TurnContextContributionInput; pub use mcp::McpServerContribution; @@ -39,6 +40,10 @@ pub use turn_lifecycle::TurnAbortInput; pub use turn_lifecycle::TurnErrorInput; pub use turn_lifecycle::TurnStartInput; pub use turn_lifecycle::TurnStopInput; +pub use world_state::PreviousWorldStateSection; +pub use world_state::RenderedWorldStateFragment; +pub use world_state::WorldStateContributionInput; +pub use world_state::WorldStateSectionContribution; /// Boxed, sendable future returned by asynchronous extension contributors. pub type ExtensionFuture<'a, T> = Pin + Send + 'a>>; @@ -92,6 +97,17 @@ pub trait ContextContributor: Send + Sync { Vec::new() }) } + + fn contribute_world_state<'a>( + &'a self, + input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let _self = self; + let _input = input; + Vec::new() + }) + } } /// Contributor for host-owned thread lifecycle gates. diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs new file mode 100644 index 0000000000..b5bb548105 --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -0,0 +1,122 @@ +use std::sync::Arc; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::TurnEnvironmentSelection; +use serde_json::Value; + +use crate::ExtensionData; + +/// Host state available while an extension contributes one sampling step's World State. +pub struct WorldStateContributionInput<'a> { + pub thread_id: ThreadId, + pub turn_id: &'a str, + pub environments: &'a [TurnEnvironmentSelection], + pub session_store: &'a ExtensionData, + pub thread_store: &'a ExtensionData, + pub turn_store: &'a ExtensionData, +} + +/// What the harness knows about the previous value of one extension-owned section. +pub enum PreviousWorldStateSection<'a> { + Absent, + Unknown, + Known(&'a Value), +} + +/// Plain model-visible data rendered by an extension-owned World State section. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RenderedWorldStateFragment { + role: &'static str, + markers: (&'static str, &'static str), + body: String, +} + +impl RenderedWorldStateFragment { + pub fn new( + role: &'static str, + markers: (&'static str, &'static str), + body: impl Into, + ) -> Self { + Self { + role, + markers, + body: body.into(), + } + } + + pub fn role(&self) -> &'static str { + self.role + } + + pub fn markers(&self) -> (&'static str, &'static str) { + self.markers + } + + pub fn body(&self) -> &str { + &self.body + } +} + +type RenderDiff = dyn for<'a> Fn(PreviousWorldStateSection<'a>) -> Option + + Send + + Sync; +type LegacyFragmentMatcher = dyn Fn(&str, &str) -> bool + Send + Sync; + +/// One extension-owned World State section captured for a sampling step. +/// +/// The extension owns the stable ID, comparison snapshot, and diff rendering. The harness owns +/// persistence and the concrete model-context fragment envelope. +#[derive(Clone)] +pub struct WorldStateSectionContribution { + id: &'static str, + snapshot: Value, + render_diff: Arc, + matches_legacy_fragment: Arc, +} + +impl WorldStateSectionContribution { + pub fn new( + id: &'static str, + snapshot: Value, + render_diff: impl for<'a> Fn( + PreviousWorldStateSection<'a>, + ) -> Option + + Send + + Sync + + 'static, + ) -> Self { + Self { + id, + snapshot, + render_diff: Arc::new(render_diff), + matches_legacy_fragment: Arc::new(|_, _| false), + } + } + + pub fn with_legacy_matcher( + mut self, + matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static, + ) -> Self { + self.matches_legacy_fragment = Arc::new(matcher); + self + } + + pub fn id(&self) -> &'static str { + self.id + } + + pub fn snapshot(&self) -> &Value { + &self.snapshot + } + + pub fn render_diff( + &self, + previous: PreviousWorldStateSection<'_>, + ) -> Option { + (self.render_diff)(previous) + } + + pub fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { + (self.matches_legacy_fragment)(role, text) + } +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index de0e0c99f8..c3873a2ed3 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -38,8 +38,10 @@ pub use contributors::ExtensionFuture; pub use contributors::McpServerContribution; pub use contributors::McpServerContributionContext; pub use contributors::McpServerContributor; +pub use contributors::PreviousWorldStateSection; pub use contributors::PromptFragment; pub use contributors::PromptSlot; +pub use contributors::RenderedWorldStateFragment; pub use contributors::ThreadIdleInput; pub use contributors::ThreadLifecycleContributor; pub use contributors::ThreadResumeInput; @@ -63,6 +65,8 @@ pub use contributors::TurnItemContributor; pub use contributors::TurnLifecycleContributor; pub use contributors::TurnStartInput; pub use contributors::TurnStopInput; +pub use contributors::WorldStateContributionInput; +pub use contributors::WorldStateSectionContribution; pub use registry::ExtensionRegistry; pub use registry::ExtensionRegistryBuilder; pub use registry::empty_extension_registry; From b80fbb70cd2f6e6ae32cca8c2eba923bf3c0055e Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Thu, 25 Jun 2026 18:44:35 -0300 Subject: [PATCH 23/33] fix(app-server): suppress TUI rollback warning (#30124) ## Why The TUI uses `thread/rollback` internally for user-facing flows such as prompt cancellation/backtracking. After `thread/rollback` was marked deprecated, those internal calls started surfacing `deprecationNotice` messages in the TUI, even though the user did not explicitly call the deprecated app-server API. The endpoint should remain deprecated for external app-server clients, but the built-in `codex-tui` client should not show this implementation-detail warning during normal interaction. ## What changed - Pass the initialized app-server client name into the `thread/rollback` request processor. - Suppress the `thread/rollback` deprecation notice only for `codex-tui`. - Preserve the existing `deprecationNotice` behavior for non-TUI clients. - Add regression coverage for the `codex-tui` suppression path. ## How to Test 1. Start Codex TUI from this branch. 2. Type text into the composer and press `Esc` to cancel/backtrack. 3. Confirm the TUI restores/cancels the prompt without showing `thread/rollback is deprecated and will be removed soon`. 4. Also verify an external app-server client that calls `thread/rollback` still receives `deprecationNotice`. Targeted tests: - `just test -p codex-app-server thread_rollback` - `just argument-comment-lint` --- codex-rs/app-server/src/message_processor.rs | 2 +- .../request_processors/thread_processor.rs | 8 +++- .../tests/suite/v2/thread_rollback.rs | 43 +++++++++++++++++++ 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 41ecbfe561..3322f24d8b 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -1205,7 +1205,7 @@ impl MessageProcessor { } ClientRequest::ThreadRollback { params, .. } => { self.thread_processor - .thread_rollback(&request_id, params) + .thread_rollback(&request_id, params, app_server_client_name.as_deref()) .await } ClientRequest::ThreadList { params, .. } => { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 7c6b0aad74..0fc37d1a4f 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -8,6 +8,7 @@ use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; +const CODEX_TUI_CLIENT_NAME: &str = "codex-tui"; const THREAD_ROLLBACK_DEPRECATION_SUMMARY: &str = "thread/rollback is deprecated and will be removed soon"; @@ -634,9 +635,12 @@ impl ThreadRequestProcessor { &self, request_id: &ConnectionRequestId, params: ThreadRollbackParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.send_thread_rollback_deprecation_notice(request_id.connection_id) - .await; + if app_server_client_name != Some(CODEX_TUI_CLIENT_NAME) { + self.send_thread_rollback_deprecation_notice(request_id.connection_id) + .await; + } self.thread_rollback_inner(request_id, params) .await .map(|()| None) diff --git a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs index 9481f69232..328ae529c8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs @@ -3,6 +3,7 @@ use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; +use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::DeprecationNoticeNotification; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; @@ -24,6 +25,48 @@ use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +#[tokio::test] +async fn thread_rollback_does_not_emit_deprecation_notice_to_codex_tui() -> Result<()> { + let codex_home = TempDir::new()?; + let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + let initialized = timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_client_info(ClientInfo { + name: "codex-tui".to_string(), + title: None, + version: "0.1.0".to_string(), + }), + ) + .await??; + let JSONRPCMessage::Response(_) = initialized else { + panic!("expected initialize response, got {initialized:?}"); + }; + mcp.clear_message_buffer(); + + let rollback_id = mcp + .send_thread_rollback_request(ThreadRollbackParams { + thread_id: "00000000-0000-0000-0000-000000000001".to_string(), + num_turns: 1, + }) + .await?; + loop { + let message = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??; + match message { + JSONRPCMessage::Notification(notification) => { + assert_ne!(notification.method, "deprecationNotice"); + } + JSONRPCMessage::Error(error) if error.id == RequestId::Integer(rollback_id) => { + break; + } + message => { + panic!("expected rollback error response, got {message:?}"); + } + } + } + + Ok(()) +} + #[tokio::test] async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<()> { // Three Codex turns hit the mock model (session start + two turn/start calls). From a6d20ed29701bd8f291e94526154dcac90367d92 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Thu, 25 Jun 2026 22:50:36 +0100 Subject: [PATCH 24/33] [codex] Surface MCP reauthentication-required startup failures (#29877) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - distinguish expired, non-refreshable stored MCP OAuth credentials from first-time missing credentials - carry a typed `failureReason: "reauthenticationRequired"` on the existing `mcpServer/startupStatus/updated` notification only when user action is required - keep the public MCP auth-status API unchanged and regenerate the app-server protocol schemas and documentation ## Why An MCP server with an expired access token and no usable refresh token currently fails startup without giving clients a reliable, typed recovery signal. The existing startup-status notification is the natural place to carry this state. Its nullable `failureReason` keeps the recovery reason attached to the failed startup transition without adding a one-off notification. Internally, Codex distinguishes first-time login from reauthentication and emits the reason only when the startup error itself requires authentication. ## User impact App clients can prompt an existing user to reconnect an MCP server when automatic recovery is impossible by handling a failed `mcpServer/startupStatus/updated` notification whose `failureReason` is `reauthenticationRequired`. Starting, ready, cancelled, unrelated failures, and first-time setup carry no reauthentication reason. ## Companion app PR - openai/openai#1069582 ## Validation - `just test -p codex-app-server-protocol` — 248 passed; schema fixture tests passed - `cargo check -p codex-app-server -p codex-tui` - `just test -p codex-rmcp-client -p codex-mcp` — 184 passed, 2 skipped - `just test -p codex-protocol -p codex-app-server-protocol -p codex-mcp` — 579 passed - `just write-app-server-schema` - `just fmt` --- .../schema/json/ServerNotification.json | 16 ++++++ .../codex_app_server_protocol.schemas.json | 16 ++++++ .../codex_app_server_protocol.v2.schemas.json | 16 ++++++ .../McpServerStatusUpdatedNotification.json | 16 ++++++ .../v2/McpServerStartupFailureReason.ts | 5 ++ .../v2/McpServerStatusUpdatedNotification.ts | 3 +- .../schema/typescript/v2/index.ts | 1 + .../src/protocol/v2/mcp.rs | 7 +++ .../src/protocol/v2/tests.rs | 29 ++++++++++ codex-rs/app-server/README.md | 4 +- .../app-server/src/bespoke_event_handling.rs | 17 +++--- .../app-server/tests/suite/v2/thread_start.rs | 2 + codex-rs/cli/src/mcp_cmd.rs | 6 +-- codex-rs/codex-mcp/src/connection_manager.rs | 39 ++++++++++++-- .../codex-mcp/src/connection_manager_tests.rs | 50 +++++++++++++++-- codex-rs/codex-mcp/src/mcp/auth.rs | 23 ++++---- codex-rs/codex-mcp/src/mcp/mod.rs | 2 +- codex-rs/codex-mcp/src/rmcp_client.rs | 40 +++++++++++++- codex-rs/config/src/mcp_requirements_tests.rs | 1 + codex-rs/protocol/src/protocol.rs | 19 ++++++- codex-rs/rmcp-client/src/auth_status.rs | 53 ++++++++++++++----- codex-rs/rmcp-client/src/lib.rs | 2 + .../tests/streamable_http_oauth_startup.rs | 29 ++++++++-- .../tui/src/app/app_server_event_targets.rs | 2 + codex-rs/tui/src/app/tests.rs | 3 ++ codex-rs/tui/src/app/thread_events.rs | 1 + .../tui/src/chatwidget/tests/mcp_startup.rs | 3 ++ 27 files changed, 350 insertions(+), 55 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/McpServerStartupFailureReason.ts diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 18435f6f2b..8a8064c0b9 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -2384,6 +2384,12 @@ ], "type": "object" }, + "McpServerStartupFailureReason": { + "enum": [ + "reauthenticationRequired" + ], + "type": "string" + }, "McpServerStartupState": { "enum": [ "starting", @@ -2401,6 +2407,16 @@ "null" ] }, + "failureReason": { + "anyOf": [ + { + "$ref": "#/definitions/McpServerStartupFailureReason" + }, + { + "type": "null" + } + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 8625c08da4..56ea39276a 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -12094,6 +12094,12 @@ "title": "McpServerRefreshResponse", "type": "object" }, + "McpServerStartupFailureReason": { + "enum": [ + "reauthenticationRequired" + ], + "type": "string" + }, "McpServerStartupState": { "enum": [ "starting", @@ -12165,6 +12171,16 @@ "null" ] }, + "failureReason": { + "anyOf": [ + { + "$ref": "#/definitions/v2/McpServerStartupFailureReason" + }, + { + "type": "null" + } + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 94fb8a9890..0f2910befd 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -8498,6 +8498,12 @@ "title": "McpServerRefreshResponse", "type": "object" }, + "McpServerStartupFailureReason": { + "enum": [ + "reauthenticationRequired" + ], + "type": "string" + }, "McpServerStartupState": { "enum": [ "starting", @@ -8569,6 +8575,16 @@ "null" ] }, + "failureReason": { + "anyOf": [ + { + "$ref": "#/definitions/McpServerStartupFailureReason" + }, + { + "type": "null" + } + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/McpServerStatusUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/McpServerStatusUpdatedNotification.json index 87dab90cff..8efb36bdad 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/McpServerStatusUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/McpServerStatusUpdatedNotification.json @@ -1,6 +1,12 @@ { "$schema": "http://json-schema.org/draft-07/schema#", "definitions": { + "McpServerStartupFailureReason": { + "enum": [ + "reauthenticationRequired" + ], + "type": "string" + }, "McpServerStartupState": { "enum": [ "starting", @@ -18,6 +24,16 @@ "null" ] }, + "failureReason": { + "anyOf": [ + { + "$ref": "#/definitions/McpServerStartupFailureReason" + }, + { + "type": "null" + } + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStartupFailureReason.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStartupFailureReason.ts new file mode 100644 index 0000000000..0373e5444b --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStartupFailureReason.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type McpServerStartupFailureReason = "reauthenticationRequired"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStatusUpdatedNotification.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStatusUpdatedNotification.ts index b5c0a9058c..fd192f22c6 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStatusUpdatedNotification.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerStatusUpdatedNotification.ts @@ -1,6 +1,7 @@ // GENERATED CODE! DO NOT MODIFY BY HAND! // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { McpServerStartupFailureReason } from "./McpServerStartupFailureReason"; import type { McpServerStartupState } from "./McpServerStartupState"; -export type McpServerStatusUpdatedNotification = { threadId: string | null, name: string, status: McpServerStartupState, error: string | null, }; +export type McpServerStatusUpdatedNotification = { threadId: string | null, name: string, status: McpServerStartupState, error: string | null, failureReason: McpServerStartupFailureReason | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index ff7df9e9ec..82b74711c6 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -238,6 +238,7 @@ export type { McpServerOauthLoginCompletedNotification } from "./McpServerOauthL export type { McpServerOauthLoginParams } from "./McpServerOauthLoginParams"; export type { McpServerOauthLoginResponse } from "./McpServerOauthLoginResponse"; export type { McpServerRefreshResponse } from "./McpServerRefreshResponse"; +export type { McpServerStartupFailureReason } from "./McpServerStartupFailureReason"; export type { McpServerStartupState } from "./McpServerStartupState"; export type { McpServerStatus } from "./McpServerStatus"; export type { McpServerStatusDetail } from "./McpServerStatusDetail"; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs index ca0e124d59..d24da0c879 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs @@ -23,6 +23,12 @@ v2_enum_from_core!( } ); +v2_enum_from_core!( + pub enum McpServerStartupFailureReason from codex_protocol::protocol::McpStartupFailureReason { + ReauthenticationRequired + } +); + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] @@ -242,6 +248,7 @@ pub struct McpServerStatusUpdatedNotification { pub name: String, pub status: McpServerStartupState, pub error: Option, + pub failure_reason: Option, } #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 7c78fee4c8..03c2bcaa5b 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::ServerNotification; use codex_protocol::approvals::ElicitationRequest as CoreElicitationRequest; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::items::AgentMessageContent; @@ -2155,6 +2156,7 @@ fn mcp_server_status_updated_accepts_missing_thread_id() { name: "optional_broken".to_string(), status: McpServerStartupState::Failed, error: Some("handshake failed".to_string()), + failure_reason: None, }; assert_eq!(notification, expected); assert_eq!( @@ -2164,6 +2166,33 @@ fn mcp_server_status_updated_accepts_missing_thread_id() { "name": "optional_broken", "status": "failed", "error": "handshake failed", + "failureReason": null, + }) + ); +} + +#[test] +fn mcp_server_status_updated_serializes_failure_reason() { + let notification = + ServerNotification::McpServerStatusUpdated(McpServerStatusUpdatedNotification { + thread_id: Some("thread-1".to_string()), + name: "expired-oauth".to_string(), + status: McpServerStartupState::Failed, + error: Some("OAuth credentials expired".to_string()), + failure_reason: Some(McpServerStartupFailureReason::ReauthenticationRequired), + }); + + assert_eq!( + serde_json::to_value(notification).expect("notification should serialize"), + json!({ + "method": "mcpServer/startupStatus/updated", + "params": { + "threadId": "thread-1", + "name": "expired-oauth", + "status": "failed", + "error": "OAuth credentials expired", + "failureReason": "reauthenticationRequired", + }, }) ); } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 1a3684a024..a4b488acdf 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1345,7 +1345,7 @@ Because audio is intentionally separate from `ThreadItem`, clients can opt out o ### MCP server startup events -- `mcpServer/startupStatus/updated` — `{ threadId, name, status, error }` when app-server observes an MCP server startup transition. `threadId` identifies the owning thread when startup is thread-scoped and is `null` when startup is app-scoped. `status` is one of `starting`, `ready`, `failed`, or `cancelled`. `error` is `null` except for `failed`. +- `mcpServer/startupStatus/updated` — `{ threadId, name, status, error, failureReason }` when app-server observes an MCP server startup transition. `threadId` identifies the owning thread when startup is thread-scoped and is `null` when startup is app-scoped. `status` is one of `starting`, `ready`, `failed`, or `cancelled`. `error` and `failureReason` are `null` except for `failed`; `failureReason` is `reauthenticationRequired` when stored OAuth credentials have expired and cannot be refreshed, so clients can prompt the user to reconnect the named server. ### Turn events @@ -1935,7 +1935,7 @@ Codex supports these authentication modes. The current mode is surfaced in `acco - `account/rateLimits/updated` (notify) — emitted whenever a user's ChatGPT rate limits change. This is a sparse rolling update; merge available values into the most recent `account/rateLimits/read` response or refetch that snapshot. - `account/sendAddCreditsNudgeEmail` — ask ChatGPT to email the workspace owner about depleted credits or a reached usage limit. - `mcpServer/oauthLogin/completed` (notify) — emitted after a `mcpServer/oauth/login` flow finishes for a server; payload includes `{ name, threadId, success, error? }`. -- `mcpServer/startupStatus/updated` (notify) — emitted when a configured MCP server's startup status changes; payload includes `{ threadId, name, status, error }`, where `threadId` is the owning thread when startup is thread-scoped and `null` when it is app-scoped, and `status` is `starting`, `ready`, `failed`, or `cancelled`. +- `mcpServer/startupStatus/updated` (notify) — emitted when a configured MCP server's startup status changes; payload includes `{ threadId, name, status, error, failureReason }`, where `threadId` is the owning thread when startup is thread-scoped and `null` when it is app-scoped, and `status` is `starting`, `ready`, `failed`, or `cancelled`. `failureReason` is `reauthenticationRequired` when stored OAuth credentials have expired and cannot be refreshed, so clients can prompt the user to reconnect the named server. ### 1) Check auth state diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 36dbdd9d3b..bcba603bcd 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -200,18 +200,20 @@ pub(crate) async fn apply_bespoke_event_handling( .await; } EventMsg::McpStartupUpdate(update) => { - let (status, error) = match update.status { + let (status, error, failure_reason) = match update.status { codex_protocol::protocol::McpStartupStatus::Starting => { - (McpServerStartupState::Starting, None) + (McpServerStartupState::Starting, None, None) } codex_protocol::protocol::McpStartupStatus::Ready => { - (McpServerStartupState::Ready, None) - } - codex_protocol::protocol::McpStartupStatus::Failed { error } => { - (McpServerStartupState::Failed, Some(error)) + (McpServerStartupState::Ready, None, None) } + codex_protocol::protocol::McpStartupStatus::Failed { error, reason } => ( + McpServerStartupState::Failed, + Some(error), + reason.map(Into::into), + ), codex_protocol::protocol::McpStartupStatus::Cancelled => { - (McpServerStartupState::Cancelled, None) + (McpServerStartupState::Cancelled, None, None) } }; let notification = McpServerStatusUpdatedNotification { @@ -219,6 +221,7 @@ pub(crate) async fn apply_bespoke_event_handling( name: update.server, status, error, + failure_reason, }; outgoing .send_server_notification(ServerNotification::McpServerStatusUpdated(notification)) diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index 93dcb47ede..683a2bed18 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -945,6 +945,7 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result< name: "optional_broken".to_string(), status: McpServerStartupState::Starting, error: None, + failure_reason: None, } ); @@ -977,6 +978,7 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result< assert_eq!(failed.thread_id, Some(start_response.thread.id)); assert_eq!(failed.name, "optional_broken"); assert_eq!(failed.status, McpServerStartupState::Failed); + assert_eq!(failed.failure_reason, None); assert!( failed .error diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 586f9c184b..fd0b1a0e4f 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -573,7 +573,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> .map(|(name, cfg)| { let auth_status = auth_statuses .get(name.as_str()) - .map(|entry| entry.auth_status) + .map(|entry| McpAuthStatus::from(entry.auth_state)) .unwrap_or(McpAuthStatus::Unsupported); let transport = match &cfg.transport { McpServerTransportConfig::Stdio { @@ -657,7 +657,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> let status = format_mcp_status(cfg); let auth_status = auth_statuses .get(name.as_str()) - .map(|entry| entry.auth_status) + .map(|entry| McpAuthStatus::from(entry.auth_state)) .unwrap_or(McpAuthStatus::Unsupported) .to_string(); stdio_rows.push([ @@ -678,7 +678,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> let status = format_mcp_status(cfg); let auth_status = auth_statuses .get(name.as_str()) - .map(|entry| entry.auth_status) + .map(|entry| McpAuthStatus::from(entry.auth_state)) .unwrap_or(McpAuthStatus::Unsupported) .to_string(); let bearer_token_display = diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 0e81549556..5902bc9607 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -55,9 +55,12 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpStartupCompleteEvent; use codex_protocol::protocol::McpStartupFailure; +use codex_protocol::protocol::McpStartupFailureReason; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::McpAuthState; +use codex_rmcp_client::McpLoginRequirement; use rmcp::model::ElicitationCapability; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; @@ -228,12 +231,16 @@ impl McpConnectionManager { Ok(_) => McpStartupStatus::Ready, Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled, Err(error) => { + let reason = mcp_startup_failure_reason(auth_entry.as_ref(), error); let error_str = mcp_init_error_display( server_name.as_str(), auth_entry.as_ref(), error, ); - McpStartupStatus::Failed { error: error_str } + McpStartupStatus::Failed { + error: error_str, + reason, + } } }; @@ -266,7 +273,7 @@ impl McpConnectionManager { match outcome { Ok(_) => summary.ready.push(server_name), Err(StartupOutcomeError::Cancelled) => summary.cancelled.push(server_name), - Err(StartupOutcomeError::Failed { error }) => { + Err(StartupOutcomeError::Failed { error, .. }) => { summary.failed.push(McpStartupFailure { server: server_name, error, @@ -909,6 +916,28 @@ async fn emit_update( .await } +fn mcp_startup_failure_reason( + entry: Option<&McpAuthStatusEntry>, + error: &StartupOutcomeError, +) -> Option { + if !error.is_authentication_required() { + return None; + } + + match entry.map(|entry| entry.auth_state) { + Some(McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)) => { + Some(McpStartupFailureReason::ReauthenticationRequired) + } + Some( + McpAuthState::Unsupported + | McpAuthState::LoggedOut(McpLoginRequirement::Login) + | McpAuthState::BearerToken + | McpAuthState::OAuth, + ) + | None => None, + } +} + fn mcp_init_error_display( server_name: &str, entry: Option<&McpAuthStatusEntry>, @@ -955,20 +984,20 @@ fn mcp_init_error_display( fn startup_outcome_error_message(error: StartupOutcomeError) -> String { match error { StartupOutcomeError::Cancelled => "MCP startup cancelled".to_string(), - StartupOutcomeError::Failed { error } => error, + StartupOutcomeError::Failed { error, .. } => error, } } fn is_mcp_client_auth_required_error(error: &StartupOutcomeError) -> bool { match error { - StartupOutcomeError::Failed { error } => error.contains("Auth required"), + StartupOutcomeError::Failed { error, .. } => error.contains("Auth required"), _ => false, } } fn is_mcp_client_startup_timeout_error(error: &StartupOutcomeError) -> bool { match error { - StartupOutcomeError::Failed { error } => { + StartupOutcomeError::Failed { error, .. } => { error.contains("request timed out") || error.contains("timed out handshaking with MCP server") } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 05caf0489c..e11fcb5917 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -25,7 +25,6 @@ use codex_protocol::ToolName; 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; @@ -927,6 +926,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() let failed_client = futures::future::ready::>(Err( StartupOutcomeError::Failed { error: "startup failed".to_string(), + is_authentication_required: false, }, )) .boxed() @@ -1236,7 +1236,7 @@ fn mcp_init_error_display_prompts_for_github_pat() { oauth_resource: None, tools: HashMap::new(), }), - auth_status: McpAuthStatus::Unsupported, + auth_state: McpAuthState::Unsupported, }; let err: StartupOutcomeError = anyhow::anyhow!("OAuth is unsupported").into(); @@ -1263,6 +1263,50 @@ fn mcp_init_error_display_prompts_for_login_when_auth_required() { assert_eq!(expected, display); } +#[test] +fn mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure() { + for (auth_state, is_authentication_required, expected) in [ + ( + Some(McpAuthState::LoggedOut( + McpLoginRequirement::Reauthentication, + )), + true, + Some(McpStartupFailureReason::ReauthenticationRequired), + ), + ( + Some(McpAuthState::LoggedOut( + McpLoginRequirement::Reauthentication, + )), + false, + None, + ), + ( + Some(McpAuthState::LoggedOut(McpLoginRequirement::Login)), + true, + None, + ), + (Some(McpAuthState::Unsupported), true, None), + (Some(McpAuthState::BearerToken), true, None), + (Some(McpAuthState::OAuth), true, None), + (None, true, None), + ] { + let entry = auth_state.map(|auth_state| McpAuthStatusEntry { + config: None, + auth_state, + }); + let error = StartupOutcomeError::Failed { + error: "startup failed".to_string(), + is_authentication_required, + }; + + assert_eq!( + mcp_startup_failure_reason(entry.as_ref(), &error), + expected, + "auth_state={auth_state:?}, is_authentication_required={is_authentication_required}" + ); + } +} + #[test] fn mcp_init_error_display_reports_generic_errors() { let server_name = "custom"; @@ -1290,7 +1334,7 @@ fn mcp_init_error_display_reports_generic_errors() { oauth_resource: None, tools: HashMap::new(), }), - auth_status: McpAuthStatus::Unsupported, + auth_state: McpAuthState::Unsupported, }; let err: StartupOutcomeError = anyhow::anyhow!("boom").into(); diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 111f04163e..fcdaefe3a9 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -9,7 +9,7 @@ use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::HttpClient; use codex_login::CodexAuth; -use codex_protocol::protocol::McpAuthStatus; +use codex_rmcp_client::McpAuthState; use codex_rmcp_client::OAuthProviderError; use codex_rmcp_client::determine_streamable_http_auth_status; use codex_rmcp_client::determine_streamable_http_auth_status_with_http_client; @@ -54,7 +54,7 @@ pub struct ResolvedMcpOAuthScopes { #[derive(Debug, Clone)] pub struct McpAuthStatusEntry { pub config: Option, - pub auth_status: McpAuthStatus, + pub auth_state: McpAuthState, } pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport { @@ -208,7 +208,7 @@ where ) }); async move { - let auth_status = match config.as_ref() { + let auth_state = match config.as_ref() { Some(config) => { match compute_auth_status( &name, @@ -225,16 +225,13 @@ where warn!( "failed to determine auth status for MCP server `{name}`: {error:?}" ); - McpAuthStatus::Unsupported + McpAuthState::Unsupported } } } - None => McpAuthStatus::Unsupported, - }; - let entry = McpAuthStatusEntry { - config, - auth_status, + None => McpAuthState::Unsupported, }; + let entry = McpAuthStatusEntry { config, auth_state }; (name, entry) } }); @@ -249,17 +246,17 @@ async fn compute_auth_status( keyring_backend_kind: AuthKeyringBackendKind, has_runtime_auth: bool, runtime_context: &McpRuntimeContext, -) -> Result { +) -> Result { if !config.enabled { - return Ok(McpAuthStatus::Unsupported); + return Ok(McpAuthState::Unsupported); } if has_runtime_auth { - return Ok(McpAuthStatus::BearerToken); + return Ok(McpAuthState::BearerToken); } match &config.transport { - McpServerTransportConfig::Stdio { .. } => Ok(McpAuthStatus::Unsupported), + McpServerTransportConfig::Stdio { .. } => Ok(McpAuthState::Unsupported), McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 7532c6466b..3f24a276e4 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -567,7 +567,7 @@ fn auth_statuses_from_entries( ) -> HashMap { auth_status_entries .iter() - .map(|(name, entry)| (name.clone(), entry.auth_status)) + .map(|(name, entry)| (name.clone(), McpAuthStatus::from(entry.auth_state))) .collect::>() } diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 73ad33df7f..ddfb3ffdba 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -63,6 +63,7 @@ use rmcp::model::InitializeRequestParams; use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; +use rmcp::transport::auth::AuthError; use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::instrument; @@ -304,13 +305,39 @@ pub(crate) enum StartupOutcomeError { // We can't store the original error here because anyhow::Error doesn't implement // `Clone`. #[error("MCP startup failed: {error}")] - Failed { error: String }, + Failed { + error: String, + is_authentication_required: bool, + }, +} + +impl StartupOutcomeError { + pub(crate) fn is_authentication_required(&self) -> bool { + match self { + Self::Cancelled => false, + Self::Failed { + is_authentication_required, + .. + } => *is_authentication_required, + } + } } impl From for StartupOutcomeError { fn from(error: anyhow::Error) -> Self { + let is_authentication_required = error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|auth_error| { + matches!( + auth_error, + AuthError::AuthorizationRequired | AuthError::TokenExpired + ) + }) + }); Self::Failed { error: error.to_string(), + is_authentication_required, } } } @@ -765,6 +792,17 @@ mod tests { use pretty_assertions::assert_eq; use rmcp::model::JsonObject; use rmcp::model::Meta; + use rmcp::transport::auth::AuthError; + + #[test] + fn startup_outcome_error_identifies_authentication_required() { + let error = anyhow::Error::new(AuthError::AuthorizationRequired) + .context("failed to initialize MCP server"); + + let error = StartupOutcomeError::from(error); + + assert!(error.is_authentication_required()); + } #[test] fn mcp_initialize_advertises_openai_form_only_when_supported() { diff --git a/codex-rs/config/src/mcp_requirements_tests.rs b/codex-rs/config/src/mcp_requirements_tests.rs index b854873568..68b02316c4 100644 --- a/codex-rs/config/src/mcp_requirements_tests.rs +++ b/codex-rs/config/src/mcp_requirements_tests.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; fn stdio_server(command: &str, args: &[&str]) -> McpServerConfig { McpServerConfig { + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: command.to_string(), args: args.iter().map(ToString::to_string).collect(), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 42b01b16f3..5f3e75dd83 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -3621,10 +3621,22 @@ pub struct McpStartupUpdateEvent { pub enum McpStartupStatus { Starting, Ready, - Failed { error: String }, + Failed { + error: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional = nullable)] + reason: Option, + }, Cancelled, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum McpStartupFailureReason { + ReauthenticationRequired, +} + #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS, Default)] pub struct McpStartupCompleteEvent { pub ready: Vec, @@ -5749,6 +5761,7 @@ mod tests { server: "srv".to_string(), status: McpStartupStatus::Failed { error: "boom".to_string(), + reason: Some(McpStartupFailureReason::ReauthenticationRequired), }, }), }; @@ -5758,6 +5771,10 @@ mod tests { assert_eq!(value["msg"]["server"], "srv"); assert_eq!(value["msg"]["status"]["state"], "failed"); assert_eq!(value["msg"]["status"]["error"], "boom"); + assert_eq!( + value["msg"]["status"]["reason"], + "reauthentication_required" + ); Ok(()) } diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 5b8378e67a..32c244af8e 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -28,8 +28,33 @@ pub struct StreamableHttpOAuthDiscovery { pub scopes_supported: Option>, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum McpLoginRequirement { + Login, + Reauthentication, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum McpAuthState { + Unsupported, + LoggedOut(McpLoginRequirement), + BearerToken, + OAuth, +} + +impl From for McpAuthStatus { + fn from(value: McpAuthState) -> Self { + match value { + McpAuthState::Unsupported => Self::Unsupported, + McpAuthState::LoggedOut(_) => Self::NotLoggedIn, + McpAuthState::BearerToken => Self::BearerToken, + McpAuthState::OAuth => Self::OAuth, + } + } +} + enum AuthStatusCheck { - Complete(McpAuthStatus), + Complete(McpAuthState), Discover(HeaderMap), } @@ -42,7 +67,7 @@ pub async fn determine_streamable_http_auth_status( env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, -) -> Result { +) -> Result { let default_headers = match auth_status_before_discovery( server_name, url, @@ -75,7 +100,7 @@ pub async fn determine_streamable_http_auth_status_with_http_client( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, -) -> Result { +) -> Result { let default_headers = match auth_status_before_discovery( server_name, url, @@ -110,20 +135,22 @@ fn auth_status_before_discovery( keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { if bearer_token_env_var.is_some() { - return Ok(AuthStatusCheck::Complete(McpAuthStatus::BearerToken)); + return Ok(AuthStatusCheck::Complete(McpAuthState::BearerToken)); } let default_headers = build_default_headers(http_headers, env_http_headers)?; if default_headers.contains_key(AUTHORIZATION) { - return Ok(AuthStatusCheck::Complete(McpAuthStatus::BearerToken)); + return Ok(AuthStatusCheck::Complete(McpAuthState::BearerToken)); } match oauth_token_status(server_name, url, store_mode, keyring_backend_kind)? { StoredOAuthTokenStatus::Usable => { - return Ok(AuthStatusCheck::Complete(McpAuthStatus::OAuth)); + return Ok(AuthStatusCheck::Complete(McpAuthState::OAuth)); } StoredOAuthTokenStatus::AuthorizationRequired => { - return Ok(AuthStatusCheck::Complete(McpAuthStatus::NotLoggedIn)); + return Ok(AuthStatusCheck::Complete(McpAuthState::LoggedOut( + McpLoginRequirement::Reauthentication, + ))); } StoredOAuthTokenStatus::Missing => {} } @@ -135,15 +162,15 @@ fn determine_auth_status_from_discovery( server_name: &str, url: &str, discovery: Result>, -) -> Result { +) -> Result { match discovery { - Ok(Some(_)) => Ok(McpAuthStatus::NotLoggedIn), - Ok(None) => Ok(McpAuthStatus::Unsupported), + Ok(Some(_)) => Ok(McpAuthState::LoggedOut(McpLoginRequirement::Login)), + Ok(None) => Ok(McpAuthState::Unsupported), Err(error) => { debug!( "failed to detect OAuth support for MCP server `{server_name}` at {url}: {error:?}" ); - Ok(McpAuthStatus::Unsupported) + Ok(McpAuthState::Unsupported) } } } @@ -336,7 +363,7 @@ mod tests { .await .expect("status should compute"); - assert_eq!(status, McpAuthStatus::BearerToken); + assert_eq!(status, McpAuthState::BearerToken); } #[tokio::test] @@ -358,7 +385,7 @@ mod tests { .await .expect("status should compute"); - assert_eq!(status, McpAuthStatus::BearerToken); + assert_eq!(status, McpAuthState::BearerToken); } #[tokio::test] diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1041fdafb..5be97a56a6 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -12,6 +12,8 @@ mod rmcp_client; mod stdio_server_launcher; mod utils; +pub use auth_status::McpAuthState; +pub use auth_status::McpLoginRequirement; pub use auth_status::StreamableHttpOAuthDiscovery; pub use auth_status::determine_streamable_http_auth_status; pub use auth_status::determine_streamable_http_auth_status_with_http_client; diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index 866daa67a4..bc5a701ed0 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -7,7 +7,8 @@ use std::time::UNIX_EPOCH; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; -use codex_rmcp_client::McpAuthStatus; +use codex_rmcp_client::McpAuthState; +use codex_rmcp_client::McpLoginRequirement; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; @@ -146,6 +147,21 @@ async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by reports_auth_status_for_persisted_credentials"] async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { + let first_login_server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", first_login_server.uri()), + "token_endpoint": format!("{}/oauth/token", first_login_server.uri()), + }))) + .expect(1) + .mount(&first_login_server) + .await; + + let status = auth_status(&format!("{}/mcp", first_login_server.uri())).await?; + assert_eq!(status, McpAuthState::LoggedOut(McpLoginRequirement::Login)); + first_login_server.verify().await; + let response = OAuthTokenResponse::new( AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), BasicTokenType::Bearer, @@ -166,7 +182,10 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { )?; let status = auth_status(UNREFRESHABLE_SERVER_URL).await?; - assert_eq!(status, McpAuthStatus::NotLoggedIn); + assert_eq!( + status, + McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication) + ); let response = OAuthTokenResponse::new( AccessToken::new("unexpired-access-token".to_string()), @@ -192,7 +211,7 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { )?; let status = auth_status(UNEXPIRED_SERVER_URL).await?; - assert_eq!(status, McpAuthStatus::OAuth); + assert_eq!(status, McpAuthState::OAuth); let mut response = OAuthTokenResponse::new( AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), @@ -215,11 +234,11 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { )?; let status = auth_status(REFRESHABLE_SERVER_URL).await?; - assert_eq!(status, McpAuthStatus::OAuth); + assert_eq!(status, McpAuthState::OAuth); Ok(()) } -async fn auth_status(server_url: &str) -> anyhow::Result { +async fn auth_status(server_url: &str) -> anyhow::Result { determine_streamable_http_auth_status( SERVER_NAME, server_url, diff --git a/codex-rs/tui/src/app/app_server_event_targets.rs b/codex-rs/tui/src/app/app_server_event_targets.rs index 212c45b4be..6c55850e4d 100644 --- a/codex-rs/tui/src/app/app_server_event_targets.rs +++ b/codex-rs/tui/src/app/app_server_event_targets.rs @@ -285,6 +285,7 @@ mod tests { name: "sentry".to_string(), status: McpServerStartupState::Failed, error: Some("sentry is not logged in".to_string()), + failure_reason: None, }); let target = server_notification_thread_target(¬ification); @@ -300,6 +301,7 @@ mod tests { name: "sentry".to_string(), status: McpServerStartupState::Failed, error: Some("sentry is not logged in".to_string()), + failure_reason: None, }); let target = server_notification_thread_target(¬ification); diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 21aca66bcd..29e29541a7 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -3496,6 +3496,7 @@ async fn primary_thread_ignores_child_mcp_startup_notifications() { name: "sentry".to_string(), status: McpServerStartupState::Failed, error: Some("sentry is not logged in".to_string()), + failure_reason: None, }), ), ) @@ -3567,6 +3568,7 @@ async fn app_scoped_mcp_startup_notifications_do_not_render_in_active_thread() { name: "sentry".to_string(), status: McpServerStartupState::Failed, error: Some("sentry is not logged in".to_string()), + failure_reason: None, }), ), ) @@ -3631,6 +3633,7 @@ async fn active_side_thread_renders_live_mcp_startup_notifications() { status, error: matches!(status, McpServerStartupState::Failed) .then(|| "sentry is not logged in".to_string()), + failure_reason: None, }), ), ) diff --git a/codex-rs/tui/src/app/thread_events.rs b/codex-rs/tui/src/app/thread_events.rs index 00b7fcdedb..8d41f1c052 100644 --- a/codex-rs/tui/src/app/thread_events.rs +++ b/codex-rs/tui/src/app/thread_events.rs @@ -619,6 +619,7 @@ mod tests { name: "sentry".to_string(), status: codex_app_server_protocol::McpServerStartupState::Failed, error: Some("sentry is not logged in".to_string()), + failure_reason: None, }, ); let mut store = ThreadEventStore::new(/*capacity*/ 8); diff --git a/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs b/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs index 60efbe1353..e2992585a4 100644 --- a/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs +++ b/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs @@ -8,6 +8,7 @@ fn notify_mcp_status(chat: &mut ChatWidget, name: &str, status: McpServerStartup name: name.to_string(), status, error: None, + failure_reason: None, }), /*replay_kind*/ None, ); @@ -20,6 +21,7 @@ fn notify_mcp_status_error(chat: &mut ChatWidget, name: &str, error: &str) { name: name.to_string(), status: McpServerStartupState::Failed, error: Some(error.to_string()), + failure_reason: None, }), /*replay_kind*/ None, ); @@ -51,6 +53,7 @@ async fn mcp_startup_ignores_status_for_other_thread() { status, error: matches!(status, McpServerStartupState::Failed) .then(|| "sentry is not logged in".to_string()), + failure_reason: None, }), /*replay_kind*/ None, ); From e2746fd7e90d4f6095da62d70f50194cc21d2d3e Mon Sep 17 00:00:00 2001 From: chiam-oai Date: Thu, 25 Jun 2026 15:30:26 -0700 Subject: [PATCH 25/33] Recognize Work web and mobile thread originators (#29988) ## Summary - recognize `codex_work_web` and `codex_work_mobile` as supported `thread/start.serviceName` values - use the recognized value as the thread-scoped originator, with the same persistence and request propagation added for `codex_work_desktop` - cover precedence over persisted and inherited originators This is the Codex consumer for the service names introduced by [openai/openai#1073178](https://github.com/openai/openai/pull/1073178). ## Rollout / Compatibility The producer is ChatGPT's app-server integration in openai/openai#1073178. This PR is the Codex app-server consumer that converts those service names into the outgoing per-thread `originator`. Until this change is deployed, the new service names are ignored and Codex continues using its fallback originator. Deploy this mapper and the matching codex-backend compatibility change in [openai/openai#1073594](https://github.com/openai/openai/pull/1073594) while the existing Flora egress overwrite remains in place. Remove that overwrite in [openai/openai#1073197](https://github.com/openai/openai/pull/1073197) only after both consumers are deployed. ## Validation - `just test -p codex-core effective_originator_prefers_thread_scoped_sources_before_env_originator` - `just fix -p codex-core` - `just fmt` --- codex-rs/core/src/thread_manager.rs | 6 ++++-- codex-rs/core/src/thread_manager_tests.rs | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 358eae3527..b606bb3f75 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -196,8 +196,10 @@ pub struct StartThreadOptions { fn originator_from_service_name(service_name: Option<&str>) -> Option { let service_name = service_name?.trim(); - if service_name.eq_ignore_ascii_case("codex_work_desktop") { - return Some("codex_work_desktop".to_string()); + for originator in ["codex_work_desktop", "codex_work_web", "codex_work_mobile"] { + if service_name.eq_ignore_ascii_case(originator) { + return Some(originator.to_string()); + } } None } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 34df66cb67..9bccf8b6c1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -122,6 +122,18 @@ fn effective_originator_prefers_thread_scoped_sources_before_env_originator() { Some("inherited_originator"), "codex_work_desktop", ), + ( + Some("codex_work_web"), + Some("persisted_originator"), + Some("inherited_originator"), + "codex_work_web", + ), + ( + Some("codex_work_mobile"), + Some("persisted_originator"), + Some("inherited_originator"), + "codex_work_mobile", + ), ( None, Some("persisted_originator"), From 6c21297bba7da6ef9c14362406eec59aba5cbfce Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 15:33:58 -0700 Subject: [PATCH 26/33] [codex] add code-mode host failure supervision hooks (#30110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why A process host should be discarded and rebuilt after critical actor or V8 failure, while the existing in-process production path must keep its current cell-error semantics. This change establishes that failure boundary without adding the host process or remote client. ## What changed - add optional task-failure supervision to the transport-neutral code-mode session runtime - report Tokio cell-actor failures and V8 runtime-thread panics to a host-provided fail-stop handler - preserve the existing handler-less in-process behavior - make host-owned cell ID allocation fail before numeric wraparound ## Follow-up The V8 panic signal surfaced here should also be consumed by the `InProcessCodeModeSession` manager in a future change so it can fail the affected cell. This PR intentionally leaves the handler-less in-process behavior unchanged while putting the required panic tracking in place. ## Stack This is **2 of 4** in the process-owned code-mode session stack. - #30108 is merged into `main` - The next PR targets this branch ## Validation - `just test -p codex-code-mode` — 53 passed - `just argument-comment-lint -p codex-code-mode` - `just fix -p codex-code-mode` --- codex-rs/Cargo.lock | 1 + codex-rs/code-mode/Cargo.toml | 1 + .../code-mode/src/cell_actor/callbacks.rs | 75 ++++++++-- .../src/cell_actor/callbacks_tests.rs | 132 ++++++++++++++++++ codex-rs/code-mode/src/cell_actor/mod.rs | 35 ++++- codex-rs/code-mode/src/cell_actor/tests.rs | 62 ++++++++ codex-rs/code-mode/src/lib.rs | 2 + codex-rs/code-mode/src/runtime/mod.rs | 65 ++++++++- codex-rs/code-mode/src/service.rs | 12 ++ codex-rs/code-mode/src/session_runtime/mod.rs | 40 ++++-- .../code-mode/src/session_runtime/tests.rs | 78 +++++++++++ .../code-mode/src/session_runtime/types.rs | 4 + 12 files changed, 482 insertions(+), 25 deletions(-) create mode 100644 codex-rs/code-mode/src/cell_actor/callbacks_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3837df62ea..b147022825 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2503,6 +2503,7 @@ dependencies = [ "codex-code-mode-protocol", "codex-protocol", "deno_core_icudata", + "futures", "pretty_assertions", "serde_json", "tokio", diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 09f1d7380d..6b91e64fe2 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -19,6 +19,7 @@ workspace = true codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } +futures = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } diff --git a/codex-rs/code-mode/src/cell_actor/callbacks.rs b/codex-rs/code-mode/src/cell_actor/callbacks.rs index 9db66ed3a7..08f7cbac99 100644 --- a/codex-rs/code-mode/src/cell_actor/callbacks.rs +++ b/codex-rs/code-mode/src/cell_actor/callbacks.rs @@ -1,11 +1,14 @@ +use std::panic::AssertUnwindSafe; use std::sync::Arc; +use futures::FutureExt; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; use super::CellHost; use super::CellToolCall; +use crate::TaskFailureHandler; use crate::runtime::RuntimeCommand; #[derive(Clone, Copy)] @@ -20,10 +23,20 @@ pub(super) fn spawn_notification( call_id: String, text: String, cancellation_token: CancellationToken, + task_failure_handler: Option, ) { tasks.spawn(async move { - if let Err(err) = host.notify(call_id, text, cancellation_token).await { - warn!("failed to deliver code mode notification: {err}"); + let callback = + AssertUnwindSafe(async move { host.notify(call_id, text, cancellation_token).await }) + .catch_unwind() + .await; + match callback { + Ok(Ok(())) => {} + Ok(Err(err)) => warn!("failed to deliver code mode notification: {err}"), + Err(_) => report_task_failure( + task_failure_handler.as_ref(), + "code mode notification task panicked".to_string(), + ), } }); } @@ -34,14 +47,32 @@ pub(super) fn spawn_tool( invocation: CellToolCall, runtime_tx: std::sync::mpsc::Sender, cancellation_token: CancellationToken, + task_failure_handler: Option, ) { tasks.spawn(async move { let id = invocation.id.clone(); - let command = match host.invoke_tool(invocation, cancellation_token).await { - Ok(result) => RuntimeCommand::ToolResponse { id, result }, - Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + let callback = + AssertUnwindSafe(async move { host.invoke_tool(invocation, cancellation_token).await }) + .catch_unwind() + .await; + let (command, failure_reason) = match callback { + Ok(Ok(result)) => (RuntimeCommand::ToolResponse { id, result }, None), + Ok(Err(error_text)) => (RuntimeCommand::ToolError { id, error_text }, None), + Err(_) => { + let failure_reason = "code mode tool task panicked".to_string(); + ( + RuntimeCommand::ToolError { + id, + error_text: failure_reason.clone(), + }, + Some(failure_reason), + ) + } }; let _ = runtime_tx.send(command); + if let Some(failure_reason) = failure_reason { + report_task_failure(task_failure_handler.as_ref(), failure_reason); + } }); } @@ -50,28 +81,48 @@ pub(super) async fn finish_callbacks( notification_tasks: &mut JoinSet<()>, tool_tasks: &mut JoinSet<()>, completion: CallbackCompletion, + task_failure_handler: Option<&TaskFailureHandler>, ) { if matches!(completion, CallbackCompletion::Cancel) { cancellation_token.cancel(); } - drain_tasks(notification_tasks, "notification").await; + drain_tasks(notification_tasks, "notification", task_failure_handler).await; cancellation_token.cancel(); - drain_tasks(tool_tasks, "tool").await; + drain_tasks(tool_tasks, "tool", task_failure_handler).await; } -pub(super) fn log_task_result( +pub(super) fn report_task_result( task_result: Option>, description: &str, + task_failure_handler: Option<&TaskFailureHandler>, ) { if let Some(Err(err)) = task_result && !err.is_cancelled() { - warn!("code mode {description} task failed: {err}"); + report_task_failure( + task_failure_handler, + format!("code mode {description} task failed: {err}"), + ); } } -async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) { - while let Some(result) = tasks.join_next().await { - log_task_result(Some(result), description); +fn report_task_failure(task_failure_handler: Option<&TaskFailureHandler>, failure_reason: String) { + warn!("{failure_reason}"); + if let Some(task_failure_handler) = task_failure_handler { + task_failure_handler(failure_reason); } } + +async fn drain_tasks( + tasks: &mut JoinSet<()>, + description: &str, + task_failure_handler: Option<&TaskFailureHandler>, +) { + while let Some(result) = tasks.join_next().await { + report_task_result(Some(result), description, task_failure_handler); + } +} + +#[cfg(test)] +#[path = "callbacks_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode/src/cell_actor/callbacks_tests.rs b/codex-rs/code-mode/src/cell_actor/callbacks_tests.rs new file mode 100644 index 0000000000..585f7808a8 --- /dev/null +++ b/codex-rs/code-mode/src/cell_actor/callbacks_tests.rs @@ -0,0 +1,132 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::mpsc as std_mpsc; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use serde_json::Value as JsonValue; +use tokio::sync::mpsc; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; + +use super::*; +use crate::cell_actor::CellState; +use crate::cell_actor::CompletionCommit; +use crate::runtime::RuntimeCommand; +use crate::session_runtime::CellEvent; +use crate::session_runtime::ToolKind; +use crate::session_runtime::ToolName; + +struct PanickingCallbackHost; + +impl CellHost for PanickingCallbackHost { + async fn invoke_tool( + &self, + _invocation: CellToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + panic!("tool callback panic probe"); + } + + async fn notify( + &self, + _call_id: String, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + panic!("notification callback panic probe"); + } + + async fn commit_completion( + &self, + _stored_value_writes: HashMap, + _event: CellEvent, + _pending_initial_yield_items: Option>, + _cell_state: Arc, + ) -> CompletionCommit { + panic!("unexpected completion commit"); + } + + async fn closed(&self) {} +} + +#[tokio::test] +async fn tool_callback_panic_rejects_the_js_promise_and_reports_failure() { + let mut tasks = JoinSet::new(); + let (runtime_tx, runtime_rx) = std_mpsc::channel(); + let (failure_tx, mut failure_rx) = mpsc::unbounded_channel(); + spawn_tool( + &mut tasks, + Arc::new(PanickingCallbackHost), + CellToolCall { + id: "tool-1".to_string(), + name: ToolName { + name: "panic".to_string(), + namespace: None, + }, + kind: ToolKind::Function, + input: None, + }, + runtime_tx, + CancellationToken::new(), + Some(Arc::new(move |reason| { + let _ = failure_tx.send(reason); + })), + ); + + tasks + .join_next() + .await + .expect("tool callback task") + .expect("tool callback wrapper"); + let command = runtime_rx + .recv_timeout(Duration::from_secs(1)) + .expect("tool error command"); + let RuntimeCommand::ToolError { id, error_text } = command else { + panic!("expected a tool error command"); + }; + assert_eq!(id, "tool-1"); + assert_eq!(error_text, "code mode tool task panicked"); + assert_eq!(failure_rx.recv().await, Some(error_text)); +} + +#[tokio::test] +async fn notification_callback_panic_reports_failure() { + let mut tasks = JoinSet::new(); + let (failure_tx, mut failure_rx) = mpsc::unbounded_channel(); + spawn_notification( + &mut tasks, + Arc::new(PanickingCallbackHost), + "notify-1".to_string(), + "hello".to_string(), + CancellationToken::new(), + Some(Arc::new(move |reason| { + let _ = failure_tx.send(reason); + })), + ); + + tasks + .join_next() + .await + .expect("notification callback task") + .expect("notification callback wrapper"); + let failure_reason = failure_rx.recv().await.expect("notification failure"); + assert_eq!(failure_reason, "code mode notification task panicked"); +} + +#[tokio::test] +async fn callback_wrapper_join_error_reports_failure() { + let task_result = tokio::spawn(async { + panic!("callback wrapper panic probe"); + }) + .await; + let (failure_tx, mut failure_rx) = mpsc::unbounded_channel(); + let task_failure_handler: TaskFailureHandler = Arc::new(move |reason| { + let _ = failure_tx.send(reason); + }); + + report_task_result(Some(task_result), "tool", Some(&task_failure_handler)); + + let failure_reason = failure_rx.recv().await.expect("wrapper failure"); + assert!(failure_reason.contains("code mode tool task failed")); +} diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 24eca22517..7533f4b0cc 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -14,7 +14,7 @@ use tokio_util::sync::CancellationToken; use self::callbacks::CallbackCompletion; use self::callbacks::finish_callbacks; -use self::callbacks::log_task_result; +use self::callbacks::report_task_result; use self::callbacks::spawn_notification; use self::callbacks::spawn_tool; use self::conversions::cell_tool_kind; @@ -30,6 +30,7 @@ pub(crate) use self::types::CellToolCall; pub(crate) use self::types::CompletionCommit; use self::types::CompletionDelivery; use self::types::ObservationDelivery; +use crate::TaskFailureHandler; use crate::runtime::PendingRuntimeMode; use crate::runtime::RuntimeCommand; use crate::runtime::RuntimeControlCommand; @@ -50,6 +51,7 @@ impl CellActor { host: Arc, initial_observe_mode: ObserveMode, cell_state: Arc, + task_failure_handler: Option, ) -> Result< ( CellHandle, @@ -66,6 +68,7 @@ impl CellActor { runtime_request(request), event_tx, PendingRuntimeMode::PauseUntilResumed, + task_failure_handler.clone(), )?; let handle = CellHandle::new(command_tx, Arc::clone(&cell_state)); let task = run_cell( @@ -82,6 +85,7 @@ impl CellActor { mode: initial_observe_mode, response_tx: initial_response_tx, }, + task_failure_handler, ); let initial_response = Box::pin(async move { initial_response_rx.await.unwrap_or(Err(CellError::Closed)) }); @@ -107,6 +111,7 @@ async fn run_cell( mut event_rx: mpsc::UnboundedReceiver, command_rx: mpsc::UnboundedReceiver, initial_observer: Observer, + task_failure_handler: Option, ) { let CellContext { runtime_tx, @@ -123,6 +128,7 @@ async fn run_cell( let mut termination = false; let mut runtime_closed = false; let mut runtime_paused = false; + let mut runtime_failure_reported = false; let mut yield_timer: Option>> = None; let mut notification_tasks = JoinSet::new(); let mut tool_tasks = JoinSet::new(); @@ -149,6 +155,7 @@ async fn run_cell( &mut notification_tasks, &mut tool_tasks, CallbackCompletion::Cancel, + task_failure_handler.as_ref(), ).await; finish_termination( &cell_state, @@ -259,6 +266,7 @@ async fn run_cell( &mut notification_tasks, &mut tool_tasks, CallbackCompletion::Cancel, + task_failure_handler.as_ref(), ).await; finish_termination( &cell_state, @@ -269,11 +277,20 @@ async fn run_cell( ); break; } + if !runtime_failure_reported + && let Some(task_failure_handler) = &task_failure_handler + { + runtime_failure_reported = true; + task_failure_handler( + "code-mode V8 runtime thread ended unexpectedly".to_string(), + ); + } finish_callbacks( &callback_cancellation_token, &mut notification_tasks, &mut tool_tasks, CallbackCompletion::DrainNotifications, + task_failure_handler.as_ref(), ) .await; let event = CellEvent::Completed { @@ -376,6 +393,7 @@ async fn run_cell( call_id, text, callback_cancellation_token.child_token(), + task_failure_handler.clone(), ); } RuntimeEvent::ToolCall { id, name, kind, input } => { @@ -394,6 +412,7 @@ async fn run_cell( }, runtime_tx.clone(), callback_cancellation_token.child_token(), + task_failure_handler.clone(), ); } RuntimeEvent::Result { stored_value_writes, error_text } => { @@ -405,6 +424,7 @@ async fn run_cell( &mut notification_tasks, &mut tool_tasks, CallbackCompletion::Cancel, + task_failure_handler.as_ref(), ).await; finish_termination( &cell_state, @@ -420,6 +440,7 @@ async fn run_cell( &mut notification_tasks, &mut tool_tasks, CallbackCompletion::DrainNotifications, + task_failure_handler.as_ref(), ) .await; let event = CellEvent::Completed { @@ -455,13 +476,20 @@ async fn run_cell( } } } + RuntimeEvent::ThreadPanicked => { + runtime_failure_reported = true; + } } } task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => { - log_task_result(task_result, "notification"); + report_task_result( + task_result, + "notification", + task_failure_handler.as_ref(), + ); } task_result = tool_tasks.join_next(), if !tool_tasks.is_empty() => { - log_task_result(task_result, "tool"); + report_task_result(task_result, "tool", task_failure_handler.as_ref()); } } } @@ -479,6 +507,7 @@ async fn run_cell( &mut notification_tasks, &mut tool_tasks, CallbackCompletion::Cancel, + task_failure_handler.as_ref(), ) .await; host.closed().await; diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index a4d93fc3e6..3612daff91 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -102,6 +102,18 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne fn spawn_cell_actor_harness_with_host( initial_observe_mode: ObserveMode, host: Arc, +) -> CellActorHarness { + spawn_cell_actor_harness_with_host_and_failure_handler( + initial_observe_mode, + host, + /*task_failure_handler*/ None, + ) +} + +fn spawn_cell_actor_harness_with_host_and_failure_handler( + initial_observe_mode: ObserveMode, + host: Arc, + task_failure_handler: Option, ) -> CellActorHarness { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); @@ -118,6 +130,7 @@ fn spawn_cell_actor_harness_with_host( }, runtime_event_tx, PendingRuntimeMode::PauseUntilResumed, + /*task_failure_handler*/ None, ) .unwrap(); let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel(); @@ -137,6 +150,7 @@ fn spawn_cell_actor_harness_with_host( mode: initial_observe_mode, response_tx: initial_event_tx, }, + task_failure_handler, )); CellActorHarness { @@ -149,6 +163,54 @@ fn spawn_cell_actor_harness_with_host( } } +#[tokio::test] +async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() { + let (failure_tx, mut failure_rx) = mpsc::unbounded_channel(); + let harness = spawn_cell_actor_harness_with_host_and_failure_handler( + ObserveMode::YieldAfter(Duration::from_secs(60)), + Arc::new(TestHost), + Some(Arc::new(move |reason| { + let _ = failure_tx.send(reason); + })), + ); + drop(harness.event_tx); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), failure_rx.recv()) + .await + .expect("runtime failure timeout") + .expect("runtime failure"), + "code-mode V8 runtime thread ended unexpectedly" + ); + assert!( + harness + .initial_event_rx + .await + .expect("initial event") + .is_ok() + ); + harness.task.await.expect("cell task"); +} + +#[tokio::test] +async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() { + let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60))); + harness + .event_tx + .send(RuntimeEvent::ThreadPanicked) + .expect("runtime panic event"); + drop(harness.event_tx); + + assert_eq!( + harness.initial_event_rx.await.expect("initial event"), + Ok(CellEvent::Completed { + content_items: Vec::new(), + error_text: Some("exec runtime ended unexpectedly".to_string()), + }) + ); + harness.task.await.expect("cell task"); +} + async fn wait_for_notification(host: &RecordingHost) { tokio::time::timeout(Duration::from_secs(1), async { while !host.notified.load(Ordering::Acquire) { diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 089bb6f78d..60c02e4fae 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -3,6 +3,8 @@ mod runtime; mod service; mod session_runtime; +pub(crate) type TaskFailureHandler = std::sync::Arc; + pub use codex_code_mode_protocol::*; pub use service::InProcessCodeModeSession; pub use service::InProcessCodeModeSessionProvider; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 42af5c0b95..0b553123f7 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -5,6 +5,8 @@ mod timers; mod value; use std::collections::HashMap; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; use std::sync::OnceLock; use std::sync::mpsc as std_mpsc; use std::thread; @@ -18,6 +20,8 @@ use codex_protocol::ToolName; use serde_json::Value as JsonValue; use tokio::sync::mpsc; +use crate::TaskFailureHandler; + const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; #[derive(Debug)] @@ -63,6 +67,7 @@ pub(crate) enum RuntimeEvent { stored_value_writes: HashMap, error_text: Option, }, + ThreadPanicked, } pub(crate) fn spawn_runtime( @@ -70,6 +75,7 @@ pub(crate) fn spawn_runtime( request: ExecuteRequest, event_tx: mpsc::UnboundedSender, pending_mode: PendingRuntimeMode, + task_failure_handler: Option, ) -> Result< ( std_mpsc::Sender, @@ -96,7 +102,7 @@ pub(crate) fn spawn_runtime( stored_values, }; - thread::spawn(move || { + spawn_supervised_runtime_thread(event_tx.clone(), task_failure_handler, move || { run_runtime( config, event_tx, @@ -114,6 +120,21 @@ pub(crate) fn spawn_runtime( Ok((command_tx, control_tx, isolate_handle)) } +fn spawn_supervised_runtime_thread( + event_tx: mpsc::UnboundedSender, + task_failure_handler: Option, + runtime: impl FnOnce() + Send + 'static, +) { + thread::spawn(move || { + if catch_unwind(AssertUnwindSafe(runtime)).is_err() { + if let Some(task_failure_handler) = task_failure_handler { + task_failure_handler("code-mode V8 runtime thread panicked".to_string()); + } + let _ = event_tx.send(RuntimeEvent::ThreadPanicked); + } + }); +} + #[derive(Clone)] struct RuntimeConfig { tool_call_id: String, @@ -336,6 +357,7 @@ mod tests { use super::RuntimeControlCommand; use super::RuntimeEvent; use super::spawn_runtime; + use super::spawn_supervised_runtime_thread; use crate::FunctionCallOutputContentItem; fn execute_request(source: &str) -> ExecuteRequest { @@ -348,6 +370,45 @@ mod tests { } } + #[tokio::test] + async fn runtime_thread_panic_before_initialization_is_reported_directly() { + let (event_tx, event_rx) = mpsc::unbounded_channel(); + drop(event_rx); + let (failure_tx, mut failure_rx) = mpsc::unbounded_channel(); + spawn_supervised_runtime_thread( + event_tx, + Some(std::sync::Arc::new(move |reason| { + let _ = failure_tx.send(reason); + })), + || panic!("runtime thread panic probe"), + ); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), failure_rx.recv()) + .await + .expect("runtime failure timeout") + .expect("runtime failure"), + "code-mode V8 runtime thread panicked" + ); + } + + #[tokio::test] + async fn runtime_thread_panic_is_forwarded_without_owner_supervision() { + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + spawn_supervised_runtime_thread( + event_tx, + /*task_failure_handler*/ None, + || panic!("runtime thread panic probe"), + ); + + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("runtime panic event timeout"), + Some(RuntimeEvent::ThreadPanicked) + )); + } + #[tokio::test] async fn terminate_execution_stops_cpu_bound_module() { let (event_tx, mut event_rx) = mpsc::unbounded_channel(); @@ -356,6 +417,7 @@ mod tests { execute_request("while (true) {}"), event_tx, PendingRuntimeMode::Continue, + /*task_failure_handler*/ None, ) .unwrap(); @@ -398,6 +460,7 @@ await new Promise(() => {}); ), event_tx, PendingRuntimeMode::PauseUntilResumed, + /*task_failure_handler*/ None, ) .unwrap(); diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 57400001f0..8c5b94bb75 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -87,6 +87,18 @@ impl InProcessCodeModeSession { } } + pub fn with_delegate_and_task_failure_handler( + delegate: Arc, + task_failure_handler: Arc, + ) -> Self { + Self { + runtime: SessionRuntime::new_with_task_failure_handler( + Arc::new(ProtocolDelegate { delegate }), + Some(task_failure_handler), + ), + } + } + pub async fn execute(&self, request: ExecuteRequest) -> Result { let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); let started = self diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 5688c04303..cc014230f9 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -24,6 +24,7 @@ pub(crate) use self::types::SessionRuntimeDelegate; pub(crate) use self::types::ToolDefinition; pub(crate) use self::types::ToolKind; pub(crate) use self::types::ToolName; +use crate::TaskFailureHandler; use crate::cell_actor::CellActor; use crate::cell_actor::CellError; use crate::cell_actor::CellEventFuture; @@ -46,11 +47,19 @@ struct Inner { cell_tasks: TaskTracker, shutdown_token: CancellationToken, delegate: Arc, + task_failure_handler: Option, next_cell_id: AtomicU64, } impl SessionRuntime { pub(crate) fn new(delegate: Arc) -> Self { + Self::new_with_task_failure_handler(delegate, /*task_failure_handler*/ None) + } + + pub(crate) fn new_with_task_failure_handler( + delegate: Arc, + task_failure_handler: Option, + ) -> Self { Self { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), @@ -58,6 +67,7 @@ impl SessionRuntime { cell_tasks: TaskTracker::new(), shutdown_token: CancellationToken::new(), delegate, + task_failure_handler, next_cell_id: AtomicU64::new(1), }), } @@ -71,7 +81,7 @@ impl SessionRuntime { if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } - let cell_id = self.allocate_cell_id(); + let cell_id = self.allocate_cell_id()?; let initial_event = self .start_cell(cell_id.clone(), request, initial_observe_mode) .await?; @@ -133,13 +143,14 @@ impl SessionRuntime { Ok(()) } - fn allocate_cell_id(&self) -> CellId { - CellId::new( - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string(), - ) + fn allocate_cell_id(&self) -> Result { + self.inner + .next_cell_id + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |next_cell_id| { + next_cell_id.checked_add(1) + }) + .map(|cell_id| CellId::new(cell_id.to_string())) + .map_err(|_| Error::CellIdSpaceExhausted) } async fn start_cell( @@ -167,10 +178,21 @@ impl SessionRuntime { host, initial_observe_mode, cell_state, + self.inner.task_failure_handler.clone(), ) .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), handle); - self.inner.cell_tasks.spawn(task); + let task = self.inner.cell_tasks.spawn(task); + if let Some(task_failure_handler) = self.inner.task_failure_handler.clone() { + let failed_cell_id = cell_id.clone(); + let _failure_watcher = self.inner.cell_tasks.spawn(async move { + if let Err(err) = task.await { + task_failure_handler(format!( + "code-mode cell {failed_cell_id} task failed: {err}" + )); + } + }); + } drop(cells); Ok(map_actor_event(cell_id, initial_event)) } diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index 969f06a22d..1d90179bb3 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -15,6 +15,8 @@ use crate::cell_actor::CompletionCommit; struct RecordingDelegate; +struct PanickingClosedDelegate; + impl SessionRuntimeDelegate for RecordingDelegate { async fn invoke_tool( &self, @@ -37,6 +39,62 @@ impl SessionRuntimeDelegate for RecordingDelegate { fn cell_closed(&self, _cell_id: &CellId) {} } +impl SessionRuntimeDelegate for PanickingClosedDelegate { + async fn invoke_tool( + &self, + _invocation: NestedToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(JsonValue::Null) + } + + async fn notify( + &self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + Ok(()) + } + + fn cell_closed(&self, _cell_id: &CellId) { + panic!("cell close panic probe"); + } +} + +#[tokio::test] +async fn reports_cell_actor_panics_to_the_owner() { + let (failure_tx, mut failure_rx) = tokio::sync::mpsc::unbounded_channel(); + let runtime = SessionRuntime::new_with_task_failure_handler( + Arc::new(PanickingClosedDelegate), + Some(Arc::new(move |reason| { + let _ = failure_tx.send(reason); + })), + ); + let started = runtime + .execute( + execute_request(r#"text("done");"#), + ObserveMode::YieldAfter(Duration::from_secs(1)), + ) + .await + .expect("start cell"); + assert_eq!( + started.initial_event().await, + Ok(CellEvent::Completed { + content_items: vec![OutputItem::Text { + text: "done".to_string(), + }], + error_text: None, + }) + ); + runtime.shutdown().await.expect("shutdown runtime"); + let failure = failure_rx + .try_recv() + .expect("shutdown should wait for the cell failure watcher"); + assert!(failure.contains("code-mode cell 1 task failed")); +} + #[tokio::test] async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_load_it() { let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); @@ -118,6 +176,26 @@ fn execute_request(source: &str) -> CreateCellRequest { } } +#[tokio::test] +async fn cell_id_allocation_fails_before_wrapping() { + let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + runtime + .inner + .next_cell_id + .store(u64::MAX, Ordering::Relaxed); + + assert_eq!( + runtime + .execute( + execute_request(r#"text("unreachable");"#), + ObserveMode::YieldAfter(Duration::from_secs(1)), + ) + .await + .err(), + Some(Error::CellIdSpaceExhausted) + ); +} + #[tokio::test] #[expect( clippy::await_holding_invalid_type, diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index b04bd0f236..a33ffa1f1f 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -138,6 +138,7 @@ pub(crate) trait SessionRuntimeDelegate: Send + Sync + 'static { #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) enum Error { ShuttingDown, + CellIdSpaceExhausted, DuplicateCell(CellId), MissingCell(CellId), BusyObserver(CellId), @@ -150,6 +151,9 @@ impl fmt::Display for Error { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::ShuttingDown => formatter.write_str("code mode session is shutting down"), + Self::CellIdSpaceExhausted => { + formatter.write_str("code mode session exhausted its cell ID space") + } Self::DuplicateCell(cell_id) => write!(formatter, "exec cell {cell_id} already exists"), Self::MissingCell(cell_id) => write!(formatter, "exec cell {cell_id} not found"), Self::BusyObserver(cell_id) => { From 5eebeb816905f44fd847124ea610922795362cbc Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 00:13:43 +0100 Subject: [PATCH 27/33] Project executor skills through World State (#30088) ## Why A selected executor environment can be unavailable in one model step and ready in the next. The model should see its skills only while that environment is ready, without rescanning stable files on every sample. The product assumption is simple: - an environment ID names one stable logical environment; - the selected root contents do not change during the thread. ## Behavior ```text E1 unavailable -> do not show E1 skills E1 ready -> discover once, cache, show through World State E1 unavailable -> hide skills, keep cache E1 ready again -> reuse cache, show skills again resume -> create a new thread cache and discover again ``` The cache key is the full `SelectedCapabilityRoot`. Availability does not invalidate it; dropping the extension's thread state does. The step supplies the ready selected roots directly. They do not have to be turn environments: ```text turn environment: laptop selected root: worker:/plugins/lint-fix worker ready -> lint-fix skills are visible ``` ## What changes - Keeps executor skill catalogs in the existing skills extension. - Passes the roots resolved as ready for the step into World State contributors. - Loads each ready selected root at most once per thread. - Contributes the executor catalog as the `skills` World State section. - Uses the exact step catalog for explicit skill selection and body reads. - Leaves host and orchestrator skill behavior where it already lives. Taking a step snapshot itself does not add an RPC. Executor filesystem calls happen only on the first discovery of a stable root for that thread. ## What does not change - No filesystem watcher or content-based invalidation. - No retry/generation framework. - No skill runtime migration into core. - No general rewrite of the skills extension. ## Stack 1. Extension-owned World State sections. 2. **This PR:** project cached executor skills through World State. 3. Pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario. --- codex-rs/core/src/session/world_state.rs | 6 + .../src/contributors/world_state.rs | 3 + codex-rs/ext/skills/src/extension.rs | 54 +++-- codex-rs/ext/skills/src/lib.rs | 1 + codex-rs/ext/skills/src/sources.rs | 5 + codex-rs/ext/skills/src/state.rs | 77 ++++++-- codex-rs/ext/skills/src/world_state.rs | 62 ++++++ codex-rs/ext/skills/tests/skills_extension.rs | 184 +++++++++++++----- 8 files changed, 322 insertions(+), 70 deletions(-) create mode 100644 codex-rs/ext/skills/src/world_state.rs diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 2541972252..10d6a5d0d1 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -36,12 +36,18 @@ impl Session { ); } let environments = step_context.environments.to_selections(); + let ready_selected_capability_roots = step_context + .selected_capability_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect::>(); for contributor in self.services.extensions.context_contributors() { for section in contributor .contribute_world_state(WorldStateContributionInput { thread_id: self.thread_id(), turn_id: turn_context.sub_id.as_str(), environments: &environments, + ready_selected_capability_roots: &ready_selected_capability_roots, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index b5bb548105..23f180826d 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::TurnEnvironmentSelection; use serde_json::Value; @@ -11,6 +12,8 @@ pub struct WorldStateContributionInput<'a> { pub thread_id: ThreadId, pub turn_id: &'a str, pub environments: &'a [TurnEnvironmentSelection], + /// Selected roots whose stable environments are ready in this sampling step. + pub ready_selected_capability_roots: &'a [SelectedCapabilityRoot], pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, pub turn_store: &'a ExtensionData, diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index de9facd8c2..0b582e9387 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -18,8 +18,9 @@ use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; +use codex_extension_api::WorldStateContributionInput; +use codex_extension_api::WorldStateSectionContribution; use codex_mcp::McpResourceClient; -use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; @@ -40,9 +41,11 @@ use crate::render::truncate_main_prompt_contents; use crate::render::truncate_utf8_to_bytes; use crate::selection::collect_explicit_skill_mentions; use crate::sources::SkillProviders; +use crate::state::ExecutorSkillsStepState; use crate::state::SkillsThreadState; use crate::state::SkillsTurnState; use crate::tools::skill_tools; +use crate::world_state::executor_skills_world_state_section; struct SkillsExtension { providers: SkillProviders, @@ -56,18 +59,12 @@ where { fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { Box::pin(async move { - let selected_roots = input - .thread_store - .get::>() - .map(|selected_roots| selected_roots.as_ref().clone()) - .unwrap_or_default(); let orchestrator_skills_available = !input .environments .iter() .any(|environment| environment.environment_id == LOCAL_ENVIRONMENT_ID); input.thread_store.insert(SkillsThreadState::new( (self.config_from_host)(input.config), - selected_roots, orchestrator_skills_available, )); }) @@ -92,7 +89,6 @@ where let orchestrator_skills_available = true; thread_store.insert(SkillsThreadState::new( next_config, - Vec::new(), orchestrator_skills_available, )); } @@ -120,7 +116,7 @@ where .list_skills( SkillListQuery { turn_id: thread_store.level_id().to_string(), - executor_roots: thread_state.selected_roots().to_vec(), + executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, include_bundled_skills: config.bundled_skills_enabled, @@ -139,6 +135,39 @@ where .collect() }) } + + fn contribute_world_state<'a>( + &'a self, + input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let Some(thread_state) = input.thread_store.get::() else { + return Vec::new(); + }; + let config = thread_state.config(); + let catalog = thread_state + .executor_catalog_snapshot( + &self.providers, + SkillListQuery { + turn_id: input.turn_id.to_string(), + executor_roots: input.ready_selected_capability_roots.to_vec(), + host_snapshot: None, + include_host_skills: false, + include_bundled_skills: config.bundled_skills_enabled, + include_orchestrator_skills: false, + mcp_resources: input.session_store.get::(), + }, + ) + .await; + input + .turn_store + .insert(ExecutorSkillsStepState(catalog.clone())); + vec![executor_skills_world_state_section( + &catalog, + config.include_instructions, + )] + }) + } } impl ToolContributor for SkillsExtension @@ -187,14 +216,17 @@ where let host_snapshot = turn_store.get::(); let query = SkillListQuery { turn_id: input.turn_id.clone(), - executor_roots: thread_state.selected_roots().to_vec(), + executor_roots: Vec::new(), host_snapshot: host_snapshot.clone(), include_host_skills: true, include_bundled_skills: config.bundled_skills_enabled, include_orchestrator_skills: thread_state.orchestrator_skills_enabled(), mcp_resources: session_store.get::(), }; - let catalog = self.list_skills(query, &thread_state).await; + let mut catalog = self.list_skills(query, &thread_state).await; + if let Some(executor_skills) = turn_store.get::() { + catalog.extend(executor_skills.0.clone()); + } for warning in &catalog.warnings { self.emit_warning(&input.turn_id, warning.clone()); } diff --git a/codex-rs/ext/skills/src/lib.rs b/codex-rs/ext/skills/src/lib.rs index 8abdf6a93e..4a004ff9ac 100644 --- a/codex-rs/ext/skills/src/lib.rs +++ b/codex-rs/ext/skills/src/lib.rs @@ -8,6 +8,7 @@ mod selection; mod sources; mod state; mod tools; +mod world_state; pub use config::SkillsExtensionConfig; pub use extension::install; diff --git a/codex-rs/ext/skills/src/sources.rs b/codex-rs/ext/skills/src/sources.rs index d558c13b3b..a75a047e09 100644 --- a/codex-rs/ext/skills/src/sources.rs +++ b/codex-rs/ext/skills/src/sources.rs @@ -135,6 +135,11 @@ impl SkillProviders { Ok(catalog) } + pub(crate) async fn list_executor_for_turn(&self, query: SkillListQuery) -> SkillCatalog { + self.list_matching(&query, |source| source.kind == SkillSourceKind::Executor) + .await + } + async fn list_matching( &self, query: &SkillListQuery, diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 0b751ebdc2..2f1a84c46b 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -18,6 +18,7 @@ use crate::catalog::SkillProviderResult; use crate::catalog::SkillReadResult; use crate::catalog::SkillResourceId; use crate::catalog::SkillSourceKind; +use crate::provider::SkillListQuery; use crate::provider::SkillReadRequest; use crate::sources::SkillProviders; @@ -26,21 +27,17 @@ const MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES: usize = 8 * 1024 * 1024; pub(crate) struct SkillsThreadState { config: Mutex, - selected_roots: Vec, orchestrator_skills_available: bool, + executor_cache: Mutex>, orchestrator_cache: Mutex>>, } impl SkillsThreadState { - pub(crate) fn new( - config: SkillsExtensionConfig, - selected_roots: Vec, - orchestrator_skills_available: bool, - ) -> Self { + pub(crate) fn new(config: SkillsExtensionConfig, orchestrator_skills_available: bool) -> Self { Self { config: Mutex::new(config), - selected_roots, orchestrator_skills_available, + executor_cache: Mutex::new(Vec::new()), orchestrator_cache: Mutex::new(None), } } @@ -59,14 +56,33 @@ impl SkillsThreadState { .unwrap_or_else(std::sync::PoisonError::into_inner) = config; } - pub(crate) fn selected_roots(&self) -> &[SelectedCapabilityRoot] { - &self.selected_roots - } - pub(crate) fn orchestrator_skills_enabled(&self) -> bool { self.orchestrator_skills_available && self.config().orchestrator_skills_enabled } + /// Returns catalogs for stable selected roots. + /// + /// The first catalog returned for a root remains cached until this thread state is dropped. + /// Environment availability only controls whether the root is projected into the current + /// step; it never invalidates the cache. There is intentionally no filesystem watcher or + /// content-based invalidation because selected environment roots are treated as stable. + pub(crate) async fn executor_catalog_snapshot( + &self, + providers: &SkillProviders, + mut query: SkillListQuery, + ) -> SkillCatalog { + let roots = std::mem::take(&mut query.executor_roots); + let mut catalog = SkillCatalog::default(); + for root in roots { + query.executor_roots = vec![root.clone()]; + catalog.extend( + self.executor_root_catalog(providers, root, query.clone()) + .await, + ); + } + catalog + } + pub(crate) async fn orchestrator_catalog_snapshot( &self, mcp_resources: Option<&McpResourceClient>, @@ -140,6 +156,42 @@ impl SkillsThreadState { *cache = Some(Arc::clone(&next_cache)); next_cache } + + async fn executor_root_catalog( + &self, + providers: &SkillProviders, + root: SelectedCapabilityRoot, + query: SkillListQuery, + ) -> SkillCatalog { + if let Some(cached) = self + .executor_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find(|cached| cached.root == root) + { + return cached.catalog.clone(); + } + + let discovered = providers.list_executor_for_turn(query).await; + let mut cache = self + .executor_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cache.iter().find(|cached| cached.root == root) { + return cached.catalog.clone(); + } + cache.push(CachedExecutorCatalog { + root, + catalog: discovered.clone(), + }); + discovered + } +} + +struct CachedExecutorCatalog { + root: SelectedCapabilityRoot, + catalog: SkillCatalog, } struct OrchestratorGenerationCache { @@ -204,3 +256,6 @@ pub(crate) struct SkillsTurnState { pub(crate) warnings: Vec, pub(crate) main_prompts_injected: bool, } + +#[derive(Clone, Debug, Default)] +pub(crate) struct ExecutorSkillsStepState(pub(crate) SkillCatalog); diff --git a/codex-rs/ext/skills/src/world_state.rs b/codex-rs/ext/skills/src/world_state.rs new file mode 100644 index 0000000000..ce36427b0c --- /dev/null +++ b/codex-rs/ext/skills/src/world_state.rs @@ -0,0 +1,62 @@ +use codex_extension_api::ContextualUserFragment; +use codex_extension_api::PreviousWorldStateSection; +use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::WorldStateSectionContribution; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; +use serde_json::json; + +use crate::catalog::SkillCatalog; +use crate::render::available_skills_fragment; + +pub(crate) const SKILLS_WORLD_STATE_ID: &str = "skills"; +const NO_EXECUTOR_SKILLS_BODY: &str = + "\n## Skills update\nNo selected-environment skills are currently available.\n"; +const HIDDEN_EXECUTOR_SKILLS_BODY: &str = "\n## Skills update\nSelected-environment skills are not listed automatically. Explicit skill mentions can still be resolved when available.\n"; + +pub(crate) fn executor_skills_world_state_section( + catalog: &SkillCatalog, + include_instructions: bool, +) -> WorldStateSectionContribution { + let body = if include_instructions { + available_skills_fragment(catalog).map(|fragment| fragment.body()) + } else { + None + }; + let snapshot = json!({ + "body": body, + "includeInstructions": include_instructions, + }); + + WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { + let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); + if let PreviousWorldStateSection::Known(previous) = &previous { + let previous_body = previous.get("body").and_then(serde_json::Value::as_str); + let previous_include_instructions = previous + .get("includeInstructions") + .and_then(serde_json::Value::as_bool); + if previous_body == body.as_deref() + && previous_include_instructions == Some(include_instructions) + { + return None; + } + } + + let body = match body.as_deref() { + Some(body) => body, + None if previous_is_absent => return None, + None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, + None => NO_EXECUTOR_SKILLS_BODY, + }; + Some(RenderedWorldStateFragment::new( + "developer", + (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), + body, + )) + }) + .with_legacy_matcher(|role, text| { + role == "developer" + && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) + && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) + }) +} diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 761dc34d71..7e9ea01c52 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -15,10 +15,12 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::NoopTurnItemEmitter; +use codex_extension_api::PreviousWorldStateSection; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolPayload; use codex_extension_api::TurnInputContext; +use codex_extension_api::WorldStateContributionInput; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; @@ -28,6 +30,7 @@ use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SkillScope; use codex_protocol::protocol::TruncationPolicy; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_skills_extension::SkillProviders; use codex_skills_extension::SkillsExtensionConfig; @@ -144,9 +147,9 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult { } #[tokio::test] -async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_input() -> TestResult -{ +async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cache() -> TestResult { let read_requests = Arc::new(Mutex::new(Vec::new())); + let list_calls = Arc::new(AtomicUsize::new(0)); let executor_provider = Arc::new(StaticSkillProvider { catalog: SkillCatalog { entries: vec![test_entry( @@ -158,7 +161,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in warnings: Vec::new(), }, read_requests: Arc::clone(&read_requests), - list_calls: None, + list_calls: Some(Arc::clone(&list_calls)), fail_first_list: false, }); let providers = SkillProviders::new().with_executor_provider(executor_provider); @@ -168,13 +171,13 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); - thread_store.insert(vec![SelectedCapabilityRoot { + let selected_roots = vec![SelectedCapabilityRoot { id: "lint-fix".to_string(), location: CapabilityRootLocation::Environment { environment_id: "env-1".to_string(), path: PathUri::parse("file:///skills/lint-fix").expect("skill root URI"), }, - }]); + }]; let session_source = SessionSource::Cli; let config = default_config(); registry.thread_lifecycle_contributors()[0] @@ -191,20 +194,36 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in let prompt_fragments = registry.context_contributors()[0] .contribute_thread_context(&session_store, &thread_store) .await; - assert_eq!(1, prompt_fragments.len()); + assert!(prompt_fragments.is_empty()); + + let turn_store = ExtensionData::new("turn-1"); + let turn_environment = TurnEnvironmentSelection { + environment_id: "turn-env".to_string(), + cwd: PathUri::parse("file:///workspace").expect("cwd URI"), + }; + let available_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-1", + environments: std::slice::from_ref(&turn_environment), + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &turn_store, + }) + .await; + assert_eq!(1, available_sections.len()); + let available_snapshot = available_sections[0].snapshot().clone(); + let available_fragment = available_sections[0] + .render_diff(PreviousWorldStateSection::Absent) + .ok_or("available skills should render")?; + assert!(available_fragment.body().contains("lint-fix")); assert!( - prompt_fragments[0] - .text() - .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) - ); - assert!(prompt_fragments[0].text().contains("lint-fix")); - assert!( - prompt_fragments[0] - .text() + available_fragment + .body() .contains("(environment resource: skill://executor/lint-fix/SKILL.md)") ); - let turn_store = ExtensionData::new("turn-1"); let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { @@ -233,30 +252,86 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in )], read_request_keys(&read_requests) ); - let rebuilt_prompt_fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) + let unavailable_turn_store = ExtensionData::new("turn-2"); + let unavailable_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-2", + environments: &[], + ready_selected_capability_roots: &[], + session_store: &session_store, + thread_store: &thread_store, + turn_store: &unavailable_turn_store, + }) .await; - assert_eq!(1, rebuilt_prompt_fragments.len()); - assert!(rebuilt_prompt_fragments[0].text().contains("lint-fix")); + let unavailable_snapshot = unavailable_sections[0].snapshot().clone(); + let unavailable_fragment = unavailable_sections[0] + .render_diff(PreviousWorldStateSection::Known(&available_snapshot)) + .ok_or("removed skills should render")?; + assert!( + unavailable_fragment + .body() + .contains("No selected-environment skills") + ); - let next_turn_store = ExtensionData::new("turn-2"); - let next_fragments = registry.turn_input_contributors()[0] - .contribute( - TurnInputContext { - turn_id: "turn-2".to_string(), - user_input: vec![UserInput::Text { - text: "no skill this time".to_string(), - text_elements: Vec::new(), - }], - environments: Vec::new(), - }, - &session_store, - &thread_store, - &next_turn_store, - ) + let restored_turn_store = ExtensionData::new("turn-3"); + let restored_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-3", + environments: &[turn_environment], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &restored_turn_store, + }) .await; + let restored_snapshot = restored_sections[0].snapshot().clone(); + let restored_fragment = restored_sections[0] + .render_diff(PreviousWorldStateSection::Known(&unavailable_snapshot)) + .ok_or("restored skills should render")?; + assert!(restored_fragment.body().contains("lint-fix")); + assert_eq!(1, list_calls.load(Ordering::Relaxed)); - assert!(next_fragments.is_empty()); + let mut listing_disabled_config = config.clone(); + listing_disabled_config.include_instructions = false; + registry.config_contributors()[0].on_config_changed( + &session_store, + &thread_store, + &config, + &listing_disabled_config, + ); + let listing_disabled_turn_store = ExtensionData::new("turn-4"); + let listing_disabled_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-4", + environments: &[], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &listing_disabled_turn_store, + }) + .await; + let listing_disabled_fragment = listing_disabled_sections[0] + .render_diff(PreviousWorldStateSection::Known(&restored_snapshot)) + .ok_or("disabled skill listing should render")?; + assert_eq!( + "\n## Skills update\nSelected-environment skills are not listed automatically. Explicit skill mentions can still be resolved when available.\n", + listing_disabled_fragment.body() + ); + let mut normalized_listing_disabled_snapshot = listing_disabled_sections[0].snapshot().clone(); + normalized_listing_disabled_snapshot + .as_object_mut() + .ok_or("skills snapshot should be an object")? + .remove("body"); + assert!( + listing_disabled_sections[0] + .render_diff(PreviousWorldStateSection::Known( + &normalized_listing_disabled_snapshot + )) + .is_none() + ); Ok(()) } @@ -488,18 +563,16 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); - thread_store.insert( - [("root-a", "/skills/root-a"), ("root-b", "/skills/root-b")] - .into_iter() - .map(|(id, path)| SelectedCapabilityRoot { - id: id.to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "env-1".to_string(), - path: PathUri::parse(&format!("file://{path}")).expect("skill root URI"), - }, - }) - .collect::>(), - ); + let selected_roots = [("root-a", "/skills/root-a"), ("root-b", "/skills/root-b")] + .into_iter() + .map(|(id, path)| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "env-1".to_string(), + path: PathUri::parse(&format!("file://{path}")).expect("skill root URI"), + }, + }) + .collect::>(); let session_source = SessionSource::Cli; let config = default_config(); registry.thread_lifecycle_contributors()[0] @@ -513,6 +586,21 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te }) .await; + let turn_store = ExtensionData::new("turn-1"); + registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-1", + environments: &[TurnEnvironmentSelection { + environment_id: "env-1".to_string(), + cwd: PathUri::parse("file:///workspace").expect("cwd URI"), + }], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &turn_store, + }) + .await; let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { @@ -525,7 +613,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te }, &session_store, &thread_store, - &ExtensionData::new("turn-1"), + &turn_store, ) .await; From 8ce931ab763541736bd08d0507b4777543532762 Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 25 Jun 2026 16:22:22 -0700 Subject: [PATCH 28/33] [codex] Propagate traces through exec-server HTTP (#30117) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes distributed trace continuity across exec-server JSON-RPC HTTP egress by adding an executor client span and injecting its W3C context through a reusable `codex-otel` helper. This preserves the caller trace across core/tool → executor → provider/MCP instead of dropping parentage at raw reqwest. Note that this doesn't include the websocket path, which is needed to really get the full story but at least we cover the basic http path with this change. --- .../src/client/reqwest_http_client.rs | 17 ++++++++-- codex-rs/otel/src/lib.rs | 1 + codex-rs/otel/src/trace_context.rs | 32 +++++++++++++++++++ 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index 2fcb7c7e6c..667e4c49d8 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -18,6 +18,7 @@ use reqwest::Url; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use reqwest::header::HeaderValue; +use tracing::Instrument; use super::HttpResponseBodyStream; use super::response_body_stream::send_body_delta; @@ -146,15 +147,26 @@ impl ReqwestHttpRequestRunner { } } - let headers = Self::build_headers(params.headers)?; + let request_span = tracing::info_span!( + "codex.exec_server.http_request", + otel.kind = "client", + http.request.method = method.as_str(), + server.address = url.host_str().unwrap_or_default(), + server.port = u64::from(url.port_or_known_default().unwrap_or_default()), + http.response.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + let mut headers = Self::build_headers(params.headers)?; + codex_otel::inject_span_w3c_trace_headers(&request_span, &mut headers); let mut request = self.client.request(method.clone(), url).headers(headers); if let Some(body) = params.body { request = request.body(body.into_inner()); } - let response = match request.send().await { + let response = match request.send().instrument(request_span.clone()).await { Ok(response) => response, Err(error) => { + request_span.record("error.type", "request"); let error_message = error.to_string(); log_send_error(&method, error); return Err(internal_error(format!( @@ -163,6 +175,7 @@ impl ReqwestHttpRequestRunner { } }; let status = response.status().as_u16(); + request_span.record("http.response.status_code", u64::from(status)); let headers = Self::response_headers(response.headers()); if params.stream_response { diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index 2a36fab16d..87ccbbfeee 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -29,6 +29,7 @@ pub use crate::provider::OtelProvider; pub use crate::trace_context::context_from_w3c_trace_context; pub use crate::trace_context::current_span_trace_id; pub use crate::trace_context::current_span_w3c_trace_context; +pub use crate::trace_context::inject_span_w3c_trace_headers; pub use crate::trace_context::set_parent_from_context; pub use crate::trace_context::set_parent_from_w3c_trace_context; pub use crate::trace_context::span_w3c_trace_context; diff --git a/codex-rs/otel/src/trace_context.rs b/codex-rs/otel/src/trace_context.rs index c625a416a7..6ae924d1d8 100644 --- a/codex-rs/otel/src/trace_context.rs +++ b/codex-rs/otel/src/trace_context.rs @@ -49,6 +49,38 @@ pub fn span_w3c_trace_context(span: &Span) -> Option { }) } +/// Injects the W3C trace context for `span` into HTTP headers. +/// +/// Existing `traceparent` and `tracestate` values are replaced so callers can +/// safely reuse a request header map while keeping the supplied span as the +/// source of truth. +pub fn inject_span_w3c_trace_headers(span: &Span, headers: &mut http::HeaderMap) -> bool { + let Some(trace) = span_w3c_trace_context(span) else { + return false; + }; + match trace.traceparent { + Some(traceparent) => { + if let Ok(value) = http::HeaderValue::from_str(&traceparent) { + headers.insert("traceparent", value); + } + } + None => { + headers.remove("traceparent"); + } + } + match trace.tracestate { + Some(tracestate) => { + if let Ok(value) = http::HeaderValue::from_str(&tracestate) { + headers.insert("tracestate", value); + } + } + None => { + headers.remove("tracestate"); + } + } + true +} + pub(crate) fn set_tracestate_entries( entries: BTreeMap>, ) -> Result<(), Box> { From ee9e0f6387b91cd14008ea59a8e4315b2a2fdacb Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 00:53:07 +0100 Subject: [PATCH 29/33] Pin MCP runtimes to model steps (#30101) ## Why An MCP refresh can replace the session's current manager while a model step is still running. The step must execute calls through the same manager whose tools it advertised. ## Boundary ```text current session MCP runtime | | capture once for this model step v StepContext.mcp - exact MCP config - exact connection manager - exact runtime environment context ``` ```rust pub struct McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, } ``` ## Example ```text step A captures runtime A and advertises A's tools refresh publishes runtime B step A tool call -> runtime A next step -> runtime B ``` Capturing the snapshot is only an `Arc` clone. It does not restart MCPs or make an RPC. ## What changes - Captures one MCP runtime in `StepContext`. - Uses it for tool planning, tool calls, resources, approvals, connector attribution, and elicitation. - Publishes replacement runtimes atomically. - Lets an old runtime live only while an in-flight step or request still holds its `Arc`. Most of this diff is mechanical routing from the session-global manager to `step_context.mcp`; it does not introduce selected-plugin discovery yet. ## What does not change - No plugin or extension migration. - No new MCP cache policy. - No environment file watching. - No client sharing between separate managers. ## Stack 1. Extension-owned World State sections. 2. Project executor skills through World State. 3. **This PR:** pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario. --- .../src/request_processors/mcp_processor.rs | 44 +++--- codex-rs/codex-mcp/src/catalog.rs | 24 ++++ codex-rs/codex-mcp/src/catalog_tests.rs | 11 ++ codex-rs/core/src/codex_delegate.rs | 62 ++++++--- codex-rs/core/src/codex_delegate_tests.rs | 22 +-- codex-rs/core/src/codex_thread.rs | 5 + codex-rs/core/src/mcp_skill_dependencies.rs | 11 +- codex-rs/core/src/mcp_tool_call.rs | 127 +++++++++--------- codex-rs/core/src/mcp_tool_call_tests.rs | 54 +++++--- codex-rs/core/src/session/handlers.rs | 4 +- codex-rs/core/src/session/mcp.rs | 124 +++++++---------- codex-rs/core/src/session/mcp_runtime.rs | 94 +++++++++++++ codex-rs/core/src/session/mod.rs | 42 ++++-- codex-rs/core/src/session/session.rs | 56 ++++---- codex-rs/core/src/session/step_context.rs | 5 + codex-rs/core/src/session/tests.rs | 121 ++++++++++++++--- codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/session/turn_context.rs | 3 +- codex-rs/core/src/state/service.rs | 38 +++++- codex-rs/core/src/tools/handlers/mcp.rs | 5 +- .../list_mcp_resource_templates.rs | 11 +- .../mcp_resource/list_mcp_resources.rs | 11 +- .../mcp_resource/read_mcp_resource.rs | 6 +- .../tools/handlers/request_plugin_install.rs | 18 ++- codex-rs/core/src/tools/router_tests.rs | 4 +- codex-rs/core/src/tools/spec_plan_tests.rs | 4 +- .../core/tests/suite/mcp_refresh_cleanup.rs | 13 +- 27 files changed, 596 insertions(+), 325 deletions(-) create mode 100644 codex-rs/core/src/session/mcp_runtime.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index a3d2191782..abefd26015 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -121,17 +121,11 @@ impl McpRequestProcessor { } = params; let auth = self.auth_manager.auth().await; - let (config, mcp_config) = match thread_id.as_deref() { + let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let thread_config = thread.config().await; - let config = self - .config_manager - .load_latest_config_for_thread(thread_config.as_ref()) - .await - .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; - let mcp_config = thread.runtime_mcp_config(&config).await; - (config, mcp_config) + let runtime = thread.current_mcp_runtime(); + (runtime.config().clone(), runtime.runtime_context().clone()) } None => { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; @@ -140,7 +134,11 @@ impl McpRequestProcessor { .mcp_manager() .runtime_config(&config) .await; - (config, mcp_config) + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) } }; let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); @@ -167,10 +165,6 @@ impl McpRequestProcessor { } }; - let runtime_context = McpRuntimeContext::new( - self.thread_manager.environment_manager(), - config.cwd.to_path_buf(), - ); let http_client = runtime_context .resolve_http_client(&name, server) .map_err(|err| { @@ -189,16 +183,16 @@ impl McpRequestProcessor { let handle = perform_oauth_login_return_url_with_http_client( &name, &url, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, http_headers, env_http_headers, &resolved_scopes.scopes, server.oauth_client_id(), server.oauth_resource.as_deref(), timeout_secs, - config.mcp_oauth_callback_port, - config.mcp_oauth_callback_url.as_deref(), + mcp_config.mcp_oauth_callback_port, + mcp_config.mcp_oauth_callback_url.as_deref(), http_client, ) .await @@ -250,18 +244,16 @@ impl McpRequestProcessor { None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; let mcp_manager = self.thread_manager.mcp_manager(); + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); + let auth = self.auth_manager.auth().await; let mcp_config = match thread { Some(thread) => thread.runtime_mcp_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 - // as the local stdio fallback; named environment stdio MCPs must - // declare their own absolute cwd. - let runtime_context = - McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf()); + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); tokio::spawn(async move { Self::list_mcp_server_status_task( diff --git a/codex-rs/codex-mcp/src/catalog.rs b/codex-rs/codex-mcp/src/catalog.rs index 3a752c82ba..689f93d813 100644 --- a/codex-rs/codex-mcp/src/catalog.rs +++ b/codex-rs/codex-mcp/src/catalog.rs @@ -371,6 +371,30 @@ impl ResolvedMcpCatalog { .collect() } + /// Replaces the resolved server set while preserving known server sources. + /// + /// Names not present in the existing catalog are treated as config-owned. + pub fn with_materialized_servers(&self, servers: HashMap) -> Self { + let mut builder = Self::builder(); + for (name, config) in servers { + let source = self + .server(&name) + .map(|server| server.source.clone()) + .unwrap_or(McpServerSource::Config); + let precedence = match &source { + McpServerSource::Plugin(_) => RegistrationPrecedence::Plugin(Reverse(0)), + McpServerSource::SelectedPlugin(_) => { + RegistrationPrecedence::SelectedPlugin(Reverse(0)) + } + McpServerSource::Config => RegistrationPrecedence::Config, + McpServerSource::Compatibility { .. } => RegistrationPrecedence::Compatibility, + McpServerSource::Extension { .. } => RegistrationPrecedence::Extension(0), + }; + builder.register(McpServerRegistration::new(name, source, config, precedence)); + } + builder.build() + } + /// Returns package attribution for each winning plugin-owned server. pub fn plugin_attributions_by_server_name(&self) -> HashMap { self.servers diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index 0b2ffef50d..6a0515db39 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -304,6 +304,17 @@ fn selected_plugins_override_discovered_plugins_but_not_config() { }] ); + let refreshed = server("https://refreshed.example/mcp"); + let catalog = + catalog.with_materialized_servers(HashMap::from([("docs".to_string(), refreshed.clone())])); + assert_eq!( + catalog.server("docs"), + Some(&super::ResolvedMcpServer { + source: selected_plugin_source("selected-alpha"), + config: refreshed, + }) + ); + let mut builder = catalog.to_builder(); let configured = server("https://config.example/mcp"); builder.register(McpServerRegistration::from_config( diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 21e8de83ca..8874aeae79 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -40,6 +40,7 @@ use crate::guardian::spawn_approval_request_review; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT_FOR_SESSION; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; +use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; use crate::mcp_tool_call::lookup_mcp_tool_metadata; @@ -60,6 +61,12 @@ use codex_protocol::protocol::MultiAgentVersion; #[cfg(test)] use crate::session::completed_session_loop_termination; +#[derive(Clone)] +struct PendingMcpInvocation { + invocation: McpInvocation, + metadata: Option, +} + /// Start an interactive sub-Codex thread and return IO channels. /// /// The returned `events_rx` yields non-approval events emitted by the sub-agent. @@ -149,10 +156,10 @@ pub(crate) async fn run_codex_thread_interactive( let parent_session_clone = Arc::clone(&parent_session); let parent_ctx_clone = Arc::clone(&parent_ctx); let codex_for_events = Arc::clone(&codex); - // Cache delegated MCP invocations so guardian can recover the full tool call - // context when the later legacy RequestUserInput approval event only carries - // a call_id plus approval question metadata. - let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::::new())); + // Cache the child call's MCP metadata at begin time. The later legacy + // RequestUserInput approval event only carries a call_id and question metadata. + let pending_mcp_invocations = + Arc::new(Mutex::new(HashMap::::new())); tokio::spawn(async move { forward_events( codex_for_events, @@ -270,7 +277,7 @@ async fn forward_events( tx_sub: Sender, parent_session: Arc, parent_ctx: Arc, - pending_mcp_invocations: Arc>>, + pending_mcp_invocations: Arc>>, cancel_token: CancellationToken, ) { let cancelled = cancel_token.cancelled(); @@ -357,10 +364,34 @@ async fn forward_events( id, msg: EventMsg::McpToolCallBegin(event), } => { + // Runtime refreshes are published before a request step is captured, so + // the child runtime at call begin is the one executing this invocation. + // Cache its metadata now; the later approval event has only a call ID. + let metadata = if let Some(turn_context) = + codex.session.turn_context_for_sub_id(&id).await + { + let mcp = codex.session.services.latest_mcp_runtime(); + lookup_mcp_tool_metadata( + codex.session.as_ref(), + turn_context.as_ref(), + mcp.manager(), + &event.invocation.server, + &event.invocation.tool, + ) + .await + } else { + None + }; pending_mcp_invocations .lock() .await - .insert(event.call_id.clone(), event.invocation.clone()); + .insert( + event.call_id.clone(), + PendingMcpInvocation { + invocation: event.invocation.clone(), + metadata, + }, + ); if !forward_event_or_shutdown( &codex, &tx_sub, @@ -648,7 +679,7 @@ async fn handle_request_user_input( id: String, parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: RequestUserInputEvent, cancel_token: &CancellationToken, ) { @@ -686,12 +717,12 @@ async fn handle_request_user_input( /// programmatically after running the guardian review. /// /// The RequestUserInput event only carries `call_id` plus approval question -/// metadata, so this helper joins it back to the cached `McpToolCallBegin` -/// invocation in order to rebuild the full guardian review request. +/// metadata, so this helper joins it back to the child runtime metadata cached at +/// `McpToolCallBegin` in order to rebuild the full guardian review request. async fn maybe_auto_review_mcp_request_user_input( parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: &RequestUserInputEvent, cancel_token: &CancellationToken, ) -> Option { @@ -702,18 +733,13 @@ async fn maybe_auto_review_mcp_request_user_input( .questions .iter() .find(|question| is_mcp_tool_approval_question_id(&question.id))?; - let invocation = pending_mcp_invocations + let pending = pending_mcp_invocations .lock() .await .get(&event.call_id) .cloned()?; - let metadata = lookup_mcp_tool_metadata( - parent_session.as_ref(), - parent_ctx.as_ref(), - &invocation.server, - &invocation.tool, - ) - .await; + let invocation = pending.invocation; + let metadata = pending.metadata; let approvals_reviewer = mcp_approvals_reviewer(parent_ctx, &invocation.server, metadata.as_ref()); if !routes_approval_to_guardian_with_reviewer(parent_ctx, approvals_reviewer) { diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 6d51a0032b..2791c887e6 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -411,10 +411,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: "custom_server".to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: "custom_server".to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); @@ -460,10 +463,13 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { crate::session::tests::make_session_and_context_with_rx().await; let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: CODEX_APPS_MCP_SERVER_NAME.to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: CODEX_APPS_MCP_SERVER_NAME.to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 25f47ce45d..46b38af151 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -594,6 +594,11 @@ impl CodexThread { self.codex.session.runtime_mcp_config(config).await } + /// Returns the exact MCP config, environment bindings, and manager most recently published. + pub fn current_mcp_runtime(&self) -> Arc { + self.codex.session.services.latest_mcp_runtime() + } + pub fn multi_agent_version(&self) -> Option { self.codex.session.multi_agent_version() } diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 9b2b85aa2d..61d5c54a59 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -199,15 +199,8 @@ pub(crate) async fn maybe_install_mcp_dependencies( warn!("failed to refresh MCP dependencies for mentioned skills: {err}"); return; } - let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await; - sess.refresh_mcp_servers_now( - turn_context, - refresh_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - elicitation_reviewer, - ) - .await; + sess.refresh_mcp_servers_now(turn_context, &refresh_config, elicitation_reviewer) + .await; } async fn should_install_mcp_dependencies( diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index a270af7851..19b7c1a6ad 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -18,6 +18,7 @@ use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; use crate::mcp_tool_approval_templates::render_mcp_tool_approval_template; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::PermissionRequestPayload; @@ -35,6 +36,7 @@ use codex_features::Feature; use codex_hooks::PermissionRequestDecision; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; +use codex_mcp::McpConnectionManager; use codex_mcp::McpPermissionPromptAutoApproveContext; use codex_mcp::SandboxState; use codex_mcp::auth_elicitation_completed_result; @@ -111,13 +113,15 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; /// item lifecycle events to the `Session`. pub(crate) async fn handle_mcp_tool_call( sess: Arc, - turn_context: &Arc, + step_context: &Arc, call_id: String, server: String, tool_name: String, hook_tool_name: HookToolName, arguments: String, ) -> HandledMcpToolCall { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); // Parse the `arguments` as JSON. An empty string is OK, but invalid JSON // is not. let arguments_value = if arguments.trim().is_empty() { @@ -141,8 +145,14 @@ pub(crate) async fn handle_mcp_tool_call( arguments: arguments_value.clone(), }; - let metadata = - lookup_mcp_tool_metadata(sess.as_ref(), turn_context.as_ref(), &server, &tool_name).await; + let metadata = lookup_mcp_tool_metadata( + sess.as_ref(), + turn_context.as_ref(), + manager, + &server, + &tool_name, + ) + .await; let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref()); let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME { let annotations = metadata @@ -169,7 +179,6 @@ pub(crate) async fn handle_mcp_tool_call( } else if let Some(approval_mode) = { // Selected-plugin registrations are absent from config.toml and the legacy plugin manager, // so their resolved catalog entry is the authoritative source for tool approval policy. - let manager = sess.services.mcp_connection_manager.load(); manager .is_selected_plugin_mcp_server(&server) .then(|| manager.tool_approval_mode(&server, &tool_name)) @@ -226,7 +235,7 @@ pub(crate) async fn handle_mcp_tool_call( if let Some(decision) = maybe_request_mcp_tool_approval( &sess, - turn_context, + step_context, &call_id, &invocation, &hook_tool_name, @@ -241,7 +250,7 @@ pub(crate) async fn handle_mcp_tool_call( | McpToolApprovalDecision::AcceptAndRemember => { return handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -298,7 +307,7 @@ pub(crate) async fn handle_mcp_tool_call( handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -340,24 +349,21 @@ impl McpToolCallItemMetadata { async fn handle_approved_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: McpInvocation, metadata: Option<&McpToolApprovalMetadata>, item_metadata: McpToolCallItemMetadata, ) -> HandledMcpToolCall { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let server = invocation.server.clone(); - maybe_mark_thread_memory_mode_polluted(sess, turn_context, &server).await; + maybe_mark_thread_memory_mode_polluted(sess, turn_context, manager, &server).await; let tool_name = invocation.tool.clone(); let arguments_value = invocation.arguments.clone(); let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref()); let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref()); - let server_origin = { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - mcp_connection_manager - .server_origin(&server) - .map(str::to_string) - }; + let server_origin = manager.server_origin(&server).map(str::to_string); let start = Instant::now(); let rewrite = rewrite_mcp_tool_arguments_for_openai_files( @@ -380,7 +386,7 @@ async fn handle_approved_mcp_tool_call( build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata); execute_mcp_tool_call( sess, - turn_context, + step_context, call_id, &invocation, rewritten_arguments, @@ -420,7 +426,7 @@ async fn handle_approved_mcp_tool_call( truncate_mcp_tool_result_for_event(&result), ) .await; - maybe_track_codex_app_used(sess, turn_context, &server, &tool_name).await; + maybe_track_codex_app_used(sess, turn_context, manager, &server, &tool_name).await; let outcome = mcp_call_metric_outcome(&result); emit_mcp_call_metrics( @@ -547,17 +553,19 @@ fn truncate_str_to_char_boundary(value: &str, max_chars: usize) -> &str { async fn execute_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: &McpInvocation, rewritten_arguments: Option, metadata: Option<&McpToolApprovalMetadata>, request_meta: Option, ) -> Result { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let request_meta = with_mcp_tool_call_thread_id_meta(request_meta, &sess.thread_id.to_string()); let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( - sess, - turn_context, + step_context, + manager, &invocation.server, request_meta, ) @@ -568,7 +576,7 @@ async fn execute_mcp_tool_call( .rollout_thread_trace .start_mcp_call_trace(call_id); let request_meta = mcp_call_trace.add_request_meta(request_meta); - let result = sess + let result = manager .call_tool( &invocation.server, &invocation.tool, @@ -587,6 +595,7 @@ async fn execute_mcp_tool_call( Ok(maybe_request_codex_apps_auth_elicitation( sess, turn_context, + manager, call_id, &invocation.server, metadata, @@ -598,17 +607,13 @@ async fn execute_mcp_tool_call( async fn maybe_request_codex_apps_auth_elicitation( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, call_id: &str, server: &str, metadata: Option<&McpToolApprovalMetadata>, result: CallToolResult, ) -> CallToolResult { - if !sess - .services - .mcp_connection_manager - .load_full() - .is_host_owned_codex_apps_server(server) - { + if !manager.is_host_owned_codex_apps_server(server) { return result; } @@ -666,15 +671,16 @@ async fn maybe_request_codex_apps_auth_elicitation( return result; } - refresh_codex_apps_after_connector_auth(sess, turn_context).await; + refresh_codex_apps_after_connector_auth(sess, turn_context, manager).await; auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) { - let mcp_tools_result = { - let manager = sess.services.mcp_connection_manager.load_full(); - manager.hard_refresh_codex_apps_tools_cache().await - }; +async fn refresh_codex_apps_after_connector_auth( + sess: &Session, + turn_context: &TurnContext, + manager: &McpConnectionManager, +) { + let mcp_tools_result = manager.hard_refresh_codex_apps_tools_cache().await; match mcp_tools_result { Ok(mcp_tools) => { @@ -692,13 +698,13 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & } async fn augment_mcp_tool_request_meta_with_sandbox_state( - sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, + manager: &McpConnectionManager, server: &str, mut meta: Option, ) -> anyhow::Result> { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - let supports_sandbox_state_meta = mcp_connection_manager + let turn_context = step_context.turn.as_ref(); + let supports_sandbox_state_meta = manager .server_supports_sandbox_state_meta_capability(server) .await .unwrap_or(false); @@ -706,18 +712,18 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( return Ok(meta); } - let server_environment_id = mcp_connection_manager + let server_environment_id = manager .server_environment_id(server) .unwrap_or(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID); - let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(turn_context, server_environment_id) else { + let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(step_context, server_environment_id) else { return Ok(meta); }; let permission_profile = turn_context.permission_profile(); let sandbox_state = serde_json::to_value(SandboxState { permission_profile, - codex_linux_sandbox_exe: turn_context.config.codex_linux_sandbox_exe.clone(), + codex_linux_sandbox_exe: step_context.mcp.config().codex_linux_sandbox_exe.clone(), sandbox_cwd, - use_legacy_landlock: turn_context.config.features.use_legacy_landlock(), + use_legacy_landlock: step_context.mcp.config().use_legacy_landlock, })?; match meta.as_mut() { @@ -741,8 +747,8 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( Ok(meta) } -fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) -> Option { - if let Some(environment) = turn_context +fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) -> Option { + if let Some(environment) = step_context .environments .turn_environments .iter() @@ -753,7 +759,7 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) if environment_id == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID { #[allow(deprecated)] - return Some(PathUri::from_abs_path(&turn_context.cwd)); + return Some(PathUri::from_abs_path(&step_context.turn.cwd)); } None @@ -762,16 +768,13 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) async fn maybe_mark_thread_memory_mode_polluted( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, ) { if !turn_context.config.memories.disable_on_external_context { return; } - let pollutes_memory = sess - .services - .mcp_connection_manager - .load_full() - .server_pollutes_memory(server); + let pollutes_memory = manager.server_pollutes_memory(server); if !pollutes_memory { return; } @@ -937,13 +940,14 @@ struct McpAppUsageMetadata { async fn maybe_track_codex_app_used( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) { if server != CODEX_APPS_MCP_SERVER_NAME { return; } - let metadata = lookup_mcp_app_usage_metadata(sess, server, tool_name).await; + let metadata = lookup_mcp_app_usage_metadata(manager, server, tool_name).await; let (connector_id, app_name) = metadata .map(|metadata| (metadata.connector_id, metadata.app_name)) .unwrap_or((None, None)); @@ -983,6 +987,7 @@ enum McpToolApprovalDecision { Cancel, } +#[derive(Clone)] pub(crate) struct McpToolApprovalMetadata { annotations: Option, connector_id: Option, @@ -1176,13 +1181,15 @@ fn mcp_tool_approval_prompt_options( async fn maybe_request_mcp_tool_approval( sess: &Arc, - turn_context: &Arc, + step_context: &Arc, call_id: &str, invocation: &McpInvocation, hook_tool_name: &HookToolName, metadata: Option<&McpToolApprovalMetadata>, approval_mode: AppToolApproval, ) -> Option { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); let approvals_reviewer = mcp_approvals_reviewer(turn_context, &invocation.server, metadata); if mcp_permission_prompt_is_auto_approved( turn_context.approval_policy.value(), @@ -1201,12 +1208,7 @@ async fn maybe_request_mcp_tool_approval( } let session_approval_key = session_mcp_tool_approval_key(invocation, metadata, approval_mode); - let persistent_approval_key = if sess - .services - .mcp_connection_manager - .load() - .is_selected_plugin_mcp_server(&invocation.server) - { + let persistent_approval_key = if manager.is_selected_plugin_mcp_server(&invocation.server) { None } else { persistent_mcp_tool_approval_key(invocation, metadata, approval_mode) @@ -1452,10 +1454,10 @@ async fn mcp_tool_approval_decision_from_guardian( pub(crate) async fn lookup_mcp_tool_metadata( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let manager = sess.services.mcp_connection_manager.load_full(); let plugin_id = manager .plugin_id_for_mcp_server_name(server) .map(str::to_string); @@ -1568,16 +1570,11 @@ fn get_mcp_app_resource_uri( } async fn lookup_mcp_app_usage_metadata( - sess: &Session, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let tools = sess - .services - .mcp_connection_manager - .load_full() - .list_all_tools() - .await; + let tools = manager.list_all_tools().await; tools.into_iter().find_map(|tool_info| { if tool_info.server_name == server && tool_info.tool.name == tool_name { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 41701d557c..8c8d82037b 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::ConfigBuilder; use crate::config::ManagedFeatures; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::tests::make_session_and_context_with_rx; use crate::session::turn_context::TurnEnvironment; @@ -157,10 +158,12 @@ async fn execute_mcp_tool_call_records_replayable_correlation() -> anyhow::Resul }) }); assert!(dispatch_trace.is_enabled()); + let turn_context = Arc::new(turn_context); + let step_context = StepContext::for_test(Arc::clone(&turn_context)); let result = execute_mcp_tool_call( &session, - &turn_context, + step_context.as_ref(), "mcp-call", &McpInvocation { server: "docs".to_string(), @@ -1147,7 +1150,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul /*shell*/ None, )); - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, Some(secondary_cwd)); Ok(()) @@ -1157,7 +1161,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul async fn mcp_sandbox_cwd_is_none_for_unselected_server_environment() -> anyhow::Result<()> { let (_, turn_context) = make_session_and_context().await; - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, None); Ok(()) @@ -1370,7 +1375,10 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { ) } -async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { +async fn host_owned_codex_apps_manager( + session: &Session, + turn_context: &TurnContext, +) -> Arc { let auth = session.services.auth_manager.auth().await; let startup_cancellation_token = CancellationToken::new(); startup_cancellation_token.cancel(); @@ -1410,22 +1418,20 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: /*elicitation_reviewer*/ None, ) .await; - session - .services - .mcp_connection_manager - .store(Arc::new(manager)); + Arc::new(manager) } #[tokio::test] async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() { let (session, turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); let returned = maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1446,10 +1452,12 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); + let manager = session.services.latest_mcp_runtime().manager_arc(); let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1464,7 +1472,7 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1479,6 +1487,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1493,7 +1502,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul #[tokio::test] async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1514,6 +1523,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1528,7 +1538,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); @@ -1542,10 +1552,12 @@ async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let request_task = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let manager = Arc::clone(&manager); async move { maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -2466,7 +2478,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-1", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2542,7 +2554,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2601,7 +2613,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2663,7 +2675,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook-no-metadata", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2745,7 +2757,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-remembered", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2828,7 +2840,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian-deny", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2888,7 +2900,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval tokio::spawn(async move { maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-prompt", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2946,7 +2958,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { ] { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-2", &invocation, &HookToolName::new("mcp__test__tool"), @@ -3035,7 +3047,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-3", &invocation, &HookToolName::new("mcp__test__tool"), diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 5c68893e7f..75fbca4c86 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -598,8 +598,8 @@ async fn shutdown_session_runtime(sess: &Arc) { warn!("failed to shutdown code mode session: {err}"); } sess.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .shutdown() .await; sess.guardian_review_session.shutdown().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 249a6994b1..b9e87f0c89 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -105,8 +105,8 @@ impl Session { ) -> McpServerElicitationOutcome { if self .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .elicitations_auto_deny() { return McpServerElicitationOutcome { @@ -201,44 +201,20 @@ impl Session { } self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .resolve_elicitation(server_name, id, response) .await } - pub async fn list_resources( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .mcp_connection_manager - .load_full() - .list_resources(server, params) - .await - } - - pub async fn list_resource_templates( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .mcp_connection_manager - .load_full() - .list_resource_templates(server, params) - .await - } - pub async fn read_resource( &self, server: &str, params: ReadResourceRequestParams, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .read_resource(server, params) .await } @@ -251,8 +227,8 @@ impl Session { meta: Option, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .call_tool(server, tool, arguments, meta) .await } @@ -260,17 +236,17 @@ impl Session { async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + mut mcp_config: McpConfig, + configured_mcp_servers: HashMap, elicitation_reviewer: Option, ) { + mcp_config.mcp_server_catalog = mcp_config + .mcp_server_catalog + .with_materialized_servers(configured_mcp_servers); + let mcp_config = Arc::new(mcp_config); let auth = self.services.auth_manager.auth().await; - let config = self.get_config().await; - let mcp_config = self.runtime_mcp_config(config.as_ref()).await; let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); - let mcp_servers = - effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); + let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. @@ -286,35 +262,36 @@ impl Session { let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd); let auth_statuses = compute_auth_statuses( mcp_servers.iter(), - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth.as_ref(), &mcp_runtime_context, ) .await; let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - guard.cancel(); + // The previous runtime owns the old token and may still be serving an in-flight step. + // Its manager cancels that token when the last runtime handle is dropped. let cancellation_token = CancellationToken::new(); *guard = cancellation_token.clone(); cancellation_token }; let refreshed_manager = McpConnectionManager::new( &mcp_servers, - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth_statuses, &turn_context.approval_policy, turn_context.sub_id.clone(), self.get_tx_event(), mcp_startup_cancellation_token, turn_context.permission_profile(), - mcp_runtime_context, - config.codex_home.to_path_buf(), + mcp_runtime_context.clone(), + mcp_config.codex_home.clone(), 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, + mcp_config.client_elicitation_capability.clone(), self.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -324,14 +301,12 @@ impl Session { ) .await; { - let current_manager = self.services.mcp_connection_manager.load_full(); - refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); + let current_manager = self.services.latest_mcp_runtime(); + refreshed_manager + .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); } - let superseded_manager = self - .services - .mcp_connection_manager - .swap(Arc::new(refreshed_manager)); - superseded_manager.shutdown().await; + self.services + .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); } pub(crate) async fn refresh_mcp_servers_if_requested( @@ -376,14 +351,23 @@ impl Session { } }; - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + let mut refresh_config = self.get_config().await.as_ref().clone(); + refresh_config.mcp_oauth_credentials_store_mode = store_mode; + let secret_auth_storage_enabled = match keyring_backend_kind { + AuthKeyringBackendKind::Direct => false, + AuthKeyringBackendKind::Secrets => true, + }; + if let Err(err) = refresh_config + .features + .set_enabled(Feature::SecretAuthStorage, secret_auth_storage_enabled) + { + warn!("failed to apply MCP auth keyring backend refresh config: {err}"); + return; + } + + let mcp_config = self.runtime_mcp_config(&refresh_config).await; + self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + .await; } pub(crate) async fn set_openai_form_elicitation_support( @@ -417,19 +401,13 @@ impl Session { pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + refresh_config: &Config, elicitation_reviewer: Option, ) { - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + let mcp_config = self.runtime_mcp_config(refresh_config).await; + let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); + self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + .await; } #[cfg(test)] diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs new file mode 100644 index 0000000000..dc5ea011ca --- /dev/null +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -0,0 +1,94 @@ +use std::fmt; +use std::sync::Arc; + +use codex_mcp::McpConfig; +use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; + +/// MCP config, exact environment bindings, and manager used by one model request. +pub struct McpRuntimeSnapshot { + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, +} + +impl McpRuntimeSnapshot { + pub(crate) fn new( + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, + ) -> Self { + Self { + config, + manager, + runtime_context, + } + } + + pub fn config(&self) -> &McpConfig { + self.config.as_ref() + } + + pub fn manager(&self) -> &McpConnectionManager { + self.manager.as_ref() + } + + pub(crate) fn manager_arc(&self) -> Arc { + Arc::clone(&self.manager) + } + + pub fn runtime_context(&self) -> &McpRuntimeContext { + &self.runtime_context + } + + #[cfg(test)] + pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { + use codex_exec_server::EnvironmentManager; + use codex_features::Feature; + use codex_mcp::ResolvedMcpCatalog; + use rmcp::model::ElicitationCapability; + + let mcp_config = McpConfig { + chatgpt_base_url: config.chatgpt_base_url.clone(), + apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), + codex_home: config.codex_home.to_path_buf(), + mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, + auth_keyring_backend_kind: config.auth_keyring_backend_kind(), + mcp_oauth_callback_port: config.mcp_oauth_callback_port, + mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), + skill_mcp_dependency_install_enabled: config + .features + .enabled(Feature::SkillMcpDependencyInstall), + approval_policy: config.permissions.approval_policy.clone(), + codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), + use_legacy_landlock: config.features.use_legacy_landlock(), + apps_enabled: config.features.enabled(Feature::Apps), + prefix_mcp_tool_names: config.prefix_mcp_tool_names(), + client_elicitation_capability: ElicitationCapability::default(), + mcp_server_catalog: ResolvedMcpCatalog::default(), + connector_snapshot: codex_connectors::ConnectorSnapshot::default(), + }; + let manager = McpConnectionManager::new_uninitialized_with_permission_profile( + &config.permissions.approval_policy, + config.permissions.permission_profile(), + config.prefix_mcp_tool_names(), + ); + let runtime_context = McpRuntimeContext::new( + Arc::new(EnvironmentManager::default_for_tests()), + config.cwd.to_path_buf(), + ); + Arc::new(Self::new( + Arc::new(mcp_config), + Arc::new(manager), + runtime_context, + )) + } +} + +impl fmt::Debug for McpRuntimeSnapshot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("McpRuntimeSnapshot") + .finish_non_exhaustive() + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6fe61222cf..9421e07a66 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -158,15 +158,9 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ElicitationCapability; -use rmcp::model::FormElicitationCapability; -use rmcp::model::ListResourceTemplatesResult; -use rmcp::model::ListResourcesResult; -use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; -use rmcp::model::UrlElicitationCapability; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -214,6 +208,7 @@ mod handlers; mod inject; mod input_queue; mod mcp; +mod mcp_runtime; pub(crate) mod multi_agents; mod review; mod rollout_budget; @@ -235,6 +230,7 @@ use self::handlers::submission_loop; pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; +pub use self::mcp_runtime::McpRuntimeSnapshot; use self::review::spawn_review_thread; use self::session::AppServerClientMetadata; use self::session::Session; @@ -340,7 +336,7 @@ use codex_core_plugins::RecommendedPluginCandidatesInput; use codex_git_utils::get_git_repo_root; use codex_mcp::McpConfig; use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers_from_configured; +use codex_mcp::effective_mcp_servers; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; @@ -840,8 +836,11 @@ impl Codex { ..Default::default() }) .await?; - let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full(); - mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny); + self.session + .services + .latest_mcp_runtime() + .manager() + .set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -2855,10 +2854,12 @@ impl Session { &environments.captured_environments(), ) .await; + let mcp = self.services.latest_mcp_runtime(); Arc::new(StepContext::new( turn_context, environments, selected_capability_roots, + mcp, loaded_agents_md, )) } @@ -3125,6 +3126,17 @@ impl Session { &self, turn_context: &TurnContext, world_state: &WorldState, + ) -> Vec { + let mcp = self.services.latest_mcp_runtime(); + self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp) + .await + } + + async fn build_initial_context_with_world_state_and_mcp( + &self, + turn_context: &TurnContext, + world_state: &WorldState, + mcp: &McpRuntimeSnapshot, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3218,10 +3230,9 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( - &mcp_connection_manager, + mcp.manager(), &turn_context.config, ) .await; @@ -3328,7 +3339,8 @@ impl Session { if turn_context.config.features.enabled(Feature::TokenBudget) && turn_context.model_context_window().is_some() { - let mcp_result = self + let mcp_result = mcp + .manager() .call_tool( "notes", "thread_hint", @@ -3543,7 +3555,11 @@ impl Session { // Full initial context resets the baseline; later turns persist only its changes. let (mut context_items, world_state_item) = if should_inject_full_context { let context_items = self - .build_initial_context_with_world_state(turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + turn_context, + world_state.as_ref(), + step_context.mcp.as_ref(), + ) .await; let snapshot = world_state.snapshot(); self.state diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f385a8c983..2c9e57036d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -661,10 +661,15 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; - let mcp_runtime_context_for_auth = McpRuntimeContext::new( - Arc::clone(&environment_manager), - session_configuration.cwd().to_path_buf(), - ); + let mcp_runtime_cwd = session_configuration + .environment_selections() + .first() + .and_then(|environment| environment.cwd.to_abs_path().ok()) + .map(|cwd| cwd.to_path_buf()) + .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); + let mcp_runtime_context = + McpRuntimeContext::new(Arc::clone(&environment_manager), mcp_runtime_cwd); + let mcp_runtime_context_for_auth = mcp_runtime_context.clone(); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp @@ -680,7 +685,13 @@ impl Session { &mcp_runtime_context_for_auth, ) .await; - (auth, mcp_servers, auth_statuses, tool_plugin_provenance) + ( + auth, + mcp_config, + mcp_servers, + auth_statuses, + tool_plugin_provenance, + ) } .instrument(info_span!( "session_init.auth_mcp", @@ -691,7 +702,7 @@ impl Session { let ( thread_persistence_result, state_db_ctx, - (auth, mcp_servers, auth_statuses, tool_plugin_provenance), + (auth, mcp_config, mcp_servers, auth_statuses, tool_plugin_provenance), ) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let mut live_thread_init = @@ -1027,6 +1038,7 @@ impl Session { // changing this to use Option or OnceCell, though the current // setup is straightforward enough and performs well. mcp_connection_manager, + mcp_runtime: arc_swap::ArcSwapOption::empty(), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1159,14 +1171,6 @@ impl Session { sess.send_event_raw(event).await; } - let client_elicitation_capability = if config.features.enabled(Feature::AuthElicitation) { - ElicitationCapability { - form: Some(FormElicitationCapability::default()), - url: Some(UrlElicitationCapability::default()), - } - } else { - ElicitationCapability::default() - }; let mcp_startup_cancellation_token = { let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; cancel_guard.cancel(); @@ -1174,20 +1178,6 @@ impl Session { *cancel_guard = cancel_token.clone(); cancel_token }; - let mcp_runtime_context = { - let turn_environments = sess.services.turn_environments.snapshot().await; - // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment - // cwd values can be used without falling back to the session host cwd. - let cwd = turn_environments - .primary() - .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) - .map(|cwd| cwd.to_path_buf()) - .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); - McpRuntimeContext::new( - sess.services.turn_environments.environment_manager(), - cwd, - ) - }; let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, @@ -1198,12 +1188,12 @@ impl Session { tx_event.clone(), mcp_startup_cancellation_token, session_configuration.permission_profile(), - mcp_runtime_context, + mcp_runtime_context.clone(), 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, + mcp_config.client_elicitation_capability.clone(), sess.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -1217,7 +1207,11 @@ impl Session { )) .await; sess.services - .install_mcp_connection_manager(mcp_connection_manager) + .install_mcp_connection_manager( + Arc::new(mcp_config), + mcp_runtime_context, + mcp_connection_manager, + ) .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index e463015459..3210eb0286 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use codex_exec_server::ResolvedSelectedCapabilityRoot; @@ -12,6 +13,8 @@ pub(crate) struct StepContext { pub(crate) environments: TurnEnvironmentSnapshot, /// Capability roots bound to ready environments in this exact step. pub(crate) selected_capability_roots: Vec, + /// The exact MCP config and manager used to advertise and execute tools for this step. + pub(crate) mcp: Arc, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -21,12 +24,14 @@ impl StepContext { turn: Arc, environments: TurnEnvironmentSnapshot, selected_capability_roots: Vec, + mcp: Arc, loaded_agents_md: Option>, ) -> Self { Self { turn, environments, selected_capability_roots, + mcp, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 525fbc3b99..d0aac04249 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -18,6 +18,7 @@ use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; use codex_config::ConfigLayerStack; use codex_config::ConfigLayerStackOrdering; +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::LoaderOverrides; use codex_config::NetworkConstraints; use codex_config::NetworkDomainPermissionToml; @@ -25,6 +26,8 @@ use codex_config::NetworkDomainPermissionsToml; use codex_config::RequirementSource; use codex_config::Sourced; use codex_config::loader::project_trust_key; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerTransportConfig; use codex_config::types::ToolSuggestDisabledTool; use codex_core_skills::HostSkillsSnapshot; use core_test_support::test_codex::local_selections; @@ -188,9 +191,10 @@ impl StepContext { pub(crate) fn for_test(turn: Arc) -> Arc { let environments = turn.environments.clone(); Arc::new(Self::new( - turn, + Arc::clone(&turn), environments, Vec::new(), + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )) } @@ -379,8 +383,8 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() let (session, turn_context, rx) = make_session_and_context_with_rx().await; session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .set_elicitations_auto_deny(/*auto_deny*/ true); let response = session @@ -5363,14 +5367,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7440,14 +7441,11 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7603,8 +7601,14 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn refresh_mcp_servers_is_deferred_until_next_turn() { +async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let (session, turn_context) = make_session_and_context().await; + let turn_context = Arc::new(turn_context); + let old_runtime = session.services.latest_mcp_runtime(); + let step_context = session + .capture_step_context(Arc::clone(&turn_context)) + .await; + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); let old_token = session.mcp_startup_cancellation_token().await; assert!(!old_token.is_cancelled()); @@ -7612,8 +7616,16 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { serde_json::to_value(OAuthCredentialsStoreMode::Auto).expect("serialize store mode"); let auth_keyring_backend_kind = serde_json::to_value(AuthKeyringBackendKind::Secrets).expect("serialize keyring backend"); + let refreshed_mcp_servers = serde_json::from_value::>(json!({ + "refreshed": { + "url": "https://refreshed.example/mcp", + "enabled": false + } + })) + .expect("parse refreshed MCP servers"); let refresh_config = McpServerRefreshConfig { - mcp_servers: json!({}), + mcp_servers: serde_json::to_value(&refreshed_mcp_servers) + .expect("serialize refreshed MCP servers"), mcp_oauth_credentials_store_mode, auth_keyring_backend_kind, }; @@ -7635,7 +7647,7 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { .refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None) .await; - assert!(old_token.is_cancelled()); + assert!(!old_token.is_cancelled()); assert!( session .pending_mcp_server_refresh_config @@ -7645,6 +7657,71 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { ); let new_token = session.mcp_startup_cancellation_token().await; assert!(!new_token.is_cancelled()); + let new_runtime = session.services.latest_mcp_runtime(); + assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); + assert_eq!( + codex_mcp::configured_mcp_servers(new_runtime.config()), + refreshed_mcp_servers + ); +} + +#[tokio::test] +async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { + let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + let step_context = session.capture_step_context(turn_context).await; + + let mut refresh_config = step_context.turn.config.as_ref().clone(); + refresh_config.mcp_servers.set(HashMap::from([( + "newer".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::Stdio { + command: "missing-test-mcp-server".to_string(), + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + )]))?; + session + .refresh_mcp_servers_now( + step_context.turn.as_ref(), + &refresh_config, + /*elicitation_reviewer*/ None, + ) + .await; + + let router = crate::session::turn::built_tools( + session.as_ref(), + step_context.as_ref(), + &CancellationToken::new(), + ) + .await?; + assert!( + !router + .registered_tool_names_for_test() + .iter() + .any(|name| name.to_string() == "list_mcp_resources") + ); + Ok(()) } #[tokio::test] @@ -10161,8 +10238,8 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { let tools = { session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await }; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 9b3588fc1e..2cb1f2cada 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1177,7 +1177,7 @@ pub(crate) async fn built_tools( cancellation_token: &CancellationToken, ) -> CodexResult> { let turn_context = step_context.turn.as_ref(); - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); + let mcp_connection_manager = step_context.mcp.manager(); let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_tools() diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a5ae458957..5efbeb8ef3 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -686,7 +686,8 @@ impl Session { .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); + let mcp_runtime = self.services.latest_mcp_runtime(); + let mcp_connection_manager = mcp_runtime.manager(); mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); mcp_connection_manager .set_permission_profile(session_configuration.permission_profile()); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 3721cb1848..1bee45cce0 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; +use crate::session::McpRuntimeSnapshot; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; @@ -30,7 +31,9 @@ use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; +use codex_mcp::McpConfig; use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -44,8 +47,10 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - /// The latest manager; callers retain an owned handle while performing MCP I/O. + /// Mirror of the latest manager for extension resource clients that predate runtime snapshots. pub(crate) mcp_connection_manager: Arc>, + /// The latest atomically published MCP config and manager pair. + pub(crate) mcp_runtime: ArcSwapOption, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -99,12 +104,33 @@ impl SessionServices { /// resolve through the session's manager while validation waits. pub(crate) async fn install_mcp_connection_manager( &self, + config: Arc, + runtime_context: McpRuntimeContext, manager: McpConnectionManager, ) -> Result<()> { - self.mcp_connection_manager.store(Arc::new(manager)); - self.mcp_connection_manager - .load_full() - .validate_required_servers() - .await + let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + runtime.manager().validate_required_servers().await + } + + pub(crate) fn publish_mcp_runtime( + &self, + config: Arc, + runtime_context: McpRuntimeContext, + manager: McpConnectionManager, + ) -> Arc { + let manager = Arc::new(manager); + // Publish the manager for legacy resource clients first. Once the paired snapshot is + // visible, every model-scoped consumer observes this exact manager. + self.mcp_connection_manager.store(Arc::clone(&manager)); + let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); + self.mcp_runtime.store(Some(Arc::clone(&runtime))); + runtime + } + + pub(crate) fn latest_mcp_runtime(&self) -> Arc { + let Some(runtime) = self.mcp_runtime.load_full() else { + unreachable!("MCP runtime must be installed before handling requests"); + }; + runtime } } diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index b92a2a5655..c9533fd223 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -124,11 +124,12 @@ impl McpHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = Arc::clone(&step_context.turn); let payload = match payload { ToolPayload::Function { arguments } => arguments, @@ -143,7 +144,7 @@ impl McpHandler { // TODO(sayan): Use StepContext for MCP file arguments when MCP follows dynamic environments. let result = handle_mcp_tool_call( Arc::clone(&session), - &turn, + &step_context, call_id.clone(), self.tool_info.server_name.clone(), self.tool_info.tool.name.to_string(), diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index ce762a0f02..8c7f21fd17 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -53,11 +53,13 @@ impl ListMcpResourceTemplatesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourceTemplatesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resource_templates(&server_name, params) .await .map_err(|err| { @@ -108,10 +110,7 @@ impl ListMcpResourceTemplatesHandler { )); } - let templates = session - .services - .mcp_connection_manager - .load_full() + let templates = manager .list_all_resource_templates(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index 564fb440fd..f5e9bae7e3 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -53,11 +53,13 @@ impl ListMcpResourcesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourcesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resources(&server_name, params) .await .map_err(|err| { @@ -106,10 +108,7 @@ impl ListMcpResourcesHandler { )); } - let resources = session - .services - .mcp_connection_manager - .load_full() + let resources = manager .list_all_resources(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 3909011eb0..5e9b85241c 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -52,11 +52,13 @@ impl ReadMcpResourceHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -84,7 +86,7 @@ impl ReadMcpResourceHandler { let payload_result: Result = async { ensure_model_can_access_mcp_server(turn.as_ref(), &server)?; - let result = session + let result = manager .read_resource(&server, ReadResourceRequestParams::new(uri.clone())) .await .map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index ce7ce10da8..e44ed41fa4 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::sync::Arc; use codex_analytics::PluginInstallRequestSource; use codex_analytics::PluginInstallRequested; @@ -94,10 +95,12 @@ impl RequestPluginInstallHandler { let ToolInvocation { payload, session, - turn, + step_context, call_id, .. } = invocation; + let turn = Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -227,7 +230,8 @@ impl RequestPluginInstallHandler { let auth = session.services.auth_manager.auth().await; let completed = if user_confirmed { - verify_request_plugin_install_completed(&session, &turn, &tool, auth.as_ref()).await + verify_request_plugin_install_completed(&session, &turn, manager, &tool, auth.as_ref()) + .await } else { false }; @@ -345,13 +349,14 @@ fn disabled_install_request(tool: &DiscoverableTool) -> ToolSuggestDisabledTool async fn verify_request_plugin_install_completed( session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, tool: &DiscoverableTool, auth: Option<&codex_login::CodexAuth>, ) -> bool { match tool { DiscoverableTool::Connector(connector) => refresh_missing_requested_connectors( - session, turn, + manager, auth, std::slice::from_ref(&connector.id), connector.id.as_str(), @@ -370,8 +375,8 @@ async fn verify_request_plugin_install_completed( plugin.id.as_str(), ), refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -393,8 +398,8 @@ async fn verify_request_plugin_install_completed( session.services.plugins_manager.as_ref(), ); let _ = refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -435,8 +440,8 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { } async fn refresh_missing_requested_connectors( - session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, auth: Option<&codex_login::CodexAuth>, expected_connector_ids: &[String], tool_id: &str, @@ -445,7 +450,6 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let manager = session.services.mcp_connection_manager.load_full(); let mcp_tools = manager.list_all_tools().await; let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index db1209cc03..6f399fd5fc 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -110,8 +110,8 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow let step_context = StepContext::for_test(Arc::clone(&turn)); let mcp_tools = session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await; let router = ToolRouter::from_context( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 6fc5adcb11..498a690c58 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -686,10 +686,12 @@ async fn environment_tools_follow_the_step_context() { let environments = turn.environments.clone(); turn.environments.turn_environments.clear(); + let turn = Arc::new(turn); let step_context = Arc::new(StepContext::new( - Arc::new(turn), + Arc::clone(&turn), environments, Vec::new(), + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs index 7df18727b9..620b6c734c 100644 --- a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -16,7 +16,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_mcp_server; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> { +async fn refresh_keeps_superseded_mcp_server_alive_for_in_flight_calls() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -84,7 +84,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> "sync", Some(serde_json::json!({ "barrier": barrier, - "sleep_after_ms": 30_000 + "sleep_after_ms": 300_000 })), /*meta*/ None, ) @@ -119,9 +119,16 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> let replacement_pid = wait_for_pid_file(&pid_file).await?; assert_ne!(replacement_pid, superseded_pid); + assert!(process_is_alive(&superseded_pid)?); + long_call.abort(); + assert!( + long_call + .await + .expect_err("call should be aborted") + .is_cancelled() + ); wait_for_process_exit(&superseded_pid).await?; assert!(process_is_alive(&replacement_pid)?); - assert!(long_call.await?.is_err()); fixture.codex.shutdown_and_wait().await?; wait_for_process_exit(&replacement_pid).await From 5044062704af7667e460cc401221a081054ef3ab Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 25 Jun 2026 17:01:43 -0700 Subject: [PATCH 30/33] ci: narrow Windows test skips (#30134) ## Why The Windows cross-build skip used the broad `powershell` substring, which hid unrelated Windows tests. Narrowing it exposed the same ConPTY Ctrl-C timeout that is breaking `main`; that test is not reliable in either cross-built or native Windows Bazel CI yet. ## What changed - scope the cross-build PowerShell carve-out to the dedicated parser-process test module - exclude the exact ConPTY Ctrl-C test from Bazel CI while leaving local Windows runs enabled - repeat the exact exclusion in the cross-build config because it replaces the base skip list ## Manual validation - `just test-github-scripts` - queried the PTY test target under both `ci-windows` and `ci-windows-cross` - verified the matcher excludes parser-process and ConPTY tests without excluding unrelated PowerShell tests - [Windows shard 4/4](https://github.com/openai/codex/actions/runs/28204844286/job/83553063859) reproduced the `main` ConPTY timeout before the exact CI-only exclusion was applied --- .bazelrc | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/.bazelrc b/.bazelrc index c58667b46e..de01a87443 100644 --- a/.bazelrc +++ b/.bazelrc @@ -107,8 +107,9 @@ common:ci-bazel --config=ci common:ci-bazel --build_metadata=TAG_workflow=bazel # Keep code-mode integration cases out of ordinary Bazel legs. The # Windows-cross config below re-enables them after generating its Windows V8 -# snapshot on the Windows runner. -common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode:: +# snapshot on the Windows runner. Also skip the ConPTY Ctrl-C integration test +# until Windows Bazel CI can reliably interrupt foreground processes. +common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode::,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child # Shared config for Bazel-backed Rust linting. build:clippy --aspects=@rules_rust//rust:defs.bzl%rust_clippy_aspect @@ -191,11 +192,12 @@ common:ci-windows-cross --strategy=TestRunner=local common:ci-windows-cross --strategy=V8Mksnapshot=local common:ci-windows-cross --local_test_jobs=4 common:ci-windows-cross --test_env=RUST_TEST_THREADS=1 -# Native Windows CI still covers the PowerShell tests. The cross-built gnullvm -# binaries currently hang in PowerShell AST parser tests when those binaries are -# run on the Windows runner. Keep V8-backed code-mode tests enabled except for -# the hidden dynamic-tool callback test, which currently times out on Windows. -common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=powershell,suite::code_mode::code_mode_can_call_hidden_dynamic_tools +# Native Windows CI still covers the PowerShell parser-process tests. The +# cross-built gnullvm binaries currently hang in those tests when run on the +# Windows runner. Keep V8-backed code-mode tests enabled except for the hidden +# dynamic-tool callback test, which currently times out on Windows. This config +# replaces the base skip list, so repeat the exact ConPTY Ctrl-C test exclusion. +common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=command_safety::powershell_parser::tests::,suite::code_mode::code_mode_can_call_hidden_dynamic_tools,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child common:ci-windows-cross --platforms=//:windows_x86_64_gnullvm common:ci-windows-cross --extra_execution_platforms=//:rbe,//:windows_x86_64_msvc common:ci-windows-cross --extra_toolchains=//:windows_gnullvm_tests_on_msvc_host_toolchain From 3095ea9c3d155bfc89197d2628eb818a55c2755d Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 01:36:44 +0100 Subject: [PATCH 31/33] Project selected plugin runtime by environment availability (#30093) ## Why Selected plugin metadata is stable, but MCP processes are live runtime state. They need different lifetimes: - the MCP extension caches manifest, MCP, and connector declarations for each stable selected root; - each model step projects that cached metadata through the roots that resolved as ready for that exact step; - the MCP manager is rebuilt only when that availability projection changes. This matches executor skills: both features consume the same resolved step roots instead of inferring readiness from the turn's selected environments. ## Behavior ```text E1 not ready for this step -> no E1 MCP servers or connectors -> cached plugin metadata stays in ext/mcp E1 becomes ready -> reuse cached metadata -> publish one MCP runtime containing E1 capabilities same ready roots on the next step -> reuse the exact runtime; no rediscovery and no MCP restart resume -> create new extension thread state and a new MCP runtime ``` All model-facing consumers use the same step snapshot: ```text resolved selected roots | v extension MCP/connector projection | v { MCP config, connector snapshot, MCP manager } | +-> advertise model tools +-> build app/connector tools +-> execute MCP calls ``` ## Cache contract The existing MCP extension owns a cache keyed by the full `SelectedCapabilityRoot`: ```rust let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default); ``` The cache lives with extension thread state. Environment availability filters projection but does not invalidate metadata. Resume creates new thread state. There is no file watcher or executor generation because contents behind a stable environment/root are assumed stable. ## What changes - Keeps executor plugin discovery and cached metadata in `ext/mcp`. - Caches MCP and connector declarations together per selected root. - Uses the step's already-resolved capability roots, including lazy environments that are not turn environments. - Reuses the current MCP runtime when the ready-root projection is unchanged. - Uses the same step MCP manager and connector snapshot for model-visible tools and execution. - Resolves direct thread-scoped MCP requests from the current selected-root projection. ## Deliberately out of scope - `app/list` remains based on the latest global host-plugin state; this PR does not make its response or notifications thread-specific. - `required = true` startup semantics do not apply to delayed executor MCP activation. - No filesystem/content invalidation. - No transport-disconnect watcher. - No executor generations or environment replacement semantics. - No client sharing across complete manager replacements. ## Stack 1. Extension-owned World State sections. 2. Project executor skills through World State. 3. Pin one MCP runtime to each model step. 4. **This PR:** project selected MCP and connector state from extension-owned metadata. 5. Integration coverage for selected capability availability and resume. ## Verification - `selected_plugin_servers_use_managed_requirements_for_the_selected_root_id` - The stacked integration PR covers unavailable to ready activation, unchanged-runtime reuse, skills, MCP tools, connector attribution, and cold resume. --- codex-rs/Cargo.lock | 1 + codex-rs/Cargo.toml | 1 + .../src/request_processors/mcp_processor.rs | 23 +- .../request_processors/thread_processor.rs | 1 - .../app-server/tests/suite/v2/executor_mcp.rs | 1 - codex-rs/core/src/codex_thread.rs | 20 +- codex-rs/core/src/mcp.rs | 45 +++- codex-rs/core/src/prompt_debug.rs | 4 +- codex-rs/core/src/session/mcp.rs | 196 ++++++++++++++---- codex-rs/core/src/session/mcp_runtime.rs | 8 + codex-rs/core/src/session/mod.rs | 19 +- codex-rs/core/src/session/session.rs | 10 +- codex-rs/core/src/session/tests.rs | 3 + codex-rs/core/src/session/turn.rs | 27 +-- codex-rs/core/src/state/service.rs | 14 +- codex-rs/core/src/thread_manager_tests.rs | 24 ++- .../ext/extension-api/src/contributors/mcp.rs | 47 ++++- codex-rs/ext/mcp/Cargo.toml | 1 + codex-rs/ext/mcp/src/executor_plugin.rs | 154 ++++++++++---- codex-rs/ext/mcp/src/lib.rs | 7 - codex-rs/ext/mcp/tests/executor_plugin_mcp.rs | 12 +- 21 files changed, 455 insertions(+), 163 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b147022825..75c02adfb5 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3359,6 +3359,7 @@ name = "codex-mcp-extension" version = "0.0.0" dependencies = [ "codex-config", + "codex-connectors-extension", "codex-core", "codex-core-plugins", "codex-exec-server", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index ef17cb9ce5..0824f7db7b 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -166,6 +166,7 @@ codex-code-mode-protocol = { path = "code-mode-protocol" } codex-home = { path = "codex-home" } codex-config = { path = "config" } codex-connectors = { path = "connectors" } +codex-connectors-extension = { path = "ext/connectors" } codex-context-fragments = { path = "context-fragments" } codex-core = { path = "core" } codex-core-api = { path = "core-api" } diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index abefd26015..ac2125c2cf 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -124,7 +124,7 @@ impl McpRequestProcessor { let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let runtime = thread.current_mcp_runtime(); + let runtime = thread.current_mcp_runtime().await; (runtime.config().clone(), runtime.runtime_context().clone()) } None => { @@ -246,14 +246,21 @@ impl McpRequestProcessor { let mcp_manager = self.thread_manager.mcp_manager(); let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; - let mcp_config = match thread { - Some(thread) => thread.runtime_mcp_config(&config).await, - None => mcp_manager.runtime_config(&config).await, + let (mcp_config, runtime_context) = match thread { + Some(thread) => { + let mcp_config = thread.runtime_mcp_config(&config).await; + let runtime = thread.current_mcp_runtime().await; + (mcp_config, runtime.runtime_context().clone()) + } + None => { + let mcp_config = mcp_manager.runtime_config(&config).await; + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) + } }; - let runtime_context = McpRuntimeContext::new( - self.thread_manager.environment_manager(), - config.cwd.to_path_buf(), - ); tokio::spawn(async move { Self::list_mcp_server_status_task( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 0fc37d1a4f..39283bb35c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1155,7 +1155,6 @@ impl ThreadRequestProcessor { let mut thread_extension_init = ExtensionDataInit::new(); if !selected_capability_roots.is_empty() { thread_extension_init.insert(selected_capability_roots); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_extension_init); } let create_thread_started_at = std::time::Instant::now(); let NewThread { diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 67d97bc248..aa650c3ecf 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -191,7 +191,6 @@ HTTP_PROXY = {http_proxy} ) .await?; - std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?; let config_path = codex_home.path().join("config.toml"); let mut config = std::fs::read_to_string(&config_path)?; config.push_str(&format!( diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 46b38af151..90555db911 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -595,8 +595,14 @@ impl CodexThread { } /// Returns the exact MCP config, environment bindings, and manager most recently published. - pub fn current_mcp_runtime(&self) -> Arc { - self.codex.session.services.latest_mcp_runtime() + pub async fn current_mcp_runtime(&self) -> Arc { + let turn_context = self.codex.session.new_default_turn().await; + self.codex + .session + .capture_step_context(turn_context) + .await + .mcp + .clone() } pub fn multi_agent_version(&self) -> Option { @@ -620,8 +626,9 @@ impl CodexThread { uri: &str, ) -> anyhow::Result { let result = self - .codex - .session + .current_mcp_runtime() + .await + .manager_arc() .read_resource(server, ReadResourceRequestParams::new(uri)) .await?; @@ -635,8 +642,9 @@ impl CodexThread { arguments: Option, meta: Option, ) -> anyhow::Result { - self.codex - .session + self.current_mcp_runtime() + .await + .manager_arc() .call_tool(server, tool, arguments, meta) .await } diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 658c5471ee..45ca0d563d 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,7 +3,10 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; +use codex_connectors::ConnectorSnapshot; +use codex_connectors::PluginConnectorSource; use codex_core_plugins::PluginsManager; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_extension_api::McpServerContribution; @@ -18,6 +21,7 @@ use codex_mcp::McpServerRegistration; use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; +use codex_plugin::AppConnectorId; const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps"; @@ -69,28 +73,32 @@ impl McpManager { /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { - self.runtime_config_with_context(config, /*thread_init*/ None) + self.runtime_config_with_context(McpServerContributionContext::global(config)) .await } - pub(crate) async fn runtime_config_for_thread( + pub(crate) async fn runtime_config_for_step( &self, config: &Config, thread_init: &ExtensionDataInit, + thread_store: &ExtensionData, + available_environment_ids: &[String], ) -> McpConfig { - self.runtime_config_with_context(config, Some(thread_init)) - .await + self.runtime_config_with_context(McpServerContributionContext::for_step( + config, + thread_init, + thread_store, + available_environment_ids, + )) + .await } async fn runtime_config_with_context( &self, - config: &Config, - thread_init: Option<&ExtensionDataInit>, + context: McpServerContributionContext<'_, Config>, ) -> McpConfig { - let context = match thread_init { - Some(thread_init) => McpServerContributionContext::for_thread(config, thread_init), - None => McpServerContributionContext::global(config), - }; + let config = context.config(); + let mut selected_plugin_connector_sources = Vec::new(); let mut selected_plugin_registrations = Vec::new(); let mut overlays = Vec::new(); // A contributor can emit multiple ordered actions, so order each action globally rather @@ -121,6 +129,17 @@ impl McpManager { *config, ), ), + McpServerContribution::SelectedPluginConnectors { + plugin_id, + plugin_display_name, + connector_ids, + } => selected_plugin_connector_sources.push( + PluginConnectorSource::from_connector_ids( + plugin_id, + plugin_display_name, + connector_ids.into_iter().map(AppConnectorId), + ), + ), McpServerContribution::Remove { name } => { overlays.push(OrderedMcpOverlay::Remove { contributor_id: contributor.id(), @@ -186,6 +205,12 @@ impl McpManager { ); } mcp_config.mcp_server_catalog = catalog; + mcp_config.connector_snapshot = + mcp_config + .connector_snapshot + .merged_with(&ConnectorSnapshot::from_plugin_sources( + selected_plugin_connector_sources, + )); mcp_config } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index a866116a24..cac9bb41ae 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -64,7 +64,7 @@ pub async fn build_prompt_input( ); let thread = thread_manager.start_thread(config).await?; - let output = build_prompt_input_from_session(thread.thread.codex.session.as_ref(), input).await; + let output = build_prompt_input_from_session(&thread.thread.codex.session, input).await; let shutdown = thread.thread.shutdown_and_wait().await; let _removed = thread_manager.remove_thread(&thread.thread_id).await; @@ -73,7 +73,7 @@ pub async fn build_prompt_input( } pub(crate) async fn build_prompt_input_from_session( - sess: &Session, + sess: &Arc, input: Vec, ) -> CodexResult> { let turn_context = sess.new_default_turn().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index b9e87f0c89..50921f1591 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,7 +1,9 @@ use super::*; +use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; use codex_mcp::ElicitationReviewerHandle; +use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY as MCP_ELICITATION_APPROVAL_KIND_KEY; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_MCP_TOOL_CALL as MCP_ELICITATION_APPROVAL_KIND_MCP_TOOL_CALL; @@ -75,9 +77,20 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer { impl Session { pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig { + let environments = self.services.turn_environments.snapshot().await; + let selected_capability_roots = self + .resolve_selected_capability_roots_for_step(&environments) + .await; + let available_environment_ids = + Self::available_selected_environment_ids(&selected_capability_roots); self.services .mcp_manager - .runtime_config_for_thread(config, &self.services.mcp_thread_init) + .runtime_config_for_step( + config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await } @@ -88,6 +101,62 @@ impl Session { codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime comparison and publication must remain serialized" + )] + pub(crate) async fn mcp_runtime_for_step( + self: &Arc, + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, + selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Arc { + let available_environment_ids = + Self::available_selected_environment_ids(selected_capability_roots); + let current = self.services.latest_mcp_runtime(); + if current.available_environment_ids() == available_environment_ids { + return current; + } + + let _guard = self.services.mcp_projection_lock.lock().await; + let current = self.services.latest_mcp_runtime(); + if current.available_environment_ids() == available_environment_ids { + return current; + } + let mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + &turn_context.config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) + .await; + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + environments, + &available_environment_ids, + Some(self.mcp_elicitation_reviewer()), + ) + .await + } + + pub(crate) async fn resolve_selected_capability_roots_for_step( + &self, + environments: &TurnEnvironmentSnapshot, + ) -> Vec { + self.services + .turn_environments + .environment_manager() + .resolve_selected_capability_roots( + &self.services.selected_capability_roots, + &environments.captured_environments(), + ) + .await + } + pub(crate) fn mcp_elicitation_reviewer(self: &Arc) -> ElicitationReviewerHandle { Arc::new(GuardianMcpElicitationReviewer::new(self)) } @@ -207,51 +276,22 @@ impl Session { .await } - pub async fn read_resource( - &self, - server: &str, - params: ReadResourceRequestParams, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .read_resource(server, params) - .await - } - - pub async fn call_tool( - &self, - server: &str, - tool: &str, - arguments: Option, - meta: Option, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .call_tool(server, tool, arguments, meta) - .await - } - async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, - mut mcp_config: McpConfig, - configured_mcp_servers: HashMap, + mcp_config: McpConfig, + environments: &TurnEnvironmentSnapshot, + available_environment_ids: &[String], elicitation_reviewer: Option, - ) { - mcp_config.mcp_server_catalog = mcp_config - .mcp_server_catalog - .with_materialized_servers(configured_mcp_servers); - let mcp_config = Arc::new(mcp_config); + ) -> Arc { let auth = self.services.auth_manager.auth().await; + let mcp_config = Arc::new(mcp_config); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. - let cwd = turn_context - .environments + let cwd = environments .primary() .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) .map(|cwd| cwd.to_path_buf()) @@ -305,10 +345,18 @@ impl Session { refreshed_manager .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); } - self.services - .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); + self.services.publish_mcp_runtime( + mcp_config, + mcp_runtime_context, + available_environment_ids.to_vec(), + refreshed_manager, + ) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime refresh and publication must remain serialized" + )] pub(crate) async fn refresh_mcp_servers_if_requested( &self, turn_context: &TurnContext, @@ -365,9 +413,33 @@ impl Session { return; } - let mcp_config = self.runtime_mcp_config(&refresh_config).await; - self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + let _guard = self.services.mcp_projection_lock.lock().await; + let available_environment_ids = self + .services + .latest_mcp_runtime() + .available_environment_ids() + .to_vec(); + let mut mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + &refresh_config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await; + mcp_config.mcp_server_catalog = mcp_config + .mcp_server_catalog + .with_materialized_servers(mcp_servers); + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + &turn_context.environments, + &available_environment_ids, + elicitation_reviewer, + ) + .await; } pub(crate) async fn set_openai_form_elicitation_support( @@ -398,16 +470,54 @@ impl Session { Ok(()) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime refresh and publication must remain serialized" + )] pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, refresh_config: &Config, elicitation_reviewer: Option, ) { - let mcp_config = self.runtime_mcp_config(refresh_config).await; - let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); - self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + let _guard = self.services.mcp_projection_lock.lock().await; + let available_environment_ids = self + .services + .latest_mcp_runtime() + .available_environment_ids() + .to_vec(); + let mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + refresh_config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await; + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + &turn_context.environments, + &available_environment_ids, + elicitation_reviewer, + ) + .await; + } + + fn available_selected_environment_ids( + selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Vec { + let mut available = Vec::new(); + for root in selected_capability_roots { + let CapabilityRootLocation::Environment { environment_id, .. } = + &root.selected_root().location; + if !available.contains(environment_id) { + available.push(environment_id.clone()); + } + } + available } #[cfg(test)] diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index dc5ea011ca..34cf07cc35 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -10,6 +10,7 @@ pub struct McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, } impl McpRuntimeSnapshot { @@ -17,11 +18,13 @@ impl McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, ) -> Self { Self { config, manager, runtime_context, + available_environment_ids, } } @@ -41,6 +44,10 @@ impl McpRuntimeSnapshot { &self.runtime_context } + pub(crate) fn available_environment_ids(&self) -> &[String] { + &self.available_environment_ids + } + #[cfg(test)] pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { use codex_exec_server::EnvironmentManager; @@ -81,6 +88,7 @@ impl McpRuntimeSnapshot { Arc::new(mcp_config), Arc::new(manager), runtime_context, + Vec::new(), )) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9421e07a66..0008257542 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -100,7 +100,6 @@ use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; -use codex_protocol::mcp::CallToolResult; use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::BaseInstructions; @@ -158,8 +157,6 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ReadResourceRequestParams; -use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; use serde_json::Value; use tokio::sync::Mutex; @@ -2825,7 +2822,7 @@ impl Session { /// `run_turn` and pass the result down; standalone request or history boundaries may capture /// their own step. pub(crate) async fn capture_step_context( - &self, + self: &Arc, turn_context: Arc, ) -> Arc { let deferred_executor_enabled = turn_context @@ -2846,15 +2843,15 @@ impl Session { } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; let selected_capability_roots = self - .services - .turn_environments - .environment_manager() - .resolve_selected_capability_roots( - &self.services.selected_capability_roots, - &environments.captured_environments(), + .resolve_selected_capability_roots_for_step(&environments) + .await; + let mcp = self + .mcp_runtime_for_step( + turn_context.as_ref(), + &environments, + &selected_capability_roots, ) .await; - let mcp = self.services.latest_mcp_runtime(); Arc::new(StepContext::new( turn_context, environments, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2c9e57036d..0e1609207a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -661,6 +661,7 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; + let thread_extension_data_for_mcp = &thread_extension_data; let mcp_runtime_cwd = session_configuration .environment_selections() .first() @@ -673,7 +674,12 @@ impl Session { let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp - .runtime_config_for_thread(&config_for_mcp, mcp_thread_init_for_startup) + .runtime_config_for_step( + &config_for_mcp, + mcp_thread_init_for_startup, + thread_extension_data_for_mcp, + /*available_environment_ids*/ &[], + ) .await; let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); @@ -1039,6 +1045,7 @@ impl Session { // setup is straightforward enough and performs well. mcp_connection_manager, mcp_runtime: arc_swap::ArcSwapOption::empty(), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1210,6 +1217,7 @@ impl Session { .install_mcp_connection_manager( Arc::new(mcp_config), mcp_runtime_context, + /*available_environment_ids*/ Vec::new(), mcp_connection_manager, ) .await?; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d0aac04249..57f1ce9081 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5372,6 +5372,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7446,6 +7447,7 @@ where let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7603,6 +7605,7 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( #[tokio::test] async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); let turn_context = Arc::new(turn_context); let old_runtime = session.services.latest_mcp_runtime(); let step_context = session diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2cb1f2cada..7b509ae42e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -171,8 +171,13 @@ pub(crate) async fn run_turn( .record_context_updates_and_set_reference_context_item(first_step_context.as_ref()) .await; - let Some((injection_items, explicitly_enabled_connectors)) = - build_skills_and_plugins(&sess, turn_context.as_ref(), &input, &cancellation_token).await + let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins( + &sess, + first_step_context.as_ref(), + &input, + &cancellation_token, + ) + .await else { return Ok(None); }; @@ -504,10 +509,11 @@ async fn run_hooks_and_record_inputs( #[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, - turn_context: &TurnContext, + step_context: &StepContext, input: &[TurnInput], cancellation_token: &CancellationToken, ) -> Option<(Vec, HashSet)> { + let turn_context = step_context.turn.as_ref(); // Guardian input embeds the parent transcript as untrusted evidence. Do not interpret skill or // plugin mentions from that generated prompt as requests to inject additional instructions. if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { @@ -538,17 +544,14 @@ async fn build_skills_and_plugins( // enabled plugins, then converted into turn-scoped guidance below. let mentioned_plugins = collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries()); - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = step_context.mcp.config().connector_snapshot.clone(); let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { // Plugin mentions need raw MCP/app inventory even when app tools // are normally hidden so we can describe the plugin's currently // usable capabilities for this turn. - match sess - .services - .mcp_connection_manager - .load_full() + match step_context + .mcp + .manager_arc() .list_all_tools() .or_cancel(cancellation_token) .await @@ -1189,9 +1192,7 @@ pub(crate) async fn built_tools( .plugins_for_config(&turn_context.config.plugins_config_input()) .instrument(trace_span!("built_tools.load_plugins")) .await; - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = step_context.mcp.config().connector_snapshot.clone(); let apps_enabled = turn_context.apps_enabled(); let accessible_connectors = diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 1bee45cce0..a611b94887 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -51,6 +51,8 @@ pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, /// The latest atomically published MCP config and manager pair. pub(crate) mcp_runtime: ArcSwapOption, + /// Serializes environment-driven runtime rebuilds. + pub(crate) mcp_projection_lock: Mutex<()>, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -106,9 +108,11 @@ impl SessionServices { &self, config: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, manager: McpConnectionManager, ) -> Result<()> { - let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + let runtime = + self.publish_mcp_runtime(config, runtime_context, available_environment_ids, manager); runtime.manager().validate_required_servers().await } @@ -116,13 +120,19 @@ impl SessionServices { &self, config: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, manager: McpConnectionManager, ) -> Arc { let manager = Arc::new(manager); // Publish the manager for legacy resource clients first. Once the paired snapshot is // visible, every model-scoped consumer observes this exact manager. self.mcp_connection_manager.store(Arc::clone(&manager)); - let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); + let runtime = Arc::new(McpRuntimeSnapshot::new( + config, + manager, + runtime_context, + available_environment_ids, + )); self.mcp_runtime.store(Some(Arc::clone(&runtime))); runtime } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 9bccf8b6c1..5148f16259 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -584,8 +584,28 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() }) .await .expect("start second thread"); - let first_resolved = first_thread.thread.runtime_mcp_config(&config).await; - let second_resolved = second_thread.thread.runtime_mcp_config(&config).await; + let first_session = &first_thread.thread.codex.session; + let first_resolved = first_session + .services + .mcp_manager + .runtime_config_for_step( + &config, + &first_session.services.mcp_thread_init, + &first_session.services.thread_extension_data, + /*available_environment_ids*/ &[], + ) + .await; + let second_session = &second_thread.thread.codex.session; + let second_resolved = second_session + .services + .mcp_manager + .runtime_config_for_step( + &config, + &second_session.services.mcp_thread_init, + &second_session.services.thread_extension_data, + /*available_environment_ids*/ &[], + ) + .await; assert_eq!( *lifecycle_observed diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 91e9cb459c..499657f556 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,17 +1,22 @@ use codex_config::McpServerConfig; +use crate::ExtensionData; use crate::ExtensionDataInit; /// Input supplied while resolving MCP server contributions. /// -/// Thread-scoped implementations can read the immutable host-seeded inputs -/// through [`Self::thread_init`]. Implementations should not retain borrowed -/// context after contribution completes. +/// Thread-scoped implementations can read stable host inputs through [`Self::thread_init`] and +/// keep their cache in [`Self::thread_store`]. Implementations should not retain borrowed context +/// after contribution completes. pub struct McpServerContributionContext<'a, C> { /// Host configuration visible during MCP resolution. config: &'a C, - /// Initial inputs for the active thread, when resolution is thread-scoped. + /// Extension-owned data for the active thread, when resolution is thread-scoped. + thread_store: Option<&'a ExtensionData>, + /// Stable host inputs for the active thread, when resolution is thread-scoped. thread_init: Option<&'a ExtensionDataInit>, + /// Environment IDs whose selected roots may contribute to this exact step. + available_environment_ids: Option<&'a [String]>, } impl Clone for McpServerContributionContext<'_, C> { @@ -27,15 +32,24 @@ impl<'a, C> McpServerContributionContext<'a, C> { pub fn global(config: &'a C) -> Self { Self { config, + thread_store: None, thread_init: None, + available_environment_ids: None, } } - /// Creates context for one active thread runtime. - pub fn for_thread(config: &'a C, thread_init: &'a ExtensionDataInit) -> Self { + /// Creates context for one model step using only currently available environments. + pub fn for_step( + config: &'a C, + thread_init: &'a ExtensionDataInit, + thread_store: &'a ExtensionData, + available_environment_ids: &'a [String], + ) -> Self { Self { config, + thread_store: Some(thread_store), thread_init: Some(thread_init), + available_environment_ids: Some(available_environment_ids), } } @@ -44,10 +58,23 @@ impl<'a, C> McpServerContributionContext<'a, C> { self.config } - /// Returns the frozen initial inputs when resolving for a running thread. + /// Returns extension-owned state when resolving for a running thread. + pub fn thread_store(&self) -> Option<&'a ExtensionData> { + self.thread_store + } + + /// Returns stable host inputs when resolving for a running thread. pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> { self.thread_init } + + /// Returns the exact environment availability projection for a model step. + /// + /// `Some` means contributors must omit selected roots whose environment ID is absent from the + /// slice. Global resolution returns `None` because it has no thread environments. + pub fn available_environment_ids(&self) -> Option<&'a [String]> { + self.available_environment_ids + } } /// One extension-owned overlay for the runtime MCP server configuration. @@ -66,6 +93,12 @@ pub enum McpServerContribution { selection_order: usize, config: Box, }, + /// Adds connector IDs declared by a plugin selected for this thread. + SelectedPluginConnectors { + plugin_id: String, + plugin_display_name: String, + connector_ids: Vec, + }, /// Removes a named MCP server. Remove { name: String }, } diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index b98c0905db..d918bd5bd9 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -16,6 +16,7 @@ workspace = true codex-core = { workspace = true } codex-core-plugins = { workspace = true } codex-config = { workspace = true } +codex-connectors-extension = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 69d1a47b31..32637b4f47 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -1,44 +1,47 @@ +use codex_connectors_extension::ExecutorPluginConnectorProvider; use codex_core::config::Config; use codex_core_plugins::ExecutorPluginProvider; use codex_exec_server::EnvironmentManager; -use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; +use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; use std::sync::Arc; -use tokio::sync::OnceCell; +use std::sync::Mutex; use self::provider::ExecutorPluginMcpProvider; mod provider; -/// Frozen MCP declarations for one selected package. +/// Frozen MCP and connector declarations for one selected package. /// /// Each server config retains the stable logical environment ID. Reconnection may replace the /// concrete environment instance without changing that authority. #[derive(Clone)] -struct SelectedPluginMcpServers { +struct SelectedPluginMetadata { plugin_id: String, plugin_display_name: String, - selection_order: usize, servers: Vec<(String, codex_config::McpServerConfig)>, + connector_ids: Vec, } #[derive(Default)] pub(crate) struct SelectedExecutorPluginMcpState { - snapshot: OnceCell>, + cache: Mutex>, } -pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { - thread_init.insert(SelectedExecutorPluginMcpState::default()); +struct CachedSelectedRoot { + root: SelectedCapabilityRoot, + metadata: Option, } pub(crate) struct SelectedExecutorPluginMcpContributor { plugin_provider: ExecutorPluginProvider, mcp_provider: ExecutorPluginMcpProvider, + connector_provider: ExecutorPluginConnectorProvider, } impl SelectedExecutorPluginMcpContributor { @@ -46,46 +49,87 @@ impl SelectedExecutorPluginMcpContributor { Self { plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)), mcp_provider: ExecutorPluginMcpProvider, + connector_provider: ExecutorPluginConnectorProvider, } } - async fn resolve_snapshot( + /// Returns metadata for one stable selected root. + /// + /// Successful resolution, including a root that is not a plugin or declares no capabilities, + /// is cached until the thread state is dropped. Environment availability never invalidates + /// this cache; it only controls whether the cached metadata is projected into a model step. + async fn metadata_for_root( &self, - selected_roots: &[SelectedCapabilityRoot], - ) -> Vec { - let mut snapshot = Vec::new(); + state: &SelectedExecutorPluginMcpState, + selected_root: &SelectedCapabilityRoot, + ) -> Option { + if let Some(cached) = state + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find(|cached| cached.root == *selected_root) + { + return cached.metadata.clone(); + } - for (selection_order, selected_root) in selected_roots.iter().enumerate() { - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { - Ok(Some(plugin)) => plugin, - Ok(None) => continue, - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to resolve selected executor plugin for MCP discovery" - ); - continue; - } - }; - match self.mcp_provider.load(&plugin).await { - Ok(servers) => snapshot.push(SelectedPluginMcpServers { - plugin_id: plugin.plugin().selected_root_id().to_string(), - plugin_display_name: plugin.plugin().manifest().display_name().to_string(), - selection_order, - servers, - }), - Err(err) => { + let plugin = match self.plugin_provider.resolve_bound(selected_root).await { + Ok(plugin) => plugin, + Err(err) => { + tracing::warn!( + selected_root = selected_root.id, + error = %err, + "failed to resolve selected executor plugin" + ); + return None; + } + }; + let metadata = match plugin { + Some(plugin) => { + let servers = self.mcp_provider.load(&plugin).await.unwrap_or_else(|err| { tracing::warn!( selected_root = selected_root.id, error = %err, "failed to load selected executor plugin MCP servers" ); - } + Vec::new() + }); + let connector_ids = self + .connector_provider + .load(&plugin) + .await + .unwrap_or_else(|err| { + tracing::warn!( + selected_root = selected_root.id, + error = %err, + "failed to load selected executor plugin connectors" + ); + Vec::new() + }) + .into_iter() + .map(|declaration| declaration.connector_id.0) + .collect(); + Some(SelectedPluginMetadata { + plugin_id: plugin.plugin().selected_root_id().to_string(), + plugin_display_name: plugin.plugin().manifest().display_name().to_string(), + servers, + connector_ids, + }) } + None => None, + }; + let mut cache = state + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cache.iter().find(|cached| cached.root == *selected_root) { + return cached.metadata.clone(); } - - snapshot + cache.push(CachedSelectedRoot { + root: selected_root.clone(), + metadata: metadata.clone(), + }); + metadata } } @@ -102,20 +146,31 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { let Some(thread_init) = context.thread_init() else { return Vec::new(); }; + let Some(thread_store) = context.thread_store() else { + return Vec::new(); + }; let Some(selected_roots) = thread_init.get::>() else { return Vec::new(); }; - let Some(state) = thread_init.get::() else { - tracing::warn!("selected executor plugin MCP state was not initialized"); - return Vec::new(); - }; - let snapshot = state - .snapshot - .get_or_init(|| self.resolve_snapshot(selected_roots.as_ref())) - .await; + let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default); let mut contributions = Vec::new(); - for plugin in snapshot { + for (selection_order, selected_root) in selected_roots.iter().enumerate() { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + if context + .available_environment_ids() + .is_some_and(|available| { + !available + .iter() + .any(|available| available == environment_id) + }) + { + continue; + } + let Some(plugin) = self.metadata_for_root(&state, selected_root).await else { + continue; + }; let mut servers = plugin.servers.iter().cloned().collect::>(); context .config() @@ -127,10 +182,17 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { name, plugin_id: plugin.plugin_id.clone(), plugin_display_name: plugin.plugin_display_name.clone(), - selection_order: plugin.selection_order, + selection_order, config: Box::new(config), } })); + if !plugin.connector_ids.is_empty() { + contributions.push(McpServerContribution::SelectedPluginConnectors { + plugin_id: plugin.plugin_id, + plugin_display_name: plugin.plugin_display_name, + connector_ids: plugin.connector_ids, + }); + } } contributions diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index e32316ac36..9c22bad5a0 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -51,10 +51,3 @@ pub fn install_executor_plugins( executor_plugin::SelectedExecutorPluginMcpContributor::new(environment_manager), )); } - -/// Seeds the per-thread snapshot used by selected executor plugin MCP discovery. -pub fn initialize_executor_plugin_thread_data( - thread_init: &mut codex_extension_api::ExtensionDataInit, -) { - executor_plugin::seed_thread_state(thread_init); -} diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs index 7bdbc83d62..65a4248a32 100644 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs @@ -3,6 +3,7 @@ use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; @@ -109,12 +110,15 @@ async fn selected_plugin_contributions( path: PathUri::from_host_native_path(plugin_root)?, }, }]); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_init); + let thread_store = ExtensionData::new_with_init("test-thread", thread_init.clone()); + let available_environment_ids = vec![LOCAL_ENVIRONMENT_ID.to_string()]; Ok(registry.mcp_server_contributors()[0] - .contribute(McpServerContributionContext::for_thread( + .contribute(McpServerContributionContext::for_step( config, &thread_init, + &thread_store, + &available_environment_ids, )) .await .into_iter() @@ -132,7 +136,9 @@ async fn selected_plugin_contributions( selection_order, enabled: config.enabled, }, - McpServerContribution::Set { .. } | McpServerContribution::Remove { .. } => { + McpServerContribution::Set { .. } + | McpServerContribution::SelectedPluginConnectors { .. } + | McpServerContribution::Remove { .. } => { panic!("expected selected plugin contribution") } }) From 8ebf71ec257fcb95d72a71c7d3bedbbe9a861e45 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 01:47:00 +0100 Subject: [PATCH 32/33] Reuse walk inventory for environment skill metadata (#30145) ## Why Environment skill discovery already asks the executor to run one `fs/walk`. That response contains every regular file path found under the selected root, including any `agents/openai.yaml` files. Today Core keeps the discovered `SKILL.md` paths but discards the rest of that file inventory. It then sends one `fs/getMetadata` request per skill just to ask whether `agents/openai.yaml` exists. A root with 66 skills and no metadata therefore pays for 66 unnecessary network round trips. ## What changes - Keep the `fs/walk` file and directory inventory for the duration of the scan. - Associate each discovered `SKILL.md` with metadata that is known present, known absent, or still requires a fallback probe. - Read a known `agents/openai.yaml` directly instead of statting it first. - Skip the metadata request entirely when a complete walk shows that the skill has no `agents` directory. - Read a known `SKILL.md` and `agents/openai.yaml` concurrently. - Keep parsing and validation in `core-skills`. The inventory is scan-local. This does not add another cache or change cache lifetime. ## Network impact For a complete scan of 66 valid skills with no `agents/openai.yaml`, and one root `.codex-plugin/plugin.json`: | Operation | Current | After this PR | | --- | ---: | ---: | | `fs/walk` | 1 | 1 | | Read `SKILL.md` | 66 | 66 | | Stat `agents/openai.yaml` | 66 | 0 | | Read `agents/openai.yaml` | 0 | 0 | | Stat plugin manifest | 1 | 1 | | Read plugin manifest | 1 | 1 | | **Total executor RPCs** | **135** | **69** | This removes exactly 66 request/response exchanges from the common cold scan. Warm scans remain at zero discovery RPCs because the thread-level executor catalog cache is unchanged. When metadata exists, each file still requires one read. This PR removes only the preceding existence check; it does not batch file contents into a new RPC. ## Correctness fallbacks Absence is trusted only when the walk is complete and the metadata directory was not present. Core keeps the existing `getMetadata` fallback when: - the walk was truncated; - the walk reported an error; or - an `agents` directory was observed but `openai.yaml` was not, which preserves support for file symlinks and traversal boundaries. ## Deliberate scope This PR changes only the environment skill loader and its existing filesystem-call regression coverage. It does not: - change `fs/walk` or any exec-server protocol; - add `readFiles` or a skills-list endpoint; - change thread caching; - change local skill discovery; - change exec-server request concurrency; or - optimize plugin-manifest lookup. The plugin-manifest stat is intentionally left in place, which is why this PR reaches 69 calls rather than the broader 68-call estimate. That lookup has separate alternate-path, ancestor, and symlink semantics and should not be mixed into this change. --- .../core-skills/src/loader/environment.rs | 152 ++++++++++++++---- .../core-skills/tests/environment_loader.rs | 121 +++++++++++++- 2 files changed, 236 insertions(+), 37 deletions(-) diff --git a/codex-rs/core-skills/src/loader/environment.rs b/codex-rs/core-skills/src/loader/environment.rs index 93ab14e9aa..96caa28eea 100644 --- a/codex-rs/core-skills/src/loader/environment.rs +++ b/codex-rs/core-skills/src/loader/environment.rs @@ -23,7 +23,6 @@ use super::ParsedSkillFrontmatter; use super::SKILLS_FILENAME; use super::SKILLS_METADATA_DIR; use super::SKILLS_METADATA_FILENAME; -use super::SkillFileDiscovery; use super::SkillMetadataFile; use super::parse_skill_frontmatter_metadata_inner; use super::resolve_dependencies; @@ -34,6 +33,24 @@ use super::validate_len; const MAX_SKILLS_ENTRIES_PER_ROOT: usize = 20_000; const MAX_CONCURRENT_SKILL_LOADS: usize = 64; +struct EnvironmentSkillDiscovery { + skills: Vec, + plugin_roots: HashSet, + namespace_roots: HashSet, + warnings: Vec, +} + +struct DiscoveredEnvironmentSkill { + path: PathUri, + metadata: SkillMetadataDiscovery, +} + +enum SkillMetadataDiscovery { + Present(PathUri), + Absent, + Probe(PathUri), +} + /// URI-native metadata for one skill owned by an execution environment. #[derive(Clone, Debug, PartialEq, Eq)] pub struct EnvironmentSkillMetadata { @@ -67,28 +84,44 @@ impl EnvironmentSkillMetadata { async fn parse( file_system: &dyn ExecutorFileSystem, - path: &PathUri, + skill: &DiscoveredEnvironmentSkill, plugin_namespace: Option<&str>, ) -> Result { - let contents = file_system - .read_file_text(path, /*sandbox*/ None) - .await - .map_err(|err| format!("failed to read file: {err}"))?; + let (contents, discovered_metadata) = match &skill.metadata { + SkillMetadataDiscovery::Present(metadata_path) => { + let (contents, metadata) = tokio::join!( + read_skill_contents(file_system, &skill.path), + read_skill_metadata(file_system, metadata_path), + ); + (contents?, metadata) + } + SkillMetadataDiscovery::Absent | SkillMetadataDiscovery::Probe(_) => ( + read_skill_contents(file_system, &skill.path).await?, + (None, None), + ), + }; let ParsedSkillFrontmatter { name: base_name, description, short_description, - } = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(path)) + } = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(&skill.path)) .map_err(|err| err.to_string())?; let name = plugin_namespace .map(|namespace| format!("{namespace}:{base_name}")) .unwrap_or(base_name); validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name") .map_err(|err| err.to_string())?; - let (dependencies, policy) = load_skill_metadata(file_system, path).await; + let (dependencies, policy) = match &skill.metadata { + SkillMetadataDiscovery::Present(_) | SkillMetadataDiscovery::Absent => { + discovered_metadata + } + SkillMetadataDiscovery::Probe(metadata_path) => { + probe_skill_metadata(file_system, metadata_path).await + } + }; Ok(Self { - path_to_skills_md: path.clone(), + path_to_skills_md: skill.path.clone(), name, description, short_description, @@ -125,6 +158,7 @@ pub async fn load_environment_skills_from_root( .await { Ok(walk) => { + let inventory_complete = !walk.truncated && walk.errors.is_empty(); let mut warnings = walk .errors .into_iter() @@ -141,10 +175,13 @@ pub async fn load_environment_skills_from_root( )); } let mut skill_files = Vec::new(); + let mut file_paths = HashSet::new(); + let mut directory_paths = HashSet::new(); let mut plugin_roots = HashSet::new(); for entry in walk.entries { match entry.kind { WalkEntryKind::Directory => { + directory_paths.insert(entry.path.clone()); if DISCOVERABLE_PLUGIN_MANIFEST_PATHS .iter() .any(|path| path.split('/').next() == entry.path.basename().as_deref()) @@ -154,40 +191,53 @@ pub async fn load_environment_skills_from_root( } } WalkEntryKind::File => { + file_paths.insert(entry.path.clone()); if entry.path.basename().as_deref() == Some(SKILLS_FILENAME) { skill_files.push(entry.path); } } } } - SkillFileDiscovery { - skill_files, + let skills = skill_files + .into_iter() + .map(|path| DiscoveredEnvironmentSkill { + metadata: discover_skill_metadata( + &path, + &file_paths, + &directory_paths, + inventory_complete, + ), + path, + }) + .collect(); + EnvironmentSkillDiscovery { + skills, plugin_roots, namespace_roots: HashSet::from([root.clone()]), warnings, } } - Err(error) if error.kind() == io::ErrorKind::NotFound => SkillFileDiscovery { - skill_files: Vec::new(), + Err(error) if error.kind() == io::ErrorKind::NotFound => EnvironmentSkillDiscovery { + skills: Vec::new(), plugin_roots: HashSet::new(), namespace_roots: HashSet::new(), warnings: Vec::new(), }, - Err(error) => SkillFileDiscovery { - skill_files: Vec::new(), + Err(error) => EnvironmentSkillDiscovery { + skills: Vec::new(), plugin_roots: HashSet::new(), namespace_roots: HashSet::new(), warnings: vec![format!("failed to walk skills root {root}: {error:#}")], }, }; outcome.warnings.extend(discovery.warnings); - if discovery.skill_files.is_empty() { + if discovery.skills.is_empty() { return outcome; } let mut skill_ancestors = HashSet::new(); - for skill_path in &discovery.skill_files { - let mut ancestor = skill_path.parent(); + for skill in &discovery.skills { + let mut ancestor = skill.path.parent(); while let Some(path) = ancestor { skill_ancestors.insert(path.clone()); ancestor = path.parent(); @@ -224,9 +274,9 @@ pub async fn load_environment_skills_from_root( // Remote executors can multiplex these independent per-skill reads, so polling a bounded // number together allows the I/O for each skill and its metadata to happen concurrently. - let skill_results = futures::stream::iter(discovery.skill_files) - .map(|path| { - let mut ancestor = path.parent(); + let skill_results = futures::stream::iter(discovery.skills) + .map(|skill| { + let mut ancestor = skill.path.parent(); let plugin_namespace = loop { let Some(current) = ancestor else { break None; @@ -236,9 +286,10 @@ pub async fn load_environment_skills_from_root( } ancestor = current.parent(); }; + let path = skill.path.clone(); async move { let result = - EnvironmentSkillMetadata::parse(file_system, &path, plugin_namespace).await; + EnvironmentSkillMetadata::parse(file_system, &skill, plugin_namespace).await; (path, result) } }) @@ -267,20 +318,48 @@ pub async fn load_environment_skills_from_root( outcome } -async fn load_skill_metadata( +fn discover_skill_metadata( + skill_path: &PathUri, + file_paths: &HashSet, + directory_paths: &HashSet, + inventory_complete: bool, +) -> SkillMetadataDiscovery { + let Some(skill_dir) = skill_path.parent() else { + return SkillMetadataDiscovery::Absent; + }; + let Ok(metadata_dir) = skill_dir.join(SKILLS_METADATA_DIR) else { + return SkillMetadataDiscovery::Absent; + }; + let Ok(metadata_path) = metadata_dir.join(SKILLS_METADATA_FILENAME) else { + return SkillMetadataDiscovery::Absent; + }; + if file_paths.contains(&metadata_path) { + SkillMetadataDiscovery::Present(metadata_path) + } else if inventory_complete && !directory_paths.contains(&metadata_dir) { + SkillMetadataDiscovery::Absent + } else { + // The walk can omit entries after an error or traversal limit. It also omits file + // symlinks, so keep the existing probe when the metadata directory itself was observed. + SkillMetadataDiscovery::Probe(metadata_path) + } +} + +async fn read_skill_contents( file_system: &dyn ExecutorFileSystem, skill_path: &PathUri, +) -> Result { + file_system + .read_file_text(skill_path, /*sandbox*/ None) + .await + .map_err(|err| format!("failed to read file: {err}")) +} + +async fn probe_skill_metadata( + file_system: &dyn ExecutorFileSystem, + metadata_path: &PathUri, ) -> (Option, Option) { - let Some(skill_dir) = skill_path.parent() else { - return (None, None); - }; - let Ok(metadata_path) = - skill_dir.join(&format!("{SKILLS_METADATA_DIR}/{SKILLS_METADATA_FILENAME}")) - else { - return (None, None); - }; match file_system - .get_metadata(&metadata_path, /*sandbox*/ None) + .get_metadata(metadata_path, /*sandbox*/ None) .await { Ok(metadata) if metadata.is_file => {} @@ -291,8 +370,15 @@ async fn load_skill_metadata( return (None, None); } } + read_skill_metadata(file_system, metadata_path).await +} + +async fn read_skill_metadata( + file_system: &dyn ExecutorFileSystem, + metadata_path: &PathUri, +) -> (Option, Option) { let contents = match file_system - .read_file_text(&metadata_path, /*sandbox*/ None) + .read_file_text(metadata_path, /*sandbox*/ None) .await { Ok(contents) => contents, diff --git a/codex-rs/core-skills/tests/environment_loader.rs b/codex-rs/core-skills/tests/environment_loader.rs index 42d2446083..c17802bedb 100644 --- a/codex-rs/core-skills/tests/environment_loader.rs +++ b/codex-rs/core-skills/tests/environment_loader.rs @@ -1,5 +1,7 @@ use std::fs; use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use codex_core_skills::loader::EnvironmentSkillMetadata; use codex_core_skills::loader::load_environment_skills_from_root; @@ -13,6 +15,8 @@ use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::LOCAL_FS; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; +use codex_exec_server::WalkOptions; +use codex_exec_server::WalkOutcome; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; use tempfile::tempdir; @@ -20,6 +24,15 @@ use tempfile::tempdir; struct RecordingFileSystem<'a> { inner: &'a dyn ExecutorFileSystem, read_files: Mutex>, + metadata_files: Mutex>, + walks: AtomicUsize, +} + +#[derive(Debug, PartialEq, Eq)] +struct FileSystemCalls { + walks: usize, + read_files: Vec, + metadata_files: Vec, } impl<'a> RecordingFileSystem<'a> { @@ -27,14 +40,29 @@ impl<'a> RecordingFileSystem<'a> { Self { inner, read_files: Mutex::new(Vec::new()), + metadata_files: Mutex::new(Vec::new()), + walks: AtomicUsize::new(0), } } - fn read_files(&self) -> Vec { - self.read_files + fn calls(&self) -> FileSystemCalls { + let mut read_files = self + .read_files .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() + .clone(); + read_files.sort_by_key(ToString::to_string); + let mut metadata_files = self + .metadata_files + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + metadata_files.sort_by_key(ToString::to_string); + FileSystemCalls { + walks: self.walks.load(Ordering::Relaxed), + read_files, + metadata_files, + } } } @@ -90,6 +118,10 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { path: &'a PathUri, sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, FileMetadata> { + self.metadata_files + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(path.clone()); self.inner.get_metadata(path, sandbox) } @@ -101,6 +133,16 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { self.inner.read_directory(path, sandbox) } + fn walk<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { + self.walks.fetch_add(1, Ordering::Relaxed); + self.inner.walk(path, options, sandbox) + } + fn remove<'a>( &'a self, path: &'a PathUri, @@ -200,7 +242,8 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest ); let mut manifest_reads = file_system - .read_files() + .calls() + .read_files .into_iter() .filter(|path| path.basename().as_deref() == Some("plugin.json")) .collect::>(); @@ -214,3 +257,73 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest expected_manifest_reads.sort_by_key(ToString::to_string); assert_eq!(manifest_reads, expected_manifest_reads); } + +#[tokio::test] +async fn reuses_walk_inventory_for_missing_skill_metadata() { + const SKILL_COUNT: usize = 66; + + let root = tempdir().expect("tempdir"); + let manifest_path = root.path().join(".codex-plugin/plugin.json"); + fs::create_dir_all(manifest_path.parent().expect("manifest parent")).expect("manifest dir"); + fs::write(&manifest_path, r#"{"name":"inventory"}"#).expect("manifest"); + + let mut skill_paths = Vec::new(); + for index in 0..SKILL_COUNT { + let name = format!("skill-{index}"); + let skill_path = root.path().join(&name).join("SKILL.md"); + fs::create_dir_all(skill_path.parent().expect("skill parent")).expect("skill dir"); + fs::write( + &skill_path, + format!("---\nname: {name}\ndescription: {name} skill.\n---\n"), + ) + .expect("skill"); + skill_paths.push(skill_path); + } + + let file_system = RecordingFileSystem::new(LOCAL_FS.as_ref()); + let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI"); + let outcome = load_environment_skills_from_root( + &file_system, + &root_uri, + /*restriction_product*/ None, + ) + .await; + + let mut expected_skills = skill_paths + .iter() + .enumerate() + .map(|(index, skill_path)| EnvironmentSkillMetadata { + path_to_skills_md: PathUri::from_host_native_path(skill_path).unwrap(), + name: format!("inventory:skill-{index}"), + description: format!("skill-{index} skill."), + short_description: None, + dependencies: None, + policy: None, + }) + .collect::>(); + expected_skills.sort_by(|left, right| { + left.name.cmp(&right.name).then_with(|| { + left.path_to_skills_md + .to_string() + .cmp(&right.path_to_skills_md.to_string()) + }) + }); + assert_eq!(outcome.skills, expected_skills); + assert_eq!(outcome.warnings, Vec::::new()); + + let mut expected_read_files = skill_paths + .iter() + .map(|path| PathUri::from_host_native_path(path).unwrap()) + .collect::>(); + let manifest_uri = PathUri::from_host_native_path(manifest_path).unwrap(); + expected_read_files.push(manifest_uri.clone()); + expected_read_files.sort_by_key(ToString::to_string); + assert_eq!( + file_system.calls(), + FileSystemCalls { + walks: 1, + read_files: expected_read_files, + metadata_files: vec![manifest_uri], + } + ); +} From da78d5fdc5d8b30a80b402100d8b25ed8031f8c4 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 18:00:39 -0700 Subject: [PATCH 33/33] [codex] implement standalone code-mode process host (#30111) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - implement the standalone `codex-code-mode-host` stdio service - route sessions, cells, delegate requests, responses, and cancellation through a bounded host peer - supervise request, writer, cell-forwarding, actor, and V8 failure boundaries - bound request/session tombstones and fail-stop the connection on invalid protocol state - add host-only duplex protocol tests and local Cargo/Bazel run recipes ## Why This stage makes the host process independently runnable and reviewable before exposing any remote client in Codex. Transport or runtime failure closes the connection and relies on process replacement rather than transactional recovery. ## Stack This is **3 of 4** in the process-owned code-mode session stack. - Depends on #30110 - The final client PR targets this branch ## Validation - `just test -p codex-code-mode-host` — 7 host-only tests passed - `just fix -p codex-code-mode-host` - `just bazel-lock-update` - `just bazel-lock-check` - `just fmt` --- codex-rs/Cargo.lock | 8 + codex-rs/code-mode-host/Cargo.toml | 15 + codex-rs/code-mode-host/src/delegate.rs | 85 +++ codex-rs/code-mode-host/src/host_tests.rs | 486 +++++++++++++++++ codex-rs/code-mode-host/src/lib.rs | 605 ++++++++++++++++++++++ codex-rs/code-mode-host/src/main.rs | 5 +- codex-rs/code-mode-host/src/peer.rs | 541 +++++++++++++++++++ codex-rs/code-mode-host/src/peer_tests.rs | 95 ++++ justfile | 14 + 9 files changed, 1853 insertions(+), 1 deletion(-) create mode 100644 codex-rs/code-mode-host/src/delegate.rs create mode 100644 codex-rs/code-mode-host/src/host_tests.rs create mode 100644 codex-rs/code-mode-host/src/lib.rs create mode 100644 codex-rs/code-mode-host/src/peer.rs create mode 100644 codex-rs/code-mode-host/src/peer_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 75c02adfb5..17c51f58a9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2515,6 +2515,14 @@ dependencies = [ [[package]] name = "codex-code-mode-host" version = "0.0.0" +dependencies = [ + "anyhow", + "codex-code-mode", + "codex-code-mode-protocol", + "pretty_assertions", + "tokio", + "tokio-util", +] [[package]] name = "codex-code-mode-protocol" diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a2c384d010..d37bd41012 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -8,5 +8,20 @@ license.workspace = true name = "codex-code-mode-host" path = "src/main.rs" +[lib] +doctest = false +name = "codex_code_mode_host" +path = "src/lib.rs" + [lints] workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "sync", "time"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-host/src/delegate.rs b/codex-rs/code-mode-host/src/delegate.rs new file mode 100644 index 0000000000..50d45e7de7 --- /dev/null +++ b/codex-rs/code-mode-host/src/delegate.rs @@ -0,0 +1,85 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::SessionId; +use tokio_util::sync::CancellationToken; + +use crate::peer::HostPeer; + +pub(super) struct RemoteDelegate { + session_id: SessionId, + peer: Arc, +} + +impl RemoteDelegate { + pub(super) fn new(session_id: SessionId, peer: Arc) -> Self { + Self { session_id, peer } + } +} + +impl CodeModeSessionDelegate for RemoteDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::InvokeTool { + invocation: invocation.into(), + }, + cancellation_token, + ) + .await? + { + DelegateResponse::ToolResult { result } => Ok(result), + DelegateResponse::NotificationDelivered => { + Err("code-mode client returned an invalid tool result".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::Notify { + call_id, + cell_id: cell_id.into(), + text, + }, + cancellation_token, + ) + .await? + { + DelegateResponse::NotificationDelivered => Ok(()), + DelegateResponse::ToolResult { .. } => { + Err("code-mode client returned an invalid notification result".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.peer + .close_cell(self.session_id.clone(), cell_id.clone()); + } +} diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs new file mode 100644 index 0000000000..b189b89e09 --- /dev/null +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -0,0 +1,486 @@ +use codex_code_mode_protocol::host::Capability; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireExecuteRequest; +use codex_code_mode_protocol::host::WireResult; +use pretty_assertions::assert_eq; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use super::HostState; +use super::MAX_ACTIVE_CELLS; +use super::MAX_IN_FLIGHT_REQUESTS; +use super::MAX_RECENT_REQUEST_IDS; +use super::RequestKind; +use super::RequestRegistry; +use super::SeenSessionIds; +use super::peer::HostPeer; +use super::run; + +fn client_hello( + versions: impl IntoIterator, + required_capabilities: CapabilitySet, +) -> ClientToHost { + ClientToHost::ClientHello( + ClientHello::new( + SupportedProtocolVersions::try_new(versions).expect("supported versions"), + required_capabilities, + CapabilitySet::empty(), + ) + .expect("client hello"), + ) +} + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +async fn decode_frame(frame: EncodedFrame) -> HostToClient { + let (reader, writer) = tokio::io::duplex(/*max_buf_size*/ 4096); + let writer = tokio::spawn(async move { + FramedWriter::new(writer) + .write_frame(&frame) + .await + .expect("write encoded frame"); + }); + let message = FramedReader::new(reader) + .read() + .await + .expect("read encoded frame") + .expect("encoded frame message"); + writer.await.expect("frame writer task"); + message +} + +fn execute_request(source: &str) -> WireExecuteRequest { + WireExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: Some(60_000), + max_output_tokens: Some(1_000), + } +} + +#[tokio::test] +async fn handshake_and_multiple_session_lifecycles_are_ordered() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("read hello"), + Some(HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::empty(), + ))) + ); + + for (request_id, id) in [ + (request_id(/*value*/ 1), "session-1"), + (request_id(/*value*/ 2), "session-2"), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::OpenSession { + session_id: session_id(id), + }, + }) + .await + .expect("open session"); + assert_eq!( + reader.read::().await.expect("session ready"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id(id), + }, + }, + }) + ); + } + + for (request_id, id) in [ + (request_id(/*value*/ 3), "session-1"), + (request_id(/*value*/ 4), "session-2"), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::ShutdownSession { + session_id: session_id(id), + }, + }) + .await + .expect("shutdown session"); + assert_eq!( + reader.read::().await.expect("session closed"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id(id), + }, + }, + }) + ); + } + + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn incompatible_or_invalid_handshake_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version"); + + writer + .write(&client_hello([version_two], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { + supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .expect("host versions"), + }, + }) + ); + host.await.expect("host task").expect("host connection"); + + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&ClientToHost::Request { + id: request_id(/*value*/ 1), + request: HostRequest::OpenSession { + session_id: session_id("session-1"), + }, + }) + .await + .expect("write invalid first message"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn unsupported_required_capability_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let capability = Capability::new("required").expect("capability"); + + writer + .write(&client_hello( + [ProtocolVersion::V1], + CapabilitySet::try_new([capability.clone()]).expect("capabilities"), + )) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { capability }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn session_id_cannot_be_reused_after_shutdown() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + reader + .read::() + .await + .expect("read hello") + .expect("host hello"); + + let id = session_id("session-1"); + for (request_id, request) in [ + ( + request_id(/*value*/ 1), + HostRequest::OpenSession { + session_id: id.clone(), + }, + ), + ( + request_id(/*value*/ 2), + HostRequest::ShutdownSession { + session_id: id.clone(), + }, + ), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request, + }) + .await + .expect("session request"); + reader + .read::() + .await + .expect("session response") + .expect("session response message"); + } + writer + .write(&ClientToHost::Request { + id: request_id(/*value*/ 3), + request: HostRequest::OpenSession { session_id: id }, + }) + .await + .expect("reuse session ID"); + assert_eq!( + reader.read::().await.expect("reuse response"), + Some(HostToClient::Response { + id: request_id(/*value*/ 3), + result: WireResult::Err { + message: "code-mode session ID `session-1` was reused".to_string(), + }, + }) + ); + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[test] +fn request_cancellation_tombstones_are_bounded() { + let mut requests = RequestRegistry::default(); + let duplicate = request_id(/*value*/ -1); + requests + .start(duplicate, RequestKind::OpenSession) + .expect("start duplicate probe"); + assert!(requests.start(duplicate, RequestKind::OpenSession).is_err()); + requests.finish(duplicate); + for value in 1..=MAX_RECENT_REQUEST_IDS as i64 + 100 { + let id = request_id(value); + requests + .start(id, RequestKind::Wait) + .expect("start request"); + requests.cancel(id); + requests.finish(id); + } + for value in 10_000..20_000 { + requests.cancel(request_id(value)); + } + + assert!(requests.active.is_empty()); + assert_eq!(requests.recent.len(), MAX_RECENT_REQUEST_IDS); + assert_eq!(requests.recent_order.len(), MAX_RECENT_REQUEST_IDS); +} + +#[tokio::test] +async fn request_task_panic_disconnects_host() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }; + let task = state.request_tasks.spawn(async { + panic!("request panic probe"); + }); + state.supervise_request_task(task); + + tokio::time::timeout(Duration::from_secs(1), peer.disconnected()) + .await + .expect("request panic should disconnect host"); + assert!( + peer.failure() + .expect("request failure") + .contains("request task failed") + ); +} + +#[tokio::test] +async fn execute_request_id_remains_active_until_initial_response() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer, + }); + let session_id = session_id("session-1"); + state + .open_session(session_id.clone()) + .expect("open session"); + let request_id = request_id(/*value*/ 1); + + state + .spawn_request( + request_id, + HostRequest::Execute { + session_id: session_id.clone(), + request: execute_request("await new Promise(() => {});"), + }, + ) + .expect("spawn execute request"); + let started = decode_frame(outgoing_rx.recv().await.expect("execution started frame")).await; + let HostToClient::Response { + id, + result: + WireResult::Ok { + value: HostResponse::ExecutionStarted { cell_id }, + }, + } = started + else { + panic!("expected execution started response"); + }; + assert_eq!(id, request_id); + assert!( + state + .requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .active + .contains_key(&request_id) + ); + + state + .session(&session_id) + .expect("session") + .terminate(cell_id.into()) + .await + .expect("terminate cell"); + state.disconnect().await; +} + +#[tokio::test] +async fn active_cell_limit_rejects_execute_without_disconnecting() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(/*permits*/ 0)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }; + let session_id = session_id("session-1"); + state + .open_session(session_id.clone()) + .expect("open session"); + let request_id = request_id(/*value*/ 1); + + state + .handle_request( + request_id, + HostRequest::Execute { + session_id, + request: execute_request("text(\"hello\");"), + }, + CancellationToken::new(), + ) + .await; + + assert_eq!( + decode_frame(outgoing_rx.recv().await.expect("execute response frame")).await, + HostToClient::Response { + id: request_id, + result: WireResult::Err { + message: "code-mode host has too many active cells".to_string(), + }, + } + ); + assert!(!peer.is_disconnected()); + state.disconnect().await; +} + +#[tokio::test] +async fn cell_forwarding_panic_disconnects_host() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + peer.spawn_critical("cell forwarding", async { + panic!("cell forwarding panic probe"); + }); + + tokio::time::timeout(Duration::from_secs(1), peer.disconnected()) + .await + .expect("cell panic should disconnect host"); + assert!( + peer.failure() + .expect("cell failure") + .contains("cell forwarding task failed") + ); +} +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::time::Duration; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs new file mode 100644 index 0000000000..d2c1f507b4 --- /dev/null +++ b/codex-rs/code-mode-host/src/lib.rs @@ -0,0 +1,605 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::InProcessCodeModeSession; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use self::delegate::RemoteDelegate; +use self::peer::HostPeer; + +mod delegate; +mod peer; + +const MAX_IN_FLIGHT_REQUESTS: usize = 256; +const MAX_ACTIVE_CELLS: usize = 128; +const MAX_RECENT_REQUEST_IDS: usize = 4096; +const MAX_RECENT_SESSION_IDS: usize = 4096; +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); + +/// Runs one code-mode host connection over the process standard streams. +pub async fn run_stdio() -> Result<()> { + run(tokio::io::stdin(), tokio::io::stdout()).await +} + +/// Runs one code-mode host connection over an ordered input/output pair. +async fn run(reader: R, writer: W) -> Result<()> +where + R: AsyncRead + Send + Unpin + 'static, + W: AsyncWrite + Send + Unpin + 'static, +{ + let mut reader = FramedReader::new(reader); + let mut writer = FramedWriter::new(writer); + if !negotiate(&mut reader, &mut writer).await? { + return Ok(()); + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(/*max_capacity*/ 128); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }); + let writer_disconnected = peer.disconnection_token(); + let writer_task = tokio::spawn(async move { + loop { + tokio::select! { + _ = writer_disconnected.cancelled() => return Ok::<(), anyhow::Error>(()), + frame = outgoing_rx.recv() => { + let Some(frame) = frame else { + return Ok(()); + }; + if let Err(err) = writer.write_frame(&frame).await { + return Err( + anyhow::Error::new(err) + .context("failed to write code-mode host message") + ); + } + } + } + } + }); + let writer_peer = Arc::clone(&peer); + let writer_supervisor = tokio::spawn(async move { + match writer_task.await { + Ok(Ok(())) if !writer_peer.is_disconnected() => { + writer_peer.fail("code-mode writer task exited unexpectedly".to_string()); + } + Ok(Ok(())) => {} + Ok(Err(err)) => { + writer_peer.fail(format!("code-mode writer task failed: {err:#}")); + } + Err(err) => { + writer_peer.fail(format!("code-mode writer task failed: {err}")); + } + } + }); + + let input_result = async { + loop { + let message = tokio::select! { + _ = peer.disconnected() => break, + message = reader.read::() => message + .context("failed to read code-mode client message")?, + }; + let Some(message) = message else { + break; + }; + match message { + ClientToHost::ClientHello(_) => { + anyhow::bail!("received a second code-mode client hello"); + } + ClientToHost::Request { id, request } => { + state.spawn_request(id, request)?; + } + ClientToHost::CancelRequest { id } => { + state.cancel_request(id); + } + ClientToHost::DelegateResponse { id, result } => { + peer.complete(id, result.into_result()).await; + } + } + } + Ok::<(), anyhow::Error>(()) + } + .await; + + peer.disconnect(); + if tokio::time::timeout(SHUTDOWN_TIMEOUT, state.disconnect()) + .await + .is_err() + { + peer.fail("timed out shutting down code-mode host state".to_string()); + } + drop(state); + tokio::time::timeout(SHUTDOWN_TIMEOUT, writer_supervisor) + .await + .context("timed out supervising code-mode writer task")? + .context("code-mode writer supervisor task failed")?; + let failure = peer.failure(); + drop(peer); + input_result?; + if let Some(failure) = failure { + anyhow::bail!(failure); + } + Ok(()) +} + +async fn negotiate(reader: &mut FramedReader, writer: &mut FramedWriter) -> Result +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let Some(first_message) = reader + .read::() + .await + .context("failed to read code-mode client hello")? + else { + return Ok(false); + }; + let ClientToHost::ClientHello(client_hello) = first_message else { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + .await + .context("failed to reject invalid code-mode client hello")?; + return Ok(false); + }; + + let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?; + if !client_hello + .supported_versions() + .contains(ProtocolVersion::V1) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions }, + }) + .await + .context("failed to reject incompatible code-mode client")?; + return Ok(false); + } + + let host_capabilities = CapabilitySet::empty(); + if let Some(capability) = client_hello + .required_capabilities() + .iter() + .find(|capability| !host_capabilities.contains(capability)) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { + capability: capability.clone(), + }, + }) + .await + .context("failed to reject unsupported code-mode capability")?; + return Ok(false); + } + + writer + .write(&HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + host_capabilities, + ))) + .await + .context("failed to write code-mode host hello")?; + Ok(true) +} + +struct HostState { + sessions: Mutex>>, + seen_session_ids: Mutex, + requests: Mutex, + request_tasks: TaskTracker, + request_permits: Arc, + active_cell_permits: Arc, + closing: AtomicBool, + peer: Arc, +} + +impl HostState { + fn spawn_request( + self: &Arc, + request_id: RequestId, + request: HostRequest, + ) -> Result<(), anyhow::Error> { + let cancellation = self + .requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .start(request_id, RequestKind::from(&request))?; + let Ok(permit) = Arc::clone(&self.request_permits).try_acquire_owned() else { + self.respond( + request_id, + Err("code-mode host has too many in-flight requests".to_string()), + ); + self.finish_request(request_id); + return Ok(()); + }; + let state = Arc::clone(self); + let request_task = self.request_tasks.spawn(async move { + let _permit = permit; + state + .handle_request(request_id, request, cancellation) + .await; + state.finish_request(request_id); + }); + self.supervise_request_task(request_task); + Ok(()) + } + + fn supervise_request_task(&self, task: tokio::task::JoinHandle<()>) { + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if let Err(err) = task.await { + peer.fail(format!("code-mode request task failed: {err}")); + } + }); + } + + async fn handle_request( + &self, + request_id: RequestId, + request: HostRequest, + cancellation: CancellationToken, + ) { + if self.closing.load(Ordering::Acquire) { + self.respond( + request_id, + Err("code-mode host is shutting down".to_string()), + ); + return; + } + match request { + HostRequest::OpenSession { session_id } => { + let result = self + .open_session(session_id.clone()) + .map(|()| HostResponse::SessionReady { session_id }); + self.respond(request_id, result); + } + HostRequest::Execute { + session_id, + request, + } => { + if cancellation.is_cancelled() { + self.respond(request_id, Err("code-mode request cancelled".to_string())); + return; + } + let request = match request.try_into() { + Ok(request) => request, + Err(err) => { + self.respond( + request_id, + Err(format!("invalid code-mode execute request: {err}")), + ); + return; + } + }; + let session = match self.session(&session_id) { + Ok(session) => session, + Err(err) => { + self.respond(request_id, Err(err)); + return; + } + }; + let Ok(active_cell_permit) = + Arc::clone(&self.active_cell_permits).try_acquire_owned() + else { + self.respond( + request_id, + Err("code-mode host has too many active cells".to_string()), + ); + return; + }; + let result = session.execute(request).await; + match result { + Ok(started) => { + let cell_id = started.cell_id.clone(); + self.respond( + request_id, + Ok(HostResponse::ExecutionStarted { + cell_id: cell_id.into(), + }), + ); + let initial_response_sent = self.peer.start_cell( + session_id, + request_id, + started, + active_cell_permit, + ); + let _ = initial_response_sent.await; + } + Err(err) => self.respond(request_id, Err(err)), + } + } + HostRequest::Wait { + session_id, + request, + } => { + let result = match self.session(&session_id) { + Ok(session) => { + tokio::select! { + biased; + _ = cancellation.cancelled() => { + Err("code-mode request cancelled".to_string()) + } + result = session.wait(request.into()) => result.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + } + } + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::Terminate { + session_id, + cell_id, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::ShutdownSession { session_id } => { + let session = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&session_id); + let result = match session { + Some(session) => match session.shutdown().await { + Ok(()) => { + self.peer.wait_for_session_cells(&session_id).await; + Ok(HostResponse::SessionClosed { session_id }) + } + Err(err) => Err(err), + }, + None => Err(format!("unknown code-mode session {session_id}")), + }; + self.respond(request_id, result); + } + } + } + + fn open_session(&self, session_id: SessionId) -> Result<(), String> { + let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner); + if sessions.contains_key(&session_id) { + return Err(format!( + "code-mode session ID `{session_id}` is already open" + )); + } + if self.closing.load(Ordering::Acquire) { + return Err("code-mode host is shutting down".to_string()); + } + if !self + .seen_session_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remember(session_id.clone()) + { + return Err(format!("code-mode session ID `{session_id}` was reused")); + } + let delegate = Arc::new(RemoteDelegate::new( + session_id.clone(), + Arc::clone(&self.peer), + )); + let peer = Arc::downgrade(&self.peer); + let task_failure_handler = Arc::new(move |reason| { + if let Some(peer) = peer.upgrade() { + peer.fail(reason); + } + }); + sessions.insert( + session_id, + Arc::new( + InProcessCodeModeSession::with_delegate_and_task_failure_handler( + delegate, + task_failure_handler, + ), + ), + ); + Ok(()) + } + + fn session(&self, session_id: &SessionId) -> Result, String> { + self.sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(session_id) + .cloned() + .ok_or_else(|| format!("unknown code-mode session {session_id}")) + } + + fn respond(&self, id: RequestId, result: Result) { + self.peer.respond(id, result); + } + + fn cancel_request(&self, request_id: RequestId) { + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .cancel(request_id); + } + + fn finish_request(&self, request_id: RequestId) { + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .finish(request_id); + } + + async fn disconnect(&self) { + self.closing.store(true, Ordering::Release); + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .cancel_all(); + self.request_tasks.close(); + self.request_tasks.wait().await; + let sessions = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .drain() + .map(|(_, session)| session) + .collect::>(); + for session in sessions { + let _ = session.shutdown().await; + } + } +} + +#[derive(Clone, Copy)] +enum RequestKind { + OpenSession, + Execute, + Wait, + Terminate, + ShutdownSession, +} + +impl RequestKind { + fn from(request: &HostRequest) -> Self { + match request { + HostRequest::OpenSession { .. } => Self::OpenSession, + HostRequest::Execute { .. } => Self::Execute, + HostRequest::Wait { .. } => Self::Wait, + HostRequest::Terminate { .. } => Self::Terminate, + HostRequest::ShutdownSession { .. } => Self::ShutdownSession, + } + } + + fn is_cancellable(self) -> bool { + matches!(self, Self::Execute | Self::Wait) + } +} + +struct ActiveRequest { + kind: RequestKind, + cancellation: CancellationToken, +} + +#[derive(Default)] +struct RequestRegistry { + active: HashMap, + recent: HashSet, + recent_order: VecDeque, +} + +impl RequestRegistry { + fn start( + &mut self, + request_id: RequestId, + kind: RequestKind, + ) -> Result { + if self.active.contains_key(&request_id) || self.recent.contains(&request_id) { + anyhow::bail!("duplicate code-mode request ID {request_id:?}"); + } + let cancellation = CancellationToken::new(); + self.active.insert( + request_id, + ActiveRequest { + kind, + cancellation: cancellation.clone(), + }, + ); + Ok(cancellation) + } + + fn cancel(&self, request_id: RequestId) { + if let Some(request) = self.active.get(&request_id) + && request.kind.is_cancellable() + { + request.cancellation.cancel(); + } + } + + fn finish(&mut self, request_id: RequestId) { + if self.active.remove(&request_id).is_none() { + return; + } + self.recent.insert(request_id); + self.recent_order.push_back(request_id); + while self.recent_order.len() > MAX_RECENT_REQUEST_IDS { + if let Some(expired) = self.recent_order.pop_front() { + self.recent.remove(&expired); + } + } + } + + fn cancel_all(&self) { + for request in self.active.values() { + request.cancellation.cancel(); + } + } +} + +#[derive(Default)] +struct SeenSessionIds { + ids: HashSet, + order: VecDeque, +} + +impl SeenSessionIds { + fn remember(&mut self, session_id: SessionId) -> bool { + if !self.ids.insert(session_id.clone()) { + return false; + } + self.order.push_back(session_id); + while self.order.len() > MAX_RECENT_SESSION_IDS { + if let Some(expired) = self.order.pop_front() { + self.ids.remove(&expired); + } + } + true + } +} + +#[cfg(test)] +#[path = "host_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index f328e4d9d0..b215869c8a 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1 +1,4 @@ -fn main() {} +#[tokio::main(flavor = "current_thread")] +async fn main() -> anyhow::Result<()> { + codex_code_mode_host::run_stdio().await +} diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs new file mode 100644 index 0000000000..14ff068e7b --- /dev/null +++ b/codex-rs/code-mode-host/src/peer.rs @@ -0,0 +1,541 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireResult; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +const CELL_MESSAGE_CAPACITY: usize = 128; +const MAX_PENDING_DELEGATE_CALLS: usize = 256; + +pub(super) struct HostPeer { + outgoing_tx: mpsc::Sender, + pending: Mutex>, + delegate_permits: Arc, + cell_routes: StdMutex>, + cell_routes_changed: Notify, + next_request_id: AtomicI64, + disconnected: CancellationToken, + failure: StdMutex>, +} + +struct PendingDelegate { + response_tx: oneshot::Sender>, + dispatched: bool, + _permit: OwnedSemaphorePermit, +} + +enum CellRoute { + Pending(VecDeque), + Active(mpsc::Sender), +} + +enum CellMessage { + Delegate { + id: DelegateRequestId, + request: DelegateRequest, + dispatched_tx: oneshot::Sender>, + }, + Closed, +} + +impl HostPeer { + pub(super) fn new(outgoing_tx: mpsc::Sender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)), + cell_routes: StdMutex::new(HashMap::new()), + cell_routes_changed: Notify::new(), + next_request_id: AtomicI64::new(1), + disconnected: CancellationToken::new(), + failure: StdMutex::new(None), + } + } + + pub(super) fn send(&self, message: HostToClient) -> Result<(), PeerSendError> { + let frame = EncodedFrame::encode(&message) + .map_err(|err| PeerSendError::Payload(err.to_string()))?; + self.send_frame(frame) + } + + pub(super) fn respond( + &self, + id: RequestId, + result: Result, + ) { + let message = HostToClient::Response { + id, + result: WireResult::from_result(result), + }; + if let Err(PeerSendError::Payload(err)) = self.send(message) { + let _ = self.send(HostToClient::Response { + id, + result: WireResult::Err { + message: format!("code-mode host response exceeds the IPC frame limit: {err}"), + }, + }); + } + } + + fn initial_response( + &self, + id: RequestId, + result: Result, + ) { + let message = HostToClient::InitialResponse { + id, + result: WireResult::from_result(result), + }; + if let Err(PeerSendError::Payload(err)) = self.send(message) { + let _ = self.send(HostToClient::InitialResponse { + id, + result: WireResult::Err { + message: format!( + "code-mode initial response exceeds the IPC frame limit: {err}" + ), + }, + }); + } + } + + pub(super) async fn call( + self: &Arc, + session_id: SessionId, + request: DelegateRequest, + cancellation_token: CancellationToken, + ) -> Result { + if self.disconnected.is_cancelled() { + return Err("code-mode client connection closed".to_string()); + } + let Ok(permit) = Arc::clone(&self.delegate_permits).try_acquire_owned() else { + return Err("code-mode host has too many pending delegate calls".to_string()); + }; + let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + self.pending.lock().await.insert( + id, + PendingDelegate { + response_tx, + dispatched: false, + _permit: permit, + }, + ); + let mut pending = PendingDelegateRequest::new(Arc::clone(self), id); + let cell_id = match &request { + DelegateRequest::InvokeTool { invocation } => invocation.cell_id.clone().into(), + DelegateRequest::Notify { cell_id, .. } => cell_id.clone().into(), + }; + let (dispatched_tx, dispatched_rx) = oneshot::channel(); + if let Err(err) = self.route_cell_message( + (session_id, cell_id), + CellMessage::Delegate { + id, + request, + dispatched_tx, + }, + ) { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err(err); + } + + let dispatched = tokio::select! { + dispatched = dispatched_rx => dispatched.map_err(|_| { + "code-mode cell route closed before dispatching delegate request".to_string() + })?, + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err("code-mode client connection closed".to_string()); + } + }; + if let Err(err) = dispatched { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err(err); + } + + tokio::select! { + response = response_rx => { + pending.disarm(); + response.map_err(|_| { + "code-mode client closed before returning delegate output".to_string() + })? + } + _ = cancellation_token.cancelled() => { + if self.remove_pending(id).await.is_some() { + let _ = self.send(HostToClient::CancelDelegateRequest { id }); + } + pending.disarm(); + Err("code mode delegate request cancelled".to_string()) + } + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + Err("code-mode client connection closed".to_string()) + } + } + } + + pub(super) async fn complete( + &self, + id: DelegateRequestId, + response: Result, + ) { + if let Some(pending) = self.remove_pending(id).await { + let _ = pending.response_tx.send(response); + } + } + + pub(super) fn start_cell( + self: &Arc, + session_id: SessionId, + request_id: RequestId, + started: StartedCell, + active_cell_permit: OwnedSemaphorePermit, + ) -> oneshot::Receiver<()> { + let (initial_response_sent_tx, initial_response_sent_rx) = oneshot::channel(); + let key = (session_id, started.cell_id.clone()); + let (messages_tx, messages_rx) = mpsc::channel(CELL_MESSAGE_CAPACITY); + let previous = self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(key.clone(), CellRoute::Active(messages_tx.clone())); + match previous { + Some(CellRoute::Pending(messages)) => { + for message in messages { + if messages_tx.try_send(message).is_err() { + self.disconnect(); + return initial_response_sent_rx; + } + } + } + Some(CellRoute::Active(_)) => { + self.disconnect(); + return initial_response_sent_rx; + } + None => {} + } + let peer = Arc::clone(self); + self.spawn_critical("cell forwarding", async move { + drive_cell( + peer, + key, + request_id, + started, + messages_rx, + initial_response_sent_tx, + active_cell_permit, + ) + .await; + }); + initial_response_sent_rx + } + + pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) { + let _ = self.route_cell_message((session_id, cell_id), CellMessage::Closed); + } + + pub(super) fn disconnect(&self) { + self.disconnected.cancel(); + } + + pub(super) fn fail(&self, reason: String) { + let mut failure = self.failure.lock().unwrap_or_else(PoisonError::into_inner); + if failure.is_none() { + *failure = Some(reason); + } + drop(failure); + self.disconnect(); + } + + pub(super) fn failure(&self) -> Option { + self.failure + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + pub(super) fn is_disconnected(&self) -> bool { + self.disconnected.is_cancelled() + } + + pub(super) async fn disconnected(&self) { + self.disconnected.cancelled().await; + } + + pub(super) fn disconnection_token(&self) -> CancellationToken { + self.disconnected.clone() + } + + pub(super) async fn wait_for_session_cells(&self, session_id: &SessionId) { + loop { + let changed = self.cell_routes_changed.notified(); + if !self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .keys() + .any(|(route_session_id, _)| route_session_id == session_id) + { + return; + } + tokio::select! { + _ = changed => {} + _ = self.disconnected.cancelled() => return, + } + } + } + + async fn send_delegate_if_pending( + &self, + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + dispatched_tx: oneshot::Sender>, + ) { + let result = { + let mut pending = self.pending.lock().await; + let Some(pending) = pending.get_mut(&id) else { + let _ = dispatched_tx.send(Err( + "code-mode delegate request was cancelled before dispatch".to_string(), + )); + return; + }; + match self.send(HostToClient::DelegateRequest { + id, + session_id, + request, + }) { + Ok(()) => { + pending.dispatched = true; + Ok(()) + } + Err(err) => Err(err.to_string()), + } + }; + let _ = dispatched_tx.send(result); + } + + fn route_cell_message( + &self, + key: (SessionId, CellId), + message: CellMessage, + ) -> Result<(), String> { + use std::collections::hash_map::Entry; + + let result = match self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(key) + { + Entry::Occupied(mut entry) => match entry.get_mut() { + CellRoute::Pending(messages) if messages.len() < CELL_MESSAGE_CAPACITY => { + messages.push_back(message); + Ok(()) + } + CellRoute::Pending(_) => Err("code-mode cell message queue is full".to_string()), + CellRoute::Active(sender) => sender + .try_send(message) + .map_err(|_| "code-mode cell message queue is unavailable".to_string()), + }, + Entry::Vacant(entry) => { + entry.insert(CellRoute::Pending(VecDeque::from([message]))); + Ok(()) + } + }; + if result.is_err() { + self.disconnect(); + } + result + } + + async fn remove_pending(&self, id: DelegateRequestId) -> Option { + self.pending.lock().await.remove(&id) + } + + pub(super) fn spawn_critical(self: &Arc, task_name: &'static str, future: F) + where + F: std::future::Future + Send + 'static, + { + let task = tokio::spawn(future); + let peer = Arc::clone(self); + tokio::spawn(async move { + if let Err(err) = task.await { + peer.fail(format!("code-mode {task_name} task failed: {err}")); + } + }); + } + + fn send_frame(&self, frame: EncodedFrame) -> Result<(), PeerSendError> { + match self.outgoing_tx.try_send(frame) { + Ok(()) => Ok(()), + Err(mpsc::error::TrySendError::Full(_)) => { + self.disconnect(); + Err(PeerSendError::Unavailable( + "code-mode host outgoing queue is full".to_string(), + )) + } + Err(mpsc::error::TrySendError::Closed(_)) => { + self.disconnect(); + Err(PeerSendError::Unavailable( + "code-mode client connection closed".to_string(), + )) + } + } + } +} + +async fn drive_cell( + peer: Arc, + key: (SessionId, CellId), + request_id: RequestId, + started: StartedCell, + mut messages_rx: mpsc::Receiver, + initial_response_sent_tx: oneshot::Sender<()>, + _active_cell_permit: OwnedSemaphorePermit, +) { + let mut initial_response_sent_tx = Some(initial_response_sent_tx); + let initial_response = started.initial_response(); + tokio::pin!(initial_response); + let closed = loop { + tokio::select! { + biased; + result = &mut initial_response => { + peer.initial_response(request_id, result.map(Into::into)); + if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() { + let _ = initial_response_sent_tx.send(()); + } + break false; + } + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { + id, + request, + dispatched_tx, + }) => { + peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await; + } + Some(CellMessage::Closed) | None => break true, + }, + _ = peer.disconnected.cancelled() => { + peer.remove_cell_route(&key); + return; + } + } + }; + + if closed { + peer.initial_response(request_id, initial_response.await.map(Into::into)); + if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() { + let _ = initial_response_sent_tx.send(()); + } + } else { + loop { + tokio::select! { + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { + id, + request, + dispatched_tx, + }) => { + peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await; + } + Some(CellMessage::Closed) | None => break, + }, + _ = peer.disconnected.cancelled() => { + peer.remove_cell_route(&key); + return; + } + } + } + } + let _ = peer.send(HostToClient::CellClosed { + session_id: key.0.clone(), + cell_id: (&key.1).into(), + }); + peer.remove_cell_route(&key); +} + +impl HostPeer { + fn remove_cell_route(&self, key: &(SessionId, CellId)) { + let removed = self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(key); + if removed.is_some() { + self.cell_routes_changed.notify_waiters(); + } + } +} + +pub(super) enum PeerSendError { + Payload(String), + Unavailable(String), +} + +impl std::fmt::Display for PeerSendError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Payload(message) | Self::Unavailable(message) => formatter.write_str(message), + } + } +} + +struct PendingDelegateRequest { + peer: Arc, + id: Option, +} + +impl PendingDelegateRequest { + fn new(peer: Arc, id: DelegateRequestId) -> Self { + Self { peer, id: Some(id) } + } + + fn disarm(&mut self) { + self.id = None; + } +} + +impl Drop for PendingDelegateRequest { + fn drop(&mut self) { + let Some(id) = self.id.take() else { + return; + }; + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if let Some(pending) = peer.remove_pending(id).await + && pending.dispatched + { + let _ = peer.send(HostToClient::CancelDelegateRequest { id }); + } + }); + } +} + +#[cfg(test)] +#[path = "peer_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/peer_tests.rs b/codex-rs/code-mode-host/src/peer_tests.rs new file mode 100644 index 0000000000..a926027ce9 --- /dev/null +++ b/codex-rs/code-mode-host/src/peer_tests.rs @@ -0,0 +1,95 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use pretty_assertions::assert_eq; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::sync::oneshot::error::TryRecvError; +use tokio_util::sync::CancellationToken; + +use super::HostPeer; +use super::MAX_PENDING_DELEGATE_CALLS; + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +#[tokio::test] +async fn start_cell_reports_when_initial_response_is_enqueued() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let cell_id = CellId::new("cell-1".to_string()); + let (response_tx, response_rx) = oneshot::channel(); + let started = StartedCell::new(cell_id.clone(), response_rx); + let active_cell_permits = Arc::new(Semaphore::new(/*permits*/ 1)); + let active_cell_permit = Arc::clone(&active_cell_permits) + .try_acquire_owned() + .expect("active cell permit"); + + let mut initial_response_sent = peer.start_cell( + session_id("session-1"), + RequestId::new(/*value*/ 1), + started, + active_cell_permit, + ); + assert_eq!(initial_response_sent.try_recv(), Err(TryRecvError::Empty)); + + response_tx + .send(RuntimeResponse::Result { + cell_id: cell_id.clone(), + content_items: Vec::new(), + error_text: None, + }) + .expect("initial response receiver"); + initial_response_sent + .await + .expect("initial response completion"); + outgoing_rx.recv().await.expect("initial response frame"); + assert_eq!(active_cell_permits.available_permits(), 0); + + peer.close_cell(session_id("session-1"), cell_id); + let permit = tokio::time::timeout( + Duration::from_secs(1), + Arc::clone(&active_cell_permits).acquire_owned(), + ) + .await + .expect("cell permit should be released") + .expect("cell permit semaphore should remain open"); + drop(permit); +} + +#[tokio::test] +async fn pending_delegate_limit_rejects_call_without_disconnecting() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let permits = Arc::clone(&peer.delegate_permits) + .acquire_many_owned(MAX_PENDING_DELEGATE_CALLS as u32) + .await + .expect("delegate permits"); + + let result = peer + .call( + session_id("session-1"), + DelegateRequest::Notify { + call_id: "call-1".to_string(), + cell_id: CellId::new("cell-1".to_string()).into(), + text: "hello".to_string(), + }, + CancellationToken::new(), + ) + .await; + + assert_eq!( + result, + Err("code-mode host has too many pending delegate calls".to_string()) + ); + assert!(!peer.is_disconnected()); + drop(permits); +} diff --git a/justfile b/justfile index 03f87391ac..d4a0e6566c 100644 --- a/justfile +++ b/justfile @@ -31,6 +31,10 @@ tui-with-exec-server *args: file-search *args: cargo run --bin codex-file-search -- {args} +# Run the standalone code-mode host from source. +code-mode-host *args: + cargo run --bin codex-code-mode-host -- {args} + # Build the CLI and run the app-server test client app-server-test-client *args: cargo build -p codex-cli @@ -107,6 +111,16 @@ bazel-codex *args: bazel-codex *args: bazel run //codex-rs/cli:codex --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) +# Build and run the standalone code-mode host from source using Bazel. +[no-cd] +[unix] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under="cd $PWD &&" -- "$@" + +[windows] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) + [no-cd] bazel-lock-update: bazel mod deps --lockfile_mode=update