Merge pull request #51 from CodingThrust/fix-run-frictions

Fix four run frictions hit by real benchmark sessions

Author: GiggleLiu

Makefile

@@ -6,7 +6,7 @@
 #   test-unit            Run only unit tests (no real repo/pred needed)
 #   verify-calibration   Test the verifier against known fixtures (no AI needed)
 #   preflight            Validate submission.env with one tiny real call before a full run
-#   run                  Run the benchmark via Docker → out/submission.json (does NOT upload)
+#   run                  Run the benchmark via Docker → out/<stamp>/submission.json (does NOT upload)
 #   run-local            Clone/verify a repo and run a local headless CLI agent
 #
 # Model/auth configuration lives in submission.env (see submission.env.example). Local
@@ -21,6 +21,7 @@ JOBS     ?= 1
 SUBS_DIR ?= submissions
 SCORED   ?= results/scored
 ENV_FILE ?= submission.env
+STAMP    ?= $(shell date +%Y%m%d-%H%M%S)
 LOCAL_BACKEND ?= codex
 LOCAL_REPO_DIR ?=
 LOCAL_OUTPUT ?=
@@ -69,16 +70,18 @@ preflight:
 	  echo "No $(ENV_FILE) — copy submission.env.example and fill it in first"; exit 1; fi
 	docker run --rm --env-file "$(ENV_FILE)" $(IMAGE) --preflight
 
-## Run the bug-finding agent via Docker → writes ./out/submission.json.
+## Run the bug-finding agent via Docker → writes ./out/<stamp>/submission.json (the
+## whole-repo trajectory lands alongside it). Each run gets its own timestamped dir so
+## successive runs never overwrite each other (override with STAMP=... to reproduce).
 ## This RUNS the benchmark locally; it does NOT submit — submitting is a separate step
 ## (open a GitHub PR adding the file, see CONTRIBUTING.md). Config lives in submission.env
 ## (copy submission.env.example); run `make preflight` first to validate it.
 run:
 	@if [ ! -f "$(ENV_FILE)" ]; then \
 	  echo "No $(ENV_FILE) — copy submission.env.example and fill it in (then: make preflight)"; exit 1; fi
 	mkdir -p out
-	docker run --rm --env-file "$(ENV_FILE)" -v "$(PWD)/out:/out" $(IMAGE)
-	@echo "Wrote ./out/submission.json — now submit it with 'python -m benchmark.submit' (see CONTRIBUTING.md)."
+	docker run --rm --env-file "$(ENV_FILE)" -e OUTPUT=/out/$(STAMP)/submission.json -v "$(PWD)/out:/out" $(IMAGE)
+	@echo "Wrote out/$(STAMP)/submission.json (+ trajectory) — now submit it with 'python -m benchmark.submit' (see CONTRIBUTING.md)."
 
 ## Lightweight host run through an installed headless agent CLI. The checkout is cloned at
 ## PR_REF when LOCAL_REPO_DIR is absent; an existing checkout must already match exactly.

benchmark/preflight.py

@@ -6,7 +6,8 @@
 
   1. verify the `pred` binary + version,
   2. confirm the library's reduction rules are present,
-  3. make ONE tiny real model call to validate credentials / endpoint.
+  3. make ONE tiny real model call to validate credentials / endpoint,
+  4. exercise the per-call cost accounting the agent runs on every step.
 
 Exits 0 only if every check passes — so you can launch the full batch with confidence
 instead of discovering a typo'd key or wrong base URL 20 rules in. This is a user-facing
@@ -55,6 +56,7 @@ def run_checks(
 
     # 3. one real model call through the exact batch model config (validates key, endpoint,
     #    model name, model_kwargs).
+    model = response = None
     try:
         model = _build_model(model_name, api_base, max_tokens,
                              model_kwargs=model_kwargs, api_key=api_key)
@@ -71,6 +73,18 @@ def run_checks(
     except Exception as e:
         results.append(("model call", False, f"{type(e).__name__}: {e}"))
 
+    # 4. the run prices EVERY call — exercise that same accounting path. A raw _query passes
+    #    even for models missing from litellm's price map, but the agent then aborts the
+    #    whole session on the cost-calculation RuntimeError (unless ignore_errors is set).
+    if response is not None:
+        try:
+            model._calculate_cost(response)
+            results.append(("cost tracking", True, "per-call pricing works (or ignore_errors set)"))
+        except Exception as e:
+            results.append(("cost tracking", False,
+                            f"{type(e).__name__}: {e} — set MSWEA_COST_TRACKING=ignore_errors "
+                            "or use a priced model before launching"))
+
     return results
 
 

benchmark/run_mini.py

@@ -45,6 +45,32 @@ def save_trajectory(messages: list, path: Path) -> None:
                                      "content": _message_text(msg)}) + "\n")
 
 
+class TrajectoryWriter:
+    """Append-only incremental trajectory writer — live progress during a session.
+
+    Without this, the trajectory only exists after the session ends, so a healthy slow
+    run and a dead hang look identical from outside (and a previous crash's leftover
+    file masquerades as the current run's progress). flush() appends any new messages
+    after each agent step; save_trajectory() still makes the authoritative final write
+    at session end, in the identical format.
+    """
+
+    def __init__(self, path: Path):
+        self._path = path
+        self._written = 0
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.unlink(missing_ok=True)  # fresh session — never append onto a stale file
+
+    def flush(self, messages: list) -> None:
+        if len(messages) <= self._written:
+            return
+        with self._path.open("a", encoding="utf-8") as handle:
+            for msg in messages[self._written:]:
+                handle.write(json.dumps({"role": msg.get("role", ""),
+                                         "content": _message_text(msg)}) + "\n")
+        self._written = len(messages)
+
+
 def _session_usage(agent):
     usage = extract_usage(agent.messages)
     total_tokens = usage.total_tokens or extract_total_tokens(agent.messages)
@@ -57,7 +83,9 @@ def _build_model(model_name: str, api_base: str | None, max_tokens: int,
                  format_error_template: str | None = None):
     from minisweagent.models.litellm_model import LitellmModel
 
-    kwargs = dict(model_kwargs or {})
+    # A hung API call must fail fast and retry — never freeze a whole-repo session
+    # indefinitely. User-supplied MODEL_KWARGS still wins on any key.
+    kwargs = {"timeout": 300, "num_retries": 2, **dict(model_kwargs or {})}
     if max_tokens:
         kwargs["max_tokens"] = max_tokens
     if api_base:
@@ -108,11 +136,25 @@ def run_repo_session(
     agent_config, model_config, strategy = _load_agent_config(
         config_path, CONFIG_FILE, strategy)
     safe_model = safe_model_label(model_name)
+    writer = None
     if trajectory_dir is not None:
         trajectory_dir = Path(trajectory_dir).resolve()
         trajectory_dir.mkdir(parents=True, exist_ok=True)
+        writer = TrajectoryWriter(trajectory_dir / f"{safe_model}_whole-repo.jsonl")
+
+    class _StepFlushingAgent(DefaultAgent):
+        """DefaultAgent that flushes the trajectory after every step (best-effort)."""
+
+        def step(self):
+            result = super().step()
+            if writer is not None:
+                try:
+                    writer.flush(self.messages)
+                except OSError:
+                    pass  # observability must never kill the session
+            return result
 
-    agent = DefaultAgent(
+    agent = _StepFlushingAgent(
         _build_model(
             model_name, api_base, max_tokens, model_kwargs=model_kwargs, api_key=api_key,
             observation_template=model_config.get("observation_template"),

benchmark/tests/test_preflight.py

@@ -12,11 +12,12 @@
 
 class _FakeModel:
     """Mimics mini-swe-agent's LitellmModel low-level path used by preflight: _query returns
-    a litellm-shaped response with a .usage block."""
-    def __init__(self, prompt=6, completion=4, raise_exc=None):
+    a litellm-shaped response with a .usage block; _calculate_cost is the per-call pricing."""
+    def __init__(self, prompt=6, completion=4, raise_exc=None, cost_exc=None):
         self._usage = SimpleNamespace(prompt_tokens=prompt, completion_tokens=completion,
                                       prompt_tokens_details=None)
         self._raise = raise_exc
+        self._cost_exc = cost_exc
 
     def _prepare_messages_for_api(self, messages):
         return messages
@@ -27,6 +28,11 @@ def _query(self, messages):
         return SimpleNamespace(usage=self._usage,
                                choices=[SimpleNamespace(message=SimpleNamespace(content="OK"))])
 
+    def _calculate_cost(self, response):
+        if self._cost_exc is not None:
+            raise self._cost_exc
+        return {"cost": 0.001}
+
 
 def _patch(monkeypatch, *, ver="0.6.0", rules=("a", "b"), model=None, build_exc=None):
     monkeypatch.setattr(pf, "verify_pred_version", lambda *a, **k: ver)
@@ -43,7 +49,7 @@ class TestRunChecks:
     def test_all_pass(self, monkeypatch):
         _patch(monkeypatch)
         results = pf.run_checks("anthropic/x", repo_dir="/repo")
-        assert [ok for _, ok, _ in results] == [True, True, True]
+        assert [ok for _, ok, _ in results] == [True, True, True, True]
         assert pf.format_report(results) is True
 
     def test_pred_failure_reported(self, monkeypatch):
@@ -71,3 +77,17 @@ def test_model_call_tokens_shown(self, monkeypatch):
         results = pf.run_checks("anthropic/x", repo_dir="/repo")
         detail = dict((n, d) for n, _, d in results)["model call"]
         assert "2000000 tokens" in detail
+
+    def test_cost_tracking_failure_fails(self, monkeypatch):
+        """A model missing from litellm's price map passes _query but kills the real run —
+        preflight must catch that (the openai/kimi-for-coding failure mode)."""
+        _patch(monkeypatch, model=_FakeModel(cost_exc=RuntimeError("model isn't mapped yet")))
+        results = pf.run_checks("anthropic/x", repo_dir="/repo")
+        call = dict((n, (ok, d)) for n, ok, d in results)["cost tracking"]
+        assert call[0] is False and "ignore_errors" in call[1]
+        assert pf.format_report(results) is False
+
+    def test_no_cost_check_when_model_call_fails(self, monkeypatch):
+        _patch(monkeypatch, model=_FakeModel(raise_exc=RuntimeError("401")))
+        results = pf.run_checks("anthropic/x", repo_dir="/repo")
+        assert [n for n, _, _ in results] == ["pred binary", "library rules", "model call"]

benchmark/tests/test_run_mini.py

@@ -0,0 +1,44 @@
+"""
+Tests for benchmark/run_mini.py model construction.
+
+The real minisweagent is stubbed via sys.modules so these run anywhere.
+"""
+import sys
+import types
+
+import pytest
+
+from benchmark.run_mini import _build_model
+
+
+class _FakeLitellmModel:
+    def __init__(self, model_name, model_kwargs, **config):
+        self.model_name = model_name
+        self.model_kwargs = model_kwargs
+        self.config = config
+
+
+@pytest.fixture
+def _fake_litellm(monkeypatch):
+    fake = types.ModuleType("minisweagent.models.litellm_model")
+    fake.LitellmModel = _FakeLitellmModel
+    monkeypatch.setitem(sys.modules, "minisweagent.models.litellm_model", fake)
+
+
+class TestBuildModel:
+    def test_defaults_timeout_and_retries(self, _fake_litellm):
+        """A hung API call must fail fast, not freeze the whole-repo session."""
+        model = _build_model("openai/x", None, 8192)
+        assert model.model_kwargs["timeout"] == 300
+        assert model.model_kwargs["num_retries"] == 2
+
+    def test_user_kwargs_override_defaults(self, _fake_litellm):
+        model = _build_model("openai/x", None, 8192, model_kwargs={"timeout": 60})
+        assert model.model_kwargs["timeout"] == 60
+        assert model.model_kwargs["num_retries"] == 2
+
+    def test_endpoint_config_passthrough(self, _fake_litellm):
+        model = _build_model("openai/x", "https://api.example/v1", 4096, api_key="k")
+        assert model.model_kwargs["api_base"] == "https://api.example/v1"
+        assert model.model_kwargs["api_key"] == "k"
+        assert model.model_kwargs["max_tokens"] == 4096

benchmark/tests/test_trajectory.py

@@ -54,3 +54,43 @@ def test_creates_parent_dirs(self, tmp_path):
         out = tmp_path / "deep" / "nested" / "traj.jsonl"
         save_trajectory(_fake_messages(), out)
         assert out.exists()
+
+
+class TestTrajectoryWriter:
+    """Incremental (per-step) trajectory flushing — the live view of a running session."""
+
+    def test_incremental_appends(self, tmp_path):
+        from benchmark.run_mini import TrajectoryWriter
+        out = tmp_path / "traj.jsonl"
+        w = TrajectoryWriter(out)
+        msgs = _fake_messages()
+        w.flush(msgs[:2])
+        w.flush(msgs)
+        lines = out.read_text(encoding="utf-8").strip().splitlines()
+        assert len(lines) == len(msgs)
+
+    def test_no_duplicate_lines_on_repeat_flush(self, tmp_path):
+        from benchmark.run_mini import TrajectoryWriter
+        out = tmp_path / "traj.jsonl"
+        w = TrajectoryWriter(out)
+        msgs = _fake_messages()
+        w.flush(msgs)
+        w.flush(msgs)
+        assert len(out.read_text(encoding="utf-8").strip().splitlines()) == len(msgs)
+
+    def test_starts_fresh_not_appending_stale(self, tmp_path):
+        from benchmark.run_mini import TrajectoryWriter, save_trajectory
+        out = tmp_path / "traj.jsonl"
+        save_trajectory(_fake_messages(), out)  # leftover from a previous run
+        w = TrajectoryWriter(out)
+        w.flush(_fake_messages()[:1])
+        assert len(out.read_text(encoding="utf-8").strip().splitlines()) == 1
+
+    def test_matches_final_save_format(self, tmp_path):
+        from benchmark.run_mini import TrajectoryWriter, save_trajectory
+        msgs = _fake_messages()
+        incremental = tmp_path / "inc.jsonl"
+        TrajectoryWriter(incremental).flush(msgs)
+        final = tmp_path / "final.jsonl"
+        save_trajectory(msgs, final)
+        assert incremental.read_text(encoding="utf-8") == final.read_text(encoding="utf-8")