tower-mcp 0.24.0

//! Integration tests for HttpClientTransport.
//!
//! These tests spin up an in-process HTTP server and connect to it
//! using HttpClientTransport + McpClient.

#![cfg(all(feature = "http", feature = "http-client"))]

use async_trait::async_trait;
use schemars::JsonSchema;
use serde::Deserialize;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tower_mcp::{
    CallToolResult, ClientHandler, CompleteParams, CompleteResult, CompletionReference, Content,
    ContentRole, CreateMessageParams, CreateMessageResult, ElicitFieldValue, ElicitRequestParams,
    ElicitResult, GetPromptResult, HttpTransport, LogLevel, LoggingMessageParams, McpClientBuilder,
    McpRouter, NoParams, NotificationHandler, PromptBuilder, PromptMessage, PromptRole,
    ReadResourceResult, ResourceBuilder, ResourceContent, ResourceTemplateBuilder, Root,
    SamplingContent, SamplingContentOrArray, SamplingMessage, SubscriptionFilter, ToolBuilder,
    client::{HttpClientConfig, HttpClientTransport, HttpStatusDetails, McpClient},
    extract::{Context, RawArgs},
    transport::http::SessionConfig,
};
use tower_mcp_types::JsonRpcError;

#[derive(Debug, Deserialize, JsonSchema)]
struct EchoInput {
    message: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct AddInput {
    a: i64,
    b: i64,
}

/// Create a test router with basic tools and resources.
fn test_router() -> McpRouter {
    let echo = ToolBuilder::new("echo")
        .description("Echo a message")
        .handler(|input: EchoInput| async move { Ok(CallToolResult::text(input.message)) })
        .build();

    let add = ToolBuilder::new("add")
        .description("Add two numbers")
        .handler(|input: AddInput| async move {
            Ok(CallToolResult::text(format!("{}", input.a + input.b)))
        })
        .build();

    let config = ResourceBuilder::new("config://app")
        .name("App Config")
        .description("Application configuration")
        .text(r#"{"debug": true}"#);

    McpRouter::new()
        .server_info("test-http-server", "1.0.0")
        .tool(echo)
        .tool(add)
        .resource(config)
}

/// Start an HTTP server on a random available port and return the URL.
async fn start_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    // Give the server a moment to start
    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

#[tokio::test]
async fn test_http_client_initialize() {
    let (url, _server) = start_server().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");
    assert_eq!(info.server_info.version, "1.0.0");
    assert!(client.is_initialized());

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_list_tools() {
    let (url, _server) = start_server().await;

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

    let tools = client.list_tools().await.unwrap();
    assert_eq!(tools.tools.len(), 2);

    let tool_names: Vec<&str> = tools.tools.iter().map(|t| t.name.as_str()).collect();
    assert!(tool_names.contains(&"echo"));
    assert!(tool_names.contains(&"add"));

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_call_tool() {
    let (url, _server) = start_server().await;

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

    // Call echo
    let result = client
        .call_tool(
            "echo",
            serde_json::json!({"message": "Hello from HTTP client!"}),
        )
        .await
        .unwrap();
    assert!(!result.content.is_empty());

    // Call add
    let result = client
        .call_tool("add", serde_json::json!({"a": 42, "b": 58}))
        .await
        .unwrap();
    assert!(!result.content.is_empty());

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_list_resources() {
    let (url, _server) = start_server().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClient::connect(transport).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, "config://app");

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_read_resource() {
    let (url, _server) = start_server().await;

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

    let result = client.read_resource("config://app").await.unwrap();
    assert_eq!(result.contents.len(), 1);
    assert!(
        result.contents[0]
            .text
            .as_deref()
            .unwrap()
            .contains("debug")
    );

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_ping() {
    let (url, _server) = start_server().await;

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

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

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_shutdown_sends_delete() {
    let (url, _server) = start_server().await;

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

    // Shutdown should send DELETE and close cleanly
    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_with_custom_config() {
    let (url, _server) = start_server().await;

    let config = tower_mcp::client::HttpClientConfig {
        request_timeout: Duration::from_secs(10),
        sse_reconnect: false,
        ..Default::default()
    };
    let transport = HttpClientTransport::with_config(&url, config);
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    client.shutdown().await.unwrap();
}

/// Start an HTTP server with bearer token auth middleware on a random port.
async fn start_auth_server(valid_token: &str) -> (String, tokio::task::JoinHandle<()>) {
    use axum::{extract::Request, http::StatusCode, middleware, response::Response};

    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let mcp_router = transport.into_router();

    let token = valid_token.to_string();
    let app = mcp_router.layer(middleware::from_fn(
        move |request: Request, next: middleware::Next| {
            let token = token.clone();
            async move {
                let auth = request
                    .headers()
                    .get("Authorization")
                    .and_then(|v| v.to_str().ok())
                    .map(|s| s.to_string());
                let expected = format!("Bearer {}", token);
                if auth.as_deref() == Some(&expected) {
                    Ok::<Response, (StatusCode, String)>(next.run(request).await)
                } else {
                    Err((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))
                }
            }
        },
    ));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

#[tokio::test]
async fn test_http_client_bearer_auth() {
    let (url, _server) = start_auth_server("sk-valid-key").await;

    let transport = HttpClientTransport::new(&url).bearer_token("sk-valid-key");
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_bearer_auth_rejected() {
    let (url, _server) = start_auth_server("sk-valid-key").await;

    let transport = HttpClientTransport::new(&url).bearer_token("sk-wrong-key");
    let client = McpClient::connect(transport).await.unwrap();

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

// ---------------------------------------------------------------------------
// Extended test router with prompts, resource templates, and completions
// ---------------------------------------------------------------------------

/// Create a test router with prompts, resource templates, and completions.
fn extended_test_router() -> McpRouter {
    let echo = ToolBuilder::new("echo")
        .description("Echo a message")
        .handler(|input: EchoInput| async move { Ok(CallToolResult::text(input.message)) })
        .build();

    let add = ToolBuilder::new("add")
        .description("Add two numbers")
        .handler(|input: AddInput| async move {
            Ok(CallToolResult::text(format!("{}", input.a + input.b)))
        })
        .build();

    let config = ResourceBuilder::new("config://app")
        .name("App Config")
        .description("Application configuration")
        .text(r#"{"debug": true}"#);

    let greeting = PromptBuilder::new("greeting")
        .description("Generate a greeting")
        .required_arg("name", "Name to greet")
        .handler(|args: HashMap<String, String>| async move {
            let name = args.get("name").cloned().unwrap_or_default();
            Ok(GetPromptResult {
                description: Some("A greeting prompt".to_string()),
                messages: vec![PromptMessage {
                    role: PromptRole::User,
                    content: Content::text(format!("Hello, {}!", name)),
                    meta: None,
                }],
                meta: None,
            })
        })
        .build();

    let farewell = PromptBuilder::new("farewell")
        .description("Generate a farewell")
        .optional_arg("name", "Name to bid farewell")
        .handler(|args: HashMap<String, String>| async move {
            let name = args
                .get("name")
                .cloned()
                .unwrap_or_else(|| "friend".to_string());
            Ok(GetPromptResult {
                description: Some("A farewell prompt".to_string()),
                messages: vec![PromptMessage {
                    role: PromptRole::User,
                    content: Content::text(format!("Goodbye, {}!", name)),
                    meta: None,
                }],
                meta: None,
            })
        })
        .build();

    let file_template = ResourceTemplateBuilder::new("file:///{+path}")
        .name("Project Files")
        .description("Read project files by path")
        .mime_type("text/plain")
        .handler(|uri: String, vars: HashMap<String, String>| async move {
            let path = vars.get("path").cloned().unwrap_or_default();
            Ok(ReadResourceResult {
                contents: vec![ResourceContent {
                    uri,
                    mime_type: Some("text/plain".to_string()),
                    text: Some(format!("Contents of {}", path)),
                    blob: None,
                    meta: None,
                }],
                meta: None,
                ..Default::default()
            })
        });

    McpRouter::new()
        .server_info("test-extended-server", "1.0.0")
        .tool(echo)
        .tool(add)
        .resource(config)
        .prompt(greeting)
        .prompt(farewell)
        .resource_template(file_template)
        .completion_handler(|params: CompleteParams| async move {
            match &params.reference {
                CompletionReference::Prompt { name } if name == "greeting" => {
                    let prefix = &params.argument.value;
                    let names = vec!["Alice", "Bob", "Charlie"];
                    let matches: Vec<String> = names
                        .into_iter()
                        .filter(|n| n.to_lowercase().starts_with(&prefix.to_lowercase()))
                        .map(String::from)
                        .collect();
                    Ok(CompleteResult::new(matches))
                }
                CompletionReference::Resource { uri } if uri.contains("file") => {
                    let prefix = &params.argument.value;
                    let paths = vec!["src/main.rs", "src/lib.rs", "Cargo.toml"];
                    let matches: Vec<String> = paths
                        .into_iter()
                        .filter(|p| p.starts_with(prefix.as_str()))
                        .map(String::from)
                        .collect();
                    Ok(CompleteResult::new(matches))
                }
                _ => Ok(CompleteResult::new(vec![])),
            }
        })
}

/// Start an extended HTTP server with prompts, templates, and completions.
async fn start_extended_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = extended_test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

// ---------------------------------------------------------------------------
// Prompts E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_list_prompts() {
    let (url, _server) = start_extended_server().await;

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

    let prompts = client.list_prompts().await.unwrap();
    assert_eq!(prompts.prompts.len(), 2);

    let names: Vec<&str> = prompts.prompts.iter().map(|p| p.name.as_str()).collect();
    assert!(names.contains(&"greeting"));
    assert!(names.contains(&"farewell"));

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_get_prompt_with_args() {
    let (url, _server) = start_extended_server().await;

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

    let mut args = HashMap::new();
    args.insert("name".to_string(), "World".to_string());
    let result = client.get_prompt("greeting", Some(args)).await.unwrap();

    assert_eq!(result.description.as_deref(), Some("A greeting prompt"));
    assert_eq!(result.messages.len(), 1);
    assert!(matches!(result.messages[0].role, PromptRole::User));

    // Verify the message content contains our argument
    let text = match &result.messages[0].content {
        Content::Text { text, .. } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    assert_eq!(text, "Hello, World!");

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_get_prompt_optional_arg() {
    let (url, _server) = start_extended_server().await;

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

    // Call farewell without the optional name arg
    let result = client.get_prompt("farewell", None).await.unwrap();
    assert_eq!(result.messages.len(), 1);

    let text = match &result.messages[0].content {
        Content::Text { text, .. } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    assert_eq!(text, "Goodbye, friend!");

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Resource Templates E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_list_resource_templates() {
    let (url, _server) = start_extended_server().await;

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

    let templates = client.list_resource_templates().await.unwrap();
    assert_eq!(templates.resource_templates.len(), 1);
    assert_eq!(
        templates.resource_templates[0].uri_template,
        "file:///{+path}"
    );
    assert_eq!(templates.resource_templates[0].name, "Project Files");

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_read_resource_template() {
    let (url, _server) = start_extended_server().await;

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

    let result = client.read_resource("file:///src/main.rs").await.unwrap();
    assert_eq!(result.contents.len(), 1);
    assert_eq!(
        result.contents[0].text.as_deref().unwrap(),
        "Contents of src/main.rs"
    );
    assert_eq!(result.contents[0].mime_type.as_deref(), Some("text/plain"));

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// call_tool_text() E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_call_tool_text() {
    let (url, _server) = start_extended_server().await;

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

    let text = client
        .call_tool_text("echo", serde_json::json!({"message": "hello"}))
        .await
        .unwrap();
    assert_eq!(text, "hello");

    let text = client
        .call_tool_text("add", serde_json::json!({"a": 10, "b": 20}))
        .await
        .unwrap();
    assert_eq!(text, "30");

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Completion API E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_complete_prompt_arg() {
    let (url, _server) = start_extended_server().await;

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

    // Complete with "A" prefix -- should match "Alice"
    let result = client
        .complete_prompt_arg("greeting", "name", "A")
        .await
        .unwrap();
    assert_eq!(result.completion.values, vec!["Alice"]);

    // Complete with "b" prefix -- should match "Bob" (case-insensitive)
    let result = client
        .complete_prompt_arg("greeting", "name", "b")
        .await
        .unwrap();
    assert_eq!(result.completion.values, vec!["Bob"]);

    // Complete with "z" prefix -- no matches
    let result = client
        .complete_prompt_arg("greeting", "name", "z")
        .await
        .unwrap();
    assert!(result.completion.values.is_empty());

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_complete_resource_uri() {
    let (url, _server) = start_extended_server().await;

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

    // Complete with "src/" prefix
    let result = client
        .complete_resource_uri("file:///{+path}", "path", "src/")
        .await
        .unwrap();
    assert_eq!(result.completion.values.len(), 2);
    assert!(
        result
            .completion
            .values
            .contains(&"src/main.rs".to_string())
    );
    assert!(result.completion.values.contains(&"src/lib.rs".to_string()));

    // Complete with "C" prefix
    let result = client
        .complete_resource_uri("file:///{+path}", "path", "C")
        .await
        .unwrap();
    assert_eq!(result.completion.values, vec!["Cargo.toml"]);

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Pagination E2E
// ---------------------------------------------------------------------------

/// Create a test router with many items and small page size for pagination testing.
fn pagination_test_router() -> McpRouter {
    let mut router = McpRouter::new()
        .server_info("test-pagination-server", "1.0.0")
        .page_size(2);

    for i in 0..5 {
        let tool = ToolBuilder::new(format!("tool_{i}"))
            .description(format!("Tool {i}"))
            .handler(|_: NoParams| async move { Ok(CallToolResult::text("ok")) })
            .build();
        router = router.tool(tool);
    }

    for i in 0..3 {
        let resource = ResourceBuilder::new(format!("res://item_{i}"))
            .name(format!("Resource {i}"))
            .text(format!("content_{i}"));
        router = router.resource(resource);
    }

    for i in 0..3 {
        let prompt = PromptBuilder::new(format!("prompt_{i}"))
            .description(format!("Prompt {i}"))
            .handler(|_: HashMap<String, String>| async move {
                Ok(GetPromptResult {
                    description: None,
                    messages: vec![],
                    meta: None,
                })
            })
            .build();
        router = router.prompt(prompt);
    }

    router
}

async fn start_pagination_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = pagination_test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

#[tokio::test]
async fn test_http_client_pagination_list_all_tools() {
    let (url, _server) = start_pagination_server().await;

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

    let tools = client.list_all_tools().await.unwrap();
    assert_eq!(tools.len(), 5);

    let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
    for i in 0..5 {
        assert!(names.contains(&format!("tool_{i}").as_str()));
    }

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_pagination_manual_cursor() {
    let (url, _server) = start_pagination_server().await;

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

    // Page 1: 2 tools + cursor
    let page1 = client.list_tools_with_cursor(None).await.unwrap();
    assert_eq!(page1.tools.len(), 2);
    assert!(page1.next_cursor.is_some());

    // Page 2: 2 tools + cursor
    let page2 = client
        .list_tools_with_cursor(page1.next_cursor)
        .await
        .unwrap();
    assert_eq!(page2.tools.len(), 2);
    assert!(page2.next_cursor.is_some());

    // Page 3: 1 tool + no cursor
    let page3 = client
        .list_tools_with_cursor(page2.next_cursor)
        .await
        .unwrap();
    assert_eq!(page3.tools.len(), 1);
    assert!(page3.next_cursor.is_none());

    // Verify all 5 unique tools were returned
    let mut all_names: Vec<String> = Vec::new();
    all_names.extend(page1.tools.iter().map(|t| t.name.clone()));
    all_names.extend(page2.tools.iter().map(|t| t.name.clone()));
    all_names.extend(page3.tools.iter().map(|t| t.name.clone()));
    all_names.sort();
    assert_eq!(all_names.len(), 5);

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_pagination_list_all_prompts() {
    let (url, _server) = start_pagination_server().await;

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

    let prompts = client.list_all_prompts().await.unwrap();
    assert_eq!(prompts.len(), 3);

    let names: Vec<&str> = prompts.iter().map(|p| p.name.as_str()).collect();
    for i in 0..3 {
        assert!(names.contains(&format!("prompt_{i}").as_str()));
    }

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Bidirectional communication infrastructure
// ---------------------------------------------------------------------------

/// Create a test router with tools that exercise sampling, elicitation, and logging.
fn bidirectional_test_router() -> McpRouter {
    let sampling_tool = ToolBuilder::new("test_sampling")
        .description("Calls ctx.sample() and returns the LLM response")
        .extractor_handler((), |ctx: Context, _: RawArgs| async move {
            let params =
                CreateMessageParams::new(vec![SamplingMessage::user("Test sampling request")], 100);
            match ctx.sample(params).await {
                Ok(result) => {
                    let text = result.first_text().unwrap_or("no text").to_string();
                    Ok(CallToolResult::text(text))
                }
                Err(e) => Ok(CallToolResult::error(format!("sampling failed: {e}"))),
            }
        })
        .build();

    let confirm_tool = ToolBuilder::new("test_confirm")
        .description("Calls ctx.confirm() and returns the result")
        .extractor_handler((), |ctx: Context, _: RawArgs| async move {
            match ctx.confirm("proceed?").await {
                Ok(true) => Ok(CallToolResult::text("confirmed")),
                Ok(false) => Ok(CallToolResult::text("declined")),
                Err(e) => Ok(CallToolResult::error(format!("elicitation failed: {e}"))),
            }
        })
        .build();

    let log_tool = ToolBuilder::new("test_log")
        .description("Sends a log notification and returns")
        .extractor_handler((), |ctx: Context, _: RawArgs| async move {
            ctx.send_log(LoggingMessageParams::new(
                LogLevel::Info,
                serde_json::json!("test log message"),
            ));
            Ok(CallToolResult::text("logged"))
        })
        .build();

    let roots_tool = ToolBuilder::new("test_roots")
        .description("Requests the client's roots")
        .extractor_handler((), |ctx: Context, _: RawArgs| async move {
            match ctx.request_raw("roots/list", serde_json::json!({})).await {
                Ok(result) => {
                    let uri = result
                        .pointer("/roots/0/uri")
                        .and_then(serde_json::Value::as_str)
                        .unwrap_or("no root");
                    Ok(CallToolResult::text(uri))
                }
                Err(error) => Ok(CallToolResult::error(format!(
                    "roots request failed: {error}"
                ))),
            }
        })
        .build();

    McpRouter::new()
        .server_info("test-bidirectional-server", "1.0.0")
        .tool(sampling_tool)
        .tool(confirm_tool)
        .tool(log_tool)
        .tool(roots_tool)
}

/// Start a bidirectional HTTP server with sampling enabled.
async fn start_bidirectional_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = bidirectional_test_router();
    let transport = HttpTransport::new(router)
        .with_sampling()
        .disable_origin_validation();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

// ---------------------------------------------------------------------------
// Notification Handler via SSE E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_notification_log_message() {
    let (url, _server) = start_bidirectional_server().await;

    let captured: Arc<Mutex<Vec<LoggingMessageParams>>> = Arc::new(Mutex::new(Vec::new()));
    let captured_clone = captured.clone();

    let handler = NotificationHandler::new().on_log_message(move |params| {
        captured_clone.lock().unwrap().push(params);
    });

    let transport = HttpClientTransport::new(&url);
    let client = McpClientBuilder::new()
        .with_sampling()
        .connect(transport, handler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Call the tool that sends a log notification
    let result = client
        .call_tool("test_log", serde_json::json!({}))
        .await
        .unwrap();
    assert_eq!(result.first_text(), Some("logged"));

    // Wait for the notification to arrive via SSE
    let received = tokio::time::timeout(Duration::from_secs(3), async {
        loop {
            if !captured.lock().unwrap().is_empty() {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    })
    .await
    .unwrap_or(false);

    assert!(received, "Expected to receive log notification via SSE");

    {
        let logs = captured.lock().unwrap();
        assert_eq!(logs.len(), 1);
        assert!(matches!(logs[0].level, LogLevel::Info));
        assert_eq!(logs[0].data, serde_json::json!("test log message"));
    }

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Server-initiated Sampling E2E
// ---------------------------------------------------------------------------

struct MockSamplingHandler;

#[async_trait]
impl ClientHandler for MockSamplingHandler {
    async fn handle_create_message(
        &self,
        _params: CreateMessageParams,
    ) -> Result<CreateMessageResult, JsonRpcError> {
        Ok(CreateMessageResult {
            content: SamplingContentOrArray::Single(SamplingContent::Text {
                text: "mock-llm-response".into(),
                annotations: None,
                meta: None,
            }),
            model: "test-model".into(),
            role: ContentRole::Assistant,
            stop_reason: Some("end_turn".into()),
            meta: None,
        })
    }
}

#[tokio::test]
async fn test_http_client_sampling_round_trip() {
    let (url, _server) = start_bidirectional_server().await;

    let transport = HttpClientTransport::with_config(
        &url,
        HttpClientConfig {
            auto_sse: false,
            ..Default::default()
        },
    );
    let client = McpClientBuilder::new()
        .with_sampling()
        .connect(transport, MockSamplingHandler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Call the tool that triggers server-side sampling
    // Flow: client -> call_tool -> server ctx.sample() -> originating POST SSE ->
    //       client MockSamplingHandler -> POST response -> server returns result
    let result = client
        .call_tool("test_sampling", serde_json::json!({}))
        .await
        .unwrap();

    assert_eq!(result.first_text(), Some("mock-llm-response"));

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Server-initiated Elicitation E2E
// ---------------------------------------------------------------------------

struct MockElicitationHandler;

#[async_trait]
impl ClientHandler for MockElicitationHandler {
    async fn handle_elicit(
        &self,
        _params: ElicitRequestParams,
    ) -> Result<ElicitResult, JsonRpcError> {
        let mut content = HashMap::new();
        content.insert("confirm".to_string(), ElicitFieldValue::Boolean(true));
        Ok(ElicitResult::accept(content))
    }
}

#[tokio::test]
async fn test_http_client_elicitation_confirm() {
    let (url, _server) = start_bidirectional_server().await;

    let transport = HttpClientTransport::with_config(
        &url,
        HttpClientConfig {
            auto_sse: false,
            ..Default::default()
        },
    );
    let client = McpClientBuilder::new()
        .with_elicitation()
        .with_sampling()
        .connect(transport, MockElicitationHandler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Call the tool that triggers server-side elicitation via ctx.confirm()
    // Flow: client -> call_tool -> server ctx.confirm("proceed?") -> originating POST SSE ->
    //       client MockElicitationHandler -> POST response -> server returns "confirmed"
    let result = client
        .call_tool("test_confirm", serde_json::json!({}))
        .await
        .unwrap();

    assert_eq!(result.first_text(), Some("confirmed"));

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_roots_request_uses_originating_post() {
    let (url, _server) = start_bidirectional_server().await;

    let transport = HttpClientTransport::with_config(
        &url,
        HttpClientConfig {
            auto_sse: false,
            ..Default::default()
        },
    );
    let client = McpClientBuilder::new()
        .with_roots(vec![Root::new("file:///associated")])
        .connect_simple(transport)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    let result = client
        .call_tool("test_roots", serde_json::json!({}))
        .await
        .unwrap();
    assert_eq!(result.first_text(), Some("file:///associated"));

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn concurrent_associated_posts_use_unique_server_request_ids() {
    let (url, _server) = start_bidirectional_server().await;

    let transport = HttpClientTransport::with_config(
        &url,
        HttpClientConfig {
            auto_sse: false,
            ..Default::default()
        },
    );
    let client = McpClientBuilder::new()
        .with_sampling()
        .connect(transport, MockSamplingHandler)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    let (first, second) = tokio::time::timeout(Duration::from_secs(3), async {
        tokio::join!(
            client.call_tool("test_sampling", serde_json::json!({})),
            client.call_tool("test_sampling", serde_json::json!({})),
        )
    })
    .await
    .expect("concurrent associated requests must not cross-talk or hang");

    assert_eq!(first.unwrap().first_text(), Some("mock-llm-response"));
    assert_eq!(second.unwrap().first_text(), Some("mock-llm-response"));

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Root Management E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_root_management() {
    let (url, _server) = start_server().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClientBuilder::new()
        .with_roots(vec![Root::new("file:///project")])
        .connect_simple(transport)
        .await
        .unwrap();
    client.initialize("test-client", "1.0.0").await.unwrap();

    // Initial roots
    let roots = client.list_roots().await;
    assert_eq!(roots.roots.len(), 1);
    assert_eq!(roots.roots[0].uri, "file:///project");

    // Add a root
    client.add_root(Root::new("file:///other")).await.unwrap();
    let roots = client.list_roots().await;
    assert_eq!(roots.roots.len(), 2);

    // Remove the first root
    let removed = client.remove_root("file:///project").await.unwrap();
    assert!(removed);
    let roots = client.list_roots().await;
    assert_eq!(roots.roots.len(), 1);
    assert_eq!(roots.roots[0].uri, "file:///other");

    // Remove nonexistent root returns false
    let removed = client.remove_root("file:///nonexistent").await.unwrap();
    assert!(!removed);

    // Set roots replaces all
    client
        .set_roots(vec![Root::new("file:///new_a"), Root::new("file:///new_b")])
        .await
        .unwrap();
    let roots = client.list_roots().await;
    assert_eq!(roots.roots.len(), 2);
    let uris: Vec<&str> = roots.roots.iter().map(|r| r.uri.as_str()).collect();
    assert!(uris.contains(&"file:///new_a"));
    assert!(uris.contains(&"file:///new_b"));

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Token Provider E2E
// ---------------------------------------------------------------------------

#[cfg(feature = "oauth-client")]
mod token_provider_tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use tower_mcp::{
        OAuthClientError, OAuthScopeEscalationConfig, OAuthScopeEscalationHandler,
        OAuthScopeEscalationRequest, TokenProvider,
    };

    struct StaticTokenProvider(String);

    #[async_trait]
    impl TokenProvider for StaticTokenProvider {
        async fn get_token(&self) -> Result<String, OAuthClientError> {
            Ok(self.0.clone())
        }
    }

    struct EscalatingTokenProvider {
        token: Arc<tokio::sync::RwLock<String>>,
        requests: Arc<Mutex<Vec<OAuthScopeEscalationRequest>>>,
        reauthorization_delay: Duration,
    }

    #[async_trait]
    impl TokenProvider for EscalatingTokenProvider {
        async fn get_token(&self) -> Result<String, OAuthClientError> {
            Ok(self.token.read().await.clone())
        }
    }

    #[async_trait]
    impl OAuthScopeEscalationHandler for EscalatingTokenProvider {
        async fn reauthorize(
            &self,
            request: OAuthScopeEscalationRequest,
        ) -> Result<(), OAuthClientError> {
            self.requests.lock().unwrap().push(request);
            tokio::time::sleep(self.reauthorization_delay).await;
            *self.token.write().await = "escalated-token".to_string();
            Ok(())
        }
    }

    async fn start_scope_challenge_server(
        always_reject: bool,
    ) -> (String, Arc<AtomicUsize>, tokio::task::JoinHandle<()>) {
        use axum::{
            extract::Request,
            http::{HeaderValue, StatusCode, header::WWW_AUTHENTICATE},
            middleware,
            response::IntoResponse,
        };

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://127.0.0.1:{}", addr.port());
        let challenge = format!(
            r#"Bearer error="insufficient_scope", scope="tools.read", resource_metadata="{url}/.well-known/oauth-protected-resource", error_description="tools/list requires tools.read""#
        );
        let request_count = Arc::new(AtomicUsize::new(0));

        let transport = HttpTransport::new(test_router()).disable_origin_validation();
        let count = request_count.clone();
        let app = transport.into_router().layer(middleware::from_fn(
            move |request: Request, next: middleware::Next| {
                let count = count.clone();
                let challenge = challenge.clone();
                async move {
                    let method = request
                        .headers()
                        .get("mcp-method")
                        .and_then(|value| value.to_str().ok());
                    let is_tools_list = method == Some("tools/list");
                    if is_tools_list {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                    let escalated = request
                        .headers()
                        .get("authorization")
                        .and_then(|value| value.to_str().ok())
                        == Some("Bearer escalated-token");

                    if is_tools_list && (always_reject || !escalated) {
                        let mut response = StatusCode::FORBIDDEN.into_response();
                        response
                            .headers_mut()
                            .insert(WWW_AUTHENTICATE, HeaderValue::from_str(&challenge).unwrap());
                        response
                    } else {
                        next.run(request).await
                    }
                }
            },
        ));

        let handle = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        (url, request_count, handle)
    }

    #[tokio::test]
    async fn test_http_client_with_token_provider() {
        let (url, _server) = start_auth_server("dynamic-test-token").await;

        let transport = HttpClientTransport::new(&url)
            .with_token_provider(StaticTokenProvider("dynamic-test-token".into()));
        let client = McpClient::connect(transport).await.unwrap();

        let info = client.initialize("test-client", "1.0.0").await.unwrap();
        assert_eq!(info.server_info.name, "test-http-server");

        // Verify we can make further requests with the dynamic token
        let tools = client.list_tools().await.unwrap();
        assert_eq!(tools.tools.len(), 2);

        client.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_http_client_token_provider_rejected() {
        let (url, _server) = start_auth_server("dynamic-test-token").await;

        let transport = HttpClientTransport::new(&url)
            .with_token_provider(StaticTokenProvider("wrong-token".into()));
        let client = McpClient::connect(transport).await.unwrap();

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

    #[tokio::test]
    async fn scope_challenge_reauthorizes_and_retries_operation() {
        let (url, request_count, _server) = start_scope_challenge_server(false).await;
        let token = Arc::new(tokio::sync::RwLock::new("initial-token".to_string()));
        let requests = Arc::new(Mutex::new(Vec::new()));
        let provider = EscalatingTokenProvider {
            token,
            requests: requests.clone(),
            reauthorization_delay: Duration::ZERO,
        };
        let transport = HttpClientTransport::new(&url)
            .with_scope_aware_token_provider(provider, OAuthScopeEscalationConfig::new(["openid"]));
        let client = McpClient::connect(transport).await.unwrap();

        client.initialize("test-client", "1.0.0").await.unwrap();
        let tools = client.list_tools().await.unwrap();
        assert_eq!(tools.tools.len(), 2);
        assert_eq!(request_count.load(Ordering::SeqCst), 2);

        {
            let recorded = requests.lock().unwrap();
            assert_eq!(recorded.len(), 1);
            assert_eq!(recorded[0].resource, url);
            assert_eq!(recorded[0].operation, "tools/list");
            assert_eq!(recorded[0].previous_scopes, ["openid"]);
            assert_eq!(recorded[0].requested_scopes, ["openid", "tools.read"]);
            assert_eq!(recorded[0].challenge.required_scopes, ["tools.read"]);
            assert!(recorded[0].challenge.resource_metadata.is_some());
            assert_eq!(recorded[0].attempt, 1);
        }

        client.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn repeated_scope_challenge_stops_at_configured_retry_limit() {
        let (url, request_count, _server) = start_scope_challenge_server(true).await;
        let requests = Arc::new(Mutex::new(Vec::new()));
        let provider = EscalatingTokenProvider {
            token: Arc::new(tokio::sync::RwLock::new("initial-token".to_string())),
            requests: requests.clone(),
            reauthorization_delay: Duration::ZERO,
        };
        let transport = HttpClientTransport::new(&url).with_scope_aware_token_provider(
            provider,
            OAuthScopeEscalationConfig::new(["openid"]).max_attempts(3),
        );
        let client = McpClient::connect(transport).await.unwrap();

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

        assert!(error.contains("insufficient_scope"));
        assert!(error.contains("tools.read"));
        assert_eq!(request_count.load(Ordering::SeqCst), 4);
        assert_eq!(requests.lock().unwrap().len(), 3);
    }

    #[tokio::test]
    async fn concurrent_scope_challenges_share_one_reauthorization() {
        let (url, request_count, _server) = start_scope_challenge_server(false).await;
        let requests = Arc::new(Mutex::new(Vec::new()));
        let provider = EscalatingTokenProvider {
            token: Arc::new(tokio::sync::RwLock::new("initial-token".to_string())),
            requests: requests.clone(),
            reauthorization_delay: Duration::from_millis(50),
        };
        let transport = HttpClientTransport::new(&url)
            .with_scope_aware_token_provider(provider, OAuthScopeEscalationConfig::new(["openid"]));
        let client = McpClient::connect(transport).await.unwrap();
        client.initialize("test-client", "1.0.0").await.unwrap();

        let (first, second) = tokio::join!(client.list_tools(), client.list_tools());

        assert_eq!(first.unwrap().tools.len(), 2);
        assert_eq!(second.unwrap().tools.len(), 2);
        assert_eq!(requests.lock().unwrap().len(), 1);
        assert!((3..=4).contains(&request_count.load(Ordering::SeqCst)));
        client.shutdown().await.unwrap();
    }
}

// ---------------------------------------------------------------------------
// Auth: API key header E2E
// ---------------------------------------------------------------------------

/// Start an HTTP server that validates a custom header.
async fn start_custom_header_auth_server(
    header_name: &str,
    expected_value: &str,
) -> (String, tokio::task::JoinHandle<()>) {
    use axum::{extract::Request, http::StatusCode, middleware, response::Response};

    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let mcp_router = transport.into_router();

    let name = header_name.to_string();
    let value = expected_value.to_string();
    let app = mcp_router.layer(middleware::from_fn(
        move |request: Request, next: middleware::Next| {
            let name = name.clone();
            let value = value.clone();
            async move {
                let actual = request
                    .headers()
                    .get(&name)
                    .and_then(|v| v.to_str().ok())
                    .map(|s| s.to_string());
                if actual.as_deref() == Some(&value) {
                    Ok::<Response, (StatusCode, String)>(next.run(request).await)
                } else {
                    Err((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))
                }
            }
        },
    ));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

#[tokio::test]
async fn test_http_client_api_key_header() {
    let (url, _server) = start_custom_header_auth_server("X-API-Key", "my-secret-key").await;

    let transport = HttpClientTransport::new(&url).api_key_header("X-API-Key", "my-secret-key");
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    let tools = client.list_tools().await.unwrap();
    assert_eq!(tools.tools.len(), 2);

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_api_key_header_rejected() {
    let (url, _server) = start_custom_header_auth_server("X-API-Key", "my-secret-key").await;

    let transport = HttpClientTransport::new(&url).api_key_header("X-API-Key", "wrong-key");
    let client = McpClient::connect(transport).await.unwrap();

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

// ---------------------------------------------------------------------------
// Auth: Basic auth E2E
// ---------------------------------------------------------------------------

/// Start an HTTP server that validates Basic auth credentials.
async fn start_basic_auth_server(
    expected_user: &str,
    expected_pass: &str,
) -> (String, tokio::task::JoinHandle<()>) {
    use axum::{extract::Request, http::StatusCode, middleware, response::Response};
    use base64::Engine;

    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let mcp_router = transport.into_router();

    let expected = format!(
        "Basic {}",
        base64::engine::general_purpose::STANDARD
            .encode(format!("{}:{}", expected_user, expected_pass))
    );
    let app = mcp_router.layer(middleware::from_fn(
        move |request: Request, next: middleware::Next| {
            let expected = expected.clone();
            async move {
                let auth = request
                    .headers()
                    .get("Authorization")
                    .and_then(|v| v.to_str().ok())
                    .map(|s| s.to_string());
                if auth.as_deref() == Some(&expected) {
                    Ok::<Response, (StatusCode, String)>(next.run(request).await)
                } else {
                    Err((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))
                }
            }
        },
    ));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;

    (url, handle)
}

#[tokio::test]
async fn test_http_client_basic_auth() {
    let (url, _server) = start_basic_auth_server("admin", "password123").await;

    let transport = HttpClientTransport::new(&url).basic_auth("admin", "password123");
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_basic_auth_rejected() {
    let (url, _server) = start_basic_auth_server("admin", "password123").await;

    let transport = HttpClientTransport::new(&url).basic_auth("admin", "wrong");
    let client = McpClient::connect(transport).await.unwrap();

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

// ---------------------------------------------------------------------------
// Auth: Custom header E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_custom_header() {
    let (url, _server) = start_custom_header_auth_server("X-Custom-Auth", "custom-value-123").await;

    let transport = HttpClientTransport::new(&url).header("X-Custom-Auth", "custom-value-123");
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Auth: OAuthClientCredentials E2E
// ---------------------------------------------------------------------------

#[cfg(feature = "oauth-client")]
mod oauth_client_credentials_tests {
    use super::*;
    use tower_mcp::OAuthClientCredentials;

    /// Start a mock OAuth token endpoint that returns access tokens.
    async fn start_mock_token_server(
        expected_client_id: &str,
        expected_client_secret: &str,
        issued_token: &str,
    ) -> (String, tokio::task::JoinHandle<()>) {
        use axum::{Router, extract::Request, http::StatusCode, routing::post};
        use base64::Engine;

        let expected_auth = format!(
            "Basic {}",
            base64::engine::general_purpose::STANDARD
                .encode(format!("{}:{}", expected_client_id, expected_client_secret))
        );
        let token = issued_token.to_string();

        let app = Router::new().route(
            "/token",
            post(move |request: Request| {
                let expected_auth = expected_auth.clone();
                let token = token.clone();
                async move {
                    let auth = request
                        .headers()
                        .get("Authorization")
                        .and_then(|v| v.to_str().ok())
                        .map(|s| s.to_string());
                    if auth.as_deref() != Some(&expected_auth) {
                        return (
                            StatusCode::UNAUTHORIZED,
                            axum::Json(serde_json::json!({"error": "invalid_client"})),
                        );
                    }
                    (
                        StatusCode::OK,
                        axum::Json(serde_json::json!({
                            "access_token": token,
                            "token_type": "bearer",
                            "expires_in": 3600
                        })),
                    )
                }
            }),
        );

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://127.0.0.1:{}", addr.port());

        let handle = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        tokio::time::sleep(Duration::from_millis(50)).await;

        (url, handle)
    }

    #[tokio::test]
    async fn test_http_client_oauth_client_credentials() {
        let issued_token = "oauth-access-token-xyz";

        // Start the mock token server
        let (token_url, _token_server) =
            start_mock_token_server("my-client", "my-secret", issued_token).await;

        // Start the MCP server that validates bearer tokens
        let (mcp_url, _mcp_server) = start_auth_server(issued_token).await;

        // Build the OAuth provider pointing at our mock token endpoint
        let provider = OAuthClientCredentials::builder()
            .client_id("my-client")
            .client_secret("my-secret")
            .token_endpoint(format!("{}/token", token_url))
            .resource(&mcp_url)
            .build()
            .unwrap();

        let transport = HttpClientTransport::new(&mcp_url).with_token_provider(provider);
        let client = McpClient::connect(transport).await.unwrap();

        let info = client.initialize("test-client", "1.0.0").await.unwrap();
        assert_eq!(info.server_info.name, "test-http-server");

        // Verify subsequent requests also work (token is cached and reused)
        let tools = client.list_tools().await.unwrap();
        assert_eq!(tools.tools.len(), 2);

        client.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_http_client_oauth_client_credentials_bad_secret() {
        let (token_url, _token_server) =
            start_mock_token_server("my-client", "my-secret", "token").await;

        let (mcp_url, _mcp_server) = start_auth_server("token").await;

        // Wrong secret should fail at token acquisition
        let provider = OAuthClientCredentials::builder()
            .client_id("my-client")
            .client_secret("wrong-secret")
            .token_endpoint(format!("{}/token", token_url))
            .resource(&mcp_url)
            .build()
            .unwrap();

        let transport = HttpClientTransport::new(&mcp_url).with_token_provider(provider);
        let client = McpClient::connect(transport).await.unwrap();

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

// ---------------------------------------------------------------------------
// Config: auto_sse disabled E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_auto_sse_disabled() {
    let (url, _server) = start_server().await;

    let config = tower_mcp::client::HttpClientConfig {
        auto_sse: false,
        ..Default::default()
    };
    let transport = HttpClientTransport::with_config(&url, config);
    let client = McpClient::connect(transport).await.unwrap();

    let info = client.initialize("test-client", "1.0.0").await.unwrap();
    assert_eq!(info.server_info.name, "test-http-server");

    // Basic request/response still works without SSE
    let tools = client.list_tools().await.unwrap();
    assert_eq!(tools.tools.len(), 2);

    let result = client
        .call_tool("echo", serde_json::json!({"message": "no sse"}))
        .await
        .unwrap();
    assert_eq!(result.first_text(), Some("no sse"));

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Edge Cases: Concurrent requests E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_concurrent_requests() {
    let (url, _server) = start_server().await;

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

    // Fire off multiple call_tool requests concurrently
    let mut handles = Vec::new();
    for i in 0..5 {
        let client_ref = &client;
        handles.push(async move {
            client_ref
                .call_tool("add", serde_json::json!({"a": i, "b": i * 10}))
                .await
        });
    }

    let results: Vec<_> = futures::future::join_all(handles).await;
    for result in &results {
        assert!(result.is_ok(), "Concurrent request failed: {:?}", result);
    }

    // Verify all results are correct (order may vary)
    let mut texts: Vec<String> = results
        .into_iter()
        .map(|r| r.unwrap().first_text().unwrap().to_string())
        .collect();
    texts.sort();
    assert_eq!(texts, vec!["0", "11", "22", "33", "44"]);

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Edge Cases: Error responses E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_tool_not_found_error() {
    let (url, _server) = start_server().await;

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

    let result = client
        .call_tool("nonexistent_tool", serde_json::json!({}))
        .await;
    assert!(result.is_err());

    let err = result.unwrap_err();
    let err_msg = err.to_string();
    assert!(
        err_msg.contains("not found") || err_msg.contains("Unknown tool"),
        "Expected tool-not-found error, got: {}",
        err_msg
    );

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_invalid_tool_arguments() {
    let (url, _server) = start_server().await;

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

    // "add" expects {a: i64, b: i64}, send wrong types
    let result = client
        .call_tool("add", serde_json::json!({"a": "not_a_number", "b": 1}))
        .await;

    // The server should return an error (either JSON-RPC error or isError in CallToolResult)
    match result {
        Ok(call_result) => {
            assert!(
                call_result.is_error,
                "Expected error flag in CallToolResult"
            );
        }
        Err(_) => {
            // JSON-RPC error is also acceptable
        }
    }

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_prompt_not_found_error() {
    let (url, _server) = start_extended_server().await;

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

    let result = client.get_prompt("nonexistent_prompt", None).await;
    assert!(result.is_err());

    client.shutdown().await.unwrap();
}

#[tokio::test]
async fn test_http_client_resource_not_found_error() {
    let (url, _server) = start_server().await;

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

    let result = client.read_resource("nonexistent://resource").await;
    assert!(result.is_err());

    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Edge Cases: Shutdown without initialize E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_shutdown_without_initialize() {
    let (url, _server) = start_server().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClient::connect(transport).await.unwrap();

    // Shutdown without ever calling initialize -- should not panic
    client.shutdown().await.unwrap();
}

// ---------------------------------------------------------------------------
// Edge Cases: Requests before initialize E2E
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_request_before_initialize() {
    let (url, _server) = start_server().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClient::connect(transport).await.unwrap();

    // Attempting to list tools before initialize should fail
    let result = client.list_tools().await;
    assert!(result.is_err());

    let err_msg = result.unwrap_err().to_string();
    assert!(
        err_msg.contains("not initialized"),
        "Expected not-initialized error, got: {}",
        err_msg
    );
}

// ---------------------------------------------------------------------------
// SSE Event Replay on Reconnect E2E (#509)
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_http_client_sse_event_replay_on_reconnect() {
    let (url, _server) = start_bidirectional_server().await;
    let client = reqwest::Client::new();

    // 1. Initialize a session
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": { "sampling": {} },
                    "clientInfo": { "name": "test-sse", "version": "1.0.0" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let session_id = resp
        .headers()
        .get("mcp-session-id")
        .unwrap()
        .to_str()
        .unwrap()
        .to_string();

    // 1b. Send notifications/initialized before any other requests (spec requirement)
    client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized"
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    // 2. Open SSE stream and start consuming in a background task.
    //    This ensures the server's broadcast subscriber is active before we
    //    trigger notifications.
    let sse_resp = client
        .get(&url)
        .header("Accept", "text/event-stream")
        .header("mcp-session-id", &session_id)
        .send()
        .await
        .unwrap();
    assert_eq!(sse_resp.status(), 200);

    let (collected_tx, collected_rx) = tokio::sync::oneshot::channel();
    let sse_task = tokio::spawn(async move {
        let mut text = String::new();
        let mut stream = sse_resp.bytes_stream();
        use futures::StreamExt;
        while let Some(chunk) = stream.next().await {
            let chunk = chunk.unwrap();
            text.push_str(&String::from_utf8_lossy(&chunk));
            // Wait until we have at least 2 events with IDs
            let id_count = text.lines().filter(|l| l.starts_with("id:")).count();
            if id_count >= 2 {
                break;
            }
        }
        let _ = collected_tx.send(text);
    });

    // Give the SSE subscription time to be established
    tokio::time::sleep(Duration::from_millis(100)).await;

    // 3. Send two tool calls to trigger two log notifications
    for req_id in [2, 3] {
        let _resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .header("mcp-session-id", &session_id)
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "method": "tools/call",
                    "params": { "name": "test_log", "arguments": {} }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();
    }

    // 4. Wait for background task to collect both events
    let sse_text = tokio::time::timeout(Duration::from_secs(5), collected_rx)
        .await
        .expect("timed out waiting for SSE events")
        .expect("SSE task dropped sender");

    // Parse event IDs from the collected SSE data
    let event_ids: Vec<u64> = sse_text
        .lines()
        .filter(|l| l.starts_with("id:"))
        .filter_map(|l| l.trim_start_matches("id:").trim().parse().ok())
        .collect();

    assert!(
        event_ids.len() >= 2,
        "Expected at least 2 events, got: {event_ids:?}"
    );

    let first_id = event_ids[0];

    // 5. SSE task has finished; stream is dropped (connection closed)
    let _ = sse_task.await;

    // 6. Open new SSE stream with Last-Event-ID = first event ID.
    //    The server should replay buffered events after that ID.
    let sse_resp2 = client
        .get(&url)
        .header("Accept", "text/event-stream")
        .header("mcp-session-id", &session_id)
        .header("Last-Event-ID", first_id.to_string())
        .send()
        .await
        .unwrap();
    assert_eq!(sse_resp2.status(), 200);

    // 7. Read replayed events
    let replay_text = tokio::time::timeout(Duration::from_secs(2), async {
        let mut text = String::new();
        let mut stream = sse_resp2.bytes_stream();
        use futures::StreamExt;
        while let Some(chunk) = stream.next().await {
            let chunk = chunk.unwrap();
            text.push_str(&String::from_utf8_lossy(&chunk));
            if text.lines().any(|l| l.starts_with("id:")) {
                break;
            }
        }
        text
    })
    .await
    .unwrap_or_default();

    // Parse replayed event IDs -- they should all be > first_id
    let replayed_ids: Vec<u64> = replay_text
        .lines()
        .filter(|l| l.starts_with("id:"))
        .filter_map(|l| l.trim_start_matches("id:").trim().parse().ok())
        .collect();

    assert!(
        !replayed_ids.is_empty(),
        "Expected replayed events after reconnection with Last-Event-ID"
    );
    for id in &replayed_ids {
        assert!(
            *id > first_id,
            "Replayed event ID {id} should be > last seen ID {first_id}",
        );
    }
}

// ---------------------------------------------------------------------------
// Session Expiry E2E (#510)
// ---------------------------------------------------------------------------

/// Start a server with very short session TTL for expiry testing.
async fn start_expiring_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = test_router();
    let config = SessionConfig::with_ttl(Duration::from_millis(100))
        .cleanup_interval(Duration::from_millis(20));
    let transport = HttpTransport::new(router)
        .disable_origin_validation()
        .session_config(config);
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;
    (url, handle)
}

#[tokio::test]
async fn test_http_client_session_expiry_error() {
    let (url, _server) = start_expiring_server().await;
    let client = reqwest::Client::new();

    // Initialize a session
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": { "name": "test-expiry", "version": "1.0.0" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let session_id = resp
        .headers()
        .get("mcp-session-id")
        .unwrap()
        .to_str()
        .unwrap()
        .to_string();

    // Send notifications/initialized before any other requests (spec requirement)
    client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized"
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    // Verify the session works before expiry
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert!(
        body.get("result").is_some(),
        "Expected successful tools/list response, got: {body}"
    );

    // Wait for session to expire and cleanup to run
    tokio::time::sleep(Duration::from_millis(200)).await;

    // Session should be expired -- next request should return a JSON-RPC error
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/list",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    let body: serde_json::Value = resp.json().await.unwrap();
    let error = body
        .get("error")
        .expect("Expected JSON-RPC error after session expiry");
    assert_eq!(
        error["code"].as_i64().unwrap(),
        -32005,
        "Expected SessionNotFound error code (-32005)"
    );
}

// ---------------------------------------------------------------------------
// Session handling -- optional by default, strict via require_sessions()
// ---------------------------------------------------------------------------

/// Start a server with require_sessions() for strict session enforcement.
async fn start_required_sessions_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = test_router();
    let transport = HttpTransport::new(router)
        .disable_origin_validation()
        .require_sessions();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;
    (url, handle)
}

/// Start a server with optional sessions (the default).
async fn start_optional_sessions_server() -> (String, tokio::task::JoinHandle<()>) {
    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let axum_router = transport.into_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let handle = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;
    (url, handle)
}

/// Simulate Codex CLI behavior: initialize + tools/list during setup,
/// then tools/call without session ID on a fresh connection.
#[tokio::test]
async fn test_optional_sessions_codex_pattern() {
    let (url, _server) = start_optional_sessions_server().await;
    let client = reqwest::Client::new();

    // Phase 1: Setup -- Codex initializes and discovers tools (normal flow)
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": { "name": "codex-cli", "version": "0.1.0" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let session_id = resp
        .headers()
        .get("mcp-session-id")
        .unwrap()
        .to_str()
        .unwrap()
        .to_string();

    // Send notifications/initialized before any other requests (spec requirement)
    client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized"
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    // tools/list with session ID (setup phase)
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    let body: serde_json::Value = resp.json().await.unwrap();
    let tools = body["result"]["tools"].as_array().unwrap();
    assert_eq!(tools.len(), 2, "Should discover echo and add tools");

    // Phase 2: Execution -- Codex sends tools/call WITHOUT session ID
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        // No mcp-session-id header!
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "echo",
                    "arguments": { "message": "hello from codex" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert!(
        body.get("error").is_none(),
        "tools/call without session ID should succeed with optional_sessions, got: {body}"
    );
    let content = &body["result"]["content"][0]["text"];
    assert_eq!(content.as_str().unwrap(), "hello from codex");
}

/// Simulate Cursor behavior: tools/list without session ID.
#[tokio::test]
async fn test_optional_sessions_cursor_tools_list_no_session() {
    let (url, _server) = start_optional_sessions_server().await;
    let client = reqwest::Client::new();

    // Cursor sends tools/list directly without initializing or session ID
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        // No mcp-session-id header!
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/list",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert!(
        body.get("error").is_none(),
        "tools/list without session ID should succeed with optional_sessions, got: {body}"
    );
    let tools = body["result"]["tools"].as_array().unwrap();
    assert_eq!(tools.len(), 2);
}

/// Multiple sessionless requests should each work independently.
#[tokio::test]
async fn test_optional_sessions_multiple_stateless_calls() {
    let (url, _server) = start_optional_sessions_server().await;
    let client = reqwest::Client::new();

    // Send multiple tools/call requests without session IDs
    for i in 0..3 {
        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": i + 1,
                    "method": "tools/call",
                    "params": {
                        "name": "add",
                        "arguments": { "a": i, "b": 10 }
                    }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let body: serde_json::Value = resp.json().await.unwrap();
        assert!(
            body.get("error").is_none(),
            "Stateless call {i} failed: {body}"
        );
        let expected = format!("{}", i + 10);
        assert_eq!(
            body["result"]["content"][0]["text"].as_str().unwrap(),
            expected
        );
    }
}

/// With require_sessions(), missing session ID should be rejected.
#[tokio::test]
async fn test_sessions_required_rejects_without_session_id() {
    let (url, _server) = start_required_sessions_server().await;
    let client = reqwest::Client::new();

    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        // No mcp-session-id header!
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/list",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    let body: serde_json::Value = resp.json().await.unwrap();
    let error = body
        .get("error")
        .expect("Expected JSON-RPC error when session ID is missing");
    assert_eq!(
        error["code"].as_i64().unwrap(),
        -32006,
        "Expected SessionRequired error code (-32006)"
    );
}

/// With optional_sessions, clients that DO send session IDs should still work.
#[tokio::test]
async fn test_optional_sessions_with_session_id_still_works() {
    let (url, _server) = start_optional_sessions_server().await;
    let client = reqwest::Client::new();

    // Initialize normally
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": { "name": "good-client", "version": "1.0.0" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    let session_id = resp
        .headers()
        .get("mcp-session-id")
        .unwrap()
        .to_str()
        .unwrap()
        .to_string();

    // Send initialized notification
    let _resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized",
                "params": {}
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    // Use session normally
    let resp = client
        .post(&url)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json")
        .header("mcp-session-id", &session_id)
        .body(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/call",
                "params": {
                    "name": "echo",
                    "arguments": { "message": "session works" }
                }
            })
            .to_string(),
        )
        .send()
        .await
        .unwrap();

    assert_eq!(resp.status(), 200);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert!(
        body.get("error").is_none(),
        "Session-based call should work: {body}"
    );
    assert_eq!(
        body["result"]["content"][0]["text"].as_str().unwrap(),
        "session works"
    );
}

// =============================================================================
// SEP-1442 Stateless mode tests
// =============================================================================

#[cfg(feature = "stateless")]
mod stateless_tests {
    use super::*;
    use axum::extract::State;
    use axum::http::{HeaderMap, StatusCode};
    use axum::response::sse::{Event, Sse};
    use axum::response::{IntoResponse, Response};
    use axum::routing::post;
    use axum::{Json, Router};
    use std::convert::Infallible;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use tower_mcp::stateless::StatelessConfig;

    /// Start a server with stateless mode enabled.
    async fn start_stateless_server() -> (String, tokio::task::JoinHandle<()>) {
        let router = test_router();
        let transport = HttpTransport::new(router)
            .disable_origin_validation()
            .stateless(StatelessConfig::new());
        let axum_router = transport.into_router();

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://127.0.0.1:{}", addr.port());

        let handle = tokio::spawn(async move {
            axum::serve(listener, axum_router).await.unwrap();
        });

        tokio::time::sleep(Duration::from_millis(50)).await;
        (url, handle)
    }

    #[derive(Clone, Default)]
    struct StaleSchemaState {
        list_calls: Arc<AtomicUsize>,
        tool_calls: Arc<AtomicUsize>,
        mirrored_headers: Arc<Mutex<Vec<(bool, bool)>>>,
    }

    async fn stale_schema_endpoint(
        State(state): State<StaleSchemaState>,
        headers: HeaderMap,
        Json(body): Json<serde_json::Value>,
    ) -> Response {
        let id = body.get("id").cloned().unwrap_or(serde_json::Value::Null);
        match body
            .get("method")
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default()
        {
            "server/discover" => (
                StatusCode::OK,
                Json(serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "result": {
                        "resultType": "complete",
                        "supportedVersions": ["2026-07-28"],
                        "capabilities": {},
                        "ttlMs": 0,
                        "cacheScope": "private"
                    }
                })),
            )
                .into_response(),
            "tools/list" => {
                let suffix = if state.list_calls.fetch_add(1, Ordering::SeqCst) == 0 {
                    "Region"
                } else {
                    "Zone"
                };
                (
                    StatusCode::OK,
                    Json(serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": id,
                        "result": {
                            "resultType": "complete",
                            "tools": [{
                                "name": "changing-tool",
                                "inputSchema": {
                                    "type": "object",
                                    "properties": {
                                        "region": {
                                            "type": "string",
                                            "x-mcp-header": suffix
                                        }
                                    }
                                }
                            }],
                            "ttlMs": 60_000,
                            "cacheScope": "private"
                        }
                    })),
                )
                    .into_response()
            }
            "tools/call" => {
                state.mirrored_headers.lock().unwrap().push((
                    headers.contains_key("mcp-param-region"),
                    headers.contains_key("mcp-param-zone"),
                ));
                if state.tool_calls.fetch_add(1, Ordering::SeqCst) == 0 {
                    (
                        StatusCode::BAD_REQUEST,
                        Json(serde_json::json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "error": {
                                "code": -32020,
                                "message": "cached tool schema is stale"
                            }
                        })),
                    )
                        .into_response()
                } else {
                    (
                        StatusCode::OK,
                        Json(serde_json::json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "result": {
                                "resultType": "complete",
                                "content": [{
                                    "type": "text",
                                    "text": "fresh headers"
                                }]
                            }
                        })),
                    )
                        .into_response()
                }
            }
            _ => (
                StatusCode::NOT_FOUND,
                Json(serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "error": {"code": -32601, "message": "unknown method"}
                })),
            )
                .into_response(),
        }
    }

    #[derive(Clone, Default)]
    struct SubscriptionClientState {
        active_streams: Arc<AtomicUsize>,
        dropped_streams: Arc<AtomicUsize>,
        tools_list_calls: Arc<AtomicUsize>,
        cancellation_posts: Arc<AtomicUsize>,
    }

    enum SubscriptionStreamMode {
        LongLived,
        Graceful,
        Abrupt,
        BeforeAcknowledgment,
        MismatchedId,
    }

    struct SubscriptionStream {
        phase: u8,
        id: serde_json::Value,
        notifications: serde_json::Value,
        mode: SubscriptionStreamMode,
        state: SubscriptionClientState,
    }

    impl Drop for SubscriptionStream {
        fn drop(&mut self) {
            self.state.active_streams.fetch_sub(1, Ordering::SeqCst);
            self.state.dropped_streams.fetch_add(1, Ordering::SeqCst);
        }
    }

    async fn subscription_client_endpoint(
        State(state): State<SubscriptionClientState>,
        headers: HeaderMap,
        Json(body): Json<serde_json::Value>,
    ) -> Response {
        let id = body.get("id").cloned().unwrap_or(serde_json::Value::Null);
        let method = body
            .get("method")
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default();
        match method {
            "server/discover" => (
                StatusCode::OK,
                Json(serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "result": {
                        "resultType": "complete",
                        "supportedVersions": ["2026-07-28"],
                        "capabilities": {},
                        "ttlMs": 0,
                        "cacheScope": "private"
                    }
                })),
            )
                .into_response(),
            "tools/list" => {
                state.tools_list_calls.fetch_add(1, Ordering::SeqCst);
                (
                    StatusCode::OK,
                    Json(serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": id,
                        "result": {
                            "resultType": "complete",
                            "tools": [],
                            "ttlMs": 0,
                            "cacheScope": "private"
                        }
                    })),
                )
                    .into_response()
            }
            "subscriptions/listen" => {
                if headers
                    .get("mcp-protocol-version")
                    .and_then(|value| value.to_str().ok())
                    != Some("2026-07-28")
                    || headers
                        .get("mcp-method")
                        .and_then(|value| value.to_str().ok())
                        != Some("subscriptions/listen")
                    || body.pointer("/params/_meta/io.modelcontextprotocol~1protocolVersion")
                        != Some(&serde_json::json!("2026-07-28"))
                {
                    return StatusCode::BAD_REQUEST.into_response();
                }

                let notifications = body["params"]["notifications"].clone();
                let mode = if notifications["resourcesListChanged"] == serde_json::Value::Bool(true)
                {
                    SubscriptionStreamMode::Graceful
                } else if notifications["promptsListChanged"] == serde_json::Value::Bool(true) {
                    SubscriptionStreamMode::Abrupt
                } else if notifications["resourceSubscriptions"][0]
                    == serde_json::json!("before-ack")
                {
                    SubscriptionStreamMode::BeforeAcknowledgment
                } else if notifications["resourceSubscriptions"][0] == serde_json::json!("wrong-id")
                {
                    SubscriptionStreamMode::MismatchedId
                } else {
                    SubscriptionStreamMode::LongLived
                };
                state.active_streams.fetch_add(1, Ordering::SeqCst);
                let stream = futures::stream::unfold(
                    SubscriptionStream {
                        phase: 0,
                        id,
                        notifications,
                        mode,
                        state,
                    },
                    |mut stream| async move {
                        let message = match stream.phase {
                            0 => match stream.mode {
                                SubscriptionStreamMode::BeforeAcknowledgment => {
                                    serde_json::json!({
                                        "jsonrpc": "2.0",
                                        "method": "notifications/tools/list_changed",
                                        "params": {
                                            "_meta": {
                                                "io.modelcontextprotocol/subscriptionId": stream.id
                                            }
                                        }
                                    })
                                }
                                SubscriptionStreamMode::MismatchedId => serde_json::json!({
                                    "jsonrpc": "2.0",
                                    "method": "notifications/subscriptions/acknowledged",
                                    "params": {
                                        "_meta": {
                                            "io.modelcontextprotocol/subscriptionId": 999
                                        },
                                        "notifications": stream.notifications
                                    }
                                }),
                                _ => serde_json::json!({
                                    "jsonrpc": "2.0",
                                    "method": "notifications/subscriptions/acknowledged",
                                    "params": {
                                        "_meta": {
                                            "io.modelcontextprotocol/subscriptionId": stream.id
                                        },
                                        "notifications": stream.notifications
                                    }
                                }),
                            },
                            1 => match stream.mode {
                                SubscriptionStreamMode::LongLived => serde_json::json!({
                                    "jsonrpc": "2.0",
                                    "method": "notifications/tools/list_changed",
                                    "params": {
                                        "_meta": {
                                            "io.modelcontextprotocol/subscriptionId": stream.id
                                        }
                                    }
                                }),
                                SubscriptionStreamMode::Graceful => serde_json::json!({
                                    "jsonrpc": "2.0",
                                    "id": stream.id,
                                    "result": {
                                        "resultType": "complete",
                                        "_meta": {
                                            "io.modelcontextprotocol/subscriptionId": stream.id
                                        }
                                    }
                                }),
                                SubscriptionStreamMode::Abrupt => return None,
                                SubscriptionStreamMode::BeforeAcknowledgment
                                | SubscriptionStreamMode::MismatchedId => {
                                    std::future::pending::<()>().await;
                                    unreachable!()
                                }
                            },
                            _ => match stream.mode {
                                SubscriptionStreamMode::LongLived => {
                                    std::future::pending::<()>().await;
                                    unreachable!()
                                }
                                SubscriptionStreamMode::Graceful
                                | SubscriptionStreamMode::Abrupt => return None,
                                SubscriptionStreamMode::BeforeAcknowledgment
                                | SubscriptionStreamMode::MismatchedId => unreachable!(),
                            },
                        };
                        stream.phase += 1;
                        Some((
                            Ok::<_, Infallible>(Event::default().data(message.to_string())),
                            stream,
                        ))
                    },
                );
                Sse::new(stream).into_response()
            }
            "notifications/cancelled" => {
                state.cancellation_posts.fetch_add(1, Ordering::SeqCst);
                StatusCode::ACCEPTED.into_response()
            }
            _ => StatusCode::NOT_FOUND.into_response(),
        }
    }

    async fn wait_for_counter(counter: &AtomicUsize, expected: usize) {
        tokio::time::timeout(Duration::from_secs(2), async {
            while counter.load(Ordering::SeqCst) != expected {
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn final_http_client_manages_concurrent_subscription_streams() {
        let state = SubscriptionClientState::default();
        let app = Router::new()
            .route("/", post(subscription_client_endpoint))
            .with_state(state.clone());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let server = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        let tools_changed = Arc::new(AtomicUsize::new(0));
        let handler = NotificationHandler::new().on_tools_changed({
            let tools_changed = tools_changed.clone();
            move || {
                tools_changed.fetch_add(1, Ordering::SeqCst);
            }
        });
        let config = HttpClientConfig {
            request_timeout: Duration::from_millis(100),
            sse_reconnect: false,
            ..Default::default()
        };
        let client = McpClient::builder()
            .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
            .connect(HttpClientTransport::with_config(url, config), handler)
            .await
            .unwrap();
        client.discover("test-client", "1.0.0").await.unwrap();

        let filter = SubscriptionFilter {
            tools_list_changed: Some(true),
            ..Default::default()
        };
        let mut first = client.listen_subscriptions(filter.clone()).await.unwrap();
        let mut second = client.listen_subscriptions(filter).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)
        );
        wait_for_counter(&state.active_streams, 2).await;
        wait_for_counter(&tools_changed, 2).await;

        tokio::time::sleep(Duration::from_millis(200)).await;
        assert_eq!(state.dropped_streams.load(Ordering::SeqCst), 0);
        client.list_tools().await.unwrap();

        first.cancel().await.unwrap();
        wait_for_counter(&state.active_streams, 1).await;
        client.list_tools().await.unwrap();
        second.cancel().await.unwrap();
        wait_for_counter(&state.active_streams, 0).await;
        assert_eq!(state.dropped_streams.load(Ordering::SeqCst), 2);
        assert_eq!(state.tools_list_calls.load(Ordering::SeqCst), 2);
        assert_eq!(state.cancellation_posts.load(Ordering::SeqCst), 0);

        let mut graceful = client
            .listen_subscriptions(SubscriptionFilter {
                resources_list_changed: Some(true),
                ..Default::default()
            })
            .await
            .unwrap();
        assert_eq!(
            graceful
                .acknowledged()
                .await
                .unwrap()
                .resources_list_changed,
            Some(true)
        );
        let graceful_id = graceful.id().clone();
        let result = graceful.wait().await.unwrap();
        assert_eq!(result.meta.subscription_id, graceful_id);

        let mut abrupt = client
            .listen_subscriptions(SubscriptionFilter {
                prompts_list_changed: Some(true),
                ..Default::default()
            })
            .await
            .unwrap();
        assert_eq!(
            abrupt.acknowledged().await.unwrap().prompts_list_changed,
            Some(true)
        );
        let error = tokio::time::timeout(Duration::from_secs(2), abrupt.wait())
            .await
            .expect("abrupt close should not leave the subscription pending")
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("closed the response stream before the final reply")
        );

        for (resource, expected) in [
            (
                "before-ack",
                "subscription notification arrived before acknowledgment",
            ),
            (
                "wrong-id",
                "subscription notification carried a missing or mismatched subscription ID",
            ),
        ] {
            let invalid = client
                .listen_subscriptions(SubscriptionFilter {
                    resource_subscriptions: Some(vec![resource.to_string()]),
                    ..Default::default()
                })
                .await
                .unwrap();
            let error = tokio::time::timeout(Duration::from_secs(2), invalid.wait())
                .await
                .expect("protocol violation should end the subscription")
                .unwrap_err();
            assert!(error.to_string().contains(expected), "{error}");
        }

        client.shutdown().await.unwrap();
        server.abort();
    }

    #[tokio::test]
    async fn final_client_refreshes_headers_after_stale_schema_rejection() {
        let state = StaleSchemaState::default();
        let app = Router::new()
            .route("/", post(stale_schema_endpoint))
            .with_state(state.clone());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let server = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        let client = McpClient::builder()
            .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
            .connect_simple(HttpClientTransport::new(url))
            .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!({"region": "us-west1"}))
            .await
            .unwrap();
        assert_eq!(result.first_text(), Some("fresh headers"));
        assert_eq!(state.list_calls.load(Ordering::SeqCst), 2);
        assert_eq!(state.tool_calls.load(Ordering::SeqCst), 2);
        assert_eq!(
            *state.mirrored_headers.lock().unwrap(),
            [(true, false), (false, true)]
        );

        server.abort();
    }

    /// A 404 on a 2026-07-28 request carries no session meaning: the server
    /// answers an unknown method with 404 and -32601, and that must surface
    /// as method-not-found rather than start a re-initialize (#1556).
    #[tokio::test]
    async fn final_unknown_method_404_is_method_not_found_not_session_expiry() {
        let app = Router::new()
            .route("/", post(stale_schema_endpoint))
            .with_state(StaleSchemaState::default());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let server = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        let client = McpClient::builder()
            .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
            .connect_simple(HttpClientTransport::new(url))
            .await
            .unwrap();
        client.discover("test-client", "1.0.0").await.unwrap();

        let error = tokio::time::timeout(Duration::from_secs(5), client.list_resources())
            .await
            .expect("list_resources hung")
            .expect_err("the unknown method must fail");
        match &error {
            tower_mcp::Error::JsonRpc(rpc) => assert_eq!(rpc.code, -32601, "{error:?}"),
            other => panic!("expected method-not-found, got: {other:?}"),
        }

        tokio::time::timeout(Duration::from_secs(5), client.list_tools())
            .await
            .expect("list_tools after the 404 hung")
            .expect("the transport must stay open after an in-band 404");

        server.abort();
    }

    /// With no session to expire, a bare 404 on a 2026-07-28 request still
    /// means the endpoint is gone and closes the transport (#1553, #1556).
    #[tokio::test]
    async fn final_bare_404_still_closes_the_transport() {
        let url = spawn_raw_server(|req| {
            if raw_is(req, "server/discover") {
                let body = format!(
                    r#"{{"jsonrpc":"2.0","id":{},"result":{{"resultType":"complete","supportedVersions":["2026-07-28"],"capabilities":{{}},"ttlMs":0,"cacheScope":"private"}}}}"#,
                    raw_request_id(req)
                );
                raw_response("200 OK", RAW_JSON, &body)
            } else {
                raw_response("404 Not Found", "", "")
            }
        })
        .await;

        let client = McpClient::builder()
            .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
            .connect_simple(HttpClientTransport::with_config(url, raw_client_config()))
            .await
            .unwrap();
        client.discover("test-client", "1.0.0").await.unwrap();

        let error = tokio::time::timeout(Duration::from_secs(5), client.list_tools())
            .await
            .expect("list_tools hung")
            .expect_err("the 404 must fail the request");
        assert!(
            !matches!(error, tower_mcp::Error::SessionExpired),
            "a sessionless 404 is not session expiry: {error:?}"
        );

        tokio::time::timeout(Duration::from_secs(3), async {
            // The closed flag reaches the client after the failed request's
            // error, so poll until a call reports the closed connection.
            while !client
                .list_tools()
                .await
                .is_err_and(|error| error.to_string().contains("Connection closed"))
            {
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("a sessionless 404 must keep closing the transport");
    }

    /// tools/list works without initialize or session, just protocol version header.
    #[tokio::test]
    async fn test_stateless_tools_list() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-protocol-version", "2025-11-25")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        // Should echo back the protocol version
        assert_eq!(
            resp.headers()
                .get("mcp-protocol-version")
                .unwrap()
                .to_str()
                .unwrap(),
            "2025-11-25"
        );
        // Should NOT have a session ID
        assert!(resp.headers().get("mcp-session-id").is_none());

        let body: serde_json::Value = resp.json().await.unwrap();
        let tools = body["result"]["tools"].as_array().unwrap();
        assert!(tools.len() >= 2, "Should list tools without session");
    }

    /// tools/call works in stateless mode.
    #[tokio::test]
    async fn test_stateless_tools_call() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-protocol-version", "2025-11-25")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/call",
                    "params": {
                        "name": "echo",
                        "arguments": { "message": "stateless!" }
                    }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let body: serde_json::Value = resp.json().await.unwrap();
        assert_eq!(
            body["result"]["content"][0]["text"].as_str().unwrap(),
            "stateless!"
        );
    }

    /// Unsupported protocol version returns SEP-2575 error code -32022
    /// with the spec-shape data: `{ supported: [...], requested: "..." }`.
    /// (SEP-1442 draft used -32000 / `supportedVersions`; both were wrong.)
    #[tokio::test]
    async fn test_stateless_unsupported_version() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-protocol-version", "1999-01-01")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let body: serde_json::Value = resp.json().await.unwrap();
        assert_eq!(body["error"]["code"].as_i64().unwrap(), -32022);
        assert!(
            body["error"]["data"]["supported"].is_array(),
            "spec data field is 'supported': {body}"
        );
        assert_eq!(body["error"]["data"]["requested"], "1999-01-01");
    }

    /// Protocol version in _meta field (body) works as alternative to header.
    #[tokio::test]
    async fn test_stateless_version_in_meta() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            // No mcp-protocol-version header!
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list",
                    "params": {
                        "_meta": {
                            "modelcontextprotocol.io/mcpProtocolVersion": "2025-11-25"
                        }
                    }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let body: serde_json::Value = resp.json().await.unwrap();
        assert!(body["result"]["tools"].is_array());
    }

    /// Stateful clients (with session ID) still work alongside stateless.
    #[tokio::test]
    async fn test_stateless_mixed_mode() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        // Stateful: initialize and get session
        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json, text/event-stream")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {},
                        "clientInfo": { "name": "test", "version": "1.0" }
                    }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let session_id = resp
            .headers()
            .get("mcp-session-id")
            .unwrap()
            .to_str()
            .unwrap()
            .to_string();

        // Send initialized notification
        client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("mcp-session-id", &session_id)
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        // Stateful request with session ID
        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-session-id", &session_id)
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list",
                    "params": {}
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        let body: serde_json::Value = resp.json().await.unwrap();
        assert!(body["result"]["tools"].is_array());

        // Stateless request on same server (no session ID, protocol version header)
        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-protocol-version", "2025-11-25")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 3,
                    "method": "tools/call",
                    "params": {
                        "name": "add",
                        "arguments": { "a": 10, "b": 20 }
                    }
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        let body: serde_json::Value = resp.json().await.unwrap();
        assert_eq!(body["result"]["content"][0]["text"].as_str().unwrap(), "30");
    }

    /// server/discover works without any session or protocol version.
    #[tokio::test]
    async fn test_stateless_discover() {
        let (url, _server) = start_stateless_server().await;
        let client = reqwest::Client::new();

        // server/discover should work even without protocol version header
        // (it's allowed before initialization per session.rs)
        let resp = client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("Accept", "application/json")
            .header("mcp-protocol-version", "2025-11-25")
            .body(
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "server/discover",
                    "params": {}
                })
                .to_string(),
            )
            .send()
            .await
            .unwrap();

        assert_eq!(resp.status(), 200);
        let body: serde_json::Value = resp.json().await.unwrap();
        let result = &body["result"];
        assert!(result["supportedVersions"].is_array());
        assert!(result["_meta"]["io.modelcontextprotocol/serverInfo"]["name"].is_string());
        assert!(result["capabilities"].is_object());
    }
}

// ---------------------------------------------------------------------------
// Session recovery (#762) -- automatic re-initialization on session expiry
// ---------------------------------------------------------------------------

/// Start a server that returns a handle for session management.
async fn start_server_with_handle() -> (
    String,
    tower_mcp::transport::http::SessionHandle,
    tokio::task::JoinHandle<()>,
) {
    let router = test_router();
    let transport = HttpTransport::new(router).disable_origin_validation();
    let (axum_router, handle) = transport.into_router_with_handle();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let url = format!("http://127.0.0.1:{}", addr.port());

    let server = tokio::spawn(async move {
        axum::serve(listener, axum_router).await.unwrap();
    });

    tokio::time::sleep(Duration::from_millis(50)).await;
    (url, handle, server)
}

/// Verify that the client automatically recovers from session expiry.
#[tokio::test]
async fn test_session_recovery_after_expiry() {
    let (url, handle, _server) = start_server_with_handle().await;

    let transport = HttpClientTransport::new(&url);
    let client = McpClient::connect(transport).await.unwrap();
    let init_result = client.initialize("test-recovery", "1.0.0").await.unwrap();
    assert_eq!(init_result.server_info.name, "test-http-server");

    // Verify tools work
    let tools = client.list_tools().await.unwrap();
    assert_eq!(tools.tools.len(), 2);

    // Terminate the session to simulate expiry
    let sessions = handle.list_sessions().await;
    assert_eq!(sessions.len(), 1);
    let terminated = handle.terminate_session(&sessions[0].id).await;
    assert!(terminated, "Session should have been terminated");

    // The next request should fail with session expired, but auto-recover
    let tools = client.list_tools().await.unwrap();
    assert_eq!(tools.tools.len(), 2, "Should recover and list tools again");

    // Verify a new session was created
    let sessions = handle.list_sessions().await;
    assert_eq!(sessions.len(), 1, "Should have exactly one new session");
}

/// Verify that session recovery is disabled when configured.
#[tokio::test]
async fn test_session_recovery_disabled() {
    let (url, handle, _server) = start_server_with_handle().await;

    let transport = HttpClientTransport::new(&url).disable_session_recovery();
    let client = McpClient::connect(transport).await.unwrap();
    client
        .initialize("test-no-recovery", "1.0.0")
        .await
        .unwrap();

    // Terminate the session
    let sessions = handle.list_sessions().await;
    handle.terminate_session(&sessions[0].id).await;

    // Without recovery, the next request should fail
    let result = client.list_tools().await;
    assert!(result.is_err(), "Should fail without session recovery");
}

/// #967 regression: `notifications/initialized` must not race the first
/// request after `initialize()`.
///
/// Post-session request POSTs are spawned in background tasks; notification
/// POSTs must be awaited inline so a strict server (`strict_initialization`
/// defaults to true) processes them before the next request arrives. The
/// race depends on network timing, so an in-process server never exhibits
/// it; this test routes the client through a TCP proxy that delays any
/// connection carrying `notifications/initialized` by 200 ms, forcing the
/// reordering. Before the fix the subsequent `tools/list` deterministically
/// failed with -32600; after it, `initialize()` does not return until the
/// notification is acknowledged, so ordering holds.
#[tokio::test]
async fn initialized_notification_ordering_under_strict_server() {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let (url, _server) = start_server().await;
    let upstream = url.strip_prefix("http://").unwrap().to_string();

    // Delaying proxy: hold back the connection that carries the
    // initialized notification; pass everything else straight through.
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let proxy_addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        loop {
            let Ok((mut client_conn, _)) = listener.accept().await else {
                break;
            };
            let upstream = upstream.clone();
            tokio::spawn(async move {
                let mut first = vec![0u8; 16 * 1024];
                let n = match client_conn.read(&mut first).await {
                    Ok(0) | Err(_) => return,
                    Ok(n) => n,
                };
                if String::from_utf8_lossy(&first[..n]).contains("notifications/initialized") {
                    tokio::time::sleep(Duration::from_millis(200)).await;
                }
                let Ok(mut up) = tokio::net::TcpStream::connect(&upstream).await else {
                    return;
                };
                if up.write_all(&first[..n]).await.is_err() {
                    return;
                }
                let _ = tokio::io::copy_bidirectional(&mut client_conn, &mut up).await;
            });
        }
    });

    let proxied_url = format!("http://{proxy_addr}");
    let transport = HttpClientTransport::new(proxied_url);
    let client = McpClient::connect(transport).await.unwrap();
    client.initialize("race-test", "1.0.0").await.unwrap();
    let tools = client
        .list_tools()
        .await
        .expect("tools/list must not be rejected: initialized must be acknowledged first");
    assert!(!tools.tools.is_empty());
}

/// #983 regression: a 404 on the very first request (no session yet) is a
/// wrong endpoint path, not an expired session. It must surface as a
/// transport error naming the status and URL, never as SessionExpired.
#[tokio::test]
async fn pre_session_404_reports_endpoint_not_found() {
    let (url, _server) = start_server().await;
    // The MCP endpoint of the test server is the root path; /nope is a 404.
    let transport = HttpClientTransport::new(format!("{url}/nope"));
    let client = McpClient::connect(transport).await.unwrap();
    let err = client
        .initialize("wrong-path", "1.0.0")
        .await
        .expect_err("initialize against a wrong path must fail");
    let msg = err.to_string();
    assert!(
        msg.contains("404") && msg.contains("/nope") && msg.contains("endpoint"),
        "expected an endpoint-not-found transport error, got: {msg}"
    );
    assert!(
        !msg.to_lowercase().contains("session"),
        "must not be reported as a session error, got: {msg}"
    );
}

/// Read one full HTTP/1.1 request (headers + Content-Length body) from a raw
/// connection. Returns the whole request as a string, or `None` on EOF.
#[cfg(test)]
async fn read_http_request(stream: &mut tokio::net::TcpStream) -> Option<String> {
    use tokio::io::AsyncReadExt;
    let mut buf = Vec::new();
    let mut tmp = [0u8; 4096];
    loop {
        if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            let headers = String::from_utf8_lossy(&buf[..pos]).to_string();
            let content_length = headers
                .lines()
                .find_map(|l| {
                    let (k, v) = l.split_once(':')?;
                    k.trim()
                        .eq_ignore_ascii_case("content-length")
                        .then(|| v.trim().parse::<usize>().ok())
                        .flatten()
                })
                .unwrap_or(0);
            let need = pos + 4 + content_length;
            while buf.len() < need {
                let n = stream.read(&mut tmp).await.ok()?;
                if n == 0 {
                    break;
                }
                buf.extend_from_slice(&tmp[..n]);
            }
            return Some(String::from_utf8_lossy(&buf).to_string());
        }
        let n = stream.read(&mut tmp).await.ok()?;
        if n == 0 {
            return (!buf.is_empty()).then(|| String::from_utf8_lossy(&buf).to_string());
        }
        buf.extend_from_slice(&tmp[..n]);
    }
}

/// Regression: a post-session request whose background POST fails before
/// delivering a response must wake the awaiting caller with an error, never
/// hang it.
///
/// After the session is established, request POSTs are spawned in a background
/// task so the message loop can keep servicing the SSE stream. That task used
/// to drop a handful of failure modes on the floor (network error, timeout,
/// HTTP error with a non-JSON body, empty body, response stream closed with no
/// data), leaving the caller parked on the request id forever. This exercises
/// the empty-body dead-end: the raw server answers `tools/list` with a `200`
/// and no body. `Connection: close` forces a fresh connection per request so
/// each is intercepted deterministically.
#[tokio::test]
async fn post_session_request_failure_wakes_caller() {
    use tokio::io::AsyncWriteExt;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                let response = if req.contains("\"method\":\"initialize\"") {
                    let body = serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "result": {
                            "protocolVersion": "2025-11-25",
                            "capabilities": {},
                            "serverInfo": {"name": "raw", "version": "0"}
                        }
                    })
                    .to_string();
                    format!(
                        "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
                         mcp-session-id: raw-session\r\nmcp-protocol-version: 2025-11-25\r\n\
                         Content-Length: {}\r\nConnection: close\r\n\r\n{}",
                        body.len(),
                        body
                    )
                } else if req.contains("notifications/initialized") {
                    "HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
                        .to_string()
                } else {
                    // tools/list and anything else: a 2xx with an empty body,
                    // the dead-end that used to hang the caller.
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
                     Content-Length: 0\r\nConnection: close\r\n\r\n"
                        .to_string()
                };
                let _ = stream.write_all(response.as_bytes()).await;
                let _ = stream.flush().await;
            });
        }
    });

    let url = format!("http://{addr}");
    let config = HttpClientConfig {
        auto_sse: false,
        request_timeout: Duration::from_secs(5),
        ..Default::default()
    };
    let transport = HttpClientTransport::with_config(url, config);
    let client = McpClient::connect(transport).await.unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    // The bug made this future never resolve; bound it so the test fails as a
    // timeout instead of hanging the whole suite.
    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools()).await;
    let result = result.expect(
        "list_tools hung: the failed background POST never woke the caller (the bug this guards)",
    );
    assert!(
        result.is_err(),
        "an empty-body response must surface as an error, got: {result:?}"
    );
}

/// Regression: a server that stalls the `notifications/initialized` POST must
/// not freeze the client for the full `request_timeout`.
///
/// The notification is awaited inline to keep it ordered ahead of the first
/// request (#967), which blocks the single message loop for the duration.
/// `notification_timeout` bounds that block independently (#999). Here the raw
/// server reads the notification POST and never answers it; with a short
/// `notification_timeout` the client abandons it quickly.
///
/// Since #1174 the abandoned notification is also an `initialize()` error
/// rather than a silent success: the handshake did not complete, and a strict
/// server would reject every later request with -32600. The caller retries
/// initialization; it no longer proceeds into a poisoned session. The bounded
/// behavior this test originally guarded is asserted the same way, through
/// the outer timeout.
#[tokio::test]
async fn stalled_initialized_notification_does_not_freeze_client() {
    use tokio::io::AsyncWriteExt;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                if req.contains("\"method\":\"initialize\"") {
                    let body = serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "result": {
                            "protocolVersion": "2025-11-25",
                            "capabilities": {},
                            "serverInfo": {"name": "raw", "version": "0"}
                        }
                    })
                    .to_string();
                    let resp = format!(
                        "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
                         mcp-session-id: raw-session\r\nmcp-protocol-version: 2025-11-25\r\n\
                         Content-Length: {}\r\nConnection: close\r\n\r\n{}",
                        body.len(),
                        body
                    );
                    let _ = stream.write_all(resp.as_bytes()).await;
                } else if req.contains("notifications/initialized") {
                    // Stall: read the POST but never answer it. The client must
                    // give up after notification_timeout, not request_timeout.
                    tokio::time::sleep(Duration::from_secs(3600)).await;
                } else {
                    // tools/list: respond normally.
                    let body = serde_json::json!({
                        "jsonrpc": "2.0",
                        "id": 2,
                        "result": {"tools": [{"name": "echo", "inputSchema": {"type": "object"}}]}
                    })
                    .to_string();
                    let resp = format!(
                        "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
                         Content-Length: {}\r\nConnection: close\r\n\r\n{}",
                        body.len(),
                        body
                    );
                    let _ = stream.write_all(resp.as_bytes()).await;
                }
                let _ = stream.flush().await;
            });
        }
    });

    let url = format!("http://{addr}");
    let config = HttpClientConfig {
        auto_sse: false,
        request_timeout: Duration::from_secs(20),
        notification_timeout: Duration::from_millis(300),
        ..Default::default()
    };
    let transport = HttpClientTransport::with_config(url, config);
    let client = McpClient::connect(transport).await.unwrap();

    // The stalled notification is abandoned after ~300ms, so this resolves
    // quickly. Before #999 it blocked for the full 20s request_timeout; bound
    // it at 8s so a freeze regression fails as a timeout rather than hanging.
    let error = tokio::time::timeout(
        Duration::from_secs(8),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("initialize blocked behind the stalled notification (the bug this guards)")
    .expect_err("a stalled notifications/initialized must fail initialize, not be swallowed");

    assert!(
        error
            .to_string()
            .contains("failed to deliver notifications/initialized"),
        "error should name the handshake step, got: {error}"
    );
}

// ---------------------------------------------------------------------------
// Responses that carry no reply to the request (#1465)
// ---------------------------------------------------------------------------

const RAW_JSON: &str = "Content-Type: application/json\r\n";

/// A raw HTTP/1.1 server that answers every request with `respond(request)`,
/// one connection per request, so malformed replies reach the client exactly
/// as written.
async fn spawn_raw_server<F>(respond: F) -> String
where
    F: Fn(&str) -> String + Send + Sync + 'static,
{
    use tokio::io::AsyncWriteExt;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let respond = Arc::new(respond);
    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            let respond = respond.clone();
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                let _ = stream.write_all(respond(&req).as_bytes()).await;
                let _ = stream.flush().await;
            });
        }
    });
    format!("http://{addr}")
}

fn raw_response(status: &str, headers: &str, body: &str) -> String {
    format!(
        "HTTP/1.1 {status}\r\n{headers}Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    )
}

/// A response with no declared length: valid only paired with
/// `Connection: close`, which tells the client to read the body until the
/// server closes the socket rather than trusting a `Content-Length`.
fn raw_response_no_length(status: &str, headers: &str, body: &str) -> String {
    format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n{body}")
}

/// The JSON-RPC id of a raw HTTP request, as serialized JSON.
fn raw_request_id(req: &str) -> String {
    req.split_once("\r\n\r\n")
        .and_then(|(_, body)| serde_json::from_str::<serde_json::Value>(body).ok())
        .and_then(|value| value.get("id").cloned())
        .unwrap_or(serde_json::Value::Null)
        .to_string()
}

fn raw_is(req: &str, method: &str) -> bool {
    req.contains(&format!("\"method\":\"{method}\""))
}

/// The session-establishing reply to `initialize`.
fn raw_initialize_ok(req: &str) -> String {
    let body = format!(
        r#"{{"jsonrpc":"2.0","id":{},"result":{{"protocolVersion":"2025-11-25","capabilities":{{}},"serverInfo":{{"name":"raw","version":"0"}}}}}}"#,
        raw_request_id(req)
    );
    raw_response(
        "200 OK",
        &format!("{RAW_JSON}mcp-session-id: raw-session\r\nmcp-protocol-version: 2025-11-25\r\n"),
        &body,
    )
}

fn raw_client_config() -> HttpClientConfig {
    HttpClientConfig {
        auto_sse: false,
        request_timeout: Duration::from_secs(5),
        ..Default::default()
    }
}

/// On an established session, each response below used to leave `list_tools`
/// parked forever. It must fail instead, and the transport must stay usable:
/// a bad reply to one request is not a broken connection.
#[tokio::test]
async fn session_request_without_a_reply_fails_and_client_stays_usable() {
    // (case, status, headers, body); `{id}` in the body becomes the request id.
    let cases = [
        ("202 Accepted", "202 Accepted", "", "", None),
        ("empty JSON object", "200 OK", RAW_JSON, "{}", None),
        ("non-JSON body", "200 OK", RAW_JSON, "hello", None),
        (
            "error with a foreign id",
            "400 Bad Request",
            RAW_JSON,
            r#"{"jsonrpc":"2.0","id":"server-error","error":{"code":-32600,"message":"rejected by gateway"}}"#,
            Some("rejected by gateway"),
        ),
        (
            "error with the request id",
            "400 Bad Request",
            RAW_JSON,
            r#"{"jsonrpc":"2.0","id":{id},"error":{"code":-32600,"message":"rejected by server"}}"#,
            Some("rejected by server"),
        ),
    ];

    for (case, status, headers, body, expected_message) in cases {
        let url = spawn_raw_server(move |req| {
            if raw_is(req, "initialize") {
                raw_initialize_ok(req)
            } else if raw_is(req, "notifications/initialized") {
                raw_response("202 Accepted", "", "")
            } else if raw_is(req, "tools/list") {
                raw_response(status, headers, &body.replace("{id}", &raw_request_id(req)))
            } else {
                let reply = format!(
                    r#"{{"jsonrpc":"2.0","id":{},"result":{{}}}}"#,
                    raw_request_id(req)
                );
                raw_response("200 OK", RAW_JSON, &reply)
            }
        })
        .await;

        let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
            .await
            .unwrap();
        client.initialize("raw-client", "1.0.0").await.unwrap();

        let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
            .await
            .unwrap_or_else(|_| panic!("{case}: list_tools hung"));
        let error = result.expect_err(case).to_string();
        if let Some(expected) = expected_message {
            assert!(
                error.contains(expected),
                "{case}: the server's error should surface, got: {error}"
            );
        }

        tokio::time::timeout(Duration::from_secs(3), client.ping())
            .await
            .unwrap_or_else(|_| panic!("{case}: ping after the bad reply hung"))
            .unwrap_or_else(|error| panic!("{case}: client unusable after the bad reply: {error}"));
    }
}

/// Before a session exists, requests go through the synchronous path, which
/// had the same dead ends for `initialize`.
#[tokio::test]
async fn pre_session_request_without_a_reply_fails() {
    let cases = [
        ("202 Accepted", "202 Accepted", "", ""),
        ("empty JSON object", "200 OK", RAW_JSON, "{}"),
        ("empty body", "200 OK", RAW_JSON, ""),
    ];

    for (case, status, headers, body) in cases {
        let url = spawn_raw_server(move |_| raw_response(status, headers, body)).await;
        let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
            .await
            .unwrap();

        let result = tokio::time::timeout(
            Duration::from_secs(3),
            client.initialize("raw-client", "1.0.0"),
        )
        .await
        .unwrap_or_else(|_| panic!("{case}: initialize hung"));
        assert!(result.is_err(), "{case}: expected an error, got {result:?}");
    }
}

/// Gateways and the Python SDK reject `server/discover` with a JSON-RPC error
/// whose id is a literal string rather than the request's id.
#[cfg(feature = "stateless")]
#[tokio::test]
async fn discover_rejected_with_a_foreign_error_id_fails() {
    let url = spawn_raw_server(|_| {
        raw_response(
            "400 Bad Request",
            RAW_JSON,
            r#"{"jsonrpc":"2.0","id":"server-error","error":{"code":-32600,"message":"discover not supported"}}"#,
        )
    })
    .await;
    let client = McpClient::builder()
        .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect(
            HttpClientTransport::with_config(url, raw_client_config()),
            NotificationHandler::new(),
        )
        .await
        .unwrap();

    let result = tokio::time::timeout(
        Duration::from_secs(3),
        client.discover("raw-client", "1.0.0"),
    )
    .await
    .expect("discover hung on an error with a foreign id");
    let error = result.expect_err("discover should fail").to_string();
    assert!(
        error.contains("discover not supported"),
        "the server's error should surface, got: {error}"
    );
}

// ---------------------------------------------------------------------------
// Replies to server requests are not client requests (#1466)
// ---------------------------------------------------------------------------

/// Regression: the client's reply to a server-initiated request used to be
/// tracked exactly like one of the client's own requests, keyed by id. Server
/// and client request ids are independent sequences that both start at 1, so
/// a server request pushed while a client request is in flight can collide
/// with it. Before the fix, a failed reply POST injected an error frame
/// carrying that shared id, which incorrectly failed the client's own
/// still-pending request.
///
/// Here the raw server pushes a `ping` on the `tools/list` SSE stream using
/// the same id as the in-flight `tools/list` request, then rejects the
/// client's reply POST with a `500`. The real `tools/list` result is
/// withheld on the same connection until the rejected reply POST has been
/// fully handled, so a spurious error frame reaching the caller first would
/// prove the bug; the fix leaves the caller waiting for the genuine
/// completion instead.
#[tokio::test]
async fn server_request_colliding_with_client_id_reply_failure_does_not_fail_client_request() {
    colliding_reply_rejection_case(|_| raw_response("500 Internal Server Error", "", "")).await;
}

/// As above, but the reply POST is rejected with a JSON-RPC error body that
/// echoes the server request's id, which equals the client request's id. The
/// body must not be forwarded as if it answered the client's request.
#[tokio::test]
async fn reply_rejected_with_an_error_body_does_not_fail_the_colliding_client_request() {
    colliding_reply_rejection_case(|req| {
        let body = format!(
            r#"{{"jsonrpc":"2.0","id":{},"error":{{"code":-32600,"message":"reply rejected"}}}}"#,
            raw_request_id(req)
        );
        raw_response("400 Bad Request", RAW_JSON, &body)
    })
    .await;
}

async fn colliding_reply_rejection_case(reject_reply: fn(&str) -> String) {
    use tokio::io::AsyncWriteExt;

    fn sse_chunk(data: &str) -> String {
        format!("{:x}\r\n{}\r\n", data.len(), data)
    }

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let reply_handled = Arc::new(tokio::sync::Notify::new());

    {
        let reply_handled = reply_handled.clone();
        tokio::spawn(async move {
            loop {
                let Ok((mut stream, _)) = listener.accept().await else {
                    break;
                };
                let reply_handled = reply_handled.clone();
                tokio::spawn(async move {
                    let Some(req) = read_http_request(&mut stream).await else {
                        return;
                    };
                    if raw_is(&req, "initialize") {
                        let _ = stream.write_all(raw_initialize_ok(&req).as_bytes()).await;
                        let _ = stream.flush().await;
                    } else if raw_is(&req, "notifications/initialized") {
                        let _ = stream
                            .write_all(raw_response("202 Accepted", "", "").as_bytes())
                            .await;
                        let _ = stream.flush().await;
                    } else if raw_is(&req, "tools/list") {
                        let id = raw_request_id(&req);
                        let header = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                                       Transfer-Encoding: chunked\r\nConnection: close\r\n\r\n";
                        let _ = stream.write_all(header.as_bytes()).await;

                        // Push a server-initiated `ping` sharing the client's
                        // own request id.
                        let ping = format!(
                            "data: {{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"ping\",\"params\":{{}}}}\n\n"
                        );
                        let _ = stream.write_all(sse_chunk(&ping).as_bytes()).await;
                        let _ = stream.flush().await;

                        // Withhold the real result until the client's reply
                        // POST below has been received and rejected. If the
                        // bug regresses, its spurious error frame reaches the
                        // message loop first.
                        reply_handled.notified().await;

                        let result = format!(
                            "data: {{\"jsonrpc\":\"2.0\",\"id\":{id},\"result\":{{\"tools\":[]}}}}\n\n"
                        );
                        let _ = stream.write_all(sse_chunk(&result).as_bytes()).await;
                        let _ = stream.write_all(b"0\r\n\r\n").await;
                        let _ = stream.flush().await;
                    } else {
                        // The client's reply to the pushed `ping`: an id with
                        // no method. Reject it, then release the withheld
                        // `tools/list` completion.
                        let response = reject_reply(&req);
                        let _ = stream.write_all(response.as_bytes()).await;
                        let _ = stream.flush().await;
                        reply_handled.notify_one();
                    }
                });
            }
        });
    }

    let url = format!("http://{addr}");
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools hung");
    assert!(
        result.is_ok(),
        "the client's own request must complete normally despite the rejected reply POST, got: {result:?}"
    );
}

// ---------------------------------------------------------------------------
// Bounded response bodies (#1469)
// ---------------------------------------------------------------------------

/// A JSON-RPC response body for `id`, padded with an ignored `_pad` field so
/// its total length is exactly `limit` bytes when possible, or as close to
/// it as a non-negative pad allows. Padding a `tools/list` result keeps the
/// body valid: `_pad` is not a recognized field and is ignored.
fn padded_tools_result(id: &str, limit: usize) -> String {
    let unpadded = format!(r#"{{"jsonrpc":"2.0","id":{id},"result":{{"tools":[],"_pad":""}}}}"#);
    let pad_len = limit.saturating_sub(unpadded.len());
    format!(
        r#"{{"jsonrpc":"2.0","id":{id},"result":{{"tools":[],"_pad":"{}"}}}}"#,
        "x".repeat(pad_len)
    )
}

/// On an established session, a `tools/list` response whose `Content-Length`
/// declares it larger than `max_response_body_size` must fail that request
/// with a transport error naming the limit, checked before any byte of the
/// body is read. The connection itself is fine, so a following request must
/// still succeed.
#[tokio::test]
async fn session_response_oversized_with_content_length_fails_request() {
    let limit = 256;
    let url = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_initialize_ok(req)
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") {
            let body = padded_tools_result(&raw_request_id(req), limit * 4);
            raw_response("200 OK", RAW_JSON, &body)
        } else {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        }
    })
    .await;

    let config = HttpClientConfig {
        max_response_body_size: limit,
        ..raw_client_config()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools hung on an oversized response instead of failing");
    let error = result
        .expect_err("an oversized response with Content-Length should fail the request")
        .to_string();
    assert!(
        error.contains(&limit.to_string()),
        "error should mention the {limit}-byte limit, got: {error}"
    );

    tokio::time::timeout(Duration::from_secs(3), client.ping())
        .await
        .expect("ping after the oversized reply hung")
        .expect("client should stay usable after an oversized reply");
}

/// As above, but the server declares no length at all (`Connection: close`,
/// read until the socket closes), so the limit must be enforced by counting
/// bytes while streaming rather than by inspecting `Content-Length`.
#[tokio::test]
async fn session_response_oversized_without_content_length_fails_request() {
    let limit = 256;
    let url = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_initialize_ok(req)
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") {
            let body = padded_tools_result(&raw_request_id(req), limit * 4);
            raw_response_no_length("200 OK", RAW_JSON, &body)
        } else {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        }
    })
    .await;

    let config = HttpClientConfig {
        max_response_body_size: limit,
        ..raw_client_config()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools hung on a length-less oversized response instead of failing");
    let error = result
        .expect_err("an oversized response with no declared length should fail the request")
        .to_string();
    assert!(
        error.contains(&limit.to_string()),
        "error should mention the {limit}-byte limit, got: {error}"
    );

    tokio::time::timeout(Duration::from_secs(3), client.ping())
        .await
        .expect("ping after the oversized reply hung")
        .expect("client should stay usable after an oversized reply");
}

/// A response landing exactly at the limit is not oversized and must still
/// parse normally.
#[tokio::test]
async fn session_response_at_exact_limit_still_parses() {
    let limit = 256;
    let url = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_initialize_ok(req)
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") {
            let body = padded_tools_result(&raw_request_id(req), limit);
            assert_eq!(
                body.len(),
                limit,
                "test setup: body must land exactly at the limit"
            );
            raw_response("200 OK", RAW_JSON, &body)
        } else {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        }
    })
    .await;

    let config = HttpClientConfig {
        max_response_body_size: limit,
        ..raw_client_config()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools hung on an exact-limit response");
    assert!(
        result.is_ok(),
        "a response exactly at the limit should still parse, got: {result:?}"
    );
}

/// Error bodies are capped independently of `max_response_body_size` at a
/// small fixed size and truncated rather than failed: the request must still
/// fail with the HTTP error itself, not a size-limit error, and the oversized
/// body must not appear in full in the surfaced message.
#[tokio::test]
async fn pre_session_oversized_error_body_is_truncated_not_failed_by_size() {
    let huge_body = "x".repeat(200 * 1024); // well over the 64 KiB error-body cap
    let url =
        spawn_raw_server(move |_| raw_response("500 Internal Server Error", RAW_JSON, &huge_body))
            .await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();

    let result = tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("initialize hung on an oversized error body instead of failing");
    let error = result
        .expect_err("initialize should fail with the HTTP error")
        .to_string();
    assert!(
        error.contains("HTTP 500"),
        "the request should fail with the HTTP error, not a size-limit error, got: {error}"
    );
    assert!(
        error.len() < 200 * 1024,
        "the oversized error body should be truncated in the surfaced message, got {} bytes",
        error.len()
    );
}

/// Before a session exists, `initialize`'s response goes through the
/// synchronous path. An oversized response there must fail `initialize`
/// instead of hanging until the request timeout.
#[tokio::test]
async fn pre_session_oversized_initialize_response_fails_without_hanging() {
    let limit = 256;
    let url = spawn_raw_server(move |req| {
        let id = raw_request_id(req);
        let big = "x".repeat(limit * 4);
        let body = format!(
            r#"{{"jsonrpc":"2.0","id":{id},"result":{{"protocolVersion":"2025-11-25","capabilities":{{}},"serverInfo":{{"name":"raw","version":"{big}"}}}}}}"#
        );
        raw_response(
            "200 OK",
            &format!("{RAW_JSON}mcp-session-id: raw-session\r\nmcp-protocol-version: 2025-11-25\r\n"),
            &body,
        )
    })
    .await;

    let config = HttpClientConfig {
        max_response_body_size: limit,
        ..raw_client_config()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();

    let result = tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("initialize hung on an oversized response instead of failing");
    let error = result
        .expect_err("an oversized initialize response should fail initialize")
        .to_string();
    assert!(
        error.contains(&limit.to_string()),
        "error should mention the {limit}-byte limit, got: {error}"
    );
}

// ---------------------------------------------------------------------------
// Redirect policy (#1475): custom headers and mcp-session-id must not leak
// to a redirect target reqwest's own sensitive-header stripping doesn't
// cover.
// ---------------------------------------------------------------------------

/// A same-origin `307` (same host:port, different path) is transparent: the
/// client follows it and the request still succeeds.
#[tokio::test]
async fn same_origin_redirect_is_followed() {
    let base = Arc::new(std::sync::OnceLock::new());
    let base_for_server = base.clone();
    let redirected_from_root = Arc::new(std::sync::atomic::AtomicBool::new(false));
    let redirected_from_root_clone = redirected_from_root.clone();
    let url = spawn_raw_server(move |req| {
        if req.starts_with("POST / ") {
            redirected_from_root_clone.store(true, std::sync::atomic::Ordering::SeqCst);
            let base: &String = base_for_server
                .get()
                .expect("base url set before first request");
            raw_response(
                "307 Temporary Redirect",
                &format!("Location: {base}/redirected\r\n"),
                "",
            )
        } else {
            raw_initialize_ok(req)
        }
    })
    .await;
    base.set(url.clone()).unwrap();

    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();

    tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("same-origin redirect should not hang")
    .expect("a same-origin redirect must be followed");

    assert!(
        redirected_from_root.load(std::sync::atomic::Ordering::SeqCst),
        "the request should have hit the redirecting path first"
    );
}

/// A cross-origin `307` (different port, so a different origin on
/// `127.0.0.1`) is not followed by default: the target never sees the
/// request, and the caller fails instead of hanging.
#[tokio::test]
async fn cross_origin_redirect_is_not_followed_by_default() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    let seen_by_target = Arc::new(AtomicUsize::new(0));
    let seen_by_target_clone = seen_by_target.clone();
    let url_target = spawn_raw_server(move |req| {
        seen_by_target_clone.fetch_add(1, Ordering::SeqCst);
        raw_initialize_ok(req)
    })
    .await;

    let redirect_location = format!("{url_target}/redirected");
    let url_origin = spawn_raw_server(move |_req| {
        raw_response(
            "307 Temporary Redirect",
            &format!("Location: {redirect_location}\r\n"),
            "",
        )
    })
    .await;

    let client = McpClient::connect(HttpClientTransport::with_config(
        url_origin,
        raw_client_config(),
    ))
    .await
    .unwrap();

    let result = tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("a blocked cross-origin redirect must fail, not hang");
    let error = result
        .expect_err("a cross-origin redirect must not be followed by default")
        .to_string();
    assert!(
        error.contains(&format!("{url_target}/redirected")),
        "the error should name the redirect target, got: {error}"
    );
    assert_eq!(
        seen_by_target.load(Ordering::SeqCst),
        0,
        "the redirect target must never receive the request"
    );
}

/// With `follow_cross_origin_redirects` set, the same cross-origin `307` is
/// followed and the request succeeds.
#[tokio::test]
async fn cross_origin_redirect_is_followed_when_opted_in() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    let seen_by_target = Arc::new(AtomicUsize::new(0));
    let seen_by_target_clone = seen_by_target.clone();
    let url_target = spawn_raw_server(move |req| {
        seen_by_target_clone.fetch_add(1, Ordering::SeqCst);
        raw_initialize_ok(req)
    })
    .await;

    // Only `initialize` is redirected; `notifications/initialized` is
    // answered directly so the target's request count stays exact.
    let redirect_location = format!("{url_target}/redirected");
    let url_origin = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_response(
                "307 Temporary Redirect",
                &format!("Location: {redirect_location}\r\n"),
                "",
            )
        } else {
            raw_response("202 Accepted", "", "")
        }
    })
    .await;

    let config = HttpClientConfig {
        follow_cross_origin_redirects: true,
        ..raw_client_config()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url_origin, config))
        .await
        .unwrap();

    tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("an opted-in cross-origin redirect should not hang")
    .expect("an opted-in cross-origin redirect must be followed");

    assert_eq!(
        seen_by_target.load(Ordering::SeqCst),
        1,
        "the redirect target should have received exactly one request"
    );
}

/// The same block applies on the established-session (background task) send
/// path, not just the pre-session synchronous path: a redirected `tools/list`
/// fails instead of hanging. Like any other `reqwest::Error` on that path
/// (network error, timeout, ...), the blocked redirect also marks the
/// transport disconnected -- that part is unrelated to this fix, just the
/// existing behavior for a background POST that fails at the `send()` level
/// rather than with an HTTP-level error.
#[tokio::test]
async fn established_session_cross_origin_redirect_fails_without_hanging() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    let seen_by_target = Arc::new(AtomicUsize::new(0));
    let seen_by_target_clone = seen_by_target.clone();
    let url_target = spawn_raw_server(move |req| {
        seen_by_target_clone.fetch_add(1, Ordering::SeqCst);
        raw_initialize_ok(req)
    })
    .await;

    let redirect_location = format!("{url_target}/redirected");
    let url_origin = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_initialize_ok(req)
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") {
            raw_response(
                "307 Temporary Redirect",
                &format!("Location: {redirect_location}\r\n"),
                "",
            )
        } else {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        }
    })
    .await;

    let client = McpClient::connect(HttpClientTransport::with_config(
        url_origin,
        raw_client_config(),
    ))
    .await
    .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("a blocked cross-origin redirect must fail, not hang the caller");
    let error = result
        .expect_err("a cross-origin redirect must not be followed by default")
        .to_string();
    assert!(
        error.contains(&format!("{url_target}/redirected")),
        "the error should name the redirect target, got: {error}"
    );
    assert_eq!(
        seen_by_target.load(Ordering::SeqCst),
        0,
        "the redirect target must never receive the request"
    );
}

