fix: unify bubble KV loaders and converge DisplayBubble shapes

Route all bubbleId:* loads through workspace_db._parse_bubble_kv_row and load_bubble_map (including workspace tabs). Introduce DisplayBubble/BubbleMetadata and utils/display_bubble builders shared by tabs, CLI, and IDE markdown export.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: bradjin8

models/__init__.py

@@ -1,3 +1,4 @@
+from models.bubble_display import BubbleMetadata, BubbleRole, DisplayBubble
 from models.cli_session import CliSessionMeta
 from models.conversation import Bubble, Composer, Conversation, WorkspaceLocalComposer
 from models.errors import SchemaError
@@ -8,7 +9,10 @@
 
 __all__ = [
     "Bubble",
+    "BubbleMetadata",
+    "BubbleRole",
     "CliSessionMeta",
+    "DisplayBubble",
     "Composer",
     "Conversation",
     "ConversationSummary",

models/bubble_display.py

@@ -0,0 +1,47 @@
+"""Rendered bubble shapes for UI, CLI, and markdown export.
+
+Storage/KV rows (``bubbleId:*`` in ``cursorDiskKV``) are validated as
+:class:`models.conversation.Bubble` at load time via
+:func:`services.workspace_db.load_bubble_map`.  All user-facing paths emit
+:class:`DisplayBubble` with optional :class:`BubbleMetadata`.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal, TypedDict
+
+BubbleRole = Literal["user", "ai"]
+
+
+class BubbleMetadata(TypedDict, total=False):
+    """Nested fields on a :class:`DisplayBubble` (tabs, CLI, export)."""
+
+    modelName: str
+    inputTokens: int
+    outputTokens: int
+    cachedTokens: int
+    toolResultsCount: int
+    toolResults: list[Any]
+    toolCalls: list[dict[str, Any]]
+    thinking: str
+    thinkingDurationMs: int | float
+    contextWindowPercent: float
+    contextTokensUsed: int
+    contextTokenLimit: int
+    contextPctRemaining: float
+    responseTimeMs: int
+    cost: float
+
+
+class _DisplayBubbleRequired(TypedDict):
+    type: BubbleRole
+    text: str
+    timestamp: int
+
+
+class _DisplayBubbleOptional(TypedDict, total=False):
+    metadata: BubbleMetadata
+
+
+class DisplayBubble(_DisplayBubbleRequired, _DisplayBubbleOptional):
+    """One message bubble in the browser UI or an exported Markdown document."""

models/conversation.py

@@ -177,7 +177,10 @@ def from_dict(cls, raw: dict[str, Any]) -> "WorkspaceLocalComposer":
 
 @dataclass(frozen=True)
 class Bubble:
-    """One message in a composer; bubble_id comes from the row key, not the JSON value."""
+    """One message in a composer; bubble_id comes from the row key, not the JSON value.
+
+    Rendered for UI/export as :class:`models.bubble_display.DisplayBubble`.
+    """
 
     bubble_id: str
     raw: dict[str, Any] = field(default_factory=dict)

scripts/export.py

@@ -52,7 +52,7 @@
     cursor_cli_session_to_markdown,
     cursor_ide_chat_to_markdown,
 )
-from models import ExportEntry, SchemaError  # noqa: E402
+from models import Bubble, ExportEntry, SchemaError  # noqa: E402
 from services.workspace_context import (  # noqa: E402
     enrich_workspace_context_from_global_db,
     resolve_workspace_context,
@@ -221,7 +221,7 @@ def main():
 
     # ── Database reading via service layer ────────────────────────────────────
     project_layouts_map: dict = {}
-    bubble_map: dict = {}
+    bubble_map: dict[str, Bubble] = {}
     code_block_diff_map: dict = {}
     ide_composer_rows: list = []
     invalid_workspace_aliases: dict = {}

services/search.py

@@ -32,10 +32,11 @@
     "search_global_storage",
     "search_legacy_workspaces",
 ]
-from models import Bubble, Composer, ParseWarningCollector, SchemaError, SearchResult
+from models import Composer, ParseWarningCollector, SchemaError, SearchResult
 from services.workspace_db import (
     build_composer_id_to_workspace_id_cached,
     collect_workspace_entries,
+    load_bubble_map,
     open_global_db,
 )
 from utils.cli_chat_reader import list_cli_projects, messages_to_bubbles, traverse_blobs
@@ -153,38 +154,6 @@ def _build_ws_id_to_name(
     return mapping
 
 
-def _build_search_bubble_map(
-    global_db: sqlite3.Connection,
-    parse_warnings: ParseWarningCollector,
-) -> dict[str, dict[str, Any]]:
-    """Load ``bubbleId:*`` rows from an open global DB connection.
-
-    Returns ``{bubble_id: {"text": str, "raw": dict}}``.  Rows that fail
-    schema validation or JSON decoding are skipped; the skip is recorded in
-    *parse_warnings*.
-    """
-    bubble_map: dict[str, dict[str, Any]] = {}
-    for row in global_db.execute(
-        "SELECT key, value FROM cursorDiskKV WHERE key LIKE 'bubbleId:%'"
-    ):
-        parts = row["key"].split(":")
-        if len(parts) < 3:
-            continue
-        bid = parts[2]
-        try:
-            bubble = Bubble.from_dict(json.loads(row["value"]), bubble_id=bid)
-            bubble_map[bid] = {"text": extract_text_from_bubble(bubble), "raw": bubble.raw}
-        except SchemaError as exc:
-            _logger.warning(
-                "Schema drift in bubble %s: %s (%s)", bid, exc, type(exc).__name__
-            )
-            parse_warnings.record_bubble_skipped()
-        except (json.JSONDecodeError, TypeError, ValueError) as exc:
-            _logger.warning("Failed to decode Bubble from bubbleId:%s: %s", bid, exc)
-            parse_warnings.record_bubble_skipped()
-    return bubble_map
-
-
 # ---------------------------------------------------------------------------
 # Public: per-source search functions
 # ---------------------------------------------------------------------------
@@ -223,7 +192,7 @@ def search_global_storage(
         with open_global_db(workspace_path) as (conn, _db_path):
             if conn is None:
                 return results
-            bubble_map = _build_search_bubble_map(conn, parse_warnings)
+            bubble_map = load_bubble_map(conn, parse_warnings=parse_warnings)
             composer_rows = conn.execute(
                 "SELECT key, value FROM cursorDiskKV"
                 " WHERE key LIKE 'composerData:%' AND LENGTH(value) > 10"
@@ -276,15 +245,13 @@ def search_global_storage(
                     bid = header.get("bubbleId")
                     if not bid:
                         continue
-                    entry = bubble_map.get(bid)
-                    if not entry:
+                    bubble = bubble_map.get(bid)
+                    if bubble is None:
                         continue
-                    text = entry.get("text") or ""
+                    text = extract_text_from_bubble(bubble)
                     if text:
                         bubble_texts.append(text)
-                    raw_bubble = entry.get("raw")
-                    if raw_bubble:
-                        bubble_meta.append(_json_dump_safe(raw_bubble))
+                    bubble_meta.append(_json_dump_safe(bubble.raw))
 
                 exclusion_text = _build_exclusion_searchable(
                     project_name=project_name,

services/workspace_context.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass, replace
 from typing import Any
 
+from models import Bubble
 from services.workspace_db import (
     build_composer_id_to_workspace_id,
     build_composer_id_to_workspace_id_cached,
@@ -30,7 +31,7 @@ class WorkspaceContext:
     project_name_to_workspace_id: dict[str, str]
     workspace_path_to_id: dict[str, str]
     project_layouts_map: dict[str, list[str]]
-    bubble_map: dict[str, dict[str, Any]]
+    bubble_map: dict[str, Bubble]
 
 
 def _entries(

services/workspace_db.py

@@ -11,6 +11,7 @@
 
 _logger = logging.getLogger(__name__)
 
+from models import Bubble, ParseWarningCollector, SchemaError
 from utils.path_helpers import get_workspace_folder_paths
 from utils.workspace_descriptor import read_json_file
 
@@ -34,30 +35,63 @@ def safe_fetchall(
         return []
 
 
-def load_bubble_map(global_db: sqlite3.Connection) -> dict[str, dict[str, Any]]:
-    """Load all ``bubbleId:*`` KV entries into ``{bubble_id: bubble_dict}``.
+def _parse_bubble_kv_row(
+    row_key: str,
+    row_value: str | bytes,
+    *,
+    parse_warnings: ParseWarningCollector | None = None,
+) -> tuple[str, Bubble] | None:
+    """Parse one ``bubbleId:…`` row; return ``(bubble_id, Bubble)`` or skip."""
+    parts = row_key.split(":")
+    if len(parts) < 3:
+        return None
+    bid = parts[2]
+    try:
+        parsed = json.loads(row_value)
+        bubble = Bubble.from_dict(parsed, bubble_id=bid)
+        return bid, bubble
+    except SchemaError as exc:
+        _logger.warning(
+            "Schema drift in bubble %s: %s (%s)", bid, exc, type(exc).__name__
+        )
+        if parse_warnings is not None:
+            parse_warnings.record_bubble_skipped()
+    except (json.JSONDecodeError, TypeError, ValueError) as exc:
+        if parse_warnings is not None:
+            _logger.warning(
+                "Failed to decode Bubble from %s: %s", row_key, exc
+            )
+            parse_warnings.record_bubble_skipped()
+        else:
+            _logger.debug("Skipping malformed bubbleId row %s: %s", row_key, exc)
+    return None
 
-    Skips rows whose JSON value is not a dict; JSON parse errors are logged at
-    DEBUG level so a single malformed row cannot block the rest.
+
+def load_bubble_map(
+    global_db: sqlite3.Connection,
+    *,
+    parse_warnings: ParseWarningCollector | None = None,
+) -> dict[str, Bubble]:
+    """Load all ``bubbleId:*`` KV entries into ``{bubble_id: Bubble}``.
+
+    Uses the same :meth:`Bubble.from_dict` validation as search and tabs.
+    When *parse_warnings* is set, skipped rows are recorded for the API.
     """
-    bubble_map: dict[str, dict[str, Any]] = {}
+    bubble_map: dict[str, Bubble] = {}
     try:
         rows = global_db.execute(
-            "SELECT key, value FROM cursorDiskKV WHERE key LIKE 'bubbleId:%'"
+            "SELECT key, value FROM cursorDiskKV"
+            " WHERE key LIKE 'bubbleId:%' AND value IS NOT NULL"
         ).fetchall()
     except sqlite3.Error:
         return bubble_map
     for row in rows:
-        parts = row["key"].split(":")
-        if len(parts) < 3:
-            continue
-        bid = parts[2]
-        try:
-            b = json.loads(row["value"])
-            if isinstance(b, dict):
-                bubble_map[bid] = b
-        except (json.JSONDecodeError, ValueError, KeyError, TypeError) as e:
-            _logger.debug("Skipping malformed bubbleId row %s: %s", row["key"], e)
+        parsed = _parse_bubble_kv_row(
+            row["key"], row["value"], parse_warnings=parse_warnings
+        )
+        if parsed is not None:
+            bid, bubble = parsed
+            bubble_map[bid] = bubble
     return bubble_map
 
 
@@ -163,14 +197,17 @@ def load_code_block_diff_map(global_db: sqlite3.Connection) -> dict[str, list[di
 
 
 def load_bubbles_for_composer(
-    global_db: sqlite3.Connection, composer_id: str,
-) -> dict[str, dict[str, Any]]:
-    """Load ``bubbleId:{composer_id}:*`` KV entries into ``{bubble_id: bubble_dict}``.
+    global_db: sqlite3.Connection,
+    composer_id: str,
+    *,
+    parse_warnings: ParseWarningCollector | None = None,
+) -> dict[str, Bubble]:
+    """Load ``bubbleId:{composer_id}:*`` KV entries into ``{bubble_id: Bubble}``.
 
     Scoped alternative to :func:`load_bubble_map` for single-conversation assembly;
     avoids a full global ``bubbleId:%`` scan.
     """
-    bubble_map: dict[str, dict[str, Any]] = {}
+    bubble_map: dict[str, Bubble] = {}
     try:
         rows = global_db.execute(
             "SELECT key, value FROM cursorDiskKV WHERE key LIKE ?",
@@ -179,16 +216,12 @@ def load_bubbles_for_composer(
     except sqlite3.Error:
         return bubble_map
     for row in rows:
-        parts = row["key"].split(":")
-        if len(parts) < 3:
-            continue
-        bid = parts[2]
-        try:
-            b = json.loads(row["value"])
-            if isinstance(b, dict):
-                bubble_map[bid] = b
-        except (json.JSONDecodeError, ValueError, KeyError, TypeError) as e:
-            _logger.debug("Skipping malformed bubbleId row %s: %s", row["key"], e)
+        parsed = _parse_bubble_kv_row(
+            row["key"], row["value"], parse_warnings=parse_warnings
+        )
+        if parsed is not None:
+            bid, bubble = parsed
+            bubble_map[bid] = bubble
     return bubble_map
 
 

services/workspace_listing.py

@@ -18,7 +18,7 @@
     warn_workspace_json_read,
 )
 from utils.workspace_descriptor import read_json_file
-from models import ParseWarningCollector
+from models import Bubble, ParseWarningCollector
 from services.summary_cache import (
     fingerprint_workspace_storage,
     get_cached_projects,
@@ -146,7 +146,7 @@ def _build_workspace_projects_uncached(
                 if invalid_workspace_ids:
                     project_layouts_map = load_project_layouts_map(global_db)
 
-                bubble_map: dict[str, dict[str, Any]] = {}
+                bubble_map: dict[str, Bubble] = {}
                 invalid_workspace_aliases: dict[str, str] = {}
                 if invalid_workspace_ids:
                     invalid_workspace_aliases = infer_invalid_workspace_aliases(

services/workspace_resolver.py

@@ -248,7 +248,7 @@ def determine_project_for_conversation(
     project_name_to_workspace_id: dict[str, str],
     workspace_path_to_id: dict[str, str],
     workspace_entries: list[dict[str, Any]],
-    bubble_map: Mapping[str, Bubble | dict[str, Any]],
+    bubble_map: Mapping[str, Bubble],
     composer_id_to_workspace_id: dict[str, str] | None = None,
     invalid_workspace_ids: set[str] | None = None,
 ) -> str | None:
@@ -261,7 +261,7 @@ def determine_project_for_conversation(
         project_name_to_workspace_id: Basename-to-workspace-folder map.
         workspace_path_to_id: Normalized root path to workspace folder map.
         workspace_entries: Output of :func:`services.workspace_db.collect_workspace_entries`.
-        bubble_map: ``{bubble_id: Bubble | bubble_dict}`` from global KV.
+        bubble_map: ``{bubble_id: Bubble}`` from global KV loaders.
         composer_id_to_workspace_id: Definitive per-workspace composer map; when
             ``None``, layout and path heuristics are used without this shortcut.
         invalid_workspace_ids: Workspace folders marked invalid; mapped IDs in
@@ -380,7 +380,7 @@ def infer_invalid_workspace_aliases(
     project_name_map: dict[str, str],
     workspace_path_map: dict[str, str],
     workspace_entries: list[dict[str, Any]],
-    bubble_map: Mapping[str, Bubble | dict[str, Any]],
+    bubble_map: Mapping[str, Bubble],
     composer_id_to_ws: dict[str, str],
     invalid_workspace_ids: set[str],
 ) -> dict[str, str]:
@@ -396,7 +396,7 @@ def infer_invalid_workspace_aliases(
         project_name_map: Basename map for path resolution.
         workspace_path_map: Normalized path map for path resolution.
         workspace_entries: Workspace folder entries from storage scan.
-        bubble_map: ``{bubble_id: Bubble | bubble_dict}`` for path resolution.
+        bubble_map: ``{bubble_id: Bubble}`` for path resolution.
         composer_id_to_ws: Composer-to-workspace map (may point at invalid IDs).
         invalid_workspace_ids: Workspace folder names to reassign.