test_studio.py

"""Behavioral checks for billing, isolation, output preservation and preset recovery."""

import hashlib
import json
import sys
import time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock

import pyarrow as pa
import pytest

sys.path.insert(0, str(Path(__file__).parent))
import studio  # noqa: E402
from model_catalog import MODELS, model_profile  # noqa: E402
from provider import Budget, BudgetExceededError  # noqa: E402
from worker import FALLBACK, PreservingWriter, TunedPlan  # noqa: E402

from lerobot.annotations.steerable_pipeline.config import PlanConfig
from lerobot.annotations.steerable_pipeline.modules.plan_subtasks_memory import PlanSubtasksMemoryModule
from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient


@pytest.mark.parametrize(
    "features",
    [
        [f for f in studio.FEATURES if f not in {studio.HANDS, studio.SLAM}],
        ["Interjections & speech"],
        ["Memory"],
        ["Visual questions"],
    ],
)
@pytest.mark.parametrize("custom", [False, True])
def test_real_video_pipeline_and_checkbox_filtering(tmp_path, features, custom):
    import av
    import numpy as np
    from worker import annotate_episode

    from lerobot.annotations.steerable_pipeline.reader import EpisodeRecord
    from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging

    camera = "observation.images.top"
    video = tmp_path / "videos" / camera / "chunk-000" / "file-000.mp4"
    video.parent.mkdir(parents=True)
    with av.open(str(video), "w") as container:
        stream = container.add_stream("mpeg4", rate=10)
        stream.width, stream.height, stream.pix_fmt = 64, 64, "yuv420p"
        for i in range(40):
            frame = av.VideoFrame.from_ndarray(np.full((64, 64, 3), i * 5, dtype=np.uint8), format="rgb24")
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
    info = {"fps": 10, "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"}
    row = {
        "length": 40,
        f"videos/{camera}/chunk_index": 0,
        f"videos/{camera}/file_index": 0,
        f"videos/{camera}/from_timestamp": 0.0,
    }
    record = EpisodeRecord(
        0,
        "move the cup to the table",
        tuple(i / 10 for i in range(40)),
        tuple(range(40)),
        tmp_path / "unused.parquet",
        0,
        40,
    )

    from prompt_config import TEMPLATES, validate_prompts

    overrides = (
        validate_prompts(
            {key: item["default"] + f"\nCUSTOM_{key}" for key, item in TEMPLATES.items()}, features
        )
        if custom
        else {}
    )
    seen = set()

    def respond(messages):
        text = " ".join(p.get("text", "") for m in messages for p in m["content"] if isinstance(p, dict))
        for key in overrides:
            if f"CUSTOM_{key}" in text:
                seen.add(key)
        if "CUSTOM_plan_render" in text:
            return {"plan": "1. Put the cup on the table."}
        if "meaningful action phases" in text:
            return {
                "subtasks": [
                    {"text": "grasp the cup", "start": 0, "end": 2.5},
                    {"text": "place the cup on table", "start": 2.5, "end": 3.9},
                ]
            }
        if "compressed semantic memory" in text:
            return {"memory": "grasped the cup"}
        if "acknowledgement the robot" in text:
            return {"text": "On it."}
        if "Write ONE compact interjection" in text:
            return {"interjection": "place the cup on the table", "speech": "Okay."}
        if "frame-grounded visual question" in text:
            return {"question": "How many cups?", "answer": {"label": "cup", "count": 1}}
        if "rephras" in text.lower():
            return {"rephrasings": ["put the cup on the table"]}
        raise AssertionError("Unexpected pipeline prompt")

    result = annotate_episode(
        record,
        row,
        info,
        tmp_path,
        {"camera": camera, "features": features, "vqa_hz": 0.1, "prompts": overrides},
        StubVlmClient(respond),
        tmp_path / "staging",
    )
    staged = EpisodeStaging(tmp_path / "staging", 0).read_all()
    actual_styles = {r.get("style") for rows in staged.values() for r in rows}
    expected = {studio.FEATURES[f] for f in features if f != studio.HUMAN_VIDEOS}
    if "Interjections & speech" in features:
        expected.add(None)
    assert actual_styles == expected
    assert result["status"] == "completed"
    assert result["rows"] > 0
    assert set(overrides) == seen
    from worker import CUSTOM_PROMPTS

    assert CUSTOM_PROMPTS.get() is None


