diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index b07e18c..b8f37cc 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -68,13 +68,6 @@ _GPT_RE = re.compile(r"(?:databricks-)?gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") -# These models should use the Databricks ID, not the OpenAI ID, as the OpenAI -# ID is incompatible with Codex. -CODEX_OPENAI_ID_INCOMPATIBLE_MODELS = { - "databricks-gpt-5-2-codex", - "databricks-gpt-5-4-nano", -} - def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) @@ -274,30 +267,6 @@ def revert_legacy_shared_config() -> bool: return _strip_legacy_ucode_entries(_legacy_config_path()) -def _openai_model_id(model: str | None) -> str | None: - """Map Databricks GPT endpoint ids to OpenAI model ids for Codex metadata.""" - parsed = _parse_gpt(model) - if parsed is None: - return model - major, minor, patch, suffix = parsed - version = str(major) - if minor is not None: - version += f".{minor}" - if patch is not None: - version += f".{patch}" - return f"gpt-{version}{suffix}" - - -def _codex_model_id(model: str | None) -> str | None: - # UC model-services ids (`system.ai.gpt-5`) route by name through the - # gateway, so they must be sent verbatim — not rewritten to an OpenAI id. - if model and model.startswith("system.ai."): - return model - if model in CODEX_OPENAI_ID_INCOMPATIBLE_MODELS: - return model - return _openai_model_id(model) - - def _parse_gpt(model: str | None) -> tuple[int, int | None, int | None, str] | None: if not model: return None @@ -322,8 +291,12 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non workspace = state["workspace"] # With a Model Provider Service the gateway routes by header and Codex sends # its own canonical model name (e.g. `gpt-5`) — leave `model` unset so no - # Databricks endpoint id is pinned. - chosen_model = None if provider else _codex_model_id(model or default_model(state)) + # Databricks endpoint id is pinned. Otherwise pin the discovered endpoint id + # verbatim: the gateway routes by that exact name (whether `databricks-gpt-5` + # from the AI Gateway listing or `system.ai.gpt-5` from UC model-services), so + # rewriting it to an OpenAI id (`gpt-5`) makes the gateway resolve a + # non-existent `system.ai.*` alias and 404. + chosen_model = None if provider else (model or default_model(state)) databricks_profile = state.get("profile") if _use_legacy_layout(): diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3d4fb87..baf1ad0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -629,7 +629,7 @@ def _use_databricks() -> dict: return _use_databricks() if not names: # Feature is on but no service matches this tool's provider type. - print_note(f"No model provider services available for {display}; using Databricks models.") + print_note(f"Using Databricks models for {display}.") return _use_databricks() choice = prompt_for_selection( diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index fb6cdae..10d0f9e 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -109,7 +109,10 @@ 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): + def test_pins_discovered_databricks_model_id_verbatim(self, tmp_path, monkeypatch): + # The gateway routes by the discovered endpoint name, so the id is + # written as-is (not rewritten to an OpenAI id like `gpt-5.5`, which the + # gateway would resolve to a non-existent `system.ai.*` alias and 404). 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) @@ -122,11 +125,9 @@ def test_writes_openai_model_id_for_databricks_gpt_endpoint(self, tmp_path, monk ) doc = read_toml_safe(config_path) - assert doc["model"] == "gpt-5.5" + assert doc["model"] == "databricks-gpt-5-5" - def test_preserves_databricks_model_id_when_openai_id_is_incompatible( - self, tmp_path, monkeypatch - ): + def test_pins_uc_model_services_id_verbatim(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) @@ -135,12 +136,11 @@ def test_preserves_databricks_model_id_when_openai_id_is_incompatible( monkeypatch.setattr(codex, "save_state", lambda state: None) codex.write_tool_config( - {"workspace": WS, "codex_models": ["databricks-gpt-5-2-codex"]}, - "databricks-gpt-5-2-codex", + {"workspace": WS, "codex_models": ["system.ai.gpt-5", "system.ai.gpt-5-5"]} ) doc = read_toml_safe(config_path) - assert doc["model"] == "databricks-gpt-5-2-codex" + assert doc["model"] == "system.ai.gpt-5-5" def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" @@ -522,28 +522,10 @@ def test_codex_default_model_wins_over_allowlist(self): } assert codex.default_model(state) == "admin-chosen-default" - 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("databricks-gpt-4o") == "gpt-4o" - assert codex._openai_model_id("served-models/databricks-gpt-5-5") == "gpt-5.5" - assert codex._openai_model_id("gpt-5.5") == "gpt-5.5" - - def test_codex_model_id_preserves_openai_incompatible_models(self): - assert codex._codex_model_id("databricks-gpt-5-2-codex") == "databricks-gpt-5-2-codex" - assert codex._codex_model_id("databricks-gpt-5-4-nano") == "databricks-gpt-5-4-nano" - - def test_codex_model_id_passes_model_services_id_verbatim(self): - # UC model-services ids route by name, so they must not be rewritten - # to the OpenAI id form. - assert codex._codex_model_id("system.ai.gpt-5") == "system.ai.gpt-5" - assert codex._codex_model_id("system.ai.gpt-5-2-codex") == "system.ai.gpt-5-2-codex" - def test_default_model_selects_model_services_gpt(self): models = ["system.ai.gpt-5", "system.ai.gpt-5-5", "system.ai.claude-opus-4-8"] assert codex.default_model({"codex_models": models}) == "system.ai.gpt-5-5" - assert codex._codex_model_id("databricks-gpt-5-5") == "gpt-5.5" class TestCodexValidateCmd: