Fix ucode revert for legacy Codex shared-config edits (#150)

`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

Author: AarushiShah-db

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:

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(

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"]
 

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")