zcode_to_opencode.py

#!/usr/bin/env python3
"""zcode -> opencode session migrator.

Formalizes the working import pipeline:
  1. Read session/message/part rows from zcode's sqlite db (read-only).
  2. Build a JSON doc in opencode's native export format, generating
     opencode-style ids (ses_/msg_/prt_ + us-timestamp low-48-bits hex + 14
     random base62 chars) and remapping parentID references.
  3. Run `opencode import <file>`.
  4. Rebind the imported row to its real project/directory/path using the
     move-session rules:
       - path = directory minus drive letter (forward slashes)
       - non-git dir  -> project 'global'
       - git dir      -> project resolved by precedence:
           existing project row with same worktree
           > sha1("git-remote:<normalized origin>")
           > .git/opencode cached id
           > root commit hash
         (project row auto-registered if missing)

Usage:
  python zcode_to_opencode.py list
  python zcode_to_opencode.py import <ordinal> [ordinal ...] [--with-children]
  python zcode_to_opencode.py import all
  python zcode_to_opencode.py status

Requires: python 3.7+ (stdlib only, no third-party deps), `opencode` on PATH,
         `git` (only for git-bound dirs).
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import secrets
import shutil
import sqlite3
import subprocess
import sys
import tempfile
import time
import urllib.request

ZDB = (
    os.environ.get("ZCODE_DB")
    or os.path.join(os.path.expanduser("~"), ".zcode", "cli", "db", "db.sqlite")
)


def default_opencode_db() -> str:
    local = os.path.expandvars(r"%LOCALAPPDATA%\opencode\opencode.db")
    if "LOCALAPPDATA" in os.environ and os.path.exists(local):
        return local
    base = os.environ.get("XDG_DATA_HOME",
                          os.path.join(os.path.expanduser("~"), ".local", "share"))
    return os.path.join(base, "opencode", "opencode.db")


ODB = os.environ.get("OPENCODE_DB") or default_opencode_db()
STATE_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "zcode-import-state.json")

ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

# part types opencode's schema understands; zcode extras (timeline,
# compaction, ...) are dropped on import
KNOWN_PART_TYPES = {"text", "reasoning", "tool", "step-start", "step-finish"}


# ---------------------------------------------------------------- utilities

def rand(n: int) -> str:
    return "".join(secrets.choice(ALPHABET) for _ in range(n))


def oc_id(prefix: str, ts_ms: int) -> str:
    return f"{prefix}_{ts_ms * 1000 & 0xFFFFFFFFFFFF:012x}{rand(14)}"


def norm(p: str | None) -> str | None:
    if not p:
        return p
    # Windows-style path (backslashes or a drive prefix): normalize to
    # forward slashes and unwrap a leading "/C:" into "C:/". POSIX paths
    # are left untouched (a single-letter top segment like "/a/..." must
    # not be mistaken for a drive letter).
    if "\\" in p or re.match(r"^[A-Za-z]:", p):
        return re.sub(r"^/([a-zA-Z])/", r"\1:/", p.replace("\\", "/"))
    return p


def strip_drive(p: str) -> str:
    return re.sub(r"^[A-Za-z]:/", "", p)


def sha1(text: str) -> str:
    return hashlib.sha1(text.encode()).hexdigest()


def slugify(title: str) -> str:
    s = re.sub(r"[^a-zA-Z0-9]+", "-", title).strip("-").lower()
    return (s[:48] or "zcode-import").rstrip("-")


def git(args: list[str], cwd: str) -> str | None:
    try:
        r = subprocess.run(["git", "-C", cwd] + args,
                           capture_output=True, text=True, timeout=15)
        return r.stdout.strip() or None if r.returncode == 0 else None
    except Exception:
        return None


def load_state() -> dict:
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, encoding="utf-8") as f:
            return json.load(f)
    return {"imported": {}}


def save_state(state: dict) -> None:
    with open(STATE_FILE, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


def find_opencode() -> str:
    for cand in ("opencode.exe", "opencode.cmd", "opencode"):
        p = shutil.which(cand)
        if p and not p.lower().endswith(".ps1"):
            return p
    raise SystemExit("ERR: opencode executable not found on PATH")


OC_EXE = find_opencode()


def oc_version() -> str:
    try:
        out = subprocess.run([OC_EXE, "--version"],
                             capture_output=True, text=True, timeout=20).stdout.strip()
        m = re.search(r"\d+\.\d+\.\d+", out)
        if m:
            return m.group(0)
    except Exception:
        pass
    raise SystemExit(f"ERR: cannot determine opencode version ({OC_EXE})")


# ------------------------------------------------------------- project logic

def resolve_project(directory_win: str) -> tuple[str, str]:
    """Return (project_id, path) for a windows directory, move-session rules."""
    d = norm(directory_win)  # forward slashes

    git_top = git(["rev-parse", "--show-toplevel"], d)
    if not git_top:
        return "global", strip_drive(d)

    git_top = norm(git_top)

    con = sqlite3.connect(f"file:{ODB}?mode=ro", uri=True)
    for pid, wt in con.execute("SELECT id, worktree FROM project WHERE id != 'global'"):
        if norm(wt) == git_top:
            con.close()
            return pid, _rel(git_top, d)
    con.close()

    # precedence: origin remote > .git/opencode cache > root commit
    remote = git(["remote", "get-url", "origin"], d)
    pid = None
    if remote:
        m = re.match(r"^(?:[^@/:]+@)?([^/:]+):(.+)$", remote.strip())
        if m:
            key = f"{m.group(2).strip('/')}".rstrip("/")
        else:
            try:
                u = urllib.parse.urlparse(remote.strip())
                key = f"{u.hostname}/{u.path.strip('/')}"
            except Exception:
                key = None
        if key:
            pid = sha1(f"git-remote:{key}")
    if not pid:
        common = git(["rev-parse", "--git-common-dir"], d)
        if common:
            cache = os.path.join(common if os.path.isabs(common) else os.path.join(d, common),
                                 "opencode")
            if os.path.isfile(cache):
                pid = open(cache, encoding="utf-8").read().strip() or None
    if not pid:
        pid = git(["rev-list", "--max-parents=0", "HEAD"], d)

    if pid:
        register_project(pid, git_top)
        return pid, _rel(git_top, d)
    return "global", strip_drive(d)


def _rel(worktree: str, directory: str) -> str:
    nwt, nd = norm(worktree), norm(directory)
    if nd == nwt:
        return ""
    if nd.startswith(nwt + "/"):
        return nd[len(nwt) + 1:]
    return strip_drive(nd)


def register_project(pid: str, worktree: str) -> None:
    con = sqlite3.connect(ODB, timeout=30)
    con.execute("PRAGMA busy_timeout=30000")
    now = int(time.time() * 1000)
    con.execute(
        "INSERT OR IGNORE INTO project (id, worktree, vcs, time_created, time_updated, sandboxes)"
        " VALUES (?, ?, 'git', ?, ?, '[]')",
        (pid, worktree, now, now),
    )
    con.commit()
    con.close()


def rebind_session(new_sid: str, directory_win: str) -> tuple[str, str]:
    project_id, rel_path = resolve_project(directory_win)
    con = sqlite3.connect(ODB, timeout=30)
    con.execute("PRAGMA busy_timeout=30000")
    con.execute(
        "UPDATE session SET project_id=?, directory=?, path=? WHERE id=?",
        (project_id, norm(directory_win), rel_path, new_sid),
    )
    con.commit()
    row = con.execute("SELECT project_id, directory, path FROM session WHERE id=?",
                      (new_sid,)).fetchone()
    con.close()
    return row


# ------------------------------------------------------------------ convert

def fetch_session(zc: sqlite3.Connection, sid: str):
    meta = zc.execute("SELECT * FROM session WHERE id=?", (sid,)).fetchone()
    msgs = zc.execute(
        "SELECT * FROM message WHERE session_id=? ORDER BY COALESCE(sequence,0), time_created, id",
        (sid,),
    ).fetchall()
    parts = zc.execute(
        "SELECT * FROM part WHERE session_id=? ORDER BY message_id, COALESCE(sequence,0), time_created, id",
        (sid,),
    ).fetchall()
    return meta, msgs, parts


def build_doc(meta, msgs, parts, version: str, new_parent: str | None):
    by_msg: dict[str, list] = {}
    for p in parts:
        by_msg.setdefault(p["message_id"], []).append(json.loads(p["data"]))

    new_sid = oc_id("ses", meta["time_created"])
    id_map = {m["id"]: oc_id("msg", m["time_created"]) for m in msgs}

    doc = {
        "info": {
            "id": new_sid,
            "slug": slugify(meta["title"]),
            "title": meta["title"],
            "version": version,
            "cost": 0,
            "tokens": {"input": 0, "output": 0, "reasoning": 0,
                       "cache": {"read": 0, "write": 0}},
            "time": {"created": meta["time_created"],
                     "updated": max(meta["time_updated"], meta["time_created"])},
        },
        "messages": [],
    }
    # only true subagent children stay nested; user-facing sessions
    # (interactive + fork) must be roots to appear in the session list,
    # matching opencode's own fork behavior (parent_id = NULL)
    if new_parent and meta["task_type"] == "subagent_child":
        doc["info"]["parentID"] = new_parent

    for m in msgs:
        info = json.loads(m["data"])
        # zcode error objects (e.g. AiSdkModelAdapterError) fail opencode's
        # schema validation on import; they only mark failed attempts.
        info.pop("error", None)
        mid = id_map[m["id"]]
        info["id"] = mid
        info["sessionID"] = new_sid
        if "parentID" in info:
            info["parentID"] = id_map.get(info["parentID"], info["parentID"])
        plist = [pd for pd in by_msg.get(m["id"], [])
                 if pd.get("type") in KNOWN_PART_TYPES]
        for pd in plist:
            pd["id"] = oc_id("prt", m["time_created"])
            pd["sessionID"] = new_sid
            pd["messageID"] = mid
        doc["messages"].append({"info": info, "parts": plist})

    return doc, new_sid


def opencode_import(doc: dict, new_sid: str) -> str:
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False,
                                      encoding="utf-8")
    json.dump(doc, tmp, ensure_ascii=False)
    tmp.close()
    try:
        r = subprocess.run([OC_EXE, "import", tmp.name],
                           capture_output=True, text=True, timeout=300)
        out = (r.stdout or "") + (r.stderr or "")
        m = re.search(r"Imported session:\s*(ses_\S+)", out)
        if not m:
            # the importer writes incrementally: a failed run can leave a
            # partial session behind under OUR generated id — clean it up
            con = sqlite3.connect(ODB, timeout=30)
            hit = con.execute("SELECT COUNT(*) FROM session WHERE id=?",
                              (new_sid,)).fetchone()[0]
            con.close()
            if hit:
                subprocess.run([OC_EXE, "session", "delete", new_sid],
                               capture_output=True, text=True, timeout=120)
                print(f"  cleaned up partial session {new_sid}")
            raise SystemExit(f"opencode import failed:\n{out[-2000:]}")
        return m.group(1)
    finally:
        os.unlink(tmp.name)


# ------------------------------------------------------------------ commands

def cmd_list(zc: sqlite3.Connection, state: dict) -> None:
    print(f"{'#':>3}  {'when':16}  {'msgs':>4}  {'imp':>3}  title / id")
    rows = zc.execute("SELECT * FROM session ORDER BY time_created").fetchall()
    for i, s in enumerate(rows, 1):
        n = zc.execute("SELECT COUNT(*) FROM message WHERE session_id=?",
                       (s["id"],)).fetchone()[0]
        when = time.strftime("%Y-%m-%d %H:%M", time.localtime(s["time_created"] / 1000))
        done = "yes" if s["id"] in state["imported"] else "-"
        title = (s["title"] or "")[:46]
        parent = f" [child of {s['parent_id'][:13]}]" if s["parent_id"] else ""
        print(f"{i:>3}  {when}  {n:>4}  {done:>3}  {title}{parent}")
        print(f"{'':17}  {s['id']}")


def cmd_import(zc: sqlite3.Connection, state: dict, targets: list[str],
               with_children: bool, version: str) -> None:
    ok, fail = [], []
    queue = list(targets)
    seen: set[str] = set()

    while queue:
        sid = queue.pop(0)
        if sid in seen:
            continue
        seen.add(sid)

        if sid in state["imported"]:
            print(f"skip (already imported as {state['imported'][sid]}): {sid}")
            continue

        try:
            meta, msgs, parts = fetch_session(zc, sid)
            if not meta:
                print(f"ERR: not found in zcode db: {sid}")
                fail.append(sid)
                continue

            parent_new = None
            needs_parent = (meta["parent_id"]
                            and meta["task_type"] == "subagent_child")
            if needs_parent:
                parent_new = state["imported"].get(meta["parent_id"])
                if not parent_new:
                    print(f"  parent {meta['parent_id']} not imported yet; importing first...")
                    queue.insert(0, sid)
                    queue.insert(0, meta["parent_id"])
                    seen.discard(sid)
                    continue

            label = f"{meta['title'][:40]} ({sid[:18]}...)"
            doc, new_sid = build_doc(meta, msgs, parts, version, parent_new)
            opencode_import(doc, new_sid)

            # persist state immediately: opencode import is already durable at
            # this point; a crash during rebind must not duplicate the session
            state["imported"][sid] = new_sid
            save_state(state)
            ok.append((sid, new_sid))

            project_id, final_dir, rel_path = rebind_session(new_sid, meta["directory"])
            print(f"imported: {label}")
            print(f"          -> {new_sid}  project={project_id}  path={rel_path}")

            if with_children:
                kids = zc.execute(
                    "SELECT id FROM session WHERE parent_id=? ORDER BY time_created",
                    (sid,),
                ).fetchall()
                for (kid,) in kids:
                    if kid not in seen:
                        queue.append(kid)
        except SystemExit as e:
            print(f"ERR importing {sid}: {e}")
            fail.append(sid)
        except Exception as e:
            print(f"ERR importing {sid}: {e!r}")
            fail.append(sid)

    print(f"\ndone: {len(ok)} imported, {len(fail)} failed")


def main() -> None:
    if len(sys.argv) < 2 or sys.argv[1] not in {"list", "import", "status"}:
        print(__doc__)
        sys.exit(1)

    if not os.path.exists(ZDB):
        raise SystemExit(f"ERR: zcode db not found: {ZDB}")
    if not os.path.exists(ODB):
        raise SystemExit(f"ERR: opencode db not found: {ODB}")

    state = load_state()
    zc = sqlite3.connect(f"file:{ZDB}?mode=ro", uri=True)
    zc.row_factory = sqlite3.Row
    cmd = sys.argv[1]

    if cmd == "list":
        cmd_list(zc, state)
    elif cmd == "status":
        print(json.dumps(state, ensure_ascii=False, indent=2))
    elif cmd == "import":
        args = sys.argv[2:]
        with_children = "--with-children" in args
        args = [a for a in args if a != "--with-children"]
        version = oc_version()
        print(f"opencode version: {version}")

        rows = zc.execute("SELECT id FROM session ORDER BY time_created").fetchall()
        all_ids = [r["id"] for r in rows]

        targets: list[str] = []
        for a in args:
            if a == "all":
                targets.extend(all_ids)
            elif a.isdigit():
                idx = int(a)
                if not 1 <= idx <= len(all_ids):
                    raise SystemExit(f"ERR: ordinal {idx} out of range (1..{len(all_ids)})")
                targets.append(all_ids[idx - 1])
            elif re.fullmatch(r"sess_\S+|ses_\S+", a):
                targets.append(a)
            else:
                raise SystemExit(f"ERR: bad argument: {a}")
        if not targets:
            raise SystemExit("nothing to import (see: list)")

        cmd_import(zc, state, targets, with_children, version)

    zc.close()


if __name__ == "__main__":
    main()