diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index be40ac75f9..a5c17924dc 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -102,6 +102,11 @@ impl McpBinding { self.connections.has_servers() } + /// The environment owning this server in the step's connection set, even if it exposes no tools. + pub fn server_environment_id(&self, server: &str) -> Option<&str> { + self.connections.server_environment_id(server) + } + pub async fn list_resources( &self, server: &str, diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs index fbb5889d05..b42f09da89 100644 --- a/codex-rs/codex-mcp/src/connection_manager/resources.rs +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -170,6 +170,12 @@ impl McpConnectionSet { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } + pub(crate) fn server_environment_id(&self, server: &str) -> Option<&str> { + self.servers + .get(server) + .map(|view| view.metadata.environment_id.as_str()) + } + pub(crate) async fn client_by_name( &self, name: &str, diff --git a/codex-rs/core/src/guardian/permissions.rs b/codex-rs/core/src/guardian/permissions.rs index 3580d73eb9..30c70fa7e3 100644 --- a/codex-rs/core/src/guardian/permissions.rs +++ b/codex-rs/core/src/guardian/permissions.rs @@ -1,50 +1,28 @@ //! Resolves Guardian permission evidence from the host's captured environment context. //! Sync review and async scoring use the same rules and environment identity. +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; + use super::GuardianReviewContext; use crate::context::GuardianPermissionContext; use crate::session::turn_context::TurnEnvironment; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolPayload; -use crate::tools::handlers::RequestPermissionsEnvironmentArgs; -pub(crate) async fn for_tool( - invocation: &ToolInvocation, -) -> anyhow::Result { +pub(crate) fn for_tool(invocation: &ToolInvocation) -> anyhow::Result { let context = GuardianReviewContext::from(&invocation.step_context); - if !invocation.tool_name.is_default_namespace() { - return for_environment(&context, /*environment_id*/ None); - } - let environment_id = match (invocation.tool_name.name.as_str(), &invocation.payload) { - ("apply_patch", ToolPayload::Custom { input }) => { - codex_apply_patch::parse_patch(input)?.environment_id - } - ("write_stdin", ToolPayload::Function { arguments }) => { - let arguments: serde_json::Value = serde_json::from_str(arguments)?; - let process_id = arguments["session_id"] - .as_i64() - .and_then(|id| i32::try_from(id).ok()) - .ok_or_else(|| anyhow::anyhow!("missing terminal session id"))?; - Some( - invocation - .session - .services - .unified_exec_manager - .environment_id_for_process(process_id) - .await - .ok_or_else(|| anyhow::anyhow!("terminal session is unavailable"))?, - ) - } - ("request_permissions", ToolPayload::Function { arguments }) => { - serde_json::from_str::(arguments)?.environment_id - } - ("exec_command" | "view_image", ToolPayload::Function { arguments }) => { - let arguments: serde_json::Value = serde_json::from_str(arguments)?; - serde_json::from_value::>(arguments["environment_id"].clone())? - } - _ => return for_environment(&context, /*environment_id*/ None), + let environment_id = if matches!(invocation.environments, ToolEnvironments::Mcp { .. }) + && invocation.environments.single_id() == Some(DEFAULT_MCP_SERVER_ENVIRONMENT_ID) + && context.environments().local_selection().is_none() + { + // Local execution tools can only resolve to an environment present in the step. + // Host MCP tools and resources can also serve remote-only turns: their hooks use + // the server owner, while Guardian needs the turn's primary permissions. + None + } else { + invocation.environments.single_id() }; - for_environment(&context, environment_id.as_deref()) + for_environment(&context, environment_id) } pub(super) fn for_environment( diff --git a/codex-rs/core/src/hook_environment_executor.rs b/codex-rs/core/src/hook_environment_executor.rs new file mode 100644 index 0000000000..fa6b303405 --- /dev/null +++ b/codex-rs/core/src/hook_environment_executor.rs @@ -0,0 +1,255 @@ +//! Dispatches admitted hooks on the ready environments captured for a lifecycle event. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::bail; +use codex_exec_server::ExecOutputStream; +use codex_exec_server::ExecParams; +use codex_exec_server::ExecProcess; +use codex_exec_server::ProcessId; +use codex_hooks::EnvironmentHookOutput; +use codex_hooks::EnvironmentHookSource; +use codex_hooks::HookEnvironmentExecutor; +use codex_hooks::Hooks; +use codex_network_proxy::NetworkProxy; +use codex_protocol::models::PermissionProfile; +use futures::FutureExt; +use futures::future::BoxFuture; + +use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; +use crate::shell::ShellType; +use crate::tools::network_approval::execution_network_proxy; +use crate::tools::network_approval::remote_network_proxy_launch; +use crate::tools::sandboxing::ToolError; +use crate::unified_exec::exec_env_policy_from_shell_policy; + +const MAX_HOOK_OUTPUT_BYTES: usize = 1024 * 1024; + +pub(crate) fn with_environment_hooks( + session: &Arc, + turn: &Arc, + environments: &TurnEnvironmentSnapshot, +) -> Hooks { + with_tool_environment_hooks( + session, + turn, + environments, + /*target_environments*/ &[], + ) +} + +pub(crate) fn with_tool_environment_hooks( + session: &Arc, + turn: &Arc, + environments: &TurnEnvironmentSnapshot, + target_environments: &[&str], +) -> Hooks { + session.hooks().with_environment_hooks( + environment_hook_sources(environments, target_environments), + Arc::new(EnvironmentCommandExecutor { + environments: environments.clone(), + session: Arc::clone(session), + turn: Arc::clone(turn), + }), + ) +} + +pub(crate) fn environment_hook_sources( + environments: &TurnEnvironmentSnapshot, + target_environments: &[&str], +) -> Vec { + environments + .turn_environments() + .filter(|environment| { + !environment.config().environment_hooks.is_empty() + && (target_environments.is_empty() + || target_environments.contains(&environment.selection.environment_id.as_str())) + }) + .map(|environment| EnvironmentHookSource { + environment_id: environment.selection.environment_id.clone(), + cwd: environment.cwd().clone(), + hooks: environment.config().environment_hooks.clone(), + }) + .collect() +} + +struct EnvironmentCommandExecutor { + environments: TurnEnvironmentSnapshot, + session: Arc, + turn: Arc, +} + +impl EnvironmentCommandExecutor { + async fn launch_params( + &self, + environment: &TurnEnvironment, + argv: Vec, + env: HashMap, + ) -> Result<(ExecParams, Option), ToolError> { + let process_id = ProcessId::from(format!("hook-{}", uuid::Uuid::new_v4())); + let network = execution_network_proxy( + &self.session, + &self.turn, + Some(environment), + self.turn.network.clone(), + &environment.selection.environment_id, + environment.permission_profile(), + environment.config().network_policy.as_ref(), + process_id.as_str(), + /*fallback_policy_decider*/ None, + ) + .await?; + let network_proxy = if let Some(network) = &network { + let (launch, _) = remote_network_proxy_launch(environment, network, &self.turn).await?; + launch.proxy.enabled.then_some(launch) + } else { + None + }; + let sandbox = network_proxy.as_ref().map(|_| { + // Enforce managed network policy with the hook owner's filesystem access. + let mut sandbox = environment.sandbox_context(/*additional_permissions*/ None); + sandbox.permissions = PermissionProfile::Disabled; + sandbox + }); + let params = ExecParams { + process_id, + metadata: None, + argv, + cwd: environment.cwd().clone(), + env_policy: Some(exec_env_policy_from_shell_policy( + environment.shell_environment_policy(), + )), + shell_snapshot: None, + env, + tty: false, + pipe_stdin: true, + arg0: None, + enforce_managed_network: sandbox.is_some(), + sandbox, + managed_network: None, + network_proxy, + }; + Ok((params, network)) + } +} + +struct HookProcessGuard(Option>); + +impl Drop for HookProcessGuard { + fn drop(&mut self) { + // Don't let waiting for termination delay a hook timeout or error. + if let Some(process) = self.0.take() { + tokio::spawn(async move { + let _ = process.terminate().await; + }); + } + } +} + +impl HookEnvironmentExecutor for EnvironmentCommandExecutor { + fn execute( + &self, + environment_id: &str, + command: &str, + env: &HashMap, + input_json: &str, + timeout: Duration, + ) -> BoxFuture<'_, anyhow::Result> { + let environment_id = environment_id.to_owned(); + let command = command.to_owned(); + let env = env.clone(); + let input_json = input_json.to_owned(); + async move { + let environment = self + .environments + .turn_environments() + .find(|environment| environment.selection.environment_id == environment_id) + .context("hook environment was not ready for this event")?; + let shell = environment + .shell + .as_ref() + .context("hook environment did not report a shell")?; + + // Exec-server keeps pipe stdin open and older executors have no close-stdin RPC. + // JSON is one line: read precisely that line, then give the hook a pipe that reaches EOF. + let input_json = format!("{input_json}\n").into_bytes(); + let script = match shell.shell_type { + // Start the hook on a new line so a leading `(` doesn't become Bash's arithmetic `((`. + ShellType::Sh | ShellType::Bash | ShellType::Zsh => format!( + "IFS= read -r codex_hook_input; printf '%s' \"$codex_hook_input\" | (\n{command}\n)" + ), + // A script block can reset `$?`. Run the hook at the top level and preserve its + // exit code before PowerShell's -Command reduces it to 0 or 1. + ShellType::PowerShell => format!( + "$global:LASTEXITCODE = 0\n\ + [Console]::In.ReadLine() | {command}\n\ + if ($?) {{ exit 0 }}\n\ + if ($LASTEXITCODE) {{ exit $LASTEXITCODE }}\n\ + exit 1" + ), + ShellType::Cmd => bail!("environment command hooks require PowerShell on Windows"), + }; + let argv = shell.derive_exec_args(&script, environment.config().allow_login_shell); + // Keep network attribution registered until the hook and its output are complete. + let (params, _execution_network) = self + .launch_params(environment, argv, env) + .await + .map_err(|error| match error { + ToolError::Rejected(message) => anyhow::anyhow!("{message}"), + ToolError::Codex(error) => anyhow::Error::new(error), + })?; + let deadline = tokio::time::Instant::now() + timeout; + let backend = environment.environment.get_exec_backend(); + let process = tokio::time::timeout_at(deadline, backend.start(params)) + .await + .context("environment command hook timed out during launch")?? + .process; + let mut guard = HookProcessGuard(Some(Arc::clone(&process))); + let output = tokio::time::timeout_at(deadline, async { + let written = process.write(input_json).await?; + if written.status != codex_exec_server::WriteStatus::Accepted { + bail!("hook process did not accept its input: {:?}", written.status); + } + let mut after_seq = None; + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + loop { + let result = process.read(after_seq, Some(64 * 1024), Some(1000)).await?; + if let Some(failure) = result.failure { + bail!("{failure}"); + } + // Closed describes the process, not this bounded page. Keep + // reading retained output until we receive an empty final page. + let drained = result.closed && result.chunks.is_empty(); + for chunk in result.chunks { + let target = match chunk.stream { + ExecOutputStream::Stdout | ExecOutputStream::Pty => &mut stdout, + ExecOutputStream::Stderr => &mut stderr, + }; + let remaining = MAX_HOOK_OUTPUT_BYTES.saturating_sub(target.len()); + target.extend_from_slice(&chunk.chunk.0[..remaining.min(chunk.chunk.0.len())]); + } + after_seq = result.next_seq.checked_sub(1).or(after_seq); + if drained { + return Ok(EnvironmentHookOutput { + exit_code: result.exit_code.context("hook process closed without exit status")?, + stdout: String::from_utf8_lossy(&stdout).into_owned(), + stderr: String::from_utf8_lossy(&stderr).into_owned(), + }); + } + } + }) + .await + .context("environment command hook timed out")??; + guard.0 = None; + Ok(output) + } + .boxed() + } +} diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 34f3351372..3dde6e1aea 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -63,6 +63,8 @@ use crate::context::HookAdditionalContext; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::event_mapping::parse_turn_item; use crate::guardian::GuardianReviewContext; +use crate::hook_environment_executor::with_environment_hooks; +use crate::hook_environment_executor::with_tool_environment_hooks; use crate::session::TurnInput; use crate::session::session::Session; use crate::session::step_context::StepContext; @@ -160,7 +162,11 @@ pub(crate) async fn run_pending_session_start_hooks( permission_mode: hook_permission_mode(turn_context.approval_policy()), target, }; - let hooks = sess.hooks(); + let hooks = with_environment_hooks( + sess, + turn_context, + &turn_context.initial_environments.refresh_readiness(), + ); let preview_runs = hooks.preview_session_start(&request); if run_context_injecting_hook( sess, @@ -190,6 +196,7 @@ pub(crate) async fn run_pre_tool_use_hooks( tool_use_id: String, tool_name: &HookToolName, tool_input: &Value, + target_environments: &[&str], ) -> PreToolUseHookResult { let turn_context = &step_context.turn; let request = PreToolUseRequest { @@ -205,7 +212,12 @@ pub(crate) async fn run_pre_tool_use_hooks( tool_use_id, tool_input: tool_input.clone(), }; - let hooks = sess.hooks(); + let hooks = with_tool_environment_hooks( + sess, + turn_context, + &step_context.environments, + target_environments, + ); let preview_runs = hooks.preview_pre_tool_use(&request); emit_hook_started_events(sess, turn_context, preview_runs).await; @@ -259,6 +271,7 @@ pub(crate) async fn run_permission_request_hooks( review_context: &GuardianReviewContext, run_id_suffix: &str, payload: PermissionRequestPayload, + target_environment: &str, ) -> Option { let turn_context = review_context.turn(); let request = PermissionRequestRequest { @@ -274,7 +287,12 @@ pub(crate) async fn run_permission_request_hooks( run_id_suffix: run_id_suffix.to_string(), tool_input: payload.tool_input, }; - let hooks = sess.hooks(); + let hooks = with_tool_environment_hooks( + sess, + turn_context, + review_context.environments(), + &[target_environment], + ); let preview_runs = hooks.preview_permission_request(&request); emit_hook_started_events(sess, turn_context, preview_runs).await; @@ -293,6 +311,7 @@ pub(crate) async fn run_permission_request_hooks( /// already adapted by the tool handler into the stable hook contract. Passing /// raw internal tool data here would leak implementation details into user hook /// matchers and hook logs. +#[allow(clippy::too_many_arguments)] pub(crate) async fn run_post_tool_use_hooks( sess: &Arc, step_context: &StepContext, @@ -301,6 +320,7 @@ pub(crate) async fn run_post_tool_use_hooks( matcher_aliases: Vec, tool_input: Value, tool_response: Value, + target_environments: &[&str], ) -> PostToolUseOutcome { let turn_context = &step_context.turn; let request = PostToolUseRequest { @@ -317,7 +337,12 @@ pub(crate) async fn run_post_tool_use_hooks( tool_input, tool_response, }; - let hooks = sess.hooks(); + let hooks = with_tool_environment_hooks( + sess, + turn_context, + &step_context.environments, + target_environments, + ); let preview_runs = hooks.preview_post_tool_use(&request); emit_hook_started_events(sess, turn_context, preview_runs).await; @@ -416,7 +441,8 @@ pub(crate) async fn run_turn_stop_hooks( target, }; let executor_hook_sources = step_context.executor_hook_sources.clone(); - let hooks = sess.hooks().with_executor_hooks(executor_hook_sources); + let hooks = with_environment_hooks(sess, turn_context, &step_context.environments) + .with_executor_hooks(executor_hook_sources); emit_hook_started_events(sess, turn_context, hooks.preview_stop(&request)).await; let mut outcome = hooks.run_stop(request).await; @@ -426,13 +452,26 @@ pub(crate) async fn run_turn_stop_hooks( #[instrument(level = "trace", skip_all)] pub(crate) async fn run_session_end_hooks(sess: &Arc) { - let hooks = sess.hooks(); - let preview_runs = hooks.preview_session_end(); - if preview_runs.is_empty() { + // Building a turn can read plugins and skills from a remote executor. Check + // the captured hook declarations before spending any of the shutdown deadline. + let environments = sess.services.turn_environments.snapshot_now(); + let has_environment_end_hooks = environments.turn_environments().any(|environment| { + !environment + .config() + .environment_hooks + .session_end + .is_empty() + }); + if !has_environment_end_hooks && sess.hooks().preview_session_end().is_empty() { return; } let turn_context = sess.new_default_turn().await; + let hooks = with_environment_hooks(sess, &turn_context, &environments); + let preview_runs = hooks.preview_session_end(); + if preview_runs.is_empty() { + return; + } // SessionEnd is root-only; ThreadSpawn uses SubagentStart/SubagentStop and other subagents // are internal implementation details. @@ -472,7 +511,13 @@ pub(crate) async fn run_turn_interrupt_hooks( .map(|step| step.executor_hook_sources.clone()) .unwrap_or_default(); let has_executor_hooks = !executor_hook_sources.is_empty(); - let hooks = sess.hooks().with_executor_hooks(executor_hook_sources); + let hooks = last_known_step_context + .as_deref() + .map_or_else( + || (*sess.hooks()).clone(), + |step| with_environment_hooks(sess, &step.turn, &step.environments), + ) + .with_executor_hooks(executor_hook_sources); let preview_runs = hooks.preview_interrupt(); if preview_runs.is_empty() && !has_executor_hooks { return; @@ -513,10 +558,15 @@ pub(crate) async fn run_pre_compact_hooks( model: turn_context.model_info().slug.clone(), trigger: compaction_trigger_label(trigger).to_string(), }; - let preview_runs = sess.hooks().preview_pre_compact(&request); + let hooks = with_environment_hooks( + sess, + turn_context, + &turn_context.initial_environments.refresh_readiness(), + ); + let preview_runs = hooks.preview_pre_compact(&request); emit_hook_started_events(sess, turn_context, preview_runs).await; - let outcome = sess.hooks().run_pre_compact(request).await; + let outcome = hooks.run_pre_compact(request).await; emit_hook_completed_events(sess, turn_context, outcome.hook_events).await; if outcome.should_stop { PreCompactHookOutcome::Stopped @@ -550,10 +600,15 @@ pub(crate) async fn run_post_compact_hooks( model: turn_context.model_info().slug.clone(), trigger: compaction_trigger_label(trigger).to_string(), }; - let preview_runs = sess.hooks().preview_post_compact(&request); + let hooks = with_environment_hooks( + sess, + turn_context, + &turn_context.initial_environments.refresh_readiness(), + ); + let preview_runs = hooks.preview_post_compact(&request); emit_hook_started_events(sess, turn_context, preview_runs).await; - let outcome = sess.hooks().run_post_compact(request).await; + let outcome = hooks.run_post_compact(request).await; emit_hook_completed_events(sess, turn_context, outcome.hook_events).await; if outcome.should_stop { PostCompactHookOutcome::Stopped @@ -649,7 +704,11 @@ pub(crate) async fn inspect_pending_input( permission_mode: hook_permission_mode(turn_context.approval_policy()), prompt: UserMessageItem::new(content).message(), }; - let hooks = sess.hooks(); + let hooks = with_environment_hooks( + sess, + turn_context, + &turn_context.initial_environments.refresh_readiness(), + ); let preview_runs = hooks.preview_user_prompt_submit(&request); run_context_injecting_hook( sess, @@ -1032,16 +1091,21 @@ fn compaction_trigger_label(value: CompactionTrigger) -> &'static str { } pub(crate) fn hooks_for_model_request( - session: &Session, + session: &Arc, step: &StepContext, ) -> Arc { - let hooks = session.hooks(); match step.turn.session_source { SessionSource::Internal(InternalSessionSource::MemoryConsolidation) - | SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) => { - Arc::new(hooks.for_memory_consolidation(step.executor_hook_sources.clone())) - } - _ => hooks, + | SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) => Arc::new( + session + .hooks() + .for_memory_consolidation(step.executor_hook_sources.clone()), + ), + _ => Arc::new(with_environment_hooks( + session, + &step.turn, + &step.environments, + )), } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 390b658b60..2c91c8e487 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -99,6 +99,7 @@ mod exec_policy; mod git_info_tests; mod guardian; pub mod guardian_review; +mod hook_environment_executor; mod hook_mcp_executor; mod hook_runtime; mod image_preparation; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 62438b74fc..6bf51b9ab1 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -290,6 +290,7 @@ pub(crate) async fn handle_mcp_tool_call( &metadata, prepared_call.config(), prepared_call.permission_profile(), + prepared_call.server_environment_id(), approval_policy, ) .await @@ -1515,6 +1516,7 @@ async fn maybe_request_mcp_tool_approval( metadata: &McpToolApprovalMetadata, config: &codex_mcp::McpConfig, permission_profile: &PermissionProfile, + environment_id: &str, policy: McpToolApprovalPolicy, ) -> Option { let turn_context = &step_context.turn; @@ -1567,6 +1569,7 @@ async fn maybe_request_mcp_tool_approval( let action = ApprovalAction::McpToolCall { id: call_id.to_string(), + environment_id: environment_id.to_owned(), server: invocation.server.clone(), tool_name: invocation.tool.clone(), arguments: invocation.arguments.clone(), diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 039dc918fd..9f4ef92042 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -8,6 +8,7 @@ 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::tests::mcp_config_for_test; +use crate::session::tests::test_environment_id; use crate::session::tests::update_selected_settings_for_test; use crate::session::tests::update_turn_settings_for_test; use crate::session::turn_context::TurnEnvironment; @@ -2301,6 +2302,7 @@ async fn dispatched_mcp_approval_with_closed_response_is_classified_as_approval( let action = ApprovalAction::McpToolCall { id: call_id.to_string(), + environment_id: test_environment_id(&turn_context).to_owned(), server: "calendar".to_string(), tool_name: "send".to_string(), arguments: None, @@ -2852,6 +2854,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Approve), ) .await; @@ -2934,6 +2937,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2998,6 +3002,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -3076,6 +3081,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -3170,6 +3176,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -3277,6 +3284,7 @@ async fn strict_auto_review_forces_guardian_for_mcp_policy_skip() { &metadata, &captured_mcp_config, &captured_mcp_config.permission_profile, + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Approve), ) .await; @@ -3361,6 +3369,7 @@ async fn assert_mcp_user_approval_persistence( &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await @@ -3452,6 +3461,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Prompt), ) .await @@ -3519,6 +3529,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(approval_mode), ) .await; @@ -3614,6 +3625,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { &metadata, &approval_config(&turn_context), turn_context.config.permissions.permission_profile(), + test_environment_id(&turn_context), McpToolApprovalPolicy::for_server(AppToolApproval::Approve), ) .await; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index ac96683d40..5d4162a500 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -62,15 +62,20 @@ impl StepContext { pub(crate) fn environments(&self) -> Vec<(&TurnEnvironment, impl EnvironmentAccess + '_)> { self.environments .turn_environments() - .map(|environment| { - let grants = self - .turn - .granted_permissions(&environment.selection.environment_id); - (environment, environment.fs_accessor(grants)) - }) + .map(|environment| (environment, self.fs_accessor(environment))) .collect() } + pub(crate) fn fs_accessor<'a>( + &self, + environment: &'a TurnEnvironment, + ) -> impl EnvironmentAccess + 'a + use<'a> { + let grants = self + .turn + .granted_permissions(&environment.selection.environment_id); + environment.fs_accessor(grants) + } + /// Persist the context captured for this request, even after a live update. pub(crate) fn to_turn_context_item(&self) -> TurnContextItem { let mut item = self.turn.to_turn_context_item(); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6ee2ccf4d1..7ebdd51de8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -248,6 +248,15 @@ pub(crate) fn mcp_config_for_test(config: &crate::config::Config) -> Arc &str { + &turn + .initial_environments + .single_local_environment() + .expect("test turn should have one local environment") + .selection + .environment_id +} + /// Updates both initial/current views before sharing a test turn context. pub(crate) fn update_turn_settings_for_test( turn: &mut TurnContext, @@ -7796,11 +7805,13 @@ async fn request_permissions_tool_resolves_legacy_paths_against_selected_environ current_environment.shell, ) }; - turn_context_mut.initial_environments.environments[0] = - TurnEnvironmentState::Ready(environment); + turn_context_mut + .initial_environments + .environments + .push(TurnEnvironmentState::Ready(environment)); let call_id = "call-1".to_string(); - let handler = RequestPermissionsHandler; + let registry = ToolRegistry::with_handler_for_test(Arc::new(RequestPermissionsHandler)); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let step_context = StepContext::for_test(Arc::clone(&turn_context)); let handle = tokio::spawn({ @@ -7810,30 +7821,32 @@ async fn request_permissions_tool_resolves_legacy_paths_against_selected_environ let tracker = Arc::clone(&tracker); let call_id = call_id.clone(); async move { - handler - .handle(ToolInvocation { - session, - step_context, - turn: turn_context, - cancellation_token: CancellationToken::new(), - tracker, - call_id, - tool_name: codex_tools::ToolName::plain("request_permissions"), - source: ToolCallSource::Direct, - payload: ToolPayload::Function { - arguments: json!({ - "environment_id": "remote", - "reason": "need write", - "permissions": { - "file_system": { - "read": null, - "write": ["relative.txt", "~/home-relative.txt"], - }, + let invocation = ToolInvocation { + session, + step_context, + turn: turn_context, + cancellation_token: CancellationToken::new(), + tracker, + call_id, + tool_name: codex_tools::ToolName::plain("request_permissions"), + source: ToolCallSource::Direct, + environments: Default::default(), + payload: ToolPayload::Function { + arguments: json!({ + "environment_id": "remote", + "reason": "need write", + "permissions": { + "file_system": { + "read": null, + "write": ["relative.txt", "~/home-relative.txt"], }, - }) - .to_string(), - }, - }) + }, + }) + .to_string(), + }, + }; + registry + .dispatch_any_with_state(invocation, /*call_state*/ None) .await } }); @@ -7894,28 +7907,31 @@ async fn request_permissions_tool_rejects_invalid_requests( environment.selection.cwd = PathUri::parse("file:///workspace/%FF").expect("non-UTF8 cwd"); let turn_context = Arc::new(turn_context); let step_context = StepContext::for_test(Arc::clone(&turn_context)); - let result = RequestPermissionsHandler - .handle(ToolInvocation { - session: Arc::new(session), - step_context, - turn: turn_context, - cancellation_token: CancellationToken::new(), - tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), - call_id: "call-1".to_string(), - tool_name: codex_tools::ToolName::plain("request_permissions"), - source: ToolCallSource::Direct, - payload: ToolPayload::Function { - arguments: json!({ - "environment_id": environment_id, - "permissions": { - "file_system": { - "write": ["relative"], - }, + let registry = ToolRegistry::with_handler_for_test(Arc::new(RequestPermissionsHandler)); + let invocation = ToolInvocation { + session: Arc::new(session), + step_context, + turn: turn_context, + cancellation_token: CancellationToken::new(), + tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + call_id: "call-1".to_string(), + tool_name: codex_tools::ToolName::plain("request_permissions"), + source: ToolCallSource::Direct, + environments: Default::default(), + payload: ToolPayload::Function { + arguments: json!({ + "environment_id": environment_id, + "permissions": { + "file_system": { + "write": ["relative"], }, - }) - .to_string(), - }, - }) + }, + }) + .to_string(), + }, + }; + let result = registry + .dispatch_any_with_state(invocation, /*call_state*/ None) .await; let Err(FunctionCallError::RespondToModel(output)) = result else { @@ -13376,6 +13392,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index adf47840f0..142402b345 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -616,6 +616,7 @@ async fn guardian_allows_exec_command_additional_permissions_requests_past_polic call_id: "test-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", @@ -740,6 +741,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_exec_command_policy_s call_id: "strict-shell-command-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", @@ -892,6 +894,7 @@ async fn delayed_exec_command_uses_its_captured_authority_after_next_turn_starts call_id: call_id.to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": command, @@ -1069,6 +1072,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", @@ -1193,6 +1197,7 @@ async fn exec_command_allows_sticky_turn_permissions_without_inline_request_perm call_id: "sticky-turn-grant".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", diff --git a/codex-rs/core/src/tools/approvals.rs b/codex-rs/core/src/tools/approvals.rs index cdd468cc16..245a63878a 100644 --- a/codex-rs/core/src/tools/approvals.rs +++ b/codex-rs/core/src/tools/approvals.rs @@ -102,6 +102,7 @@ pub(crate) enum ApprovalAction { }, McpToolCall { id: String, + environment_id: String, server: String, tool_name: String, arguments: Option, @@ -153,6 +154,17 @@ pub(crate) enum ApprovalCacheKey { } impl ApprovalAction { + fn hook_environment_id(&self) -> &str { + match self { + Self::ExecCommand { environment_id, .. } + | Self::WriteStdin { environment_id, .. } + | Self::ApplyPatch { environment_id, .. } + | Self::McpToolCall { environment_id, .. } + | Self::NetworkAccess { environment_id, .. } + | Self::RequestPermissions { environment_id, .. } => environment_id, + } + } + pub(crate) fn permission_request_payload(&self) -> PermissionRequestPayload { match self { Self::ExecCommand { @@ -488,6 +500,7 @@ impl Session { &ctx.review_context, &permission_request_run_id, action.permission_request_payload(), + action.hook_environment_id(), ) .await { diff --git a/codex-rs/core/src/tools/approvals_tests.rs b/codex-rs/core/src/tools/approvals_tests.rs index 09baa934c8..c7a131c59a 100644 --- a/codex-rs/core/src/tools/approvals_tests.rs +++ b/codex-rs/core/src/tools/approvals_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::session::tests::make_session_and_context_with_rx; +use crate::session::tests::test_environment_id; use codex_models_manager::model_info::model_info_from_slug; use codex_protocol::approvals::NetworkPolicyAmendment; use pretty_assertions::assert_eq; @@ -190,6 +191,7 @@ async fn explicit_mcp_reviewer_override_takes_precedence_over_action_context() { let (session, turn, events) = make_session_and_context_with_rx().await; let action = ApprovalAction::McpToolCall { id: "mcp-override".to_string(), + environment_id: test_environment_id(&turn).to_owned(), server: "example".to_string(), tool_name: "dangerous".to_string(), arguments: None, diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 1b02ab7abf..c80a89ddd1 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -2,6 +2,7 @@ use crate::original_image_detail::sanitize_original_image_detail; use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; use crate::turn_diff_tracker::TurnDiffTracker; use crate::unified_exec::format_output_omission_marker; use crate::unified_exec::resolve_max_tokens; @@ -66,6 +67,47 @@ pub enum ToolCallSource { }, } +/// Targets captured from the original call. Environments follow the order declared by the tool. +#[derive(Clone, Default)] +pub(crate) enum ToolEnvironments { + #[default] + All, + Terminal(crate::unified_exec::TerminalSession), + Arguments(Vec<(&'static str, Arc)>), + Mcp { + server: Option, + owner: Option, + }, +} + +impl ToolEnvironments { + pub(crate) fn single_executor(env: &TurnEnvironment) -> Self { + Self::Arguments(vec![("environment_id", Arc::new(env.clone()))]) + } + + pub(crate) fn single_id(&self) -> Option<&str> { + match self { + Self::Arguments(args) => match args.as_slice() { + [(_, env)] => Some(&env.selection.environment_id), + _ => None, + }, + Self::Terminal(terminal) => Some(&terminal.environment_id), + Self::Mcp { owner, .. } => owner.as_deref(), + Self::All => None, + } + } + + pub(crate) fn hook_ids(&self) -> Vec<&str> { + match self { + Self::Arguments(args) => args + .iter() + .map(|(_, env)| env.selection.environment_id.as_str()) + .collect(), + _ => self.single_id().into_iter().collect(), + } + } +} + #[derive(Clone)] pub struct ToolInvocation { pub session: Arc, @@ -78,6 +120,8 @@ pub struct ToolInvocation { pub tool_name: ToolName, pub source: ToolCallSource, pub payload: ToolPayload, + /// Captured before PreToolUse and retained across input rewrites. + pub(crate) environments: ToolEnvironments, } /// Identity of the model call that started a tool invocation, retained across code-mode waits. diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 41db894b59..86d3512faf 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -19,6 +19,7 @@ use crate::session::turn_context::TurnEnvironment; use crate::tools::context::ApplyPatchToolOutput; use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; @@ -28,6 +29,7 @@ use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::apply_patch_spec::create_apply_patch_freeform_tool; use crate::tools::handlers::file_system_sandbox_policy_context_for_cwd; use crate::tools::handlers::resolve_tool_environment; +use crate::tools::handlers::selected_executor; use crate::tools::handlers::updated_hook_command; use crate::tools::hook_names::HookToolName; use crate::tools::orchestrator::ToolOrchestrator; @@ -304,6 +306,7 @@ impl ApplyPatchHandler { invocation: ToolInvocation, ) -> Result, FunctionCallError> { let ToolInvocation { + environments, session, step_context, cancellation_token, @@ -327,14 +330,7 @@ impl ApplyPatchHandler { ))); } }; - let selected_environment_id = - require_environment_id(args.environment_id.as_deref(), self.multi_environment)?; - - // Verify the parsed patch against the selected environment filesystem. - let turn_environment = resolve_tool_environment( - &step_context.environments, - selected_environment_id.as_deref(), - )?; + let turn_environment = selected_executor(&environments, &step_context.environments)?; let fs = turn_environment.environment.get_filesystem(); let sandbox = turn_environment.sandbox_context(/*additional_permissions*/ None); match codex_apply_patch::verify_apply_patch_args( @@ -383,6 +379,27 @@ impl ApplyPatchHandler { } impl CoreToolRuntime for ApplyPatchHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + let ToolPayload::Custom { input } = &invocation.payload else { + return Err(FunctionCallError::Fatal( + "expected apply_patch input".to_string(), + )); + }; + let args = codex_apply_patch::parse_patch(input).map_err(|err| { + FunctionCallError::RespondToModel(format!("apply_patch verification failed: {err}")) + })?; + let environment_id = + require_environment_id(args.environment_id.as_deref(), self.multi_environment)?; + let environment = resolve_tool_environment( + &invocation.step_context.environments, + environment_id.as_deref(), + )?; + Ok(ToolEnvironments::single_executor(environment)) + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Custom { .. }) } @@ -404,6 +421,17 @@ impl CoreToolRuntime for ApplyPatchHandler { updated_input: serde_json::Value, ) -> Result { let patch = updated_hook_command(&updated_input)?; + let args = codex_apply_patch::parse_patch(patch).map_err(|err| { + FunctionCallError::RespondToModel(format!("apply_patch verification failed: {err}")) + })?; + if let Some(id) = + require_environment_id(args.environment_id.as_deref(), self.multi_environment)? + && Some(id.as_str()) != invocation.environments.single_id() + { + return Err(FunctionCallError::RespondToModel( + "PreToolUse hook cannot change the tool environment".to_string(), + )); + } invocation.payload = match invocation.payload { ToolPayload::Custom { .. } => ToolPayload::Custom { input: patch.to_string(), diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index 32298978ec..a89fed5bf3 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -55,6 +55,7 @@ async fn invocation_for_payload(payload: ToolPayload) -> ToolInvocation { tool_name: codex_tools::ToolName::plain("apply_patch"), source: crate::tools::context::ToolCallSource::Direct, payload, + environments: Default::default(), } } diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 793bc97925..5c499ebb0c 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -19,6 +19,7 @@ use codex_utils_string::to_ascii_json_string; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::lifecycle::extension_tool_call_source; @@ -64,8 +65,18 @@ impl ToolExecutor for ExtensionToolAdapter { ToolInvocation: 'a, { Box::pin(async move { - let environment_accessors = invocation.step_context.environments(); - let environments = environment_accessors + let step = &invocation.step_context; + let selected: Vec<_> = match &invocation.environments { + ToolEnvironments::Arguments(environments) => { + environments.iter().map(|(_, env)| env.as_ref()).collect() + } + _ => step.environments.turn_environments().collect(), + }; + let accessors: Vec<_> = selected + .into_iter() + .map(|env| (env, step.fs_accessor(env))) + .collect(); + let environments = accessors .iter() .map(|(environment, fs)| { ToolEnvironment::new( @@ -83,6 +94,33 @@ impl ToolExecutor for ExtensionToolAdapter { } impl CoreToolRuntime for ExtensionToolAdapter { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + let arguments = self.0.environment_arguments(&invocation.payload)?; + if arguments.is_empty() { + return Ok(ToolEnvironments::All); + } + arguments + .into_iter() + .map(|(field, id)| { + let env = invocation + .step_context + .environments + .turn_environments() + .find(|env| env.selection.environment_id == id) + .ok_or_else(|| { + codex_tools::FunctionCallError::RespondToModel(format!( + "unknown or unavailable environment ID: {id}" + )) + })?; + Ok((field, Arc::new(env.clone()))) + }) + .collect::>() + .map(ToolEnvironments::Arguments) + } + fn is_builtin_control_tool(&self) -> bool { let tool_name = self.0.tool_name(); if tool_name.is_default_namespace() { @@ -394,6 +432,7 @@ mod tests { call_id: "call-extension".to_string(), tool_name: codex_tools::ToolName::plain("extension_echo"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "message": "hello" }).to_string(), }, @@ -493,6 +532,7 @@ mod tests { cell_id: "cell-1".to_string(), runtime_tool_call_id: "nested-call-1".to_string(), }, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "message": "hello" }).to_string(), }, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index cadcbb855a..c607b341ff 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -11,6 +11,7 @@ use crate::original_image_detail::can_request_original_image_detail; use crate::session::session::Session; use crate::tools::context::McpToolOutput; use crate::tools::context::ToolCallSource; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -306,6 +307,21 @@ impl McpHandler { } impl CoreToolRuntime for McpHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + let server = &self.tool_info().server_name; + Ok(ToolEnvironments::Mcp { + server: Some(server.clone()), + owner: invocation + .step_context + .mcp + .server_environment_id(server) + .map(str::to_owned), + }) + } + fn immutable_spec(&self) -> Option<&Arc> { Some(&self.definition.spec) } @@ -685,6 +701,7 @@ mod tests { tool_name: codex_tools::ToolName::namespaced("memory", "create_entities"), source: ToolCallSource::Direct, payload, + environments: Default::default(), }), Some(PreToolUsePayload { tool_name: HookToolName::new("mcp__memory__create_entities"), @@ -719,6 +736,7 @@ mod tests { tool_name: codex_tools::ToolName::namespaced("mcp__foo", "exec_command"), source: ToolCallSource::Direct, payload, + environments: Default::default(), }), Some(PreToolUsePayload { tool_name: HookToolName::new("mcp__foo__exec_command"), @@ -749,6 +767,7 @@ mod tests { tool_name: codex_tools::ToolName::namespaced("mcp__foo", "exec_command"), source: ToolCallSource::Direct, payload, + environments: Default::default(), }, json!({ "message": "rewritten" }), ) @@ -810,6 +829,7 @@ mod tests { tool_name: codex_tools::ToolName::namespaced("filesystem", "read_file"), source: ToolCallSource::Direct, payload, + environments: Default::default(), }; assert_eq!( handler.post_tool_use_payload(&invocation, &output), diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index 205954ec7e..7a3d9223bc 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -30,7 +30,10 @@ use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolEnvironments; +use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; use codex_protocol::protocol::McpInvocation; @@ -59,6 +62,47 @@ fn ensure_model_can_access_mcp_server( } } +enum ResourceOperation { + List, + Read, +} + +fn resource_environments( + invocation: &ToolInvocation, + operation: ResourceOperation, +) -> Result { + let arguments = resource_arguments(&invocation.payload)?; + let server = match operation { + ResourceOperation::List => ListResourceArgs::parse(arguments)?.server, + ResourceOperation::Read => Some(ReadResourceArgs::parse(arguments)?.server), + }; + // Only an aggregate list has no owning server. A named but unknown server is an error. + let Some(server) = server else { + return Ok(ToolEnvironments::Mcp { + server: None, + owner: None, + }); + }; + let step = &invocation.step_context; + ensure_model_can_access_mcp_server(&step.turn, &server)?; + let environment_id = step.mcp.server_environment_id(&server).ok_or_else(|| { + FunctionCallError::RespondToModel(format!("unknown MCP server '{server}'")) + })?; + Ok(ToolEnvironments::Mcp { + server: Some(server), + owner: Some(environment_id.to_owned()), + }) +} + +fn resource_arguments(payload: &ToolPayload) -> Result, FunctionCallError> { + let ToolPayload::Function { arguments } = payload else { + return Err(FunctionCallError::Fatal( + "expected function arguments for MCP resource operation".to_string(), + )); + }; + parse_arguments(arguments) +} + #[derive(Debug, Deserialize, Default, PartialEq, Eq)] struct ListResourceArgs { #[serde(default)] @@ -68,6 +112,10 @@ struct ListResourceArgs { } impl ListResourceArgs { + fn parse(arguments: Option) -> Result { + parse_args_with_default::(arguments).map(Self::normalized) + } + fn normalized(self) -> Self { Self { server: normalize_optional_string(self.server), @@ -102,6 +150,16 @@ struct ReadResourceArgs { uri: String, } +impl ReadResourceArgs { + fn parse(arguments: Option) -> Result { + let Self { server, uri } = parse_args(arguments)?; + Ok(Self { + server: normalize_required_string("server", server)?, + uri: normalize_required_string("uri", uri)?, + }) + } +} + #[derive(Debug, Serialize)] struct ResourceWithServer { server: 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 271b843594..8a570ad115 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 @@ -1,6 +1,6 @@ use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_list_mcp_resource_templates_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -11,9 +11,10 @@ use codex_tools::ToolSpec; use super::ListResourceArgs; use super::ListResourceTemplatesPayload; +use super::ResourceOperation; use super::model_can_access_mcp_server; -use super::parse_args_with_default; -use super::parse_arguments; +use super::resource_arguments; +use super::resource_environments; use super::run_resource_operation; pub struct ListMcpResourceTemplatesHandler { @@ -62,18 +63,8 @@ impl ListMcpResourceTemplatesHandler { let turn = std::sync::Arc::clone(&step_context.turn); let mcp = &step_context.mcp; - let arguments = match payload { - ToolPayload::Function { arguments } => arguments, - _ => { - return Err(FunctionCallError::RespondToModel( - "list_mcp_resource_templates handler received unsupported payload".to_string(), - )); - } - }; - - let arguments = parse_arguments(arguments.as_str())?; - let args: ListResourceArgs = parse_args_with_default(arguments.clone())?; - let args = args.normalized(); + let arguments = resource_arguments(&payload)?; + let args = ListResourceArgs::parse(arguments.clone())?; let invocation = McpInvocation { server: args.server.clone().unwrap_or_else(|| "codex".to_string()), @@ -108,4 +99,11 @@ impl ListMcpResourceTemplatesHandler { } } -impl CoreToolRuntime for ListMcpResourceTemplatesHandler {} +impl CoreToolRuntime for ListMcpResourceTemplatesHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + resource_environments(invocation, ResourceOperation::List) + } +} 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 650673f76b..1f3e4fdcd0 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 @@ -1,6 +1,6 @@ use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_list_mcp_resources_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -11,9 +11,10 @@ use codex_tools::ToolSpec; use super::ListResourceArgs; use super::ListResourcesPayload; +use super::ResourceOperation; use super::model_can_access_mcp_server; -use super::parse_args_with_default; -use super::parse_arguments; +use super::resource_arguments; +use super::resource_environments; use super::run_resource_operation; pub struct ListMcpResourcesHandler { @@ -62,18 +63,8 @@ impl ListMcpResourcesHandler { let turn = std::sync::Arc::clone(&step_context.turn); let mcp = &step_context.mcp; - let arguments = match payload { - ToolPayload::Function { arguments } => arguments, - _ => { - return Err(FunctionCallError::RespondToModel( - "list_mcp_resources handler received unsupported payload".to_string(), - )); - } - }; - - let arguments = parse_arguments(arguments.as_str())?; - let args: ListResourceArgs = parse_args_with_default(arguments.clone())?; - let args = args.normalized(); + let arguments = resource_arguments(&payload)?; + let args = ListResourceArgs::parse(arguments.clone())?; let invocation = McpInvocation { server: args.server.clone().unwrap_or_else(|| "codex".to_string()), @@ -106,4 +97,11 @@ impl ListMcpResourcesHandler { } } -impl CoreToolRuntime for ListMcpResourcesHandler {} +impl CoreToolRuntime for ListMcpResourcesHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + resource_environments(invocation, ResourceOperation::List) + } +} 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 e8c828bf8c..67cdb7b67b 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 @@ -1,6 +1,6 @@ use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_read_mcp_resource_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -13,10 +13,10 @@ use rmcp::model::ReadResourceRequestParams; use super::ReadResourceArgs; use super::ReadResourcePayload; +use super::ResourceOperation; use super::ensure_model_can_access_mcp_server; -use super::normalize_required_string; -use super::parse_args; -use super::parse_arguments; +use super::resource_arguments; +use super::resource_environments; use super::run_resource_operation; pub struct ReadMcpResourceHandler { @@ -65,20 +65,8 @@ impl ReadMcpResourceHandler { let turn = std::sync::Arc::clone(&step_context.turn); let mcp = &step_context.mcp; - let arguments = match payload { - ToolPayload::Function { arguments } => arguments, - _ => { - return Err(FunctionCallError::RespondToModel( - "read_mcp_resource handler received unsupported payload".to_string(), - )); - } - }; - - let arguments = parse_arguments(arguments.as_str())?; - let args: ReadResourceArgs = parse_args(arguments.clone())?; - let ReadResourceArgs { server, uri } = args; - let server = normalize_required_string("server", server)?; - let uri = normalize_required_string("uri", uri)?; + let arguments = resource_arguments(&payload)?; + let ReadResourceArgs { server, uri } = ReadResourceArgs::parse(arguments.clone())?; let invocation = McpInvocation { server: server.clone(), @@ -105,4 +93,11 @@ impl ReadMcpResourceHandler { } } -impl CoreToolRuntime for ReadMcpResourceHandler {} +impl CoreToolRuntime for ReadMcpResourceHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + resource_environments(invocation, ResourceOperation::Read) + } +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 56ae313eef..eb71db6b3e 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -36,6 +36,9 @@ mod view_image; pub(crate) mod view_image_spec; mod wait_for_environment; +use crate::tools::context::ToolEnvironments; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; use codex_file_system::FileSystemSandboxContext; use codex_sandboxing::policy_transforms::materialize_additional_permissions_with_context; use codex_sandboxing::policy_transforms::merge_permission_profiles; @@ -68,7 +71,6 @@ pub use mcp_resource::ListMcpResourcesHandler; pub use mcp_resource::ReadMcpResourceHandler; pub use new_context_window::NewContextWindowHandler; pub use plan::PlanHandler; -pub(crate) use request_permissions::RequestPermissionsEnvironmentArgs; pub use request_permissions::RequestPermissionsHandler; pub use request_plugin_install::RequestPluginInstallHandler; pub use request_user_input::RequestUserInputHandler; @@ -158,6 +160,42 @@ where parse_arguments(arguments) } +#[derive(Deserialize)] +struct ToolEnvironmentArgs { + environment_id: Option, +} + +fn function_tool_environment( + invocation: &ToolInvocation, + environment_id: impl FnOnce(T) -> Option, +) -> Result { + let ToolPayload::Function { arguments } = &invocation.payload else { + return Err(FunctionCallError::Fatal( + "expected function arguments for environment selection".to_string(), + )); + }; + let environment_id = environment_id(parse_arguments(arguments)?); + resolve_tool_environment( + &invocation.step_context.environments, + environment_id.as_deref(), + ) + .map(ToolEnvironments::single_executor) +} + +fn selected_executor<'a>( + environments: &'a ToolEnvironments, + snapshot: &'a TurnEnvironmentSnapshot, +) -> Result<&'a TurnEnvironment, FunctionCallError> { + match environments { + ToolEnvironments::Arguments(args) if let [(_, env)] = args.as_slice() => Ok(env), + // Direct handler callers may supply only the step context. + ToolEnvironments::All => resolve_tool_environment(snapshot, /*environment_id*/ None), + _ => Err(FunctionCallError::Fatal( + "expected one executor for this call".to_string(), + )), + } +} + fn resolve_tool_environment<'a>( environments: &'a TurnEnvironmentSnapshot, environment_id: Option<&str>, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index be73c65ac3..9582a977e0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -117,6 +117,7 @@ fn invocation( tool_name: codex_tools::ToolName::plain(tool_name), source: crate::tools::context::ToolCallSource::Direct, payload, + environments: Default::default(), } } diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index b393766fb1..32830f516c 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -6,15 +6,18 @@ use serde_json::Value; use crate::function_tool::FunctionCallError; use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::function_tool_environment; use crate::tools::handlers::parse_arguments; -use crate::tools::handlers::resolve_tool_environment; +use crate::tools::handlers::selected_executor; use crate::tools::handlers::shell_spec::create_request_permissions_tool; use crate::tools::handlers::shell_spec::request_permissions_tool_description; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; +use crate::tools::registry::with_updated_function_input; use codex_tools::ToolName; use codex_tools::ToolSpec; use serde::Deserialize; @@ -22,9 +25,9 @@ use serde::Deserialize; pub struct RequestPermissionsHandler; #[derive(Deserialize)] -pub(crate) struct RequestPermissionsEnvironmentArgs { - #[serde(default, rename = "environment_id", alias = "environmentId")] - pub(crate) environment_id: Option, +struct RequestPermissionsEnvironmentArgs { + #[serde(alias = "environmentId")] + environment_id: Option, } impl ToolExecutor for RequestPermissionsHandler { @@ -50,6 +53,7 @@ impl RequestPermissionsHandler { invocation: ToolInvocation, ) -> Result, FunctionCallError> { let ToolInvocation { + environments, session, step_context, cancellation_token, @@ -67,11 +71,7 @@ impl RequestPermissionsHandler { } }; - let environment_args: RequestPermissionsEnvironmentArgs = parse_arguments(&arguments)?; - let turn_environment = resolve_tool_environment( - &step_context.environments, - environment_args.environment_id.as_deref(), - )?; + let turn_environment = selected_executor(&environments, &step_context.environments)?; let sandbox_context = turn_environment.sandbox_context(/*additional_permissions*/ None); let context = sandbox_context.policy_context(); @@ -121,7 +121,33 @@ impl RequestPermissionsHandler { } } -impl CoreToolRuntime for RequestPermissionsHandler {} +impl CoreToolRuntime for RequestPermissionsHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + function_tool_environment(invocation, |args: RequestPermissionsEnvironmentArgs| { + args.environment_id + }) + } + + fn with_updated_hook_input( + &self, + invocation: ToolInvocation, + mut updated_input: Value, + ) -> Result { + // The public permission tool accepts either spelling of this selector. + if let Some(args) = updated_input.as_object_mut() + && let Some(alias) = args.remove("environmentId") + && args.insert("environment_id".to_string(), alias).is_some() + { + return Err(FunctionCallError::RespondToModel( + "duplicate field `environment_id`".to_string(), + )); + } + with_updated_function_input(invocation, updated_input) + } +} fn resolve_permission_path_strings( arguments: &mut Value, diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs index b2f6c7b477..b4cff1e0b4 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs @@ -66,6 +66,7 @@ async fn request_plugin_install_rejects_subagent_threads(presentation: ToolSugge call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_PLUGIN_INSTALL_TOOL_NAME), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: arguments.to_string(), }, diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 1c1dbfca16..cb470abada 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -48,6 +48,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "questions": [{ @@ -126,6 +127,7 @@ async fn request_user_input_sets_non_blocking_outside_plan_mode( call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "questions": [{ @@ -247,6 +249,7 @@ async fn request_user_input_sets_blocking_from_turn_mode() { call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "questions": [{ diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 6eeaf861c1..50faf5259a 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -1,20 +1,22 @@ +use crate::function_tool::FunctionCallError; use crate::sandboxing::SandboxPermissions; use crate::shell::Shell; use crate::shell::ShellInvocation; use crate::shell::get_shell_by_model_provided_path; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; +use crate::tools::handlers::parse_arguments; +use crate::tools::handlers::resolve_tool_environment; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use codex_protocol::models::AdditionalPermissionProfile; +use futures::future::BoxFuture; use serde::Deserialize; use std::path::PathBuf; use std::sync::Arc; -#[cfg(test)] -use crate::tools::handlers::parse_arguments; - mod exec_command; mod write_stdin; @@ -22,6 +24,42 @@ pub use exec_command::ExecCommandHandler; pub(crate) use exec_command::ExecCommandHandlerOptions; pub use write_stdin::WriteStdinHandler; +/// Terminal tools expose the unified-exec process ID as `session_id`. +pub(crate) fn terminal_session_for_input( + invocation: &ToolInvocation, +) -> BoxFuture<'_, Result> { + #[derive(Deserialize)] + struct TerminalSessionArgs { + session_id: i32, + } + + Box::pin(async move { + let ToolPayload::Function { arguments } = &invocation.payload else { + return Err(FunctionCallError::Fatal( + "expected function arguments for terminal session".to_string(), + )); + }; + let args: TerminalSessionArgs = parse_arguments(arguments)?; + resolve_tool_environment( + &invocation.step_context.environments, + /*environment_id*/ None, + )?; + invocation + .session + .services + .unified_exec_manager + .resolve_terminal(args.session_id) + .await + .map(ToolEnvironments::Terminal) + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "{} failed: {err}", + invocation.tool_name.name + )) + }) + }) +} + #[derive(Debug, Deserialize)] pub(crate) struct ExecCommandArgs { pub(crate) cmd: String, @@ -48,9 +86,7 @@ pub(crate) struct ExecCommandArgs { } #[derive(Debug, Deserialize)] -struct ExecCommandEnvironmentArgs { - #[serde(default)] - environment_id: Option, +struct ExecCommandWorkdirArgs { // Keep this raw until after environment selection; relative paths must be // resolved against the selected environment cwd, not the process cwd. #[serde(default)] diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index ce0a502453..ca15be3f73 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -7,19 +7,22 @@ use crate::exec_policy::prompt_is_rejected_by_policy; use crate::function_tool::FunctionCallError; use crate::maybe_emit_implicit_skill_invocation; use crate::tools::context::ExecCommandToolOutput; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::ToolEnvironmentArgs; use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::apply_patch::intercept_apply_patch; use crate::tools::handlers::file_system_sandbox_policy_context_for_cwd; +use crate::tools::handlers::function_tool_environment; use crate::tools::handlers::implicit_granted_permissions; use crate::tools::handlers::normalize_and_validate_additional_permissions; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::parse_arguments_with_base_path; use crate::tools::handlers::resolve_sandbox_permissions; -use crate::tools::handlers::resolve_tool_environment; use crate::tools::handlers::rewrite_function_string_argument; +use crate::tools::handlers::selected_executor; use crate::tools::handlers::updated_hook_command; use crate::tools::hook_names::HookToolName; use crate::tools::registry::CoreToolRuntime; @@ -48,7 +51,7 @@ use codex_utils_string::truncate_middle_chars; use super::super::shell_spec::CommandToolOptions; use super::super::shell_spec::create_exec_command_tool_with_environment_id; use super::ExecCommandArgs; -use super::ExecCommandEnvironmentArgs; +use super::ExecCommandWorkdirArgs; use super::get_command; use super::post_unified_exec_tool_use_payload; @@ -150,6 +153,7 @@ impl ExecCommandHandler { invocation: ToolInvocation, ) -> Result, FunctionCallError> { let ToolInvocation { + environments, session, turn, step_context, @@ -176,13 +180,10 @@ impl ExecCommandHandler { cancellation_token, call_id.clone(), ); - let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?; - let turn_environment = resolve_tool_environment( - &step_context.environments, - environment_args.environment_id.as_deref(), - )?; + let workdir_args: ExecCommandWorkdirArgs = parse_arguments(&arguments)?; + let turn_environment = selected_executor(&environments, &step_context.environments)?; let native_environment_cwd = turn_environment.cwd().clone(); - let cwd = environment_args + let cwd = workdir_args .workdir .as_deref() .filter(|workdir| !workdir.is_empty()) @@ -492,6 +493,13 @@ fn one_shot_exec_command_spec(spec: ToolSpec) -> ToolSpec { } impl CoreToolRuntime for ExecCommandHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + function_tool_environment(invocation, |args: ToolEnvironmentArgs| args.environment_id) + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index 1fecb98b04..ce64c7c004 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -1,10 +1,11 @@ use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; use crate::tools::handlers::parse_arguments; -use crate::tools::handlers::resolve_tool_environment; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::LiveEnvironmentResolver; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolExecutor; @@ -22,8 +23,6 @@ use super::post_unified_exec_tool_use_payload; #[derive(Debug, Deserialize)] struct WriteStdinArgs { - // The model is trained on `session_id`. - session_id: i32, #[serde(default)] chars: String, #[serde(default = "super::default_write_stdin_yield_time_ms")] @@ -32,6 +31,17 @@ struct WriteStdinArgs { max_output_tokens: Option, } +impl WriteStdinArgs { + fn parse(payload: &ToolPayload) -> Result { + let ToolPayload::Function { arguments } = payload else { + return Err(FunctionCallError::RespondToModel( + "write_stdin handler received unsupported payload".to_string(), + )); + }; + parse_arguments(arguments) + } +} + pub struct WriteStdinHandler; impl ToolExecutor for WriteStdinHandler { @@ -67,20 +77,16 @@ impl WriteStdinHandler { cancellation_token, call_id, payload, + environments, .. } = invocation; - let arguments = match payload { - ToolPayload::Function { arguments } => arguments, - _ => { - return Err(FunctionCallError::RespondToModel( - "write_stdin handler received unsupported payload".to_string(), - )); - } + let ToolEnvironments::Terminal(terminal) = environments else { + return Err(FunctionCallError::Fatal( + "write_stdin requires a terminal session".to_string(), + )); }; - - let args: WriteStdinArgs = parse_arguments(&arguments)?; - resolve_tool_environment(&step_context.environments, /*environment_id*/ None)?; + let args = WriteStdinArgs::parse(&payload)?; let context = UnifiedExecContext::new(session.clone(), step_context, cancellation_token, call_id); let response = session @@ -89,7 +95,7 @@ impl WriteStdinHandler { .write_stdin( &context, WriteStdinRequest { - process_id: args.session_id, + terminal: &terminal, input: &args.chars, yield_time_ms: args.yield_time_ms, max_output_tokens: args.max_output_tokens, @@ -124,6 +130,10 @@ impl WriteStdinHandler { } impl CoreToolRuntime for WriteStdinHandler { + fn live_environment_resolver(&self) -> Option { + Some(super::terminal_session_for_input) + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index ecff8f5039..0881f7eaff 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -42,6 +42,7 @@ async fn invocation_for_payload( tool_name: codex_tools::ToolName::plain(tool_name), source: ToolCallSource::Direct, payload, + environments: Default::default(), } } @@ -352,6 +353,7 @@ async fn exec_command_rejects_login_when_selected_environment_disallows_it() { call_id: "login-disallowed".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hello", "login": true }).to_string(), }, @@ -436,6 +438,7 @@ async fn exec_command_reuses_foreign_windows_grant() { call_id: "foreign-windows-grant".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "*** Begin Patch\n*** Add File: granted/file.txt\n+text\n*** End Patch", @@ -477,6 +480,7 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, payload, + environments: Default::default(), }), Some(crate::tools::registry::PreToolUsePayload { tool_name: HookToolName::bash(), @@ -505,6 +509,7 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { tool_name: codex_tools::ToolName::plain("write_stdin"), source: crate::tools::context::ToolCallSource::Direct, payload, + environments: Default::default(), }), None ); diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 903df019f9..1d4e4e6c4b 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -15,12 +15,15 @@ use serde::Deserialize; use crate::function_tool::FunctionCallError; use crate::original_image_detail::can_request_original_image_detail; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::ToolEnvironmentArgs; +use crate::tools::handlers::function_tool_environment; use crate::tools::handlers::parse_arguments; -use crate::tools::handlers::resolve_tool_environment; +use crate::tools::handlers::selected_executor; use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::handlers::view_image_spec::create_view_image_tool; use crate::tools::registry::CoreToolRuntime; @@ -58,8 +61,6 @@ const VIEW_IMAGE_INVALID_MESSAGE: &str = #[derive(Deserialize)] struct ViewImageArgs { path: String, - #[serde(default)] - environment_id: Option, detail: Option, } @@ -108,6 +109,7 @@ impl ViewImageHandler { } let ToolInvocation { + environments, session, turn, step_context, @@ -125,11 +127,7 @@ impl ViewImageHandler { } }; - let ViewImageArgs { - path, - environment_id, - detail, - } = parse_arguments(&arguments)?; + let ViewImageArgs { path, detail } = parse_arguments(&arguments)?; // Keep accepting previously supported detail hints after they disappear from the schema. let detail = match detail.as_deref() { None => None, @@ -142,8 +140,7 @@ impl ViewImageHandler { } }; - let turn_environment = - resolve_tool_environment(&step_context.environments, environment_id.as_deref())?; + let turn_environment = selected_executor(&environments, &step_context.environments)?; let path_uri = turn_environment.cwd().join(&path).map_err(|err| { FunctionCallError::RespondToModel(format!( "unable to resolve image path `{path}` against environment cwd `{}`: {err}", @@ -211,6 +208,13 @@ impl ViewImageHandler { } impl CoreToolRuntime for ViewImageHandler { + fn environments_for_input( + &self, + invocation: &ToolInvocation, + ) -> Result { + function_tool_environment(invocation, |args: ToolEnvironmentArgs| args.environment_id) + } + fn is_builtin_control_tool(&self) -> bool { true } @@ -393,6 +397,7 @@ mod tests { call_id: "call-view-image".to_string(), tool_name: codex_tools::ToolName::plain("view_image"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "path": "image.png" }).to_string(), }, @@ -423,6 +428,7 @@ mod tests { call_id: "call-view-image".to_string(), tool_name: codex_tools::ToolName::plain("view_image"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "path": "image.png", "detail": "low" }).to_string(), }, @@ -466,6 +472,7 @@ mod tests { call_id: "call-view-image".to_string(), tool_name: codex_tools::ToolName::plain("view_image"), source: ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "path": "image.png", "detail": "high" }).to_string(), }, @@ -507,6 +514,7 @@ mod tests { cell_id: "cell-1".to_string(), runtime_tool_call_id: "tool-1".to_string(), }, + environments: Default::default(), payload: ToolPayload::Function { arguments: json!({ "path": "not-an-image.txt" }).to_string(), }, diff --git a/codex-rs/core/src/tools/lifecycle.rs b/codex-rs/core/src/tools/lifecycle.rs index e26d83da6b..81aa6b70b9 100644 --- a/codex-rs/core/src/tools/lifecycle.rs +++ b/codex-rs/core/src/tools/lifecycle.rs @@ -50,7 +50,7 @@ pub(crate) async fn notify_tool_start( originating_item_id: originating_item_id.as_ref(), tool_name: &invocation.tool_name, mcp_tool, - permissions: Box::pin(async { tool_permission_context(invocation).await.ok() }), + permissions: Box::pin(async { tool_permission_context(invocation).ok() }), payload: &invocation.payload, conversation_history: Arc::clone(&conversation_history), source: extension_tool_call_source(invocation.source.clone()), diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 467d54e052..5971d10683 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -10,6 +10,7 @@ use crate::tools::approvals::ApprovalAction; use crate::tools::approvals::ApprovalContext; use crate::tools::events::truncate_rejection_message; use crate::tools::sandboxing::ToolError; +use codex_features::Feature; use codex_network_proxy::BlockedRequest; use codex_network_proxy::BlockedRequestObserver; use codex_network_proxy::EnvironmentNetworkPolicy; @@ -21,9 +22,12 @@ use codex_network_proxy::NetworkProxy; use codex_network_proxy::NetworkRequestCancellation; use codex_network_proxy::NetworkRequestCancellationReason; use codex_network_proxy::NetworkRequestDisconnect; +use codex_network_proxy::RemoteCredentialMasking; +use codex_network_proxy::RemoteNetworkProxyLaunchConfig; use codex_protocol::approvals::NetworkApprovalContext; use codex_protocol::approvals::NetworkApprovalProtocol; use codex_protocol::approvals::NetworkPolicyRuleAction; +use codex_protocol::error::CodexErr; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; @@ -1115,18 +1119,76 @@ pub(crate) async fn begin_network_approval( return Ok(None); } + let fallback_policy_decider = network_policy.as_ref().map(|_| { + build_network_policy_decider( + Arc::clone(&session.services.network_approval), + Arc::new(RwLock::new(Arc::downgrade(session))), + ) + }); + let registration_id = Uuid::new_v4().to_string(); + let Some(execution_proxy) = execution_network_proxy( + session, + turn, + environments + .turn_environments() + .find(|environment| environment.selection.environment_id == environment_id), + network, + &environment_id, + &permission_profile, + network_policy.as_ref(), + ®istration_id, + fallback_policy_decider, + ) + .await? + else { + return Ok(None); + }; + let cancellation_token = CancellationToken::new(); + session + .services + .network_approval + .register_call(ActiveNetworkApprovalCall { + registration_id: registration_id.clone(), + turn_id: turn.sub_id.clone(), + trigger, + tool_name, + command, + environment_id, + permission_profile, + environments: environments.clone(), + cancellation_token: cancellation_token.clone(), + }) + .await; + + Ok(Some(ActiveNetworkApproval { + registration_id: Some(registration_id), + cancellation_token, + execution_proxy, + })) +} + +/// Resolve the issuing environment's network policy under controller requirements and +/// return a proxy with unique execution attribution. Does not register an approval flow. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn execution_network_proxy( + session: &Session, + turn: &TurnContext, + environment: Option<&TurnEnvironment>, + network: Option, + environment_id: &str, + permission_profile: &PermissionProfile, + network_policy: Option<&EnvironmentNetworkPolicy>, + registration_id: &str, + fallback_policy_decider: Option>, +) -> Result, ToolError> { let controller = turn.config.permissions.network.as_ref(); - let environment = environments - .turn_environments() - .find(|environment| environment.selection.environment_id == environment_id); let owner_spec = network_policy - .as_ref() .map(|policy| { // Resolve owner policy once, where the command's actual enforcing proxy is built. NetworkProxySpec::for_environment( controller, policy, - &permission_profile, + permission_profile, session.services.exec_policy.current().as_ref(), environment.map_or( codex_network_proxy::LocalBindingPolicy::DefaultFalse, @@ -1194,49 +1256,62 @@ pub(crate) async fn begin_network_approval( return Ok(None); }; let environment_policy = owner_spec.map(|spec| spec.environment_policy()); - let fallback_policy_decider = environment_policy.as_ref().map(|_| { - build_network_policy_decider( - Arc::clone(&session.services.network_approval), - Arc::new(RwLock::new(Arc::downgrade(session))), - ) - }); - let registration_id = Uuid::new_v4().to_string(); let attribution_token = Uuid::new_v4().to_string(); let execution_proxy = network .for_execution( - &environment_id, - ®istration_id, + environment_id, + registration_id, attribution_token, environment_policy, fallback_policy_decider, ) .map_err(|err| { - ToolError::Codex(codex_protocol::error::CodexErr::Io(io::Error::other( - format!("failed to create execution-scoped network proxy: {err}"), - ))) + ToolError::Codex(CodexErr::Io(io::Error::other(format!( + "failed to create execution-scoped network proxy: {err}" + )))) })?; - let cancellation_token = CancellationToken::new(); - session - .services - .network_approval - .register_call(ActiveNetworkApprovalCall { - registration_id: registration_id.clone(), - turn_id: turn.sub_id.clone(), - trigger, - tool_name, - command, - environment_id, - permission_profile, - environments: environments.clone(), - cancellation_token: cancellation_token.clone(), - }) - .await; + Ok(Some(execution_proxy)) +} - Ok(Some(ActiveNetworkApproval { - registration_id: Some(registration_id), - cancellation_token, - execution_proxy, - })) +/// Resolve a proxy launch against the executor's capabilities, failing closed if required. +pub(crate) async fn remote_network_proxy_launch( + environment: &TurnEnvironment, + network: &NetworkProxy, + turn: &TurnContext, +) -> Result<(RemoteNetworkProxyLaunchConfig, RemoteCredentialMasking), ToolError> { + let network_proxy = environment + .environment + .info() + .await + .map_err(|err| { + ToolError::Codex(CodexErr::Io(io::Error::other(format!( + "failed to query exec-server capabilities: {err}" + )))) + })? + .capabilities + .network_proxy(); + let credential_masking = if turn.config.features.enabled(Feature::CredentialMasking) { + network_proxy.credential_masking + } else { + RemoteCredentialMasking::Disabled + }; + let launch = network + .remote_launch_config( + crate::windows_sandbox::local_binding_policy_for_sandbox( + environment.config().windows_sandbox_type, + environment.executor_platform_os.as_deref(), + ), + credential_masking, + ) + .await + .map_err(|err| ToolError::Codex(CodexErr::Io(io::Error::other(err))))?; + if launch.proxy.enabled && !network_proxy.launch { + return Err(ToolError::Rejected( + "selected exec-server does not support executor-local network proxy launches" + .to_string(), + )); + } + Ok((launch, credential_masking)) } pub(crate) async fn finish_deferred_network_approval( diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index ebde5a7bdc..c60feef0be 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -15,6 +15,7 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolCallState; +use crate::tools::context::ToolEnvironments; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -47,6 +48,9 @@ use indexmap::IndexMap; use indexmap::map::Entry; use serde_json::Value; +pub(crate) type LiveEnvironmentResolver = + fn(&ToolInvocation) -> BoxFuture<'_, Result>; + pub(crate) type ToolTelemetryTags = Vec<(&'static str, String)>; pub use codex_tools::ToolExecutor; @@ -94,6 +98,19 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { false } + /// Identify the environments selected by the original input. + fn environments_for_input( + &self, + _invocation: &ToolInvocation, + ) -> Result { + Ok(ToolEnvironments::All) + } + + /// Select a shared resolver when the environment belongs to a live resource. + fn live_environment_resolver(&self) -> Option { + None + } + /// Returns the owning server only for MCP-backed tool runtimes. fn mcp_server_name(&self) -> Option<&str> { None @@ -167,22 +184,7 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { invocation: ToolInvocation, updated_input: Value, ) -> Result { - let ToolPayload::Function { .. } = &invocation.payload else { - return Err(FunctionCallError::RespondToModel( - "hook input rewrite received unsupported function tool payload".to_string(), - )); - }; - - let arguments = serde_json::to_string(&updated_input).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize rewritten {} arguments: {err}", - flat_tool_name(&invocation.tool_name) - )) - })?; - Ok(ToolInvocation { - payload: ToolPayload::Function { arguments }, - ..invocation - }) + with_updated_function_input(invocation, updated_input) } /// Creates an optional consumer for streamed tool argument diffs. @@ -191,6 +193,80 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { } } +pub(super) fn with_updated_function_input( + invocation: ToolInvocation, + mut updated_input: Value, +) -> Result { + let ToolPayload::Function { .. } = &invocation.payload else { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite received unsupported function tool payload".to_string(), + )); + }; + + match &invocation.environments { + ToolEnvironments::Arguments(environments) => { + for (field, environment) in environments { + preserve_hook_selector( + &mut updated_input, + field, + serde_json::json!(environment.selection.environment_id), + )?; + } + } + ToolEnvironments::Mcp { server, .. } => { + if let Some(Value::String(value)) = updated_input.get_mut("server") { + *value = value.trim().to_owned(); + if value.is_empty() { + updated_input["server"] = Value::Null; + } + } + preserve_hook_selector(&mut updated_input, "server", serde_json::json!(server))?; + } + ToolEnvironments::Terminal(terminal) => { + preserve_hook_selector( + &mut updated_input, + "session_id", + serde_json::json!(terminal.process_id), + )?; + } + ToolEnvironments::All => {} + } + let arguments = serde_json::to_string(&updated_input).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize rewritten {} arguments: {err}", + flat_tool_name(&invocation.tool_name) + )) + })?; + Ok(ToolInvocation { + payload: ToolPayload::Function { arguments }, + ..invocation + }) +} + +/// A rewrite may omit the selector or repeat its original value. +fn preserve_hook_selector( + input: &mut Value, + field: &str, + expected: Value, +) -> Result<(), FunctionCallError> { + let object = input.as_object_mut().ok_or_else(|| { + FunctionCallError::RespondToModel( + "expected an object for rewritten tool arguments".to_string(), + ) + })?; + if let Some(value) = object.get(field) + && value != &expected + { + return Err(FunctionCallError::RespondToModel(format!( + "PreToolUse hook cannot change the tool environment ({field})" + ))); + } + if !expected.is_null() { + object.insert(field.to_string(), expected); + } + Ok(()) +} + /// Consumes streamed argument diffs for a tool call and emits protocol events /// derived from partial tool input. pub(crate) trait ToolArgumentDiffConsumer: Send { @@ -698,6 +774,30 @@ impl ToolRegistry { return Err(err); } + let fail_before_execution = async |invocation: &ToolInvocation, err: FunctionCallError| { + if tool.is_builtin_control_tool() { + let mut analytics = ControlToolCallGuard::new(invocation); + analytics.finish(ControlToolCallStatus::Failed); + } + dispatch_trace.record_failed(&err); + notify_tool_finish_if_unclaimed( + invocation, + call_state.as_deref(), + ToolCallOutcome::Failed { + handler_executed: false, + }, + ) + .await; + err + }; + let environments = match tool.live_environment_resolver() { + Some(resolve) => resolve(&invocation).await, + None => tool.environments_for_input(&invocation), + }; + invocation.environments = match environments { + Ok(environments) => environments, + Err(err) => return Err(fail_before_execution(&invocation, err).await), + }; if let Some(pre_tool_use_payload) = tool.pre_tool_use_payload(&invocation) { match run_pre_tool_use_hooks( &invocation.session, @@ -705,6 +805,7 @@ impl ToolRegistry { invocation.call_id.clone(), &pre_tool_use_payload.tool_name, &pre_tool_use_payload.tool_input, + &invocation.environments.hook_ids(), ) .or_cancel_tool(&cancellation_token) .await? @@ -730,22 +831,7 @@ impl ToolRegistry { Ok(updated_invocation) => { invocation = updated_invocation; } - Err(err) => { - if tool.is_builtin_control_tool() { - let mut analytics = ControlToolCallGuard::new(&invocation); - analytics.finish(ControlToolCallStatus::Failed); - } - dispatch_trace.record_failed(&err); - notify_tool_finish_if_unclaimed( - &invocation, - call_state.as_deref(), - ToolCallOutcome::Failed { - handler_executed: false, - }, - ) - .await; - return Err(err); - } + Err(err) => return Err(fail_before_execution(&invocation, err).await), }, PreToolUseHookResult::Continue { updated_input: None, @@ -826,6 +912,7 @@ impl ToolRegistry { post_tool_use_payload.tool_name.matcher_aliases().to_vec(), post_tool_use_payload.tool_input, post_tool_use_payload.tool_response, + &invocation.environments.hook_ids(), ) .or_cancel_tool(&cancellation_token) .await?, diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 895d8d4b99..0cf26b084a 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -837,6 +837,7 @@ fn test_invocation( call_id: call_id.to_string(), tool_name, source: crate::tools::context::ToolCallSource::Direct, + environments: Default::default(), payload: ToolPayload::Function { arguments: "{}".to_string(), }, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 3c01ab9818..36fbfe4427 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -385,6 +385,7 @@ impl ToolRouter { tool_name, source, payload, + environments: Default::default(), }; self.registry diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index f0ba8205d4..44a6ab9b35 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -20,6 +20,7 @@ use crate::shell_snapshot::ShellSnapshotSandbox; use crate::shell_snapshot::snapshot_read_permissions; use crate::tools::flat_tool_name; use crate::tools::network_approval::NetworkApprovalSpec; +use crate::tools::network_approval::remote_network_proxy_launch; use crate::tools::runtimes::RuntimePathPrepends; use crate::tools::runtimes::exec_env_for_sandbox_permissions; use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot; @@ -342,44 +343,16 @@ impl<'a> ToolRuntime for UnifiedExecRunt }; let (mut env, managed_network_context, network_proxy_launch) = match managed_network { Some(network) if environment_is_remote => { - let environment_info = - req.turn_environment - .environment - .info() - .await - .map_err(|err| { - ToolError::Codex(CodexErr::Io(io::Error::other(format!( - "failed to query exec-server capabilities: {err}" - )))) - })?; - let network_proxy = environment_info.capabilities.network_proxy(); - let credential_masking = if ctx - .step_context - .turn - .config - .features - .enabled(Feature::CredentialMasking) - { - network_proxy.credential_masking - } else { - RemoteCredentialMasking::Disabled - }; - let mut launch = network - .remote_launch_config( - crate::windows_sandbox::local_binding_policy_for_sandbox( - req.turn_environment.config().windows_sandbox_type, - req.turn_environment.executor_platform_os.as_deref(), - ), - credential_masking, - ) - .await - .map_err(|err| { - ToolError::Codex(CodexErr::Io(io::Error::other(err.to_string()))) - })?; + let (mut launch, credential_masking) = remote_network_proxy_launch( + &req.turn_environment, + network, + &ctx.step_context.turn, + ) + .await?; warn_if_remote_credential_masking_unavailable( ctx, req, - network_proxy.credential_masking, + credential_masking, &launch.proxy, ) .await; @@ -404,17 +377,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt ) })?); } - if !launch.proxy.enabled { - (env, None, None) - } else { - if !network_proxy.launch { - return Err(ToolError::Rejected( - "selected exec-server does not support executor-local network proxy launches" - .to_string(), - )); - } - (env, None, Some(launch)) - } + (env, None, launch.proxy.enabled.then_some(launch)) } Some(network) => { let prepared = snapshot_credential_context diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index df8c510778..4365be3df2 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -366,6 +366,7 @@ fn test_invocation_with_payload( tool_name, source, payload, + environments: Default::default(), } } diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index c4ea0991ab..4c47abb2de 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -52,6 +52,7 @@ mod head_tail_buffer; mod oneshot; mod process; mod process_manager; +pub(crate) use process_manager::exec_env_policy_from_shell_policy; mod process_state; mod shell_snapshot; mod stdin_approval; @@ -125,9 +126,18 @@ pub(crate) struct ExecCommandRequest { pub prefix_rule: Option>, } +/// The process and environment behind a terminal session ID at the time of lookup. +#[derive(Clone, Debug)] +pub(crate) struct TerminalSession { + pub(crate) process_id: i32, + process: Arc, + call_id: String, + pub(crate) environment_id: String, +} + #[derive(Debug)] pub(crate) struct WriteStdinRequest<'a> { - pub process_id: i32, + pub terminal: &'a TerminalSession, pub input: &'a str, pub yield_time_ms: u64, pub max_output_tokens: Option, diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 20f2236668..5fb49dbd25 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -315,9 +315,9 @@ async fn write_stdin( input: &str, yield_time_ms: u64, ) -> Result { - session - .services - .unified_exec_manager + let manager = &session.services.unified_exec_manager; + let terminal = manager.resolve_terminal(process_id).await?; + manager .write_stdin( &UnifiedExecContext::new( Arc::clone(session), @@ -326,7 +326,7 @@ async fn write_stdin( "write".to_string(), ), WriteStdinRequest { - process_id, + terminal: &terminal, input, yield_time_ms, max_output_tokens: None, @@ -671,8 +671,12 @@ async fn terminating_initial_exec_command_rechecks_initial_response_state() -> a Ok(()) } +#[test_case::test_case(false; "terminal removed")] +#[test_case::test_case(true; "session id reused")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Result<()> { +async fn terminating_during_stdin_poll_returns_exited_response( + reuse_id: bool, +) -> anyhow::Result<()> { let (session, turn) = test_session_and_turn().await; let manager = &session.services.unified_exec_manager; let process_id = manager.allocate_process_id().await; @@ -743,7 +747,20 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu .await .expect("poll should clone process handles"); - manager.release_process_id(process_id).await; + if reuse_id { + let replacement = super::process_tests::remote_process( + WriteStatus::Accepted, + /*terminate_error*/ None, + SandboxType::None, + ) + .await; + let mut store = manager.process_store.lock().await; + let entry = store.processes.get_mut(&process_id).unwrap(); + entry.process = Arc::new(replacement); + entry.call_id = "replacement".into(); + } else { + manager.release_process_id(process_id).await; + } allow_terminate.notify_one(); process.terminate_confirmed().await?; @@ -751,8 +768,19 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu .await .expect("poll should finish") .expect("poll task should not panic")?; - assert_eq!(output.process_id, None); - assert!(manager.process_store.lock().await.processes.is_empty()); + assert_eq!( + (output.process_id, output.event_call_id.as_str()), + (None, "call") + ); + assert_eq!( + manager + .process_store + .lock() + .await + .processes + .contains_key(&process_id), + reuse_id + ); Ok(()) } @@ -871,7 +899,7 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()> .await?; let process_id = opened.process_id.expect("running terminal"); let cwd = PathUri::parse("file:///C:/workspace")?; - let original = { + { let mut store = manager.process_store.lock().await; let entry = store.processes.get_mut(&process_id).unwrap(); entry.permissions = TerminalPermissions::for_launch( @@ -886,8 +914,9 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()> ); entry.environment_id = "unselected-executor".to_string(); entry.cwd = cwd.clone(); - Arc::clone(&entry.process) - }; + } + let original = manager.resolve_terminal(process_id).await?; + assert_eq!(original.environment_id, "unselected-executor"); // Empty polling must complete without an approval response. // The test deadline must allow the minimum empty-poll wait. tokio::time::timeout( @@ -943,7 +972,13 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()> ), ("call", Some("write"), cwd.clone().into()) ); - assert!(original.interaction_lock().try_lock_owned().is_err()); + assert!( + original + .process + .interaction_lock() + .try_lock_owned() + .is_err() + ); if input == "replace\n" { let replacement = super::process_tests::remote_process( WriteStatus::Accepted, @@ -979,9 +1014,32 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()> matches!(result, Err(UnifiedExecError::UnknownProcessId { process_id: id }) if id == process_id) ), } - assert!(original.interaction_lock().try_lock_owned().is_ok()); + assert!(original.process.interaction_lock().try_lock_owned().is_ok()); } - original.terminate(); + // Reusing the numeric ID must also reject a subsequent write to the original terminal. + let context = UnifiedExecContext::new( + Arc::clone(&session), + crate::session::step_context::StepContext::for_test(turn), + tokio_util::sync::CancellationToken::new(), + "stale".to_string(), + ); + let stale = manager + .write_stdin( + &context, + WriteStdinRequest { + terminal: &original, + input: "", + yield_time_ms: 250, + max_output_tokens: None, + truncation_policy: TruncationPolicy::Tokens(10_000), + interaction_event: None, + }, + ) + .await; + assert!( + matches!(stale, Err(UnifiedExecError::UnknownProcessId { process_id: id }) if id == process_id) + ); + original.process.terminate(); assert!(session.terminate_background_terminal(process_id).await); Ok(()) } @@ -1010,27 +1068,18 @@ async fn stdin_approval_observes_strict_review_enabled_while_queued() -> anyhow: ) .await?; let process_id = opened.process_id.expect("running terminal"); - let process = Arc::clone( - &manager - .process_store - .lock() - .await - .processes - .get(&process_id) - .unwrap() - .process, - ); + let terminal = manager.resolve_terminal(process_id).await?; let context = UnifiedExecContext::new( Arc::clone(&session), StepContext::for_test(turn), tokio_util::sync::CancellationToken::new(), "write".to_string(), ); - let interaction = process.interaction_lock().lock_owned().await; + let interaction = terminal.process.interaction_lock().lock_owned().await; let mut queued = Box::pin(manager.write_stdin( &context, WriteStdinRequest { - process_id, + terminal: &terminal, // NUL input must be rejected when strict review becomes required. input: "\0", yield_time_ms: 250, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 5c4a429b3f..0ba92cb486 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -52,6 +52,7 @@ use crate::unified_exec::MIN_EMPTY_YIELD_TIME_MS; use crate::unified_exec::MIN_YIELD_TIME_MS; use crate::unified_exec::ProcessEntry; use crate::unified_exec::ProcessStore; +use crate::unified_exec::TerminalSession; use crate::unified_exec::UnifiedExecContext; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcessManager; @@ -134,7 +135,7 @@ pub(super) fn apply_unified_exec_env(mut env: HashMap) -> HashMa env } -pub(super) fn exec_env_policy_from_shell_policy( +pub(crate) fn exec_env_policy_from_shell_policy( policy: &ShellEnvironmentPolicy, ) -> codex_exec_server::ExecEnvPolicy { let mut exclude = policy @@ -477,13 +478,26 @@ impl UnifiedExecProcessManager { } pub(crate) async fn release_process_id(&self, process_id: i32) { + self.release_process(process_id, /*expected_process*/ None) + .await; + } + + async fn release_process( + &self, + process_id: i32, + expected_process: Option<&Arc>, + ) { let removed = { let mut store = self.process_store.lock().await; - let entry = store.processes.remove(&process_id); - if !should_use_deterministic_process_ids() { - store.reserved_process_ids.remove(&process_id); + if let Some(expected) = expected_process + && !store + .processes + .get(&process_id) + .is_some_and(|entry| Arc::ptr_eq(&entry.process, expected)) + { + return; } - entry + store.remove(process_id) }; if let Some(entry) = removed { unregister_network_approval_for_entry(&entry).await; @@ -732,7 +746,10 @@ impl UnifiedExecProcessManager { } else { OutputRetention::UntilCollected }; - match self.refresh_process_state(process_id, retention).await { + match self + .refresh_process_state(process_id, &process, retention) + .await + { ProcessStatus::Alive { exit_code, process_id, @@ -871,7 +888,7 @@ impl UnifiedExecProcessManager { turn_id = trace_id(&context.step_context.turn.sub_id), call_id = trace_id(&context.call_id), original_exec_call_id = tracing::field::Empty, - unified_exec_process_id = request.process_id, + unified_exec_process_id = request.terminal.process_id, interaction = if request.input.is_empty() { "poll" } else { "write" }, outcome = tracing::field::Empty, ) @@ -897,23 +914,12 @@ impl UnifiedExecProcessManager { context: &UnifiedExecContext, request: WriteStdinRequest<'_>, ) -> Result { - let process_id = request.process_id; - - // Different terminal sessions can be polled concurrently, but reads and - // writes against one terminal must not overlap because they share a - // draining output buffer and process lifecycle. - let locked_process = { - let store = self.process_store.lock().await; - let entry = store - .processes - .get(&process_id) - .ok_or(UnifiedExecError::UnknownProcessId { process_id })?; - // Capture the original call even if this interaction is cancelled while queued. - if let Some(call_id) = trace_id(&entry.call_id) { - tracing::Span::current().record("original_exec_call_id", call_id); - } - Arc::clone(&entry.process) - }; + let process_id = request.terminal.process_id; + let locked_process = &request.terminal.process; + if let Some(call_id) = trace_id(&request.terminal.call_id) { + tracing::Span::current().record("original_exec_call_id", call_id); + } + // Serialize writes and polls on the original terminal. let _interaction_guard = locked_process.interaction_lock().lock_owned().await; // A queued write must observe strict review enabled while it was waiting. let strict_auto_review = context.step_context.turn.strict_auto_review_enabled(); @@ -922,7 +928,7 @@ impl UnifiedExecProcessManager { let entry = store .processes .get(&process_id) - .filter(|entry| Arc::ptr_eq(&entry.process, &locked_process)) + .filter(|entry| Arc::ptr_eq(&entry.process, locked_process)) .ok_or(UnifiedExecError::UnknownProcessId { process_id })?; entry.stdin_approval(context, request.input, strict_auto_review)? }; @@ -994,7 +1000,7 @@ impl UnifiedExecProcessManager { tty, .. } = self - .prepare_process_handles(process_id, &locked_process) + .prepare_process_handles(process_id, locked_process) .await?; context .session @@ -1026,7 +1032,7 @@ impl UnifiedExecProcessManager { if !process.has_exited() { if matches!(err, UnifiedExecError::ProcessFailed { .. }) { process.terminate(); - self.release_process_id(process_id).await; + self.release_process(process_id, Some(locked_process)).await; } return Err(err); } @@ -1065,7 +1071,7 @@ impl UnifiedExecProcessManager { let message = network_denial_message_for_session(session.as_ref(), network_approval.clone()) .await; - self.release_process_id(process_id).await; + self.release_process(process_id, Some(locked_process)).await; return Err(fail_process_with_message(process.as_ref(), message)); } if let Some(message) = process.failure_message() { @@ -1074,7 +1080,7 @@ impl UnifiedExecProcessManager { network_approval.clone(), ) .await; - self.release_process_id(process_id).await; + self.release_process(process_id, Some(locked_process)).await; if let Err(message) = finish_result { return Err(fail_process_with_message(process.as_ref(), message)); } @@ -1083,7 +1089,7 @@ impl UnifiedExecProcessManager { // Only finalize remote sessions after both process exit and output collection. let status = self - .refresh_process_state(process_id, OutputRetention::UntilCollected) + .refresh_process_state(process_id, locked_process, OutputRetention::UntilCollected) .await; let (process_id, exit_code, event_call_id) = match status { ProcessStatus::Alive { @@ -1105,9 +1111,7 @@ impl UnifiedExecProcessManager { if process.has_exited() { (None, process.exit_code(), call_id) } else { - return Err(UnifiedExecError::UnknownProcessId { - process_id: request.process_id, - }); + return Err(UnifiedExecError::UnknownProcessId { process_id }); } } }; @@ -1134,7 +1138,7 @@ impl UnifiedExecProcessManager { call_id: response.event_call_id.clone(), process_id: response .process_id - .unwrap_or(request.process_id) + .unwrap_or(request.terminal.process_id) .to_string(), stdin: request.input.to_string(), }; @@ -1149,10 +1153,15 @@ impl UnifiedExecProcessManager { async fn refresh_process_state( &self, process_id: i32, + expected_process: &Arc, retention: OutputRetention, ) -> ProcessStatus { let mut store = self.process_store.lock().await; - let Some(entry) = store.processes.get_mut(&process_id) else { + let Some(entry) = store + .processes + .get_mut(&process_id) + .filter(|entry| Arc::ptr_eq(&entry.process, expected_process)) + else { return ProcessStatus::Unknown; }; @@ -1847,14 +1856,21 @@ impl UnifiedExecProcessManager { } } - /// Resolves stdin's target from the host-owned terminal, not model-supplied metadata. - pub(crate) async fn environment_id_for_process(&self, process_id: i32) -> Option { - self.process_store - .lock() - .await + pub(crate) async fn resolve_terminal( + &self, + process_id: i32, + ) -> Result { + let store = self.process_store.lock().await; + let entry = store .processes .get(&process_id) - .map(|entry| entry.environment_id.clone()) + .ok_or(UnifiedExecError::UnknownProcessId { process_id })?; + Ok(TerminalSession { + process_id, + process: Arc::clone(&entry.process), + call_id: entry.call_id.clone(), + environment_id: entry.environment_id.clone(), + }) } pub(crate) async fn list_processes(&self) -> Vec { diff --git a/codex-rs/core/tests/suite/guardian_connector_trust.rs b/codex-rs/core/tests/suite/guardian_connector_trust.rs index 2c190cf68e..7aea947096 100644 --- a/codex-rs/core/tests/suite/guardian_connector_trust.rs +++ b/codex-rs/core/tests/suite/guardian_connector_trust.rs @@ -19,11 +19,15 @@ use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::SEARCH_CALENDAR_LIST_TOOL; use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; use core_test_support::apps_test_server::apps_enabled_builder; +use core_test_support::exec_server::ExecServerProcess; use core_test_support::responses; use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestEnv; +use core_test_support::test_codex::test_env; use core_test_support::wait_for_mcp_server; use serde_json::Value; use serde_json::json; +use test_case::test_case; use wiremock::Mock; use wiremock::matchers::body_partial_json; use wiremock::matchers::method; @@ -53,8 +57,12 @@ impl McpServerContributor for ExtensionOwnedAppsServer { } } +#[test_case(false; "selected executor")] +#[test_case(true; "host MCP server with only a remote executor")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn extension_owned_apps_cannot_attest_connector_to_guardian() -> Result<()> { +async fn extension_owned_apps_cannot_attest_connector_to_guardian( + use_remote_executor: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -120,6 +128,13 @@ async fn extension_owned_apps_cannot_attest_connector_to_guardian() -> Result<() extensions.mcp_server_contributor(Arc::new(ExtensionOwnedAppsServer { url: format!("{}/api/codex/ps/mcp", apps.chatgpt_base_url), })); + let executor; + let environment = if use_remote_executor { + executor = ExecServerProcess::start().await?; + TestEnv::local_with_exec_server_url(Some(executor.websocket_url.clone())).await? + } else { + test_env().await? + }; let test = apps_enabled_builder(apps.chatgpt_base_url) .with_model("test-gpt-5.1-codex") .with_extensions(Arc::new(extensions.build())) @@ -147,7 +162,7 @@ async fn extension_owned_apps_cannot_attest_connector_to_guardian() -> Result<() config.approvals_reviewer = ApprovalsReviewer::AutoReview; config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); }) - .build_with_auto_env(&server) + .build_with_environment(&server, environment) .await?; wait_for_mcp_server(&test.codex, "codex_apps").await?; test.submit_text_turn("List today's calendar events.") diff --git a/codex-rs/core/tests/suite/hooks_executor.rs b/codex-rs/core/tests/suite/hooks_executor.rs index 9d44e567af..b58e169093 100644 --- a/codex-rs/core/tests/suite/hooks_executor.rs +++ b/codex-rs/core/tests/suite/hooks_executor.rs @@ -27,9 +27,11 @@ use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::user_input::UserInput; use codex_utils_path_uri::PathUri; +use core_test_support::TestTargetOs; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::apps_enabled_builder; use core_test_support::apps_test_server::recorded_apps_tool_calls; +use core_test_support::exec_server::ExecServerProcess; use core_test_support::responses::ResponseMock; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -41,11 +43,13 @@ use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::skip_if_remote; +use core_test_support::skip_if_wine_exec; use core_test_support::submit_thread_settings; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::environment_config_for_selection; use core_test_support::test_codex::test_codex; +use core_test_support::test_target_os; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; @@ -67,6 +71,196 @@ use wiremock::matchers::method; use wiremock::matchers::path; use super::rmcp_client::remote_aware_environment_id; +use super::rmcp_client::remote_aware_stdio_server_bin; + +#[test_case( + json!({"permissionDecision": "deny", "permissionDecisionReason": "blocked by owner"}), + "blocked by owner"; "camelCase environment selects the owners blocking hook" +)] +#[test_case( + json!({"permissionDecision": "allow", "updatedInput": {"permissions": {"network": {"enabled": true}}}}), + "\"scope\":\"turn\""; "rewrite may omit the original environment" +)] +#[test_case( + json!({"permissionDecision": "allow", "updatedInput": {"environmentId": "another", "permissions": {"network": {"enabled":true}}}}), + "PreToolUse hook cannot change the tool environment"; "rewrite cannot change the original environment" +)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_tool_use_preserves_request_permissions_environment( + hook_output: Value, + expected_message: &str, +) -> Result<()> { + let server = start_mock_server().await; + let test = test_codex() + .with_config(|config| { + for feature in [Feature::CodexHooks, Feature::RequestPermissionsTool] { + config.features.enable(feature).expect("enable feature"); + } + }) + .build_with_auto_env(&server) + .await?; + let (selections, _executor) = environments_with_pre_tool_use_hook( + &test, + HookEnvironment::Secondary, + "^request_permissions$", + hook_output, + ) + .await?; + let args = json!({"environmentId": selections[1].environment_id, "permissions": {"network": {"enabled": true}}}); + let responses = mount_response_sequence( + &server, + vec![ + sse_response(sse(vec![ + ev_function_call("permissions", "request_permissions", &args.to_string()), + ev_completed("first"), + ])), + completed_turn_response("second"), + ], + ) + .await; + test.submit_turn_with_environments("request network access", Some(selections)) + .await?; + let output = responses + .function_call_output_text("permissions") + .context("expected request_permissions result")?; + assert!(output.contains(expected_message), "{output}"); + Ok(()) +} + +#[test_case(HookEnvironment::Secondary; "unrelated environment cannot block a targeted resource call")] +#[test_case(HookEnvironment::Primary; "server owner can block a targeted resource call")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_resource_hooks_follow_server_ownership( + hook_environment: HookEnvironment, +) -> Result<()> { + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + let server = start_mock_server().await; + let command = remote_aware_stdio_server_bin()?; + let test = test_codex() + .with_config(move |config| { + config + .features + .enable(Feature::CodexHooks) + .expect("enable hooks"); + let mcp_server: McpServerConfig = serde_json::from_value(json!({ + "command": command, + "environment_id": remote_aware_environment_id(), + "cwd": config.cwd, + })) + .expect("valid docs MCP server configuration"); + config + .mcp_servers + .set(HashMap::from([("docs".to_string(), mcp_server)])) + .expect("configure docs MCP server"); + }) + .build_with_auto_env(&server) + .await?; + wait_for_mcp_server(&test.codex, "docs").await?; + + let denied = "denied by resource hook"; + let expected = match hook_environment { + HookEnvironment::Primary => denied, + HookEnvironment::Secondary => "memo://codex/", + }; + let (mut selections, _executor) = environments_with_pre_tool_use_hook( + &test, + hook_environment, + "^(read|list)_mcp_resource", + json!({"permissionDecision": "deny", "permissionDecisionReason": denied}), + ) + .await?; + // Put the server's environment second so routing every call to the primary fails this test too. + selections.reverse(); + let read = json!({"server": "docs", "uri": "memo://codex/example-note"}); + let list = json!({"server": " docs "}); + let calls = [ + ("read", "read_mcp_resource", read, expected), + ("list", "list_mcp_resources", list.clone(), expected), + ("templates", "list_mcp_resource_templates", list, expected), + ("list-all", "list_mcp_resources", json!({}), denied), + ( + "unknown", + "list_mcp_resources", + json!({"server": "missing"}), + "unknown MCP server 'missing'", + ), + ]; + let mut events = vec![ev_response_created("resources")]; + events.extend( + calls + .iter() + .map(|(id, tool, args, _)| ev_function_call(id, tool, &args.to_string())), + ); + events.push(ev_completed("resources")); + let responses = mount_response_sequence( + &server, + vec![sse_response(sse(events)), completed_turn_response("done")], + ) + .await; + test.submit_turn_with_environments("inspect MCP resources", Some(selections)) + .await?; + for (id, _, _, expected) in calls { + let output = responses + .function_call_output_text(id) + .context("missing resource tool result")?; + assert!( + output.contains(expected), + "{id}: expected {expected:?} in {output}" + ); + } + Ok(()) +} + +enum HookEnvironment { + Primary, + Secondary, +} + +async fn environments_with_pre_tool_use_hook( + test: &TestCodex, + hook_environment: HookEnvironment, + matcher: &str, + mut hook_output: Value, +) -> Result<(Vec, Option)> { + let mut executor = None; + let executor_url = if let Some(url) = test.executor_environment().exec_server_url() { + url + } else { + &executor + .insert(ExecServerProcess::start().await?) + .websocket_url + }; + let mut primary = test.executor_environment().selection().clone(); + let mut secondary = primary.clone(); + secondary.environment_id = "secondary".to_string(); + test.thread_manager + .environment_manager() + .upsert_environment( + secondary.environment_id.clone(), + executor_url.to_owned(), + /*connect_timeout*/ None, + )?; + + let owner = match hook_environment { + HookEnvironment::Primary => &mut primary, + HookEnvironment::Secondary => &mut secondary, + }; + hook_output["hookEventName"] = json!("PreToolUse"); + let output = json!({"hookSpecificOutput": hook_output}); + let command = match test_target_os() { + TestTargetOs::Windows => format!("ForEach-Object {{ '{output}' }}"), + TestTargetOs::Linux | TestTargetOs::MacOs => format!("printf '%s' '{output}'"), + }; + let mut config = environment_config_for_selection(&test.config, owner); + config.environment_hooks = serde_json::from_value(json!({ + "PreToolUse": [{ + "matcher": matcher, + "hooks": [{"type": "command", "command": command}], + }], + }))?; + owner.config = EnvironmentConfigState::Ready(config); + Ok((vec![primary, secondary], executor)) +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn environment_requires_loadable_managed_hooks_at_startup_and_readiness() -> Result<()> { diff --git a/codex-rs/ext/multi-env-copy-file/src/tool.rs b/codex-rs/ext/multi-env-copy-file/src/tool.rs index a7cd0e2213..ac60a7e405 100644 --- a/codex-rs/ext/multi-env-copy-file/src/tool.rs +++ b/codex-rs/ext/multi-env-copy-file/src/tool.rs @@ -8,8 +8,10 @@ use codex_tools::JsonSchema; use codex_tools::JsonToolOutput; use codex_tools::ResponsesApiTool; use codex_tools::ToolCall; +use codex_tools::ToolEnvironment; use codex_tools::ToolExecutor; use codex_tools::ToolName; +use codex_tools::ToolPayload; use codex_tools::ToolSpec; use codex_utils_path_uri::LegacyAppPathString; use futures::StreamExt; @@ -30,6 +32,13 @@ struct CopyFileArgs { target_path: String, } +impl CopyFileArgs { + fn parse(arguments: &str) -> Result { + serde_json::from_str(arguments) + .map_err(|error| FunctionCallError::RespondToModel(error.to_string())) + } +} + impl<'call> ToolExecutor> for CopyFileTool { fn tool_name(&self) -> ToolName { ToolName::plain(TOOL_NAME) @@ -77,54 +86,61 @@ impl<'call> ToolExecutor> for CopyFileTool { }) } + fn environment_arguments( + &self, + payload: &ToolPayload, + ) -> Result, FunctionCallError> { + let ToolPayload::Function { arguments } = payload else { + return Err(FunctionCallError::Fatal( + "expected copy arguments".to_string(), + )); + }; + let args = CopyFileArgs::parse(arguments)?; + if args.source_environment_id == args.target_environment_id { + return Err(FunctionCallError::RespondToModel( + "source and target environment IDs must differ. To copy within one environment, use a shell command.".to_string(), + )); + } + Ok(vec![ + ("source_environment_id", args.source_environment_id), + ("target_environment_id", args.target_environment_id), + ]) + } + fn handle<'a>(&'a self, call: ToolCall<'call>) -> codex_tools::ToolExecutorFuture<'a> where 'call: 'a, { Box::pin(async move { - let args: CopyFileArgs = serde_json::from_str(call.function_arguments()?) - .map_err(|error| FunctionCallError::RespondToModel(error.to_string()))?; - // Leave copies within one environment to tools like `cp`. They already handle - // different paths to the same file, including through hard links and symlinks. - if args.source_environment_id == args.target_environment_id { - return Err(FunctionCallError::RespondToModel( - "source and target environment IDs must differ. To copy within one environment, use a shell command.".to_string(), + let args = CopyFileArgs::parse(call.function_arguments()?)?; + let [source, target] = call.environments.as_slice() else { + return Err(FunctionCallError::Fatal( + "copy requires a source and a target environment".to_string(), )); - } - let resolve_endpoint = |id: &str, path: &str| { - let environment = call - .environments - .iter() - .find(|environment| environment.environment_id == id) - .ok_or_else(|| { - FunctionCallError::RespondToModel(format!( - "unknown or unavailable environment ID: {id}" - )) - })?; + }; + let resolve_path = |environment: &ToolEnvironment<'_>, path: &str| { + let id = &environment.environment_id; let convention = environment.cwd.infer_path_convention().ok_or_else(|| { FunctionCallError::RespondToModel(format!( "invalid working directory path for environment {id}" )) })?; - let path = LegacyAppPathString::from_string(path) + LegacyAppPathString::from_string(path) .to_path_uri(convention) .map_err(|_| { FunctionCallError::RespondToModel(format!( "provide an absolute {convention} file path for environment {id}" )) - })?; - // The host binds this accessor to the endpoint's sandbox policy and current grants. - Ok::<_, FunctionCallError>((environment.fs(), path)) + }) }; - let (source, source_path) = - resolve_endpoint(&args.source_environment_id, &args.source_path)?; - let (target, target_path) = - resolve_endpoint(&args.target_environment_id, &args.target_path)?; + let source_path = resolve_path(source, &args.source_path)?; + let target_path = resolve_path(target, &args.target_path)?; // The target writer pulls one block at a time when it supports streaming writes. // For an older executor it buffers this stream and sends a single write instead. // Report only the I/O error kind to the model. Executor connection errors can // include credentials from a signed websocket URL in their message. let contents = source + .fs() .read_file_stream(&source_path) .await .map_err(|error| { @@ -134,6 +150,7 @@ impl<'call> ToolExecutor> for CopyFileTool { )) })?; let bytes_copied = target + .fs() .write_file_stream(&target_path, contents.boxed()) .await .map_err(|error| { diff --git a/codex-rs/ext/multi-env-copy-file/tests/suite/transfers.rs b/codex-rs/ext/multi-env-copy-file/tests/suite/transfers.rs index 5209ba2aa2..519b27f4e1 100644 --- a/codex-rs/ext/multi-env-copy-file/tests/suite/transfers.rs +++ b/codex-rs/ext/multi-env-copy-file/tests/suite/transfers.rs @@ -11,6 +11,7 @@ use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::context_snapshot::SnapshotEntry; use core_test_support::responses::ResponsesRequest; use pretty_assertions::assert_eq; +use test_case::test_case; /// The selected executor may be local, Docker, or Wine; the second executor is host-native. /// One flow covers same-environment rejection and both directions across executor boundaries. @@ -285,3 +286,151 @@ async fn copy_file_transfers_between_local_and_foreign_executors() -> Result<()> assert_eq!(fs::read(local_root.path().join("returned.bin"))?, contents); Ok(()) } + +#[test_case(false; "rewrite paths within the selected environments")] +#[test_case(true; "reject swapping source and target")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn copy_hooks_follow_source_and_target(swap_endpoints: bool) -> Result<()> { + use codex_protocol::models::PermissionProfile; + use codex_protocol::models::PermissionProfileSnapshot; + use codex_protocol::permissions::NetworkSandboxPolicy; + use core_test_support::exec_server::ExecServerProcess; + + let server = start_mock_server().await; + let mut extensions = ExtensionRegistryBuilder::new(); + codex_multi_env_copy_file::install(&mut extensions); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config + .features + .enable(Feature::CodexHooks) + .expect("enable hooks"); + }) + .build_with_auto_env(&server) + .await?; + let executor = ExecServerProcess::start().await?; + let root = TempDir::new()?; + let source = root.path().join("a/source"); + let original_target = root.path().join("b/original"); + let target = root.path().join("b/rewritten"); + let mut rewrite = json!({"source_path": source, "target_path": target}); + if swap_endpoints { + rewrite["source_environment_id"] = json!("b"); + rewrite["target_environment_id"] = json!("a"); + } + let hook_output = json!({"hookSpecificOutput": { + "hookEventName": "PreToolUse", "permissionDecision": "allow", "updatedInput": rewrite, + }}); + let manager = test.thread_manager.environment_manager(); + let mut selections = Vec::new(); + // Select the target first to check that the copy still reads from the source. + for id in ["b", "c", "a"] { + let dir = root.path().join(id); + fs::create_dir(&dir)?; + manager.upsert_environment( + id.to_string(), + executor.websocket_url.clone(), + /*connect_timeout*/ None, + )?; + manager + .get_environment(id) + .context("test executor")? + .wait_until_ready() + .await?; + let cwd = PathUri::from_host_native_path(&dir)?; + let mut selection = TurnEnvironmentSelection { + environment_id: id.to_string(), + cwd: cwd.clone(), + workspace_roots: vec![cwd], + selected_capability_roots: Default::default(), + config: EnvironmentConfigState::FromThread, + }; + let mut config = environment_config_for_selection(&test.config, &selection); + config.permission_profile = + PermissionProfileSnapshot::legacy(PermissionProfile::workspace_write_with( + &[], + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + )); + let hook = |marker, output: &serde_json::Value| { + let output = output + .to_string() + .replace('\'', if cfg!(windows) { "''" } else { "'\\''" }); + let command = if cfg!(windows) { + // The hook receives its event on the PowerShell pipeline. + format!( + "Out-Null; Add-Content -Path {marker} -Value run -ErrorAction Stop; Write-Output '{output}'" + ) + } else { + format!("echo run >> {marker}; printf '%s' '{output}'") + }; + json!([{"matcher":"multi_env_copy_file","hooks":[{"type":"command","command": command}]}]) + }; + let pass = json!({}); + config.environment_hooks = serde_json::from_value(json!({ + "PreToolUse": hook("pre", if id == "a" { &hook_output } else { &pass }), + "PostToolUse": hook("post", &pass), + }))?; + selection.config = EnvironmentConfigState::Ready(config); + selections.push(selection); + } + fs::write(&source, b"from a")?; + let args = json!({ + "source_environment_id": "a", "source_path": source, + "target_environment_id": "b", "target_path": original_target, + }); + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call_with_namespace( + "copy", + "functions", + "multi_env_copy_file", + &args.to_string(), + ), + ev_completed("copy"), + ]), + sse(vec![ + ev_assistant_message("done", "done"), + ev_completed("done"), + ]), + ], + ) + .await; + test.submit_turn_with_environments("Copy the file.", Some(selections)) + .await?; + let output = mock + .function_call_output_text("copy") + .context("copy result")?; + if swap_endpoints { + assert!( + output.contains("cannot change the tool environment"), + "{output}" + ); + assert!(!target.exists()); + } else { + assert_eq!( + serde_json::from_str::(&output) + .with_context(|| format!("expected a successful copy, got {output}"))?, + json!({"bytes_copied":6}) + ); + assert_eq!(fs::read(&target)?, b"from a"); + } + assert!(!original_target.exists()); + for id in ["a", "b", "c"] { + for marker in ["pre", "post"] { + let path = root.path().join(id).join(marker); + if id != "c" && (marker == "pre" || !swap_endpoints) { + let contents = fs::read_to_string(&path) + .with_context(|| format!("read hook marker {}", path.display()))?; + assert_eq!(contents.lines().count(), 1); + } else { + assert!(!path.exists(), "unexpected hook at {}", path.display()); + } + } + } + Ok(()) +} diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 5af16c3b7d..f2a1bca43a 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -123,6 +123,15 @@ pub trait ToolExecutor: Send + Sync { false } + /// Return the argument names and environment IDs used by this call. The host supplies their + /// filesystem accessors in the same order. An empty list keeps the turn-wide environment view. + fn environment_arguments( + &self, + _payload: &crate::ToolPayload, + ) -> Result, FunctionCallError> { + Ok(Vec::new()) + } + /// Handles one invocation without retaining capabilities borrowed by the host. fn handle<'a>(&'a self, invocation: Invocation) -> ToolExecutorFuture<'a> where