From 1ca247f0252cece585e098877bbd822318f56533 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sun, 1 Nov 2026 08:05:40 +0000
Subject: [PATCH] Add unit tests for untested pure-logic modules

Co-Authored-By: Gilad Samuels <gilad.samuels@cognition.ai>
---
 tests/agents/test_agent_utils.py              |  61 ++++
 tests/agents/test_most_basic_agent.py         | 103 ++++++
 tests/agents/test_reproducibility_agent.py    | 130 ++++++++
 tests/analyze/test_archive_studies.py         |  77 +++++
 tests/analyze/test_episode_to_html.py         | 189 +++++++++++
 tests/analyze/test_inspect_results_utils.py   | 280 ++++++++++++++++
 tests/benchmarks/test_abstract_env.py         |  63 ++++
 .../test_osworld_axtree_preprocessing.py      | 148 +++++++++
 tests/benchmarks/test_setup_benchmark.py      |  98 ++++++
 tests/experiments/test_exp_utils.py           | 131 ++++++++
 tests/llm/test_chat_api_unit.py               | 284 ++++++++++++++++
 tests/llm/test_huggingface_chat_model.py      | 143 ++++++++
 tests/llm/test_llm_utils_messages.py          | 216 ++++++++++++
 tests/llm/test_prompt_templates.py            |  78 +++++
 tests/llm/test_tracking_pricing.py            | 310 ++++++++++++++++++
 15 files changed, 2311 insertions(+)
 create mode 100644 tests/agents/test_agent_utils.py
 create mode 100644 tests/agents/test_most_basic_agent.py
 create mode 100644 tests/agents/test_reproducibility_agent.py
 create mode 100644 tests/analyze/test_archive_studies.py
 create mode 100644 tests/analyze/test_episode_to_html.py
 create mode 100644 tests/analyze/test_inspect_results_utils.py
 create mode 100644 tests/benchmarks/test_abstract_env.py
 create mode 100644 tests/benchmarks/test_osworld_axtree_preprocessing.py
 create mode 100644 tests/benchmarks/test_setup_benchmark.py
 create mode 100644 tests/experiments/test_exp_utils.py
 create mode 100644 tests/llm/test_chat_api_unit.py
 create mode 100644 tests/llm/test_huggingface_chat_model.py
 create mode 100644 tests/llm/test_llm_utils_messages.py
 create mode 100644 tests/llm/test_prompt_templates.py
 create mode 100644 tests/llm/test_tracking_pricing.py

diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py
new file mode 100644
index 000000000..5c173f564
--- /dev/null
+++ b/tests/agents/test_agent_utils.py
@@ -0,0 +1,61 @@
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image
+
+from agentlab.agents.agent_utils import (
+    draw_arrowhead,
+    draw_click_indicator,
+    draw_mouse_pointer,
+    zoom_webpage,
+)
+
+
+@pytest.fixture
+def white_image():
+    return Image.new("RGB", (100, 100), (255, 255, 255))
+
+
+def test_draw_mouse_pointer_returns_new_rgba_image(white_image):
+    out = draw_mouse_pointer(white_image, 10, 10)
+    assert out is not white_image
+    assert out.mode == "RGBA"
+    assert out.size == white_image.size
+    # pixel inside the arrow is darkened, pixel far away is untouched
+    assert out.getpixel((15, 13))[:3] != (255, 255, 255)
+    assert out.getpixel((90, 90))[:3] == (255, 255, 255)
+    # original is not modified
+    assert white_image.getpixel((15, 13)) == (255, 255, 255)
+
+
+def test_draw_click_indicator_leaves_center_gap(white_image):
+    out = draw_click_indicator(white_image, 50, 50)
+    assert out.mode == "RGBA"
+    assert out.getpixel((50, 50))[:3] == (255, 255, 255)  # gap at the center
+    assert out.getpixel((50, 40))[:3] == (0, 0, 0)  # top line
+    assert out.getpixel((60, 50))[:3] == (0, 0, 0)  # right line
+    assert out.getpixel((5, 5))[:3] == (255, 255, 255)
+
+
+def test_draw_arrowhead_draws_two_lines():
+    draw = MagicMock()
+    draw_arrowhead(draw, (0, 0), (100, 0))
+    assert draw.line.call_count == 2
+    for call in draw.line.call_args_list:
+        points = call.args[0]
+        assert points[0] == (100, 0)
+        assert call.kwargs == {"fill": "red", "width": 4}
+
+
+def test_zoom_webpage_evaluates_js():
+    page = MagicMock()
+    assert zoom_webpage(page, 2.0) is page
+    page.evaluate.assert_called_once_with("document.documentElement.style.zoom='200.0%'")
+
+
+@pytest.mark.parametrize("factor", [0, -1.5])
+def test_zoom_webpage_invalid_factor(factor):
+    page = MagicMock()
+    with pytest.raises(ValueError):
+        zoom_webpage(page, factor)
+    page.evaluate.assert_not_called()
diff --git a/tests/agents/test_most_basic_agent.py b/tests/agents/test_most_basic_agent.py
new file mode 100644
index 000000000..3717d80da
--- /dev/null
+++ b/tests/agents/test_most_basic_agent.py
@@ -0,0 +1,103 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from agentlab.agents.most_basic_agent import most_basic_agent
+from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgent, MostBasicAgentArgs
+from agentlab.llm.llm_utils import AIMessage, Discussion, ParseError
+
+OBS = {"goal": "Click the OK button", "pruned_html": "<button bid='a1'>OK</button>"}
+
+
+class FakeChat:
+    def __init__(self, responses):
+        self.responses = list(responses)
+        self.calls = []
+
+    def __call__(self, messages):
+        self.calls.append(messages)
+        return AIMessage(self.responses.pop(0))
+
+
+@dataclass
+class FakeModelArgs:
+    model_name: str = "fake-model"
+    responses: list = field(default_factory=list)
+    prepared: bool = False
+    closed: bool = False
+
+    def make_model(self):
+        self.chat = FakeChat(self.responses)
+        return self.chat
+
+    def prepare_server(self):
+        self.prepared = True
+        return "prepared"
+
+    def close_server(self):
+        self.closed = True
+        return "closed"
+
+
+def test_args_make_agent_and_lifecycle():
+    model_args = FakeModelArgs()
+    args = MostBasicAgentArgs(
+        temperature=0.7, use_chain_of_thought=True, chat_model_args=model_args
+    )
+    agent = args.make_agent()
+    assert isinstance(agent, MostBasicAgent)
+    assert agent.temperature == 0.7 and agent.use_chain_of_thought
+    assert agent.chat is model_args.chat
+
+    args.set_reproducibility_mode()
+    assert args.temperature == 0
+
+    assert args.prepare() == "prepared" and model_args.prepared
+    assert args.close() == "closed" and model_args.closed
+
+
+def test_get_action_parses_first_code_block():
+    model_args = FakeModelArgs(responses=["I will click\n```\nclick('a1')\n```\n```\nnoop()\n```"])
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+
+    action, info = agent.get_action(OBS)
+
+    assert action == "click('a1')"
+    assert info.think.startswith("I will click")
+    assert isinstance(info.chat_messages, Discussion)
+    prompt = str(info.chat_messages[1])
+    assert OBS["goal"] in prompt and OBS["pruned_html"] in prompt
+    assert "chain of thoughts" not in prompt
+    assert info.stats["prompt_length"] == 3  # system, user, assistant
+    assert info.stats["response_length"] == len(info.think)
+    assert {"input_tokens", "output_tokens", "cost"} <= set(info.stats)
+    assert info.extra_info["chat_model_args"]["model_name"] == "fake-model"
+
+
+def test_get_action_chain_of_thought_prompt():
+    model_args = FakeModelArgs(responses=["```\nclick('a1')\n```"])
+    agent = MostBasicAgentArgs(use_chain_of_thought=True, chat_model_args=model_args).make_agent()
+    _, info = agent.get_action(OBS)
+    assert "chain of thoughts" in str(info.chat_messages[1])
+
+
+def test_get_action_retries_on_missing_code_block():
+    model_args = FakeModelArgs(responses=["no code here", "```\nclick('a1')\n```"])
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+    action, _ = agent.get_action(OBS)
+    assert action == "click('a1')"
+    assert len(model_args.chat.calls) == 2
+
+
+def test_get_action_raises_after_retries():
+    model_args = FakeModelArgs(responses=["nope"] * 3)
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+    with pytest.raises(ParseError):
+        agent.get_action(OBS)
+
+
+def test_experiment_config():
+    exps = most_basic_agent.experiment_config()
+    assert len(exps) == 2
+    assert [e.agent_args.use_chain_of_thought for e in exps] == [True, False]
+    assert all(e.env_args.task_name == "miniwob.click-button" for e in exps)
diff --git a/tests/agents/test_reproducibility_agent.py b/tests/agents/test_reproducibility_agent.py
new file mode 100644
index 000000000..c81ce5af0
--- /dev/null
+++ b/tests/agents/test_reproducibility_agent.py
@@ -0,0 +1,130 @@
+from types import SimpleNamespace
+
+import pytest
+from browsergym.experiments.agent import AgentInfo
+from bgym import HighLevelActionSetArgs
+
+from agentlab.agents.generic_agent import AGENT_4o_MINI
+from agentlab.agents.generic_agent.reproducibility_agent import (
+    ReproAgentArgs,
+    ReproChatModel,
+    _diff_stats,
+    _format_messages,
+    _make_agent_stats,
+    _make_backward_compatible,
+    _make_diff,
+    _set_style,
+    make_repro_agent,
+)
+from agentlab.llm.llm_utils import AIMessage, Discussion, HumanMessage, SystemMessage
+
+OLD_MESSAGES = [
+    {"role": "system", "content": "sys"},
+    {"role": "user", "content": "do it"},
+    {"role": "assistant", "content": "<action>click('a1')</action>"},
+]
+
+
+def test_repro_chat_model_replays_next_message():
+    chat = ReproChatModel(OLD_MESSAGES, delay=0)
+    response = chat(OLD_MESSAGES[:2])
+    assert response == OLD_MESSAGES[2]
+    assert chat.new_messages == OLD_MESSAGES
+    assert chat.get_stats() == {}
+
+
+def test_repro_chat_model_does_not_mutate_input():
+    chat = ReproChatModel(OLD_MESSAGES, delay=0)
+    messages = OLD_MESSAGES[:2]
+    chat(messages)
+    assert len(messages) == 2
+
+
+def test_repro_chat_model_exhausted_returns_none_action():
+    chat = ReproChatModel(OLD_MESSAGES[:2], delay=0)
+    response = chat(OLD_MESSAGES[:2])
+    assert response["role"] == "assistant"
+    assert response["content"] == "<action>None</action>"
+
+
+def test_format_messages_list_and_discussion():
+    assert _format_messages(OLD_MESSAGES[:2]) == "system message:\nsys\n\nuser message:\ndo it\n"
+    discussion = Discussion([SystemMessage("sys"), HumanMessage("do it")])
+    assert _format_messages(discussion) == discussion.to_string()
+
+
+def test_diff_stats_identical():
+    assert _diff_stats("a\nb", "a\nb") == dict(lines_added=0, lines_removed=0, difference_ratio=0.0)
+
+
+def test_diff_stats_changes():
+    stats = _diff_stats("a\nb\nc", "a\nB\nc\nd")
+    assert stats["lines_added"] == 2
+    assert stats["lines_removed"] == 1
+    assert stats["difference_ratio"] == pytest.approx(3 / 8)
+
+
+def test_make_diff_has_style_and_no_nowrap():
+    page = _make_diff("same\nold", "same\nnew")
+    assert 'nowrap="nowrap"' not in page
+    assert "<style" in page
+    assert "Old Version" in page and "New Version" in page
+
+
+def test_set_style_appends_and_prepends():
+    html = "<html><head><style>a{}</style></head><body></body></html>"
+    assert "a{}\nb{}" in _set_style(html, "b{}")
+    assert "b{}\na{}" in _set_style(html, "b{}", prepend_previous_style=True)
+
+
+def test_set_style_creates_style_tag():
+    out = _set_style("<html><head></head><body></body></html>", "b{}")
+    assert "<style>\nb{}</style>" in out
+
+
+@pytest.mark.parametrize("agent_info", [{"think": "t"}, AgentInfo(think="t")])
+def test_make_agent_stats(agent_info):
+    old = OLD_MESSAGES
+    new = OLD_MESSAGES[:2] + [AIMessage("<action>noop()</action>")]
+    action, info = _make_agent_stats("noop()", agent_info, None, old, new)
+    assert action == "noop()"
+    assert isinstance(info, AgentInfo)
+    assert info.think == "t"
+    assert "<table" in info.html_page
+    assert info.stats["lines_added"] == 1 and info.stats["lines_removed"] == 1
+
+
+@pytest.mark.parametrize("action_set", ["bid+nav", ["bid", "nav"]])
+def test_make_backward_compatible_converts_action_set(action_set):
+    args = SimpleNamespace(
+        flags=SimpleNamespace(action=SimpleNamespace(action_set=action_set, multi_actions=True))
+    )
+    out = _make_backward_compatible(args)
+    new_set = out.flags.action.action_set
+    assert isinstance(new_set, HighLevelActionSetArgs)
+    assert list(new_set.subsets) == ["bid", "nav"]
+    assert new_set.multiaction is True
+
+
+def test_make_backward_compatible_keeps_new_format():
+    action_set = HighLevelActionSetArgs(subsets=["bid"])
+    args = SimpleNamespace(
+        flags=SimpleNamespace(action=SimpleNamespace(action_set=action_set, multi_actions=False))
+    )
+    assert _make_backward_compatible(args).flags.action.action_set is action_set
+
+
+def test_make_repro_agent(tmp_path):
+    repro = make_repro_agent(AGENT_4o_MINI, tmp_path)
+    assert isinstance(repro, ReproAgentArgs)
+    assert repro._repro_dir == tmp_path
+    assert repro.agent_name.startswith("Repro_")
+    assert repro.chat_model_args is AGENT_4o_MINI.chat_model_args
+    assert repro.flags is AGENT_4o_MINI.flags
+
+
+def test_make_repro_agent_rejects_bad_inputs(tmp_path):
+    with pytest.raises(AssertionError):
+        make_repro_agent(SimpleNamespace(agent_name="x"), tmp_path)
+    with pytest.raises(AssertionError):
+        make_repro_agent(AGENT_4o_MINI, tmp_path / "missing")
diff --git a/tests/analyze/test_archive_studies.py b/tests/analyze/test_archive_studies.py
new file mode 100644
index 000000000..076b03094
--- /dev/null
+++ b/tests/analyze/test_archive_studies.py
@@ -0,0 +1,77 @@
+import pandas as pd
+import pytest
+
+from agentlab.analyze.archive_studies import StudyInfo, search_for_reasons_to_archive
+
+
+def _make_study(result_dir, name, rows):
+    study_dir = result_dir / name
+    study_dir.mkdir()
+    pd.DataFrame(rows).to_csv(study_dir / "summary_df.csv", index=False)
+    return study_dir
+
+
+@pytest.fixture
+def result_dir(tmp_path):
+    _make_study(tmp_path, "good_study", [{"n_completed": "10/10", "n_err": 0}])
+    _make_study(tmp_path, "miniwob-tiny-test_study", [{"n_completed": "10/10", "n_err": 0}])
+    _make_study(tmp_path, "error_study", [{"n_completed": "10/10", "n_err": 8}])
+    _make_study(tmp_path, "small_study", [{"n_completed": "2/2", "n_err": 0}])
+    _make_study(tmp_path, "nothing_completed", [{"n_completed": "0/4", "n_err": 0}])
+    _make_study(
+        tmp_path,
+        "multi_agent",
+        [{"n_completed": "4/5", "n_err": 1}, {"n_completed": "5/5", "n_err": 0}],
+    )
+    (tmp_path / "not_a_dir.txt").write_text("ignored")
+    (tmp_path / "no_summary_study").mkdir()
+    return tmp_path
+
+
+def test_search_for_reasons_to_archive(result_dir):
+    infos = search_for_reasons_to_archive(result_dir, min_study_size=5)
+    by_name = {info.study_dir.name: info for info in infos}
+
+    # files and unreadable studies are skipped
+    assert "not_a_dir.txt" not in by_name
+    assert "no_summary_study" not in by_name
+    assert len(by_name) == 6
+    assert all(isinstance(info, StudyInfo) for info in infos)
+
+    assert not by_name["good_study"].should_delete
+    assert by_name["good_study"].reason == ""
+    assert by_name["good_study"].study is None  # no study.pkl.gz in the folder
+
+    assert by_name["miniwob-tiny-test_study"].should_delete
+    assert "Miniwob tiny test" in by_name["miniwob-tiny-test_study"].reason
+
+    assert by_name["error_study"].should_delete
+    assert "Less than 50%" in by_name["error_study"].reason
+
+    assert by_name["small_study"].should_delete
+    assert "Too few tasks" in by_name["small_study"].reason
+
+    assert by_name["nothing_completed"].should_delete
+    assert "No tasks completed" in by_name["nothing_completed"].reason
+
+    # rows are aggregated: 9/10 completed, 1 error -> kept
+    assert not by_name["multi_agent"].should_delete
+
+
+def test_search_uses_most_recent_summary(tmp_path):
+    import time
+
+    study_dir = _make_study(tmp_path, "study", [{"n_completed": "0/10", "n_err": 0}])
+    newer = study_dir / "summary_new.csv"
+    time.sleep(0.05)
+    pd.DataFrame([{"n_completed": "10/10", "n_err": 0}]).to_csv(newer, index=False)
+
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert info.summary_df["n_completed"].tolist() == ["10/10"]
+    assert not info.should_delete
+
+
+def test_min_study_size_default_keeps_small_study(tmp_path):
+    _make_study(tmp_path, "small", [{"n_completed": "2/2", "n_err": 0}])
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert not info.should_delete
diff --git a/tests/analyze/test_episode_to_html.py b/tests/analyze/test_episode_to_html.py
new file mode 100644
index 000000000..065073735
--- /dev/null
+++ b/tests/analyze/test_episode_to_html.py
@@ -0,0 +1,189 @@
+from types import SimpleNamespace
+
+import pytest
+from PIL import Image
+
+from agentlab.analyze.episode_to_html import (
+    _escape_html,
+    _format_chat_messages_like_xray,
+    _format_goal,
+    _format_tool_call_like_xray,
+    _image_to_html,
+    _path_to_html,
+    exp_result_to_html,
+)
+from agentlab.llm.llm_utils import HumanMessage
+
+
+def _step(action, obs=None, agent_info=None):
+    return SimpleNamespace(action=action, obs=obs or {}, agent_info=agent_info or {})
+
+
+class FakeExpResult:
+    def __init__(self, steps_info, screenshot_dir=None, som=True):
+        self.exp_args = SimpleNamespace(
+            env_args=SimpleNamespace(task_name="miniwob.click", task_seed=7)
+        )
+        self.steps_info = steps_info
+        self.screenshot_dir = screenshot_dir
+        self.som = som
+
+    def get_screenshot(self, step, som=False):
+        if som and not self.som:
+            return None
+        if step >= len(self.steps_info):
+            raise IndexError
+        return Image.new("RGB", (4, 4), (0, 0, 255) if som else (255, 0, 0))
+
+    def get_screenshot_path(self, step, som=False):
+        if self.screenshot_dir is None:
+            raise FileNotFoundError
+        return self.screenshot_dir / f"screenshot{'_som' if som else ''}_step_{step}.png"
+
+
+@pytest.fixture
+def steps():
+    return [
+        _step(None, obs={"goal_object": [{"type": "text", "text": "Click <OK>"}]}),
+        _step(
+            "click('a<1>')",
+            obs={"axtree_txt": "[a1] button 'OK'", "dom_txt": "<button>OK</button>"},
+            agent_info={
+                "think": "I'll click & go",
+                "chat_messages": [{"role": "user", "content": "hi"}],
+            },
+        ),
+        _step(None),
+    ]
+
+
+def test_escape_html():
+    assert _escape_html("<a href=\"x\">'&'</a>") == (
+        "&lt;a href=&quot;x&quot;&gt;&#x27;&amp;&#x27;&lt;/a&gt;"
+    )
+    assert _escape_html(42) == "42"
+    assert _escape_html("") == ""
+
+
+def test_format_tool_call_like_xray():
+    assert (
+        _format_tool_call_like_xray({"name": "click", "input": {"bid": "a1"}, "call_id": "c1"})
+        == "Tool Call: click `{'bid': 'a1'}` (call_id: c1)"
+    )
+    assert _format_tool_call_like_xray({}) == "Tool Call: unknown `{}` (call_id: unknown)"
+
+
+def test_image_to_html():
+    assert _image_to_html(None, "Shot") == "<p>Shot not available</p>"
+    html = _image_to_html(Image.new("RGB", (2, 2)), "Shot")
+    assert html.startswith('<img src="data:image/png;base64,')
+    assert 'alt="Shot"' in html
+
+
+def test_path_to_html(tmp_path):
+    assert _path_to_html(None, "Shot") == "<p>Shot not available</p>"
+    assert _path_to_html(tmp_path / "missing.png", "Shot") == "<p>Shot not available</p>"
+    img = tmp_path / "img.png"
+    Image.new("RGB", (2, 2)).save(img)
+    assert (
+        _path_to_html(img, "Shot")
+        == f'<img src="file://{img.resolve()}" alt="Shot" class="screenshot">'
+    )
+
+
+def test_format_goal():
+    assert "No goal specified" in _format_goal(None)
+    assert _format_goal("Do <it>") == "<div class='code-block'>Do &lt;it&gt;</div>"
+    html = _format_goal([{"type": "text", "text": "a"}, {"type": "text", "text": "b"}])
+    assert html == "<div class='code-block'>a\nb</div>"
+
+
+def test_format_chat_messages_empty():
+    assert _format_chat_messages_like_xray([]) == "<div class='code-block'>No chat messages</div>"
+    assert "No chat messages" in _format_chat_messages_like_xray(None)
+
+
+def test_format_chat_messages_variants():
+    obj_msg = SimpleNamespace(role="assistant", content="<done>")
+    messages = [
+        {"role": "system", "content": "be <brief>"},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "look"},
+                {"type": "image", "source": "..."},
+                {"type": "tool_use", "name": "click", "input": {"bid": "1"}, "call_id": "x"},
+                {"type": "other", "value": 1},
+                "raw part",
+            ],
+        },
+        {"content": "no role"},
+        obj_msg,
+        HumanMessage("dict subclass"),
+        12345,
+    ]
+    html = _format_chat_messages_like_xray(messages)
+    assert html.count('<div class="chat-message">') == len(messages)
+    assert '<div class="chat-role">SYSTEM</div>' in html
+    assert "be &lt;brief&gt;" in html
+    assert "[IMAGE]" in html
+    assert "Tool Call: click" in html
+    assert "raw part" in html
+    assert "&#x27;type&#x27;: &#x27;other&#x27;" in html
+    assert '<div class="chat-role">UNKNOWN</div>' in html
+    assert '<div class="chat-role">ASSISTANT</div>' in html and "&lt;done&gt;" in html
+    assert '<div class="chat-role">USER</div>' in html and "dict subclass" in html
+    assert "12345" in html
+
+
+def test_exp_result_to_html_embedded(steps):
+    html = exp_result_to_html(FakeExpResult(steps), som_open=True, axtree_open=True)
+    assert '<div class="episode-title">miniwob.click</div>' in html
+    assert "Seed: 7" in html
+    assert "Click &lt;OK&gt;" in html
+    # step 0 has no action and is skipped; step 2 is rendered without action
+    assert "<summary>Step 0</summary>" not in html
+    assert "<summary>Step 1</summary>" in html and "<summary>Step 2</summary>" in html
+    assert "click(&#x27;a&lt;1&gt;&#x27;)" in html
+    assert "I&#x27;ll click &amp; go" in html
+    assert "<summary>AXTree</summary>" in html
+    assert "<summary>HTML</summary>" in html
+    assert "<summary>Prompt</summary>" in html
+    assert "<summary>Screenshot_som[1]</summary>" in html
+    assert html.count("data:image/png;base64,") == 4  # main + som for steps 1 and 2
+    assert html.count("<details open>") == 2
+    assert html.endswith("</div>")
+
+
+def test_exp_result_to_html_collapsed_without_som(steps):
+    html = exp_result_to_html(FakeExpResult(steps, som=False), steps_open=False)
+    assert "<details open>" not in html
+    assert "<details >" in html
+    assert "Screenshot_som" not in html
+
+
+def test_exp_result_to_html_paths(tmp_path, steps):
+    Image.new("RGB", (2, 2)).save(tmp_path / "screenshot_step_1.png")
+    Image.new("RGB", (2, 2)).save(tmp_path / "screenshot_som_step_1.png")
+    html = exp_result_to_html(FakeExpResult(steps, screenshot_dir=tmp_path), embed_images=False)
+    assert f"file://{(tmp_path / 'screenshot_step_1.png').resolve()}" in html
+    assert "<summary>Screenshot_som[1]</summary>" in html
+    assert "Screenshot 2 not available" in html
+    assert "Screenshot_som[2]" not in html
+
+
+def test_exp_result_to_html_missing_screenshots(steps):
+    html = exp_result_to_html(FakeExpResult(steps), embed_images=False)
+    assert html.count("<p>Screenshot not available</p>") == 2
+
+
+def test_exp_result_to_html_no_steps():
+    html = exp_result_to_html(FakeExpResult([]))
+    assert "Goal:" not in html
+    assert "<summary>Step" not in html
+
+
+def test_exp_result_to_html_goal_default():
+    html = exp_result_to_html(FakeExpResult([_step("noop()")]))
+    assert "No goal specified" in html
+    assert "<summary>Step 0</summary>" in html
diff --git a/tests/analyze/test_inspect_results_utils.py b/tests/analyze/test_inspect_results_utils.py
new file mode 100644
index 000000000..b537ba97d
--- /dev/null
+++ b/tests/analyze/test_inspect_results_utils.py
@@ -0,0 +1,280 @@
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from agentlab.analyze.inspect_results import (
+    _benchmark_from_task_name,
+    _extract_ablation_study,
+    _find_diff,
+    _get_mtimes,
+    _is_stale,
+    _rename_bool_flags,
+    categorize_error,
+    flag_report,
+    get_constants_and_variables,
+    get_sample_std_err,
+    get_std_err,
+    map_err_key,
+    shrink_columns,
+    split_by_key,
+    summarize,
+    summarize_stats,
+)
+
+
+@pytest.fixture
+def episodes():
+    return pd.DataFrame(
+        {
+            "cum_reward": [1.0, 0.0, 1.0, 0.0],
+            "err_msg": [None, None, None, "boom"],
+            "truncated": [False, False, False, False],
+            "terminated": [True, True, False, False],
+            "n_steps": [3, 4, 5, 6],
+            "stats.cum_cost": [0.1, 0.2, 0.3, 0.4],
+        }
+    )
+
+
+def test_get_constants_and_variables():
+    df = pd.DataFrame({"const": [np.int64(1), np.int64(1)], "var": [1, 2], "lst": [[1], [1]]})
+    constants, variables, out = get_constants_and_variables(df)
+    assert constants == {"const": 1}
+    assert type(constants["const"]) is int
+    assert variables == ["var", "lst"]  # unhashable columns are treated as variables
+    assert out is df
+
+    _, _, dropped = get_constants_and_variables(df, drop_constants=True)
+    assert list(dropped.columns) == ["var", "lst"]
+
+
+def test_get_std_err_binary():
+    df = pd.DataFrame({"r": [1, 0, 1, 1, np.nan]})
+    mean, std_err = get_std_err(df, "r")
+    assert mean == pytest.approx(0.75)
+    assert std_err == pytest.approx(np.sqrt(0.75 * 0.25 / 4))
+
+
+def test_get_std_err_non_binary_uses_sample_std_err():
+    df = pd.DataFrame({"r": [0.5, 1.5, 2.5]})
+    assert get_std_err(df, "r") == get_sample_std_err(df, "r")
+    mean, std_err = get_sample_std_err(df, "r")
+    assert mean == pytest.approx(1.5)
+    assert std_err == pytest.approx(1 / np.sqrt(3))
+
+
+def test_get_sample_std_err_single_value():
+    mean, std_err = get_sample_std_err(pd.DataFrame({"r": [0.3]}), "r")
+    assert mean == pytest.approx(0.3)
+    assert std_err == 0
+
+
+def test_summarize(episodes):
+    record = summarize(episodes)
+    assert record["avg_reward"] == 0.5
+    assert record["std_err"] == pytest.approx(0.25)
+    assert record["avg_steps"] == 4.5
+    assert record["n_completed"] == "3/4"  # 2 terminated + 1 error
+    assert record["n_err"] == 1
+    assert record["cum_cost"] == pytest.approx(1.0)
+
+
+def test_summarize_effective_cost_replaces_cost(episodes):
+    episodes["stats.cum_effective_cost"] = [0.5, 0.5, 0.5, 0.5]
+    record = summarize(episodes)
+    assert "cum_cost" not in record
+    assert record["cum_effective_cost"] == 2.0
+
+
+def test_summarize_without_rewards():
+    record = summarize(pd.DataFrame({"n_steps": [1, 2]}))
+    assert np.isnan(record["avg_reward"])
+    assert record["n_completed"] == "0/2"
+    assert record["n_err"] == 0
+
+
+def test_summarize_nothing_completed(episodes):
+    episodes["terminated"] = False
+    episodes["err_msg"] = None
+    assert summarize(episodes) is None
+
+
+def test_summarize_error_with_reward_fails_sanity_check(episodes):
+    episodes.loc[3, "cum_reward"] = 1.0
+    with pytest.raises(AssertionError):
+        summarize(episodes)
+
+
+def test_summarize_stats(episodes):
+    episodes["stats.max_tokens"] = [10, 30, 20, 5]
+    record = summarize_stats(episodes)
+    assert record["avg_reward"] == 0.5
+    assert record["cum_cost"] == pytest.approx(1.0)
+    assert record["max_tokens"] == 30
+
+
+def test_summarize_stats_unknown_op(episodes):
+    episodes["stats.avg_tokens"] = 1
+    with pytest.raises(ValueError, match="Unknown stats operation: avg"):
+        summarize_stats(episodes)
+
+
+def test_summarize_stats_nothing_completed(episodes):
+    episodes["terminated"] = False
+    episodes["err_msg"] = None
+    assert summarize_stats(episodes) is None
+
+
+def test_find_diff():
+    assert _find_diff((1, 2, 3), (1, 0, 3)) == [1]
+    assert _find_diff((1, 2), (1, 2)) == []
+    assert _find_diff((), ()) == []
+
+
+@pytest.fixture
+def ablation_report():
+    index = pd.MultiIndex.from_tuples(
+        [(True, 1), (False, 1), (True, 2)], names=["use_cot", "n_retry"]
+    )
+    return pd.DataFrame({"avg_reward": [0.5, 0.4, 0.6]}, index=index)
+
+
+def test_extract_ablation_study(ablation_report):
+    out = _extract_ablation_study(ablation_report)
+    assert list(out.index) == ["Initial Configuration", "-use_cot", "n_retry←2"]
+    assert list(out.columns) == ["avg_reward"]
+    assert out["avg_reward"].tolist() == [0.5, 0.4, 0.6]
+
+
+def test_extract_ablation_study_progression(ablation_report):
+    out = _extract_ablation_study(ablation_report, progression=True)
+    assert list(out.index) == [
+        "Initial Configuration",
+        "↳ use_cot=False",
+        "↳ use_cot=True, n_retry=2",
+    ]
+
+
+def test_rename_bool_flags(ablation_report):
+    out = _rename_bool_flags(ablation_report.copy())
+    assert set(out.index.levels[0]) == {"✓", "-"}
+    assert list(out.index.levels[1]) == [1, 2]
+
+
+def test_rename_bool_flags_single_index():
+    df = pd.DataFrame({"a": [1]}, index=[True])
+    assert _rename_bool_flags(df).index.tolist() == [True]
+
+
+def test_flag_report():
+    index = pd.MultiIndex.from_tuples(
+        [(True, True), (True, False), (False, True), (False, False)], names=["f1", "f2"]
+    )
+    report = pd.DataFrame({"avg_reward": [1.0, 0.6, 0.5, 0.2]}, index=index)
+    out = flag_report(report)
+    assert list(out.index) == ["f1", "f2"]  # sorted by ratio
+    assert out.loc["f1", "avg_true"] == 0.8
+    assert out.loc["f1", "avg_false"] == 0.35
+    assert out.loc["f1", "ratio"] == pytest.approx(2.29)
+    assert out.loc["f2", "ratio"] == pytest.approx(1.88)
+
+
+def test_flag_report_single_level():
+    report = pd.DataFrame({"avg_reward": [1.0]}, index=pd.Index(["a"], name="agent"))
+    assert flag_report(report) is None
+
+
+def test_shrink_columns():
+    df = pd.DataFrame(
+        {"avg_reward": [0.5000, 0.125], "n_completed": ["1/2", "2/2"]},
+        index=pd.Index(["a", "b"], name="agent_name"),
+    )
+    out = shrink_columns(df)
+    assert list(out.columns) == ["avg\nreward", "n\ncompleted"]
+    assert out.index.names == ["agent\nname"]
+    assert out["avg\nreward"].tolist() == ["0.5", "0.125"]
+    assert out["n\ncompleted"].tolist() == ["1/2", "2/2"]
+    assert list(df.columns) == ["avg_reward", "n_completed"]  # original untouched
+
+    assert shrink_columns(df, also_wrap_index=False).index.names == ["agent_name"]
+
+
+def test_map_err_key():
+    assert map_err_key(None) is None
+    msg = "Error: your messages resulted in 4242 tokens\n=== logs ===\nlots of logs"
+    assert map_err_key(msg) == "Error: your messages resulted in x tokens"
+    msg = "Exception uncaught by agent or environment in task miniwob.click-test.\n=== logs ==="
+    assert map_err_key(msg) == "Exception uncaught by agent or environment in task <task_name>."
+
+
+def test_categorize_error_without_error():
+    assert categorize_error({"err_msg": None}) is None
+    assert categorize_error({}) is None
+    assert categorize_error({"err_msg": np.nan}) is None
+
+
+@pytest.mark.parametrize(
+    "task, expected",
+    [
+        ("miniwob.click-test", "miniwob"),
+        ("workarena.servicenow.x", "workarena"),
+        ("plain", "plain"),
+    ],
+)
+def test_benchmark_from_task_name(task, expected):
+    assert _benchmark_from_task_name(task) == expected
+
+
+def test_split_by_key_column():
+    df = pd.DataFrame(
+        {
+            "agent.agent_name": ["a", "b", "a"],
+            "env.task_name": ["t1", "t1", "t2"],
+            "cum_reward": [1, 0, 1],
+        }
+    )
+    out = split_by_key(df, "agent.agent_name")
+    assert set(out) == {"a", "b"}
+    assert len(out["a"]) == 2 and len(out["b"]) == 1
+
+
+def test_split_by_key_from_index():
+    df = pd.DataFrame(
+        {"agent.agent_name": ["a", "b"], "env.task_name": ["t", "t"], "cum_reward": [1, 0]}
+    ).set_index("agent.agent_name")
+    out = split_by_key(df, "agent.agent_name")
+    assert set(out) == {"a", "b"}
+    assert out["b"]["cum_reward"].tolist() == [0]
+
+
+def test_get_mtimes(tmp_path):
+    (tmp_path / "sub").mkdir()
+    (tmp_path / "sub" / "file.txt").write_text("x")
+    (tmp_path / "_hidden.json").write_text("{}")
+    (tmp_path / ".dot").write_text("")
+    (tmp_path / "summary.csv").write_text("")
+    mtimes = _get_mtimes(tmp_path, whitelist=(tmp_path / "summary.csv",))
+    assert set(mtimes) == {"sub", os.path.join("sub", "file.txt")}
+
+
+def test_is_stale(tmp_path):
+    summary = tmp_path / "summary_df.csv"
+    data = tmp_path / "exp" / "result.json"
+    data.parent.mkdir()
+    data.write_text("{}")
+
+    assert _is_stale(tmp_path, summary)  # no summary yet
+    summary.write_text("x")
+    assert not _is_stale(tmp_path, summary)
+    saved = json.loads((tmp_path / "_last_modification_times.json").read_text())
+    assert os.path.join("exp", "result.json") in saved
+
+    summary.write_text("y")  # summary changes are ignored
+    assert not _is_stale(tmp_path, summary)
+
+    os.utime(data, (1, 1))
+    assert _is_stale(tmp_path, summary)
+    assert not _is_stale(tmp_path, summary)
diff --git a/tests/benchmarks/test_abstract_env.py b/tests/benchmarks/test_abstract_env.py
new file mode 100644
index 000000000..c6a156662
--- /dev/null
+++ b/tests/benchmarks/test_abstract_env.py
@@ -0,0 +1,63 @@
+import pytest
+
+from agentlab.benchmarks.abstract_env import (
+    AbstractBenchmark,
+    AbstractEnv,
+    add_step_timing_to_env_info_decorator,
+)
+
+
+class DummyEnv(AbstractEnv):
+    def __init__(self, env_info):
+        self.env_info = env_info
+        self.actions = []
+
+    def reset(self, seed=None):
+        return {}, {}
+
+    @add_step_timing_to_env_info_decorator
+    def step(self, action):
+        self.actions.append(action)
+        return {"obs": action}, 1.0, True, False, self.env_info
+
+    def close(self):
+        pass
+
+
+def test_abstract_env_cannot_be_instantiated():
+    with pytest.raises(TypeError):
+        AbstractEnv()
+
+
+def test_decorator_adds_timing_when_env_info_none():
+    env = DummyEnv(env_info=None)
+    obs, reward, terminated, truncated, info = env.step("click")
+    assert obs == {"obs": "click"} and reward == 1.0 and terminated and not truncated
+    assert env.actions == ["click"]
+    assert info["action_exec_start"] <= info["action_exec_stop"]
+    assert info["action_exec_timeout"] == 0.0
+
+
+def test_decorator_preserves_existing_values():
+    preset = {"action_exec_start": 1, "action_exec_stop": 2, "action_exec_timeout": 3, "x": 4}
+    env = DummyEnv(env_info=dict(preset))
+    *_, info = env.step("a")
+    assert info == preset
+
+
+def test_decorator_preserves_function_metadata():
+    assert DummyEnv.step.__name__ == "step"
+
+
+def test_abstract_benchmark_defaults():
+    bench = AbstractBenchmark(name="dummy", env_args_list=[1, 2])
+    assert bench.name == "dummy"
+    assert bench.env_args_list == [1, 2]
+    assert bench.get_version() == "1"
+    assert bench.prepare_backends() is None
+    assert bench.dependency_graph_over_tasks() == {}
+
+
+def test_abstract_benchmark_requires_fields():
+    with pytest.raises(Exception):
+        AbstractBenchmark(name="missing env args")
diff --git a/tests/benchmarks/test_osworld_axtree_preprocessing.py b/tests/benchmarks/test_osworld_axtree_preprocessing.py
new file mode 100644
index 000000000..a59d4ab7a
--- /dev/null
+++ b/tests/benchmarks/test_osworld_axtree_preprocessing.py
@@ -0,0 +1,148 @@
+import io
+import xml.etree.ElementTree as ET
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from agentlab.benchmarks.osworld_axtree_preprocessing import (
+    draw_bounding_boxes,
+    filter_nodes,
+    find_leaf_nodes,
+    judge_node,
+    linearize_accessibility_tree,
+    tag_screenshot,
+    trim_accessibility_tree,
+)
+
+NS = (
+    'xmlns:st="https://accessibility.ubuntu.example.org/ns/state" '
+    'xmlns:cp="https://accessibility.ubuntu.example.org/ns/component" '
+    'xmlns:attr="https://accessibility.windows.example.org/ns/attributes"'
+)
+VISIBLE = 'st:showing="true" st:visible="true" st:enabled="true"'
+
+UBUNTU_TREE = f"""<desktop-frame {NS}>
+  <panel>
+    <push-button name="OK" {VISIBLE} cp:screencoord="(10, 10)" cp:size="(20, 20)" attr:class="btn" attr:description="confirm">OK</push-button>
+    <label name="Hidden" st:showing="false" st:visible="true" st:enabled="true" cp:screencoord="(0, 0)" cp:size="(5, 5)"/>
+  </panel>
+  <push-button name="Offscreen" {VISIBLE} cp:screencoord="(-1, 5)" cp:size="(10, 10)"/>
+  <push-button name="ZeroSize" {VISIBLE} cp:screencoord="(5, 5)" cp:size="(0, 10)"/>
+  <entry name="Quote" {VISIBLE} cp:screencoord="(40, 40)" cp:size="(30, 10)">say "hi"</entry>
+  <panel name="NotKeptTag" {VISIBLE} cp:screencoord="(1, 1)" cp:size="(1, 1)"/>
+  <image {VISIBLE} image="true" cp:screencoord="(60, 60)" cp:size="(20, 20)"/>
+</desktop-frame>"""
+
+
+@pytest.fixture
+def root():
+    return ET.fromstring(UBUNTU_TREE)
+
+
+@pytest.fixture
+def noisy_png():
+    rng = np.random.default_rng(0)
+    arr = rng.integers(0, 255, size=(100, 100, 3), dtype=np.uint8)
+    buf = io.BytesIO()
+    Image.fromarray(arr).save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_find_leaf_nodes_empty():
+    assert find_leaf_nodes("") == []
+    assert find_leaf_nodes(None) == []
+
+
+def test_find_leaf_nodes():
+    tags = [n.get("name") for n in find_leaf_nodes(UBUNTU_TREE)]
+    assert tags == ["OK", "Hidden", "Offscreen", "ZeroSize", "Quote", "NotKeptTag", None]
+
+
+def test_filter_nodes_ubuntu(root):
+    names = [n.get("name") for n in filter_nodes(root)]
+    assert names == ["OK", "Quote"]
+
+
+def test_filter_nodes_check_image(root):
+    nodes = filter_nodes(root, check_image=True)
+    assert [n.tag for n in nodes] == ["push-button", "entry", "image"]
+
+
+def test_judge_node_invalid_platform(root):
+    with pytest.raises(ValueError, match="Invalid platform"):
+        judge_node(root, platform="macos")
+
+
+def test_judge_node_windows_only_requires_visible():
+    tree = (
+        '<root xmlns:st="https://accessibility.windows.example.org/ns/state" '
+        'xmlns:cp="https://accessibility.windows.example.org/ns/component">'
+        '<button name="Win" st:visible="true" st:enabled="true" '
+        'cp:screencoord="(1, 1)" cp:size="(5, 5)"/></root>'
+    )
+    node = ET.fromstring(tree)[0]
+    assert judge_node(node, platform="windows")
+    assert not judge_node(node, platform="ubuntu")  # ubuntu also needs "showing"
+
+
+def test_linearize_accessibility_tree(root):
+    lines = linearize_accessibility_tree(UBUNTU_TREE).split("\n")
+    assert lines[0].startswith("tag\tname\ttext")
+    assert lines[1] == "push-button\tOK\tOK\tbtn\tconfirm\t(10, 10)\t(20, 20)"
+    assert lines[2] == 'entry\tQuote\t"say ""hi"""\t\t\t(40, 40)\t(30, 10)'
+    assert len(lines) == 3
+
+
+def test_linearize_invalid_platform():
+    with pytest.raises(ValueError):
+        linearize_accessibility_tree(UBUNTU_TREE, platform="bsd")
+
+
+def test_draw_bounding_boxes(root, noisy_png):
+    nodes = filter_nodes(root)
+    marks, drew_nodes, text, image_bytes = draw_bounding_boxes(nodes, noisy_png)
+    assert marks == [[10, 10, 20, 20], [40, 40, 30, 10]]
+    assert [n.get("name") for n in drew_nodes] == ["OK", "Quote"]
+    lines = text.split("\n")
+    assert lines[0] == "index\ttag\tname\ttext"
+    assert lines[1] == "1\tpush-button\tOK\tOK"
+    assert lines[2] == '2\tentry\tQuote\t"say ""hi"""'
+    assert Image.open(io.BytesIO(image_bytes)).size == (100, 100)
+
+
+def test_draw_bounding_boxes_skips_uniform_regions(root):
+    buf = io.BytesIO()
+    Image.new("RGB", (100, 100), (255, 255, 255)).save(buf, format="PNG")
+    marks, drew_nodes, text, _ = draw_bounding_boxes(filter_nodes(root), buf.getvalue())
+    assert marks == [] and drew_nodes == []
+    assert text == "index\ttag\tname\ttext"
+
+
+def test_draw_bounding_boxes_downsampling_keeps_original_marks(root, noisy_png):
+    marks, _, _, image_bytes = draw_bounding_boxes(
+        filter_nodes(root), noisy_png, down_sampling_ratio=0.5
+    )
+    assert marks[0] == [10, 10, 20, 20]
+    assert Image.open(io.BytesIO(image_bytes)).size == (50, 50)
+
+
+def test_draw_bounding_boxes_invalid_platform(noisy_png):
+    with pytest.raises(ValueError):
+        draw_bounding_boxes([], noisy_png, platform="plan9")
+
+
+def test_tag_screenshot(noisy_png):
+    marks, drew_nodes, tagged, element_list = tag_screenshot(noisy_png, UBUNTU_TREE)
+    assert len(marks) == len(drew_nodes) == 3
+    assert isinstance(tagged, bytes)
+    assert element_list.split("\n")[3].startswith("3\timage")
+
+
+def test_trim_accessibility_tree():
+    short = "tag\tname"
+    assert trim_accessibility_tree(short, max_tokens=100) == short
+    long = "word " * 200
+    trimmed = trim_accessibility_tree(long, max_tokens=10)
+    assert trimmed.endswith("[...]\n")
+    assert len(trimmed) < len(long)
diff --git a/tests/benchmarks/test_setup_benchmark.py b/tests/benchmarks/test_setup_benchmark.py
new file mode 100644
index 000000000..2c548ea86
--- /dev/null
+++ b/tests/benchmarks/test_setup_benchmark.py
@@ -0,0 +1,98 @@
+import os
+from unittest.mock import patch
+
+import pytest
+
+from agentlab.benchmarks import setup_benchmark
+from agentlab.benchmarks.setup_benchmark import (
+    _ensure_repo,
+    _write_env_kv,
+    ensure_benchmark,
+    setup_miniwob,
+)
+
+
+def test_write_env_kv_creates_file(tmp_path):
+    env = tmp_path / "sub" / ".env"
+    _write_env_kv(env, "KEY", "value")
+    assert env.read_text() == "KEY=value\n"
+
+
+def test_write_env_kv_updates_and_dedupes(tmp_path):
+    env = tmp_path / ".env"
+    env.write_text("A=1\nKEY=old\nB=2\n  KEY=older\n")
+    _write_env_kv(env, "KEY", "new")
+    assert env.read_text() == "A=1\nKEY=new\nB=2\n"
+
+
+def test_write_env_kv_appends_when_missing(tmp_path):
+    env = tmp_path / ".env"
+    env.write_text("A=1\n")
+    _write_env_kv(env, "KEY", "v")
+    _write_env_kv(env, "KEY", "v")  # idempotent
+    assert env.read_text() == "A=1\nKEY=v\n"
+
+
+def test_write_env_kv_does_not_match_prefix_keys(tmp_path):
+    env = tmp_path / ".env"
+    env.write_text("KEYS=1\n")
+    _write_env_kv(env, "KEY", "2")
+    assert env.read_text() == "KEYS=1\nKEY=2\n"
+
+
+def test_ensure_repo_clones_and_checks_out(tmp_path):
+    clone_dir = tmp_path / "repo"
+    with patch.object(setup_benchmark.os, "system") as system:
+        _ensure_repo("https://example.com/r.git", clone_dir, commit="abc")
+    # clone_dir does not exist and no .git appears, so only the clone is attempted
+    assert system.call_count == 1
+    assert "git clone 'https://example.com/r.git'" in system.call_args.args[0]
+
+
+def test_ensure_repo_existing_repo_checks_out_commit(tmp_path):
+    clone_dir = tmp_path / "repo"
+    (clone_dir / ".git").mkdir(parents=True)
+    with patch.object(setup_benchmark.os, "system") as system:
+        _ensure_repo("https://example.com/r.git", clone_dir, commit="abc")
+    cmds = [c.args[0] for c in system.call_args_list]
+    assert len(cmds) == 2
+    assert "fetch --all --tags" in cmds[0]
+    assert "checkout abc" in cmds[1]
+
+
+def test_ensure_repo_existing_without_commit_does_nothing(tmp_path):
+    (tmp_path / "repo" / ".git").mkdir(parents=True)
+    with patch.object(setup_benchmark.os, "system") as system:
+        _ensure_repo("u", tmp_path / "repo")
+    system.assert_not_called()
+
+
+def test_setup_miniwob_writes_env(tmp_path, monkeypatch):
+    monkeypatch.delenv("MINIWOB_URL", raising=False)
+    with patch.object(setup_benchmark, "_ensure_repo") as ensure_repo:
+        url = setup_miniwob(tmp_path)
+    ensure_repo.assert_called_once()
+    assert url.startswith("file://") and url.endswith("/miniwob/html/miniwob/")
+    assert os.environ["MINIWOB_URL"] == url
+    assert (tmp_path / ".env").read_text() == f"MINIWOB_URL={url}\n"
+
+
+@pytest.mark.parametrize("name", ["miniwob", " MiniWob++ ", "miniwob-plusplus"])
+def test_ensure_benchmark_runs_setup_when_url_missing(tmp_path, monkeypatch, name):
+    monkeypatch.delenv("MINIWOB_URL", raising=False)
+    with patch.object(setup_benchmark, "setup_miniwob", return_value="url") as setup:
+        assert ensure_benchmark(name, tmp_path) == "url"
+    setup.assert_called_once_with(tmp_path)
+
+
+def test_ensure_benchmark_skips_when_url_set(tmp_path, monkeypatch):
+    monkeypatch.setenv("MINIWOB_URL", "http://localhost/miniwob/")
+    with patch.object(setup_benchmark, "setup_miniwob") as setup:
+        assert ensure_benchmark("miniwob", tmp_path) is None
+    setup.assert_not_called()
+
+
+def test_ensure_benchmark_other_benchmark_is_noop(tmp_path):
+    with patch.object(setup_benchmark, "setup_miniwob") as setup:
+        assert ensure_benchmark("webarena", tmp_path) is None
+    setup.assert_not_called()
diff --git a/tests/experiments/test_exp_utils.py b/tests/experiments/test_exp_utils.py
new file mode 100644
index 000000000..602aa74d5
--- /dev/null
+++ b/tests/experiments/test_exp_utils.py
@@ -0,0 +1,131 @@
+import signal
+import time
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from agentlab.experiments import exp_utils
+from agentlab.experiments.exp_utils import (
+    MockedExpArgs,
+    _episode_timeout,
+    add_dependencies,
+    hide_some_exp,
+    make_seeds,
+    order,
+    run_exp,
+    timeout_manager,
+)
+
+
+class FakeExpArgs:
+    def __init__(self, task_name):
+        self.env_args = SimpleNamespace(task_name=task_name)
+        self.exp_id = None
+        self.depends_on = ()
+
+    def make_id(self):
+        if self.exp_id is None:
+            self.exp_id = f"id_{self.env_args.task_name}"
+
+
+def test_episode_timeout_uses_max_steps():
+    exp = SimpleNamespace(env_args=SimpleNamespace(max_steps=10))
+    assert _episode_timeout(exp, avg_step_timeout=60) == 600
+
+
+def test_episode_timeout_defaults_without_max_steps():
+    exp = SimpleNamespace(env_args=SimpleNamespace())
+    assert _episode_timeout(exp) == 10 * 60 * 60
+
+
+def test_episode_timeout_respects_smaller_exp_timeout():
+    exp = SimpleNamespace(env_args=SimpleNamespace(max_steps=10), episode_timeout=42)
+    assert _episode_timeout(exp, avg_step_timeout=60) == 42
+    exp.episode_timeout = 10_000
+    assert _episode_timeout(exp, avg_step_timeout=60) == 600
+
+
+def test_timeout_manager_none_is_noop():
+    with timeout_manager(None):
+        pass
+
+
+def test_timeout_manager_raises_and_restores_handler():
+    previous = signal.getsignal(signal.SIGALRM)
+    with pytest.raises(TimeoutError):
+        with timeout_manager(seconds=1):
+            time.sleep(3)
+    assert signal.getsignal(signal.SIGALRM) == previous
+    assert signal.alarm(0) == 0  # no pending alarm
+
+
+def test_timeout_manager_float_is_rounded_up_to_one_second():
+    start = time.time()
+    with pytest.raises(TimeoutError):
+        with timeout_manager(seconds=0.1):
+            time.sleep(3)
+    assert time.time() - start >= 0.9
+
+
+def test_timeout_manager_no_timeout_when_fast():
+    with timeout_manager(seconds=5):
+        pass
+    assert signal.alarm(0) == 0
+
+
+def test_run_exp_calls_run():
+    exp = MockedExpArgs(exp_id=1)
+    with patch.object(exp_utils, "sleep"):
+        result = run_exp(exp, "dep1", "dep2")
+    assert result is exp
+    assert exp.start_time is not None and exp.end_time >= exp.start_time
+
+
+def test_mocked_exp_args_defaults():
+    exp = MockedExpArgs(exp_id=3, depends_on=[1, 2])
+    assert exp.exp_name == "exp_3"
+    assert exp.depends_on == [1, 2]
+    assert MockedExpArgs(exp_id=4).depends_on == []
+
+
+@pytest.mark.parametrize("deps", [None, {}, {"a": [], "b": []}])
+def test_add_dependencies_noop(deps):
+    exp_list = [FakeExpArgs("a"), FakeExpArgs("b")]
+    assert add_dependencies(exp_list, deps) is exp_list
+    assert all(e.depends_on == () for e in exp_list)
+    assert all(e.exp_id is None for e in exp_list)
+
+
+def test_add_dependencies_maps_task_names_to_exp_ids():
+    exp_list = [FakeExpArgs("a"), FakeExpArgs("b"), FakeExpArgs("c")]
+    out = add_dependencies(exp_list, {"a": [], "b": ["a"], "c": ["a", "b"]})
+    by_name = {e.env_args.task_name: e for e in out}
+    assert by_name["a"].depends_on == ()
+    assert by_name["b"].depends_on == ("id_a",)
+    assert by_name["c"].depends_on == ("id_a", "id_b")
+
+
+def test_add_dependencies_duplicate_task_names():
+    exp_list = [FakeExpArgs("a"), FakeExpArgs("a")]
+    with pytest.raises(ValueError, match="not unique"):
+        add_dependencies(exp_list, {"a": ["b"]})
+
+
+def test_add_dependencies_missing_task():
+    exp_list = [FakeExpArgs("a"), FakeExpArgs("b")]
+    with pytest.raises(ValueError, match="b is missing"):
+        add_dependencies(exp_list, {"a": ["b"]})
+
+
+@pytest.mark.parametrize(
+    "fn, args",
+    [(make_seeds, (3,)), (order, ([],)), (hide_some_exp, ("dir", lambda e: True, True))],
+)
+def test_deprecated_functions_raise(fn, args):
+    with pytest.raises(DeprecationWarning):
+        fn(*args)
+
+
+def test_results_dir_exists():
+    assert exp_utils.RESULTS_DIR.exists()
diff --git a/tests/llm/test_chat_api_unit.py b/tests/llm/test_chat_api_unit.py
new file mode 100644
index 000000000..5fd0d7e77
--- /dev/null
+++ b/tests/llm/test_chat_api_unit.py
@@ -0,0 +1,284 @@
+"""Unit tests for chat_api that never hit a real API (clients are mocked)."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import openai
+import pytest
+
+from agentlab.llm import chat_api
+from agentlab.llm.chat_api import (
+    AnthropicChatModel,
+    AnthropicModelArgs,
+    ChatModel,
+    ChatModelArgs,
+    CheatMiniWoBLLM,
+    CheatMiniWoBLLMArgs,
+    OpenAIChatModel,
+    OpenAIModelArgs,
+    RetryError,
+    SelfHostedModelArgs,
+    _extract_wait_time,
+    handle_error,
+    make_assistant_message,
+    make_system_message,
+    make_user_message,
+)
+from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage
+from agentlab.llm.tracking import set_tracker
+
+MESSAGES = [make_system_message("sys"), make_user_message("hello")]
+
+
+def _completion(contents=("hi",), prompt_tokens=10, completion_tokens=5, usage=True):
+    return SimpleNamespace(
+        usage=(
+            SimpleNamespace(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
+            if usage
+            else None
+        ),
+        choices=[
+            SimpleNamespace(message=SimpleNamespace(content=c), log_probs=[-0.5]) for c in contents
+        ],
+    )
+
+
+def _make_chat_model(responses, **kwargs):
+    client = MagicMock()
+    client.chat.completions.create.side_effect = responses
+    client_class = MagicMock(return_value=client)
+    kwargs.setdefault("min_retry_wait_time", 0)
+    model = ChatModel("test-model", client_class=client_class, **kwargs)
+    return model, client, client_class
+
+
+@pytest.fixture
+def no_sleep():
+    with patch.object(chat_api.time, "sleep") as sleep:
+        yield sleep
+
+
+def test_message_constructors():
+    assert make_system_message("a") == {"role": "system", "content": "a"}
+    assert make_user_message("b") == {"role": "user", "content": "b"}
+    assert make_assistant_message("c") == {"role": "assistant", "content": "c"}
+
+
+@pytest.mark.parametrize(
+    "msg, minimum, expected",
+    [
+        ("Rate limit. Please try again in 20s.", 5, 20.0),
+        ("Please try again in 1.5s", 5, 5),
+        ("Please try again in 90.25s", 60, 90.25),
+        ("no hint here", 7, 7),
+        ("", 60, 60),
+    ],
+)
+def test_extract_wait_time(msg, minimum, expected):
+    assert _extract_wait_time(msg, min_retry_wait_time=minimum) == expected
+
+
+def test_handle_error_reraises_non_openai_errors(no_sleep):
+    with pytest.raises(KeyError):
+        handle_error(KeyError("x"), 0, 1, 3)
+    no_sleep.assert_not_called()
+
+
+def test_handle_error_sleeps_and_returns_message(no_sleep):
+    err = openai.OpenAIError("try again in 3s")
+    assert handle_error(err, 0, 1, 3) == "try again in 3s"
+    no_sleep.assert_called_once_with(3.0)
+
+
+def test_cheat_llm_with_message_list():
+    llm = CheatMiniWoBLLMArgs().make_model()
+    assert isinstance(llm, CheatMiniWoBLLM)
+    answer = llm([make_system_message("s"), make_user_message("[13] button 'Click Me!'")])
+    assert answer["role"] == "assistant"
+    assert 'click("13")' in answer["content"]
+
+
+def test_cheat_llm_with_discussion():
+    llm = CheatMiniWoBLLM()
+    answer = llm(Discussion([SystemMessage("s"), HumanMessage("  [7] BUTTON 'ok'")]))
+    assert 'click("7")' in answer["content"]
+
+
+def test_cheat_llm_without_button():
+    with pytest.raises(Exception, match="Can't find the button's bid"):
+        CheatMiniWoBLLM()([make_system_message("s"), make_user_message("[1] link 'x'")])
+
+
+def test_cheat_llm_args_server_hooks_are_noops():
+    args = CheatMiniWoBLLMArgs()
+    assert args.prepare_server() is None and args.close_server() is None
+
+
+def test_chat_model_init_reads_env_and_pricing(monkeypatch):
+    monkeypatch.setenv("MY_KEY", "secret-from-env")
+    pricing = lambda: {"test-model": {"prompt": "0.5", "completion": "1"}}
+    model, _, client_class = _make_chat_model(
+        [], api_key_env_var="MY_KEY", pricing_func=pricing, client_args={"base_url": "u"}
+    )
+    client_class.assert_called_once_with(api_key="secret-from-env", base_url="u")
+    assert (model.input_cost, model.output_cost) == (0.5, 1.0)
+
+
+def test_chat_model_unknown_pricing_defaults_to_zero():
+    model, _, _ = _make_chat_model([], pricing_func=lambda: {})
+    assert (model.input_cost, model.output_cost) == (0.0, 0.0)
+
+
+def test_chat_model_requires_positive_retry():
+    with pytest.raises(AssertionError):
+        _make_chat_model([], max_retry=0)
+
+
+def test_chat_model_call_success_and_tracking():
+    model, client, _ = _make_chat_model([_completion()], temperature=0.2, max_tokens=50)
+    model.input_cost, model.output_cost = 1.0, 2.0
+    with set_tracker() as tracker:
+        answer = model(MESSAGES)
+    assert answer["role"] == "assistant" and answer["content"] == "hi"
+    kwargs = client.chat.completions.create.call_args.kwargs
+    assert kwargs == dict(
+        model="test-model",
+        messages=MESSAGES,
+        n=1,
+        temperature=0.2,
+        max_completion_tokens=50,
+        logprobs=False,
+    )
+    assert tracker.stats == {"input_tokens": 10, "output_tokens": 5, "cost": 20.0}
+    assert model.get_stats() == {"n_retry_llm": 1}
+    assert model.success and model.error_types == []
+
+
+def test_chat_model_multiple_samples_and_temperature_override():
+    model, client, _ = _make_chat_model([_completion(contents=("a", "b"))])
+    answers = model(MESSAGES, n_samples=2, temperature=0.9)
+    assert [a["content"] for a in answers] == ["a", "b"]
+    assert client.chat.completions.create.call_args.kwargs["temperature"] == 0.9
+
+
+def test_chat_model_log_probs():
+    model, _, _ = _make_chat_model([_completion()], log_probs=True)
+    assert model(MESSAGES)["log_probs"] == [-0.5]
+
+
+def test_chat_model_retries_on_openai_errors(no_sleep):
+    model, client, _ = _make_chat_model(
+        [openai.OpenAIError("try again in 2s"), _completion(usage=False), _completion()],
+        max_retry=3,
+    )
+    assert model(MESSAGES)["content"] == "hi"
+    assert model.retries == 3
+    assert model.error_types[0] == "try again in 2s"
+    assert "usage information" in model.error_types[1]
+    assert no_sleep.call_count == 2
+
+
+def test_chat_model_raises_retry_error(no_sleep):
+    model, _, _ = _make_chat_model([openai.OpenAIError("down")] * 2, max_retry=2)
+    with pytest.raises(RetryError, match="after 2 retries"):
+        model(MESSAGES)
+    assert not model.success
+    assert model.error_types == ["down", "down"]
+
+
+def test_chat_model_non_openai_error_propagates():
+    model, _, _ = _make_chat_model([ValueError("bad")])
+    with pytest.raises(ValueError, match="bad"):
+        model(MESSAGES)
+
+
+def test_openai_model_args_make_model(monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "fake")
+    with patch.object(chat_api.tracking, "get_model_info", side_effect=Exception("offline")):
+        model = OpenAIModelArgs(model_name="gpt-x", max_new_tokens=None).make_model()
+    assert isinstance(model, OpenAIChatModel)
+    assert model.max_tokens is chat_api.NOT_GIVEN
+    assert model.api_key == "fake"
+
+
+def test_self_hosted_unknown_backend():
+    with pytest.raises(ValueError, match="Backend foo is not supported"):
+        SelfHostedModelArgs(model_name="m", backend="foo").make_model()
+
+
+def test_chat_model_args_is_deprecated():
+    with pytest.warns(DeprecationWarning):
+        args = ChatModelArgs(model_name="legacy")
+    assert args.make_model() is None
+
+
+@pytest.fixture
+def anthropic_client():
+    client = MagicMock()
+    with patch.object(chat_api.anthropic, "Anthropic", return_value=client) as cls:
+        yield client, cls
+
+
+def _anthropic_response(text="ok"):
+    return SimpleNamespace(
+        content=[SimpleNamespace(text=text)],
+        usage=SimpleNamespace(input_tokens=3, output_tokens=4),
+    )
+
+
+def test_anthropic_model_args_make_model(anthropic_client, monkeypatch):
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "fake")
+    _, cls = anthropic_client
+    model = AnthropicModelArgs(
+        model_name="claude-x", max_new_tokens=12, temperature=0.1
+    ).make_model()
+    assert isinstance(model, AnthropicChatModel)
+    assert model.max_tokens == 12 and model.temperature == 0.1
+    cls.assert_called_once_with(api_key="fake")
+
+
+def test_anthropic_call_converts_messages(anthropic_client):
+    client, _ = anthropic_client
+    client.messages.create.return_value = _anthropic_response("done")
+    model = AnthropicChatModel("claude-x", api_key="k", max_tokens=8, temperature=0.4)
+    with set_tracker() as tracker:
+        answer = model(MESSAGES + [make_assistant_message("prev")])
+    assert answer["content"] == "done"
+    kwargs = client.messages.create.call_args.kwargs
+    assert kwargs["system"] == "sys"
+    assert kwargs["messages"] == [
+        {"role": "user", "content": "hello"},
+        {"role": "assistant", "content": "prev"},
+    ]
+    assert kwargs["max_tokens"] == 8 and kwargs["temperature"] == 0.4
+    assert tracker.stats["input_tokens"] == 3 and tracker.stats["output_tokens"] == 4
+
+
+def test_anthropic_call_without_system(anthropic_client):
+    client, _ = anthropic_client
+    client.messages.create.return_value = _anthropic_response()
+    model = AnthropicChatModel("claude-x", api_key="k")
+    with set_tracker():
+        model([make_user_message("hi")], temperature=0.0)
+    kwargs = client.messages.create.call_args.kwargs
+    assert "system" not in kwargs
+    assert kwargs["temperature"] == 0.0
+
+
+def test_anthropic_retry_then_success(anthropic_client, no_sleep):
+    client, _ = anthropic_client
+    client.messages.create.side_effect = [RuntimeError("overloaded"), _anthropic_response("ok")]
+    model = AnthropicChatModel("claude-x", api_key="k", max_retry=2)
+    with set_tracker():
+        assert model(MESSAGES)["content"] == "ok"
+    no_sleep.assert_called_once_with(60)
+
+
+def test_anthropic_retry_exhausted(anthropic_client, no_sleep):
+    client, _ = anthropic_client
+    client.messages.create.side_effect = RuntimeError("overloaded")
+    model = AnthropicChatModel("claude-x", api_key="k", max_retry=3)
+    with pytest.raises(RuntimeError, match="overloaded"):
+        model(MESSAGES)
+    assert client.messages.create.call_count == 3
+    assert no_sleep.call_count == 2
diff --git a/tests/llm/test_huggingface_chat_model.py b/tests/llm/test_huggingface_chat_model.py
new file mode 100644
index 000000000..b5ea4b8a9
--- /dev/null
+++ b/tests/llm/test_huggingface_chat_model.py
@@ -0,0 +1,143 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from agentlab.llm import huggingface_utils
+from agentlab.llm.huggingface_utils import HFBaseChatModel, _prepend_system_to_first_user
+from agentlab.llm.llm_utils import AIMessage, Discussion, HumanMessage, SystemMessage
+from agentlab.llm.prompt_templates import STARCHAT_PROMPT_TEMPLATE
+
+MESSAGES = [
+    {"role": "system", "content": "be nice"},
+    {"role": "user", "content": "hi"},
+]
+
+
+class FakeGPT2Tokenizer:
+    pass
+
+
+def make_model(tokenizer=None, model_name="org/model", log_probs=False, n_retry=2):
+    tokenizer = tokenizer if tokenizer is not None else MagicMock()
+    with (
+        patch.object(huggingface_utils, "AutoTokenizer") as auto,
+        patch.object(huggingface_utils, "GPT2TokenizerFast", FakeGPT2Tokenizer),
+    ):
+        auto.from_pretrained.return_value = tokenizer
+        model = HFBaseChatModel(model_name, None, n_retry_server=n_retry, log_probs=log_probs)
+    model.temperature = 0.3
+    model.llm = MagicMock(return_value="answer")
+    return model, auto
+
+
+def test_prepend_system_to_first_user():
+    out = _prepend_system_to_first_user([dict(m) for m in MESSAGES])
+    assert out == [{"role": "user", "content": "be nice\nhi"}]
+
+
+def test_prepend_system_no_system_message():
+    messages = [{"role": "user", "content": "hi"}]
+    assert _prepend_system_to_first_user(messages) == [{"role": "user", "content": "hi"}]
+
+
+def test_prepend_system_with_column_remap():
+    messages = [{"type": "system", "text": "S"}, {"type": "human", "text": "U"}]
+    out = _prepend_system_to_first_user(
+        messages, column_remap={"HumanMessage": "human", "role": "type", "text": "text"}
+    )
+    assert out == [{"type": "human", "text": "S\nU"}]
+
+
+def test_init_uses_base_model_name_for_tokenizer():
+    with patch.object(huggingface_utils, "AutoTokenizer") as auto:
+        auto.from_pretrained.return_value = MagicMock()
+        HFBaseChatModel("org/finetuned", "org/base", n_retry_server=1, log_probs=False)
+    auto.from_pretrained.assert_called_once_with("org/base")
+
+
+def test_call_with_chat_template():
+    tokenizer = MagicMock()
+    tokenizer.apply_chat_template.return_value = "PROMPT"
+    model, auto = make_model(tokenizer)
+    auto.from_pretrained.assert_called_once_with("org/model")
+    assert model.tokenizer is tokenizer and model._llm_type() == "huggingface"
+
+    response = model(MESSAGES)
+    assert isinstance(response, AIMessage)
+    assert response["content"] == "answer"
+    tokenizer.apply_chat_template.assert_called_once_with(MESSAGES, tokenize=False)
+    model.llm.assert_called_once_with("PROMPT", temperature=0.3)
+
+
+def test_call_merges_discussion_and_overrides_temperature():
+    tokenizer = MagicMock()
+    tokenizer.apply_chat_template.return_value = "PROMPT"
+    model, _ = make_model(tokenizer)
+    msg = HumanMessage("a")
+    msg.add_text("b")
+    model(Discussion([SystemMessage("s"), msg]), temperature=0.9)
+    merged = tokenizer.apply_chat_template.call_args.args[0]
+    assert merged[1]["content"] == "a\nb"
+    model.llm.assert_called_once_with("PROMPT", temperature=0.9)
+
+
+def test_call_retries_chat_template_without_system_role():
+    tokenizer = MagicMock()
+    tokenizer.apply_chat_template.side_effect = [
+        Exception("Conversation roles must alternate user/assistant"),
+        "PROMPT",
+    ]
+    model, _ = make_model(tokenizer)
+    model([dict(m) for m in MESSAGES])
+    retried = tokenizer.apply_chat_template.call_args_list[1].args[0]
+    assert retried == [{"role": "user", "content": "be nice\nhi"}]
+
+
+def test_call_chat_template_other_error_raises():
+    tokenizer = MagicMock()
+    tokenizer.apply_chat_template.side_effect = ValueError("boom")
+    model, _ = make_model(tokenizer)
+    with pytest.raises(ValueError, match="boom"):
+        model(MESSAGES)
+
+
+def test_gpt2_tokenizer_falls_back_to_prompt_template():
+    model, _ = make_model(FakeGPT2Tokenizer(), model_name="bigcode/starcoder")
+    assert model.tokenizer is None
+    assert model.prompt_template is STARCHAT_PROMPT_TEMPLATE
+    model(MESSAGES)
+    prompt = model.llm.call_args.args[0]
+    assert prompt == STARCHAT_PROMPT_TEMPLATE.construct_prompt(MESSAGES)
+
+
+def test_multiple_samples():
+    model, _ = make_model()
+    model.llm.side_effect = ["a", "b", "c"]
+    responses = model(MESSAGES, n_samples=3)
+    assert [r["content"] for r in responses] == ["a", "b", "c"]
+
+
+def test_server_retry_then_success():
+    model, _ = make_model(n_retry=3)
+    model.llm.side_effect = [RuntimeError("down"), RuntimeError("down"), "ok"]
+    with patch.object(huggingface_utils.time, "sleep") as sleep:
+        assert model(MESSAGES)["content"] == "ok"
+    assert sleep.call_count == 2
+
+
+def test_server_retry_exhausted():
+    model, _ = make_model(n_retry=2)
+    model.llm.side_effect = RuntimeError("down")
+    with patch.object(huggingface_utils.time, "sleep"):
+        with pytest.raises(RuntimeError, match="down"):
+            model(MESSAGES)
+    assert model.llm.call_count == 2
+
+
+def test_log_probs():
+    model, _ = make_model(log_probs=True)
+    model.llm.return_value = SimpleNamespace(generated_text="txt", details={"lp": [-0.1]})
+    response = model(MESSAGES)
+    assert response["content"] == "txt"
+    assert response["log_probs"] == {"lp": [-0.1]}
diff --git a/tests/llm/test_llm_utils_messages.py b/tests/llm/test_llm_utils_messages.py
new file mode 100644
index 000000000..c6135a2e6
--- /dev/null
+++ b/tests/llm/test_llm_utils_messages.py
@@ -0,0 +1,216 @@
+import base64
+import io
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from agentlab.llm import llm_utils
+from agentlab.llm.llm_utils import (
+    AIMessage,
+    BaseMessage,
+    Discussion,
+    HumanMessage,
+    ParseError,
+    SystemMessage,
+    extract_html_tags,
+    get_tokenizer,
+    image_to_jpg_base64_url,
+    image_to_png_base64_url,
+    messages_to_dict,
+    parse_html_tags,
+    parse_html_tags_raise,
+    retry_multiple,
+    supports_tool_calling_for_openrouter,
+)
+
+
+def _decode(url, prefix):
+    assert url.startswith(prefix)
+    return Image.open(io.BytesIO(base64.b64decode(url[len(prefix) :])))
+
+
+def test_messages_to_dict():
+    out = messages_to_dict([{"role": "user", "content": "hi"}, "raw"])
+    assert isinstance(out, Discussion)
+    assert out[0] == {"role": "user", "content": "hi"}
+    assert out[1] == {"role": "<unknown role>", "content": "raw"}
+    assert all(isinstance(m, BaseMessage) for m in out)
+
+
+def test_messages_to_dict_empty_and_invalid():
+    assert len(messages_to_dict([])) == 0
+    with pytest.raises(ValueError, match="Unknown message type"):
+        messages_to_dict([42])
+
+
+def test_extract_html_tags():
+    text = "<a> one </a> noise <b>two</b><a>\nthree\n</a>"
+    assert extract_html_tags(text, ["a", "b", "c"]) == {"a": ["one", "three"], "b": ["two"]}
+    assert extract_html_tags("", ["a"]) == {}
+
+
+def test_parse_html_tags_valid_and_optional():
+    content, valid, msg = parse_html_tags(
+        "<think>t</think><action>a</action>", ["action"], ["think", "memo"]
+    )
+    assert valid and msg == ""
+    assert content == {"think": "t", "action": "a"}
+
+
+def test_parse_html_tags_missing_and_duplicate():
+    content, valid, msg = parse_html_tags(
+        "<action>a</action><action>b</action>", ["action", "think"]
+    )
+    assert not valid
+    assert "Missing the key <think>" in msg
+    assert "multiple instances of the key action" in msg
+    assert content["action"] == "a"
+
+
+def test_parse_html_tags_merge_multiple():
+    content, valid, _ = parse_html_tags("<a>1</a><a>2</a>", ["a"], merge_multiple=True)
+    assert valid and content == {"a": "1\n2"}
+
+
+def test_parse_html_tags_raise():
+    assert parse_html_tags_raise("<a>1</a>", ["a"]) == {"a": "1"}
+    with pytest.raises(ParseError, match="Missing the key <b>"):
+        parse_html_tags_raise("<a>1</a>", ["a", "b"])
+
+
+def _parser(text):
+    if "good" not in text:
+        raise ParseError(f"bad: {text}")
+    return text.upper()
+
+
+def test_retry_multiple_returns_parsed_answers():
+    chat = MagicMock(return_value=[AIMessage("bad"), AIMessage("good 1"), AIMessage("good 2")])
+    messages = Discussion([HumanMessage("q")])
+    parsed, tries = retry_multiple(chat, messages, n_retry=2, parser=_parser, num_samples=3)
+    assert parsed == ["GOOD 1", "GOOD 2"] and tries == 0
+    chat.assert_called_once_with(messages, n_samples=3)
+    assert len(messages) == 2  # only the first answer is appended
+
+
+def test_retry_multiple_single_answer_retry_then_success():
+    chat = MagicMock(side_effect=[AIMessage("bad"), AIMessage("good")])
+    messages = Discussion([HumanMessage("q")])
+    parsed, tries = retry_multiple(chat, messages, n_retry=3, parser=_parser)
+    assert parsed == ["GOOD"] and tries == 1
+    assert [m["role"] for m in messages] == ["user", "assistant", "user", "assistant"]
+    assert "bad: bad" in messages[2]["content"]
+
+
+def test_retry_multiple_exhausted():
+    chat = MagicMock(return_value=AIMessage("bad"))
+    with pytest.raises(ParseError, match="after 2 retries"):
+        retry_multiple(chat, Discussion([HumanMessage("q")]), n_retry=2, parser=_parser, log=False)
+    assert chat.call_count == 2
+
+
+def test_image_to_base64_urls_from_rgba_array():
+    arr = np.zeros((4, 6, 4), dtype=np.uint8)
+    arr[..., 0] = 255
+    arr[..., 3] = 128
+    jpg = _decode(image_to_jpg_base64_url(arr), "data:image/jpeg;base64,")
+    png = _decode(image_to_png_base64_url(arr), "data:image/png;base64,")
+    assert jpg.size == png.size == (6, 4)
+    assert png.mode == "RGB"
+    assert png.getpixel((0, 0)) == (255, 0, 0)
+
+
+def test_image_to_png_from_pil_grayscale():
+    img = Image.new("L", (3, 3), 7)
+    assert _decode(image_to_png_base64_url(img), "data:image/png;base64,").getpixel((1, 1)) == 7
+
+
+def test_base_message_str_and_content_helpers():
+    msg = HumanMessage("hello")
+    assert str(msg) == "hello"
+    assert msg.to_markdown() == "\n```\nhello\n```\n"
+    msg.add_text("world")
+    msg.add_image("data:image/png;base64,xyz", detail="low")
+    assert msg["content"][0] == {"type": "text", "text": "hello"}
+    assert msg["content"][2] == {
+        "type": "image_url",
+        "image_url": {"url": "data:image/png;base64,xyz", "detail": "low"},
+    }
+    assert str(msg) == "hello\nworld"
+    md = msg.to_markdown()
+    assert "![image](data:image/png;base64,xyz)" in md and "```\nworld\n```" in md
+
+
+def test_base_message_add_image_array_and_warning():
+    msg = SystemMessage([{"type": "text", "text": "a"}])
+    msg.add_image(np.zeros((2, 2, 3), dtype=np.uint8))
+    assert msg["content"][1]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+    assert msg.__str__(warn_if_image=True) == "a"
+    assert msg["role"] == "system"
+
+
+def test_ai_message_log_probs_and_kwargs():
+    assert AIMessage("x", log_probs=[1])["log_probs"] == [1]
+    with pytest.raises(ValueError, match="Invalid attributes"):
+        BaseMessage("tool", "x", tool_call_id="1")
+
+
+def test_discussion_api():
+    d = Discussion(SystemMessage("sys"))
+    assert len(d) == 1
+    d.add_message(role="user", content="q")
+    d.append({"role": "assistant", "content": "a"})
+    d.add_text("more")
+    d.add_content("text", "even more")
+    d.add_image("data:image/png;base64,abc")
+    assert d.last_message["role"] == "assistant"
+    assert [m["role"] for m in d] == ["system", "user", "assistant"]
+    assert isinstance(d[2], BaseMessage)
+
+    assert d.to_string() == "sys\nq\na\nmore\neven more"
+    assert d[2]["content"][0] == {"type": "text", "text": "a\nmore\neven more"}
+    assert d.to_openai() is d.messages
+    md = d.to_markdown()
+    assert md.startswith("Message 0\n") and "Message 2" in md
+
+
+def test_discussion_empty():
+    d = Discussion()
+    assert len(d) == 0 and str(d) == ""
+    with pytest.raises(IndexError):
+        d.last_message
+
+
+def test_get_tokenizer_falls_back_to_gpt4():
+    with patch.object(llm_utils, "AutoTokenizer") as auto:
+        auto.from_pretrained.side_effect = OSError("not found")
+        enc = get_tokenizer("definitely/not-a-model")
+    assert enc.name == llm_utils.tiktoken.encoding_for_model("gpt-4").name
+
+
+def test_get_tokenizer_uses_huggingface():
+    with patch.object(llm_utils, "AutoTokenizer") as auto:
+        auto.from_pretrained.return_value = "hf-tokenizer"
+        assert get_tokenizer("org/model") == "hf-tokenizer"
+
+
+@pytest.mark.parametrize(
+    "message, expected", [({"tool_calls": [{"id": "1"}]}, True), ({"content": "no"}, False)]
+)
+def test_supports_tool_calling_for_openrouter(message, expected):
+    client = MagicMock()
+    client.chat.completions.create.return_value.to_dict.return_value = {
+        "choices": [{"message": message}]
+    }
+    with patch("openai.Client", return_value=client):
+        assert supports_tool_calling_for_openrouter("org/model") is expected
+    assert client.chat.completions.create.call_args.kwargs["tool_choice"] == "required"
+
+
+def test_supports_tool_calling_for_openrouter_error_defaults_true():
+    client = MagicMock()
+    client.chat.completions.create.side_effect = RuntimeError("network down")
+    with patch("openai.Client", return_value=client):
+        assert supports_tool_calling_for_openrouter("org/model") is True
diff --git a/tests/llm/test_prompt_templates.py b/tests/llm/test_prompt_templates.py
new file mode 100644
index 000000000..62bd5e7a9
--- /dev/null
+++ b/tests/llm/test_prompt_templates.py
@@ -0,0 +1,78 @@
+import pytest
+
+from agentlab.llm.prompt_templates import (
+    MODEL_PREFIX_TO_PROMPT_TEMPLATES,
+    STARCHAT_PROMPT_TEMPLATE,
+    PromptTemplate,
+    get_prompt_template,
+)
+
+
+@pytest.fixture
+def template():
+    return PromptTemplate(
+        system="S:{input}|", human="H:{input}|", ai="A:{input}|", prompt_end="END"
+    )
+
+
+def test_format_message_each_role(template):
+    assert template.format_message({"role": "system", "content": "x"}) == "S:x|"
+    assert template.format_message({"role": "user", "content": "y"}) == "H:y|"
+    assert template.format_message({"role": "assistant", "content": "z"}) == "A:z|"
+
+
+def test_format_message_unknown_role_raises(template):
+    with pytest.raises(ValueError, match="tool"):
+        template.format_message({"role": "tool", "content": "x"})
+
+
+def test_construct_prompt_concatenates_and_appends_end(template):
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "hello"},
+    ]
+    assert template.construct_prompt(messages) == "S:sys|H:hi|A:hello|END"
+
+
+def test_construct_prompt_empty_messages(template):
+    assert template.construct_prompt([]) == "END"
+
+
+@pytest.mark.parametrize(
+    "messages",
+    [
+        ["not a dict"],
+        [{"role": "user"}],
+        [{"content": "missing role"}],
+    ],
+)
+def test_construct_prompt_rejects_non_openai_format(template, messages):
+    with pytest.raises(ValueError, match="openai format"):
+        template.construct_prompt(messages)
+
+
+def test_default_prompt_end_is_empty():
+    t = PromptTemplate(system="{input}", human="{input}", ai="{input}")
+    assert t.construct_prompt([{"role": "user", "content": "a"}]) == "a"
+
+
+def test_starchat_template_format():
+    prompt = STARCHAT_PROMPT_TEMPLATE.construct_prompt(
+        [{"role": "system", "content": "be nice"}, {"role": "user", "content": "hi"}]
+    )
+    assert prompt == "<|system|>\nbe nice<|end|>\n<|user|>\nhi<|end|>\n<|assistant|>"
+
+
+@pytest.mark.parametrize("model_name", ["bigcode/starcoder", "HuggingFaceH4/starchat-beta"])
+def test_get_prompt_template_by_substring(model_name):
+    assert get_prompt_template(model_name) is STARCHAT_PROMPT_TEMPLATE
+
+
+def test_get_prompt_template_unknown_model():
+    with pytest.raises(NotImplementedError, match="llama"):
+        get_prompt_template("meta/llama")
+
+
+def test_mapping_keys():
+    assert set(MODEL_PREFIX_TO_PROMPT_TEMPLATES) == {"starcoder", "starchat"}
diff --git a/tests/llm/test_tracking_pricing.py b/tests/llm/test_tracking_pricing.py
new file mode 100644
index 000000000..c04b551bf
--- /dev/null
+++ b/tests/llm/test_tracking_pricing.py
@@ -0,0 +1,310 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from agentlab.llm import tracking
+from agentlab.llm.tracking import (
+    LLMTracker,
+    Stats,
+    TrackAPIPricingMixin,
+    _remove_version_suffix,
+    cost_tracker_decorator,
+    get_pricing_anthropic,
+    get_pricing_litellm,
+    get_pricing_openai,
+    get_pricing_openrouter,
+    set_tracker,
+)
+
+
+class Usage:
+    """Mimics pydantic usage objects: attribute access + iteration over (key, value)."""
+
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+    def __iter__(self):
+        return iter(self.__dict__.items())
+
+
+class DummyModel(TrackAPIPricingMixin):
+    def __init__(self, model_name="dummy-model", pricing_api="openai", response=None):
+        self.model_name = model_name
+        self.response = response
+        self.init_pricing_tracker(pricing_api=pricing_api)
+
+    def _call_api(self, *args, **kwargs):
+        return self.response
+
+    def _parse_response(self, response):
+        return "parsed"
+
+
+@pytest.fixture
+def no_litellm_info():
+    with patch.object(tracking, "get_model_info", side_effect=Exception("unknown model")) as m:
+        yield m
+
+
+def test_llm_tracker_accumulates_with_suffix():
+    tracker = LLMTracker(suffix="agent")
+    tracker(10, 5, 0.1)
+    tracker(1, 1, 0.2)
+    assert tracker.stats == {
+        "input_tokens_agent": 11,
+        "output_tokens_agent": 6,
+        "cost_agent": pytest.approx(0.3),
+    }
+    other = LLMTracker()
+    other.add_tracker(tracker)
+    assert other.stats == {"input_tokens": 11, "output_tokens": 6, "cost": pytest.approx(0.3)}
+    assert "input_tokens=11" in repr(other)
+
+
+def test_set_tracker_nested_propagates_to_parent():
+    with set_tracker("outer") as outer:
+        outer(1, 1, 1.0)
+        with set_tracker("inner") as inner:
+            inner(2, 3, 0.5)
+        assert inner.stats["cost_inner"] == 0.5
+    assert outer.stats == {"input_tokens_outer": 3, "output_tokens_outer": 4, "cost_outer": 1.5}
+
+
+def test_cost_tracker_decorator_updates_agent_info_stats():
+    class Agent:
+        @cost_tracker_decorator
+        def get_action(self, obs):
+            tracking.TRACKER.instance(4, 2, 0.25)
+            return "noop()", {"stats": {"existing": 1}}
+
+    action, info = Agent().get_action({})
+    assert action == "noop()"
+    assert info["stats"] == {"existing": 1, "input_tokens": 4, "output_tokens": 2, "cost": 0.25}
+
+
+@pytest.mark.parametrize(
+    "name, expected",
+    [
+        ("anthropic.claude-3-sonnet-v1:0", "claude-3-sonnet"),
+        ("anthropic.claude-v2", "claude"),
+        ("claude-instant-v1.2", "claude-instant"),
+        ("gpt-4o", "gpt-4o"),
+    ],
+)
+def test_remove_version_suffix(name, expected):
+    assert _remove_version_suffix(name) == expected
+
+
+def test_get_pricing_openai_pairs_prompt_and_completion():
+    fake = SimpleNamespace(
+        MODEL_COST_PER_1K_TOKENS={"gpt-x": 1.0, "gpt-x-completion": 2.0, "lonely": 3.0}
+    )
+    with patch.object(tracking, "openai_info", fake):
+        assert get_pricing_openai() == {"gpt-x": {"prompt": 0.001, "completion": 0.002}}
+
+
+def test_get_pricing_openai_without_langchain():
+    with patch.object(tracking, "openai_info", None):
+        assert get_pricing_openai() == {}
+
+
+def test_get_pricing_anthropic():
+    fake = SimpleNamespace(
+        MODEL_COST_PER_1K_INPUT_TOKENS={"anthropic.claude-3-v1:0": 3.0},
+        MODEL_COST_PER_1K_OUTPUT_TOKENS={"anthropic.claude-3-v1:0": 15.0, "anthropic.other": 1.0},
+    )
+    with patch.object(tracking, "bedrock_anthropic_callback", fake):
+        assert get_pricing_anthropic() == {
+            "claude-3": {"prompt": 0.003, "completion": 0.015},
+            "other": {"completion": 0.001},
+        }
+
+
+def test_get_pricing_anthropic_without_langchain():
+    with patch.object(tracking, "bedrock_anthropic_callback", None):
+        assert get_pricing_anthropic() == {}
+
+
+def test_get_pricing_litellm():
+    info = {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6}
+    with patch.object(tracking, "get_model_info", return_value=info):
+        assert get_pricing_litellm("m") == {"m": {"prompt": 1e-6, "completion": 2e-6}}
+
+
+def test_get_pricing_litellm_unknown_model(no_litellm_info):
+    assert get_pricing_litellm("m") == {"m": {"prompt": 0.0, "completion": 0.0}}
+
+
+def test_get_pricing_openrouter(monkeypatch):
+    monkeypatch.setenv("OPENROUTER_API_KEY", "fake-key")
+    response = MagicMock(status_code=200)
+    response.json.return_value = {
+        "data": [{"id": "org/model", "pricing": {"prompt": "0.1", "completion": "0.2"}}]
+    }
+    get_pricing_openrouter.cache_clear()
+    try:
+        with patch.object(tracking.requests, "get", return_value=response) as get:
+            assert get_pricing_openrouter() == {"org/model": {"prompt": 0.1, "completion": 0.2}}
+        assert get.call_args.kwargs["headers"] == {"Authorization": "Bearer fake-key"}
+    finally:
+        get_pricing_openrouter.cache_clear()
+
+
+def test_get_pricing_openrouter_http_error(monkeypatch):
+    monkeypatch.setenv("OPENROUTER_API_KEY", "fake-key")
+    get_pricing_openrouter.cache_clear()
+    try:
+        with patch.object(tracking.requests, "get", return_value=MagicMock(status_code=500)):
+            with pytest.raises(ValueError, match="Failed to get model metadata"):
+                get_pricing_openrouter()
+    finally:
+        get_pricing_openrouter.cache_clear()
+
+
+def test_get_pricing_openrouter_requires_key(monkeypatch):
+    monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
+    get_pricing_openrouter.cache_clear()
+    with pytest.raises(AssertionError):
+        get_pricing_openrouter()
+
+
+def test_set_pricing_attributes_from_provider():
+    with patch.object(
+        tracking, "get_pricing_openai", return_value={"m": {"prompt": "1", "completion": "2"}}
+    ):
+        model = DummyModel(model_name="m", pricing_api="openai")
+    assert (model.input_cost, model.output_cost) == (1.0, 2.0)
+    assert isinstance(model.stats, Stats)
+
+
+def test_set_pricing_attributes_falls_back_to_litellm():
+    info = {"input_cost_per_token": 0.5, "output_cost_per_token": 0.25}
+    with patch.object(tracking, "get_model_info", return_value=info):
+        model = DummyModel(pricing_api="unsupported")
+    assert (model.input_cost, model.output_cost) == (0.5, 0.25)
+
+
+def test_set_pricing_attributes_unknown_everywhere(no_litellm_info):
+    with patch.object(tracking, "get_pricing_anthropic", return_value={}):
+        model = DummyModel(pricing_api="anthropic")
+    assert (model.input_cost, model.output_cost) == (0.0, 0.0)
+
+
+@pytest.fixture
+def priced_model(no_litellm_info):
+    model = DummyModel(pricing_api=None)
+    model.input_cost, model.output_cost = 1.0, 2.0
+    return model
+
+
+@pytest.mark.parametrize(
+    "response, expected",
+    [
+        (SimpleNamespace(usage=SimpleNamespace(input_tokens=3, output_tokens=4)), (3, 4)),
+        (SimpleNamespace(usage=SimpleNamespace(prompt_tokens=5, completion_tokens=6)), (5, 6)),
+        ({"usage": {"prompt_tokens": 7, "completion_tokens": 8}}, (7, 8)),
+        ({"usage": {"input_tokens": 1}}, (0, 0)),
+        (SimpleNamespace(), (0, 0)),
+    ],
+)
+def test_get_tokens_counts_from_response(priced_model, response, expected):
+    assert priced_model.get_tokens_counts_from_response(response) == expected
+
+
+def test_update_pricing_tracker(priced_model):
+    response = SimpleNamespace(usage=SimpleNamespace(input_tokens=10, output_tokens=5))
+    with set_tracker() as tracker:
+        priced_model.update_pricing_tracker(response)
+    assert tracker.stats == {"input_tokens": 10, "output_tokens": 5, "cost": 20.0}
+
+
+def test_effective_cost_anthropic(priced_model):
+    priced_model._pricing_api = "anthropic"
+    usage = SimpleNamespace(
+        input_tokens=100, output_tokens=10, cache_input_tokens=50, cache_creation_input_tokens=20
+    )
+    # 100*1 + 10*2 + 50*0.1 + 20*1.25
+    assert priced_model.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx(150)
+
+
+def test_effective_cost_openai_chat_completion(priced_model):
+    priced_model._pricing_api = "openai"
+    usage = SimpleNamespace(
+        prompt_tokens=100,
+        completion_tokens=10,
+        prompt_tokens_details=SimpleNamespace(cached_tokens=40),
+    )
+    # 60 new * 1 + 40 cached * 0.5 + 10 * 2
+    assert priced_model.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx(100)
+    usage.prompt_tokens_details = None
+    assert priced_model.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx(120)
+
+
+def test_effective_cost_openai_responses_api(priced_model):
+    priced_model._pricing_api = "openai"
+    usage = SimpleNamespace(
+        input_tokens=100, output_tokens=10, input_tokens_details=SimpleNamespace(cached_tokens=20)
+    )
+    assert priced_model.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx(110)
+    usage.input_tokens_details = None
+    assert priced_model.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx(120)
+
+
+def test_effective_cost_openai_no_usage(priced_model):
+    priced_model._pricing_api = "openai"
+    assert priced_model.get_effective_cost(SimpleNamespace()) == 0.0
+
+
+def test_effective_cost_litellm(priced_model):
+    priced_model._pricing_api = "litellm"
+    with patch.object(tracking, "completion_cost", return_value=0.42) as cost:
+        assert priced_model.get_effective_cost("resp") == 0.42
+    cost.assert_called_once_with("resp")
+
+
+def test_effective_cost_unsupported_provider(priced_model):
+    assert priced_model.get_effective_cost(SimpleNamespace()) == 0.0
+
+
+def test_call_updates_stats_and_tracker(no_litellm_info):
+    usage = Usage(
+        prompt_tokens=100,
+        completion_tokens=10,
+        total_tokens=110,
+        prompt_tokens_details=SimpleNamespace(cached_tokens=40),
+        model="ignored-non-numeric",
+    )
+    model = DummyModel(pricing_api="openai", response=SimpleNamespace(usage=usage))
+    model.input_cost, model.output_cost = 1.0, 2.0
+    with set_tracker() as tracker:
+        assert model() == "parsed"
+        assert model() == "parsed"
+    stats = model.stats.stats_dict
+    assert stats["usage_prompt_tokens"] == 200
+    assert stats["usage_completion_tokens"] == 20
+    assert stats["usage_cached_tokens"] == 80
+    assert stats["n_api_calls"] == 2
+    assert stats["effective_cost"] == pytest.approx(200)
+    assert "usage_model" not in stats
+    assert tracker.stats["cost"] == pytest.approx(240)
+
+    model.reset_stats()
+    assert dict(model.stats.stats_dict) == {}
+
+
+def test_call_with_responses_api_details(no_litellm_info):
+    usage = Usage(
+        input_tokens=10, output_tokens=1, input_tokens_details=SimpleNamespace(cached_tokens=4)
+    )
+    model = DummyModel(pricing_api=None, response=SimpleNamespace(usage=usage))
+    model()
+    assert model.stats.stats_dict["usage_cached_tokens"] == 4
+    assert model.stats.stats_dict["effective_cost"] == 0.0
+
+
+def test_stats_increment():
+    stats = Stats()
+    stats.increment_stats_dict({"a": 1, "b": 2.5})
+    stats.increment_stats_dict({"a": 1})
+    assert dict(stats.stats_dict) == {"a": 2, "b": 2.5}