diff --git a/scripts/benchmarks/terminal.py b/scripts/benchmarks/terminal.py index 0c92ba7244..49af0c8d42 100644 --- a/scripts/benchmarks/terminal.py +++ b/scripts/benchmarks/terminal.py @@ -58,6 +58,13 @@ def feed(self, chunk: str) -> None: self.stream.feed(self.pending) self.pending = "" + def prompt_ready(self) -> bool: + text = self.text() + return "agents/resume" in text or ( + any(line.strip() == ">" for line in text.splitlines()) + and any(line.lstrip().startswith("← manage") for line in text.splitlines()) + ) + def text(self) -> str: return "\n".join(row.rstrip() for row in self.screen.display) @@ -108,7 +115,7 @@ def settle(self, seconds: float) -> None: self.pump() def ready(self) -> float: - self.until(lambda display: "agents/resume" in display.text(), 30) + self.until(lambda display: display.prompt_ready(), 30) self.child.send("benchready") echoed = self.until(lambda display: "benchready" in display.text(), 5) self.child.send("\x7f" * len("benchready")) diff --git a/scripts/benchmarks/tests/test_benchmarks.py b/scripts/benchmarks/tests/test_benchmarks.py index 6a53629f6a..a66dd89dd1 100644 --- a/scripts/benchmarks/tests/test_benchmarks.py +++ b/scripts/benchmarks/tests/test_benchmarks.py @@ -30,7 +30,7 @@ write_json, ) from terminal import QUERIES, Display, Terminal -from worker import environment, install, measure +from worker import environment, install, measure, stop_agents SHA = "a" * 40 HEAD = "b" * 40 @@ -63,7 +63,7 @@ def test_real_pty_detects_injected_startup_delays_without_submitting(self): import os, sys, time, tty tty.setraw(0) time.sleep(float(sys.argv[1])) -os.write(1, b'agents/resume\\r\\n> ') +os.write(1, sys.argv[2].encode()) while True: byte = os.read(0, 1) if byte == b'\\x7f': @@ -80,9 +80,9 @@ def test_real_pty_detects_injected_startup_delays_without_submitting(self): root = Path(directory) path = root / "fixture.py" path.write_text(script) - for delay in (0.05, 0.6): + for delay, label in ((0.05, "agents/resume\r\n> "), (0.6, ">\r\n← manage")): terminal = Terminal( - [sys.executable, str(path), str(delay)], + [sys.executable, str(path), str(delay), label], root, os.environ.copy(), root / f"transcript-{delay}", @@ -93,6 +93,16 @@ def test_real_pty_detects_injected_startup_delays_without_submitting(self): terminal.close() self.assertGreater(measurements[1] - measurements[0], 0.25) + def test_recognizes_current_and_legacy_prompt_bars(self): + for text in ("agents/resume\r\n> ", ">\r\n← manage unknown 0"): + display = Display(lambda _reply: None) + display.feed(text) + self.assertTrue(display.prompt_ready()) + for text in ("Loading...", ">", "← manage"): + display = Display(lambda _reply: None) + display.feed(text) + self.assertFalse(display.prompt_ready()) + def test_queries_split_at_every_boundary(self): for query, reply in QUERIES.items(): for boundary in range(len(query) + 1): @@ -590,6 +600,33 @@ def test_dispatch_from_a_feature_branch_cannot_publish_or_clean_up(self): class MeasurementTests(unittest.TestCase): + def test_cleanup_accepts_a_session_that_exits_between_list_and_stop(self): + stale = subprocess.CalledProcessError( + 1, ["prime-agent", "stop", "session", "--json"], stderr="Error: Unknown active session: session\n" + ) + with patch( + "worker.run_as", + side_effect=['{"sessions": [{"activeSessionId": "session"}]}', stale, '{"sessions": []}'], + ) as run: + stop_agents(Path("/fixture")) + self.assertEqual(run.call_count, 3) + self.assertEqual(run.call_args.args[1], ["prime-agent", "list", "--json"]) + + def test_cleanup_preserves_real_stop_failures(self): + for error, remaining in ( + ("Error: Unknown active session: session\n", '{"sessions": [{"activeSessionId": "session"}]}'), + ("Error: connection closed\n", '{"sessions": []}'), + ("Error: Unknown active session: another\n", '{"sessions": []}'), + ): + with self.subTest(error=error, remaining=remaining): + failure = subprocess.CalledProcessError(1, ["prime-agent", "stop"], stderr=error) + with patch( + "worker.run_as", + side_effect=['{"sessions": [{"activeSessionId": "session"}]}', failure, remaining], + ): + with self.assertRaises(subprocess.CalledProcessError): + stop_agents(Path("/fixture")) + def test_disk_footprint_is_measured_after_interactive_first_use(self): order = [] terminal = Mock() diff --git a/scripts/benchmarks/worker.py b/scripts/benchmarks/worker.py index db6bb64b0c..83745625db 100644 --- a/scripts/benchmarks/worker.py +++ b/scripts/benchmarks/worker.py @@ -378,12 +378,20 @@ def record( def stop_agents(home: Path) -> None: listing = json.loads(run_as("benchmark1", ["prime-agent", "list", "--json"], home)) for session in listing["sessions"]: - if session.get("activeSessionId"): - run_as( - "benchmark1", - ["prime-agent", "stop", session["activeSessionId"], "--json"], - home, - ) + active_id = session.get("activeSessionId") + if not active_id: + continue + try: + run_as("benchmark1", ["prime-agent", "stop", active_id, "--json"], home) + except subprocess.CalledProcessError as error: + if ( + error.returncode != 1 + or (error.stderr or "").strip() != f"Error: Unknown active session: {active_id}" + ): + raise + current = json.loads(run_as("benchmark1", ["prime-agent", "list", "--json"], home)) + if any(item.get("activeSessionId") == active_id for item in current["sessions"]): + raise def measure(request: Request, side: Side, trial: int) -> None: