fix(agents): complete responses to EOS and ground critique

Continue Generator and Critic streams across chunk boundaries until model EOS, mark safety-limit truncation explicitly, and require the Critic to respect completion status and honest boundaries for open problems.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

docs/ops/distributed-prefill-kv-network.md

@@ -334,6 +334,11 @@ Type a prompt at `prompt>`. The terminal prints allens Prefill progress, then
 streams `generator>` and `critic>` tokens as they arrive, followed by KV hit
 rate, decode tok/s, generation latency and E2E tok/s. `/quit` exits; services
 remain running. Reports remain redacted and appear in the Benchmarks tab.
+Generation uses 64-token streaming chunks and automatically continues to the
+model's EOS. `--max-response-tokens` defaults to 512 as an explicit safety cap;
+reaching it marks the stage incomplete instead of presenting a truncated answer
+as successful. The Critic receives the Generator completion status and must not
+penalize an honest statement that an open problem has no accepted proof.
 
 ## Rollback
 

inference_engine/bench/prefill_fleet_report.py

@@ -65,6 +65,10 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]:
     sources = {source: 0 for source in HIT_SOURCES}
     for stage in normalized:
         sources[stage["hit_source"]] += 1
+    stop_reasons: dict[str, int] = {}
+    for stage in normalized:
+        reason = str(stage.get("stop_reason", "unknown"))
+        stop_reasons[reason] = stop_reasons.get(reason, 0) + 1
     decode = [stage["decode_tok_s"] for stage in normalized]
     prefix_tokens = sum(stage["prefix_tokens"] for stage in normalized)
     hit_tokens = sum(
@@ -83,6 +87,10 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]:
     return {
         "stages_total": len(normalized),
         "stages_failed": sum(not stage.get("ok", False) for stage in normalized),
+        "incomplete_stages": sum(
+            not stage.get("complete", True) for stage in normalized
+        ),
+        "stop_reason_counts": stop_reasons,
         "ttft_p50_s": _median(stage["ttft_s"] for stage in normalized),
         "prefill_tok_s_p50": _median(
             stage["prefill_or_restore_tok_s"] for stage in normalized

inference_engine/network/dashboard.py

@@ -49,8 +49,8 @@ def dashboard_html() -> str:
 $('createRegistration').onclick=async()=>{let r=await fetch('/v1/network/nodes/register',{method:'POST',headers:writeHeaders(),body:JSON.stringify({alias:$('alias').value,address:$('address').value,region:$('region').value,role:'hybrid'})});let j=await r.json();$('pairing').textContent=r.ok?`Pairing token: ${j.pairing_token}\nExpires: ${new Date(j.expires_at*1000).toLocaleTimeString()}`:`Error: ${j.detail||r.status}`;$('pairing').classList.remove('hidden');load()};
 $('createGroup').onclick=async()=>{await fetch('/v1/network/groups',{method:'POST',headers:writeHeaders(),body:JSON.stringify({name:$('groupName').value,node_ids:$('groupNodes').value.split(',').map(x=>x.trim()).filter(Boolean)})});load()};
 function nodePosition(i,total){let a=(i/Math.max(total,1))*Math.PI*2;return {x:50+38*Math.cos(a),y:53+35*Math.sin(a)}}
-function phaseCards(stages){$('benchmarkPhases').innerHTML=stages.map(x=>`<div class="card"><b>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</b><p>${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}</p><div class="muted">TTFT ${num(x.ttft_s)}s<br>Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s<br>Decode ${num(x.decode_tok_s)} tok/s<br>Generation ${num(x.generation_latency_ms_per_token)} ms/token<br>E2E ${num(x.e2e_tok_s)} tok/s</div></div>`).join('')||'<div class="card muted">No stages yet.</div>'}
-async function showBenchmark(id){let r=await fetch('/v1/network/benchmarks/'+encodeURIComponent(id)).then(x=>x.json());phaseCards(r.stages||[]);$('benchmarkDetail').innerHTML=`<b>${esc(r.id)} · ${esc(r.status)}</b><p class="muted">${esc(r.kind)} · ${new Date(r.started_at*1000).toLocaleString()}</p><table><thead><tr><th>Phase</th><th>Source</th><th>TTFT</th><th>Prefill/restore</th><th>Decode</th><th>Latency/token</th><th>E2E</th></tr></thead><tbody>${(r.stages||[]).map(x=>`<tr><td>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</td><td>${esc(x.hit_source)}</td><td>${num(x.ttft_s)}s</td><td>${num(x.prefill_or_restore_tok_s)}</td><td>${num(x.decode_tok_s)}</td><td>${num(x.generation_latency_ms_per_token)}ms</td><td>${num(x.e2e_tok_s)}</td></tr>`).join('')}</tbody></table>`}
+function phaseCards(stages){$('benchmarkPhases').innerHTML=stages.map(x=>`<div class="card"><b>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</b><p>${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}</p><div class="muted">Stop ${esc(x.stop_reason||'n/a')} · ${x.complete===false?'INCOMPLETE':'complete'}<br>TTFT ${num(x.ttft_s)}s<br>Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s<br>Decode ${num(x.decode_tok_s)} tok/s<br>Generation ${num(x.generation_latency_ms_per_token)} ms/token<br>E2E ${num(x.e2e_tok_s)} tok/s</div></div>`).join('')||'<div class="card muted">No stages yet.</div>'}
+async function showBenchmark(id){let r=await fetch('/v1/network/benchmarks/'+encodeURIComponent(id)).then(x=>x.json());phaseCards(r.stages||[]);$('benchmarkDetail').innerHTML=`<b>${esc(r.id)} · ${esc(r.status)}</b><p class="muted">${esc(r.kind)} · ${new Date(r.started_at*1000).toLocaleString()}</p><table><thead><tr><th>Phase</th><th>Source</th><th>Stop</th><th>TTFT</th><th>Prefill/restore</th><th>Decode</th><th>Latency/token</th><th>E2E</th></tr></thead><tbody>${(r.stages||[]).map(x=>`<tr><td>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</td><td>${esc(x.hit_source)}</td><td>${esc(x.stop_reason||'n/a')}</td><td>${num(x.ttft_s)}s</td><td>${num(x.prefill_or_restore_tok_s)}</td><td>${num(x.decode_tok_s)}</td><td>${num(x.generation_latency_ms_per_token)}ms</td><td>${num(x.e2e_tok_s)}</td></tr>`).join('')}</tbody></table>`}
 async function load(){let [s,n,g,live,runs]=await Promise.all([fetch('/v1/network/summary').then(r=>r.json()),fetch('/v1/network/nodes').then(r=>r.json()),fetch('/v1/network/groups').then(r=>r.json()),fetch('/v1/network/benchmarks/live').then(r=>r.json()),fetch('/v1/network/benchmarks?limit=20').then(r=>r.json())]);
 $('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB';let p=s.prefill||{};$('remoteJobs').textContent=fmt(p.remote_jobs);$('remoteHits').textContent=fmt(p.remote_hits);$('reusedTokens').textContent=fmt(p.tokens_reused);$('evictions').textContent=fmt(s.cache_evictions);$('publishFailures').textContent=fmt(p.publish_failures);
 let total=s.cache_bytes_used+s.cache_bytes_free,pct=total?s.cache_bytes_used/total*100:0;$('capacityLabel').textContent=`${gb(s.cache_bytes_used)} / ${gb(total)} GB`;$('capacityBar').style.width=pct+'%';

scripts/agent_gan_inference_demo.py

@@ -41,6 +41,7 @@ def _infer(
     output_tokens: int,
     get_stats,
     on_token=None,
+    max_response_tokens=None,
 ):
     before = get_stats()
     started = time.perf_counter()
@@ -50,12 +51,30 @@ def _infer(
         append_done = time.perf_counter()
         first_at = None
         generated = []
-        for token in s.generate(max_tokens=output_tokens):
-            generated.append(int(token))
-            if on_token is not None:
-                on_token(generated)
-            if first_at is None:
-                first_at = time.perf_counter()
+        response_limit = int(max_response_tokens or output_tokens)
+        stop_reason = "unknown"
+        while len(generated) < response_limit:
+            before_count = len(generated)
+            chunk = min(output_tokens, response_limit - len(generated))
+            for token in s.generate(max_tokens=chunk):
+                generated.append(int(token))
+                if on_token is not None:
+                    on_token(generated)
+                if first_at is None:
+                    first_at = time.perf_counter()
+            stop_reason = {
+                1: "max_tokens",
+                2: "eos",
+                3: "cancelled",
+                4: "truncated",
+            }.get(s.last_stop_reason, "unknown")
+            if stop_reason != "max_tokens":
+                break
+            if len(generated) == before_count:
+                stop_reason = "no_progress"
+                break
+        if len(generated) >= response_limit and stop_reason == "max_tokens":
+            stop_reason = "client_safety_limit"
         done = time.perf_counter()
     after = get_stats()
     first_at = first_at or done
@@ -67,6 +86,8 @@ def _infer(
         "decode_s": done - append_done,
         "e2e_s": done - started,
         "delta": _delta(before, after),
+        "stop_reason": stop_reason,
+        "complete": stop_reason == "eos",
     }
 
 
@@ -85,6 +106,7 @@ def main() -> int:
              "reliably; larger values require more worker memory or timeout.",
     )
     parser.add_argument("--output-tokens", type=int, default=64)
+    parser.add_argument("--max-response-tokens", type=int, default=512)
     parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json")
     parser.add_argument("--skip-ensure", action="store_true")
     args = parser.parse_args()
@@ -113,7 +135,9 @@ def main() -> int:
             "role": "system",
             "content": (
                 "You are the Generator agent. Propose a technically precise "
-                "architecture improvement. Respond with actionable reasoning."
+                "architecture improvement. Respond with actionable reasoning. "
+                "For open or unsolved problems, state the accepted boundary "
+                "honestly and never fabricate a proof."
             ),
         },
         {"role": "user", "content": task},
@@ -123,7 +147,10 @@ def main() -> int:
         "content": (
             "You are the Critic/Discriminator agent. Attack the proposal, "
             "identify false assumptions and bottlenecks, score it from 0 to "
-            "10, and demand specific corrections."
+            "10, and demand specific corrections. Do not call a response "
+            "incomplete merely because it refuses to fabricate a solution to "
+            "an open problem. Claim truncation only when completion_status is "
+            "not EOS or the text is syntactically cut off."
         ),
     }]
 
@@ -139,6 +166,7 @@ def main() -> int:
                 "agents": ["generator", "critic"],
                 "rounds": args.rounds,
                 "output_tokens": args.output_tokens,
+                "max_response_tokens": args.max_response_tokens,
             },
         },
     )
