server: add SSE ping interval (#24013)

diff --git a/common/arg.cpp b/common/arg.cpp
index feffb4c12ef8d72061522cfcddc9588ccaa3f8b9..7e4a90b5feeda64899dbe8cba84fd91be999a42b 100644 (file)
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -3027,6 +3027,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.timeout_write = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TIMEOUT"));
+    add_opt(common_arg(
+        {"--sse-ping-interval"}, "N",
+        string_format("server SSE ping interval in seconds (-1 = disabled, default: %d)", params.sse_ping_interval),
+        [](common_params & params, int value) {
+            params.sse_ping_interval = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSE_PING_INTERVAL"));
     add_opt(common_arg(
         {"--threads-http"}, "N",
         string_format("number of threads used to process HTTP requests (default: %d)", params.n_threads_http),

diff --git a/common/common.h b/common/common.h
index f2c7ee027b850231f26750a93cbe298dda4d38b2..208e3cee24099a8495652cb18e2761de04c8bd6f 100644 (file)
--- a/common/common.h
+++ b/common/common.h
@@ -592,6 +592,7 @@ struct common_params {
     bool    reuse_port          = false;         // allow multiple sockets to bind to the same port
     int32_t timeout_read        = 3600;          // http read timeout in seconds
     int32_t timeout_write       = timeout_read;  // http write timeout in seconds
+    int32_t sse_ping_interval   = 30;            // SSE ping interval in seconds
     int32_t n_threads_http      = -1;    // number of threads to process HTTP requests (TODO: support threadpool)
     int32_t n_cache_reuse       = 0;     // min chunk size to reuse from the cache via KV shifting
     bool    cache_prompt        = true;  // whether to enable prompt caching

diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index fae73f09f822111963552e4791d14b38119e20dd..c5d973cd20cf2d1844c39c44d26e9ae24ec530c4 100644 (file)
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -3693,6 +3693,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
     auto res = create_response();
     auto completion_id = gen_chatcmplid();
     auto & rd = res->rd;
+    auto & params = this->params;
 
     try {
         std::vector<server_task> tasks;
@@ -3828,7 +3829,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
         }
         res->status = 200;
         res->content_type = "text/event-stream";
-        res->next = [res_this = res.get(), res_type, &req](std::string & output) -> bool {
+        res->next = [res_this = res.get(), res_type, &req, &params](std::string & output) -> bool {
             static auto format_error = [](task_response_type res_type, const json & res_json) {
                 if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) {
                     return format_anthropic_sse({
@@ -3873,7 +3874,25 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 }
 
                 // receive subsequent results
-                auto result = rd.next(req.should_stop);
+                bool timeout = false;
+                int64_t start_time = ggml_time_ms();
+                auto result = rd.next([&timeout, &req, &start_time, &params]() {
+                    if (req.should_stop()) {
+                        return true; // should_stop condition met
+                    } else if (params.sse_ping_interval > 0 && ggml_time_ms() - start_time > (int64_t)params.sse_ping_interval * 1000) {
+                        timeout = true;
+                        return true; // timeout
+                    }
+                    return false;
+                });
+
+                if (timeout) {
+                    // some clients may time out (e.g. undici) will time out if no data is received for a while, so we need to send a ping to keep the connection alive
+                    SRV_DBG("%s", "sending SSE ping\n");
+                    output = ":\n\n";
+                    return true;
+                }
+
                 if (result == nullptr) {
                     SRV_DBG("%s", "stopping streaming due to should_stop condition\n");
                     GGML_ASSERT(req.should_stop());

diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp
index 588e1a82b1828f59f67c2f249a6a7fd2b4c69174..32cfe7830c353f6b2b7c5b10d89d29411301950c 100644 (file)
--- a/tools/server/server-queue.cpp
+++ b/tools/server/server-queue.cpp
@@ -381,10 +381,6 @@ server_task_result_ptr server_response_reader::next(const std::function<bool()>
         if (result == nullptr) {
             // timeout, check stop condition
             if (should_stop()) {
-                const int64_t time_elapsed_ms = ggml_time_ms() - time_start_ms;
-                if (time_elapsed_ms > 30000) {
-                    SRV_WRN("%s", "request cancelled after 30s, potentially a client-side timeout; please check your client's code\n");
-                }
                 return nullptr;
             }
         } else {

diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h
index 8ce32c69fb0454e656a38e8dfb423e361c23e1d0..35f010401fcbdfc820d99acd6924e6fde86020c9 100644 (file)
--- a/tools/server/server-queue.h
+++ b/tools/server/server-queue.h
@@ -169,8 +169,6 @@ struct server_response_reader {
     bool cancelled = false;
     int polling_interval_seconds;
 
-    const int64_t time_start_ms = ggml_time_ms();
-
     // tracking generation state and partial tool calls
     // only used by streaming completions
     std::vector<task_result_state> states;