Add 3-hour agent timeout, retry policy, and spend cap for enhanced experiment robustness

- Extend `AGENT_TIMEOUT_SEC` to 3 hours for long-running tasks, with per-run spend capped at $30 to maintain budget control.
- Define a default retry policy (`max_retries=2`) to handle transient API errors while excluding genuine timeouts.
- Update `run_agent.sh` with Anthropic API key and base URL support.
- Ensure user-supplied configurations override defaults for retry and spend cap.
- Expand unit tests to validate new retry logic, cost cap enforcement, and timeout adjustments.

Author: Frame17

scripts/run_agent.sh

@@ -8,6 +8,8 @@ set -a
 set +a
 export OPENAI_API_KEY="$LITELLM_API_KEY"
 export OPENAI_BASE_URL="$LITELLM_BASE_URL"
+export ANTHROPIC_API_KEY="$LITELLM_API_KEY"
+export ANTHROPIC_BASE_URL="$LITELLM_BASE_URL/anthropic"
 
 LITELLM_BASE_HOST="${LITELLM_BASE_URL#*://}"
 LITELLM_BASE_HOST="${LITELLM_BASE_HOST%%/*}"

src/kotlinai/experiment.py

@@ -44,6 +44,22 @@
 # in one job that is a confound, not just noise.
 AGENT_VERSION_PACKAGES = {"codex": "@openai/codex"}
 
+# Default retry policy for the generated Harbor job. Transient agent API errors
+# (e.g. UnknownApiError, "API Error: Failed to parse JSON") are one-off proxy
+# hiccups that a clean re-run of the trial almost always clears, so retrying
+# recovers otherwise-lost trials. We rely on Harbor's default
+# ``RetryConfig.exclude_exceptions`` (which already contains AgentTimeoutError,
+# VerifierTimeoutError, ApiUsageLimitError, ...) so genuine timeouts/limits stay
+# non-retried. A retry re-runs the whole trial from scratch, so it is capped low.
+DEFAULT_RETRY = {"max_retries": 2}
+
+# Default per-run model spend cap for claude-code (``--max-budget-usd``). The
+# agent timeout is now 3 hours (see ``kotlinai.harbor.AGENT_TIMEOUT_SEC``); this
+# bounds spend under that larger time budget. A retry re-runs the whole trial,
+# so this cap is per agent run, not per trial. Passed as a string, matching how
+# claude-code's CLI flag is resolved from kwargs.
+DEFAULT_AGENT_MAX_BUDGET_USD = "30"
+
 # Two-sided 95% Student's t critical values indexed by degrees of freedom (entry
 # 0 is unused). Beyond the table the value has converged on the normal quantile.
 # fmt: off
@@ -77,6 +93,21 @@ def pin_agent_version(agent: dict[str, Any], resolver: Callable[[str], str]) ->
         kwargs["version"] = resolver(package)
 
 
+def apply_agent_cost_limit(agent: dict[str, Any]) -> None:
+    """Cap claude-code's per-run model spend unless the caller set a limit.
+
+    Harbor passes ``kwargs.max_budget_usd`` to claude-code as ``--max-budget-usd``.
+    Now that the agent timeout is 3 hours (see ``AGENT_TIMEOUT_SEC``), this keeps
+    a single agent run's spend bounded. The flag is claude-code specific, so this
+    is a no-op for other agents, and an explicit ``agent.kwargs.max_budget_usd``
+    in the config still wins.
+    """
+    if agent["name"].strip().lower() != "claude-code":
+        return
+    kwargs = agent.setdefault("kwargs", {})
+    kwargs.setdefault("max_budget_usd", DEFAULT_AGENT_MAX_BUDGET_USD)
+
+
 @dataclass(frozen=True)
 class ComparisonSettings:
     languages: tuple[str, str]
@@ -162,6 +193,8 @@ def validate_comparison_config(config: dict[str, Any]) -> ComparisonSettings:
     conflicting_fields = sorted(set(harbor) & RESERVED_HARBOR_FIELDS)
     if conflicting_fields:
         raise ComparisonError(f"harbor contains generated field(s): {conflicting_fields}")
