Skip to content

Remove ug setup and ug publish; ug is developer-only #510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
david-siqi-liu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Remove ug setup and ug publish; ug is developer-only #510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 3 additions & 85 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,87 +246,15 @@ ucode skill add --location main.default --skills my-skill,other-skill
ucode skill add --skills main.default.my-skill,main.default.other-skill
```

### Managed config for a workspace (admins)

Author the coding config your developers pick up automatically, instead of asking each of them to
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
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
```

`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 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.
Claude Code is asked one model per family (opus/sonnet/haiku/fable), since it selects models by family
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. `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 `ug configure` step, which does configure this machine.

Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which
`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 `ug setup --from-file`.

```bash
# 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).
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.
ug publish

# Publish without the confirmation prompt (for CI).
ug publish --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 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 ug run.

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
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 `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 `ug
publish -f <file>` consumes.
directory must already exist) while stdout stays empty. The exported file is the portable
`CodingAgentConfig` proto-JSON the AI Gateway API accepts.

```bash
# Print the managed config as JSON.
Expand Down Expand Up @@ -383,16 +311,6 @@ The output looks like:
| `ug skill add --location main.default --mcp` | Add schemas to the skills MCP scope, keeping any already configured (additive; never replaces) |
| `ug skill add --location main.default` | Download a schema's skills to disk without removing existing downloads |
| `ug skill add --skills main.default.my-skill` | Download a named subset of skills (bare names need `--location`; fully-qualified names stand alone) |
| `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 <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 <file>` | 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 <agent>` launches.
Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to
Expand All @@ -413,7 +331,7 @@ 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 `ug setup` (admins) and refreshed from the workspace on launch |
| `~/.ucode/managed-state.json` | The managed config (published by an admin through the AI Gateway) 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. `ug revert` restores backups.
Expand Down
203 changes: 10 additions & 193 deletions src/ucode/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@
get_databricks_token,
install_databricks_cli,
is_model_provider_feature_unavailable,
is_workspace_admin,
list_profile_entries,
list_tool_provider_services,
normalize_workspace_url,
Expand Down Expand Up @@ -90,15 +89,6 @@
recommended_agent,
resolve_state,
)
from ucode.managed_wizard import (
publish_command,
setup_budget_policy_command,
setup_command,
setup_help_command,
setup_mcp_command,
setup_skills_command,
show_command,
)
from ucode.mcp import (
MCP_CLIENTS,
SKILLS_MCP_KIND,
Expand Down Expand Up @@ -1051,12 +1041,6 @@ def revert() -> int:
app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ug.")
skill_app = typer.Typer(add_completion=False, no_args_is_help=True)
app.add_typer(skill_app, name="skill", help="Databricks Skills for your coding tools.")
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 `ug setup help`.",
)


def _version_callback(value: bool) -> None:
Expand Down Expand Up @@ -2284,7 +2268,7 @@ def _launch_managed_default(
)
return
if not managed:
_print_no_managed_config_guidance(current, state.get("profile"))
_print_no_managed_config_guidance()
return
# The budget tier can move the org to a cheaper agent, so it outranks the config's
# default_agent. Fetched here and handed to _launch_tool so it is read once per launch.
Expand All @@ -2308,22 +2292,12 @@ def _launch_managed_default(
)


def _print_no_managed_config_guidance(workspace: str, profile: str | None) -> None:
"""Tell an admin how to publish a config, and everyone else who to ask."""
print_warning(
"No managed coding agent config was found for this workspace; using your local settings."
def _print_no_managed_config_guidance() -> None:
"""Point the developer at per-user configure when no managed config is published."""
print_note(
"No managed coding agent config is published for this workspace. Run `ug configure` to "
"set up your coding agents, then launch one with `ug <agent>` (for example `ug claude`)."
)
try:
token = get_databricks_token(workspace, profile)
except RuntimeError:
return
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 `ug setup`.")
else:
# None means the admin check itself failed; point at setup rather than a dead end.
print_note("Run `ug setup` to configure one for your workspace, then `ug publish`.")


@app.command(
Expand Down Expand Up @@ -3007,162 +2981,6 @@ def configure_tracing(
raise typer.Exit(130) from None


@setup_app.callback(invoke_without_command=True)
def setup(
ctx: typer.Context,
from_file: Annotated[
str | None,
typer.Option(
"--from-file",
help="Skip the interactive flow and load a hand-written managed config (JSON, in "
"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 `ug setup help`.
"""
if ctx.invoked_subcommand is not None:
return
# `typer.Exit` subclasses RuntimeError, so it must be raised outside the try — inside, the
# `except RuntimeError` below would swallow it and report the exit code as an error message.
try:
install_databricks_cli()
code = setup_command(from_file=from_file)
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
except KeyboardInterrupt:
print_err("Interrupted.")
raise typer.Exit(130) from None
if code:
raise typer.Exit(code)


@setup_app.command("mcps")
def setup_mcp_cmd() -> None:
"""Choose the MCP servers the managed config gives developers (admins only)."""
# Same `typer.Exit`/RuntimeError ordering trap as the `setup` callback above.
try:
install_databricks_cli()
code = setup_mcp_command()
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
except KeyboardInterrupt:
print_err("Interrupted.")
raise typer.Exit(130) from None
if code:
raise typer.Exit(code)


@setup_app.command("skills")
def setup_skills_cmd(
location: Annotated[
str | None,
typer.Option(
"--location",
help="Skill schemas to publish as `<catalog>.<schema>` (comma-separated for several). "
"Skips the prompt.",
),
] = None,
) -> None:
"""Choose the skills the managed config gives developers (admins only)."""
try:
install_databricks_cli()
# None means "prompt"; an explicit `--location` is parsed to the list to publish.
locations = None if location is None else _parse_skill_locations(location)
code = setup_skills_command(locations)
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
except KeyboardInterrupt:
print_err("Interrupted.")
raise typer.Exit(130) from None
if code:
raise typer.Exit(code)


@setup_app.command("spend-tiers")
def setup_budget_policy_cmd() -> None:
"""Route developers to cheaper agents as the workspace spends its budget (admins only)."""
try:
install_databricks_cli()
code = setup_budget_policy_command()
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
except KeyboardInterrupt:
print_err("Interrupted.")
raise typer.Exit(130) from None
if code:
raise typer.Exit(code)


@setup_app.command("help")
def setup_help_cmd() -> None:
"""Walk through the managed-config setup: every command, in order, and what's already done."""
# No auth and no CLI install: this reads the local draft only, so it works before `ucode
# configure` and on a machine without the Databricks CLI.
try:
code = setup_help_command()
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
if code:
raise typer.Exit(code)


@setup_app.command("show")
def setup_show_cmd() -> None:
"""Print the authored managed config and the payload `ug publish` would publish."""
try:
code = show_command()
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
if code:
raise typer.Exit(code)


@app.command("publish")
def publish_cmd(
file_path: Annotated[
str | None,
typer.Option(
"--file",
"-f",
help="Publish a config file exported with `ug export` instead of the locally "
"authored config. Its `workspace` must match the configured workspace.",
),
] = None,
yes: Annotated[
bool,
typer.Option("--yes", "-y", help="Publish without the confirmation prompt."),
] = False,
) -> None:
"""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: `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
# the try block or the handler below would report a successful exit as an error.
try:
install_databricks_cli()
code = publish_command(file_path=file_path, yes=yes)
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
except KeyboardInterrupt:
print_err("Interrupted.")
raise typer.Exit(130) from None
if code:
raise typer.Exit(code)


@app.command("export")
def export_cmd(
file_path: Annotated[
Expand All @@ -3177,11 +2995,10 @@ def export_cmd(
) -> None:
"""Export this workspace's managed coding-agent config as portable JSON.

Serializes the local managed config to the external `CodingAgentConfig` format that
`ug publish -f <path>` 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.
Serializes the local managed config to the external `CodingAgentConfig` proto-JSON format,
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.
"""
from ucode.managed_export import export_command

Expand Down
12 changes: 5 additions & 7 deletions src/ucode/managed_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,13 @@

- fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`),
- normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names,
- persisting it via :func:`save_managed_state` / :func:`load_managed_state` — the admin-write side
(``managed_setup`` / ``managed_wizard``) authors the manifest here, and the launch path pulls the
published copy back into the same file, and
- persisting it via :func:`save_managed_state` / :func:`load_managed_state`, which the launch path
uses to pull the published copy into the local file, and
- re-reading it on each launch, falling back to the persisted copy when the read fails.

There is deliberately one file, not a separate authored ``managed-settings.json``: the workspace is
the source of truth, so an authored draft and the pulled copy are the same shape and coexist in
``managed-state.json``. ``ucode setup`` authors the draft; ``ucode publish`` publishes it; a launch
then pulls the published copy back into the same file.
The workspace is the source of truth: an admin authors the ``CodingAgentConfig`` through the AI
Gateway API or UI, and each launch pulls the published copy into ``managed-state.json``. ``ucode``
only reads and applies it; it never authors or publishes.

:func:`refresh_managed_config` is the launch path's entry point. It is called before model discovery,
because the manifest decides whether that discovery is needed at all; the launch path then hands the
Expand Down
Loading
Loading