def selection():
    return {
        "episodes": [{"episode_index": 0, "length": 300}, {"episode_index": 1, "length": 600}],
        "seconds": 30,
        "info": {"fps": 30},
        "dataset": "source/robot",
        "revision": "a" * 40,
        "start": 0,
        "count": 2,
    }


def plan():
    return {
        "selection": selection(),
        "features": ["Subtasks"],
        "camera": "observation.images.top",
        "vqa_hz": 0.1,
        "mode": studio.COPY,
        "created": time.time(),
        "nonce": "request123",
        "costs": studio.estimate(selection(), ["Subtasks"], 0.1, 4, 0.0005),
    }


def api_mock():
    api = Mock()
    api.whoami.return_value = {"name": "alice"}
    api.list_jobs.return_value = []
    api.repo_exists.return_value = False
    api.dataset_info.return_value = SimpleNamespace(id="source/robot", sha="a" * 40)
    api.space_info.return_value = SimpleNamespace(sha="b" * 40)
    api.list_jobs_hardware.return_value = [
        SimpleNamespace(name="cpu-upgrade", unit_label="minute", unit_cost_usd=0.0005)
    ]
    api.model_info.return_value = SimpleNamespace(
        inference_provider_mapping={studio.PROVIDER: SimpleNamespace(status="live")}
    )
    api.run_job.return_value = SimpleNamespace(id="job123")
    return api


def test_cost_features_and_timeout():
    base = studio.estimate(selection(), ["Subtasks"], 0.1, 1, 0.0005)
    plans = studio.estimate(selection(), ["Subtasks", "Plans"], 0.1, 1, 0.0005)
    memory = studio.estimate(selection(), ["Subtasks", "Memory"], 0.1, 1, 0.0005)
    vqa = studio.estimate(selection(), ["Visual questions"], 1.0, 4, 0.0005)
    assert base["inference"] == plans["inference"] < memory["inference"]
    assert base["cpu_max"] < vqa["cpu_max"]
    assert base["low"] < base["high"]
    assert vqa["calls"] == 32


@pytest.mark.parametrize("value", ["../secret", "/tmp/file", "videos/../../secret", "a\\b"])
def test_reject_metadata_path_traversal(value):
    with pytest.raises(ValueError):
        studio.relative_path(value)


def test_budget_unknown_failure_is_never_refunded():
    budget = Budget(0.09, 0.15, 0.5)
    first = budget.reserve(4096)
    budget.reconcile(first, None)
    second = budget.reserve(4096)
    budget.reconcile(second, {"prompt_tokens": 1000, "completion_tokens": 100})
    assert budget.reserved == pytest.approx(first + 0.0002)
    budget.reserve(4096)
    with pytest.raises(BudgetExceededError):
        budget.reserve(4096)


@pytest.mark.parametrize("mapping_format", ["dict", "list"])
def test_launch_identity_secret_and_immutable_manifest(mapping_format):
    api = api_mock()
    if mapping_format == "list":
        api.model_info.return_value.inference_provider_mapping = [
            SimpleNamespace(provider=studio.PROVIDER, status="live")
        ]
    configured = plan()
    configured["prompts"] = {"plan_subtasks": "custom instructions preserved verbatim"}
    factory = Mock(return_value=api)
    result = studio.launch(
        configured,
        "alice/annotations",
        5,
        True,
        "sensitive-token",
        space="lerobot/annotation-studio",
        api_factory=factory,
        expected_owner="alice",
    )
    factory.assert_called_once_with(token="sensitive-token")
    args = api.run_job.call_args.kwargs
    assert args["namespace"] == "alice"
    assert args["labels"]["annotation_output"] == "annotations"
    assert all("/" not in value for value in args["labels"].values())
    assert args["secrets"] == {"HF_TOKEN": "sensitive-token"}
    assert "sensitive-token" not in json.dumps(args["env"]) + json.dumps(args["command"])
    manifest = json.loads(args["env"]["ANNOTATION_MANIFEST"])
    assert manifest["episode_ids_sha256"] == hashlib.sha256(json.dumps([0, 1]).encode()).hexdigest()
    assert manifest["space_revision"] == "b" * 40
    assert manifest["prompts"] == configured["prompts"]
    assert manifest["model"] == "Qwen/Qwen3.8-27B:novita"
    assert manifest["request_concurrency"] == manifest["episode_parallelism"] == 32
    assert manifest["costs"]["context_length"] == 1_000_000
    assert "prompt-templates.json" in str(args["command"])
    assert manifest["inference_budget"] + plan()["costs"]["cpu_max"] == 5
    assert result["id"] == "job123"


