From 1795ccb6e346c4f17cabd38269407dabd18abd93 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 20:23:30 +0000 Subject: [PATCH] Control Codex scoped model discovery --- README.md | 9 ++-- src/ucode/agents/codex.py | 32 +++++++++---- src/ucode/cli.py | 18 ++++--- src/ucode/constants.py | 4 ++ tests/test_agent_codex.py | 72 ++++++++++++++++++++++++++++ tests/test_cli.py | 66 ++++++++++++++++++++++++- tests/test_codex_smart_routing_v2.py | 20 ++++++++ 7 files changed, 197 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 68eb2c42..a63286e7 100644 --- a/README.md +++ b/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. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 572d9898..c560a1de 100644 --- a/src/ucode/agents/codex.py +++ b/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 diff --git a/src/ucode/cli.py b/src/ucode/cli.py index f6bc5cb7..2e952245 100644 --- a/src/ucode/cli.py +++ b/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: diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 03465b45..2943b762 100644 --- a/src/ucode/constants.py +++ b/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( *, diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index a95a1fe6..4c2b0842 100644 --- a/tests/test_agent_codex.py +++ b/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"}]} diff --git a/tests/test_cli.py b/tests/test_cli.py index 34936f51..84d53a16 100644 --- a/tests/test_cli.py +++ b/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") diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index f5876013..beb3bde1 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/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", [