diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a78c68a..aad4f58 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -33,7 +33,7 @@ ) from ucode.agents.codex import revert_legacy_shared_config from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH -from ucode.config_io import restore_file, set_dry_run +from ucode.config_io import is_dry_run, restore_file, set_dry_run from ucode.databricks import ( apply_pat_environment, build_shared_base_urls, @@ -62,6 +62,7 @@ render_budget_panel, ) from ucode.managed_config import ( + MANAGED_CONFIG_ENV_VAR, get_model_recommendation, load_managed_state, managed_agent_config_enabled, @@ -1346,16 +1347,14 @@ def _reject_disabled_agent(managed: dict | None, tool: str) -> None: ) -def _fetch_managed_config(state: dict, *, skip_preflight: bool) -> dict | None: +def _fetch_managed_config(state: dict) -> dict | None: """The workspace's managed config for this launch, or None when there is none. - ``skip_preflight`` mirrors the launch flag: it reads the last persisted copy instead of - re-fetching, so the config can be stale until a normal launch refreshes it. + Returns None 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). """ if not managed_agent_config_enabled(): return None - if skip_preflight: - return load_managed_state(state.get("workspace")) or None with spinner("Checking for a managed coding agent config..."): return refresh_managed_config(state) @@ -1379,15 +1378,13 @@ def _note_recommended_agent(recommendation: dict | None, tool: str) -> None: ) -def _fetch_budget_recommendation( - state: dict, managed: dict | None, *, skip_preflight: bool -) -> dict | None: +def _fetch_budget_recommendation(state: dict, managed: dict | None) -> dict | None: """The agent and model the caller's budget tier allows, or None when there is no budget to read. Enforcement is server-side, so a failed read only costs the recommendation: the config's own ``default_model`` still applies and the launch proceeds. """ - if managed is None or skip_preflight: + if managed is None or is_dry_run(): return None reason: str | None = None recommendation = None @@ -1532,7 +1529,7 @@ def _launch_tool( # 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: - managed = _fetch_managed_config(state, skip_preflight=skip_preflight) + managed = _fetch_managed_config(state) # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast. _reject_disabled_agent(managed, tool) # Discovery exists to find models and isn't needed for managed config that already names them. @@ -1554,9 +1551,7 @@ def _launch_tool( # 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: - recommendation = _fetch_budget_recommendation( - state, managed, skip_preflight=skip_preflight - ) + recommendation = _fetch_budget_recommendation(state, managed) _note_recommended_agent(recommendation, tool) if managed is not None: state = resolve_state(managed, state, tool) @@ -1732,9 +1727,9 @@ def _launch_tool( _print_budget_panel(recommendation, tool, managed) # 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. Skipped under --skip-preflight (deliberately - # unmanaged). - if managed is not None and not skip_preflight: + # workspace-published server never shows up. `managed` is already None when the config is + # skipped (--skip-managed-config / feature off); --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) print_success(f"Starting {TOOL_SPECS[tool]['display']}") @@ -1750,17 +1745,42 @@ 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. +# which skips the model smoke test, and from `--skip-managed-config`, which +# controls whether the workspace's managed config is applied. SkipPreflightOption = Annotated[ bool, typer.Option( "--skip-preflight", help="Skip the per-launch Databricks auth + AI Gateway re-validation, trusting a " - "prior `ucode configure`. Launches with your own local settings, ignoring any " - "workspace managed config.", + "prior `ucode configure`.", ), ] +# 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[ + 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.", + ), +] + + +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. + """ + if skip_managed_config: + os.environ.pop(MANAGED_CONFIG_ENV_VAR, 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`. WorkspaceOption = Annotated[ @@ -1786,7 +1806,16 @@ def default( is_eager=True, ), ] = False, + dry_run: Annotated[ + bool, + typer.Option( + "--dry-run", + help="Print config files without writing them. Uses the last saved managed " + "config instead of fetching a fresh one.", + ), + ] = False, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, ) -> None: """Configure and launch coding agents through Databricks AI Gateway. @@ -1795,8 +1824,12 @@ def default( """ if ctx.invoked_subcommand is not None: return + set_dry_run(dry_run) + _disable_managed_config_if_requested(skip_managed_config) try: - _launch_managed_default(ctx, skip_preflight=skip_preflight, workspace=workspace) + _launch_managed_default( + ctx, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace + ) except typer.Exit: # `typer.Exit` subclasses RuntimeError, so it has to be re-raised ahead of the handler # below. Otherwise a launch that already reported its own error is followed by @@ -1810,6 +1843,7 @@ def default( def _launch_managed_default( ctx: typer.Context, *, + dry_run: bool, skip_preflight: bool, workspace: str | None, ) -> None: @@ -1825,20 +1859,18 @@ def _launch_managed_default( if not current: raise RuntimeError("No workspace configured. Run `ucode configure` first.") apply_pat_environment(state) - if skip_preflight: - # Deliberately unmanaged, so no config is read at all — and there is none to name an agent. - raise RuntimeError( - "--skip-preflight launches with your own settings, so `ucode` has no managed config " - "to pick an agent from. Run `ucode --skip-preflight` instead." - ) - with spinner("Checking for a managed coding agent config..."): - managed = refresh_managed_config(state) + # --dry-run avoids the fetch but still applies the last saved config. + if dry_run: + managed = load_managed_state(current) + else: + with spinner("Checking for a managed coding agent config..."): + managed = refresh_managed_config(state) if not managed: _print_no_managed_config_guidance(current, state.get("profile")) 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, skip_preflight=skip_preflight) + recommendation = _fetch_budget_recommendation(state, managed) tool = recommended_agent(recommendation, managed) or next( iter(managed.get("enabled_agents") or {}), None ) @@ -1889,6 +1921,7 @@ def codex_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -1906,6 +1939,7 @@ def codex_cmd( ] = False, ) -> None: """Launch Codex via Databricks.""" + _disable_managed_config_if_requested(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) @@ -1946,6 +1980,7 @@ def claude_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -1963,6 +1998,7 @@ def claude_cmd( ] = False, ) -> None: """Launch Claude Code via Databricks.""" + _disable_managed_config_if_requested(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) @@ -1982,28 +2018,48 @@ def claude_cmd( @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def gemini_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: +def gemini_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, +) -> None: """Launch Gemini CLI via Databricks.""" + _disable_managed_config_if_requested(skip_managed_config) _launch_tool("gemini", ctx, skip_preflight=skip_preflight) @app.command( "opencode", context_settings={"allow_extra_args": True, "ignore_unknown_options": True} ) -def opencode_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: +def opencode_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, +) -> None: """Launch OpenCode via Databricks.""" + _disable_managed_config_if_requested(skip_managed_config) _launch_tool("opencode", ctx, skip_preflight=skip_preflight) @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def copilot_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: +def copilot_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, +) -> None: """Launch GitHub Copilot CLI via Databricks.""" + _disable_managed_config_if_requested(skip_managed_config) _launch_tool("copilot", ctx, skip_preflight=skip_preflight) @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def pi_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: +def pi_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, +) -> None: """Launch Pi coding agent via Databricks.""" + _disable_managed_config_if_requested(skip_managed_config) _launch_tool("pi", ctx, skip_preflight=skip_preflight) @@ -2144,6 +2200,7 @@ def configure( "still applied.", ), ] = False, + skip_managed_config: SkipManagedConfigOption = False, verbose: Annotated[ str, typer.Option( @@ -2156,6 +2213,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) if verbose not in ("normal", "low"): print_err("--verbose must be one of: normal, low.") raise typer.Exit(2) diff --git a/tests/test_cli.py b/tests/test_cli.py index 181e935..efaf2e0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2209,10 +2209,10 @@ class TestFetchManagedConfig: """The launch path's managed-config read, which gates both the allowlist and model discovery.""" @staticmethod - def _fetch(state, *, skip_preflight=False): + def _fetch(state): import ucode.cli as cli_mod - return cli_mod._fetch_managed_config(state, skip_preflight=skip_preflight) + return cli_mod._fetch_managed_config(state) def test_fetches_fresh_when_enabled(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") @@ -2235,21 +2235,17 @@ def test_disabled_reads_nothing_at_all(self, monkeypatch, env_value): ) assert self._fetch({"workspace": "https://w"}) is None - def test_skip_preflight_reads_the_cache_without_fetching(self, monkeypatch): - # Headless launchers pass --skip-preflight to avoid per-launch network calls. + 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") ) - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: {"enabled_agents": {}}) - assert self._fetch({"workspace": "https://w"}, skip_preflight=True) == { - "enabled_agents": {} - } + import ucode.cli as cli_mod - def test_skip_preflight_with_an_empty_cache_is_none(self, monkeypatch): - monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") - monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: {}) - assert self._fetch({"workspace": "https://w"}, skip_preflight=True) is None + cli_mod._disable_managed_config_if_requested(True) + assert self._fetch({"workspace": "https://w"}) is None class TestManagedConfigDecidesDiscoveryFromFreshRead: @@ -2584,25 +2580,60 @@ def test_non_admin_without_a_config_is_told_to_ask(self, monkeypatch): assert launched == [] assert "Ask a workspace admin" in result.output - def test_skip_preflight_has_no_config_to_pick_an_agent_from(self, monkeypatch): - # --skip-preflight is deliberately unmanaged, so bare `ucode` cannot resolve an agent. It - # must say that rather than report "no config found", which would be wrong. + def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") 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.refresh_managed_config", - lambda state: pytest.fail("--skip-preflight must not fetch"), + lambda state: pytest.fail("--dry-run must not fetch"), ) + monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: self.MANAGED) + launched: list[tuple] = [] monkeypatch.setattr( - "ucode.cli.load_managed_state", - lambda ws: pytest.fail("--skip-preflight must not read the cache"), + "ucode.cli._launch_tool", lambda tool, ctx, **kw: launched.append((tool, kw)) + ) + result = runner.invoke(app, ["--dry-run"]) + assert result.exit_code == 0, result.output + # The config bare `ucode` already read is handed down, so the launch path does not refetch. + assert launched[0][1]["managed"] == self.MANAGED + + def test_skip_preflight_still_resolves_an_agent_from_the_managed_config(self, monkeypatch): + # --skip-preflight is now only about auth/gateway re-validation, decoupled from managed + # config, so bare `ucode --skip-preflight` still fetches the config and picks its agent. + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + 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"}) + managed = { + "default_agent": "claude", + "enabled_agents": {"claude": {"model_config": {"default_model": "m"}}}, + } + monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: managed) + monkeypatch.setattr("ucode.cli._fetch_budget_recommendation", lambda state, m: None) + monkeypatch.setattr("ucode.cli._print_managed_summary", lambda *a, **k: None) + seen: dict = {} + monkeypatch.setattr( + "ucode.cli._launch_tool", + lambda tool, ctx, **kw: seen.update({"tool": tool, **kw}), ) result = runner.invoke(app, ["--skip-preflight"]) - assert result.exit_code == 1 - assert "ucode --skip-preflight" in result.output - assert "No managed coding agent config was found" not in result.output + assert result.exit_code == 0, result.output + 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): @@ -2629,6 +2660,29 @@ 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_e2e.py b/tests/test_e2e.py index ebd2ab3..76b5087 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -573,9 +573,14 @@ def _first_service(tool: str, workspace: str, token: str) -> str: return names[0] @staticmethod - def _skip_if_no_permission(combined: str, provider: str) -> None: + def _skip_if_provider_unusable(combined: str, provider: str) -> None: + # Environmental provider-account conditions, not ucode bugs: the test only proves routing + # reaches the provider, so skip (rather than fail) when the account lacks a grant on the + # connection or has run out of credits — state outside the code under test. if "USE CONNECTION" in combined or "EXECUTE" in combined: pytest.skip(f"no permission on provider {provider}: {combined[:200]}") + if "Credit balance is too low" in combined: + pytest.skip(f"provider {provider} account is out of credits: {combined[:200]}") def test_launch_claude_through_provider( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token @@ -613,7 +618,7 @@ def test_launch_claude_through_provider( } result = _run_agent(claude.validate_cmd("claude"), env=env, timeout=90) combined = (result.stdout + result.stderr).strip() - self._skip_if_no_permission(combined, provider) + self._skip_if_provider_unusable(combined, provider) assert result.returncode == 0 and combined, ( f"provider={provider} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" @@ -650,7 +655,7 @@ def test_launch_codex_through_provider( except subprocess.TimeoutExpired: pytest.fail(f"provider={provider} timed out after {timeout_seconds}s") combined = (result.stdout + result.stderr).strip() - self._skip_if_no_permission(combined, provider) + self._skip_if_provider_unusable(combined, provider) assert result.returncode == 0 and combined, ( f"provider={provider} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"