From d4d358ef0db408e67f4c2c1c30811d624a169ce1 Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Wed, 12 Aug 2026 20:35:39 +0000 Subject: [PATCH 1/3] Add --skip-managed-config and decouple managed config from --skip-preflight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --skip-managed-config (on `ucode configure` and every launcher) ignores the workspace's managed coding-agent config for that run, behaving as though ENABLE_MANAGED_AGENT_CONFIG were unset — it clears the env var for the process, so the single managed_agent_config_enabled() gate short-circuits every managed path (fetch/apply, budget read, MCP registration, the configure reject flow, and the bare-`ucode` agent picker). Per-invocation only. --skip-preflight no longer touches managed config: it now means only "skip the per-launch auth + AI Gateway re-validation." Dropped its cache-instead-of-refresh branch, its budget-read and MCP-registration gating, and the bare-`ucode` "no config to pick an agent" error. Co-authored-by: Isaac --- src/ucode/cli.py | 105 ++++++++++++++++++++++++++++++++-------------- tests/test_cli.py | 81 +++++++++++++++++++++++++---------- 2 files changed, 132 insertions(+), 54 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ffaf0ec9..4f6303bc 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -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, @@ -1338,16 +1339,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) @@ -1371,17 +1370,16 @@ 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. """ - # --dry-run resolves the agent from the last saved config alone, so it must not reach the + # No managed config (feature off, --skip-managed-config, or none published) means no budget to + # read. --dry-run resolves the agent from the last saved config alone, so it must not reach the # control plane — mirroring the managed-config read, which is likewise skipped under --dry-run. - if managed is None or skip_preflight or is_dry_run(): + if managed is None or is_dry_run(): return None reason: str | None = None recommendation = None @@ -1526,7 +1524,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. @@ -1548,9 +1546,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) @@ -1726,9 +1722,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) and --dry-run (writes nothing). - if managed is not None and not skip_preflight and not is_dry_run(): + # 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']}") @@ -1744,17 +1740,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[ @@ -1789,6 +1810,7 @@ def default( ), ] = False, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, ) -> None: """Configure and launch coding agents through Databricks AI Gateway. @@ -1798,6 +1820,7 @@ 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, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace @@ -1831,12 +1854,6 @@ 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." - ) # --dry-run avoids the fetch but still applies the last saved config. if dry_run: managed = load_managed_state(current) @@ -1857,7 +1874,7 @@ def _launch_managed_default( 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 ) @@ -1908,6 +1925,7 @@ def codex_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -1925,6 +1943,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) @@ -1965,6 +1984,7 @@ def claude_cmd( ), ] = None, skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, enable_smart_routing_flag: Annotated[ bool, @@ -1982,6 +2002,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) @@ -2001,28 +2022,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) @@ -2163,6 +2204,7 @@ def configure( "still applied.", ), ] = False, + skip_managed_config: SkipManagedConfigOption = False, verbose: Annotated[ str, typer.Option( @@ -2175,6 +2217,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 d0049373..989afc5a 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: @@ -2576,25 +2572,41 @@ def test_dry_run_uses_the_cache_and_does_not_fetch(self, monkeypatch): # 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_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_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.refresh_managed_config", - lambda state: pytest.fail("--skip-preflight must not fetch"), + "ucode.cli._launch_tool", + lambda tool, ctx, **kw: seen.update({"tool": tool, **kw}), ) + result = runner.invoke(app, ["--skip-preflight"]) + 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.load_managed_state", - lambda ws: pytest.fail("--skip-preflight must not read the cache"), + "ucode.cli.refresh_managed_config", + lambda state: pytest.fail("--skip-managed-config must not fetch"), ) - 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 + 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): @@ -2621,6 +2633,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.""" From 8d7152edc7aa587265a83a3dc0d443591523544c Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Thu, 13 Aug 2026 23:29:04 +0000 Subject: [PATCH 2/3] test(e2e): skip provider-launch tests when the account is out of credits The provider-launch e2e tests only prove routing reaches the Model Provider Service; a real "Credit balance is too low" from the provider account is an environmental condition, not a ucode bug, so skip it like the existing no-permission case rather than failing CI. Generalize _skip_if_no_permission -> _skip_if_provider_unusable to cover both. Co-authored-by: Isaac --- tests/test_e2e.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index ebd2ab3f..76b5087e 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}" From c2fadb69481b55e18081ef8c5138dd78a3c45fed Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Fri, 14 Aug 2026 05:46:15 +0000 Subject: [PATCH 3/3] cli: drop the budget-recommendation guard comment --- src/ucode/cli.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index cabe4a4f..aad4f58a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1384,9 +1384,6 @@ def _fetch_budget_recommendation(state: dict, managed: dict | None) -> dict | No Enforcement is server-side, so a failed read only costs the recommendation: the config's own ``default_model`` still applies and the launch proceeds. """ - # No managed config (feature off, --skip-managed-config, or none published) means no budget to - # read. --dry-run resolves the agent from the last saved config alone, so it must not reach the - # control plane — mirroring the managed-config read, which is likewise skipped under --dry-run. if managed is None or is_dry_run(): return None reason: str | None = None