openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-08 15:50:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Capture tool search pipeline diagnostics in feedback

Browse Source
This commit is contained in:
Steven Lee
2026-07-07 20:43:08 +00:00
parent a6b99ee5c4
commit f3a679e055
16 changed files with 1161 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
codex-rs/app-server/src/request_processors/feedback_processor.rs
View File

@@ -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 || {

77
codex-rs/codex-mcp/src/connection_manager.rs
View File

@@ -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<BTreeMap<String, ToolSearchToolsListResponseSnapshot>>,
}
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<ToolInfo> {
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<ToolInfo>, 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<String, ToolSearchToolsListResponseSnapshot>,
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 (

4
codex-rs/codex-mcp/src/connection_manager_tests.rs
View File

@@ -119,6 +119,10 @@ async fn create_test_managed_client(tools: Vec<ToolInfo>) -> 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,
}

5
codex-rs/codex-mcp/src/lib.rs
View File

@@ -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;

35
codex-rs/codex-mcp/src/rmcp_client.rs
View File

@@ -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<Duration>,
pub(crate) server_instructions: Option<String>,
pub(crate) server_supports_sandbox_state_meta_capability: bool,
pub(crate) tools_list_response: ToolSearchToolsListResponseSnapshot,
pub(crate) codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
}
@@ -516,23 +518,26 @@ impl AsyncManagedClient {
.map(|tools| filter_tools(tools, &self.tool_filter))
}
pub(crate) async fn listed_tools(&self) -> Option<Vec<ToolInfo>> {
pub(crate) async fn listed_tools(
&self,
) -> Option<(Vec<ToolInfo>, Option<ToolSearchToolsListResponseSnapshot>)> {
// 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<RmcpClient>,
timeout: Option<Duration>,
server_instructions: Option<&str>,
) -> Result<Vec<ToolInfo>> {
) -> Result<ListedToolsFetch> {
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<ToolInfo>,
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,
};

164
codex-rs/codex-mcp/src/tool_search_diagnostics.rs Normal file
View File

@@ -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<String>,
pub callable_name: Option<String>,
pub connector_id: Option<String>,
pub connector_name: Option<String>,
pub plugin_display_names: Vec<String>,
}
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<ToolSearchDiagnosticIdentity>,
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<ToolSearchDiagnosticIdentity>,
) -> 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<ToolSearchToolsListResponseSnapshot>,
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::<String>();
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())
}

14
codex-rs/core/src/codex_thread.rs
View File

@@ -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<Vec<u8>> {
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()
}

57
codex-rs/core/src/mcp_tool_exposure.rs
View File

@@ -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<McpToolInfo>,
pub(crate) deferred_tools: Option<Vec<McpToolInfo>>,
pub(crate) diagnostic_decisions: Vec<ToolExposureDecision>,
}
#[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<ToolExposureDecision> {
let allowed = connectors
.into_iter()
.flatten()
.map(|connector| connector.id.as_str())
.collect::<HashSet<_>>();
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;

19
codex-rs/core/src/session/step_context.rs
View File

@@ -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<ToolInfo>,
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<ResolvedSelectedCapabilityRoot>,
/// The exact MCP config and manager used to advertise and execute tools for this step.
pub(crate) mcp: Arc<McpRuntimeSnapshot>,
/// The fixed MCP tool list used for this exact sampling request.
mcp_tool_snapshot: OnceCell<Vec<ToolInfo>>,
/// The fixed MCP tool list and compact diagnostics used for this exact sampling request.
mcp_tool_snapshot: OnceCell<McpToolSnapshot>,
/// The canonical AGENTS.md value observed with this environment snapshot.
pub(crate) loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
}
@@ -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
}
}

14
codex-rs/core/src/session/turn.rs
View File

