Copy open Codex rollouts to closed objects before a deploy

A restart loses whatever the FUSE writer has not flushed. Snapshot each open
rollout out through the bucket API, verify it by readback, and reconcile on the
way up.

scripts/migrate-open-rollouts.sh

@@ -0,0 +1,330 @@
+#!/usr/bin/env bash
+# Get every open Codex rollout onto the bucket as a CLOSED object, before a
+# deploy replaces the container.
+#
+#   scripts/migrate-open-rollouts.sh                # snapshot, upload, verify
+#   scripts/migrate-open-rollouts.sh --dry-run      # report, touch nothing
+#   scripts/migrate-open-rollouts.sh --restore      # boot side: reconcile
+#
+# WHY THIS EXISTS
+#
+# Codex appends to $CODEX_HOME/sessions/YYYY/MM/DD/rollout-<ts>-<uuid>.jsonl and
+# holds that descriptor open for the whole life of a resumed session. That path
+# is a symlink onto the /data bucket (entrypoint.sh), which is FUSE over object
+# storage, and the mount's streaming writer can hold the entire open epoch in
+# memory until close. Reads inside this container see the buffer, so the file
+# looks complete from here. The OBJECT does not have it.
+#
+# Measured on prod, 2026-08-13, three live Codex sessions:
+#
+#     rollout-2026-08-05T14-33-11-…   local 9 671 486 B   object     767 B
+#     rollout-2026-08-07T15-30-46-…   local 8 051 317 B   object     767 B
+#     rollout-2026-08-13T14-09-11-…   local   971 492 B   object  absent
+#
+# 767 bytes is the session header written at open. Everything after it exists
+# only in this container. Kill the container and ~18 MB of conversation across
+# three agents is gone — which is exactly the loss mode PR #25 fixes, and
+# exactly what deploying that fix would trigger one last time on the way in.
+#
+# So: copy the writer's own view out through the bucket API (HTTP, not the
+# mount), verify it landed, and only then let the deploy proceed.
+#
+# WHAT IT DOES NOT DO
+#
+# It does not stop, signal, or touch the Codex processes. It reads. A rollout is
+# append-only JSONL, so a snapshot taken while the writer runs is a valid prefix
+# of the session — we trim a partial trailing line and record how many bytes
+# that cost. Losing the last half-written line is not the failure mode anyone is
+# worried about here.
+#
+# It does not write the canonical path. If it did, the dying writer could flush
+# its 767-byte view back over a full copy — the clobber this is meant to
+# prevent. Copies go to a separate migration prefix, and --restore reconciles
+# them on the way up, once no writer holds the file.
+set -euo pipefail
+
+MODE=snapshot
+DRY=0
+for a in "$@"; do
+  case "$a" in
+    --restore)  MODE=restore ;;
+    --dry-run)  DRY=1 ;;
+    -h|--help)  sed -n '2,6p' "$0"; exit 0 ;;
+    *) echo "unknown argument: $a" >&2; exit 2 ;;
+  esac
+done
+
+[ -n "${HF_TOKEN:-}" ] || { echo "HF_TOKEN is not set" >&2; exit 1; }
+
+DATA_DIR="${DATA_DIR:-/data}"
+CODEX_DURABLE="${CODEX_DURABLE:-$DATA_DIR/state/codex}"
+# The dev deploy script names a Space's bucket "<space>-data"; prod follows the
+# same rule (lvwerra/agent-manager -> lvwerra/agent-manager-data). Override with
+# AM_BUCKET when running against something else.
+BUCKET="${AM_BUCKET:-${SPACE_ID:?SPACE_ID unset and AM_BUCKET not given}-data}"
+# Staging is local POSIX disk, never the bucket: the whole point is to hold a
+# closed byte-exact copy somewhere the FUSE writer has no opinion about.
+STAGE="${AM_LOCAL:-/tmp}/rollout-migration"
+
+export CODEX_DURABLE BUCKET STAGE MODE DRY
+
+python3 - <<'PY'
+import hashlib, json, os, re, shutil, sys, time
+from pathlib import Path
+
+from huggingface_hub import (
+    download_bucket_files,
+    list_bucket_tree,
+    sync_bucket,
+)
+
+TOKEN   = os.environ["HF_TOKEN"]
+BUCKET  = os.environ["BUCKET"]
+DURABLE = Path(os.environ["CODEX_DURABLE"])
+STAGE   = Path(os.environ["STAGE"])
+DRY     = os.environ["DRY"] == "1"
+RESTORE = os.environ["MODE"] == "restore"
+
+DATA_DIR = Path(os.environ.get("DATA_DIR", "/data"))
+PREFIX   = "state/codex/migration"          # remote home for closed copies
+ROLLOUT  = re.compile(r"/sessions/.*/rollout-[^/]+\.jsonl$")
+
+
+def remote_of(p: Path) -> str:
+    """Bucket key for a path under the mount: /data/state/x -> state/x."""
+    return str(p.relative_to(DATA_DIR))
+
+
+def object_size(key: str):
+    """Size of one bucket object, or None if there is no object.
+
+    Deliberately NOT get_bucket_file_metadata(): on these xet-backed objects its
+    .size is wrong. Measured 2026-08-13 against known-length uploads — a
+    9 796 319 B rollout reported 767, a 2 329 B manifest reported 694 — while
+    list_bucket_tree reported both correctly and download_bucket_files returned
+    byte-exact content. Do not "simplify" this back to a metadata call: every
+    number this script prints, and its entire safe/unsafe verdict, depends on
+    the size being real.
+    """
+    parent, _, name = key.rpartition("/")
+    for f in list_bucket_tree(BUCKET, parent, recursive=False, token=TOKEN):
+        if Path(f.path).name == name:
+            return getattr(f, "size", None)
+    return None
+
+
+def open_rollouts():
+    """Every rollout held open for WRITING, with the writer's own file position.
+
+    /proc is the only honest source here. An agent's rollout is identified by
+    the descriptor a live Codex process holds, not by anything on disk: mtime
+    on the mount is stale (it tracks the last object commit, not the last
+    append), so a find -newermt sweep misses precisely the files at risk.
+    """
+    out = {}
+    for pid_dir in Path("/proc").iterdir():
+        if not pid_dir.name.isdigit():
+            continue
+        fd_dir = pid_dir / "fd"
+        try:
+            fds = list(fd_dir.iterdir())
+        except OSError:
+            continue                        # not ours, or exited mid-scan
+        for fd in fds:
+            try:
+                target = os.readlink(fd)
+            except OSError:
+                continue
+            if not ROLLOUT.search(target):
+                continue
+            try:
+                info = (pid_dir / "fdinfo" / fd.name).read_text()
+            except OSError:
+                continue
+            flags = pos = 0
+            for line in info.splitlines():
+                if line.startswith("flags:"):
+                    flags = int(line.split()[1], 8)
+                elif line.startswith("pos:"):
+                    pos = int(line.split()[1])
+            # O_RDONLY is 0 in the low two bits. Readers (the manager parsing a
+            # trace, a repin hook) are not at risk and must not be snapshotted
+            # at their seek position.
+            if flags & 0o3 == 0:
+                continue
+            prev = out.get(target)
+            if prev is None or pos > prev["pos"]:
+                out[target] = {"pos": pos, "pid": int(pid_dir.name)}
+    return out
+
+
+def trim_to_last_record(raw: bytes):
+    """Drop a half-written trailing line. Returns (kept, dropped, bad_lines)."""
+    cut = raw.rfind(b"\n")
+    kept, dropped = (raw[: cut + 1], len(raw) - cut - 1) if cut >= 0 else (b"", len(raw))
+    bad = 0
+    for line in kept.splitlines():
+        if not line.strip():
+            continue
+        try:
+            json.loads(line)
+        except Exception:
+            bad += 1
+    return kept, dropped, bad
+
+
+def snapshot():
+    live = open_rollouts()
+    if not live:
+        print("no Codex rollout is open for writing — nothing to migrate")
+        return 0
+
+    run_id = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
+    run_dir = STAGE / run_id
+    entries, concerns = [], 0
+
+    print(f"==> {len(live)} open rollout(s); bucket {BUCKET}")
+    for path, fd in sorted(live.items()):
+        p = Path(path)
+        raw = p.read_bytes()                # the writer's view, buffer included
+        kept, dropped, bad = trim_to_last_record(raw)
+        digest = hashlib.sha256(kept).hexdigest()
+
+        remote_size = object_size(remote_of(p))
+
+        at_risk = len(kept) - (remote_size or 0)
+        staged = run_dir / remote_of(p)
+        if not DRY:
+            staged.parent.mkdir(parents=True, exist_ok=True)
+            staged.write_bytes(kept)
+
+        entries.append({
+            "path": str(p),
+            "remote": remote_of(p),
+            "pid": fd["pid"],
+            "writer_pos": fd["pos"],
+            "bytes": len(kept),
+            "sha256": digest,
+            "partial_tail_dropped": dropped,
+            "invalid_lines": bad,
+            "object_bytes_before": remote_size,
+            "bytes_at_risk": at_risk,
+        })
+        if bad:
+            concerns += 1
+        flag = f"  !! {bad} unparseable line(s)" if bad else ""
+        print(f"    {p.name}")
+        print(f"      local {len(kept):>10}  object {str(remote_size):>10}"
+              f"  at risk {at_risk:>10}  dropped {dropped}{flag}")
+
+    if DRY:
+        print("==> --dry-run: nothing staged, nothing uploaded")
+        return 0
+
+    manifest = {
+        "run_id": run_id,
+        "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "bucket": BUCKET,
+        "space": os.environ.get("SPACE_ID"),
+        "entries": entries,
+    }
+    (run_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
+
+    dest = f"hf://buckets/{BUCKET}/{PREFIX}/{run_id}"
+    print(f"==> uploading {run_dir} -> {dest}")
+    sync_bucket(source=str(run_dir), dest=dest, token=TOKEN, quiet=True)
+
+    # Verify by reading the objects back, not by trusting the writer. A copy
+    # nobody checked is the same bet we are trying to get off.
+    print("==> verifying")
+    for e in entries:
+        key = f"{PREFIX}/{run_id}/{e['remote']}"
+        size = object_size(key)
+        if size is None:
+            print(f"    !! {e['remote']}: no object after upload")
+            concerns += 1
+            continue
+        if size != e["bytes"]:
+            print(f"    !! {e['remote']}: object {size} B, staged {e['bytes']} B")
+            concerns += 1
+            continue
+        check = run_dir / "verify" / e["remote"]
+        check.parent.mkdir(parents=True, exist_ok=True)
+        download_bucket_files(BUCKET, [(key, str(check))], token=TOKEN)
+        got = hashlib.sha256(check.read_bytes()).hexdigest()
+        if got != e["sha256"]:
+            print(f"    !! {e['remote']}: sha256 mismatch on readback")
+            concerns += 1
+        else:
+            print(f"    ok {Path(e['remote']).name}  {e['bytes']} B  {got[:12]}")
+        check.unlink(missing_ok=True)
+    shutil.rmtree(run_dir / "verify", ignore_errors=True)
+
+    print(f"==> run {run_id} at hf://buckets/{BUCKET}/{PREFIX}/{run_id}")
+    if concerns:
+        print(f"!!  {concerns} problem(s) — do NOT restart until these are understood",
+              file=sys.stderr)
+        return 1
+    print("==> every open rollout is a verified closed object; safe to deploy")
+    return 0
+
+
+def restore():
+    """Boot side. Put back anything the dying writer failed to flush.
+
+    Runs before agents start, when no descriptor is held. A migrated copy is
+    only ever restored when it is strictly longer than what is on the canonical
+    path, so a container that shut down cleanly — and therefore flushed a
+    complete rollout — is left completely alone.
+    """
+    runs = sorted(
+        {Path(f.path).parts[3] for f in list_bucket_tree(BUCKET, PREFIX, recursive=True, token=TOKEN)
+         if Path(f.path).name == "manifest.json"}
+    )
+    if not runs:
+        print("no migration runs on the bucket — nothing to restore")
+        return 0
+
+    run_id = runs[-1]
+    local_manifest = STAGE / f"manifest-{run_id}.json"
+    local_manifest.parent.mkdir(parents=True, exist_ok=True)
+    download_bucket_files(BUCKET, [(f"{PREFIX}/{run_id}/manifest.json", str(local_manifest))],
+                          token=TOKEN)
+    manifest = json.loads(local_manifest.read_text())
+    print(f"==> newest migration run {run_id}, {len(manifest['entries'])} rollout(s)")
+
+    restored = failed = 0
+    for e in manifest["entries"]:
+        target = Path(e["path"])
+        have = target.stat().st_size if target.exists() else 0
+        if have >= e["bytes"]:
+            print(f"    skip {target.name}: on disk {have} B >= migrated {e['bytes']} B")
+            continue
+        if DRY:
+            print(f"    would restore {target.name}: {have} B -> {e['bytes']} B")
+            restored += 1
+            continue
+        target.parent.mkdir(parents=True, exist_ok=True)
+        tmp = STAGE / "restore" / e["remote"]
+        tmp.parent.mkdir(parents=True, exist_ok=True)
+        download_bucket_files(BUCKET, [(f"{PREFIX}/{run_id}/{e['remote']}", str(tmp))],
+                              token=TOKEN)
+        if hashlib.sha256(tmp.read_bytes()).hexdigest() != e["sha256"]:
+            print(f"    !! {target.name}: migrated copy fails its own checksum; left alone")
+            failed += 1
+            continue
+        # Stage locally, then one close()d copy onto the mount. Never stream a
+        # download straight through FUSE — that is the write pattern this whole
+        # script exists to work around.
+        shutil.copyfile(tmp, target)
+        tmp.unlink(missing_ok=True)
+        print(f"    restored {target.name}: {have} B -> {e['bytes']} B")
+        restored += 1
+
+    shutil.rmtree(STAGE / "restore", ignore_errors=True)
+    print(f"==> {restored} restored, {failed} failed")
+    return 1 if failed else 0
+
+
+sys.exit(restore() if RESTORE else snapshot())
+PY