Cut ug startup cost: fast --version entry point and lazy heavy imports

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`.

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.

src/ucode/agents/claude.py

@@ -16,7 +16,6 @@
 from collections.abc import Callable, Iterable, Sequence
 from pathlib import Path
 
-from ucode import gateway_proxy
 from ucode.config_io import (
     APP_DIR,
     ToolSpec,
@@ -2193,6 +2192,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

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,
@@ -1074,6 +1073,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)
     )

src/ucode/cli.py

@@ -4181,7 +4181,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:
@@ -4190,15 +4193,15 @@ 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,
         )
         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 '.'}"
             )
 

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,
@@ -191,6 +189,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(

src/ucode/entrypoint.py

@@ -0,0 +1,72 @@
+"""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 NotADirectoryError:
+            # Windows puts the console-script launcher (e.g. Scripts\ug.exe) on sys.path; it holds
+            # no distribution metadata. Other files may be zip archives, so they fall back.
+            if entry.lower().endswith(".exe"):
+                continue
+            return None
+        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()

src/ucode/smart_routing/v2.py

@@ -51,7 +51,7 @@
     release_file_lock,
 )
 from ucode.skills import SMART_ROUTER_ORCHESTRATOR_SKILL, SMART_ROUTER_SKILL, install_skill
-from ucode.smart_routing import claude_routing, codex_interposer, orchestrator, routing
+from ucode.smart_routing import claude_routing, orchestrator, routing
 from ucode.smart_routing.claude_hooks import (
     FIRST_PROMPT_SOCKET_ENV,
     sync_first_prompt_hook,
@@ -729,6 +729,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(

tests/README.md

@@ -68,8 +68,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

tests/integration/README.md

@@ -72,6 +72,13 @@ timeout. Offline PTY checks cover this ordering.
 The shared wait also runs before CUJ4's Claude preset sessions exit, after delegated-task
 evidence is checked, because Claude can resume a child after its first completed answer.
 
+`../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.
+
 Agent-specific bundled instruction selection and prompt/compaction hook content
 are covered in `../test_skills.py`; those component checks do not run live agents.
 

tests/test_agent_claude.py

@@ -14,7 +14,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
@@ -2888,7 +2888,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",

tests/test_agent_codex.py

@@ -10,7 +10,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
@@ -1504,7 +1504,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:

tests/test_cli.py

@@ -195,8 +195,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:
@@ -267,13 +267,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,