knot-server 0.8.8

mod cleanup;
mod config;
mod fs_utils;
mod git;
mod handlers;
mod local_sync;
mod locking;
mod metrics;
mod models;
mod progress_store;
mod registry;
mod scheduler;
mod startup_guard;
mod telemetry;
mod time_utils;
mod webhook;
mod worker;

use std::path::Path;
use std::sync::{Arc, Mutex};

use axum::Router;
use axum::extract::{Request, State};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Redirect, Response};
use axum::routing::get;
use knot::db::graph::{ConnectExt, GraphDb};
use knot::db::vector::{VectorConnectExt, VectorDb};
use knot::pipeline::embed::Embedder;
use models::AppState;
use registry::Registry;
use std::collections::HashMap;
use std::str::FromStr as _;
use tokio::signal;
use utoipa::OpenApi;
use utoipa_axum::router::OpenApiRouter;
use utoipa_axum::routes;
use utoipa_swagger_ui::SwaggerUi;

#[tokio::main]
#[expect(clippy::too_many_lines, reason = "deferred refactoring")]
async fn main() -> anyhow::Result<()> {
    // Order matters: parse config first (clap may print help/errors to stderr,
    // which is fine), build the tracer provider if enabled, then install the
    // subscriber. No `tracing::` macro runs before the subscriber is installed.
    let cfg = config::ServerConfig::from_env();

    let tracer_provider = if cfg.tracing_enabled {
        Some(telemetry::init_tracer_provider(&cfg)?)
    } else {
        None
    };
    telemetry::init_subscriber(tracer_provider.as_ref());

    let metrics_handle = if cfg.metrics_enabled {
        Some(metrics::init()?)
    } else {
        None
    };

    tracing::info!("Starting knot-server v{}", env!("CARGO_PKG_VERSION"));
    tracing::info!("Binding to {}:{}", cfg.bind_addr, cfg.port);

    // Resolve the embedding model, the derived dimension and the derived
    // collection. The dimension is model-derived (KNOT_SERVER_EMBED_DIM is
    // deprecated); an explicitly supplied collection always wins, otherwise
    // the model's suffix is applied to the base default.
    let embed_model = config::resolved_embed_model();
    let embed_dim = config::resolve_embed_dim(&embed_model, cfg.embed_dim)?;
    let embed_choice = knot::pipeline::embed::EmbedModelChoice::from_str(&embed_model)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let qdrant_collection = cfg.resolved_collection(&embed_choice);
    tracing::info!(
        "Embedding model: {embed_model} (dim {embed_dim}), collection: {qdrant_collection}"
    );

    // After model resolution so the build-info gauge carries the active model.
    if metrics_handle.is_some() {
        metrics::set_build_info(&embed_model);
    }

    setup_rayon(cfg.rayon_threads);

    let graph_db = setup_neo4j(&cfg.neo4j_uri, &cfg.neo4j_user, &cfg.neo4j_password).await;

    let fastembed_cache_dir = setup_fastembed_cache(&cfg.workspace_dir)?;

    // Marker-aware startup guard (S4): before the collection is touched, the
    // guard ladder aborts on a dimension mismatch, warns when repositories
    // indexed with another model would be invisible to semantic search, and
    // stays silent on a fresh deployment.
    startup_guard::verify_server_startup(
        &cfg.qdrant_url,
        &qdrant_collection,
        (&cfg.neo4j_uri, &cfg.neo4j_user, &cfg.neo4j_password),
        &embed_model,
        &embed_choice,
    )
    .await?;

    let vector_db =
        setup_qdrant(&cfg.qdrant_url, &qdrant_collection, embed_dim, &embed_model).await?;

    let embedder = setup_embedder(fastembed_cache_dir)?;

    tracing::info!("Loading repository registry from {}...", cfg.workspace_dir);
    let mut registry = Registry::load_or_create(Path::new(&cfg.workspace_dir))?;
    tracing::info!("Registry loaded: {} repositories", registry.list().len());

    // Create job queue channel
    let (job_tx, job_rx) = tokio::sync::mpsc::channel::<models::IndexJob>(cfg.queue_capacity);

    recover_stuck_repos(&mut registry, &job_tx);

    let start_time = std::time::Instant::now();

    let state = Arc::new(AppState {
        vector_db: Arc::new(vector_db),
        graph_db: Arc::new(graph_db),
        embedder: Some(Arc::new(Mutex::new(embedder))),
        workspace_dir: cfg.workspace_dir.clone(),
        registry: Arc::new(Mutex::new(registry)),
        job_tx: job_tx.clone(),
        qdrant_url: cfg.qdrant_url.clone(),
        qdrant_collection: qdrant_collection.clone(),
        neo4j_uri: cfg.neo4j_uri.clone(),
        neo4j_user: cfg.neo4j_user.clone(),
        neo4j_password: cfg.neo4j_password.clone(),
        embed_dim,
        embed_model,
        rayon_threads: cfg.rayon_threads,
        batch_size: cfg.batch_size,
        ingest_concurrency: cfg.ingest_concurrency,
        start_time,
        progress_trackers: Arc::new(Mutex::new(HashMap::new())),
    });

    // Spawn the worker loop
    let worker_state = state.clone();
    tokio::spawn(async move {
        worker::worker_loop(job_rx, worker_state).await;
    });
    tracing::info!("Indexing worker started (concurrency: 1)");

    // Spawn the background scheduler
    let scheduler_state = state.clone();
    let poll_interval = cfg.poll_interval_secs;
    let stale_lock_timeout = cfg.stale_lock_timeout_secs;
    let max_index_age = cfg.max_index_age_secs;
    tokio::spawn(async move {
        scheduler::scheduler_loop(
            scheduler_state,
            poll_interval,
            stale_lock_timeout,
            max_index_age,
        )
        .await;
    });
    tracing::info!(
        "Background scheduler started (poll: {}s, stale lock timeout: {}s, max index age: {}s)",
        poll_interval,
        stale_lock_timeout,
        max_index_age
    );

    #[derive(OpenApi)]
    #[openapi(
        info(
            title = "knot-server",
            version = env!("CARGO_PKG_VERSION"),
            description = "REST API for managing and indexing Git repositories. Provides semantic search, caller analysis, file exploration, dependency graphs, and interactive 3D visualization.",
            license(name = "MIT", url = "https://github.com/raultov/knot-server/blob/master/LICENSE"),
        ),
        tags(
            (name = "Repositories", description = "Register, list, inspect, and delete Git repositories"),
            (name = "Search", description = "Semantic search, caller analysis, file exploration, and dependency lookup"),
            (name = "Graph", description = "Entity relationship subgraph queries"),
            (name = "Indexing", description = "Trigger manual sync and re-indexing"),
            (name = "Webhooks", description = "Git provider webhook endpoints (GitHub, GitLab, Bitbucket)"),
            (name = "Health", description = "Server health and statistics"),
            (name = "MCP", description = "Stateless MCP (Model Context Protocol) endpoint at /mcp (JSON-RPC over HTTP). Try the POST operation with a tools/list body to check connectivity, the flag, and the tool surface."),
        ),
    )]
    struct ApiDoc;

    // Build API routes with automatic OpenAPI spec collection
    let mut openapi_router = OpenApiRouter::with_openapi(ApiDoc::openapi())
        .routes(routes!(
            handlers::repo::list_repos_handler,
            handlers::register_repo_handler
        ))
        .routes(routes!(
            handlers::get_repo_handler,
            handlers::delete_repo_handler
        ))
        .routes(routes!(handlers::sync_repo_handler))
        .routes(routes!(handlers::progress_handler))
        .routes(routes!(handlers::batch_progress_handler))
        .routes(routes!(handlers::search_handler))
        .routes(routes!(handlers::callers_handler))
        .routes(routes!(handlers::search_all_handler))
        .routes(routes!(handlers::callers_all_handler))
        .routes(routes!(handlers::explore_handler))
        .routes(routes!(handlers::deps_handler))
        .routes(routes!(handlers::graph_handler))
        .routes(routes!(handlers::graph_expand_handler))
        .routes(routes!(handlers::repo_graph_handler))
        .routes(routes!(handlers::webhook_handler))
        .routes(routes!(handlers::health_handler));

    // MCP surface (MCP_ENDPOINT_PLAN.md): stateless JSON-RPC over HTTP at
    // /mcp, additive to the REST API. Registered on the OpenApiRouter (rather
    // than a plain axum route) so the three operations show up in /docs with
    // "Try it out" — the fastest way to diagnose MCP connectivity and flag
    // state. Mounted before the metrics/trace layers, so it inherits
    // `track_http` and `trace_http` with no extra work. When disabled the
    // route is simply absent and requests hit the 404 fallback.
    if cfg.mcp_enabled {
        tracing::info!("MCP endpoint enabled at /mcp (stateless, no session affinity)");
        openapi_router = openapi_router.routes(routes!(
            handlers::mcp_post_handler,
            handlers::mcp_get_handler,
            handlers::mcp_delete_handler
        ));
    }

    let (api_router, api) = openapi_router.split_for_parts();

    // Merge the generated API routes with non-API routes and Swagger UI
    let app = api_router
        .route("/favicon.ico", get(handlers::favicon_handler))
        .route("/graph", get(handlers::graph_viewer_handler));

    let app = if let Some(ref handle) = metrics_handle {
        let mh = handle.clone();
        app.route(
            "/metrics",
            get(move |State(s): State<Arc<AppState>>| metrics::metrics_handler(s, mh.clone())),
        )
    } else {
        app
    };

    let app = app.merge(
        Router::from(SwaggerUi::new("/docs").url("/api-docs/openapi.json", api))
            .layer(middleware::from_fn(docs_html_override)),
    );

    let app = if metrics_handle.is_some() {
        app.layer(middleware::from_fn(metrics::track_http))
    } else {
        app
    };

    // Layered *after* (i.e. outside) the metrics layer so the metrics timing
    // runs inside the trace span. Added last => outermost => sees the request
    // first and owns the root span for its whole lifetime.
    let app = if tracer_provider.is_some() {
        app.layer(middleware::from_fn(telemetry::trace_http))
    } else {
        app
    };

    let app = app.with_state(state);

    let listener = tokio::net::TcpListener::bind(format!("{}:{}", cfg.bind_addr, cfg.port)).await?;
    tracing::info!("knot-server listening on {}:{}", cfg.bind_addr, cfg.port);

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    // Flush the batch span processor so the last spans aren't dropped.
    if let Some(provider) = tracer_provider {
        telemetry::shutdown(provider).await;
    }

    tracing::info!("knot-server shut down gracefully");
    Ok(())
}

