From e5037452381d54b35226f32d1a7f2c3d6f2e8465 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 16 Sep 2026 21:36:09 -0400 Subject: [PATCH] Polish managed configure output --- README.md | 8 ++ src/ucode/cli.py | 52 ++++---- .../integration/test_ug_configure_managed.py | 2 - tests/test_cli.py | 111 +++++++++++++++++- 4 files changed, 135 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index d64d972c..606e2097 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,14 @@ ug configure --workspace https://first.databricks.com `ug` logs into and saves state for that workspace. +Set `UG_WORKSPACE` to use the same workspace without repeating `--workspace`. An explicit +`--workspace` or `--profile` takes precedence over the environment variable. + +```bash +export UG_WORKSPACE=https://first.databricks.com +ug configure +``` + Alternatively, pass an existing Databricks CLI profile (from `~/.databrickscfg`) instead of a workspace URL — the profile's host supplies the workspace URL: ```bash diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 7a8749fa..81bb4774 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -23,7 +23,6 @@ configure_selected_tools, configure_single_tool, configure_tool, - configured_paths, ensure_bootstrap_dependencies, ensure_provider_state, explicit_model_arg_value, @@ -212,7 +211,7 @@ def _print_managed_summary( enabled = [t for t in (managed.get("enabled_agents") or {}) if t in TOOL_SPECS] if enabled: lines.append( - f"[bold]Enabled agents:[/bold] {', '.join(TOOL_SPECS[t]['display'] for t in enabled)}" + f"[bold]Coding Agents:[/bold] {', '.join(TOOL_SPECS[t]['display'] for t in enabled)}" ) if tool is not None: provider = managed_provider_service(managed, tool) @@ -239,9 +238,7 @@ def _print_managed_summary( else: lines.append("[bold]Skills:[/bold] [dim]none configured[/dim]") lines.extend(_policy_summary_lines(managed)) - console.print( - Panel("\n".join(lines), title="Workspace-managed config", style="green", expand=False) - ) + console.print(Panel("\n".join(lines), title="Configuration", style="green", expand=False)) def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None) -> None: @@ -265,27 +262,10 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None ) -def _announce_managed_config(managed: dict) -> None: - """Tell the developer, before configuring, that the admin's config drives this setup. - - Printed up front so the skipped agent selector reads as intended, not as a surprise.""" - print_success("A managed config is published for your workspace.") - enabled = [TOOL_SPECS[t]["display"] for t in managed_enabled_tools(managed) if t in TOOL_SPECS] - if enabled: - print_note(f"Applying it to the agents your admin enabled: {', '.join(enabled)}.") - - -def _print_configured_files(tool: str, state: dict) -> None: - """Name the config file(s) ug just wrote for ``tool``, so the developer sees what changed.""" - paths = configured_paths(tool, state) - if paths: - print_note(f"Updated {TOOL_SPECS[tool]['display']}: {', '.join(paths)}") - - def _summarize_managed_config(managed: dict, workspace: str) -> None: """Show the resulting managed setup once every enabled agent has been configured.""" _print_managed_summary(managed, {"workspace": workspace}, tool=None) - print_note("You're all set — run `ug` to launch with your managed settings.") + print_success("Configuration complete — launch with [bold cyan]ug[/bold cyan].") def _print_discovery_diagnostics(state: dict) -> None: @@ -584,7 +564,7 @@ def configure_shared_state( # search (claude only) still needs one Responses-capable model, so fetch # just that with a single call. if want_claude: - with spinner("Fetching web search model..."): + with spinner("Fetching available models..."): ws_models, _ = discover_codex_models(workspace, token) if ws_models: web_search_model = ws_models[0] @@ -810,10 +790,13 @@ def configure_workspace_command( # A published managed config means the admin dictates the setup: apply it to every enabled agent # now rather than prompting the developer to pick. managed, _ = refresh_managed_config(state) - if managed is not None: - _announce_managed_config(managed) - for tool_name in managed_enabled_tools(managed): + managed_tools = managed_enabled_tools(managed) if managed is not None else [] + if managed is not None and managed_tools: + configured_tools: list[str] = [] + for tool_name in managed_tools: if check_gateway_endpoint(state, tool_name): + if not install_tool_binary(tool_name, strict=False): + continue configured = configure_selected_tools( resolve_state(managed, state, tool_name), [tool_name], @@ -825,7 +808,12 @@ def configure_workspace_command( state["available_tools"] = configured.get("available_tools") or state.get( "available_tools" ) - _print_configured_files(tool_name, configured) + configured_tools.append(tool_name) + if not configured_tools: + raise RuntimeError( + "None of the coding agents enabled by your workspace configuration " + "are available on this workspace." + ) if not is_dry_run(): _configure_managed_mcp_servers(managed) _summarize_managed_config(managed, state["workspace"]) @@ -2215,7 +2203,6 @@ def _launch_tool( _note_recommended_agent(recommendation, tool) if managed is not None: state = resolve_state(managed, state, tool) - print_note("Applying your workspace's managed coding agent config...") unservable = managed_unservable_models(managed, tool) if unservable: print_warning( @@ -2357,8 +2344,6 @@ def _launch_tool( ctx.args = ["--model", relayed_forward_model, *ctx.args] forwarded_model = relayed_forward_model print_section(_launch_title(tool)) - if managed is not None: - print_kv("Config", "workspace-managed") if provider: print_kv("Provider", provider) if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and not provider: @@ -2882,7 +2867,8 @@ def configure( str | None, typer.Option( "--workspace", - help="Configure a single workspace without prompting.", + help="Configure a single workspace without prompting. " + "Defaults to the UG_WORKSPACE environment variable when set.", ), ] = None, workspaces: Annotated[ @@ -3029,6 +3015,8 @@ def configure( raise RuntimeError("Use either --profile or --profiles, not both.") workspace = workspace if workspace is not None else workspaces profile = profile if profile is not None else profiles + if workspace is None and profile is None: + workspace = os.environ.get("UG_WORKSPACE") or None if workspace is not None and profile is not None: raise RuntimeError("Use either --workspace or --profile, not both.") if use_pat and profile is None: diff --git a/tests/integration/test_ug_configure_managed.py b/tests/integration/test_ug_configure_managed.py index ef1b2c2b..d2fdf803 100644 --- a/tests/integration/test_ug_configure_managed.py +++ b/tests/integration/test_ug_configure_managed.py @@ -34,7 +34,6 @@ def test_ug_configure_managed_claude(live_session, workspace): session = live_session result = session.run("configure", "--workspace", workspace, "--skip-upgrade", timeout=240) assert "Select coding agents to configure:" not in result.stdout, result.stdout - assert "managed config is published" in result.stdout, result.stdout settings = json.loads((session.home / ".claude" / "ucode-settings.json").read_text()) assert settings.get("availableModels") == MANAGED_CLAUDE_MODELS, settings @@ -61,7 +60,6 @@ def test_ug_configure_managed_codex(live_session, workspace): session = live_session result = session.run("configure", "--workspace", workspace, "--skip-upgrade", timeout=240) assert "Select coding agents to configure:" not in result.stdout, result.stdout - assert "managed config is published" in result.stdout, result.stdout catalog = json.loads((session.home / ".ucode" / "codex-model-catalog.json").read_text()) listed = [ diff --git a/tests/test_cli.py b/tests/test_cli.py index 112e450b..cbb5afd2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1415,8 +1415,8 @@ def test_status_shows_managed_config_box_when_present_and_enabled(self, monkeypa result = runner.invoke(app, ["status"]) assert result.exit_code == 0, result.output - assert "Workspace-managed config" in result.output - assert "Enabled agents:" in result.output + assert "Configuration" in result.output + assert "Coding Agents:" in result.output assert "github-mcp" in result.output assert "debug-ci" in result.output @@ -1428,7 +1428,7 @@ def test_status_hides_managed_config_box_when_none_present(self, monkeypatch): result = runner.invoke(app, ["status"]) assert result.exit_code == 0, result.output - assert "Workspace-managed config" not in result.output + assert "Configuration" not in result.output class TestConfigureSkillsCommand: @@ -2805,6 +2805,12 @@ def test_managed_config_applies_all_enabled_and_skips_selection(self, monkeypatc lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False), ) monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True) + installed: list[str] = [] + monkeypatch.setattr( + cli_mod, + "install_tool_binary", + lambda tool, **kwargs: installed.append(tool) or True, + ) monkeypatch.setattr(cli_mod, "resolve_state", lambda managed, s, tool: s) monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) configured: list[str] = [] @@ -2820,8 +2826,57 @@ def test_managed_config_applies_all_enabled_and_skips_selection(self, monkeypatc ) assert cli_mod.configure_workspace_command(workspaces=[("https://w.com", None)]) == 0 + assert installed == ["claude", "codex"] assert configured == ["claude", "codex"] + def test_managed_config_fails_when_no_enabled_agent_is_available(self, monkeypatch): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr( + cli_mod, + "refresh_managed_config", + lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False), + ) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: False) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda *args, **kwargs: pytest.fail("must not configure unavailable agents"), + ) + + with pytest.raises(RuntimeError, match="None of the coding agents enabled"): + cli_mod.configure_workspace_command(workspaces=[("https://w.com", None)]) + + def test_budget_only_managed_config_uses_requested_agents(self, monkeypatch): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr( + cli_mod, + "refresh_managed_config", + lambda s: ({"budget_policy": {"policy_id": "budget"}}, False), + ) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: t == "claude") + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda managed: None) + configured: list[str] = [] + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda s, tools, **kwargs: configured.extend(tools) or s, + ) + + assert ( + cli_mod.configure_workspace_command( + selected_tools=["claude"], workspaces=[("https://w.com", None)] + ) + == 0 + ) + assert configured == ["claude"] + def test_managed_config_registers_mcp_servers_after_configuring_agents(self, monkeypatch): # The managed branch registers the config's MCP servers for the enabled agents once they are # configured — after the per-agent configure loop, so the agents' MCP configs already exist. @@ -2835,6 +2890,7 @@ def test_managed_config_registers_mcp_servers_after_configuring_agents(self, mon } monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s: (managed, False)) monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) monkeypatch.setattr(cli_mod, "resolve_state", lambda m, s, tool: s) monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) order: list[str] = [] @@ -2868,6 +2924,7 @@ def test_managed_configure_accumulates_available_tools_for_all_agents(self, monk lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False), ) monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) # Mirror production: resolve_state hands each iteration a fresh copy of `state`. monkeypatch.setattr(cli_mod, "resolve_state", lambda m, s, tool: dict(s)) monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) @@ -3120,6 +3177,52 @@ def test_profiles_flag_resolves_workspaces(self): workspaces=[("https://first.databricks.com", "DEFAULT")], ) + def test_ug_workspace_env_skips_workspace_prompt(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + ["configure"], + env={"UG_WORKSPACE": "https://env.databricks.com"}, + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + workspaces=[("https://env.databricks.com", None)], + ) + + def test_explicit_profile_overrides_ug_workspace_env(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.list_profile_entries", return_value=self.PROFILE_ENTRIES), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + ["configure", "--profile", "DEFAULT"], + env={"UG_WORKSPACE": "https://env.databricks.com"}, + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + workspaces=[("https://first.databricks.com", "DEFAULT")], + ) + + def test_explicit_workspace_overrides_ug_workspace_env(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, + ["configure", "--workspace", "https://explicit.databricks.com"], + env={"UG_WORKSPACE": "https://env.databricks.com"}, + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + workspaces=[("https://explicit.databricks.com", None)], + ) + def test_deprecated_profiles_alias_resolves_single_profile(self): # `--profiles` is a hidden alias of `--profile` and takes one profile. with ( @@ -4008,7 +4111,7 @@ def test_launch_banner_is_abridged_not_the_full_box(self, monkeypatch): assert "launching Claude Code as the default agent" in result.output assert "system.ai.opus" in result.output # The full box's per-config enumeration is left to `ucode status`. - assert "Enabled agents:" not in result.output + assert "Coding Agents:" not in result.output assert "system.ai.slack" not in result.output assert "main.default.my_skill" not in result.output