server.cpp

#include "server-common.h"
#include "server-context.h"
#include "server-http.h"
#include "server-models.h"
#include "server-cors-proxy.h"
#include "server-tools.h"
#include "worker_isolation.h"

#include "arg.h"
#include "build-info.h"
#include "common.h"
#include "fit.h"
#include "llama.h"
#include "log.h"

#include <atomic>
#include <clocale>
#include <cstdlib>
#include <exception>
#include <memory>
#include <signal.h>
#include <thread> // for std::thread::hardware_concurrency

#if defined(_WIN32)
#include <windows.h>
#endif

static std::function<void(int)> shutdown_handler;
static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT;

static inline void signal_handler(int signal) {
    if (is_terminating.test_and_set()) {
        // in case it hangs, we can force terminate the server by hitting Ctrl+C twice
        // this is for better developer experience, we can remove when the server is stable enough
        fprintf(stderr, "Received second interrupt, terminating immediately.\n");
        exit(1);
    }

    shutdown_handler(signal);
}

// wrapper function that handles exceptions and logs errors
// this is to make sure handler_t never throws exceptions; instead, it returns an error response
static server_http_context::handler_t ex_wrapper(server_http_context::handler_t func) {
    return [func = std::move(func)](const server_http_req & req) -> server_http_res_ptr {
        std::string message;
        error_type error;
        try {
            return func(req);
        } catch (const std::invalid_argument & e) {
            // treat invalid_argument as invalid request (400)
            error = ERROR_TYPE_INVALID_REQUEST;
            message = e.what();
        } catch (const std::exception & e) {
            // treat other exceptions as server error (500)
            error = ERROR_TYPE_SERVER;
            message = e.what();
        } catch (...) {
            error = ERROR_TYPE_SERVER;
            message = "unknown error";
        }

        auto res = std::make_unique<server_http_res>();
        res->status = 500;
        try {
            json error_data = format_error_response(message, error);
            res->status = json_value(error_data, "code", 500);
            res->data = safe_json_to_str({{ "error", error_data }});
            SRV_WRN("got exception: %s\n", res->data.c_str());
        } catch (const std::exception & e) {
            SRV_ERR("got another exception: %s | while handling exception: %s\n", e.what(), message.c_str());
            res->data = "Internal Server Error";
        }
        return res;
    };
}

// satisfies -Wmissing-declarations
int llama_server(int argc, char ** argv);

