mcp-hub 0.3.0

//! Mock upstream MCP server used by integration tests.

use std::{future::Future, io::Write, sync::Arc, time::Duration};

use anyhow::{Context, Result};
#[path = "../telemetry.rs"]
mod telemetry;

use rmcp::{
    ErrorData as McpError, Peer, ServerHandler, ServiceExt,
    model::{
        CallToolRequestParams, CallToolResult, CancelledNotificationParam, ContentBlock,
        Extensions, Implementation, ListToolsResult, Meta, PaginatedRequestParams,
        ProgressNotification, ProgressNotificationParam, ProgressToken, ProtocolVersion, Resource,
        ServerCapabilities, ServerInfo, ServerNotification, TaskSupport, Tool, ToolAnnotations,
        ToolExecution, ToolsCapability,
    },
    service::{MaybeSendFuture, NotificationContext, RequestContext, RoleServer},
    transport::stdio,
};
use serde_json::{Value, json};
use tokio::sync::Mutex;

/// Uses `mimalloc` as the global allocator when the default feature is enabled.
#[cfg(feature = "mimalloc")]
#[global_allocator]
static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc;

/// Mock tool server with deterministic tool behavior for integration tests.
#[derive(Clone)]
struct MockUpstreamServer {
    server_name: String,
    list_tools_delay: Option<Duration>,
    protocol_version: ProtocolVersion,
    extra_star_tool_name: Option<String>,
    termination_tool_enabled: bool,
    tool_list_change_notification_tool_enabled: bool,
    lifecycle_tools_enabled: bool,
    session_counter: Arc<Mutex<u64>>,
    cancellation: Arc<Mutex<Option<CancelledNotificationParam>>>,
}

impl MockUpstreamServer {
    /// Builds the mock server from CLI args, environment variables, and defaults.
    fn from_env() -> Self {
        Self {
            server_name: configured_server_name(),
            list_tools_delay: configured_list_tools_delay(),
            protocol_version: configured_protocol_version(),
            extra_star_tool_name: configured_star_tool_name(),
            termination_tool_enabled: termination_tool_enabled(),
            tool_list_change_notification_tool_enabled: tool_list_change_notification_tool_enabled(
            ),
            lifecycle_tools_enabled: lifecycle_tools_enabled(),
            session_counter: Arc::new(Mutex::new(0)),
            cancellation: Arc::new(Mutex::new(None)),
        }
    }

