diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index e4519817a9..2f98e9e173 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -4,6 +4,8 @@ use crate::tools::context::ToolPayload; use crate::tools::context::boxed_tool_output; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; +use codex_prompts::ResolvedModelMessages; +use codex_tools::IndirectNamespacePrefixes; use codex_tools::ToolName; use codex_tools::ToolSpec; use std::sync::Arc; @@ -66,6 +68,13 @@ impl CodeModeExecuteHandler { } enabled_tools.sort_by(|left, right| left.name.cmp(&right.name)); enabled_tools.dedup_by(|left, right| left.name == right.name); + let model_messages = ResolvedModelMessages::from_model(&step_context.settings.model_info); + IndirectNamespacePrefixes::new( + model_messages.indirect_description_prefixes(), + step_context.tool_router.mcp_namespaces(), + ) + .map_err(|error| FunctionCallError::Fatal(error.to_string()))? + .apply_code_mode(&mut enabled_tools); let started_at = std::time::Instant::now(); let started_cell = exec .session diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 5cbe6e085a..67b782927e 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -12,6 +12,8 @@ use bm25::Document; use bm25::Language; use bm25::SearchEngine; use bm25::SearchEngineBuilder; +use codex_prompts::ResolvedModelMessages; +use codex_tools::IndirectNamespacePrefixes; use codex_tools::LoadableToolSpec; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; @@ -192,7 +194,11 @@ impl ToolSearchHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { - let ToolInvocation { payload, .. } = invocation; + let ToolInvocation { + payload, + step_context, + .. + } = invocation; let args = match payload { ToolPayload::ToolSearch { arguments } => arguments, @@ -221,7 +227,14 @@ impl ToolSearchHandler { return Ok(boxed_tool_output(ToolSearchOutput { tools: Vec::new() })); } - let tools = self.search(query, limit)?; + let mut tools = self.search(query, limit)?; + let model_messages = ResolvedModelMessages::from_model(&step_context.settings.model_info); + IndirectNamespacePrefixes::new( + model_messages.indirect_description_prefixes(), + step_context.tool_router.mcp_namespaces(), + ) + .map_err(|error| FunctionCallError::Fatal(error.to_string()))? + .apply_search(&mut tools); Ok(boxed_tool_output(ToolSearchOutput { tools })) } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 4d3c50c0f9..e70223c39a 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -430,6 +430,13 @@ impl ToolRegistry { self.tools.values_mut() } + /// Returns configured MCP server names and their registered callable namespaces. + pub(crate) fn mcp_namespaces(&self) -> impl Iterator { + self.tools.iter().filter_map(|(name, tool)| { + Some((tool.runtime.mcp_server_name()?, name.namespace.as_deref()?)) + }) + } + pub(crate) fn deferred_tool_namespaces(&self) -> BTreeMap { let mut namespaces = BTreeMap::::new(); for (name, tool) in &self.tools { diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index c9e6db431b..d9f1f63eed 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -210,6 +210,10 @@ impl ToolRouter { self.registry.deferred_tool_namespaces() } + pub(crate) fn mcp_namespaces(&self) -> impl Iterator { + self.registry.mcp_namespaces() + } + #[cfg(test)] pub(crate) fn registered_tool_names_for_test(&self) -> Vec { self.registry.tool_names_for_test() diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 61fdad23f6..ddac78f3cd 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -80,6 +80,7 @@ use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::MultiAgentVersion; +use codex_tools::IndirectNamespacePrefixes; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::TOOL_SEARCH_TOOL_NAME; use codex_tools::ToolCall as ExtensionToolCall; @@ -403,8 +404,14 @@ pub(crate) fn finalize_tool_router( append_tool_search_executor(turn_context, &mut registry, tool_search_handler_cache); } + let model_messages = ResolvedModelMessages::from_model(model_info); + let indirect_prefixes = IndirectNamespacePrefixes::new( + model_messages.indirect_description_prefixes(), + registry.mcp_namespaces(), + ) + .map_err(|error| CodexErrorDetails::InvalidRequest(error.to_string()))?; let code_mode_tool_names = - register_code_mode_executors(turn_context, model_info, &mut registry); + register_code_mode_executors(turn_context, model_info, &mut registry, &indirect_prefixes); let include_tool_namespaces_info = turn_context .config .tool_registry @@ -810,6 +817,7 @@ fn register_code_mode_executors( turn_context: &TurnContext, model_info: &ModelInfo, registry: &mut ToolRegistry, + indirect_prefixes: &IndirectNamespacePrefixes<'_>, ) -> BTreeMap { let tool_mode = effective_tool_mode(turn_context, model_info); if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { @@ -884,15 +892,18 @@ fn register_code_mode_executors( code_mode_nested_tool_specs.push((spec, cached_runtime)); } - let namespace_descriptions = code_mode_namespace_descriptions(&exec_prompt_tool_specs); + let mut namespace_descriptions = code_mode_namespace_descriptions(&exec_prompt_tool_specs); let mut enabled_tools = collect_code_mode_exec_prompt_tool_definitions(exec_prompt_tool_specs.iter()); let deferred_tools = collect_code_mode_exec_prompt_tool_definitions( deferred_exec_prompt_tool_specs.iter().map(Arc::as_ref), ); + let model_messages = ResolvedModelMessages::from_model(model_info); + if tool_mode == ToolMode::CodeModeOnly { + indirect_prefixes.apply_exec_prompt(&mut enabled_tools, &mut namespace_descriptions); + } enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); - let model_messages = ResolvedModelMessages::from_model(model_info); let execute_handler = CodeModeExecuteHandler::new( create_code_mode_tool( &enabled_tools, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 67c77fb907..a70da0f93f 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -211,7 +211,10 @@ fn custom_tool_output_body_and_success( (output, success) } -fn custom_tool_output_last_non_empty_text(req: &ResponsesRequest, call_id: &str) -> Option { +pub(super) fn custom_tool_output_last_non_empty_text( + req: &ResponsesRequest, + call_id: &str, +) -> Option { match req.custom_tool_call_output(call_id).get("output") { Some(Value::String(text)) if !text.trim().is_empty() => Some(text.clone()), Some(Value::Array(items)) => items diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index a355bf966d..46ca14959f 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -85,6 +85,9 @@ mod agent_message_board; #[path = "scenarios_guardian_extra_policy.rs"] mod guardian_extra_policy; +#[path = "scenarios_indirect_namespace_prefixes.rs"] +mod indirect_namespace_prefixes; + #[path = "scenarios_preparation.rs"] mod preparation; diff --git a/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs b/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs new file mode 100644 index 0000000000..f1b96caa31 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs @@ -0,0 +1,274 @@ +//! Catalog namespace and MCP server prefixes follow the current model without changing dispatch. + +use super::super::code_mode::custom_tool_output_last_non_empty_text; +use super::super::rmcp_client::remote_aware_environment_id; +use super::super::rmcp_client::remote_aware_stdio_server_bin; +use codex_core::TurnInputRequest; +use codex_features::Feature; +use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::context_snapshot; +use core_test_support::context_snapshot::ContextSnapshotOptions; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::collections::HashMap; +use test_case::test_case; + +#[test_case(ToolMode::CodeModeOnly; "code_mode_only")] +#[test_case(ToolMode::CodeMode; "code_mode_with_direct_tools")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn catalog_namespace_prefixes_follow_the_selected_model( + tool_mode: ToolMode, +) -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + skip_if_no_network!(Ok(())); + let include_search = tool_mode == ToolMode::CodeMode; + let requests_per_turn = if include_search { 3 } else { 2 }; + let server = responses::start_mock_server().await; + let command = remote_aware_stdio_server_bin()?; + let environment_id = remote_aware_environment_id(); + let mut builder = test_codex().with_config(move |config| { + super::configure_scenario_catalog(config); + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable MAv2"); + config.multi_agent_v2.non_code_mode_only = false; + config.multi_agent_v2.root_agent_usage_hint_text = Some("Coordinate sub-agents.".into()); + config + .mcp_servers + .set(HashMap::from([( + "reports.one".to_string(), + serde_json::from_value(json!({ + "command": command, + "environment_id": environment_id, + "cwd": config.cwd, + "env": {"MCP_TEST_SERVER_INSTRUCTIONS": "Reports."}, + "enabled_tools": ["echo"], + "default_tools_approval_mode": "approve", + })) + .expect("MCP fixture config"), + )])) + .expect("MCP fixture allowed"); + }); + for (slug, prefixes, mcp_prefix) in [ + ( + "prefixes-conflicting", + json!({"mcp__reports_one": "Different report context."}), + "Reports context.", + ), + ( + "prefixes-cleared", + json!({"functions": "", "collaboration": ""}), + "", + ), + ( + "gpt-5.5", + json!({ + "functions": " Functions context. ", + "collaboration": "Agents context.", + }), + "Reports context.", + ), + ] { + builder = builder.with_model_info_override(slug, move |model| { + model.tool_mode = Some(tool_mode); + model.supports_search_tool = include_search; + model.model_messages.as_mut().expect("model messages").tools = Some( + serde_json::from_value(json!({ + "indirect_description_prefixes": { + "namespaces": prefixes, + "mcp_servers": {"reports.one": mcp_prefix}, + }, + })) + .expect("sparse catalog tool messages"), + ); + }); + } + let test = builder.build_with_auto_env(&server).await?; + wait_for_mcp_server(&test.codex, "reports.one").await?; + let script = r#" +const descriptions = ["view_image", "collaboration__list_agents", "mcp__reports_one__echo"] + .map(name => ALL_TOOLS.find(t => t.name === name).description); +const report = await tools.mcp__reports_one__echo({message: "ping"}); +const agents = await tools.collaboration__list_agents({}); +text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0].agent_name}); +"#; + let mock = responses::mount_sse_sequence( + &server, + [ + ("find", "run", "done"), + ("find-cleared", "run-cleared", "done-cleared"), + ] + .into_iter() + .flat_map(|(find, run, done)| { + let mut responses = Vec::new(); + if include_search { + responses.push(responses::sse(vec![ + responses::ev_tool_search_call(find, &json!({"query": "echo"})), + responses::ev_completed(find), + ])); + } + responses.extend([ + responses::sse(vec![ + responses::ev_custom_tool_call(run, "exec", script), + responses::ev_completed(run), + ]), + responses::sse(vec![ + responses::ev_assistant_message(done, "Done."), + responses::ev_completed(done), + ]), + ]); + responses + }) + .collect(), + ) + .await; + test.submit_turn("Find echo, call it, and list the agents.") + .await?; + test.codex + .submit(Op::ThreadSettings { + thread_settings: ThreadSettingsOverrides { + model: Some("prefixes-cleared".to_string()), + ..Default::default() + }, + }) + .await?; + test.submit_text_turn("Find echo again, call it, and list the agents.") + .await?; + let requests = mock.requests(); + assert_eq!(requests.len(), requests_per_turn * 2); + let prefixes = ["Functions context.", "Agents context.", "Reports context."]; + let direct_specs = [0, requests_per_turn].map(|index| { + let mut tools = requests[index].body_json()["tools"] + .as_array() + .expect("tools") + .clone(); + assert_eq!( + tools.iter().any(|tool| tool["type"] == "tool_search"), + include_search, + ); + let exec = tools.remove( + tools + .iter() + .position(|tool| tool["name"] == "exec") + .expect("exec"), + ); + let description = exec["description"].as_str().expect("description"); + for &prefix in &prefixes { + assert_eq!( + description.matches(prefix).count(), + usize::from(!include_search && index == 0), + "{description}" + ); + } + assert!( + prefixes + .iter() + .all(|prefix| !json!(tools).to_string().contains(*prefix)) + ); + tools + }); + assert_eq!(direct_specs[0], direct_specs[1]); + if include_search { + let loaded = requests[1].tool_search_output("find"); + let mut cleared = requests[requests_per_turn + 1].tool_search_output("find-cleared"); + assert_eq!(cleared["tools"][0]["description"], "Reports."); + cleared["tools"][0]["description"] = json!("Reports context.\n\nReports."); + assert_eq!(loaded["tools"], cleared["tools"]); + } + let mut outputs = Vec::new(); + for (index, call_id) in [ + (requests_per_turn - 1, "run"), + (requests_per_turn * 2 - 1, "run-cleared"), + ] { + let output: Value = serde_json::from_str( + &custom_tool_output_last_non_empty_text(&requests[index], call_id) + .expect("exec JSON output"), + )?; + assert_eq!( + (&output["echo"], &output["agent"]), + (&json!("ECHOING: ping"), &json!("/root")), + ); + outputs.push(output); + } + let mut expected = outputs[1].clone(); + for (description, prefix) in expected["descriptions"] + .as_array_mut() + .expect("tool descriptions") + .iter_mut() + .zip(prefixes) + { + *description = json!(format!( + "{prefix}\n\n{}", + description.as_str().expect("tool description") + )); + } + assert_eq!(outputs[0], expected); + if tool_mode == ToolMode::CodeModeOnly { + insta::assert_snapshot!( + "catalog_indirect_namespace_prefixes", + context_snapshot::format_request_history_snapshot( + "Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch.", + &requests, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .include_request_settings(), + ) + ); + } + if include_search { + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Find echo with conflicting prefixes.".to_string(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + model: Some("prefixes-conflicting".to_string()), + ..Default::default() + }), + ) + .await?; + let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await + else { + unreachable!("event predicate guarantees turn completion"); + }; + assert_eq!( + completed + .error + .expect("conflicting prefixes must fail the turn") + .message, + "Conflicting indirect description prefixes for namespace `mcp__reports_one`: `namespaces.mcp__reports_one` and `mcp_servers.reports.one`", + ); + assert_eq!( + server + .received_requests() + .await + .expect("received requests") + .iter() + .filter(|request| request.url.path() == "/v1/responses") + .count(), + requests.len(), + "conflicting prefixes must fail before another model request", + ); + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap new file mode 100644 index 0000000000..e7bb189843 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap @@ -0,0 +1,92 @@ +--- +source: core/tests/suite/scenarios_indirect_namespace_prefixes.rs +expression: "context_snapshot::format_request_history_snapshot(\"Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + instructions: + model: "gpt-5.5" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"effort":"medium"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" + tools (4; hash=da1f22cdbb122728): + - custom/exec: Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided JavaScript code in...; hash=bd4bc63f8dac64cb + - function/wait: Waits on a yielded `exec` cell and returns new output or completion. - Use `wait` only after `exe...; args=[cell_id, max_tokens, terminate, yield_time_ms]; hash=6692d8f6c55be489 + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - web_search; hash=e214d4db1bc49e89 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer: + Coordinate sub-agents. +02:message/developer: + +03:message/user: + +04:message/user: + Find echo, call it, and list the agents. +-- request 2 (turn) -- +05:custom_tool_call/exec: + const descriptions = ["view_image", "collaboration__list_agents", "mcp__reports_one__echo"] + .map(name => ALL_TOOLS.find(t => t.name === name).description); + const report = await tools.mcp__reports_one__echo({message: "ping"}); + const agents = await tools.collaboration__list_agents({}); + text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0].agent_name}); + +06:custom_tool_call_output:Script completed + Wall time seconds + Output: + | {"descriptions":["Functions context.\\n\\nView a local image file from the filesystem when visual inspecti...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=22aad8f340003b22] + +## Window 2 (after request 2) +Settings: relative to window 1 + model: "prefixes-cleared" + tools (4; hash=e51735744389fb0e): + ~ custom/exec: Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided JavaScript code in...; hash=1ef6e3872d1dfe3f +-- request 3 (turn) -- +00:message/developer: + +01:message/developer: + Coordinate sub-agents. +02:message/developer: + +03:message/user: + +04:message/user: + Find echo, call it, and list the agents. +05:custom_tool_call/exec: + const descriptions = ["view_image", "collaboration__list_agents", "mcp__reports_one__echo"] + .map(name => ALL_TOOLS.find(t => t.name === name).description); + const report = await tools.mcp__reports_one__echo({message: "ping"}); + const agents = await tools.collaboration__list_agents({}); + text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0].agent_name}); + +06:custom_tool_call_output:Script completed + Wall time seconds + Output: + | {"descriptions":["Functions context.\\n\\nView a local image file from the filesystem when visual inspecti...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=22aad8f340003b22] +07:message/assistant: + Done. +08:message/developer: + +09:message/user: + Find echo again, call it, and list the agents. +-- request 4 (turn) -- +10:custom_tool_call/exec: + const descriptions = ["view_image", "collaboration__list_agents", "mcp__reports_one__echo"] + .map(name => ALL_TOOLS.find(t => t.name === name).description); + const report = await tools.mcp__reports_one__echo({message: "ping"}); + const agents = await tools.collaboration__list_agents({}); + text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0].agent_name}); + +11:custom_tool_call_output:Script completed + Wall time seconds + Output: + | {"descriptions":["View a local image file from the filesystem when visual inspection is needed. Use this...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=6ecfde0f75e5f051] diff --git a/codex-rs/core/tests/suite/step_settings.rs b/codex-rs/core/tests/suite/step_settings.rs index 28f0d4a9a5..04e2bea4c2 100644 --- a/codex-rs/core/tests/suite/step_settings.rs +++ b/codex-rs/core/tests/suite/step_settings.rs @@ -2328,6 +2328,7 @@ async fn tool_messages_follow_mid_turn_model_changes() -> Result<()> { }), ..Default::default() }), + ..Default::default() }); } }) diff --git a/codex-rs/prompts/src/model_messages.rs b/codex-rs/prompts/src/model_messages.rs index ac6ec4042c..d7b7ef8298 100644 --- a/codex-rs/prompts/src/model_messages.rs +++ b/codex-rs/prompts/src/model_messages.rs @@ -5,6 +5,7 @@ use codex_protocol::openai_models::CodeModeToolMessages; use codex_protocol::openai_models::ConfirmationPolicies; +use codex_protocol::openai_models::IndirectDescriptionPrefixes; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelMessages; use codex_protocol::openai_models::ToolMessage; @@ -181,6 +182,15 @@ impl<'a> ResolvedModelMessages<'a> { tool.as_ref() } + /// Selects indirect tool guidance; tool rendering owns namespace mapping and normalization. + pub fn indirect_description_prefixes(&self) -> Option<&'a IndirectDescriptionPrefixes> { + self.catalog_messages? + .tools + .as_ref()? + .indirect_description_prefixes + .as_ref() + } + /// Selects Code Mode messages; bundled text and runtime composition belong to the tool owner. pub fn code_mode(&self) -> Option<&'a CodeModeToolMessages> { self.catalog_messages?.tools.as_ref()?.code_mode.as_ref() diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index b917fb2fc4..7f4578180f 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -575,9 +575,12 @@ pub struct ConfirmationPolicies { pub computer_use: Option, } -/// Model-owned messages for built-in tools. +/// Model-owned tool messages and indirect namespace guidance. #[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] pub struct ToolMessages { + /// Optional guidance for indirectly presented tools; missing or empty adds no prefix. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub indirect_description_prefixes: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub send_user_message_async: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -586,6 +589,21 @@ pub struct ToolMessages { pub code_mode: Option, } +/// Plain-text prefixes for Code Mode documentation, ALL_TOOLS, and loaded tool-search namespaces. +/// Values are trimmed, and selectors for the same final namespace must agree, including empty values. +/// Empty values add no prefix; unregistered targets are ignored. +#[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] +pub struct IndirectDescriptionPrefixes { + /// Exact rendered namespaces, including `functions` for plain tools. + /// Overlap with an MCP server prefix is allowed only when the trimmed values match. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub namespaces: Option>, + /// Configured MCP server names, before callable namespace normalization. + /// Each prefix applies to all namespaces exposed by that server. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mcp_servers: Option>, +} + /// Model-owned messages for a built-in tool. #[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] pub struct ToolMessage { diff --git a/codex-rs/tools/src/indirect_namespace_prefixes.rs b/codex-rs/tools/src/indirect_namespace_prefixes.rs new file mode 100644 index 0000000000..4261520627 --- /dev/null +++ b/codex-rs/tools/src/indirect_namespace_prefixes.rs @@ -0,0 +1,169 @@ +//! Adds namespace guidance to indirect tool descriptions after selection. +//! Conflicting prefixes for the same final namespace are rejected before rendering. +//! Direct specifications, search indexes, and cached definitions remain unchanged. + +use crate::LoadableToolSpec; +use crate::ToolName; +use codex_code_mode::ToolDefinition; +use codex_code_mode::ToolNamespaceDescription; +use codex_protocol::DEFAULT_FUNCTION_NAMESPACE; +use codex_protocol::openai_models::IndirectDescriptionPrefixes; +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::collections::btree_map::Entry; + +/// Conflicting catalog selectors for the same namespace, without including their prompt text. +#[derive(Debug, thiserror::Error)] +#[error( + "Conflicting indirect description prefixes for namespace `{namespace}`: `{first_source}` and `{second_source}`" +)] +pub struct IndirectNamespacePrefixConflict { + namespace: String, + first_source: String, + second_source: String, +} + +/// Borrowed model text applied only to fresh indirect presentation data. +pub struct IndirectNamespacePrefixes<'a> { + prefixes: BTreeMap, +} + +impl<'a> IndirectNamespacePrefixes<'a> { + /// Expands configured server names using registered tools' final namespaces. + /// Overlapping selectors must agree after trimming, including explicit empty values. + pub fn new<'s>( + prefixes: Option<&'a IndirectDescriptionPrefixes>, + namespaces: impl IntoIterator, + ) -> Result { + let mut resolved = BTreeMap::new(); + if let Some(namespaces) = prefixes.and_then(|prefixes| prefixes.namespaces.as_ref()) { + for (namespace, prefix) in namespaces { + resolved.insert( + namespace.clone(), + (prefix.trim(), format!("namespaces.{namespace}")), + ); + } + } + if let Some(servers) = prefixes + .and_then(|prefixes| prefixes.mcp_servers.as_ref()) + .filter(|servers| !servers.is_empty()) + { + for (server, namespace) in namespaces { + if let Some(prefix) = servers.get(server) { + let prefix = prefix.trim(); + match resolved.entry(namespace.to_owned()) { + Entry::Vacant(entry) => { + entry.insert((prefix, format!("mcp_servers.{server}"))); + } + Entry::Occupied(entry) if entry.get().0 != prefix => { + return Err(IndirectNamespacePrefixConflict { + namespace: namespace.to_owned(), + first_source: entry.get().1.clone(), + second_source: format!("mcp_servers.{server}"), + }); + } + Entry::Occupied(_) => {} + } + } + } + } + Ok(Self { + prefixes: resolved + .into_iter() + .filter(|(_, (prefix, _))| !prefix.is_empty()) + .map(|(namespace, (prefix, _))| (namespace, prefix)) + .collect(), + }) + } + + /// ALL_TOOLS is flat, so each entry carries its namespace's guidance. + /// Its eventual model-visible output is subject to normal tool-output truncation. + pub fn apply_code_mode(&self, tools: &mut [ToolDefinition]) { + if self.prefixes.is_empty() { + return; + } + for tool in tools { + if let Some(prefix) = self.prefix(namespace_name(&tool.tool_name)) { + prepend(prefix, &mut tool.description); + } + } + } + + /// Adds one prefix per namespace actually embedded in the exec description. + pub fn apply_exec_prompt( + &self, + tools: &mut [ToolDefinition], + namespaces: &mut BTreeMap, + ) { + if self.prefixes.is_empty() { + return; + } + let mut has_default_prefix = false; + let rendered_namespaces = tools + .iter() + .map(|tool| namespace_name(&tool.tool_name)) + .collect::>(); + for name in rendered_namespaces { + if let Some(prefix) = self.prefix(name) { + let namespace = namespaces.entry(name.to_string()).or_insert_with(|| { + ToolNamespaceDescription { + name: name.to_string(), + description: String::new(), + } + }); + prepend(prefix, &mut namespace.description); + has_default_prefix |= name == DEFAULT_FUNCTION_NAMESPACE; + } + } + if has_default_prefix { + // Only temporary prompt definitions join this group; runtime names stay unchanged. + for tool in tools { + if tool + .tool_name + .namespace + .as_deref() + .is_none_or(str::is_empty) + { + tool.tool_name.namespace = Some(DEFAULT_FUNCTION_NAMESPACE.to_string()); + } + } + } + } + + /// Applies prefixes after search ranking and namespace coalescing. + pub fn apply_search(&self, tools: &mut [LoadableToolSpec]) { + if self.prefixes.is_empty() { + return; + } + for spec in tools { + if let LoadableToolSpec::Namespace(namespace) = spec + && let Some(prefix) = self.prefix(&namespace.name) + { + prepend(prefix, &mut namespace.description); + } + } + } + + fn prefix(&self, namespace: &str) -> Option<&'a str> { + self.prefixes.get(namespace).copied() + } +} + +fn prepend(prefix: &str, description: &mut String) { + *description = if description.is_empty() { + prefix.to_string() + } else { + format!("{prefix}\n\n{description}") + }; +} + +fn namespace_name(name: &ToolName) -> &str { + name.namespace + .as_deref() + .filter(|namespace| !namespace.is_empty()) + .unwrap_or(DEFAULT_FUNCTION_NAMESPACE) +} + +#[cfg(test)] +#[path = "indirect_namespace_prefixes_tests.rs"] +mod tests; diff --git a/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs b/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs new file mode 100644 index 0000000000..e9281cb3e6 --- /dev/null +++ b/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs @@ -0,0 +1,205 @@ +//! Checks prefix normalization and rendering without exercising tool execution. + +use super::*; +use codex_code_mode::CodeModeToolKind; +use codex_protocol::openai_models::IndirectDescriptionPrefixes; +use codex_protocol::openai_models::ToolMessages; +use pretty_assertions::assert_eq; +use serde_json::json; + +#[test] +fn plain_tools_join_an_existing_functions_group_only_with_a_prefix() { + let original_tools = [ + ToolName::plain("plain"), + ToolName::namespaced("functions", "namespaced"), + ] + .map(|tool_name| ToolDefinition { + name: tool_name.name.clone(), + tool_name, + description: "Tool guidance.".to_string(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }); + let original_namespaces = BTreeMap::from([( + "functions".to_string(), + ToolNamespaceDescription { + name: "functions".to_string(), + description: "Original namespace guidance.".to_string(), + }, + )]); + let render = |tools: &[ToolDefinition], + namespaces: &BTreeMap| { + codex_code_mode::build_exec_tool_description( + tools, + &[], + namespaces, + codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS, + /*code_mode_only*/ true, + codex_code_mode::ImageDetailVisibility::Visible, + /*messages*/ None, + ) + }; + let original = render(&original_tools, &original_namespaces); + for messages in [ + json!({}), + json!({"indirect_description_prefixes": null}), + json!({"indirect_description_prefixes": {}}), + json!({"indirect_description_prefixes": {"namespaces": null, "mcp_servers": null}}), + json!({"indirect_description_prefixes": {"namespaces": {}, "mcp_servers": {}}}), + json!({"indirect_description_prefixes": {"namespaces": {"functions": " \n"}}}), + json!({"indirect_description_prefixes": { + "namespaces": {"unavailable": "Unused guidance."}, + }}), + ] { + let configured: ToolMessages = serde_json::from_value(messages).unwrap(); + let mut tools = original_tools.clone(); + let mut namespaces = original_namespaces.clone(); + IndirectNamespacePrefixes::new(configured.indirect_description_prefixes.as_ref(), []) + .unwrap() + .apply_exec_prompt(&mut tools, &mut namespaces); + assert_eq!(render(&tools, &namespaces), original); + } + let configured = IndirectDescriptionPrefixes { + namespaces: Some(BTreeMap::from([( + "functions".to_string(), + "Catalog guidance.".to_string(), + )])), + ..Default::default() + }; + let mut tools = original_tools; + let mut namespaces = original_namespaces; + IndirectNamespacePrefixes::new(Some(&configured), []) + .unwrap() + .apply_exec_prompt(&mut tools, &mut namespaces); + let rendered = render(&tools, &namespaces); + assert_eq!(rendered.matches("## functions").count(), 1); + assert!(rendered.contains("## functions\nCatalog guidance.\n\nOriginal namespace guidance.")); + assert!(rendered.find("## functions").unwrap() < rendered.find("### `plain`").unwrap()); +} + +#[test] +fn flat_tool_descriptions_use_trimmed_namespace_prefixes() { + let configured = IndirectDescriptionPrefixes { + namespaces: Some(BTreeMap::from([( + "functions".to_string(), + " \nFunction guidance.\t ".to_string(), + )])), + ..Default::default() + }; + let original = ToolDefinition { + name: "inspect".to_string(), + tool_name: ToolName::plain("inspect"), + description: "Inspect.".to_string(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }; + let mut tools = [original.clone()]; + IndirectNamespacePrefixes::new(Some(&configured), []) + .unwrap() + .apply_code_mode(&mut tools); + assert_eq!( + tools, + [ToolDefinition { + description: "Function guidance.\n\nInspect.".to_string(), + ..original + }] + ); +} + +#[test] +fn matching_prefixes_cover_each_namespace_once() { + let servers = BTreeMap::from([ + ("codex_apps".to_string(), " App guidance. ".to_string()), + ("apps_alias".to_string(), "App guidance.\n".to_string()), + ]); + let exact = BTreeMap::from([( + "mcp__codex_apps__gmail".to_string(), + "App guidance.".to_string(), + )]); + let configured = IndirectDescriptionPrefixes { + namespaces: Some(exact), + mcp_servers: Some(servers), + }; + let original = ["gmail", "calendar", "drive", "slack"].map(|app| ToolDefinition { + name: format!("mcp__codex_apps__{app}__search"), + tool_name: ToolName::namespaced(format!("mcp__codex_apps__{app}"), "search"), + description: "Search.".to_string(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }); + let mut tools = original.clone(); + let namespaces = original + .iter() + .map(|tool| ("codex_apps", tool.tool_name.namespace.as_deref().unwrap())) + .chain([ + ("codex_apps", "mcp__codex_apps__gmail"), + ("apps_alias", "mcp__codex_apps__gmail"), + ]); + IndirectNamespacePrefixes::new(Some(&configured), namespaces) + .unwrap() + .apply_code_mode(&mut tools); + let expected = original + .into_iter() + .map(|tool| ToolDefinition { + description: "App guidance.\n\nSearch.".to_string(), + ..tool + }) + .collect::>(); + assert_eq!(tools.as_slice(), expected); +} + +#[test] +fn conflicting_namespace_and_server_prefixes_report_keys_without_values() { + for (namespace_prefix, server_prefix) in [ + ("Namespace guidance.", "Server guidance."), + ("", "Server guidance."), + ("Namespace guidance.", " \n"), + ] { + let configured = IndirectDescriptionPrefixes { + namespaces: Some(BTreeMap::from([( + "mcp__reports_one".to_string(), + namespace_prefix.to_string(), + )])), + mcp_servers: Some(BTreeMap::from([( + "reports.one".to_string(), + server_prefix.to_string(), + )])), + }; + let error = IndirectNamespacePrefixes::new( + Some(&configured), + [("reports.one", "mcp__reports_one")], + ) + .err() + .expect("conflicting prefixes") + .to_string(); + assert_eq!( + error, + "Conflicting indirect description prefixes for namespace `mcp__reports_one`: `namespaces.mcp__reports_one` and `mcp_servers.reports.one`", + ); + } +} + +#[test] +fn conflicting_mcp_servers_report_the_shared_namespace() { + let configured = IndirectDescriptionPrefixes { + mcp_servers: Some(BTreeMap::from([ + ("reports.one".to_string(), "First guidance.".to_string()), + ("reports.two".to_string(), "Second guidance.".to_string()), + ])), + ..Default::default() + }; + let error = IndirectNamespacePrefixes::new( + Some(&configured), + [("reports.one", "shared"), ("reports.two", "shared")], + ) + .err() + .expect("conflicting prefixes") + .to_string(); + assert_eq!( + error, + "Conflicting indirect description prefixes for namespace `shared`: `mcp_servers.reports.one` and `mcp_servers.reports.two`", + ); +} diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 820ee5413c..b07c7b0cd2 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -5,6 +5,7 @@ mod code_mode; mod dynamic_tool; mod function_call_error; mod image_detail; +mod indirect_namespace_prefixes; mod json_schema; mod mcp_tool; mod output_schema; @@ -32,6 +33,8 @@ pub use function_call_error::FunctionCallError; pub use image_detail::can_request_original_image_detail; pub use image_detail::normalize_output_image_detail; pub use image_detail::sanitize_original_image_detail; +pub use indirect_namespace_prefixes::IndirectNamespacePrefixConflict; +pub use indirect_namespace_prefixes::IndirectNamespacePrefixes; pub use json_schema::AdditionalProperties; pub use json_schema::JsonSchema; pub use json_schema::JsonSchemaPrimitiveType;