llm-kernel 0.25.0

//! HTTP/SSE remote transport for MCP.
//!
//! Exposes an [`McpServer`] over HTTP: a JSON-RPC endpoint (`POST /mcp`) and an
//! SSE endpoint (`POST /mcp/sse`) that streams the response as a server-sent
//! event. Both reuse the server's `Authorization` (Bearer) check, so a server
//! secured for stdio is secured identically over HTTP.
//!
//! The transport holds the server behind an `Arc` (shared across request
//! tasks) and dispatches `tools/call` via [`McpServer::call_tool_async`], so
//! async handlers work transparently over HTTP.
//!
//! Requires the `mcp-http` feature (axum + tokio).

use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;

use axum::Json;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::routing::post;
use serde_json::Value;
use tokio_stream::wrappers::UnboundedReceiverStream;

use crate::mcp::McpServer;

/// Shared MCP server state for the HTTP transport.
#[derive(Clone)]
pub struct HttpTransport {
    server: Arc<McpServer>,
}

impl HttpTransport {
    /// Wrap a shared MCP server for HTTP serving.
    pub fn new(server: Arc<McpServer>) -> Self {
        Self { server }
    }

    /// Build the axum router with JSON-RPC and SSE routes.
    pub fn router(&self) -> axum::Router {
        axum::Router::new()
            .route("/mcp", post(rpc_handler))
            .route("/mcp/sse", post(sse_handler))
            .with_state(self.clone())
    }
}

/// Run the MCP HTTP transport on `addr` until the server is stopped.
pub async fn serve(server: Arc<McpServer>, addr: SocketAddr) -> std::io::Result<()> {
    let transport = HttpTransport::new(server);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, transport.router()).await?;
    Ok(())
}

/// JSON-RPC code for "method not found".
const ERR_METHOD_NOT_FOUND: i32 = -32601;
/// JSON-RPC code for invalid params (unknown tool / prompt / resource).
const ERR_INVALID_PARAMS: i32 = -32602;
/// JSON-RPC code for a tool-execution / internal error.
const ERR_INTERNAL: i32 = -32603;
/// JSON-RPC code for unauthorized access.
const ERR_UNAUTHORIZED: i32 = -32001;

/// Dispatch a single JSON-RPC request against the server (async path).
///
/// `tools/call` is awaited via [`McpServer::call_tool_async`]; `initialize`,
/// `ping`, `tools/list`, `resources/list`, `resources/templates/list`,
/// `prompts/list`, `prompts/get`, and `resources/read` are handled
/// synchronously. Notifications (no `id`) return `None`.
async fn dispatch_async(server: &McpServer, req: &Value) -> Option<Value> {
    // Notifications (no id) get no response.
    let id = req.get("id")?.clone();
    let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");

    let result: Result<Value, (i32, String)> = match method {
        "initialize" => {
            let requested = req
                .pointer("/params/protocolVersion")
                .and_then(|v| v.as_str());
            Ok(server.initialize_response(requested))
        }
        "ping" => Ok(serde_json::json!({})),
        "tools/list" => Ok(serde_json::json!({ "tools": server.tools() })),
        "resources/list" => Ok(serde_json::json!({ "resources": server.resources() })),
        "resources/templates/list" => Ok(serde_json::json!({ "resourceTemplates": [] })),
        "prompts/list" => Ok(serde_json::json!({ "prompts": server.prompts() })),
        "prompts/get" => {
            let name = req
                .pointer("/params/name")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            let args = req
                .pointer("/params/arguments")
                .cloned()
                .unwrap_or(serde_json::json!({}));
            server
                .get_prompt(name, args)
                .map_err(|e| (ERR_INVALID_PARAMS, e.to_string()))
        }
        "resources/read" => {
            let uri = req
                .pointer("/params/uri")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            server
                .read_resource(uri, serde_json::json!({}))
                .map(|content| {
                    serde_json::json!({
                        "contents": [{ "uri": uri, "text": content.to_string() }]
                    })
                })
                .map_err(|e| (ERR_INTERNAL, e.to_string()))
        }
        "tools/call" => {
            let name = req
                .pointer("/params/name")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            let params = req
                .pointer("/params/arguments")
                .cloned()
                .unwrap_or(serde_json::json!(null));
            if !server.has_tool(name) {
                Err((ERR_INVALID_PARAMS, format!("Unknown tool: {name}")))
            } else if let Err(e) = server.validate_tool_args(name, &params) {
                Err((ERR_INVALID_PARAMS, e))
            } else {
                // Execution failures are reported in-band with isError: true.
                match server.call_tool_async(name, params).await {
                    Ok(r) => Ok(serde_json::json!({
                        "content": [{ "type": "text", "text": r.to_string() }],
                        "isError": false
                    })),
                    Err(e) => Ok(serde_json::json!({
                        "content": [{ "type": "text", "text": e.to_string() }],
                        "isError": true
                    })),
                }
            }
        }
        _ => Err((ERR_METHOD_NOT_FOUND, format!("Method not found: {method}"))),
    };

    Some(match result {
        Ok(value) => serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": value }),
        Err((code, message)) => serde_json::json!({
            "jsonrpc": "2.0", "id": id,
            "error": { "code": code, "message": message }
        }),
    })
}

/// Extract and validate the `Authorization` header. Returns `true` if the
/// request may proceed.
fn authorized(server: &McpServer, headers: &HeaderMap) -> bool {
    let auth = headers
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");
    server.check_auth(auth)
}

/// Reject cross-origin browser requests (MCP spec: servers MUST validate
/// `Origin` to prevent DNS rebinding). A page on any website can POST to a
/// loopback MCP server; without this, that page executes tools.
///
/// Non-browser clients send no `Origin` and are unaffected.
fn origin_allowed(headers: &HeaderMap) -> bool {
    let Some(origin) = headers.get("origin").and_then(|v| v.to_str().ok()) else {
        return true; // no Origin — not a browser-initiated request
    };
    if origin == "null" {
        return false;
    }
    // Only loopback origins may drive a local MCP server. Parse the host
    // bracket-aware: an IPv6 origin is `http://[::1]:3000`, where a naive
    // `split(':').next()` yields "[" and rejects a legitimate loopback.
    origin
        .split_once("://")
        .map(|(_, host_port)| {
            if let Some(rest) = host_port.strip_prefix('[') {
                rest.split(']').next().unwrap_or("") // "::1" (no brackets)
            } else {
                host_port.split(':').next().unwrap_or("")
            }
        })
        .is_some_and(|host| host == "localhost" || host == "127.0.0.1" || host == "::1")
}

fn forbidden_response(id: Option<Value>) -> Json<Value> {
    Json(serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": { "code": ERR_UNAUTHORIZED, "message": "Forbidden origin" }
    }))
}

fn unauthorized_response(id: Option<Value>) -> Json<Value> {
    Json(serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": { "code": ERR_UNAUTHORIZED, "message": "Unauthorized" }
    }))
}

/// Dispatch a single request or a JSON-RPC batch (array). Returns `None` only
/// when nothing needs answering (all notifications).
async fn dispatch_any(server: &McpServer, req: &Value) -> Option<Value> {
    let Some(batch) = req.as_array() else {
        return dispatch_async(server, req).await;
    };
    let mut out = Vec::with_capacity(batch.len());
    for item in batch {
        if let Some(resp) = dispatch_async(server, item).await {
            out.push(resp);
        }
    }
    if out.is_empty() {
        None
    } else {
        Some(Value::Array(out))
    }
}

async fn rpc_handler(
    State(state): State<HttpTransport>,
    headers: HeaderMap,
    Json(req): Json<Value>,
) -> impl IntoResponse {
    let id = req.get("id").cloned();
    if !origin_allowed(&headers) {
        return (StatusCode::FORBIDDEN, forbidden_response(id));
    }
    if !authorized(&state.server, &headers) {
        return (StatusCode::UNAUTHORIZED, unauthorized_response(id));
    }
    match dispatch_any(&state.server, &req).await {
        Some(resp) => (StatusCode::OK, Json(resp)),
        // Notification — acknowledge with 204 No Content.
        None => (StatusCode::NO_CONTENT, Json(serde_json::Value::Null)),
    }
}

