serve: ACCEPT frame — oversized prompt is a clean 400, not a broken 200 stream (#597 item 6)

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(prefix<sc->len && prefix<nt && sc->hist[prefix]==tmp[prefix]) prefix++;
     if(prefix<sc->len){ sc->len=prefix; if(m->has_mtp) m->kv_start[m->c.n_layers]=-1;
         kv_disk_truncate(m,sc->len); }

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})

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()