Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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`.
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 3 additions & 1 deletion src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion src/ucode/agents/codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
)
Expand Down
9 changes: 6 additions & 3 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -3886,15 +3889,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 '.'}"
)

Expand Down
5 changes: 3 additions & 2 deletions src/ucode/custom_oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down
66 changes: 66 additions & 0 deletions src/ucode/entrypoint.py
Original file line number Diff line number Diff line change
@@ -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()
5 changes: 4 additions & 1 deletion src/ucode/smart_routing/v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down
9 changes: 8 additions & 1 deletion tests/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions tests/integration/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions tests/test_agent_claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down
4 changes: 2 additions & 2 deletions tests/test_agent_codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading