From 0511f86adc920540c38982c6d06a6c6526680809 Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Wed, 10 Jun 2026 00:17:51 +0000 Subject: [PATCH] Fix ucode revert for legacy Codex shared-config edits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode revert` only restored the per-profile ~/.codex/ucode.config.toml, so the legacy in-place edits to the user's shared ~/.codex/config.toml (profile = "ucode", [profiles.ucode], [model_providers.ucode-databricks]) were left behind — every bare `codex` invocation kept routing through the gateway. - revert now surgically strips the three ucode-namespaced keys from the shared config (safer than a whole-file backup restore, which would clobber later edits) - _remove_legacy_ucode_profile also drops the inert ucode-databricks provider block - default_model only considers GPT-parseable ids, so a non-GPT workspace (e.g. moonshotai/kimi-k2.5) no longer pins an unroutable model Co-authored-by: Isaac --- src/ucode/agents/codex.py | 94 +++++++++++++++++++++++++++++++-------- src/ucode/cli.py | 6 +++ tests/test_agent_codex.py | 89 ++++++++++++++++++++++++++++++++++++ tests/test_agents_init.py | 9 ++-- 4 files changed, 176 insertions(+), 22 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 3719eea..e0bb64b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -160,30 +160,84 @@ def _legacy_backup_path() -> Path: return CODEX_BACKUP_PATH.with_name("codex-legacy-config.backup.toml") -def _remove_legacy_ucode_profile() -> None: - """Remove ucode's old [profiles.ucode] entry from shared Codex config.""" - path = _legacy_config_path() - if path == CODEX_CONFIG_PATH or not path.exists(): - return +def _has_legacy_ucode_entries(doc: dict) -> bool: + profiles = doc.get("profiles") + providers = doc.get("model_providers") + return ( + doc.get("profile") == CODEX_PROFILE_NAME + or (isinstance(profiles, dict) and CODEX_PROFILE_NAME in profiles) + or (isinstance(providers, dict) and CODEX_MODEL_PROVIDER_NAME in providers) + ) + + +def _strip_legacy_ucode_entries(path: Path) -> bool: + """Surgically remove ucode's keys from a shared Codex config. + + Drops the top-level ``profile = "ucode"`` selector, ``[profiles.ucode]``, + and ``[model_providers.ucode-databricks]`` while leaving everything else the + user has in the file untouched. Returns True if anything was removed. + + Surgical removal beats restoring the backup: ``backup_existing_file`` only + keeps the first-ever snapshot, so a whole-file restore would clobber edits + made since ucode first ran. + """ + if not path.exists(): + return False doc = read_toml_safe(path) changed = False + if doc.get("profile") == CODEX_PROFILE_NAME: + doc.pop("profile", None) + changed = True + profiles = doc.get("profiles") if isinstance(profiles, dict) and CODEX_PROFILE_NAME in profiles: - backup_existing_file(path, _legacy_backup_path()) profiles.pop(CODEX_PROFILE_NAME, None) if not profiles: doc.pop("profiles", None) changed = True - if doc.get("profile") == CODEX_PROFILE_NAME: - backup_existing_file(path, _legacy_backup_path()) - doc.pop("profile", None) + providers = doc.get("model_providers") + if isinstance(providers, dict) and CODEX_MODEL_PROVIDER_NAME in providers: + providers.pop(CODEX_MODEL_PROVIDER_NAME, None) + if not providers: + doc.pop("model_providers", None) changed = True if changed: write_toml_file(path, doc) + return changed + + +def _remove_legacy_ucode_profile() -> None: + """Remove ucode's old shared-config entries when configuring modern Codex. + + Strips the legacy ``profile``/``[profiles.ucode]`` selector and the + ``[model_providers.ucode-databricks]`` provider block that older ucode + versions deep-merged into ``~/.codex/config.toml``. + """ + path = _legacy_config_path() + if path == CODEX_CONFIG_PATH or not path.exists(): + return + + if _has_legacy_ucode_entries(read_toml_safe(path)): + backup_existing_file(path, _legacy_backup_path()) + _strip_legacy_ucode_entries(path) + + +def revert_legacy_shared_config() -> bool: + """Undo legacy in-place edits to ``~/.codex/config.toml`` on revert. + + Codex CLI < 0.134.0 had ucode deep-merge ``profile = "ucode"``, + ``[profiles.ucode]``, and ``[model_providers.ucode-databricks]`` into the + user's real shared config, which routes every bare ``codex`` invocation + through the workspace gateway. ``ucode revert`` only restored the + per-profile file, leaving those edits in place. Surgically strip them here. + + Returns True if anything was removed. + """ + return _strip_legacy_ucode_entries(_legacy_config_path()) def _openai_model_id(model: str | None) -> str | None: @@ -256,22 +310,26 @@ def default_model(state: dict) -> str | None: The discovery list is alphabetically sorted, which can put "databricks-gpt-5" ahead of "databricks-gpt-5-5". Prefer the - highest semantic version instead. Falls back to the first - discovered entry when parsing fails. + highest semantic version instead. + + Only GPT-parseable ids are considered. Codex routes the chosen ``model`` + through the gateway as-is, so a non-GPT entry (e.g. ``moonshotai/kimi-k2.5``) + would be rejected with a Unity Catalog endpoint-name error. When no + candidate parses as GPT we return None rather than pinning an unroutable id. """ codex_models = state.get("codex_models") or [] - if not codex_models: + parsed: list[tuple[str, tuple[int, int | None, int | None, str]]] = [ + (mid, gpt) for mid in codex_models if (gpt := _parse_gpt(mid)) is not None + ] + if not parsed: return None - def _gpt_version_key(mid: str) -> tuple[int, int, int, int]: - parsed = _parse_gpt(mid) - if parsed is None: - return (0, 0, 0, 0) - major, minor, patch, suffix = parsed + def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]): + major, minor, patch, suffix = entry[1] base_bonus = 1 if not suffix else 0 return (major, minor or 0, patch or 0, base_bonus) - return max(codex_models, key=_gpt_version_key) + return max(parsed, key=_gpt_version_key)[0] def launch(state: dict, tool_args: list[str]) -> None: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d934020..c363e22 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -25,6 +25,7 @@ from ucode.agents import ( launch as launch_agent, ) +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 restore_file, set_dry_run from ucode.databricks import ( @@ -440,12 +441,17 @@ def revert() -> int: pi_settings_restored = restore_file( PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) ) + # Older Codex (< 0.134.0) had ucode edit the shared ~/.codex/config.toml in + # place; restoring the per-profile file above does not undo that. + legacy_codex_stripped = revert_legacy_shared_config() clear_state() print_heading("Revert") print_kv("Workspace", state.get("workspace") or "none") for tool, spec in TOOL_SPECS.items(): print_kv(f"{spec['display']} config", "restored" if results[tool] else "unchanged") + if legacy_codex_stripped: + print_kv("Codex shared config", "ucode entries removed") print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") for client, spec in MCP_CLIENTS.items(): print_kv( diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 6c2d64a..b84b667 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -219,6 +219,83 @@ def test_unknown_version_uses_modern_layout(self, monkeypatch): assert codex._use_legacy_layout() is False +class TestCodexRemoveLegacyProfile: + def test_drops_provider_block_on_modern_path(self, tmp_path, monkeypatch): + config_dir = tmp_path / ".codex" + config_dir.mkdir() + profile_path = config_dir / "ucode.config.toml" + legacy_path = config_dir / "config.toml" + legacy_path.write_text( + 'profile = "ucode"\n\n' + "[profiles.ucode]\n" + 'model_provider = "ucode-databricks"\n\n' + "[model_providers.ucode-databricks]\n" + 'name = "Databricks AI Gateway"\n\n' + "[model_providers.other]\n" + 'name = "keep"\n', + encoding="utf-8", + ) + backup_path = tmp_path / "codex-ucode-config.backup.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) + + doc = read_toml_safe(legacy_path) + assert "profile" not in doc + assert "ucode" not in doc.get("profiles", {}) + assert "ucode-databricks" not in doc["model_providers"] + assert doc["model_providers"]["other"]["name"] == "keep" + + +class TestCodexRevertLegacySharedConfig: + def test_strips_all_ucode_entries(self, tmp_path, monkeypatch): + config_dir = tmp_path / ".codex" + config_dir.mkdir() + profile_path = config_dir / "ucode.config.toml" + legacy_path = config_dir / "config.toml" + legacy_path.write_text( + 'profile = "ucode"\n\n' + "[profiles.ucode]\n" + 'model_provider = "ucode-databricks"\n\n' + "[profiles.other]\n" + 'model_provider = "keep"\n\n' + "[model_providers.ucode-databricks]\n" + 'name = "Databricks AI Gateway"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + + assert codex.revert_legacy_shared_config() is True + + doc = read_toml_safe(legacy_path) + assert "profile" not in doc + assert "ucode" not in doc["profiles"] + assert doc["profiles"]["other"]["model_provider"] == "keep" + assert "model_providers" not in doc + + def test_returns_false_when_no_ucode_entries(self, tmp_path, monkeypatch): + config_dir = tmp_path / ".codex" + config_dir.mkdir() + profile_path = config_dir / "ucode.config.toml" + legacy_path = config_dir / "config.toml" + legacy_path.write_text('[profiles.other]\nmodel_provider = "keep"\n', encoding="utf-8") + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + + assert codex.revert_legacy_shared_config() is False + + doc = read_toml_safe(legacy_path) + assert doc["profiles"]["other"]["model_provider"] == "keep" + + def test_returns_false_when_no_shared_config(self, tmp_path, monkeypatch): + profile_path = tmp_path / ".codex" / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + + assert codex.revert_legacy_shared_config() is False + + class TestCodexDefaultModel: def test_picks_highest_semver_over_alpha(self): state = {"codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"]} @@ -228,6 +305,18 @@ def test_picks_highest_semver_over_alpha(self): def test_none_when_no_models(self): assert codex.default_model({}) is None + def test_none_when_no_gpt_parseable_models(self): + # A workspace whose responses-capable models aren't GPT (e.g. kimi) + # must not pin an unroutable id as the Codex model. + state = {"codex_models": ["moonshotai/kimi-k2.5", "claude-sonnet-4"]} + + assert codex.default_model(state) is None + + def test_ignores_non_gpt_candidates(self): + state = {"codex_models": ["moonshotai/kimi-k2.5", "databricks-gpt-5-5"]} + + assert codex.default_model(state) == "databricks-gpt-5-5" + def test_prefers_base_over_suffixed_same_version(self): models = ["gpt-5-5-mini", "gpt-5-5", "gpt-5"] diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index a60255b..83299b1 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -107,8 +107,9 @@ def test_pi_unavailable_when_no_models(self): class TestDefaultModelForTool: - def test_codex_returns_first_model(self): - assert default_model_for_tool("codex", {"codex_models": ["c1", "c2"]}) == "c1" + def test_codex_returns_highest_gpt_model(self): + models = ["databricks-gpt-5", "databricks-gpt-5-5"] + assert default_model_for_tool("codex", {"codex_models": models}) == "databricks-gpt-5-5" def test_codex_returns_none_when_no_models(self): assert default_model_for_tool("codex", {}) is None @@ -161,9 +162,9 @@ def test_pi_returns_none_when_no_models(self): class TestResolveLaunchModel: def test_codex_default_model_used_when_no_explicit(self): - state = {"codex_models": ["c1"]} + state = {"codex_models": ["databricks-gpt-5"]} _, model = resolve_launch_model("codex", state, None) - assert model == "c1" + assert model == "databricks-gpt-5" def test_explicit_model_used_when_provided(self): _, model = resolve_launch_model("claude", {}, "my-model")