smg-mcp 2.2.0

//! MCP Tool Session — bundles all MCP execution state for a single request.
//!
//! Instead of threading `orchestrator`, `request_ctx`, `mcp_servers`,
//! and `mcp_tools` through every function, callers create one `McpToolSession` and
//! pass `&session` everywhere. When an MCP parameter changes (e.g. `mcp_servers`
//! representation), only this struct and its constructor need updating — not every
//! router function signature.

use std::collections::{HashMap, HashSet};

use futures::stream::{self, StreamExt};
use openai_protocol::responses::ResponseTool;

use super::{
    orchestrator::{McpOrchestrator, McpRequestContext, ToolExecutionInput, ToolExecutionOutput},
    UNKNOWN_SERVER_KEY,
};
use crate::{
    approval::ApprovalMode,
    inventory::{QualifiedToolName, ToolCategory, ToolEntry},
    responses_bridge::{
        build_chat_function_tools_with_names, build_function_tools_json_with_names,
        build_mcp_list_tools_item, build_mcp_list_tools_json, build_response_tools_with_names,
    },
    tenant::TenantContext,
    transform::ResponseFormat,
};

/// Default user-facing label for MCP servers when no explicit label is provided.
pub const DEFAULT_SERVER_LABEL: &str = "mcp";

/// Named pair of `(label, server_key)` for a connected MCP server.
///
/// Replaces the opaque `(String, String)` tuple that was threaded through
/// ~20 call sites, improving readability and preventing field-swap bugs.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpServerBinding {
    /// User-facing label (e.g. the `server_label` from the request).
    pub label: String,
    /// Internal key used to look up the server in the orchestrator.
    pub server_key: String,
    /// Optional per-server tool allowlist.
    ///
    /// When `Some`, only the listed tool names are exposed for this server.
    /// When `None`, all tools from the server are exposed.
    pub allowed_tools: Option<Vec<String>>,
}

#[derive(Debug, Clone)]
struct ExposedToolBinding {
    server_key: String,
    server_label: String,
    resolved_tool_name: String,
    response_format: ResponseFormat,
}

/// Bundles all MCP execution state for a single request.
///
/// Created once per request, then passed by reference to every function
/// that needs MCP infrastructure. This eliminates repeated parameter
/// threading of `orchestrator`, `request_ctx`, `mcp_servers`,
/// and `mcp_tools`.
pub struct McpToolSession<'a> {
    orchestrator: &'a McpOrchestrator,
    request_ctx: McpRequestContext<'a>,
    /// All MCP servers in this session (including builtin).
    all_mcp_servers: Vec<McpServerBinding>,
    /// Non-builtin MCP servers only — used for `mcp_list_tools` output.
    mcp_servers: Vec<McpServerBinding>,
    mcp_tools: Vec<ToolEntry>,
    exposed_name_map: HashMap<String, ExposedToolBinding>,
    exposed_name_by_qualified: HashMap<QualifiedToolName, String>,
}

impl<'a> McpToolSession<'a> {
    /// Create a new session by performing the setup every path currently repeats:
    /// 1. Create request context with default tenant and policy-only approval
    /// 2. List tools for the selected servers
    /// 3. Apply per-server allowed_tools filtering from bindings
    pub fn new(
        orchestrator: &'a McpOrchestrator,
        mcp_servers: Vec<McpServerBinding>,
        request_id: impl Into<String>,
    ) -> Self {
        let request_ctx = orchestrator.create_request_context(
            request_id,
            TenantContext::default(),
            ApprovalMode::PolicyOnly,
        );
        let server_keys: Vec<String> = mcp_servers.iter().map(|b| b.server_key.clone()).collect();
        let mut mcp_tools = Self::collect_visible_mcp_tools(orchestrator, &server_keys);

        // Build per-server allowlists from bindings that specify allowed_tools.
        let allowed_tools_by_server_key: HashMap<&str, HashSet<&str>> = mcp_servers
            .iter()
            .filter_map(|b| {
                b.allowed_tools.as_ref().map(|tools| {
                    let set: HashSet<&str> = tools
                        .iter()
                        .map(|s| s.trim())
                        .filter(|s| !s.is_empty())
                        .collect();
                    (b.server_key.as_str(), set)
                })
            })
            .collect();

        if !allowed_tools_by_server_key.is_empty() {
            mcp_tools.retain(|entry| {
                match allowed_tools_by_server_key.get(Self::associated_server_key(entry)) {
                    None => true,
                    Some(allowed) => Self::matches_allowed_tool_name(entry, allowed),
                }
            });
        }
        let (exposed_name_map, exposed_name_by_qualified) =
            Self::build_exposed_function_tools(&mcp_tools, &mcp_servers);

        // Filter out servers configured with builtin_type from the visible list.
        let builtin_names = orchestrator.builtin_server_names();
        let visible_mcp_servers: Vec<McpServerBinding> = mcp_servers
            .iter()
            .filter(|b| !builtin_names.contains(&b.server_key))
            .cloned()
            .collect();

        Self {
            orchestrator,
            request_ctx,
            all_mcp_servers: mcp_servers,
            mcp_servers: visible_mcp_servers,
            mcp_tools,
            exposed_name_map,
            exposed_name_by_qualified,
        }
    }

    // --- Accessors ---

    pub fn orchestrator(&self) -> &McpOrchestrator {
        self.orchestrator
    }

    pub fn request_ctx(&self) -> &McpRequestContext<'a> {
        &self.request_ctx
    }

    /// Returns only non-builtin MCP servers
    pub fn mcp_servers(&self) -> &[McpServerBinding] {
        &self.mcp_servers
    }

    /// Returns all MCP servers including builtin ones.
    pub fn all_mcp_servers(&self) -> &[McpServerBinding] {
        &self.all_mcp_servers
    }

    pub fn mcp_tools(&self) -> &[ToolEntry] {
        &self.mcp_tools
    }

    /// Returns true if the name is exposed to the model for this session.
    pub fn has_exposed_tool(&self, tool_name: &str) -> bool {
        self.exposed_name_map.contains_key(tool_name)
    }

    /// Returns the session's qualified-name -> exposed-name mapping.
    ///
    /// Router adapters should use this with response bridge builders.
    pub fn exposed_name_by_qualified(&self) -> &HashMap<QualifiedToolName, String> {
        &self.exposed_name_by_qualified
    }

    // --- Delegation methods ---

    /// Execute multiple tools concurrently using this session's exposed-name mapping.
    ///
    /// Uses `buffered()` to cap in-flight requests while preserving input ordering.
    pub async fn execute_tools(&self, inputs: Vec<ToolExecutionInput>) -> Vec<ToolExecutionOutput> {
        const MAX_IN_FLIGHT_TOOL_CALLS: usize = 8;
        stream::iter(inputs)
            .map(|input| self.execute_tool(input))
            .buffered(MAX_IN_FLIGHT_TOOL_CALLS)
            .collect()
            .await
    }

    /// Execute a single tool using this session's exposed-name mapping.
    pub async fn execute_tool(&self, input: ToolExecutionInput) -> ToolExecutionOutput {
        let invoked_name = input.tool_name.clone();

        if let Some(binding) = self.exposed_name_map.get(&invoked_name) {
            let resolved_tool_name = binding.resolved_tool_name.clone();
            let mut output = self
                .orchestrator
                .execute_tool_resolved(
                    ToolExecutionInput {
                        call_id: input.call_id,
                        tool_name: resolved_tool_name.clone(),
                        arguments: input.arguments,
                    },
                    &binding.server_key,
                    &binding.server_label,
                    &self.request_ctx,
                )
                .await;

            output.tool_name = invoked_name;
            output
        } else {
            let fallback_label = self
                .all_mcp_servers
                .first()
                .map(|b| b.label.as_str())
                .unwrap_or(DEFAULT_SERVER_LABEL)
                .to_string();
            let err = format!("Tool '{invoked_name}' is not in this session's exposed tool map");
            ToolExecutionOutput {
                call_id: input.call_id,
                tool_name: invoked_name.clone(),
                server_key: UNKNOWN_SERVER_KEY.to_string(),
                server_label: fallback_label,
                arguments_str: input.arguments.to_string(),
                output: serde_json::json!({ "error": &err }),
                is_error: true,
                error_message: Some(err),
                response_format: ResponseFormat::Passthrough,
                duration: std::time::Duration::default(),
            }
        }
    }

    /// Resolve the user-facing server label for a tool.
    ///
    /// Uses the orchestrator inventory to find the tool's server key, then maps
    /// it to the request's MCP server label. Falls back to the first server
    /// label (or [`DEFAULT_SERVER_LABEL`]).
    pub fn resolve_tool_server_label(&self, tool_name: &str) -> String {
        let fallback_label = self
            .all_mcp_servers
            .first()
            .map(|b| b.label.as_str())
            .unwrap_or(DEFAULT_SERVER_LABEL);

        self.exposed_name_map
            .get(tool_name)
            .map(|binding| binding.server_label.clone())
            .unwrap_or_else(|| fallback_label.to_string())
    }

    /// List tools for a single server key.
    ///
    /// Useful for emitting per-server `mcp_list_tools` items.
    pub fn list_tools_for_server(&self, server_key: &str) -> Vec<ToolEntry> {
        // Use the session's pre-filtered tool snapshot for consistency.
        self.mcp_tools
            .iter()
            .filter(|entry| Self::associated_server_key(entry) == server_key)
            .cloned()
            .collect()
    }

    /// Look up the response format for a tool.
    ///
    /// Convenience method that returns `Passthrough` if the tool is not found.
    pub fn tool_response_format(&self, tool_name: &str) -> ResponseFormat {
        self.exposed_name_map
            .get(tool_name)
            .map(|binding| binding.response_format.clone())
            .unwrap_or(ResponseFormat::Passthrough)
    }

    /// Build function-tool JSON payloads for upstream model calls.
    pub fn build_function_tools_json(&self) -> Vec<serde_json::Value> {
        build_function_tools_json_with_names(&self.mcp_tools, Some(&self.exposed_name_by_qualified))
    }

    /// Build Chat API `Tool` structs for chat completions.
    pub fn build_chat_function_tools(&self) -> Vec<openai_protocol::common::Tool> {
        build_chat_function_tools_with_names(&self.mcp_tools, Some(&self.exposed_name_by_qualified))
    }

    /// Build Responses API `ResponseTool` structs.
    pub fn build_response_tools(&self) -> Vec<ResponseTool> {
        build_response_tools_with_names(&self.mcp_tools, Some(&self.exposed_name_by_qualified))
    }

    /// Build `mcp_list_tools` JSON for a specific server.
    pub fn build_mcp_list_tools_json(
        &self,
        server_label: &str,
        server_key: &str,
    ) -> serde_json::Value {
        let tools = self.list_tools_for_server(server_key);
        build_mcp_list_tools_json(server_label, &tools)
    }

    /// Build typed `mcp_list_tools` output item for a specific server.
    pub fn build_mcp_list_tools_item(
        &self,
        server_label: &str,
        server_key: &str,
    ) -> openai_protocol::responses::ResponseOutputItem {
        let tools = self.list_tools_for_server(server_key);
        build_mcp_list_tools_item(server_label, &tools)
    }

    /// Inject MCP metadata into a response output array.
    ///
    /// Standardized ordering:
    /// 1. `mcp_list_tools` items (one per server) — prepended
    /// 2. `tool_call_items` (mcp_call / web_search_call / etc.) — after list_tools
    /// 3. Existing items (messages, etc.) — remain at end
    pub fn inject_mcp_output_items(
        &self,
        output: &mut Vec<openai_protocol::responses::ResponseOutputItem>,
        tool_call_items: Vec<openai_protocol::responses::ResponseOutputItem>,
    ) {
        // Modify the vector in-place: take existing items, then rebuild
        // with the correct ordering without allocating a temporary Vec.
        let existing = std::mem::take(output);
        output.reserve(self.mcp_servers.len() + tool_call_items.len() + existing.len());

        // 1. mcp_list_tools items (one per server)
        for binding in &self.mcp_servers {
            output.push(self.build_mcp_list_tools_item(&binding.label, &binding.server_key));
        }

        // 2. Tool call items (mcp_call / web_search_call / etc.)
        output.extend(tool_call_items);

        // 3. Existing items (messages, etc.)
        output.extend(existing);
    }

    fn build_exposed_function_tools(
        tools: &[ToolEntry],
        mcp_servers: &[McpServerBinding],
    ) -> (
        HashMap<String, ExposedToolBinding>,
        HashMap<QualifiedToolName, String>,
    ) {
        let server_labels: HashMap<&str, &str> = mcp_servers
            .iter()
            .map(|b| (b.server_key.as_str(), b.label.as_str()))
            .collect();

        let mut name_counts: HashMap<&str, usize> = HashMap::new();
        for entry in tools {
            *name_counts.entry(entry.tool_name()).or_insert(0) += 1;
        }

        let mut used_exposed_names: HashSet<String> = HashSet::with_capacity(tools.len());
        let mut name_suffixes: HashMap<String, usize> = HashMap::with_capacity(tools.len());
        let mut exposed_name_map: HashMap<String, ExposedToolBinding> =
            HashMap::with_capacity(tools.len());
        let mut exposed_name_by_qualified: HashMap<QualifiedToolName, String> =
            HashMap::with_capacity(tools.len());

        for entry in tools {
            let server_key = entry.server_key().to_string();
            let associated_server_key = Self::associated_server_key(entry);
            let server_label = server_labels
                .get(associated_server_key)
                .copied()
                .unwrap_or(associated_server_key)
                .to_string();
            let resolved_tool_name = entry.tool_name().to_string();

            let base_exposed_name = if name_counts.get(entry.tool_name()).copied().unwrap_or(0) <= 1
            {
                resolved_tool_name.clone()
            } else {
                format!(
                    "mcp_{}_{}",
                    sanitize_tool_token(&server_label),
                    sanitize_tool_token(&resolved_tool_name)
                )
            };

            let suffix = name_suffixes.entry(base_exposed_name.clone()).or_insert(0);
            let mut exposed_name = if *suffix == 0 {
                base_exposed_name.clone()
            } else {
                format!("{base_exposed_name}_{suffix}")
            };
            while used_exposed_names.contains(&exposed_name) {
                *suffix += 1;
                exposed_name = format!("{base_exposed_name}_{suffix}");
            }
            used_exposed_names.insert(exposed_name.clone());

            exposed_name_by_qualified.insert(entry.qualified_name.clone(), exposed_name.clone());

            exposed_name_map.insert(
                exposed_name,
                ExposedToolBinding {
                    server_key,
                    server_label,
                    resolved_tool_name,
                    response_format: entry.response_format.clone(),
                },
            );
        }

        (exposed_name_map, exposed_name_by_qualified)
    }

    fn collect_visible_mcp_tools(
        orchestrator: &McpOrchestrator,
        server_keys: &[String],
    ) -> Vec<ToolEntry> {
        let direct_tools = orchestrator.list_tools_for_servers(server_keys);
        let server_key_set: HashSet<&str> = server_keys.iter().map(String::as_str).collect();

        let mut aliases_by_target: HashMap<QualifiedToolName, Vec<ToolEntry>> = HashMap::new();
        for alias_entry in orchestrator
            .tool_inventory()
            .list_by_category(ToolCategory::Alias)
        {
            let Some(target) = alias_entry
                .alias_target
                .as_ref()
                .map(|alias| alias.target.clone())
            else {
                continue;
            };
            if !server_key_set.contains(target.server_key()) {
                continue;
            }
            aliases_by_target
                .entry(target)
                .or_default()
                .push(alias_entry);
        }

        let mut visible_tools = Vec::with_capacity(
            direct_tools.len() + aliases_by_target.values().map(Vec::len).sum::<usize>(),
        );

        for direct_entry in direct_tools {
            if let Some(mut alias_entries) = aliases_by_target.remove(&direct_entry.qualified_name)
            {
                visible_tools.append(&mut alias_entries);
            } else {
                visible_tools.push(direct_entry);
            }
        }

        for (_, mut alias_entries) in aliases_by_target {
            visible_tools.append(&mut alias_entries);
        }

        visible_tools
    }

    fn associated_server_key(entry: &ToolEntry) -> &str {
        entry
            .alias_target
            .as_ref()
            .map(|alias| alias.target.server_key())
            .unwrap_or_else(|| entry.server_key())
    }

    fn matches_allowed_tool_name(entry: &ToolEntry, allowed: &HashSet<&str>) -> bool {
        allowed.contains(entry.tool_name())
            || entry
                .alias_target
                .as_ref()
                .is_some_and(|alias| allowed.contains(alias.target.tool_name()))
    }
}

fn sanitize_tool_token(input: &str) -> String {
    let mut out = String::with_capacity(input.len().max(1));
    for ch in input.chars() {
        if ch.is_ascii_alphanumeric() {
            out.push(ch.to_ascii_lowercase());
        } else if ch == '_' {
            out.push(ch);
        } else {
            out.push('_');
        }
    }
    let out = out.trim_matches('_');
    if out.is_empty() {
        "tool".to_string()
    } else {
        out.to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::config::Tool as McpTool;

    #[test]
    fn test_session_creation_keeps_servers() {
        let orchestrator = McpOrchestrator::new_test();
        let mcp_servers = vec![
            McpServerBinding {
                label: "label1".to_string(),
                server_key: "key1".to_string(),
                allowed_tools: None,
            },
            McpServerBinding {
                label: "label2".to_string(),
                server_key: "key2".to_string(),
                allowed_tools: None,
            },
        ];

        let session = McpToolSession::new(&orchestrator, mcp_servers, "test-request");

        assert_eq!(session.mcp_servers().len(), 2);
        assert_eq!(session.mcp_servers()[0].label, "label1");
        assert_eq!(session.mcp_servers()[0].server_key, "key1");
    }

    #[test]
    fn test_session_empty_servers() {
        let orchestrator = McpOrchestrator::new_test();
        let session = McpToolSession::new(&orchestrator, vec![], "test-request");

        assert!(session.mcp_servers().is_empty());
        assert!(session.mcp_tools().is_empty());
    }

    #[test]
    fn test_resolve_tool_server_label_fallback() {
        let orchestrator = McpOrchestrator::new_test();
        let mcp_servers = vec![McpServerBinding {
            label: "my_label".to_string(),
            server_key: "my_key".to_string(),
            allowed_tools: None,
        }];
        let session = McpToolSession::new(&orchestrator, mcp_servers, "test-request");

        // Tool doesn't exist, should fall back to first label
        let label = session.resolve_tool_server_label("nonexistent_tool");
        assert_eq!(label, "my_label");
    }

    #[test]
    fn test_resolve_tool_server_label_no_servers() {
        let orchestrator = McpOrchestrator::new_test();
        let session = McpToolSession::new(&orchestrator, vec![], "test-request");

        // No servers, should fall back to DEFAULT_SERVER_LABEL
        let label = session.resolve_tool_server_label("nonexistent_tool");
        assert_eq!(label, DEFAULT_SERVER_LABEL);
    }

    #[test]
    fn test_tool_response_format_default() {
        let orchestrator = McpOrchestrator::new_test();
        let session = McpToolSession::new(&orchestrator, vec![], "test-request");

        let format = session.tool_response_format("nonexistent");
        assert!(matches!(format, ResponseFormat::Passthrough));
    }

    fn create_test_tool(name: &str) -> McpTool {
        use std::{borrow::Cow, sync::Arc};

        McpTool {
            name: Cow::Owned(name.to_string()),
            title: None,
            description: Some(Cow::Owned(format!("Test tool: {name}"))),
            input_schema: Arc::new(serde_json::Map::new()),
            output_schema: None,
            annotations: None,
            icons: None,
        }
    }

    #[test]
    fn test_has_exposed_tool_with_inventory() {
        let orchestrator = McpOrchestrator::new_test();

        // Insert a tool into the inventory
        let tool = create_test_tool("test_tool");
        let entry = ToolEntry::from_server_tool("server1", tool);
        orchestrator.tool_inventory().insert_entry(entry);

        let mcp_servers = vec![McpServerBinding {
            label: "label1".to_string(),
            server_key: "server1".to_string(),
            allowed_tools: None,
        }];
        let session = McpToolSession::new(&orchestrator, mcp_servers, "test-request");

        assert!(session.has_exposed_tool("test_tool"));
        assert_eq!(session.mcp_tools().len(), 1);
    }

    #[test]
    fn test_resolve_label_with_inventory() {
        let orchestrator = McpOrchestrator::new_test();

        let tool = create_test_tool("test_tool");
        let entry = ToolEntry::from_server_tool("server1", tool);
        orchestrator.tool_inventory().insert_entry(entry);

        let mcp_servers = vec![McpServerBinding {
            label: "my_server".to_string(),
            server_key: "server1".to_string(),
            allowed_tools: None,
        }];
        let session = McpToolSession::new(&orchestrator, mcp_servers, "test-request");

        let label = session.resolve_tool_server_label("test_tool");
        assert_eq!(label, "my_server");
    }

    #[test]
    fn test_exposed_names_are_unique_for_tool_name_collisions() {
        let orchestrator = McpOrchestrator::new_test();

        let tool_a = create_test_tool("shared_tool");
        let tool_b = create_test_tool("shared_tool");
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool("server1", tool_a));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool("server2", tool_b));

        let session = McpToolSession::new(
            &orchestrator,
            vec![
                McpServerBinding {
                    label: "alpha".to_string(),
                    server_key: "server1".to_string(),
                    allowed_tools: None,
                },
                McpServerBinding {
                    label: "beta".to_string(),
                    server_key: "server2".to_string(),
                    allowed_tools: None,
                },
            ],
            "test-request",
        );

        let name_a = session
            .exposed_name_by_qualified()
            .get(&QualifiedToolName::new("server1", "shared_tool"))
            .cloned()
            .expect("missing exposed name for server1 tool");
        let name_b = session
            .exposed_name_by_qualified()
            .get(&QualifiedToolName::new("server2", "shared_tool"))
            .cloned()
            .expect("missing exposed name for server2 tool");

        assert_ne!(name_a, name_b);
        assert_ne!(name_a, "shared_tool");
        assert_ne!(name_b, "shared_tool");
        assert!(session.has_exposed_tool(&name_a));
        assert!(session.has_exposed_tool(&name_b));
    }

    #[test]
    fn test_exposed_names_handle_pre_suffixed_name_conflicts() {
        let orchestrator = McpOrchestrator::new_test();

        let tool_base = create_test_tool("foo");
        let tool_suffixed = create_test_tool("foo_1");
        let tool_dup = create_test_tool("foo");

        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool("s1", tool_base));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool("s2", tool_suffixed));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool("s3", tool_dup));

        let session = McpToolSession::new(
            &orchestrator,
            vec![
                McpServerBinding {
                    label: "a".to_string(),
                    server_key: "s1".to_string(),
                    allowed_tools: None,
                },
                McpServerBinding {
                    label: "b".to_string(),
                    server_key: "s2".to_string(),
                    allowed_tools: None,
                },
                McpServerBinding {
                    label: "c".to_string(),
                    server_key: "s3".to_string(),
                    allowed_tools: None,
                },
            ],
            "test-request",
        );

        let exposed_names: HashSet<String> = session
            .exposed_name_by_qualified()
            .values()
            .cloned()
            .collect();
        assert_eq!(exposed_names.len(), 3);
    }

    // --- Builtin server filtering tests ---

    fn create_builtin_orchestrator() -> McpOrchestrator {
        use crate::core::config::{BuiltinToolType, McpConfig, McpServerConfig, McpTransport};

        let config = McpConfig {
            servers: vec![
                McpServerConfig {
                    name: "brave-builtin".to_string(),
                    transport: McpTransport::Sse {
                        url: "http://localhost:8001/sse".to_string(),
                        token: None,
                        headers: HashMap::new(),
                    },
                    proxy: None,
                    required: false,
                    tools: None,
                    builtin_type: Some(BuiltinToolType::WebSearchPreview),
                    builtin_tool_name: Some("brave_web_search".to_string()),
                },
                McpServerConfig {
                    name: "regular-server".to_string(),
                    transport: McpTransport::Sse {
                        url: "http://localhost:3000/sse".to_string(),
                        token: None,
                        headers: HashMap::new(),
                    },
                    proxy: None,
                    required: false,
                    tools: None,
                    builtin_type: None,
                    builtin_tool_name: None,
                },
            ],
            ..Default::default()
        };

        McpOrchestrator::new_test_with_config(config)
    }

    #[test]
    fn test_mcp_servers_filters_builtin() {
        let orchestrator = create_builtin_orchestrator();
        let mcp_servers = vec![
            McpServerBinding {
                label: "brave".to_string(),
                server_key: "brave-builtin".to_string(),
                allowed_tools: None,
            },
            McpServerBinding {
                label: "regular".to_string(),
                server_key: "regular-server".to_string(),
                allowed_tools: None,
            },
        ];

        let session = McpToolSession::new(&orchestrator, mcp_servers, "test-request");

        // mcp_servers() should only return non-builtin servers
        let visible = session.mcp_servers();
        assert_eq!(visible.len(), 1);
        assert_eq!(visible[0].label, "regular");
        assert_eq!(visible[0].server_key, "regular-server");

        // all_mcp_servers() should return everything
        assert_eq!(session.all_mcp_servers().len(), 2);
    }

    #[test]
    fn test_allowed_tools_filters_inventory_and_list_tools() {
        let orchestrator = McpOrchestrator::new_test();

        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_web_search"),
            ));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_local_search"),
            ));

        let session = McpToolSession::new(
            &orchestrator,
            vec![McpServerBinding {
                label: "mock".to_string(),
                server_key: "server1".to_string(),
                allowed_tools: Some(vec!["brave_web_search".to_string()]),
            }],
            "test-request",
        );

        assert!(session.has_exposed_tool("brave_web_search"));
        assert!(!session.has_exposed_tool("brave_local_search"));
        assert_eq!(session.mcp_tools().len(), 1);

        let listed = session.list_tools_for_server("server1");
        assert_eq!(listed.len(), 1);
        assert_eq!(listed[0].tool_name(), "brave_web_search");
    }

    #[test]
    fn test_alias_tools_replace_target_tool_in_session_inventory() {
        let orchestrator = McpOrchestrator::new_test();

        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_web_search"),
            ));

        orchestrator
            .register_alias(
                "web_search",
                "server1",
                "brave_web_search",
                Some(
                    crate::inventory::ArgMapping::new()
                        .with_override("enable_brave", serde_json::json!(false)),
                ),
                ResponseFormat::WebSearchCall,
            )
            .expect("alias registration should succeed");

        let session = McpToolSession::new(
            &orchestrator,
            vec![McpServerBinding {
                label: "brave".to_string(),
                server_key: "server1".to_string(),
                allowed_tools: None,
            }],
            "test-request",
        );

        assert!(session.has_exposed_tool("web_search"));
        assert!(!session.has_exposed_tool("brave_web_search"));
        assert_eq!(session.mcp_tools().len(), 1);
        assert_eq!(session.mcp_tools()[0].tool_name(), "web_search");
        assert_eq!(session.resolve_tool_server_label("web_search"), "brave");
        assert_eq!(
            session.tool_response_format("web_search"),
            ResponseFormat::WebSearchCall
        );

        let listed = session.list_tools_for_server("server1");
        assert_eq!(listed.len(), 1);
        assert_eq!(listed[0].tool_name(), "web_search");
    }

    #[test]
    fn test_allowed_tools_accepts_alias_name() {
        let orchestrator = McpOrchestrator::new_test();

        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_web_search"),
            ));

        orchestrator
            .register_alias(
                "web_search",
                "server1",
                "brave_web_search",
                None,
                ResponseFormat::WebSearchCall,
            )
            .expect("alias registration should succeed");

        let session = McpToolSession::new(
            &orchestrator,
            vec![McpServerBinding {
                label: "brave".to_string(),
                server_key: "server1".to_string(),
                allowed_tools: Some(vec!["web_search".to_string()]),
            }],
            "test-request",
        );

        assert!(session.has_exposed_tool("web_search"));
        assert_eq!(session.mcp_tools().len(), 1);
        assert_eq!(session.mcp_tools()[0].tool_name(), "web_search");
    }

    /// Verify that `inject_mcp_output_items` produces the exact ordering:
    ///   1. mcp_list_tools items (one per server, in server order)
    ///   2. tool_call_items (in their original order)
    ///   3. existing output items (in their original order)
    ///
    /// This is a regression test so future perf refactors cannot
    /// accidentally change the output ordering contract.
    #[test]
    fn test_inject_mcp_output_items_ordering() {
        use openai_protocol::responses::ResponseOutputItem;

        let orchestrator = McpOrchestrator::new_test();

        // Register one tool per server so build_mcp_list_tools_item has
        // something to return.
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "srv_a",
                create_test_tool("tool_a"),
            ));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "srv_b",
                create_test_tool("tool_b"),
            ));

        let session = McpToolSession::new(
            &orchestrator,
            vec![
                McpServerBinding {
                    label: "Server A".to_string(),
                    server_key: "srv_a".to_string(),
                    allowed_tools: None,
                },
                McpServerBinding {
                    label: "Server B".to_string(),
                    server_key: "srv_b".to_string(),
                    allowed_tools: None,
                },
            ],
            "test-ordering",
        );

        // Pre-existing output items (e.g. assistant message).
        let existing_1 = ResponseOutputItem::Message {
            id: "msg_existing_1".to_string(),
            role: "assistant".to_string(),
            content: vec![],
            status: "completed".to_string(),
        };
        let existing_2 = ResponseOutputItem::Message {
            id: "msg_existing_2".to_string(),
            role: "assistant".to_string(),
            content: vec![],
            status: "completed".to_string(),
        };

        // Tool call items injected by the router.
        let call_1 = ResponseOutputItem::McpCall {
            id: "call_1".to_string(),
            status: "completed".to_string(),
            approval_request_id: None,
            arguments: "{}".to_string(),
            error: None,
            name: "tool_a".to_string(),
            output: "result_a".to_string(),
            server_label: "Server A".to_string(),
        };
        let call_2 = ResponseOutputItem::McpCall {
            id: "call_2".to_string(),
            status: "completed".to_string(),
            approval_request_id: None,
            arguments: "{}".to_string(),
            error: None,
            name: "tool_b".to_string(),
            output: "result_b".to_string(),
            server_label: "Server B".to_string(),
        };

        let mut output = vec![existing_1, existing_2];
        let tool_call_items = vec![call_1, call_2];

        session.inject_mcp_output_items(&mut output, tool_call_items);

        // Expected ordering: 2 mcp_list_tools + 2 mcp_call + 2 messages = 6
        assert_eq!(output.len(), 6, "expected 6 items in output");

        // Serialize to JSON values for easier field-level assertions.
        let items: Vec<serde_json::Value> = output
            .iter()
            .map(|item| serde_json::to_value(item).expect("serialization failed"))
            .collect();

        // [0..2] mcp_list_tools — one per server, in server order
        assert_eq!(items[0]["type"], "mcp_list_tools");
        assert_eq!(items[0]["server_label"], "Server A");
        assert_eq!(items[1]["type"], "mcp_list_tools");
        assert_eq!(items[1]["server_label"], "Server B");

        // [2..4] tool call items in original order
        assert_eq!(items[2]["type"], "mcp_call");
        assert_eq!(items[2]["id"], "call_1");
        assert_eq!(items[3]["type"], "mcp_call");
        assert_eq!(items[3]["id"], "call_2");

        // [4..6] existing items in original order
        assert_eq!(items[4]["type"], "message");
        assert_eq!(items[4]["id"], "msg_existing_1");
        assert_eq!(items[5]["type"], "message");
        assert_eq!(items[5]["id"], "msg_existing_2");
    }

    #[test]
    fn test_allowed_tools_filters_only_target_server() {
        let orchestrator = McpOrchestrator::new_test();

        // server1 has two tools
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_web_search"),
            ));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server1",
                create_test_tool("brave_local_search"),
            ));

        // server2 has two tools
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server2",
                create_test_tool("deepwiki_search"),
            ));
        orchestrator
            .tool_inventory()
            .insert_entry(ToolEntry::from_server_tool(
                "server2",
                create_test_tool("deepwiki_read"),
            ));

        let session = McpToolSession::new(
            &orchestrator,
            vec![
                McpServerBinding {
                    label: "brave".to_string(),
                    server_key: "server1".to_string(),
                    allowed_tools: Some(vec!["brave_web_search".to_string()]),
                },
                McpServerBinding {
                    label: "deepwiki".to_string(),
                    server_key: "server2".to_string(),
                    allowed_tools: None,
                },
            ],
            "test-request",
        );

        // server1 is filtered
        assert!(session.has_exposed_tool("brave_web_search"));
        assert!(!session.has_exposed_tool("brave_local_search"));
        let listed_server1 = session.list_tools_for_server("server1");
        assert_eq!(listed_server1.len(), 1);
        assert_eq!(listed_server1[0].tool_name(), "brave_web_search");

        // server2 is unfiltered
        assert!(session.has_exposed_tool("deepwiki_search"));
        assert!(session.has_exposed_tool("deepwiki_read"));
        let listed_server2 = session.list_tools_for_server("server2");
        assert_eq!(listed_server2.len(), 2);
    }
}