async fn sse_handler(
    State(state): State<HttpTransport>,
    headers: HeaderMap,
    Json(req): Json<Value>,
) -> impl IntoResponse {
    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
    let server = state.server.clone();

    // Produce the response for this request, then stream it as one SSE event.
    tokio::spawn(async move {
        let event = if !origin_allowed(&headers) {
            Event::default().event("error").data(
                serde_json::to_string(&forbidden_response(req.get("id").cloned()).0)
                    .unwrap_or_default(),
            )
        } else if !authorized(&server, &headers) {
            Event::default().event("error").data(
                serde_json::to_string(&unauthorized_response(req.get("id").cloned()).0)
                    .unwrap_or_default(),
            )
        } else if let Some(resp) = dispatch_any(&server, &req).await {
            let data = serde_json::to_string(&resp).unwrap_or_default();
            Event::default().event("message").data(data)
        } else {
            // Notification — no response event.
            Event::default().event("noop")
        };
        let _ = tx.send(Ok::<_, Infallible>(event));
    });

    Sse::new(UnboundedReceiverStream::new(rx)).keep_alive(KeepAlive::default())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::schema::{ResourceDescription, ToolDescription};

    fn server_with_echo() -> McpServer {
        let mut server = McpServer::new("http-test", "1.0.0");
        server.register_tool(ToolDescription {
            name: "echo".into(),
            description: "Echo".into(),
            input_schema: serde_json::json!({"type": "object"}),
        });
        server.set_async_handler("echo", |params| async move { Ok(params) });
        server
    }

    #[tokio::test]
    async fn dispatch_initialize() {
        let server = server_with_echo();
        let req = serde_json::json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{}});
        let resp = dispatch_async(&server, &req).await.unwrap();
        assert_eq!(resp["result"]["serverInfo"]["name"], "http-test");
    }

    #[tokio::test]
    async fn dispatch_tools_call_async() {
        let server = server_with_echo();
        let req = serde_json::json!({
            "jsonrpc": "2.0", "id": 2, "method": "tools/call",
            "params": { "name": "echo", "arguments": { "msg": "hello" } }
        });
        let resp = dispatch_async(&server, &req).await.unwrap();
        let text = resp["result"]["content"][0]["text"].as_str().unwrap();
        assert!(text.contains("hello"));
    }

    #[tokio::test]
    async fn dispatch_unknown_method() {
        let server = server_with_echo();
        let req = serde_json::json!({"jsonrpc":"2.0","id":3,"method":"nope"});
        let resp = dispatch_async(&server, &req).await.unwrap();
        assert_eq!(resp["error"]["code"], ERR_METHOD_NOT_FOUND);
    }

    /// AC2: HTTP dispatch also serves `resources/read`, not just tools.
    #[tokio::test]
    async fn dispatch_resources_read() {
        let mut server = McpServer::new("http-test", "1.0.0");
        server.register_resource(ResourceDescription {
            uri: "docs://x".into(),
            name: "X".into(),
            description: None,
            mime_type: None,
        });
        server.set_resource_handler("docs://x", |_| Ok(serde_json::json!("# body")));
        let req = serde_json::json!({
            "jsonrpc": "2.0", "id": 4, "method": "resources/read",
            "params": { "uri": "docs://x" }
        });
        let resp = dispatch_async(&server, &req).await.unwrap();
        let text = resp["result"]["contents"][0]["text"].as_str().unwrap();
        assert!(text.contains("body"));
    }

    #[test]
    fn origin_validation_blocks_cross_site_browsers() {
        let mut h = HeaderMap::new();
        assert!(origin_allowed(&h), "no Origin (non-browser client) passes");
        h.insert("origin", "http://localhost:3000".parse().unwrap());
        assert!(origin_allowed(&h));
        h.insert("origin", "http://127.0.0.1:8080".parse().unwrap());
        assert!(origin_allowed(&h));
        h.insert("origin", "https://evil.example.com".parse().unwrap());
        assert!(!origin_allowed(&h), "DNS-rebinding origin must be rejected");
        h.insert("origin", "null".parse().unwrap());
        assert!(!origin_allowed(&h));
        // Suffix trickery must not pass.
        h.insert("origin", "https://localhost.evil.com".parse().unwrap());
        assert!(!origin_allowed(&h));
        // IPv6 loopback — a naive `split(':')` would see "[" and reject it.
        h.insert("origin", "http://[::1]:3000".parse().unwrap());
        assert!(origin_allowed(&h), "IPv6 loopback must pass: {h:?}");
        h.insert("origin", "http://[::1]".parse().unwrap());
        assert!(origin_allowed(&h), "IPv6 loopback (no port) must pass");
        h.insert("origin", "http://[fe80::1]:3000".parse().unwrap());
        assert!(!origin_allowed(&h), "non-loopback IPv6 must be rejected");
    }

    #[tokio::test]
    async fn batch_requests_get_a_batch_response() {
        let server = server_with_echo();
        let batch = serde_json::json!([
            {"jsonrpc":"2.0","id":1,"method":"ping"},
            {"jsonrpc":"2.0","id":2,"method":"tools/call",
             "params":{"name":"echo","arguments":{"v":1}}}
        ]);
        let resp = dispatch_any(&server, &batch)
            .await
            .expect("batch must not be dropped");
        let arr = resp.as_array().expect("array response");
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0]["id"], 1);
        assert_eq!(arr[1]["result"]["isError"], false);
    }

    /// AC2: a full HTTP round-trip — bind an ephemeral port, POST a tools/call,
    /// and read the JSON-RPC response off the wire.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn http_round_trip_calls_tool() {
        let server = Arc::new(server_with_echo());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        // Hand the listener to axum in a background task.
        let transport = HttpTransport::new(server);
        tokio::spawn(async move {
            let _ = axum::serve(listener, transport.router()).await;
        });

        let body = serde_json::to_string(&serde_json::json!({
            "jsonrpc": "2.0", "id": 9, "method": "tools/call",
            "params": { "name": "echo", "arguments": { "v": 42 } }
        }))
        .unwrap();
        let req = format!(
            "POST /mcp HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            body.len(),
            body
        );

        let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        stream.write_all(req.as_bytes()).await.unwrap();
        let mut buf = Vec::new();
        stream.read_to_end(&mut buf).await.unwrap();
        let response = String::from_utf8_lossy(&buf);
        assert!(response.contains("200 OK"), "response: {response}");
        // The tool result is JSON-encoded inside the `text` field, so its quotes
        // are escaped on the wire — assert on the unescaped value + content shape.
        assert!(response.contains("\"content\""), "response: {response}");
        assert!(response.contains("\\\"v\\\":42"), "response: {response}");
    }
}