//! Protocol [`Client`]: handshake, correlation, routing, shutdown.
//!
//! The runner consumes this Client through whatever transport it
//! brought along (stdio, streamable HTTP, SSE). Responsibilities split
//! across three async tasks:
//!
//! 1. The Client owns an in-flight map keyed by request ID. Sending a
//!    request inserts a `oneshot` sender; receiving a response or error
//!    response pulls the sender out and completes the request.
//! 2. A background `pump` task drains the transport's inbound channel,
//!    classifies each envelope, and routes it to the right consumer:
//!    response into the in-flight map, notification to the registered
//!    handler, server-initiated request to the registered handler.
//! 3. `close()` shuts the pump down and aborts every outstanding
//!    `oneshot` with [`ProtocolError::Aborted`].
//!
//! The Client is `Send + Sync` so the runner can hand clones to its
//! per-test workers and still drive `close()` from the orchestrator.

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde_json::{json, Value};
use tokio::sync::{oneshot, Mutex};
use tokio::task::{JoinHandle, JoinSet};
use tokio::time::timeout;

use super::message::MessageKind;
use super::violations::{
    frame_snippet, ProtocolViolationKind, ProtocolViolationSummary, ViolationRecorder,
};
use super::{
    JsonRpcErrorBody, JsonRpcId, JsonRpcMessage, ProtocolError, ProtocolVersion, Transport,
    TransportReceiver, TransportSender,
};

/// Client-side capability block sent during `initialize`.
///
/// Treated as an opaque JSON value here so the engine does not have to
/// re-implement the upstream schema before  lands. Callers build
/// the block from `serde_json::json!` or the rust-mcp-schema types.
#[derive(Debug, Clone, PartialEq)]
pub struct ClientCapabilities {
    /// Free-form name advertised to the server. Defaults to `"mcptest"`.
    pub name: String,
    /// Free-form version string advertised to the server. Defaults to
    /// the crate version.
    pub version: String,
    /// Opaque capabilities object. Defaults to `{}` so a minimal
    /// handshake works without manual fields.
    pub capabilities: Value,
}

impl Default for ClientCapabilities {
    fn default() -> Self {
        Self {
            name: "mcptest".into(),
            // MCP servers parse `clientInfo.version` as a SemVer string, so
            // hand them the bare Cargo package version (no `+sha` suffix
            // that would trip strict parsers).
            version: crate::version::semver().into(),
            capabilities: json!({}),
        }
    }
}

/// Server-side capability block parsed from the `initialize` response.
///
/// Fields are opaque pass-through values until  brings the
/// upstream schema. Callers inspect `capabilities` via `serde_json` until
/// then.
#[derive(Debug, Clone, PartialEq)]
pub struct ServerCapabilities {
    /// `serverInfo.name` from the response, if present.
    pub name: Option<String>,
    /// `serverInfo.version` from the response, if present.
    pub version: Option<String>,
    /// Opaque capabilities object.
    pub capabilities: Value,
}

/// Outcome of [`Client::initialize`]: negotiated version plus server
/// capabilities.
#[derive(Debug, Clone, PartialEq)]
pub struct InitializeOutcome {
    /// Protocol version actually negotiated. May differ from the
    /// version the Client advertised when the server pinned a different
    /// supported revision.
    pub negotiated_version: ProtocolVersion,
    /// Capability block returned by the server.
    pub server: ServerCapabilities,
}

/// User-provided handler for server-initiated requests.
///
/// Boxed because async closures are still unstable. The handler returns
/// the result payload that the Client wraps in a JSON-RPC response and
/// sends back to the server. Errors propagate as JSON-RPC errors with
/// code `-32603` (internal error); richer mapping lands in a follow-up.
pub type ServerRequestHandler = Arc<
    dyn Fn(String, Value) -> Pin<Box<dyn Future<Output = Result<Value, String>> + Send>>
        + Send
        + Sync,
>;

/// Builder for [`Client`] so optional knobs (handlers, timeouts) do not
/// require a 5-arg constructor.
pub struct ClientBuilder {
    transport: Transport,
    client_capabilities: ClientCapabilities,
    request_timeout: Duration,
    server_request_handler: Option<ServerRequestHandler>,
    notification_handler: Option<Arc<dyn Fn(String, Value) + Send + Sync>>,
}

impl ClientBuilder {
    /// Start a builder over `transport` with default knobs.
    pub fn new(transport: Transport) -> Self {
        Self {
            transport,
            client_capabilities: ClientCapabilities::default(),
            request_timeout: Duration::from_secs(30),
            server_request_handler: None,
            notification_handler: None,
        }
    }

    /// Override the client capability block sent during `initialize`.
    pub fn with_client_capabilities(mut self, caps: ClientCapabilities) -> Self {
        self.client_capabilities = caps;
        self
    }

    /// Override the per-request timeout. Default is 30 seconds.
    pub fn with_request_timeout(mut self, d: Duration) -> Self {
        self.request_timeout = d;
        self
    }

    /// Register a handler for server-initiated requests (for example
    /// `sampling/createMessage`). When unset, every server request
    /// receives a `-32601 method not found` error response.
    pub fn with_server_request_handler(mut self, handler: ServerRequestHandler) -> Self {
        self.server_request_handler = Some(handler);
        self
    }

    /// Register a handler for server-initiated notifications. When
    /// unset, notifications are recorded internally and returned by
    /// [`Client::drain_notifications`].
    pub fn with_notification_handler(
        mut self,
        handler: Arc<dyn Fn(String, Value) + Send + Sync>,
    ) -> Self {
        self.notification_handler = Some(handler);
        self
    }

    /// Build the Client and spawn its inbound pump task.
    pub fn build(self) -> Client {
        Client::spawn(
            self.transport,
            self.client_capabilities,
            self.request_timeout,
            self.server_request_handler,
            self.notification_handler,
        )
    }
}

type InflightMap = Arc<Mutex<HashMap<i64, oneshot::Sender<Result<Value, ProtocolError>>>>>;
type NotificationLog = Arc<Mutex<Vec<(String, Value)>>>;
type ViolationLog = Arc<Mutex<ViolationRecorder>>;
/// Tracker for spawned server-initiated request handler tasks
/// (WOR-1314). Shared between the pump (which spawns into it) and
/// `close()` (which aborts and drains it) so a panicked handler is
/// logged via its join error instead of vanishing, and no handler task
/// outlives the Client.
type ServerTaskSet = Arc<Mutex<JoinSet<()>>>;

/// JSON-RPC client that drives the MCP handshake and correlates in-flight
/// requests.
///
/// Cloning is cheap (an `Arc` clone) and safe; every handle shares the
/// same in-flight map and pump task.
#[derive(Clone)]
pub struct Client {
    inner: Arc<ClientInner>,
}

struct ClientInner {
    outbound: TransportSender,
    inflight: InflightMap,
    next_id: AtomicI64,
    request_timeout: Duration,
    client_capabilities: ClientCapabilities,
    pump: Mutex<Option<JoinHandle<()>>>,
    notifications: NotificationLog,
    violations: ViolationLog,
    server_tasks: ServerTaskSet,
    /// Set by the pump when the transport's inbound channel closes
    /// (WOR-1326). Once true, no response can ever arrive, so new
    /// requests fail immediately instead of waiting out their timeout.
    transport_closed: Arc<AtomicBool>,
    /// Raw `initialize` result captured during the handshake (WOR-1347).
    /// Re-issuing `initialize` mid-session is not valid MCP, so consumers
    /// that need to inspect the handshake (compliance `check: initialize`)
    /// read this retained copy instead of sending a second request.
    initialize_response: Mutex<Option<Value>>,
}

impl ClientInner {
    /// Build (and log) the timeout error, enriched with the connection's
    /// protocol-violation hint when one exists (WOR-1311), so a timeout on a
    /// violating connection points at the findings, not the timeout setting.
    async fn timeout_error(&self, id: i64, method: &str, timeout_ms: u64) -> ProtocolError {
        let hint = self.violations.lock().await.summary().timeout_hint();
        match &hint {
            Some(hint) => tracing::warn!(
                target: "mcptest_core::protocol",
                id,
                method = %method,
                timeout_ms,
                "request timed out (note: {hint})"
            ),
            None => tracing::warn!(
                target: "mcptest_core::protocol",
                id,
                method = %method,
                timeout_ms,
                "request timed out"
            ),
        }
        ProtocolError::request_timeout_with_hint(id, timeout_ms, hint)
    }
}

/// A request sent on the wire but not yet awaited (WOR-1407), so a caller can
/// read the id immediately and await the reply separately. The
/// cancel-after-completion race probe uses this: send a call, cancel it
/// mid-flight, then observe whether the server still replies.
pub struct PendingRequest {
    id: i64,
    method: String,
    rx: oneshot::Receiver<Result<Value, ProtocolError>>,
    inner: Arc<ClientInner>,
}