+    if "retry" in harbor and not isinstance(harbor["retry"], dict):
+        raise ComparisonError("harbor.retry must be an object")
     return ComparisonSettings(
         languages=languages,
         parallelism=parallelism,
@@ -194,6 +227,9 @@ def build_comparison_job_config(
     an IndexError partway through. Datasets keep that registry populated.
     """
     config = copy.deepcopy(settings.harbor)
+    # A default retry only when the user has not set one, so an explicit
+    # ``harbor.retry`` in the config still fully overrides it.
+    config.setdefault("retry", copy.deepcopy(DEFAULT_RETRY))
     config["job_name"] = job_name
     config["jobs_dir"] = str(jobs_dir.resolve())
     config["n_concurrent_trials"] = settings.parallelism
@@ -253,6 +289,8 @@ def run_comparison_experiment(
     # CLI is the confound this run exists to avoid. Pin explicitly in the config
     # to run without registry access.
     pin_agent_version(settings.agent, version_resolver)
+    # Bound claude-code's per-run spend now that the agent timeout is 3 hours.
+    apply_agent_cost_limit(settings.agent)
     output_dir.mkdir(parents=True)
     jobs_dir = output_dir / "jobs"
     repo_root = runner.parent.parent

src/kotlinai/harbor.py

@@ -38,7 +38,10 @@
 COMPILE_TIMEOUT_SEC = 900.0
 BRANCH_TIMEOUT_SEC = 3600.0
 BUILD_TIMEOUT_SEC = 1800.0
-AGENT_TIMEOUT_SEC = 3600.0
+# 3 hours. These reverse-engineering tasks need more than the earlier 1-hour cap
+# to finish; a per-run cost limit (see experiment.py) keeps the larger budget
+# bounded. AgentTimeoutError stays in Harbor's exclude_exceptions (not retried).
+AGENT_TIMEOUT_SEC = 10800.0
 KOTLIN_VERSION = "2.4.10"
 KOTLIN_COMPILER_SHA256 = "473dd66c7a3ef4b182065b3da670466c1bf2773a9dbb0ed8b33a39fe9d4f876d"
 # The agent phase reaches the model API and nothing else: fetching upstream

tests/test_experiment.py

@@ -13,7 +13,9 @@
 import yaml
 
 from kotlinai.experiment import (
+    DEFAULT_AGENT_MAX_BUDGET_USD,
     ComparisonError,
+    apply_agent_cost_limit,
     build_comparison_job_config,
     build_comparison_report,
     render_markdown_report,
@@ -65,6 +67,67 @@ def test_comparison_job_config_contains_both_languages_and_shared_parameters(tmp
     ]
 
 
+def test_comparison_job_config_injects_default_retry_when_unset(tmp_path):
+    base = _base_config()
+    base["harbor"].pop("retry")
+    settings = validate_comparison_config(base)
+    config = build_comparison_job_config(
+        settings,
+        task_dirs={
+            "java": tmp_path / "tasks" / "java",
+            "kotlin": tmp_path / "tasks" / "kotlin",
+        },
+        jobs_dir=tmp_path / "jobs",
+        job_name="comparison",
+    )
+
+    # Transient agent API errors are retried; genuine timeouts/limits stay
+    # non-retried via Harbor's default exclude_exceptions.
+    assert config["retry"] == {"max_retries": 2}
+
+
+def test_comparison_job_config_keeps_user_supplied_retry(tmp_path):
+    base = _base_config()
+    base["harbor"]["retry"] = {"max_retries": 5}
+    settings = validate_comparison_config(base)
+    config = build_comparison_job_config(
+        settings,
+        task_dirs={
+            "java": tmp_path / "tasks" / "java",
+            "kotlin": tmp_path / "tasks" / "kotlin",
+        },
+        jobs_dir=tmp_path / "jobs",
+        job_name="comparison",
+    )
+
+    assert config["retry"] == {"max_retries": 5}
+
+
+def test_validate_comparison_config_rejects_non_object_retry():
+    config = _base_config()
+    config["harbor"]["retry"] = 3
+    with pytest.raises(ComparisonError, match="harbor.retry"):
+        validate_comparison_config(config)
+
+
+def test_apply_agent_cost_limit_defaults_claude_code_budget():
+    agent = {"name": "claude-code", "model_name": "claude-sonnet-5"}
+    apply_agent_cost_limit(agent)
+    assert agent["kwargs"]["max_budget_usd"] == DEFAULT_AGENT_MAX_BUDGET_USD
+
+
+def test_apply_agent_cost_limit_keeps_user_supplied_budget():
+    agent = {"name": "claude-code", "model_name": "claude-sonnet-5", "kwargs": {"max_budget_usd": "50"}}
+    apply_agent_cost_limit(agent)
+    assert agent["kwargs"]["max_budget_usd"] == "50"
+
+
+def test_apply_agent_cost_limit_is_noop_for_other_agents():
+    agent = {"name": "codex", "model_name": "openai/test-model"}
+    apply_agent_cost_limit(agent)
+    assert "max_budget_usd" not in agent.get("kwargs", {})
+
+
 @pytest.mark.parametrize("languages", [[], ["Java"], ["Java", "Kotlin", "Go"], ["Java", "java"], ["", "Kotlin"]])
 def test_validate_comparison_config_requires_two_distinct_languages(languages):
     config = _base_config()

tests/test_harbor.py

@@ -156,6 +156,16 @@ def test_convert_instance_task_toml_enforces_programbench_policy(tmp_path):
     assert "rm -f ./executable" in (harbor.HARBOR_DATA / "test.sh").read_text()
 
 
+def test_convert_instance_bakes_three_hour_agent_timeout(tmp_path):
+    out = convert_instance(_calc_instance(), tmp_path, target_language="Kotlin")
+    cfg = tomllib.loads((out / "task.toml").read_text())
+
+    # The agent gets a 3-hour budget so long reverse-engineering tasks finish;
+    # a per-run cost cap (experiment.py) keeps spend bounded under that budget.
+    assert harbor.AGENT_TIMEOUT_SEC == 10800.0
+    assert cfg["agent"]["timeout_sec"] == harbor.AGENT_TIMEOUT_SEC
+
+
 def test_agent_image_ships_inspection_tools_and_drops_sudo(tmp_path):
     out = convert_instance(_calc_instance(), tmp_path, target_language="Kotlin")
     dockerfile = (out / "environment" / "Dockerfile").read_text()