    /// Returns the static tool inventory exposed by the mock server.
    fn tools(&self) -> Vec<Tool> {
        let mut tools = vec![
            Tool::new(
                "echo",
                "Echo the provided message with the server name.",
                single_string_arg_schema("message", true),
            ),
            Tool::new(
                "duplicate",
                "Return a stable response used to test duplicate tool names.",
                empty_object_schema(),
            ),
            Tool::new(
                "session_counter",
                "Increment a session-local counter and return its new value.",
                empty_object_schema(),
            ),
            Tool::new(
                "read_snapshot",
                "Read a stable snapshot of mock state without mutating anything.",
                empty_object_schema(),
            )
            .with_title("Read Snapshot")
            .with_raw_output_schema(report_output_schema())
            .with_annotations(
                ToolAnnotations::with_title("Read Snapshot")
                    .read_only(true)
                    .destructive(false)
                    .idempotent(true)
                    .open_world(false),
            )
            .with_meta(tool_meta("read")),
            Tool::new(
                "delete_record",
                "Delete a record from the mock state.",
                single_string_arg_schema("record_id", true),
            )
            .with_title("Delete Record")
            .with_annotations(
                ToolAnnotations::with_title("Delete Record")
                    .read_only(false)
                    .destructive(true)
                    .idempotent(false)
                    .open_world(false),
            )
            .with_meta(tool_meta("destructive")),
            Tool::new(
                "publish_webhook",
                "Send an outbound webhook to an external destination.",
                single_string_arg_schema("url", true),
            )
            .with_title("Publish Webhook")
            .with_annotations(
                ToolAnnotations::with_title("Publish Webhook")
                    .read_only(false)
                    .destructive(false)
                    .idempotent(false)
                    .open_world(true),
            )
            .with_meta(tool_meta("external")),
            Tool::new(
                "structured_report",
                "Return structured content and metadata for protocol-level forwarding tests.",
                empty_object_schema(),
            )
            .with_title("Structured Report")
            .with_raw_output_schema(report_output_schema()),
            Tool::new(
                "resource_bundle",
                "Return mixed content including resource links and embedded resources.",
                empty_object_schema(),
            )
            .with_title("Resource Bundle")
            .with_meta(tool_meta("bundle")),
            Tool::new(
                "error_result",
                "Return a tool-level error result without turning it into a transport error.",
                empty_object_schema(),
            )
            .with_title("Error Result")
            .with_meta(tool_meta("error")),
            Tool::new(
                "task_only",
                "A task-required tool that the hub must filter from v1.",
                empty_object_schema(),
            )
            .with_execution(ToolExecution::new().with_task_support(TaskSupport::Required)),
        ];

        if let Some(tool_name) = &self.extra_star_tool_name {
            tools.push(Tool::new(
                tool_name.clone(),
                "A deliberately invalid tool name used to test startup rejection.",
                empty_object_schema(),
            ));
        }

        if self.termination_tool_enabled {
            tools.push(Tool::new(
                "terminate_after_startup",
                "Terminate the mock process after startup for routing-failure tests.",
                empty_object_schema(),
            ));
        }

        if self.tool_list_change_notification_tool_enabled {
            tools.push(Tool::new(
                "emit_tool_list_changed",
                "Emit a tool-list change notification for hub integration tests.",
                empty_object_schema(),
            ));
        }

        if self.lifecycle_tools_enabled {
            tools.extend([
                Tool::new(
                    "wait_for_cancel",
                    "Wait until the MCP request is cancelled for lifecycle forwarding tests.",
                    empty_object_schema(),
                ),
                Tool::new(
                    "cancellation_status",
                    "Return the latest cancellation observed by this mock upstream.",
                    empty_object_schema(),
                ),
                Tool::new(
                    "emit_progress",
                    "Emit one standard MCP progress notification for lifecycle forwarding tests.",
                    empty_object_schema(),
                ),
                Tool::new(
                    "emit_late_progress",
                    "Emit one progress notification after its final result for stale-route tests.",
                    empty_object_schema(),
                ),
                Tool::new(
                    "emit_unknown_progress",
                    "Emit progress with an uncorrelated token for lifecycle forwarding tests.",
                    empty_object_schema(),
                ),
            ]);
        }

        tools
    }
}

impl ServerHandler for MockUpstreamServer {
    /// Describes the mock server as a tools-only MCP peer for integration tests.
    fn get_info(&self) -> ServerInfo {
        let mut tools_capability = ToolsCapability::default();
        tools_capability.list_changed = self
            .tool_list_change_notification_tool_enabled
            .then_some(true);

        ServerInfo::new(
            ServerCapabilities::builder()
                .enable_tools_with(tools_capability)
                .build(),
        )
        .with_server_info(Implementation::new(
            format!("mock-upstream-{}", self.server_name),
            env!("CARGO_PKG_VERSION"),
        ))
        .with_protocol_version(self.protocol_version.clone())
    }

    /// Returns the mock tool inventory used by integration coverage.
    async fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, McpError> {
        if let Some(delay) = self.list_tools_delay {
            tokio::time::sleep(delay).await;
        }

        Ok(ListToolsResult {
            tools: self.tools(),
            next_cursor: None,
            meta: None,
        })
    }

    /// Executes deterministic mock tool behavior for the requested tool name.
    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        context: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, McpError> {
        match request.name.as_ref() {
            "echo" => {
                let message = request
                    .arguments
                    .as_ref()
                    .and_then(|arguments| arguments.get("message"))
                    .and_then(Value::as_str)
                    .unwrap_or_default();
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "{}:{}",
                    self.server_name, message
                ))]))
            }
            "duplicate" => Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                "duplicate:{}",
                self.server_name
            ))])),
            "session_counter" => {
                let mut counter = self.session_counter.lock().await;
                *counter += 1;
                Ok(CallToolResult::success(vec![ContentBlock::text(
                    counter.to_string(),
                )]))
            }
            "read_snapshot" => Ok(CallToolResult::structured(json!({
                "server": self.server_name,
                "kind": "read_snapshot",
                "readOnly": true
            }))
            .with_meta(Some(result_meta("read_snapshot")))),
            "delete_record" => {
                let record_id = string_argument(&request, "record_id");
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "deleted:{}:{}",
                    self.server_name, record_id
                ))]))
            }
            "publish_webhook" => {
                let url = string_argument(&request, "url");
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "published:{}:{}",
                    self.server_name, url
                ))]))
            }
            "structured_report" => Ok(CallToolResult::structured(json!({
                "server": self.server_name,
                "kind": "structured_report",
                "version": self.protocol_version.as_str(),
            }))
            .with_meta(Some(result_meta("structured_report")))),
            "resource_bundle" => {
                let mut resource_link = Resource::new(
                    format!("memory://{}/bundle.json", self.server_name),
                    format!("{} bundle", self.server_name),
                );
                resource_link.description =
                    Some("Structured bundle exported by the upstream".to_string());
                resource_link.mime_type = Some("application/json".to_string());
                resource_link.size = Some(128);

                Ok(CallToolResult::success(vec![
                    ContentBlock::text(format!("bundle:{}", self.server_name)),
                    ContentBlock::resource_link(resource_link),
                    ContentBlock::embedded_text(
                        format!("memory://{}/embedded.txt", self.server_name),
                        format!("embedded:{}", self.server_name),
                    ),
                    ContentBlock::image("ZmFrZS1pbWFnZS1kYXRh", "image/png"),
                ])
                .with_meta(Some(result_meta("resource_bundle"))))
            }
            "error_result" => Ok(CallToolResult::error(vec![ContentBlock::text(format!(
                "upstream-error:{}",
                self.server_name
            ))])
            .with_meta(Some(result_meta("error_result")))),
            "terminate_after_startup" if self.termination_tool_enabled => {
                std::process::exit(0);
            }
            "emit_tool_list_changed" if self.tool_list_change_notification_tool_enabled => {
                context
                    .peer
                    .notify_tool_list_changed()
                    .await
                    .map_err(|error| {
                        McpError::internal_error(
                            "failed to emit mock tool-list change notification",
                            Some(json!({ "reason": error.to_string() })),
                        )
                    })?;
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "tool-list-change-notification:{}",
                    self.server_name
                ))]))
            }
            "wait_for_cancel" if self.lifecycle_tools_enabled => {
                context.ct.cancelled().await;
                emit_progress_notification(&context.peer, context.meta.get_progress_token())
                    .await?;
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "cancelled:{}",
                    self.server_name
                ))]))
            }
            "cancellation_status" if self.lifecycle_tools_enabled => {
                let cancellation = self.cancellation.lock().await.clone();
                Ok(CallToolResult::structured(json!({
                    "server": self.server_name,
                    "cancellation": cancellation,
                })))
            }
            "emit_progress" if self.lifecycle_tools_enabled => {
                emit_progress_notification(&context.peer, context.meta.get_progress_token())
                    .await?;
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "progress:{}",
                    self.server_name
                ))]))
            }
            "emit_late_progress" if self.lifecycle_tools_enabled => {
                let peer = context.peer.clone();
                let progress_token = context.meta.get_progress_token();
                tokio::spawn(async move {
                    tokio::time::sleep(Duration::from_millis(25)).await;
                    let _ = emit_progress_notification(&peer, progress_token).await;
                });
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "late-progress:{}",
                    self.server_name
                ))]))
            }
            "emit_unknown_progress" if self.lifecycle_tools_enabled => {
                let unknown_token = serde_json::from_value(json!("mock-unknown-progress-token"))
                    .expect("mock unknown progress token must be valid");
                emit_progress_notification(&context.peer, Some(unknown_token)).await?;
                Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                    "unknown-progress:{}",
                    self.server_name
                ))]))
            }
            "task_only" => Ok(CallToolResult::success(vec![ContentBlock::text(format!(
                "task-only:{}",
                self.server_name
            ))])),
            other => Err(McpError::invalid_params(
                if self.extra_star_tool_name.as_deref() == Some(other) {
                    format!("invalid-star-tool:{}:{}", self.server_name, other)
                } else {
                    format!("unknown mock tool '{other}'")
                },
                None,
            )),
        }
    }

    /// Records standard MCP cancellation notifications received by this mock upstream.
    fn on_cancelled(
        &self,
        mut notification: CancelledNotificationParam,
        context: NotificationContext<RoleServer>,
    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
        let cancellation = self.cancellation.clone();
        async move {
            notification.meta = notification_context_meta(&context);
            *cancellation.lock().await = Some(notification);
        }
    }
}

