Upload an SDK run

import argparse
import hashlib
import json
import math
import re
import uuid
from datetime import datetime
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("--session", type=Path, required=True)
parser.add_argument("--task", required=True)
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--rewards", type=Path)
args = parser.parse_args()

def stop(message):
    parser.error(message)

if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}", args.task):
    stop("Use the task directory name: letters, digits, dots, _ or -.")
raw = args.session.expanduser().read_bytes()
records = []
for line in raw.decode("utf-8").splitlines():
    try:
        record = json.loads(line)
    except json.JSONDecodeError:
        continue  # Keep non-JSON output in the original file.
    if isinstance(record, dict):
        records.append(record)

metadata = [r["_meta"] for r in records if "_meta" in r]
prompts = [r["_prompt"] for r in records if "_prompt" in r]
if len(metadata) != 1 or len(prompts) != 1:
    stop("Use a local log with one SDK run; do not combine runs or commands.")
if not any("_sessionEnd" in r for r in records):
    stop("Finish the run and call kill() before packaging its log.")
meta, prompt = metadata[0], prompts[0]
if not isinstance(meta, dict) or not isinstance(prompt, dict):
    stop("Expected local SDK metadata and prompt records.")
harness = meta.get("agent")
if harness not in ("codex", "claude"):
    stop("This example supports local Codex and Claude SDK logs.")
if not isinstance(prompt.get("text"), str) or not prompt["text"]:
    stop("The log has no recorded prompt text.")
if harness == "codex":
    completed = sum(r.get("type") == "turn.completed" for r in records)
    failed = any(r.get("type") == "turn.failed" for r in records)
    if completed != 1 or failed:
        stop("Expected one successfully completed Codex turn.")
else:
    outcomes = [r for r in records if r.get("type") == "result"]
    if len(outcomes) != 1 or outcomes[0].get("is_error") is not False:
        stop("Expected one successful Claude result.")

rewards = None
if args.rewards:
    rewards = json.loads(args.rewards.read_text())
    if not isinstance(rewards, dict) or not rewards or not all(
        isinstance(v, (int, float)) and not isinstance(v, bool)
        and math.isfinite(v) for v in rewards.values()
    ):
        stop("Rewards must be a nonempty JSON object of recorded numeric scores.")

source_id = str(uuid.uuid5(
    uuid.NAMESPACE_URL, hashlib.sha256(raw).hexdigest() + ":" + args.task
))
trial_name = args.task + "__sdk"
agent_info = {"name": harness, "version": ""}
if meta.get("model"):
    agent_info["model_info"] = {"name": meta["model"]}
result = {
    "task_name": args.task, "trial_name": trial_name, "trial_uri": "",
    "task_id": {}, "task_checksum": "", "config": {},
    "agent_info": agent_info,
}
if rewards is not None:
    result["verifier_result"] = {"rewards": rewards}
trajectory = {
    "schema_version": "ATIF-v1.7", "session_id": meta["tag"],
    "agent": {"name": harness, "version": ""},
    "steps": [{"step_id": 1, "source": "user", "message": prompt["text"]}],
    "notes": "Opening SDK prompt; agent activity is in the native transcript.",
}
if harness == "codex":
    date = datetime.fromisoformat(meta["timestamp"].replace("Z", "+00:00"))
    native = Path(".codex/sessions") / date.strftime("%Y/%m/%d") / f"rollout-{source_id}.jsonl"
else:
    native = Path(".claude/projects/sdk-upload") / f"{source_id}.jsonl"

args.out.mkdir(parents=True, exist_ok=False)
trial = args.out / trial_name
destination = trial / "agent" / native
destination.parent.mkdir(parents=True)
destination.write_bytes(raw)
for path, value in (
    (args.out / "config.json", {"job_name": args.out.name}),
    (args.out / "result.json", {"id": source_id}),
    (trial / "result.json", result),
    (trial / "agent/trajectory.json", trajectory),
):
    path.write_text(json.dumps(value, indent=2) + "\n")
print(f"Ready to upload: {args.out}")
print("Recorded scores included." if rewards is not None else "No score supplied.")