From 33ca468aba05abe68b030075a2d3d1f481acb485 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 07:39:50 +0000 Subject: [PATCH 1/4] Add managed model location support --- src/ucode/cli.py | 51 ++++++++++++++++++++++-------------- src/ucode/managed_resolve.py | 16 ++++++++--- 2 files changed, 44 insertions(+), 23 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 52b58965..19efc9dc 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -82,6 +82,7 @@ managed_default_model, managed_enabled_tools, managed_launch_model, + managed_model_location, managed_provider_family_models, managed_provider_service, managed_static_models, @@ -797,11 +798,15 @@ def configure_workspace_command( clear_custom_oauth=custom_oauth is None, ) state = states[0] - managed = None + managed, _ = refresh_managed_config(state) if model_location is not None: - managed, _ = refresh_managed_config(state) _reject_configure_model_location(managed, [tool]) - if model_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: + admin_location = managed_model_location(managed or {}, tool) + if admin_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: + state = _configure_location_backed_tool( + resolve_state(managed or {}, state, tool), tool, admin_location + ) + elif model_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: if managed is not None: state = resolve_state(managed, state, tool) state = _configure_tools_with_model_location( @@ -860,6 +865,10 @@ def configure_workspace_command( ] for tool_name in tools_to_configure: resolved = resolve_state(managed, developer_state, tool_name) + admin_location = managed_model_location(managed, tool_name) + location_backed = ( + admin_location is not None and tool_name in CAN_USE_CACHED_CONFIG_AGENTS + ) if tool_name in fallback_location_tools: configured = _configure_tools_with_model_location( resolved, @@ -867,6 +876,10 @@ def configure_workspace_command( model_location, install_ai_tools=not is_dry_run(), ) + elif location_backed: + configured = _configure_location_backed_tool(resolved, tool_name, admin_location) + if not is_dry_run(): + install_databricks_ai_tools_for_agents([tool_name], configured) elif check_gateway_endpoint(developer_state, tool_name): configured = configure_selected_tools( resolved, [tool_name], install_ai_tools=not is_dry_run() @@ -967,12 +980,12 @@ def _state_with_model_location(state: dict, tool: str, location: str | None) -> return candidate -def _configure_model_location(state: dict, tools: list[str], location: str | None) -> dict: - """Rewrite selected Claude/Codex configs with the persisted model-location scope.""" - if location is None: - return state - for tool in tools: - state = configure_tool(tool, state, parent_schema=location) +def _configure_location_backed_tool(state: dict, tool: str, location: str) -> dict: + """Configure one agent at ``location`` and persist only ordinary developer state.""" + state = configure_tool(tool, state, parent_schema=location) + existing = state.get("available_tools") or [] + state["available_tools"] = sorted(set(existing) | {tool}) + save_state(state) return state @@ -1004,10 +1017,7 @@ def _configure_tools_with_model_location( state = configure_selected_tools(state, regular_tools, install_ai_tools=False) for tool in scoped_tools: candidate = _state_with_model_location(state, tool, location) - state = _configure_model_location(candidate, [tool], location) - existing = state.get("available_tools") or [] - state["available_tools"] = sorted(set(existing) | {tool}) - save_state(state) + state = _configure_location_backed_tool(candidate, tool, location) if install_ai_tools: install_databricks_ai_tools_for_agents(tools, state) return state @@ -1930,6 +1940,7 @@ def _auto_configure_tool( tool: str, custom_oauth: CustomOAuthConfig | None = None, model_location: str | None = None, + managed_config: dict | None = None, explicit_provider: str | None = None, ) -> tuple[dict | None, bool]: """Configure a tool for launch without sending a separate validation prompt. @@ -1952,7 +1963,7 @@ def _auto_configure_tool( configure_kwargs["persist"] = False state = configure_shared_state(workspace, profile=profile, tools=[tool], **configure_kwargs) - managed = None + managed = managed_config coding_agent_config_feature_disabled = False if prompted_first_run: managed, coding_agent_config_feature_disabled = refresh_managed_config(state) @@ -1963,16 +1974,16 @@ def _auto_configure_tool( explicit_provider=explicit_provider, explicit_model_location=model_location is not None, ) - - if model_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: + admin_location = managed_model_location(managed or {}, tool) + effective_location = admin_location or model_location + if effective_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: # This is a launch-scoped choice, not an explicit `ug configure` preference. # Write the agent config needed by the imminent session and remember only # that the agent is available; a later bare launch must not inherit this # one-shot location. - state = configure_tool(tool, state, parent_schema=model_location) - existing_tools = state.get("available_tools") or [] - state["available_tools"] = sorted(set(existing_tools) | {tool}) - save_state(state) + if admin_location is not None: + state = resolve_state(managed or {}, state, tool) + state = _configure_location_backed_tool(state, tool, effective_location) else: state = configure_single_tool(tool, state) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 81f1d029..d0e06e8e 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -169,11 +169,16 @@ def managed_supplies_models(managed: dict | None, tool: str) -> bool: """True when the managed config already says which models ``tool`` should use. Lets the launch path skip Databricks model discovery, whose whole purpose is to find the models - the config has now specified. Any of the three counts: a provider (the agent routes by header and - pins no Databricks model), a ``default_model``, or at least one entry in ``models``. + the config has now specified. Any of the four counts: a provider or Unity Catalog location (the + agent routes by header and pins no Databricks model), a ``default_model``, or at least one entry + in ``models``. """ model_config = _agent_model_config(managed or {}, tool) - if _str(model_config.get("model_provider_service")) or _str(model_config.get("default_model")): + if ( + _str(model_config.get("model_provider_service")) + or _str(model_config.get("unity_catalog_location")) + or _str(model_config.get("default_model")) + ): return True models = model_config.get("models") if isinstance(models, dict): @@ -188,6 +193,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_location(managed: dict, tool: str) -> str | None: + """Return the admin-selected Unity Catalog model location for ``tool``, if any.""" + return _str(_agent_model_config(managed, tool).get("unity_catalog_location")) + + 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``. From 0af9877bbf0fb48c3ea72256a746a5fac1a08f31 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 08:29:17 +0000 Subject: [PATCH 2/4] Complete managed model location launch support --- src/ucode/cli.py | 66 ++++++-- src/ucode/managed_resolve.py | 11 +- tests/test_cli.py | 306 ++++++++++++++++++++++++++++++++++ tests/test_managed_resolve.py | 32 ++++ 4 files changed, 395 insertions(+), 20 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 19efc9dc..29c1a91c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1974,7 +1974,11 @@ def _auto_configure_tool( explicit_provider=explicit_provider, explicit_model_location=model_location is not None, ) - admin_location = managed_model_location(managed or {}, tool) + admin_location = ( + managed_model_location(managed or {}, tool) + if tool in CAN_USE_CACHED_CONFIG_AGENTS + else None + ) effective_location = admin_location or model_location if effective_location is not None and tool in CAN_USE_CACHED_CONFIG_AGENTS: # This is a launch-scoped choice, not an explicit `ug configure` preference. @@ -2294,11 +2298,7 @@ def _managed_smart_routing_enabled(managed: dict | None, tool: str) -> bool: def _managed_controls_model_source(managed: dict | None, tool: str) -> bool: - """Whether the managed config selects a provider or Hosted/static models for ``tool``. - - Managed ``unity_catalog_location`` intentionally remains outside this PR; the downstream - managed-location change owns interpreting and enforcing that source. - """ + """Whether managed config selects a provider, location, or Hosted/static models.""" if managed is None: return False return managed_supplies_models(managed, tool) or bool(managed_static_models(managed, tool)) @@ -2329,8 +2329,18 @@ def _reject_managed_source_override( return display = TOOL_SPECS[tool]["display"] managed_provider = managed_provider_service(managed or {}, tool) + managed_location = ( + managed_model_location(managed or {}, tool) + if tool in CAN_USE_CACHED_CONFIG_AGENTS + else None + ) if explicit_model_location: - source = f"provider {managed_provider}" if managed_provider else "Hosted/static models" + if managed_provider: + source = f"provider {managed_provider}" + elif managed_location: + source = f"model location {managed_location}" + else: + source = "Hosted/static models" raise RuntimeError( f"You cannot launch {display} with --model-location because your admin has " f"specified managed {source}." @@ -2341,6 +2351,11 @@ def _reject_managed_source_override( f"You cannot launch {display} with provider {explicit_provider} because your " f"admin has specified managed provider {managed_provider}." ) + if managed_location: + raise RuntimeError( + f"You cannot launch {display} with provider {explicit_provider} because your " + f"admin has specified managed model location {managed_location}." + ) raise RuntimeError( f"You cannot launch {display} with provider {explicit_provider} because your admin " "has specified managed Hosted/static models." @@ -2422,24 +2437,35 @@ def _launch_tool( if target_workspace is not None: set_current_workspace(target_workspace) if needs_auto_configure: + managed_auto_kwargs = {"managed_config": managed} if managed is not None else {} if custom_oauth is not None and parent_schema is not None: auto_managed = _auto_configure_tool( - tool, custom_oauth=custom_oauth, model_location=parent_schema + tool, + custom_oauth=custom_oauth, + model_location=parent_schema, + **managed_auto_kwargs, ) elif custom_oauth is not None and explicit_provider is not None: auto_managed = _auto_configure_tool( tool, custom_oauth=custom_oauth, explicit_provider=explicit_provider, + **managed_auto_kwargs, ) elif custom_oauth is not None: - auto_managed = _auto_configure_tool(tool, custom_oauth=custom_oauth) + auto_managed = _auto_configure_tool( + tool, custom_oauth=custom_oauth, **managed_auto_kwargs + ) elif parent_schema is not None: - auto_managed = _auto_configure_tool(tool, model_location=parent_schema) + auto_managed = _auto_configure_tool( + tool, model_location=parent_schema, **managed_auto_kwargs + ) elif explicit_provider is not None: - auto_managed = _auto_configure_tool(tool, explicit_provider=explicit_provider) + auto_managed = _auto_configure_tool( + tool, explicit_provider=explicit_provider, **managed_auto_kwargs + ) else: - auto_managed = _auto_configure_tool(tool) + auto_managed = _auto_configure_tool(tool, **managed_auto_kwargs) if not existing.get("workspace"): managed, coding_agent_config_feature_disabled = auto_managed managed_config_checked = True @@ -2502,11 +2528,19 @@ 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) + managed_location = ( + managed_model_location(managed, tool) + if tool in CAN_USE_CACHED_CONFIG_AGENTS + else None + ) if _managed_controls_model_source(managed, tool): - # The managed source outranks saved developer preferences. Managed - # unity_catalog_location remains intentionally out of scope. - provider = managed_provider - parent_schema = None + # The managed source outranks saved developer preferences. + if managed_location is not None: + provider = None + parent_schema = managed_location + else: + provider = managed_provider + parent_schema = None 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 diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index d0e06e8e..8eb864b3 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -169,14 +169,17 @@ def managed_supplies_models(managed: dict | None, tool: str) -> bool: """True when the managed config already says which models ``tool`` should use. Lets the launch path skip Databricks model discovery, whose whole purpose is to find the models - the config has now specified. Any of the four counts: a provider or Unity Catalog location (the - agent routes by header and pins no Databricks model), a ``default_model``, or at least one entry - in ``models``. + the config has now specified. Any of the four counts: a provider, a supported Claude/Codex + Unity Catalog location (the agent routes by header and pins no Databricks model), a + ``default_model``, or at least one entry in ``models``. """ model_config = _agent_model_config(managed or {}, tool) + managed_location = ( + _str(model_config.get("unity_catalog_location")) if tool in ("claude", "codex") else None + ) if ( _str(model_config.get("model_provider_service")) - or _str(model_config.get("unity_catalog_location")) + or managed_location or _str(model_config.get("default_model")) ): return True diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c09ed71..0d9df75a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1572,6 +1572,122 @@ def test_managed_hosted_source_overrides_saved_provider_for_bare_launch( assert state["provider_services"][tool] == saved_provider +class TestManagedModelLocationLaunch: + @staticmethod + def _managed(tool: str) -> dict: + return { + "enabled_agents": { + tool: {"model_config": {"unity_catalog_location": "main.managed_models"}} + } + } + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_admin_location_overrides_saved_source_and_forces_scoped_discovery( + self, monkeypatch, tool + ): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + state = { + **MINIMAL_STATE, + "provider_services": {tool: f"main.user.{tool}"}, + "model_locations": {tool: "main.user_models"}, + "claude_models": {}, + "codex_models": [], + } + original_state = json.loads(json.dumps(state)) + managed = self._managed(tool) + 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=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.configure_shared_state", return_value=state) as mock_shared, + patch("ucode.cli.resolve_launch_model") as mock_resolve_model, + patch("ucode.cli.resolve_provider_models") as mock_resolve_provider, + patch( + "ucode.cli.configure_tool", side_effect=lambda *args, **kwargs: args[1] + ) as mock_configure, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, [tool]) + + assert result.exit_code == 0, result.output + assert mock_shared.call_args.kwargs["skip_model_discovery"] is True + mock_resolve_model.assert_not_called() + mock_resolve_provider.assert_not_called() + assert mock_configure.call_args.args[2] is None + assert mock_configure.call_args.kwargs["provider"] is None + assert mock_configure.call_args.kwargs["parent_schema"] == "main.managed_models" + launch_state = mock_launch.call_args.args[1] + assert launch_state[f"_{tool}_launch_parent_schema"] == "main.managed_models" + assert launch_state[f"_{tool}_scoped_model_discovery"] is True + assert not launch_state.get("provider_services", {}).get(tool) + assert launch_state["model_locations"][tool] == "main.user_models" + assert "main.managed_models" not in launch_state.get("model_locations", {}).values() + assert "_managed_overlay" not in launch_state + assert state == original_state + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + @pytest.mark.parametrize( + ("option", "expected"), + [ + (["--provider", "main.user.provider"], "provider main.user.provider"), + (["--model-location", "main.user_models"], "--model-location"), + ], + ) + def test_admin_location_rejects_explicit_source_when_discovery_is_disabled( + self, monkeypatch, tool, option, expected + ): + monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + managed = self._managed(tool) + 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.configure_shared_state") as mock_shared, + patch("ucode.cli.configure_tool") as mock_configure, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, [tool, *option]) + + output = _strip_ansi(result.output) + assert result.exit_code == 1 + assert expected in output + assert "managed model location main.managed_models" in output + mock_shared.assert_not_called() + mock_configure.assert_not_called() + mock_launch.assert_not_called() + + def test_gemini_location_does_not_skip_global_discovery_or_create_scope(self): + state = dict(MINIMAL_STATE) + managed = self._managed("gemini") + 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=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.configure_shared_state", return_value=state) as mock_shared, + patch( + "ucode.cli.resolve_launch_model", return_value=(state, "system.ai.gemini") + ) as mock_resolve, + patch( + "ucode.cli.configure_tool", side_effect=lambda *args, **kwargs: args[1] + ) as mock_configure, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["gemini"]) + + assert result.exit_code == 0, result.output + assert mock_shared.call_args.kwargs["skip_model_discovery"] is False + mock_resolve.assert_called_once() + assert mock_configure.call_args.kwargs["parent_schema"] is None + launch_state = mock_launch.call_args.args[1] + assert not any("scoped_model_discovery" in key for key in launch_state) + assert not any("launch_parent_schema" in key for key in launch_state) + + class TestGeminiProviderLaunch: @staticmethod def _launch(monkeypatch, resolve_provider_models): @@ -2783,6 +2899,107 @@ def test_model_location_autoconfigure_is_launch_scoped(self, tool): assert saved_states[-1]["available_tools"] == [tool] assert "model_locations" not in saved_states[-1] + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_unconfigured_agent_reuses_fetched_managed_location_snapshot(self, tool): + existing_state = {**MINIMAL_STATE, "available_tools": []} + configured_state = { + **MINIMAL_STATE, + "available_tools": [tool], + "claude_models": {}, + "codex_models": [], + } + managed = { + "enabled_agents": { + tool: {"model_config": {"unity_catalog_location": "main.managed_models"}} + } + } + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=existing_state), + patch("ucode.cli._fetch_managed_config", return_value=(managed, False)) as mock_fetch, + patch("ucode.cli._auto_configure_tool", return_value=(managed, False)) as mock_auto, + patch("ucode.cli.ensure_provider_state", return_value=configured_state), + patch("ucode.cli.configure_shared_state", return_value=configured_state), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.resolve_launch_model") as mock_resolve, + patch( + "ucode.cli.configure_tool", side_effect=lambda *args, **kwargs: args[1] + ) as mock_configure, + patch("ucode.cli.launch_agent"), + ): + result = runner.invoke(app, [tool]) + + assert result.exit_code == 0, result.output + mock_fetch.assert_called_once_with(existing_state) + mock_auto.assert_called_once_with(tool, managed_config=managed) + mock_resolve.assert_not_called() + assert mock_configure.call_args.kwargs["parent_schema"] == "main.managed_models" + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_true_first_run_configures_managed_location_without_persisting_policy(self, tool): + configured_state = { + **MINIMAL_STATE, + "available_tools": [tool], + "claude_models": {}, + "codex_models": [], + } + original_state = json.loads(json.dumps(configured_state)) + managed = { + "enabled_agents": { + tool: {"model_config": {"unity_catalog_location": "main.managed_models"}} + } + } + persisted: list[dict] = [] + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value={}), + patch( + "ucode.cli._prompt_for_configuration", + return_value=(MINIMAL_STATE["workspace"], None), + ), + patch( + "ucode.cli.configure_shared_state", + side_effect=[configured_state, configured_state], + ) as mock_shared, + patch( + "ucode.cli.refresh_managed_config", return_value=(managed, False) + ) as mock_refresh, + patch("ucode.cli._fetch_managed_config") as mock_fetch, + patch("ucode.cli.ensure_provider_state", return_value=configured_state), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.resolve_launch_model") as mock_resolve, + patch( + "ucode.cli.configure_tool", side_effect=lambda *args, **kwargs: args[1] + ) as mock_configure, + patch( + "ucode.cli.save_state", + side_effect=lambda state: persisted.append( + cli_mod.developer_state_from_resolved(state) + ), + ), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, [tool]) + + assert result.exit_code == 0, result.output + assert mock_refresh.call_count == 1 + mock_fetch.assert_not_called() + first_shared, launch_shared = mock_shared.call_args_list + assert first_shared.kwargs["persist"] is False + assert "skip_model_discovery" not in first_shared.kwargs + assert launch_shared.kwargs["skip_model_discovery"] is True + assert [call.kwargs["parent_schema"] for call in mock_configure.call_args_list] == [ + "main.managed_models", + "main.managed_models", + ] + mock_resolve.assert_not_called() + assert persisted[-1]["available_tools"] == [tool] + assert "model_locations" not in persisted[-1] + assert "_managed_overlay" not in persisted[-1] + assert configured_state == original_state + launch_state = mock_launch.call_args.args[1] + assert launch_state[f"_{tool}_launch_parent_schema"] == "main.managed_models" + def test_triggers_when_no_workspace(self): """Auto-configure runs when state has no workspace.""" empty_state = {} @@ -3918,6 +4135,95 @@ def test_model_location_configures_without_globally_discovered_models(self, monk assert configured_state["available_tools"] == [tool] assert "model_locations" not in state + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_managed_model_location_configures_without_globally_discovered_models( + self, monkeypatch, tool + ): + state = { + **MINIMAL_STATE, + "available_tools": [], + "claude_models": {}, + "codex_models": [], + } + original_state = json.loads(json.dumps(state)) + managed = { + "enabled_agents": { + tool: {"model_config": {"unity_catalog_location": "main.managed_models"}} + } + } + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda current: (managed, False)) + monkeypatch.setattr( + cli_mod, + "check_gateway_endpoint", + lambda *a, **k: pytest.fail("global availability must not gate a managed location"), + ) + monkeypatch.setattr(cli_mod, "_announce_managed_config", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_summarize_managed_config", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", lambda *a, **k: None) + configure_calls = [] + monkeypatch.setattr( + cli_mod, + "configure_tool", + lambda *args, **kwargs: configure_calls.append((args, kwargs)) or args[1], + ) + + assert ( + cli_mod.configure_workspace_command( + selected_tools=[tool], + workspaces=[("https://example.databricks.com", None)], + ) + == 0 + ) + + assert len(configure_calls) == 1 + args, kwargs = configure_calls[0] + assert args[0] == tool + assert kwargs["parent_schema"] == "main.managed_models" + assert args[1]["available_tools"] == [tool] + assert "model_locations" not in args[1] + assert state == original_state + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_single_agent_configure_uses_managed_model_location(self, monkeypatch, tool): + state = { + **MINIMAL_STATE, + "available_tools": [], + "claude_models": {}, + "codex_models": [], + } + original_state = json.loads(json.dumps(state)) + managed = { + "enabled_agents": { + tool: {"model_config": {"unity_catalog_location": "main.managed_models"}} + } + } + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda current: (managed, False)) + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", lambda *a, **k: None) + configure_calls = [] + monkeypatch.setattr( + cli_mod, + "configure_tool", + lambda *args, **kwargs: configure_calls.append((args, kwargs)) or args[1], + ) + + assert ( + cli_mod.configure_workspace_command( + tool=tool, + workspaces=[("https://example.databricks.com", None)], + ) + == 0 + ) + + assert len(configure_calls) == 1 + args, kwargs = configure_calls[0] + assert args[0] == tool + assert kwargs["parent_schema"] == "main.managed_models" + assert args[1]["available_tools"] == [tool] + assert "model_locations" not in args[1] + assert state == original_state + def test_managed_config_detection_does_not_reset_saved_locations(self, monkeypatch): state = { **MINIMAL_STATE, diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index c4a5ef13..54467a33 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_location, managed_otel_tracing_enabled, managed_provider_service, managed_state_overrides, @@ -176,6 +177,22 @@ def test_none_for_agent_not_in_manifest(self): assert managed_provider_service(MANAGED, "gemini") is None +class TestManagedModelLocation: + def test_returns_only_the_normalized_manifest_value(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"unity_catalog_location": " main.models "}} + } + } + before = json.loads(json.dumps(managed)) + + assert managed_model_location(managed, "claude") == "main.models" + assert managed == before + + def test_does_not_fall_back_to_developer_state(self): + assert managed_model_location({}, "claude") is None + + class TestResolveState: def test_does_not_mutate_input_state(self): # managed-state.json and state.json stay separate files: resolution is per-write and @@ -403,6 +420,21 @@ def test_true_for_a_provider(self): } assert managed_supplies_models(managed, "claude") is True + @pytest.mark.parametrize("tool", ["claude", "codex"]) + def test_true_for_a_supported_managed_model_location(self, tool): + managed = { + "enabled_agents": {tool: {"model_config": {"unity_catalog_location": "main.models"}}} + } + assert managed_supplies_models(managed, tool) is True + + def test_false_for_an_unsupported_managed_model_location(self): + managed = { + "enabled_agents": { + "gemini": {"model_config": {"unity_catalog_location": "main.models"}} + } + } + assert managed_supplies_models(managed, "gemini") is False + def test_true_for_a_flat_model_list(self): managed = {"enabled_agents": {"opencode": {"model_config": {"models": ["a", "b"]}}}} assert managed_supplies_models(managed, "opencode") is True From b52e2f48d6b5d4cee582ea6fa07b19b95fba494d Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 08:34:41 +0000 Subject: [PATCH 3/4] Isolate custom OAuth configure test from managed fetch --- tests/test_custom_oauth.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index ed0d945a..33238961 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -258,6 +258,7 @@ def test_configure_without_custom_options_resets_custom_oauth(self): state = {"workspace": WS, "available_tools": ["claude"]} with ( patch("ucode.cli._configure_shared_workspace_states", return_value=[state]) as shared, + patch("ucode.cli.refresh_managed_config", return_value=(None, False)), patch("ucode.cli.configure_single_tool", return_value=state), patch("ucode.cli.install_databricks_ai_tools_for_agents"), ): From 64e02bfa27b3d95259f3a41d1e91e124f6c2107f Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 22:24:47 +0000 Subject: [PATCH 4/4] Keep managed model locations free of global Claude pins --- src/ucode/agents/claude.py | 5 ++- tests/test_agent_claude.py | 22 ++++++++++++ tests/test_cli.py | 69 +++++++++++++++++++++++++++++++++----- 3 files changed, 86 insertions(+), 10 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 2a6333be..1e1cd358 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -869,7 +869,10 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: _reconcile_managed_settings( state, - lambda base: _compose(base, enforce_model_default_hierarchy=provider is None), + lambda base: _compose( + base, + enforce_model_default_hierarchy=provider is None and parent_schema is None, + ), managed_file_keys, relayed, ) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index c5acbe5b..af3531f7 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1027,6 +1027,28 @@ def test_managed_file_omits_workspace_defaults_for_provider(self, monkeypatch): env = json.loads(managed_writes[0][1])["env"] assert not set(claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.values()) & env.keys() + def test_managed_file_omits_workspace_defaults_for_model_location(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": {"ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8"} + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = { + "workspace": WS, + "claude_models": { + "opus": "system.ai.claude-opus-4-8", + "haiku": "system.ai.claude-haiku-4-6", + }, + } + + claude.write_tool_config(state, None, parent_schema="main.managed_models") + + env = json.loads(managed_writes[0][1])["env"] + assert not set(claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.values()) & env.keys() + def test_managed_file_keeps_provider_model_pins(self, monkeypatch): private_writes: list = [] managed_writes: list = [] diff --git a/tests/test_cli.py b/tests/test_cli.py index 0d9df75a..db245d39 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1582,10 +1582,7 @@ def _managed(tool: str) -> dict: } @pytest.mark.parametrize("tool", ["claude", "codex"]) - def test_admin_location_overrides_saved_source_and_forces_scoped_discovery( - self, monkeypatch, tool - ): - monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + def test_admin_location_overrides_saved_source(self, tool): state = { **MINIMAL_STATE, "provider_services": {tool: f"main.user.{tool}"}, @@ -1620,13 +1617,70 @@ def test_admin_location_overrides_saved_source_and_forces_scoped_discovery( assert mock_configure.call_args.kwargs["parent_schema"] == "main.managed_models" launch_state = mock_launch.call_args.args[1] assert launch_state[f"_{tool}_launch_parent_schema"] == "main.managed_models" - assert launch_state[f"_{tool}_scoped_model_discovery"] is True assert not launch_state.get("provider_services", {}).get(tool) assert launch_state["model_locations"][tool] == "main.user_models" assert "main.managed_models" not in launch_state.get("model_locations", {}).values() assert "_managed_overlay" not in launch_state assert state == original_state + def test_admin_claude_location_scopes_native_discovery(self, monkeypatch): + key = cli_mod.claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR + monkeypatch.delenv(key, raising=False) + observed = [] + state = {**MINIMAL_STATE, "claude_models": {}} + managed = self._managed("claude") + + def launch(*_args, **_kwargs): + observed.append(os.environ.get(key)) + + 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=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli.launch_agent", side_effect=launch), + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + assert observed == ["1"] + assert key not in os.environ + + def test_admin_claude_location_blocks_old_native_discovery_version(self, monkeypatch): + key = cli_mod.claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR + monkeypatch.delenv(key, raising=False) + install_environments = [] + state = {**MINIMAL_STATE, "claude_models": {}} + managed = self._managed("claude") + + def install(_tool, *, strict): + install_environments.append((strict, os.environ.get(key))) + raise RuntimeError("Claude Code 2.1.247 could not be upgraded") + + 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=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.claude_agent.agent_version", return_value="2.1.247"), + patch("ucode.cli.install_tool_binary", side_effect=install), + patch("ucode.cli.configure_tool") as mock_configure, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 1 + assert "2.1.247 could not be upgraded" in _strip_ansi(result.output) + assert install_environments == [(True, "1")] + mock_configure.assert_not_called() + mock_launch.assert_not_called() + assert key not in os.environ + @pytest.mark.parametrize("tool", ["claude", "codex"]) @pytest.mark.parametrize( ("option", "expected"), @@ -1635,10 +1689,7 @@ def test_admin_location_overrides_saved_source_and_forces_scoped_discovery( (["--model-location", "main.user_models"], "--model-location"), ], ) - def test_admin_location_rejects_explicit_source_when_discovery_is_disabled( - self, monkeypatch, tool, option, expected - ): - monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0") + def test_admin_location_rejects_explicit_source(self, tool, option, expected): managed = self._managed(tool) with ( patch("ucode.cli.ensure_bootstrap_dependencies"),