diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py
new file mode 100644
index 000000000..26339af90
--- /dev/null
+++ b/tests/agents/test_agent_utils.py
@@ -0,0 +1,67 @@
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image, ImageDraw
+
+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), color=(255, 255, 255))
+
+
+def test_draw_mouse_pointer_returns_new_rgba_image(white_image):
+    result = draw_mouse_pointer(white_image, 10, 10)
+
+    assert result is not white_image
+    assert result.mode == "RGBA"
+    assert result.size == white_image.size
+    # the pointer tip is darkened, a far away pixel is untouched
+    assert result.getpixel((12, 11))[:3] != (255, 255, 255)
+    assert result.getpixel((90, 90))[:3] == (255, 255, 255)
+    # the original image is not modified
+    assert white_image.getpixel((12, 11)) == (255, 255, 255)
+
+
+def test_draw_click_indicator_draws_plus_with_gap(white_image):
+    result = draw_click_indicator(white_image, 50, 50)
+
+    assert result.mode == "RGBA"
+    # center is left blank (gap), arms are drawn
+    assert result.getpixel((50, 50))[:3] == (255, 255, 255)
+    assert result.getpixel((50, 40))[:3] == (0, 0, 0)
+    assert result.getpixel((50, 60))[:3] == (0, 0, 0)
+    assert result.getpixel((40, 50))[:3] == (0, 0, 0)
+    assert result.getpixel((60, 50))[:3] == (0, 0, 0)
+    assert result.getpixel((5, 5))[:3] == (255, 255, 255)
+
+
+def test_draw_arrowhead_draws_two_lines():
+    draw = MagicMock(spec=ImageDraw.ImageDraw)
+    draw_arrowhead(draw, (0, 0), (50, 0))
+
+    assert draw.line.call_count == 2
+    for call in draw.line.call_args_list:
+        points = call.args[0]
+        assert points[0] == (50, 0)
+        # arrowhead points go back towards the start
+        assert points[1][0] < 50
+        assert call.kwargs["fill"] == "red"
+
+
+def test_zoom_webpage_sets_zoom():
+    page = MagicMock()
+    assert zoom_webpage(page, 2) is page
+    page.evaluate.assert_called_once_with("document.documentElement.style.zoom='200%'")
+
+
+@pytest.mark.parametrize("factor", [0, -1])
+def test_zoom_webpage_rejects_invalid_factor(factor):
+    with pytest.raises(ValueError):
+        zoom_webpage(MagicMock(), factor)
diff --git a/tests/agents/test_debug_agent.py b/tests/agents/test_debug_agent.py
new file mode 100644
index 000000000..f7b557c44
--- /dev/null
+++ b/tests/agents/test_debug_agent.py
@@ -0,0 +1,71 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+
+from agentlab.agents import debug_agent
+from agentlab.agents.debug_agent import DEBUG_AGENT, DebugAgent, DebugAgentArgs
+
+
+def test_debug_agent_args_defaults():
+    args = DebugAgentArgs()
+    assert args.agent_name == "debug"
+    assert args.use_html is False
+    assert DEBUG_AGENT.agent_name == "debug"
+
+
+def test_set_benchmark_enables_html_for_miniwob():
+    args = DebugAgentArgs()
+    action_set_args = MagicMock()
+    args.set_benchmark(
+        SimpleNamespace(name="miniwob_all", high_level_action_set_args=action_set_args), False
+    )
+    assert args.use_html is True
+    assert args.action_set_args is action_set_args
+
+    args = DebugAgentArgs()
+    args.set_benchmark(
+        SimpleNamespace(name="workarena", high_level_action_set_args=action_set_args), False
+    )
+    assert args.use_html is False
+
+
+def test_make_agent():
+    agent = DebugAgentArgs().make_agent()
+    assert isinstance(agent, DebugAgent)
+    assert agent.action_set is not None
+
+
+def test_obs_preprocessor(monkeypatch):
+    monkeypatch.setattr(debug_agent, "flatten_dom_to_str", lambda dom, **kw: "<p>dom</p>")
+    monkeypatch.setattr(debug_agent, "flatten_axtree_to_str", lambda ax, **kw: "axtree")
+    monkeypatch.setattr(debug_agent, "prune_html", lambda html: "pruned:" + html)
+    monkeypatch.setattr(debug_agent, "overlay_som", lambda img, extra_properties: "som")
+    obs = {
+        "dom_object": {},
+        "axtree_object": {},
+        "extra_element_properties": {},
+        "screenshot": np.zeros((2, 2, 3)),
+    }
+    out = DebugAgentArgs().make_agent().obs_preprocessor(obs)
+    assert out["dom_txt"] == "<p>dom</p>"
+    assert out["axtree_txt"] == "axtree"
+    assert out["pruned_html"] == "pruned:<p>dom</p>"
+    assert out["screenshot_som"] == "som"
+    assert "dom_txt" not in obs  # input is not mutated
+
+
+def test_get_action_prompts_user(monkeypatch):
+    prompts = []
+    monkeypatch.setattr("builtins.input", lambda p: prompts.append(p) or "noop()")
+    obs = {"pruned_html": "<html/>", "axtree_txt": "tree"}
+
+    agent = DebugAgentArgs().make_agent()
+    action, info = agent.get_action(obs)
+    assert action == "noop()"
+    assert info.think == "nope"
+    assert prompts[-1].startswith("tree")
+
+    agent.use_html = True
+    agent.get_action(obs)
+    assert prompts[-1].startswith("<html/>")
diff --git a/tests/agents/test_most_basic_agent.py b/tests/agents/test_most_basic_agent.py
new file mode 100644
index 000000000..6c091c09d
--- /dev/null
+++ b/tests/agents/test_most_basic_agent.py
@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgentArgs
+from agentlab.llm.llm_utils import ParseError
+
+
+@dataclass
+class _FakeModelArgs:
+    model_name: str = "fake"
+    answers: tuple = ()
+
+    def make_model(self):
+        answers = iter(self.answers)
+        return lambda messages: {"role": "assistant", "content": next(answers)}
+
+    def prepare_server(self):
+        return "prepared"
+
+    def close_server(self):
+        return "closed"
+
+
+OBS = {"goal": "Click the button", "pruned_html": "<button bid='a1'>OK</button>"}
+
+
+def test_args_lifecycle():
+    args = MostBasicAgentArgs(chat_model_args=_FakeModelArgs())
+    assert args.prepare() == "prepared"
+    assert args.close() == "closed"
+    args.set_reproducibility_mode()
+    assert args.temperature == 0
+
+
+@pytest.mark.parametrize("cot", [True, False])
+def test_get_action_parses_code_block(cot):
+    answer = "I will click.\n```\nclick('a1')\n```"
+    agent = MostBasicAgentArgs(
+        use_chain_of_thought=cot, chat_model_args=_FakeModelArgs(answers=(answer,))
+    ).make_agent()
+
+    action, info = agent.get_action(OBS)
+
+    assert action == "click('a1')"
+    assert info.think == answer
+    prompt = str(info.chat_messages[1])
+    assert "Click the button" in prompt and "bid='a1'" in prompt
+    assert ("chain of thoughts" in prompt) == cot
+    assert info.extra_info["chat_model_args"]["model_name"] == "fake"
+    assert info.stats["response_length"] == len(answer)
+
+
+def test_get_action_retries_until_code_block():
+    answers = ("no code", "```\nnoop()\n```")
+    agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=answers)).make_agent()
+    action, info = agent.get_action(OBS)
+    assert action == "noop()"
+    assert any("No code block found" in str(m) for m in info.chat_messages)
+
+
+def test_get_action_fails_after_retries():
+    agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=("x",) * 3)).make_agent()
+    with pytest.raises(ParseError):
+        agent.get_action(OBS)
diff --git a/tests/agents/test_tmlr_config.py b/tests/agents/test_tmlr_config.py
new file mode 100644
index 000000000..ef8cfe5cb
--- /dev/null
+++ b/tests/agents/test_tmlr_config.py
@@ -0,0 +1,26 @@
+from agentlab.agents.generic_agent import tmlr_config
+from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
+
+LLM = "openai/gpt-4o-mini-2024-07-18"
+
+
+def test_base_agent():
+    agent = tmlr_config.get_base_agent(LLM)
+    assert agent.chat_model_args is CHAT_MODEL_ARGS_DICT[LLM]
+    assert agent.flags.obs.use_screenshot is False
+    assert agent.flags.obs.use_som is False
+
+
+def test_vision_agent_does_not_mutate_base_flags():
+    agent = tmlr_config.get_vision_agent(LLM)
+    assert agent.flags.obs.use_screenshot is True
+    assert agent.flags.obs.use_som is False
+    assert agent.agent_name.endswith("_vision")
+    assert tmlr_config.BASE_FLAGS.obs.use_screenshot is False
+
+
+def test_som_agent():
+    agent = tmlr_config.get_som_agent(LLM)
+    assert agent.flags.obs.use_screenshot is True
+    assert agent.flags.obs.use_som is True
+    assert tmlr_config.BASE_FLAGS.obs.use_som is False
diff --git a/tests/analyze/test_archive_studies.py b/tests/analyze/test_archive_studies.py
new file mode 100644
index 000000000..abbd13936
--- /dev/null
+++ b/tests/analyze/test_archive_studies.py
@@ -0,0 +1,46 @@
+import pandas as pd
+
+from agentlab.analyze import archive_studies
+from agentlab.analyze.archive_studies import search_for_reasons_to_archive
+
+
+def _write_summary(study_dir, rows):
+    study_dir.mkdir()
+    pd.DataFrame(rows).to_csv(study_dir / "summary_2024.csv", index=False)
+
+
+def test_search_for_reasons_to_archive(tmp_path, monkeypatch):
+    monkeypatch.setattr(archive_studies.Study, "load", lambda d: None)
+
+    _write_summary(tmp_path / "good_study", [{"n_completed": "10/10", "n_err": 0}])
+    _write_summary(tmp_path / "errors_study", [{"n_completed": "10/10", "n_err": 8}])
+    _write_summary(tmp_path / "2024_miniwob-tiny-test", [{"n_completed": "4/4", "n_err": 0}])
+    _write_summary(tmp_path / "small_study", [{"n_completed": "2/3", "n_err": 0}])
+    (tmp_path / "not_a_dir.txt").write_text("x")
+    # no summary and no results -> load_result_df returns None, skipped
+    (tmp_path / "broken_study").mkdir()
+
+    infos = {i.study_dir.name: i for i in search_for_reasons_to_archive(tmp_path, min_study_size=5)}
+
+    assert set(infos) == {"good_study", "errors_study", "2024_miniwob-tiny-test", "small_study"}
+    assert not infos["good_study"].should_delete
+    assert infos["good_study"].reason == ""
+    assert infos["errors_study"].should_delete
+    assert "Less than 50%" in infos["errors_study"].reason
+    assert "Miniwob tiny test" in infos["2024_miniwob-tiny-test"].reason
+    assert "Too few tasks" in infos["small_study"].reason
+
+
+def test_search_uses_most_recent_summary(tmp_path, monkeypatch):
+    import os
+
+    monkeypatch.setattr(archive_studies.Study, "load", lambda d: None)
+    study_dir = tmp_path / "study"
+    _write_summary(study_dir, [{"n_completed": "0/10", "n_err": 0}])
+    newer = study_dir / "summary_2025.csv"
+    pd.DataFrame([{"n_completed": "10/10", "n_err": 0}]).to_csv(newer, index=False)
+    old = study_dir / "summary_2024.csv"
+    os.utime(old, (1, 1))
+
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert info.summary_df["n_completed"].iloc[0] in ("10/10", "0/10")
diff --git a/tests/analyze/test_episode_to_html.py b/tests/analyze/test_episode_to_html.py
new file mode 100644
index 000000000..6399aaf78
--- /dev/null
+++ b/tests/analyze/test_episode_to_html.py
@@ -0,0 +1,137 @@
+from pathlib import Path
+from types import SimpleNamespace
+
+from PIL import Image
+
+from agentlab.analyze import episode_to_html as e2h
+
+
+class _FakeExpResult:
+    def __init__(self, steps_info, screenshots=None, paths=None, som=None):
+        self.exp_args = SimpleNamespace(
+            env_args=SimpleNamespace(task_name="miniwob.click", task_seed=7)
+        )
+        self.steps_info = steps_info
+        self._screenshots = screenshots or {}
+        self._paths = paths or {}
+        self._som = som or {}
+
+    def get_screenshot(self, step, som=False):
+        store = self._som if som else self._screenshots
+        if step not in store:
+            raise FileNotFoundError
+        return store[step]
+
+    def get_screenshot_path(self, step, som=False):
+        if step not in self._paths:
+            raise FileNotFoundError
+        return self._paths[step]
+
+
+def _step(action=None, obs=None, agent_info=None):
+    return SimpleNamespace(action=action, obs=obs or {}, agent_info=agent_info or {})
+
+
+def test_exp_result_to_html_full_episode():
+    img = Image.new("RGB", (4, 4))
+    steps = [
+        _step(obs={"goal_object": [{"type": "text", "text": "Click <b>"}]}),
+        _step(
+            action="click('a')",
+            obs={"axtree_txt": "[1] button", "dom_txt": "<div>x</div>"},
+            agent_info={
+                "think": "I should click",
+                "chat_messages": [{"role": "user", "content": "hi"}],
+            },
+        ),
+    ]
+    exp_result = _FakeExpResult(steps, screenshots={1: img}, som={1: img})
+
+    html = e2h.exp_result_to_html(exp_result, steps_open=False, axtree_open=True)
+
+    assert "miniwob.click" in html and "Seed: 7" in html
+    assert "Click &lt;b&gt;" in html
+    assert "Step 0" not in html  # initial reset step without action is skipped
+    assert "<summary>Step 1</summary>" in html
+    assert "<details >" in html  # steps closed
+    assert "click(&#x27;a&#x27;)" in html
+    assert "I should click" in html
+    assert "Screenshot_som[1]" in html
+    assert "AXTree" in html and "&lt;div&gt;x&lt;/div&gt;" in html
+    assert "Prompt" in html and "USER" in html
+    assert html.count("data:image/png;base64,") == 2
+
+
+def test_exp_result_to_html_without_images_and_goal():
+    steps = [_step(action="noop()")]
+    html = e2h.exp_result_to_html(_FakeExpResult(steps))
+    assert "Goal:" in html  # default "No goal specified"
+    assert "Screenshot not available" in html
+    assert "Screenshot_som" not in html
+
+
+def test_exp_result_to_html_empty_episode():
+    html = e2h.exp_result_to_html(_FakeExpResult([]))
+    assert "Goal:" not in html
+    assert "<summary>" not in html
+
+
+def test_exp_result_to_html_with_file_paths(tmp_path):
+    path = tmp_path / "screenshot_step_1.png"
+    Image.new("RGB", (2, 2)).save(path)
+    steps = [_step(action="noop()"), _step(action="noop()")]
+    exp_result = _FakeExpResult(steps, paths={0: path, 1: tmp_path / "missing.png"})
+    html = e2h.exp_result_to_html(exp_result, embed_images=False)
+    assert f"file://{path.resolve()}" in html
+    assert "Screenshot 1 not available" in html
+    assert "Screenshot_som[0]" in html
+
+
+def test_image_and_path_helpers_handle_none():
+    assert e2h._image_to_html(None, "alt") == "<p>alt not available</p>"
+    assert e2h._path_to_html(None, "alt") == "<p>alt not available</p>"
+    assert e2h._path_to_html(Path("/does/not/exist.png"), "alt") == "<p>alt not available</p>"
+
+
+def test_format_goal():
+    assert "No goal specified" in e2h._format_goal(None)
+    assert "buy milk" in e2h._format_goal("buy milk")
+
+
+def test_format_chat_messages_like_xray():
+    assert "No chat messages" in e2h._format_chat_messages_like_xray([])
+
+    msg_obj = SimpleNamespace(role="system", content="be nice")
+    messages = [
+        msg_obj,
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "<hello>"},
+                {"type": "image"},
+                {"type": "tool_use", "name": "click", "input": {"bid": "1"}, "call_id": "c1"},
+                {"type": "other", "x": 1},
+                "raw part",
+            ],
+        },
+        {"content": "no role"},
+        42,
+    ]
+    html = e2h._format_chat_messages_like_xray(messages)
+    assert html.count('class="chat-message"') == 4
+    assert "SYSTEM" in html and "ASSISTANT" in html and "UNKNOWN" in html
+    assert "&lt;hello&gt;" in html
+    assert "[IMAGE]" in html
+    assert "Tool Call: click" in html and "call_id: c1" in html
+    assert "raw part" in html and "42" in html
+
+
+def test_format_tool_call_defaults():
+    assert e2h._format_tool_call_like_xray({}) == "Tool Call: unknown `{}` (call_id: unknown)"
+
+
+def test_escape_html():
+    assert e2h._escape_html("<a href=\"x\">'&'</a>") == (
+        "&lt;a href=&quot;x&quot;&gt;&#x27;&amp;&#x27;&lt;/a&gt;"
+    )
+    assert e2h._escape_html(5) == "5"
diff --git a/tests/analyze/test_inspect_results_unit.py b/tests/analyze/test_inspect_results_unit.py
new file mode 100644
index 000000000..97d8b19e9
--- /dev/null
+++ b/tests/analyze/test_inspect_results_unit.py
@@ -0,0 +1,251 @@
+import json
+import time
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from agentlab.analyze import inspect_results as ir
+
+
+def _result_df():
+    return pd.DataFrame(
+        {
+            "env.task_name": ["miniwob.a", "miniwob.a", "miniwob.b", "miniwob.b"] * 2,
+            "agent.agent_name": ["ag1"] * 4 + ["ag2"] * 4,
+            "agent.flags.use_x": [True] * 4 + [False] * 4,
+            "agent._private": [1, 2, 3, 4, 5, 6, 7, 8],
+            "cum_reward": [1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0],
+            "n_steps": [1, 2, 3, 4, 5, 6, 7, 8],
+            "err_msg": [None] * 7 + [None],
+            "stack_trace": [None] * 8,
+            "truncated": [False] * 8,
+            "terminated": [True] * 8,
+            "stats.cum_cost": [0.1] * 8,
+            "constant": ["c"] * 8,
+        }
+    )
+
+
+def test_get_constants_and_variables():
+    df = pd.DataFrame({"a": [1, 1], "b": [1, 2], "c": [np.int64(3)] * 2, "d": [[1], [2]]})
+    constants, variables, out = ir.get_constants_and_variables(df, drop_constants=True)
+    assert constants == {"a": 1, "c": 3}
+    assert isinstance(constants["c"], int)
+    assert variables == ["b", "d"]
+    assert list(out.columns) == ["b", "d"]
+
+
+def test_set_index_from_variables():
+    df = _result_df()
+    df.loc[0, "agent.flags.use_x"] = None
+    ir.set_index_from_variables(df)
+    assert df.index.names == [
+        "env.task_name",
+        "agent.agent_name",
+        "env.benchmark",
+        "agent.flags.use_x",
+    ]
+    assert set(df.index.get_level_values("env.benchmark")) == {"miniwob"}
+    # None values in index variables are replaced by the string "None"
+    assert "None" in set(df.index.get_level_values("agent.flags.use_x"))
+
+
+def _indexed():
+    df = _result_df()
+    ir.set_index_from_variables(df)
+    return df
+
+
+def test_summarize_with_and_without_rewards():
+    df = _result_df()
+    s = ir.summarize(df[df["agent.agent_name"] == "ag1"])
+    assert s["avg_reward"] == 0.75
+    assert s["n_completed"] == "4/4"
+    assert s["n_err"] == 0
+    assert s["cum_cost"] == 0.4
+
+    s = ir.summarize(pd.DataFrame({"x": [1, 2]}))
+    assert s["n_completed"] == "0/2" and np.isnan(s["avg_reward"])
+
+    not_done = df.assign(terminated=False)
+    assert ir.summarize(not_done) is None
+
+
+def test_summarize_effective_cost_replaces_cost():
+    df = _result_df().assign(**{"stats.cum_effective_cost": 0.2})
+    s = ir.summarize(df)
+    assert "cum_cost" not in s and s["cum_effective_cost"] == 1.6
+
+
+def test_summarize_stats():
+    df = _result_df().assign(**{"stats.max_tokens": range(8)})
+    s = ir.summarize_stats(df)
+    assert s["cum_cost"] == pytest.approx(0.8)
+    assert s["max_tokens"] == 7
+    assert ir.summarize_stats(df.assign(terminated=False)) is None
+    with pytest.raises(ValueError, match="Unknown stats"):
+        ir.summarize_stats(df.assign(**{"stats.avg_x": 1}))
+
+
+def test_std_err_helpers():
+    df = pd.DataFrame({"m": [1, 0, 1, 0]})
+    mean, err = ir.get_std_err(df, "m")
+    assert mean == 0.5 and err == pytest.approx(0.25)
+
+    df = pd.DataFrame({"m": [0.5, 1.5]})
+    mean, err = ir.get_std_err(df, "m")
+    assert mean == 1.0 and err == pytest.approx(0.5)
+
+    _, err = ir.get_sample_std_err(pd.DataFrame({"m": [0.3]}), "m")
+    assert err == 0
+
+
+def test_global_report_multiple_agents_and_flag_report():
+    report = ir.global_report(_indexed())
+    assert list(report.index.names) == ["agent.agent_name", "env.benchmark", "use_x"]
+    assert report["avg_reward"].iloc[0] >= report["avg_reward"].iloc[1]
+
+    flags = ir.flag_report(report)
+    assert list(flags.index) == ["use_x"]
+    assert flags.loc["use_x", "avg_true"] == 0.75
+    assert flags.loc["use_x", "avg_false"] == 0.25
+    assert flags.loc["use_x", "ratio"] == 3.0
+
+
+def test_global_report_single_configuration():
+    df = _result_df()
+    df = df[df["agent.agent_name"] == "ag1"].set_index("env.task_name")
+    report = ir.global_report(df)
+    assert "[ALL TASKS]" in report.index
+    assert report.loc["[ALL TASKS]", "avg_reward"] == 0.75
+    assert ir.flag_report(report) is None
+
+
+def test_summarize_study_and_split_by_key():
+    df = _indexed()
+    summary = ir.summarize_study(df)
+    assert len(summary) == 2
+
+    parts = ir.split_by_key(df, "agent.agent_name")
+    assert set(parts) == {"ag1", "ag2"}
+    assert len(parts["ag1"]) == 4
+
+
+def test_report_constant_and_variables(capsys):
+    df = _result_df().assign(stack_trace=["t"] * 7 + ["u"])
+    ir.report_constant_and_variables(df, show_stack_traces=False)
+    out = capsys.readouterr().out
+    assert "constant: c" in out
+    assert "agent._private: n_unique=8" in out
+    assert "stack_trace" not in out
+
+
+def test_rename_bool_flags_and_shrink_columns():
+    report = ir.global_report(_indexed())
+    renamed = ir._rename_bool_flags(report.copy())
+    assert set(renamed.index.get_level_values("use_x")) == {"✓", "-"}
+
+    shrunk = ir.shrink_columns(
+        pd.DataFrame({"avg_reward": [0.5000]}, index=pd.Index([0], name="a_b"))
+    )
+    assert list(shrunk.columns) == ["avg\nreward"]
+    assert shrunk.index.names == ["a\nb"]
+    assert shrunk.iloc[0, 0] == "0.5"
+
+
+def test_find_diff():
+    assert ir._find_diff((1, 2, 3), (1, 0, 4)) == [1, 2]
+
+
+def test_map_err_key():
+    assert ir.map_err_key(None) is None
+    msg = (
+        "Error: your messages resulted in 12345 tokens\n"
+        "Exception uncaught by agent or environment in task miniwob.abc\n=== logs ===\nblah"
+    )
+    key = ir.map_err_key(msg)
+    assert "x tokens" in key
+    assert "<task_name>." in key
+    assert "blah" not in key
+
+
+def test_categorize_error_without_error():
+    assert ir.categorize_error({"err_msg": None}) is None
+    assert ir.categorize_error({}) is None
+
+
+def test_format_log():
+    assert ir._format_log(type("R", (), {"logs": None})()) == "No log found"
+    short = type("R", (), {"logs": "a\nb"})()
+    assert ir._format_log(short) == "a\nb"
+
+    lines = [f"line {i}" for i in range(100)]
+    long_log = type("R", (), {"logs": "\n".join(lines)})()
+    out = ir._format_log(long_log, head_lines=2, tail_lines=3)
+    assert out.startswith("line 0\nline 1\n...")
+    assert out.endswith("line 97\nline 98\nline 99")
+
+    with_tb = "\n".join(lines + ["action: click", "Traceback (most recent call last):", "boom"])
+    out = ir._format_log(type("R", (), {"logs": with_tb})(), head_lines=2, tail_lines=3)
+    assert out.endswith("action: click\nTraceback (most recent call last):\nboom")
+
+
+def test_benchmark_from_task_name():
+    assert ir._benchmark_from_task_name("workarena.servicenow.x") == "workarena"
+
+
+def test_is_stale_tracks_modification_times(tmp_path):
+    (tmp_path / "exp").mkdir()
+    f = tmp_path / "exp" / "file.txt"
+    f.write_text("a")
+    summary = tmp_path / "study_summary.csv"
+
+    assert ir._is_stale(tmp_path, summary)  # no summary yet
+    summary.write_text("x")
+    # the summary itself is excluded from the tracked modification times
+    assert not ir._is_stale(tmp_path, summary)
+
+    mtimes = json.loads((tmp_path / "_last_modification_times.json").read_text())
+    assert "exp/file.txt" in mtimes and "study_summary.csv" not in mtimes
+
+    time.sleep(0.01)
+    f.write_text("b")
+    import os
+
+    os.utime(f, (time.time() + 10, time.time() + 10))
+    assert ir._is_stale(tmp_path, summary)
+
+
+def test_get_study_summary_uses_cache(tmp_path, monkeypatch):
+    df = _indexed()
+    calls = []
+
+    def fake_load(study_dir, progress_fn=None):
+        calls.append(study_dir)
+        return df
+
+    monkeypatch.setattr(ir, "load_result_df", fake_load)
+    sentinel = {}
+    ir.get_study_summary(tmp_path, sentinel=sentinel)
+    assert sentinel["from_cache"] is False
+    ir.get_study_summary(tmp_path, sentinel=sentinel)
+    assert sentinel["from_cache"] is True
+    ir.get_study_summary(tmp_path, ignore_cache=True, ignore_stale=True, sentinel=sentinel)
+    assert sentinel["from_cache"] is False
+    assert len(calls) == 2
+
+
+def test_get_study_summary_none_and_get_all_summaries(tmp_path, monkeypatch):
+    df = _indexed()
+    (tmp_path / "_hidden").mkdir()
+    (tmp_path / "empty").mkdir()
+    (tmp_path / "study").mkdir()
+    monkeypatch.setattr(
+        ir,
+        "load_result_df",
+        lambda d, progress_fn=None: df if d.name == "study" else None,
+    )
+    assert ir.get_study_summary(tmp_path / "empty") is None
+    summaries = ir.get_all_summaries(tmp_path)
+    assert set(summaries.index) == {"study"}
diff --git a/tests/benchmarks/test_abstract_env.py b/tests/benchmarks/test_abstract_env.py
new file mode 100644
index 000000000..4c9b17683
--- /dev/null
+++ b/tests/benchmarks/test_abstract_env.py
@@ -0,0 +1,52 @@
+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=None):
+        self._env_info = env_info
+
+    def reset(self, seed=None):
+        return {}, {}
+
+    @add_step_timing_to_env_info_decorator
+    def step(self, action: str):
+        return {"action": action}, 1.0, True, False, self._env_info
+
+    def close(self):
+        pass
+
+
+def test_abstract_benchmark_defaults():
+    bench = AbstractBenchmark(name="dummy", 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_env_cannot_be_instantiated():
+    with pytest.raises(TypeError):
+        AbstractEnv()
+
+
+def test_step_timing_added_when_env_info_is_none():
+    obs, reward, terminated, truncated, env_info = _DummyEnv(None).step("noop()")
+    assert obs == {"action": "noop()"}
+    assert reward == 1.0 and terminated and not truncated
+    assert env_info["action_exec_start"] <= env_info["action_exec_stop"]
+    assert env_info["action_exec_timeout"] == 0.0
+
+
+def test_step_timing_does_not_override_existing_values():
+    existing = {"action_exec_start": 1, "action_exec_stop": 2, "action_exec_timeout": 3}
+    *_, env_info = _DummyEnv(dict(existing)).step("noop()")
+    assert env_info == existing
+
+
+def test_step_timing_preserves_wrapped_name():
+    assert _DummyEnv.step.__name__ == "step"
diff --git a/tests/experiments/test_exp_utils.py b/tests/experiments/test_exp_utils.py
new file mode 100644
index 000000000..049244a90
--- /dev/null
+++ b/tests/experiments/test_exp_utils.py
@@ -0,0 +1,115 @@
+import signal
+import time
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlab.experiments import exp_utils
+from agentlab.experiments.exp_utils import (
+    MockedExpArgs,
+    _episode_timeout,
+    add_dependencies,
+    run_exp,
+    timeout_manager,
+)
+
+
+def test_run_exp_calls_run():
+    exp_arg = MagicMock()
+    exp_arg.run.return_value = "result"
+    assert run_exp(exp_arg, "dep1", "dep2") == "result"
+    exp_arg.run.assert_called_once()
+
+
+def test_episode_timeout_without_max_steps():
+    exp_arg = SimpleNamespace(env_args=SimpleNamespace())
+    assert _episode_timeout(exp_arg) == 10 * 60 * 60
+
+
+def test_episode_timeout_with_max_steps():
+    exp_arg = SimpleNamespace(env_args=SimpleNamespace(max_steps=5))
+    assert _episode_timeout(exp_arg, avg_step_timeout=10) == 50
+
+
+def test_episode_timeout_uses_smaller_exp_timeout():
+    exp_arg = SimpleNamespace(env_args=SimpleNamespace(max_steps=5), episode_timeout=20)
+    assert _episode_timeout(exp_arg, avg_step_timeout=10) == 20
+
+
+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(0.5):  # floats are rounded up to at least 1 second
+            time.sleep(3)
+    assert signal.getsignal(signal.SIGALRM) == previous
+    assert signal.alarm(0) == 0
+
+
+def test_timeout_manager_no_timeout():
+    with timeout_manager(5):
+        pass
+    assert signal.alarm(0) == 0
+
+
+def _exp_args(task_name):
+    exp_id = f"id_{task_name}"
+    ea = SimpleNamespace(env_args=SimpleNamespace(task_name=task_name), exp_id=None, depends_on=())
+
+    def make_id():
+        ea.exp_id = exp_id
+
+    ea.make_id = make_id
+    return ea
+
+
+@pytest.mark.parametrize("deps", [None, {"a": [], "b": []}])
+def test_add_dependencies_noop(deps):
+    exp_args_list = [_exp_args("a"), _exp_args("b")]
+    assert add_dependencies(exp_args_list, deps) is exp_args_list
+    assert all(ea.exp_id is None for ea in exp_args_list)
+
+
+def test_add_dependencies_maps_task_names_to_exp_ids():
+    exp_args_list = [_exp_args("a"), _exp_args("b"), _exp_args("c")]
+    add_dependencies(exp_args_list, {"a": [], "b": ["a"], "c": ["a", "b"]})
+    assert exp_args_list[0].depends_on == ()
+    assert exp_args_list[1].depends_on == ("id_a",)
+    assert exp_args_list[2].depends_on == ("id_a", "id_b")
+
+
+def test_add_dependencies_duplicate_task_names():
+    with pytest.raises(ValueError, match="not unique"):
+        add_dependencies([_exp_args("a"), _exp_args("a")], {"a": ["a"]})
+
+
+def test_add_dependencies_missing_task():
+    with pytest.raises(ValueError, match="missing"):
+        add_dependencies([_exp_args("a"), _exp_args("b")], {"a": ["b"]})
+
+
+def test_mocked_exp_args(monkeypatch):
+    monkeypatch.setattr(exp_utils, "sleep", lambda s: None)
+    exp = MockedExpArgs("1", depends_on=["0"])
+    assert exp.exp_name == "exp_1" and exp.depends_on == ["0"]
+    assert MockedExpArgs("2").depends_on == []
+    assert exp.run() is exp
+    assert exp.start_time <= exp.end_time
+
+
+@pytest.mark.parametrize(
+    "fn, args",
+    [
+        (exp_utils.make_seeds, (3,)),
+        (exp_utils.order, ([],)),
+        (exp_utils.hide_some_exp, (".", lambda e: True, True)),
+    ],
+)
+def test_deprecated_functions_raise(fn, args):
+    with pytest.raises(DeprecationWarning):
+        fn(*args)
diff --git a/tests/experiments/test_reproducibility_util_extra.py b/tests/experiments/test_reproducibility_util_extra.py
new file mode 100644
index 000000000..666aa2f73
--- /dev/null
+++ b/tests/experiments/test_reproducibility_util_extra.py
@@ -0,0 +1,114 @@
+import logging
+
+import pandas as pd
+import pytest
+
+from agentlab.experiments import reproducibility_util as ru
+
+
+def _report(rows):
+    return pd.DataFrame(rows).set_index("agent.agent_name")
+
+
+GOOD_ROWS = [
+    {"agent.agent_name": "a1", "avg_reward": 0.5, "std_err": 0.1, "n_err": 0, "n_completed": "4/4"},
+    {"agent.agent_name": "a2", "avg_reward": 0.7, "std_err": 0.2, "n_err": 0, "n_completed": "4/4"},
+]
+
+
+def _info():
+    return {
+        "git_user": "me",
+        "agent_names": ["a1", "a2"],
+        "benchmark": "miniwob",
+        "date": "today",
+    }
+
+
+def test_assert_compatible_ignores_result_keys():
+    info = {"benchmark": "b", "date": "1", "avg_reward": 0.1}
+    old = {"benchmark": "b", "date": "2", "avg_reward": 0.9}
+    ru.assert_compatible(info, old)
+
+
+def test_assert_compatible_raises_or_warns(caplog):
+    info, old = {"benchmark": "b1"}, {"benchmark": "b2"}
+    with pytest.raises(ValueError, match="Key benchmark has changed"):
+        ru.assert_compatible(info, old)
+    with caplog.at_level(logging.WARNING):
+        ru.assert_compatible(info, old, raise_if_incompatible=False)
+    assert "Key benchmark has changed" in caplog.text
+
+
+def test_verify_report_ok():
+    df = ru._verify_report(_report(GOOD_ROWS), ["a1", "a2"])
+    assert list(df.index) == ["a1", "a2"]
+
+
+def test_verify_report_agent_name_mismatch():
+    with pytest.raises(ValueError, match="do not match"):
+        ru._verify_report(_report(GOOD_ROWS), ["a1"])
+
+
+def test_verify_report_duplicate_agent_names():
+    with pytest.raises(ValueError, match="Duplicate"):
+        ru._verify_report(_report(GOOD_ROWS), ["a1", "a2", "a2"])
+
+
+@pytest.mark.parametrize(
+    "override, match",
+    [({"n_err": 1}, "has 1 errors"), ({"n_completed": "3/4"}, "3 completed tasks out of 4")],
+)
+def test_verify_report_incomplete(override, match, caplog):
+    rows = [dict(GOOD_ROWS[0], **override)]
+    with pytest.raises(ValueError, match=match):
+        ru._verify_report(_report(rows), ["a1"])
+    with caplog.at_level(logging.WARNING):
+        ru._verify_report(_report(rows), ["a1"], strict_reproducibility=False)
+    assert "a1" in caplog.text
+
+
+def test_get_csv_headers(tmp_path):
+    empty = tmp_path / "empty.csv"
+    empty.write_text("")
+    assert ru._get_csv_headers(empty) is None
+    f = tmp_path / "f.csv"
+    f.write_text("a,b\n1,2\n")
+    assert ru._get_csv_headers(f) == ["a", "b"]
+
+
+def test_add_result_to_info():
+    info = {"agent_name": "a2"}
+    ru._add_result_to_info(info, _report(GOOD_ROWS))
+    assert info == {
+        "agent_name": "a2",
+        "avg_reward": 0.7,
+        "std_err": 0.2,
+        "n_err": 0,
+        "n_completed": "4/4",
+    }
+
+
+def test_append_to_journal_creates_and_appends(tmp_path):
+    journal = tmp_path / "journal.csv"
+    ru.append_to_journal(_info(), _report(GOOD_ROWS), journal_path=journal)
+    ru.append_to_journal(_info(), _report(GOOD_ROWS), journal_path=journal)
+
+    df = pd.read_csv(journal)
+    assert list(df.columns) == ["git_user", "agent_name", "benchmark", "date"]
+    assert list(df["agent_name"]) == ["a1", "a2", "a1", "a2"]
+
+
+def test_append_to_journal_mismatched_agents(tmp_path):
+    with pytest.raises(ValueError, match="Mismatch"):
+        ru.append_to_journal(_info(), _report(GOOD_ROWS[:1]), journal_path=tmp_path / "j.csv")
+
+
+def test_append_to_journal_default_path_without_git(tmp_path, monkeypatch):
+    def no_repo(module):
+        raise ru.InvalidGitRepositoryError()
+
+    monkeypatch.setattr(ru, "_get_repo", no_repo)
+    monkeypatch.setattr(ru, "RESULTS_DIR", tmp_path)
+    ru.append_to_journal(_info(), _report(GOOD_ROWS))
+    assert (tmp_path / "reproducibility_journal.csv").exists()
diff --git a/tests/experiments/test_study_utils.py b/tests/experiments/test_study_utils.py
new file mode 100644
index 000000000..8c2f55a50
--- /dev/null
+++ b/tests/experiments/test_study_utils.py
@@ -0,0 +1,65 @@
+from dataclasses import asdict
+
+import bgym
+import pytest
+
+from agentlab.experiments import study as study_mod
+from agentlab.experiments.loop import EnvArgs
+from agentlab.experiments.study import (
+    _convert_env_args,
+    _make_study_name,
+    get_most_recent_study,
+    set_demo_mode,
+)
+
+
+def test_make_study_name_single_agent_and_benchmark():
+    assert _make_study_name(["agent"], ["miniwob"], "v1") == "agent-on-miniwob-v1"
+
+
+def test_make_study_name_multiple():
+    name = _make_study_name(["a", "b", "a"], ["x", "y", "z"])
+    assert name == "2-agents-on-3-benchmarks"
+
+
+def test_get_most_recent_study(tmp_path):
+    for name in [
+        "2024-01-01_10-00-00_old",
+        "2024-06-01_10-00-00_newer_foo",
+        "2024-03-01_10-00-00_foo",
+        "_2025-01-01_10-00-00_hidden",
+        "not_a_date",
+    ]:
+        (tmp_path / name).mkdir()
+    (tmp_path / "2026-01-01_10-00-00_file").write_text("x")
+
+    assert get_most_recent_study(tmp_path).name == "2024-06-01_10-00-00_newer_foo"
+    assert get_most_recent_study(tmp_path, contains="old").name == "2024-01-01_10-00-00_old"
+    assert get_most_recent_study(tmp_path, contains="missing") is None
+
+
+def test_get_most_recent_study_defaults_to_results_dir(tmp_path, monkeypatch):
+    (tmp_path / "2024-01-01_10-00-00_x").mkdir()
+    monkeypatch.setattr(study_mod, "RESULTS_DIR", tmp_path)
+    assert get_most_recent_study().name == "2024-01-01_10-00-00_x"
+
+
+def test_set_demo_mode():
+    env_args = EnvArgs(task_name="miniwob.click-test")
+    set_demo_mode([env_args])
+    assert env_args.viewport == {"width": 1280, "height": 720}
+    assert env_args.record_video is True
+    assert env_args.wait_for_user_message is False
+    assert env_args.slow_mo == 1000
+
+
+def test_convert_env_args():
+    new = EnvArgs(task_name="miniwob.a")
+    old = bgym.EnvArgs(task_name="miniwob.b", task_seed=3)
+    converted = _convert_env_args([new, old])
+    assert converted[0] is new
+    assert isinstance(converted[1], EnvArgs)
+    assert asdict(converted[1])["task_seed"] == 3
+
+    with pytest.raises(TypeError, match="Unexpected type"):
+        _convert_env_args(["nope"])
diff --git a/tests/llm/test_chat_api_unit.py b/tests/llm/test_chat_api_unit.py
new file mode 100644
index 000000000..dc6c28afd
--- /dev/null
+++ b/tests/llm/test_chat_api_unit.py
@@ -0,0 +1,246 @@
+"""Offline unit tests for chat_api (no network calls)."""
+
+import threading
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import openai
+import pytest
+
+from agentlab.llm import chat_api
+from agentlab.llm.chat_api import (
+    AnthropicChatModel,
+    AnthropicModelArgs,
+    ChatModel,
+    ChatModelArgs,
+    CheatMiniWoBLLM,
+    CheatMiniWoBLLMArgs,
+    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
+
+
+@pytest.fixture(autouse=True)
+def fresh_tracker(monkeypatch):
+    # other tests may leave TRACKER.instance set to None in this thread
+    monkeypatch.setattr(chat_api.tracking, "TRACKER", threading.local())
+
+
+@pytest.fixture(autouse=True)
+def fake_api_keys(monkeypatch):
+    for var in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "OPENROUTER_API_KEY"):
+        monkeypatch.setenv(var, "test-key")
+
+
+@pytest.fixture
+def no_sleep(monkeypatch):
+    sleeps = []
+    monkeypatch.setattr(chat_api.time, "sleep", lambda s: sleeps.append(s))
+    return sleeps
+
+
+def test_make_messages():
+    assert make_system_message("s") == {"role": "system", "content": "s"}
+    assert make_user_message("u") == {"role": "user", "content": "u"}
+    assert make_assistant_message("a") == {"role": "assistant", "content": "a"}
+
+
+PROMPT = "Some header\n[12] button 'Click me'\n"
+
+
+@pytest.mark.parametrize(
+    "messages",
+    [
+        [make_system_message("sys"), make_user_message(PROMPT)],
+        Discussion([SystemMessage("sys"), HumanMessage(PROMPT)]),
+    ],
+)
+def test_cheat_miniwob_llm(messages):
+    answer = CheatMiniWoBLLM()(messages)
+    assert answer["role"] == "assistant"
+    assert 'click("12")' in answer["content"]
+
+
+def test_cheat_miniwob_llm_waits_and_fails_without_button(monkeypatch, no_sleep):
+    llm = CheatMiniWoBLLMArgs(wait_time=2).make_model()
+    with pytest.raises(Exception, match="Can't find the button"):
+        llm([make_system_message("sys"), make_user_message("no buttons here")])
+    assert no_sleep == [2]
+
+
+def test_cheat_miniwob_llm_args_server_noops():
+    args = CheatMiniWoBLLMArgs()
+    assert args.prepare_server() is None
+    assert args.close_server() is None
+
+
+@pytest.mark.parametrize(
+    "msg, expected",
+    [
+        ("Rate limit. Please try again in 90.5s.", 90.5),
+        ("Rate limit. Please try again in 5s.", 60),
+        ("unrelated error", 60),
+    ],
+)
+def test_extract_wait_time(msg, expected):
+    assert _extract_wait_time(msg) == expected
+
+
+def test_handle_error(no_sleep):
+    err = openai.OpenAIError("please try again in 2s")
+    assert handle_error(err, 0, min_retry_wait_time=1, max_retry=3) == "please try again in 2s"
+    assert no_sleep == [2.0]
+
+
+def test_handle_error_reraises_non_openai_errors():
+    with pytest.raises(KeyError):
+        handle_error(KeyError("x"), 0, 1, 3)
+
+
+def _completion(contents, usage=True):
+    return SimpleNamespace(
+        usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5) if usage else None,
+        choices=[
+            SimpleNamespace(message=SimpleNamespace(content=c), log_probs="lp") for c in contents
+        ],
+    )
+
+
+def _chat_model(create, **kwargs):
+    client = MagicMock()
+    client.chat.completions.create.side_effect = create
+    model = ChatModel("m", client_class=MagicMock(return_value=client), **kwargs)
+    return model, client
+
+
+def test_chat_model_pricing():
+    model, _ = _chat_model(None, pricing_func=lambda: {"m": {"prompt": "0.1", "completion": "0.2"}})
+    assert (model.input_cost, model.output_cost) == (0.1, 0.2)
+
+    model, _ = _chat_model(None, pricing_func=lambda: {})
+    assert (model.input_cost, model.output_cost) == (0.0, 0.0)
+
+
+def test_chat_model_api_key_from_env(monkeypatch):
+    monkeypatch.setenv("MY_KEY", "secret")
+    model, _ = _chat_model(None, api_key_env_var="MY_KEY")
+    assert model.api_key == "secret"
+
+
+def test_chat_model_single_sample_with_log_probs():
+    model, client = _chat_model([_completion(["hello"])], log_probs=True, temperature=0.3)
+    res = model([make_user_message("hi")])
+    assert res["content"] == "hello" and res["log_probs"] == "lp"
+    kwargs = client.chat.completions.create.call_args.kwargs
+    assert kwargs["temperature"] == 0.3 and kwargs["n"] == 1
+    assert model.get_stats() == {"n_retry_llm": 1}
+
+
+def test_chat_model_multiple_samples():
+    model, _ = _chat_model([_completion(["a", "b"])])
+    res = model([make_user_message("hi")], n_samples=2, temperature=0)
+    assert [r["content"] for r in res] == ["a", "b"]
+
+
+def test_chat_model_retries_on_missing_usage(no_sleep):
+    model, _ = _chat_model(
+        [_completion(["x"], usage=False), _completion(["ok"])], min_retry_wait_time=0
+    )
+    assert model([make_user_message("hi")])["content"] == "ok"
+    assert model.retries == 2 and model.success
+    assert len(model.error_types) == 1
+
+
+def test_chat_model_raises_retry_error(no_sleep):
+    model, _ = _chat_model(openai.OpenAIError("down"), max_retry=2, min_retry_wait_time=0)
+    with pytest.raises(RetryError, match="after 2 retries"):
+        model([make_user_message("hi")])
+    assert model.error_types == ["down", "down"]
+
+
+def test_chat_model_rejects_zero_retry():
+    with pytest.raises(AssertionError):
+        _chat_model(None, max_retry=0)
+
+
+def test_azure_chat_model_requires_env(monkeypatch):
+    monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False)
+    with pytest.raises(AssertionError, match="AZURE_OPENAI_API_KEY"):
+        chat_api.AzureChatModel("gpt-4o")
+    with pytest.raises(AssertionError, match="AZURE_OPENAI_ENDPOINT"):
+        chat_api.AzureChatModel("gpt-4o", api_key="k")
+
+
+def test_azure_chat_model_client_args(monkeypatch):
+    monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.azure.com")
+    monkeypatch.setattr(chat_api.tracking, "get_pricing_openai", lambda: {})
+    model = chat_api.AzureChatModel("gpt-4o", api_key="k", deployment_name="old")
+    assert str(model.client.base_url).startswith("https://example.azure.com")
+
+
+def test_openrouter_and_vllm_models(monkeypatch):
+    monkeypatch.setattr(chat_api.tracking, "get_pricing_openrouter", lambda: {})
+    router = chat_api.OpenRouterChatModel("some/model", api_key="k")
+    assert "openrouter.ai" in str(router.client.base_url)
+
+    vllm = chat_api.VLLMChatModel("some/model", api_key="k", n_retry_server=2)
+    assert "0.0.0.0:8000" in str(vllm.client.base_url)
+    assert vllm.max_retry == 2
+
+
+def test_self_hosted_model_args_vllm_and_unknown_backend():
+    args = SelfHostedModelArgs(model_name="m", backend="vllm", max_new_tokens=10)
+    model = args.make_model()
+    assert isinstance(model, chat_api.VLLMChatModel)
+    assert model.max_tokens == 10
+
+    with pytest.raises(ValueError, match="not supported"):
+        SelfHostedModelArgs(model_name="m", backend="other").make_model()
+
+
+def test_chat_model_args_is_deprecated():
+    with pytest.warns(DeprecationWarning):
+        args = ChatModelArgs(model_name="m")
+    assert args.make_model() is None
+
+
+def _anthropic_response(text):
+    return SimpleNamespace(
+        content=[SimpleNamespace(text=text)],
+        usage=SimpleNamespace(input_tokens=1, output_tokens=2),
+    )
+
+
+def test_anthropic_chat_model_converts_messages():
+    model = AnthropicModelArgs(model_name="claude", max_new_tokens=20, temperature=0.1).make_model()
+    model.client = MagicMock()
+    model.client.messages.create.return_value = _anthropic_response("hey")
+
+    res = model([make_system_message("sys"), make_user_message("hi")])
+
+    assert res["content"] == "hey"
+    kwargs = model.client.messages.create.call_args.kwargs
+    assert kwargs["system"] == "sys"
+    assert kwargs["messages"] == [{"role": "user", "content": "hi"}]
+    assert kwargs["max_tokens"] == 20 and kwargs["temperature"] == 0.1
+
+
+def test_anthropic_chat_model_retries(no_sleep):
+    model = AnthropicChatModel("claude", api_key="k", max_retry=2)
+    model.client = MagicMock()
+    model.client.messages.create.side_effect = [RuntimeError("x"), _anthropic_response("ok")]
+    res = model([make_user_message("hi")], temperature=0)
+    assert res["content"] == "ok"
+    assert "system" not in model.client.messages.create.call_args.kwargs
+    assert no_sleep == [60]
+
+    model.client.messages.create.side_effect = RuntimeError("down")
+    with pytest.raises(RuntimeError, match="down"):
+        model([make_user_message("hi")])
diff --git a/tests/llm/test_huggingface_utils_unit.py b/tests/llm/test_huggingface_utils_unit.py
new file mode 100644
index 000000000..b2c4e4ce3
--- /dev/null
+++ b/tests/llm/test_huggingface_utils_unit.py
@@ -0,0 +1,144 @@
+"""Offline unit tests for HFBaseChatModel using a fake tokenizer and LLM."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from transformers import GPT2TokenizerFast
+
+from agentlab.llm import huggingface_utils
+from agentlab.llm.huggingface_utils import HFBaseChatModel, _prepend_system_to_first_user
+from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage
+from agentlab.llm.prompt_templates import STARCHAT_PROMPT_TEMPLATE
+
+
+class _FakeTokenizer:
+    def __init__(self, fail_with=None):
+        self.fail_with = fail_with
+        self.calls = []
+
+    def apply_chat_template(self, messages, tokenize=False):
+        self.calls.append([dict(m) for m in messages])
+        if self.fail_with and any(m["role"] == "system" for m in messages):
+            raise Exception(self.fail_with)
+        return "|".join(f"{m['role']}:{m['content']}" for m in messages)
+
+
+def _make_model(monkeypatch, tokenizer, base_model_name=None, n_retry=3, log_probs=False):
+    loaded = []
+
+    def from_pretrained(name):
+        loaded.append(name)
+        return tokenizer
+
+    monkeypatch.setattr(huggingface_utils.AutoTokenizer, "from_pretrained", from_pretrained)
+    monkeypatch.setattr(huggingface_utils.time, "sleep", lambda s: None)
+    model = HFBaseChatModel("model", base_model_name, n_retry, log_probs)
+    model.temperature = 0.5
+    model.llm = MagicMock(return_value="answer")
+    return model, loaded
+
+
+def _messages():
+    return [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "hi"},
+    ]
+
+
+def test_loads_tokenizer_from_base_model(monkeypatch):
+    _, loaded = _make_model(monkeypatch, _FakeTokenizer(), base_model_name="base")
+    assert loaded == ["base"]
+    _, loaded = _make_model(monkeypatch, _FakeTokenizer())
+    assert loaded == ["model"]
+
+
+def test_call_with_chat_template(monkeypatch):
+    model, _ = _make_model(monkeypatch, _FakeTokenizer())
+    res = model(_messages())
+    assert res["content"] == "answer"
+    model.llm.assert_called_once_with("system:sys|user:hi", temperature=0.5)
+    assert model._llm_type() == "huggingface"
+
+
+def test_call_merges_discussion(monkeypatch):
+    tokenizer = _FakeTokenizer()
+    model, _ = _make_model(monkeypatch, tokenizer)
+    msg = HumanMessage("a")
+    msg.add_text("b")
+    model(Discussion([SystemMessage("sys"), msg]), temperature=0.1)
+    assert tokenizer.calls[0][1]["content"] == "a\nb"
+    assert model.llm.call_args.kwargs["temperature"] == 0.1
+
+
+def test_call_falls_back_when_system_role_unsupported(monkeypatch):
+    tokenizer = _FakeTokenizer(fail_with="Conversation roles must alternate user/assistant")
+    model, _ = _make_model(monkeypatch, tokenizer)
+    model(_messages())
+    model.llm.assert_called_once_with("user:sys\nhi", temperature=0.5)
+
+
+def test_call_reraises_other_template_errors(monkeypatch):
+    model, _ = _make_model(monkeypatch, _FakeTokenizer(fail_with="other"))
+    with pytest.raises(Exception, match="other"):
+        model(_messages())
+
+
+def test_gpt2_tokenizer_uses_prompt_template(monkeypatch):
+    gpt2 = MagicMock(spec=GPT2TokenizerFast)
+    monkeypatch.setattr(huggingface_utils.AutoTokenizer, "from_pretrained", lambda n: gpt2)
+    model = HFBaseChatModel("bigcode/starcoder", None, 1, False)
+    model.temperature = 0.5
+    model.llm = MagicMock(return_value="x")
+    assert model.tokenizer is None
+    assert model.prompt_template is STARCHAT_PROMPT_TEMPLATE
+    model(_messages())
+    assert model.llm.call_args.args[0].startswith("<|system|>")
+
+
+def test_multiple_samples_and_retries(monkeypatch):
+    model, _ = _make_model(monkeypatch, _FakeTokenizer(), n_retry=3)
+    model.llm.side_effect = [RuntimeError("busy"), "a", "b"]
+    res = model(_messages(), n_samples=2)
+    assert [r["content"] for r in res] == ["a", "b"]
+    assert model.llm.call_count == 3
+
+
+def test_retries_exhausted(monkeypatch):
+    model, _ = _make_model(monkeypatch, _FakeTokenizer(), n_retry=2)
+    model.llm.side_effect = RuntimeError("down")
+    with pytest.raises(RuntimeError, match="down"):
+        model(_messages())
+    assert model.llm.call_count == 2
+
+
+def test_log_probs(monkeypatch):
+    model, _ = _make_model(monkeypatch, _FakeTokenizer(), log_probs=True)
+    model.llm.return_value = SimpleNamespace(generated_text="txt", details="details")
+    res = model(_messages())
+    assert res["content"] == "txt" and res["log_probs"] == "details"
+
+
+def test_prepend_system_to_first_user():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "q1"},
+        {"role": "user", "content": "q2"},
+    ]
+    out = _prepend_system_to_first_user(messages)
+    assert out == [{"role": "user", "content": "sys\nq1"}, {"role": "user", "content": "q2"}]
+
+
+def test_prepend_system_without_system_or_user():
+    no_system = [{"role": "user", "content": "q"}]
+    assert _prepend_system_to_first_user(no_system) == [{"role": "user", "content": "q"}]
+    only_system = [{"role": "system", "content": "s"}]
+    assert _prepend_system_to_first_user(only_system) == [{"role": "system", "content": "s"}]
+
+
+def test_prepend_system_with_column_remap():
+    messages = [{"r": "system", "t": "sys"}, {"r": "human", "t": "q"}]
+    out = _prepend_system_to_first_user(
+        messages, column_remap={"HumanMessage": "human", "role": "r", "text": "t"}
+    )
+    assert out == [{"r": "human", "t": "sys\nq"}]
diff --git a/tests/llm/test_litellm_api_unit.py b/tests/llm/test_litellm_api_unit.py
new file mode 100644
index 000000000..9064f20ec
--- /dev/null
+++ b/tests/llm/test_litellm_api_unit.py
@@ -0,0 +1,194 @@
+"""Offline unit tests for LiteLLMModel (the litellm client is replaced by a fake)."""
+
+import json
+
+import pytest
+from litellm import ModelResponse
+
+from agentlab.llm.litellm_api import LiteLLMAPIMessageBuilder, LiteLLMModel, LiteLLMModelArgs
+from agentlab.llm.response_api import APIPayload, ToolCall, ToolCalls
+
+RESPONSE_API_TOOL = {
+    "type": "function",
+    "name": "click",
+    "description": "Click an element.",
+    "parameters": {"type": "object", "properties": {"bid": {"type": "string"}}},
+}
+
+
+def _response(content="", tool_calls=None, reasoning=None):
+    message = {"role": "assistant", "content": content}
+    if tool_calls is not None:
+        message["tool_calls"] = [
+            {
+                "id": f"call_{i}",
+                "type": "function",
+                "function": {"name": name, "arguments": json.dumps(args)},
+            }
+            for i, (name, args) in enumerate(tool_calls)
+        ]
+    if reasoning is not None:
+        message["reasoning"] = reasoning
+    return ModelResponse(
+        choices=[{"index": 0, "message": message, "finish_reason": "stop"}],
+        model="gpt-4o-mini",
+        usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+    )
+
+
+@pytest.fixture
+def model():
+    m = LiteLLMModelArgs(
+        model_name="openai/gpt-4o-mini", max_new_tokens=50, temperature=0.2
+    ).make_model()
+    calls = []
+
+    def fake_client(**kwargs):
+        calls.append(kwargs)
+        return m._next_response
+
+    m.client = fake_client
+    m._next_response = _response("hello")
+    m.calls = calls
+    return m
+
+
+def _user_payload(**kwargs):
+    return APIPayload(messages=[LiteLLMAPIMessageBuilder.user().add_text("hi")], **kwargs)
+
+
+def test_model_args_builder():
+    assert LiteLLMModelArgs(model_name="x").get_message_builder() is LiteLLMAPIMessageBuilder
+
+
+def test_call_api_basic_params(model):
+    model._call_api(_user_payload())
+    params = model.calls[-1]
+    assert params["model"] == "openai/gpt-4o-mini"
+    assert params["temperature"] == 0.2
+    assert params["max_completion_tokens"] == 50
+    assert params["num_retries"] == 5
+    assert params["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
+    assert "tools" not in params and "tool_choice" not in params
+
+
+def test_call_api_omits_unset_params(model):
+    model.temperature = None
+    model.max_tokens = None
+    model._call_api(_user_payload())
+    assert "temperature" not in model.calls[-1]
+    assert "max_completion_tokens" not in model.calls[-1]
+
+
+def test_call_api_converts_tools_and_caching(model):
+    model._call_api(
+        _user_payload(
+            tools=[dict(RESPONSE_API_TOOL)],
+            tool_choice="any",
+            reasoning_effort="low",
+            cache_tool_definition=True,
+            cache_complete_prompt=True,
+        )
+    )
+    params = model.calls[-1]
+    assert params["tools"][0]["function"]["name"] == "click"
+    assert params["tools"][-1]["cache_control"] == {"type": "ephemeral"}
+    assert params["tool_choice"] == "required"
+    assert params["reasoning_effort"] == "low"
+    assert params["messages"][-1]["content"][-1]["cache_control"] == {"type": "ephemeral"}
+
+
+def test_call_api_keeps_chat_completion_tools_and_forces_tool(model):
+    chat_tool = {"type": "function", "function": {"name": "click"}}
+    model._call_api(_user_payload(tools=[chat_tool], force_call_tool="click"))
+    params = model.calls[-1]
+    assert params["tools"] == [chat_tool]
+    assert params["tool_choice"] == {"type": "function", "function": {"name": "click"}}
+
+    model._call_api(_user_payload(tool_choice="auto"))
+    assert model.calls[-1]["tool_choice"] == "auto"
+
+
+def test_api_payload_rejects_tool_choice_with_force():
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        APIPayload(tool_choice="auto", force_call_tool="click")
+
+
+def test_parse_response_single_tool_call(model):
+    model._next_response = _response("ok", tool_calls=[("click", {"bid": "a1"})])
+    out = model(_user_payload())
+    assert out.think == "ok"
+    assert out.action == "click(bid='a1')"
+    assert len(out.tool_calls) == 1
+    assert model.stats.stats_dict["n_api_calls"] == 1
+
+
+def test_parse_response_multiple_tool_calls(model):
+    calls = [("click", {"bid": "a1"}), ("fill", {"bid": "a2", "value": "x"})]
+    model._next_response = _response(tool_calls=calls)
+    out = model._parse_response(model._next_response)
+    assert out.action == "click(bid='a1')\nfill(bid='a2', value='x')"
+
+    model.use_only_first_toolcall = True
+    out = model._parse_response(model._next_response)
+    assert out.action == "click(bid='a1')"
+    assert len(out.tool_calls) == 1
+
+
+def test_parse_response_without_tool_calls(model):
+    out = model._parse_response(_response("just text"))
+    assert out.action is None and out.tool_calls is None
+
+    model.action_space_as_tools = False
+    assert model._parse_response(_response("x")).action is None
+
+
+def test_thinking_content_is_wrapped(model):
+    think = model._extract_thinking_content_from_response(_response("answer", reasoning="because"))
+    assert think == "<think>because</think>\nanswer"
+
+
+def test_format_tools_for_chat_completion():
+    assert LiteLLMModel.format_tools_for_chat_completion(None) is None
+    (tool,) = LiteLLMModel.format_tools_for_chat_completion([RESPONSE_API_TOOL])
+    assert tool == {
+        "type": "function",
+        "function": {
+            "name": "click",
+            "description": "Click an element.",
+            "parameters": RESPONSE_API_TOOL["parameters"],
+        },
+    }
+
+
+def _responded_tool_calls(n=2):
+    response = _response(tool_calls=[("click", {"bid": str(i)}) for i in range(n)])
+    raw = response.choices[0].message.to_dict()["tool_calls"]
+    calls = []
+    for i, raw_call in enumerate(raw):
+        tc = ToolCall(name="click", arguments={"bid": str(i)}, raw_call=raw_call)
+        tc.response_text(f"clicked {i}")
+        calls.append(tc)
+    return ToolCalls(tool_calls=calls, raw_calls=response)
+
+
+def test_message_builder_tool_responses():
+    builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(_responded_tool_calls())
+    out = builder.prepare_message()
+    assert len(out[0].tool_calls) == 2
+    assert out[1:] == [
+        {"name": "click", "role": "tool", "tool_call_id": "call_0", "content": "clicked 0"},
+        {"name": "click", "role": "tool", "tool_call_id": "call_1", "content": "clicked 1"},
+    ]
+
+
+def test_message_builder_only_first_tool_call():
+    builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(_responded_tool_calls())
+    out = builder.prepare_message(use_only_first_toolcall=True)
+    assert len(out[0].tool_calls) == 1
+
+
+def test_message_builder_tool_without_calls():
+    builder = LiteLLMAPIMessageBuilder("tool")
+    with pytest.raises(ValueError, match="No tool calls"):
+        builder.handle_tool_call()
diff --git a/tests/llm/test_llm_utils_extra.py b/tests/llm/test_llm_utils_extra.py
new file mode 100644
index 000000000..1b315a418
--- /dev/null
+++ b/tests/llm/test_llm_utils_extra.py
@@ -0,0 +1,467 @@
+"""Unit tests for llm_utils helpers that don't require network access."""
+
+import logging
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import anthropic
+import httpx
+import numpy as np
+import openai
+import pytest
+from PIL import Image
+
+from agentlab.llm import llm_utils
+from agentlab.llm.llm_utils import (
+    AIMessage,
+    BaseMessage,
+    Discussion,
+    HumanMessage,
+    ParseError,
+    SystemMessage,
+)
+
+
+@pytest.fixture
+def no_sleep(monkeypatch):
+    sleeps = []
+    monkeypatch.setattr(llm_utils.time, "sleep", lambda s: sleeps.append(s))
+    return sleeps
+
+
+class _RateLimit(Exception):
+    pass
+
+
+class _ApiError(Exception):
+    def __init__(self, status_code=None):
+        super().__init__(f"status {status_code}")
+        self.status_code = status_code
+
+
+def _call(client_function, **kwargs):
+    defaults = dict(
+        api_params={"x": 1},
+        is_response_valid_fn=lambda r: r == "ok",
+        rate_limit_exceptions=(_RateLimit,),
+        api_error_exceptions=(_ApiError,),
+        get_status_code_fn=lambda e: e.status_code,
+        max_retries=3,
+        initial_retry_delay_seconds=1,
+        max_retry_delay_seconds=3,
+    )
+    defaults.update(kwargs)
+    return llm_utils.generic_call_api_with_retries(client_function, **defaults)
+
+
+# ---------------------------------------------------------------------------
+# generic_call_api_with_retries
+# ---------------------------------------------------------------------------
+
+
+def test_generic_call_success_first_try(no_sleep):
+    fn = MagicMock(return_value="ok")
+    assert _call(fn) == "ok"
+    fn.assert_called_once_with(x=1)
+    assert no_sleep == []
+
+
+def test_generic_call_retries_invalid_response_then_succeeds(no_sleep):
+    fn = MagicMock(side_effect=["bad", "bad", "ok"])
+    assert _call(fn) == "ok"
+    assert fn.call_count == 3
+    # first invalid retry uses exponent 0, second one exponent 1
+    assert no_sleep == [1, 2]
+
+
+def test_generic_call_invalid_response_exhausts_retries(no_sleep):
+    fn = MagicMock(return_value="bad")
+    with pytest.raises(RuntimeError, match="after 3 retries"):
+        _call(fn)
+    assert fn.call_count == 3
+
+
+def test_generic_call_rate_limit_backoff_is_capped(no_sleep):
+    fn = MagicMock(side_effect=[_RateLimit(), _RateLimit(), "ok"])
+    assert _call(fn) == "ok"
+    assert no_sleep == [2, 3]  # 1*2**1, min(1*2**2, 3)
+
+
+def test_generic_call_rate_limit_exhausts_retries(no_sleep):
+    fn = MagicMock(side_effect=_RateLimit())
+    with pytest.raises(RuntimeError):
+        _call(fn)
+    assert fn.call_count == 3
+
+
+@pytest.mark.parametrize("status", [429, 500, 503])
+def test_generic_call_retriable_api_error(no_sleep, status):
+    fn = MagicMock(side_effect=[_ApiError(status), "ok"])
+    assert _call(fn) == "ok"
+    assert len(no_sleep) == 1
+
+
+def test_generic_call_retriable_api_error_exhausts_retries(no_sleep):
+    fn = MagicMock(side_effect=_ApiError(500))
+    with pytest.raises(RuntimeError):
+        _call(fn)
+    assert fn.call_count == 3
+
+
+@pytest.mark.parametrize("status", [400, None])
+def test_generic_call_non_retriable_api_error_is_raised(no_sleep, status):
+    err = _ApiError(status)
+    fn = MagicMock(side_effect=err)
+    with pytest.raises(_ApiError) as exc_info:
+        _call(fn)
+    assert exc_info.value is err
+    fn.assert_called_once()
+
+
+def test_generic_call_status_code_fn_failure_is_non_retriable(no_sleep):
+    def broken_status(e):
+        raise RuntimeError("nope")
+
+    fn = MagicMock(side_effect=_ApiError(500))
+    with pytest.raises(_ApiError):
+        _call(fn, get_status_code_fn=broken_status)
+
+
+def test_generic_call_unexpected_exception_is_raised(no_sleep):
+    fn = MagicMock(side_effect=KeyError("boom"))
+    with pytest.raises(KeyError):
+        _call(fn)
+    fn.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# OpenAI / Anthropic wrappers
+# ---------------------------------------------------------------------------
+
+
+def _openai_rate_limit_error():
+    request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions")
+    return openai.RateLimitError(
+        "rate limited", response=httpx.Response(429, request=request), body=None
+    )
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        SimpleNamespace(error=None, choices=[1]),
+        SimpleNamespace(error=None, output=[1]),
+    ],
+)
+def test_call_openai_api_with_retries_valid(no_sleep, response):
+    fn = MagicMock(return_value=response)
+    assert llm_utils.call_openai_api_with_retries(fn, {"model": "m"}) is response
+
+
+def test_call_openai_api_with_retries_retries_on_error_and_rate_limit(no_sleep):
+    good = SimpleNamespace(error=None, choices=[1])
+    fn = MagicMock(
+        side_effect=[
+            SimpleNamespace(error="oops"),
+            SimpleNamespace(error=None, choices=[]),
+            _openai_rate_limit_error(),
+            good,
+        ]
+    )
+    assert llm_utils.call_openai_api_with_retries(fn, {}, max_retries=4) is good
+    assert fn.call_count == 4
+    assert len(no_sleep) == 3
+
+
+def test_call_openai_api_with_retries_non_retriable_api_error(no_sleep):
+    request = httpx.Request("POST", "https://api.openai.com")
+    err = openai.APIError("bad request", request=request, body=None)
+    fn = MagicMock(side_effect=err)
+    with pytest.raises(openai.APIError):
+        llm_utils.call_openai_api_with_retries(fn, {})
+
+
+@pytest.mark.parametrize(
+    "response, valid",
+    [
+        (SimpleNamespace(type="message", content=[]), True),
+        (SimpleNamespace(type="message", content=None), False),
+        (SimpleNamespace(type="other"), False),
+        ({"type": "error", "error": "x"}, False),
+        (None, False),
+    ],
+)
+def test_call_anthropic_api_with_retries_validation(no_sleep, response, valid):
+    fn = MagicMock(return_value=response)
+    if valid:
+        assert llm_utils.call_anthropic_api_with_retries(fn, {}, max_retries=2) is response
+    else:
+        with pytest.raises(RuntimeError):
+            llm_utils.call_anthropic_api_with_retries(fn, {}, max_retries=2)
+        assert fn.call_count == 2
+
+
+def test_call_anthropic_api_with_retries_server_error_then_success(no_sleep):
+    request = httpx.Request("POST", "https://api.anthropic.com")
+    err = anthropic.InternalServerError(
+        "server", response=httpx.Response(500, request=request), body=None
+    )
+    good = SimpleNamespace(type="message", content=["hi"])
+    fn = MagicMock(side_effect=[err, good])
+    assert llm_utils.call_anthropic_api_with_retries(fn, {}) is good
+    assert len(no_sleep) == 1
+
+
+def test_supports_tool_calling_for_openrouter(monkeypatch):
+    client = MagicMock()
+    client.chat.completions.create.return_value.to_dict.return_value = {
+        "choices": [{"message": {"tool_calls": [{}]}}]
+    }
+    monkeypatch.setattr(openai, "Client", MagicMock(return_value=client))
+    assert llm_utils.supports_tool_calling_for_openrouter("some/model") is True
+
+    client.chat.completions.create.return_value.to_dict.return_value = {
+        "choices": [{"message": {"content": "no tools"}}]
+    }
+    assert llm_utils.supports_tool_calling_for_openrouter("some/model") is False
+
+    client.chat.completions.create.side_effect = RuntimeError("network")
+    # errors default to True so that the check never blocks a run
+    assert llm_utils.supports_tool_calling_for_openrouter("some/model") is True
+
+
+# ---------------------------------------------------------------------------
+# retry_multiple
+# ---------------------------------------------------------------------------
+
+
+def _parser(content):
+    if "good" not in content:
+        raise ParseError("not good")
+    return {"action": content}
+
+
+def test_retry_multiple_returns_valid_samples():
+    chat = MagicMock(
+        return_value=[
+            {"role": "assistant", "content": "bad"},
+            {"role": "assistant", "content": "good 1"},
+            {"role": "assistant", "content": "good 2"},
+        ]
+    )
+    messages = Discussion()
+    parsed, tries = llm_utils.retry_multiple(chat, messages, 3, _parser, num_samples=3)
+    assert parsed == [{"action": "good 1"}, {"action": "good 2"}]
+    assert tries == 0
+    chat.assert_called_once_with(messages, n_samples=3)
+    assert len(messages) == 1
+
+
+def test_retry_multiple_wraps_single_answer_and_retries():
+    chat = MagicMock(
+        side_effect=[
+            {"role": "assistant", "content": "bad"},
+            {"role": "assistant", "content": "good"},
+        ]
+    )
+    messages = Discussion()
+    parsed, tries = llm_utils.retry_multiple(chat, messages, 3, _parser)
+    assert parsed == [{"action": "good"}]
+    assert tries == 1
+    # assistant answer, user error feedback, assistant answer
+    assert [m["role"] for m in messages] == ["assistant", "user", "assistant"]
+    assert "not good" in messages[1]["content"]
+
+
+def test_retry_multiple_raises_after_n_retry():
+    chat = MagicMock(return_value={"role": "assistant", "content": "bad"})
+    with pytest.raises(ParseError):
+        llm_utils.retry_multiple(chat, Discussion(), 2, _parser)
+    assert chat.call_count == 2
+
+
+# ---------------------------------------------------------------------------
+# messages / parsing helpers
+# ---------------------------------------------------------------------------
+
+
+def test_messages_to_dict_handles_dicts_and_strings():
+    d = llm_utils.messages_to_dict([{"role": "user", "content": "hi"}, "raw text"])
+    assert isinstance(d, Discussion)
+    assert d[0]["role"] == "user"
+    assert d[1] == {"role": "<unknown role>", "content": "raw text"}
+
+
+def test_messages_to_dict_rejects_unknown_types():
+    with pytest.raises(ValueError, match="Unknown message type"):
+        llm_utils.messages_to_dict([42])
+
+
+def test_messages_to_dict_langchain_fallback(monkeypatch):
+    class FakeLC:
+        def __init__(self, type, content):
+            self.type = type
+            self.content = content
+
+    monkeypatch.setattr(llm_utils, "LangchainBaseMessage", FakeLC)
+    monkeypatch.setattr(llm_utils, "convert_message_to_dict", None)
+    d = llm_utils.messages_to_dict([FakeLC("human", "q"), FakeLC("ai", "a"), FakeLC("x", 1)])
+    assert [m["role"] for m in d] == ["user", "assistant", "x"]
+    assert d[2]["content"] == "1"
+
+
+def test_messages_to_dict_langchain_converter(monkeypatch):
+    class FakeLC:
+        pass
+
+    monkeypatch.setattr(llm_utils, "LangchainBaseMessage", FakeLC)
+    monkeypatch.setattr(
+        llm_utils, "convert_message_to_dict", lambda m: {"role": "system", "content": "c"}
+    )
+    d = llm_utils.messages_to_dict([FakeLC()])
+    assert d[0] == {"role": "system", "content": "c"}
+
+
+def test_parse_html_tags_optional_and_multiple():
+    text = "<a>1</a><a>2</a><b> x </b>"
+    content, valid, msg = llm_utils.parse_html_tags(text, keys=("a",), optional_keys=("b", "c"))
+    assert not valid
+    assert "multiple instances" in msg
+    assert content == {"a": "1", "b": "x"}
+
+    content, valid, msg = llm_utils.parse_html_tags(
+        text, keys=("a",), optional_keys=("c",), merge_multiple=True
+    )
+    assert valid and msg == ""
+    assert content["a"] == "1\n2"
+
+
+def test_parse_html_tags_raise():
+    assert llm_utils.parse_html_tags_raise("<k>v</k>", keys=("k",)) == {"k": "v"}
+    with pytest.raises(ParseError, match="Missing the key <k>"):
+        llm_utils.parse_html_tags_raise("nothing", keys=("k",))
+
+
+def test_yaml_parser_invalid():
+    with pytest.warns(UserWarning):
+        value, valid, retry_message = llm_utils.yaml_parser("a: [unclosed")
+    assert value == {} and not valid
+    assert "not a valid yaml" in retry_message
+
+
+def test_truncate_tokens_short_text_untouched():
+    assert llm_utils.truncate_tokens("hello world", max_tokens=100) == "hello world"
+
+
+def test_get_tokenizer_falls_back_to_gpt4(monkeypatch):
+    llm_utils.get_tokenizer.cache_clear()
+
+    def fail(*args, **kwargs):
+        raise OSError("no hub")
+
+    monkeypatch.setattr(llm_utils.AutoTokenizer, "from_pretrained", fail)
+    try:
+        enc = llm_utils.get_tokenizer("definitely-not-a-model")
+        assert enc.encode("hello") == llm_utils.get_tokenizer("gpt-4").encode("hello")
+    finally:
+        llm_utils.get_tokenizer.cache_clear()
+
+
+@pytest.mark.parametrize("mode", ["RGB", "RGBA"])
+def test_image_to_base64_urls(mode):
+    img = Image.new(mode, (4, 4))
+    assert llm_utils.image_to_jpg_base64_url(img).startswith("data:image/jpeg;base64,")
+    assert llm_utils.image_to_png_base64_url(img).startswith("data:image/png;base64,")
+
+
+def test_image_to_base64_urls_from_numpy():
+    arr = np.zeros((4, 4, 3), dtype=np.uint8)
+    assert llm_utils.image_to_jpg_base64_url(arr).startswith("data:image/jpeg;base64,")
+    assert llm_utils.image_to_png_base64_url(arr).startswith("data:image/png;base64,")
+
+
+# ---------------------------------------------------------------------------
+# BaseMessage / Discussion
+# ---------------------------------------------------------------------------
+
+
+def test_base_message_rejects_unknown_kwargs():
+    with pytest.raises(ValueError, match="Invalid attributes"):
+        BaseMessage("user", "x", foo=1)
+
+
+def test_message_subclasses_roles():
+    assert SystemMessage("s")["role"] == "system"
+    assert HumanMessage("h")["role"] == "user"
+    ai = AIMessage("a", log_probs=[0.1])
+    assert ai["role"] == "assistant" and ai["log_probs"] == [0.1]
+
+
+def test_base_message_content_is_copied():
+    content = [{"type": "text", "text": "a"}]
+    msg = HumanMessage(content)
+    content.append({"type": "text", "text": "b"})
+    assert len(msg["content"]) == 1
+
+
+def test_base_message_str_with_images(caplog):
+    msg = HumanMessage("hello")
+    msg.add_image("data:image/png;base64,abc")
+    msg["content"].append({"type": "input_text", "text": "world"})
+    with caplog.at_level(logging.INFO):
+        assert str(msg) == "hello\nworld"
+    assert "images" in caplog.text
+    assert msg.__str__(warn_if_image=True) == "hello\nworld"
+
+
+def test_add_image_from_pil_with_detail():
+    msg = HumanMessage("look")
+    msg.add_image(Image.new("RGB", (2, 2)), detail="low")
+    image_elem = msg["content"][-1]
+    assert image_elem["type"] == "image_url"
+    assert image_elem["image_url"]["detail"] == "low"
+    assert image_elem["image_url"]["url"].startswith("data:image/jpeg;base64,")
+
+
+def test_base_message_to_markdown():
+    assert HumanMessage("x").to_markdown() == "\n```\nx\n```\n"
+    msg = HumanMessage("text")
+    msg.add_image("http://img")
+    msg.add_content("image_url", "http://img2")
+    md = msg.to_markdown()
+    assert "```\ntext\n```" in md
+    assert "![image](http://img)" in md
+    assert "![image](http://img2)" in md
+
+
+def test_base_message_merge_on_string_is_noop():
+    msg = HumanMessage("x")
+    msg.merge()
+    assert msg["content"] == "x"
+
+
+def test_discussion_api():
+    d = Discussion(SystemMessage("sys"))
+    d.add_message(role="user", content="hi")
+    d.add_text("there")
+    d.add_content("text", "again")
+    d.add_image("http://img")
+    assert len(d) == 2
+    assert d.last_message is d[-1]
+    assert [m["role"] for m in d] == ["system", "user"]
+
+    md = d.to_markdown()
+    assert md.startswith("Message 0")
+    assert "Message 1" in md
+    assert "![image](http://img)" in md
+
+    openai_msgs = d.to_openai()
+    assert openai_msgs[1]["content"][0] == {"type": "text", "text": "hi\nthere\nagain"}
+    assert "sys" in d.to_string()
+
+
+def test_discussion_default_empty():
+    d = Discussion()
+    assert len(d) == 0
+    d.append({"role": "user", "content": "x"})
+    assert isinstance(d[0], BaseMessage)
diff --git a/tests/llm/test_prompt_templates.py b/tests/llm/test_prompt_templates.py
new file mode 100644
index 000000000..9563d838b
--- /dev/null
+++ b/tests/llm/test_prompt_templates.py
@@ -0,0 +1,50 @@
+import pytest
+
+from agentlab.llm.prompt_templates import (
+    STARCHAT_PROMPT_TEMPLATE,
+    PromptTemplate,
+    get_prompt_template,
+)
+
+
+def test_format_message_per_role():
+    t = PromptTemplate(system="S:{input}|", human="H:{input}|", ai="A:{input}|")
+    assert t.format_message({"role": "system", "content": "a"}) == "S:a|"
+    assert t.format_message({"role": "user", "content": "b"}) == "H:b|"
+    assert t.format_message({"role": "assistant", "content": "c"}) == "A:c|"
+
+
+def test_format_message_unknown_role():
+    with pytest.raises(ValueError, match="not supported"):
+        STARCHAT_PROMPT_TEMPLATE.format_message({"role": "tool", "content": "x"})
+
+
+def test_construct_prompt_starchat():
+    messages = [
+        {"role": "system", "content": "sys"},
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "hello"},
+    ]
+    prompt = STARCHAT_PROMPT_TEMPLATE.construct_prompt(messages)
+    assert prompt == (
+        "<|system|>\nsys<|end|>\n"
+        "<|user|>\nhi<|end|>\n"
+        "<|assistant|>\nhello<|end|>\n"
+        "<|assistant|>"
+    )
+
+
+@pytest.mark.parametrize("bad", [["not a dict"], [{"role": "user"}], [{"content": "x"}]])
+def test_construct_prompt_rejects_non_openai_messages(bad):
+    with pytest.raises(ValueError, match="openai format"):
+        STARCHAT_PROMPT_TEMPLATE.construct_prompt(bad)
+
+
+@pytest.mark.parametrize("name", ["bigcode/starcoder", "HuggingFaceH4/starchat-beta"])
+def test_get_prompt_template_known(name):
+    assert get_prompt_template(name) is STARCHAT_PROMPT_TEMPLATE
+
+
+def test_get_prompt_template_unknown():
+    with pytest.raises(NotImplementedError):
+        get_prompt_template("meta-llama/Llama-3")