nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 8 Packages Projects Releases Wiki Activity

server + ui: SSE Replay Buffer (#23226)

Browse Source 
* server: SSE replay buffer, survives client disconnect

Opt in on POST /v1/chat/completions when the client sends
X-Stream-Resume: 1 and a non empty X-Conversation-Id. The conv id is
the session identity end to end, no extra opaque token. The drain
runs detached server side and buffers SSE bytes, the generation
survives HTTP disconnect, F5, or lets users switch from iOS Safari
to another app without losing the actively generated response.

Routes:
  GET    /v1/stream/<conv_id>?from=N       replay
  GET    /v1/streams[?conversation_id=X]   list, drives sidebar spinners
  DELETE /v1/stream/<conv_id>              Stop, idempotent

Router parent fans out to children for list and delete, probes on GET
to route to the owner, fans out DELETE on POST so "one session per
conv" holds across model swaps.

WebUI: the layout snapshots /v1/streams at mount and on
visibilitychange, the sidebar reflects live inferences across all
convs. The chat page reattaches on mount, append vs fresh is detected
from existing content so continue mid stream keeps its prefix.

update_slots: on llama_memory_seq_rm refusal at a deep position, full
clear of the seq and reprefill from zero instead of GGML_ABORT.

OAI strict path unchanged when the opt in headers are absent.

* server: create stream session only after post_tasks succeeds

* server, ui: drop X-Stream-Resume, X-Conversation-Id alone enables the replay buffer

* server: drop magic 17, derive the X-Conversation-Id header length from sizeof at build time

* refactor: address review feedback from ngxson

* server-context: cleaning

* server-stream: fix use-after-free on rd

Guard stop_producer with a shared alive flag, flipped by on_stream_end
before rd dies. Prevents a late cancel (session eviction by a later
POST on the same conv_id, or a DELETE arriving after the producer
ended) from touching a destroyed rd.

* ui: fix cross-conversation contamination

Scope streaming flags per conv so one finishing does not unflag the
others, guard discoverActiveStream against concurrent runs to avoid
duplicate attaches, and stop racing syncRemoteRunningStreams for the
sidebar set.

* server-http: keep request alive in detached SSE drain

The response next() lambda may reach into *request via &req long
after on_complete reset the request shared_ptr. Capture request in
the detached thread so it outlives the drain.

* ui: address review feedback from coder543

Forward Authorization to /v1/stream and /v1/streams fetches, the resumable routes
must obey --api-key like the rest of the API.

Wrap reader.read() in a try/catch, the underlying connection drop rejects with
TypeError instead of resolving done=true, treat it as a premature end of stream
so the existing resume loop kicks in.

Freeze the model at session start in chatStreamingStates.model and thread it
through cancel and resume, the dropdown selection may have changed since the
POST and the server side identity is fixed at that time.

* format

* ui: remove unused selectedModelName

* server-stream: poll session->is_cancelled() in stream_aware_should_stop

Address review feedback from coder543. The cancel propagation through
rd.stop() relies on the slot eventually processing the cancel task and
posting a result that notifies the recv condvar, remove_waiting_task_ids
does not notify directly. Add a defensive poll on session->is_cancelled()
so the producer-side next() loop exits on its next iteration after
cancel() without waiting for the cancel task to round trip through a slot.

* server-stream, ui: replace GET /v1/streams with POST /v1/streams/lookup

Address review feedback from coder543. Listing live sessions leaks the
conversation_id of every concurrent user, which defeats the random UUID
unguessability. The new route takes {conversation_ids: [...]} in the
body and returns matches only for the ids the caller already owns, so
foreign UUIDs stay private. The router fans out the same POST to every
child and aggregates, the WebUI passes the convs visible in its sidebar.

* ui: read conv ids from IndexedDB in syncRemoteRunningStreams

The conversations store is not hydrated yet at +layout onMount, so the
sidebar spinners stayed off for background convs until the user clicked
on them. Read straight from the DB to dodge the init race.

* server-models: deduplicate stream lookup timeouts behind one constant

* ui: extract visibility kick grace into a stream constant, bump to 1000 ms

* make it safer & more simple

* server-stream: survive client disconnect via stream_pipe::finish_producer

After the RAII rewrite the generation stopped the moment the client
disconnected. httplib bails its content provider on the is_peer_alive
check at the top of write_content_chunked, so returning true from the
provider never keeps it producing: the response resets, rd is destroyed
and its task gets cancelled.

Reinstate the disconnect survival inside the pipe. stream_pipe gains
finish_producer, which pumps the response next() into the ring buffer
until the generation ends, and mark_producer_done for the clean wire
end. server-http only triggers them: mark before sink.done on a clean
close, finish in on_complete when the peer left early. No detach, no
stream logic in server-http beyond the trigger, and the strict OAI path
is untouched when no pipe is attached.

Known limitation: finish_producer pumps synchronously on the http
worker, so a disconnected stream keeps its worker busy until the
generation ends. A follow-up will move the drain off the http worker so
no worker is held.

* server-stream: drain disconnected streams on a manager owned thread

The previous commit pumped the post disconnect drain synchronously in
on_complete, on the http worker, so a disconnected stream kept its
worker busy until the generation ended. Under a wave of reloads or tab
closes that pins workers from the pool.

Move the drain off the http worker. on_complete now hands the response
to stream_session_manager::adopt_orphan, which pumps it to completion on
a manager owned thread and releases the worker at once. One thread per
disconnected stream still generating, stored in a list, joined and
reaped on the next adopt, by the GC, and at shutdown. No detach, the
thread lifecycle is fully owned by the manager. needs_drain gates the
handoff so a cleanly finished stream never spawns a thread, and the
strict OAI path stays untouched when no pipe is attached.

stop_gc now cancels sessions before finalizing them, so an in flight
drain sees is_cancelled and exits instead of blocking the shutdown join
until the generation ends naturally.

* ui: add missing JSDoc

* server-stream: drain on the http worker, drop the manager thread

Address @ngxson review: httplib runs a large dynamic pool and a worker
blocked in next() sits on a condvar instead of burning cpu, so draining
the rest of the generation on that worker is fine and much simpler than
a dedicated thread.

on_complete calls finish_producer directly again. Removes adopt_orphan,
the orphan thread list and its reaping, the stop_gc session cancel that
only existed to unblock those threads, and the now dead drain_shutdown
flag.

* server-stream: split stream_pipe into producer and consumer classes

Address @ngxson review: one class covering both ends was messy. stream_pipe
is now a base holding the session and is_cancelled, with stream_pipe_producer
(write, mark_producer_done, finish_producer, cleanup, finalizes on destruct)
and stream_pipe_consumer (read only, no finalize) deriving from it.

Drops the is_producer_ discriminator and its runtime guards, the type now
encodes the role. res.spipe is retyped to shared_ptr<stream_pipe_producer>
since it is only ever a producer. No behavior change.

* server-stream: rename producer methods to unix pipe semantics

Address @ngxson review: mark_producer_done becomes done(), finish_producer
becomes close(), matching a unix pipe write end. The producer_done_ member
follows as done_. write() is unchanged. No behavior change.

* server, ui: route resumable streams via a conv map, persist resume identity

Address ngxson review: drop the polling probe, proxy_post records a conv_id ->
model map and the stream routes resolve the owning child with one lookup. The
map is the single source of truth, the ::model suffix stays for child session
uniqueness but the router never parses it.

UI: the server keys a session by the POST time identity (conv::model), but reload
probed with the bare conv id and missed model tagged sessions, so F5 stopped the
stream and sidebar spinners stayed off. Persist the model and rebuild the exact
identity on resume, single conv and bulk sidebar both send it.

Add unit coverage for the identity round trip.

* ui: resolve continue target by id to stop cross-conversation flash on switch

* ui: skip stream resume when the abort is intentional

* server: move the conv id to model map into a self contained tracker

Address review from ngxson: server_models held two mutexes side by side, the
global one and a bare conv_model_mu guarding a loose map, which made the locking
hard to follow. Wrap the map and its lock in a small conv_model_tracker struct
that owns its mutex, one mutex per struct. The remember, lookup and forget
methods move inline into the tracker, server_models exposes a single conv_models
member and the routes call models.conv_models.lookup and friends. No behavior
change, the map stays the single source of truth for routing resumable streams
to a child.

* ui: replace stream magic values with enums and shared constants

Address review from allozaur: lift the inline literals around the resumable
stream code into named symbols so the intent is explicit and reusable.

* ui: fold the stream resume and discovery helpers into ChatService

Address review from allozaur: drop the two standalone stream-*.service files.
They were used only by the chat service and store, carried no shared state, and
did not follow the static class pattern the other services use, so a separate
abstraction was not warranted. Move the helpers onto ChatService as static
methods. No behavior change, tests now exercise them through ChatService.

* docs: document the SSE replay buffer in server README-dev

Add the resumable streaming section, list stream_session_manager in the
backend component inventory, and link PR 23226 in the related PRs.

* ui: align attachServerStream call with onCompletionId param in handleStreamResponse

* server-http: rename del_ to del to match get and post

* ui: address review feedback from allozaur

* ui: drop duplicate SSE constants, keep sse.ts canonical

* ui: use svelte:document for the visibilitychange listener

address review from allozaur: replace the manual document.addEventListener
in onMount with a declarative <svelte:document onvisibilitychange>. svelte
handles attach, detach and SSR, so the typeof document guard and the onMount
cleanup go away. onMount keeps only the first load snapshot.

* server: trim redundant stream drain comments

Address review from ngxson

* server: balance and clean up stream comments

remove redundant comments and tighten the verbose ones across the resumable
stream code, keeping the concurrency and lifetime rationale that is not obvious
from the code. also fix two stale comments in server.cpp and server-models.h
that still described the old ::model suffix probe and fan out routing, now
replaced by the conv_id -> model map

Address review from ngxson

* ui: balance and clean up stream comments

dedup repeated rationale (frozen conv::model identity, the lookup privacy note,
the abort patterns) down to one canonical spot, tighten the verbose blocks, and
keep the concurrency and resume-offset reasoning. fix stale comments in
stream-identity.ts and chat.service.ts that still described the old loopback
probe and fan out routing, now the conv_id -> model map.

---------

Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This commit is contained in:
Pascal
2026-06-26 09:31:29 +02:00
committed by GitHub
co-authored by Xuan Son Nguyen
parent e7e3f35090
commit 1a87dcdc45
33 changed files with 2374 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/CMakeLists.txt
+2
View File
@@ -15,6 +15,8 @@ add_library(${TARGET} STATIC
server-common.h
server-context.cpp
server-context.h
server-stream.cpp
server-stream.h
server-tools.cpp
server-tools.h
server-schema.cpp
tools/server/README-dev.md
+54
View File
@@ -57,6 +57,7 @@ The core architecture consists of the following components:
- `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`.
- `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation.
- `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`.
- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
```mermaid
graph TD
@@ -117,6 +118,58 @@ Here is an example trace of an API request for text completion:
- As the response is stateless, `server_res_generator` calls `response->update()` to update the response with the current state.
- `server_res_generator` then calls `response->to_json()` and passes the response to the HTTP layer.
### Resumable streaming (SSE replay buffer)
By default a streaming generation is bound to its HTTP socket: when the socket drops (refresh, tab close, mobile background, transient network) the generation aborts and the live stream is lost. This feature keeps the generation running server side and lets a client reattach.
It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`. Without the header the OAI strict path is unchanged. The conversation id is the only identity end to end (server map key, client localStorage key, route path), with an optional `::model` suffix for direct routing in router mode.
The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL.
- `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.
Consumer side: `GET /v1/stream/<conv_id>?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400.
Routes:
- `GET /v1/stream/:conv_id?from=N`: replay or live reattach.
- `POST /v1/streams/lookup` with `{"conversation_ids": [...]}`: returns session status only for ids the caller already owns. There is no listing route, so live sessions cannot be enumerated (an earlier `GET /v1/streams` was removed for exactly this reason).
- `DELETE /v1/stream/:conv_id`: explicit Stop, idempotent (`evict_and_cancel`).
Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.
Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
| Constant | Value | Role |
| --- | --- | --- |
| `STREAM_SESSION_TTL_SECONDS` | 300 | retention of a completed session before GC |
| `STREAM_SESSION_MAX_BYTES` | 4 MiB | ring cap per session |
| `STREAM_SESSION_GC_INTERVAL_SECONDS` | 60 | GC tick |
| `STREAM_READ_WAKE_INTERVAL_MS` | 200 | read_from wake to recheck should_stop |
| `STREAM_LOOKUP_TIMEOUT_MS` | 250 | router to child loopback budget |
```mermaid
graph TD
Client -- "POST + X-Conversation-Id" --> RG[server_res_generator]
RG -- attach --> Prod[stream_pipe_producer]
Prod -- "write, drain on peer drop" --> Sess
subgraph g_stream_sessions
Sess[stream_session: ring buffer, 4 MiB]
GC[GC thread] -- drop after TTL --> Sess
end
Sess -- read_from offset --> Cons[stream_pipe_consumer]
Cons -- "GET /v1/stream/:id?from=N" --> Client
DEL[DELETE /v1/stream/:id] -- evict_and_cancel --> Sess
```
The diagram shows the buffer touch points. The live wire (chunks streamed to the original client during a normal generation) is the producer's default output, described under "Producer side" above.
### Testing
`llama-server` includes an automated test suite based on `pytest`.
@@ -223,6 +276,7 @@ The flow for downloading a new model:
- Speculative decoding: https://github.com/ggml-org/llama.cpp/pull/17808 and rework in https://github.com/ggml-org/llama.cpp/pull/17808
- INI presets: https://github.com/ggml-org/llama.cpp/pull/17859 (+ refactoring: https://github.com/ggml-org/llama.cpp/pull/18169)
- Sleeping mode: https://github.com/ggml-org/llama.cpp/pull/18228
- Resumable streaming (SSE replay buffer): https://github.com/ggml-org/llama.cpp/pull/23226
tools/server/server-context.cpp
+20 -4
View File
@@ -5,6 +5,7 @@
#include "server-task.h"
#include "server-queue.h"
#include "server-schema.h"
#include "server-stream.h"
#include "build-info.h"
#include "common.h"
@@ -4022,6 +4023,15 @@ struct server_res_generator : server_http_res {
queue_tasks.wait_until_no_sleep();
}
}
~server_res_generator() override {
// cleanup() must run while rd is still alive (rd is destroyed after this body returns)
if (spipe) {
spipe->cleanup();
}
}
void stop() override {
rd.stop();
}
void ok(const json & response_data) {
status = 200;
data = safe_json_to_str(response_data);
@@ -4210,8 +4220,10 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
}
};
auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop);
try {
if (req.should_stop()) {
if (effective_should_stop()) {
SRV_DBG("%s", "stopping streaming due to should_stop condition\n");
return false; // should_stop condition met
}
@@ -4245,8 +4257,8 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
// receive subsequent results
bool timeout = false;
int64_t start_time = ggml_time_ms();
auto result = rd.next([&timeout, &req, &start_time, &params]() {
if (req.should_stop()) {
auto result = rd.next([&timeout, &start_time, &params, &effective_should_stop]() {
if (effective_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;
@@ -4264,7 +4276,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
if (result == nullptr) {
SRV_DBG("%s", "stopping streaming due to should_stop condition\n");
GGML_ASSERT(req.should_stop());
GGML_ASSERT(effective_should_stop());
return false; // should_stop condition met
}
@@ -4302,6 +4314,10 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
};
}
// attach a producer pipe to the response when X-Conversation-Id is present.
// the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
stream_session_attach_pipe(*res, req.headers);
return res;
}
tools/server/server-http.cpp
+49 -6
View File
@@ -1,5 +1,6 @@
#include "common.h"
#include "server-http.h"
#include "server-stream.h"
#include "server-common.h"
#include "ui.h"
@@ -456,13 +457,40 @@ static void set_headers(httplib::Response & res, const std::map<std::string, std
}
}
// percent-decode a path component (%XX). path params arrive raw from httplib, unlike query
// params, so a conv id like "conv::model" sent as "conv%3A%3Amodel" must be decoded here to
// match the value the client put in the X-Conversation-Id header
static std::string decode_path_component(const std::string & in) {
std::string out;
out.reserve(in.size());
for (size_t i = 0; i < in.size(); i++) {
if (in[i] == '%' && i + 2 < in.size()) {
auto hex = [](char c) -> int {
if (c >= '0' && c <= '9') return c - '0';
if (c >= 'a' && c <= 'f') return c - 'a' + 10;
if (c >= 'A' && c <= 'F') return c - 'A' + 10;
return -1;
};
int hi = hex(in[i + 1]);
int lo = hex(in[i + 2]);
if (hi >= 0 && lo >= 0) {
out.push_back(char((hi << 4) | lo));
i += 2;
continue;
}
}
out.push_back(in[i]);
}
return out;
}
static std::map<std::string, std::string> get_params(const httplib::Request & req) {
std::map<std::string, std::string> params;
for (const auto & [key, value] : req.params) {
params[key] = value;
}
for (const auto & [key, value] : req.path_params) {
params[key] = value;
params[key] = decode_path_component(value);
}
return params;
}
@@ -497,26 +525,41 @@ static void process_handler_response(server_http_req_ptr && request, server_http
set_headers(res, response->headers);
const std::string content_type = response->content_type;
// convert to shared_ptr as both chunked_content_provider() and on_complete() need to use it
std::shared_ptr q_ptr = std::move(request);
std::shared_ptr r_ptr = std::move(response);
const auto chunked_content_provider = [response = r_ptr](size_t, const httplib::DataSink & sink) -> bool {
std::shared_ptr<server_http_req> q_ptr = std::move(request);
std::shared_ptr<server_http_res> r_ptr = std::move(response);
const auto chunked_content_provider = [response = r_ptr](size_t, httplib::DataSink & sink) -> bool {
std::string chunk;
const bool has_next = response->next(chunk);
if (!chunk.empty()) {
// mirror into the ring buffer first, the session must reflect every SSE chunk
// whether or not the wire write below succeeds
if (response->spipe) {
response->spipe->write(chunk.data(), chunk.size());
}
if (!sink.write(chunk.data(), chunk.size())) {
// peer is gone, stop the wire path here
return false;
}
SRV_DBG("http: streamed chunk: %s\n", chunk.c_str());
}
if (!has_next) {
// producer reached its natural end on the wire, a later close() skips the drain
if (response->spipe) {
response->spipe->done();
}
sink.done();
SRV_DBG("%s", "http: stream ended\n");
}
return has_next;
};
const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable {
response.reset(); // trigger the destruction of the response object
request.reset(); // trigger the destruction of the request object
// on a dropped peer, close() drains the rest of the generation into the ring buffer
if (response->spipe) {
response->spipe->close();
}
response.reset(); // spipe destructor finalizes the session if attached
request.reset();
};
res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete);
} else {
tools/server/server-http.h
+11 -1
View File
@@ -3,6 +3,7 @@
#include <atomic>
#include <functional>
#include <map>
#include <memory>
#include <string>
#include <thread>
#include <vector>
@@ -10,6 +11,7 @@
#include <unordered_map>
struct common_params;
struct stream_pipe_producer; // defined in server-stream.h
// generator-like API for HTTP response generation
// this object response with one of the 2 modes:
@@ -23,12 +25,20 @@ struct server_http_res {
std::string data;
std::map<std::string, std::string> headers;
// TODO: move this to a virtual function once we have proper polymorphism support
// if set, the stream survives a client disconnect: the producer pipe keeps draining into the
// ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed.
// shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here.
std::shared_ptr<stream_pipe_producer> spipe;
std::function<bool(std::string &)> next = nullptr;
bool is_stream() const {
return next != nullptr;
}
// called when the session is cancelled (e.g. DELETE /v1/stream/<conv_id>).
// server_res_generator overrides this to stop its reader; the default is a no-op.
virtual void stop() {}
virtual ~server_http_res() = default;
};
tools/server/server-models.cpp
+172
View File
@@ -1,12 +1,14 @@
#include "server-common.h"
#include "server-models.h"
#include "server-context.h"
#include "server-stream.h"
#include "build-info.h"
#include "preset.h"
#include "download.h"
#include <cpp-httplib/httplib.h> // TODO: remove this once we use HTTP client from download.h
#include <optional>
#include <sheredom/subprocess.h>
#include <functional>
@@ -92,6 +94,9 @@ struct server_subproc {
}
};
// short loopback budget for the resumable stream router to child JSON calls (probe, lookup,
// delete). distinct from params.timeout_read/write which only applies to the generation proxy
static constexpr int STREAM_LOOKUP_TIMEOUT_MS = 250;
static std::filesystem::path get_server_exec_path() {
#if defined(_WIN32)
@@ -1580,6 +1585,45 @@ static bool is_autoload(const common_params & params, const server_http_req & re
}
}
// percent encode one query or path component, covers reserved chars without pulling in
// httplib::detail. used by the stream routes to forward conversation_id to children safely
static std::string encode_qs(const std::string & in) {
std::string out;
out.reserve(in.size() * 3);
for (unsigned char c : in) {
bool safe = (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9')
|| c == '-' || c == '_' || c == '.' || c == '~';
if (safe) {
out.push_back(char(c));
} else {
char buf[4];
std::snprintf(buf, sizeof(buf), "%%%02X", c);
out.append(buf, 3);
}
}
return out;
}
// resolve the child that owns a conversation's stream session via the conv_id -> model map
// populated when the POST was routed. single map lookup then a meta lookup, no polling, no
// parsing of the conv id. returns nullopt when nothing maps, the caller answers not found and
// the client recovers
static std::optional<server_model_meta> resolve_child_for_conv(
server_models & models, const std::string & conversation_id) {
if (conversation_id.empty()) {
return std::nullopt;
}
auto tracked = models.conv_models.lookup(conversation_id);
if (!tracked.has_value()) {
return std::nullopt;
}
auto meta = models.get_meta(*tracked);
if (meta.has_value() && meta->is_ready()) {
return meta;
}
return std::nullopt;
}
void server_models_routes::init_routes() {
this->get_router_props = [this](const server_http_req & req) {
std::string name = req.get_param("model");
@@ -1628,6 +1672,12 @@ void server_models_routes::init_routes() {
if (!router_validate_model(name, models, autoload, error_res)) {
return error_res;
}
// remember which child serves this conversation so the stream routes can route straight
// to it without polling, keyed on the exact conv id from the header
std::string conv_id = stream_conv_id_from_headers(req.headers);
if (!conv_id.empty()) {
models.conv_models.remember(conv_id, name);
}
return models.proxy_request(req, method, name, true); // update last usage for POST request only
};
@@ -1819,6 +1869,128 @@ void server_models_routes::init_routes() {
res_ok(res, {{"success", true}});
return res;
};
this->router_stream_get = [this](const server_http_req & req) {
// GET /v1/stream/<conv_id>?from=N. resolve the owning child from the conv_id -> model
// map, 404 when nothing maps
auto res = std::make_unique<server_http_res>();
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST));
return res;
}
std::optional<server_model_meta> owner = resolve_child_for_conv(models, conv_id);
if (!owner.has_value()) {
res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND));
return res;
}
std::string from = req.get_param("from");
std::string child_path = "/v1/stream/" + encode_qs(conv_id);
if (!from.empty()) {
child_path += "?from=" + from;
}
SRV_INF("proxying stream resume to model %s on port %d, path=%s\n",
owner->name.c_str(), owner->port, child_path.c_str());
auto proxy = std::make_unique<server_http_proxy>(
"GET",
"http",
CHILD_ADDR,
owner->port,
child_path,
req.headers,
req.body,
req.files,
req.should_stop,
params.timeout_read,
params.timeout_write);
return std::unique_ptr<server_http_res>(std::move(proxy));
};
this->router_streams_lookup = [this](const server_http_req & req) {
// POST /v1/streams/lookup. resolve each requested conv id to its owning child via the
// map, group the ids per child, and query only the children that actually own some of
// them instead of fanning out to every ready child. a child only answers for the ids
// it owns, never lists anything else
auto res = std::make_unique<server_http_res>();
std::vector<std::string> requested;
try {
json body = json::parse(req.body);
if (body.contains("conversation_ids") && body["conversation_ids"].is_array()) {
for (const auto & v : body["conversation_ids"]) {
if (v.is_string() && !v.get<std::string>().empty()) {
requested.push_back(v.get<std::string>());
}
}
}
} catch (const std::exception &) {
res_ok(res, json::array());
return res;
}
// group requested ids by the child port that owns them, drop ids that map to nothing
std::unordered_map<int, json> per_child;
for (const auto & cid : requested) {
auto owner = resolve_child_for_conv(models, cid);
if (!owner.has_value()) {
continue;
}
per_child[owner->port].push_back(cid);
}
json aggregated = json::array();
for (auto & [port, ids] : per_child) {
json child_body = {{"conversation_ids", ids}};
httplib::Client cli(CHILD_ADDR, port);
cli.set_connection_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
cli.set_read_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
auto resp = cli.Post("/v1/streams/lookup", child_body.dump(), "application/json");
if (!resp || resp->status != 200) {
continue;
}
try {
json child_arr = json::parse(resp->body);
if (!child_arr.is_array()) {
continue;
}
for (auto & entry : child_arr) {
if (entry.is_object()) {
aggregated.push_back(entry);
}
}
} catch (const std::exception &) {
continue;
}
}
res_ok(res, aggregated);
return res;
};
this->router_stream_delete = [this](const server_http_req & req) {
// DELETE /v1/stream/<conv_id>. resolve the owning child via the map and forward only to
// it, evict_and_cancel is idempotent on the child
auto res = std::make_unique<server_http_res>();
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST));
return res;
}
std::string child_path = "/v1/stream/" + encode_qs(conv_id);
auto owner = resolve_child_for_conv(models, conv_id);
if (owner.has_value()) {
httplib::Client cli(CHILD_ADDR, owner->port);
cli.set_connection_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
cli.set_read_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
auto resp = cli.Delete(child_path.c_str());
(void) resp; // best effort, 404 and network errors are equivalent to no op
}
// drop the tracking entry, the session is being torn down
models.conv_models.forget(conv_id);
res->status = 204;
res->content_type = "application/json";
return res;
};
}
tools/server/server-models.h
+50
View File
@@ -11,7 +11,10 @@
#include <condition_variable>
#include <functional>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <unordered_map>
/**
* state diagram:
@@ -126,6 +129,44 @@ private:
// if true, the next get_meta() will trigger a reload of model list
bool need_reload = false;
// conv_id -> model name that currently serves its stream session, lets the resumable stream
// routes go straight to the owning child instead of polling every one. populated when
// proxy_request forwards a POST carrying an X-Conversation-Id. best effort: a stale entry just
// makes the child answer not found and the client recovers. owns its lock, one mutex per struct
struct conv_model_tracker {
void remember(const std::string & conv_id, const std::string & model) {
if (conv_id.empty() || model.empty()) {
return;
}
std::lock_guard<std::mutex> lock(mu);
map[conv_id] = model;
}
std::optional<std::string> lookup(const std::string & conv_id) {
if (conv_id.empty()) {
return std::nullopt;
}
std::lock_guard<std::mutex> lock(mu);
auto it = map.find(conv_id);
if (it == map.end()) {
return std::nullopt;
}
return it->second;
}
void forget(const std::string & conv_id) {
if (conv_id.empty()) {
return;
}
std::lock_guard<std::mutex> lock(mu);
map.erase(conv_id);
}
private:
std::mutex mu;
std::unordered_map<std::string, std::string> map;
};
common_preset_context ctx_preset;
common_params base_params;
@@ -145,6 +186,9 @@ private:
void notify_sse(const std::string & event, const std::string & model_id, const json & data = nullptr);
public:
// conv_id -> model tracker for the resumable stream routes, owns its lock
conv_model_tracker conv_models;
server_models(const common_params & params, int argc, char ** argv);
server_response sse; // for real-time updates via SSE endpoint
@@ -268,6 +312,12 @@ struct server_models_routes {
server_http_context::handler_t get_router_models_sse;
server_http_context::handler_t post_router_models;
server_http_context::handler_t del_router_models;
// router side handlers for the resumable streaming routes. each resolves the child that owns
// a conversation through the conv_id -> model map, no probing or fan out
server_http_context::handler_t router_stream_get;
server_http_context::handler_t router_streams_lookup;
server_http_context::handler_t router_stream_delete;
};
/**
tools/server/server-stream.cpp
+569
View File
@@ -0,0 +1,569 @@
#include "server-stream.h"
#include "server-common.h"
#include "server-http.h"
#include "server-queue.h"
#include <chrono>
#include <memory>
#include <utility>
namespace {
constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300;
constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024;
constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;
constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200;
// returns unix time in seconds
int64_t now_seconds() {
return std::chrono::duration_cast<std::chrono::seconds>(
std::chrono::system_clock::now().time_since_epoch()
).count();
}
}
stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
: conversation_id(std::move(conversation_id_))
, started_ts(now_seconds())
, prefix_dropped(0)
, cap_bytes(max_bytes_)
, done(false)
, cancelled(false)
, completed_ts(0) {
buffer.reserve(64 * 1024);
}
bool stream_session::append(const char * data, size_t len) {
if (len == 0) {
return true;
}
{
std::lock_guard<std::mutex> lock(mu);
if (done.load(std::memory_order_relaxed)) {
return false;
}
if (len >= cap_bytes) {
// single chunk bigger than the cap, keep only the tail that fits
size_t skip = len - cap_bytes;
prefix_dropped += buffer.size() + skip;
buffer.clear();
buffer.insert(buffer.end(), data + skip, data + len);
} else {
size_t needed = buffer.size() + len;
if (needed > cap_bytes) {
size_t to_drop = needed - cap_bytes;
buffer.erase(buffer.begin(), buffer.begin() + to_drop);
prefix_dropped += to_drop;
}
buffer.insert(buffer.end(), data, data + len);
}
}
cv.notify_all();
return true;
}
void stream_session::finalize() {
bool was_done = done.exchange(true, std::memory_order_acq_rel);
if (was_done) {
return;
}
completed_ts.store(now_seconds(), std::memory_order_release);
cv.notify_all();
}
stream_read_status stream_session::read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop) {
std::unique_lock<std::mutex> lock(mu);
while (true) {
if (should_stop && should_stop()) {
return stream_read_status::OK;
}
if (offset < prefix_dropped) {
return stream_read_status::OFFSET_LOST;
}
size_t logical_end = prefix_dropped + buffer.size();
if (offset < logical_end) {
size_t local_off = offset - prefix_dropped;
size_t n = buffer.size() - local_off;
// copy the available chunk under the lock, release before calling the sink
std::vector<char> chunk(buffer.begin() + local_off, buffer.begin() + local_off + n);
offset += n;
lock.unlock();
bool keep_going = sink(chunk.data(), chunk.size());
if (!keep_going) {
return stream_read_status::OK;
}
lock.lock();
continue;
}
if (done.load(std::memory_order_acquire)) {
return stream_read_status::OK;
}
// wait for new bytes, finalize, or a periodic wake to re check should_stop
cv.wait_for(lock, std::chrono::milliseconds(STREAM_READ_WAKE_INTERVAL_MS));
}
}
bool stream_session::is_done() const {
return done.load(std::memory_order_acquire);
}
size_t stream_session::total_size() const {
std::lock_guard<std::mutex> lock(mu);
return prefix_dropped + buffer.size();
}
size_t stream_session::dropped_prefix() const {
std::lock_guard<std::mutex> lock(mu);
return prefix_dropped;
}
int64_t stream_session::completed_at() const {
return completed_ts.load(std::memory_order_acquire);
}
void stream_session::set_stop_producer(std::function<void()> fn) {
std::lock_guard<std::mutex> lock(mu);
stop_producer = std::move(fn);
}
void stream_session::cancel() {
// flip cancelled first so the producer-side stream_aware_should_stop can break out of the
// recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
// posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
cancelled.store(true, std::memory_order_release);
// copy the hook under the lock then invoke outside, the producer side may grab queue locks
// and we do not want to hold our mu across that path
std::function<void()> fn;
{
std::lock_guard<std::mutex> lock(mu);
fn = stop_producer;
}
if (fn) {
fn();
}
}
bool stream_session::is_cancelled() const {
return cancelled.load(std::memory_order_acquire);
}
stream_session_manager::stream_session_manager()
: running(false) {
}
stream_session_manager::~stream_session_manager() {
stop_gc();
}
stream_session_ptr stream_session_manager::create_or_replace(const std::string & conversation_id) {
// evict any previous session on the same conv, this guarantees the invariant
// "one conv = at most one live session" and propagates cancel to its producer
stream_session_ptr previous;
auto fresh = std::make_shared<stream_session>(conversation_id, STREAM_SESSION_MAX_BYTES);
{
std::unique_lock<std::shared_mutex> lock(map_mu);
auto it = sessions.find(conversation_id);
if (it != sessions.end()) {
previous = it->second;
it->second = fresh;
} else {
sessions.emplace(conversation_id, fresh);
}
}
if (previous) {
previous->cancel();
previous->finalize();
}
return fresh;
}
stream_session_ptr stream_session_manager::get(const std::string & conversation_id) {
std::shared_lock<std::shared_mutex> lock(map_mu);
auto it = sessions.find(conversation_id);
if (it == sessions.end()) {
return nullptr;
}
return it->second;
}
std::vector<stream_session_ptr> stream_session_manager::list_all() const {
std::vector<stream_session_ptr> out;
std::shared_lock<std::shared_mutex> lock(map_mu);
out.reserve(sessions.size());
for (auto & kv : sessions) {
out.push_back(kv.second);
}
return out;
}
void stream_session_manager::evict(const std::string & conversation_id) {
stream_session_ptr s;
{
std::unique_lock<std::shared_mutex> lock(map_mu);
auto it = sessions.find(conversation_id);
if (it == sessions.end()) {
return;
}
s = it->second;
sessions.erase(it);
}
// finalize outside the map lock so any pending readers wake up and exit
s->finalize();
}
void stream_session_manager::evict_and_cancel(const std::string & conversation_id) {
stream_session_ptr s;
{
std::unique_lock<std::shared_mutex> lock(map_mu);
auto it = sessions.find(conversation_id);
if (it == sessions.end()) {
return;
}
s = it->second;
sessions.erase(it);
}
// signal the producer side first so the inference is cancelled at the queue level,
// then finalize, which wakes any pending HTTP reader and lets the drain exit naturally
s->cancel();
s->finalize();
}
void stream_session_manager::start_gc() {
if (running.exchange(true)) {
return;
}
gc_thread = std::thread([this] { gc_loop(); });
}
void stream_session_manager::stop_gc() {
bool was_running = running.exchange(false);
if (was_running) {
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
}
gc_wake_cv.notify_all();
if (gc_thread.joinable()) {
gc_thread.join();
}
}
// finalize all live sessions so no reader ever hangs
std::vector<stream_session_ptr> snapshot;
{
std::unique_lock<std::shared_mutex> lock(map_mu);
snapshot.reserve(sessions.size());
for (auto & kv : sessions) {
snapshot.push_back(kv.second);
}
sessions.clear();
}
for (auto & s : snapshot) {
s->finalize();
}
}
void stream_session_manager::gc_loop() {
while (running.load(std::memory_order_acquire)) {
{
std::unique_lock<std::mutex> lock(gc_wake_mu);
gc_wake_cv.wait_for(lock,
std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),
[this] { return !running.load(std::memory_order_acquire); });
}
if (!running.load(std::memory_order_acquire)) {
return;
}
int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;
std::vector<stream_session_ptr> to_drop;
{
std::unique_lock<std::shared_mutex> lock(map_mu);
for (auto it = sessions.begin(); it != sessions.end(); ) {
int64_t completed = it->second->completed_at();
if (completed != 0 && completed <= cutoff) {
to_drop.push_back(it->second);
it = sessions.erase(it);
} else {
++it;
}
}
}
// finalize outside the map lock, idempotent if the session was already done
for (auto & s : to_drop) {
s->finalize();
}
}
}
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
stream_session_manager g_stream_sessions;
// stream_pipe ---------------------------------------------------------------------------------
stream_pipe::stream_pipe(stream_session_ptr session)
: session_(std::move(session)) {
}
bool stream_pipe::is_cancelled() const {
return session_->is_cancelled();
}
// stream_pipe_producer
stream_pipe_producer::stream_pipe_producer(stream_session_ptr session)
: stream_pipe(std::move(session)) {
}
stream_pipe_producer::~stream_pipe_producer() {
cleanup();
session_->finalize();
}
void stream_pipe_producer::cleanup() {
if (!alive_) {
return;
}
alive_->store(false, std::memory_order_release);
session_->set_stop_producer(nullptr);
alive_.reset();
}
bool stream_pipe_producer::write(const char * data, size_t len) {
return session_->append(data, len);
}
void stream_pipe_producer::done() {
done_ = true;
}
void stream_pipe_producer::close() {
// httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest
// of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short
if (done_ || session_->is_cancelled()) {
SRV_INF("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n",
done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str());
return;
}
SRV_INF("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str());
size_t drained = 0;
std::string chunk;
while (true) {
chunk.clear();
bool has_next = res_->next(chunk);
if (!chunk.empty()) {
write(chunk.data(), chunk.size());
drained += chunk.size();
}
if (!has_next) {
break;
}
}
SRV_INF("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained);
}
std::shared_ptr<stream_pipe_producer> stream_pipe_producer::create(stream_session_ptr session,
server_http_res & res) {
auto alive = std::make_shared<std::atomic<bool>>(true);
auto * res_ptr = &res;
session->set_stop_producer([alive, res_ptr]() {
if (alive->load(std::memory_order_acquire)) {
res_ptr->stop();
}
});
auto pipe = std::shared_ptr<stream_pipe_producer>(new stream_pipe_producer(std::move(session)));
pipe->alive_ = std::move(alive);
pipe->res_ = res_ptr;
return pipe;
}
// stream_pipe_consumer
stream_pipe_consumer::stream_pipe_consumer(stream_session_ptr session)
: stream_pipe(std::move(session)) {
}
stream_read_status stream_pipe_consumer::read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop) {
return session_->read_from(offset, sink, should_stop);
}
std::shared_ptr<stream_pipe_consumer> stream_pipe_consumer::create(stream_session_ptr session) {
return std::shared_ptr<stream_pipe_consumer>(new stream_pipe_consumer(std::move(session)));
}
// helper, builds the standard error response and assigns it to a brand new http_res
static server_http_res_ptr make_error_response(int status, const std::string & message, error_type type) {
auto res = std::make_unique<server_http_res>();
json err = format_error_response(message, type);
res->status = json_value(err, "code", status);
res->content_type = "application/json; charset=utf-8";
res->data = safe_json_to_str({{"error", err}});
return res;
}
server_http_context::handler_t make_stream_get_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// GET /v1/stream/<conv_id>?from=N replays the SSE bytes already buffered for the
// session, blocks for more bytes when the session is still running, returns when
// the session is finalized. the body is streamed back as text/event-stream so the
// browser EventSource can attach to it like a fresh request
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
}
auto session = g_stream_sessions.get(conv_id);
if (!session) {
return make_error_response(404, "Stream not found or expired", ERROR_TYPE_NOT_FOUND);
}
size_t from = 0;
std::string from_str = req.get_param("from");
if (!from_str.empty()) {
try {
from = static_cast<size_t>(std::stoull(from_str));
} catch (const std::exception &) {
return make_error_response(400, "Invalid 'from' offset", ERROR_TYPE_INVALID_REQUEST);
}
}
if (from < session->dropped_prefix()) {
return make_error_response(400, "Stream offset lost, please restart", ERROR_TYPE_INVALID_REQUEST);
}
auto res = std::make_unique<server_http_res>();
res->status = 200;
res->content_type = "text/event-stream";
// the next closure reads from the ring buffer at the requested offset, blocks until
// bytes arrive or the session finalizes. exit each call after draining the available
// chunk so set_chunked_content_provider gets a chance to flush to the socket
auto offset_ptr = std::make_shared<size_t>(from);
// consumer pipe: read-only, does not finalize the session on destruction
auto pipe = stream_pipe_consumer::create(session);
res->next = [pipe, offset_ptr, &req](std::string & output) -> bool {
bool got_any = false;
pipe->read(*offset_ptr,
[&](const char * d, size_t n) {
output.append(d, n);
*offset_ptr += n;
got_any = true;
return false;
},
req.should_stop);
return got_any;
};
return res;
};
}
server_http_context::handler_t make_streams_lookup_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the
// matching sessions, only for ids the caller already knows. each id matches the exact key
// and any "<id>::<model>" variant, so one lookup covers every per model session for a conv
std::vector<std::string> requested;
try {
json body = json::parse(req.body);
if (body.contains("conversation_ids") && body["conversation_ids"].is_array()) {
for (const auto & v : body["conversation_ids"]) {
if (v.is_string()) {
std::string id = v.get<std::string>();
if (!id.empty()) {
requested.push_back(std::move(id));
}
}
}
}
} catch (const std::exception & e) {
auto res = std::make_unique<server_http_res>();
res->status = 400;
res->content_type = "application/json; charset=utf-8";
res->data = safe_json_to_str({{"error", {{"message", std::string("invalid body: ") + e.what()},
{"type", "invalid_request_error"}}}});
return res;
}
std::vector<stream_session_ptr> sessions;
if (!requested.empty()) {
auto all = g_stream_sessions.list_all();
for (const auto & rid : requested) {
const std::string with_sep = rid + "::";
for (auto & s : all) {
if (s->conversation_id == rid ||
s->conversation_id.compare(0, with_sep.size(), with_sep) == 0) {
sessions.push_back(s);
}
}
}
}
json arr = json::array();
for (auto & s : sessions) {
arr.push_back({
{"conversation_id", s->conversation_id},
{"is_done", s->is_done()},
{"total_bytes", s->total_size()},
{"started_at", s->started_ts},
{"completed_at", s->completed_at()},
});
}
auto res = std::make_unique<server_http_res>();
res->status = 200;
res->content_type = "application/json; charset=utf-8";
res->data = safe_json_to_str(arr);
return res;
};
}
server_http_context::handler_t make_stream_delete_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer hook
// wired by handle_completions_impl and evicts the buffer. idempotent, a session that
// already finalized or was never created returns 204 either way
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
}
SRV_INF("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str());
g_stream_sessions.evict_and_cancel(conv_id);
auto res = std::make_unique<server_http_res>();
res->status = 204;
res->content_type = "application/json";
return res;
};
}
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
// case-insensitive scan for x-conversation-id
static constexpr char target[] = "x-conversation-id";
static constexpr size_t target_len = sizeof(target) - 1;
for (const auto & [hk, hv] : headers) {
if (hk.size() != target_len) continue;
bool match = true;
for (size_t i = 0; i < target_len; ++i) {
char c = hk[i];
if (c >= 'A' && c <= 'Z') c = char(c + 32);
if (c != target[i]) { match = false; break; }
}
if (match) {
return hv;
}
}
return std::string();
}
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
std::string conversation_id = stream_conv_id_from_headers(headers);
SRV_INF("stream_session_attach_pipe: conv_id=%s (empty=%d)\n",
conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
if (conversation_id.empty()) {
return;
}
auto session = g_stream_sessions.create_or_replace(conversation_id);
res.spipe = stream_pipe_producer::create(session, res);
}
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
return [res, fallback = std::move(fallback)]() -> bool {
if (res->spipe) {
return res->spipe->is_cancelled();
}
return fallback();
};
}
tools/server/server-stream.h
+203
View File
@@ -0,0 +1,203 @@
#pragma once
#include "server-http.h"
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
enum class stream_read_status {
OK,
OFFSET_LOST,
};
// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE
// bytes, readers drain from any offset via read_from and block until more bytes or finalize.
// keyed by conversation_id: one conv = at most one live session
struct stream_session {
std::string conversation_id;
int64_t started_ts; // unix seconds at construction, used by /v1/streams listing
stream_session(std::string conversation_id_, size_t max_bytes_);
stream_session(const stream_session &) = delete;
stream_session & operator=(const stream_session &) = delete;
// append raw bytes, drops from the front if the cap is reached.
// returns false if the session is already finalized
bool append(const char * data, size_t len);
// mark the session as complete, wakes all pending readers
void finalize();
// drain bytes from offset, calling sink for each chunk. blocks until more
// bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST
// if offset falls below the dropped prefix
stream_read_status read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
bool is_done() const;
bool is_cancelled() const;
size_t total_size() const; // bytes that ever entered the session
size_t dropped_prefix() const; // bytes evicted from the front due to cap
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
// attach the producer stop hook used to cancel its reader, pass an empty function to detach
void set_stop_producer(std::function<void()> fn);
// signal the producer to abort its inference asap via the stop hook, idempotent
void cancel();
private:
mutable std::mutex mu;
std::condition_variable cv;
std::vector<char> buffer;
size_t prefix_dropped;
size_t cap_bytes;
std::atomic<bool> done;
std::atomic<bool> cancelled;
std::atomic<int64_t> completed_ts;
std::function<void()> stop_producer; // protected by mu
};
using stream_session_ptr = std::shared_ptr<stream_session>;
// one end of a stream_session pipe. the base holds the session and the shared query, the
// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown:
// the producer finalizes the session, the consumer leaves it untouched
struct stream_pipe {
virtual ~stream_pipe() = default;
// true if the session was cancelled (e.g. via DELETE /v1/stream/<conv_id>)
bool is_cancelled() const;
protected:
explicit stream_pipe(stream_session_ptr session);
stream_session_ptr session_;
};
// producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it
// on destruction.
//
// lifetime safety: holds a shared_ptr<atomic<bool>> alive also captured by the session's
// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the
// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly.
struct stream_pipe_producer : stream_pipe {
~stream_pipe_producer() override;
// append raw bytes to the session's ring buffer, returns false if already finalized
bool write(const char * data, size_t len);
// mark the natural end on the wire so a later close() is a no-op
void done();
// on a peer drop, pump the response next() into the ring buffer until done. runs on the http
// worker from on_complete, no-op after done() or cancel
void close();
// disarm the stop hook and drop the alive guard, must run while the response the hook
// references is still alive. idempotent, the destructor calls it too
void cleanup();
// res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not
// called after cleanup() has run
static std::shared_ptr<stream_pipe_producer> create(stream_session_ptr session, server_http_res & res);
private:
explicit stream_pipe_producer(stream_session_ptr session);
bool done_ = false;
std::shared_ptr<std::atomic<bool>> alive_;
server_http_res * res_ = nullptr;
};
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
struct stream_pipe_consumer : stream_pipe {
// drain bytes from offset, calling sink for each available chunk. blocks until more data
// arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset
// fell below the dropped prefix
stream_read_status read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
private:
explicit stream_pipe_consumer(stream_session_ptr session);
};
// owns all live sessions, runs a periodic GC to evict expired ones.
// the map is keyed by conversation_id, so the invariant "one conv = at most one
// live session" is enforced at the type level
class stream_session_manager {
public:
stream_session_manager();
~stream_session_manager();
stream_session_manager(const stream_session_manager &) = delete;
stream_session_manager & operator=(const stream_session_manager &) = delete;
// install a new session for this conversation, evicting and cancelling any previous one.
// the conversation_id must be non empty, the caller is responsible for that check.
// returns the new session
stream_session_ptr create_or_replace(const std::string & conversation_id);
// lookup, returns null if unknown or already evicted
stream_session_ptr get(const std::string & conversation_id);
// list every live or recently completed session, used by GET /v1/streams without filter
std::vector<stream_session_ptr> list_all() const;
// remove from the map and finalize, wakes any pending readers
void evict(const std::string & conversation_id);
// signal the producer to cancel asap then evict, used by the explicit user Stop path
void evict_and_cancel(const std::string & conversation_id);
void start_gc();
void stop_gc();
private:
void gc_loop();
mutable std::shared_mutex map_mu;
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
std::thread gc_thread;
std::atomic<bool> running;
std::mutex gc_wake_mu;
std::condition_variable gc_wake_cv;
};
// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives
// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit
// is safe whether or not the GC thread ran
extern stream_session_manager g_stream_sessions;
// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp.
// keeps the resumable stream surface confined to server-stream
server_http_context::handler_t make_stream_get_handler();
server_http_context::handler_t make_streams_lookup_handler();
server_http_context::handler_t make_stream_delete_handler();
// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so
// the router can track which child serves a forwarded POST
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to
// res. no-op when absent, called from the server_res_generator constructor
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
// delegates to fallback, the legacy non-resumable flow
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
tools/server/server.cpp
+32
View File
@@ -2,6 +2,7 @@
#include "server-http.h"
#include "server-models.h"
#include "server-cors-proxy.h"
#include "server-stream.h"
#include "server-tools.h"
#include "arg.h"
@@ -82,6 +83,10 @@ int llama_server(int argc, char ** argv) {
common_init();
// start the stream session manager GC right after common init, before any HTTP route can
// touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
g_stream_sessions.start_gc();
if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
return 1;
}
@@ -239,6 +244,29 @@ int llama_server(int argc, char ** argv) {
ctx_http.get ("/slots", ex_wrapper(routes.get_slots));
ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots));
// resumable streaming, the conversation_id is the session identity end to end. router and
// child wire different handlers under the same paths: a child binds the local g_stream_sessions
// backed factories, the router binds proxies that resolve the owning child through the
// conv_id -> model map
server_http_context::handler_t stream_get_h;
server_http_context::handler_t streams_lookup_h;
server_http_context::handler_t stream_delete_h;
if (is_router_server) {
stream_get_h = models_routes->router_stream_get;
streams_lookup_h = models_routes->router_streams_lookup;
stream_delete_h = models_routes->router_stream_delete;
} else {
stream_get_h = make_stream_get_handler();
streams_lookup_h = make_streams_lookup_handler();
stream_delete_h = make_stream_delete_handler();
}
ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h));
// POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids
// you already own (the WebUI passes the convs visible in its sidebar). the server never
// lists ids it has not been asked about, so a random caller cannot enumerate live sessions
ctx_http.post("/v1/streams/lookup", ex_wrapper(streams_lookup_h));
ctx_http.del ("/v1/stream/:conv_id", ex_wrapper(stream_delete_h));
// Google Cloud Platform (Vertex AI) compat
ctx_http.register_gcp_compat();
@@ -314,6 +342,8 @@ int llama_server(int argc, char ** argv) {
clean_up = [&models_routes]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
if (models_routes.has_value()) {
models_routes->stopping.store(true); // maybe redundant, but just to be safe
models_routes->models.unload_all();
@@ -340,6 +370,8 @@ int llama_server(int argc, char ** argv) {
// setup clean up function, to be called before exit
clean_up = [&ctx_http, &ctx_server]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
ctx_http.stop();
ctx_server.terminate();
llama_backend_free();