From 60e588222376322e8dbedf0228d75d0015ff169f Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 21:57:03 +0000 Subject: [PATCH 1/2] Derive model discovery from managed providers --- src/ucode/cli.py | 50 ++++++++++++++----- src/ucode/constants.py | 3 ++ src/ucode/managed_resolve.py | 5 ++ tests/test_cli.py | 90 +++++++++++++++++++++++++++++++++++ tests/test_managed_resolve.py | 40 ++++++++++++++++ 5 files changed, 175 insertions(+), 13 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index f14d9a109..575430bce 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, @@ -2100,6 +2102,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, @@ -2173,6 +2193,19 @@ 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 and 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 managed_provider and 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}." + ) + 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) @@ -2211,20 +2244,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: @@ -2381,7 +2402,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 c129da93e..28790cb63 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 81f1d0296..41e1965f4 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 cbb5afd27..eeb2e0322 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4053,6 +4053,96 @@ 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"]}}} + } + + @staticmethod + def _run(monkeypatch, argv, managed, *, launch_error=None, local_provider=False): + state = dict(MINIMAL_STATE) + if local_provider: + state["provider_services"] = {"claude": "main.default.local"} + seen: dict[str, str | None] = {} + + def configure_shared(*args, **kwargs): + seen["configure"] = os.environ.get("UG_ENABLE_MODEL_DISCOVERY") + return state + + def launch(*args, **kwargs): + seen["launch"] = os.environ.get("UG_ENABLE_MODEL_DISCOVERY") + if launch_error is not None: + raise launch_error + + monkeypatch.setattr(cli_mod, "ensure_bootstrap_dependencies", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "load_state", lambda: state) + monkeypatch.setattr(cli_mod, "ensure_provider_state", lambda tool: state) + monkeypatch.setattr(cli_mod, "configure_shared_state", configure_shared) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda current: (managed, False)) + monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda current, cfg: None) + monkeypatch.setattr(cli_mod, "resolve_provider_models", lambda *a, **k: (None, None, False)) + monkeypatch.setattr(cli_mod, "resolve_launch_model", lambda *a, **k: (state, "model")) + monkeypatch.setattr(cli_mod, "configure_tool", lambda *a, **k: state) + mock_launch = MagicMock(side_effect=launch) + monkeypatch.setattr(cli_mod, "launch_agent", mock_launch) + + return runner.invoke(app, argv), seen, mock_launch + + @pytest.mark.parametrize(("managed", "expected"), [(MPS, "1"), (STATIC, "0")]) + def test_sets_literal_value_only_during_launch_and_restores_prior( + self, monkeypatch, managed, expected + ): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "prior-value") + + result, seen, _ = self._run(monkeypatch, ["claude"], managed) + + assert result.exit_code == 0, result.output + assert seen == {"configure": "prior-value", "launch": 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) + + result, seen, _ = self._run( + monkeypatch, ["claude"], self.MPS, launch_error=RuntimeError("launch failed") + ) + + assert result.exit_code == 1 + assert seen == {"configure": None, "launch": "1"} + 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, monkeypatch, flag, value): + result, _seen, mock_launch = self._run(monkeypatch, ["claude", flag, value], self.MPS) + + assert result.exit_code == 1 + assert flag in _strip_ansi(result.output) + mock_launch.assert_not_called() + + @pytest.mark.parametrize( + ("flag", "value"), + [("--provider", "main.default.other"), ("--model-location", "main.models")], + ) + def test_static_config_allows_explicit_routing_flags(self, monkeypatch, flag, value): + result, _seen, mock_launch = self._run( + monkeypatch, + ["claude", flag, value], + self.STATIC, + local_provider=flag == "--model-location", + ) + + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + + 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 c4a5ef137..6c9a260dd 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 From a3b48e932716e295034f4acb643450c4adc650d0 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 17 Sep 2026 16:08:03 +0000 Subject: [PATCH 2/2] Address model discovery review feedback --- src/ucode/cli.py | 22 +++++++------ tests/test_cli.py | 78 ++++++++++++----------------------------------- 2 files changed, 31 insertions(+), 69 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 575430bce..e21b9a9e8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2194,16 +2194,18 @@ def _launch_tool( # 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 and 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 managed_provider and 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}." - ) + 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index eeb2e0322..7ffbe7976 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4063,84 +4063,44 @@ class TestManagedModelDiscoveryLaunch: "enabled_agents": {"claude": {"model_config": {"model_services": ["system.ai.claude"]}}} } - @staticmethod - def _run(monkeypatch, argv, managed, *, launch_error=None, local_provider=False): - state = dict(MINIMAL_STATE) - if local_provider: - state["provider_services"] = {"claude": "main.default.local"} - seen: dict[str, str | None] = {} - - def configure_shared(*args, **kwargs): - seen["configure"] = os.environ.get("UG_ENABLE_MODEL_DISCOVERY") - return state - - def launch(*args, **kwargs): - seen["launch"] = os.environ.get("UG_ENABLE_MODEL_DISCOVERY") - if launch_error is not None: - raise launch_error - - monkeypatch.setattr(cli_mod, "ensure_bootstrap_dependencies", lambda *a, **k: None) - monkeypatch.setattr(cli_mod, "load_state", lambda: state) - monkeypatch.setattr(cli_mod, "ensure_provider_state", lambda tool: state) - monkeypatch.setattr(cli_mod, "configure_shared_state", configure_shared) - monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda current: (managed, False)) - monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda current, cfg: None) - monkeypatch.setattr(cli_mod, "resolve_provider_models", lambda *a, **k: (None, None, False)) - monkeypatch.setattr(cli_mod, "resolve_launch_model", lambda *a, **k: (state, "model")) - monkeypatch.setattr(cli_mod, "configure_tool", lambda *a, **k: state) - mock_launch = MagicMock(side_effect=launch) - monkeypatch.setattr(cli_mod, "launch_agent", mock_launch) - - return runner.invoke(app, argv), seen, mock_launch - @pytest.mark.parametrize(("managed", "expected"), [(MPS, "1"), (STATIC, "0")]) - def test_sets_literal_value_only_during_launch_and_restores_prior( - self, monkeypatch, managed, expected - ): + def test_sets_literal_value_and_restores_prior(self, monkeypatch, managed, expected): monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "prior-value") - result, seen, _ = self._run(monkeypatch, ["claude"], managed) + with cli_mod._managed_model_discovery_environment(managed, "claude"): + assert os.environ["UG_ENABLE_MODEL_DISCOVERY"] == expected - assert result.exit_code == 0, result.output - assert seen == {"configure": "prior-value", "launch": 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) - result, seen, _ = self._run( - monkeypatch, ["claude"], self.MPS, launch_error=RuntimeError("launch failed") - ) + with pytest.raises(RuntimeError, match="launch failed"): + with cli_mod._managed_model_discovery_environment(self.MPS, "claude"): + raise RuntimeError("launch failed") - assert result.exit_code == 1 - assert seen == {"configure": None, "launch": "1"} 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, monkeypatch, flag, value): - result, _seen, mock_launch = self._run(monkeypatch, ["claude", flag, value], self.MPS) + 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) - mock_launch.assert_not_called() - - @pytest.mark.parametrize( - ("flag", "value"), - [("--provider", "main.default.other"), ("--model-location", "main.models")], - ) - def test_static_config_allows_explicit_routing_flags(self, monkeypatch, flag, value): - result, _seen, mock_launch = self._run( - monkeypatch, - ["claude", flag, value], - self.STATIC, - local_provider=flag == "--model-location", - ) - - assert result.exit_code == 0, result.output - mock_launch.assert_called_once() class TestBareUcode: