mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
core: use current step environments for tools
This commit is contained in:
@@ -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?;
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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())),
|
||||
|
||||
@@ -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())),
|
||||
|
||||
@@ -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())),
|
||||
|
||||
@@ -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())),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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())),
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user