fastmcp-rs 0.2.0

use std::sync::Arc;

use serde_json::Value;
use tokio::io::{self as aio, AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tracing::{error, info};

use crate::command::{ServerCommand, execute_command};
use crate::error::{FastMcpError, Result};
use crate::server::FastMcpServer;

/// Runs the server over STDIO using a simple JSON command protocol.
pub async fn run_stdio(server: Arc<FastMcpServer>) -> Result<()> {
    let stdin = aio::stdin();
    let stdout = aio::stdout();

    run_stdio_with_io(server, BufReader::new(stdin), stdout).await
}

pub async fn run_stdio_with_io<R, W>(
    server: Arc<FastMcpServer>,
    reader: R,
    mut writer: W,
) -> Result<()>
where
    R: AsyncBufRead + Unpin,
    W: AsyncWrite + Unpin,
{
    log_startup(&server, "stdio", None);

    // Use fully-qualified trait call to avoid generic method resolution issues during packaging
    let mut lines = tokio::io::AsyncBufReadExt::lines(reader);

    loop {
        let line = match lines.next_line().await {
            Ok(Some(line)) => line,
            Ok(None) => break,
            Err(err) => {
                error!("Failed reading stdin: {err}");
                break;
            }
        };

        if line.trim().is_empty() {
            continue;
        }

        let command: ServerCommand = match serde_json::from_str(&line) {
            Ok(cmd) => cmd,
            Err(err) => {
                let error = FastMcpError::InvalidInvocation(format!("invalid command: {err}"));
                write_response(
                    &mut writer,
                    &serde_json::json!({"type": "error", "message": error.to_string()}),
                )
                .await?;
                continue;
            }
        };

        let (result, shutdown) = match execute_command(&server, command).await {
            Ok(res) => res,
            Err(err) => {
                write_response(
                    &mut writer,
                    &serde_json::json!({"type": "error", "message": err.to_string()}),
                )
                .await?;
                continue;
            }
        };

        let payload = serde_json::to_value(&result)?;
        write_response(&mut writer, &payload).await?;

        if shutdown {
            break;
        }
    }

    Ok(())
}

async fn write_response<W>(writer: &mut W, payload: &Value) -> Result<()>
where
    W: AsyncWrite + Unpin,
{
    let mut json = serde_json::to_vec(payload)?;
    json.push(b'\n');
    tokio::io::AsyncWriteExt::write_all(&mut *writer, &json).await?;
    tokio::io::AsyncWriteExt::flush(&mut *writer).await?;
    Ok(())
}

fn log_startup(server: &FastMcpServer, transport: &str, url: Option<&str>) {
    let metadata = server.metadata();
    let tools = server
        .list_tools()
        .into_iter()
        .map(|tool| tool.name)
        .collect::<Vec<_>>();
    let resources = server
        .list_resources()
        .into_iter()
        .map(|resource| resource.uri)
        .collect::<Vec<_>>();
    let prompts = server
        .list_prompts()
        .into_iter()
        .map(|prompt| prompt.name)
        .collect::<Vec<_>>();

    let mut lines = Vec::new();
    lines.push(format!(
        "FastMCP '{}' (id: {}) using {transport} transport{}",
        metadata.name,
        metadata.id,
        url.map(|u| format!(" at {u}")).unwrap_or_default()
    ));
    lines.push("  Transport URI: mcp+stdio://localhost".to_string());
    lines.push("  Transport note: STDIO runs over stdin/stdout (no network socket)".to_string());
    lines.push(format!(
        "  Instructions: {}",
        metadata
            .instructions
            .as_deref()
            .unwrap_or("No instructions configured")
    ));
    lines.push(format!(
        "  Registered tools: {}",
        if tools.is_empty() {
            "none".into()
        } else {
            tools.join(", ")
        }
    ));
    lines.push(format!(
        "  Registered resources: {}",
        if resources.is_empty() {
            "none".into()
        } else {
            resources.join(", ")
        }
    ));
    lines.push(format!(
        "  Registered prompts: {}",
        if prompts.is_empty() {
            "none".into()
        } else {
            prompts.join(", ")
        }
    ));

    emit_startup_lines(lines);
}

fn emit_startup_lines(lines: Vec<String>) {
    for line in lines {
        info!("{}", line);
        println!("{line}");
    }
}