bashkit-cli 0.1.21

//! MCP (Model Context Protocol) server implementation
//!
//! Implements a JSON-RPC 2.0 server that exposes bashkit as an MCP tool.
//! Optionally registers ScriptedTool instances as additional MCP tools.
//!
//! Protocol:
//! - Input: JSON-RPC requests on stdin
//! - Output: JSON-RPC responses on stdout

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::io::{BufRead, Read, Write};

/// Hard cap for one JSON-RPC request line (including trailing '\n').
// THREAT[TM-DOS-037]: Bound request buffering to prevent stdin memory DoS.
const MAX_REQUEST_LINE_BYTES: usize = 1024 * 1024; // 1 MiB

/// JSON-RPC 2.0 request
#[derive(Debug, Deserialize)]
struct JsonRpcRequest {
    #[allow(dead_code)]
    jsonrpc: String,
    id: serde_json::Value,
    method: String,
    #[serde(default)]
    params: serde_json::Value,
}

/// JSON-RPC 2.0 response
#[derive(Debug, Serialize)]
struct JsonRpcResponse {
    jsonrpc: String,
    id: serde_json::Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<JsonRpcError>,
}

/// JSON-RPC 2.0 error
#[derive(Debug, Serialize)]
struct JsonRpcError {
    code: i32,
    message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    data: Option<serde_json::Value>,
}

impl JsonRpcResponse {
    fn success(id: serde_json::Value, result: serde_json::Value) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: Some(result),
            error: None,
        }
    }

    fn error(id: serde_json::Value, code: i32, message: String) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: None,
            error: Some(JsonRpcError {
                code,
                message,
                data: None,
            }),
        }
    }
}

/// MCP tool definition
#[derive(Debug, Serialize)]
struct McpTool {
    name: String,
    description: String,
    #[serde(rename = "inputSchema")]
    input_schema: serde_json::Value,
}

/// MCP server capabilities
#[derive(Debug, Serialize)]
struct ServerCapabilities {
    tools: serde_json::Value,
}

/// MCP server info
#[derive(Debug, Serialize)]
struct ServerInfo {
    name: String,
    version: String,
}

/// MCP initialize result
#[derive(Debug, Serialize)]
struct InitializeResult {
    #[serde(rename = "protocolVersion")]
    protocol_version: String,
    capabilities: ServerCapabilities,
    #[serde(rename = "serverInfo")]
    server_info: ServerInfo,
}

/// Tool call arguments for bash execution
#[derive(Debug, Deserialize)]
struct BashToolArgs {
    script: String,
}

/// Tool call result
#[derive(Debug, Serialize)]
struct ToolResult {
    content: Vec<ContentItem>,
    #[serde(rename = "isError", skip_serializing_if = "Option::is_none")]
    is_error: Option<bool>,
}

#[derive(Debug, Serialize)]
struct ContentItem {
    #[serde(rename = "type")]
    content_type: String,
    text: String,
}

/// MCP server with optional ScriptedTool registrations.
///
/// Accepts a factory function that produces configured `Bash` instances,
/// ensuring CLI execution limits (max_commands, etc.) are applied to every
/// MCP `tools/call` invocation.
///
/// DESIGN: Session-level counters (session_commands, session_exec_calls) are
/// tracked cumulatively here because each tools/call creates a fresh Bash
/// instance. Without this, an LLM agent could make unlimited sequential calls
/// without session limits ever triggering. See issue #1193.
pub struct McpServer {
    bash_factory: Box<dyn Fn() -> bashkit::Bash + Send>,
    /// Cumulative session command count across all MCP tool calls.
    cumulative_commands: u64,
    /// Cumulative session exec call count across all MCP tool calls.
    cumulative_exec_calls: u64,
    /// Rate limiter: max tool-call requests per minute (0 = unlimited).
    // THREAT[TM-DOS-036]: Prevent resource exhaustion from rapid MCP requests.
    max_requests_per_minute: u32,
    /// Timestamps of recent tool-call requests for rate limiting.
    request_timestamps: Vec<std::time::Instant>,
    #[cfg(feature = "scripted_tool")]
    scripted_tools: Vec<bashkit::ScriptedTool>,
}

impl McpServer {
    /// Create a new MCP server with only the default `bash` tool.
    /// Each `tools/call` will create a `Bash` via the provided factory,
    /// inheriting whatever limits/configuration the caller sets up.
    pub fn new(bash_factory: impl Fn() -> bashkit::Bash + Send + 'static) -> Self {
        Self {
            bash_factory: Box::new(bash_factory),
            cumulative_commands: 0,
            cumulative_exec_calls: 0,
            max_requests_per_minute: 0,
            request_timestamps: Vec::new(),
            #[cfg(feature = "scripted_tool")]
            scripted_tools: Vec::new(),
        }
    }

    /// Set the maximum number of tool-call requests per minute.
    /// 0 means unlimited.
    pub fn max_requests_per_minute(mut self, limit: u32) -> Self {
        self.max_requests_per_minute = limit;
        self
    }

    /// Register a ScriptedTool. It will appear in `tools/list` and route
    /// `tools/call` to `ScriptedTool::execute()`.
    #[cfg(feature = "scripted_tool")]
    #[allow(dead_code)] // Public API for external consumers; used in tests
    pub fn register_scripted_tool(&mut self, tool: bashkit::ScriptedTool) {
        self.scripted_tools.push(tool);
    }

    /// Run the server, reading JSON-RPC from stdin and writing responses to stdout.
    pub async fn run(&mut self) -> Result<()> {
        let stdin = std::io::stdin();
        let mut stdin_lock = stdin.lock();
        let mut stdout = std::io::stdout();
        let mut line_buf = Vec::new();

        loop {
            let read = read_bounded_request_line(&mut stdin_lock, &mut line_buf)
                .context("Failed to read line from stdin")?;
            let Some(read) = read else {
                break;
            };
            if read == RequestLineRead::TooLarge {
                let response = JsonRpcResponse::error(
                    serde_json::Value::Null,
                    -32600,
                    "Invalid request: request line too large".to_string(),
                );
                writeln!(stdout, "{}", serde_json::to_string(&response)?)?;
                stdout.flush()?;
                continue;
            }

            let line = String::from_utf8_lossy(&line_buf);
            if line.trim().is_empty() {
                continue;
            }

            let request: JsonRpcRequest = match serde_json::from_str(&line) {
                Ok(req) => req,
                Err(e) => {
                    eprintln!("MCP parse error: {}", e);
                    let response = JsonRpcResponse::error(
                        serde_json::Value::Null,
                        -32700,
                        "Parse error: invalid JSON".to_string(),
                    );
                    writeln!(stdout, "{}", serde_json::to_string(&response)?)?;
                    stdout.flush()?;
                    continue;
                }
            };

            let response = self.handle_request(request).await;
            writeln!(stdout, "{}", serde_json::to_string(&response)?)?;
            stdout.flush()?;
        }

        Ok(())
    }

    async fn handle_request(&mut self, request: JsonRpcRequest) -> JsonRpcResponse {
        // THREAT[TM-DOS-036]: Rate-limit tools/call to prevent resource exhaustion.
        if request.method == "tools/call"
            && let Some(err) = self.check_rate_limit()
        {
            return JsonRpcResponse::error(request.id, -32000, err);
        }

        match request.method.as_str() {
            "initialize" => Self::handle_initialize(request.id),
            "initialized" => JsonRpcResponse::success(request.id, serde_json::Value::Null),
            "tools/list" => self.handle_tools_list(request.id),
            "tools/call" => self.handle_tools_call(request.id, request.params).await,
            "shutdown" => JsonRpcResponse::success(request.id, serde_json::Value::Null),
            _ => JsonRpcResponse::error(request.id, -32601, "Method not found".to_string()),
        }
    }

    /// Check rate limit. Returns an error message if exceeded, None if OK.
    fn check_rate_limit(&mut self) -> Option<String> {
        if self.max_requests_per_minute == 0 {
            return None;
        }
        let now = std::time::Instant::now();
        let window = std::time::Duration::from_secs(60);
        // Remove timestamps older than 1 minute
        self.request_timestamps
            .retain(|t| now.duration_since(*t) < window);
        if self.request_timestamps.len() >= self.max_requests_per_minute as usize {
            return Some(format!(
                "Rate limit exceeded: max {} requests per minute",
                self.max_requests_per_minute
            ));
        }
        self.request_timestamps.push(now);
        None
    }

    fn handle_initialize(id: serde_json::Value) -> JsonRpcResponse {
        let result = InitializeResult {
            protocol_version: "2024-11-05".to_string(),
            capabilities: ServerCapabilities {
                tools: serde_json::json!({}),
            },
            server_info: ServerInfo {
                name: "bashkit".to_string(),
                version: env!("CARGO_PKG_VERSION").to_string(),
            },
        };

        JsonRpcResponse::success(id, serde_json::to_value(result).expect("serialize init"))
    }

    fn handle_tools_list(&self, id: serde_json::Value) -> JsonRpcResponse {
        #[allow(unused_mut)]
        let mut tools = vec![McpTool {
            name: "bash".to_string(),
            description: "Execute a bash script in a virtual environment".to_string(),
            input_schema: serde_json::json!({
                "type": "object",
                "properties": {
                    "script": {
                        "type": "string",
                        "description": "The bash script to execute"
                    }
                },
                "required": ["script"]
            }),
        }];

        #[cfg(feature = "scripted_tool")]
        {
            use bashkit::tool::Tool;
            for st in &self.scripted_tools {
                tools.push(McpTool {
                    name: st.name().to_string(),
                    description: st.short_description().to_string(),
                    input_schema: serde_json::json!({
                        "type": "object",
                        "properties": {
                            "commands": {
                                "type": "string",
                                "description": st.description()
                            }
                        },
                        "required": ["commands"]
                    }),
                });
            }
        }

        JsonRpcResponse::success(id, serde_json::json!({ "tools": tools }))
    }

    async fn handle_tools_call(
        &mut self,
        id: serde_json::Value,
        params: serde_json::Value,
    ) -> JsonRpcResponse {
        let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
        let arguments = params.get("arguments").cloned().unwrap_or_default();

        #[cfg(feature = "scripted_tool")]
        {
            if let Some(st) = self.scripted_tools.iter_mut().find(|t| {
                use bashkit::tool::Tool;
                t.name() == tool_name
            }) {
                return Self::handle_scripted_tool_call(id, st, arguments).await;
            }
        }

        if tool_name != "bash" {
            return JsonRpcResponse::error(id, -32602, format!("Unknown tool: {}", tool_name));
        }

        let args: BashToolArgs = match serde_json::from_value(arguments) {
            Ok(a) => a,
            Err(e) => {
                eprintln!("MCP invalid arguments: {}", e);
                return JsonRpcResponse::error(id, -32602, "Invalid arguments".to_string());
            }
        };

        let mut bash = (self.bash_factory)();
        // Restore cumulative session counters so limits persist across MCP calls
        bash.restore_session_counters(self.cumulative_commands, self.cumulative_exec_calls);
        let result = match bash.exec(&args.script).await {
            Ok(r) => r,
            Err(e) => {
                // Update cumulative counters even on error (commands were still counted)
                let (cmds, execs) = bash.session_counters();
                self.cumulative_commands = cmds;
                self.cumulative_exec_calls = execs;
                let tool_result = ToolResult {
                    content: vec![ContentItem {
                        content_type: "text".to_string(),
                        text: format!("Error: {}", e),
                    }],
                    is_error: Some(true),
                };
                return JsonRpcResponse::success(
                    id,
                    serde_json::to_value(tool_result).expect("serialize"),
                );
            }
        };
        // Update cumulative counters after successful execution
        let (cmds, execs) = bash.session_counters();
        self.cumulative_commands = cmds;
        self.cumulative_exec_calls = execs;

        let mut output = result.stdout;
        if !result.stderr.is_empty() {
            output.push_str("\n[stderr]\n");
            output.push_str(&result.stderr);
        }
        if result.exit_code != 0 {
            output.push_str(&format!("\n[exit code: {}]", result.exit_code));
        }

        let tool_result = ToolResult {
            content: vec![ContentItem {
                content_type: "text".to_string(),
                text: output,
            }],
            is_error: if result.exit_code != 0 {
                Some(true)
            } else {
                None
            },
        };

        JsonRpcResponse::success(id, serde_json::to_value(tool_result).expect("serialize"))
    }

    #[cfg(feature = "scripted_tool")]
    async fn handle_scripted_tool_call(
        id: serde_json::Value,
        tool: &mut bashkit::ScriptedTool,
        arguments: serde_json::Value,
    ) -> JsonRpcResponse {
        use bashkit::tool::{Tool, ToolRequest};

        let commands = arguments
            .get("commands")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        let resp = tool
            .execute(ToolRequest {
                commands: commands.to_string(),
                timeout_ms: None,
            })
            .await;

        let mut output = resp.stdout;
        if !resp.stderr.is_empty() {
            output.push_str("\n[stderr]\n");
            output.push_str(&resp.stderr);
        }
        if resp.exit_code != 0 {
            output.push_str(&format!("\n[exit code: {}]", resp.exit_code));
        }

        let tool_result = ToolResult {
            content: vec![ContentItem {
                content_type: "text".to_string(),
                text: output,
            }],
            is_error: if resp.exit_code != 0 {
                Some(true)
            } else {
                None
            },
        };

        JsonRpcResponse::success(id, serde_json::to_value(tool_result).expect("serialize"))
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum RequestLineRead {
    Complete,
    TooLarge,
}

fn read_bounded_request_line<R: BufRead>(
    reader: &mut R,
    line_buf: &mut Vec<u8>,
) -> Result<Option<RequestLineRead>> {
    line_buf.clear();
    let bytes_read = reader
        .by_ref()
        .take((MAX_REQUEST_LINE_BYTES + 1) as u64)
        .read_until(b'\n', line_buf)?;
    if bytes_read == 0 {
        return Ok(None);
    }
    if bytes_read > MAX_REQUEST_LINE_BYTES {
        if !line_buf.ends_with(b"\n") {
            discard_until_newline(reader)?;
        }
        return Ok(Some(RequestLineRead::TooLarge));
    }
    Ok(Some(RequestLineRead::Complete))
}

fn discard_until_newline<R: BufRead>(reader: &mut R) -> Result<()> {
    loop {
        let available = reader.fill_buf()?;
        if available.is_empty() {
            return Ok(());
        }
        if let Some(pos) = available.iter().position(|&b| b == b'\n') {
            reader.consume(pos + 1);
            return Ok(());
        }
        let len = available.len();
        reader.consume(len);
    }
}

/// Run the MCP server with a factory that produces configured `Bash` instances.
#[allow(dead_code)] // Public API; used by external callers / tests
pub async fn run(bash_factory: impl Fn() -> bashkit::Bash + Send + 'static) -> Result<()> {
    let mut server = McpServer::new(bash_factory);
    server.run().await
}

/// Run the MCP server with rate limiting.
pub async fn run_with_rate_limit(
    bash_factory: impl Fn() -> bashkit::Bash + Send + 'static,
    max_requests_per_minute: u32,
) -> Result<()> {
    let mut server = McpServer::new(bash_factory).max_requests_per_minute(max_requests_per_minute);
    server.run().await
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[tokio::test]
    async fn test_initialize() {
        let mut server = McpServer::new(bashkit::Bash::new);
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "initialize".to_string(),
            params: serde_json::json!({}),
        };
        let resp = server.handle_request(req).await;
        let result = resp.result.expect("should have result");
        assert_eq!(result["protocolVersion"], "2024-11-05");
        assert_eq!(result["serverInfo"]["name"], "bashkit");
    }

    #[tokio::test]
    async fn test_tools_list_default() {
        let mut server = McpServer::new(bashkit::Bash::new);
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/list".to_string(),
            params: serde_json::json!({}),
        };
        let resp = server.handle_request(req).await;
        let result = resp.result.expect("should have result");
        let tools = result["tools"].as_array().expect("tools array");
        assert!(tools.iter().any(|t| t["name"] == "bash"));
    }

    #[tokio::test]
    async fn test_tools_call_bash() {
        let mut server = McpServer::new(bashkit::Bash::new);
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo hello" }
            }),
        };
        let resp = server.handle_request(req).await;
        let result = resp.result.expect("should have result");
        let text = result["content"][0]["text"].as_str().expect("text");
        assert!(text.contains("hello"));
    }

    #[tokio::test]
    async fn test_tools_call_unknown() {
        let mut server = McpServer::new(bashkit::Bash::new);
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "nonexistent",
                "arguments": {}
            }),
        };
        let resp = server.handle_request(req).await;
        assert!(resp.error.is_some());
    }

    #[tokio::test]
    async fn test_method_not_found() {
        let mut server = McpServer::new(bashkit::Bash::new);
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "unknown/method".to_string(),
            params: serde_json::json!({}),
        };
        let resp = server.handle_request(req).await;
        assert!(resp.error.is_some());
        assert_eq!(resp.error.expect("error").code, -32601);
    }

    #[tokio::test]
    async fn test_tools_call_respects_max_commands() {
        // Factory that creates a Bash with max_commands=2
        let mut server = McpServer::new(|| {
            bashkit::Bash::builder()
                .limits(bashkit::ExecutionLimits::new().max_commands(2))
                .build()
        });

        // Script with 3 commands should hit the limit
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo a; echo b; echo c" }
            }),
        };
        let resp = server.handle_request(req).await;
        let result = resp.result.expect("should have result");
        let text = result["content"][0]["text"].as_str().expect("text");
        // Should report the limit was exceeded
        assert!(
            text.contains("limit") || text.contains("exceeded") || result["isError"] == true,
            "expected execution limit error, got: {text}"
        );
    }

    #[tokio::test]
    async fn test_session_limits_accumulate_across_mcp_calls() {
        // Session limit: max 3 total commands across all calls, max 2 exec calls.
        let mut server = McpServer::new(|| {
            bashkit::Bash::builder()
                .session_limits(
                    bashkit::SessionLimits::new()
                        .max_total_commands(3)
                        .max_exec_calls(2),
                )
                .build()
        });

        // First call: 2 commands. Should succeed.
        let req1 = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo a; echo b" }
            }),
        };
        let resp1 = server.handle_request(req1).await;
        let result1 = resp1.result.expect("should have result");
        let text1 = result1["content"][0]["text"].as_str().expect("text");
        assert!(
            text1.contains('a') && text1.contains('b'),
            "first call should succeed, got: {text1}"
        );

        // Second call: 2 more commands -> cumulative 4 > limit of 3. Should fail.
        let req2 = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(2),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo c; echo d" }
            }),
        };
        let resp2 = server.handle_request(req2).await;
        let result2 = resp2.result.expect("should have result");
        let text2 = result2["content"][0]["text"].as_str().expect("text");
        assert!(
            text2.contains("session") || text2.contains("limit") || result2["isError"] == true,
            "second call should hit session limit, got: {text2}"
        );
    }

    #[tokio::test]
    async fn test_session_exec_calls_accumulate_across_mcp_calls() {
        // Session limit: max 2 exec calls.
        let mut server = McpServer::new(|| {
            bashkit::Bash::builder()
                .session_limits(bashkit::SessionLimits::new().max_exec_calls(2))
                .build()
        });

        // First two calls should succeed.
        for i in 1..=2 {
            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(i),
                method: "tools/call".to_string(),
                params: serde_json::json!({
                    "name": "bash",
                    "arguments": { "script": "echo ok" }
                }),
            };
            let resp = server.handle_request(req).await;
            let result = resp.result.expect("should have result");
            let text = result["content"][0]["text"].as_str().expect("text");
            assert!(text.contains("ok"), "call {i} should succeed, got: {text}");
        }

        // Third call should hit session exec call limit.
        let req3 = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(3),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo should_fail" }
            }),
        };
        let resp3 = server.handle_request(req3).await;
        let result3 = resp3.result.expect("should have result");
        let text3 = result3["content"][0]["text"].as_str().expect("text");
        assert!(
            text3.contains("session") || text3.contains("limit") || result3["isError"] == true,
            "third call should hit session exec call limit, got: {text3}"
        );
    }

    #[tokio::test]
    async fn test_parse_error_does_not_leak_internal_details() {
        // Simulate malformed JSON parse and verify the response that run() would
        // produce doesn't contain serde/internal details.
        let malformed = "not valid json {{{";
        let err = serde_json::from_str::<JsonRpcRequest>(malformed).unwrap_err();
        // This is what we'd send in run() — verify generic message
        let response = JsonRpcResponse::error(
            serde_json::Value::Null,
            -32700,
            "Parse error: invalid JSON".to_string(),
        );
        let serialized = serde_json::to_string(&response).unwrap();
        // Must NOT contain serde error details like "expected value" or struct names
        let err_str = err.to_string();
        assert!(
            !serialized.contains(&err_str),
            "response should not contain raw serde error: {err_str}"
        );
        assert!(!serialized.contains("expected"));
        assert!(!serialized.contains("line "));
        assert!(!serialized.contains("column "));
        // Must contain the generic message
        assert!(serialized.contains("invalid JSON"));
    }

    #[tokio::test]
    async fn test_invalid_arguments_does_not_leak_struct_names() {
        let mut server = McpServer::new(bashkit::Bash::new);
        // Send a tools/call with arguments missing the required "script" field
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "wrong_field": 123 }
            }),
        };
        let resp = server.handle_request(req).await;
        let err = resp.error.expect("should be an error");
        assert_eq!(err.code, -32602);
        assert_eq!(err.message, "Invalid arguments");
        // Must NOT leak field names, struct names, or serde details
        assert!(!err.message.contains("script"));
        assert!(!err.message.contains("missing field"));
        assert!(!err.message.contains("BashToolArgs"));
    }

    #[cfg(feature = "scripted_tool")]
    mod scripted_tool_tests {
        use super::*;
        use bashkit::{ScriptedTool, ToolArgs, ToolDef};

        fn make_test_tool() -> ScriptedTool {
            ScriptedTool::builder("test_api")
                .short_description("Test API tool")
                .tool_fn(ToolDef::new("greet", "Greet someone"), |args: &ToolArgs| {
                    let name = args.param_str("name").unwrap_or("world");
                    Ok(format!("hello {name}\n"))
                })
                .build()
        }

        #[tokio::test]
        async fn test_tools_list_includes_scripted_tool() {
            let mut server = McpServer::new(bashkit::Bash::new);
            server.register_scripted_tool(make_test_tool());

            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(1),
                method: "tools/list".to_string(),
                params: serde_json::json!({}),
            };
            let resp = server.handle_request(req).await;
            let result = resp.result.expect("should have result");
            let tools = result["tools"].as_array().expect("tools array");
            assert!(tools.iter().any(|t| t["name"] == "bash"));
            assert!(tools.iter().any(|t| t["name"] == "test_api"));
        }

        #[tokio::test]
        async fn test_tools_call_scripted_tool() {
            let mut server = McpServer::new(bashkit::Bash::new);
            server.register_scripted_tool(make_test_tool());

            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(1),
                method: "tools/call".to_string(),
                params: serde_json::json!({
                    "name": "test_api",
                    "arguments": { "commands": "greet --name Alice" }
                }),
            };
            let resp = server.handle_request(req).await;
            let result = resp.result.expect("should have result");
            let text = result["content"][0]["text"].as_str().expect("text");
            assert!(text.contains("hello Alice"));
        }

        #[tokio::test]
        async fn test_tools_call_scripted_tool_error() {
            let mut server = McpServer::new(bashkit::Bash::new);
            let tool = ScriptedTool::builder("err_api")
                .short_description("Error API")
                .tool_fn(ToolDef::new("fail", "Always fails"), |_args: &ToolArgs| {
                    Err("service down".to_string())
                })
                .build();
            server.register_scripted_tool(tool);

            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(1),
                method: "tools/call".to_string(),
                params: serde_json::json!({
                    "name": "err_api",
                    "arguments": { "commands": "fail" }
                }),
            };
            let resp = server.handle_request(req).await;
            let result = resp.result.expect("should have result");
            assert_eq!(result["isError"], true);
        }

        #[tokio::test]
        async fn test_full_jsonrpc_roundtrip() {
            let mut server = McpServer::new(bashkit::Bash::new);
            server.register_scripted_tool(make_test_tool());

            // Step 1: initialize
            let init_resp = server
                .handle_request(JsonRpcRequest {
                    jsonrpc: "2.0".to_string(),
                    id: serde_json::json!(1),
                    method: "initialize".to_string(),
                    params: serde_json::json!({}),
                })
                .await;
            assert!(init_resp.result.is_some());

            // Step 2: tools/list
            let list_resp = server
                .handle_request(JsonRpcRequest {
                    jsonrpc: "2.0".to_string(),
                    id: serde_json::json!(2),
                    method: "tools/list".to_string(),
                    params: serde_json::json!({}),
                })
                .await;
            let list_result = list_resp.result.expect("result");
            let tools = list_result["tools"].as_array().expect("tools");
            assert!(tools.len() >= 2);

            // Step 3: tools/call
            let call_resp = server
                .handle_request(JsonRpcRequest {
                    jsonrpc: "2.0".to_string(),
                    id: serde_json::json!(3),
                    method: "tools/call".to_string(),
                    params: serde_json::json!({
                        "name": "test_api",
                        "arguments": { "commands": "greet --name MCP" }
                    }),
                })
                .await;
            let call_result = call_resp.result.expect("result");
            let text = call_result["content"][0]["text"].as_str().expect("text");
            assert!(text.contains("hello MCP"));
        }
    }

    // THREAT[TM-DOS-036]: Rate limiting tests

    #[tokio::test]
    async fn test_rate_limit_blocks_excess_requests() {
        let mut server = McpServer::new(bashkit::Bash::new).max_requests_per_minute(2);

        // First two requests should succeed
        for i in 1..=2 {
            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(i),
                method: "tools/call".to_string(),
                params: serde_json::json!({
                    "name": "bash",
                    "arguments": { "script": "echo ok" }
                }),
            };
            let resp = server.handle_request(req).await;
            assert!(resp.error.is_none(), "request {i} should succeed");
        }

        // Third request should be rate-limited
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(3),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo should_fail" }
            }),
        };
        let resp = server.handle_request(req).await;
        let err = resp.error.expect("should be rate limited");
        assert_eq!(err.code, -32000);
        assert!(err.message.contains("Rate limit exceeded"));
    }

    #[tokio::test]
    async fn test_rate_limit_zero_means_unlimited() {
        let mut server = McpServer::new(bashkit::Bash::new).max_requests_per_minute(0);

        for i in 1..=10 {
            let req = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: serde_json::json!(i),
                method: "tools/call".to_string(),
                params: serde_json::json!({
                    "name": "bash",
                    "arguments": { "script": "echo ok" }
                }),
            };
            let resp = server.handle_request(req).await;
            assert!(
                resp.error.is_none(),
                "request {i} should succeed with no limit"
            );
        }
    }

    #[tokio::test]
    async fn test_rate_limit_does_not_block_non_tool_calls() {
        let mut server = McpServer::new(bashkit::Bash::new).max_requests_per_minute(1);

        // Use the one allowed tool call
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            method: "tools/call".to_string(),
            params: serde_json::json!({
                "name": "bash",
                "arguments": { "script": "echo ok" }
            }),
        };
        server.handle_request(req).await;

        // tools/list should still work even though tool call limit is reached
        let req = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(2),
            method: "tools/list".to_string(),
            params: serde_json::json!({}),
        };
        let resp = server.handle_request(req).await;
        assert!(
            resp.error.is_none(),
            "tools/list should not be rate-limited"
        );
    }

    #[test]
    fn test_read_bounded_request_line_rejects_oversized_line() {
        let oversized = vec![b'a'; MAX_REQUEST_LINE_BYTES + 64];
        let mut input = Cursor::new(oversized);
        let mut line_buf = Vec::new();

        let read = read_bounded_request_line(&mut input, &mut line_buf).expect("read");
        assert_eq!(read, Some(RequestLineRead::TooLarge));
        assert!(line_buf.len() > MAX_REQUEST_LINE_BYTES);
    }

    #[test]
    fn test_read_bounded_request_line_discards_until_newline_after_oversized_line() {
        let mut payload = vec![b'a'; MAX_REQUEST_LINE_BYTES + 32];
        payload.push(b'\n');
        payload.extend_from_slice(br#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#);
        payload.push(b'\n');
        let mut input = Cursor::new(payload);
        let mut line_buf = Vec::new();

        let first = read_bounded_request_line(&mut input, &mut line_buf).expect("first read");
        assert_eq!(first, Some(RequestLineRead::TooLarge));

        let second = read_bounded_request_line(&mut input, &mut line_buf).expect("second read");
        assert_eq!(second, Some(RequestLineRead::Complete));
        assert_eq!(
            std::str::from_utf8(&line_buf).expect("utf8"),
            "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}\n"
        );
    }
}