From b4244f1bbc03bad9cfdae2322542c9d0fdddb074 Mon Sep 17 00:00:00 2001 From: JustVugg Date: Sat, 25 Jul 2026 13:12:15 +0200 Subject: [PATCH] serve: ACCEPT frame so an oversized prompt is a clean 400, not a broken 200 stream (#597 item 6) A streaming chat request committed its HTTP 200 (SSE headers) before the engine tokenised the prompt and checked it against CTX. An over-long prompt was only then found -- ERROR CONTEXT_EXCEEDED -- far too late to send the 400 the client needs, so the client saw a broken/empty stream instead of a clean context_length_exceeded. (Non-streaming was already correct: it sends its 200 only after generation.) Add an acceptance frame to the engine<->gateway protocol (mkelcb's design in #597): - c/colibri.c mux_submit(): after all cheap validation (encode, non-empty, fits CTX) and BEFORE prefill, emit `ACCEPT `. Every early failure already returns first with an ERROR, so a request yields exactly one of ACCEPT or an early ERROR. - Engine._dispatch_stdout(): parse ACCEPT -> ("accept", {...}) without popping pending. - Engine.generate(): new on_accept callback; commit exactly once on the first of ACCEPT / DATA / DONE. DATA/DONE before an ACCEPT still commit (an older engine with no ACCEPT frame keeps working); a second ACCEPT is a protocol error. - APIHandler.generation() streaming: defer the SSE 200 + keepalive into start_stream(), fired via on_accept. An ERROR that arrives before ACCEPT now propagates as an APIError with nothing committed, so do_POST returns a normal JSON 400. Tests: 4 Engine-level cases over a fake engine process (ACCEPT-before-DATA fires on_accept first; ERROR-before-ACCEPT raises 400 and never commits; DATA-before-ACCEPT still commits for an old engine; duplicate ACCEPT is a protocol error) + 1 end-to-end HTTP case (streaming context-exceeded returns a clean 400 with context_length_exceeded, not a 200 stream). Full server suite 83/83; the C engine builds clean. Completes the #597 series (items 1/2/4/5 already on dev via #535/#606/#607). Co-Authored-By: Claude Opus 4.8 (1M context) --- c/colibri.c | 5 ++ c/openai_server.py | 97 ++++++++++++++++++++++-------- c/tests/test_openai_server.py | 110 ++++++++++++++++++++++++++++++++-- 3 files changed, 183 insertions(+), 29 deletions(-) diff --git a/c/colibri.c b/c/colibri.c index bb5f7aae..3b0f1371 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -5562,6 +5562,11 @@ static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, GrDraft *g printf("ERROR %llu CONTEXT_EXCEEDED %d %d\n",sub.id,nt,maxctx-2); fflush(stdout); return 0; } + /* #597: the submission is valid (encoded, non-empty, fits the context). Announce it before + * prefill so the gateway commits the streaming 200 only now. Every failure above returned + * first with an ERROR, so a request yields exactly one of ACCEPT or an early ERROR -- which + * lets a CONTEXT_EXCEEDED become a clean HTTP 400 instead of a broken already-200 stream. */ + printf("ACCEPT %llu %d\n",sub.id,nt); fflush(stdout); int prefix=0; while(prefixlen && prefixhist[prefix]==tmp[prefix]) prefix++; if(prefixlen){ sc->len=prefix; if(m->has_mtp) m->kv_start[m->c.n_layers]=-1; kv_disk_truncate(m,sc->len); } diff --git a/c/openai_server.py b/c/openai_server.py index e5b87c29..84a26574 100644 --- a/c/openai_server.py +++ b/c/openai_server.py @@ -1163,6 +1163,15 @@ def _dispatch_stdout(self): events = self.pending.get(request_id) if events is not None: events.put(("data", data)) + 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 + # HTTP stream only now, so an earlier CONTEXT_EXCEEDED stays a clean 400. + request_id = fields[1] + with self.pending_lock: + events = self.pending.get(request_id) + if events is not None: + events.put(("accept", {"prompt_tokens": int(fields[2])})) elif kind == "DONE" and len(fields) >= 7: request_id = fields[1] stats = self._stats(fields[2:]) @@ -1215,7 +1224,7 @@ def _dispatch_stdout(self): self._fail_pending(error) def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None): + cancelled=None, grammar=None, stopped=None, on_accept=None): 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") payload = prompt.encode("utf-8") @@ -1258,9 +1267,27 @@ def decode(data): 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) + while True: kind, value = events.get() - if kind == "data": + 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(): @@ -1274,6 +1301,7 @@ def decode(data): self.process.stdin.write(f"CANCEL {request_id}\n".encode()) self.process.stdin.flush() elif kind == "done": + _accept({"prompt_tokens": None}) tail = decoder.decode(b"", final=True) if tail: on_text(tail) @@ -1634,22 +1662,19 @@ def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=Non return stream_object = "chat.completion.chunk" if chat else object_name - 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("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 + # #597 item 6: DO NOT commit the 200 yet. The engine validates the prompt against the + # context AFTER we would have sent headers, so an oversized prompt used to be a + # CONTEXT_EXCEEDED discovered too late to send a clean 400. Defer the SSE headers into + # start_stream(), fired on the engine's ACCEPT frame (before prefill); an ERROR that + # arrives before ACCEPT propagates as an APIError with nothing committed -> proper 400. + connected = False + stream_started = [False] + ka_thread = [None] # KEEPALIVE: engine.generate() blocks SILENTLY during the (minutes-long) cold # prefill, and the client drops the socket after its idle timeout. A background pump - # emits a reasoning_content "." delta (the channel that reliably resets the client's - # timer and lands in the thinking panel, so answer content stays clean) whenever no - # event has been written for KA_GAP seconds. All wfile writes share ka_lock so the - # pump and event() never interleave; last_write gates the pump so it stays quiet - # while real tokens are flowing (e.g. during decode). + # emits a keepalive delta whenever no event has been written for KA_GAP seconds. All + # wfile writes share ka_lock so the pump and event() never interleave; last_write + # gates the pump so it stays quiet while real tokens are flowing (e.g. during decode). ka_lock = threading.Lock() last_write = [time.time()] ka_stop = threading.Event() @@ -1689,10 +1714,6 @@ def _keepalive(): if time.time() - last_write[0] >= KA_GAP: event(ping) - if chat: - event([{"index": 0, "delta": {"role": "assistant", "content": ""}, - "logprobs": None, "finish_reason": None}]) - def emit(text): choice = ({"index": 0, "delta": {"content": text}, "logprobs": None, "finish_reason": None} if chat else @@ -1710,8 +1731,29 @@ def emit_reasoning(text): # thinking → reasoning_content deltas (chat only # into visible content or the tool-call buffer. glm_think = chat and ARCH != "inkling" - ka_thread = threading.Thread(target=_keepalive, daemon=True) - ka_thread.start() + def start_stream(_accept_info=None): + # #597 item 6: commit the streaming 200 (and start the keepalive) exactly once, + # only after the engine ACCEPTs the prompt. Idempotent: generate() also calls this + # on the first DATA/DONE so an older engine with no ACCEPT frame still streams. + nonlocal connected + if stream_started[0]: + return + stream_started[0] = True + 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("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 + last_write[0] = time.time() + if chat: + event([{"index": 0, "delta": {"role": "assistant", "content": ""}, + "logprobs": None, "finish_reason": None}]) + ka_thread[0] = threading.Thread(target=_keepalive, daemon=True) + ka_thread[0].start() if chat and tools: # Suppress tool-call markers from the streamed content and parse the authoritative # calls from the FULL reply after generation. Hold back a marker-length tail so a @@ -1747,7 +1789,8 @@ def emit_tools(chunk): stop_filter = StopFilter(stop_sequences, emit_tools, ignore_leading_stop) stats = self.server.engine.generate( prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, - lambda: not connected, grammar=grammar, stopped=stop_filter.stopped) + lambda: not connected, grammar=grammar, stopped=stop_filter.stopped, + on_accept=start_stream) stop_filter.finish() if think: think.finish() @@ -1775,13 +1818,17 @@ def emit_plain(chunk): stop_filter = StopFilter(stop_sequences, emit_plain, ignore_leading_stop) stats = self.server.engine.generate( prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, - lambda: not connected, grammar=grammar, stopped=stop_filter.stopped) + lambda: not connected, grammar=grammar, stopped=stop_filter.stopped, + on_accept=start_stream) stop_filter.finish() if content_split: content_split.close() finish = "length" if stats["length_limited"] else "stop" + # generate() returned, so the prompt was ACCEPTed and start_stream() ran; guard anyway. + start_stream() ka_stop.set() # generation done: stop the keepalive pump - ka_thread.join(timeout=2) + if ka_thread[0] is not None: + ka_thread[0].join(timeout=2) final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish} if chat else {"index": 0, "text": "", "logprobs": None, "finish_reason": finish}) diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index ab92acb4..6cdf3cb9 100644 --- a/c/tests/test_openai_server.py +++ b/c/tests/test_openai_server.py @@ -24,8 +24,10 @@ def __init__(self): self.stop_requests = 0 def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None): + cancelled=None, grammar=None, stopped=None, on_accept=None): self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + if on_accept is not None: # simulate the engine's ACCEPT frame (#597) + on_accept({"prompt_tokens": 7}) for chunk in ("Hé", "llo"): on_text(chunk) if stopped and stopped(): @@ -41,11 +43,11 @@ 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): + cancelled=None, grammar=None, stopped=None, on_accept=None): self.entered.set() self.release.wait(2) return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot, - cancelled, grammar, stopped) + cancelled, grammar, stopped, on_accept) class TemplateTest(unittest.TestCase): @@ -1041,8 +1043,10 @@ def __init__(self, chunks): self.chunks = list(chunks) def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, - cancelled=None, grammar=None, stopped=None): + cancelled=None, grammar=None, stopped=None, on_accept=None): self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + if on_accept is not None: # simulate the engine's ACCEPT frame (#597) + on_accept({"prompt_tokens": 7}) for chunk in self.chunks: on_text(chunk) if stopped and stopped(): @@ -1145,5 +1149,103 @@ def test_nonstreaming_splits_reasoning(self): self.assertEqual(message["content"], "final answer") +class AcceptFrameTest(unittest.TestCase): + """#597 item 6: the engine's ACCEPT frame gates the HTTP commit, and its invariants.""" + + def _engine(self, respond): + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + return Engine("glm", "model") + + def test_accept_fires_before_any_data(self): + def respond(process, frame): + rid = frame.split()[1] + process.stdout.feed(b"ACCEPT " + rid + b" 42\n") + process.stdout.feed(b"DATA " + rid + b" 3\nHi!\n") + process.stdout.feed(b"DONE " + rid + b" STAT 1 2.5 0 1.0 4 0\n") + engine = self._engine(respond) + seen = [] + engine.generate("hi", 8, 0.7, 0.9, lambda t: seen.append(("text", t)), + on_accept=lambda info: seen.append(("accept", info))) + engine.close() + self.assertEqual(seen[0], ("accept", {"prompt_tokens": 42})) + self.assertEqual("".join(t for k, t in seen if k == "text"), "Hi!") + + def test_error_before_accept_never_commits(self): + def respond(process, frame): + rid = frame.split()[1] + process.stdout.feed(b"ERROR " + rid + b" CONTEXT_EXCEEDED 5000 4094\n") + engine = self._engine(respond) + accepts = [] + with self.assertRaises(APIError) as caught: + engine.generate("hi", 8, 0.7, 0.9, lambda _: None, + on_accept=lambda info: accepts.append(info)) + engine.close() + self.assertEqual(caught.exception.status, 400) + self.assertEqual(caught.exception.code, "context_length_exceeded") + self.assertEqual(accepts, []) # nothing committed -> HTTP layer can send a clean 400 + + def test_data_before_accept_still_commits_for_old_engine(self): + def respond(process, frame): + rid = frame.split()[1] + process.stdout.feed(b"DATA " + rid + b" 3\nHey\n") + process.stdout.feed(b"DONE " + rid + b" STAT 1 2.5 0 1.0 4 0\n") + engine = self._engine(respond) + accepts, chunks = [], [] + engine.generate("hi", 8, 0.7, 0.9, chunks.append, + on_accept=lambda info: accepts.append(info)) + engine.close() + self.assertEqual("".join(chunks), "Hey") + self.assertEqual(len(accepts), 1) # first DATA implies acceptance (no ACCEPT frame) + self.assertIsNone(accepts[0]["prompt_tokens"]) + + def test_duplicate_accept_is_a_protocol_error(self): + def respond(process, frame): + rid = frame.split()[1] + process.stdout.feed(b"ACCEPT " + rid + b" 10\n") + process.stdout.feed(b"ACCEPT " + rid + b" 10\n") + engine = self._engine(respond) + with self.assertRaisesRegex(RuntimeError, "duplicate ACCEPT"): + engine.generate("hi", 8, 0.7, 0.9, lambda _: None, on_accept=lambda _: None) + engine.close() + + +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") + + +class StreamingContextRejectTest(unittest.TestCase): + """#597 item 6: an oversized prompt on a *streaming* request must return a clean HTTP 400, + not a committed 200 SSE stream that only later discovers the overflow.""" + + def setUp(self): + self.server = APIServer(("127.0.0.1", 0), _ContextExceededEngine(), "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_streaming_context_exceeded_is_clean_400(self): + req = Request(self.base + "/v1/chat/completions", + data=json.dumps({"model": "test-model", "stream": True, + "messages": [{"role": "user", "content": "x" * 100}]}).encode(), + headers={"Content-Type": "application/json"}) + with self.assertRaises(HTTPError) as caught: + urlopen(req, timeout=3) + self.assertEqual(caught.exception.code, 400) # a real 400, not a 200 stream + body = json.load(caught.exception) + self.assertEqual(body["error"]["code"], "context_length_exceeded") + self.assertEqual(body["error"]["param"], "messages") + + if __name__ == "__main__": unittest.main()