diff --git a/libs/cli/deepagents_cli/app.py b/libs/cli/deepagents_cli/app.py index 65f95bdac7e..a3ef8560d46 100644 --- a/libs/cli/deepagents_cli/app.py +++ b/libs/cli/deepagents_cli/app.py @@ -532,6 +532,13 @@ async def on_mount(self) -> None: ) self._ui_adapter.set_token_tracker(self._token_tracker) + # Prewarm `/threads` cache in the background so first open is faster. + self.run_worker( + self._prewarm_threads_cache, + exclusive=True, + group="startup-thread-prewarm", + ) + # Focus the input (autocomplete is now built into ChatInput) self._chat_input.focus_input() @@ -586,6 +593,15 @@ def on_resize(self, _event: Resize) -> None: except NoMatches: pass # Spacer already removed, no action needed + async def _prewarm_threads_cache(self) -> None: # noqa: PLR6301 # Worker hook kept as instance method + """Prewarm thread selector cache without blocking app startup.""" + from deepagents_cli.sessions import ( + get_thread_limit, + prewarm_thread_message_counts, + ) + + await prewarm_thread_message_counts(limit=get_thread_limit()) + def on_scroll_up(self, _event: ScrollUp) -> None: """Handle scroll up to check if we need to hydrate older messages.""" self._check_hydration_needed() @@ -2316,7 +2332,11 @@ def handle_result(result: tuple[str, str] | None) -> None: async def _show_thread_selector(self) -> None: """Show interactive thread selector as a modal screen.""" + from deepagents_cli.sessions import get_cached_threads, get_thread_limit + current = self._session_state.thread_id if self._session_state else None + thread_limit = get_thread_limit() + initial_threads = get_cached_threads(limit=thread_limit) def handle_result(result: str | None) -> None: """Handle the thread selector result.""" @@ -2325,9 +2345,36 @@ def handle_result(result: str | None) -> None: if self._chat_input: self._chat_input.focus_input() - screen = ThreadSelectorScreen(current_thread=current) + screen = ThreadSelectorScreen( + current_thread=current, + thread_limit=thread_limit, + initial_threads=initial_threads, + ) self.push_screen(screen, handle_result) + def _update_welcome_banner( + self, + thread_id: str, + *, + missing_message: str, + warn_if_missing: bool, + ) -> None: + """Update the welcome banner thread ID when the banner is mounted. + + Args: + thread_id: Thread ID to display on the banner. + missing_message: Log message template when banner is missing. + warn_if_missing: Whether to log missing-banner cases at warning level. + """ + try: + banner = self.query_one("#welcome-banner", WelcomeBanner) + banner.update_thread_id(thread_id) + except NoMatches: + if warn_if_missing: + logger.warning(missing_message, thread_id) + else: + logger.debug(missing_message, thread_id) + async def _resume_thread(self, thread_id: str) -> None: """Resume a previously saved thread. @@ -2366,11 +2413,36 @@ async def _resume_thread(self, thread_id: str) -> None: if self._chat_input: self._chat_input.set_cursor_active(active=False) + prefetched_history: list[MessageData] | None = None try: - try: - self._update_status(f"Loading thread: {thread_id}") - prefetched_history = await self._fetch_thread_history_data(thread_id) - except Exception as exc: + self._update_status(f"Loading thread: {thread_id}") + prefetched_history = await self._fetch_thread_history_data(thread_id) + + # Clear conversation (similar to /clear, without creating a new thread) + self._pending_messages.clear() + self._queued_widgets.clear() + await self._clear_messages() + if self._token_tracker: + self._token_tracker.reset() + self._update_status("") + + # Switch to the selected thread + self._session_state.thread_id = thread_id + self._lc_thread_id = thread_id + + self._update_welcome_banner( + thread_id, + missing_message="Welcome banner not found during thread switch to %s", + warn_if_missing=False, + ) + + # Load thread history + await self._load_thread_history( + thread_id=thread_id, + preloaded_data=prefetched_history, + ) + except Exception as exc: + if prefetched_history is None: logger.exception("Failed to prefetch history for thread %s", thread_id) await self._mount_message( AppMessage( @@ -2379,68 +2451,35 @@ async def _resume_thread(self, thread_id: str) -> None: ) ) return - + logger.exception("Failed to switch to thread %s", thread_id) + # Restore previous thread IDs so the user can retry + self._session_state.thread_id = prev_session_thread + self._lc_thread_id = prev_thread_id + self._update_welcome_banner( + prev_session_thread, + missing_message=( + "Welcome banner not found during rollback to thread %s; " + "banner may display stale thread ID" + ), + warn_if_missing=True, + ) + rollback_restore_failed = False + # Attempt to restore the previous thread's visible history try: - # Clear conversation (similar to /clear, without creating a new thread) - self._pending_messages.clear() - self._queued_widgets.clear() await self._clear_messages() - if self._token_tracker: - self._token_tracker.reset() - self._update_status("") - - # Switch to the selected thread - self._session_state.thread_id = thread_id - self._lc_thread_id = thread_id - - # Update welcome banner - try: - banner = self.query_one("#welcome-banner", WelcomeBanner) - banner.update_thread_id(thread_id) - except NoMatches: - logger.debug( - "Welcome banner not found during thread switch to %s", - thread_id, - ) - - # Load thread history - await self._load_thread_history( - thread_id=thread_id, - preloaded_data=prefetched_history, + await self._load_thread_history(thread_id=prev_session_thread) + except Exception: # Resilient session state saving + rollback_restore_failed = True + msg = ( + "Could not restore previous thread history after failed " + "switch to %s" ) - - except Exception as exc: - logger.exception("Failed to switch to thread %s", thread_id) - # Restore previous thread IDs so the user can retry - self._session_state.thread_id = prev_session_thread - self._lc_thread_id = prev_thread_id - try: - banner = self.query_one("#welcome-banner", WelcomeBanner) - banner.update_thread_id(prev_session_thread) - except NoMatches: - logger.warning( - "Welcome banner not found during rollback to thread %s; " - "banner may display stale thread ID", - prev_session_thread, - ) - rollback_restore_failed = False - # Attempt to restore the previous thread's visible history - try: - await self._clear_messages() - await self._load_thread_history(thread_id=prev_session_thread) - except Exception: # Resilient session state saving - rollback_restore_failed = True - logger.warning( - "Could not restore previous thread history after " - "failed switch to %s", - thread_id, - exc_info=True, - ) - error_message = f"Failed to switch to thread {thread_id}: {exc}." - if rollback_restore_failed: - error_message += " Previous thread history could not be restored." - error_message += " Use /threads to try again." - await self._mount_message(AppMessage(error_message)) + logger.warning(msg, thread_id, exc_info=True) + error_message = f"Failed to switch to thread {thread_id}: {exc}." + if rollback_restore_failed: + error_message += " Previous thread history could not be restored." + error_message += " Use /threads to try again." + await self._mount_message(AppMessage(error_message)) finally: self._thread_switching = False self._update_status("") diff --git a/libs/cli/deepagents_cli/sessions.py b/libs/cli/deepagents_cli/sessions.py index 119b63883d1..8ff2846c870 100644 --- a/libs/cli/deepagents_cli/sessions.py +++ b/libs/cli/deepagents_cli/sessions.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import logging +import sqlite3 import uuid from contextlib import asynccontextmanager from datetime import datetime @@ -19,6 +21,11 @@ logger = logging.getLogger(__name__) _aiosqlite_patched = False +_jsonplus_serializer: JsonPlusSerializer | None = None +_message_count_cache: dict[str, tuple[str | None, int]] = {} +_MAX_MESSAGE_COUNT_CACHE = 4096 +_recent_threads_cache: dict[tuple[str | None, int], list[ThreadInfo]] = {} +_MAX_RECENT_THREADS_CACHE_KEYS = 16 def _patch_aiosqlite() -> None: @@ -84,6 +91,9 @@ class ThreadInfo(TypedDict): message_count: NotRequired[int] """Number of messages in the thread.""" + latest_checkpoint_id: NotRequired[str | None] + """Most recent checkpoint ID for cache invalidation.""" + def format_timestamp(iso_timestamp: str | None) -> str: """Format ISO timestamp for display (e.g., 'Dec 30, 6:10pm'). @@ -169,7 +179,8 @@ async def list_threads( query = """ SELECT thread_id, json_extract(metadata, '$.agent_name') as agent_name, - MAX(json_extract(metadata, '$.updated_at')) as updated_at + MAX(json_extract(metadata, '$.updated_at')) as updated_at, + MAX(checkpoint_id) as latest_checkpoint_id FROM checkpoints WHERE json_extract(metadata, '$.agent_name') = ? GROUP BY thread_id @@ -181,7 +192,8 @@ async def list_threads( query = """ SELECT thread_id, json_extract(metadata, '$.agent_name') as agent_name, - MAX(json_extract(metadata, '$.updated_at')) as updated_at + MAX(json_extract(metadata, '$.updated_at')) as updated_at, + MAX(checkpoint_id) as latest_checkpoint_id FROM checkpoints GROUP BY thread_id ORDER BY updated_at DESC @@ -192,22 +204,208 @@ async def list_threads( async with conn.execute(query, params) as cursor: rows = await cursor.fetchall() threads: list[ThreadInfo] = [ - ThreadInfo(thread_id=r[0], agent_name=r[1], updated_at=r[2]) + ThreadInfo( + thread_id=r[0], + agent_name=r[1], + updated_at=r[2], + latest_checkpoint_id=r[3], + ) for r in rows ] # Fetch message counts if requested if include_message_count and threads: - from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + await _populate_message_counts(conn, threads) - serde = JsonPlusSerializer() - for thread in threads: - thread["message_count"] = await _count_messages_from_checkpoint( - conn, thread["thread_id"], serde - ) + _cache_recent_threads(agent_name, limit, threads) + return threads + + +async def populate_thread_message_counts(threads: list[ThreadInfo]) -> list[ThreadInfo]: + """Populate `message_count` for an existing thread list. + + This is used by the `/threads` modal to render rows quickly, then backfill + counts in the background without issuing a second thread-list query. + + Args: + threads: Thread rows to enrich in place. + Returns: + The same list object with `message_count` values populated. + """ + if not threads: return threads + async with _connect() as conn: + await _populate_message_counts(conn, threads) + return threads + + +async def prewarm_thread_message_counts(limit: int | None = None) -> None: + """Prewarm thread message-count cache for faster `/threads` open. + + Fetches a bounded list of recent threads and populates counts into the + in-memory cache. Intended to run in a background worker during app startup. + + Args: + limit: Maximum threads to prewarm. Uses `get_thread_limit()` when `None`. + """ + thread_limit = limit if limit is not None else get_thread_limit() + if thread_limit < 1: + return + + try: + threads = await list_threads(limit=thread_limit, include_message_count=False) + if threads: + await populate_thread_message_counts(threads) + _cache_recent_threads(None, thread_limit, threads) + except (OSError, sqlite3.Error): + logger.debug("Could not prewarm thread message counts", exc_info=True) + except Exception: + logger.warning( + "Unexpected error while prewarming thread message counts", + exc_info=True, + ) + + +def get_cached_threads( + agent_name: str | None = None, + limit: int | None = None, +) -> list[ThreadInfo] | None: + """Get cached recent threads, if available. + + Args: + agent_name: Optional agent-name filter key. + limit: Maximum rows requested. Uses `get_thread_limit()` when `None`. + + Returns: + Copy of cached rows when available, otherwise `None`. + """ + + def _copy_with_cached_counts(rows: list[ThreadInfo]) -> list[ThreadInfo]: + copied_rows = _copy_threads(rows) + apply_cached_thread_message_counts(copied_rows) + return copied_rows + + thread_limit = limit if limit is not None else get_thread_limit() + if thread_limit < 1: + return None + + exact = _recent_threads_cache.get((agent_name, thread_limit)) + if exact is not None: + return _copy_with_cached_counts(exact) + + best_key: tuple[str | None, int] | None = None + for key in _recent_threads_cache: + cache_agent, cache_limit = key + if cache_agent != agent_name or cache_limit < thread_limit: + continue + if best_key is None or cache_limit < best_key[1]: + best_key = key + + if best_key is None: + return None + + return _copy_with_cached_counts(_recent_threads_cache[best_key][:thread_limit]) + + +def apply_cached_thread_message_counts(threads: list[ThreadInfo]) -> int: + """Apply cached message counts onto thread rows when freshness matches. + + Args: + threads: Thread rows to mutate in place. + + Returns: + Number of rows that were populated from cache. + """ + populated = 0 + for thread in threads: + if "message_count" in thread: + continue + thread_id = thread["thread_id"] + freshness = _thread_freshness(thread) + cached = _message_count_cache.get(thread_id) + if cached is None or cached[0] != freshness: + continue + thread["message_count"] = cached[1] + populated += 1 + return populated + + +async def _populate_message_counts( + conn: aiosqlite.Connection, + threads: list[ThreadInfo], +) -> None: + """Fill `message_count` on thread rows with cache-aware lookup.""" + serde = await _get_jsonplus_serializer() + for thread in threads: + thread_id = thread["thread_id"] + freshness = _thread_freshness(thread) + cached = _message_count_cache.get(thread_id) + if cached is not None and cached[0] == freshness: + thread["message_count"] = cached[1] + continue + + count = await _count_messages_from_checkpoint(conn, thread_id, serde) + thread["message_count"] = count + _cache_message_count(thread_id, freshness, count) + + +async def _get_jsonplus_serializer() -> JsonPlusSerializer: + """Return a cached JsonPlus serializer, loading it off the UI loop.""" + global _jsonplus_serializer # noqa: PLW0603 # Module-level cache requires global statement + if _jsonplus_serializer is not None: + return _jsonplus_serializer + + loop = asyncio.get_running_loop() + _jsonplus_serializer = await loop.run_in_executor(None, _create_jsonplus_serializer) + return _jsonplus_serializer + + +def _create_jsonplus_serializer() -> JsonPlusSerializer: + """Import and create a JsonPlus serializer. + + Returns: + A ready `JsonPlusSerializer` instance. + """ + from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + + return JsonPlusSerializer() + + +def _cache_message_count(thread_id: str, freshness: str | None, count: int) -> None: + """Cache a thread's message count with a freshness token.""" + if len(_message_count_cache) >= _MAX_MESSAGE_COUNT_CACHE and ( + thread_id not in _message_count_cache + ): + oldest = next(iter(_message_count_cache)) + _message_count_cache.pop(oldest, None) + _message_count_cache[thread_id] = (freshness, count) + + +def _thread_freshness(thread: ThreadInfo) -> str | None: + """Return a cache freshness token for a thread row.""" + return thread.get("latest_checkpoint_id") or thread.get("updated_at") + + +def _cache_recent_threads( + agent_name: str | None, + limit: int, + threads: list[ThreadInfo], +) -> None: + """Store a copy of recent thread rows for fast selector startup.""" + key = (agent_name, max(1, limit)) + if len(_recent_threads_cache) >= _MAX_RECENT_THREADS_CACHE_KEYS and ( + key not in _recent_threads_cache + ): + _recent_threads_cache.clear() + _recent_threads_cache[key] = _copy_threads(threads) + + +def _copy_threads(threads: list[ThreadInfo]) -> list[ThreadInfo]: + """Return shallow-copied thread rows.""" + return [ThreadInfo(**thread) for thread in threads] + async def _count_messages_from_checkpoint( conn: aiosqlite.Connection, @@ -366,6 +564,11 @@ async def delete_thread(thread_id: str) -> bool: if await _table_exists(conn, "writes"): await conn.execute("DELETE FROM writes WHERE thread_id = ?", (thread_id,)) await conn.commit() + if deleted: + _message_count_cache.pop(thread_id, None) + for key, rows in list(_recent_threads_cache.items()): + filtered = [row for row in rows if row["thread_id"] != thread_id] + _recent_threads_cache[key] = filtered return deleted diff --git a/libs/cli/deepagents_cli/widgets/thread_selector.py b/libs/cli/deepagents_cli/widgets/thread_selector.py index 3da41a6b45f..dc92b8184aa 100644 --- a/libs/cli/deepagents_cli/widgets/thread_selector.py +++ b/libs/cli/deepagents_cli/widgets/thread_selector.py @@ -20,14 +20,13 @@ from textual.app import ComposeResult from textual.events import Click - from deepagents_cli.sessions import ThreadInfo - from deepagents_cli.config import ( CharsetMode, _detect_charset_mode, build_langsmith_thread_url, get_glyphs, ) +from deepagents_cli.sessions import ThreadInfo from deepagents_cli.widgets._links import open_style_link logger = logging.getLogger(__name__) @@ -179,17 +178,40 @@ class ThreadSelectorScreen(ModalScreen[str | None]): } """ - def __init__(self, current_thread: str | None = None) -> None: + def __init__( + self, + current_thread: str | None = None, + *, + thread_limit: int | None = None, + initial_threads: list[ThreadInfo] | None = None, + ) -> None: """Initialize the `ThreadSelectorScreen`. Args: current_thread: The currently active thread ID (to highlight). + thread_limit: Maximum number of rows to fetch when querying DB. + initial_threads: Optional preloaded rows to render immediately. """ super().__init__() self._current_thread = current_thread - self._threads: list[ThreadInfo] = [] + self._thread_limit = thread_limit + self._threads: list[ThreadInfo] = ( + [ThreadInfo(**thread) for thread in initial_threads] + if initial_threads is not None + else [] + ) + self._has_initial_threads = initial_threads is not None self._selected_index = 0 self._option_widgets: list[ThreadOption] = [] + self._sync_selected_index() + + def _sync_selected_index(self) -> None: + """Select the current thread when it exists in the loaded rows.""" + self._selected_index = 0 + for i, thread in enumerate(self._threads): + if thread["thread_id"] == self._current_thread: + self._selected_index = i + break def _build_title(self, thread_url: str | None = None) -> str | Text: """Build the title, optionally with a clickable thread ID link. @@ -226,11 +248,21 @@ def compose(self) -> ComposeResult: yield Static(self._format_header(), classes="thread-list-header") with VerticalScroll(classes="thread-list"): - yield Static( - "[dim]Loading threads...[/dim]", - classes="thread-empty", - id="thread-loading", - ) + if self._has_initial_threads: + if self._threads: + self._option_widgets, _ = self._create_option_widgets() + yield from self._option_widgets + else: + yield Static( + "[dim]No threads found[/dim]", + classes="thread-empty", + ) + else: + yield Static( + "[dim]Loading threads...[/dim]", + classes="thread-empty", + id="thread-loading", + ) help_text = ( f"{glyphs.arrow_up}/{glyphs.arrow_down}/tab navigate " @@ -244,12 +276,49 @@ async def on_mount(self) -> None: container = self.query_one(Vertical) container.styles.border = ("ascii", "green") - from deepagents_cli.sessions import get_thread_limit, list_threads + self.focus() + if self._has_initial_threads: + self.call_after_refresh(self._scroll_selected_into_view) + self._schedule_message_count_load() + if self._current_thread: + self._resolve_thread_url() + # Cached rows are only a startup snapshot; refresh from SQLite. + self.run_worker( + self._load_threads, exclusive=True, group="thread-selector-load" + ) + return - try: - self._threads = await list_threads( - limit=get_thread_limit(), include_message_count=True + # Defer DB work to a worker so modal paints immediately. + self.run_worker( + self._load_threads, exclusive=True, group="thread-selector-load" + ) + + def _schedule_message_count_load(self) -> None: + """Schedule background message-count loading when counts are missing.""" + has_missing_counts = self._threads and any( + "message_count" not in thread for thread in self._threads + ) + if has_missing_counts: + self.run_worker( + self._load_message_counts, + exclusive=True, + group="thread-selector-counts", ) + + async def _load_threads(self) -> None: + """Load thread rows first, then kick off background message counts.""" + from deepagents_cli.sessions import ( + apply_cached_thread_message_counts, + list_threads, + ) + + try: + limit = self._thread_limit + if limit is None: + from deepagents_cli.sessions import get_thread_limit + + limit = get_thread_limit() + self._threads = await list_threads(limit=limit, include_message_count=False) except (OSError, sqlite3.Error) as exc: logger.exception("Failed to load threads for thread selector") await self._show_mount_error(str(exc)) @@ -259,24 +328,68 @@ async def on_mount(self) -> None: await self._show_mount_error(str(exc)) return - for i, t in enumerate(self._threads): - if t["thread_id"] == self._current_thread: - self._selected_index = i - break + self._sync_selected_index() + + # Reuse startup-prewarmed counts before first list paint. + apply_cached_thread_message_counts(self._threads) await self._build_list() + # Populate message counts after first paint. + self._schedule_message_count_load() + if self._current_thread: self._resolve_thread_url() - self.focus() + async def _load_message_counts(self) -> None: + """Populate thread message counts in background and refresh labels.""" + from deepagents_cli.sessions import populate_thread_message_counts + + if not self._threads: + return + + try: + await populate_thread_message_counts(self._threads) + except (OSError, sqlite3.Error): + logger.debug( + "Could not load message counts for thread selector", + exc_info=True, + ) + return + except Exception: + logger.warning( + "Unexpected error loading message counts for thread selector", + exc_info=True, + ) + return + + self._refresh_message_count_labels() + + def _refresh_message_count_labels(self) -> None: + """Refresh only row labels after background message counts complete.""" + if not self._threads or not self._option_widgets: + return + + for index, thread in enumerate(self._threads): + if index >= len(self._option_widgets): + break + widget = self._option_widgets[index] + widget.update( + self._format_option_label( + thread, + selected=index == self._selected_index, + current=thread["thread_id"] == self._current_thread, + ) + ) def _resolve_thread_url(self) -> None: """Start exclusive background worker to resolve LangSmith thread URL. `exclusive=True` so repeated calls cancel any in-flight resolution. """ - self.run_worker(self._fetch_thread_url, exclusive=True) + self.run_worker( + self._fetch_thread_url, exclusive=True, group="thread-selector-url" + ) async def _fetch_thread_url(self) -> None: """Resolve the LangSmith URL and update the title with a clickable link. @@ -341,9 +454,9 @@ async def _build_list(self) -> None: """Build the thread option widgets.""" scroll = self.query_one(".thread-list", VerticalScroll) await scroll.remove_children() - self._option_widgets = [] if not self._threads: + self._option_widgets = [] await scroll.mount( Static( "[dim]No threads found[/dim]", @@ -352,6 +465,19 @@ async def _build_list(self) -> None: ) return + self._option_widgets, selected_widget = self._create_option_widgets() + await scroll.mount(*self._option_widgets) + + if selected_widget: + self._scroll_selected_into_view() + + def _create_option_widgets(self) -> tuple[list[ThreadOption], ThreadOption | None]: + """Build option widgets from loaded threads without mounting. + + Returns: + Tuple of all option widgets and the currently selected widget. + """ + widgets: list[ThreadOption] = [] selected_widget: ThreadOption | None = None for i, thread in enumerate(self._threads): @@ -373,18 +499,27 @@ async def _build_list(self) -> None: index=i, classes=classes, ) - self._option_widgets.append(widget) - + widgets.append(widget) if is_selected: selected_widget = widget - await scroll.mount(*self._option_widgets) + return widgets, selected_widget - if selected_widget: - if self._selected_index == 0: - scroll.scroll_home(animate=False) - else: - selected_widget.scroll_visible(animate=False) + def _scroll_selected_into_view(self) -> None: + """Scroll selected option into view without animation.""" + if not self._option_widgets: + return + if self._selected_index >= len(self._option_widgets): + return + try: + scroll = self.query_one(".thread-list", VerticalScroll) + except NoMatches: + return + + if self._selected_index == 0: + scroll.scroll_home(animate=False) + else: + self._option_widgets[self._selected_index].scroll_visible(animate=False) @staticmethod def _format_header() -> str: @@ -421,7 +556,8 @@ def _format_option_label( cursor = f"{glyphs.cursor} " if selected else " " tid = thread["thread_id"][:_COL_TID] agent = (thread.get("agent_name") or "unknown")[:_COL_AGENT] - msgs = str(thread.get("message_count", 0)) + raw_count = thread.get("message_count") + msgs = str(raw_count) if raw_count is not None else "..." timestamp = format_timestamp(thread.get("updated_at")) label = ( diff --git a/libs/cli/tests/unit_tests/test_app.py b/libs/cli/tests/unit_tests/test_app.py index 3c4f8f66d35..0ddaeafbbab 100644 --- a/libs/cli/tests/unit_tests/test_app.py +++ b/libs/cli/tests/unit_tests/test_app.py @@ -78,6 +78,62 @@ async def test_app_css_validates_on_mount(self) -> None: assert app.is_running +class TestThreadCachePrewarm: + """Tests for startup thread-cache prewarming.""" + + @pytest.mark.asyncio + async def test_prewarm_uses_current_thread_limit(self) -> None: + """Prewarm helper should pass the resolved thread limit through.""" + app = DeepAgentsApp(agent=MagicMock(), thread_id="thread-123") + + with ( + patch("deepagents_cli.sessions.get_thread_limit", return_value=7), + patch( + "deepagents_cli.sessions.prewarm_thread_message_counts", + new_callable=AsyncMock, + ) as mock_prewarm, + ): + await app._prewarm_threads_cache() + + mock_prewarm.assert_awaited_once_with(limit=7) + + @pytest.mark.asyncio + async def test_show_thread_selector_uses_cached_rows(self) -> None: + """Thread selector should receive prefetched rows when available.""" + cached_threads = [ + { + "thread_id": "thread-abc", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + "message_count": 2, + } + ] + app = DeepAgentsApp() + + async with app.run_test() as pilot: + await pilot.pause() + with ( + patch("deepagents_cli.sessions.get_thread_limit", return_value=9), + patch( + "deepagents_cli.sessions.get_cached_threads", + return_value=cached_threads, + ), + patch("deepagents_cli.app.ThreadSelectorScreen") as mock_screen_cls, + patch.object(app, "push_screen") as mock_push_screen, + ): + mock_screen = MagicMock() + mock_screen_cls.return_value = mock_screen + await app._show_thread_selector() + + assert app._session_state is not None + mock_screen_cls.assert_called_once_with( + current_thread=app._session_state.thread_id, + thread_limit=9, + initial_threads=cached_threads, + ) + mock_push_screen.assert_called_once() + + class TestAppBindings: """Test app keybindings.""" diff --git a/libs/cli/tests/unit_tests/test_sessions.py b/libs/cli/tests/unit_tests/test_sessions.py index 32e54022ed5..336f63fe2e9 100644 --- a/libs/cli/tests/unit_tests/test_sessions.py +++ b/libs/cli/tests/unit_tests/test_sessions.py @@ -5,7 +5,7 @@ import sqlite3 from datetime import UTC, datetime from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer @@ -402,6 +402,249 @@ def test_no_message_count_by_default(self, temp_db_with_messages: Path) -> None: assert len(threads) == 1 assert "message_count" not in threads[0] + def test_message_count_uses_cache_for_unchanged_thread( + self, temp_db_with_messages: Path + ) -> None: + """Second call should reuse cached count for unchanged checkpoint.""" + sessions._message_count_cache.clear() + try: + with ( + patch.object( + sessions, "get_db_path", return_value=temp_db_with_messages + ), + patch.object( + sessions, + "_get_jsonplus_serializer", + new_callable=AsyncMock, + return_value=object(), + ), + patch.object( + sessions, + "_count_messages_from_checkpoint", + new_callable=AsyncMock, + return_value=3, + ) as mock_count, + ): + first = asyncio.run(sessions.list_threads(include_message_count=True)) + second = asyncio.run(sessions.list_threads(include_message_count=True)) + + assert first[0]["message_count"] == 3 + assert second[0]["message_count"] == 3 + assert mock_count.await_count == 1 + finally: + sessions._message_count_cache.clear() + + def test_message_count_cache_invalidates_on_new_checkpoint( + self, temp_db_with_messages: Path + ) -> None: + """A newer checkpoint should invalidate cached message count.""" + sessions._message_count_cache.clear() + try: + with ( + patch.object( + sessions, "get_db_path", return_value=temp_db_with_messages + ), + patch.object( + sessions, + "_get_jsonplus_serializer", + new_callable=AsyncMock, + return_value=object(), + ), + patch.object( + sessions, + "_count_messages_from_checkpoint", + new_callable=AsyncMock, + side_effect=[3, 4], + ) as mock_count, + ): + first = asyncio.run(sessions.list_threads(include_message_count=True)) + assert first[0]["message_count"] == 3 + + conn = sqlite3.connect(str(temp_db_with_messages)) + type_str, checkpoint_blob, metadata = conn.execute( + "SELECT type, checkpoint, metadata FROM checkpoints " + "WHERE thread_id = ? AND checkpoint_id = ?", + ("thread1", "cp_1"), + ).fetchone() + conn.execute( + "INSERT INTO checkpoints " + "(thread_id, checkpoint_ns, checkpoint_id, type, checkpoint, " + "metadata) " + "VALUES (?, '', ?, ?, ?, ?)", + ("thread1", "cp_2", type_str, checkpoint_blob, metadata), + ) + conn.commit() + conn.close() + + second = asyncio.run(sessions.list_threads(include_message_count=True)) + assert second[0]["message_count"] == 4 + assert mock_count.await_count == 2 + finally: + sessions._message_count_cache.clear() + + +class TestApplyCachedThreadMessageCounts: + """Tests for applying cached thread counts to rows.""" + + def test_populates_rows_from_cache(self) -> None: + """Rows with matching freshness should get counts from cache.""" + sessions._message_count_cache.clear() + try: + sessions._message_count_cache["thread-a"] = ("cp_1", 7) + threads: list[sessions.ThreadInfo] = [ + { + "thread_id": "thread-a", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + "latest_checkpoint_id": "cp_1", + }, + { + "thread_id": "thread-b", + "agent_name": "agent2", + "updated_at": "2024-01-01T00:00:00+00:00", + "latest_checkpoint_id": "cp_1", + }, + ] + + populated = sessions.apply_cached_thread_message_counts(threads) + + assert populated == 1 + assert threads[0]["message_count"] == 7 + assert "message_count" not in threads[1] + finally: + sessions._message_count_cache.clear() + + def test_skips_stale_cache_entries(self) -> None: + """Rows should not use cache when freshness token changes.""" + sessions._message_count_cache.clear() + try: + sessions._message_count_cache["thread-a"] = ("cp_1", 7) + threads: list[sessions.ThreadInfo] = [ + { + "thread_id": "thread-a", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + "latest_checkpoint_id": "cp_2", + } + ] + + populated = sessions.apply_cached_thread_message_counts(threads) + + assert populated == 0 + assert "message_count" not in threads[0] + finally: + sessions._message_count_cache.clear() + + +class TestGetCachedThreads: + """Tests for cached thread snapshot retrieval.""" + + def test_returns_exact_cached_limit(self) -> None: + """Exact cache key should return copied rows.""" + sessions._recent_threads_cache.clear() + try: + sessions._recent_threads_cache[None, 5] = [ + { + "thread_id": "thread-a", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + "message_count": 3, + } + ] + rows = sessions.get_cached_threads(limit=5) + assert rows is not None + assert len(rows) == 1 + assert rows[0]["thread_id"] == "thread-a" + rows[0]["thread_id"] = "mutated" + assert sessions._recent_threads_cache[None, 5][0]["thread_id"] == "thread-a" + finally: + sessions._recent_threads_cache.clear() + + def test_uses_larger_cached_limit(self) -> None: + """Larger cached window should satisfy smaller requested limit.""" + sessions._recent_threads_cache.clear() + try: + sessions._recent_threads_cache[None, 20] = [ + { + "thread_id": "thread-1", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + }, + { + "thread_id": "thread-2", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + }, + ] + rows = sessions.get_cached_threads(limit=1) + assert rows is not None + assert len(rows) == 1 + assert rows[0]["thread_id"] == "thread-1" + finally: + sessions._recent_threads_cache.clear() + + def test_applies_cached_message_counts_to_snapshot(self) -> None: + """Returned snapshot should hydrate counts from message-count cache.""" + sessions._recent_threads_cache.clear() + sessions._message_count_cache.clear() + try: + sessions._recent_threads_cache[None, 5] = [ + { + "thread_id": "thread-a", + "agent_name": "agent1", + "updated_at": "2024-01-01T00:00:00+00:00", + "latest_checkpoint_id": "cp_1", + } + ] + sessions._message_count_cache["thread-a"] = ("cp_1", 9) + + rows = sessions.get_cached_threads(limit=5) + + assert rows is not None + assert rows[0]["message_count"] == 9 + assert "message_count" not in sessions._recent_threads_cache[None, 5][0] + finally: + sessions._recent_threads_cache.clear() + sessions._message_count_cache.clear() + + +class TestPrewarmThreadMessageCounts: + """Tests for prewarm_thread_message_counts error handling.""" + + @pytest.mark.asyncio + async def test_unexpected_errors_log_warning(self) -> None: + """Unexpected prewarm failures should be visible at warning level.""" + with ( + patch( + "deepagents_cli.sessions.list_threads", + new_callable=AsyncMock, + side_effect=RuntimeError("unexpected type mismatch"), + ), + patch.object(sessions.logger, "warning") as mock_warning, + ): + await sessions.prewarm_thread_message_counts(limit=3) + + mock_warning.assert_called_once() + + +class TestCacheMessageCount: + """Tests for message-count cache eviction behavior.""" + + def test_overflow_evicts_oldest_entry_only(self) -> None: + """Cache overflow should evict only the oldest key, not clear all keys.""" + sessions._message_count_cache.clear() + try: + with patch.object(sessions, "_MAX_MESSAGE_COUNT_CACHE", 2): + sessions._cache_message_count("thread-1", "cp_1", 1) + sessions._cache_message_count("thread-2", "cp_2", 2) + sessions._cache_message_count("thread-3", "cp_3", 3) + + assert "thread-1" not in sessions._message_count_cache + assert sessions._message_count_cache["thread-2"] == ("cp_2", 2) + assert sessions._message_count_cache["thread-3"] == ("cp_3", 3) + finally: + sessions._message_count_cache.clear() + class TestMessageCountFromCheckpointBlob: """Tests for counting messages from checkpoint blob (not writes table). diff --git a/libs/cli/tests/unit_tests/test_thread_selector.py b/libs/cli/tests/unit_tests/test_thread_selector.py index 710b7824c2e..249cf60e241 100644 --- a/libs/cli/tests/unit_tests/test_thread_selector.py +++ b/libs/cli/tests/unit_tests/test_thread_selector.py @@ -1,5 +1,6 @@ """Tests for ThreadSelectorScreen.""" +import asyncio from typing import Any, ClassVar from unittest.mock import AsyncMock, MagicMock, patch @@ -8,6 +9,7 @@ from textual.app import App, ComposeResult from textual.binding import Binding, BindingType from textual.containers import Container, Vertical +from textual.css.query import NoMatches from textual.screen import ModalScreen from textual.widgets import Static @@ -569,6 +571,18 @@ def test_includes_message_count(self) -> None: ) assert "5" in label + def test_missing_message_count_shows_placeholder(self) -> None: + """Rows without loaded counts should show an explicit placeholder.""" + thread = ThreadInfo( + thread_id="abc12345", + agent_name="my-agent", + updated_at="2025-01-15T10:30:00", + ) + label = ThreadSelectorScreen._format_option_label( + thread, selected=False, current=False + ) + assert "..." in label + def test_columns_align_with_header(self) -> None: """Option labels should align with the column header.""" header = ThreadSelectorScreen._format_header() @@ -858,7 +872,247 @@ async def test_custom_limit_is_forwarded(self) -> None: app.show_selector() await pilot.pause() - mock_lt.assert_awaited_once_with(limit=5, include_message_count=True) + mock_lt.assert_awaited_once_with(limit=5, include_message_count=False) + + @pytest.mark.asyncio + async def test_message_counts_are_loaded_in_background(self) -> None: + """Missing counts should be populated asynchronously after list render.""" + threads_without_counts: list[ThreadInfo] = [ + { + "thread_id": "abc12345", + "agent_name": "my-agent", + "updated_at": "2025-01-15T10:30:00", + } + ] + + async def _populate(threads: list[ThreadInfo]) -> list[ThreadInfo]: + await asyncio.sleep(0) + for thread in threads: + thread["message_count"] = 9 + return threads + + with ( + patch( + "deepagents_cli.sessions.list_threads", + new_callable=AsyncMock, + return_value=threads_without_counts, + ) as mock_lt, + patch( + "deepagents_cli.sessions.populate_thread_message_counts", + new_callable=AsyncMock, + side_effect=_populate, + ) as mock_populate, + ): + app = ThreadSelectorTestApp() + async with app.run_test() as pilot: + app.show_selector() + await pilot.pause() + + for _ in range(10): + if mock_populate.await_count >= 1: + break + await pilot.pause(0.05) + + mock_lt.assert_awaited_once_with(limit=20, include_message_count=False) + mock_populate.assert_awaited_once() + + screen = app.screen + assert isinstance(screen, ThreadSelectorScreen) + assert screen._threads[0]["message_count"] == 9 + + @pytest.mark.asyncio + async def test_cached_counts_skip_background_population(self) -> None: + """If cache fills counts before paint, background populate is skipped.""" + threads_without_counts: list[ThreadInfo] = [ + { + "thread_id": "abc12345", + "agent_name": "my-agent", + "updated_at": "2025-01-15T10:30:00", + "latest_checkpoint_id": "cp_1", + } + ] + + def _apply_cached(threads: list[ThreadInfo]) -> int: + threads[0]["message_count"] = 11 + return 1 + + with ( + patch( + "deepagents_cli.sessions.list_threads", + new_callable=AsyncMock, + return_value=threads_without_counts, + ), + patch( + "deepagents_cli.sessions.apply_cached_thread_message_counts", + side_effect=_apply_cached, + ) as mock_apply_cached, + patch( + "deepagents_cli.sessions.populate_thread_message_counts", + new_callable=AsyncMock, + ) as mock_populate, + ): + app = ThreadSelectorTestApp() + async with app.run_test() as pilot: + app.show_selector() + await pilot.pause() + await pilot.pause(0.1) + + mock_apply_cached.assert_called_once() + mock_populate.assert_not_awaited() + + screen = app.screen + assert isinstance(screen, ThreadSelectorScreen) + assert screen._threads[0]["message_count"] == 11 + + +class TestThreadSelectorMessageCountErrors: + """Tests for thread selector message-count load error handling.""" + + @pytest.mark.asyncio + async def test_unexpected_message_count_error_logs_warning(self) -> None: + """Unexpected count-load errors should be visible at warning level.""" + screen = ThreadSelectorScreen( + initial_threads=[ + { + "thread_id": "abc12345", + "agent_name": "my-agent", + "updated_at": "2025-01-15T10:30:00", + } + ] + ) + + with ( + patch( + "deepagents_cli.sessions.populate_thread_message_counts", + new_callable=AsyncMock, + side_effect=RuntimeError("unexpected type mismatch"), + ), + patch( + "deepagents_cli.widgets.thread_selector.logger.warning" + ) as mock_warning, + ): + await screen._load_message_counts() + + mock_warning.assert_called_once() + + +class TestThreadSelectorPrefetchedRows: + """Tests for rendering with prefetched rows from startup cache.""" + + @pytest.mark.asyncio + async def test_prefetched_rows_render_without_loading_state(self) -> None: + """Prefetched rows should render immediately, then refresh from SQLite.""" + prefetched: list[ThreadInfo] = [ + { + "thread_id": "abc12345", + "agent_name": "my-agent", + "updated_at": "2025-01-15T10:30:00", + "message_count": 5, + } + ] + refreshed: list[ThreadInfo] = [ + { + "thread_id": "new12345", + "agent_name": "my-agent", + "updated_at": "2025-01-16T12:00:00", + "message_count": 6, + }, + { + "thread_id": "abc12345", + "agent_name": "my-agent", + "updated_at": "2025-01-15T10:30:00", + "message_count": 5, + }, + ] + app = ThreadSelectorTestApp(current_thread="abc12345") + + # Use an Event gate so the mock cannot resolve until we allow it, + # avoiding race conditions across Python versions (3.13 in particular). + gate = asyncio.Event() + + async def _list_threads(*_args: object, **_kwargs: object) -> list[ThreadInfo]: + await gate.wait() + return refreshed + + with patch( + "deepagents_cli.sessions.list_threads", + new_callable=AsyncMock, + side_effect=_list_threads, + ) as mock_list_threads: + async with app.run_test() as pilot: + app.push_screen( + ThreadSelectorScreen( + current_thread="abc12345", + thread_limit=20, + initial_threads=prefetched, + ) + ) + await pilot.pause() + + screen = app.screen + assert isinstance(screen, ThreadSelectorScreen) + assert len(screen._option_widgets) == 1 + with pytest.raises(NoMatches): + screen.query_one("#thread-loading", Static) + + # Release the mock so the background refresh can complete. + gate.set() + + for _ in range(10): + if mock_list_threads.await_count >= 1 and len(screen._threads) == 2: + break + await pilot.pause(0.05) + + mock_list_threads.assert_awaited_once_with( + limit=20, + include_message_count=False, + ) + assert len(screen._threads) == 2 + assert screen._threads[0]["thread_id"] == "new12345" + + @pytest.mark.asyncio + async def test_empty_prefetched_snapshot_still_refreshes(self) -> None: + """An empty cached snapshot should still hydrate from SQLite in background.""" + refreshed: list[ThreadInfo] = [ + { + "thread_id": "new12345", + "agent_name": "my-agent", + "updated_at": "2025-01-16T12:00:00", + "message_count": 6, + } + ] + app = ThreadSelectorTestApp(current_thread="abc12345") + with patch( + "deepagents_cli.sessions.list_threads", + new_callable=AsyncMock, + return_value=refreshed, + ) as mock_list_threads: + async with app.run_test() as pilot: + app.push_screen( + ThreadSelectorScreen( + current_thread="abc12345", + thread_limit=20, + initial_threads=[], + ) + ) + await pilot.pause() + + screen = app.screen + assert isinstance(screen, ThreadSelectorScreen) + with pytest.raises(NoMatches): + screen.query_one("#thread-loading", Static) + + for _ in range(10): + if mock_list_threads.await_count >= 1 and len(screen._threads) == 1: + break + await pilot.pause(0.05) + + mock_list_threads.assert_awaited_once_with( + limit=20, + include_message_count=False, + ) + assert len(screen._threads) == 1 + assert screen._threads[0]["thread_id"] == "new12345" def _get_widget_text(widget: Static) -> str: