Enforce lossless Prefill token budgets

autoresearch/prefill/program.md

@@ -57,6 +57,14 @@ benchmark here. Model/tokenizer/quantization/rope/window/cache-format changes
 must be deployed outside this supervisor. Cold benchmarks are explicit,
 separate invocations of `scripts/benchmark_prefill_architecture.py`.
 
+Prefill budgets are hard admission limits, never truncation instructions.
+Strategy input must fit 4096 tokens by carrying the complete active leaf
+ancestry and its exact experiment records. Generator and Critic inputs must fit
+6144 tokens; the Critic always receives the complete current Generator output.
+If any complete semantic unit exceeds its budget, reject it before remote
+Prefill and preserve the checkpoint. Never slice, sample, summarize, or drop
+the tail of an over-budget input.
+
 Do not optimize output wording, scores, prizes, or other proof-irrelevant
 content. Prefill performance is a tertiary objective after mathematical
 decomposition progress, while preserving the complete semantic contract.

autoresearch/prefill/supervisor.py

@@ -6,6 +6,7 @@
 import ast
 import csv
 import hashlib
+import io
 import json
 import os
 import re
@@ -328,7 +329,7 @@ def _extract_json(text: str) -> dict:
             if line.strip()
             and not line.strip().startswith(("#", "```"))
         ]
-        objective = " ".join(prose[:3])[:1200]
+        objective = " ".join(prose)
     if not objective:
         raise ValueError("strategy agent returned no usable candidate")
     digest = hashlib.sha256(stripped.encode()).hexdigest()[:12]
@@ -338,7 +339,7 @@ def _extract_json(text: str) -> dict:
             target_matches[-1] if target_matches else ""
         ),
         "hypothesis": objective,
-        "plan": {"steps": steps[:8]},
+        "plan": {"steps": steps},
         "strategy_parse_mode": "prose",
     }
 
@@ -417,6 +418,106 @@ def _pending_leaf_ids(ledger: dict) -> list[str]:
     return sorted(unresolved - unresolved_parents)
 
 