/// Boots tracing, builds the Tokio runtime explicitly, and runs the mock upstream server.
fn main() -> Result<()> {
    emit_startup_stderr()?;
    let _tracing = telemetry::init_tracing("info");
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime for mock upstream server")?;

    runtime
        .block_on(run())
        .context("mock upstream server terminated with an error")
}

/// Runs the mock upstream MCP server over stdio for integration testing.
async fn run() -> Result<()> {
    let service = MockUpstreamServer::from_env()
        .serve(stdio())
        .await
        .context("failed to initialize mock upstream over stdio")?;
    service
        .waiting()
        .await
        .context("mock upstream service terminated unexpectedly")?;
    Ok(())
}

/// Parses the configured protocol version override for the mock upstream server.
fn configured_protocol_version() -> ProtocolVersion {
    std::env::var("MOCK_SERVER_PROTOCOL_VERSION")
        .map(|value| parse_protocol_version(&value))
        .unwrap_or_else(|_| ProtocolVersion::default())
}

/// Parses one optional tool-list delay used to simulate slow startup inventory responses.
fn configured_list_tools_delay() -> Option<Duration> {
    std::env::var("MOCK_SERVER_LIST_TOOLS_DELAY_MS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .filter(|value| *value > 0)
        .map(Duration::from_millis)
}

/// Writes configured direct stderr output before the mock starts its MCP handshake.
fn emit_startup_stderr() -> Result<()> {
    let mut stderr = std::io::stderr().lock();

    if let Ok(message) = std::env::var("MOCK_SERVER_STARTUP_STDERR") {
        writeln!(stderr, "{message}").context("failed to write mock startup stderr line")?;
    }

    let byte_count = std::env::var("MOCK_SERVER_STARTUP_STDERR_BYTES")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|value| *value > 0);
    if let Some(byte_count) = byte_count {
        let bytes = vec![b'x'; byte_count];
        stderr
            .write_all(&bytes)
            .context("failed to write mock startup stderr bytes")?;
    }

    stderr
        .flush()
        .context("failed to flush mock startup stderr")?;
    Ok(())
}

/// Resolves one optional extra tool name used to test invalid literal `*` discovery.
fn configured_star_tool_name() -> Option<String> {
    std::env::var("MOCK_SERVER_STAR_TOOL_NAME")
        .ok()
        .filter(|value| !value.is_empty())
}

/// Returns whether the test-only process-termination tool should be exposed.
fn termination_tool_enabled() -> bool {
    std::env::var("MOCK_SERVER_ENABLE_TERMINATION_TOOL").is_ok_and(|value| value == "1")
}

/// Returns whether the mock-only tool-list change notification tool should be exposed.
fn tool_list_change_notification_tool_enabled() -> bool {
    std::env::var("MOCK_SERVER_ENABLE_TOOL_LIST_CHANGED_NOTIFICATION_TOOL")
        .is_ok_and(|value| value == "1")
}

/// Returns whether the mock-only lifecycle tools should be exposed.
fn lifecycle_tools_enabled() -> bool {
    std::env::var("MOCK_SERVER_ENABLE_LIFECYCLE_TOOLS").is_ok_and(|value| value == "1")
}

/// Resolves the configured server name from CLI args, environment, or the default name.
fn configured_server_name() -> String {
    configured_cli_value("--server-name")
        .or_else(|| std::env::var("MOCK_SERVER_NAME").ok())
        .unwrap_or_else(|| "mock".to_string())
}

