Control Claude scoped model discovery with UG_ENABLE_MODEL_DISCOVERY

README.md

@@ -81,6 +81,13 @@ ug codex --full-auto
 
 All agents route through Databricks AI Gateway using your workspace credentials — no API keys required.
 
+For a scoped Claude launch, ug enables Claude Code's native gateway discovery; Claude Code owns
+any model-cache update after it starts. ug does not fetch or rewrite Claude's private model cache
+before launch. Set `UG_ENABLE_MODEL_DISCOVERY=0` to keep the routing header while using Claude's
+native picker catalog instead. This switch does not disable normal Databricks `system.ai` model
+discovery. A workspace-managed scoped source still enables discovery because administrator policy
+takes precedence over the developer environment switch.
+
 Codex uses the provider ID `Databricks` while keeping the `ucode` profile name.
 Re-run `ug configure --agents codex` to update existing generated configurations.
 This reuses history stored under the exact case-sensitive `Databricks` ID; it does not merge

src/ucode/agents/claude.py

@@ -11,7 +11,8 @@
 import socket
 import subprocess
 import threading
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from pathlib import Path
 from typing import cast
 
@@ -25,11 +26,13 @@
     write_json_file,
 )
 from ucode.constants import (
+    CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY,
     LOOPBACK_HOST,
     MCP_CLEANUP_SCOPES,
     MCP_USER_SCOPE,
     MODEL_PROVIDER_SERVICE_HEADER,
     MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+    scoped_model_discovery_enabled,
 )
 from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command
 from ucode.databricks import (
@@ -60,14 +63,21 @@
     remove_smart_routing_hooks,
     sync_smart_routing_hooks,
 )
-from ucode.state import MANAGED_OVERLAY_KEY, is_tool_managed, mark_tool_managed, save_state
+from ucode.state import (
+    MANAGED_OVERLAY_KEY,
+    get_provider_service,
+    is_tool_managed,
+    mark_tool_managed,
+    save_state,
+)
 from ucode.telemetry import agent_version, ug_version
 from ucode.tracing import tracing_env
 from ucode.ui import print_note, print_success, print_warning
 
 from .args import LaunchOptions, has_explicit_model_arg
 
 GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
+CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
 # If set, Claude Code launches in headless mode instead of the interactive login flow.
 CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR = "CLAUDE_CODE_OAUTH_TOKEN"
 CLAUDE_CONFIG_DIR = Path.home() / ".claude"
@@ -88,6 +98,35 @@
     "backup_path": CLAUDE_BACKUP_PATH,
 }
 
+_MISSING_ENV_VALUE = object()
+
+
+@contextmanager
+def launch_discovery_environment(enabled: bool | None) -> Iterator[None]:
+    """Apply Claude discovery variables only for the lifetime of one launch.
+
+    POSIX launches replace this process, so the successful production path
+    never reaches restoration. Spawned, mocked, and failed launches do, and
+    must leave the caller's environment exactly as it was.
+    """
+    names = (GATEWAY_MODEL_DISCOVERY_ENV_VAR, CLAUDE_GATEWAY_MODEL_DISCOVERY_ENV_VAR)
+    previous = {name: os.environ.get(name, _MISSING_ENV_VALUE) for name in names}
+    try:
+        if enabled is not None:
+            for name in names:
+                if enabled:
+                    os.environ[name] = "1"
+                else:
+                    os.environ.pop(name, None)
+        yield
+    finally:
+        for name, value in previous.items():
+            if value is _MISSING_ENV_VALUE:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = cast(str, value)
+
+
 # Retained only to identify and remove state written by the legacy persisted opt-in.
 SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY
 
@@ -108,6 +147,18 @@ def _minimum_version_requirement_message(version: str) -> str:
     )
 
 
+def _model_discovery_minimum_version_error() -> str | None:
+    """Return the discovery blocker after the effective launch source is known."""
+    version = agent_version(SPEC["binary"])
+    parsed_version = _parse_version(version)
+    if parsed_version is None or parsed_version >= MINIMUM_CLAUDE_VERSION:
+        return None
+    return (
+        f"Model discovery requires Claude Code {MINIMUM_CLAUDE_VERSION_TEXT} or newer. "
+        f"Your current version is Claude Code {version}. Upgrade Claude Code and try again."
+    )
+
+
 def minimum_version_error() -> str | None:
     if (
         os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1"
@@ -1424,7 +1475,14 @@ def _rewrite_relayed_port(state: dict, port: int) -> None:
     a different port than the cached one. Keeps ANTHROPIC_BASE_URL (which Claude
     Code reads) in sync with the live proxy so requests reach it."""
     state["relayed_proxy_port"] = port
-    save_state(state)
+    persisted_state = dict(state)
+    for key in (
+        CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY,
+        "_claude_launch_provider",
+        "_claude_launch_parent_schema",
+    ):
+        persisted_state.pop(key, None)
+    save_state(persisted_state)
     settings = read_json_safe(CLAUDE_SETTINGS_PATH)
     env = settings.get("env")
     if isinstance(env, dict):
@@ -1480,44 +1538,63 @@ def launch(
 ) -> None:
     binary = SPEC["binary"]
     workspace = state.get("workspace")
-    if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
-        # Discovery is launch-scoped. Pass it in the process environment rather
-        # than persisting it in Claude's private or OS-managed settings.
-        os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
-    if state.get("claude_relayed"):
-        _launch_relayed(state, binary, tool_args)
-        return
-    # Smart routing needs Unix PTY support, which Windows does not provide.
-    if options.launch_smart_routing and os.name == "nt":
+    scoped_model_source = bool(
+        state.get("_claude_launch_provider")
+        or get_provider_service(state, "claude")
+        or state.get("_claude_launch_parent_schema")
+    )
+    override = state.get(CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY)
+    scoped_discovery = (
+        override if isinstance(override, bool) else scoped_model_discovery_enabled()
+    )
+    if options.launch_smart_routing and scoped_model_source:
         raise RuntimeError(
-            "Smart routing in Claude Code is currently not supported on Windows. "
-            "Please use Codex or disable smart routing."
-        )
-    if options.launch_smart_routing:
-        smart_routing_v2.launch_claude(
-            state,
-            tool_args,
-            binary=binary,
-            user_settings_path=CLAUDE_USER_SETTINGS_PATH,
-            # With no user pin, let Claude resolve its starting model from its own settings.
-            launch_model=options.user_pinned_model,
-            compose_settings=_compose_v2_settings,
-            launch_model_args=_launch_model_args,
-            model_name=_maybe_add_1m_suffix,
+            "Claude Code smart routing cannot be used with a Model Provider Service or model "
+            "location. Disable smart routing or remove the scoped model source and try again."
         )
-        return
-    if workspace:
-        os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
-    settings_override = None
-    launch_args = list(tool_args)
-    if options.user_pinned_model:
-        os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model
-        settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}}
-        launch_args = [
-            *_launch_model_args(tool_args, options.user_pinned_model),
-            *tool_args,
-        ]
-    exec_or_spawn(_build_claude_argv(binary, launch_args, settings_override=settings_override))
+    requested_discovery = os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1"
+    launch_discovery = scoped_discovery if scoped_model_source else requested_discovery
+    manage_discovery_env = launch_discovery if scoped_model_source or requested_discovery else None
+
+    with launch_discovery_environment(manage_discovery_env):
+        if workspace and launch_discovery:
+            version_error = _model_discovery_minimum_version_error()
+            if version_error:
+                raise RuntimeError(version_error)
+        if state.get("claude_relayed"):
+            _launch_relayed(state, binary, tool_args)
+            return
+        # Smart routing needs Unix PTY support, which Windows does not provide.
+        if options.launch_smart_routing and os.name == "nt":
+            raise RuntimeError(
+                "Smart routing in Claude Code is currently not supported on Windows. "
+                "Please use Codex or disable smart routing."
+            )
+        if options.launch_smart_routing:
+            smart_routing_v2.launch_claude(
+                state,
+                tool_args,
+                binary=binary,
+                user_settings_path=CLAUDE_USER_SETTINGS_PATH,
+                # With no user pin, let Claude resolve its starting model from its own settings.
+                launch_model=options.user_pinned_model,
+                compose_settings=_compose_v2_settings,
+                launch_model_args=_launch_model_args,
+                model_name=_maybe_add_1m_suffix,
+            )
+            return
+        if workspace:
+            os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+        settings_override = None
+        launch_args = list(tool_args)
+        if options.user_pinned_model:
+            os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model
+            settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}}
+            launch_args = [
+                *_launch_model_args(tool_args, options.user_pinned_model),
+                *tool_args,
+            ]
+        exec_or_spawn(_build_claude_argv(binary, launch_args, settings_override=settings_override))
 
 
 def validate_cmd(binary: str) -> list[str]:

src/ucode/cli.py

@@ -43,6 +43,10 @@
 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 (
+    CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY,
+    scoped_model_discovery_enabled,
+)
 from ucode.databricks import (
     apply_pat_environment,
     build_shared_base_urls,
@@ -2221,6 +2225,7 @@ def _launch_tool(
                 )
         elif not coding_agent_config_feature_disabled:
             print_note("No managed coding agent config found; using your own settings")
+        managed_provider = None
         if managed is not None:
             managed_provider = managed_provider_service(managed, tool)
             if explicit_provider and managed_provider and managed_provider != explicit_provider:
@@ -2236,8 +2241,18 @@ def _launch_tool(
                 provider = managed_provider
         if provider and parent_schema is not None:
             raise RuntimeError("--provider and --model-location cannot be used together.")
+        claude_scoped_model_source = tool == "claude" and bool(provider or parent_schema)
+        claude_scoped_model_discovery = _scoped_model_discovery_enabled(
+            managed_config_exists=managed is not None
+        )
         # 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 smart_routing_enabled and claude_scoped_model_source:
+            raise RuntimeError(
+                f"{TOOL_SPECS[tool]['display']} smart routing cannot be used with a Model "
+                "Provider Service or model location. Disable smart routing or remove the scoped "
+                "model source and try again."
+            )
         if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider:
             raise RuntimeError(
                 f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with "
@@ -2378,6 +2393,10 @@ def _launch_tool(
         if tool == "claude":
             if provider:
                 state["_claude_launch_provider"] = provider
+            elif parent_schema:
+                state["_claude_launch_parent_schema"] = parent_schema
+            if claude_scoped_model_source:
+                state[CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY] = claude_scoped_model_discovery
         elif tool == "codex":
             if provider:
                 state["_codex_launch_provider"] = provider
@@ -2567,6 +2586,17 @@ def _print_no_managed_config_guidance() -> None:
     )
 
 
+def _scoped_model_discovery_enabled(*, managed_config_exists: bool = False) -> bool:
+    """Whether an agent should refresh a provider/location-scoped catalog.
+
+    ``UG_ENABLE_MODEL_DISCOVERY=0`` affects only an agent's scoped picker
+    catalog. It does not suppress the normal workspace discovery that populates
+    ``system.ai`` models. Managed config may force discovery because workspace
+    policy outranks a developer environment variable.
+    """
+    return scoped_model_discovery_enabled(managed_config_exists=managed_config_exists)
+
+
 @app.command(
     "codex",
     cls=_PromptAwareCommand,
@@ -2750,21 +2780,25 @@ def claude_cmd(
         claude_agent.disable_smart_routing(load_state())
         print_success("Claude Code smart routing disabled; ug routing hooks removed")
         return
-    if enable_model_discovery or (model_location is not None and provider is None):
-        os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1"
-    with _smart_routing_v2_flag(enable_smart_routing_flag):
-        with _disable_smart_routing_for_subcommand("claude", ctx):
-            _launch_tool(
-                "claude",
-                ctx,
-                provider=provider,
-                model=model,
-                refresh=refresh,
-                skip_preflight=skip_preflight,
-                workspace_url=workspace,
-                parent_schema=model_location,
-                custom_oauth=custom_oauth,
-            )
+    requested_discovery: bool | None = None
+    if enable_model_discovery:
+        requested_discovery = True
+    elif provider is not None or model_location is not None:
+        requested_discovery = _scoped_model_discovery_enabled()
+    with claude_agent.launch_discovery_environment(requested_discovery):
+        with _smart_routing_v2_flag(enable_smart_routing_flag):
+            with _disable_smart_routing_for_subcommand("claude", ctx):
+                _launch_tool(
+                    "claude",
+                    ctx,
+                    provider=provider,
+                    model=model,
+                    refresh=refresh,
+                    skip_preflight=skip_preflight,
+                    workspace_url=workspace,
+                    parent_schema=model_location,
+                    custom_oauth=custom_oauth,
+                )
 
 
 @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})

src/ucode/constants.py

@@ -11,6 +11,10 @@
 # Controls provider- and location-scoped agent catalogs.
 MODEL_DISCOVERY_ENV_VAR = "UG_ENABLE_MODEL_DISCOVERY"
 
+# Launch-only state handed from the CLI to Claude. A managed source sets this
+# true so workspace policy can override a developer's environment variable.
+CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY = "_claude_scoped_model_discovery"
+
 
 def scoped_model_discovery_enabled(
     *,

tests/conftest.py

@@ -60,6 +60,7 @@ def reject_privileged_write(path, _desired_text):
     monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write)
     monkeypatch.delenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", raising=False)
     monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False)
+    monkeypatch.delenv("UG_ENABLE_MODEL_DISCOVERY", raising=False)
     # A developer's ambient managed-config stub would otherwise short-circuit every fetch in the suite.
     monkeypatch.delenv("UCODE_MANAGED_CONFIG_STUB", raising=False)
     # The model-services listing is memoized for the life of the process, so without this a cached