diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 7c3ef73..38e0c9e 100644 --- a/src/ucode/agents/codex.py +++ b/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: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 3f6ac35..2729776 100644 --- a/tests/test_agent_codex.py +++ b/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):