diff --git a/models/__init__.py b/models/__init__.py index ee0eec5..7c3b742 100644 --- a/models/__init__.py +++ b/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", ] diff --git a/models/session.py b/models/session.py index cf0e608..4c878ca 100644 --- a/models/session.py +++ b/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. diff --git a/tests/mypy_types/message_dict_invalid.py b/tests/mypy_types/message_dict_invalid.py new file mode 100644 index 0000000..78b1e62 --- /dev/null +++ b/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 diff --git a/tests/mypy_types/message_dict_valid.py b/tests/mypy_types/message_dict_valid.py new file mode 100644 index 0000000..264c7d5 --- /dev/null +++ b/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"] diff --git a/tests/test_message_dict_mypy.py b/tests/test_message_dict_mypy.py new file mode 100644 index 0000000..205c6ed --- /dev/null +++ b/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}" diff --git a/utils/json_exporter.py b/utils/json_exporter.py index 3439732..2fa1688 100644 --- a/utils/json_exporter.py +++ b/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) diff --git a/utils/jsonl_parser.py b/utils/jsonl_parser.py index 337ba27..a1dd7a9 100644 --- a/utils/jsonl_parser.py +++ b/utils/jsonl_parser.py @@ -10,12 +10,18 @@ from models.record_data import RecordDataUnion from models.session import ( + AssistantMessageDict, + BaseMessageDict, MessageDict, + ProgressMessageDict, + ResultMessageDict, RoleLiteral, SessionDict, SessionMetadataBuilderDict, SessionMetadataDict, + SystemMessageDict, ToolUseDict, + UserMessageDict, ) from models.tool_results import ToolResultUnion, is_tool_result_dict from utils.jsonl_helpers import ( @@ -50,6 +56,16 @@ def _coerce_role(raw: str) -> RoleLiteral: return "system" +def _fallback_common_fields(entry: dict[str, Any]) -> BaseMessageDict: + """uuid/parent/timestamp/is_sidechain shared by every fallback message shape.""" + return { + "uuid": entry.get("uuid"), + "parent_uuid": entry.get("parentUuid"), + "timestamp": entry.get("timestamp"), + "is_sidechain": entry.get("isSidechain", False), + } + + def _fallback_message(entry: dict[str, Any], role: RoleLiteral) -> MessageDict: """Minimal message for JSONL entry types without a dedicated processor.""" raw_content = entry.get("content", "") @@ -59,15 +75,57 @@ def _fallback_message(entry: dict[str, Any], role: RoleLiteral) -> MessageDict: text = _extract_text(_normalize_content(raw_content)) else: text = "" - return { - "role": role, - "uuid": entry.get("uuid"), - "parent_uuid": entry.get("parentUuid"), - "timestamp": entry.get("timestamp"), + base = _fallback_common_fields(entry) + if role == "user": + user_msg: UserMessageDict = { + "role": "user", + "uuid": base["uuid"], + "parent_uuid": base["parent_uuid"], + "timestamp": base["timestamp"], + "is_sidechain": base["is_sidechain"], + "text": text, + } + return user_msg + if role == "assistant": + assistant_msg: AssistantMessageDict = { + "role": "assistant", + "uuid": base["uuid"], + "parent_uuid": base["parent_uuid"], + "timestamp": base["timestamp"], + "is_sidechain": base["is_sidechain"], + "text": text, + } + return assistant_msg + if role == "system": + system_msg: SystemMessageDict = { + "role": "system", + "uuid": base["uuid"], + "parent_uuid": base["parent_uuid"], + "timestamp": base["timestamp"], + "is_sidechain": base["is_sidechain"], + "text": text, + "content": text, + } + return system_msg + if role == "progress": + progress_msg: ProgressMessageDict = { + "role": "progress", + "uuid": base["uuid"], + "parent_uuid": base["parent_uuid"], + "timestamp": base["timestamp"], + "is_sidechain": base["is_sidechain"], + } + return progress_msg + result_msg: ResultMessageDict = { + "role": "result", + "uuid": base["uuid"], + "parent_uuid": base["parent_uuid"], + "timestamp": base["timestamp"], + "is_sidechain": base["is_sidechain"], "text": text, "content": text, - "is_sidechain": entry.get("isSidechain", False), } + return result_msg def _safe_int(val: Any) -> int: @@ -289,20 +347,19 @@ def _process_user( if tr_images: images = (images or []) + tr_images - messages.append( - { - "role": "user", - "uuid": entry.get("uuid"), - "parent_uuid": entry.get("parentUuid"), - "timestamp": entry.get("timestamp"), - "text": text, - "images": images if images else None, - "is_sidechain": entry.get("isSidechain", False), - "tool_result": tool_result, - "tool_result_parsed": tool_result_parsed, - "slug": entry.get("slug"), - } - ) + user_msg: UserMessageDict = { + "role": "user", + "uuid": entry.get("uuid"), + "parent_uuid": entry.get("parentUuid"), + "timestamp": entry.get("timestamp"), + "text": text, + "images": images if images else None, + "is_sidechain": entry.get("isSidechain", False), + "tool_result": tool_result, + "tool_result_parsed": tool_result_parsed, + "slug": entry.get("slug"), + } + messages.append(user_msg) def _process_assistant( @@ -377,28 +434,27 @@ def _process_assistant( tool_uses.append(tool_use) track_tool_file_activity(tool_name, safe_input, metadata) - messages.append( - { - "role": "assistant", - "uuid": entry.get("uuid"), - "parent_uuid": entry.get("parentUuid"), - "timestamp": entry.get("timestamp"), - "model": model, - "stop_reason": stop_reason, - "text": "\n".join(text_parts), - "thinking": "\n\n".join(thinking_parts) if thinking_parts else None, - "tool_uses": tool_uses if tool_uses else None, - "is_sidechain": entry.get("isSidechain", False), - "is_api_error": entry.get("isApiErrorMessage", False), - "usage": { - "input_tokens": _safe_int(usage.get("input_tokens")), - "output_tokens": _safe_int(usage.get("output_tokens")), - "cache_read": _safe_int(usage.get("cache_read_input_tokens")), - "cache_creation": _safe_int(usage.get("cache_creation_input_tokens")), - "service_tier": tier if isinstance(tier, str) else None, - }, - } - ) + assistant_msg: AssistantMessageDict = { + "role": "assistant", + "uuid": entry.get("uuid"), + "parent_uuid": entry.get("parentUuid"), + "timestamp": entry.get("timestamp"), + "model": model, + "stop_reason": stop_reason, + "text": "\n".join(text_parts), + "thinking": "\n\n".join(thinking_parts) if thinking_parts else None, + "tool_uses": tool_uses if tool_uses else None, + "is_sidechain": entry.get("isSidechain", False), + "is_api_error": entry.get("isApiErrorMessage", False), + "usage": { + "input_tokens": _safe_int(usage.get("input_tokens")), + "output_tokens": _safe_int(usage.get("output_tokens")), + "cache_read": _safe_int(usage.get("cache_read_input_tokens")), + "cache_creation": _safe_int(usage.get("cache_creation_input_tokens")), + "service_tier": tier if isinstance(tier, str) else None, + }, + } + messages.append(assistant_msg) def _process_system( @@ -419,17 +475,16 @@ def _process_system( } ) - messages.append( - { - "role": "system", - "uuid": entry.get("uuid"), - "parent_uuid": entry.get("parentUuid"), - "timestamp": entry.get("timestamp"), - "subtype": subtype, - "content": entry.get("content", ""), - "is_sidechain": entry.get("isSidechain", False), - } - ) + system_msg: SystemMessageDict = { + "role": "system", + "uuid": entry.get("uuid"), + "parent_uuid": entry.get("parentUuid"), + "timestamp": entry.get("timestamp"), + "subtype": subtype, + "content": entry.get("content", ""), + "is_sidechain": entry.get("isSidechain", False), + } + messages.append(system_msg) def _process_progress(entry: dict[str, Any], messages: list[MessageDict]) -> None: @@ -439,16 +494,15 @@ def _process_progress(entry: dict[str, Any], messages: list[MessageDict]) -> Non data: RecordDataUnion = raw_data if isinstance(raw_data, dict) else {} progress_type = str(data.get("type", "")) - messages.append( - { - "role": "progress", - "uuid": entry.get("uuid"), - "parent_uuid": entry.get("parentUuid"), - "timestamp": entry.get("timestamp"), - "progress_type": progress_type, - "data": data, - "tool_use_id": entry.get("toolUseID"), - "parent_tool_use_id": entry.get("parentToolUseID"), - "is_sidechain": entry.get("isSidechain", False), - } - ) + progress_msg: ProgressMessageDict = { + "role": "progress", + "uuid": entry.get("uuid"), + "parent_uuid": entry.get("parentUuid"), + "timestamp": entry.get("timestamp"), + "progress_type": progress_type, + "data": data, + "tool_use_id": entry.get("toolUseID"), + "parent_tool_use_id": entry.get("parentToolUseID"), + "is_sidechain": entry.get("isSidechain", False), + } + messages.append(progress_msg) diff --git a/utils/md_exporter.py b/utils/md_exporter.py index ab76d71..fb9c488 100644 --- a/utils/md_exporter.py +++ b/utils/md_exporter.py @@ -5,7 +5,14 @@ from datetime import datetime from typing import Any -from models.session import MessageDict, SessionDict, ToolUseDict +from models.session import ( + AssistantMessageDict, + MessageDict, + SessionDict, + SystemMessageDict, + ToolUseDict, + UserMessageDict, +) from models.stats import SessionStatsDict from utils.jsonl_helpers import strip_system_tags from utils.session_stats import format_duration @@ -189,18 +196,17 @@ def _build_summary(session: SessionDict, stats: SessionStatsDict) -> str: def _build_body(messages: list[MessageDict]) -> str: parts = [] for msg in messages: - role = msg["role"] - if role == "user": + if msg["role"] == "user": parts.append(_render_user(msg)) - elif role == "assistant": + elif msg["role"] == "assistant": parts.append(_render_assistant(msg)) - elif role == "system": + elif msg["role"] == "system": parts.append(_render_system(msg)) # Skip progress messages in MD (too noisy) return "\n".join(parts) -def _render_user(msg: MessageDict) -> str: +def _render_user(msg: UserMessageDict) -> str: lines = [] lines.append("### User\n") @@ -233,7 +239,7 @@ def _render_user(msg: MessageDict) -> str: return "\n".join(lines) -def _render_assistant(msg: MessageDict) -> str: +def _render_assistant(msg: AssistantMessageDict) -> str: lines = [] lines.append("### Assistant\n") @@ -440,7 +446,7 @@ def _render_tool_result(parsed: dict[str, Any]) -> str: return "\n".join(lines) -def _render_system(msg: MessageDict) -> str: +def _render_system(msg: SystemMessageDict) -> str: lines = [] subtype = msg.get("subtype", "") content = msg.get("content", "") diff --git a/utils/search_index.py b/utils/search_index.py index 5e74f7e..a09beae 100644 --- a/utils/search_index.py +++ b/utils/search_index.py @@ -371,13 +371,33 @@ def combine_searchable_text( def message_searchable_text(msg: MessageDict) -> str: """Searchable text for one parsed session message (live-scan path).""" - text = msg.get("text", "") or msg.get("content", "") - if not isinstance(text, str): - text = "" + role = msg["role"] + text = "" + content = "" + tool_result: object = None + progress_data: object = None + + if role == "user": + raw_text = msg.get("text", "") + text = raw_text if isinstance(raw_text, str) else "" + tool_result = msg.get("tool_result") + elif role == "assistant": + raw_text = msg.get("text", "") + text = raw_text if isinstance(raw_text, str) else "" + elif role == "system": + raw_text = msg.get("text", "") or msg.get("content", "") or "" + text = raw_text if isinstance(raw_text, str) else "" + elif role == "result": + raw_text = msg.get("text", "") or msg.get("content", "") or "" + text = raw_text if isinstance(raw_text, str) else "" + elif role == "progress": + progress_data = msg.get("data") + return combine_searchable_text( text=text, - tool_result=msg.get("tool_result"), - progress_data=msg.get("data") if msg.get("role") == "progress" else None, + content=content, + tool_result=tool_result, + progress_data=progress_data, ) diff --git a/utils/session_stats.py b/utils/session_stats.py index 965d753..682d333 100644 --- a/utils/session_stats.py +++ b/utils/session_stats.py @@ -60,11 +60,10 @@ def _compute_commands_run(messages: list[MessageDict]) -> list[dict[str, Any]]: """Walk through messages and match up Bash tool_use calls with their subsequent tool_result entries to get exit codes and error status.""" commands = [] - # Build a map of tool_use_id -> command from assistant messages - pending_commands = {} + pending_commands: dict[str, dict[str, Any]] = {} for msg in messages: - tool_uses = msg.get("tool_uses") or [] - if msg["role"] == "assistant" and tool_uses: + if msg["role"] == "assistant": + tool_uses = msg.get("tool_uses") or [] for tu in tool_uses: if tu["name"] == "Bash": cmd = tu["input"].get("command", "") @@ -73,11 +72,13 @@ def _compute_commands_run(messages: list[MessageDict]) -> list[dict[str, Any]]: "command": cmd, "timestamp": msg.get("timestamp"), } + continue + + if msg["role"] != "user": + continue - # Match tool results back to commands trp = msg.get("tool_result_parsed") - if msg["role"] == "user" and trp and trp.get("result_type") == "bash": - # Try to find matching command by sequential order + if trp and trp.get("result_type") == "bash": if pending_commands: first_id = next(iter(pending_commands)) entry = pending_commands.pop(first_id)