diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index c6cee3d..de4bb23 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -342,6 +342,15 @@ Generator completion status and must not penalize an honest statement that an open problem has no accepted proof. The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop. +The REPL input boundary is a strict state machine. Before a goal exists, raw +text sets the initial immutable goal. Afterwards only `/continue`, +`/steer `, `/new `, and `/quit` are accepted; bare text and copied +runtime output are rejected. No stdin is read while inference is running. +Completed Generator/Critic text is atomically checkpointed with mode `0600` at +`~/.kakeya/agent_gan_state.json`, allowing an exact `/continue` after restart. +`--recover-run --recover-log --auto-continue` reconstructs a +checkpoint from a complete timestamped run and resumes without using any +partial Terminal echo. Generator output is always streamed in full to Terminal and passed verbatim to the Gemma Critic. Sampling, truncation, summarization, independent chunk scores, and semantic fallback are forbidden. A global Critic score is valid only when diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index a597cf3..161f607 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -13,7 +13,9 @@ import threading import time import uuid +from dataclasses import asdict, dataclass from datetime import datetime +from enum import Enum from pathlib import Path from scripts.agent_gan_inference_demo import ( @@ -120,6 +122,135 @@ def is_runtime_artifact_prompt(text: str) -> bool: return bool(lines) and bool(_RUNTIME_ARTIFACT.match(lines[0])) +class ReplPhase(str, Enum): + WAITING_FOR_GOAL = "waiting_for_goal" + READY = "ready" + RUNNING = "running" + + +@dataclass(frozen=True) +class ReplCommand: + action: str + payload: str = "" + + +@dataclass +class ReplCheckpoint: + research_goal: str + previous_generator: str = "" + previous_critic: str = "" + last_run_id: str = "" + schema_version: int = 1 + + +def parse_repl_command(raw: str, phase: ReplPhase) -> ReplCommand: + text = raw.strip() + lower = text.lower() + if lower in {"/quit", "/exit"}: + return ReplCommand("quit") + if lower.startswith("/new"): + goal = text[4:].strip() + if not goal: + raise ValueError("usage: /new ") + return ReplCommand("new", goal) + if phase == ReplPhase.WAITING_FOR_GOAL: + if text.startswith("/"): + raise ValueError("set a goal with /new ") + if not text: + raise ValueError("research goal must be non-empty") + return ReplCommand("new", text) + if phase == ReplPhase.RUNNING: + raise ValueError("inference is running; input is disabled") + if lower == "/continue": + return ReplCommand("continue") + if lower.startswith("/steer"): + steering = text[6:].strip() + if not steering: + raise ValueError("usage: /steer ") + return ReplCommand("steer", steering) + raise ValueError( + "command rejected; use /continue, /steer , " + "/new , or /quit", + ) + + +def save_checkpoint(path: Path, checkpoint: ReplCheckpoint) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(asdict(checkpoint), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + os.chmod(temporary, 0o600) + temporary.replace(path) + + +def load_checkpoint(path: Path) -> ReplCheckpoint | None: + if not path.exists(): + return None + raw = json.loads(path.read_text(encoding="utf-8")) + checkpoint = ReplCheckpoint(**raw) + if checkpoint.schema_version != 1 or not checkpoint.research_goal: + raise ValueError("invalid Agent GAN checkpoint") + return checkpoint + + +_TIMESTAMP_PREFIX = re.compile(r"^\[[^\]]+\]\s?") + + +def recover_checkpoint_from_log( + log_path: Path, + run_id: str, +) -> ReplCheckpoint: + lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines() + goal = "" + active = False + section = "" + generator: list[str] = [] + critic: list[str] = [] + for raw_line in lines: + line = _TIMESTAMP_PREFIX.sub("", raw_line, count=1) + if line.startswith("[goal] anchored:"): + goal = line.split(":", 1)[1].strip() + if line.startswith("[goal] reset:"): + goal = line.split(":", 1)[1].strip() + if f"[inference-start]" in line and f"run={run_id}" in line: + active = True + section = "" + generator = [] + critic = [] + continue + if not active: + continue + if line.startswith("generator>"): + section = "generator" + generator.append(line.removeprefix("generator>").lstrip()) + continue + if line.startswith("[allens] Critic Prefill:"): + section = "" + continue + if line.startswith("critic>"): + section = "critic" + critic.append(line.removeprefix("critic>").lstrip()) + continue + if line.startswith("[metrics]"): + break + if section == "generator": + generator.append(line) + elif section == "critic": + critic.append(line) + generator_text = "\n".join(generator).strip() + critic_text = "\n".join(critic).strip() + if not goal or not generator_text or not critic_text: + raise ValueError(f"complete run {run_id!r} not found in transcript") + return ReplCheckpoint( + research_goal=goal, + previous_generator=generator_text, + previous_critic=critic_text, + last_run_id=run_id, + ) + + def build_generator_messages( goal: str, *, @@ -339,6 +470,17 @@ def main() -> int: default="~/.kakeya/logs/agent_gan_repl.log", help="Timestamped local transcript log.", ) + parser.add_argument( + "--state-file", + default="~/.kakeya/agent_gan_state.json", + help="Private resumable Generator/Critic checkpoint.", + ) + parser.add_argument("--recover-run", default="") + parser.add_argument( + "--recover-log", + default="~/.kakeya/logs/agent_gan_repl.log", + ) + parser.add_argument("--auto-continue", action="store_true") args = parser.parse_args() if args.output_tokens <= 0: raise SystemExit("output-tokens must be > 0") @@ -376,58 +518,75 @@ def get_stats(): telemetry_state["last_stats"] = stats return stats + state_path = Path(args.state_file).expanduser() + if args.recover_run: + recovered = recover_checkpoint_from_log( + Path(args.recover_log).expanduser(), + args.recover_run, + ) + save_checkpoint(state_path, recovered) + print( + f"[state-recovered] run={recovered.last_run_id} " + f"generator_chars={len(recovered.previous_generator)} " + f"critic_chars={len(recovered.previous_critic)}", + flush=True, + ) + checkpoint = load_checkpoint(state_path) + research_goal = checkpoint.research_goal if checkpoint else "" + previous_generator = checkpoint.previous_generator if checkpoint else "" + previous_critic = checkpoint.previous_critic if checkpoint else "" + phase = ReplPhase.READY if checkpoint else ReplPhase.WAITING_FOR_GOAL + pending_input = "/continue" if args.auto_continue and checkpoint else "" print( - "Kakeya Agent GAN REPL ready. First prompt sets the immutable goal.\n" - "Use /continue to apply Critic feedback, /new to reset, " - "and /quit to exit.\n" + "Kakeya Agent GAN REPL ready.\n" + "Commands: /new , /continue, /steer , /quit.\n" + "Raw text is accepted only as the initial goal; inference output can " + "never become implicit steering.\n" "Each turn runs allens Prefill → Primary hot Generator → " "allens Prefill → Primary hot Critic.", flush=True, ) print(f"[log] {transcript.log_path}", flush=True) - research_goal = "" - previous_generator = "" - previous_critic = "" + print(f"[state] {state_path} phase={phase.value}", flush=True) + if checkpoint: + print( + f"[state-restored] run={checkpoint.last_run_id or '(none)'} " + f"goal={hashlib.sha256(research_goal.encode()).hexdigest()}", + flush=True, + ) with Client(args.address) as client: while True: + if pending_input: + raw_input = pending_input + pending_input = "" + print(f"\nprompt> {raw_input}", flush=True) + else: + try: + raw_input = input("\nprompt> ") + except EOFError: + print("\n[bye]") + break + transcript.log_only(f"[input] {raw_input or '(empty)'}") try: - prompt = input("\nprompt> ").strip() - except EOFError: - print("\n[bye]") - break - transcript.log_only(f"[input] {prompt or '(empty)'}") - if not prompt: + command = parse_repl_command(raw_input, phase) + except ValueError as exc: + print(f"[input-rejected] {exc}", flush=True) continue - if prompt.lower() in {"/quit", "/exit"}: + if command.action == "quit": print("[bye]") break - if prompt.lower().startswith("/new"): - new_goal = prompt[4:].strip() - research_goal = new_goal + if command.action == "new": + research_goal = command.payload previous_generator = "" previous_critic = "" - if not research_goal: - print("[goal] cleared; enter a new research goal", flush=True) - continue - prompt = research_goal - print(f"[goal] reset: {research_goal}", flush=True) - elif prompt.lower() == "/continue": - if not research_goal: - print("[goal-error] no active research goal", flush=True) - continue - prompt = "" - elif is_runtime_artifact_prompt(prompt): - print( - "[input-rejected] runtime output cannot become a research " - "prompt; use /continue or /new ", - flush=True, + save_checkpoint( + state_path, + ReplCheckpoint(research_goal=research_goal), ) - continue - elif not research_goal: - research_goal = prompt - print(f"[goal] anchored: {research_goal}", flush=True) - prompt = "" - steering = prompt + phase = ReplPhase.READY + print(f"[goal] reset: {research_goal}", flush=True) + steering = command.payload if command.action == "steer" else "" + phase = ReplPhase.RUNNING run_nonce = uuid.uuid4().hex telemetry_state["degraded"] = False run = _telemetry_request( @@ -607,6 +766,16 @@ def get_stats(): f"run={run_id}", flush=True, ) + save_checkpoint( + state_path, + ReplCheckpoint( + research_goal=research_goal, + previous_generator=previous_generator, + previous_critic=previous_critic, + last_run_id=run_id, + ), + ) + phase = ReplPhase.READY except Exception as exc: if remote_run: _telemetry_request( @@ -621,6 +790,7 @@ def get_stats(): f"run={run_id} error={type(exc).__name__}: {exc}", flush=True, ) + phase = ReplPhase.READY transcript.log_only("[session-end]") return 0 diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index 93b6846..5c58b78 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -5,6 +5,8 @@ from scripts.agent_gan_repl import ( PrefillHeartbeat, + ReplCheckpoint, + ReplPhase, TimestampedTee, TokenPrinter, _gate_failure, @@ -14,6 +16,10 @@ build_generator_messages, install_signal_protection, is_runtime_artifact_prompt, + load_checkpoint, + parse_repl_command, + recover_checkpoint_from_log, + save_checkpoint, ) @@ -259,3 +265,72 @@ def test_runtime_output_cannot_replace_research_goal(): ): assert is_runtime_artifact_prompt(text) assert not is_runtime_artifact_prompt("证明黎曼猜想") + + +def test_command_state_machine_requires_explicit_ready_commands(): + assert parse_repl_command( + "证明黎曼猜想", + ReplPhase.WAITING_FOR_GOAL, + ).action == "new" + assert parse_repl_command("/continue", ReplPhase.READY).action == "continue" + steering = parse_repl_command( + "/steer analyze the explicit formula", + ReplPhase.READY, + ) + assert steering.action == "steer" + assert steering.payload == "analyze the explicit formula" + assert parse_repl_command("/new new goal", ReplPhase.READY).payload == ( + "new goal" + ) + assert parse_repl_command("/quit", ReplPhase.READY).action == "quit" + for raw, phase in ( + ("Operator output fragment", ReplPhase.READY), + ("critic> copied output", ReplPhase.READY), + ("/continue", ReplPhase.WAITING_FOR_GOAL), + ("/continue", ReplPhase.RUNNING), + ("/steer", ReplPhase.READY), + ("/new", ReplPhase.READY), + ): + try: + parse_repl_command(raw, phase) + except ValueError: + pass + else: + raise AssertionError(f"expected rejection for {raw!r} in {phase}") + + +def test_checkpoint_round_trip_is_private(tmp_path): + path = tmp_path / "state.json" + expected = ReplCheckpoint( + research_goal="prove RH", + previous_generator="full generator", + previous_critic="full critic", + last_run_id="br_ok", + ) + save_checkpoint(path, expected) + assert load_checkpoint(path) == expected + assert path.stat().st_mode & 0o777 == 0o600 + assert load_checkpoint(tmp_path / "missing.json") is None + + +def test_recover_complete_checkpoint_from_timestamped_log(tmp_path): + path = tmp_path / "agent.log" + path.write_text( + "\n".join(( + "[t] [goal] anchored: prove RH", + "[t] [inference-start] time=t run=br_good goal=hash", + "[t] generator> first generator line", + "[t] second generator line", + "[t] [allens] Critic Prefill: 100 tokens...", + "[t] critic> first critic line", + "[t] second critic line", + "[t] [metrics] run=br_good", + )), + ) + recovered = recover_checkpoint_from_log(path, "br_good") + assert recovered.research_goal == "prove RH" + assert recovered.previous_generator == ( + "first generator line\nsecond generator line" + ) + assert recovered.previous_critic == "first critic line\nsecond critic line" + assert recovered.last_run_id == "br_good"