AI-Studio/runtime/src/runtime_api.rs

use log::info;
use once_cell::sync::Lazy;
use axum::extract::DefaultBodyLimit;
use axum::routing::{delete, get, post};
use axum::Router;
use axum_server::tls_rustls::RustlsConfig;
use std::net::SocketAddr;
use std::sync::Once;
use crate::runtime_certificate::{CERTIFICATE, CERTIFICATE_PRIVATE_KEY};
use crate::environment::is_dev;
use crate::network::get_available_port;

static RUSTLS_CRYPTO_PROVIDER_INIT: Once = Once::new();

/// The request body limit for one batch of prompt injection filtering. The app caps the text
/// it returns to a model well below this, so the limit is headroom, not a working constraint.
const PROMPT_INJECTION_BATCH_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024;

/// The port used for the runtime API server. In the development environment, we use a fixed
/// port, in the production environment we use the next available port. This differentiation
/// is necessary because we cannot communicate the port to the .NET server in the development
/// environment.
pub static API_SERVER_PORT: Lazy<u16> = Lazy::new(|| {
    if is_dev() {
        5000
    } else {
        get_available_port().unwrap()
    }
});

/// Starts the runtime API server. The server is used to communicate with the .NET server and
/// to provide additional functionality to the Tauri app.
pub fn start_runtime_api() {
    let api_port = *API_SERVER_PORT;
    info!("Try to start the API server on 'http://localhost:{api_port}'...");

    let app = Router::new()
        .route("/system/dotnet/port", get(crate::dotnet::dotnet_port))
        .route("/system/dotnet/ready", get(crate::dotnet::dotnet_ready))
        .route("/system/qdrant-edge/info", get(crate::qdrant_edge_database::qdrant_edge_info))
        .route("/system/qdrant-edge/ensure", post(crate::qdrant_edge_database::ensure_qdrant_edge_store))
        .route("/system/qdrant-edge/insert", post(crate::qdrant_edge_database::insert_qdrant_edge_embedding))
        .route("/system/qdrant-edge/search", post(crate::qdrant_edge_database::search_qdrant_edge_embeddings))
        .route("/system/qdrant-edge/delete-file", post(crate::qdrant_edge_database::delete_qdrant_edge_embedding_by_file))
        .route("/system/qdrant-edge/optimize", post(crate::qdrant_edge_database::optimize_qdrant_edge_store))
        .route("/system/qdrant-edge/delete-store", post(crate::qdrant_edge_database::delete_qdrant_edge_store))
        .route("/clipboard/set", post(crate::clipboard::set_clipboard))
        .route("/share/file", post(crate::share_sheet::share_file))
        .route("/events", get(crate::app_window::get_event_stream))
        .route("/updates/check", get(crate::app_window::check_for_update))
        .route("/updates/install", get(crate::app_window::install_update))
        .route("/app/exit", post(crate::app_window::exit_app))
        .route("/select/directory", post(crate::file_actions::select_directory))
        .route("/select/file", post(crate::file_actions::select_file))
        .route("/select/files", post(crate::file_actions::select_files))
        .route("/save/file", post(crate::file_actions::save_file))
        .route("/open/path", post(crate::file_actions::open_path_in_file_manager))
        .route("/open/document", post(crate::file_actions::open_document))
        .route("/secrets/get", post(crate::secret::get_secret))
        .route("/secrets/store", post(crate::secret::store_secret))
        .route("/secrets/delete", post(crate::secret::delete_secret))
        .route("/system/directories/config", get(crate::environment::get_config_directory))
        .route("/system/directories/data", get(crate::environment::get_data_directory))
        .route("/system/runtime/info", get(crate::environment::get_runtime_info))
        .route("/system/language", get(crate::environment::read_user_language))
        .route("/system/username", get(crate::environment::read_user_name))
        .route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
        .route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
        .route("/retrieval/fs/extract", get(crate::file_data::extract_data))
        .route("/security/prompt-injection/sanitize", post(crate::prompt_injection::api::sanitize))
        //
        // A batch carries every text of one tool call, which is far more than Axum's 2 MB
        // default allows. Exceeding that limit would answer 413, and the app treats a failed
        // filter call as "cannot filter" and uses the text unfiltered — the protection would
        // drop out silently on exactly the largest results. Hence the explicit limit.
        //
        .route("/security/prompt-injection/sanitize-batch", post(crate::prompt_injection::api::sanitize_batch)
            .layer(DefaultBodyLimit::max(PROMPT_INJECTION_BATCH_BODY_LIMIT_BYTES)))
        .route("/media/jobs", post(crate::media::create_job))
        .route("/media/jobs/{id}/events", get(crate::media::get_job_events))
        .route("/media/jobs/{id}", delete(crate::media::cancel_job))
        .route("/image/prepare", post(crate::image::prepare_image))
        .route("/log/paths", get(crate::log::get_log_paths))
        .route("/log/event", post(crate::log::log_event))
        .route("/tokenizer/count", post(crate::tokenizer::token_count))
        .route("/tokenizer/validate", post(crate::tokenizer::validate_tokenizer))
        .route("/tokenizer/store", post(crate::tokenizer::store_tokenizer))
        .route("/tokenizer/delete", post(crate::tokenizer::delete_tokenizer))
        .route("/shortcuts/register", post(crate::app_window::register_shortcut))
        .route("/shortcuts/validate", post(crate::app_window::validate_shortcut))
        .route("/shortcuts/suspend", post(crate::app_window::suspend_shortcuts))
        .route("/shortcuts/resume", post(crate::app_window::resume_shortcuts));

    tauri::async_runtime::spawn(async move {
        install_rustls_crypto_provider();

        let cert = CERTIFICATE.get().unwrap().clone();
        let key = CERTIFICATE_PRIVATE_KEY.get().unwrap().clone();
        let tls_config = RustlsConfig::from_pem(cert, key).await.unwrap();
        let addr = SocketAddr::from(([127, 0, 0, 1], api_port));

        axum_server::bind_rustls(addr, tls_config)
            .serve(app.into_make_service())
            .await
            .unwrap();
    });
}

