keep/src/modes/server/api/mcp.rs

use axum::{
    extract::State,
    response::sse::{Event, KeepAlive, Sse},
    http::StatusCode,
};
use futures::stream::{self, Stream};
use log::{debug, error, info};
use std::convert::Infallible;
use std::time::Duration;
use tokio_stream::StreamExt as _;

use crate::modes::server::common::AppState;
use crate::modes::server::mcp::KeepMcpServer;
use rmcp::ServiceExt;

#[utoipa::path(
    get,
    path = "/mcp/sse",
    operation_id = "mcp_sse",
    summary = "Model Context Protocol SSE endpoint",
    description = "Server-Sent Events endpoint for Model Context Protocol communication. This endpoint allows AI tools and clients to interact with Keep's functionality through the standardized MCP protocol. Supports saving items, retrieving content, searching by tags and metadata, and listing stored items.",
    responses(
        (status = 200, description = "SSE stream established for MCP communication"),
        (status = 401, description = "Unauthorized - Invalid or missing authentication credentials"),
        (status = 500, description = "Internal server error - Failed to establish MCP connection")
    ),
    security(
        ("bearerAuth" = [])
    ),
    tag = "mcp"
)]
pub async fn handle_mcp_sse(
    State(state): State<AppState>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, StatusCode> {
    debug!("MCP: Starting SSE endpoint");
    
    let mcp_server = KeepMcpServer::new(state);
    
    // Create a simple message channel for SSE communication
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
    
    // Send initial connection message
    let _ = tx.send("data: {\"type\":\"connection\",\"status\":\"connected\"}\n\n".to_string());
    
    // For now, create a simple stream that sends periodic keep-alive messages
    // In a full implementation, this would integrate with the rmcp transport layer
    let stream = stream::unfold((rx, tx), |(mut rx, tx)| async move {
        tokio::select! {
            msg = rx.recv() => {
                match msg {
                    Some(data) => {
                        let event = Event::default().data(data);
                        Some((Ok(event), (rx, tx)))
                    }
                    None => None,
                }
            }
            _ = tokio::time::sleep(Duration::from_secs(30)) => {
                let event = Event::default()
                    .event("keep-alive")
                    .data("ping");
                Some((Ok(event), (rx, tx)))
            }
        }
    });
    
    info!("MCP: SSE endpoint established");
    
    Ok(Sse::new(stream).keep_alive(
        KeepAlive::new()
            .interval(Duration::from_secs(30))
            .text("keep-alive"),
    ))
}