@@ -164,6 +192,7 @@ def execute_agent(client, name, round_index, history):
             token_ids,
             args.output_tokens,
             get_stats,
+            max_response_tokens=args.max_response_tokens,
         )
         text = tokenizer.decode(generated, skip_special_tokens=True)
         delta = actual["delta"]
@@ -174,7 +203,7 @@ def execute_agent(client, name, round_index, history):
             "agent": name,
             "round": round_index,
             "hit_source": "primary_hot" if delta["local_hits"] else "unknown",
-            "ok": ok,
+            "ok": ok and actual["complete"],
             "warmup_prefix_tokens": warm["prefix_tokens"],
             "warmup_tokens_reused": (
                 warm["delta"]["tokens_reused"]
@@ -194,7 +223,7 @@ def execute_agent(client, name, round_index, history):
         )
         all_stages.append(stage)
         print(f"\n[{name.upper()} round {round_index}]\n{text}\n", flush=True)
-        return text
+        return text, stage
 
     try:
         with Client(args.address) as client:
@@ -208,17 +237,19 @@ def execute_agent(client, name, round_index, history):
                             + critic_feedback
                         ),
                     })
-                proposal = execute_agent(
+                proposal, generator_stage = execute_agent(
                     client, "generator", round_index, generator_history,
                 )
                 generator_history.append({"role": "assistant", "content": proposal})
                 critic_history.append({
                     "role": "user",
                     "content": (
                         f"Architecture task:\n{task}\n\nGenerator proposal:\n{proposal}"
+                        f"\n\ncompletion_status={generator_stage['stop_reason']}; "
+                        f"complete={generator_stage['complete']}"
                     ),
                 })
-                critic_feedback = execute_agent(
+                critic_feedback, _critic_stage = execute_agent(
                     client, "critic", round_index, critic_history,
                 )
                 critic_history.append({

scripts/agent_gan_repl.py

@@ -42,7 +42,7 @@ def _stage(name: str, warm: dict, actual: dict, text: str) -> dict:
         "agent": name,
         "round": 1,
         "hit_source": "primary_hot" if delta["local_hits"] else "unknown",
-        "ok": _agent_cache_gate(warm["delta"], delta),
+        "ok": _agent_cache_gate(warm["delta"], delta) and actual["complete"],
         "warmup_prefix_tokens": warm["prefix_tokens"],
         "warmup_tokens_reused": (
             warm["delta"]["tokens_reused"]
@@ -63,6 +63,7 @@ def main() -> int:
     parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key")
     parser.add_argument("--tokenizer-id", required=True)
     parser.add_argument("--output-tokens", type=int, default=64)
+    parser.add_argument("--max-response-tokens", type=int, default=512)
     parser.add_argument("--skip-ensure", action="store_true")
     args = parser.parse_args()
 
@@ -121,7 +122,10 @@ def get_stats():
                         "role": "system",
                         "content": (
                             "You are the Generator agent. Produce a concrete, "
-                            "technically rigorous answer. Internal run "
+                            "technically rigorous answer. For open or unsolved "
+                            "problems, state the accepted boundary honestly, "
+                            "provide rigorous context, and never fabricate a "
+                            "proof. Internal run "
                             f"{run_nonce}."
                         ),
                     },
@@ -149,6 +153,7 @@ def get_stats():
                     args.output_tokens,
                     get_stats,
                     on_token=generator_printer,
+                    max_response_tokens=args.max_response_tokens,
                 )
                 generator_printer.finish()
                 generator_text = tokenizer.decode(
@@ -176,15 +181,22 @@ def get_stats():
                         "content": (
                             "You are the Critic/Discriminator. Score the answer "
                             "0-10, identify false assumptions, and propose "
-                            "specific corrections. Internal run "
+                            "specific corrections. Do not penalize a correct "
+                            "statement that an open problem has no accepted "
+                            "proof. Call an answer incomplete only when its "
+                            "completion status is not EOS or its syntax is "
+                            "visibly cut off. Internal run "
                             f"{run_nonce}."
                         ),
                     },
                     {
                         "role": "user",
                         "content": (
                             f"Original task:\n{prompt}\n\n"
-                            f"Generator answer:\n{generator_text}"
+                            f"Generator answer:\n{generator_text}\n\n"
+                            "Generator completion status: "
+                            f"{generator_actual['stop_reason']}; "
+                            f"complete={generator_actual['complete']}"
                         ),
                     },
                 ]
@@ -210,6 +222,7 @@ def get_stats():
                     args.output_tokens,
                     get_stats,
                     on_token=critic_printer,
+                    max_response_tokens=args.max_response_tokens,
                 )
                 critic_printer.finish()
                 critic_text = tokenizer.decode(

tests/inference_engine/bench/test_prefill_fleet_report.py

@@ -36,10 +36,17 @@ def test_summary_aggregates_sources_and_medians():
     summary = summarize_stages([
         _stage(),
         _stage("primary_hot_hit", "primary_hot"),
-        {**_stage("allens_cold_restore", "allens_offload"), "ok": False},
+        {
+            **_stage("allens_cold_restore", "allens_offload"),
+            "ok": False,
+            "complete": False,
+            "stop_reason": "client_safety_limit",
+        },
     ])
     assert summary["stages_total"] == 3
     assert summary["stages_failed"] == 1
+    assert summary["incomplete_stages"] == 1
+    assert summary["stop_reason_counts"]["client_safety_limit"] == 1
     assert summary["hit_source_counts"]["remote_worker"] == 1
     assert summary["hit_source_counts"]["primary_hot"] == 1
     assert summary["bytes_received"] == 3000

tests/inference_engine/bridge/test_agent_gan_demo.py

@@ -1,5 +1,6 @@
 from scripts.agent_gan_inference_demo import (
     _agent_cache_gate,
+    _infer,
     _output_metadata,
 )
 
@@ -23,3 +24,66 @@ def test_agent_output_report_is_redacted():
     assert result["output_chars"] == 20
     assert len(result["output_hash"]) == 64
     assert "output" not in result
+
+
+class Session:
+    def __init__(self, chunks):
+        self.chunks = list(chunks)
+        self.last_stop_reason = None
+        self.calls = 0
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_args):
+        pass
+
+    def append(self, token_ids):
+        self.appended = list(token_ids)
+
+    def generate(self, *, max_tokens):
+        tokens, reason = self.chunks[self.calls]
+        self.calls += 1
+        assert len(tokens) <= max_tokens
+        yield from tokens
+        self.last_stop_reason = reason
+
+
+class Client:
+    def __init__(self, session):
+        self.session = session
+
+    def create_session(self, **_kwargs):
+        return self.session
+
+
+def test_infer_continues_chunks_until_eos():
+    session = Session([([1, 2], 1), ([3], 2)])
+    streamed = []
+    tokens, metrics = _infer(
+        Client(session),
+        [],
+        [9],
+        2,
+        lambda: {},
+        on_token=lambda values: streamed.append(list(values)),
+        max_response_tokens=10,
+    )
+    assert tokens == [1, 2, 3]
+    assert streamed == [[1], [1, 2], [1, 2, 3]]
+    assert metrics["stop_reason"] == "eos"
+    assert metrics["complete"] is True
+
+
+def test_infer_reports_explicit_client_safety_limit():
+    tokens, metrics = _infer(
+        Client(Session([([1, 2], 1)])),
+        [],
+        [9],
+        2,
+        lambda: {},
+        max_response_tokens=2,
+    )
+    assert tokens == [1, 2]
+    assert metrics["stop_reason"] == "client_safety_limit"
+    assert metrics["complete"] is False

tests/inference_engine/bridge/test_agent_gan_repl.py

@@ -31,6 +31,8 @@ def test_repl_stage_is_redacted_and_passes_cache_gate():
         "ttft_s": 0.2,
         "decode_s": 0.3,
         "e2e_s": 0.4,
+        "stop_reason": "eos",
+        "complete": True,
         "delta": {
             "local_hits": 1,
             "remote_jobs": 0,
@@ -43,3 +45,9 @@ def test_repl_stage_is_redacted_and_passes_cache_gate():
     assert stage["output_chars"] == 14
     assert len(stage["output_hash"]) == 64
     assert "output" not in stage
+    assert not _stage(
+        "generator",
+        warm,
+        {**actual, "complete": False, "stop_reason": "client_safety_limit"},
+        "cut off",
+    )["ok"]