/// Intercepts the HTML and redirect routes served by `utoipa-swagger-ui` for
/// `/docs` and `/docs/`, returning a custom Swagger UI page that injects a
/// badge with the `knot` library version between the `knot-server` version
/// stamp and the `OAS 3.1` stamp. All other requests (CSS, JS, favicon, the
/// OpenAPI JSON, …) pass through unchanged.
async fn docs_html_override(req: Request, next: Next) -> Response {
    match req.uri().path() {
        "/docs" => Redirect::to("/docs/").into_response(),
        "/docs/" => handlers::docs_handler().await,
        _ => next.run(req).await,
    }
}

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        () = ctrl_c => {
            tracing::info!("Received SIGINT, shutting down...");
        },
        () = terminate => {
            tracing::info!("Received SIGTERM, shutting down...");
        },
    }
}

fn setup_rayon(threads: Option<usize>) {
    if let Some(t) = threads {
        rayon::ThreadPoolBuilder::new()
            .num_threads(t)
            .build_global()
            .expect("Failed to configure Rayon thread pool");
        tracing::info!("Rayon thread pool configured: {} threads", t);
    }
}

async fn setup_neo4j(uri: &str, user: &str, pass: &str) -> GraphDb {
    tracing::info!("Connecting to Neo4j at {}...", uri);
    let graph_db = loop {
        match GraphDb::connect(uri, user, pass).await {
            Ok(db) => break db,
            Err(e) => {
                tracing::warn!("Neo4j connection attempt failed: {e}");
                tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
            }
        }
    };
    loop {
        match graph_db.ensure_indexes().await {
            Ok(()) => break,
            Err(e) => {
                tracing::warn!("Neo4j index creation failed: {e}");
                tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
            }
        }
    }
    tracing::info!("Neo4j connection established");
    graph_db
}

