Carry parsed Markdown through saved transcript rendering

src/toad/markdown_preparation.py

@@ -1,7 +1,11 @@
 """Process-safe conversation parsing and code-fence highlighting data."""
 
+from copy import deepcopy
+from collections.abc import Iterable, Iterator
 from dataclasses import dataclass, field
 from functools import cached_property
+from hashlib import sha256
+import pickle
 
 from markdown_it.token import Token
 from textual.content import Content
@@ -13,9 +17,82 @@
 
 @dataclass(frozen=True)
 class PreparedMarkdownPart:
-    """One bounded native source part, without an invented transcript role."""
-
-    text: str
+    """Source and syntax acquired through the conversation parser declaration.
+
+    Syntax precedes project-link resolution. Every document acquires its own
+    mutable tokens; neither partitioning nor presentation mutates this source.
+    Custom parser grammars retain their native acquisition instead.
+    """
+
+    text: str = field(compare=False)
+    tokens: tuple[Token, ...] = field(compare=False, repr=False)
+    revision: bytes = field(init=False)
+
+    def __post_init__(self) -> None:
+        # Source acquisition/partitioning runs in preparation. Native refresh
+        # compares this revision without walking an arbitrarily nested tree.
+        object.__setattr__(self, "revision", sha256(pickle.dumps((self.text, self.tokens), protocol=5)).digest())
+
+    @classmethod
+    def capture(cls, text: str) -> "PreparedMarkdownPart":
+        from toad.conversation_markdown import parse_markdown_syntax
+
+        return cls(text, tuple(parse_markdown_syntax(text)))
+
+    def acquire_tokens(self) -> list[Token]:
+        return deepcopy(list(self.tokens))
+
+    def select_blocks(self, text: str, first: int, tokens: slice) -> "PreparedMarkdownPart":
+        """Select complete top-level blocks, retaining their closing tokens."""
+        if text == self.text and first == 0 and tokens.start == 0 and tokens.stop == len(self.tokens):
+            return self
+        selected = deepcopy(list(self.tokens[tokens]))
+        for token in self._descendants(selected):
+            if token.map is not None:
+                token.map = [line - first for line in token.map]
+        return PreparedMarkdownPart(text, tuple(selected))
+
+    def select_table_rows(self, ranges: Iterable[tuple[int, int]]) -> Iterator["PreparedMarkdownPart"]:
+        """Repeat the original parsed header, never parse a synthetic table."""
+        lines = self.text.splitlines(keepends=True)
+        rows = {}
+        body_first = body_stop = 0
+        for index, token in enumerate(self.tokens):
+            if token.type == "tbody_open":
+                body_first = index + 1
+            elif token.type == "tbody_close":
+                body_stop = index
+                break
+            elif body_first and token.type == "tr_open":
+                assert token.map is not None
+                rows[token.map[0]] = index
+        table = self.tokens[0]
+        assert table.type == "table_open" and table.map is not None
+        table_stop = table.map[1]
+        for first, stop in ranges:
+            row_first = rows.get(first, body_stop)
+            row_stop = rows.get(stop, body_stop)
+            tokens = deepcopy(list((*self.tokens[:body_first],
+                                    *self.tokens[row_first:row_stop],
+                                    *self.tokens[body_stop:])))
+            length = 2 + min(stop, table_stop) - first
+            for token in self._descendants(tokens):
+                if token.map is None:
+                    continue
+                if token.type == "table_open":
+                    token.map = [0, length]
+                elif token.type == "tbody_open":
+                    token.map = [2, length]
+                elif token.map[0] >= 2:
+                    token.map = [line - first + 2 for line in token.map]
+            yield PreparedMarkdownPart("".join((*lines[:2], *lines[first:stop])), tuple(tokens))
+
+    @staticmethod
+    def _descendants(tokens):
+        for token in tokens:
+            yield token
+            if token.children is not None:
+                yield from PreparedMarkdownPart._descendants(token.children)
 
     @cached_property
     def retained_bytes(self) -> int:
@@ -33,24 +110,20 @@ class PreparedFence:
 class PreparedMarkdown:
     tokens: list[Token]
     fences: dict[FenceKey, PreparedFence]
-    inlines: dict[int, Content] = field(default_factory=dict)
 
     def acquire_inline_content(self, tokens: list[Token]) -> "PreparedMarkdown":
         """Bind native content to the actual delivered, link-resolved tokens.
 
-        Token identities belong to this one acquisition, not the reusable
-        parser result or another document. Native block construction consumes
-        these same token objects and retains the resulting Content itself.
+        The parser's token metadata owns its acquired content. Native block
+        construction consumes these same tokens, without an identity registry.
         """
-        def descendants(source):
-            for token in source:
-                yield token
-                if token.children is not None:
-                    yield from descendants(token.children)
-
-        inlines = {id(token): MarkdownBlock._token_to_content(token)
-                   for token in descendants(tokens) if token.type == "inline"}
-        return PreparedMarkdown(tokens, self.fences, inlines)
+        for token in PreparedMarkdownPart._descendants(tokens):
+            if token.type == "inline":
+                token.meta[PreparedMarkdown] = MarkdownBlock._token_to_content(token)
+        return PreparedMarkdown(tokens, self.fences)
+
+    def inline_content(self, token: Token) -> Content:
+        return token.meta[PreparedMarkdown]
 
 
 def prepare_tokens(tokens: list[Token], ansi: bool, dark: bool) -> PreparedMarkdown:

src/toad/render_tasks.py

@@ -18,6 +18,7 @@
 
 if TYPE_CHECKING:
     from toad.widgets.patch_diff import PreparedPatch
+    from toad.work_preparation import PreparationRuntime, RenderPreparation, WorkKey
 
 
 @dataclass(frozen=True)
@@ -59,16 +60,25 @@ def accept_result(self, result: object) -> PreparedPatch:
 
 @dataclass(frozen=True)
 class MarkdownRenderTask(ReusableRenderTask[PreparedMarkdown]):
-    """Parse and highlight one detached body before its independent delivery."""
+    """Acquire syntax once, then highlight before independent delivery."""
 
-    source: str
+    source: str | PreparedMarkdownPart
     ansi: bool
     dark: bool
 
-    def execute(self) -> PreparedMarkdown:
-        from toad.conversation_markdown import parse_markdown_syntax
+    async def preparation_identity(self, work: RenderPreparation, runtime: PreparationRuntime) -> WorkKey:
+        if isinstance(self.source, PreparedMarkdownPart):
+            from toad.work_preparation import WorkKey
 
-        return prepare_tokens(parse_markdown_syntax(self.source), self.ansi, self.dark)
+            # Acquired syntax owns its revision; accounting and delivery copies
+            # cannot change the original preparation identity.
+            return WorkKey(type(work), (type(self), self.source.revision, self.ansi, self.dark), work.scope)
+        return await super().preparation_identity(work, runtime)
+
+    def execute(self) -> PreparedMarkdown:
+        part = (PreparedMarkdownPart.capture(self.source)
+                if isinstance(self.source, str) else self.source)
+        return prepare_tokens(part.acquire_tokens(), self.ansi, self.dark)
 
     def accept_result(self, result: object) -> PreparedMarkdown:
         if not isinstance(result, PreparedMarkdown):
@@ -80,12 +90,21 @@ def accept_result(self, result: object) -> PreparedMarkdown:
 class MarkdownPartsTask(ReusableRenderTask[tuple[PreparedMarkdownPart, ...]]):
     """Partition one original message with the shared Markdown block budget."""
 
-    source: str
+    source: str | PreparedMarkdownPart
+
+    async def preparation_identity(self, work: RenderPreparation, runtime: PreparationRuntime) -> WorkKey:
+        if isinstance(self.source, PreparedMarkdownPart):
+            from toad.work_preparation import WorkKey
+
+            return WorkKey(type(work), (type(self), self.source.revision), work.scope)
+        return await super().preparation_identity(work, runtime)
 
     def execute(self) -> tuple[PreparedMarkdownPart, ...]:
         from toad.widgets.transcript_fragments import RenderBudget
 
-        parts = tuple(PreparedMarkdownPart(text) for text in RenderBudget().split(self.source))
+        part = (PreparedMarkdownPart.capture(self.source)
+                if isinstance(self.source, str) else self.source)
+        parts = tuple(RenderBudget().split(part))
         for part in parts:
             part.retained_bytes
         return parts

src/toad/widgets/agent_response.py

@@ -8,6 +8,7 @@
 from toad.widgets.message_divider import MessageClock, LiveMessageClock, MessageDivider
 from toad.widgets.route_header import RouteHeader
 from toad import response_delivery
