//! MCP server integration using the official rmcp SDK.

use std::collections::HashMap;
use std::ffi::OsString;
use std::future::Future;
use std::sync::{Arc, LazyLock};

use agent_client_protocol as acp;
use futures::StreamExt;
use regex::Regex;
use tokio::{
    io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader},
    process::{ChildStderr, Command},
    sync::{Mutex, Notify},
};

use rmcp::{
    ClientHandler, ServiceExt,
    model::{
        CallToolRequestParams, ClientCapabilities, ClientInfo, Implementation,
        PaginatedRequestParams,
    },
    service::{
        ClientInitializeError, NotificationContext, RoleClient, RunningService, ServiceError,
    },
    service::{RxJsonRpcMessage, TxJsonRpcMessage},
    transport::{
        StreamableHttpClientTransport, Transport,
        streamable_http_client::StreamableHttpClientTransportConfig,
    },
};

use crate::oauth_config::McpOAuthConfig;

use kigi_tools::types::{
    output::{MCPOutput, MCPOutputDetails, ToolOutput},
    tool::{ToolKind, ToolNamespace},
    tool_metadata::ToolMetadata,
};
use kigi_tools::util::ProcessGroup;

/// MCP tool name delimiter: server names are qualified as `"server__tool"`.
/// Canonical definition lives in `kigi_workspace_types`; re-exported here
/// for callers that historically imported it from this module.
pub use kigi_workspace_types::MCP_TOOL_NAME_DELIMITER;

/// Normalize an MCP server URL for comparison: strip trailing slashes.
/// Must match the normalization the host's managed-config layer uses
/// (e.g. shell's `session::managed_mcp::normalize_url`) so URL
/// lookup keys agree.
fn normalize_url(url: &str) -> String {
    url.trim_end_matches('/').to_string()
}

/// Regex for strictest cross-provider tool name validation.
///
/// Requirements across providers:
/// - Anthropic/OpenAI: `^[a-zA-Z0-9_-]{1,64}$` (allows starting with digit/hyphen)
/// - Google Gemini: `^[a-zA-Z_][a-zA-Z0-9_.-]{0,63}$` (must start with letter/underscore, allows dots)
///
/// Strictest common denominator: must start with letter/underscore, only alphanumeric/_/- allowed, max 64 chars.
static TOOL_NAME_REGEX: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_-]{0,63}$").unwrap());

/// Validate that a tool name matches the strictest cross-provider LLM API requirements.
///
/// Pattern: `^[a-zA-Z_][a-zA-Z0-9_-]{0,63}$`
/// - Must start with a letter or underscore (Gemini requirement)
/// - Only letters, digits, underscores, hyphens allowed (no dots — Anthropic/OpenAI requirement)
/// - Maximum 64 characters
///
/// Returns `Ok(())` if valid, or `Err(reason)` if invalid.
pub fn validate_tool_name(name: &str) -> Result<(), String> {
    if name.is_empty() {
        return Err("tool name cannot be empty".to_string());
    }
    if !TOOL_NAME_REGEX.is_match(name) {
        return Err(format!(
            "tool name '{}' is invalid — must match ^[a-zA-Z_][a-zA-Z0-9_-]{{0,63}}$ (start with letter/underscore, max 64 chars)",
            name
        ));
    }
    Ok(())
}

/// Sanitize an MCP server or tool name into a single safe path segment
/// (e.g. `"user-Hugging Face"` becomes `user-Hugging_Face`). Shared so the
/// per-server folder advertised in the prompt matches the tool files on disk.
pub fn sanitize_descriptor_segment(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' {
            out.push(c);
        } else {
            out.push('_');
        }
    }
    if out.is_empty() {
        out.push('_');
    }
    out
}

/// Result of a diff-based MCP config update.
#[derive(Debug)]
pub struct McpConfigDiff {
    /// Server names that are new or had their config changed.
    pub added: Vec<McpServerName>,
    /// Server names that were removed or had their config changed (old instance torn down).
    pub removed: Vec<McpServerName>,
    /// Server names whose config is identical — clients kept alive.
    pub retained: Vec<McpServerName>,
}

/// MCP server name used as the key in client/tool maps (e.g. `"github"`).
pub type McpServerName = String;

/// Unqualified MCP tool name (e.g. `"create_issue"`, without the `server__` prefix).
type ToolName = String;

/// Typed state machine for MCP-pool initialization.
///
/// Replaces the previous trio of correlated fields — `initialized: bool`,
/// `initializing: bool`, `initializing_servers: HashSet<McpServerName>` —
/// whose product space could represent nonsensical combinations such as
/// "initialized AND initializing" or "no init started AND per-server
/// handshakes outstanding". With one enum field, every legal state has
/// exactly one representation and the compiler enforces exhaustiveness
/// at every match site.
///
/// Lifecycle:
///
/// ```text
///   ┌─────────────┐  try_start  ┌──────────────────────┐
///   │  NotStarted │ ──────────▶ │  Starting{handshakes}│
///   └─────────────┘ ◀── cancel ─┴──────────┬───────────┘
///         ▲                                │ finish
///         │ cancel                         ▼
///         │                  ┌──────────────────────────┐
///         └──────────────────┤  Finished{handshakes}    │
///                            └──────────────────────────┘
/// ```
///
/// `Starting` is the pre-`finish_init` window; `Finished` is the post-
/// `finish_init` window where per-server background handshakes may still
/// be draining. `is_complete()` requires `Finished` with an empty
/// handshaking set.
#[derive(Debug, Default)]
pub enum InitProgress {
    /// Init has never been started, or was cancelled / reset by a
    /// config change.
    #[default]
    NotStarted,
    /// `try_start_init` was called; per-server tasks may be spawning;
    /// `finish_init` has NOT yet fired. `handshaking` tracks the set of
    /// servers whose background handshake is in flight.
    Starting {
        handshaking: std::collections::HashSet<McpServerName>,
    },
    /// `finish_init` fired (deliberately early, so the session is not
    /// blocked on MCP for non-MCP work). Background per-server
    /// handshakes may still be running; `handshaking` shrinks as each
    /// completes. `is_complete()` returns `true` only when it is empty.
    Finished {
        handshaking: std::collections::HashSet<McpServerName>,
    },
}

impl InitProgress {
    /// True iff every per-server handshake has settled and `finish_init`
    /// has fired. Pairs with [`Self::is_in_progress`].
    pub fn is_complete(&self) -> bool {
        matches!(self, Self::Finished { handshaking } if handshaking.is_empty())
    }

    /// True iff any init work is outstanding — either we are pre-
    /// `finish_init`, or per-server handshakes are still in flight in
    /// the background.
    pub fn is_in_progress(&self) -> bool {
        match self {
            Self::Starting { .. } => true,
            Self::Finished { handshaking } => !handshaking.is_empty(),
            Self::NotStarted => false,
        }
    }

    /// True iff `finish_init` has fired, regardless of whether
    /// background handshakes are still draining. Used for diagnostic
    /// logging where the caller wants to distinguish pre-finish from
    /// post-finish-with-bg-work.
    pub fn has_finished_init(&self) -> bool {
        matches!(self, Self::Finished { .. })
    }

    /// True iff the named server is currently handshaking.
    pub fn is_server_handshaking(&self, name: &str) -> bool {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => {
                handshaking.contains(name)
            }
            Self::NotStarted => false,
        }
    }

    /// Iterate over server names whose background handshake is still
    /// in flight. Empty when [`Self::NotStarted`] or fully complete.
    pub fn handshaking_servers(&self) -> impl Iterator<Item = &McpServerName> {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => {
                Some(handshaking.iter())
            }
            Self::NotStarted => None,
        }
        .into_iter()
        .flatten()
    }

    /// Number of in-flight per-server handshakes.
    pub fn handshaking_count(&self) -> usize {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => handshaking.len(),
            Self::NotStarted => 0,
        }
    }

    /// Transition `NotStarted` → `Starting { ∅ }`. Returns `true` on
    /// successful transition, `false` if init was already started or
    /// finished (mirrors the pre-refactor `try_start_init` contract).
    pub fn try_start(&mut self) -> bool {
        if matches!(self, Self::NotStarted) {
            *self = Self::Starting {
                handshaking: std::collections::HashSet::new(),
            };
            true
        } else {
            false
        }
    }

    /// Transition `Starting { hs }` → `Finished { hs }`, preserving the
    /// handshaking set. No-op if already `Finished`; no-op-with-log if
    /// called from `NotStarted` (defensive — that would be a caller bug).
    pub fn finish(&mut self) {
        match self {
            Self::Starting { handshaking } => {
                // Move only the inner set, leaving the outer `&mut self`
                // ready to be reassigned without going through
                // `mem::take(self)` (which would force a `NotStarted`
                // placeholder and a redundant put-back in the
                // already-Finished arm below).
                let handshaking = std::mem::take(handshaking);
                *self = Self::Finished { handshaking };
            }
            // Idempotent: already past the finish boundary. Per-server
            // handshakes continue draining via `mark_handshake_complete`.
            Self::Finished { .. } => {}
            Self::NotStarted => {
                tracing::warn!(
                    "InitProgress::finish called from NotStarted; staying in NotStarted"
                );
            }
        }
    }

    /// Transition any state → `NotStarted`. Clears all per-server
    /// progress. Used on generation mismatch (config change racing
    /// active init) and on full reset.
    pub fn cancel(&mut self) {
        *self = Self::NotStarted;
    }

    /// Add names to the handshaking set. Only meaningful in `Starting`
    /// or `Finished`; warns if called from `NotStarted` (that would
    /// mean a per-server handshake started without a `try_start_init`,
    /// which is a caller bug).
    pub fn mark_handshaking(&mut self, names: impl IntoIterator<Item = McpServerName>) {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => {
                handshaking.extend(names);
            }
            Self::NotStarted => {
                tracing::warn!("InitProgress::mark_handshaking called from NotStarted; ignoring");
            }
        }
    }

    /// Remove a server from the handshaking set (on success or failure
    /// of its handshake). No-op if not present or if `NotStarted`.
    pub fn mark_handshake_complete(&mut self, name: &str) {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => {
                handshaking.remove(name);
            }
            Self::NotStarted => {}
        }
    }

    /// Clear the handshaking set entirely. Used by the proxy-mode and
    /// bg-handshake completion paths as a defensive sweep after the
    /// per-server `mark_handshake_complete` calls — ensures the set is
    /// empty before/after `finish_init` fires.
    pub fn clear_handshaking(&mut self) {
        match self {
            Self::Starting { handshaking } | Self::Finished { handshaking } => {
                handshaking.clear();
            }
            Self::NotStarted => {}
        }
    }
}

/// One in-process SDK MCP server registration: its tool-namespace name and the
/// SDK-side id echoed back in `kigi/mcp/sdk_call`. A named struct (rather than a
/// `(String, String)` tuple) so callers can't transpose the two strings.
///
/// `Deserialize`d straight from a `_meta["kigi/mcp/servers"]` entry, so the
/// `serverId` wire field name is declared (and serde-checked) exactly once here.
#[derive(Debug, Clone, serde::Deserialize)]
pub struct AcpServerEntry {
    pub name: McpServerName,
    #[serde(rename = "serverId")]
    pub server_id: String,
}

/// The session's in-process SDK MCP servers (declared via `_meta["kigi/mcp/servers"]`,
/// reached over the ACP reverse channel), bundled with the shared reverse-RPC invoker.
/// Held as `McpState::acp_mcp: Option<_>` so the set is one atom — present together or
/// absent, never "servers without an invoker" — and survives `update_configs` clears
/// (config reloads only touch `configs`/`owned_clients`). Per-server config.toml overrides
/// are NOT cached here — they are re-resolved per init (see [`McpState::build_pending_acp_clients`]).
struct AcpMcpRegistry {
    /// Registered servers (`name -> serverId`).
    servers: Vec<AcpServerEntry>,
    /// Shared reverse-RPC invoker all these servers' tools are called through (emits
    /// `kigi/mcp/sdk_call` over the ACP connection).
    invoker: Arc<dyn crate::acp_transport::AcpReverseInvoker>,
}

/// Consolidated MCP state behind a single lock. Generation counter detects stale inits.
pub struct McpState {
    pub configs: Vec<acp::McpServer>,
    pub meta_config_map: McpMetaConfigMap,
    /// Clients owned by this session; cleared on config changes.
    pub owned_clients: HashMap<McpServerName, Arc<McpClient>>,
    /// Clients inherited from parent via `SharedMcpPool`; never cleared by config changes.
    pub shared_clients: HashMap<McpServerName, Arc<McpClient>>,
    /// The session's in-process SDK MCP servers + their shared invoker/overrides; `None`
    /// when the session has none. See [`AcpMcpRegistry`]. Kept out of `configs` (the closed
    /// `acp::McpServer` enum) so it survives `update_configs` clears.
    acp_mcp: Option<AcpMcpRegistry>,
    /// Encapsulated init lifecycle. Access via [`Self::is_initialized`],
    /// [`Self::is_initializing`], [`Self::try_start_init`],
    /// [`Self::finish_init`], etc. — those route through a single
    /// [`InitProgress`] state machine that rules out nonsensical
    /// combinations like "initialized AND initializing".
    ///
    /// Private on purpose: external callers must go through the typed
    /// transition methods, not poke the variant directly.
    init_progress: InitProgress,
    pub generation: u64,
    /// Qualified tool name → `_meta` from MCP tools/list. Populated during init.
    pub mcp_tool_meta: HashMap<String, serde_json::Value>,
    /// HTTP servers that support OAuth but haven't been authenticated yet.
    pub auth_required: std::collections::HashSet<McpServerName>,
    /// Servers whose background init failed (handshake error, `tools/list`
    /// error, or overall init timeout) even though a client object exists,
    /// mapped to a short failure cause surfaced to the model in the MCP
    /// reminder. Surfaced as `Unavailable` in status snapshots so a server
    /// that connected but never finished initializing — e.g. wedged on
    /// `tools/list` and registered zero tools — does not misleadingly show
    /// as `Ready`. Cleared when the server begins a fresh init attempt.
    pub init_failed: std::collections::HashMap<McpServerName, String>,
    /// Per-server set of unqualified tool names that the user has disabled.
    /// Persisted to `~/.kigi/config.toml` under `[mcp_servers.<name>].disabled_tools`.
    pub disabled_tools: HashMap<McpServerName, std::collections::HashSet<ToolName>>,
    /// Stashed registrations for disabled tools so they can be re-enabled
    /// without a full MCP re-init (no need to call `list_tools` again).
    pub disabled_tool_registrations: HashMap<String, McpToolRegistration>,
    event_writer: kigi_file_utils::events::EventWriter,
    /// Sender wired by the session actor to its `StatusDispatcher`
    /// task.  When `Some`, the state — and every [`McpClient`] reached
    /// through [`Self::all_clients`] / [`Self::get_client`] — forwards
    /// [`McpClientEvent`]s here for coalescing and fan-out as ACP
    /// `kigi/mcp/server_status` notifications.
    ///
    /// Intentionally `None` in subagent-pool / shared-pool snapshots
    /// ([`SharedMcpPool`]) where the **parent** session is the
    /// single owner of liveness/notification flow. Clients in those
    /// snapshots inherit the parent's `Arc<McpClient>` (with the
    /// parent's `notify_tx` slot still pointing at the parent), so
    /// duplicating event flow into a subagent would just double-push
    /// every event.
    ///
    /// Populated by [`Self::set_client_event_tx`], which fans the
    /// sender into every existing client's `notify_tx` slot.
    ///
    /// **Private on purpose.** Callers MUST go through
    /// [`Self::set_client_event_tx`] so the sender is fanned out into
    /// every existing `owned_clients` entry; a direct field write
    /// (`state.client_event_tx = Some(tx)`) would leave all
    /// already-owned clients with `notify_tx = None`, silently
    /// dropping `tools/list_changed`, `Ready`, and `HandshakeFailed`
    /// emits for them. Read access is via [`Self::client_event_tx`].
    client_event_tx: Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>>,
}

impl McpState {
    pub fn new(configs: Vec<acp::McpServer>) -> Self {
        Self::new_with_meta(configs, McpMetaConfigMap::new())
    }

    pub fn new_with_meta(configs: Vec<acp::McpServer>, meta_config_map: McpMetaConfigMap) -> Self {
        Self {
            configs,
            meta_config_map,
            owned_clients: HashMap::new(),
            shared_clients: HashMap::new(),
            acp_mcp: None,
            init_progress: InitProgress::default(),
            generation: 0,
            mcp_tool_meta: HashMap::new(),
            auth_required: std::collections::HashSet::new(),
            init_failed: HashMap::new(),
            disabled_tools: HashMap::new(),
            disabled_tool_registrations: HashMap::new(),
            event_writer: kigi_file_utils::events::EventWriter::noop(),
            client_event_tx: None,
        }
    }

    /// Install (or remove) the [`McpClientEvent`] sender owned by the
    /// session-actor `StatusDispatcher`.
    ///
    /// Synchronous: the per-client slot is a `parking_lot::Mutex`, so
    /// the iteration no longer holds `&mut McpState` across `.await`.
    ///
    /// Side effect: clones the sender into every existing client's
    /// shared `notify_tx` slot. New clients added later (e.g. on a
    /// config diff that re-spawns a server) MUST be wired by the
    /// caller post-construction — typically by calling
    /// [`McpClient::set_event_tx`] **before**
    /// `get_tool_registrations` (so `ensure_initialized`'s
    /// `Ready`/`HandshakeFailed` emit fires with `Some(tx)` and the
    /// `KigiClientHandler` cloned during `try_handshake` reads
    /// through the same Arc).
    pub fn set_client_event_tx(
        &mut self,
        tx: Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>>,
    ) {
        self.client_event_tx = tx.clone();
        for client in self.owned_clients.values() {
            client.set_event_tx(tx.clone());
        }
        // Shared clients are intentionally NOT wired here: see the
        // `client_event_tx` doc-comment for why a subagent must not
        // duplicate the parent's event flow.
    }

    /// Read-only access to the installed [`McpClientEvent`] sender.
    ///
    /// Returns a clone of the sender wired by
    /// [`Self::set_client_event_tx`], or `None` if no dispatcher is
    /// attached (subagent / shared-pool snapshot). Exposed as a getter
    /// rather than a `pub` field so the fan-out contract documented on
    /// `client_event_tx` cannot be bypassed by a direct assignment.
    pub fn client_event_tx(&self) -> Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>> {
        self.client_event_tx.clone()
    }

    pub fn set_event_writer(&mut self, writer: kigi_file_utils::events::EventWriter) {
        self.event_writer = writer;
    }

    pub fn event_writer(&self) -> &kigi_file_utils::events::EventWriter {
        &self.event_writer
    }

    /// Register the session's in-process SDK MCP servers (`name -> serverId`) plus the
    /// reverse-RPC invoker. Held across `update_configs` clears so each init re-adds them.
    pub fn set_acp_servers(
        &mut self,
        servers: Vec<AcpServerEntry>,
        invoker: Arc<dyn crate::acp_transport::AcpReverseInvoker>,
    ) {
        self.acp_mcp = Some(AcpMcpRegistry { servers, invoker });
    }

    /// Whether any in-process SDK MCP servers are registered (so the session knows to
    /// run MCP init even with no `configs`).
    pub fn has_acp_servers(&self) -> bool {
        self.acp_mcp
            .as_ref()
            .is_some_and(|acp| !acp.servers.is_empty())
    }

    /// Registered SDK servers not yet connected (no owned/shared client) — the ones an
    /// init pass should build. Shared by [`build_pending_acp_clients`] and
    /// [`pending_acp_server_names`] so the "what to build" filter lives in one place.
    fn pending_acp_entries(&self) -> impl Iterator<Item = &AcpServerEntry> {
        self.acp_mcp.iter().flat_map(|acp| {
            acp.servers.iter().filter(|entry| {
                !self.owned_clients.contains_key(&entry.name)
                    && !self.shared_clients.contains_key(&entry.name)
            })
        })
    }

    /// Names of the SDK servers [`build_pending_acp_clients`] will build — used to mark
    /// them initializing before the (async) build.
    pub fn pending_acp_server_names(&self) -> Vec<String> {
        self.pending_acp_entries()
            .map(|entry| entry.name.to_string())
            .collect()
    }

    /// Build [`McpClient`]s for registered ACP servers not already connected. Appended to the
    /// init handshake batch so they register tools + land in `owned_clients` on the SAME path
    /// as HTTP/stdio servers.
    ///
    /// `overrides` is the per-server config.toml tuning (keyed by server name), resolved by
    /// the caller per init — kept caller-side so this method stays pure (no file I/O under
    /// the `McpState` lock).
    pub fn build_pending_acp_clients(
        &self,
        overrides: &HashMap<String, McpClientTimeoutOverrides>,
    ) -> Vec<McpClient> {
        let Some(acp) = &self.acp_mcp else {
            return Vec::new();
        };
        self.pending_acp_entries()
            .map(|entry| {
                McpClient::new_acp(
                    entry.name.clone(),
                    entry.server_id.clone(),
                    acp.invoker.clone(),
                    overrides.get(&entry.name),
                    self.meta_config_map.get(&entry.name),
                )
            })
            .collect()
    }

    /// Check if a specific tool is disabled for a server (unqualified tool name).
    pub fn is_tool_disabled(&self, server_name: &str, tool_name: &str) -> bool {
        self.disabled_tools
            .get(server_name)
            .is_some_and(|set| set.contains(tool_name))
    }

    /// Update configs and reset initialization state.
    /// Returns true if the configs actually changed, false if they were identical.
    pub fn update_configs(&mut self, new_configs: Vec<acp::McpServer>) -> bool {
        if mcp_servers_equal(&self.configs, &new_configs) {
            tracing::debug!("MCP configs unchanged, skipping update");
            return false;
        }

        // Clear owned clients only — shared (inherited) clients are untouched.
        self.owned_clients.clear();
        self.mcp_tool_meta.clear();
        self.disabled_tool_registrations.clear();
        self.configs = new_configs;
        self.init_progress.cancel();
        self.auth_required.clear();
        self.generation = self.generation.wrapping_add(1);
        true
    }

    /// Diff-based config update: only tears down servers whose config changed
    /// or were removed, keeps healthy unchanged servers alive.
    ///
    /// Returns `None` if configs are identical (no work needed), or `Some(diff)`
    /// describing which servers to add/remove.
    pub fn update_configs_diff(
        &mut self,
        new_configs: Vec<acp::McpServer>,
    ) -> Option<McpConfigDiff> {
        if mcp_servers_equal(&self.configs, &new_configs) {
            tracing::debug!("MCP configs unchanged, skipping update");
            return None;
        }

        let old_by_name: HashMap<&str, String> = self
            .configs
            .iter()
            .filter_map(|c| match serde_json::to_string(c) {
                Ok(json) => Some((mcp_server_name(c), json)),
                Err(e) => {
                    tracing::warn!(server = mcp_server_name(c), error = %e, "Failed to serialize MCP server config for diff");
                    None
                }
            })
            .collect();

        let new_by_name: HashMap<&str, String> = new_configs
            .iter()
            .filter_map(|c| match serde_json::to_string(c) {
                Ok(json) => Some((mcp_server_name(c), json)),
                Err(e) => {
                    tracing::warn!(server = mcp_server_name(c), error = %e, "Failed to serialize MCP server config for diff");
                    None
                }
            })
            .collect();

        let mut removed = Vec::new();
        let mut added = Vec::new();
        let mut retained = Vec::new();

        for (name, old_json) in &old_by_name {
            match new_by_name.get(name) {
                None => removed.push(name.to_string()),
                Some(new_json) if new_json != old_json => {
                    removed.push(name.to_string());
                }
                Some(_) => retained.push(name.to_string()),
            }
        }

        for (name, new_json) in &new_by_name {
            match old_by_name.get(name) {
                None => added.push(name.to_string()),
                Some(old_json) if old_json != new_json => {
                    added.push(name.to_string());
                }
                Some(_) => {}
            }
        }

        for name in &removed {
            self.owned_clients.remove(name);
            self.auth_required.remove(name);
            self.init_progress.mark_handshake_complete(name);
            let prefix = format!("{}{}", name, MCP_TOOL_NAME_DELIMITER);
            self.mcp_tool_meta.retain(|k, _| !k.starts_with(&prefix));
            self.disabled_tool_registrations
                .retain(|k, _| !k.starts_with(&prefix));
        }

        tracing::info!(
            retained = retained.len(),
            added = added.len(),
            removed = removed.len(),
            "MCP config diff: {} retained, {} added, {} removed",
            retained.len(),
            added.len(),
            removed.len(),
        );

        self.configs = new_configs;
        self.init_progress.cancel();
        self.generation = self.generation.wrapping_add(1);

        Some(McpConfigDiff {
            added,
            removed,
            retained,
        })
    }

    /// Returns `true` only when MCP setup is fully complete: the
    /// init lifecycle reached [`InitProgress::Finished`] AND every
    /// per-server background handshake has settled (success or failure).
    ///
    /// The strict per-server check matters because session actors call
    /// [`Self::finish_init`] **early** (right after spawning processes,
    /// before any handshake completes) so the session isn't blocked on
    /// MCP for non-MCP work. Callers that gate MCP-tool dispatch on
    /// "is MCP actually ready" — e.g. the Blocking-strategy waits in
    /// `prepare_tool_definitions_timed`, `wait_for_mcp_initialized`,
    /// and the tool-dispatch fast path — therefore need the *combined*
    /// check or they'd race the in-flight per-server handshakes and the
    /// first tool call would land inside the
    /// [`ClientState::Initializing`] window.
    ///
    /// Delegates to [`InitProgress::is_complete`]; see that doc for the
    /// full state machine.
    pub fn is_initialized(&self) -> bool {
        self.init_progress.is_complete()
    }

    /// Returns `true` whenever any initialization work is still
    /// outstanding: pre-`finish_init` OR at least one per-server
    /// handshake is still running in the background.
    ///
    /// Pairs with [`Self::is_initialized`]: during the window between
    /// the early [`Self::finish_init`] and the background task draining
    /// the per-server handshaking set, `is_initialized()` is still
    /// `false` (per-server work remains) AND `is_initializing()` is
    /// `true` (so wait-loops keep waiting instead of kicking off a
    /// second init).
    pub fn is_initializing(&self) -> bool {
        self.init_progress.is_in_progress()
    }

    /// Returns `true` once `finish_init` has fired, regardless of
    /// whether per-server background handshakes are still draining.
    /// Used for diagnostic logging where the caller wants to
    /// distinguish "pre-finish window" from "post-finish, bg work
    /// outstanding".
    pub fn has_finished_init(&self) -> bool {
        self.init_progress.has_finished_init()
    }

    /// Borrow the underlying [`InitProgress`] state machine, primarily
    /// for tests that want to assert against the discriminant directly.
    pub fn init_progress(&self) -> &InitProgress {
        &self.init_progress
    }

    /// Try to start initialization. Returns `true` if we transitioned
    /// from [`InitProgress::NotStarted`] to [`InitProgress::Starting`];
    /// returns `false` if init is already in progress or finished.
    pub fn try_start_init(&mut self) -> bool {
        self.init_progress.try_start()
    }

    /// Transition [`InitProgress::Starting`] → [`InitProgress::Finished`],
    /// preserving the per-server handshaking set. Called early (before
    /// per-server handshakes complete) so the session is unblocked for
    /// non-MCP work — `is_initialized()` still returns `false` until
    /// every handshake has reported via [`Self::mark_server_ready`].
    pub fn finish_init(&mut self) {
        self.init_progress.finish();
    }

    /// Cancel initialization back to [`InitProgress::NotStarted`].
    /// Used when generation changed during init (config change races
    /// with active init) and on full reset.
    pub fn cancel_init(&mut self) {
        self.init_progress.cancel();
    }

    /// Add server names to the handshaking set. Call after filtering
    /// `configs_to_start`, before spawning per-server tasks. Only
    /// meaningful in [`InitProgress::Starting`] / [`InitProgress::Finished`];
    /// logs a warning otherwise.
    pub fn mark_servers_initializing(&mut self, names: impl IntoIterator<Item = McpServerName>) {
        let names: Vec<McpServerName> = names.into_iter().collect();
        // A fresh init attempt clears any prior failure for these servers so
        // a server that recovers on retry stops showing as `Unavailable`.
        for name in &names {
            self.init_failed.remove(name);
        }
        self.init_progress.mark_handshaking(names);
    }

    /// Remove a server from the handshaking set (on success or failure
    /// of its handshake). Safe if not present.
    pub fn mark_server_ready(&mut self, name: &str) {
        self.init_progress.mark_handshake_complete(name);
    }

    /// Record a per-server background-init failure for status reporting,
    /// routing it to the correct set so the two stay disjoint.
    ///
    /// `needs_auth` failures are owned by the auth state machine: its recovery
    /// paths (`handle_mcp_auth_trigger`, `retry_auth_required_servers`)
    /// re-handshake and clear `auth_required`. Such servers must therefore NOT
    /// also land in `init_failed`, or a server that successfully authenticates
    /// would stay reported as `Unavailable` with zero tools. Every other
    /// failure (handshake / `tools/list` error or init timeout) goes to
    /// `init_failed` so the server surfaces as `Unavailable`.
    ///
    /// `detail` is a short cause stored for non-auth failures (the value in
    /// [`Self::init_failed`]); ignored for `needs_auth`.
    pub fn record_init_failure(&mut self, name: &str, needs_auth: bool, detail: Option<String>) {
        if needs_auth {
            self.auth_required.insert(name.to_string());
        } else {
            self.init_failed
                .insert(name.to_string(), detail.unwrap_or_default());
        }
    }

    /// Clear a prior init failure for `name` (symmetric with
    /// [`Self::record_init_failure`]). Used by the reactive managed re-auth
    /// path so a server that recovers is no longer reported as `Unavailable`
    /// with a stale non-auth `detail`.
    pub fn clear_init_failed(&mut self, name: &str) {
        self.init_failed.remove(name);
    }

    /// Clear the entire handshaking set in one shot. Used by the
    /// proxy-mode "init complete" path and the bg-handshake completion
    /// path as a defensive sweep after the per-server
    /// [`Self::mark_server_ready`] calls; cheap no-op if already empty.
    pub fn mark_all_servers_ready(&mut self) {
        self.init_progress.clear_handshaking();
    }

    /// True iff the named server's handshake is still in flight.
    /// Used by status snapshots and tool-dispatch gating that need to
    /// know per-server progress without cloning the whole set.
    pub fn is_server_handshaking(&self, name: &str) -> bool {
        self.init_progress.is_server_handshaking(name)
    }

    /// Iterate over server names whose background handshake is still
    /// in flight. Empty when init has not started or has fully
    /// completed.
    pub fn handshaking_servers_iter(&self) -> impl Iterator<Item = &McpServerName> {
        self.init_progress.handshaking_servers()
    }

    /// Snapshot of the handshaking set as a cloned `HashSet`. Used by
    /// the status snapshot API where the caller wants an owned copy
    /// that survives lock release.
    pub fn handshaking_servers_cloned(&self) -> std::collections::HashSet<McpServerName> {
        self.init_progress.handshaking_servers().cloned().collect()
    }

    /// Number of in-flight per-server handshakes.
    pub fn handshaking_servers_count(&self) -> usize {
        self.init_progress.handshaking_count()
    }

    /// Get current generation (for stale check after async init)
    pub fn generation(&self) -> u64 {
        self.generation
    }

    /// Replace managed MCP clients whose URL matches a fresh config entry.
    ///
    /// Caller passes `(endpoint, headers)` pairs from whatever source it uses
    /// (e.g. shell's cli-chat-proxy `ManagedMcpConfig` cache). The MCP crate
    /// stays free of the host's managed-config schema.
    ///
    /// Old `Arc<McpClient>` holders (in-flight tool calls) finish naturally;
    /// new calls look up the fresh client from the map.
    pub fn refresh_managed_clients<'a, I>(&mut self, fresh_configs: I)
    where
        I: IntoIterator<Item = (&'a str, &'a HashMap<String, String>)>,
    {
        let fresh_by_url: HashMap<String, (&'a str, &'a HashMap<String, String>)> = fresh_configs
            .into_iter()
            .map(|(endpoint, headers)| (normalize_url(endpoint), (endpoint, headers)))
            .collect();

        for (client_name, client) in &mut self.owned_clients {
            let Some(client_url) = self.configs.iter().find_map(|cfg| match cfg {
                acp::McpServer::Http(acp::McpServerHttp { name, url, .. })
                | acp::McpServer::Sse(acp::McpServerSse { name, url, .. })
                    if name == client_name =>
                {
                    Some(normalize_url(url))
                }
                _ => None,
            }) else {
                continue;
            };

            let Some(&(fresh_endpoint, fresh_headers)) = fresh_by_url.get(&client_url) else {
                continue;
            };
            if fresh_headers.is_empty() {
                continue;
            }
            // Rebuilding drops the warm connection and forces a full
            // re-handshake on next use; skip it when the token is unchanged.
            if client.http_headers_match(fresh_headers) {
                continue;
            }

            let headers = fresh_headers
                .iter()
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect();
            *client = Arc::new(McpClient::new_http(
                client_name.clone(),
                HttpConfig {
                    url: fresh_endpoint.to_string(),
                    headers,
                },
                None,
                self.meta_config_map.get(client_name.as_str()),
            ));
            tracing::info!(server = %client_name, "Refreshed managed MCP client with fresh token");
        }
    }

    /// Look up a client by server name.
    /// Owned clients take priority (they can override inherited ones).
    pub fn get_client(&self, name: &str) -> Option<&Arc<McpClient>> {
        self.owned_clients
            .get(name)
            .or_else(|| self.shared_clients.get(name))
    }

    /// Iterate over all clients (owned first, then shared — skipping shared
    /// entries whose name is overridden by an owned client).
    pub fn all_clients(&self) -> impl Iterator<Item = (&McpServerName, &Arc<McpClient>)> {
        self.owned_clients.iter().chain(
            self.shared_clients
                .iter()
                .filter(|(name, _)| !self.owned_clients.contains_key(name.as_str())),
        )
    }

    /// Import shared clients from a parent pool snapshot.
    /// Clients whose name collides with an agent-definition-owned server
    /// are skipped (the owned server takes priority).
    pub fn import_shared_clients(&mut self, pool: &SharedMcpPool) {
        let config_names: std::collections::HashSet<&str> =
            self.configs.iter().map(mcp_server_name).collect();
        for (name, client) in &pool.clients {
            if !config_names.contains(name.as_str()) {
                self.shared_clients.insert(name.clone(), Arc::clone(client));
            }
        }
    }
}

/// Snapshot of an MCP connection pool, taken at subagent spawn time.
///
/// The HashMap is cloned (cheap — values are `Arc<McpClient>`), so the
/// subagent's map is independent of the parent's. The `Arc<McpClient>`
/// entries are shared — both parent and child use the same transport.
/// This is intentionally snapshot-based, not live-updating.
#[derive(Clone)]
pub struct SharedMcpPool {
    clients: HashMap<McpServerName, Arc<McpClient>>,
    configs: Vec<acp::McpServer>,
    meta_config_map: McpMetaConfigMap,
}

impl SharedMcpPool {
    /// Create a snapshot from an existing `McpState`.
    /// Captures both owned and shared clients (deduped — owned wins).
    pub fn from_state(state: &McpState) -> Self {
        Self {
            clients: state
                .all_clients()
                .map(|(k, v)| (k.clone(), Arc::clone(v)))
                .collect(),
            configs: state.configs.clone(),
            meta_config_map: state.meta_config_map.clone(),
        }
    }

    pub fn get_client(&self, name: &str) -> Option<&Arc<McpClient>> {
        self.clients.get(name)
    }

    pub fn len(&self) -> usize {
        self.clients.len()
    }

    pub fn is_empty(&self) -> bool {
        self.clients.is_empty()
    }

    pub fn server_names(&self) -> impl Iterator<Item = &str> {
        self.clients.keys().map(String::as_str)
    }

    pub fn configs(&self) -> &[acp::McpServer] {
        &self.configs
    }

    pub fn meta_config_map(&self) -> &McpMetaConfigMap {
        &self.meta_config_map
    }

    /// Retain only clients whose name satisfies `predicate`.
    ///
    /// Only filters the `clients` map. `configs` and `meta_config_map` are
    /// left unchanged — callers that need config-level consistency should
    /// filter those separately. In the subagent inheritance path this is
    /// fine because `import_shared_clients` only iterates `clients`.
    pub fn retain_clients(&mut self, predicate: impl Fn(&str) -> bool) {
        self.clients.retain(|name, _| predicate(name));
    }
}

/// Compare two MCP server config lists for equality.
///
/// Since `acp::McpServer` may not implement PartialEq, we serialize to JSON and compare.
/// This is only called during config updates, so the overhead is acceptable.
pub(crate) fn mcp_servers_equal(a: &[acp::McpServer], b: &[acp::McpServer]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    // Compare JSON serializations
    match (serde_json::to_string(a), serde_json::to_string(b)) {
        (Ok(a_json), Ok(b_json)) => a_json == b_json,
        _ => false, // If serialization fails, assume not equal
    }
}

/// Default timeout for an MCP server's `initialize` handshake & initial tool
/// listing, used when no override is supplied. 30s is generous enough that
/// cold-start `uvx` / `uv run --with` stdio servers that download deps on
/// first launch aren't killed mid-handshake. The shell resolves env / config /
/// requirements / remote overrides and injects them via `McpClientTimeoutOverrides`.
const DEFAULT_STARTUP_TIMEOUT_SECS: u64 = 30;

/// Default timeout for individual tool calls.
const DEFAULT_TOOL_TIMEOUT_SECS: u64 = 6000;

/// How long a stdio server gets to exit after its transport closes before
/// its process group is killed.
const STDIO_SHUTDOWN_GRACE: std::time::Duration = std::time::Duration::from_secs(3);

/// Timeout for OAuth metadata discovery when building an HTTP transport.
/// Bounds transport setup for servers without OAuth support.
const OAUTH_DISCOVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// Per-MCP-server config overrides from `_meta.mcpConfig` in session/new or session/load.
#[derive(Debug, Clone, Default, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpServerMetaConfig {
    /// Init handshake timeout in ms. Overrides config.toml `startup_timeout_sec`.
    #[serde(default)]
    pub startup_timeout_ms: Option<u64>,
    /// Per-tool-call timeout in ms. Overrides config.toml `tool_timeout_sec`.
    #[serde(default)]
    pub tool_timeout_ms: Option<u64>,
    /// Per-tool timeout overrides in ms: `{ "create_issue": 120000, "search": 30000 }`.
    /// Overrides config.toml `tool_timeouts` (and `tool_timeout_sec`) for matching tools.
    #[serde(default)]
    pub tool_timeouts_ms: Option<HashMap<ToolName, u64>>,
    /// Also keep the raw base64 in tool-result text (in addition to the
    /// vision-token rendering) so the agent can decode + forward it via
    /// path-based tools like `send_file`. Costs ~2× tokens per image.
    /// Default `false`. See [`format_mcp_image`].
    #[serde(default)]
    pub expose_image_base64: Option<bool>,
}

/// MCP server name → per-server config overrides from `_meta.mcpConfig`.
pub type McpMetaConfigMap = HashMap<McpServerName, McpServerMetaConfig>;

/// Parse `mcpConfig` from a session request's `_meta`. Empty map if absent/invalid.
pub fn parse_mcp_meta_config(
    meta: Option<&serde_json::Map<String, serde_json::Value>>,
) -> McpMetaConfigMap {
    meta.and_then(|m| m.get("mcpConfig"))
        .and_then(|v| serde_json::from_value::<McpMetaConfigMap>(v.clone()).ok())
        .unwrap_or_default()
}

/// MCP initialization strategy
#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum McpInitStrategy {
    /// Wait for MCP initialization before first LLM call
    #[default]
    Blocking,
    /// Start immediately, advertise tools as they become available
    Progressive,
}

impl<S: AsRef<str>> From<S> for McpInitStrategy {
    fn from(s: S) -> Self {
        match s.as_ref() {
            "progressive" => McpInitStrategy::Progressive,
            _ => McpInitStrategy::Blocking,
        }
    }
}

/// Parse MCP tool name in format "server__tool"
/// Returns (server_name, tool_name) if valid MCP tool, None otherwise
pub fn parse_mcp_tool_name(name: &str) -> Option<(String, String)> {
    let parts: Vec<&str> = name.splitn(2, MCP_TOOL_NAME_DELIMITER).collect();
    if parts.len() == 2 {
        Some((parts[0].to_string(), parts[1].to_string()))
    } else {
        None
    }
}

#[derive(Debug, thiserror::Error)]
pub enum McpError {
    #[error("MCP client error: {0}")]
    ClientError(String),

    #[error("MCP server '{server}' timed out after {timeout_secs}s")]
    Timeout { server: String, timeout_secs: u64 },

    #[error("Failed to spawn MCP server '{server}': {source}")]
    SpawnFailed {
        server: String,
        source: std::io::Error,
    },

    #[error("MCP server '{server}' handshake failed: {source}")]
    HandshakeFailed {
        server: String,
        source: Box<ClientInitializeError>,
    },

    /// Pre-spawn gate: server needs OAuth but this session cannot complete interactive auth.
    #[error(
        "MCP server '{server}': Auth required (non-interactive session; authenticate in TUI or set Authorization header)"
    )]
    AuthRequired { server: String },

    #[error("MCP service error: {0}")]
    ServiceError(#[from] ServiceError),
}

impl McpError {
    fn timeout(server: &str, duration: std::time::Duration) -> Self {
        Self::Timeout {
            server: server.to_string(),
            timeout_secs: duration.as_secs(),
        }
    }

    pub fn is_timeout(&self) -> bool {
        matches!(self, Self::Timeout { .. })
    }

    pub fn error_category(&self) -> kigi_file_utils::events::McpErrorCategory {
        use kigi_file_utils::events::McpErrorCategory;
        match self {
            Self::SpawnFailed { .. } => McpErrorCategory::SpawnFailed,
            Self::Timeout { .. } => McpErrorCategory::Timeout,
            Self::HandshakeFailed { .. } => McpErrorCategory::HandshakeFailed,
            Self::AuthRequired { .. } => McpErrorCategory::AuthRequired,
            Self::ClientError(_) | Self::ServiceError(_) => McpErrorCategory::ClientError,
        }
    }

    pub fn server_name(&self) -> Option<&str> {
        match self {
            Self::SpawnFailed { server, .. }
            | Self::Timeout { server, .. }
            | Self::HandshakeFailed { server, .. }
            | Self::AuthRequired { server } => Some(server),
            Self::ClientError(_) | Self::ServiceError(_) => None,
        }
    }

    /// True if this error indicates the server rejected us for auth reasons (a
    /// credential re-fetch could help). Timeout/spawn failures can't be cured by
    /// re-fetching credentials, so they're never auth.
    pub fn is_auth_rejection(&self) -> bool {
        match self {
            Self::AuthRequired { .. } => true,
            Self::HandshakeFailed { source, .. } => is_auth_rejection_message(&source.to_string()),
            Self::ServiceError(e) => is_auth_rejection_message(&e.to_string()),
            Self::ClientError(s) => is_auth_rejection_message(s),
            Self::Timeout { .. } | Self::SpawnFailed { .. } => false,
        }
    }
}

/// True if an MCP error *message* indicates an auth rejection (vs. a transport
/// drop, timeout, or protocol error), so host recovery can decide whether a
/// credential re-fetch would help.
///
/// Matches auth wording and context-anchored 401 patterns only, so a bare digit
/// ("took 401ms", ports) can't trip it. Excludes 403/forbidden — a non-auth
/// policy denial here, not a credential problem.
pub fn is_auth_rejection_message(s: &str) -> bool {
    let l = s.to_ascii_lowercase();
    // Auth wording has no numeric component, so plain substrings are safe.
    if l.contains("auth required")
        || l.contains("authorizationrequired")
        || l.contains("authrequired")
        || l.contains("authentication")
        || l.contains("unauthorized")
    {
        return true;
    }
    // Require a non-alphanumeric (or end) after "401" so "http 401" matches but
    // "http 4012" (other status) and "http 401ms" (a duration) do not.
    [
        "status: 401",
        "status code 401",
        "http status 401",
        "http 401",
        "error 401",
    ]
    .iter()
    .any(|token| token_at_word_boundary(&l, token))
}

/// Whether `haystack` contains `token` at a right word boundary, so a
/// digit-terminated token (`...401`) doesn't match a longer run (`4012`) or an
/// adjacent unit (`401ms`).
///
/// Invariant: `token` must be ASCII (all callers pass ASCII literals). A
/// non-ASCII token could advance `from` mid-UTF-8 and panic on the next slice.
fn token_at_word_boundary(haystack: &str, token: &str) -> bool {
    debug_assert!(
        token.is_ascii(),
        "token_at_word_boundary requires an ASCII token"
    );
    let mut from = 0;
    while let Some(idx) = haystack[from..].find(token) {
        let end = from + idx + token.len();
        if !haystack[end..].starts_with(|c: char| c.is_ascii_alphanumeric()) {
            return true;
        }
        from += idx + 1;
    }
    false
}

#[derive(Clone)]
pub struct McpTool {
    name: String,
    description: String,
    server_name: String,
    mcp_state: Arc<Mutex<McpState>>,
    schema: serde_json::Value,
    meta: Option<serde_json::Value>,
}

/// Data needed to register an MCP tool via `register_erased()`.
///
/// MCP tools have two visibility audiences controlled by `_meta.ui.visibility`:
///
/// - **Model-visible** (default, or `["model", "app"]`): registered in `ToolBridge`
///   so the LLM can invoke them during a conversation.
/// - **App-visible only** (`["app"]`): not registered in `ToolBridge`, so the LLM
///   never sees them. These are UI-only actions (e.g. refresh buttons) surfaced to
///   the frontend via `kigi/mcp/tools_changed` notifications and callable via
///   `kigi/mcp/call`.
pub struct McpToolRegistration {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
    pub tool: McpErasedTool,
    pub meta: Option<serde_json::Value>,
    pub model_visible: bool,
}

impl McpTool {
    /// Reconstruct an `McpTool` from its constituent parts. Used when stashing
    /// a disabled tool at runtime so it can be re-enabled without a full re-init.
    pub fn new(
        name: String,
        description: String,
        server_name: String,
        mcp_state: Arc<Mutex<McpState>>,
        schema: serde_json::Value,
        meta: Option<serde_json::Value>,
    ) -> Self {
        Self {
            name,
            description,
            server_name,
            mcp_state,
            schema,
            meta,
        }
    }

    /// Convert into the data needed for `ToolBridge::register_erased()`.
    ///
    /// Returns `None` if the tool name is invalid (doesn't match LLM API requirements).
    /// Invalid tools are logged and skipped — fix the upstream connector.
    ///
    /// Also rejects qualified names that contain the delimiter
    /// (`MCP_TOOL_NAME_DELIMITER`) more than once. The underlying tool-name
    /// regex permits underscores in each segment, so a server like
    /// `"foo__bar"`, a tool like `"my__thing"`, or even a `"foo_"`/`"_bar"`
    /// pair (which concatenates to `"foo___bar"` — two valid `__`
    /// positions) would produce a qualified name that downstream
    /// `split_once("__")` consumers would split at the wrong boundary.
    /// The "exactly one delimiter" check covers all three cases with a
    /// single rule.
    pub fn into_registration(self) -> Option<McpToolRegistration> {
        // Qualify MCP tool name with server name: "server__tool"
        let qualified_name = format!(
            "{}{}{}",
            self.server_name, MCP_TOOL_NAME_DELIMITER, self.name
        );

        // Reject ambiguous qualified names — see doc-comment above.
        if qualified_name.matches(MCP_TOOL_NAME_DELIMITER).count() != 1 {
            tracing::error!(
                server = %self.server_name,
                tool = %self.name,
                qualified = %qualified_name,
                "Skipping MCP tool: qualified name contains '{MCP_TOOL_NAME_DELIMITER}' more than once (server, tool, or their boundary collides with the reserved delimiter)"
            );
            return None;
        }

        if let Err(reason) = validate_tool_name(&qualified_name) {
            tracing::error!(
                tool_name = %qualified_name,
                server = %self.server_name,
                reason = %reason,
                "Skipping MCP tool with invalid name"
            );
            return None;
        }

        let description = self.description.clone();
        let input_schema = self.schema.clone();
        let meta = self.meta.clone();

        let model_visible = meta
            .as_ref()
            .and_then(|m| m.get("ui"))
            .and_then(|ui| ui.get("visibility"))
            .and_then(|v| v.as_array())
            .map(|arr| arr.iter().any(|s| s.as_str() == Some("model")))
            .unwrap_or(true); // default: visible to model

        Some(McpToolRegistration {
            name: qualified_name,
            description,
            input_schema,
            tool: McpErasedTool { tool: self },
            meta,
            model_visible,
        })
    }
}

/// MCP tool wrapper for runtime dispatch.
///
/// MCP tools are already untyped (JSON → JSON), so they implement
/// `kigi_tool_runtime::Tool` directly instead of going through typed wrappers.
pub struct McpErasedTool {
    tool: McpTool,
}

impl std::fmt::Debug for McpErasedTool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpErasedTool")
            .field("name", &self.tool.name)
            .field("server", &self.tool.server_name)
            .finish()
    }
}

impl ToolMetadata for McpErasedTool {
    fn kind(&self) -> ToolKind {
        ToolKind::Other
    }

    fn tool_namespace(&self) -> ToolNamespace {
        ToolNamespace::MCP
    }

    fn description_template(&self) -> &str {
        &self.tool.description
    }
}

impl kigi_tool_runtime::Tool for McpErasedTool {
    type Args = serde_json::Value;
    type Output = ToolOutput;

    fn id(&self) -> kigi_tool_protocol::ToolId {
        // Use the qualified name (server__tool) so that two MCP servers
        // exposing the same raw tool name get distinct LocalRegistry entries.
        let qualified = format!(
            "{}{}{}",
            self.tool.server_name, MCP_TOOL_NAME_DELIMITER, self.tool.name
        );
        kigi_tool_protocol::ToolId::new(&qualified)
            .unwrap_or_else(|_| kigi_tool_protocol::ToolId::new("mcp_tool").expect("valid"))
    }

    fn description(
        &self,
        _ctx: &::kigi_tool_runtime::ListToolsContext,
    ) -> kigi_tool_types::ToolDescription {
        kigi_tool_types::ToolDescription::new(&self.tool.name, &self.tool.description)
    }

    async fn run(
        &self,
        _ctx: kigi_tool_runtime::ToolCallContext,
        raw: serde_json::Value,
    ) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
        let mcp_call_start = std::time::Instant::now();
        let (client, event_writer) = {
            let state = self.tool.mcp_state.lock().await;
            let c = Arc::clone(state.get_client(&self.tool.server_name).ok_or_else(|| {
                kigi_tool_runtime::ToolError::custom(
                    "process_manager",
                    format!("MCP server '{}' not found", self.tool.server_name),
                )
            })?);
            (c, state.event_writer().clone())
        };

        let server = &self.tool.server_name;
        let tool = &self.tool.name;
        let tool_timeout = client.tool_timeout_for(tool);
        let qualified_name = format!("{}{}{}", server, MCP_TOOL_NAME_DELIMITER, tool);
        event_writer.emit(kigi_file_utils::events::Event::McpToolCallStarted {
            server_name: server.clone(),
            tool_name: tool.clone(),
            call_id: qualified_name.clone(),
            timeout_sec: tool_timeout,
        });

        let mut auth_retry_attempted = false;
        let mut reconnect_attempted = false;
        let mut is_timeout = false;
        let ew = &event_writer;
        let dispatch_result = match self
            .try_call_tool(&client, &raw, &mut reconnect_attempted, &mut is_timeout, ew)
            .await
        {
            Ok(result) => Ok(result),
            Err(first_err) if client.has_auth() => {
                auth_retry_attempted = true;
                let reauth_ok = client.force_reauth(false).await;
                ew.emit(kigi_file_utils::events::Event::McpAuthRetry {
                    server_name: server.clone(),
                    trigger: "tool_call_failed".to_string(),
                    success: reauth_ok,
                });
                if reauth_ok {
                    self.try_call_tool(&client, &raw, &mut reconnect_attempted, &mut is_timeout, ew)
                        .await
                        .map_err(|e| {
                            kigi_tool_runtime::ToolError::custom("process_manager", e.to_string())
                        })
                } else {
                    Err(first_err)
                }
            }
            Err(e) => Err(e),
        };

        let call_result = match dispatch_result {
            Ok(result) => result,
            Err(e) => {
                ew.emit(kigi_file_utils::events::Event::McpToolCallCompleted {
                    server_name: server.clone(),
                    tool_name: tool.clone(),
                    call_id: qualified_name,
                    duration_ms: mcp_call_start.elapsed().as_millis() as u64,
                    success: false,
                    is_timeout,
                    error: Some(e.to_string()),
                    reconnect_attempted,
                    auth_retry_attempted,
                });
                return Err(e);
            }
        };

        let is_error = call_result.is_error.unwrap_or(false);
        let mut output = if is_error {
            let error_msg = call_result
                .content
                .iter()
                .filter_map(|c| match c {
                    rmcp::model::ContentBlock::Text(t) => Some(t.text.clone()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("\n");
            ToolOutput::MCP(MCPOutput::errored(tool.clone(), server.clone(), error_msg))
        } else {
            let expose_base64 = client.expose_image_base64();
            let parts: Vec<String> = call_result
                .content
                .into_iter()
                .filter_map(|c| match c {
                    rmcp::model::ContentBlock::Text(t) => Some(t.text),
                    rmcp::model::ContentBlock::Image(img) => {
                        Some(format_mcp_image(&img.mime_type, &img.data, expose_base64))
                    }
                    rmcp::model::ContentBlock::Resource(r) => match &r.resource {
                        rmcp::model::ResourceContents::BlobResourceContents {
                            mime_type,
                            blob,
                            ..
                        } if mime_type
                            .as_deref()
                            .is_some_and(|m| m.starts_with("image/")) =>
                        {
                            let mime = mime_type.as_deref().unwrap();
                            Some(format_mcp_image(mime, blob, expose_base64))
                        }
                        _ => serde_json::to_string(&r).ok(),
                    },
                    _ => None,
                })
                .collect();
            let text = parts.join("\n");
            ToolOutput::MCP(MCPOutput::okay_output(tool.clone(), server.clone(), text))
        };

        if let ToolOutput::MCP(ref mut mcp_out) = output {
            mcp_out.auth_retry_attempted = auth_retry_attempted;
            mcp_out.reconnect_attempted = reconnect_attempted;
            mcp_out.is_timeout = is_timeout;
        }

        let success = !is_error;
        let duration_ms = mcp_call_start.elapsed().as_millis() as u64;
        let error_text = if is_error {
            match &output {
                ToolOutput::MCP(mcp) => match mcp.output() {
                    MCPOutputDetails::Error(e) => Some(e.clone()),
                    _ => None,
                },
                _ => None,
            }
        } else {
            None
        };
        event_writer.emit(kigi_file_utils::events::Event::McpToolCallCompleted {
            server_name: server.clone(),
            tool_name: tool.clone(),
            call_id: qualified_name,
            duration_ms,
            success,
            is_timeout,
            error: error_text,
            reconnect_attempted,
            auth_retry_attempted,
        });
        Ok(output)
    }
}

/// Render an MCP image content block. The data URI is consumed by the
/// session-layer `extract_base64_images` and rendered as vision tokens.
/// When `expose_base64`, also emit a `<mcp_image_base64>` wrapper that
/// survives extraction (wrapper has no `data:image/` prefix → regex skips
/// it), exposing the raw bytes to the agent for path-based forwarding.
fn format_mcp_image(mime: &str, base64_data: &str, expose_base64: bool) -> String {
    if expose_base64 {
        format!(
            "data:{mime};base64,{base64_data}\n\
             <mcp_image_base64 mime=\"{mime}\">\n\
             {base64_data}\n\
             </mcp_image_base64>"
        )
    } else {
        format!("data:{mime};base64,{base64_data}")
    }
}

/// Check whether a `ServiceError` indicates the underlying transport has died
/// and a fresh connection could recover it.
fn is_retriable_transport_error(err: &ServiceError) -> bool {
    matches!(
        err,
        ServiceError::TransportClosed | ServiceError::TransportSend(_)
    )
}

/// Recover for every JSON-RPC code except the deterministic client set
/// {-32700, -32600, -32601, -32602} (those mean the request was wrong, not the session).
fn should_recover_mcp_error(code: i32) -> bool {
    use rmcp::model::ErrorCode;
    let deterministic_client_error = code == ErrorCode::PARSE_ERROR.0
        || code == ErrorCode::INVALID_REQUEST.0
        || code == ErrorCode::METHOD_NOT_FOUND.0
        || code == ErrorCode::INVALID_PARAMS.0;
    !deterministic_client_error
}

/// Recovers transport errors, and an HTTP `McpError` once per dispatch except
/// deterministic client codes and auth-class errors — a rebuild reuses stale
/// creds, so auth is routed to the re-auth paths instead.
fn should_recover_service_error(
    err: &ServiceError,
    is_http: bool,
    reconnect_attempted: bool,
) -> bool {
    is_retriable_transport_error(err)
        || matches!(
            err,
            ServiceError::McpError(e)
                if is_http
                    && !reconnect_attempted
                    && should_recover_mcp_error(e.code.0)
                    && !is_auth_rejection_message(e.message.as_ref())
        )
}

impl McpErasedTool {
    async fn try_call_tool(
        &self,
        client: &Arc<McpClient>,
        raw: &serde_json::Value,
        reconnect_attempted: &mut bool,
        is_timeout: &mut bool,
        ew: &kigi_file_utils::events::EventWriter,
    ) -> Result<rmcp::model::CallToolResult, kigi_tool_runtime::ToolError> {
        let mcp_service = client
            .ensure_initialized()
            .await
            .map_err(|e| kigi_tool_runtime::ToolError::custom("process_manager", e.to_string()))?;
        let tool_timeout = client.tool_timeout_for(&self.tool.name);
        let timeout_duration = std::time::Duration::from_secs(tool_timeout);
        let mut params = CallToolRequestParams::new(self.tool.name.clone());
        params.arguments = raw.as_object().cloned();

        let result =
            tokio::time::timeout(timeout_duration, mcp_service.call_tool(params.clone())).await;

        match result {
            Ok(Ok(call_result)) => Ok(call_result),
            Ok(Err(service_err))
                if should_recover_service_error(
                    &service_err,
                    client.is_http(),
                    *reconnect_attempted,
                ) =>
            {
                self.recover_and_retry(
                    client,
                    params,
                    timeout_duration,
                    tool_timeout,
                    service_err,
                    reconnect_attempted,
                    is_timeout,
                    ew,
                )
                .await
            }
            Ok(Err(e)) => Err(kigi_tool_runtime::ToolError::custom(
                "process_manager",
                e.to_string(),
            )),
            Err(_) => {
                *is_timeout = true;
                // Reset for the next call but don't retry — a slow side-effecting tool must not run twice.
                if client.is_http() && !*reconnect_attempted {
                    client.reset_transport().await;
                    *reconnect_attempted = true;
                }
                Err(kigi_tool_runtime::ToolError::custom(
                    "process_manager",
                    format!(
                        "MCP tool '{}' timed out after {} seconds",
                        self.tool.name, tool_timeout
                    ),
                ))
            }
        }
    }

    /// On `recover()` failure surface the original error, else the retry error
    /// (preserves the auth signal managed re-auth reads from the string).
    #[allow(clippy::too_many_arguments)]
    async fn recover_and_retry(
        &self,
        client: &Arc<McpClient>,
        params: CallToolRequestParams,
        timeout_duration: std::time::Duration,
        tool_timeout: u64,
        original_err: ServiceError,
        reconnect_attempted: &mut bool,
        is_timeout: &mut bool,
        ew: &kigi_file_utils::events::EventWriter,
    ) -> Result<rmcp::model::CallToolResult, kigi_tool_runtime::ToolError> {
        *reconnect_attempted = true;
        tracing::warn!(
            server = self.tool.server_name.as_str(),
            tool = self.tool.name.as_str(),
            error = %original_err,
            "MCP transport error, attempting reconnect"
        );
        ew.emit(kigi_file_utils::events::Event::McpTransportError {
            server_name: self.tool.server_name.clone(),
            tool_name: self.tool.name.clone(),
            error: original_err.to_string(),
        });
        let mcp_service = match client.recover().await {
            Ok(service) => {
                ew.emit(kigi_file_utils::events::Event::McpTransportReconnect {
                    server_name: self.tool.server_name.clone(),
                    success: true,
                    error: None,
                });
                service
            }
            Err(e) => {
                ew.emit(kigi_file_utils::events::Event::McpTransportReconnect {
                    server_name: self.tool.server_name.clone(),
                    success: false,
                    error: Some(e.to_string()),
                });
                return Err(kigi_tool_runtime::ToolError::custom(
                    "process_manager",
                    original_err.to_string(),
                ));
            }
        };
        match tokio::time::timeout(timeout_duration, mcp_service.call_tool(params)).await {
            Ok(Ok(call_result)) => Ok(call_result),
            Ok(Err(retry_err)) => Err(kigi_tool_runtime::ToolError::custom(
                "process_manager",
                retry_err.to_string(),
            )),
            Err(_) => {
                *is_timeout = true;
                Err(kigi_tool_runtime::ToolError::custom(
                    "process_manager",
                    format!(
                        "MCP tool '{}' timed out after {} seconds",
                        self.tool.name, tool_timeout
                    ),
                ))
            }
        }
    }
}

/// Whether this session can complete an interactive (browser) OAuth flow.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OauthInteractivity {
    Interactive,
    NonInteractive,
}

impl OauthInteractivity {
    /// Headless/SDK sessions set `non_interactive = true` and cannot complete browser OAuth.
    pub fn from_non_interactive(non_interactive: bool) -> Self {
        if non_interactive {
            Self::NonInteractive
        } else {
            Self::Interactive
        }
    }
}

/// Outcome of probing whether an HTTP/SSE MCP server needs OAuth and whether
/// we have credentials usable without an interactive browser flow.
enum HttpOauthPrep {
    /// Server does not advertise OAuth (or discovery failed conservatively).
    NoOauthSupport,
    /// Ready to connect with an auth manager (stored token works, or interactive deferred auth).
    ManagerReady(Arc<tokio::sync::Mutex<rmcp::transport::auth::AuthorizationManager>>),
    /// OAuth is required but cannot complete in non-interactive mode — do not start unauthenticated.
    NeedsInteractiveLogin,
}

impl HttpOauthPrep {
    /// Inconclusive OAuth probe (manager-create error, discovery error, or timeout):
    /// interactive proceeds as plain HTTP; non-interactive fails closed to avoid rmcp
    /// auth-worker stderr noise.
    fn on_probe_failure(mode: OauthInteractivity) -> Self {
        match mode {
            OauthInteractivity::Interactive => Self::NoOauthSupport,
            OauthInteractivity::NonInteractive => Self::NeedsInteractiveLogin,
        }
    }
}

/// Proactive OAuth discovery per RFC 8414 + 9728.
///
/// Creates an `AuthorizationManager` with our `CredentialStoreAdapter`,
/// discovers server metadata, and loads stored tokens if available.
///
/// With no stored tokens but server OAuth support, behavior splits on `mode`:
/// `Interactive` spawns the browser flow in the background (non-blocking; the
/// first tool call picks up tokens via `force_reauth` → `initialize_from_store`
/// once the user consents), while `NonInteractive` fails closed
/// (`NeedsInteractiveLogin`) rather than start an unauthenticated worker that
/// fatals with `Auth(AuthorizationRequired)` on stderr while the prompt still
/// succeeds.
///
/// Known gap: rmcp `get_access_token` returns stored tokens as-is when expiry
/// metadata is absent, so an expiry-less revoked token can still reach
/// `ManagerReady`.
async fn discover_and_prepare_auth(
    server_name: &str,
    server_url: &str,
    mode: OauthInteractivity,
) -> HttpOauthPrep {
    let Ok(parsed_url) = url::Url::parse(server_url) else {
        return HttpOauthPrep::NoOauthSupport;
    };
    let adapter =
        crate::credentials::McpCredentialStoreAdapter::new(server_name.to_string(), parsed_url);

    let mut manager = match rmcp::transport::auth::AuthorizationManager::new(server_url).await {
        Ok(m) => m,
        Err(e) => {
            tracing::warn!(server = server_name, %e, "Failed to create OAuth manager");
            // Non-interactive: fail closed — unauthenticated HTTP may still fatal in rmcp.
            return HttpOauthPrep::on_probe_failure(mode);
        }
    };
    manager.set_credential_store(adapter);

    if let Ok(true) = manager.initialize_from_store().await {
        // Stored creds may be expired/unrefreshable; in non-interactive mode that
        // still yields rmcp worker fatal AuthorizationRequired on stderr. This probe
        // shares the caller's 5s discovery timeout budget.
        if mode == OauthInteractivity::NonInteractive
            && let Err(e) = manager.get_access_token().await
        {
            tracing::warn!(
                server = server_name,
                error = %e,
                "Skipping OAuth MCP in non-interactive mode (stored credentials unusable); re-authenticate in TUI"
            );
            return HttpOauthPrep::NeedsInteractiveLogin;
        }
        tracing::info!(server = server_name, "Loaded stored OAuth credentials");
        return HttpOauthPrep::ManagerReady(Arc::new(tokio::sync::Mutex::new(manager)));
    }

    match manager.discover_metadata().await {
        Ok(metadata) => {
            manager.set_metadata(metadata);
            if mode == OauthInteractivity::NonInteractive {
                tracing::warn!(
                    server = server_name,
                    "Skipping OAuth MCP in non-interactive mode (no stored tokens); authenticate in TUI or set an Authorization header"
                );
                return HttpOauthPrep::NeedsInteractiveLogin;
            }
            tracing::info!(
                server = server_name,
                "Server supports OAuth but has no stored tokens"
            );
            HttpOauthPrep::ManagerReady(Arc::new(tokio::sync::Mutex::new(manager)))
        }
        Err(rmcp::transport::auth::AuthError::NoAuthorizationSupport) => {
            tracing::debug!(server = server_name, "Server does not support OAuth");
            HttpOauthPrep::NoOauthSupport
        }
        Err(e) => {
            tracing::warn!(server = server_name, %e, "OAuth discovery failed");
            HttpOauthPrep::on_probe_failure(mode)
        }
    }
}

/// Configuration for HTTP MCP server connection.
#[derive(Clone)]
pub struct HttpConfig {
    pub url: String,
    pub headers: Vec<(String, String)>,
}

/// Newline-delimited JSON-RPC stdio transport whose read side survives a
/// single undecodable line.
///
/// Used instead of rmcp's `AsyncRwTransport` for two reasons:
/// - **Wire silence:** a bad line is skipped without replying, whereas rmcp
///   answers shape-mismatched JSON with a -32600 error — a reply an off-spec
///   server could echo back as more invalid input.
/// - **Telemetry:** each skip emits an `McpTransportDecodeError` event (with a
///   truncated sample of the offending line) so the failure is visible in the
///   session trace — rmcp's own tracing is not captured there.
///
/// We read lines ourselves (rather than via `FramedRead` + rmcp's codec) so
/// reading continues after a bad line; only a genuine end-of-stream returns
/// `None`. A stray non-JSON stdout line, a JSON-RPC batch array, or an
/// off-spec response therefore never collapses the transport ("Transport
/// closed" failing every in-flight request — the "connector shows but doesn't
/// work" report).
///
/// Generic over `R`/`W` so it can be unit-tested with in-memory pipes; the
/// production transport binds `ChildStdout`/`ChildStdin`.
struct ResilientRwTransport<R, W>
where
    R: AsyncRead,
    W: AsyncWrite,
{
    read: BufReader<R>,
    /// `Arc<Mutex<Option<…>>>` so `send` can return a `Send + 'static` future
    /// (the `Transport` contract) without borrowing `self`, and so `close` can
    /// drop the writer — mirrors rmcp's own `AsyncRwTransport`.
    write: Arc<Mutex<Option<W>>>,
    server_name: String,
    event_writer: kigi_file_utils::events::EventWriter,
}

/// Max bytes of an offending line copied into the decode-error event.
const DECODE_ERROR_SAMPLE_LEN: usize = 200;

/// A line that failed to deserialize but is a JSON *notification* (an object
/// with a `method` and no `id`) is benign — many servers emit non-MCP / unknown
/// notifications (e.g. LSP-style). Skip those quietly instead of flagging a
/// decode error, mirroring rmcp's compatibility handling.
fn is_ignorable_notification(line: &[u8]) -> bool {
    match serde_json::from_slice::<serde_json::Value>(line) {
        Ok(v) => v.get("id").is_none() && v.get("method").and_then(|m| m.as_str()).is_some(),
        Err(_) => false,
    }
}

impl<R, W> ResilientRwTransport<R, W>
where
    R: AsyncRead + Send + Unpin,
    W: AsyncWrite + Send + Unpin + 'static,
{
    fn new(
        read: R,
        write: W,
        server_name: String,
        event_writer: kigi_file_utils::events::EventWriter,
    ) -> Self {
        Self {
            read: BufReader::new(read),
            write: Arc::new(Mutex::new(Some(write))),
            server_name,
            event_writer,
        }
    }

    /// Record a skipped, undecodable stdout line: a `warn!` log plus an
    /// `McpTransportDecodeError` event carrying the serde error and a truncated
    /// sample of the raw line (the diagnostic the untagged-enum serde error
    /// alone lacks).
    fn record_decode_error(&self, line: &[u8], err: &serde_json::Error) {
        let sample: String = String::from_utf8_lossy(line)
            .chars()
            .take(DECODE_ERROR_SAMPLE_LEN)
            .collect();
        tracing::warn!(
            server = %self.server_name,
            error = %err,
            sample = %sample,
            "Skipping undecodable MCP stdout line; keeping transport alive",
        );
        self.event_writer
            .emit(kigi_file_utils::events::Event::McpTransportDecodeError {
                server_name: self.server_name.clone(),
                error: err.to_string(),
                sample,
            });
    }
}

impl<R, W> Transport<RoleClient> for ResilientRwTransport<R, W>
where
    R: AsyncRead + Send + Unpin,
    W: AsyncWrite + Send + Unpin + 'static,
{
    type Error = std::io::Error;

    fn send(
        &mut self,
        item: TxJsonRpcMessage<RoleClient>,
    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
        let lock = self.write.clone();
        async move {
            let mut bytes = serde_json::to_vec(&item).map_err(std::io::Error::other)?;
            bytes.push(b'\n');
            let mut guard = lock.lock().await;
            match guard.as_mut() {
                Some(write) => {
                    write.write_all(&bytes).await?;
                    write.flush().await
                }
                None => Err(std::io::Error::new(
                    std::io::ErrorKind::NotConnected,
                    "transport is closed",
                )),
            }
        }
    }

    async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
        loop {
            let mut line = Vec::new();
            match self.read.read_until(b'\n', &mut line).await {
                Ok(0) => return None, // genuine end-of-stream
                Ok(_) => {}
                Err(e) => {
                    tracing::debug!(
                        server = %self.server_name,
                        error = %e,
                        "MCP stdio read error; closing transport",
                    );
                    return None;
                }
            }
            if line.last() == Some(&b'\n') {
                line.pop();
            }
            if line.last() == Some(&b'\r') {
                line.pop();
            }
            if line.is_empty() {
                continue;
            }

            match serde_json::from_slice::<RxJsonRpcMessage<RoleClient>>(&line) {
                Ok(msg) => return Some(msg),
                // The whole point: a single undecodable line must not
                // collapse the transport — skip it and keep reading.
                Err(err) => {
                    if is_ignorable_notification(&line) {
                        tracing::trace!(
                            server = %self.server_name,
                            "Ignoring unrecognized MCP notification",
                        );
                    } else {
                        self.record_decode_error(&line, &err);
                    }
                    continue;
                }
            }
        }
    }

    async fn close(&mut self) -> Result<(), Self::Error> {
        let mut guard = self.write.lock().await;
        drop(guard.take());
        Ok(())
    }
}

/// Stdio MCP transport with a non-panicking cleanup path.
///
/// Unlike `rmcp`'s `TokioChildProcess` (which `tokio::spawn`s from `Drop` and so
/// panics when dropped without an entered runtime), this wrapper's `Drop` is
/// best-effort: it reaps via the current runtime if present, else a short-lived
/// cleanup thread, so the child never leaks as a zombie. Since the caller's
/// `detach_command` `setsid`s the child into its own group, teardown also
/// `killpg`s the whole group via [`ProcessGroup`] to avoid orphaning
/// grandchildren (e.g. `npx` -> `node`) before reaping the leader.
pub struct SafeTokioChildProcess {
    child: Option<tokio::process::Child>,
    process_group: Option<ProcessGroup>,
    transport: ResilientRwTransport<tokio::process::ChildStdout, tokio::process::ChildStdin>,
}

impl SafeTokioChildProcess {
    /// `server_name` + `event_writer` are threaded into the transport so a
    /// skipped (undecodable) stdout line emits an `McpTransportDecodeError`
    /// event for that server.
    fn spawn(
        mut cmd: Command,
        server_name: String,
        event_writer: kigi_file_utils::events::EventWriter,
    ) -> std::io::Result<(Self, Option<ChildStderr>)> {
        cmd.stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped());

        let mut child = cmd.spawn()?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| std::io::Error::other("stdin was already taken"))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| std::io::Error::other("stdout was already taken"))?;
        let stderr = child.stderr.take();

        // Best-effort: a missing group just degrades to direct-child-only cleanup.
        let process_group = match ProcessGroup::new() {
            Ok(mut group) => match group.attach(&child) {
                Ok(()) => Some(group),
                Err(e) => {
                    tracing::warn!("Failed to attach MCP child to process group: {e}");
                    None
                }
            },
            Err(e) => {
                tracing::warn!("Failed to create MCP child process group: {e}");
                None
            }
        };

        Ok((
            Self {
                child: Some(child),
                process_group,
                transport: ResilientRwTransport::new(stdout, stdin, server_name, event_writer),
            },
            stderr,
        ))
    }

    fn id(&self) -> Option<u32> {
        self.child.as_ref()?.id()
    }

    /// SIGKILLs the whole process group (child + grandchildren). Synchronous, so
    /// it's safe from `Drop`; the leader still needs reaping afterwards.
    fn kill_process_group(&self) {
        if let Some(group) = &self.process_group
            && let Err(e) = group.kill()
        {
            tracing::warn!("Error killing MCP child process group: {e}");
        }
    }

    async fn graceful_shutdown(&mut self) -> std::io::Result<()> {
        let Some(mut child) = self.child.take() else {
            return Ok(());
        };
        self.transport.close().await?;

        let result = tokio::select! {
            _ = tokio::time::sleep(STDIO_SHUTDOWN_GRACE) => {
                self.kill_process_group();
                match child.kill().await {
                    Ok(()) => Ok(()),
                    Err(e) => {
                        tracing::warn!("Error killing MCP child: {e}");
                        Err(e)
                    }
                }
            }
            res = child.wait() => {
                // Reap any grandchildren now, while the pgid is still kept alive
                // by them and before the reaped leader's pid can be reused.
                self.kill_process_group();
                match res {
                    Ok(status) => {
                        tracing::info!("MCP child exited gracefully {status}");
                        Ok(())
                    }
                    Err(e) => {
                        tracing::warn!("Error waiting for MCP child: {e}");
                        Err(e)
                    }
                }
            }
        };

        // Leader is reaped; drop the group so `Drop` can't `killpg` a now-reusable pid.
        self.process_group = None;
        result
    }
}

impl Drop for SafeTokioChildProcess {
    fn drop(&mut self) {
        // Group teardown is synchronous, so do it first regardless of runtime.
        self.kill_process_group();

        let Some(mut child) = self.child.take() else {
            return;
        };

        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn(async move {
                if let Err(e) = child.kill().await {
                    tracing::warn!("Error killing MCP child process: {e}");
                }
            });
        } else if let Err(e) = std::thread::Builder::new()
            .name("mcp-stdio-child-cleanup".to_string())
            .spawn(move || {
                let rt = match tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                {
                    Ok(rt) => rt,
                    Err(e) => {
                        tracing::warn!("Error creating runtime to clean up MCP child process: {e}");
                        if let Err(e) = child.start_kill() {
                            tracing::warn!("Error signaling MCP child process during drop: {e}");
                        }
                        return;
                    }
                };

                rt.block_on(async move {
                    if let Err(e) = child.kill().await {
                        tracing::warn!("Error killing MCP child process: {e}");
                    }
                });
            })
        {
            tracing::warn!("Error spawning MCP child cleanup thread: {e}");
        }
    }
}

impl Transport<RoleClient> for SafeTokioChildProcess {
    type Error = std::io::Error;

    fn send(
        &mut self,
        item: TxJsonRpcMessage<RoleClient>,
    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
        self.transport.send(item)
    }

    fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<RoleClient>>> + Send {
        self.transport.receive()
    }

    async fn close(&mut self) -> Result<(), Self::Error> {
        self.graceful_shutdown().await
    }
}

/// Transport configuration before connection is established.
enum PendingTransport {
    Stdio(Box<SafeTokioChildProcess>),
    Http(HttpConfig),
    HttpAuth {
        config: HttpConfig,
        auth_manager: Arc<tokio::sync::Mutex<rmcp::transport::auth::AuthorizationManager>>,
    },
    /// In-process SDK MCP server reached over the ACP reverse channel
    /// (`kigi/mcp/sdk_call`). Rebuildable from its `server_id` + invoker, so handshake
    /// failures restore like Http (unlike the consumed Stdio child).
    Acp {
        server_id: String,
        invoker: Arc<dyn crate::acp_transport::AcpReverseInvoker>,
    },
}

/// A connected MCP service (rmcp's RunningService wrapped in Arc).
/// Uses [`KigiClientHandler`] rather than rmcp's default `ClientInfo`
/// handler: rmcp 2.1 parameterizes `RunningService` over the handler
/// type, and `ClientInfo` is only a `ClientHandler` impl with no
/// notification routing. The custom handler keeps the same protocol
/// behavior (same `get_info`) while plumbing
/// `tools/list_changed` / `resources/list_changed` notifications
/// through to the session-actor dispatcher.
pub type McpService = Arc<RunningService<RoleClient, KigiClientHandler>>;

/// MCP client connection state machine.
///
/// Single-flight handshake invariant: at most one task at a time may run
/// the handshake. While the handshake is in flight the state is
/// [`ClientState::Initializing`]; the holder owns the transport for the
/// duration of [`McpClient::try_handshake`]. Concurrent callers of
/// [`McpClient::ensure_initialized`] observe [`ClientState::Initializing`]
/// and park on [`McpClient::init_done`] until the holder publishes a
/// result, instead of failing fast with
/// `"MCP client already initializing"` as in earlier versions.
enum ClientState {
    /// No transport configured. Reachable from:
    /// - [`McpClient::stub`] (test placeholder; `ensure_initialized`
    ///   returns a configuration error).
    /// - Stdio handshake failure (the spawned child process is consumed
    ///   by `client.serve` and cannot be reused — Http/HttpAuth keep
    ///   their `HttpConfig` clone and transition back to `Pending`).
    Empty,
    /// Transport is configured and ready for the next handshake.
    Pending(PendingTransport),
    /// A caller is currently inside [`McpClient::try_handshake`] and owns
    /// the transport. New callers MUST park on
    /// [`McpClient::init_done`] (with a bounded timeout) rather than
    /// attempt a parallel handshake. If the holder is cancelled or
    /// panics before publishing a result, an [`InitGuard`] restores the
    /// transport on a best-effort basis so other callers can retry.
    Initializing,
    /// Handshake completed; the service is reference-counted via `Arc`.
    Ready(McpService),
}

/// `Copy` projection of [`ClientState`] used for cheap state-machine
/// inspection (see [`McpClient::state_kind`]). Mirrors the variants
/// 1:1, dropping the payloads so callers can pattern-match without
/// borrowing the state mutex's inner data.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientStateKind {
    Empty,
    Pending,
    Initializing,
    Ready,
}

/// Classification used by [`crate::liveness::spawn_transport_liveness`].
///
/// Returned by [`McpClient::liveness_check`] under a single state-mutex
/// acquisition, so the watcher's per-tick predicate is atomic.
///
/// `Transient` covers states the watcher should silently exit on
/// (re-handshake races, externally-reset transports, post-failure
/// `Empty` slots). Only `Ready + transport closed` produces a
/// `TransportClosed` ACP push.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LivenessCheck {
    /// `Ready` + `is_transport_closed() == false`. Keep polling.
    Healthy,
    /// `Ready` + `is_transport_closed() == true`. Emit + exit.
    TransportClosed,
    /// Anything else (`Initializing`, `Pending`, `Empty`). The
    /// watcher exits silently — the new state is being managed
    /// externally; if it returns to `Ready` the owner can re-arm.
    Transient,
}

/// Events emitted by a live MCP client to its session-side dispatcher.
///
/// Produced by three sources:
///
/// 1. [`crate::liveness::spawn_transport_liveness`] when an `is_healthy`
///    poll observes that the rmcp service loop has shut down its receiver
///    (`TransportClosed`).
/// 2. [`KigiClientHandler`] when the server pushes a notification we
///    care about — currently `notifications/tools/list_changed` and
///    `notifications/resources/list_changed`.
/// 3. The session/managed-config layer when a server is added, removed,
///    or successfully (re-)initialized.
///
/// Consumers fan these out to ACP `kigi/mcp/server_status` after 50 ms
/// of tumbling-window coalescing keyed by `(server, kind)`; see the
/// session-actor `StatusDispatcher`.
#[derive(Debug, Clone)]
pub enum McpClientEvent {
    /// The rmcp service loop has terminated; the client is no longer
    /// usable for tool calls and must be torn down (or restarted).
    TransportClosed {
        server: McpServerName,
        /// Identity of the client whose transport closed (see
        /// [`McpClient::client_id`]). A mismatch with the client
        /// currently registered under `server` marks the event stale —
        /// it must not tear down the replacement. Every emitter holds the
        /// closing `McpClient`, so the id is always known.
        client_id: u64,
    },
    /// `ensure_initialized` returned `Err(_)`; `reason` is the full
    /// stringified error, surfaced verbatim to the client (no
    /// sanitization) so failures are easy to debug.
    HandshakeFailed {
        server: McpServerName,
        reason: String,
    },
    /// Server pushed `notifications/tools/list_changed`.
    ToolsChanged { server: McpServerName },
    /// Server pushed `notifications/resources/list_changed`.
    ResourcesChanged { server: McpServerName },
    /// Client transitioned to [`ClientState::Ready`]; dispatcher uses
    /// this to surface "ready" status without polling. Emitted from
    /// `ensure_initialized`; the dispatcher maps it to
    /// `reason=initialized` (NOT `reason=restart_succeeded`, which is
    /// reserved for the restart path).
    Ready { server: McpServerName },
    /// Managed/local config diff resolved. The dispatcher fans this
    /// out into one [`Self::ConfigAdded`] / [`Self::ConfigRemoved`]
    /// event per affected server before buffering.
    ConfigDiff {
        added: Vec<McpServerName>,
        removed: Vec<McpServerName>,
    },
    /// Per-server `(server, ConfigAdded)` fan-out variant produced by
    /// the dispatcher from a [`Self::ConfigDiff`]. Keeps the
    /// `kind ↔ event payload` invariant: storing a fake `Ready`
    /// payload at a `ConfigAdded` key would be a footgun whenever a
    /// real `Ready` and a `ConfigDiff` collided in the same coalesce
    /// window.
    ConfigAdded { server: McpServerName },
    /// Per-server `(server, ConfigRemoved)` fan-out variant — the
    /// dispatched analogue of [`Self::ConfigAdded`] for the removed
    /// set of a [`Self::ConfigDiff`].
    ConfigRemoved { server: McpServerName },
}

/// Discriminant for [`McpClientEvent`], used as the second half of the
/// coalescing key `(server, kind)`. Two events with the same
/// `(server, kind)` collapse into the latest one inside the
/// dispatcher's 50 ms window.
///
/// Distinct from [`McpClientEvent`] because the latter carries
/// payload (e.g. `reason` on `HandshakeFailed`) that we don't want
/// participating in equality / hashing.
#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
pub enum McpClientEventKind {
    TransportClosed,
    HandshakeFailed,
    ToolsChanged,
    ResourcesChanged,
    Ready,
    ConfigAdded,
    ConfigRemoved,
}

impl McpClientEvent {
    /// Server name carried by the event, if any.
    ///
    /// Returns `None` only for [`McpClientEvent::ConfigDiff`] — that
    /// variant is fanned out per-server by the dispatcher into
    /// [`Self::ConfigAdded`] / [`Self::ConfigRemoved`], where each
    /// fan-out child has a single server name.
    pub fn server_name(&self) -> Option<&str> {
        match self {
            Self::TransportClosed { server, .. }
            | Self::HandshakeFailed { server, .. }
            | Self::ToolsChanged { server }
            | Self::ResourcesChanged { server }
            | Self::Ready { server }
            | Self::ConfigAdded { server }
            | Self::ConfigRemoved { server } => Some(server.as_str()),
            Self::ConfigDiff { .. } => None,
        }
    }
}

/// RAII guard that restores [`ClientState::Pending`] if the
/// [`McpClient::ensure_initialized`] holder is dropped before publishing
/// its handshake result (task cancellation, panic). Without this guard a
/// cancellation mid-handshake would leave `state` stuck in
/// [`ClientState::Initializing`] and every subsequent caller would block
/// until the wait-timeout fallback fires, then return an error — the
/// caller would have to call [`McpClient::reset_transport`]
/// manually to recover.
///
/// On the success path the holder calls [`Self::disarm`] before storing
/// `Ready`/`Empty` under the state lock, which converts the `Drop` into
/// a no-op. The guard never restores on the success path.
///
/// Drop uses [`tokio::sync::Mutex::try_lock`] because `Drop` runs
/// synchronously and we cannot block the runtime here. If the lock is
/// contended (extremely rare — the only competing locker is another
/// `ensure_initialized` caller which holds the lock for the duration of
/// a match arm, microseconds), the restore is skipped and the
/// inflight-wait timeout in `ensure_initialized` becomes the
/// last-resort recovery path.
struct InitGuard<'a> {
    state: &'a Mutex<ClientState>,
    init_done: &'a Notify,
    /// `Some` until [`Self::disarm`] is called. Holds the restorable
    /// transport (HTTP / HttpAuth) or `None` for Stdio (whose child
    /// process is consumed by `client.serve` and cannot be reused).
    restore: Option<PendingTransport>,
}

impl InitGuard<'_> {
    /// Mark the guard as having published a result. Subsequent `Drop`
    /// becomes a no-op so it doesn't fight with the holder's own
    /// state-store-under-the-lock or wake waiters twice.
    fn disarm(&mut self) {
        self.restore = None;
    }
}

impl Drop for InitGuard<'_> {
    fn drop(&mut self) {
        let Some(restore) = self.restore.take() else {
            return;
        };
        // Best-effort restore. `try_lock` cannot block the runtime from
        // inside Drop; on contention the slot stays Initializing and the
        // inflight-wait timeout becomes the recovery path.
        if let Ok(mut guard) = self.state.try_lock()
            && matches!(&*guard, ClientState::Initializing)
        {
            *guard = ClientState::Pending(restore);
        }
        // Notify whether or not we managed to restore — parked waiters
        // need to wake up and either retry against the restored
        // transport or hit the wait-timeout error path.
        self.init_done.notify_waiters();
    }
}

/// Build a restorable handle for a pending transport, or `None` if the
/// transport cannot be reused after a handshake failure.
///
/// `PendingTransport` deliberately does not implement `Clone`: the
/// `Stdio` variant owns a [`tokio::process::Child`] that is consumed by
/// `client.serve`, and a "restored" Stdio entry would be a dead handle.
/// HTTP and HttpAuth, by contrast, only carry config + an `Arc` to a
/// shared auth manager, so a clone is the canonical way to retry.
fn restorable_transport(pending: &PendingTransport) -> Option<PendingTransport> {
    match pending {
        PendingTransport::Http(cfg) => Some(PendingTransport::Http(cfg.clone())),
        PendingTransport::HttpAuth {
            config,
            auth_manager,
        } => Some(PendingTransport::HttpAuth {
            config: config.clone(),
            auth_manager: auth_manager.clone(),
        }),
        PendingTransport::Acp { server_id, invoker } => Some(PendingTransport::Acp {
            server_id: server_id.clone(),
            invoker: invoker.clone(),
        }),
        PendingTransport::Stdio(_) => None,
    }
}

/// Monotonic source for [`McpClient::client_id`]. Process-global so every
/// client instance — including test stubs — gets a unique identity.
static NEXT_CLIENT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);

fn next_client_id() -> u64 {
    NEXT_CLIENT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
}

pub struct McpClient {
    /// Unique identity of this client *instance*. See [`Self::client_id`].
    client_id: u64,
    server_name: McpServerName,
    state: Mutex<ClientState>,
    /// Wakes [`Self::ensure_initialized`] callers that observed
    /// [`ClientState::Initializing`] and parked. Notified after each
    /// handshake attempt finishes (success **or** failure) and `state`
    /// has been updated. See [`ClientState`] for the single-flight
    /// invariant this preserves.
    ///
    /// Replaces the previous fail-fast
    /// `McpError::ClientError("MCP client already initializing")` branch
    /// which leaked into model-visible tool results whenever the model's
    /// first tool dispatch raced the session actor's background
    /// `get_tool_registrations` handshake.
    init_done: Notify,
    startup_timeout_sec: u64,
    tool_timeout_sec: u64,
    /// Per-tool timeout overrides in seconds. Looked up by tool name;
    /// falls back to `tool_timeout_sec` when a tool isn't listed.
    tool_timeouts: HashMap<ToolName, u64>,
    /// See [`McpServerMetaConfig::expose_image_base64`].
    expose_image_base64: bool,
    /// Shared `AuthorizationManager` for OAuth-enabled servers. `AuthClient`
    /// inside the transport holds a clone of this Arc so token updates are
    /// visible to both the transport and the re-auth path.
    auth_manager: Option<Arc<tokio::sync::Mutex<rmcp::transport::auth::AuthorizationManager>>>,
    /// Stored for OAuth clients so we can rebuild the transport after re-auth.
    http_config: Option<HttpConfig>,
    /// BYO OAuth config for the full browser flow fallback (when refresh fails).
    byo_oauth_config: Option<McpOAuthConfig>,
    /// Rate limit on this server's reconnect warnings; passed to each HTTP
    /// transport so rebuilds keep the limit.
    warn_budget: crate::mcp_http_client::WarnBudget,
    /// The transport to rebuild on a dead connection — see
    /// [`McpClient::reset_transport`]. `None` for transports that can't
    /// reconnect, e.g. Stdio (whose child process is consumed by the
    /// handshake and can't be restarted from here).
    reconnect: Option<PendingTransport>,
    /// Event sink for transport-closed pollers, server-pushed
    /// `tools/list_changed` / `resources/list_changed` notifications,
    /// and handshake failures.
    ///
    /// The slot is `Some` after [`Self::set_event_tx`] is called and
    /// `None` otherwise. The `Arc<Mutex<...>>` is **shared with
    /// [`KigiClientHandler`]** constructed by
    /// [`Self::make_client_handler`]: the handler holds a clone of
    /// the same Arc and reads through it on every notification.
    /// Snapshotting the slot at handshake time instead would mean any
    /// session that wired `notify_tx` post-handshake silently lost
    /// every `tools/list_changed` and `resources/list_changed` for the
    /// life of the connection.
    ///
    /// `None` in three cases:
    /// 1. Test stubs and standalone-pool fixtures that don't need
    ///    cross-component event flow.
    /// 2. Subagent / shared-pool snapshots — only the **parent** session
    ///    is the owner of these events. A subagent that inherits a
    ///    shared `Arc<McpClient>` reads tools through it but does not
    ///    install its own dispatcher; the parent's
    ///    [`crate::liveness::TransportLivenessHandle`] already covers it.
    /// 3. Brand-new clients before the session's per-server task has
    ///    called [`Self::set_event_tx`].
    ///
    /// `parking_lot::Mutex` is sufficient (and lighter than the previous
    /// `tokio::sync::Mutex`): the lock is never held across an
    /// `.await`, and the handler's `emit` path is short and
    /// allocation-free.
    notify_tx: SharedEventTx,
    /// RAII handle for the per-client transport-liveness poller.
    ///
    /// `Some` after [`Self::arm_liveness_watcher`] succeeds; `None`
    /// initially. The slot is also cleared by the poller itself when
    /// it exits (whether on `TransportClosed` or because the state
    /// machine drifted out of `Ready` during a re-handshake — see
    /// [`crate::liveness::spawn_transport_liveness`]) so subsequent
    /// `arm_liveness_watcher` calls aren't silently blocked by a
    /// dead-but-still-present handle.
    ///
    /// `parking_lot::Mutex` is sufficient because the lock is only ever
    /// held for the duration of a slot swap. The poller task uses an
    /// internal `Arc` clone of this mutex (the same memory) so it
    /// can clear the slot before `break`.
    liveness_handle: Arc<parking_lot::Mutex<Option<crate::liveness::TransportLivenessHandle>>>,
}

/// Shared sender slot type — the same Arc lives on the [`McpClient`]
/// and the [`KigiClientHandler`] it constructs during
/// [`McpClient::try_handshake`]. Mutating the slot via
/// [`McpClient::set_event_tx`] is observed by the live rmcp service
/// loop on the next notification, so there's no "snapshot at
/// handshake" hazard.
pub type SharedEventTx =
    Arc<parking_lot::Mutex<Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>>>>;

/// External-config overrides for an MCP server, surfaced to kigi-mcp
/// from whatever loader the host crate uses (e.g. the host's `config.toml` parser).
///
/// All fields are pre-precedence: [`McpClient::load_timeouts`] (and
/// [`McpClient::load_expose_image_base64`]) still apply the
/// `_meta > overrides > default` precedence on top. Owning these types here
/// keeps MCP transport state free of the host's TOML schema.
///
/// Name retained for call-site stability; struct now carries non-timeout
/// config too (e.g. [`Self::expose_image_base64`]).
#[derive(Default, Debug, Clone)]
pub struct McpClientTimeoutOverrides {
    /// Server startup timeout in seconds.
    pub startup_timeout_sec: Option<u64>,
    /// Default per-tool timeout in seconds (used when a tool has no entry in `tool_timeouts`).
    pub tool_timeout_sec: Option<u64>,
    /// Per-tool timeout overrides in seconds, keyed by tool name.
    pub tool_timeouts: Option<HashMap<String, u64>>,
    /// See [`McpServerMetaConfig::expose_image_base64`].
    pub expose_image_base64: Option<bool>,
}

impl McpClient {
    fn load_timeouts(
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> (u64, u64, HashMap<ToolName, u64>) {
        // _meta > overrides > default; env / config / requirements / remote are
        // resolved by the shell and injected via `overrides.startup_timeout_sec`.
        let startup = meta_config
            .and_then(|mc| mc.startup_timeout_ms)
            .map(|ms| ms.div_ceil(1000))
            .or_else(|| overrides.and_then(|o| o.startup_timeout_sec))
            .unwrap_or(DEFAULT_STARTUP_TIMEOUT_SECS);

        let tool = meta_config
            .and_then(|mc| mc.tool_timeout_ms)
            .map(|ms| ms.div_ceil(1000))
            .or_else(|| overrides.and_then(|o| o.tool_timeout_sec))
            .unwrap_or(DEFAULT_TOOL_TIMEOUT_SECS);

        // Per-tool overrides: external base, _meta overrides on top.
        // Precedence: _meta per-tool > overrides per-tool > (falls back to server-level `tool`)
        let mut tool_timeouts = HashMap::new();

        // Layer 1: overrides (already in seconds)
        if let Some(o) = overrides
            && let Some(ref tt) = o.tool_timeouts
        {
            tool_timeouts.extend(tt.iter().map(|(k, v)| (k.clone(), *v)));
        }

        // Layer 2: _meta tool_timeouts_ms (milliseconds → seconds), overrides external config
        if let Some(mc) = meta_config
            && let Some(ref tt) = mc.tool_timeouts_ms
        {
            for (k, v) in tt {
                tool_timeouts.insert(k.clone(), v.div_ceil(1000));
            }
        }

        (startup, tool, tool_timeouts)
    }

    /// `_meta > overrides > default(false)`, mirroring [`Self::load_timeouts`].
    fn load_expose_image_base64(
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> bool {
        meta_config
            .and_then(|mc| mc.expose_image_base64)
            .or_else(|| overrides.and_then(|o| o.expose_image_base64))
            .unwrap_or(false)
    }

    /// The ONLY place that writes the `McpClient { .. }` struct literal.
    /// Every constructor funnels through here so adding a field touches one
    /// site. `reconnect` is snapshotted from the transport before it is
    /// moved into [`ClientState::Pending`] (`None` for non-reconnectable
    /// transports like Stdio — see [`restorable_transport`]).
    #[allow(clippy::too_many_arguments)]
    fn new_with_transport(
        server_name: String,
        transport: PendingTransport,
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
        auth_manager: Option<Arc<tokio::sync::Mutex<rmcp::transport::auth::AuthorizationManager>>>,
        http_config: Option<HttpConfig>,
        byo_oauth_config: Option<McpOAuthConfig>,
    ) -> Self {
        let reconnect = restorable_transport(&transport);
        let (startup_timeout_sec, tool_timeout_sec, tool_timeouts) =
            Self::load_timeouts(overrides, meta_config);
        let expose_image_base64 = Self::load_expose_image_base64(overrides, meta_config);
        Self {
            client_id: next_client_id(),
            server_name,
            state: Mutex::new(ClientState::Pending(transport)),
            init_done: Notify::new(),
            startup_timeout_sec,
            tool_timeout_sec,
            tool_timeouts,
            expose_image_base64,
            auth_manager,
            http_config,
            byo_oauth_config,
            warn_budget: crate::mcp_http_client::WarnBudget::default(),
            reconnect,
            notify_tx: Arc::new(parking_lot::Mutex::new(None)),
            liveness_handle: Arc::new(parking_lot::Mutex::new(None)),
        }
    }

    pub fn new_http_auth(
        server_name: String,
        config: HttpConfig,
        auth_manager: Arc<tokio::sync::Mutex<rmcp::transport::auth::AuthorizationManager>>,
        byo_oauth_config: Option<McpOAuthConfig>,
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> Self {
        Self::new_with_transport(
            server_name,
            PendingTransport::HttpAuth {
                config: config.clone(),
                auth_manager: auth_manager.clone(),
            },
            overrides,
            meta_config,
            Some(auth_manager),
            Some(config),
            byo_oauth_config,
        )
    }

    pub fn has_auth(&self) -> bool {
        self.auth_manager.is_some()
    }

    /// Try to recover tokens from disk or via refresh — no browser flow.
    ///
    /// Returns true if valid tokens were found (from another session/process
    /// writing to the credential store, or a successful token refresh).
    /// Used by `retry_auth_required_servers` on overlay refresh.
    pub async fn try_reauth_from_disk(&self) -> bool {
        let (Some(auth_mgr), Some(config)) = (&self.auth_manager, &self.http_config) else {
            return false;
        };

        // Token-changed gate: rmcp's `initialize_from_store` returns Ok(true)
        // for any disk-resident creds regardless of expiry, so without
        // comparing against the in-memory token we'd claim "fresh tokens from
        // disk" on the same stale token that triggered this retry. The
        // downstream handshake would catch it, but the log line would lie
        // during incident debugging — and the divergence from `force_reauth`'s
        // gate is the exact invariant drift we just fixed there.
        {
            use oauth2::TokenResponse as _;
            let mut mgr = auth_mgr.lock().await;
            let token_before = mgr
                .get_credentials()
                .await
                .ok()
                .and_then(|(_, tok)| tok)
                .map(|t| t.access_token().secret().to_string());
            if let Ok(true) = mgr.initialize_from_store().await {
                let token_after = mgr
                    .get_credentials()
                    .await
                    .ok()
                    .and_then(|(_, tok)| tok)
                    .map(|t| t.access_token().secret().to_string());
                if token_after.is_some() && token_after != token_before {
                    tracing::info!(
                        server = self.server_name.as_str(),
                        "Loaded fresh tokens from disk"
                    );
                    drop(mgr);
                    self.replace_state(ClientState::Pending(PendingTransport::HttpAuth {
                        config: config.clone(),
                        auth_manager: auth_mgr.clone(),
                    }))
                    .await;
                    return true;
                }
            }
        }

        let refresh_ok = {
            let mgr = auth_mgr.lock().await;
            mgr.refresh_token().await.is_ok()
        };

        if refresh_ok {
            self.replace_state(ClientState::Pending(PendingTransport::HttpAuth {
                config: config.clone(),
                auth_manager: auth_mgr.clone(),
            }))
            .await;
            return true;
        }

        false
    }

    /// Force token acquisition and reset the transport so the next
    /// `ensure_initialized` rebuilds it with the fresh token.
    ///
    /// Tries in order:
    /// 1. Reload from disk (picks up tokens from background auth task)
    /// 2. Refresh via refresh_token grant
    /// 3. Full browser-based OAuth flow
    pub async fn force_reauth(&self, force: bool) -> bool {
        let (Some(auth_mgr), Some(config)) = (&self.auth_manager, &self.http_config) else {
            return false;
        };

        // Check if another process/session wrote *fresh* tokens to disk. We
        // must compare against the token we already had in memory — rmcp's
        // `initialize_from_store` returns Ok(true) for any disk-resident
        // credentials regardless of expiry, so without the token-changed
        // check we'd short-circuit on the same stale token that triggered
        // this re-auth in the first place (real bug: pressing the auth
        // shortcut on a server with an expired bearer + no refresh_token
        // would no-op and then 401 on the next handshake).
        //
        // Hold a single lock guard across `token_before` → `initialize_from_store`
        // → `token_after` so the comparison's invariant ("snapshot, reload,
        // re-read") can't be torn by an interleaved mutation.
        {
            use oauth2::TokenResponse as _;
            let mut mgr = auth_mgr.lock().await;
            let token_before = mgr
                .get_credentials()
                .await
                .ok()
                .and_then(|(_, tok)| tok)
                .map(|t| t.access_token().secret().to_string());
            if let Ok(true) = mgr.initialize_from_store().await {
                let token_after = mgr
                    .get_credentials()
                    .await
                    .ok()
                    .and_then(|(_, tok)| tok)
                    .map(|t| t.access_token().secret().to_string());
                if token_after.is_some() && token_after != token_before {
                    tracing::info!(
                        server = self.server_name.as_str(),
                        "Loaded fresh tokens from disk (background auth or other process)"
                    );
                    drop(mgr);
                    self.replace_state(ClientState::Pending(PendingTransport::HttpAuth {
                        config: config.clone(),
                        auth_manager: auth_mgr.clone(),
                    }))
                    .await;
                    return true;
                }
            }
        }

        // Try token refresh.
        let refresh_ok = {
            let mgr = auth_mgr.lock().await;
            mgr.refresh_token().await.is_ok()
        };

        if refresh_ok {
            tracing::info!(
                server = self.server_name.as_str(),
                "Token refreshed successfully (no browser)"
            );
            self.replace_state(ClientState::Pending(PendingTransport::HttpAuth {
                config: config.clone(),
                auth_manager: auth_mgr.clone(),
            }))
            .await;
            return true;
        }

        // Full browser-based OAuth flow.
        {
            tracing::info!(
                server = self.server_name.as_str(),
                "Falling back to browser auth"
            );
            if let Err(e) = crate::oauth::authenticate_mcp_server_dedup(
                &self.server_name,
                &config.url,
                auth_mgr,
                self.byo_oauth_config.as_ref(),
                force,
            )
            .await
            {
                tracing::warn!(
                    server = self.server_name.as_str(),
                    %e,
                    "Full re-authentication failed"
                );
                return false;
            }
        }

        self.replace_state(ClientState::Pending(PendingTransport::HttpAuth {
            config: config.clone(),
            auth_manager: auth_mgr.clone(),
        }))
        .await;
        true
    }

    /// Reset the transport so the next `ensure_initialized` rebuilds it with a
    /// fresh connection.
    ///
    /// Called when a tool call fails with a transport error (`TransportClosed`,
    /// `TransportSend`) — the underlying connection is dead but the server's
    /// addressing (URL/headers for HTTP, `server_id`/invoker for ACP) is still
    /// valid.
    ///
    /// Returns `true` if the transport was reset: HTTP/HttpAuth/ACP rebuild
    /// from the `reconnect` snapshot taken at construction. Returns `false`
    /// for clients whose `reconnect` is `None` (e.g. Stdio — dead child
    /// processes can't be restarted from here).
    async fn reset_transport(&self) -> bool {
        let Some(t) = self.reconnect.as_ref().and_then(restorable_transport) else {
            return false;
        };
        self.replace_state(ClientState::Pending(t)).await;
        tracing::info!(
            server = %self.server_name,
            "Reset transport for reconnect after transport failure"
        );
        true
    }

    /// `true` if this client has an HTTP/SSE transport. The explicit predicate
    /// for recovery gates (the proactive path only recovers HTTP clients).
    pub fn is_http(&self) -> bool {
        self.http_config.is_some()
    }

    /// `true` for an in-process SDK client reached over the ACP reverse channel
    /// (rather than HTTP/stdio). Gates liveness watching — see
    /// [`Self::arm_liveness_watcher`].
    pub fn is_acp(&self) -> bool {
        matches!(self.reconnect, Some(PendingTransport::Acp { .. }))
    }

    /// Read-only: do `headers` equal this client's current HTTP transport
    /// headers? Compares the full set order-insensitively (the caller's
    /// headers originate from a `HashMap`). Returns `false` for a client
    /// with no HTTP config.
    pub fn http_headers_match(&self, headers: &HashMap<String, String>) -> bool {
        let Some(config) = &self.http_config else {
            return false;
        };
        // Materialize into a map so a duplicate stored key collapses to one
        // entry, keeping the length comparison honest. HTTP header names are
        // case-insensitive, so normalize names to lowercase on both sides (the
        // crate already does this for `authorization`) and avoid a needless
        // rebuild on a pure casing difference. Values stay case-sensitive.
        let stored: HashMap<String, &str> = config
            .headers
            .iter()
            .map(|(k, v)| (k.to_ascii_lowercase(), v.as_str()))
            .collect();
        stored.len() == headers.len()
            && headers
                .iter()
                .all(|(k, v)| stored.get(&k.to_ascii_lowercase()) == Some(&v.as_str()))
    }

    /// Recover a dead transport in place: reset → re-handshake → re-arm the
    /// liveness watcher. Returns the live [`McpService`].
    ///
    /// The single recovery path for both the proactive HTTP recovery
    /// (`SessionActor::reset_http_client`, gated on [`Self::is_http`]) and the
    /// lazy `try_call_tool` retry. Rebuilds from the `reconnect` snapshot, so it
    /// covers HTTP/HttpAuth/ACP; `arm_liveness_watcher` self-gates for ACP.
    ///
    /// `Err` for a client with no restorable transport (e.g. Stdio — its child
    /// was consumed by the handshake).
    pub async fn recover(self: &Arc<Self>) -> Result<McpService, McpError> {
        // Coalesce concurrent recoveries: reset only when Ready; if already
        // non-Ready a recovery is in flight, so join its single-flight
        // ensure_initialized instead of racing a reset.
        if matches!(self.state_kind().await, ClientStateKind::Ready)
            && !self.reset_transport().await
        {
            return Err(McpError::ClientError(format!(
                "MCP client {} has no transport to recover",
                self.server_name,
            )));
        }
        let service = self.ensure_initialized().await?;
        // Re-arm liveness so the next close is detected again. A `false` return
        // with a wired sender is unexpected only for watched (non-ACP) clients.
        if !self
            .arm_liveness_watcher(crate::liveness::DEFAULT_POLL_INTERVAL)
            .await
            && !self.is_acp()
            && self.event_tx_clone().is_some()
        {
            tracing::warn!(
                server = %self.server_name,
                "recovery: liveness watcher not re-armed despite a wired event sender",
            );
        }
        Ok(service)
    }

    /// Replace [`Self::state`] under the lock and wake any
    /// [`Self::ensure_initialized`] callers parked on [`Self::init_done`]
    /// so they re-check the new state on their next loop iteration.
    ///
    /// Use this for *external* state transitions that need to invalidate
    /// in-flight waits (re-auth completions, transport resets) so a parked
    /// waiter doesn't sit on a stale [`ClientState::Initializing`] view of
    /// the world. `ensure_initialized` itself doesn't go through this
    /// helper because it already mints the new state under the lock it's
    /// holding and notifies waiters once at the end of the handshake
    /// attempt.
    async fn replace_state(&self, new_state: ClientState) {
        {
            let mut guard = self.state.lock().await;
            *guard = new_state;
        }
        self.init_done.notify_waiters();
    }

    pub fn new_stdio(
        server_name: String,
        transport: SafeTokioChildProcess,
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> Self {
        Self::new_with_transport(
            server_name,
            PendingTransport::Stdio(Box::new(transport)),
            overrides,
            meta_config,
            None,
            None,
            None,
        )
    }

    /// Build a client for an in-process SDK MCP server reached over the ACP reverse
    /// channel. `server_id` is the id the agent echoes back in `kigi/mcp/sdk_call`; the
    /// `invoker` performs the reverse request. Same downstream path as HTTP/stdio.
    pub fn new_acp(
        server_name: String,
        server_id: String,
        invoker: Arc<dyn crate::acp_transport::AcpReverseInvoker>,
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> Self {
        Self::new_with_transport(
            server_name,
            PendingTransport::Acp { server_id, invoker },
            overrides,
            meta_config,
            None,
            None,
            None,
        )
    }

    pub fn new_http(
        server_name: String,
        config: HttpConfig,
        overrides: Option<&McpClientTimeoutOverrides>,
        meta_config: Option<&McpServerMetaConfig>,
    ) -> Self {
        Self::new_with_transport(
            server_name,
            PendingTransport::Http(config.clone()),
            overrides,
            meta_config,
            None,
            Some(config),
            None,
        )
    }

    pub fn server_name(&self) -> &str {
        &self.server_name
    }

    /// Unique identity of this client *instance*. Two clients for the
    /// same server name (e.g. a dead client and its replacement after
    /// a config remove+re-add) have different ids. Carried on
    /// [`McpClientEvent::TransportClosed`] so consumers can tell a
    /// death event for the current client from a stale predecessor's.
    pub fn client_id(&self) -> u64 {
        self.client_id
    }

    pub fn startup_timeout_sec(&self) -> u64 {
        self.startup_timeout_sec
    }

    pub fn tool_timeout_sec(&self) -> u64 {
        self.tool_timeout_sec
    }

    /// See [`McpServerMetaConfig::expose_image_base64`].
    pub fn expose_image_base64(&self) -> bool {
        self.expose_image_base64
    }

    /// Resolve the timeout for a specific tool.
    ///
    /// Precedence (highest → lowest):
    /// 1. `_meta.mcpConfig.<server>.toolTimeoutsMs.<tool>`
    /// 2. `config.toml [mcp_servers.<server>].tool_timeouts.<tool>`
    /// 3. `_meta.mcpConfig.<server>.toolTimeoutMs`
    /// 4. `config.toml [mcp_servers.<server>].tool_timeout_sec`
    /// 5. Default (60s)
    ///
    /// Steps 1–2 are already merged into `self.tool_timeouts` at construction;
    /// steps 3–5 are already resolved into `self.tool_timeout_sec`.
    pub fn tool_timeout_for(&self, tool_name: &str) -> u64 {
        self.tool_timeouts
            .get(tool_name)
            .copied()
            .unwrap_or(self.tool_timeout_sec)
    }

    /// Drive the MCP handshake to completion (or return the cached
    /// service if one is already established), with single-flight
    /// semantics that are safe under arbitrary concurrent callers.
    ///
    /// ## Concurrency contract
    ///
    /// At most one task at a time runs [`Self::try_handshake`]; that task
    /// holds the transport and observes [`ClientState::Initializing`].
    /// Other concurrent callers park on [`Self::init_done`] (with a
    /// bounded timeout) instead of issuing parallel handshakes or
    /// failing immediately. When the holder publishes a result, all
    /// parked waiters re-check `state` and either:
    ///
    /// - return the freshly-stored [`McpService`] (handshake succeeded),
    /// - take ownership of the freshly-restored transport and run their
    ///   own handshake (handshake failed but transport is restorable),
    /// - or surface the error (Stdio handshake failed → no restorable
    ///   transport → [`ClientState::Empty`]).
    ///
    /// This replaces the pre-fix behavior where concurrent callers got
    /// an immediate `McpError::ClientError("MCP client already
    /// initializing")` — surfaced inside model-visible tool results
    /// whenever the model's first tool call landed inside the session
    /// actor's background `get_tool_registrations` handshake, causing
    /// repeated retries that exhausted prompt budgets without ever
    /// reaching the actual MCP server.
    ///
    /// ## Cancellation safety
    ///
    /// If the holder is dropped (parent task cancelled, panic) before
    /// publishing a result, [`InitGuard`]'s `Drop` impl best-effort
    /// restores the transport (so future callers can retry without an
    /// explicit `reset_transport`) and wakes parked waiters. The
    /// restore uses [`tokio::sync::Mutex::try_lock`] because `Drop` is
    /// synchronous; on the rare contention case the slot stays
    /// `Initializing` and the wait-timeout fallback below surfaces a
    /// clear error rather than blocking forever.
    pub async fn ensure_initialized(&self) -> Result<McpService, McpError> {
        // Bound how long a parked caller waits on `init_done` before
        // surfacing an error. `try_handshake` is itself bounded by
        // `startup_timeout_sec`, so anything beyond that plus a 1 s margin
        // means the holder was dropped without restoring the transport
        // (cancellation under heavy contention) — wedging silently would
        // turn this into the exact "stuck client" failure mode the rest of
        // this rewrite is designed to eliminate.
        let inflight_wait =
            std::time::Duration::from_secs(self.startup_timeout_sec.saturating_add(1));

        // Drive the loop body until we either return directly or break
        // out with an owned `PendingTransport`. We deliberately use a
        // labelled `loop` with a `break <expr>` so the compiler proves
        // every arm of the inner match either diverges (return /
        // continue) or yields the transport — no `unreachable!()`
        // escape hatch needed.
        let pending: PendingTransport = loop {
            // Subscribe to `init_done` BEFORE inspecting `state` so a
            // wake-up fired between our state check and the wait can't be
            // lost. `tokio::sync::Notify` only delivers a permit to
            // notify-futures that exist at the time of `notify_waiters`.
            let notified = self.init_done.notified();
            tokio::pin!(notified);

            let mut guard = self.state.lock().await;
            // Swap the current state for `Initializing` up front and
            // match on the OWNED previous value. This avoids the
            // `match-by-ref → mem::replace → re-match → unreachable!()`
            // dance — the compiler can bind `ClientState::Pending(t)`
            // directly from an owned value with no irrefutable-let
            // hole. Non-Pending arms restore their original variant
            // before falling through; the lock is held the entire
            // window so the brief `Initializing` placeholder is
            // invisible to other callers. Cost is one trivial unit-
            // variant write per non-Pending call (plus an `Arc::clone`
            // on the Ready path) — negligible.
            match std::mem::replace(&mut *guard, ClientState::Initializing) {
                ClientState::Ready(service) => {
                    *guard = ClientState::Ready(service.clone());
                    return Ok(service);
                }
                ClientState::Empty => {
                    *guard = ClientState::Empty;
                    return Err(McpError::ClientError(format!(
                        "MCP client {} has no transport configured",
                        self.server_name,
                    )));
                }
                ClientState::Initializing => {
                    // Another caller already owns the slot. Restore the
                    // placeholder we just swapped in (semantically a
                    // no-op since `Initializing` is a unit variant),
                    // drop the lock, park on `init_done`.
                    *guard = ClientState::Initializing;
                    drop(guard);
                    match tokio::time::timeout(inflight_wait, notified.as_mut()).await {
                        Ok(()) => continue,
                        Err(_) => {
                            return Err(McpError::ClientError(format!(
                                "MCP client {} init still in progress after {}s",
                                self.server_name,
                                inflight_wait.as_secs(),
                            )));
                        }
                    }
                }
                // The single arm that KEEPS the `Initializing`
                // placeholder we swapped in — this caller becomes the
                // single-flight handshake holder for the duration of
                // `try_handshake` below.
                ClientState::Pending(transport) => break transport,
            }
        };

        // Lock released. Run the handshake outside the lock so other
        // callers can park on `init_done` instead of stalling on
        // `state.lock()`.

        // Clone the transport's restorable handle twice — once for
        // the failure-path retry below, once for the drop guard.
        // `PendingTransport` is intentionally not `Clone` (Stdio's
        // `TokioChildProcess` is unique), so the helper returns
        // `None` for Stdio (whose handshake failures cannot be
        // recovered without a fresh spawn).
        let restore = restorable_transport(&pending);
        let restore_for_guard = restorable_transport(&pending);

        // Drop guard: if `try_handshake` panics or is cancelled before
        // we publish a result, restore `Pending(restore)` so other
        // callers don't stall on `Initializing` forever. Disarm via
        // `disarm()` immediately before storing the real result.
        let mut init_guard = InitGuard {
            state: &self.state,
            init_done: &self.init_done,
            restore: restore_for_guard,
        };

        let handshake_start = std::time::Instant::now();
        let mut result = self.try_handshake(pending).await;

        let handshake_elapsed = handshake_start.elapsed().as_micros() as u64;
        tracing::info!(target: kigi_log::instrumentation::TARGET, event = "timing", name = "mcp_try_handshake", elapsed_us = handshake_elapsed);
        // On handshake failure, if we have an auth_manager, try
        // refreshing the token and retrying once. Handles expired
        // access tokens loaded from disk — the handshake fails at the
        // transport layer before rmcp's transparent 401 refresh can
        // kick in. We attempt refresh on any failure (not just auth
        // errors) because the cost is low and error strings from
        // different MCP servers are not reliable to match.
        if result.is_err()
            && let (Some(auth_mgr), Some(config)) = (&self.auth_manager, &self.http_config)
        {
            tracing::info!(
                server = %self.server_name,
                "Handshake failed, attempting token refresh and retry"
            );
            let refresh_ok = {
                let mgr = auth_mgr.lock().await;
                mgr.refresh_token().await.is_ok()
            };
            if refresh_ok {
                let retry_transport = PendingTransport::HttpAuth {
                    config: config.clone(),
                    auth_manager: auth_mgr.clone(),
                };
                result = self.try_handshake(retry_transport).await;
            }
        }

        // Disarm before publishing the result so the drop guard
        // doesn't double-restore on the success path or fight with
        // the failure-path assignment below.
        init_guard.disarm();

        // Snapshot the event sender before we commit Ready/Pending/Empty
        // under the lock. We want to emit `HandshakeFailed` (on `Err`) or
        // signal the dispatcher to set status=ready (on `Ok`) AFTER
        // releasing the state lock, so a `state.lock().await` inside the
        // dispatcher (should one ever exist — none today) can't deadlock.
        //
        // The snapshot reads through the SHARED `Arc<Mutex<...>>`
        // slot. If the per-server task wired [`Self::set_event_tx`]
        // BEFORE invoking `get_tool_registrations` (the pattern in
        // `acp_session.rs`), this snapshot picks up the sender even
        // for the very first handshake.
        let event_tx = self.event_tx_clone();

        let outcome = {
            let mut guard = self.state.lock().await;
            match result {
                Ok(service) => {
                    let service = Arc::new(service);
                    *guard = ClientState::Ready(service.clone());
                    tracing::info!(
                        server = %self.server_name,
                        "MCP server initialized successfully"
                    );
                    Ok(service)
                }
                Err(e) => {
                    *guard = match restore {
                        Some(transport) => ClientState::Pending(transport),
                        None => ClientState::Empty,
                    };
                    tracing::warn!(
                        server = %self.server_name,
                        error = %e,
                        "MCP server init failed"
                    );
                    Err(e)
                }
            }
        };
        // Wake parked callers AFTER releasing the state lock so they
        // observe the freshly-stored Ready/Pending/Empty value.
        self.init_done.notify_waiters();

        // Notify the session-actor StatusDispatcher of the handshake
        // outcome, AFTER releasing the state lock. Best-effort: if the
        // receiver is gone (dispatcher torn down, subagent without
        // wiring) the send fails silently. The dispatcher is the only
        // path that turns these into ACP pushes — see the
        // `client_event_tx` field on `McpState`.
        if let Some(tx) = &event_tx {
            match &outcome {
                Ok(_) => {
                    let _ = tx.send(McpClientEvent::Ready {
                        server: self.server_name.clone(),
                    });
                }
                Err(e) => {
                    let _ = tx.send(McpClientEvent::HandshakeFailed {
                        server: self.server_name.clone(),
                        reason: e.to_string(),
                    });
                }
            }
        }
        outcome
    }

    /// Run the MCP handshake (no lock held).
    async fn try_handshake(
        &self,
        pending: PendingTransport,
    ) -> Result<rmcp::service::RunningService<RoleClient, KigiClientHandler>, McpError> {
        let timeout = std::time::Duration::from_secs(self.startup_timeout_sec);
        let name = &self.server_name;

        match pending {
            PendingTransport::Stdio(process) => {
                let handler = self.make_client_handler();
                tokio::time::timeout(timeout, handler.serve(*process))
                    .await
                    .map_err(|_| McpError::timeout(name, timeout))?
                    .map_err(|e| McpError::HandshakeFailed {
                        server: name.to_string(),
                        source: Box::new(e),
                    })
            }
            PendingTransport::Http(config) => {
                let transport =
                    Self::build_http_transport(&config, name, self.warn_budget.clone())?;
                let handler = self.make_client_handler();
                tokio::time::timeout(timeout, handler.serve(transport))
                    .await
                    .map_err(|_| McpError::timeout(name, timeout))?
                    .map_err(|e| McpError::HandshakeFailed {
                        server: name.to_string(),
                        source: Box::new(e),
                    })
            }
            PendingTransport::HttpAuth {
                config,
                auth_manager,
            } => {
                let mut headers = reqwest::header::HeaderMap::new();
                for (key, value) in &config.headers {
                    if key.eq_ignore_ascii_case("Authorization") {
                        continue;
                    }
                    if let (Ok(n), Ok(v)) = (
                        reqwest::header::HeaderName::from_bytes(key.as_bytes()),
                        value.parse::<reqwest::header::HeaderValue>(),
                    ) {
                        headers.insert(n, v);
                    }
                }
                ensure_figma_user_agent(&mut headers, name, &config.url);
                let http_client = reqwest::Client::builder()
                    .default_headers(headers)
                    .build()
                    .map_err(|e| {
                        McpError::ClientError(format!("Failed to build HTTP client: {e}"))
                    })?;
                // `AuthClient::new` wants an owned manager, but ours is shared
                // (`Arc`) with the OAuth flow; the struct is non_exhaustive, so
                // build with a throwaway manager and swap in the shared one.
                let placeholder_manager =
                    rmcp::transport::auth::AuthorizationManager::new(config.url.as_str())
                        .await
                        .map_err(|e| {
                            McpError::ClientError(format!("Failed to build OAuth client: {e}"))
                        })?;
                let mut auth_client =
                    rmcp::transport::auth::AuthClient::new(http_client, placeholder_manager);
                auth_client.auth_manager = auth_manager.clone();
                let mcp_http_client = crate::mcp_http_client::McpHttpClient::new(
                    auth_client,
                    name.as_str(),
                    self.warn_budget.clone(),
                );
                let transport_config =
                    StreamableHttpClientTransportConfig::with_uri(config.url.as_str());
                let transport =
                    StreamableHttpClientTransport::with_client(mcp_http_client, transport_config);
                let handler = self.make_client_handler();
                tokio::time::timeout(timeout, handler.serve(transport))
                    .await
                    .map_err(|_| McpError::timeout(name, timeout))?
                    .map_err(|e| McpError::HandshakeFailed {
                        server: name.to_string(),
                        source: Box::new(e),
                    })
            }
            PendingTransport::Acp { server_id, invoker } => {
                // Per-reverse-call backstop on `kigi/mcp/sdk_call`: the larger of the
                // startup and tool timeouts, so it never undercuts the real outer bound
                // (the handshake `initialize` is bounded by the serve `timeout` below;
                // tool calls by `tool_timeout_for` in `try_call_tool`). The bridge
                // forwards raw JSON-RPC without the tool name, so per-TOOL overrides
                // aren't applied here in v1; the HTTP path still honors them.
                let invoke_timeout = std::time::Duration::from_secs(
                    self.startup_timeout_sec.max(self.tool_timeout_sec),
                );
                let transport =
                    crate::acp_transport::acp_bridge_transport(server_id, invoker, invoke_timeout);
                let handler = self.make_client_handler();
                tokio::time::timeout(timeout, handler.serve(transport))
                    .await
                    .map_err(|_| McpError::timeout(name, timeout))?
                    .map_err(|e| McpError::HandshakeFailed {
                        server: name.to_string(),
                        source: Box::new(e),
                    })
            }
        }
    }

    fn make_client_info(server_name: &str) -> ClientInfo {
        let mut extensions = rmcp::model::ExtensionCapabilities::new();
        extensions.insert(
            "io.modelcontextprotocol/ui".to_string(),
            serde_json::from_value(serde_json::json!({
                "mimeTypes": ["text/html;profile=mcp-app"]
            }))
            .unwrap_or_default(),
        );
        let mut capabilities = ClientCapabilities::default();
        capabilities.extensions = Some(extensions);
        ClientInfo::new(
            capabilities,
            Implementation::new(
                format!("kigi-shell-{server_name}"),
                kigi_version::VERSION.to_string(),
            ),
        )
        // rmcp's default `ProtocolVersion` tracks its LATEST; pin explicitly
        // so the advertised protocol only changes deliberately, never as a
        // side effect of an rmcp bump.
        .with_protocol_version(rmcp::model::ProtocolVersion::V_2025_06_18)
    }

    /// Build the [`KigiClientHandler`] that drives `client.serve(...)`.
    ///
    /// The handler holds a **clone of `Arc<Mutex<Option<Sender>>>`**,
    /// not a snapshot — so any subsequent call to
    /// [`Self::set_event_tx`] is observed by the live rmcp service
    /// loop on its next notification.
    fn make_client_handler(&self) -> KigiClientHandler {
        KigiClientHandler {
            info: Self::make_client_info(&self.server_name),
            server_name: self.server_name.clone(),
            notify_tx: Arc::clone(&self.notify_tx),
        }
    }

    /// Wire a sender for [`McpClientEvent`]s emitted by this client.
    ///
    /// Mutates the shared slot synchronously. All previously-cloned
    /// references (the [`KigiClientHandler`] handed to
    /// `client.serve`, the [`crate::liveness::spawn_transport_liveness`]
    /// task) read through the same Arc, so this is observed
    /// session-wide on the next event.
    pub fn set_event_tx(&self, tx: Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>>) {
        *self.notify_tx.lock() = tx;
    }

    /// Snapshot the current event sender, if any.
    ///
    /// Used by [`crate::liveness::spawn_transport_liveness`] (which
    /// captures a `Sender` clone at spawn time) and by
    /// [`Self::ensure_initialized`]'s post-handshake emit. Synchronous
    /// because the shared slot is a `parking_lot::Mutex`.
    pub fn event_tx_clone(&self) -> Option<tokio::sync::mpsc::UnboundedSender<McpClientEvent>> {
        self.notify_tx.lock().clone()
    }

    /// Install or replace this client's transport-liveness handle.
    /// Dropping the previous handle (if any) cancels its task; the
    /// new handle starts polling on its own schedule. Pass `None` to
    /// stop watching without installing a new one.
    ///
    /// Synchronous: the slot is a `parking_lot::Mutex`. The poller
    /// task uses an `Arc` clone of this same mutex so it can clear the
    /// slot from inside the task before exiting.
    pub fn set_liveness_handle(&self, handle: Option<crate::liveness::TransportLivenessHandle>) {
        *self.liveness_handle.lock() = handle;
    }

    /// Arm the per-client transport-liveness watcher.
    ///
    /// Idempotent and gated:
    /// - Returns `false` for in-process SDK ([`Self::is_acp`]) clients: the
    ///   watcher's only output is `TransportClosed`, which the dispatcher can't
    ///   recover for ACP (not in `configs`), so it would evict the client. ACP
    ///   recovers lazily via [`Self::reset_transport`] instead. Gated here so no
    ///   caller can forget it.
    /// - Returns `false` if there's no `notify_tx` wired (subagent
    ///   snapshot or pre-dispatcher state) — nothing to do.
    /// - Returns `false` if the client isn't `Ready` — armed pollers
    ///   would just exit silently on their first poll, but skipping
    ///   the spawn entirely is cheaper.
    /// - Returns `false` if a live handle is already installed.
    /// - Otherwise spawns the poller and stores the handle.
    ///
    /// **TOCTOU note**: the state check is performed before the
    /// liveness lock is acquired. A concurrent re-handshake could move
    /// the state to `Initializing` between the check and the spawn.
    /// This is benign — the poller's first tick observes the
    /// non-`Ready` state and exits silently without emitting. So the
    /// worst case under TOCTOU is "the poller starts and immediately
    /// stops"; it never produces a spurious `TransportClosed`.
    ///
    /// Lifecycle: when the watcher emits `TransportClosed` it clears
    /// the slot itself; the next `arm_liveness_watcher` call can
    /// install a fresh handle without a manual
    /// [`Self::set_liveness_handle`] reset.
    pub async fn arm_liveness_watcher(
        self: &Arc<Self>,
        poll_interval: std::time::Duration,
    ) -> bool {
        if self.is_acp() {
            return false;
        }
        let Some(event_tx) = self.event_tx_clone() else {
            return false;
        };
        if !matches!(self.state_kind().await, ClientStateKind::Ready) {
            return false;
        }
        let mut slot = self.liveness_handle.lock();
        if slot.is_some() {
            return false;
        }
        let handle = crate::liveness::spawn_transport_liveness(
            self.server_name.clone(),
            Arc::clone(self),
            poll_interval,
            event_tx,
            Arc::clone(&self.liveness_handle),
        );
        *slot = Some(handle);
        true
    }

    fn build_http_transport(
        config: &HttpConfig,
        server_name: &str,
        warn_budget: crate::mcp_http_client::WarnBudget,
    ) -> Result<
        StreamableHttpClientTransport<crate::mcp_http_client::McpHttpClient<reqwest::Client>>,
        McpError,
    > {
        let mut headers = reqwest::header::HeaderMap::new();
        for (key, value) in &config.headers {
            match (
                reqwest::header::HeaderName::from_bytes(key.as_bytes()),
                value.parse::<reqwest::header::HeaderValue>(),
            ) {
                (Ok(name), Ok(val)) => {
                    headers.insert(name, val);
                }
                _ => {
                    tracing::warn!("Skipping invalid MCP HTTP header: {key}");
                }
            }
        }
        ensure_figma_user_agent(&mut headers, server_name, &config.url);
        let client = reqwest::Client::builder()
            .default_headers(headers)
            .build()
            .map_err(|e| McpError::ClientError(format!("Failed to build HTTP client: {e}")))?;
        let mcp_http_client =
            crate::mcp_http_client::McpHttpClient::new(client, server_name, warn_budget);
        let transport_config = StreamableHttpClientTransportConfig::with_uri(config.url.as_str());
        Ok(StreamableHttpClientTransport::with_client(
            mcp_http_client,
            transport_config,
        ))
    }

    /// Cheap, non-blocking liveness predicate.
    ///
    /// Inspects the current [`ClientState`] under the state mutex only —
    /// it MUST NOT call [`Self::ensure_initialized`] or any other path
    /// that can trigger a network round-trip. The previous implementation
    /// went through `ensure_initialized`, which could block UI callers
    /// (e.g. an MCP status modal) for up to `startup_timeout_sec` seconds
    /// on a dead stdio server.
    ///
    /// Semantics:
    /// - `Ready(service)` with an open transport → `true`.
    /// - `Ready(service)` whose receiver-side has been dropped (typically
    ///   because the rmcp service loop terminated) → `false`. rmcp 2.1
    ///   `Peer::is_transport_closed` reports `self.tx.is_closed()` at
    ///   `service.rs:703-705`; `RunningService` derefs to `Peer` at
    ///   `service.rs:716-722`.
    /// - Any other variant (`Empty`, `Pending`, `Initializing`) →
    ///   `false`.
    ///
    /// HTTP idle caveat: for [`StreamableHttpClientTransport`] the rmcp
    /// service loop only terminates on an outgoing send failure or an
    /// explicit shutdown. A long-idle HTTP server therefore keeps
    /// `is_transport_closed()` returning `false`, and this method
    /// continues to report `true`. That is the desired semantics — a
    /// liveness probe would belong in a separate watcher, not here.
    pub async fn is_healthy(&self) -> bool {
        let guard = self.state.lock().await;
        match &*guard {
            ClientState::Ready(service) => !service.is_transport_closed(),
            _ => false,
        }
    }

    /// Atomic classification for the liveness watcher.
    ///
    /// Reads `state` once and projects onto
    /// [`LivenessCheck`]: distinguishes "transport actually closed"
    /// (emit + exit) from "state moved out of `Ready`" (exit
    /// silently). The watcher depends on this distinction: a plain
    /// `is_healthy`-based predicate cannot tell the cases apart and
    /// would false-fire `TransportClosed` on re-handshake transitions.
    pub async fn liveness_check(&self) -> LivenessCheck {
        let guard = self.state.lock().await;
        match &*guard {
            ClientState::Ready(service) => {
                if service.is_transport_closed() {
                    LivenessCheck::TransportClosed
                } else {
                    LivenessCheck::Healthy
                }
            }
            _ => LivenessCheck::Transient,
        }
    }

    /// State-machine snapshot for diagnostics and downstream UI.
    ///
    /// Like [`Self::is_healthy`], this is a cheap state inspection
    /// (no handshake, no network I/O). Maps [`ClientState`] onto a
    /// `Copy` enum so callers can match without holding a reference to
    /// the inner [`McpService`] / [`PendingTransport`].
    pub async fn state_kind(&self) -> ClientStateKind {
        let guard = self.state.lock().await;
        match &*guard {
            ClientState::Empty => ClientStateKind::Empty,
            ClientState::Pending(_) => ClientStateKind::Pending,
            ClientState::Initializing => ClientStateKind::Initializing,
            ClientState::Ready(_) => ClientStateKind::Ready,
        }
    }

    /// Materialize this server's tool descriptors as JSON files under
    /// `<server_dir>/tools/`.
    ///
    /// The model reads these before issuing an MCP tool call.
    /// Each tool becomes `<server_dir>/tools/<sanitized_tool_name>.json`
    /// with `{name, description, inputSchema}`. Resources are intentionally
    /// not materialized: this harness exposes only MCP tool calls, so resource
    /// descriptors would advertise MCP-resource tools the model can't use.
    ///
    /// Best-effort: errors writing individual descriptors are logged but
    /// don't abort the materialization. Returns the number of files written.
    pub async fn materialize_descriptors(
        &self,
        server_dir: &std::path::Path,
    ) -> Result<usize, McpError> {
        let mcp_service = self.ensure_initialized().await?;

        // Collect descriptors via the async MCP API, then defer all filesystem
        // work to a single `spawn_blocking` so the executor is never blocked on
        // `std::fs` (this runs on every MCP tool-set change, not just startup).
        let mut files: Vec<(String, Vec<u8>)> = Vec::new();
        let mut cursor: Option<String> = None;
        loop {
            let result = mcp_service
                .list_tools(Some(
                    PaginatedRequestParams::default().with_cursor(cursor.clone()),
                ))
                .await?;
            for tool in result.tools {
                let descriptor = serde_json::json!({
                    "name": tool.name.as_ref(),
                    "description": tool.description.as_deref(),
                    "inputSchema": tool.input_schema.as_ref(),
                });
                match serde_json::to_vec_pretty(&descriptor) {
                    Ok(bytes) => files.push((
                        format!("{}.json", sanitize_descriptor_segment(tool.name.as_ref())),
                        bytes,
                    )),
                    Err(e) => tracing::warn!(
                        tool = %tool.name.as_ref(),
                        error = %e,
                        "failed to serialize MCP tool descriptor",
                    ),
                }
            }
            match result.next_cursor {
                Some(next) => cursor = Some(next),
                None => break,
            }
        }

        // Write each descriptor atomically (temp file + rename) so a concurrent
        // reader never sees a half-written JSON and overlapping writers converge
        // without a lock.
        let tools_dir = server_dir.join("tools");
        tokio::task::spawn_blocking(move || -> Result<usize, McpError> {
            std::fs::create_dir_all(&tools_dir).map_err(|e| {
                McpError::ClientError(format!(
                    "failed to create MCP tools descriptor dir {}: {e}",
                    tools_dir.display()
                ))
            })?;
            let mut written = 0usize;
            for (file_name, bytes) in files {
                let path = tools_dir.join(&file_name);
                let write_result =
                    tempfile::NamedTempFile::new_in(&tools_dir).and_then(|mut tmp| {
                        std::io::Write::write_all(&mut tmp, &bytes)?;
                        tmp.persist(&path).map_err(|e| e.error)
                    });
                match write_result {
                    Ok(_) => written += 1,
                    Err(e) => tracing::warn!(
                        path = %path.display(),
                        error = %e,
                        "failed to write MCP tool descriptor",
                    ),
                }
            }
            Ok(written)
        })
        .await
        .map_err(|e| McpError::ClientError(format!("descriptor write task panicked: {e}")))?
    }

    /// Read the server's `instructions` from the MCP initialize handshake.
    /// Returns `None` if the client isn't ready yet.
    pub async fn server_instructions(&self) -> Option<String> {
        let guard = self.state.lock().await;
        if let ClientState::Ready(service) = &*guard {
            service
                .peer_info()?
                .instructions
                .as_deref()
                .filter(|s| !s.trim().is_empty())
                .map(String::from)
        } else {
            None
        }
    }

    pub async fn get_tool_registrations(
        &self,
        mcp_state: Arc<Mutex<McpState>>,
    ) -> Result<Vec<McpToolRegistration>, McpError> {
        let _ensure_init_timer = kigi_log::instrumentation::timer("mcp_ensure_initialized");
        let mcp_service = self.ensure_initialized().await?;

        let mut all_tools = Vec::new();
        let mut cursor: Option<String> = None;

        let _list_tools_timer = kigi_log::instrumentation::timer("mcp_list_tools");
        loop {
            let list_tools_result = mcp_service
                .list_tools(Some(
                    PaginatedRequestParams::default().with_cursor(cursor.clone()),
                ))
                .await?;

            all_tools.extend(list_tools_result.tools);

            match list_tools_result.next_cursor {
                Some(next) => cursor = Some(next),
                None => break,
            }
        }

        let registrations: Vec<_> = all_tools
            .into_iter()
            .filter_map(|tool| {
                let meta: Option<serde_json::Value> = tool
                    .meta
                    .as_ref()
                    .and_then(|m| serde_json::to_value(m).ok());

                let name = tool.name.to_string();
                let description = tool.description.map(|d| d.to_string()).unwrap_or_default();
                let mut schema = serde_json::to_value(tool.input_schema.as_ref())
                    .unwrap_or_else(|_| serde_json::json!({"type": "object"}));
                // Ensure the schema has "type": "object" — some MCP servers
                // (e.g., VSCode) send `inputSchema: {}` for tools with no
                // parameters. Azure's OpenAI API rejects schemas without a
                // `type` field with: 'schema must be a JSON Schema of type:
                // "object", got type: "None"'.
                if let Some(obj) = schema.as_object_mut() {
                    obj.entry("type")
                        .or_insert_with(|| serde_json::json!("object"));
                    obj.entry("properties")
                        .or_insert_with(|| serde_json::json!({}));
                }

                let mcp_tool = McpTool {
                    name,
                    description,
                    server_name: self.server_name.clone(),
                    mcp_state: Arc::clone(&mcp_state),
                    schema,
                    meta,
                };
                // Invalid tools (bad names) return None and are skipped
                mcp_tool.into_registration()
            })
            .collect();

        // Warn about tool_timeouts keys that don't match any discovered tool.
        // This catches typos like `creat_issue` instead of `create_issue`.
        if !self.tool_timeouts.is_empty() {
            // Registration names are qualified ("server__tool"); tool_timeouts
            // keys are raw tool names. Strip the server prefix for comparison.
            let prefix = format!("{}{}", self.server_name, MCP_TOOL_NAME_DELIMITER);
            let raw_names: Vec<&str> = registrations
                .iter()
                .map(|r| r.name.strip_prefix(prefix.as_str()).unwrap_or(&r.name))
                .collect();
            let discovered: std::collections::HashSet<&str> = raw_names.iter().copied().collect();
            for key in self.tool_timeouts.keys() {
                if !discovered.contains(key.as_str()) {
                    tracing::info!(
                        server = %self.server_name,
                        tool_timeout_key = %key,
                        "tool_timeouts entry '{}' does not match any tool exposed by MCP server '{}' \
                         (available: {}). The per-tool timeout will have no effect — check for typos.",
                        key,
                        self.server_name,
                        raw_names.join(", "),
                    );
                }
            }
        }

        Ok(registrations)
    }

    /// Call a tool directly on the MCP server (for testing/debugging).
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> Result<rmcp::model::CallToolResult, McpError> {
        let mcp_service = self.ensure_initialized().await?;
        let result = mcp_service
            .call_tool({
                let mut params = CallToolRequestParams::new(tool_name.to_string());
                params.arguments = arguments.as_object().cloned();
                params
            })
            .await?;
        Ok(result)
    }
}

fn contains_session_placeholder(value: &str) -> bool {
    value.contains("{{session_id}}") || value.contains("${session_id}")
}

/// Sanitize an MCP server name into a safe filename component.
fn sanitize_mcp_log_filename(name: &str) -> String {
    let sanitized: String = name
        .chars()
        .take(96)
        .map(|c| match c {
            c if c.is_ascii_alphanumeric() => c,
            '.' | '_' | '-' => c,
            _ => '_',
        })
        .collect();
    if sanitized.is_empty() {
        "server".into()
    } else {
        sanitized
    }
}

/// Copy an MCP server's stderr to `~/.kigi/logs/mcp/<server>.stderr.log`
/// in a background task. Truncated per spawn.
fn drain_mcp_stderr_to_log(server_name: &str, mut stderr: tokio::process::ChildStderr) {
    let log_dir = kigi_config::kigi_home().join("logs").join("mcp");
    if let Err(e) = std::fs::create_dir_all(&log_dir) {
        tracing::warn!("MCP stderr drain: failed to create log dir: {e}");
        return;
    }
    let log_path = log_dir.join(format!(
        "{}.stderr.log",
        sanitize_mcp_log_filename(server_name)
    ));

    let file = match std::fs::OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&log_path)
    {
        Ok(f) => f,
        Err(e) => {
            tracing::warn!(
                "MCP stderr drain: failed to open {}: {e}",
                log_path.display()
            );
            return;
        }
    };
    let mut file = tokio::fs::File::from_std(file);
    let server_name = server_name.to_string();
    tokio::spawn(async move {
        if let Err(e) = tokio::io::copy(&mut stderr, &mut file).await {
            tracing::warn!("MCP stderr drain '{server_name}': {e}");
        }
    });
}

fn expand_session_id_headers(
    headers: Vec<acp::HttpHeader>,
    session_id: Option<&str>,
) -> Vec<(String, String)> {
    headers
        .into_iter()
        .filter_map(|header| {
            let value = header.value;
            if let Some(session_id) = session_id {
                let expanded = value
                    .replace("{{session_id}}", session_id)
                    .replace("${session_id}", session_id);
                Some((header.name, expanded))
            } else if contains_session_placeholder(&value) {
                None
            } else {
                Some((header.name, value))
            }
        })
        .collect()
}

/// Decide the actual (program, args) to spawn for a stdio MCP server.
///
/// On Windows, npm ships launchers like `npx`/`npm`/`pnpm`/`yarn` as `.cmd`
/// batch shims (there is no `npx.exe`). `CreateProcessW` only appends `.exe`
/// and ignores `PATHEXT`, so `Command::new("npx")` fails with "file not
/// found". We resolve the bare name on `PATH` (honoring `PATHEXT`, via the
/// `resolve` closure) so std spawns the real launcher path (e.g. `npx.cmd`) —
/// std then runs `.cmd`/`.bat` through `cmd.exe` with hardened arg escaping. On
/// non-Windows we never touch the command (verified working). A command
/// containing a path separator is used as-is. The resolved path is returned as
/// an `OsString` so it reaches `Command::new` without a lossy UTF-8 round-trip.
fn plan_stdio_spawn(
    command: &str,
    args: &[String],
    is_windows: bool,
    resolve: impl Fn(&str) -> Option<std::path::PathBuf>,
) -> (OsString, Vec<String>) {
    if is_windows
        && !command.contains('/')
        && !command.contains('\\')
        && let Some(resolved) = resolve(command)
    {
        return (resolved.into_os_string(), args.to_vec());
    }
    (OsString::from(command), args.to_vec())
}

fn is_figma_mcp(server_name: &str, url: &str) -> bool {
    if server_name.eq_ignore_ascii_case("figma") {
        return true;
    }
    reqwest::Url::parse(url)
        .ok()
        .and_then(|u| u.host_str().map(|h| h.to_ascii_lowercase()))
        .is_some_and(|h| h == "figma.com" || h.ends_with(".figma.com"))
}

fn ensure_figma_user_agent(headers: &mut reqwest::header::HeaderMap, server_name: &str, url: &str) {
    if !is_figma_mcp(server_name, url) {
        return;
    }
    if headers.contains_key(reqwest::header::USER_AGENT) {
        return;
    }
    headers.insert(
        reqwest::header::USER_AGENT,
        reqwest::header::HeaderValue::from_static("kigi-cli"),
    );
}

fn stdio_path_override(env: &[acp::EnvVariable]) -> Option<&str> {
    env.iter()
        .find(|e| e.name.eq_ignore_ascii_case("PATH"))
        .map(|e| e.value.as_str())
}

pub async fn start_mcp_server(
    mcp_server: acp::McpServer,
    session_id: Option<&str>,
    overrides: Option<&McpClientTimeoutOverrides>,
    meta_config: Option<&McpServerMetaConfig>,
    byo_config: Option<&McpOAuthConfig>,
    event_writer: &kigi_file_utils::events::EventWriter,
    mode: OauthInteractivity,
) -> Result<McpClient, McpError> {
    let _per_server_timer = kigi_log::instrumentation::timer("mcp_start_one_server");
    match mcp_server {
        acp::McpServer::Stdio(acp::McpServerStdio {
            name,
            command,
            args,
            env,
            ..
        }) => {
            if let Some(mc) = meta_config {
                tracing::info!(server = %name, ?mc, "MCP stdio: meta config override");
            }

            let command_str = command.to_string_lossy().into_owned();
            let _stdio_spawn_timer = kigi_log::instrumentation::timer("mcp_stdio_spawn");
            let path_override = stdio_path_override(&env);
            let (program, spawn_args) = plan_stdio_spawn(&command_str, &args, cfg!(windows), |c| {
                if let Some(path) = path_override
                    && let Ok(cwd) = std::env::current_dir()
                {
                    which::which_in(c, Some(path), cwd).ok()
                } else {
                    which::which(c).ok()
                }
            });
            let mut cmd = Command::new(&program);
            cmd.kill_on_drop(true).args(&spawn_args);
            for env_variable in &env {
                cmd.env(&env_variable.name, &env_variable.value);
            }
            kigi_tools::util::detach_command(&mut cmd);

            let (transport, stderr_handle) =
                SafeTokioChildProcess::spawn(cmd, name.clone(), event_writer.clone()).map_err(
                    |e| {
                        tracing::error!("Failed to spawn MCP server '{}': {}", name, e);
                        McpError::SpawnFailed {
                            server: name.clone(),
                            source: e,
                        }
                    },
                )?;

            tracing::debug!("MCP server '{}' spawned: PID={:?}", name, transport.id());

            if let Some(stderr) = stderr_handle {
                drain_mcp_stderr_to_log(&name, stderr);
            }

            Ok(McpClient::new_stdio(
                name.clone(),
                transport,
                overrides,
                meta_config,
            ))
        }
        acp::McpServer::Http(acp::McpServerHttp {
            name, url, headers, ..
        })
        | acp::McpServer::Sse(acp::McpServerSse {
            name, url, headers, ..
        }) => {
            if let Some(mc) = meta_config {
                tracing::info!(server = %name, %url, ?mc, "MCP http: meta config override");
            }

            let headers = expand_session_id_headers(headers, session_id);
            let http_config = HttpConfig {
                url: url.clone(),
                headers,
            };

            let has_existing_auth = http_config
                .headers
                .iter()
                .any(|(k, _)| k.eq_ignore_ascii_case("authorization"));

            let auth_prep = if has_existing_auth {
                tracing::debug!(
                    server = %name,
                    "Skipping OAuth discovery: server already has Authorization header"
                );
                HttpOauthPrep::NoOauthSupport
            } else {
                let _auth_discovery_timer =
                    kigi_log::instrumentation::timer("mcp_http_auth_discovery");
                match tokio::time::timeout(
                    OAUTH_DISCOVERY_TIMEOUT,
                    discover_and_prepare_auth(&name, &url, mode),
                )
                .await
                {
                    Ok(result) => result,
                    Err(_) => {
                        tracing::warn!(
                            server = %name,
                            url = %url,
                            ?mode,
                            timeout_secs = OAUTH_DISCOVERY_TIMEOUT.as_secs(),
                            "OAuth discovery timed out"
                        );
                        event_writer.emit(
                            kigi_file_utils::events::Event::McpOAuthDiscoveryTimeout {
                                server_name: name.clone(),
                                url: url.clone(),
                            },
                        );
                        HttpOauthPrep::on_probe_failure(mode)
                    }
                }
            };
            match auth_prep {
                HttpOauthPrep::ManagerReady(auth_mgr) => Ok(McpClient::new_http_auth(
                    name.clone(),
                    http_config,
                    auth_mgr,
                    byo_config.cloned(),
                    overrides,
                    meta_config,
                )),
                HttpOauthPrep::NoOauthSupport => Ok(McpClient::new_http(
                    name.clone(),
                    http_config,
                    overrides,
                    meta_config,
                )),
                // Avoid starting an unauthenticated HTTP worker that fatals on server OAuth challenge.
                HttpOauthPrep::NeedsInteractiveLogin => Err(McpError::AuthRequired {
                    server: name.clone(),
                }),
            }
        }
        // TODO(acp-0.10): `McpServer` is #[non_exhaustive]; reject unknown transports.
        other => Err(McpError::ClientError(format!(
            "unsupported MCP server transport: {other:?}"
        ))),
    }
}

pub async fn start_mcp_servers(
    mcp_servers: Vec<acp::McpServer>,
    session_id: Option<&str>,
    overrides_map: &HashMap<String, McpClientTimeoutOverrides>,
    meta_config_map: &McpMetaConfigMap,
    oauth_config_map: &crate::oauth_config::McpOAuthConfigMap,
    event_writer: &kigi_file_utils::events::EventWriter,
    mode: OauthInteractivity,
) -> Vec<Result<McpClient, McpError>> {
    let _mcp_start_timer = kigi_log::instrumentation::timer("mcp_start_servers");

    if !meta_config_map.is_empty() {
        tracing::info!(
            count = mcp_servers.len(),
            overrides = ?meta_config_map.keys().collect::<Vec<_>>(),
            "Starting MCP servers with meta config"
        );
    }

    futures::stream::iter(mcp_servers)
        .map(|server| {
            let server_name = mcp_server_name(&server);
            let overrides = overrides_map.get(server_name);
            let mc = meta_config_map.get(server_name);
            let byo = oauth_config_map.get(server_name);
            start_mcp_server(server, session_id, overrides, mc, byo, event_writer, mode)
        })
        .buffer_unordered(8)
        .collect::<Vec<_>>()
        .await
}

/// Extract the name from an MCP server enum variant.
pub fn mcp_server_name(server: &acp::McpServer) -> &str {
    match server {
        acp::McpServer::Stdio(stdio) => &stdio.name,
        acp::McpServer::Http(http) => &http.name,
        acp::McpServer::Sse(sse) => &sse.name,
        // TODO(acp-0.10): `McpServer` is #[non_exhaustive].
        _ => "unknown",
    }
}

pub fn mcp_transport_str(server: &acp::McpServer) -> &'static str {
    match server {
        acp::McpServer::Stdio(_) => "stdio",
        acp::McpServer::Http(_) => "http",
        acp::McpServer::Sse(_) => "sse",
        // TODO(acp-0.10): `McpServer` is #[non_exhaustive].
        _ => "unknown",
    }
}

pub fn mcp_target_str(server: &acp::McpServer) -> String {
    match server {
        acp::McpServer::Stdio(acp::McpServerStdio { command, args, .. }) => {
            let cmd = command.to_string_lossy();
            if args.is_empty() {
                cmd.to_string()
            } else {
                format!("{} {}", cmd, args.join(" "))
            }
        }
        acp::McpServer::Http(acp::McpServerHttp { url, .. })
        | acp::McpServer::Sse(acp::McpServerSse { url, .. }) => url.clone(),
        // TODO(acp-0.10): `McpServer` is #[non_exhaustive].
        _ => String::new(),
    }
}

impl McpClient {
    /// Minimal stub for unit tests in dependent crates. Hidden from rustdoc;
    /// not gated behind `#[cfg(test)]` so cross-crate test code can construct
    /// it without the host crate enabling a feature.
    #[doc(hidden)]
    pub fn stub(name: &str) -> Self {
        // Route through the single constructor (so new fields never need
        // touching here), then downgrade to the no-transport placeholder:
        // `Empty` state makes `ensure_initialized` error, and `reconnect =
        // None` makes `reset_transport` return false — i.e. a client that
        // can't reconnect, like a dead Stdio child. Overrides preserve the
        // historical stub timeouts (10s startup / 60s tool).
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(10),
            tool_timeout_sec: Some(60),
            ..Default::default()
        };
        let mut client = Self::new_with_transport(
            name.to_string(),
            PendingTransport::Http(HttpConfig {
                url: String::new(),
                headers: Vec::new(),
            }),
            Some(&overrides),
            None,
            None,
            None,
            None,
        );
        *client.state.get_mut() = ClientState::Empty;
        client.reconnect = None;
        client
    }
}

/// rmcp [`ClientHandler`] used by all MCP transports.
///
/// Replaces the previous bare [`ClientInfo`] handler at the three
/// `client.serve(...)` call sites in [`McpClient::try_handshake`].
/// Plumbs server-pushed notifications through an
/// [`tokio::sync::mpsc::UnboundedSender<McpClientEvent>`] so the
/// session-actor dispatcher can fan them out as ACP
/// `kigi/mcp/server_status` events.
///
/// ## RPIT, not `#[async_trait]`
///
/// rmcp 2.1's [`ClientHandler`] declares its async methods as
/// return-position `impl Future` (see
/// `~/.cargo/registry/src/.../rmcp-2.1.0/src/handler/client.rs`,
/// lines 202–217). Applying `#[async_trait]` here would produce
/// methods whose signature mismatches the trait, and the impl would
/// not satisfy the bound. The macro path is also unnecessary — the
/// trait already supports `async fn` syntax indirectly via
/// `impl Future<Output = ()> + Send + '_`, which is what we mirror.
///
/// Future contributor reading this: do **not** add `#[async_trait]`.
/// The methods below intentionally return `impl Future` directly.
///
/// ## Notification routing
///
/// `on_tool_list_changed` / `on_resource_list_changed` push an
/// [`McpClientEvent`] into [`Self::notify_tx`]. If the receiver has
/// been dropped (subagent teardown, session shutdown, or the field
/// was `None` to begin with — see [`McpClient::notify_tx`] doc), the
/// send fails silently; rmcp must not see an error from a
/// notification handler or the service loop tears down.
#[derive(Debug)]
pub struct KigiClientHandler {
    /// Static `ClientInfo` returned by [`Self::get_info`]; built once
    /// at handshake time and stored to avoid re-allocating per call.
    info: ClientInfo,
    /// MCP server name this handler is bound to. Cloned into emitted
    /// events so the dispatcher can route per-server.
    server_name: McpServerName,
    /// **Shared** event sink — the same Arc lives on the owning
    /// [`McpClient`]. Mutating the slot via [`McpClient::set_event_tx`]
    /// is observed here on the next read, so wiring the sender
    /// post-handshake is supported without restarting the rmcp
    /// service loop.
    notify_tx: SharedEventTx,
}

impl KigiClientHandler {
    /// Best-effort event emit. Reads the shared `notify_tx` slot on
    /// every call (so the handler picks up any post-handshake wiring
    /// done by [`McpClient::set_event_tx`]). Drops the send error: if
    /// the receiver is gone, the consumer has shut down and there's
    /// nothing useful to do here. Splitting this out keeps the trait
    /// methods short.
    fn emit(&self, ev: McpClientEvent) {
        let sender = self.notify_tx.lock().clone();
        if let Some(tx) = sender {
            let _ = tx.send(ev);
        }
    }
}

impl ClientHandler for KigiClientHandler {
    // NOTE: `async fn` here is sugar for the trait's
    // `-> impl Future<Output = ()> + Send + '_`. We INTENTIONALLY do
    // not use `#[async_trait]` — rmcp 2.1's `ClientHandler` declares
    // its notification methods as return-position `impl Future`, and
    // async_trait would produce a different (incompatible) signature.
    // See the [`KigiClientHandler`] doc-comment for the full RPIT
    // contract.
    async fn on_tool_list_changed(&self, _context: NotificationContext<RoleClient>) {
        self.emit(McpClientEvent::ToolsChanged {
            server: self.server_name.clone(),
        });
    }

    async fn on_resource_list_changed(&self, _context: NotificationContext<RoleClient>) {
        self.emit(McpClientEvent::ResourcesChanged {
            server: self.server_name.clone(),
        });
    }

    fn get_info(&self) -> ClientInfo {
        self.info.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    /// A single undecodable line on an MCP stdio server's stdout must NOT
    /// collapse the transport: if the decode error surfaced as `None`, the
    /// service would read it as EOF → "Transport closed" → `tools/list` fails
    /// and the connector "shows but doesn't work". The resilient transport
    /// skips the bad line and keeps reading, so a stray stdout log line never
    /// takes the whole server down.
    #[tokio::test]
    async fn resilient_transport_skips_undecodable_line_and_keeps_stream_alive() {
        // `server_out` is the writer half (the fake server's stdout); the
        // transport reads framed JSON-RPC from `client_in`.
        let (mut server_out, client_in) = tokio::io::duplex(64 * 1024);
        let mut transport = ResilientRwTransport::new(
            client_in,
            tokio::io::sink(),
            "fwbuild".to_string(),
            kigi_file_utils::events::EventWriter::noop(),
        );

        let valid = r#"{"jsonrpc":"2.0","method":"notifications/tools/list_changed"}"#;
        // A stray non-JSON log line — the shape that, under rmcp's stock
        // transport, decodes to an error and closes the connection.
        let garbage = "info: fwbuild started, listening on stdio";
        server_out
            .write_all(format!("{valid}\n{garbage}\n{valid}\n").as_bytes())
            .await
            .unwrap();
        // Dropping the writer half signals a clean end-of-stream.
        drop(server_out);

        assert!(
            transport.receive().await.is_some(),
            "first valid message must be received"
        );
        assert!(
            transport.receive().await.is_some(),
            "the undecodable line must be skipped and the next valid message delivered"
        );
        assert!(
            transport.receive().await.is_none(),
            "only a genuine end-of-stream yields None"
        );
    }

    fn make_stdio_server(name: &str, command: &str) -> acp::McpServer {
        acp::McpServer::Stdio(acp::McpServerStdio::new(name, PathBuf::from(command)))
    }

    fn make_http_server(name: &str, url: &str) -> acp::McpServer {
        acp::McpServer::Http(acp::McpServerHttp::new(name, url))
    }

    #[test]
    fn plan_stdio_spawn_windows_resolves_bare_launcher_to_cmd_shim() {
        let args = vec!["-y".to_string(), "@scope/pkg".to_string()];
        let (program, spawn_args) = plan_stdio_spawn("npx", &args, true, |c| {
            assert_eq!(c, "npx");
            Some(PathBuf::from(r"C:\path\npx.cmd"))
        });
        assert_eq!(program, OsString::from(r"C:\path\npx.cmd"));
        assert_eq!(spawn_args, args);
    }

    #[test]
    fn plan_stdio_spawn_windows_unresolved_falls_back_to_raw_command() {
        let args = vec!["-y".to_string(), "@scope/pkg".to_string()];
        let (program, spawn_args) = plan_stdio_spawn("npx", &args, true, |_| None);
        assert_eq!(program, OsString::from("npx"));
        assert_eq!(spawn_args, args);
    }

    #[test]
    fn plan_stdio_spawn_windows_backslash_path_command_used_as_is_without_resolving() {
        let args = vec!["--config".to_string(), "x.json".to_string()];
        let (program, spawn_args) = plan_stdio_spawn(r"C:\tools\server.exe", &args, true, |_| {
            panic!("resolver must not be consulted for a command with a backslash separator")
        });
        assert_eq!(program, OsString::from(r"C:\tools\server.exe"));
        assert_eq!(spawn_args, args);
    }

    #[test]
    fn plan_stdio_spawn_windows_forward_slash_path_command_used_as_is_without_resolving() {
        let args = vec!["--port".to_string(), "8080".to_string()];
        let (program, spawn_args) = plan_stdio_spawn("C:/tools/server.exe", &args, true, |_| {
            panic!("resolver must not be consulted for a command with a forward-slash separator")
        });
        assert_eq!(program, OsString::from("C:/tools/server.exe"));
        assert_eq!(spawn_args, args);
    }

    #[test]
    fn plan_stdio_spawn_non_windows_never_resolves() {
        let args = vec!["-y".to_string(), "pkg".to_string()];
        let (program, spawn_args) = plan_stdio_spawn("npx", &args, false, |_| {
            panic!("resolver must not be consulted on non-Windows")
        });
        assert_eq!(program, OsString::from("npx"));
        assert_eq!(spawn_args, args);
    }

    #[test]
    fn stdio_path_override_matches_path_case_insensitively() {
        let mk = |name: &str, value: &str| acp::EnvVariable::new(name, value);

        let env = vec![mk("FOO", "bar"), mk("Path", r"C:\node")];
        assert_eq!(stdio_path_override(&env), Some(r"C:\node"));

        let env_upper = vec![mk("PATH", "/custom/bin")];
        assert_eq!(stdio_path_override(&env_upper), Some("/custom/bin"));

        let env_none = vec![mk("FOO", "bar")];
        assert_eq!(stdio_path_override(&env_none), None);
    }

    #[test]
    fn is_figma_mcp_matches_name_and_host() {
        assert!(is_figma_mcp("figma", "https://example.com/mcp"));
        assert!(is_figma_mcp("Figma", "https://example.com/mcp"));
        assert!(is_figma_mcp("other", "https://mcp.figma.com/mcp"));
        assert!(is_figma_mcp("other", "https://figma.com/mcp"));
        assert!(!is_figma_mcp("linear", "https://mcp.linear.app/mcp"));
        assert!(!is_figma_mcp("figma_extra", "https://example.com/mcp"));
        assert!(!is_figma_mcp("linear", "not-a-url"));
        assert!(!is_figma_mcp("linear", "https://notfigma.com/mcp"));
        assert!(!is_figma_mcp("linear", "https://figma.com.evil/mcp"));
    }

    #[test]
    fn ensure_figma_user_agent_sets_kigi_cli_when_missing() {
        let mut headers = reqwest::header::HeaderMap::new();
        ensure_figma_user_agent(&mut headers, "figma", "https://mcp.figma.com/mcp");
        assert_eq!(
            headers.get(reqwest::header::USER_AGENT).unwrap(),
            "kigi-cli"
        );

        let mut host_only = reqwest::header::HeaderMap::new();
        ensure_figma_user_agent(&mut host_only, "other", "https://mcp.figma.com/mcp");
        assert_eq!(
            host_only.get(reqwest::header::USER_AGENT).unwrap(),
            "kigi-cli"
        );
    }

    #[test]
    fn ensure_figma_user_agent_does_not_overwrite_existing() {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert(
            reqwest::header::USER_AGENT,
            reqwest::header::HeaderValue::from_static("custom-ua"),
        );
        ensure_figma_user_agent(&mut headers, "figma", "https://mcp.figma.com/mcp");
        assert_eq!(
            headers.get(reqwest::header::USER_AGENT).unwrap(),
            "custom-ua"
        );
    }

    #[test]
    fn ensure_figma_user_agent_skips_non_figma() {
        let mut headers = reqwest::header::HeaderMap::new();
        ensure_figma_user_agent(&mut headers, "linear", "https://mcp.linear.app/mcp");
        assert!(!headers.contains_key(reqwest::header::USER_AGENT));

        let mut invalid_url = reqwest::header::HeaderMap::new();
        ensure_figma_user_agent(&mut invalid_url, "linear", "not-a-url");
        assert!(!invalid_url.contains_key(reqwest::header::USER_AGENT));
    }

    #[cfg(unix)]
    #[test]
    fn safe_stdio_child_drop_without_entered_runtime_reaps_child() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("test runtime");

        let (transport, pid) = rt.block_on(async {
            let mut cmd = Command::new("sleep");
            cmd.arg("30").kill_on_drop(true);
            kigi_tools::util::detach_command(&mut cmd);
            let (transport, _stderr) = SafeTokioChildProcess::spawn(
                cmd,
                "test".to_string(),
                kigi_file_utils::events::EventWriter::noop(),
            )
            .expect("spawn test child");
            let pid = transport.id().expect("spawned child pid");
            (transport, pid)
        });

        drop(rt);
        drop(transport);

        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
        while std::time::Instant::now() < deadline {
            if !unix_process_exists(pid) {
                return;
            }
            std::thread::sleep(std::time::Duration::from_millis(25));
        }

        panic!("MCP child process {pid} was not reaped after no-runtime drop");
    }

    #[cfg(unix)]
    fn unix_process_exists(pid: u32) -> bool {
        let result = unsafe { libc::kill(pid as libc::pid_t, 0) };
        if result == 0 {
            return true;
        }
        std::io::Error::last_os_error().raw_os_error() != Some(libc::ESRCH)
    }

    #[test]
    fn test_mcp_state_new() {
        let configs = vec![make_stdio_server("test", "/bin/test")];
        let state = McpState::new(configs.clone());

        assert_eq!(state.configs.len(), 1);
        assert!(state.owned_clients.is_empty());
        assert!(!state.is_initialized());
        assert!(!state.is_initializing());
        assert!(!state.has_finished_init());
        assert!(matches!(state.init_progress(), InitProgress::NotStarted));
        assert_eq!(state.generation, 0);
    }

    #[test]
    fn test_mcp_state_update_configs_returns_false_when_unchanged() {
        let configs = vec![make_stdio_server("test", "/bin/test")];
        let mut state = McpState::new(configs.clone());

        // Same configs should return false
        let changed = state.update_configs(configs.clone());
        assert!(!changed);
        assert_eq!(state.generation, 0); // Generation should not change
    }

    #[test]
    fn test_mcp_state_update_configs_returns_true_when_changed() {
        let configs = vec![make_stdio_server("test", "/bin/test")];
        let mut state = McpState::new(configs);

        // Different configs should return true
        let new_configs = vec![make_stdio_server("test2", "/bin/test2")];
        let changed = state.update_configs(new_configs);
        assert!(changed);
        assert_eq!(state.generation, 1); // Generation should increment
    }

    #[test]
    fn test_mcp_state_update_configs_resets_initialized() {
        let configs = vec![make_stdio_server("test", "/bin/test")];
        let mut state = McpState::new(configs);
        // Drive the state machine into Finished{handshaking:{"a"}} so
        // the reset path has both the lifecycle flag AND a per-server
        // entry to clear.
        assert!(state.try_start_init());
        state.mark_servers_initializing(["a".to_string()]);
        state.finish_init();
        assert!(state.has_finished_init());
        assert!(state.is_server_handshaking("a"));

        let new_configs = vec![make_stdio_server("test2", "/bin/test2")];
        let changed = state.update_configs(new_configs);
        assert!(changed);
        // update_configs must drop us back to NotStarted — neither
        // lifecycle flag set nor any per-server progress carried over.
        assert!(!state.is_initialized());
        assert!(!state.is_initializing());
        assert!(!state.has_finished_init());
        assert!(matches!(state.init_progress(), InitProgress::NotStarted));
    }

    #[tokio::test]
    async fn acp_servers_survive_update_configs_clear() {
        use crate::acp_transport::AcpReverseInvoker;
        use std::time::Duration;

        struct NoopInvoker;
        #[async_trait::async_trait]
        impl AcpReverseInvoker for NoopInvoker {
            async fn invoke(
                &self,
                _server_id: &str,
                _message: serde_json::Value,
                _timeout: Duration,
            ) -> Result<serde_json::Value, String> {
                Ok(serde_json::Value::Null)
            }
        }

        let mut state = McpState::new(vec![make_http_server("http-srv", "http://localhost")]);
        state.set_acp_servers(
            vec![AcpServerEntry {
                name: "sdk-tools".to_string(),
                server_id: "srv_0".to_string(),
            }],
            Arc::new(NoopInvoker),
        );
        assert!(state.has_acp_servers());
        assert_eq!(state.build_pending_acp_clients(&HashMap::new()).len(), 1);

        // A config change clears owned clients/configs (proven by the generation bump)
        // but must NOT drop the separately-held acp servers — otherwise the in-process
        // SDK tools would silently vanish on every `update_configs`.
        let changed = state.update_configs(vec![make_http_server("other", "http://other")]);
        assert!(changed);
        assert_eq!(state.generation, 1);
        assert!(
            state.has_acp_servers(),
            "acp servers must survive update_configs"
        );
        let pending = state.build_pending_acp_clients(&HashMap::new());
        assert_eq!(pending.len(), 1, "acp clients rebuild after the clear");
        assert_eq!(pending[0].server_name(), "sdk-tools");
    }

    #[tokio::test]
    async fn acp_overrides_apply_to_built_clients() {
        use crate::acp_transport::AcpReverseInvoker;
        use std::time::Duration;

        struct NoopInvoker;
        #[async_trait::async_trait]
        impl AcpReverseInvoker for NoopInvoker {
            async fn invoke(
                &self,
                _server_id: &str,
                _message: serde_json::Value,
                _timeout: Duration,
            ) -> Result<serde_json::Value, String> {
                Ok(serde_json::Value::Null)
            }
        }

        let mut overrides = HashMap::new();
        overrides.insert(
            "sdk-tools".to_string(),
            McpClientTimeoutOverrides {
                tool_timeout_sec: Some(123),
                ..Default::default()
            },
        );

        let mut state = McpState::new(vec![]);
        state.set_acp_servers(
            vec![AcpServerEntry {
                name: "sdk-tools".to_string(),
                server_id: "srv_0".to_string(),
            }],
            Arc::new(NoopInvoker),
        );

        let pending = state.build_pending_acp_clients(&overrides);
        assert_eq!(pending.len(), 1);
        assert_eq!(
            pending[0].tool_timeout_sec(),
            123,
            "config.toml tool_timeout_sec override must reach the SDK client"
        );
    }

    /// In-process SDK (ACP) clients must never get a liveness watcher: the
    /// dispatcher can't recover them (no `configs` entry), so a proactive
    /// `TransportClosed` would evict the client with no recovery. Guards both
    /// the `is_acp` predicate (across transports) and the `arm_liveness_watcher`
    /// self-gate that depends on it. HTTP/stdio must report `false` so they
    /// keep their watchers.
    #[tokio::test]
    async fn acp_clients_are_not_liveness_watched() {
        use crate::acp_transport::AcpReverseInvoker;
        use std::time::Duration;

        struct NoopInvoker;
        #[async_trait::async_trait]
        impl AcpReverseInvoker for NoopInvoker {
            async fn invoke(
                &self,
                _server_id: &str,
                _message: serde_json::Value,
                _timeout: Duration,
            ) -> Result<serde_json::Value, String> {
                Ok(serde_json::Value::Null)
            }
        }

        let acp = McpClient::new_acp(
            "sdk".to_string(),
            "srv_0".to_string(),
            Arc::new(NoopInvoker),
            None,
            None,
        );
        assert!(acp.is_acp());
        assert!(!acp.is_http());

        let http = McpClient::new_http(
            "http".to_string(),
            HttpConfig {
                url: "http://localhost/api/mcp".to_string(),
                headers: vec![],
            },
            None,
            None,
        );
        assert!(!http.is_acp());

        // Stub stands in for a no-transport / Stdio client (reconnect = None).
        assert!(!McpClient::stub("stdio").is_acp());

        // The gate that prevents the evict-on-close bug: arming is a no-op for ACP.
        assert!(
            !Arc::new(acp)
                .arm_liveness_watcher(Duration::from_millis(500))
                .await
        );
    }

    #[test]
    fn test_mark_servers_initializing_clears_prior_init_failure() {
        // A server that failed a previous init is recorded in `init_failed`
        // (so the status snapshot reports it Unavailable). Starting a fresh
        // init attempt for that server must clear the failure flag so a
        // successful retry can surface as Ready again.
        let mut state = McpState::new(vec![make_stdio_server("a", "/bin/a")]);
        state.init_failed.insert("a".to_string(), String::new());
        state.init_failed.insert("b".to_string(), String::new());

        state.mark_servers_initializing(["a".to_string()]);

        assert!(
            !state.init_failed.contains_key("a"),
            "fresh init attempt must clear the prior failure for that server",
        );
        assert!(
            state.init_failed.contains_key("b"),
            "servers not in this init attempt must keep their failure flag",
        );
    }

    #[test]
    fn test_record_init_failure_keeps_auth_and_init_failed_disjoint() {
        let mut state = McpState::new(vec![make_stdio_server("a", "/bin/a")]);

        // Auth failures are owned by `auth_required` only — never `init_failed` —
        // so a later successful authentication (which clears `auth_required` and
        // registers tools) is not left stuck as Unavailable with zero tools.
        state.record_init_failure("auth-srv", true, None);
        assert!(state.auth_required.contains("auth-srv"));
        assert!(
            !state.init_failed.contains_key("auth-srv"),
            "auth-required failures must not also be flagged init_failed",
        );

        // Non-auth failures (handshake/`tools/list` error or timeout) → init_failed,
        // and their cause is retained for the model-facing reminder.
        state.record_init_failure(
            "dead-srv",
            false,
            Some("tools/list failed: boom".to_string()),
        );
        assert!(!state.auth_required.contains("dead-srv"));
        assert_eq!(
            state.init_failed.get("dead-srv").map(String::as_str),
            Some("tools/list failed: boom"),
        );

        // A fresh init attempt clears the failure entry and its cause.
        state.mark_servers_initializing(["dead-srv".to_string()]);
        assert!(!state.init_failed.contains_key("dead-srv"));
    }

    #[test]
    fn test_clear_init_failed_removes_entry() {
        let mut state = McpState::new(vec![make_stdio_server("a", "/bin/a")]);
        state.record_init_failure("dead-srv", false, Some("boom".to_string()));
        assert!(state.init_failed.contains_key("dead-srv"));

        // Symmetric with record_init_failure: the reactive re-auth path clears
        // a prior failure so a recovered server is not stuck Unavailable.
        state.clear_init_failed("dead-srv");
        assert!(!state.init_failed.contains_key("dead-srv"));
        // Idempotent: clearing an absent entry is a no-op.
        state.clear_init_failed("never-seen");
    }

    #[test]
    fn test_mcp_state_update_configs_increments_generation() {
        let mut state = McpState::new(vec![]);

        // Each change should increment generation
        state.update_configs(vec![make_stdio_server("a", "/bin/a")]);
        assert_eq!(state.generation, 1);

        state.update_configs(vec![make_stdio_server("b", "/bin/b")]);
        assert_eq!(state.generation, 2);

        state.update_configs(vec![make_stdio_server("c", "/bin/c")]);
        assert_eq!(state.generation, 3);
    }

    #[test]
    fn test_mcp_servers_equal_empty_lists() {
        let a: Vec<acp::McpServer> = vec![];
        let b: Vec<acp::McpServer> = vec![];
        assert!(mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_mcp_servers_equal_identical_configs() {
        let a = vec![make_stdio_server("test", "/bin/test")];
        let b = vec![make_stdio_server("test", "/bin/test")];
        assert!(mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_mcp_servers_equal_different_names() {
        let a = vec![make_stdio_server("test1", "/bin/test")];
        let b = vec![make_stdio_server("test2", "/bin/test")];
        assert!(!mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_mcp_servers_equal_different_lengths() {
        let a = vec![make_stdio_server("test", "/bin/test")];
        let b = vec![
            make_stdio_server("test", "/bin/test"),
            make_stdio_server("test2", "/bin/test2"),
        ];
        assert!(!mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_mcp_servers_equal_different_types() {
        let a = vec![make_stdio_server("test", "/bin/test")];
        let b = vec![make_http_server("test", "http://localhost")];
        assert!(!mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_mcp_servers_equal_order_matters() {
        let a = vec![
            make_stdio_server("a", "/bin/a"),
            make_stdio_server("b", "/bin/b"),
        ];
        let b = vec![
            make_stdio_server("b", "/bin/b"),
            make_stdio_server("a", "/bin/a"),
        ];
        // Order matters since we're comparing JSON serialization
        assert!(!mcp_servers_equal(&a, &b));
    }

    #[test]
    fn test_try_start_init_prevents_concurrent_init() {
        let mut state = McpState::new(vec![make_stdio_server("test", "/bin/test")]);

        // First call should succeed
        assert!(state.try_start_init());
        assert!(state.is_initializing());
        assert!(!state.is_initialized());

        // Second call should fail (already initializing)
        assert!(!state.try_start_init());
    }

    #[test]
    fn test_try_start_init_fails_when_initialized() {
        let mut state = McpState::new(vec![make_stdio_server("test", "/bin/test")]);
        // Drive to Finished{empty} via the typed API.
        assert!(state.try_start_init());
        state.finish_init();
        assert!(state.is_initialized());

        // Second `try_start_init` must be rejected: we're already done.
        assert!(!state.try_start_init());
        assert!(!state.is_initializing());
        assert!(state.is_initialized(), "is_initialized stays true");
    }

    #[test]
    fn test_finish_init_clears_initializing() {
        let mut state = McpState::new(vec![make_stdio_server("test", "/bin/test")]);

        state.try_start_init();
        assert!(state.is_initializing());
        assert!(!state.is_initialized());

        state.finish_init();
        assert!(!state.is_initializing());
        assert!(state.is_initialized());
    }

    #[test]
    fn test_cancel_init_clears_initializing() {
        let mut state = McpState::new(vec![make_stdio_server("test", "/bin/test")]);

        state.try_start_init();
        assert!(state.is_initializing());

        state.cancel_init();
        assert!(!state.is_initializing());
        assert!(!state.is_initialized()); // Should NOT be marked as initialized
    }

    #[test]
    fn test_update_configs_resets_initializing() {
        let mut state = McpState::new(vec![make_stdio_server("test", "/bin/test")]);
        state.try_start_init();
        assert!(state.is_initializing());

        // Updating configs should reset initializing flag
        state.update_configs(vec![make_stdio_server("test2", "/bin/test2")]);
        assert!(!state.is_initializing());
        assert!(!state.is_initialized());
    }

    #[test]
    fn test_parse_mcp_meta_config_with_tool_timeouts_ms() {
        let meta = serde_json::json!({
            "mcpConfig": {
                "github": {
                    "toolTimeoutMs": 60000,
                    "toolTimeoutsMs": {
                        "create_issue": 120000,
                        "search": 30000
                    }
                }
            }
        })
        .as_object()
        .cloned()
        .unwrap();
        let map = parse_mcp_meta_config(Some(&meta));
        let github = map.get("github").unwrap();
        assert_eq!(github.tool_timeout_ms, Some(60000));
        let tt = github.tool_timeouts_ms.as_ref().unwrap();
        assert_eq!(tt.get("create_issue"), Some(&120000));
        assert_eq!(tt.get("search"), Some(&30000));
    }

    #[test]
    fn test_parse_mcp_meta_config_without_tool_timeouts_ms() {
        let meta = serde_json::json!({
            "mcpConfig": {
                "github": {
                    "toolTimeoutMs": 60000
                }
            }
        })
        .as_object()
        .cloned()
        .unwrap();
        let map = parse_mcp_meta_config(Some(&meta));
        let github = map.get("github").unwrap();
        assert_eq!(github.tool_timeout_ms, Some(60000));
        assert!(github.tool_timeouts_ms.is_none());
        assert!(github.expose_image_base64.is_none());
    }

    /// Locks in the `exposeImageBase64` camelCase wire-format contract.
    #[test]
    fn test_parse_mcp_meta_config_with_expose_image_base64() {
        let meta = serde_json::json!({
            "mcpConfig": {
                "grafana": { "exposeImageBase64": true },
                "linear":  { "exposeImageBase64": false },
            }
        })
        .as_object()
        .cloned()
        .unwrap();
        let map = parse_mcp_meta_config(Some(&meta));
        assert_eq!(map.get("grafana").unwrap().expose_image_base64, Some(true));
        assert_eq!(map.get("linear").unwrap().expose_image_base64, Some(false));
    }

    #[test]
    fn test_tool_timeout_for_returns_per_tool_override() {
        let mut tool_timeouts = HashMap::new();
        tool_timeouts.insert("create_issue".to_string(), 120u64);
        tool_timeouts.insert("search".to_string(), 30u64);

        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(10),
            tool_timeout_sec: Some(60),
            tool_timeouts: Some(tool_timeouts),
            ..Default::default()
        };
        let client = McpClient::new_http(
            "github".to_string(),
            HttpConfig {
                url: String::new(),
                headers: vec![],
            },
            Some(&overrides),
            None,
        );

        // Per-tool overrides
        assert_eq!(client.tool_timeout_for("create_issue"), 120);
        assert_eq!(client.tool_timeout_for("search"), 30);
        // Falls back to server-level default
        assert_eq!(client.tool_timeout_for("list_repos"), 60);
        assert_eq!(client.tool_timeout_for(""), 60);
    }

    #[test]
    fn test_tool_timeout_for_empty_map_returns_default() {
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(10),
            tool_timeout_sec: Some(45),
            ..Default::default()
        };
        let client = McpClient::new_http(
            "test".to_string(),
            HttpConfig {
                url: String::new(),
                headers: vec![],
            },
            Some(&overrides),
            None,
        );

        // All tools should get the server-level default
        assert_eq!(client.tool_timeout_for("any_tool"), 45);
        assert_eq!(client.tool_timeout_sec(), 45);
    }

    #[test]
    fn test_load_timeouts_startup_precedence() {
        // No override -> the standalone default (env/config resolved by the shell).
        assert_eq!(
            McpClient::load_timeouts(None, None).0,
            DEFAULT_STARTUP_TIMEOUT_SECS
        );

        // A per-server `startup_timeout_sec` (injected by the shell) wins over the default...
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(7),
            ..Default::default()
        };
        assert_eq!(McpClient::load_timeouts(Some(&overrides), None).0, 7);

        // ...and `_meta.startup_timeout_ms` wins over that.
        let meta = McpServerMetaConfig {
            startup_timeout_ms: Some(12_000),
            ..Default::default()
        };
        assert_eq!(
            McpClient::load_timeouts(Some(&overrides), Some(&meta)).0,
            12
        );
    }

    #[test]
    fn test_update_configs_diff_no_change() {
        let configs = vec![make_stdio_server("test", "/bin/test")];
        let mut state = McpState::new(configs.clone());
        assert!(state.update_configs_diff(configs).is_none());
        assert_eq!(state.generation, 0);
    }

    #[test]
    fn test_update_configs_diff_added() {
        let configs = vec![make_stdio_server("a", "/bin/a")];
        let mut state = McpState::new(configs);

        let new_configs = vec![
            make_stdio_server("a", "/bin/a"),
            make_stdio_server("b", "/bin/b"),
        ];
        let diff = state
            .update_configs_diff(new_configs)
            .expect("should detect change");
        assert_eq!(diff.retained, vec!["a"]);
        assert_eq!(diff.added, vec!["b"]);
        assert!(diff.removed.is_empty());
        assert_eq!(state.generation, 1);
    }

    #[test]
    fn test_update_configs_diff_removed() {
        let configs = vec![
            make_stdio_server("a", "/bin/a"),
            make_stdio_server("b", "/bin/b"),
        ];
        let mut state = McpState::new(configs);

        let new_configs = vec![make_stdio_server("a", "/bin/a")];
        let diff = state
            .update_configs_diff(new_configs)
            .expect("should detect change");
        assert_eq!(diff.retained, vec!["a"]);
        assert!(diff.added.is_empty());
        assert_eq!(diff.removed, vec!["b"]);
    }

    #[test]
    fn test_update_configs_diff_changed() {
        let configs = vec![make_stdio_server("a", "/bin/a")];
        let mut state = McpState::new(configs);

        let new_configs = vec![make_stdio_server("a", "/bin/a_v2")];
        let diff = state
            .update_configs_diff(new_configs)
            .expect("should detect change");
        assert!(diff.retained.is_empty());
        assert_eq!(diff.added, vec!["a"]);
        assert_eq!(diff.removed, vec!["a"]);
    }

    #[test]
    fn test_update_configs_diff_auth_required_cleanup() {
        let configs = vec![
            make_stdio_server("keep", "/bin/keep"),
            make_stdio_server("remove", "/bin/remove"),
        ];
        let mut state = McpState::new(configs);
        state.auth_required.insert("remove".to_string());
        state.auth_required.insert("keep".to_string());

        let new_configs = vec![make_stdio_server("keep", "/bin/keep")];
        let diff = state
            .update_configs_diff(new_configs)
            .expect("should detect change");
        assert_eq!(diff.retained, vec!["keep"]);
        assert_eq!(diff.removed, vec!["remove"]);
        assert!(state.auth_required.contains("keep"));
        assert!(!state.auth_required.contains("remove"));
    }

    #[test]
    fn test_update_configs_diff_empty_to_nonempty() {
        let mut state = McpState::new(vec![]);
        let new_configs = vec![make_stdio_server("a", "/bin/a")];
        let diff = state
            .update_configs_diff(new_configs)
            .expect("should detect change");
        assert!(diff.retained.is_empty());
        assert_eq!(diff.added, vec!["a"]);
        assert!(diff.removed.is_empty());
    }

    #[test]
    fn test_update_configs_diff_nonempty_to_empty() {
        let configs = vec![make_stdio_server("a", "/bin/a")];
        let mut state = McpState::new(configs);
        let diff = state
            .update_configs_diff(vec![])
            .expect("should detect change");
        assert!(diff.retained.is_empty());
        assert!(diff.added.is_empty());
        assert_eq!(diff.removed, vec!["a"]);
    }

    /// Two MCP servers exposing a tool with the same raw name must produce
    /// `McpErasedTool` instances with **distinct** `ToolId`s (qualified with
    /// the server name). Regression test for a bug where `McpErasedTool::id()`
    /// returned the unqualified name, causing the second registration to
    /// silently overwrite the first in the `LocalRegistry`.
    #[test]
    fn test_mcp_erased_tool_id_is_qualified() {
        use kigi_tool_runtime::Tool;

        let mcp_state = Arc::new(Mutex::new(McpState::new(vec![])));

        let tool_a = McpErasedTool {
            tool: McpTool::new(
                "SearchUsers".to_string(),
                "Search users".to_string(),
                "calendar".to_string(),
                Arc::clone(&mcp_state),
                serde_json::json!({"type": "object"}),
                None,
            ),
        };
        let tool_b = McpErasedTool {
            tool: McpTool::new(
                "SearchUsers".to_string(),
                "Search users".to_string(),
                "teams".to_string(),
                Arc::clone(&mcp_state),
                serde_json::json!({"type": "object"}),
                None,
            ),
        };

        let id_a = tool_a.id();
        let id_b = tool_b.id();

        // IDs must be qualified with the server name.
        assert_eq!(id_a.as_str(), "calendar__SearchUsers");
        assert_eq!(id_b.as_str(), "teams__SearchUsers");

        // And therefore distinct.
        assert_ne!(id_a, id_b);
    }

    /// Registering two MCP tools with the same raw name from different servers
    /// into a `LocalRegistry` must preserve both entries (no silent overwrite).
    #[test]
    fn test_same_raw_name_different_servers_no_local_registry_collision() {
        use kigi_tool_runtime::LocalRegistry;
        use kigi_tool_runtime::Tool;

        let mcp_state = Arc::new(Mutex::new(McpState::new(vec![])));
        let registry = LocalRegistry::new();

        let tool_a = McpErasedTool {
            tool: McpTool::new(
                "SearchUsers".to_string(),
                "Search users on calendar".to_string(),
                "calendar".to_string(),
                Arc::clone(&mcp_state),
                serde_json::json!({"type": "object"}),
                None,
            ),
        };
        let tool_b = McpErasedTool {
            tool: McpTool::new(
                "SearchUsers".to_string(),
                "Search users on teams".to_string(),
                "teams".to_string(),
                Arc::clone(&mcp_state),
                serde_json::json!({"type": "object"}),
                None,
            ),
        };

        let id_a = tool_a.id();
        let id_b = tool_b.id();

        // First registration should not displace anything.
        let displaced_a = registry.register(tool_a);
        assert!(
            displaced_a.is_none(),
            "first registration should not displace"
        );

        // Second registration should also not displace anything (distinct IDs).
        let displaced_b = registry.register(tool_b);
        assert!(
            displaced_b.is_none(),
            "second registration must not overwrite first"
        );

        // Both tools must be independently resolvable.
        assert!(
            registry.find(&id_a).is_some(),
            "calendar tool must be found"
        );
        assert!(registry.find(&id_b).is_some(), "teams tool must be found");
        assert_eq!(registry.len(), 2);
    }

    fn make_test_client(name: &str) -> Arc<McpClient> {
        // Same shape as the no-transport placeholder.
        Arc::new(McpClient::stub(name))
    }

    #[test]
    fn test_shared_mcp_pool_from_empty_state() {
        let state = McpState::new(vec![]);
        let pool = SharedMcpPool::from_state(&state);
        assert_eq!(pool.len(), 0);
        assert_eq!(pool.server_names().count(), 0);
        assert!(pool.configs().is_empty());
        assert!(pool.meta_config_map().is_empty());
        assert!(pool.get_client("anything").is_none());
    }

    #[test]
    fn test_shared_mcp_pool_len_matches_client_count() {
        let mut state = McpState::new(vec![]);
        for name in ["alpha", "beta", "gamma"] {
            state
                .owned_clients
                .insert(name.to_string(), make_test_client(name));
        }
        let pool = SharedMcpPool::from_state(&state);
        assert_eq!(pool.len(), 3);
        assert_eq!(pool.len(), pool.server_names().count());
    }

    #[test]
    fn test_shared_mcp_pool_snapshot_shares_arc_clients() {
        let mut state = McpState::new(vec![make_stdio_server("github", "/bin/gh")]);
        let client = make_test_client("github");
        state
            .owned_clients
            .insert("github".to_string(), Arc::clone(&client));

        let pool = SharedMcpPool::from_state(&state);
        let pool_client = pool.get_client("github").expect("should find client");

        // Must point to the same allocation (shared transport)
        assert!(Arc::ptr_eq(&client, pool_client));
    }

    #[test]
    fn test_shared_mcp_pool_get_client_missing() {
        let mut state = McpState::new(vec![]);
        state
            .owned_clients
            .insert("a".to_string(), make_test_client("a"));
        let pool = SharedMcpPool::from_state(&state);

        assert!(pool.get_client("a").is_some());
        assert!(pool.get_client("nonexistent").is_none());
        assert!(pool.get_client("").is_none());
    }

    #[test]
    fn test_shared_mcp_pool_server_names() {
        let mut state = McpState::new(vec![]);
        for name in ["alpha", "beta", "gamma"] {
            state
                .owned_clients
                .insert(name.to_string(), make_test_client(name));
        }

        let pool = SharedMcpPool::from_state(&state);
        let mut names: Vec<&str> = pool.server_names().collect();
        names.sort();
        assert_eq!(names, vec!["alpha", "beta", "gamma"]);
    }

    #[test]
    fn test_shared_mcp_pool_snapshot_independent_of_state_mutations() {
        let mut state = McpState::new(vec![make_stdio_server("srv", "/bin/srv")]);
        state
            .owned_clients
            .insert("srv".to_string(), make_test_client("srv"));

        let pool = SharedMcpPool::from_state(&state);

        // Mutate state after snapshot
        state.owned_clients.clear();
        state.configs.clear();

        // Pool retains original data
        assert_eq!(pool.server_names().count(), 1);
        assert!(pool.get_client("srv").is_some());
        assert_eq!(pool.configs().len(), 1);
    }

    #[test]
    fn test_shared_mcp_pool_meta_config_preserved() {
        let mut meta = McpMetaConfigMap::new();
        meta.insert(
            "github".to_string(),
            McpServerMetaConfig {
                startup_timeout_ms: Some(5000),
                tool_timeout_ms: Some(120000),
                tool_timeouts_ms: None,
                expose_image_base64: None,
            },
        );
        let state =
            McpState::new_with_meta(vec![make_http_server("github", "http://gh.local")], meta);
        let pool = SharedMcpPool::from_state(&state);

        let mc = pool
            .meta_config_map()
            .get("github")
            .expect("should have meta config");
        assert_eq!(mc.startup_timeout_ms, Some(5000));
        assert_eq!(mc.tool_timeout_ms, Some(120000));
    }

    #[test]
    fn test_shared_mcp_pool_clone_shares_arcs() {
        let mut state = McpState::new(vec![]);
        let client = make_test_client("svc");
        state
            .owned_clients
            .insert("svc".to_string(), Arc::clone(&client));

        let pool = SharedMcpPool::from_state(&state);
        let pool2 = pool.clone();

        // Both clones share the same Arc<McpClient>
        let c1 = pool.get_client("svc").unwrap();
        let c2 = pool2.get_client("svc").unwrap();
        assert!(Arc::ptr_eq(c1, c2));
    }

    // ── owned/shared split behavioral tests ─────────────────────────

    #[test]
    fn test_get_client_owned_overrides_shared() {
        let mut state = McpState::new(vec![]);
        let shared = make_test_client("srv");
        let owned = make_test_client("srv");
        state
            .shared_clients
            .insert("srv".to_string(), Arc::clone(&shared));
        state
            .owned_clients
            .insert("srv".to_string(), Arc::clone(&owned));

        let got = state.get_client("srv").unwrap();
        assert!(Arc::ptr_eq(got, &owned));
        assert!(!Arc::ptr_eq(got, &shared));
    }

    #[test]
    fn test_get_client_falls_through_to_shared() {
        let mut state = McpState::new(vec![]);
        let shared = make_test_client("srv");
        state
            .shared_clients
            .insert("srv".to_string(), Arc::clone(&shared));

        let got = state.get_client("srv").unwrap();
        assert!(Arc::ptr_eq(got, &shared));
        assert!(state.get_client("missing").is_none());
    }

    #[test]
    fn test_all_clients_deduplicates_shared_by_owned() {
        let mut state = McpState::new(vec![]);
        state
            .owned_clients
            .insert("a".to_string(), make_test_client("a"));
        state
            .shared_clients
            .insert("a".to_string(), make_test_client("a-shared"));
        state
            .shared_clients
            .insert("b".to_string(), make_test_client("b-shared"));

        let all: Vec<_> = state.all_clients().map(|(n, _)| n.as_str()).collect();
        // "a" appears once (from owned), "b" from shared
        assert_eq!(all.iter().filter(|&&n| n == "a").count(), 1);
        assert!(all.contains(&"b"));
        assert_eq!(all.len(), 2);

        // The "a" entry must be the owned client, not the shared one
        let (_, a_client) = state.all_clients().find(|(n, _)| *n == "a").unwrap();
        assert!(Arc::ptr_eq(a_client, state.owned_clients.get("a").unwrap()));
    }

    #[test]
    fn test_import_shared_clients_skips_config_collisions() {
        // Child has a config entry named "github" — importing a shared
        // client with the same name must be skipped.
        let mut state = McpState::new(vec![make_stdio_server("github", "/bin/gh")]);
        let mut pool_clients = HashMap::new();
        pool_clients.insert("github".to_string(), make_test_client("github"));
        pool_clients.insert("linear".to_string(), make_test_client("linear"));
        let pool = SharedMcpPool {
            clients: pool_clients,
            configs: vec![],
            meta_config_map: McpMetaConfigMap::new(),
        };

        state.import_shared_clients(&pool);

        assert!(
            !state.shared_clients.contains_key("github"),
            "github should be skipped — collides with child config"
        );
        assert!(
            state.shared_clients.contains_key("linear"),
            "linear should be imported — no collision"
        );
    }

    #[test]
    fn test_update_configs_preserves_shared_clients() {
        let mut state = McpState::new(vec![make_stdio_server("old", "/bin/old")]);
        state
            .owned_clients
            .insert("old".to_string(), make_test_client("old"));
        let shared = make_test_client("inherited");
        state
            .shared_clients
            .insert("inherited".to_string(), Arc::clone(&shared));

        let changed = state.update_configs(vec![make_stdio_server("new", "/bin/new")]);

        assert!(changed);
        assert!(state.owned_clients.is_empty(), "owned should be cleared");
        assert_eq!(state.shared_clients.len(), 1, "shared should be untouched");
        assert!(Arc::ptr_eq(
            state.shared_clients.get("inherited").unwrap(),
            &shared
        ));
    }

    #[test]
    fn test_update_configs_diff_preserves_shared_clients() {
        let mut state = McpState::new(vec![
            make_stdio_server("keep", "/bin/keep"),
            make_stdio_server("drop", "/bin/drop"),
        ]);
        state
            .owned_clients
            .insert("keep".to_string(), make_test_client("keep"));
        state
            .owned_clients
            .insert("drop".to_string(), make_test_client("drop"));
        let shared = make_test_client("inherited");
        state
            .shared_clients
            .insert("inherited".to_string(), Arc::clone(&shared));

        // New config removes "drop", keeps "keep"
        let diff = state
            .update_configs_diff(vec![make_stdio_server("keep", "/bin/keep")])
            .expect("configs changed");

        assert!(diff.removed.contains(&"drop".to_string()));
        assert!(diff.retained.contains(&"keep".to_string()));
        assert!(!state.owned_clients.contains_key("drop"));
        assert!(state.owned_clients.contains_key("keep"));
        // Shared clients must be completely untouched
        assert!(Arc::ptr_eq(
            state.shared_clients.get("inherited").unwrap(),
            &shared
        ));
    }

    #[test]
    fn test_from_state_captures_both_owned_and_shared() {
        let mut state = McpState::new(vec![]);
        let owned = make_test_client("owned-srv");
        let shared = make_test_client("shared-srv");
        state
            .owned_clients
            .insert("owned-srv".to_string(), Arc::clone(&owned));
        state
            .shared_clients
            .insert("shared-srv".to_string(), Arc::clone(&shared));

        let pool = SharedMcpPool::from_state(&state);

        assert!(Arc::ptr_eq(pool.get_client("owned-srv").unwrap(), &owned));
        assert!(Arc::ptr_eq(pool.get_client("shared-srv").unwrap(), &shared));
        assert_eq!(pool.server_names().count(), 2);
    }

    #[test]
    fn test_retain_clients_keeps_matching() {
        let mut state = McpState::new(vec![]);
        for name in ["github", "linear", "slack"] {
            state
                .owned_clients
                .insert(name.to_string(), make_test_client(name));
        }
        let mut pool = SharedMcpPool::from_state(&state);

        pool.retain_clients(|name| name == "github" || name == "slack");

        assert!(pool.get_client("github").is_some());
        assert!(pool.get_client("slack").is_some());
        assert!(pool.get_client("linear").is_none());
        assert_eq!(pool.server_names().count(), 2);
    }

    #[test]
    fn test_retain_clients_remove_all() {
        let mut state = McpState::new(vec![]);
        state
            .owned_clients
            .insert("srv".to_string(), make_test_client("srv"));
        let mut pool = SharedMcpPool::from_state(&state);

        pool.retain_clients(|_| false);

        assert_eq!(pool.server_names().count(), 0);
        assert!(pool.get_client("srv").is_none());
    }

    #[test]
    fn test_retain_clients_keep_all() {
        let mut state = McpState::new(vec![]);
        for name in ["a", "b", "c"] {
            state
                .owned_clients
                .insert(name.to_string(), make_test_client(name));
        }
        let mut pool = SharedMcpPool::from_state(&state);

        pool.retain_clients(|_| true);

        assert_eq!(pool.server_names().count(), 3);
    }

    #[test]
    fn test_retain_clients_preserves_arc_identity() {
        let mut state = McpState::new(vec![]);
        let client = make_test_client("keep");
        state
            .owned_clients
            .insert("keep".to_string(), Arc::clone(&client));
        state
            .owned_clients
            .insert("drop".to_string(), make_test_client("drop"));
        let mut pool = SharedMcpPool::from_state(&state);

        pool.retain_clients(|name| name == "keep");

        assert!(Arc::ptr_eq(pool.get_client("keep").unwrap(), &client));
    }

    fn make_mcp_tool(server_name: &str, name: &str) -> McpTool {
        McpTool::new(
            name.to_string(),
            "test desc".to_string(),
            server_name.to_string(),
            Arc::new(Mutex::new(McpState::new(vec![]))),
            serde_json::json!({}),
            None,
        )
    }

    #[test]
    fn into_registration_accepts_well_formed_segments() {
        // Positive guard: the count check rejects `__`-anywhere-but-the-delimiter
        // names but must not reject legitimate ones. If the rejection rule is
        // ever tightened too far, this test breaks before any of the negative
        // cases below.
        let tool = make_mcp_tool("linear", "list_issues");
        let reg = tool.into_registration().expect("should register");
        assert_eq!(reg.name, "linear__list_issues");
    }

    #[test]
    fn into_registration_rejects_boundary_ambiguity() {
        // `"foo_"` + `"__"` + `"_bar"` => `"foo___bar"` has two valid
        // `__` positions (indices 3 and 4), so `split_once("__")` would
        // misparse it as `("foo", "_bar")` and silently auto-allow a
        // future legitimate `"foo"` server. The naïve per-segment check
        // (each side individually has no `__`) misses this — the count
        // check catches it. Same rule also rejects "__-in-segment" cases
        // (`"weird__server"` + `"list"`, `"linear"` + `"my__weird__tool"`)
        // which are covered by the same `count() != 1` line of code.
        let tool = make_mcp_tool("foo_", "_bar");
        assert!(tool.into_registration().is_none());
    }

    // ── is_retriable_transport_error tests ───────────────────────────

    #[test]
    fn test_is_retriable_transport_closed() {
        assert!(is_retriable_transport_error(&ServiceError::TransportClosed));
    }

    #[test]
    fn test_is_retriable_transport_send() {
        let err = ServiceError::TransportSend(rmcp::transport::DynamicTransportError::from_parts(
            "test",
            std::any::TypeId::of::<()>(),
            Box::new(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "connection reset",
            )),
        ));
        assert!(is_retriable_transport_error(&err));
    }

    #[test]
    fn test_not_retriable_unexpected_response() {
        assert!(!is_retriable_transport_error(
            &ServiceError::UnexpectedResponse
        ));
    }

    #[test]
    fn test_not_retriable_cancelled() {
        assert!(!is_retriable_transport_error(&ServiceError::Cancelled {
            reason: Some("shutdown".to_string()),
        }));
    }

    #[test]
    fn test_not_retriable_timeout() {
        assert!(!is_retriable_transport_error(&ServiceError::Timeout {
            timeout: std::time::Duration::from_secs(30),
        }));
    }

    fn mcp_service_err(code: i32) -> ServiceError {
        ServiceError::McpError(rmcp::ErrorData::new(
            rmcp::model::ErrorCode(code),
            "boom",
            None,
        ))
    }

    #[test]
    fn should_recover_mcp_error_recovers_everything_outside_excluded_set() {
        assert!(should_recover_mcp_error(-32603));
        assert!(should_recover_mcp_error(-32002));
        assert!(should_recover_mcp_error(-32000));
        assert!(should_recover_mcp_error(-32099));
        assert!(should_recover_mcp_error(-32100));
        assert!(should_recover_mcp_error(0));
        assert!(should_recover_mcp_error(1));
        assert!(should_recover_mcp_error(i32::MIN));
        assert!(should_recover_mcp_error(i32::MAX));
    }

    #[test]
    fn should_recover_mcp_error_skips_deterministic_client_errors() {
        assert!(!should_recover_mcp_error(-32700));
        assert!(!should_recover_mcp_error(-32600));
        assert!(!should_recover_mcp_error(-32601));
        assert!(!should_recover_mcp_error(-32602));
    }

    #[test]
    fn should_recover_service_error_http_mcperror_recoverable() {
        assert!(should_recover_service_error(
            &mcp_service_err(-32603),
            true,
            false,
        ));
    }

    #[test]
    fn should_recover_service_error_http_mcperror_invalid_params_skipped() {
        assert!(!should_recover_service_error(
            &mcp_service_err(-32602),
            true,
            false,
        ));
    }

    #[test]
    fn should_recover_service_error_stdio_mcperror_not_recovered() {
        assert!(!should_recover_service_error(
            &mcp_service_err(-32603),
            false,
            false,
        ));
    }

    #[test]
    fn should_recover_service_error_mcperror_at_most_once_per_dispatch() {
        assert!(!should_recover_service_error(
            &mcp_service_err(-32603),
            true,
            true,
        ));
    }

    #[test]
    fn should_recover_service_error_http_mcperror_auth_rejection_not_recovered() {
        let auth_err = ServiceError::McpError(rmcp::ErrorData::new(
            rmcp::model::ErrorCode(-32603),
            "Unauthorized: token expired",
            None,
        ));
        assert!(!should_recover_service_error(&auth_err, true, false));
        let session_err = ServiceError::McpError(rmcp::ErrorData::new(
            rmcp::model::ErrorCode(-32603),
            "session not found",
            None,
        ));
        assert!(should_recover_service_error(&session_err, true, false));
    }

    #[test]
    fn should_recover_service_error_transport_errors_always_recover() {
        assert!(should_recover_service_error(
            &ServiceError::TransportClosed,
            true,
            false
        ));
        assert!(should_recover_service_error(
            &ServiceError::TransportClosed,
            false,
            false
        ));
        assert!(should_recover_service_error(
            &ServiceError::TransportClosed,
            true,
            true
        ));
    }

    #[test]
    fn should_recover_service_error_other_non_transport_not_recovered() {
        assert!(!should_recover_service_error(
            &ServiceError::UnexpectedResponse,
            true,
            false
        ));
        assert!(!should_recover_service_error(
            &ServiceError::Timeout {
                timeout: std::time::Duration::from_secs(30),
            },
            true,
            false
        ));
    }

    #[tokio::test]
    async fn recover_and_retry_surfaces_original_error_when_recover_fails() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(1),
            ..Default::default()
        };
        let client = Arc::new(McpClient::new_http(
            "wedged".to_string(),
            config,
            Some(&overrides),
            None,
        ));

        let tool = McpErasedTool {
            tool: McpTool::new(
                "do_thing".to_string(),
                "desc".to_string(),
                "wedged".to_string(),
                Arc::new(Mutex::new(McpState::new(vec![]))),
                serde_json::json!({"type": "object"}),
                None,
            ),
        };

        let original = mcp_service_err(-32603);
        let expected = original.to_string();
        let params = CallToolRequestParams::new("do_thing");

        let mut reconnect_attempted = false;
        let mut is_timeout = false;
        let ew = kigi_file_utils::events::EventWriter::noop();

        let err = tool
            .recover_and_retry(
                &client,
                params,
                std::time::Duration::from_secs(1),
                1,
                original,
                &mut reconnect_attempted,
                &mut is_timeout,
                &ew,
            )
            .await
            .expect_err("recover must fail against an unreachable host");

        assert_eq!(err.to_string(), expected, "original error must be surfaced");
        assert!(reconnect_attempted, "reconnect attempt must be flagged");
        assert!(!is_timeout, "a recover failure is not a tool timeout");
    }

    use std::sync::atomic::{AtomicUsize, Ordering};

    #[derive(Clone, Copy)]
    enum CallToolBehavior {
        ErrorThenOk { code: i32 },
        AlwaysError { code: i32 },
        HangThenOk { hang_ms: u64 },
        ErrorThenHang { code: i32, hang_ms: u64 },
    }

    #[derive(Clone)]
    struct FakeMcpState {
        behavior: CallToolBehavior,
        inits: Arc<AtomicUsize>,
        calls: Arc<AtomicUsize>,
    }

    async fn fake_handle_post(
        axum::extract::State(state): axum::extract::State<FakeMcpState>,
        axum::Json(req): axum::Json<serde_json::Value>,
    ) -> axum::response::Response {
        use axum::response::IntoResponse;
        let id = req["id"].clone();
        let ok = || {
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": id.clone(),
                "result": {"content": [{"type": "text", "text": "ok"}], "isError": false},
            })
        };
        let err = |code: i32, msg: String| {
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": id.clone(),
                "error": {"code": code, "message": msg},
            })
        };
        match req["method"].as_str() {
            Some("initialize") => {
                state.inits.fetch_add(1, Ordering::Relaxed);
                let result = serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": id.clone(),
                    "result": {
                        "protocolVersion": req["params"]["protocolVersion"].clone(),
                        "capabilities": {},
                        "serverInfo": {"name": "fake", "version": "0.0.0"},
                    },
                });
                ([("mcp-session-id", "fake-session")], axum::Json(result)).into_response()
            }
            Some("tools/list") => axum::Json(serde_json::json!({
                "jsonrpc": "2.0",
                "id": id.clone(),
                "result": {"tools": [{"name": "echo", "inputSchema": {"type": "object"}}]},
            }))
            .into_response(),
            Some("tools/call") => {
                let n = state.calls.fetch_add(1, Ordering::Relaxed);
                match state.behavior {
                    CallToolBehavior::ErrorThenOk { code } => {
                        if n == 0 {
                            axum::Json(err(code, "session expired".to_string())).into_response()
                        } else {
                            axum::Json(ok()).into_response()
                        }
                    }
                    CallToolBehavior::AlwaysError { code } => {
                        axum::Json(err(code, format!("attempt {}", n + 1))).into_response()
                    }
                    CallToolBehavior::HangThenOk { hang_ms } => {
                        if n == 0 {
                            tokio::time::sleep(std::time::Duration::from_millis(hang_ms)).await;
                        }
                        axum::Json(ok()).into_response()
                    }
                    CallToolBehavior::ErrorThenHang { code, hang_ms } => {
                        if n == 0 {
                            axum::Json(err(code, "session expired".to_string())).into_response()
                        } else {
                            tokio::time::sleep(std::time::Duration::from_millis(hang_ms)).await;
                            axum::Json(ok()).into_response()
                        }
                    }
                }
            }
            _ => axum::http::StatusCode::ACCEPTED.into_response(),
        }
    }