fn setup_fastembed_cache(workspace_dir: &str) -> anyhow::Result<std::path::PathBuf> {
    let fastembed_cache_dir = Path::new(workspace_dir).join("fastembed_cache");
    let cache_str = fastembed_cache_dir
        .to_str()
        .expect("workspace_dir contains invalid UTF-8");
    std::fs::create_dir_all(cache_str)?;
    // The indexing pipeline builds its own embedder inside knot
    // (`pipeline::runner`, via `pipeline::state::fastembed_cache_dir`), and the
    // env var is the only channel knot exposes to point it at the shared
    // workspace cache instead of a per-repo `.knot/fastembed_cache`. Our own
    // search-side embedder does not need it: `Embedder::init` takes the
    // returned path directly.
    //
    // Caveat, and the reason this is only an audited exception rather than a
    // clean solution: the call happens from `main` with the Tokio runtime
    // already started, so worker threads exist and a concurrent `getenv`
    // elsewhere would be UB. Nothing reads this variable before knot's pipeline
    // does, but the guarantee rests on that timing, not on the type system.
    // Proper fix is upstream: a cache-dir field on `knot::config::Config`.
    #[expect(
        unsafe_code,
        reason = "knot only exposes the fastembed cache dir through KNOT_FASTEMBED_CACHE_DIR; \
                  set once at startup before the indexing pipeline reads it"
    )]
    unsafe {
        std::env::set_var("KNOT_FASTEMBED_CACHE_DIR", cache_str);
    }
    tracing::info!("Fastembed cache dir: {cache_str}");
    Ok(fastembed_cache_dir)
}