// ---------------------------------------------------------------------------
// Drop aborts background tasks and terminates the session (#1479)
// ---------------------------------------------------------------------------

/// A raw server for the drop tests below: answers `initialize` and
/// `notifications/initialized` normally, holds every SSE `GET` open until the
/// client side closes it (reported on `get_closed_tx`), and records every
/// `DELETE` it receives, including the `mcp-session-id` header line, on
/// `delete_tx`. Returns the URL and both receivers.
async fn spawn_session_drop_server() -> (
    String,
    tokio::sync::mpsc::Receiver<()>,
    tokio::sync::mpsc::Receiver<String>,
) {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (get_closed_tx, get_closed_rx) = tokio::sync::mpsc::channel::<()>(4);
    let (delete_tx, delete_rx) = tokio::sync::mpsc::channel::<String>(4);

    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            let get_closed_tx = get_closed_tx.clone();
            let delete_tx = delete_tx.clone();
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                if req.starts_with("DELETE") {
                    let session_header = req
                        .lines()
                        .find(|l| l.to_lowercase().starts_with("mcp-session-id"))
                        .unwrap_or_default()
                        .to_string();
                    let _ = delete_tx.send(session_header).await;
                    let _ = stream
                        .write_all(raw_response("204 No Content", "", "").as_bytes())
                        .await;
                    let _ = stream.flush().await;
                    return;
                }
                if raw_is(&req, "initialize") {
                    let _ = stream.write_all(raw_initialize_ok(&req).as_bytes()).await;
                    let _ = stream.flush().await;
                    return;
                }
                if raw_is(&req, "notifications/initialized") {
                    let _ = stream
                        .write_all(raw_response("202 Accepted", "", "").as_bytes())
                        .await;
                    let _ = stream.flush().await;
                    return;
                }
                if req.starts_with("GET") {
                    // Hold the SSE stream open (no Content-Length, no body)
                    // and block on a read; the client's Drop-driven close is
                    // what unblocks it with EOF.
                    let header = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                                   Connection: keep-alive\r\n\r\n";
                    let _ = stream.write_all(header.as_bytes()).await;
                    let _ = stream.flush().await;
                    let mut buf = [0u8; 16];
                    loop {
                        match stream.read(&mut buf).await {
                            Ok(0) | Err(_) => {
                                let _ = get_closed_tx.send(()).await;
                                break;
                            }
                            Ok(_) => continue,
                        }
                    }
                    return;
                }
                let _ = stream
                    .write_all(raw_response("202 Accepted", "", "").as_bytes())
                    .await;
                let _ = stream.flush().await;
            });
        }
    });

    (format!("http://{addr}"), get_closed_rx, delete_rx)
}

