arborist/tests/test_journal.py

"""Tests for the unfirehose-compatible session journal writer.

The contract being tested is the unfirehose/1.0 schema:
- session header line first
- one message per subsequent line
- session_end system message on close
- ``$schema: "unfirehose/1.0"`` on every record
- arborist-specific extras under namespaced ``arborist_meta``
"""

from __future__ import annotations

import json
from pathlib import Path

from arborist.journal import (
    HARNESS_NAME,
    UNFIREHOSE_SCHEMA,
    SessionWriter,
    new_session_id,
    slugify_cwd,
)


def _read_lines(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def test_slugify_cwd_basic():
    assert slugify_cwd("/home/fox/git/arborist") == "-home-fox-git-arborist"


def test_slugify_cwd_dots_become_hyphens():
    assert slugify_cwd("/home/fox/git/my.app") == "-home-fox-git-my-app"


def test_slugify_cwd_uses_cwd_when_none(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    expected = str(tmp_path).replace("/", "-").replace(".", "-")
    assert slugify_cwd() == expected


def test_session_header_first_line(tmp_path):
    with SessionWriter(first_prompt="hello?", root=tmp_path, cwd="/home/fox/git/myproj"):
        pass
    files = list(tmp_path.glob("**/*.jsonl"))
    assert len(files) == 1
    rows = _read_lines(files[0])
    header = rows[0]
    assert header["$schema"] == UNFIREHOSE_SCHEMA
    assert header["type"] == "session"
    assert header["projectId"] == "-home-fox-git-myproj"
    assert header["firstPrompt"] == "hello?"
    assert header["harness"] == HARNESS_NAME
    assert header["status"] == "active"


def test_session_messages_chained_via_parent_id(tmp_path):
    with SessionWriter(first_prompt="q", root=tmp_path) as s:
        s.user_message("first user")
        s.assistant_message("first assistant", model="hermes-3")
        s.user_message("second user")
    rows = _read_lines(list(tmp_path.glob("**/*.jsonl"))[0])
    # Look at the user/assistant chain — session_end is a separate
    # system message we exclude here.
    msgs = [r for r in rows if r["type"] == "message" and r["role"] != "system"]
    assert len(msgs) == 3
    # First message has no parent.
    assert msgs[0]["parentId"] is None
    # Subsequent messages chain to the previous.
    assert msgs[1]["parentId"] == msgs[0]["id"]
    assert msgs[2]["parentId"] == msgs[1]["id"]


def test_close_writes_session_end(tmp_path):
    with SessionWriter(first_prompt="q", root=tmp_path) as s:
        s.user_message("hi")
    rows = _read_lines(list(tmp_path.glob("**/*.jsonl"))[0])
    last = rows[-1]
    assert last["type"] == "message"
    assert last["role"] == "system"
    assert last["subtype"] == "session_end"
    assert isinstance(last["durationMs"], int)
    assert last["durationMs"] >= 0


def test_arborist_meta_passes_through(tmp_path):
    meta = {"audit_mode": "STRICT", "n_verified": 2, "cache_key": "abc123"}
    with SessionWriter(first_prompt="q", root=tmp_path) as s:
        s.assistant_message("answer", arborist_meta=meta)
    rows = _read_lines(list(tmp_path.glob("**/*.jsonl"))[0])
    asst = next(r for r in rows if r["type"] == "message" and r["role"] == "assistant")
    assert asst["arborist_meta"]["audit_mode"] == "STRICT"
    assert asst["arborist_meta"]["n_verified"] == 2
    assert asst["arborist_meta"]["cache_key"] == "abc123"


def test_session_id_used_as_filename(tmp_path):
    sid = new_session_id()
    with SessionWriter(first_prompt="q", root=tmp_path, session_id=sid):
        pass
    files = list(tmp_path.glob("**/*.jsonl"))
    assert files[0].stem == sid


def test_assistant_usage_block(tmp_path):
    usage = {"inputTokens": 100, "outputTokens": 50, "totalTokens": 150}
    with SessionWriter(first_prompt="q", root=tmp_path) as s:
        s.assistant_message("a", model="hermes-3", usage=usage, duration_ms=1234)
    rows = _read_lines(list(tmp_path.glob("**/*.jsonl"))[0])
    asst = next(r for r in rows if r["type"] == "message" and r["role"] == "assistant")
    assert asst["usage"] == usage
    assert asst["durationMs"] == 1234
    assert asst["model"] == "hermes-3"


def test_writer_is_idempotent_on_double_close(tmp_path):
    s = SessionWriter(first_prompt="q", root=tmp_path)
    s.close()
    s.close()  # should be a no-op, not raise