cursor_bubbles.py

"""Read Cursor agent chat bubbles from globalStorage/state.vscdb.

This is the same source cursor-history uses (cursorDiskKV bubbleId / composerData).
Unlike ~/.cursor/projects/*/agent-transcripts/*.jsonl, bubbles retain thinking.text
and per-step toolFormerData in chronological order.
"""

from __future__ import annotations

import json
import os
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.parse import unquote, urlparse

THINKING_MAX = 16_000
TOOL_IO_MAX = 8_000


def cursor_user_dir() -> Path:
    """Cursor User data root (platform-aware)."""
    custom = (os.environ.get("CURSOR_DATA_PATH") or "").strip()
    if custom:
        # CURSOR_DATA_PATH usually points at workspaceStorage; parent is User/
        p = Path(custom).expanduser()
        if p.name.lower() == "workspacestorage":
            return p.parent
        if (p / "workspaceStorage").is_dir():
            return p
        return p
    if os.name == "nt":
        appdata = os.environ.get("APPDATA") or str(Path.home() / "AppData" / "Roaming")
        return Path(appdata) / "Cursor" / "User"
    if Path.home().joinpath("Library", "Application Support").is_dir():
        mac = Path.home() / "Library" / "Application Support" / "Cursor" / "User"
        if mac.is_dir():
            return mac
    return Path.home() / ".config" / "Cursor" / "User"


def global_state_db_path() -> Path:
    return cursor_user_dir() / "globalStorage" / "state.vscdb"


def workspace_storage_dir() -> Path:
    return cursor_user_dir() / "workspaceStorage"


def _norm_path(path: str | Path) -> str:
    try:
        return str(Path(path).expanduser().resolve()).replace("\\", "/").lower()
    except OSError:
        return str(path).replace("\\", "/").lower()


def _uri_to_fs_path(uri: str | None) -> str | None:
    if not uri or not isinstance(uri, str):
        return None
    raw = uri.strip()
    if raw.startswith("file:"):
        parsed = urlparse(raw)
        path = unquote(parsed.path or "")
        # file:///c%3A/Users/... → /c:/Users/... on Windows
        if len(path) >= 3 and path[0] == "/" and path[2] == ":":
            path = path[1:]
        return path.replace("\\", "/")
    return raw.replace("\\", "/")