/// #1479: dropping a connected session client with the SSE stream open must
/// close that GET connection instead of leaving its task detached. Before the
/// `Drop` impl this connection stayed open until the server's own idle
/// timeout, since nothing ever aborted the SSE background task.
#[tokio::test]
async fn drop_closes_the_open_sse_stream() {
    let (url, mut get_closed_rx, _delete_rx) = spawn_session_drop_server().await;
    let config = HttpClientConfig {
        auto_sse: true,
        request_timeout: Duration::from_secs(5),
        ..Default::default()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    // Let the background SSE GET actually reach the server before dropping;
    // otherwise there is nothing yet open to close.
    tokio::time::sleep(Duration::from_millis(100)).await;

    drop(client);

    tokio::time::timeout(Duration::from_secs(3), get_closed_rx.recv())
        .await
        .expect("server never observed the GET connection close (the bug this guards)")
        .expect("get_closed channel closed unexpectedly");
}

/// #1479: dropping a connected session client must send a best-effort
/// session-termination `DELETE` carrying the session id, the same as an
/// explicit `close()`.
#[tokio::test]
async fn drop_sends_a_session_delete() {
    let (url, _get_closed_rx, mut delete_rx) = spawn_session_drop_server().await;
    let config = HttpClientConfig {
        auto_sse: true,
        request_timeout: Duration::from_secs(5),
        ..Default::default()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    drop(client);

    let session_header = tokio::time::timeout(Duration::from_secs(3), delete_rx.recv())
        .await
        .expect("server never received a session-termination DELETE (the bug this guards)")
        .expect("delete channel closed unexpectedly");
    assert!(
        session_header.to_lowercase().contains("raw-session"),
        "DELETE should carry the session id, got header line: {session_header}"
    );
}

/// #1479: dropping after an explicit `close()` (reached here through
/// `McpClient::shutdown()`, which awaits `close()` before the transport is
/// dropped) must not send a second session-termination `DELETE`. `close()`
/// clears the session id once its own `DELETE` is sent, which is what makes
/// `Drop` a no-op afterward.
#[tokio::test]
async fn drop_after_shutdown_does_not_send_a_second_delete() {
    let (url, _get_closed_rx, mut delete_rx) = spawn_session_drop_server().await;
    let config = HttpClientConfig {
        auto_sse: true,
        request_timeout: Duration::from_secs(5),
        ..Default::default()
    };
    let client = McpClient::connect(HttpClientTransport::with_config(url, config))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    client.shutdown().await.unwrap();

    let first = tokio::time::timeout(Duration::from_secs(3), delete_rx.recv())
        .await
        .expect("shutdown's own close() never sent a DELETE")
        .expect("delete channel closed unexpectedly");
    assert!(first.to_lowercase().contains("raw-session"));

    // `shutdown()` already consumed the client, and its message loop dropped
    // the transport right after `close()` returned, so there is nothing left
    // to drop explicitly here: this bounded wait is what proves the absence
    // of a second DELETE rather than triggering one.
    let second = tokio::time::timeout(Duration::from_millis(300), delete_rx.recv()).await;
    assert!(
        second.is_err(),
        "a second DELETE was sent after close() already terminated the session"
    );
}

/// #1479: the 2026-07-28 stateless path never establishes a session (no
/// `mcp-session-id` is ever assigned; see the `!is_modern_request` guard on
/// every session_id assignment in `send_http_request`), so dropping a
/// stateless connection must not attempt a session-termination `DELETE`.
#[cfg(feature = "stateless")]
#[tokio::test]
async fn drop_sends_no_delete_for_a_stateless_connection() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    let delete_count = Arc::new(AtomicUsize::new(0));
    let delete_count_clone = delete_count.clone();
    let url = spawn_raw_server(move |req| {
        if req.starts_with("DELETE") {
            delete_count_clone.fetch_add(1, Ordering::SeqCst);
            return raw_response("204 No Content", "", "");
        }
        if raw_is(req, "server/discover") {
            let body = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{"supportedVersions":["2026-07-28"],"capabilities":{{}}}}}}"#,
                raw_request_id(req)
            );
            return raw_response("200 OK", RAW_JSON, &body);
        }
        raw_response("202 Accepted", "", "")
    })
    .await;

    let client = McpClient::builder()
        .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect(
            HttpClientTransport::with_config(url, raw_client_config()),
            NotificationHandler::new(),
        )
        .await
        .unwrap();
    client.discover("raw-client", "1.0.0").await.unwrap();

    drop(client);

    // Bound the wait so a regression (a wrongly-spawned DELETE) would show up
    // as a nonzero count rather than the test racing ahead of it.
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(
        delete_count.load(Ordering::SeqCst),
        0,
        "a stateless connection must never send a session-termination DELETE"
    );
}

// ---------------------------------------------------------------------------
// A transient HTTP status fails one request, not the transport (#1553)
// ---------------------------------------------------------------------------

/// A raw server that answers the first `tools/list` with `status`, `headers`
/// and a plain-text `body`, then answers everything normally.
async fn spawn_flaky_list_server(
    status: &'static str,
    headers: &'static str,
    body: &'static str,
) -> String {
    use std::sync::atomic::{AtomicBool, Ordering};

    let failed_once = Arc::new(AtomicBool::new(false));
    spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            raw_initialize_ok(req)
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") && !failed_once.swap(true, Ordering::SeqCst) {
            raw_response(status, headers, body)
        } else {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{"tools":[]}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        }
    })
    .await
}

async fn flaky_list_outcome(
    status: &'static str,
    headers: &'static str,
    body: &'static str,
) -> (tower_mcp::Error, McpClient) {
    let url = spawn_flaky_list_server(status, headers, body).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let error = tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("the failed list_tools hung")
        .expect_err("the first list_tools must fail");
    (error, client)
}

#[tokio::test]
async fn gateway_503_fails_one_request_and_client_stays_usable() {
    let (error, client) = flaky_list_outcome(
        "503 Service Unavailable",
        "Content-Type: text/plain\r\n",
        "upstream unavailable",
    )
    .await;

    let details = HttpStatusDetails::from_error(&error)
        .unwrap_or_else(|| panic!("the error should carry the HTTP status, got: {error:?}"));
    assert_eq!(details.status, 503);
    assert_eq!(details.retry_after, None);
    assert_eq!(details.retry_after_delay(), None);

    tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools after the 503 hung")
        .expect("the same client must serve the next request");
}

#[tokio::test]
async fn rate_limit_429_exposes_retry_after() {
    let (error, client) = flaky_list_outcome(
        "429 Too Many Requests",
        "Content-Type: text/plain\r\nRetry-After: 7\r\n",
        "slow down",
    )
    .await;

    let details = HttpStatusDetails::from_error(&error)
        .unwrap_or_else(|| panic!("the error should carry the HTTP status, got: {error:?}"));
    assert_eq!(details.status, 429);
    assert_eq!(details.retry_after.as_deref(), Some("7"));
    assert_eq!(details.retry_after_delay(), Some(Duration::from_secs(7)));

    tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools after the 429 hung")
        .expect("the same client must serve the next request");
}

#[tokio::test]
async fn retry_after_http_date_is_passed_through_raw() {
    let (error, _client) = flaky_list_outcome(
        "503 Service Unavailable",
        "Retry-After: Wed, 21 Oct 2026 07:28:00 GMT\r\n",
        "",
    )
    .await;

    let details = HttpStatusDetails::from_error(&error).expect("status details");
    assert_eq!(
        details.retry_after.as_deref(),
        Some("Wed, 21 Oct 2026 07:28:00 GMT")
    );
    assert_eq!(details.retry_after_delay(), None);
}

// ---------------------------------------------------------------------------
// A 401, 403 or 429 with a JSON-RPC error body keeps its HTTP status (#1555)
// ---------------------------------------------------------------------------

