Import UI modules lazily in the CLI; pass categories to preparation

src/toad/cli.py

@@ -1,13 +1,17 @@
+from __future__ import annotations
+
 import os
 import sys
+from typing import TYPE_CHECKING
 
 import click
-from toad.app import ToadApp
-from toad.application_lifetime import ApplicationLifetime
-from toad.agent_schema import AgentDefinition
-from toad.render_backend import Renderer
 from toad.render_choices import RendererChoice
 
+if TYPE_CHECKING:
+    from toad.app import ToadApp
+    from toad.agent_schema import AgentDefinition
+    from toad.render_backend import Renderer
+
 
 def persist_terminal_failure(error: Exception) -> None:
     """Keep Textual's terminal failure in the existing diagnostic directory."""
@@ -181,6 +185,7 @@ def run(
     if serve:
         import shlex
 
+        from toad.application_lifetime import ApplicationLifetime
         from toad.web_server import ToadWebServer
 
         command_args = list(sys.argv)
@@ -202,6 +207,8 @@ def run(
         ApplicationLifetime.set_process_title("toad --serve")
         server.serve()
     else:
+        from toad.app import ToadApp
+
         app = ToadApp(
             mode=None if agent_data else "store",
             agent_data=agent_data,
@@ -286,6 +293,7 @@ def acp(
     if serve:
         import shlex
 
+        from toad.application_lifetime import ApplicationLifetime
         from toad.web_server import ToadWebServer
 
         command_components = [sys.argv[0], "acp", command]
@@ -307,6 +315,8 @@ def acp(
         server.serve()
 
     else:
+        from toad.app import ToadApp
+
         app = ToadApp(
             agent_data=agent_data, project_dir=project_dir, agent_session_id=session_id,
             renderer=renderer_from_cli(renderer),
@@ -363,6 +373,7 @@ def replay(path: str) -> None:
 )
 def serve(port: int, host: str, public_url: str | None = None) -> None:
     """Serve Toad as a web application."""
+    from toad.application_lifetime import ApplicationLifetime
     from toad.web_server import ToadWebServer
 
     server = ToadWebServer(
@@ -377,6 +388,7 @@ def about() -> None:
     """Show about information."""
 
     from toad import about
+    from toad.app import ToadApp
 
     app = ToadApp()
 

src/toad/transcript_publication.py

@@ -196,7 +196,8 @@ async def publish_page(self) -> None:
         history = TranscriptHistory(self.page, self.agent.get_transcript_page,
                                     fragments=fragments, committed=False)
         self.owner.prepare_reader(history)
-        await history.prepare_body(partial(self.source_current, self.page.after))
+        await history.prepare_body(partial(self.source_current, self.page.after),
+                                   selected=view.visible_categories)
         async with AsyncExitStack() as retirement:
             async with self.window.history_lock:
                 # Preparation may yield to another accepted source publication.
@@ -447,7 +448,7 @@ async def publish_source_page(self, page: TranscriptPage) -> bool:
             replacement = TranscriptHistory(
                 page, self.agent.get_transcript_page, fragments=prepared.fragments, committed=False,
             )
-            await replacement.prepare_body(is_current)
+            await replacement.prepare_body(is_current, selected=view.visible_categories)
         async with AsyncExitStack() as retirement:
             async with window.history_lock:
                 retired = retirement_candidates(contents.children, evidence)

src/toad/widgets/conversation.py

@@ -1171,7 +1171,11 @@ async def action_provider_login(self) -> None:
     async def on_change_model(self, event: CoreEventMessage) -> None:
         if (agent := self.agent) is None:
             return
-        error = await agent.set_model(event.event.model_id)
+        # Selecting the current model opens its thinking choices. Reapplying it
+        # would ask native SetModel to restore that model's default effort.
+        error = None
+        if event.event.model_id != agent.configuration.model.current:
+            error = await agent.set_model(event.event.model_id)
         if agent is not self.agent:
             return
         if error is not None:
@@ -1209,7 +1213,7 @@ async def on_change_model(self, event: CoreEventMessage) -> None:
                     return
             self.flash(
                 Content.from_markup(
-                    "Model changed to [b]$model[/] · thinking [b]$level",
+                    "Model [b]$model[/] · thinking [b]$level",
                     model=model.name,
                     level=agent.configuration.thinking.current or "unavailable",
                 ),

src/toad/widgets/prepared_markdown.py

@@ -86,40 +86,51 @@ def update_slice(self, fragments, follow: bool) -> slice:
         stop = min(self.stop, len(fragments))
         return slice(min(self.start, stop), stop)
 
-    async def extend(self, older: bool, current: Callable[[], bool]) -> bool:
+    async def extend(self, older: bool, current: Callable[[], bool], *, prefix=()) -> bool:
         admission = self.capture_admission()
-        selected = self.extension_slice(older)
-        added = tuple(self._body(fragment) for fragment in self.fragments[selected])
-        previous = self.fragment_views
-        with ExitStack() as acquisition:
-            if added:
-                acquisition.callback(self.remove_children, added)
-                await self.mount_all(added, before=previous[0] if older and previous else None)
-            if not current() or self.capture_admission() != admission:
-                return False
-            self._fragment_views = (*added, *previous) if older else (*previous, *added)
-            if older:
-                self.start = selected.start
-            else:
-                self.stop = selected.stop
-            acquisition.pop_all()
-        return True
+        extension = self.extension_slice(older)
+        selected = slice(extension.start if older else self.start,
+                         self.stop if older else extension.stop)
+        previous = {self.start + index: child for index, child in enumerate(self.fragment_views)}
+        return await self.replace_range(
+            self.fragments, selected, previous,
+            lambda: current() and self.capture_admission() == admission,
+            prefix=prefix,
+        )
 
     async def replace_range(self, fragments, selected, previous, current, *, prefix=()) -> bool:
         """Publish ordered native parts inside their owner's source custody."""
-        ordered = tuple(previous[index] if index in previous else self._body(fragments[index])
-                        for index in range(selected.start, selected.stop))
-        added = tuple(child for child in ordered if child not in previous.values())
+        ordered = []
+        added = []
         with ExitStack() as acquisition:
-            if added:
-                acquisition.callback(self.remove_children, added)
-                await self.mount_all(added)
+            acquisition.callback(lambda: self.remove_children(added))
+            # Source demand may span a viewport. Native construction/mount
+            # remains a bounded burst, including after worker preparation.
+            # Keep the original admission until every batch has joined.
+            for first in range(selected.start, selected.stop, self.BATCH):
+                if not current():
+                    return False
+                batch = []
+                for index in range(first, min(selected.stop, first + self.BATCH)):
+                    if index in previous:
+                        child = previous[index]
+                    else:
+                        child = self._body(fragments[index])
+                        added.append(child)
+                        batch.append(child)
+                    ordered.append(child)
+                if batch:
+                    await self.mount_all(batch)
+                    if not current():
+                        return False
             if not current():
                 return False
-            self.fragments, self._fragment_views = fragments, ordered
+            self.fragments, self._fragment_views = fragments, tuple(ordered)
             self.start, self.stop = selected.start, selected.stop
-            rank = {child: index for index, child in enumerate((*prefix, *ordered))}
-            self.sort_children(key=lambda child: rank.get(child, len(rank)))
+            desired = (*prefix, *ordered)
+            if tuple(self.children) != desired:
+                rank = {child: index for index, child in enumerate(desired)}
+                self.sort_children(key=lambda child: rank.get(child, len(rank)))
             self.remove_children(tuple(child for child in previous.values() if child not in ordered))
             acquisition.pop_all()
         return True

src/toad/widgets/presentation_window.py

@@ -87,16 +87,16 @@ def admit_paint(self, candidates, required, max_bytes: int):
         """Retain message paint independently of its disposable native tree.
 
         The native widget limit still governs lookahead materialization. Paint
-        admission bounds original message resources and their source/row bytes;
-        a live tree awaiting capture must not price its eventual retained rows.
+        admission bounds original message resources by their source/row bytes.
+        Native item counts price source slots and preparation, not immutable
+        paint; a live tree awaiting capture must not price its eventual rows.
         """
         admitted = set(required)
         source_bytes = 0
-        limit = self.item_limit(len(admitted))
         for owner in candidates:
             size = owner.retained_source_bytes + owner.retained_paint_bytes
             if owner not in admitted:
-                if len(admitted) >= limit or source_bytes + size > max_bytes:
+                if source_bytes + size > max_bytes:
                     continue
                 admitted.add(owner)
             source_bytes += size

src/toad/widgets/streaming_markdown.py

@@ -228,7 +228,7 @@ def current():
                         self.fragments, selected, previous, current, prefix=self._prefix,
                     ):
                         return False
-                elif not await self.extend(older, current):
+                elif not await self.extend(older, current, prefix=self._prefix):
                     return False
                 visible = self.screen._compositor.visible_widgets
                 protected = protected_presentations(self.fragment_views, self.screen._interaction_widgets())

src/toad/widgets/transcript_fragments.py

@@ -16,7 +16,7 @@
 from toad.render_backend import ReusableRenderTask
 from toad.acp.status import ToolCallStatus
 from toad.widgets.agent_activity import AgentActivityBoundary
-from toad.widgets.message_filter import event_category
+from toad.widgets.message_filter import event_category, keep_events
 
 if TYPE_CHECKING:
     from toad.render_backend import Renderer
@@ -246,8 +246,9 @@ def transcript_fragments(
 class TranscriptBodyPreparation(MroDispatch):
     """Pure body work for declared transcript cases, without native widgets."""
 
-    def __init__(self, renderer, ansi: bool, dark: bool):
+    def __init__(self, renderer, ansi: bool, dark: bool, *, selected=None):
         self.renderer, self.ansi, self.dark = renderer, ansi, dark
+        self.selected = selected
 
     async def prepare_fragments(self, fragments, keep_going, *, batch_size: int) -> None:
         """Warm a bounded source range in shared workers, without native mounts.
@@ -260,6 +261,7 @@ async def prepare_fragments(self, fragments, keep_going, *, batch_size: int) ->
                 return
             await asyncio.gather(*(self.dispatch(event)
                                    for fragment in fragments[first:first + batch_size]
+                                   if self.selected is None or keep_events(fragment.events, self.selected)
                                    for event in fragment.events))
 
     @handles(TranscriptEvent)

src/toad/widgets/transcript_history.py

@@ -165,10 +165,17 @@ def message_category(self) -> type[MessageCategory]:
     def set_categories(self, selected: frozenset[type[MessageCategory]]) -> None:
         if getattr(self, "selected_categories", None) == selected:
             return
+        visible = self.message_category in getattr(self, "selected_categories", ())
         self.selected_categories = selected
         apply_block_filter(self, selected)
+        if self.is_attached and not visible and self.message_category in selected and not self.children:
+            self.publish_body(self.materialize_native_body).call_next(self)
 
     def compose(self) -> ComposeResult:
+        # A filtered source slot retains identity/coverage, not hidden native
+        # bodies. Its existing body publisher materializes on first reveal.
+        if self.message_category not in self.selected_categories:
+            return
         # A semantic fragment may be one oversized paragraph, list, or fence.
         # It is already a page leaf: re-paging it would recursively remount the
         # same indivisible block forever without producing visible Markdown.
@@ -241,13 +248,19 @@ async def acquire(
         batch_size: int, before: Widget, current: Callable[[], bool], newest: bool = True,
     ) -> AsyncIterator[TranscriptPageView]:
         """Acquire a page until its original source transfers native custody."""
-        view = cls(page, fragments=fragments, batch_size=batch_size, newest=newest)
+        view = cls(page, fragments=fragments, batch_size=min(batch_size, cls.BATCH), newest=newest)
         view.visible_categories = owner.selected_categories
         with ExitStack() as acquisition:
             acquisition.callback(owner.remove_children, (view,))
             await owner.mount(view, before=before)
             if not current():
                 raise _PublicationRetired
+            selected = cls.initial_slice(fragments, batch_size, newest)
+            if selected != slice(view.start, view.stop):
+                previous = {view.start + index: child for index, child in enumerate(view.fragment_views)}
+                if not await view.replace_range(fragments, selected, previous, current):
+                    raise _PublicationRetired
+            view.batch_size = batch_size
             yield view
             acquisition.pop_all()
 
@@ -332,19 +345,20 @@ def __init__(self, page: TranscriptPage, loader: Callable[..., Awaitable[Transcr
         self._fragment_budget = self.budget.max_items
         self.window: HistoryWindow
 
-    async def prepare_fragments(self, fragments, current: Callable[[], bool]) -> None:
+    async def prepare_fragments(self, fragments, current: Callable[[], bool], *, selected=None) -> None:
         """Prepare this admitted source range before acquiring native custody."""
         from toad.widgets.transcript_fragments import TranscriptBodyPreparation
 
         preparation = TranscriptBodyPreparation(
             self.app.render_processes, self.app.native_ansi_color, self.app.current_theme.dark,
+            selected=self.selected_categories if selected is None else selected,
         )
         await preparation.prepare_fragments(fragments, current, batch_size=self.budget.admission_items)
 
-    async def prepare_body(self, current: Callable[[], bool]) -> None:
+    async def prepare_body(self, current: Callable[[], bool], *, selected=None) -> None:
         """Warm the actual page admissions, including a restored reader range."""
         for page in self.pages:
-            await self.prepare_fragments(page.fragments[page.start:page.stop], current)
+            await self.prepare_fragments(page.fragments[page.start:page.stop], current, selected=selected)
 
     @property
     def fragment_views(self) -> tuple[TranscriptFragmentView, ...]:
@@ -628,6 +642,7 @@ async def prepare_lookahead(self, intent) -> None:
         from toad.widgets.transcript_fragments import TranscriptBodyPreparation
         preparation = TranscriptBodyPreparation(
             self.app.render_processes, self.app.native_ansi_color, self.app.current_theme.dark,
+            selected=self.selected_categories,
         )
         for page in pages:
             await page.prepare_adjacent(preparation, demand, count, current)

tests/model_picker_pilot.py

@@ -150,9 +150,10 @@ def ids():
             await pilot.press(*"glm")
             await pilot.pause()
             assert await pilot.click(picker.option_list, offset=(2, 0))
-            await until(lambda: agent.calls and agent.calls[-1] == glm)
+            await until(lambda: not picker.is_open)
             await pilot.pause()
             assert not picker.is_open
+            assert not agent.calls, "Opening the current model must not reapply it"
 
             await CommandCatalog.execute("/model", conversation)
             await pilot.pause()