Control Codex scoped model discovery with UG_ENABLE_MODEL_DISCOVERY

README.md

@@ -83,10 +83,11 @@ All agents route through Databricks AI Gateway using your workspace credentials
 
 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.
+before launch. Codex continues to receive a launch-scoped model catalog from ug. Set
+`UG_ENABLE_MODEL_DISCOVERY=0` to keep the routing header while using the agent'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.

src/ucode/agents/codex.py

@@ -32,8 +32,10 @@
     write_toml_file,
 )
 from ucode.constants import (
+    CODEX_SCOPED_MODEL_DISCOVERY_STATE_KEY,
     MODEL_PROVIDER_SERVICE_HEADER,
     MODEL_SERVICE_PARENT_SCHEMA_HEADER,
+    scoped_model_discovery_enabled,
 )
 from ucode.custom_oauth import (
     CUSTOM_OAUTH_TIMEOUT_MS,
@@ -747,12 +749,6 @@ def launch(
     *,
     options: LaunchOptions,
 ) -> None:
-    if options.launch_smart_routing:
-        _launch_smart_routing(state, tool_args)
-        return
-    clear_model_preferences(state)
-    binary = SPEC["binary"]
-    workspace = state.get("workspace")
     launch_provider = state.get("_codex_launch_provider")
     provider = (
         launch_provider.strip()
@@ -765,7 +761,23 @@ def launch(
         if isinstance(launch_parent_schema, str) and launch_parent_schema.strip()
         else None
     )
-    if workspace and (provider or parent_schema):
+    scoped_model_source = bool(provider or parent_schema)
+    if options.launch_smart_routing and scoped_model_source:
+        raise RuntimeError(
+            "Codex 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 options.launch_smart_routing:
+        _launch_smart_routing(state, tool_args)
+        return
+    clear_model_preferences(state)
+    binary = SPEC["binary"]
+    workspace = state.get("workspace")
+    override = state.get(CODEX_SCOPED_MODEL_DISCOVERY_STATE_KEY)
+    scoped_model_discovery = (
+        override if isinstance(override, bool) else scoped_model_discovery_enabled()
+    )
+    if workspace and scoped_model_source and scoped_model_discovery:
         _reject_managed_model_catalog()
     otel_args: list[str] = []
     token = None
@@ -794,7 +806,11 @@ def launch(
         )
     _set_provider_header(profile_doc, provider)
     _set_parent_schema_header(profile_doc, parent_schema if not provider else None)
-    if workspace and token and (provider or parent_schema):
+    if scoped_model_source and not scoped_model_discovery:
+        # A persisted/static catalog must not defeat the native-catalog policy
+        # for this invocation. This changes only the launch overlay, not disk.
+        profile_doc.pop("model_catalog_json", None)
+    if workspace and token and scoped_model_source and scoped_model_discovery:
         try:
             if provider is not None:
                 catalog_source = CodexCatalogSource.PROVIDER

src/ucode/cli.py

@@ -45,6 +45,7 @@
 from ucode.config_io import is_dry_run, restore_file, set_dry_run
 from ucode.constants import (
     CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY,
+    CODEX_SCOPED_MODEL_DISCOVERY_STATE_KEY,
     scoped_model_discovery_enabled,
 )
 from ucode.databricks import (
@@ -2241,23 +2242,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(
+        scoped_model_source = bool(provider or parent_schema)
+        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:
+        if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and 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 "
-                "--provider. Launch without a Model Provider Service and try again."
-            )
         # Validate the provider service before launching — it must exist, be a
         # provider type this tool can route to (e.g. claude can't use an OpenAI
         # or Foundry service), and, for Bedrock, expose Claude models to pin.
@@ -2395,9 +2391,11 @@ def _launch_tool(
                 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
+            if scoped_model_source:
+                state[CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY] = scoped_model_discovery
         elif tool == "codex":
+            if scoped_model_source:
+                state[CODEX_SCOPED_MODEL_DISCOVERY_STATE_KEY] = scoped_model_discovery
             if provider:
                 state["_codex_launch_provider"] = provider
             elif parent_schema:

src/ucode/constants.py

@@ -15,6 +15,10 @@
 # true so workspace policy can override a developer's environment variable.
 CLAUDE_SCOPED_MODEL_DISCOVERY_STATE_KEY = "_claude_scoped_model_discovery"
 
+# Launch-only state handed from the CLI to Codex. Routing headers remain active
+# when this is false; only the scoped ``model_catalog_json`` refresh is skipped.
+CODEX_SCOPED_MODEL_DISCOVERY_STATE_KEY = "_codex_scoped_model_discovery"
+
 
 def scoped_model_discovery_enabled(
     *,

tests/test_agent_codex.py

@@ -777,6 +777,78 @@ def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypat
         )
         assert 'Databricks-Model-Provider-Service = "main.default.openai"' in provider_arg
 
+    @pytest.mark.parametrize(
+        ("scope_state", "expected_header"),
+        [
+            (
+                {"_codex_launch_provider": "main.default.openai"},
+                'Databricks-Model-Provider-Service = "main.default.openai"',
+            ),
+            (
+                {"_codex_launch_parent_schema": "main.default"},
+                'Databricks-Model-Service-Parent-Schema = "main.default"',
+            ),
+        ],
+    )
+    def test_disabled_scoped_discovery_keeps_routing_without_catalog(
+        self, tmp_path, monkeypatch, scope_state, expected_header
+    ):
+        launches = self._patch(tmp_path, monkeypatch)
+        monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0")
+        with codex.CODEX_CONFIG_PATH.open("a", encoding="utf-8") as config:
+            config.write('model_catalog_json = "/persisted/models.json"\n')
+
+        def unexpected(*_args, **_kwargs):
+            raise AssertionError("scoped catalog fetch must be disabled")
+
+        monkeypatch.setattr(codex, "_fetch_codex_model_catalog", unexpected)
+        monkeypatch.setattr(codex, "_reject_managed_model_catalog", unexpected)
+
+        codex.launch(
+            {
+                "workspace": WS,
+                **scope_state,
+            },
+            [],
+            options=LaunchOptions(),
+        )
+
+        assert launches
+        assert not any(arg.startswith("model_catalog_json=") for arg in launches[0])
+        provider_arg = next(
+            arg for arg in launches[0] if arg.startswith("model_providers.Databricks=")
+        )
+        assert expected_header in provider_arg
+
+    def test_managed_scoped_override_forces_discovery_when_public_policy_disables_it(
+        self, tmp_path, monkeypatch
+    ):
+        launches = self._patch(tmp_path, monkeypatch)
+        monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0")
+        catalog_path = tmp_path / "managed-models.json"
+        fetch_calls: list[tuple] = []
+
+        def fetch(*args, **kwargs):
+            fetch_calls.append((args, kwargs))
+            return {"models": [{"slug": "managed-gpt"}]}
+
+        monkeypatch.setattr(codex, "_fetch_codex_model_catalog", fetch)
+        monkeypatch.setattr(codex, "_model_catalog_path", lambda *_args: catalog_path)
+
+        codex.launch(
+            {
+                "workspace": WS,
+                "_codex_launch_parent_schema": "managed.models",
+                "_codex_scoped_model_discovery": True,
+            },
+            [],
+            options=LaunchOptions(),
+        )
+
+        assert len(fetch_calls) == 1
+        assert catalog_path.exists()
+        assert f'model_catalog_json="{catalog_path}"' in launches[0]
+
     def test_provider_pins_first_catalog_model(self, tmp_path, monkeypatch):
         launches = self._patch(tmp_path, monkeypatch)
         catalog = {"models": [{"slug": "gpt-primary"}, {"slug": "gpt-secondary"}]}

tests/test_cli.py

@@ -537,6 +537,24 @@ def test_claude_scoped_source_rejects_smart_routing(self, source_args):
         assert "smart routing cannot be used" in result.output
         assert "Disable smart routing or remove the scoped model source" in result.output
 
+    @pytest.mark.parametrize(
+        "source_args",
+        [
+            ["--provider", "main.default.provider"],
+            ["--model-location", "main.default"],
+        ],
+    )
+    def test_codex_scoped_source_rejects_smart_routing(self, source_args):
+        patches = _patch_launch("codex")
+        with contextlib.ExitStack() as stack:
+            for item in patches:
+                stack.enter_context(item)
+            result = runner.invoke(app, ["codex", "--enable-smart-routing", *source_args])
+
+        assert result.exit_code == 1
+        assert "smart routing cannot be used" in result.output
+        assert "Disable smart routing or remove the scoped model source" in result.output
+
     @pytest.mark.parametrize(
         ("args", "forwarded", "has_separator"),
         [
@@ -1206,7 +1224,9 @@ def test_provider_sets_transient_codex_launch_marker(self):
             result = runner.invoke(app, ["codex", "--provider", "main.default.openai"])
 
         assert result.exit_code == 0, result.output
-        assert mock_launch.call_args.args[1]["_codex_launch_provider"] == "main.default.openai"
+        launch_state = mock_launch.call_args.args[1]
+        assert launch_state["_codex_launch_provider"] == "main.default.openai"
+        assert launch_state["_codex_scoped_model_discovery"] is True
 
     def test_model_location_sets_transient_codex_launch_marker(self):
         state = dict(MINIMAL_STATE)
@@ -1223,7 +1243,9 @@ def test_model_location_sets_transient_codex_launch_marker(self):
             result = runner.invoke(app, ["codex", "--model-location", "main.default"])
 
         assert result.exit_code == 0, result.output
-        assert mock_launch.call_args.args[1]["_codex_launch_parent_schema"] == "main.default"
+        launch_state = mock_launch.call_args.args[1]
+        assert launch_state["_codex_launch_parent_schema"] == "main.default"
+        assert launch_state["_codex_scoped_model_discovery"] is True
 
     @pytest.mark.parametrize(
         ("tool", "source_args", "expected_marker"),
@@ -1268,6 +1290,46 @@ def test_claude_discovery_disable_keeps_scope_but_suppresses_native_discovery(
             # ordinary system.ai discovery still refreshes for this launch.
             assert mock_shared.call_args.kwargs["skip_model_discovery"] is False
 
+    @pytest.mark.parametrize(
+        ("source_args", "expected_marker"),
+        [
+            (
+                ["--provider", "main.default.openai"],
+                ("_codex_launch_provider", "main.default.openai"),
+            ),
+            (
+                ["--model-location", "main.default"],
+                ("_codex_launch_parent_schema", "main.default"),
+            ),
+        ],
+    )
+    def test_codex_discovery_disable_keeps_scope_but_suppresses_catalog(
+        self, monkeypatch, source_args, expected_marker
+    ):
+        monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0")
+        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.configure_shared_state", return_value=state) as mock_shared,
+            patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)),
+            patch("ucode.cli.resolve_launch_model", return_value=(state, "system.ai.model")),
+            patch("ucode.cli.configure_tool", return_value=state),
+            patch("ucode.cli._fetch_managed_config", return_value=(None, False)),
+            patch("ucode.cli.launch_agent") as mock_launch,
+        ):
+            result = runner.invoke(app, ["codex", *source_args])
+
+        assert result.exit_code == 0, result.output
+        launch_state = mock_launch.call_args.args[1]
+        assert launch_state[expected_marker[0]] == expected_marker[1]
+        assert launch_state["_codex_scoped_model_discovery"] is False
+        if source_args[0] == "--model-location":
+            # The policy flag suppresses only the agent-native scoped catalog;
+            # ordinary system.ai discovery still refreshes for this launch.
+            assert mock_shared.call_args.kwargs["skip_model_discovery"] is False
+
     def test_managed_config_overrides_developer_discovery_disable(self, monkeypatch):
         monkeypatch.setenv("UG_ENABLE_MODEL_DISCOVERY", "0")
 

tests/test_codex_smart_routing_v2.py

@@ -84,6 +84,26 @@ def launch_v2(state, tool_args, **kwargs):
             )
         ]
 
+    @pytest.mark.parametrize(
+        "scope_state",
+        [
+            {"_codex_launch_provider": "main.default.openai"},
+            {"_codex_launch_parent_schema": "main.default"},
+        ],
+    )
+    def test_direct_scoped_launch_rejects_smart_routing_before_v2(self, monkeypatch, scope_state):
+        launch_v2 = []
+        monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: launch_v2.append(kwargs))
+
+        with pytest.raises(RuntimeError, match="cannot be used.*model location"):
+            codex.launch(
+                {"workspace": WS, **scope_state},
+                [],
+                options=LaunchOptions(launch_smart_routing=True),
+            )
+
+        assert launch_v2 == []
+
     @pytest.mark.parametrize(
         "tool_args",
         [