async fn setup_qdrant(
    url: &str,
    collection: &str,
    dim: u64,
    _embed_model: &str,
) -> anyhow::Result<VectorDb> {
    tracing::info!("Connecting to Qdrant at {}...", url);
    // The dimension/model/collection consistency was already verified by the
    // marker-aware startup guard (`startup_guard::verify_server_startup`),
    // which runs before this and inspects the real stored vector size plus
    // the Neo4j embed markers.
    let vector_db = VectorDb::connect(url, collection, dim).await?;
    vector_db.ensure_collection().await?;
    tracing::info!("Qdrant connection established");
    Ok(vector_db)
}

fn setup_embedder(cache_dir: std::path::PathBuf) -> anyhow::Result<Embedder> {
    tracing::info!("Initializing embedding model...");
    let embedder = Embedder::init(cache_dir)?;
    tracing::info!("Embedding model ready");
    Ok(embedder)
}

fn recover_stuck_repos(
    registry: &mut Registry,
    job_tx: &tokio::sync::mpsc::Sender<models::IndexJob>,
) {
    let stuck: Vec<models::RepoEntry> = registry
        .list()
        .iter()
        .filter(|r| {
            matches!(
                r.status,
                models::RepoStatus::Pending
                    | models::RepoStatus::Queued
                    | models::RepoStatus::Cloning
                    | models::RepoStatus::Pulling
                    | models::RepoStatus::Indexing
            )
        })
        .cloned()
        .collect();

    for repo in stuck {
        let git_dir = Path::new(&repo.local_path).join(".git");
        let job = if git_dir.exists() {
            tracing::warn!(
                "Recovering stuck repo '{}' (was {}): .git exists, enqueuing Pull",
                repo.id,
                repo.status
            );
            models::IndexJob::Pull {
                repo_id: repo.id.clone(),
            }
        } else {
            tracing::warn!(
                "Recovering stuck repo '{}' (was {}): .git missing, enqueuing Clone (worker will wipe + fresh-clone)",
                repo.id,
                repo.status
            );
            models::IndexJob::Clone {
                repo_id: repo.id.clone(),
            }
        };
        if job_tx.try_send(job).is_err() {
            tracing::warn!(
                "Failed to enqueue recovery job for '{}', leaving Pending",
                repo.id
            );
            let _ = registry.update_status(&repo.id, models::RepoStatus::Pending);
        } else {
            let _ = registry.update_status(&repo.id, models::RepoStatus::Queued);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::{Request as HttpRequest, StatusCode};
    use tower::util::ServiceExt;

    async fn not_found_handler() -> Response {
        (StatusCode::NOT_FOUND, "not-found").into_response()
    }

    fn build_test_app() -> Router {
        Router::new()
            .fallback(not_found_handler)
            .layer(middleware::from_fn(docs_html_override))
    }

    #[tokio::test]
    async fn docs_root_redirects_to_slash() {
        let app = build_test_app();
        let req = HttpRequest::builder()
            .uri("/docs")
            .body(Body::empty())
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::SEE_OTHER);
        assert_eq!(
            response
                .headers()
                .get("location")
                .unwrap()
                .to_str()
                .unwrap(),
            "/docs/"
        );
    }

    #[tokio::test]
    async fn docs_slash_returns_custom_html() {
        let app = build_test_app();
        let req = HttpRequest::builder()
            .uri("/docs/")
            .body(Body::empty())
            .unwrap();
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get("content-type")
                .unwrap()
                .to_str()
                .unwrap(),
            "text/html; charset=utf-8"
        );

        let body_bytes = axum::body::to_bytes(response.into_body(), 1024 * 1024)
            .await
            .unwrap();
        let body = String::from_utf8(body_bytes.to_vec()).unwrap();
        assert!(body.contains(env!("KNOT_VERSION")));
        assert!(!body.contains("{{KNOT_VERSION}}"));
    }

    #[tokio::test]
    async fn docs_assets_pass_through() {
        let app = build_test_app();
        let req = HttpRequest::builder()
            .uri("/docs/swagger-ui.css")
            .body(Body::empty())
            .unwrap();
        // The middleware should not intercept asset paths; they reach the
        // fallback that responds 404 (proving the override did not fire).
        let response = app.oneshot(req).await.unwrap();
        assert_eq!(response.status(), StatusCode::NOT_FOUND);
    }
}