impl PendingRequest {
    /// The JSON-RPC request id allocated for this call, available before the
    /// response arrives.
    pub fn id(&self) -> i64 {
        self.id
    }

    /// Await the correlated response, honoring the client's per-request timeout
    /// (the same semantics as [`Client::request`]).
    pub async fn response(self) -> Result<Value, ProtocolError> {
        let timeout_ms = self.inner.request_timeout.as_millis() as u64;
        match timeout(self.inner.request_timeout, self.rx).await {
            Ok(Ok(Ok(value))) => Ok(value),
            Ok(Ok(Err(err))) => Err(err),
            Ok(Err(_canceled)) => Err(ProtocolError::Aborted { id: self.id }),
            Err(_elapsed) => {
                let _ = self.inner.inflight.lock().await.remove(&self.id);
                Err(self
                    .inner
                    .timeout_error(self.id, &self.method, timeout_ms)
                    .await)
            }
        }
    }

    /// Observe the response with a bounded wait: `Ok(None)` means no reply
    /// arrived within `within`, `Ok(Some(_))` is the reply. After cancelling an
    /// in-flight request a conforming server sends nothing (`None`); a
    /// non-conforming one replies anyway (`Some`).
    pub async fn response_within(self, within: Duration) -> Result<Option<Value>, ProtocolError> {
        match timeout(within, self.rx).await {
            Ok(Ok(Ok(value))) => Ok(Some(value)),
            Ok(Ok(Err(err))) => Err(err),
            Ok(Err(_canceled)) => Err(ProtocolError::Aborted { id: self.id }),
            Err(_elapsed) => {
                // No reply within the window. Drop our slot so the in-flight
                // map stays honest; a later reply for this id is harmless (the
                // pump finds no sender and discards it).
                let _ = self.inner.inflight.lock().await.remove(&self.id);
                Ok(None)
            }
        }
    }
}

impl Client {
    /// Start a [`ClientBuilder`] over `transport`.
    pub fn builder(transport: Transport) -> ClientBuilder {
        ClientBuilder::new(transport)
    }

    fn spawn(
        transport: Transport,
        client_capabilities: ClientCapabilities,
        request_timeout: Duration,
        server_request_handler: Option<ServerRequestHandler>,
        notification_handler: Option<Arc<dyn Fn(String, Value) + Send + Sync>>,
    ) -> Self {
        let Transport {
            outbound_tx,
            inbound_rx,
        } = transport;
        let inflight: InflightMap = Arc::new(Mutex::new(HashMap::new()));
        let notifications: NotificationLog = Arc::new(Mutex::new(Vec::new()));
        let violations: ViolationLog = Arc::new(Mutex::new(ViolationRecorder::default()));
        let server_tasks: ServerTaskSet = Arc::new(Mutex::new(JoinSet::new()));
        let transport_closed = Arc::new(AtomicBool::new(false));

        let pump_handle = spawn_pump(
            inbound_rx,
            outbound_tx.clone(),
            PumpCtx {
                inflight: Arc::clone(&inflight),
                notifications: Arc::clone(&notifications),
                violations: Arc::clone(&violations),
                server_tasks: Arc::clone(&server_tasks),
                transport_closed: Arc::clone(&transport_closed),
                server_request_handler,
                notification_handler,
            },
        );

        Self {
            inner: Arc::new(ClientInner {
                outbound: outbound_tx,
                inflight,
                next_id: AtomicI64::new(1),
                request_timeout,
                client_capabilities,
                pump: Mutex::new(Some(pump_handle)),
                notifications,
                violations,
                server_tasks,
                transport_closed,
                initialize_response: Mutex::new(None),
            }),
        }
    }

    /// Allocate the next monotonic request ID. Cast to `i64` for spec
    /// compatibility; the spec accepts any JSON number so this is safe
    /// for billions of requests per run.
    fn next_id(&self) -> i64 {
        self.inner.next_id.fetch_add(1, Ordering::Relaxed)
    }

    /// Send `initialize`, wait for the response, send the `initialized`
    /// notification. Returns the negotiated version and the server's
    /// capability block.
    pub async fn initialize(
        &self,
        client_version: ProtocolVersion,
    ) -> Result<InitializeOutcome, ProtocolError> {
        let caps = &self.inner.client_capabilities;
        let params = json!({
            "protocolVersion": client_version.as_wire(),
            "capabilities": caps.capabilities,
            "clientInfo": {
                "name": caps.name,
                "version": caps.version,
            },
        });
        let response = self.request("initialize", params).await?;
        let outcome = parse_capability_block(&response)?;
        // Retain the raw result so [`Self::initialize_response`] can hand
        // the handshake to consumers that must not re-issue `initialize`.
        *self.inner.initialize_response.lock().await = Some(response);
        self.notify("initialized", json!({})).await?;
        Ok(outcome)
    }

    /// Return the raw `initialize` result captured during the handshake.
    ///
    /// Why this exists: a compliance `check: initialize` must evaluate the
    /// handshake, but re-issuing `initialize` on an established session is
    /// not valid MCP. The connector runs the handshake once at connect
    /// time; this accessor replays its raw result. `None` until
    /// [`Client::initialize`] has completed successfully.
    pub async fn initialize_response(&self) -> Option<Value> {
        self.inner.initialize_response.lock().await.clone()
    }

    /// Send `server/discover` (SEP-2575) and parse the same
    /// capability shape `initialize` returns.
    ///
    /// `server/discover` is the stateless replacement for the
    /// `initialize` handshake on the 2026-07-28 transport: no session,
    /// no `initialized` notification. The wire shape mirrors the
    /// `initialize` response (`protocolVersion`, `serverInfo`,
    /// `capabilities`), so the parsing logic is shared.
    pub async fn discover(&self) -> Result<InitializeOutcome, ProtocolError> {
        let response = self.request("server/discover", json!({})).await?;
        parse_capability_block(&response)
    }

    /// Send a request, wait for the response, return the `result`
    /// payload. Honors the per-request timeout configured on the
    /// builder. Concurrent in-flight requests are correlated by ID; the
    /// pump task fans responses out to the right `oneshot`.
    pub async fn request(
        &self,
        method: impl Into<String>,
        params: Value,
    ) -> Result<Value, ProtocolError> {
        self.request_with_id(method, params)
            .await
            .map(|(_id, value)| value)
    }

    /// Like [`Client::request`], but also return the request id that was
    /// sent on the wire.
    ///
    /// Why the id surfaces at all: the URL transports keep per-request
    /// side channels keyed by JSON-RPC request id (response header
    /// capture, WOR-1287), and only the Client knows which id a given
    /// call was allocated. Callers that do not consume a side channel
    /// should stay on [`Client::request`].
    // WOR-1519: one span per outbound MCP request. `skip_all` keeps params (which
    // can be large or secret-bearing) out of the span; the span is the unit a
    // future OTLP exporter and the OpenInference reporter hang latency on.
    #[tracing::instrument(name = "mcp.request", skip_all)]
    pub async fn request_with_id(
        &self,
        method: impl Into<String>,
        params: Value,
    ) -> Result<(i64, Value), ProtocolError> {
        let pending = self.send_request(method, params).await?;
        let id = pending.id();
        let value = pending.response().await?;
        Ok((id, value))
    }

    /// Send a request and return a [`PendingRequest`] handle without awaiting
    /// the reply (WOR-1407), so the id is available before the response. A
    /// transport that already closed, or a send that times out, fails here
    /// rather than handing back a handle that can never complete.
    pub async fn send_request(
        &self,
        method: impl Into<String>,
        params: Value,
    ) -> Result<PendingRequest, ProtocolError> {
        let id = self.next_id();
        let method_str: String = method.into();
        tracing::debug!(
            target: "mcptest_core::protocol",
            id,
            method = %method_str,
            "send request"
        );
        let (tx, rx) = oneshot::channel();
        {
            let mut map = self.inner.inflight.lock().await;
            map.insert(id, tx);
        }
        // Fail fast on a transport that already closed (WOR-1326): no
        // response can ever arrive. Checked after the insert so this
        // request cannot slip between the pump's closing drain and the
        // flag flip; one of the two paths always sees it.
        if self.inner.transport_closed.load(Ordering::SeqCst) {
            let _ = self.inner.inflight.lock().await.remove(&id);
            return Err(ProtocolError::TransportClosed(format!(
                "transport closed before request {id} (`{method_str}`) could complete"
            )));
        }
        let msg = JsonRpcMessage::request(JsonRpcId::Number(id), method_str.clone(), params);
        let timeout_ms = self.inner.request_timeout.as_millis() as u64;
        // `send_wait` applies backpressure when the transport runs a
        // bounded outbound queue (WOR-1313); the timeout keeps a wedged
        // writer from parking this request forever.
        let sent = timeout(
            self.inner.request_timeout,
            self.inner.outbound.send_wait(msg),
        )
        .await;
        match sent {
            Ok(Ok(())) => {}
            Ok(Err(_closed)) => {
                // Pull the sender back out so we do not leak an entry. The
                // receiver is about to be dropped on its own, but draining
                // keeps the map honest.
                let _ = self.inner.inflight.lock().await.remove(&id);
                return Err(ProtocolError::TransportClosed(
                    "outbound channel closed before request was sent".into(),
                ));
            }
            Err(_elapsed) => {
                let _ = self.inner.inflight.lock().await.remove(&id);
                return Err(self.inner.timeout_error(id, &method_str, timeout_ms).await);
            }
        }
        Ok(PendingRequest {
            id,
            method: method_str,
            rx,
            inner: Arc::clone(&self.inner),
        })
    }

