diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index f273e25d7..4825c978b 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -524,9 +524,11 @@ def _availability_failure_detail(tool: str, state: dict) -> str: return " (" + "; ".join(parts) + ")" -def configure_single_tool(tool: str, state: dict, *, parent_schema: str | None = None) -> dict: +def configure_single_tool( + tool: str, state: dict, *, parent_schema: str | None = None, provider: str | None = None +) -> dict: """Check availability, configure, and persist state for one tool only.""" - provider = None if parent_schema else get_provider_service(state, tool) + provider = None if parent_schema else provider or get_provider_service(state, tool) # A Model Provider Service or parent schema routes through the same gateway and pins no # globally discovered Databricks model, so the availability check doesn't apply. if not provider and not parent_schema: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 11f53876e..3304af050 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2447,7 +2447,9 @@ def claude_router_hook_cmd( sys.stdout.write(json.dumps(output)) -def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = None) -> None: +def _auto_configure_tool( + tool: str, custom_oauth: CustomOAuthConfig | None = None, provider: str | None = None +) -> None: """Configure a tool for launch without sending a separate validation prompt. The real agent session follows immediately; explicit configure retains the @@ -2460,8 +2462,14 @@ def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = Non workspace, profile = _prompt_for_configuration(tool) configure_kwargs = {"custom_oauth": custom_oauth} if custom_oauth is not None else {} state = configure_shared_state(workspace, profile=profile, tools=[tool], **configure_kwargs) + # A managed config rejects --provider later in the launch; resolving it here would replace that + # error with a lookup failure. + if provider is not None: + managed, _ = _fetch_managed_config(state) + if managed is not None: + provider = None - state = configure_single_tool(tool, state) + state = configure_single_tool(tool, state, provider=provider) spec = TOOL_SPECS[tool] console.print( @@ -2862,10 +2870,12 @@ def _launch_tool( skip_cli_version_check=skip_preflight, ) if needs_auto_configure: + # Configure through this launch's --provider: a default-gateway setup would write the + # managed auth keys that a relayed provider's launch then refuses to start over. if custom_oauth is None: - _auto_configure_tool(tool) + _auto_configure_tool(tool, provider=provider) else: - _auto_configure_tool(tool, custom_oauth=custom_oauth) + _auto_configure_tool(tool, custom_oauth=custom_oauth, provider=provider) state = ensure_provider_state(tool) # Remembered before the fallback below collapses the two cases: a managed config may not # silently override a provider the user typed on the command line (it errors instead). diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index fd234e8b4..fd27a1747 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -206,6 +206,24 @@ def test_managed_parent_skips_global_availability_and_writes_header(self, monkey configure.assert_called_once_with("claude", state, parent_schema="main.default") + def test_explicit_relayed_provider_configures_relayed(self, monkeypatch): + state = {"workspace": "https://x.databricks.com"} + monkeypatch.setattr( + agents_mod, + "check_gateway_endpoint", + lambda *_a: pytest.fail("a provider must not require global model availability"), + ) + monkeypatch.setattr(agents_mod, "resolve_provider_models", lambda *_a: (None, None, True)) + configure = MagicMock(return_value=state) + monkeypatch.setattr(agents_mod, "configure_tool", configure) + monkeypatch.setattr(agents_mod, "save_state", lambda _state: None) + + agents_mod.configure_single_tool("claude", state, provider="main.default.relay") + + configure.assert_called_once_with( + "claude", state, None, provider="main.default.relay", provider_models=None, relayed=True + ) + def test_configure_selected_tools_triggers_install(self, monkeypatch): captured = self._stub_configure(monkeypatch) agents_mod.configure_selected_tools( diff --git a/tests/test_cli.py b/tests/test_cli.py index e58df6640..1d26f0f75 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -25,7 +25,7 @@ import ucode.databricks as db_mod from ucode.cli import app from ucode.databricks import GatewayProbe -from ucode.managed_config import normalize_managed_config +from ucode.managed_config import ManagedConfigResult, normalize_managed_config _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -1153,6 +1153,23 @@ def test_managed_config_rejects_launch_source_options(self, tool, option, value) assert "`--provider` or `--model-location` is not allowed" in _strip_ansi(result.output) calls["launch"].assert_not_called() + def test_fresh_state_still_rejects_provider_under_managed_config(self): + with ( + _launch_policy_patches({}) as calls, + patch("ucode.cli.load_state", return_value={}), + patch( + "ucode.cli._prompt_for_configuration", + return_value=(MINIMAL_STATE["workspace"], None), + ), + patch("ucode.cli.configure_single_tool", return_value=MINIMAL_STATE) as mock_configure, + ): + result = runner.invoke(app, ["claude", "--provider", "main.default.provider"]) + + assert result.exit_code == 1 + assert "`--provider` or `--model-location` is not allowed" in _strip_ansi(result.output) + mock_configure.assert_called_once_with("claude", MINIMAL_STATE, provider=None) + calls["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 calls: result = runner.invoke(app, ["claude"]) @@ -3043,7 +3060,7 @@ def test_launch_autoconfigures_without_test_prompt(self, tool, has_workspace): result = runner.invoke(app, [tool]) assert result.exit_code == 0, result.output - mock_configure.assert_called_once_with(tool, configured_state) + mock_configure.assert_called_once_with(tool, configured_state, provider=None) mock_restore.assert_not_called() mock_launch.assert_called_once() assert mock_launch.call_args.args[:2] == (tool, configured_state) @@ -3072,7 +3089,7 @@ def test_triggers_when_no_workspace(self): result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output mock_bootstrap.assert_called_once_with("claude", skip_cli_version_check=False) - mock_auto.assert_called_once_with("claude") + mock_auto.assert_called_once_with("claude", provider=None) def test_triggers_when_tool_not_in_available_tools(self): """Auto-configure runs when workspace exists but the tool wasn't configured.""" @@ -3097,7 +3114,41 @@ def test_triggers_when_tool_not_in_available_tools(self): result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output mock_bootstrap.assert_called_once_with("claude", skip_cli_version_check=False) - mock_auto.assert_called_once_with("claude") + mock_auto.assert_called_once_with("claude", provider=None) + + def test_first_run_configures_through_launch_provider(self): + """A default-gateway first-run setup would write the managed auth keys that a relayed + provider's launch then refuses to start over.""" + with ( + _launch_policy_patches(None) as calls, + patch("ucode.cli.load_state", return_value={}), + patch("ucode.cli._auto_configure_tool") as mock_auto, + ): + calls["resolve_provider"].return_value = (None, None, True) + result = runner.invoke(app, ["claude", "--provider", "main.default.relay"]) + + assert result.exit_code == 0, result.output + mock_auto.assert_called_once_with("claude", provider="main.default.relay") + assert calls["configure"].call_args.kwargs["relayed"] is True + + @pytest.mark.parametrize( + ("manifest", "expected_provider"), + [(None, "main.default.relay"), ({"claude": {}}, None)], + ) + def test_auto_configure_forwards_provider_unless_managed(self, manifest, expected_provider): + # A managed config must still reject --provider with its own error later in the launch. + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), + patch( + "ucode.cli._fetch_managed_config", + return_value=ManagedConfigResult(manifest, False), + ), + patch("ucode.cli.configure_single_tool", return_value=MINIMAL_STATE) as mock_configure, + ): + cli_mod._auto_configure_tool("claude", provider="main.default.relay") + + mock_configure.assert_called_once_with("claude", MINIMAL_STATE, provider=expected_provider) def test_skipped_when_already_configured(self): """Auto-configure is skipped when workspace and tool are already set up."""