Generalize turn tool availability policy

This commit is contained in:
shijie-openai
2026-06-09 23:40:34 -07:00
parent 7249c9dfe7
commit 203ecbdb0b
8 changed files with 144 additions and 47 deletions
@@ -11,7 +11,7 @@ use crate::tools::handlers::request_user_input_spec::request_user_input_tool_des
use crate::tools::handlers::request_user_input_spec::request_user_input_unavailable_message;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::ToolExecutor;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ToolAvailability;
use codex_protocol::config_types::ModeKind;
use codex_protocol::request_user_input::RequestUserInputArgs;
use codex_tools::ToolName;
@@ -58,9 +58,9 @@ impl ToolExecutor<ToolInvocation> for RequestUserInputHandler {
));
}
if let Some(suppression) = turn.extension_data.get::<RequestUserInputSuppression>() {
if let Some(unavailable) = tool_unavailable_for_turn(turn.as_ref()) {
return Err(FunctionCallError::RespondToModel(
suppression.unavailable_message().to_string(),
unavailable.into_model_message(),
));
}
@@ -96,6 +96,14 @@ impl ToolExecutor<ToolInvocation> for RequestUserInputHandler {
impl CoreToolRuntime for RequestUserInputHandler {}
fn tool_unavailable_for_turn(
turn: &crate::session::turn_context::TurnContext,
) -> Option<codex_extension_api::ToolUnavailable> {
turn.extension_data
.get::<ToolAvailability>()?
.unavailable_reason(&ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME))
}
#[cfg(test)]
#[path = "request_user_input_tests.rs"]
mod tests;
@@ -3,7 +3,8 @@ use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ToolAvailability;
use codex_extension_api::ToolUnavailable;
use codex_protocol::ThreadId;
use codex_protocol::config_types::ModeKind;
use codex_protocol::protocol::SessionSource;
@@ -13,11 +14,17 @@ use serde_json::json;
use std::sync::Arc;
use tokio::sync::Mutex;
const REQUEST_USER_INPUT_UNAVAILABLE_MESSAGE: &str = "request_user_input unavailable for test";
#[tokio::test]
async fn request_user_input_rejects_suppressed_turns() {
let (session, turn) = make_session_and_context().await;
turn.extension_data
.insert(RequestUserInputSuppression::ActiveDefaultModeGoal);
.get_or_init(ToolAvailability::default)
.mark_unavailable(
codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME),
ToolUnavailable::model_message(REQUEST_USER_INPUT_UNAVAILABLE_MESSAGE),
);
let result = RequestUserInputHandler {
available_modes: vec![ModeKind::Default],
@@ -58,11 +65,7 @@ async fn request_user_input_rejects_suppressed_turns() {
};
assert_eq!(
err,
FunctionCallError::RespondToModel(
RequestUserInputSuppression::ActiveDefaultModeGoal
.unavailable_message()
.to_string(),
)
FunctionCallError::RespondToModel(REQUEST_USER_INPUT_UNAVAILABLE_MESSAGE.to_string())
);
}
+14 -5
View File
@@ -43,6 +43,7 @@ use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHand
use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2;
use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2;
use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2;
use crate::tools::handlers::request_user_input_spec::REQUEST_USER_INPUT_TOOL_NAME;
use crate::tools::handlers::view_image_spec::ViewImageToolOptions;
use crate::tools::hosted_spec::WebSearchToolOptions;
use crate::tools::hosted_spec::create_image_generation_tool;
@@ -53,7 +54,7 @@ use crate::tools::registry::ToolRegistry;
use crate::tools::registry::override_tool_exposure;
use crate::tools::router::ToolRouter;
use crate::tools::router::ToolRouterParams;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ToolAvailability;
use codex_features::Feature;
use codex_login::AuthManager;
use codex_mcp::ToolInfo;
@@ -647,7 +648,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
planned_tools.add(PlanHandler);
if turn_context.config.experimental_request_user_input_enabled
&& !request_user_input_suppressed_for_turn(turn_context)
&& !request_user_input_unavailable_for_turn(turn_context)
{
planned_tools.add(RequestUserInputHandler {
available_modes: request_user_input_available_modes(features),
@@ -696,11 +697,19 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
}
}
fn request_user_input_suppressed_for_turn(turn_context: &TurnContext) -> bool {
fn request_user_input_unavailable_for_turn(turn_context: &TurnContext) -> bool {
tool_unavailable_for_turn(turn_context, &ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME))
.is_some()
}
fn tool_unavailable_for_turn(
turn_context: &TurnContext,
tool_name: &ToolName,
) -> Option<codex_extension_api::ToolUnavailable> {
turn_context
.extension_data
.get::<RequestUserInputSuppression>()
.is_some()
.get::<ToolAvailability>()?
.unavailable_reason(tool_name)
}
fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) {
+8 -2
View File
@@ -1,7 +1,8 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ToolAvailability;
use codex_extension_api::ToolUnavailable;
use codex_features::Feature;
use codex_login::AuthManager;
use codex_login::CodexAuth;
@@ -34,6 +35,7 @@ use serde_json::json;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE;
use crate::tools::handlers::request_user_input_spec::REQUEST_USER_INPUT_TOOL_NAME;
use crate::tools::router::ToolRouter;
use crate::tools::router::ToolRouterParams;
@@ -452,7 +454,11 @@ async fn request_user_input_tool_respects_experimental_config_gate() {
async fn request_user_input_tool_respects_turn_suppression() {
let suppressed = probe(|turn| {
turn.extension_data
.insert(RequestUserInputSuppression::ActiveDefaultModeGoal);
.get_or_init(ToolAvailability::default)
.mark_unavailable(
ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME),
ToolUnavailable::model_message("request_user_input unavailable for test"),
);
})
.await;
suppressed.assert_visible_lacks(&["request_user_input"]);
+2 -1
View File
@@ -63,4 +63,5 @@ pub use registry::ExtensionRegistryBuilder;
pub use registry::empty_extension_registry;
pub use state::ExtensionData;
pub use state::ExtensionDataInit;
pub use tool_policy::RequestUserInputSuppression;
pub use tool_policy::ToolAvailability;
pub use tool_policy::ToolUnavailable;
+70 -16
View File
@@ -1,21 +1,75 @@
/// Per-turn policy that suppresses the request_user_input tool.
use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::PoisonError;
use codex_tools::ToolName;
/// Turn-scoped availability policy for tools.
///
/// Extensions insert this marker into turn-scoped [`ExtensionData`](crate::ExtensionData)
/// when they own context that makes blocking user input inappropriate for the
/// current turn.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RequestUserInputSuppression {
/// Suppress request_user_input while Default mode is working on an active goal.
ActiveDefaultModeGoal,
/// Extensions can attach this to turn-scoped [`ExtensionData`](crate::ExtensionData)
/// when they own context that makes one or more tools inappropriate for the
/// current turn. Core tool planning can omit unavailable tools, while handlers
/// can use the same policy for defensive rejection if an unavailable tool is
/// invoked anyway.
#[derive(Debug, Default)]
pub struct ToolAvailability {
unavailable: Mutex<HashMap<ToolName, ToolUnavailable>>,
}
impl RequestUserInputSuppression {
/// Returns the model-facing message to use if the suppressed tool is invoked.
pub fn unavailable_message(self) -> &'static str {
match self {
Self::ActiveDefaultModeGoal => {
"request_user_input is unavailable while the current Default mode turn is working on an active goal"
}
}
impl ToolAvailability {
/// Marks a tool unavailable for the current turn.
pub fn mark_unavailable(
&self,
tool_name: ToolName,
unavailable: ToolUnavailable,
) -> Option<ToolUnavailable> {
self.unavailable().insert(tool_name, unavailable)
}
/// Returns the unavailable reason for `tool_name`, if one exists.
pub fn unavailable_reason(&self, tool_name: &ToolName) -> Option<ToolUnavailable> {
self.unavailable().get(tool_name).cloned()
}
/// Returns true when `tool_name` is marked unavailable for the current turn.
pub fn is_unavailable(&self, tool_name: &ToolName) -> bool {
self.unavailable().contains_key(tool_name)
}
/// Returns true when `tool_name` has not been marked unavailable.
pub fn is_available(&self, tool_name: &ToolName) -> bool {
!self.is_unavailable(tool_name)
}
fn unavailable(&self) -> std::sync::MutexGuard<'_, HashMap<ToolName, ToolUnavailable>> {
self.unavailable
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
}
/// Explains why a tool is unavailable for a turn.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolUnavailable {
model_message: String,
}
impl ToolUnavailable {
/// Creates an unavailable reason whose message is returned to the model if
/// the unavailable tool is invoked anyway.
pub fn model_message(message: impl Into<String>) -> Self {
Self {
model_message: message.into(),
}
}
/// Returns the model-facing unavailable message.
pub fn model_message_text(&self) -> &str {
&self.model_message
}
/// Consumes this reason and returns the model-facing unavailable message.
pub fn into_model_message(self) -> String {
self.model_message
}
}
+13 -2
View File
@@ -7,18 +7,20 @@ use codex_extension_api::ConfigContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ThreadIdleInput;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadResumeInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ThreadStopInput;
use codex_extension_api::TokenUsageContributor;
use codex_extension_api::ToolAvailability;
use codex_extension_api::ToolCallOutcome;
use codex_extension_api::ToolContributor;
use codex_extension_api::ToolFinishInput;
use codex_extension_api::ToolLifecycleContributor;
use codex_extension_api::ToolLifecycleFuture;
use codex_extension_api::ToolName;
use codex_extension_api::ToolUnavailable;
use codex_extension_api::TurnAbortInput;
use codex_extension_api::TurnErrorInput;
use codex_extension_api::TurnLifecycleContributor;
@@ -45,6 +47,9 @@ use crate::spec::UPDATE_GOAL_TOOL_NAME;
use crate::steering::budget_limit_steering_item;
use crate::tool::GoalToolExecutor;
const REQUEST_USER_INPUT_TOOL_NAME: &str = "request_user_input";
const ACTIVE_GOAL_REQUEST_USER_INPUT_UNAVAILABLE_MESSAGE: &str = "request_user_input is unavailable while the current Default mode turn is working on an active goal";
#[derive(Clone, Debug)]
pub struct GoalExtensionConfig {
pub enabled: bool,
@@ -224,7 +229,13 @@ where
{
input
.turn_store
.insert(RequestUserInputSuppression::ActiveDefaultModeGoal);
.get_or_init(ToolAvailability::default)
.mark_unavailable(
ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME),
ToolUnavailable::model_message(
ACTIVE_GOAL_REQUEST_USER_INPUT_UNAVAILABLE_MESSAGE,
),
);
}
if matches!(
@@ -9,15 +9,16 @@ use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::FunctionCallError;
use codex_extension_api::NoopTurnItemEmitter;
use codex_extension_api::RequestUserInputSuppression;
use codex_extension_api::ThreadResumeInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ThreadStopInput;
use codex_extension_api::ToolAvailability;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolCallOutcome;
use codex_extension_api::ToolCallSource;
use codex_extension_api::ToolExecutor;
use codex_extension_api::ToolFinishInput;
use codex_extension_api::ToolName;
use codex_extension_api::ToolPayload;
use codex_extension_api::TurnErrorInput;
use codex_extension_api::TurnStartInput;
@@ -736,10 +737,11 @@ async fn request_user_input_is_suppressed_for_active_default_goal_turn() -> anyh
.start_turn_with_mode("turn-default", ModeKind::Default, &TokenUsage::default())
.await;
assert_eq!(
Some(RequestUserInputSuppression::ActiveDefaultModeGoal),
default_turn_store
.get::<RequestUserInputSuppression>()
.map(|suppression| *suppression)
Some(
"request_user_input is unavailable while the current Default mode turn is working on an active goal"
.to_string()
),
request_user_input_unavailable_message(&default_turn_store)
);
let plan_turn_store = harness
@@ -747,9 +749,7 @@ async fn request_user_input_is_suppressed_for_active_default_goal_turn() -> anyh
.await;
assert_eq!(
None,
plan_turn_store
.get::<RequestUserInputSuppression>()
.map(|suppression| *suppression)
request_user_input_unavailable_message(&plan_turn_store)
);
Ok(())
}
@@ -784,9 +784,7 @@ async fn request_user_input_is_not_suppressed_for_inactive_default_goal_turns()
.await;
assert_eq!(
None,
turn_store
.get::<RequestUserInputSuppression>()
.map(|suppression| *suppression),
request_user_input_unavailable_message(&turn_store),
"status {status:?} should not suppress request_user_input"
);
}
@@ -1413,6 +1411,13 @@ fn tool_by_name<'a>(
.unwrap_or_else(|| panic!("missing tool {name}"))
}
fn request_user_input_unavailable_message(turn_store: &ExtensionData) -> Option<String> {
turn_store
.get::<ToolAvailability>()?
.unavailable_reason(&ToolName::plain("request_user_input"))
.map(|unavailable| unavailable.into_model_message())
}
fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> ToolCall {
ToolCall {
turn_id: "turn-1".to_string(),