refactor: split MessageDict by role for mypy narrowing#126
Conversation
📝 WalkthroughWalkthrough
ChangesMessageDict Discriminated Union Refactor
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant JSONLParser
participant MessageDict
participant JSONExporter
participant MDExporter
participant SearchIndex
participant SessionStats
JSONLParser->>MessageDict: build role-specific message shapes
JSONExporter->>MessageDict: serialize list[MessageDict]
MDExporter->>MessageDict: render user/assistant/system messages
SearchIndex->>MessageDict: extract role-specific searchable text
SessionStats->>MessageDict: inspect user tool results and assistant messages
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/mypy_types/message_dict_valid.py (1)
5-9: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winVerify mypy resolves the literal-to-union assignment without ambiguity.
{"role": "user", "text": "hello"}is assigned directly to alist[MessageDict](5-wayUnion[TypedDict]) annotation. Several of the union members (UserMessageDict,AssistantMessageDict,SystemMessageDict,ResultMessageDict) share an identically-typedtextfield, which is exactly the pattern that historically triggered mypy's "Type of TypedDict is ambiguous" error even with distinctLiteraltags (mypy#8533), later addressed by a separate feature PR (#13274/#14505). Worth confirming the project's pinned mypy version actually resolves this cleanly, since a false ambiguity error here would break the "should pass" assertion intest_message_dict_mypy.py.A more defensive fixture would avoid relying on this inference by annotating the literal with its concrete role type first:
user_only: UserMessageDict = {"role": "user", "text": "hello"} messages: list[MessageDict] = [user_only]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/mypy_types/message_dict_valid.py` around lines 5 - 9, The mypy fixture relies on direct inference of a dict literal into list[MessageDict], which can still trigger ambiguous TypedDict union resolution in some versions. Update the test in message_dict_valid.py to avoid depending on that inference by first binding the literal to the concrete UserMessageDict type, then building the list from that value; keep the rest of the narrowing check with msg, narrowed, and get("slug") unchanged so the test still verifies the intended path.utils/jsonl_parser.py (1)
59-130: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReduce duplication in per-role fallback construction.
Each role branch repeats
base["uuid"],base["parent_uuid"],base["timestamp"],base["is_sidechain"]verbatim. Since every role-specific dict is a subtype ofBaseMessageDict, spreadingbasewould remove the repetition.♻️ Proposed refactor using `**base` spread
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, - } + user_msg: UserMessageDict = {**base, "role": "user", "text": text} return user_msgApply similarly for the
assistant,system,progress, andresultbranches.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@utils/jsonl_parser.py` around lines 59 - 130, The per-role fallback builders in _fallback_message repeat the same BaseMessageDict fields in every branch. Refactor the user, assistant, system, progress, and result dict construction to reuse the common values from _fallback_common_fields via unpacking instead of spelling out uuid, parent_uuid, timestamp, and is_sidechain repeatedly. Keep the role-specific fields unchanged while simplifying the branch bodies.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/mypy_types/message_dict_valid.py`:
- Around line 5-9: The mypy fixture relies on direct inference of a dict literal
into list[MessageDict], which can still trigger ambiguous TypedDict union
resolution in some versions. Update the test in message_dict_valid.py to avoid
depending on that inference by first binding the literal to the concrete
UserMessageDict type, then building the list from that value; keep the rest of
the narrowing check with msg, narrowed, and get("slug") unchanged so the test
still verifies the intended path.
In `@utils/jsonl_parser.py`:
- Around line 59-130: The per-role fallback builders in _fallback_message repeat
the same BaseMessageDict fields in every branch. Refactor the user, assistant,
system, progress, and result dict construction to reuse the common values from
_fallback_common_fields via unpacking instead of spelling out uuid, parent_uuid,
timestamp, and is_sidechain repeatedly. Keep the role-specific fields unchanged
while simplifying the branch bodies.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b03c7a3f-666e-4d57-908f-d6f07c1b6b2e
📒 Files selected for processing (10)
models/__init__.pymodels/session.pytests/mypy_types/message_dict_invalid.pytests/mypy_types/message_dict_valid.pytests/test_message_dict_mypy.pyutils/json_exporter.pyutils/jsonl_parser.pyutils/md_exporter.pyutils/search_index.pyutils/session_stats.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/test_message_dict_mypy.py (1)
19-35: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winAdd a timeout to the mypy subprocess
subprocess.run(...)here can hang the test suite indefinitely if mypy stalls. Set atimeout=so the test fails fast instead of blocking CI.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_message_dict_mypy.py` around lines 19 - 35, The _run_mypy_on helper currently calls subprocess.run without any timeout, so the test can hang indefinitely if mypy stalls. Update the _run_mypy_on function to pass a timeout argument to subprocess.run so it fails fast in CI, keeping the existing CompletedProcess behavior and preserving the current mypy invocation details.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tests/test_message_dict_mypy.py`:
- Around line 19-35: The _run_mypy_on helper currently calls subprocess.run
without any timeout, so the test can hang indefinitely if mypy stalls. Update
the _run_mypy_on function to pass a timeout argument to subprocess.run so it
fails fast in CI, keeping the existing CompletedProcess behavior and preserving
the current mypy invocation details.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 17177c63-8e2c-418f-923a-79eec8d4be67
📒 Files selected for processing (1)
tests/test_message_dict_mypy.py
There was a problem hiding this comment.
🧹 Nitpick comments (2)
utils/search_index.py (1)
372-401: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
contentlocal is now always empty — dead parameter.
contentis initialized at Line 376 and never reassigned in any role branch (system/result already inline the content fallback intotextat Lines 388-392). Thecontent=contentargument at Line 398 is therefore always"", making the parameter and variable vestigial.♻️ Suggested cleanup
role = msg["role"] text = "" - content = "" tool_result: object = None progress_data: object = None @@ return combine_searchable_text( text=text, - content=content, tool_result=tool_result, progress_data=progress_data, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@utils/search_index.py` around lines 372 - 401, `message_searchable_text` has a dead `content` local that is never populated, so remove the unused variable and stop passing it to `combine_searchable_text`. Keep the role-specific text extraction logic in `message_searchable_text` intact, and update the `combine_searchable_text` call to rely only on the meaningful fields (`text`, `tool_result`, `progress_data`) so the function reflects the actual data flow.tests/test_message_dict_mypy.py (1)
19-35: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winIsolate mypy's cache directory to avoid repo pollution/races.
Without
--cache-dir, mypy writes.mypy_cacheunderREPO_ROOT(thecwd) on every parametrized invocation. This pollutes the working tree and can race with parallel test execution (e.g. pytest-xdist) or a developer's own mypy cache, causing flaky results.♻️ Suggested fix: use an isolated temp cache dir per invocation
+import tempfile + 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, - ) + with tempfile.TemporaryDirectory() as cache_dir: + return subprocess.run( + [ + sys.executable, + "-m", + "mypy", + "--strict", + str(path), + "--config-file", + str(REPO_ROOT / "pyproject.toml"), + "--cache-dir", + cache_dir, + ], + cwd=REPO_ROOT, + capture_output=True, + text=True, + check=False, + timeout=60.0, + )Note: the static analysis "OS command injection" flag on this block is a false positive —
pathcomes fromMYPY_TYPES_DIR / fixture_name, a fixed internal test path, not external input.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_message_dict_mypy.py` around lines 19 - 35, The `_run_mypy_on` helper is letting mypy use the default `.mypy_cache` under `REPO_ROOT`, which can pollute the repo and cause races in parallel runs. Update the `subprocess.run` invocation in `_run_mypy_on` to pass an isolated temporary cache directory for each call, so every `mypy` execution uses its own cache and does not share state with other tests or developer runs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/test_message_dict_mypy.py`:
- Around line 19-35: The `_run_mypy_on` helper is letting mypy use the default
`.mypy_cache` under `REPO_ROOT`, which can pollute the repo and cause races in
parallel runs. Update the `subprocess.run` invocation in `_run_mypy_on` to pass
an isolated temporary cache directory for each call, so every `mypy` execution
uses its own cache and does not share state with other tests or developer runs.
In `@utils/search_index.py`:
- Around line 372-401: `message_searchable_text` has a dead `content` local that
is never populated, so remove the unused variable and stop passing it to
`combine_searchable_text`. Keep the role-specific text extraction logic in
`message_searchable_text` intact, and update the `combine_searchable_text` call
to rely only on the meaningful fields (`text`, `tool_result`, `progress_data`)
so the function reflects the actual data flow.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e0039686-7f37-4fe9-b6c7-1ccd194b45fa
📒 Files selected for processing (5)
tests/mypy_types/message_dict_invalid.pytests/mypy_types/message_dict_valid.pytests/test_message_dict_mypy.pyutils/jsonl_parser.pyutils/search_index.py
💤 Files with no reviewable changes (1)
- utils/jsonl_parser.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/mypy_types/message_dict_invalid.py
3 participants
Closes #117
Split MessageDict into per-role TypedDicts for strict mypy
MessageDict was one flat TypedDict with optional fields for every role in the same bag. Strict mypy had no way to know that
thinkingonly shows up on assistant messages, somsg["thinking"]on a user message type-checked and failed at runtime.This branch adds
BaseMessageDict, five role-specific TypedDicts, and aMessageDictunion keyed onrole. The JSONL parser builds the right shape per role. Exporters, session stats, and the search index branch onmsg["role"]before touching fields that belong to one role.api/sessions.pyandutils/tool_dispatch.pyare unchanged. They route or pass messages without reading role-specific fields, so there was nothing to narrow there.Regression tests live in
tests/mypy_types/and run throughtests/test_message_dict_mypy.py. The invalid fixture fails mypy forthinkingon both aUserMessageDictand an un-narrowedMessageDict. The valid fixture checks that narrowing afterrole == "user"allows safe reads. Serialized JSON and runtime behavior should match master.Sanity check:
mypy -p models -p utils -p api, thenpytest tests/test_message_dict_mypy.py tests/test_jsonl_parser.py tests/test_real_session_fixtures.py -q, then fullpytestandruff check ..Summary by CodeRabbit