collossal 0.1.0__py3-none-any.whl

main_mcp/capture/history.py

@@ -0,0 +1,349 @@
+"""Past sessions, uploaded from the transcripts already on this machine.
+
+Claude Code writes every session to ``~/.claude/projects/<folder>/<id>.jsonl``
+and keeps them for weeks. The live stream only sees a session from the moment a
+hook attaches to it, so on a machine that has just installed Collossal all of
+that history is invisible — the dashboard shows a teammate with two sessions
+when they have run eighty. This uploads it: every session, whatever it ran in.
+
+Three rules keep it honest:
+
+* **Nothing twice.** Before any transcript is read in full, the server is asked
+  which of them it wants. One the live stream already carried is left alone
+  rather than duplicated, and one imported before only sends what it gained.
+* **Real time.** Every event keeps the moment it happened, so an imported
+  session sits in the history where it happened, not where it arrived.
+* **Same redaction.** Events are rendered by the tailer's own function, so
+  nothing leaves the machine by this route that the live stream would have
+  scrubbed.
+
+Run by ``collossal install`` and ``collossal import`` over everything, and by the
+SessionEnd hook over the one session that just ended — which is what completes
+a session that began before the hooks were installed.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Iterator
+
+from ..agents import claude_dir
+from ..client import CollossalClient, CollossalError
+from ..config import ConfigError, load_config
+from ..git_context import read_git_context
+from .checkouts import Tally
+from .hook import _log
+from .tailer import render_event
+from .transcript import SessionMeta, read_events
+
+#: A transcript that moved more recently than this may still be running — the
+#: tailer's own idle horizon, so the two never race over one session. It is
+#: picked up when it ends.
+ACTIVE_SECONDS = 15 * 60
+
+#: Per request: well inside the server's limits (500 events, 8 MB), and small
+#: enough that one slow chunk is not a stall.
+CHUNK_EVENTS = 200
+CHUNK_CHARS = 1_500_000
+
+#: The server plans this many sessions per call.
+PLAN_BATCH = 1000
+
+#: A rate-limited request is retried this many times before the session is
+#: given up on (and left for the next run).
+RETRIES = 6
+
+
+@dataclass
+class Transcript:
+    path: Path
+    session_id: str
+    #: The directory the session ran in — which is what says which repository.
+    cwd: str | None
+    modified: float
+
+
+def _first_cwd(path: Path) -> str | None:
+    """A session's working directory, from its opening records."""
+    try:
+        with path.open("r", encoding="utf-8", errors="replace") as fh:
+            for count, line in enumerate(fh):
+                if count > 200:
+                    break
+                try:
+                    record = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                cwd = record.get("cwd") if isinstance(record, dict) else None
+                if isinstance(cwd, str) and cwd:
+                    return cwd
+    except OSError:
+        return None
+    return None
+
+
+def find_transcripts(root: Path | None = None) -> list[Transcript]:
+    """Every session transcript on this machine, newest first.
+
+    The top level of each project folder only: the folders beside the
+    transcripts hold subagents' runs, which belong to their parent session and
+    are not sessions of their own.
+    """
+    projects = (root or claude_dir()) / "projects"
+    try:
+        folders = [p for p in projects.iterdir() if p.is_dir()]
+    except OSError:
+        return []
+
+    found: list[Transcript] = []
+    for folder in folders:
+        for path in folder.glob("*.jsonl"):
+            try:
+                modified = path.stat().st_mtime
+            except OSError:
+                continue
+            found.append(Transcript(path, path.stem, _first_cwd(path), modified))
+    return sorted(found, key=lambda t: t.modified, reverse=True)
+
+
+def touched_checkouts(transcript: Transcript) -> list[Path]:
+    """The checkouts a finished session worked in, most-worked first."""
+    parsed, _ = read_events(transcript.path, start_line=0)
+    tally = Tally(Path(transcript.cwd) if transcript.cwd else None)
+    tally.add(parsed)
+    return tally.ranked()
+
+
+def build(transcript: Transcript, strip_dir: Path | None) -> dict[str, Any] | None:
+    """A transcript as the import endpoint takes it; None when it holds nothing.
+
+    `strip_dir` is taken off the front of every path the session quotes: its
+    repository's root, or — for a session that began in a folder — that
+    folder, so a path still says which repository it is in.
+    """
+    meta = SessionMeta()
+    parsed, _ = read_events(transcript.path, start_line=0, meta=meta)
+
+    events: list[dict[str, Any]] = []
+    last_at = meta.first_at
+    for e in parsed:
+        item = render_event(e, strip_dir)
+        if item is None:
+            continue
+        # A record without a time of its own happened no earlier than the one
+        # before it.
+        at = e.timestamp or last_at
+        if at is None:
+            continue
+        last_at = at
+        events.append({"seq": len(events) + 1, **item, "at": at})
+
+    if not events or not meta.first_at:
+        return None
+
+    payload: dict[str, Any] = {
+        "startedAt": meta.first_at,
+        "endedAt": meta.last_at or meta.first_at,
+        "events": events,
+    }
+    if meta.title:
+        payload["title"] = meta.title[:200]
+    if meta.model:
+        payload["model"] = meta.model[:100]
+    if meta.branch:
+        payload["branch"] = meta.branch[:200]
+    return payload
+
+
+def _chunks(events: list[dict[str, Any]]) -> Iterator[list[dict[str, Any]]]:
+    chunk: list[dict[str, Any]] = []
+    size = 0
+    for event in events:
+        n = len(event["body"])
+        if chunk and (len(chunk) >= CHUNK_EVENTS or size + n > CHUNK_CHARS):
+            yield chunk
+            chunk, size = [], 0
+        chunk.append(event)
+        size += n
+    if chunk:
+        yield chunk
+
+
+async def _send(client: CollossalClient, payload: dict[str, Any]) -> dict[str, Any]:
+    """One upload, waiting out a rate limit rather than losing the session to it."""
+    for attempt in range(RETRIES):
+        try:
+            return await client.import_session(payload)
+        except CollossalError as exc:
+            if exc.status != 429 or attempt == RETRIES - 1:
+                raise
+            await asyncio.sleep(min(max(exc.retry_after or 10.0, 1.0), 60.0))
+    raise AssertionError("unreachable")
+
+
+@dataclass
+class Summary:
+    found: int = 0
+    imported: int = 0
+    events: int = 0
+    streamed: int = 0
+    up_to_date: int = 0
+    active: int = 0
+    by_activity: int = 0
+    folder_only: int = 0
+    empty: int = 0
+    failed: int = 0
+
+
+#: Called as each session starts uploading: (index, total, title).
+Progress = Callable[[int, int, str], None]
+
+
+@dataclass
+class _Candidate:
+    transcript: Transcript
+    #: What the session is labelled with — every part optional.
+    place: dict[str, str]
+    #: Taken off the front of the paths it quotes.
+    strip_dir: Path | None
+    #: How far the server already has it.
+    reached: int = 0
+
+
+async def import_history(
+    client: CollossalClient,
+    transcripts: list[Transcript],
+    *,
+    progress: Progress | None = None,
+    include_active: bool = False,
+) -> Summary:
+    summary = Summary(found=len(transcripts))
+    now = time.time()
+
+    # Where each one ran — asked of git once per directory. Every session is
+    # uploaded; this only decides what it is labelled with.
+    contexts: dict[Path, Any] = {}
+
+    def context(directory: Path):
+        if directory not in contexts:
+            contexts[directory] = read_git_context(directory)
+        return contexts[directory]
+
+    candidates: list[_Candidate] = []
+    for t in transcripts:
+        if not include_active and now - t.modified < ACTIVE_SECONDS:
+            summary.active += 1
+            continue
+        cwd = Path(t.cwd) if t.cwd else None
+        place: dict[str, str] = {"folder": cwd.name} if cwd and cwd.name else {}
+
+        git = context(cwd) if cwd and cwd.is_dir() else None
+        if git and git.usable:
+            place["remoteUrl"] = git.remote_url
+            candidates.append(_Candidate(t, place, git.root))
+            continue
+
+        # Began in a folder that is not a checkout — most often one holding
+        # several. Labelled by the repository it worked in most, if any.
+        chosen = next(
+            (ctx for root in touched_checkouts(t) if (ctx := context(root)).usable),
+            None,
+        )
+        if chosen:
+            place["remoteUrl"] = chosen.remote_url
+            summary.by_activity += 1
+        else:
+            summary.folder_only += 1
+        candidates.append(_Candidate(t, place, cwd))
+
+    wanted: list[_Candidate] = []
+    for start in range(0, len(candidates), PLAN_BATCH):
+        batch = candidates[start : start + PLAN_BATCH]
+        plan = await client.import_plan([c.transcript.session_id for c in batch])
+        # The server answers in the order it was asked.
+        for candidate, verdict in zip(batch, plan.get("sessions") or []):
+            if verdict.get("action") == "import":
+                candidate.reached = int(verdict.get("lastSeq") or 0)
+                wanted.append(candidate)
+            else:
+                summary.streamed += 1
+
+    for index, c in enumerate(wanted, start=1):
+        t = c.transcript
+        payload = build(t, c.strip_dir)
+        if payload is None:
+            summary.empty += 1
+            continue
+        # Events are numbered by their place in the transcript, so what the
+        # server already holds is exactly a prefix: send only what follows it.
+        fresh = [e for e in payload["events"] if e["seq"] > c.reached]
+        if not fresh:
+            summary.up_to_date += 1
+            continue
+        if progress:
+            progress(index, len(wanted), payload.get("title") or t.session_id[:8])
+
+        head = {k: v for k, v in payload.items() if k != "events"}
+        try:
+            accepted = 0
+            status = "imported"
+            for chunk in _chunks(fresh):
+                result = await _send(
+                    client,
+                    {"externalId": t.session_id, **c.place, **head, "events": chunk},
+                )
+                status = result.get("status", status)
+                if status == "streamed":
+                    # It began streaming between the plan and now.
+                    break
+                accepted += int(result.get("accepted") or 0)
+        except CollossalError as exc:
+            _log(f"history: {t.session_id[:8]} failed ({exc})")
+            summary.failed += 1
+            continue
+
+        if status == "streamed":
+            summary.streamed += 1
+        else:
+            summary.imported += 1
+            summary.events += accepted
+
+    return summary
+
+
+def main() -> None:
+    """The SessionEnd hook's detached worker: the session that just ended.
+
+    A session that began before the hooks were installed was never streamed,
+    and this is what completes it. One the live stream carried costs a single
+    small request, which the plan answers with "streamed".
+    """
+    try:
+        raw = sys.stdin.read()
+        event = json.loads(raw) if raw.strip() else {}
+        path = Path(event.get("transcript_path") or "")
+        session_id = event.get("session_id")
+        if not session_id or not path.is_file():
+            return
+        transcript = Transcript(
+            path,
+            session_id,
+            event.get("cwd") or _first_cwd(path),
+            path.stat().st_mtime,
+        )
+        client = CollossalClient(load_config())
+        asyncio.run(import_history(client, [transcript], include_active=True))
+    except (ConfigError, CollossalError, OSError, ValueError) as exc:
+        _log(f"history: {exc}")
+    except Exception as exc:  # noqa: BLE001 - must never surface to the session
+        _log(f"history error: {exc}")
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()