fix(autoresearch): warm and retry Lean typechecks

Prewarm a reduced mathlib environment, classify 45/120-second typecheck retries, kill timed-out process groups, and stop whitespace-only model decode loops.

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

Author: 2 people

KakeyaLeanGate/Prelude.lean

@@ -1,4 +1,8 @@
-import Mathlib
+import Mathlib.Analysis.Complex.Basic
+import Mathlib.Analysis.Complex.Hadamard
+import Mathlib.Analysis.Complex.JensenFormula
+import Mathlib.Analysis.Complex.LocallyUniformLimit
+import Mathlib.Analysis.Complex.Order
 
 /-!
 Minimal import target for AutoResearch theorem-signature validation.

autoresearch/prefill/lean_gate.py

@@ -2,9 +2,12 @@
 from __future__ import annotations
 
 import hashlib
+import os
 import re
+import signal
 import subprocess
 import tempfile
+import time
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -28,7 +31,19 @@ class LeanSignatureResult:
     source: str
     signature_hash: str
     ok: bool
+    status: str = "FORMALIZED"
     error: str = ""
+    attempts: int = 1
+    elapsed_s: float = 0.0
+    output: str = ""
+
+
+@dataclass(frozen=True)
+class _LeanRun:
+    returncode: int | None
+    timed_out: bool
+    elapsed_s: float
+    output: str
 
 
 def extract_lean_signature_blocks(text: str) -> list[tuple[str, str]]:
@@ -43,38 +58,159 @@ def _signature_only(source: str) -> str:
     return source[:match.start()].strip() if match else source.strip()
 
 
+def _run_lean(
+    content: str,
+    *,
+    project_root: Path,
+    timeout_s: float,
+) -> _LeanRun:
+    started = time.monotonic()
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".lean",
+        encoding="utf-8",
+        delete=False,
+    ) as handle:
+        handle.write(content)
+        path = Path(handle.name)
+    process = None
+    try:
+        process = subprocess.Popen(
+            ["lake", "env", "lean", str(path)],
+            cwd=project_root,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            start_new_session=True,
+        )
+        try:
+            output, _ = process.communicate(timeout=timeout_s)
+            return _LeanRun(
+                process.returncode,
+                False,
+                time.monotonic() - started,
+                output or "",
+            )
+        except subprocess.TimeoutExpired as exc:
+            partial = (
+                exc.stdout.decode(errors="replace")
+                if isinstance(exc.stdout, bytes)
+                else (exc.stdout or "")
+            )
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except (OSError, ProcessLookupError):
+                process.kill()
+            remainder, _ = process.communicate()
+            return _LeanRun(
+                None,
+                True,
+                time.monotonic() - started,
+                partial + (remainder or ""),
+            )
+    except OSError as exc:
+        return _LeanRun(
+            None,
+            False,
+            time.monotonic() - started,
+            f"{type(exc).__name__}: {exc}",
+        )
+    finally:
+        if process is not None and process.poll() is None:
+            process.kill()
+            process.wait()
+        path.unlink(missing_ok=True)
+
+
+def warm_lean_environment(
+    project_root: Path,
+    *,
+    timeout_s: float = 120.0,
+) -> LeanSignatureResult:
+    source = "theorem kakeyaLeanWarmup : True := by trivial"
+    content = (
+        "import KakeyaLeanGate\n\n"
+        "set_option autoImplicit false\n\n"
+        + source
+        + "\n"
+    )
+    run = _run_lean(
+        content,
+        project_root=project_root,
+        timeout_s=timeout_s,
+    )
+    if run.timed_out:
+        return LeanSignatureResult(
+            source,
+            "",
+            False,
+            status="TYPECHECK_TIMEOUT",
+            error=f"Lean warmup timed out after {timeout_s:.1f}s",
+            elapsed_s=run.elapsed_s,
+            output=run.output,
+        )
+    if run.returncode != 0:
+        return LeanSignatureResult(
+            source,
+            "",
+            False,
+            status="ENVIRONMENT_FAILED",
+            error=f"Lean warmup failed: {run.output[-2000:]}",
+            elapsed_s=run.elapsed_s,
+            output=run.output,
+        )
+    return LeanSignatureResult(
+        source,
+        "",
+        True,
+        status="ENVIRONMENT_READY",
+        elapsed_s=run.elapsed_s,
+        output=run.output,
+    )
+
+
 def validate_lean_signature(
     source: str,
     *,
     project_root: Path,
-    timeout_s: float = 30.0,
+    timeout_s: float = 45.0,
+    retry_timeout_s: float = 120.0,
 ) -> LeanSignatureResult:
     source = source.strip()
     if not source:
-        return LeanSignatureResult("", "", False, "empty Lean signature")
+        return LeanSignatureResult(
+            "", "", False, status="TYPECHECK_FAILED",
+            error="empty Lean signature",
+        )
     if len(source) > 12_000:
-        return LeanSignatureResult("", "", False, "Lean signature too large")
+        return LeanSignatureResult(
+            "", "", False, status="TYPECHECK_FAILED",
+            error="Lean signature too large",
+        )
     if _FORBIDDEN.search(source):
         return LeanSignatureResult(
             source,
             "",
             False,
-            "forbidden Lean command in generated signature",
+            status="UNSAFE_REJECTED",
+            error="forbidden Lean command in generated signature",
         )
     declarations = re.findall(r"^\s*theorem\s+([A-Za-z_][\w']*)", source, re.MULTILINE)
     if len(declarations) != 1:
         return LeanSignatureResult(
             source,
             "",
             False,
-            "expected exactly one theorem declaration",
+            status="TYPECHECK_FAILED",
+            error="expected exactly one theorem declaration",
         )
     if not re.search(r"\s*:=\s*by\b", source):
         return LeanSignatureResult(
             source,
             "",
             False,
-            "theorem signature must end with `:= by` proof scaffold",
+            status="TYPECHECK_FAILED",
+            error="theorem signature must end with `:= by` proof scaffold",
         )
     signature = " ".join(_signature_only(source).split())
     signature_hash = hashlib.sha256(signature.encode()).hexdigest()
@@ -84,39 +220,72 @@ def validate_lean_signature(
         + source
         + "\n"
     )
-    try:
-        with tempfile.NamedTemporaryFile(
-            mode="w",
-            suffix=".lean",
-            encoding="utf-8",
-            delete=False,
-        ) as handle:
-            handle.write(content)
-            path = Path(handle.name)
-        completed = subprocess.run(
-            ["lake", "env", "lean", str(path)],
-            cwd=project_root,
-            capture_output=True,
-            text=True,
-            timeout=timeout_s,
-            check=False,
+    first = _run_lean(
+        content,
+        project_root=project_root,
+        timeout_s=timeout_s,
+    )
+    attempts = 1
+    total_elapsed = first.elapsed_s
+    output = first.output
+    run = first
+    if first.timed_out:
+        warmup = warm_lean_environment(
+            project_root,
+            timeout_s=retry_timeout_s,
         )
-    except (OSError, subprocess.TimeoutExpired) as exc:
+        total_elapsed += warmup.elapsed_s
+        output += warmup.output
+        if not warmup.ok:
+            return LeanSignatureResult(
+                source,
+                signature_hash,
+                False,
+                status=warmup.status,
+                error=warmup.error,
+                attempts=1,
+                elapsed_s=total_elapsed,
+                output=output,
+            )
+        run = _run_lean(
+            content,
+            project_root=project_root,
+            timeout_s=retry_timeout_s,
+        )
+        attempts = 2
+        total_elapsed += run.elapsed_s
+        output += run.output
+    if run.timed_out:
         return LeanSignatureResult(
             source,
             signature_hash,
             False,
-            f"Lean invocation failed: {type(exc).__name__}: {exc}",
+            status="TYPECHECK_TIMEOUT",
+            error=(
+                f"Lean typecheck timed out after {attempts} attempts "
+                f"({timeout_s:.1f}s/{retry_timeout_s:.1f}s)"
+            ),
+            attempts=attempts,
+            elapsed_s=total_elapsed,
+            output=output,
         )
-    finally:
-        if "path" in locals():
-            path.unlink(missing_ok=True)
-    if completed.returncode != 0:
-        error = (completed.stderr or completed.stdout).strip()
+    if run.returncode != 0:
         return LeanSignatureResult(
             source,
             signature_hash,
             False,
-            f"Lean typecheck failed: {error[-2000:]}",
+            status="TYPECHECK_FAILED",
+            error=f"Lean typecheck failed: {run.output[-2000:]}",
+            attempts=attempts,
+            elapsed_s=total_elapsed,
+            output=output,
         )
-    return LeanSignatureResult(source, signature_hash, True)
+    return LeanSignatureResult(
+        source,
+        signature_hash,
+        True,
+        status="FORMALIZED",
+        attempts=attempts,
+        elapsed_s=total_elapsed,
+        output=output,
+    )

autoresearch/prefill/program.md

@@ -92,6 +92,14 @@ against pinned Lean/mathlib before persistence and records `FORMALIZED` plus a
 signature hash. Missing, unsafe, ill-typed, or duplicate signatures reject the
 child. `FORMALIZED` is not `PROVED`: closure still requires a separate proof
 with no `sorry` and no added axioms.
+The supervisor prewarms Lean. Signature checks use a 45-second first attempt;
+on timeout the entire Lean process group is killed, the environment is warmed
+again, and one 120-second retry is allowed. Distinguish `TYPECHECK_FAILED`,
+`TYPECHECK_TIMEOUT`, `UNSAFE_REJECTED`, and `ENVIRONMENT_FAILED`.
+
+Generator/Critic decode must also make semantic progress. Three consecutive
+chunks containing only whitespace or empty decoded text terminate the turn as
+`semantic_stall`; never accept an unterminated partial Lean block.
 
 Do not optimize output wording, scores, prizes, or other proof-irrelevant
 content. Prefill performance is a tertiary objective after mathematical

autoresearch/prefill/supervisor.py

@@ -19,6 +19,7 @@
 from pathlib import Path
 
 from autoresearch.prefill.prepare import _load_candidate, evaluate
+from autoresearch.prefill.lean_gate import warm_lean_environment
 
 
 REQUIRED_CANDIDATE_FIELDS = (
@@ -1315,6 +1316,18 @@ def main() -> int:
         raise SystemExit("strategy-max-prefill-tokens must be > 0")
     if args.strategy_stagnation_rounds <= 0:
         raise SystemExit("strategy-stagnation-rounds must be > 0")
+    lean_warmup = warm_lean_environment(
+        Path(__file__).resolve().parents[2],
+    )
+    print(
+        "[autoresearch] phase=lean-warmup "
+        f"status={lean_warmup.status} "
+        f"elapsed_s={lean_warmup.elapsed_s:.2f} "
+        f"error={lean_warmup.error or '(none)'}",
+        flush=True,
+    )
+    if not lean_warmup.ok:
+        raise SystemExit(lean_warmup.error)
     for iteration in range(args.iterations):
         row = run_iteration(args, iteration)
         print(json.dumps(row, indent=2, sort_keys=True))

scripts/agent_gan_inference_demo.py

@@ -64,7 +64,11 @@ def _infer(
     get_stats,
     on_token=None,
     max_response_tokens=None,
+    semantic_progress=None,
+    max_semantic_stall_chunks: int = 3,
 ):
+    if max_semantic_stall_chunks <= 0:
+        raise ValueError("max_semantic_stall_chunks must be > 0")
     before = get_stats()
     started = time.perf_counter()
     with client.create_session(eos_token_ids=eos_ids, client_label="agent-gan") as s:
@@ -79,6 +83,7 @@ def _infer(
             else int(max_response_tokens) or None
         )
         stop_reason = "unknown"
+        stalled_chunks = 0
         while response_limit is None or len(generated) < response_limit:
             before_count = len(generated)
             chunk = (
@@ -92,12 +97,21 @@ def _infer(
                     on_token(generated)
                 if first_at is None:
                     first_at = time.perf_counter()
+            new_tokens = generated[before_count:]
+            if semantic_progress is not None and new_tokens:
+                if semantic_progress(new_tokens):
+                    stalled_chunks = 0
+                else:
+                    stalled_chunks += 1
             stop_reason = {
                 1: "max_tokens",
                 2: "eos",
                 3: "cancelled",
                 4: "truncated",
             }.get(s.last_stop_reason, "unknown")
+            if stalled_chunks >= max_semantic_stall_chunks:
+                stop_reason = "semantic_stall"
+                break
             if stop_reason != "max_tokens":
                 break
             if len(generated) == before_count: