tower-mcp 0.25.0

//! Unit tests for [`McpClient`](super::McpClient) and its message loop.
//!
//! Moved out in #1256. A sibling module rather than an integration
//! test because these reach private items.

use super::*;
use async_trait::async_trait;
use std::sync::Mutex;

/// Mock transport for testing that auto-responds to requests.
///
/// When the client sends a request via `send()`, the mock extracts the
/// request ID, pairs it with the next preconfigured response, and feeds
/// it back through a channel that `recv()` awaits on. This ensures
/// `recv()` blocks when no messages are available (instead of returning
/// EOF), keeping the background message loop alive.
struct MockTransport {
    /// Pre-configured result or error replies (not full envelopes).
    responses: Arc<Mutex<Vec<MockReply>>>,
    /// Index of the next response to use.
    response_idx: Arc<std::sync::atomic::AtomicUsize>,
    /// Channel sender for feeding responses back to `recv()`.
    incoming_tx: mpsc::Sender<String>,
    /// Channel receiver for `recv()` to await on.
    incoming_rx: mpsc::Receiver<String>,
    /// Collected outgoing messages from `send()`.
    outgoing: Arc<Mutex<Vec<String>>>,
    connected: Arc<AtomicBool>,
    /// When set, `send()` fails for notifications (messages without an
    /// `id`), simulating a transport that could not deliver them.
    fail_notification_sends: Arc<AtomicBool>,
}

enum MockReply {
    Result(serde_json::Value),
    Error(JsonRpcError),
}

#[allow(dead_code)]
impl MockTransport {
    fn new() -> Self {
        let (tx, rx) = mpsc::channel(32);
        Self {
            responses: Arc::new(Mutex::new(Vec::new())),
            response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
            incoming_tx: tx,
            incoming_rx: rx,
            outgoing: Arc::new(Mutex::new(Vec::new())),
            connected: Arc::new(AtomicBool::new(true)),
            fail_notification_sends: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Create a mock that auto-responds with the given result payloads.
    ///
    /// When `send()` receives a JSON-RPC request, it extracts the request
    /// ID and pairs it with the next response from this list, sending the
    /// complete JSON-RPC response through the channel for `recv()`.
    fn with_responses(responses: Vec<serde_json::Value>) -> Self {
        let (tx, rx) = mpsc::channel(32);
        Self {
            responses: Arc::new(Mutex::new(
                responses.into_iter().map(MockReply::Result).collect(),
            )),
            response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
            incoming_tx: tx,
            incoming_rx: rx,
            outgoing: Arc::new(Mutex::new(Vec::new())),
            connected: Arc::new(AtomicBool::new(true)),
            fail_notification_sends: Arc::new(AtomicBool::new(false)),
        }
    }

    fn with_replies(responses: Vec<MockReply>) -> Self {
        let (tx, rx) = mpsc::channel(32);
        Self {
            responses: Arc::new(Mutex::new(responses)),
            response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
            incoming_tx: tx,
            incoming_rx: rx,
            outgoing: Arc::new(Mutex::new(Vec::new())),
            connected: Arc::new(AtomicBool::new(true)),
            fail_notification_sends: Arc::new(AtomicBool::new(false)),
        }
    }
}

#[async_trait]
impl ClientTransport for MockTransport {
    async fn send(&mut self, message: &str) -> Result<()> {
        self.outgoing.lock().unwrap().push(message.to_string());

        // Parse the outgoing message to extract the request ID
        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(message) {
            if parsed.get("id").is_none() && self.fail_notification_sends.load(Ordering::Relaxed) {
                return Err(Error::Transport(
                    "mock transport dropped the notification".to_string(),
                ));
            }
            // Only respond to requests (messages with an id and method)
            if let Some(id) = parsed.get("id") {
                let idx = self
                    .response_idx
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                let responses = self.responses.lock().unwrap();
                if let Some(reply) = responses.get(idx) {
                    let response = match reply {
                        MockReply::Result(result) => serde_json::json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "result": result
                        }),
                        MockReply::Error(error) => serde_json::json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "error": error
                        }),
                    };
                    let _ = self.incoming_tx.try_send(response.to_string());
                }
            }
        }

        Ok(())
    }

    async fn recv(&mut self) -> Result<Option<String>> {
        // Await on the channel -- blocks until a message is available
        // or the sender is dropped (returns None = EOF).
        match self.incoming_rx.recv().await {
            Some(msg) => Ok(Some(msg)),
            None => Ok(None),
        }
    }

    fn is_connected(&self) -> bool {
        self.connected.load(Ordering::Relaxed)
    }

    async fn close(&mut self) -> Result<()> {
        self.connected.store(false, Ordering::Relaxed);
        Ok(())
    }
}

fn mock_initialize_response() -> serde_json::Value {
    serde_json::json!({
        "protocolVersion": "2025-11-25",
        "serverInfo": {
            "name": "test-server",
            "version": "1.0.0"
        },
        "capabilities": {
            "tools": {}
        }
    })
}

/// #1174: a notification the transport fails to deliver must fail
/// `initialize()`. Reporting success left the handshake incomplete and a
/// strict server rejecting every subsequent request with -32600.
#[tokio::test]
async fn initialize_fails_when_initialized_notification_is_not_delivered() {
    let transport = MockTransport::with_responses(vec![mock_initialize_response()]);
    let fail_notifications = transport.fail_notification_sends.clone();
    let outgoing = transport.outgoing.clone();
    // The initialize request itself succeeds; only the follow-up
    // notification is dropped.
    fail_notifications.store(true, Ordering::Relaxed);

    let client = McpClient::connect(transport).await.unwrap();
    let error = client
        .initialize("test-client", "1.0.0")
        .await
        .expect_err("undelivered notifications/initialized must fail initialize");

    assert!(
        error
            .to_string()
            .contains("failed to deliver notifications/initialized"),
        "error should name the handshake step, got: {error}"
    );
    // The notification was attempted, not skipped.
    assert!(
        outgoing
            .lock()
            .unwrap()
            .iter()
            .any(|message| message.contains("notifications/initialized")),
        "the client must have tried to send the notification"
    );
    // The client must not consider itself initialized.
    assert!(!client.is_initialized());
}

#[tokio::test]
async fn notification_delivery_errors_reach_the_caller() {
    let transport = MockTransport::with_responses(vec![mock_initialize_response()]);
    let fail_notifications = transport.fail_notification_sends.clone();

    let client = McpClient::connect(transport).await.unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Healthy so far; now the transport starts dropping notifications.
    fail_notifications.store(true, Ordering::Relaxed);
    let error = client
        .notify("notifications/progress", &serde_json::json!({}))
        .await
        .expect_err("a dropped notification must surface as an error");
    assert!(error.to_string().contains("dropped the notification"));
}

#[tokio::test]
async fn test_client_not_initialized() {
    let client = McpClient::connect(MockTransport::with_responses(vec![]))
        .await
        .unwrap();

    let result = client.list_tools().await;
    assert!(result.is_err());
    assert!(result.unwrap_err().to_string().contains("not initialized"));
}

#[tokio::test]
async fn test_client_initialize() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
    ]))
    .await
    .unwrap();

    assert!(!client.is_initialized());

    let result = client.initialize("test-client", "1.0.0").await;
    assert!(result.is_ok());
    assert!(client.is_initialized());

    let server_info = client.server_info().await.unwrap();
    assert_eq!(server_info.server_info.name, "test-server");
}

fn mock_initialize_response_with_version(protocol_version: &str) -> serde_json::Value {
    serde_json::json!({
        "protocolVersion": protocol_version,
        "serverInfo": {
            "name": "test-server",
            "version": "1.0.0"
        },
        "capabilities": {
            "tools": {}
        }
    })
}

/// #1473: `initialize` must validate the version the server selected, not
/// just accept whatever comes back. A discover-lifecycle version here means
/// the server answered the legacy handshake with a version that removed it
/// -- accepting it would leave the client holding a session the server never
/// intends to service that way.
#[tokio::test]
async fn initialize_rejects_a_discover_lifecycle_version() {
    let transport =
        MockTransport::with_responses(vec![mock_initialize_response_with_version("2026-07-28")]);
    let client = McpClient::connect(transport).await.unwrap();

    let outcome = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.initialize("test-client", "1.0.0"),
    )
    .await
    .expect("initialize must not hang");

    let error = outcome
        .expect_err("a discover-lifecycle version must not complete the initialize handshake");
    assert!(
        error.to_string().contains("2026-07-28"),
        "error should name the rejected version, got: {error}"
    );
    assert!(!client.is_initialized());
}

/// A server that selects 2025-06-18 completes the handshake: it is an
/// initialize-lifecycle version in the default client policy.
#[tokio::test]
async fn initialize_accepts_2025_06_18() {
    let transport =
        MockTransport::with_responses(vec![mock_initialize_response_with_version("2025-06-18")]);
    let client = McpClient::connect(transport).await.unwrap();

    let result = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.initialize("test-client", "1.0.0"),
    )
    .await
    .expect("initialize must not hang")
    .expect("a server-selected 2025-06-18 must complete initialize");

    assert_eq!(result.protocol_version, "2025-06-18");
    assert!(client.is_initialized());
}

/// Run `initialize` against a server that selects `version`, with the
/// client's version policy set to `support`.
async fn initialize_answered_with(
    version: &str,
    support: ProtocolSupport,
) -> (McpClient, Result<InitializeResult>) {
    let transport =
        MockTransport::with_responses(vec![mock_initialize_response_with_version(version)]);
    let client = McpClient::builder()
        .protocol_support(support)
        .connect_simple(transport)
        .await
        .unwrap();

    let outcome = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.initialize("test-client", "1.0.0"),
    )
    .await
    .expect("initialize must not hang");
    (client, outcome)
}

/// #1544: many stdio servers still answer 2024-11-05. Every built-in
/// version policy accepts it.
#[tokio::test]
async fn initialize_accepts_2024_11_05_with_built_in_policies() {
    for support in [
        ProtocolSupport::default(),
        ProtocolSupport::compiled(),
        ProtocolSupport::stable(),
    ] {
        let (client, outcome) = initialize_answered_with("2024-11-05", support.clone()).await;
        let result = outcome.unwrap_or_else(|error| {
            panic!("2024-11-05 must complete initialize under {support:?}, got: {error}")
        });
        assert_eq!(result.protocol_version, "2024-11-05");
        assert!(client.is_initialized());
    }
}

/// 2024-11-05 follows 2025-03-26: a client narrowed away from the oldest
/// version this crate implements rejects it too.
#[tokio::test]
async fn initialize_rejects_2024_11_05_when_2025_03_26_is_excluded() {
    let support = ProtocolSupport::try_new(["2025-11-25", "2025-06-18"]).unwrap();
    let (client, outcome) = initialize_answered_with("2024-11-05", support).await;

    let error = outcome.expect_err("2024-11-05 must be rejected without 2025-03-26");
    assert!(
        error.to_string().contains("2024-11-05"),
        "error should name the rejected version, got: {error}"
    );
    assert!(!client.is_initialized());
}

/// Accepting 2024-11-05 does not open the check to other pre-2025 versions.
#[tokio::test]
async fn initialize_rejects_2024_10_07() {
    let (client, outcome) =
        initialize_answered_with("2024-10-07", ProtocolSupport::default()).await;

    let error = outcome.expect_err("2024-10-07 must not complete initialize");
    assert!(
        error.to_string().contains("2024-10-07"),
        "error should name the rejected version, got: {error}"
    );
    assert!(!client.is_initialized());
}

/// A version the client has never heard of is rejected the same way as a
/// known-but-wrong-lifecycle one: `initialize` has no basis for trusting it.
#[tokio::test]
async fn initialize_rejects_an_unknown_version() {
    let transport =
        MockTransport::with_responses(vec![mock_initialize_response_with_version("1999-01-01")]);
    let client = McpClient::connect(transport).await.unwrap();

    let outcome = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.initialize("test-client", "1.0.0"),
    )
    .await
    .expect("initialize must not hang");

    let error = outcome.expect_err("an unrecognized protocol version must not complete initialize");
    assert!(
        error.to_string().contains("1999-01-01"),
        "error should name the rejected version, got: {error}"
    );
    assert!(!client.is_initialized());
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_discover_injects_metadata_on_every_request() {
    let transport = MockTransport::with_responses(vec![
        serde_json::json!({
            "resultType": "complete",
            "supportedVersions": ["2026-07-28"],
            "capabilities": {}
        }),
        serde_json::json!({
            "resultType": "complete",
            "tools": [],
            "ttlMs": 0,
            "cacheScope": "private"
        }),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .with_elicitation()
        .connect_simple(transport)
        .await
        .unwrap();

    client.discover("test-client", "1.0.0").await.unwrap();
    client.list_tools().await.unwrap();
    assert_eq!(
        client.selected_protocol_version().await.as_deref(),
        Some("2026-07-28")
    );

    let messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect();
    assert_eq!(messages.len(), 2);
    assert_eq!(messages[0]["method"], "server/discover");
    assert_eq!(messages[1]["method"], "tools/list");
    for message in messages {
        let meta = &message["params"]["_meta"];
        assert_eq!(
            meta["io.modelcontextprotocol/protocolVersion"],
            "2026-07-28"
        );
        assert_eq!(
            meta["io.modelcontextprotocol/clientInfo"]["name"],
            "test-client"
        );
        assert!(meta["io.modelcontextprotocol/clientCapabilities"].is_object());
        assert!(
            meta.get("io.modelcontextprotocol/logLevel").is_none(),
            "final clients must not opt into logging by default"
        );
    }
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
fn final_discover_result() -> serde_json::Value {
    serde_json::json!({
        "resultType": "complete",
        "supportedVersions": ["2026-07-28"],
        "capabilities": {},
        "ttlMs": 0,
        "cacheScope": "private"
    })
}

/// #1473: a server that rejects `server/discover` with `-32022` and offers
/// only legacy `initialize`-lifecycle versions in `supported` has nothing
/// discover can retry with. Picking a legacy version there would leave the
/// client believing it's on the sessionless path while it never sends
/// `initialize` or the `notifications/initialized` handshake the server
/// actually expects.
#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn discover_retry_fails_when_only_legacy_versions_are_mutual() {
    let error = JsonRpcError {
        code: -32022,
        message: "Unsupported protocol version".to_string(),
        data: Some(serde_json::json!({ "supported": ["2025-11-25"] })),
    };
    let transport = MockTransport::with_replies(vec![MockReply::Error(error)]);
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::default())
        .connect_simple(transport)
        .await
        .unwrap();

    let outcome = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.discover("test-client", "1.0.0"),
    )
    .await
    .expect("discover must not hang waiting on a mismatched retry");

    let error = outcome.expect_err(
        "a server offering only legacy versions must fail discover, not retry into one",
    );
    assert!(
        error.to_string().contains("fall back to initialize"),
        "error should point the caller at the initialize fallback, got: {error}"
    );
    assert!(!client.is_initialized());
}

/// The counterpart to the above: when `supported` mixes a legacy version
/// with a discover-lifecycle one the client also enables, the retry must
/// still pick the discover-lifecycle version and succeed.
#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn discover_retry_succeeds_with_a_mutual_discover_lifecycle_version() {
    let error = JsonRpcError {
        code: -32022,
        message: "Unsupported protocol version".to_string(),
        data: Some(serde_json::json!({ "supported": ["2025-11-25", "2026-07-28"] })),
    };
    let transport = MockTransport::with_replies(vec![
        MockReply::Error(error),
        MockReply::Result(final_discover_result()),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::default())
        .connect_simple(transport)
        .await
        .unwrap();

    let outcome = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        client.discover("test-client", "1.0.0"),
    )
    .await
    .expect("discover must not hang");

    outcome.expect("a mutual discover-lifecycle version must let the retry succeed");
    assert_eq!(
        client.selected_protocol_version().await.as_deref(),
        Some("2026-07-28")
    );

    let messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect();
    assert_eq!(
        messages.len(),
        2,
        "expected the original request plus exactly one retry"
    );
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_request_log_level_seeds_typed_calls_and_preserves_explicit_overrides() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        serde_json::json!({
            "resultType": "complete",
            "tools": [],
            "ttlMs": 0,
            "cacheScope": "private"
        }),
        serde_json::json!({
            "resultType": "complete",
            "content": [{"type": "text", "text": "done"}]
        }),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .request_log_level(LogLevel::Info)
        .connect_simple(transport)
        .await
        .unwrap();

    client.discover("test-client", "1.0.0").await.unwrap();
    client.list_tools().await.unwrap();
    let params = CallToolParams {
        name: "work".to_string(),
        arguments: serde_json::json!({}),
        input_responses: None,
        request_state: None,
        meta: Some(RequestMeta {
            log_level: Some(LogLevel::Warning),
            ..Default::default()
        }),
        task: None,
    };
    let _: CallToolResult = client.request("tools/call", &params).await.unwrap();
    client
        .notify("notifications/custom", &serde_json::json!({}))
        .await
        .unwrap();

    let messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect();
    assert_eq!(messages.len(), 4);
    assert_eq!(messages[0]["method"], "server/discover");
    assert_eq!(messages[1]["method"], "tools/list");
    assert_eq!(messages[2]["method"], "tools/call");
    assert_eq!(messages[3]["method"], "notifications/custom");
    for message in &messages[..2] {
        assert_eq!(
            message["params"]["_meta"]["io.modelcontextprotocol/logLevel"],
            "info"
        );
    }
    assert_eq!(
        messages[2]["params"]["_meta"]["io.modelcontextprotocol/logLevel"], "warning",
        "an explicit per-call threshold must override the client default"
    );
    assert!(
        messages[3]["params"]["_meta"]
            .get("io.modelcontextprotocol/logLevel")
            .is_none(),
        "request-only log metadata must not leak onto notifications"
    );
}

#[tokio::test]
async fn request_log_level_does_not_change_stable_wire() {
    let transport = MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "content": [{"type": "text", "text": "done"}]
        }),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .request_log_level(LogLevel::Info)
        .connect_simple(transport)
        .await
        .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    client.set_request_log_level(Some(LogLevel::Warning)).await;
    client
        .call_tool("work", serde_json::json!({}))
        .await
        .unwrap();
    client
        .notify("notifications/custom", &serde_json::json!({}))
        .await
        .unwrap();

    let messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect();
    assert_eq!(messages.len(), 4);
    assert!(
        messages.iter().all(|message| message["params"]["_meta"]
            .get("io.modelcontextprotocol/logLevel")
            .is_none()),
        "a final-only request default must not alter stable traffic: {messages:?}"
    );
    assert!(
        messages[2]["params"].get("_meta").is_none(),
        "a stable typed call must retain its ordinary params shape"
    );
}

const TEST_TRACEPARENT: &str = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";

fn test_trace_context() -> crate::TraceContext {
    crate::TraceContext::parse(TEST_TRACEPARENT, Some("vendor=value")).unwrap()
}

fn sent_messages(outgoing: &Arc<Mutex<Vec<String>>>) -> Vec<serde_json::Value> {
    outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect()
}

#[tokio::test]
async fn trace_context_provider_marks_stable_requests_but_not_notifications() {
    let transport = MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({}),
        serde_json::json!({"content": [{"type": "text", "text": "done"}]}),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .trace_context_provider(|| Some(test_trace_context()))
        .connect_simple(transport)
        .await
        .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    // A request with no arguments still gets a `_meta` object to carry it.
    client.ping().await.unwrap();
    client
        .call_tool("work", serde_json::json!({}))
        .await
        .unwrap();
    client
        .notify("notifications/custom", &serde_json::json!({}))
        .await
        .unwrap();

    let messages = sent_messages(&outgoing);
    let methods: Vec<_> = messages
        .iter()
        .map(|message| message["method"].as_str().unwrap())
        .collect();
    assert_eq!(
        methods,
        [
            "initialize",
            "notifications/initialized",
            "ping",
            "tools/call",
            "notifications/custom"
        ]
    );
    for index in [0, 2, 3] {
        let meta = &messages[index]["params"]["_meta"];
        assert_eq!(meta["traceparent"], TEST_TRACEPARENT, "{}", messages[index]);
        assert_eq!(meta["tracestate"], "vendor=value", "{}", messages[index]);
    }
    for index in [1, 4] {
        assert!(
            messages[index]["params"].get("_meta").is_none(),
            "a notification must not carry trace context: {}",
            messages[index]
        );
    }
}