/// Reads one optional `--flag value` pair from the current mock-server command line.
fn configured_cli_value(flag: &str) -> Option<String> {
    let mut args = std::env::args().skip(1);
    args.find(|argument| argument == flag)?;
    args.next()
}

/// Parses one protocol version string using `rmcp`'s own deserializer.
fn parse_protocol_version(value: &str) -> ProtocolVersion {
    serde_json::from_value(json!(value)).expect("mock protocol version must be valid JSON")
}

/// Extracts one string argument from a tool request, defaulting to the empty string.
fn string_argument(request: &CallToolRequestParams, name: &str) -> String {
    request
        .arguments
        .as_ref()
        .and_then(|arguments| arguments.get(name))
        .and_then(Value::as_str)
        .unwrap_or_default()
        .to_string()
}

/// Sends one deterministic standard progress notification for lifecycle integration tests.
async fn emit_progress_notification(
    peer: &Peer<RoleServer>,
    progress_token: Option<ProgressToken>,
) -> Result<(), McpError> {
    let progress_token = progress_token.ok_or_else(|| {
        McpError::invalid_params("progress token is required for mock lifecycle tool", None)
    })?;
    let progress = ProgressNotificationParam::new(progress_token, 3.0)
        .with_total(5.0)
        .with_message("mock progress");
    let mut extensions = Extensions::new();
    extensions.insert(progress_meta());
    let mut notification = ProgressNotification::new(progress);
    notification.extensions = extensions;
    let notification = ServerNotification::ProgressNotification(notification);
    peer.send_notification(notification).await.map_err(|error| {
        McpError::internal_error(
            "failed to emit mock progress notification",
            Some(json!({ "reason": error.to_string() })),
        )
    })
}

/// Returns metadata retained by `rmcp` for one incoming notification.
fn notification_context_meta(context: &NotificationContext<RoleServer>) -> Option<Meta> {
    (!context.meta.0.is_empty())
        .then(|| context.meta.clone())
        .or_else(|| context.extensions.get::<Meta>().cloned())
}

/// Builds an empty JSON object schema used by mock tools without arguments.
fn empty_object_schema() -> Arc<rmcp::model::JsonObject> {
    Arc::new(
        serde_json::from_value(json!({
            "type": "object",
            "properties": {},
            "additionalProperties": false
        }))
        .expect("mock schema must be valid"),
    )
}

/// Builds a simple JSON object schema with one named string property.
fn single_string_arg_schema(name: &str, required: bool) -> Arc<rmcp::model::JsonObject> {
    Arc::new(
        serde_json::from_value(json!({
            "type": "object",
            "properties": {
                name: {
                    "type": "string"
                }
            },
            "required": if required { vec![name] } else { Vec::<&str>::new() },
            "additionalProperties": false
        }))
        .expect("mock schema must be valid"),
    )
}

/// Builds a structured output schema used by metadata-preservation tests.
fn report_output_schema() -> Arc<rmcp::model::JsonObject> {
    Arc::new(
        serde_json::from_value(json!({
            "type": "object",
            "properties": {
                "server": { "type": "string" },
                "kind": { "type": "string" }
            },
            "required": ["server", "kind"],
            "additionalProperties": false
        }))
        .expect("mock output schema must be valid"),
    )
}

/// Builds tool-level metadata used by outward inventory preservation tests.
fn tool_meta(kind: &str) -> Meta {
    let mut meta = Meta::new();
    meta.insert("mockKind".to_string(), json!(kind));
    meta.insert("source".to_string(), json!("mock-upstream"));
    meta
}

/// Builds result-level metadata used by forwarded tool-result tests.
fn result_meta(kind: &str) -> Meta {
    let mut meta = Meta::new();
    meta.insert("resultKind".to_string(), json!(kind));
    meta.insert("emittedBy".to_string(), json!("mock-upstream"));
    meta
}

/// Builds progress-notification metadata used by lifecycle forwarding tests.
fn progress_meta() -> Meta {
    let mut meta = Meta::new();
    meta.insert("progressKind".to_string(), json!("mock-progress"));
    meta.insert("emittedBy".to_string(), json!("mock-upstream"));
    meta
}