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::tools::handlers::McpHandler;
use crate::tools::registry::CoreToolRuntime;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::ResponsesApiTool;
use codex_extension_api::ToolCall as ExtensionToolCall;
use codex_extension_api::ToolExecutor;
use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceTool;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_tools::ResponsesApiNamespace;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_tools::default_namespace_description;
use core_test_support::responses::strip_response_item_ids_from_json;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio_util::sync::CancellationToken;
use super::ToolCall;
use super::ToolCallSource;
use super::ToolRouter;
use super::ToolRouterParams;
use super::extension_tool_executors;
struct ExtensionEchoContributor;
impl codex_extension_api::ToolContributor for ExtensionEchoContributor {
fn tools(
&self,
_session_store: &ExtensionData,
_thread_store: &ExtensionData,
_step_store: &ExtensionData,
) -> Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>> {
vec![Arc::new(ExtensionEchoExecutor)]
}
}
struct ExtensionEchoExecutor;
impl ToolExecutor<ExtensionToolCall> for ExtensionEchoExecutor {
fn tool_name(&self) -> ToolName {
ToolName::namespaced("extension/", "echo")
}
fn spec(&self) -> ToolSpec {
ToolSpec::Namespace(ResponsesApiNamespace {
name: "extension/".to_string(),
description: default_namespace_description("extension/"),
tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool {
name: "echo".to_string(),
description: "Echoes arguments through an extension tool.".to_string(),
strict: true,
parameters: codex_extension_api::parse_tool_input_schema(&json!({
"type": "object",
"properties": {
"message": { "type": "string" },
},
"required": ["message"],
"additionalProperties": false,
}))
.expect("extension schema should parse"),
output_schema: None,
defer_loading: None,
})],
})
}
fn handle(&self, call: ExtensionToolCall) -> codex_tools::ToolExecutorFuture<'_> {
Box::pin(self.handle_call(call))
}
}
impl ExtensionEchoExecutor {
async fn handle_call(
&self,
call: ExtensionToolCall,
) -> Result<Box<dyn codex_tools::ToolOutput>, codex_tools::FunctionCallError> {
let arguments: serde_json::Value =
serde_json::from_str(call.function_arguments()?).expect("test arguments should parse");
Ok(Box::new(codex_tools::JsonToolOutput::new(json!({
"arguments": arguments,
"callId": call.call_id,
"conversationHistory": call.conversation_history.items(),
"ok": true,
}))) as Box<dyn codex_tools::ToolOutput>)
}
}
fn extension_tool_test_registry() -> Arc<ExtensionRegistry<Config>> {
let mut builder = ExtensionRegistryBuilder::new();
builder.tool_contributor(Arc::new(ExtensionEchoContributor));
Arc::new(builder.build())
}
#[tokio::test]
async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let router = ToolRouter::from_context(
step_context.as_ref(),
ToolRouterParams {
tool_suggest_candidates: None,
tool_runtimes: Vec::new(),
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
},
&Default::default(),
);
let parallel_tool_name = ["exec_command", "shell_command"]
.into_iter()
.find(|name| {
router.tool_supports_parallel(&ToolCall {
tool_name: ToolName::plain(*name),
call_id: "call-parallel-tool".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
})
})
.expect("test session should expose a parallel shell-like tool");
assert!(!router.tool_supports_parallel(&ToolCall {
tool_name: ToolName::namespaced("mcp__server__", parallel_tool_name),
call_id: "call-namespaced-tool".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
}));
Ok(())
}
#[tokio::test]
async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<()> {
let tool_name = "create_event".to_string();
let call = ToolRouter::build_tool_call(ResponseItem::FunctionCall {
id: None,
name: tool_name.clone(),
namespace: Some("mcp__codex_apps__calendar".to_string()),
arguments: "{}".to_string(),
call_id: "call-namespace".to_string(),
internal_chat_message_metadata_passthrough: None,
})?
.expect("function_call should produce a tool call");
assert_eq!(
call.tool_name,
ToolName::namespaced("mcp__codex_apps__calendar", tool_name)
);
assert_eq!(call.call_id, "call-namespace");
match call.payload {
ToolPayload::Function { arguments } => {
assert_eq!(arguments, "{}");
}
other => panic!("expected function payload, got {other:?}"),
}
Ok(())
}
#[tokio::test]
async fn build_custom_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<()> {
let tool_name = "exec".to_string();
let call = ToolRouter::build_tool_call(ResponseItem::CustomToolCall {
id: None,
status: None,
call_id: "call-namespace".to_string(),
name: tool_name.clone(),
namespace: Some("mcp__python".to_string()),
input: "print('hello')".to_string(),
internal_chat_message_metadata_passthrough: None,
})?
.expect("custom_tool_call should produce a tool call");
assert_eq!(
call,
ToolCall {
tool_name: ToolName::namespaced("mcp__python", tool_name),
call_id: "call-namespace".to_string(),
payload: ToolPayload::Custom {
input: "print('hello')".to_string(),
},
}
);
Ok(())
}
#[tokio::test]
async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let router = ToolRouter::from_context(
step_context.as_ref(),
ToolRouterParams {
tool_suggest_candidates: None,
tool_runtimes: vec![
mcp_runtime(mcp_tool_info(
"echo",
true,
"mcp__echo__",
"query_with_delay",
)),
mcp_runtime(mcp_tool_info(
"hello_echo",
false,
"mcp__hello_echo__",
"query_with_delay",
)),
],
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
},
&Default::default(),
);
let call = ToolCall {
tool_name: ToolName::namespaced("mcp__echo__", "query_with_delay"),
call_id: "call-handler".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
assert!(router.tool_supports_parallel(&call));
let different_server_call = ToolCall {
tool_name: ToolName::namespaced("mcp__hello_echo__", "query_with_delay"),
call_id: "call-other-server".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
assert!(!router.tool_supports_parallel(&different_server_call));
Ok(())
}
#[tokio::test]
async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let router = ToolRouter::from_context(
step_context.as_ref(),
ToolRouterParams {
tool_suggest_candidates: None,
tool_runtimes: Vec::new(),
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
},
&Default::default(),
);
assert!(!router.tool_supports_parallel(&ToolCall {
tool_name: ToolName::plain("web_search"),
call_id: "call-web-search".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
}));
Ok(())
}
#[tokio::test]
async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let hidden_tool = "hidden_dynamic_tool";
let visible_tool = "visible_dynamic_tool";
let dynamic_tools = vec![DynamicToolSpec::Namespace(DynamicToolNamespaceSpec {
name: "codex_app".to_string(),
description: "Codex app tools.".to_string(),
tools: vec![
DynamicToolNamespaceTool::Function(DynamicToolFunctionSpec {
name: hidden_tool.to_string(),
description: "Hidden until discovered.".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"additionalProperties": false,
}),
defer_loading: true,
}),
DynamicToolNamespaceTool::Function(DynamicToolFunctionSpec {
name: visible_tool.to_string(),
description: "Visible immediately.".to_string(),
input_schema: json!({
"type": "object",
"properties": {},
"additionalProperties": false,
}),
defer_loading: false,
}),
],
})];
let router = ToolRouter::from_context(
step_context.as_ref(),
ToolRouterParams {
tool_suggest_candidates: None,
tool_runtimes: Vec::new(),
extension_tool_executors: Vec::new(),
dynamic_tools: &dynamic_tools,
},
&Default::default(),
);
assert_eq!(
namespace_function_names(&router.model_visible_specs(), "codex_app"),
vec![visible_tool.to_string()]
);
Ok(())
}
fn mcp_tool_info(
server_name: &str,
supports_parallel_tool_calls: bool,
callable_namespace: &str,
tool_name: &str,
) -> codex_mcp::ToolInfo {
codex_mcp::ToolInfo {
server_name: server_name.to_string(),
supports_parallel_tool_calls,
server_origin: None,
callable_name: tool_name.to_string(),
callable_namespace: callable_namespace.to_string(),
namespace_description: None,
tool: rmcp::model::Tool::new(
tool_name.to_string(),
"Test MCP tool",
Arc::new(rmcp::model::object(json!({
"type": "object",
}))),
),
openai_file_input_optional_fields: Default::default(),
connector_id: None,
connector_name: None,
plugin_display_names: Vec::new(),
}
}
fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> Arc<dyn CoreToolRuntime> {
Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build"))
}
#[tokio::test]
async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow::Result<()> {
let (mut session, turn) = make_session_and_context().await;
session.services.extensions = extension_tool_test_registry();
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let history_item = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "extension history".to_string(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
};
session
.record_conversation_items(&turn, std::slice::from_ref(&history_item))
.await;
let mut expected_history_item = history_item.clone();
expected_history_item.set_turn_id_if_missing(&turn.sub_id);
let router = ToolRouter::from_context(
step_context.as_ref(),
ToolRouterParams {
tool_suggest_candidates: None,
tool_runtimes: Vec::new(),
extension_tool_executors: extension_tool_executors(&session, step_context.as_ref()),
dynamic_tools: turn.dynamic_tools.as_slice(),
},
&Default::default(),
);
assert!(
router.model_visible_specs().iter().any(
|spec| matches!(spec, ToolSpec::Namespace(namespace)
if namespace.name == "extension/"
&& namespace.tools.iter().any(|tool| matches!(
tool,
ResponsesApiNamespaceTool::Function(tool) if tool.name == "echo"
)))
),
"expected extension-provided tool to be visible to the model"
);
let call = ToolRouter::build_tool_call(ResponseItem::FunctionCall {
id: None,
name: "echo".to_string(),
namespace: Some("extension/".to_string()),
arguments: json!({ "message": "hello" }).to_string(),
call_id: "call-extension".to_string(),
internal_chat_message_metadata_passthrough: None,
})?
.expect("function_call should produce a tool call");
let result = router
.dispatch_tool_call_with_code_mode_result(
Arc::new(session),
step_context,
CancellationToken::new(),
Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call,
ToolCallSource::Direct,
)
.await?;
let response = result.into_response();
match response {
ResponseInputItem::FunctionCallOutput { call_id, output } => {
assert_eq!(call_id, "call-extension");
let FunctionCallOutputBody::Text(text) = output.body else {
panic!("expected text function call output")
};
let value: serde_json::Value =
serde_json::from_str(&text).expect("extension tool output should be json");
assert_eq!(
strip_response_item_ids_from_json(value),
strip_response_item_ids_from_json(json!({
"arguments": { "message": "hello" },
"callId": "call-extension",
"conversationHistory": [expected_history_item],
"ok": true,
}))
);
}
other => panic!("expected function call output, got {other:?}"),
}
Ok(())
}
fn namespace_function_names(specs: &[ToolSpec], namespace_name: &str) -> Vec<String> {
specs
.iter()
.find_map(|spec| match spec {
ToolSpec::Namespace(namespace) if namespace.name == namespace_name => Some(
namespace
.tools
.iter()
.map(|tool| match tool {
ResponsesApiNamespaceTool::Function(tool) => tool.name.clone(),
})
.collect(),
),
ToolSpec::Function(_)
| ToolSpec::Freeform(_)
| ToolSpec::ToolSearch { .. }
| ToolSpec::WebSearch { .. }
| ToolSpec::Namespace(_) => None,
})
.unwrap_or_default()
}