fn install_rustls_crypto_provider() {
    RUSTLS_CRYPTO_PROVIDER_INIT.call_once(|| {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
    });
}

#[cfg(test)]
pub(crate) mod test_support {
    use std::sync::Mutex;
    use std::time::{Duration, Instant};

    /// How long the lock stays held while the request waits for it.
    const LOCK_DURATION: Duration = Duration::from_secs(1);

    /// How long the probe may wait for the runtime. Far below the lock duration, so a request which
    /// blocks the worker while it waits fails the test every time.
    const PROBE_LIMIT: Duration = Duration::from_millis(200);

    /// Starts a request while another thread holds the given lock, as if a long operation were
    /// running, and checks that a probe still gets through the runtime meanwhile. Returns the
    /// response once the lock is free again.
    ///
    /// Call it from a test with a single worker, `#[tokio::test(flavor = "multi_thread",
    /// worker_threads = 1)]`. A request which blocks that worker while it waits then holds up the
    /// probe, just as in the app it holds up the one HTTP/2 connection which all calls share.
    pub(crate) async fn assert_runtime_stays_free<T, F>(lock: &'static Mutex<T>, request: F) -> F::Output
    where
        T: Send + 'static,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let (locked_sender, locked_receiver) = std::sync::mpsc::channel();
        let lock_holder = std::thread::spawn(move || {
            let guard = lock.lock().unwrap();
            locked_sender.send(()).unwrap();
            std::thread::sleep(LOCK_DURATION);
            drop(guard);
        });

        locked_receiver.recv().unwrap();
        let request = tokio::spawn(request);

        let probe_started = Instant::now();
        tokio::spawn(async {}).await.unwrap();
        let probe_duration = probe_started.elapsed();

        lock_holder.join().unwrap();
        let response = request.await.unwrap();

        assert!(probe_duration < PROBE_LIMIT, "The probe waited {probe_duration:?} for the runtime.");
        response
    }
}