A helper script for syncing or exporting Codex sessions

codex_helper.py

#!/usr/bin/env python3
"""Codex session export/sync tool."""

import json
import shlex
from pathlib import Path

import typer

app = typer.Typer(help="Export OpenAI Codex sessions.")

CODEX_DIR = Path.home() / ".codex"
SESSION_INDEX = CODEX_DIR / "session_index.jsonl"
SESSIONS_DIR = CODEX_DIR / "sessions"
ARCHIVED_SESSIONS_DIR = CODEX_DIR / "archived_sessions"
SESSION_DIRS = (SESSIONS_DIR, ARCHIVED_SESSIONS_DIR)


def parse_jsonl(path: Path) -> list[dict]:
    records = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    return records


def _session_ids_for_file(path: Path) -> set[str]:
    ids = set()
    stem = path.stem
    if stem:
        ids.add(stem)
        if stem.startswith("rollout-"):
            parts = stem.split("-", 3)
            if len(parts) == 4:
                ids.add(parts[3])

    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                continue
            payload = record.get("payload", {})
            if record.get("type") == "session_meta" and isinstance(payload, dict):
                sid = payload.get("id")
                if sid:
                    ids.add(str(sid))
                    break
    return ids


def _matches_session_id(candidate: str, query: str) -> bool:
    candidate = candidate.strip()
    query = query.strip()
    if not candidate or not query:
        return False
    return candidate == query or candidate.endswith(query) or query in candidate


def _find_session_file(session_id: str, roots: tuple[Path, ...] = SESSION_DIRS) -> Path | None:
    existing_roots = [root for root in roots if root.exists()]

    # Fast path: modern Codex filenames include the session id suffix.
    filename_matches: list[Path] = []
    for root in existing_roots:
        filename_matches.extend(root.rglob(f"*{session_id}*.jsonl"))
    if len(filename_matches) == 1:
        return filename_matches[0]
    if len(filename_matches) > 1:
        exact = [
            path for path in filename_matches
            if any(_matches_session_id(sid, session_id) for sid in _session_ids_for_file(path))
        ]
        if len(exact) == 1:
            return exact[0]
        raise typer.BadParameter(
            f"Session id {session_id!r} matched multiple files:\n"
            + "\n".join(f"  {path}" for path in filename_matches[:20])
        )

    # Fallback: scan metadata so older or moved files still work.
    for root in existing_roots:
        for candidate in root.rglob("*.jsonl"):
            try:
                ids = _session_ids_for_file(candidate)
            except OSError:
                continue
            if any(_matches_session_id(sid, session_id) for sid in ids):
                return candidate
    return None


def _session_roots_description() -> str:
    return ", ".join(str(path) for path in SESSION_DIRS)


@app.command()
def list_sessions(
    limit: int = typer.Option(20, "--limit", "-n", help="Max sessions to show."),
    host: str = typer.Option(None, "--host", help="Read from remote host instead of local."),
):
    """List recent sessions from session_index.jsonl."""
    import subprocess

    if host:
        result = subprocess.run(
            ["ssh", host, f"cat {SESSION_INDEX}"],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            typer.echo(f"Failed to read remote index:\n{result.stderr}", err=True)
            raise typer.Exit(1)
        records = []
        for line in result.stdout.splitlines():
            line = line.strip()
            if line:
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    else:
        if not SESSION_INDEX.exists():
            typer.echo(f"Index not found: {SESSION_INDEX}", err=True)
            raise typer.Exit(1)
        records = parse_jsonl(SESSION_INDEX)
    records.sort(key=lambda r: r.get("updated_at", ""), reverse=True)
    records = records[:limit]

    if not records:
        typer.echo("No sessions found.")
        return

    id_w = max(len(r.get("id", "")) for r in records)
    id_w = max(id_w, 8)

    typer.echo(f"{'ID':<{id_w}}  {'UPDATED_AT':<26}  THREAD_NAME")
    typer.echo("-" * (id_w + 2 + 26 + 2 + 40))
    for r in records:
        sid = r.get("id", "?")
        updated = r.get("updated_at", "")
        name = r.get("thread_name", "(untitled)")
        typer.echo(f"{sid:<{id_w}}  {updated:<26}  {name}")


@app.command()
def export(
    session_id: str = typer.Option(..., "--id", help="Session ID to export."),
    output: Path = typer.Option(None, "--output", "-o", help="Output .md file (default: stdout)."),
    include_events: bool = typer.Option(False, "--include-events", help="Include function_call / token_count events."),
):
    """Export a session to Markdown."""
    if not any(root.exists() for root in SESSION_DIRS):
        typer.echo(f"No Codex session directories found under: {_session_roots_description()}", err=True)
        raise typer.Exit(1)

    target_file = _find_session_file(session_id)

    if not target_file:
        typer.echo(
            f"Session not found: {session_id}\nSearched: {_session_roots_description()}",
            err=True,
        )
        raise typer.Exit(1)

    records = parse_jsonl(target_file)
    lines = _render_markdown(session_id, records, include_events)
    content = "\n".join(lines)

    if output:
        output.write_text(content, encoding="utf-8")
        typer.echo(f"Exported to {output}")
    else:
        typer.echo(content)


def _render_markdown(session_id: str, records: list[dict], include_events: bool) -> list[str]:
    out = []

    # --- Header from session_meta ---
    meta_payload = {}
    for r in records:
        if r.get("type") == "session_meta":
            meta_payload = r.get("payload", {})
            break

    started = meta_payload.get("timestamp", "")
    cwd = meta_payload.get("cwd", "")
    cli_version = meta_payload.get("cli_version", "")
    source = meta_payload.get("source", "")

    out.append(f"# Codex Session: `{session_id}`")
    out.append("")
    if started:
        out.append(f"- **Started:** {started}")
    if cwd:
        out.append(f"- **Working dir:** `{cwd}`")
    if cli_version:
        out.append(f"- **CLI version:** {cli_version}")
    if source:
        out.append(f"- **Source:** {source}")
    out.append("")
    out.append("---")
    out.append("")

    # --- Messages ---
    ROLE_LABEL = {
        "user": "**User**",
        "assistant": "**Assistant**",
        "developer": "**Developer**",
    }

    for r in records:
        rtype = r.get("type")
        payload = r.get("payload", {})
        ts = r.get("timestamp", "")

        if rtype == "response_item":
            role = payload.get("role", "")
            label = ROLE_LABEL.get(role, f"**{role}**") if role else "**Message**"
            content_blocks = payload.get("content", [])
            text = _extract_text(content_blocks)
            if text:
                ts_note = f" <sup>{ts}</sup>" if ts else ""
                out.append(f"### {label}{ts_note}")
                out.append("")
                out.append(text)
                out.append("")

        elif rtype == "event_msg" and include_events:
            etype = payload.get("type", "")
            if etype in ("function_call", "user_message", "agent_message", "token_count"):
                ts_note = f" `{ts}`" if ts else ""
                out.append(f"<!-- event:{etype}{ts_note} -->")
                out.append(f"```json\n{json.dumps(payload, ensure_ascii=False, indent=2)}\n```")
                out.append("")

    return out


def _extract_text(content_blocks) -> str:
    """Extract readable text from a content block array."""
    if not content_blocks:
        return ""
    if isinstance(content_blocks, str):
        return content_blocks

    parts = []
    for block in content_blocks:
        if isinstance(block, str):
            parts.append(block)
        elif isinstance(block, dict):
            for key in ("input_text", "output_text", "text"):
                val = block.get(key)
                if val:
                    parts.append(str(val))
                    break
    return "\n\n".join(p.strip() for p in parts if p.strip())


@app.command()
def sync(
    session_id: str = typer.Option(..., "--id", help="Session ID to sync."),
    host: str = typer.Option(..., "--host", help="Remote host (user@host or host)."),
    dry_run: bool = typer.Option(False, "--dry-run", help="Print actions without executing."),
):
    """Sync a session from a remote host to this machine."""
    import subprocess

    remote_index = f"{host}:{SESSION_INDEX}"

    def run(cmd: list[str]) -> subprocess.CompletedProcess:
        typer.echo(f"  $ {' '.join(cmd)}")
        if dry_run:
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")
        return subprocess.run(cmd, capture_output=True, text=True)

    # --- Step 1: verify local dirs exist ---
    typer.echo("Checking local directories...")
    for d in (CODEX_DIR,):
        if not d.exists():
            typer.echo(f"  Directory not found: {d}", err=True)
            raise typer.Exit(1)
        typer.echo(f"  OK {d}")
    for d in SESSION_DIRS:
        if not dry_run:
            d.mkdir(parents=True, exist_ok=True)
        typer.echo(f"  OK {d}")

    # --- Step 2: fetch remote index, find the session entry ---
    typer.echo(f"\nFetching remote index from {remote_index} ...")
    if not dry_run:
        result = subprocess.run(
            ["ssh", host, f"cat {SESSION_INDEX}"],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            typer.echo(f"Failed to read remote index:\n{result.stderr}", err=True)
            raise typer.Exit(1)

        remote_records = []
        for line in result.stdout.splitlines():
            line = line.strip()
            if line:
                try:
                    remote_records.append(json.loads(line))
                except json.JSONDecodeError:
                    pass

        entry = next((r for r in remote_records if r.get("id") == session_id), None)
        if not entry:
            typer.echo(f"Session {session_id!r} not found in remote index.", err=True)
            raise typer.Exit(1)
        typer.echo(f"  Found: {entry.get('thread_name', '(untitled)')}  {entry.get('updated_at', '')}")
    else:
        typer.echo(f"  [dry-run] would search remote index for id={session_id}")
        entry = {"id": session_id}

    # --- Step 3: find the session file on the remote ---
    typer.echo(f"\nLocating session file on {host} ...")
    if not dry_run:
        quoted_id = shlex.quote(session_id)
        quoted_roots = " ".join(shlex.quote(str(root)) for root in SESSION_DIRS)
        result = subprocess.run(
            [
                "ssh",
                host,
                (
                    f"find {quoted_roots} -type f -name '*'{quoted_id}'*.jsonl' -print "
                    "2>/dev/null | head -n 1"
                ),
            ],
            capture_output=True, text=True,
        )
        remote_paths = [p.strip() for p in result.stdout.splitlines() if p.strip()]
        if not remote_paths:
            result = subprocess.run(
                [
                    "ssh",
                    host,
                    (
                        f"grep -Erl '\"id\"[[:space:]]*:[[:space:]]*\"{session_id}\"' "
                        f"{quoted_roots} 2>/dev/null | head -n 1"
                    ),
                ],
                capture_output=True, text=True,
            )
            remote_paths = [p.strip() for p in result.stdout.splitlines() if p.strip()]
        if not remote_paths:
            typer.echo(
                f"No session file found for session {session_id} on remote.\n"
                f"Searched: {_session_roots_description()}",
                err=True,
            )
            raise typer.Exit(1)
        remote_file = remote_paths[0]
        typer.echo(f"  Remote file: {remote_file}")
        if remote_file.startswith(str(SESSIONS_DIR) + "/"):
            local_file = SESSIONS_DIR / Path(remote_file).relative_to(SESSIONS_DIR)
        elif remote_file.startswith(str(ARCHIVED_SESSIONS_DIR) + "/"):
            local_file = ARCHIVED_SESSIONS_DIR / Path(remote_file).relative_to(ARCHIVED_SESSIONS_DIR)
        else:
            local_file = ARCHIVED_SESSIONS_DIR / Path(remote_file).name
    else:
        remote_file = str(SESSIONS_DIR / "YYYY" / "MM" / "DD" / f"rollout-...-{session_id}.jsonl")
        local_file = SESSIONS_DIR / "YYYY" / "MM" / "DD" / f"rollout-...-{session_id}.jsonl"
        typer.echo(f"  [dry-run] would locate remote file, e.g. {remote_file}")

    # --- Step 4: check if already exists locally ---
    typer.echo(f"\nChecking local session file {local_file} ...")
    if not dry_run and local_file.exists():
        typer.echo("  Already exists locally, skipping scp.")
    else:
        if dry_run:
            typer.echo(f"  [dry-run] would scp {host}:{remote_file} -> {local_file}")
        else:
            local_file.parent.mkdir(parents=True, exist_ok=True)
            typer.echo(f"  Copying {host}:{remote_file} -> {local_file}")
        result = run(["scp", f"{host}:{remote_file}", str(local_file)])
        if not dry_run and result.returncode != 0:
            typer.echo(f"scp failed:\n{result.stderr}", err=True)
            raise typer.Exit(1)

    # --- Step 5: merge index entry ---
    typer.echo(f"\nMerging index entry into {SESSION_INDEX} ...")
    if not dry_run:
        existing_ids: set[str] = set()
        if SESSION_INDEX.exists():
            for r in parse_jsonl(SESSION_INDEX):
                if "id" in r:
                    existing_ids.add(r["id"])

        if session_id in existing_ids:
            typer.echo("  Already in local index, skipping.")
        else:
            with SESSION_INDEX.open("a", encoding="utf-8") as f:
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
            typer.echo("  Appended.")
    else:
        typer.echo(f"  [dry-run] would append id={session_id} to local index if not present")

    typer.echo("\nDone." if not dry_run else "\nDry run complete.")


if __name__ == "__main__":
    app()