From 9b8ec21377a667876e9224618662e72c3c0f005b Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 17 Sep 2026 23:31:26 +0000 Subject: [PATCH] Reject launch source flags under managed config --- src/ucode/cli.py | 35 ++++++++++++------- tests/test_cli.py | 86 ++++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 99 insertions(+), 22 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 287181a9..83d193f8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2046,6 +2046,19 @@ def _reject_disabled_agent(managed: dict | None, tool: str) -> None: ) +def _reject_managed_launch_source_options( + managed: dict | None, + *, + provider: str | None, + parent_schema: str | None, +) -> None: + if managed is not None and (provider is not None or parent_schema is not None): + raise RuntimeError( + "`--provider` or `--model-location` is not allowed when a managed config exists " + "for the workspace; the managed config controls the model source." + ) + + def _fetch_managed_config(state: dict) -> ManagedConfigResult: """The workspace's managed config for this launch, plus whether the feature is disabled. @@ -2267,10 +2280,6 @@ def _launch_tool( # every ug line from here on must go to stderr instead. if _child_owns_stdout(tool, ctx.args): redirect_output_to_stderr() - if provider is not None and parent_schema is not None: - raise RuntimeError("--provider and --model-location cannot be used together.") - if parent_schema is not None and not is_valid_catalog_schema(parent_schema): - raise RuntimeError("--model-location must be `.`.") explicit_prompt = _has_explicit_prompt(ctx) smart_routing_enabled = smart_routing_v2.smart_routing_enabled() # Launchers such as isaac put their harness arguments after `--`, so the harness's own @@ -2317,6 +2326,15 @@ def _launch_tool( coding_agent_config_feature_disabled = False if managed is None: managed, coding_agent_config_feature_disabled = _fetch_managed_config(state) + _reject_managed_launch_source_options( + managed, + provider=explicit_provider, + parent_schema=parent_schema, + ) + if explicit_provider is not None and parent_schema is not None: + raise RuntimeError("--provider and --model-location cannot be used together.") + if parent_schema is not None and not is_valid_catalog_schema(parent_schema): + raise RuntimeError("--model-location must be `.`.") # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast. _reject_disabled_agent(managed, tool) # The environment switch remains a developer override; managed config is the workspace @@ -2358,15 +2376,6 @@ def _launch_tool( print_note("No managed coding agent config found; using your own settings") if managed is not None: managed_provider = managed_provider_service(managed, tool) - if explicit_provider and managed_provider and managed_provider != explicit_provider: - # An explicit --provider that disagrees with the admin's is a hard error rather - # than a silent override: the user asked for something the managed config forbids, - # and quietly routing them elsewhere would hide it. - raise RuntimeError( - f"You cannot launch {TOOL_SPECS[tool]['display']} with provider " - f"{explicit_provider} because your admin has specified managed provider " - f"{managed_provider}." - ) if managed_provider: provider = managed_provider if provider and parent_schema is not None: diff --git a/tests/test_cli.py b/tests/test_cli.py index c138a160..88aedbfd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -446,6 +446,32 @@ def _patch_launch(tool: str): ] +@contextlib.contextmanager +def _launch_policy_patches( + managed: dict | None, + *, + persisted_provider: str | None = None, +): + 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._fetch_managed_config", return_value=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.get_provider_service", return_value=persisted_provider), + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), + patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)), + patch("ucode.cli.resolve_gemini_provider_model", return_value=("gemini-2.0-flash", None)), + 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.launch_agent") as launch, + ): + yield launch + + class TestSubcommandRouting: @pytest.mark.parametrize("tool", TOOLS) def test_subcommand_calls_correct_tool(self, tool): @@ -710,26 +736,41 @@ def test_codex_model_location_is_forwarded(self): assert mock_launch.call_args.args[1].args == [] def test_codex_provider_and_model_location_are_mutually_exclusive(self): - result = runner.invoke( - app, - ["codex", "--provider", "main.default.provider", "--model-location", "main.default"], - ) + with _launch_policy_patches(None): + result = runner.invoke( + app, + [ + "codex", + "--provider", + "main.default.provider", + "--model-location", + "main.default", + ], + ) assert result.exit_code == 1 assert "--provider and --model-location cannot be used together" in result.output def test_claude_provider_and_model_location_are_mutually_exclusive(self): - result = runner.invoke( - app, - ["claude", "--provider", "main.default.provider", "--model-location", "main.default"], - ) + with _launch_policy_patches(None): + result = runner.invoke( + app, + [ + "claude", + "--provider", + "main.default.provider", + "--model-location", + "main.default", + ], + ) assert result.exit_code == 1 assert "--provider and --model-location cannot be used together" in result.output @pytest.mark.parametrize("tool", ["claude", "codex"]) def test_invalid_model_location_is_rejected(self, tool): - result = runner.invoke(app, [tool, "--model-location", "main"]) + with _launch_policy_patches(None): + result = runner.invoke(app, [tool, "--model-location", "main"]) assert result.exit_code == 1 assert "--model-location must be `.`." in _strip_ansi(result.output) @@ -926,6 +967,33 @@ def test_claude_v2_subagent_hook_uses_v2_router(self, monkeypatch): mock_v2_route.assert_called_once() +class TestManagedConfigLaunchSourceGuard: + @pytest.mark.parametrize( + ("tool", "option", "value"), + [ + ("claude", "--provider", "main.default.provider"), + ("codex", "--provider", "main.default.provider"), + ("gemini", "--provider", "main.default.provider"), + ("claude", "--model-location", "main.default"), + ("codex", "--model-location", "main.default"), + ], + ) + def test_managed_config_rejects_launch_source_options(self, tool, option, value): + with _launch_policy_patches({}) as launch: + result = runner.invoke(app, [tool, option, value]) + + assert result.exit_code == 1 + assert "`--provider` or `--model-location` is not allowed" in _strip_ansi(result.output) + launch.assert_not_called() + + def test_persisted_provider_is_not_mistaken_for_an_explicit_option(self): + with _launch_policy_patches({}, persisted_provider="main.default.provider") as launch: + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + launch.assert_called_once() + + class TestClaudeModelFlag: """`ucode claude --model ` pins the id into the family aliases so the gateway resolves any Databricks model id, instead of Claude Code's own --model flag rejecting non-catalog ids."""