kiteml-cli 0.5.0__tar.gz → 0.7.0__tar.gz

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: kiteml-cli
-Version: 0.5.0
+Version: 0.7.0
 Summary: Kite CLI and MCP server: train robot policies with reinforcement learning in simulation, augment LeRobot datasets, and build MuJoCo digital twins, from your terminal or an AI agent.
 Project-URL: Homepage, https://kiteml.com
 Project-URL: Documentation, https://docs.kiteml.com/tools/cli-and-mcp
@@ -138,7 +138,7 @@ dataset's folder in a public bucket (`s3://bucket/dataset/`).
 Build a digital twin (beta, enabled per account):
 
 ```bash
-kite twin create lerobot/svla_so101_pickplace --out ./twins   # -> ./twins/twin_.../scene.xml
+kite twin create lerobot/svla_so101_pickplace --robot so101 --out ./twins   # -> ./twins/twin_.../scene.xml
 ```
 
 Evaluate a policy on a Kite SDK spec, on Kite's machines (beta):
@@ -150,7 +150,9 @@ kite eval run specs/smoke.yaml --out ./runs    # -> ./runs/eval_.../scorecard.js
 
 A file too big to send, such as a checkpoint's weights, can come from the Hugging Face Hub or a public S3 bucket:
 `--hf PATH=OWNER/NAME@COMMIT:FILE` or `--hf PATH=<the file's huggingface.co URL>`, and `--s3 PATH=s3://BUCKET/KEY`,
-PATH being the path the spec names.
+PATH being the path the spec names. A checkpoint of one of your Kite training runs goes as
+`--training-run PATH=trn_...[@STEP]` (its latest without `@STEP`), PATH being the checkpoint folder the spec names:
+Kite reads it from its own storage and publishes nothing.
 
 Every command prints JSON (`{"ok": true, "data": ...}`) so scripts and agents can parse it. Run `kite --help` for
 the full list.

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/README.md

@@ -96,7 +96,7 @@ dataset's folder in a public bucket (`s3://bucket/dataset/`).
 Build a digital twin (beta, enabled per account):
 
 ```bash
-kite twin create lerobot/svla_so101_pickplace --out ./twins   # -> ./twins/twin_.../scene.xml
+kite twin create lerobot/svla_so101_pickplace --robot so101 --out ./twins   # -> ./twins/twin_.../scene.xml
 ```
 
 Evaluate a policy on a Kite SDK spec, on Kite's machines (beta):
@@ -108,7 +108,9 @@ kite eval run specs/smoke.yaml --out ./runs    # -> ./runs/eval_.../scorecard.js
 
 A file too big to send, such as a checkpoint's weights, can come from the Hugging Face Hub or a public S3 bucket:
 `--hf PATH=OWNER/NAME@COMMIT:FILE` or `--hf PATH=<the file's huggingface.co URL>`, and `--s3 PATH=s3://BUCKET/KEY`,
-PATH being the path the spec names.
+PATH being the path the spec names. A checkpoint of one of your Kite training runs goes as
+`--training-run PATH=trn_...[@STEP]` (its latest without `@STEP`), PATH being the checkpoint folder the spec names:
+Kite reads it from its own storage and publishes nothing.
 
 Every command prints JSON (`{"ok": true, "data": ...}`) so scripts and agents can parse it. Run `kite --help` for
 the full list.

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/src/kite_cli/__init__.py

@@ -1,3 +1,3 @@
 """Kite CLI and MCP server for the Kite robotics training platform."""
 
-__version__ = "0.5.0"
+__version__ = "0.7.0"

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/src/kite_cli/cli/commands/evaluations.py

@@ -17,7 +17,8 @@ SPEC is the file `kite run` takes (kite_spec_version "2"). The command sends its
 names (`files_required`), and sends each from this disk: a file or folder Kite ships by its sha256 alone, anything
 else inline (about 15 MB in all). A file on the Hugging Face Hub, such as a checkpoint's weights, goes as a
 reference: `--hf PATH=OWNER/NAME@COMMIT:FILE`, PATH as the spec names it (a folder's file as FOLDER/FILE); an object
-in a public S3 bucket as `--s3 PATH=s3://BUCKET/KEY`.
+in a public S3 bucket as `--s3 PATH=s3://BUCKET/KEY`; a checkpoint of one of your Kite training runs as
+`--training-run PATH=trn_...[@STEP]`, PATH being the folder the spec names (`policy.asset_dir`).
 Where runs go: --out DIR, else $KITE_EVAL_RUNS_DIR, else ./runs (`~` works). Paths printed are absolute. A run
 directory extracts to <out>/<run_id>/; the SDK's `kite verify <out>/<run_id> --evals DIR` checks it, DIR being the
 evaluators the spec named (the run records where its job read them, not your folder).
@@ -55,12 +56,16 @@ _BODY_MAX = 24_000_000  # base64 characters in all, under the 32 MiB a request c
 _ROUNDS = 64  # validate calls while the files are worked out: one per missing path or unshipped hash, at most
 _HF = re.compile(r"(?P<repo>[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*)@(?P<revision>[0-9a-f]{40})"
                  r":(?P<path>.+)")
+_TRN = re.compile(r"(?P<run>trn_[0-9A-HJKMNP-TV-Z]{26})(@(?P<step>[0-9]+))?")  # a training run's id, as the API has it
 
 _SET_HELP = "FIELD=VALUE, as `kite run --set` (repeatable)"
 _HF_HELP = ("PATH=OWNER/NAME@COMMIT:FILE, or PATH=URL (https://huggingface.co/OWNER/NAME/blob/REVISION/FILE; Kite pins a "
             "branch to its commit): the spec's PATH is that file of a public Hugging Face repo (repeatable)")
 _S3_HELP = ("PATH=s3://BUCKET/KEY (or its https://BUCKET.s3.amazonaws.com/KEY URL): the spec's PATH is that object of a "
             "public S3 bucket, pinned to the object it is when the run is created (repeatable)")
+_TRN_HELP = ("PATH=trn_...[@STEP]: the spec's PATH (a folder, e.g. policy.asset_dir) is the pretrained_model/ folder "
+             "of one of your training runs' checkpoints, at STEP or its latest; read from Kite's storage, nothing is "
+             "published (repeatable)")
 
 
 @app.command()
@@ -69,12 +74,13 @@ def validate(
     set_: Optional[List[str]] = typer.Option(None, "--set", help=_SET_HELP),
     hf: Optional[List[str]] = typer.Option(None, "--hf", help=_HF_HELP),
     s3: Optional[List[str]] = typer.Option(None, "--s3", help=_S3_HELP),
+    training_run: Optional[List[str]] = typer.Option(None, "--training-run", help=_TRN_HELP),
     kite_version: Optional[str] = typer.Option(None, "--kite-version", help="The Kite release you expect"),
 ):
     """Free dry run of a run: the spec plans, every file it names is sent, and the price. Creates nothing."""
     try:
         client = KiteClient(timeout=_SLOW_S)
-        _, validation, sent = request(client, spec, set_ or [], {**hf_refs(hf or []), **s3_refs(s3 or [])}, kite_version)
+        _, validation, sent = request(client, spec, set_ or [], remote_refs(hf, s3, training_run), kite_version)
         print_json({"ok": True, "data": validation, "files": sent})
     except KiteAPIError as e:
         print_error(e.message, code=e.code, details=e.details)
@@ -86,6 +92,7 @@ def run(
     set_: Optional[List[str]] = typer.Option(None, "--set", help=_SET_HELP),
     hf: Optional[List[str]] = typer.Option(None, "--hf", help=_HF_HELP),
     s3: Optional[List[str]] = typer.Option(None, "--s3", help=_S3_HELP),
+    training_run: Optional[List[str]] = typer.Option(None, "--training-run", help=_TRN_HELP),
     max_minutes: Optional[int] = typer.Option(
         None, "--max-minutes", help="The run's cap, 1-720: what it reserves and the most it can cost "
                                     "(default 60 on CPU, 120 on an L4)"),
@@ -106,7 +113,7 @@ def run(
     key = idempotency_key or str(uuid.uuid4())
     try:
         client = KiteClient(timeout=_SLOW_S)
-        body, _, _ = request(client, spec, set_ or [], {**hf_refs(hf or []), **s3_refs(s3 or [])}, kite_version)
+        body, _, _ = request(client, spec, set_ or [], remote_refs(hf, s3, training_run), kite_version)
         options = {k: v for k, v in (("max_minutes", max_minutes), ("no_eval", no_eval),
                                      ("verify_determinism", verify_determinism)) if v}
         body = {**body, **({"options": options} if options else {}), **({"label": label} if label else {})}
@@ -251,8 +258,8 @@ def resume(run_id: str = typer.Argument(..., help="Evaluation id (eval_...)")):
 def request(client: KiteClient, spec: Path, overrides: List[str], remote: Dict[str, dict],
             kite_version: Optional[str] = None) -> Tuple[dict, dict, Dict[str, str]]:
     """The request body for SPEC with every path it names, the validation that accepted it, and how each path went
-    (`sha256`, `inline`, `hf`, `s3`). `remote` maps the paths that are an `hf` file or an `s3` object to their
-    reference ({"hf": …} or {"s3": …}).
+    (`sha256`, `inline`, `hf`, `s3`, `training_run`). `remote` maps the paths that are an `hf` file, an `s3` object
+    or a training run's checkpoint folder to their reference ({"hf": …}, {"s3": …} or {"training_run": …}).
 
     The API names what it lacks: `files_required` lists the paths the spec names that the request does not carry,
     and `file_unavailable` on `files.<path>` says a hash Kite does not ship. A path first goes by its hash alone, so
@@ -263,7 +270,8 @@ def request(client: KiteClient, spec: Path, overrides: List[str], remote: Dict[s
     for key, ref in remote.items():
         files[key] = dict(ref)
         local = _local(base, key)
-        if local.is_file():  # pin the bytes this machine has: the Hub's file or the object must be the same
+        # pin the bytes this machine has: the Hub's file or the object must be the same (Kite pins a checkpoint's)
+        if local.is_file() and "training_run" not in ref:
             files[key]["sha256"] = _file_sha256(local)
         sent[key] = next(iter(ref))
     body: dict = {"spec_yaml": text, "spec_path": spec_path, "files": files}
@@ -364,7 +372,8 @@ def _fits(key: str, raw_bytes: int, files: Dict[str, dict], biggest: List[Tuple[
         hint = "".join(f", {posixpath.join(key, rel)} {p.stat().st_size >> 20} MB" for rel, p in biggest)
         raise KiteAPIError(413, f"{key} is not a file Kite ships, and it is too big to send inline "
                                 f"({raw_bytes >> 20} MB{hint}): put its large files on the Hugging Face Hub or in a public "
-                                "S3 bucket and pass each with --hf PATH=OWNER/NAME@COMMIT:FILE or --s3 PATH=s3://BUCKET/KEY", code="file_too_large",
+                                "S3 bucket and pass each with --hf PATH=OWNER/NAME@COMMIT:FILE or --s3 PATH=s3://BUCKET/KEY (a "
+                                "checkpoint Kite trained for you: --training-run PATH=trn_…@STEP)", code="file_too_large",
                            details={"path": key})
 
 
@@ -433,6 +442,26 @@ def s3_refs(pairs: List[str]) -> Dict[str, dict]:
     return refs
 
 
+def training_run_refs(pairs: List[str]) -> Dict[str, dict]:
+    """PATH=trn_…[@STEP] pairs as {"training_run": id, "step": STEP} references, the step left out for the run's
+    latest checkpoint; the API checks the run is the caller's."""
+    refs: Dict[str, dict] = {}
+    for pair in pairs:
+        key, _, ref = pair.partition("=")
+        m = _TRN.fullmatch(ref.strip())
+        if not key.strip() or not m:
+            raise KiteAPIError(400, f"--training-run {pair!r}: expected PATH=trn_…[@STEP], a training run's id "
+                                    "(trn_ and 26 characters) and optionally the step of one of its checkpoints",
+                               code="parameter_invalid", details={"param": "training_run"})
+        refs[key.strip()] = {"training_run": m["run"], **({"step": int(m["step"])} if m["step"] else {})}
+    return refs
+
+
+def remote_refs(hf: Optional[List[str]], s3: Optional[List[str]], training_run: Optional[List[str]]) -> Dict[str, dict]:
+    """The --hf, --s3 and --training-run pairs as one map of path to reference."""
+    return {**hf_refs(hf or []), **s3_refs(s3 or []), **training_run_refs(training_run or [])}
+
+
 # --- downloads ------------------------------------------------------------------------------------------------------
 
 def _page(limit: int, after: Optional[str]) -> dict:

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/src/kite_cli/cli/commands/twin.py

@@ -2,6 +2,7 @@
 MuJoCo scene out (about 100 minutes per twin).
 
   kite twin validate lerobot/svla_so101_pickplace [--camera side]    free: cameras, episodes, time and price
+  kite twin validate s3://my-bucket/my_dataset --robot so101          a dataset in a public S3 bucket
   kite twin create   lerobot/svla_so101_pickplace --wait --out ./twins    -> ./twins/twin_.../scene.xml
   kite twin status   twin_01K...   [--wait]
   kite twin list     [--limit 20] [--after twin_...]
@@ -10,7 +11,8 @@ MuJoCo scene out (about 100 minutes per twin).
   kite twin resume   twin_01K...
   kite twin download twin_01K...   [--out DIR] [--no-extract]
 
-DATASET is a Hugging Face repo id; the API takes it as {"source": {"type": "huggingface", "repo_id": DATASET}}.
+DATASET is a Hugging Face repo id or URL, which the API takes as {"source": {"type": "huggingface", "repo_id": DATASET}},
+or an S3 link to the dataset's folder in a public bucket, taken as {"source": {"type": "s3", "uri": DATASET}}.
 Where twins go: --out DIR, else $KITE_TWINS_DIR, else ./twins (`~` works). Paths printed are absolute.
 
 A signed storage URL is fetched with a plain HTTP client, never with the Kite key; the API's own archive
@@ -21,6 +23,7 @@ and unpacked only inside <out>/<twin_id>/.
 from __future__ import annotations
 
 import os
+import re
 import uuid
 from pathlib import Path, PurePosixPath
 from typing import Optional
@@ -38,18 +41,21 @@ app = typer.Typer(no_args_is_help=True, help="Digital twins: one robot episode -
 _TERMINAL = {"succeeded", "failed", "canceled"}  # not canceling: the cancel was accepted, the build has not stopped
 _POLL_S = 30.0
 _SLOW_S = 180.0  # validate/create check the dataset server-side (the API allows the builder 120 s)
+_DATASET_HELP = ("Hugging Face dataset id or URL, e.g. lerobot/svla_so101_pickplace, or an S3 link to the dataset's "
+                 "folder in a public bucket, s3://BUCKET/PREFIX")
 
 
 @app.command()
 def validate(
-    dataset: str = typer.Argument(..., help="Hugging Face dataset id or URL, e.g. lerobot/svla_so101_pickplace"),
+    dataset: str = typer.Argument(..., help=_DATASET_HELP),
     camera: Optional[str] = typer.Option(None, "--camera", help="Substring of the camera that seeds the room"),
     wrist_camera: Optional[str] = typer.Option(None, "--wrist-camera", help="Substring of the wrist camera"),
+    robot: Optional[str] = typer.Option(None, "--robot", help="Also check this robot can drive the dataset"),
 ):
     """Free dry run of a create: cameras, episodes, fps, and a time estimate. Creates nothing."""
     try:
         client = KiteClient(timeout=_SLOW_S)
-        body = _body(source=_source(dataset), camera=camera, wrist_camera=wrist_camera)
+        body = _body(source=twin_source(dataset), camera=camera, wrist_camera=wrist_camera, robot=robot)
         print_json(client.post("/v1/twins/validate", json_data=body))
     except KiteAPIError as e:
         print_error(e.message, code=e.code, details=e.details)
@@ -57,7 +63,8 @@ def validate(
 
 @app.command()
 def create(
-    dataset: str = typer.Argument(..., help="Hugging Face dataset id or URL, e.g. lerobot/svla_so101_pickplace"),
+    dataset: str = typer.Argument(..., help=_DATASET_HELP),
+    robot: str = typer.Option(..., "--robot", help="The robot in the episode: so100, so101, koch, aloha, wx250s, wxai_biarm, piper, xarm7, panda, panda_ee_state, g1_dex3, vega_1p"),
     name: Optional[str] = typer.Option(None, "--name", help="A short name for the twin"),
     task: Optional[str] = typer.Option(None, "--task", help="Which objects to author: the ones the robot manipulates"),
     prompt: Optional[str] = typer.Option(None, "--prompt", help="How to describe the room to the reconstructor"),
@@ -69,16 +76,16 @@ def create(
     idempotency_key: Optional[str] = typer.Option(
         None, "--idempotency-key", help="Safe retries: same key + body -> same twin (default: a generated key)"
     ),
-    wait: bool = typer.Option(False, "--wait", help="Poll until succeeded | failed | canceled (about 100 min)"),
+    wait: bool = typer.Option(False, "--wait", help="Poll until succeeded | failed | canceled (a few hours)"),
     out: Optional[Path] = typer.Option(
         None, "-o", "--out", help="When it succeeds, download, verify and unpack the scene into DIR (implies --wait); "
                                   "default location for `kite twin download` is $KITE_TWINS_DIR, else ./twins"
     ),
 ):
-    """Build a twin from one episode of a LeRobot dataset (queued; about 100 minutes; 5,000 tokens, refunded if it
-    builds nothing)."""
+    """Build a twin from one episode of a LeRobot dataset and its robot (queued; a few hours; 25,000 tokens, refunded
+    if it ends before authoring)."""
     key = idempotency_key or str(uuid.uuid4())
-    body = _body(source=_source(dataset), name=name, task=task, prompt=prompt, camera=camera,
+    body = _body(source=twin_source(dataset), robot=robot, name=name, task=task, prompt=prompt, camera=camera,
                  wrist_camera=wrist_camera, ceiling_m=ceiling_m)
     try:
         client = KiteClient(timeout=_SLOW_S)
@@ -177,7 +184,7 @@ def cancel(twin_id: str = typer.Argument(..., help="Twin id (twin_...)")):
 @app.command()
 def resume(twin_id: str = typer.Argument(..., help="Twin id (twin_...)")):
     """Re-queue a failed or canceled twin; only the stages that did not finish run again. Free while it's paid for;
-    a twin that was refunded is charged 5,000 tokens again."""
+    a twin that was refunded is charged 25,000 tokens again."""
     try:
         print_json(KiteClient().post(f"/v1/twins/{twin_id}/resume"))
     except KiteAPIError as e:
@@ -204,8 +211,11 @@ def _body(**fields) -> dict:
     return {k: v for k, v in fields.items() if v not in (None, "")}
 
 
-def _source(dataset: str) -> dict:
-    """The positional DATASET as the API takes it: the `source` object /v1/augmentations takes too."""
+def twin_source(dataset: str) -> dict:
+    """DATASET as the API's `source`: an S3 link (s3://BUCKET/PREFIX, or its https://….amazonaws.com URL) is its
+    `uri`, anything else a Hugging Face `repo_id`. The API checks either."""
+    if re.match(r"(?i)\s*(s3://|https?://[^/]*amazonaws\.com/)", dataset):
+        return {"type": "s3", "uri": dataset.strip()}
     return {"type": "huggingface", "repo_id": dataset}
 
 

{kiteml_cli-0.5.0 → kiteml_cli-0.7.0}/src/kite_cli/mcp/server.py

@@ -51,11 +51,6 @@ def _client() -> KiteClient:
     return KiteClient(api_url=os.environ.get("KITE_API_URL"), token=token, surface="kite-mcp")
 
 
-def _twin_source(dataset: str) -> dict:
-    """The tools take `dataset`, the simplest thing to ask an agent for; the API takes it as a `source`."""
-    return {"type": "huggingface", "repo_id": dataset}
-
-
 def _fetch_command(d: dict, fname: str, client: KiteClient) -> str:
     """The command that fetches a finished resource's archive. From Kite-Version 2026-10-12 `download.url` is this
     API's path, which needs the key: ask it for the fresh signed link it redirects to, so an agent without the key
@@ -140,10 +135,10 @@ def _eval_request(args: dict, client: KiteClient) -> tuple[dict, dict | None]:
         raise ValueError("`spec` only works with the local (stdio) Kite MCP server; send `spec_yaml` and `files`")
     from pathlib import Path
 
-    from kite_cli.cli.commands.evaluations import hf_refs, request, s3_refs
+    from kite_cli.cli.commands.evaluations import remote_refs, request
 
-    remote = {**hf_refs([f"{k}={v}" for k, v in (args.get("hf") or {}).items()]),
-              **s3_refs([f"{k}={v}" for k, v in (args.get("s3") or {}).items()])}
+    pairs = {m: [f"{k}={v}" for k, v in (args.get(m) or {}).items()] for m in ("hf", "s3", "training_runs")}
+    remote = remote_refs(pairs["hf"], pairs["s3"], pairs["training_runs"])
     body, validation, _ = request(client, Path(args["spec"]), args.get("overrides") or [], remote, args.get("kite_version"))
     return body, validation
 
@@ -204,14 +199,19 @@ _EVAL_SPEC = {
     "spec": {"type": "string", "description": "A Kite SDK spec file on this machine, best as an absolute path (the local server only): every path it names is sent from this disk, a file Kite ships by its hash alone"},
     "spec_yaml": {"type": "string", "description": "The spec's text, instead of `spec` (the hosted server): send `spec_path` and `files` with it"},
     "spec_path": {"type": "string", "description": "With spec_yaml: its path in your workspace, e.g. specs/smoke.yaml. Paths in the spec are read relative to it"},
-    "files": {"type": "object", "description": "With spec_yaml: each path the spec names, as it names it, to where its bytes come from: {\"kind\": \"directory\", \"shipped\": \"evals\"} (a path Kite ships), {\"sha256\", \"kind\"} (a shipped file by hash), {\"text\"} or {\"base64\"} (a folder as a .tar), or {\"hf\": {\"repo\", \"revision\" (a full commit), \"path\"}} or {\"hf\": \"<the file's huggingface.co URL>\"} (a branch is pinned to its commit), or {\"s3\": \"s3://BUCKET/KEY\"} (an object in a public bucket, pinned to its ETag). A missing one is files_required, listing each"},
+    "files": {"type": "object", "description": "With spec_yaml: each path the spec names, as it names it, to where its bytes come from: {\"kind\": \"directory\", \"shipped\": \"evals\"} (a path Kite ships), {\"sha256\", \"kind\"} (a shipped file by hash), {\"text\"} or {\"base64\"} (a folder as a .tar), or {\"hf\": {\"repo\", \"revision\" (a full commit), \"path\"}} or {\"hf\": \"<the file's huggingface.co URL>\"} (a branch is pinned to its commit), or {\"s3\": \"s3://BUCKET/KEY\"} (an object in a public bucket, pinned to its ETag), or {\"training_run\": \"trn_...\", \"step\"} (a folder: the pretrained_model/ of one of your training runs' checkpoints, step left out for its latest; nothing is published). A missing one is files_required, listing each"},
     "overrides": {"type": "array", "items": {"type": "string"}, "description": "FIELD=VALUE pairs, as `kite run --set`, e.g. episodes.count=20. Optional"},
     "hf": {"type": "object", "description": "With spec: {PATH: \"OWNER/NAME@COMMIT:FILE\" or the file's huggingface.co URL}, the paths that are a public Hugging Face file (a checkpoint's weights). Optional"},
     "s3": {"type": "object", "description": "With spec: {PATH: \"s3://BUCKET/KEY\"}, the paths that are an object in a public S3 bucket. Optional"},
+    "training_runs": {"type": "object", "description": "With spec: {PATH: \"trn_...\" or \"trn_...@STEP\"}, the folders (policy.asset_dir) that are the pretrained_model/ of one of your training runs' checkpoints, its latest without @STEP; read from Kite's storage, nothing is published. Optional"},
     "kite_version": {"type": "string", "description": "The Kite release you expect; refused if it is not the one served. Optional"},
 }
 _EVAL_ID = {"run_id": {"type": "string", "description": "Evaluation id (eval_...)"}}
 
+# The twin tools take `dataset`, the simplest thing to ask an agent for, and send it as the API's `source`.
+_TWIN_DATASET = ("HF dataset id or URL, e.g. lerobot/svla_so101_pickplace (the API's source.repo_id), or an S3 link to "
+                 "the dataset's folder in a public bucket, s3://BUCKET/PREFIX (source.uri)")
+
 # ---------------------------------------------------------------------------
 # Tool definitions — every tool has title + annotations per Claude directory
 # ---------------------------------------------------------------------------
@@ -265,20 +265,22 @@ TOOLS = [
 
     # Digital twins (public /v1 API): one robot episode -> an interactable MuJoCo scene, ~100 min each
     Tool(name="kite_twin_validate", title="Validate Digital Twin Dataset",
-         description="Free check before building a twin: confirms a HuggingFace LeRobot dataset resolves and has camera video, says which camera will seed the room, and returns the episodes, fps, typical duration and price (estimate.tokens). Creates nothing. Call this first; errors name what to fix (dataset_not_found, camera_not_found with the available cameras, no_suitable_camera).",
+         description="Free check before building a twin: confirms a LeRobot dataset (on HuggingFace, or in a public S3 bucket) resolves and has camera video, says which camera will seed the room, and returns the episodes, fps, typical duration and price (estimate.tokens). Creates nothing. Call this first; errors name what to fix (dataset_not_found, dataset_private for an S3 bucket that is not public, camera_not_found with the available cameras, no_suitable_camera).",
          inputSchema={"type": "object", "properties": {
-             "dataset": {"type": "string", "description": "HF dataset id or URL, e.g. lerobot/svla_so101_pickplace (the API's source.repo_id)"},
-             "camera": {"type": "string", "description": "Substring of the camera that should seed the room. Optional"}},
+             "dataset": {"type": "string", "description": _TWIN_DATASET},
+             "camera": {"type": "string", "description": "Substring of the camera that should seed the room. Optional"},
+             "robot": {"type": "string", "description": "Also check this robot can drive the dataset (see kite_twin_create). Optional"}},
           "required": ["dataset"]}, annotations=_READ),
     Tool(name="kite_twin_create", title="Create Digital Twin",
-         description="Build an interactable MuJoCo digital twin from one episode of a HuggingFace LeRobot dataset: the room is reconstructed, the manipulated objects are authored with working joints, and every stage is reviewed by measurement. Only `dataset` is required; the twin reports it back as source.repo_id. Takes about 100 minutes and costs 5,000 tokens, charged now and refunded in full if the build ends before any stage finishes (insufficient_tokens when the balance is short); returns a queued twin (id twin_...) to poll with kite_twin_status.",
+         description="Build an interactable MuJoCo digital twin from one episode of a LeRobot dataset, on HuggingFace or in a public S3 bucket: the room is reconstructed, the manipulated objects are authored with working joints, and every stage is reviewed by measurement. `dataset` and `robot` are required; the twin reports the dataset back as source.repo_id (source.uri for S3). Takes a few hours and costs 25,000 tokens, charged now and refunded in full if the build ends before authoring begins (insufficient_tokens when the balance is short); returns a queued twin (id twin_...) to poll with kite_twin_status.",
          inputSchema={"type": "object", "properties": {
-             "dataset": {"type": "string", "description": "HF dataset id or URL, e.g. lerobot/svla_so101_pickplace (the API's source.repo_id)"},
+             "dataset": {"type": "string", "description": _TWIN_DATASET},
+             "robot": {"type": "string", "description": "The robot in the episode: so100, so101, koch, aloha, wx250s, wxai_biarm, piper, xarm7, panda, panda_ee_state, g1_dex3 or vega_1p (robot_mismatch lists the ones that fit the dataset)"},
              "name": {"type": "string", "description": "Short name for the twin. Optional"},
              "task": {"type": "string", "description": "Which objects to author (the things the robot manipulates). Optional"},
              "camera": {"type": "string", "description": "Substring of the camera that seeds the room. Optional"},
              "idempotency_key": {"type": "string", "description": "Reuse the same key when retrying, so a retry never builds twice. Optional"}},
-          "required": ["dataset"]}, annotations=_WRITE),
+          "required": ["dataset", "robot"]}, annotations=_WRITE),
     Tool(name="kite_twin_status", title="Get Digital Twin Status",
          description="Get a twin's status (queued/processing/canceling/succeeded/failed/canceled), stage, progress, each finished stage's measurement, and a `next` instruction. canceling means a cancel was accepted but the build has not stopped yet; it becomes canceled. When succeeded, `download.url` is a direct link to the whole scene as one .tar.gz and `next` is the exact command to fetch and extract it (the folder contains scene.xml for MuJoCo).",
          inputSchema={"type": "object", "properties": {
@@ -294,7 +296,7 @@ TOOLS = [
              "twin_id": {"type": "string", "description": "Twin id (twin_...)"}},
           "required": ["twin_id"]}, annotations=_DESTRUCTIVE),
     Tool(name="kite_twin_resume", title="Resume Digital Twin",
-         description="Re-queue a failed or canceled twin. Every stage it finished is cached, so only what did not finish runs again. Free while the twin is paid for; one that was refunded (it built nothing) is charged 5,000 tokens again.",
+         description="Re-queue a failed or canceled twin. Every stage it finished is cached, so only what did not finish runs again. Free while the twin is paid for; one that was refunded (it ended before authoring) is charged 25,000 tokens again.",
          inputSchema={"type": "object", "properties": {
              "twin_id": {"type": "string", "description": "Twin id (twin_...)"}},
           "required": ["twin_id"]}, annotations=_WRITE),
@@ -464,12 +466,14 @@ async def _dispatch(name: str, args: dict, client: KiteClient) -> list[TextConte
         return _ok(client.post(f"/v1/augmentations/{args['aug_id']}/cancel"))
 
     # --- Digital twins (public /v1 API) ---
+    from kite_cli.cli.commands.twin import twin_source
+
     if name == "kite_twin_validate":
-        body = {"source": _twin_source(args["dataset"]), **{k: args[k] for k in ("camera",) if args.get(k)}}
+        body = {"source": twin_source(args["dataset"]), **{k: args[k] for k in ("camera", "robot") if args.get(k)}}
         return _ok(client.post("/v1/twins/validate", json_data=body))
     if name == "kite_twin_create":
-        body = {"source": _twin_source(args["dataset"]),
-                **{k: args[k] for k in ("name", "task", "camera") if args.get(k)}}
+        body = {"source": twin_source(args["dataset"]),
+                **{k: args[k] for k in ("robot", "name", "task", "camera") if args.get(k)}}
         headers = {"Idempotency-Key": args["idempotency_key"]} if args.get("idempotency_key") else None
         return _ok(_twin_view(client.post("/v1/twins", json_data=body, extra_headers=headers), client))
     if name == "kite_twin_status":