fix(chat): continue gRPC responses until natural EOS

scripts/chat_grpc.py

@@ -49,6 +49,7 @@
 
 import argparse
 import sys
+import time
 from typing import List, Optional
 
 
@@ -92,6 +93,7 @@ def _generate_and_print(
     tokenizer,
     new_tokens: List[int],
     max_tokens: int,
+    max_response_tokens: int = 0,
 ) -> int:
     """Drive one append + generate cycle. Streams tokens to stdout
     as they arrive, returns the count emitted. The generator's
@@ -103,36 +105,74 @@ def _generate_and_print(
     print("kakeya> ", end="", flush=True)
     n = 0
     accumulated = []
+    started = time.perf_counter()
+    server_elapsed = 0.0
+    stop_reason = "unknown"
     try:
-        for token_id in session.generate(max_tokens=max_tokens):
-            n += 1
-            accumulated.append(token_id)
-            # Decode incrementally — tokenizer.decode on the running
-            # buffer gives the right text including BPE merges that
-            # span multiple tokens. We re-decode the full buffer
-            # each time (Qwen3-family tokenizers re-decode in <1ms
-            # for a 64-token buffer; per-token decoding loses some
-            # whitespace correctness on the tokenizer level).
-            text_so_far = tokenizer.decode(
-                accumulated, skip_special_tokens=True,
+        while True:
+            before = n
+            remaining = (
+                min(max_tokens, max_response_tokens - n)
+                if max_response_tokens > 0 else max_tokens
             )
-            # Print only the suffix that's new since last frame.
-            if hasattr(_generate_and_print, "_last_text"):
-                last = _generate_and_print._last_text
-            else:
-                last = ""
-            new_text = text_so_far[len(last):]
-            print(new_text, end="", flush=True)
-            _generate_and_print._last_text = text_so_far
+            if remaining <= 0:
+                stop_reason = "client_safety_limit"
+                break
+            for token_id in session.generate(max_tokens=remaining):
+                n += 1
+                accumulated.append(token_id)
+                # Decode incrementally — tokenizer.decode on the running
+                # buffer gives the right text including BPE merges that
+                # span multiple tokens.
+                text_so_far = tokenizer.decode(
+                    accumulated, skip_special_tokens=True,
+                )
+                if hasattr(_generate_and_print, "_last_text"):
+                    last = _generate_and_print._last_text
+                else:
+                    last = ""
+                new_text = text_so_far[len(last):]
+                print(new_text, end="", flush=True)
+                _generate_and_print._last_text = text_so_far
+            server_elapsed += float(
+                getattr(session, "last_total_duration_seconds", 0.0) or 0.0
+            )
+            stop_reason = {
+                1: "max_tokens",
+                2: "eos",
+                3: "cancelled",
+                4: "truncated",
+            }.get(getattr(session, "last_stop_reason", None), "unknown")
+            if stop_reason != "max_tokens":
+                break
+            if n == before:
+                stop_reason = "no_progress"
+                break
     except KeyboardInterrupt:
         print("\n[interrupted]", file=sys.stderr)
+        stop_reason = "interrupted"
     finally:
         # Reset the per-call decoder state so the next turn starts
         # fresh.
         if hasattr(_generate_and_print, "_last_text"):
             del _generate_and_print._last_text
 
     print()  # final newline
+    elapsed = max(time.perf_counter() - started, 1e-9)
+    measured = server_elapsed or elapsed
+    print(
+        f"[{n} tokens · {measured:.2f}s · {n / measured:.2f} tok/s "
+        f"· stop={stop_reason}]",
+        file=sys.stderr,
+        flush=True,
+    )
+    if stop_reason == "client_safety_limit":
+        print(
+            f"[response reached optional --max-response-tokens "
+            f"{max_response_tokens}]",
+            file=sys.stderr,
+            flush=True,
+        )
     return n
 
 
@@ -161,7 +201,11 @@ def main() -> int:
     )
     ap.add_argument(
         "--max-tokens", type=int, default=64,
-        help="max_tokens per turn",
+        help="tokens per streaming Generate RPC; max_tokens continues automatically",
+    )
+    ap.add_argument(
+        "--max-response-tokens", type=int, default=0,
+        help="optional client safety cap per answer; 0 means continue until EOS",
     )
     ap.add_argument(
         "--system-prompt", default="You are a helpful assistant.",
@@ -250,6 +294,7 @@ def _make_session(client):
                         tokenizer=tokenizer,
                         new_tokens=new_tokens,
                         max_tokens=args.max_tokens,
+                        max_response_tokens=args.max_response_tokens,
                     )
                 except KakeyaError as exc:
                     print(f"[runtime error: {exc}]", file=sys.stderr)

tests/scripts/test_chat_grpc.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from scripts.chat_grpc import _generate_and_print
+
+
+class Tokenizer:
+    def decode(self, token_ids, *, skip_special_tokens=True):
+        assert skip_special_tokens
+        return " ".join(str(token) for token in token_ids)
+
+
+class Session:
+    def __init__(self, chunks):
+        self.chunks = list(chunks)
+        self.calls = 0
+        self.last_stop_reason = None
+        self.last_total_duration_seconds = 0.0
+        self.appended = None
+
+    def append(self, token_ids):
+        self.appended = list(token_ids)
+
+    def generate(self, *, max_tokens):
+        tokens, reason, seconds = self.chunks[self.calls]
+        self.calls += 1
+        assert len(tokens) <= max_tokens
+        yield from tokens
+        self.last_stop_reason = reason
+        self.last_total_duration_seconds = seconds
+
+
+def test_continues_max_token_chunks_until_eos(capsys):
+    session = Session([
+        ([11, 12], 1, 1.0),
+        ([21, 22], 2, 2.0),
+    ])
+    count = _generate_and_print(session, Tokenizer(), [9], max_tokens=2)
+    output = capsys.readouterr()
+    assert count == 4
+    assert session.calls == 2
+    assert session.appended == [9]
+    assert "11 12 21 22" in output.out
+    assert "4 tokens" in output.err
+    assert "1.33 tok/s" in output.err
+    assert "stop=eos" in output.err
+
+
+def test_optional_response_cap_is_explicit(capsys):
+    session = Session([
+        ([1, 2], 1, 1.0),
+        ([3, 4], 1, 1.0),
+    ])
+    count = _generate_and_print(
+        session,
+        Tokenizer(),
+        [9],
+        max_tokens=2,
+        max_response_tokens=4,
+    )
+    output = capsys.readouterr()
+    assert count == 4
+    assert "stop=client_safety_limit" in output.err
+    assert "--max-response-tokens 4" in output.err
+
+
+def test_no_progress_breaks_continuation_loop(capsys):
+    session = Session([
+        ([], 1, 0.1),
+    ])
+    assert _generate_and_print(session, Tokenizer(), [9], max_tokens=2) == 0
+    assert "stop=no_progress" in capsys.readouterr().err