huskies/server/src/http/mcp/mod.rs

//! HTTP MCP server module.

use crate::http::context::AppContext;
use poem::handler;
use poem::http::StatusCode;
use poem::web::Data;
use poem::{Body, Request, Response};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::sync::Arc;

pub mod agent_tools;
pub mod diagnostics;
pub mod git_tools;
pub mod merge_tools;
pub mod qa_tools;
pub mod shell_tools;
pub mod status_tools;
pub mod story_tools;
pub mod wizard_tools;

mod dispatch;
mod tools_list;

use dispatch::handle_tools_call;
use tools_list::handle_tools_list;

/// Returns true when the Accept header includes text/event-stream.
fn wants_sse(req: &Request) -> bool {
    req.header("accept")
        .unwrap_or("")
        .contains("text/event-stream")
}

// ── JSON-RPC structs ──────────────────────────────────────────────

#[derive(Deserialize)]
struct JsonRpcRequest {
    jsonrpc: String,
    id: Option<Value>,
    method: String,
    #[serde(default)]
    params: Value,
}

#[derive(Serialize)]
pub(super) struct JsonRpcResponse {
    jsonrpc: &'static str,
    #[serde(skip_serializing_if = "Option::is_none")]
    id: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<JsonRpcError>,
}

#[derive(Serialize)]
struct JsonRpcError {
    code: i64,
    message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    data: Option<Value>,
}

impl JsonRpcResponse {
    pub(super) fn success(id: Option<Value>, result: Value) -> Self {
        Self {
            jsonrpc: "2.0",
            id,
            result: Some(result),
            error: None,
        }
    }

    pub(super) fn error(id: Option<Value>, code: i64, message: String) -> Self {
        Self {
            jsonrpc: "2.0",
            id,
            result: None,
            error: Some(JsonRpcError {
                code,
                message,
                data: None,
            }),
        }
    }
}

// ── Poem handlers ─────────────────────────────────────────────────

#[handler]
pub async fn mcp_get_handler() -> Response {
    Response::builder()
        .status(StatusCode::METHOD_NOT_ALLOWED)
        .body(Body::empty())
}

#[handler]
pub async fn mcp_post_handler(req: &Request, body: Body, ctx: Data<&Arc<AppContext>>) -> Response {
    // Validate Content-Type
    let content_type = req.header("content-type").unwrap_or("");
    if !content_type.is_empty() && !content_type.contains("application/json") {
        return json_rpc_error_response(
            None,
            -32700,
            "Unsupported Content-Type; expected application/json".into(),
        );
    }

    let bytes = match body.into_bytes().await {
        Ok(b) => b,
        Err(_) => return json_rpc_error_response(None, -32700, "Parse error".into()),
    };

    let rpc: JsonRpcRequest = match serde_json::from_slice(&bytes) {
        Ok(r) => r,
        Err(_) => return json_rpc_error_response(None, -32700, "Parse error".into()),
    };

    if rpc.jsonrpc != "2.0" {
        return json_rpc_error_response(rpc.id, -32600, "Invalid JSON-RPC version".into());
    }

    // Notifications (no id) — accept silently
    if rpc.id.is_none() || rpc.id.as_ref() == Some(&Value::Null) {
        if rpc.method.starts_with("notifications/") {
            return Response::builder()
                .status(StatusCode::ACCEPTED)
                .body(Body::empty());
        }
        return json_rpc_error_response(None, -32600, "Missing id".into());
    }

    let sse = wants_sse(req);

    // Streaming agent output over SSE
    if sse && rpc.method == "tools/call" {
        let tool_name = rpc
            .params
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        if tool_name == "run_command" {
            return shell_tools::handle_run_command_sse(rpc.id, &rpc.params, &ctx);
        }
    }

    let resp = match rpc.method.as_str() {
        "initialize" => handle_initialize(rpc.id, &rpc.params),
        "tools/list" => handle_tools_list(rpc.id),
        "tools/call" => handle_tools_call(rpc.id, &rpc.params, &ctx).await,
        _ => JsonRpcResponse::error(rpc.id, -32601, format!("Unknown method: {}", rpc.method)),
    };

    if sse {
        to_sse_response(resp)
    } else {
        to_json_response(resp)
    }
}

fn json_rpc_error_response(id: Option<Value>, code: i64, message: String) -> Response {
    to_json_response(JsonRpcResponse::error(id, code, message))
}

fn to_json_response(resp: JsonRpcResponse) -> Response {
    let body = serde_json::to_vec(&resp).unwrap_or_default();
    Response::builder()
        .status(StatusCode::OK)
        .header("Content-Type", "application/json")
        .body(Body::from(body))
}

pub(super) fn to_sse_response(resp: JsonRpcResponse) -> Response {
    let json = serde_json::to_string(&resp).unwrap_or_default();
    let sse_body = format!("data: {json}\n\n");
    Response::builder()
        .status(StatusCode::OK)
        .header("Content-Type", "text/event-stream")
        .header("Cache-Control", "no-cache")
        .body(Body::from_string(sse_body))
}

// ── MCP protocol handlers ─────────────────────────────────────────

fn handle_initialize(id: Option<Value>, params: &Value) -> JsonRpcResponse {
    let _protocol_version = params
        .get("protocolVersion")
        .and_then(|v| v.as_str())
        .unwrap_or("2025-03-26");

    JsonRpcResponse::success(
        id,
        json!({
            "protocolVersion": "2025-03-26",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": "huskies",
                "version": "1.0.0"
            }
        }),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::test_helpers::test_ctx;

    #[test]
    fn json_rpc_response_serializes_success() {
        let resp = JsonRpcResponse::success(Some(json!(1)), json!({"ok": true}));
        let s = serde_json::to_string(&resp).unwrap();
        assert!(s.contains("\"result\""));
        assert!(!s.contains("\"error\""));
    }

    #[test]
    fn json_rpc_response_serializes_error() {
        let resp = JsonRpcResponse::error(Some(json!(1)), -32600, "bad".into());
        let s = serde_json::to_string(&resp).unwrap();
        assert!(s.contains("\"error\""));
        assert!(!s.contains("\"result\""));
    }

    #[test]
    fn initialize_returns_capabilities() {
        let resp = handle_initialize(
            Some(json!(1)),
            &json!({"protocolVersion": "2025-03-26", "capabilities": {}, "clientInfo": {"name": "test", "version": "1.0"}}),
        );
        let result = resp.result.unwrap();
        assert_eq!(result["protocolVersion"], "2025-03-26");
        assert!(result["capabilities"]["tools"].is_object());
        assert_eq!(result["serverInfo"]["name"], "huskies");
    }

    #[test]
    fn to_sse_response_wraps_in_data_prefix() {
        let resp = JsonRpcResponse::success(Some(json!(1)), json!({"ok": true}));
        let http_resp = to_sse_response(resp);
        assert_eq!(
            http_resp.headers().get("content-type").unwrap(),
            "text/event-stream"
        );
    }

    #[test]
    fn wants_sse_detects_accept_header() {
        // Can't easily construct a Request in tests without TestClient,
        // so test the logic indirectly via to_sse_response format
        let resp = JsonRpcResponse::success(Some(json!(1)), json!("ok"));
        let json_resp = to_json_response(resp);
        assert_eq!(
            json_resp.headers().get("content-type").unwrap(),
            "application/json"
        );
    }

    #[test]
    fn json_rpc_error_response_builds_json_response() {
        let resp = json_rpc_error_response(Some(json!(42)), -32600, "test error".into());
        assert_eq!(resp.status(), poem::http::StatusCode::OK);
        assert_eq!(
            resp.headers().get("content-type").unwrap(),
            "application/json"
        );
    }

    // ── HTTP handler tests (TestClient) ───────────────────────────

    fn test_mcp_app(ctx: std::sync::Arc<AppContext>) -> impl poem::Endpoint {
        use poem::EndpointExt;
        poem::Route::new()
            .at("/mcp", poem::post(mcp_post_handler).get(mcp_get_handler))
            .data(ctx)
    }

    async fn read_body_json(resp: poem::test::TestResponse) -> Value {
        let body = resp.0.into_body().into_string().await.unwrap();
        serde_json::from_str(&body).unwrap()
    }

    async fn post_json_mcp<E: poem::Endpoint>(
        cli: &poem::test::TestClient<E>,
        payload: &str,
    ) -> Value {
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .body(payload.to_string())
            .send()
            .await;
        read_body_json(resp).await
    }

    #[tokio::test]
    async fn mcp_get_handler_returns_405() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli.get("/mcp").send().await;
        assert_eq!(resp.0.status(), poem::http::StatusCode::METHOD_NOT_ALLOWED);
    }

    #[tokio::test]
    async fn mcp_post_invalid_content_type_returns_error() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "text/plain")
            .body("{}")
            .send()
            .await;
        let body = read_body_json(resp).await;
        assert!(body.get("error").is_some(), "expected error field: {body}");
    }

    #[tokio::test]
    async fn mcp_post_invalid_json_returns_parse_error() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .body("not-valid-json")
            .send()
            .await;
        let body = read_body_json(resp).await;
        assert!(body.get("error").is_some(), "expected error field: {body}");
    }

    #[tokio::test]
    async fn mcp_post_wrong_jsonrpc_version_returns_error() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let body = post_json_mcp(
            &cli,
            r#"{"jsonrpc":"1.0","id":1,"method":"initialize","params":{}}"#,
        )
        .await;
        assert!(
            body["error"]["message"]
                .as_str()
                .unwrap_or("")
                .contains("version"),
            "expected version error: {body}"
        );
    }

    #[tokio::test]
    async fn mcp_post_notification_with_null_id_returns_accepted() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized","params":{}}"#)
            .send()
            .await;
        assert_eq!(resp.0.status(), poem::http::StatusCode::ACCEPTED);
    }

    #[tokio::test]
    async fn mcp_post_notification_with_explicit_null_id_returns_accepted() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .body(r#"{"jsonrpc":"2.0","id":null,"method":"notifications/initialized","params":{}}"#)
            .send()
            .await;
        assert_eq!(resp.0.status(), poem::http::StatusCode::ACCEPTED);
    }

    #[tokio::test]
    async fn mcp_post_missing_id_non_notification_returns_error() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let body = post_json_mcp(
            &cli,
            r#"{"jsonrpc":"2.0","method":"initialize","params":{}}"#,
        )
        .await;
        assert!(body.get("error").is_some(), "expected error: {body}");
    }

    #[tokio::test]
    async fn mcp_post_unknown_method_returns_error() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let body = post_json_mcp(
            &cli,
            r#"{"jsonrpc":"2.0","id":1,"method":"bogus/method","params":{}}"#,
        )
        .await;
        assert!(
            body["error"]["message"]
                .as_str()
                .unwrap_or("")
                .contains("Unknown method"),
            "expected unknown method error: {body}"
        );
    }

    #[tokio::test]
    async fn mcp_post_initialize_returns_capabilities() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let body = post_json_mcp(
            &cli,
            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#,
        )
        .await;
        assert_eq!(body["result"]["protocolVersion"], "2025-03-26");
        assert_eq!(body["result"]["serverInfo"]["name"], "huskies");
    }

    #[tokio::test]
    async fn mcp_post_tools_list_returns_tools() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let body = post_json_mcp(
            &cli,
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#,
        )
        .await;
        assert!(body["result"]["tools"].is_array());
    }

    #[tokio::test]
    async fn mcp_post_sse_returns_event_stream_content_type() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#)
            .send()
            .await;
        assert_eq!(
            resp.0.headers().get("content-type").unwrap(),
            "text/event-stream"
        );
    }

    #[tokio::test]
    async fn mcp_post_sse_get_agent_output_missing_story_id() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get_agent_output","arguments":{}}}"#)
            .send()
            .await;
        assert_eq!(
            resp.0.headers().get("content-type").unwrap(),
            "text/event-stream",
            "expected SSE content-type"
        );
    }

    #[tokio::test]
    async fn mcp_post_sse_get_agent_output_without_agent_name_returns_disk_content() {
        // Without agent_name the SSE live-streaming intercept is skipped and
        // the disk-based handler runs. The transport still wraps the result in
        // SSE format (data: …\n\n) because the client sent Accept: text/event-stream,
        // but the content should be a valid JSON-RPC result, not a subscribe error.
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get_agent_output","arguments":{"story_id":"1_test"}}}"#)
            .send()
            .await;
        let body = resp.0.into_body().into_string().await.unwrap();
        // Body is SSE-wrapped: "data: {…}\n\n" — strip the prefix and verify it's
        // a valid JSON-RPC result (not an error about missing agent_name).
        let json_part = body
            .trim_start_matches("data: ")
            .trim_end_matches("\n\n")
            .trim();
        let parsed: serde_json::Value = serde_json::from_str(json_part)
            .unwrap_or_else(|_| panic!("expected JSON-RPC in SSE body, got: {body}"));
        assert!(
            parsed.get("result").is_some(),
            "expected JSON-RPC result (disk-based handler ran): {parsed}"
        );
        // Must NOT be an error about missing agent_name (agent_name is now optional)
        assert!(
            parsed.get("error").is_none(),
            "unexpected error when agent_name omitted: {parsed}"
        );
    }

    #[tokio::test]
    async fn mcp_post_sse_get_agent_output_no_agent_no_logs_returns_not_found() {
        // Agent not in pool and no log files → SSE success with "No log files found" message.
        let tmp = tempfile::tempdir().unwrap();
        let ctx = std::sync::Arc::new(test_ctx(tmp.path()));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get_agent_output","arguments":{"story_id":"99_nope","agent_name":"bot"}}}"#)
            .send()
            .await;
        assert_eq!(
            resp.0.headers().get("content-type").unwrap(),
            "text/event-stream"
        );
        let body = resp.0.into_body().into_string().await.unwrap();
        assert!(body.contains("data:"), "expected SSE data prefix: {body}");
        // Must NOT return isError — should be a success result with "No log files found"
        assert!(
            !body.contains("isError"),
            "expected no isError for missing agent: {body}"
        );
        assert!(
            body.contains("No log files found"),
            "expected not-found message: {body}"
        );
    }

    #[tokio::test]
    async fn mcp_post_sse_get_agent_output_exited_agent_reads_disk_logs() {
        use crate::agent_log::AgentLogWriter;
        use crate::agents::AgentEvent;
        // Agent has exited (not in pool) but wrote logs to disk.
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        let mut writer = AgentLogWriter::new(root, "42_story_foo", "coder-1", "sess-sse").unwrap();
        writer
            .write_event(&AgentEvent::Output {
                story_id: "42_story_foo".to_string(),
                agent_name: "coder-1".to_string(),
                text: "disk output".to_string(),
            })
            .unwrap();
        drop(writer);

        let ctx = std::sync::Arc::new(test_ctx(root));
        let cli = poem::test::TestClient::new(test_mcp_app(ctx));
        let resp = cli
            .post("/mcp")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get_agent_output","arguments":{"story_id":"42_story_foo","agent_name":"coder-1"}}}"#)
            .send()
            .await;
        let body = resp.0.into_body().into_string().await.unwrap();
        assert!(
            body.contains("disk output"),
            "expected disk log content in SSE response: {body}"
        );
        assert!(
            !body.contains("isError"),
            "expected no error for exited agent with logs: {body}"
        );
    }
}