+from toad.markdown_preparation import PreparedMarkdownPart
 from toad.widgets.message_filter import (
     CategorizedBlock,
     MessageCategory,
@@ -32,10 +33,12 @@ def message_category(self) -> type[MessageCategory]:
 
     def __init__(self, markdown: str | None = None, *, delivery: response_delivery.ResponseDelivery = response_delivery.UnroutedResponse(),
                  category: type[MessageCategory] | None = None,
+                 markdown_part: PreparedMarkdownPart | None = None,
                  paginate: bool = True, show_divider: bool = True, clock: MessageClock = LiveMessageClock()) -> None:
         self._message_category = category or delivery.category
         super().__init__(
             markdown,
+            markdown_part=markdown_part,
             paginate=paginate,
             **self.dispatch_sync(delivery, clock, show_divider),
         )

src/toad/widgets/incoming_message.py

@@ -12,6 +12,7 @@
 from toad.widgets.committed_presentation import CommitParticipant, SequenceClaim
 from toad.widgets.wire_message_handling import WireMessageHandling
 from toad.widgets.message_notifications import MessageNotifications
+from toad.markdown_preparation import PreparedMarkdownPart
 
 
 
@@ -26,10 +27,12 @@ def message_category(self) -> type[MessageCategory]:
     def commit_claim(self) -> SequenceClaim:
         return SequenceClaim(self.sequence)
 
-    def __init__(self, event: IncomingTranscript, *, show_header: bool = True) -> None:
+    def __init__(self, event: IncomingTranscript, *, show_header: bool = True,
+                 markdown_part: PreparedMarkdownPart | None = None) -> None:
         super().__init__()
         self.event = event
         self.show_header = show_header
+        self.markdown_part = markdown_part
 
     @property
     def message_reference(self):
@@ -46,7 +49,8 @@ def compose(self) -> ComposeResult:
             yield MessageDivider(f"Inbound · @{self.event.route.sender}",
                                  clock=MessageClock.recorded(self.event.timestamp))
             yield RouteHeader(self.event.route, incoming=True)
-        yield AgentResponse(self.event.text, show_divider=False).add_class("routed-body")
+        yield AgentResponse(self.event.text, show_divider=False,
+                            markdown_part=self.markdown_part).add_class("routed-body")
         if self.handling_references:
             yield MessageNotifications()
 

src/toad/widgets/outgoing_message.py

@@ -11,15 +11,18 @@
 from toad.widgets.message_filter import CategorizedBlock, OutboundCategory
 from toad.widgets.message_notifications import MessageNotifications
 from toad.widgets.wire_message_handling import WireMessageHandling
+from toad.markdown_preparation import PreparedMarkdownPart
 
 
 class OutgoingMessage(WireMessageHandling, ConversationBlock, CategorizedBlock, VerticalGroup):
     DEFAULT_CLASSES = "block"
 
-    def __init__(self, event: SentTranscript, *, show_header: bool = True):
+    def __init__(self, event: SentTranscript, *, show_header: bool = True,
+                 markdown_part: PreparedMarkdownPart | None = None):
         super().__init__()
         self.event = event
         self.show_header = show_header
+        self.markdown_part = markdown_part
 
     @property
     def message_reference(self):
@@ -31,6 +34,7 @@ def message_category(self):
 
     def compose(self) -> ComposeResult:
         yield AgentResponse(self.event.text,
+                            markdown_part=self.markdown_part,
                             delivery=ResponseDelivery.from_route(self.event.routing.reply),
                             show_divider=self.show_header,
                             clock=MessageClock.recorded(self.event.timestamp))

src/toad/widgets/prepared_markdown.py

@@ -20,10 +20,9 @@
 
 from toad.app import ToadApp
 from toad.conversation_markdown import ConversationCodeFence, ConversationMarkdown, _ThreadLocalPathParser
-from toad.markdown_preparation import PreparedMarkdown
+from toad.markdown_preparation import PreparedMarkdown, PreparedMarkdownPart
 from toad.block_content import MarkdownBlockContent
 from toad.render_tasks import MarkdownRenderTask
-from toad.widgets.transcript_fragments import RenderBudget
 from toad.widgets.viewport_body import MeasuredViewportBody
 from toad.widgets.worker_static import WorkerStatic
 
@@ -227,12 +226,32 @@ def __init__(
         self, markdown: str | None = None, *, name: str | None = None,
         id: str | None = None, classes: str | None = None,
         parser_factory: Callable[[], MarkdownIt] | None = None, open_links: bool = False,
+        markdown_part: PreparedMarkdownPart | None = None,
     ) -> None:
         self._prepared_markdown: PreparedMarkdown | None = None
+        self._markdown_part = markdown_part
         factory = self._make_parser if parser_factory is None else parser_factory
         super().__init__(markdown, name=name, id=id, classes=classes,
                          parser_factory=factory, open_links=open_links)
 
+    @property
+    def partitionable_syntax(self) -> bool:
+        """Only the original conversation factory declares this grammar.
+
+        An overridden factory or arbitrary custom parser has no declaration
+        permitting conversation-syntax partitioning, even if today's output
+        happens to look the same.
+        """
+        return getattr(self._parser_factory, "__func__", None) is ConversationMarkdown._make_parser
+
+    def acquired_source(self, markdown: str) -> str | PreparedMarkdownPart:
+        part = self._markdown_part
+        return part if part is not None and part.text == markdown else markdown
+
+    def update_part(self, part: PreparedMarkdownPart) -> AwaitComplete:
+        self._markdown_part = part
+        return self.update(part.text)
+
     def reconstructible_children(self) -> tuple[Widget, ...]:
         """Native resources rebuilt from this document's original source."""
         return tuple(child for child in self.children if isinstance(child, MarkdownBlock))
@@ -259,9 +278,13 @@ def append(self, markdown: str) -> AwaitComplete:
         return self.publish_body(partial(self._append_body_source, markdown))
 
     def _update_body_source(self, markdown: str) -> AwaitComplete:
+        if self._markdown_part is not None and self._markdown_part.text != markdown:
+            self._markdown_part = None
         return super().update(markdown)
 
     def _append_body_source(self, markdown: str) -> AwaitComplete:
+        if self._markdown_part is not None and self._markdown_part.text != self.source + markdown:
+            self._markdown_part = None
         return super().append(markdown)
 
     def on_unmount(self) -> None:
@@ -273,7 +296,7 @@ async def _parse_tokens(
         # Custom parser factories retain their native semantics. This process
         # function implements the declaration-owned conversation/path parser.
         parent = self.parent
-        if not isinstance(parser, _ThreadLocalPathParser):
+        if type(parser) is not _ThreadLocalPathParser:
             tokens = await super()._parse_tokens(parser, markdown, use_thread=use_thread)
             if tokens is not None:
                 prepared = await asyncio.to_thread(
@@ -284,22 +307,24 @@ async def _parse_tokens(
                 self._prepared_markdown = prepared
             return tokens
         if not isinstance(self.app, ToadApp):
-            tokens = (await asyncio.to_thread(parser.parse, markdown)
-                      if use_thread else parser.parse(markdown))
-            prepared = await asyncio.to_thread(
-                PreparedMarkdown(tokens, {}).acquire_inline_content, tokens,
-            )
+            source = self.acquired_source(markdown)
+            def acquire_local():
+                tokens = (parser.resolve_tokens(source.acquire_tokens())
+                          if isinstance(source, PreparedMarkdownPart) else parser.parse(source))
+                return PreparedMarkdown(tokens, {}).acquire_inline_content(tokens)
+
+            prepared = await asyncio.to_thread(acquire_local)
             if self._closing or self._pruning or self.parent is not parent:
                 return None
             self._prepared_markdown = prepared
-            return tokens
+            return prepared.tokens
 
         # Runway and delivery use the same pure preparation key. Filesystem
         # links are resolved only in each independent delivered token resource,
         # after highlighting; they never become reusable renderer inputs.
         while not self._closing and self.is_attached and not self._pruning:
             theme = (self.app.native_ansi_color, self.app.current_theme.dark)
-            request = self.app.render_processes.submit(MarkdownRenderTask(markdown, *theme))
+            request = self.app.render_processes.submit(MarkdownRenderTask(self.acquired_source(markdown), *theme))
             worker = self.run_worker(request, group="markdown-preparation", exit_on_error=False)
             try:
                 prepared = await worker.wait()
@@ -338,7 +363,7 @@ def _get_token_content(self, token: Token, *, block: MarkdownBlock) -> Content:
             # Its declared answer cannot be replaced by the native pure result.
             return super()._get_token_content(token, block=block)
         assert self._prepared_markdown is not None
-        return self._prepared_markdown.inlines[id(token)]
+        return self._prepared_markdown.inline_content(token)
 
     def get_block_class(self, block_name: str) -> type[MarkdownBlock]:
         if block_name in {"fence", "code_block"}: