"""Traces record a run, replay it offline, and count how often the choice was right.""" import json import pytest from tiny_computer_agent import ( Action, ActionError, AgentLoop, DecisionMenu, Element, MockBrain, PlannerBrain, ) from tiny_computer_agent.trace import ( DONE, NOTHING, TRACE_VERSION, Scenario, Scoreboard, Trace, TraceStep, benchmark, benchmark_traces, load_traces, run_scenario, score, score_replay, score_trace, ) # -- fixtures -------------------------------------------------------------- def frame(*labels): return tuple(Element(label, x=i * 100, y=10, width=80, height=24) for i, label in enumerate(labels)) def topbar(*labels): return tuple(Element(label, x=i * 100, y=6, width=80, height=20) for i, label in enumerate(labels)) def button(label, x, y=270): return Element(label, x=x, y=y, width=90, height=30) SAVE = Scenario( "click 'Save' to save the document", frames=[topbar("Files", "Edit", "View"), (Element("Name", x=200, y=195, width=200, height=26), button("Save", 200)), topbar("Files", "Edit", "View") + (button("Save", 200),)], expected="Save", name="save", ) CANCEL = Scenario( "click 'Cancel' to discard the change", frames=[(button("Save", 200), button("Cancel", 400))], expected="Cancel", name="cancel", ) DENIED = Scenario( "click 'Terminal' to open a shell", frames=[topbar("Terminal", "Files")], expected="Terminal", name="terminal (denied by policy)", ) class FirstItemBrain: """A baseline: always click the first item in reading order. No model, no idea.""" name = "first-item" def decide(self, goal, observation): menu = DecisionMenu(observation) if not menu: return Action("done", reason="nothing to click") return Action("click", label=menu.items[0].label, reason="first item in reading order") def oracle(scenario): return MockBrain([ {"kind": "click", "label": scenario.expected}, {"kind": "done", "reason": "clicked the right thing"}, ]) # -- recording ------------------------------------------------------------- def test_a_trace_records_the_frames_and_the_decision_without_pixels(): episode = AgentLoop(screen=SAVE.screen(), brain=PlannerBrain()).run(SAVE.goal) trace = Trace.record(episode, expected="Save", brain="planner", scenario="save") assert trace.expected == "Save" and trace.brain == "planner" and trace.scenario == "save" assert trace.frames == len(episode.steps) assert trace.choice == "Save" # the element index survives; the pixels deliberately do not assert trace.steps[-1].elements and all(isinstance(e, Element) for e in trace.steps[-1].elements) assert all(step.observation().frame is None for step in trace.steps) def test_waits_are_not_choices(): episode = AgentLoop( screen=SAVE.screen(), brain=MockBrain([{"kind": "wait", "seconds": 0.1}, {"kind": "click", "label": "Save"}]), ).run(SAVE.goal) trace = Trace.record(episode, expected="Save") assert trace.steps[0].action.kind == "wait" assert trace.choice == "Save" # the click, not the wait assert trace.decided == "Save" def test_a_brain_that_says_done_before_choosing_is_recorded_bluntly(): episode = AgentLoop( screen=SAVE.screen(), brain=MockBrain([{"kind": "done", "reason": "assumed"}]), ).run(SAVE.goal) trace = Trace.record(episode, expected="Save") assert trace.choice is None assert trace.decided == DONE assert score_trace(trace).correct is False def test_a_trace_that_never_committed_reads_as_nothing(): trace = Trace(goal="click something") assert trace.decided == NOTHING and trace.choice is None assert trace.frames == 0 # -- persistence ----------------------------------------------------------- def test_a_trace_round_trips_through_json(): trace = run_scenario(SAVE, PlannerBrain()) again = Trace.from_json(trace.to_json()) assert again.to_dict() == trace.to_dict() assert again.to_json() == trace.to_json() assert json.loads(trace.to_json())["trace_version"] == TRACE_VERSION def test_a_trace_saves_and_loads_from_a_directory(tmp_path): run_scenario(SAVE, PlannerBrain()).save(tmp_path / "b.json") run_scenario(CANCEL, PlannerBrain()).save(tmp_path / "a.json") names = [t.name for t in load_traces(tmp_path)] # directories load in name order assert names == ["cancel", "save"] assert load_traces(tmp_path / "a.json")[0].name == "cancel" def test_a_newer_trace_version_is_refused_rather_than_misread(): with pytest.raises(ValueError, match="newer than this build"): Trace.from_dict({"trace_version": TRACE_VERSION + 1, "goal": "g"}) def test_a_corrupt_action_in_a_trace_is_reported_on_load(): data = {"goal": "g", "steps": [{"index": 0, "action": {"kind": "format_disk"}}]} with pytest.raises(ActionError, match="unknown action"): Trace.from_dict(data) def test_unknown_keys_are_ignored_so_a_trace_stays_forward_compatible(): trace = Trace.from_dict({"goal": "g", "expected": "Save", "extra": 1, "steps": [{"index": 0, "future_field": True}]}) assert trace.goal == "g" and trace.frames == 1 # -- scoring --------------------------------------------------------------- def test_score_matches_the_choice_to_the_expected_answer(): board = benchmark([SAVE, CANCEL], {"planner": lambda s: PlannerBrain()}) assert board.accuracy == 1.0 assert all(v.correct for v in board) def test_score_forgives_ocr_spacing_the_way_the_index_does(): trace = Trace(goal="g", expected="Save", steps=[TraceStep(0, action=Action("click", label="Save ."))]) assert score_trace(trace).correct is True def test_a_wrong_choice_is_scored_with_the_alternative_it_picked(): verdict = score(SAVE, FirstItemBrain()) assert verdict.correct is False assert verdict.chosen == "Files" assert "not 'Save'" in verdict.detail def test_a_trace_with_no_expectation_cannot_be_scored_and_says_so(): verdict = score_trace(run_scenario(SAVE, PlannerBrain())) # run_scenario always records ground truth, so drop it to check the guard unjudged = score_trace(Trace(goal="g", steps=[TraceStep(0, action=Action("click", label="Save"))])) assert verdict.correct is True assert unjudged.correct is False and "no expected answer" in unjudged.detail def test_a_refused_decision_is_scored_on_the_choice_but_flagged(): """A denied click is still what the brain chose - and it must not look clean.""" verdict = score(DENIED, FirstItemBrain()) assert verdict.chosen == "Terminal" and verdict.correct is True assert verdict.blocked is True assert "refused" in verdict.line() def test_the_scoreboard_counts_blocked_decisions_separately_from_accuracy(): board = benchmark([SAVE, DENIED], {"first-item": lambda s: FirstItemBrain()}) assert board.total == 2 and board.blocked == 1 assert "stopped by the harness" in board.render() # -- replay ---------------------------------------------------------------- def test_replay_answers_the_recorded_frames_with_a_different_brain(): recorded = run_scenario(SAVE, PlannerBrain()) assert recorded.choice == "Save" replayed = recorded.replay(FirstItemBrain()) assert replayed.brain == "first-item" assert replayed.choice == "Files" # same frames, a worse brain assert replayed.frames == recorded.frames assert replayed.steps[0].elements == recorded.steps[0].elements def test_replay_keeps_the_ground_truth_so_the_replay_can_be_scored(): recorded = run_scenario(CANCEL, FirstItemBrain()) replayed = recorded.replay(PlannerBrain()) assert replayed.expected == "Cancel" assert score_trace(replayed).correct is True def test_a_brain_that_explodes_during_replay_is_a_scored_error_not_a_crash(): class Exploding: name = "exploding" def decide(self, goal, observation): raise RuntimeError("model died") replayed = run_scenario(SAVE, PlannerBrain()).replay(Exploding()) assert replayed.decided == NOTHING assert "RuntimeError: model died" in replayed.steps[0].error assert score_trace(replayed).correct is False def test_score_replay_is_the_single_trace_form_of_the_roadmap_step(): recorded = run_scenario(SAVE, PlannerBrain()) assert score_replay(recorded, PlannerBrain()).correct is True assert score_replay(recorded, FirstItemBrain()).correct is False def test_benchmark_traces_replays_a_recorded_suite_against_one_brain(): traces = [run_scenario(SAVE, PlannerBrain())] # a trace of a run that got it right assert benchmark_traces(traces, PlannerBrain()).accuracy == 1.0 # the same frames, a worse brain: the trace is the fixed part, the brain is the variable replayed = benchmark_traces(traces, FirstItemBrain()) assert replayed.brains() == ["first-item"] assert replayed.accuracy == 0.0 # -- the suite, and its table --------------------------------------------- def test_benchmark_scores_every_brain_on_every_scenario_with_a_fresh_brain_each_time(): seen = [] def make(scenario): seen.append(scenario.name) return PlannerBrain() board = benchmark([SAVE, CANCEL], {"planner": make, "oracle": oracle, "first-item": lambda s: FirstItemBrain()}) assert board.total == 6 assert board.brains() == ["planner", "oracle", "first-item"] assert board.by_brain() == {"planner": 1.0, "oracle": 1.0, "first-item": 0.0} assert seen == ["save", "cancel"] # a fresh brain per scenario assert {row["correct"] for row in board.to_rows()} == {True, False} def test_the_scoreboard_table_has_the_columns_a_notebook_wants(): board = benchmark([SAVE], {"planner": lambda s: PlannerBrain()}) assert board.to_rows() == [ {"brain": "planner", "scenario": "save", "expected": "Save", "chose": "Save", "correct": True} ] assert board.confusion() == {("Save", "Save"): 1} def test_an_empty_scoreboard_is_zero_not_a_division_error(): board = Scoreboard() assert board.total == 0 and board.accuracy == 0.0 and board.brains() == [] assert board.to_rows() == [] and "0/0" in board.render() # -- scenarios ------------------------------------------------------------- def test_a_scenario_needs_at_least_one_frame(): with pytest.raises(ValueError, match="at least one frame"): Scenario("g", frames=[], expected="Save") def test_a_scenario_hands_out_a_fresh_screen_every_time(): scenario = Scenario("g", frames=[frame("Save")], expected="Save") first, second = scenario.screen(), scenario.screen() assert first is not second assert first.captures == second.captures == 0 assert scenario.frames[0][0].label == "Save" def test_render_is_a_replayable_summary_of_the_trace(): trace = run_scenario(DENIED, FirstItemBrain()) text = trace.render() assert "trace: terminal (denied by policy)" in text assert "brain: first-item" in text assert "expected: 'Terminal'" in text and "chose: 'Terminal'" in text assert "the harness refused it" in text assert "click" in text # -- the command line ------------------------------------------------------ def test_cli_records_a_trace_and_scores_it(tmp_path, capsys): from tiny_computer_agent import __main__ as cli target = tmp_path / "files.json" assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]', "--trace", str(target), "--expected", "Files"]) == 0 out = capsys.readouterr().out assert "trace: wrote" in out and "2 frames" in out assert "correct" in out trace = Trace.load(target) assert trace.expected == "Files" and trace.choice == "Files" assert score_trace(trace).correct is True def test_cli_writes_a_trace_even_without_an_expected_answer(tmp_path, capsys): from tiny_computer_agent import __main__ as cli target = tmp_path / "unscored.json" assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]', "--trace", str(target)]) == 0 assert "score:" not in capsys.readouterr().out trace = Trace.load(target) assert trace.expected is None and score_trace(trace).correct is False