diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b94e855..63fe1ed 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -304,7 +304,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = { "claude": ("claude",), - "opencode": ("claude", "gemini"), + "opencode": ("claude", "codex", "gemini"), "codex": ("codex",), "gemini": ("gemini",), "copilot": ("claude", "codex"), diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 8792625..8b2f319 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -40,6 +40,7 @@ PROVIDER_KEYS: list[list[str]] = [ ["provider", "databricks-anthropic"], + ["provider", "databricks-openai"], ["provider", "databricks-google"], ] @@ -50,13 +51,21 @@ def is_update_available() -> tuple[str, str] | None: def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str: """Return an OpenCode model selector in provider/model form when possible.""" - if model.startswith("databricks-anthropic/") or model.startswith("databricks-google/"): + if ( + model.startswith("databricks-anthropic/") + or model.startswith("databricks-openai/") + or model.startswith("databricks-google/") + ): return model anthropic_models = opencode_models.get("anthropic") or [] if model in anthropic_models: return f"databricks-anthropic/{model}" + openai_models = opencode_models.get("openai") or [] + if model in openai_models: + return f"databricks-openai/{model}" + gemini_models = opencode_models.get("gemini") or [] if model in gemini_models: return f"databricks-google/{model}" @@ -81,6 +90,7 @@ def render_overlay( } anthropic_models = opencode_models.get("anthropic") or [] + openai_models = opencode_models.get("openai") or [] gemini_models = opencode_models.get("gemini") or [] providers: dict = {} @@ -105,6 +115,21 @@ def render_overlay( "models": dict.fromkeys(anthropic_models, anthropic_model_overlay), } keys.append(["provider", "databricks-anthropic"]) + if openai_models: + # @ai-sdk/openai points at the Databricks codex gateway + # (/ai-gateway/codex/v1), the same OpenAI Responses-API path Pi's + # `databricks-openai` provider uses. The AI SDK's openai provider + # negotiates the Responses API there, so GPT models route correctly. + providers["databricks-openai"] = { + "npm": "@ai-sdk/openai", + "options": { + "baseURL": opencode_base_urls["openai"], + "apiKey": token, + "headers": auth_headers, + }, + "models": {m: {"headers": ua_header} for m in openai_models}, + } + keys.append(["provider", "databricks-openai"]) if gemini_models: providers["databricks-google"] = { "npm": "@ai-sdk/google", @@ -192,10 +217,16 @@ def remove_mcp_server_config(name: str) -> bool: def default_model(state: dict) -> str | None: + # Preference order mirrors Pi (`agents/pi.py:default_model`): + # Claude → OpenAI → Gemini. Picks the first id in each bucket; bucket + # population order in `cli.py` decides which specific id wins. opencode_models = state.get("opencode_models") or {} anthropic = opencode_models.get("anthropic") or [] if anthropic: return anthropic[0] + openai = opencode_models.get("openai") or [] + if openai: + return openai[0] gemini = opencode_models.get("gemini") or [] return gemini[0] if gemini else None diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0d60916..7343cbb 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -249,7 +249,9 @@ def configure_shared_state( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools ) want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools - want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools + want_codex = ( + fetch_all or "codex" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools + ) claude_reason: str | None = None gemini_reason: str | None = None @@ -278,6 +280,8 @@ def configure_shared_state( opencode_models: dict[str, list[str]] = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) + if codex_models: + opencode_models["openai"] = codex_models if gemini_models: opencode_models["gemini"] = gemini_models diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 574d906..4ed0cea 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1586,6 +1586,11 @@ def build_tool_base_url(tool: str, workspace: str) -> str: def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", + # codex gateway already includes /v1, no extra suffix needed. This is + # the same path Pi's `databricks-openai` provider uses (see + # `build_pi_base_urls`); @ai-sdk/openai negotiates the responses API + # against it. + "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", } diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932..71737d0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -60,6 +60,8 @@ def e2e_state(e2e_workspace, e2e_token): opencode_models: dict = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) + if codex_models: + opencode_models["openai"] = codex_models if gemini_models: opencode_models["gemini"] = gemini_models diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0f32f4d..d836e0f 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -13,6 +13,7 @@ def _base_urls() -> dict[str, str]: return { "anthropic": f"{WS}/ai-gateway/anthropic/v1", + "openai": f"{WS}/ai-gateway/codex/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", } @@ -54,6 +55,22 @@ def test_both_providers_when_both_present(self): assert "databricks-anthropic" in overlay["provider"] assert "databricks-google" in overlay["provider"] + def test_openai_provider_added_when_models_present(self): + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + assert "databricks-openai" in overlay["provider"] + + def test_all_three_providers_when_all_present(self): + models = { + "anthropic": ["claude-sonnet"], + "openai": ["databricks-gpt-5"], + "gemini": ["gemini-2"], + } + overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + assert "databricks-anthropic" in overlay["provider"] + assert "databricks-openai" in overlay["provider"] + assert "databricks-google" in overlay["provider"] + def test_no_provider_key_when_no_models(self): overlay, _ = opencode.render_overlay("model", "tok", _base_urls(), {}) assert "provider" not in overlay @@ -70,6 +87,20 @@ def test_gemini_base_url(self): options = overlay["provider"]["databricks-google"]["options"] assert options["baseURL"] == f"{WS}/ai-gateway/gemini/v1beta" + def test_openai_base_url_points_at_codex_gateway(self): + # Matches Pi's `databricks-openai` provider path so OpenCode's + # @ai-sdk/openai negotiates the Responses API against the same + # Databricks codex endpoint that already works for Pi. + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + options = overlay["provider"]["databricks-openai"]["options"] + assert options["baseURL"] == f"{WS}/ai-gateway/codex/v1" + + def test_openai_uses_ai_sdk_openai_npm_package(self): + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + assert overlay["provider"]["databricks-openai"]["npm"] == "@ai-sdk/openai" + def test_token_in_api_key(self): models = {"anthropic": ["claude-sonnet"]} overlay, _ = opencode.render_overlay("claude-sonnet", "mytoken", _base_urls(), models) @@ -111,6 +142,16 @@ def test_user_agent_header_gemini(self, monkeypatch): model_headers = overlay["provider"]["databricks-google"]["models"]["gemini-2"]["headers"] assert model_headers["User-Agent"] == "ucode/0.1.0 opencode/0.74.0" + def test_user_agent_header_openai(self, monkeypatch): + monkeypatch.setattr(opencode, "ucode_version", lambda: "0.1.0") + monkeypatch.setattr(opencode, "agent_version", lambda binary: "0.74.0") + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + model_headers = overlay["provider"]["databricks-openai"]["models"]["databricks-gpt-5"][ + "headers" + ] + assert model_headers["User-Agent"] == "ucode/0.1.0 opencode/0.74.0" + def test_provider_level_headers_only_authorization(self, monkeypatch): # Sanity: provider-level headers should NOT include User-Agent (since # it's clobbered there) — only Authorization. @@ -134,6 +175,11 @@ def test_managed_keys_include_gemini_provider(self): _, keys = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert ["provider", "databricks-google"] in keys + def test_managed_keys_include_openai_provider(self): + models = {"openai": ["databricks-gpt-5"]} + _, keys = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + assert ["provider", "databricks-openai"] in keys + def test_anthropic_models_listed(self): models = {"anthropic": ["claude-sonnet", "claude-haiku"]} overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) @@ -151,6 +197,18 @@ def test_prefixes_gemini_model_with_provider_id(self): overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert overlay["model"] == "databricks-google/gemini-2" + def test_prefixes_openai_model_with_provider_id(self): + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5", "tok", _base_urls(), models) + assert overlay["model"] == "databricks-openai/databricks-gpt-5" + + def test_preserves_existing_provider_prefix_openai(self): + models = {"openai": ["databricks-gpt-5"]} + overlay, _ = opencode.render_overlay( + "databricks-openai/databricks-gpt-5", "tok", _base_urls(), models + ) + assert overlay["model"] == "databricks-openai/databricks-gpt-5" + class TestMcpServerConfig: def test_builds_remote_server_entry_with_oauth_token_env_header(self): @@ -264,6 +322,17 @@ def test_prefers_anthropic(self): state = {"opencode_models": {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "claude-sonnet" + def test_falls_back_to_openai_before_gemini(self): + # Mirrors Pi's preference: Claude > OpenAI > Gemini. + state = { + "opencode_models": { + "anthropic": [], + "openai": ["databricks-gpt-5"], + "gemini": ["gemini-2"], + } + } + assert opencode.default_model(state) == "databricks-gpt-5" + def test_falls_back_to_gemini(self): state = {"opencode_models": {"anthropic": [], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "gemini-2" diff --git a/tests/test_cli.py b/tests/test_cli.py index cf156bc..7006f2f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1113,3 +1113,46 @@ def test_skips_purge_when_workspace_unchanged(self, monkeypatch): cli_mod.configure_shared_state("https://same.databricks.com") assert purge_calls == [] + + +class TestConfigureSharedStateOpencodeDiscovery: + """`ucode configure --tools opencode` must surface GPT models alongside + Claude and Gemini. That requires running codex discovery for the opencode + tool and bucketing the result under `opencode_models["openai"]` so + `agents/opencode.py:render_overlay` can wire the `databricks-openai` provider.""" + + WS = "https://workspace.databricks.com" + + @staticmethod + def _stub_with_codex_models(monkeypatch, codex_models): + import ucode.cli as cli_mod + + monkeypatch.setattr(cli_mod, "load_state", lambda: {}) + monkeypatch.setattr(cli_mod, "save_state", lambda s: None) + monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) + monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) + monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], None)) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) + monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: (codex_models, None)) + monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) + return cli_mod + + def test_opencode_tools_triggers_codex_discovery(self, monkeypatch): + cli_mod = self._stub_with_codex_models(monkeypatch, ["databricks-gpt-5"]) + + state = cli_mod.configure_shared_state(self.WS, tools=["opencode"]) + + assert state["codex_models"] == ["databricks-gpt-5"] + assert state["opencode_models"]["openai"] == ["databricks-gpt-5"] + + def test_opencode_skips_openai_bucket_when_no_codex_models(self, monkeypatch): + cli_mod = self._stub_with_codex_models(monkeypatch, []) + + state = cli_mod.configure_shared_state(self.WS, tools=["opencode"]) + + assert "openai" not in state.get("opencode_models", {})