arborist/arborist/journal.py

"""Unfirehose-compatible session journal.

Emits one JSONL file per query/bench-cycle to
``~/.arborist/unfirehose/{project-slug}/{session-uuid}.jsonl``.
Unfirehose's native-harness auto-discovery picks up any
``~/.{name}/unfirehose/`` directory without registration; once a
session lands, the unfirehose watcher debounces, ingests, and
exposes it in the dashboard alongside Claude Code, Fetch, and
uncloseai sessions.

Schema reference: ``~/git/unfirehose-nextjs-logger/docs/unfirehose-schema.md``
(``unfirehose/1.0``). Every record carries ``$schema:
"unfirehose/1.0"``; consumers ignore unknown fields, so arborist-
specific extras (audit_mode, cache_key, source list, verifier
timings) ride along under namespaced keys.

Per-query layout:

    session header line          type=session
    user message                 role=user, content=[text]
    assistant message            role=assistant, content=[text],
                                 usage=..., model=...,
                                 arborist_meta={audit_mode, cache_key,
                                               sources, timings}
    system session_end message   subtype=session_end, durationMs

Per-bench-cycle layout (bench_emergent):

    session header line          (firstPrompt = generated question)
    system init message          subtype=init, arborist_meta={words}
    user message                 the generated question
    assistant message            the student answer + arborist_meta
    system session_end message   subtype=session_end

Everything arborist-specific lives under ``arborist_meta`` so the
canonical fields stay clean for off-the-shelf unfirehose consumers
that expect the strict schema.
"""

from __future__ import annotations

import json
import os
import time
import uuid
from pathlib import Path
from typing import Any


# Default journal root. Override via env var or constructor.
DEFAULT_JOURNAL_ROOT = Path.home() / ".arborist" / "unfirehose"

# Schema literal pinned to a single string — matches what unfirehose's
# native-harness ingestion checks for.
UNFIREHOSE_SCHEMA = "unfirehose/1.0"

# Harness identity. Any arborist process emitting JSONL claims this
# `harness` name. Unfirehose surfaces it in its dashboard alongside
# claude-code / fetch / uncloseai.
HARNESS_NAME = "arborist"


def _arborist_version() -> str:
    """Best-effort version tag for the harnessVersion field."""
    try:
        from importlib.metadata import version as _v
        return _v("arborist")
    except Exception:
        return "0.0.0+dev"


def slugify_cwd(cwd: Path | str | None = None) -> str:
    """Convert a working-directory path to an unfirehose project slug.

    Rules (mirror Claude Code's encoding):
      - leading ``/`` becomes leading ``-``
      - path separators (``/``) become ``-``
      - dots (``.``) become ``-``

    ``/home/fox/git/arborist``  → ``-home-fox-git-arborist``
    ``/home/fox/git/my.app``   → ``-home-fox-git-my-app``
    """
    p = Path(cwd) if cwd else Path.cwd()
    s = str(p)
    return s.replace("/", "-").replace(".", "-")


def now_iso() -> str:
    """ISO 8601 UTC timestamp with millisecond precision."""
    t = time.time()
    base = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(t))
    ms = int((t - int(t)) * 1000)
    return f"{base}.{ms:03d}Z"


def new_session_id() -> str:
    """UUID4 stringification for session_id / message_id."""
    return str(uuid.uuid4())


def _git_branch(cwd: Path | str | None = None) -> str | None:
    """Best-effort current git branch. Returns None on failure."""
    try:
        import subprocess
        out = subprocess.check_output(
            ["git", "-C", str(cwd or "."), "branch", "--show-current"],
            stderr=subprocess.DEVNULL,
            timeout=2,
        )
        return out.decode().strip() or None
    except Exception:
        return None


class SessionWriter:
    """One open .jsonl file representing one arborist session.

    Use as a context manager so the close timestamp is set & the file
    flushed on exit:

        with SessionWriter(first_prompt=question, cwd=Path.cwd()) as s:
            s.user_message(question)
            s.assistant_message(answer, model=..., usage=..., arborist_meta=...)

    The constructor writes the session header line; subsequent calls
    append message lines; ``__exit__`` writes the session_end line.
    """

    def __init__(
        self,
        *,
        first_prompt: str,
        cwd: Path | str | None = None,
        root: Path | str | None = None,
        session_id: str | None = None,
        sidechain: bool = False,
    ) -> None:
        self.session_id = session_id or new_session_id()
        self.cwd = Path(cwd or Path.cwd())
        self.root = Path(root or os.environ.get("ARBORIST_JOURNAL_ROOT") or DEFAULT_JOURNAL_ROOT)
        self.project_slug = slugify_cwd(self.cwd)
        self.sidechain = sidechain
        self.git_branch = _git_branch(self.cwd)
        self.harness_version = _arborist_version()

        # File path: {root}/{slug}/{session_id}.jsonl
        self.session_dir = self.root / self.project_slug
        self.session_dir.mkdir(parents=True, exist_ok=True)
        self.path = self.session_dir / f"{self.session_id}.jsonl"

        self._fp = self.path.open("a", encoding="utf-8")
        self._t_start = time.time()
        self._closed = False
        self._message_count = 0
        self._last_message_id: str | None = None

        # Session header
        self._write({
            "$schema": UNFIREHOSE_SCHEMA,
            "type": "session",
            "id": self.session_id,
            "projectId": self.project_slug,
            "status": "active",
            "createdAt": now_iso(),
            "updatedAt": now_iso(),
            "closedAt": None,
            "firstPrompt": (first_prompt or "")[:500],
            "summary": "",
            "gitBranch": self.git_branch,
            "cwd": str(self.cwd),
            "sidechain": self.sidechain,
            "harness": HARNESS_NAME,
            "harnessVersion": self.harness_version,
        })

    # ─────────────────────────────────────────────────────────── messages

    def _write(self, obj: dict[str, Any]) -> None:
        if self._closed:
            raise RuntimeError("session already closed")
        self._fp.write(json.dumps(obj, ensure_ascii=False) + "\n")
        self._fp.flush()

    def _new_message_id(self) -> str:
        return f"msg_{uuid.uuid4().hex[:24]}"

    def _base_message(self, role: str) -> dict[str, Any]:
        msg_id = self._new_message_id()
        record: dict[str, Any] = {
            "$schema": UNFIREHOSE_SCHEMA,
            "type": "message",
            "id": msg_id,
            "sessionId": self.session_id,
            "parentId": self._last_message_id,
            "role": role,
            "timestamp": now_iso(),
            "sidechain": self.sidechain,
            "cwd": str(self.cwd),
            "gitBranch": self.git_branch,
            "harness": HARNESS_NAME,
            "harnessVersion": self.harness_version,
        }
        self._last_message_id = msg_id
        self._message_count += 1
        return record

    def user_message(
        self,
        text: str,
        *,
        arborist_meta: dict[str, Any] | None = None,
    ) -> None:
        rec = self._base_message("user")
        rec["content"] = [{"type": "text", "text": text}]
        if arborist_meta:
            rec["arborist_meta"] = arborist_meta
        self._write(rec)

    def assistant_message(
        self,
        text: str,
        *,
        model: str | None = None,
        provider: str | None = None,
        stop_reason: str | None = None,
        usage: dict[str, Any] | None = None,
        duration_ms: int | None = None,
        arborist_meta: dict[str, Any] | None = None,
    ) -> None:
        rec = self._base_message("assistant")
        rec["content"] = [{"type": "text", "text": text}]
        if model:
            rec["model"] = model
        if provider:
            rec["provider"] = provider
        if stop_reason:
            rec["stopReason"] = stop_reason
        if usage:
            rec["usage"] = usage
        if duration_ms is not None:
            rec["durationMs"] = int(duration_ms)
        if arborist_meta:
            rec["arborist_meta"] = arborist_meta
        self._write(rec)

    def system_message(
        self,
        text: str = "",
        *,
        subtype: str | None = None,
        duration_ms: int | None = None,
        arborist_meta: dict[str, Any] | None = None,
    ) -> None:
        rec = self._base_message("system")
        if text:
            rec["content"] = [{"type": "text", "text": text}]
        if subtype:
            rec["subtype"] = subtype
        if duration_ms is not None:
            rec["durationMs"] = int(duration_ms)
        if arborist_meta:
            rec["arborist_meta"] = arborist_meta
        self._write(rec)

    # ──────────────────────────────────────────────────────────── lifecycle

    def close(self, *, status: str = "closed") -> None:
        if self._closed:
            return
        elapsed_ms = int((time.time() - self._t_start) * 1000)
        # session_end system message — gives unfirehose a clean signal
        # that the session is complete + the wall-clock duration.
        self.system_message(
            subtype="session_end",
            duration_ms=elapsed_ms,
            arborist_meta={"status": status, "messageCount": self._message_count},
        )
        self._fp.close()
        self._closed = True

    def __enter__(self) -> "SessionWriter":
        return self

    def __exit__(self, *exc) -> None:
        self.close()