diff --git a/c/openai_server.py b/c/openai_server.py index 7476fc087..0a6032839 100644 --- a/c/openai_server.py +++ b/c/openai_server.py @@ -2660,12 +2660,22 @@ def generation_options(body, limit): if choice != "none" and not (body.get("tools") or body.get("functions")): raise APIError(400, "`tool_choice` requires `tools`.", "tool_choice", "invalid_value") stop_sequences = parse_stop_sequences(body) - if body.get("logprobs"): - raise APIError(400, "Log probabilities are not supported yet.", "logprobs", "unsupported_parameter") + # `logprobs`/`echo`/`top_logprobs` validation lives in logprobs_options() + # (called from generation(), which knows chat vs completions and the + # launched engine's capability) -- not here, so wiring the numeric + # channel can never leak chat-templating behavior into the raw + # completions path through this shared check. if body.get("frequency_penalty", 0) or body.get("presence_penalty", 0): raise APIError(400, "Token penalties are not supported yet.", None, "unsupported_parameter") - if body.get("seed") is not None: - raise APIError(400, "Per-request seeds are not supported yet.", "seed", "unsupported_parameter") + # `seed`: accepted for OpenAI-API request-shape compatibility, then + # silently discarded -- it has NO effect on any code path today, at any + # temperature: no engine and no wire field reads a per-request seed. + # glm and inkling each seed a process-global RNG from SEED once, at + # launch, never per request; no other engine reads SEED or a + # per-request seed at all. (At temperature 0 the question is moot + # anyway -- greedy decoding has no distribution to seed.) + # Accept-and-discard stays the whole of the seed behavior for this + # build; docs/api.md documents the no-op honestly. # response_format -> optional per-request grammar for the engine's grammar-forced # draft source (#70/#148). NEVER a sampling constraint: drafts are verified, so a # schema the engine cannot compile degrades to "no speedup", not to an error and @@ -2736,6 +2746,344 @@ def generation_options(body, limit): return maximum, float(temperature), float(top_p), grammar, stop_sequences +# The per-token top-k emission cap the engine's numeric logprobs channel +# (U7a's SUBMIT `logprobs=`) supports, mirrored in c/decode_batch.h as +# COLI_SUBMIT_TOPK_MAX -- run_ablate_score's existing top-32 read-out +# ceiling. The public request range is bound to exactly this engine +# interface: 1..32, refused with a named 400 above it. +LOGPROBS_TOP_K_CAP = 32 + +# Engine.generate()'s bound on waiting for ACCEPT after a SUBMIT that +# carries the extended key=value namespace (logprobs=k, ids=1, or both) -- +# see the comment beside accept_deadline in Engine.generate() for why an +# engine that predates the U7a extension can otherwise wedge this wait +# forever. +LOGPROBS_ACCEPT_TIMEOUT = float(os.environ.get("COLI_LOGPROBS_ACCEPT_TIMEOUT", "30")) + +# The most prompts one /v1/completions request may batch via an array +# `prompt`. A calibration marker, not a measured limit -- a batch above +# this is a named 400, never a silent truncation. +PROMPT_BATCH_CAP = 128 + +# The aggregate-token accounting boundary for a `prompt` array batch: total +# prompt tokens across a batch's members, named 400 above it -- a batch +# summing to exactly this many is admitted, one more is refused. Bounds +# worst-case echo-record memory when logprobs is also requested. For string +# batches the server has no tokenizer, so the same budget is applied to +# total UTF-8 bytes instead: every token carries at least one byte, so byte +# count is an exact upper bound on token count -- conservative for typical +# text, never under-protective. +PROMPT_BATCH_TOKEN_BUDGET = 65536 + +# The aggregate-token accounting boundary for a batch's GENERATED side: +# members multiplied by the request's effective max_tokens (the single +# `max_tokens`/`max_completion_tokens` value generation_options() returns +# after its own clamp to the operator's --max-tokens/--ngen), named 400 +# above it. Every member of a batch shares one generation_options() call, +# so this is members * maximum, not a per-member sum. Symmetric with +# PROMPT_BATCH_TOKEN_BUDGET, but bounds the response the batch will hold in +# memory before its single write rather than the prompt it submits. +PROMPT_BATCH_COMPLETION_BUDGET = 65536 + + +def logprobs_options(body, chat, engine_supports): + """Validate the client's logprobs/echo/top_logprobs request and translate + it into (engine_k, echo, display_k): + + - engine_k: the SUBMIT `logprobs=` value to send to the engine (0 = the + per-token numeric channel stays off entirely -- no ECHO/extended DATA + frames). + - echo: whether the response should include the prompt-echo positions. + Chat has no echo concept (OpenAI's chat schema doesn't have one) -- + always False for chat, and `echo: true` on a chat request is a named + 400, not a silent ignore. + - display_k: how many top_logprobs alternatives the CLIENT asked to see + (0..LOGPROBS_TOP_K_CAP) -- may be less than engine_k when a chat + request asked for logprobs with `top_logprobs: 0`. + + Completions' `logprobs` is the legacy integer top-k count; chat's is a + boolean gate plus a separate `top_logprobs` count -- two different + OpenAI conventions on the two endpoints, both accepted here. + + The zero/false/null semantics are explicit, never a truthiness + accident. On completions, `null`, `false`, and `0` all mean "no + logprobs" -- the request succeeds with `choices[].logprobs: null` -- + while `true` is a named 400 (the legacy field is an integer count, and + a boolean carries no count). On chat, `null` and `false` mean "no + logprobs" and any integer is a named 400 (the chat field is a boolean + gate). Integers outside 0..LOGPROBS_TOP_K_CAP get a named 400 on both + endpoints; the engine's top-k interface is the ceiling. + + On chat, `top_logprobs` is TYPE- and RANGE-checked even when + `logprobs` is false/absent (validation tightened during review): a + malformed `top_logprobs` is a named 400 whether or not the gate that + would use it is open, so a client never has a field silently ignored + outright because a sibling field made it moot. A valid `top_logprobs` + with `logprobs` off remains a documented no-op -- it still returns + (0, False, 0). `top_logprobs: null` is normalized to absent -- same + as `logprobs: null` -- before that type/range check runs, so it never + itself triggers a 400; the check applies only to a PRESENT, non-null + value. `top_logprobs` on completions is not read at all -- it is a + chat-only field in the OpenAI request shape, silently ignored there + exactly like any other unrecognized field. + + `echo`'s type is validated with the same isinstance(bool) check on + both endpoints before either endpoint decides what a valid value + means: completions accepts any boolean, chat still refuses `True` + with a named 400 (chat has no echo concept at all), but a non-bool + `echo` gets the shared "must be a boolean" 400 on either endpoint, + never chat's truthiness-based refusal message for a value that was + never a valid echo request shape to begin with. + """ + if chat: + echo = body.get("echo", False) + if not isinstance(echo, bool): + raise APIError(400, "`echo` must be a boolean.", "echo", "invalid_value") + if echo: + raise APIError(400, "`echo` is not supported for chat completions.", + "echo", "unsupported_parameter") + logprobs = body.get("logprobs", False) + if logprobs is None: + logprobs = False # null == absent == no logprobs + if not isinstance(logprobs, bool): + raise APIError(400, "`logprobs` must be a boolean.", "logprobs", "invalid_value") + display_k, param, echo = body.get("top_logprobs", 0), "top_logprobs", False + if display_k is None: + display_k = 0 # null == absent, same as `logprobs` + else: + echo = body.get("echo", False) + if not isinstance(echo, bool): + raise APIError(400, "`echo` must be a boolean.", "echo", "invalid_value") + logprobs = body.get("logprobs") + if logprobs is None or logprobs is False: + return 0, echo, 0 # echo without logprobs: documented no-op + display_k, param = logprobs, "logprobs" + if (isinstance(display_k, bool) or not isinstance(display_k, int) or + not 0 <= display_k <= LOGPROBS_TOP_K_CAP): + raise APIError(400, f"`{param}` must be an integer between 0 and " + f"{LOGPROBS_TOP_K_CAP}.", param, "invalid_value") + if chat and not logprobs: + return 0, False, 0 # top_logprobs validated above; logprobs off is still a no-op + if not chat and display_k == 0: + return 0, echo, 0 # 0 = no logprobs, documented + if not engine_supports: + # Conservative sender-side capability gate: the server never emits + # SUBMIT logprobs= to an engine that does not implement the numeric + # per-token channel (only the glm engine's mux loop reads + # sub.logprobs at all). Rejecting here, before generate() ever + # builds the extension header, keeps a rejected-SUBMIT + # payload-drain hazard unreachable. + raise APIError(400, f"Log probabilities are not supported by the {ARCH} engine.", + "logprobs", "unsupported_parameter") + return max(1, display_k), echo, display_k + + +def _json_float(value): + """A logprob value that arrives non-finite (nan/inf/-inf, which the + engine's numeric channel can emit for a degenerate logit row) must + serialize as JSON `null`, matching what OpenAI clients expect -- + never the invalid-JSON NaN/Infinity literals json.dumps would + otherwise write, and never clamped to a made-up finite number.""" + return value if math.isfinite(value) else None + + +def _encode_token_id_prompt(ids): + """A `prompt` array of token ids, encoded into the exact ASCII decimal + wire form c/decode_batch.h's coli_ids_parse expects -- straight + passthrough, no detokenize/re-encode round trip. Structural validation + only; an id the engine's own vocab rejects still comes back as a named + BAD_REQUEST from the engine itself (coli_ids_parse), not silently + accepted here.""" + if not isinstance(ids, list) or not ids: + raise APIError(400, "`prompt` array must be a non-empty array of token ids.", + "prompt", "invalid_value") + for value in ids: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise APIError(400, "`prompt` array must contain only non-negative integer " + "token ids.", "prompt", "invalid_value") + return " ".join(str(value) for value in ids) + + +def _order_echo_records(prompt_records): + """Place each prompt-echo record at its own wire `pos` index rather + than trusting the order the frames arrived in -- arrival order + happens to match position order for today's single-threaded engine + loop, but nothing here depends on that holding. A duplicate, + negative, out-of-range, or (by pigeonhole) missing position among the + N records that must fill exactly slots 0..N-1 raises a named + RuntimeError, the same class every other malformed-engine-output path + already raises, so it surfaces as a clean 500 rather than a hang or a + silently corrupted response.""" + ordered = [None] * len(prompt_records) + for pos, data, record in prompt_records: + if (not isinstance(pos, int) or isinstance(pos, bool) + or not 0 <= pos < len(ordered) or ordered[pos] is not None): + raise RuntimeError( + f"invalid engine ECHO position {pos!r} (expected each of " + f"0..{len(ordered) - 1} exactly once, got {len(prompt_records)} records)") + ordered[pos] = (data, record) + return ordered + + +def _logprob_positions(prompt_records, generated_records): + """Merge prompt-echo and generated-token logprob records into one + ordered list of {"text", "raw_lp", "topk", "bytes"} dicts: echoed + positions first, reassembled by wire `pos` via _order_echo_records, + then generated tokens in emission order. + + Text is decoded with a single stateful incremental UTF-8 decoder + spanning the whole sequence -- the same pattern the main generation + path already uses -- rather than decoding each frame's bytes + independently. A multi-byte codepoint can arrive split across two + adjacent byte-level vocabulary pieces; decoding per-frame would turn + each half into its own replacement character instead of the true + joined text. A frame that only completes a pending byte sequence + contributes "" as its own text (the resolved character lands on + whichever position finished it), while `bytes` always stays that + frame's own raw payload regardless of what text, if any, it decoded + to.""" + raw = [(data, record) for data, record in _order_echo_records(prompt_records)] + raw += [(data, record) for data, record in generated_records] + decoder = codecs.getincrementaldecoder("utf-8")("replace") + positions = [{"text": decoder.decode(data), "raw_lp": record["lp"], + "topk": record["topk"], "bytes": data} for data, record in raw] + tail = decoder.decode(b"", final=True) + if tail and positions: + positions[-1]["text"] += tail + return positions + + +def _own_token_label(entry, topk, idx): + """The wire's top-k table carries only raw candidate token ids, never + decoded text -- there is no server-side tokenizer, and the wire's + logprob_tail record carries no chosen-token id either (only `lp` and + the table). The one candidate that CAN be labeled correctly without + decoding is the position's own chosen token: its top-k entry and its + own `lp` come from the exact same computation, so an exact float + match identifies it without ever comparing token ids. + + The engine prints logprobs to 6 decimal digits, so two distinct + candidates can legitimately share the printed value the chosen token + also carries -- a real tie, not a bug in either side. With no id to + break it, the match is resolved deterministically: only the FIRST + table entry (in wire order, up to `idx`) whose value exactly equals + the chosen logprob is labeled as the chosen token; any later entry + that also matches is labeled by its raw id like any other + unidentified candidate, same as `topk` argmax rank 1 falls back + to id when it never matches at all. This is a documented limitation, + not a decode -- and separately, the top-k table is unsorted on the + wire, so this never assumes the first entry overall is the argmax.""" + tid, tlp = topk[idx] + if tlp != entry["raw_lp"]: + return f"" + if any(topk[j][1] == entry["raw_lp"] for j in range(idx)): + return f"" + return entry["text"] + + +def _completions_logprobs_object(prompt_records, generated_records, display_k): + """Build the legacy `/v1/completions` `logprobs` object: `tokens[]`, + `token_logprobs[]`, `top_logprobs[]` (one dict per position, keyed by + token text via _own_token_label), and `text_offset[]` (character + offsets into the reconstructed text, derived from the same decode -- + no tokenizer required, always counted from 0). The first prompt + position's `token_logprobs` entry comes out `null` for free: the + engine's own echo position 0 already carries a non-finite sentinel + logprob (there is nothing to condition the first token on), which + _json_float maps to null.""" + positions = _logprob_positions(prompt_records, generated_records) + tokens = [p["text"] for p in positions] + token_logprobs = [_json_float(p["raw_lp"]) for p in positions] + top_logprobs = [] + for p in positions: + table = {} + displayed = p["topk"][:display_k] + for idx, (tid, tlp) in enumerate(displayed): + table[_own_token_label(p, displayed, idx)] = _json_float(tlp) + top_logprobs.append(table) + text_offset = [] + offset = 0 + for text in tokens: + text_offset.append(offset) + offset += len(text) + return {"tokens": tokens, "token_logprobs": token_logprobs, + "top_logprobs": top_logprobs, "text_offset": text_offset} + + +def _chat_logprobs_content(generated_records, display_k): + """Build chat-completions `choices[].logprobs.content[]`: one + `{token, logprob, bytes, top_logprobs}` entry per generated token. + Chat has no echo concept, so there is no `pos` field to reassemble -- + generated records are used in emission order. Each `top_logprobs` + entry is `{token, logprob, bytes}`, using the same id-based labeling + limitation as the legacy object's dict keys (_own_token_label). + + Token text for every generated record is decoded FIRST, in one pass, + with a single stateful incremental UTF-8 decoder spanning the whole + generated sequence (same reasoning as _logprob_positions: a frame + that only completes a prior pending multi-byte codepoint contributes + "" as its own token text) -- including the final flush of any + trailing incomplete sequence, before any `content` entry is built. + Building entries only after every position's final text is known + keeps the last entry's own `top_logprobs` label consistent with its + `token` field; building them interleaved with decoding would let the + last entry's `token` gain the flushed text while the label computed + from its pre-flush text stayed stale. `bytes` always stays each + frame's own raw payload regardless of what text, if any, it decoded + to.""" + decoder = codecs.getincrementaldecoder("utf-8")("replace") + texts = [decoder.decode(data) for data, _record in generated_records] + tail = decoder.decode(b"", final=True) + if tail and texts: + texts[-1] += tail + content = [] + for (data, record), text in zip(generated_records, texts): + entry = {"text": text, "raw_lp": record["lp"], "topk": record["topk"]} + alternatives = [] + displayed = record["topk"][:display_k] + for idx, (tid, tlp) in enumerate(displayed): + label = _own_token_label(entry, displayed, idx) + alternatives.append({"token": label, "logprob": _json_float(tlp), + "bytes": list(data) if label == text else None}) + content.append({"token": text, "logprob": _json_float(record["lp"]), + "bytes": list(data), "top_logprobs": alternatives}) + return content + + +def _trim_generated_records_to_text(generated_records, text): + """Drop trailing generated-token records whose bytes were filtered out of + the text actually returned to the client (a matched stop sequence, most + commonly) -- the logprobs arrays must correspond to what `text` holds, + never a superset of it. Decodes each record's bytes with the same + incremental-UTF-8-decoder-over-the-whole-sequence approach the response + assembly helpers use, and keeps a record only while the text decoded so + far stays a prefix of `text`; the first record whose decoded text would + NOT be a prefix, and every record after it, is dropped. + + This only closes the stop-sequence case (the common one, and a small, + local fix): a filtered thinking/reasoning split or a tool-call parse can + ALSO diverge the returned `text` from the raw generated stream in ways + this prefix check does not attempt to align -- callers pass this + function the raw, pre-split, pre-tool-parse text for exactly that + reason, and the divergence that remains afterward is a documented + limitation, not silently shipped.""" + decoder = codecs.getincrementaldecoder("utf-8")("replace") + kept = [] + offset = 0 + for data, record in generated_records: + piece = decoder.decode(data) + # `text[:offset]` is already established as a prefix of `text` by + # every prior iteration, so checking `piece` against `text` at + # `offset` is equivalent to rebuilding and checking the whole + # accumulated candidate, without the O(len(acc)) rebuild -- this + # loop is O(len(text)) total instead of O(len(text)^2). + if not text.startswith(piece, offset): + break + offset += len(piece) + kept.append((data, record)) + return kept + + def read_engine_turn(stream, sentinel, on_bytes): pending = b"" while True: @@ -2927,6 +3275,35 @@ class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure): return None # never let process bookkeeping break starting the engine +def _write_all(stream, data, frame): + """Write every byte of `data` to `stream`, looping on short writes. + + The production engine stdin is a raw, unbuffered pipe (bufsize=0 -> + io.FileIO), whose write() is a single os.write() and may transfer fewer + bytes than it was given (a signal landing mid-write, a full pipe buffer + on a large IMAGE frame). Discarding the return value would leave the + tail of a frame unsent and desynchronize the engine's stdin framing, so + the remainder is re-offered until it is all consumed. + + Neither `None` nor 0 is progress. `RawIOBase.write` answers `None` when + the stream is non-blocking and could not take a single byte, and 0 says + the same thing with a count; re-offering the buffer after either would + spin forever, so both fail closed as the named engine-write error a + broken pipe raises. The production stdin is a blocking raw pipe whose + write() always returns a positive int, so neither ever runs there.""" + written = 0 + total = len(data) + while written < total: + sent = stream.write(data[written:]) + # None is RawIOBase's "not one byte went out", not an uncounted + # full write, so it fails closed exactly as a zero count does. + if sent is None or sent <= 0: + raise RuntimeError( + f"failed to write {frame} to the engine " + f"(stdin took {written} of {total} bytes)") + written += sent + + class Engine: # cap=None = "not explicitly set": a glm-arch model's engine resolves the # 0 sentinel (8 historically, 1 on Metal+darwin+fast SSD -- colibri.c @@ -2946,6 +3323,14 @@ def __init__(self, executable, model, cap=None, max_tokens=1024, env=None, kv_sl arch = family.id self.family = family self.model_dir = str(model) + # Capability gating of the extended SUBMIT namespace (logprobs=/ids=), + # established once here at launch, not per request, and never via an + # engine-version handshake -- only the glm engine (c/colibri.c) + # implements the U7a numeric channel and token-id intake; every other + # engine's mux_data/prefill loop never reads sub.logprobs or + # sub.tok_ids at all, so sending them would be silently wrong (an + # accepted-but-ignored request), not just unsupported. + self.supports_logprobs_echo = (arch == "glm") child_env = dict(env or os.environ, SNAP=str(model), SERVE="1", SERVE_BATCH="1", NGEN=str(max_tokens), KV_SLOTS=str(kv_slots)) tune_child_env(child_env, arch) @@ -3003,17 +3388,70 @@ def _fail_pending(self, error): for events in requests: events.put(("error", error)) - def _read_exact(self, size): + def _write_frame(self, data, frame): + """Checked server->engine protocol write for CANCEL/STOP: the write + and its flush happen under one write_lock acquisition, and a failed + write is re-raised as a named RuntimeError rather than left as the + OSError it started as. BrokenPipeError is a ConnectionError + subclass, so an unwrapped failure here would fall into do_POST's + client-hangup handler (`except ConnectionError: pass`) and the + client would see a silent connection close instead of the 500 + engine_error the failure actually is.""" + try: + with self.write_lock: + _write_all(self.process.stdin, data, frame) + self.process.stdin.flush() + except OSError as error: + raise RuntimeError(f"failed to write {frame} to the engine ({error})") from error + + def _read_exact(self, size, kind="DATA"): chunks = [] remaining = size while remaining: chunk = self.process.stdout.read(remaining) if chunk == b"": - raise RuntimeError("truncated engine DATA payload") + raise RuntimeError(f"truncated engine {kind} payload") chunks.append(chunk) remaining -= len(chunk) return b"".join(chunks) + @staticmethod + def _parse_logprob_tail(fields, i): + """Parse the numeric tail shared by opted-in DATA/ECHO frames + (c/colibri.c logprob_tail): " [ ]*k". float() + parses the engine's numeric wire tokens directly -- "nan"/"inf"/ + "-inf" and any %g precision alike (an echo's position 0 -- nothing + to condition on -- carries "nan 0", mux_prefill_echo). The table is + UNSORTED on the wire (logprob_tail selects by raw logit) -- callers + must not assume the first pair is argmax. Every malformed shape -- + a short or over-long field list (trailing garbage included), + non-numeric fields, an out-of-range k, or a negative token id -- + is a named RuntimeError, never a silent partial record.""" + if len(fields) < i + 2: + raise RuntimeError("invalid engine logprob tail: missing lp/k") + try: + lp = float(fields[i]) + k = int(fields[i + 1]) + except ValueError as error: + raise RuntimeError(f"invalid engine logprob tail: {error}") from error + if not 0 <= k <= LOGPROBS_TOP_K_CAP: + raise RuntimeError(f"invalid engine logprob tail: k={k} out of range") + if len(fields) != i + 2 + 2 * k: + raise RuntimeError("invalid engine logprob tail: field count mismatch") + topk = [] + j = i + 2 + for _ in range(k): + try: + tid = int(fields[j]) + tlp = float(fields[j + 1]) + except ValueError as error: + raise RuntimeError(f"invalid engine logprob tail: {error}") from error + if tid < 0: + raise RuntimeError(f"invalid engine logprob tail: negative token id {tid}") + topk.append((tid, tlp)) + j += 2 + return {"lp": lp, "topk": topk} + def _dispatch_stdout(self): try: while True: @@ -3029,9 +3467,15 @@ def _dispatch_stdout(self): # numeric channel ("DATA [tid tlp]*k"), # emitted only for requests that opted in via the SUBMIT # logprobs field. The payload framing is identical; the - # numeric fields are consumed by the server feature half - # (U7b) -- accepted here so the frame never kills the - # dispatcher (and with it every in-flight request). + # numeric tail is parsed into a record (None for legacy + # frames) and carried alongside the payload bytes so a + # later consumer can thread it into the response. A + # malformed tail now raises and kills the dispatcher + # (failing every in-flight request) rather than being + # tolerated as an unrecognized trailing field -- a + # deliberate reversal of this arm's older, more lenient + # stance, matching the fail-closed frame-validation + # policy every other frame kind here already follows. request_id = fields[1] size = int(fields[2]) if not 0 <= size <= 65536: @@ -3039,10 +3483,15 @@ def _dispatch_stdout(self): data = self._read_exact(size) if self._read_exact(1) != b"\n": raise RuntimeError("invalid engine DATA terminator") + record = self._parse_logprob_tail(fields, 3) if len(fields) > 3 else None with self.pending_lock: events = self.pending.get(request_id) if events is not None: - events.put(("data", data)) + # A bare-bytes put on the (far more common) + # non-opted-in path costs zero extra allocation per + # generated token, same as before this channel + # existed. + events.put(("data", data if record is None else (data, record))) elif kind == "TOOL" and len(fields) == 3: # Opaque, request-scoped structured output. K3 emits an # initial zero-byte frame before generation so DATA marker @@ -3061,16 +3510,43 @@ def _dispatch_stdout(self): elif kind == "ECHO" and len(fields) >= 6: # U7a prefill read-out: "ECHO # [tid tlp]*k" plus a DATA-framed payload (n bytes + LF). - # Emitted only for opted-in requests; no current request - # path opts in, so the frame is read (to keep the stream - # in sync) and dropped -- U7b delivers it to the response - # assembly when it wires the opt-in. + # Emitted only for opted-in requests. Delivered to the + # pending request the same way DATA is; a request with + # no pending entry (the id was never admitted, or already + # finished) drops it after fully reading the frame, same + # as every other frame kind here. + request_id = fields[1] size = int(fields[2]) if not 0 <= size <= 65536: - raise RuntimeError("invalid engine DATA size") - self._read_exact(size) - if self._read_exact(1) != b"\n": - raise RuntimeError("invalid engine DATA terminator") + raise RuntimeError("invalid engine ECHO size") + data = self._read_exact(size, "ECHO") + if self._read_exact(1, "ECHO") != b"\n": + raise RuntimeError("invalid engine ECHO terminator") + pos = int(fields[3]) + record = self._parse_logprob_tail(fields, 4) + with self.pending_lock: + events = self.pending.get(request_id) + if events is not None: + events.put(("echo", (pos, data, record))) + elif ((kind == "GRPP" and len(fields) >= 6) or + (kind == "GRPG" and len(fields) >= 7)): + # Group-scoring read-outs a future engine may emit + # alongside a batch of ordinary requests. This server has + # no group-response contract yet, so both are drained + # like any other payload-framed kind (byte count in + # field 2, same bound and terminator as DATA) and never + # reach a request's event queue. + size = int(fields[2]) + if not 0 <= size <= 65536: + raise RuntimeError(f"invalid engine {kind} size") + self._read_exact(size, kind) + if self._read_exact(1, kind) != b"\n": + raise RuntimeError(f"invalid engine {kind} terminator") + elif ((kind == "GRPS" and len(fields) == 5) or + (kind == "GRPE" and len(fields) == 6)): + # Header-only group-member boundaries: no payload to + # read and nothing for this server to act on. + pass elif kind == "ACCEPT" and len(fields) >= 3: # #597: the engine validated the submission (fits context) before prefill. # Keep it pending — DATA/DONE still follow — and let generate() commit the @@ -3133,9 +3609,16 @@ def _dispatch_stdout(self): def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0, cancelled=None, grammar=None, stopped=None, on_accept=None, audio=None, - on_tool=None, image=None): + on_tool=None, image=None, logprobs=0, echo=False, tok_ids=False): if isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or not 0 <= cache_slot < self.kv_slots: raise APIError(400, "Invalid cache slot.", "cache_slot") + if (logprobs or tok_ids) and not self.supports_logprobs_echo: + # Defense in depth: APIHandler's logprobs_options()/generation()/ + # completion() already refuse the HTTP request with a named 400 + # before ever reaching here. Reaching this with the gate false is + # a caller bug, not a bad request. + raise RuntimeError("logprobs/token-id intake requested against an engine " + "that does not support the U7a extension") payload = prompt.encode("utf-8") if b"\0" in payload: raise APIError(400, "NUL bytes are not supported in prompts.", "messages") @@ -3174,139 +3657,230 @@ def decode_tool(data): request_id = str(self.next_request_id) self.next_request_id += 1 self.pending[request_id] = events - xpayload = gpayload or apayload - # DeepSeek V4 prefix hint (optional 8th header field): the byte length of - # the rendered prompt up to the first user/assistant turn marker — the - # stable system prefix. The engine snapshots its attention state at that - # token boundary during the prefill, so the FIRST request of the first - # conversation already seeds the shared-prefix checkpoint that every later - # conversation (opencode session) restores in seconds; without the hint - # the engine only discovers the boundary on the second fresh prompt. - # Older engines parse six or seven fields and ignore the eighth. - prefix_field = "" - if ARCH == "deepseek_v4": - cut = min((i for i in (prompt.find("<\uff5cUser\uff5c>"), - prompt.find("<\uff5cAssistant\uff5c>")) if i > 0), - default=0) - if cut > 0: - prefix_field = f" {len(xpayload)} {len(prompt[:cut].encode('utf-8'))}" - header = (f"SUBMIT {request_id} {cache_slot} {len(payload)} {max_tokens} " - f"{temperature:.8g} {top_p:.8g}" - + (prefix_field if prefix_field else (f" {len(xpayload)}" if xpayload else "")) - + "\n").encode() try: + xpayload = gpayload or apayload + # DeepSeek V4 prefix hint (optional 8th header field): the byte length of + # the rendered prompt up to the first user/assistant turn marker — the + # stable system prefix. The engine snapshots its attention state at that + # token boundary during the prefill, so the FIRST request of the first + # conversation already seeds the shared-prefix checkpoint that every later + # conversation (opencode session) restores in seconds; without the hint + # the engine only discovers the boundary on the second fresh prompt. + # Older engines parse six or seven fields and ignore the eighth. + prefix_field = "" + if ARCH == "deepseek_v4": + cut = min((i for i in (prompt.find("<\uff5cUser\uff5c>"), + prompt.find("<\uff5cAssistant\uff5c>")) if i > 0), + default=0) + if cut > 0: + prefix_field = f" {len(xpayload)} {len(prompt[:cut].encode('utf-8'))}" + # SUBMIT's key=value extension namespace (decode_batch.h + # coli_submit_ext) -- logprobs=k opts into U7a's per-token numeric + # channel (ECHO + extended DATA frames), ids=1 marks the payload as + # pre-tokenized ASCII decimal ids rather than raw text. Only ever + # built when the capability gate above already passed, so this never + # reaches an engine that would mis-parse or silently ignore it. The + # extension arm requires the 7th (gbytes) field to be present even + # when it's 0 -- coli_submit_parse expects exactly 7 numeric fields + # before the first key=value token. + ext_parts = [] + if logprobs: + ext_parts.append(f"logprobs={logprobs}") + if tok_ids: + ext_parts.append("ids=1") + ext_field = (" " + " ".join(ext_parts)) if ext_parts else "" + gbytes_field = f" {len(xpayload)}" if (xpayload or ext_parts) else "" + header = (f"SUBMIT {request_id} {cache_slot} {len(payload)} {max_tokens} " + f"{temperature:.8g} {top_p:.8g}" + + (prefix_field if prefix_field else gbytes_field) + + ext_field + + "\n").encode() with self.write_lock: if self.process.poll() is not None: raise RuntimeError("colibri engine is not running") - # Le patch sono binarie e grosse: viaggiano in un frame loro, - # annunciato subito prima del SUBMIT a cui appartengono. Deve - # partire dentro lo stesso lock, o un'altra richiesta potrebbe - # infilarsi in mezzo e prendersi l'immagine di questa. - if image is not None: - patches, grid_h, grid_w = image - blob = patches.tobytes() if hasattr(patches, "tobytes") else patches - self.process.stdin.write( - f"IMAGE {request_id} {len(blob)} {grid_h} {grid_w}\n".encode() - + blob + b"\n") - self.process.stdin.write(header + payload + xpayload + b"\n") - self.process.stdin.flush() - except Exception: - with self.pending_lock: - self.pending.pop(request_id, None) - raise - - cancel_sent = False - stop_sent = False - accepted = False - - def _accept(info): - # #597: commit exactly once, on the first of ACCEPT / DATA / DONE. A new engine sends - # ACCEPT before any output, so on_accept fires before prefill and a preceding - # CONTEXT_EXCEEDED never reaches here (it propagates as a 400 with nothing committed). - # An older engine that never sends ACCEPT still commits on its first DATA/DONE. - nonlocal accepted - if not accepted: - accepted = True - if on_accept is not None: - on_accept(info) + try: + # Le patch sono binarie e grosse: viaggiano in un frame loro, + # annunciato subito prima del SUBMIT a cui appartengono. Deve + # partire dentro lo stesso lock, o un'altra richiesta potrebbe + # infilarsi in mezzo e prendersi l'immagine di questa. + if image is not None: + patches, grid_h, grid_w = image + blob = patches.tobytes() if hasattr(patches, "tobytes") else patches + _write_all( + self.process.stdin, + f"IMAGE {request_id} {len(blob)} {grid_h} {grid_w}\n".encode() + + blob + b"\n", "SUBMIT") + _write_all(self.process.stdin, + header + payload + xpayload + b"\n", "SUBMIT") + self.process.stdin.flush() + except OSError as error: + raise RuntimeError( + f"failed to write SUBMIT to the engine ({error})") from error + + cancel_sent = False + stop_sent = False + accepted = False + # U7a's ECHO frames (whenever logprobs>0, the engine ALWAYS runs the + # full prefill read-out -- c/colibri.c mux_submit: `echo = + # sub.logprobs>0`, unconditionally, there is no separate wire bit for + # "logprobs but no echo") land here in position order; DATA's numeric + # tail lands here in emission order. `generated_logprobs` stays + # empty when logprobs=0. `prompt_logprobs` ALSO stays empty unless + # the caller opted into `echo` too -- the engine still sends every + # ECHO frame regardless (there is no wire bit for "logprobs but no + # echo"), but retaining a full prompt-length echo table for every + # opted-in request that never asked to see it (the common chat case, + # and any completions request with echo=False) would hold it in + # memory for the whole request lifetime for nothing: at k=32 and a + # 32k-token prompt that table is on the order of 100+ MB. Frames not + # kept are still fully drained off the wire above (read_engine_turn + # already consumed the payload bytes before this event is queued), + # so this never desyncs the dispatcher -- it only decides whether the + # record is RETAINED past this iteration. + prompt_logprobs = [] + generated_logprobs = [] + # Bound how long a request that opted into the extended SUBMIT + # namespace (logprobs=k, ids=1, or both) waits for the engine's + # ACCEPT before concluding the engine silently rejected the extended + # header -- c/decode_batch.h: "An OLD engine rejects ANY extended + # header ... the engine answers ERROR 0 BAD_REQUEST", an id that + # never matches this (or any) pending request, so that reply is + # otherwise invisible to this loop and it would wait here forever. + # A legacy (non-opted-in) request keeps the old, unbounded wait -- + # every engine build ever shipped handles a plain SUBMIT. + accept_deadline = (time.monotonic() + LOGPROBS_ACCEPT_TIMEOUT + if (logprobs or tok_ids) else None) + + def _accept(info): + # #597: commit exactly once, on the first of ACCEPT / DATA / DONE. A new engine sends + # ACCEPT before any output, so on_accept fires before prefill and a preceding + # CONTEXT_EXCEEDED never reaches here (it propagates as a 400 with nothing committed). + # An older engine that never sends ACCEPT still commits on its first DATA/DONE. + nonlocal accepted + if not accepted: + accepted = True + if on_accept is not None: + on_accept(info) - while True: - try: - kind, value = events.get(timeout=0.05) - except queue.Empty: - # #908: cancelled() is only polled in the "data" branch, so a - # client that disconnects before the engine's first DATA frame - # (it is still prefilling) never cancels: the CANCEL never went - # out, the turn ran to its token limit, and this thread stayed - # blocked until the engine emitted something. Poll the callback - # while idle so a pre-first-frame disconnect cancels too. - # - # Do NOT raise here: this thread holds the scheduler admission, - # and releasing it before the engine confirms the cancel lets - # the next request SUBMIT into a pipe the busy engine is not - # reading — every later request then hangs silently behind the - # orphaned generation. Wait for the engine's ERROR CANCELLED / - # DONE frame; ClientCancelled is raised when it arrives. - if not cancel_sent and not stop_sent and cancelled and cancelled(): - cancel_sent = True - with self.write_lock: - self.process.stdin.write(f"CANCEL {request_id}\n".encode()) - self.process.stdin.flush() - continue - if kind == "accept": - if accepted: - raise RuntimeError("engine sent a duplicate ACCEPT frame") - _accept(value) - elif kind == "data": - _accept({"prompt_tokens": None}) - if not cancel_sent and not stop_sent: - decode(value) - if stopped and stopped(): - stop_sent = True - with self.write_lock: - self.process.stdin.write(f"STOP {request_id}\n".encode()) - self.process.stdin.flush() - elif cancelled and cancelled(): - # Same admission-holding rule as the idle branch above: - # send CANCEL, then keep consuming frames until the - # engine acknowledges with ERROR CANCELLED or DONE. - cancel_sent = True - with self.write_lock: - self.process.stdin.write(f"CANCEL {request_id}\n".encode()) - self.process.stdin.flush() - elif kind == "tool": - _accept({"prompt_tokens": None}) - if not cancel_sent and not stop_sent: - decode_tool(value) - if stopped and stopped(): - stop_sent = True - with self.write_lock: - self.process.stdin.write(f"STOP {request_id}\n".encode()) - self.process.stdin.flush() - elif cancelled and cancelled(): + while True: + try: + kind, value = events.get(timeout=0.05) + except queue.Empty: + # #908: cancelled() is only polled in the "data" branch, so a + # client that disconnects before the engine's first DATA frame + # (it is still prefilling) never cancels: the CANCEL never went + # out, the turn ran to its token limit, and this thread stayed + # blocked until the engine emitted something. Poll the callback + # while idle so a pre-first-frame disconnect cancels too. + # + # Do NOT raise here: this thread holds the scheduler admission, + # and releasing it before the engine confirms the cancel lets + # the next request SUBMIT into a pipe the busy engine is not + # reading — every later request then hangs silently behind the + # orphaned generation. Wait for the engine's ERROR CANCELLED / + # DONE frame; ClientCancelled is raised when it arrives. + if (accept_deadline is not None and not accepted + and time.monotonic() > accept_deadline): + # logprobs takes priority in the message/param/code when + # both were requested together -- an established, + # already-tested shape this must not disturb; tok_ids + # alone gets its own named error instead of borrowing + # the logprobs one. + if logprobs: + raise APIError( + 503, "The colibri engine did not accept a per-token logprobs " + "request in time; it may not support the per-token logprobs " + "extension.", "logprobs", "engine_logprobs_unsupported", + "server_error") + raise APIError( + 503, "The colibri engine did not accept a token-id prompt " + "request in time; it may not support the pre-tokenized " + "prompt extension.", "prompt", "engine_tok_ids_unsupported", + "server_error") + if not cancel_sent and not stop_sent and cancelled and cancelled(): cancel_sent = True - with self.write_lock: - self.process.stdin.write(f"CANCEL {request_id}\n".encode()) - self.process.stdin.flush() - elif kind == "done": - _accept({"prompt_tokens": None}) - if cancel_sent: - # The engine finished the turn before seeing the CANCEL - # (or honored it at a token boundary and still framed a - # DONE). Either way the client is gone: the ack is what - # mattered, the output is not deliverable. + self._write_frame(f"CANCEL {request_id}\n".encode(), "CANCEL") + continue + if kind == "accept": + if accepted: + raise RuntimeError("engine sent a duplicate ACCEPT frame") + _accept(value) + elif kind == "data": + _accept({"prompt_tokens": None}) + # The dispatcher only wraps in a tuple when a logprob record + # rides along; a bare-bytes value is the (far more common) + # non-opted-in shape. + data, record = value if isinstance(value, tuple) else (value, None) + if record is not None: + generated_logprobs.append((data, record)) + if not cancel_sent and not stop_sent: + decode(data) + if stopped and stopped(): + stop_sent = True + self._write_frame(f"STOP {request_id}\n".encode(), "STOP") + elif cancelled and cancelled(): + # Same admission-holding rule as the idle branch above: + # send CANCEL, then keep consuming frames until the + # engine acknowledges with ERROR CANCELLED or DONE. + cancel_sent = True + self._write_frame(f"CANCEL {request_id}\n".encode(), "CANCEL") + elif kind == "echo": + # Prompt-position readout: recorded (only when the caller + # asked to see it -- see the comment above prompt_logprobs), + # nothing emitted -- no text is decoded for a prompt echo, + # unlike "data"/"tool". + pos, data, record = value + if echo: + prompt_logprobs.append((pos, data, record)) + elif kind == "tool": + _accept({"prompt_tokens": None}) + if not cancel_sent and not stop_sent: + decode_tool(value) + if stopped and stopped(): + stop_sent = True + self._write_frame(f"STOP {request_id}\n".encode(), "STOP") + elif cancelled and cancelled(): + cancel_sent = True + self._write_frame(f"CANCEL {request_id}\n".encode(), "CANCEL") + elif kind == "done": + _accept({"prompt_tokens": None}) + if cancel_sent: + # The engine finished the turn before seeing the CANCEL + # (or honored it at a token boundary and still framed a + # DONE). Either way the client is gone: the ack is what + # mattered, the output is not deliverable. + raise ClientCancelled() + tail = decoder.decode(b"", final=True) + if tail: + on_text(tail) + tool_tail = tool_decoder.decode(b"", final=True) + if tool_tail and on_tool is not None: + on_tool(tool_tail) + if logprobs: + value["logprobs"] = {"prompt": prompt_logprobs, + "generated": generated_logprobs} + return value + elif cancel_sent and isinstance(value, RuntimeError) and str(value) == "CANCELLED": raise ClientCancelled() - tail = decoder.decode(b"", final=True) - if tail: - on_text(tail) - tool_tail = tool_decoder.decode(b"", final=True) - if tool_tail and on_tool is not None: - on_tool(tool_tail) - return value - elif cancel_sent and isinstance(value, RuntimeError) and str(value) == "CANCELLED": - raise ClientCancelled() - else: - raise value + elif (tok_ids and isinstance(value, RuntimeError) + and str(value) == "BAD_REQUEST"): + # coli_ids_parse (c/decode_batch.h) reports a malformed or + # out-of-vocabulary token id as "ERROR BAD_REQUEST", + # carrying THIS request's own id (unlike the "ERROR 0 + # BAD_REQUEST" an old engine sends for an unrecognized + # extended header, which never reaches here at all -- id 0 + # never matches a pending request, see accept_deadline + # above). This is the client's fault, not the server's: a + # named 400 on `prompt`, not the generic 500 engine_error + # every other unexpected RuntimeError still becomes. + raise APIError(400, "The engine rejected this token-id prompt (a " + "malformed or out-of-vocabulary token id).", + "prompt", "invalid_value") + else: + raise value + finally: + with self.pending_lock: + self.pending.pop(request_id, None) def close(self): with self.pending_lock: @@ -3839,7 +4413,7 @@ def error_body(self, error): return {"type": "error", "error": {"type": error.error_type, "message": error.message}} def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=None, - enable_thinking=False, audio=None, image=None): + enable_thinking=False, audio=None, image=None, tok_ids=False): # COLI_DEBUG tees the engine transaction to stderr: 1 = decoded output stream only, # 2 = both sides (rendered prompt + output). render_chat already folds prior turns and # tool results into `prompt`, so level 2 is the full conversation the engine saw. @@ -3858,6 +4432,8 @@ def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=Non # grammar payload extension would desync its stdin framing. raise APIError(400, f"`response_format` grammars are not supported by the {ARCH} " "engine yet.", "response_format", "unsupported_parameter") + engine_k, echo, display_k = logprobs_options( + body, chat, getattr(self.server.engine, "supports_logprobs_echo", False)) stop_sequences, ignore_leading_stop = stop_policy(body, chat) # tools and tool_choice come from chat_completion() already processed/filtered if chat and tool_choice == "none": @@ -3878,6 +4454,13 @@ def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=Non stream = body.get("stream", False) if not isinstance(stream, bool): raise APIError(400, "`stream` must be a boolean.", "stream") + if engine_k and stream: + # Streaming + logprobs together is a named 400, not a silent drop + # of the numeric channel: streamed per-delta logprobs are not + # built, and nothing downstream of this point ever threads + # logprobs= into a streaming call. + raise APIError(400, "`logprobs` is not supported together with `stream`.", + "logprobs", "unsupported_parameter") stream_options = body.get("stream_options") if stream else None if stream and stream_options is not None and not isinstance(stream_options, dict): raise APIError(400, "`stream_options` must be an object.", "stream_options") @@ -3904,10 +4487,18 @@ def generation_stopped(): self.client_disconnected, grammar=grammar, stopped=generation_stopped, **({"on_tool": sideband.feed} if sideband.enabled else {}), **({"audio": audio} if audio else {}), - **({"image": image} if image is not None else {})) + **({"image": image} if image is not None else {}), + **({"logprobs": engine_k, "echo": echo} if engine_k else {}), + **({"tok_ids": True} if tok_ids else {})) stop_filter.finish() sideband.finish() text = "".join(output) + # The raw, pre-split, pre-tool-parse text -- what the stop + # filter actually emitted -- is what the logprobs arrays are + # trimmed against below, BEFORE the thinking/inkling split or + # tool-call parsing can further diverge `text` from them (see + # _trim_generated_records_to_text's own docstring). + raw_text = text reasoning = "" if ARCH == "inkling": text, reasoning = split_inkling(text) @@ -3917,6 +4508,44 @@ def generation_stopped(): # into the visible answer / tool-call parser. reasoning, text = split_thinking_reply(text, enable_thinking) length_finish = "length" if stats["length_limited"] else "stop" + logprobs_obj = None + if engine_k: + channel = stats.get("logprobs") or {"prompt": [], "generated": []} + # A matched stop sequence withholds its own trailing + # text from `output` (StopFilter._emit is never called on + # it); that token's own logprob record must not survive + # into the response either, or the logprobs arrays would + # describe tokens the client never received. The + # thinking-split/tool-call divergence this does NOT close + # is a documented limitation (see docs/api.md). + generated = _trim_generated_records_to_text(channel["generated"], raw_text) + if chat: + logprobs_obj = {"content": _chat_logprobs_content(generated, display_k), + "refusal": None} + else: + logprobs_obj = _completions_logprobs_object( + channel["prompt"] if echo else [], generated, display_k) + if echo: + # The OpenAI legacy shape returns the prompt AND + # the completion concatenated in `text` when + # `echo` is true, and `text_offset` indexes into + # exactly that concatenation. `logprobs_obj + # ["tokens"]` already reconstructs prompt-then- + # generated text with ONE shared incremental + # decoder spanning both (_logprob_positions), so + # `text` is built from that same join here + # rather than by string-concatenating the + # prompt's reconstruction with the completion + # text decoded separately (by generate()'s own, + # independent decoder) -- a codepoint split + # across the last prompt byte and the first + # generated byte must be decoded as ONE + # character by ONE decoder, or the two halves + # decode as mojibake/replacement characters on + # either side of the seam and `text` and + # `text_offset` disagree about how long that + # character is. + text = "".join(logprobs_obj["tokens"]) if chat and tools: content, calls = parse_arch_tool_calls(text, tools, sideband.reply()) message = {"role": "assistant", "content": content or None, "refusal": None} @@ -3925,14 +4554,16 @@ def generation_stopped(): if calls: message["tool_calls"] = calls finish = "tool_calls" if calls else length_finish - choice = {"index": 0, "message": message, "logprobs": None, "finish_reason": finish} + choice = {"index": 0, "message": message, "logprobs": logprobs_obj, + "finish_reason": finish} else: _msg = {"role": "assistant", "content": text, "refusal": None} if reasoning: _msg["reasoning_content"] = reasoning choice = ({"index": 0, "message": _msg, - "logprobs": None, "finish_reason": length_finish} if chat else - {"index": 0, "text": text, "logprobs": None, "finish_reason": length_finish}) + "logprobs": logprobs_obj, "finish_reason": length_finish} if chat else + {"index": 0, "text": text, "logprobs": logprobs_obj, + "finish_reason": length_finish}) self.send_json(200, {"id": completion_id, "object": object_name, "created": created, "model": self.server.model_id, "choices": [choice], "usage": self.usage(stats)}, request_id, queue_headers) @@ -4086,7 +4717,8 @@ def generation_stopped(): self.client_disconnected, grammar=grammar, stopped=generation_stopped, **({"on_tool": sideband.feed} if sideband.enabled else {}), on_accept=start_stream, **({"audio": audio} if audio else {}), - **({"image": image} if image is not None else {})) + **({"image": image} if image is not None else {}), + **({"tok_ids": True} if tok_ids else {})) stop_filter.finish() sideband.finish() if think: @@ -4118,7 +4750,8 @@ def emit_plain(chunk): prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, self.client_disconnected, grammar=grammar, stopped=stop_filter.stopped, on_accept=start_stream, **({"audio": audio} if audio else {}), - **({"image": image} if image is not None else {})) + **({"image": image} if image is not None else {}), + **({"tok_ids": True} if tok_ids else {})) stop_filter.finish() if content_split: content_split.close() @@ -4349,18 +4982,13 @@ def generation_stopped(): request_id, queue_headers) return - self.send_response(200) - self.send_header("Content-Type", "text/event-stream") - self.send_header("Cache-Control", "no-cache") - self.send_header("X-Accel-Buffering", "no") - self.send_header("Connection", "close") # see the OpenAI path: SSE is close-framed - self.close_connection = True - self.send_header("x-request-id", request_id) - for name, value in queue_headers.items(): - self.send_header(name, value) - self.send_cors_headers() - self.end_headers() - connected = [True] + # Defer the HTTP 200 until the engine ACCEPTs the prompt: a refusal before + # ACCEPT -- today, CONTEXT_EXCEEDED -- must surface with its mapped HTTP status + # in the Anthropic error envelope rather than as a committed 200 followed by a + # truncated stream -- the same contract the OpenAI-style path already follows. + connected = [False] + stream_started = [False] + ka_thread = [None] write_lock = threading.Lock() last_write = [time.time()] ka_stop = threading.Event() @@ -4386,21 +5014,51 @@ def keepalive(): if time.time() - last_write[0] >= 10.0: send_event("ping", {"type": "ping"}) - send_event("message_start", {"type": "message_start", "message": { - "id": message_id, "type": "message", "role": "assistant", - "model": self.server.model_id, "content": [], "stop_reason": None, - "stop_sequence": None, "usage": {"input_tokens": 0, "output_tokens": 0}}}) text_index = 1 if enable_thinking else 0 stream_state = {"thinking_closed": not enable_thinking, "text_started": not enable_thinking} - if enable_thinking: - send_event("content_block_start", {"type": "content_block_start", "index": 0, - "content_block": {"type": "thinking", "thinking": "", "signature": ""}}) - else: - send_event("content_block_start", {"type": "content_block_start", "index": 0, - "content_block": {"type": "text", "text": ""}}) - ka_thread = threading.Thread(target=keepalive, daemon=True) - ka_thread.start() + + def start_stream(_accept_info=None): + # Commit the streaming 200 (and start the keepalive) exactly once, only after + # the engine ACCEPTs the prompt. Idempotent: also called after generate() + # returns, so an older engine with no ACCEPT frame still streams. + if stream_started[0]: + return + stream_started[0] = True + try: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("X-Accel-Buffering", "no") + self.send_header("Connection", "close") # see the OpenAI path: SSE is close-framed + self.close_connection = True + self.send_header("x-request-id", request_id) + for name, value in queue_headers.items(): + self.send_header(name, value) + self.send_cors_headers() + self.end_headers() + except OSError: + # The client vanished at the exact moment the engine accepted. Do not let + # this unwind out of generate()'s dispatch loop -- that would skip the + # CANCEL and leave the request stuck in the engine's pending map. Marking + # disconnected instead makes the cancelled() callback below fire on the + # loop's very next poll, so the normal cancel path takes it from here. + connected[0] = False + return + connected[0] = True + last_write[0] = time.time() + send_event("message_start", {"type": "message_start", "message": { + "id": message_id, "type": "message", "role": "assistant", + "model": self.server.model_id, "content": [], "stop_reason": None, + "stop_sequence": None, "usage": {"input_tokens": 0, "output_tokens": 0}}}) + if enable_thinking: + send_event("content_block_start", {"type": "content_block_start", "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}}) + else: + send_event("content_block_start", {"type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}}) + ka_thread[0] = threading.Thread(target=keepalive, daemon=True) + ka_thread[0].start() raw = [] sideband = ToolSideband(ARCH == "kimi" and bool(tools), stop_sequences, @@ -4475,8 +5133,14 @@ def generation_stopped(): stats = self.server.engine.generate( prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, - lambda: not connected[0], grammar=grammar, stopped=generation_stopped, + # Before the 200 commits, disconnect detection is the socket's (same as + # the OpenAI path); after it, a failed SSE write flips connected[0]. + lambda: (not connected[0]) if stream_started[0] else self.client_disconnected(), + grammar=grammar, stopped=generation_stopped, + on_accept=start_stream, **({"on_tool": sideband.feed} if sideband.enabled else {})) + # generate() returned, so the prompt was ACCEPTed and start_stream() ran; guard anyway. + start_stream() stop_filter.finish() sideband.finish() if split: @@ -4485,7 +5149,8 @@ def generation_stopped(): if tools and not state["in_tool"] and state["buf"]: emit_text(state["buf"]) ka_stop.set() - ka_thread.join(timeout=2) + if ka_thread[0] is not None: + ka_thread[0].join(timeout=2) if stream_state["text_started"]: send_event("content_block_stop", {"type": "content_block_stop", "index": text_index}) @@ -4510,13 +5175,346 @@ def generation_stopped(): send_event("message_stop", {"type": "message_stop"}) # close_connection was already set when the 200 was committed (#597 item 3). + def _completion_prompt_array(self, prompt_field): + """Classify an array `prompt` into its batch members. + + Accepted array forms (the two OpenAI legacy batch shapes, plus the + flat tokenized single prompt): + - flat [int, ...] -> ONE tokenized prompt (unchanged behavior) + - [str, ...] -> N string prompts + - [[int, ...], ...] -> N tokenized prompts (N=1 keeps the + batch-of-one unwrap byte-identical) + Mixed element types, an empty array, a batch above + PROMPT_BATCH_CAP, and a batch over the PROMPT_BATCH_TOKEN_BUDGET + aggregate boundary are each a named 400 with param `prompt` -- + whole-array conditions carry no member attribution. + + Returns (members, tok_ids): members is a non-empty list of prompts + (strings, or token-id lists still to be validated per member), + tok_ids says which kind. A real (N>1) batch is dispatched by the + caller through batch_completion() once this validation passes, so + a batch that is too large or over budget is refused here for + THAT specific reason, distinctly from anything batch_completion() + itself might later refuse. + """ + if not prompt_field: + raise APIError(400, "`prompt` array must not be empty.", "prompt", + "invalid_value") + if all(isinstance(item, str) for item in prompt_field): + members, tok_ids = list(prompt_field), False + elif all(isinstance(item, list) for item in prompt_field): + members, tok_ids = list(prompt_field), True + elif all(isinstance(item, int) and not isinstance(item, bool) + for item in prompt_field): + # One flat token-id prompt (a tokenized single shape sent as a + # bare array), not a batch: no cap or budget, and the + # single-prompt path downstream applies the glm gate and + # per-id validation exactly as before. + return [list(prompt_field)], True + else: + raise APIError(400, "`prompt` array must be homogeneous: all strings, " + "all token ids, or all token-id arrays.", "prompt", + "invalid_value") + if len(members) > PROMPT_BATCH_CAP: + # Its own code (prompt_batch_cap_exceeded) so a client or test + # can tell a size violation apart from every other array-prompt + # refusal. + raise APIError(400, f"`prompt` accepts at most {PROMPT_BATCH_CAP} prompts " + f"per request (got {len(members)}).", "prompt", + "prompt_batch_cap_exceeded") + if tok_ids and not getattr(self.server.engine, "supports_logprobs_echo", False): + # Same capability gate and message as the single-prompt path -- + # a nested batch must not smuggle token ids past it. + raise APIError(400, f"Token-ID array prompts are only supported by the glm " + f"engine (current engine: {ARCH}).", "prompt", + "unsupported_parameter") + if len(members) > 1: + # The aggregate-token boundary applies to real batches (N>1) + # only: a length-1 nested array must stay byte-identical to the + # flat single-prompt path, whose oversize handling remains the + # engine's own CONTEXT_EXCEEDED. + if tok_ids: + total, unit = sum(len(member) for member in members), "tokens" + else: + total = sum(len(member.encode("utf-8")) for member in members + if isinstance(member, str)) + unit = "UTF-8 bytes (an upper bound on tokens)" + if total > PROMPT_BATCH_TOKEN_BUDGET: + # Its own code (prompt_batch_token_budget_exceeded), same + # reasoning as the cap check above. This budget bounds the + # prompt side of the request only -- nothing here bounds + # the tokens a batch can be asked to generate. + raise APIError(400, f"`prompt` batch exceeds the total prompt budget: " + f"{total} {unit} across {len(members)} prompts " + f"(limit {PROMPT_BATCH_TOKEN_BUDGET}).", "prompt", + "prompt_batch_token_budget_exceeded") + return members, tok_ids + + def batch_completion(self, body, members, request_id, tok_ids): + """Multi-prompt batch dispatch for /v1/completions: one request, N + prompts, N indexed choices (choices[i]["index"] == i), each choice + carrying exactly what a single-prompt request for that same prompt + would have produced (text, the logprobs object under echo/logprobs, + finish_reason), and `usage` summed across prompts. + + A batch is all-or-nothing: every member's shape is validated before + the first engine submit, so a malformed member seven is a clean 400 + naming `prompt[7]` with no engine submits made. A member the engine + itself rejects (rather than a shape defect caught here) can still + fail after earlier members have already been submitted; see + submit_member below for how that failure is still member-named. + Engine dispatch, once it starts, is strictly sequential per-prompt + submits in request order, and stops the moment the client is gone; + nothing goes on the wire until every remaining member finished. + """ + # generation_options also rejects n != 1 with its own named 400: + # nothing here ever fans one prompt out to n samples. + maximum, temperature, top_p, grammar, _requested_stop_sequences = generation_options( + body, self.server.max_tokens) + family = family_by_id(ARCH) + if grammar is not None and not family.capabilities.grammar_payload: + raise APIError(400, f"`response_format` grammars are not supported by the {ARCH} " + "engine yet.", "response_format", "unsupported_parameter") + # Validate every member before the first engine submit: a malformed + # member seven must be a clean 400 with no engine submits made, + # never a failure discovered mid-batch -- and the 400 names the + # member (param "prompt[7]"). + encoded = [] + for index, member in enumerate(members): + if tok_ids: + try: + encoded.append(_encode_token_id_prompt(member)) + except APIError as error: + raise APIError(error.status, f"prompt[{index}]: {error.message}", + f"prompt[{index}]", error.code, error.error_type, + error.headers) + else: + if not member: + raise APIError(400, f"prompt[{index}]: `prompt` must not be empty.", + f"prompt[{index}]") + encoded.append(member) + # Bound the GENERATED side before any engine submit: every member + # shares this one `maximum` (already clamped to the operator's + # --max-tokens/--ngen above), so members * maximum is the batch's + # worst-case held-in-memory generated total. A batch of one member + # never reaches this function (it unwraps to the flat path), and the + # flat path itself is never subject to this budget. Checked after + # every member's own shape, not before: a malformed member still + # gets its own prompt[i] refusal even on a batch that is also over + # this budget, matching the shape-before-budget order the prompt-side + # budget already uses. + completion_total = len(members) * maximum + if completion_total > PROMPT_BATCH_COMPLETION_BUDGET: + if body.get("max_completion_tokens") is None and body.get("max_tokens") is None: + omitted_note = (f" `max_tokens` was not set, so this used the server's " + f"configured cap ({maximum}); set an explicit `max_tokens` " + f"on the request to fit this batch inside the budget.") + else: + omitted_note = "" + raise APIError(400, f"`prompt` batch would generate too many tokens: " + f"{len(members)} members x {maximum} max_tokens = " + f"{completion_total} (limit {PROMPT_BATCH_COMPLETION_BUDGET})." + f"{omitted_note}", + "max_tokens", "batch_completion_budget_exceeded") + engine_k, echo, display_k = logprobs_options( + body, False, getattr(self.server.engine, "supports_logprobs_echo", False)) + stop_sequences, ignore_leading_stop = stop_policy(body, False) + cache_slot = body.get("cache_slot") + if (cache_slot is not None and + (isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or + not 0 <= cache_slot < self.server.kv_slots)): + raise APIError(400, f"`cache_slot` must be an integer between 0 and {self.server.kv_slots - 1}.", + "cache_slot") + # Checked last, same position as the flat path's own stream check: + # a batch validates everything else about the request before + # refusing the one thing only a batch refuses outright. + stream = body.get("stream", False) + if not isinstance(stream, bool): + raise APIError(400, "`stream` must be a boolean.", "stream") + if stream: + # No interleaved multi-prompt SSE -- a named 400, not a silently + # non-streamed 200. + raise APIError(400, "`stream` is not supported with an array of prompts; " + "send one prompt per request to stream.", "stream", + "unsupported_parameter") + completion_id = "cmpl-" + uuid.uuid4().hex + created = int(time.time()) + try: + dbg = int(os.environ.get("COLI_DEBUG", "0")) + except ValueError: + dbg = 0 + + with self.server.scheduler.admit(self.client_disconnected, cache_slot) as admission: + queue_wait, cache_slot = admission + queue_headers = {"x-colibri-queue-wait-ms": str(round(queue_wait * 1000))} + + def submit_one(prompt): + """One engine submit plus choice-field assembly: the same + non-streaming single-prompt recipe generation() uses -- a + fresh StopFilter, and (inside _completions_logprobs_object) + a fresh per-sequence echo reassembly and stateful UTF-8 + decoder for this member alone. Sharing one incremental + decoder across members would smuggle one prompt's + trailing partial codepoint into the next prompt's first + token text. + """ + if dbg >= 2: + sys.stderr.write(f"\n===== PROMPT [{request_id}] =====\n{prompt}\n" + f"===== OUTPUT [{request_id}] =====\n") + sys.stderr.flush() + output = [] + stop_filter = StopFilter(stop_sequences, output.append, ignore_leading_stop) + stats = self.server.engine.generate( + prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, + self.client_disconnected, grammar=grammar, stopped=stop_filter.stopped, + **({"logprobs": engine_k, "echo": echo} if engine_k else {}), + **({"tok_ids": True} if tok_ids else {})) + stop_filter.finish() + raw_text = "".join(output) + text = raw_text + if ARCH == "inkling": + text, _reasoning = split_inkling(text) + finish = "length" if stats["length_limited"] else "stop" + logprobs_obj = None + if engine_k: + channel = stats.get("logprobs") or {"prompt": [], "generated": []} + generated = _trim_generated_records_to_text(channel["generated"], raw_text) + logprobs_obj = _completions_logprobs_object( + channel["prompt"] if echo else [], generated, display_k) + if echo: + # Same "text is the logprobs reconstruction" rule + # generation() applies -- one shared decoder spans + # the prompt-then-generated join for THIS member, + # never one string-concatenated from two separate + # decodes. + text = "".join(logprobs_obj["tokens"]) + return {"text": text, "logprobs": logprobs_obj, "finish_reason": finish}, stats + + def submit_member(index, prompt): + """A batch is all-or-nothing, so when one member sinks the + request the client must learn which one -- but only when + the member is actually at fault. + + A client-fault APIError from a member submit (the engine + rejecting that member's own content as malformed or + out-of-vocabulary) is re-raised with the failing member + named in both the message and `param` ("prompt[index]"); + status/code/type pass through unchanged. This already + covers a token-id member the engine rejects: for a + tok_ids submit, Engine.generate() itself turns the wire's + "ERROR BAD_REQUEST" frame into a client-fault + APIError(400, param="prompt") before this function ever + sees it. + + A server-fault APIError (error_type "server_error" -- the + engine failing to accept a per-token-logprobs or token-id + request in time) is not this member's doing: it keeps its + own status, code and `param` unchanged, and only its + message gains the member index, so a client is not told + to fix a prompt that was never the problem. + + Every other exception (protocol corruption, dispatcher + death, hostile frames, a matching-id engine BAD_REQUEST + that never became an APIError) propagates unattributed, + landing on the generic 500 engine_error the flat + single-prompt path already uses for the same failures -- + ONE clean engine_error for the whole batch, never a hang + and never a partial response. + """ + try: + return submit_one(prompt) + except APIError as error: + if error.error_type == "server_error": + raise APIError(error.status, f"prompt[{index}]: {error.message}", + error.param, error.code, error.error_type, + error.headers) + raise APIError(error.status, f"prompt[{index}]: {error.message}", + f"prompt[{index}]", error.code, error.error_type, + error.headers) + + # Sequential, in request order: member i+1 is never submitted + # until member i has fully finished and the client is still + # there to read the answer -- a departed client stops the + # batch at whichever member is in flight, and nothing is + # written to the client until every remaining member has + # finished, so a mid-batch failure on any member fails the + # whole request with no partial choices ever assembled. + outcomes = [] + for index, prompt in enumerate(encoded): + if self.client_disconnected(): + raise ClientCancelled() + outcomes.append(submit_member(index, prompt)) + + choices = [] + prompt_tokens = completion_tokens = 0 + for index, (choice, stats) in enumerate(outcomes): + choices.append({"index": index, **choice}) + prompt_tokens += stats["prompt_tokens"] + completion_tokens += stats["completion_tokens"] + usage = self.usage({"prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens}) + + # The scheduler admission (and this batch's hold on the engine) is + # released here, before the response is serialized and written -- + # a slow client reading a large batched response no longer holds + # the engine, or the clients queued behind it, hostage on socket + # I/O. A batch still occupies the engine for the sum of every + # member's generation time while the `with` block above is open; + # other clients queue behind it exactly as they would behind one + # long single-prompt request. + self.send_json(200, { + "id": completion_id, "object": "text_completion", "created": created, + "model": self.server.model_id, "choices": choices, "usage": usage}, + request_id, queue_headers) + def completion(self, body, request_id): - prompt = body.get("prompt") - if not isinstance(prompt, str): - raise APIError(400, "Colibri currently requires `prompt` to be a string.", "prompt") + # Group scoring is not wired up in this build: no engine or wire + # path routes it, and its contract changes the RESPONSE SHAPE + # (continuation-only logprob arrays), so silently ignoring the + # opt-in would hand a client a differently shaped answer than it + # asked for -- fail closed with a named 400 instead. Checked first, + # before either request shape below, so the array path can never + # reach batch_completion() with the opt-in still set. + group_score = body.get("group_score") + if group_score is not None and group_score is not False: + raise APIError(400, "`group_score` is not supported by this build.", + "group_score", "unsupported_value") + prompt_field = body.get("prompt") + tok_ids = False + if isinstance(prompt_field, list): + # Array intake: validate shape, cap, the glm gate, and the + # aggregate budget regardless of what happens next. A length-1 + # array unwraps to the single-prompt path below -- for + # [[ids]] that keeps the batch-of-one unwrap byte-identical + # (a tokenized-request client that always wraps its token-id + # array in an outer batch-of-one list, even at batch size 1, + # gets the same response the flat prompt would have produced). + # A real batch (N>1) dispatches through batch_completion(). + members, batch_tok_ids = self._completion_prompt_array(prompt_field) + if len(members) > 1: + self.batch_completion(body, members, request_id, batch_tok_ids) + return + prompt_field = members[0] + if isinstance(prompt_field, list): + # A tokenized single shape sends a token-ID array. Only the glm + # engine's mux_submit reads sub.tok_ids at all (coli_ids_parse) + # -- every other engine would tok_encode the decimal-digit + # string as if it were literal text, a silent wrong-answer, not + # a crash. Named 400 instead. + if not getattr(self.server.engine, "supports_logprobs_echo", False): + raise APIError(400, f"Token-ID array prompts are only supported by the glm " + f"engine (current engine: {ARCH}).", "prompt", + "unsupported_parameter") + prompt = _encode_token_id_prompt(prompt_field) + tok_ids = True + elif isinstance(prompt_field, str): + prompt = prompt_field + else: + raise APIError(400, "Colibri currently requires `prompt` to be a string or an " + "array of token ids.", "prompt") if not prompt: raise APIError(400, "`prompt` must not be empty.", "prompt") - self.generation(body, prompt, request_id, False) + self.generation(body, prompt, request_id, False, tok_ids=tok_ids) def serve(model, host="127.0.0.1", port=8000, model_id=None, api_key=None, diff --git a/c/tests/fixtures/captured_lmeval_request.json b/c/tests/fixtures/captured_lmeval_request.json new file mode 100644 index 000000000..909739f52 --- /dev/null +++ b/c/tests/fixtures/captured_lmeval_request.json @@ -0,0 +1,132 @@ +{ + "method": "POST", + "path": "/v1/completions", + "headers": { + "Host": "127.0.0.1:8812", + "User-Agent": "python-requests/2.34.2", + "Accept-Encoding": "gzip, deflate", + "Accept": "*/*", + "Connection": "keep-alive", + "Authorization": "Bearer ", + "Content-Length": "725", + "Content-Type": "application/json" + }, + "body_text": "{\"model\": \"glm-5.2-colibri\", \"prompt\": [[154822, 154824, 14572, 25, 55638, 614, 1012, 2952, 311, 5656, 20683, 504, 264, 34184, 1119, 279, 32616, 315, 264, 12722, 52938, 13, 3197, 279, 52938, 87894, 11, 432, 18611, 14065, 429, 5610, 264, 12827, 1730, 304, 34184, 80264, 13, 1096, 12827, 646, 1221, 387, 27341, 323, 1483, 311, 1281, 3746, 714, 29037, 7236, 369, 17409, 15768, 96457, 11, 77920, 11, 323, 14126, 13, 576, 2297, 1558, 537, 5240, 279, 52938, 894, 6646, 11, 714, 1045, 1251, 1744, 429, 432, 374, 537, 1290, 311, 912, 20683, 504, 279, 32616, 315, 825, 9859, 311, 2441, 13, 15908, 3409, 1850, 16539, 419, 4643, 5267, 16127, 25, 220, 6955]], \"temperature\": 0, \"max_tokens\": 1, \"logprobs\": 1, \"seed\": 1234, \"echo\": true}", + "body_json": { + "model": "glm-5.2-colibri", + "prompt": [ + [ + 154822, + 154824, + 14572, + 25, + 55638, + 614, + 1012, + 2952, + 311, + 5656, + 20683, + 504, + 264, + 34184, + 1119, + 279, + 32616, + 315, + 264, + 12722, + 52938, + 13, + 3197, + 279, + 52938, + 87894, + 11, + 432, + 18611, + 14065, + 429, + 5610, + 264, + 12827, + 1730, + 304, + 34184, + 80264, + 13, + 1096, + 12827, + 646, + 1221, + 387, + 27341, + 323, + 1483, + 311, + 1281, + 3746, + 714, + 29037, + 7236, + 369, + 17409, + 15768, + 96457, + 11, + 77920, + 11, + 323, + 14126, + 13, + 576, + 2297, + 1558, + 537, + 5240, + 279, + 52938, + 894, + 6646, + 11, + 714, + 1045, + 1251, + 1744, + 429, + 432, + 374, + 537, + 1290, + 311, + 912, + 20683, + 504, + 279, + 32616, + 315, + 825, + 9859, + 311, + 2441, + 13, + 15908, + 3409, + 1850, + 16539, + 419, + 4643, + 5267, + 16127, + 25, + 220, + 6955 + ] + ], + "temperature": 0, + "max_tokens": 1, + "logprobs": 1, + "seed": 1234, + "echo": true + } +} \ No newline at end of file diff --git a/c/tests/golden_fixture_capture.py b/c/tests/golden_fixture_capture.py index 9ae6424e1..e86521340 100644 --- a/c/tests/golden_fixture_capture.py +++ b/c/tests/golden_fixture_capture.py @@ -22,8 +22,9 @@ python3 tests/golden_fixture_capture.py diff fixtures_pre fixtures_post Battery: chat, chat+tools, chat streaming, completions without logprobs, and -the error cases whose behavior must not move (seed 400, array-prompt 400, -logprobs 400, out-of-range temperature 400) plus /v1/models. +the error/no-op cases whose behavior must not move (seed accepted-and-ignored, +an array `prompt` classified as one flat token-id prompt, logprobs served on +the glm engine, out-of-range temperature 400) plus /v1/models. Normalization: every "id"/"created" field (recursively, and per SSE event) is replaced with a constant; nothing else is touched. Generation-bearing requests @@ -76,11 +77,11 @@ def battery(model): ("completions_stop", "POST", "/v1/completions", {"model": model, "prompt": "Count: one, two,", "max_tokens": 16, "temperature": 0, "stop": ["five"]}), - ("err_seed", "POST", "/v1/completions", + ("seed_accepted", "POST", "/v1/completions", {"model": model, "prompt": "hello", "max_tokens": 1, "seed": 1234}), - ("err_array_prompt", "POST", "/v1/completions", + ("array_prompt_token_ids", "POST", "/v1/completions", {"model": model, "prompt": [1, 2, 3], "max_tokens": 1}), - ("err_logprobs", "POST", "/v1/completions", + ("logprobs_served", "POST", "/v1/completions", {"model": model, "prompt": "hello", "max_tokens": 1, "logprobs": 1}), ("err_bad_temperature", "POST", "/v1/chat/completions", {"model": model, "messages": chat_messages, diff --git a/c/tests/test_anthropic_messages.py b/c/tests/test_anthropic_messages.py index b8006b441..7e309deec 100644 --- a/c/tests/test_anthropic_messages.py +++ b/c/tests/test_anthropic_messages.py @@ -34,6 +34,8 @@ def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, on_tool=None): self.prompts.append(prompt) self.emitted = 0 + if on_accept is not None: # simulate the engine's ACCEPT frame (#597) + on_accept({"prompt_tokens": 11}) for chunk in self.script: on_text(chunk) self.emitted += 1 diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index b08d5dff7..217c01b89 100644 --- a/c/tests/test_openai_server.py +++ b/c/tests/test_openai_server.py @@ -1,9 +1,13 @@ import http.client +import inspect import io import json import math import os +import queue +import signal import socket +import subprocess import tempfile import threading import sys @@ -17,24 +21,44 @@ from openai_server import (APIError, APIHandler, APIServer, ClientCancelled, DEFAULT_CHAT_STOP_SEQUENCES, END, GenerationScheduler, - READY, Engine, InklingStreamSplit, StopFilter, ThinkingStreamSplit, - _engine_error, _image_bytes_from_url, cap_for_arch, conversation_cache_slot, model_arch, - generation_options, parse_tool_calls, parse_dsv4_tool_calls, + LOGPROBS_TOP_K_CAP, PROMPT_BATCH_CAP, PROMPT_BATCH_COMPLETION_BUDGET, + PROMPT_BATCH_TOKEN_BUDGET, + READY, Engine, InklingStreamSplit, StopFilter, + ThinkingStreamSplit, + _engine_error, _image_bytes_from_url, cap_for_arch, + conversation_cache_slot, model_arch, + generation_options, logprobs_options, parse_tool_calls, + parse_dsv4_tool_calls, parse_arch_tool_calls, parse_k3_tool_calls, parse_qwen38_tool_calls, read_engine_turn, render_chat, render_chat_kimi, render_chat_olmoe, render_chat_qwen38, render_chat_v4, _dsv4_tool_calls, serve, split_thinking_reply, - stop_policy, tune_child_env) + stop_policy, tune_child_env, + _chat_logprobs_content, _completions_logprobs_object, + _encode_token_id_prompt, + _json_float, _order_echo_records, _own_token_label, + _trim_generated_records_to_text) class FakeEngine: + # The per-token logprobs capability gate: glm-only in production + # (Engine.__init__ sets it from `arch == "glm"`). Tests run under the + # module's default ARCH="glm" unless a test patches it, so True is the + # representative default for this double; capability-gate tests override + # it explicitly (see NonGlmEngine below) rather than relying on this. + supports_logprobs_echo = True + def __init__(self): self.calls = [] self.stop_requests = 0 def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None, on_accept=None): + cancelled=None, grammar=None, stopped=None, on_accept=None, logprobs=0, + echo=False, tok_ids=False): self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + self.last_logprobs = logprobs + self.last_echo = echo + self.last_tok_ids = tok_ids if on_accept is not None: # simulate the engine's ACCEPT frame (#597) on_accept({"prompt_tokens": 7}) for chunk in ("Hé", "llo"): @@ -42,7 +66,34 @@ def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, if stopped and stopped(): self.stop_requests += 1 break - return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False} + stats = {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False} + if logprobs: + stats["logprobs"] = self.logprobs_channel(logprobs) + return stats + + def logprobs_channel(self, engine_k): + """Canned U7a logprob records for HTTP-level response-shape tests -- + shaped exactly like Engine.generate()'s real return value: "prompt" + is a list of (pos, bytes, record), "generated" a list of (bytes, + record), record = {"lp": float, "topk": [(tid, tlp), ...]}. Position + 0 carries the engine's own "nothing to condition on" sentinel (nan, + empty table) -- the real wire behavior mux_prefill_echo always sends. + The tail values are NON-dyadic (e.g. -0.3, -2.7) so a `%.6f`-shaped + fixture is distinguishable from a `%.17g`-shaped one in tests. "H"=72 + "\\xc3\\xa9"="é" pretend token ids; the generated tokens' own lp is + bit-identical to one of their own topk entries (mux's logprob_tail + invariant), which is what the bit-identity property test and + _own_token_label depend on.""" + k = min(engine_k, 2) + prompt = [ + (0, b"H", {"lp": float("nan"), "topk": []}), + (1, b"\xc3\xa9", {"lp": -0.3, "topk": [(72, -0.3), (100, -1.7)][:k]}), + ] + generated = [ + (b"H", {"lp": -0.2, "topk": [(72, -0.2), (200, -2.4)][:k]}), + (b"\xc3\xa9", {"lp": -0.4, "topk": [(101, -0.4), (300, -3.6)][:k]}), + ] + return {"prompt": prompt, "generated": generated} class BlockingEngine(FakeEngine): @@ -52,11 +103,12 @@ def __init__(self): self.release = threading.Event() def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None, on_accept=None): + cancelled=None, grammar=None, stopped=None, on_accept=None, logprobs=0, + echo=False, tok_ids=False): self.entered.set() self.release.wait(2) return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot, - cancelled, grammar, stopped, on_accept) + cancelled, grammar, stopped, on_accept, logprobs, echo, tok_ids) class TemplateTest(unittest.TestCase): @@ -374,6 +426,11 @@ def test_validates_generation_limits(self): opts = generation_options({"response_format": {"type": "gbnf", "grammar": "not a grammar ::="}}, 8) self.assertEqual(opts[3], "not a grammar ::=") + def test_seed_no_longer_rejected_by_generation_options(self): + # generation_options() used to 400 on any `seed`; it is now a silent + # accept-and-discard (documented no-op). + generation_options({"seed": 1234, "prompt": "hi"}, 16) # must not raise + def test_coli_temp_is_the_default_for_requests_that_omit_temperature(self): with patch.dict("openai_server.os.environ", {"COLI_TEMP": "0.25"}): self.assertEqual(generation_options({}, 8)[1], 0.25) @@ -406,6 +463,482 @@ def test_glm_chat_defaults_role_stops_without_changing_other_policies(self): stop_policy({"x_colibri_ignore_leading_stop": "yes"}, True) +class LogprobsOptionsTest(unittest.TestCase): + """logprobs_options(): pure validation/translation, no HTTP or engine. + + Covers the range checks and the zero/false/null semantics, at the unit + level -- fast, and independent of any fixture. + """ + + def test_completions_valid_integer_logprobs(self): + self.assertEqual(logprobs_options({"logprobs": 3}, False, True), (3, False, 3)) + self.assertEqual(logprobs_options({"logprobs": 3, "echo": True}, False, True), + (3, True, 3)) + self.assertEqual(logprobs_options({"logprobs": 1}, False, True), (1, False, 1)) + self.assertEqual(logprobs_options({"logprobs": LOGPROBS_TOP_K_CAP}, False, True), + (LOGPROBS_TOP_K_CAP, False, LOGPROBS_TOP_K_CAP)) + + def test_completions_zero_false_null_mean_no_logprobs(self): + # The zero semantics are explicit -- `0`, `false`, and `null` all + # mean "no logprobs" on completions, never a truthiness accident + # and never an engine channel floored on at k=1. + for off in ({"logprobs": 0}, {"logprobs": False}, {"logprobs": None}, {}): + with self.subTest(off=off): + self.assertEqual(logprobs_options(off, False, True), (0, False, 0)) + self.assertEqual(logprobs_options({"logprobs": 0, "echo": True}, False, True), + (0, True, 0)) + + def test_completions_true_is_a_named_400(self): + # The legacy completions field is an integer COUNT; a boolean + # `true` carries no count and is a named 400, not a guess at k. + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": True}, False, True) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.param, "logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + + def test_break_it_battery_non_integer_negative_huge(self): + for bad in (1.5, "5", -1, LOGPROBS_TOP_K_CAP + 1): + with self.subTest(bad=bad): + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": bad}, False, True) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.param, "logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + + def test_completions_echo_requires_a_boolean(self): + with self.assertRaises(APIError) as caught: + logprobs_options({"echo": "yes"}, False, True) + self.assertEqual(caught.exception.param, "echo") + + def test_chat_logprobs_requires_a_boolean(self): + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": 1}, True, True) + self.assertEqual(caught.exception.param, "logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + + def test_chat_false_and_null_mean_no_logprobs(self): + # Chat's boolean gate treats `null` like `false` -- explicitly. + for off in ({"logprobs": False}, {"logprobs": None}, {}): + with self.subTest(off=off): + self.assertEqual(logprobs_options(off, True, True), (0, False, 0)) + + def test_chat_echo_is_always_rejected(self): + # Chat has no echo concept at all -- a named 400, not a silent + # ignore, whether or not logprobs was also requested. + with self.assertRaises(APIError) as caught: + logprobs_options({"echo": True}, True, True) + self.assertEqual(caught.exception.param, "echo") + with self.assertRaises(APIError) as caught: + logprobs_options({"echo": True, "logprobs": True}, True, True) + self.assertEqual(caught.exception.param, "echo") + + def test_chat_top_logprobs_default_and_cap(self): + self.assertEqual(logprobs_options({"logprobs": True}, True, True), (1, False, 0)) + self.assertEqual( + logprobs_options({"logprobs": True, "top_logprobs": 5}, True, True), (5, False, 5)) + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": True, "top_logprobs": LOGPROBS_TOP_K_CAP + 1}, + True, True) + self.assertEqual(caught.exception.param, "top_logprobs") + + def test_capability_gate_rejects_unsupported_engine(self): + # The server never emits logprobs= to an engine that does not + # implement the numeric per-token channel -- a named 400, not a + # silent downgrade to "no logprobs". + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": 1}, False, False) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.code, "unsupported_parameter") + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": True}, True, False) + self.assertEqual(caught.exception.status, 400) + # Absent/zero logprobs never reach the capability check at all. + self.assertEqual(logprobs_options({}, False, False), (0, False, 0)) + self.assertEqual(logprobs_options({"logprobs": 0}, False, False), (0, False, 0)) + + def test_range_error_precedes_capability_error(self): + # The named 400 above the cap fires even on a non-supporting + # engine -- the range check is about the public API surface, not + # the engine. + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": LOGPROBS_TOP_K_CAP + 1}, False, False) + self.assertEqual(caught.exception.code, "invalid_value") + + def test_cap_pinned_to_engine_topk_maximum(self): + # Mirrors the engine's per-request top-k maximum, the + # COLI_SUBMIT_TOPK_MAX constant defined in c/decode_batch.h:19. A + # literal check (not the LOGPROBS_TOP_K_CAP symbol) so a change to + # either side is caught, instead of the test drifting in lockstep + # with the constant it exists to pin. + self.assertEqual(LOGPROBS_TOP_K_CAP, 32) + + def test_cap_boundary_literals_both_endpoints(self): + # Literal 32/33, not LOGPROBS_TOP_K_CAP +/- 1: a boundary check + # that stays meaningful even if the cap constant itself drifts. + self.assertEqual(logprobs_options({"logprobs": 32}, False, True), (32, False, 32)) + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": 33}, False, True) + self.assertEqual(caught.exception.param, "logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + self.assertEqual( + logprobs_options({"logprobs": True, "top_logprobs": 32}, True, True), + (32, False, 32)) + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": True, "top_logprobs": 33}, True, True) + self.assertEqual(caught.exception.param, "top_logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + + def test_chat_top_logprobs_validated_even_when_logprobs_is_off(self): + # Pre-existing defect found by review: with `logprobs` + # false/absent, `top_logprobs` used to be ignored outright -- + # 999 and "x" both passed silently. It is now TYPE- and + # RANGE-checked regardless, with the same named 400s as when + # `logprobs` is true; a genuinely valid value stays a documented + # no-op. + for bad in ("x", -1, 999): + with self.subTest(bad=bad): + with self.assertRaises(APIError) as caught: + logprobs_options({"top_logprobs": bad}, True, True) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.param, "top_logprobs") + self.assertEqual(caught.exception.code, "invalid_value") + self.assertEqual(logprobs_options({"top_logprobs": 5}, True, True), (0, False, 0)) + self.assertEqual(logprobs_options({}, True, True), (0, False, 0)) + + def test_echo_non_bool_rejected_same_way_both_endpoints(self): + # Completions already checked echo's type with isinstance(bool); + # chat's refusal used truthiness. Both endpoints now validate + # echo's TYPE first, with the shared "must be a boolean" 400 -- + # chat's "not supported for chat completions" refusal now only + # ever fires for an actual boolean True. + for bad in (1, "true"): + with self.subTest(endpoint="completions", bad=bad): + with self.assertRaises(APIError) as caught: + logprobs_options({"echo": bad}, False, True) + self.assertEqual(caught.exception.param, "echo") + self.assertEqual(caught.exception.code, "invalid_value") + with self.subTest(endpoint="chat", bad=bad): + with self.assertRaises(APIError) as caught: + logprobs_options({"echo": bad}, True, True) + self.assertEqual(caught.exception.param, "echo") + self.assertEqual(caught.exception.code, "invalid_value") + + def test_range_error_message_states_the_cap_value(self): + with self.assertRaises(APIError) as caught: + logprobs_options({"logprobs": 999}, False, True) + self.assertIn("32", caught.exception.message) + + def test_chat_top_logprobs_null_normalizes_to_absent(self): + # `top_logprobs: null` means the same thing as the field being + # absent -- exactly like `logprobs: null` -- on both sides of the + # `logprobs` gate. It never itself reaches the type/range check. + self.assertEqual( + logprobs_options({"logprobs": False, "top_logprobs": None}, True, True), + (0, False, 0)) + self.assertEqual( + logprobs_options({"logprobs": True, "top_logprobs": None}, True, True), + (1, False, 0)) + + def test_completions_ignores_top_logprobs_entirely(self): + # `top_logprobs` is a chat-only field in the OpenAI request shape; + # completions never reads it, so a nonsense value there changes + # nothing -- the result is exactly the completions tuple for + # `logprobs: 2` alone. + self.assertEqual( + logprobs_options({"logprobs": 2, "top_logprobs": 999}, False, True), + (2, False, 2)) + + +class ResponseAssemblyTest(unittest.TestCase): + """Pure response-shape builders: _json_float, _order_echo_records, + _own_token_label, _completions_logprobs_object, _chat_logprobs_content. + No HTTP, no engine -- fixtures hand-build the (pos, bytes, record) and + (bytes, record) tuples the wire dispatcher would otherwise produce.""" + + def test_json_float_maps_non_finite_to_none(self): + # Non-finite becomes JSON null on the wire out, never a clamped + # number and never a raw float (json.dumps would emit the + # invalid-JSON NaN/Infinity literal for that). + self.assertIsNone(_json_float(float("nan"))) + self.assertIsNone(_json_float(float("inf"))) + self.assertIsNone(_json_float(float("-inf"))) + self.assertEqual(_json_float(-0.5), -0.5) + self.assertEqual(json.dumps({"x": _json_float(float("nan"))}), '{"x": null}') + + def test_own_token_label_matches_by_value_not_rank_or_id(self): + # The chosen token need not sit first in the top-k table (the + # table is unsorted on the wire) -- the label must be found by + # exact float match against the position's own logprob, not by + # assuming table position 0 and not by any id-to-text mapping. + entry = {"text": "café", "raw_lp": -0.8} + topk = [(999, -0.05), (42, -0.8)] + self.assertEqual(_own_token_label(entry, topk, 0), "") + self.assertEqual(_own_token_label(entry, topk, 1), "café") + + def test_own_token_label_tie_break_deterministic_first_match(self): + # The engine prints logprobs to 6 decimal digits, so two distinct + # candidates can share the exact printed value the chosen token + # also carries -- a genuine tie the record's own shape (lp + a + # topk table of raw ids, no chosen-token id) cannot resolve by + # identity. Documented behavior: the FIRST table entry (in wire + # order) whose value matches is labeled as the chosen token; a + # later entry that also matches is labeled by its raw id like any + # other unidentified candidate -- deterministic, not a claim that + # the first entry is provably the true chosen one. + entry = {"text": "cat", "raw_lp": -0.223144} + topk = [(1, -0.223144), (2, -0.223144)] + self.assertEqual(_own_token_label(entry, topk, 0), "cat") + self.assertEqual(_own_token_label(entry, topk, 1), "") + + def test_completions_logprobs_object_first_prompt_token_is_null(self): + prompt = [(0, b"The", {"lp": float("nan"), "topk": []}), + (1, b" cat", {"lp": -0.3, "topk": [(7, -0.3), (8, -1.1)]})] + obj = _completions_logprobs_object(prompt, [], display_k=2) + self.assertEqual(obj["tokens"], ["The", " cat"]) + self.assertIsNone(obj["token_logprobs"][0]) + self.assertEqual(obj["token_logprobs"][1], -0.3) + self.assertEqual(obj["top_logprobs"][0], {}) + + def test_text_offset_reconstruction_and_monotonic(self): + # "".join(tokens) reproduces the prompt text, and text_offset + # matches the exact per-position character counts -- a literal + # expected list, not a self-referential sortedness/start-at-0 + # check that every possible offset sequence satisfies by + # construction (offsets are a running sum of non-negative token + # lengths, so both of those would hold even for a wrong sequence). + prompt_text = "Hé said éé" + pieces = ["H", "é", " said ", "é", "é"] + prompt = [(i, p.encode("utf-8"), {"lp": float("nan") if i == 0 else -0.1, "topk": []}) + for i, p in enumerate(pieces)] + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual("".join(obj["tokens"]), prompt_text) + self.assertEqual(obj["text_offset"], [0, 1, 2, 8, 9]) + + def test_trailing_incomplete_multibyte_sequence_flushed_in_completions_path(self): + # The same stateful-decoder trailing flush _chat_logprobs_content + # relies on lives in the shared _logprob_positions helper -- prove + # it from the completions/echo side too, not only via chat, so a + # regression in the shared helper that happens to leave the chat + # test green cannot slip through. + prompt = [(0, b"H", {"lp": float("nan"), "topk": []}), + (1, b"\xc3", {"lp": -0.1, "topk": []})] # first byte of 'é', never completed + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual(obj["tokens"][0], "H") + self.assertIn("\ufffd", obj["tokens"][1]) + + def test_bit_identity_when_chosen_token_is_argmax(self): + # token_logprobs[i] equals top_logprobs[i][tokens[i]] exactly when + # token i is itself the argmax of its own table -- checked against + # the fixture's own argmax logprob literal (-0.05) on both sides, + # not by comparing two fields of the same call's output to each + # other (which a shared, consistently-wrong source would pass). + generated = [(b"cat", {"lp": -0.05, "topk": [(7, -0.05), (8, -3.0)]})] + obj = _completions_logprobs_object([], generated, display_k=2) + self.assertEqual(obj["token_logprobs"][0], -0.05) + self.assertEqual(obj["top_logprobs"][0]["cat"], -0.05) + + def test_display_k_truncates_alternatives_independent_of_engine_k(self): + generated = [(b"cat", {"lp": -0.05, "topk": [(7, -0.05), (8, -3.0), (9, -4.0)]})] + obj = _completions_logprobs_object([], generated, display_k=1) + self.assertEqual(len(obj["top_logprobs"][0]), 1) + obj0 = _completions_logprobs_object([], generated, display_k=0) + self.assertEqual(obj0["top_logprobs"][0], {}) + self.assertEqual(obj0["token_logprobs"][0], -0.05) # still reported + + def test_chat_content_shape_and_no_echo_field(self): + generated = [(b"cat", {"lp": -0.05, "topk": [(7, -0.05), (8, -3.0)]})] + content = _chat_logprobs_content(generated, display_k=2) + self.assertEqual(len(content), 1) + entry = content[0] + self.assertEqual(set(entry), {"token", "logprob", "bytes", "top_logprobs"}) + self.assertEqual(entry["token"], "cat") + self.assertEqual(entry["logprob"], -0.05) + self.assertEqual(entry["bytes"], [99, 97, 116]) + self.assertEqual(len(entry["top_logprobs"]), 2) + for alt in entry["top_logprobs"]: + self.assertEqual(set(alt), {"token", "logprob", "bytes"}) + own = [a for a in entry["top_logprobs"] if a["token"] == "cat"][0] + self.assertEqual(own["logprob"], -0.05) + self.assertEqual(own["bytes"], [99, 97, 116]) + other = [a for a in entry["top_logprobs"] if a["token"] != "cat"][0] + self.assertIsNone(other["bytes"]) + self.assertNotIn("echo", json.dumps(content)) + + def test_chat_content_nan_serializes_as_null(self): + generated = [(b"x", {"lp": float("-inf"), "topk": [(1, float("-inf"))]})] + content = _chat_logprobs_content(generated, display_k=1) + self.assertIsNone(content[0]["logprob"]) + self.assertIsNone(content[0]["top_logprobs"][0]["logprob"]) + + def test_encode_token_id_prompt_round_trips_and_validates(self): + self.assertEqual(_encode_token_id_prompt([1, 2, 30000]), "1 2 30000") + for bad in ([], [1, -1], [1, 2.5], [1, True], "not-a-list", None): + with self.subTest(bad=bad): + with self.assertRaises(APIError) as caught: + _encode_token_id_prompt(bad) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.param, "prompt") + + # Adversarial top-k order: the chosen token's own candidate sits + # second in the table and is not even the highest logprob present (a + # sampled-path shape) -- no code here may assume table position 0. + + def test_token_logprobs_sourced_from_chosen_record_not_table_position(self): + generated = [(b"cat", {"lp": -0.8, "topk": [(999, -0.05), (7, -0.8)]})] + obj = _completions_logprobs_object([], generated, display_k=2) + self.assertEqual(obj["token_logprobs"][0], -0.8) + self.assertEqual(obj["top_logprobs"][0]["cat"], -0.8) + self.assertNotEqual(obj["token_logprobs"][0], -0.05, + "token_logprobs must not be sourced from topk[0]") + + def test_chat_content_logprob_sourced_from_chosen_record_not_table_position(self): + generated = [(b"cat", {"lp": -0.8, "topk": [(999, -0.05), (7, -0.8)]})] + content = _chat_logprobs_content(generated, display_k=2) + self.assertEqual(content[0]["logprob"], -0.8) + own = [a for a in content[0]["top_logprobs"] if a["token"] == "cat"][0] + self.assertEqual(own["logprob"], -0.8) + + # Numeric text_offset: exact values, not just monotonic/starts-at-0, + # over a fixture with a multi-byte character ("é": one Python + # character, two UTF-8 bytes) so a byte-vs-character-count confusion + # would fail. + + def test_text_offset_exact_values_with_multibyte_character(self): + prompt = [(0, b"H", {"lp": float("nan"), "topk": []}), + (1, b"\xc3\xa9", {"lp": -0.1, "topk": []}), # "é", complete on its own + (2, b" cat", {"lp": -0.2, "topk": []})] + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual(obj["tokens"], ["H", "é", " cat"]) + self.assertEqual(obj["text_offset"], [0, 1, 2]) + + # Multi-byte character split across two adjacent frames -- must + # reconstruct via the stateful incremental decoder instead of + # mangling into two replacement-character halves. + + def test_multibyte_character_split_across_adjacent_tokens_reconstructs(self): + prompt_text = "café" + prompt = [(0, b"c", {"lp": float("nan"), "topk": []}), + (1, b"a", {"lp": -0.1, "topk": []}), + (2, b"f", {"lp": -0.1, "topk": []}), + (3, b"\xc3", {"lp": -0.1, "topk": []}), # first byte of 'é' + (4, b"\xa9", {"lp": -0.1, "topk": []})] # second byte of 'é' + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual("".join(obj["tokens"]), prompt_text) + self.assertEqual(obj["tokens"], ["c", "a", "f", "", "é"]) + + def test_multibyte_character_split_across_generated_tokens_reconstructs(self): + generated = [(b"\xc3", {"lp": -0.1, "topk": []}), (b"\xa9", {"lp": -0.1, "topk": []})] + content = _chat_logprobs_content(generated, display_k=0) + self.assertEqual("".join(c["token"] for c in content), "é") + self.assertEqual([c["token"] for c in content], ["", "é"]) + # `bytes` always stays the frame's own raw payload, independent + # of what text (if any) it resolved to. + self.assertEqual(content[0]["bytes"], [0xC3]) + self.assertEqual(content[1]["bytes"], [0xA9]) + + def test_trailing_incomplete_multibyte_sequence_still_flushed(self): + # A dangling partial sequence at the very end of the stream (no + # more data ever completes it) must still surface via the final + # flush, not silently vanish. + generated = [(b"cat", {"lp": -0.1, "topk": []}), (b"\xc3", {"lp": -0.1, "topk": []})] + content = _chat_logprobs_content(generated, display_k=0) + self.assertEqual(content[0]["token"], "cat") + self.assertIn("�", content[1]["token"]) + + def test_chat_content_tail_flush_consistent_with_own_alternative_label(self): + # The trailing-flush text must reach the chosen token's OWN + # top_logprobs entry, not just the outer `token` field -- decoding + # is done in one pass over all positions (tail included) before + # any content entry is built, so the two can never disagree about + # what the last position's text actually is. + generated = [(b"\xc3", {"lp": -0.1, "topk": [(1, -0.1)]})] # never completed + content = _chat_logprobs_content(generated, display_k=1) + self.assertIn("�", content[0]["token"]) + self.assertEqual(content[0]["top_logprobs"][0]["token"], content[0]["token"]) + + def test_completions_top_logprobs_tie_break_deterministic_first_match(self): + # Two candidates print the identical 6-decimal logprob the chosen + # token also carries. The FIRST exact match in wire order is + # labeled as the chosen token; the second gets its own distinct + # id-labeled entry rather than being silently merged into the + # first (a dict keyed only by label would otherwise collapse two + # genuinely different candidates into one entry). + generated = [(b"cat", {"lp": -0.223144, "topk": [(1, -0.223144), (2, -0.223144)]})] + obj = _completions_logprobs_object([], generated, display_k=2) + self.assertEqual(set(obj["top_logprobs"][0]), {"cat", ""}) + + # ECHO's wire `pos` field, not arrival order, decides placement. + + def test_echo_positions_reassembled_by_wire_pos_not_arrival_order(self): + prompt = [(1, b"b", {"lp": -1.0, "topk": []}), # delivered out of order + (0, b"a", {"lp": float("nan"), "topk": []})] + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual(obj["tokens"], ["a", "b"]) + self.assertIsNone(obj["token_logprobs"][0]) + self.assertEqual(obj["token_logprobs"][1], -1.0) + + def test_order_echo_records_repeated_token_prompt_not_fooled_by_join_check(self): + # A join-only check ("abab") can pass even when positions are + # corrupted, because repeated tokens make the wrong order look + # identical to the right one under "".join(). Position-indexed + # assembly must get the per-position values right regardless. + prompt = [(0, b"a", {"lp": float("nan"), "topk": []}), + (2, b"a", {"lp": -2.0, "topk": []}), + (1, b"b", {"lp": -1.0, "topk": []}), + (3, b"b", {"lp": -3.0, "topk": []})] + obj = _completions_logprobs_object(prompt, [], display_k=0) + self.assertEqual(obj["tokens"], ["a", "b", "a", "b"]) + self.assertEqual("".join(obj["tokens"]), "abab") + self.assertEqual(obj["token_logprobs"], [None, -1.0, -2.0, -3.0]) + + def test_duplicate_echo_position_raises_named_error(self): + prompt = [(0, b"a", {"lp": float("nan"), "topk": []}), + (0, b"a2", {"lp": -1.0, "topk": []})] + with self.assertRaisesRegex(RuntimeError, "invalid engine ECHO position"): + _order_echo_records(prompt) + with self.assertRaises(RuntimeError): + _completions_logprobs_object(prompt, [], display_k=0) + + def test_out_of_range_echo_position_raises_named_error(self): + prompt = [(0, b"a", {"lp": float("nan"), "topk": []}), + (5, b"b", {"lp": -1.0, "topk": []})] + with self.assertRaisesRegex(RuntimeError, "invalid engine ECHO position"): + _order_echo_records(prompt) + + def test_negative_echo_position_raises_named_error(self): + prompt = [(-1, b"a", {"lp": float("nan"), "topk": []})] + with self.assertRaises(RuntimeError): + _order_echo_records(prompt) + + +class TrimGeneratedRecordsTest(unittest.TestCase): + """_trim_generated_records_to_text: a helper the flat and batch paths + both call to drop trailing generated-token records whose bytes were + filtered out of `text` (a matched stop sequence, most commonly), so + the logprobs table stays aligned with what is actually returned. Each + candidate piece must be checked against `text` at the RUNNING offset + -- the byte count of every record kept so far -- never from the start + of `text`, and that offset must advance by exactly the decoded piece + length, not off by one.""" + + def test_all_records_kept_when_text_holds_every_piece(self): + # Three records of unequal length, none trimmed: proves the + # running offset lands on the true boundary between every pair + # (2, then 5), not merely that the whole concatenation matches. + records = [(b"abc", {"r": 1}), (b"de", {"r": 2}), (b"fgh", {"r": 3})] + self.assertEqual(_trim_generated_records_to_text(records, "abcdefgh"), records) + + def test_records_after_a_stop_sequence_truncation_are_dropped(self): + # `text` was cut short by a matched stop sequence after "ab" -- + # the second and third records' bytes are no longer a prefix of + # what is actually returned, so both are dropped, not just the + # record whose own bytes fail to match. + records = [(b"ab", {"r": 1}), (b"cd", {"r": 2}), (b"ef", {"r": 3})] + self.assertEqual(_trim_generated_records_to_text(records, "ab"), + [(b"ab", {"r": 1})]) + + class StopFilterTest(unittest.TestCase): def test_explicit_stop_composes_with_inkling_stream_split(self): content = [] @@ -946,9 +1479,11 @@ def test_accepts_u7a_echo_and_extended_data_frames(self): # U7a forward-compat: the engine's opt-in per-token numeric channel -- # ECHO frames for echoed prompt positions and DATA frames extended # with " [tid tlp]*k" -- must NOT trip the dispatcher's - # catch-all (which kills every in-flight request). Text delivery and - # the DONE stats stay exactly as for legacy frames; the numeric - # fields are consumed by the server feature half (U7b). + # catch-all (which kills every in-flight request), and must not + # change the legacy response shape: text delivery and the DONE + # stats stay exactly as for legacy frames. The numeric records + # themselves are collected internally; response assembly reading + # them back out is separate, later work. def respond(process, frame): request_id = frame.split()[1] process.stdout.feed(b"ACCEPT " + request_id + b" 3\n") @@ -976,6 +1511,330 @@ def respond(process, frame): self.assertIsNone(engine.dispatcher_error) engine.close() + def test_legacy_data_frame_dispatches_bare_bytes_no_record(self): + # A DATA frame WITHOUT a tail dispatches exactly as on the + # predecessor dispatcher -- the same event tuple, same bytes, no + # record object allocated. Guards against a record being allocated + # unconditionally on the (far more common) non-opted-in path. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed(b"DATA " + request_id.encode() + b" 2\nok\n") + kind, value = events.get(timeout=1) + self.assertEqual(kind, "data") + self.assertIs(type(value), bytes) + self.assertEqual(value, b"ok") + engine.close() + + def test_echo_position_zero_nan_tail_parses_as_float_nan(self): + # Position 0 of an echoed prompt carries no preceding token to + # condition on, so the engine's tail there is always "nan 0" -- + # this must parse to float("nan"), not 0.0. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed(b"ECHO " + request_id.encode() + b" 2 0 nan 0\nHi\n") + kind, (pos, data, record) = events.get(timeout=1) + self.assertEqual((kind, pos, data), ("echo", 0, b"Hi")) + self.assertTrue(math.isnan(record["lp"])) + self.assertEqual(record["topk"], []) + engine.close() + + def test_nan_and_inf_logprob_tail_values_parse_cleanly(self): + # A degenerate logit row (all -inf after grammar masking, say) + # produces a well-formed record, not a parse failure. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed( + b"DATA " + request_id.encode() + b" 1 inf 2 5 inf 6 -inf\nx\n") + kind, (data, record) = events.get(timeout=1) + self.assertEqual((kind, data), ("data", b"x")) + self.assertEqual(record["lp"], float("inf")) + self.assertEqual(record["topk"], [(5, float("inf")), (6, float("-inf"))]) + engine.close() + + def test_high_precision_and_negative_infinity_tail_values_parse(self): + # The parser must not assume any particular float rendering -- a + # 17-significant-digit value (as a higher-precision engine build + # might emit, versus the shipped build's %.6f) and a bare "-inf" + # both parse through the same float() call as any other value. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed( + b"DATA " + request_id.encode() + + b" 1 -0.10536051565782628 2 3 -1.7987654321098765 8 -inf\nx\n") + kind, (data, record) = events.get(timeout=1) + self.assertEqual((kind, data), ("data", b"x")) + self.assertEqual(record["lp"], -0.10536051565782628) + self.assertEqual(record["topk"], [(3, -1.7987654321098765), (8, float("-inf"))]) + engine.close() + + def test_top_k_stays_in_wire_order_never_sorted(self): + # The table is unsorted on the wire; a pair earlier in the wire + # order but numerically/id-smaller later must stay first -- catches + # an accidental sort by id or by log-probability. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed( + b"DATA " + request_id.encode() + b" 1 -1.0 2 9 -1.0 2 -0.1\nx\n") + kind, (data, record) = events.get(timeout=1) + self.assertEqual((kind, data), ("data", b"x")) + self.assertEqual(record["topk"], [(9, -1.0), (2, -0.1)]) + engine.close() + + def test_short_logprob_tail_raises_rather_than_silently_truncating(self): + # A tail whose k claims more pairs than are actually present on the + # wire must fail loudly, never silently truncate to the pairs that + # happen to be there. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 1 -0.5 2 3 -0.5\nx\n") + # A DONE trails the malformed frame so a dispatcher that fails to + # catch the mismatch completes normally instead of hanging -- + # keeping this a clean assertion failure, not a stuck test. + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_long_logprob_tail_raises_rather_than_ignoring_trailing_fields(self): + # The reverse of the short-tail case: a tail with MORE fields than + # its own k accounts for must also fail loudly, not silently + # ignore the trailing garbage. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed( + b"DATA " + request_id + b" 1 -0.5 1 3 -0.5 9 -9.0\nx\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_non_numeric_tail_fields_raise_named_error(self): + # A non-numeric lp, token id, or per-candidate log-probability must + # be a named engine-protocol error, not an uncaught ValueError. + cases = ( + b"DATA {id} 1 nope 1 3 -0.5\nx\n", # lp + b"DATA {id} 1 -0.5 1 abc -0.5\nx\n", # tid + b"DATA {id} 1 -0.5 1 3 nope\nx\n", # tlp (not "nan"/"inf") + ) + for malformed in cases: + with self.subTest(frame=malformed): + def respond(process, frame, malformed=malformed): + request_id = frame.split()[1] + process.stdout.feed(malformed.replace(b"{id}", request_id)) + process.stdout.feed( + b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_negative_or_oversized_k_raises_named_error(self): + # k selects how many candidate pairs follow; a negative k or one + # past the engine's own top-32 cap is malformed, not a huge/empty + # read. + for k in (-1, LOGPROBS_TOP_K_CAP + 1): + with self.subTest(k=k): + def respond(process, frame, k=k): + request_id = frame.split()[1] + process.stdout.feed( + b"DATA " + request_id + f" 1 -0.5 {k}\n".encode() + b"x\n") + process.stdout.feed( + b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_negative_token_id_raises_named_error(self): + # A candidate's token id is a vocabulary index -- never negative. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed( + b"DATA " + request_id + b" 1 -0.5 1 -3 -0.5\nx\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_echo_frame_bad_terminator_is_a_named_error(self): + # The byte after an ECHO frame's payload must be the LF terminator; + # anything else is a named protocol error, matching DATA/TOOL. The + # stream is closed right after the bad terminator so a check that + # is missing or weakened surfaces "colibri engine exited + # unexpectedly" from the next (never-arriving) frame instead of + # hanging forever. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ECHO " + request_id + b" 2 0 nan 0\nHiX") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine ECHO terminator"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_data_size_bound_checked_before_any_payload_read(self): + # The 65536-byte size bound must be enforced BEFORE any payload + # byte is read. The fake stream here is closed right after a + # too-large claimed size and a couple of payload bytes -- a bound + # check that fires first raises immediately; one that is missing or + # weakened would instead try to read 70000 bytes from a stream that + # only ever offers 2 and then closes, surfacing "truncated engine + # DATA payload" (a different, misleading error) rather than + # hanging forever waiting for bytes that will never arrive. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 70000\nxy") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine DATA size"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_data_frame_truncated_payload_is_a_named_data_error(self): + # _read_exact's kind defaults to "DATA" for the legacy DATA/TOOL + # call sites -- that default is the only thing keeping this message + # byte-identical to what it said before GRPP/GRPG (and now ECHO) + # learned to name their own kind. A declared size larger than what + # the stream ever offers, followed by close(), drives _read_exact's + # chunk == b"" branch instead of the separate size-bound or + # terminator checks. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 10\nZZ") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "^truncated engine DATA payload$"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_echo_frame_errors_are_named_echo_not_data(self): + # A malformed ECHO frame's error message must say ECHO, not a + # copy-pasted DATA -- a wrong frame name in a dispatcher-killing + # error is actively misleading to whoever reads it. The stream is + # closed right after the oversized header so a missing/weakened + # bound check surfaces a truncation error instead of hanging. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ECHO " + request_id + b" 99999 0 nan 0\n") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine ECHO size"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_echo_frame_truncated_payload_names_the_echo_kind(self): + # ECHO's size and terminator checks already name ECHO explicitly + # (the test above); the payload read inside _read_exact still fell + # back to its "DATA" default, because the ECHO call sites never + # threaded kind through the way GRPP/GRPG's do. A declared size + # larger than what the stream ever offers, followed by close(), + # drives _read_exact's chunk == b"" branch instead of the separate + # size-bound or terminator checks. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ECHO " + request_id + b" 10 0 nan 0\nZZ") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "^truncated engine ECHO payload$"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_dispatcher_drops_echo_frames_with_no_pending_request(self): + # An ECHO frame for an id with no pending entry (already DONE, or + # never admitted) stays droppable, exactly like DATA/ACCEPT already + # do -- it must not raise or wedge the dispatcher for the NEXT + # request on the same connection. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DONE " + request_id + b" STAT 0 1 0 1 3 0\n") + # A stray ECHO for an id that is no longer pending (this one + # just finished) must be read and dropped, not raise. + process.stdout.feed(b"ECHO " + request_id + b" 1 0 nan 0\nx\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + # The dispatcher must still be alive and able to serve a second + # request -- proof the stray frame didn't wedge it. + engine.generate("hello2", 4, 0.0, 1.0, lambda _: None) + self.assertIsNone(engine.dispatcher_error) + engine.close() + + def test_supports_logprobs_echo_flag_set_once_by_arch(self): + # The capability flag is set once at launch from the arch id, + # glm only. + process_glm = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process_glm): + engine = Engine("glm", "model") + self.assertTrue(engine.supports_logprobs_echo) + engine.close() + + process_other = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process_other), \ + patch("openai_server.ARCH", "inkling"): + engine = Engine("inkling", "model") + self.assertFalse(engine.supports_logprobs_echo) + engine.close() + def test_unknown_frame_still_stops_dispatcher(self): # The catch-all that makes an unrecognized frame a hard failure is # load-bearing for the U7a compatibility asymmetry (a new engine's @@ -1093,6 +1952,35 @@ def generate(): self.assertEqual(outcome, ["cancelled"]) self.assertEqual(process.writes[-1].split(), [b"CANCEL", request_id]) + def test_generate_drops_its_pending_entry_when_the_cancel_write_fails(self): + # Every raise out of generate() after admission must still clear its + # own self.pending[request_id] slot. The DATA/ERROR dispatcher arms + # pop that slot themselves on the frames they own, but a raise from + # anywhere else in generate() -- here, a broken engine stdin on the + # CANCEL write an already-disconnected client triggers -- has no + # other code path clearing it. A leaked slot sits in self.pending + # until an unrelated dispatcher failure clears the whole map via + # _fail_pending, which can be arbitrarily far in the future. + request_id = None + + def respond(process, frame): + nonlocal request_id + fields = frame.split() + if fields[0] == b"SUBMIT": + request_id = fields[1] + elif fields[0] == b"CANCEL": + raise BrokenPipeError("synthetic engine stdin failure") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "failed to write CANCEL"): + engine.generate("hello", 8, 0.7, 0.9, lambda _: None, + cancelled=lambda: True) + with engine.pending_lock: + self.assertEqual(engine.pending, {}) + engine.close() + def test_stops_generation_through_successful_done_path(self): request_id = None @@ -1117,6 +2005,221 @@ def respond(process, frame): self.assertEqual(stats["completion_tokens"], 1) self.assertEqual(process.writes[-1].split(), [b"STOP", request_id]) + def test_group_frames_between_data_frames_drain_leaving_second_data_intact(self): + # A future engine may interleave the four group-scoring frame kinds + # with an ordinary request's own DATA frames on the same pipe. Each + # kind must be fully drained without disturbing frame sync, so the + # request's own second DATA frame still arrives byte-for-byte. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 1\nA\n") + process.stdout.feed(b"GRPP 99 2 0 nan 0\nZZ\n") + process.stdout.feed(b"GRPG 99 2 0 1 nan 0\nZZ\n") + process.stdout.feed(b"GRPS 99 0 3 2\n") + process.stdout.feed(b"GRPE 99 0 -1.5 2 1\n") + process.stdout.feed(b"DATA " + request_id + b" 1\nB\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 2 2.5 0 1.0 4 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + stats = engine.generate("hello", 8, 0.7, 0.9, chunks.append) + engine.close() + self.assertEqual(chunks, ["A", "B"]) + self.assertEqual(stats["completion_tokens"], 2) + + def test_group_payload_never_reaches_the_matching_requests_queue(self): + # Even when a group frame's own id field collides with an active + # request's id, group payload bytes must never land in that + # request's event queue -- this server has no group-response + # contract to hand them to. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"GRPP " + request_id + b" 2 0 nan 0\nZZ\n") + process.stdout.feed(b"DATA " + request_id + b" 1\nA\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 2.5 0 1.0 3 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + stats = engine.generate("hello", 4, 0.7, 0.9, chunks.append) + engine.close() + self.assertEqual(chunks, ["A"]) + self.assertEqual(stats["completion_tokens"], 1) + + def test_group_payload_size_bound_checked_before_any_read(self): + # The 65536-byte size bound applies to GRPP/GRPG exactly as it does + # to DATA, and must be enforced before any payload byte is read. + def respond(process, frame): + process.stdout.feed(b"GRPP 99 70000 0 nan 0\nxy") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine GRPP size"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_group_frame_missing_terminator_is_a_named_error(self): + # The byte after a GRPG frame's payload must be LF, same as DATA; + # the error names the frame kind that was actually malformed. + def respond(process, frame): + process.stdout.feed(b"GRPG 99 2 0 1 nan 0\nZZX") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine GRPG terminator"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_group_frame_truncated_payload_names_the_group_kind(self): + # A GRPP/GRPG truncation must name its own kind, not fall back to + # _read_exact's "DATA" default -- the counterpart to the legacy + # DATA/TOOL/ECHO case above, which pins that the default is still + # "DATA" for them. A declared size larger than what the stream ever + # offers, followed by close(), drives _read_exact's chunk == b"" + # branch instead of the separate size-bound or terminator checks. + def respond(process, frame): + process.stdout.feed(b"GRPP 99 10 0 nan 0\nZZ") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "^truncated engine GRPP payload$"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def _assert_wrong_field_count_falls_to_unknown_frame_handling(self, frame_bytes, + expected_regex): + # A group frame whose field count misses its kind's guard falls + # through to the same catch-all every other unrecognized frame hits + # (pinned by test_unknown_frame_still_stops_dispatcher above). The + # fake stdout is closed right after feeding the hostile frame: if a + # regression widens a guard so the frame is silently consumed + # instead of raising, the dispatcher's next readline() sees a + # closed, empty stream and fails with "colibri engine exited + # unexpectedly" rather than blocking forever on a frame that will + # never arrive. The generate() call itself also runs on a worker + # thread joined with a timeout, so even an unforeseen hang fails + # this test with a clear message instead of hanging the suite. + def respond(process, frame): + process.stdout.feed(frame_bytes) + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + outcome = {} + + def run(): + try: + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + except Exception as error: # noqa: BLE001 - captured for the main thread + outcome["error"] = error + + thread = threading.Thread(target=run) + thread.start() + thread.join(timeout=10) + still_running = thread.is_alive() + engine.close() + if still_running: + thread.join(timeout=2) + self.assertFalse(still_running, + "dispatcher did not fail within the timeout; the " + "wrong-field-count guard likely let the frame " + "through instead of falling to the catch-all") + self.assertIsInstance(outcome.get("error"), RuntimeError) + self.assertRegex(str(outcome["error"]), expected_regex) + + def test_group_frame_with_wrong_field_count_falls_to_unknown_frame_handling(self): + # A GRPS frame with a short field count doesn't match the group + # branch's exact-5 guard. + self._assert_wrong_field_count_falls_to_unknown_frame_handling( + b"GRPS 99 0 3\n", "invalid engine response: GRPS") + + def test_group_frame_grpp_wrong_field_count_falls_to_unknown_frame_handling(self): + # A GRPP frame with 5 fields is one short of the payload kind's + # minimum-6 guard. + self._assert_wrong_field_count_falls_to_unknown_frame_handling( + b"GRPP 99 2 0 nan\n", "invalid engine response: GRPP 99 2 0 nan") + + def test_group_frame_grpg_wrong_field_count_falls_to_unknown_frame_handling(self): + # A GRPG frame with 6 fields is one short of the payload kind's + # minimum-7 guard. + self._assert_wrong_field_count_falls_to_unknown_frame_handling( + b"GRPG 99 2 0 1 nan\n", "invalid engine response: GRPG 99 2 0 1 nan") + + def test_group_frame_grps_wrong_field_count_falls_to_unknown_frame_handling(self): + # A GRPS frame with 6 fields is one over the header-only kind's + # exact-5 guard. + self._assert_wrong_field_count_falls_to_unknown_frame_handling( + b"GRPS 99 0 3 2 1\n", "invalid engine response: GRPS 99 0 3 2 1") + + def test_group_frame_grpe_wrong_field_count_falls_to_unknown_frame_handling(self): + # A GRPE frame with 5 fields is one short of the header-only kind's + # exact-6 guard. + self._assert_wrong_field_count_falls_to_unknown_frame_handling( + b"GRPE 99 0 -1.5 2\n", "invalid engine response: GRPE 99 0 -1.5 2") + + +class SeedWireFrameTest(unittest.TestCase): + """C1: `seed` is accepted and ignored. A stub-response equality check alone + is vacuous here (FakeEngine always returns the same canned text regardless + of any request field) -- the real proof is that the byte-exact SUBMIT frame + the dispatcher writes to the engine process (see DispatcherTest above) never + carries the seed value at all, seeded or not. + """ + + def _completion(self, body): + frames = [] + + def respond(process, frame): + frames.append(frame) + rid = frame.split()[1] + process.stdout.feed(b"DATA " + rid + b" 5\nHello\n") + process.stdout.feed(b"DONE " + rid + b" STAT 1 2.5 0 1.0 4 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + server = APIServer(("127.0.0.1", 0), engine, "test-model", "secret", 16) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + data = json.dumps(body).encode() + headers = {"Authorization": "Bearer secret", "Content-Type": "application/json"} + request = Request(f"http://127.0.0.1:{server.server_port}/v1/completions", + data=data, headers=headers) + with urlopen(request, timeout=2) as response: + status = response.status + parsed = json.load(response) + finally: + server.scheduler.close() + server.shutdown() + server.server_close() + thread.join(timeout=2) + engine.close() + return status, parsed, frames[0] + + def test_seed_accepted_and_absent_from_submit_frame(self): + base = {"model": "test-model", "prompt": "Complete me", "temperature": 0, "max_tokens": 4} + status_plain, body_plain, frame_plain = self._completion(base) + status_seeded, body_seeded, frame_seeded = self._completion({**base, "seed": 1234}) + self.assertEqual(status_plain, 200) + self.assertEqual(status_seeded, 200) + self.assertEqual(body_seeded["choices"][0], body_plain["choices"][0]) + # Each call uses a freshly-constructed Engine, so both first requests are + # assigned request id "1" -- the wire frames are directly byte-comparable, + # no field needs normalizing. If `seed` ever leaked onto the SUBMIT + # header or into an extension field, this equality would break. + self.assertEqual(frame_seeded, frame_plain) + class CapSentinelShimTest(unittest.TestCase): # #379 cap-sentinel shim, arch-keyed (#386 r2, F3): an absent cap is @@ -1362,6 +2465,24 @@ def test_chat_completion(self): self.assertIn("<|user|>Hi<|assistant|>", self.engine.calls[-1][0]) self.assertEqual(self.engine.calls[-1][4], 1) + def test_group_score_is_not_read_by_chat_or_anthropic_surfaces(self): + # The group_score guard lives only in completion(); the chat and + # Anthropic surfaces never call it, so the opt-in is ignored, not + # refused, on those two surfaces. Pinned here so the scope choice + # is visible rather than an untested accident. + with self.request("/v1/chat/completions", { + "model": "test-model", "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 4, "group_score": True, + }) as response: + body = json.load(response) + self.assertEqual(body["object"], "chat.completion") + with self.request("/v1/messages", { + "model": "test-model", "max_tokens": 4, "group_score": True, + "messages": [{"role": "user", "content": "Hi"}], + }) as response: + body = json.load(response) + self.assertEqual(body["type"], "message") + def test_kimi_chat_completion_uses_multiturn_wire_payload(self): with patch("openai_server.ARCH", "kimi"): with self.request("/v1/chat/completions", { @@ -2222,12 +3343,112 @@ def respond(process, frame): engine.generate("hi", 8, 0.7, 0.9, lambda _: None, on_accept=lambda _: None) engine.close() + def _streaming_commit_probe(self): + """Real Engine + FakeProcess + a real APIServer/socket, wired so a CANCEL frame + written to the (fake) engine is observable deterministically (no sleep-and-hope): + `cancel_seen` fires the instant the CANCEL write happens, synchronously in the + request-handling thread; `pending_empty()` bounded-polls for the dispatcher + thread's async pop of the request out of `engine.pending`, which is genuinely + asynchronous relative to the CANCEL write.""" + cancel_seen = threading.Event() -class _ContextExceededEngine(FakeEngine): - """Engine that rejects the prompt before ACCEPT — on_accept is never called.""" - def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None, on_accept=None): - raise APIError(400, "This model's maximum context length is 4094 tokens.", + def respond(process, frame): + parts = frame.split() + if parts[0] == b"SUBMIT": + process.stdout.feed(b"ACCEPT " + parts[1] + b" 7\n") + elif parts[0] == b"CANCEL": + process.stdout.feed(b"ERROR " + parts[1] + b" CANCELLED\n") + cancel_seen.set() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + server = APIServer(("127.0.0.1", 0), engine, "test-model") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join, 2) + self.addCleanup(server.server_close) + self.addCleanup(server.shutdown) + self.addCleanup(server.scheduler.close) + self.addCleanup(engine.close) + + def pending_empty(timeout=3.0): + deadline = time.time() + timeout + while time.time() < deadline: + if "1" not in engine.pending: + return True + time.sleep(0.01) + return "1" not in engine.pending + + return process, engine, server, cancel_seen, pending_empty + + def _post_streaming_request(self, server): + payload = json.dumps({"model": "test-model", "stream": True, "max_tokens": 16, + "messages": [{"role": "user", "content": "Hi"}]}) + request = (f"POST /v1/messages HTTP/1.1\r\nHost: 127.0.0.1\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(payload)}\r\n\r\n{payload}").encode() + sock = socket.create_connection(("127.0.0.1", server.server_port), timeout=3) + self.addCleanup(sock.close) + sock.sendall(request) + return sock + + def test_anthropic_end_headers_failure_at_commit_does_not_orphan_the_pending_request(self): + """If the client vanishes at the exact moment the engine ACCEPTs, the real header + flush -- end_headers(), where BaseHTTPRequestHandler actually hits the socket -- + can raise. That must not unwind out of generate()'s dispatch loop with nothing + sent: the request would sit in the engine's pending map forever, with no CANCEL + ever going out to free it. Patches end_headers() itself (not send_response), so + send_response's real _committed/close_connection bookkeeping still runs first, + exactly as it would for a real dropped socket.""" + process, engine, server, cancel_seen, pending_empty = self._streaming_commit_probe() + with patch.object(APIHandler, "end_headers", + side_effect=BrokenPipeError("client vanished exactly at ACCEPT")): + self._post_streaming_request(server) + self.assertTrue(cancel_seen.wait(timeout=3), + "the commit failure never triggered a CANCEL") + self.assertTrue(any(w.startswith(b"CANCEL ") for w in process.writes), + "the commit failure never triggered a CANCEL") + self.assertTrue(pending_empty(), + "the request was never removed from the engine's pending map") + + def test_anthropic_first_sse_write_failure_after_commit_does_not_orphan_the_pending_request(self): + """Companion to the end_headers() variant above: a write failure on the FIRST SSE + body write (message_start, right after headers commit cleanly) is already caught + by send_event()'s own try/except -- this is REGRESSION-COVERAGE that pathway keeps + working, not a new defect. Patches the handler's wfile so only that first write + after headers raises; the header commit itself goes through for real.""" + process, engine, server, cancel_seen, pending_empty = self._streaming_commit_probe() + real_end_headers = APIHandler.end_headers + state = {"headers_done": False} + + def committing_end_headers(handler): + real_end_headers(handler) + state["headers_done"] = True + real_write = handler.wfile.write + + def failing_write(data): + if state["headers_done"]: + state["headers_done"] = False # only the first post-header write fails + raise BrokenPipeError("client vanished on the first SSE write") + return real_write(data) + handler.wfile.write = failing_write + + with patch.object(APIHandler, "end_headers", committing_end_headers): + self._post_streaming_request(server) + self.assertTrue(cancel_seen.wait(timeout=3), + "the first-SSE-write failure never triggered a CANCEL") + self.assertTrue(any(w.startswith(b"CANCEL ") for w in process.writes), + "the first-SSE-write failure never triggered a CANCEL") + self.assertTrue(pending_empty(), + "the request was never removed from the engine's pending map") + + +class _ContextExceededEngine(FakeEngine): + """Engine that rejects the prompt before ACCEPT — on_accept is never called.""" + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None, on_accept=None): + raise APIError(400, "This model's maximum context length is 4094 tokens.", "messages", "context_length_exceeded") @@ -2260,6 +3481,194 @@ def test_streaming_context_exceeded_is_clean_400(self): self.assertEqual(body["error"]["code"], "context_length_exceeded") self.assertEqual(body["error"]["param"], "messages") + def test_anthropic_streaming_context_exceeded_is_clean_400(self): + # The Anthropic streaming path must defer its 200 the same way the OpenAI path + # above does: a refusal discovered before the engine accepts the prompt is a + # real HTTP 400 in the Anthropic error envelope, never a committed 200 followed + # by a truncated (or empty) SSE body. + req = Request(self.base + "/v1/messages", + data=json.dumps({"model": "test-model", "stream": True, + "max_tokens": 16, + "messages": [{"role": "user", "content": "x" * 100}]}).encode(), + headers={"Content-Type": "application/json"}) + with self.assertRaises(HTTPError) as caught: + urlopen(req, timeout=3) + self.addCleanup(caught.exception.close) + self.assertEqual(caught.exception.code, 400) # a real 400, not a 200 stream + raw = caught.exception.read() + self.assertNotIn(b"event:", raw, "a pre-accept refusal must not emit any SSE bytes") + body = json.loads(raw) + self.assertEqual(body["type"], "error") # the Anthropic envelope + self.assertEqual(body["error"]["type"], "invalid_request_error") + self.assertIn("maximum context length", body["error"]["message"]) + + +class AnthropicStreamCommitTest(unittest.TestCase): + """On a healthy engine the deferred commit is invisible -- the stream's byte-level + framing (headers, message_start through message_stop order) is unchanged, and the + 200 is provably committed from the engine's ACCEPT, not before generate() is even + called.""" + + def setUp(self): + self.engine = FakeEngine() + self.server = APIServer(("127.0.0.1", 0), self.engine, "test-model") + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + + def tearDown(self): + self.server.scheduler.close() + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + def _raw(self, body): + payload = json.dumps(body) + request = (f"POST /v1/messages HTTP/1.1\r\nHost: 127.0.0.1\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(payload)}\r\nConnection: close\r\n\r\n" + f"{payload}").encode() + with socket.create_connection(("127.0.0.1", self.server.server_port), + timeout=3) as sock: + sock.sendall(request) + chunks = [] + while True: + chunk = sock.recv(65536) + if not chunk: + break + chunks.append(chunk) + return b"".join(chunks).decode("utf-8", "replace") + + def test_success_stream_framing_is_unchanged(self): + raw = self._raw({"model": "test-model", "stream": True, "max_tokens": 16, + "messages": [{"role": "user", "content": "Hi"}]}) + head, body = raw.split("\r\n\r\n", 1) + self.assertIn("HTTP/1.1 200", head) + self.assertIn("content-type: text/event-stream", head.lower()) + self.assertIn("connection: close", head.lower()) + # Exact sequence, not mere presence: two chunks ("Hé", "llo") from FakeEngine's + # default script must produce two content_block_delta events in position, not + # just an unordered set of the six event names. + names = [line[len("event: "):] for line in body.splitlines() + if line.startswith("event: ")] + self.assertEqual(names, ["message_start", "content_block_start", "content_block_delta", + "content_block_delta", "content_block_stop", "message_delta", + "message_stop"]) + deltas = "".join(json.loads(line[len("data: "):])["delta"]["text"] + for line in body.splitlines() + if line.startswith("data: ") and '"text_delta"' in line) + self.assertEqual(deltas, "Héllo") + + def test_stream_commits_only_after_accept(self): + # The mechanism itself: the 200 must be sent from the engine's on_accept + # callback, not unconditionally before generate() is even called. + committed_at = {} + commit_flags = [] + + def handler_committed(): + return bool(commit_flags) + + class CommitProbeEngine(FakeEngine): + def generate(self, prompt, maximum, temperature, top_p, on_text, + cache_slot=0, cancelled=None, grammar=None, stopped=None, + on_accept=None): + committed_at["before_accept"] = handler_committed() + if on_accept is not None: + on_accept({"prompt_tokens": 7}) + committed_at["after_accept"] = handler_committed() + on_text("ok") + return {"prompt_tokens": 7, "completion_tokens": 1, "length_limited": False} + + original = APIHandler.send_response + + def recording_send_response(handler, code, message=None): + commit_flags.append(code) + return original(handler, code, message) + + self.server.engine = CommitProbeEngine() + with patch.object(APIHandler, "send_response", recording_send_response): + raw = self._raw({"model": "test-model", "stream": True, "max_tokens": 16, + "messages": [{"role": "user", "content": "Hi"}]}) + self.assertIn("HTTP/1.1 200", raw) + self.assertFalse(committed_at["before_accept"], + "the Anthropic stream committed its 200 before engine ACCEPT") + self.assertTrue(committed_at["after_accept"]) + + +class AnthropicColdPrefillTest(unittest.TestCase): + """docs/api.md's Anthropic-streaming section: "Until the engine accepts, no + bytes are sent at all -- a request queued behind another generation waits + silently, exactly as the OpenAI-style streaming path already does." Against + an engine binary old enough never to send ACCEPT, the first accept-equivalent + event is the engine's first DATA or DONE instead, so a cold multi-minute + prefill sends zero bytes where `dev` sent `message_start` plus a periodic + `ping` for the same window -- traced through the code, never exercised end + to end against a real socket. + + `BlockingEngine` (used elsewhere for scheduler-queueing tests) stands in + for exactly that: `generate()` blocks -- simulating the prefill window -- + before it ever calls `on_accept`, the same callback boundary a real + Engine collapses ACCEPT and "first DATA/DONE from an old engine" onto. + This test reads the raw socket while the engine is still blocked and + proves nothing at all has arrived, then releases it and proves the + deferred 200 and SSE stream still show up once the engine finally + accepts. + + Note: unlike the opt-in per-token-logprobs channel (`COLI_LOGPROBS_ + ACCEPT_TIMEOUT`, a named 503 on timeout), the Anthropic endpoint never + passes `logprobs`/`tok_ids` to `Engine.generate()` (see + `anthropic_generation()`), so `accept_deadline` is `None` for this path + and the wait this test pins is genuinely unbounded on a real engine that + never answers at all -- there is no timeout or 503 documented or + observed for this case, only the silent wait docs/api.md describes.""" + + def setUp(self): + self.engine = BlockingEngine() + self.server = APIServer(("127.0.0.1", 0), self.engine, "test-model") + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + + def tearDown(self): + self.engine.release.set() + self.server.scheduler.close() + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + def test_no_bytes_reach_the_client_until_the_engine_finally_accepts(self): + payload = json.dumps({"model": "test-model", "stream": True, "max_tokens": 16, + "messages": [{"role": "user", "content": "Hi"}]}) + request = (f"POST /v1/messages HTTP/1.1\r\nHost: 127.0.0.1\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(payload)}\r\n\r\n{payload}").encode() + sock = socket.create_connection(("127.0.0.1", self.server.server_port), timeout=3) + self.addCleanup(sock.close) + sock.sendall(request) + + self.assertTrue(self.engine.entered.wait(2), + "the request never reached the engine (admitted elsewhere?)") + sock.settimeout(0.5) + with self.assertRaises(socket.timeout): + sock.recv(4096) + # Confirmed: while the (simulated) cold prefill is in progress the client + # sees literally nothing -- no status line, no headers, no SSE bytes -- + # unlike `dev`, which sent `message_start` plus periodic pings here. + + self.engine.release.set() # the engine "accepts" now (ACCEPT, or an old + # engine's first DATA/DONE -- indistinguishable + # from here on) + sock.settimeout(3) + chunks = [] + while True: + chunk = sock.recv(65536) + if not chunk: + break + chunks.append(chunk) + raw = b"".join(chunks).decode("utf-8", "replace") + head, body = raw.split("\r\n\r\n", 1) + self.assertIn("HTTP/1.1 200", head) + self.assertIn("event: message_start", body) + self.assertIn("event: message_stop", body) + class _ExplodingEngine(FakeEngine): """ACCEPTs the prompt (committing the streaming 200), then dies mid-generation.""" @@ -2430,6 +3839,59 @@ def test_engine_failure_after_commit_does_not_splice_a_second_response(self): self.assertIn("partial", raw) # the events sent before the failure survive self.assertNotIn("", raw) + def test_write_failure_reaching_the_committed_stream_ends_it_cleanly(self): + """docs/api.md, "Engine protocol contract: checked writes and SIGPIPE": "for a + request whose response is already committed as a stream, a failed write ends + the stream instead of producing a 500." + + `test_engine_failure_after_commit_does_not_splice_a_second_response` above + pins the same `_fail()`/`_committed` branch with a generic engine + RuntimeError; `test_generate_drops_its_pending_entry_when_the_cancel_write_fails` + pins the checked STOP/CANCEL write itself, but only at the `Engine.generate()` + level -- no HTTP handler, no socket. Neither proves what a real client sees + when that specific checked-write failure reaches an already-committed HTTP + stream. This test drives a real Engine + a fake engine subprocess whose stdin + raises on the STOP write (the same injection those tests use) through a real + streaming HTTP request, and reads the raw socket.""" + request_id = None + + def respond(process, frame): + nonlocal request_id + parts = frame.split() + if parts[0] == b"SUBMIT": + request_id = parts[1] + process.stdout.feed(b"ACCEPT " + request_id + b" 7\n") + process.stdout.feed(b"DATA " + request_id + b" 11\nhello STOP!\n") + elif parts[0] == b"STOP": + raise BrokenPipeError("synthetic engine stdin failure") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + server = self._server(engine) + + log = io.StringIO() + with patch("sys.stderr", log): + raw = self._raw(server, self._request_bytes( + dict(self.CHAT, stream=True, stop=["STOP!"]))) + + self.assertEqual(raw.count("HTTP/1."), 1, + "a failed write must not splice a second status line into " + "the committed stream") + head, body = raw.split("\r\n\r\n", 1) + self.assertIn("HTTP/1.1 200", head) + self.assertIn("hello ", body) # the text sent before the failed write survives + self.assertNotIn("STOP!", body) # the matched stop sequence itself stays withheld + self.assertNotIn("data: [DONE]", body, + "the stream must end at the failure, not run to a normal finish") + self.assertNotIn('"type": "error"', body, + "no error body may be spliced into an already-committed stream") + self.assertNotIn("", raw) # the connection actually closed, not hung + self.assertIn("failed to write STOP to the engine", log.getvalue(), + "the write failure must be logged (do_POST's `except Exception` -> " + "log_error), not silently swallowed") + def test_non_streaming_response_still_reuses_the_connection(self): """The fix must not turn every response into a close: plain JSON stays persistent.""" server = self._server() @@ -2693,5 +4155,2525 @@ def test_missing_fields_do_not_crash_the_message(self): self.assertIn("the context", text) +class NonGlmEngine(FakeEngine): + """The per-token logprobs / token-id-prompt capability gate's negative + case: an engine that predates the U7a extension (or isn't glm) never + gets the extension fields, never mind what it would do with them.""" + supports_logprobs_echo = False + + +class LogprobsHTTPTest(unittest.TestCase): + """End-to-end acceptance tests for the per-token logprobs and + token-id-prompt capability gate, against a real APIServer + APIHandler, + with FakeEngine standing in for the engine subprocess (its + logprobs_channel() returns the canned U7a records documented on + FakeEngine.generate() above). + + On the predecessor head, ANY truthy `logprobs` (an integer, or `True`) + unconditionally 400s before this server's own validation ever runs, and + `choices[].logprobs` is otherwise always null. That predecessor check + happens to also 400 several of this class's "reject this" cases for the + WRONG reason (a blanket rejection, not this server's specific named + validation), so only 6 of these 12 methods are regression pins (fail + outright on the predecessor head): test_chat_logprobs_content_shape, + test_chat_echo_is_rejected, test_completions_logprobs_true_is_named_400 + (predecessor 400s but with the wrong error code), + test_bit_identity_end_to_end, test_echo_reconstructs_prompt_text, and + test_nan_logprob_serializes_as_json_null_over_the_wire. The other 6 + already pass on the predecessor head (a truthy `logprobs`/`echo` request + there either already 400s for the coincidentally-matching blanket reason, + or a falsy/absent one already no-ops to `logprobs: null`) and are + REGRESSION-COVERAGE, pinning that this server's own validation reaches the + identical observable result: test_completions_logprobs_zero_is_no_logprobs_end_to_end, + test_break_it_logprobs_out_of_range, test_break_it_echo_without_logprobs_is_a_documented_noop, + test_break_it_streaming_plus_logprobs_is_named_400, + test_break_it_logprobs_rejected_for_non_glm_engine, and + test_golden_fixture_style_plain_request_is_unaffected. [RAN] verified by + running this class against the predecessor head's server file: 6 + failed, 6 passed. (That 6/6/12 count covers only the original + per-token-logprobs methods above; the array/token-id-prompt methods + added later in this class are a separate capability and are not + included in it.) + """ + + @classmethod + def setUpClass(cls): + cls.engine = FakeEngine() + cls.server = APIServer(("127.0.0.1", 0), cls.engine, "test-model", "secret", 16, + kv_slots=2) + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + cls.base = f"http://127.0.0.1:{cls.server.server_port}" + + @classmethod + def tearDownClass(cls): + cls.server.scheduler.close() + cls.server.shutdown() + cls.server.server_close() + cls.thread.join(timeout=2) + + def request(self, path, body=None): + headers = {"Authorization": "Bearer secret"} + data = None + if body is not None: + data = json.dumps(body).encode() + headers["Content-Type"] = "application/json" + return urlopen(Request(self.base + path, data=data, headers=headers), timeout=2) + + def _temp_server(self, engine): + """A second, throwaway server backed by a different fake engine -- + for the capability-gate negative cases, which must not share + cls.engine/cls.server with the rest of this class.""" + server = APIServer(("127.0.0.1", 0), engine, "test-model", "secret", 16, kv_slots=1) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.scheduler.close) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + self.addCleanup(thread.join, timeout=2) + return f"http://127.0.0.1:{server.server_port}" + + # ---- chat logprobs shape ------------------------------------------------ + + def test_chat_logprobs_content_shape(self): + with self.request("/v1/chat/completions", { + "model": "test-model", "messages": [{"role": "user", "content": "Hi"}], + "logprobs": True, "top_logprobs": 2, + }) as response: + body = json.load(response) + content = body["choices"][0]["logprobs"]["content"] + self.assertGreater(len(content), 0) + for entry in content: + self.assertEqual(set(entry), {"token", "logprob", "bytes", "top_logprobs"}) + for alt in entry["top_logprobs"]: + self.assertEqual(set(alt), {"token", "logprob", "bytes"}) + self.assertNotIn("echo", body["choices"][0]) + + def test_chat_echo_is_rejected(self): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/chat/completions", { + "model": "test-model", "messages": [{"role": "user", "content": "Hi"}], + "echo": True, + }) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "echo") + + # ---- end to end ----------------------------------------------------------- + + def test_completions_logprobs_zero_is_no_logprobs_end_to_end(self): + # `logprobs: 0` behaves exactly like an omitted field -- the engine + # channel stays off and the choice carries logprobs: null. + with self.request("/v1/completions", { + "model": "test-model", "prompt": "hi", "logprobs": 0, + }) as response: + body = json.load(response) + self.assertIsNone(body["choices"][0]["logprobs"]) + self.assertEqual(self.engine.last_logprobs, 0) + + def test_completions_logprobs_true_is_named_400(self): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": "hi", "logprobs": True, + }) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "logprobs") + self.assertEqual(error["code"], "invalid_value") + + # ---- bit-identity, end to end ------------------------------------------- + + def test_bit_identity_end_to_end(self): + with self.request("/v1/completions", { + "model": "test-model", "prompt": "Hé", "echo": True, "logprobs": 2, + "max_tokens": 1, + }) as response: + body = json.load(response) + logprobs = body["choices"][0]["logprobs"] + for i, tok in enumerate(logprobs["tokens"]): + lp = logprobs["token_logprobs"][i] + table = logprobs["top_logprobs"][i] + if lp is not None and tok in table: + self.assertEqual(table[tok], lp, + f"position {i}: token_logprobs != top_logprobs[tokens[i]]") + + # ---- text_offset / echoed-prompt reconstruction ------------------------- + + def test_echo_reconstructs_prompt_text(self): + with self.request("/v1/completions", { + "model": "test-model", "prompt": "Hé", "echo": True, "logprobs": 1, + "max_tokens": 1, + }) as response: + body = json.load(response) + logprobs = body["choices"][0]["logprobs"] + # The canned engine always echoes exactly "H", "é" for the prompt + # positions (FakeEngine.logprobs_channel) -- the first two tokens + # reconstruct the prompt exactly; anything after that is generated. + self.assertEqual("".join(logprobs["tokens"][:2]), "Hé") + offsets = logprobs["text_offset"] + self.assertEqual(offsets, sorted(offsets)) + self.assertEqual(offsets[0], 0) + self.assertIsNone(logprobs["token_logprobs"][0]) # first prompt token: null by convention + + # ---- array / token-id prompt intake -------------------------------------- + + # The LITERAL request body a genuine unmodified lm-eval run sent over + # the wire against a real bridge server, captured by a passive logging + # proxy (not hand-constructed). + LMEVAL_FIXTURE = Path(__file__).parent / "fixtures" / "captured_lmeval_request.json" + LMEVAL_FIXTURE_SHA256 = \ + "8242860518586177bba0dbe4d85a41a183c60b32ae56953be9e2e1e09251fe69" + + def test_lm_eval_fixture_bytes_match_the_bound_blob(self): + import hashlib + # The committed blob uses LF line endings. A checkout with line-ending + # conversion (git's autocrlf on Windows) hands the test CRLF bytes for the + # same blob, so normalise CRLF back to LF before comparing to the bound + # digest; the digest and length below are those of the committed bytes. + raw = self.LMEVAL_FIXTURE.read_bytes().replace(b"\r\n", b"\n") + digest = hashlib.sha256(raw).hexdigest() + self.assertEqual(digest, self.LMEVAL_FIXTURE_SHA256) + self.assertEqual(len(raw), 2732) + + def test_lm_eval_captured_request_replay(self): + # The captured request itself (a single nested token-id member -- + # lm-eval's tokenized loglikelihood shape wraps even one prompt in + # an outer batch list) replayed byte-for-byte against a real + # server, past both array intake and the batch dispatch path it + # would take if lm-eval ever grew to N>1: still just the + # single-prompt unwrap at N=1, producing a normal response. + captured = json.loads(self.LMEVAL_FIXTURE.read_text())["body_json"] + captured["model"] = "test-model" + with self.request("/v1/completions", captured) as response: + replayed = json.load(response) + self.assertEqual(len(replayed["choices"]), 1) + self.assertIsNotNone(replayed["choices"][0]["logprobs"]) + self.assertIn("prompt_tokens", replayed["usage"]) + self.assertEqual(self.engine.last_tok_ids, True) + + def test_nested_batch_of_one_prompt_is_identical_to_flat(self): + # An unmodified lm-eval-style client's tokenized loglikelihood path + # always wraps its token-id array in an outer batch-of-one list, + # even at batch size 1 -- it must unwrap to the identical flat + # behavior. + ids = [72, 233, 108] + with self.request("/v1/completions", { + "model": "test-model", "prompt": ids, "max_tokens": 1, + }) as response: + flat = json.load(response) + self.assertEqual(self.engine.calls[-1][0], "72 233 108") + with self.request("/v1/completions", { + "model": "test-model", "prompt": [ids], "max_tokens": 1, + }) as response: + nested = json.load(response) + self.assertEqual(self.engine.calls[-1][0], "72 233 108") + self.assertEqual(flat["choices"], nested["choices"]) + self.assertEqual(self.engine.last_tok_ids, True) + + def test_nested_single_prompt_with_non_int_element_is_a_named_400(self): + # After unwrapping the length-1 outer list, this must hit the SAME + # existing validation _encode_token_id_prompt already does for a + # flat list with a bad element -- no new error path. + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": [[1, "bad", 3]], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt") + + def test_array_prompt_over_the_batch_cap_is_a_named_400(self): + # Pinned against the documented cap value itself (not just derived + # from PROMPT_BATCH_CAP), so a mutation that raises the constant + # cannot silently widen this test's own boundary along with it. + self.assertEqual(PROMPT_BATCH_CAP, 128) + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": ["hi"] * 129, + "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt") + # A distinct code from the "batch dispatch not yet implemented" + # refusal (see test_array_prompt_of_more_than_one_member_is_a_named_400 + # below) -- this asserts the CAP is what fired, not that other, + # separate reason (both are named 400/param=prompt, so param alone + # cannot tell them apart). + self.assertEqual(error["code"], "prompt_batch_cap_exceeded") + + def test_array_prompt_at_exactly_the_batch_cap_is_not_a_cap_refusal(self): + # An off-by-one guard on the cap boundary itself: exactly + # PROMPT_BATCH_CAP (128) members must NOT trip the cap check (the + # one-over case above pins that 129 does) -- the batch dispatches + # and is admitted, one choice per member. + self.assertEqual(PROMPT_BATCH_CAP, 128) + with self.request("/v1/completions", { + "model": "test-model", "prompt": ["hi"] * 128, + "max_tokens": 1, + }) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 128) + + def test_array_prompt_over_the_token_budget_is_a_named_400(self): + # Pinned against the documented budget value itself (not just + # derived from PROMPT_BATCH_TOKEN_BUDGET), so a mutation that + # widens the constant cannot silently widen this test's own + # boundary along with it. + self.assertEqual(PROMPT_BATCH_TOKEN_BUDGET, 65536) + big = "x" * 32769 + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": [big, big], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt") + self.assertEqual(error["code"], "prompt_batch_token_budget_exceeded") + self.assertIn("budget", error["message"]) + + def test_array_prompt_token_budget_is_counted_in_utf8_bytes_not_characters(self): + # For string batches (no tokenizer available) the docs and the + # code both say the budget is applied to total UTF-8 BYTES, an + # upper bound on tokens -- not characters. "e-acute" is one + # character but two UTF-8 bytes, so this batch's byte total and + # character total straddle the budget on opposite sides: over by + # bytes, comfortably under by characters. If the accounting were + # ever swapped to count characters, this request would wrongly be + # ADMITTED (200, not 400). + self.assertEqual(PROMPT_BATCH_TOKEN_BUDGET, 65536) + member = "\u00e9" * 16385 # 16,385 chars, 32,770 UTF-8 bytes + self.assertEqual(len(member.encode("utf-8")), 32770) + total_bytes = 2 * len(member.encode("utf-8")) + total_chars = 2 * len(member) + self.assertGreater(total_bytes, PROMPT_BATCH_TOKEN_BUDGET) + self.assertLessEqual(total_chars, PROMPT_BATCH_TOKEN_BUDGET) + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": [member, member], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt") + self.assertEqual(error["code"], "prompt_batch_token_budget_exceeded") + self.assertIn(str(total_bytes), error["message"]) + + def test_array_prompt_token_budget_counts_actual_tokens_for_token_id_batches(self): + # docs/api.md (isolated-batch-limits section): "token-id batches + # count actual tokens, string batches count total UTF-8 bytes as + # an upper bound on tokens (no tokenizer is available + # server-side)." Every existing budget test (this class, above) + # uses STRING members, so only the byte-counting half of that + # sentence was ever exercised; the token-id half + # (`c/openai_server.py`'s `_completion_prompt_array`, `sum(len(member) + # for member in members)` when `tok_ids`) was unexercised. A member + # with 32769 ints has 32769 UTF-8-encoded-repr bytes that are + # irrelevant here -- if the aggregate check ever counted bytes (or + # anything else) instead of list length for a token-id batch, this + # request would be silently ADMITTED instead of refused. Two such + # members = 65538 actual tokens, one over PROMPT_BATCH_TOKEN_BUDGET + # (65536). + self.assertEqual(PROMPT_BATCH_TOKEN_BUDGET, 65536) + member = list(range(32769)) + before = len(self.engine.calls) + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": [member, member], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt") + self.assertEqual(error["code"], "prompt_batch_token_budget_exceeded") + self.assertIn("65538", error["message"]) + self.assertIn(str(PROMPT_BATCH_TOKEN_BUDGET), error["message"]) + self.assertIn("tokens", error["message"]) + # The defining property, matching the string-prompt budget test's + # shape: no engine work started on a refused batch. + self.assertEqual(len(self.engine.calls), before) + + def test_array_prompt_of_more_than_one_member_dispatches(self): + # A real (N>1) batch that passes shape/cap/budget validation + # dispatches -- one choice per member, in order. + with self.request("/v1/completions", { + "model": "test-model", "prompt": ["hi", "there"], "max_tokens": 1, + }) as response: + body = json.load(response) + self.assertEqual([choice["index"] for choice in body["choices"]], [0, 1]) + + def test_array_prompt_mixed_shapes_is_a_named_400(self): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": ["hi", [1, 2]], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt") + + def test_empty_array_prompt_is_a_named_400(self): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": [], "max_tokens": 1, + }) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt") + + # ---- break-it battery --------------------------------------------------- + + def test_break_it_logprobs_out_of_range(self): + for bad in (-1, 33, 1.5): + with self.subTest(bad=bad): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": "hi", "logprobs": bad, + }) + self.assertEqual(caught.exception.code, 400) + + def test_break_it_echo_without_logprobs_is_a_documented_noop(self): + with self.request("/v1/completions", { + "model": "test-model", "prompt": "hi", "echo": True, + }) as response: + body = json.load(response) + self.assertEqual(response.status, 200) + self.assertIsNone(body["choices"][0]["logprobs"]) + + def test_break_it_streaming_plus_logprobs_is_named_400(self): + with self.assertRaises(HTTPError) as caught: + self.request("/v1/completions", { + "model": "test-model", "prompt": "hi", "logprobs": 1, "stream": True, + }) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "logprobs") + + def test_break_it_logprobs_rejected_for_non_glm_engine(self): + base = self._temp_server(NonGlmEngine()) + request = Request(base + "/v1/chat/completions", method="POST", + headers={"Authorization": "Bearer secret", + "Content-Type": "application/json"}, + data=json.dumps({"model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "logprobs": True}).encode()) + with self.assertRaises(HTTPError) as caught: + urlopen(request, timeout=2) + self.assertEqual(caught.exception.code, 400) + + def test_break_it_array_prompt_rejected_for_non_glm_engine(self): + base = self._temp_server(NonGlmEngine()) + request = Request(base + "/v1/completions", method="POST", + headers={"Authorization": "Bearer secret", + "Content-Type": "application/json"}, + data=json.dumps({"model": "test-model", "prompt": [1, 2, 3], + "max_tokens": 1}).encode()) + with self.assertRaises(HTTPError) as caught: + urlopen(request, timeout=2) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt") + + # ---- non-finite serializes as JSON null, end to end --------------------- + + def test_nan_logprob_serializes_as_json_null_over_the_wire(self): + with self.request("/v1/completions", { + "model": "test-model", "prompt": "Hé", "echo": True, "logprobs": 1, + "max_tokens": 1, + }) as response: + raw = response.read() + self.assertNotIn(b"NaN", raw) + self.assertNotIn(b"Infinity", raw) + body = json.loads(raw) + self.assertIsNone(body["choices"][0]["logprobs"]["token_logprobs"][0]) + + # ---- a request that never touches logprobs is unaffected ---------------- + + def test_golden_fixture_style_plain_request_is_unaffected(self): + with self.request("/v1/chat/completions", { + "model": "test-model", "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 4, + }) as response: + body = json.load(response) + self.assertIsNone(body["choices"][0]["logprobs"]) + self.assertEqual(self.engine.last_logprobs, 0) + + +class LogprobsSubmitHeaderRegressionTest(unittest.TestCase): + """REGRESSION-COVERAGE: a legacy request (no logprobs asked at all) must + produce a byte-identical SUBMIT header to the predecessor -- the + extension namespace must never appear unless the client opted in.""" + + def test_legacy_request_submit_header_is_byte_identical(self): + request_id = "1" + expected = f"SUBMIT {request_id} 0 5 4 0.25 0.9\n".encode() + b"hello\n" + + def respond(process, frame): + self.assertEqual(frame, expected) + process.stdout.feed( + b"DATA " + request_id.encode() + b" 2\nok\n" + b"DONE " + request_id.encode() + b" STAT 1 2.5 0 1.0 5 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + engine.generate("hello", 4, 0.25, 0.9, chunks.append) + engine.close() + self.assertEqual(process.writes, [expected]) + + def test_opted_in_request_submit_header_carries_the_extension(self): + # The mutation this pins: sending logprobs= on every request (not + # only opted-in ones) would make this test's own legacy sibling + # above fail -- the extension field would show up unconditionally. + request_id = "1" + expected = (f"SUBMIT {request_id} 0 5 4 0.25 0.9 0 logprobs=2\n".encode() + + b"hello\n") + + def respond(process, frame): + self.assertEqual(frame, expected) + process.stdout.feed( + b"ACCEPT " + request_id.encode() + b" 5\n" + b"ECHO " + request_id.encode() + b" 1 0 nan 0\nh\n" + b"DATA " + request_id.encode() + + b" 2 -0.223144 1 3 -0.223144\nok\n" + b"DONE " + request_id.encode() + b" STAT 1 2.5 0 1.0 5 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + engine.generate("hello", 4, 0.25, 0.9, chunks.append, logprobs=2) + engine.close() + self.assertEqual(process.writes, [expected]) + + def test_token_id_prompt_submit_header_carries_ids_extension(self): + # The mutation this pins: dropping `ids=1` from the extension + # namespace when `tok_ids=True` -- the engine would then tok_encode + # the decimal-digit payload as literal text instead of reading it + # as pre-tokenized ids (coli_ids_parse never runs). + request_id = "1" + payload = b"72 233 108" + expected = (f"SUBMIT {request_id} 0 {len(payload)} 4 0.25 0.9 0 ids=1\n".encode() + + payload + b"\n") + + def respond(process, frame): + self.assertEqual(frame, expected) + process.stdout.feed( + b"DATA " + request_id.encode() + b" 2\nok\n" + b"DONE " + request_id.encode() + b" STAT 1 2.5 0 1.0 5 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + engine.generate("72 233 108", 4, 0.25, 0.9, chunks.append, tok_ids=True) + engine.close() + self.assertEqual(process.writes, [expected]) + + def test_token_id_prompt_with_logprobs_submit_header_carries_both(self): + # Both extension fields together, ordered logprobs= before ids=1 + # (the order this server always emits them in). + request_id = "1" + payload = b"72 233 108" + expected = (f"SUBMIT {request_id} 0 {len(payload)} 4 0.25 0.9 0 " + f"logprobs=2 ids=1\n".encode() + payload + b"\n") + + def respond(process, frame): + self.assertEqual(frame, expected) + process.stdout.feed( + b"ACCEPT " + request_id.encode() + b" 3\n" + b"ECHO " + request_id.encode() + b" 1 0 nan 0\nh\n" + b"DATA " + request_id.encode() + + b" 2 -0.223144 1 3 -0.223144\nok\n" + b"DONE " + request_id.encode() + b" STAT 1 2.5 0 1.0 5 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + engine.generate("72 233 108", 4, 0.25, 0.9, chunks.append, + logprobs=2, tok_ids=True) + engine.close() + self.assertEqual(process.writes, [expected]) + + +class LogprobsGoldenResponseRegressionTest(unittest.TestCase): + """REGRESSION-COVERAGE: a golden plain (non-logprobs) request's response + must be byte-identical to the predecessor's -- this feature must not + perturb any response field for a request that never asked for + logprobs.""" + + def setUp(self): + self.engine = FakeEngine() + self.server = APIServer(("127.0.0.1", 0), self.engine, "test-model") + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + self.base = f"http://127.0.0.1:{self.server.server_port}" + + def tearDown(self): + self.server.scheduler.close() + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + def test_golden_plain_completions_response_is_byte_identical(self): + req = Request(self.base + "/v1/completions", + data=json.dumps({"model": "test-model", "prompt": "hi", + "max_tokens": 4, "temperature": 0}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + choice = body["choices"][0] + self.assertEqual(set(choice), {"index", "text", "logprobs", "finish_reason"}) + self.assertIsNone(choice["logprobs"]) + self.assertEqual(choice["text"], "Héllo") + self.assertEqual(choice["finish_reason"], "stop") + + +class LogprobsTailTerminatorAndRangeTest(unittest.TestCase): + """Two mutation survivors carried from the wire-dispatch work: a + wrong (not-LF) DATA terminator byte, and a k outside 0..32 presented + WITH a matching field count -- so the field-count check alone could + never catch it, only the dedicated range check.""" + + def test_data_frame_wrong_terminator_byte_is_a_named_error(self): + # The byte after a DATA frame's tail-extended payload must be LF; + # a wrong byte here (not a closed stream) is the same class of + # protocol error ECHO's terminator check already has a dedicated + # test for -- DATA's own wrong-byte case had none. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed( + b"DATA " + request_id + b" 2 -0.223144 1 3 -0.223144\nokX") + process.stdout.close() + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine DATA terminator"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_k_out_of_range_with_matching_field_count_is_a_named_error(self): + # k = LOGPROBS_TOP_K_CAP + 1, with exactly that many (tid, tlp) + # pairs actually present -- the field-count check passes cleanly, + # so only the dedicated `0 <= k <= LOGPROBS_TOP_K_CAP` range check + # can catch this. + def respond(process, frame): + request_id = frame.split()[1] + k = LOGPROBS_TOP_K_CAP + 1 + pairs = " ".join(f"{i} -0.1" for i in range(k)) + process.stdout.feed( + b"DATA " + request_id + f" 1 -0.5 {k} {pairs}\n".encode() + b"x\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "invalid engine logprob tail: k=.* out of range"): + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + engine.close() + + def test_g17_precision_tail_value_parses_through_the_same_float_call(self): + # The shipped engine prints tail numbers as %.6f; a %.17g + # value (a higher-precision build's own extension, not exactly + # representable at 6 decimals) must parse through the same + # float() call, not a format assumption. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + events = queue.Queue() + request_id = "1" + with engine.pending_lock: + engine.pending[request_id] = events + process.stdout.feed( + b"DATA " + request_id.encode() + + b" 1 -0.30000000000000004 1 3 -0.30000000000000004\nx\n") + kind, (data, record) = events.get(timeout=1) + self.assertEqual((kind, data), ("data", b"x")) + self.assertEqual(record["lp"], -0.30000000000000004) + self.assertEqual(record["topk"], [(3, -0.30000000000000004)]) + engine.close() + + +class _DistinctEchoEngine(FakeEngine): + """Prompt-echo text ("PQ") and generated text ("gen") are chosen so + they share no character -- unlike LogprobsHTTPTest's shared canned + fixture, where the generated text happens to start with the same two + characters as the reconstructed prompt, a coincidence that would make a + `text.startswith(prompt_text)` check pass even with the prompt never + actually prepended.""" + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None, on_accept=None, logprobs=0, + echo=False): + self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + self.last_logprobs = logprobs + self.last_echo = echo + if on_accept is not None: + on_accept({"prompt_tokens": 2}) + on_text("gen") + stats = {"prompt_tokens": 2, "completion_tokens": 1, "length_limited": False} + if logprobs: + stats["logprobs"] = { + "prompt": [ + (0, b"P", {"lp": float("nan"), "topk": []}), + (1, b"Q", {"lp": -0.1, "topk": [(1, -0.1)]}), + ], + "generated": [(b"gen", {"lp": -0.2, "topk": [(2, -0.2)]})], + } + return stats + + +class EchoTextPrependTest(unittest.TestCase): + """`echo: true` must return prompt+completion in `text` itself (the + OpenAI legacy shape), not the completion alone, with `text_offset` + indexing that same concatenation.""" + + def _server(self): + engine = _DistinctEchoEngine() + server = APIServer(("127.0.0.1", 0), engine, "test-model") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.scheduler.close) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + self.addCleanup(thread.join, timeout=2) + return f"http://127.0.0.1:{server.server_port}" + + def test_echo_true_returns_prompt_plus_completion_in_text(self): + base = self._server() + req = Request(base + "/v1/completions", + data=json.dumps({"model": "test-model", "prompt": "PQ", + "echo": True, "logprobs": 1, + "max_tokens": 1}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + choice = body["choices"][0] + logprobs = choice["logprobs"] + self.assertEqual(choice["text"], "PQgen") + self.assertEqual(logprobs["tokens"], ["P", "Q", "gen"]) + self.assertEqual(logprobs["text_offset"], [0, 1, 2]) + for offset in logprobs["text_offset"]: + self.assertLessEqual(offset, len(choice["text"])) + + def test_echo_false_returns_completion_only_in_text(self): + # The control case: without echo, `text` stays completion-only, as + # before this fix. + base = self._server() + req = Request(base + "/v1/completions", + data=json.dumps({"model": "test-model", "prompt": "PQ", + "logprobs": 1, "max_tokens": 1}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + self.assertEqual(body["choices"][0]["text"], "gen") + + +class _SeamSplitEngine(FakeEngine): + """A 3-byte UTF-8 character (the Euro sign, "\u20ac") is split 1+2 + across the prompt/completion seam: its leading byte rides the last + prompt ECHO frame's own bytes, and its two trailing bytes ride the + first generated DATA frame's own bytes. `on_text` is fed exactly what + an incremental UTF-8 decoder with no leading-byte context produces for + those two trailing bytes alone -- two replacement characters -- which + is what a caller decoding the generated bytes independently of the + prompt bytes (the seam bug) would see; a decoder that instead sees the + whole byte stream as one sequence reconstructs the real character.""" + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None, on_accept=None, logprobs=0, + echo=False): + self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + self.last_logprobs = logprobs + self.last_echo = echo + if on_accept is not None: + on_accept({"prompt_tokens": 2}) + on_text("\ufffd\ufffd") + stats = {"prompt_tokens": 2, "completion_tokens": 1, "length_limited": False} + if logprobs: + stats["logprobs"] = { + "prompt": [ + (0, b"A", {"lp": float("nan"), "topk": []}), + (1, b"\xe2", {"lp": -0.1, "topk": [(1, -0.1)]}), + ], + "generated": [(b"\x82\xac", {"lp": -0.2, "topk": [(2, -0.2)]})], + } + return stats + + +class EchoSeamDecodingTest(unittest.TestCase): + """A UTF-8 codepoint split across the prompt/completion seam must be + decoded as one character by one decoder spanning both sides, not as + two independently decoded halves.""" + + def test_split_codepoint_at_the_seam_decodes_as_one_character(self): + engine = _SeamSplitEngine() + server = APIServer(("127.0.0.1", 0), engine, "test-model") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.scheduler.close) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + self.addCleanup(thread.join, timeout=2) + base = f"http://127.0.0.1:{server.server_port}" + req = Request(base + "/v1/completions", + data=json.dumps({"model": "test-model", "prompt": "A\u20ac", + "echo": True, "logprobs": 1, + "max_tokens": 1}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + choice = body["choices"][0] + logprobs = choice["logprobs"] + self.assertEqual(choice["text"], "A\u20ac") + self.assertEqual(choice["text"].count("\u20ac"), 1) + self.assertNotIn("\ufffd", choice["text"]) + offsets = logprobs["text_offset"] + self.assertEqual(offsets, sorted(offsets), "text_offset must be monotonic") + for offset in offsets: + self.assertLessEqual(offset, len(choice["text"])) + + +class _StopTokenLogprobsEngine(FakeEngine): + """Emits three generated chunks and a matching logprob record for each; + a `stop` sequence matching the second chunk exactly withholds it (and + everything after) from `text` -- its own record, and the record after + it, must not survive into the response either.""" + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None, on_accept=None, logprobs=0, + echo=False): + self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + self.last_logprobs = logprobs + self.last_echo = echo + if on_accept is not None: + on_accept({"prompt_tokens": 3}) + for chunk in ("ok ", "STOP", " more"): + on_text(chunk) + if stopped and stopped(): + self.stop_requests += 1 + break + stats = {"prompt_tokens": 3, "completion_tokens": 3, "length_limited": False} + if logprobs: + stats["logprobs"] = {"prompt": [], "generated": [ + (b"ok ", {"lp": -0.1, "topk": [(1, -0.1)]}), + (b"STOP", {"lp": -0.2, "topk": [(2, -0.2)]}), + (b" more", {"lp": -0.3, "topk": [(3, -0.3)]}), + ]} + return stats + + +class LogprobsDroppedStopTokenTest(unittest.TestCase): + """A matched stop sequence withholds its own (and any later) text from + the response -- the logprobs arrays must not still describe a token the + client never received.""" + + def test_filtered_stop_token_record_is_dropped(self): + engine = _StopTokenLogprobsEngine() + server = APIServer(("127.0.0.1", 0), engine, "test-model") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.scheduler.close) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + self.addCleanup(thread.join, timeout=2) + base = f"http://127.0.0.1:{server.server_port}" + req = Request(base + "/v1/completions", + data=json.dumps({"model": "test-model", "prompt": "hi", + "max_tokens": 8, "stop": ["STOP"], + "logprobs": 1}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + choice = body["choices"][0] + self.assertEqual(choice["text"], "ok ") + logprobs = choice["logprobs"] + # Only the ONE record whose bytes are a prefix of "ok " may survive; + # the "STOP" record (and the " more" record after it) must not. + self.assertEqual(logprobs["tokens"], ["ok "]) + self.assertEqual(len(logprobs["token_logprobs"]), 1) + self.assertEqual(len(logprobs["top_logprobs"]), 1) + + +class LogprobsOldEngineAcceptTimeoutTest(unittest.TestCase): + """An engine build that silently rejects the extended SUBMIT header + must not wedge the caller forever.""" + + def test_old_engine_rejection_times_out_with_a_named_503(self): + # engine.generate() must run on its own thread here: if the + # accept-deadline check is missing or broken, the call blocks + # forever on this stub (the "ERROR 0 ..." reply never resolves the + # real pending request -- see the module docstring above), and a + # direct call in this test's own thread would hang the whole suite + # rather than fail this one test. A bounded join turns that failure + # mode into a normal, fast test failure instead. + def respond(process, frame): + process.stdout.feed(b"ERROR 0 BAD_REQUEST\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process), \ + patch("openai_server.LOGPROBS_ACCEPT_TIMEOUT", 0.2): + engine = Engine("glm", "model") + outcome = {} + + def run(): + try: + engine.generate("hello", 4, 0.0, 1.0, lambda _: None, logprobs=1) + except Exception as error: # noqa: BLE001 -- captured, not swallowed + outcome["error"] = error + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=5) + self.assertFalse(thread.is_alive(), + "engine.generate() did not return within the bound -- " + "the accept-deadline check did not fire") + caught = outcome.get("error") + self.assertIsInstance(caught, APIError, f"wrong exception: {caught!r}") + self.assertEqual(caught.status, 503) + self.assertEqual(caught.param, "logprobs") + self.assertEqual(caught.code, "engine_logprobs_unsupported") + engine.close() + + def test_legacy_non_opted_in_request_is_unaffected_by_the_timeout(self): + # The bound only applies when logprobs was requested; a legacy + # request keeps waiting exactly as before -- proven here by a very + # short LOGPROBS_ACCEPT_TIMEOUT that would fire immediately if it + # (wrongly) applied to every request. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 2\nok\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process), \ + patch("openai_server.LOGPROBS_ACCEPT_TIMEOUT", 0.01): + engine = Engine("glm", "model") + chunks = [] + stats = engine.generate("hello", 4, 0.0, 1.0, chunks.append) + self.assertEqual(chunks, ["ok"]) + self.assertEqual(stats["completion_tokens"], 1) + engine.close() + + def test_old_engine_rejects_a_token_id_prompt_times_out_with_a_named_503(self): + # A token-id prompt (`tok_ids=True`, no logprobs) carries its own + # extended SUBMIT field (`ids=1`) -- an old engine that predates + # this extension rejects it the exact same silent way ("ERROR 0 + # BAD_REQUEST", an id that never matches this request), so the + # SAME accept-deadline bound must apply here too, not only when + # logprobs was requested. + def respond(process, frame): + process.stdout.feed(b"ERROR 0 BAD_REQUEST\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process), \ + patch("openai_server.LOGPROBS_ACCEPT_TIMEOUT", 0.2): + engine = Engine("glm", "model") + outcome = {} + + def run(): + try: + engine.generate("72 233 108", 4, 0.0, 1.0, lambda _: None, + tok_ids=True) + except Exception as error: # noqa: BLE001 + outcome["error"] = error + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=5) + self.assertFalse(thread.is_alive(), + "engine.generate() did not return within the bound -- " + "the accept-deadline check did not fire for a token-id " + "prompt") + caught = outcome.get("error") + self.assertIsInstance(caught, APIError, f"wrong exception: {caught!r}") + self.assertEqual(caught.status, 503) + self.assertEqual(caught.param, "prompt") + self.assertEqual(caught.code, "engine_tok_ids_unsupported") + # The pending entry is dropped (treated as cancelled), not left + # for a stray late frame to resolve against. + self.assertNotIn("1", engine.pending) + engine.close() + + +class LogprobsAcceptTimeoutEnvVarTest(unittest.TestCase): + """The tests above patch the module attribute `LOGPROBS_ACCEPT_TIMEOUT` + directly, which proves the accept-deadline logic reacts to that + attribute but proves nothing about the documented public knob: the + `COLI_LOGPROBS_ACCEPT_TIMEOUT` environment variable and its 30-second + default are only read once, at import time. A typo in the env var + name, or a changed default, would ship silently and green under a + patch()-only suite. These tests import the module fresh in a real + subprocess so the env var is read for real, by its real name.""" + + SERVER_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ENV_VAR = "COLI_LOGPROBS_ACCEPT_TIMEOUT" + + def _read_timeout_in_subprocess(self, value=None): + # Build the child's environment from scratch for this one variable: + # start from a copy of ours, drop the real name unconditionally, + # then set it back only if a value was requested -- so neither + # branch is at the mercy of whatever happens to be in this + # process's own environment. + env = dict(os.environ) + env.pop(self.ENV_VAR, None) + if value is not None: + env[self.ENV_VAR] = value + result = subprocess.run( + [sys.executable, "-c", + "import openai_server; print(openai_server.LOGPROBS_ACCEPT_TIMEOUT)"], + cwd=self.SERVER_DIR, env=env, + capture_output=True, text=True, timeout=30) + self.assertEqual(result.returncode, 0, result.stderr) + return float(result.stdout.strip()) + + def test_env_var_by_its_real_name_overrides_the_default(self): + self.assertEqual(self._read_timeout_in_subprocess("5"), 5.0) + + def test_unset_env_var_defaults_to_30(self): + self.assertEqual(self._read_timeout_in_subprocess(), 30.0) + + +class TokenIdPromptWireErrorTest(unittest.TestCase): + """The engine's own vocabulary/structural refusal of a token-id prompt + (`coli_ids_parse` returning -1 for a malformed or out-of-vocabulary id, + c/decode_batch.h) must surface as a named 400 on `prompt`, not the + generic 500 `engine_error` every other unexpected engine RuntimeError + still becomes.""" + + def test_engine_bad_request_for_a_token_id_prompt_is_a_named_400(self): + # Unlike the old-engine-rejects-the-header case ("ERROR 0 + # BAD_REQUEST", never matching a pending request), a real engine + # that understands `ids=1` but rejects THIS id list answers with + # the request's own id -- that is what distinguishes the two wire + # shapes here. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ERROR " + request_id + b" BAD_REQUEST\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaises(APIError) as caught: + engine.generate("1 2 999999999", 4, 0.0, 1.0, lambda _: None, + tok_ids=True) + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.param, "prompt") + engine.close() + + def test_engine_bad_request_without_tok_ids_still_a_generic_error(self): + # The BAD_REQUEST-to-400 mapping is scoped to token-id-prompt + # requests specifically -- a plain-text request that somehow drew + # a matching-id BAD_REQUEST (should not happen in practice; NUL + # bytes and the cache-slot range are both already rejected before + # SUBMIT is ever sent) is NOT silently reinterpreted as a prompt + # validation failure it did not have. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ERROR " + request_id + b" BAD_REQUEST\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaises(RuntimeError) as caught: + engine.generate("hello", 4, 0.0, 1.0, lambda _: None) + self.assertNotIsInstance(caught.exception, APIError) + engine.close() + + def test_out_of_vocabulary_looking_id_is_not_rejected_before_submit(self): + # This server has no accessible vocabulary size at request time + # (FamilyDescriptor carries no vocab_size field, and the real + # embedding-table bound the engine enforces -- c/colibri.c's + # m->c.vocab, read by coli_ids_parse -- is derived per-arch inside + # family_registry.py's geometry functions, not surfaced to the + # HTTP layer). _encode_token_id_prompt is therefore structural + # validation only (non-negative integers), by design: an + # implausibly large id is NOT rejected here -- it is the engine's + # own BAD_REQUEST (mapped above) that is authoritative. + self.assertEqual(_encode_token_id_prompt([1, 2, 999999999]), + "1 2 999999999") + + +class LogprobsEchoBufferingTest(unittest.TestCase): + """Prompt-echo records must not be retained when the caller never asked + to see them, even though the engine still sends every ECHO frame (there + is no wire bit for "logprobs but no echo").""" + + def test_prompt_records_stay_empty_without_echo(self): + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ACCEPT " + request_id + b" 2\n") + process.stdout.feed(b"ECHO " + request_id + b" 1 0 nan 0\nh\n") + process.stdout.feed( + b"ECHO " + request_id + b" 1 1 -0.1 1 3 -0.1\ni\n") + process.stdout.feed( + b"DATA " + request_id + b" 2 -0.223144 1 3 -0.223144\nok\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 2 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + stats = engine.generate("hi", 4, 0.0, 1.0, lambda _: None, logprobs=1, echo=False) + self.assertEqual(stats["logprobs"]["prompt"], []) + self.assertEqual(len(stats["logprobs"]["generated"]), 1) + engine.close() + + def test_prompt_records_are_kept_with_echo(self): + # The control case: the same wire traffic, but the caller DID ask + # to see the echo table -- the records must still be retained. + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ACCEPT " + request_id + b" 2\n") + process.stdout.feed(b"ECHO " + request_id + b" 1 0 nan 0\nh\n") + process.stdout.feed( + b"ECHO " + request_id + b" 1 1 -0.1 1 3 -0.1\ni\n") + process.stdout.feed( + b"DATA " + request_id + b" 2 -0.223144 1 3 -0.223144\nok\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 2 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + stats = engine.generate("hi", 4, 0.0, 1.0, lambda _: None, logprobs=1, echo=True) + self.assertEqual(len(stats["logprobs"]["prompt"]), 2) + engine.close() + + +def _spawn_test_server(case, engine, kv_slots=1, max_tokens=16): + """A throwaway APIServer on an ephemeral port, torn down with the test + case -- the shared harness for the batch-dispatch batteries below. + max_tokens is the operator's --max-tokens/--ngen cap that + generation_options() clamps a request's max_tokens to; it defaults to + 16 (the value every other test in this file was written against) and + is only raised where a test needs headroom above that to reach the + generated-side completion budget.""" + server = APIServer(("127.0.0.1", 0), engine, "test-model", "secret", max_tokens, + kv_slots=kv_slots) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + case.addCleanup(server.scheduler.close) + case.addCleanup(server.shutdown) + case.addCleanup(server.server_close) + case.addCleanup(thread.join, timeout=2) + return f"http://127.0.0.1:{server.server_port}" + + +def _post_completions(base, body, timeout=5): + return urlopen(Request(base + "/v1/completions", + data=json.dumps(body).encode(), + headers={"Authorization": "Bearer secret", + "Content-Type": "application/json"}), + timeout=timeout) + + +class ScriptedEngine(FakeEngine): + """A deterministic pure function of the prompt: the same prompt always + produces the same text, stats, and logprob frames, so a batched-vs- + single comparison can assert exact equality instead of a tolerance.""" + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None, on_accept=None, audio=None, + on_tool=None, image=None, logprobs=0, echo=False, tok_ids=False): + self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + self.last_logprobs = logprobs + self.last_echo = echo + self.last_tok_ids = tok_ids + if on_accept is not None: + on_accept({"prompt_tokens": len(prompt)}) + for chunk in ("out<", prompt, ">"): + on_text(chunk) + stats = {"prompt_tokens": len(prompt), "completion_tokens": 3, + "length_limited": len(prompt) % 2 == 0} + if logprobs: + stats["logprobs"] = self.scripted_channel(prompt, logprobs) + return stats + + def scripted_channel(self, prompt, engine_k): + """One echo record per whitespace piece of the (encoded) prompt -- + position 0 carrying the real wire's nothing-to-condition-on + sentinel -- with logprob values derived from the piece itself, plus + one generated record. Same record shapes as FakeEngine's canned + channel, but prompt-dependent.""" + echoed = [] + for pos, piece in enumerate(prompt.split()): + if pos == 0: + echoed.append((pos, piece.encode(), {"lp": float("nan"), "topk": []})) + continue + lp = -float(len(piece)) - pos / 8.0 + topk = [(pos, lp), (pos + 1000, lp - 1.0)][:min(engine_k, 2)] + echoed.append((pos, piece.encode(), {"lp": lp, "topk": topk})) + generated = [(b"G", {"lp": -0.25, "topk": [(9, -0.25)][:min(engine_k, 1)]})] + return {"prompt": echoed, "generated": generated} + + +class SecondSubmitErrorEngine(ScriptedEngine): + """The first engine submit in a batch succeeds normally; the second + raises an APIError (an engine-side rejection reached mid-dispatch, as + opposed to a member that never gets this far because pre-submit + validation already rejected it) -- exercises the "a later member's + engine-level failure must not leave an earlier member's output + on the wire" guarantee.""" + + def __init__(self): + super().__init__() + self.calls_made = 0 + + def generate(self, *args, **kwargs): + self.calls_made += 1 + if self.calls_made == 2: + raise APIError(400, "the engine rejected this prompt.", "prompt") + return super().generate(*args, **kwargs) + + +class SecondSubmitBadRequestEngine(ScriptedEngine): + """The second engine submit in a batch raises a bare + RuntimeError("BAD_REQUEST") -- the matching-id engine rejection that + Engine.generate() only converts into an APIError when `tok_ids` is + set. A string batch member never sets tok_ids, so this is the shape a + non-tok_ids engine rejection actually takes on the wire, and it must + land exactly where the flat single-prompt path lands the same + condition: a generic, un-attributed 500, never a member-named 400 + batch_completion() invents on its own.""" + + def __init__(self): + super().__init__() + self.calls_made = 0 + + def generate(self, *args, **kwargs): + self.calls_made += 1 + if self.calls_made == 2: + raise RuntimeError("BAD_REQUEST") + return super().generate(*args, **kwargs) + + +class SecondSubmitContextExceededEngine(ScriptedEngine): + """The second engine submit in a batch raises the CONTEXT_EXCEEDED + APIError exactly as `_engine_error` maps it (client fault, `param` + "messages", `code` "context_length_exceeded") -- the batch must fail as + ONE 400 attributed to the failing member, `prompt[1]`.""" + + def __init__(self): + super().__init__() + self.calls_made = 0 + + def generate(self, *args, **kwargs): + self.calls_made += 1 + if self.calls_made == 2: + raise APIError(400, "This model's maximum context length is 4094 tokens, " + "however your prompt resulted in at least 5000 tokens.", + "prompt", "context_length_exceeded") + return super().generate(*args, **kwargs) + + +class SecondSubmitServerFaultEngine(ScriptedEngine): + """The second engine submit in a batch raises the accept-deadline + APIError Engine.generate() raises when the engine build does not + accept a per-token-logprobs request in time -- a server_error: the + engine's own capability gap, never any one member's content.""" + + def __init__(self): + super().__init__() + self.calls_made = 0 + + def generate(self, *args, **kwargs): + self.calls_made += 1 + if self.calls_made == 2: + raise APIError( + 503, "The colibri engine did not accept a per-token logprobs " + "request in time; it may not support the per-token logprobs " + "extension.", "logprobs", "engine_logprobs_unsupported", + "server_error") + return super().generate(*args, **kwargs) + + +class BatchCompletionHTTPTest(unittest.TestCase): + """Acceptance tests for real multi-prompt batches on /v1/completions + against a real APIServer, with ScriptedEngine standing in for the + engine subprocess: batch admission, prompt[i] member attribution, and + (once dispatch is wired) the assembled response itself.""" + + @classmethod + def setUpClass(cls): + cls.engine = ScriptedEngine() + cls.server = APIServer(("127.0.0.1", 0), cls.engine, "test-model", "secret", 16, + kv_slots=1) + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + cls.base = f"http://127.0.0.1:{cls.server.server_port}" + + @classmethod + def tearDownClass(cls): + cls.server.scheduler.close() + cls.server.shutdown() + cls.server.server_close() + cls.thread.join(timeout=2) + + def _json(self, body): + with _post_completions(self.base, {"model": "test-model", **body}) as response: + return json.load(response) + + def _reject(self, body, param): + with self.assertRaises(HTTPError) as caught: + _post_completions(self.base, {"model": "test-model", **body}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], param) + return error + + def test_break_it_mixed_element_types_are_named_400(self): + # Whole-array shape defects carry param "prompt" with NO member + # attribution (the array as a whole is malformed). + for bad in (["a", [1, 2]], [[1, 2], "a"], [1, "a"], ["a", 1], + [None], [{"p": 1}], [True, False], [1.5, 2]): + with self.subTest(bad=bad): + error = self._reject({"prompt": bad, "max_tokens": 1}, "prompt") + self.assertEqual(error["code"], "invalid_value") + + def test_break_it_bad_member_elements_are_member_attributed(self): + # A member-attributable 400 carries error.param = "prompt[i]" -- + # the literal indexed param, not a bare "prompt" with the index + # buried in prose. + for bad_member in ([3, "x"], [3, -4], [3, True], []): + with self.subTest(bad_member=bad_member): + error = self._reject({"prompt": [[1, 2], bad_member], + "max_tokens": 1}, "prompt[1]") + self.assertEqual(error["code"], "invalid_value") + self.assertIn("prompt[1]:", error["message"]) + + def test_break_it_empty_forms_are_named_400(self): + error = self._reject({"prompt": [], "max_tokens": 1}, "prompt") + self.assertEqual(error["code"], "invalid_value") + # An empty STRING member is member-attributed (param "prompt[1]", + # code null -- the single-prompt empty error's own code) and the + # message carries the same "prompt[1]:" prefix every other + # member-attributed message carries. + error = self._reject({"prompt": ["a", ""], "max_tokens": 1}, "prompt[1]") + self.assertIsNone(error["code"]) + self.assertIn("prompt[1]:", error["message"]) + error = self._reject({"prompt": [[1], []], "max_tokens": 1}, "prompt[1]") + self.assertEqual(error["code"], "invalid_value") + + def test_group_score_opt_in_is_refused_fail_closed(self): + # No group-scoring routing exists in this build, and its contract + # changes the response shape -- silently ignoring the opt-in would + # be a semantic surprise. Named 400, checked on both the array and + # the flat request shape, before any prompt intake or engine work. + error = self._reject({"prompt": ["a", "b"], "group_score": True, + "max_tokens": 1}, "group_score") + self.assertEqual(error["code"], "unsupported_value") + error = self._reject({"prompt": "hi", "group_score": True, + "max_tokens": 1}, "group_score") + self.assertEqual(error["code"], "unsupported_value") + # `false` and `null` are accepted as absent. + body = self._json({"prompt": ["a", "b"], "group_score": False, "max_tokens": 1}) + self.assertEqual(len(body["choices"]), 2) + body = self._json({"prompt": "a", "group_score": None, "max_tokens": 1}) + self.assertEqual(body["object"], "text_completion") + + def test_group_score_only_literal_absence_or_false_is_safe(self): + # Only literal absence, `None`, or `False` are accepted -- every + # other value is refused, including values a truthiness check + # would fold into "absent" (`0`, `0.0`) or into "present" without + # being a real opt-in (`""`, `"false"`, `[]`, `{}`, `"0"`). This + # pins the guard to identity comparison, not `in (None, False)`, + # which Python's `==` folds `0`/`0.0` into `False`. + for bad in (0, 0.0, "", "false", [], {}, "0"): + with self.subTest(bad=bad): + error = self._reject({"prompt": ["a", "b"], "group_score": bad, + "max_tokens": 1}, "group_score") + self.assertEqual(error["code"], "unsupported_value") + error = self._reject({"prompt": "hi", "group_score": bad, + "max_tokens": 1}, "group_score") + self.assertEqual(error["code"], "unsupported_value") + + def test_group_score_true_precedes_a_malformed_array_prompt(self): + # The guard runs before array intake, so a request that is BOTH + # shape-malformed AND carries the opt-in fails for `group_score`, + # never for `prompt` -- proof the guard precedes intake rather + # than merely preceding a successful dispatch. + error = self._reject({"prompt": [1, "a"], "group_score": True, + "max_tokens": 1}, "group_score") + self.assertEqual(error["code"], "unsupported_value") + + def test_stream_and_n_together_defer_to_generation_options_first(self): + # Batch validation checks stream last, the same position the flat + # path checks it in -- so a request carrying both `stream: true` + # and `n: 2` is refused for `n`, not `stream`, matching the flat + # path's own precedence. + error = self._reject({"prompt": ["a", "b"], "stream": True, "n": 2}, "n") + self.assertEqual(error["code"], "unsupported_value") + + def test_budget_does_not_apply_to_a_single_tokenized_prompt(self): + # The flat and batch-of-one shapes keep their pre-batch behavior -- + # oversize single prompts stay the engine's own CONTEXT_EXCEEDED + # business, not the batch budget's (a length-1 array never reaches + # batch_completion() at all). + over = 65536 + 1000 + body = self._json({"prompt": [[3] * over], "max_tokens": 1}) + self.assertEqual(len(body["choices"]), 1) + + def test_break_it_streaming_with_array_prompt_is_named_400(self): + error = self._reject({"prompt": ["a", "b"], "stream": True}, "stream") + self.assertEqual(error["code"], "unsupported_parameter") + + def test_break_it_n_above_one_with_array_prompt_is_named_400(self): + error = self._reject({"prompt": ["a", "b"], "n": 2}, "n") + self.assertEqual(error["code"], "unsupported_value") + + def test_rejected_batch_makes_no_engine_submits(self): + # A malformed member fails the WHOLE request before any engine work + # starts -- never a partial batch. + before = len(self.engine.calls) + self._reject({"prompt": ["a", ""], "max_tokens": 1}, "prompt[1]") + self._reject({"prompt": [[1], []], "max_tokens": 1}, "prompt[1]") + self._reject({"prompt": ["a", "b"], "stream": True}, "stream") + self.assertEqual(len(self.engine.calls), before) + + def test_mid_batch_engine_failure_never_emits_a_partial_response(self): + # A member that fails only once engine dispatch is already under + # way (as opposed to a member that fails pre-submit validation) + # must still fail the WHOLE request -- never a 200 carrying just + # the members that happened to finish before it. + engine = SecondSubmitErrorEngine() + base = _spawn_test_server(self, engine) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", + "prompt": ["a", "b", "c"], "max_tokens": 1}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt[1]") + self.assertIn("prompt[1]:", error["message"]) + # This is the engine's own client-fault rejection of one member's + # content, not a shape defect caught pre-submit -- it carries no + # `code` at all, the same `code: null` the single-prompt path's + # own engine rejections carry, never `invalid_value`. + self.assertIsNone(error["code"]) + # The third member must never have been submitted either -- one + # named failure ends the whole batch, not just the failing member. + self.assertEqual(engine.calls_made, 2) + + def test_context_exceeded_mid_batch_names_member_in_param(self): + # The engine's own CONTEXT_EXCEEDED rejection of one member's + # prompt length is a client-actionable fault: the batch fails as + # one named 400 attributed to that member, the same way any other + # engine-side client-fault rejection mid-batch is attributed. + engine = SecondSubmitContextExceededEngine() + base = _spawn_test_server(self, engine) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", + "prompt": ["a", "b", "c"], "max_tokens": 1}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt[1]") + self.assertEqual(error["code"], "context_length_exceeded") + self.assertIn("prompt[1]:", error["message"]) + # The third member must never have been submitted either. + self.assertEqual(engine.calls_made, 2) + + def test_id_batch_rejected_for_non_glm_engine(self): + base = _spawn_test_server(self, NonGlmEngine()) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", + "prompt": [[1, 2], [3, 4]], "max_tokens": 1}) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt") + + def test_string_batch_works_on_non_glm_engine(self): + # Sequential text submits carry no numeric-logprobs/token-id + # extension fields, so a string batch is engine-agnostic. + engine = NonGlmEngine() + base = _spawn_test_server(self, engine) + with _post_completions(base, {"model": "test-model", + "prompt": ["a", "b"], "max_tokens": 4}) as response: + body = json.load(response) + self.assertEqual([choice["text"] for choice in body["choices"]], + ["Héllo", "Héllo"]) + self.assertEqual(engine.last_logprobs, 0) + self.assertEqual(engine.last_tok_ids, False) + + def test_engine_bad_request_without_tok_ids_is_the_same_generic_500_as_flat(self): + # A non-tok_ids engine BAD_REQUEST is NOT subsumed by the + # member-attributed mapping -- it lands as the same generic 500 + # engine_error the flat single-prompt path already uses for this + # condition, with no partial output and no member attribution + # invented for it. + engine = SecondSubmitBadRequestEngine() + base = _spawn_test_server(self, engine) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", + "prompt": ["a", "b", "c"], "max_tokens": 1}) + self.assertEqual(caught.exception.code, 500) + error = json.load(caught.exception)["error"] + self.assertEqual(error["code"], "engine_error") + self.assertIsNone(error["param"]) + # The third member must never have been submitted either. + self.assertEqual(engine.calls_made, 2) + + def test_server_fault_mid_batch_keeps_status_and_param_names_member_in_message_only(self): + # A server-fault APIError (the engine's own capability gap, not + # any one member's content) keeps its original status, code and + # param; only its message gains the failing member's index, so a + # client is never told the wrong prompt is the problem. + engine = SecondSubmitServerFaultEngine() + base = _spawn_test_server(self, engine) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", + "prompt": ["a", "b", "c"], "max_tokens": 1, + "logprobs": 1}) + self.assertEqual(caught.exception.code, 503) + error = json.load(caught.exception)["error"] + self.assertEqual(error["code"], "engine_logprobs_unsupported") + self.assertEqual(error["param"], "logprobs") + self.assertIn("prompt[1]:", error["message"]) + self.assertEqual(engine.calls_made, 2) + + def test_batch_usage_is_built_by_the_shared_usage_helper(self): + # batch_completion() must produce `usage` by calling + # APIHandler.usage(), not a hand-rolled duplicate dict that could + # silently stop matching it if a field is ever added there. + original = APIHandler.usage + + def usage_with_marker(stats): + result = original(stats) + result["_shared_usage_helper_marker"] = True + return result + + with patch.object(APIHandler, "usage", staticmethod(usage_with_marker)): + body = self._json({"prompt": ["a", "b"], "max_tokens": 1}) + self.assertIn("_shared_usage_helper_marker", body["usage"]) + + def test_client_disconnect_mid_batch_stops_further_submits(self): + # Built, not hoped for (#1329): the engine signals an Event only + # once the third member has actually finished, and then blocks on + # a second Event that the test sets only after confirming the + # socket is closed. client_disconnected() is wrapped so the + # moment it first observes the closed socket is itself an Event + # -- the batch loop raises ClientCancelled synchronously inside + # that same call, before any further engine submit, so waiting + # for this Event (rather than a fixed sleep) is enough to know + # engine.calls has already reached its final count. + member_three_done = threading.Event() + resume = threading.Event() + disconnect_observed = threading.Event() + + class DisconnectAfterThirdEngine(ScriptedEngine): + def generate(self, *args, **kwargs): + stats = super().generate(*args, **kwargs) + if len(self.calls) == 3: + member_three_done.set() + resume.wait(5) + return stats + + engine = DisconnectAfterThirdEngine() + server = APIServer(("127.0.0.1", 0), engine, "test-model", "secret", 16, + kv_slots=1) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.scheduler.close) + self.addCleanup(server.shutdown) + self.addCleanup(server.server_close) + self.addCleanup(thread.join, 2) + + original_client_disconnected = APIHandler.client_disconnected + + def watched_client_disconnected(self): + seen = original_client_disconnected(self) + if seen: + disconnect_observed.set() + return seen + + body = json.dumps({"model": "test-model", + "prompt": [f"p{i}" for i in range(8)], + "max_tokens": 1}).encode() + sock = socket.create_connection(("127.0.0.1", server.server_port), 5) + request = (f"POST /v1/completions HTTP/1.1\r\nHost: 127.0.0.1\r\n" + f"Authorization: Bearer secret\r\nContent-Type: application/json\r\n" + f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n").encode() + body + with patch.object(APIHandler, "client_disconnected", watched_client_disconnected): + sock.sendall(request) + self.assertTrue(member_three_done.wait(5), "engine never reached the third member") + # An RST rather than a clean FIN, same idiom ClientHangupTest + # uses, so client_disconnected()'s recv() sees the closure + # immediately. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + sock.close() + resume.set() + self.assertTrue(disconnect_observed.wait(5), + "server's own client_disconnected() never observed " + "the closed socket") + self.assertEqual(len(engine.calls), 3, + "batch kept submitting members after the client left") + + def test_batch_admission_is_released_before_the_response_write(self): + # Built, not hoped for: send_json is wrapped so a batch's (more + # than one choice) response parks on an Event the instant its + # write begins. While it is parked, a second client's flat + # request goes to the same kv_slots=1 server -- if the scheduler + # admission the batch holds were still held during that write, + # the second request would queue behind it instead of completing + # promptly. It must complete first: the admission is released + # before send_json runs. + writing = threading.Event() + release_write = threading.Event() + original_send_json = APIHandler.send_json + + def blocking_send_json(self, status, body, request_id=None, headers=None): + if isinstance(body, dict) and len(body.get("choices", [])) > 1: + writing.set() + release_write.wait(5) + return original_send_json(self, status, body, request_id, headers) + + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, kv_slots=1) + first_status = [] + + def first_request(): + with _post_completions(base, {"model": "test-model", + "prompt": ["a", "b"], "max_tokens": 1}) as response: + first_status.append(response.status) + + with patch.object(APIHandler, "send_json", blocking_send_json): + first = threading.Thread(target=first_request) + first.start() + self.assertTrue(writing.wait(5), "batch response write never started") + + second_start = time.monotonic() + with _post_completions(base, {"model": "test-model", "prompt": "c", + "max_tokens": 1}, timeout=2) as response: + second_elapsed = time.monotonic() - second_start + self.assertEqual(response.status, 200) + self.assertLess(second_elapsed, 1.0, + "second client's admission waited behind the first " + "client's batch response write") + + release_write.set() + first.join(5) + + self.assertEqual(first_status, [200]) + + # ---- batch shapes, N in {1, 2, 5, 128}, both forms ---------------------- + + def test_string_batch_shapes(self): + for n in (1, 2, 5, 128): + with self.subTest(n=n): + prompts = [f"p{i}" for i in range(n)] + body = self._json({"prompt": prompts, "max_tokens": 4}) + self.assertEqual(len(body["choices"]), n) + for i, choice in enumerate(body["choices"]): + self.assertEqual(choice["index"], i) + self.assertEqual(choice["text"], f"out") + self.assertEqual(body["usage"]["prompt_tokens"], + sum(len(p) for p in prompts)) + self.assertEqual(body["usage"]["completion_tokens"], 3 * n) + self.assertEqual(body["usage"]["total_tokens"], + body["usage"]["prompt_tokens"] + + body["usage"]["completion_tokens"]) + + def test_id_array_batch_shapes(self): + for n in (1, 2, 5, 128): + with self.subTest(n=n): + prompts = [[100 + i, 200 + i] for i in range(n)] + body = self._json({"prompt": prompts, "max_tokens": 4}) + self.assertEqual(len(body["choices"]), n) + for i, choice in enumerate(body["choices"]): + self.assertEqual(choice["index"], i) + self.assertEqual(choice["text"], f"out<{100 + i} {200 + i}>") + self.assertEqual(self.engine.last_tok_ids, True) + + # ---- batched == N single-prompt requests, field for field --------------- + + def test_batched_id_choices_bit_identical_to_singles(self): + prompts = [[7, 8, 9], [7, 8, 10, 11], [42]] + base = {"max_tokens": 1, "echo": True, "logprobs": 2, "temperature": 0} + batched = self._json({**base, "prompt": prompts}) + singles = [self._json({**base, "prompt": [p]}) for p in prompts] + for i, single in enumerate(singles): + expect = dict(single["choices"][0]) + got = dict(batched["choices"][i]) + self.assertEqual(got.pop("index"), i) + expect.pop("index") + self.assertEqual(got, expect) + self.assertEqual(batched["usage"]["prompt_tokens"], + sum(s["usage"]["prompt_tokens"] for s in singles)) + self.assertEqual(batched["usage"]["completion_tokens"], + sum(s["usage"]["completion_tokens"] for s in singles)) + + def test_batched_string_choices_bit_identical_to_singles(self): + prompts = ["alpha", "béta gamma", "delta!"] + base = {"max_tokens": 4, "temperature": 0} + batched = self._json({**base, "prompt": prompts}) + for i, prompt in enumerate(prompts): + single = self._json({**base, "prompt": prompt}) + expect = dict(single["choices"][0]) + got = dict(batched["choices"][i]) + self.assertEqual(got.pop("index"), i) + expect.pop("index") + self.assertEqual(got, expect) + + +class BoundaryUtf8Engine(FakeEngine): + """Odd calls end the echo stream with the FIRST byte of a two-byte + UTF-8 codepoint; even calls begin with the SECOND byte (0xC3 / 0xA9 -- + a clean 'e-acute' if wrongly joined). The per-member decoder contract + requires each batch member to surface its own replacement character; + a decoder shared across members would join the halves and leak one + prompt's bytes into the next prompt's token text.""" + + def __init__(self): + super().__init__() + self.channel_calls = 0 + + def logprobs_channel(self, engine_k): + self.channel_calls += 1 + if self.channel_calls % 2 == 1: + echoed = [(0, b"A", {"lp": float("nan"), "topk": []}), + (1, b"\xc3", {"lp": -0.5, "topk": []})] + else: + echoed = [(0, b"\xa9", {"lp": float("nan"), "topk": []}), + (1, b"B", {"lp": -0.5, "topk": []})] + return {"prompt": echoed, "generated": []} + + +class HostileEchoEngine(FakeEngine): + """Well-formed frames for the first member, then a duplicate wire + `pos` on the second member's echo frames: the batch must fail as ONE + clean engine_error -- never a partial response, never a hang.""" + + def __init__(self): + super().__init__() + self.channel_calls = 0 + + def logprobs_channel(self, engine_k): + self.channel_calls += 1 + if self.channel_calls == 1: + return super().logprobs_channel(engine_k) + return {"prompt": [(0, b"A", {"lp": float("nan"), "topk": []}), + (0, b"B", {"lp": -0.5, "topk": []})], + "generated": []} + + +class BatchSequenceIsolationTest(unittest.TestCase): + """Per-member echo reassembly and per-member stateful UTF-8 decoding + across a batch's members.""" + + def test_utf8_codepoint_split_across_members_stays_per_sequence(self): + base = _spawn_test_server(self, BoundaryUtf8Engine()) + with _post_completions(base, {"model": "test-model", "prompt": [[1, 2], [3, 4]], + "max_tokens": 1, "echo": True, + "logprobs": 1}) as response: + body = json.load(response) + first = body["choices"][0]["logprobs"]["tokens"] + second = body["choices"][1]["logprobs"]["tokens"] + self.assertEqual(first, ["A", "�"]) + self.assertEqual(second, ["�", "B"]) + self.assertNotIn("é", "".join(first) + "".join(second)) + + def test_hostile_echo_positions_mid_batch_fail_clean_not_hang(self): + base = _spawn_test_server(self, HostileEchoEngine()) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", "prompt": [[1, 2], [3, 4]], + "max_tokens": 1, "echo": True, "logprobs": 1}, + timeout=2) + self.assertEqual(caught.exception.code, 500) + self.assertEqual(json.load(caught.exception)["error"]["code"], "engine_error") + + +class BatchCompletionBudgetHTTPTest(unittest.TestCase): + """The GENERATED-side batch budget: members * the request's effective + max_tokens must not exceed PROMPT_BATCH_COMPLETION_BUDGET, or the + whole batch is refused before any engine submit. Each test spawns its + own server so it can set the operator's --max-tokens cap + independently of the shared 16 every other batch test in this file + uses.""" + + def test_over_budget_batch_is_refused_before_any_engine_submit(self): + # 2 members * 32769 max_tokens = 65538, one over + # PROMPT_BATCH_COMPLETION_BUDGET (65536): without this budget, + # nothing bounds the generated side and the batch would dispatch + # (200, two engine submits) instead of a named 400 with zero + # engine submits. + self.assertEqual(PROMPT_BATCH_COMPLETION_BUDGET, 65536) + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + before = len(engine.calls) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", "prompt": ["a", "b"], + "max_tokens": 32769}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "max_tokens") + self.assertEqual(error["code"], "batch_completion_budget_exceeded") + self.assertIn("2", error["message"]) + self.assertIn("32769", error["message"]) + self.assertIn("65538", error["message"]) + self.assertIn(str(PROMPT_BATCH_COMPLETION_BUDGET), error["message"]) + # The defining property: no engine work started on a refused batch. + self.assertEqual(len(engine.calls), before) + + def test_product_exactly_at_the_budget_is_admitted(self): + # Boundary companion to the regression-pin test above: 2 * 32768 = + # 65536, exactly the budget, must be ADMITTED, not refused -- the + # budget is an inclusive ceiling, same convention as + # PROMPT_BATCH_TOKEN_BUDGET's own boundary test. + self.assertEqual(PROMPT_BATCH_COMPLETION_BUDGET, 65536) + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + with _post_completions(base, {"model": "test-model", "prompt": ["a", "b"], + "max_tokens": 32768}) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 2) + self.assertEqual(len(engine.calls), 2) + + def test_product_one_above_the_budget_is_refused(self): + # The other half of the boundary pair: one token over (32769) trips + # the refusal -- pinned separately from the regression-pin test so a + # future edit to that test's other assertions cannot silently lose + # boundary coverage. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", "prompt": ["a", "b"], + "max_tokens": 32769}) + self.assertEqual(caught.exception.code, 400) + self.assertEqual(json.load(caught.exception)["error"]["code"], + "batch_completion_budget_exceeded") + + def test_budget_is_checked_against_the_clamped_max_tokens_not_the_raw_request(self): + # The check must use `maximum` -- generation_options()'s return + # value AFTER its clamp to the operator's --max-tokens/--ngen -- + # not the client's raw requested max_tokens. Server cap is 100 + # here; the client asks for 1,000,000 (which alone, times 2 + # members, would be 2,000,000 and hugely over budget), but the + # clamped value is 100, and 2 * 100 = 200 is comfortably inside + # the budget, so the batch must be ADMITTED. engine.calls records + # the actual `maximum` each submit carried, so this also confirms + # which value the engine itself received. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=100) + with _post_completions(base, {"model": "test-model", "prompt": ["a", "b"], + "max_tokens": 1000000}) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 2) + self.assertEqual(engine.calls[0][1], 100) + self.assertEqual(engine.calls[1][1], 100) + + def test_budget_does_not_apply_to_a_flat_single_prompt(self): + # A flat (non-array) prompt never reaches batch_completion() at + # all: the same huge max_tokens that trips the batch budget above + # must be unaffected on the flat path, whose own oversize handling + # (the operator's --max-tokens/--ngen clamp) is unchanged here. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + with _post_completions(base, {"model": "test-model", "prompt": "a", + "max_tokens": 1000000}) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 1) + self.assertEqual(engine.calls[0][1], 1000000) + + def test_budget_does_not_apply_to_a_batch_of_one(self): + # A batch-of-one array unwraps to the flat single-prompt path and + # never reaches batch_completion() either, so the same huge + # max_tokens in a length-1 array must also be unaffected. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + with _post_completions(base, {"model": "test-model", "prompt": ["a"], + "max_tokens": 1000000}) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 1) + self.assertEqual(engine.calls[0][1], 1000000) + + def test_a_malformed_member_wins_over_the_completion_budget(self): + # Regression pin: 2 members, one empty, with max_tokens large + # enough that the product also exceeds PROMPT_BATCH_COMPLETION_BUDGET + # (2 * 40000 = 80000 > 65536). Member validation runs before the + # budget check, so the empty member wins the refusal -- the client + # is told to fix prompt[1], never told to shrink max_tokens for a + # request that was never going to reach the engine over that + # member anyway. At 187f770 the budget check ran first, so this + # request was refused with param "max_tokens" / code + # "batch_completion_budget_exceeded" instead. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=1 << 20) + before = len(engine.calls) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", "prompt": ["a", ""], + "max_tokens": 40000}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "prompt[1]") + self.assertIsNone(error["code"]) + self.assertEqual(len(engine.calls), before) + + def test_budget_measures_an_omitted_max_tokens_at_the_server_cap(self): + # An omitted max_tokens is not exempt from the budget: generation_ + # options() returns the operator's own --max-tokens/--ngen cap in + # that case, and the budget check uses that same clamped value. + # 3 members * a 32768 server cap = 98304, over budget; 2 members * + # 32768 = 65536, exactly at it and admitted. The refusal message + # also tells the client that max_tokens was not set, and to set + # one, since the client cannot see the operator's cap otherwise. + engine = ScriptedEngine() + base = _spawn_test_server(self, engine, max_tokens=32768) + with _post_completions(base, {"model": "test-model", + "prompt": ["a", "b"]}) as response: + body = json.load(response) + self.assertEqual(len(body["choices"]), 2) + before = len(engine.calls) + with self.assertRaises(HTTPError) as caught: + _post_completions(base, {"model": "test-model", "prompt": ["a", "b", "c"]}) + self.assertEqual(caught.exception.code, 400) + error = json.load(caught.exception)["error"] + self.assertEqual(error["param"], "max_tokens") + self.assertEqual(error["code"], "batch_completion_budget_exceeded") + self.assertIn("max_tokens", error["message"]) + self.assertIn("was not set", error["message"]) + self.assertEqual(len(engine.calls), before) + + +class SlotAwareBlockingEngine(ScriptedEngine): + """Blocks exactly the request naming `blocked_prompt` until told to + proceed, while any other request dispatched concurrently -- a batch, + in particular -- runs to completion in the meantime. Lets a test hold + one KV slot open on a live admission so the scheduler's only + remaining free slot is forced onto whatever runs next.""" + + def __init__(self): + super().__init__() + self.blocked_prompt = None + self.entered = threading.Event() + self.release = threading.Event() + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + *args, **kwargs): + if prompt == self.blocked_prompt: + self.entered.set() + self.release.wait(2) + return super().generate(prompt, maximum, temperature, top_p, on_text, + cache_slot, *args, **kwargs) + + +class BatchCacheSlotHTTPTest(unittest.TestCase): + """docs/api.md's isolated-KV-contexts section (a batch holds the + engine's scheduler admission -- and so the engine itself -- for the + sum of every member's generation time; validation order for a batch + "matches the flat path for `stream`, `cache_slot`..."): no test + anywhere set `cache_slot` on a batch before this one. A batch that + omits `cache_slot` is admitted ONCE -- the + scheduler picks a single free slot for the whole batch -- and every + member's engine.generate() call must carry that SAME slot, not a + fresh scheduler pick per member and not the unresolved + pre-admission value (`c/openai_server.py`'s `batch_completion`: + `with self.server.scheduler.admit(self.client_disconnected, cache_slot) + as admission: queue_wait, cache_slot = admission`, then every + `submit_one` closes over that same rebound `cache_slot`).""" + + def test_batch_without_cache_slot_shares_the_one_slot_the_scheduler_picked(self): + # kv_slots=2, and slot 0 is held open by a concurrent, deliberately + # blocked single-prompt request that pins cache_slot=0 explicitly. + # With slot 0 unavailable, the scheduler's admit(slot=None) for the + # batch has exactly one candidate: slot 1. If the admitted slot + # were not correctly threaded through to every member's engine + # call -- e.g. a bug that re-read the pre-admission `cache_slot` + # (None) instead of the tuple admit() returned, or that called + # admit() fresh per member -- at least one member would show a + # slot other than 1 (or None), instead of every member agreeing. + engine = SlotAwareBlockingEngine() + engine.blocked_prompt = "hold this slot" + base = _spawn_test_server(self, engine, kv_slots=2) + + holder_errors = [] + + def hold_slot_zero(): + try: + _post_completions(base, {"model": "test-model", + "prompt": "hold this slot", + "cache_slot": 0, "max_tokens": 1}).read() + except Exception as error: + holder_errors.append(error) + + holder = threading.Thread(target=hold_slot_zero) + holder.start() + self.assertTrue(engine.entered.wait(2), + "the slot-0 holder never reached generate()") + try: + with _post_completions(base, {"model": "test-model", + "prompt": ["a", "b", "c"], + "max_tokens": 1}) as response: + body = json.load(response) + finally: + engine.release.set() + holder.join(timeout=2) + self.assertFalse(holder.is_alive()) + self.assertEqual(holder_errors, []) + + self.assertEqual(len(body["choices"]), 3) + batch_slots = [call[4] for call in engine.calls if call[0] in ("a", "b", "c")] + self.assertEqual(len(batch_slots), 3) + # The defining property: every member of the batch agrees on ONE slot. + self.assertEqual(len(set(batch_slots)), 1, + f"batch members landed on different slots: {batch_slots}") + # And it is the only slot free while slot 0 is held -- not None, + # not 0, not a value that never went through admission. + self.assertEqual(batch_slots[0], 1) + # The concurrent single request on its own explicit slot is + # unaffected by the batch: it completed independently, on the + # different slot it asked for. + held_call = next(call for call in engine.calls if call[0] == "hold this slot") + self.assertEqual(held_call[4], 0) + + +class FlushTrackingProcess(FakeProcess): + """A FakeProcess that also counts stdin.flush() calls, so a test can + pin that a write is followed by a flush rather than only that the bytes + landed in `writes`.""" + + def __init__(self, on_write): + super().__init__(on_write) + self.flushes = 0 + + def flush(self): + self.flushes += 1 + + +class DeadStdinProcess(FakeProcess): + """A process whose stdin write always raises BrokenPipeError, standing + in for an engine child that has already died: the pipe is broken, so + the write itself is what surfaces the failure.""" + + def write(self, data): + raise BrokenPipeError(32, "Broken pipe") + + +class ShortWriteProcess(FakeProcess): + """A FakeProcess whose stdin.write() accepts only `chunk` bytes per + call, standing in for the production engine's raw, unbuffered pipe (a + raw FileIO whose write() is a single os.write() and can transfer fewer + bytes than given). Every partial write is recorded in `self.writes` in + order, so a test can reassemble the frame and prove the retry loop + delivered every byte.""" + + def __init__(self, on_write=None, chunk=3): + super().__init__(on_write or (lambda _process, _chunk: None)) + self.chunk = chunk + + def write(self, data): + n = min(self.chunk, len(data)) + self.writes.append(bytes(data[:n])) + self.on_write(self, self.writes[-1]) + return n + + +class StalledWriteProcess(FakeProcess): + """A FakeProcess whose stdin takes a few bytes and then reports that it + took none. A raw pipe may legitimately return a short count, but a + count of zero is no progress: re-offering the same bytes forever is a + hang, so the writer has to give up and raise instead.""" + + def __init__(self, chunk=3): + super().__init__(lambda _process, _chunk: None) + self.chunk = chunk + + def write(self, data): + if self.writes: + return 0 + n = min(self.chunk, len(data)) + self.writes.append(bytes(data[:n])) + return n + + +class UncountedWriteProcess(FakeProcess): + """A FakeProcess whose stdin.write() returns None instead of a count. + Under the RawIOBase contract that answer means the stream is + non-blocking and could not take a single byte, so nothing is recorded + as written -- the writer has to fail closed rather than treat the + missing count as a full write.""" + + def __init__(self): + super().__init__(lambda _process, _chunk: None) + self.offered = [] + + def write(self, data): + self.offered.append(bytes(data)) + return None + + +class SlowStdin: + """A stdin stand-in that appends one byte at a time, with a scheduling + yield between bytes, into a single shared buffer -- so a test can prove + two threads calling _write_frame never interleave their bytes: without + write_lock serializing the two calls, a concurrent writer's bytes land + in the middle of the other frame while this one is mid-write.""" + + def __init__(self): + self.buffer = bytearray() + self.started = threading.Event() # set once a write is under way + + def write(self, data): + for byte in data: + self.buffer.append(byte) + self.started.set() + time.sleep(0.002) + return len(data) + + def flush(self): + pass + + +class EngineWriteCheckingTest(unittest.TestCase): + """The server half of the checked-write contract: every SUBMIT/CANCEL/ + STOP write onto the engine's stdin is checked, and a failed write + surfaces as a named RuntimeError -- never silence. + + Every call below that could wait -- on a write loop, on an engine + response, on an HTTP round trip -- is driven through _bounded, so that + a defect in the code under test or in a fake is reported as a failure + inside the bound instead of stalling the run.""" + + def _bounded(self, call, *args, timeout=10): + """Run `call(*args)` on a daemon thread and hand back whatever it + raised (None if it returned). A call that has not come back within + `timeout` seconds fails the test rather than hanging it: nothing + here is allowed to wait on an engine that may never answer.""" + raised = [] + thread = threading.Thread( + target=lambda: raised.append(self._capture(call, *args)), + daemon=True) + thread.start() + thread.join(timeout) + self.assertFalse(thread.is_alive(), + f"the call under test never returned within {timeout}s") + return raised[0] + + @staticmethod + def _capture(call, *args): + """Run `call` and hand back whatever it raised, so a test can drive + it on a worker thread and still assert on the exception.""" + try: + call(*args) + except BaseException as error: # noqa: BLE001 - handed to the caller + return error + return None + + def _wait_until(self, predicate, what, timeout=10): + """Poll `predicate` until it holds, and fail the test if it has not + held within `timeout` seconds -- a bounded stand-in for waiting on + a side effect that a background thread produces.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.01) + self.fail(f"{what} did not happen within {timeout}s") + + def test_dead_engine_submit_is_a_named_500_engine_error_not_silence(self): + # Regression pin: at the pre-fix revision this write's BrokenPipeError + # (a ConnectionError subclass) falls straight into do_POST's + # client-hangup handler and the client sees a silent connection + # close instead of an answer. + process = DeadStdinProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + base = _spawn_test_server(self, engine) + raised = self._bounded( + lambda: _post_completions(base, {"model": "test-model", "prompt": "hi", + "max_tokens": 1}, timeout=3)) + self.assertIsInstance(raised, HTTPError) + self.assertEqual(raised.code, 500) + error = json.load(raised)["error"] + self.assertEqual(error["code"], "engine_error") + # and the failed request must not leak a pending entry + self.assertEqual(engine.pending, {}) + + def test_submit_write_failure_wraps_oserror_and_names_the_frame(self): + process = DeadStdinProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + raised = self._bounded(engine.generate, "hi", 8, 0.7, 0.9, lambda _: None) + self.assertIsInstance(raised, RuntimeError) + self.assertNotIsInstance(raised, ConnectionError) + self.assertIn("failed to write SUBMIT", str(raised)) + + def test_write_frame_wraps_oserror_and_names_the_frame(self): + process = DeadStdinProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + raised = self._bounded(engine._write_frame, b"CANCEL 7\n", "CANCEL") + self.assertIsInstance(raised, RuntimeError) + self.assertNotIsInstance(raised, ConnectionError) + self.assertIn("failed to write CANCEL", str(raised)) + + def test_submit_refuses_before_any_write_when_the_process_has_exited(self): + # require_running semantics: a process already dead when generate() + # is first called is refused before request_id registration. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + process.returncode = 1 + raised = self._bounded(engine.generate, "hi", 8, 0.7, 0.9, lambda _: None) + self.assertIsInstance(raised, RuntimeError) + self.assertIn("colibri engine is not running", str(raised)) + self.assertEqual(process.writes, []) + + def test_submit_refuses_before_the_write_lock_write_when_process_exits_late(self): + # require_running semantics, at the write_lock's own check: the + # process is still alive at request_id registration but has exited + # by the time the write is about to happen -- the check inside the + # lock must catch this too, before any byte reaches stdin. A stub + # poll() answers None on the registration-time call and non-None + # thereafter, standing in for the engine dying in between. + class LateExitProcess(FakeProcess): + def __init__(self, on_write): + super().__init__(on_write) + self.poll_calls = 0 + + def poll(self): + self.poll_calls += 1 + return None if self.poll_calls == 1 else 1 + + def respond(process, frame): + # A frame reaching stdin at all means the require_running check + # was skipped; answer immediately so a defect here is a fast, + # clean test failure rather than a hang waiting on a response + # that a correct implementation would never let through. + if frame.startswith(b"SUBMIT"): + process.stdout.feed(b"ERROR " + frame.split()[1] + b" CANCELLED\n") + + process = LateExitProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + raised = self._bounded(engine.generate, "hi", 8, 0.7, 0.9, lambda _: None) + self.assertIsInstance(raised, RuntimeError) + self.assertIn("colibri engine is not running", str(raised)) + self.assertEqual(process.writes, []) + + def test_submit_and_cancel_writes_are_each_followed_by_a_flush(self): + def respond(process, frame): + fields = frame.split() + if fields[0] == b"CANCEL": + process.stdout.feed(b"ERROR " + fields[1] + b" CANCELLED\n") + + process = FlushTrackingProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + flag = {"cancelled": False} + outcome = [] + + def generate(): + try: + engine.generate("hello", 8, 0.7, 0.9, lambda _: None, + cancelled=lambda: flag["cancelled"]) + except ClientCancelled: + outcome.append("cancelled") + + thread = threading.Thread(target=generate, daemon=True) + thread.start() + # One flush for the SUBMIT frame before the CANCEL is sent. Waiting + # for the flush rather than for the write keeps the poll bounded + # and free of the write/flush race a writes-only wait would have. + self._wait_until(lambda: process.flushes >= 1, "the SUBMIT frame was flushed") + self.assertEqual(process.flushes, 1) + flag["cancelled"] = True + thread.join(timeout=10) + self.assertFalse(thread.is_alive(), "the cancelled generate never returned") + self.assertEqual(outcome, ["cancelled"]) + # A second flush for the CANCEL frame. + self.assertEqual(process.flushes, 2) + + def test_write_frame_loops_until_a_short_write_delivers_every_byte(self): + # The production stdin is a raw, unbuffered pipe (bufsize=0): + # write() is one os.write() and may transfer fewer bytes than + # given. Bite: remove _write_all's retry loop in _write_frame and + # this fails -- only the first `chunk` bytes ever reach stdin. + process = ShortWriteProcess(chunk=3) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + data = b"CANCEL 1234567\n" + self.assertIsNone(self._bounded(engine._write_frame, data, "CANCEL")) + self.assertEqual(b"".join(process.writes), data) + self.assertGreater(len(process.writes), 1) + + def test_submit_write_loops_until_a_short_write_delivers_every_byte(self): + # Same hazard, in the SUBMIT block, with an IMAGE frame ahead of + # the header: both must survive short writes, and the IMAGE bytes + # must still precede the SUBMIT header in the reassembled stream. + # Bite: remove _write_all's retry loop in the SUBMIT block and + # this fails. + # The responder answers the very first partial write so a writer + # that stops early still lets generate() return quickly, but the + # bound does not depend on it: generate() runs through _bounded, so + # a responder that never fires is a failure, not a hang. + answered = [] + + def respond(process, _chunk): + if not answered: + answered.append(True) + process.stdout.feed(b"DONE 1 STAT 1 2.500 50.0 1.25 2 0\n") + + process = ShortWriteProcess(respond, chunk=3) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + image = (b"\x01\x02\x03\x04", 2, 2) + self.assertIsNone(self._bounded( + lambda: engine.generate("hi", 8, 0.7, 0.9, lambda _: None, image=image))) + expected = (b"IMAGE 1 4 2 2\n\x01\x02\x03\x04\n" + b"SUBMIT 1 0 2 8 0.7 0.9\nhi\n") + self.assertEqual(b"".join(process.writes), expected) + self.assertGreater(len(process.writes), 2) + + def test_write_frame_raises_when_stdin_takes_no_bytes(self): + # A short write is retried; a write that takes zero bytes is not + # progress, and retrying it is an unbreakable spin. It has to fail + # closed as the same named engine-write error a broken pipe gives. + # Bite: drop the zero-count check and this test times out instead + # of passing, which _bounded reports as a failure. + process = StalledWriteProcess() + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + raised = self._bounded(engine._write_frame, b"CANCEL 7\n", "CANCEL") + self.assertIsInstance(raised, RuntimeError) + self.assertIn("CANCEL", str(raised)) + self.assertIn("3 of 9 bytes", str(raised)) + + def test_write_frame_raises_when_stdin_reports_no_count(self): + # RawIOBase.write answers None when a non-blocking stream could not + # take a byte -- no progress, exactly like a zero count, and not an + # uncounted full write. Bite: treat None as a completed write and + # the frame is silently dropped, which is the stdin desynchronization + # this whole group exists to prevent. + process = UncountedWriteProcess() + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + raised = self._bounded(engine._write_frame, b"STOP 9\n", "STOP") + self.assertIsInstance(raised, RuntimeError) + self.assertIn("STOP", str(raised)) + self.assertIn("0 of 7 bytes", str(raised)) + # and the writer gave up on the first refusal rather than spinning + self.assertEqual(process.offered, [b"STOP 9\n"]) + + def test_write_frame_lock_prevents_interleaved_writes(self): + # Bite: drop `with self.write_lock:` from _write_frame -- two + # threads writing through a slow stdin then interleave their + # bytes and this fails. + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + slow_stdin = SlowStdin() + engine.process.stdin = slow_stdin + + first = b"CANCEL 111\n" + second = b"STOP 222\n" + errors = [] + threads = [ + threading.Thread( + target=lambda d=data, f=frame: errors.append( + self._capture(engine._write_frame, d, f)), + daemon=True) + for data, frame in ((first, "CANCEL"), (second, "STOP")) + ] + threads[0].start() + # Start the second writer only once the first is demonstrably mid + # frame, so the test really does put two writes in flight at once. + self.assertTrue(slow_stdin.started.wait(timeout=10), + "the first frame write never started") + threads[1].start() + for thread in threads: + thread.join(timeout=10) + self.assertFalse(thread.is_alive(), "a frame write never finished") + self.assertEqual(errors, [None, None]) + # Serialized by write_lock: the second write cannot start until the + # first completes, so the buffer holds each frame whole, in order. + self.assertEqual(bytes(slow_stdin.buffer), first + second) + + +class WireTranscriptTest(unittest.TestCase): + """The server->engine stdin bytes for a fixed request sequence (a plain + SUBMIT, a SUBMIT carrying an IMAGE frame, a CANCEL, and a STOP) must be + byte-identical to the pre-existing wire format. + + BASE_TRANSCRIPT below was captured by running the same four + Engine.generate() calls as _run_fixed_sequence against the commit + before this change, in an isolated scratch checkout, with a FakeProcess + responder recording every frame written to stdin and concatenating + them; the captured bytes are pasted here verbatim as the expectation. + This test reproduces the identical call sequence against the current + code and asserts the two byte strings are equal -- any change to frame + order, content, or a stray extra/missing byte fails it.""" + + BASE_TRANSCRIPT = ( + b"SUBMIT 1 0 5 8 0.7 0.9\nhello\n" + b"IMAGE 2 4 2 2\n\x01\x02\x03\x04\nSUBMIT 2 0 8 8 0.7 0.9\ndescribe\n" + b"SUBMIT 3 0 9 8 0.7 0.9\ncancel-me\nCANCEL 3\n" + b"SUBMIT 4 0 7 8 0.7 0.9\nstop-me\nSTOP 4\n" + ) + + def _run_fixed_sequence(self): + writes = [] + + def respond_submit_only(process, frame): + writes.append(frame) + fields = frame.split() + if fields[0] == b"SUBMIT": + request_id = fields[1] + process.stdout.feed(b"DATA " + request_id + b" 1\nx\n" + b"DONE " + request_id + b" STAT 1 1 0 1 2 0\n") + + process = FakeProcess(respond_submit_only) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + self.addCleanup(engine.close) + + # 1) a plain SUBMIT, no IMAGE frame. + engine.generate("hello", 8, 0.7, 0.9, lambda _: None) + + # 2) a SUBMIT with an IMAGE frame ahead of it, one lock acquisition. + image = (b"\x01\x02\x03\x04", 2, 2) + engine.generate("describe", 8, 0.7, 0.9, lambda _: None, image=image) + + # 3) a CANCEL sent before the first frame arrives (mirrors + # test_cancels_generation_before_first_frame above). + def respond_cancel(process, frame): + writes.append(frame) + fields = frame.split() + if fields[0] == b"CANCEL": + process.stdout.feed(b"ERROR " + fields[1] + b" CANCELLED\n") + + process.on_write = respond_cancel + flag = {"cancelled": False} + + def run_cancel(): + try: + engine.generate("cancel-me", 8, 0.7, 0.9, lambda _: None, + cancelled=lambda: flag["cancelled"]) + except ClientCancelled: + pass + + thread = threading.Thread(target=run_cancel) + thread.start() + for _ in range(200): + if any(frame.startswith(b"SUBMIT 3") for frame in writes): + break + time.sleep(0.01) + time.sleep(0.05) + flag["cancelled"] = True + thread.join(timeout=2) + + # 4) a STOP sent after one DATA frame (mirrors + # test_stops_generation_through_successful_done_path above). + def respond_stop(process, frame): + writes.append(frame) + fields = frame.split() + if fields[0] == b"SUBMIT": + request_id = fields[1] + process.stdout.feed(b"DATA " + request_id + b" 1\nx\n") + elif fields[0] == b"STOP": + process.stdout.feed(b"DONE " + fields[1] + b" STAT 1 1 0 1 2 0\n") + + process.on_write = respond_stop + output = [] + engine.generate("stop-me", 8, 0.7, 0.9, output.append, + stopped=lambda: output == ["x"]) + + return b"".join(writes) + + def test_wire_transcript_is_byte_identical_to_base(self): + self.assertEqual(self._run_fixed_sequence(), self.BASE_TRANSCRIPT) + + +@unittest.skipUnless(os.name == "posix", + "SIGPIPE disposition and a real fork/pipe child are POSIX-only; " + "the policy under test does not exist on Windows.") +class SigpipeDispositionTest(unittest.TestCase): + """The disconnected-consumer policy on the engine side of the pipe, and + the precondition that makes it hold: the engine child is launched under + the default POSIX SIGPIPE disposition.""" + + # The stand-in writer restores SIG_DFL explicitly because CPython + # re-ignores SIGPIPE at interpreter startup -- SIG_DFL is the exec-time + # disposition the real engine inherits from Popen(restore_signals=True) + # and never changes. Frames are protocol-shaped DATA frames; COMPLETED + # on stderr marks a writer that outlived the disconnect (must never + # appear in the default-disposition arm). + WRITER = ( + "import os, signal, sys\n" + "signal.signal(signal.SIGPIPE, {disposition})\n" + "out = os.fdopen(1, 'wb', buffering=0)\n" + "try:\n" + " for i in range(1000000):\n" + " out.write(b'DATA 1 2\\nok\\n')\n" + "except BrokenPipeError:\n" + " sys.stderr.write('EPIPE\\n')\n" + " sys.exit(3)\n" + "sys.stderr.write('COMPLETED\\n')\n" + ) + + def _run_writer(self, disposition): + process = subprocess.Popen( + [sys.executable, "-c", self.WRITER.format(disposition=disposition)], + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + try: + # Consume a couple of well-formed frames, then disconnect the + # consumer: closing the read end is exactly what a dying server + # does to the engine's stdout pipe. + head = process.stdout.read(24) + self.assertEqual(head, b"DATA 1 2\nok\nDATA 1 2\nok\n") + process.stdout.close() + stderr = process.stderr.read() + process.stderr.close() + returncode = process.wait(timeout=10) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + return returncode, stderr + + def test_default_disposition_terminates_the_writer_with_no_evidence_after_failure(self): + # The default SIGPIPE disposition (SIG_DFL): a disconnected consumer + # kills the writer with signal 13 at its next write -- fail-closed, + # no DONE/ERROR/PROF evidence records after the failure point. + returncode, stderr = self._run_writer("signal.SIG_DFL") + self.assertEqual(returncode, -signal.SIGPIPE) + self.assertNotIn(b"COMPLETED", stderr) + self.assertNotIn(b"EPIPE", stderr) # the SIG_IGN/EPIPE path never ran + + def test_sigpipe_ignored_fails_closed_on_epipe_instead(self): + # With SIGPIPE ignored, the same write instead sees EPIPE and must + # fail closed: nonzero exit, a named diagnostic, no completion + # record. + returncode, stderr = self._run_writer("signal.SIG_IGN") + self.assertEqual(returncode, 3) + self.assertIn(b"EPIPE", stderr) + self.assertNotIn(b"COMPLETED", stderr) + + def test_engine_launch_uses_the_deployment_default_signal_disposition(self): + # The policy's precondition: the server does not opt the engine out + # of signal restoration -- it passes no restore_signals kwarg to + # Popen at all, so the child inherits Popen's own default rather + # than a value this server chooses. Asserting `.get(..., True)` + # is truthy would pass whether or not the real call ever expresses + # this policy -- it only reflects the probe's own fallback. Assert + # the kwarg is genuinely absent instead, and pin Python's own + # default separately so this test would still catch it if a + # future stdlib version changed that default. + captured = {} + + class _PopenProbe: + def __init__(self, *args, **kwargs): + captured.update(kwargs) + raise RuntimeError("probe stop") + + with patch("openai_server.subprocess.Popen", _PopenProbe): + with self.assertRaisesRegex(RuntimeError, "probe stop"): + Engine("glm", "model") + self.assertNotIn("restore_signals", captured) + default = inspect.signature(subprocess.Popen.__init__).parameters[ + "restore_signals"].default + self.assertIs(default, True) + + if __name__ == "__main__": unittest.main() diff --git a/docs/api.md b/docs/api.md index 8ceeb3ac7..4a5fbde4a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -42,10 +42,285 @@ sequences. The extension The server serves one generation at a time: the model stays in one persistent process, so concurrent HTTP requests queue instead of loading duplicate model copies. Tool calling depends on the active engine; see the support matrix below. -Images, log probabilities, and token penalties return an explicit error rather -than being silently ignored. Audio is accepted only by Inkling checkpoints with -audio support. The default bind address is localhost; set `COLI_API_KEY` before -exposing the server beyond the machine. +Images and token penalties return an explicit error rather than being silently +ignored, with one documented exception: `seed` is accepted and ignored rather +than rejected (see below). Log probabilities are served on the glm engine (see +below) and refused with a named error on every other engine, never silently +ignored. Audio is accepted only by Inkling checkpoints with audio support. The +default bind address is localhost; set `COLI_API_KEY` before exposing the +server beyond the machine. + +### `seed` + +`seed` is accepted (not rejected) for OpenAI-API request-shape compatibility. +It currently has **no effect on any code path, at any temperature**: no +engine, and no field on the wire protocol, reads a per-request seed. The +`glm` and `inkling` engines seed their process-global RNG once from the +`SEED` environment variable at launch, never per request; no other engine +reads `SEED` or any per-request seed at all. Either way the request's +`seed` value goes nowhere. At `temperature: 0` this is moot anyway (greedy +decoding has no distribution to seed), but the same "no effect" is equally +true at `temperature > 0`, where a client might otherwise expect the value +to matter. A true per-request seed is out of scope for this build +regardless. + +### Log probabilities and prompt echo (glm engine only) + +`/v1/completions` accepts the legacy integer `logprobs` (**0–32**; 0 means no +log probabilities at all, see below; the upper bound is bound to the +engine's top-32 read-out interface, and anything above 32 is a named 400) +and boolean `echo`; `/v1/chat/completions` accepts boolean `logprobs` plus +integer `top_logprobs` (0–32) and returns +`choices[].logprobs.content[]` (`{token, logprob, bytes, top_logprobs}` per +generated token) — chat has no `echo` concept and rejects one with a 400. A +non-boolean `echo` is a named 400 (`invalid_value`) on both endpoints, +independent of whether `logprobs` is requested at all. On chat, +`top_logprobs` is type- and range-checked even when `logprobs` is +false or absent, so a malformed `top_logprobs` is a named 400 whether or +not the gate it would feed is open; a valid `top_logprobs` with `logprobs` +off remains a documented no-op. + +The zero semantics are explicit, not a truthiness accident: on +`/v1/completions`, `logprobs: 0`, `false`, and `null` all mean **no log +probabilities** (the request succeeds with `choices[].logprobs: null`, +exactly as if the field were omitted), while boolean `true` is a named 400 — +the legacy field is an integer count, and a boolean carries no count. On +`/v1/chat/completions` the field is a boolean gate (`null` behaves like +`false`; any integer is a named 400). + +`/v1/completions` with `echo: true` returns the full legacy `logprobs` object +(`tokens`, `token_logprobs`, `top_logprobs`, `text_offset`) covering the +echoed prompt plus any generated tokens, and `text` itself is the +reconstructed prompt followed by the completion (the standard OpenAI legacy +behavior for `echo: true`) rather than the completion alone; `echo` without +`logprobs` is a documented no-op. `text_offset` is a character offset into +that same returned `text` string, always counted from 0 — including when +`echo` is false, where `text` holds only the completion and the offsets +describe only that text, not a position within the (unreturned) prompt. The +requested top-k table is **unsorted** on the wire — do not assume the first +entry is the argmax. Per-token values are printed by the engine to six +decimal digits of precision. Non-finite values (a degenerate all-`-inf` +logit row, say) serialize as JSON `null`, never a clamped number. + +Only the glm engine implements this channel; every other engine returns a +named 400 rather than silently ignoring the request. + +Known limitations, current build: + +- **Cost.** Requesting `logprobs` at all — completions or chat, `echo` or + not — makes the engine re-run the ENTIRE prompt through a full read-out + pass to score every position (the wire has one opt-in bit, not a separate + "echo" bit), forfeiting prefix-cache reuse for that request. There is no + long-echo cap; a very long prompt pays a correspondingly large one-shot + activation buffer. +- **Cancellation.** `CANCEL` is not honored while a logprobs-opted-in request + is inside its prefill read-out — the un-cancellable window widens by the + read-out's own wall time on long prompts. +- **Alternative-token labels.** `top_logprobs` entries for candidate token + ids other than the position's own actual token are not decoded text (no + server-side tokenizer exists, by design) — they are labeled + ``. Only the position's own token (identified by an exact + logprob match, not by id) gets its real decoded text. +- **The sampled token is not guaranteed to appear in its own + `top_logprobs` table.** The engine's numeric channel reports the top-k + candidates by its own read-out; if the actually chosen token falls + outside that table, no entry represents it, and the response's + `token_logprobs`/`logprob` field is still the chosen token's own value + read from the DATA/ECHO frame directly, not looked up in the table. +- **A filtered stop token's own record is dropped; a reasoning/tool-call + split's is not.** A matched `stop` sequence is withheld from the + returned text, and its own logprob record is dropped along with it, but + chat's ``/answer split and tool-call parsing can still remove or + rewrite text that a generated-token logprob record continues to + describe — the two are not realigned in this build. +- **An engine build older than this server's per-token logprobs extension + is refused, not silently ignored, but only after a bounded wait.** Such + an engine rejects the whole opted-in request at the wire level in a way + this server cannot see as a rejection of THIS specific request; after + `COLI_LOGPROBS_ACCEPT_TIMEOUT` seconds (default 30) with no + acknowledgment, the request fails with a named 503 rather than hanging. + An opted-in request that never reaches ACCEPT within that window is + treated as cancelled — the server stops waiting on it and answers the + named 503 — rather than left pending indefinitely. +- **Server-side buffering.** The gateway holds a logprobs-opted-in request's + full echo table in memory for the whole request lifetime (no streaming is + allowed together with `logprobs` — the combination is a named 400). + +### Array `prompt` intake on `/v1/completions` (glm engine for token ids) + +`group_score` (a future opt-in that would change the response shape to +continuation-only log-probability arrays) is not implemented yet, and on +`/v1/completions` — flat or array `prompt` alike — is refused outright with +a named 400 (`param: "group_score"`, `code: "unsupported_value"`) rather +than silently ignored; `false` and `null` are accepted as absent. This +guard applies to `/v1/completions` only: `/v1/chat/completions` and +`/v1/messages` do not read the field, so the same request sent to either +of those endpoints is accepted and the opt-in is silently ignored. + +`prompt` also accepts a flat array of non-negative integers — a single +pre-tokenized prompt, sent as ASCII decimal token ids straight to the +engine rather than re-tokenized from text — and, structurally, the two +OpenAI legacy batch forms: an array of strings, or an array of token-id +arrays (this is what unmodified `lm-eval` sends for its tokenized +loglikelihood requests). Token-id prompts, flat or nested, are +glm-engine-only; every other engine returns a named 400 rather than +silently mis-tokenizing the decimal digits as literal text. A malformed or +out-of-vocabulary token id is refused with a named 400 on `prompt` as well +(the engine itself is the source of truth for its vocabulary; this server +does not re-validate ids against it before sending them). + +A single-member array (one string, or one token-id list) is accepted and +produces exactly what the same prompt would have produced as a +single-prompt request — including a nested batch-of-one token-id array, +which unwraps to the identical flat behavior. Each array is validated +structurally regardless of size: it must not be empty, and its element +types must be homogeneous — all strings, all token ids, or all +token-id arrays, never mixed — each a named 400 on `prompt`. + +**A real batch (more than one member) dispatches.** The response carries +one choice per member, `choices[i].index == i`, in the same order the +members were sent — each choice holding exactly what a single-prompt +request for that member alone would have produced (text, the `logprobs` +object under `echo`/`logprobs`, `finish_reason`), and `usage` summed +across every member. Members are submitted to the engine one at a time, +in order; each carries its own independent UTF-8 decoder and its own +echo-position reassembly, so one member's trailing partial character (or +its logprobs table) can never leak into another member's text. + +A batch is all-or-nothing: every member's shape is validated before the +first engine submit — array structure, homogeneity, and (for strings) +non-emptiness — so a malformed member — say, the eighth — is a clean 400 +naming `prompt[7]` with no engine submits made. A member that passes that +shape check but is rejected by the engine itself (a NUL byte, an +out-of-vocabulary token id) can still fail after earlier members have +already been submitted, and that failure is still a clean 400 naming the +member the same way. A failure that is the engine's own fault rather than +any one member's content takes one of two shapes. An engine failure the +server cannot name as a specific typed error (a shutdown, a protocol +error, a matching-id rejection the engine never turned into an `APIError`) +is a single un-attributed 500 for the whole batch, the same failure mode +the flat single-prompt path already uses for the same conditions. A +failure the engine reports as its own capability gap through a typed +`APIError` (it did not accept a per-token logprobs or token-id request in +time, for example) keeps that error's own status, code, and `param` +unchanged — only its message gains the failing member's index, so the +client is never told to fix a prompt that was never the problem. Either +way, nothing is ever written to the client until +every remaining member has finished, so a request never receives a +partial or truncated set of choices; a client that disconnects partway +through a batch stops it at whichever member is in flight, without +submitting the rest. + +A batch holds the engine's scheduler admission — and so the engine +itself — for the sum of every member's generation time; other clients +queue behind it exactly as they would behind one long single-prompt +request. The admission is released before the response is written back to +this client, so a slow reader does not additionally hold the engine +hostage on socket I/O. + +A batch is capped at `PROMPT_BATCH_CAP` members (128) and at +`PROMPT_BATCH_TOKEN_BUDGET` total prompt tokens (65,536) summed across +members — token-id batches count actual tokens, string batches count +total UTF-8 bytes as an upper bound on tokens (no tokenizer is available +server-side). This budget bounds the prompt side of the request only. +The generated side has its own budget, `PROMPT_BATCH_COMPLETION_BUDGET` +(also 65,536): members multiplied by the request's effective `max_tokens` +(after the server's own clamp to its configured `--max-tokens`/`--ngen`) +must not exceed it, or the whole batch is refused before any engine +submit — a batch's assembled response is held in memory in full until +its single write, so this bounds the generated side of what that hold +can grow to. With `echo` and `logprobs` the same hold additionally +retains every member's echoed prompt positions, bounded instead by +`PROMPT_BATCH_TOKEN_BUDGET` (65,536) times the per-position top-k table +(`LOGPROBS_TOP_K_CAP`, 32) — roughly double the generated-side figure +in the worst case, not covered by `PROMPT_BATCH_COMPLETION_BUDGET` +alone. Both boundaries apply to real batches only: a length-1 array is +exempt and +keeps the flat single-prompt path's own oversize handling (the engine's +`CONTEXT_EXCEEDED`, and no cap on requested `max_tokens` beyond the +server's own clamp). `stream: true` and `n` other than 1 are both +refused with a named 400 when `prompt` is an array with more than one +member. For `n` this is the same refusal used everywhere else on this +endpoint. For `stream` it is not: a batch refuses `stream: true` +outright, `param: "stream"`, where the flat single-prompt path only +refuses the narrower `logprobs`-plus-`stream` combination, `param: +"logprobs"` — a client keying off `error.param` to decide which field +to drop is told to drop a different field depending on which path +answered. Validation order +matches the flat path for `stream`, `cache_slot`, and the shape and +homogeneity checks a malformed array trips, with one exception: +`generation_options` (`max_tokens`, `n`, `temperature`, `top_p`, `stop`) +is checked before any member's own shape on a batch, where the flat path +checks `prompt` first — so a request combining both defects (an empty +second member and `n: 2`, say) is refused for `n` on a batch and for +`prompt` on the flat path. The generated-side completion budget (below) +is checked *after* every member's own shape, the opposite order from +`generation_options`: a malformed member always wins over an also-over- +budget batch, so a request combining both defects (an empty second +member and a `max_tokens` that alone would exceed the budget, say) is +refused for `prompt[1]`, never for `max_tokens`. + +A request that omits both `max_tokens` and `max_completion_tokens` is +measured against the server's own configured cap (`--max-tokens`/ +`--ngen`), the same value the flat path uses for an omitted `max_tokens` +— not an arbitrary client-side default — so a large operator cap can +make an ordinary-sized batch exceed the completion budget with no +`max_tokens` in the request at all. The refusal message says so and +tells the client to set an explicit `max_tokens` when that is why the +batch failed. + +Refusals whose `param` is `prompt`, `prompt[i]`, or `max_tokens` carry +one of the following `code` values, or no code at all (`code: null`), +plus any client-fault code the engine itself raises for a member (for +example `context_length_exceeded`, when the engine rejects one member +mid-batch as too long for its context — `param` is rewritten to +`prompt[i]` but the engine's own `code` is preserved, so this list is +not exhaustive): `invalid_value` (a malformed shape, an empty array, or +a malformed or out-of-vocabulary token-id member — whether the shape +check catches it before any submit or the engine itself rejects it +after earlier members have already been submitted — the +member-specific cases carry `param` set to `prompt[i]`), +`unsupported_parameter` (a token-id array on a non-glm engine), +`prompt_batch_cap_exceeded`, `prompt_batch_token_budget_exceeded`, +`batch_completion_budget_exceeded` (`param: "max_tokens"`, the +generated-side budget above), and `engine_tok_ids_unsupported` (`param: +"prompt"`; the engine did not accept a token-id prompt request in time +— a 503 server-fault code, unlike the others in this list, which are +400 client-fault refusals). `code: null` — the single-prompt path's own +code for these conditions — covers an empty STRING member, whether +caught before any submit or rejected by the engine after earlier +members have already been submitted (a NUL byte); it is still +member-attributed as `prompt[i]` even though it carries no code. Other +refusals on the same request — `stream: true` or `n` other than 1 with +more than one member, an invalid `cache_slot`, and so on — carry their +own `param` and `code`, independent of this list. + +### Engine protocol contract: checked writes and SIGPIPE + +The server↔engine stdio protocol is fail-closed in both directions: + +- **Server→engine writes are checked.** Every `SUBMIT`/`CANCEL`/`STOP` + frame write is verified; a failed write (the engine died, its stdin pipe + broke) surfaces as a named HTTP 500 `engine_error` on the affected + request, as long as the response has not already been committed. `CANCEL` + and `STOP` writes only ever happen after `SUBMIT` has been written, mid + response; for a request whose response is already committed as a stream, + a failed write ends the stream instead of producing a 500. +- **Engine→server frames are strictly validated.** The dispatcher checks + the frame grammar of every `ACCEPT`/`DATA`/`ECHO`/`TOOL`/`GRPP`/`GRPG`/ + `GRPS`/`GRPE`/`ERROR`/`PROF`/`DONE` (and telemetry) line; a malformed + frame is a protocol failure that fails every in-flight request with a + 500 and stops the dispatcher, rather than desynchronizing the stream. + No engine in this tree emits `GRPP`/`GRPG`/`GRPS`/`GRPE` yet — the + dispatcher validates and drains them for a group-scoring wire channel + proposed separately. +- **Disconnected consumer (frozen POSIX policy).** The engine child runs + under the **default** SIGPIPE disposition (the server launches it with + signal restoration on, and the engine installs no handler). If the + server-side reader goes away, the engine is terminated by signal 13 + (wait status 141) at its next stdout write — fail-closed: no completion + or error "evidence" can be fabricated after the failure point. This + policy is pinned by a real fork/pipe test in the server suite. ### Tool-calling support @@ -137,6 +412,16 @@ own `{"type":"error","error":{...}}` envelope on this path. Architecture-local features that have not been wired to this protocol are likewise rejected with an explicit error. +Streaming commits its HTTP 200 only once the engine has accepted the prompt, +the same rule the OpenAI-protocol endpoints follow: a refusal discovered +before acceptance (an oversized prompt over the context limit, say, which +is reported as HTTP 400) surfaces with its own mapped HTTP status in the +Anthropic error envelope, not a committed 200 whose event stream then ends +abruptly. On a healthy engine nothing observable changes — the SSE framing +and event order are exactly as documented above. Until the engine accepts, +no bytes are sent at all — a request queued behind another generation +waits silently, exactly as the OpenAI-style streaming path already does. + > The prefill warning below applies here too, and applies *hardest* to Claude Code: > its system prompt and tool catalog are large, and on a disk-streaming CPU path > that is a long silent wait before the first token. Read it before you connect.