    async fn fake_handle_get() -> axum::response::Response {
        use axum::response::IntoResponse;
        let body = axum::body::Body::from_stream(futures::stream::pending::<
            Result<String, std::io::Error>,
        >());
        (
            [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
            body,
        )
            .into_response()
    }

    async fn spawn_fake_mcp(
        behavior: CallToolBehavior,
    ) -> (String, Arc<AtomicUsize>, Arc<AtomicUsize>) {
        let inits = Arc::new(AtomicUsize::new(0));
        let calls = Arc::new(AtomicUsize::new(0));
        let app = axum::Router::new()
            .route(
                "/mcp",
                axum::routing::get(fake_handle_get).post(fake_handle_post),
            )
            .with_state(FakeMcpState {
                behavior,
                inits: inits.clone(),
                calls: calls.clone(),
            });
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("bind");
        let addr = listener.local_addr().expect("addr");
        tokio::spawn(async move {
            let _ = axum::serve(listener, app).await;
        });
        (format!("http://{addr}/mcp"), inits, calls)
    }

    fn fake_http_client(url: &str, tool_timeout_sec: u64) -> Arc<McpClient> {
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(5),
            tool_timeout_sec: Some(tool_timeout_sec),
            ..Default::default()
        };
        Arc::new(McpClient::new_http(
            "fake".to_string(),
            HttpConfig {
                url: url.to_string(),
                headers: vec![],
            },
            Some(&overrides),
            None,
        ))
    }

