tftio-kb 2.5.4

#![cfg_attr(
    not(test),
    deny(clippy::unwrap_used, clippy::panic, clippy::indexing_slicing)
)]
//! `kb-mcp` executable: the MCP stdio server for the kb knowledge base.
//!
//! Reads JSON-RPC messages from stdin, writes responses to stdout, and
//! keeps protocol noise off the protocol channel by sending diagnostics
//! to stderr. Configuration is by environment variable so the server is
//! trivial to launch from an MCP host (Claude Desktop, etc.):
//!
//! - `KB_DB_PATH` — `SQLite` database file (default
//!   `$HOME/.local/share/kb/kb.db`). Shares the
//!   database with `kb-server` via `SQLite` WAL mode (one writer, many
//!   readers concurrently across processes).
//! - `KB_EMBEDDING_BASE_URL` / `KB_EMBEDDING_MODEL` /
//!   `KB_EMBEDDING_API_KEY` — when present, enables hybrid (FTS +
//!   vector) search and write-path embedding regeneration; absent means
//!   FTS-only.
//!
//! Mirrors the Haskell `app/kb-mcp/Main.hs`.

use std::process::ExitCode;
use std::sync::{Arc, Mutex};

use kb::auth::BearerToken;

use kb::embedding;
use kb::mcp::{ToolsState, builtin_methods, run_mcp_with, tools_methods};
use kb::storage;

#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
#[allow(
    clippy::disallowed_methods,
    reason = "kb-mcp reads KB_TOKEN/KB_DB_PATH once at the binary edge (REPO_INVARIANTS.md #5)"
)]
async fn main() -> ExitCode {
    // Require KB_TOKEN env var for MCP access control (same as kb-server).
    let token = match std::env::var("KB_TOKEN") {
        Ok(t) if !t.trim().is_empty() => t,
        _ => {
            eprintln!("kb-mcp: KB_TOKEN env var not set — refusing to start");
            return ExitCode::FAILURE;
        }
    };
    // Validate token format at startup (no request parsing needed).
    let _auth = BearerToken::new(&token);

    let db_path = std::env::var("KB_DB_PATH")
        .unwrap_or_else(|_| storage::default_db_path().to_string_lossy().into_owned());
    eprintln!("kb-mcp opening {db_path}");

    let conn = match storage::open_db(&db_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("kb-mcp: failed to open database {db_path}: {e}");
            return ExitCode::FAILURE;
        }
    };
    if let Err(e) = storage::init_db(&conn) {
        eprintln!("kb-mcp: failed to init database schema: {e}");
        return ExitCode::FAILURE;
    }

    let (embedding_client, embedding_model, runtime) =
        if let Some(cfg) = embedding::read_embedding_config_from_env() {
            eprintln!(
                "kb-mcp: embeddings via {} using model {}",
                cfg.base_url, cfg.model
            );
            let model = cfg.model.clone();
            let client: Arc<dyn kb::embedding::EmbeddingClient> =
                Arc::new(embedding::http_embedding_client(cfg));
            (
                Some(client),
                Some(model),
                Some(tokio::runtime::Handle::current()),
            )
        } else {
            eprintln!("kb-mcp: embeddings disabled (KB_EMBEDDING_BASE_URL unset)");
            (None, None, None)
        };

    let state = Arc::new(ToolsState::new(
        Arc::new(Mutex::new(conn)),
        embedding_client,
        embedding_model,
        runtime,
    ));

    let mut methods = builtin_methods();
    for (k, v) in tools_methods(state) {
        methods.insert(k, v);
    }

    eprintln!("kb-mcp ready (stdin → JSON-RPC, stdout → responses)");

    // The MCP stdio loop is synchronous (line-buffered stdin/stdout).
    // Run it on a dedicated blocking thread so its embedding handlers
    // can drive the tokio Handle via `block_on` without re-entering the
    // worker pool.
    let result = match tokio::task::spawn_blocking(move || run_mcp_with(&methods)).await {
        Ok(inner) => inner,
        Err(e) => {
            eprintln!("kb-mcp: blocking task failed to join: {e}");
            return ExitCode::FAILURE;
        }
    };

    if let Err(e) = result {
        eprintln!("kb-mcp: stdio loop terminated: {e}");
        return ExitCode::FAILURE;
    }
    ExitCode::SUCCESS
}