alkhttp/tests/from_mcp_integration.rs

//! Integration test for `FromMCP`: spins up a real rmcp streamable HTTP MCP
//! server, imports its tools via `FromMCP::import()`, and invokes a
//! forwarding handler end-to-end. Verifies the handler calls the remote MCP
//! tool via rmcp and reads `context.capabilities` (not `std::env::var`).

#![cfg(feature = "mcp")]

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use alkcall::client::OperationAdapter;
use alkcall::core::types::Capabilities;
use alkcall::protocol::wire::ResponseEnvelope;
use alkcall::registry::context::{AbortPolicy, OperationContext, ScopedPeerEnv};
use alkcall::registry::env::OperationEnv;
use alkcall::registry::registration::{HandlerKind, OperationProvenance};
use alkhttp::adapters::FromMCP;
use axum::Router;
use rmcp::model::{
    CallToolRequestParams, CallToolResult, Content, ListToolsResult, PaginatedRequestParams, Tool,
};
use rmcp::service::RequestContext;
use rmcp::transport::{
    streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService},
    StreamableHttpServerConfig,
};
use rmcp::{RoleServer, ServerHandler};
use serde_json::Value;

struct NoopEnv;

#[async_trait::async_trait]
impl OperationEnv for NoopEnv {
    async fn invoke_with_policy(
        &self,
        _ns: &str,
        _op: &str,
        _input: Value,
        parent: &OperationContext,
        _policy: AbortPolicy,
    ) -> ResponseEnvelope {
        ResponseEnvelope::ok(parent.request_id.clone(), Value::Null)
    }

    fn contains(&self, _name: &str) -> bool {
        false
    }
}

fn test_context(request_id: &str, caps: Capabilities) -> OperationContext {
    OperationContext {
        request_id: request_id.to_string(),
        parent_request_id: None,
        identity: None,
        handler_identity: None,
        forwarded_for: None,
        capabilities: caps,
        metadata: HashMap::new(),
        scoped_env: ScopedPeerEnv::empty(),
        env: Arc::new(NoopEnv),
        abort_policy: AbortPolicy::default(),
        deadline: Some(Instant::now() + Duration::from_secs(30)),
        internal: true,
        ownership: None,
    }
}

struct EchoServer;

impl ServerHandler for EchoServer {
    fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>>
           + rmcp::service::MaybeSendFuture
           + '_ {
        let tools = vec![
            Tool::new_with_raw(
                "echo",
                Some("Echo the input back as structured content".into()),
                Arc::new(serde_json::Map::new()),
            )
            .with_raw_output_schema(Arc::new(serde_json::Map::from_iter([(
                "type".to_string(),
                Value::String("object".into()),
            )]))),
            Tool::new_with_raw(
                "legacy",
                Some("Legacy tool returning text content blocks".into()),
                Arc::new(serde_json::Map::new()),
            ),
        ];
        std::future::ready(Ok(ListToolsResult {
            meta: None,
            next_cursor: None,
            tools,
        }))
    }

    fn call_tool(
        &self,
        request: CallToolRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<CallToolResult, rmcp::ErrorData>>
           + rmcp::service::MaybeSendFuture
           + '_ {
        let name = request.name.to_string();
        std::future::ready(Ok(match name.as_str() {
            "echo" => {
                let args = request.arguments.map(Value::Object).unwrap_or(Value::Null);
                CallToolResult::structured(serde_json::json!({ "echoed": args }))
            }
            "legacy" => CallToolResult::success(vec![Content::text("plain text result")]),
            other => CallToolResult::error(vec![Content::text(format!("unknown tool: {other}"))]),
        }))
    }

    fn get_info(&self) -> rmcp::model::ServerInfo {
        rmcp::model::ServerInfo::default()
    }
}

/// A paginating `tools/list` server (CON-01): three pages behind a
/// cursor chain; the importer must follow `next_cursor` to see all tools.
struct PagingServer;

impl ServerHandler for PagingServer {
    fn list_tools(
        &self,
        request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>>
           + rmcp::service::MaybeSendFuture
           + '_ {
        let cursor = request.and_then(|p| p.cursor);
        let page_tools: Vec<(&str, &str)> = match cursor.as_deref() {
            None => vec![("p1_a", "page 1"), ("p1_b", "page 1")],
            Some("page1") => vec![("p2_a", "page 2")],
            Some("page2") => vec![("p3_a", "page 3")],
            Some(_) => vec![],
        };
        let next_cursor = match cursor.as_deref() {
            None => Some("page1".to_string()),
            Some("page1") => Some("page2".to_string()),
            _ => None,
        };
        let tools = page_tools
            .into_iter()
            .map(|(name, desc)| {
                Tool::new_with_raw(
                    name.to_string(),
                    Some(desc.into()),
                    Arc::new(serde_json::Map::new()),
                )
            })
            .collect();
        std::future::ready(Ok(ListToolsResult {
            meta: None,
            next_cursor,
            tools,
        }))
    }

    fn get_info(&self) -> rmcp::model::ServerInfo {
        rmcp::model::ServerInfo::default()
    }
}

/// A hostile paging `tools/list` server (review-002 CON-14): always
/// responds with the same single tool and `next_cursor: Some("a")` — the
/// cursor cycles and never clears. The importer must bound the walk and
/// fail with a clean `DiscoveryFailed` rather than hang forever.
struct CyclingCursorServer;

impl ServerHandler for CyclingCursorServer {
    fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>>
           + rmcp::service::MaybeSendFuture
           + '_ {
        let tools = vec![Tool::new_with_raw(
            "hostile",
            Some("served on every page, forever".into()),
            Arc::new(serde_json::Map::new()),
        )];
        std::future::ready(Ok(ListToolsResult {
            meta: None,
            next_cursor: Some("a".to_string()),
            tools,
        }))
    }

    fn get_info(&self) -> rmcp::model::ServerInfo {
        rmcp::model::ServerInfo::default()
    }
}

async fn spawn_server_for<S: ServerHandler + 'static>(
    server: impl Fn() -> Result<S, std::io::Error> + Send + Sync + 'static,
) -> (String, tokio::task::JoinHandle<()>) {
    let mcp_service: StreamableHttpService<S, LocalSessionManager> = StreamableHttpService::new(
        server,
        LocalSessionManager::default().into(),
        StreamableHttpServerConfig::default(),
    );
    let app = Router::new().nest_service("/mcp", mcp_service);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let handle = tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    (format!("http://{addr}/mcp"), handle)
}

async fn spawn_server() -> (String, tokio::task::JoinHandle<()>) {
    spawn_server_for(|| Ok(EchoServer)).await
}

#[tokio::test]
async fn import_discovers_tools_and_builds_registrations() {
    let (endpoint, _handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter
        .import()
        .await
        .expect("import succeeds against running server");
    assert_eq!(bundles.len(), 2);
    let names: Vec<&str> = bundles.iter().map(|b| b.spec.name.as_str()).collect();
    assert!(names.contains(&"echo/echo"));
    assert!(names.contains(&"echo/legacy"));
    for b in &bundles {
        assert_eq!(b.provenance, OperationProvenance::FromMCP);
        assert!(b.composition_authority.is_none());
        assert!(b.scoped_env.is_none());
    }
}

#[tokio::test]
async fn forwarding_handler_calls_echo_and_returns_structured_content() {
    let (endpoint, _handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let echo = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/echo")
        .expect("echo tool present");

    let caps = Capabilities::new().with_http_token("mcp", "unused-on-server".to_string());
    let ctx = test_context("req-echo", caps);
    let input = serde_json::json!({ "msg": "hello" });
    let response = match &echo.handler {
        HandlerKind::Once(h) => h(input, ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => {
            panic!("expected Once handler for echo tool")
        }
    };

    assert_eq!(response.request_id, "req-echo");
    match response.result {
        Ok(v) => {
            let obj = v.as_object().expect("structured object");
            assert!(obj.contains_key("echoed"));
        }
        Err(e) => panic!("expected Ok, got Err: {e:?}"),
    }
}

#[tokio::test]
async fn forwarding_handler_calls_legacy_and_returns_content_blocks() {
    let (endpoint, _handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let legacy = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/legacy")
        .expect("legacy tool present");

    let ctx = test_context("req-legacy", Capabilities::new());
    let response = match &legacy.handler {
        HandlerKind::Once(h) => h(serde_json::json!({}), ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => {
            panic!("expected Once handler for legacy tool")
        }
    };

    match response.result {
        Ok(Value::Array(blocks)) => {
            assert_eq!(blocks.len(), 1);
            assert_eq!(blocks[0]["type"], "text");
            assert_eq!(blocks[0]["text"], "plain text result");
        }
        other => panic!("expected array of content blocks, got {other:?}"),
    }
}

#[tokio::test]
async fn forwarding_handler_does_not_read_env_vars() {
    std::env::set_var("MCP_TOKEN", "should-not-be-used");
    let (endpoint, _handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let echo = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/echo")
        .expect("echo tool present");

    let ctx = test_context("req-noenv", Capabilities::new());
    let response = match &echo.handler {
        HandlerKind::Once(h) => h(serde_json::json!({ "x": 1 }), ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => {
            panic!("expected Once handler for echo tool")
        }
    };
    assert!(response.result.is_ok(), "handler works without env var");
    std::env::remove_var("MCP_TOKEN");
}

#[tokio::test]
async fn import_unreachable_server_returns_discovery_failed() {
    let adapter = FromMCP::new("http://127.0.0.1:1/mcp", "x");
    match adapter.import().await {
        Ok(_) => panic!("expected Err for unreachable server"),
        Err(alkcall::client::AdapterError::DiscoveryFailed { .. }) => {}
        Err(alkcall::client::AdapterError::Transport { .. }) => {}
        Err(other) => panic!("expected DiscoveryFailed or Transport, got {other}"),
    }
}

#[tokio::test]
async fn import_follows_tools_list_pagination() {
    let (endpoint, _handle) = spawn_server_for(|| Ok(PagingServer)).await;
    let adapter = FromMCP::new(endpoint, "pg");
    let bundles = adapter
        .import()
        .await
        .expect("import follows every tools/list page (CON-01)");
    let mut names: Vec<&str> = bundles.iter().map(|b| b.spec.name.as_str()).collect();
    names.sort();
    assert_eq!(
        names,
        vec!["pg/p1_a", "pg/p1_b", "pg/p2_a", "pg/p3_a"],
        "all three pages must be imported"
    );
}

#[tokio::test]
async fn import_cycling_cursor_fails_bounded_with_discovery_failed() {
    let (endpoint, _handle) = spawn_server_for(|| Ok(CyclingCursorServer)).await;
    let adapter = FromMCP::new(endpoint, "x");
    let started = Instant::now();
    let result = tokio::time::timeout(Duration::from_secs(10), adapter.import()).await;
    let elapsed = started.elapsed();

    let outcome = result.expect("import must terminate, not hang (CON-14)");
    match outcome {
        Ok(_) => panic!("expected Err for a cycling-cursor server"),
        Err(alkcall::client::AdapterError::DiscoveryFailed { message }) => {
            assert!(
                message.contains("pagination exceeded budget"),
                "error must name the budget, got: {message}"
            );
            assert!(
                message.contains("page"),
                "error must name pages fetched, got: {message}"
            );
            assert!(
                message.contains("tool"),
                "error must name tools accumulated, got: {message}"
            );
        }
        Err(other) => panic!("expected DiscoveryFailed, got {other}"),
    }
    assert!(
        elapsed < Duration::from_secs(10),
        "bounded walk must trip the page cap well inside the outer guard, took {elapsed:?}"
    );
}

#[tokio::test]
async fn import_refuses_tool_name_containing_slash() {
    struct SlashToolServer;

    impl ServerHandler for SlashToolServer {
        fn list_tools(
            &self,
            _request: Option<PaginatedRequestParams>,
            _context: RequestContext<RoleServer>,
        ) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>>
               + rmcp::service::MaybeSendFuture
               + '_ {
            let tools = vec![Tool::new_with_raw(
                "weird/tool",
                Some("a tool name with a slash".into()),
                Arc::new(serde_json::Map::new()),
            )];
            std::future::ready(Ok(ListToolsResult {
                meta: None,
                next_cursor: None,
                tools,
            }))
        }

        fn get_info(&self) -> rmcp::model::ServerInfo {
            rmcp::model::ServerInfo::default()
        }
    }

    let (endpoint, _handle) = spawn_server_for(|| Ok(SlashToolServer)).await;
    let adapter = FromMCP::new(endpoint, "ns");
    match adapter.import().await {
        Ok(_) => panic!("expected Err for remote tool name containing `/`"),
        Err(alkcall::client::AdapterError::SchemaParse { message }) => {
            assert!(message.contains('/'), "error names the offending tool");
        }
        Err(other) => panic!("expected SchemaParse, got {other}"),
    }
}

#[tokio::test]
async fn forwarding_handler_wraps_non_object_input_as_value_field() {
    // A scalar/array tool argument cannot be a JSON object on the MCP
    // wire; the adapter wraps it as {"value": <input>} (review-002
    // from_mcp :460-468 arm). Proven over the real rmcp round trip: the
    // echo server reflects the arguments object back.
    let (endpoint, _handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let echo = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/echo")
        .expect("echo tool present");

    let ctx = test_context("req-wrap", Capabilities::new());
    let response = match &echo.handler {
        HandlerKind::Once(h) => h(serde_json::json!("bare-scalar"), ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => panic!("expected Once handler"),
    };
    match response.result {
        Ok(Value::Object(obj)) => {
            assert_eq!(
                obj.get("echoed"),
                Some(&serde_json::json!({ "value": "bare-scalar" })),
                "non-object input must reach the wire as {{\"value\": …}}: got {obj:?}"
            );
        }
        other => panic!("expected object structured content, got {other:?}"),
    }
}

#[tokio::test]
async fn forwarding_handler_maps_json_rpc_tool_error_with_code_fidelity() {
    // A server whose `call_tool` returns a JSON-RPC error (rmcp
    // `ErrorData`): the forwarding handler must surface the remote
    // error code (MCP_JRPC_<code>), not a flattened INTERNAL.
    struct JsonRpcErrorServer;

    impl ServerHandler for JsonRpcErrorServer {
        fn list_tools(
            &self,
            _request: Option<PaginatedRequestParams>,
            _context: RequestContext<RoleServer>,
        ) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>>
               + rmcp::service::MaybeSendFuture
               + '_ {
            let tools = vec![Tool::new_with_raw(
                "boom",
                Some("this tool always fails with a JSON-RPC error".into()),
                Arc::new(serde_json::Map::new()),
            )];
            std::future::ready(Ok(ListToolsResult {
                meta: None,
                next_cursor: None,
                tools,
            }))
        }

        fn call_tool(
            &self,
            _request: CallToolRequestParams,
            _context: RequestContext<RoleServer>,
        ) -> impl std::future::Future<Output = Result<CallToolResult, rmcp::ErrorData>>
               + rmcp::service::MaybeSendFuture
               + '_ {
            std::future::ready(Err(rmcp::ErrorData::resource_not_found(
                "tool not found on the remote server",
                Some(serde_json::json!({ "detail": "unknown-tool" })),
            )))
        }

        fn get_info(&self) -> rmcp::model::ServerInfo {
            rmcp::model::ServerInfo::default()
        }
    }

    let (endpoint, _handle) = spawn_server_for(|| Ok(JsonRpcErrorServer)).await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let any = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/boom")
        .expect("boom tool present");

    let ctx = test_context("req-jrpc", Capabilities::new());
    let response = match &any.handler {
        HandlerKind::Once(h) => h(serde_json::json!({}), ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => panic!("expected Once handler"),
    };
    match response.result {
        Err(e) => {
            assert_eq!(
                e.code, "MCP_JRPC_-32002",
                "JSON-RPC error code preserved in the call-error code, got {e:?}"
            );
            assert_eq!(e.message, "tool not found on the remote server");
            let details = e.details.expect("JSON-RPC data preserved as details");
            assert_eq!(details["detail"], "unknown-tool");
        }
        Ok(_) => panic!("expected Err for the JSON-RPC error path"),
    }
}

#[tokio::test]
async fn transport_call_failure_maps_to_declared_mcp_transport_error() {
    // Import against a live server, kill it, then call: the in-flight
    // handler surfaces MCP_TRANSPORT_ERROR (declared), not INTERNAL.
    let (endpoint, handle) = spawn_server().await;
    let adapter = FromMCP::new(endpoint, "echo");
    let bundles = adapter.import().await.expect("import succeeds");
    let echo = bundles
        .into_iter()
        .find(|b| b.spec.name == "echo/echo")
        .expect("echo tool present");

    handle.abort();

    let ctx = test_context("req-transport", Capabilities::new());
    let response = match &echo.handler {
        HandlerKind::Once(h) => h(serde_json::json!({ "x": 1 }), ctx).await,
        HandlerKind::Stream(_) | HandlerKind::Sink(_) => panic!("expected Once handler"),
    };
    match response.result {
        Err(e) => {
            assert_eq!(
                e.code, "MCP_TRANSPORT_ERROR",
                "declared transport failure mode (CON-11), got {e:?}"
            );
            assert!(e.retryable, "transport failure is retryable");
        }
        Ok(_) => panic!("expected Err after server shutdown"),
    }
}