+def build_strategy_research_state(
+    *,
+    current: dict,
+    ledger: dict,
+    results_text: str,
+) -> dict:
+    target_id = _select_repair_target(current, ledger)
+    obligations = {
+        str(item.get("obligation_id", "")): item
+        for item in ledger.get("obligations", [])
+    }
+    ancestry = []
+    cursor = target_id
+    visited = set()
+    while cursor and cursor not in visited:
+        visited.add(cursor)
+        item = obligations[cursor]
+        ancestry.append({
+            "obligation_id": cursor,
+            "statement": item.get("statement", ""),
+            "status": item.get("status", ""),
+            "parent_id": item.get("parent_id", ""),
+            "last_run_id": item.get("last_run_id", ""),
+            "last_evidence": item.get("last_evidence", ""),
+        })
+        cursor = str(item.get("parent_id", ""))
+    ancestry.reverse()
+    ancestry_ids = {
+        item["obligation_id"] for item in ancestry
+    }
+    relevant_results = []
+    if results_text.strip():
+        for row in csv.DictReader(
+            io.StringIO(results_text),
+            delimiter="\t",
+        ):
+            if row.get("target_obligation_id") not in ancestry_ids:
+                continue
+            relevant_results.append({
+                "candidate_id": row.get("candidate_id", ""),
+                "target_obligation_id": row.get(
+                    "target_obligation_id",
+                    "",
+                ),
+                "hypothesis_sha256": row.get("hypothesis_sha256", ""),
+                "research_outcome": row.get("research_outcome", ""),
+                "research_evidence": row.get("research_evidence", ""),
+                "new_frontier": row.get("new_frontier", ""),
+                "kept": row.get("kept", ""),
+                "error": row.get("error", ""),
+            })
+    return {
+        "target_leaf_id": target_id,
+        "target_ancestry": ancestry,
+        "relevant_experiments": relevant_results,
+        "current_candidate": {
+            "candidate_id": current.get("candidate_id", ""),
+            "target_obligation_id": current.get(
+                "target_obligation_id",
+                "",
+            ),
+            "hypothesis": current.get("hypothesis", ""),
+            "prefill_compute_chunk_tokens": current.get(
+                "prefill_compute_chunk_tokens",
+            ),
+        },
+    }
+
+
+def build_strategy_prompt(
+    *,
+    program: str,
+    current: dict,
+    results_text: str,
+    ledger: dict,
+) -> str:
+    research_state = build_strategy_research_state(
+        current=current,
+        ledger=ledger,
+        results_text=results_text,
+    )
+    return (
+        "You are the AutoResearch strategy agent. Follow the human-owned "
+        "program exactly. Attack TARGET_LEAF_ID and propose "
+        "one falsifiable GAN strategy experiment. Return JSON only with keys: "
+        + ", ".join(REQUIRED_CANDIDATE_FIELDS)
+        + ". prefill_compute_chunk_tokens is immutable and must equal "
+        f"{current['prefill_compute_chunk_tokens']}. "
+        "Do not weaken full context, final-only snapshots, or no-fallback rules."
+        " The hypothesis must not repeat any hypothesis hash in RESEARCH_STATE. "
+        "It must either construct a concrete object or attempt a concrete "
+        "counterexample for the target leaf. target_obligation_id must equal "
+        "TARGET_LEAF_ID. Every statement and evidence item below is complete; "
+        "do not infer omitted text from unrelated branches."
+        f"\n\nPROGRAM:\n{program}"
+        "\n\nRESEARCH_STATE:\n"
+        f"{json.dumps(research_state, ensure_ascii=False)}"
+    )
+
+
 class StrategyPrefillHeartbeat:
     def __init__(
         self,
@@ -495,28 +596,18 @@ def propose_candidate(
     current: dict,
     results_text: str,
     ledger: dict,
+    max_prefill_tokens: int = 4096,
 ) -> dict:
     from kakeya import Client
     from transformers import AutoTokenizer
     from scripts.chat_grpc import _resolve_eos_token_ids
 
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    prompt = (
-        "You are the AutoResearch strategy agent. Follow the human-owned "
-        "program exactly. Select one unresolved proof obligation and propose "
-        "one falsifiable GAN strategy experiment. Return JSON only with keys: "
-        + ", ".join(REQUIRED_CANDIDATE_FIELDS)
-        + ". prefill_compute_chunk_tokens is immutable and must equal "
-        f"{current['prefill_compute_chunk_tokens']}. "
-        "Do not weaken full context, final-only snapshots, or no-fallback rules."
-        " The hypothesis must not repeat any hypothesis or hash in RESULTS. "
-        "It must either construct a concrete object or attempt a concrete "
-        "counterexample for one smaller decomposition leaf. The "
-        "target_obligation_id must be one of PENDING_LEAF_IDS."
-        f"\n\nPROGRAM:\n{program}\n\nCURRENT:\n{json.dumps(current)}"
-        f"\n\nRESULTS:\n{results_text[-12000:]}"
-        f"\n\nLEDGER:\n{json.dumps(ledger)}"
-        f"\n\nPENDING_LEAF_IDS:\n{json.dumps(_pending_leaf_ids(ledger))}"
+    prompt = build_strategy_prompt(
+        program=program,
+        current=current,
+        results_text=results_text,
+        ledger=ledger,
     )
     ids = tokenizer.apply_chat_template(
         [{"role": "user", "content": prompt}],
@@ -525,6 +616,11 @@ def propose_candidate(
         return_dict=False,
         enable_thinking=False,
     )
+    if len(ids) > max_prefill_tokens:
+        raise ValueError(
+            "Strategy Prefill token budget exceeded without truncation: "
+            f"{len(ids)} > {max_prefill_tokens}",
+        )
     generated: list[int] = []
     print(
         f"[autoresearch] Strategy Prefill: 0/{len(ids)} tokens (0.0%)",
@@ -820,6 +916,7 @@ def run_iteration(args, iteration: int) -> dict:
                     results_path.read_text() if results_path.exists() else ""
                 ),
                 ledger=ledger_data,
+                max_prefill_tokens=args.strategy_max_prefill_tokens,
             )
             if proposed["target_obligation_id"] not in _pending_leaf_ids(
                 ledger_data,
@@ -987,6 +1084,11 @@ def main() -> int:
     )
     parser.add_argument("--address", default="127.0.0.1:51051")
     parser.add_argument("--dashboard", default="http://127.0.0.1:8090")
+    parser.add_argument(
+        "--strategy-max-prefill-tokens",
+        type=int,
+        default=4096,
+    )
     parser.add_argument(
         "--tokenizer-id",
         default=str(
@@ -1014,6 +1116,8 @@ def main() -> int:
     args = parser.parse_args()
     if args.iterations <= 0:
         raise SystemExit("iterations must be > 0")
+    if args.strategy_max_prefill_tokens <= 0:
+        raise SystemExit("strategy-max-prefill-tokens must be > 0")
     for iteration in range(args.iterations):
         row = run_iteration(args, iteration)
         print(json.dumps(row, indent=2, sort_keys=True))

scripts/agent_gan_repl.py

@@ -670,6 +670,19 @@ def recover_checkpoint_from_log(
     )
 
 
+def enforce_prefill_token_budget(
+    stage: str,
+    token_ids,
+    max_tokens: int,
+) -> None:
+    token_count = len(token_ids)
+    if token_count > max_tokens:
+        raise ValueError(
+            f"{stage} Prefill token budget exceeded without truncation: "
+            f"{token_count} > {max_tokens}",
+        )
+
+
 def build_generator_messages(
     goal: str,
     *,
@@ -914,6 +927,12 @@ 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-prefill-tokens",
+        type=int,
+        default=6144,
+        help="Hard per-stage Prefill budget; over-budget input is rejected.",
+    )
     parser.add_argument(
         "--max-response-tokens",
         type=int,
@@ -968,6 +987,8 @@ def main() -> int:
     args = parser.parse_args()
     if args.output_tokens <= 0:
         raise SystemExit("output-tokens must be > 0")
+    if args.max_prefill_tokens <= 0:
+        raise SystemExit("max-prefill-tokens must be > 0")
     if args.auto_loop_boundary_wait_s < 0:
         raise SystemExit("auto-loop-boundary-wait-s must be >= 0")
     research_candidate = None
@@ -1245,6 +1266,11 @@ def get_stats():
                     return_dict=False,
                     enable_thinking=False,
                 )
+                enforce_prefill_token_budget(
+                    "Generator",
+                    generator_ids,
+                    args.max_prefill_tokens,
+                )
                 print(
                     f"[allens] Generator Prefill: {len(generator_ids)} tokens...",
                     flush=True,
@@ -1334,6 +1360,11 @@ def get_stats():
                     return_dict=False,
                     enable_thinking=False,
                 )
+                enforce_prefill_token_budget(
+                    "Critic",
+                    critic_ids,
+                    args.max_prefill_tokens,
+                )
                 print(
                     f"[allens] Critic Prefill: {len(critic_ids)} tokens...",
                     flush=True,

tests/inference_engine/bench/test_autoresearch_supervisor.py

@@ -1,6 +1,7 @@
 from autoresearch.prefill.supervisor import (
     append_result,
     best_kept,
+    build_strategy_research_state,
     check_runtime_health,
     parse_research_verdict,
     read_results,
@@ -281,6 +282,51 @@ def test_pending_leaf_ids_excludes_unresolved_parents():
     assert _pending_leaf_ids(ledger) == ["RH-C1-child", "RH-C2"]
 
 
+def test_strategy_state_keeps_complete_active_ancestry_only():
+    ledger = {"obligations": [
+        {
+            "obligation_id": "RH-C1",
+            "statement": "unrelated root",
+            "status": "UNRESOLVED",
+            "parent_id": "",
+            "last_evidence": "unrelated evidence",
+        },
+        {
+            "obligation_id": "RH-C2",
+            "statement": "exact root statement",
+            "status": "UNRESOLVED",
+            "parent_id": "",
+            "last_evidence": "exact root evidence",
+        },
+        {
+            "obligation_id": "RH-C2-child",
+            "statement": "exact child statement",
+            "status": "UNRESOLVED",
+            "parent_id": "RH-C2",
+            "last_evidence": "exact child evidence",
+        },
+    ]}
+    results = (
+        "candidate_id\ttarget_obligation_id\tresearch_outcome\t"
+        "research_evidence\tnew_frontier\tkept\terror\t"
+        "hypothesis_sha256\n"
+        "c1\tRH-C1\tINCONCLUSIVE\tunrelated result\tx\tFalse\t\th1\n"
+        "c2\tRH-C2-child\tDECOMPOSED\texact result\tfrontier\tTrue\t\th2\n"
+    )
+    state = build_strategy_research_state(
+        current={**_candidate(), "target_obligation_id": "RH-C2"},
+        ledger=ledger,
+        results_text=results,
+    )
+    assert state["target_leaf_id"] == "RH-C2-child"
+    serialized = str(state)
+    assert "exact root statement" in serialized
+    assert "exact child evidence" in serialized
+    assert "exact result" in serialized
+    assert "unrelated root" not in serialized
+    assert "unrelated result" not in serialized
+
+
 def test_results_are_append_only_and_best_is_selected(tmp_path):
     path = tmp_path / "results.tsv"
     common = {
@@ -390,6 +436,7 @@ def test_supervisor_preserves_runtime_and_cache_across_iterations():
     assert "clear_primary_cache" not in source
     assert "launchctl" not in source
     assert "bootout" not in source
+    assert "results_text[-" not in source
 
 
 def test_gan_subprocess_output_is_streamed_not_captured():

tests/inference_engine/bridge/test_agent_gan_repl.py

@@ -20,6 +20,7 @@
     build_critic_messages,
     build_generator_messages,
     extract_obligation_history,
+    enforce_prefill_token_budget,
     install_signal_protection,
     is_runtime_artifact_prompt,
     consume_critic_issue_batch,
@@ -324,6 +325,19 @@ def test_generator_history_is_scoped_to_target_leaf():
     assert "unrelated operator branch" not in prompt
 
 
+def test_prefill_budget_rejects_whole_input_without_truncation():
+    token_ids = list(range(7))
+    enforce_prefill_token_budget("Generator", token_ids, 7)
+    try:
+        enforce_prefill_token_budget("Critic", token_ids, 6)
+    except ValueError as exc:
+        assert "without truncation" in str(exc)
+        assert "7 > 6" in str(exc)
+    else:
+        raise AssertionError("over-budget Prefill must be rejected")
+    assert token_ids == list(range(7))
+
+
 def test_runtime_output_cannot_replace_research_goal():
     for text in (
         "critic> ### Central Claim",