    /// Send a fire-and-forget notification. Returns once the envelope
    /// has been handed to the outbound channel; on a bounded transport
    /// this awaits queue space (backpressure, WOR-1313) rather than
    /// failing on a momentarily full queue.
    pub async fn notify(
        &self,
        method: impl Into<String>,
        params: Value,
    ) -> Result<(), ProtocolError> {
        let msg = JsonRpcMessage::notification(method, params);
        self.inner.outbound.send_wait(msg).await.map_err(|_| {
            ProtocolError::TransportClosed(
                "outbound channel closed before notification was sent".into(),
            )
        })
    }

    /// Drain and return every server-initiated notification recorded
    /// since the last call. Useful when no notification handler was
    /// registered on the builder.
    pub async fn drain_notifications(&self) -> Vec<(String, Value)> {
        let mut log = self.inner.notifications.lock().await;
        std::mem::take(&mut *log)
    }

    /// Snapshot the protocol violations observed on this connection
    /// (WOR-1311).
    ///
    /// Why: mcptest is conformance tooling, so evidence of peer
    /// misbehavior (string-id echoes, orphan responses, malformed
    /// frames) must not be eaten by the pump. The runner attaches this
    /// snapshot to reports so a failing run carries its diagnosis.
    pub async fn protocol_violations(&self) -> ProtocolViolationSummary {
        self.inner.violations.lock().await.summary()
    }

    /// Shut down the Client. Aborts every outstanding request with
    /// [`ProtocolError::Aborted`], stops the pump task, and closes the
    /// outbound channel. Idempotent.
    pub async fn close(&self) -> Result<(), ProtocolError> {
        // Drop the outbound channel by replacing it with an unused one
        // is not possible without `&mut self`; instead we drain the
        // in-flight map so every pending request fails fast, then abort
        // the pump.
        let pending = {
            let mut map = self.inner.inflight.lock().await;
            std::mem::take(&mut *map)
        };
        for (id, tx) in pending {
            let _ = tx.send(Err(ProtocolError::Aborted { id }));
        }
        let mut slot = self.inner.pump.lock().await;
        if let Some(handle) = slot.take() {
            handle.abort();
            // We deliberately swallow the join error: the pump was
            // aborted on purpose so the JoinError carries no caller-
            // useful state.
            let _ = handle.await;
        }
        drop(slot);
        // With the pump stopped, no new handler tasks can be spawned.
        // Abort and drain the outstanding ones (WOR-1314) so nothing
        // outlives the Client, logging any panic the drain surfaces.
        let mut tasks = self.inner.server_tasks.lock().await;
        tasks.abort_all();
        while let Some(joined) = tasks.join_next().await {
            log_handler_join_error(joined);
        }
        Ok(())
    }
}

/// Log the join outcome of a server-initiated request handler task
/// (WOR-1314). A panic must leave a trace in the logs; a cancellation
/// is the expected result of `close()` aborting the set and stays
/// silent.
fn log_handler_join_error(joined: Result<(), tokio::task::JoinError>) {
    let Err(err) = joined else {
        return;
    };
    if err.is_panic() {
        tracing::error!(
            target: "mcptest_core::protocol",
            error = %err,
            "server-initiated request handler panicked"
        );
    }
}

/// Parse the capability block that both `initialize` and
/// `server/discover` return. Pulled out as a free
/// function so the two entry points stay in sync.
///
/// Shape:
///
/// ```json
/// {
///   "protocolVersion": "2026-07-28",
///   "serverInfo": { "name": "...", "version": "..." },
///   "capabilities": { ... }
/// }
/// ```
///
/// A missing or unparseable `protocolVersion` is a protocol error;
/// `serverInfo` is optional (some servers omit it on `server/discover`).
fn parse_capability_block(response: &Value) -> Result<InitializeOutcome, ProtocolError> {
    let server_version = response
        .get("protocolVersion")
        .and_then(Value::as_str)
        .ok_or_else(|| {
            ProtocolError::MalformedMessage("capability block missing `protocolVersion`".into())
        })?;
    let negotiated = ProtocolVersion::from_wire(server_version)
        .ok_or_else(|| ProtocolError::ProtocolVersionMismatch(server_version.to_owned()))?;
    let server = ServerCapabilities {
        name: response
            .get("serverInfo")
            .and_then(|v| v.get("name"))
            .and_then(Value::as_str)
            .map(str::to_owned),
        version: response
            .get("serverInfo")
            .and_then(|v| v.get("version"))
            .and_then(Value::as_str)
            .map(str::to_owned),
        capabilities: response
            .get("capabilities")
            .cloned()
            .unwrap_or_else(|| json!({})),
    };
    Ok(InitializeOutcome {
        negotiated_version: negotiated,
        server,
    })
}

/// Shared state the pump task routes envelopes against. Bundled into a
/// struct so [`spawn_pump`] stays under the argument-count lint now
/// that the handler task set (WOR-1314) rides along.
struct PumpCtx {
    inflight: InflightMap,
    notifications: NotificationLog,
    violations: ViolationLog,
    server_tasks: ServerTaskSet,
    transport_closed: Arc<AtomicBool>,
    server_request_handler: Option<ServerRequestHandler>,
    notification_handler: Option<Arc<dyn Fn(String, Value) + Send + Sync>>,
}

fn spawn_pump(
    mut inbound_rx: TransportReceiver,
    outbound_tx: TransportSender,
    ctx: PumpCtx,
) -> JoinHandle<()> {
    tokio::spawn(async move {
        let PumpCtx {
            inflight,
            notifications,
            violations,
            server_tasks,
            transport_closed,
            server_request_handler,
            notification_handler,
        } = ctx;
        while let Some(msg) = inbound_rx.recv().await {
            let kind = match msg.kind() {
                Ok(k) => k,
                Err(err) => {
                    // WOR-1311: malformed frames are evidence of a broken
                    // server, so they are recorded as violations (with a
                    // bounded snippet) instead of being silently skipped.
                    let snippet = frame_snippet(&msg);
                    violations.lock().await.record(
                        ProtocolViolationKind::MalformedFrame,
                        format!("{err}; frame: {snippet}"),
                    );
                    continue;
                }
            };
            match kind {
                MessageKind::Response { id, result } => {
                    if let Some(tx) = take_inflight_sender(&id, &inflight, &violations).await {
                        let _ = tx.send(Ok(result));
                    }
                }
                MessageKind::ErrorResponse { id, error } => {
                    if let Some(tx) = take_inflight_sender(&id, &inflight, &violations).await {
                        let _ = tx.send(Err(ProtocolError::ServerError {
                            code: error.code,
                            message: error.message,
                            data: error.data,
                        }));
                    }
                }
                MessageKind::Notification { method } => {
                    let params = msg.params.clone().unwrap_or_else(|| json!({}));
                    if let Some(handler) = &notification_handler {
                        handler(method, params);
                    } else {
                        let mut log = notifications.lock().await;
                        log.push((method, params));
                    }
                }
                MessageKind::Request { id, method } => {
                    let params = msg.params.clone().unwrap_or_else(|| json!({}));
                    let outbound = outbound_tx.clone();
                    if let Some(handler) = server_request_handler.clone() {
                        let mut tasks = server_tasks.lock().await;
                        // Reap already-finished handlers while we are
                        // here so a panic surfaces in the logs promptly
                        // (WOR-1314) instead of waiting for close().
                        while let Some(joined) = tasks.try_join_next() {
                            log_handler_join_error(joined);
                        }
                        tasks.spawn(async move {
                            let fut = handler(method, params);
                            let reply = match fut.await {
                                Ok(value) => JsonRpcMessage::response(id, value),
                                Err(message) => JsonRpcMessage::error_response(
                                    id,
                                    JsonRpcErrorBody {
                                        code: -32603,
                                        message,
                                        data: None,
                                    },
                                ),
                            };
                            // Best-effort: if the outbound channel is
                            // closed the Client is being torn down.
                            let _ = outbound.send_wait(reply).await;
                        });
                    } else {
                        // No handler: reply with method-not-found so the
                        // server's request does not hang.
                        let reply = JsonRpcMessage::error_response(
                            id,
                            JsonRpcErrorBody {
                                code: -32601,
                                message: format!("method `{method}` not handled by client"),
                                data: None,
                            },
                        );
                        let _ = outbound.send_wait(reply).await;
                    }
                }
            }
        }
        // The transport's inbound channel closed: the child exited, the
        // socket dropped, or the replay finished. No in-flight request
        // can ever complete, so fail them all now (WOR-1326) instead of
        // letting each wait out its own timeout while the server is
        // already gone. The flag flips before the drain so a request
        // inserted after the drain still fails fast through the
        // post-insert check in `Client::request`.
        transport_closed.store(true, Ordering::SeqCst);
        let pending = {
            let mut map = inflight.lock().await;
            std::mem::take(&mut *map)
        };
        for (id, tx) in pending {
            // The receiver may already be gone (its request timed out
            // on its own); best-effort delivery is fine.
            let _ = tx.send(Err(ProtocolError::TransportClosed(format!(
                "transport closed while request {id} was waiting for its response"
            ))));
        }
    })
}