def workspace_ids_for_root(root: str | Path) -> list[str]:
    """Map a scanned repo root → Cursor workspaceStorage folder id(s)."""
    target = _norm_path(root)
    if not target:
        return []
    ws = workspace_storage_dir()
    if not ws.is_dir():
        return []
    found: list[str] = []
    for d in ws.iterdir():
        if not d.is_dir():
            continue
        wj = d / "workspace.json"
        if not wj.is_file():
            continue
        try:
            data = json.loads(wj.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        folder = data.get("folder") or data.get("workspace") or ""
        fs = _uri_to_fs_path(str(folder)) if folder else None
        if not fs:
            continue
        if _norm_path(fs) == target:
            found.append(d.name)
    return found


def _open_ro(db_path: Path) -> sqlite3.Connection | None:
    if not db_path.is_file():
        return None
    try:
        return sqlite3.connect(f"file:{db_path.as_posix()}?mode=ro", uri=True)
    except sqlite3.Error:
        return None


def _parse_json(value: Any) -> Any:
    if isinstance(value, (dict, list)):
        return value
    if not isinstance(value, str) or not value.strip():
        return None
    try:
        return json.loads(value)
    except json.JSONDecodeError:
        return None


def _iso_at(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        # Cursor often stores ms epochs
        ms = float(value)
        if ms > 1e12:
            ms = ms / 1000.0
        try:
            return (
                datetime.fromtimestamp(ms, tz=timezone.utc)
                .replace(microsecond=0)
                .isoformat()
                .replace("+00:00", "Z")
            )
        except (OSError, OverflowError, ValueError):
            return None
    if isinstance(value, str):
        s = value.strip()
        if not s:
            return None
        if s.isdigit():
            return _iso_at(int(s))
        # Already ISO-ish
        if "T" in s:
            return s.replace("+00:00", "Z")
        return s
    return None


def _truncate(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    return text[: limit - 1] + "…"


def _composer_workspace_path(composer: dict[str, Any]) -> str | None:
    wid = composer.get("workspaceIdentifier")
    if isinstance(wid, dict):
        uri = wid.get("uri")
        if isinstance(uri, dict):
            fs = uri.get("fsPath") or uri.get("path")
            if isinstance(fs, str) and fs.strip():
                return fs.replace("\\", "/")
            external = uri.get("external")
            got = _uri_to_fs_path(external) if isinstance(external, str) else None
            if got:
                return got
        elif isinstance(uri, str):
            got = _uri_to_fs_path(uri)
            if got:
                return got
    wuri = composer.get("workspaceUri")
    if isinstance(wuri, str):
        return _uri_to_fs_path(wuri)
    return None


def _composer_matches_root(composer: dict[str, Any], root: str | Path, workspace_ids: set[str]) -> bool:
    path = _composer_workspace_path(composer)
    if path and _norm_path(path) == _norm_path(root):
        return True
    wid = composer.get("workspaceIdentifier")
    if isinstance(wid, dict):
        cid = wid.get("id")
        if isinstance(cid, str) and cid in workspace_ids:
            return True
    return False


def _list_composer_ids_for_root(
    con: sqlite3.Connection,
    root: str | Path,
    *,
    limit: int,
) -> list[str]:
    """Prefer composerHeaders (fast), fall back to scanning composerData keys."""
    workspace_ids = set(workspace_ids_for_root(root))
    ids: list[str] = []
    seen: set[str] = set()

    # Fast path: composerHeaders.value JSON includes workspaceIdentifier
    try:
        rows = con.execute(
            "SELECT composerId, value, lastUpdatedAt, createdAt, recency "
            "FROM composerHeaders WHERE COALESCE(isArchived,0)=0 "
            "ORDER BY COALESCE(recency, lastUpdatedAt, createdAt) DESC"
        ).fetchall()
    except sqlite3.Error:
        rows = []

    for composer_id, value, *_rest in rows:
        if not composer_id or composer_id in seen:
            continue
        if composer_id in ("empty-state-draft",):
            continue
        head = _parse_json(value) or {}
        if not isinstance(head, dict):
            head = {}
        # Headers sometimes only stamp workspace id — also load full composerData when unsure
        if workspace_ids and isinstance(head.get("workspaceIdentifier"), dict):
            wid = head["workspaceIdentifier"].get("id")
            if isinstance(wid, str) and wid in workspace_ids:
                seen.add(composer_id)
                ids.append(composer_id)
                if len(ids) >= limit:
                    return ids
                continue
        # Need full composer for path match
        raw = con.execute(
            "SELECT value FROM cursorDiskKV WHERE key = ?",
            (f"composerData:{composer_id}",),
        ).fetchone()
        if not raw:
            continue
        composer = _parse_json(raw[0])
        if not isinstance(composer, dict):
            continue
        if _composer_matches_root(composer, root, workspace_ids):
            seen.add(composer_id)
            ids.append(composer_id)
            if len(ids) >= limit:
                return ids

    if ids:
        return ids

    # Slow fallback: newest composerData rows
    try:
        crow = con.execute(
            "SELECT key, value FROM cursorDiskKV WHERE key LIKE 'composerData:%' "
            "ORDER BY rowid DESC LIMIT ?",
            (max(limit * 8, 80),),
        ).fetchall()
    except sqlite3.Error:
        return []
    for key, value in crow:
        composer_id = key.split(":", 1)[1] if ":" in key else ""
        if not composer_id or composer_id in seen:
            continue
        composer = _parse_json(value)
        if not isinstance(composer, dict):
            continue
        if _composer_matches_root(composer, root, workspace_ids):
            seen.add(composer_id)
            ids.append(composer_id)
            if len(ids) >= limit:
                break
    return ids


def _load_composer(con: sqlite3.Connection, composer_id: str) -> dict[str, Any] | None:
    raw = con.execute(
        "SELECT value FROM cursorDiskKV WHERE key = ?",
        (f"composerData:{composer_id}",),
    ).fetchone()
    if not raw:
        return None
    data = _parse_json(raw[0])
    return data if isinstance(data, dict) else None


def _load_bubbles(con: sqlite3.Connection, composer_id: str) -> dict[str, dict[str, Any]]:
    out: dict[str, dict[str, Any]] = {}
    try:
        rows = con.execute(
            "SELECT key, value FROM cursorDiskKV WHERE key LIKE ?",
            (f"bubbleId:{composer_id}:%",),
        ).fetchall()
    except sqlite3.Error:
        return out
    for key, value in rows:
        data = _parse_json(value)
        if not isinstance(data, dict):
            continue
        bid = data.get("bubbleId")
        if not isinstance(bid, str) or not bid:
            bid = key.rsplit(":", 1)[-1]
        out[bid] = data
    return out


def _ordered_bubbles(composer: dict[str, Any], by_id: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
    headers = composer.get("fullConversationHeadersOnly")
    ordered: list[dict[str, Any]] = []
    seen: set[str] = set()
    if isinstance(headers, list):
        for h in headers:
            if not isinstance(h, dict):
                continue
            bid = h.get("bubbleId")
            if not isinstance(bid, str) or bid in seen:
                continue
            data = by_id.get(bid)
            if not data:
                continue
            seen.add(bid)
            # Prefer header createdAt when bubble lacks one
            if not data.get("createdAt") and h.get("createdAt"):
                data = {**data, "createdAt": h.get("createdAt")}
            ordered.append(data)
    # Append any leftover bubbles (subagents / orphans) by createdAt
    leftovers = [b for bid, b in by_id.items() if bid not in seen]
    leftovers.sort(key=lambda b: str(b.get("createdAt") or ""))
    ordered.extend(leftovers)
    return ordered


def _tool_payload(tf: dict[str, Any]) -> dict[str, Any]:
    name = str(tf.get("name") or tf.get("tool") or "tool")
    params = tf.get("params")
    raw_args = tf.get("rawArgs")
    input_obj: Any = _parse_json(params) if isinstance(params, str) else params
    if input_obj is None and raw_args is not None:
        input_obj = _parse_json(raw_args) if isinstance(raw_args, str) else raw_args
    if input_obj is None and isinstance(params, str):
        input_obj = params

    result = tf.get("result")
    output: Any = _parse_json(result) if isinstance(result, str) else result
    if output is None and isinstance(result, str):
        output = result
    if isinstance(output, str):
        output = _truncate(output, TOOL_IO_MAX)
    elif output is not None:
        try:
            dumped = json.dumps(output, ensure_ascii=False, default=str)
            if len(dumped) > TOOL_IO_MAX:
                output = _truncate(dumped, TOOL_IO_MAX)
        except TypeError:
            output = _truncate(str(output), TOOL_IO_MAX)

    tool: dict[str, Any] = {
        "id": tf.get("toolCallId") or tf.get("modelCallId"),
        "name": name,
        "status": tf.get("status"),
    }
    if input_obj is not None:
        tool["input"] = input_obj
    if output is not None:
        tool["output"] = output
    return tool


def _bubble_to_messages(bubble: dict[str, Any]) -> list[dict[str, Any]]:
    """Expand one Cursor bubble into zero or more stored chat messages (chrono atoms)."""
    at = _iso_at(bubble.get("createdAt"))
    bubble_type = bubble.get("type")
    # type 1 = user, type 2 = assistant (cursor-history convention)
    is_user = bubble_type == 1

    tf = bubble.get("toolFormerData") if isinstance(bubble.get("toolFormerData"), dict) else None
    thinking_obj = bubble.get("thinking") if isinstance(bubble.get("thinking"), dict) else None
    thinking_text = ""
    if thinking_obj and isinstance(thinking_obj.get("text"), str):
        thinking_text = thinking_obj["text"].strip()

    text = bubble.get("text")
    text = text.strip() if isinstance(text, str) else ""

    out: list[dict[str, Any]] = []

    if is_user and text:
        msg: dict[str, Any] = {"role": "user", "content": text}
        if at:
            msg["at"] = at
        out.append(msg)
        return out

    # Thinking-only (or thinking before tool/text on same bubble)
    if thinking_text:
        msg = {
            "role": "assistant",
            "content": "",
            "meta": {"thinking": [_truncate(thinking_text, THINKING_MAX)]},
        }
        if at:
            msg["at"] = at
        out.append(msg)

    if tf and (tf.get("name") or tf.get("tool")):
        tool = _tool_payload(tf)
        msg = {
            "role": "tool",
            "name": tool.get("name") or "tool",
            "tool": tool,
            "content": "",
        }
        if at:
            msg["at"] = at
        out.append(msg)
        # Tool bubbles rarely also carry visible assistant prose; still allow it below.

    if text:
        msg = {"role": "assistant", "content": text}
        if at:
            msg["at"] = at
        # Avoid duplicating a thinking-only empty assistant we already emitted
        out.append(msg)

    return out


def composer_to_thread(
    composer_id: str,
    composer: dict[str, Any],
    bubbles_by_id: dict[str, dict[str, Any]],
    *,
    repo_root: str | None = None,
    cursor_project: str | None = None,
) -> dict[str, Any] | None:
    ordered = _ordered_bubbles(composer, bubbles_by_id)
    messages: list[dict[str, Any]] = []
    for bubble in ordered:
        messages.extend(_bubble_to_messages(bubble))

    if not messages:
        return None

    title = ""
    for key in ("name", "title"):
        val = composer.get(key)
        if isinstance(val, str) and val.strip():
            title = val.strip()[:120]
            break
    if not title:
        for m in messages:
            if m.get("role") == "user" and m.get("content"):
                line = str(m["content"]).split("\n", 1)[0].strip()
                title = (line[:64] + "…") if len(line) > 64 else line
                break
    if not title:
        title = f"Cursor session {composer_id[:8]}"

    first_at = next((m.get("at") for m in messages if m.get("at")), None)
    last_at = next((m.get("at") for m in reversed(messages) if m.get("at")), None)
    created = _iso_at(composer.get("createdAt")) or first_at
    updated = _iso_at(composer.get("lastUpdatedAt") or composer.get("updatedAt")) or last_at or created

    tool_calls = [
        m["tool"]
        for m in messages
        if isinstance(m.get("tool"), dict)
    ]
    thinking_n = sum(
        1
        for m in messages
        if isinstance((m.get("meta") or {}).get("thinking"), list)
        and (m.get("meta") or {}).get("thinking")
    )

    out: dict[str, Any] = {
        "id": "",  # filled by importer
        "title": title or "Imported Cursor thread",
        "source": "cursor",
        "externalId": f"cursor:composer:{composer_id}",
        "composerId": composer_id,
        "createdAt": created or datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z"),
        "updatedAt": updated or created,
        "messages": messages,
        "trace": {
            "runs": [],
            "tool_calls": tool_calls,
            "source": "cursor-bubbles",
            "thinkingCount": thinking_n,
        },
    }
    if repo_root:
        out["repoRoot"] = repo_root
    if cursor_project:
        out["cursorProject"] = cursor_project
    return out


def list_threads_for_root(
    root: str | Path,
    *,
    limit: int = 20,
    db_path: str | Path | None = None,
) -> dict[str, Any]:
    """Load up to `limit` Cursor composer threads for a workspace root from state.vscdb."""
    path = Path(db_path) if db_path else global_state_db_path()
    if not path.is_file():
        return {
            "ok": False,
            "error": f"Cursor global DB not found at {path}",
            "threads": [],
            "dbPath": str(path),
            "workspaceIds": workspace_ids_for_root(root),
        }

    con = _open_ro(path)
    if con is None:
        return {
            "ok": False,
            "error": f"Could not open Cursor DB (is Cursor locking it?): {path}",
            "threads": [],
            "dbPath": str(path),
        }

    threads: list[dict[str, Any]] = []
    try:
        composer_ids = _list_composer_ids_for_root(con, root, limit=limit)
        for cid in composer_ids:
            composer = _load_composer(con, cid)
            if not composer:
                continue
            if composer.get("isDraft") and not composer.get("fullConversationHeadersOnly"):
                continue
            bubbles = _load_bubbles(con, cid)
            thread = composer_to_thread(cid, composer, bubbles)
            if thread:
                threads.append(thread)
    finally:
        con.close()

    return {
        "ok": True,
        "threads": threads,
        "dbPath": str(path),
        "workspaceIds": workspace_ids_for_root(root),
        "scanned": len(threads),
    }