From 0b87364e68a9b9db758c7a37061a8e6993e3d3d3 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Mon, 24 Aug 2026 10:49:48 -0400 Subject: [PATCH] Add admin-only local managed config testing --- README.md | 28 +++- src/ucode/cli.py | 174 +++++++++++++++++------- src/ucode/managed_resolve.py | 9 +- src/ucode/managed_wizard.py | 26 ++-- tests/test_cli.py | 239 ++++++++++++++++++++++++++------- tests/test_managed_wizard.py | 13 +- tests/test_managed_workflow.py | 133 ++++++++++++++++++ 7 files changed, 499 insertions(+), 123 deletions(-) create mode 100644 tests/test_managed_workflow.py diff --git a/README.md b/README.md index 79d4b02d..00186ced 100644 --- a/README.md +++ b/README.md @@ -222,7 +222,9 @@ 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 apply # publish it to the workspace +ucode claude --local # optionally test the local draft +ucode codex --local # edit and test as often as needed +ucode apply # publish the latest draft to the workspace ``` `ucode setup` walks through the agents to enable and which one bare `ucode` launches, then per agent: @@ -238,10 +240,10 @@ that switches the default agent and model as the workspace burns through a budge 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. -Everything is written to `~/.ucode/managed-state.json` — the one local managed-config file — which -`ucode apply` publishes. Re-running `ucode 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`. +Everything is written to `~/.ucode/managed-state.json` — the editable local draft — which +`ucode apply` publishes. Re-running `ucode 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`. ```bash # Review the manifest and the exact payload `ucode apply` would publish. @@ -251,6 +253,18 @@ ucode setup show ucode setup --from-file ./managed-config.json ``` +To test before publishing, add `--local` to a launch. The override lasts for that invocation only: + +```bash +ucode --local # test the draft's configured default agent +ucode claude --local # or test a specific enabled agent +``` + +Without `--local`, launches always use the workspace-published config when one exists. The local +draft is never turned into a persistent mode and ordinary launches never overwrite it, so an admin +can edit and test it repeatedly. Spend-tier recommendations begin after publication because they +are resolved server-side. + Once the manifest looks right, publish it: ```bash @@ -304,7 +318,9 @@ their next ucode run. | `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 apply` would publish | | `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | -| `ucode apply` | Publish the authored managed config to the workspace, after a diff and confirmation (admins only) | +| `ucode --local` | Launch the draft's default agent for an optional local test | +| `ucode --local` | Launch a specific agent using the local draft for this invocation only | +| `ucode apply` | Publish the latest local draft to the workspace, whether or not it was tested | | `ucode apply --yes` | Publish without the confirmation prompt | ## Managed Local Files diff --git a/src/ucode/cli.py b/src/ucode/cli.py index fd6bbfd2..d2468457 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -63,9 +63,9 @@ render_budget_panel, ) from ucode.managed_config import ( - MANAGED_CONFIG_ENV_VAR, get_model_recommendation, load_managed_cache, + load_managed_state, managed_agent_config_enabled, refresh_managed_config, ) @@ -1428,8 +1428,7 @@ def _reject_disabled_agent(managed: dict | None, tool: str) -> None: def _fetch_managed_config(state: dict) -> tuple[dict | None, bool]: """The workspace's managed config for this launch, or ``(None, _)`` when there is none. - Returns ``(None, False)`` when managed configs are switched off — either the feature is disabled - or the launch passed ``--skip-managed-config`` (which clears the enabling env var for the process). + Returns ``(None, False)`` when the managed-config feature is disabled. """ if not managed_agent_config_enabled(): @@ -1597,6 +1596,37 @@ def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: _download_managed_skills(managed, state) +def _require_local_config_admin(state: dict) -> None: + """Fail closed unless the current user is an admin of the draft's workspace. + + Unlike ``setup`` and ``apply``, a local launch has no workspace API operation that can enforce + authorization server-side. An unknown SCIM result therefore has to block the launch rather + than optimistically continuing. + """ + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace: + raise RuntimeError("No workspace configured. Run `ucode configure` first.") + try: + token = get_databricks_token(workspace, state.get("profile")) + with spinner("Checking workspace admin permissions..."): + admin = is_workspace_admin(workspace, token) + except RuntimeError as exc: + raise RuntimeError( + "Could not verify workspace admin permissions for `--local`: " + f"{exc}. Check your Databricks authentication and retry." + ) from exc + if admin is False: + raise RuntimeError( + f"You are not an admin of {workspace}. `--local` runs an unpublished workspace-wide " + "coding config, so it is restricted to workspace admins." + ) + if admin is None: + raise RuntimeError( + "Could not verify workspace admin permissions for `--local`. Check your Databricks " + "authentication and workspace connectivity, then retry." + ) + + def _launch_tool( tool_name: str, ctx: typer.Context, @@ -1607,6 +1637,8 @@ def _launch_tool( managed: dict | None = None, recommendation: dict | None = None, model: str | None = None, + local_config: bool = False, + local_admin_verified: bool = False, ) -> None: try: tool = normalize_tool(tool_name) @@ -1642,7 +1674,16 @@ def _launch_tool( # at all and whether the model discovery below can be skipped. # Bare `ucode` already fetched one to choose the agent; refetching would double the # control-plane round trip and any fallback warning it printed. - if managed is None: + if local_config: + if managed is None: + managed = load_managed_state(state["workspace"]) + if not managed: + raise RuntimeError( + "No local managed config was found for this workspace. Run `ucode setup` first." + ) + if not local_admin_verified: + _require_local_config_admin(state) + elif managed is None: managed, _coding_agent_config_feature_disabled = _fetch_managed_config(state) # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast. _reject_disabled_agent(managed, tool) @@ -1664,12 +1705,20 @@ def _launch_tool( # `configure_shared_state`, whose returned state it overrides, and before the provider and # model are settled below — the two state files are never merged on disk. # Bare `ucode` already read one to choose the agent; refetching would double the round trip. - if recommendation is None: + if recommendation is None and not local_config: recommendation = _fetch_budget_recommendation(state, managed) + if local_config and managed is not None and managed.get("budget_policy"): + print_note( + "The local config uses its baseline agent and model; spend-tier routing is " + "resolved from the published workspace config." + ) _note_recommended_agent(recommendation, tool) if managed is not None: state = resolve_state(managed, state, tool) - print_success("Applied your workspace's managed coding agent config") + if local_config: + print_success("Applied your local managed coding agent config") + else: + print_success("Applied your workspace's managed coding agent config") unservable = managed_unservable_models(managed, tool) if unservable: print_warning( @@ -1818,7 +1867,7 @@ def _launch_tool( ) print_section(f"ucode with {TOOL_SPECS[tool]['display']}") if managed is not None: - print_kv("Config", "workspace-managed") + print_kv("Config", "local" if local_config else "workspace-managed") if provider: print_kv("Provider", provider) elif model and tool == "claude": @@ -1849,7 +1898,7 @@ def _launch_tool( # Register the managed config's MCP servers so they reach the agent's `/mcp` list. Nothing # else on this path does it — the config only lists them — so without this a # workspace-published server never shows up. `managed` is already None when the config is - # skipped (--skip-managed-config / feature off); --dry-run writes nothing. + # disabled by the feature gate; --dry-run writes nothing. if managed is not None and not is_dry_run(): _register_managed_mcp_servers(managed, tool, state) _apply_managed_skills(managed, tool, state) @@ -1866,8 +1915,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 # 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. +# which skips the model smoke test. SkipPreflightOption = Annotated[ bool, typer.Option( @@ -1877,30 +1925,34 @@ 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 -# feature is off, so a headless launcher can always pass it. -SkipManagedConfigOption = Annotated[ +# Reserve the removed flag so launchers do not forward it to the underlying agent as an unknown +# argument. It no longer provides any managed-config bypass. +RemovedSkipManagedConfigOption = Annotated[ bool, typer.Option( "--skip-managed-config", - help="Ignore your workspace's managed coding-agent config for this run, as if managed " - "configs were switched off — use your own local settings instead.", + help="Removed: workspace managed configs can no longer be bypassed.", hidden=True, ), ] +LocalManagedConfigOption = Annotated[ + bool, + typer.Option( + "--local", + help="Workspace admins only: use the latest local config authored by `ucode setup` for " + "this launch instead of the published workspace config.", + ), +] -def _disable_managed_config_if_requested(skip_managed_config: bool) -> None: - """Make this process behave as though ``ENABLE_MANAGED_AGENT_CONFIG`` were never set. - ``managed_agent_config_enabled()`` reads the env var live and gates every managed-config path - (the launch fetch/apply, the budget read, MCP registration, the bare-``ucode`` agent picker, and - the ``configure`` reject-under-managed flow), so clearing it once here short-circuits them all - without threading a flag through each. Per-invocation only: it affects just the current command. - """ +def _reject_removed_skip_managed_config(skip_managed_config: bool) -> None: if skip_managed_config: - os.environ.pop(MANAGED_CONFIG_ENV_VAR, None) + print_err( + "`--skip-managed-config` has been removed. Normal launches use the workspace's " + "managed config; workspace admins can test a local draft with `--local`." + ) + raise typer.Exit(2) # Target this launch at a specific workspace, auto-configuring (and logging in) @@ -1937,7 +1989,8 @@ def default( ), ] = False, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, workspace: WorkspaceOption = None, ) -> None: """Configure and launch coding agents through Databricks AI Gateway. @@ -1947,10 +2000,14 @@ def default( if ctx.invoked_subcommand is not None: return set_dry_run(dry_run) - _disable_managed_config_if_requested(skip_managed_config) + _reject_removed_skip_managed_config(skip_managed_config) try: _launch_managed_default( - ctx, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace + ctx, + dry_run=dry_run, + skip_preflight=skip_preflight, + workspace=workspace, + local_config=local_config, ) except typer.Exit: # `typer.Exit` subclasses RuntimeError, so it has to be re-raised ahead of the handler @@ -1968,9 +2025,10 @@ def _launch_managed_default( dry_run: bool, skip_preflight: bool, workspace: str | None, + local_config: bool, ) -> None: """Route bare ``ucode`` by whether the workspace publishes a managed config.""" - if not managed_agent_config_enabled(): + if not managed_agent_config_enabled() and not local_config: console.print(ctx.get_help()) return if workspace: @@ -1981,19 +2039,27 @@ def _launch_managed_default( if not current: raise RuntimeError("No workspace configured. Run `ucode configure` first.") apply_pat_environment(state) - # --dry-run avoids the fetch but still applies the last saved config. - if dry_run: + coding_agent_config_feature_disabled = False + # Both avoid a control-plane fetch, but they intentionally read different files: --local uses + # the admin's editable draft, while an ordinary dry-run stays on the published-config cache. + if local_config: + managed = load_managed_state(current) + elif dry_run: managed = load_managed_cache(current) else: with spinner("Loading..."): managed, coding_agent_config_feature_disabled = refresh_managed_config(state) + if local_config and not managed: + raise RuntimeError("No local managed config was found. Run `ucode setup` first.") + if local_config: + _require_local_config_admin(state) if not managed and not coding_agent_config_feature_disabled: _print_no_managed_config_guidance(current, state.get("profile")) if not managed: 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. - recommendation = _fetch_budget_recommendation(state, managed) + recommendation = None if local_config else _fetch_budget_recommendation(state, managed) tool = recommended_agent(recommendation, managed) or next( iter(managed.get("enabled_agents") or {}), None ) @@ -2010,6 +2076,8 @@ def _launch_managed_default( workspace=workspace, managed=managed, recommendation=recommendation, + local_config=local_config, + local_admin_verified=local_config, ) @@ -2044,7 +2112,8 @@ def codex_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -2062,7 +2131,7 @@ def codex_cmd( ] = False, ) -> None: """Launch Codex via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) + _reject_removed_skip_managed_config(skip_managed_config) if enable_smart_routing_flag and disable_smart_routing_flag: print_err("Use only one of --enable-smart-routing or --disable-smart-routing.") raise typer.Exit(1) @@ -2077,6 +2146,7 @@ def codex_cmd( skip_preflight=skip_preflight, workspace=workspace, enable_smart_routing_flag=enable_smart_routing_flag, + local_config=local_config, ) @@ -2103,7 +2173,8 @@ def claude_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -2121,7 +2192,7 @@ def claude_cmd( ] = False, ) -> None: """Launch Claude Code via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) + _reject_removed_skip_managed_config(skip_managed_config) if enable_smart_routing_flag and disable_smart_routing_flag: print_err("Use only one of --enable-smart-routing or --disable-smart-routing.") raise typer.Exit(1) @@ -2137,6 +2208,7 @@ def claude_cmd( skip_preflight=skip_preflight, workspace=workspace, enable_smart_routing_flag=enable_smart_routing_flag, + local_config=local_config, ) @@ -2144,11 +2216,12 @@ def claude_cmd( def gemini_cmd( ctx: typer.Context, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, ) -> None: """Launch Gemini CLI via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("gemini", ctx, skip_preflight=skip_preflight) + _reject_removed_skip_managed_config(skip_managed_config) + _launch_tool("gemini", ctx, skip_preflight=skip_preflight, local_config=local_config) @app.command( @@ -2157,33 +2230,36 @@ def gemini_cmd( def opencode_cmd( ctx: typer.Context, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, ) -> None: """Launch OpenCode via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("opencode", ctx, skip_preflight=skip_preflight) + _reject_removed_skip_managed_config(skip_managed_config) + _launch_tool("opencode", ctx, skip_preflight=skip_preflight, local_config=local_config) @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def copilot_cmd( ctx: typer.Context, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, ) -> None: """Launch GitHub Copilot CLI via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("copilot", ctx, skip_preflight=skip_preflight) + _reject_removed_skip_managed_config(skip_managed_config) + _launch_tool("copilot", ctx, skip_preflight=skip_preflight, local_config=local_config) @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def pi_cmd( ctx: typer.Context, skip_preflight: SkipPreflightOption = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, + local_config: LocalManagedConfigOption = False, ) -> None: """Launch Pi coding agent via Databricks.""" - _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("pi", ctx, skip_preflight=skip_preflight) + _reject_removed_skip_managed_config(skip_managed_config) + _launch_tool("pi", ctx, skip_preflight=skip_preflight, local_config=local_config) @app.command("cursor", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -2334,7 +2410,7 @@ def configure( "still applied.", ), ] = False, - skip_managed_config: SkipManagedConfigOption = False, + skip_managed_config: RemovedSkipManagedConfigOption = False, verbose: Annotated[ str, typer.Option( @@ -2347,7 +2423,7 @@ def configure( """Configure workspace URL and AI Gateway.""" if ctx.invoked_subcommand is not None: return - _disable_managed_config_if_requested(skip_managed_config) + _reject_removed_skip_managed_config(skip_managed_config) if verbose not in ("normal", "low"): print_err("--verbose must be one of: normal, low.") raise typer.Exit(2) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 811b9f1f..c7b00492 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -1,10 +1,9 @@ """Resolve the effective agent settings from the managed config plus local ucode state. -The managed config (``~/.ucode/managed-state.json`` — authored by ``ucode setup`` and refreshed -from the workspace at launch, both through :mod:`ucode.managed_config`) and the developer's own -ucode state (``~/.ucode/state.json``) stay separate files — they are never merged on disk. Instead -this module resolves them *per key* at config-write time: whatever the manifest specifies wins, and -anything it leaves unset falls back to +The selected managed config (the workspace-published config normally, or the local setup draft when +``--local`` is passed) and the developer's own ucode state (``~/.ucode/state.json``) are never merged +on disk. Instead this module resolves them *per key* at config-write time: whatever the manifest +specifies wins, and anything it leaves unset falls back to the developer's ucode state. The resolved view is what gets rendered into the agent config files (e.g. ``~/.claude/ucode-settings.json``), so managed settings take precedence for every ``ucode`` command without either file being rewritten. diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 5f7be5d9..bd3199f9 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -1,9 +1,9 @@ """Interactive `ucode 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 apply`` (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`). +it with ``ucode apply`` (a separate command, so the manifest can be reviewed or tested first). The +editable draft lives at ``~/.ucode/managed-state.json`` and is never overwritten by ordinary +launches. 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`` @@ -1447,6 +1447,7 @@ def _print_next_steps(manifest: dict) -> None: # would report "run `ucode setup` first". print_note("Dry run — nothing was saved. Re-run without --dry-run to author the config.") return + print_note("Local draft: [bold]~/.ucode/managed-state.json[/bold]") # These sections aren't required to publish — call them out as optional so an admin doesn't read # a config with none configured as unfinished. print_note("[dim]Optional — configure any of these, or skip straight to publishing:[/dim]") @@ -1454,7 +1455,10 @@ def _print_next_steps(manifest: dict) -> None: console.print(line) print_panel( "All done?", - ["Publish with [bold]ucode apply[/bold] so all developers use this configuration."], + [ + "Optionally test the draft with [bold]ucode --local[/bold] or " + "[bold]ucode --local[/bold]. Publish it with [bold]ucode apply[/bold]." + ], ) @@ -1845,8 +1849,8 @@ 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 apply`, neither of which a flag listing conveys. Reads - the manifest but never authenticates, so it works before `ucode configure`. + that nothing reaches developers until `ucode apply`, neither of which a flag listing conveys. + Reads the manifest but never authenticates, so it works before `ucode configure`. """ print_section("ucode setup") print_note( @@ -1854,8 +1858,7 @@ def setup_help_command() -> int: "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 apply`." + "Each command below edits your local draft; nothing reaches developers until `ucode apply`." ) workspace = load_state().get("workspace") or managed_state_workspace() @@ -1864,6 +1867,7 @@ def setup_help_command() -> int: # 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("ucode --local")) print_heading("1. Start here") console.print( @@ -1882,11 +1886,13 @@ def setup_help_command() -> int: for line in _section_status_lines(manifest, width): console.print(line) - print_heading("3. Review and publish") + print_heading("3. Review, optionally test, and publish") console.print( _command_line("ucode setup show", "The draft, and the payload `apply` sends", width=width) ) - console.print(_command_line("ucode apply", "Publish it to the workspace", width=width)) + console.print(_command_line("ucode --local", "Test the configured default agent", width=width)) + console.print(_command_line("ucode --local", "Test a specific agent", width=width)) + console.print(_command_line("ucode apply", "Publish it for every developer", width=width)) print_heading("Also") console.print( diff --git a/tests/test_cli.py b/tests/test_cli.py index f84ef454..d5739ece 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2414,6 +2414,149 @@ def test_absent_flag_defaults_false(self, tool): assert cfg.call_args.kwargs["skip_preflight"] is False +class TestLocalManagedConfigFlag: + LAUNCH_TOOLS = ["codex", "claude", "gemini", "opencode", "copilot", "pi"] + + @pytest.mark.parametrize("tool", LAUNCH_TOOLS) + def test_launch_flag_is_forwarded_for_every_managed_agent(self, tool): + with patch("ucode.cli._launch_tool") as launch: + result = runner.invoke(app, [tool, "--local"]) + + assert result.exit_code == 0, result.output + assert launch.call_args.kwargs["local_config"] is True + + def test_absent_flag_keeps_the_workspace_source(self): + with patch("ucode.cli._launch_tool") as launch: + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + assert launch.call_args.kwargs["local_config"] is False + + def test_help_identifies_local_as_admin_only(self): + result = runner.invoke(app, ["claude", "--help"]) + + assert result.exit_code == 0, result.output + assert "Workspace admins only" in _strip_ansi(result.output) + + @pytest.mark.parametrize( + "args", + [ + ["--skip-managed-config"], + ["claude", "--skip-managed-config"], + ["configure", "--skip-managed-config"], + ], + ) + def test_removed_skip_managed_config_is_rejected(self, args): + with patch("ucode.cli._launch_tool") as launch: + result = runner.invoke(app, args) + + assert result.exit_code == 2 + assert "has been removed" in result.output + launch.assert_not_called() + + def test_local_launch_reads_the_draft_and_never_fetches_workspace_config(self): + authored = { + "enabled_agents": { + "claude": {"model_config": {"default_model": "databricks-claude-sonnet-4"}} + } + } + fetch = MagicMock(side_effect=AssertionError("--local must not fetch the workspace config")) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.load_managed_state", return_value=authored) as load_draft, + patch("ucode.cli._require_local_config_admin") as require_admin, + patch("ucode.cli._fetch_managed_config", fetch), + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), + patch( + "ucode.cli.resolve_launch_model", + return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), + ), + patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch("ucode.cli._register_managed_mcp_servers"), + patch("ucode.cli._apply_managed_skills"), + patch("ucode.cli.launch_agent"), + ): + result = runner.invoke(app, ["claude", "--local"]) + + assert result.exit_code == 0, result.output + load_draft.assert_called_once_with(MINIMAL_STATE["workspace"]) + require_admin.assert_called_once_with(MINIMAL_STATE) + fetch.assert_not_called() + assert "local managed coding agent config" in result.output + + def test_specific_local_launch_without_a_draft_is_actionable(self): + fetch = MagicMock(side_effect=AssertionError("--local must not fetch the workspace config")) + require_admin = MagicMock() + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.load_managed_state", return_value=None), + patch("ucode.cli._require_local_config_admin", require_admin), + patch("ucode.cli._fetch_managed_config", fetch), + ): + result = runner.invoke(app, ["claude", "--local"]) + + assert result.exit_code == 1 + assert "setup" in result.output + require_admin.assert_not_called() + fetch.assert_not_called() + + def test_non_admin_cannot_launch_a_local_draft(self): + authored = {"enabled_agents": {"claude": {}}} + launch = MagicMock() + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.load_managed_state", return_value=authored), + patch("ucode.cli.get_databricks_token", return_value="token"), + patch("ucode.cli.is_workspace_admin", return_value=False), + patch("ucode.cli.launch_agent", launch), + ): + result = runner.invoke(app, ["claude", "--local"]) + + assert result.exit_code == 1 + assert "not an admin" in result.output + launch.assert_not_called() + + +class TestRequireLocalConfigAdmin: + def test_verified_admin_is_allowed(self): + with ( + patch("ucode.cli.get_databricks_token", return_value="token") as get_token, + patch("ucode.cli.is_workspace_admin", return_value=True) as is_admin, + ): + import ucode.cli as cli_mod + + cli_mod._require_local_config_admin(MINIMAL_STATE) + + get_token.assert_called_once_with(MINIMAL_STATE["workspace"], None) + is_admin.assert_called_once_with(MINIMAL_STATE["workspace"], "token") + + @pytest.mark.parametrize("admin", [False, None]) + def test_non_admin_or_unverifiable_status_is_blocked(self, admin): + with ( + patch("ucode.cli.get_databricks_token", return_value="token"), + patch("ucode.cli.is_workspace_admin", return_value=admin), + ): + import ucode.cli as cli_mod + + with pytest.raises(RuntimeError, match="admin"): + cli_mod._require_local_config_admin(MINIMAL_STATE) + + def test_authentication_failure_is_blocked_with_actionable_error(self): + with patch( + "ucode.cli.get_databricks_token", side_effect=RuntimeError("credentials expired") + ): + import ucode.cli as cli_mod + + with pytest.raises(RuntimeError, match="authentication.*retry"): + cli_mod._require_local_config_admin(MINIMAL_STATE) + + class TestRejectDisabledAgent: """`enabled_agents` is an allowlist: an agent the admin didn't enable would launch unmanaged.""" @@ -2461,25 +2604,13 @@ def test_disabled_reads_nothing_at_all(self, monkeypatch, env_value): monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False) else: monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", env_value) - for name in ("refresh_managed_config", "load_managed_cache"): + for name in ("refresh_managed_config", "load_managed_state"): monkeypatch.setattr( f"ucode.cli.{name}", lambda *a, called=name, **k: pytest.fail(f"{called} must not run when disabled"), ) assert self._fetch({"workspace": "https://w"}) == (None, False) - def test_skip_managed_config_makes_the_fetch_a_no_op(self, monkeypatch): - # --skip-managed-config clears the enabling env var, so the read behaves as feature-off: - # no fetch, no cache read, no network — just None. - monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", lambda state: pytest.fail("should not fetch") - ) - import ucode.cli as cli_mod - - cli_mod._disable_managed_config_if_requested(True) - assert self._fetch({"workspace": "https://w"}) == (None, False) - class TestManagedConfigDecidesDiscoveryFromFreshRead: def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): @@ -2550,7 +2681,7 @@ def test_fetches_the_config_rather_than_reading_a_cold_cache(self, monkeypatch): monkeypatch.setattr("ucode.cli.set_current_workspace", lambda ws: None) monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) # Cold cache — a cache read would wrongly fall through to the local configure flow. - monkeypatch.setattr("ucode.cli.load_managed_cache", lambda ws: None) + monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: None) monkeypatch.setattr( "ucode.cli.refresh_managed_config", lambda state: ({"enabled_agents": {"claude": {}}}, False), @@ -2678,7 +2809,7 @@ def test_passes_entries_through_when_the_env_var_is_off(self, monkeypatch, capsy else: monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", env_value) monkeypatch.setattr( - "ucode.cli.load_managed_cache", + "ucode.cli.load_managed_state", lambda ws: pytest.fail("must not read the config when disabled"), ) monkeypatch.setattr( @@ -2761,7 +2892,6 @@ def _run( managed, is_admin=False, args=None, - cached=None, coding_agent_config_feature_disabled=False, ): launched: list[tuple] = [] @@ -2775,7 +2905,6 @@ def _run( else: monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (managed, False)) - monkeypatch.setattr("ucode.cli.load_managed_cache", lambda ws: cached) monkeypatch.setattr("ucode.cli.get_databricks_token", lambda *a, **k: "tok") monkeypatch.setattr("ucode.cli.is_workspace_admin", lambda *a, **k: is_admin) monkeypatch.setattr( @@ -2792,6 +2921,47 @@ def test_launches_the_managed_default_agent(self, monkeypatch): assert "paved" not in result.output # no policy set in this config assert "Claude Code" in result.output + def test_local_flag_launches_the_authored_default_without_the_feature_flag(self, monkeypatch): + monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False) + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.apply_pat_environment", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) + monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: self.MANAGED) + require_admin = MagicMock() + monkeypatch.setattr("ucode.cli._require_local_config_admin", require_admin) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", + lambda state: pytest.fail("--local must not fetch the workspace config"), + ) + launched: list[tuple] = [] + monkeypatch.setattr( + "ucode.cli._launch_tool", lambda tool, ctx, **kw: launched.append((tool, kw)) + ) + + result = runner.invoke(app, ["--local"]) + + assert result.exit_code == 0, result.output + assert launched[0][0] == "claude" + assert launched[0][1]["managed"] == self.MANAGED + assert launched[0][1]["local_config"] is True + assert launched[0][1]["local_admin_verified"] is True + require_admin.assert_called_once_with({"workspace": "https://w"}) + + def test_local_flag_without_an_authored_config_is_actionable(self, monkeypatch): + monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False) + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.apply_pat_environment", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) + monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: None) + require_admin = MagicMock() + monkeypatch.setattr("ucode.cli._require_local_config_admin", require_admin) + + result = runner.invoke(app, ["--local"]) + + assert result.exit_code == 1 + assert "setup" in result.output + require_admin.assert_not_called() + def test_falls_back_to_the_first_enabled_agent(self, monkeypatch): managed = {"enabled_agents": {"opencode": {}}} result, launched = self._run(monkeypatch, managed=managed) @@ -2869,18 +3039,6 @@ def test_skip_preflight_still_resolves_an_agent_from_the_managed_config(self, mo assert seen["tool"] == "claude" assert seen["skip_preflight"] is True - def test_skip_managed_config_behaves_as_feature_off(self, monkeypatch): - # --skip-managed-config clears the enabling env var, so bare `ucode` has no config to pick an - # agent from and just prints help — exactly the feature-off behavior, no fetch. - monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: pytest.fail("--skip-managed-config must not fetch"), - ) - result = runner.invoke(app, ["--skip-managed-config"]) - assert result.exit_code == 0, result.output - assert "Usage:" in result.output - @pytest.mark.parametrize("env_value", [None, "", "0"]) def test_prints_help_when_the_env_var_is_off(self, monkeypatch, env_value): if env_value is None: @@ -2906,29 +3064,6 @@ def test_subcommands_still_work(self, monkeypatch): result = runner.invoke(app, ["status"]) assert result.exit_code == 0, result.output - def test_launcher_skip_managed_config_does_not_fetch(self, monkeypatch): - # `ucode claude --skip-managed-config` clears the env var, so the launch never reads the - # workspace's managed config and falls back to the developer's own settings. - monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") - monkeypatch.setattr( - "ucode.cli.refresh_managed_config", - lambda state: pytest.fail("--skip-managed-config must not fetch"), - ) - state = dict(MINIMAL_STATE) - with ( - patch("ucode.cli.load_state", return_value=state), - patch("ucode.cli.apply_pat_environment"), - patch("ucode.cli.ensure_bootstrap_dependencies"), - patch("ucode.cli.ensure_provider_state", return_value=state), - patch("ucode.cli.configure_shared_state", return_value=state), - patch("ucode.cli.configure_tool", return_value=state), - patch("ucode.cli.get_databricks_token", return_value="tok"), - patch("ucode.cli.launch_agent"), - ): - result = runner.invoke(app, ["claude", "--skip-managed-config"]) - assert result.exit_code == 0, result.output - assert "managed coding agent config" not in result.output - class TestBudgetRecommendationAtLaunch: """The budget read informs the launch; it never blocks it.""" diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 8cdd9a28..2ee299ea 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2255,6 +2255,8 @@ def test_lists_every_setup_command(self, capsys): "ucode setup skills", "ucode setup spend-tiers", "ucode setup show", + "ucode --local", + "ucode --local", "ucode apply", ): assert command in out @@ -2579,7 +2581,7 @@ def test_apply_is_registered(self): assert result.exit_code == 0 assert "apply" in result.output - def test_apply_declares_yes_and_no_dry_run(self): + def test_apply_declares_yes_and_no_mode_flags(self): # `--dry-run` was removed: apply always validates before publishing, so a separate # validate-only mode is redundant. Asserted on declared options rather than rendered help, # which Rich ellipsizes at narrow widths (see test_setup_help_lists_from_file). @@ -2587,6 +2589,15 @@ def test_apply_declares_yes_and_no_dry_run(self): declared = {opt for param in command.params for opt in param.opts} assert "--yes" in declared assert "--dry-run" not in declared + assert "--local" not in declared + assert "--workspace" not in declared + + def test_apply_publishes_the_latest_draft(self): + with patch.object(cli_mod, "apply_command", return_value=0) as apply: + result = runner.invoke(app, ["apply", "--yes"]) + + assert result.exit_code == 0, result.output + assert apply.call_args.kwargs == {"yes": True} def test_apply_error_exits_nonzero_with_a_message(self): with patch.object(cli_mod, "apply_command", side_effect=RuntimeError("no config authored")): diff --git a/tests/test_managed_workflow.py b/tests/test_managed_workflow.py new file mode 100644 index 00000000..bce20923 --- /dev/null +++ b/tests/test_managed_workflow.py @@ -0,0 +1,133 @@ +"""End-to-end managed-config workflow with only external systems stubbed. + +Exercises the real CLI orchestration, draft/cache persistence, manifest serialization, +managed-state resolution, and Claude config writer across setup, local testing, publishing, and a +normal workspace-managed launch. +""" + +from __future__ import annotations + +import contextlib +import json +from unittest.mock import patch + +from typer.testing import CliRunner + +import ucode.managed_config as managed_config +import ucode.managed_wizard as managed_wizard +from ucode.agents import claude +from ucode.cli import app + +runner = CliRunner() + +WORKSPACE = "https://workspace.example.com" +PUBLISHED_MODEL = "system.ai.claude-opus-4-8" +UNPUBLISHED_MODEL = "system.ai.claude-opus-4-9" + + +def _manifest(model: str) -> dict: + return { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": model, + "models": {"default_opus_model": model}, + } + } + }, + } + + +def _configured_opus(settings: dict) -> str: + return settings["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"].removesuffix("[1m]") + + +def test_setup_local_apply_and_normal_launch_keep_sources_separate(tmp_path, monkeypatch): + state = { + "workspace": WORKSPACE, + "profile": "DEFAULT", + "available_tools": ["claude"], + "base_urls": {"claude": f"{WORKSPACE}/ai-gateway/anthropic"}, + "claude_models": {"opus": PUBLISHED_MODEL}, + "all_claude_models": [PUBLISHED_MODEL, UNPUBLISHED_MODEL], + "managed_configs": {}, + } + manifest_path = tmp_path / "managed-config.json" + manifest_path.write_text(json.dumps(_manifest(PUBLISHED_MODEL)), encoding="utf-8") + settings_path = tmp_path / "claude" / "ucode-settings.json" + + monkeypatch.setattr(managed_config, "MANAGED_STATE_PATH", tmp_path / "managed-state.json") + monkeypatch.setattr(managed_config, "MANAGED_CACHE_DIR", tmp_path / "managed-cache") + monkeypatch.setattr( + managed_config, "LEGACY_MANAGED_CACHE_PATH", tmp_path / "legacy-managed-cache.json" + ) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "claude-settings.backup.json") + + published_payload: dict = {} + launches: list[str] = [] + + def publish(_workspace, _token, payload): + published_payload.update(payload) + return {"name": "coding-agent-configs/test"}, None + + patches = [ + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.load_state", return_value=state), + patch.object(managed_wizard, "load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli._require_local_config_admin"), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli._register_managed_mcp_servers"), + patch("ucode.cli._apply_managed_skills"), + patch("ucode.cli.launch_agent", side_effect=lambda tool, *_args: launches.append(tool)), + patch.object(claude, "managed_settings_model_overrides", return_value=None), + patch.object(claude, "agent_version", return_value="test"), + patch.object(claude, "ucode_version", return_value="test"), + patch.object(managed_wizard, "ensure_databricks_auth"), + patch.object(managed_wizard, "get_databricks_token", return_value="token"), + patch.object(managed_wizard, "is_workspace_admin", return_value=True), + patch.object(managed_wizard, "get_managed_config", return_value=(None, None)), + patch.object(managed_wizard, "create_coding_agent_config", side_effect=publish), + ] + with contextlib.ExitStack() as stack: + for managed_patch in patches: + stack.enter_context(managed_patch) + + setup_result = runner.invoke(app, ["setup", "--from-file", str(manifest_path)]) + assert setup_result.exit_code == 0, setup_result.output + assert managed_config.load_managed_state(WORKSPACE) == _manifest(PUBLISHED_MODEL) + + local_result = runner.invoke(app, ["claude", "--local"]) + assert local_result.exit_code == 0, local_result.output + local_settings = json.loads(settings_path.read_text(encoding="utf-8")) + assert _configured_opus(local_settings) == PUBLISHED_MODEL + + apply_result = runner.invoke(app, ["apply", "--yes"]) + assert apply_result.exit_code == 0, apply_result.output + published = managed_config.normalize_managed_config(published_payload) + assert published["enabled_agents"]["claude"]["model_config"]["default_model"] == ( + PUBLISHED_MODEL + ) + + # The admin continues editing after publication. A normal developer launch must still use + # the workspace-published snapshot, not this newer local draft. + managed_config.save_managed_state(WORKSPACE, _manifest(UNPUBLISHED_MODEL)) + with ( + patch.dict("os.environ", {managed_config.MANAGED_CONFIG_ENV_VAR: "1"}), + patch.object(managed_config, "get_databricks_token", return_value="token"), + patch.object(managed_config, "get_managed_config", return_value=(published, None)), + ): + normal_result = runner.invoke(app, ["claude"]) + + assert normal_result.exit_code == 0, normal_result.output + + normal_settings = json.loads(settings_path.read_text(encoding="utf-8")) + assert _configured_opus(normal_settings) == PUBLISHED_MODEL + assert managed_config.load_managed_state(WORKSPACE) == _manifest(UNPUBLISHED_MODEL) + assert managed_config.load_managed_cache(WORKSPACE) == published + assert launches == ["claude", "claude"]