#[tokio::test]
async fn trace_context_without_tracestate_sends_only_traceparent() {
    let transport = MockTransport::with_responses(vec![serde_json::json!({})]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .trace_context_provider(|| crate::TraceContext::parse(TEST_TRACEPARENT, None))
        .connect_simple(transport)
        .await
        .unwrap();

    client.ping().await.unwrap();

    let meta = &sent_messages(&outgoing)[0]["params"]["_meta"];
    assert_eq!(meta["traceparent"], TEST_TRACEPARENT);
    assert!(meta.get("tracestate").is_none());
}

#[tokio::test]
async fn no_provider_or_no_context_sends_no_trace_keys() {
    for provider in [false, true] {
        let transport = MockTransport::with_responses(vec![
            mock_initialize_response(),
            serde_json::json!({"content": [{"type": "text", "text": "done"}]}),
        ]);
        let outgoing = transport.outgoing.clone();
        let mut builder = McpClient::builder();
        if provider {
            // A provider with nothing to report, such as outside any trace.
            builder = builder.trace_context_provider(|| None);
        }
        let client = builder.connect_simple(transport).await.unwrap();

        client.initialize("test-client", "1.0.0").await.unwrap();
        client
            .call_tool("work", serde_json::json!({}))
            .await
            .unwrap();

        for message in sent_messages(&outgoing) {
            assert!(
                message["params"]["_meta"].get("traceparent").is_none()
                    && message["params"]["_meta"].get("tracestate").is_none(),
                "provider={provider}: {message}"
            );
        }
        // The typed call keeps its ordinary params shape.
        assert!(sent_messages(&outgoing)[2]["params"].get("_meta").is_none());
    }
}

#[tokio::test]
async fn invalid_trace_context_is_never_sent() {
    // `TraceContext` has no unchecked constructor, so a provider that tries
    // to hand over a malformed value ends up with `None`.
    let transport = MockTransport::with_responses(vec![serde_json::json!({})]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .trace_context_provider(|| {
            crate::TraceContext::parse(
                "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
                Some("vendor=value"),
            )
        })
        .connect_simple(transport)
        .await
        .unwrap();

    client.ping().await.unwrap();

    let message = &sent_messages(&outgoing)[0];
    assert!(message["params"]["_meta"].get("traceparent").is_none());
    assert!(message["params"]["_meta"].get("tracestate").is_none());
}

#[tokio::test]
async fn trace_context_does_not_overwrite_caller_meta() {
    let transport = MockTransport::with_responses(vec![
        serde_json::json!({"content": [{"type": "text", "text": "done"}]}),
        serde_json::json!({"content": [{"type": "text", "text": "done"}]}),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .trace_context_provider(|| Some(test_trace_context()))
        .connect_simple(transport)
        .await
        .unwrap();

    let caller_traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00";
    let _: CallToolResult = client
        .request(
            "tools/call",
            &serde_json::json!({
                "name": "work",
                "arguments": {},
                "_meta": {"traceparent": caller_traceparent, "other": 1},
            }),
        )
        .await
        .unwrap();
    // Either key is enough: the provider's pair is not mixed with it.
    let _: CallToolResult = client
        .request(
            "tools/call",
            &serde_json::json!({
                "name": "work",
                "arguments": {},
                "_meta": {"tracestate": "caller=1"},
            }),
        )
        .await
        .unwrap();

    let messages = sent_messages(&outgoing);
    let meta = &messages[0]["params"]["_meta"];
    assert_eq!(meta["traceparent"], caller_traceparent);
    assert_eq!(meta["other"], 1);
    assert!(meta.get("tracestate").is_none(), "{meta}");
    let meta = &messages[1]["params"]["_meta"];
    assert_eq!(meta["tracestate"], "caller=1");
    assert!(meta.get("traceparent").is_none(), "{meta}");
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_requests_carry_trace_context_beside_protocol_meta() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        serde_json::json!({
            "resultType": "complete",
            "tools": [],
            "ttlMs": 0,
            "cacheScope": "private"
        }),
        serde_json::json!({
            "resultType": "complete",
            "content": [{"type": "text", "text": "done"}]
        }),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .trace_context_provider(|| Some(test_trace_context()))
        .connect_simple(transport)
        .await
        .unwrap();

    client.discover("test-client", "1.0.0").await.unwrap();
    client.list_tools().await.unwrap();
    let caller_traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00";
    let _: CallToolResult = client
        .request(
            "tools/call",
            &serde_json::json!({
                "name": "work",
                "arguments": {},
                "_meta": {"traceparent": caller_traceparent},
            }),
        )
        .await
        .unwrap();

    let messages = sent_messages(&outgoing);
    assert_eq!(messages.len(), 3);
    for message in &messages {
        let meta = &message["params"]["_meta"];
        assert_eq!(
            meta["io.modelcontextprotocol/protocolVersion"], "2026-07-28",
            "{message}"
        );
    }
    for message in &messages[..2] {
        let meta = &message["params"]["_meta"];
        assert_eq!(meta["traceparent"], TEST_TRACEPARENT, "{message}");
        assert_eq!(meta["tracestate"], "vendor=value", "{message}");
    }
    assert_eq!(
        messages[2]["params"]["_meta"]["traceparent"], caller_traceparent,
        "a caller-set traceparent must win on the final protocol too"
    );
    assert!(messages[2]["params"]["_meta"].get("tracestate").is_none());
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_subscription_listen_carries_trace_context() {
    let transport = SubscriptionTestTransport::new();
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .trace_context_provider(|| Some(test_trace_context()))
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let handle = client
        .listen_subscriptions(SubscriptionFilter {
            prompts_list_changed: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    drop(handle);

    let messages = sent_messages(&outgoing);
    let listen = messages
        .iter()
        .find(|message| message["method"] == "subscriptions/listen")
        .expect("subscriptions/listen was sent");
    assert_eq!(listen["params"]["_meta"]["traceparent"], TEST_TRACEPARENT);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
struct SubscriptionTestTransport {
    incoming_tx: mpsc::Sender<String>,
    incoming_rx: mpsc::Receiver<String>,
    outgoing: Arc<Mutex<Vec<String>>>,
    connected: Arc<AtomicBool>,
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
impl SubscriptionTestTransport {
    fn new() -> Self {
        let (incoming_tx, incoming_rx) = mpsc::channel(32);
        Self {
            incoming_tx,
            incoming_rx,
            outgoing: Arc::new(Mutex::new(Vec::new())),
            connected: Arc::new(AtomicBool::new(true)),
        }
    }
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[async_trait]
impl ClientTransport for SubscriptionTestTransport {
    async fn send(&mut self, message: &str) -> Result<()> {
        self.outgoing.lock().unwrap().push(message.to_string());
        let value: serde_json::Value =
            serde_json::from_str(message).map_err(|error| Error::Transport(error.to_string()))?;
        let Some(id) = value.get("id").cloned() else {
            return Ok(());
        };
        match value
            .get("method")
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default()
        {
            "server/discover" => {
                self.incoming_tx
                    .send(
                        serde_json::json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "result": final_discover_result()
                        })
                        .to_string(),
                    )
                    .await
                    .map_err(|_| Error::Transport("test channel closed".to_string()))?;
            }
            "subscriptions/listen" => {
                let notifications = value["params"]["notifications"].clone();
                self.incoming_tx
                    .send(
                        serde_json::json!({
                            "jsonrpc": "2.0",
                            "method": notifications::SUBSCRIPTIONS_ACKNOWLEDGED,
                            "params": {
                                "_meta": {
                                    "io.modelcontextprotocol/subscriptionId": id
                                },
                                "notifications": notifications
                            }
                        })
                        .to_string(),
                    )
                    .await
                    .map_err(|_| Error::Transport("test channel closed".to_string()))?;

                if value["params"]["notifications"]["promptsListChanged"]
                    == serde_json::Value::Bool(true)
                {
                    self.incoming_tx
                        .send(
                            serde_json::json!({
                                "jsonrpc": "2.0",
                                "id": id,
                                "result": {
                                    "resultType": "complete",
                                    "_meta": {
                                        "io.modelcontextprotocol/subscriptionId": id
                                    }
                                }
                            })
                            .to_string(),
                        )
                        .await
                        .map_err(|_| Error::Transport("test channel closed".to_string()))?;
                } else {
                    self.incoming_tx
                        .send(
                            serde_json::json!({
                                "jsonrpc": "2.0",
                                "method": notifications::TOOLS_LIST_CHANGED,
                                "params": {
                                    "_meta": {
                                        "io.modelcontextprotocol/subscriptionId": id
                                    }
                                }
                            })
                            .to_string(),
                        )
                        .await
                        .map_err(|_| Error::Transport("test channel closed".to_string()))?;
                }
            }
            _ => {}
        }
        Ok(())
    }

    async fn recv(&mut self) -> Result<Option<String>> {
        Ok(self.incoming_rx.recv().await)
    }

    fn is_connected(&self) -> bool {
        self.connected.load(Ordering::Acquire)
    }

    async fn close(&mut self) -> Result<()> {
        self.connected.store(false, Ordering::Release);
        Ok(())
    }
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_subscriptions_correlate_and_cancel_over_message_transport() {
    let transport = SubscriptionTestTransport::new();
    let outgoing = transport.outgoing.clone();
    let incoming = transport.incoming_tx.clone();
    let received = Arc::new(Mutex::new(Vec::new()));

    struct RecordingHandler(Arc<Mutex<Vec<ServerNotification>>>);

    #[async_trait]
    impl ClientHandler for RecordingHandler {
        async fn on_notification(&self, notification: ServerNotification) {
            self.0.lock().unwrap().push(notification);
        }
    }

    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .request_log_level(LogLevel::Info)
        .connect(transport, RecordingHandler(received.clone()))
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let requested = SubscriptionFilter {
        tools_list_changed: Some(true),
        ..Default::default()
    };
    let mut first = client
        .listen_subscriptions(requested.clone())
        .await
        .unwrap();
    let mut second = client.listen_subscriptions(requested).await.unwrap();
    assert_ne!(first.id(), second.id());
    assert_eq!(
        first.acknowledged().await.unwrap().tools_list_changed,
        Some(true)
    );
    assert_eq!(
        second.acknowledged().await.unwrap().tools_list_changed,
        Some(true)
    );
    let listen_messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .filter(|message: &serde_json::Value| message["method"] == "subscriptions/listen")
        .collect();
    assert_eq!(listen_messages.len(), 2);
    for message in listen_messages {
        assert_eq!(
            message["params"]["_meta"]["io.modelcontextprotocol/logLevel"],
            "info"
        );
    }

    for _ in 0..100 {
        if received
            .lock()
            .unwrap()
            .iter()
            .filter(|notification| {
                matches!(
                    notification,
                    ServerNotification::Subscription {
                        notification,
                        ..
                    } if matches!(notification.as_ref(), ServerNotification::ToolsListChanged)
                )
            })
            .count()
            == 2
        {
            break;
        }
        tokio::task::yield_now().await;
    }
    let subscription_ids: Vec<RequestId> = received
        .lock()
        .unwrap()
        .iter()
        .filter_map(|notification| match notification {
            ServerNotification::Subscription {
                subscription_id,
                notification,
            } if matches!(notification.as_ref(), ServerNotification::ToolsListChanged) => {
                Some(subscription_id.clone())
            }
            _ => None,
        })
        .collect();
    assert_eq!(subscription_ids, [first.id().clone(), second.id().clone()]);

    incoming
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": notifications::CANCELLED,
                "params": {
                    "requestId": 999,
                    "reason": "not a subscription"
                }
            })
            .to_string(),
        )
        .await
        .unwrap();
    tokio::task::yield_now().await;
    assert!(
        !received.lock().unwrap().iter().any(|notification| matches!(
            notification,
            ServerNotification::SubscriptionCancelled { subscription_id, .. }
                if subscription_id == &RequestId::Number(999)
        ))
    );

    let first_id = first.id().clone();
    let second_id = second.id().clone();
    first.cancel().await.unwrap();
    second.cancel().await.unwrap();
    let cancellation_ids: Vec<RequestId> = outgoing
        .lock()
        .unwrap()
        .iter()
        .filter_map(|message| {
            let value: serde_json::Value = serde_json::from_str(message).unwrap();
            (value["method"] == notifications::CANCELLED)
                .then(|| serde_json::from_value(value["params"]["requestId"].clone()).unwrap())
        })
        .collect();
    assert_eq!(cancellation_ids, [first_id, second_id]);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_subscription_observes_graceful_completion() {
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect_simple(SubscriptionTestTransport::new())
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let mut handle = client
        .listen_subscriptions(SubscriptionFilter {
            prompts_list_changed: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    assert_eq!(
        handle.acknowledged().await.unwrap().prompts_list_changed,
        Some(true)
    );
    let expected_id = handle.id().clone();
    let result = handle.wait().await.unwrap();
    assert!(result.result_type.is_complete());
    assert_eq!(result.meta.subscription_id, expected_id);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_subscription_accepts_server_cancellation_only_for_active_listen() {
    let transport = SubscriptionTestTransport::new();
    let incoming = transport.incoming_tx.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let mut handle = client
        .listen_subscriptions(SubscriptionFilter {
            tools_list_changed: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    handle.acknowledged().await.unwrap();
    let id = handle.id().clone();
    incoming
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": notifications::CANCELLED,
                "params": {
                    "requestId": id,
                    "reason": "server shutdown"
                }
            })
            .to_string(),
        )
        .await
        .unwrap();

    let error = handle.wait().await.unwrap_err();
    assert!(
        error
            .to_string()
            .contains("subscription cancelled by server: server shutdown")
    );
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
fn cacheable_tools_result(name: &str) -> serde_json::Value {
    serde_json::json!({
        "resultType": "complete",
        "tools": [{
            "name": name,
            "inputSchema": {
                "type": "object",
                "properties": {}
            }
        }],
        "ttlMs": 60_000,
        "cacheScope": "private"
    })
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_cache_serves_a_fresh_list_without_a_round_trip() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        cacheable_tools_result("cached"),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let first = client.list_tools().await.unwrap();
    let second = client.list_tools().await.unwrap();

    assert_eq!(first.tools[0].name, "cached");
    assert_eq!(second.tools[0].name, "cached");
    assert_eq!(outgoing.lock().unwrap().len(), 2);
    assert_eq!(client.response_cache_len().await, 1);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn disabling_final_cache_forces_each_list_request() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        cacheable_tools_result("first"),
        cacheable_tools_result("second"),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .disable_response_cache()
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    assert_eq!(client.list_tools().await.unwrap().tools[0].name, "first");
    assert_eq!(client.list_tools().await.unwrap().tools[0].name, "second");
    assert_eq!(outgoing.lock().unwrap().len(), 3);
    assert_eq!(client.response_cache_len().await, 0);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn list_changed_notification_invalidates_a_fresh_entry() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        cacheable_tools_result("first"),
        cacheable_tools_result("second"),
    ]);
    let incoming = transport.incoming_tx.clone();
    let outgoing = transport.outgoing.clone();
    let notification_seen = Arc::new(AtomicBool::new(false));
    let handler = NotificationHandler::new().on_tools_changed({
        let notification_seen = notification_seen.clone();
        move || notification_seen.store(true, Ordering::Release)
    });
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect(transport, handler)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();
    assert_eq!(client.list_tools().await.unwrap().tools[0].name, "first");

    incoming
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": notifications::TOOLS_LIST_CHANGED,
                "params": {}
            })
            .to_string(),
        )
        .await
        .unwrap();
    for _ in 0..100 {
        if notification_seen.load(Ordering::Acquire) {
            break;
        }
        tokio::task::yield_now().await;
    }
    assert!(notification_seen.load(Ordering::Acquire));

    assert_eq!(client.list_tools().await.unwrap().tools[0].name, "second");
    assert_eq!(outgoing.lock().unwrap().len(), 3);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn rotating_private_partition_refetches_resource() {
    let resource_result = |text: &str| {
        serde_json::json!({
            "resultType": "complete",
            "contents": [{
                "uri": "config://app",
                "text": text
            }],
            "ttlMs": 60_000,
            "cacheScope": "private"
        })
    };
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        resource_result("principal-a"),
        resource_result("principal-b"),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .response_cache(ClientCacheConfig::default().with_partition("principal-a"))
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    assert_eq!(
        client
            .read_resource("config://app")
            .await
            .unwrap()
            .first_text(),
        Some("principal-a")
    );
    client.set_cache_partition("principal-b").await;
    assert_eq!(
        client
            .read_resource("config://app")
            .await
            .unwrap()
            .first_text(),
        Some("principal-b")
    );
    assert_eq!(outgoing.lock().unwrap().len(), 3);
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_tool_call_refreshes_stale_schema_and_retries_once() {
    let transport = MockTransport::with_replies(vec![
        MockReply::Result(final_discover_result()),
        MockReply::Result(cacheable_tools_result("changing-tool")),
        MockReply::Error(JsonRpcError::header_mismatch("stale x-mcp-header mapping")),
        MockReply::Result(cacheable_tools_result("changing-tool")),
        MockReply::Result(serde_json::json!({
            "resultType": "complete",
            "content": [{"type": "text", "text": "retried"}]
        })),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .request_log_level(LogLevel::Warning)
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();
    client.list_tools().await.unwrap();

    let result = client
        .call_tool("changing-tool", serde_json::json!({}))
        .await
        .unwrap();
    assert_eq!(result.first_text(), Some("retried"));

    let messages: Vec<serde_json::Value> = outgoing
        .lock()
        .unwrap()
        .iter()
        .map(|message| serde_json::from_str(message).unwrap())
        .collect();
    let methods: Vec<&str> = messages
        .iter()
        .map(|message| message["method"].as_str().unwrap())
        .collect();
    assert_eq!(
        methods,
        [
            "server/discover",
            "tools/list",
            "tools/call",
            "tools/list",
            "tools/call"
        ]
    );
    for message in messages {
        assert_eq!(
            message["params"]["_meta"]["io.modelcontextprotocol/logLevel"], "warning",
            "every request in the retry flow must retain the configured threshold"
        );
    }
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_call_tool_as_task_sends_no_legacy_task_parameter() {
    let transport = MockTransport::with_responses(vec![
        final_discover_result(),
        serde_json::json!({
            "resultType": "task",
            "taskId": "task-final",
            "status": "working",
            "createdAt": "2026-07-31T00:00:00Z",
            "lastUpdatedAt": "2026-07-31T00:00:00Z",
            "ttlMs": null,
            "pollIntervalMs": 50
        }),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .with_tasks()
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();

    let created = client
        .call_tool_as_task("long-tool", serde_json::json!({}), None)
        .await
        .unwrap();
    assert_eq!(created.task.task_id, "task-final");

    let messages = outgoing.lock().unwrap();
    let call: serde_json::Value = serde_json::from_str(&messages[1]).unwrap();
    assert_eq!(call["method"], "tools/call");
    assert!(
        call["params"].get("task").is_none(),
        "final tools/call leaked the legacy task parameter: {call}"
    );
    assert!(
        call["params"]["_meta"]["io.modelcontextprotocol/clientCapabilities"]["extensions"]
            .get(crate::protocol::TASKS_EXTENSION_ID)
            .is_some()
    );
}

#[cfg(any(feature = "protocol-2026-07-28", feature = "stateless"))]
#[tokio::test]
async fn final_stale_schema_retry_is_bounded() {
    let transport = MockTransport::with_replies(vec![
        MockReply::Result(final_discover_result()),
        MockReply::Result(cacheable_tools_result("changing-tool")),
        MockReply::Error(JsonRpcError::header_mismatch("first rejection")),
        MockReply::Result(cacheable_tools_result("changing-tool")),
        MockReply::Error(JsonRpcError::header_mismatch("second rejection")),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::builder()
        .protocol_support(ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect_simple(transport)
        .await
        .unwrap();
    client.discover("test-client", "1.0.0").await.unwrap();
    client.list_tools().await.unwrap();

    let error = client
        .call_tool("changing-tool", serde_json::json!({}))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        Error::JsonRpc(error) if error.code == McpErrorCode::HeaderMismatch.code()
    ));
    assert_eq!(outgoing.lock().unwrap().len(), 5);
}

#[tokio::test]
async fn legacy_tool_call_does_not_retry_a_header_mismatch() {
    let transport = MockTransport::with_replies(vec![
        MockReply::Result(mock_initialize_response()),
        MockReply::Error(JsonRpcError::header_mismatch("legacy rejection")),
    ]);
    let outgoing = transport.outgoing.clone();
    let client = McpClient::connect(transport).await.unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    let error = client
        .call_tool("changing-tool", serde_json::json!({}))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        Error::JsonRpc(error) if error.code == McpErrorCode::HeaderMismatch.code()
    ));
    let messages = outgoing.lock().unwrap();
    assert_eq!(messages.len(), 3);
    assert!(
        !messages
            .iter()
            .any(|message| message.contains("tools/list"))
    );
}

#[test]
fn stale_tool_schema_errors_are_pre_execution_protocol_errors() {
    for error in [
        Error::JsonRpc(JsonRpcError::header_mismatch("mismatch")),
        Error::JsonRpc(JsonRpcError::method_not_found("tool")),
        Error::JsonRpc(JsonRpcError::invalid_params("arguments")),
    ] {
        assert!(is_stale_tool_schema_error(&error));
    }
    assert!(!is_stale_tool_schema_error(&Error::JsonRpc(
        JsonRpcError::internal_error("executed")
    )));
}

#[tokio::test]
async fn test_list_tools() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "tools": [
                {
                    "name": "test_tool",
                    "description": "A test tool",
                    "inputSchema": {
                        "type": "object",
                        "properties": {}
                    }
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let tools = client.list_tools().await.unwrap();

    assert_eq!(tools.tools.len(), 1);
    assert_eq!(tools.tools[0].name, "test_tool");
}

#[tokio::test]
async fn test_call_tool() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "content": [
                {
                    "type": "text",
                    "text": "Tool result"
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client
        .call_tool("test_tool", serde_json::json!({"arg": "value"}))
        .await
        .unwrap();

    assert!(!result.content.is_empty());
}

#[tokio::test]
async fn test_list_resources() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "resources": [
                {
                    "uri": "file://test.txt",
                    "name": "Test File"
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let resources = client.list_resources().await.unwrap();

    assert_eq!(resources.resources.len(), 1);
    assert_eq!(resources.resources[0].uri, "file://test.txt");
}

#[tokio::test]
async fn test_read_resource() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "contents": [
                {
                    "uri": "file://test.txt",
                    "text": "File contents"
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client.read_resource("file://test.txt").await.unwrap();

    assert_eq!(result.contents.len(), 1);
    assert_eq!(result.contents[0].text.as_deref(), Some("File contents"));
}

#[tokio::test]
async fn test_list_prompts() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "prompts": [
                {
                    "name": "test_prompt",
                    "description": "A test prompt"
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let prompts = client.list_prompts().await.unwrap();

    assert_eq!(prompts.prompts.len(), 1);
    assert_eq!(prompts.prompts[0].name, "test_prompt");
}

#[tokio::test]
async fn test_get_prompt() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "messages": [
                {
                    "role": "user",
                    "content": {
                        "type": "text",
                        "text": "Prompt message"
                    }
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client.get_prompt("test_prompt", None).await.unwrap();

    assert_eq!(result.messages.len(), 1);
}

#[tokio::test]
async fn test_ping() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({}),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client.ping().await;
    assert!(result.is_ok());
}

#[tokio::test]
async fn test_with_roots() {
    let roots = vec![Root::new("file:///test")];
    let client = McpClient::builder()
        .with_roots(roots)
        .connect_simple(MockTransport::with_responses(vec![]))
        .await
        .unwrap();

    let current_roots = client.roots().await;
    assert_eq!(current_roots.len(), 1);
}

#[tokio::test]
async fn test_roots_management() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
    ]))
    .await
    .unwrap();

    // Initially no roots
    assert!(client.roots().await.is_empty());

    // Add a root before initialization (no notification sent)
    client.add_root(Root::new("file:///project")).await.unwrap();
    assert_eq!(client.roots().await.len(), 1);

    // Initialize
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Remove a root
    let removed = client.remove_root("file:///project").await.unwrap();
    assert!(removed);
    assert!(client.roots().await.is_empty());

    // Try to remove non-existent root
    let not_removed = client.remove_root("file:///nonexistent").await.unwrap();
    assert!(!not_removed);
}

#[tokio::test]
async fn test_list_roots() {
    let roots = vec![
        Root::new("file:///project1"),
        Root::with_name("file:///project2", "Project 2"),
    ];
    let client = McpClient::builder()
        .with_roots(roots)
        .connect_simple(MockTransport::with_responses(vec![]))
        .await
        .unwrap();

    let result = client.list_roots().await;
    assert_eq!(result.roots.len(), 2);
    assert_eq!(result.roots[1].name, Some("Project 2".to_string()));
}

#[test]
fn test_builder_with_sampling() {
    let builder = McpClientBuilder::new().with_sampling();
    assert!(builder.capabilities.sampling.is_some());
}

#[test]
fn test_builder_with_elicitation() {
    let builder = McpClientBuilder::new().with_elicitation();
    assert!(builder.capabilities.elicitation.is_some());
}

#[test]
fn builder_adds_protocol_extension_without_replacing_other_capabilities() {
    let extension = crate::ExtensionDeclaration::new(
        "com.example/rendering",
        serde_json::json!({"formats": ["html"]}),
    )
    .unwrap();
    let builder = McpClientBuilder::new()
        .with_sampling()
        .with_protocol_extension(extension);

    assert!(builder.capabilities.sampling.is_some());
    assert_eq!(
        builder.capabilities.extensions.as_ref().unwrap()["com.example/rendering"]["formats"][0],
        "html"
    );
}

#[test]
fn test_builder_chaining() {
    let builder = McpClientBuilder::new()
        .with_sampling()
        .with_elicitation()
        .with_roots(vec![Root::new("file:///project")]);
    assert!(builder.capabilities.sampling.is_some());
    assert!(builder.capabilities.elicitation.is_some());
    assert!(builder.capabilities.roots.is_some());
}

#[tokio::test]
async fn test_bidirectional_sampling_round_trip() {
    use crate::protocol::{
        ContentRole, CreateMessageParams, CreateMessageResult, SamplingContent,
        SamplingContentOrArray,
    };

    // A handler that records whether handle_create_message was called
    struct RecordingHandler {
        called: Arc<AtomicBool>,
    }

    #[async_trait]
    impl ClientHandler for RecordingHandler {
        async fn handle_create_message(
            &self,
            _params: CreateMessageParams,
        ) -> std::result::Result<CreateMessageResult, tower_mcp_types::JsonRpcError> {
            self.called.store(true, Ordering::SeqCst);
            Ok(CreateMessageResult {
                content: SamplingContentOrArray::Single(SamplingContent::Text {
                    text: "test response".to_string(),
                    annotations: None,
                    meta: None,
                }),
                model: "test-model".to_string(),
                role: ContentRole::Assistant,
                stop_reason: Some("end_turn".to_string()),
                meta: None,
            })
        }
    }

    let called = Arc::new(AtomicBool::new(false));
    let handler = RecordingHandler {
        called: called.clone(),
    };

    // Build a mock transport, keeping a clone of incoming_tx so we can
    // inject a server-initiated request after the transport is consumed.
    let (inject_tx, rx) = mpsc::channel::<String>(32);
    let responses = vec![mock_initialize_response()];
    let inject_tx_clone = inject_tx.clone();

    let transport = MockTransport {
        responses: Arc::new(Mutex::new(
            responses.into_iter().map(MockReply::Result).collect(),
        )),
        response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        incoming_tx: inject_tx,
        incoming_rx: rx,
        outgoing: Arc::new(Mutex::new(Vec::new())),
        connected: Arc::new(AtomicBool::new(true)),
        fail_notification_sends: Arc::new(AtomicBool::new(false)),
    };

    let client = McpClient::builder()
        .with_sampling()
        .connect(transport, handler)
        .await
        .unwrap();

    // Initialize the client (this sends initialize request + notification)
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Inject a server-initiated sampling/createMessage request
    let sampling_request = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 100,
        "method": "sampling/createMessage",
        "params": {
            "messages": [
                {
                    "role": "user",
                    "content": {
                        "type": "text",
                        "text": "Hello"
                    }
                }
            ],
            "maxTokens": 100
        }
    });
    inject_tx_clone
        .send(sampling_request.to_string())
        .await
        .unwrap();

    // Give the background loop time to process
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // Verify the handler was called
    assert!(
        called.load(Ordering::SeqCst),
        "handle_create_message should have been called"
    );
}

// =========================================================================
// Server-initiated request cancellation (#1483)
// =========================================================================

#[tokio::test]
async fn cancelled_server_request_gets_no_reply_and_the_handler_observes_it() {
    use crate::protocol::{ElicitRequestParams, ElicitResult};

    // Waits on its own CancellationToken instead of completing on its
    // own, so the test proves the reply is suppressed rather than merely
    // winning a race against a fast handler.
    struct WaitingHandler {
        started: Arc<tokio::sync::Notify>,
        saw_cancel: Arc<AtomicBool>,
    }

    #[async_trait]
    impl ClientHandler for WaitingHandler {
        async fn handle_elicit(
            &self,
            _params: ElicitRequestParams,
        ) -> std::result::Result<ElicitResult, JsonRpcError> {
            self.started.notify_one();
            let token = self
                .current_cancellation_token()
                .expect("token available while handling elicitation/create");
            token.cancelled().await;
            self.saw_cancel.store(true, Ordering::SeqCst);
            Err(JsonRpcError::internal_error("cancelled"))
        }
    }

    let started = Arc::new(tokio::sync::Notify::new());
    let saw_cancel = Arc::new(AtomicBool::new(false));
    let handler = WaitingHandler {
        started: started.clone(),
        saw_cancel: saw_cancel.clone(),
    };

    let (inject_tx, rx) = mpsc::channel::<String>(32);
    let outgoing = Arc::new(Mutex::new(Vec::new()));
    let transport = MockTransport {
        responses: Arc::new(Mutex::new(vec![MockReply::Result(
            mock_initialize_response(),
        )])),
        response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        incoming_tx: inject_tx.clone(),
        incoming_rx: rx,
        outgoing: outgoing.clone(),
        connected: Arc::new(AtomicBool::new(true)),
        fail_notification_sends: Arc::new(AtomicBool::new(false)),
    };

    let client = McpClient::builder()
        .with_elicitation()
        .connect(transport, handler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    inject_tx
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 100,
                "method": "elicitation/create",
                "params": {
                    "message": "need input",
                    "requestedSchema": {"type": "object", "properties": {}, "required": []}
                }
            })
            .to_string(),
        )
        .await
        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(5), started.notified())
        .await
        .expect("handler should have started");

    inject_tx
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/cancelled",
                "params": {"requestId": 100}
            })
            .to_string(),
        )
        .await
        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(5), async {
        while !saw_cancel.load(Ordering::SeqCst) {
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("handler should have observed cancellation via current_cancellation_token()");

    // Give the outcome time to reach the loop and (if it were going to)
    // get written to the transport.
    tokio::time::sleep(std::time::Duration::from_millis(50)).await;

    let sent = outgoing.lock().unwrap().clone();
    assert!(
        !sent.iter().any(|msg| {
            serde_json::from_str::<serde_json::Value>(msg)
                .ok()
                .and_then(|value| value.get("id").cloned())
                == Some(serde_json::json!(100))
        }),
        "a reply for the cancelled request must not be sent: {sent:?}"
    );
}

#[tokio::test]
async fn client_handler_can_call_back_into_the_client_without_deadlocking() {
    // Before #1483, the message loop awaited a server-initiated request's
    // handler inline, so a handler that called back into the same client
    // (e.g. `ping()`) would deadlock: the reentrant call's own response
    // could never reach the loop that was blocked waiting on the handler.
    struct ReentrantHandler {
        client: Arc<tokio::sync::OnceCell<McpClient>>,
        completed: Arc<AtomicBool>,
    }

    #[async_trait]
    impl ClientHandler for ReentrantHandler {
        async fn handle_list_roots(&self) -> std::result::Result<ListRootsResult, JsonRpcError> {
            let client = self
                .client
                .get()
                .expect("client set before the server request arrives");
            client
                .ping()
                .await
                .map_err(|error| JsonRpcError::internal_error(error.to_string()))?;
            self.completed.store(true, Ordering::SeqCst);
            Ok(ListRootsResult {
                roots: vec![],
                meta: None,
            })
        }
    }

    let client_cell = Arc::new(tokio::sync::OnceCell::new());
    let completed = Arc::new(AtomicBool::new(false));
    let handler = ReentrantHandler {
        client: client_cell.clone(),
        completed: completed.clone(),
    };

    let (inject_tx, rx) = mpsc::channel::<String>(32);
    let transport = MockTransport {
        responses: Arc::new(Mutex::new(vec![
            MockReply::Result(mock_initialize_response()),
            MockReply::Result(serde_json::json!({})),
        ])),
        response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        incoming_tx: inject_tx.clone(),
        incoming_rx: rx,
        outgoing: Arc::new(Mutex::new(Vec::new())),
        connected: Arc::new(AtomicBool::new(true)),
        fail_notification_sends: Arc::new(AtomicBool::new(false)),
    };

    let client = McpClient::builder()
        .connect(transport, handler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();
    client_cell
        .set(client)
        .unwrap_or_else(|_| panic!("client cell is set exactly once"));

    inject_tx
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 200,
                "method": "roots/list"
            })
            .to_string(),
        )
        .await
        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(5), async {
        while !completed.load(Ordering::SeqCst) {
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("reentrant call should complete without deadlocking the message loop");
}

#[tokio::test]
async fn cancellation_id_matching_is_exact_not_stringified() {
    // A cancellation whose requestId is the string "1" must not cancel a
    // server request with numeric id 1: JSON-RPC treats them as distinct
    // ids, unlike the deliberately lenient fallback response correlation
    // uses (rmcp-compat note in handle_response).
    struct WaitingHandler {
        started: Arc<tokio::sync::Notify>,
        finished: Arc<AtomicBool>,
    }

    #[async_trait]
    impl ClientHandler for WaitingHandler {
        async fn handle_list_roots(&self) -> std::result::Result<ListRootsResult, JsonRpcError> {
            self.started.notify_one();
            // Give a wrongly-matched cancellation a chance to land before
            // this returns.
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            let cancelled = self
                .current_cancellation_token()
                .map(|token| token.is_cancelled())
                .unwrap_or(false);
            self.finished.store(true, Ordering::SeqCst);
            if cancelled {
                Err(JsonRpcError::internal_error(
                    "wrongly cancelled by a stringified id",
                ))
            } else {
                Ok(ListRootsResult {
                    roots: vec![],
                    meta: None,
                })
            }
        }
    }

    let started = Arc::new(tokio::sync::Notify::new());
    let finished = Arc::new(AtomicBool::new(false));
    let handler = WaitingHandler {
        started: started.clone(),
        finished: finished.clone(),
    };

    let (inject_tx, rx) = mpsc::channel::<String>(32);
    let outgoing = Arc::new(Mutex::new(Vec::new()));
    let transport = MockTransport {
        responses: Arc::new(Mutex::new(vec![MockReply::Result(
            mock_initialize_response(),
        )])),
        response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        incoming_tx: inject_tx.clone(),
        incoming_rx: rx,
        outgoing: outgoing.clone(),
        connected: Arc::new(AtomicBool::new(true)),
        fail_notification_sends: Arc::new(AtomicBool::new(false)),
    };

    let client = McpClient::builder()
        .connect(transport, handler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Server request with numeric id 1.
    inject_tx
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "roots/list"
            })
            .to_string(),
        )
        .await
        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(5), started.notified())
        .await
        .expect("handler should have started");

    // A cancellation for the *string* id "1" must not match it.
    inject_tx
        .send(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/cancelled",
                "params": {"requestId": "1"}
            })
            .to_string(),
        )
        .await
        .unwrap();

    tokio::time::timeout(std::time::Duration::from_secs(5), async {
        while !finished.load(Ordering::SeqCst) {
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("handler should finish on its own");

    // Give the reply time to reach the transport.
    tokio::time::sleep(std::time::Duration::from_millis(50)).await;

    let sent = outgoing.lock().unwrap().clone();
    assert!(
        sent.iter().any(|msg| {
            let value: serde_json::Value = serde_json::from_str(msg).unwrap();
            value.get("id") == Some(&serde_json::json!(1)) && value.get("result").is_some()
        }),
        "the reply for numeric id 1 should still be sent: {sent:?}"
    );
}

// =========================================================================
// Server request with an invalid id (#1558)
// =========================================================================

#[tokio::test]
async fn server_request_with_a_non_integer_non_string_id_gets_an_invalid_request_error() {
    struct CountingHandler {
        roots_calls: Arc<std::sync::atomic::AtomicUsize>,
    }

    #[async_trait]
    impl ClientHandler for CountingHandler {
        async fn handle_list_roots(&self) -> std::result::Result<ListRootsResult, JsonRpcError> {
            self.roots_calls.fetch_add(1, Ordering::SeqCst);
            Ok(ListRootsResult {
                roots: vec![],
                meta: None,
            })
        }
    }

    let roots_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let (inject_tx, rx) = mpsc::channel::<String>(32);
    let outgoing = Arc::new(Mutex::new(Vec::new()));
    let transport = MockTransport {
        responses: Arc::new(Mutex::new(vec![MockReply::Result(
            mock_initialize_response(),
        )])),
        response_idx: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        incoming_tx: inject_tx.clone(),
        incoming_rx: rx,
        outgoing: outgoing.clone(),
        connected: Arc::new(AtomicBool::new(true)),
        fail_notification_sends: Arc::new(AtomicBool::new(false)),
    };
    let client = McpClient::builder()
        .connect(
            transport,
            CountingHandler {
                roots_calls: roots_calls.clone(),
            },
        )
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();
    let before = outgoing.lock().unwrap().len();

    // A float, a boolean, an object, an array, null, and an integer too
    // large for i64: none is a string or a signed integer.
    let bad_ids = [
        "1.5",
        "true",
        "{\"a\":1}",
        "[1]",
        "null",
        "18446744073709551615",
    ];
    for id in bad_ids {
        for method in ["ping", "roots/list"] {
            inject_tx
                .send(format!(
                    r#"{{"jsonrpc":"2.0","id":{id},"method":"{method}"}}"#
                ))
                .await
                .unwrap();
        }
    }
    // A well-formed integer id still gets its normal reply.
    inject_tx
        .send(r#"{"jsonrpc":"2.0","id":7,"method":"ping"}"#.to_string())
        .await
        .unwrap();

    let expected = bad_ids.len() * 2 + 1;
    tokio::time::timeout(std::time::Duration::from_secs(5), async {
        while outgoing.lock().unwrap().len() - before < expected {
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("every request should be answered");

    let sent: Vec<serde_json::Value> = outgoing.lock().unwrap()[before..]
        .iter()
        .map(|msg| serde_json::from_str(msg).unwrap())
        .collect();
    let (pongs, errors): (Vec<_>, Vec<_>) =
        sent.iter().partition(|msg| msg.get("result").is_some());
    assert_eq!(errors.len(), bad_ids.len() * 2, "sent: {sent:?}");
    for error in errors {
        assert_eq!(error["id"], serde_json::Value::Null, "{error}");
        assert_eq!(error["error"]["code"], -32600, "{error}");
    }
    assert_eq!(pongs.len(), 1, "sent: {sent:?}");
    assert_eq!(pongs[0]["id"], 7);
    assert_eq!(pongs[0]["result"], serde_json::json!({}));
    assert_eq!(
        roots_calls.load(Ordering::SeqCst),
        0,
        "the handler must not run for an invalid request"
    );
}

#[tokio::test]
async fn test_list_resource_templates() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "resourceTemplates": [
                {
                    "uriTemplate": "file:///{path}",
                    "name": "File Template",
                    "description": "A file template"
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client.list_resource_templates().await.unwrap();

    assert_eq!(result.resource_templates.len(), 1);
    assert_eq!(result.resource_templates[0].name, "File Template");
}

#[tokio::test]
async fn test_list_all_tools_single_page() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "tools": [
                {
                    "name": "tool_a",
                    "description": "Tool A",
                    "inputSchema": { "type": "object", "properties": {} }
                },
                {
                    "name": "tool_b",
                    "description": "Tool B",
                    "inputSchema": { "type": "object", "properties": {} }
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let tools = client.list_all_tools().await.unwrap();

    assert_eq!(tools.len(), 2);
    assert_eq!(tools[0].name, "tool_a");
    assert_eq!(tools[1].name, "tool_b");
}

#[tokio::test]
async fn test_list_all_tools_paginated() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        // First page with a next_cursor
        serde_json::json!({
            "tools": [
                {
                    "name": "tool_a",
                    "description": "Tool A",
                    "inputSchema": { "type": "object", "properties": {} }
                }
            ],
            "nextCursor": "page2"
        }),
        // Second page with no next_cursor
        serde_json::json!({
            "tools": [
                {
                    "name": "tool_b",
                    "description": "Tool B",
                    "inputSchema": { "type": "object", "properties": {} }
                }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let tools = client.list_all_tools().await.unwrap();

    assert_eq!(tools.len(), 2);
    assert_eq!(tools[0].name, "tool_a");
    assert_eq!(tools[1].name, "tool_b");
}

#[tokio::test]
async fn test_call_tool_text_success() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "content": [
                { "type": "text", "text": "Hello " },
                { "type": "text", "text": "World" }
            ]
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let text = client
        .call_tool_text("test_tool", serde_json::json!({}))
        .await
        .unwrap();

    assert_eq!(text, "Hello World");
}

#[tokio::test]
async fn test_call_tool_text_error() {
    let client = McpClient::connect(MockTransport::with_responses(vec![
        mock_initialize_response(),
        serde_json::json!({
            "content": [
                { "type": "text", "text": "something went wrong" }
            ],
            "isError": true
        }),
    ]))
    .await
    .unwrap();

    client.initialize("test-client", "1.0.0").await.unwrap();
    let result = client
        .call_tool_text("test_tool", serde_json::json!({}))
        .await;

    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.to_string().contains("something went wrong"),
        "Error message should contain tool error text, got: {}",
        err
    );
}

#[tokio::test]
async fn test_server_notification_parsing() {
    let notification = parse_server_notification("notifications/tools/list_changed", None);
    assert!(matches!(notification, ServerNotification::ToolsListChanged));

    let notification = parse_server_notification("notifications/resources/list_changed", None);
    assert!(matches!(
        notification,
        ServerNotification::ResourcesListChanged
    ));

    let notification = parse_server_notification(
        "notifications/resources/updated",
        Some(serde_json::json!({"uri": "file:///test"})),
    );
    match notification {
        ServerNotification::ResourceUpdated { uri } => {
            assert_eq!(uri, "file:///test");
        }
        _ => panic!("Expected ResourceUpdated"),
    }

    let notification =
        parse_server_notification("custom/notification", Some(serde_json::json!({"data": 42})));
    match notification {
        ServerNotification::Unknown { method, params } => {
            assert_eq!(method, "custom/notification");
            assert!(params.is_some());
        }
        _ => panic!("Expected Unknown"),
    }

    let notification = parse_server_notification(
        notifications::SUBSCRIPTIONS_ACKNOWLEDGED,
        Some(serde_json::json!({
            "_meta": {
                "io.modelcontextprotocol/subscriptionId": 7
            },
            "notifications": {
                "toolsListChanged": true
            }
        })),
    );
    assert!(matches!(
        notification,
        ServerNotification::SubscriptionAcknowledged {
            subscription_id: RequestId::Number(7),
            ..
        }
    ));

    let notification = parse_server_notification(
        notifications::TOOLS_LIST_CHANGED,
        Some(serde_json::json!({
            "_meta": {
                "io.modelcontextprotocol/subscriptionId": "stream-a"
            }
        })),
    );
    assert!(matches!(
        notification,
        ServerNotification::Subscription {
            subscription_id: RequestId::String(id),
            notification,
        } if id == "stream-a"
            && matches!(notification.as_ref(), ServerNotification::ToolsListChanged)
    ));

    let notification = parse_server_notification(
        notifications::CANCELLED,
        Some(serde_json::json!({
            "requestId": 7,
            "reason": "done"
        })),
    );
    assert!(matches!(
        notification,
        ServerNotification::SubscriptionCancelled {
            subscription_id: RequestId::Number(7),
            reason: Some(reason),
        } if reason == "done"
    ));

    let notification = parse_server_notification(
        notifications::TASK_STATUS_CHANGED,
        Some(serde_json::json!({
            "taskId": "legacy-task",
            "status": "completed",
            "createdAt": "2026-08-02T00:00:00Z",
            "lastUpdatedAt": "2026-08-02T00:00:01Z",
            "ttl": null
        })),
    );
    assert!(matches!(
        notification,
        ServerNotification::TaskStatusChanged(TaskStatusParams {
            task_id,
            status: crate::protocol::TaskStatus::Completed,
            ..
        }) if task_id == "legacy-task"
    ));

    let notification = parse_server_notification(
        notifications::TASK_STATUS_CHANGED,
        Some(serde_json::json!({
            "taskId": "final-task",
            "status": "cancelled",
            "createdAt": "2026-08-02T00:00:00Z",
            "lastUpdatedAt": "2026-08-02T00:00:01Z",
            "ttlMs": null,
            "_meta": {
                "io.modelcontextprotocol/subscriptionId": "task-stream"
            }
        })),
    );
    assert!(matches!(
        notification,
        ServerNotification::Subscription {
            subscription_id: RequestId::String(id),
            notification,
        } if id == "task-stream"
            && matches!(
                notification.as_ref(),
                ServerNotification::FinalTaskStatusChanged(params)
                    if params.task.task_id() == "final-task"
                        && params.task.status() == crate::protocol::TaskStatus::Cancelled
            )
    ));
}

// =========================================================================
// handle_response ID correlation
// =========================================================================

fn pending_with(
    ids: &[RequestId],
) -> (
    HashMap<RequestId, PendingRequest>,
    Vec<oneshot::Receiver<Result<serde_json::Value>>>,
) {
    let mut map = HashMap::new();
    let mut rxs = Vec::new();
    for id in ids {
        let (tx, rx) = oneshot::channel();
        map.insert(
            id.clone(),
            PendingRequest {
                method: "test".to_string(),
                response_tx: tx,
                acknowledgment_tx: None,
            },
        );
        rxs.push(rx);
    }
    (map, rxs)
}

#[tokio::test]
async fn test_stringified_numeric_response_id_correlates() {
    // rmcp #1021 analog: a numeric request id 42 answered with a
    // stringified id "42" still correlates.
    let (mut pending, mut rxs) = pending_with(&[RequestId::Number(42)]);

    let response = serde_json::json!({
        "jsonrpc": "2.0",
        "id": "42",
        "result": {"ok": true}
    });
    handle_response(&response, &mut pending);

    assert!(pending.is_empty(), "pending request should be resolved");
    let result = rxs.remove(0).await.unwrap().unwrap();
    assert_eq!(result, serde_json::json!({"ok": true}));
}

#[tokio::test]
async fn test_integral_float_response_id_correlates() {
    // rmcp #1345 analog: a server that writes every number in float form
    // answers request 42 with 42.0.
    let (mut pending, mut rxs) = pending_with(&[RequestId::Number(42)]);

    let response = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 42.0,
        "result": {"ok": true}
    });
    handle_response(&response, &mut pending);

    assert!(pending.is_empty(), "pending request should be resolved");
    let result = rxs.remove(0).await.unwrap().unwrap();
    assert_eq!(result, serde_json::json!({"ok": true}));
}

#[tokio::test]
async fn test_exact_string_id_takes_precedence() {
    // A genuine string id "42" must match exactly and win over the
    // numeric interpretation when both are pending.
    let (mut pending, mut rxs) =
        pending_with(&[RequestId::String("42".to_string()), RequestId::Number(42)]);

    let response = serde_json::json!({
        "jsonrpc": "2.0",
        "id": "42",
        "result": {"which": "string"}
    });
    handle_response(&response, &mut pending);

    // The string entry resolved; the numeric entry is still pending.
    assert_eq!(pending.len(), 1);
    assert!(pending.contains_key(&RequestId::Number(42)));
    let result = rxs.remove(0).await.unwrap().unwrap();
    assert_eq!(result, serde_json::json!({"which": "string"}));
}

#[tokio::test]
async fn test_non_numeric_string_id_does_not_correlate() {
    // A string id that is not the string form of the pending numeric
    // id must not resolve it.
    let (mut pending, _rxs) = pending_with(&[RequestId::Number(42)]);

    let response = serde_json::json!({
        "jsonrpc": "2.0",
        "id": "not-a-number",
        "result": {}
    });
    handle_response(&response, &mut pending);

    assert_eq!(pending.len(), 1, "numeric request should stay pending");
}

#[tokio::test]
async fn test_response_with_both_result_and_error_fails_as_malformed() {
    // JSON-RPC 2.0 requires exactly one of `result`/`error`. `error` is
    // checked first below, so without the #1481 guard this would resolve
    // as the server's error response instead of failing the call outright.
    let (mut pending, mut rxs) = pending_with(&[RequestId::Number(1)]);

    let response = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "result": {"ok": true},
        "error": {"code": -32603, "message": "x"}
    });
    handle_response(&response, &mut pending);

    assert!(pending.is_empty(), "pending request should be resolved");
    let result = rxs.remove(0).await.unwrap();
    match result {
        Err(Error::Transport(msg)) => {
            assert!(
                msg.contains("both result and error"),
                "unexpected message: {msg}"
            );
        }
        other => panic!("expected Error::Transport, got: {other:?}"),
    }
}