/// Resolve a response id against the in-flight map and pull the waiting
/// sender out, recording protocol violations along the way (WOR-1311).
///
/// Why this exists: this engine only allocates numeric ids, but real
/// servers echo them back as strings ("1" instead of 1) often enough
/// that silently dropping those responses turned a diagnosable server
/// bug into an opaque RequestTimeout. The coercion resolves the request
/// anyway and records [`ProtocolViolationKind::StringIdEcho`]; ids that
/// cannot be correlated at all record
/// [`ProtocolViolationKind::OrphanResponse`].
async fn take_inflight_sender(
    id: &JsonRpcId,
    inflight: &InflightMap,
    violations: &ViolationLog,
) -> Option<oneshot::Sender<Result<Value, ProtocolError>>> {
    let numeric = match id.as_i64() {
        Some(n) => n,
        None => {
            let coerced = match id {
                JsonRpcId::String(s) => s.parse::<i64>().ok(),
                JsonRpcId::Number(n) => Some(*n),
            };
            match coerced {
                Some(n) => {
                    violations.lock().await.record(
                        ProtocolViolationKind::StringIdEcho,
                        format!("server echoed numeric request id {n} as the string \"{n}\""),
                    );
                    n
                }
                None => {
                    violations.lock().await.record(
                        ProtocolViolationKind::OrphanResponse,
                        format!("response carried non-numeric id {id:?} that matches no in-flight request"),
                    );
                    return None;
                }
            }
        }
    };
    let sender = {
        let mut map = inflight.lock().await;
        map.remove(&numeric)
    };
    if sender.is_none() {
        violations.lock().await.record(
            ProtocolViolationKind::OrphanResponse,
            format!("response with id {numeric} matches no in-flight request"),
        );
    }
    sender
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::mpsc;

    #[tokio::test]
    async fn inflight_request_fails_fast_when_the_transport_closes() {
        let (out_tx, _out_rx) = mpsc::unbounded_channel();
        let (in_tx, in_rx) = mpsc::unbounded_channel();
        let client = Client::builder(Transport::new(out_tx, in_rx)).build();

        let started = tokio::time::Instant::now();
        let request = client.request("initialize", json!({}));
        // Close the transport while the request is (or is about to be)
        // in flight, the way a dead stdio child closes its stdout.
        drop(in_tx);
        let err = request.await.expect_err("no response can arrive");
        assert!(
            matches!(err, ProtocolError::TransportClosed(_)),
            "expected TransportClosed, got {err:?}"
        );
        assert!(
            started.elapsed() < Duration::from_secs(5),
            "request waited {:?} instead of failing fast",
            started.elapsed()
        );
        let _ = client.close().await;
    }

    #[tokio::test]
    async fn request_after_transport_close_fails_without_waiting_for_the_timeout() {
        let (out_tx, _out_rx) = mpsc::unbounded_channel();
        let (in_tx, in_rx) = mpsc::unbounded_channel();
        let client = Client::builder(Transport::new(out_tx, in_rx)).build();

        drop(in_tx);
        // Give the pump a moment to observe the close and flip the flag.
        tokio::time::sleep(Duration::from_millis(50)).await;
        let started = tokio::time::Instant::now();
        let err = client
            .request("tools/list", json!({}))
            .await
            .expect_err("transport is gone");
        assert!(
            matches!(err, ProtocolError::TransportClosed(_)),
            "expected TransportClosed, got {err:?}"
        );
        assert!(
            started.elapsed() < Duration::from_secs(5),
            "request waited {:?} instead of failing fast",
            started.elapsed()
        );
        let _ = client.close().await;
    }

    #[tokio::test]
    async fn request_with_id_returns_the_id_that_was_sent_on_the_wire() {
        // WOR-1287: transport side channels (response header capture) are
        // keyed by request id, so the id handed back must be the one the
        // envelope carried, not a fresh allocation.
        let (out_tx, mut out_rx) = mpsc::unbounded_channel();
        let (in_tx, in_rx) = mpsc::unbounded_channel();
        let client = Client::builder(Transport::new(out_tx, in_rx)).build();

        let request = client.request_with_id("tools/list", json!({}));
        let echo = async {
            let sent = out_rx.recv().await.expect("outbound envelope");
            let wire_id = sent.id.clone().expect("request carries an id");
            let numeric = wire_id.as_i64().expect("engine allocates numeric ids");
            in_tx
                .send(JsonRpcMessage::response(wire_id, json!({"ok": true})))
                .expect("inbound channel open");
            numeric
        };
        let (outcome, wire_id) = tokio::join!(request, echo);
        let (returned_id, value) = outcome.expect("response resolves");
        assert_eq!(returned_id, wire_id);
        assert_eq!(value, json!({"ok": true}));
        let _ = client.close().await;
    }

    #[tokio::test]
    async fn initialize_retains_the_raw_handshake_result_for_later_inspection() {
        // WOR-1347: compliance `check: initialize` evaluates the handshake
        // without re-issuing it, so the raw result must survive the parse.
        let (out_tx, mut out_rx) = mpsc::unbounded_channel();
        let (in_tx, in_rx) = mpsc::unbounded_channel();
        let client = Client::builder(Transport::new(out_tx, in_rx)).build();

        assert!(
            client.initialize_response().await.is_none(),
            "nothing retained before the handshake"
        );
        let block = json!({
            "protocolVersion": "2025-03-26",
            "serverInfo": { "name": "demo", "version": "1.0.0" },
            "capabilities": { "tools": {} }
        });
        let handshake = client.initialize(ProtocolVersion::V2025_03_26);
        let echo = async {
            let sent = out_rx.recv().await.expect("outbound initialize");
            let wire_id = sent.id.clone().expect("request carries an id");
            in_tx
                .send(JsonRpcMessage::response(wire_id, block.clone()))
                .expect("inbound channel open");
        };
        let (outcome, ()) = tokio::join!(handshake, echo);
        outcome.expect("handshake succeeds");
        assert_eq!(client.initialize_response().await, Some(block));
        let _ = client.close().await;
    }

    #[test]
    fn parse_capability_block_pulls_version_and_capabilities() {
        // `server/discover` (and initialize) share the same
        // capability shape; the parser pulls every documented field.
        let response = json!({
            "protocolVersion": "2026-07-28",
            "serverInfo": { "name": "demo-mcp", "version": "1.2.3" },
            "capabilities": { "tools": { "listChanged": true } }
        });
        let outcome = parse_capability_block(&response).expect("well-formed block parses");
        assert_eq!(outcome.negotiated_version, ProtocolVersion::V2026_07_28);
        assert_eq!(outcome.server.name.as_deref(), Some("demo-mcp"));
        assert_eq!(outcome.server.version.as_deref(), Some("1.2.3"));
        assert_eq!(
            outcome.server.capabilities,
            json!({ "tools": { "listChanged": true } })
        );
    }

    #[test]
    fn parse_capability_block_treats_missing_server_info_as_optional() {
        // 2026-07-28 server/discover responses sometimes omit
        // serverInfo; the parser must still succeed.
        let response = json!({
            "protocolVersion": "2026-07-28",
            "capabilities": {}
        });
        let outcome = parse_capability_block(&response).expect("missing serverInfo is fine");
        assert!(outcome.server.name.is_none());
        assert!(outcome.server.version.is_none());
    }

    #[test]
    fn parse_capability_block_rejects_missing_protocol_version() {
        let response = json!({ "capabilities": {} });
        let err = parse_capability_block(&response).expect_err("must require protocolVersion");
        assert!(matches!(err, ProtocolError::MalformedMessage(_)));
    }

    #[test]
    fn parse_capability_block_rejects_unknown_protocol_version() {
        let response = json!({ "protocolVersion": "2099-12-31", "capabilities": {} });
        let err = parse_capability_block(&response).expect_err("unknown version is a mismatch");
        match err {
            ProtocolError::ProtocolVersionMismatch(s) => assert_eq!(s, "2099-12-31"),
            other => panic!("expected ProtocolVersionMismatch, got {other:?}"),
        }
    }
}
