core: use current step environments for tools

This commit is contained in:
Sayan Sisodiya
2026-06-23 01:04:36 -07:00
parent f394ec2f7a
commit a24d9864e4
34 changed files with 281 additions and 56 deletions
+2
View File
@@ -18,6 +18,7 @@ use crate::hook_runtime::run_pre_compact_hooks;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::responses_metadata::CompactionTurnMetadata;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn::built_tools;
use crate::session::turn_context::TurnContext;
use codex_analytics::CompactionImplementation;
@@ -222,6 +223,7 @@ async fn run_remote_compact_task_inner_impl(
let tool_router = built_tools(
sess.as_ref(),
turn_context.as_ref(),
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
&CancellationToken::new(),
)
.await?;
+2
View File
@@ -21,6 +21,7 @@ use crate::responses_metadata::CompactionTurnMetadata;
use crate::responses_retry::ResponsesStreamRequest;
use crate::responses_retry::handle_retryable_response_stream_error;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn::built_tools;
use crate::session::turn_context::TurnContext;
use codex_analytics::CompactionImplementation;
@@ -230,6 +231,7 @@ async fn run_remote_compact_task_inner_impl(
let tool_router = built_tools(
sess.as_ref(),
turn_context.as_ref(),
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
&CancellationToken::new(),
)
.await?;
+8 -1
View File
@@ -14,6 +14,7 @@ use tokio_util::sync::CancellationToken;
use crate::config::Config;
use crate::resolve_installation_id;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn::build_prompt;
use crate::session::turn::built_tools;
use crate::state_db_bridge::StateDbHandle;
@@ -90,7 +91,13 @@ pub(crate) async fn build_prompt_input_from_session(
.clone_history()
.await
.for_prompt(&turn_context.model_info.input_modalities);
let router = built_tools(sess, turn_context.as_ref(), &CancellationToken::new()).await?;
let router = built_tools(
sess,
turn_context.as_ref(),
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
&CancellationToken::new(),
)
.await?;
let base_instructions = sess.get_base_instructions().await;
let prompt = build_prompt(
prompt_input,
+18 -1
View File
@@ -387,7 +387,6 @@ use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TurnModerationMetadataEvent;
use codex_protocol::protocol::WarningEvent;
use codex_protocol::user_input::UserInput;
use codex_tools::ToolEnvironmentMode;
use codex_tools::UnifiedExecShellMode;
use codex_utils_absolute_path::AbsolutePathBuf;
#[cfg(test)]
@@ -2788,6 +2787,24 @@ impl Session {
world_state
}
pub(crate) async fn capture_step_context(
&self,
turn_context: &TurnContext,
) -> Arc<StepContext> {
let step_context = if turn_context
.config
.features
.enabled(Feature::DeferredExecutor)
{
StepContext {
environments: self.services.turn_environments.snapshot().await,
}
} else {
StepContext::from_turn_context(turn_context)
};
Arc::new(step_context)
}
pub(crate) async fn record_inter_agent_communication(
&self,
turn_context: &TurnContext,
@@ -1,7 +1,16 @@
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::session::turn_context::TurnContext;
/// Request-scoped state that may change between model sampling requests.
#[derive(Debug)]
pub(crate) struct StepContext {
pub(crate) environments: TurnEnvironmentSnapshot,
}
impl StepContext {
pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self {
Self {
environments: turn_context.environments.clone(),
}
}
}
+9 -2
View File
@@ -8,6 +8,7 @@ use crate::context::ContextualUserFragment;
use crate::context::TurnAborted;
use crate::environment_selection::ThreadEnvironments;
use crate::function_tool::FunctionCallError;
use crate::session::step_context::StepContext;
use crate::shell::default_user_shell;
use crate::shell_snapshot::ShellSnapshot;
use crate::skills::SkillRenderSideEffects;
@@ -595,8 +596,9 @@ async fn preview_session_start_hooks(
}
fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> ToolCallRuntime {
let router = Arc::new(ToolRouter::from_turn_context(
let router = Arc::new(ToolRouter::from_contexts(
&turn_context,
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
crate::tools::router::ToolRouterParams {
tool_suggest_candidates: None,
mcp_tools: None,
@@ -5899,6 +5901,7 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir
handler
.handle(ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
turn: turn_context,
cancellation_token: CancellationToken::new(),
tracker,
@@ -5973,6 +5976,7 @@ async fn request_permissions_tool_rejects_unknown_environment_id() {
let result = RequestPermissionsHandler
.handle(ToolInvocation {
session: Arc::new(session),
step_context: Arc::new(StepContext::from_turn_context(&turn_context)),
turn: Arc::new(turn_context),
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
@@ -9898,8 +9902,9 @@ async fn fatal_tool_error_stops_turn_and_reports_error() {
.await
};
let deferred_mcp_tools = Some(tools.clone());
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn_context,
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
crate::tools::router::ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools,
@@ -10155,6 +10160,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() {
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::clone(&turn_diff_tracker),
call_id,
@@ -10308,6 +10314,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request()
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::clone(&tracker),
call_id: "exec-call".to_string(),
@@ -3,6 +3,7 @@ use crate::compact::InitialContextInjection;
use crate::exec_policy::ExecPolicyManager;
use crate::guardian::GUARDIAN_REVIEWER_NAME;
use crate::sandboxing::SandboxPermissions;
use crate::session::step_context::StepContext;
use crate::test_support::models_manager_with_provider;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolOutput;
@@ -333,6 +334,7 @@ async fn guardian_allows_shell_command_additional_permissions_requests_past_poli
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call_id: "test-call".to_string(),
@@ -437,6 +439,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call_id: "strict-shell-command-call".to_string(),
@@ -484,6 +487,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::clone(&tracker),
call_id: "exec-call".to_string(),
@@ -612,6 +616,7 @@ async fn shell_command_allows_sticky_turn_permissions_without_inline_request_per
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
step_context: Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call_id: "sticky-turn-grant".to_string(),
+13 -5
View File
@@ -244,14 +244,12 @@ pub(crate) async fn run_turn(
)
.await?;
let step_context = sess.capture_step_context(turn_context.as_ref()).await;
if turn_context
.config
.features
.enabled(Feature::DeferredExecutor)
{
let step_context = StepContext {
environments: sess.services.turn_environments.snapshot().await,
};
world_state = sess
.record_step_environment_context_if_changed(
turn_context.as_ref(),
@@ -278,6 +276,7 @@ pub(crate) async fn run_turn(
run_sampling_request(
Arc::clone(&sess),
Arc::clone(&turn_context),
step_context,
Arc::clone(&turn_extension_data),
Arc::clone(&turn_diff_tracker),
&mut client_session,
@@ -1130,6 +1129,7 @@ pub(crate) fn build_prompt(
async fn run_sampling_request(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
step_context: Arc<StepContext>,
turn_store: Arc<codex_extension_api::ExtensionData>,
turn_diff_tracker: SharedTurnDiffTracker,
client_session: &mut ModelClientSession,
@@ -1137,7 +1137,13 @@ async fn run_sampling_request(
input: Vec<ResponseItem>,
cancellation_token: CancellationToken,
) -> CodexResult<(SamplingRequestResult, Vec<ResponseItem>)> {
let router = built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token).await?;
let router = built_tools(
sess.as_ref(),
turn_context.as_ref(),
step_context,
&cancellation_token,
)
.await?;
let base_instructions = sess.get_base_instructions().await;
@@ -1234,6 +1240,7 @@ async fn run_sampling_request(
pub(crate) async fn built_tools(
sess: &Session,
turn_context: &TurnContext,
step_context: Arc<StepContext>,
cancellation_token: &CancellationToken,
) -> CodexResult<Arc<ToolRouter>> {
let mcp_connection_manager = sess.services.mcp_connection_manager.load_full();
@@ -1351,8 +1358,9 @@ pub(crate) async fn built_tools(
);
let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools);
let deferred_mcp_tools = mcp_tool_exposure.deferred_tools;
Ok(Arc::new(ToolRouter::from_turn_context(
Ok(Arc::new(ToolRouter::from_contexts(
turn_context,
step_context,
ToolRouterParams {
mcp_tools,
deferred_mcp_tools,
@@ -207,10 +207,6 @@ impl TurnContext {
&& self.config.orchestrator_mcp_enabled
}
pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode {
ToolEnvironmentMode::from_count(self.environments.turn_environments.len())
}
pub(crate) async fn with_model(
&self,
model: String,
@@ -14,6 +14,7 @@ use crate::guardian::routes_approval_to_guardian;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::session::INITIAL_SUBMIT_ID;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn::build_prompt;
use crate::session::turn::built_tools;
use codex_otel::STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC;
@@ -261,6 +262,9 @@ async fn schedule_startup_prewarm_inner(
let startup_router = built_tools(
session.as_ref(),
startup_turn_context.as_ref(),
Arc::new(StepContext::from_turn_context(
startup_turn_context.as_ref(),
)),
&startup_cancellation_token,
)
.await?;
@@ -8,6 +8,7 @@ use super::image_generation_artifact_path;
use super::last_assistant_message_from_item;
use super::response_item_may_include_external_context;
use super::save_image_generation_result;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::ToolRouter;
use crate::tools::parallel::ToolCallRuntime;
@@ -279,8 +280,9 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() {
session.services.extensions = Arc::new(builder.build());
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let router = Arc::new(ToolRouter::from_turn_context(
let router = Arc::new(ToolRouter::from_contexts(
&turn_context,
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
crate::tools::router::ToolRouterParams {
tool_suggest_candidates: None,
mcp_tools: None,
+2
View File
@@ -1,6 +1,7 @@
use crate::context_manager::truncate_function_output_payload;
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::tools::TELEMETRY_PREVIEW_MAX_BYTES;
use crate::tools::TELEMETRY_PREVIEW_MAX_LINES;
@@ -54,6 +55,7 @@ pub enum ToolCallSource {
pub struct ToolInvocation {
pub session: Arc<Session>,
pub turn: Arc<TurnContext>,
pub(crate) step_context: Arc<StepContext>,
pub cancellation_token: CancellationToken,
pub tracker: SharedTurnDiffTracker,
pub call_id: String,
@@ -348,6 +348,7 @@ impl ApplyPatchHandler {
let ToolInvocation {
session,
turn,
step_context,
tracker,
call_id,
tool_name,
@@ -372,8 +373,10 @@ impl ApplyPatchHandler {
require_environment_id(args.environment_id.as_deref(), self.multi_environment)?;
// Verify the parsed patch against the selected environment filesystem.
let Some(turn_environment) =
resolve_tool_environment(turn.as_ref(), selected_environment_id.as_deref())?
let Some(turn_environment) = resolve_tool_environment(
&step_context.environments,
selected_environment_id.as_deref(),
)?
else {
return Err(FunctionCallError::RespondToModel(
"apply_patch is unavailable in this session".to_string(),
@@ -13,6 +13,7 @@ use std::sync::Arc;
use tempfile::TempDir;
use tokio::sync::Mutex;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolInvocation;
use crate::tools::hook_names::HookToolName;
@@ -31,6 +32,7 @@ async fn invocation_for_payload(payload: ToolPayload) -> ToolInvocation {
let (session, turn) = make_session_and_context().await;
ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -112,8 +112,9 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {
let conversation_history =
ConversationHistory::new(invocation.session.clone_history().await.into_raw_items());
let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len());
for environment in &invocation.turn.environments.turn_environments {
let mut environments =
Vec::with_capacity(invocation.step_context.environments.turn_environments.len());
for environment in &invocation.step_context.environments.turn_environments {
// TODO(anp): Migrate extension ToolEnvironment and granted-permission lookup to PathUri
// so extensions can receive foreign environment cwd values.
let Ok(native_cwd) = environment.cwd().to_abs_path() else {
@@ -175,6 +176,7 @@ mod tests {
use super::CoreTurnItemEmitter;
use super::ExtensionToolAdapter;
use crate::session::step_context::StepContext;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
@@ -274,6 +276,7 @@ mod tests {
let (session, turn) = crate::session::tests::make_session_and_context().await;
let invocation = ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
@@ -343,6 +346,7 @@ mod tests {
assert_eq!(raw_history_item.item, expected_history_item);
let invocation = ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())),
turn,
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
@@ -544,6 +548,7 @@ mod tests {
);
let invocation = ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())),
turn,
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
+5
View File
@@ -317,6 +317,7 @@ mod search_tests;
#[cfg(test)]
mod tests {
use super::*;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolCallSource;
use crate::tools::hook_names::HookToolName;
@@ -345,6 +346,7 @@ mod tests {
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -377,6 +379,7 @@ mod tests {
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -405,6 +408,7 @@ mod tests {
.with_updated_hook_input(
ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -452,6 +456,7 @@ mod tests {
.expect("MCP tool spec should build");
let invocation = ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
+4 -4
View File
@@ -46,10 +46,10 @@ use serde_json::Map;
use serde_json::Value;
use std::path::Path;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::function_tool::FunctionCallError;
use crate::sandboxing::SandboxPermissions;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::session::turn_context::TurnEnvironment;
pub(crate) use crate::tools::code_mode::CodeModeExecuteHandler;
pub(crate) use crate::tools::code_mode::CodeModeWaitHandler;
@@ -153,13 +153,13 @@ fn resolve_workdir_base_path(
}
fn resolve_tool_environment<'a>(
turn: &'a TurnContext,
environments: &'a TurnEnvironmentSnapshot,
environment_id: Option<&str>,
) -> Result<Option<&'a TurnEnvironment>, FunctionCallError> {
environment_id.map_or_else(
|| Ok(turn.environments.primary()),
|| Ok(environments.primary()),
|environment_id| {
turn.environments
environments
.turn_environments
.iter()
.find(|environment| environment.environment_id == environment_id)
@@ -4,6 +4,7 @@ use crate::config::AgentRoleConfig;
use crate::config::DEFAULT_AGENT_MAX_DEPTH;
use crate::function_tool::FunctionCallError;
use crate::init_state_db;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session_prefix::format_inter_agent_completion_message;
use crate::thread_manager::thread_store_from_config;
@@ -75,6 +76,7 @@ fn invocation(
) -> ToolInvocation {
ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())),
turn,
cancellation_token: CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::default())),
@@ -47,6 +47,7 @@ impl RequestPermissionsHandler {
let ToolInvocation {
session,
turn,
step_context,
cancellation_token,
call_id,
payload,
@@ -63,8 +64,10 @@ impl RequestPermissionsHandler {
};
let environment_args: RequestPermissionsEnvironmentArgs = parse_arguments(&arguments)?;
let Some(turn_environment) =
resolve_tool_environment(turn.as_ref(), environment_args.environment_id.as_deref())?
let Some(turn_environment) = resolve_tool_environment(
&step_context.environments,
environment_args.environment_id.as_deref(),
)?
else {
return Err(FunctionCallError::RespondToModel(
"request_permissions requires a primary environment".to_string(),
@@ -1,4 +1,5 @@
use super::*;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
@@ -27,6 +28,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() {
}
.handle(ToolInvocation {
session: Arc::new(session),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: Arc::new(turn),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::default())),
+3 -5
View File
@@ -8,6 +8,7 @@ use crate::exec::ExecParams;
use crate::exec_policy::ExecApprovalRequest;
use crate::function_tool::FunctionCallError;
use crate::session::turn_context::TurnContext;
use crate::session::turn_context::TurnEnvironment;
use crate::shell::ShellType;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolPayload;
@@ -53,6 +54,7 @@ struct RunExecLikeArgs {
prefix_rule: Option<Vec<String>>,
session: Arc<crate::session::session::Session>,
turn: Arc<TurnContext>,
turn_environment: TurnEnvironment,
tracker: crate::tools::context::SharedTurnDiffTracker,
call_id: String,
shell_runtime_backend: ShellRuntimeBackend,
@@ -69,16 +71,12 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
prefix_rule,
session,
turn,
turn_environment,
tracker,
call_id,
shell_runtime_backend,
} = args;
let Some(turn_environment) = turn.environments.primary() else {
return Err(FunctionCallError::RespondToModel(
"shell is unavailable in this session".to_string(),
));
};
let fs = turn_environment.environment.get_filesystem();
let explicit_env_overrides = turn
@@ -160,6 +160,7 @@ impl ShellCommandHandler {
let ToolInvocation {
session,
turn,
step_context,
cancellation_token,
tracker,
call_id,
@@ -174,6 +175,12 @@ impl ShellCommandHandler {
)));
};
let Some(turn_environment) = step_context.environments.primary().cloned() else {
return Err(FunctionCallError::RespondToModel(
"shell is unavailable in this session".to_string(),
));
};
#[allow(deprecated)]
let cwd = resolve_workdir_base_path(&arguments, &turn.cwd)?;
let params: ShellCommandToolCallParams = parse_arguments_with_base_path(&arguments, &cwd)?;
@@ -205,6 +212,7 @@ impl ShellCommandHandler {
prefix_rule,
session,
turn,
turn_environment,
tracker,
call_id,
shell_runtime_backend: self.shell_runtime_backend(),
@@ -6,6 +6,7 @@ use pretty_assertions::assert_eq;
use crate::exec_env::create_env;
use crate::sandboxing::SandboxPermissions;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::shell::Shell;
use crate::shell::ShellType;
@@ -201,6 +202,7 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() {
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -230,6 +232,7 @@ async fn build_post_tool_use_payload_uses_tool_output_wire_value() {
let (session, turn) = make_session_and_context().await;
let invocation = ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -110,6 +110,7 @@ impl ExecCommandHandler {
let ToolInvocation {
session,
turn,
step_context,
tracker,
call_id,
payload,
@@ -128,8 +129,10 @@ impl ExecCommandHandler {
let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager;
let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone());
let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?;
let Some(turn_environment) =
resolve_tool_environment(turn.as_ref(), environment_args.environment_id.as_deref())?
let Some(turn_environment) = resolve_tool_environment(
&step_context.environments,
environment_args.environment_id.as_deref(),
)?
else {
return Err(FunctionCallError::RespondToModel(
"unified exec is unavailable in this session".to_string(),
@@ -9,6 +9,7 @@ use codex_utils_output_truncation::TruncationPolicy;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ExecCommandToolOutput;
use crate::tools::context::ToolCallSource;
@@ -29,6 +30,7 @@ async fn invocation_for_payload(
let (session, turn) = make_session_and_context().await;
ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -241,6 +243,7 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() {
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -267,6 +270,7 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() {
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
session: session.into(),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -101,6 +101,7 @@ impl ViewImageHandler {
let ToolInvocation {
session,
turn,
step_context,
payload,
call_id,
..
@@ -134,7 +135,7 @@ impl ViewImageHandler {
};
let Some(turn_environment) =
resolve_tool_environment(turn.as_ref(), environment_id.as_deref())?
resolve_tool_environment(&step_context.environments, environment_id.as_deref())?
else {
return Err(FunctionCallError::RespondToModel(
"view_image is unavailable in this session".to_string(),
@@ -253,6 +254,7 @@ impl ToolOutput for ViewImageOutput {
#[cfg(test)]
mod tests {
use super::*;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnEnvironment;
use crate::tools::context::ToolCallSource;
@@ -326,6 +328,7 @@ mod tests {
let result = ViewImageHandler::default()
.handle(ToolInvocation {
session: Arc::new(session),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: Arc::new(turn),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -354,6 +357,7 @@ mod tests {
let result = ViewImageHandler::default()
.handle(ToolInvocation {
session: Arc::new(session),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: Arc::new(turn),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
@@ -389,6 +393,7 @@ mod tests {
let result = ViewImageHandler::default()
.handle(ToolInvocation {
session: Arc::new(session),
step_context: Arc::new(StepContext::from_turn_context(&turn)),
turn: Arc::new(turn),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
+3
View File
@@ -240,6 +240,7 @@ mod tests {
use super::*;
use std::time::Duration;
use crate::session::step_context::StepContext;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolInvocation;
use crate::tools::registry::CoreToolRuntime;
@@ -426,6 +427,7 @@ mod tests {
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, turn_context, tracker);
@@ -498,6 +500,7 @@ mod tests {
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
Arc::new(StepContext::from_turn_context(turn_context.as_ref())),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, turn_context, tracker);
@@ -1,4 +1,5 @@
use super::*;
use crate::session::step_context::StepContext;
use pretty_assertions::assert_eq;
struct TestHandler {
@@ -459,6 +460,7 @@ fn test_invocation(
) -> ToolInvocation {
ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())),
turn,
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(
+17 -3
View File
@@ -1,5 +1,6 @@
use crate::function_tool::FunctionCallError;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::context::ToolInvocation;
@@ -35,6 +36,7 @@ pub struct ToolCall {
pub struct ToolRouter {
registry: ToolRegistry,
model_visible_specs: Vec<ToolSpec>,
step_context: Arc<StepContext>,
}
pub(crate) struct ToolRouterParams<'a> {
@@ -58,18 +60,29 @@ pub(crate) struct ToolSuggestCandidates {
}
impl ToolRouter {
pub(crate) fn from_turn_context(
pub(crate) fn from_contexts(
turn_context: &TurnContext,
step_context: Arc<StepContext>,
params: ToolRouterParams<'_>,
tool_search_handler_cache: &ToolSearchHandlerCache,
) -> Self {
build_tool_router(turn_context, params, tool_search_handler_cache)
build_tool_router(
turn_context,
step_context,
params,
tool_search_handler_cache,
)
}
pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec<ToolSpec>) -> Self {
pub(crate) fn from_parts(
registry: ToolRegistry,
model_visible_specs: Vec<ToolSpec>,
step_context: Arc<StepContext>,
) -> Self {
Self {
registry,
model_visible_specs,
step_context,
}
}
@@ -226,6 +239,7 @@ impl ToolRouter {
let invocation = ToolInvocation {
session,
turn,
step_context: Arc::clone(&self.step_context),
cancellation_token,
tracker,
call_id,
+11 -5
View File
@@ -1,6 +1,7 @@
use std::sync::Arc;
use crate::config::Config;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolPayload;
use crate::turn_diff_tracker::TurnDiffTracker;
@@ -111,8 +112,9 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow
.load_full()
.list_all_tools()
.await;
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools: None,
@@ -179,8 +181,9 @@ async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<()
#[tokio::test]
async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools: None,
@@ -228,8 +231,9 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
#[tokio::test]
async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools: None,
@@ -283,8 +287,9 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> {
],
})];
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools: None,
@@ -348,8 +353,9 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow
let mut expected_history_item = history_item.clone();
expected_history_item.set_turn_id_if_missing(&turn.sub_id);
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: None,
deferred_mcp_tools: None,
+27 -8
View File
@@ -1,5 +1,6 @@
use crate::agent::exceeds_thread_spawn_depth_limit;
use crate::agent::next_thread_spawn_depth;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::tools::code_mode::execute_spec::create_code_mode_tool;
use crate::tools::context::ToolInvocation;
@@ -144,6 +145,7 @@ impl PlannedTools {
#[derive(Clone, Copy)]
struct CoreToolPlanContext<'a> {
turn_context: &'a TurnContext,
step_context: &'a StepContext,
mcp_tools: Option<&'a [ToolInfo]>,
deferred_mcp_tools: Option<&'a [ToolInfo]>,
tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>,
@@ -157,17 +159,23 @@ struct CoreToolPlanContext<'a> {
#[instrument(level = "trace", skip_all)]
pub(crate) fn build_tool_router(
turn_context: &TurnContext,
step_context: Arc<StepContext>,
params: ToolRouterParams<'_>,
tool_search_handler_cache: &ToolSearchHandlerCache,
) -> ToolRouter {
let (model_visible_specs, registry) =
build_tool_specs_and_registry(turn_context, params, tool_search_handler_cache);
ToolRouter::from_parts(registry, model_visible_specs)
let (model_visible_specs, registry) = build_tool_specs_and_registry(
turn_context,
&step_context,
params,
tool_search_handler_cache,
);
ToolRouter::from_parts(registry, model_visible_specs, step_context)
}
#[instrument(level = "trace", skip_all)]
fn build_tool_specs_and_registry(
turn_context: &TurnContext,
step_context: &StepContext,
params: ToolRouterParams<'_>,
tool_search_handler_cache: &ToolSearchHandlerCache,
) -> (Vec<ToolSpec>, ToolRegistry) {
@@ -182,6 +190,7 @@ fn build_tool_specs_and_registry(
crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new());
let context = CoreToolPlanContext {
turn_context,
step_context,
mcp_tools: mcp_tools.as_deref(),
deferred_mcp_tools: deferred_mcp_tools.as_deref(),
tool_suggest_candidates: tool_suggest_candidates.as_ref(),
@@ -625,11 +634,15 @@ fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool {
.enabled(Feature::StandaloneWebSearch))
}
fn tool_environment_mode(step_context: &StepContext) -> ToolEnvironmentMode {
ToolEnvironmentMode::from_count(step_context.environments.turn_environments.len())
}
#[instrument(level = "trace", skip_all)]
fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) {
let turn_context = context.turn_context;
let features = turn_context.config.features.get();
let environment_mode = turn_context.tool_environment_mode();
let environment_mode = tool_environment_mode(context.step_context);
if !environment_mode.has_environment() {
return;
}
@@ -649,7 +662,10 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne
allow_login_shell,
exec_permission_approvals_enabled,
include_environment_id,
include_shell_parameter: unified_exec_should_include_shell_parameter(turn_context),
include_shell_parameter: unified_exec_should_include_shell_parameter(
turn_context,
context.step_context,
),
}));
planned_tools.add(WriteStdinHandler);
@@ -666,11 +682,14 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne
}
}
fn unified_exec_should_include_shell_parameter(turn_context: &TurnContext) -> bool {
fn unified_exec_should_include_shell_parameter(
turn_context: &TurnContext,
step_context: &StepContext,
) -> bool {
!matches!(
&turn_context.unified_exec_shell_mode,
UnifiedExecShellMode::ZshFork(_)
) || turn_context
) || step_context
.environments
.turn_environments
.iter()
@@ -690,7 +709,7 @@ fn add_mcp_resource_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) {
let turn_context = context.turn_context;
let features = turn_context.config.features.get();
let environment_mode = turn_context.tool_environment_mode();
let environment_mode = tool_environment_mode(context.step_context);
planned_tools.add(PlanHandler);
+32 -3
View File
@@ -30,6 +30,7 @@ use codex_tools::ToolSpec;
use pretty_assertions::assert_eq;
use serde_json::json;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tools::handlers::ToolSearchHandlerCache;
@@ -178,8 +179,9 @@ async fn probe_with(
) -> ToolPlanProbe {
let (_session, mut turn) = make_session_and_context().await;
configure_turn(&mut turn);
let router = ToolRouter::from_turn_context(
let router = ToolRouter::from_contexts(
&turn,
Arc::new(StepContext::from_turn_context(&turn)),
ToolRouterParams {
tool_suggest_candidates: inputs.tool_suggest_candidates,
mcp_tools: inputs.mcp_tools,
@@ -639,6 +641,31 @@ async fn environment_count_controls_environment_backed_tools() {
));
}
#[tokio::test]
async fn environment_tools_follow_the_step_context() {
let (_session, mut turn) = make_session_and_context().await;
set_feature(&mut turn, Feature::UnifiedExec, /*enabled*/ true);
turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform);
let step_context = Arc::new(StepContext::from_turn_context(&turn));
turn.environments.turn_environments.clear();
let plan = ToolPlanProbe::from_router(ToolRouter::from_contexts(
&turn,
step_context,
ToolRouterParams {
mcp_tools: None,
deferred_mcp_tools: None,
tool_suggest_candidates: None,
extension_tool_executors: Vec::new(),
dynamic_tools: &[],
},
&Default::default(),
));
plan.assert_visible_contains(&["exec_command", "apply_patch", "view_image"]);
}
#[tokio::test]
async fn host_context_gates_agent_job_tools() {
let normal_agent_job = probe(|turn| {
@@ -773,8 +800,9 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
let (_session, mut first_turn) = make_session_and_context().await;
first_turn.model_info.supports_search_tool = true;
let first_router = ToolRouter::from_turn_context(
let first_router = ToolRouter::from_contexts(
&first_turn,
Arc::new(StepContext::from_turn_context(&first_turn)),
ToolRouterParams {
mcp_tools: None,
deferred_mcp_tools: Some(vec![mcp_tool("first", "mcp__first", "lookup")]),
@@ -788,8 +816,9 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
let (_session, mut second_turn) = make_session_and_context().await;
second_turn.model_info.supports_search_tool = true;
let second_router = ToolRouter::from_turn_context(
let second_router = ToolRouter::from_contexts(
&second_turn,
Arc::new(StepContext::from_turn_context(&second_turn)),
ToolRouterParams {
mcp_tools: None,
deferred_mcp_tools: Some(vec![mcp_tool("second", "mcp__second", "lookup")]),
@@ -13,6 +13,7 @@ use tokio_util::sync::CancellationToken;
use crate::function_tool::FunctionCallError;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tools::code_mode::CodeModeWaitHandler;
@@ -271,6 +272,7 @@ fn test_invocation_with_payload(
) -> ToolInvocation {
ToolInvocation {
session,
step_context: Arc::new(StepContext::from_turn_context(turn.as_ref())),
turn,
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
+51 -5
View File
@@ -547,7 +547,7 @@ async fn serve_environment_info(listener: TcpListener) {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let server = start_mock_server().await;
let user_input_call_id = "wait-for-startup";
@@ -580,11 +580,13 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
sse(vec![
ev_response_created("resp-2"),
ev_function_call(
"update-plan",
"update_plan",
"request-permissions",
"request_permissions",
&json!({
"explanation": "Continue after startup.",
"plan": [{"step": "Finish", "status": "completed"}]
"reason": "Verify that the ready environment is used.",
"permissions": {
"network": { "enabled": true }
}
})
.to_string(),
),
@@ -601,7 +603,17 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
let mut builder = test_codex()
.with_exec_server_url(format!("ws://{}", listener.local_addr()?))
.with_config(|config| {
config.use_experimental_unified_exec_tool = true;
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config.approvals_reviewer = ApprovalsReviewer::User;
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
assert!(config.features.enable(Feature::UnifiedExec).is_ok());
assert!(
config
.features
.enable(Feature::RequestPermissionsTool)
.is_ok()
);
assert!(
config
.features
@@ -645,6 +657,30 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
},
})
.await?;
let event = wait_for_event(&test.codex, |event| {
matches!(
event,
EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_)
)
})
.await;
let EventMsg::RequestPermissions(permission_request) = event else {
panic!("ready environment should be available to request_permissions: {event:?}");
};
assert_eq!(
permission_request.environment_id.as_deref(),
Some(REMOTE_ENVIRONMENT_ID)
);
test.codex
.submit(Op::RequestPermissionsResponse {
id: permission_request.call_id,
response: RequestPermissionsResponse {
permissions: RequestPermissionProfile::default(),
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
@@ -652,6 +688,16 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
let requests = response_mock.requests();
assert_eq!(requests.len(), 3);
let tool_names = |request_index: usize| {
requests[request_index].body_json()["tools"]
.as_array()
.expect("tools array")
.iter()
.filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_owned))
.collect::<Vec<_>>()
};
assert!(!tool_names(0).contains(&"exec_command".to_string()));
assert!(tool_names(1).contains(&"exec_command".to_string()));
assert!(
requests[0]
.message_input_texts("user")