From f3a679e05510e7e653f72ff23b9bd2632eb46b2a Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Tue, 7 Jul 2026 20:43:08 +0000 Subject: [PATCH] Capture tool search pipeline diagnostics in feedback --- .../request_processors/feedback_processor.rs | 14 + codex-rs/codex-mcp/src/connection_manager.rs | 77 ++- .../codex-mcp/src/connection_manager_tests.rs | 4 + codex-rs/codex-mcp/src/lib.rs | 5 + codex-rs/codex-mcp/src/rmcp_client.rs | 35 +- .../codex-mcp/src/tool_search_diagnostics.rs | 164 +++++ codex-rs/core/src/codex_thread.rs | 14 + codex-rs/core/src/mcp_tool_exposure.rs | 57 ++ codex-rs/core/src/session/step_context.rs | 19 +- codex-rs/core/src/session/turn.rs | 14 +- codex-rs/core/src/session/world_state.rs | 2 +- .../core/src/tools/handlers/tool_search.rs | 51 +- codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/spec_plan.rs | 16 + .../core/src/tools/tool_search_diagnostics.rs | 561 ++++++++++++++++++ .../tools/tool_search_diagnostics_tests.rs | 152 +++++ 16 files changed, 1161 insertions(+), 25 deletions(-) create mode 100644 codex-rs/codex-mcp/src/tool_search_diagnostics.rs create mode 100644 codex-rs/core/src/tools/tool_search_diagnostics.rs create mode 100644 codex-rs/core/src/tools/tool_search_diagnostics_tests.rs diff --git a/codex-rs/app-server/src/request_processors/feedback_processor.rs b/codex-rs/app-server/src/request_processors/feedback_processor.rs index 5f488f72c3..d44b4530f0 100644 --- a/codex-rs/app-server/src/request_processors/feedback_processor.rs +++ b/codex-rs/app-server/src/request_processors/feedback_processor.rs @@ -1,8 +1,10 @@ use super::*; +use codex_feedback::FeedbackAttachment; #[cfg(target_os = "windows")] use codex_feedback::WINDOWS_SANDBOX_LOG_ATTACHMENT_FILENAME; const MAX_FEEDBACK_TREE_THREADS: usize = 8; +const TOOL_SEARCH_PIPELINE_ATTACHMENT_FILENAME: &str = "tool-search-pipeline.json"; #[derive(Clone)] pub(crate) struct FeedbackRequestProcessor { @@ -220,6 +222,18 @@ impl FeedbackRequestProcessor { } } + if include_logs + && let Some(conversation_id) = conversation_id + && let Ok(conversation) = self.thread_manager.get_thread(conversation_id).await + && let Some(buffer) = conversation.tool_search_pipeline_feedback_json() + { + extra_attachments.push(FeedbackAttachment { + filename: TOOL_SEARCH_PIPELINE_ATTACHMENT_FILENAME.to_string(), + content_type: Some("application/json".to_string()), + buffer, + }); + } + let session_source = self.thread_manager.session_source(); let upload_result = tokio::task::spawn_blocking(move || { diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index c21a9f1d10..66137f69ca 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -6,9 +6,11 @@ //! calls to the right client, and exposes the public manager API used by //! `codex-core`. +use std::collections::BTreeMap; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; @@ -25,6 +27,7 @@ use crate::mcp::ToolPluginProvenance; use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC; use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; +use crate::rmcp_client::ListedToolsFetch; use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; @@ -33,6 +36,8 @@ use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; +use crate::tool_search_diagnostics::ToolSearchMcpDiagnosticsSnapshot; +use crate::tool_search_diagnostics::ToolSearchToolsListResponseSnapshot; use crate::tools::ToolInfo; use crate::tools::filter_tools; use crate::tools::normalize_tools_for_model_with_prefix; @@ -119,6 +124,7 @@ pub struct McpConnectionManager { prefix_mcp_tool_names: bool, elicitation_requests: ElicitationRequestManager, startup_cancellation_token: CancellationToken, + latest_tools_list_responses: Mutex>, } impl McpConnectionManager { @@ -293,6 +299,7 @@ impl McpConnectionManager { prefix_mcp_tool_names, elicitation_requests: elicitation_requests.clone(), startup_cancellation_token: startup_cancellation_token.clone(), + latest_tools_list_responses: Mutex::new(BTreeMap::new()), }; tokio::spawn(async move { let outcomes = join_set.join_all().await; @@ -384,6 +391,7 @@ impl McpConnectionManager { ElicitationRequestRouter::default(), ), startup_cancellation_token: CancellationToken::new(), + latest_tools_list_responses: Mutex::new(BTreeMap::new()), } } @@ -503,16 +511,29 @@ impl McpConnectionManager { /// Returns all tools with model-visible names normalized. #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] pub async fn list_all_tools(&self) -> Vec { + self.list_all_tools_with_diagnostics().await.0 + } + + /// Returns normalized tools plus compact MCP inputs for thread-scoped feedback diagnostics. + #[doc(hidden)] + pub async fn list_all_tools_with_diagnostics( + &self, + ) -> (Vec, ToolSearchMcpDiagnosticsSnapshot) { let mut tools = Vec::new(); let mut available_server_count = 0; let mut unavailable_server_count = 0; + let mut cached_server_count = 0; + let mut startup_complete_server_count = 0; + let mut fallback_responses = BTreeMap::new(); for (server_name, managed_client) in &self.clients { managed_client.reconnect_failed_startup().await; let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete .load(std::sync::atomic::Ordering::Acquire); - let Some(server_tools) = managed_client + cached_server_count += usize::from(has_cached_tools); + startup_complete_server_count += usize::from(startup_complete); + let Some((server_tools, response)) = managed_client .listed_tools() .instrument(trace_span!( "list_tools_for_server", @@ -529,8 +550,23 @@ impl McpConnectionManager { startup_complete, "MCP server tools unavailable while building tool list" ); + fallback_responses.insert( + server_name.clone(), + ToolSearchToolsListResponseSnapshot::not_observed(server_name), + ); continue; }; + if let Some(response) = response { + self.observe_tools_list_response(response); + } else { + fallback_responses.insert( + server_name.clone(), + ToolSearchToolsListResponseSnapshot::from_cached_tools( + server_name, + &server_tools, + ), + ); + } available_server_count += 1; tools.extend( server_tools @@ -545,7 +581,41 @@ impl McpConnectionManager { tool_count = tools.len(), "built MCP tool list" ); - tools + let diagnostics = self.mcp_diagnostics_snapshot( + fallback_responses, + cached_server_count, + startup_complete_server_count, + ); + (tools, diagnostics) + } + + fn observe_tools_list_response(&self, response: ToolSearchToolsListResponseSnapshot) { + let mut responses = match self.latest_tools_list_responses.lock() { + Ok(responses) => responses, + Err(poisoned) => poisoned.into_inner(), + }; + responses.insert(response.server_name.clone(), response); + } + + fn mcp_diagnostics_snapshot( + &self, + mut fallback_responses: BTreeMap, + cached_server_count: usize, + startup_complete_server_count: usize, + ) -> ToolSearchMcpDiagnosticsSnapshot { + let responses = match self.latest_tools_list_responses.lock() { + Ok(responses) => responses, + Err(poisoned) => poisoned.into_inner(), + }; + for (server_name, response) in responses.iter() { + fallback_responses.insert(server_name.clone(), response.clone()); + } + ToolSearchMcpDiagnosticsSnapshot { + tools_list_responses: fallback_responses.into_values().collect(), + server_count: self.clients.len(), + cached_server_count, + startup_complete_server_count, + } } /// Force-refresh codex apps tools by bypassing the in-process cache. @@ -568,7 +638,7 @@ impl McpConnectionManager { .codex_apps_tools_cache_context .as_ref() .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh)); - let tools = list_tools_for_client_uncached( + let ListedToolsFetch { tools, diagnostics } = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, /*codex_apps_refresh_trigger*/ "explicit", @@ -580,6 +650,7 @@ impl McpConnectionManager { .with_context(|| { format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'") })?; + self.observe_tools_list_response(diagnostics); let tools = match ( diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d8a1f1863e..2e43898205 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -119,6 +119,10 @@ async fn create_test_managed_client(tools: Vec) -> ManagedClient { tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, + tools_list_response: + crate::tool_search_diagnostics::ToolSearchToolsListResponseSnapshot::not_observed( + "test", + ), server_supports_sandbox_state_meta_capability: false, codex_apps_tools_cache_context: None, } diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 77128d7970..2c8be3961e 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -12,6 +12,10 @@ pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntimeContext; pub use runtime::SandboxState; +#[doc(hidden)] +pub use tool_search_diagnostics::ToolSearchDiagnosticIdentity; +#[doc(hidden)] +pub use tool_search_diagnostics::ToolSearchMcpDiagnosticsSnapshot; pub use tools::ToolInfo; pub use catalog::McpCatalogBuilder; @@ -87,4 +91,5 @@ mod resource_client; pub(crate) mod rmcp_client; pub(crate) mod runtime; pub(crate) mod server; +mod tool_search_diagnostics; pub(crate) mod tools; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 6fe820b4fe..034070db2a 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -31,6 +31,7 @@ use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; use crate::server::McpServerLaunch; +use crate::tool_search_diagnostics::ToolSearchToolsListResponseSnapshot; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; @@ -106,6 +107,7 @@ pub(crate) struct ManagedClient { pub(crate) tool_timeout: Option, pub(crate) server_instructions: Option, pub(crate) server_supports_sandbox_state_meta_capability: bool, + pub(crate) tools_list_response: ToolSearchToolsListResponseSnapshot, pub(crate) codex_apps_tools_cache_context: Option, } @@ -516,23 +518,26 @@ impl AsyncManagedClient { .map(|tools| filter_tools(tools, &self.tool_filter)) } - pub(crate) async fn listed_tools(&self) -> Option> { + pub(crate) async fn listed_tools( + &self, + ) -> Option<(Vec, Option)> { // Keep cache payloads raw; plugin provenance is resolved per-session at read time. - let tools = if !self.startup_complete.load(Ordering::Acquire) + let (tools, diagnostics) = if !self.startup_complete.load(Ordering::Acquire) && let Some(startup_tools) = self.cached_tools() { - Some(startup_tools) + (startup_tools, None) } else { match self.client().await { - Ok(client) => Some(client.listed_tools()), - Err(_) => self.cached_tools(), + Ok(client) => (client.listed_tools(), Some(client.tools_list_response)), + Err(_) => (self.cached_tools()?, None), } - }?; - Some(if self.is_codex_apps_mcp_server { + }; + let tools = if self.is_codex_apps_mcp_server { prepare_codex_apps_tools_for_model(tools, &self.tool_plugin_provenance) } else { prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) - }) + }; + Some((tools, diagnostics)) } } @@ -579,11 +584,13 @@ pub(crate) async fn list_tools_for_client_uncached( client: &Arc, timeout: Option, server_instructions: Option<&str>, -) -> Result> { +) -> Result { let fetch_start = Instant::now(); let resp = client .list_tools_with_connector_ids(/*params*/ None, timeout) .await?; + let diagnostics = + ToolSearchToolsListResponseSnapshot::from_live_response(server_name, &resp.tools); let tools = resp .tools .into_iter() @@ -609,7 +616,12 @@ pub(crate) async fn list_tools_for_client_uncached( &[], ); } - Ok(tools) + Ok(ListedToolsFetch { tools, diagnostics }) +} + +pub(crate) struct ListedToolsFetch { + pub(crate) tools: Vec, + pub(crate) diagnostics: ToolSearchToolsListResponseSnapshot, } /// Presents declared Codex Apps file parameters to the model as local-path inputs and adds plugin @@ -849,7 +861,7 @@ async fn start_server_task( let fetch_ticket = codex_apps_tools_cache_context .as_ref() .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); - let tools = list_tools_for_client_uncached( + let ListedToolsFetch { tools, diagnostics } = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, /*codex_apps_refresh_trigger*/ "initial", @@ -884,6 +896,7 @@ async fn start_server_task( tool_filter, server_instructions: initialize_result.instructions, server_supports_sandbox_state_meta_capability, + tools_list_response: diagnostics, codex_apps_tools_cache_context, }; diff --git a/codex-rs/codex-mcp/src/tool_search_diagnostics.rs b/codex-rs/codex-mcp/src/tool_search_diagnostics.rs new file mode 100644 index 0000000000..74d0a51fb8 --- /dev/null +++ b/codex-rs/codex-mcp/src/tool_search_diagnostics.rs @@ -0,0 +1,164 @@ +//! Compact MCP-side inputs for per-session tool-search feedback diagnostics. +//! +//! These records deliberately carry identities only. Schemas, descriptions, +//! metadata, auth state, and searchable text never enter this path. + +use crate::tools::ToolInfo; +use codex_rmcp_client::ToolWithConnectorId; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; + +const MAX_MCP_DIAGNOSTIC_IDENTITIES: usize = 256; +const MAX_IDENTITY_CHARS: usize = 160; + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)] +pub struct ToolSearchDiagnosticIdentity { + pub server_name: String, + pub raw_tool_name: String, + pub callable_namespace: Option, + pub callable_name: Option, + pub connector_id: Option, + pub connector_name: Option, + pub plugin_display_names: Vec, +} + +impl ToolSearchDiagnosticIdentity { + pub(crate) fn from_raw_tool(server_name: &str, tool: &ToolWithConnectorId) -> Self { + Self { + server_name: bounded(server_name.to_string()), + raw_tool_name: bounded(tool.tool.name.to_string()), + callable_namespace: None, + callable_name: None, + connector_id: tool.connector_id.clone().map(bounded), + connector_name: tool.connector_name.clone().map(bounded), + plugin_display_names: Vec::new(), + } + } + + pub fn from_tool_info(tool: &ToolInfo) -> Self { + Self { + server_name: bounded(tool.server_name.clone()), + raw_tool_name: bounded(tool.tool.name.to_string()), + callable_namespace: Some(bounded(tool.callable_namespace.clone())), + callable_name: Some(bounded(tool.callable_name.clone())), + connector_id: tool.connector_id.clone().map(bounded), + connector_name: tool.connector_name.clone().map(bounded), + plugin_display_names: tool + .plugin_display_names + .iter() + .take(4) + .cloned() + .map(bounded) + .collect(), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolSearchToolsListSource { + LiveResponse, + CachedToolsWithoutLiveResponse, + NotObserved, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ToolSearchToolsListResponseSnapshot { + pub server_name: String, + pub source: ToolSearchToolsListSource, + pub total_count: usize, + pub fingerprint: String, + pub identities: Vec, + pub identities_truncated: bool, +} + +impl ToolSearchToolsListResponseSnapshot { + pub(crate) fn from_live_response(server_name: &str, tools: &[ToolWithConnectorId]) -> Self { + Self::new( + server_name, + ToolSearchToolsListSource::LiveResponse, + tools + .iter() + .map(|tool| ToolSearchDiagnosticIdentity::from_raw_tool(server_name, tool)) + .collect(), + ) + } + + pub(crate) fn from_cached_tools(server_name: &str, tools: &[ToolInfo]) -> Self { + Self::new( + server_name, + ToolSearchToolsListSource::CachedToolsWithoutLiveResponse, + tools + .iter() + .map(ToolSearchDiagnosticIdentity::from_tool_info) + .collect(), + ) + } + + pub(crate) fn not_observed(server_name: &str) -> Self { + Self::new( + server_name, + ToolSearchToolsListSource::NotObserved, + Vec::new(), + ) + } + + fn new( + server_name: &str, + source: ToolSearchToolsListSource, + mut identities: Vec, + ) -> Self { + identities.sort(); + identities.dedup(); + let total_count = identities.len(); + let fingerprint = fingerprint(&identities); + let identities_truncated = identities.len() > MAX_MCP_DIAGNOSTIC_IDENTITIES; + identities.truncate(MAX_MCP_DIAGNOSTIC_IDENTITIES); + Self { + server_name: bounded(server_name.to_string()), + source, + total_count, + fingerprint, + identities, + identities_truncated, + } + } +} + +#[derive(Clone, Debug, Default, Serialize)] +pub struct ToolSearchMcpDiagnosticsSnapshot { + pub tools_list_responses: Vec, + pub server_count: usize, + pub cached_server_count: usize, + pub startup_complete_server_count: usize, +} + +fn bounded(value: String) -> String { + let mut chars = value.chars(); + let bounded = chars.by_ref().take(MAX_IDENTITY_CHARS).collect::(); + if chars.next().is_some() { + format!("{bounded}…") + } else { + bounded + } +} + +fn fingerprint(identities: &[ToolSearchDiagnosticIdentity]) -> String { + let mut hasher = Sha1::new(); + for identity in identities { + hasher.update(identity.server_name.as_bytes()); + hasher.update([0]); + hasher.update(identity.raw_tool_name.as_bytes()); + hasher.update([0]); + hasher.update( + identity + .connector_id + .as_deref() + .unwrap_or_default() + .as_bytes(), + ); + hasher.update([0]); + } + format!("{:x}", hasher.finalize()) +} diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 257a32733b..dfb405623d 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -503,6 +503,20 @@ impl CodexThread { self.session_configured.clone() } + /// Returns the bounded live tool-search snapshot consumed by normal feedback uploads. + #[doc(hidden)] + pub fn tool_search_pipeline_feedback_json(&self) -> Option> { + self.codex + .session + .services + .tool_search_handler_cache + .diagnostics() + .feedback_json( + self.session_configured.thread_id, + self.session_configured.session_id, + ) + } + pub(crate) fn is_running(&self) -> bool { !self.codex.tx_sub.is_closed() } diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index fda2b04c9b..05a69c3780 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -9,10 +9,13 @@ use tracing::instrument; use crate::config::Config; use crate::connectors; +use crate::tools::tool_search_diagnostics::DiagnosticToolIdentity; +use crate::tools::tool_search_diagnostics::ToolExposureDecision; pub(crate) struct McpToolExposure { pub(crate) direct_tools: Vec, pub(crate) deferred_tools: Option>, + pub(crate) diagnostic_decisions: Vec, } #[instrument(level = "trace", skip_all)] @@ -30,17 +33,20 @@ pub(crate) fn build_mcp_tool_exposure( config, )); } + let diagnostic_decisions = tool_exposure_decisions(all_mcp_tools, connectors, config); if !search_tool_enabled { return McpToolExposure { direct_tools: deferred_tools, deferred_tools: None, + diagnostic_decisions, }; } McpToolExposure { direct_tools: Vec::new(), deferred_tools: (!deferred_tools.is_empty()).then_some(deferred_tools), + diagnostic_decisions, } } @@ -95,6 +101,57 @@ fn filter_codex_apps_mcp_tools( .collect() } +fn tool_exposure_decisions( + all_mcp_tools: &[McpToolInfo], + connectors: Option<&[connectors::AppInfo]>, + config: &Config, +) -> Vec { + let allowed = connectors + .into_iter() + .flatten() + .map(|connector| connector.id.as_str()) + .collect::>(); + let app_tool_policy = AppToolPolicyEvaluator::new(&config.config_layer_stack); + + all_mcp_tools + .iter() + .map(|tool| { + let is_codex_apps_tool = tool.server_name == CODEX_APPS_MCP_SERVER_NAME; + let model_visible = tool_is_model_visible(tool); + let connector_id_present = tool.connector_id.is_some(); + let connector_allowed = !is_codex_apps_tool + || (connectors.is_some() + && tool + .connector_id + .as_deref() + .is_some_and(|connector_id| allowed.contains(connector_id))); + let policy_enabled = !is_codex_apps_tool + || tool.connector_id.as_deref().is_some_and(|connector_id| { + let annotations = tool.tool.annotations.as_ref(); + app_tool_policy + .policy(AppToolPolicyInput { + connector_id: Some(connector_id), + tool_name: &tool.tool.name, + tool_title: tool.tool.title.as_deref(), + destructive_hint: annotations + .and_then(|annotations| annotations.destructive_hint), + open_world_hint: annotations + .and_then(|annotations| annotations.open_world_hint), + }) + .enabled + }); + ToolExposureDecision { + identity: DiagnosticToolIdentity::from_tool_info(tool), + included: model_visible && connector_allowed && policy_enabled, + model_visible, + connector_id_present, + connector_allowed, + policy_enabled, + } + }) + .collect() +} + #[cfg(test)] #[path = "mcp_tool_exposure_test.rs"] mod tests; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 1d7316d914..c881449add 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -6,8 +6,15 @@ use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::ToolInfo; +use codex_mcp::ToolSearchMcpDiagnosticsSnapshot; use tokio::sync::OnceCell; +#[derive(Debug)] +pub(crate) struct McpToolSnapshot { + pub(crate) tools: Vec, + pub(crate) diagnostics: ToolSearchMcpDiagnosticsSnapshot, +} + /// Request-scoped state that may change between model sampling requests. #[derive(Debug)] pub(crate) struct StepContext { @@ -17,8 +24,8 @@ pub(crate) struct StepContext { pub(crate) selected_capability_roots: Vec, /// The exact MCP config and manager used to advertise and execute tools for this step. pub(crate) mcp: Arc, - /// The fixed MCP tool list used for this exact sampling request. - mcp_tool_snapshot: OnceCell>, + /// The fixed MCP tool list and compact diagnostics used for this exact sampling request. + mcp_tool_snapshot: OnceCell, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -41,9 +48,13 @@ impl StepContext { } } - pub(crate) async fn mcp_tools(&self) -> &[ToolInfo] { + pub(crate) async fn mcp_tools(&self) -> &McpToolSnapshot { self.mcp_tool_snapshot - .get_or_init(|| self.mcp.manager().list_all_tools()) + .get_or_init(|| async { + let (tools, diagnostics) = + self.mcp.manager().list_all_tools_with_diagnostics().await; + McpToolSnapshot { tools, diagnostics } + }) .await } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index bafb368b32..93eba114ca 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1222,10 +1222,15 @@ pub(crate) async fn built_tools( let turn_context = step_context.turn.as_ref(); let mcp_connection_manager = step_context.mcp.manager(); let has_mcp_servers = mcp_connection_manager.has_servers(); - let all_mcp_tools = step_context + let mcp_tool_snapshot = step_context .mcp_tools() .or_cancel(cancellation_token) .await?; + let all_mcp_tools = mcp_tool_snapshot.tools.as_slice(); + sess.services + .tool_search_handler_cache + .diagnostics() + .record_mcp_snapshot(&mcp_tool_snapshot.diagnostics, all_mcp_tools); let loaded_plugins = sess .services .plugins_manager @@ -1328,12 +1333,17 @@ pub(crate) async fn built_tools( .instrument(trace_span!("built_tools.load_discoverable_tools")) .await }; + let search_tool_enabled = search_tool_enabled(turn_context); let mcp_tool_exposure = build_mcp_tool_exposure( all_mcp_tools, connectors.as_deref(), &turn_context.config, - search_tool_enabled(turn_context), + search_tool_enabled, ); + sess.services + .tool_search_handler_cache + .diagnostics() + .record_exposure(mcp_tool_exposure.diagnostic_decisions, search_tool_enabled); 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_context( diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index c7743c3fe1..44ed51dd59 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -41,7 +41,7 @@ impl Session { } let apps_available = if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let tools = step_context.mcp_tools().await; + let tools = &step_context.mcp_tools().await.tools; connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(tools), &turn_context.config, diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 272da2f226..3ff1001932 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -6,6 +6,9 @@ use crate::tools::context::boxed_tool_output; use crate::tools::handlers::tool_search_spec::create_tool_search_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; +use crate::tools::tool_search_diagnostics::DiagnosticToolIdentity; +use crate::tools::tool_search_diagnostics::SearchResultDiagnostic; +use crate::tools::tool_search_diagnostics::ToolSearchPipelineDiagnostics; use bm25::Document; use bm25::Language; use bm25::SearchEngine; @@ -26,11 +29,13 @@ pub struct ToolSearchHandler { search_infos: Vec, spec: ToolSpec, search_engine: SearchEngine, + diagnostics: Arc, } #[derive(Default)] pub(crate) struct ToolSearchHandlerCache { cached: Mutex>>, + diagnostics: Arc, } impl ToolSearchHandlerCache { @@ -41,22 +46,33 @@ impl ToolSearchHandlerCache { if let Some(cached) = cached.as_ref() && cached.search_infos == search_infos { + self.diagnostics.record_index(&search_infos, "hit"); return Arc::clone(cached); } } - let handler = Arc::new(ToolSearchHandler::new(search_infos)); + let handler = Arc::new(ToolSearchHandler::new_with_diagnostics( + search_infos, + Arc::clone(&self.diagnostics), + )); let mut cached = self.cached(); if let Some(cached) = cached.as_ref() && cached.search_infos == handler.search_infos { + self.diagnostics + .record_index(&handler.search_infos, "hit_after_build"); return Arc::clone(cached); } *cached = Some(Arc::clone(&handler)); + self.diagnostics.record_index(&handler.search_infos, "miss"); handler } + pub(crate) fn diagnostics(&self) -> &ToolSearchPipelineDiagnostics { + self.diagnostics.as_ref() + } + fn cached(&self) -> std::sync::MutexGuard<'_, Option>> { match self.cached.lock() { Ok(cached) => cached, @@ -66,12 +82,20 @@ impl ToolSearchHandlerCache { } impl ToolSearchHandler { + #[cfg(test)] #[instrument( level = "trace", skip_all, fields(search_info_count = search_infos.len()) )] pub(crate) fn new(search_infos: Vec) -> Self { + Self::new_with_diagnostics(search_infos, Arc::new(Default::default())) + } + + fn new_with_diagnostics( + search_infos: Vec, + diagnostics: Arc, + ) -> Self { let search_source_infos = search_infos .iter() .filter_map(|search_info| search_info.source_info.clone()) @@ -90,6 +114,7 @@ impl ToolSearchHandler { search_infos, spec, search_engine, + diagnostics, } } } @@ -143,6 +168,8 @@ impl ToolSearchHandler { } if self.search_infos.is_empty() { + self.diagnostics + .record_search(query, limit, /*returned_count*/ 0, Vec::new()); return Ok(boxed_tool_output(ToolSearchOutput { tools: Vec::new() })); } @@ -160,9 +187,25 @@ impl ToolSearchHandler { query: &str, limit: usize, ) -> Result, FunctionCallError> { - let results = self - .search_engine - .search(query, limit) + let results = self.search_engine.search(query, limit); + let diagnostics = results + .iter() + .enumerate() + .filter_map(|(rank, result)| { + let document_index = result.document.id; + self.search_infos + .get(document_index) + .map(|search_info| SearchResultDiagnostic { + rank: rank + 1, + score: result.score, + document_index, + identity: DiagnosticToolIdentity::from_search_info(search_info), + }) + }) + .collect(); + self.diagnostics + .record_search(query, limit, results.len(), diagnostics); + let results = results .into_iter() .map(|result| result.document.id) .filter_map(|id| self.search_infos.get(id)) diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 9bca3d3a56..4c14a84d83 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -14,6 +14,7 @@ pub(crate) mod runtimes; pub(crate) mod sandboxing; pub(crate) mod spec_plan; pub(crate) mod tool_dispatch_trace; +pub(crate) mod tool_search_diagnostics; use std::borrow::Cow; diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index f6e50984ac..f553cb77e3 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -963,6 +963,14 @@ fn append_tool_search_executor( ) { let turn_context = context.step_context.turn.as_ref(); if !search_tool_enabled(turn_context) { + context + .tool_search_handler_cache + .diagnostics() + .record_deferred(&[], "search_disabled"); + context + .tool_search_handler_cache + .diagnostics() + .record_index(&[], "not_built"); return; } @@ -972,7 +980,15 @@ fn append_tool_search_executor( .filter(|executor| executor.exposure() == ToolExposure::Deferred) .filter_map(|executor| executor.search_info()) .collect::>(); + context + .tool_search_handler_cache + .diagnostics() + .record_deferred(&search_infos, "not_checked"); if search_infos.is_empty() { + context + .tool_search_handler_cache + .diagnostics() + .record_index(&[], "not_built"); return; } diff --git a/codex-rs/core/src/tools/tool_search_diagnostics.rs b/codex-rs/core/src/tools/tool_search_diagnostics.rs new file mode 100644 index 0000000000..7238824df8 --- /dev/null +++ b/codex-rs/core/src/tools/tool_search_diagnostics.rs @@ -0,0 +1,561 @@ +//! Always-on, bounded tool-search diagnostics retained for immediate feedback. +//! +//! The recorder stores only compact identities and latest stage snapshots. It is +//! independent of tracing so the normal feedback subscriber cannot accidentally +//! turn on verbose payloads. + +use codex_mcp::ToolInfo; +use codex_mcp::ToolSearchDiagnosticIdentity; +use codex_mcp::ToolSearchMcpDiagnosticsSnapshot; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; +use codex_tools::LoadableToolSpec; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ToolSearchInfo; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::collections::VecDeque; +use std::sync::Mutex; + +const MAX_STAGE_IDENTITIES: usize = 256; +const MAX_STAGE_CHANGES: usize = 128; +const MAX_EXCLUDED_IDENTITIES: usize = 256; +const MAX_SEARCHES: usize = 32; +const MAX_SEARCH_RESULTS: usize = 100; +const MAX_IDENTITY_CHARS: usize = 160; +const MAX_QUERY_CHARS: usize = 256; + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)] +pub(crate) struct DiagnosticToolIdentity { + server_name: String, + raw_tool_name: String, + tool_name: Option, + connector_id: Option, + connector_name: Option, + plugin_display_names: Vec, +} + +impl DiagnosticToolIdentity { + pub(crate) fn from_tool_info(tool: &ToolInfo) -> Self { + Self::from_mcp_identity(ToolSearchDiagnosticIdentity::from_tool_info(tool)) + } + + fn from_mcp_identity(identity: ToolSearchDiagnosticIdentity) -> Self { + let tool_name = match (identity.callable_namespace, identity.callable_name) { + (Some(namespace), Some(name)) => Some(format!("{namespace}.{name}")), + (None, None) => None, + (Some(namespace), None) => Some(namespace), + (None, Some(name)) => Some(name), + }; + Self { + server_name: bounded(identity.server_name, MAX_IDENTITY_CHARS), + raw_tool_name: bounded(identity.raw_tool_name, MAX_IDENTITY_CHARS), + tool_name: tool_name.map(|name| bounded(name, MAX_IDENTITY_CHARS)), + connector_id: identity + .connector_id + .map(|connector_id| bounded(connector_id, MAX_IDENTITY_CHARS)), + connector_name: identity + .connector_name + .map(|connector_name| bounded(connector_name, MAX_IDENTITY_CHARS)), + plugin_display_names: identity + .plugin_display_names + .into_iter() + .take(4) + .map(|name| bounded(name, MAX_IDENTITY_CHARS)) + .collect(), + } + } + + pub(crate) fn from_search_info(search_info: &ToolSearchInfo) -> Self { + Self { + server_name: String::new(), + raw_tool_name: String::new(), + tool_name: Some(bounded( + search_info_output_identity(search_info), + MAX_IDENTITY_CHARS, + )), + connector_id: None, + connector_name: search_info + .source_info + .as_ref() + .map(|source| bounded(source.name.clone(), MAX_IDENTITY_CHARS)), + plugin_display_names: Vec::new(), + } + } + + fn key(&self) -> String { + format!( + "{}|{}|{}|{}", + self.server_name, + self.raw_tool_name, + self.tool_name.as_deref().unwrap_or_default(), + self.connector_id.as_deref().unwrap_or_default() + ) + } +} + +#[derive(Clone, Debug, Serialize)] +pub(crate) struct ToolExposureDecision { + pub(crate) identity: DiagnosticToolIdentity, + pub(crate) included: bool, + pub(crate) model_visible: bool, + pub(crate) connector_id_present: bool, + pub(crate) connector_allowed: bool, + pub(crate) policy_enabled: bool, +} + +#[derive(Clone, Debug, Default, Serialize)] +struct StageMetadata { + #[serde(skip_serializing_if = "Vec::is_empty")] + observation_sources: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + server_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + cached_server_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + startup_complete_server_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + included_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + search_tool_enabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + cache_state: Option, + #[serde(skip)] + total_count_override: Option, + #[serde(skip)] + upstream_identities_truncated: bool, + #[serde(skip)] + fingerprint_override: Option, +} + +#[derive(Clone, Debug, Serialize)] +struct StageSnapshot { + stage: String, + inventory_generation: String, + fingerprint: String, + total_count: usize, + identities: Vec, + identities_truncated: bool, + added: Vec, + removed: Vec, + changes_truncated: bool, + #[serde(skip_serializing_if = "Vec::is_empty")] + excluded: Vec, + excluded_truncated: bool, + #[serde(flatten)] + metadata: StageMetadata, + #[serde(skip_serializing)] + retained_keys: Vec, +} + +#[derive(Clone, Debug, Serialize)] +pub(crate) struct SearchResultDiagnostic { + pub(crate) rank: usize, + pub(crate) score: f32, + pub(crate) document_index: usize, + pub(crate) identity: DiagnosticToolIdentity, +} + +#[derive(Clone, Debug, Serialize)] +struct SearchSnapshot { + stage: &'static str, + inventory_generation: String, + query: String, + requested_limit: usize, + returned_count: usize, + results: Vec, + results_truncated: bool, +} + +#[derive(Default)] +struct RecorderState { + inventory_generation: String, + stages: BTreeMap, + identity_by_tool_name: BTreeMap, + searches: VecDeque, +} + +#[derive(Default)] +pub(crate) struct ToolSearchPipelineDiagnostics { + state: Mutex, +} + +impl ToolSearchPipelineDiagnostics { + pub(crate) fn record_mcp_snapshot( + &self, + mcp: &ToolSearchMcpDiagnosticsSnapshot, + inventory: &[ToolInfo], + ) { + let raw_identities = mcp + .tools_list_responses + .iter() + .flat_map(|response| response.identities.iter().cloned()) + .map(DiagnosticToolIdentity::from_mcp_identity) + .collect::>(); + let raw_total_count = mcp + .tools_list_responses + .iter() + .map(|response| response.total_count) + .sum(); + let raw_identities_truncated = mcp + .tools_list_responses + .iter() + .any(|response| response.identities_truncated); + let raw_fingerprint = { + let mut hasher = Sha1::new(); + for response in &mcp.tools_list_responses { + hasher.update(response.server_name.as_bytes()); + hasher.update([0]); + hasher.update(response.fingerprint.as_bytes()); + hasher.update([0]); + } + format!("{:x}", hasher.finalize()) + }; + let inventory_identities = inventory + .iter() + .map(DiagnosticToolIdentity::from_tool_info) + .collect::>(); + let generation = fingerprint(&inventory_identities); + let observation_sources = mcp + .tools_list_responses + .iter() + .map(|response| format!("{}={:?}", response.server_name, response.source)) + .collect(); + + let mut state = self.state(); + state.inventory_generation = generation.clone(); + state.identity_by_tool_name.clear(); + record_stage( + &mut state, + "mcp_tools_list_response", + &generation, + raw_identities, + StageMetadata { + observation_sources, + server_count: Some(mcp.server_count), + cached_server_count: Some(mcp.cached_server_count), + startup_complete_server_count: Some(mcp.startup_complete_server_count), + total_count_override: Some(raw_total_count), + upstream_identities_truncated: raw_identities_truncated, + fingerprint_override: Some(raw_fingerprint), + ..Default::default() + }, + Vec::new(), + ); + record_stage( + &mut state, + "mcp_inventory_snapshot", + &generation, + inventory_identities, + StageMetadata { + server_count: Some(mcp.server_count), + cached_server_count: Some(mcp.cached_server_count), + startup_complete_server_count: Some(mcp.startup_complete_server_count), + ..Default::default() + }, + Vec::new(), + ); + } + + pub(crate) fn record_exposure( + &self, + decisions: Vec, + search_tool_enabled: bool, + ) { + let identities = decisions + .iter() + .map(|decision| decision.identity.clone()) + .collect(); + let included_count = decisions + .iter() + .filter(|decision| decision.included) + .count(); + let excluded = decisions + .into_iter() + .filter(|decision| !decision.included) + .collect(); + let mut state = self.state(); + let generation = state.inventory_generation.clone(); + record_stage( + &mut state, + "mcp_exposure_filter", + &generation, + identities, + StageMetadata { + included_count: Some(included_count), + search_tool_enabled: Some(search_tool_enabled), + ..Default::default() + }, + excluded, + ); + } + + pub(crate) fn record_deferred(&self, search_infos: &[ToolSearchInfo], cache_state: &str) { + self.record_search_info_stage("deferred_tool_search_info", search_infos, cache_state); + } + + pub(crate) fn record_index(&self, search_infos: &[ToolSearchInfo], cache_state: &str) { + self.record_search_info_stage("bm25_index_cache", search_infos, cache_state); + } + + pub(crate) fn record_search( + &self, + query: &str, + requested_limit: usize, + returned_count: usize, + mut results: Vec, + ) { + let mut state = self.state(); + for result in &mut results { + enrich_identity(&state, &mut result.identity); + } + let results_truncated = results.len() > MAX_SEARCH_RESULTS; + results.truncate(MAX_SEARCH_RESULTS); + let generation = state.inventory_generation.clone(); + state.searches.push_back(SearchSnapshot { + stage: "bm25_search_results", + inventory_generation: generation, + query: bounded(query.to_string(), MAX_QUERY_CHARS), + requested_limit, + returned_count, + results, + results_truncated, + }); + while state.searches.len() > MAX_SEARCHES { + state.searches.pop_front(); + } + } + + pub(crate) fn feedback_json( + &self, + thread_id: ThreadId, + session_id: SessionId, + ) -> Option> { + let state = self.state(); + if state.stages.is_empty() && state.searches.is_empty() { + return None; + } + let snapshot = FeedbackSnapshot { + version: 1, + thread_id: thread_id.to_string(), + session_id: session_id.to_string(), + inventory_generation: state.inventory_generation.clone(), + limits: FeedbackLimits { + max_stage_identities: MAX_STAGE_IDENTITIES, + max_stage_changes: MAX_STAGE_CHANGES, + max_excluded_identities: MAX_EXCLUDED_IDENTITIES, + max_searches: MAX_SEARCHES, + max_search_results: MAX_SEARCH_RESULTS, + max_identity_chars: MAX_IDENTITY_CHARS, + max_query_chars: MAX_QUERY_CHARS, + }, + stages: state.stages.values().cloned().collect(), + searches: state.searches.iter().cloned().collect(), + }; + serde_json::to_vec(&snapshot).ok() + } + + fn record_search_info_stage( + &self, + stage: &str, + search_infos: &[ToolSearchInfo], + cache_state: &str, + ) { + let mut state = self.state(); + let mut identities = search_infos + .iter() + .map(DiagnosticToolIdentity::from_search_info) + .collect::>(); + for identity in &mut identities { + enrich_identity(&state, identity); + } + let generation = state.inventory_generation.clone(); + record_stage( + &mut state, + stage, + &generation, + identities, + StageMetadata { + cache_state: Some(cache_state.to_string()), + ..Default::default() + }, + Vec::new(), + ); + } + + fn state(&self) -> std::sync::MutexGuard<'_, RecorderState> { + match self.state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + } + } +} + +#[derive(Serialize)] +struct FeedbackSnapshot { + version: u8, + thread_id: String, + session_id: String, + inventory_generation: String, + limits: FeedbackLimits, + stages: Vec, + searches: Vec, +} + +#[derive(Serialize)] +struct FeedbackLimits { + max_stage_identities: usize, + max_stage_changes: usize, + max_excluded_identities: usize, + max_searches: usize, + max_search_results: usize, + max_identity_chars: usize, + max_query_chars: usize, +} + +fn record_stage( + state: &mut RecorderState, + stage: &str, + generation: &str, + mut identities: Vec, + metadata: StageMetadata, + mut excluded: Vec, +) { + identities.sort(); + identities.dedup(); + let total_count = metadata.total_count_override.unwrap_or(identities.len()); + let stage_fingerprint = metadata + .fingerprint_override + .clone() + .unwrap_or_else(|| fingerprint(&identities)); + let retained_keys = identities + .iter() + .take(MAX_STAGE_IDENTITIES) + .map(DiagnosticToolIdentity::key) + .collect::>(); + let identities_truncated = + metadata.upstream_identities_truncated || identities.len() > MAX_STAGE_IDENTITIES; + identities.truncate(MAX_STAGE_IDENTITIES); + let excluded_truncated = excluded.len() > MAX_EXCLUDED_IDENTITIES; + excluded.truncate(MAX_EXCLUDED_IDENTITIES); + + if let Some(existing) = state.stages.get_mut(stage) + && existing.fingerprint == stage_fingerprint + { + existing.inventory_generation = generation.to_string(); + existing.total_count = total_count; + existing.identities_truncated = identities_truncated; + existing.metadata = metadata; + existing.excluded = excluded; + existing.excluded_truncated = excluded_truncated; + return; + } + + let previous_keys = state + .stages + .get(stage) + .map(|existing| { + existing + .retained_keys + .iter() + .cloned() + .collect::>() + }) + .unwrap_or_default(); + let current_keys = retained_keys.iter().cloned().collect::>(); + let mut added = current_keys + .difference(&previous_keys) + .cloned() + .collect::>(); + let mut removed = previous_keys + .difference(¤t_keys) + .cloned() + .collect::>(); + let changes_truncated = metadata.upstream_identities_truncated + || added.len() > MAX_STAGE_CHANGES + || removed.len() > MAX_STAGE_CHANGES; + added.truncate(MAX_STAGE_CHANGES); + removed.truncate(MAX_STAGE_CHANGES); + + for identity in &identities { + if let Some(tool_name) = identity.tool_name.as_ref() { + state + .identity_by_tool_name + .insert(tool_name.clone(), identity.clone()); + } + } + state.stages.insert( + stage.to_string(), + StageSnapshot { + stage: stage.to_string(), + inventory_generation: generation.to_string(), + fingerprint: stage_fingerprint, + total_count, + identities, + identities_truncated, + added, + removed, + changes_truncated, + excluded, + excluded_truncated, + metadata, + retained_keys, + }, + ); +} + +fn enrich_identity(state: &RecorderState, identity: &mut DiagnosticToolIdentity) { + let Some(tool_name) = identity.tool_name.as_ref() else { + return; + }; + let Some(known) = state.identity_by_tool_name.get(tool_name) else { + return; + }; + let source_name = identity.connector_name.clone(); + *identity = known.clone(); + if identity.connector_name.is_none() { + identity.connector_name = source_name; + } +} + +fn fingerprint(identities: &[DiagnosticToolIdentity]) -> String { + let mut hasher = Sha1::new(); + for identity in identities { + hasher.update(identity.key().as_bytes()); + hasher.update([0]); + } + format!("{:x}", hasher.finalize()) +} + +fn search_info_output_identity(search_info: &ToolSearchInfo) -> String { + match &search_info.entry.output { + LoadableToolSpec::Function(tool) => tool.name.clone(), + LoadableToolSpec::Namespace(namespace) => { + let Some(ResponsesApiNamespaceTool::Function(tool)) = namespace.tools.first() else { + return namespace.name.clone(); + }; + if namespace.tools.len() == 1 { + format!("{}.{}", namespace.name, tool.name) + } else { + format!("{} (+{} tools)", namespace.name, namespace.tools.len()) + } + } + } +} + +fn bounded(value: String, max_chars: usize) -> String { + let mut chars = value.chars(); + let bounded = chars.by_ref().take(max_chars).collect::(); + if chars.next().is_some() { + format!("{bounded}…") + } else { + bounded + } +} + +#[cfg(test)] +#[path = "tool_search_diagnostics_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/tool_search_diagnostics_tests.rs b/codex-rs/core/src/tools/tool_search_diagnostics_tests.rs new file mode 100644 index 0000000000..d7dbecdbea --- /dev/null +++ b/codex-rs/core/src/tools/tool_search_diagnostics_tests.rs @@ -0,0 +1,152 @@ +use super::*; +use codex_protocol::ThreadId; +use pretty_assertions::assert_eq; +use serde_json::Value; + +#[test] +fn feedback_snapshot_is_bounded_and_omits_sensitive_payloads() { + let recorder = ToolSearchPipelineDiagnostics::default(); + let identities = (0..(MAX_STAGE_IDENTITIES + 20)) + .map(identity) + .collect::>(); + { + let mut state = recorder.state(); + state.inventory_generation = "generation".to_string(); + record_stage( + &mut state, + "mcp_inventory_snapshot", + "generation", + identities, + StageMetadata::default(), + Vec::new(), + ); + } + for search_index in 0..(MAX_SEARCHES + 4) { + let results = (0..(MAX_SEARCH_RESULTS + 7)) + .map(|result_index| SearchResultDiagnostic { + rank: result_index + 1, + score: result_index as f32, + document_index: result_index, + identity: identity(result_index), + }) + .collect(); + recorder.record_search( + &format!("{} sensitive-search-text", "q".repeat(MAX_QUERY_CHARS + 40)), + search_index + 1, + MAX_SEARCH_RESULTS + 7, + results, + ); + } + + let bytes = recorder + .feedback_json(ThreadId::new(), ThreadId::new().into()) + .expect("snapshot should be present"); + let text = String::from_utf8(bytes.clone()).expect("json should be utf8"); + let json: Value = serde_json::from_slice(&bytes).expect("json should parse"); + + assert_eq!( + json["stages"][0]["identities"] + .as_array() + .expect("identities should be an array") + .len(), + MAX_STAGE_IDENTITIES + ); + assert_eq!( + json["searches"] + .as_array() + .expect("searches should be an array") + .len(), + MAX_SEARCHES + ); + assert_eq!( + json["searches"][0]["results"] + .as_array() + .expect("results should be an array") + .len(), + MAX_SEARCH_RESULTS + ); + assert!(!text.contains("input_schema")); + assert!(!text.contains("output_schema")); + assert!(!text.contains("search_text")); +} + +#[test] +fn changed_fingerprint_records_added_and_removed_canonical_identities() { + let recorder = ToolSearchPipelineDiagnostics::default(); + { + let mut state = recorder.state(); + state.inventory_generation = "generation".to_string(); + record_stage( + &mut state, + "mcp_inventory_snapshot", + "generation", + vec![identity(/*index*/ 1)], + StageMetadata::default(), + Vec::new(), + ); + record_stage( + &mut state, + "mcp_inventory_snapshot", + "generation", + vec![identity(/*index*/ 2)], + StageMetadata::default(), + Vec::new(), + ); + } + + let bytes = recorder + .feedback_json(ThreadId::new(), ThreadId::new().into()) + .expect("snapshot should be present"); + let json: Value = serde_json::from_slice(&bytes).expect("json should parse"); + let stage = &json["stages"][0]; + + assert_eq!( + stage["added"], + serde_json::json!(["server_2|raw_tool_2|namespace.tool_2|connector_2"]) + ); + assert_eq!( + stage["removed"], + serde_json::json!(["server_1|raw_tool_1|namespace.tool_1|connector_1"]) + ); +} + +#[test] +fn trace_subscriber_does_not_enable_verbose_tool_search_payloads() { + let recorder = ToolSearchPipelineDiagnostics::default(); + let subscriber = tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .finish(); + + tracing::subscriber::with_default(subscriber, || { + let mut state = recorder.state(); + state.inventory_generation = "generation".to_string(); + record_stage( + &mut state, + "mcp_inventory_snapshot", + "generation", + vec![identity(/*index*/ 1)], + StageMetadata::default(), + Vec::new(), + ); + }); + + let bytes = recorder + .feedback_json(ThreadId::new(), ThreadId::new().into()) + .expect("always-on snapshot should not depend on RUST_LOG"); + let text = String::from_utf8(bytes).expect("json should be utf8"); + assert!(text.contains("mcp_inventory_snapshot")); + assert!(!text.contains("input_schema")); + assert!(!text.contains("output_schema")); + assert!(!text.contains("search_text")); +} + +fn identity(index: usize) -> DiagnosticToolIdentity { + DiagnosticToolIdentity { + server_name: format!("server_{index}"), + raw_tool_name: format!("raw_tool_{index}"), + tool_name: Some(format!("namespace.tool_{index}")), + connector_id: Some(format!("connector_{index}")), + connector_name: Some(format!("connector name {index}")), + plugin_display_names: vec![format!("plugin_{index}")], + } +}