From 865aba6547b1b4f3500407cd557b83b36c8bf02f Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:13:00 +0000 Subject: [PATCH 1/3] [AIGTWY-4565] Use MPS model catalog in Codex --- src/ucode/agents/codex.py | 19 +++++++++++++++++-- src/ucode/cli.py | 2 ++ src/ucode/databricks.py | 24 ++++++++++++++++++++---- tests/test_agent_codex.py | 20 ++++++++++++++++++++ tests/test_cli.py | 17 +++++++++++++++++ tests/test_databricks.py | 25 +++++++++++++++++++++++++ 6 files changed, 101 insertions(+), 6 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index f0f10429..a92c9c4e 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -18,12 +18,14 @@ backup_existing_file, deep_merge_dict, read_toml_safe, + write_json_file, write_toml_file, ) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_token_argv from ucode.databricks import ( build_auth_token_argv, build_tool_base_url, + fetch_codex_mps_model_catalog, get_databricks_token, ) from ucode.launcher import exec_or_spawn @@ -46,7 +48,7 @@ sync_smart_routing_hooks, ) from ucode.smart_routing.codex_routing import codex_model_id -from ucode.state import mark_tool_managed, save_state +from ucode.state import get_provider_service, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version from ucode.ui import print_warning_err @@ -56,6 +58,7 @@ CODEX_PROFILE_NAME = "ucode" CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / f"{CODEX_PROFILE_NAME}.config.toml" CODEX_BACKUP_PATH = APP_DIR / "codex-ucode-config.backup.toml" +CODEX_MPS_MODEL_CATALOG_PATH = APP_DIR / "codex-mps-model-catalog.json" LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml" LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" @@ -509,8 +512,16 @@ def launch( clear_model_preferences(state) binary = SPEC["binary"] workspace = state.get("workspace") + launch_provider = state.get("_codex_launch_provider") + provider = ( + launch_provider.strip() + if isinstance(launch_provider, str) and launch_provider.strip() + else get_provider_service(state, "codex") + ) + token = None if workspace: - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + token = get_databricks_token(workspace, state.get("profile")) + os.environ["OAUTH_TOKEN"] = token # Layer ucode's named profile as ordinary config overrides. Unlike # `--profile`, `--config` is accepted by runtime, utility, and server # commands, so every invocation keeps the same Databricks settings without @@ -521,6 +532,10 @@ def launch( f"Cannot launch Codex with the ucode profile because {CODEX_CONFIG_PATH} " "is missing or empty. Run `ucode configure --agents codex` first." ) + if workspace and token and provider: + catalog = fetch_codex_mps_model_catalog(workspace, token, provider) + write_json_file(CODEX_MPS_MODEL_CATALOG_PATH, catalog) + profile_doc["model_catalog_json"] = str(CODEX_MPS_MODEL_CATALOG_PATH) exec_or_spawn([binary, *codex_config_args(profile_doc), *tool_args]) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d5739e2e..f48d7256 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2213,6 +2213,8 @@ def _launch_tool( state["_claude_launch_model"] = launch_model if provider: state["_claude_launch_provider"] = provider + elif tool == "codex" and provider: + state["_codex_launch_provider"] = provider launch_options = _launch_options( tool, ctx.args, diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b4caae66..1d0c3c72 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -252,6 +252,7 @@ def _http_get_json( *, timeout: int = 10, max_retries: int = 0, + headers: dict[str, str] | None = None, ) -> tuple[dict | list | None, str | None]: """GET a JSON endpoint. Returns (payload, None) on success, (None, reason) on failure. @@ -264,10 +265,9 @@ def _http_get_json( if max_retries < 0: raise ValueError("max_retries must be non-negative") - request = urllib_request.Request( - url, - headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}, - ) + request_headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"} + request_headers.update(headers or {}) + request = urllib_request.Request(url, headers=request_headers) for attempt in range(max_retries + 1): try: with urllib_request.urlopen(request, timeout=timeout) as response: @@ -3258,6 +3258,22 @@ def build_tool_base_url(tool: str, workspace: str) -> str: raise RuntimeError(f"Unsupported tool '{tool}'.") +def fetch_codex_mps_model_catalog(workspace: str, token: str, provider: str) -> dict: + payload, reason = _http_get_json( + f"{build_tool_base_url('codex', workspace)}/models", + token, + max_retries=2, + headers={"Databricks-Model-Provider-Service": provider}, + ) + if reason: + raise RuntimeError(f"Could not discover Codex models for {provider}: {reason}") + if not isinstance(payload, dict) or not isinstance(payload.get("models"), list): + raise RuntimeError(f"Provider {provider} returned an invalid Codex model catalog.") + if not payload["models"]: + raise RuntimeError(f"Provider {provider} returned no Codex models.") + return payload + + def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 4f4af970..163bea3e 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -611,6 +611,26 @@ def test_sets_oauth_token(self, tmp_path, monkeypatch): assert os.environ["OAUTH_TOKEN"] == "fresh-token" assert launches[0][-1] == "--search" + def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypatch): + launches = self._patch(tmp_path, monkeypatch) + catalog_path = tmp_path / "models.json" + catalog = {"models": [{"slug": "gpt-mps"}]} + monkeypatch.setattr(codex, "CODEX_MPS_MODEL_CATALOG_PATH", catalog_path) + monkeypatch.setattr( + codex, + "fetch_codex_mps_model_catalog", + lambda workspace, token, provider: catalog, + ) + + codex.launch( + {"workspace": WS, "_codex_launch_provider": "main.default.openai"}, + [], + options=LaunchOptions(), + ) + + assert catalog_path.exists() + assert f'model_catalog_json="{catalog_path}"' in launches[0] + @pytest.mark.parametrize( "tool_args", [ diff --git a/tests/test_cli.py b/tests/test_cli.py index 11e4aa09..df14f377 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -985,6 +985,23 @@ def test_provider_sets_transient_claude_launch_marker(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_claude_launch_provider"] == "main.default.anthropic" + def test_provider_sets_transient_codex_launch_marker(self): + 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), + patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)), + 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", "--provider", "main.default.openai"]) + + assert result.exit_code == 0, result.output + assert mock_launch.call_args.args[1]["_codex_launch_provider"] == "main.default.openai" + class TestGeminiProviderLaunch: @staticmethod diff --git a/tests/test_databricks.py b/tests/test_databricks.py index baccc7b6..e1bed471 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -62,6 +62,31 @@ def read(self): return self._body +class TestFetchCodexMpsModelCatalog: + def test_sends_provider_header(self, monkeypatch): + seen = {} + + def fake_get(url, token, **kwargs): + seen.update(url=url, token=token, **kwargs) + return {"models": [{"slug": "gpt-mps"}]}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + result = db_mod.fetch_codex_mps_model_catalog(WS, "tok", "main.default.openai") + + assert result["models"][0]["slug"] == "gpt-mps" + assert seen["url"] == f"{WS}/ai-gateway/codex/v1/models" + assert seen["headers"] == {"Databricks-Model-Provider-Service": "main.default.openai"} + + def test_rejects_empty_catalog(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda *args, **kwargs: ({"models": []}, None) + ) + + with pytest.raises(RuntimeError, match="returned no Codex models"): + db_mod.fetch_codex_mps_model_catalog(WS, "tok", "main.default.openai") + + class TestWorkspaceHostname: def test_extracts_hostname(self): assert workspace_hostname(WS) == "example.databricks.com" From 410489fe2bc423ff93477d97193018a47fa43222 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:04:41 +0000 Subject: [PATCH 2/3] [AIGTWY-4565] Scope Codex MPS discovery to launch --- src/ucode/agents/codex.py | 78 ++++++++++++++++++-- tests/test_agent_codex.py | 148 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 217 insertions(+), 9 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index a92c9c4e..33b55d35 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -3,8 +3,10 @@ from __future__ import annotations import copy +import hashlib import os import re +import tempfile from collections.abc import Callable from pathlib import Path @@ -21,7 +23,11 @@ write_json_file, write_toml_file, ) -from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_token_argv +from ucode.custom_oauth import ( + CustomOAuthConfig, + build_custom_auth_token_argv, + get_custom_client_token, +) from ucode.databricks import ( build_auth_token_argv, build_tool_base_url, @@ -62,6 +68,7 @@ LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml" LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" +MPS_HEADER = "Databricks-Model-Provider-Service" MINIMUM_CODEX_VERSION = (0, 134, 0) MINIMUM_CODEX_VERSION_TEXT = "0.134.0" MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0) @@ -163,7 +170,7 @@ def _provider_block( # Route to an external Model Provider Service; the gateway selects the # provider from this header on every request. if provider: - http_headers["Databricks-Model-Provider-Service"] = provider + http_headers[MPS_HEADER] = provider return { "name": "Databricks AI Gateway", "base_url": base_url, @@ -348,6 +355,7 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non ): for key in ("model", "model_reasoning_effort"): profiles[CODEX_PROFILE_NAME].pop(key, None) + _set_provider_header(doc, None) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS) save_state(state) @@ -370,6 +378,7 @@ def compose(base: dict) -> dict: if chosen_model is None: for key in ("model", "model_reasoning_effort"): base.pop(key, None) + _set_provider_header(base, None) return base doc = read_toml_safe(CODEX_CONFIG_PATH) @@ -500,6 +509,63 @@ def clear_model_preferences(state: dict) -> bool: return changed +def _set_provider_header(config: dict, provider: str | None) -> None: + model_providers = config.get("model_providers") + if not isinstance(model_providers, dict): + return + provider_block = model_providers.get(CODEX_MODEL_PROVIDER_NAME) + if not isinstance(provider_block, dict): + return + headers = provider_block.get("http_headers") + if not isinstance(headers, dict): + provider_block["http_headers"] = {} + headers = provider_block["http_headers"] + if provider: + headers[MPS_HEADER] = provider + else: + headers.pop(MPS_HEADER, None) + + +def _model_catalog_path(workspace: str, provider: str) -> Path: + key = f"{workspace.rstrip('/')}\0{provider}".encode() + digest = hashlib.sha256(key).hexdigest()[:16] + base = CODEX_MPS_MODEL_CATALOG_PATH + return base.with_name(f"{base.stem}-{digest}{base.suffix}") + + +def _write_model_catalog(path: Path, catalog: dict) -> None: + temp_path = None + try: + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_temp_path = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + os.close(fd) + temp_path = Path(raw_temp_path) + write_json_file(temp_path, catalog) + os.replace(temp_path, path) + except OSError as exc: + raise RuntimeError(f"Could not write Codex model catalog at {path}.") from exc + finally: + if temp_path is not None: + try: + temp_path.unlink(missing_ok=True) + except OSError: + pass + + +def _launch_token(state: dict, workspace: str) -> str: + custom_oauth = state.get("custom_oauth") + if isinstance(custom_oauth, dict): + return get_custom_client_token( + workspace, + custom_oauth["client_id"], + custom_oauth["redirect_url"], + scopes=custom_oauth["scopes"], + ) + return get_databricks_token(workspace, state.get("profile")) + + def launch( state: dict, tool_args: list[str], @@ -520,7 +586,7 @@ def launch( ) token = None if workspace: - token = get_databricks_token(workspace, state.get("profile")) + token = _launch_token(state, workspace) os.environ["OAUTH_TOKEN"] = token # Layer ucode's named profile as ordinary config overrides. Unlike # `--profile`, `--config` is accepted by runtime, utility, and server @@ -532,10 +598,12 @@ def launch( f"Cannot launch Codex with the ucode profile because {CODEX_CONFIG_PATH} " "is missing or empty. Run `ucode configure --agents codex` first." ) + _set_provider_header(profile_doc, provider) if workspace and token and provider: catalog = fetch_codex_mps_model_catalog(workspace, token, provider) - write_json_file(CODEX_MPS_MODEL_CATALOG_PATH, catalog) - profile_doc["model_catalog_json"] = str(CODEX_MPS_MODEL_CATALOG_PATH) + catalog_path = _model_catalog_path(workspace, provider) + _write_model_catalog(catalog_path, catalog) + profile_doc["model_catalog_json"] = str(catalog_path) exec_or_spawn([binary, *codex_config_args(profile_doc), *tool_args]) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 163bea3e..096669fd 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os from pathlib import Path @@ -154,6 +155,7 @@ def test_user_agent_set_on_provider(self, monkeypatch): def test_managed_keys_include_http_headers(self): # Revert must clean up the new key. assert ["model_providers", "ucode-databricks", "http_headers"] in codex.MANAGED_KEYS + assert ["model_catalog_json"] not in codex.MANAGED_KEYS class TestCodexWriteConfig: @@ -203,7 +205,7 @@ def test_removes_uc_model_services_id(self, tmp_path, monkeypatch): doc = read_toml_safe(config_path) assert "model" not in doc - def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatch): + def test_provider_drops_stale_model_without_persisting_header(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" backup_path = tmp_path / "codex-ucode-config.backup.toml" monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) @@ -215,7 +217,7 @@ def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatc codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) assert "model" not in read_toml_safe(config_path) - # A provider run must clear it and add the routing header. + # The routing header is added only to the launch config. codex.write_tool_config( {"workspace": WS, "codex_models": ["gpt-5"]}, provider="main.aarushi.aarushi-openai", @@ -224,7 +226,21 @@ def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatc doc = read_toml_safe(config_path) assert "model" not in doc headers = doc["model_providers"]["ucode-databricks"]["http_headers"] - assert headers["Databricks-Model-Provider-Service"] == "main.aarushi.aarushi-openai" + assert "Databricks-Model-Provider-Service" not in headers + + def test_non_provider_write_removes_stale_provider_header(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + state = {"workspace": WS, "codex_models": ["gpt-5"]} + codex.write_tool_config(state, provider="main.default.openai") + codex.write_tool_config(state) + + headers = read_toml_safe(config_path)["model_providers"]["ucode-databricks"]["http_headers"] + assert "Databricks-Model-Provider-Service" not in headers def test_clears_profile_model_preferences_before_launch(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" @@ -615,7 +631,7 @@ def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypat launches = self._patch(tmp_path, monkeypatch) catalog_path = tmp_path / "models.json" catalog = {"models": [{"slug": "gpt-mps"}]} - monkeypatch.setattr(codex, "CODEX_MPS_MODEL_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "_model_catalog_path", lambda workspace, provider: catalog_path) monkeypatch.setattr( codex, "fetch_codex_mps_model_catalog", @@ -630,6 +646,106 @@ def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypat assert catalog_path.exists() assert f'model_catalog_json="{catalog_path}"' in launches[0] + provider_arg = next( + arg for arg in launches[0] if arg.startswith("model_providers.ucode-databricks=") + ) + assert 'Databricks-Model-Provider-Service = "main.default.openai"' in provider_arg + + def test_non_provider_launch_removes_stale_provider_header(self, tmp_path, monkeypatch): + launches = self._patch(tmp_path, monkeypatch) + profile_path = tmp_path / "ucode.config.toml" + profile_path.write_text( + profile_path.read_text(encoding="utf-8") + + "\n[model_providers.ucode-databricks.http_headers]\n" + + 'Databricks-Model-Provider-Service = "main.default.old"\n', + encoding="utf-8", + ) + + codex.launch({"workspace": WS}, [], options=LaunchOptions()) + + provider_arg = next( + arg for arg in launches[0] if arg.startswith("model_providers.ucode-databricks=") + ) + assert "Databricks-Model-Provider-Service" not in provider_arg + + def test_provider_discovery_uses_custom_oauth_token(self, tmp_path, monkeypatch): + self._patch(tmp_path, monkeypatch) + catalog_path = tmp_path / "models.json" + seen = {} + monkeypatch.setattr(codex, "_model_catalog_path", lambda workspace, provider: catalog_path) + monkeypatch.setattr( + codex, + "get_databricks_token", + lambda *args, **kwargs: pytest.fail("standard token used"), + ) + monkeypatch.setattr( + codex, + "get_custom_client_token", + lambda workspace, client_id, redirect_url, *, scopes: "custom-token", + ) + + def fetch(workspace, token, provider): + seen.update(workspace=workspace, token=token, provider=provider) + return {"models": [{"slug": "gpt-mps"}]} + + monkeypatch.setattr(codex, "fetch_codex_mps_model_catalog", fetch) + state = { + "workspace": WS, + "_codex_launch_provider": "main.default.openai", + "custom_oauth": { + "client_id": "client", + "redirect_url": "http://localhost:8020", + "scopes": ["all-apis", "offline_access"], + }, + } + + codex.launch(state, [], options=LaunchOptions()) + + assert seen == { + "workspace": WS, + "token": "custom-token", + "provider": "main.default.openai", + } + assert os.environ["OAUTH_TOKEN"] == "custom-token" + + def test_catalog_paths_are_provider_scoped(self, tmp_path, monkeypatch): + monkeypatch.setattr(codex, "CODEX_MPS_MODEL_CATALOG_PATH", tmp_path / "models.json") + + first = codex._model_catalog_path(WS, "main.default.first") + second = codex._model_catalog_path(WS, "main.default.second") + + assert first != second + assert first == codex._model_catalog_path(WS, "main.default.first") + + def test_catalog_write_is_complete_and_atomic(self, tmp_path): + path = tmp_path / "models.json" + catalog = {"models": [{"slug": "gpt-mps"}]} + + codex._write_model_catalog(path, catalog) + + assert json.loads(path.read_text(encoding="utf-8")) == catalog + assert list(tmp_path.glob(".models.json.*.tmp")) == [] + + def test_catalog_write_reports_path_on_failure(self, tmp_path, monkeypatch): + path = tmp_path / "models.json" + monkeypatch.setattr(codex.os, "replace", lambda *args: (_ for _ in ()).throw(OSError())) + + with pytest.raises(RuntimeError, match=str(path)): + codex._write_model_catalog(path, {"models": [{"slug": "gpt-mps"}]}) + + assert list(tmp_path.glob(".models.json.*.tmp")) == [] + + def test_catalog_cleanup_does_not_mask_write_failure(self, tmp_path, monkeypatch): + path = tmp_path / "models.json" + monkeypatch.setattr(codex.os, "replace", lambda *args: (_ for _ in ()).throw(OSError())) + monkeypatch.setattr( + codex.Path, + "unlink", + lambda *args, **kwargs: (_ for _ in ()).throw(OSError()), + ) + + with pytest.raises(RuntimeError, match=str(path)): + codex._write_model_catalog(path, {"models": [{"slug": "gpt-mps"}]}) @pytest.mark.parametrize( "tool_args", @@ -714,6 +830,30 @@ def test_managed_config_preserves_other_keys(self, tmp_path, monkeypatch): assert doc["approval_policy"] == "on-request" assert "model" not in doc + def test_provider_settings_stay_launch_scoped(self, tmp_path, monkeypatch): + config_path, managed_path = self._patch(tmp_path, monkeypatch) + managed_path.parent.mkdir(parents=True, exist_ok=True) + managed_path.write_text( + 'model_catalog_json = "/tmp/stale.json"\n\n' + "[model_providers.ucode-databricks.http_headers]\n" + 'Databricks-Model-Provider-Service = "main.default.stale"\n', + encoding="utf-8", + ) + + codex.write_tool_config( + {"workspace": WS, "codex_models": ["gpt-5"]}, + provider="main.default.openai", + ) + + local_headers = read_toml_safe(config_path)["model_providers"]["ucode-databricks"][ + "http_headers" + ] + managed = read_toml_safe(managed_path) + managed_headers = managed["model_providers"]["ucode-databricks"]["http_headers"] + assert codex.MPS_HEADER not in local_headers + assert codex.MPS_HEADER not in managed_headers + assert managed["model_catalog_json"] == "/tmp/stale.json" + def test_noninteractive_uses_local_config_when_managed_config_is_compatible( self, tmp_path, monkeypatch ): From 96ede79acf87d0488471b829b77ae13df7811679 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:12:00 +0000 Subject: [PATCH 3/3] [AIGTWY-4565] Reject managed Codex catalog conflicts --- src/ucode/agents/codex.py | 20 ++++++++++++++++++++ tests/test_agent_codex.py | 15 +++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 33b55d35..76192ce8 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -566,6 +566,24 @@ def _launch_token(state: dict, workspace: str) -> str: return get_databricks_token(workspace, state.get("profile")) +def _reject_managed_mps_catalog() -> None: + path = _managed_config_path() + if path is None: + return + text = read_managed_file(path) + if text is None: + return + try: + managed = _parse_managed_config(text) + except RuntimeError as exc: + raise RuntimeError(f"Cannot read Codex managed settings at {path}: {exc}") from exc + if "model_catalog_json" in managed: + raise RuntimeError( + f"Codex managed settings at {path} define model_catalog_json, which overrides MPS " + "discovery. Remove it or contact your administrator." + ) + + def launch( state: dict, tool_args: list[str], @@ -584,6 +602,8 @@ def launch( if isinstance(launch_provider, str) and launch_provider.strip() else get_provider_service(state, "codex") ) + if workspace and provider: + _reject_managed_mps_catalog() token = None if workspace: token = _launch_token(state, workspace) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 096669fd..37e8e8d6 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -651,6 +651,21 @@ def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypat ) assert 'Databricks-Model-Provider-Service = "main.default.openai"' in provider_arg + def test_provider_rejects_managed_model_catalog(self, tmp_path, monkeypatch): + launches = self._patch(tmp_path, monkeypatch) + managed_path = tmp_path / "managed_config.toml" + managed_path.write_text('model_catalog_json = "/admin/models.json"\n', encoding="utf-8") + monkeypatch.setattr(codex, "_managed_config_path", lambda: managed_path) + + with pytest.raises(RuntimeError, match="overrides MPS discovery"): + codex.launch( + {"workspace": WS, "_codex_launch_provider": "main.default.openai"}, + [], + options=LaunchOptions(), + ) + + assert launches == [] + def test_non_provider_launch_removes_stale_provider_header(self, tmp_path, monkeypatch): launches = self._patch(tmp_path, monkeypatch) profile_path = tmp_path / "ucode.config.toml"