    fn fake_echo_tool() -> McpErasedTool {
        McpErasedTool {
            tool: McpTool::new(
                "echo".to_string(),
                "echo desc".to_string(),
                "fake".to_string(),
                Arc::new(Mutex::new(McpState::new(vec![]))),
                serde_json::json!({"type": "object"}),
                None,
            ),
        }
    }

    fn event_types(jsonl: &str) -> Vec<serde_json::Value> {
        jsonl
            .lines()
            .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
            .collect()
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn try_call_tool_http_mcperror_recovers_then_retry_succeeds() {
        let (url, inits, calls) =
            spawn_fake_mcp(CallToolBehavior::ErrorThenOk { code: -32603 }).await;
        let client = fake_http_client(&url, 5);
        let tool = fake_echo_tool();
        let tmp = tempfile::tempdir().unwrap();
        let ew = kigi_file_utils::events::EventWriter::open(tmp.path());

        let mut reconnect = false;
        let mut is_timeout = false;
        let raw = serde_json::json!({});
        let out = tool
            .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew)
            .await
            .expect("recovered call should succeed");

        assert!(
            !out.is_error.unwrap_or(false),
            "retry should return a success result"
        );
        assert!(reconnect, "reconnect_attempted must be set");
        assert!(!is_timeout);
        assert_eq!(
            calls.load(Ordering::Relaxed),
            2,
            "one failed + one retried tools/call"
        );
        assert_eq!(
            inits.load(Ordering::Relaxed),
            2,
            "initial handshake + one recovery re-init"
        );

        let jsonl = std::fs::read_to_string(tmp.path().join("events.jsonl")).unwrap();
        let events = event_types(&jsonl);
        assert!(
            events.iter().any(|e| e["type"] == "mcp_transport_error"),
            "expected mcp_transport_error in {jsonl}"
        );
        assert!(
            events
                .iter()
                .any(|e| e["type"] == "mcp_transport_reconnect" && e["success"] == true),
            "expected a successful mcp_transport_reconnect in {jsonl}"
        );
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn try_call_tool_http_retry_failure_surfaces_retry_error() {
        let (url, _inits, calls) =
            spawn_fake_mcp(CallToolBehavior::AlwaysError { code: -32603 }).await;
        let client = fake_http_client(&url, 5);
        let tool = fake_echo_tool();
        let ew = kigi_file_utils::events::EventWriter::noop();

        let mut reconnect = false;
        let mut is_timeout = false;
        let raw = serde_json::json!({});
        let err = tool
            .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew)
            .await
            .expect_err("both attempts fail");

        let msg = err.to_string();
        assert!(msg.contains("attempt 2"), "want retry error, got: {msg}");
        assert!(
            !msg.contains("attempt 1"),
            "must not surface the original error: {msg}"
        );
        assert!(reconnect);
        assert!(!is_timeout);
        assert_eq!(
            calls.load(Ordering::Relaxed),
            2,
            "one failed + one retried tools/call"
        );
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn try_call_tool_http_invalid_params_not_recovered() {
        let (url, inits, calls) =
            spawn_fake_mcp(CallToolBehavior::AlwaysError { code: -32602 }).await;
        let client = fake_http_client(&url, 5);
        let tool = fake_echo_tool();
        let ew = kigi_file_utils::events::EventWriter::noop();

        let mut reconnect = false;
        let mut is_timeout = false;
        let raw = serde_json::json!({});
        let err = tool
            .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew)
            .await
            .expect_err("invalid params surfaced as-is");

        assert!(err.to_string().contains("attempt 1"), "got: {err}");
        assert!(!reconnect, "invalid-params must not trigger recovery");
        assert!(!is_timeout);
        assert_eq!(calls.load(Ordering::Relaxed), 1, "no retry POST");
        assert_eq!(inits.load(Ordering::Relaxed), 1, "no recovery re-init");
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn try_call_tool_http_outer_timeout_resets_transport_no_retry() {
        let (url, inits, calls) =
            spawn_fake_mcp(CallToolBehavior::HangThenOk { hang_ms: 3000 }).await;
        let client = fake_http_client(&url, 1);
        let tool = fake_echo_tool();
        let ew = kigi_file_utils::events::EventWriter::noop();

        let mut reconnect = false;
        let mut is_timeout = false;
        let raw = serde_json::json!({});
        let err = tool
            .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew)
            .await
            .expect_err("call must time out");

        assert!(err.to_string().contains("timed out"), "got: {err}");
        assert!(is_timeout, "is_timeout must be set");
        assert!(reconnect, "timeout arm flags the reconnect after resetting");
        assert_eq!(
            calls.load(Ordering::Relaxed),
            1,
            "timed-out call is NOT retried"
        );
        assert!(matches!(
            client.state_kind().await,
            ClientStateKind::Pending
        ));
        assert_eq!(
            inits.load(Ordering::Relaxed),
            1,
            "no re-init during the timed-out dispatch"
        );

        let mut reconnect2 = false;
        let mut is_timeout2 = false;
        let out = tool
            .try_call_tool(&client, &raw, &mut reconnect2, &mut is_timeout2, &ew)
            .await
            .expect("second dispatch should re-init and succeed");
        assert!(!out.is_error.unwrap_or(false));
        assert!(!is_timeout2);
        assert_eq!(
            inits.load(Ordering::Relaxed),
            2,
            "second dispatch re-initialized the session"
        );
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn try_call_tool_http_retry_timeout_surfaces_timeout() {
        let (url, inits, calls) = spawn_fake_mcp(CallToolBehavior::ErrorThenHang {
            code: -32603,
            hang_ms: 3000,
        })
        .await;
        let client = fake_http_client(&url, 1);
        let tool = fake_echo_tool();
        let ew = kigi_file_utils::events::EventWriter::noop();

        let mut reconnect = false;
        let mut is_timeout = false;
        let raw = serde_json::json!({});
        let err = tool
            .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew)
            .await
            .expect_err("the retried call must time out");

        assert!(err.to_string().contains("timed out"), "got: {err}");
        assert!(is_timeout, "retry-timeout must set is_timeout");
        assert!(reconnect, "recovery was attempted");
        assert_eq!(
            calls.load(Ordering::Relaxed),
            2,
            "the retry tools/call was attempted"
        );
        assert_eq!(
            inits.load(Ordering::Relaxed),
            2,
            "recovery re-initialized before the retry"
        );
    }

    // ── new_http stores http_config tests ────────────────────────────

    #[test]
    fn test_new_http_stores_http_config() {
        let config = HttpConfig {
            url: "http://localhost:5000/api/mcp".to_string(),
            headers: vec![("x-token".to_string(), "abc".to_string())],
        };
        let client = McpClient::new_http("example-mcp".to_string(), config, None, None);
        let stored = client
            .http_config
            .as_ref()
            .expect("http_config should be Some");
        assert_eq!(stored.url, "http://localhost:5000/api/mcp");
        assert_eq!(stored.headers.len(), 1);
        assert_eq!(stored.headers[0].0, "x-token");
    }

    #[test]
    fn test_new_stdio_has_no_http_config() {
        // Stdio clients must NOT have http_config — they can't reconnect via HTTP.
        let client = McpClient::stub("stdio-srv");
        assert!(client.http_config.is_none());
    }

    // ── http_headers_match / refresh_managed_clients guard tests ─────

    #[test]
    fn http_headers_match_compares_full_set_order_insensitively() {
        let config = HttpConfig {
            url: "http://localhost:5000/api/mcp".to_string(),
            headers: vec![
                ("authorization".to_string(), "Bearer t".to_string()),
                ("x-scope".to_string(), "read".to_string()),
            ],
        };
        let client = McpClient::new_http("managed".to_string(), config, None, None);

        let equal: HashMap<String, String> = [
            ("x-scope".to_string(), "read".to_string()),
            ("authorization".to_string(), "Bearer t".to_string()),
        ]
        .into_iter()
        .collect();
        assert!(client.http_headers_match(&equal));

        let changed_value: HashMap<String, String> = [
            ("authorization".to_string(), "Bearer NEW".to_string()),
            ("x-scope".to_string(), "read".to_string()),
        ]
        .into_iter()
        .collect();
        assert!(!client.http_headers_match(&changed_value));

        let missing_key: HashMap<String, String> =
            [("authorization".to_string(), "Bearer t".to_string())]
                .into_iter()
                .collect();
        assert!(!client.http_headers_match(&missing_key));
    }

    #[test]
    fn http_headers_match_handles_duplicate_stored_keys() {
        // Duplicate stored key must not mask a missing fresh key by inflating
        // the stored length to match.
        let config = HttpConfig {
            url: "http://localhost:5000/api/mcp".to_string(),
            headers: vec![
                ("authorization".to_string(), "Bearer t".to_string()),
                ("authorization".to_string(), "Bearer t".to_string()),
            ],
        };
        let client = McpClient::new_http("managed".to_string(), config, None, None);

        let two_distinct: HashMap<String, String> = [
            ("authorization".to_string(), "Bearer t".to_string()),
            ("x-scope".to_string(), "read".to_string()),
        ]
        .into_iter()
        .collect();
        assert!(!client.http_headers_match(&two_distinct));

        let single: HashMap<String, String> =
            [("authorization".to_string(), "Bearer t".to_string())]
                .into_iter()
                .collect();
        assert!(client.http_headers_match(&single));
    }

    #[test]
    fn http_headers_match_false_for_non_http_client() {
        let client = McpClient::stub("stdio-srv");
        let headers: HashMap<String, String> =
            [("authorization".to_string(), "Bearer t".to_string())]
                .into_iter()
                .collect();
        assert!(!client.http_headers_match(&headers));
    }

    #[test]
    fn refresh_managed_clients_keeps_arc_when_headers_unchanged() {
        let url = "http://localhost:5000/api/mcp";
        let mut state = McpState::new(vec![make_http_server("managed", url)]);
        let config = HttpConfig {
            url: url.to_string(),
            headers: vec![("authorization".to_string(), "Bearer t".to_string())],
        };
        state.owned_clients.insert(
            "managed".to_string(),
            Arc::new(McpClient::new_http(
                "managed".to_string(),
                config,
                None,
                None,
            )),
        );
        let before = Arc::clone(state.owned_clients.get("managed").unwrap());

        let fresh: HashMap<String, String> =
            [("authorization".to_string(), "Bearer t".to_string())]
                .into_iter()
                .collect();
        state.refresh_managed_clients(std::iter::once((url, &fresh)));

        let after = state.owned_clients.get("managed").unwrap();
        assert!(
            Arc::ptr_eq(&before, after),
            "unchanged headers must not rebuild the client"
        );
    }

    #[test]
    fn refresh_managed_clients_installs_new_arc_when_headers_differ() {
        let url = "http://localhost:5000/api/mcp";
        let mut state = McpState::new(vec![make_http_server("managed", url)]);
        let config = HttpConfig {
            url: url.to_string(),
            headers: vec![("authorization".to_string(), "Bearer old".to_string())],
        };
        state.owned_clients.insert(
            "managed".to_string(),
            Arc::new(McpClient::new_http(
                "managed".to_string(),
                config,
                None,
                None,
            )),
        );
        let before = Arc::clone(state.owned_clients.get("managed").unwrap());

        let fresh: HashMap<String, String> =
            [("authorization".to_string(), "Bearer new".to_string())]
                .into_iter()
                .collect();
        state.refresh_managed_clients(std::iter::once((url, &fresh)));

        let after = state.owned_clients.get("managed").unwrap();
        assert!(
            !Arc::ptr_eq(&before, after),
            "changed headers must install a fresh client"
        );
        assert!(after.http_headers_match(&fresh));
    }

    // ── reset_transport tests ────────────────────────────────────────

    #[tokio::test]
    async fn test_reset_transport_succeeds_for_http_client() {
        let config = HttpConfig {
            url: "http://127.0.0.1:9/api/mcp".to_string(),
            headers: vec![],
        };
        let client = McpClient::new_http("example-mcp".to_string(), config, None, None);
        assert!(client.reset_transport().await);
    }

    #[tokio::test]
    async fn test_reset_transport_fails_for_stub() {
        // Stub has `reconnect = None`, simulating a Stdio client.
        let client = McpClient::stub("stdio-srv");
        assert!(!client.reset_transport().await);
    }

    #[tokio::test]
    async fn test_reset_transport_is_idempotent() {
        let config = HttpConfig {
            url: "http://127.0.0.1:9/api/mcp".to_string(),
            headers: vec![],
        };
        let client = McpClient::new_http("example-mcp".to_string(), config, None, None);

        // Multiple resets should all succeed.
        assert!(client.reset_transport().await);
        assert!(client.reset_transport().await);
        assert!(client.reset_transport().await);
    }

    #[tokio::test]
    async fn test_reset_transport_makes_ensure_initialized_retry_handshake() {
        // Port 1 on loopback refuses immediately (ECONNREFUSED -> HandshakeFailed),
        // so each handshake fails fast instead of waiting out the connect timeout.
        let config = HttpConfig {
            url: "http://127.0.0.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let client = McpClient::new_http("test".to_string(), config, None, None);

        // First ensure_initialized will fail (unreachable server) but proves
        // the client attempts a handshake from the Pending state.
        let err1 = client.ensure_initialized().await.unwrap_err();
        assert!(
            matches!(
                err1,
                McpError::Timeout { .. } | McpError::HandshakeFailed { .. }
            ),
            "first init should fail: {err1}"
        );

        // Reset puts the client back into Pending with a fresh transport.
        assert!(client.reset_transport().await);

        // Second ensure_initialized should attempt another handshake (not
        // return a cached error). It will fail again with the same kind of
        // error, proving the reset restored the transport.
        let err2 = client.ensure_initialized().await.unwrap_err();
        assert!(
            matches!(
                err2,
                McpError::Timeout { .. } | McpError::HandshakeFailed { .. }
            ),
            "second init after reset should also attempt handshake: {err2}"
        );
    }

    #[tokio::test]
    async fn recover_errors_for_client_with_no_restorable_transport() {
        // A stub has `reconnect = None` (like Stdio): `recover` can't rebuild it.
        let err = Arc::new(McpClient::stub("stdio"))
            .recover()
            .await
            .unwrap_err();
        assert!(matches!(err, McpError::ClientError(_)), "got {err}");
    }

    #[tokio::test]
    async fn reset_transport_rebuilds_acp_client() {
        use crate::acp_transport::AcpReverseInvoker;
        use std::time::Duration;

        struct NoopInvoker;
        #[async_trait::async_trait]
        impl AcpReverseInvoker for NoopInvoker {
            async fn invoke(
                &self,
                _server_id: &str,
                _message: serde_json::Value,
                _timeout: Duration,
            ) -> Result<serde_json::Value, String> {
                Ok(serde_json::Value::Null)
            }
        }

        let client = McpClient::new_acp(
            "sdk-tools".to_string(),
            "srv_0".to_string(),
            Arc::new(NoopInvoker),
            None,
            None,
        );

        // ACP clients restore from `reconnect`, unlike Stdio.
        assert!(client.reset_transport().await);
        assert!(
            matches!(
                &*client.state.lock().await,
                ClientState::Pending(PendingTransport::Acp { .. })
            ),
            "reset_transport should restore the ACP transport to Pending"
        );
    }

    /// End-to-end reconnect-THEN-SUCCEED for the `try_call_tool` retry arm: the one
    /// piece otherwise covered only by its parts (`is_retriable_transport_error`,
    /// `reset_transport_*`, `ensure_initialized_*`).
    ///
    /// Drives the REAL `McpErasedTool::try_call_tool` against a real
    /// `McpClient`. The first `call_tool` hits a real `RunningService`
    /// whose transport is already closed, so it returns a genuine,
    /// retriable `ServiceError::TransportClosed`; the arm must then flag
    /// `reconnect_attempted`, run the real `reset_transport` +
    /// `ensure_initialized` re-handshake (rebuilding the ACP transport
    /// against a working echo server), and return the SECOND attempt's
    /// `Ok` result.
    ///
    /// Why a separately-built dead service instead of failing the initial
    /// connection: the ACP bridge transport can only be torn down from the
    /// rmcp side, so a fresh real service is built over a raw duplex whose
    /// server answers `initialize` then drops — closing the transport so
    /// the first `call_tool` observes `TransportClosed`. Everything from
    /// the retriable-error gate through the successful retry is real code.
    #[tokio::test]
    async fn try_call_tool_reconnects_then_succeeds_after_retriable_transport_error() {
        use crate::acp_transport::AcpReverseInvoker;
        use std::time::Duration;
        use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};

        // Working in-process echo server for the post-reconnect retry.
        struct EchoSdkServer;
        #[async_trait::async_trait]
        impl AcpReverseInvoker for EchoSdkServer {
            async fn invoke(
                &self,
                _server_id: &str,
                message: serde_json::Value,
                _timeout: Duration,
            ) -> Result<serde_json::Value, String> {
                let id = message
                    .get("id")
                    .cloned()
                    .unwrap_or(serde_json::Value::Null);
                let method = message
                    .get("method")
                    .and_then(|m| m.as_str())
                    .unwrap_or_default();
                let result = match method {
                    "initialize" => serde_json::json!({
                        "protocolVersion": message["params"]["protocolVersion"],
                        "capabilities": { "tools": {} },
                        "serverInfo": { "name": "echo", "version": "0.0.0" },
                    }),
                    "tools/call" => serde_json::json!({
                        "content": [{
                            "type": "text",
                            "text": message["params"]["arguments"]["text"]
                                .as_str()
                                .unwrap_or_default(),
                        }],
                        "isError": false,
                    }),
                    other => return Err(format!("unexpected method {other}")),
                };
                Ok(serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": result }))
            }
        }

        // A real `RunningService` whose transport is already closed: the
        // server answers `initialize`, consumes the `initialized`
        // notification (so the client's handshake send succeeds), then drops
        // its duplex ends. The next `call_tool` therefore observes a real
        // `ServiceError::TransportClosed`.
        async fn dead_service() -> McpService {
            let (client_read, server_write) = tokio::io::duplex(64 * 1024); // server -> client
            let (server_read, client_write) = tokio::io::duplex(64 * 1024); // client -> server
            tokio::spawn(async move {
                let mut reader = BufReader::new(server_read);
                let mut writer = server_write;
                let mut line = String::new();
                loop {
                    line.clear();
                    if reader.read_line(&mut line).await.unwrap_or(0) == 0 {
                        return;
                    }
                    let Ok(msg) = serde_json::from_str::<serde_json::Value>(line.trim()) else {
                        continue;
                    };
                    if msg.get("method").and_then(|m| m.as_str()) == Some("initialize") {
                        let id = msg.get("id").cloned().unwrap_or(serde_json::Value::Null);
                        let resp = serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": {
                            "protocolVersion": msg["params"]["protocolVersion"],
                            "capabilities": { "tools": {} },
                            "serverInfo": { "name": "dead", "version": "0.0.0" },
                        }});
                        let mut encoded = serde_json::to_string(&resp).unwrap();
                        encoded.push('\n');
                        let _ = writer.write_all(encoded.as_bytes()).await;
                        let _ = writer.flush().await;
                        // Drain the `initialized` notification, then drop to close.
                        let _ = reader.read_line(&mut line).await;
                        return;
                    }
                }
            });
            let handler = KigiClientHandler {
                info: McpClient::make_client_info("dead"),
                server_name: "dead".to_string(),
                notify_tx: Arc::new(parking_lot::Mutex::new(None)),
            };
            let transport = rmcp::transport::async_rw::AsyncRwTransport::<RoleClient, _, _>::new(
                client_read,
                client_write,
            );
            Arc::new(
                handler
                    .serve(transport)
                    .await
                    .expect("dead-service handshake"),
            )
        }

        // ACP client whose `reconnect` snapshot rebuilds against the echo server.
        let client = Arc::new(McpClient::new_acp(
            "sdk".to_string(),
            "srv_0".to_string(),
            Arc::new(EchoSdkServer),
            None,
            None,
        ));
        // Inject the closed real service so the FIRST `call_tool` fails retriably.
        let dead = dead_service().await;
        *client.state.lock().await = ClientState::Ready(dead);

        let erased = McpErasedTool {
            tool: McpTool::new(
                "echo".to_string(),
                "echo".to_string(),
                "sdk".to_string(),
                Arc::new(Mutex::new(McpState::new(vec![]))),
                serde_json::json!({}),
                None,
            ),
        };

        let raw = serde_json::json!({ "text": "after reconnect" });
        let mut reconnect_attempted = false;
        let mut is_timeout = false;
        let ew = kigi_file_utils::events::EventWriter::noop();
        let result = erased
            .try_call_tool(
                &client,
                &raw,
                &mut reconnect_attempted,
                &mut is_timeout,
                &ew,
            )
            .await
            .expect("retry after reconnect should succeed");

        // The Ok came from the SECOND attempt — the dead service cannot echo,
        // so this text proves the rebuilt transport served the retry.
        assert_eq!(
            result.content[0].as_text().expect("text content").text,
            "after reconnect"
        );
        assert!(
            reconnect_attempted,
            "retriable transport error must set reconnect_attempted"
        );
        assert!(
            !is_timeout,
            "successful retry must not be flagged as timeout"
        );
        // reset_transport + re-handshake replaced the dead service with a live one.
        assert!(matches!(&*client.state.lock().await, ClientState::Ready(_)));
    }

    #[test]
    fn is_auth_rejection_message_matches_auth_signals() {
        // The verbatim string captured in production for a managed handshake.
        assert!(is_auth_rejection_message(
            "MCP server 'notion' handshake failed: Auth required, when send initialize request"
        ));
        assert!(is_auth_rejection_message("401 Unauthorized"));
        assert!(is_auth_rejection_message("unauthorized"));
        assert!(is_auth_rejection_message("Authentication required"));
        assert!(is_auth_rejection_message("authentication failed"));
        assert!(is_auth_rejection_message("status: 401"));
        assert!(is_auth_rejection_message("HTTP status 401"));
        assert!(is_auth_rejection_message("server returned status code 401"));
        assert!(is_auth_rejection_message("HTTP 401"));
        assert!(is_auth_rejection_message("error 401"));
        // rmcp worker fatal context uses Debug form without spaces.
        assert!(is_auth_rejection_message(
            "worker quit with fatal: Transport channel closed, when Auth(AuthorizationRequired)"
        ));
        let auth_req = McpError::AuthRequired {
            server: "clickhouse".into(),
        };
        assert!(auth_req.is_auth_rejection());
        assert_eq!(auth_req.server_name(), Some("clickhouse"));
    }

    #[test]
    fn auth_required_records_as_auth_not_init_failed_and_maps_category() {
        // Pre-spawn gate is owned by the auth state machine: it lands in
        // `auth_required` (recoverable via re-auth) and never `init_failed`.
        let mut state = McpState::new(vec![]);
        state.record_init_failure("oauth-srv", true, None);
        assert!(state.auth_required.contains("oauth-srv"));
        assert!(!state.init_failed.contains_key("oauth-srv"));

        // AuthRequired carries the AuthRequired telemetry category, not ClientError.
        let err = McpError::AuthRequired {
            server: "oauth-srv".into(),
        };
        assert!(matches!(
            err.error_category(),
            kigi_file_utils::events::McpErrorCategory::AuthRequired
        ));
    }

    #[test]
    fn is_auth_rejection_message_rejects_non_auth() {
        // Transport / timeout / spawn wording is never an auth rejection.
        assert!(!is_auth_rejection_message("Transport closed"));
        assert!(!is_auth_rejection_message(
            "MCP server 'x' timed out after 30s"
        ));
        assert!(!is_auth_rejection_message(
            "Failed to spawn MCP server 'x': No such file or directory"
        ));
        // 403/forbidden is a non-auth policy denial in this stack, not auth.
        assert!(!is_auth_rejection_message("403 Forbidden"));
        assert!(!is_auth_rejection_message("forbidden"));
        // Incidental digits must not trip the status-anchored 401 patterns.
        assert!(!is_auth_rejection_message("request took 401ms"));
        assert!(!is_auth_rejection_message("connect 10.0.4.01:443"));
        assert!(!is_auth_rejection_message("read 401 bytes"));
        // A status literal followed by another alphanumeric is a different
        // token: a longer number (4012) or an adjacent unit (401ms).
        assert!(!is_auth_rejection_message("http 4012"));
        assert!(!is_auth_rejection_message("error 4012"));
        assert!(!is_auth_rejection_message("status: 4012"));
        assert!(!is_auth_rejection_message("http 401ms"));
        assert!(!is_auth_rejection_message("error 401ms"));
        // ...but a trailing punctuation/whitespace still matches.
        assert!(is_auth_rejection_message("http 401."));
        assert!(is_auth_rejection_message("error 401: token expired"));
    }

    #[test]
    fn mcp_error_is_auth_rejection_delegates() {
        assert!(McpError::ClientError("Auth required".to_string()).is_auth_rejection());
        assert!(!McpError::ClientError("Transport closed".to_string()).is_auth_rejection());
        assert!(
            !McpError::Timeout {
                server: "x".to_string(),
                timeout_secs: 30,
            }
            .is_auth_rejection()
        );
        assert!(
            !McpError::SpawnFailed {
                server: "x".to_string(),
                source: std::io::Error::new(std::io::ErrorKind::NotFound, "401 Unauthorized"),
            }
            .is_auth_rejection()
        );
        // HandshakeFailed is the production carrier: its `source` Display must
        // surface the auth substring for the delegation to fire.
        assert!(
            McpError::HandshakeFailed {
                server: "x".to_string(),
                source: Box::new(ClientInitializeError::ConnectionClosed(
                    "Auth required, when send initialize request".to_string()
                )),
            }
            .is_auth_rejection()
        );
        assert!(
            !McpError::HandshakeFailed {
                server: "x".to_string(),
                source: Box::new(ClientInitializeError::ConnectionClosed(
                    "transport closed".to_string()
                )),
            }
            .is_auth_rejection()
        );
    }

    #[test]
    fn format_mcp_image_default_emits_only_data_uri() {
        let out = format_mcp_image("image/png", "AAAA", false);
        assert_eq!(out, "data:image/png;base64,AAAA");
        assert!(!out.contains("<mcp_image_base64"));
    }

    #[test]
    fn format_mcp_image_expose_emits_data_uri_and_raw_block() {
        let out = format_mcp_image("image/png", "AAAA", true);
        assert!(out.contains("data:image/png;base64,AAAA"));
        assert!(out.contains("<mcp_image_base64 mime=\"image/png\">\nAAAA\n</mcp_image_base64>"));
    }

    /// Wrapper must not re-match the extractor regex, else the raw copy gets stripped too.
    #[test]
    fn format_mcp_image_expose_raw_block_has_no_data_prefix() {
        let out = format_mcp_image("image/jpeg", "ZZZZ", true);
        assert_eq!(out.matches("data:image/").count(), 1);
    }

    #[test]
    fn load_expose_image_base64_defaults_to_false() {
        assert!(!McpClient::load_expose_image_base64(None, None));
    }

    #[test]
    fn load_expose_image_base64_uses_overrides_when_meta_unset() {
        let overrides = McpClientTimeoutOverrides {
            expose_image_base64: Some(true),
            ..Default::default()
        };
        assert!(McpClient::load_expose_image_base64(Some(&overrides), None));
    }

    #[test]
    fn load_expose_image_base64_meta_wins_over_overrides() {
        let overrides = McpClientTimeoutOverrides {
            expose_image_base64: Some(true),
            ..Default::default()
        };
        let meta = McpServerMetaConfig {
            expose_image_base64: Some(false),
            ..Default::default()
        };
        assert!(!McpClient::load_expose_image_base64(
            Some(&overrides),
            Some(&meta)
        ));
    }

    #[test]
    fn load_expose_image_base64_meta_falls_through_when_none() {
        let overrides = McpClientTimeoutOverrides {
            expose_image_base64: Some(true),
            ..Default::default()
        };
        let meta = McpServerMetaConfig::default(); // expose_image_base64 = None
        assert!(McpClient::load_expose_image_base64(
            Some(&overrides),
            Some(&meta)
        ));
    }

    /// End-to-end: override → constructor → public getter.
    /// New constructors should add a similar assertion.
    #[test]
    fn new_http_propagates_expose_image_base64_override_to_getter() {
        let config = HttpConfig {
            url: "http://localhost/api/mcp".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            expose_image_base64: Some(true),
            ..Default::default()
        };
        let client = McpClient::new_http(
            "grafana".to_string(),
            config.clone(),
            Some(&overrides),
            None,
        );
        assert!(client.expose_image_base64());

        let client_default = McpClient::new_http("grafana".to_string(), config, None, None);
        assert!(!client_default.expose_image_base64());
    }

    // ------------------------------------------------------------------
    // ensure_initialized single-flight + Notify behavior (regression
    // suite for the "MCP client already initializing" doom-loop).
    // ------------------------------------------------------------------

    /// `ensure_initialized` on a stub (no transport) must surface a
    /// clear, actionable configuration error — never the legacy
    /// "already initializing" sentinel which leaked into model-visible
    /// tool results and triggered retry loops that exhausted the
    /// per-tick prompt budget.
    #[tokio::test]
    async fn ensure_initialized_on_empty_client_returns_no_transport_error() {
        let client = McpClient::stub("test-server");

        let err = client.ensure_initialized().await.unwrap_err();
        let msg = err.to_string();

        assert!(
            msg.contains("no transport configured"),
            "expected clear 'no transport configured' error, got: {msg}"
        );
        assert!(
            !msg.contains("already initializing"),
            "regression: legacy fast-fail sentinel surfaced: {msg}"
        );
    }

    /// Drive `N` `ensure_initialized` calls concurrently against an
    /// unreachable HTTP server with a tight startup timeout. Every
    /// caller must surface a real handshake error (`Timeout` or
    /// `HandshakeFailed`); none may surface the legacy
    /// "MCP client already initializing" sentinel which the
    /// pre-fix branch emitted whenever a caller observed
    /// `Pending(None)` while another caller was running the handshake.
    ///
    /// The race window is intentionally widened by using an unreachable
    /// host (`192.0.2.1:1` — TEST-NET-1, guaranteed unrouteable) so the
    /// handshake stalls for `startup_timeout_sec` and every concurrent
    /// caller spawned after the first observes `Initializing` instead
    /// of `Pending`.
    #[tokio::test]
    async fn ensure_initialized_concurrent_callers_never_see_legacy_fast_fail() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(1),
            ..Default::default()
        };
        let client = Arc::new(McpClient::new_http(
            "test-server".to_string(),
            config,
            Some(&overrides),
            None,
        ));

        let mut handles = Vec::new();
        for _ in 0..5 {
            let c = Arc::clone(&client);
            handles.push(tokio::spawn(async move { c.ensure_initialized().await }));
        }

        for (idx, handle) in handles.into_iter().enumerate() {
            let result = handle.await.expect("task did not panic");
            let err = result.expect_err("unreachable host must fail");
            let msg = err.to_string();
            assert!(
                !msg.contains("MCP client already initializing"),
                "caller {idx}: legacy fast-fail sentinel surfaced: {msg}"
            );
            assert!(
                matches!(
                    err,
                    McpError::Timeout { .. } | McpError::HandshakeFailed { .. }
                ),
                "caller {idx}: expected handshake failure, got: {err}"
            );
        }
    }

    /// A caller that finds `ClientState::Initializing` must park on
    /// `init_done` and wake up when the holder publishes a new state,
    /// then take the freshly-restored transport for its own retry.
    ///
    /// We exercise the wake path directly (without an actual concurrent
    /// handshake) by manually transitioning state to `Initializing`,
    /// spawning a parker, then transitioning back to `Pending` and
    /// firing `notify_waiters`. The parker should retry against the
    /// restored (still-unreachable) transport and surface a normal
    /// handshake error rather than the wait-timeout error.
    #[tokio::test]
    async fn ensure_initialized_parked_caller_retries_after_notify() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(1),
            ..Default::default()
        };
        let client = Arc::new(McpClient::new_http(
            "test-server".to_string(),
            config.clone(),
            Some(&overrides),
            None,
        ));

        // Simulate an in-flight handshake by another task: pretend
        // that task took the transport and entered Initializing.
        *client.state.lock().await = ClientState::Initializing;

        // Spawn the parker. It must observe Initializing and park on
        // `init_done` rather than fail-fast.
        let parker_client = Arc::clone(&client);
        let parker = tokio::spawn(async move { parker_client.ensure_initialized().await });

        // Give the parker a chance to reach the await on `init_done`.
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Publish a fresh Pending transport and notify — simulates the
        // holder's failure-path restore.
        *client.state.lock().await = ClientState::Pending(PendingTransport::Http(config.clone()));
        client.init_done.notify_waiters();

        // The parker should wake, take the transport, run its own
        // handshake (which fails against the unreachable host), and
        // surface a regular handshake error — never the wait-timeout
        // error and never the legacy fast-fail.
        let err = parker
            .await
            .expect("parker did not panic")
            .expect_err("unreachable host must fail");
        let msg = err.to_string();
        assert!(
            !msg.contains("MCP client already initializing"),
            "regression: legacy fast-fail sentinel: {msg}"
        );
        assert!(
            !msg.contains("init still in progress"),
            "parker should not hit wait-timeout when notified: {msg}"
        );
        assert!(
            matches!(
                err,
                McpError::Timeout { .. } | McpError::HandshakeFailed { .. }
            ),
            "expected handshake failure, got: {err}"
        );
    }

    /// If a caller is parked on `Initializing` and the holder is
    /// dropped without notifying (cancellation-storm edge case), the
    /// parker must eventually surface a clear `init still in progress`
    /// timeout error rather than block indefinitely.
    ///
    /// Without the inflight-wait timeout, a wedged client (one whose
    /// drop guard couldn't acquire the lock to restore) would silently
    /// stall every future `ensure_initialized` caller until process
    /// restart. The 1 s margin past `startup_timeout_sec` keeps the
    /// happy path snappy while still bounding the worst case.
    #[tokio::test]
    async fn ensure_initialized_inflight_wait_times_out_when_holder_silent() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(0),
            ..Default::default()
        };
        let client = McpClient::new_http("test-server".to_string(), config, Some(&overrides), None);

        // Wedge the slot in Initializing with no live holder.
        *client.state.lock().await = ClientState::Initializing;

        let err = client.ensure_initialized().await.unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("init still in progress"),
            "expected wait-timeout error, got: {msg}"
        );
        assert!(
            !msg.contains("already initializing"),
            "regression: legacy fast-fail sentinel: {msg}"
        );
    }

    /// When the holder task is cancelled (`abort()`) mid-handshake, the
    /// `InitGuard` drop impl restores `Pending(transport)` on a
    /// best-effort basis so a follow-on caller can retry without
    /// requiring an explicit `reset_transport`.
    #[tokio::test]
    async fn ensure_initialized_drop_guard_restores_state_after_holder_aborted() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let overrides = McpClientTimeoutOverrides {
            // Long enough that the holder is guaranteed to still be
            // inside try_handshake when we abort it.
            startup_timeout_sec: Some(10),
            ..Default::default()
        };
        let client = Arc::new(McpClient::new_http(
            "test-server".to_string(),
            config,
            Some(&overrides),
            None,
        ));

        let holder_client = Arc::clone(&client);
        let holder = tokio::spawn(async move { holder_client.ensure_initialized().await });

        // Wait for the holder to enter Initializing.
        let started = std::time::Instant::now();
        loop {
            if matches!(&*client.state.lock().await, ClientState::Initializing) {
                break;
            }
            assert!(
                started.elapsed() < std::time::Duration::from_secs(2),
                "holder never reached Initializing"
            );
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }

        // Cancel the holder mid-handshake. The drop guard should
        // restore Pending so the next caller can retry.
        holder.abort();
        let _ = holder.await;

        // The drop guard restores best-effort via `try_lock` and notifies.
        // Wait briefly for it to settle.
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        match &*client.state.lock().await {
            ClientState::Pending(_) => {} // expected
            other => panic!(
                "expected Pending after holder abort + drop guard, found {}",
                state_label(other)
            ),
        }
    }

    /// `McpState::is_initialized()` MUST require both the early
    /// `finish_init` flag AND an empty `initializing_servers` set.
    ///
    /// The session actor's `start_mcp_servers` path calls `finish_init`
    /// **early** (right after spawning processes, before any handshake
    /// completes) so non-MCP work can proceed in parallel. Tool dispatch
    /// and the Blocking-strategy prompt guard, however, must NOT
    /// observe "initialized" until every per-server handshake is done —
    /// otherwise the model's first tool call races the background
    /// `get_tool_registrations` handshake and the
    /// `McpClient::ensure_initialized` window described above triggers.
    #[test]
    fn test_mcp_state_is_initialized_requires_empty_initializing_servers() {
        let mut state = McpState::new(vec![make_stdio_server("a", "/bin/a")]);

        // NotStarted: neither flag set, no per-server work.
        assert!(!state.is_initialized());
        assert!(!state.is_initializing());
        assert!(!state.has_finished_init());
        assert!(matches!(state.init_progress(), InitProgress::NotStarted));

        // Starting: try_start_init fired, per-server names registered,
        // finish_init has NOT yet fired. is_initializing() is true.
        assert!(state.try_start_init());
        state.mark_servers_initializing(["a".to_string()]);
        assert!(!state.is_initialized());
        assert!(state.is_initializing());
        assert!(!state.has_finished_init());
        assert!(matches!(
            state.init_progress(),
            InitProgress::Starting { .. }
        ));

        // Finished + handshakes outstanding: actor called finish_init
        // early but the per-server background handshake is still in
        // flight. is_initialized() must be FALSE during this window.
        state.finish_init();
        assert!(
            !state.is_initialized(),
            "is_initialized() must wait for per-server handshakes"
        );
        assert!(
            state.is_initializing(),
            "is_initializing() must report in-flight per-server work"
        );
        assert!(state.has_finished_init());
        assert!(state.is_server_handshaking("a"));
        assert_eq!(state.handshaking_servers_count(), 1);

        // Finished + empty: background task has reported the handshake
        // complete. Now and only now is the pool fully initialized.
        state.mark_server_ready("a");
        assert!(state.is_initialized());
        assert!(!state.is_initializing());
        assert!(state.has_finished_init());
        assert!(!state.is_server_handshaking("a"));
        assert_eq!(state.handshaking_servers_count(), 0);
    }

    /// Locks in the typed-state contract: the `init_progress` field
    /// makes nonsensical combinations like "initialized AND
    /// initializing" structurally unrepresentable. Every legal state
    /// has exactly one [`InitProgress`] variant; every transition is
    /// driven through the typed methods.
    #[test]
    fn test_init_progress_state_machine_invariants() {
        let mut state = McpState::new(vec![make_stdio_server("a", "/bin/a")]);

        // Invariant: try_start_init is one-shot per cycle.
        assert!(state.try_start_init());
        assert!(!state.try_start_init(), "double try_start_init is rejected");

        // Invariant: mark_all_servers_ready clears handshaking in
        // both Starting and Finished states; never resurrects them.
        state.mark_servers_initializing(["a".to_string(), "b".to_string()]);
        assert_eq!(state.handshaking_servers_count(), 2);
        state.mark_all_servers_ready();
        assert_eq!(state.handshaking_servers_count(), 0);
        assert!(
            matches!(state.init_progress(), InitProgress::Starting { .. }),
            "mark_all_servers_ready preserves the lifecycle variant"
        );

        // Invariant: finish_init from Starting → Finished preserves
        // (or in this case, the now-empty) handshaking set.
        state.finish_init();
        assert!(state.is_initialized());
        assert!(matches!(
            state.init_progress(),
            InitProgress::Finished { .. }
        ));

        // Invariant: cancel_init returns us cleanly to NotStarted,
        // ready for a new try_start_init.
        state.cancel_init();
        assert!(matches!(state.init_progress(), InitProgress::NotStarted));
        assert!(state.try_start_init(), "cancel_init re-enables init");
    }

    fn state_label(s: &ClientState) -> &'static str {
        match s {
            ClientState::Empty => "Empty",
            ClientState::Pending(_) => "Pending",
            ClientState::Initializing => "Initializing",
            ClientState::Ready(_) => "Ready",
        }
    }

    // -- is_healthy / state_kind --------------------------------------
    //
    // These tests cover the cheap, non-blocking predicate. They focus
    // on the state-machine inspection: any
    // non-`Ready` variant returns `false` for `is_healthy`, and
    // `state_kind` projects every variant onto the matching
    // [`ClientStateKind`].
    //
    // The two `Ready` cases
    // (`is_healthy_ready_open_returns_true` and
    // `is_healthy_transport_closed_returns_false`) require a real
    // `RunningService<RoleClient, InitializeRequestParams>`, which can
    // only be constructed through rmcp's `serve_client` path. That
    // path needs a peer that responds to the MCP initialize
    // handshake, and this crate intentionally does NOT enable rmcp's
    // `server` feature (see `Cargo.toml`). Wiring up a hand-rolled
    // JSON-RPC responder over `tokio::io::duplex` would balloon the
    // test scaffolding far beyond what these tests need. We therefore
    // exercise the `Ready` arm indirectly: the cheap predicate is a
    // single `match` on the state mutex plus
    // `Peer::is_transport_closed`, which is upstream-tested in rmcp
    // itself (`rmcp-2.1.0/tests/test_close_connection.rs`).

    #[tokio::test]
    async fn is_healthy_empty_returns_false() {
        let client = McpClient::stub("empty");
        // `stub` starts in `ClientState::Empty`.
        assert!(matches!(*client.state.lock().await, ClientState::Empty));
        assert!(!client.is_healthy().await);
        assert_eq!(client.state_kind().await, ClientStateKind::Empty);
    }

    #[tokio::test]
    async fn is_healthy_pending_returns_false() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        let client = McpClient::new_http("pending".to_string(), config, None, None);
        // `new_http` constructs with `ClientState::Pending(_)`.
        assert!(matches!(
            *client.state.lock().await,
            ClientState::Pending(_)
        ));
        assert!(!client.is_healthy().await);
        assert_eq!(client.state_kind().await, ClientStateKind::Pending);
    }

    #[tokio::test]
    async fn is_healthy_initializing_returns_false() {
        let client = McpClient::stub("initializing");
        *client.state.lock().await = ClientState::Initializing;
        assert!(!client.is_healthy().await);
        assert_eq!(client.state_kind().await, ClientStateKind::Initializing);
    }

    /// `is_healthy` MUST NOT trigger a handshake. Regression guard:
    /// the previous implementation called `ensure_initialized`, which
    /// for a `Pending` HTTP client pointing at an unreachable host
    /// would block for `startup_timeout_sec` seconds. The cheap
    /// predicate must return immediately.
    #[tokio::test]
    async fn is_healthy_pending_does_not_block_on_handshake() {
        let config = HttpConfig {
            url: "http://192.0.2.1:1/unreachable".to_string(),
            headers: vec![],
        };
        // Force a generous startup timeout — if the predicate
        // regressed to going through ensure_initialized, this test
        // would hang for ~10 s. We assert it completes in well under
        // a second.
        let overrides = McpClientTimeoutOverrides {
            startup_timeout_sec: Some(10),
            ..Default::default()
        };
        let client = McpClient::new_http(
            "pending-unreachable".to_string(),
            config,
            Some(&overrides),
            None,
        );
        let start = std::time::Instant::now();
        let healthy = client.is_healthy().await;
        let elapsed = start.elapsed();
        assert!(!healthy);
        // 1 s bound: the cheap path is microseconds, so this is a 10×
        // safety margin against cold-runtime / contended-CI jitter while
        // still firing well inside the 10 s blocking window that a
        // regressed predicate (back through `ensure_initialized`) would
        // sit in.
        assert!(
            elapsed < std::time::Duration::from_secs(1),
            "is_healthy must be a cheap state inspection, took {elapsed:?}"
        );
    }

    // -- KigiClientHandler --------------------------------------
    //
    // The handler's notification routing is the only behavior worth
    // unit-testing here; `get_info` is a literal `info.clone()` and
    // doesn't merit a test. `NotificationContext` is non-trivial to
    // construct outside of an rmcp `RunningService`, so we exercise
    // the routing through the `emit` helper that the trait methods
    // call. If the trait wiring (one-line `async move { self.emit(...) }`)
    // ever regresses, the integration tests against a real MCP
    // server will catch it.

    #[tokio::test]
    async fn client_handler_routes_tools_changed() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<McpClientEvent>();
        let handler = KigiClientHandler {
            info: McpClient::make_client_info("test"),
            server_name: "test".to_string(),
            notify_tx: Arc::new(parking_lot::Mutex::new(Some(tx))),
        };
        handler.emit(McpClientEvent::ToolsChanged {
            server: handler.server_name.clone(),
        });
        let ev = rx.recv().await.expect("event arrived");
        match ev {
            McpClientEvent::ToolsChanged { server } => assert_eq!(server, "test"),
            other => panic!("expected ToolsChanged, got {other:?}"),
        }
    }

    /// Contract: when `notify_tx` is `None` (subagent snapshot,
    /// no dispatcher), `emit` is a no-op and the trait methods
    /// must not panic.
    #[tokio::test]
    async fn client_handler_no_dispatcher_is_silent() {
        let handler = KigiClientHandler {
            info: McpClient::make_client_info("test"),
            server_name: "test".to_string(),
            notify_tx: Arc::new(parking_lot::Mutex::new(None)),
        };
        handler.emit(McpClientEvent::ToolsChanged {
            server: "test".to_string(),
        });
        // No assertion needed — reaching this line means no panic.
    }

    /// Contract: get_info returns a clone of the stored ClientInfo.
    #[tokio::test]
    async fn client_handler_get_info_round_trips() {
        let info = McpClient::make_client_info("test-srv");
        let handler = KigiClientHandler {
            info: info.clone(),
            server_name: "test-srv".to_string(),
            notify_tx: Arc::new(parking_lot::Mutex::new(None)),
        };
        let got = handler.get_info();
        // ClientInfo doesn't derive PartialEq; check the visible
        // fields the constructor sets.
        assert_eq!(got.client_info.name, info.client_info.name);
        assert_eq!(got.client_info.version, info.client_info.version);
    }

    // A sender wired *after* the handler is constructed must still
    // reach the live rmcp service loop. This test exercises the
    // post-construction wiring path: build a handler from a client
    // whose slot is `None`, then install a sender via
    // `client.set_event_tx` and verify the handler picks it up (the
    // handler holds a clone of the same shared Arc slot).
    #[tokio::test]
    async fn client_handler_observes_post_handshake_set_event_tx() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<McpClientEvent>();
        // McpClient::stub initializes notify_tx as `Arc<Mutex<None>>`.
        let client = Arc::new(McpClient::stub("test"));

        // Build the handler BEFORE wiring the sender — emulates
        // the production flow where `make_client_handler` is called
        // during `try_handshake` and the dispatcher is wired
        // separately.
        let handler = client.make_client_handler();

        // Confirm the slot is `None` at handler-construction time.
        assert!(handler.notify_tx.lock().is_none());

        // Now wire the sender on the client. Because the handler
        // holds a CLONE OF THE SAME ARC, this mutation is observed
        // by the handler's next `emit`.
        client.set_event_tx(Some(tx));

        handler.emit(McpClientEvent::ToolsChanged {
            server: "test".to_string(),
        });
        let ev = rx.recv().await.expect("event arrived");
        match ev {
            McpClientEvent::ToolsChanged { server } => assert_eq!(server, "test"),
            other => panic!("expected ToolsChanged, got {other:?}"),
        }
    }

    // Mirrors the post-construction wiring on the `ensure_initialized`
    // emit path: even though `Ready` / `HandshakeFailed` fire from
    // inside `try_handshake`, the slot is read at emit time through the
    // SAME shared Arc, so wiring `set_event_tx` BEFORE the handshake is
    // sufficient to capture these events.
    #[tokio::test]
    async fn event_tx_clone_observes_set_event_tx() {
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel::<McpClientEvent>();
        let client = McpClient::stub("test");
        assert!(client.event_tx_clone().is_none());
        client.set_event_tx(Some(tx));
        assert!(client.event_tx_clone().is_some());
        client.set_event_tx(None);
        assert!(client.event_tx_clone().is_none());
    }

    // An `ensure_initialized`-emitted `Ready` event must NOT be
    // conflated with a restart. This unit test exercises the event
    // level; the wire-level mapping ("Ready → reason=initialized, NOT
    // restart_succeeded") is covered by host integration tests.
    #[test]
    fn config_added_kind_carries_correct_server_name() {
        let ev = McpClientEvent::ConfigAdded {
            server: "srv".to_string(),
        };
        assert_eq!(ev.server_name(), Some("srv"));
    }
}
