Skip to content

Commit f294fb9

Browse files
refactor: priority-based tool result dispatch selection (#127)
* refactor: priority-based tool result dispatch selection * fix: harden dispatch id validation and priority docs * fix: task_message dispatch uses order, not priority boost
1 parent 964813f commit f294fb9

7 files changed

Lines changed: 230 additions & 128 deletions

CONTRIBUTING.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ See [`docs/architecture.md`](docs/architecture.md) for data flow, export state m
145145

146146
Claude Code assistant `tool_use` blocks carry a `name` string (e.g. `"Read"`, `"Bash"`). The browser coordinates that name across four sites; drift is caught by `tests/test_tool_dispatch_sync.py`.
147147

148-
1. **`utils/tool_dispatch.py`** — add the name to `_FILE_ACTIVITY_HANDLERS` (`None` if no file/bash/web side effects); `KNOWN_TOOL_TYPES` is derived from its keys. If the tool has a distinct `toolUseResult` JSON shape, add `(predicate, builder)` to `_TOOL_RESULT_DISPATCH` (respect ordering — see module docstring and `tests/test_tool_dispatch_ordering.py`).
148+
1. **`utils/tool_dispatch.py`** — add the name to `_FILE_ACTIVITY_HANDLERS` (`None` if no file/bash/web side effects); `KNOWN_TOOL_TYPES` is derived from its keys. If the tool has a distinct `toolUseResult` JSON shape, add a `ToolResultDispatchEntry` to `_TOOL_RESULT_DISPATCH` (set `priority` when overlapping another predicate — see module docstring and `tests/test_tool_dispatch_ordering.py`).
149149
2. **`models/tool_results.py`** — add the name to `ToolNameLiteral` and, when the tool has a distinct result payload, add the TypedDict, type guard (`is_*_tool_result`), and union member on `ToolResultUnion`.
150150
3. **`utils/md_exporter.py`** — add an `elif name == "…"` branch in `_render_tool_use` (sync test parses these branches).
151151
4. **`static/js/render/registry.js`** — add a `TOOL_USE_RENDERERS` entry (and a `tool_use/*.js` renderer module).

docs/architecture.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,13 +71,13 @@
7171

7272
## Dispatch table
7373

74-
In `utils/tool_dispatch.py`, tool results are classified through `_parse_tool_result`, a **predicate-ordered dispatch table** (not a simple `if tool_name == ...` chain). **Order is load-bearing**: the first matching predicate wins. Tests in `tests/test_jsonl_parser.py` and `tests/test_real_session_fixtures.py` guard ordering regressions.
74+
In `utils/tool_dispatch.py`, tool results are classified through `_parse_tool_result`, a **priority-based dispatch table**. Every matching predicate is considered; the entry with the highest `priority` wins (ties favor earlier registration). Overlap exceptions are explicit priorities, not tuple position — see `tests/test_tool_dispatch_ordering.py` and `tests/test_tool_dispatch_adversarial.py`.
7575

7676
When adding a new tool renderer:
7777

78-
1. Add a `(predicate, builder)` pair to `_TOOL_RESULT_DISPATCH` in `utils/tool_dispatch.py`, preserving existing predicate order unless you also update fixtures and ordering tests (`tests/test_jsonl_parser.py`, `tests/test_real_session_fixtures.py`). Order is **not** “specific before generic” in general — the first match wins. `is_task_message_tool_result` is the intentional broad-before-narrow exception (`task_id` or `message` before retrieval/completed/async).
78+
1. Add a `ToolResultDispatchEntry` to `_TOOL_RESULT_DISPATCH` in `utils/tool_dispatch.py`. Set `priority` higher than any overlapping predicate it must beat without relying on registration order, and add a row to `ORDERING_INVARIANTS` in `tests/test_tool_dispatch_ordering.py` when overlaps exist. Documented exception: `plan` (priority 1) over `file_write` (0). Broad predicates like `task_message` use registration order instead of elevated priority.
7979
2. Add or extend a JSONL fixture under `tests/fixtures/` (especially for overlaps with existing predicates).
80-
3. Run `pytest tests/test_jsonl_parser.py tests/test_real_session_fixtures.py -v`.
80+
3. Run `pytest tests/test_tool_dispatch_ordering.py tests/test_tool_dispatch_adversarial.py tests/test_jsonl_parser.py tests/test_real_session_fixtures.py -v`.
8181

8282
## Export state machine
8383

@@ -105,7 +105,7 @@ The UI is a **hash-routed** SPA with ES modules under `static/js/`:
105105

106106
- `app.js` — routing and boot
107107
- `projects.js`, `sessions.js`, `search.js`, `export.js` — route handlers
108-
- `render/registry.js`**tool dispatch registry** for session UI: `TOOL_USE_RENDERERS` and `TOOL_RESULT_RENDERERS` map tool name / `result_type` → render function (one module per type under `render/tool_use/` and `render/tool_result/`). Parallels backend `utils/tool_dispatch.py` (backend uses ordered predicates; frontend uses direct key lookup + fallback).
108+
- `render/registry.js`**tool dispatch registry** for session UI: `TOOL_USE_RENDERERS` and `TOOL_RESULT_RENDERERS` map tool name / `result_type` → render function (one module per type under `render/tool_use/` and `render/tool_result/`). Parallels backend `utils/tool_dispatch.py` (backend uses priority-based predicate matching; frontend uses direct key lookup + fallback).
109109
- `static/tool_types.json` — generated manifest of backend tool-use names (`python scripts/gen_tool_types_manifest.py` from `KNOWN_TOOL_TYPES`). Fetched non-blocking at boot by `render/tool_types_manifest.js` (`void initToolTypesManifest()` in `app.js`), which cross-checks `TOOL_USE_RENDERERS` and logs `console.warn` on drift.
110110
- `shared/markdown.js` — markdown + **DOMPurify** sanitization (do not render raw LLM HTML)
111111
- `shared/state.js`, `shared/utils.js`, `shared/theme.js` — shared UI state and helpers

tests/test_jsonl_parser.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,7 +241,7 @@ def test_plan_result(self):
241241
assert r["result_type"] == "plan"
242242

243243
def test_plan_with_content_not_classified_as_file_write(self):
244-
"""plan is registered before file_write in _TOOL_RESULT_DISPATCH."""
244+
"""plan outranks file_write on overlapping toolUseResult blobs."""
245245
r = _parse_tool_result(
246246
{
247247
"plan": [],

tests/test_real_session_fixtures.py

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import pytest
1414

1515
from utils.jsonl_parser import parse_session
16-
from utils.tool_dispatch import _TOOL_RESULT_DISPATCH, _parse_tool_result
16+
from utils.tool_dispatch import _TOOL_RESULT_DISPATCH, _parse_tool_result, _winning_dispatch_entry
1717

1818
FIXTURES_DIR = os.path.join(os.path.dirname(__file__), "fixtures")
1919

@@ -89,7 +89,7 @@ def test_real_session_minimal_has_bash_tool_result() -> None:
8989

9090

9191
def test_real_session_all_tool_types_covers_dispatch_predicates() -> None:
92-
hit: set[int] = set()
92+
hit: set[str] = set()
9393
path = _fixture_path("real_session_all_tool_types.jsonl")
9494
with open(path, encoding="utf-8") as f:
9595
for line in f:
@@ -100,14 +100,10 @@ def test_real_session_all_tool_types_covers_dispatch_predicates() -> None:
100100
tr = entry.get("toolUseResult")
101101
if not isinstance(tr, dict):
102102
continue
103-
matched = False
104-
for i, (pred, _) in enumerate(_TOOL_RESULT_DISPATCH):
105-
if pred(tr):
106-
hit.add(i)
107-
matched = True
108-
break
109-
assert matched, f"toolUseResult matched no predicate: {list(tr.keys())}"
110-
assert hit == set(range(len(_TOOL_RESULT_DISPATCH)))
103+
winner = _winning_dispatch_entry(tr)
104+
assert winner is not None, f"toolUseResult matched no predicate: {list(tr.keys())}"
105+
hit.add(winner.id)
106+
assert hit == {entry.id for entry in _TOOL_RESULT_DISPATCH}
111107

112108

113109
def test_real_session_nested_tools_has_sidechain_and_tool_use() -> None:
@@ -146,12 +142,13 @@ def test_task_retrieval_not_misclassified_as_task_message() -> None:
146142

147143

148144
def test_task_completed_with_message_key_matches_task_message_first() -> None:
149-
"""Legacy dispatch: broad task_message runs before task_completed when ``message`` present.
145+
"""task_message outranks task_completed when ``message`` key is present.
150146
151147
``is_task_message_tool_result`` matches any dict with a ``message`` or ``task_id``
152-
key. Future tool shapes that add ``message`` for status text (e.g. web-fetch) would
153-
be misclassified as task until dispatch order is refined — this test locks that
154-
known false-positive surface.
148+
key and is registered before the narrower task predicates, so registration order
149+
picks it on overlap. Future tool shapes that add ``message`` for status text
150+
(e.g. web-fetch) would be misclassified as task unless dispatch is refined —
151+
this test locks that known false-positive surface.
155152
"""
156153
tr = {
157154
"agentId": "agent-sanitized",

tests/test_tool_dispatch_adversarial.py

Lines changed: 44 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,55 +1,42 @@
11
"""Behavioral adversarial fixtures for ``_TOOL_RESULT_DISPATCH`` predicate overlap.
22
3-
Structural tuple-position guards live in ``test_tool_dispatch_ordering.py``.
4-
These tests construct ``toolUseResult`` JSON that satisfies multiple predicates
5-
and assert the classified winner via ``_parse_tool_result`` (first match wins).
3+
Overlap blobs and invariant tables live in ``test_tool_dispatch_ordering.py``.
4+
These tests assert classified output via ``_parse_tool_result``.
65
"""
76

87
from __future__ import annotations
98

109
from collections.abc import Callable
10+
from dataclasses import replace
1111

1212
import pytest
1313

14-
from tests.test_tool_dispatch_ordering import ORDERING_INVARIANT_IDS, ORDERING_INVARIANTS
14+
from models.tool_results import is_file_write_tool_result, is_task_message_tool_result
15+
from tests.test_tool_dispatch_ordering import (
16+
ORDERING_INVARIANT_IDS,
17+
ORDERING_INVARIANTS,
18+
OVERLAP_BLOBS,
19+
)
1520
from utils import tool_dispatch
16-
from utils.tool_dispatch import _TOOL_RESULT_DISPATCH, _parse_tool_result
21+
from utils.tool_dispatch import _TOOL_RESULT_DISPATCH, _parse_tool_result, _winning_dispatch_entry
1722

18-
# Overlap blobs: keys chosen so both predicates in each ORDERING_INVARIANTS pair match.
19-
PLAN_FILE_WRITE_OVERLAP: dict[str, object] = {
20-
"plan": {"name": "sprint-plan", "steps": ["index", "search", "render"]},
21-
"filePath": ".cursor/plans/week28.md",
22-
"content": "Plan body that would also satisfy file_write.",
23-
}
24-
25-
TASK_MESSAGE_RETRIEVAL_OVERLAP: dict[str, object] = {
26-
"task_id": "task-overlap-retrieval",
27-
"message": "polling retrieval",
28-
"retrieval_status": "found",
29-
"task": {"task_id": "task-overlap-retrieval", "description": "subagent scan"},
30-
}
31-
32-
TASK_MESSAGE_COMPLETED_OVERLAP: dict[str, object] = {
33-
"agentId": "agent-overlap-completed",
34-
"totalDurationMs": 3200,
35-
"totalTokens": 900,
36-
"status": "completed",
37-
"message": "task finished",
38-
}
39-
40-
TASK_MESSAGE_ASYNC_OVERLAP: dict[str, object] = {
41-
"agentId": "agent-overlap-async",
42-
"isAsync": True,
43-
"description": "explore auth handlers",
44-
"message": "task launched",
45-
}
23+
PLAN_FILE_WRITE_OVERLAP = OVERLAP_BLOBS["plan_before_file_write"]
24+
TASK_MESSAGE_RETRIEVAL_OVERLAP = OVERLAP_BLOBS["task_message_before_task_retrieval"]
25+
TASK_MESSAGE_COMPLETED_OVERLAP = OVERLAP_BLOBS["task_message_before_task_completed"]
26+
TASK_MESSAGE_ASYNC_OVERLAP = OVERLAP_BLOBS["task_message_before_task_async"]
4627

4728
# Narrow retrieval shape: only the downstream predicate matches (no task_id/message).
4829
TASK_RETRIEVAL_NARROW: dict[str, object] = {
4930
"retrieval_status": "pending",
5031
"task": {"task_id": "task-narrow", "description": "wait for result"},
5132
}
5233

34+
FILE_WRITE_TASK_MESSAGE_OVERLAP: dict[str, object] = {
35+
"message": "status line on a write result",
36+
"filePath": "/tmp/example.txt",
37+
"content": "file body",
38+
}
39+
5340
AssertWinner = Callable[[dict[str, object]], None]
5441

5542

@@ -118,19 +105,32 @@ def test_task_retrieval_narrow_shape_without_task_message_keys() -> None:
118105
assert result.get("task_id") == "task-narrow"
119106

120107

121-
def test_inverted_plan_file_write_dispatch_misclassifies_overlap(
108+
def test_file_write_beats_task_message_when_both_match() -> None:
109+
"""Regression: broad task_message must not outrank earlier shapes via priority."""
110+
blob = FILE_WRITE_TASK_MESSAGE_OVERLAP
111+
assert is_file_write_tool_result(blob)
112+
assert is_task_message_tool_result(blob)
113+
winner = _winning_dispatch_entry(blob)
114+
assert winner is not None
115+
assert winner.id == "file_write"
116+
result = _parse_tool_result(blob)
117+
assert result is not None
118+
assert result["result_type"] == "file_write"
119+
120+
121+
def test_inverted_plan_file_write_priority_misclassifies_overlap(
122122
monkeypatch: pytest.MonkeyPatch,
123123
) -> None:
124-
"""Regression: swapping plan below file_write flips the overlap winner."""
125-
table = list(_TOOL_RESULT_DISPATCH)
126-
plan_idx = next(
127-
i for i, (pred, _) in enumerate(table) if pred.__name__ == "is_plan_tool_result"
128-
)
129-
write_idx = next(
130-
i for i, (pred, _) in enumerate(table) if pred.__name__ == "is_file_write_tool_result"
124+
"""Regression: giving file_write higher priority than plan flips the overlap winner."""
125+
table = tuple(
126+
replace(entry, priority=1)
127+
if entry.id == "file_write"
128+
else replace(entry, priority=0)
129+
if entry.id == "plan"
130+
else entry
131+
for entry in _TOOL_RESULT_DISPATCH
131132
)
132-
table[plan_idx], table[write_idx] = table[write_idx], table[plan_idx]
133-
monkeypatch.setattr(tool_dispatch, "_TOOL_RESULT_DISPATCH", tuple(table))
133+
monkeypatch.setattr(tool_dispatch, "_TOOL_RESULT_DISPATCH", table)
134134

135135
result = _parse_tool_result(PLAN_FILE_WRITE_OVERLAP)
136136
assert result is not None
@@ -145,3 +145,4 @@ def test_ordering_invariants_have_adversarial_coverage() -> None:
145145
assert len(ORDERING_INVARIANTS) == len(ORDERING_INVARIANT_IDS)
146146
assert len(ORDERING_INVARIANT_IDS) == len(_INVARIANT_BEHAVIOR)
147147
assert set(_INVARIANT_BEHAVIOR.keys()) == set(ORDERING_INVARIANT_IDS)
148+
assert set(OVERLAP_BLOBS.keys()) == set(ORDERING_INVARIANT_IDS)
Lines changed: 62 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
"""Structural ordering invariants for ``_TOOL_RESULT_DISPATCH``.
1+
"""Structural overlap invariants for ``_TOOL_RESULT_DISPATCH``.
22
3-
First matching predicate wins; misordering silently misclassifies tool results.
4-
Invariants are declared as ``(before, after, reason)`` triples — add a row to
5-
``ORDERING_INVARIANTS`` when inserting a predicate that must sit above another.
3+
When multiple predicates match, the highest ``priority`` wins; equal priority
4+
favors earlier registration. Invariants are declared as ``(before, after, reason)``
5+
triples with a shared overlap blob — add a row to ``ORDERING_INVARIANTS`` when a
6+
new predicate must outrank another on overlap.
67
"""
78

89
from collections.abc import Callable
@@ -17,30 +18,34 @@
1718
is_task_message_tool_result,
1819
is_task_retrieval_tool_result,
1920
)
20-
from utils.tool_dispatch import _TOOL_RESULT_DISPATCH
21+
from utils.tool_dispatch import (
22+
_TOOL_RESULT_DISPATCH,
23+
ToolResultDispatchEntry,
24+
_winning_dispatch_entry,
25+
)
2126

2227
Predicate = Callable[..., bool]
2328

2429
ORDERING_INVARIANTS: list[tuple[Predicate, Predicate, str]] = [
2530
(
2631
is_plan_tool_result,
2732
is_file_write_tool_result,
28-
"plan blobs may carry filePath + content; plan must win before file_write",
33+
"plan blobs may carry filePath + content; plan must outrank file_write",
2934
),
3035
(
3136
is_task_message_tool_result,
3237
is_task_retrieval_tool_result,
33-
"task_message is broad (task_id or message); must precede narrower task_retrieval",
38+
"task_message is broad (task_id or message); must outrank task_retrieval",
3439
),
3540
(
3641
is_task_message_tool_result,
3742
is_task_completed_tool_result,
38-
"task_message is broad (task_id or message); must precede narrower task_completed",
43+
"task_message is broad (task_id or message); must outrank task_completed",
3944
),
4045
(
4146
is_task_message_tool_result,
4247
is_task_async_tool_result,
43-
"task_message is broad (task_id or message); must precede narrower task_async",
48+
"task_message is broad (task_id or message); must outrank task_async",
4449
),
4550
]
4651

@@ -51,30 +56,64 @@
5156
"task_message_before_task_async",
5257
]
5358

59+
# Overlap blobs: keys chosen so both predicates in each ORDERING_INVARIANTS pair match.
60+
OVERLAP_BLOBS: dict[str, dict[str, object]] = {
61+
"plan_before_file_write": {
62+
"plan": {"name": "sprint-plan", "steps": ["index", "search", "render"]},
63+
"filePath": ".cursor/plans/week28.md",
64+
"content": "Plan body that would also satisfy file_write.",
65+
},
66+
"task_message_before_task_retrieval": {
67+
"task_id": "task-overlap-retrieval",
68+
"message": "polling retrieval",
69+
"retrieval_status": "found",
70+
"task": {"task_id": "task-overlap-retrieval", "description": "subagent scan"},
71+
},
72+
"task_message_before_task_completed": {
73+
"agentId": "agent-overlap-completed",
74+
"totalDurationMs": 3200,
75+
"totalTokens": 900,
76+
"status": "completed",
77+
"message": "task finished",
78+
},
79+
"task_message_before_task_async": {
80+
"agentId": "agent-overlap-async",
81+
"isAsync": True,
82+
"description": "explore auth handlers",
83+
"message": "task launched",
84+
},
85+
}
86+
5487

55-
def _predicate_index(predicate: Predicate) -> int:
56-
for i, entry in enumerate(_TOOL_RESULT_DISPATCH):
57-
pred = entry[0]
58-
# Identity match: dispatch table must store bare function refs (not wrappers).
59-
if pred is predicate:
60-
return i
88+
def _entry_for(predicate: Predicate) -> ToolResultDispatchEntry:
89+
for entry in _TOOL_RESULT_DISPATCH:
90+
if entry.predicate is predicate:
91+
return entry
6192
raise ValueError(f"predicate {predicate.__name__} not found in _TOOL_RESULT_DISPATCH")
6293

6394

6495
@pytest.mark.parametrize(
65-
"before,after,reason",
66-
ORDERING_INVARIANTS,
96+
"before,after,reason,fixture_id",
97+
[
98+
(*row, invariant_id)
99+
for row, invariant_id in zip(ORDERING_INVARIANTS, ORDERING_INVARIANT_IDS, strict=True)
100+
],
67101
ids=ORDERING_INVARIANT_IDS,
68102
)
69103
def test_tool_dispatch_ordering_invariant(
70104
before: Predicate,
71105
after: Predicate,
72106
reason: str,
107+
fixture_id: str,
73108
) -> None:
74-
before_idx = _predicate_index(before)
75-
after_idx = _predicate_index(after)
76-
assert before_idx < after_idx, (
77-
f"_TOOL_RESULT_DISPATCH ordering violation: "
78-
f"{before.__name__} (index {before_idx}) must precede "
79-
f"{after.__name__} (index {after_idx}). Reason: {reason}"
109+
blob = OVERLAP_BLOBS[fixture_id]
110+
before_entry = _entry_for(before)
111+
after_entry = _entry_for(after)
112+
assert before(blob), f"{fixture_id}: fixture no longer matches {before.__name__}"
113+
assert after(blob), f"{fixture_id}: fixture no longer matches {after.__name__}"
114+
winner = _winning_dispatch_entry(blob)
115+
assert winner is not None
116+
assert winner.id == before_entry.id, (
117+
f"_TOOL_RESULT_DISPATCH overlap violation: expected {before_entry.id!r} "
118+
f"to beat {after_entry.id!r} on {fixture_id}. Reason: {reason}"
80119
)

0 commit comments

Comments
 (0)