Support managed Unity Catalog model locations

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

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,13 +865,21 @@ 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,
                     [tool_name],
                     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,20 @@ 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)
+        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.
         # 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)
 
@@ -2283,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))
@@ -2318,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}."
@@ -2330,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."
@@ -2411,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
@@ -2491,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

src/ucode/managed_resolve.py

@@ -169,11 +169,19 @@ 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, 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)
-    if _str(model_config.get("model_provider_service")) or _str(model_config.get("default_model")):
+    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 managed_location
+        or _str(model_config.get("default_model"))
+    ):
         return True
     models = model_config.get("models")
     if isinstance(models, dict):
@@ -188,6 +196,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``.
 

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 = []