Derive UG_ENABLE_MODEL_DISCOVERY from managed MPS

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))

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.

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``.
 

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."""
 

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