From 65bc8ca094d19e169990d716d5ad0ec270c66436 Mon Sep 17 00:00:00 2001
From: Stevengre <zhaojianhong96@gmail.com>
Date: Mon, 22 Jun 2026 15:33:22 +0800
Subject: [PATCH] feat: add `--step-timeout` option for kontrol prove

Exposes pyk's per-step timeout through `kontrol prove`. When
`--step-timeout S` is set, each rewrite step is given a wall-clock budget
of `S` whole seconds; on timeout the backend request is interrupted, the
execution depth is halved, and the step is retried (enforced by
`APRProver`, threaded via `run_prover`).

Default `None` disables the timeout, leaving prior behavior unchanged.
---
 src/kontrol/cli.py     | 7 +++++++
 src/kontrol/options.py | 2 ++
 src/kontrol/prove.py   | 1 +
 3 files changed, 10 insertions(+)

diff --git a/src/kontrol/cli.py b/src/kontrol/cli.py
index 55739f513..9b4f17a4b 100644
--- a/src/kontrol/cli.py
+++ b/src/kontrol/cli.py
@@ -576,6 +576,13 @@ def parse(s: str) -> list[T]:
         action='store_true',
         help='Generate a Solidity test contract with concrete counterexample values when proofs fail.',
     )
+    prove_args.add_argument(
+        '--step-timeout',
+        type=int,
+        default=None,
+        dest='step_timeout',
+        help='Per-step wall-clock budget in whole seconds; on timeout the backend request is interrupted and the execution depth is halved before retrying. Disabled by default.',
+    )
 
     show_args = command_parser.add_parser(
         'show',
diff --git a/src/kontrol/options.py b/src/kontrol/options.py
index 046326bc8..6b58edeae 100644
--- a/src/kontrol/options.py
+++ b/src/kontrol/options.py
@@ -350,6 +350,7 @@ class ProveOptions(
     extra_module: str | None
     symbolic_caller: bool
     generate_counterexample: bool
+    step_timeout: int | None
 
     def __init__(self, args: dict[str, Any]) -> None:
         super().__init__(args)
@@ -383,6 +384,7 @@ def default() -> dict[str, Any]:
             'extra_module': None,
             'symbolic_caller': False,
             'generate_counterexample': False,
+            'step_timeout': None,
         }
 
     @staticmethod
diff --git a/src/kontrol/prove.py b/src/kontrol/prove.py
index 22b293176..6b682eb4d 100644
--- a/src/kontrol/prove.py
+++ b/src/kontrol/prove.py
@@ -461,6 +461,7 @@ def create_kcfg_explore() -> KCFGExplore:
                 assume_defined=options.assume_defined,
                 extra_module=lemmas_module,
                 optimize_kcfg=options.optimize_kcfg,
+                step_timeout=options.step_timeout,
             )
 
             if progress is not None and task is not None: