Server enhancements + adopted upstream fixes (Metal, security, KV fidelity)#1
Merged
Conversation
Prefills resumed from an unaligned cached position (disk hits, live text continuations) never land on exact multiples of the continued-store interval, so waypoint snapshots silently stopped firing and disk recovery could re-prefill 12-20K tokens instead of <=1 interval. Fire a continued store whenever a full interval has passed since the last store, and restart the cadence from the continuation point on a cache hit so the first chunk after a hit does not store a near duplicate of the snapshot just loaded.
Prefill previously always ran to completion: a streaming client that vanished was only detected on the next SSE write after sync returned, and a non-streaming client was never checked at all, so an agent timeout or retry could burn many minutes of GPU on a dead socket, with duplicate jobs queueing behind each other. Wire ds4_session_set_cancel() around every server-side sync so prefill stops at the next chunk boundary on shutdown, on a failed keepalive write, or on a closed client socket. Interrupted syncs keep the valid session prefix and do not discard the loaded disk snapshot. Also skip queued jobs whose client already disconnected and poll the socket during non-streaming decode.
Loading a text-prefix snapshot longer than cold_max_tokens unlinked the file immediately, before the tail prefill that extends it had run. A transient sync failure (Metal OOM under memory pressure) or a cancelled prefill then left neither live state nor the snapshot, turning the next attempt into a multi-minute cold prefill. Keep the file on disk until the request's sync completes and unlink it at that point; interrupted syncs already skip the failed-entry discard, so a retry can hit the same snapshot again.
The server exposed no machine-readable operational state: every metric (cache hit source, prefill/decode t/s, queue behavior) existed only as stderr log lines, and the job queue was unbounded and invisible, so an agent with a client-side timeout shorter than a turn could silently stack duplicate jobs that each run for minutes. Add GET /health and GET /stats, answered on the client thread so they respond even mid-generation. /stats reports uptime, queue depth, busy flag, live/ctx tokens, per-source cache hit counters, token totals, cancellation counters, and last prefill/decode t/s. Add --max-queue N to reject requests with 429 once N jobs wait (default 0 keeps the old unbounded behavior), log when a request queues behind others, and give 429/503 their proper HTTP reason strings.
kv_cache_restore_tool_memory_for_messages() opened and parsed the trailer of every checkpoint file in the cache directory on every tools request, even though tool_memory_remember() already holds the replayed DSML blocks in RAM at generation time. Filter the wanted ids against the in-memory map first and only walk the directory for ids that are actually missing, which normally only happens right after a restart.
tool_memory_attach_to_messages() attached a remembered DSML block when every id in the message resolved to the same block, but never checked that the message covers all of the block's invokes. A client that splits one turn's parallel tool calls across two assistant messages (or drops one call) got the full multi-invoke block attached to each message, rendering duplicate ghost invokes into the prompt. Count the invokes per remembered block and require the replaying message to carry exactly that many calls; partial messages fall back to canonical JSON rendering. Adds a regression test.
parse_generated_message_ex() anchored the raw block at the marker match, which includes the separator only when the model sampled the canonical "\n\n". Any other separator (a single newline, spaces) was trimmed from the content and excluded from raw_dsml, so those bytes existed in live KV but in no re-rendered prompt, silently breaking byte-exact tool replay and forcing a cache round trip. Anchor the raw block at the trimmed-content boundary instead, so content + raw block always reassemble the sampled bytes exactly. Adds a regression test.
Chat/completions and Anthropic tool-call turns had no live binding to the next request: the sampled tail (hidden reasoning, exact DSML spelling, BPE drift) never token-matches the client replay, and the byte-exact memory-text path rarely matches either, so every agent turn paid a full-state evict store (~650 MiB) plus a disk snapshot restore and tail re-prefill. The Responses API already solves this with a predicted visible transcript key. Reuse the thinking_live visible-key mechanism for tool turns: after finish=tool_calls, remember prompt_text plus the same canonical suffix that canonicalize_tool_checkpoint() builds (its byte-equality with the next rendered prompt is already covered by test_tool_checkpoint_suffix_is_future_prompt_canonical). The next request then continues from live KV, tokenizing only the new tool results. Hits are labeled tool-visible in logs and /stats.
- Answer Expect: 100-continue so curl-style clients do not stall a
second before sending large POST bodies, and reject chunked
transfer encoding explicitly instead of parsing an empty body and
returning a misleading JSON error.
- Set TCP_NODELAY on client sockets; per-token SSE writes are tiny and
Nagle adds latency jitter off loopback.
- Give /v1/messages clients Anthropic-shaped error envelopes
({"type":"error",...}) for parse, queue, continuation, and
prefill failures, and classify OpenAI error types by status code
(api_error for 5xx, rate_limit_error for 429).
- Reject tool_choice "required" and forced function targets on chat
completions like the Responses endpoint already does, instead of
silently downgrading to auto.
- Log once per request when thinking mode overrides client sampling
parameters instead of ignoring them silently.
Eviction scored a never-hit entry as (0+1) * density from the moment it was written, while cold/evict/shutdown anchors kept an unconditional 2x reason bonus forever. Under budget pressure this evicted the newest waypoints of the live conversation first (they are always hits=0, they were just written) while stale anchors from long-dead conversations survived. Observed in production: the waypoint one position below a mid-history divergence was evicted seconds before it would have been hit, costing ~12K tokens of avoidable re-prefill. Score never-hit entries with a freshness grace equal to one hit that decays on the existing hit half-life, and scale the anchor bonus by the same activity term so a never-hit anchor older than the half-life competes on density alone. Adds a regression test; updates the fresh floor test to the new (1+1) * density value.
Post-review fixes for the session-state transaction: - Rollback now clears protocol live bindings (responses/anthropic/ thinking) exactly as the pre-transaction epilogue did on error finishes, so a failed turn cannot leave a binding pointing at a frontier the session no longer keeps. Regression test added (test_production_settle_rollback_clears_live_bindings). - The shutdown KV persist result is no longer ignored; failure logs a warning so an operator knows the next start will re-prefill. - Duplicate terminal publish logs loudly instead of die(): the terminalizer is idempotent and the first outcome already reached the client, so aborting a loaded server was worse than dropping the dup. - test_txn_event is compiled only under DS4_SERVER_TEST. - Spec text now matches the implementation: a queued disconnect is recorded as client gone decided in the restore phase. Tested on the M5 Max reference machine (Metal, q2-q4-imatrix quant): warning-free `make ds4-server ds4_test`; `./ds4_test --server` passes including the new settle-rollback regression test. GPU-bound suites not rerun (production server holds the device); no inference-path changes.
In the SNAPSHOT_LOAD_BEGIN worker path, token_count is only bounded to ~UINT32_MAX/4 by the header check (expected_token_bytes <= UINT32_MAX). The code then malloc()s token_count*4 bytes and runs the full token read loop, and only afterwards does the matching-state check reject token_count > ctx_size. An untrusted peer (the distributed protocol is unauthenticated) can therefore drive a multi-GB allocation and a full-length socket read before the request is rejected. Move the token_count > ctx_size test ahead of the allocation and read loop, discarding the body bytes, so an over-large count is rejected up front. (cherry picked from commit e7628ea)
agent_kv_read_text() and agent_kv_read_title_trailer() read a uint32 length from a persisted agent session/cache file and immediately xmalloc that many bytes. On 64-bit (size_t)len + 1 does not wrap, so there is no overflow, but a 1-byte file declaring 0xFFFFFFFF drives a ~4 GiB allocation before the following fread() fails -- an attacker-influenced cache file becomes a memory-pressure DoS. Add agent_fp_remaining() and reject a length larger than the bytes left in the file before allocating, at both sites. (cherry picked from commit 3c56797)
agent_bash_start() creates a per-job output file with mkstemp() and stores it in job->path. Only that function's error paths unlink it; the normal cleanup path agent_bash_job_free() closes the fds and frees memory but never unlinks job->path. A process that runs many bash-tool jobs therefore leaves a /tmp/ds4_agent_output_* file behind for each one, accumulating without bound for the lifetime of the process. Unlink job->path in agent_bash_job_free() before freeing the job. (cherry picked from commit d5d1c8f)
parse_metadata() and parse_tensors() calloc() a table sized directly from the header-declared m->n_kv / m->n_tensors, with no check against the mapped file. A 24-byte file declaring n_kv = 2^32 makes parse_metadata() request ~137 GB up front -- a load-time DoS from a tiny crafted model. Each metadata/tensor entry occupies at least one byte in the file, so a count larger than the bytes remaining (c->size - c->pos) cannot be real; reject it before the allocation. (cherry picked from commit c127fb4)
vocab_load() validates the token table with tokens.len > INT32_MAX but omits the same bound on the merges table. The merge loop then stores each rank with table_put(&vocab->merge_rank, merge, (int)i), narrowing the uint64 index to int. A GGUF declaring more than INT32_MAX merges makes (int)i wrap negative and stores corrupt ranks. Mirror the tokens sibling's bound on merges.len. (cherry picked from commit 8ce9a67)
shard_load() (safetensors header) and read_gguf_string_fp() (GGUF string) both read a uint64 length straight from a downloaded model file and then do xmalloc((size_t)len + 1) followed by fread(len). A crafted length near SIZE_MAX makes (size_t)len + 1 wrap to 0, so xmalloc() returns a 1-byte buffer (xmalloc maps 0 to malloc(1)); the subsequent fread() of the file's real bytes then overflows it -- a heap-buffer-overflow write. A large but non-wrapping length instead requests a multi-GB allocation from a tiny file. Add read_checked_len_fp(), which reads the length and rejects it if it exceeds the bytes actually remaining in the file, and use it at both sites. Valid lengths (<= file size, so no size_t wrap) are unchanged. (cherry picked from commit dab7707)
A safetensors tensor's shape and its data_offsets byte range are independent header fields. db_read() sizes the buffer (st_value.nbytes) from data_offsets, while dequant_fp8_weight() / dequant_fp4_weight() index w->data and scale->data by shape. They were never cross-validated, so a weight that declares a large shape but a tiny data_offsets range (e.g. shape [128,128] = 16384 elements with only 128 bytes of data) makes the loops read past the end of the heap buffer -- a heap-buffer-overflow read, leaking adjacent heap bytes into the produced GGUF. Before the loops, require nbytes to cover the shape-driven access: one byte per element for the F8_E4M3 / I8 weights and F8_E8M0 scales. (cherry picked from commit 8007e04)
Launching ds4 binaries from outside the repo silently falls back to the CPU backend because metal/*.metal only resolves against the current directory. Try the executable's own directory (symlinks followed) as a last candidate, so the binaries work from any working directory or a PATH symlink without --chdir. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit 5095214)
Reader threads inherit the parent QoS: launched via taskpolicy -b or a background launchd job they end up on the throttled IO tier and every expert miss pays for it. Pin them to user-initiated QoS and IOPOL_IMPORTANT; DS4_METAL_DISABLE_STREAMING_EXPERT_PREAD_QOS opts out. (cherry picked from commit 66ca6ef)
One extra line next to the streaming expert timing report (same DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY gate): dispatches, tasks, bytes, average dispatch wall time, effective queue depth (qd_avg = sum of per-task pread ms / pool wall ms) and delivered bandwidth. Makes streaming IO regressions measurable without dtrace. (cherry picked from commit a1afb82)
json_string() returns false on a non-string or malformed value without
writing *out. Several callers reparse in place with
`free(x); json_string(&p, &x)` -- the "model" field in
parse_chat_request / parse_completion_request / parse_anthropic_request,
and duplicate JSON keys for type/name/description across the tool-schema
and message parsers. On a failed reparse, *out keeps the just-freed
pointer; the request's cleanup (request_free) then frees it again.
This is reachable pre-auth with a single request, e.g.
POST /v1/chat/completions {"messages":[...],"model":0}
Initialize *out = NULL at the top of json_string so every failure path
leaves the caller's pointer defined, closing all the free-then-reparse
sites at the root rather than patching each call site. The success path
is unchanged (*out is still set by buf_take).
(cherry picked from commit 9fe8faf)
json_number() is a thin strtod() wrapper, which per C99 accepts "NaN" and "Infinity". In json_int() the clamps `v < 0` and `v > INT_MAX` are both false for NaN (every NaN comparison is false), so NaN reaches the (int)v cast, which is undefined behavior (on x86-64 it yields INT_MIN). Replace `if (v < 0)` with `if (!(v >= 0))`, which folds NaN and negatives to 0; +/-Infinity are still handled by the two clamps. (cherry picked from commit 19232c6)
responses_validate_tool_outputs() and anthropic_validate_tool_results() resolved each tool output's call_id with responses_find_prior_call_msg(), a backward scan over all preceding messages. Run once per id per tool message over an attacker-supplied, uncapped message array, that is O(n^2): ~10^6 messages is ~10^12 comparisons, pinning a CPU core pre-auth. Only assistant messages declare call ids (chat_msg_has_call_id matched assistant .calls entries), so build a call_id -> nearest-preceding assistant message map incrementally with a rax while scanning forward, and resolve each id with an O(1) lookup. Registering assistant messages as they are passed and looking up on tool messages reproduces the previous nearest-preceding-before-i result exactly. Drop the now-unused responses_find_prior_call_msg() and chat_msg_has_call_id(). Verified: the validator's accept/reject results are unchanged, and wall clock now scales ~linearly with n instead of ~4x per doubling. (cherry picked from commit 1be59d5)
When thinking mode is on and generation stops before </think> (e.g. the response is truncated at max_tokens mid-thought), parse_generated_message_ex dumped the partial reasoning into content with an empty reasoning_content, so clients received raw chain-of-thought as the assistant's answer. Route the accumulated text to reasoning_content and leave content empty, matching what the live streaming classifier already does. This composes with the existing tool-call handling in the same branch: unclosed DSML is still treated as reasoning rather than executed. Adds a regression test (test_thinking_unclosed_is_reasoning_not_content). Fixes antirez#509. (cherry picked from commit 76573b0)
(cherry picked from commit dc7f2ca)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Brings the local server-enhancements work together with a curated set of
adopted fixes from upstream
antirez/ds4open PRs into the fork's main line.Part 1 — Server enhancements (existing local work)
Server + KV-cache improvements built around token-exact KV replay for a single
live session (pi backend): session transaction lifecycle, visible live
checkpoint for chat/Anthropic tool-call turns, prefill/decode cancellation on
client disconnect, threshold-based continued-snapshot cadence, deferred
consume-unlink,
GET /health+/stats, a bounded request queue (--max-queue),and an eviction-freshness grace so fresh never-hit waypoints are not evicted
before stale anchors.
Part 2 — Adopted from upstream open PRs (15 commits, 14 PRs)
Cherry-picked cleanly (3-way), original authorship preserved.
Server correctness / security
json_int()(UB double→int cast)reasoning_content, notcontenttoken_countbefore allocatingMetal (Apple Silicon)
Deliberately not included
ds4_server.c:18245/18352, which require a COLD prefix to outlive a conversation tip). The local8477b8ceviction-grace passes those tests while giving equivalent fresh-waypoint protection, so it is kept instead.Verification
make../ds4_test --serverOK,./ds4_agent_testOK,./ds4-eval --self-test-extractorsOK.