Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
19 kB · 449 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450from __future__ import annotations
import hashlibimport jsonimport osimport tempfileimport timeimport unittestfrom datetime import UTC, datetimefrom pathlib import Pathfrom types import SimpleNamespacefrom typing import get_argsfrom unittest.mock import Mock, patch
import uifrom report import UI_METRICS, renderfrom schema import UI_METRIC_KEYS, Config, Metric, Observation, ProcessMemory, Report, Request, Sidefrom ui import ( cpu_delta, expand_subagents, expected_roster_inactive, fixture_spec, move_down, open_agents_view, session_id, session_name, spawn_ledger_path, tail_marker, total_cpu, ui_measure, wait_for_roster, write_fixtures,)
def small_fixture(**overrides): """Fixture generation with tiny counts so tests stay fast.""" constants = { "MEDIUM_COUNT": 3, "MEDIUM_MESSAGES": 2, "LARGE_COUNT": 2, "LARGE_MESSAGES": 3, "HUGE_MESSAGES": 4, "FANOUT_COUNT": 2, "FANOUT_MESSAGES": 2, "SUBAGENT_DEPTH": 4, "SUBAGENT_MESSAGES": 2, } constants.update(overrides) return [patch(f"ui.{name}", value) for name, value in constants.items()]
def ui_request() -> Request: """A minimal valid request for ui_measure entry checks.""" return Request( repository="PrimeIntellect-ai/prime-agent", source_repository="PrimeIntellect-ai/prime-agent", sha="a" * 40, harness_sha="a" * 40, pr=42, run_id=100, attempt=1, role="main", config=Config.load(), )
class FixtureTests(unittest.TestCase): def setUp(self): temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.root = Path(temporary.name) self.agent_dir = self.root / "agent" self.workspace = self.root / "workspace" self.workspace.mkdir() self.spec: object = None
def small_fixture(self, **overrides) -> None: for patcher in small_fixture(**overrides): patcher.start() self.addCleanup(patcher.stop) self.spec = fixture_spec()
def write(self): write_fixtures(self.agent_dir, self.workspace, self.spec) return self.agent_dir / "sessions"
def test_ids_and_names_are_deterministic(self): self.assertEqual(session_id("large", 1), session_id("large", 1)) self.assertEqual(session_id("sub", 6), fixture_spec().subagents[5]) self.assertNotEqual(session_id("large", 1), session_id("large", 2)) self.assertEqual(session_name("large", 1), "ui-bench-large-01") self.assertEqual(tail_marker("sub", 6), "ui-bench-tail-sub-06")
def test_write_fixtures_regenerates_identical_files_and_clears_stale_state(self): self.small_fixture() sessions = self.write() first = {path.name: path.read_bytes() for path in sorted(self.agent_dir.rglob("*")) if path.is_file()} stale = sessions / "stale.jsonl" stale.write_text("stale") (self.agent_dir / "session-artifacts").mkdir(parents=True) (self.agent_dir / "session-artifacts" / "leftover.json").write_text("{}") (self.agent_dir / "session-leases").mkdir(parents=True) write_fixtures(self.agent_dir, self.workspace, self.spec) second = { path.name: path.read_bytes() for path in sorted(self.agent_dir.rglob("*")) if path.is_file() } self.assertEqual(first, second) self.assertFalse(stale.exists()) self.assertFalse((self.agent_dir / "session-artifacts" / "leftover.json").exists()) expected = ui.LARGE_COUNT + ui.MEDIUM_COUNT + 1 + ui.SUBAGENT_DEPTH + ui.FANOUT_COUNT self.assertEqual(len(list(sessions.glob("*.jsonl"))), expected)
def test_session_file_shape_and_tree_linking(self): self.small_fixture(LARGE_MESSAGES=4) sessions = self.write() identifier = self.spec.large[0] entries = [json.loads(line) for line in (sessions / f"{identifier}.jsonl").read_text().splitlines()] header = entries[0] self.assertEqual(header["type"], "session") self.assertEqual(header["version"], 3) self.assertEqual(header["id"], identifier) self.assertEqual(header["cwd"], str(self.workspace)) self.assertEqual(header["rlmDepth"], 0) names = [entry for entry in entries if entry.get("type") == "session_info"] self.assertEqual(names[0]["name"], session_name("large", 1)) previous = None for entry in entries[1:]: self.assertEqual(entry["parentId"], previous) previous = entry["id"] assistants = [entry for entry in entries if entry.get("message", {}).get("role") == "assistant"] self.assertTrue(any(tail_marker("large", 1) in json.dumps(entry) for entry in assistants[-1:])) self.assertGreater(len(entries), ui.LARGE_MESSAGES)
def test_subagent_chain_depths_and_spawn_ledger(self): self.small_fixture() sessions = self.write() ledger = spawn_ledger_path(self.agent_dir, sessions) expected_digest = hashlib.sha256(str(Path(os.path.realpath(sessions))).encode()).hexdigest()[:16] self.assertEqual(ledger.name, f"{expected_digest}.jsonl") records = [json.loads(line) for line in ledger.read_text().splitlines()] self.assertEqual(records[0]["op"], "meta") self.assertEqual(records[0]["sessionsDir"], str(Path(os.path.realpath(sessions)))) spawns = records[1:] self.assertEqual(len(spawns), len(self.spec.subagents) + len(self.spec.fanout)) chain = [record for record in spawns if record["name"].startswith("ui-bench-sub-")] self.assertEqual(len(chain), len(self.spec.subagents)) fanout = [record for record in spawns if record["name"].startswith("ui-bench-fan-")] self.assertEqual(len(fanout), len(self.spec.fanout)) parent = str(sessions / f"{self.spec.root}.jsonl") for record in fanout: self.assertEqual(record["depth"], 1) self.assertEqual(record["parent"], parent) for depth, (record, identifier) in enumerate(zip(chain, self.spec.subagents, strict=True), start=1): self.assertEqual(record["op"], "spawn") self.assertEqual(record["v"], 1) self.assertEqual(record["depth"], depth) self.assertEqual(record["parent"], parent) self.assertEqual(record["child"], str(sessions / f"{identifier}.jsonl")) self.assertEqual(record["name"], session_name("sub", depth)) parent = record["child"] header = json.loads((sessions / f"{identifier}.jsonl").read_text().splitlines()[0]) self.assertEqual(header["rlmDepth"], depth)
class ProbeLogicTests(unittest.TestCase): def test_total_cpu_sums_process_tree(self): processes = [ ProcessMemory(pid=1, name="daemon", rss=10, cpu=1.25), ProcessMemory(pid=2, name="worker", rss=20, cpu=0.75), ProcessMemory(pid=3, name="kernel", rss=30, cpu=2.0), ] self.assertEqual(total_cpu(processes), 4.0) self.assertEqual(total_cpu([]), 0.0)
def test_open_agents_view_waits_for_fresh_output_before_clearing_search(self): terminal = FakeTerminal( [ "Ctrl+N new\n> old-query", "session still rendering", "Ctrl+N new\n> old-query", "Search sessions", ] ) open_agents_view(terminal, clear=True) self.assertEqual(terminal.sent, [ui.LEFT_ARROW, ui.BACKSPACE * 24]) self.assertEqual(terminal.display.text(), "Search sessions")
def test_wait_for_roster_requires_a_settled_count(self): frames = [ "agents 0 running, 1 idle, 5 inactive", "agents 0 running, 1 idle, 40 inactive", "agents 0 running, 1 idle, 70 inactive", "agents 0 running, 1 idle, 70 inactive", "agents 0 running, 1 idle, 70 inactive", "agents 0 running, 1 idle, 70 inactive", ] terminal = FakeTerminal(frames) with patch("time.perf_counter", side_effect=[float(i) for i in range(200)]): seconds = wait_for_roster(terminal, minimum=70, settle=3.0, timeout=90) self.assertGreaterEqual(seconds, 0.0) self.assertIn("70 inactive", terminal.display.text())
with self.assertRaises(TimeoutError): wait_for_roster( FakeTerminal(["agents 0 running, 1 idle, 1 inactive"] * 8), minimum=70, settle=3.0, timeout=0.05, )
def test_wait_for_roster_never_settles_below_the_expected_fixture_count(self): # A splash stuck at an empty or partial roster must not count as hydrated, # even when its count is unchanged for the whole settle window. terminal = FakeTerminal(["agents 0 running, 1 idle, 5 inactive"] * 8) with patch("time.perf_counter", side_effect=[float(i) for i in range(400)]): with self.assertRaises(TimeoutError): wait_for_roster(terminal, minimum=70, settle=3.0, timeout=90)
def test_expected_roster_inactive_counts_top_level_fixtures_minus_the_live_target(self): self.assertEqual(expected_roster_inactive(), ui.LARGE_COUNT + ui.MEDIUM_COUNT)
def test_ui_trial_requires_a_successful_first_installation(self): side = Side(sha="a" * 40) with self.assertRaisesRegex(RuntimeError, "first installation"): ui_measure(ui_request(), side, 0, results=Path("/unused"), homes=Path("/unused"), user="bench")
def test_move_down_sends_separate_keypresses(self): terminal = FakeTerminal(["initial"]) move_down(terminal, 3) self.assertEqual(terminal.sent, [ui.DOWN_ARROW] * 3)
def test_navigation_survives_persisted_agents_search(self): terminal = FakeTerminal([""]) terminal.started = time.perf_counter() terminal.ready = Mock() terminal.close = Mock() query = "" returned_queries = []
def send(data): nonlocal query terminal.sent.append(data) if data == ui.LEFT_ARROW: returned_queries.append(query) # The state-aware navigation consumes a fresh agents-view render on entry. terminal.frames.append( "agents 0 running, 1 idle, 153 inactive\n" + (query or "Search sessions") + "\nCtrl+N new" ) elif data == ui.RIGHT_ARROW or data.startswith("/resume"): terminal.display.set("chat editor") return elif data.startswith(ui.BACKSPACE): query = query[: -len(data)] elif not data.startswith("\x1b"): query += data terminal.display.set( "agents 0 running, 1 idle, 153 inactive\n" + (query or "Search sessions") + "\nui-bench-large-01\nui-bench-root-00" )
def stats(_uid): # End after navigation so this test cannot launch subsequent benchmark workloads. if "subagent_open_cpu" in side.metrics: raise RuntimeError("navigation test complete") return [ProcessMemory(pid=1, name="daemon", rss=100)]
terminal.child.send = send with ( tempfile.TemporaryDirectory() as directory, patch("ui.Terminal", return_value=terminal), patch("ui.input_ready", side_effect=lambda *args: time.perf_counter()), patch("ui.expand_subagents"), patch("ui.wait_for_roster", return_value=0.1), patch("ui.write_fixtures"), patch("ui.pwd.getpwnam", return_value=SimpleNamespace(pw_uid=123)), patch("ui.process_stats", side_effect=stats), patch("worker.environment", return_value={}), patch("worker.stop_processes"), ): side = Side(sha="a" * 40, metrics={"install": [Observation(trial=0, value=1)]}) root = Path(directory) ui_measure(ui_request(), side, 0, results=root, homes=root, user="benchmark1") for metric in ("agents_reopen", "parent_open", "subagent_open"): self.assertIsNone(side.metrics[metric][0].error) self.assertEqual(side.metrics["ui_rss"][0].error, "RuntimeError: navigation test complete") spec = fixture_spec() self.assertEqual(returned_queries, ["", spec.open_id[:8], spec.open_id[:8], spec.root[:8]])
def test_expand_subagents_requires_a_newly_expanded_row(self): # An ancestor already shows "▾", so only a growing count proves the selected row expanded; # the first expand keystroke is eaten and the retry lands. frames = [ "▾ ui-bench-sub-00", "▾ ui-bench-sub-00", "▾ ui-bench-sub-00\n▾ ui-bench-sub-01", ] terminal = FakeTerminal(frames) expand_subagents(terminal) self.assertEqual(terminal.sent, [ui.EXPAND_ARROW, ui.EXPAND_ARROW])
def test_expand_subagents_times_out_when_nothing_expands(self): frames = ["▸ ui-bench-sub-00"] * 8 terminal = FakeTerminal(frames) with self.assertRaises(TimeoutError): expand_subagents(terminal) self.assertEqual(terminal.sent, [ui.EXPAND_ARROW] * 4)
class FakeDisplay: def __init__(self, text: str = ""): self._text = text
def text(self) -> str: return self._text
def set(self, text: str) -> None: self._text = text
class FakeChild: """Records sent keystrokes the way pexpect.spawn.send would receive them."""
def __init__(self, sent: list[str]): self.sent = sent
def send(self, data: str) -> None: self.sent.append(data)
class FakeTerminal: """Scripted terminal: until() and settle() advance the frames one at a time."""
def __init__(self, frames: list[str]): self.display = FakeDisplay(frames[0] if frames else "") self.frames = list(frames[1:]) self.sent: list[str] = [] self.child = FakeChild(self.sent) self.bytes = 0
def pump(self) -> bool: if not self.frames: return False self.display.set(self.frames.pop(0)) self.bytes += len(self.display.text()) return True
def until(self, predicate, timeout: float) -> None: # type: ignore[no-untyped-def] while not predicate(self.display): if not self.pump(): raise TimeoutError("Timed out waiting for the expected terminal state")
def until_output(self, predicate, timeout: float) -> None: # type: ignore[no-untyped-def] output = "" while not predicate(output): if not self.frames: raise TimeoutError("Timed out waiting for fresh terminal output") self.pump() output += self.display.text()
def settle(self, seconds: float) -> None: self.pump()
class SchemaTests(unittest.TestCase): def test_ui_metric_keys_match_report_definitions_and_metric_literal(self): self.assertEqual(UI_METRIC_KEYS, {definition.key for definition in UI_METRICS}) for definition in UI_METRICS: self.assertIn(definition.key, get_args(Metric))
def test_config_accepts_and_defaults_ui_trials(self): config = Config.load() self.assertEqual(config.ui_trials, 3) raw = json.loads(Path(__file__).resolve().parents[1].joinpath("config.json").read_text()) self.assertIn("ui_trials", raw)
def test_report_validator_uses_ui_trial_counts(self): started = datetime(2026, 9, 9, tzinfo=UTC) config = Config.load() report = Report( repository="PrimeIntellect-ai/prime-agent", head_repository="PrimeIntellect-ai/prime-agent", pr=42, run_id=100, harness_sha="a" * 40, base_sha="a" * 40, head_sha="b" * 40, started_at=started, config=config, main=Side(sha="a" * 40), pr_head=Side(sha="b" * 40), ) for side in (report.main, report.pr_head): for definition in UI_METRICS: side.metrics[definition.key] = [ Observation(trial=trial, value=1.0) for trial in range(config.ui_trials) ] report.model_validate(report.model_dump())
overflowing = report.pr_head.model_copy(deep=True) overflowing.metrics["ui_rss"].append(Observation(trial=config.ui_trials, value=1.0)) with self.assertRaises(ValueError): Report.model_validate({**report.model_dump(), "pr_head": overflowing.model_dump()})
def test_render_includes_ui_section(self): started = datetime(2026, 9, 9, tzinfo=UTC) config = Config.load() report = Report( repository="PrimeIntellect-ai/prime-agent", head_repository="PrimeIntellect-ai/prime-agent", pr=42, run_id=100, harness_sha="a" * 40, base_sha="a" * 40, head_sha="b" * 40, started_at=started, config=config, main=Side(sha="a" * 40), pr_head=Side(sha="b" * 40), ) for definition in UI_METRICS: report.main.metrics[definition.key] = [ Observation(trial=trial, value=2.0) for trial in range(config.ui_trials) ] report.pr_head.metrics[definition.key] = [ Observation(trial=trial, value=2.0) for trial in range(config.ui_trials) ] report.status = "completed" text = render(report) self.assertIn("**UI interactions**", text) self.assertIn("Resume large session (cold)", text) self.assertIn("Open subagent session at depth 6", text) self.assertIn("UI memory after interactions", text)
if __name__ == "__main__": unittest.main()
class CpuDeltaTests(unittest.TestCase): def test_reports_cpu_consumed_since_start(self): self.assertAlmostEqual(cpu_delta(1.0, 1.5), 0.5)
def test_clamps_negative_deltas_from_exited_processes_to_zero(self): self.assertEqual(cpu_delta(5.0, 4.9996), 0.0) Observation(trial=0, value=cpu_delta(5.0, 4.9996))