Fix Codex GPT model selection

src/ucode/agents/codex.py

@@ -177,9 +177,25 @@ def _remove_legacy_ucode_profile() -> None:
         write_toml_file(path, doc)
 
 
+def _openai_model_id(model: str | None) -> str | None:
+    """Map Databricks GPT endpoint ids to OpenAI model ids for Codex metadata."""
+    if not model:
+        return model
+    match = re.fullmatch(r"databricks-gpt-(\d+)(?:-(\d+))?(?:-(\d+))?((?:-.+)?)", model)
+    if not match:
+        return model
+    major, minor, patch, suffix = match.groups()
+    version = major
+    if minor is not None:
+        version += f".{minor}"
+    if patch is not None:
+        version += f".{patch}"
+    return f"gpt-{version}{suffix or ''}"
+
+
 def write_tool_config(state: dict, model: str | None = None) -> dict:
     workspace = state["workspace"]
-    chosen_model = model or default_model(state)
+    chosen_model = _openai_model_id(model or default_model(state))
     databricks_profile = state.get("profile")
 
     if _use_legacy_layout():
@@ -208,8 +224,36 @@ def write_tool_config(state: dict, model: str | None = None) -> dict:
 
 
 def default_model(state: dict) -> str | None:
+    """Pick the newest GPT model when multiple are available.
+
+    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.
+    """
     codex_models = state.get("codex_models") or []
-    return codex_models[0] if codex_models else None
+    if not codex_models:
+        return None
+
+    def _gpt_version_key(mid: str) -> tuple[int, int, int, int]:
+        try:
+            name = mid.split("/")[-1]
+            m = re.search(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?", name)
+            if not m:
+                return (0, 0, 0, 0)
+            major = int(m.group(1) or 0)
+            minor = int(m.group(2) or 0)
+            patch = int(m.group(3) or 0)
+            suffix = name[m.end() :]
+            base_bonus = 1 if not suffix else 0
+            return (major, minor, patch, base_bonus)
+        except Exception:
+            return (0, 0, 0, 0)
+
+    try:
+        return max(codex_models, key=_gpt_version_key)
+    except ValueError:
+        return codex_models[0]
 
 
 def launch(state: dict, tool_args: list[str]) -> None:

tests/test_agent_codex.py

@@ -91,6 +91,21 @@ def test_writes_ucode_profile_config_file(self, tmp_path, monkeypatch):
         assert doc["model"] == "gpt-5"
         assert "profiles" not in doc
 
+    def test_writes_openai_model_id_for_databricks_gpt_endpoint(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)
+        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": ["databricks-gpt-5", "databricks-gpt-5-5"]}
+        )
+
+        doc = read_toml_safe(config_path)
+        assert doc["model"] == "gpt-5.5"
+
     def test_removes_legacy_ucode_profile_from_shared_config(self, tmp_path, monkeypatch):
         config_dir = tmp_path / ".codex"
         config_dir.mkdir()
@@ -187,12 +202,24 @@ def test_unknown_version_uses_modern_layout(self, monkeypatch):
 
 
 class TestCodexDefaultModel:
-    def test_returns_first_codex_model(self):
-        assert codex.default_model({"codex_models": ["gpt-5", "gpt-4o"]}) == "gpt-5"
+    def test_picks_highest_semver_over_alpha(self):
+        state = {"codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"]}
+
+        assert codex.default_model(state) == "databricks-gpt-5-5"
 
     def test_none_when_no_models(self):
         assert codex.default_model({}) is None
 
+    def test_prefers_base_over_suffixed_same_version(self):
+        models = ["gpt-5-5-mini", "gpt-5-5", "gpt-5"]
+
+        assert codex.default_model({"codex_models": models}) == "gpt-5-5"
+
+    def test_openai_model_id_maps_databricks_naming(self):
+        assert codex._openai_model_id("databricks-gpt-5-5") == "gpt-5.5"
+        assert codex._openai_model_id("databricks-gpt-5-5-mini") == "gpt-5.5-mini"
+        assert codex._openai_model_id("gpt-5.5") == "gpt-5.5"
+
 
 class TestCodexValidateCmd:
     def test_starts_with_binary(self):