server: cancel resumable streams before first token

tools/server/server-context.cpp

@@ -4298,10 +4298,12 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
         // in streaming mode, the first error must be treated as non-stream response
         // this is to match the OAI API behavior
         // ref: https://github.com/ggml-org/llama.cpp/pull/16486#discussion_r2419657309
-        auto first_result = rd.next(req.should_stop);
+        auto first_result = rd.next([res_this = res.get()]() {
+            return res_this->should_stop();
+        });
         if (first_result == nullptr) {
-            GGML_ASSERT(req.should_stop());
-            return res; // connection is closed
+            GGML_ASSERT(res->should_stop());
+            return res; // connection is closed or the resumable stream was explicitly stopped
         }
 
         if (first_result->is_error()) {

tools/server/tests/unit/test_stream.py

@@ -98,6 +98,50 @@ def fire_post():
     assert res.status_code == 404
 
 
+def test_stream_stop_during_prompt_processing():
+    global server
+    server.n_ctx = 8192
+    server.start()
+
+    thread_error: list[ServerError] = []
+    thread_done = threading.Event()
+
+    def fire_post():
+        try:
+            for _ in server.make_stream_request("POST", "/chat/completions", data={
+                "model": MODEL,
+                "stream": True,
+                "max_tokens": 512,
+                "messages": [{"role": "user", "content": "hello " * 7000}],
+            }, headers={"X-Conversation-Id": STREAM_ID}):
+                pass
+        except ServerError as e:
+            thread_error.append(e)
+        finally:
+            thread_done.set()
+
+    t = threading.Thread(target=fire_post)
+    t.start()
+
+    # Wait until the child has installed the stream session. This happens before the first
+    # result, while the request can still be tokenizing or processing its prompt.
+    deadline = time.time() + 60.0
+    while time.time() < deadline:
+        res = server.make_request("POST", "/v1/streams/lookup", data={
+            "conversation_ids": [STREAM_ID],
+        })
+        if res.body:
+            break
+        time.sleep(0.01)
+    else:
+        pytest.fail("stream session was not installed")
+
+    res = server.make_request("DELETE", f"/v1/stream?{QS}")
+    assert res.status_code == 204
+    assert thread_done.wait(timeout=10), "generation continued after explicit stop"
+    t.join()
+
+
 def test_stream_resumes_after_reload_during_model_load():
     global server
     server.start()