chore: import upstream snapshot with attribution

tests/packaged_modules/test_json.py

@@ -0,0 +1,809 @@
+import json
+import textwrap
+
+import pyarrow as pa
+import pytest
+
+from datasets import Features, Value, load_dataset
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.json.json import AGENT_TRACES_FEATURES, Json, JsonConfig
+
+from ..utils import require_teich
+
+
+@pytest.fixture
+def json_file(tmp_path):
+    filename = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "col_1": 1,
+            "col_2": 2
+        }
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+# ndjson format is no longer maintained (see: https://github.com/ndjson/ndjson-spec/issues/35#issuecomment-1285673417)
+@pytest.fixture
+def ndjson_file(tmp_path):
+    filename = tmp_path / "file.ndjson"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_utf16_encoded(tmp_path):
+    filename = tmp_path / "file_utf16_encoded.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w", encoding="utf-16") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts(tmp_path):
+    filename = tmp_path / "file_with_list_of_dicts.json"
+    data = textwrap.dedent(
+        """\
+        [
+            {"col_1": -1},
+            {"col_1": 1, "col_2": 2},
+            {"col_1": 10, "col_2": 20}
+        ]
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_strings(tmp_path):
+    filename = tmp_path / "file_with_list_of_strings.json"
+    data = textwrap.dedent(
+        """\
+        [
+            "First text.",
+            "Second text.",
+            "Third text."
+        ]
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_field(tmp_path):
+    filename = tmp_path / "file_with_list_of_dicts_field.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                {"col_1": -1},
+                {"col_1": 1, "col_2": 2},
+                {"col_1": 10, "col_2": 20}
+            ]
+        }
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_strings_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                "First text.",
+                "Second text.",
+                "Third text."
+            ]
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_dict_of_lists_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": {
+                "col_1": [-1, 1, 10],
+                "col_2": [null, 2, 20]
+            }
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_with_sorted_columns(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        [
+            {"ID": 0, "Language": "Language-0", "Topic": "Topic-0"},
+            {"ID": 1, "Language": "Language-1", "Topic": "Topic-1"},
+            {"ID": 2, "Language": "Language-2", "Topic": "Topic-2"}
+        ]
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_with_sorted_columns_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                {"ID": 0, "Language": "Language-0", "Topic": "Topic-0"},
+                {"ID": 1, "Language": "Language-1", "Topic": "Topic-1"},
+                {"ID": 2, "Language": "Language-2", "Topic": "Topic-2"}
+            ]
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def jsonl_file_with_mix_of_str_and_int(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1}
+        {"col_1": "foo"}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_dicts_of_varying_keys(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": {"a": 0}}
+        {"col_1": {"b": 0}}
+        {"col_1": {"c": 0}}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_lists_of_dicts_of_varying_keys(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": [{"a": 0}, {"b": 0}]}
+        {"col_1": [{"c": 0}, {"d": 0}]}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_lists_of_dicts_of_varying_keys_and_bom(tmp_path):
+    # Same content as jsonl_file_with_lists_of_dicts_of_varying_keys, but the file
+    # starts with a UTF-8 BOM (written via the utf-8-sig codec).
+    filename = tmp_path / "file_with_bom.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": [{"a": 0}, {"b": 0}]}
+        {"col_1": [{"c": 0}, {"d": 0}]}
+        """
+    )
+    with open(filename, "w", encoding="utf-8-sig") as f:
+        f.write(data)
+    return str(filename)
+
+
+_messages = [
+    {"role": "user", "content": "Turn on the living room lights and play my electronic music playlist."},
+    {
+        "role": "assistant",
+        "tool_calls": [
+            {
+                "type": "function",
+                "function": {"name": "control_light", "arguments": {"room": "living room", "state": "on"}},
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "play_music",
+                    "arguments": {
+                        "playlist": "electronic"
+                    },  # mixed-type here since keys ["playlist"] and ["room", "state"] are different
+                },
+            },
+        ],
+    },
+    {"role": "tool", "name": "control_light", "content": "The lights in the living room are now on."},
+    {"role": "tool", "name": "play_music", "content": "The music is now playing."},
+    {"role": "assistant", "content": "Done!"},
+]
+
+EXPECTED_ONE = {"col_1": [1], "col_2": [2]}
+EXPECTED_THREE = {"col_1": [-1, 1, 10], "col_2": [None, 2, 20]}
+EXPECTED_LIST_OF_STRINGS = {"text": ["First text.", "Second text.", "Third text."]}
+EXPECTED_MIX = {"col_1": [-1, 1, "foo"]}
+EXPECTED_DICTS_WITH_VARYING_KEYS = {"col_1": [{"a": 0}, {"b": 0}, {"c": 0}]}
+EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS = {"col_1": [[{"a": 0}, {"b": 0}], [{"c": 0}, {"d": 0}]]}
+EXPECTED_MESSAGES = {"messages": [_messages]}
+
+
+@pytest.fixture
+def jsonl_file_with_messages(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = json.dumps({"messages": _messages})
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+def write_jsonl(path, rows):
+    with open(path, "w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+    return str(path)
+
+
+def generate_agent_traces_output(trace_file):
+    json_builder = Json(features=AGENT_TRACES_FEATURES)
+    base_files = [trace_file]
+    files_iterables = [[trace_file]]
+    original_files = list(base_files)
+    generator = json_builder._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    return Features.from_arrow_schema(pa_table.schema).decode_batch(pa_table.to_pydict())
+
+
+def assert_agent_traces_output(tmp_path, filename, rows, expected, num_sessions=1):
+    trace_file = write_jsonl(tmp_path / filename, rows)
+    out = generate_agent_traces_output(trace_file)
+    for key, value in zip(AGENT_TRACE_FIELD_NAMES_TO_CHECK, expected):
+        assert out[key] == [value] * num_sessions
+    assert out["file_path"] == [trace_file] * num_sessions
+    assert isinstance(out["messages"], list)
+    assert out["messages"]
+    assert all(
+        isinstance(message["role"], str) and isinstance(message["content"], str)
+        for messages in out["messages"]
+        for message in messages
+    )
+    assert isinstance(out["tools"], list)
+    assert isinstance(out["metadata"], (dict, list))
+    assert isinstance(out["trace"], (dict, list))
+    return trace_file, out
+
+
+AGENT_TRACE_FIELD_NAMES_TO_CHECK = (
+    "harness",
+    "session_id",
+    "prompt",
+    "sent_at",
+    "num_user_messages",
+    "num_tool_calls",
+)
+CODEX_AGENT_TRACE_ROWS = [
+    {"type": "session_meta", "payload": {"id": "codex-session"}},
+    {
+        "type": "response_item",
+        "payload": {
+            "type": "message",
+            "role": "user",
+            "content": [{"type": "input_text", "text": "context-wrapped codex prompt"}],
+        },
+    },
+    {
+        "type": "response_item",
+        "payload": {
+            "type": "message",
+            "role": "user",
+            "content": [{"type": "input_text", "text": "actual codex prompt"}],
+        },
+    },
+    {
+        "timestamp": "2026-04-01T10:01:00.000Z",
+        "type": "event_msg",
+        "payload": {"type": "user_message", "message": "actual codex prompt"},
+    },
+    {
+        "type": "response_item",
+        "payload": {"type": "function_call", "name": "exec_command", "call_id": "call_1"},
+    },
+    {
+        "type": "response_item",
+        "payload": {"type": "function_call_output", "call_id": "call_1", "output": "done"},
+    },
+]
+CODEX_EXPECTED_AGENT_TRACE_FIELDS = (
+    "codex",
+    "codex-session",
+    "actual codex prompt",
+    "2026-04-01T10:01:00.000Z",
+    1,
+    1,
+)
+
+HERMES_SESSION = {
+    "id": "20260605_092247_d018ec",
+    "source": "cli",
+    "model": "Qwen/Qwen3.5-35B-A3B",
+    "system_prompt": "You are Hermes.",
+    "started_at": 1_780_665_768.307,
+    "message_count": 4,
+    "tool_call_count": 1,
+    "messages": [
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "user",
+            "content": "Run pwd and date.",
+            "timestamp": 1_780_665_767.307,
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "assistant",
+            "content": "",
+            "reasoning_content": "The user asked for two shell commands.",
+            "timestamp": 1_780_665_767.308,
+            "tool_calls": [
+                {
+                    "id": "call_677f321e2b3047b4b8c7a1e1",
+                    "type": "function",
+                    "function": {
+                        "name": "terminal",
+                        "arguments": json.dumps({"command": "pwd && date -u +%Y-%m-%dT%H:%M:%SZ"}),
+                    },
+                },
+            ],
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "tool",
+            "content": json.dumps({"output": "/tmp/work\n2026-06-05T13:22:47Z", "exit_code": 0, "error": None}),
+            "tool_call_id": "call_677f321e2b3047b4b8c7a1e1",
+            "timestamp": 1_780_665_767.309,
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "assistant",
+            "content": "Working directory: `/tmp/work`.",
+            "reasoning": "The commands executed successfully.",
+            "timestamp": 1_780_665_767.31,
+        },
+    ],
+}
+
+
+DROID_SESSION = [
+    {
+        "type": "session_start",
+        "id": "droid-session",
+        "title": "inspect the project",
+        "sessionTitle": "Inspect project files",
+        "owner": "caleb",
+        "version": 2,
+        "cwd": "/workspace/project",
+    },
+    {
+        "type": "message",
+        "id": "context-1",
+        "timestamp": "2026-06-02T18:55:29.000Z",
+        "message": {
+            "role": "user",
+            "visibility": "llm_only",
+            "content": [{"type": "text", "text": "<system-reminder>injected context</system-reminder>"}],
+        },
+        "parentId": None,
+    },
+    {
+        "type": "message",
+        "id": "message-1",
+        "timestamp": "2026-06-02T18:55:30.274Z",
+        "message": {
+            "role": "user",
+            "content": [{"type": "text", "text": "Inspect the project"}],
+        },
+        "parentId": "context-1",
+    },
+    {
+        "type": "message",
+        "id": "message-2",
+        "timestamp": "2026-06-02T18:55:35.000Z",
+        "message": {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "thinking",
+                    "thinking": "I should list the files first.",
+                    "signature": "reasoning_content",
+                    "signatureProvider": "generic-chat-completion-api",
+                    "durationMs": 1200,
+                },
+                {"type": "text", "text": "I'll list the files."},
+                {"type": "tool_use", "id": "LS_0", "name": "LS", "input": {"directory_path": "/workspace/project"}},
+            ],
+            "chatCompletionReasoningField": "reasoning_content",
+            "chatCompletionReasoningContent": "I should list the files first.",
+        },
+        "parentId": "message-1",
+    },
+    {
+        "type": "message",
+        "id": "message-3",
+        "timestamp": "2026-06-02T18:55:36.000Z",
+        "message": {
+            "role": "user",
+            "content": [
+                {"type": "tool_result", "tool_use_id": "LS_0", "is_error": False, "content": "README.md\nsrc"},
+            ],
+        },
+        "parentId": "message-2",
+    },
+]
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = JsonConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = JsonConfig(name="name", data_files=data_files)
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs, expected",
+    [
+        ("json_file", {}, EXPECTED_ONE),
+        ("jsonl_file", {}, EXPECTED_THREE),
+        ("ndjson_file", {}, EXPECTED_THREE),
+        ("jsonl_file_utf16_encoded", {"encoding": "utf-16"}, EXPECTED_THREE),
+        ("json_file_with_list_of_dicts", {}, EXPECTED_THREE),
+        ("json_file_with_list_of_dicts_field", {"field": "field3"}, EXPECTED_THREE),
+        ("json_file_with_list_of_strings", {}, EXPECTED_LIST_OF_STRINGS),
+        ("json_file_with_list_of_strings_field", {"field": "field3"}, EXPECTED_LIST_OF_STRINGS),
+        ("json_file_with_dict_of_lists_field", {"field": "field3"}, EXPECTED_THREE),
+        ("jsonl_file_with_mix_of_str_and_int", {}, EXPECTED_MIX),
+        ("jsonl_file_with_dicts_of_varying_keys", {}, EXPECTED_DICTS_WITH_VARYING_KEYS),
+        ("jsonl_file_with_lists_of_dicts_of_varying_keys", {}, EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS),
+        (
+            "jsonl_file_with_lists_of_dicts_of_varying_keys_and_bom",
+            {},
+            EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS,
+        ),
+        ("jsonl_file_with_messages", {}, EXPECTED_MESSAGES),
+    ],
+)
+def test_json_generate_tables(file_fixture, config_kwargs, expected, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    out = Features.from_arrow_schema(pa_table.schema).decode_batch(pa_table.to_pydict())
+    assert out == expected
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs",
+    [
+        (
+            "jsonl_file",
+            {"features": Features({"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")})},
+        ),
+        (
+            "json_file_with_list_of_dicts",
+            {"features": Features({"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")})},
+        ),
+        (
+            "json_file_with_list_of_dicts_field",
+            {
+                "field": "field3",
+                "features": Features(
+                    {"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")}
+                ),
+            },
+        ),
+    ],
+)
+def test_json_generate_tables_with_missing_features(file_fixture, config_kwargs, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20], "missing_col": [None, None, None]}
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs",
+    [
+        ("json_file_with_list_of_dicts_with_sorted_columns", {}),
+        ("json_file_with_list_of_dicts_with_sorted_columns_field", {"field": "field3"}),
+    ],
+)
+def test_json_generate_tables_with_sorted_columns(file_fixture, config_kwargs, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.column_names == ["ID", "Language", "Topic"]
+
+
+@pytest.mark.parametrize(
+    "filename, rows, expected",
+    [
+        pytest.param("codex.jsonl", CODEX_AGENT_TRACE_ROWS, CODEX_EXPECTED_AGENT_TRACE_FIELDS, id="codex"),
+        pytest.param(
+            "codex_response_item_only.jsonl",
+            [
+                {"type": "session_meta", "payload": {"id": "codex-session"}},
+                {
+                    "type": "response_item",
+                    "payload": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [{"text": "codex response item prompt"}],
+                    },
+                },
+                {
+                    "type": "response_item",
+                    "payload": {"type": "function_call", "name": "exec_command", "call_id": "call_1"},
+                },
+            ],
+            ("codex", "codex-session", None, None, 0, 1),
+            id="codex-response-item-only",
+        ),
+        pytest.param(
+            "claude.jsonl",
+            [
+                {
+                    "timestamp": "2026-04-02T10:00:00.000Z",
+                    "type": "user",
+                    "sessionId": "claude-session",
+                    "message": {"role": "user", "content": "claude prompt"},
+                },
+                {
+                    "type": "assistant",
+                    "sessionId": "claude-session",
+                    "message": {
+                        "role": "assistant",
+                        "content": [{"type": "tool_use", "id": "toolu_1"}, {"type": "tool_use", "id": "toolu_2"}],
+                    },
+                },
+                {
+                    "type": "user",
+                    "sessionId": "claude-session",
+                    "message": {"role": "user", "content": [{"type": "tool_result", "content": "done"}]},
+                },
+            ],
+            ("claude_code", "claude-session", "claude prompt", "2026-04-02T10:00:00.000Z", 1, 2),
+            id="claude",
+        ),
+        pytest.param(
+            "pi.jsonl",
+            [
+                {"type": "session", "id": "pi-session"},
+                {
+                    "timestamp": "2026-04-03T10:01:00.000Z",
+                    "type": "message",
+                    "message": {"role": "user", "content": "pi prompt"},
+                },
+                {
+                    "type": "message",
+                    "message": {
+                        "role": "assistant",
+                        "content": [{"type": "toolCall", "id": "call_1"}, {"type": "toolCall", "id": "call_2"}],
+                    },
+                },
+                {
+                    "type": "message",
+                    "message": {"role": "toolResult", "content": [{"type": "text", "text": "done"}]},
+                },
+                {
+                    "type": "message",
+                    "message": {"role": "user", "content": [{"type": "text", "text": "second pi prompt"}]},
+                },
+            ],
+            ("pi", "pi-session", "pi prompt", "2026-04-03T10:01:00.000Z", 2, 2),
+            id="pi",
+        ),
+        pytest.param(
+            "openclaw.jsonl",
+            [
+                {
+                    "type": "session",
+                    "id": "openclaw-session",
+                    "cwd": "/Users/test/.openclaw/agents/main",
+                },
+                {
+                    "timestamp": "2026-04-03T10:01:00.000Z",
+                    "type": "message",
+                    "message": {"role": "user", "content": "openclaw prompt"},
+                },
+            ],
+            ("openclaw", "openclaw-session", "openclaw prompt", "2026-04-03T10:01:00.000Z", 1, 0),
+            id="openclaw",
+        ),
+        pytest.param(
+            "hermes.jsonl",
+            [HERMES_SESSION],
+            ("hermes", "20260605_092247_d018ec", "Run pwd and date.", "2026-06-05T13:22:48.307Z", 1, 1),
+        ),
+        pytest.param(
+            "hermes_two_sessions.jsonl",
+            [HERMES_SESSION] * 2,
+            ("hermes", "20260605_092247_d018ec", "Run pwd and date.", "2026-06-05T13:22:48.307Z", 1, 1),
+        ),
+        pytest.param(
+            "droid.jsonl",
+            DROID_SESSION,
+            ("droid", "droid-session", "Inspect the project", "2026-06-02T18:55:30.274Z", 1, 1),
+            id="droid",
+        ),
+        pytest.param(
+            "missing_prompt.jsonl",
+            [
+                {"type": "session_meta", "payload": {"id": "codex-session"}},
+                {
+                    "type": "response_item",
+                    "payload": {"type": "message", "role": "assistant", "content": [{"text": "assistant response"}]},
+                },
+            ],
+            ("codex", "codex-session", None, None, 0, 0),
+            id="missing-prompt",
+        ),
+    ],
+)
+@require_teich
+def test_json_generate_tables_with_agent_trace_metadata(tmp_path, filename, rows, expected):
+    num_sessions = 2 if filename == "hermes_two_sessions.jsonl" else 1
+    _, out = assert_agent_traces_output(tmp_path, filename, rows, expected, num_sessions=num_sessions)
+    if filename == "droid.jsonl":
+        assert out["metadata"][0]["trace_type"] == "droid"
+    assert "models" not in out
+
+
+@require_teich
+@pytest.mark.parametrize(
+    "filename, rows, expected",
+    [
+        pytest.param("codex.jsonl", CODEX_AGENT_TRACE_ROWS, CODEX_EXPECTED_AGENT_TRACE_FIELDS, id="codex"),
+        pytest.param(
+            "droid.jsonl",
+            DROID_SESSION,
+            ("droid", "droid-session", "Inspect the project", "2026-06-02T18:55:30.274Z", 1, 1),
+            id="droid",
+        ),
+    ],
+)
+def test_json_load_dataset_with_agent_trace_metadata(tmp_path, filename, rows, expected):
+    trace_file = write_jsonl(tmp_path / filename, rows)
+
+    dataset = load_dataset("json", data_files=trace_file, split="train", cache_dir=str(tmp_path / "cache"))
+    row = dataset[0]
+
+    assert dataset.column_names == [
+        "harness",
+        "session_id",
+        "prompt",
+        "messages",
+        "tools",
+        "metadata",
+        "sent_at",
+        "num_user_messages",
+        "num_tool_calls",
+        "trace",
+        "file_path",
+    ]
+    for key, value in zip(AGENT_TRACE_FIELD_NAMES_TO_CHECK, expected):
+        assert row[key] == value, key
+    if filename == "droid.jsonl":
+        assert row["metadata"]["trace_type"] == "droid"
+        assert json.loads(row["trace"].splitlines()[0])["type"] == "session_start"
+
+
+def test_json_load_dataset_without_droid_marker_stays_ordinary_json(tmp_path):
+    trace_file = write_jsonl(
+        tmp_path / "droid_missing_marker.jsonl",
+        [
+            {"type": "session_start", "id": "droid-session", "version": 2},
+            {
+                "type": "message",
+                "id": "message-1",
+                "timestamp": "2026-06-02T18:55:30.274Z",
+                "message": {"role": "user", "content": [{"type": "text", "text": "Inspect the project"}]},
+            },
+        ],
+    )
+
+    dataset = load_dataset("json", data_files=trace_file, split="train", cache_dir=str(tmp_path / "cache"))
+
+    assert dataset.column_names == ["type", "id", "version", "timestamp", "message"]
+    assert dataset[0]["type"] == "session_start"