@pytest.mark.parametrize("quote_owner,expected_owner", [("bob", "alice"), ("alice", "bob")])
def test_launch_binds_token_identity_to_session_and_quote(monkeypatch, quote_owner, expected_owner):
    monkeypatch.setenv("HF_TOKEN", "finerobotics-owner-token")
    api = api_mock()
    factory = Mock(return_value=api)
    configured = plan()
    configured["quote_owner"] = quote_owner
    with pytest.raises(ValueError):
        studio.launch(
            configured,
            "alice/output",
            5,
            True,
            "alice-token",
            expected_owner=expected_owner,
            api_factory=factory,
        )
    factory.assert_called_once_with(token="alice-token")
    api.run_job.assert_not_called()


@pytest.mark.parametrize(
    "output,consent,token",
    [("bob/output", True, "token"), ("alice/output", False, "token"), ("alice/output", True, False)],
)
def test_launch_rejects_unauthorized_or_unconfirmed(output, consent, token):
    api = api_mock()
    with pytest.raises(ValueError):
        studio.launch(
            plan(), output, 5, consent, token, space="lerobot/annotation-studio", api_factory=lambda **_: api
        )
    api.run_job.assert_not_called()


def test_duplicate_click_recovers_existing_job():
    api = api_mock()
    api.list_jobs.return_value = [
        SimpleNamespace(id="first-job", labels={"annotation_request": "request123"})
    ]
    assert (
        studio.launch(plan(), "alice/output", 5, True, "token", api_factory=lambda **_: api)["id"]
        == "first-job"
    )
    api.run_job.assert_not_called()


def test_source_changes_block_launch():
    api = api_mock()
    api.dataset_info.return_value.sha = "c" * 40
    with pytest.raises(ValueError, match="source changed"):
        studio.launch(plan(), "alice/output", 5, True, "token", api_factory=lambda **_: api)
    api.run_job.assert_not_called()


def test_compact_quote_preserves_pinned_episode_selection():
    api = api_mock()
    request = plan()
    digest = hashlib.sha256(json.dumps([0, 1]).encode()).hexdigest()
    request["selection"].pop("episodes")
    request["selection"]["episode_ids_sha256"] = digest
    studio.launch(
        request,
        "alice/output",
        5,
        True,
        "token",
        space="lerobot/annotation-studio",
        api_factory=lambda **_: api,
    )
    manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"])
    assert manifest["episode_ids_sha256"] == digest


def test_preserve_unselected_episodes_and_styles():
    def row(text, style):
        return {
            "role": "assistant",
            "content": text,
            "style": style,
            "timestamp": 0.0,
            "camera": None,
            "tool_calls": None,
        }

    old = [[row("old subtask", "subtask"), row("keep memory", "memory")], [row("untouched", "subtask")]]
    table = pa.table({"episode_index": [0, 1], "timestamp": [0.0, 0.0], "language_persistent": old})
    writer = PreservingWriter([0], {"subtask"})
    result = writer._materialize_table(table, [[row("new subtask", "subtask")], []], [[], []], drop_old=False)
    rows = result.column("language_persistent").to_pylist()
    assert {r["content"] for r in rows[0]} == {"new subtask", "keep memory"}
    assert rows[1] == old[1]


def test_recovery_uses_original_prompt_and_does_not_swallow_budget(monkeypatch):
    calls = []

    def generate(self, record, task=None):
        calls.append(FALLBACK.get())
        if not FALLBACK.get():
            raise ValueError("Invalid model response")
        return [{"text": "grasp cup", "start": 0, "end": 1}]

    monkeypatch.setattr(PlanSubtasksMemoryModule, "_generate_subtasks", generate)
    model = TunedPlan(vlm=StubVlmClient(lambda _: {}), config=PlanConfig())
    assert model._generate_subtasks(None)[0]["text"] == "grasp cup"
    assert calls == [False, True] and not FALLBACK.get()
    monkeypatch.setattr(
        PlanSubtasksMemoryModule, "_generate_subtasks", Mock(side_effect=BudgetExceededError())
    )
    with pytest.raises(BudgetExceededError):
        model._generate_subtasks(None)
    assert not model.used_fallback


def test_prompt_defaults_and_validation():
    from prompt_config import TEMPLATES, validate_prompts
    from worker import ACTION_PROMPT, ORIGINAL_LOAD

    for key, item in TEMPLATES.items():
        default = item["default"]
        assert validate_prompts({key: default}) == {}
        if key == "plan_subtasks":
            assert default == ACTION_PROMPT
        elif key != "plan_render":
            assert default == ORIGINAL_LOAD(key)
    default = TEMPLATES["plan_subtasks"]["default"]
    for bad in [
        "",
        default.replace("{episode_task}", ""),
        default + "{episode_task.__class__}",
        default + "{unknown}",
        default + "{episode_task:999999999}",
        "a" * 20001,
    ]:
        with pytest.raises(ValueError):
            validate_prompts({"plan_subtasks": bad})
    with pytest.raises(ValueError):
        validate_prompts({"unknown": "text"})
    edits = {"plan_subtasks": default + "\nUse concise labels."}
    assert validate_prompts(edits, ["Memory"]) == edits
    assert validate_prompts(edits, ["Visual questions"]) == {}


def test_custom_plan_adds_estimated_calls_and_reset_restores_baseline():
    from prompt_config import TEMPLATES

    features = ["Plans"]
    base = studio.estimate(selection(), features, 0.1, 1, 0.0005)
    template = TEMPLATES["plan_render"]["default"]
    custom = studio.estimate(
        selection(), features, 0.1, 1, 0.0005, {"plan_render": template + "\nBe concise."}
    )
    reset = studio.estimate(selection(), features, 0.1, 1, 0.0005, {"plan_render": template})
    assert custom["calls"] > base["calls"]
    assert custom["inference"] > base["inference"]
    assert custom["custom_prompts"] == 1
    assert reset == base


def test_custom_prompt_context_resets_on_failure_and_is_thread_isolated(monkeypatch):
    from concurrent.futures import ThreadPoolExecutor
    from threading import Barrier

    import worker
    from prompt_config import TEMPLATES

    barrier = Barrier(2)

    def fail(*args):
        barrier.wait(timeout=5)
        text = worker.load_prompt("plan_subtasks")
        fallback = FALLBACK.set(True)
        try:
            assert worker.load_prompt("plan_subtasks") == text
        finally:
            FALLBACK.reset(fallback)
        raise ValueError(text.splitlines()[-1])

    monkeypatch.setattr(worker, "run_runtime_episode", fail)

    def run(name):
        custom = TEMPLATES["plan_subtasks"]["default"] + "\n" + name
        with pytest.raises(ValueError, match=name):
            worker.annotate_episode(
                None,
                None,
                None,
                None,
                {"features": ["Subtasks"], "prompts": {"plan_subtasks": custom}},
                None,
                None,
            )
        assert worker.CUSTOM_PROMPTS.get() is None

    with ThreadPoolExecutor(max_workers=2) as pool:
        list(pool.map(run, ["FIRST_USER", "SECOND_USER"]))


def test_custom_plan_only_uses_remaining_actions_and_default_avoids_model():
    from prompt_config import TEMPLATES
    from worker import CUSTOM_PROMPTS

    record = SimpleNamespace(episode_task="move the cup")
    spans = [{"start": 0, "text": "reach the cup"}, {"start": 1, "text": "grasp the cup"}]
    vlm = Mock()
    module = TunedPlan(vlm=vlm, config=PlanConfig())
    assert module._generate_plan(record, spans, refresh_t=1) == "1. grasp the cup"
    vlm.generate_json.assert_not_called()
    vlm.generate_json.return_value = [{"plan": "1. Grasp the cup."}]
    ctx = CUSTOM_PROMPTS.set({"plan_render": TEMPLATES["plan_render"]["default"] + "\nBe brief."})
    try:
        assert module._generate_plan(record, spans, refresh_t=1) == "1. Grasp the cup."
        sent = str(vlm.generate_json.call_args)
        assert "grasp the cup" in sent and "reach the cup" not in sent
    finally:
        CUSTOM_PROMPTS.reset(ctx)


@pytest.mark.parametrize("token", [False, "private-reader-token"])
def test_camera_discovery_reads_only_pinned_info_and_uses_request_identity(tmp_path, monkeypatch, token):
    import tempfile

    path = tmp_path / "info.json"
    path.write_text(
        json.dumps(
            {
                "codebase_version": "v3.0",
                "fps": 30,
                "features": {
                    "observation.images.top": {"dtype": "video"},
                    "observation.images.left_wrist": {"dtype": "video"},
                    "observation.images.right_wrist": {"dtype": "video"},
                    "observation.images.depth": {"dtype": "video"},
                    "observation.state": {"dtype": "float32"},
                    "observation.images.still": {"dtype": "image"},
                },
            }
        )
    )
    api = Mock()
    api.dataset_info.return_value = SimpleNamespace(sha="c" * 40)
    factory = Mock(return_value=api)
    download = Mock(return_value=str(path))
    monkeypatch.setattr(studio, "HfApi", factory)
    monkeypatch.setattr(studio, "hf_hub_download", download)
    result = studio.read_cameras("owner/dataset", token)
    assert result["cameras"] == [
        "observation.images.top",
        "observation.images.left_wrist",
        "observation.images.right_wrist",
    ]
    factory.assert_called_once_with(token=token)
    api.list_repo_tree.assert_not_called()
    download.assert_called_once()
    args = download.call_args
    assert args.args == ("owner/dataset", "meta/info.json")
    assert args.kwargs["revision"] == "c" * 40
    assert args.kwargs["token"] == token
    assert args.kwargs["cache_dir"].startswith(tempfile.gettempdir())
    assert not Path(args.kwargs["cache_dir"]).exists()


def test_camera_discovery_rejects_unsupported_sources():
    with pytest.raises(ValueError, match="RGB video camera"):
        studio.video_cameras({"codebase_version": "v3.0", "fps": 30, "features": {}})
    with pytest.raises(ValueError, match="v3"):
        studio.video_cameras({"codebase_version": "v2.1"})


def test_usage_costs_use_reported_runtime_and_keep_unknown_charges_separate():
    from reporting import usage_costs

    job = SimpleNamespace(durations=SimpleNamespace(running_secs=61))
    result = usage_costs(
        job,
        {
            "manifest": {"costs": {"cpu_rate": 0.0005}},
            "usage": {"measured_inference_usd": 0.2, "reserved_including_uncertain_usd": 0.24},
        },
    )
    assert result["cpu_usd"] == 0.001
    assert result["total_usd"] == pytest.approx(0.201)
    assert result["uncertain_inference_usd"] == pytest.approx(0.04)
    assert usage_costs(SimpleNamespace(), {})["total_usd"] is None


def test_visualizer_links_only_target_materialized_copy_annotations():
    from reporting import annotation_url

    report = {
        "manifest": {"mode": studio.COPY, "output": "alice/out"},
        "status": "partial",
        "written_episode_ids": [12, 5],
    }
    assert annotation_url(report) == "https://lerobot-visualize-dataset.hf.space/alice/out/5?tab=annotations"
    report["manifest"]["mode"] = studio.EXPORT
    assert annotation_url(report) is None
    report["manifest"]["mode"] = studio.COPY
    report["written_episode_ids"] = []
    assert annotation_url(report) is None


def test_cached_usage_and_deadline_stop_before_a_paid_request(monkeypatch):
    import provider

    budget = Budget(1, 1, 2, cached_input_rate=0.2)
    reserved = budget.reserve(100)
    budget.reconcile(reserved, {"prompt_tokens": 10, "completion_tokens": 20, "cached_tokens": 6})
    assert budget.summary()["measured_inference_usd"] == pytest.approx(0.0000452)
    assert budget.summary()["input_tokens"] == 10
    assert budget.summary()["cached_input_tokens"] == 6
    post = Mock()
    monkeypatch.setattr(provider.requests, "post", post)
    client = provider.Provider("token", studio.MODEL, budget, deadline=time.time() - 1)
    with pytest.raises(provider.DeadlineExceededError):
        client.generate([[{"role": "user", "content": "hello"}]], 100, 0.2)
    post.assert_not_called()


def test_five_minute_run_has_job_and_worker_deadlines():
    api = api_mock()
    studio.launch(
        plan(),
        "alice/output",
        30,
        True,
        "token",
        space="lerobot/annotation-studio",
        api_factory=lambda **_: api,
        timeout_minutes=5,
    )
    kwargs = api.run_job.call_args.kwargs
    assert kwargs["timeout"] == "5m"
    assert "ANNOTATION_DEADLINE_EPOCH" in str(kwargs["command"])
    assert "270" in str(kwargs["command"])


def test_parallel_reservations_respect_hosted_context_and_small_budget():
    from provider import Provider

    budget = Budget(5, 0.42, 3, context_length=1_000_000)
    client = Provider("token", studio.MODEL, budget, concurrency=32)
    assert client.concurrency == 11
    reservations = [budget.reserve(4096) for _ in range(client.concurrency)]
    assert budget.reserved < budget.limit
    with pytest.raises(BudgetExceededError):
        budget.reserve(4096)
    budget.reconcile(reservations[0], {"prompt_tokens": 100, "completion_tokens": 40})
    budget.reserve(4096)
    assert (
        Provider(
            "token", studio.MODEL, Budget(30, 0.42, 3, context_length=1_000_000), concurrency=32
        ).concurrency
        == 32
    )


@pytest.mark.parametrize("key", list(MODELS))
def test_selected_model_pins_route_prices_and_generation_parameters(key, monkeypatch):
    import provider

    profile = model_profile(key)
    configured = plan()
    configured["costs"] = studio.estimate(selection(), ["Subtasks"], 0.1, 1, 0.0005, model_key=key)
    api = api_mock()
    api.model_info.return_value.inference_provider_mapping = {
        profile["provider"]: SimpleNamespace(status="live")
    }
    studio.launch(
        configured,
        "alice/selected",
        30,
        True,
        "secret",
        space="lerobot/annotation-studio",
        api_factory=lambda **_: api,
    )
    manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"])
    assert manifest["model"] == profile["route"]
    assert manifest["costs"]["input_rate"] == profile["input_rate"]
    assert manifest["request_concurrency"] == profile["concurrency"]
    api.model_info.assert_called_once_with(profile["model_id"], expand=["inferenceProviderMapping"])
    post = Mock(
        return_value=SimpleNamespace(
            status_code=200,
            json=lambda: {
                "choices": [{"message": {"content": "{}"}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": 100, "completion_tokens": 10},
            },
        )
    )
    monkeypatch.setattr(provider.requests, "post", post)
    client = provider.Provider(
        "secret",
        manifest["model"],
        Budget(10, profile["input_rate"], profile["output_rate"]),
        generation_options=manifest["generation_options"],
    )
    client.generate([[{"role": "user", "content": "JSON please"}]], 100, 0.2)
    body = post.call_args.kwargs["json"]
    assert body["model"] == profile["route"]
    assert {
        k: v for k, v in body.items() if k not in {"model", "messages", "max_tokens", "temperature"}
    } == profile["generation_options"]


def test_tiered_prices_reserve_maximum_and_reconcile_actual_prompt_tier():
    profile = model_profile("minimax-m3")
    budget = Budget(10, 0.3, 1.2, 0.06, 1_000_000, profile["pricing_tiers"])
    assert budget.reservation_cost(1000) == pytest.approx(0.6024)
    first = budget.reserve(1000)
    budget.reconcile(first, {"prompt_tokens": 1000, "completion_tokens": 1000, "cached_tokens": 500})
    assert budget.measured == pytest.approx(0.00138)
    second = budget.reserve(1000)
    budget.reconcile(second, {"prompt_tokens": 524288, "completion_tokens": 1000})
    assert budget.measured == pytest.approx(0.00138 + 524288 * 0.6 / 1e6 + 0.0024)


def test_unknown_model_and_stale_model_prices_cannot_launch():
    with pytest.raises(ValueError, match="supported vision model"):
        studio.prepare(
            "source/robot", 0, 1, ["Subtasks"], None, 0.1, 1, studio.COPY, model_key="invented/model"
        )
    api = api_mock()
    configured = plan()
    configured["costs"]["input_rate"] = 0
    with pytest.raises(ValueError, match="pricing changed"):
        studio.launch(
            configured,
            "alice/output",
            30,
            True,
            "secret",
            space="lerobot/annotation-studio",
            api_factory=lambda **_: api,
        )
    api.run_job.assert_not_called()


@pytest.mark.parametrize("visibility", ["private", "public"])
def test_launch_pins_output_visibility(visibility):
    api = api_mock()
    studio.launch(
        plan(),
        "alice/new-output",
        5,
        True,
        "token",
        space="lerobot/annotation-studio",
        api_factory=lambda **_: api,
        visibility=visibility,
    )
    manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"])
    assert manifest["visibility"] == visibility


@pytest.mark.parametrize("visibility", [None, "private", "public"])
def test_output_visibility_applies_to_parquet_dataset(visibility):
    import worker

    api = Mock()
    api.dataset_info.return_value.sha = "pinned"
    manifest = {
        "dataset": "alice/source",
        "dataset_revision": "pinned",
        "output": "alice/new",
        "mode": studio.COPY,
    }
    if visibility is not None:
        manifest["visibility"] = visibility
    worker.create_output_repo(api, manifest)
    api.create_repo.assert_called_once_with(
        "alice/new", repo_type="dataset", private=visibility != "public", exist_ok=False
    )
    api.duplicate_repo.assert_not_called()
    api.update_repo_settings.assert_not_called()


def test_invalid_visibility_cannot_launch_or_create_output():
    from worker import create_output_repo

    api = api_mock()
    with pytest.raises(ValueError, match="visibility"):
        studio.launch(
            plan(), "alice/out", 5, True, "token", api_factory=lambda **_: api, visibility="unknown"
        )
    with pytest.raises(ValueError, match="visibility"):
        create_output_repo(api, {"visibility": "unknown"})
    api.run_job.assert_not_called()
    api.create_repo.assert_not_called()
    api.duplicate_repo.assert_not_called()


def test_export_cannot_start_new_jobs_or_create_repositories():
    from worker import create_output_repo

    api = api_mock()
    legacy = plan()
    legacy["mode"] = studio.EXPORT
    with pytest.raises(ValueError, match="Parquet"):
        studio.prepare("alice/source", 0, 1, ["Subtasks"], None, 0.1, 1, studio.EXPORT)
    with pytest.raises(ValueError, match="Parquet"):
        studio.launch(legacy, "alice/out", 5, True, "token", api_factory=lambda **_: api)
    with pytest.raises(ValueError, match="Parquet"):
        create_output_repo(api, {"mode": studio.EXPORT})
    api.run_job.assert_not_called()
    api.create_repo.assert_not_called()