chore: import upstream snapshot with attribution

tests/dev/test_check_function_signatures.py

@@ -0,0 +1,230 @@
+import ast
+
+from dev.check_function_signatures import check_signature_compatibility
+
+
+def test_no_changes():
+    old_code = "def func(a, b=1): pass"
+    new_code = "def func(a, b=1): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 0
+
+
+def test_positional_param_removed():
+    old_code = "def func(a, b, c): pass"
+    new_code = "def func(a, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "Positional param 'c' was removed."
+    assert errors[0].param_name == "c"
+
+
+def test_positional_param_renamed():
+    old_code = "def func(a, b): pass"
+    new_code = "def func(x, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert "Positional param order/name changed: 'a' -> 'x'." in errors[0].message
+    assert errors[0].param_name == "x"
+
+
+def test_only_first_positional_rename_flagged():
+    old_code = "def func(a, b, c, d): pass"
+    new_code = "def func(x, y, z, w): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert "Positional param order/name changed: 'a' -> 'x'." in errors[0].message
+
+
+def test_optional_positional_became_required():
+    old_code = "def func(a, b=1): pass"
+    new_code = "def func(a, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "Optional positional param 'b' became required."
+    assert errors[0].param_name == "b"
+
+
+def test_multiple_optional_became_required():
+    old_code = "def func(a, b=1, c=2): pass"
+    new_code = "def func(a, b, c): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 2
+    assert errors[0].message == "Optional positional param 'b' became required."
+    assert errors[1].message == "Optional positional param 'c' became required."
+
+
+def test_new_required_positional_param():
+    old_code = "def func(a): pass"
+    new_code = "def func(a, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "New required positional param 'b' added."
+    assert errors[0].param_name == "b"
+
+
+def test_new_optional_positional_param_allowed():
+    old_code = "def func(a): pass"
+    new_code = "def func(a, b=1): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 0
+
+
+def test_keyword_only_param_removed():
+    old_code = "def func(*, a, b): pass"
+    new_code = "def func(*, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "Keyword-only param 'a' was removed."
+    assert errors[0].param_name == "a"
+
+
+def test_multiple_keyword_only_removed():
+    old_code = "def func(*, a, b, c): pass"
+    new_code = "def func(*, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 2
+    error_messages = {e.message for e in errors}
+    assert "Keyword-only param 'a' was removed." in error_messages
+    assert "Keyword-only param 'c' was removed." in error_messages
+
+
+def test_optional_keyword_only_became_required():
+    old_code = "def func(*, a=1): pass"
+    new_code = "def func(*, a): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "Keyword-only param 'a' became required."
+    assert errors[0].param_name == "a"
+
+
+def test_new_required_keyword_only_param():
+    old_code = "def func(*, a): pass"
+    new_code = "def func(*, a, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "New required keyword-only param 'b' added."
+    assert errors[0].param_name == "b"
+
+
+def test_new_optional_keyword_only_allowed():
+    old_code = "def func(*, a): pass"
+    new_code = "def func(*, a, b=1): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 0
+
+
+def test_complex_mixed_violations():
+    old_code = "def func(a, b=1, *, c, d=2): pass"
+    new_code = "def func(x, b, *, c=3, e): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 3
+    error_messages = [e.message for e in errors]
+    assert any("Positional param order/name changed: 'a' -> 'x'." in msg for msg in error_messages)
+    assert any("Keyword-only param 'd' was removed." in msg for msg in error_messages)
+    assert any("New required keyword-only param 'e' added." in msg for msg in error_messages)
+
+
+def test_parameter_error_has_location_info():
+    old_code = "def func(a): pass"
+    new_code = "def func(b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].lineno == 1
+    assert errors[0].col_offset > 0
+
+
+def test_async_function_compatibility():
+    old_code = "async def func(a, b=1): pass"
+    new_code = "async def func(a, b): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert errors[0].message == "Optional positional param 'b' became required."
+
+
+def test_positional_only_compatibility():
+    old_code = "def func(a, /): pass"
+    new_code = "def func(b, /): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert "Positional param order/name changed: 'a' -> 'b'." in errors[0].message
+
+
+def test_rename_stops_further_positional_checks():
+    old_code = "def func(a, b=1, c=2): pass"
+    new_code = "def func(x, b, c): pass"
+
+    old_tree = ast.parse(old_code)
+    new_tree = ast.parse(new_code)
+    errors = check_signature_compatibility(old_tree.body[0], new_tree.body[0])
+
+    assert len(errors) == 1
+    assert "Positional param order/name changed: 'a' -> 'x'." in errors[0].message

tests/dev/test_check_init_py.py

@@ -0,0 +1,226 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+def get_check_init_py_script() -> Path:
+    return Path(__file__).resolve().parents[2] / "dev" / "check_init_py.py"
+
+
+@pytest.fixture
+def temp_git_repo(tmp_path: Path) -> Path:
+    subprocess.check_call(["git", "init"], cwd=tmp_path)
+    subprocess.check_call(["git", "config", "user.email", "test@example.com"], cwd=tmp_path)
+    subprocess.check_call(["git", "config", "user.name", "Test User"], cwd=tmp_path)
+    return tmp_path
+
+
+def test_exits_with_0_when_all_directories_have_init_py(temp_git_repo: Path) -> None:
+    mlflow_dir = temp_git_repo / "mlflow"
+    test_package_dir = mlflow_dir / "test_package"
+    test_package_dir.mkdir(parents=True)
+
+    (mlflow_dir / "__init__.py").touch()
+    (test_package_dir / "__init__.py").touch()
+    (test_package_dir / "test_module.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 0
+    assert result.stdout == ""
+
+
+def test_exits_with_1_when_directories_missing_init_py(temp_git_repo: Path) -> None:
+    mlflow_dir = temp_git_repo / "mlflow"
+    test_package_dir = mlflow_dir / "test_package"
+    test_package_dir.mkdir(parents=True)
+
+    (test_package_dir / "test_module.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 1
+    assert (
+        "Error: The following directories contain Python files but lack __init__.py:"
+        in result.stdout
+    )
+    assert "mlflow" in result.stdout
+    assert "mlflow/test_package" in result.stdout
+
+
+def test_exits_with_0_when_no_python_files_exist(temp_git_repo: Path) -> None:
+    mlflow_dir = temp_git_repo / "mlflow"
+    js_dir = mlflow_dir / "server" / "js"
+    js_dir.mkdir(parents=True)
+
+    (js_dir / "main.js").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 0
+    assert result.stdout == ""
+
+
+def test_identifies_only_directories_missing_init_py(temp_git_repo: Path) -> None:
+    mlflow_dir = temp_git_repo / "mlflow"
+    package1_dir = mlflow_dir / "package1"
+    package2_dir = mlflow_dir / "package2"
+    package1_dir.mkdir(parents=True)
+    package2_dir.mkdir(parents=True)
+
+    (mlflow_dir / "__init__.py").touch()
+    (package1_dir / "__init__.py").touch()
+
+    (package1_dir / "module1.py").touch()
+    (package2_dir / "module2.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 1
+    assert (
+        "Error: The following directories contain Python files but lack __init__.py:"
+        in result.stdout
+    )
+    assert "mlflow/package2" in result.stdout
+    assert "mlflow/package1" not in result.stdout
+
+
+def test_checks_tests_directory_for_missing_init_py(temp_git_repo: Path) -> None:
+    tests_dir = temp_git_repo / "tests"
+    test_package_dir = tests_dir / "test_package"
+    test_package_dir.mkdir(parents=True)
+
+    # Only test files (starting with test_) are checked
+    (test_package_dir / "test_module.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 1
+    assert (
+        "Error: The following directories contain Python files but lack __init__.py:"
+        in result.stdout
+    )
+    assert "tests/test_package" in result.stdout
+    assert "tests" in result.stdout  # Parent directory also needs __init__.py
+
+
+def test_exits_with_0_when_tests_directories_have_init_py(temp_git_repo: Path) -> None:
+    tests_dir = temp_git_repo / "tests"
+    test_package_dir = tests_dir / "test_package"
+    test_package_dir.mkdir(parents=True)
+
+    (tests_dir / "__init__.py").touch()
+    (test_package_dir / "__init__.py").touch()
+    (test_package_dir / "test_module.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 0
+    assert result.stdout == ""
+
+
+def test_ignores_non_test_files_in_tests_directory(temp_git_repo: Path) -> None:
+    tests_dir = temp_git_repo / "tests"
+    test_package_dir = tests_dir / "test_package"
+    test_package_dir.mkdir(parents=True)
+
+    # Non-test file (doesn't start with test_) should be ignored
+    (test_package_dir / "helper.py").touch()
+    (test_package_dir / "utils.py").touch()
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    # Should pass since no test files exist
+    assert result.returncode == 0
+    assert result.stdout == ""
+
+
+def test_checks_all_parent_directories(temp_git_repo: Path) -> None:
+    mlflow_dir = temp_git_repo / "mlflow"
+    deep_dir = mlflow_dir / "level1" / "level2" / "level3"
+    deep_dir.mkdir(parents=True)
+
+    # Create a Python file deep in the hierarchy
+    (deep_dir / "module.py").touch()
+
+    # Only add __init__.py to some directories
+    (mlflow_dir / "__init__.py").touch()
+    (mlflow_dir / "level1" / "__init__.py").touch()
+    # Missing: level2 and level3 __init__.py files
+
+    subprocess.check_call(["git", "add", "."], cwd=temp_git_repo)
+    subprocess.check_call(["git", "commit", "-m", "Initial commit"], cwd=temp_git_repo)
+
+    result = subprocess.run(
+        [sys.executable, get_check_init_py_script()],
+        capture_output=True,
+        text=True,
+        cwd=temp_git_repo,
+    )
+
+    assert result.returncode == 1
+    assert (
+        "Error: The following directories contain Python files but lack __init__.py:"
+        in result.stdout
+    )
+    # Both level2 and level3 should be reported as missing __init__.py
+    assert "mlflow/level1/level2" in result.stdout
+    assert "mlflow/level1/level2/level3" in result.stdout

tests/dev/test_dev_stubs.py

@@ -0,0 +1,122 @@
+import importlib.util
+import json
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlflow.assistant.providers.claude_code import ClaudeCodeProvider
+from mlflow.assistant.types import EventType
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+DEV_STUBS = REPO_ROOT / "dev" / "dev_stubs"
+CLAUDE_CLI = DEV_STUBS / "claude_cli.py"
+
+
+def _load(name: str, path: Path):
+    # dev/ is not an installed package; load this module directly by path. It only
+    # imports stdlib and references siblings as files, so standalone load works.
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    # Register before exec so @dataclass can resolve the module's globals by name.
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+dev_stubs = _load("mlflow_dev_stubs", DEV_STUBS / "__init__.py")
+
+
+# --- claude CLI stub ---------------------------------------------------------
+
+
+def run_claude(*args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, CLAUDE_CLI, *args], capture_output=True, text=True, timeout=30
+    )
+
+
+def test_claude_auth_probe_exits_zero_with_valid_json():
+    result = run_claude("-p", "hi", "--max-turns", "1", "--output-format", "json")
+    assert result.returncode == 0
+    payload = json.loads(result.stdout)
+    assert payload["type"] == "result"
+    assert payload["is_error"] is False
+
+
+def test_claude_stream_json_parses_into_message_and_done_events():
+    result = run_claude(
+        "-p",
+        "hello",
+        "--output-format",
+        "stream-json",
+        "--verbose",
+        "--append-system-prompt",
+        "ignored prompt with --output-format text and --resume decoys",
+    )
+    assert result.returncode == 0
+
+    provider = ClaudeCodeProvider()
+    events = [
+        event
+        for line in result.stdout.splitlines()
+        if line.strip()
+        if (event := provider._parse_message_to_event(json.loads(line))) is not None
+    ]
+    assert [e.type for e in events] == [EventType.MESSAGE, EventType.DONE]
+    assert events[0].data["message"]["content"][0]["text"]
+    assert events[1].data["session_id"]
+
+
+def test_claude_resume_reuses_session_id():
+    result = run_claude("-p", "hi", "--output-format", "stream-json", "--resume", "sess-abc")
+    assert result.returncode == 0
+    events = [json.loads(line) for line in result.stdout.splitlines() if line.strip()]
+    assert {e["session_id"] for e in events if "session_id" in e} == {"sess-abc"}
+
+
+# --- stub registry -----------------------------------------------------------
+
+
+def test_install_stubs_rejects_unknown_name():
+    with pytest.raises(ValueError, match="Unknown stub"):
+        dev_stubs.install_stubs(["not-a-stub"])
+
+
+def test_install_stubs_cleans_up_temp_dirs_on_partial_failure(monkeypatch):
+    tmp_root = Path(tempfile.gettempdir())
+    before = set(tmp_root.glob("mlflow-dev-stub-bin-*"))
+
+    def boom(_result):
+        raise RuntimeError("boom")
+
+    monkeypatch.setitem(dev_stubs._INSTALLERS, "boom", boom)
+    with pytest.raises(RuntimeError, match="boom"):
+        # claude stages a temp dir, then `boom` fails -- the dir must not leak.
+        dev_stubs.install_stubs(["claude", "boom"])
+
+    assert set(tmp_root.glob("mlflow-dev-stub-bin-*")) == before
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="the claude shim is a POSIX shell script; run_dev_server (its consumer) is POSIX-only",
+)
+def test_install_claude_stages_working_shim():
+    result = dev_stubs.install_stubs(["claude"])
+    try:
+        (shim_dir,) = result.path_prepend
+        probe = subprocess.run(
+            [shim_dir / "claude", "-p", "hi", "--output-format", "json"],
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        assert probe.returncode == 0
+        assert json.loads(probe.stdout)["type"] == "result"
+    finally:
+        for path in result.cleanup_paths:
+            shutil.rmtree(path, ignore_errors=True)

tests/dev/test_remove_experimental_decorators.py

@@ -0,0 +1,206 @@
+import subprocess
+import sys
+from pathlib import Path
+
+SCRIPT_PATH = "dev/remove_experimental_decorators.py"
+
+
+def test_script_with_specific_file(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0")
+def func():
+    pass
+""")
+
+    output = subprocess.check_output(
+        [sys.executable, SCRIPT_PATH, "--dry-run", test_file], text=True
+    )
+
+    assert "Would remove" in output
+    assert "experimental(version='1.0.0')" in output
+    assert (
+        test_file.read_text()
+        == """
+@experimental(version="1.0.0")
+def func():
+    pass
+"""
+    )
+
+
+def test_script_without_files() -> None:
+    subprocess.check_call([sys.executable, SCRIPT_PATH, "--dry-run"])
+
+
+def test_script_removes_decorator_without_dry_run(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0")
+def func():
+    pass
+""")
+
+    subprocess.check_call([sys.executable, SCRIPT_PATH, test_file])
+    content = test_file.read_text()
+    assert (
+        content
+        == """
+def func():
+    pass
+"""
+    )
+
+
+def test_script_with_multiline_decorator(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(
+    version="1.0.0",
+)
+def func():
+    pass
+""")
+
+    output = subprocess.check_output([sys.executable, SCRIPT_PATH, test_file], text=True)
+    assert "Removed" in output
+    assert (
+        test_file.read_text()
+        == """
+def func():
+    pass
+"""
+    )
+
+
+def test_script_with_multiple_decorators(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0")
+def func1():
+    pass
+
+@experimental(version="1.1.0")
+class MyClass:
+    @experimental(version="1.2.0")
+    def method(self):
+        pass
+
+def regular_func():
+    pass
+""")
+
+    output = subprocess.check_output([sys.executable, SCRIPT_PATH, test_file], text=True)
+    assert output.count("Removed") == 3  # Should remove all 3 decorators
+    content = test_file.read_text()
+    assert (
+        content
+        == """
+def func1():
+    pass
+
+class MyClass:
+    def method(self):
+        pass
+
+def regular_func():
+    pass
+"""
+    )
+
+
+def test_script_with_cutoff_days_argument(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0")
+def func():
+    pass
+""")
+
+    # Test with a very large cutoff (should not remove anything)
+    output = subprocess.check_output(
+        [sys.executable, SCRIPT_PATH, "--cutoff-days", "9999", "--dry-run", test_file], text=True
+    )
+    assert "Would remove" not in output
+
+    # Test with default cutoff (180 days, should remove old decorators)
+    output = subprocess.check_output(
+        [sys.executable, SCRIPT_PATH, "--dry-run", test_file], text=True
+    )
+    assert "Would remove" in output
+
+    # Test with explicit cutoff of 180 days
+    output = subprocess.check_output(
+        [sys.executable, SCRIPT_PATH, "--cutoff-days", "180", "--dry-run", test_file], text=True
+    )
+    assert "Would remove" in output
+
+
+def test_skip_preserves_decorator(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0", skip=True)
+def func():
+    pass
+""")
+    original_content = test_file.read_text()
+
+    subprocess.check_call([sys.executable, SCRIPT_PATH, test_file])
+    assert test_file.read_text() == original_content
+
+
+def test_skip_dry_run_shows_skipped(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0", skip=True)
+def func():
+    pass
+""")
+
+    output = subprocess.check_output(
+        [sys.executable, SCRIPT_PATH, "--dry-run", test_file], text=True
+    )
+    assert "Skipped (skip=True)" in output
+    assert "Would remove" not in output
+
+
+def test_skip_false_still_removed(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0", skip=False)
+def func():
+    pass
+""")
+
+    subprocess.check_call([sys.executable, SCRIPT_PATH, test_file])
+    content = test_file.read_text()
+    assert "@experimental" not in content
+
+
+def test_skip_mixed_file(tmp_path: Path) -> None:
+    test_file = tmp_path / "test.py"
+    test_file.write_text("""
+@experimental(version="1.0.0", skip=True)
+def func_keep():
+    pass
+
+@experimental(version="1.0.0")
+def func_remove():
+    pass
+""")
+
+    subprocess.check_call([sys.executable, SCRIPT_PATH, test_file])
+    content = test_file.read_text()
+    assert "func_keep" in content
+    assert "func_remove" in content
+    assert (
+        content
+        == """
+@experimental(version="1.0.0", skip=True)
+def func_keep():
+    pass
+
+def func_remove():
+    pass
+"""
+    )

tests/dev/test_update_mlflow_versions.py

@@ -0,0 +1,170 @@
+import difflib
+import re
+from pathlib import Path
+
+import pytest
+from packaging.version import Version
+
+from dev.update_mlflow_versions import (
+    get_current_py_version,
+    replace_java,
+    replace_java_pom_xml,
+    replace_pyproject_toml,
+    replace_python,
+    replace_r,
+    replace_ts,
+)
+
+# { filename: expected lines changed }
+_JAVA_FILES = {}
+
+_JAVA_XML_FILES = {
+    "mlflow/java/pom.xml": {
+        6: "  <version>{new_version}</version>",
+        52: "    <mlflow-version>{new_version}</mlflow-version>",
+    },
+    "mlflow/java/client/pom.xml": {
+        8: "    <version>{new_version}</version>",
+    },
+    "mlflow/java/spark_2.12/pom.xml": {
+        4: "  <version>{new_version}</version>",
+        18: "    <version>{new_version}</version>",
+    },
+    "mlflow/java/spark_2.13/pom.xml": {
+        4: "  <version>{new_version}</version>",
+        18: "    <version>{new_version}</version>",
+    },
+}
+
+_TS_FILES = {
+    "mlflow/server/js/src/common/constants.tsx": {
+        12: "export const Version = '{new_version}';",
+    },
+    "docs/src/constants.ts": {
+        1: "export const Version = '{new_version}';",
+    },
+}
+
+_PYTHON_FILES = {
+    "mlflow/version.py": {
+        5: 'VERSION = "{new_version}"',
+    }
+}
+
+_PYPROJECT_TOML_FILES = {
+    "pyproject.toml": {
+        12: 'version = "{new_version}"',
+    },
+    "pyproject.release.toml": {
+        12: 'version = "{new_version}"',
+        30: '  "mlflow-skinny=={new_version}",',
+        31: '  "mlflow-tracing=={new_version}",',
+    },
+    "libs/skinny/pyproject.toml": {
+        10: 'version = "{new_version}"',
+    },
+    "libs/tracing/pyproject.toml": {
+        10: 'version = "{new_version}"',
+    },
+}
+
+_R_FILES = {
+    "mlflow/R/mlflow/DESCRIPTION": {
+        4: "Version: {new_version}",
+    }
+}
+
+_DIFF_REGEX = re.compile(r"--- (\d+,?\d*) ----")
+
+old_version = Version(get_current_py_version())
+_NEW_PY_VERSION = f"{old_version.major}.{old_version.minor}.{old_version.micro + 1}"
+
+
+def copy_and_run_change_func(monkeypatch, tmp_path, paths_to_copy, replace_func, new_version):
+    for path in paths_to_copy:
+        copy_path = tmp_path / path
+        copy_path.parent.mkdir(parents=True, exist_ok=True)
+        copy_path.write_text(path.read_text())
+
+    with monkeypatch.context() as m:
+        m.chdir(tmp_path)
+
+        # pyproject.toml replace doesn't search for the old version,
+        # it just replaces the version line with the new version.
+        if replace_func == replace_pyproject_toml:
+            replace_func(new_version, paths_to_copy)
+        else:
+            replace_func(str(old_version), new_version, paths_to_copy)
+
+
+@pytest.mark.parametrize(
+    ("replace_func", "expect_dict", "new_py_version", "expected_new_version"),
+    [
+        (replace_java, _JAVA_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+        (replace_java, _JAVA_FILES, _NEW_PY_VERSION + ".dev0", _NEW_PY_VERSION + "-SNAPSHOT"),
+        (replace_java, _JAVA_FILES, _NEW_PY_VERSION + "rc1", _NEW_PY_VERSION + "-SNAPSHOT"),
+        (replace_java_pom_xml, _JAVA_XML_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+        (
+            replace_java_pom_xml,
+            _JAVA_XML_FILES,
+            _NEW_PY_VERSION + ".dev0",
+            _NEW_PY_VERSION + "-SNAPSHOT",
+        ),
+        (
+            replace_java_pom_xml,
+            _JAVA_XML_FILES,
+            _NEW_PY_VERSION + "rc1",
+            _NEW_PY_VERSION + "-SNAPSHOT",
+        ),
+        (replace_ts, _TS_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+        (replace_python, _PYTHON_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+        (replace_pyproject_toml, _PYPROJECT_TOML_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+        (replace_r, _R_FILES, _NEW_PY_VERSION, _NEW_PY_VERSION),
+    ],
+)
+def test_update_mlflow_versions(
+    monkeypatch, tmp_path, replace_func, expect_dict, new_py_version, expected_new_version
+):
+    paths_to_change = [Path(filename) for filename in expect_dict]
+    copy_and_run_change_func(
+        monkeypatch,
+        tmp_path,
+        # always copy version.py since we need it in get_current_py_version()
+        paths_to_change + [Path("mlflow/version.py")],
+        replace_func,
+        new_py_version,
+    )
+
+    # diff files
+    for filename, expected_changes in expect_dict.items():
+        old_file = Path(filename).read_text().splitlines()
+        new_file = (tmp_path / filename).read_text().splitlines()
+        diff = list(difflib.context_diff(old_file, new_file, n=0))
+        changed_lines = _parse_diff_line(diff)
+
+        formatted_expected_changes = {
+            line_num: change.format(new_version=expected_new_version)
+            for line_num, change in expected_changes.items()
+        }
+
+        assert changed_lines == formatted_expected_changes
+
+
+def _parse_diff_line(diff: list[str]) -> dict[int, str]:
+    diff_lines = {}
+    for idx, line in enumerate(diff):
+        match = _DIFF_REGEX.search(line)
+        if not match:
+            continue
+
+        if "," in match.group(1):
+            # multi-line change is represented as [(start,end), line1, line2, ...]
+            start, end = map(int, match.group(1).split(","))
+            for i in range(start, end + 1):
+                # the [2:] is to cut out the "! " at the beginning of diff lines
+                diff_lines[i] = diff[idx + (i - start) + 1][2:]
+        else:
+            # single-line change
+            diff_lines[int(match.group(1))] = diff[idx + 1][2:]
+
+    return diff_lines

tests/dev/test_update_model_catalog.py

@@ -0,0 +1,795 @@
+import json
+import sys
+from datetime import date
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "dev"))
+
+from update_model_catalog import (
+    _extract_long_context_pricing,
+    _extract_modality_pricing,
+    _extract_service_tiers,
+    _extract_tool_pricing,
+    _is_deprecated,
+    _migrate_legacy_pricing,
+    _normalize_provider,
+    _transform_entry,
+    convert,
+)
+
+
+@pytest.mark.parametrize(
+    ("provider", "expected"),
+    [
+        ("openai", "openai"),
+        ("anthropic", "anthropic"),
+        ("vertex_ai", "vertex_ai"),
+        ("vertex_ai-anthropic", "vertex_ai"),
+        ("vertex_ai-llama_models", "vertex_ai"),
+        ("vertex_ai-chat-models", "vertex_ai"),
+        ("bedrock", "bedrock"),
+    ],
+)
+def test_normalize_provider(provider, expected):
+    assert _normalize_provider(provider) == expected
+
+
+def test_transform_entry_chat_model():
+    info = {
+        "mode": "chat",
+        "input_cost_per_token": 3e-6,
+        "output_cost_per_token": 1.5e-5,
+        "cache_read_input_token_cost": 3e-7,
+        "cache_creation_input_token_cost": 3.75e-6,
+        "max_input_tokens": 200000,
+        "max_output_tokens": 64000,
+        "max_tokens": 64000,
+        "supports_function_calling": True,
+        "supports_vision": True,
+        "supports_reasoning": True,
+        "supports_prompt_caching": True,
+        "supports_response_schema": True,
+    }
+    result = _transform_entry(info)
+    assert result == {
+        "mode": "chat",
+        "context_window": {"max_input": 200000, "max_output": 64000, "max_tokens": 64000},
+        "pricing": {
+            "input_per_million_tokens": 3.0,
+            "output_per_million_tokens": 15.0,
+            "cache_read_per_million_tokens": 0.3,
+            "cache_write_per_million_tokens": 3.75,
+        },
+        "capabilities": {
+            "function_calling": True,
+            "vision": True,
+            "reasoning": True,
+            "prompt_caching": True,
+            "response_schema": True,
+        },
+    }
+
+
+def test_transform_entry_includes_image_generation():
+    info = {"mode": "image_generation", "input_cost_per_token": 1e-6}
+    result = _transform_entry(info)
+    assert result is not None
+    assert result["mode"] == "image_generation"
+
+
+def test_transform_entry_includes_video_generation():
+    info = {"mode": "video_generation", "input_cost_per_token": 1e-6}
+    result = _transform_entry(info)
+    assert result is not None
+    assert result["mode"] == "video_generation"
+
+
+def test_transform_entry_includes_future_deprecation_date():
+    info = {
+        "mode": "chat",
+        "deprecation_date": "2099-01-01",
+    }
+    result = _transform_entry(info)
+    assert result is not None
+    assert result["deprecation_date"] == "2099-01-01"
+
+
+def test_transform_entry_skips_past_deprecation_date():
+    info = {
+        "mode": "chat",
+        "deprecation_date": "2020-01-01",
+    }
+    assert _transform_entry(info) is None
+
+
+def test_is_deprecated_past_date():
+    assert _is_deprecated({"deprecation_date": "2020-01-01"}) is True
+
+
+def test_is_deprecated_future_date():
+    assert _is_deprecated({"deprecation_date": "2099-01-01"}) is False
+
+
+def test_is_deprecated_no_date():
+    assert _is_deprecated({}) is False
+
+
+def test_is_deprecated_invalid_date():
+    assert _is_deprecated({"deprecation_date": "not-a-date"}) is False
+
+
+def test_transform_entry_with_service_tiers():
+    info = {
+        "mode": "chat",
+        "input_cost_per_token": 2e-6,
+        "output_cost_per_token": 8e-6,
+        "cache_read_input_token_cost": 5e-7,
+        "input_cost_per_token_flex": 1e-6,
+        "output_cost_per_token_flex": 4e-6,
+        "cache_read_input_token_cost_flex": 2.5e-7,
+        "input_cost_per_token_priority": 3.5e-6,
+        "output_cost_per_token_priority": 1.4e-5,
+        "input_cost_per_token_batches": 1e-6,
+        "output_cost_per_token_batches": 4e-6,
+    }
+    result = _transform_entry(info)
+    tiers = result["pricing"]["service_tiers"]
+    assert tiers["flex"] == {
+        "input_per_million_tokens": 1.0,
+        "output_per_million_tokens": 4.0,
+        "cache_read_per_million_tokens": 0.25,
+    }
+    assert tiers["priority"] == {
+        "input_per_million_tokens": 3.5,
+        "output_per_million_tokens": 14.0,
+    }
+    assert tiers["batch"] == {
+        "input_per_million_tokens": 1.0,
+        "output_per_million_tokens": 4.0,
+    }
+
+
+def test_transform_entry_with_long_context():
+    info = {
+        "mode": "chat",
+        "input_cost_per_token": 1.25e-6,
+        "output_cost_per_token": 1e-5,
+        "input_cost_per_token_above_200k_tokens": 2.5e-6,
+        "output_cost_per_token_above_200k_tokens": 1.5e-5,
+        "cache_read_input_token_cost_above_200k_tokens": 2.5e-7,
+    }
+    result = _transform_entry(info)
+    long_ctx = result["pricing"]["long_context"]
+    assert len(long_ctx) == 1
+    assert long_ctx[0] == {
+        "threshold_tokens": 200000,
+        "input_per_million_tokens": 2.5,
+        "output_per_million_tokens": 15.0,
+        "cache_read_per_million_tokens": 0.25,
+    }
+
+
+def test_extract_long_context_multiple_thresholds():
+    info = {
+        "input_cost_per_token_above_128k_tokens": 1e-6,
+        "output_cost_per_token_above_128k_tokens": 2e-6,
+        "input_cost_per_token_above_256k_tokens": 2e-6,
+        "output_cost_per_token_above_256k_tokens": 4e-6,
+    }
+    result = _extract_long_context_pricing(info)
+    assert len(result) == 2
+    assert result[0]["threshold_tokens"] == 128000
+    assert result[1]["threshold_tokens"] == 256000
+
+
+def test_extract_service_tiers_empty_when_no_tiers():
+    info = {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6}
+    assert _extract_service_tiers(info) == {}
+
+
+def test_extract_modality_pricing():
+    info = {
+        "input_cost_per_audio_token": 7e-7,
+        "output_cost_per_audio_token": 1.1e-6,
+        "cache_read_input_audio_token_cost": 2e-7,
+        "cache_creation_input_audio_token_cost": 4e-7,
+    }
+    assert _extract_modality_pricing(info) == {
+        "audio": {
+            "input_per_million_tokens": 0.7,
+            "output_per_million_tokens": 1.1,
+            "cache_read_per_million_tokens": 0.2,
+            "cache_write_per_million_tokens": 0.4,
+        }
+    }
+
+
+def test_extract_modality_pricing_skips_reasoning():
+    info = {
+        "input_cost_per_audio_token": 7e-7,
+        "output_cost_per_reasoning_token": 4e-7,
+    }
+    assert _extract_modality_pricing(info) == {"audio": {"input_per_million_tokens": 0.7}}
+
+
+def test_extract_modality_pricing_mixed():
+    info = {
+        "input_cost_per_audio_token": 7e-7,
+        "input_cost_per_video_per_second": 0.0007,
+        "output_cost_per_video_per_second": 0.0014,
+    }
+    assert _extract_modality_pricing(info) == {
+        "audio": {"input_per_million_tokens": 0.7},
+        "video": {"input_per_second": 0.0007, "output_per_second": 0.0014},
+    }
+
+
+def test_extract_tool_pricing():
+    info = {
+        "computer_use_input_cost_per_1k_tokens": 0.00225,
+        "computer_use_output_cost_per_1k_tokens": 0.009,
+        "search_context_cost_per_query": {
+            "search_context_size_low": 0.01,
+            "search_context_size_medium": 0.01,
+            "search_context_size_high": 0.01,
+        },
+        "tool_use_system_prompt_tokens": 159,
+    }
+    assert _extract_tool_pricing(info) == {
+        "computer_use": {
+            "input_per_million_tokens": 2.25,
+            "output_per_million_tokens": 9.0,
+        },
+        "search_context_per_query": {
+            "search_context_size_low": 0.01,
+            "search_context_size_medium": 0.01,
+            "search_context_size_high": 0.01,
+        },
+        "tool_use_system_prompt_tokens": 159,
+    }
+
+
+def test_transform_entry_with_modality_and_tool_pricing():
+    info = {
+        "mode": "chat",
+        "input_cost_per_token": 1e-7,
+        "output_cost_per_token": 4e-7,
+        "input_cost_per_audio_token": 7e-7,
+        "computer_use_input_cost_per_1k_tokens": 0.00225,
+        "tool_use_system_prompt_tokens": 159,
+    }
+    result = _transform_entry(info)
+    assert result["pricing"]["modality"] == {"audio": {"input_per_million_tokens": 0.7}}
+    assert result["pricing"]["tooling"] == {
+        "computer_use": {"input_per_million_tokens": 2.25},
+        "tool_use_system_prompt_tokens": 159,
+    }
+
+
+def test_convert_end_to_end(tmp_path):
+    input_data = {
+        "sample_spec": {"mode": "chat"},
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+            "output_cost_per_token": 1e-5,
+            "max_input_tokens": 128000,
+            "max_output_tokens": 16384,
+            "supports_function_calling": True,
+            "supports_vision": True,
+        },
+        "openai/gpt-4o-mini": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 1.5e-7,
+            "output_cost_per_token": 6e-7,
+            "max_input_tokens": 128000,
+            "max_output_tokens": 16384,
+        },
+        "claude-3-5-sonnet": {
+            "litellm_provider": "anthropic",
+            "mode": "chat",
+            "input_cost_per_token": 3e-6,
+            "output_cost_per_token": 1.5e-5,
+        },
+        "dall-e-3": {
+            "litellm_provider": "openai",
+            "mode": "image_generation",
+        },
+        "sora": {
+            "litellm_provider": "openai",
+            "mode": "video_generation",
+        },
+        "ft:gpt-4o:org::id": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+        },
+        "bedrock_converse/model": {
+            "litellm_provider": "bedrock_converse",
+            "mode": "chat",
+        },
+    }
+
+    output_dir = tmp_path / "output"
+
+    stats = convert(input_data, output_dir)
+
+    assert stats == {"anthropic": 1, "bedrock": 1, "openai": 4}
+    assert (output_dir / "openai.json").exists()
+    assert (output_dir / "anthropic.json").exists()
+    assert (output_dir / "bedrock.json").exists()
+    assert not (output_dir / "bedrock_converse.json").exists()
+
+    openai_catalog = json.loads((output_dir / "openai.json").read_text())
+    assert openai_catalog["schema_version"] == "1.0"
+    assert "gpt-4o" in openai_catalog["models"]
+    assert "gpt-4o-mini" in openai_catalog["models"]
+    # Fine-tuned models should be excluded
+    assert "ft:gpt-4o:org::id" not in openai_catalog["models"]
+    assert "dall-e-3" in openai_catalog["models"]
+    assert "sora" in openai_catalog["models"]
+
+
+def test_convert_preserves_existing_models(tmp_path):
+
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with a manually-added model
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "custom-model": {
+                "mode": "chat",
+                "pricing": {"input_per_million_tokens": 1.0},
+                "capabilities": {
+                    "function_calling": False,
+                    "vision": False,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    stats = convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    # Both upstream and manually-added models should be present
+    assert "gpt-4o" in catalog["models"]
+    assert "custom-model" in catalog["models"]
+    assert stats["openai"] == 2
+
+
+def test_convert_preserves_community_provider(tmp_path):
+
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with a community-maintained provider
+    community_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "my-model": {
+                "mode": "chat",
+                "capabilities": {
+                    "function_calling": False,
+                    "vision": False,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+            }
+        },
+    }
+    (output_dir / "custom_provider.json").write_text(json.dumps(community_catalog))
+
+    stats = convert(input_data, output_dir)
+
+    # Community provider should be preserved
+    assert (output_dir / "custom_provider.json").exists()
+    assert "custom_provider" in stats
+    assert stats["custom_provider"] == 1
+
+
+def test_convert_skips_deprecated_models(tmp_path):
+    input_data = {
+        "old-model": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "deprecation_date": "2020-01-01",
+        },
+        "new-model": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "deprecation_date": "2099-12-31",
+        },
+    }
+
+    output_dir = tmp_path / "output"
+
+    stats = convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    assert "old-model" not in catalog["models"]
+    assert "new-model" in catalog["models"]
+    assert stats["openai"] == 1
+
+
+def test_convert_upstream_overrides_existing_model(tmp_path):
+
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 9.99e-6,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with old pricing
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "gpt-4o": {
+                "mode": "chat",
+                "pricing": {"input_per_million_tokens": 1.0},
+                "capabilities": {
+                    "function_calling": False,
+                    "vision": False,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    # Upstream price should win
+    assert catalog["models"]["gpt-4o"]["pricing"]["input_per_million_tokens"] == pytest.approx(9.99)
+
+
+def test_convert_sets_last_updated_at_for_new_models(tmp_path):
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    assert catalog["models"]["gpt-4o"]["last_updated_at"] == date.today().isoformat()
+
+
+def test_convert_preserves_last_updated_at_when_entry_unchanged(tmp_path):
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+            "max_input_tokens": 128000,
+            "max_output_tokens": 16384,
+            "supports_function_calling": True,
+            "supports_vision": True,
+            "supports_reasoning": False,
+            "supports_prompt_caching": False,
+            "supports_response_schema": False,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with the same data and an existing last_updated_at
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "gpt-4o": {
+                "mode": "chat",
+                "context_window": {"max_input": 128000, "max_output": 16384},
+                "pricing": {"input_per_million_tokens": 2.5},
+                "capabilities": {
+                    "function_calling": True,
+                    "vision": True,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+                "last_updated_at": "2025-01-01",
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    assert catalog["models"]["gpt-4o"]["last_updated_at"] == "2025-01-01"
+
+
+def test_convert_updates_last_updated_at_when_entry_changes(tmp_path):
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 9.99e-6,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with different pricing and an old last_updated_at
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "gpt-4o": {
+                "mode": "chat",
+                "pricing": {"input_per_million_tokens": 1.0},
+                "capabilities": {
+                    "function_calling": False,
+                    "vision": False,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+                "last_updated_at": "2025-01-01",
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    assert catalog["models"]["gpt-4o"]["last_updated_at"] == date.today().isoformat()
+
+
+def test_convert_sets_last_updated_at_for_unchanged_entry_without_existing_date(tmp_path):
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+            "max_input_tokens": 128000,
+            "max_output_tokens": 16384,
+            "supports_function_calling": True,
+            "supports_vision": True,
+            "supports_reasoning": False,
+            "supports_prompt_caching": False,
+            "supports_response_schema": False,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with the same data but NO last_updated_at (simulates pre-feature catalog)
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "gpt-4o": {
+                "mode": "chat",
+                "context_window": {"max_input": 128000, "max_output": 16384},
+                "pricing": {"input_per_million_tokens": 2.5},
+                "capabilities": {
+                    "function_calling": True,
+                    "vision": True,
+                    "reasoning": False,
+                    "prompt_caching": False,
+                    "response_schema": False,
+                },
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    assert catalog["models"]["gpt-4o"]["last_updated_at"] == date.today().isoformat()
+
+    entry = {
+        "mode": "chat",
+        "pricing": {
+            "input_per_token": 3e-6,
+            "output_per_token": 1.5e-5,
+            "cache_read_per_token": 3e-7,
+            "cache_write_per_token": 3.75e-6,
+        },
+    }
+    result = _migrate_legacy_pricing(entry)
+    assert result["pricing"] == {
+        "input_per_million_tokens": 3.0,
+        "output_per_million_tokens": 15.0,
+        "cache_read_per_million_tokens": 0.3,
+        "cache_write_per_million_tokens": 3.75,
+    }
+
+
+def test_migrate_legacy_pricing_service_tiers():
+    entry = {
+        "mode": "chat",
+        "pricing": {
+            "input_per_token": 2e-6,
+            "output_per_token": 8e-6,
+            "service_tiers": {
+                "batch": {
+                    "input_per_token": 1e-6,
+                    "output_per_token": 4e-6,
+                },
+                "priority": {
+                    "input_per_token": 3e-6,
+                    "output_per_token": 1.2e-5,
+                    "cache_read_per_token": 3e-7,
+                },
+            },
+        },
+    }
+    result = _migrate_legacy_pricing(entry)
+    assert result["pricing"]["input_per_million_tokens"] == pytest.approx(2.0)
+    tiers = result["pricing"]["service_tiers"]
+    assert tiers["batch"] == {
+        "input_per_million_tokens": 1.0,
+        "output_per_million_tokens": 4.0,
+    }
+    assert tiers["priority"] == {
+        "input_per_million_tokens": 3.0,
+        "output_per_million_tokens": 12.0,
+        "cache_read_per_million_tokens": 0.3,
+    }
+
+
+def test_migrate_legacy_pricing_long_context():
+    entry = {
+        "mode": "chat",
+        "pricing": {
+            "input_per_token": 1e-6,
+            "output_per_token": 4e-6,
+            "long_context": [
+                {
+                    "threshold_tokens": 200000,
+                    "input_per_token": 2e-6,
+                    "output_per_token": 8e-6,
+                    "cache_read_per_token": 2e-7,
+                }
+            ],
+        },
+    }
+    result = _migrate_legacy_pricing(entry)
+    ctx = result["pricing"]["long_context"]
+    assert len(ctx) == 1
+    assert ctx[0] == {
+        "threshold_tokens": 200000,
+        "input_per_million_tokens": 2.0,
+        "output_per_million_tokens": 8.0,
+        "cache_read_per_million_tokens": 0.2,
+    }
+
+
+def test_migrate_legacy_pricing_modality():
+    entry = {
+        "mode": "chat",
+        "pricing": {
+            "input_per_token": 1e-7,
+            "output_per_token": 4e-7,
+            "modality": {
+                "audio": {
+                    "input_per_token": 7e-7,
+                    "output_per_token": 1.1e-6,
+                }
+            },
+        },
+    }
+    result = _migrate_legacy_pricing(entry)
+    assert result["pricing"]["modality"] == {
+        "audio": {
+            "input_per_million_tokens": 0.7,
+            "output_per_million_tokens": 1.1,
+        }
+    }
+
+
+def test_migrate_legacy_pricing_noop_when_already_normalized():
+    entry = {
+        "mode": "chat",
+        "pricing": {
+            "input_per_million_tokens": 2.5,
+            "output_per_million_tokens": 10.0,
+        },
+    }
+    result = _migrate_legacy_pricing(entry)
+    assert result["pricing"] == {
+        "input_per_million_tokens": 2.5,
+        "output_per_million_tokens": 10.0,
+    }
+
+
+def test_migrate_legacy_pricing_noop_when_no_pricing():
+    entry = {"mode": "chat", "capabilities": {}}
+    assert _migrate_legacy_pricing(entry) is entry
+
+
+def test_convert_migrates_legacy_pricing_in_preserved_models(tmp_path):
+    input_data = {
+        "gpt-4o": {
+            "litellm_provider": "openai",
+            "mode": "chat",
+            "input_cost_per_token": 2.5e-6,
+        },
+    }
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    # Pre-populate with a community model that uses the legacy per-token format
+    existing_catalog = {
+        "schema_version": "1.0",
+        "models": {
+            "legacy-model": {
+                "mode": "chat",
+                "pricing": {
+                    "input_per_token": 3e-6,
+                    "output_per_token": 1.5e-5,
+                    "cache_read_per_token": 3e-7,
+                    "service_tiers": {
+                        "batch": {
+                            "input_per_token": 1.5e-6,
+                            "output_per_token": 7.5e-6,
+                        }
+                    },
+                },
+                "capabilities": {
+                    "function_calling": False,
+                    "vision": False,
+                    "reasoning": False,
+                    "prompt_caching": True,
+                    "response_schema": False,
+                },
+            }
+        },
+    }
+    (output_dir / "openai.json").write_text(json.dumps(existing_catalog))
+
+    convert(input_data, output_dir)
+
+    catalog = json.loads((output_dir / "openai.json").read_text())
+    legacy_pricing = catalog["models"]["legacy-model"]["pricing"]
+    assert "input_per_token" not in legacy_pricing
+    assert legacy_pricing["input_per_million_tokens"] == pytest.approx(3.0)
+    assert legacy_pricing["output_per_million_tokens"] == pytest.approx(15.0)
+    assert legacy_pricing["cache_read_per_million_tokens"] == pytest.approx(0.3)
+    batch = legacy_pricing["service_tiers"]["batch"]
+    assert "input_per_token" not in batch
+    assert batch["input_per_million_tokens"] == pytest.approx(1.5)
+    assert batch["output_per_million_tokens"] == pytest.approx(7.5)