int llama_server(int argc, char ** argv) {
    std::setlocale(LC_NUMERIC, "C");

    // own arguments required by this example
    common_params params;

    common_init();

    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
        return 1;
    }

    // worker-isolation parent mode: opt-in via LLAMA_WORKER_ISOLATION=1.
    // Parent stays CUDA-free, owns the HTTP port, fork+execv's a child
    // on a private port that owns the model + GPU. SIGKILLing the child
    // via POST /v1/admin/unload reclaims ALL VRAM (primary CUDA context
    // teardown). The child sees LLAMA_WORKER_ISOLATION_CHILD=1 and runs
    // as a normal single-process llama-server. See worker_isolation.h.
    auto getenv_truthy = [](const char * name) {
        const char * v = std::getenv(name);
        return v && v[0] && v[0] != '0';
    };
    const bool worker_iso_parent =
        getenv_truthy("LLAMA_WORKER_ISOLATION") &&
        !getenv_truthy("LLAMA_WORKER_ISOLATION_CHILD");

    llama_backend_init();
    llama_numa_init(params.numa);

    // router server never loads a model and must not touch the GPU
    // skip device enumeration so the CUDA primary context stays uncreated
    const bool is_router_server = params.model.path.empty();
    // worker-isolation parent is router-like: doesn't load a model in
    // this process, holds the inbound port + proxies to a forked child.
    const bool is_parent_only = is_router_server || worker_iso_parent;
    common_params_print_info(params, !is_parent_only);

    // validate batch size for embeddings
    // embeddings require all tokens to be processed in a single ubatch
    // see https://github.com/ggml-org/llama.cpp/issues/12836
    if (params.embedding && params.n_batch > params.n_ubatch) {
        SRV_WRN("embeddings enabled with n_batch (%d) > n_ubatch (%d)\n", params.n_batch, params.n_ubatch);
        SRV_WRN("setting n_batch = n_ubatch = %d to avoid assertion failure\n", params.n_ubatch);
        params.n_batch = params.n_ubatch;
    }

    if (params.n_parallel < 0) {
        SRV_INF("%s", "n_parallel is set to auto, using n_parallel = 4 and kv_unified = true\n");

        params.n_parallel = 4;
        params.kv_unified = true;
    }

    // for consistency between server router mode and single-model mode, we set the same model name as alias
    if (params.model_alias.empty() && !params.model.name.empty()) {
        params.model_alias.insert(params.model.name);
    }

    // struct that contains llama context and inference
    server_context ctx_server;

    server_http_context ctx_http;
    if (!ctx_http.init(params)) {
        SRV_ERR("%s", "failed to initialize HTTP server\n");
        return 1;
    }

    //
    // Router
    //

    // register API routes
    server_routes routes(params, ctx_server);
    server_tools tools;

    // worker-isolation parent state. Spawns child lazily on first
    // chat/completion/embedding request via routes overrides below.
    std::unique_ptr<worker_isolation> worker_iso;
    if (worker_iso_parent) {
        SRV_INF("%s", "worker-isolation: parent mode active (LLAMA_WORKER_ISOLATION=1)\n");
        SRV_INF("%s", "  parent stays CUDA-free, child holds model + GPU\n");
        SRV_INF("%s", "  POST /v1/admin/unload SIGKILLs child → all VRAM reclaimed\n");
        worker_iso = std::make_unique<worker_isolation>(params, argc, argv);

        // Default GPU for the child: env WORKER_DEFAULT_GPU (a
        // CUDA_VISIBLE_DEVICES value — index or "GPU-..." UUID). Empty →
        // inherit the container CUDA_VISIBLE_DEVICES (backward compatible).
        // Per-request override is the X-Worker-Gpu header (see worker_isolation).
        if (const char * dg = std::getenv("WORKER_DEFAULT_GPU"); dg && dg[0]) {
            worker_iso->set_default_gpu(dg);
            SRV_INF("  WORKER_DEFAULT_GPU=%s (per-request override: X-Worker-Gpu header)\n", dg);
        }

        // Replace all model-facing handlers with proxies through the
        // child subprocess. Drain check is done inside each lambda so
        // streaming responses already in flight aren't aborted.
        auto wi = worker_iso.get();
        auto proxy_post_with_drain = [wi](const server_http_req & req) -> server_http_res_ptr {
            if (wi->is_draining()) {
                auto res = std::make_unique<server_http_res>();
                res->status = 503;
                res->data = safe_json_to_str({{"error", format_error_response(
                    "service is draining; new requests rejected (admin/drain in effect)",
                    ERROR_TYPE_UNAVAILABLE)}});
                return res;
            }
            return wi->proxy(req, "POST");
        };
        auto proxy_post_no_drain = [wi](const server_http_req & req) -> server_http_res_ptr {
            return wi->proxy_noflight(req, "POST");
        };
        auto proxy_get_no_drain  = [wi](const server_http_req & req) -> server_http_res_ptr {
            return wi->proxy_noflight(req, "GET");
        };

        // Model-bearing endpoints: drain-gated + counted toward in_flight.
        routes.post_completions            = proxy_post_with_drain;
        routes.post_completions_oai        = proxy_post_with_drain;
        routes.post_chat_completions       = proxy_post_with_drain;
        routes.post_responses_oai          = proxy_post_with_drain;
        routes.post_transcriptions_oai     = proxy_post_with_drain;
        routes.post_anthropic_messages     = proxy_post_with_drain;
        routes.post_anthropic_count_tokens = proxy_post_with_drain;
        routes.post_infill                 = proxy_post_with_drain;
        routes.post_embeddings             = proxy_post_with_drain;
        routes.post_embeddings_oai         = proxy_post_with_drain;
        routes.post_rerank                 = proxy_post_with_drain;
        // Metadata + tokenize endpoints — proxy but don't gate (cheap,
        // useful while draining; only fail when child is fully down).
        routes.get_metrics                 = proxy_get_no_drain;
        routes.get_slots                   = proxy_get_no_drain;
        routes.post_slots                  = proxy_post_no_drain;
        routes.get_props                   = proxy_get_no_drain;
        routes.post_props                  = proxy_post_no_drain;
        routes.get_models                  = proxy_get_no_drain;
        routes.post_tokenize               = proxy_post_no_drain;
        routes.post_detokenize             = proxy_post_no_drain;
        routes.post_apply_template         = proxy_post_no_drain;
        routes.get_lora_adapters           = proxy_get_no_drain;
        routes.post_lora_adapters          = proxy_post_no_drain;
        // /health stays parent-side and reports loaded/busy/draining/in_flight.
        routes.get_health = [wi](const server_http_req &) -> server_http_res_ptr {
            auto res = std::make_unique<server_http_res>();
            res->status = 200;
            json body = {
                {"status",    "ok"},
                {"loaded",    wi->is_loaded()},
                {"busy",      wi->in_flight() > 0},
                {"draining",  wi->is_draining()},
                {"in_flight", wi->in_flight()},
            };
            res->data = safe_json_to_str(body);
            return res;
        };
    }

    std::optional<server_models_routes> models_routes{};
    if (is_router_server) {
        // setup server instances manager
        try {
            models_routes.emplace(params, argc, argv);
        } catch (const std::exception & e) {
            SRV_ERR("failed to initialize router models: %s\n", e.what());
            return 1;
        }

        // proxy handlers
        // note: routes.get_health stays the same
        routes.get_metrics                 = models_routes->proxy_get;
        routes.post_props                  = models_routes->proxy_post;
        routes.post_completions            = models_routes->proxy_post;
        routes.post_completions_oai        = models_routes->proxy_post;
        routes.post_chat_completions       = models_routes->proxy_post;
        routes.post_control                = models_routes->proxy_post;
        routes.post_responses_oai          = models_routes->proxy_post;
        routes.post_transcriptions_oai     = models_routes->proxy_post;
        routes.post_anthropic_messages     = models_routes->proxy_post;
        routes.post_anthropic_count_tokens = models_routes->proxy_post;
        routes.post_infill                 = models_routes->proxy_post;
        routes.post_embeddings             = models_routes->proxy_post;
        routes.post_embeddings_oai         = models_routes->proxy_post;
        routes.post_rerank                 = models_routes->proxy_post;
        routes.post_tokenize               = models_routes->proxy_post;
        routes.post_detokenize             = models_routes->proxy_post;
        routes.post_apply_template         = models_routes->proxy_post;
        routes.post_chat_completions_tok   = models_routes->proxy_post;
        routes.post_responses_tok_oai      = models_routes->proxy_post;
        routes.get_lora_adapters           = models_routes->proxy_get;
        routes.post_lora_adapters          = models_routes->proxy_post;
        routes.get_slots                   = models_routes->proxy_get;
        routes.post_slots                  = models_routes->proxy_post;

        // custom routes for router
        routes.get_props                   = models_routes->get_router_props;
        routes.get_models                  = models_routes->get_router_models;

        ctx_http.post("/models/load",          ex_wrapper(models_routes->post_router_models_load));
        ctx_http.post("/models/unload",        ex_wrapper(models_routes->post_router_models_unload));
    }

    ctx_http.get ("/health",                   ex_wrapper(routes.get_health)); // public endpoint (no API key check)
    ctx_http.get ("/v1/health",                ex_wrapper(routes.get_health)); // public endpoint (no API key check)
    ctx_http.get ("/metrics",                  ex_wrapper(routes.get_metrics));
    ctx_http.get ("/props",                    ex_wrapper(routes.get_props));
    ctx_http.post("/props",                    ex_wrapper(routes.post_props));
    ctx_http.get ("/models",                   ex_wrapper(routes.get_models)); // public endpoint (no API key check)
    ctx_http.get ("/v1/models",                ex_wrapper(routes.get_models)); // public endpoint (no API key check)
    ctx_http.post("/completion",               ex_wrapper(routes.post_completions)); // legacy
    ctx_http.post("/completions",              ex_wrapper(routes.post_completions));
    ctx_http.post("/v1/completions",           ex_wrapper(routes.post_completions_oai));
    ctx_http.post("/chat/completions",         ex_wrapper(routes.post_chat_completions));
    ctx_http.post("/v1/chat/completions",      ex_wrapper(routes.post_chat_completions));
    ctx_http.post("/v1/chat/completions/control", ex_wrapper(routes.post_control));
    ctx_http.post("/v1/responses",             ex_wrapper(routes.post_responses_oai));
    ctx_http.post("/responses",                ex_wrapper(routes.post_responses_oai));
    ctx_http.post("/v1/audio/transcriptions",  ex_wrapper(routes.post_transcriptions_oai));
    ctx_http.post("/audio/transcriptions",     ex_wrapper(routes.post_transcriptions_oai));
    ctx_http.post("/v1/messages",              ex_wrapper(routes.post_anthropic_messages)); // anthropic messages API
    ctx_http.post("/infill",                   ex_wrapper(routes.post_infill));
    ctx_http.post("/embedding",                ex_wrapper(routes.post_embeddings)); // legacy
    ctx_http.post("/embeddings",               ex_wrapper(routes.post_embeddings));
    ctx_http.post("/v1/embeddings",            ex_wrapper(routes.post_embeddings_oai));
    ctx_http.post("/rerank",                   ex_wrapper(routes.post_rerank));
    ctx_http.post("/reranking",                ex_wrapper(routes.post_rerank));
    ctx_http.post("/v1/rerank",                ex_wrapper(routes.post_rerank));
    ctx_http.post("/v1/reranking",             ex_wrapper(routes.post_rerank));
    ctx_http.post("/tokenize",                 ex_wrapper(routes.post_tokenize));
    ctx_http.post("/detokenize",               ex_wrapper(routes.post_detokenize));
    ctx_http.post("/apply-template",           ex_wrapper(routes.post_apply_template));
    // token counting
    ctx_http.post("/chat/completions/input_tokens",    ex_wrapper(routes.post_chat_completions_tok));
    ctx_http.post("/v1/chat/completions/input_tokens", ex_wrapper(routes.post_chat_completions_tok));
    ctx_http.post("/responses/input_tokens",           ex_wrapper(routes.post_responses_tok_oai));
    ctx_http.post("/v1/responses/input_tokens",        ex_wrapper(routes.post_responses_tok_oai));
    ctx_http.post("/v1/messages/count_tokens",         ex_wrapper(routes.post_anthropic_count_tokens)); // anthropic token counting
    // LoRA adapters hotswap
    ctx_http.get ("/lora-adapters",            ex_wrapper(routes.get_lora_adapters));
    ctx_http.post("/lora-adapters",            ex_wrapper(routes.post_lora_adapters));
    // Save & load slots
    ctx_http.get ("/slots",                    ex_wrapper(routes.get_slots));
    ctx_http.post("/slots/:id_slot",           ex_wrapper(routes.post_slots));

    // worker-isolation admin endpoints — register only in parent mode.
    // These match kob-gpu-gate's expected paths (see koblibs/kob-gpu-gate
    // src/lib.rs::acquire_for_avatar — drain → wait → unload sequence).
    if (worker_iso_parent) {
        auto wi = worker_iso.get();
        // POST /v1/admin/drain — block new requests; in-flight finish.
        auto drain_handler = [wi](const server_http_req &) -> server_http_res_ptr {
            wi->set_draining(true);
            auto res = std::make_unique<server_http_res>();
            res->status = 200;
            res->data = safe_json_to_str({
                {"status",    "ok"},
                {"draining",  true},
                {"in_flight", wi->in_flight()},
            });
            return res;
        };
        // POST /v1/admin/unload — SIGKILL child, reclaim all VRAM.
        // Also clears the draining flag so the next /v1/chat/completions
        // request lazily re-spawns the child (avatar-release → llm-chat
        // path; the kob-gpu-gate owner can still POST /v1/admin/load
        // first if it wants a pre-warm).
        auto unload_handler = [wi](const server_http_req &) -> server_http_res_ptr {
            bool killed = wi->shutdown_child();
            wi->set_draining(false);
            auto res = std::make_unique<server_http_res>();
            res->status = 200;
            res->data = safe_json_to_str({
                {"status",  killed ? "ok" : "idle"},
                {"loaded",  wi->is_loaded()},
                {"unloaded", killed},
            });
            return res;
        };
        // POST /v1/admin/load — pre-warm + clear drain flag.
        auto load_handler = [wi](const server_http_req &) -> server_http_res_ptr {
            wi->set_draining(false);
            bool was_loaded = wi->is_loaded();
            bool ok = wi->ensure_loaded();
            auto res = std::make_unique<server_http_res>();
            if (!ok) {
                res->status = 500;
                res->data = safe_json_to_str({{"error", format_error_response(
                    "failed to load worker child", ERROR_TYPE_SERVER)}});
                return res;
            }
            res->status = 200;
            res->data = safe_json_to_str({
                {"status", "ok"},
                {"loaded", wi->is_loaded()},
                {"was_loaded", was_loaded},
            });
            return res;
        };
        ctx_http.post("/v1/admin/drain",  ex_wrapper(drain_handler));
        ctx_http.post("/v1/admin/unload", ex_wrapper(unload_handler));
        ctx_http.post("/v1/admin/load",   ex_wrapper(load_handler));
        // legacy aliases without /v1 prefix (some clients use bare paths).
        ctx_http.post("/admin/drain",     ex_wrapper(drain_handler));
        ctx_http.post("/admin/unload",    ex_wrapper(unload_handler));
        ctx_http.post("/admin/load",      ex_wrapper(load_handler));
    }

    // Google Cloud Platform (Vertex AI) compat
    ctx_http.register_gcp_compat();

    // CORS proxy (EXPERIMENTAL, only used by the Web UI for MCP)
    // Supports both new ui_mcp_proxy and deprecated webui_mcp_proxy fields
    if (params.ui_mcp_proxy || params.webui_mcp_proxy) {
        SRV_WRN("%s", "-----------------\n");
        SRV_WRN("%s", "CORS proxy is enabled, do not expose server to untrusted environments\n");
        SRV_WRN("%s", "This feature is EXPERIMENTAL and may be removed or changed in future versions\n");
        SRV_WRN("%s", "-----------------\n");
        ctx_http.get ("/cors-proxy",      ex_wrapper(proxy_handler_get));
        ctx_http.post("/cors-proxy",      ex_wrapper(proxy_handler_post));
    }
    // EXPERIMENTAL built-in tools
    if (!params.server_tools.empty()) {
        try {
            tools.setup(params.server_tools);
        } catch (const std::exception & e) {
            SRV_ERR("tools setup failed: %s\n", e.what());
            return 1;
        }
        SRV_WRN("%s", "-----------------\n");
        SRV_WRN("%s", "Built-in tools are enabled, do not expose server to untrusted environments\n");
        SRV_WRN("%s", "This feature is EXPERIMENTAL and may be changed in the future\n");
        SRV_WRN("%s", "-----------------\n");
        ctx_http.get ("/tools",           ex_wrapper(tools.handle_get));
        ctx_http.post("/tools",           ex_wrapper(tools.handle_post));
    }

    //
    // Start the server
    //

    std::function<void()> clean_up;

    if (is_router_server) {
        SRV_INF("%s", "starting router server, no model will be loaded in this process\n");

        clean_up = [&models_routes]() {
            SRV_INF("%s: cleaning up before exit...\n", __func__);
            if (models_routes.has_value()) {
                models_routes->models.unload_all();
            }
            llama_backend_free();
        };

        if (!ctx_http.start()) {
            clean_up();
            SRV_ERR("%s", "exiting due to HTTP server error\n");
            return 1;
        }
        ctx_http.is_ready.store(true);

        shutdown_handler = [&](int) {
            ctx_http.stop();
        };

    } else if (worker_iso_parent) {
        SRV_INF("%s", "starting worker-isolation parent (no model loaded in this process)\n");
        SRV_INF("%s", "  child will be spawned lazily on first model-bound request\n");

        clean_up = [&ctx_http, &worker_iso]() {
            SRV_INF("%s: cleaning up before exit...\n", __func__);
            if (worker_iso) {
                worker_iso->shutdown_child();
            }
            ctx_http.stop();
            llama_backend_free();
        };

        if (!ctx_http.start()) {
            clean_up();
            SRV_ERR("%s", "exiting due to HTTP server error\n");
            return 1;
        }
        ctx_http.is_ready.store(true);

        shutdown_handler = [&](int) {
            if (worker_iso) worker_iso->shutdown_child();
            ctx_http.stop();
        };

    } else {
        // setup clean up function, to be called before exit
        clean_up = [&ctx_http, &ctx_server]() {
            SRV_INF("%s: cleaning up before exit...\n", __func__);
            ctx_http.stop();
            ctx_server.terminate();
            llama_backend_free();
        };

        // start the HTTP server before loading the model to be able to serve /health requests
        if (!ctx_http.start()) {
            clean_up();
            SRV_ERR("%s", "exiting due to HTTP server error\n");
            return 1;
        }

        // load the model
        SRV_INF("%s", "loading model\n");

        if (server_models::is_child_server()) {
            ctx_server.on_sleeping_changed([&](bool sleeping) {
                server_models::notify_router_sleeping_state(sleeping);
            });
        }

        if (!ctx_server.load_model(params)) {
            clean_up();
            if (ctx_http.thread.joinable()) {
                ctx_http.thread.join();
            }
            SRV_ERR("%s", "exiting due to model loading error\n");
            return 1;
        }

        routes.update_meta(ctx_server);
        ctx_http.is_ready.store(true);

        SRV_INF("%s", "model loaded\n");

        shutdown_handler = [&](int) {
            // this will unblock start_loop()
            ctx_server.terminate();
        };
    }

    // TODO: refactor in common/console
#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
    struct sigaction sigint_action;
    sigint_action.sa_handler = signal_handler;
    sigemptyset (&sigint_action.sa_mask);
    sigint_action.sa_flags = 0;
    sigaction(SIGINT, &sigint_action, NULL);
    sigaction(SIGTERM, &sigint_action, NULL);
#elif defined (_WIN32)
    auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
        return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
    };
    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
#endif

    if (is_router_server) {
        SRV_INF("router server is listening on %s\n", ctx_http.listening_address.c_str());
        SRV_WRN("%s", "NOTE: router mode is experimental\n");
        SRV_WRN("%s", "      it is not recommended to use this mode in untrusted environments\n");
        if (ctx_http.thread.joinable()) {
            ctx_http.thread.join(); // keep the main thread alive
        }

        // when the HTTP server stops, clean up and exit
        clean_up();
    } else if (worker_iso_parent) {
        SRV_INF("worker-isolation parent listening on %s\n", ctx_http.listening_address.c_str());
        if (ctx_http.thread.joinable()) {
            ctx_http.thread.join();
        }
        clean_up();
    } else {
        SRV_INF("server is listening on %s\n", ctx_http.listening_address.c_str());

        // optionally, notify router server that this instance is ready
        std::thread monitor_thread;
        if (server_models::is_child_server()) {
            json model_info = routes.get_model_info();
            monitor_thread = server_models::setup_child_server(shutdown_handler, model_info);
        }

        // this call blocks the main thread until queue_tasks.terminate() is called
        ctx_server.start_loop();

        clean_up();
        if (ctx_http.thread.joinable()) {
            ctx_http.thread.join();
        }
        if (monitor_thread.joinable()) {
            monitor_thread.join();
        }

        auto * ll_ctx = ctx_server.get_llama_context();
        if (ll_ctx != nullptr) {
            common_memory_breakdown_print(ll_ctx);
        }
    }

    return 0;
}