Skip to content

Commit 7c6fa96

Browse files
feat: narrow MessageDict.role to RoleLiteral
1 parent 0435b2d commit 7c6fa96

6 files changed

Lines changed: 94 additions & 9 deletions

File tree

models/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from models.session import (
99
MessageDict,
1010
QuickSessionInfoDict,
11+
RoleLiteral,
1112
SessionDict,
1213
SessionMetadataDict,
1314
ToolUseDict,
@@ -23,6 +24,7 @@
2324
"ProjectDict",
2425
"ProjectSessionRowDict",
2526
"QuickSessionInfoDict",
27+
"RoleLiteral",
2628
"SearchHitDict",
2729
"SessionDict",
2830
"SessionListItemDict",

models/search.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22

33
from typing import TypedDict
44

5+
from models.session import RoleLiteral
6+
57

68
class SearchHitDict(TypedDict):
79
project: str
810
session_id: str
911
title: str
10-
role: str
12+
role: RoleLiteral
1113
timestamp: str | None
1214
snippet: str

models/session.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
from models.record_data import RecordDataUnion
66
from models.tool_results import ToolNameLiteral, ToolResultUnion
77

8+
RoleLiteral = Literal["user", "assistant", "system", "result", "progress"]
9+
810

911
class ToolUseDict(TypedDict, total=False):
1012
id: str
@@ -25,7 +27,7 @@ class MessageUsageDict(TypedDict, total=False):
2527

2628

2729
class MessageDict(TypedDict):
28-
role: str
30+
role: RoleLiteral
2931
uuid: NotRequired[str | None]
3032
parent_uuid: NotRequired[str | None]
3133
timestamp: NotRequired[str | None]

tests/test_jsonl_parser.py

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -737,16 +737,19 @@ def test_empty_file_returns_skeleton(self):
737737
finally:
738738
os.unlink(path)
739739

740-
def test_unknown_entry_type_silently_ignored(self):
740+
def test_unknown_entry_type_maps_to_system(self, caplog):
741741
path = _write_jsonl(
742742
[
743743
{"type": "custom", "timestamp": "2026-01-01T00:00:00Z"},
744744
]
745745
)
746746
try:
747-
s = parse_session(path)
748-
assert s["messages"] == []
747+
with caplog.at_level("WARNING"):
748+
s = parse_session(path)
749+
assert len(s["messages"]) == 1
750+
assert s["messages"][0]["role"] == "system"
749751
assert s["metadata"]["entry_counts"].get("custom") == 1
752+
assert "Unknown message role" in caplog.text
750753
finally:
751754
os.unlink(path)
752755

@@ -1034,3 +1037,48 @@ def test_tool_use_result_string_returns_none(self):
10341037
meta,
10351038
)
10361039
assert messages[0]["tool_result_parsed"] is None
1040+
1041+
1042+
# ---------------------------------------------------------------------------
1043+
# Unknown role coercion
1044+
# ---------------------------------------------------------------------------
1045+
1046+
1047+
class TestUnknownRoleCoercion:
1048+
def test_unknown_entry_type_maps_to_system_with_warning(self, caplog):
1049+
path = _write_jsonl(
1050+
[
1051+
{
1052+
"type": "mystery_future_type",
1053+
"timestamp": "2026-01-01T00:00:00Z",
1054+
"content": "forward-compat payload",
1055+
}
1056+
]
1057+
)
1058+
try:
1059+
with caplog.at_level("WARNING"):
1060+
s = parse_session(path)
1061+
assert len(s["messages"]) == 1
1062+
assert s["messages"][0]["role"] == "system"
1063+
assert s["messages"][0]["content"] == "forward-compat payload"
1064+
assert "Unknown message role" in caplog.text
1065+
assert "mystery_future_type" in caplog.text
1066+
finally:
1067+
os.unlink(path)
1068+
1069+
def test_valid_unhandled_result_type_emits_result_role(self):
1070+
path = _write_jsonl(
1071+
[
1072+
{
1073+
"type": "result",
1074+
"timestamp": "2026-01-01T00:00:00Z",
1075+
"content": "task outcome",
1076+
}
1077+
]
1078+
)
1079+
try:
1080+
s = parse_session(path)
1081+
assert len(s["messages"]) == 1
1082+
assert s["messages"][0]["role"] == "result"
1083+
finally:
1084+
os.unlink(path)

tests/test_parser_fuzz.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,8 +240,10 @@ def test_unknown_record_type_is_graceful(tmp_path: Path) -> None:
240240
path = _write_jsonl(tmp_path / "unknown.jsonl", lines)
241241
session = parse_session(path)
242242
assert session["metadata"]["entry_counts"].get("totally-new-claude-record") == 1
243-
# Unknown type produces no message; only the valid user line does.
244-
assert len(session["messages"]) == 1
243+
# Unknown type is coerced to a system message; the valid user line follows.
244+
assert len(session["messages"]) == 2
245+
assert session["messages"][0]["role"] == "system"
246+
assert session["messages"][1]["role"] == "user"
245247

246248

247249
def test_non_numeric_usage_tokens_do_not_crash(tmp_path: Path) -> None:

utils/jsonl_parser.py

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,14 @@
22
actually work with -- messages, tool calls, token counts, file activity, etc."""
33

44
import json
5+
import logging
56
import math
67
import os
78
from datetime import datetime
8-
from typing import Any
9+
from typing import Any, cast, get_args
910

1011
from models.record_data import RecordDataUnion
11-
from models.session import MessageDict, SessionDict, ToolUseDict
12+
from models.session import MessageDict, RoleLiteral, SessionDict, ToolUseDict
1213
from models.tool_results import ToolResultUnion, is_tool_result_dict
1314
from utils.jsonl_helpers import (
1415
entry_message as _entry_message,
@@ -40,6 +41,30 @@
4041
"_track_file_activity",
4142
]
4243

44+
_HANDLED_ENTRY_TYPES = frozenset({"user", "assistant", "system", "progress"})
45+
_VALID_ROLES = frozenset(get_args(RoleLiteral))
46+
47+
48+
def _coerce_role(raw: str) -> RoleLiteral:
49+
if raw in _VALID_ROLES:
50+
return cast(RoleLiteral, raw)
51+
logging.warning("Unknown message role %r; mapping to 'system'", raw)
52+
return "system"
53+
54+
55+
def _fallback_message(entry: dict[str, Any], role: RoleLiteral) -> MessageDict:
56+
"""Minimal message for JSONL entry types without a dedicated processor."""
57+
content = entry.get("content", "")
58+
text = content if isinstance(content, str) else (str(content) if content is not None else "")
59+
return {
60+
"role": role,
61+
"uuid": entry.get("uuid"),
62+
"parent_uuid": entry.get("parentUuid"),
63+
"timestamp": entry.get("timestamp"),
64+
"content": text,
65+
"is_sidechain": entry.get("isSidechain", False),
66+
}
67+
4368

4469
def _safe_int(val: Any) -> int:
4570
"""Coerce a value to a non-negative int for token accounting; non-numeric,
@@ -144,6 +169,10 @@ def parse_session(filepath: str) -> SessionDict:
144169
_process_system(entry, messages, metadata)
145170
elif entry_type == "progress":
146171
_process_progress(entry, messages)
172+
elif entry_type:
173+
type_str = entry_type if isinstance(entry_type, str) else str(entry_type)
174+
if type_str not in _HANDLED_ENTRY_TYPES:
175+
messages.append(_fallback_message(entry, _coerce_role(type_str)))
147176

148177
metadata["models_used"] = sorted(metadata["models_used"])
149178
metadata["service_tiers"] = sorted(metadata["service_tiers"])

0 commit comments

Comments
 (0)