@@ -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(

2
codex-rs/core/src/session/world_state.rs
View File

@@ -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,

51
codex-rs/core/src/tools/handlers/tool_search.rs
View File

@@ -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<ToolSearchInfo>,
spec: ToolSpec,
search_engine: SearchEngine<usize>,
diagnostics: Arc<ToolSearchPipelineDiagnostics>,
}
#[derive(Default)]
pub(crate) struct ToolSearchHandlerCache {
cached: Mutex<Option<Arc<ToolSearchHandler>>>,
diagnostics: Arc<ToolSearchPipelineDiagnostics>,
}
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<Arc<ToolSearchHandler>>> {
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<ToolSearchInfo>) -> Self {
Self::new_with_diagnostics(search_infos, Arc::new(Default::default()))
}
fn new_with_diagnostics(
search_infos: Vec<ToolSearchInfo>,
diagnostics: Arc<ToolSearchPipelineDiagnostics>,
) -> 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<Vec<LoadableToolSpec>, 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))

1
codex-rs/core/src/tools/mod.rs
View File

@@ -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;

16
codex-rs/core/src/tools/spec_plan.rs
View File

@@ -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::<Vec<_>>();
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;
}

561
codex-rs/core/src/tools/tool_search_diagnostics.rs Normal file
View File

@@ -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<String>,
connector_id: Option<String>,
connector_name: Option<String>,
plugin_display_names: Vec<String>,
}
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<String>,
#[serde(skip_serializing_if = "Option::is_none")]
server_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
cached_server_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
startup_complete_server_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
included_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
search_tool_enabled: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
cache_state: Option<String>,
#[serde(skip)]
total_count_override: Option<usize>,
#[serde(skip)]
upstream_identities_truncated: bool,
#[serde(skip)]
fingerprint_override: Option<String>,
}
#[derive(Clone, Debug, Serialize)]
struct StageSnapshot {
stage: String,
inventory_generation: String,
fingerprint: String,
total_count: usize,
identities: Vec<DiagnosticToolIdentity>,
identities_truncated: bool,
added: Vec<String>,
removed: Vec<String>,
changes_truncated: bool,
#[serde(skip_serializing_if = "Vec::is_empty")]
excluded: Vec<ToolExposureDecision>,
excluded_truncated: bool,
#[serde(flatten)]
metadata: StageMetadata,
#[serde(skip_serializing)]
retained_keys: Vec<String>,
}
#[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<SearchResultDiagnostic>,
results_truncated: bool,
}
#[derive(Default)]
struct RecorderState {
inventory_generation: String,
stages: BTreeMap<String, StageSnapshot>,
identity_by_tool_name: BTreeMap<String, DiagnosticToolIdentity>,
searches: VecDeque<SearchSnapshot>,
}
#[derive(Default)]
pub(crate) struct ToolSearchPipelineDiagnostics {
state: Mutex<RecorderState>,
}
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::<Vec<_>>();
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::<Vec<_>>();
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<ToolExposureDecision>,
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<SearchResultDiagnostic>,
) {
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<Vec<u8>> {
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::<Vec<_>>();
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<StageSnapshot>,
searches: Vec<SearchSnapshot>,
}
#[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<DiagnosticToolIdentity>,
metadata: StageMetadata,
mut excluded: Vec<ToolExposureDecision>,
) {
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::<Vec<_>>();
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::<BTreeSet<_>>()
})
.unwrap_or_default();
let current_keys = retained_keys.iter().cloned().collect::<BTreeSet<_>>();
let mut added = current_keys
.difference(&previous_keys)
.cloned()
.collect::<Vec<_>>();
let mut removed = previous_keys
.difference(&current_keys)
.cloned()
.collect::<Vec<_>>();
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::<String>();
if chars.next().is_some() {
format!("{bounded}…")
} else {
bounded
}
}
#[cfg(test)]
#[path = "tool_search_diagnostics_tests.rs"]
mod tests;

152
codex-rs/core/src/tools/tool_search_diagnostics_tests.rs Normal file
View File

@@ -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::<Vec<_>>();
{
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}")],
}
}