diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3a5d8c40..bfda0b68 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -42,6 +42,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 is_dry_run, restore_file, set_dry_run +from ucode.constants import MODEL_DISCOVERY_ENV_VAR from ucode.custom_oauth import ( CUSTOM_OAUTH_CLI_ENV_VAR, custom_oauth_cli_enabled, @@ -87,6 +88,7 @@ managed_default_model, managed_enabled_tools, managed_launch_model, + managed_model_discovery_enabled, managed_provider_family_models, managed_provider_service, managed_supplies_models, @@ -2138,6 +2140,24 @@ def _managed_smart_routing_enabled(managed: dict | None, tool: str) -> bool: return agent_config.get("smart_routing_enabled") is True +@contextmanager +def _managed_model_discovery_environment(managed: dict | None, tool: str) -> Iterator[None]: + """Expose managed model discovery only to the launched agent.""" + existed = MODEL_DISCOVERY_ENV_VAR in os.environ + previous = os.environ.get(MODEL_DISCOVERY_ENV_VAR) + os.environ[MODEL_DISCOVERY_ENV_VAR] = ( + "1" if managed_model_discovery_enabled(managed, tool) else "0" + ) + try: + yield + finally: + if existed: + assert previous is not None + os.environ[MODEL_DISCOVERY_ENV_VAR] = previous + else: + os.environ.pop(MODEL_DISCOVERY_ENV_VAR, None) + + def _launch_tool( tool_name: str, ctx: typer.Context, @@ -2211,6 +2231,21 @@ def _launch_tool( 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) + managed_provider = managed_provider_service(managed or {}, tool) + if managed_provider: + if explicit_provider is not None: + raise RuntimeError( + f"--provider cannot be used for {TOOL_SPECS[tool]['display']} because your admin " + f"has configured managed provider {managed_provider}." + ) + if parent_schema is not None: + raise RuntimeError( + f"--model-location cannot be used for {TOOL_SPECS[tool]['display']} because your " + f"admin has configured managed provider {managed_provider}." + ) + # Admin config cannot contain both model sources; this only clears local state. + if parent_schema is not None: + provider = None # The environment switch remains a developer override; managed config is the workspace # policy equivalent and must take effect before launch options are computed. managed_smart_routing_enabled = _managed_smart_routing_enabled(managed, tool) @@ -2249,20 +2284,8 @@ def _launch_tool( elif not coding_agent_config_feature_disabled: 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: - raise RuntimeError("--provider and --model-location cannot be used together.") # Checked after the managed config settles `provider`: an admin-set provider must trip this # guard too, or routing would be persisted as on while a provider is active. if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider: @@ -2419,7 +2442,10 @@ def _launch_tool( provider=provider, ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") - with _managed_smart_routing_environment(managed, tool): + with ( + _managed_smart_routing_environment(managed, tool), + _managed_model_discovery_environment(managed, tool), + ): launch_agent(tool, state, ctx.args, options=launch_options) except RuntimeError as exc: print_err(str(exc)) diff --git a/src/ucode/constants.py b/src/ucode/constants.py index c129da93..28790cb6 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -6,6 +6,9 @@ MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" MODEL_SERVICE_PARENT_SCHEMA_HEADER = "Databricks-Model-Service-Parent-Schema" +MODEL_DISCOVERY_ENV_VAR = "UG_ENABLE_MODEL_DISCOVERY" + + # MCP server registration scopes. Claude Code supports local/project/user; the # other CLIs only take the user-scope name. Kept here (a leaf module) so both # `ucode.mcp` and `ucode.agents.claude` can import them without an import cycle. diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 81f1d029..41e1965f 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -188,6 +188,11 @@ def managed_provider_service(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("model_provider_service")) +def managed_model_discovery_enabled(managed: dict | None, tool: str) -> bool: + """Whether managed config enables model discovery for ``tool`` through an MPS.""" + return bool(managed_provider_service(managed or {}, tool)) + + def managed_static_models(managed: dict, tool: str) -> list[str] | None: """The explicit model allow-list (``model_config.model_services``) the config sets for ``tool``. diff --git a/tests/test_cli.py b/tests/test_cli.py index 8effbbd4..6ccd7042 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4114,6 +4114,56 @@ def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): assert mock_shared.call_args.kwargs["skip_model_discovery"] is False +class TestManagedModelDiscoveryLaunch: + MPS = { + "enabled_agents": { + "claude": {"model_config": {"model_provider_service": "main.default.mps"}} + } + } + STATIC = { + "enabled_agents": {"claude": {"model_config": {"model_services": ["system.ai.claude"]}}} + } + + @pytest.mark.parametrize(("managed", "expected"), [(MPS, "1"), (STATIC, "0")]) + def test_sets_literal_value_and_restores_prior(self, monkeypatch, managed, expected): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "prior-value") + + with cli_mod._managed_model_discovery_environment(managed, "claude"): + assert os.environ["UG_ENABLE_MODEL_DISCOVERY"] == expected + + assert os.environ["UG_ENABLE_MODEL_DISCOVERY"] == "prior-value" + + def test_restores_absent_environment_after_launch_error(self, monkeypatch): + monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False) + + with pytest.raises(RuntimeError, match="launch failed"): + with cli_mod._managed_model_discovery_environment(self.MPS, "claude"): + raise RuntimeError("launch failed") + + assert "UG_ENABLE_MODEL_DISCOVERY" not in os.environ + + @pytest.mark.parametrize( + ("flag", "value"), + [("--provider", "main.default.other"), ("--model-location", "main.models")], + ) + def test_managed_provider_rejects_explicit_routing_flags(self, flag, value): + state = dict(MINIMAL_STATE) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(self.MPS, False)), + patch( + "ucode.cli.configure_shared_state", + side_effect=AssertionError("managed source flag was not rejected"), + ), + ): + result = runner.invoke(app, ["claude", flag, value]) + + assert result.exit_code == 1 + assert flag in _strip_ansi(result.output) + + class TestBareUcode: """Bare `ucode` launches the managed default agent, or explains why it can't.""" diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index c4a5ef13..6c9a260d 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -14,6 +14,7 @@ managed_default_model, managed_enabled_tools, managed_launch_model, + managed_model_discovery_enabled, managed_otel_tracing_enabled, managed_provider_service, managed_state_overrides, @@ -176,6 +177,45 @@ def test_none_for_agent_not_in_manifest(self): assert managed_provider_service(MANAGED, "gemini") is None +class TestManagedModelDiscovery: + def test_enabled_for_selected_agent_with_provider_service(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"model_provider_service": "main.default.mps"}} + } + } + assert managed_model_discovery_enabled(managed, "claude") is True + + @pytest.mark.parametrize( + "managed", + [ + None, + {}, + {"enabled_agents": []}, + {"enabled_agents": {"claude": []}}, + {"enabled_agents": {"claude": {"model_config": []}}}, + {"enabled_agents": {"claude": {"model_config": {"model_provider_service": " "}}}}, + ], + ) + def test_disabled_for_missing_or_malformed_config(self, managed): + assert managed_model_discovery_enabled(managed, "claude") is False + + def test_disabled_for_static_model_config(self): + managed = { + "enabled_agents": {"claude": {"model_config": {"model_services": ["system.ai.claude"]}}} + } + assert managed_model_discovery_enabled(managed, "claude") is False + + def test_provider_for_another_agent_does_not_enable_discovery(self): + managed = { + "enabled_agents": { + "codex": {"model_config": {"model_provider_service": "main.default.mps"}}, + "claude": {"model_config": {"model_services": ["system.ai.claude"]}}, + } + } + assert managed_model_discovery_enabled(managed, "claude") is False + + class TestResolveState: def test_does_not_mutate_input_state(self): # managed-state.json and state.json stay separate files: resolution is per-write and