From ef46f850fc6de19f5e909a18f384ee660e012421 Mon Sep 17 00:00:00 2001 From: Anthony Ivan Date: Sun, 2 Aug 2026 19:14:53 +0800 Subject: [PATCH] Align Pi and OpenCode model families --- src/ucode/agents/__init__.py | 5 ++- src/ucode/agents/opencode.py | 29 ++++++++++++- src/ucode/agents/pi.py | 48 ++++++++++++++++++---- src/ucode/cli.py | 8 +++- src/ucode/databricks.py | 13 +++--- tests/conftest.py | 7 ++++ tests/test_agent_opencode.py | 71 ++++++++++++++++++++++++++++++-- tests/test_agent_pi.py | 80 +++++++++++++++++++++++++++++++----- tests/test_agents_init.py | 21 ++++++++++ tests/test_cli.py | 51 +++++++++++++++++++++++ tests/test_databricks.py | 19 ++++++++- tests/test_e2e.py | 8 +++- 12 files changed, 324 insertions(+), 36 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 27b0458f..213f63f6 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -363,17 +363,18 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: bool(state.get("claude_models")) or bool(state.get("codex_models")) or bool(state.get("gemini_models")) + or bool(state.get("oss_models")) ) return False _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = { "claude": ("claude",), - "opencode": ("claude", "gemini", "oss"), + "opencode": ("claude", "codex", "gemini", "oss"), "codex": ("codex",), "gemini": ("gemini",), "copilot": ("claude", "codex"), - "pi": ("claude", "codex", "gemini"), + "pi": ("claude", "codex", "gemini", "oss"), } diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b89536e5..6a9bb1b1 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"], ["provider", "databricks-oss"], ] @@ -51,13 +52,24 @@ 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/", "databricks-google/", "databricks-oss/")): + if model.startswith( + ( + "databricks-anthropic/", + "databricks-openai/", + "databricks-google/", + "databricks-oss/", + ) + ): 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}" @@ -100,6 +112,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 [] oss_models = opencode_models.get("oss") or [] @@ -125,6 +138,17 @@ def render_overlay( "models": dict.fromkeys(anthropic_models, anthropic_model_overlay), } keys.append(["provider", "databricks-anthropic"]) + if openai_models: + 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", @@ -232,6 +256,9 @@ def default_model(state: dict) -> str | None: 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 [] if gemini: return gemini[0] diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c17609..4102a865 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -7,6 +7,7 @@ - `databricks-claude` (api: anthropic-messages) → /ai-gateway/anthropic - `databricks-openai` (api: openai-responses) → /ai-gateway/codex/v1 - `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta +- `databricks-oss` (api: openai-responses) → /ai-gateway/mlflow/v1 Per-provider `compat` flags work around fields the gateway translators reject: @@ -16,10 +17,10 @@ sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header instead, which the gateway accepts. -OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via -pi today — they live behind /ai-gateway/mlflow/v1 with per-model -`max_tokens` caps that pi has no global way to honor without per-model -config we don't currently maintain. +- OSS models use the MLflow Responses route because its chat-completions stream + can end without a `finish_reason`, which Pi treats as an error. They also + carry per-model `contextWindow` and `maxTokens` from the shared token-limits + table. The bearer token is baked into the file and refreshed by a background thread while the session runs (same pattern as OpenCode/Copilot). @@ -45,6 +46,7 @@ TOKEN_REFRESH_INTERVAL_SECONDS, build_pi_base_urls, get_databricks_token, + model_token_limits, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -68,13 +70,14 @@ "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-oss", ) PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] # Old provider names earlier ucode versions wrote; cleaned up on each write so # users don't end up with stale entries pointing at routes that 400. -LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss") +LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-kimi") def is_update_available() -> tuple[str, str] | None: @@ -86,6 +89,7 @@ def _resolve_model_selector( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], ) -> str: """Return a Pi model selector in `/` form when possible.""" for name in PROVIDER_NAMES: @@ -97,9 +101,20 @@ def _resolve_model_selector( return f"databricks-openai/{model}" if model in gemini_models: return f"databricks-gemini/{model}" + if model in oss_models: + return f"databricks-oss/{model}" return model +def _oss_model_entry(model: str) -> dict: + entry: dict = {"id": model} + limits = model_token_limits(model) + if limits is not None: + entry["contextWindow"] = limits["context"] + entry["maxTokens"] = limits["output"] + return entry + + def render_overlay( model: str, token: str, @@ -107,6 +122,7 @@ def render_overlay( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.""" providers: dict = {} @@ -150,8 +166,20 @@ def render_overlay( "models": [{"id": m} for m in gemini_models], } keys.append(["providers", "databricks-gemini"]) + if oss_models: + providers["databricks-oss"] = { + "baseUrl": pi_base_urls["oss"], + "api": "openai-responses", + "apiKey": token, + "authHeader": True, + "headers": ua_headers, + "models": [_oss_model_entry(m) for m in oss_models], + } + keys.append(["providers", "databricks-oss"]) overlay: dict = { - "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models), + "model": _resolve_model_selector( + model, claude_models, codex_models, gemini_models, oss_models + ), } if providers: overlay["providers"] = providers @@ -178,6 +206,7 @@ def write_tool_config( state.get("claude_models") or {}, state.get("codex_models") or [], state.get("gemini_models") or [], + state.get("oss_models") or [], ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") @@ -206,7 +235,7 @@ def _write_settings(model_selector: str) -> None: def default_model(state: dict) -> str | None: - """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini.""" + """Prefer Claude opus → sonnet → haiku; then codex, gemini, OSS.""" claude_models = state.get("claude_models") or {} for family in ("opus", "sonnet", "haiku"): if claude_models.get(family): @@ -215,7 +244,10 @@ def default_model(state: dict) -> str | None: if codex_models: return codex_models[0] gemini_models = state.get("gemini_models") or [] - return gemini_models[0] if gemini_models else None + if gemini_models: + return gemini_models[0] + oss_models = state.get("oss_models") or [] + return oss_models[0] if oss_models else None def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b7b7a6e6..1bc518b6 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -370,8 +370,10 @@ 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_oss = fetch_all or "opencode" in tools + want_codex = ( + fetch_all or "codex" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools + ) + want_oss = fetch_all or "opencode" in tools or "pi" in tools claude_reason: str | None = None gemini_reason: str | None = None @@ -425,6 +427,8 @@ def configure_shared_state( oss_models, oss_reason = ms_oss, ms_reason 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 if oss_models: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1d32f319..22e8cfdc 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1333,7 +1333,9 @@ def discover_model_services( - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the newest matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). - - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. + - ``codex_models`` is the list of Responses-compatible ``system.ai.*gpt-*`` + ids; ``gpt-oss`` is excluded because its gateway rejects the session and + prompt-caching fields sent by Pi and OpenCode. - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. @@ -1354,7 +1356,7 @@ def discover_model_services( if candidates: claude_models[family] = candidates[0] - codex_models = [m for m in ids if "gpt-" in m] + codex_models = [m for m in ids if "gpt-" in m and "gpt-oss-" not in m] gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] @@ -2370,6 +2372,7 @@ 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", + "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", "oss": f"{workspace}/ai-gateway/mlflow/v1", } @@ -2380,17 +2383,15 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: # path (verified end-to-end). Each `api` type appends its own path suffix: # # - anthropic-messages appends `/v1/messages` - # - openai-responses appends `/responses` + # - openai-responses appends `/responses` (codex and OSS providers) # - google-generative-ai appends `/v1beta/models/{id}:streamGenerateContent` - # - openai-completions appends `/chat/completions` # # So the baseUrls below stop just before the suffix Pi will tack on. - # Compat flags applied per-provider in agents/pi.py; required for `oss` - # only (MLflow rejects `store` and `tools[].function.strict`). return { "claude": build_tool_base_url("claude", workspace), "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932b..8ee2bb13 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ from ucode.databricks import ( build_shared_base_urls, + discover_model_services, fetch_ai_gateway_claude_models, fetch_codex_models, fetch_gemini_models, @@ -56,18 +57,24 @@ def e2e_state(e2e_workspace, e2e_token): claude_models = fetch_ai_gateway_claude_models(e2e_workspace, e2e_token) gemini_models = fetch_gemini_models(e2e_workspace, e2e_token) codex_models = fetch_codex_models(e2e_workspace, e2e_token) + _, _, _, oss_models, _ = discover_model_services(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 + if oss_models: + opencode_models["oss"] = oss_models return { "workspace": e2e_workspace, "claude_models": claude_models, "gemini_models": gemini_models, "codex_models": codex_models, + "oss_models": oss_models, "opencode_models": opencode_models, "base_urls": build_shared_base_urls(e2e_workspace), "managed_configs": {}, diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 71ca2bc6..a8bb8499 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", "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -49,6 +50,11 @@ def test_gemini_provider_added_when_models_present(self): overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert "databricks-google" in overlay["provider"] + def test_openai_provider_added_when_models_present(self): + models = {"openai": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + assert "databricks-openai" in overlay["provider"] + def test_oss_provider_added_when_models_present(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} overlay, _ = opencode.render_overlay( @@ -63,11 +69,20 @@ def test_oss_provider_uses_ai_sdk_openai_package(self): ) assert overlay["provider"]["databricks-oss"]["npm"] == "@ai-sdk/openai" - def test_both_providers_when_both_present(self): - models = {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]} + def test_all_four_providers_when_all_present(self): + models = { + "anthropic": ["claude-sonnet"], + "openai": ["system.ai.gpt-5"], + "gemini": ["gemini-2"], + "oss": ["system.ai.kimi-k2-7-code"], + } overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) - assert "databricks-anthropic" in overlay["provider"] - assert "databricks-google" in overlay["provider"] + assert set(overlay["provider"]) == { + "databricks-anthropic", + "databricks-openai", + "databricks-google", + "databricks-oss", + } def test_no_provider_key_when_no_models(self): overlay, _ = opencode.render_overlay("model", "tok", _base_urls(), {}) @@ -85,6 +100,17 @@ 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(self): + models = {"openai": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + options = overlay["provider"]["databricks-openai"]["options"] + assert options["baseURL"] == f"{WS}/ai-gateway/codex/v1" + + def test_openai_provider_uses_ai_sdk_openai_package(self): + models = {"openai": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + assert overlay["provider"]["databricks-openai"]["npm"] == "@ai-sdk/openai" + def test_oss_base_url(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} overlay, _ = opencode.render_overlay( @@ -149,6 +175,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": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + model_headers = overlay["provider"]["databricks-openai"]["models"]["system.ai.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. @@ -172,6 +208,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": ["system.ai.gpt-5"]} + _, keys = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + assert ["provider", "databricks-openai"] in keys + def test_managed_keys_include_oss_provider(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} _, keys = opencode.render_overlay("system.ai.kimi-k2-7-code", "tok", _base_urls(), models) @@ -194,6 +235,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": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-5", "tok", _base_urls(), models) + assert overlay["model"] == "databricks-openai/system.ai.gpt-5" + + def test_preserves_existing_openai_provider_prefix(self): + models = {"openai": ["system.ai.gpt-5"]} + overlay, _ = opencode.render_overlay( + "databricks-openai/system.ai.gpt-5", "tok", _base_urls(), models + ) + assert overlay["model"] == "databricks-openai/system.ai.gpt-5" + def test_prefixes_oss_model_with_provider_id(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} overlay, _ = opencode.render_overlay( @@ -314,6 +367,16 @@ def test_falls_back_to_gemini(self): state = {"opencode_models": {"anthropic": [], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "gemini-2" + def test_falls_back_to_openai_before_gemini(self): + state = { + "opencode_models": { + "anthropic": [], + "openai": ["system.ai.gpt-5"], + "gemini": ["gemini-2"], + } + } + assert opencode.default_model(state) == "system.ai.gpt-5" + def test_falls_back_to_oss(self): state = { "opencode_models": { diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb3..d8c8f704 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -17,6 +17,7 @@ def _base_urls() -> dict[str, str]: "claude": f"{WS}/ai-gateway/anthropic", "openai": f"{WS}/ai-gateway/codex/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -26,6 +27,7 @@ def _empty() -> dict: "claude_models": {}, "codex_models": [], "gemini_models": [], + "oss_models": [], } @@ -39,6 +41,7 @@ def _overlay(model: str, token: str = "tok", **kwargs): bundle["claude_models"], bundle["codex_models"], bundle["gemini_models"], + bundle["oss_models"], ) @@ -81,22 +84,30 @@ def test_gemini_provider_uses_google_generative_ai(self): assert provider["api"] == "google-generative-ai" assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta" - def test_all_three_providers_when_all_present(self): + def test_oss_provider_uses_openai_responses(self): + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + provider = overlay["providers"]["databricks-oss"] + assert provider["api"] == "openai-responses" + assert provider["baseUrl"] == f"{WS}/ai-gateway/mlflow/v1" + + def test_all_four_providers_when_all_present(self): overlay, _ = _overlay( "claude-sonnet", claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) assert set(overlay["providers"].keys()) == { "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-oss", } class TestRenderOverlayUserAgent: - def test_user_agent_set_on_all_three_providers(self, monkeypatch): + def test_user_agent_set_on_all_four_providers(self, monkeypatch): monkeypatch.setattr(pi, "ucode_version", lambda: "0.1.0") monkeypatch.setattr(pi, "agent_version", lambda binary: "0.74.0") overlay, _ = _overlay( @@ -104,9 +115,15 @@ def test_user_agent_set_on_all_three_providers(self, monkeypatch): claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) expected = "ucode/0.1.0 pi/0.74.0" - for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + for name in ( + "databricks-claude", + "databricks-openai", + "databricks-gemini", + "databricks-oss", + ): assert overlay["providers"][name]["headers"]["User-Agent"] == expected @@ -119,15 +136,17 @@ def test_claude_disables_eager_tool_input_streaming(self): compat = overlay["providers"]["databricks-claude"]["compat"] assert compat["supportsEagerToolInputStreaming"] is False - def test_openai_and_gemini_have_no_compat_flags(self): + def test_openai_gemini_and_oss_have_no_compat_flags(self): # Their gateway routes accept pi's request shape as-is. overlay, _ = _overlay( "gpt-5", codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) assert "compat" not in overlay["providers"]["databricks-openai"] assert "compat" not in overlay["providers"]["databricks-gemini"] + assert "compat" not in overlay["providers"]["databricks-oss"] class TestRenderOverlayAuthAndModels: @@ -143,8 +162,14 @@ def test_auth_header_flag_set_on_all_providers(self): claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) - for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + for name in ( + "databricks-claude", + "databricks-openai", + "databricks-gemini", + "databricks-oss", + ): assert overlay["providers"][name]["authHeader"] is True def test_claude_models_listed(self): @@ -163,6 +188,17 @@ def test_gemini_models_listed(self): ids = {m["id"] for m in overlay["providers"]["databricks-gemini"]["models"]} assert ids == {"gemini-2", "gemini-2-pro"} + def test_oss_models_listed_with_known_limits(self): + overlay, _ = _overlay( + "system.ai.glm-5-2", + oss_models=["system.ai.kimi-k2-7-code", "system.ai.glm-5-2"], + ) + entries = {m["id"]: m for m in overlay["providers"]["databricks-oss"]["models"]} + assert set(entries) == {"system.ai.kimi-k2-7-code", "system.ai.glm-5-2"} + assert entries["system.ai.glm-5-2"]["contextWindow"] == 200_000 + assert entries["system.ai.glm-5-2"]["maxTokens"] == 25_000 + assert "maxTokens" not in entries["system.ai.kimi-k2-7-code"] + class TestRenderOverlayManagedKeys: def test_managed_keys_include_model(self): @@ -175,8 +211,14 @@ def test_managed_keys_include_each_provider_present(self): claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.kimi-k2-7-code"], ) - for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + for name in ( + "databricks-claude", + "databricks-openai", + "databricks-gemini", + "databricks-oss", + ): assert ["providers", name] in keys @@ -193,6 +235,10 @@ def test_prefixes_gemini_model(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) assert overlay["model"] == "databricks-gemini/gemini-2" + def test_prefixes_oss_model(self): + overlay, _ = _overlay("system.ai.kimi-k2-7-code", oss_models=["system.ai.kimi-k2-7-code"]) + assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" + def test_preserves_already_prefixed_model(self): overlay, _ = _overlay( "databricks-claude/claude-sonnet", @@ -228,10 +274,22 @@ def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert pi.default_model(state) == "gemini-2" + def test_falls_back_to_oss(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.kimi-k2-7-code"], + } + assert pi.default_model(state) == "system.ai.kimi-k2-7-code" + def test_returns_none_when_empty(self): assert pi.default_model({}) is None assert ( - pi.default_model({"claude_models": {}, "codex_models": [], "gemini_models": []}) is None + pi.default_model( + {"claude_models": {}, "codex_models": [], "gemini_models": [], "oss_models": []} + ) + is None ) @@ -283,6 +341,7 @@ def _state(self, **overrides) -> dict: "claude_models": {"sonnet": "claude-sonnet"}, "codex_models": [], "gemini_models": [], + "oss_models": [], "managed_configs": {}, } state.update(overrides) @@ -296,6 +355,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc "databricks-claude": {"old": True}, "databricks-openai": {"old": True}, "databricks-gemini": {"old": True}, + "databricks-oss": {"old": True}, "user-provider": {"keep": True}, } } @@ -315,7 +375,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): """Earlier ucode versions wrote `databricks-anthropic`, `databricks-codex`, - and `databricks-oss` providers. They must be stripped on the next write + and `databricks-kimi` providers. They must be stripped on the next write so users don't end up with stale entries pointing at routes that 400.""" pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) @@ -325,7 +385,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): "providers": { "databricks-anthropic": {"api": "anthropic-messages"}, "databricks-codex": {"api": "openai-responses"}, - "databricks-oss": {"api": "openai-completions"}, + "databricks-kimi": {"api": "openai-responses"}, } } ), @@ -339,7 +399,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") written_providers = json.loads(config_file.read_text()).get("providers", {}) - for legacy in ("databricks-anthropic", "databricks-codex", "databricks-oss"): + for legacy in ("databricks-anthropic", "databricks-codex", "databricks-kimi"): assert legacy not in written_providers assert "databricks-claude" in written_providers diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 97acf307..fce37226 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -189,6 +189,9 @@ def test_pi_available_with_codex(self): def test_pi_available_with_gemini(self): assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "pi") is True + def test_pi_available_with_oss(self): + assert check_gateway_endpoint({"oss_models": ["system.ai.kimi-k2-7-code"]}, "pi") is True + def test_pi_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "pi") is False @@ -231,6 +234,15 @@ def test_opencode_falls_back_to_gemini(self): state = {"opencode_models": {"gemini": ["gemini-2"]}} assert default_model_for_tool("opencode", state) == "gemini-2" + def test_opencode_falls_back_to_openai_before_gemini(self): + state = { + "opencode_models": { + "openai": ["system.ai.gpt-5"], + "gemini": ["gemini-2"], + } + } + assert default_model_for_tool("opencode", state) == "system.ai.gpt-5" + def test_pi_prefers_claude_opus(self): state = {"claude_models": {"opus": "o4", "sonnet": "s4"}, "codex_models": ["c"]} assert default_model_for_tool("pi", state) == "o4" @@ -243,6 +255,15 @@ def test_pi_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert default_model_for_tool("pi", state) == "gemini-2" + def test_pi_falls_back_to_oss(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.kimi-k2-7-code"], + } + assert default_model_for_tool("pi", state) == "system.ai.kimi-k2-7-code" + def test_pi_returns_none_when_no_models(self): assert default_model_for_tool("pi", {}) is None diff --git a/tests/test_cli.py b/tests/test_cli.py index f5119c2d..9e8006a9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1858,6 +1858,57 @@ def test_skips_purge_when_workspace_unchanged(self, monkeypatch): assert purge_calls == [] +class TestConfigureSharedStateAlignedModelFamilies: + WS = "https://workspace.databricks.com" + + @staticmethod + def _stub(monkeypatch): + import ucode.cli as cli_mod + + monkeypatch.setattr(cli_mod, "load_state", lambda: {}) + monkeypatch.setattr(cli_mod, "save_state", lambda state: None) + monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda workspace: workspace) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda workspace: None) + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda workspace, profile: "token") + monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda workspace, token: None) + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda workspace, token: ( + {"sonnet": "system.ai.claude-sonnet-4-6"}, + ["system.ai.gpt-5"], + ["system.ai.gemini-2-5-pro"], + ["system.ai.kimi-k2-7-code"], + None, + ), + ) + monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda workspace: {}) + return cli_mod + + def test_opencode_discovers_openai_alongside_existing_families(self, monkeypatch): + cli_mod = self._stub(monkeypatch) + + state = cli_mod.configure_shared_state(self.WS, tools=["opencode"]) + + assert state["opencode_models"] == { + "anthropic": ["system.ai.claude-sonnet-4-6"], + "openai": ["system.ai.gpt-5"], + "gemini": ["system.ai.gemini-2-5-pro"], + "oss": ["system.ai.kimi-k2-7-code"], + } + + def test_pi_discovers_oss_alongside_existing_families(self, monkeypatch): + cli_mod = self._stub(monkeypatch) + + state = cli_mod.configure_shared_state(self.WS, tools=["pi"]) + + assert state["claude_models"] == {"sonnet": "system.ai.claude-sonnet-4-6"} + assert state["codex_models"] == ["system.ai.gpt-5"] + assert state["gemini_models"] == ["system.ai.gemini-2-5-pro"] + assert state["oss_models"] == ["system.ai.kimi-k2-7-code"] + + class TestConfigureSharedStateSkipDiscovery: """With skip_model_discovery (provider mode), the heavy family discovery is skipped; only a single web-search model is fetched, and existing model lists diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7e1a73a1..e5f6e2f7 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -20,6 +20,7 @@ build_auth_token_argv, build_databricks_cli_env, build_opencode_base_urls, + build_pi_base_urls, build_shared_base_urls, build_skills_mcp_url, build_tool_base_url, @@ -96,9 +97,19 @@ def test_unsupported_tool_raises(self): class TestBuildOpencodeBaseUrls: - def test_returns_anthropic_gemini_and_oss(self): + def test_returns_all_supported_model_families(self): urls = build_opencode_base_urls(WS) assert urls["anthropic"] == f"{WS}/ai-gateway/anthropic/v1" + assert urls["openai"] == f"{WS}/ai-gateway/codex/v1" + assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" + assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" + + +class TestBuildPiBaseUrls: + def test_returns_all_supported_model_families(self): + urls = build_pi_base_urls(WS) + assert urls["claude"] == f"{WS}/ai-gateway/anthropic" + assert urls["openai"] == f"{WS}/ai-gateway/codex/v1" assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" @@ -197,6 +208,7 @@ def test_buckets_families_by_name(self, monkeypatch): _model_service("system.ai.claude-opus-4-8"), _model_service("system.ai.claude-sonnet-4-6"), _model_service("system.ai.gpt-5"), + _model_service("system.ai.gpt-oss-120b"), _model_service("system.ai.gemini-2-5-flash"), _model_service("system.ai.gemini-3-5-flash"), _model_service("system.ai.kimi-k2-7-code"), @@ -217,6 +229,8 @@ def test_buckets_families_by_name(self, monkeypatch): "opus": "system.ai.claude-opus-4-8", "sonnet": "system.ai.claude-sonnet-4-6", } + # gpt-oss is not Responses-compatible with Pi/OpenCode and must not + # leak into the Codex bucket just because its name contains `gpt-`. assert codex == ["system.ai.gpt-5"] # Gemini ordered newest-first via the shared sort key. assert gemini[0] == "system.ai.gemini-3-5-flash" @@ -224,11 +238,12 @@ def test_buckets_families_by_name(self, monkeypatch): assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] def test_oss_allowlist_drops_unsupported_families(self, monkeypatch): - # Only kimi/glm are allowlisted; other families are dropped. + # Only kimi/glm are allowlisted; other families, including gpt-oss, are dropped. payload = { "model_services": [ _model_service("system.ai.glm-5-2"), _model_service("system.ai.kimi-k2-7-code"), + _model_service("system.ai.gpt-oss-20b"), _model_service("system.ai.qwen-3-coder"), _model_service("system.ai.deepseek-v3"), _model_service("system.ai.gte-large-embed"), diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 025ced4d..e520b7d3 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -711,7 +711,7 @@ def test_does_not_write_settings_json_for_auth( class TestOpencodeLaunch: - """Run opencode against every available opencode model (anthropic + gemini).""" + """Run opencode against every available model across all four providers.""" # Models that hang opencode well past 180s on the staging gateway with # no stderr beyond the initial `> build · ` line, while every @@ -899,6 +899,8 @@ def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: out.append(("codex", model)) for model in e2e_state.get("gemini_models") or []: out.append(("gemini", model)) + for model in e2e_state.get("oss_models") or []: + out.append(("oss", model)) return out def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token): @@ -916,10 +918,14 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa pi_home = tmp_path / "pi-home" pi_dir = pi_home / ".pi" / "agent" config_path = pi_dir / "models.json" + settings_path = pi_dir / "settings.json" backup_path = tmp_path / "pi-models.backup.json" + settings_backup_path = tmp_path / "pi-settings.backup.json" monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", settings_path) monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", settings_backup_path) failures = [] for family, model in models: