From 1c935085400b4c6c1381ab13c1cf2dd7125eaa1d Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 13:37:00 -0400 Subject: [PATCH 1/6] Introduce ug command and migration path --- README.md | 246 ++++++++++++++------------- pyproject.toml | 1 + src/ucode/cli.py | 166 +++++++++--------- src/ucode/managed_export.py | 4 +- src/ucode/managed_wizard.py | 196 ++++++++++----------- src/ucode/smart_routing/routing.py | 4 +- tests/test_cli.py | 101 +++++++---- tests/test_codex_routing.py | 2 +- tests/test_codex_smart_routing_v2.py | 2 +- tests/test_managed_wizard.py | 66 ++++--- 10 files changed, 421 insertions(+), 367 deletions(-) diff --git a/README.md b/README.md index 661ab410..f231ab4c 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,10 @@ -# Unity AI Gateway Coding CLI (ucode) +# Unity Gateway (`ug`) -`ucode` is a lightweight launcher for running Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Pi through Databricks. +Existing `ucode` commands continue to work unchanged. Going forward, the CLI is named Unity +Gateway and its primary command is `ug`; `ucode` remains a supported alias. + +Unity Gateway is a lightweight launcher for running Codex, Claude Code, Gemini CLI, OpenCode, +GitHub Copilot CLI, and Pi through Databricks. ## Requirements @@ -13,7 +17,7 @@ uv tool install git+https://github.com/databricks/ucode ``` -Check your version with `ucode --version`. Between releases this looks like +Check your version with `ug --version`. Between releases this looks like `0.1.0+14.g93986a8` — the trailing `g` is the exact commit the build came from, so include it when reporting a bug. @@ -24,22 +28,22 @@ from, so include it when reporting a bug. Just run the tool you want: ```bash -ucode codex # OpenAI Codex -ucode claude # Claude Code -ucode gemini # Gemini CLI -ucode opencode # OpenCode -ucode copilot # GitHub Copilot CLI -ucode pi # Pi -ucode cursor # Cursor Agent (MCP only — see below) +ug codex # OpenAI Codex +ug claude # Claude Code +ug gemini # Gemini CLI +ug opencode # OpenCode +ug copilot # GitHub Copilot CLI +ug pi # Pi +ug cursor # Cursor Agent (MCP only — see below) ``` -On first launch, `ucode` will prompt for your Databricks workspace URL, authenticate, and configure that tool automatically. Subsequent Claude and Codex launches use the generated local settings directly. Use `ucode claude --refresh` or `ucode codex --refresh` when you want to re-check Databricks and update the model/configuration. +On first launch, `ug` will prompt for your Databricks workspace URL, authenticate, and configure that tool automatically. Subsequent Claude and Codex launches use the generated local settings directly. Use `ug claude --refresh` or `ug codex --refresh` when you want to re-check Databricks and update the model/configuration. Pass flags directly to the underlying tool: ```bash -ucode claude -r # resume last session -ucode codex --full-auto +ug claude -r # resume last session +ug codex --full-auto ``` All agents route through Databricks AI Gateway using your workspace credentials — no API keys required. @@ -49,8 +53,8 @@ router to select models for that session and its subagents. Codex may require on the launch-scoped hooks through `/hooks`. ```bash -ucode codex --enable-smart-routing -ucode claude --enable-smart-routing +ug codex --enable-smart-routing +ug claude --enable-smart-routing ``` The flag applies only to that launch; later launches use normal model selection unless the flag is @@ -59,13 +63,13 @@ passed again. To configure all tools at once: ```bash -ucode configure +ug configure ``` To configure specific tools without the picker, pass a comma-separated list: ```bash -ucode configure --agents claude,codex +ug configure --agents claude,codex ``` Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, and `pi`. `cursor` is also accepted (MCP-only — it registers Databricks MCP servers but configures no models). @@ -73,7 +77,7 @@ Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, an Naming agents explicitly is treated as a request for all of them: if any one isn't available on the workspace, the run fails without configuring the others. Add `--skip-unavailable` to configure the available subset instead and skip the rest with a warning: ```bash -ucode configure --agents claude,codex,pi --skip-unavailable +ug configure --agents claude,codex,pi --skip-unavailable ``` This is useful in CI against a mix of workspaces — on a workspace whose AI Gateway exposes no OpenAI models, the command above still configures `claude` and `pi`, and reports Codex as skipped. It exits non-zero only when none of the requested agents are available. @@ -81,29 +85,29 @@ This is useful in CI against a mix of workspaces — on a workspace whose AI Gat To configure without the workspace picker, pass a comma-separated list of workspaces: ```bash -ucode configure --workspaces https://first.databricks.com,https://second.databricks.com +ug configure --workspaces https://first.databricks.com,https://second.databricks.com ``` -When multiple workspaces are provided, `ucode` logs into and saves state for each workspace. Launch commands such as `ucode codex` use the first workspace in the list. +When multiple workspaces are provided, `ug` logs into and saves state for each workspace. Launch commands such as `ug codex` use the first workspace in the list. Alternatively, pass existing Databricks CLI profiles (from `~/.databrickscfg`) instead of workspace URLs — each profile's host supplies the workspace URL: ```bash -ucode configure --profiles DEFAULT --agents claude,codex +ug configure --profiles DEFAULT --agents claude,codex ``` Auth behaves the same as `--workspaces`: an OAuth `databricks auth login` is forced by default. -For CI or headless environments where the profile holds a personal access token (`auth_type = pat` in `~/.databrickscfg`), add `--use-pat`. It must be combined with `--profiles` — ucode never picks up a PAT implicitly — and runs no interactive login: the profile's token is used for the whole setup (and by launched agents afterwards), with workspace access verified against the AI Gateway. `--skip-validate` additionally skips the post-configure test message sent through each agent, so configure only writes config files with the freshly discovered models. Together these make setup fully non-interactive: +For CI or headless environments where the profile holds a personal access token (`auth_type = pat` in `~/.databrickscfg`), add `--use-pat`. It must be combined with `--profiles` — ug never picks up a PAT implicitly — and runs no interactive login: the profile's token is used for the whole setup (and by launched agents afterwards), with workspace access verified against the AI Gateway. `--skip-validate` additionally skips the post-configure test message sent through each agent, so configure only writes config files with the freshly discovered models. Together these make setup fully non-interactive: ```bash -ucode configure --profiles DEFAULT --agents claude,codex --use-pat --skip-validate --skip-upgrade +ug configure --profiles DEFAULT --agents claude,codex --use-pat --skip-validate --skip-upgrade ``` ### MCP servers (optional) ```bash -ucode configure mcp +ug configure mcp ``` Add Databricks MCP servers to installed MCP-capable tools: Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Cursor Agent. @@ -116,22 +120,22 @@ Options are shown in this order: Discovered external MCP connections are listed directly. -Every Databricks MCP server is registered as a local **stdio** server that runs `ucode mcp-proxy` -— a small bridge (shipped with `ucode`) between the coding tool and the Databricks +Every Databricks MCP server is registered as a local **stdio** server that runs `ug mcp-proxy` +— a small bridge (shipped with `ug`) between the coding tool and the Databricks streamable-HTTP MCP endpoint. The proxy mints a fresh OAuth token from your Databricks CLI profile on every request, so MCP auth is handled uniformly for every client and never expires mid-session. The coding tool starts and stops the proxy as a child process; there's nothing extra to run. -**Cursor** is MCP-only: `cursor-agent` runs models on your own Cursor account, so `ucode` +**Cursor** is MCP-only: `cursor-agent` runs models on your own Cursor account, so `ug` configures no models for it — it only registers Databricks MCP servers in `~/.cursor/mcp.json` -(via the same proxy). Include it with `ucode configure --agents cursor` or pick it in -`ucode configure mcp`, then launch with `ucode cursor`. +(via the same proxy). Include it with `ug configure --agents cursor` or pick it in +`ug configure mcp`, then launch with `ug cursor`. To set up an agent and its MCP server(s) in one command, pass `--mcp` with fully-qualified -service name(s) to `ucode configure`: +service name(s) to `ug configure`: ```bash -ucode configure --agents claude --mcp system.ai.slack +ug configure --agents claude --mcp system.ai.slack ``` `--mcp` also works without `--agents` for MCP-only clients (it configures just the workspace, @@ -139,22 +143,22 @@ then registers the servers); pass a comma-separated list to register several at #### Add servers without replacing existing ones -`ucode configure mcp` **replaces** the registered MCP servers with your selection — anything +`ug configure mcp` **replaces** the registered MCP servers with your selection — anything outside a `--location`/`--services` scope (or left unchecked in the picker) is removed. To -**add** servers while leaving everything already configured in place, use `ucode mcp add`: +**add** servers while leaving everything already configured in place, use `ug mcp add`: ```bash # Register a whole schema's services, keeping any servers already configured. -ucode mcp add --location system.ai +ug mcp add --location system.ai # Register just a subset (same name rules as `configure mcp --services`). -ucode mcp add --services system.ai.slack,system.ai.github +ug mcp add --services system.ai.slack,system.ai.github # No arguments launches the same interactive picker, but never removes servers. -ucode mcp add +ug mcp add ``` -`ucode mcp add` takes the same `--location` and `--services` options as `ucode configure mcp`; +`ug mcp add` takes the same `--location` and `--services` options as `ug configure mcp`; the only difference is that it never removes servers outside the selection. In the interactive picker, servers you already have configured are shown as `(already configured)` and can't be toggled off — you only pick new ones to add. @@ -164,24 +168,24 @@ configured first (workspace + models), so this doubles as one-command setup: ```bash # Set up Claude Code (if needed) and register the server for it, in one command. -ucode mcp add --agents claude --services system.ai.slack +ug mcp add --agents claude --services system.ai.slack # Target several agents at once. -ucode mcp add --agents claude,codex --location system.ai +ug mcp add --agents claude,codex --location system.ai ``` Without `--agents`, the server is registered for every already-configured agent. #### Remove configured servers -To unregister servers you've already configured, use `ucode mcp remove`: +To unregister servers you've already configured, use `ug mcp remove`: ```bash -ucode mcp remove +ug mcp remove # Remove only from specific agents. A server registered on several agents is # unregistered from the named ones and kept on the rest. -ucode mcp remove --agents codex +ug mcp remove --agents codex ``` It shows the servers you currently have configured — each with the coding tools it's registered @@ -189,20 +193,20 @@ on — and removes the ones you select from those tools. It needs no Databricks ### Skills (optional) -Configure Unity Catalog Skills for your coding tools with `ucode configure skills`: +Configure Unity Catalog Skills for your coding tools with `ug configure skills`: ```bash # Utility tools only: register the schema-less skills MCP connection, no download. -ucode configure skills +ug configure skills # Download mode: fetch every skill in the schema to disk (and register the connection). -ucode configure skills --location main.default --path /abs/project/dir +ug configure skills --location main.default --path /abs/project/dir # Download a named subset of the schema's skills instead of all of them. -ucode configure skills --location main.default --skill my-skill +ug configure skills --location main.default --skill my-skill # MCP mode: expose the schema's skills as MCP tools instead of downloading. -ucode configure skills --location main.default,ml.prod --mcp +ug configure skills --location main.default,ml.prod --mcp ``` - **Bare command** (no `--location`) registers the schema-less skills MCP connection — the @@ -220,27 +224,27 @@ ucode configure skills --location main.default,ml.prod --mcp (override-only) and rebuilds its `?schema=` URL; no files are downloaded and `--path` is rejected. Each run prints the registered server, its URL, the configured agents, and its tools, and reminds -you to run `ucode ` (existing agent sessions need a restart before the MCP tools load). +you to run `ug ` (existing agent sessions need a restart before the MCP tools load). ### Managed config for a workspace (admins) Author the coding config your developers pick up automatically, instead of asking each of them to -run `ucode configure` by hand. Restricted to workspace admins. `ucode setup help` prints the whole +run `ug configure` by hand. Restricted to workspace admins. `ug setup help` prints the whole sequence; the short version is one command for the agents and models, then a command per optional section, then publish: ```bash -ucode setup # agents and models (start here) -ucode setup mcps # managed MCP servers -ucode setup skills # managed skills -ucode setup spend-tiers # spend-based routing -ucode publish # publish it to the workspace +ug setup # agents and models (start here) +ug setup mcps # managed MCP servers +ug setup skills # managed skills +ug setup spend-tiers # spend-based routing +ug publish # publish it to the workspace ``` -`ucode setup` walks through the agents to enable and which one bare `ucode` launches, then per agent: +`ug setup` walks through the agents to enable and which one bare `ug` launches, then per agent: Databricks-hosted models or an external Model Provider Service and the models to expose. Interactive Claude Code and Codex configuration installs gateway-critical values in the OS-managed settings -scope so enterprise settings cannot silently override ucode. Non-interactive and CI runs use local +scope so enterprise settings cannot silently override Unity Gateway. Non-interactive and CI runs use local files without invoking `sudo`, and stop with an actionable error if an existing managed value conflicts. Claude subscription relay is local-only because its loopback proxy exists only for that session. @@ -248,47 +252,47 @@ Claude Code is asked one model per family (opus/sonnet/haiku/fable), since it se alias; any family can be skipped. The optional sections each edit their own part of the same config, so you can add an MCP server or -change a spend tier later without walking the whole flow. `ucode setup skills --location -main.default,other.schema` skips the prompt. `ucode setup spend-tiers` sets a tiered spend policy +change a spend tier later without walking the whole flow. `ug setup skills --location +main.default,other.schema` skips the prompt. `ug setup spend-tiers` sets a tiered spend policy that switches the default agent and model as the workspace burns through a budget. Each section command also offers to publish right away, so you can apply changes incrementally; answering the -section prompts also runs the matching `ucode configure` step, which does configure this machine. +section prompts also runs the matching `ug configure` step, which does configure this machine. Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which -`ucode publish` publishes. Re-running `ucode setup` keeps the MCP servers, skills, tracing table, and +`ug publish` publishes. Re-running `ug setup` keeps the MCP servers, skills, tracing table, and tiered spend policy already authored, rather than clearing them; to drop one, edit the file and reload -it with `ucode setup --from-file`. +it with `ug setup --from-file`. ```bash -# Review the manifest and the exact payload `ucode publish` would publish. -ucode setup show +# Review the manifest and the exact payload `ug publish` would publish. +ug setup show # Skip the prompts and load a hand-written config instead (validated before saving). -ucode setup --from-file ./managed-config.json +ug setup --from-file ./managed-config.json ``` Once the manifest looks right, publish it: ```bash # Validate, show a diff against what's live, and ask before publishing. -ucode publish +ug publish # Publish without the confirmation prompt (for CI). -ucode publish --yes +ug publish --yes -# Publish a config file exported with `ucode export` instead of the locally authored one. -ucode publish -f ./managed-config.json -ucode publish --file ./managed-config.json --yes +# Publish a config file exported with `ug export` instead of the locally authored one. +ug publish -f ./managed-config.json +ug publish --file ./managed-config.json --yes ``` `publish` updates the workspace's existing config in place rather than replacing it, so a failed publish leaves the current config intact. It shows a diff of exactly what changes against the published config before asking to confirm, and does nothing when the two already match. It is a -whole-manifest write — every field ucode authors is sent — but because `ucode setup` carries the +whole-manifest write — every field ug authors is sent — but because `ug setup` carries the other sections forward, a re-run no longer silently drops them. Developers pick the new config up on -their next ucode run. +their next ug run. -With `-f`/`--file`, `publish` reads a config file produced by `ucode export` and publishes it through +With `-f`/`--file`, `publish` reads a config file produced by `ug export` and publishes it through the same validation, diff, and confirmation flow. The file's `workspace` must match the configured workspace (it can never redirect publication elsewhere) and its `spec_version` must be a supported integer; server-owned fields (resource name, workspace ids, timestamps, user ids) and unknown fields @@ -296,20 +300,20 @@ are rejected rather than silently dropped. ### Exporting the config -Any user (not only admins) can print the workspace's managed config as portable JSON with `ucode +Any user (not only admins) can print the workspace's managed config as portable JSON with `ug export`. The output leads with the source `workspace` URL and a `spec_version` (the export format version), followed by the canonical external config; credentials and server-assigned fields (the resource name, timestamps, user ids) are excluded. Without `--file` the JSON is written to stdout; with `--file`/`-f` the same bytes are written to a file (atomically, and the destination's parent -directory must already exist) while stdout stays empty. The exported file is exactly what `ucode +directory must already exist) while stdout stays empty. The exported file is exactly what `ug publish -f ` consumes. ```bash # Print the managed config as JSON. -ucode export +ug export # Write it to a file; stdout stays empty. -ucode export --file ./managed-config.json +ug export --file ./managed-config.json ``` The output looks like: @@ -329,56 +333,56 @@ The output looks like: | Command | Description | |---------|-------------| -| `ucode status` | Show current workspace, base URLs, managed config files, and selected models | -| `ucode export` | Print the workspace's managed config as portable JSON (`--file ` / `-f` to write a file) | -| `ucode doctor` | Diagnose local issues (uv, npm, Databricks CLI, workspace, credentials, agent CLIs, tracing) and offer to fix any problems found | -| `ucode usage` | Show AI Gateway usage summary, plus your budget spend against its alert threshold when the workspace reports one | -| `ucode usage --warehouse-id ` | Query a specific SQL warehouse instead of discovering one | -| `ucode revert` | Clear saved state and restore backed-up config files | -| `ucode configure --dry-run` | Preview config files without writing them | -| `ucode configure --agents claude,codex` | Configure specific agents without the interactive picker | -| `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker | -| `ucode configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) | -| `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | -| `ucode codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents | -| `ucode codex --refresh` | Re-check Databricks, refresh models/configuration, and launch Codex | -| `ucode claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents | -| `ucode claude --refresh` | Re-check Databricks, refresh models/configuration, and launch Claude Code | -| `ucode configure --skip-validate` | Write configs without sending a test message through each agent | -| `ucode configure --agents claude,codex,pi --skip-unavailable` | Configure the requested agents that are available; skip the rest with a warning | -| `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | -| `ucode mcp add --location system.ai` | Register a schema's MCP servers, keeping any already configured (additive; never removes) | -| `ucode mcp add --services system.ai.slack` | Register specific MCP server(s) without removing existing ones | -| `ucode mcp add --agents claude --services system.ai.slack` | Set up the agent(s) if needed and register the server for them | -| `ucode mcp remove` | Interactively unregister configured MCP servers from your coding tools | -| `ucode mcp remove --agents codex` | Unregister selected servers from specific agents only | -| `ucode configure skills` | Register the skills MCP connection (utility tools only); no skills download | -| `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | -| `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | -| `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | -| `ucode setup` | Author the managed config's agents and models (workspace admins only) | -| `ucode setup mcps` | Add or change the managed config's MCP servers | -| `ucode setup skills [--location a.b,c.d]` | Add or change the managed config's skills | -| `ucode setup spend-tiers` | Set the managed config's tiered spend routing policy | -| `ucode setup help` | Walk through the whole setup sequence, marking what's already configured | -| `ucode setup show` | Print the authored config and the payload `ucode publish` would publish | -| `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | -| `ucode publish` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | -| `ucode publish -f ` | Publish a config file exported with `ucode export` instead of the locally authored one | -| `ucode publish --yes` | Publish without the confirmation prompt | - -Databricks AI Tools are installed only by `ucode configure`, never by `ucode ` launches. -Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ucode configure` to +| `ug status` | Show current workspace, base URLs, managed config files, and selected models | +| `ug export` | Print the workspace's managed config as portable JSON (`--file ` / `-f` to write a file) | +| `ug doctor` | Diagnose local issues (uv, npm, Databricks CLI, workspace, credentials, agent CLIs, tracing) and offer to fix any problems found | +| `ug usage` | Show AI Gateway usage summary, plus your budget spend against its alert threshold when the workspace reports one | +| `ug usage --warehouse-id ` | Query a specific SQL warehouse instead of discovering one | +| `ug revert` | Clear saved state and restore backed-up config files | +| `ug configure --dry-run` | Preview config files without writing them | +| `ug configure --agents claude,codex` | Configure specific agents without the interactive picker | +| `ug configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker | +| `ug configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) | +| `ug configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | +| `ug codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents | +| `ug codex --refresh` | Re-check Databricks, refresh models/configuration, and launch Codex | +| `ug claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents | +| `ug claude --refresh` | Re-check Databricks, refresh models/configuration, and launch Claude Code | +| `ug configure --skip-validate` | Write configs without sending a test message through each agent | +| `ug configure --agents claude,codex,pi --skip-unavailable` | Configure the requested agents that are available; skip the rest with a warning | +| `ug configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | +| `ug mcp add --location system.ai` | Register a schema's MCP servers, keeping any already configured (additive; never removes) | +| `ug mcp add --services system.ai.slack` | Register specific MCP server(s) without removing existing ones | +| `ug mcp add --agents claude --services system.ai.slack` | Set up the agent(s) if needed and register the server for them | +| `ug mcp remove` | Interactively unregister configured MCP servers from your coding tools | +| `ug mcp remove --agents codex` | Unregister selected servers from specific agents only | +| `ug configure skills` | Register the skills MCP connection (utility tools only); no skills download | +| `ug configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | +| `ug configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | +| `ug configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | +| `ug setup` | Author the managed config's agents and models (workspace admins only) | +| `ug setup mcps` | Add or change the managed config's MCP servers | +| `ug setup skills [--location a.b,c.d]` | Add or change the managed config's skills | +| `ug setup spend-tiers` | Set the managed config's tiered spend routing policy | +| `ug setup help` | Walk through the whole setup sequence, marking what's already configured | +| `ug setup show` | Print the authored config and the payload `ug publish` would publish | +| `ug setup --from-file ` | Load a hand-written managed config instead of running the prompts | +| `ug publish` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | +| `ug publish -f ` | Publish a config file exported with `ug export` instead of the locally authored one | +| `ug publish --yes` | Publish without the confirmation prompt | + +Databricks AI Tools are installed only by `ug configure`, never by `ug ` launches. +Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to control the installation. ## Managed Local Files -`ucode` manages these files: +`ug` manages these files: | File | Tool | |------|------| | `~/.codex/ucode.config.toml` (or legacy `~/.codex/config.toml`) | Codex | -| `~/.claude/ucode-settings.json` | Claude Code settings generated by ucode | +| `~/.claude/ucode-settings.json` | Claude Code settings generated by ug | | `/etc/claude-code/managed-settings.json` (Linux) or `/Library/Application Support/ClaudeCode/managed-settings.json` (macOS) | Claude Code OS-managed settings | | `/etc/codex/managed_config.toml` | Codex OS-managed settings | | `~/.gemini/.env` | Gemini CLI | @@ -386,10 +390,10 @@ control the installation. | `~/.copilot/.env` | GitHub Copilot CLI | | `~/.pi/agent/models.json` | Pi | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | -| `~/.ucode/managed-state.json` | The managed config — authored by `ucode setup` (admins) and refreshed from the workspace on launch | -| `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ucode | +| `~/.ucode/managed-state.json` | The managed config — authored by `ug setup` (admins) and refreshed from the workspace on launch | +| `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ug | -Existing files are backed up before being overwritten. `ucode revert` restores backups. +Existing files are backed up before being overwritten. `ug revert` restores backups. ## Documentation diff --git a/pyproject.toml b/pyproject.toml index 52564f5c..f69b8990 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ tracing = ["mlflow[databricks]>=3.4"] [project.scripts] +ug = "ucode.cli:main" ucode = "ucode.cli:main" [tool.hatch.build.targets.wheel] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..21ef7144 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""CLI entry point for ucode.""" +"""CLI entry point for ``ug``.""" from __future__ import annotations @@ -184,11 +184,11 @@ def _print_managed_summary( """Show which of the admin's settings are in force. With ``tool`` set (launch path) the per-agent Agent/Provider/Model lines are included; - with ``tool=None`` (e.g. ``ucode configure`` under a managed config) they are skipped + with ``tool=None`` (e.g. ``ug configure`` under a managed config) they are skipped since no single agent has been chosen yet. ``abridged`` prints only what changes launch-to-launch — the agent and model this run will use, - and the policy in force — with a pointer to ``ucode status`` for the rest. Bare ``ucode`` runs + and the policy in force — with a pointer to ``ug status`` for the rest. Bare ``ucode`` runs every session, so re-enumerating the workspace's full MCP/skills/tier list each time is noise; the full box stays for ``status`` and ``configure``, where the reader asked to see it. """ @@ -237,7 +237,7 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None """One-line launch banner: which agent (and model) this managed run is launching. Bare ``ucode`` runs every session, so the full box's MCP/skills/policy enumeration is noise - each time; ``ucode status`` still shows all of it. See ``_print_managed_summary``'s ``abridged`` + each time; ``ug status`` still shows all of it. See ``_print_managed_summary``'s ``abridged`` note. ``tool`` is always set on the launch path, but is guarded for callers that pass None.""" if tool is None: print_note("Using managed config.") @@ -257,25 +257,25 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None def _confirm_managed_config_applied(managed: dict, workspace: str) -> None: print_success("A managed config is published for your workspace — you're all set.") _print_managed_summary(managed, {"workspace": workspace}, tool=None) - print_note("Run `ucode` to launch with your managed settings.") + print_note("Run `ug` to launch with your managed settings.") def _resolve_workspace_then_maybe_reject( workspace_entries: list[tuple[str, str | None]] | None, ) -> list[tuple[str, str | None]] | None: - """Resolve the workspace ``ucode configure`` targets, then branch on role + managed config. + """Resolve the workspace ``ug configure`` targets, then branch on role + managed config. Enablement is both client- and server-side: the client-side ``ENABLE_MANAGED_AGENT_CONFIG`` env var must be set for ``ucode`` to run any of this (the opt-in bug-bash gate below), and the workspace's gateway must not report the feature disabled (``FEATURE_DISABLED``) — a config only exists to adopt when the server side is on too. - When managed coding-agent configs are enabled, ``ucode configure`` must still let a developer + When managed coding-agent configs are enabled, ``ug configure`` must still let a developer switch workspaces — so resolve the target workspace up front (prompting when the interactive path gave no ``--workspaces``/``--profiles``) and make it current *before* deciding what to do. Then, gated by the client-side ``ENABLE_MANAGED_AGENT_CONFIG``, the four role/config paths are: - * **No managed config** → a workspace admin is dropped straight into the ``ucode setup`` + * **No managed config** → a workspace admin is dropped straight into the ``ug setup`` authoring flow (``configure`` is replacing ``setup``) and the command exits with its code; a non-admin's own ``configure`` proceeds, with the resolved entries returned so the caller reuses them instead of re-prompting. @@ -320,7 +320,7 @@ def _resolve_workspace_then_maybe_reject( def _maybe_run_admin_setup(workspace: str, profile: str | None) -> None: """When a workspace admin runs ``configure`` on a workspace with no managed config, drop straight - into the ``ucode setup`` authoring flow — ``configure`` is replacing ``setup``, so the admin + into the ``ug setup`` authoring flow — ``configure`` is replacing ``setup``, so the admin never has to invoke it themselves. On completion, exit with setup's own status code. A plain developer (and any caller whose admin status can't be verified) instead falls through to @@ -339,13 +339,13 @@ def _maybe_run_admin_setup(workspace: str, profile: str | None) -> None: print_note( "You're a workspace admin, and no managed coding agent config exists for this workspace " "yet — let's set one up. Choose the agents, models, MCPs, and skills once and every " - "developer inherits them when they run `ucode`." + "developer inherits them when they run `ug`." ) _run_setup_and_exit(workspace, profile, token) def _run_setup_and_exit(workspace: str, profile: str | None, token: str | None = None) -> None: - """Launch the ``ucode setup`` authoring flow in place, then exit with its status code. + """Launch the ``ug setup`` authoring flow in place, then exit with its status code. Reuses the workspace/profile ``configure`` already resolved and authenticated against so setup doesn't prompt for them again, and hands setup the same ``token`` the admin check already used @@ -355,13 +355,13 @@ def _run_setup_and_exit(workspace: str, profile: str | None, token: str | None = mapped to clean exit codes rather than bubbling up as unhandled errors. """ try: - # Brand the flow as "Configure Unity Gateway": it was reached through `ucode configure`, - # not a bare `ucode setup`, so its section headers use the product name rather than the + # Brand the flow as "Configure ug": it was reached through `ug configure`, + # not a bare `ug setup`, so its section headers use the product name rather than the # bare command. code = setup_command( workspace=workspace, profile=profile, - command_label="Configure Unity Gateway", + command_label="Configure ug", token=token, ) except RuntimeError as exc: @@ -530,7 +530,7 @@ def configure_shared_state( don't error out. If ``None``, we resolve it from the host after login. If skip_preflight is True, skip the entire preflight block below — auth validation, the AI Gateway probe, and model discovery — trusting a prior - ``ucode configure``. The PAT/bearer is already exported (``apply_pat_environment`` + ``ug configure``. The PAT/bearer is already exported (``apply_pat_environment`` in ``_launch_tool``) and the gateway was verified by that earlier configure. Only the local profile resolution and the shared state assembly still run; the saved model lists are preserved. @@ -561,7 +561,7 @@ def configure_shared_state( # each path below resolves it once, where a host->profile lookup is reliable # (the skip branch trusts the prior configure; the preflight resolves after # login). --skip-preflight persists exactly this and returns, trusting a prior - # `ucode configure` — it already validated auth + the AI Gateway and saved the + # `ug configure` — it already validated auth + the AI Gateway and saved the # model lists (carried over by load_state, left untouched). state = load_state() state["workspace"] = workspace @@ -587,7 +587,7 @@ def configure_shared_state( state["base_urls"] = build_shared_base_urls(workspace) if skip_preflight: - # A prior `ucode configure` created the profile; resolve it locally (no + # A prior `ug configure` created the profile; resolve it locally (no # login needed) and persist it so launches disambiguate. if profile is None: profile = find_profile_name_for_host(workspace) @@ -847,7 +847,7 @@ def configure_workspace_command( raise RuntimeError("Use either --agent or --agents, not both.") # The Databricks-vs-Model-Provider-Service picker is shown only on the fully - # interactive path (`ucode configure` with no --agent/--agents). Naming agents + # interactive path (`ug configure` with no --agent/--agents). Naming agents # explicitly signals the non-interactive flow, which stays on Databricks. offer_provider = tool is None and selected_tools is None @@ -993,7 +993,7 @@ def status() -> int: mcp_servers = state.get("mcp_servers") or [] configured_tools = set(state.get("available_tools") or managed_configs.keys()) - console.print(heading("ucode status")) + console.print(heading("ug status")) console.print( f" {status_badge('Configured', 'ok') if workspace else status_badge('Not Configured', 'warn')}" ) @@ -1040,7 +1040,7 @@ def status() -> int: print_kv("MCP list command", str(MCP_CLIENTS[tool]["list_command"])) print_kv( "MCP servers", - ", ".join(tool_mcp_servers) if tool_mcp_servers else "none saved by ucode", + ", ".join(tool_mcp_servers) if tool_mcp_servers else "none saved by ug", ) print_kv("Config file", str(config_path) if config_path.exists() else "missing") if tool == "claude": @@ -1094,15 +1094,13 @@ def status() -> int: print_heading("State") print_kv("State file", str(STATE_PATH) if STATE_PATH.exists() else "missing") - print_note("Use `ucode configure` to update workspace settings or configure new tools.") + print_note("Use `ug configure` to update workspace settings or configure new tools.") + print_note("Use `ug configure mcp` to add Databricks MCP servers to configured coding tools.") print_note( - "Use `ucode configure mcp` to add Databricks MCP servers to configured coding tools." + "Use `ug configure skills` to set up Unity Catalog Skills for configured coding tools." ) - print_note( - "Use `ucode configure skills` to set up Unity Catalog Skills for configured coding tools." - ) - print_note("Use `ucode configure tracing` to log coding sessions to an MLflow experiment.") - print_note("Use `ucode revert` to clear managed configs and restore prior files.") + print_note("Use `ug configure tracing` to log coding sessions to an MLflow experiment.") + print_note("Use `ug revert` to clear managed configs and restore prior files.") return 0 @@ -1141,7 +1139,7 @@ def revert() -> int: f"{spec['display']} MCP config", "restored" if mcp_results.get(client) else "unchanged", ) - print_success("ucode state cleared") + print_success("ug state cleared") return 0 @@ -1158,12 +1156,12 @@ def revert() -> int: configure_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) -app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ug.") setup_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer( setup_app, name="setup", - help="Author the workspace's managed coding config (admins only). See `ucode setup help`.", + help="Author the workspace's managed coding config (admins only). See `ug setup help`.", ) @@ -1179,8 +1177,8 @@ def _configure_agents_for_mcp( requested: list[str], *, prompt_optional_updates: bool = True ) -> set[str]: """Ensure the named coding agents are set up (workspace + models) so a - subsequent `ucode mcp add` has them as targets, and return their canonical - names. Mirrors `ucode configure --agents`: model agents go through + subsequent `ug mcp add` has them as targets, and return their canonical + names. Mirrors `ug configure --agents`: model agents go through configure_workspace_command (which installs binaries and configures models); Cursor is MCP-only, so it just needs workspace state established and rides along via MCP_ONLY_CLIENTS. Interactive — prompts for the workspace URL on @@ -1250,7 +1248,7 @@ def mcp_add( ) -> None: """Add Databricks MCP servers to installed coding tools. - Like `ucode configure mcp`, but purely additive: it never removes MCP servers + Like `ug configure mcp`, but purely additive: it never removes MCP servers that are already configured, only registers new ones. Pass --agents to target (and, if needed, set up) specific agents. """ @@ -1333,14 +1331,14 @@ def mcp_proxy_cmd( "--use-pat", help="Authenticate with the profile's static personal access token (from " "~/.databrickscfg) instead of OAuth. Set automatically for workspaces configured " - "with `ucode configure --profiles --use-pat`.", + "with `ug configure --profiles --use-pat`.", ), ] = False, ) -> None: """Bridge a coding agent's stdio MCP transport to a Databricks MCP endpoint. Each configured client spawns this as a local stdio MCP server (see - `ucode configure mcp`); it forwards messages to ``--url`` and injects a + `ug configure mcp`); it forwards messages to ``--url`` and injects a freshly-minted token on every upstream request, so it never expires mid-session. Not meant for interactive use — the agent manages this process's lifecycle.""" @@ -1349,7 +1347,7 @@ def mcp_proxy_cmd( state = load_state() workspace = host or state.get("workspace") if not workspace: - print_err("No workspace configured. Run `ucode configure` first.") + print_err("No workspace configured. Run `ug configure` first.") raise typer.Exit(1) profile = profile or state.get("profile") serve(url, workspace, profile, use_pat=use_pat or bool(state.get("use_pat"))) @@ -1379,7 +1377,7 @@ def auth_token_cmd( state = load_state() workspace = host or state.get("workspace") if not workspace: - print_err("No workspace configured. Run `ucode configure` first.") + print_err("No workspace configured. Run `ug configure` first.") raise typer.Exit(1) profile = profile or state.get("profile") if use_pat or state.get("use_pat"): @@ -1392,7 +1390,7 @@ def auth_token_cmd( f"--use-pat: no personal access token available for profile " f"'{profile or ''}'. Add a `token = ` entry under " f"[{profile or 'your-profile'}] in ~/.databrickscfg, or re-run " - "`ucode configure` without --use-pat to use OAuth." + "`ug configure` without --use-pat to use OAuth." ) raise typer.Exit(1) try: @@ -1735,7 +1733,7 @@ def _fetch_budget_recommendation(state: dict, managed: dict | None) -> dict | No def _launch_title(tool: str) -> str: - return f"Launching {TOOL_SPECS[tool]['display']} with Unity Gateway" + return f"Launching {TOOL_SPECS[tool]['display']} with ug" def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = None) -> None: @@ -1914,8 +1912,8 @@ def _launch_tool( # which tier the service offers to launch on, rather than being rejected — see the provider # branch below. # An explicit --workspace targets that workspace for this launch (and - # auto-configures it if unseen), so `ucode claude --provider ... --workspace ...` - # works without a prior `ucode configure`. + # auto-configures it if unseen), so `ug claude --provider ... --workspace ...` + # works without a prior `ug configure`. if workspace_url: set_current_workspace(normalize_workspace_url(workspace_url)) existing = load_state() @@ -1934,7 +1932,7 @@ def _launch_tool( # silently override a provider the user typed on the command line (it errors instead). explicit_provider = provider # An explicit --provider overrides the persisted choice; otherwise fall - # back to whatever `ucode configure` saved for this tool. + # back to whatever `ug configure` saved for this tool. provider = provider or get_provider_service(state, tool) state = _migrate_legacy_smart_routing(state) if _can_launch_from_cached_config( @@ -1962,7 +1960,7 @@ def _launch_tool( # Discovery exists to find models and isn't needed for managed config that already names them. managed_models_known = managed_supplies_models(managed, tool) # Re-fetch model lists on every launch so newly-added Databricks - # endpoints show up without a manual `ucode configure` (and so that + # endpoints show up without a manual `ug configure` (and so that # tools like pi which read multiple model bundles never run on # stale state from before a tool added a new bundle). Under a provider # this heavy discovery is skipped (only a web-search model is fetched). @@ -2115,7 +2113,7 @@ def _launch_tool( print_kv("Smart routing", "enabled") print_note( f"{TOOL_SPECS[tool]['display']} may require one-time hook review. Open " - "`/hooks` and trust the ucode routing hooks if prompted." + "`/hooks` and trust the ug routing hooks if prompted." ) if tool in ("gemini", "opencode", "copilot", "pi"): print_note( @@ -2152,7 +2150,7 @@ def _launch_tool( # Launch-only escape hatch for managed/headless launchers (e.g. omnigent) that -# have already run `ucode configure`: skip the ~5-10s per-launch auth + AI +# have already run `ug configure`: skip the ~5-10s per-launch auth + AI # Gateway re-validation. Distinct from the configure-only `--skip-validate`, # which skips the model smoke test, and from `--skip-managed-config`, which # controls whether the workspace's managed config is applied. @@ -2161,7 +2159,7 @@ def _launch_tool( typer.Option( "--skip-preflight", help="Skip the per-launch Databricks auth + AI Gateway re-validation, trusting a " - "prior `ucode configure`.", + "prior `ug configure`.", ), ] @@ -2171,7 +2169,7 @@ def _launch_tool( ) # Ignore the workspace's managed coding-agent config for this one command, on both -# `ucode configure` and the launchers. Accepted (and no-op) even when the managed-config +# `ug configure` and the launchers. Accepted (and no-op) even when the managed-config # feature is off, so a headless launcher can always pass it. SkipManagedConfigOption = Annotated[ bool, @@ -2197,7 +2195,7 @@ def _disable_managed_config_if_requested(skip_managed_config: bool) -> None: # Target this launch at a specific workspace, auto-configuring (and logging in) -# if it hasn't been set up yet — so a launch needs no prior `ucode configure`. +# if it hasn't been set up yet — so a launch needs no prior `ug configure`. WorkspaceOption = Annotated[ str | None, typer.Option( @@ -2216,7 +2214,7 @@ def default( typer.Option( "--version", "-V", - help="Show the ucode version and exit.", + help="Show the ug version and exit.", callback=_version_callback, is_eager=True, ), @@ -2235,6 +2233,8 @@ def default( ) -> None: """Configure and launch coding agents through Databricks AI Gateway. + The primary command is `ug`; `ucode` remains supported as an alias. + With no subcommand, launches the agent your workspace's managed config selects. """ if ctx.invoked_subcommand is not None: @@ -2272,7 +2272,7 @@ def _launch_managed_default( state = load_state() current = state.get("workspace") if not current: - raise RuntimeError("No workspace configured. Run `ucode configure` first.") + raise RuntimeError("No workspace configured. Run `ug configure` first.") apply_pat_environment(state) # --dry-run avoids the fetch but still applies the last saved config. if dry_run: @@ -2293,7 +2293,7 @@ def _launch_managed_default( if not isinstance(tool, str) or not tool: raise RuntimeError( "Your workspace's managed config names no agent to launch. Ask an admin to set a " - "default agent, or run `ucode ` directly." + "default agent, or run `ug ` directly." ) _print_managed_summary(managed, state, tool, abridged=True) _launch_tool( @@ -2318,10 +2318,10 @@ def _print_no_managed_config_guidance(workspace: str, profile: str | None) -> No with spinner("Checking your workspace permissions..."): is_admin = is_workspace_admin(workspace, token) if is_admin is False: - print_note("Ask a workspace admin to set one up with `ucode setup`.") + print_note("Ask a workspace admin to set one up with `ug setup`.") else: # None means the admin check itself failed; point at setup rather than a dead end. - print_note("Run `ucode setup` to configure one for your workspace, then `ucode publish`.") + print_note("Run `ug setup` to configure one for your workspace, then `ug publish`.") @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -2358,7 +2358,7 @@ def codex_cmd( typer.Option( "--disable-smart-routing", hidden=True, - help="Disable smart routing and remove ucode's Codex routing hooks.", + help="Disable smart routing and remove ug's Codex routing hooks.", ), ] = False, ) -> None: @@ -2369,7 +2369,7 @@ def codex_cmd( raise typer.Exit(1) if disable_smart_routing_flag: codex_agent.disable_smart_routing(load_state()) - print_success("Codex smart routing disabled; ucode routing hooks removed") + print_success("Codex smart routing disabled; ug routing hooks removed") return with _smart_routing_v2_flag(enable_smart_routing_flag): _launch_tool( @@ -2435,7 +2435,7 @@ def claude_cmd( typer.Option( "--disable-smart-routing", hidden=True, - help="Disable smart routing and remove ucode's Claude Code routing hooks.", + help="Disable smart routing and remove ug's Claude Code routing hooks.", ), ] = False, ) -> None: @@ -2446,7 +2446,7 @@ def claude_cmd( raise typer.Exit(1) if disable_smart_routing_flag: claude_agent.disable_smart_routing(load_state()) - print_success("Claude Code smart routing disabled; ucode routing hooks removed") + print_success("Claude Code smart routing disabled; ug routing hooks removed") return if enable_model_discovery: os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" @@ -2513,10 +2513,10 @@ def cursor_cmd(ctx: typer.Context) -> None: """Launch Cursor Agent. Cursor is MCP-only: `cursor-agent` runs models on your own Cursor account, so - ucode configures no models for it. Its Databricks MCP servers (added via - `ucode configure mcp`) run `ucode mcp-proxy`, which authenticates itself — so + ug configures no models for it. Its Databricks MCP servers (added via + `ug configure mcp`) run `ug mcp-proxy`, which authenticates itself — so this command is a thin convenience wrapper over `cursor-agent`, kept for - symmetry with the other `ucode ` launchers. + symmetry with the other `ug ` launchers. """ from ucode.agents import cursor @@ -2524,12 +2524,12 @@ def cursor_cmd(ctx: typer.Context) -> None: if not shutil.which(cursor.CURSOR_BINARY): raise RuntimeError( f"`{cursor.CURSOR_BINARY}` was not found on PATH. Install Cursor Agent " - "(https://cursor.com/cli), then re-run `ucode cursor`." + "(https://cursor.com/cli), then re-run `ug cursor`." ) - print_section("ucode with Cursor") + print_section("ug with Cursor") print_note( "Cursor runs models on your Cursor account; its Databricks MCP servers " - "authenticate through `ucode mcp-proxy`." + "authenticate through `ug mcp-proxy`." ) print_success("Starting Cursor Agent") cursor.launch(load_state(), ctx.args) @@ -2685,7 +2685,7 @@ def configure( if use_pat and profiles is None: raise RuntimeError( "--use-pat requires --profiles. Pass the PAT-backed Databricks CLI " - "profile(s) explicitly, e.g. `ucode configure --profiles DEFAULT --use-pat`." + "profile(s) explicitly, e.g. `ug configure --profiles DEFAULT --use-pat`." ) # Skipping only has meaning against an explicit agent list: the interactive # picker already offers just the available agents, and --agent names a @@ -2693,7 +2693,7 @@ def configure( if skip_unavailable and agents is None: raise RuntimeError( "--skip-unavailable requires --agents. It selects the available subset " - "of an explicit agent list, e.g. `ucode configure --agents claude,codex,pi " + "of an explicit agent list, e.g. `ug configure --agents claude,codex,pi " "--skip-unavailable`." ) workspace_entries = _parse_workspaces_option(workspaces) if workspaces is not None else None @@ -2750,7 +2750,7 @@ def configure( # Cursor is MCP-only (no model routing), so it can't go through the # model-agent configure path. Split it out: model agents configure # normally; cursor only needs workspace state established here, and - # its MCP servers are added separately via `ucode configure mcp` + # its MCP servers are added separately via `ug configure mcp` # (which picks cursor up through MCP_ONLY_CLIENTS). If cursor is the # only agent, do a workspace-only configure so that later `configure # mcp` run has a current workspace to target. @@ -2850,7 +2850,7 @@ def configure( if bare: raise RuntimeError( "--mcp names must be fully qualified `..` " - f"(got: {', '.join(bare)}). Use `ucode configure mcp` for the " + f"(got: {', '.join(bare)}). Use `ug configure mcp` for the " "interactive picker." ) configure_mcp_command(services=services) @@ -3009,13 +3009,13 @@ def setup( typer.Option( "--from-file", help="Skip the interactive flow and load a hand-written managed config (JSON, in " - "ucode's manifest shape) instead. Validated before it is saved.", + "ug's manifest shape) instead. Validated before it is saved.", ), ] = None, ) -> None: """Choose the agents and models for your workspace's managed config (admins only). - MCP servers, skills, and the tiered spend policy have their own commands — see `ucode setup help`. + MCP servers, skills, and the tiered spend policy have their own commands — see `ug setup help`. """ if ctx.invoked_subcommand is not None: return @@ -3110,7 +3110,7 @@ def setup_help_cmd() -> None: @setup_app.command("show") def setup_show_cmd() -> None: - """Print the authored managed config and the payload `ucode publish` would publish.""" + """Print the authored managed config and the payload `ug publish` would publish.""" try: code = show_command() except RuntimeError as exc: @@ -3127,7 +3127,7 @@ def publish_cmd( typer.Option( "--file", "-f", - help="Publish a config file exported with `ucode export` instead of the locally " + help="Publish a config file exported with `ug export` instead of the locally " "authored config. Its `workspace` must match the configured workspace.", ), ] = None, @@ -3139,7 +3139,7 @@ def publish_cmd( """Publish this workspace's managed coding config (workspace admins only). Always validates the manifest before publishing (and shows what would change, then confirms), so - there is no separate dry-run: `ucode setup` only ever writes a valid manifest, and a + there is no separate dry-run: `ug setup` only ever writes a valid manifest, and a hand-editing admin sees any error here before anything reaches the workspace. """ # See the `setup` callback: `typer.Exit` subclasses RuntimeError, so it must be raised after @@ -3172,7 +3172,7 @@ def export_cmd( """Export this workspace's managed coding-agent config as portable JSON. Serializes the local managed config to the external `CodingAgentConfig` format that - `ucode publish -f ` consumes, with credentials and server-owned fields (resource name, + `ug publish -f ` consumes, with credentials and server-owned fields (resource name, workspace id, timestamps, user ids) excluded. Any user can run it; it makes no network calls and mutates no workspace or local state. Without --file the JSON is printed to stdout; diagnostics and errors go to stderr. @@ -3198,7 +3198,7 @@ def status_cmd() -> None: @app.command("revert") def revert_cmd() -> None: - """Clear ucode state and restore backed-up agent config files.""" + """Clear ug state and restore backed-up agent config files.""" try: revert() except RuntimeError as exc: @@ -3208,7 +3208,7 @@ def revert_cmd() -> None: @app.command("doctor") def doctor_cmd() -> None: - """Diagnose the local ucode setup and offer to fix any problems found.""" + """Diagnose the local ug setup and offer to fix any problems found.""" from ucode.doctor import doctor try: @@ -3236,24 +3236,32 @@ def usage_cmd( @app.command("upgrade") def upgrade_cmd() -> None: - """Upgrade ucode to the latest version from GitHub.""" + """Upgrade ug to the latest version from GitHub.""" import subprocess git_url = "git+https://github.com/databricks/ucode" print_section("Upgrade") print_kv("Source", git_url) try: + # uv tracks tools by distribution name, so it cannot upgrade the old + # `ucode` distribution across the future rename to `unity-gateway`. + # Removing the legacy registration is best-effort: users who are + # already migrated will not have a tool registered under that name. + subprocess.run( + ["uv", "tool", "uninstall", "ucode"], + check=False, + ) subprocess.run( - ["uv", "tool", "install", "--reinstall", git_url], + ["uv", "tool", "install", "--force", git_url], check=True, ) except FileNotFoundError: - print_err("`uv` was not found on PATH. Install uv to upgrade ucode.") + print_err("`uv` was not found on PATH. Install uv to upgrade ug.") raise typer.Exit(1) from None except subprocess.CalledProcessError as exc: print_err(f"Upgrade failed (exit code {exc.returncode}).") raise typer.Exit(1) from None - print_success("ucode upgraded") + print_success("ug upgraded; `ucode` remains available as an alias") def main() -> None: diff --git a/src/ucode/managed_export.py b/src/ucode/managed_export.py index 511d331d..670ff661 100644 --- a/src/ucode/managed_export.py +++ b/src/ucode/managed_export.py @@ -40,8 +40,8 @@ def build_export_payload() -> dict: manifest = load_managed_state(workspace) if not manifest: raise RuntimeError( - "No managed coding-agent config found locally. Run `ucode setup` to author one, or run " - "`ucode` against a workspace that publishes one, then re-run `ucode export`." + "No managed coding-agent config found locally. Run `ug setup` to author one, or run " + "`ug` against a workspace that publishes one, then re-run `ug export`." ) errors = validate_manifest(manifest, None) if errors: diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 6854f833..6ab4e623 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -1,18 +1,18 @@ -"""Interactive `ucode setup`: author the workspace's managed coding-agent config. +"""Interactive `ug setup`: author the workspace's managed coding-agent config. Workspace admins run this to build the ``CodingAgentConfig`` their developers will pull, then publish -it with ``ucode publish`` (a separate command, so the manifest can be reviewed first). The config lives +it with ``ug publish`` (a separate command, so the manifest can be reviewed first). The config lives at ``~/.ucode/managed-state.json`` (the one local managed-config file, owned by :mod:`ucode.managed_config`). Authoring is split across commands so an admin can change one part without walking the whole flow: -``ucode setup`` picks the agents and models, and ``ucode setup mcps`` / ``skills`` / ``spend-tiers`` -each edit their own section of the same manifest. ``ucode setup`` carries the other sections forward -untouched (:func:`_carry_forward_sections`), and ``ucode setup help`` prints the whole sequence. +``ug setup`` picks the agents and models, and ``ug setup mcps`` / ``skills`` / ``spend-tiers`` +each edit their own section of the same manifest. ``ug setup`` carries the other sections forward +untouched (:func:`_carry_forward_sections`), and ``ug setup help`` prints the whole sequence. Serialization, validation, and the per-agent model catalogs live in :mod:`ucode.managed_setup`; this module is the interaction layer on top of them. Sub-flows an admin already knows — MCP, skills — are -delegated to the existing ``ucode configure `` commands and their results read back out of +delegated to the existing ``ug configure `` commands and their results read back out of ``state.json``, so there is exactly one picker per concern in the codebase. """ @@ -84,7 +84,7 @@ # Shown whenever the workspace's coding-agent-config APIs return FEATURE_DISABLED. CODING_AGENT_CONFIGS_DISABLED_MESSAGE = ( "Workspace-managed coding agent configuration is not available on this workspace. Use " - "`ucode configure` to set up agents for individual users instead." + "`ug configure` to set up agents for individual users instead." ) BUDGET_POLICY_BLURB = ( @@ -95,14 +95,14 @@ "budget's own hard block is what actually caps spend." ) -# Agents not offered in `ucode setup`'s picker, even when the workspace serves their models. -# `ucode gemini` still works as a launch target; it's just not part of the managed config authored +# Agents not offered in `ug setup`'s picker, even when the workspace serves their models. +# `ug gemini` still works as a launch target; it's just not part of the managed config authored # here. Serialize/validate keep supporting it, so a `--from-file` manifest can still name it. SETUP_EXCLUDED_AGENTS = frozenset({"gemini"}) def _tracing_table_from_state(state: dict) -> str | None: - """The UC table `ucode configure tracing` wired up, or None when tracing is off. + """The UC table `ug configure tracing` wired up, or None when tracing is off. ``configure tracing`` records the destination as ``uc_destination``; the managed config calls the same thing ``tracing.table``. @@ -118,11 +118,11 @@ def _mcp_server_from_url(url: str) -> tuple[str, str] | None: """Derive a managed-config ``(name, type)`` entry from a registered server's resolved URL. ``state.json`` stores each MCP server's resolved URL but not its type, while the managed config - stores ``{name, type}`` and lets the developer's ucode rebuild the URL. So map the URL back to the + stores ``{name, type}`` and lets the developer's ug rebuild the URL. So map the URL back to the type *and* the identifier the ai-gateway ``McpServer.name`` field is meant to hold for that type (a UC name for a UC service, a Genie space id for a genie space, a `.` for vector-search / uc-functions, a connection name for external). Deriving ``name`` from the URL — - rather than reusing the local display slug — is what lets the developer's ucode reconstruct the + rather than reusing the local display slug — is what lets the developer's ug reconstruct the URL on launch. Returns None for a URL that matches nothing reconstructable (e.g. an app's off-workspace host), so those are skipped rather than published unusably. """ @@ -178,7 +178,7 @@ def _mcp_servers_from_state(state: dict) -> list[dict]: resolved = _mcp_server_from_url(url) if resolved is None: print_warning( - f"Skipping MCP server '{name}': ucode can't publish it to a managed config " + f"Skipping MCP server '{name}': ug can't publish it to a managed config " f"(unrecognized or app-hosted URL: {url})." ) continue @@ -400,7 +400,7 @@ def _prompt_models_for_agent(tool: str, state: dict, provider_service: dict | No # Gemini and Copilot do declare `repeated string models`, but their config writers take one model # (`gemini.write_tool_config(state, model)` / `copilot.write_tool_config(state, model)`) and write a # single env var — so a published list would be read by nothing. Offering one keeps the manifest -# honest about what ucode can apply; widen this when those writers grow a picker. +# honest about what ug can apply; widen this when those writers grow a picker. SINGLE_MODEL_AGENTS = frozenset({"codex", "gemini", "copilot"}) # Skip sentinel for a Claude family prompt. Every `ClaudeDefaultModels` slot is optional, and an @@ -567,7 +567,7 @@ def _prompt_claude_provider_family_models(targets: list[str], service_name: str) } # Quick setup: fill each family with the service's newest id (highest version, broadest region), - # the same pick a developer's own `ucode configure` would make. The alternative is choosing a + # the same pick a developer's own `ug configure` would make. The alternative is choosing a # specific id per family — e.g. to pin an older, validated version or a particular region. # map_claude_family_models covers opus/sonnet/haiku but not fable, so a fable-only service has # nothing to quick-fill — only offer quick setup when it would actually populate a slot. @@ -893,8 +893,8 @@ def _prompt_budget_policy( model it wasn't given, which neither this validation nor the server's would reject: the tier would activate and hand the developer a model their agent doesn't have. - Asks no "set up a tiered spend policy?" gate — running `ucode setup spend-tiers` is the answer to - that question, the same way `ucode configure ` needs no confirmation. + Asks no "set up a tiered spend policy?" gate — running `ug setup spend-tiers` is the answer to + that question, the same way `ug configure ` needs no confirmation. """ print_section("Tiered Spend Policy") @@ -908,7 +908,7 @@ def _prompt_budget_policy( print_warning_panel( "No AI Gateway budgets are visible for this workspace, so there is nothing to attach a " "policy to. Create a budget in the Databricks console first, then re-run " - "`ucode setup spend-tiers`. Currently, only AI Gateway budgets with hard blocks are " + "`ug setup spend-tiers`. Currently, only AI Gateway budgets with hard blocks are " "eligible to be associated with Tiered Spend Policies." ) return None @@ -923,7 +923,7 @@ def _prompt_budget_policy( "None of this workspace's AI Gateway budgets have a per-user threshold with a usage " "block configured, which spend routing enforces. Add a per-user alert threshold with a " "block action to a budget in the Databricks console, then re-run " - "`ucode setup spend-tiers`." + "`ug setup spend-tiers`." ) return None @@ -1194,7 +1194,7 @@ def _render_config_diff(existing: dict | None, incoming: dict, workspace: str) - Returns True when there is a difference. Lists only what changes — labelled ADD, DELETE, or CHANGE (``old → new``) — since the full config was just printed by :func:`_render_summary` above; - repeating the unchanged rows here would bury the actual delta. Both configs are in ucode's + repeating the unchanged rows here would bury the actual delta. Both configs are in ug's normalized shape (the caller round-trips the local manifest through serialize/normalize first), so the comparison is field-for-field with what the workspace holds. """ @@ -1238,12 +1238,12 @@ def _require_admin(workspace: str, token: str) -> None: admin = is_workspace_admin(workspace, token) if admin is False: raise RuntimeError( - f"You are not an admin of {workspace}. `ucode setup` authors the workspace-wide " + f"You are not an admin of {workspace}. `ug setup` authors the workspace-wide " "coding config, so it is restricted to workspace admins." ) if admin is None: print_warning( - "Could not verify workspace admin permissions. Continuing — `ucode publish` will fail " + "Could not verify workspace admin permissions. Continuing — `ug publish` will fail " "if you lack them." ) else: @@ -1257,10 +1257,10 @@ def _handle_existing_config(workspace: str, token: str) -> tuple[bool, dict | No replaces the existing config) and False to stop (the admin chose to delete it instead). ``existing`` is the published config when one was read, so the caller can carry its MCP servers / skills / tracing / budget policy forward — the local draft may be missing on a fresh machine or after - ``ucode revert``, and without this those sections would be silently dropped on the next publish. + ``ug revert``, and without this those sections would be silently dropped on the next publish. Deliberately doesn't itemize what the existing config holds. The admin doesn't need an inventory - to act on this, and `ucode setup show` prints the real thing for anyone who wants to compare. + to act on this, and `ug setup show` prints the real thing for anyone who wants to compare. """ with spinner("Checking for an existing managed config..."): existing, reason = get_managed_config(workspace, token) @@ -1309,17 +1309,17 @@ def _delete_existing_config(workspace: str, token: str, existing: dict) -> None: """Delete the workspace's published config after confirming. Raises RuntimeError on failure. Deleting leaves the workspace with no managed config, so every developer falls back to their own - settings on their next ucode run — confirm before doing it. + settings on their next ug run — confirm before doing it. """ name = existing.get("name") if not isinstance(name, str): raise RuntimeError( "This workspace has a managed config but the API didn't return its resource name, so " - "ucode can't delete it. Delete it in the workspace directly." + "ug can't delete it. Delete it in the workspace directly." ) print_warning( "Deleting removes the managed config entirely. Every developer falls back to their own " - "settings on their next ucode run." + "settings on their next ug run." ) if not prompt_yes_no_default("Delete the existing managed config?", default=False): print_note("Nothing was deleted.") @@ -1335,7 +1335,7 @@ def setup_from_file(path: str) -> int: """Validate an admin-written manifest and save it, skipping the interactive flow. The non-interactive path for CI and for admins who'd rather keep the JSON in version control. - Reads ucode's own manifest shape (the same thing the wizard writes), not proto-JSON. + Reads ug's own manifest shape (the same thing the wizard writes), not proto-JSON. """ manifest_path = Path(path).expanduser() try: @@ -1355,7 +1355,7 @@ def setup_from_file(path: str) -> int: workspace = state.get("workspace") if not workspace: raise RuntimeError( - "No workspace is configured. Run `ucode configure` first so ucode knows which " + "No workspace is configured. Run `ug configure` first so ug knows which " "workspace this manifest is for." ) @@ -1373,13 +1373,13 @@ def setup_from_file(path: str) -> int: return 0 -# The sections that have their own `ucode setup ` command, in the order the checklist lists +# The sections that have their own `ug setup ` command, in the order the checklist lists # them: the command, the label the summary uses, and how to tell whether the manifest has one. SETUP_SECTIONS: list[tuple[str, str, Callable[[dict], bool]]] = [ - ("ucode setup mcps", "MCP servers", lambda m: bool(m.get("mcp_servers"))), - ("ucode setup skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))), + ("ug setup mcps", "MCP servers", lambda m: bool(m.get("mcp_servers"))), + ("ug setup skills", "Skills", lambda m: bool((m.get("skills") or {}).get("names"))), ( - "ucode setup spend-tiers", + "ug setup spend-tiers", "Tiered Spend Policy", lambda m: isinstance(m.get("budget_policy"), dict), ), @@ -1391,12 +1391,12 @@ def _command_line(command: str, description: str, *, marker: str = " ", width: i return f" {marker} [bold]{command.ljust(width)}[/bold] {description}" -# `ucode setup` walks these phases in order; the banners announce each one so the admin can see how +# `ug setup` walks these phases in order; the banners announce each one so the admin can see how # far along the flow they are, the way a multi-page form numbers its pages. SETUP_STEP_TITLES = ["Coding agents", "Models & settings", "Default agent"] -def _step_banner(index: int, title: str, command_label: str = "ucode setup") -> None: +def _step_banner(index: int, title: str, command_label: str = "ug setup") -> None: """Announce one phase of the flow as `step N of M`, branded to the invoking command.""" print_section(f"{command_label} · step {index} of {len(SETUP_STEP_TITLES)} · {title}") @@ -1445,7 +1445,7 @@ def _print_next_steps(manifest: dict) -> None: if config_io.is_dry_run(): # Under --dry-run nothing was written, so the section commands (which read the saved draft) # and `publish` have nothing to act on. Say so rather than send the admin to commands that - # would report "run `ucode setup` first". + # would report "run `ug setup` first". print_note("Dry run — nothing was saved. Re-run without --dry-run to author the config.") return # These sections aren't required to publish — call them out as optional so an admin doesn't read @@ -1455,41 +1455,41 @@ def _print_next_steps(manifest: dict) -> None: console.print(line) print_panel( "All done?", - ["Publish with [bold]ucode publish[/bold] so all developers use this configuration."], + ["Publish with [bold]ug publish[/bold] so all developers use this configuration."], ) def _offer_publish() -> None: """Offer to publish the saved draft right away, so an admin can apply changes incrementally. - Each `ucode setup` command only writes a local draft. Without this an admin has to remember to run - `ucode publish` separately, and a `ucode setup` re-run in the meantime is easy to mistake for having + Each `ug setup` command only writes a local draft. Without this an admin has to remember to run + `ug publish` separately, and a `ug setup` re-run in the meantime is easy to mistake for having lost the change. Answering yes runs `publish_command`, which shows the diff against the published - config as it publishes; declining leaves the draft for a later `ucode publish`. Skipped under + config as it publishes; declining leaves the draft for a later `ug publish`. Skipped under --dry-run, where nothing was saved to publish. """ if config_io.is_dry_run(): return console.print() if not prompt_yes_no_default( - "Publish these changes to the workspace now? (runs `ucode publish`)", default=False + "Publish these changes to the workspace now? (runs `ug publish`)", default=False ): - print_note("Draft saved. Run `ucode publish` when you're ready to publish.") + print_note("Draft saved. Run `ug publish` when you're ready to publish.") return publish_command(yes=True) -# The sections `ucode setup` carries forward instead of prompting for, and how to rebuild each one. +# The sections `ug setup` carries forward instead of prompting for, and how to rebuild each one. CARRIED_SECTIONS: list[tuple[str, str, str]] = [ - ("mcp_servers", "MCP servers", "ucode setup mcps"), - ("skills", "Skills", "ucode setup skills"), - ("tracing_table", "Tracing table", "ucode setup --from-file"), - ("budget_policy", "Tiered Spend Policy", "ucode setup spend-tiers"), + ("mcp_servers", "MCP servers", "ug setup mcps"), + ("skills", "Skills", "ug setup skills"), + ("tracing_table", "Tracing table", "ug setup --from-file"), + ("budget_policy", "Tiered Spend Policy", "ug setup spend-tiers"), ] def _carry_forward_sections(previous: dict, manifest: dict) -> None: - """Copy the sections `ucode setup` no longer prompts for out of a previously authored config. + """Copy the sections `ug setup` no longer prompts for out of a previously authored config. `setup` writes the whole manifest, so without this a re-run would silently clear the MCP servers, skills, tracing table, and budget policy an admin authored with the other commands — they'd have @@ -1499,7 +1499,7 @@ def _carry_forward_sections(previous: dict, manifest: dict) -> None: fits. Otherwise a carried section could make the manifest invalid and block the save outright, with no way out: the commands that could repair a section read the very manifest that can't be written. The live case is a budget-policy tier naming an agent the admin just de-selected, but - hand-edited drafts and configs authored by an older ucode can trip the others the same way. + hand-edited drafts and configs authored by an older ug can trip the others the same way. """ # Validating against no inventory keeps this to structural checks, which is all that's at stake # here: the models were just picked from the workspace's own catalog a few prompts ago. @@ -1530,27 +1530,27 @@ def setup_command( *, workspace: str | None = None, profile: str | None = None, - command_label: str = "ucode setup", + command_label: str = "ug setup", token: str | None = None, ) -> int: """Author the agents and models half of the workspace's managed coding config interactively. Agents and per-agent models only. MCP servers, skills, and the tiered spend policy each have their - own command (`ucode setup mcps` / `skills` / `spend-tiers`), so an admin changing one of them doesn't + own command (`ug setup mcps` / `skills` / `spend-tiers`), so an admin changing one of them doesn't have to walk the whole flow again — and this command carries whatever they already authored forward untouched rather than clearing it (:func:`_carry_forward_sections`). ``workspace``/``profile`` let a caller that has already resolved (and authenticated against) a - workspace hand it in so the admin isn't prompted to pick one again — e.g. `ucode configure` + workspace hand it in so the admin isn't prompted to pick one again — e.g. `ug configure` launching setup after its admin offer. When ``workspace`` is None the flow prompts as usual. - ``command_label`` brands the section headers to the invoking command: `ucode configure` passes - "Configure Unity Gateway" so a user who never typed `ucode setup` isn't jarred by it (the - standalone `ucode setup` command keeps the default). References to specific sub-commands (`ucode - setup mcps`, `ucode apply`, …) stay verbatim — those are real command names, not branding. + ``command_label`` brands the section headers to the invoking command: `ug configure` passes + "Configure ug" so a user who never typed `ug setup` isn't jarred by it (the + standalone `ug setup` command keeps the default). References to specific sub-commands (`ucode + setup mcps`, `ug apply`, …) stay verbatim — those are real command names, not branding. ``token`` lets a caller that already authenticated and admin-checked the workspace (e.g. - `ucode configure`) hand its token in, so setup's admin gate uses the *same* token as the routing + `ug configure`) hand its token in, so setup's admin gate uses the *same* token as the routing decision — a second fetch here could resolve a different identity right after a credential switch and reject a caller configure just treated as an admin. When None, setup authenticates and fetches its own token as usual. @@ -1567,7 +1567,7 @@ def setup_command( print_section(command_label) print_note("Choose the coding agents and models for this workspace's managed config.") - print_note("Developers pull it automatically when they run ucode.") + print_note("Developers pull it automatically when they run ug.") if workspace is None: workspace, profile = _prompt_for_configuration() @@ -1601,7 +1601,7 @@ def setup_command( ) # The local draft is the carry-forward source, falling back to what's published on the workspace: - # a fresh machine (or one after `ucode revert`) has no draft, and without the fallback the next + # a fresh machine (or one after `ug revert`) has no draft, and without the fallback the next # publish would silently wipe the workspace's MCP servers, skills, tracing, and budget policy. previous = load_managed_state(workspace) or published or {} previously_enabled = [ @@ -1650,7 +1650,7 @@ def setup_command( # Tracing is intentionally not prompted here: the managed-tracing path isn't working yet, so # asking would author a `tracing_table` the workspace can't honor. The manifest field and its # serialize/validate support stay in place, so a hand-written `--from-file` config can still set - # it once the backend is ready. Re-add a `ucode setup tracing` command when it is. + # it once the backend is ready. Re-add a `ug setup tracing` command when it is. _carry_forward_sections(previous, manifest) errors = validate_manifest(manifest, state) @@ -1674,11 +1674,11 @@ def setup_command( def _resolve_admin_workspace() -> tuple[str, str | None, str]: """Resolve the workspace a section command edits, authenticate, and gate on admin. - Returns ``(workspace, profile, token)``. Unlike `ucode setup`, this doesn't prompt for a workspace + Returns ``(workspace, profile, token)``. Unlike `ug setup`, this doesn't prompt for a workspace and takes it strictly from local state rather than falling back to the draft file's workspace: the MCP and skills pickers re-read ``current_workspace`` themselves (via ``setup_mcp_clients``), so a mismatch would have them operate against one workspace while the manifest is saved for another. - Requiring ``ucode configure`` to have set the current workspace keeps the two in lockstep. It also + Requiring ``ug configure`` to have set the current workspace keeps the two in lockstep. It also skips :func:`_handle_existing_config` — the create-or-delete choice belongs to authoring a config, not to changing one section of it. """ @@ -1686,7 +1686,7 @@ def _resolve_admin_workspace() -> tuple[str, str | None, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError( - "No workspace is configured. Run `ucode configure` first, then `ucode setup` to author " + "No workspace is configured. Run `ug configure` first, then `ug setup` to author " "this workspace's managed config." ) profile = state.get("profile") @@ -1697,7 +1697,7 @@ def _resolve_admin_workspace() -> tuple[str, str | None, str]: def _manifest_for_edit(workspace: str) -> dict: - """The authored manifest a section command edits. Raises when `ucode setup` hasn't run. + """The authored manifest a section command edits. Raises when `ug setup` hasn't run. An empty ``enabled_agents`` counts as "hasn't run": a launch records ``{}`` for a workspace with no managed config (see ``refresh_managed_config``), so the file existing is not proof an admin @@ -1707,7 +1707,7 @@ def _manifest_for_edit(workspace: str) -> dict: manifest = load_managed_state(workspace) if not (manifest or {}).get("enabled_agents"): raise RuntimeError( - f"No managed config has been authored for {workspace} yet. Run `ucode setup` first to " + f"No managed config has been authored for {workspace} yet. Run `ug setup` first to " "pick the agents and models, then re-run this command." ) return cast(dict, manifest) @@ -1720,7 +1720,7 @@ def _save_section_update(workspace: str, manifest: dict) -> int: availability. That's deliberate: a section command doesn't touch agents or models, so re-checking them would only reject a legitimately pinned older Claude model (`load_state` keeps just the newest per family) or, worse, wrongly flag a codex/gemini model whenever the re-fetched inventory happens - to be Claude-only. `ucode publish` runs the full model check against the live catalog before + to be Claude-only. `ug publish` runs the full model check against the live catalog before publishing, which is where it belongs. """ errors = validate_manifest(manifest, None) @@ -1740,12 +1740,12 @@ def _save_section_update(workspace: str, manifest: dict) -> int: def setup_mcp_command() -> int: - """Author the managed config's MCP servers (`ucode setup mcps`).""" + """Author the managed config's MCP servers (`ug setup mcps`).""" workspace, _, _ = _resolve_admin_workspace() manifest = _manifest_for_edit(workspace) print_section("Managed MCP servers") - print_note("Developers get these MCP servers registered automatically when they run ucode.") + print_note("Developers get these MCP servers registered automatically when they run ug.") from ucode.mcp import configure_mcp_command # Snapshot the managed-shaped servers before the picker so a cancelled run on an empty local @@ -1758,7 +1758,7 @@ def setup_mcp_command() -> int: configure_mcp_command(exclude_sources={"apps"}) after = _mcp_servers_from_state(load_state()) - # `after == before` isn't enough to call this a no-op: an admin who ran `ucode configure mcp` + # `after == before` isn't enough to call this a no-op: an admin who ran `ug configure mcp` # first arrives with those servers already registered, so confirming the picker leaves local state # unchanged even though the manifest doesn't carry them yet. Also sync when local state already # holds servers the manifest is missing — but only when servers are actually registered, so an Esc @@ -1781,7 +1781,7 @@ def setup_mcp_command() -> int: def setup_skills_command(locations: list[str] | None = None) -> int: - """Author the managed config's skills (`ucode setup skills`). + """Author the managed config's skills (`ug setup skills`). ``locations`` comes from ``--location`` (already parsed to `.` refs); when None the admin is prompted and the answer is parsed the same way. @@ -1790,7 +1790,7 @@ def setup_skills_command(locations: list[str] | None = None) -> int: manifest = _manifest_for_edit(workspace) print_section("Managed skills") - print_note("Developers get these skills downloaded automatically when they run ucode.") + print_note("Developers get these skills downloaded automatically when they run ug.") if locations is None: answer = prompt_for_text( "Skill schemas to publish, comma-separated `catalog.schema` (blank to leave unchanged)", @@ -1818,7 +1818,7 @@ def setup_skills_command(locations: list[str] | None = None) -> int: def setup_budget_policy_command() -> int: - """Author the managed config's tiered spend policy (`ucode setup spend-tiers`).""" + """Author the managed config's tiered spend policy (`ug setup spend-tiers`).""" workspace, _, token = _resolve_admin_workspace() manifest = _manifest_for_edit(workspace) @@ -1851,17 +1851,17 @@ def setup_help_command() -> int: """Walk through the whole managed-config setup, marking what this machine has authored. Hand-written rather than left to `--help`: the point is the *order* of the commands and the fact - that nothing reaches developers until `ucode publish`, neither of which a flag listing conveys. Reads - the manifest but never authenticates, so it works before `ucode configure`. + that nothing reaches developers until `ug publish`, neither of which a flag listing conveys. Reads + the manifest but never authenticates, so it works before `ug configure`. """ - print_section("ucode setup") + print_section("ug setup") print_note( - "A managed config is the coding setup your developers pull automatically — they run ucode " + "A managed config is the coding setup your developers pull automatically — they run ug " "and get the agents, models, MCP servers, and skills you chose here. Admins only." ) print_note( "Each command below edits your local draft; nothing reaches the workspace until " - "`ucode publish`." + "`ug publish`." ) workspace = load_state().get("workspace") or managed_state_workspace() @@ -1869,12 +1869,12 @@ def setup_help_command() -> int: agents_done = bool(manifest.get("enabled_agents")) # One column width across all three groups, so the commands line up as a single list. width = max(len(command) for command, _, _ in SETUP_SECTIONS) - width = max(width, len("ucode setup --from-file ")) + width = max(width, len("ug setup --from-file ")) print_heading("1. Start here") console.print( _command_line( - "ucode setup", + "ug setup", "Agents and models — " + ("[green]configured[/green]" if agents_done else "[yellow]not configured[/yellow]"), marker="[green]✔[/green]" if agents_done else "[yellow]○[/yellow]", @@ -1890,14 +1890,14 @@ def setup_help_command() -> int: print_heading("3. Review and publish") console.print( - _command_line("ucode setup show", "The draft, and the payload `publish` sends", width=width) + _command_line("ug setup show", "The draft, and the payload `publish` sends", width=width) ) - console.print(_command_line("ucode publish", "Publish it to the workspace", width=width)) + console.print(_command_line("ug publish", "Publish it to the workspace", width=width)) print_heading("Also") console.print( _command_line( - "ucode setup --from-file ", + "ug setup --from-file ", "Load a hand-written manifest instead of prompting", width=width, ) @@ -1906,24 +1906,24 @@ def setup_help_command() -> int: f"The draft lives in ~/.ucode/managed-state.json (workspace: {workspace or 'none'})." ) print_note( - "Re-running `ucode setup` keeps the sections in step 2; to drop one, edit the draft and " - "reload it with `ucode setup --from-file`." + "Re-running `ug setup` keeps the sections in step 2; to drop one, edit the draft and " + "reload it with `ug setup --from-file`." ) return 0 def show_command() -> int: - """Print the authored manifest and the proto-JSON `ucode publish` would publish.""" - # Fall back to the workspace the on-disk file was authored for, so `ucode setup --show` still - # works before `ucode configure` has put a workspace in local state. + """Print the authored manifest and the proto-JSON `ug publish` would publish.""" + # Fall back to the workspace the on-disk file was authored for, so `ug setup --show` still + # works before `ug configure` has put a workspace in local state. workspace = load_state().get("workspace") or managed_state_workspace() manifest = load_managed_state(workspace) if manifest is None: - print_note("No managed config has been authored yet. Run `ucode setup` to create one.") + print_note("No managed config has been authored yet. Run `ug setup` to create one.") return 0 _render_summary(workspace or "unknown", manifest) console.print() - print_heading("Payload for `ucode publish`") + print_heading("Payload for `ug publish`") console.print(json.dumps(serialize_managed_config(manifest), indent=2)) return 0 @@ -1942,8 +1942,8 @@ def _explain_publish_failure(reason: str) -> str: ) if "already_exists" in lowered: return ( - "This workspace already has a managed config, but ucode couldn't read it to update in " - "place. Run `ucode publish` again — if it keeps failing, the existing config may need to " + "This workspace already has a managed config, but ug couldn't read it to update in " + "place. Run `ug publish` again — if it keeps failing, the existing config may need to " "be deleted by hand." ) if "invalid_parameter_value" in lowered: @@ -1956,7 +1956,7 @@ def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> """``state`` plus the full Claude listing, for validating a manifest against the workspace. ``state["claude_models"]`` holds only the newest id per family (the launch path pins one model - per family alias), but `ucode setup` deliberately offers the older versions too — pinning + per family alias), but `ug setup` deliberately offers the older versions too — pinning ``default_opus_model`` to a known-good ``claude-opus-4-8`` is a normal thing for an admin to want. Validating against ``claude_models`` alone therefore rejected a model the wizard itself had just offered: @@ -1964,7 +1964,7 @@ def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> claude: model 'system.ai.claude-opus-4-8' is not available on this workspace. The wizard stashes the full listing on ``state["all_claude_models"]`` mid-run, but that is never - persisted — `setup` saves the manifest, not the state — so a separate `ucode publish` process + persisted — `setup` saves the manifest, not the state — so a separate `ug publish` process starts from a fresh ``load_state()`` without it. Re-fetching here makes the check independent of what the wizard happened to leave behind, which also covers a hand-edited or ``--from-file`` manifest authored on another machine. @@ -1990,7 +1990,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: """Publish a managed config to the workspace. With no ``file_path`` the locally authored manifest is published; with one, the config file - (produced by ``ucode export``) is published instead. Both routes are validated against the + (produced by ``ug export``) is published instead. Both routes are validated against the configured workspace and canonicalized before anything is sent. Updates the existing config in place when there is one, rather than deleting and recreating it: a failed recreate would leave the workspace with no managed config at all, and every developer would silently fall back to @@ -2003,7 +2003,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: from ucode.cli import _prompt_for_configuration from ucode.managed_publish import load_publish_payload, parse_publish_payload - print_section("ucode publish") + print_section("ug publish") state = load_state() workspace = state.get("workspace") @@ -2026,9 +2026,9 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: for error in errors: print_note(error) if file_path is None: - print_note("Re-run `ucode setup` to fix it, or edit ~/.ucode/managed-state.json.") + print_note("Re-run `ug setup` to fix it, or edit ~/.ucode/managed-state.json.") else: - print_note("Fix the config file and re-run `ucode publish -f`.") + print_note("Fix the config file and re-run `ug publish -f`.") return 1 token = get_databricks_token(workspace, profile) @@ -2052,7 +2052,7 @@ def publish_command(*, file_path: str | None = None, yes: bool = False) -> int: if existing is not None and not isinstance(existing_name, str): raise RuntimeError( "This workspace has a managed config but the API didn't return its resource name, so " - "ucode can't update it in place. Delete it in the workspace and re-run `ucode publish`." + "ug can't update it in place. Delete it in the workspace and re-run `ug publish`." ) console.print() diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index e27401fd..7c871d54 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -31,7 +31,7 @@ def format_switch_message(model: str, reason: str | None) -> str: """Format the first-prompt routed-model notice.""" lines = [ - "Using Unity Gateway Smart Router.", + "Using ug Smart Router.", f"Selected Model : {model}", *([f"Reason : {reason}"] if reason else []), SUBAGENT_ROUTING_DISCLAIMER, @@ -42,7 +42,7 @@ def format_switch_message(model: str, reason: str | None) -> str: def format_subagent_message(model: str, reason: str | None) -> str: """Format a routed-subagent notice without the first-prompt disclaimer.""" lines = [ - "Using Unity Gateway Smart Router - Subagent", + "Using ug Smart Router - Subagent", f"Selected Model : {model}", *([f"Reason : {reason}"] if reason else []), ] diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..b048cd5f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -7,8 +7,11 @@ import json import os import re +import subprocess import time -from unittest.mock import MagicMock, patch +import tomllib +from pathlib import Path +from unittest.mock import MagicMock, call, patch import pytest from typer.testing import CliRunner @@ -81,6 +84,16 @@ def test_help_lists_all_agent_subcommands(self): for tool in TOOLS: assert tool in result.output + @pytest.mark.parametrize("prog_name", ["ug", "ucode"]) + def test_help_uses_invoked_name_and_names_ucode_as_an_alias(self, prog_name): + result = runner.invoke(app, ["--help"], prog_name=prog_name) + output = _strip_ansi(result.output) + + assert result.exit_code == 0 + assert f"Usage: {prog_name}" in output + assert "primary command is `ug`" in output + assert "`ucode` remains supported as an alias" in output + @pytest.mark.parametrize("tool", TOOLS) def test_subcommand_help(self, tool): result = runner.invoke(app, [tool, "--help"]) @@ -99,6 +112,52 @@ def test_configure_help_lists_agents_flag(self): assert "--workspaces" in output +class TestProjectScripts: + def test_ug_and_ucode_are_equivalent_entry_points(self): + scripts = tomllib.loads((Path(__file__).parent.parent / "pyproject.toml").read_text())[ + "project" + ]["scripts"] + + assert scripts["ug"] == "ucode.cli:main" + assert scripts["ucode"] == "ucode.cli:main" + + +class TestUpgrade: + def test_uninstalls_legacy_distribution_then_force_installs(self): + git_url = "git+https://github.com/databricks/ucode" + with patch("subprocess.run") as mock_run: + mock_run.side_effect = [ + subprocess.CompletedProcess([], 1), + subprocess.CompletedProcess([], 0), + ] + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 0, result.output + assert mock_run.call_args_list == [ + call(["uv", "tool", "uninstall", "ucode"], check=False), + call(["uv", "tool", "install", "--force", git_url], check=True), + ] + assert "ug upgraded" in result.output + + def test_install_failure_is_actionable(self): + with patch("subprocess.run") as mock_run: + mock_run.side_effect = [ + subprocess.CompletedProcess([], 0), + subprocess.CalledProcessError(7, ["uv", "tool", "install"]), + ] + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 1 + assert "Upgrade failed (exit code 7)" in result.output + + def test_missing_uv_is_actionable(self): + with patch("subprocess.run", side_effect=FileNotFoundError): + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 1 + assert "uv" in result.output.lower() + + class TestVersion: @pytest.mark.parametrize("flag", ["--version", "-V"]) def test_prints_version_and_exits(self, flag): @@ -1380,12 +1439,12 @@ def test_skipped_when_already_configured(self): @pytest.mark.parametrize( ("tool", "expected"), [ - ("claude", "Launching Claude Code with Unity Gateway"), - ("codex", "Launching Codex with Unity Gateway"), - ("gemini", "Launching Gemini CLI with Unity Gateway"), - ("opencode", "Launching OpenCode with Unity Gateway"), - ("copilot", "Launching GitHub Copilot CLI with Unity Gateway"), - ("pi", "Launching Pi with Unity Gateway"), + ("claude", "Launching Claude Code with ug"), + ("codex", "Launching Codex with ug"), + ("gemini", "Launching Gemini CLI with ug"), + ("opencode", "Launching OpenCode with ug"), + ("copilot", "Launching GitHub Copilot CLI with ug"), + ("pi", "Launching Pi with ug"), ], ) def test_launch_title(tool, expected): @@ -1842,22 +1901,6 @@ def test_agent_flag_normalizes_alias(self): assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with("claude") - def test_upgrade_runs_uv_tool_install(self): - with patch("subprocess.run") as mock_run: - result = runner.invoke(app, ["upgrade"]) - assert result.exit_code == 0, result.output - mock_run.assert_called_once() - cmd = mock_run.call_args[0][0] - assert cmd[:3] == ["uv", "tool", "install"] - assert "--reinstall" in cmd - assert any("github.com/databricks/ucode" in s for s in cmd) - - def test_upgrade_handles_uv_missing(self): - with patch("subprocess.run", side_effect=FileNotFoundError): - result = runner.invoke(app, ["upgrade"]) - assert result.exit_code != 0 - assert "uv" in result.output.lower() - def test_agent_flag_rejects_unknown(self): with ( patch("ucode.cli.install_databricks_cli"), @@ -3141,7 +3184,7 @@ def test_non_admin_with_config_confirms_and_exits(self, monkeypatch, capsys): assert exc.value.exit_code == 0 out = capsys.readouterr().out assert "you're all set" in out - assert "Run `ucode`" in out + assert "Run `ug`" in out def test_fetches_the_config_rather_than_reading_a_cold_cache(self, monkeypatch): # The gap this guards: on a fresh machine the local cache is empty until the first launch, @@ -3182,7 +3225,7 @@ def test_admin_with_config_runs_setup(self, monkeypatch): { "workspace": "https://w", "profile": None, - "command_label": "Configure Unity Gateway", + "command_label": "Configure ug", "token": "tok", } ] @@ -3289,7 +3332,7 @@ def test_admin_with_no_config_runs_setup_in_place(self, monkeypatch): { "workspace": "https://w", "profile": None, - "command_label": "Configure Unity Gateway", + "command_label": "Configure ug", "token": "tok", } ] @@ -3536,7 +3579,7 @@ def test_admin_without_a_config_is_pointed_at_setup(self, monkeypatch): result, launched = self._run(monkeypatch, managed=None, is_admin=True) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" in result.output + assert "ug setup" in result.output def test_non_admin_without_a_config_is_told_to_ask(self, monkeypatch): result, launched = self._run(monkeypatch, managed=None, is_admin=False) @@ -3550,7 +3593,7 @@ def test_admin_without_a_config_sees_no_setup_when_feature_disabled(self, monkey ) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" not in result.output + assert "ug setup" not in result.output def test_non_admin_without_a_config_sees_no_setup_when_feature_disabled(self, monkeypatch): result, launched = self._run( @@ -3558,7 +3601,7 @@ def test_non_admin_without_a_config_sees_no_setup_when_feature_disabled(self, mo ) assert result.exit_code == 0, result.output assert launched == [] - assert "ucode setup" not in result.output + assert "ug setup" not in result.output def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 8a41e2a6..8a51429d 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -163,7 +163,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): "gpt-5.5", "Review needs deeper reasoning." ) assert output["systemMessage"] == expected_message - assert "Using Unity Gateway Smart Router - Subagent" in expected_message + assert "Using ug Smart Router - Subagent" in expected_message assert codex_routing.routing.SUBAGENT_ROUTING_DISCLAIMER not in expected_message assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 1d618e5e..79720a76 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -40,7 +40,7 @@ def test_smart_routing_switch_message_is_boxed(): assert message == ( "┌───────────────────────────────────────────────────────────────────────────┐\n" - "│ Using Unity Gateway Smart Router. │\n" + "│ Using ug Smart Router. │\n" "│ Selected Model : model-x │\n" "│ Reason : Because X. │\n" "│ Spawned subagents are routed independently based on their own complexity. │\n" diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 013cdc9d..013890d8 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -1,4 +1,4 @@ -"""Tests for the interactive `ucode setup` flow and its CLI wiring. +"""Tests for the interactive `ug setup` flow and its CLI wiring. The wizard is mostly orchestration, so these focus on the parts where it can silently produce a wrong manifest: reading tracing/MCP/skills back out of ``state.json``, classifying MCP URLs into @@ -264,7 +264,7 @@ def test_feature_disabled_blocks_setup_with_an_actionable_error(self): assert not select.called message = str(exc_info.value) assert message == wizard.CODING_AGENT_CONFIGS_DISABLED_MESSAGE - assert "`ucode configure`" in message + assert "`ug configure`" in message # The raw 404 / JSON body must not leak into the message. assert "404" not in message assert "FEATURE_DISABLED" not in message @@ -287,7 +287,7 @@ def test_choosing_create_continues_authoring(self): def test_warning_does_not_itemize_the_existing_config(self): # The warning is the same whatever the config holds: an inventory doesn't change what the - # admin should do, and `ucode setup show` prints the real thing for comparison. + # admin should do, and `ug setup show` prints the real thing for comparison. with ( patch.object(wizard, "get_managed_config", return_value=(self.RICH_CONFIG, None)), patch.object(wizard, "prompt_for_selection", return_value="create"), @@ -372,22 +372,22 @@ def test_cancelling_the_picker_aborts(self): class TestStepBanner: - """The step headers brand themselves to the invoking command, so a `ucode configure` run - doesn't show `ucode setup` headers.""" + """The step headers brand themselves to the invoking command, so a `ug configure` run + doesn't show `ug setup` headers.""" def test_defaults_to_ucode_setup(self): with patch.object(wizard, "print_section") as section: wizard._step_banner(1, "Agents") - assert section.call_args.args[0].startswith("ucode setup · step 1 of ") + assert section.call_args.args[0].startswith("ug setup · step 1 of ") def test_uses_the_command_label_when_given(self): with patch.object(wizard, "print_section") as section: - wizard._step_banner(2, "Models", "ucode configure") - assert section.call_args.args[0].startswith("ucode configure · step 2 of ") + wizard._step_banner(2, "Models", "ug configure") + assert section.call_args.args[0].startswith("ug configure · step 2 of ") class TestSetupCommandToken: - """A caller (e.g. `ucode configure`) can hand setup a token so its admin gate uses the same + """A caller (e.g. `ug configure`) can hand setup a token so its admin gate uses the same identity as the routing decision, instead of fetching a second time.""" def test_reuses_a_passed_token_and_skips_a_second_fetch(self): @@ -1485,7 +1485,7 @@ def test_malformed_targets_yield_nothing(self): class TestBudgetPolicy: def test_no_up_front_gate(self): - # Running `ucode setup spend-tiers` is the consent, so the flow asks no "set up a policy?" + # Running `ug setup spend-tiers` is the consent, so the flow asks no "set up a policy?" # question — it goes straight to listing budgets. (The only yes/no it asks is "add another # tier?", after a tier is built.) with ( @@ -2132,7 +2132,7 @@ def fake_sel(prompt, options, **kwargs): assert any("model" in p for p in searchable_prompts), searchable_prompts -# A minimal authored manifest (agents + models only), the shape `ucode setup` now writes. +# A minimal authored manifest (agents + models only), the shape `ug setup` now writes. AGENTS_ONLY = { "default_agent": "claude", "enabled_agents": {"claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}}, @@ -2199,21 +2199,21 @@ def test_marks_configured_and_unconfigured_sections(self, capsys): manifest = {**AGENTS_ONLY, "skills": {"names": ["main.default"]}} wizard._print_next_steps(manifest) out = capsys.readouterr().out - assert "ucode setup mcps" in out - assert "ucode setup skills" in out - assert "ucode setup spend-tiers" in out - assert "ucode publish" in out + assert "ug setup mcps" in out + assert "ug setup skills" in out + assert "ug setup spend-tiers" in out + assert "ug publish" in out def test_dry_run_says_nothing_was_saved(self, capsys, monkeypatch): monkeypatch.setattr(config_io_mod, "_dry_run", True) wizard._print_next_steps(AGENTS_ONLY) out = capsys.readouterr().out assert "Dry run" in out - assert "ucode publish" not in out + assert "ug publish" not in out class TestSectionCommands: - """The `ucode setup mcps` / `skills` / `spend-tiers` section commands.""" + """The `ug setup mcps` / `skills` / `spend-tiers` section commands.""" @staticmethod def _admin(**overrides): @@ -2242,13 +2242,13 @@ def _run(self, fn, *, admin_overrides=None, **patches): def test_mcp_requires_an_authored_config(self): # No manifest on disk → the command can't edit a section that doesn't exist. - with pytest.raises(RuntimeError, match="ucode setup"): + with pytest.raises(RuntimeError, match="ug setup"): self._run(wizard.setup_mcp_command) def test_mcp_requires_enabled_agents(self): # A launch stores `{}` to mean "no managed config"; that must not count as authored. managed_config_mod.save_managed_state(WORKSPACE, {}) - with pytest.raises(RuntimeError, match="ucode setup"): + with pytest.raises(RuntimeError, match="ug setup"): self._run(wizard.setup_mcp_command) def test_mcp_writes_only_its_section(self): @@ -2281,7 +2281,7 @@ def test_mcp_cancel_is_a_no_op(self): assert not save.called def test_mcp_carries_forward_preregistered_servers(self): - # An admin who ran `ucode configure mcp` first arrives with those servers already registered, + # An admin who ran `ug configure mcp` first arrives with those servers already registered, # so the picker leaves local state unchanged (before == after). The manifest doesn't carry them # yet, so `setup mcps` must still save them rather than report "no changes" and drop them. managed_config_mod.save_managed_state(WORKSPACE, AGENTS_ONLY) @@ -2373,12 +2373,12 @@ def test_lists_every_setup_command(self, capsys): wizard.setup_help_command() out = capsys.readouterr().out for command in ( - "ucode setup", - "ucode setup mcps", - "ucode setup skills", - "ucode setup spend-tiers", - "ucode setup show", - "ucode publish", + "ug setup", + "ug setup mcps", + "ug setup skills", + "ug setup spend-tiers", + "ug setup show", + "ug publish", ): assert command in out @@ -2465,7 +2465,7 @@ def _config_file(tmp_path, manifest, *, workspace=WORKSPACE, spec_version=1, **e def test_unauthored_config_is_an_actionable_error(self): with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): - with pytest.raises(RuntimeError, match="ucode setup"): + with pytest.raises(RuntimeError, match="ug setup"): wizard.publish_command() def test_creates_when_no_config_exists(self): @@ -2478,7 +2478,7 @@ def fake_create(workspace, token, payload): assert self._run(create_coding_agent_config=fake_create) == 0 assert created["workspace"] == WORKSPACE - # What goes over the wire is proto-JSON, not ucode's manifest shape. + # What goes over the wire is proto-JSON, not ug's manifest shape. assert created["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" def test_updates_in_place_when_a_config_exists(self): @@ -2784,7 +2784,7 @@ def test_feature_disabled_uses_the_shared_message(self): 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED","message":"..."}' ) assert message == wizard.CODING_AGENT_CONFIGS_DISABLED_MESSAGE - assert "`ucode configure`" in message + assert "`ug configure`" in message def test_permission_denied_says_admin_is_required(self): message = wizard._explain_publish_failure( @@ -2954,7 +2954,7 @@ def test_setup_skills_location_is_parsed_to_a_list(self): assert fn.call_args.args[0] == ["main.a", "main.b"] def test_setup_help_needs_no_auth(self): - # `ucode setup help` reads the local draft only — it must not shell out to install the CLI. + # `ug setup help` reads the local draft only — it must not shell out to install the CLI. with ( patch("ucode.cli.install_databricks_cli") as install, patch("ucode.cli.setup_help_command", return_value=0) as fn, @@ -2967,13 +2967,11 @@ def test_setup_help_needs_no_auth(self): def test_section_command_runtime_error_exits_1(self): with ( patch("ucode.cli.install_databricks_cli"), - patch( - "ucode.cli.setup_mcp_command", side_effect=RuntimeError("run `ucode setup` first") - ), + patch("ucode.cli.setup_mcp_command", side_effect=RuntimeError("run `ug setup` first")), ): result = runner.invoke(app, ["setup", "mcps"]) assert result.exit_code == 1 - assert "ucode setup" in _out(result) + assert "ug setup" in _out(result) def test_section_command_interrupt_exits_130(self): with ( From 2f5ecf8e8c689f9c66c28926955d014435102b14 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 13:39:05 -0400 Subject: [PATCH 2/6] Keep Unity Gateway smart router branding --- src/ucode/smart_routing/routing.py | 4 ++-- tests/test_codex_routing.py | 2 +- tests/test_codex_smart_routing_v2.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 7c871d54..e27401fd 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -31,7 +31,7 @@ def format_switch_message(model: str, reason: str | None) -> str: """Format the first-prompt routed-model notice.""" lines = [ - "Using ug Smart Router.", + "Using Unity Gateway Smart Router.", f"Selected Model : {model}", *([f"Reason : {reason}"] if reason else []), SUBAGENT_ROUTING_DISCLAIMER, @@ -42,7 +42,7 @@ def format_switch_message(model: str, reason: str | None) -> str: def format_subagent_message(model: str, reason: str | None) -> str: """Format a routed-subagent notice without the first-prompt disclaimer.""" lines = [ - "Using ug Smart Router - Subagent", + "Using Unity Gateway Smart Router - Subagent", f"Selected Model : {model}", *([f"Reason : {reason}"] if reason else []), ] diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 8a51429d..8a41e2a6 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -163,7 +163,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): "gpt-5.5", "Review needs deeper reasoning." ) assert output["systemMessage"] == expected_message - assert "Using ug Smart Router - Subagent" in expected_message + assert "Using Unity Gateway Smart Router - Subagent" in expected_message assert codex_routing.routing.SUBAGENT_ROUTING_DISCLAIMER not in expected_message assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 79720a76..1d618e5e 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -40,7 +40,7 @@ def test_smart_routing_switch_message_is_boxed(): assert message == ( "┌───────────────────────────────────────────────────────────────────────────┐\n" - "│ Using ug Smart Router. │\n" + "│ Using Unity Gateway Smart Router. │\n" "│ Selected Model : model-x │\n" "│ Reason : Because X. │\n" "│ Spawned subagents are routed independently based on their own complexity. │\n" From bed2212c588d245cab0311525d105bc75ca56c41 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 13:46:09 -0400 Subject: [PATCH 3/6] Preserve Unity Gateway launch branding --- src/ucode/cli.py | 4 ++-- tests/test_cli.py | 24 ++++++++++++++++++------ 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 21ef7144..40c7e07f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1733,7 +1733,7 @@ def _fetch_budget_recommendation(state: dict, managed: dict | None) -> dict | No def _launch_title(tool: str) -> str: - return f"Launching {TOOL_SPECS[tool]['display']} with ug" + return f"Launching {TOOL_SPECS[tool]['display']} with Unity Gateway" def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = None) -> None: @@ -2526,7 +2526,7 @@ def cursor_cmd(ctx: typer.Context) -> None: f"`{cursor.CURSOR_BINARY}` was not found on PATH. Install Cursor Agent " "(https://cursor.com/cli), then re-run `ug cursor`." ) - print_section("ug with Cursor") + print_section("Unity Gateway with Cursor") print_note( "Cursor runs models on your Cursor account; its Databricks MCP servers " "authenticate through `ug mcp-proxy`." diff --git a/tests/test_cli.py b/tests/test_cli.py index b048cd5f..2b61a390 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1439,12 +1439,12 @@ def test_skipped_when_already_configured(self): @pytest.mark.parametrize( ("tool", "expected"), [ - ("claude", "Launching Claude Code with ug"), - ("codex", "Launching Codex with ug"), - ("gemini", "Launching Gemini CLI with ug"), - ("opencode", "Launching OpenCode with ug"), - ("copilot", "Launching GitHub Copilot CLI with ug"), - ("pi", "Launching Pi with ug"), + ("claude", "Launching Claude Code with Unity Gateway"), + ("codex", "Launching Codex with Unity Gateway"), + ("gemini", "Launching Gemini CLI with Unity Gateway"), + ("opencode", "Launching OpenCode with Unity Gateway"), + ("copilot", "Launching GitHub Copilot CLI with Unity Gateway"), + ("pi", "Launching Pi with Unity Gateway"), ], ) def test_launch_title(tool, expected): @@ -1453,6 +1453,18 @@ def test_launch_title(tool, expected): assert _launch_title(tool) == expected +def test_cursor_launch_uses_unity_gateway_branding(): + with ( + patch("ucode.cli.shutil.which", return_value="/usr/local/bin/cursor-agent"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.agents.cursor.launch"), + ): + result = runner.invoke(app, ["cursor"]) + + assert result.exit_code == 0, result.output + assert "Unity Gateway with Cursor" in result.output + + class TestCachedConfigPredicate: @staticmethod def _kwargs(**overrides): From 1bee50725fb30f5a588216779d8a3e59cb432358 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 13:51:16 -0400 Subject: [PATCH 4/6] Smoke test ug and ucode entry points --- tests/test_entry_points.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 tests/test_entry_points.py diff --git a/tests/test_entry_points.py b/tests/test_entry_points.py new file mode 100644 index 00000000..b20e1331 --- /dev/null +++ b/tests/test_entry_points.py @@ -0,0 +1,32 @@ +"""Smoke tests for the installed ``ug`` and ``ucode`` console scripts.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("command", ["ug", "ucode"]) +def test_installed_console_script_runs_with_its_invoked_name(command: str) -> None: + """Both scripts installed by ``uv run pytest`` execute the same CLI successfully.""" + bin_dir = Path(sys.executable).parent + script = shutil.which(command, path=str(bin_dir)) + assert script is not None, f"{command} was not installed in {bin_dir}" + + result = subprocess.run( + [script, "--help"], + cwd=Path(__file__).parent.parent, + env={**os.environ, "NO_COLOR": "1"}, + text=True, + capture_output=True, + check=False, + ) + + output = result.stdout + result.stderr + assert result.returncode == 0, output + assert f"Usage: {command} " in output From a2590b46aa39851a771a3fbfd289e022d6b13d33 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 13:59:03 -0400 Subject: [PATCH 5/6] Make entry point smoke tests ANSI-safe --- tests/test_entry_points.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_entry_points.py b/tests/test_entry_points.py index b20e1331..bb33a86c 100644 --- a/tests/test_entry_points.py +++ b/tests/test_entry_points.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import re import shutil import subprocess import sys @@ -10,6 +11,8 @@ import pytest +_ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") + @pytest.mark.parametrize("command", ["ug", "ucode"]) def test_installed_console_script_runs_with_its_invoked_name(command: str) -> None: @@ -29,4 +32,4 @@ def test_installed_console_script_runs_with_its_invoked_name(command: str) -> No output = result.stdout + result.stderr assert result.returncode == 0, output - assert f"Usage: {command} " in output + assert f"Usage: {command} " in _ANSI_RE.sub("", output) From 5f0732187b9b91bd5fbc7d3718b456234fa87726 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 2 Sep 2026 16:18:52 -0400 Subject: [PATCH 6/6] Make ug upgrade distribution-aware --- src/ucode/cli.py | 124 +++++++++++++++++++++++++++---- tests/test_cli.py | 186 ++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 281 insertions(+), 29 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d12f3f7c..eff2979c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -5,8 +5,10 @@ import os import shutil +import subprocess from collections.abc import Iterator from contextlib import contextmanager +from importlib import metadata from typing import Annotated import typer @@ -3250,31 +3252,125 @@ def usage_cmd( @app.command("upgrade") def upgrade_cmd() -> None: """Upgrade ug to the latest version from GitHub.""" - import subprocess - + legacy_distribution = "ucode" + current_distribution = "unity-gateway" git_url = "git+https://github.com/databricks/ucode" + installed_distribution = _installed_cli_distribution() + migrated = False + legacy_removed = False + print_section("Upgrade") print_kv("Source", git_url) + print_kv("Installed distribution", installed_distribution) try: - # uv tracks tools by distribution name, so it cannot upgrade the old - # `ucode` distribution across the future rename to `unity-gateway`. - # Removing the legacy registration is best-effort: users who are - # already migrated will not have a tool registered under that name. - subprocess.run( - ["uv", "tool", "uninstall", "ucode"], + result = subprocess.run( + ["uv", "tool", "upgrade", installed_distribution], check=False, + capture_output=True, + text=True, ) - subprocess.run( - ["uv", "tool", "install", "--force", git_url], - check=True, - ) + if result.returncode != 0: + if installed_distribution == legacy_distribution and _is_distribution_cutover(result): + print_note( + "The package is now distributed as `unity-gateway`; migrating this installation." + ) + subprocess.run( + ["uv", "tool", "uninstall", legacy_distribution], + check=True, + ) + legacy_removed = True + subprocess.run( + ["uv", "tool", "install", "--force", git_url], + check=True, + ) + migrated = True + installed_distribution = current_distribution + else: + detail = _upgrade_failure_detail(result) + print_err( + f"Upgrade failed (exit code {result.returncode})" + f"{f': {detail}' if detail else '.'}" + ) + print_note("The existing installation was left unchanged.") + raise typer.Exit(1) + + _verify_upgraded_commands() except FileNotFoundError: print_err("`uv` was not found on PATH. Install uv to upgrade ug.") raise typer.Exit(1) from None except subprocess.CalledProcessError as exc: - print_err(f"Upgrade failed (exit code {exc.returncode}).") + if legacy_removed: + print_err( + "The legacy `ucode` tool was removed, but installing `unity-gateway` failed " + f"(exit code {exc.returncode})." + ) + print_note(f"Recover by running `uv tool install --force {git_url}`.") + else: + print_err(f"Upgrade failed (exit code {exc.returncode}); `ucode` was not removed.") + raise typer.Exit(1) from None + except RuntimeError as exc: + print_err(str(exc)) raise typer.Exit(1) from None - print_success("ug upgraded; `ucode` remains available as an alias") + + if migrated: + print_success("Migrated to `unity-gateway`; both `ug` and `ucode` are working") + else: + print_success(f"{installed_distribution} upgraded; both `ug` and `ucode` are working") + + +def _installed_cli_distribution() -> str: + """Return the uv tool identity, preferring the post-cutover distribution.""" + for distribution_name in ("unity-gateway", "ucode"): + try: + metadata.version(distribution_name) + except metadata.PackageNotFoundError: + continue + return distribution_name + # Source checkouts and unusual installers may expose neither distribution. + # The repository is still named ucode when this bridge ships, so use the + # non-destructive legacy upgrade path and let uv report an actionable error. + return "ucode" + + +def _is_distribution_cutover(result: subprocess.CompletedProcess[str]) -> bool: + """Recognize uv's specific error when source metadata changes distribution name.""" + output = f"{result.stdout or ''}\n{result.stderr or ''}".lower() + return all( + marker in output + for marker in ( + "metadata name", + "unity-gateway", + "does not match given name", + "ucode", + ) + ) + + +def _upgrade_failure_detail(result: subprocess.CompletedProcess[str]) -> str: + return (result.stderr or result.stdout or "").strip() + + +def _verify_upgraded_commands() -> None: + """Ensure both compatibility entry points were installed and can start.""" + for command in ("ug", "ucode"): + executable = shutil.which(command) + if executable is None: + raise RuntimeError( + f"Upgrade completed, but `{command}` is not available on PATH. " + "Reinstall Unity Gateway and ensure the uv tool bin directory is on PATH." + ) + result = subprocess.run( + [executable, "--version"], + 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"{f': {detail}' if detail else '.'}" + ) def main() -> None: diff --git a/tests/test_cli.py b/tests/test_cli.py index 956f647b..b383d4e8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -10,6 +10,7 @@ import subprocess import time import tomllib +from importlib import metadata from pathlib import Path from unittest.mock import MagicMock, call, patch @@ -127,40 +128,195 @@ def test_ug_and_ucode_are_equivalent_entry_points(self): class TestUpgrade: - def test_uninstalls_legacy_distribution_then_force_installs(self): + @staticmethod + def _ok() -> subprocess.CompletedProcess[str]: + return subprocess.CompletedProcess([], 0, stdout="", stderr="") + + @staticmethod + def _which(command: str) -> str: + return f"/tools/{command}" + + def test_before_cutover_upgrades_ucode_normally_and_verifies_commands(self): + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("ucode.cli.shutil.which", side_effect=self._which), + patch("subprocess.run", side_effect=[self._ok(), self._ok(), self._ok()]) as run, + ): + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 0, result.output + assert run.call_args_list == [ + call( + ["uv", "tool", "upgrade", "ucode"], + check=False, + capture_output=True, + text=True, + ), + call( + ["/tools/ug", "--version"], + check=False, + capture_output=True, + text=True, + ), + call( + ["/tools/ucode", "--version"], + check=False, + capture_output=True, + text=True, + ), + ] + assert "ucode upgraded" in result.output + + def test_cutover_migrates_legacy_distribution_and_verifies_commands(self): git_url = "git+https://github.com/databricks/ucode" - with patch("subprocess.run") as mock_run: - mock_run.side_effect = [ - subprocess.CompletedProcess([], 1), - subprocess.CompletedProcess([], 0), - ] + rename_failure = subprocess.CompletedProcess( + [], + 1, + stdout="", + stderr=("Package metadata name `unity-gateway` does not match given name `ucode`"), + ) + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("ucode.cli.shutil.which", side_effect=self._which), + patch( + "subprocess.run", + side_effect=[ + rename_failure, + self._ok(), + self._ok(), + self._ok(), + self._ok(), + ], + ) as run, + ): result = runner.invoke(app, ["upgrade"]) assert result.exit_code == 0, result.output - assert mock_run.call_args_list == [ - call(["uv", "tool", "uninstall", "ucode"], check=False), + assert run.call_args_list == [ + call( + ["uv", "tool", "upgrade", "ucode"], + check=False, + capture_output=True, + text=True, + ), + call(["uv", "tool", "uninstall", "ucode"], check=True), call(["uv", "tool", "install", "--force", git_url], check=True), + call( + ["/tools/ug", "--version"], + check=False, + capture_output=True, + text=True, + ), + call( + ["/tools/ucode", "--version"], + check=False, + capture_output=True, + text=True, + ), ] - assert "ug upgraded" in result.output + assert "Migrated to `unity-gateway`" in result.output + + def test_after_cutover_upgrades_unity_gateway_normally_and_verifies_commands(self): + with ( + patch("ucode.cli._installed_cli_distribution", return_value="unity-gateway"), + patch("ucode.cli.shutil.which", side_effect=self._which), + patch("subprocess.run", side_effect=[self._ok(), self._ok(), self._ok()]) as run, + ): + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 0, result.output + assert run.call_args_list[0] == call( + ["uv", "tool", "upgrade", "unity-gateway"], + check=False, + capture_output=True, + text=True, + ) + assert all("uninstall" not in invocation.args[0] for invocation in run.call_args_list) + assert "unity-gateway upgraded" in result.output + + def test_unrelated_legacy_upgrade_failure_does_not_uninstall_ucode(self): + failure = subprocess.CompletedProcess( + [], 7, stdout="", stderr="Could not resolve host: github.com" + ) + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("subprocess.run", return_value=failure) as run, + ): + result = runner.invoke(app, ["upgrade"]) - def test_install_failure_is_actionable(self): - with patch("subprocess.run") as mock_run: - mock_run.side_effect = [ - subprocess.CompletedProcess([], 0), + assert result.exit_code == 1 + run.assert_called_once_with( + ["uv", "tool", "upgrade", "ucode"], + check=False, + capture_output=True, + text=True, + ) + assert "left unchanged" in result.output + + def test_cutover_install_failure_has_recovery_command(self): + rename_failure = subprocess.CompletedProcess( + [], + 1, + stdout="", + stderr=("Package metadata name `unity-gateway` does not match given name `ucode`"), + ) + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("subprocess.run") as run, + ): + run.side_effect = [ + rename_failure, + self._ok(), subprocess.CalledProcessError(7, ["uv", "tool", "install"]), ] result = runner.invoke(app, ["upgrade"]) assert result.exit_code == 1 - assert "Upgrade failed (exit code 7)" in result.output + assert "legacy `ucode` tool was removed" in result.output + assert "uv tool install --force git+https://github.com/databricks/ucode" in re.sub( + r"\s+", " ", result.output + ) + + def test_verification_failure_is_actionable(self): + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("ucode.cli.shutil.which", return_value=None), + patch("subprocess.run", return_value=self._ok()), + ): + result = runner.invoke(app, ["upgrade"]) + + assert result.exit_code == 1 + assert "`ug` is not available on PATH" in result.output def test_missing_uv_is_actionable(self): - with patch("subprocess.run", side_effect=FileNotFoundError): + with ( + patch("ucode.cli._installed_cli_distribution", return_value="ucode"), + patch("subprocess.run", side_effect=FileNotFoundError), + ): result = runner.invoke(app, ["upgrade"]) assert result.exit_code == 1 assert "uv" in result.output.lower() + def test_installed_distribution_prefers_unity_gateway(self): + with patch("ucode.cli.metadata.version", return_value="1.0.0") as package_version: + from ucode.cli import _installed_cli_distribution + + assert _installed_cli_distribution() == "unity-gateway" + + package_version.assert_called_once_with("unity-gateway") + + def test_installed_distribution_falls_back_to_ucode(self): + def package_version(distribution_name: str) -> str: + if distribution_name == "unity-gateway": + raise metadata.PackageNotFoundError + return "1.0.0" + + with patch("ucode.cli.metadata.version", side_effect=package_version): + from ucode.cli import _installed_cli_distribution + + assert _installed_cli_distribution() == "ucode" + class TestVersion: @pytest.mark.parametrize("flag", ["--version", "-V"])