ai-agent/sidecar/fork.py

"""
``POST /fork`` — branch a new session off an existing one at a message
boundary, plus the transcript surgery it needs (``_fork_transcript``).
"""

import json
import pathlib

from fastapi import APIRouter, Header, HTTPException
from pydantic import BaseModel

from claude_accounts import _account_projects, _require_login, _valid_account
from config import PI_DEFAULT_MODEL, RUNNER_BIN
from launch import _claude_args, _launch, _remote_control, _run_cwd
from validate import _auth, _valid_harness, _valid_model, _valid_uuid

router = APIRouter()

class ForkBody(BaseModel):
    sessionId: str                    # fresh uuid for the fork (backend-minted)
    srcSessionId: str                 # the session to branch off
    prompt: str                       # the edited message — the fork's next turn
    cutLine: int                      # raw JSONL line index of the fork point
    cutText: str | None = None        # collapsed prefix of that message (canary)
    cutUserOrd: int | None = None     # its 0-based ordinal among user msgs (pi)
    model: str | None = None
    cwd: str | None = None            # must match the source session's cwd
    harness: str | None = None        # "claude" (default) | "pi"
    account: str | None = None        # the source session's account


def _collapse(s: str) -> str:
    return " ".join((s or "").split())


def _record_texts(line: str) -> list[str]:
    """Every text payload in one transcript record (message content as a plain
    string, or the text blocks of a content list) — what the fork-point canary
    compares against."""
    try:
        o = json.loads(line)
    except ValueError:
        return []
    c = (o.get("message") or {}).get("content")
    if isinstance(c, str):
        return [c]
    if isinstance(c, list):
        return [b.get("text", "") for b in c
                if isinstance(b, dict) and b.get("type") == "text"]
    return []


def _fork_transcript(src: pathlib.Path, dst: pathlib.Path, new_sid: str,
                     cut_line: int, cut_text: str | None) -> None:
    """Write the first ``cut_line`` lines of ``src`` as a new session file.

    This is the same shape ``claude --resume --fork-session`` produces — the
    identical records under a new ``sessionId`` (record uuids stay; they're
    only unique within a file's parent chain) — just cut short at the fork
    point instead of cloning the whole session. ``cut_text`` is the message
    the viewer showed at that ordinal: if the record at ``cut_line`` doesn't
    contain it, the archive the backend counted lines in has diverged from
    the live file (it never should — the import is append-only) and we refuse
    rather than fork at the wrong message."""
    try:
        lines = src.read_text(encoding="utf-8", errors="ignore") \
                   .splitlines(keepends=True)
    except OSError as e:
        raise HTTPException(502, f"cannot read source transcript: {e}")
    if not 0 < cut_line < len(lines):
        raise HTTPException(409, f"fork point out of range: line {cut_line} "
                                 f"of {len(lines)}")
    if cut_text:
        want = _collapse(cut_text)
        if not any(_collapse(t)[:len(want)] == want
                   for t in _record_texts(lines[cut_line])):
            raise HTTPException(409, "fork point mismatch — the transcript "
                                     "changed since the viewer loaded it; "
                                     "reload the conversation and retry")
    keep = lines[:cut_line]
    # The records just before a user message are that turn's bookkeeping
    # (queue-operation enqueue/dequeue carrying the message text, last-prompt).
    # They belong to the message being edited away, so a cut at the user
    # record would otherwise leak its text into the fork — trim them.
    def _rec_type(ln: str) -> str | None:
        try:
            return json.loads(ln).get("type")
        except ValueError:
            return None
    while keep and _rec_type(keep[-1]) in ("queue-operation", "last-prompt"):
        keep.pop()
    if not keep:
        raise HTTPException(400, "nothing before the fork point")
    out = []
    for ln in keep:
        s = ln.strip()
        if s.startswith("{"):
            try:
                o = json.loads(s)
                if "sessionId" in o:
                    o["sessionId"] = new_sid
                out.append(json.dumps(o, ensure_ascii=False) + "\n")
                continue
            except ValueError:
                pass
        out.append(ln if ln.endswith("\n") else ln + "\n")
    try:
        dst.write_text("".join(out), encoding="utf-8")
    except OSError as e:
        raise HTTPException(502, f"cannot write forked transcript: {e}")


@router.post("/fork")
def fork(body: ForkBody, authorization: str | None = Header(default=None)) -> dict:
    """Branch a new session off an existing one at a message boundary.

    Claude Code has no message-level fork (``--fork-session`` clones whole
    sessions), so: copy the source session file truncated at the fork point
    under a fresh session id (see ``_fork_transcript``), then ``claude -p
    <prompt> --resume <newSid>`` — claude loads the truncated history and
    appends the edited turn. The pi harness forks natively: pi sessions are
    trees, so the runner branches pi's own session file (``--fork-from``) and
    mirrors the truncated prefix itself."""
    _auth(authorization)
    prompt = (body.prompt or "").strip()
    if not prompt:
        raise HTTPException(400, "empty prompt")
    new_sid = _valid_uuid(body.sessionId)
    src_sid = _valid_uuid(body.srcSessionId)
    harness = _valid_harness(body.harness)
    model = _valid_model(body.model, harness)
    cwd = _run_cwd(body.cwd)

    if harness == "pi":
        model = model or PI_DEFAULT_MODEL
        args = [RUNNER_BIN, "-p", prompt, "--session-id", new_sid,
                "--fork-from", src_sid, "--cut-line", str(body.cutLine),
                "--cut-user-ord", str(body.cutUserOrd or 0), "--model", model]
        if body.cutText:
            args += ["--cut-text", body.cutText]
        return _launch(args, new_sid, cwd, model=model)

    # A fork stays on the source's account: the source transcript lives only
    # in that account's projects/, and so must the fork for --resume to find it.
    account = _valid_account(body.account)
    _require_login(account)
    projects = _account_projects(account)
    src = next(projects.glob(f"*/{src_sid}.jsonl"), None)
    if src is None:
        raise HTTPException(404, f"source session {src_sid} not found under "
                                 f"{projects}")
    dst = src.parent / f"{new_sid}.jsonl"
    _fork_transcript(src, dst, new_sid, body.cutLine, body.cutText)
    args = _claude_args(prompt, "--resume", new_sid)
    if model:
        args += ["--model", model]
    args += _remote_control(account)
    try:
        return _launch(args, new_sid, cwd, model=model, account=account)
    except HTTPException:
        dst.unlink(missing_ok=True)  # no headless truncated session left behind
        raise