fn jsonrpc_error_body(data: Option<&str>) -> String {
    let data = data.map_or(String::new(), |data| format!(r#","data":{data}"#));
    format!(
        r#"{{"jsonrpc":"2.0","id":null,"error":{{"code":-32001,"message":"denied by policy"{data}}}}}"#
    )
}

/// Leak a test-built body to the `'static` the flaky-server helper takes.
fn leak(body: String) -> &'static str {
    Box::leak(body.into_boxed_str())
}

fn rpc_error(error: &tower_mcp::Error) -> &JsonRpcError {
    match error {
        tower_mcp::Error::JsonRpc(rpc) => rpc,
        other => panic!("expected a JSON-RPC error, got: {other:?}"),
    }
}

#[tokio::test]
async fn forbidden_jsonrpc_error_body_exposes_status_and_keeps_server_data() {
    let body = leak(jsonrpc_error_body(Some(r#"{"reason":"missing scope"}"#)));
    let (error, client) = flaky_list_outcome("403 Forbidden", RAW_JSON, body).await;

    let details = HttpStatusDetails::from_error(&error)
        .unwrap_or_else(|| panic!("the error should carry the HTTP status, got: {error:?}"));
    assert_eq!(details.status, 403);
    assert_eq!(details.www_authenticate, None);

    // The server's own error is what the caller sees, with its data intact.
    let rpc = rpc_error(&error);
    assert_eq!(rpc.code, -32001);
    assert_eq!(rpc.message, "denied by policy");
    let data = rpc.data.as_ref().unwrap();
    assert_eq!(data["reason"], "missing scope");
    assert_eq!(data["httpStatus"], 403);

    tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("list_tools after the 403 hung")
        .expect("the same client must serve the next request");
}

#[tokio::test]
async fn unauthorized_jsonrpc_error_body_exposes_the_challenge() {
    let challenge = r#"Bearer error="invalid_token", resource_metadata="https://example.test/.well-known/oauth-protected-resource""#;
    let headers = leak(format!(
        "{RAW_JSON}WWW-Authenticate: {challenge}
Retry-After: 3
"
    ));
    let (error, _client) =
        flaky_list_outcome("401 Unauthorized", headers, leak(jsonrpc_error_body(None))).await;

    let details = HttpStatusDetails::from_error(&error).expect("status details");
    assert_eq!(details.status, 401);
    assert_eq!(details.www_authenticate.as_deref(), Some(challenge));
    assert_eq!(details.retry_after.as_deref(), Some("3"));
    assert_eq!(rpc_error(&error).code, -32001);
    assert_eq!(rpc_error(&error).message, "denied by policy");
}

#[tokio::test]
async fn unauthorized_without_a_jsonrpc_body_exposes_the_challenge() {
    let (error, _client) = flaky_list_outcome(
        "401 Unauthorized",
        "Content-Type: text/plain\r\nWWW-Authenticate: Bearer realm=\"mcp\"\r\n",
        "unauthorized",
    )
    .await;

    let details = HttpStatusDetails::from_error(&error).expect("status details");
    assert_eq!(details.status, 401);
    assert_eq!(
        details.www_authenticate.as_deref(),
        Some(r#"Bearer realm="mcp""#)
    );
}

#[tokio::test]
async fn server_status_keys_in_error_data_are_not_overwritten() {
    let body = leak(jsonrpc_error_body(Some(r#"{"httpStatus":"upstream"}"#)));
    let (error, _client) = flaky_list_outcome("403 Forbidden", RAW_JSON, body).await;

    let data = rpc_error(&error).data.as_ref().unwrap();
    assert_eq!(data["httpStatus"], "upstream");
}

#[tokio::test]
async fn non_object_error_data_moves_under_server_data() {
    let body = leak(jsonrpc_error_body(Some(r#""policy 7""#)));
    let (error, _client) = flaky_list_outcome("403 Forbidden", RAW_JSON, body).await;

    let details = HttpStatusDetails::from_error(&error).expect("status details");
    assert_eq!(details.status, 403);
    let data = rpc_error(&error).data.as_ref().unwrap();
    assert_eq!(data["serverData"], "policy 7");
}

/// Statuses that are not about authorization or rate limits leave the
/// server's error exactly as sent.
#[tokio::test]
async fn in_band_error_on_a_400_gets_no_status_added() {
    let body = leak(jsonrpc_error_body(None));
    let (error, _client) = flaky_list_outcome("400 Bad Request", RAW_JSON, body).await;

    assert_eq!(rpc_error(&error).code, -32001);
    assert_eq!(rpc_error(&error).data, None);
    assert_eq!(HttpStatusDetails::from_error(&error), None);
}

// ---------------------------------------------------------------------------
// A 404 on a session request recovers the session, whatever the body (#1556)
// ---------------------------------------------------------------------------

/// A raw server that issues a numbered session on each `initialize`, answers
/// the first `tools/list` with a 404 carrying `body`, and answers `tools/list`
/// only for the second session afterwards. Returns the URL and the count of
/// `initialize` requests seen.
async fn spawn_expiring_session_server(
    body: &'static str,
) -> (String, Arc<std::sync::atomic::AtomicUsize>) {
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    let initializes = Arc::new(AtomicUsize::new(0));
    let expired_once = Arc::new(AtomicBool::new(false));
    let counter = initializes.clone();
    let url = spawn_raw_server(move |req| {
        if raw_is(req, "initialize") {
            let n = counter.fetch_add(1, Ordering::SeqCst) + 1;
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{"protocolVersion":"2025-11-25","capabilities":{{}},"serverInfo":{{"name":"raw","version":"0"}}}}}}"#,
                raw_request_id(req)
            );
            raw_response(
                "200 OK",
                &format!(
                    "{RAW_JSON}mcp-session-id: session-{n}\r\nmcp-protocol-version: 2025-11-25\r\n"
                ),
                &reply,
            )
        } else if raw_is(req, "notifications/initialized") {
            raw_response("202 Accepted", "", "")
        } else if raw_is(req, "tools/list") && !expired_once.swap(true, Ordering::SeqCst) {
            raw_response("404 Not Found", if body.is_empty() { "" } else { RAW_JSON }, body)
        } else if req.to_lowercase().contains("mcp-session-id: session-2") {
            let reply = format!(
                r#"{{"jsonrpc":"2.0","id":{},"result":{{"tools":[]}}}}"#,
                raw_request_id(req)
            );
            raw_response("200 OK", RAW_JSON, &reply)
        } else {
            raw_response("404 Not Found", "", "")
        }
    })
    .await;
    (url, initializes)
}

async fn assert_session_404_recovers(body: &'static str) {
    let (url, initializes) = spawn_expiring_session_server(body).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    tokio::time::timeout(Duration::from_secs(5), client.list_tools())
        .await
        .expect("list_tools hung")
        .expect("the request must succeed on the re-initialized session");
    assert_eq!(initializes.load(std::sync::atomic::Ordering::SeqCst), 2);

    tokio::time::timeout(Duration::from_secs(3), client.list_tools())
        .await
        .expect("the next list_tools hung")
        .expect("the recovered session must keep serving requests");
}

#[tokio::test]
async fn bare_404_on_a_session_request_recovers_the_session() {
    assert_session_404_recovers("").await;
}

#[tokio::test]
async fn session_404_with_a_32001_body_recovers_the_session() {
    assert_session_404_recovers(
        r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32001,"message":"Session not found"}}"#,
    )
    .await;
}

#[tokio::test]
async fn session_404_with_a_32005_body_still_recovers_the_session() {
    assert_session_404_recovers(
        r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32005,"message":"Session not found"}}"#,
    )
    .await;
}

/// With recovery switched off the 404 is not turned into a session signal:
/// the request fails with the status, and the transport closes as before.
#[tokio::test]
async fn session_404_without_recovery_fails_and_closes_the_transport() {
    let (url, initializes) = spawn_expiring_session_server("").await;
    let client = McpClient::connect(HttpClientTransport::with_config(
        url,
        HttpClientConfig {
            session_recovery: false,
            ..raw_client_config()
        },
    ))
    .await
    .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    let error = tokio::time::timeout(Duration::from_secs(5), client.list_tools())
        .await
        .expect("list_tools hung")
        .expect_err("the 404 must fail the request");
    assert_eq!(
        HttpStatusDetails::from_error(&error).map(|details| details.status),
        Some(404),
        "{error:?}"
    );
    assert_eq!(initializes.load(std::sync::atomic::Ordering::SeqCst), 1);

    tokio::time::timeout(Duration::from_secs(3), async {
        // The closed flag reaches the client after the failed request's
        // error, so poll until a call reports the closed connection.
        while !client
            .list_tools()
            .await
            .is_err_and(|error| error.to_string().contains("Connection closed"))
        {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("a 404 without recovery must close the transport");
}

// =============================================================================
// Edge cases from #1558
// =============================================================================

/// A server that answers `initialize` with an SSE stream, pings the client on
/// that stream, and withholds the `InitializeResult` until the pong arrives.
/// The client has to answer the ping while the initialize response is still
/// open, which it cannot do if it reads that response to EOF before
/// dispatching anything.
#[tokio::test]
async fn initialize_response_stream_is_dispatched_as_frames_arrive() {
    use tokio::io::AsyncWriteExt;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let pong_received = Arc::new(tokio::sync::Notify::new());
    let pong_seen = Arc::new(std::sync::atomic::AtomicBool::new(false));

    let (pong_received_server, pong_seen_server) = (pong_received.clone(), pong_seen.clone());
    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            let pong_received = pong_received_server.clone();
            let pong_seen = pong_seen_server.clone();
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                if raw_is(&req, "initialize") {
                    let headers = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                                   mcp-session-id: raw-session\r\n\
                                   mcp-protocol-version: 2025-11-25\r\n\
                                   Connection: close\r\n\r\n";
                    let _ = stream.write_all(headers.as_bytes()).await;
                    let ping =
                        "data: {\"jsonrpc\":\"2.0\",\"id\":\"srv-ping\",\"method\":\"ping\"}\n\n";
                    let _ = stream.write_all(ping.as_bytes()).await;
                    let _ = stream.flush().await;

                    // Hold the InitializeResult until the pong has arrived.
                    let _ = tokio::time::timeout(Duration::from_secs(3), pong_received.notified())
                        .await;
                    if pong_seen.load(std::sync::atomic::Ordering::SeqCst) {
                        let result = format!(
                            "data: {{\"jsonrpc\":\"2.0\",\"id\":{},\"result\":{{\"protocolVersion\":\"2025-11-25\",\"capabilities\":{{}},\"serverInfo\":{{\"name\":\"raw\",\"version\":\"0\"}}}}}}\n\n",
                            raw_request_id(&req)
                        );
                        let _ = stream.write_all(result.as_bytes()).await;
                    }
                    let _ = stream.flush().await;
                    return;
                }
                if req.contains("\"id\":\"srv-ping\"") && req.contains("\"result\"") {
                    pong_seen.store(true, std::sync::atomic::Ordering::SeqCst);
                    pong_received.notify_one();
                }
                let _ = stream
                    .write_all(raw_response("202 Accepted", "", "").as_bytes())
                    .await;
                let _ = stream.flush().await;
            });
        }
    });

    let client = McpClient::connect(HttpClientTransport::with_config(
        format!("http://{addr}"),
        HttpClientConfig {
            request_timeout: Duration::from_secs(10),
            ..raw_client_config()
        },
    ))
    .await
    .unwrap();

    let started = std::time::Instant::now();
    let result = tokio::time::timeout(
        Duration::from_secs(5),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("initialize hung waiting on the ping it never answered")
    .expect("initialize must succeed once the ping is answered");
    assert_eq!(result.server_info.name, "raw");
    assert!(
        started.elapsed() < Duration::from_secs(2),
        "initialize should complete as soon as the pong lands, took {:?}",
        started.elapsed()
    );
    assert!(pong_seen.load(std::sync::atomic::Ordering::SeqCst));
}

/// A raw server for the cancellation tests: answers `server/discover` and
/// notifications normally, never answers `tools/call`, and reports the JSON
/// body of every POST it receives on the returned channel. `initialize` is
/// answered with a session, or with an SSE stream that issues the session and
/// then stays silent when `hold_initialize` is set.
async fn spawn_recording_server(
    hold_initialize: bool,
) -> (String, tokio::sync::mpsc::UnboundedReceiver<String>) {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (body_tx, body_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            let body_tx = body_tx.clone();
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                if let Some((_, body)) = req.split_once("\r\n\r\n") {
                    let _ = body_tx.send(body.to_string());
                }
                let hold =
                    (raw_is(&req, "initialize") && hold_initialize) || raw_is(&req, "tools/call");
                if hold {
                    if raw_is(&req, "initialize") {
                        let headers = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                                       mcp-session-id: raw-session\r\n\
                                       mcp-protocol-version: 2025-11-25\r\n\
                                       Connection: close\r\n\r\n";
                        let _ = stream.write_all(headers.as_bytes()).await;
                        let _ = stream.flush().await;
                    }
                    // Never answer; hold the connection until the client
                    // closes it.
                    let mut buf = [0u8; 16];
                    while matches!(stream.read(&mut buf).await, Ok(n) if n > 0) {}
                    return;
                }
                if raw_is(&req, "initialize") {
                    let _ = stream.write_all(raw_initialize_ok(&req).as_bytes()).await;
                } else if raw_is(&req, "server/discover") {
                    let body = format!(
                        r#"{{"jsonrpc":"2.0","id":{},"result":{{"supportedVersions":["2026-07-28"],"capabilities":{{}}}}}}"#,
                        raw_request_id(&req)
                    );
                    let _ = stream
                        .write_all(raw_response("200 OK", RAW_JSON, &body).as_bytes())
                        .await;
                } else {
                    let _ = stream
                        .write_all(raw_response("202 Accepted", "", "").as_bytes())
                        .await;
                }
                let _ = stream.flush().await;
            });
        }
    });
    (format!("http://{addr}"), body_rx)
}

/// Read bodies off `rx` until `matches` accepts one, within `limit`.
async fn next_body_matching(
    rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
    limit: Duration,
    matches: impl Fn(&serde_json::Value) -> bool,
) -> Option<serde_json::Value> {
    tokio::time::timeout(limit, async {
        while let Some(body) = rx.recv().await {
            if let Ok(value) = serde_json::from_str::<serde_json::Value>(&body)
                && matches(&value)
            {
                return Some(value);
            }
        }
        None
    })
    .await
    .ok()
    .flatten()
}

/// On a 2025-11-25 session a closed connection is not a cancellation, so
/// abandoning a request has to say so with `notifications/cancelled`, as it
/// does on stdio.
#[tokio::test]
async fn dropping_a_session_request_sends_notifications_cancelled() {
    let (url, mut bodies) = spawn_recording_server(false).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    // Abandon the call once the server has received it.
    let client = Arc::new(client);
    let call = tokio::spawn({
        let client = client.clone();
        async move {
            let _ = client.call_tool("slow", serde_json::json!({})).await;
        }
    });
    let request = next_body_matching(&mut bodies, Duration::from_secs(3), |v| {
        v["method"] == "tools/call"
    })
    .await
    .expect("the server never received the tools/call");
    let request_id = request["id"].clone();
    call.abort();

    let cancelled = next_body_matching(&mut bodies, Duration::from_secs(3), |v| {
        v["method"] == "notifications/cancelled"
    })
    .await
    .expect("no notifications/cancelled reached the server for the abandoned request");
    assert_eq!(cancelled["params"]["requestId"], request_id);
}

/// The spec forbids cancelling `initialize`, so abandoning one that has
/// already issued a session sends nothing.
#[tokio::test]
async fn dropping_initialize_sends_no_notifications_cancelled() {
    let (url, mut bodies) = spawn_recording_server(true).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();

    // The server issues the session and then goes quiet, so this times out
    // and drops the initialize future with its request in flight.
    let outcome = tokio::time::timeout(
        Duration::from_millis(300),
        client.initialize("raw-client", "1.0.0"),
    )
    .await;
    assert!(outcome.is_err(), "initialize should still be waiting");

    let cancelled = next_body_matching(&mut bodies, Duration::from_millis(500), |v| {
        v["method"] == "notifications/cancelled"
    })
    .await;
    assert!(cancelled.is_none(), "initialize must not be cancelled");
}

/// On the 2026-07-28 path a disconnect is the cancellation, and there is no
/// session to address a notification to, so none is sent.
#[cfg(feature = "stateless")]
#[tokio::test]
async fn dropping_a_final_request_sends_no_notifications_cancelled() {
    let (url, mut bodies) = spawn_recording_server(false).await;
    let client = McpClient::builder()
        .protocol_support(tower_mcp::ProtocolSupport::try_new(["2026-07-28"]).unwrap())
        .connect(
            HttpClientTransport::with_config(url, raw_client_config()),
            NotificationHandler::new(),
        )
        .await
        .unwrap();
    client.discover("raw-client", "1.0.0").await.unwrap();

    let client = Arc::new(client);
    let call = tokio::spawn({
        let client = client.clone();
        async move {
            let _ = client.call_tool("slow", serde_json::json!({})).await;
        }
    });
    next_body_matching(&mut bodies, Duration::from_secs(3), |v| {
        v["method"] == "tools/call"
    })
    .await
    .expect("the server never received the tools/call");
    call.abort();

    let cancelled = next_body_matching(&mut bodies, Duration::from_millis(500), |v| {
        v["method"] == "notifications/cancelled"
    })
    .await;
    assert!(
        cancelled.is_none(),
        "the final protocol cancels by disconnecting, not by notification"
    );
}

/// A raw server that answers `initialize` with `reply` (which issues the
/// session `raw-session`) and reports the `mcp-session-id` header line of
/// every `DELETE` it receives.
async fn spawn_initialize_server(
    reply: fn(&str) -> String,
) -> (String, tokio::sync::mpsc::UnboundedReceiver<String>) {
    use tokio::io::AsyncWriteExt;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (delete_tx, delete_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
    tokio::spawn(async move {
        loop {
            let Ok((mut stream, _)) = listener.accept().await else {
                break;
            };
            let delete_tx = delete_tx.clone();
            tokio::spawn(async move {
                let Some(req) = read_http_request(&mut stream).await else {
                    return;
                };
                let response = if req.starts_with("DELETE") {
                    let session_header = req
                        .lines()
                        .find(|l| l.to_lowercase().starts_with("mcp-session-id"))
                        .unwrap_or_default()
                        .to_string();
                    let _ = delete_tx.send(session_header);
                    raw_response("204 No Content", "", "")
                } else if raw_is(&req, "initialize") {
                    reply(&req)
                } else {
                    raw_response("202 Accepted", "", "")
                };
                let _ = stream.write_all(response.as_bytes()).await;
                let _ = stream.flush().await;
            });
        }
    });
    (format!("http://{addr}"), delete_rx)
}

fn session_headers() -> String {
    format!("{RAW_JSON}mcp-session-id: raw-session\r\nmcp-protocol-version: 2025-11-25\r\n")
}

/// Run a failing `initialize` against `reply`, and require exactly one
/// `DELETE` for the session it issued, even after the client is dropped.
async fn assert_failed_initialize_deletes_the_session(reply: fn(&str) -> String) {
    let (url, mut deletes) = spawn_initialize_server(reply).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();

    tokio::time::timeout(
        Duration::from_secs(3),
        client.initialize("raw-client", "1.0.0"),
    )
    .await
    .expect("initialize hung")
    .expect_err("the handshake must fail");

    let header = tokio::time::timeout(Duration::from_secs(3), deletes.recv())
        .await
        .expect("the failed initialize's session was never deleted")
        .expect("delete channel closed unexpectedly");
    assert!(
        header.to_lowercase().contains("raw-session"),
        "DELETE should carry the session id, got header line: {header}"
    );

    // The session id was cleared, so dropping the client sends no second one.
    drop(client);
    assert!(
        tokio::time::timeout(Duration::from_millis(300), deletes.recv())
            .await
            .is_err(),
        "a second DELETE was sent for the same session"
    );
}

/// #1558: a session issued alongside a JSON-RPC error result is deleted when
/// `initialize` fails, not left until shutdown or the server's idle timeout.
#[tokio::test]
async fn failed_initialize_deletes_the_session_it_was_issued() {
    assert_failed_initialize_deletes_the_session(|req| {
        let body = format!(
            r#"{{"jsonrpc":"2.0","id":{},"error":{{"code":-32603,"message":"init failed"}}}}"#,
            raw_request_id(req)
        );
        raw_response("200 OK", &session_headers(), &body)
    })
    .await;
}

/// #1558: a server that picks a protocol version the client does not accept
/// has `initialize` rejected after the transport already took the session.
#[tokio::test]
async fn initialize_rejecting_the_protocol_version_deletes_the_session() {
    assert_failed_initialize_deletes_the_session(|req| {
        let body = format!(
            r#"{{"jsonrpc":"2.0","id":{},"result":{{"protocolVersion":"1999-01-01","capabilities":{{}},"serverInfo":{{"name":"raw","version":"0"}}}}}}"#,
            raw_request_id(req)
        );
        raw_response("200 OK", &session_headers(), &body)
    })
    .await;
}

/// A successful `initialize` keeps its session until shutdown.
#[tokio::test]
async fn successful_initialize_does_not_delete_the_session() {
    let (url, mut deletes) = spawn_initialize_server(raw_initialize_ok).await;
    let client = McpClient::connect(HttpClientTransport::with_config(url, raw_client_config()))
        .await
        .unwrap();
    client.initialize("raw-client", "1.0.0").await.unwrap();

    assert!(
        tokio::time::timeout(Duration::from_millis(300), deletes.recv())
            .await
            .is_err(),
        "a healthy session was deleted"
    );
}