From fe129354fc93972532c6049afe64582054d85c7a Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 14 Oct 2026 15:38:05 +0000
Subject: [PATCH] test: feature unit tests for episode_to_html, agent_xray
 helpers, overlay_utils

Co-Authored-By: Daniella Grimberg <daniella@cognition.ai>
---
 tests/analyze/test_agent_xray_helpers.py | 867 +++++++++++++++++++++++
 tests/analyze/test_episode_to_html.py    | 346 +++++++++
 tests/analyze/test_overlay_utils_unit.py | 284 ++++++++
 3 files changed, 1497 insertions(+)
 create mode 100644 tests/analyze/test_agent_xray_helpers.py
 create mode 100644 tests/analyze/test_episode_to_html.py
 create mode 100644 tests/analyze/test_overlay_utils_unit.py

diff --git a/tests/analyze/test_agent_xray_helpers.py b/tests/analyze/test_agent_xray_helpers.py
new file mode 100644
index 000000000..1af23abfd
--- /dev/null
+++ b/tests/analyze/test_agent_xray_helpers.py
@@ -0,0 +1,867 @@
+import base64
+import json
+import math
+import os
+import re
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import pytest  # noqa: E402
+from openai.types.responses import ResponseFunctionToolCall  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from agentlab.analyze import agent_xray as ax  # noqa: E402
+from agentlab.experiments.loop import ExpResult, StepInfo, StepTimestamps  # noqa: E402
+from agentlab.llm.llm_utils import Discussion  # noqa: E402
+from agentlab.llm.response_api import MessageBuilder, ToolCall, ToolCalls  # noqa: E402
+
+GR_UPDATE = {"__type__": "update"}
+
+
+@pytest.fixture(autouse=True)
+def fresh_info(monkeypatch):
+    """Give every test its own module-global `info`; monkeypatch restores the original."""
+    new_info = ax.Info()
+    monkeypatch.setattr(ax, "info", new_info)
+    yield new_info
+    plt.close("all")
+
+
+def _img_array(color, size=(8, 6)):
+    arr = np.zeros((size[1], size[0], 3), dtype=np.uint8)
+    arr[:, :] = color
+    return arr
+
+
+def _make_exp_result(tmp_path, steps, summary=None, logs=None, name="exp"):
+    exp_dir = tmp_path / "study" / name
+    exp_dir.mkdir(parents=True)
+    for s in steps:
+        s.save_step_info(exp_dir, save_som=True)
+    if summary is not None:
+        (exp_dir / "summary_info.json").write_text(json.dumps(summary))
+    if logs is not None:
+        (exp_dir / "experiment.log").write_text(logs)
+    exp = ExpResult(exp_dir)
+    exp._exp_args = SimpleNamespace(env_args=SimpleNamespace(task_name="task.a", task_seed=3))
+    return exp
+
+
+GOAL = [{"type": "text", "text": "Do it"}]
+
+
+def _step(i, action=None, obs=None, agent_info=None, screens=True, goal=True, **kw):
+    obs = dict(obs or {})
+    if goal:
+        obs.setdefault("goal_object", GOAL)
+    if screens:
+        obs["screenshot"] = _img_array((20 * i, 0, 0))
+        obs["screenshot_som"] = _img_array((0, 20 * i, 0))
+    return StepInfo(step=i, obs=obs, action=action, agent_info=agent_info or {}, **kw)
+
+
+@pytest.fixture
+def episode(tmp_path, fresh_info):
+    steps = [
+        _step(0, task_info={"k": 1}),
+        _step(
+            1,
+            action="click('a1')",
+            obs={
+                "axtree_txt": "AX",
+                "dom_txt": "<html/>",
+                "pruned_html": "<p/>",
+                "extra_element_properties": {"a1": {"bbox": (1, 1, 3, 3)}},
+            },
+            agent_info={
+                "think": "thinking",
+                "chat_messages": [
+                    {"role": "system", "content": [{"type": "text", "text": "sys"}]},
+                    {"role": "user", "content": [{"type": "text", "text": "user prompt"}]},
+                ],
+                "markdown_page": "# md",
+                "html_page": "<b>hi</b>",
+            },
+            stats={"n_tokens": 5, "cost": 0.1},
+        ),
+        _step(2, action="noop()", obs={"last_action_error": "boom"}, terminated=True),
+    ]
+    exp = _make_exp_result(
+        tmp_path,
+        steps,
+        summary={"cum_reward": 1.0, "err_msg": None, "stack_trace": "Traceback..."},
+        logs="log line",
+    )
+    fresh_info.exp_result = exp
+    fresh_info.step = 1
+    return exp
+
+
+# --------------------------------------------------------------------------- tables
+
+
+def test_display_table_wraps_dotted_names_without_mutating():
+    df = pd.DataFrame({"agent.name": [1], "b": [2]}).set_index("agent.name")
+    df2 = pd.DataFrame({"env.task_name": [1]}, index=pd.Index([0], name="x.y"))
+    out = ax.display_table(df2)
+    assert list(out.columns) == ["env.\ntask_name"]
+    assert list(out.index.names) == ["x.\ny"]
+    assert list(df2.columns) == ["env.task_name"]
+    assert list(ax.display_table(df).index.names) == ["agent.\nname"]
+
+
+def test_remove_args_from_col_in_place():
+    df = pd.DataFrame({"agent_args.x": [1], "env_args.task_name": [2]})
+    df.index.name = "agent_args.name"
+    out = ax.remove_args_from_col(df)
+    assert out is df
+    assert list(df.columns) == ["agent.x", "env.task_name"]
+    assert list(df.index.names) == ["agent.name"]
+
+
+@pytest.mark.parametrize(
+    "cols, expected",
+    [([], []), (["a"], ["a"]), (["a.b.c", "d_args"], ["a.\nb.\nc", "d_args"])],
+)
+def test_clean_column_names(cols, expected):
+    assert ax.clean_column_names(cols) == expected
+
+
+def test_click_mapper_maps_pixels_to_steps():
+    fig, ax_ = plt.subplots()
+    ax_.set_xlim(0, 10)
+    fig.canvas.draw()
+    mapper = ax.ClickMapper(ax_, step_times=[2.0, 5.0, 8.0])
+    for t, expected_step in [(1.0, 0), (3.0, 1), (6.0, 2), (9.0, 3)]:
+        x_pix, _ = ax_.transData.transform((t, 0))
+        assert mapper.to_time(x_pix) == pytest.approx(t)
+        assert mapper.to_step(x_pix) == expected_step
+
+
+def test_dataclass_defaults():
+    assert ax.EpisodeId() == ax.EpisodeId(None, None, None, None)
+    assert ax.StepId().step is None
+    info = ax.Info()
+    assert info.active_tab == "Screenshot"
+    assert info.exp_result is None
+
+
+# --------------------------------------------------------------------------- Info
+
+
+def _result_df(tmp_path):
+    return pd.DataFrame(
+        {
+            ax.TASK_NAME_KEY: ["t1", "t1", "t2"],
+            ax.TASK_SEED_KEY: [0, 1, 0],
+            "agent.agent_name": ["A", "A", "B"],
+            "cum_reward": [1.0, 0.0, 0.5],
+            "err_msg": [None, "err", None],
+            "n_steps": [3, 4, 5],
+            "exp_dir": [tmp_path / "e0", tmp_path / "e1", tmp_path / "e2"],
+        }
+    )
+
+
+class TestInfo:
+    def test_get_agent_id(self):
+        info = ax.Info(agent_id_keys=["agent.agent_name", "x"])
+        row = pd.Series({"agent.agent_name": "A", "x": 2, "y": 3})
+        assert info.get_agent_id(row) == [("agent.agent_name", "A"), ("x", 2)]
+
+    def test_filter_agent_id_preserves_row_index(self, tmp_path):
+        info = ax.Info(result_df=_result_df(tmp_path))
+        info.filter_agent_id([("agent.\nagent_name", "A")])
+        assert list(info.agent_df.index) == ["t1", "t1"]
+        assert list(info.agent_df["_row_index"]) == [0, 1]
+        info.filter_agent_id([("agent.agent_name", "B")])
+        assert list(info.agent_df["_row_index"]) == [2]
+
+    def test_update_exp_result(self, tmp_path):
+        info = ax.Info(result_df=_result_df(tmp_path))
+        info.filter_agent_id([("agent.agent_name", "A")])
+        info.step = 5
+        info.update_exp_result(ax.EpisodeId("A", "t1", 1, row_index=1))
+        assert isinstance(info.exp_result, ExpResult)
+        assert info.exp_result.exp_dir == tmp_path / "e1"
+        assert info.step == 0
+
+    def test_update_exp_result_unknown_row(self, tmp_path):
+        info = ax.Info(result_df=_result_df(tmp_path))
+        info.filter_agent_id([("agent.agent_name", "A")])
+        with pytest.raises(ValueError, match="_row_index: 2"):
+            info.update_exp_result(ax.EpisodeId("A", "t2", 0, row_index=2))
+        assert info.exp_result is None
+
+    def test_update_exp_result_duplicate_rows_warns(self, tmp_path, monkeypatch):
+        info = ax.Info(result_df=_result_df(tmp_path))
+        info.filter_agent_id([("agent.agent_name", "A")])
+        info.agent_df["_row_index"] = 0
+        warn = MagicMock()
+        monkeypatch.setattr(ax, "warning", warn)
+        info.update_exp_result(ax.EpisodeId("A", "t1", 0, row_index=0))
+        assert "multiple rows" in warn.call_args[0][0]
+        assert info.exp_result.exp_dir == tmp_path / "e0"
+
+    @pytest.mark.xfail(
+        strict=True,
+        reason="BUG: update_exp_result does not return after the result_df-is-None guard",
+    )
+    def test_update_exp_result_without_results_clears_selection(self):
+        info = ax.Info()
+        info.update_exp_result(ax.EpisodeId())
+        assert info.exp_result is None
+
+
+# --------------------------------------------------------------------------- markdown
+
+
+class TestDictToMarkdown:
+    def test_tool_calls_render_empty(self):
+        assert ax.dict_to_markdown(ToolCalls()) == ""
+
+    def test_non_dict_returns_repr_and_warns(self, monkeypatch):
+        warn = MagicMock()
+        monkeypatch.setattr(ax, "warning", warn)
+        assert ax.dict_to_markdown([1, "a"]) == "[1, 'a']"
+        warn.assert_called_once()
+
+    def test_empty(self):
+        assert ax.dict_to_markdown({}) == "No Data"
+
+    def test_nested(self):
+        d = {"a": 1, "b": {"c": "x"}, "l": [{"k": 2}, "s"], "e": {}}
+        assert ax.dict_to_markdown(d) == (
+            "- **a**: 1\n"
+            "### b\n- **c**: x\n\n"
+            "### l\n#### Item 0\n- **k**: 2\n\n- s\n"
+            "### e\nNo Data\n"
+        )
+
+
+class TestDictMsgToMarkdown:
+    def test_without_role_falls_back(self):
+        assert ax.dict_msg_to_markdown({"x": 1}) == "- **x**: 1\n"
+
+    def test_content_types(self):
+        class PydanticLike:
+            def dict(self):
+                return {"type": "text", "text": "from dict()"}
+
+        msg = {
+            "role": "assistant",
+            "content": [
+                {"type": "image", "image": "data:img"},
+                {"type": "text", "text": "hello"},
+                {"type": "tool_use", "name": "click", "input": {"bid": "a"}, "id": "t1"},
+                {"type": "weird", "v": 1},
+                PydanticLike(),
+            ],
+        }
+        out = ax.dict_msg_to_markdown(msg)
+        assert out.startswith("### Assistant\n")
+        assert "![Image](data:img)" in out
+        assert "\n```\nhello\n```\n" in out
+        assert "Tool Call: click  `{'bid': 'a'}` (call_id: t1)" in out
+        assert "{'type': 'weird', 'v': 1}" in out
+        assert "from dict()" in out
+
+    @pytest.mark.xfail(
+        strict=True,
+        reason="BUG: dict messages with plain-string content crash (iterates characters)",
+    )
+    def test_string_content(self):
+        out = ax.dict_msg_to_markdown({"role": "user", "content": "hello"})
+        assert "hello" in out
+
+
+def test_format_tool_call():
+    assert ax._format_tool_call("f", '{"a": 1}', "c9") == 'Tool Call: f  `{"a": 1}` (call_id: c9)'
+
+
+class TestFormatChatMessage:
+    def test_langchain_message(self, monkeypatch):
+        class FakeLCMessage:
+            def __init__(self, content):
+                self.content = content
+
+        monkeypatch.setattr(ax, "BaseMessage", FakeLCMessage)
+        assert ax.format_chat_message(FakeLCMessage("lc content")) == "lc content"
+
+    def test_message_builder(self):
+        mb = MessageBuilder.assistant().add_text("yo")
+        assert ax.format_chat_message(mb) == "### Assistant\n\n```\nyo\n```\n"
+
+    def test_message_builder_with_tool_calls(self):
+        tc = ToolCalls([ToolCall(name="click", arguments={"bid": "a1"}).response_text("ok")])
+        md = ax.format_chat_message(MessageBuilder.add_responded_tool_calls(tc))
+        assert "**Tool Call 1**: click(bid='a1')" in md
+        assert "ok" in md
+
+    def test_dict(self):
+        out = ax.format_chat_message({"role": "user", "content": [{"type": "text", "text": "q"}]})
+        assert out == "### User\n\n```\nq\n```\n"
+
+    def test_response_function_tool_call(self):
+        call = ResponseFunctionToolCall(
+            name="fill", arguments='{"v": 1}', call_id="c1", type="function_call"
+        )
+        assert ax.format_chat_message(call) == (
+            '### Tool Use\n```\nTool Call: fill  `{"v": 1}` (call_id: c1)\n```\n'
+        )
+
+    def test_other(self):
+        assert ax.format_chat_message(12) == "12"
+        assert ax.format_chat_message("No Chat Messages") == "No Chat Messages"
+
+
+def test_page_to_iframe_roundtrip():
+    page = "<h1>héllo</h1>"
+    out = ax._page_to_iframe(page)
+    b64 = re.search(r'src="data:text/html;base64,([^"]+)"', out).group(1)
+    assert base64.b64decode(b64).decode("utf-8") == page
+    assert out.strip().startswith("<iframe") and out.strip().endswith("</iframe>")
+
+
+def test_code():
+    assert ax.code("x") == "```\nx\n```"
+
+
+# --------------------------------------------------------------------------- tabs
+
+
+class TestIfActive:
+    def test_runs_when_active(self, fresh_info):
+        fresh_info.active_tab = "Logs"
+        fn = MagicMock(return_value="ran")
+        assert ax.if_active("Logs")(fn)(1, k=2) == "ran"
+        fn.assert_called_once_with(1, k=2)
+
+    @pytest.mark.parametrize("n_out, expected", [(1, GR_UPDATE), (3, (GR_UPDATE,) * 3)])
+    def test_skips_when_inactive(self, fresh_info, n_out, expected):
+        fresh_info.active_tab = "Screenshot"
+        fn = MagicMock()
+        assert ax.if_active("Logs", n_out)(fn)() == expected
+        fn.assert_not_called()
+
+    def test_tab_select_sets_active_tab(self, fresh_info):
+        ax.tab_select(SimpleNamespace(value="AXTree"))
+        assert fresh_info.active_tab == "AXTree"
+
+
+class TestStepNavigation:
+    def test_handle_key_event(self, fresh_info, episode):
+        ep = ax.EpisodeId("A", "t", 0, 0)
+        assert ax.handle_key_event("Cmd+Left", ax.StepId(ep, 1)) == ("", ax.StepId(ep, 0))
+        assert ax.handle_key_event("Cmd+Left", ax.StepId(ep, 0))[1].step == 0
+        assert ax.handle_key_event("Cmd+Right", ax.StepId(ep, 0))[1].step == 1
+        # clamps to len(steps) - 2
+        assert ax.handle_key_event("Cmd+Right", ax.StepId(ep, 1))[1].step == 1
+        assert fresh_info.step == 1
+
+    def test_handle_key_event_noop(self, fresh_info):
+        sid = ax.StepId(None, 4)
+        assert ax.handle_key_event("", sid) == ("", sid)
+        assert ax.handle_key_event("Shift+Up", sid) == GR_UPDATE
+
+    def test_gallery_step_change(self, fresh_info):
+        ep = ax.EpisodeId("A")
+        assert ax.gallery_step_change(SimpleNamespace(index=2), ep) == ax.StepId(ep, 2)
+        assert fresh_info.step == 2
+
+    def test_select_step_uses_click_mapper(self, fresh_info):
+        fresh_info.click_mapper = MagicMock()
+        fresh_info.click_mapper.to_step.return_value = 3
+        out = ax.select_step(ax.EpisodeId("A"), SimpleNamespace(index=(120, 5)))
+        fresh_info.click_mapper.to_step.assert_called_once_with(120)
+        assert out.step == 3 and fresh_info.step == 3
+
+
+# --------------------------------------------------------------------------- screenshots
+
+
+class TestScreenshots:
+    def test_get_screenshot_raw_and_som(self, fresh_info, episode):
+        img, action = ax.get_screenshot(fresh_info, step=2)
+        assert img.getpixel((0, 0)) == (40, 0, 0) and action is None
+        img, _ = ax.get_screenshot(fresh_info, step=2, som_or_not="SOM Screenshots")
+        assert img.getpixel((0, 0)) == (0, 40, 0)
+
+    def test_get_screenshot_annotates_with_current_step(self, fresh_info, episode):
+        img, action = ax.get_screenshot(fresh_info, annotate=True)
+        assert "<span" in action and "'a1'" in action
+        assert img.getpixel((1, 1)) != (20, 0, 0)  # bbox overlay drawn
+
+    def test_annotation_failure_falls_back_to_raw_action(self, fresh_info, tmp_path, monkeypatch):
+        fresh_info.exp_result = _make_exp_result(tmp_path, [_step(0, action="click('zz')")])
+        warn = MagicMock()
+        monkeypatch.setattr(ax, "warning", warn)
+        _, action = ax.get_screenshot(fresh_info, step=0, annotate=True)
+        assert action == "click('zz')"
+        assert "Failed to annotate" in warn.call_args[0][0]
+
+    def test_out_of_range_and_missing(self, fresh_info, episode, tmp_path):
+        assert ax.get_screenshot(fresh_info, step=10) == (None, None)
+        fresh_info.exp_result = _make_exp_result(
+            tmp_path, [_step(0, screens=False)], name="noscreens"
+        )
+        assert ax.get_screenshot(fresh_info, step=0) == (None, None)
+
+    def test_update_screenshot_and_pair(self, fresh_info, episode):
+        assert ax.update_screenshot("Raw Screenshots").getpixel((0, 0)) == (20, 0, 0)
+        s1, s2 = ax.update_screenshot_pair("SOM Screenshots")
+        assert s1.getpixel((5, 5)) == (0, 20, 0)
+        assert s2.getpixel((0, 0)) == (0, 40, 0)
+
+    def test_update_screenshot_gallery(self, fresh_info, episode):
+        gallery = ax.update_screenshot_gallery("Raw Screenshots")
+        assert gallery.selected_index == 1
+        assert gallery.columns == 2
+        assert [item["caption"] for item in gallery.value] == ["Step 0", "Step 1", "Step 2"]
+
+
+# --------------------------------------------------------------------------- tab contents
+
+
+class TestTabUpdates:
+    def test_update_episode_escapes_html_into_srcdoc(self, fresh_info, episode):
+        out = ax.update_episode()
+        assert out.startswith('<iframe srcdoc="')
+        srcdoc = out.split('srcdoc="', 1)[1].split('"', 1)[0]
+        assert "<" not in srcdoc and "&lt;div class=&quot;episode-container&quot;&gt;" in srcdoc
+
+    def test_obs_based_tabs(self, fresh_info, episode):
+        assert ax.update_html() == "<html/>"
+        assert ax.update_pruned_html() == "<p/>"
+        assert ax.update_axtree() == "```\nAX\n```"
+        fresh_info.step = 2
+        assert ax.update_html() == "No DOM HTML"
+        assert ax.update_pruned_html() == "No Pruned HTML"
+        assert ax.update_axtree() == "```\nNo AXTree\n```"
+
+    def test_update_chat_messages_list(self, fresh_info, episode):
+        out = ax.update_chat_messages()
+        assert out == "### System\n\n```\nsys\n```\n\n\n### User\n\n```\nuser prompt\n```\n"
+
+    def test_update_chat_messages_default_and_discussion(self, fresh_info, episode):
+        fresh_info.step = 2
+        assert ax.update_chat_messages() == "No Chat Messages"
+        steps = fresh_info.exp_result.steps_info
+        steps[2].agent_info["chat_messages"] = Discussion([ax.AgentLabBaseMessage("user", "d")])
+        assert ax.update_chat_messages() == "Message 0\n\n```\nd\n```\n\n"
+
+    def test_task_error_and_logs(self, fresh_info, episode, tmp_path):
+        assert ax.update_task_error() == "```\nTraceback...\n```"
+        assert ax.update_logs() == "log line"
+        fresh_info.exp_result = _make_exp_result(tmp_path, [_step(0)], name="bare")
+        assert ax.update_task_error() == "No Task Error"
+        assert ax.update_logs() == "No Logs"
+
+    def test_update_stats(self, fresh_info, episode):
+        df = ax.update_stats()
+        assert df.to_dict("records") == [
+            {"name": "n_tokens", "value": 5},
+            {"name": "cost", "value": 0.1},
+        ]
+        fresh_info.step = 9
+        assert ax.update_stats() is None
+
+    def test_update_agent_info_md(self, fresh_info, episode):
+        assert ax.update_agent_info_md() == "# md"
+        fresh_info.step = 2
+        assert "Fill up markdown_page" in ax.update_agent_info_md()
+        fresh_info.exp_result.steps_info[2].agent_info["markup_page"] = "legacy"
+        assert ax.update_agent_info_md() == "legacy"
+        fresh_info.step = 9
+        assert ax.update_agent_info_md() is None
+
+    def test_update_agent_info_html(self, fresh_info, episode):
+        page, s1, s2 = ax.update_agent_info_html()
+        b64 = re.search(r"base64,([^\"]+)", page).group(1)
+        assert base64.b64decode(b64).decode() == "<b>hi</b>"
+        assert s1.getpixel((0, 0)) == (20, 0, 0) and s2.getpixel((0, 0)) == (40, 0, 0)
+
+        fresh_info.exp_result.steps_info[1].agent_info["html_page"] = None
+        page, _, _ = ax.update_agent_info_html()
+        assert "Fill up html_page" in page
+
+        fresh_info.step = 9
+        assert ax.update_agent_info_html() == (None, None, None)
+
+    @pytest.mark.xfail(
+        strict=True,
+        reason="BUG: missing html_page defaults to a list, which _page_to_iframe can't encode",
+    )
+    def test_update_agent_info_html_without_page(self, fresh_info, episode):
+        fresh_info.step = 2
+        page, _, _ = ax.update_agent_info_html()
+        assert isinstance(page, str)
+
+    def test_update_prompt_tests(self, fresh_info, episode):
+        prompt = [{"type": "text", "text": "user prompt"}]
+        assert ax.update_prompt_tests() == (prompt, prompt)
+        fresh_info.exp_result.steps_info[1].agent_info["chat_messages"] = ["a", "plain"]
+        assert ax.update_prompt_tests() == ("plain", "plain")
+
+
+class TestSubmitAction:
+    def _client(self, monkeypatch):
+        client = MagicMock()
+        client.chat.completions.create.return_value = SimpleNamespace(
+            choices=[SimpleNamespace(message=SimpleNamespace(content="LLM says hi"))]
+        )
+        factory = MagicMock(return_value=client)
+        monkeypatch.setattr(ax, "OpenAI", factory)
+        return client
+
+    def test_dict_messages(self, fresh_info, episode, monkeypatch):
+        client = self._client(monkeypatch)
+        original = fresh_info.exp_result.steps_info[1].agent_info["chat_messages"]
+        assert ax.submit_action("new prompt") == "LLM says hi"
+        kwargs = client.chat.completions.create.call_args.kwargs
+        assert kwargs["model"] == "gpt-4o-mini"
+        assert kwargs["messages"][1] == {"role": "user", "content": "new prompt"}
+        assert len(kwargs["messages"]) == 2
+        assert original[1]["content"] == [{"type": "text", "text": "user prompt"}]  # deep-copied
+
+    def test_langchain_messages_converted(self, fresh_info, episode, monkeypatch):
+        class FakeBase:
+            def __init__(self, content):
+                self.content = content
+
+        class FakeHuman(FakeBase):
+            pass
+
+        monkeypatch.setattr(ax, "BaseMessage", FakeBase)
+        monkeypatch.setattr(ax, "HumanMessage", FakeHuman)
+        fresh_info.exp_result.steps_info[1].agent_info["chat_messages"] = [
+            FakeBase("sys"),
+            FakeHuman("usr"),
+        ]
+        client = self._client(monkeypatch)
+        ax.submit_action("edited")
+        msgs = client.chat.completions.create.call_args.kwargs["messages"]
+        assert msgs[0] == {"role": "system", "content": "sys"}
+        assert msgs[1]["role"] == "user" and msgs[1]["content"] == "edited"
+
+    def test_rejects_non_user_second_message(self, fresh_info, episode, monkeypatch):
+        self._client(monkeypatch)
+        fresh_info.exp_result.steps_info[1].agent_info["chat_messages"] = [
+            {"role": "system", "content": "a"},
+            {"role": "assistant", "content": "b"},
+        ]
+        with pytest.raises(AssertionError, match="Second message should be user"):
+            ax.submit_action("x")
+
+    def test_rejects_unknown_message_types(self, fresh_info, episode, monkeypatch):
+        self._client(monkeypatch)
+        fresh_info.exp_result.steps_info[1].agent_info["chat_messages"] = ["a", "b"]
+        with pytest.raises(ValueError, match="list of BaseMessage or dict"):
+            ax.submit_action("x")
+
+
+# --------------------------------------------------------------------------- step info panels
+
+
+class TestStepInfoPanels:
+    def test_update_step_info(self, fresh_info, episode):
+        ep_info, action_info, state_error = ax.update_step_info()
+        assert "### task.a (seed: 3)" in ep_info
+        assert "### Step 1 / 2 (Reward: 1.0)" in ep_info
+        assert "Do it" in ep_info
+        assert "Terminated: None, Truncated: None" in ep_info
+        assert "study/exp" in ep_info
+        assert action_info.startswith("**Action:**")
+        assert "<span" in action_info
+        assert "**Think:**\n\n```\nthinking\n```" in action_info
+        assert "```\nboom\n```" in state_error
+
+    def test_get_episode_info_clamps_step_and_handles_missing_summary(self, fresh_info, tmp_path):
+        fresh_info.exp_result = _make_exp_result(tmp_path, [_step(0, task_info=None)])
+        fresh_info.step = 5
+        out = ax.get_episode_info(fresh_info)
+        assert fresh_info.step == 0
+        assert "Step 0 / 0 (Reward: nan)" in out
+
+    @pytest.mark.xfail(
+        strict=True,
+        reason="BUG: get_episode_info errors when obs has no goal_object (str(BaseMessage(None)))",
+    )
+    def test_get_episode_info_no_goal_key(self, fresh_info, tmp_path):
+        fresh_info.exp_result = _make_exp_result(
+            tmp_path, [_step(0, goal=False)], summary={"cum_reward": 0.0}
+        )
+        fresh_info.step = 0
+        assert "### task.a (seed: 3)" in ax.get_episode_info(fresh_info)
+
+    def test_get_episode_info_error_is_reported(self, fresh_info):
+        out = ax.get_episode_info(fresh_info)  # exp_result is None
+        assert out.startswith("**Error while getting episode info**")
+        assert "AttributeError" in out
+
+    def test_get_action_info_without_think(self, fresh_info, episode):
+        fresh_info.step = 2
+        out = ax.get_action_info(fresh_info)
+        assert "noop()" in out and "Think" not in out
+
+    def test_get_action_info_out_of_bounds(self, fresh_info, episode):
+        fresh_info.step = 7
+        assert ax.get_action_info(fresh_info) == (
+            "Step 7 is out of bounds. The episode has 3 steps."
+        )
+
+    def test_get_action_info_no_steps(self, fresh_info):
+        fresh_info.exp_result = MagicMock(steps_info=[])
+        fresh_info.step = 0
+        assert ax.get_action_info(fresh_info) == "No steps were taken"
+
+    @pytest.mark.parametrize("step", [0, 2, 10])
+    def test_get_state_error_none(self, fresh_info, episode, step):
+        fresh_info.step = step
+        assert "```\nNo Error\n```" in ax.get_state_error(fresh_info)
+
+    def test_get_state_error_without_exp_result(self, fresh_info):
+        fresh_info.step = 0
+        assert "No Error" in ax.get_state_error(fresh_info)
+
+
+# --------------------------------------------------------------------------- seed/agent selection
+
+
+class TestSelection:
+    def test_get_seeds_df(self, tmp_path, fresh_info):
+        fresh_info.result_df = _result_df(tmp_path)
+        fresh_info.filter_agent_id([("agent.agent_name", "A")])
+        seeds = ax.get_seeds_df(fresh_info.agent_df, "t1")
+        assert list(seeds.columns) == ["seed", "reward", "err", "n_steps", "idx"]
+        assert seeds.to_dict("records") == [
+            {"seed": 0, "reward": 1.0, "err": False, "n_steps": 3, "idx": 0},
+            {"seed": 1, "reward": 0.0, "err": True, "n_steps": 4, "idx": 1},
+        ]
+
+    def test_update_seeds_selects_first(self, tmp_path, fresh_info):
+        fresh_info.result_df = _result_df(tmp_path)
+        fresh_info.filter_agent_id([("agent.agent_name", "A")])
+        seeds, ep = ax.update_seeds(("agentA", "t1"))
+        assert len(seeds) == 2
+        assert ep == ax.EpisodeId(agent_id="agentA", task_name="t1", seed=0, row_index=0)
+
+    def test_on_select_agent(self, fresh_info):
+        fresh_info.agent_id_keys = ["agent.agent_name"]
+        df = pd.DataFrame({"agent.agent_name": ["A", "B"]})
+        evt = SimpleNamespace(index=(1, 0))
+        assert ax.on_select_agent(evt, df) == [("agent.agent_name", "B")]
+
+    def test_on_select_task(self):
+        df = pd.DataFrame({"x": [1], ax.TASK_NAME_KEY: ["t9"]})
+        evt = SimpleNamespace(row_value=[1, "t9"])
+        assert ax.on_select_task(evt, df, [("a", 1)]) == ([("a", 1)], "t9")
+
+    def test_on_select_seed(self):
+        df = pd.DataFrame(columns=["seed", "reward", "err", "n_steps", "idx"])
+        evt = SimpleNamespace(row_value=[4, 1.0, False, 3, 12])
+        ep = ax.on_select_seed(evt, df, ("ag", "t"))
+        assert ep == ax.EpisodeId(agent_id="ag", task_name="t", seed=4, row_index=12)
+
+
+# --------------------------------------------------------------------------- directory listing
+
+
+def test_get_directory_contents(tmp_path, monkeypatch):
+    (tmp_path / "file.txt").write_text("not a dir")
+    (tmp_path / "2024-01-01_no_summary").mkdir()
+
+    good = tmp_path / "2024-02-01_good"
+    good.mkdir()
+    pd.DataFrame(
+        {"avg_reward": [0.25, 0.5], "n_completed": ["4/4", "3/4"], "n_err": [0, 1]}
+    ).to_csv(good / "summary_df.csv", index=False)
+
+    empty = tmp_path / "2024-03-01_empty"
+    empty.mkdir()
+    pd.DataFrame(columns=["avg_reward", "n_completed", "n_err"]).to_csv(
+        empty / "summary_df.csv", index=False
+    )
+
+    broken = tmp_path / "2024-04-01_broken"
+    broken.mkdir()
+    pd.DataFrame({"other": [1]}).to_csv(broken / "summary_df.csv", index=False)
+
+    out = ax.get_directory_contents(tmp_path)
+    assert out == [
+        ax.select_dir_instructions,
+        "2024-04-01_broken",
+        "2024-02-01_good - avg-reward: 50.0% - completed: 3/4 - errors: 1",
+        "2024-01-01_no_summary",
+    ]
+
+
+def test_get_directory_contents_picks_most_recent_summary(tmp_path, monkeypatch):
+    d = tmp_path / "study"
+    d.mkdir()
+    old, new = d / "summary_old.csv", d / "summary_new.csv"
+    pd.DataFrame({"avg_reward": [0.1], "n_completed": ["1/2"], "n_err": [0]}).to_csv(
+        old, index=False
+    )
+    pd.DataFrame({"avg_reward": [0.9], "n_completed": ["2/2"], "n_err": [3]}).to_csv(
+        new, index=False
+    )
+    monkeypatch.setattr(os.path, "getctime", lambda p: 2.0 if "new" in str(p) else 1.0)
+    out = ax.get_directory_contents(tmp_path)
+    assert out[1] == "study - avg-reward: 90.0% - completed: 2/2 - errors: 3"
+
+
+# --------------------------------------------------------------------------- profiling plot
+
+
+def _prof(t0, **kw):
+    p = StepTimestamps(env_start=t0, env_stop=t0 + 1, agent_start=t0 + 1, agent_stop=t0 + 2)
+    p.action_exec_start, p.action_exec_stop = t0 + 0.2, t0 + 0.4
+    for k, v in kw.items():
+        setattr(p, k, v)
+    return p
+
+
+def _passthrough_progress(it, desc=None):
+    return it
+
+
+class TestProfiling:
+    def test_fig_to_pil(self):
+        fig, _ = plt.subplots(figsize=(2, 1), dpi=50)
+        img = ax.fig_to_pil(fig)
+        assert isinstance(img, Image.Image) and img.size == (100, 50)
+        assert not plt.fignum_exists(fig.number)
+
+    @pytest.mark.parametrize("edge", [False, True])
+    def test_add_patch(self, edge):
+        _, ax_ = plt.subplots()
+        ax.add_patch(ax_, 1, 4, "red", "lbl", edge=edge)
+        (rect,) = ax_.patches
+        assert (rect.get_x(), rect.get_width(), rect.get_height()) == (1, 3, 1)
+        assert rect.get_label() == "lbl"
+        assert rect.get_fill() is (not edge)
+
+    def test_plot_profiling_empty(self, monkeypatch):
+        warn = MagicMock()
+        monkeypatch.setattr(ax, "warning", warn)
+        _, ax_ = plt.subplots()
+        assert ax.plot_profiling(ax_, [], {}, _passthrough_progress) is None
+        warn.assert_called_once_with("No step info to plot")
+
+    def test_plot_profiling_full(self):
+        t0 = 1000.0
+        steps = [
+            StepInfo(step=0, obs={}, action="a0", profiling=_prof(t0)),
+            StepInfo(
+                step=1,
+                obs={},
+                action="a1",
+                profiling=_prof(
+                    t0 + 2,
+                    wait_for_page_loading_start=t0 + 2.5,
+                    wait_for_page_loading_stop=t0 + 2.6,
+                    validation_start=t0 + 2.6,
+                    validation_stop=t0 + 2.7,
+                    get_observation_start=t0 + 2.7,
+                    get_observation_stop=t0 + 2.8,
+                ),
+            ),
+            StepInfo(
+                step=2,
+                obs={"last_action_error": "oops"},
+                action=None,
+                terminated=True,
+                profiling=_prof(t0 + 4, wait_for_page_loading_start=None),
+            ),
+        ]
+        _, ax_ = plt.subplots()
+        times = ax.plot_profiling(
+            ax_, steps, {"cum_reward": 1.0, "err_msg": None}, _passthrough_progress
+        )
+        assert times == pytest.approx([2.0, 4.0])
+        labels = {p.get_label() for p in ax_.patches} - {"_nolegend_", None}
+        assert {
+            "reset",
+            "env",
+            "exec action",
+            "wait for page",
+            "validation",
+            "get observation",
+            "agent",
+        } <= labels
+        texts = [t.get_text() for t in ax_.texts]
+        assert texts == ["1", "2", "R:1.0"]
+        assert ax_.get_xlim() == pytest.approx((0, 7.0))
+
+    @pytest.mark.parametrize(
+        "kw, summary, expected_text",
+        [
+            ({"truncated": True}, {"cum_reward": 0.0, "err_msg": None}, "R:0.0"),
+            ({"terminated": True}, {"cum_reward": 0.0, "err_msg": None}, "R:0.0"),
+            ({"terminated": True}, {"err_msg": "crash"}, "Err"),
+        ],
+    )
+    def test_plot_profiling_end_markers(self, kw, summary, expected_text):
+        steps = [
+            StepInfo(step=0, obs={}, action=None, profiling=_prof(0.0)),
+            StepInfo(step=1, obs=None, action=None, profiling=_prof(2.0), **kw),
+        ]
+        _, ax_ = plt.subplots()
+        ax.plot_profiling(ax_, steps, summary, _passthrough_progress)
+        assert ax_.texts[-1].get_text() == expected_text
+        # truncated / failed episodes are drawn with a black or red marker, not green
+        assert ax_.lines[-1].get_color() in ("black",)
+
+    def test_plot_profiling_rejects_bad_step_type(self):
+        _, ax_ = plt.subplots()
+        bad = SimpleNamespace(profiling=_prof(0.0), step=0)
+        with pytest.raises(AssertionError, match="Expected StepInfo"):
+            ax.plot_profiling(ax_, [bad], {}, _passthrough_progress)
+
+    def test_generate_profiling(self, fresh_info, tmp_path):
+        assert ax.generate_profiling(_passthrough_progress) is None
+
+        steps = [
+            _step(0, profiling=_prof(10.0)),
+            _step(1, action="noop()", profiling=_prof(12.0), terminated=True),
+        ]
+        fresh_info.exp_result = _make_exp_result(tmp_path, steps)  # no summary_info.json
+        with pytest.raises(KeyError, match="err_msg"):
+            # summary falls back to {} and plot_profiling indexes summary_info["err_msg"]
+            ax.generate_profiling(_passthrough_progress)
+
+    def test_generate_profiling_with_summary(self, fresh_info, episode):
+        steps = fresh_info.exp_result.steps_info
+        for i, s in enumerate(steps):
+            s.profiling = _prof(100.0 + 2 * i)
+        img = ax.generate_profiling(_passthrough_progress)
+        assert isinstance(img, Image.Image)
+        assert isinstance(fresh_info.click_mapper, ax.ClickMapper)
+        assert fresh_info.click_mapper.step_times == pytest.approx([4.0, 6.0])
+        assert fresh_info.exp_result.progress_fn is None
+
+
+def test_format_constant_and_variables(tmp_path, fresh_info):
+    df = pd.DataFrame(
+        {
+            "const": ["c"] * 5,
+            "var": ["a", "a", "b", "c", "d"],
+            "stack_trace": ["s1", "s2", "s3", "s4", "s5"],
+        }
+    )
+    fresh_info.result_df = df
+    constants, variables = ax.format_constant_and_variables()
+    assert constants.to_dict("records") == [{"name": "const", "value": "c"}]
+    var_rows = variables[variables["Name"] == "var"]
+    assert list(var_rows["i"]) == [0, 1, 2, "..."]
+    assert var_rows.iloc[0]["count"] == "2/5" and var_rows.iloc[0]["value"] == "a"
+    assert "stack_trace" not in set(variables["Name"])
+    assert math.isnan(variables.iloc[-1]["n unique"])
diff --git a/tests/analyze/test_episode_to_html.py b/tests/analyze/test_episode_to_html.py
new file mode 100644
index 000000000..b7a6e4af0
--- /dev/null
+++ b/tests/analyze/test_episode_to_html.py
@@ -0,0 +1,346 @@
+import base64
+import re
+from io import BytesIO
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from agentlab.analyze import episode_to_html as e2h
+from agentlab.experiments.loop import ExpResult, StepInfo
+from agentlab.llm.llm_utils import BaseMessage
+from agentlab.llm.response_api import MessageBuilder
+
+XSS = "<script>alert('x')</script>"
+
+
+def _img_array(color, size=(8, 6)):
+    arr = np.zeros((size[1], size[0], 3), dtype=np.uint8)
+    arr[:, :] = color
+    return arr
+
+
+def _make_exp_result(tmp_path, steps, task_name="miniwob.click-test", seed=7, som=True):
+    """Write StepInfo pickles + screenshots to tmp_path and return an ExpResult on it."""
+    exp_dir = tmp_path / "exp"
+    exp_dir.mkdir()
+    for i, (obs, action, agent_info) in enumerate(steps):
+        StepInfo(step=i, obs=obs, action=action, agent_info=agent_info).save_step_info(
+            exp_dir, save_som=som
+        )
+    exp_result = ExpResult(exp_dir)
+    exp_result._exp_args = SimpleNamespace(
+        env_args=SimpleNamespace(task_name=task_name, task_seed=seed)
+    )
+    return exp_result
+
+
+def _obs(i, goal=None, with_screens=True, **extra):
+    obs = dict(extra)
+    if goal is not None:
+        obs["goal_object"] = goal
+    if with_screens:
+        obs["screenshot"] = _img_array((10 * i, 0, 0))
+        obs["screenshot_som"] = _img_array((0, 10 * i, 0))
+    return obs
+
+
+@pytest.fixture
+def full_exp(tmp_path):
+    goal = [{"type": "text", "text": f"Click the button {XSS}"}]
+    steps = [
+        (_obs(0, goal=goal), None, {}),
+        (
+            _obs(1, axtree_txt="[a1] button 'OK' <b>", dom_txt="<div id='x'>hi</div>"),
+            "click('a1')",
+            {
+                "think": "I should <click> & go",
+                "chat_messages": [
+                    {"role": "system", "content": "sys"},
+                    {"role": "user", "content": [{"type": "text", "text": XSS}]},
+                ],
+            },
+        ),
+        (_obs(2), "noop()", {}),
+    ]
+    return _make_exp_result(tmp_path, steps)
+
+
+class TestExpResultToHtml:
+    def test_full_episode_embedded(self, full_exp):
+        html = e2h.exp_result_to_html(full_exp)
+
+        assert '<div class="episode-title">miniwob.click-test</div>' in html
+        assert '<div class="episode-meta">Seed: 7</div>' in html
+        assert "Click the button &lt;script&gt;alert(&#x27;x&#x27;)&lt;/script&gt;" in html
+        assert XSS not in html
+
+        # reset step with no action skipped, others rendered
+        assert "<summary>Step 0</summary>" not in html
+        assert "<summary>Step 1</summary>" in html
+        assert "<summary>Step 2</summary>" in html
+        assert "click(&#x27;a1&#x27;)" in html
+        assert "I should &lt;click&gt; &amp; go" in html
+
+        # screenshots embedded for step 1 and 2 (main + som each)
+        assert html.count('src="data:image/png;base64,') == 4
+        assert "file://" not in html
+        assert "<summary>Screenshot_som[1]</summary>" in html
+
+        assert "<summary>AXTree</summary>" in html
+        assert "[a1] button &#x27;OK&#x27; &lt;b&gt;" in html
+        assert "<summary>HTML</summary>" in html
+        assert "&lt;div id=&#x27;x&#x27;&gt;hi&lt;/div&gt;" in html
+        assert "<summary>Prompt</summary>" in html
+        assert '<div class="chat-role">SYSTEM</div>' in html
+        assert '<div class="chat-role">USER</div>' in html
+
+        # step 2 has no think/axtree/prompt -> sections only appear once
+        assert html.count("<summary>AXTree</summary>") == 1
+        assert html.count("<strong>Think:</strong>") == 1
+        assert html.count("<summary>Prompt</summary>") == 1
+
+    def test_embedded_image_roundtrips(self, full_exp):
+        html = e2h.exp_result_to_html(full_exp)
+        b64 = (
+            re.search(r'alt="Screenshot 1"', html)
+            and re.findall(r'src="data:image/png;base64,([^"]+)" alt="Screenshot 1"', html)[0]
+        )
+        img = Image.open(BytesIO(base64.b64decode(b64)))
+        assert img.size == (8, 6)
+        assert img.convert("RGB").getpixel((0, 0)) == (10, 0, 0)
+
+    def test_paths_instead_of_embedding(self, full_exp):
+        html = e2h.exp_result_to_html(full_exp, embed_images=False)
+        assert "data:image/png" not in html
+        main = full_exp.exp_dir.resolve() / "screenshot_step_1.png"
+        som = full_exp.exp_dir.resolve() / "screenshot_som_step_1.png"
+        assert f'src="file://{main}" alt="Screenshot 1"' in html
+        assert f'src="file://{som}" alt="SOM Screenshot 1"' in html
+
+    @pytest.mark.parametrize(
+        "kwargs, n_open_nested, step_open",
+        [
+            ({}, 0, True),
+            ({"steps_open": False}, 0, False),
+            ({"som_open": True}, 2, True),
+            ({"axtree_open": True, "html_open": True, "prompt_open": True}, 3, True),
+            (
+                {"som_open": True, "axtree_open": True, "html_open": True, "prompt_open": True},
+                5,
+                True,
+            ),
+        ],
+    )
+    def test_open_flags(self, full_exp, kwargs, n_open_nested, step_open):
+        html = e2h.exp_result_to_html(full_exp, **kwargs)
+        n_open = len(re.findall(r'<details class="nested-details" open>', html))
+        assert n_open == n_open_nested
+        assert html.count("<details open>") == (2 if step_open else 0)
+        assert html.count("<details >") == (0 if step_open else 2)
+
+    def test_missing_screenshots(self, tmp_path):
+        steps = [
+            (_obs(0, goal="g", with_screens=False), None, {}),
+            (_obs(1, with_screens=False), "noop()", {}),
+        ]
+        exp = _make_exp_result(tmp_path, steps, som=False)
+        for embed in (True, False):
+            html = e2h.exp_result_to_html(exp, embed_images=embed)
+            assert "<p>Screenshot not available</p>" in html
+            assert "Screenshot_som" not in html
+
+    def test_first_step_with_action_is_rendered(self, tmp_path):
+        exp = _make_exp_result(tmp_path, [(_obs(0, goal="g"), "noop()", {})])
+        html = e2h.exp_result_to_html(exp)
+        assert "<summary>Step 0</summary>" in html
+        assert "noop()" in html
+
+    def test_no_steps(self):
+        exp = MagicMock(spec=ExpResult)
+        exp.exp_args.env_args = SimpleNamespace(task_name="t", task_seed=1)
+        exp.steps_info = []
+        html = e2h.exp_result_to_html(exp)
+        assert "goal-section" not in html.split("</style>")[1]
+        assert "<details" not in html
+        assert html.rstrip().endswith("</div>")
+
+    def test_goal_lookup_failure_is_tolerated(self):
+        exp = MagicMock(spec=ExpResult)
+        exp.exp_args.env_args = SimpleNamespace(task_name="t", task_seed=1)
+        exp.steps_info = [StepInfo(step=0, obs=None, action=None)]
+        html = e2h.exp_result_to_html(exp)
+        assert '<div class="goal-section">' not in html
+
+    def test_missing_goal_key_uses_default(self):
+        exp = MagicMock(spec=ExpResult)
+        exp.exp_args.env_args = SimpleNamespace(task_name="t", task_seed=1)
+        exp.steps_info = [StepInfo(step=0, obs={}, action=None)]
+        html = e2h.exp_result_to_html(exp)
+        assert "<div class='code-block'>No goal specified</div>" in html
+
+    @pytest.mark.xfail(strict=True, reason="BUG: task_name is interpolated into HTML unescaped")
+    def test_task_name_is_escaped(self):
+        exp = MagicMock(spec=ExpResult)
+        exp.exp_args.env_args = SimpleNamespace(task_name=XSS, task_seed=1)
+        exp.steps_info = []
+        assert XSS not in e2h.exp_result_to_html(exp)
+
+
+class TestScreenshotHelpers:
+    def test_screenshot_index_error(self):
+        exp = MagicMock()
+        exp.get_screenshot.side_effect = IndexError
+        exp.get_screenshot_path.side_effect = IndexError
+        assert e2h._get_screenshot_html(exp, 3, True) == "<p>Screenshot not available</p>"
+        assert e2h._get_screenshot_html(exp, 3, False) == "<p>Screenshot not available</p>"
+        assert e2h._get_som_screenshot_html(exp, 3, True, True) == ""
+        assert e2h._get_som_screenshot_html(exp, 3, True, False) == ""
+
+    def test_som_none_image_returns_empty(self):
+        exp = MagicMock()
+        exp.get_screenshot.return_value = None
+        assert e2h._get_som_screenshot_html(exp, 0, False, True) == ""
+
+    def test_som_nonexistent_path_returns_empty(self, tmp_path):
+        exp = MagicMock()
+        exp.get_screenshot_path.return_value = tmp_path / "nope.png"
+        assert e2h._get_som_screenshot_html(exp, 0, False, False) == ""
+        exp.get_screenshot_path.return_value = None
+        assert e2h._get_som_screenshot_html(exp, 0, False, False) == ""
+
+    def test_som_embedded_open(self):
+        exp = MagicMock()
+        exp.get_screenshot.return_value = Image.new("RGB", (2, 2))
+        html = e2h._get_som_screenshot_html(exp, 4, True, True)
+        exp.get_screenshot.assert_called_once_with(4, som=True)
+        assert '<details class="nested-details" open>' in html
+        assert "<summary>Screenshot_som[4]</summary>" in html
+        assert 'alt="SOM Screenshot 4"' in html
+
+    def test_image_to_html(self):
+        assert e2h._image_to_html(None, "Pic") == "<p>Pic not available</p>"
+        img = Image.new("RGB", (3, 2), (1, 2, 3))
+        out = e2h._image_to_html(img, "Pic")
+        m = re.fullmatch(
+            r'<img src="data:image/png;base64,(.+)" alt="Pic" class="screenshot">', out
+        )
+        assert m
+        decoded = Image.open(BytesIO(base64.b64decode(m.group(1))))
+        assert decoded.size == (3, 2) and decoded.getpixel((0, 0)) == (1, 2, 3)
+
+    def test_path_to_html(self, tmp_path, monkeypatch):
+        assert e2h._path_to_html(None, "P") == "<p>P not available</p>"
+        assert e2h._path_to_html(tmp_path / "missing.png", "P") == "<p>P not available</p>"
+        f = tmp_path / "a.png"
+        f.write_bytes(b"x")
+        monkeypatch.chdir(tmp_path)
+        out = e2h._path_to_html(Path("a.png"), "P")
+        assert out == f'<img src="file://{f.resolve()}" alt="P" class="screenshot">'
+
+
+class TestFormatGoal:
+    def test_none(self):
+        assert e2h._format_goal(None) == "<div class='code-block'>No goal specified</div>"
+
+    def test_string_escaped(self):
+        assert e2h._format_goal("a < b") == "<div class='code-block'>a &lt; b</div>"
+
+    def test_multipart_keeps_only_text(self):
+        goal = [
+            {"type": "text", "text": "line1"},
+            {"type": "image_url", "image_url": {"url": "data:..."}},
+            {"type": "text", "text": "line2"},
+        ]
+        assert e2h._format_goal(goal) == "<div class='code-block'>line1\nline2</div>"
+
+
+class TestFormatChatMessages:
+    def test_empty(self):
+        assert e2h._format_chat_messages_like_xray([]) == (
+            "<div class='code-block'>No chat messages</div>"
+        )
+        assert "No chat messages" in e2h._format_chat_messages_like_xray(None)
+
+    def test_object_with_role_and_content(self):
+        msg = SimpleNamespace(role="assistant", content="<hi>")
+        out = e2h._format_chat_messages_like_xray([msg])
+        assert out == (
+            '<div class="chat-message"><div class="chat-role">ASSISTANT</div>'
+            '<div class="code-block">&lt;hi&gt;</div></div>'
+        )
+
+    def test_message_builder(self):
+        mb = MessageBuilder.user().add_text("hello")
+        out = e2h._format_chat_messages_like_xray([mb])
+        assert '<div class="chat-role">USER</div>' in out
+        assert "hello" in out
+
+    def test_agentlab_base_message_is_dict(self):
+        out = e2h._format_chat_messages_like_xray([BaseMessage("user", "plain")])
+        assert out == (
+            '<div class="chat-message"><div class="chat-role">USER</div>'
+            '<div class="code-block">plain</div></div>'
+        )
+
+    def test_dict_multipart(self):
+        msg = {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "a&b"},
+                {"type": "image", "image": "data:png"},
+                {"type": "tool_use", "name": "click", "input": {"bid": "a1"}, "call_id": "c1"},
+                {"type": "other", "x": 1},
+                "raw <part>",
+            ],
+        }
+        out = e2h._format_chat_messages_like_xray([msg])
+        blocks = re.findall(r'<div class="code-block">(.*?)</div>', out)
+        assert blocks == [
+            "a&amp;b",
+            "[IMAGE]",
+            "Tool Call: click `{&#x27;bid&#x27;: &#x27;a1&#x27;}` (call_id: c1)",
+            "{&#x27;type&#x27;: &#x27;other&#x27;, &#x27;x&#x27;: 1}",
+            "raw &lt;part&gt;",
+        ]
+
+    def test_dict_defaults(self):
+        out = e2h._format_chat_messages_like_xray([{"foo": 1}])
+        assert '<div class="chat-role">UNKNOWN</div>' in out
+        assert "{&#x27;foo&#x27;: 1}" in out
+
+    def test_other_types(self):
+        out = e2h._format_chat_messages_like_xray(["<str msg>", 42])
+        assert out.count('<div class="chat-message">') == 2
+        assert "&lt;str msg&gt;" in out and ">42<" in out
+
+    @pytest.mark.xfail(strict=True, reason="BUG: chat message role is not HTML-escaped")
+    def test_role_is_escaped(self):
+        out = e2h._format_chat_messages_like_xray([{"role": XSS, "content": "x"}])
+        assert "<SCRIPT>" not in out
+
+
+def test_format_tool_call_defaults():
+    assert e2h._format_tool_call_like_xray({}) == "Tool Call: unknown `{}` (call_id: unknown)"
+    assert (
+        e2h._format_tool_call_like_xray({"name": "f", "input": {"a": 1}, "call_id": "z"})
+        == "Tool Call: f `{'a': 1}` (call_id: z)"
+    )
+
+
+@pytest.mark.parametrize(
+    "text, expected",
+    [
+        ("", ""),
+        ("plain", "plain"),
+        ("<a href=\"x\">'&'</a>", "&lt;a href=&quot;x&quot;&gt;&#x27;&amp;&#x27;&lt;/a&gt;"),
+        ("&lt;", "&amp;lt;"),
+        (123, "123"),
+        (None, "None"),
+    ],
+)
+def test_escape_html(text, expected):
+    assert e2h._escape_html(text) == expected
diff --git a/tests/analyze/test_overlay_utils_unit.py b/tests/analyze/test_overlay_utils_unit.py
new file mode 100644
index 000000000..5173b10ac
--- /dev/null
+++ b/tests/analyze/test_overlay_utils_unit.py
@@ -0,0 +1,284 @@
+import ast
+import inspect
+from types import SimpleNamespace
+
+import pytest
+from PIL import Image
+
+from agentlab.analyze import overlay_utils
+from agentlab.analyze.overlay_utils import ArgInfo
+
+
+def _arg(name, value, start=0, stop=1, fn="f"):
+    return ArgInfo(
+        function_name=fn,
+        name=name,
+        value=value,
+        type=type(value).__name__,
+        start_index=start,
+        stop_index=stop,
+    )
+
+
+def _non_white_pixels(img):
+    return [
+        (x, y)
+        for x in range(img.width)
+        for y in range(img.height)
+        if img.getpixel((x, y)) != (255, 255, 255)
+    ]
+
+
+class TestParseFunctionCalls:
+    def test_positional_args_are_named_from_signature(self):
+        code = 'fill("a12", "hello")'
+        args = overlay_utils.parse_function_calls(code)
+        assert [(a.function_name, a.name, a.value, a.type) for a in args] == [
+            ("fill", "bid", "a12", "str"),
+            ("fill", "value", "hello", "str"),
+        ]
+        assert code[args[0].start_index : args[0].stop_index] == '"a12"'
+        assert code[args[1].start_index : args[1].stop_index] == '"hello"'
+
+    def test_keyword_arg_span_starts_at_keyword_name(self):
+        code = 'click("b1", button="right")'
+        args = overlay_utils.parse_function_calls(code)
+        kw = [a for a in args if a.name == "button"][0]
+        assert kw.value == "right"
+        assert code[kw.start_index : kw.stop_index] == 'button="right"'
+
+    def test_unknown_functions_and_attribute_calls_are_ignored(self):
+        code = 'foo(1, 2)\nobj.click("a1")\nprint("x")'
+        assert overlay_utils.parse_function_calls(code) == []
+
+    def test_extra_positional_args_beyond_signature_are_dropped(self):
+        n_params = len(inspect.signature(overlay_utils.BGYM_FUNCTION_MAP["focus"]).parameters)
+        code = "focus(" + ", ".join(f"'a{i}'" for i in range(n_params + 3)) + ")"
+        args = overlay_utils.parse_function_calls(code)
+        assert len(args) == n_params
+        assert args[0].value == "a0"
+
+    def test_syntax_error_returns_empty_list(self, capsys):
+        assert overlay_utils.parse_function_calls("click('a1'") == []
+        assert "Syntax error" in capsys.readouterr().out
+
+    def test_empty_string(self):
+        assert overlay_utils.parse_function_calls("") == []
+
+    def test_multiline_indices_account_for_previous_lines(self):
+        code = "noop()\n\nmouse_click(12.5, 30)\n"
+        args = overlay_utils.parse_function_calls(code)
+        assert [a.value for a in args] == [12.5, 30]
+        assert [a.type for a in args] == ["float", "int"]
+        assert code[args[0].start_index : args[0].stop_index] == "12.5"
+        assert code[args[1].start_index : args[1].stop_index] == "30"
+
+    def test_signature_failure_is_skipped_with_warning(self, monkeypatch, capsys):
+        def boom(_):
+            raise ValueError("no sig")
+
+        monkeypatch.setattr(overlay_utils.inspect, "signature", boom)
+        assert overlay_utils.parse_function_calls('click("a1")') == []
+        assert "Could not process function click" in capsys.readouterr().out
+
+
+class TestExtractValue:
+    @pytest.mark.parametrize(
+        "src, expected",
+        [
+            ("'abc'", "abc"),
+            ("42", 42),
+            ("3.5", 3.5),
+            ("None", None),
+            ("True", True),
+            ("['a', 1, ['b']]", ["a", 1, ["b"]]),
+            ("some_var", "some_var"),
+            ("1 + 2", "1 + 2"),
+            ("{'k': 1}", "{'k': 1}"),
+            ("-5", "-5"),
+        ],
+    )
+    def test_values(self, src, expected):
+        node = ast.parse(src, mode="eval").body
+        assert overlay_utils._extract_value(node) == expected
+
+
+class TestGetNodeIndices:
+    def test_positional_node(self):
+        src = "a\nfoo(bar)"
+        call = ast.parse(src).body[1].value
+        assert overlay_utils._get_node_indices(src, call.args[0]) == (6, 9)
+
+    def test_keyword_node_uses_keyword_start(self):
+        src = "foo(x=12)"
+        kw = ast.parse(src).body[0].value.keywords[0]
+        start, stop = overlay_utils._get_node_indices(src, kw.value, kw)
+        assert src[start:stop] == "x=12"
+
+    def test_fallback_when_no_end_positions(self):
+        src = "foo(abc)"
+        node = SimpleNamespace(lineno=1, col_offset=4)
+        # get_source_segment returns None without end positions -> one-char fallback
+        assert overlay_utils._get_node_indices(src, node) == (4, 5)
+
+    def test_fallback_uses_source_segment_length(self, monkeypatch):
+        src = "foo(abc)"
+        node = SimpleNamespace(lineno=1, col_offset=4)
+        monkeypatch.setattr(overlay_utils.ast, "get_source_segment", lambda s, n: "abc")
+        assert overlay_utils._get_node_indices(src, node) == (4, 7)
+
+
+class TestFindBidsAndXYPairs:
+    def test_bids_kept_and_xy_merged(self):
+        args = [
+            _arg("bid", "a1", 0, 4),
+            _arg("value", "text"),
+            _arg("x", 10, 10, 12, fn="mouse_click"),
+            _arg("y", "20", 14, 16, fn="mouse_click"),
+        ]
+        out = overlay_utils.find_bids_and_xy_pairs(args)
+        assert [a.name for a in out] == ["bid", "xy"]
+        merged = out[1]
+        assert merged.value == (10.0, 20.0)
+        assert merged.type == "tuple"
+        assert (merged.start_index, merged.stop_index) == (10, 16)
+        assert merged.function_name == "mouse_click"
+
+    def test_prefixed_pairs(self):
+        args = [
+            _arg("from_x", 1, fn="d"),
+            _arg("from_y", 2, fn="d"),
+            _arg("to_x", 3, fn="d"),
+            _arg("to_y", 4, fn="d"),
+        ]
+        out = overlay_utils.find_bids_and_xy_pairs(args)
+        assert [(a.name, a.value) for a in out] == [("from_xy", (1.0, 2.0)), ("to_xy", (3.0, 4.0))]
+
+    def test_mismatched_prefix_not_merged(self):
+        args = [_arg("from_x", 1), _arg("to_y", 2)]
+        assert overlay_utils.find_bids_and_xy_pairs(args) == []
+
+    def test_different_function_not_merged(self):
+        args = [_arg("x", 1, fn="a"), _arg("y", 2, fn="b")]
+        assert overlay_utils.find_bids_and_xy_pairs(args) == []
+
+    def test_non_numeric_values_kept_as_is(self):
+        args = [_arg("x", "left"), _arg("y", None)]
+        (merged,) = overlay_utils.find_bids_and_xy_pairs(args)
+        assert merged.value == ("left", None)
+
+    def test_trailing_x_without_y(self):
+        assert overlay_utils.find_bids_and_xy_pairs([_arg("x", 1)]) == []
+
+    def test_from_bid_and_to_bid(self):
+        code = 'drag_and_drop("a1", "b2")'
+        out = overlay_utils.find_bids_and_xy_pairs(overlay_utils.parse_function_calls(code))
+        assert [(a.name, a.value) for a in out] == [("from_bid", "a1"), ("to_bid", "b2")]
+
+
+class TestDrawing:
+    def test_overlay_cross_draws_centered_plus(self):
+        img = Image.new("RGB", (20, 20), "white")
+        out = overlay_utils.overlay_cross(img, (10, 10), color=(255, 0, 0), length=7, width=1)
+        assert out is img
+        assert img.getpixel((10, 10)) == (255, 0, 0)
+        assert img.getpixel((7, 10)) == (255, 0, 0)
+        assert img.getpixel((13, 10)) == (255, 0, 0)
+        assert img.getpixel((10, 7)) == (255, 0, 0)
+        assert img.getpixel((10, 13)) == (255, 0, 0)
+        assert img.getpixel((7, 7)) == (255, 255, 255)
+        assert img.getpixel((15, 10)) == (255, 255, 255)
+
+    def test_overlay_rectangle_solid(self):
+        img = Image.new("RGB", (30, 30), "white")
+        out = overlay_utils.overlay_rectangle(img, (5, 5, 10, 10), color=(0, 0, 255), dashed=False)
+        assert out is img
+        for p in [(5, 5), (15, 5), (5, 15), (15, 15), (10, 5), (5, 10)]:
+            assert img.getpixel(p) == (0, 0, 255)
+        assert img.getpixel((10, 10)) == (255, 255, 255)
+
+    def test_overlay_rectangle_dashed_has_gaps(self):
+        img = Image.new("RGB", (40, 40), "white")
+        overlay_utils.overlay_rectangle(img, (2, 2, 30, 30), color=(0, 0, 255), dashed=True)
+        top_edge = [img.getpixel((x, 2)) for x in range(2, 33)]
+        assert (0, 0, 255) in top_edge
+        assert (255, 255, 255) in top_edge
+        assert img.getpixel((15, 15)) == (255, 255, 255)
+
+    def test_linedashed_zero_length_draws_nothing(self):
+        img = Image.new("RGB", (10, 10), "white")
+        draw = overlay_utils.ImageDraw.Draw(img)
+        overlay_utils.linedashed(draw, 5, 5, 5, 5, (0, 0, 0), 1)
+        assert _non_white_pixels(img) == []
+
+    def test_linedashed_pattern(self):
+        img = Image.new("RGB", (40, 3), "white")
+        draw = overlay_utils.ImageDraw.Draw(img)
+        overlay_utils.linedashed(draw, 0, 1, 30, 1, (0, 0, 0), 1, dash_length=4, nodash_length=8)
+        row = [img.getpixel((x, 1)) == (0, 0, 0) for x in range(40)]
+        assert row[0:5] == [True] * 5
+        assert row[5:12] == [False] * 7
+        assert row[12] and row[24]
+        assert not any(row[31:])
+
+    def test_linedashed_negative_direction(self):
+        img = Image.new("RGB", (3, 20), "white")
+        draw = overlay_utils.ImageDraw.Draw(img)
+        overlay_utils.linedashed(draw, 1, 18, 1, 0, (0, 0, 0), 1)
+        assert img.getpixel((1, 18)) == (0, 0, 0)
+
+
+class TestCreateColoredHtml:
+    def test_spans_sorted_by_start_and_hex_colors(self):
+        s = 'click("a1") then fill("b2", "x")'
+        a1 = _arg("bid", "a1", s.index('"a1"'), s.index('"a1"') + 4)
+        b2 = _arg("bid", "b2", s.index('"b2"'), s.index('"b2"') + 4)
+        html = overlay_utils.create_colored_html(s, [b2, a1], [(0, 0, 1.0), (1.0, 0, 0)])
+        assert html == (
+            'click(<span style="color: #ff0000; font-weight: bold;">"a1"</span>) then fill('
+            '<span style="color: #0000ff; font-weight: bold;">"b2"</span>, "x")'
+        )
+
+    def test_no_args_returns_original(self):
+        assert overlay_utils.create_colored_html("noop()", [], []) == "noop()"
+
+    @pytest.mark.xfail(
+        strict=True,
+        reason="BUG: create_colored_html inserts raw action text into HTML without escaping",
+    )
+    def test_action_text_is_html_escaped(self):
+        s = 'fill("a1", "<script>alert(1)</script>")'
+        a1 = _arg("bid", "a1", 5, 9)
+        html = overlay_utils.create_colored_html(s, [a1], [(1.0, 0, 0)])
+        assert "<script>" not in html
+
+
+class TestAnnotateAction:
+    def test_bid_and_click_overlays_and_html(self):
+        img = Image.new("RGB", (200, 200), "white")
+        action = 'click("btn")\nmouse_click(150, 150)'
+        props = {"btn": {"bbox": (10, 10, 50, 30)}}
+        html = overlay_utils.annotate_action(img, action, props, colormap="tab10")
+
+        assert html.count("<span") == 2
+        assert '"btn"</span>' in html
+        assert "150, 150</span>" in html
+        # tab10 colors 0 and 1
+        assert "#1f77b4" in html and "#ff7f0e" in html
+        # rectangle drawn around (10,10)-(60,40); cross at (150,150)
+        assert img.getpixel((10, 10)) != (255, 255, 255)
+        assert img.getpixel((150, 150)) != (255, 255, 255)
+        assert img.getpixel((100, 100)) == (255, 255, 255)
+
+    def test_unknown_bid_and_empty_bbox_not_drawn(self):
+        img = Image.new("RGB", (50, 50), "white")
+        html = overlay_utils.annotate_action(
+            img, 'click("missing")\nhover("nobox")', {"nobox": {"bbox": None}}
+        )
+        assert html.count("<span") == 2
+        assert _non_white_pixels(img) == []
+
+    def test_no_overlayable_args(self):
+        img = Image.new("RGB", (10, 10), "white")
+        assert overlay_utils.annotate_action(img, "noop()", {}) == "noop()"
+        assert _non_white_pixels(img) == []