refactor: split MessageDict by role for mypy narrowing (#126)

* refactor: split MessageDict by role for mypy narrowing

* test: require typeddict-item on both invalid MessageDict fixture lines

* test: harden MessageDict mypy fixtures and restore search parity

Author: clean6378-max-it

models/__init__.py

@@ -6,12 +6,17 @@
 from models.record_data import RecordDataUnion
 from models.search import SearchHitDict
 from models.session import (
+    AssistantMessageDict,
     MessageDict,
+    ProgressMessageDict,
     QuickSessionInfoDict,
+    ResultMessageDict,
     RoleLiteral,
     SessionDict,
     SessionMetadataDict,
+    SystemMessageDict,
     ToolUseDict,
+    UserMessageDict,
 )
 from models.stats import FilesTouchedDict, SessionStatsDict
 from models.tool_results import ToolResultUnion
@@ -20,17 +25,22 @@
     "ErrorResponse",
     "ExportStateDict",
     "FilesTouchedDict",
+    "AssistantMessageDict",
     "MessageDict",
+    "ProgressMessageDict",
     "ProjectDict",
     "ProjectSessionRowDict",
     "QuickSessionInfoDict",
+    "ResultMessageDict",
     "RoleLiteral",
     "SearchHitDict",
     "SessionDict",
     "SessionListItemDict",
     "SessionMetadataDict",
     "SessionStatsDict",
+    "SystemMessageDict",
     "RecordDataUnion",
     "ToolResultUnion",
     "ToolUseDict",
+    "UserMessageDict",
 ]

models/session.py

@@ -1,6 +1,6 @@
 """Parsed session shapes from jsonl_parser."""
 
-from typing import Any, Literal, NotRequired, TypedDict
+from typing import Any, Literal, NotRequired, TypedDict, Union
 
 from models.record_data import RecordDataUnion
 from models.tool_results import ToolNameLiteral, ToolResultUnion
@@ -26,32 +26,66 @@ class MessageUsageDict(TypedDict, total=False):
 SystemSubtypeLiteral = Literal["compact_boundary", "init"]
 
 
-class MessageDict(TypedDict):
-    role: RoleLiteral
-    uuid: NotRequired[str | None]
-    parent_uuid: NotRequired[str | None]
-    timestamp: NotRequired[str | None]
+class BaseMessageDict(TypedDict, total=False):
+    """Fields shared across every parsed message role."""
+
+    uuid: str | None
+    parent_uuid: str | None
+    timestamp: str | None
+    is_sidechain: bool
+
+
+class UserMessageDict(BaseMessageDict):
+    role: Literal["user"]
     text: NotRequired[str]
-    content: NotRequired[str]
     images: NotRequired[list[Any] | None]
-    is_sidechain: NotRequired[bool]
     tool_result: NotRequired[ToolResultUnion | None]
     tool_result_parsed: NotRequired[dict[str, object] | None]
     slug: NotRequired[str | None]
+
+
+class AssistantMessageDict(BaseMessageDict):
+    role: Literal["assistant"]
+    text: NotRequired[str]
     model: NotRequired[str]
     stop_reason: NotRequired[str]
     thinking: NotRequired[str | None]
     tool_uses: NotRequired[list[ToolUseDict] | None]
     is_api_error: NotRequired[bool]
     usage: NotRequired[MessageUsageDict]
+
+
+class SystemMessageDict(BaseMessageDict):
+    role: Literal["system"]
+    text: NotRequired[str]
     subtype: NotRequired[str]
+    content: NotRequired[str]
     level: NotRequired[str]
-    data: NotRequired[RecordDataUnion]
+
+
+class ResultMessageDict(BaseMessageDict):
+    role: Literal["result"]
+    text: NotRequired[str]
+    content: NotRequired[str]
+
+
+class ProgressMessageDict(BaseMessageDict):
+    role: Literal["progress"]
     progress_type: NotRequired[str]
+    data: NotRequired[RecordDataUnion]
     tool_use_id: NotRequired[str | None]
     parent_tool_use_id: NotRequired[str | None]
 
 
+MessageDict = Union[
+    UserMessageDict,
+    AssistantMessageDict,
+    SystemMessageDict,
+    ResultMessageDict,
+    ProgressMessageDict,
+]
+
+
 class SessionMetadataDict(TypedDict):
     """Metadata accumulated while parsing a Claude Code JSONL session.
 

tests/mypy_types/message_dict_invalid.py

@@ -0,0 +1,9 @@
+"""Negative fixture: role-inappropriate MessageDict access must fail mypy strict."""
+
+from models.session import MessageDict, UserMessageDict
+
+user_msg: UserMessageDict = {"role": "user", "text": "hello"}
+_ = user_msg["thinking"]  # expect: typeddict-item thinking
+
+union_msg: MessageDict = {"role": "user", "text": "hello"}
+_ = union_msg["thinking"]  # expect: typeddict-item thinking

tests/mypy_types/message_dict_valid.py

@@ -0,0 +1,15 @@
+"""Positive fixture: discriminant narrowing allows role-specific fields."""
+
+from models.session import AssistantMessageDict, MessageDict, UserMessageDict
+
+messages: list[MessageDict] = [
+    {"role": "user", "text": "hello"},
+    {"role": "assistant", "text": "hi", "thinking": "hmm"},
+]
+for msg in messages:
+    if msg["role"] == "user":
+        narrowed: UserMessageDict = msg
+        _ = narrowed["slug"]
+    elif msg["role"] == "assistant":
+        narrowed_asst: AssistantMessageDict = msg
+        _ = narrowed_asst["thinking"]

tests/test_message_dict_mypy.py

@@ -0,0 +1,78 @@
+"""Type-level regression: MessageDict union rejects role-inappropriate access."""
+
+from __future__ import annotations
+
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+MYPY_TYPES_DIR = REPO_ROOT / "tests" / "mypy_types"
+
+_INVALID_FIXTURE = "message_dict_invalid.py"
+_EXPECT_MARKER = "expect: typeddict-item"
+
+
+def _run_mypy_on(path: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "mypy",
+            "--strict",
+            str(path),
+            "--config-file",
+            str(REPO_ROOT / "pyproject.toml"),
+        ],
+        cwd=REPO_ROOT,
+        capture_output=True,
+        text=True,
+        check=False,
+        timeout=60.0,
+    )
+
+
+def _typeddict_item_error_lines(output: str, fixture_name: str) -> set[int]:
+    pattern = rf"{re.escape(fixture_name)}:(\d+): error:.*\[typeddict-item\]"
+    return {int(match.group(1)) for match in re.finditer(pattern, output)}
+
+
+def _invalid_fixture_expectations(fixture_path: Path) -> tuple[set[int], list[str]]:
+    expected_lines: set[int] = set()
+    expected_tokens: list[str] = []
+    for lineno, line in enumerate(fixture_path.read_text(encoding="utf-8").splitlines(), start=1):
+        if _EXPECT_MARKER not in line:
+            continue
+        expected_lines.add(lineno)
+        suffix = line.split(_EXPECT_MARKER, 1)[1].strip()
+        if suffix:
+            expected_tokens.append(suffix)
+    return expected_lines, expected_tokens
+
+
+@pytest.mark.parametrize(
+    ("fixture_name", "should_pass"),
+    [
+        (_INVALID_FIXTURE, False),
+        ("message_dict_valid.py", True),
+    ],
+)
+def test_message_dict_mypy_fixtures(fixture_name: str, should_pass: bool) -> None:
+    fixture_path = MYPY_TYPES_DIR / fixture_name
+    result = _run_mypy_on(fixture_path)
+    output = result.stdout + result.stderr
+    if should_pass:
+        assert result.returncode == 0, output
+    else:
+        assert result.returncode != 0
+        error_lines = _typeddict_item_error_lines(output, fixture_name)
+        expected_lines, expected_tokens = _invalid_fixture_expectations(fixture_path)
+        assert error_lines >= expected_lines, (
+            f"expected typeddict-item on lines {sorted(expected_lines)}, "
+            f"got {sorted(error_lines)}: {output}"
+        )
+        for token in expected_tokens:
+            assert token in output, f"expected mypy output to mention {token!r}: {output}"

utils/json_exporter.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timezone
 from typing import Any
 
-from models.session import SessionDict, SessionMetadataDict
+from models.session import MessageDict, SessionDict, SessionMetadataDict
 from models.stats import SessionStatsDict
 
 
@@ -39,11 +39,11 @@ def _serialize_metadata(meta: SessionMetadataDict) -> dict[str, Any]:
     return result
 
 
-def _serialize_messages(messages: list[Any]) -> list[dict[str, Any]]:
+def _serialize_messages(messages: list[MessageDict]) -> list[dict[str, Any]]:
     """Same set-to-list cleanup, but for each message dict."""
-    out = []
+    out: list[dict[str, Any]] = []
     for msg in messages:
-        clean = {}
+        clean: dict[str, Any] = {}
         for key, val in msg.items():
             if isinstance(val, set):
                 clean[key] = sorted(val)