diff --git a/README.md b/README.md index e19175612..85aa3ddcd 100644 --- a/README.md +++ b/README.md @@ -18,8 +18,14 @@ package is still named `ucode` for compatibility. ```bash uv tool install git+https://github.com/databricks/unity-gateway ug --version +ug --help ``` +`--version` reports installed package metadata without loading the CLI. Use +`ug --help` (or `ucode --help`) to verify that the CLI's startup imports work; +version output alone is not an installation health check. Help does not validate +workspace authentication or agent runtime behavior. + ### Migrating from ucode Reinstall under the new distribution name: @@ -29,6 +35,8 @@ uv tool uninstall ucode uv tool install git+https://github.com/databricks/unity-gateway ug --version ucode --version +ug --help +ucode --help ``` Future upgrades can use `ug upgrade` or `uv tool upgrade unity-gateway`. diff --git a/pyproject.toml b/pyproject.toml index 39fbd9cf9..4ae481525 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,8 +31,8 @@ dependencies = [ ] [project.scripts] -ug = "ucode.cli:main" -ucode = "ucode.cli:main" +ug = "ucode.entrypoint:main" +ucode = "ucode.entrypoint:main" [tool.hatch.build.targets.wheel] # Keep the Python module and existing ~/.ucode configuration paths compatible. diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ee0d4d013..12dd03239 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -16,7 +16,6 @@ from collections.abc import Callable from pathlib import Path -from ucode import gateway_proxy from ucode.config_io import ( APP_DIR, ToolSpec, @@ -2032,6 +2031,9 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: def token_provider(force_refresh: bool) -> str: return get_databricks_token(workspace, profile, force_refresh=force_refresh) + # Only relayed launches need the proxy (and its `httpx` import). + from ucode import gateway_proxy + server, cache, client = gateway_proxy.start_relay_proxy(workspace, token_provider, port) # start_relay_proxy falls back to an OS-assigned port when the cached one is taken # (stale proxy from a killed session). Reconcile settings + state to whatever diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 3c04a93ff..79a73df6d 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -15,7 +15,6 @@ import tomlkit from tomlkit.exceptions import ParseError -from ucode import gateway_proxy from ucode.codex_config import ( catalog_slugs, codex_config_args, @@ -1079,6 +1078,9 @@ def _launch_codex_with_otel_proxy( exec-replacing this process; mirrors Claude's relayed launch. The proxy binds an OS-assigned port and tears everything down when Codex exits (or fails to spawn). """ + # Only tracing launches need the proxy (and its `httpx` import). + from ucode import gateway_proxy + server, cache, client = gateway_proxy.start_otel_proxy( workspace, _otel_token_provider(state, workspace) ) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index cfbe67eea..3f1aecb9f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -3877,7 +3877,10 @@ def _upgrade_failure_detail(result: subprocess.CompletedProcess[str]) -> str: def _verify_upgraded_commands() -> None: - """Ensure both compatibility entry points were installed and can start.""" + """Ensure both compatibility entry points were installed and can start. + + Uses ``--help`` because the entry point answers a bare ``--version`` without importing the CLI. + """ for command in ("ug", "ucode"): executable = shutil.which(command) if executable is None: @@ -3886,7 +3889,7 @@ def _verify_upgraded_commands() -> None: "Reinstall Unity Gateway and ensure the uv tool bin directory is on PATH." ) result = subprocess_cross_os.run( - [executable, "--version"], + [executable, "--help"], check=False, capture_output=True, text=True, @@ -3894,7 +3897,7 @@ def _verify_upgraded_commands() -> None: if result.returncode != 0: detail = _upgrade_failure_detail(result) raise RuntimeError( - f"Upgrade completed, but `{command} --version` failed" + f"Upgrade completed, but `{command} --help` failed" f"{f': {detail}' if detail else '.'}" ) diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index 14b571591..08b1c7cc3 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -12,8 +12,6 @@ from typing import NotRequired, TypedDict from urllib.parse import urlparse -from databricks.sdk import oauth - from ucode.constants import LOCALHOST, LOOPBACK_HOST from ucode.databricks import ( build_auth_token_argv, @@ -187,6 +185,9 @@ def get_custom_client_token( if custom_oauth_cli_enabled(config): profile = profile or _custom_cli_profile(workspace, config["client_id"]) return get_databricks_token(workspace, profile, force_refresh=force_refresh) + # `databricks.sdk` eagerly imports every service module (~0.7s); only the SDK backend needs it. + from databricks.sdk import oauth + try: endpoints = oauth.get_workspace_endpoints(workspace) cache = oauth.TokenCache( diff --git a/src/ucode/entrypoint.py b/src/ucode/entrypoint.py new file mode 100644 index 000000000..c7f774e65 --- /dev/null +++ b/src/ucode/entrypoint.py @@ -0,0 +1,66 @@ +"""Console-script entry point for ``ug`` / ``ucode``. + +Launchers such as isaac run ``ug --version`` on every start. Importing ``ucode.cli`` (Typer, Rich, +questionary, every subcommand module) costs far more than printing the version, so a bare version +request is answered here without importing it. Every other invocation goes to ``ucode.cli.main``. +""" + +from __future__ import annotations + +import os +import sys + +_DISTRIBUTION = "unity_gateway" +_VERSION_ARGV = (["--version"], ["-V"]) + + +def _read_installed_version() -> str | None: + """Return the installed ``unity-gateway`` version from its dist-info ``METADATA``. + + Mirrors the lookup ``importlib.metadata.version`` does for a plain directory ``sys.path`` + (first matching ``*.dist-info`` wins) without importing ``importlib.metadata``, which alone + costs several times a bare interpreter start. Returns None for anything it does not handle + (zip/egg entries, missing or unreadable metadata) so the caller uses the standard lookup. + """ + for entry in sys.path: + try: + names = os.listdir(entry or ".") + except (FileNotFoundError, PermissionError): + continue + except OSError: + return None + for name in names: + lowered = name.lower() + if lowered.endswith((".egg-info", ".egg")) and lowered.startswith(_DISTRIBUTION): + return None + if not lowered.endswith(".dist-info"): + continue + if lowered.rpartition(".")[0].partition("-")[0].replace(".", "_") != _DISTRIBUTION: + continue + try: + with open(os.path.join(entry, name, "METADATA"), encoding="utf-8") as metadata: + for line in metadata: + if not line.strip(): + break + key, _, value = line.partition(":") + if key.strip().lower() == "version": + return value.strip() or None + except OSError: + return None + return None + return None + + +def main() -> None: + if sys.argv[1:] in _VERSION_ARGV: + version = _read_installed_version() + if version is None: + from ucode.telemetry import ug_version + + version = ug_version() + print(version) + return + + from ucode.cli import main as cli_main + + cli_main() diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 50790dfac..d6b11ddf5 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -48,7 +48,7 @@ release_file_lock, ) from ucode.skills import SMART_ROUTER_SKILL, install_skill -from ucode.smart_routing import claude_routing, codex_interposer, routing +from ucode.smart_routing import claude_routing, routing from ucode.smart_routing.claude_hooks import ( FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook, @@ -677,6 +677,9 @@ def launch_codex( app_port = _free_port() app_server_url = _loopback_websocket_url(app_port) + # Codex-only, and pulls in `websockets`; keep it off the Claude launch path. + from ucode.smart_routing import codex_interposer + # Preserve the user's normal CODEX_HOME (including MCP servers, skills, and # preferences) and layer only ucode's gateway settings at CLI precedence. app_server = subprocess_cross_os.popen( diff --git a/tests/README.md b/tests/README.md index dcfc7497f..ce8a3d85a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -28,8 +28,15 @@ preconfigured workspace is assigned. `test_entry_points.py` also runs both installed console scripts (`ug` and `ucode`) and checks their version output against the `unity-gateway` distribution metadata. +`test_startup_imports.py` guards the bare `--version` / `-V` fast path and deferred +command-specific imports in fresh interpreters. Component tests cover metadata lookup, +version suffixes, fallback, and unchanged dispatch for other arguments; they do not +claim a live agent launch or end-to-end startup timing. An isolated interpreter +without site packages also proves that intact metadata can produce version output +while `--help` fails on a missing CLI startup dependency. `TestUpgrade` in `test_cli.py` covers both command names before, during, and after -the distribution rename with mocked installer calls, including failure recovery guidance. +the distribution rename with mocked installer calls, including `--help` verification +of the full CLI import path and failure recovery guidance. `test_subprocess_cross_os.py` covers Windows npm shim resolution, native and Node targets, literal argument preservation, and the shared `subprocess_cross_os.run` / `subprocess_cross_os.popen` entry points, including UTF-8 text decoding, explicit encoding/error overrides, and unchanged diff --git a/tests/integration/README.md b/tests/integration/README.md index c6497d503..665423023 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -18,6 +18,13 @@ The existing unit tests keep their fixtures. Integration has an independent pytest configuration and uses `--confcutdir` so those fixtures cannot leak in. It is not collected by the default `uv run pytest` command. +`../test_startup_imports.py` covers the bare version fast path, deferred imports, +metadata fallback, and other-argument dispatch as component checks. Installation +checks still exercise real console-script help and version output; neither suite +asserts end-to-end agent startup latency. The component suite also checks that +version output alone cannot establish CLI health when startup dependencies are +absent; use `--help` for the startup-import probe. + `TestChildStdoutLaunch` in `../test_cli.py` covers clean Claude print-mode and Codex exec/app-server stdout, early launch errors, and forwarding through ug's `--`. Token membership before the agent's `--` intentionally also matches option values and prompt diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index ebed95c38..e05eb8129 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -13,7 +13,7 @@ import pytest from ucode import databricks as db_mod -from ucode import managed_files +from ucode import gateway_proxy, managed_files from ucode.agents import LaunchOptions, claude from ucode.smart_routing import claude_routing, v2 from ucode.state import MANAGED_OVERLAY_KEY @@ -2527,7 +2527,7 @@ def start_relay_proxy(workspace, token_provider, port): monkeypatch.setattr(claude, "_managed_relayed_conflicts", lambda: None) monkeypatch.setattr(claude, "_ensure_subscription_login", lambda: None) - monkeypatch.setattr(claude.gateway_proxy, "start_relay_proxy", start_relay_proxy) + monkeypatch.setattr(gateway_proxy, "start_relay_proxy", start_relay_proxy) monkeypatch.setattr( claude, "get_databricks_token", diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 935e9712e..485e410f3 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -9,7 +9,7 @@ import pytest -from ucode import managed_files +from ucode import gateway_proxy, managed_files from ucode.agents import LaunchOptions, codex from ucode.config_io import read_toml_safe from ucode.smart_routing import codex_routing @@ -1498,7 +1498,7 @@ def start_otel_proxy(workspace, token_provider): assert token_provider(False) == "tok" return server, cache, client - monkeypatch.setattr(codex.gateway_proxy, "start_otel_proxy", start_otel_proxy) + monkeypatch.setattr(gateway_proxy, "start_otel_proxy", start_otel_proxy) monkeypatch.setattr(codex.subprocess_cross_os, "popen", popen) with pytest.raises(SystemExit) as exc: diff --git a/tests/test_cli.py b/tests/test_cli.py index 06c6f9c9c..a500855de 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -193,8 +193,8 @@ def test_ug_and_ucode_are_equivalent_entry_points(self): "project" ]["scripts"] - assert scripts["ug"] == "ucode.cli:main" - assert scripts["ucode"] == "ucode.cli:main" + assert scripts["ug"] == "ucode.entrypoint:main" + assert scripts["ucode"] == "ucode.entrypoint:main" class TestUpgrade: @@ -265,13 +265,13 @@ def test_cutover_migrates_legacy_distribution_and_verifies_commands(self, prog_n call(["uv", "tool", "uninstall", "ucode"], check=True), call(["uv", "tool", "install", "--force", git_url], check=True), call( - ["/tools/ug", "--version"], + ["/tools/ug", "--help"], check=False, capture_output=True, text=True, ), call( - ["/tools/ucode", "--version"], + ["/tools/ucode", "--help"], check=False, capture_output=True, text=True, diff --git a/tests/test_startup_imports.py b/tests/test_startup_imports.py new file mode 100644 index 000000000..d6a759797 --- /dev/null +++ b/tests/test_startup_imports.py @@ -0,0 +1,233 @@ +"""Component guards for startup imports and compatible ug/ucode version reporting.""" + +from __future__ import annotations + +import json +import subprocess +import sys +import zipfile +from importlib.metadata import version +from pathlib import Path + +import pytest + +from ucode import entrypoint + +_LAZY_MODULES = ( + "databricks.sdk", + "httpx", + "mcp", + "websockets", + "ucode.gateway_proxy", + "ucode.smart_routing.codex_interposer", +) + + +def _modules_loaded_after(code: str, modules: tuple[str, ...]) -> tuple[str, list[str]]: + """Run code in a fresh interpreter and report the guarded imports.""" + probe = f"{code}\nimport json, sys\nprint(json.dumps([m for m in {list(modules)!r} if m in sys.modules]))" + result = subprocess.run( + [sys.executable, "-c", probe], + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + check=False, + timeout=60, + ) + assert result.returncode == 0, result.stderr + *output, loaded = result.stdout.splitlines() + return "\n".join(output), json.loads(loaded) + + +def _write_metadata(root: Path, reported_version: str, distribution: str = "unity_gateway") -> Path: + metadata = root / f"{distribution}-{reported_version}.dist-info" / "METADATA" + metadata.parent.mkdir(parents=True) + metadata.write_text( + f"Metadata-Version: 2.4\nName: unity-gateway\nVersion: {reported_version}\n", + encoding="utf-8", + ) + return metadata + + +@pytest.mark.parametrize("command", ["ug", "ucode"]) +@pytest.mark.parametrize("flag", ["--version", "-V"]) +def test_version_entrypoint_skips_cli_imports(command: str, flag: str) -> None: + output, loaded = _modules_loaded_after( + f"import sys\nsys.argv = [{command!r}, {flag!r}]\nfrom ucode.entrypoint import main\nmain()", + ("ucode.cli", "typer", "ucode.telemetry", "importlib.metadata", *_LAZY_MODULES), + ) + + assert output == version("unity-gateway") + assert loaded == [] + + +@pytest.mark.parametrize("command", ["ug", "ucode"]) +@pytest.mark.parametrize("flag", ["--version", "-V"]) +def test_version_output_does_not_establish_cli_health( + tmp_path: Path, command: str, flag: str +) -> None: + _write_metadata(tmp_path, "1.2.3") + source = Path(entrypoint.__file__).resolve().parents[1] + probe = ( + f"import sys\nsys.path[:0] = {[str(tmp_path), str(source)]!r}\n" + f"sys.argv = [{command!r}, sys.argv[1]]\n" + "from ucode.entrypoint import main\nmain()" + ) + # Disable site packages so metadata survives while startup dependencies are absent. + command_line = [sys.executable, "-I", "-S", "-c", probe] + reported = subprocess.run( + [*command_line, flag], + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + check=False, + timeout=60, + ) + health = subprocess.run( + [*command_line, "--help"], + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + check=False, + timeout=60, + ) + + assert reported.returncode == 0, reported.stderr + assert reported.stdout == "1.2.3\n" + assert reported.stderr == "" + assert health.returncode != 0 + assert "ModuleNotFoundError: No module named 'typer'" in health.stderr + + +def test_installed_version_matches_importlib_metadata() -> None: + assert entrypoint._read_installed_version() == version("unity-gateway") + + +def test_cli_import_defers_subcommand_specific_modules() -> None: + _, loaded = _modules_loaded_after("import ucode.cli\nimport ucode.agents.claude", _LAZY_MODULES) + + assert loaded == [] + + +@pytest.mark.parametrize("reported_version", ["1.2.3", "1.2.3.post4", "1.2.3+abc123"]) +def test_metadata_reader_preserves_version_suffixes( + tmp_path, monkeypatch, reported_version +) -> None: + _write_metadata(tmp_path, reported_version) + monkeypatch.setattr(sys, "path", [str(tmp_path)]) + + assert entrypoint._read_installed_version() == version("unity-gateway") == reported_version + + +@pytest.mark.parametrize("distribution", ["unity_gateway", "UNITY_GATEWAY", "unity.gateway"]) +def test_metadata_reader_normalizes_distribution_name(tmp_path, monkeypatch, distribution) -> None: + _write_metadata(tmp_path, "1.2.3", distribution) + _write_metadata(tmp_path, "9.9.9", "unity_gateway_extra") + monkeypatch.setattr(sys, "path", [str(tmp_path)]) + + assert entrypoint._read_installed_version() == version("unity-gateway") == "1.2.3" + + +def test_metadata_reader_uses_first_sys_path_match(tmp_path, monkeypatch) -> None: + first, second = tmp_path / "first", tmp_path / "second" + _write_metadata(first, "1.2.3") + _write_metadata(second, "4.5.6") + monkeypatch.setattr(sys, "path", [str(first), str(second)]) + + assert entrypoint._read_installed_version() == version("unity-gateway") == "1.2.3" + + +def test_metadata_reader_handles_empty_sys_path_entry(tmp_path, monkeypatch) -> None: + _write_metadata(tmp_path, "1.2.3") + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", [""]) + + assert entrypoint._read_installed_version() == version("unity-gateway") == "1.2.3" + + +def test_metadata_reader_skips_missing_sys_path_directory(tmp_path, monkeypatch) -> None: + _write_metadata(tmp_path, "1.2.3") + monkeypatch.setattr(sys, "path", [str(tmp_path / "missing"), str(tmp_path)]) + + assert entrypoint._read_installed_version() == "1.2.3" + + +@pytest.mark.parametrize( + "contents", + [None, "Name: unity-gateway\nVersion: \n", "Name: unity-gateway\n\nVersion: 1.2.3\n"], +) +def test_metadata_reader_falls_back_for_incomplete_metadata( + tmp_path, monkeypatch, contents +) -> None: + metadata = _write_metadata(tmp_path, "1.2.3") + if contents is None: + metadata.unlink() + else: + metadata.write_text(contents, encoding="utf-8") + monkeypatch.setattr(sys, "path", [str(tmp_path)]) + + assert entrypoint._read_installed_version() is None + + +def test_metadata_reader_falls_back_for_zip_path_before_directory(tmp_path, monkeypatch) -> None: + archive = tmp_path / "installed.zip" + with zipfile.ZipFile(archive, "w") as package: + package.writestr( + "unity_gateway-1.2.3.dist-info/METADATA", "Name: unity-gateway\nVersion: 1.2.3\n" + ) + directory = tmp_path / "later" + _write_metadata(directory, "4.5.6") + monkeypatch.setattr(sys, "path", [str(archive), str(directory)]) + + assert entrypoint._read_installed_version() is None + assert version("unity-gateway") == "1.2.3" + + +def test_metadata_reader_falls_back_for_egg_info(tmp_path, monkeypatch) -> None: + egg = tmp_path / "unity_gateway.egg-info" + egg.mkdir() + (egg / "PKG-INFO").write_text("Name: unity-gateway\nVersion: 1.2.3\n", encoding="utf-8") + monkeypatch.setattr(sys, "path", [str(tmp_path)]) + + assert entrypoint._read_installed_version() is None + assert version("unity-gateway") == "1.2.3" + + +@pytest.mark.parametrize("flag", ["--version", "-V"]) +def test_version_entrypoint_falls_back_to_telemetry(monkeypatch, capsys, flag) -> None: + from ucode import telemetry + + monkeypatch.setattr(sys, "argv", ["ug", flag]) + monkeypatch.setattr(entrypoint, "_read_installed_version", lambda: None) + monkeypatch.setattr(telemetry, "ug_version", lambda: "1.2.3.post4+abc123") + + entrypoint.main() + + assert capsys.readouterr().out == "1.2.3.post4+abc123\n" + + +@pytest.mark.parametrize( + "argv", + [ + [], + ["--help"], + ["--dry-run", "--version"], + ["--version", "foo"], + ["--version", "--help"], + ["-V", "claude"], + ["claude", "--version"], + ], +) +def test_other_arguments_delegate_to_cli_unchanged(monkeypatch, argv) -> None: + from ucode import cli + + calls = [] + monkeypatch.setattr(sys, "argv", ["ug", *argv]) + monkeypatch.setattr(cli, "main", lambda: calls.append(sys.argv.copy())) + monkeypatch.setattr( + entrypoint, "_read_installed_version", lambda: pytest.fail("unexpected version fast path") + ) + + entrypoint.main() + + assert calls == [["ug", *argv]]